refactoring grid lookup (step 3)

src/mol-math/geometry.ts

@@ -0,0 +1,12 @@
+/*
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+export * from './geometry/common'
+export * from './geometry/symmetry-operator'
+export * from './geometry/lookup3d/common'
+export * from './geometry/lookup3d/grid'
+export * from './geometry/primitives/box3d'
+export * from './geometry/primitives/sphere3d'

src/mol-math/geometry/_spec/lookup3d.spec.ts

@@ -0,0 +1,51 @@
+/*
+* Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+*
+* @author David Sehnal <david.sehnal@gmail.com>
+*/
+
+import { GridLookup3D } from '../../geometry';
+import { sortArray } from 'mol-data/util';
+import { OrderedSet } from 'mol-data/int';
+
+const xs = [0, 0, 1];
+const ys = [0, 1, 0];
+const zs = [0, 0, 0];
+const rs = [0, 0.5, 1/3];
+
+describe('GridLookup3d', () => {
+    it('basic', () => {
+        const grid = GridLookup3D({ x: xs, y: ys, z: zs, indices: OrderedSet.ofBounds(0, 3) });
+
+        let r = grid.find(0, 0, 0, 0);
+        expect(r.count).toBe(1);
+        expect(r.indices[0]).toBe(0);
+
+        r = grid.find(0, 0, 0, 1);
+        expect(r.count).toBe(3);
+        expect(sortArray(r.indices)).toEqual([0, 1, 2]);
+    });
+
+    it('radius', () => {
+        const grid = GridLookup3D({ x: xs, y: ys, z: zs, radius: [0, 0.5, 1 / 3], indices: OrderedSet.ofBounds(0, 3) });
+
+        let r = grid.find(0, 0, 0, 0);
+        expect(r.count).toBe(1);
+        expect(r.indices[0]).toBe(0);
+
+        r = grid.find(0, 0, 0, 0.5);
+        expect(r.count).toBe(2);
+        expect(sortArray(r.indices)).toEqual([0, 1]);
+    });
+
+    it('indexed', () => {
+        const grid = GridLookup3D({ x: xs, y: ys, z: zs, indices: OrderedSet.ofSingleton(1), radius: rs });
+
+        let r = grid.find(0, 0, 0, 0);
+        expect(r.count).toBe(0);
+
+        r = grid.find(0, 0, 0, 0.5);
+        expect(r.count).toBe(1);
+        expect(sortArray(r.indices)).toEqual([0]);
+    });
+});

src/mol-math/geometry/common.ts

@@ -1,14 +1,17 @@
 /*
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author David Sehnal <david.sehnal@gmail.com>
- */
+* Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+*
+* @author David Sehnal <david.sehnal@gmail.com>
+*/
+
+import { OrderedSet } from 'mol-data/int'
 
 export interface PositionData {
     x: ArrayLike<number>,
     y: ArrayLike<number>,
     z: ArrayLike<number>,
-    radius?: ArrayLike<number>,
     // subset indices into the x/y/z/radius arrays
-    indices?: ArrayLike<number>,
+    indices: OrderedSet,
+    // optional element radius
+    radius?: ArrayLike<number>
 }

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

