mol-model: Unit.Atomic.tryRemapBonds

src/mol-math/graph/int-adjacency-graph.ts

@@ -17,13 +17,14 @@ import { AssignableArrayLike } from '../../mol-util/type-helpers';
  *
  * Edge properties are indexed same as in the arrays a and b.
  */
-export interface IntAdjacencyGraph<VertexIndex extends number, EdgeProps extends IntAdjacencyGraph.EdgePropsBase> {
+export interface IntAdjacencyGraph<VertexIndex extends number, EdgeProps extends IntAdjacencyGraph.EdgePropsBase, Props = any> {
     readonly offset: ArrayLike<number>,
     readonly a: ArrayLike<VertexIndex>,
     readonly b: ArrayLike<VertexIndex>,
     readonly vertexCount: number,
     readonly edgeCount: number,
-    readonly edgeProps: Readonly<EdgeProps>
+    readonly edgeProps: Readonly<EdgeProps>,
+    readonly props?: Props
 
     /**
      * Get the edge index between i-th and j-th vertex.
@@ -49,6 +50,8 @@ export namespace IntAdjacencyGraph {
     export type EdgePropsBase = { [name: string]: ArrayLike<any> }
 
     export function areEqual<I extends number, P extends IntAdjacencyGraph.EdgePropsBase>(a: IntAdjacencyGraph<I, P>, b: IntAdjacencyGraph<I, P>) {
+        if (a === b) return true;
+
         if (a.vertexCount !== b.vertexCount || a.edgeCount !== b.edgeCount) return false;
 
         const { a: aa, b: ab, offset: ao } = a;
@@ -78,7 +81,7 @@ export namespace IntAdjacencyGraph {
         return true;
     }
 
-    class IntGraphImpl<VertexIndex extends number, EdgeProps extends IntAdjacencyGraph.EdgePropsBase> implements IntAdjacencyGraph<VertexIndex, EdgeProps> {
+    class IntGraphImpl<VertexIndex extends number, EdgeProps extends IntAdjacencyGraph.EdgePropsBase, Props> implements IntAdjacencyGraph<VertexIndex, EdgeProps, Props> {
         readonly vertexCount: number;
         readonly edgeProps: EdgeProps;
 
@@ -106,14 +109,14 @@ export namespace IntAdjacencyGraph {
             return this.offset[i + 1] - this.offset[i];
         }
 
-        constructor(public offset: ArrayLike<number>, public a: ArrayLike<VertexIndex>, public b: ArrayLike<VertexIndex>, public edgeCount: number, edgeProps?: EdgeProps) {
+        constructor(public offset: ArrayLike<number>, public a: ArrayLike<VertexIndex>, public b: ArrayLike<VertexIndex>, public edgeCount: number, edgeProps?: EdgeProps, public props?: Props) {
             this.vertexCount = offset.length - 1;
             this.edgeProps = (edgeProps || {}) as EdgeProps;
         }
     }
 
-    export function create<VertexIndex extends number, EdgeProps extends IntAdjacencyGraph.EdgePropsBase>(offset: ArrayLike<number>, a: ArrayLike<VertexIndex>, b: ArrayLike<VertexIndex>, edgeCount: number, edgeProps?: EdgeProps): IntAdjacencyGraph<VertexIndex, EdgeProps> {
-        return new IntGraphImpl(offset, a, b, edgeCount, edgeProps) as IntAdjacencyGraph<VertexIndex, EdgeProps>;
+    export function create<VertexIndex extends number, EdgeProps extends IntAdjacencyGraph.EdgePropsBase, Props>(offset: ArrayLike<number>, a: ArrayLike<VertexIndex>, b: ArrayLike<VertexIndex>, edgeCount: number, edgeProps?: EdgeProps, props?: Props): IntAdjacencyGraph<VertexIndex, EdgeProps, Props> {
+        return new IntGraphImpl(offset, a, b, edgeCount, edgeProps, props) as IntAdjacencyGraph<VertexIndex, EdgeProps, Props>;
     }
 
     export class EdgeBuilder<VertexIndex extends number> {
@@ -129,8 +132,8 @@ export namespace IntAdjacencyGraph {
         a: AssignableArrayLike<VertexIndex>;
         b: AssignableArrayLike<VertexIndex>;
 
-        createGraph<EdgeProps extends IntAdjacencyGraph.EdgePropsBase>(edgeProps: EdgeProps) {
-            return create<VertexIndex, EdgeProps>(this.offsets, this.a, this.b, this.edgeCount, edgeProps);
+        createGraph<EdgeProps extends IntAdjacencyGraph.EdgePropsBase, Props>(edgeProps: EdgeProps, props?: Props) {
+            return create<VertexIndex, EdgeProps, Props>(this.offsets, this.a, this.b, this.edgeCount, edgeProps, props);
         }
 
         /**

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

@@ -22,6 +22,7 @@ import { PrincipalAxes } from '../../../mol-math/linear-algebra/matrix/principal
 import { getPrincipalAxes } from './util/principal-axes';
 import { Boundary, getBoundary, tryAdjustBoundary } from '../../../mol-math/geometry/boundary';
 import { Mat4 } from '../../../mol-math/linear-algebra';
+import { IndexPairBonds } from '../../../mol-model-formats/structure/property/bonds/index-pair';
 
 /**
  * A building block of a structure that corresponds to an atomic or
@@ -209,8 +210,8 @@ namespace Unit {
                 const { x, y, z } = this.model.atomicConformation;
                 boundary = tryAdjustBoundary({ x, y, z, indices: this.elements }, boundary);
             }
-            const bonds = this.props.bonds && canRemapBonds(this, model) ? this.props.bonds : void 0;
-            const props = { ...this.props, bonds, boundary, lookup3d: undefined, principalAxes: undefined };
+            // TODO: add element set based bond caching?
+            const props = { ...this.props, bonds: tryRemapBonds(this, this.props.bonds, model), boundary, lookup3d: undefined, principalAxes: undefined };
             const conformation = this.model.atomicConformation !== model.atomicConformation
                 ? SymmetryOperator.createMapping(this.conformation.operator, model.atomicConformation)
                 : this.conformation;
@@ -457,7 +458,25 @@ namespace Unit {
         return true;
     }
 
-    function canRemapBonds(a: Atomic, model: Model) {
+    function tryRemapBonds(a: Atomic, old: IntraUnitBonds | undefined, model: Model) {
+        // TODO: should include additional checks?
+
+        if (!old) return void 0;
+        if (a.model.atomicConformation.id === model.atomicConformation.id) return old;
+
+        const oldIndex = IndexPairBonds.Provider.get(a.model);
+        if (oldIndex) {
+            const newIndex = IndexPairBonds.Provider.get(model);
+            // TODO: check the actual indices instead of just reference equality?
+            if (!newIndex || oldIndex === newIndex) return old;
+            return void 0;
+        }
+
+        if (old.props?.canRemap) return old;
+        return isSameConformation(a, model) ? old : void 0;
+    }
+
+    function isSameConformation(a: Atomic, model: Model) {
         const xs = a.elements;
         const { x: xa, y: ya, z: za } = a.conformation.coordinates;
         const { x: xb, y: yb, z: zb } = model.atomicConformation;

src/mol-model/structure/structure/unit/bonds/data.ts

@@ -12,7 +12,7 @@ import StructureElement from '../../element';
 import { Bond } from '../bonds';
 import { InterUnitGraph } from '../../../../../mol-math/graph/inter-unit-graph';
 
-type IntraUnitBonds = IntAdjacencyGraph<StructureElement.UnitIndex, { readonly order: ArrayLike<number>, readonly flags: ArrayLike<BondType.Flag> }>
+type IntraUnitBonds = IntAdjacencyGraph<StructureElement.UnitIndex, { readonly order: ArrayLike<number>, readonly flags: ArrayLike<BondType.Flag> }, { readonly canRemap?: boolean }>
 
 namespace IntraUnitBonds {
     export const Empty: IntraUnitBonds = IntAdjacencyGraph.create([], [], [], 0, { flags: [], order: [] });

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

@@ -20,7 +20,7 @@ import { Vec3 } from '../../../../../mol-math/linear-algebra';
 import { ElementIndex } from '../../../model/indexing';
 import { equalEps } from '../../../../../mol-math/linear-algebra/3d/common';
 
-function getGraph(atomA: StructureElement.UnitIndex[], atomB: StructureElement.UnitIndex[], _order: number[], _flags: number[], atomCount: number): IntraUnitBonds {
+function getGraph(atomA: StructureElement.UnitIndex[], atomB: StructureElement.UnitIndex[], _order: number[], _flags: number[], atomCount: number, canRemap: boolean): IntraUnitBonds {
     const builder = new IntAdjacencyGraph.EdgeBuilder(atomCount, atomA, atomB);
     const flags = new Uint16Array(builder.slotCount);
     const order = new Int8Array(builder.slotCount);
@@ -30,7 +30,7 @@ function getGraph(atomA: StructureElement.UnitIndex[], atomB: StructureElement.U
         builder.assignProperty(order, _order[i]);
     }
 
-    return builder.createGraph({ flags, order });
+    return builder.createGraph({ flags, order }, { canRemap });
 }
 
 const tmpDistVecA = Vec3();
@@ -43,7 +43,44 @@ function getDistance(unit: Unit.Atomic, indexA: ElementIndex, indexB: ElementInd
 
 const __structConnAdded = new Set<StructureElement.UnitIndex>();
 
-function _computeBonds(unit: Unit.Atomic, props: BondComputationProps): IntraUnitBonds {
+function findIndexPairBonds(unit: Unit.Atomic) {
+    const indexPairs = IndexPairBonds.Provider.get(unit.model)!;
+    const { elements: atoms } = unit;
+    const { type_symbol } = unit.model.atomicHierarchy.atoms;
+    const atomCount = unit.elements.length;
+    const { edgeProps } = indexPairs;
+
+    const atomA: StructureElement.UnitIndex[] = [];
+    const atomB: StructureElement.UnitIndex[] = [];
+    const flags: number[] = [];
+    const order: number[] = [];
+
+    for (let _aI = 0 as StructureElement.UnitIndex; _aI < atomCount; _aI++) {
+        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];
+            if (aI >= bI) continue;
+
+            const _bI = SortedArray.indexOf(unit.elements, bI) as StructureElement.UnitIndex;
+            if (_bI < 0) continue;
+            if (isHa && type_symbol.value(bI) === 'H') continue;
+
+            const d = edgeProps.distance[i];
+            if (d === -1 || d === void 0 || equalEps(getDistance(unit, aI, bI), d, 0.5)) {
+                atomA[atomA.length] = _aI;
+                atomB[atomB.length] = _bI;
+                order[order.length] = edgeProps.order[i];
+                flags[flags.length] = edgeProps.flag[i];
+            }
+        }
+    }
+
+    return getGraph(atomA, atomB, order, flags, atomCount, false);
+}
+
+function findBonds(unit: Unit.Atomic, props: BondComputationProps): IntraUnitBonds {
     const MAX_RADIUS = 4;
 
     const { x, y, z } = unit.model.atomicConformation;
@@ -57,7 +94,6 @@ function _computeBonds(unit: Unit.Atomic, props: BondComputationProps): IntraUni
 
     const structConn = StructConn.Provider.get(unit.model);
     const component = ComponentBond.Provider.get(unit.model);
-    const indexPairs = IndexPairBonds.Provider.get(unit.model);
 
     const atomA: StructureElement.UnitIndex[] = [];
     const atomB: StructureElement.UnitIndex[] = [];
@@ -67,31 +103,15 @@ function _computeBonds(unit: Unit.Atomic, props: BondComputationProps): IntraUni
     let lastResidue = -1;
     let componentMap: Map<string, Map<string, { flags: number, order: number }>> | undefined = void 0;
 
+    let isWatery = true, isDictionaryBased = true, isSequenced = true;
+
     const structConnAdded = __structConnAdded;
 
     for (let _aI = 0 as StructureElement.UnitIndex; _aI < atomCount; _aI++) {
         const aI =  atoms[_aI];
 
-        if (!props.forceCompute && indexPairs) {
-            const { edgeProps } = indexPairs;
-            for (let i = indexPairs.offset[aI], il = indexPairs.offset[aI + 1]; i < il; ++i) {
-                const bI = indexPairs.b[i];
-                if (aI >= bI) continue;
-
-                const _bI = SortedArray.indexOf(unit.elements, bI) as StructureElement.UnitIndex;
-                if (_bI < 0) continue;
-                if (type_symbol.value(aI) === 'H' && type_symbol.value(bI) === 'H') continue;
-
-                const d = edgeProps.distance[i];
-                if (d === -1 || equalEps(getDistance(unit, aI, bI), d, 0.5)) {
-                    atomA[atomA.length] = _aI;
-                    atomB[atomB.length] = _bI;
-                    order[order.length] = edgeProps.order[i];
-                    flags[flags.length] = edgeProps.flag[i];
-                }
-            }
-            continue; // assume `indexPairs` supplies all bonds
-        }
+        const elemA = type_symbol.value(aI);
+        if (isWatery && (elemA !== 'H' || elemA !== 'O')) isWatery = false;
 
         const structConnEntries = props.forceCompute ? void 0 : structConn && structConn.byAtomIndex.get(aI);
         let hasStructConn = false;
@@ -117,12 +137,13 @@ function _computeBonds(unit: Unit.Atomic, props: BondComputationProps): IntraUni
         }
 
         const raI = residueIndex[aI];
+        const seqIdA = label_seq_id.value(raI);
         const compId = label_comp_id.value(aI);
 
         if (!props.forceCompute && raI !== lastResidue) {
             if (!!component && component.entries.has(compId)) {
                 const entitySeq = byEntityKey[index.getEntityFromChain(chainIndex[aI])];
-                if (entitySeq && entitySeq.sequence.microHet.has(label_seq_id.value(raI))) {
+                if (entitySeq && entitySeq.sequence.microHet.has(seqIdA)) {
                     // compute for sequence positions with micro-heterogeneity
                     componentMap = void 0;
                 } else {
@@ -134,7 +155,7 @@ function _computeBonds(unit: Unit.Atomic, props: BondComputationProps): IntraUni
         }
         lastResidue = raI;
 
-        const aeI = getElementIdx(type_symbol.value(aI));
+        const aeI = getElementIdx(elemA);
         const atomIdA = label_atom_id.value(aI);
         const componentPairs = componentMap ? componentMap.get(atomIdA) : void 0;
 
@@ -194,11 +215,17 @@ function _computeBonds(unit: Unit.Atomic, props: BondComputationProps): IntraUni
                 atomB[atomB.length] = _bI;
                 order[order.length] = getIntraBondOrderFromTable(compId, atomIdA, label_atom_id.value(bI));
                 flags[flags.length] = (isMetal ? BondType.Flag.MetallicCoordination : BondType.Flag.Covalent) | BondType.Flag.Computed;
+
+                const seqIdB = label_seq_id.value(rbI);
+
+                if (seqIdA === seqIdB) isDictionaryBased = false;
+                if (Math.abs(seqIdA - seqIdB) > 1) isSequenced = false;
             }
         }
     }
 
-    return getGraph(atomA, atomB, order, flags, atomCount);
+    const canRemap = isWatery || (isDictionaryBased && isSequenced);
+    return getGraph(atomA, atomB, order, flags, atomCount, canRemap);
 }
 
 function computeIntraUnitBonds(unit: Unit.Atomic, props?: Partial<BondComputationProps>) {
@@ -208,7 +235,12 @@ function computeIntraUnitBonds(unit: Unit.Atomic, props?: Partial<BondComputatio
         //      and avoid using unit.lookup
         return IntraUnitBonds.Empty;
     }
-    return _computeBonds(unit, p);
+
+    if (!p.forceCompute && IndexPairBonds.Provider.get(unit.model)!) {
+        return findIndexPairBonds(unit);
+    } else {
+        return findBonds(unit, p);
+    }
 }
 
 export { computeIntraUnitBonds };