fix multi-model loading & intra bonds caching

- moved atomSourceIndex to a separate property outside the data so diffing works

src/mol-model-formats/structure/basic/atomic.ts

@@ -62,8 +62,7 @@ function createHierarchyData(atom_site: AtomSite, sourceIndex: Column<number>, o
         label_alt_id: atom_site.label_alt_id,
         label_comp_id: atom_site.label_comp_id,
         auth_comp_id: atom_site.auth_comp_id,
-        pdbx_formal_charge: atom_site.pdbx_formal_charge,
-        sourceIndex
+        pdbx_formal_charge: atom_site.pdbx_formal_charge
     });
 
     const residues = Table.view(atom_site, ResiduesSchema, offsets.residues);
@@ -94,7 +93,7 @@ function createHierarchyData(atom_site: AtomSite, sourceIndex: Column<number>, o
     substUndefinedColumn(residues, 'label_seq_id', 'auth_seq_id');
     substUndefinedColumn(chains, 'label_asym_id', 'auth_asym_id');
 
-    return { atoms, residues, chains };
+    return { atoms, residues, chains, atomSourceIndex: sourceIndex };
 }
 
 function getConformation(atom_site: AtomSite): AtomicConformation {
@@ -111,7 +110,6 @@ function getConformation(atom_site: AtomSite): AtomicConformation {
 }
 
 function isHierarchyDataEqual(a: AtomicData, b: AtomicData) {
-    // TODO need to cast because of how TS handles type resolution for interfaces https://github.com/Microsoft/TypeScript/issues/15300
     return Table.areEqual(a.chains, b.chains)
         && Table.areEqual(a.residues, b.residues)
         && Table.areEqual(a.atoms, b.atoms);

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

@@ -26,7 +26,7 @@ export function trajectoryFromPDB(pdb: PdbFile): Task<Trajectory> {
             //      would need to do model splitting again
             if (models.frameCount === 1) {
                 const first = models.representative;
-                const srcIndex = first.atomicHierarchy.atoms.sourceIndex;
+                const srcIndex = first.atomicHierarchy.atomSourceIndex;
                 const isIdentity = Column.isIdentity(srcIndex);
                 const srcIndexArray = isIdentity ? void 0 : srcIndex.toArray({ array: Int32Array });
 

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

@@ -92,7 +92,7 @@ export namespace Model {
         const trajectory: Model[] = [];
         const { frames } = coordinates;
 
-        const srcIndex = model.atomicHierarchy.atoms.sourceIndex;
+        const srcIndex = model.atomicHierarchy.atomSourceIndex;
         const isIdentity = Column.isIdentity(srcIndex);
         const srcIndexArray = isIdentity ? void 0 : srcIndex.toArray({ array: Int32Array });
 

src/mol-model/structure/model/properties/atomic/hierarchy.ts

@@ -50,12 +50,6 @@ export const AtomsSchema = {
      */
     pdbx_formal_charge: mmCIF.atom_site.pdbx_formal_charge,
 
-    /**
-     * The index of this atom in the input data.
-     * Required because of sorting of atoms.
-     */
-    sourceIndex: Column.Schema.int
-
     // id, occupancy and B_iso_or_equiv are part of conformation
 };
 
@@ -108,6 +102,11 @@ export type Chains = Table<ChainsSchema>
 
 export interface AtomicData {
     atoms: Atoms,
+    /**
+     * The index of this atom in the input data.
+     * Required because of sorting of atoms.
+     */
+    atomSourceIndex: Column<number>,
     residues: Residues,
     chains: Chains
 }

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

@@ -529,7 +529,7 @@ export namespace Loci {
 
     function sourceIndex(unit: Unit, element: ElementIndex) {
         return Unit.isAtomic(unit)
-            ? unit.model.atomicHierarchy.atoms.sourceIndex.value(element)
+            ? unit.model.atomicHierarchy.atomSourceIndex.value(element)
             // TODO: when implemented, this should map to the source index.
             : element;
     }

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

@@ -42,7 +42,7 @@ const atom = {
     occupancy: p(l => !Unit.isAtomic(l.unit) ?  notAtomic() : l.unit.model.atomicConformation.occupancy.value(l.element)),
     B_iso_or_equiv: p(l => !Unit.isAtomic(l.unit) ?  notAtomic() : l.unit.model.atomicConformation.B_iso_or_equiv.value(l.element)),
     sourceIndex: p(l => Unit.isAtomic(l.unit)
-        ? l.unit.model.atomicHierarchy.atoms.sourceIndex.value(l.element)
+        ? l.unit.model.atomicHierarchy.atomSourceIndex.value(l.element)
         // TODO: when implemented, this should map to the source index.
         : l.element),
 

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

@@ -23,6 +23,7 @@ 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';
+import { ElementSetIntraBondCache } from './unit/bonds/element-set-intra-bond-cache';
 
 /**
  * A building block of a structure that corresponds to an atomic or
@@ -210,7 +211,6 @@ namespace Unit {
                 const { x, y, z } = this.model.atomicConformation;
                 boundary = tryAdjustBoundary({ x, y, z, indices: this.elements }, boundary);
             }
-            // 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)
@@ -240,7 +240,14 @@ namespace Unit {
 
         get bonds() {
             if (this.props.bonds) return this.props.bonds;
-            this.props.bonds = computeIntraUnitBonds(this);
+
+            const cache = ElementSetIntraBondCache.get(this.model);
+            let bonds = cache.get(this.elements);
+            if (!bonds) {
+                bonds = computeIntraUnitBonds(this);
+                cache.set(this.elements, bonds);
+            }
+            this.props.bonds = bonds;
             return this.props.bonds;
         }
 
@@ -472,7 +479,9 @@ namespace Unit {
             return void 0;
         }
 
-        if (old.props?.canRemap) return old;
+        if (old.props?.canRemap) {
+            return old;
+        }
         return isSameConformation(a, model) ? old : void 0;
     }
 

src/mol-model/structure/structure/unit/bonds/element-set-intra-bond-cache.ts

@@ -0,0 +1,45 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import StructureElement from '../../element';
+import { IntraUnitBonds } from './data';
+import { SortedArray } from '../../../../../mol-data/int';
+import { Model } from '../../../model';
+
+export class ElementSetIntraBondCache {
+    private data = new Map<number, [StructureElement.Set, IntraUnitBonds][]>();
+
+    get(xs: StructureElement.Set): IntraUnitBonds | undefined {
+        const hash = SortedArray.hashCode(xs);
+        if (!this.data.has(hash)) return void 0;
+        for (const [s, b] of this.data.get(hash)!) {
+            if (SortedArray.areEqual(xs, s)) return b;
+        }
+    }
+
+    set(xs: StructureElement.Set, bonds: IntraUnitBonds) {
+        const hash = SortedArray.hashCode(xs);
+        if (this.data.has(hash)) {
+            const es = this.data.get(hash)!;
+            for (const e of es) {
+                if (SortedArray.areEqual(xs, e[0])) {
+                    e[1] = bonds;
+                    return;
+                }
+            }
+            es.push([xs, bonds]);
+        } else {
+            this.data.set(hash, [[xs, bonds]]);
+        }
+    }
+
+    static get(model: Model): ElementSetIntraBondCache {
+        if (!model._dynamicPropertyData.ElementSetIntraBondCache) {
+            model._dynamicPropertyData.ElementSetIntraBondCache = new ElementSetIntraBondCache();
+        }
+        return model._dynamicPropertyData.ElementSetIntraBondCache;
+    }
+}