@@ -1,291 +1,291 @@
-/*
- * Copyright (c) 2017 MolQL contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author David Sehnal <david.sehnal@gmail.com>
- * @author Alexander Rose <alexander.rose@weirdbyte.de>
- */
-
-// TODO: 3d grid lookup (use single cell for small sets), make bounding sphere part of the structure
-// TODO: assign radius to points?
-
-/**
- * A "masked" 3D spatial lookup structure.
- */
-
-import Mask from 'mol-util/mask'
-
-export type FindFunc = (x: number, y: number, z: number, radius: number) => Result
-export type CheckFunc = (x: number, y: number, z: number, radius: number) => boolean
-
-export interface Result {
-    readonly count: number,
-    readonly indices: number[],
-    readonly squaredDistances: number[]
-}
-
-export interface Positions { x: ArrayLike<number>, y: ArrayLike<number>, z: ArrayLike<number> }
-
-interface GridLookup { find: (mask: Mask) => FindFunc, check: (mask: Mask) => CheckFunc }
-
-function GridLookup(positions: Positions): GridLookup {
-    const tree = build(createInputData(positions));
-    return {
-        find(mask) {
-            const ctx = QueryContext.create(tree, mask, false);
-            return function (x: number, y: number, z: number, radius: number) {
-                QueryContext.update(ctx, x, y, z, radius);
-                nearest(ctx);
-                return ctx.buffer;
-            }
-        },
-        check(mask) {
-            const ctx = QueryContext.create(tree, mask, true);
-            return function (x: number, y: number, z: number, radius: number) {
-                QueryContext.update(ctx, x, y, z, radius);
-                return nearest(ctx);
-            }
-        }
-    }
-}
-
-interface InputData {
-    bounds: Box3D,
-    count: number,
-    positions: Positions
-}
-
-interface Box3D {
-    min: number[],
-    max: number[]
-}
-
-namespace Box3D {
-    export function createInfinite(): Box3D {
-        return {
-            min: [Number.MAX_VALUE, Number.MAX_VALUE, Number.MAX_VALUE],
-            max: [-Number.MAX_VALUE, -Number.MAX_VALUE, -Number.MAX_VALUE]
-        }
-    }
-}
-
-/**
-* Query context. Handles the actual querying.
-*/
-interface QueryContext<T> {
-    structure: T,
-    pivot: number[],
-    radius: number,
-    radiusSq: number,
-    buffer: QueryContext.Buffer,
-    mask: Mask,
-    isCheck: boolean
-}
-
-namespace QueryContext {
-    export interface Buffer extends Result {
-        count: number;
-        indices: any[];
-        squaredDistances: number[];
-    }
-
-    export function add<T>(ctx: QueryContext<T>, distSq: number, index: number) {
-        const buffer = ctx.buffer;
-        buffer.squaredDistances[buffer.count] = distSq;
-        buffer.indices[buffer.count++] = index;
-    }
-
-    function resetBuffer(buffer: Buffer) { buffer.count = 0; }
-
-    function createBuffer(): Buffer {
-        return {
-            indices: [],
-            count: 0,
-            squaredDistances: []
-        }
-    }
-
-    /**
-     * Query the tree and store the result to this.buffer. Overwrites the old result.
-     */
-    export function update<T>(ctx: QueryContext<T>, x: number, y: number, z: number, radius: number) {
-        ctx.pivot[0] = x;
-        ctx.pivot[1] = y;
-        ctx.pivot[2] = z;
-        ctx.radius = radius;
-        ctx.radiusSq = radius * radius;
-        resetBuffer(ctx.buffer);
-    }
-
-    export function create<T>(structure: T, mask: Mask, isCheck: boolean): QueryContext<T> {
-        return {
-            structure,
-            buffer: createBuffer(),
-            pivot: [0.1, 0.1, 0.1],
-            radius: 1.1,
-            radiusSq: 1.1 * 1.1,
-            mask,
-            isCheck
-        }
-    }
-}
-
-function createInputData(positions: { x: ArrayLike<number>, y: ArrayLike<number>, z: ArrayLike<number> }): InputData {
-    const { x, y, z } = positions;
-    const bounds = Box3D.createInfinite();
-    const count = x.length;
-    const { min, max } = bounds;
-
-    for (let i = 0; i < count; i++) {
-        min[0] = Math.min(x[i], min[0]);
-        min[1] = Math.min(y[i], min[1]);
-        min[2] = Math.min(z[i], min[2]);
-        max[0] = Math.max(x[i], max[0]);
-        max[1] = Math.max(y[i], max[1]);
-        max[2] = Math.max(z[i], max[2]);
-    }
-
-    return { positions, bounds, count };
-}
-
-/**
- * Adapted from https://github.com/arose/ngl
- * MIT License Copyright (C) 2014+ Alexander Rose
- */
-
-interface SpatialHash {
-    size: number[],
-    min: number[],
-    grid: Uint32Array,
-    bucketOffset: Uint32Array,
-    bucketCounts: Int32Array,
-    bucketArray: Int32Array,
-    positions: Positions
-}
-
-interface State {
-    size: number[],
-    positions: Positions,
-    bounds: Box3D,
-    count: number
-}
-
-const enum Constants { Exp = 3 }
-
-function nearest(ctx: QueryContext<SpatialHash>): boolean {
-    const { min: [minX, minY, minZ], size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, positions: { x: px, y: py, z: pz } } = ctx.structure;
-    const { radius: r, radiusSq: rSq, pivot: [x, y, z], isCheck, mask } = ctx;
-
-    const loX = Math.max(0, (x - r - minX) >> Constants.Exp);
-    const loY = Math.max(0, (y - r - minY) >> Constants.Exp);
-    const loZ = Math.max(0, (z - r - minZ) >> Constants.Exp);
-
-    const hiX = Math.min(sX, (x + r - minX) >> Constants.Exp);
-    const hiY = Math.min(sY, (y + r - minY) >> Constants.Exp);
-    const hiZ = Math.min(sZ, (z + r - minZ) >> Constants.Exp);
-
-    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) {
-                    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 = bucketArray[i];
-                        if (!mask.has(idx)) continue;
-
-                        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;
-                            QueryContext.add(ctx, distSq, idx)
-                        }
-                    }
-                }
-            }
-        }
-    }
-    return ctx.buffer.count > 0;
-}
-
-function _build(state: State): SpatialHash {
-    const { bounds, size: [sX, sY, sZ], positions: { x: px, y: py, z: pz }, count } = state;
-    const n = sX * sY * sZ;
-    const { min: [minX, minY, minZ] } = bounds;
-
-    let bucketCount = 0;
-    const grid = new Uint32Array(n);
-    const bucketIndex = new Int32Array(count);
-    for (let i = 0; i < count; i++) {
-        const x = (px[i] - minX) >> Constants.Exp;
-        const y = (py[i] - minY) >> Constants.Exp;
-        const z = (pz[i] - minZ) >> Constants.Exp;
-        const idx = (((x * sY) + y) * sZ) + z;
-        if ((grid[idx] += 1) === 1) {
-            bucketCount += 1
-        }
-        bucketIndex[i] = idx;
-    }
-
-    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(count);
-    for (let i = 1; i < count; ++i) {
-        bucketOffset[i] += bucketOffset[i - 1] + bucketCounts[i - 1];
-    }
-
-    const bucketFill = new Int32Array(bucketCount);
-    const bucketArray = new Int32Array(count);
-    for (let i = 0; i < count; 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,
-        min: state.bounds.min,
-        positions: state.positions
-    }
-}
-
-function build({ positions, bounds, count }: InputData): SpatialHash {
-    const size = [
-        ((bounds.max[0] - bounds.min[0]) >> Constants.Exp) + 1,
-        ((bounds.max[1] - bounds.min[1]) >> Constants.Exp) + 1,
-        ((bounds.max[2] - bounds.min[2]) >> Constants.Exp) + 1
-    ];
-
-    const state: State = {
-        size,
-        positions,
-        bounds,
-        count
-    }
-
-    return _build(state);
-}
-
-export default GridLookup;
+// /*
+//  * Copyright (c) 2017 MolQL contributors, licensed under MIT, See LICENSE file for more info.
+//  *
+//  * @author David Sehnal <david.sehnal@gmail.com>
+//  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+//  */
+
+// // TODO: 3d grid lookup (use single cell for small sets), make bounding sphere part of the structure
+// // TODO: assign radius to points?
+
+// /**
+//  * A "masked" 3D spatial lookup structure.
+//  */
+
+// import Mask from 'mol-util/mask'
+
+// export type FindFunc = (x: number, y: number, z: number, radius: number) => Result
+// export type CheckFunc = (x: number, y: number, z: number, radius: number) => boolean
+
+// export interface Result {
+//     readonly count: number,
+//     readonly indices: number[],
+//     readonly squaredDistances: number[]
+// }
+
+// export interface Positions { x: ArrayLike<number>, y: ArrayLike<number>, z: ArrayLike<number> }
+
+// interface GridLookup { find: (mask: Mask) => FindFunc, check: (mask: Mask) => CheckFunc }
+
+// function GridLookup(positions: Positions): GridLookup {
+//     const tree = build(createInputData(positions));
+//     return {
+//         find(mask) {
+//             const ctx = QueryContext.create(tree, mask, false);
+//             return function (x: number, y: number, z: number, radius: number) {
+//                 QueryContext.update(ctx, x, y, z, radius);
+//                 nearest(ctx);
+//                 return ctx.buffer;
+//             }
+//         },
+//         check(mask) {
+//             const ctx = QueryContext.create(tree, mask, true);
+//             return function (x: number, y: number, z: number, radius: number) {
+//                 QueryContext.update(ctx, x, y, z, radius);
+//                 return nearest(ctx);
+//             }
+//         }
+//     }
+// }
+
+// interface InputData {
+//     bounds: Box3D,
+//     count: number,
+//     positions: Positions
+// }
+
+// interface Box3D {
+//     min: number[],
+//     max: number[]
+// }
+
+// namespace Box3D {
+//     export function createInfinite(): Box3D {
+//         return {
+//             min: [Number.MAX_VALUE, Number.MAX_VALUE, Number.MAX_VALUE],
+//             max: [-Number.MAX_VALUE, -Number.MAX_VALUE, -Number.MAX_VALUE]
+//         }
+//     }
+// }
+
+// /**
+// * Query context. Handles the actual querying.
+// */
+// interface QueryContext<T> {
+//     structure: T,
+//     pivot: number[],
+//     radius: number,
+//     radiusSq: number,
+//     buffer: QueryContext.Buffer,
+//     mask: Mask,
+//     isCheck: boolean
+// }
+
+// namespace QueryContext {
+//     export interface Buffer extends Result {
+//         count: number;
+//         indices: any[];
+//         squaredDistances: number[];
+//     }
+
+//     export function add<T>(ctx: QueryContext<T>, distSq: number, index: number) {
+//         const buffer = ctx.buffer;
+//         buffer.squaredDistances[buffer.count] = distSq;
+//         buffer.indices[buffer.count++] = index;
+//     }
+
+//     function resetBuffer(buffer: Buffer) { buffer.count = 0; }
+
+//     function createBuffer(): Buffer {
+//         return {
+//             indices: [],
+//             count: 0,
+//             squaredDistances: []
+//         }
+//     }
+
+//     /**
+//      * Query the tree and store the result to this.buffer. Overwrites the old result.
+//      */
+//     export function update<T>(ctx: QueryContext<T>, x: number, y: number, z: number, radius: number) {
+//         ctx.pivot[0] = x;
+//         ctx.pivot[1] = y;
+//         ctx.pivot[2] = z;
+//         ctx.radius = radius;
+//         ctx.radiusSq = radius * radius;
+//         resetBuffer(ctx.buffer);
+//     }
+
+//     export function create<T>(structure: T, mask: Mask, isCheck: boolean): QueryContext<T> {
+//         return {
+//             structure,
+//             buffer: createBuffer(),
+//             pivot: [0.1, 0.1, 0.1],
+//             radius: 1.1,
+//             radiusSq: 1.1 * 1.1,
+//             mask,
+//             isCheck
+//         }
+//     }
+// }
+
+// function createInputData(positions: { x: ArrayLike<number>, y: ArrayLike<number>, z: ArrayLike<number> }): InputData {
+//     const { x, y, z } = positions;
+//     const bounds = Box3D.createInfinite();
+//     const count = x.length;
+//     const { min, max } = bounds;
+
+//     for (let i = 0; i < count; i++) {
+//         min[0] = Math.min(x[i], min[0]);
+//         min[1] = Math.min(y[i], min[1]);
+//         min[2] = Math.min(z[i], min[2]);
+//         max[0] = Math.max(x[i], max[0]);
+//         max[1] = Math.max(y[i], max[1]);
+//         max[2] = Math.max(z[i], max[2]);
+//     }
+
+//     return { positions, bounds, count };
+// }
+
+// /**
+//  * Adapted from https://github.com/arose/ngl
+//  * MIT License Copyright (C) 2014+ Alexander Rose
+//  */
+
+// interface SpatialHash {
+//     size: number[],
+//     min: number[],
+//     grid: Uint32Array,
+//     bucketOffset: Uint32Array,
+//     bucketCounts: Int32Array,
+//     bucketArray: Int32Array,
+//     positions: Positions
+// }
+
+// interface State {
+//     size: number[],
+//     positions: Positions,
+//     bounds: Box3D,
+//     count: number
+// }
+
+// const enum Constants { Exp = 3 }
+
+// function nearest(ctx: QueryContext<SpatialHash>): boolean {
+//     const { min: [minX, minY, minZ], size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, positions: { x: px, y: py, z: pz } } = ctx.structure;
+//     const { radius: r, radiusSq: rSq, pivot: [x, y, z], isCheck, mask } = ctx;
+
+//     const loX = Math.max(0, (x - r - minX) >> Constants.Exp);
+//     const loY = Math.max(0, (y - r - minY) >> Constants.Exp);
+//     const loZ = Math.max(0, (z - r - minZ) >> Constants.Exp);
+
+//     const hiX = Math.min(sX, (x + r - minX) >> Constants.Exp);
+//     const hiY = Math.min(sY, (y + r - minY) >> Constants.Exp);
+//     const hiZ = Math.min(sZ, (z + r - minZ) >> Constants.Exp);
+
+//     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) {
+//                     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 = bucketArray[i];
+//                         if (!mask.has(idx)) continue;
+
+//                         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;
+//                             QueryContext.add(ctx, distSq, idx)
+//                         }
+//                     }
+//                 }
+//             }
+//         }
+//     }
+//     return ctx.buffer.count > 0;
+// }
+
+// function _build(state: State): SpatialHash {
+//     const { bounds, size: [sX, sY, sZ], positions: { x: px, y: py, z: pz }, count } = state;
+//     const n = sX * sY * sZ;
+//     const { min: [minX, minY, minZ] } = bounds;
+
+//     let bucketCount = 0;
+//     const grid = new Uint32Array(n);
+//     const bucketIndex = new Int32Array(count);
+//     for (let i = 0; i < count; i++) {
+//         const x = (px[i] - minX) >> Constants.Exp;
+//         const y = (py[i] - minY) >> Constants.Exp;
+//         const z = (pz[i] - minZ) >> Constants.Exp;
+//         const idx = (((x * sY) + y) * sZ) + z;
+//         if ((grid[idx] += 1) === 1) {
+//             bucketCount += 1
+//         }
+//         bucketIndex[i] = idx;
+//     }
+
+//     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(count);
+//     for (let i = 1; i < count; ++i) {
+//         bucketOffset[i] += bucketOffset[i - 1] + bucketCounts[i - 1];
+//     }
+
+//     const bucketFill = new Int32Array(bucketCount);
+//     const bucketArray = new Int32Array(count);
+//     for (let i = 0; i < count; 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,
+//         min: state.bounds.min,
+//         positions: state.positions
+//     }
+// }
+
+// function build({ positions, bounds, count }: InputData): SpatialHash {
+//     const size = [
+//         ((bounds.max[0] - bounds.min[0]) >> Constants.Exp) + 1,
+//         ((bounds.max[1] - bounds.min[1]) >> Constants.Exp) + 1,
+//         ((bounds.max[2] - bounds.min[2]) >> Constants.Exp) + 1
+//     ];
+
+//     const state: State = {
+//         size,
+//         positions,
+//         bounds,
+//         count
+//     }
+
+//     return _build(state);
+// }
+
+// export default GridLookup;

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

