wip, strongly type model indexing

src/mol-data/int/segmentation.ts

@@ -9,25 +9,25 @@ import OrderedSet from './ordered-set'
 import * as Impl from './impl/segmentation'
 
 namespace Segmentation {
-    export interface Segment<T extends number = number> { index: number, start: T, end: T }
+    export interface Segment<T extends number = number, I extends number = number> { index: number, start: T, end: T }
 
-    export const create: <T extends number = number>(segs: ArrayLike<T>) => Segmentation<T> = Impl.create as any;
-    export const ofOffsets: <T extends number = number>(offsets: ArrayLike<T>, bounds: Interval) => Segmentation<T> = Impl.ofOffsets as any;
+    export const create: <T extends number = number, I extends number = number>(segs: ArrayLike<T>) => Segmentation<T, I> = Impl.create as any;
+    export const ofOffsets: <T extends number = number, I extends number = number>(offsets: ArrayLike<T>, bounds: Interval) => Segmentation<T, I> = Impl.ofOffsets as any;
 
-    export const count: <T extends number = number>(segs: Segmentation<T>) => number = Impl.count as any;
-    export const getSegment: <T extends number = number>(segs: Segmentation<T>, value: T) => number = Impl.getSegment as any;
-    export const projectValue: <T extends number = number>(segs: Segmentation<T>, set: OrderedSet<T>, value: T) => Interval = Impl.projectValue as any;
+    export const count: <T extends number = number, I extends number = number>(segs: Segmentation<T, I>) => number = Impl.count as any;
+    export const getSegment: <T extends number = number, I extends number = number>(segs: Segmentation<T, I>, value: T) => number = Impl.getSegment as any;
+    export const projectValue: <T extends number = number, I extends number = number>(segs: Segmentation<T, I>, set: OrderedSet<T>, value: T) => Interval = Impl.projectValue as any;
 
     // Segment iterator that mutates a single segment object to mark all the segments.
-    export const transientSegments: <T extends number = number>(segs: Segmentation<T>, set: OrderedSet<T>, segment?: Segment<T>) => Impl.SegmentIterator<T> = Impl.segments as any;
+    export const transientSegments: <T extends number = number, I extends number = number>(segs: Segmentation<T, I>, set: OrderedSet<T>, segment?: Segment<T>) => Impl.SegmentIterator<T> = Impl.segments as any;
 }
 
