wip, interactions refiners

- saltBridgeRefiner
- piStackingRefiner
- metalCoordinationRefiner

src/mol-math/graph/inter-unit-graph.ts

@@ -44,11 +44,11 @@ class InterUnitGraph<Unit extends InterUnitGraph.UnitBase, VertexIndex extends n
         return this.vertexKeyIndex.get(InterUnitGraph.getVertexKey(index, unit)) || []
     }
 
-    constructor(private map: Map<number, InterUnitGraph.UnitPairEdges<Unit, VertexIndex, EdgeProps>[]>) {
+    constructor(protected readonly map: Map<number, InterUnitGraph.UnitPairEdges<Unit, VertexIndex, EdgeProps>[]>) {
         let count = 0
         const edges: (InterUnitGraph.Edge<Unit, VertexIndex, EdgeProps>)[] = []
         const edgeKeyIndex = new Map<string, number>()
-        const elementKeyIndex = new Map<string, number[]>()
+        const vertexKeyIndex = new Map<string, number[]>()
 
         this.map.forEach(pairEdgesArray => {
             pairEdgesArray.forEach(pairEdges => {
@@ -57,12 +57,12 @@ class InterUnitGraph<Unit extends InterUnitGraph.UnitBase, VertexIndex extends n
                     pairEdges.getEdges(indexA).forEach(edgeInfo => {
                         const { unitA, unitB } = pairEdges
 
-                        const edgeKey = InterUnitGraph.getEdgeKey<Unit, VertexIndex>(indexA, unitA, edgeInfo.indexB, unitB)
+                        const edgeKey = InterUnitGraph.getEdgeKey(indexA, unitA, edgeInfo.indexB, unitB)
                         edgeKeyIndex.set(edgeKey, edges.length)
 
-                        const elementKey = InterUnitGraph.getVertexKey(indexA, unitA)
-                        const e = elementKeyIndex.get(elementKey)
-                        if (e === undefined) elementKeyIndex.set(elementKey, [edges.length])
+                        const vertexKey = InterUnitGraph.getVertexKey(indexA, unitA)
+                        const e = vertexKeyIndex.get(vertexKey)
+                        if (e === undefined) vertexKeyIndex.set(vertexKey, [edges.length])
                         else e.push(edges.length)
 
                         edges.push({ ...edgeInfo, indexA, unitA, unitB })
@@ -74,7 +74,7 @@ class InterUnitGraph<Unit extends InterUnitGraph.UnitBase, VertexIndex extends n
         this.edgeCount = count
         this.edges = edges
         this.edgeKeyIndex = edgeKeyIndex
-        this.vertexKeyIndex = elementKeyIndex
+        this.vertexKeyIndex = vertexKeyIndex
     }
 }
 

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

@@ -9,16 +9,131 @@ import { InterUnitGraph } from '../../../mol-math/graph/inter-unit-graph'
 import { Unit } from '../../../mol-model/structure'
 import { AssignableArrayLike } from '../../../mol-util/type-helpers'
 import { Features } from './features'
+import { StructureElement } from '../../../mol-model/structure/structure'
+import { IntMap } from '../../../mol-data/int'
 
 type IntraProps = {
     readonly type: ArrayLike<InteractionType>
     readonly flag: AssignableArrayLike<InteractionFlag>
 }
-export type InteractionsIntraContacts = IntAdjacencyGraph<Features.FeatureIndex, IntraProps>
+export { InteractionsIntraContacts }
+interface InteractionsIntraContacts extends IntAdjacencyGraph<Features.FeatureIndex, IntraProps> {
+    readonly elementsIndex: InteractionsIntraContacts.ElementsIndex
+}
+namespace InteractionsIntraContacts {
+    /** maps unit elements to contacts, range for unit element i is offsets[i] to offsets[i + 1] */
+    export type ElementsIndex = {
+        /** intra contact indices */
+        readonly indices: ArrayLike<number>
+        /** range for unit element i is offsets[i] to offsets[i + 1] */
+        readonly offsets: ArrayLike<number>
+    }
+
+    /**
+     * Note: assumes that feature members of a contact are non-overlapping
+     */
+    export function createElementsIndex(contacts: IntAdjacencyGraph<Features.FeatureIndex, IntraProps>, features: Features, elementsCount: number) {
+        const offsets = new Int32Array(elementsCount + 1)
+        const bucketFill = new Int32Array(elementsCount)
+        const bucketSizes = new Int32Array(elementsCount)
+        const { members, offsets: featureOffsets } = features
+
+        for (let i = 0, il = contacts.edgeCount * 2; i < il; ++i) {
+            const aI = contacts.a[i]
+            const bI = contacts.b[i]
+            if (aI > bI) continue
+
+            for (let j = featureOffsets[aI], jl = featureOffsets[aI + 1]; j < jl; ++j) {
+                ++bucketSizes[members[j]]
+            }
+            for (let j = featureOffsets[bI], jl = featureOffsets[bI + 1]; j < jl; ++j) {
+                ++bucketSizes[members[j]]
+            }
+        }
+
+        let offset = 0
+        for (let i = 0; i < elementsCount; i++) {
+            offsets[i] = offset
+            offset += bucketSizes[i]
+        }
+        offsets[elementsCount] = offset
+
+        const indices = new Int32Array(offset)
+        for (let i = 0, il = contacts.edgeCount * 2; i < il; ++i) {
+            const aI = contacts.a[i]
+            const bI = contacts.b[i]
+            if (aI > bI) continue
+
+            for (let j = featureOffsets[aI], jl = featureOffsets[aI + 1]; j < jl; ++j) {
+                const m = members[j]
+                const om = offsets[m] + bucketFill[m]
+                indices[om] = i
+                ++bucketFill[m]
+            }
+            for (let j = featureOffsets[bI], jl = featureOffsets[bI + 1]; j < jl; ++j) {
+                const m = members[j]
+                const om = offsets[m] + bucketFill[m]
+                indices[om] = i
+                ++bucketFill[m]
+            }
+        }
+
+        return { indices, offsets }
+    }
+}
 
 export { InteractionsInterContacts }
 type InterProps = { type: InteractionType, flag: InteractionFlag }
-type InteractionsInterContacts = InterUnitGraph<Unit, Features.FeatureIndex, InterProps>
+class InteractionsInterContacts extends InterUnitGraph<Unit, Features.FeatureIndex, InterProps> {
+    private readonly elementKeyIndex: Map<string, number[]>
+
+    getContactIndicesForElement(index: StructureElement.UnitIndex, unit: Unit): ReadonlyArray<number> {
+        return this.elementKeyIndex.get(this.getElementKey(index, unit)) || []
+    }
+
+    private getElementKey(index: StructureElement.UnitIndex, unit: Unit): string {
+        return `${index}|${unit.id}`
+    }
+
+    constructor(map: Map<number, InterUnitGraph.UnitPairEdges<Unit, Features.FeatureIndex, InterProps>[]>, unitsFeatures: IntMap<Features>) {
+        super(map)
+
+        let count = 0
+        const elementKeyIndex = new Map<string, number[]>()
+
+        const add = (index: StructureElement.UnitIndex, unit: Unit) => {
+            const vertexKey = this.getElementKey(index, unit)
+            const e = elementKeyIndex.get(vertexKey)
+            if (e === undefined) elementKeyIndex.set(vertexKey, [count])
+            else e.push(count)
+        }
+
+        this.map.forEach(pairEdgesArray => {
+            pairEdgesArray.forEach(pairEdges => {
+                pairEdges.connectedIndices.forEach(indexA => {
+                    pairEdges.getEdges(indexA).forEach(edgeInfo => {
+                        const { unitA, unitB } = pairEdges
+
+                        const { offsets: offsetsA, members: membersA } = unitsFeatures.get(unitA.id)
+                        for (let j = offsetsA[indexA], jl = offsetsA[indexA + 1]; j < jl; ++j) {
+                            add(membersA[j], unitA)
+                        }
+
+                        const { indexB } = edgeInfo
+                        const { offsets: offsetsB, members: membersB } = unitsFeatures.get(unitB.id)
+                        for (let j = offsetsB[indexB], jl = offsetsB[indexB + 1]; j < jl; ++j) {
+                            add(membersB[j], unitB)
+                        }
+
+                        count += 1
+                    })
+                })
+            })
+        })
+
+        this.elementKeyIndex = elementKeyIndex
+    }
+}
 namespace InteractionsInterContacts {
     export class Pair extends InterUnitGraph.UnitPairEdges<Unit, Features.FeatureIndex, InterProps> {}
     export type Info = InterUnitGraph.EdgeInfo<Features.FeatureIndex, InterProps>

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

@@ -8,9 +8,9 @@ import { Features } from './features';
 import { IntAdjacencyGraph } from '../../../mol-math/graph';
 import { InteractionType, InteractionsIntraContacts, InteractionsInterContacts, InteractionFlag } from './common';
 import { Unit } from '../../../mol-model/structure/structure';
-import { InterUnitGraph } from '../../../mol-math/graph/inter-unit-graph';
 import { UniqueArray } from '../../../mol-data/generic';
 import { NumberArray } from '../../../mol-util/type-helpers';
+import { IntMap } from '../../../mol-data/int';
 
 export { IntraContactsBuilder }
 
@@ -39,7 +39,15 @@ namespace IntraContactsBuilder {
                     builder.addNextEdge()
                     builder.assignProperty(type, types[i])
                 }
-                return builder.createGraph({ type, flag })
+                const graph = builder.createGraph({ type, flag })
+
+                let elementsIndex: InteractionsIntraContacts.ElementsIndex
+                const contacts: InteractionsIntraContacts = Object.defineProperty(graph, 'elementsIndex', {
+                    get: () => {
+                        return elementsIndex || (elementsIndex = InteractionsIntraContacts.createElementsIndex(graph, features, elementsCount))
+                    }
+                })
+                return contacts
             }
         }
     }
@@ -51,7 +59,7 @@ interface InterContactsBuilder {
     startUnitPair: (unitA: Unit, unitB: Unit) => void
     finishUnitPair: () => void
     add: (indexA: Features.FeatureIndex, indexB: Features.FeatureIndex, type: InteractionType) => void
-    getContacts: () => InteractionsInterContacts
+    getContacts: (unitsFeatures: IntMap<Features>) => InteractionsInterContacts
 }
 
 namespace InterContactsBuilder {
@@ -93,8 +101,8 @@ namespace InterContactsBuilder {
                 UniqueArray.add(linkedB, indexB, indexB)
                 linkCount += 1
             },
-            getContacts() {
-                return new InterUnitGraph(map);
+            getContacts(unitsFeatures: IntMap<Features>) {
+                return new InteractionsInterContacts(map, unitsFeatures);
             }
         }
     }

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

@@ -235,5 +235,5 @@ function findInterUnitContacts(structure: Structure, unitsFeatures: IntMap<Featu
         }
     }
 
-    return builder.getContacts()
+    return builder.getContacts(unitsFeatures)
 }

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

@@ -7,7 +7,7 @@
  */
 
 import { Interactions } from './interactions';
-import { InteractionType, InteractionFlag, InteractionsIntraContacts, FeatureType } from './common';
+import { InteractionType, InteractionFlag, InteractionsIntraContacts, FeatureType, InteractionsInterContacts } from './common';
 import { Vec3 } from '../../../mol-math/linear-algebra';
 import { Unit, Structure } from '../../../mol-model/structure';
 import { Features } from './features';
@@ -15,7 +15,7 @@ import { Features } from './features';
 interface ContactRefiner {
     isApplicable: (type: InteractionType) => boolean
     handleInterContact: (index: number, infoA: Features.Info, infoB: Features.Info) => void
-    startUnit: (contacts: InteractionsIntraContacts) => void
+    startUnit: (unit: Unit.Atomic, contacts: InteractionsIntraContacts, features: Features) => void
     handleIntraContact: (index: number, infoA: Features.Info, infoB: Features.Info) => void
 }
 
@@ -25,6 +25,9 @@ export function refineInteractions(structure: Structure, interactions: Interacti
     const contactRefiners: ContactRefiner[] = [
         hydrophobicRefiner(structure, interactions),
         weakHydrogenBondsRefiner(structure, interactions),
+        saltBridgeRefiner(structure, interactions),
+        piStackingRefiner(structure, interactions),
+        metalCoordinationRefiner(structure, interactions),
     ]
 
     for (let i = 0, il = contacts.edgeCount; i < il; ++i) {
@@ -56,11 +59,12 @@ export function refineInteractions(structure: Structure, interactions: Interacti
         const infoA = Features.Info(structure, unit, features)
         const infoB = Features.Info(structure, unit, features)
 
-        for (const refiner of contactRefiners) refiner.startUnit(contacts)
+        for (const refiner of contactRefiners) refiner.startUnit(unit, contacts, features)
 
         for (let i = 0, il = contacts.edgeCount * 2; i < il; ++i) {
             infoA.feature = contacts.a[i]
             infoB.feature = contacts.b[i]
+            // console.log(i, contacts.a[i], contacts.b[i])
 
             for (const refiner of contactRefiners) {
                 if (refiner.isApplicable(contacts.edgeProps.type[i])) refiner.handleIntraContact(i, infoA, infoB)
@@ -118,7 +122,7 @@ function hydrophobicRefiner(structure: Structure, interactions: Interactions): C
         handleInterContact: (index: number, infoA: Features.Info, infoB: Features.Info) => {
             handleEdge(index, infoA, infoB, residueInterMap, setInterFiltered)
         },
-        startUnit: (contacts: InteractionsIntraContacts) => {
+        startUnit: (unit: Unit.Atomic, contacts: InteractionsIntraContacts, features: Features) => {
             residueIntraMap = new Map<string, [number, number]>()
             setIntraFiltered = (i: number) => contacts.edgeProps.flag[i] = InteractionFlag.Filtered
         },
@@ -162,7 +166,7 @@ function weakHydrogenBondsRefiner(structure: Structure, interactions: Interactio
                 contacts.edges[index].props.flag = InteractionFlag.Filtered
             }
         },
-        startUnit: (contacts: InteractionsIntraContacts) => { },
+        startUnit: () => {},
         handleIntraContact: (index: number, infoA: Features.Info, infoB: Features.Info) => {
             if (hasHydrogenBond(infoA, infoB)) {
                 const { flag } = interactions.unitsContacts.get(infoA.unit.id).edgeProps
@@ -170,4 +174,96 @@ function weakHydrogenBondsRefiner(structure: Structure, interactions: Interactio
             }
         }
     }
+}
+
+//
+
+function filterInter(type: InteractionType, index: number, infoA: Features.Info, infoB: Features.Info, contacts: InteractionsInterContacts) {
+    const aI = infoA.members[infoA.offsets[infoA.feature]]
+    const bI = infoB.members[infoB.offsets[infoB.feature]]
+    const indices = contacts.getContactIndicesForElement(aI, infoA.unit)
+
+    for (let i = 0, il = indices.length; i < il; ++i) {
+        if (contacts.edges[i].props.type === type) {
+            if (contacts.getContactIndicesForElement(bI, infoB.unit).includes(i)) {
+                contacts.edges[index].props.flag = InteractionFlag.Filtered
+                return
+            }
+        }
+    }
+}
+
+function filterIntra(type: InteractionType, index: number, infoA: Features.Info, infoB: Features.Info, contacts: InteractionsIntraContacts) {
+    const { edgeProps: { type: _type, flag }, elementsIndex: { offsets, indices } } = contacts
+    const aI = infoA.members[infoA.offsets[infoA.feature]]
+    const bI = infoB.members[infoB.offsets[infoB.feature]]
+
+    for (let i = offsets[aI], il = offsets[aI + 1]; i < il; ++i) {
+        const cI = indices[i]
+        if (_type[cI] === type) {
+            for (let j = offsets[bI], jl = offsets[bI + 1]; j < jl; ++j) {
+                if (cI === indices[j]) {
+                    flag[index] = InteractionFlag.Filtered
+                    return
+                }
+            }
+        }
+    }
+}
+
+/**
+ * Remove hydrogen bonds (normal and weak) between groups that also form
+ * an ionic interaction between each other
+ */
+function saltBridgeRefiner(structure: Structure, interactions: Interactions): ContactRefiner {
+    const { contacts } = interactions
+
+    return {
+        isApplicable: (type: InteractionType) => type === InteractionType.HydrogenBond || type === InteractionType.WeakHydrogenBond,
+        handleInterContact: (index: number, infoA: Features.Info, infoB: Features.Info) => {
+            filterInter(InteractionType.Ionic, index, infoA, infoB, contacts)
+        },
+        startUnit: () => {},
+        handleIntraContact: (index: number, infoA: Features.Info, infoB: Features.Info) => {
+            filterIntra(InteractionType.Ionic, index, infoA, infoB, interactions.unitsContacts.get(infoA.unit.id))
+        }
+    }
+}
+
+/**
+ * Remove hydrophobic and cation-pi interactions between groups that also form
+ * a pi-stacking interaction between each other
+ */
+function piStackingRefiner(structure: Structure, interactions: Interactions): ContactRefiner {
+    const { contacts } = interactions
+
+    return {
+        isApplicable: (type: InteractionType) => type === InteractionType.Hydrophobic || type === InteractionType.CationPi,
+        handleInterContact: (index: number, infoA: Features.Info, infoB: Features.Info) => {
+            filterInter(InteractionType.PiStacking, index, infoA, infoB, contacts)
+        },
+        startUnit: () => {},
+        handleIntraContact: (index: number, infoA: Features.Info, infoB: Features.Info) => {
+            filterIntra(InteractionType.PiStacking, index, infoA, infoB, interactions.unitsContacts.get(infoA.unit.id))
+        }
+    }
+}
+
+/**
+ * Remove ionic interactions between groups that also form
+ * a metal coordination between each other
+ */
+function metalCoordinationRefiner(structure: Structure, interactions: Interactions): ContactRefiner {
+    const { contacts } = interactions
+
+    return {
+        isApplicable: (type: InteractionType) => type === InteractionType.Ionic,
+        handleInterContact: (index: number, infoA: Features.Info, infoB: Features.Info) => {
+            filterInter(InteractionType.MetalCoordination, index, infoA, infoB, contacts)
+        },
+        startUnit: () => {},
+        handleIntraContact: (index: number, infoA: Features.Info, infoB: Features.Info) => {
+            filterIntra(InteractionType.MetalCoordination, index, infoA, infoB, interactions.unitsContacts.get(infoA.unit.id))
+        }
+    }
 }