Fix IndexPairBonds for structures with re-ordered atoms

CHANGELOG.md

@@ -7,6 +7,7 @@ Note that since we don't clearly distinguish between a public and private interf
 
 - Add glTF (GLB) and STL support to ``geo-export`` extension.
 - Change O-S bond distance to allow for NOS bridges (doi:10.1038/s41586-021-03513-3)
+- Fix #178: ``IndexPairBonds`` for non-single residue structures (bug due to atom reordering).
 
 ## [v2.0.5] - 2021-04-26
 

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

@@ -131,32 +131,14 @@ export namespace Model {
         return new ArrayTrajectory(_trajectoryFromModelAndCoordinates(model, coordinates).trajectory);
     }
 
-    export function invertIndex(xs: ArrayLike<number>) {
-        const ret = new Int32Array(xs.length);
-        for (let i = 0, _i = xs.length; i < _i; i++) {
-            ret[xs[i]] = i;
-        }
-        return ret;
-    }
-
     export function trajectoryFromTopologyAndCoordinates(topology: Topology, coordinates: Coordinates): Task<Trajectory> {
         return Task.create('Create Trajectory', async ctx => {
             const models = await createModels(topology.basic, topology.sourceData, ctx);
             if (models.frameCount === 0) throw new Error('found no model');
             const model = models.representative;
-            const { trajectory, srcIndexArray } = _trajectoryFromModelAndCoordinates(model, coordinates);
-
-            // TODO: cache the inverted index somewhere?
-            const invertedIndex = srcIndexArray ? invertIndex(srcIndexArray) : void 0;
-            const pairs = srcIndexArray
-                ? {
-                    indexA: Column.ofIntArray(Column.mapToArray(topology.bonds.indexA, i => invertedIndex![i], Int32Array)),
-                    indexB: Column.ofIntArray(Column.mapToArray(topology.bonds.indexB, i => invertedIndex![i], Int32Array)),
-                    order: topology.bonds.order
-                }
-                : topology.bonds;
-
-            const bondData = { pairs, count: model.atomicHierarchy.atoms._rowCount };
+            const { trajectory } = _trajectoryFromModelAndCoordinates(model, coordinates);
+
+            const bondData = { pairs: topology.bonds, count: model.atomicHierarchy.atoms._rowCount };
             const indexPairBonds = IndexPairBonds.fromData(bondData);
 
             let index = 0;
@@ -176,6 +158,25 @@ export namespace Model {
         return center;
     }
 
+    function invertIndex(xs: Column<number>) {
+        const invertedIndex = new Int32Array(xs.rowCount);
+        let isIdentity = false;
+        for (let i = 0, _i = xs.rowCount; i < _i; i++) {
+            const x = xs.value(i);
+            if (x !== i) isIdentity = false;
+            invertedIndex[x] = i;
+        }
+        return { isIdentity, invertedIndex: invertedIndex as unknown as ArrayLike<ElementIndex> };
+    }
+
+    const InvertedAtomSrcIndexProp = '__InvertedAtomSrcIndex__';
+    export function getInvertedAtomSourceIndex(model: Model): { isIdentity: boolean, invertedIndex: ArrayLike<ElementIndex> } {
+        if (model._staticPropertyData[InvertedAtomSrcIndexProp]) return model._staticPropertyData[InvertedAtomSrcIndexProp];
+        const index = invertIndex(model.atomicHierarchy.atomSourceIndex);
+        model._staticPropertyData[InvertedAtomSrcIndexProp] = index;
+        return index;
+    }
+
     const TrajectoryInfoProp = '__TrajectoryInfo__';
     export type TrajectoryInfo = { readonly index: number, readonly size: number }
     export const TrajectoryInfo = {

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

@@ -19,6 +19,7 @@ import { IndexPairBonds } from '../../../../../mol-model-formats/structure/prope
 import { InterUnitGraph } from '../../../../../mol-math/graph/inter-unit-graph';
 import { StructConn } from '../../../../../mol-model-formats/structure/property/bonds/struct_conn';
 import { equalEps } from '../../../../../mol-math/linear-algebra/3d/common';
+import { Model } from '../../../model';
 
 const MAX_RADIUS = 4;
 
@@ -48,7 +49,10 @@ function findPairBonds(unitA: Unit.Atomic, unitB: Unit.Atomic, props: BondComput
     const hasOccupancy = occupancyA.isDefined && occupancyB.isDefined;
 
     const structConn = unitA.model === unitB.model && StructConn.Provider.get(unitA.model);
-    const indexPairs = unitA.model === unitB.model && IndexPairBonds.Provider.get(unitA.model);
+    const indexPairs = !props.forceCompute && unitA.model === unitB.model && IndexPairBonds.Provider.get(unitA.model);
+
+    const { atomSourceIndex: sourceIndex } = unitA.model.atomicHierarchy;
+    const { invertedIndex } = indexPairs ? Model.getInvertedAtomSourceIndex(unitB.model) : { invertedIndex: void 0 };
 
     const structConnExhaustive = unitA.model === unitB.model && StructConn.isExhaustive(unitA.model);
 
@@ -70,8 +74,10 @@ function findPairBonds(unitA: Unit.Atomic, unitB: Unit.Atomic, props: BondComput
 
         if (!props.forceCompute && indexPairs) {
             const { order, distance, flag } = indexPairs.edgeProps;
-            for (let i = indexPairs.offset[aI], il = indexPairs.offset[aI + 1]; i < il; ++i) {
-                const bI = indexPairs.b[i];
+
+            const srcA = sourceIndex.value(aI);
+            for (let i = indexPairs.offset[srcA], il = indexPairs.offset[srcA + 1]; i < il; ++i) {
+                const bI = invertedIndex![indexPairs.b[i]];
 
                 const _bI = SortedArray.indexOf(unitB.elements, bI) as StructureElement.UnitIndex;
                 if (_bI < 0) continue;

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

@@ -51,6 +51,9 @@ function findIndexPairBonds(unit: Unit.Atomic) {
     const atomCount = unit.elements.length;
     const { edgeProps } = indexPairs;
 
+    const { atomSourceIndex: sourceIndex } = unit.model.atomicHierarchy;
+    const { invertedIndex } = Model.getInvertedAtomSourceIndex(unit.model);
+
     const atomA: StructureElement.UnitIndex[] = [];
     const atomB: StructureElement.UnitIndex[] = [];
     const flags: number[] = [];
@@ -60,8 +63,10 @@ function findIndexPairBonds(unit: Unit.Atomic) {
         const aI =  atoms[_aI];
         const isHa = type_symbol.value(aI) === 'H';
 
-        for (let i = indexPairs.offset[aI], il = indexPairs.offset[aI + 1]; i < il; ++i) {
-            const bI = indexPairs.b[i];
+        const srcA = sourceIndex.value(aI);
+
+        for (let i = indexPairs.offset[srcA], il = indexPairs.offset[srcA + 1]; i < il; ++i) {
+            const bI = invertedIndex[indexPairs.b[i]];
             if (aI >= bI) continue;
 
             const _bI = SortedArray.indexOf(unit.elements, bI) as StructureElement.UnitIndex;