Browse Source

wip, refactored interaction computation

Alexander Rose 5 years ago
parent
commit
af700c1481

+ 28 - 1
src/mol-model-props/computed/interactions/features.ts

@@ -4,11 +4,12 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { StructureElement } from '../../../mol-model/structure/structure';
+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 { FeatureGroup, FeatureType } from './common';
+import { ValenceModelProvider } from '../valence-model';
 
 export { Features }
 
@@ -97,6 +98,32 @@ namespace Features {
             },
         }
     }
+
+    export interface Info {
+        unit: Unit.Atomic,
+        types: ArrayLike<FeatureType>,
+        feature: number,
+        members: ArrayLike<StructureElement.UnitIndex>,
+        offsets: ArrayLike<number>,
+        idealGeometry: Int8Array
+    }
+    export function Info(structure: Structure, unit: Unit.Atomic, features: Features) {
+        const valenceModel = ValenceModelProvider.getValue(structure).value
+        if (!valenceModel || !valenceModel.has(unit.id)) throw new Error('valence model required')
+
+        return {
+            unit,
+            types: features.types,
+            members: features.members,
+            offsets: features.offsets,
+            idealGeometry: valenceModel.get(unit.id)!.idealGeometry
+        } as Info
+    }
+
+    export interface Provider {
+        name: string
+        add: (structure: Structure, unit: Unit.Atomic, featuresBuilder: FeaturesBuilder) => void
+    }
 }
 
 export { FeaturesBuilder }

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

@@ -16,8 +16,7 @@ import { typeSymbol, altLoc, eachBondedAtom } from '../chemistry/util';
 import { Elements } from '../../../mol-model/structure/model/properties/atomic/types';
 import { degToRad } from '../../../mol-math/misc';
 import { FeatureType, FeatureGroup, InteractionType } from './common';
