molstar
Version:
A comprehensive macromolecular library.
448 lines (447 loc) • 18.4 kB
JavaScript
/**
* Copyright (c) 2018-2025 mol* contributors, licensed under MIT, See LICENSE file for more info.
*
* @author David Sehnal <david.sehnal@gmail.com>
* @author Alexander Rose <alexander.rose@weirdbyte.de>
* @author Gianluca Tomasello <giagitom@gmail.com>
*/
import { Result } from './common';
import { Box3D } from '../primitives/box3d';
import { Vec3 } from '../../linear-algebra';
import { OrderedSet } from '../../../mol-data/int';
import { FibonacciHeap } from '../../../mol-util/fibonacci-heap';
import { memoize1 } from '../../../mol-util/memoize';
function GridLookup3D(data, boundary, cellSizeOrCount) {
return new GridLookup3DImpl(data, boundary, cellSizeOrCount);
}
export { GridLookup3D };
class GridLookup3DImpl {
find(x, y, z, radius, result) {
this.ctx.x = x;
this.ctx.y = y;
this.ctx.z = z;
this.ctx.radius = radius;
this.ctx.isCheck = false;
const ret = result !== null && result !== void 0 ? result : this.result;
query(this.ctx, ret);
return ret;
}
nearest(x, y, z, k = 1, stopIf, result) {
this.ctx.x = x;
this.ctx.y = y;
this.ctx.z = z;
this.ctx.k = k;
this.ctx.stopIf = stopIf;
const ret = result !== null && result !== void 0 ? result : this.result;
queryNearest(this.ctx, ret);
return ret;
}
check(x, y, z, radius) {
this.ctx.x = x;
this.ctx.y = y;
this.ctx.z = z;
this.ctx.radius = radius;
this.ctx.isCheck = true;
return query(this.ctx, this.result);
}
approxNearest(x, y, z, radius, result) {
this.ctx.x = x;
this.ctx.y = y;
this.ctx.z = z;
this.ctx.radius = radius;
this.ctx.isCheck = false;
const ret = result !== null && result !== void 0 ? result : this.result;
approxQueryNearest(this.ctx, ret);
return ret;
}
constructor(data, boundary, cellSizeOrCount) {
const structure = build(data, boundary, cellSizeOrCount);
this.ctx = createContext(structure);
this.boundary = { box: structure.boundingBox, sphere: structure.boundingSphere };
this.buckets = { offset: structure.bucketOffset, count: structure.bucketCounts, array: structure.bucketArray };
this.result = Result.create();
}
}
function _build(state) {
const { expandedBox, size: [sX, sY, sZ], data: { x: px, y: py, z: pz, radius, indices }, elementCount, delta } = state;
const n = sX * sY * sZ;
const { min: [minX, minY, minZ] } = expandedBox;
let maxRadius = 0;
let bucketCount = 0;
const grid = new Uint32Array(n);
const bucketIndex = new Int32Array(elementCount);
for (let t = 0; t < elementCount; t++) {
const i = OrderedSet.getAt(indices, t);
const x = Math.floor((px[i] - minX) / delta[0]);
const y = Math.floor((py[i] - minY) / delta[1]);
const z = Math.floor((pz[i] - minZ) / delta[2]);
const idx = (((x * sY) + y) * sZ) + z;
if ((grid[idx] += 1) === 1) {
bucketCount += 1;
}
bucketIndex[t] = idx;
}
if (radius) {
for (let t = 0; t < elementCount; t++) {
const i = OrderedSet.getAt(indices, t);
if (radius[i] > maxRadius)
maxRadius = radius[i];
}
}
const bucketCounts = new Int32Array(bucketCount);
for (let i = 0, j = 0; i < n; i++) {
const c = grid[i];
if (c > 0) {
grid[i] = j + 1;
bucketCounts[j] = c;
j += 1;
}
}
const bucketOffset = new Uint32Array(bucketCount);
for (let i = 1; i < bucketCount; ++i) {
bucketOffset[i] += bucketOffset[i - 1] + bucketCounts[i - 1];
}
const bucketFill = new Int32Array(bucketCount);
const bucketArray = new Int32Array(elementCount);
for (let i = 0; i < elementCount; i++) {
const bucketIdx = grid[bucketIndex[i]];
if (bucketIdx > 0) {
const k = bucketIdx - 1;
bucketArray[bucketOffset[k] + bucketFill[k]] = i;
bucketFill[k] += 1;
}
}
return {
size: state.size,
bucketArray,
bucketCounts,
bucketOffset,
grid,
delta,
min: state.expandedBox.min,
data: state.data,
maxRadius,
expandedBox: state.expandedBox,
boundingBox: state.boundingBox,
boundingSphere: state.boundingSphere
};
}
const MaxVolume = 2 ** 24;
function build(data, boundary, cellSizeOrCount) {
// need to expand the grid bounds to avoid rounding errors
const expandedBox = Box3D.expand(Box3D(), boundary.box, Vec3.create(0.5, 0.5, 0.5));
const { indices } = data;
const S = Box3D.size(Vec3(), expandedBox);
let delta, size;
const elementCount = OrderedSet.size(indices);
const cellCount = typeof cellSizeOrCount === 'number' ? cellSizeOrCount : 32;
const cellSize = Array.isArray(cellSizeOrCount) && cellSizeOrCount;
if (cellSize && !Vec3.isZero(cellSize)) {
size = [Math.ceil(S[0] / cellSize[0]), Math.ceil(S[1] / cellSize[1]), Math.ceil(S[2] / cellSize[2])];
delta = cellSize;
}
else if (elementCount > 0) {
// size of the box
// required "grid volume" so that each cell contains on average 'cellCount' elements.
const V = Math.ceil(elementCount / cellCount);
const f = Math.pow(V / (S[0] * S[1] * S[2]), 1 / 3);
size = [Math.ceil(S[0] * f), Math.ceil(S[1] * f), Math.ceil(S[2] * f)];
delta = [S[0] / size[0], S[1] / size[1], S[2] / size[2]];
}
else {
delta = S;
size = [1, 1, 1];
}
// guard against overly large grids
const volume = size[0] * size[1] * size[2];
if (volume > MaxVolume) {
const f = Math.cbrt(volume / MaxVolume);
size = [Math.ceil(size[0] / f), Math.ceil(size[1] / f), Math.ceil(size[2] / f)];
delta = [S[0] / size[0], S[1] / size[1], S[2] / size[2]];
}
const inputData = {
x: data.x,
y: data.y,
z: data.z,
indices,
radius: data.radius
};
const state = {
size,
data: inputData,
expandedBox,
boundingBox: boundary.box,
boundingSphere: boundary.sphere,
elementCount,
delta
};
return _build(state);
}
function createContext(grid) {
return { grid, x: 0.1, y: 0.1, z: 0.1, k: 1, stopIf: undefined, radius: 0.1, isCheck: false };
}
function query(ctx, result) {
const { min, size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, data: { x: px, y: py, z: pz, indices, radius }, delta, maxRadius } = ctx.grid;
const { radius: inputRadius, isCheck, x, y, z } = ctx;
const r = inputRadius + maxRadius;
const rSq = r * r;
Result.reset(result);
const loX = Math.max(0, Math.floor((x - r - min[0]) / delta[0]));
const loY = Math.max(0, Math.floor((y - r - min[1]) / delta[1]));
const loZ = Math.max(0, Math.floor((z - r - min[2]) / delta[2]));
const hiX = Math.min(sX - 1, Math.floor((x + r - min[0]) / delta[0]));
const hiY = Math.min(sY - 1, Math.floor((y + r - min[1]) / delta[1]));
const hiZ = Math.min(sZ - 1, Math.floor((z + r - min[2]) / delta[2]));
if (loX > hiX || loY > hiY || loZ > hiZ)
return false;
for (let ix = loX; ix <= hiX; ix++) {
for (let iy = loY; iy <= hiY; iy++) {
for (let iz = loZ; iz <= hiZ; iz++) {
const bucketIdx = grid[(((ix * sY) + iy) * sZ) + iz];
if (bucketIdx === 0)
continue;
const k = bucketIdx - 1;
const offset = bucketOffset[k];
const count = bucketCounts[k];
const end = offset + count;
for (let i = offset; i < end; i++) {
const idx = OrderedSet.getAt(indices, bucketArray[i]);
const dx = px[idx] - x;
const dy = py[idx] - y;
const dz = pz[idx] - z;
const distSq = dx * dx + dy * dy + dz * dz;
if (distSq <= rSq) {
if (maxRadius > 0 && Math.sqrt(distSq) - radius[idx] > inputRadius)
continue;
if (isCheck)
return true;
Result.add(result, bucketArray[i], distSq);
}
}
}
}
}
return result.count > 0;
}
function _insideOut(r) {
const cells = [];
const n = r * 2 + 1;
for (let x = 0; x < n; ++x) {
for (let y = 0; y < n; ++y) {
for (let z = 0; z < n; ++z) {
cells.push(Vec3.create(x - r, y - r, z - r));
}
}
}
cells.sort((a, b) => Vec3.squaredMagnitude(a) - Vec3.squaredMagnitude(b));
return cells.flat();
}
const insideOut = memoize1(_insideOut);
/**
* The maximum error is on the order of cell size + max radius (if the grid has radii).
*/
function approxQueryNearest(ctx, result) {
const { min, size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, data: { x: px, y: py, z: pz, indices }, delta } = ctx.grid;
const { radius, x, y, z } = ctx;
const rSq = radius * radius;
Result.reset(result);
const loX = Math.max(0, Math.floor((x - radius - min[0]) / delta[0]));
const loY = Math.max(0, Math.floor((y - radius - min[1]) / delta[1]));
const loZ = Math.max(0, Math.floor((z - radius - min[2]) / delta[2]));
const hiX = Math.min(sX - 1, Math.floor((x + radius - min[0]) / delta[0]));
const hiY = Math.min(sY - 1, Math.floor((y + radius - min[1]) / delta[1]));
const hiZ = Math.min(sZ - 1, Math.floor((z + radius - min[2]) / delta[2]));
if (loX > hiX || loY > hiY || loZ > hiZ)
return false;
const miX = Math.floor((x - min[0]) / delta[0]);
const miY = Math.floor((y - min[1]) / delta[1]);
const miZ = Math.floor((z - min[2]) / delta[2]);
const cells = insideOut(Math.max(hiX - loX, hiY - loY, hiZ - loZ) + 1);
for (let i = 0, _i = cells.length; i < _i; i += 3) {
const ix = miX + cells[i];
const iy = miY + cells[i + 1];
const iz = miZ + cells[i + 2];
if (ix < loX || ix > hiX || iy < loY || iy > hiY || iz < loZ || iz > hiZ)
continue;
const bucketIdx = grid[(((ix * sY) + iy) * sZ) + iz];
if (bucketIdx === 0)
continue;
const k = bucketIdx - 1;
const offset = bucketOffset[k];
const count = bucketCounts[k];
const end = offset + count;
let minDistSq = Number.MAX_VALUE;
for (let i = offset; i < end; i++) {
const idx = OrderedSet.getAt(indices, bucketArray[i]);
const dx = px[idx] - x;
const dy = py[idx] - y;
const dz = pz[idx] - z;
const distSq = dx * dx + dy * dy + dz * dz;
if (distSq <= rSq && distSq < minDistSq) {
Result.add(result, bucketArray[i], distSq);
minDistSq = distSq;
}
}
if (minDistSq !== Number.MAX_VALUE)
return true;
}
return result.count > 0;
}
const tmpDirVec = Vec3();
const tmpVec = Vec3();
const tmpSetG = new Set();
const tmpSetG2 = new Set();
const tmpArrG1 = [0.1];
const tmpArrG2 = [0.1];
const tmpArrG3 = [0.1];
const tmpHeapG = new FibonacciHeap();
function queryNearest(ctx, result) {
const { min, expandedBox: box, boundingSphere: { center }, size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, data: { x: px, y: py, z: pz, indices, radius }, delta, maxRadius } = ctx.grid;
const { x, y, z, k, stopIf } = ctx;
const indicesCount = OrderedSet.size(indices);
Result.reset(result);
if (indicesCount === 0 || k <= 0)
return false;
let gX, gY, gZ, stop = false, gCount = 1, expandGrid = true, nextGCount = 0, arrG = tmpArrG1, nextArrG = tmpArrG2, maxRange = 0, expandRange = true, gridId, gridPointsFinished = false;
const expandedArrG = tmpArrG3, sqMaxRadius = maxRadius * maxRadius;
arrG.length = 0;
expandedArrG.length = 0;
tmpSetG.clear();
tmpHeapG.clear();
Vec3.set(tmpVec, x, y, z);
if (!Box3D.containsVec3(box, tmpVec)) {
// intersect ray pointing to box center
Box3D.nearestIntersectionWithRay(tmpVec, box, tmpVec, Vec3.normalize(tmpDirVec, Vec3.sub(tmpDirVec, center, tmpVec)));
gX = Math.max(0, Math.min(sX - 1, Math.floor((tmpVec[0] - min[0]) / delta[0])));
gY = Math.max(0, Math.min(sY - 1, Math.floor((tmpVec[1] - min[1]) / delta[1])));
gZ = Math.max(0, Math.min(sZ - 1, Math.floor((tmpVec[2] - min[2]) / delta[2])));
}
else {
gX = Math.floor((x - min[0]) / delta[0]);
gY = Math.floor((y - min[1]) / delta[1]);
gZ = Math.floor((z - min[2]) / delta[2]);
}
const dX = maxRadius !== 0 ? Math.max(1, Math.min(sX - 1, Math.ceil(maxRadius / delta[0]))) : 1;
const dY = maxRadius !== 0 ? Math.max(1, Math.min(sY - 1, Math.ceil(maxRadius / delta[1]))) : 1;
const dZ = maxRadius !== 0 ? Math.max(1, Math.min(sZ - 1, Math.ceil(maxRadius / delta[2]))) : 1;
arrG.push(gX, gY, gZ, (((gX * sY) + gY) * sZ) + gZ);
while (result.count < indicesCount) {
const arrGLen = gCount * 4;
for (let ig = 0; ig < arrGLen; ig += 4) {
gridId = arrG[ig + 3];
if (!tmpSetG.has(gridId)) {
tmpSetG.add(gridId);
gridPointsFinished = tmpSetG.size >= grid.length;
const bucketIdx = grid[gridId];
if (bucketIdx !== 0) {
const _maxRange = maxRange;
const ki = bucketIdx - 1;
const offset = bucketOffset[ki];
const count = bucketCounts[ki];
const end = offset + count;
for (let i = offset; i < end; i++) {
const bIdx = bucketArray[i];
const idx = OrderedSet.getAt(indices, bIdx);
const dx = px[idx] - x;
const dy = py[idx] - y;
const dz = pz[idx] - z;
let distSq = dx * dx + dy * dy + dz * dz;
if (maxRadius !== 0) {
const r = radius[idx];
distSq -= r * r;
}
if (expandRange && distSq > maxRange) {
maxRange = distSq;
}
tmpHeapG.insert(distSq, bIdx);
}
if (_maxRange < maxRange)
expandRange = false;
}
}
}
// find next grid points
nextArrG.length = 0;
nextGCount = 0;
tmpSetG2.clear();
for (let ig = 0; ig < arrGLen; ig += 4) {
gX = arrG[ig];
gY = arrG[ig + 1];
gZ = arrG[ig + 2];
// fill grid points array with valid adiacent positions
for (let ix = -dX; ix <= dX; ix++) {
const xPos = gX + ix;
if (xPos < 0 || xPos >= sX)
continue;
for (let iy = -dY; iy <= dY; iy++) {
const yPos = gY + iy;
if (yPos < 0 || yPos >= sY)
continue;
for (let iz = -dZ; iz <= dZ; iz++) {
const zPos = gZ + iz;
if (zPos < 0 || zPos >= sZ)
continue;
gridId = (((xPos * sY) + yPos) * sZ) + zPos;
if (tmpSetG2.has(gridId))
continue; // already scanned
tmpSetG2.add(gridId);
if (tmpSetG.has(gridId))
continue; // already visited
if (!expandGrid) {
const xP = min[0] + xPos * delta[0] - x;
const yP = min[1] + yPos * delta[1] - y;
const zP = min[2] + zPos * delta[2] - z;
const distSqG = (xP * xP) + (yP * yP) + (zP * zP) - sqMaxRadius; // is sqMaxRadius necessary?
if (distSqG > maxRange) {
expandedArrG.push(xPos, yPos, zPos, gridId);
continue;
}
}
nextArrG.push(xPos, yPos, zPos, gridId);
nextGCount++;
}
}
}
}
expandGrid = false;
if (nextGCount === 0) {
if (k === 1) {
const node = tmpHeapG.findMinimum();
if (node) {
const { key: squaredDistance, value: index } = node;
// const squaredDistance = node!.key, index = node!.value;
Result.add(result, index, squaredDistance);
return true;
}
}
else {
while (!tmpHeapG.isEmpty() && (gridPointsFinished || tmpHeapG.findMinimum().key <= maxRange) && result.count < k) {
const node = tmpHeapG.extractMinimum();
const squaredDistance = node.key, index = node.value;
Result.add(result, index, squaredDistance);
if (stopIf && !stop) {
stop = stopIf(index, squaredDistance);
}
}
}
if (result.count >= k || stop || result.count >= indicesCount)
return result.count > 0;
expandGrid = true;
expandRange = true;
if (expandedArrG.length > 0) {
for (let i = 0, l = expandedArrG.length; i < l; i++) {
arrG.push(expandedArrG[i]);
}
expandedArrG.length = 0;
gCount = arrG.length;
}
}
else {
const tmp = arrG;
arrG = nextArrG;
nextArrG = tmp;
gCount = nextGCount;
}
}
return result.count > 0;
}