molstar
Version:
A comprehensive macromolecular library.
139 lines (138 loc) • 7.56 kB
JavaScript
;
/**
* Copyright (c) 2018-2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
*
* @author Alexander Rose <alexander.rose@weirdbyte.de>
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.GaussianDensityCPU = void 0;
var tslib_1 = require("tslib");
var geometry_1 = require("../../geometry");
var linear_algebra_1 = require("../../linear-algebra");
var int_1 = require("../../../mol-data/int");
var approx_1 = require("../../approx");
function GaussianDensityCPU(ctx, position, box, radius, props) {
return tslib_1.__awaiter(this, void 0, void 0, function () {
function accumulateRange(begI, endI) {
for (var i = begI; i < endI; ++i) {
var j = int_1.OrderedSet.getAt(indices, i);
var vx = x[j], vy = y[j], vz = z[j];
var rad = radii[i];
var rSq = rad * rad;
var rSqInv = 1 / rSq;
var r2 = rad * 2;
var r2sq = r2 * r2;
// Number of grid points, round this up...
var ng = Math.ceil(r2 * scaleFactor);
// Center of the atom, mapped to grid points (take floor)
var iax = Math.floor(scaleFactor * (vx - min[0]));
var iay = Math.floor(scaleFactor * (vy - min[1]));
var iaz = Math.floor(scaleFactor * (vz - min[2]));
// Extents of grid to consider for this atom
var begX = Math.max(0, iax - ng);
var begY = Math.max(0, iay - ng);
var begZ = Math.max(0, iaz - ng);
// Add two to these points:
// - iax are floor'd values so this ensures coverage
// - these are loop limits (exclusive)
var endX = Math.min(dimX, iax + ng + 2);
var endY = Math.min(dimY, iay + ng + 2);
var endZ = Math.min(dimZ, iaz + ng + 2);
for (var xi = begX; xi < endX; ++xi) {
var dx = gridx[xi] - vx;
var xIdx = xi * iuv;
for (var yi = begY; yi < endY; ++yi) {
var dy = gridy[yi] - vy;
var dxySq = dx * dx + dy * dy;
var xyIdx = yi * iu + xIdx;
for (var zi = begZ; zi < endZ; ++zi) {
var dz = gridz[zi] - vz;
var dSq = dxySq + dz * dz;
if (dSq <= r2sq) {
var dens = (0, approx_1.fasterExp)(-alpha * (dSq * rSqInv));
var idx = zi + xyIdx;
data[idx] += dens;
if (dens > densData[idx]) {
densData[idx] = dens;
idData[idx] = id ? id[i] : i;
}
}
}
}
}
}
}
function accumulate() {
return tslib_1.__awaiter(this, void 0, void 0, function () {
var i;
return tslib_1.__generator(this, function (_a) {
switch (_a.label) {
case 0:
i = 0;
_a.label = 1;
case 1:
if (!(i < n)) return [3 /*break*/, 4];
accumulateRange(i, Math.min(i + updateChunk, n));
if (!ctx.shouldUpdate) return [3 /*break*/, 3];
return [4 /*yield*/, ctx.update({ message: 'filling density grid', current: i, max: n })];
case 2:
_a.sent();
_a.label = 3;
case 3:
i += updateChunk;
return [3 /*break*/, 1];
case 4: return [2 /*return*/];
}
});
});
}
var resolution, radiusOffset, smoothness, scaleFactor, indices, x, y, z, id, n, radii, maxRadius, i, r, pad, expandedBox, min, scaledBox, dim, space, data, field, idData, idField, dimX, dimY, dimZ, iu, iv, iuv, gridx, gridy, gridz, densData, alpha, updateChunk, transform;
return tslib_1.__generator(this, function (_a) {
switch (_a.label) {
case 0:
resolution = props.resolution, radiusOffset = props.radiusOffset, smoothness = props.smoothness;
scaleFactor = 1 / resolution;
indices = position.indices, x = position.x, y = position.y, z = position.z, id = position.id;
n = int_1.OrderedSet.size(indices);
radii = new Float32Array(n);
maxRadius = 0;
for (i = 0; i < n; ++i) {
r = radius(int_1.OrderedSet.getAt(indices, i)) + radiusOffset;
if (maxRadius < r)
maxRadius = r;
radii[i] = r;
}
pad = maxRadius * 2 + resolution;
expandedBox = geometry_1.Box3D.expand((0, geometry_1.Box3D)(), box, linear_algebra_1.Vec3.create(pad, pad, pad));
min = expandedBox.min;
scaledBox = geometry_1.Box3D.scale((0, geometry_1.Box3D)(), expandedBox, scaleFactor);
dim = geometry_1.Box3D.size((0, linear_algebra_1.Vec3)(), scaledBox);
linear_algebra_1.Vec3.ceil(dim, dim);
space = linear_algebra_1.Tensor.Space(dim, [0, 1, 2], Float32Array);
data = space.create();
field = linear_algebra_1.Tensor.create(space, data);
idData = space.create();
idData.fill(-1);
idField = linear_algebra_1.Tensor.create(space, idData);
dimX = dim[0], dimY = dim[1], dimZ = dim[2];
iu = dimZ, iv = dimY, iuv = iu * iv;
gridx = (0, geometry_1.fillGridDim)(dim[0], min[0], resolution);
gridy = (0, geometry_1.fillGridDim)(dim[1], min[1], resolution);
gridz = (0, geometry_1.fillGridDim)(dim[2], min[2], resolution);
densData = space.create();
alpha = smoothness;
updateChunk = Math.ceil(100000 / ((Math.pow(Math.pow(maxRadius, 3), 3) * scaleFactor)));
// console.time('gaussian density cpu')
return [4 /*yield*/, accumulate()];
case 1:
// console.time('gaussian density cpu')
_a.sent();
transform = linear_algebra_1.Mat4.identity();
linear_algebra_1.Mat4.fromScaling(transform, linear_algebra_1.Vec3.create(resolution, resolution, resolution));
linear_algebra_1.Mat4.setTranslation(transform, expandedBox.min);
return [2 /*return*/, { field: field, idField: idField, transform: transform, radiusFactor: 1, resolution: resolution, maxRadius: maxRadius }];
}
});
});
}
exports.GaussianDensityCPU = GaussianDensityCPU;