@@ -32,6 +32,5 @@ export interface Lookup3D {
     // The result is mutated with each call to find.
     find(x: number, y: number, z: number, radius: number): Result,
     check(x: number, y: number, z: number, radius: number): boolean,
-    boundingBox: Box3D,
-    boundingSphere: Sphere3D
+    boundary: { box: Box3D, sphere: Sphere3D }
 }

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

@@ -5,26 +5,47 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Result/*, Lookup3D*/ } from './common'
+import { Result, Lookup3D } from './common'
 import { Box3D } from '../primitives/box3d';
 import { Sphere3D } from '../primitives/sphere3d';
 import { PositionData } from '../common';
 import { Vec3 } from '../../linear-algebra';
+import { OrderedSet } from 'mol-data/int';
 
-// class GridLookup3D implements Lookup3D {
-//     private result = Result.create();
-//     private pivot = [0.1, 0.1, 0.1];
-//     private radiusSq = 0.1;
-
-//     find(x: number, y: number, z: number, radius: number): Result {
-//         throw new Error("Method not implemented.");
-//     }
-//     check(x: number, y: number, z: number, radius: number): boolean {
-//         throw new Error("Method not implemented.");
-//     }
-//     boundingBox: Box3D;
-//     boundingSphere: Sphere3D;
-// }
+function GridLookup3D(data: PositionData): Lookup3D {
+    return new GridLookup3DImpl(data);
+}
+
+export { GridLookup3D }
+
+class GridLookup3DImpl implements Lookup3D {
+    private ctx: QueryContext;
+    boundary: Lookup3D['boundary'];
+
+    find(x: number, y: number, z: number, radius: number): Result {
+        this.ctx.x = x;
+        this.ctx.y = y;
+        this.ctx.z = z;
+        this.ctx.radius = radius;
+        this.ctx.isCheck = false;
+        query(this.ctx);
+        return this.ctx.result;
+    }
+    check(x: number, y: number, z: number, radius: number): boolean {
+        this.ctx.x = x;
+        this.ctx.y = y;
+        this.ctx.z = z;
+        this.ctx.radius = radius;
+        this.ctx.isCheck = true;
+        return query(this.ctx);
+    }
+
+    constructor(data: PositionData) {
+        const structure = build(data);
+        this.ctx = createContext(structure);
+        this.boundary = { box: structure.boundingBox, sphere: structure.boundingSphere };
+    }
+}
 
 /**
  * Adapted from https://github.com/arose/ngl
@@ -35,11 +56,11 @@ interface InputData {
     x: ArrayLike<number>,
     y: ArrayLike<number>,
     z: ArrayLike<number>,
+    indices: OrderedSet,
     radius?: ArrayLike<number>,
-    indices: ArrayLike<number>
 }
 
-interface SpatialHash {
+interface Grid3D {
     size: number[],
     min: number[],
     grid: Uint32Array,
@@ -64,7 +85,7 @@ interface BuildState {
     count: number
 }
 
-function _build(state: BuildState): SpatialHash {
+function _build(state: BuildState): Grid3D {
     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] } = expandedBox;
@@ -73,8 +94,8 @@ function _build(state: BuildState): SpatialHash {
     let bucketCount = 0;
     const grid = new Uint32Array(n);
     const bucketIndex = new Int32Array(count);
-    for (let t = 0, _t = indices.length; t < _t; t++) {
-        const i = indices[t];
+    for (let t = 0, _t = OrderedSet.size(indices); t < _t; 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]);
@@ -87,8 +108,8 @@ function _build(state: BuildState): SpatialHash {
     }
 
     if (radius) {
-        for (let t = 0, _t = indices.length; t < _t; t++) {
-            const i = indices[t];
+        for (let t = 0, _t = OrderedSet.size(indices); t < _t; t++) {
+            const i = OrderedSet.getAt(indices, t);
             if (radius[i] > maxRadius) maxRadius = radius[i];
         }
     }
@@ -135,24 +156,22 @@ function _build(state: BuildState): SpatialHash {
     }
 }
 
-export function build(data: PositionData) {
+function build(data: PositionData) {
     const boundingBox = Box3D.computeBounding(data);
+    // need to expand the grid bounds to avoid rounding errors
     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?
-    const indices = data.indices ? data.indices as number[] : new Int32Array(data.x.length) as any as number[];
-    if (!data.indices) {
-        for (let i = 0, _i = data.x.length; i < _i; i++) indices[i] = i;
-    }
+    const { indices } = data;
 
     const S = Vec3.sub(Vec3.zero(), expandedBox.max, expandedBox.min);
     let delta, size;
 
-    if (indices.length > 0) {
+    const elementCount = OrderedSet.size(indices);
+
+    if (elementCount > 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 V = Math.ceil(elementCount / 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]];
@@ -175,19 +194,35 @@ export function build(data: PositionData) {
         expandedBox,
         boundingBox,
         boundingSphere,
-        count: indices.length,
+        count: elementCount,
         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 }, delta } = structure;
-    //const { radius: r, radiusSq: rSq, pivot: [x, y, z], isCheck, mask } = ctx;
-    const r = radius, rSq = r * r;
+interface QueryContext {
+    structure: Grid3D,
+    x: number,
+    y: number,
+    z: number,
+    radius: number,
+    result: Result,
+    isCheck: boolean
+}
+
+function createContext(structure: Grid3D): QueryContext {
+    return { structure, x: 0.1, y: 0.1, z: 0.1, radius: 0.1, result: Result.create(), isCheck: false }
+}
+
+function query(ctx: QueryContext): boolean {
+    const { min, size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, data: { x: px, y: py, z: pz, indices, radius }, delta, maxRadius } = ctx.structure;
+    const { radius: inputRadius, isCheck, x, y, z, result } = ctx;
+
+    const r = inputRadius + maxRadius;
+    const rSq = r * r;
 
-    const result = Result.create();
+    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]));
@@ -197,7 +232,7 @@ export function inSphere(structure: SpatialHash, x: number, y: number, z: number
     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 result;
+    if (loX > hiX || loY > hiY || loZ > hiZ) return false;
 
     for (let ix = loX; ix <= hiX; ix++) {
         for (let iy = loY; iy <= hiY; iy++) {
@@ -211,7 +246,7 @@ export function inSphere(structure: SpatialHash, x: number, y: number, z: number
                 const end = offset + count;
 
                 for (let i = offset; i < end; i++) {
-                    const idx = indices[bucketArray[i]];
+                    const idx = OrderedSet.getAt(indices, bucketArray[i]);
 
                     const dx = px[idx] - x;
                     const dy = py[idx] - y;
@@ -219,12 +254,13 @@ export function inSphere(structure: SpatialHash, x: number, y: number, z: number
                     const distSq = dx * dx + dy * dy + dz * dz;
 
                     if (distSq <= rSq) {
-                        // if (isCheck) return true;
+                        if (maxRadius > 0 && Math.sqrt(distSq) - radius![idx] > inputRadius) continue;
+                        if (isCheck) return true;
                         Result.add(result, bucketArray[i], distSq);
                     }
                 }
             }
         }
     }
-    return result;
+    return result.count > 0;
 }

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

@@ -6,6 +6,7 @@
 
 import { Vec3 } from '../../linear-algebra'
 import { PositionData } from '../common'
+import { OrderedSet } from 'mol-data/int';
 
 interface Box3D { min: Vec3, max: Vec3 }
 
@@ -15,28 +16,15 @@ namespace Box3D {
         const max = [-Number.MAX_VALUE, -Number.MAX_VALUE, -Number.MAX_VALUE];
 
         const { x, y, z, indices } = data;
-
-        if (indices) {
-            for (let t = 0, _t = indices.length; t < _t; t++) {
-                const i = indices[t];
-                min[0] = Math.min(x[i], min[0]);
-                min[1] = Math.min(y[i], min[1]);
-                min[2] = Math.min(z[i], min[2]);
-                max[0] = Math.max(x[i], max[0]);
-                max[1] = Math.max(y[i], max[1]);
-                max[2] = Math.max(z[i], max[2]);
-            }
-        } else {
-            for (let i = 0, _i = x.length; i < _i; i++) {
-                min[0] = Math.min(x[i], min[0]);
-                min[1] = Math.min(y[i], min[1]);
-                min[2] = Math.min(z[i], min[2]);
-                max[0] = Math.max(x[i], max[0]);
-                max[1] = Math.max(y[i], max[1]);
-                max[2] = Math.max(z[i], max[2]);
-            }
+        for (let t = 0, _t = OrderedSet.size(indices); t < _t; t++) {
+            const i = OrderedSet.getAt(indices, t);
+            min[0] = Math.min(x[i], min[0]);
+            min[1] = Math.min(y[i], min[1]);
+            min[2] = Math.min(z[i], min[2]);
+            max[0] = Math.max(x[i], max[0]);
+            max[1] = Math.max(y[i], max[1]);
+            max[2] = Math.max(z[i], max[2]);
         }
-
         return { min: Vec3.create(min[0], min[1], min[2]), max: Vec3.create(max[0], max[1], max[2]) }
     }
 

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

@@ -6,6 +6,7 @@
 
 import { Vec3 } from '../../linear-algebra'
 import { PositionData } from '../common'
+import { OrderedSet } from 'mol-data/int';
 
 interface Sphere3D { center: Vec3, radius: number }
 
@@ -15,44 +16,25 @@ namespace Sphere3D {
         let cx = 0, cy = 0, cz = 0;
         let radiusSq = 0;
 
-        if (indices) {
-            for (let t = 0, _t = indices.length; t < _t; t++) {
-                const i = indices[t];
-                cx += x[i];
-                cy += y[i];
-                cz += z[i];
-            }
-
-            if (indices.length > 0) {
-                cx /= indices.length;
-                cy /= indices.length;
-                cz /= indices.length;
-            }
-
-            for (let t = 0, _t = indices.length; t < _t; t++) {
-                const i = indices[t];
-                const dx = x[i] - cx, dy = y[i] - cy, dz = z[i] - cz;
-                const d = dx * dx + dy * dy + dz * dz;
-                if (d > radiusSq) radiusSq = d;
-            }
-        } else {
-            for (let i = 0, _i = x.length; i < _i; i++) {
-                cx += x[i];
-                cy += y[i];
-                cz += z[i];
-            }
-
-            if (x.length > 0) {
-                cx /= x.length;
-                cy /= x.length;
-                cz /= x.length;
-            }
-
-            for (let i = 0, _i = x.length; i < _i; i++) {
-                const dx = x[i] - cx, dy = y[i] - cy, dz = z[i] - cz;
-                const d = dx * dx + dy * dy + dz * dz;
-                if (d > radiusSq) radiusSq = d;
-            }
+        const size = OrderedSet.size(indices);
+        for (let t = 0; t < size; t++) {
+            const i = OrderedSet.getAt(indices, t);
+            cx += x[i];
+            cy += y[i];
+            cz += z[i];
+        }
+
+        if (size > 0) {
+            cx /= size;
+            cy /= size;
+            cz /= size;
+        }
+
+        for (let t = 0; t < size; t++) {
+            const i = OrderedSet.getAt(indices, t);
+            const dx = x[i] - cx, dy = y[i] - cy, dz = z[i] - cz;
+            const d = dx * dx + dy * dy + dz * dz;
+            if (d > radiusSq) radiusSq = d;
         }
 
         return { center: Vec3.create(cx, cy, cz), radius: Math.sqrt(radiusSq) };

src/mol-math/geometry/symmetry-operator.ts

@@ -61,7 +61,7 @@ namespace SymmetryOperator {
     }
 }
 
-export default SymmetryOperator
+export { SymmetryOperator }
 
 interface Projections { x(index: number): number, y(index: number): number, z(index: number): number }
 

src/mol-model/structure/model/formats/mmcif/assembly.ts

@@ -5,7 +5,7 @@
  */
 
 import { Mat4, Tensor } from 'mol-math/linear-algebra'
