refactoring grid lookup (step 1)

src/mol-math/geometry/common.ts

@@ -0,0 +1,14 @@
+/*
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+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>,
+}

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

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Box3D } from '../primitives/box3d'
+import { Sphere3D } from '../primitives/sphere3d'
+
+export interface Result {
+    count: number,
+    indices: number[],
+    squaredDistances: number[]
+}
+
+export namespace Result {
+    export function add(result: Result, index: number, distSq: number) {
+        result.squaredDistances[result.count] = distSq;
+        result.indices[result.count++] = index;
+    }
+
+    export function reset(result: Result) {
+        result.count = 0;
+    }
+
+    export function create(): Result {
+        return { count: 0, indices: [], squaredDistances: [] };
+    }
+}
+
+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
+}

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

@@ -0,0 +1,220 @@
+/*
+ * Copyright (c) 2018 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>
+ */
+
+import { Result/*, Lookup3D*/ } from './common'
+import { Box3D } from '../primitives/box3d';
+import { Sphere3D } from '../primitives/sphere3d';
+import { PositionData } from '../common';
+import { Vec3 } from '../../linear-algebra';
+
+// 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;
+// }
+
+/**
+ * Adapted from https://github.com/arose/ngl
+ * MIT License Copyright (C) 2014+ Alexander Rose
+ */
+
+interface InputData {
+    x: ArrayLike<number>,
+    y: ArrayLike<number>,
+    z: ArrayLike<number>,
+    radius?: ArrayLike<number>,
+    indices: ArrayLike<number>
+}
+
+interface SpatialHash {
+    size: number[],
+    min: number[],
+    grid: Uint32Array,
+    divisor: number[],
+    bucketOffset: Uint32Array,
+    bucketCounts: Int32Array,
+    bucketArray: Int32Array,
+    data: InputData,
+    maxRadius: number,
+    boundingBox: Box3D,
+    boundingSphere: Sphere3D
+}
+
+interface BuildState {
+    size: number[],
+    divisor: number[],
+    data: InputData,
+    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 n = sX * sY * sZ;
+    const { min: [minX, minY, minZ] } = boundingBox;
+
+    let maxRadius = 0;
+    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];
+        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 idx = (((x * sY) + y) * sZ) + z;
+        if ((grid[idx] += 1) === 1) {
+            bucketCount += 1
+        }
+        bucketIndex[t] = idx;
+    }
+
+    if (radius) {
+        for (let t = 0, _t = indices.length; t < _t; t++) {
+            const i = 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(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,
+        divisor,
+        min: state.boundingBox.min,
+        data: state.data,
+        maxRadius,
+        boundingBox: state.boundingBox,
+        boundingSphere: state.boundingSphere
+    }
+}
+
+export function build(data: PositionData) {
+    const boundingBox = Box3D.computeBounding(data);
+    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;
+    }
+
+    // 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 inputData: InputData = {
+        x: data.x,
+        y: data.y,
+        z: data.z,
+        indices,
+        radius: data.radius
+    };
+
+    const state: BuildState = {
+        size,
+        data: inputData,
+        boundingBox,
+        boundingSphere,
+        count: indices.length,
+        divisor
+    }
+
+    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 { 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 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]));
+
+    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 = 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 (isCheck) return true;
+                            Result.add(result, idx, distSq);
+                        }
+                    }
+                }
+            }
+        }
+    }
+    return result;
+}

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Vec3 } from '../../linear-algebra'
+import { PositionData } from '../common'
+
+interface Box3D { min: Vec3, max: Vec3 }
+
+namespace Box3D {
+    export function computeBounding(data: PositionData): Box3D {
+        const min = [Number.MAX_VALUE, Number.MAX_VALUE, Number.MAX_VALUE];
+        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]);
+            }
+        }
+
+        return { min: Vec3.create(min[0], min[1], min[2]), max: Vec3.create(max[0], max[1], max[2]) }
+    }
+}
+
+export { Box3D }

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

@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Vec3 } from '../../linear-algebra'
+import { PositionData } from '../common'
+
+interface Sphere3D { center: Vec3, radius: number }
+
+namespace Sphere3D {
+    export function computeBounding(data: PositionData): Sphere3D {
+        const { x, y, z, indices } = data;
+        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;
+            }
+        }
+
+        return { center: Vec3.create(cx, cy, cz), radius: Math.sqrt(radiusSq) };
+    }
+}
+
+export { Sphere3D }

src/mol-math/geometry/sphere.ts

@@ -1 +0,0 @@
-// TODO: rebranded vec4

src/perf-tests/lookup3d.ts

@@ -0,0 +1,51 @@
+import * as util from 'util'
+import * as fs from 'fs'
+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';
+
+require('util.promisify').shim();
+const readFileAsync = util.promisify(fs.readFile);
+
+async function readData(path: string) {
+    if (path.match(/\.bcif$/)) {
+        const input = await readFileAsync(path)
+        const data = new Uint8Array(input.byteLength);
+        for (let i = 0; i < input.byteLength; i++) data[i] = input[i];
+        return data;
+    } else {
+        return readFileAsync(path, 'utf8');
+    }
+}
+
+
+export async function readCIF(path: string) {
+    const input = await readData(path)
+    const comp = typeof input === 'string' ? CIF.parseText(input) : CIF.parseBinary(input);
+    const parsed = await Run(comp);
+    if (parsed.isError) {
+        throw parsed;
+    }
+
+    const data = parsed.result.blocks[0];
+    const mmcif = CIF.schema.mmCIF(data);
+    const models = Model.create({ kind: 'mmCIF', data: mmcif });
+    const structures = models.map(Structure.ofModel);
+
+    return { mmcif, models, structures };
+}
+
+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 result = inSphere(data, -30.07, 8.178, -13.897, 10);
+    console.log(result.count);
+}
+
+test();