UNPKG

molstar

Version:

A comprehensive macromolecular library.

448 lines (447 loc) 18.4 kB
/** * 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; }