-import SymmetryOperator from 'mol-math/geometry/symmetry-operator'
+import { SymmetryOperator } from 'mol-math/geometry/symmetry-operator'
 import Format from '../../format'
 import { Assembly, OperatorGroup, OperatorGroups } from '../../properties/symmetry'
 import { Queries as Q } from '../../../query'

src/mol-model/structure/model/model.ts

@@ -5,16 +5,12 @@
  */
 
 import UUID from 'mol-util/uuid'
-import GridLookup from 'mol-math/geometry/grid-lookup'
 import Format from './format'
 import Hierarchy from './properties/hierarchy'
 import Conformation from './properties/conformation'
 import Symmetry from './properties/symmetry'
-import Bonds from './properties/bonds'
 import CoarseGrained from './properties/coarse-grained'
 
-import computeBonds from './utils/compute-bonds'
-
 import from_gro from './formats/gro'
 import from_mmCIF from './formats/mmcif'
 
@@ -38,8 +34,7 @@ interface Model extends Readonly<{
 
     atomCount: number,
 }> {
-    '@spatialLookup'?: GridLookup,
-    '@bonds'?: Bonds
+
 } { }
 
 namespace Model {
@@ -49,18 +44,18 @@ namespace Model {
             case 'mmCIF': return from_mmCIF(format);
         }
     }
