wip, interactions

- improved performance by searching for links in only the subset of relevant Features

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

@@ -378,11 +378,11 @@ export const AromaticRingProvider = Features.Provider([FeatureType.AromaticRing]
 export const IonicProvider: LinkProvider<IonicParams> = {
     name: 'ionic',
     params: IonicParams,
-    requiredFeatures: [FeatureType.NegativeCharge, FeatureType.PositiveCharge],
     createTester: (props: IonicProps) => {
         const opts = getIonicOptions(props)
         return {
-            maxDistanceSq: opts.distanceMaxSq,
+            maxDistance: props.distanceMax,
+            requiredFeatures: new Set([FeatureType.NegativeCharge, FeatureType.PositiveCharge]),
             getType: (structure, infoA, infoB, distanceSq) => testIonic(structure, infoA, infoB, distanceSq, opts)
         }
     }
@@ -391,11 +391,11 @@ export const IonicProvider: LinkProvider<IonicParams> = {
 export const PiStackingProvider: LinkProvider<PiStackingParams> = {
     name: 'pi-stacking',
     params: PiStackingParams,
-    requiredFeatures: [FeatureType.AromaticRing],
     createTester: (props: PiStackingProps) => {
         const opts = getPiStackingOptions(props)
         return {
-            maxDistanceSq: props.distanceMax * props.distanceMax,
+            maxDistance: props.distanceMax,
+            requiredFeatures: new Set([FeatureType.AromaticRing]),
             getType: (structure, infoA, infoB, distanceSq) => testPiStacking(structure, infoA, infoB, distanceSq, opts)
         }
     }
@@ -404,11 +404,11 @@ export const PiStackingProvider: LinkProvider<PiStackingParams> = {
 export const CationPiProvider: LinkProvider<CationPiParams> = {
     name: 'cation-pi',
     params: CationPiParams,
-    requiredFeatures: [FeatureType.AromaticRing, FeatureType.PositiveCharge],
     createTester: (props: CationPiProps) => {
         const opts = getCationPiOptions(props)
         return {
-            maxDistanceSq: props.distanceMax * props.distanceMax,
+            maxDistance: props.distanceMax,
+            requiredFeatures: new Set([FeatureType.AromaticRing, FeatureType.PositiveCharge]),
             getType: (structure, infoA, infoB, distanceSq) => testCationPi(structure, infoA, infoB, distanceSq, opts)
         }
     }

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

@@ -7,7 +7,7 @@
 import { StructureElement, Unit, Structure } from '../../../mol-model/structure/structure';
 import { ChunkedArray } from '../../../mol-data/util';
 import { GridLookup3D } from '../../../mol-math/geometry';
-import { OrderedSet } from '../../../mol-data/int';
+import { OrderedSet, SortedArray } from '../../../mol-data/int';
 import { FeatureGroup, FeatureType } from './common';
 import { ValenceModelProvider } from '../valence-model';
 
@@ -32,6 +32,8 @@ interface Features {
     readonly lookup3d: GridLookup3D
     /** maps unit elements to features, range for unit element i is offsets[i] to offsets[i + 1] */
     readonly elementsIndex: Features.ElementsIndex
+
+    subset(types: ReadonlySet<FeatureType>): Features.Subset
 }
 
 namespace Features {
@@ -57,6 +59,11 @@ namespace Features {
         members: ArrayLike<StructureElement.UnitIndex>
     }
 
+    export type Subset = {
+        readonly indices: OrderedSet
+        readonly lookup3d: GridLookup3D
+    }
+
     export function createElementsIndex(data: Data, elementsCount: number): ElementsIndex {
         const offsets = new Int32Array(elementsCount + 1)
         const bucketFill = new Int32Array(elementsCount)
@@ -96,6 +103,26 @@ namespace Features {
             get elementsIndex() {
                 return elementsIndex || (elementsIndex = createElementsIndex(data, elementsCount))
             },
+
+            subset: (types: Set<FeatureType>) => createSubset(data, types)
+        }
+    }
+
+    export function createSubset(data: Data, types: ReadonlySet<FeatureType>): Subset {
+        let lookup3d: GridLookup3D
+
+        const { count, types: _types } = data
+        const _indices = []
+        for (let i = 0; i < count; ++i) {
+            if (types.has(_types[i])) _indices.push(i)
+        }
+        const indices = SortedArray.ofSortedArray(_indices)
+
+        return {
+            indices,
+            get lookup3d() {
+                return lookup3d || (lookup3d = GridLookup3D({ x: data.x, y: data.y, z: data.z, indices }))
+            }
         }
     }
 

src/mol-model-props/computed/interactions/halogen-bonds.ts

@@ -118,11 +118,11 @@ export const HalogenAcceptorProvider = Features.Provider([FeatureType.HalogenAcc
 export const HalogenBondsProvider: LinkProvider<HalogenBondsParams> = {
     name: 'halogen-bonds',
     params: HalogenBondsParams,
-    requiredFeatures: [FeatureType.HalogenDonor, FeatureType.HalogenAcceptor],
     createTester: (props: HalogenBondsProps) => {
         const opts = getOptions(props)
         return {
-            maxDistanceSq: props.distanceMax * props.distanceMax,
+            maxDistance: props.distanceMax,
+            requiredFeatures: new Set([FeatureType.HalogenDonor, FeatureType.HalogenAcceptor]),
             getType: (structure, infoA, infoB) => testHalogenBond(structure, infoA, infoB, opts)
         }
     }

src/mol-model-props/computed/interactions/hydrogen-bonds.ts

@@ -310,12 +310,12 @@ export const HydrogenAcceptorProvider = Features.Provider([FeatureType.HydrogenA
 export const HydrogenBondsProvider: LinkProvider<HydrogenBondsParams> = {
     name: 'hydrogen-bonds',
     params: HydrogenBondsParams,
-    requiredFeatures: [FeatureType.HydrogenDonor, FeatureType.HydrogenAcceptor],
     createTester: (props: HydrogenBondsProps) => {
         const maxDistance = Math.max(props.distanceMax, props.sulfurDistanceMax)
         const opts = getHydrogenBondsOptions(props)
         return {
-            maxDistanceSq: maxDistance * maxDistance,
+            maxDistance,
+            requiredFeatures: new Set([FeatureType.HydrogenDonor, FeatureType.HydrogenAcceptor]),
             getType: (structure, infoA, infoB, distanceSq) => testHydrogenBond(structure, infoA, infoB, distanceSq, opts)
         }
     }
@@ -324,11 +324,11 @@ export const HydrogenBondsProvider: LinkProvider<HydrogenBondsParams> = {
 export const WeakHydrogenBondsProvider: LinkProvider<WeakHydrogenBondsParams> = {
     name: 'weak-hydrogen-bonds',
     params: WeakHydrogenBondsParams,
-    requiredFeatures: [FeatureType.WeakHydrogenDonor, FeatureType.HydrogenAcceptor],
     createTester: (props: WeakHydrogenBondsProps) => {
         const opts = getGeometryOptions(props)
         return {
-            maxDistanceSq: props.distanceMax * props.distanceMax,
+            maxDistance: props.distanceMax,
+            requiredFeatures: new Set([FeatureType.WeakHydrogenDonor, FeatureType.HydrogenAcceptor]),
             getType: (structure, infoA, infoB, distanceSq) => testWeakHydrogenBond(structure, infoA, infoB, distanceSq, opts)
         }
     }

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

@@ -73,10 +73,10 @@ export const HydrophobicAtomProvider = Features.Provider([FeatureType.Hydrophobi
 export const HydrophobicProvider: LinkProvider<HydrophobicParams> = {
     name: 'hydrophobic',
     params: HydrophobicParams,
-    requiredFeatures: [FeatureType.HydrophobicAtom],
     createTester: (props: HydrophobicProps) => {
         return {
-            maxDistanceSq: props.distanceMax * props.distanceMax,
+            maxDistance: props.distanceMax,
+            requiredFeatures: new Set([FeatureType.HydrophobicAtom]),
             getType: (structure, infoA, infoB, distanceSq) => testHydrophobic(structure, infoA, infoB, distanceSq)
         }
     }

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

@@ -138,7 +138,7 @@ export const InteractionsParams = {
         'halogen-bonds',
         // 'hydrophobic',
         'metal-coordination',
-        'weak-hydrogen-bonds',
+        // 'weak-hydrogen-bonds',
     ], PD.objectToOptions(LinkProviders)),
     ...getProvidersParams()
 }
@@ -156,8 +156,8 @@ export async function computeInteractions(runtime: RuntimeContext, structure: St
     const linkTesters = linkProviders.map(l => l.createTester(p[l.name as keyof InteractionsProps]))
 
     const requiredFeatures = new Set<FeatureType>()
-    linkProviders.forEach(l => { for (const f of l.requiredFeatures) requiredFeatures.add(f) })
-    const featureProviders = FeatureProviders.filter(f => SetUtils.areIntersecting(requiredFeatures, new Set(f.types)))
+    linkTesters.forEach(l => SetUtils.add(requiredFeatures, l.requiredFeatures))
+    const featureProviders = FeatureProviders.filter(f => SetUtils.areIntersecting(requiredFeatures, f.types))
 
     const unitsFeatures = IntMap.Mutable<Features>()
     const unitsLinks = IntMap.Mutable<InteractionsIntraLinks>()
@@ -203,7 +203,7 @@ function findIntraUnitLinks(structure: Structure, unit: Unit, features: Features
 function findInterUnitLinks(structure: Structure, unitsFeatures: IntMap<Features>, linkTesters: ReadonlyArray<LinkTester>) {
     const builder = InterLinksBuilder.create()
 
-    const maxDistance = Math.sqrt(Math.max(...linkTesters.map(t => t.maxDistanceSq)))
+    const maxDistance = Math.max(...linkTesters.map(t => t.maxDistance))
 
     const lookup = structure.lookup3d;
     const imageCenter = Vec3.zero();

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

@@ -11,19 +11,20 @@ import { InteractionType, FeatureType } from './common';
 import { IntraLinksBuilder, InterLinksBuilder } from './links-builder';
 import { Mat4, Vec3 } from '../../../mol-math/linear-algebra';
 import { altLoc, connectedTo } from '../chemistry/util';
+import { OrderedSet } from '../../../mol-data/int';
 
 const MAX_DISTANCE = 5
 
 export interface LinkProvider<P extends PD.Params> {
     readonly name: string
     readonly params: P
-    readonly requiredFeatures: ReadonlyArray<FeatureType>
     createTester(props: PD.Values<P>): LinkTester
 }
 
 export interface LinkTester {
-    maxDistanceSq: number
-    getType: (structure: Structure, infoA: Features.Info, infoB: Features.Info, distanceSq: number, ) => InteractionType | undefined
+    readonly maxDistance: number
+    readonly requiredFeatures: ReadonlySet<FeatureType>
+    getType: (structure: Structure, infoA: Features.Info, infoB: Features.Info, distanceSq: number) => InteractionType | undefined
 }
 
 function validPair(structure: Structure, infoA: Features.Info, infoB: Features.Info): boolean {
@@ -44,36 +45,35 @@ function validPair(structure: Structure, infoA: Features.Info, infoB: Features.I
  * Add all intra-unit links, i.e. pairs of features
  */
 export function addUnitLinks(structure: Structure, unit: Unit.Atomic, features: Features, builder: IntraLinksBuilder, testers: ReadonlyArray<LinkTester>) {
-    const { x, y, z, count, lookup3d } = features
+
+    for (const tester of testers) {
+        _addUnitLinks(structure, unit, features, builder, tester)
+    }
+}
+
+function _addUnitLinks(structure: Structure, unit: Unit.Atomic, features: Features, builder: IntraLinksBuilder, tester: LinkTester) {
+    const { x, y, z } = features
+    const { lookup3d, indices: subsetIndices } = features.subset(tester.requiredFeatures)
 
     const infoA = Features.Info(structure, unit, features)
     const infoB = { ...infoA }
 
-    const maxDistanceSq = Math.max(...testers.map(t => t.maxDistanceSq))
-
-    for (let i = 0; i < count; ++i) {
-        const { count, indices, squaredDistances } = lookup3d.find(x[i], y[i], z[i], maxDistanceSq)
+    for (let t = 0, tl = OrderedSet.size(subsetIndices); t < tl; ++t) {
+        const i = OrderedSet.getAt(subsetIndices, t)
+        const { count, indices, squaredDistances } = lookup3d.find(x[i], y[i], z[i], tester.maxDistance)
         if (count === 0) continue
 
         infoA.feature = i
 
         for (let r = 0; r < count; ++r) {
-            const j = indices[r]
+            const j = OrderedSet.getAt(subsetIndices, indices[r])
             if (j <= i) continue
 
-            const distanceSq = squaredDistances[r]
             infoB.feature = j
             if (!validPair(structure, infoA, infoB)) continue
 
-            for (const tester of testers) {
-                if (distanceSq < tester.maxDistanceSq) {
-                    const type = tester.getType(structure, infoA, infoB, distanceSq)
-                    if (type) {
-                        builder.add(i, j, type)
-                        break
-                    }
-                }
-            }
+            const type = tester.getType(structure, infoA, infoB, squaredDistances[r])
+            if (type) builder.add(i, j, type)
         }
     }
 }
@@ -93,9 +93,9 @@ export function addStructureLinks(structure: Structure, unitA: Unit.Atomic, feat
     const isNotIdentity = !Mat4.isIdentity(imageTransform)
     const imageA = Vec3()
 
-    const maxDistanceSq = Math.max(...testers.map(t => t.maxDistanceSq))
+    const maxDistance = Math.max(...testers.map(t => t.maxDistance))
     const { center, radius } = lookup3d.boundary.sphere;
-    const testDistanceSq = (radius + maxDistanceSq) * (radius + maxDistanceSq);
+    const testDistanceSq = (radius + maxDistance) * (radius + maxDistance);
 
     const infoA = Features.Info(structure, unitA, featuresA)
     const infoB = Features.Info(structure, unitB, featuresB)
@@ -114,12 +114,12 @@ export function addStructureLinks(structure: Structure, unitA: Unit.Atomic, feat
 
         for (let r = 0; r < count; ++r) {
             const j = indices[r]
-            const distanceSq = squaredDistances[r]
             infoB.feature = j
             if (!validPair(structure, infoA, infoB)) continue
 
+            const distanceSq = squaredDistances[r]
             for (const tester of testers) {
-                if (distanceSq < tester.maxDistanceSq) {
+                if (distanceSq < tester.maxDistance * tester.maxDistance) {
                     const type = tester.getType(structure, infoA, infoB, distanceSq)
                     if (type) {
                         builder.add(i, j, type)

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

@@ -155,10 +155,10 @@ export const MetalBindingProvider = Features.Provider([FeatureType.IonicTypePart
 export const MetalCoordinationProvider: LinkProvider<MetalCoordinationParams> = {
     name: 'metal-coordination',
     params: MetalCoordinationParams,
-    requiredFeatures: [FeatureType.IonicTypeMetal, FeatureType.TransitionMetal, FeatureType.IonicTypePartner, FeatureType.DativeBondPartner],
     createTester: (props: MetalCoordinationProps) => {
         return {
-            maxDistanceSq: props.distanceMax * props.distanceMax,
+            maxDistance: props.distanceMax,
+            requiredFeatures: new Set([FeatureType.IonicTypeMetal, FeatureType.TransitionMetal, FeatureType.IonicTypePartner, FeatureType.DativeBondPartner]),
             getType: (structure, infoA, infoB, distanceSq) => testMetalCoordination(structure, infoA, infoB, distanceSq)
         }
     }