fix struct conn for waters (#803)

CHANGELOG.md

@@ -7,6 +7,7 @@ Note that since we don't clearly distinguish between a public and private interf
 ## [Unreleased]
 
 - Add a uniform color theme for NtC tube that still paints residue and segment dividers in a different color
+- Fix bond assignments `struct_conn` records referencing waters
 
 ## [v3.34.0] - 2023-04-16
 

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

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2017-2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2017-2023 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
@@ -1225,8 +1225,10 @@ namespace Structure {
             const closeUnits = lookup.findUnitIndices(imageCenter[0], imageCenter[1], imageCenter[2], bs.radius + maxRadius);
             for (let i = 0; i < closeUnits.count; i++) {
                 const other = structure.units[closeUnits.indices[i]];
+                if (unit.id >= other.id) continue;
+
                 if (other.elements.length > 3 && !Box3D.overlaps(bbox, other.boundary.box)) continue;
-                if (!validUnit(other) || unit.id >= other.id || !validUnitPair(unit, other)) continue;
+                if (!validUnit(other) || !validUnitPair(unit, other)) continue;
 
                 if (other.elements.length >= unit.elements.length) callback(unit, other);
                 else callback(other, unit);

src/mol-model/structure/structure/unit/bonds/inter-compute.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2017-2022 Mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2017-2023 Mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
@@ -243,6 +243,11 @@ function computeInterUnitBonds(structure: Structure, props?: Partial<InterBondCo
         ...p,
         validUnit: (props && props.validUnit) || (u => Unit.isAtomic(u)),
         validUnitPair: (props && props.validUnitPair) || ((s, a, b) => {
+            // In case both units have a struct conn record, ignore other criteria
+            if (hasCommonStructConnRecord(a, b)) {
+                return Structure.validUnitPair(s, a, b);
+            }
+
             const mtA = a.model.atomicHierarchy.derived.residue.moleculeType;
             const mtB = b.model.atomicHierarchy.derived.residue.moleculeType;
             const notWater = (
@@ -250,25 +255,34 @@ function computeInterUnitBonds(structure: Structure, props?: Partial<InterBondCo
                 (!Unit.isAtomic(b) || mtB[b.residueIndex[b.elements[0]]] !== MoleculeType.Water)
             );
 
-            const sameModel = a.model === b.model;
-            const notIonA = (!Unit.isAtomic(a) || mtA[a.residueIndex[a.elements[0]]] !== MoleculeType.Ion) || (sameModel && hasStructConnRecord(a));
-            const notIonB = (!Unit.isAtomic(b) || mtB[b.residueIndex[b.elements[0]]] !== MoleculeType.Ion) || (sameModel && hasStructConnRecord(b));
+            const notIonA = (!Unit.isAtomic(a) || mtA[a.residueIndex[a.elements[0]]] !== MoleculeType.Ion);
+            const notIonB = (!Unit.isAtomic(b) || mtB[b.residueIndex[b.elements[0]]] !== MoleculeType.Ion);
             const notIon = notIonA && notIonB;
             return Structure.validUnitPair(s, a, b) && (notWater || !p.ignoreWater) && (notIon || !p.ignoreIon);
         }),
     });
 }
 
-function hasStructConnRecord(unit: Unit) {
-    if (!Unit.isAtomic(unit)) return false;
+function hasCommonStructConnRecord(unitA: Unit, unitB: Unit) {
+    if (unitA.model !== unitB.model || !Unit.isAtomic(unitA) || !Unit.isAtomic(unitB)) return false;
+    const structConn = StructConn.Provider.get(unitA.model);
+    if (!structConn) return false;
 
-    const elements = unit.elements;
-    const structConn = StructConn.Provider.get(unit.model);
-    if (structConn) {
-        for (let i = 0, _i = elements.length; i < _i; i++) {
-            if (structConn.byAtomIndex.get(elements[i])) {
-                return true;
-            }
+    const smaller = unitA.elements.length < unitB.elements.length ? unitA : unitB;
+    const bigger = unitA.elements.length >= unitB.elements.length ? unitA : unitB;
+
+    const { elements: xs } = smaller;
+    const { elements: ys } = bigger;
+    const { indexOf } = SortedArray;
+
+    for (let i = 0, _i = xs.length; i < _i; i++) {
+        const aI = xs[i];
+        const entries = structConn.byAtomIndex.get(aI);
+        if (!entries?.length) continue;
+
+        for (const e of entries) {
+            const bI = e.partnerA.atomIndex === aI ? e.partnerB.atomIndex : e.partnerA.atomIndex;
+            if (indexOf(ys, bI) >= 0) return true;
         }
     }
     return false;