Refactored AtomicKeys to AtomicIndex and added more types of lookups

src/mol-model-props/pdbe/structure-quality-report.ts

@@ -71,7 +71,7 @@ function createIssueMapFromJson(modelData: Model, data: any): IssueMap | undefin
 
                 for (const residue of model.residues) {
                     const auth_seq_id = residue.author_residue_number, ins_code = residue.author_insertion_code || '';
-                    const idx = modelData.atomicHierarchy.findResidueKey(entity_id, asym_id, '', auth_seq_id, ins_code);
+                    const idx = modelData.atomicHierarchy.index.findResidue(entity_id, asym_id, auth_seq_id, ins_code);
                     ret.set(idx, residue.outlier_types);
                 }
             }
@@ -86,7 +86,7 @@ function createIssueMapFromCif(modelData: Model, data: Table<typeof StructureQua
     const { label_entity_id, label_asym_id, auth_seq_id, pdbx_PDB_ins_code, issues, _rowCount } = data;
 
     for (let i = 0; i < _rowCount; i++) {
-        const idx = modelData.atomicHierarchy.findResidueKey(label_entity_id.value(i), label_asym_id.value(i), '', auth_seq_id.value(i), pdbx_PDB_ins_code.value(i));
+        const idx = modelData.atomicHierarchy.index.findResidue(label_entity_id.value(i), label_asym_id.value(i), auth_seq_id.value(i), pdbx_PDB_ins_code.value(i));
         ret.set(idx, issues.value(i));
     }
 

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

@@ -12,7 +12,7 @@ import { ElementIndex } from '../../../../structure';
 import Format from '../../format';
 import { Model } from '../../model';
 import { AtomicConformation, AtomicData, AtomicHierarchy, AtomicSegments, AtomsSchema, ChainsSchema, ResiduesSchema } from '../../properties/atomic';
-import { getAtomicKeys } from '../../properties/utils/atomic-keys';
+import { getAtomicIndex } from '../../properties/utils/atomic-index';
 import { ElementSymbol } from '../../types';
 import { Entities } from '../../properties/common';
 
@@ -98,8 +98,8 @@ export function getAtomicHierarchyAndConformation(format: mmCIF_Format, atom_sit
         chainAtomSegments: Segmentation.ofOffsets(hierarchyOffsets.chains, Interval.ofBounds(0, atom_site._rowCount)),
     }
 
-    const hierarchyKeys = getAtomicKeys(hierarchyData, entities, hierarchySegments);
+    const index = getAtomicIndex(hierarchyData, entities, hierarchySegments);
     const hierarchyRanges = getAtomicRanges(hierarchyData, hierarchySegments, conformation, formatData.chemicalComponentMap);
-    const hierarchy: AtomicHierarchy = { ...hierarchyData, ...hierarchyKeys, ...hierarchySegments, ...hierarchyRanges };
+    const hierarchy: AtomicHierarchy = { ...hierarchyData, ...hierarchySegments, ...hierarchyRanges, index };
     return { sameAsPrevious: false, hierarchy, conformation };
 }

src/mol-model/structure/model/formats/mmcif/bonds/struct_conn.ts