-    export function spatialLookup(model: Model): GridLookup {
-        if (model['@spatialLookup']) return model['@spatialLookup']!;
-        const lookup = GridLookup(model.conformation);
-        model['@spatialLookup'] = lookup;
-        return lookup;
-    }
-    export function bonds(model: Model): Bonds {
-        if (model['@bonds']) return model['@bonds']!;
-        const bonds = computeBonds(model);
-        model['@bonds'] = bonds;
-        return bonds;
-    }
+    // export function spatialLookup(model: Model): GridLookup {
+    //     if (model['@spatialLookup']) return model['@spatialLookup']!;
+    //     const lookup = GridLookup(model.conformation);
+    //     model['@spatialLookup'] = lookup;
+    //     return lookup;
+    // }
+    // export function bonds(model: Model): Bonds {
+    //     if (model['@bonds']) return model['@bonds']!;
+    //     const bonds = computeBonds(model);
+    //     model['@bonds'] = bonds;
+    //     return bonds;
+    // }
 }
 
 export default Model

src/mol-model/structure/model/properties/symmetry.ts

@@ -4,7 +4,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import SymmetryOperator from 'mol-math/geometry/symmetry-operator'
+import { SymmetryOperator } from 'mol-math/geometry/symmetry-operator'
 import { arrayFind } from 'mol-data/util'
 import { Query } from '../../query'
 import { Model } from '../../model'

