Structure lookup 3D

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

@@ -16,7 +16,7 @@ namespace Element {
     export const create: (unit: number, index: number) => Element = Tuple.create;
     export const is: (x: any) => x is Element = Tuple.is;
     export const unitId: (e: Element) => number = Tuple.fst;
-    export const index: (e: Element) => number = Tuple.snd;
+    export const elementIndex: (e: Element) => number = Tuple.snd;
     export const areEqual: (e: Element, b: Element) => boolean = Tuple.areEqual;
     export const hashCode: (e: Element) => number = Tuple.hashCode;
 
@@ -30,7 +30,7 @@ namespace Element {
 
     export function updateLocation(structure: Structure, l: Location, element: Element) {
         l.unit = structure.units[unitId(element)];
-        l.element = index(element);
+        l.element = elementIndex(element);
         return l;
     }
 

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

@@ -12,7 +12,6 @@ import { sortArray, sort, arraySwap, hash1 } from 'mol-data/util';
 import Element from './element'
 import Unit from './unit'
 import { StructureLookup3D } from './util/lookup3d';
-import { computeStructureBoundary } from './util/boundary';
 
 class Structure {
     readonly unitMap: IntMap<Unit>;
@@ -48,7 +47,14 @@ class Structure {
     }
 
     get boundary() {
-        return computeStructureBoundary(this);
+        return this.lookup3d.boundary;
+    }
+
+    private _lookup3d?: StructureLookup3D = void 0;
+    get lookup3d() {
+        if (this._lookup3d) return this._lookup3d;
+        this._lookup3d = StructureLookup3D.create(this);
+        return this._lookup3d;
     }
 
     constructor(units: ArrayLike<Unit>) {

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

@@ -6,93 +6,85 @@
 
 import Structure from '../structure'
 import Element from '../element'
-import { Lookup3D, /*GridLookup3D,*/ Result, Box3D, Sphere3D } from 'mol-math/geometry';
-//import { ElementSet, Unit } from '../../structure';
-//import { Vec3 } from 'mol-math/linear-algebra';
-//import { OrderedSet } from 'mol-data/int';
-//import { computeStructureBoundary } from './boundary';
+import { Lookup3D, GridLookup3D, Result, Box3D, Sphere3D } from 'mol-math/geometry';
+import { Vec3 } from 'mol-math/linear-algebra';
+import { computeStructureBoundary } from './boundary';
+import { OrderedSet } from 'mol-data/int';
 
 interface StructureLookup3D extends Lookup3D<Element> {}
 
 namespace StructureLookup3D {
     class Impl implements StructureLookup3D {
-        //private unitLookup: Lookup3D;
+        private unitLookup: Lookup3D;
         private result = Result.create<Element>();
-        //private pivot = Vec3.zero();
+        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;
-
-            // for (let t = 0, _t = closeUnits.count; t < _t; t++) {
-            //     const i = closeUnits.indices[t];
-            //     const unitId = ElementSet.groupUnitIndex(elements, i);
-            //     const group = ElementSet.groupAt(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 = Unit.getLookup3d(unit, group);
-            //     const groupResult = groupLookup.find(this.pivot[0], this.pivot[1], this.pivot[2], radius);
-            //     for (let j = 0, _j = groupResult.count; j < _j; j++) {
-            //         Result.add(this.result, Element.create(unitId, groupResult.indices[j]), groupResult.squaredDistances[j]);
-            //     }
-            // }
+            const { units } = this.structure;
+            const closeUnits = this.unitLookup.find(x, y, z, radius);
+            if (closeUnits.count === 0) return this.result;
+
+            for (let t = 0, _t = closeUnits.count; t < _t; t++) {
+                const unit = units[closeUnits.indices[t]];
+                Vec3.set(this.pivot, x, y, z);
+                if (!unit.conformation.operator.isIdentity) {
+                    Vec3.transformMat4(this.pivot, this.pivot, unit.conformation.operator.inverse);
+                }
+                const unitLookup = unit.lookup3d;
+                const groupResult = unitLookup.find(this.pivot[0], this.pivot[1], this.pivot[2], radius);
+                for (let j = 0, _j = groupResult.count; j < _j; j++) {
+                    Result.add(this.result, Element.create(unit.id, groupResult.indices[j]), groupResult.squaredDistances[j]);
+                }
+            }
 
             return this.result;
         }
 
         check(x: number, y: number, z: number, radius: number): boolean {
-            // const { units, elements } = this.structure;
-            // const closeUnits = this.unitLookup.find(x, y, z, radius);
-            // if (closeUnits.count === 0) return false;
-
-            // for (let t = 0, _t = closeUnits.count; t < _t; t++) {
-            //     const i = closeUnits.indices[t];
-            //     const unitId = ElementSet.groupUnitIndex(elements, i);
-            //     const group = ElementSet.groupAt(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 = Unit.getLookup3d(unit, group);
-            //     if (groupLookup.check(this.pivot[0], this.pivot[1], this.pivot[2], radius)) return true;
-            // }
+            const { units } = this.structure;
+            const closeUnits = this.unitLookup.find(x, y, z, radius);
+            if (closeUnits.count === 0) return false;
+
+            for (let t = 0, _t = closeUnits.count; t < _t; t++) {
+                const unit = units[closeUnits.indices[t]];
+                Vec3.set(this.pivot, x, y, z);
+                if (!unit.conformation.operator.isIdentity) {
+                    Vec3.transformMat4(this.pivot, this.pivot, unit.conformation.operator.inverse);
+                }
+                const groupLookup = unit.lookup3d;
+                if (groupLookup.check(this.pivot[0], this.pivot[1], this.pivot[2], radius)) return true;
+            }
 
             return false;
         }
 
         boundary: { box: Box3D; sphere: Sphere3D; };
 
-        constructor(structure: Structure) {
-            // const { units, elements } = structure;
-            // const unitCount = ElementSet.groupCount(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.groupAt(elements, i);
-            //     const unit = units[ElementSet.groupUnitIndex(elements, i)];
-            //     const lookup = Unit.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;
-            // }
-
-            // this.unitLookup = GridLookup3D({ x: xs, y: ys, z: zs, radius, indices: OrderedSet.ofBounds(0, unitCount) });
-            // this.boundary = computeStructureBoundary(structure);
+        constructor(private structure: Structure) {
+            const { units } = structure;
+            const unitCount = units.length;
+            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 unit = units[i];
+                const lookup = unit.lookup3d;
+                const s = lookup.boundary.sphere;
+
+                Vec3.transformMat4(center, s.center, unit.conformation.operator.matrix);
+
+                xs[i] = center[0];
+                ys[i] = center[1];
+                zs[i] = center[2];
+                radius[i] = s.radius;
+            }
+
+            this.unitLookup = GridLookup3D({ x: xs, y: ys, z: zs, radius, indices: OrderedSet.ofBounds(0, unitCount) });
+            this.boundary = computeStructureBoundary(structure);
         }
     }