support nested Lookup3D queries

- fixes non-covalent interactions bug

CHANGELOG.md

@@ -11,6 +11,7 @@ Note that since we don't clearly distinguish between a public and private interf
     - Added NOS-bridges query & improved disulfide-bridges query
 - Fix #178: ``IndexPairBonds`` for non-single residue structures (bug due to atom reordering).
 - Add volumetric color smoothing for MolecularSurface and GaussianSurface representations (#173)
+- Fix nested 3d grid lookup that caused results being overwritten in non-covalent interactions computation.
 
 ## [v2.0.5] - 2021-04-26
 

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

@@ -40,7 +40,7 @@ export namespace Result {
 
 export interface Lookup3D<T = number> {
     // The result is mutated with each call to find.
-    find(x: number, y: number, z: number, radius: number): Result<T>,
+    find(x: number, y: number, z: number, radius: number, result?: Result<T>): Result<T>,
     check(x: number, y: number, z: number, radius: number): boolean,
     readonly boundary: { readonly box: Box3D, readonly sphere: Sphere3D }
     /** transient result */

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

@@ -24,19 +24,20 @@ function GridLookup3D<T extends number = number>(data: PositionData, boundary: B
 export { GridLookup3D };
 
 class GridLookup3DImpl<T extends number = number> implements GridLookup3D<T> {
-    private ctx: QueryContext<T>;
+    private ctx: QueryContext;
     boundary: Lookup3D['boundary'];
     buckets: GridLookup3D['buckets'];
     result: Result<T>
 
-    find(x: number, y: number, z: number, radius: number): Result<T> {
+    find(x: number, y: number, z: number, radius: number, result?: Result<T>): Result<T> {
         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;
+        const ret = result ?? this.result;
+        query(this.ctx, ret);
+        return ret;
     }
 
     check(x: number, y: number, z: number, radius: number): boolean {
@@ -45,15 +46,15 @@ class GridLookup3DImpl<T extends number = number> implements GridLookup3D<T> {
         this.ctx.z = z;
         this.ctx.radius = radius;
         this.ctx.isCheck = true;
-        return query(this.ctx);
+        return query(this.ctx, this.result);
     }
 
     constructor(data: PositionData, boundary: Boundary, cellSizeOrCount?: Vec3 | number) {
         const structure = build(data, boundary, cellSizeOrCount);
-        this.ctx = createContext<T>(structure);
+        this.ctx = createContext(structure);
         this.boundary = { box: structure.boundingBox, sphere: structure.boundingSphere };
         this.buckets = { offset: structure.bucketOffset, count: structure.bucketCounts, array: structure.bucketArray };
-        this.result = this.ctx.result;
+        this.result = Result.create();
     }
 }
 
@@ -215,23 +216,22 @@ function build(data: PositionData, boundary: Boundary, cellSizeOrCount?: Vec3 |
     return _build(state);
 }
 
-interface QueryContext<T extends number = number> {
+interface QueryContext {
     grid: Grid3D,
     x: number,
     y: number,
     z: number,
     radius: number,
-    result: Result<T>,
     isCheck: boolean
 }
 
-function createContext<T extends number = number>(grid: Grid3D): QueryContext<T> {
-    return { grid, x: 0.1, y: 0.1, z: 0.1, radius: 0.1, result: Result.create(), isCheck: false };
+function createContext(grid: Grid3D): QueryContext {
+    return { grid, x: 0.1, y: 0.1, z: 0.1, radius: 0.1, isCheck: false };
 }
 
-function query<T extends number = number>(ctx: QueryContext<T>): boolean {
+function query<T extends number = number>(ctx: QueryContext, result: Result<T>): boolean {
     const { min, size: [sX, sY, sZ], bucketOffset, bucketCounts, bucketArray, grid, data: { x: px, y: py, z: pz, indices, radius }, delta, maxRadius } = ctx.grid;
-    const { radius: inputRadius, isCheck, x, y, z, result } = ctx;
+    const { radius: inputRadius, isCheck, x, y, z } = ctx;
 
     const r = inputRadius + maxRadius;
     const rSq = r * r;

src/mol-model-props/computed/interactions/contacts.ts

@@ -4,16 +4,17 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { ParamDefinition as PD } from '../../../mol-util/param-definition';
-import { Structure, Unit, StructureElement } from '../../../mol-model/structure';
-import { Features } from './features';
-import { InteractionType, FeatureType } from './common';
-import { IntraContactsBuilder, InterContactsBuilder } from './contacts-builder';
-import { Mat4, Vec3 } from '../../../mol-math/linear-algebra';
-import { altLoc, connectedTo, typeSymbol } from '../chemistry/util';
 import { OrderedSet } from '../../../mol-data/int';
+import { Mat4, Vec3 } from '../../../mol-math/linear-algebra';
+import { Structure, StructureElement, Unit } from '../../../mol-model/structure';
 import { VdwRadius } from '../../../mol-model/structure/model/properties/atomic';
 import { Elements } from '../../../mol-model/structure/model/properties/atomic/types';
+import { StructureLookup3DResultContext } from '../../../mol-model/structure/structure/util/lookup3d';
+import { ParamDefinition as PD } from '../../../mol-util/param-definition';
+import { altLoc, connectedTo, typeSymbol } from '../chemistry/util';
+import { FeatureType, InteractionType } from './common';
+import { InterContactsBuilder, IntraContactsBuilder } from './contacts-builder';
+import { Features } from './features';
 
 export const ContactsParams = {
     lineOfSightDistFactor: PD.Numeric(1.0, { min: 0, max: 3, step: 0.1 }),
@@ -53,7 +54,7 @@ function validPair(structure: Structure, infoA: Features.Info, infoB: Features.I
 
 //
 
-function invalidAltLoc (unitA: Unit.Atomic, indexA: StructureElement.UnitIndex, unitB: Unit.Atomic, indexB: StructureElement.UnitIndex) {
+function invalidAltLoc(unitA: Unit.Atomic, indexA: StructureElement.UnitIndex, unitB: Unit.Atomic, indexB: StructureElement.UnitIndex) {
     const altA = altLoc(unitA, indexA);
     const altB = altLoc(unitB, indexB);
     return altA && altB && altA !== altB;
@@ -71,6 +72,9 @@ const tmpVec = Vec3();
 const tmpVecA = Vec3();
 const tmpVecB = Vec3();
 
+// need to use a separate context for structure.lookup3d.find because of nested queries
+const lineOfSightLookupCtx = StructureLookup3DResultContext();
+
 function checkLineOfSight(structure: Structure, infoA: Features.Info, infoB: Features.Info, distFactor: number) {
     const featureA = infoA.feature;
     const featureB = infoB.feature;
@@ -83,7 +87,7 @@ function checkLineOfSight(structure: Structure, infoA: Features.Info, infoB: Fea
 
     const distMax = distFactor * MAX_LINE_OF_SIGHT_DISTANCE;
 
-    const { count, indices, units, squaredDistances } = structure.lookup3d.find(tmpVec[0], tmpVec[1], tmpVec[2], distMax);
+    const { count, indices, units, squaredDistances } = structure.lookup3d.find(tmpVec[0], tmpVec[1], tmpVec[2], distMax, lineOfSightLookupCtx);
     if (count === 0) return true;
 
     for (let r = 0; r < count; ++r) {

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

@@ -12,6 +12,7 @@ import { OrderedSet } from '../../../../mol-data/int';
 import { StructureUniqueSubsetBuilder } from './unique-subset-builder';
 import { StructureElement } from '../element';
 import { Unit } from '../unit';
+import { UnitIndex } from '../element/util';
 
 export interface StructureResult extends Result<StructureElement.UnitIndex> {
     units: Unit[]
@@ -40,6 +41,16 @@ export namespace StructureResult {
     }
 }
 
+export interface StructureLookup3DResultContext {
+    result: StructureResult,
+    closeUnitsResult: Result<number>,
+    unitGroupResult: Result<UnitIndex>,
+}
+
+export function StructureLookup3DResultContext(): StructureLookup3DResultContext {
+    return { result: StructureResult.create(), closeUnitsResult: Result.create(), unitGroupResult: Result.create() };
+}
+
 export class StructureLookup3D {
     private unitLookup: Lookup3D;
     private pivot = Vec3();
@@ -48,12 +59,17 @@ export class StructureLookup3D {
         return this.unitLookup.find(x, y, z, radius);
     }
 
-    private result = StructureResult.create();
-    find(x: number, y: number, z: number, radius: number): StructureResult {
-        Result.reset(this.result);
+    private findContext = StructureLookup3DResultContext();
+
+    find(x: number, y: number, z: number, radius: number, ctx?: StructureLookup3DResultContext): StructureResult {
+        return this._find(x, y, z, radius, ctx ?? this.findContext);
+    }
+
+    private _find(x: number, y: number, z: number, radius: number, ctx: StructureLookup3DResultContext): StructureResult {
+        Result.reset(ctx.result);
         const { units } = this.structure;
-        const closeUnits = this.unitLookup.find(x, y, z, radius);
-        if (closeUnits.count === 0) return this.result;
+        const closeUnits = this.unitLookup.find(x, y, z, radius, ctx.closeUnitsResult);
+        if (closeUnits.count === 0) return ctx.result;
 
         for (let t = 0, _t = closeUnits.count; t < _t; t++) {
             const unit = units[closeUnits.indices[t]];
@@ -62,12 +78,12 @@ export class StructureLookup3D {
                 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);
+            const groupResult = unitLookup.find(this.pivot[0], this.pivot[1], this.pivot[2], radius, ctx.unitGroupResult);
             for (let j = 0, _j = groupResult.count; j < _j; j++) {
-                StructureResult.add(this.result, unit, groupResult.indices[j], groupResult.squaredDistances[j]);
+                StructureResult.add(ctx.result, unit, groupResult.indices[j], groupResult.squaredDistances[j]);
             }
         }
-        return this.result;
+        return ctx.result;
     }
 
     findIntoBuilder(x: number, y: number, z: number, radius: number, builder: StructureUniqueSubsetBuilder) {