@@ -163,7 +163,6 @@ export namespace StructConn {
         const { conn_type_id, pdbx_dist_value, pdbx_value_order } = struct_conn;
         const p1 = {
             label_asym_id: struct_conn.ptnr1_label_asym_id,
-            label_comp_id: struct_conn.ptnr1_label_comp_id,
             label_seq_id: struct_conn.ptnr1_label_seq_id,
             auth_seq_id: struct_conn.ptnr1_auth_seq_id,
             label_atom_id: struct_conn.ptnr1_label_atom_id,
@@ -173,7 +172,6 @@ export namespace StructConn {
         };
         const p2: typeof p1 = {
             label_asym_id: struct_conn.ptnr2_label_asym_id,
-            label_comp_id: struct_conn.ptnr2_label_comp_id,
             label_seq_id: struct_conn.ptnr2_label_seq_id,
             auth_seq_id: struct_conn.ptnr2_auth_seq_id,
             label_atom_id: struct_conn.ptnr2_label_atom_id,
@@ -184,11 +182,10 @@ export namespace StructConn {
 
         const _p = (row: number, ps: typeof p1) => {
             if (ps.label_asym_id.valueKind(row) !== Column.ValueKind.Present) return void 0;
-            const asymId = ps.label_asym_id.value(row)
-            const residueIndex = model.atomicHierarchy.findResidueKey(
+            const asymId = ps.label_asym_id.value(row);
+            const residueIndex = model.atomicHierarchy.index.findResidue(
                 findEntityIdByAsymId(model, asymId),
                 asymId,
-                ps.label_comp_id.value(row),
                 ps.auth_seq_id.value(row),
                 ps.ins_code.value(row)
             );

src/mol-model/structure/model/formats/mmcif/pair-restraints/cross-links.ts

@@ -11,9 +11,9 @@ import { findAtomIndexByLabelName } from '../util';
 import { Unit } from '../../../../structure';
 import { ElementIndex } from '../../../indexing';
 
-function findAtomIndex(model: Model, entityId: string, asymId: string, compId: string, seqId: number, atomId: string) {
+function findAtomIndex(model: Model, entityId: string, asymId: string, seqId: number, atomId: string) {
     if (!model.atomicHierarchy.atoms.auth_atom_id.isDefined) return -1
-    const residueIndex = model.atomicHierarchy.findResidueKey(entityId, compId, asymId, seqId, '')
+    const residueIndex = model.atomicHierarchy.index.findResidue(entityId, asymId, seqId)
     if (residueIndex < 0) return -1
     return findAtomIndexByLabelName(model, residueIndex, atomId, '') as ElementIndex
 }
@@ -35,7 +35,6 @@ export namespace IHMCrossLinkRestraint {
         const p1 = {
             entity_id: ihm_cross_link_restraint.entity_id_1,
             asym_id: ihm_cross_link_restraint.asym_id_1,
-            comp_id: ihm_cross_link_restraint.comp_id_1,
             seq_id: ihm_cross_link_restraint.seq_id_1,
             atom_id: ihm_cross_link_restraint.atom_id_1,
         }
@@ -43,7 +42,6 @@ export namespace IHMCrossLinkRestraint {
         const p2: typeof p1 = {
             entity_id: ihm_cross_link_restraint.entity_id_2,
             asym_id: ihm_cross_link_restraint.asym_id_2,
-            comp_id: ihm_cross_link_restraint.comp_id_2,
             seq_id: ihm_cross_link_restraint.seq_id_2,
             atom_id: ihm_cross_link_restraint.atom_id_2,
         }
@@ -60,7 +58,7 @@ export namespace IHMCrossLinkRestraint {
             const seqId = ps.seq_id.value(row)
 
             if (ihm_cross_link_restraint.model_granularity.value(row) === 'by-atom') {
-                const atomicElement = findAtomIndex(model, entityId, asymId, ps.comp_id.value(row), seqId, ps.atom_id.value(row))
+                const atomicElement = findAtomIndex(model, entityId, asymId, seqId, ps.atom_id.value(row))
                 if (atomicElement >= 0) _add(atomicElementMap, atomicElement as ElementIndex, row)
             } else if (model.coarseHierarchy.isDefined) {
                 const sphereElement = model.coarseHierarchy.spheres.findSequenceKey(entityId, asymId, seqId)

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

@@ -119,22 +119,6 @@ export interface AtomicSegments {
     // TODO: include entity segments?
 }
 
-export interface AtomicKeys {
-    // TODO: include (lazily computed) "entity/chain/residue" indices?
-
-    /** @returns index or -1 if not present. */
-    getEntityKey(cI: ChainIndex): EntityIndex,
-
-    /** @returns index or -1 if not present. */
-    findChainKey(entityId: string, label_asym_id: string): ChainIndex,
-
-    /**
-     * Unique number for each of the residue. Also the index of the 1st occurence of this residue.
-     * @returns index or -1 if not present.
-     */
-    findResidueKey(entityId: string, label_asym_id: string, label_comp_id: string, auth_seq_id: number, pdbx_PDB_ins_code: string): ResidueIndex
-}
-
 export interface AtomicIndex {
     /** @returns index or -1 if not present. */
     getEntityFromChain(cI: ChainIndex): EntityIndex,
@@ -157,6 +141,7 @@ export interface AtomicIndex {
      * @returns index or -1 if not present.
      */
     findResidue(key: AtomicIndex.ResidueKey): ResidueIndex,
+    findResidue(label_entity_id: string, label_asym_id: string, auth_seq_id: number, pdbx_PDB_ins_code?: string): ResidueIndex,
 
     /**
      * Index of the 1st occurence of this residue.
@@ -181,6 +166,8 @@ export namespace AtomicIndex {
     export interface ChainAuthKey { auth_asym_id: string }
 
     export interface ResidueKey { label_entity_id: string, label_asym_id: string, auth_seq_id: number, pdbx_PDB_ins_code?: string }
+    export function EmptyResidueKey(): ResidueKey { return { label_entity_id: '', label_asym_id: '', auth_seq_id: 0, pdbx_PDB_ins_code: void 0 }; }
+
     export interface ResidueAuthKey { auth_asym_id: string, auth_comp_id: string, auth_seq_id: number, pdbx_PDB_ins_code?: string }
     export interface ResidueLabelKey { label_entity_id: string, label_asym_id: string, label_seq_id: number, pdbx_PDB_ins_code?: string }
 }
@@ -191,8 +178,10 @@ export interface AtomicRanges {
     cyclicPolymerMap: Map<ResidueIndex, ResidueIndex>
 }
 
-type _Hierarchy = AtomicData & AtomicSegments & AtomicKeys & AtomicRanges
-export interface AtomicHierarchy extends _Hierarchy { }
+type _Hierarchy = AtomicData & AtomicSegments & AtomicRanges
+export interface AtomicHierarchy extends _Hierarchy {
+    index: AtomicIndex
+}
 
 export namespace AtomicHierarchy {
     /** Start residue inclusive */

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

@@ -32,13 +32,13 @@ namespace StructureSequence {
         const sequences: StructureSequence.Entity[] = [];
 
         for (let cI = 0 as ChainIndex, _cI = hierarchy.chains._rowCount; cI < _cI; cI++) {
-            const entityKey = hierarchy.getEntityKey(cI);
+            const entityKey = hierarchy.index.getEntityFromChain(cI);
             // Only for polymers, trying to mirror _entity_poly_seq
             if (byEntityKey[entityKey] !== void 0 || entities.data.type.value(entityKey) !== 'polymer') continue;
 
             let start = cI;
             cI++;
-            while (cI < _cI && entityKey === hierarchy.getEntityKey(cI) && entities.data.type.value(entityKey) !== 'polymer') {
+            while (cI < _cI && entityKey === hierarchy.index.getEntityFromChain(cI) && entities.data.type.value(entityKey) !== 'polymer') {
                 cI++;
             }
             cI--;

src/mol-model/structure/model/properties/utils/atomic-index.ts

@@ -62,6 +62,7 @@ function createMapping(entities: Entities, data: AtomicData, segments: AtomicSeg
     };
 }
 
+const _tempResidueKey = AtomicIndex.EmptyResidueKey();
 class Index implements AtomicIndex {
     private entityIndex: Entities['getEntityIndex'];
 
@@ -81,7 +82,19 @@ class Index implements AtomicIndex {
         return this.map.auth_asym_id.has(key.auth_asym_id) ? this.map.auth_asym_id.get(key.auth_asym_id)! : -1 as ChainIndex;
     }
 
-    findResidue(key: AtomicIndex.ResidueKey): ResidueIndex {
+    findResidue(label_entity_id: string, label_asym_id: string, auth_seq_id: number, pdbx_PDB_ins_code?: string): ResidueIndex
+    findResidue(key: AtomicIndex.ResidueKey): ResidueIndex
+    findResidue(label_entity_id_or_key: string | AtomicIndex.ResidueKey, label_asym_id?: string, auth_seq_id?: number, pdbx_PDB_ins_code?: string): ResidueIndex {
+        let key: AtomicIndex.ResidueKey;
+        if (arguments.length === 1) {
+            key = label_entity_id_or_key as AtomicIndex.ResidueKey
+        } else {
+            _tempResidueKey.label_entity_id = label_entity_id_or_key as string;
+            _tempResidueKey.label_asym_id = label_asym_id!;
+            _tempResidueKey.auth_seq_id = auth_seq_id!;
+            _tempResidueKey.pdbx_PDB_ins_code = pdbx_PDB_ins_code;
+            key = _tempResidueKey;
+        }
         const cI = this.findChainLabel(key);
         if (cI < 0) return -1 as ResidueIndex;
         const rm = this.map.chain_index_auth_seq_id.get(cI)!;
@@ -117,7 +130,7 @@ class Index implements AtomicIndex {
     }
 }
 
-export function getAtomicKeys(data: AtomicData, entities: Entities, segments: AtomicSegments): AtomicIndex {
+export function getAtomicIndex(data: AtomicData, entities: Entities, segments: AtomicSegments): AtomicIndex {
     const map = createMapping(entities, data, segments);
 
     const { label_seq_id, auth_seq_id, pdbx_PDB_ins_code } = data.residues;

src/mol-model/structure/model/properties/utils/atomic-keys.ts

@@ -1,99 +0,0 @@
-/**
- * Copyright (c) 2017 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author David Sehnal <david.sehnal@gmail.com>
- */
-
-import { AtomicData, AtomicSegments, AtomicKeys } from '../atomic'
-import { Interval, Segmentation } from 'mol-data/int'
-import { Entities } from '../common'
-import { ChainIndex, ResidueIndex, EntityIndex } from '../../indexing';
-
-function getResidueId(comp_id: string, seq_id: number, ins_code: string) {
-    // TODO: add new index that support comp_id again?
-    return `${seq_id} ${ins_code}`;
-    //return `${comp_id} ${seq_id} ${ins_code}`;
-}
-
-function getElementKey(map: Map<string, number>, key: string, counter: { index: number }) {
-    if (map.has(key)) return map.get(key)!;
-    const ret = counter.index++;
-    map.set(key, ret);
-    return ret;
-}
-
-function getElementSubstructureKeyMap(map: Map<number, Map<string, number>>, key: number) {
-    if (map.has(key)) return map.get(key)!;
-    const ret = new Map<string, number>();
-    map.set(key, ret);
-    return ret;
-}
-
-function createLookUp(entities: Entities, chain: Map<number, Map<string, number>>, residue: Map<number, Map<string, number>>) {
-    const getEntKey = entities.getEntityIndex;
-    const findChainKey: AtomicKeys['findChainKey'] = (e, c) => {
-        let eKey = getEntKey(e);
-        if (eKey < 0) return -1 as ChainIndex;
-        const cm = chain.get(eKey)!;
-        if (!cm.has(c)) return -1 as ChainIndex;
-        return cm.get(c)! as ChainIndex;
-    }
-    const findResidueKey: AtomicKeys['findResidueKey'] = (e, c, name, seq, ins) => {
-        let eKey = getEntKey(e);
-        if (eKey < 0) return -1 as ResidueIndex;
-        const cm = chain.get(eKey)!;
-        if (!cm.has(c)) return -1 as ResidueIndex;
-        const rm = residue.get(cm.get(c)!)!
-        const id = getResidueId(name, seq, ins);
-        if (!rm.has(id)) return -1 as ResidueIndex;
-        return rm.get(id)! as ResidueIndex;
-    }
-    return { findChainKey, findResidueKey };
-}
-
-function missingEntity(k: string) {
-    throw new Error(`Missing entity entry for entity id '${k}'.`);
-}
-
-export function getAtomicKeys(data: AtomicData, entities: Entities, segments: AtomicSegments): AtomicKeys {
-    const { chains, residues } = data;
-
-    const chainMaps = new Map<number, Map<string, number>>(), chainCounter = { index: 0 };
-    const residueMaps = new Map<number, Map<string, number>>(), residueCounter = { index: 0 };
-
-    const residueKey = new Int32Array(residues._rowCount);
-    const chainKey = new Int32Array(chains._rowCount);
-    const entityKey = new Int32Array(chains._rowCount);
-
-    const { label_comp_id, auth_seq_id, pdbx_PDB_ins_code } = data.residues;
-    const { label_entity_id, label_asym_id } = data.chains;
-
-    const atomSet = Interval.ofBounds(0, data.atoms._rowCount);
-
-    const chainsIt = Segmentation.transientSegments(segments.chainAtomSegments, atomSet);
-    while (chainsIt.hasNext) {
-        const chainSegment = chainsIt.move();
-        const cI = chainSegment.index;
-
-        let eKey = entities.getEntityIndex(label_entity_id.value(cI));
-        if (eKey < 0) missingEntity(label_entity_id.value(cI));
-        const chainMap = getElementSubstructureKeyMap(chainMaps, eKey);
-        const cKey = getElementKey(chainMap, label_asym_id.value(cI), chainCounter);
-
-        chainKey[cI] = cKey;
-        entityKey[cI] = eKey;
-
-        const residueMap = getElementSubstructureKeyMap(residueMaps, cKey);
-        const residuesIt = Segmentation.transientSegments(segments.residueAtomSegments, atomSet, chainSegment);
-        while (residuesIt.hasNext) {
-            const residueSegment = residuesIt.move();
-            const rI = residueSegment.index;
-            const residueId = getResidueId(label_comp_id.value(rI), auth_seq_id.value(rI), pdbx_PDB_ins_code.value(rI));
-            residueKey[rI] = getElementKey(residueMap, residueId, residueCounter);
-        }
-    }
-
-    const { findChainKey, findResidueKey } = createLookUp(entities, chainMaps, residueMaps);
-
-    return { getEntityKey: cI => entityKey[cI] as EntityIndex, findChainKey, findResidueKey };
-}

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

@@ -99,7 +99,7 @@ namespace StructureElement {
     export function entityIndex(l: StructureElement) {
         switch (l.unit.kind) {
             case Unit.Kind.Atomic:
-                return l.unit.model.atomicHierarchy.getEntityKey(l.unit.chainIndex[l.element])
+                return l.unit.model.atomicHierarchy.index.getEntityFromChain(l.unit.chainIndex[l.element])
             case Unit.Kind.Spheres:
                 return l.unit.model.coarseHierarchy.spheres.entityKey[l.element]
             case Unit.Kind.Gaussians: