Structure level lookup3d (part 1)

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

@@ -7,30 +7,30 @@
 import { Box3D } from '../primitives/box3d'
 import { Sphere3D } from '../primitives/sphere3d'
 
-export interface Result {
+export interface Result<T> {
     count: number,
-    indices: number[],
+    indices: T[],
     squaredDistances: number[]
 }
 
 export namespace Result {
-    export function add(result: Result, index: number, distSq: number) {
+    export function add<T>(result: Result<T>, index: T, distSq: number) {
         result.squaredDistances[result.count] = distSq;
         result.indices[result.count++] = index;
     }
 
-    export function reset(result: Result) {
+    export function reset(result: Result<any>) {
         result.count = 0;
     }
 
-    export function create(): Result {
+    export function create<T>(): Result<T> {
         return { count: 0, indices: [], squaredDistances: [] };
     }
 }
 
-export interface Lookup3D {
+export interface Lookup3D<T = number> {
     // The result is mutated with each call to find.
-    find(x: number, y: number, z: number, radius: number): Result,
+    find(x: number, y: number, z: number, radius: number): Result<T>,
     check(x: number, y: number, z: number, radius: number): boolean,
     boundary: { box: Box3D, sphere: Sphere3D }
 }

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

@@ -18,11 +18,11 @@ function GridLookup3D(data: PositionData): Lookup3D {
 
 export { GridLookup3D }
 
-class GridLookup3DImpl implements Lookup3D {
+class GridLookup3DImpl implements Lookup3D<number> {
     private ctx: QueryContext;
     boundary: Lookup3D['boundary'];
 
-    find(x: number, y: number, z: number, radius: number): Result {
+    find(x: number, y: number, z: number, radius: number): Result<number> {
         this.ctx.x = x;
         this.ctx.y = y;
         this.ctx.z = z;
@@ -207,7 +207,7 @@ interface QueryContext {
     y: number,
     z: number,
     radius: number,
-    result: Result,
+    result: Result<number>,
     isCheck: boolean
 }
 

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

@@ -12,12 +12,15 @@ import Unit from './unit'
 import ElementSet from './element/set'
 import ElementGroup from './element/group'
 import Element from './element'
+import { StructureLookup3D } from './util/lookup3d';
 
 // A structure is a pair of "units" and an element set.
 // Each unit contains the data and transformation of its corresponding elements.
 interface Structure {
     readonly units: ReadonlyArray<Unit>,
-    readonly elements: ElementSet
+    readonly elements: ElementSet,
+
+    __lookup3d__?: StructureLookup3D
 }
 
 namespace Structure {
@@ -82,6 +85,12 @@ namespace Structure {
         return arr.array;
     }
 
+    export function getLookup3d(s: Structure) {
+        if (s.__lookup3d__) return s.__lookup3d__;
+        s.__lookup3d__ = StructureLookup3D.create(s);
+        return s.__lookup3d__;
+    }
+
     // TODO: "lift" atom set operators?
     // TODO: "diff"
 }

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

@@ -0,0 +1,89 @@
+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import Structure from '../structure'
+import Element from '../element'
+import { Lookup3D, GridLookup3D, Result, Box3D, Sphere3D } from 'mol-math/geometry';
+import { ElementGroup, ElementSet } from '../../structure';
+import { Vec3 } from 'mol-math/linear-algebra';
+import { OrderedSet } from 'mol-data/int';
+
+interface StructureLookup3D extends Lookup3D<Element> {}
+
+namespace StructureLookup3D {
+    class Impl implements StructureLookup3D {
+        private unitLookup: Lookup3D;
+        private result = Result.create<Element>();
+        private pivot = Vec3.zero();
+
+        find(x: number, y: number, z: number, radius: number): Result<Element> {
+            Result.reset(this.result);
+            const { units, elements } = this.structure;
+            const closeUnits = this.unitLookup.find(x, y, z, radius);
+            if (closeUnits.count === 0) return this.result;
+
+            console.log({ closeUnits });
+
+            for (let t = 0, _t = closeUnits.count; t < _t; t++) {
+                const i = closeUnits.indices[t];
+                const unitId = ElementSet.unitGetId(elements, i);
+                const group = ElementSet.unitGetByIndex(elements, i);
+                const unit = units[unitId];
+                Vec3.set(this.pivot, x, y, z);
+                if (!unit.operator.isIdentity) {
+                    Vec3.transformMat4(this.pivot, this.pivot, unit.operator.inverse);
+                }
+                const groupLookup = ElementGroup.getLookup3d(unit, group);
+                const groupResult = groupLookup.find(this.pivot[0], this.pivot[1], this.pivot[2], radius);
+                //console.log(groupLookup);
+                //console.log({ groupCount: groupResult.count });
+                for (let j = 0, _j = groupResult.count; j < _j; j++) {
+                    Result.add(this.result, Element.create(unitId, groupResult.indices[j]), groupResult.squaredDistances[j]);
+                }
+            }
+
+            return this.result;
+        }
+        check(x: number, y: number, z: number, radius: number): boolean {
+            throw new Error("Method not implemented.");
+        }
+        boundary: { box: Box3D; sphere: Sphere3D; } = 0 as any;
+
+        constructor(private structure: Structure) {
+            const { units, elements } = structure;
+            const unitCount = ElementSet.unitCount(elements);
+            const xs = new Float32Array(unitCount);
+            const ys = new Float32Array(unitCount);
+            const zs = new Float32Array(unitCount);
+            const radius = new Float32Array(unitCount);
+
+            const center = Vec3.zero();
+            for (let i = 0; i < unitCount; i++) {
+                const group = ElementSet.unitGetByIndex(elements, i);
+                const unit = units[ElementSet.unitGetId(elements, i)];
+                const lookup = ElementGroup.getLookup3d(unit, group);
+                const s = lookup.boundary.sphere;
+
+                Vec3.transformMat4(center, s.center, unit.operator.matrix);
+
+                xs[i] = center[0];
+                ys[i] = center[1];
+                zs[i] = center[2];
+                radius[i] = s.radius;
+            }
+
+            console.log({ xs, ys, zs, radius, unitCount })
+
+            this.unitLookup = GridLookup3D({ x: xs, y: ys, z: zs, radius, indices: OrderedSet.ofBounds(0, unitCount) });
+        }
+    }
+
+    export function create(s: Structure): StructureLookup3D {
+        return new Impl(s);
+    }
+}
+
+export { StructureLookup3D }

src/perf-tests/lookup3d.ts

@@ -41,17 +41,21 @@ export async function readCIF(path: string) {
 }
 
 export async function test() {
-    const { mmcif  } = await readCIF('e:/test/quick/1tqn_updated.cif');
+    const { mmcif, structures } = await readCIF('e:/test/quick/1tqn_updated.cif');
 
     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: OrderedSet.ofBounds(0, mmcif.atom_site._rowCount),
+        //radius: [1, 1, 1, 1]
         //indices: [1]
     });
     console.log(lookup.boundary.box, lookup.boundary.sphere);
 
-    const result = lookup.find(-30.07, 8.178, -13.897, 1);
-    console.log(result.count, sortArray(result.indices));
+    const result = lookup.find(-30.07, 8.178, -13.897, 10);
+    console.log(result.count)//, sortArray(result.indices));
+
+    const sl = Structure.getLookup3d(structures[0]);
+    const result1 = sl.find(-30.07, 8.178, -13.897, 10);
+    console.log(result1.count);//, result1.indices);
 }
 
 test();