src/mol-model/structure/structure/element/group.ts

@@ -5,12 +5,15 @@
  */
 
 import { OrderedSet } from 'mol-data/int'
+import { Lookup3D, GridLookup3D } from 'mol-math/geometry';
 import Unit from '../unit'
 
 interface ElementGroup {
     elements: OrderedSet,
     // Unique identifier of the group, usable as partial key for various "caches".
-    key: number
+    key: number,
+
+    __lookup3d__?: Lookup3D
 }
 
 namespace ElementGroup {
@@ -61,6 +64,17 @@ namespace ElementGroup {
         return createNew(set);
     }
 
+    export function getLookup3d(unit: Unit, group: ElementGroup) {
+        if (group.__lookup3d__)  return group.__lookup3d__;
+        if (Unit.isAtomic(unit)) {
+            const { x, y, z } = unit.model.conformation;
+            group.__lookup3d__ = GridLookup3D({ x, y, z, indices: group.elements });
+            return group.__lookup3d__;
+        }
+
+        throw 'not implemented';
+    }
+
     let _id = 0;
     function nextKey() {
         const ret = _id;

src/mol-model/structure/structure/structure.ts

@@ -6,7 +6,7 @@
 
 import { OrderedSet, Iterator } from 'mol-data/int'
 import { UniqueArray } from 'mol-data/generic'
-import SymmetryOperator from 'mol-math/geometry/symmetry-operator'
+import { SymmetryOperator } from 'mol-math/geometry/symmetry-operator'
 import { Model, Format } from '../model'
 import Unit from './unit'
 import ElementSet from './element/set'

src/mol-model/structure/structure/unit.ts

@@ -4,7 +4,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import SymmetryOperator from 'mol-math/geometry/symmetry-operator'
+import { SymmetryOperator } from 'mol-math/geometry/symmetry-operator'
 import ElementGroup from './element/group'
 import { Model } from '../model'
 

src/perf-tests/lookup3d.ts

@@ -5,8 +5,9 @@ import CIF from 'mol-io/reader/cif'
 import { Structure, Model } from 'mol-model/structure'
 
 import { Run } from 'mol-task';
-import { build, inSphere } from 'mol-math/geometry/lookup3d/grid';
+import { GridLookup3D } from 'mol-math/geometry';
 import { sortArray } from 'mol-data/util';
+import { OrderedSet } from 'mol-data/int';
 
 require('util.promisify').shim();
 const readFileAsync = util.promisify(fs.readFile);
@@ -42,13 +43,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: [0, 1, 2, 3]
+    const lookup = GridLookup3D({ x: mmcif.atom_site.Cartn_x.toArray(), y: mmcif.atom_site.Cartn_y.toArray(), z: mmcif.atom_site.Cartn_z.toArray(),
+        indices: OrderedSet.ofBounds(0, 4),
+        radius: [1, 1, 1, 1]
         //indices: [1]
     });
-    console.log(data.boundingBox, data.boundingSphere);
+    console.log(lookup.boundary.box, lookup.boundary.sphere);
 
-    const result = inSphere(data, -30.07, 8.178, -13.897, 10);
+    const result = lookup.find(-30.07, 8.178, -13.897, 1);
     console.log(result.count, sortArray(result.indices));
 }