molstar
Version:
A comprehensive macromolecular library.
135 lines • 7.1 kB
JavaScript
/**
* Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
*
* @author Alexander Rose <alexander.rose@weirdbyte.de>
*/
import { __awaiter, __generator } from "tslib";
import { Box3D, fillGridDim } from '../../geometry';
import { Vec3, Mat4, Tensor } from '../../linear-algebra';
import { OrderedSet } from '../../../mol-data/int';
import { fasterExp } from '../../approx';
export function GaussianDensityCPU(ctx, position, box, radius, props) {
return __awaiter(this, void 0, void 0, function () {
function accumulateRange(begI, endI) {
for (var i = begI; i < endI; ++i) {
var j = 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 = 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 __awaiter(this, void 0, void 0, function () {
var i;
return __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 __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 = OrderedSet.size(indices);
radii = new Float32Array(n);
maxRadius = 0;
for (i = 0; i < n; ++i) {
r = radius(OrderedSet.getAt(indices, i)) + radiusOffset;
if (maxRadius < r)
maxRadius = r;
radii[i] = r;
}
pad = maxRadius * 2 + resolution;
expandedBox = Box3D.expand(Box3D(), box, Vec3.create(pad, pad, pad));
min = expandedBox.min;
scaledBox = Box3D.scale(Box3D(), expandedBox, scaleFactor);
dim = Box3D.size(Vec3(), scaledBox);
Vec3.ceil(dim, dim);
space = Tensor.Space(dim, [0, 1, 2], Float32Array);
data = space.create();
field = Tensor.create(space, data);
idData = space.create();
idData.fill(-1);
idField = Tensor.create(space, idData);
dimX = dim[0], dimY = dim[1], dimZ = dim[2];
iu = dimZ, iv = dimY, iuv = iu * iv;
gridx = fillGridDim(dim[0], min[0], resolution);
gridy = fillGridDim(dim[1], min[1], resolution);
gridz = 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 = Mat4.identity();
Mat4.fromScaling(transform, Vec3.create(resolution, resolution, resolution));
Mat4.setTranslation(transform, expandedBox.min);
return [2 /*return*/, { field: field, idField: idField, transform: transform, radiusFactor: 1, resolution: resolution }];
}
});
});
}
//# sourceMappingURL=cpu.js.map