-interface Segmentation<T extends number = number> {
+interface Segmentation<T extends number = number, I extends number = number> {
     '@type': 'segmentation',
     /** All segments are defined by offsets [offsets[i], offsets[i + 1]) for i \in [0, count - 1] */
     readonly offsets: ArrayLike<T>,
     /** Segment index of the i-th element */
-    readonly index: ArrayLike<number>,
+    readonly index: ArrayLike<I>,
     readonly count: number
 }
 

src/mol-model/structure/model.ts

@@ -10,4 +10,5 @@ import Format from './model/format'
 import { ModelSymmetry } from './model/properties/symmetry'
 import StructureSequence from './model/properties/sequence'
 
+export * from './model/indexing'
 export { Model, Types, Format, ModelSymmetry, StructureSequence }

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

@@ -10,6 +10,7 @@ import { SecondaryStructureType } from '../../types';
 import { AtomicHierarchy } from '../../properties/atomic';
 import { SecondaryStructure } from '../../properties/seconday-structure';
 import { Column } from 'mol-data/db';
+import { ChainIndex, ResidueIndex } from '../../indexing';
 
 export function getSecondaryStructureMmCif(data: mmCIF_Database, hierarchy: AtomicHierarchy): SecondaryStructure {
     const map: SecondaryStructureMap = new Map();
@@ -130,7 +131,7 @@ function addSheets(cat: mmCIF['struct_sheet_range'], map: SecondaryStructureMap,
     return;
 }
 
-function assignSecondaryStructureEntry(hierarchy: AtomicHierarchy, entry: SecondaryStructureEntry, resStart: number, resEnd: number, data: SecondaryStructureData) {
+function assignSecondaryStructureEntry(hierarchy: AtomicHierarchy, entry: SecondaryStructureEntry, resStart: ResidueIndex, resEnd: ResidueIndex, data: SecondaryStructureData) {
     const { label_seq_id, pdbx_PDB_ins_code } = hierarchy.residues;
     const { endSeqNumber, endInsCode, key, type } = entry;
 
@@ -154,7 +155,7 @@ function assignSecondaryStructureRanges(hierarchy: AtomicHierarchy, map: Seconda
     const { label_asym_id } = hierarchy.chains;
     const { label_seq_id, pdbx_PDB_ins_code } = hierarchy.residues;
 
-    for (let cI = 0; cI < chainCount; cI++) {
+    for (let cI = 0 as ChainIndex; cI < chainCount; cI++) {
         const resStart = AtomicHierarchy.chainStartResidueIndex(hierarchy, cI), resEnd = AtomicHierarchy.chainEndResidueIndexExcl(hierarchy, cI);
         const asymId = label_asym_id.value(cI);
         if (map.has(asymId)) {

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

@@ -0,0 +1,10 @@
+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+export type ElementIndex = { readonly '@type': 'element-index' } & number
+export type ResidueIndex = { readonly '@type': 'residue-index' } & number
+export type ChainIndex = { readonly '@type': 'chain-index' } & number
+export type EntityIndex = { readonly '@type': 'entity-index' } & number

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

@@ -71,10 +71,4 @@ export namespace Model {
             case 'mmCIF': return from_mmCIF(format);
         }
     }
-}
-
-
-export type ElementIndex = { readonly '@type': 'element-index' } & number
-export type ResidueIndex = { readonly '@type': 'residue-index' } & number
-export type ChainIndex = { readonly '@type': 'chain-index' } & number
-export type EntityIndex = { readonly '@type': 'entity-index' } & number
+}

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

@@ -9,6 +9,7 @@ import { Segmentation } from 'mol-data/int'
 import { mmCIF_Schema as mmCIF } from 'mol-io/reader/cif/schema/mmcif'
 import { ElementSymbol } from '../../types'
 import { Element } from '../../../structure'
+import { ChainIndex, EntityIndex, ResidueIndex } from '../../indexing';
 
 export const AtomsSchema = {
     type_symbol: Column.Schema.Aliased<ElementSymbol>(mmCIF.atom_site.type_symbol),
@@ -49,7 +50,7 @@ export interface AtomicData {
 
 export interface AtomicSegments {
     /** Maps residueIndex to a range of atoms [segments[rI], segments[rI + 1]) */
-    residueAtomSegments: Segmentation<Element>,
+    residueAtomSegments: Segmentation<Element, ResidueIndex>,
     /**
      * Maps chainIndex to a range of atoms [segments[cI], segments[cI + 1]),
      *
@@ -58,7 +59,7 @@ export interface AtomicSegments {
      * const start = rI[offsets[i]], const end = rI[offsets[i + 1] - 1] + 1;
      * for (let j = start; j < end; i++) { }
      */
-    chainAtomSegments: Segmentation<Element>,
+    chainAtomSegments: Segmentation<Element, ChainIndex>,
     /**
      * bonded/connected stretches of polymer chains, i.e. a chain will be
      * broken into multiple polymer segments if there are missing residues
@@ -69,27 +70,19 @@ export interface AtomicSegments {
 }
 
 export interface AtomicKeys {
-    // TODO: since Atoms must be sorted now, get rid of keys
     // TODO: include (lazily computed) "entity/chain/residue" indices?
 
-    // assign a key to each residue index.
-    residueKey: ArrayLike<number>,
-    // assign a key to each chain index
-    chainKey: ArrayLike<number>,
-    // assigne a key to each chain index
-    // also index to the Entities table.
-    entityKey: ArrayLike<number>,
+    /** @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): number,
+    /** @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): number
+    findResidueKey(entityId: string, label_asym_id: string, label_comp_id: string, auth_seq_id: number, pdbx_PDB_ins_code: string): ResidueIndex
 }
 
 type _Hierarchy = AtomicData & AtomicSegments & AtomicKeys
@@ -97,12 +90,12 @@ export interface AtomicHierarchy extends _Hierarchy { }
 
 export namespace AtomicHierarchy {
     /** Start residue inclusive */
-    export function chainStartResidueIndex(segs: AtomicSegments, cI: number) {
+    export function chainStartResidueIndex(segs: AtomicSegments, cI: ChainIndex) {
         return segs.residueAtomSegments.index[segs.chainAtomSegments.offsets[cI]];
     }
 
     /** End residue exclusive */
-    export function chainEndResidueIndexExcl(segs: AtomicSegments, cI: number) {
-        return segs.residueAtomSegments.index[segs.chainAtomSegments.offsets[cI + 1] - 1] + 1;
+    export function chainEndResidueIndexExcl(segs: AtomicSegments, cI: ChainIndex) {
+        return segs.residueAtomSegments.index[segs.chainAtomSegments.offsets[cI + 1] - 1] + 1 as ResidueIndex;
     }
 }

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

@@ -5,8 +5,9 @@
  */
 
 import { mmCIF_Database as mmCIF } from 'mol-io/reader/cif/schema/mmcif'
+import { EntityIndex } from '../indexing';
 
 export interface Entities {
     data: mmCIF['entity'],
-    getEntityIndex(id: string): import('../model').EntityIndex
+    getEntityIndex(id: string): EntityIndex
 }

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

@@ -8,6 +8,7 @@ import { Column } from 'mol-data/db'
 import { AtomicHierarchy } from './atomic/hierarchy';
 import { Entities } from './common';
 import { Sequence } from '../../../sequence';
+import { ChainIndex } from '../indexing';
 
 interface StructureSequence {
     readonly sequences: ReadonlyArray<StructureSequence.Entity>,
@@ -30,14 +31,14 @@ namespace StructureSequence {
         const byEntityKey: StructureSequence['byEntityKey'] = { };
         const sequences: StructureSequence.Entity[] = [];
 
-        for (let cI = 0, _cI = hierarchy.chains._rowCount; cI < _cI; cI++) {
-            const entityKey = hierarchy.entityKey[cI];
+        for (let cI = 0 as ChainIndex, _cI = hierarchy.chains._rowCount; cI < _cI; cI++) {
+            const entityKey = hierarchy.getEntityKey(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.entityKey[cI] && entities.data.type.value(entityKey) !== 'polymer') {
+            while (cI < _cI && entityKey === hierarchy.getEntityKey(cI) && entities.data.type.value(entityKey) !== 'polymer') {
                 cI++;
             }
             cI--;

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

@@ -7,6 +7,7 @@
 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) {
     return `${comp_id} ${seq_id} ${ins_code}`;
@@ -30,20 +31,20 @@ function createLookUp(entities: Entities, chain: Map<number, Map<string, number>
     const getEntKey = entities.getEntityIndex;
     const findChainKey: AtomicKeys['findChainKey'] = (e, c) => {
         let eKey = getEntKey(e);
-        if (eKey < 0) return -1;
+        if (eKey < 0) return -1 as ChainIndex;
         const cm = chain.get(eKey)!;
-        if (!cm.has(c)) return -1;
-        return cm.get(c)!;
+        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;
+        if (eKey < 0) return -1 as ResidueIndex;
         const cm = chain.get(eKey)!;
-        if (!cm.has(c)) return -1;
+        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;
-        return rm.get(id)!;
+        if (!rm.has(id)) return -1 as ResidueIndex;
+        return rm.get(id)! as ResidueIndex;
     }
     return { findChainKey, findResidueKey };
 }
@@ -92,5 +93,5 @@ export function getAtomicKeys(data: AtomicData, entities: Entities, segments: At
 
     const { findChainKey, findResidueKey } = createLookUp(entities, chainMaps, residueMaps);
 
-    return { residueKey, chainKey, entityKey, findChainKey, findResidueKey };
+    return { getEntityKey: cI => entityKey[cI] as EntityIndex, findChainKey, findResidueKey };
 }

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

@@ -48,7 +48,7 @@ const atom = {
 }
 
 const residue = {
-    key: Element.property(l => !Unit.isAtomic(l.unit) ? notAtomic() : l.unit.model.atomicHierarchy.residueKey[l.unit.residueIndex[l.element]]),
+    key: Element.property(l => !Unit.isAtomic(l.unit) ? notAtomic() : l.unit.residueIndex[l.element]),
 
     group_PDB: Element.property(l => !Unit.isAtomic(l.unit) ? notAtomic() : l.unit.model.atomicHierarchy.residues.group_PDB.value(l.unit.residueIndex[l.element])),
     label_comp_id: Element.property(l => !Unit.isAtomic(l.unit) ? notAtomic() : l.unit.model.atomicHierarchy.residues.label_comp_id.value(l.unit.residueIndex[l.element])),
@@ -63,7 +63,7 @@ const residue = {
 }
 
 const chain = {
-    key: Element.property(l => !Unit.isAtomic(l.unit) ? notAtomic() : l.unit.model.atomicHierarchy.chainKey[l.unit.chainIndex[l.element]]),
+    key: Element.property(l => !Unit.isAtomic(l.unit) ? notAtomic() : l.unit.chainIndex[l.element]),
 
     label_asym_id: Element.property(l => !Unit.isAtomic(l.unit) ? notAtomic() : l.unit.model.atomicHierarchy.chains.label_asym_id.value(l.unit.chainIndex[l.element])),
     auth_asym_id: Element.property(l => !Unit.isAtomic(l.unit) ? notAtomic() : l.unit.model.atomicHierarchy.chains.auth_asym_id.value(l.unit.chainIndex[l.element])),
@@ -92,7 +92,7 @@ const coarse = {
 function eK(l: Element.Location) {
     switch (l.unit.kind) {
         case Unit.Kind.Atomic:
-            return l.unit.model.atomicHierarchy.entityKey[l.unit.chainIndex[l.element]]
+            return l.unit.model.atomicHierarchy.getEntityKey(l.unit.chainIndex[l.element])
         case Unit.Kind.Spheres:
             return l.unit.model.coarseHierarchy.spheres.entityKey[l.element]
         case Unit.Kind.Gaussians:

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

@@ -13,6 +13,7 @@ import { CoarseElements, CoarseSphereConformation, CoarseGaussianConformation }
 import { ValueRef } from 'mol-util';
 import { UnitRings } from './unit/rings';
 import Element from './element'
+import { ChainIndex, ResidueIndex } from '../model/indexing';
 
 // A building block of a structure that corresponds to an atomic or a coarse grained representation
 // 'conveniently grouped together'.
@@ -82,8 +83,8 @@ namespace Unit {
         readonly conformation: SymmetryOperator.ArrayMapping;
 
         // Reference some commonly accessed things for faster access.
-        readonly residueIndex: ArrayLike<number>;
-        readonly chainIndex: ArrayLike<number>;
+        readonly residueIndex: ArrayLike<ResidueIndex>;
+        readonly chainIndex: ArrayLike<ChainIndex>;
 
         private props: AtomicProperties;