refactoring grid lookup (step 2)

src/mol-math/geometry/lookup3d/grid.ts

@@ -43,29 +43,31 @@ interface SpatialHash {
     size: number[],
     min: number[],
     grid: Uint32Array,
-    divisor: number[],
+    delta: number[],
     bucketOffset: Uint32Array,
     bucketCounts: Int32Array,
     bucketArray: Int32Array,
     data: InputData,
     maxRadius: number,
+    expandedBox: Box3D,
     boundingBox: Box3D,
     boundingSphere: Sphere3D
 }
 
 interface BuildState {
     size: number[],
-    divisor: number[],
+    delta: number[],
     data: InputData,
+    expandedBox: Box3D,
     boundingBox: Box3D,
     boundingSphere: Sphere3D,
     count: number
 }
 
 function _build(state: BuildState): SpatialHash {
-    const { boundingBox, size: [sX, sY, sZ], data: { x: px, y: py, z: pz, radius, indices }, count, divisor } = state;
+    const { expandedBox, size: [sX, sY, sZ], data: { x: px, y: py, z: pz, radius, indices }, count, delta } = state;
     const n = sX * sY * sZ;
-    const { min: [minX, minY, minZ] } = boundingBox;
+    const { min: [minX, minY, minZ] } = expandedBox;
 
     let maxRadius = 0;
     let bucketCount = 0;
@@ -73,10 +75,11 @@ function _build(state: BuildState): SpatialHash {
     const bucketIndex = new Int32Array(count);
     for (let t = 0, _t = indices.length; t < _t; t++) {
         const i = indices[t];
-        const x = Math.floor((px[i] - minX) / divisor[0]);
-        const y = Math.floor((py[i] - minY) / divisor[1]);
-        const z = Math.floor((pz[i] - minZ) / divisor[2]);
+        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
         }
@@ -122,10 +125,11 @@ function _build(state: BuildState): SpatialHash {
         bucketCounts,
         bucketOffset,
         grid,
-        divisor,
-        min: state.boundingBox.min,
+        delta,
+        min: state.expandedBox.min,
         data: state.data,
         maxRadius,
+        expandedBox: state.expandedBox,
         boundingBox: state.boundingBox,
         boundingSphere: state.boundingSphere
     }
@@ -133,6 +137,7 @@ function _build(state: BuildState): SpatialHash {
 
 export function build(data: PositionData) {
     const boundingBox = Box3D.computeBounding(data);
+    const expandedBox = Box3D.expand(boundingBox, Vec3.create(0.5, 0.5, 0.5));
     const boundingSphere = Sphere3D.computeBounding(data);
 
     // TODO: specialized code that does not require the indices annotation?
@@ -141,17 +146,20 @@ export function build(data: PositionData) {
         for (let i = 0, _i = data.x.length; i < _i; i++) indices[i] = i;
     }
 
-    // size of the box
-    const S = Vec3.sub(Vec3.zero(), boundingBox.max, boundingBox.min);
-    // required "grid volume" so that each cell contains on average 32 elements.
-    const V = Math.ceil(indices.length / 32);
-    const f = Math.pow(V / (S[0] * S[1] * S[2]), 1 / 3);
-
-    const divisor = [Math.ceil(S[0] * f), Math.ceil(S[1] * f), Math.ceil(S[2] * f)];
-
-    const size = [Math.floor(S[0] / divisor[0]) + 1, Math.floor(S[1] / divisor[1]) + 1, Math.floor(S[2] / divisor[2]) + 1];
-
-    console.log({ divisor, size })
+    const S = Vec3.sub(Vec3.zero(), expandedBox.max, expandedBox.min);
+    let delta, size;
+
+    if (indices.length > 0) {
+        // size of the box
+        // required "grid volume" so that each cell contains on average 32 elements.
+        const V = Math.ceil(indices.length / 32);
+        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];
+    }
 
     const inputData: InputData = {
         x: data.x,
@@ -164,53 +172,55 @@ export function build(data: PositionData) {
     const state: BuildState = {
         size,
         data: inputData,
+        expandedBox,
         boundingBox,
         boundingSphere,
         count: indices.length,
-        divisor
+        delta
     }
 
     return _build(state);
 }
 
 export function inSphere(structure: SpatialHash, x: number, y: number, z: number, radius: number): Result {
-    const { min, size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, data: { x: px, y: py, z: pz, indices }, divisor } = structure;
+    const { min, size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, data: { x: px, y: py, z: pz, indices }, delta } = structure;
     //const { radius: r, radiusSq: rSq, pivot: [x, y, z], isCheck, mask } = ctx;
     const r = radius, rSq = r * r;
 
     const result = Result.create();
 
-    const loX = Math.max(0, Math.floor((x - r - min[0]) / divisor[0]));
-    const loY = Math.max(0, Math.floor((y - r - min[1]) / divisor[1]));
-    const loZ = Math.max(0, Math.floor((z - r - min[2]) / divisor[2]));
+    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]));
 
-    const hiX = Math.min(sX, Math.floor((x + r - min[0]) / divisor[0]));
-    const hiY = Math.min(sY, Math.floor((y + r - min[1]) / divisor[1]));
-    const hiZ = Math.min(sZ, Math.floor((z + r - min[2]) / divisor[0]));
+    if (loX > hiX || loY > hiY || loZ > hiZ) return result;
 
     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;
 
-                if (bucketIdx > 0) {
-                    const k = bucketIdx - 1;
-                    const offset = bucketOffset[k];
-                    const count = bucketCounts[k];
-                    const end = offset + count;
+                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 = indices[bucketArray[i]];
+                for (let i = offset; i < end; i++) {
+                    const idx = 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;
+                    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 (isCheck) return true;
-                            Result.add(result, idx, distSq);
-                        }
+                    if (distSq <= rSq) {
+                        // if (isCheck) return true;
+                        Result.add(result, bucketArray[i], distSq);
                     }
                 }
             }

src/mol-math/geometry/primitives/box3d.ts

@@ -39,6 +39,13 @@ namespace Box3D {
 
         return { min: Vec3.create(min[0], min[1], min[2]), max: Vec3.create(max[0], max[1], max[2]) }
     }
+
+    export function expand(box: Box3D, delta: Vec3): Box3D {
+        return {
+            min: Vec3.sub(Vec3.zero(), box.min, delta),
+            max: Vec3.add(Vec3.zero(), box.max, delta)
+        }
+    }
 }
 
 export { Box3D }