-import { IntraLinksBuilder, InterLinksBuilder } from './builder';
-import { Mat4, Vec3 } from '../../../mol-math/linear-algebra';
+import { LinkProvider } from './links';
 
 export const HalogenBondsParams = {
     distanceMax: PD.Numeric(4.0, { min: 1, max: 5, step: 0.1 }),
@@ -136,77 +135,14 @@ function testHalogenBond(structure: Structure, infoA: Info, infoB: Info, opts: O
     return InteractionType.HalogenBond
 }
 
-/**
- * All intra-unit pairs of halogen donor and acceptor atoms
- */
-export function addUnitHalogenBonds(structure: Structure, unit: Unit.Atomic, features: Features, builder: IntraLinksBuilder, props: HalogenBondsProps) {
-    const opts = getOptions(props)
-    const { x, y, z, count, lookup3d } = features
-
-    const infoA = Info(structure, unit, features)
-    const infoB = { ...infoA }
-
-    for (let i = 0; i < count; ++i) {
-        const { count, indices } = lookup3d.find(x[i], y[i], z[i], opts.distanceMax)
-        if (count === 0) continue
-
-        infoA.feature = i
-
-        for (let r = 0; r < count; ++r) {
-            const j = indices[r]
-            if (j <= i) continue
-
-            infoB.feature = j
-            const bondType = testHalogenBond(structure, infoA, infoB, opts)
-            if (bondType) builder.add(i, j, bondType)
+export const HalogenBondsProvider: LinkProvider<HalogenBondsParams> = {
+    name: 'halogen-bonds',
+    params: HalogenBondsParams,
+    createTester: (props: HalogenBondsProps) => {
+        const opts = getOptions(props)
+        return {
+            maxDistanceSq: props.distanceMax * props.distanceMax,
+            getType: (structure, infoA, infoB) => testHalogenBond(structure, infoA, infoB, opts)
         }
     }
-}
-
-const _imageTransform = Mat4()
-
-//
-
-/**
- * All inter-unit pairs of halogen donor and acceptor atoms
- */
-export function addStructureHalogenBonds(structure: Structure, unitA: Unit.Atomic, featuresA: Features, unitB: Unit.Atomic, featuresB: Features, builder: InterLinksBuilder, props: HalogenBondsProps) {
-    const opts = getOptions(props)
-
-    const { count, x: xA, y: yA, z: zA } = featuresA;
-    const { lookup3d } = featuresB;
-
-    // the lookup queries need to happen in the "unitB space".
-    // that means imageA = inverseOperB(operA(i))
-    const imageTransform = Mat4.mul(_imageTransform, unitB.conformation.operator.inverse, unitA.conformation.operator.matrix)
-    const isNotIdentity = !Mat4.isIdentity(imageTransform)
-    const imageA = Vec3()
-
-    const { center: bCenter, radius: bRadius } = lookup3d.boundary.sphere;
-    const testDistanceSq = (bRadius + opts.distanceMax) * (bRadius + opts.distanceMax);
-
-    const infoA = Info(structure, unitA, featuresA)
-    const infoB = Info(structure, unitB, featuresB)
-
-    builder.startUnitPair(unitA, unitB)
-
-    for (let i = 0; i < count; ++i) {
-        Vec3.set(imageA, xA[i], yA[i], zA[i])
-        if (isNotIdentity) Vec3.transformMat4(imageA, imageA, imageTransform)
-        if (Vec3.squaredDistance(imageA, bCenter) > testDistanceSq) continue
-
-        const { indices, count } = lookup3d.find(imageA[0], imageA[1], imageA[2], opts.distanceMax)
-        if (count === 0) continue
-
-        infoA.feature = i
-
-        for (let r = 0; r < count; ++r) {
-            const j = indices[r]
-            infoB.feature = j
-            const bondType = testHalogenBond(structure, infoA, infoB, opts)
-            if (bondType) builder.add(i, j, bondType)
-        }
-    }
-
-    builder.finishUnitPair()
 }

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

@@ -17,8 +17,7 @@ import { Elements } from '../../../mol-model/structure/model/properties/atomic/t
 import { ValenceModelProvider } from '../valence-model';
 import { degToRad } from '../../../mol-math/misc';
 import { FeatureType, FeatureGroup, InteractionType } from './common';
-import { IntraLinksBuilder, InterLinksBuilder } from './builder';
-import { Mat4, Vec3 } from '../../../mol-math/linear-algebra';
+import { LinkProvider } from './links';
 
 export const HydrogenBondsParams = {
     distanceMax: PD.Numeric(3.5, { min: 1, max: 5, step: 0.1 }),
@@ -211,27 +210,6 @@ function getHydrogenBondType(unitA: Unit.Atomic, indexA: StructureElement.UnitIn
     }
 }
 
-interface Info {
-    unit: Unit.Atomic,
-    types: ArrayLike<FeatureType>,
-    feature: number,
-    members: ArrayLike<StructureElement.UnitIndex>,
-    offsets: ArrayLike<number>,
-    idealGeometry: Int8Array
-}
-function Info(structure: Structure, unit: Unit.Atomic, features: Features) {
-    const valenceModel = ValenceModelProvider.getValue(structure).value
-    if (!valenceModel || !valenceModel.has(unit.id)) throw new Error('valence model required')
-
-    return {
-        unit,
-        types: features.types,
-        members: features.members,
-        offsets: features.offsets,
-        idealGeometry: valenceModel.get(unit.id)!.idealGeometry
-    } as Info
-}
-
 function getOptions(props: HydrogenBondsProps) {
     return {
         maxAccAngleDev: degToRad(props.accAngleDevMax),
@@ -246,7 +224,7 @@ type Options = ReturnType<typeof getOptions>
 
 const deg120InRad = degToRad(120)
 
-function testHydrogenBond(dSq: number, structure: Structure, infoA: Info, infoB: Info, opts: Options): InteractionType | undefined {
+function testHydrogenBond(structure: Structure, infoA: Features.Info, infoB: Features.Info, distanceSq: number, opts: Options): InteractionType | undefined {
     const typeA = infoA.types[infoA.feature]
     const typeB = infoB.types[infoB.feature]
 
@@ -267,7 +245,7 @@ function testHydrogenBond(dSq: number, structure: Structure, infoA: Info, infoB:
     if (don.unit.residueIndex[don.unit.elements[donIndex]] === acc.unit.residueIndex[acc.unit.elements[accIndex]]) return // same residue
 
     // check if distance is ok for non-sulfur-containing hbond
-    if (typeSymbol(don.unit, donIndex) !== Elements.S && typeSymbol(acc.unit, accIndex) !== Elements.S && dSq > opts.maxHbondDistSq) return
+    if (typeSymbol(don.unit, donIndex) !== Elements.S && typeSymbol(acc.unit, accIndex) !== Elements.S && distanceSq > opts.maxHbondDistSq) return
 
     // no hbond if donor and acceptor are bonded
     if (connectedTo(structure, don.unit, donIndex, acc.unit, accIndex)) return
@@ -295,77 +273,15 @@ function testHydrogenBond(dSq: number, structure: Structure, infoA: Info, infoB:
     return isWeak ? InteractionType.WeakHydrogenBond : getHydrogenBondType(don.unit, donIndex, acc.unit, accIndex)
 }
 
-/**
- * All intra-unit pairs of hydrogen donor and acceptor atoms
- */
-export function addUnitHydrogenBonds(structure: Structure, unit: Unit.Atomic, features: Features, builder: IntraLinksBuilder, props: HydrogenBondsProps) {
-    const opts = getOptions(props)
-    const { x, y, z, count, lookup3d } = features
-
-    const infoA = Info(structure, unit, features)
-    const infoB = { ...infoA }
-
-    for (let i = 0; i < count; ++i) {
-        const { count, indices, squaredDistances } = lookup3d.find(x[i], y[i], z[i], opts.maxDist)
-        if (count === 0) continue
-
-        infoA.feature = i
-
-        for (let r = 0; r < count; ++r) {
-            const j = indices[r]
-            if (j <= i) continue
-
-            infoB.feature = j
-            const bondType = testHydrogenBond(squaredDistances[r], structure, infoA, infoB, opts)
-            if (bondType) builder.add(i, j, bondType)
-        }
-    }
-}
-
-//
-
-const _imageTransform = Mat4()
-
-/**
- * All inter-unit pairs of hydrogen donor and acceptor atoms
- */
-export function addStructureHydrogenBonds(structure: Structure, unitA: Unit.Atomic, featuresA: Features, unitB: Unit.Atomic, featuresB: Features, builder: InterLinksBuilder, props: HydrogenBondsProps) {
-    const opts = getOptions(props)
-
-    const { count, x: xA, y: yA, z: zA } = featuresA;
-    const { lookup3d } = featuresB;
-
-    // the lookup queries need to happen in the "unitB space".
-    // that means imageA = inverseOperB(operA(i))
-    const imageTransform = Mat4.mul(_imageTransform, unitB.conformation.operator.inverse, unitA.conformation.operator.matrix)
-    const isNotIdentity = !Mat4.isIdentity(imageTransform)
-    const imageA = Vec3()
-
-    const { center: bCenter, radius: bRadius } = lookup3d.boundary.sphere;
-    const testDistanceSq = (bRadius + opts.maxDist) * (bRadius + opts.maxDist);
-
-    const infoA = Info(structure, unitA, featuresA)
-    const infoB = Info(structure, unitB, featuresB)
-
-    builder.startUnitPair(unitA, unitB)
-
-    for (let i = 0; i < count; ++i) {
-        Vec3.set(imageA, xA[i], yA[i], zA[i])
-        if (isNotIdentity) Vec3.transformMat4(imageA, imageA, imageTransform)
-        if (Vec3.squaredDistance(imageA, bCenter) > testDistanceSq) continue
-
-        const { indices, count, squaredDistances } = lookup3d.find(imageA[0], imageA[1], imageA[2], opts.maxDist)
-        if (count === 0) continue
-
-        infoA.feature = i
-
-        for (let r = 0; r < count; ++r) {
-            const j = indices[r]
-            infoB.feature = j
-            const bondType = testHydrogenBond(squaredDistances[r], structure, infoA, infoB, opts)
-            if (bondType) builder.add(i, j, bondType)
+export const HydrogenBondsProvider: LinkProvider<HydrogenBondsParams> = {
+    name: 'hydrogen-bonds',
+    params: HydrogenBondsParams,
+    createTester: (props: HydrogenBondsProps) => {
+        const maxDistance = Math.max(props.distanceMax, props.sulfurDistanceMax)
+        const opts = getOptions(props)
+        return {
+            maxDistanceSq: maxDistance * maxDistance,
+            getType: (structure, infoA, infoB, distanceSq) => testHydrogenBond(structure, infoA, infoB, distanceSq, opts)
         }
     }
-
-    builder.finishUnitPair()
 }

+ 32 - 23
src/mol-model-props/computed/interactions/interactions.ts

@@ -7,14 +7,15 @@
 import { ParamDefinition as PD } from '../../../mol-util/param-definition';
 import { Structure, Unit } from '../../../mol-model/structure';
 import { RuntimeContext } from '../../../mol-task';
-import { addUnitHydrogenDonors, addUnitWeakHydrogenDonors, addUnitHydrogenAcceptors, addUnitHydrogenBonds, HydrogenBondsParams, addStructureHydrogenBonds } from './hydrogen-bonds';
 import { Features, FeaturesBuilder } from './features';
 import { ValenceModelProvider } from '../valence-model';
 import { InteractionsIntraLinks, InteractionsInterLinks } from './common';
 import { IntraLinksBuilder, InterLinksBuilder } from './builder';
 import { IntMap } from '../../../mol-data/int';
 import { Vec3 } from '../../../mol-math/linear-algebra';
-import { addUnitHalogenBonds, HalogenBondsParams, addUnitHalogenDonors, addUnitHalogenAcceptors, addStructureHalogenBonds } from './halogen-bonds';
+import { addUnitLinks, LinkTester, addStructureLinks } from './links';
+import { addUnitHalogenDonors, addUnitHalogenAcceptors, HalogenBondsProvider } from './halogen-bonds';
+import { addUnitHydrogenDonors, addUnitWeakHydrogenDonors, addUnitHydrogenAcceptors, HydrogenBondsProvider } from './hydrogen-bonds';
 
 export { Interactions }
 
@@ -98,8 +99,8 @@ namespace Interactions {
 }
 
 export const InteractionsParams = {
-    hydrogenBonds: PD.Group(HydrogenBondsParams),
-    halogenBonds: PD.Group(HalogenBondsParams),
+    hydrogenBonds: PD.Group(HydrogenBondsProvider.params),
+    halogenBonds: PD.Group(HalogenBondsProvider.params),
 }
 export type InteractionsParams = typeof InteractionsParams
 export type InteractionsProps = PD.Values<InteractionsParams>
@@ -108,12 +109,17 @@ export async function computeInteractions(runtime: RuntimeContext, structure: St
     const p = { ...PD.getDefaultValues(InteractionsParams), ...props }
     await ValenceModelProvider.attach(structure).runInContext(runtime)
 
+    const linkTesters: LinkTester[] = [
+        HydrogenBondsProvider.createTester(p.hydrogenBonds),
+        HalogenBondsProvider.createTester(p.halogenBonds)
+    ]
+
     const unitsFeatures = IntMap.Mutable<Features>()
     const unitsLinks = IntMap.Mutable<InteractionsIntraLinks>()
 
     for (let i = 0, il = structure.unitSymmetryGroups.length; i < il; ++i) {
         const group = structure.unitSymmetryGroups[i]
-        const d = findIntraUnitLinksAndFeatures(structure, group.units[0], p)
+        const d = findIntraUnitLinksAndFeatures(structure, group.units[0], linkTesters)
         for (let j = 0, jl = group.units.length; j < jl; ++j) {
             const u = group.units[j]
             unitsFeatures.set(u.id, d.features)
@@ -121,38 +127,43 @@ export async function computeInteractions(runtime: RuntimeContext, structure: St
         }
     }
 
-    const links = findInterUnitLinks(structure, unitsFeatures, p)
+    const links = findInterUnitLinks(structure, unitsFeatures, linkTesters)
 
     return { unitsFeatures, unitsLinks, links }
 }
 
-function findIntraUnitLinksAndFeatures(structure: Structure, unit: Unit, props: InteractionsProps) {
+const FeatureProviders: Features.Provider[] = [
+    { name: 'hydrogen-donors', add: addUnitHydrogenDonors },
+    { name: 'weak-hydrogen-donors', add: addUnitWeakHydrogenDonors },
+    { name: 'hydrogen-acceptors', add: addUnitHydrogenAcceptors },
+
+    { name: 'halogen-donors', add: addUnitHalogenDonors },
+    { name: 'halogen-acceptors', add: addUnitHalogenAcceptors },
+]
+
+function findIntraUnitLinksAndFeatures(structure: Structure, unit: Unit, linkTesters: ReadonlyArray<LinkTester>) {
     const count = unit.elements.length
     const featuresBuilder = FeaturesBuilder.create(count, count / 2)
     if (Unit.isAtomic(unit)) {
-        addUnitHydrogenDonors(structure, unit, featuresBuilder)
-        addUnitWeakHydrogenDonors(structure, unit, featuresBuilder)
-        addUnitHydrogenAcceptors(structure, unit, featuresBuilder)
-
-        addUnitHalogenDonors(structure, unit, featuresBuilder)
-        addUnitHalogenAcceptors(structure, unit, featuresBuilder)
+        for (const featureProvider of FeatureProviders) {
+            featureProvider.add(structure, unit, featuresBuilder)
+        }
     }
     const features = featuresBuilder.getFeatures(count)
 
     const linksBuilder = IntraLinksBuilder.create(features, count)
     if (Unit.isAtomic(unit)) {
-        addUnitHydrogenBonds(structure, unit, features, linksBuilder, props.hydrogenBonds)
-        addUnitHalogenBonds(structure, unit, features, linksBuilder, props.halogenBonds)
+        addUnitLinks(structure, unit, features, linksBuilder, linkTesters)
     }
 
     return { features, links: linksBuilder.getLinks() }
 }
 
-const MAX_RADIUS = 5
-
-function findInterUnitLinks(structure: Structure, unitsFeatures: IntMap<Features>, props: InteractionsProps) {
+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 lookup = structure.lookup3d;
     const imageCenter = Vec3.zero();
 
@@ -163,7 +174,7 @@ function findInterUnitLinks(structure: Structure, unitsFeatures: IntMap<Features
 
         const bs = unitA.lookup3d.boundary.sphere;
         Vec3.transformMat4(imageCenter, bs.center, unitA.conformation.operator.matrix);
-        const closeUnits = lookup.findUnitIndices(imageCenter[0], imageCenter[1], imageCenter[2], bs.radius + MAX_RADIUS);
+        const closeUnits = lookup.findUnitIndices(imageCenter[0], imageCenter[1], imageCenter[2], bs.radius + maxDistance);
 
         for (let i = 0; i < closeUnits.count; i++) {
             const unitB = structure.units[closeUnits.indices[i]];
@@ -172,11 +183,9 @@ function findInterUnitLinks(structure: Structure, unitsFeatures: IntMap<Features
             const featuresB = unitsFeatures.get(unitB.id)
 
             if (unitB.elements.length >= unitA.elements.length) {
-                addStructureHydrogenBonds(structure, unitA, featuresA, unitB, featuresB, builder, props.hydrogenBonds)
-                addStructureHalogenBonds(structure, unitA, featuresA, unitB, featuresB, builder, props.halogenBonds)
+                addStructureLinks(structure, unitA, featuresA, unitB, featuresB, builder, linkTesters)
             } else {
-                addStructureHydrogenBonds(structure, unitB, featuresB, unitA, featuresA, builder, props.hydrogenBonds)
-                addStructureHalogenBonds(structure, unitB, featuresB, unitA, featuresA, builder, props.halogenBonds)
+                addStructureLinks(structure, unitB, featuresB, unitA, featuresA, builder, linkTesters)
             }
         }
     }

+ 108 - 0
src/mol-model-props/computed/interactions/links.ts

@@ -0,0 +1,108 @@
+/**
+ * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { ParamDefinition as PD } from '../../../mol-util/param-definition';
+import { Structure, Unit } from '../../../mol-model/structure';
+import { Features } from './features';
+import { InteractionType } from './common';
+import { IntraLinksBuilder, InterLinksBuilder } from './builder';
+import { Mat4, Vec3 } from '../../../mol-math/linear-algebra';
+
+const MAX_DISTANCE = 5
+
+export interface LinkProvider<P extends PD.Params> {
+    name: string
+    params: P
+    createTester(props: PD.Values<P>): LinkTester
+}
+
+export interface LinkTester {
+    maxDistanceSq: number
+    getType: (structure: Structure, infoA: Features.Info, infoB: Features.Info, distanceSq: number, ) => InteractionType | undefined
+}
+
+/**
+ * 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
+
+    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)
+        if (count === 0) continue
+
+        infoA.feature = i
+
+        for (let r = 0; r < count; ++r) {
+            const j = indices[r]
+            if (j <= i) continue
+
+            const distanceSq = squaredDistances[r]
+            infoB.feature = j
+            for (const tester of testers) {
+                if (distanceSq < tester.maxDistanceSq) {
+                    const type = tester.getType(structure, infoA, infoB, distanceSq)
+                    if (type) builder.add(i, j, type)
+                }
+            }
+        }
+    }
+}
+
+const _imageTransform = Mat4()
+
+/**
+ * Add all inter-unit links, i.e. pairs of features
+ */
+export function addStructureLinks(structure: Structure, unitA: Unit.Atomic, featuresA: Features, unitB: Unit.Atomic, featuresB: Features, builder: InterLinksBuilder, testers: ReadonlyArray<LinkTester>) {
+    const { count, x: xA, y: yA, z: zA } = featuresA;
+    const { lookup3d } = featuresB;
+
+    // the lookup queries need to happen in the "unitB space".
+    // that means imageA = inverseOperB(operA(i))
+    const imageTransform = Mat4.mul(_imageTransform, unitB.conformation.operator.inverse, unitA.conformation.operator.matrix)
+    const isNotIdentity = !Mat4.isIdentity(imageTransform)
+    const imageA = Vec3()
+
+    const maxDistanceSq = Math.max(...testers.map(t => t.maxDistanceSq))
+    const { center, radius } = lookup3d.boundary.sphere;
+    const testDistanceSq = (radius + maxDistanceSq) * (radius + maxDistanceSq);
+
+    const infoA = Features.Info(structure, unitA, featuresA)
+    const infoB = Features.Info(structure, unitB, featuresB)
+
+    builder.startUnitPair(unitA, unitB)
+
+    for (let i = 0; i < count; ++i) {
+        Vec3.set(imageA, xA[i], yA[i], zA[i])
+        if (isNotIdentity) Vec3.transformMat4(imageA, imageA, imageTransform)
+        if (Vec3.squaredDistance(imageA, center) > testDistanceSq) continue
+
+        const { indices, count, squaredDistances } = lookup3d.find(imageA[0], imageA[1], imageA[2], MAX_DISTANCE)
+        if (count === 0) continue
+
+        infoA.feature = i
+
+        for (let r = 0; r < count; ++r) {
+            const j = indices[r]
+            const distanceSq = squaredDistances[r]
+            infoB.feature = j
+            for (const tester of testers) {
+                if (distanceSq < tester.maxDistanceSq) {
+                    const type = tester.getType(structure, infoA, infoB, distanceSq)
+                    if (type) builder.add(i, j, type)
+                }
+            }
+        }
+    }
+
+    builder.finishUnitPair()
+}