src/perf-tests/lookup3d.ts

@@ -6,6 +6,7 @@ import { Structure, Model } from 'mol-model/structure'
 
 import { Run } from 'mol-task';
 import { build, inSphere } from 'mol-math/geometry/lookup3d/grid';
+import { sortArray } from 'mol-data/util';
 
 require('util.promisify').shim();
 const readFileAsync = util.promisify(fs.readFile);
@@ -41,11 +42,14 @@ export async function readCIF(path: string) {
 export async function test() {
     const { mmcif  } = await readCIF('e:/test/quick/1tqn_updated.cif');
 
-    const data = build({ x: mmcif.atom_site.Cartn_x.toArray(), y: mmcif.atom_site.Cartn_y.toArray(), z: mmcif.atom_site.Cartn_z.toArray(), indices: void 0 });
-    console.log(data.boundingBox);
+    const data = build({ x: mmcif.atom_site.Cartn_x.toArray(), y: mmcif.atom_site.Cartn_y.toArray(), z: mmcif.atom_site.Cartn_z.toArray(),
+        //indices: [0, 1, 2, 3]
+        //indices: [1]
+    });
+    console.log(data.boundingBox, data.boundingSphere);
 
     const result = inSphere(data, -30.07, 8.178, -13.897, 10);
-    console.log(result.count);
+    console.log(result.count, sortArray(result.indices));
 }
 
 test();