mol-model: added sourceIndex to AtomicHierarchy

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

@@ -16,11 +16,8 @@ import { ElementSymbol } from 'mol-model/structure/model/types';
 import { Entities } from 'mol-model/structure/model/properties/common';
 import { getAtomicRanges } from 'mol-model/structure/model/properties/utils/atomic-ranges';
 import { getAtomicDerivedData } from 'mol-model/structure/model/properties/utils/atomic-derived';
-import { ModelFormat } from '../format';
-import mmCIF_Format = ModelFormat.mmCIF
 import { FormatData } from './parser';
 
-
 type AtomSite = mmCIF_Database['atom_site']
 
 function findHierarchyOffsets(atom_site: AtomSite) {
@@ -45,13 +42,14 @@ function findHierarchyOffsets(atom_site: AtomSite) {
     return { residues, chains };
 }
 
-function createHierarchyData(atom_site: AtomSite, offsets: { residues: ArrayLike<number>, chains: ArrayLike<number> }): AtomicData {
+function createHierarchyData(atom_site: AtomSite, sourceIndex: Column<number>, offsets: { residues: ArrayLike<number>, chains: ArrayLike<number> }): AtomicData {
     const atoms = Table.ofColumns(AtomsSchema, {
         type_symbol: Column.ofArray({ array: Column.mapToArray(atom_site.type_symbol, ElementSymbol), schema: Column.Schema.Aliased<ElementSymbol>(Column.Schema.str) }),
         label_atom_id: atom_site.label_atom_id,
         auth_atom_id: atom_site.auth_atom_id,
         label_alt_id: atom_site.label_alt_id,
-        pdbx_formal_charge: atom_site.pdbx_formal_charge
+        pdbx_formal_charge: atom_site.pdbx_formal_charge,
+        sourceIndex
     });
     const residues = Table.view(atom_site, ResiduesSchema, offsets.residues);
     // Optimize the numeric columns
@@ -80,9 +78,9 @@ function isHierarchyDataEqual(a: AtomicData, b: AtomicData) {
         && Table.areEqual(a.atoms as Table<AtomsSchema>, b.atoms as Table<AtomsSchema>)
 }
 
-export function getAtomicHierarchyAndConformation(format: mmCIF_Format, atom_site: AtomSite, entities: Entities, formatData: FormatData, previous?: Model) {
+export function getAtomicHierarchyAndConformation(atom_site: AtomSite, sourceIndex: Column<number>, entities: Entities, formatData: FormatData, previous?: Model) {
     const hierarchyOffsets = findHierarchyOffsets(atom_site);
-    const hierarchyData = createHierarchyData(atom_site, hierarchyOffsets);
+    const hierarchyData = createHierarchyData(atom_site, sourceIndex, hierarchyOffsets);
 
     if (previous && isHierarchyDataEqual(previous.atomicHierarchy, hierarchyData)) {
         return {

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

@@ -21,6 +21,7 @@ export interface IHMData {
     model_name: string,
     entities: Entities,
     atom_site: mmCIF['atom_site'],
+    atom_site_sourceIndex: Column<number>,
     ihm_sphere_obj_site: mmCIF['ihm_sphere_obj_site'],
     ihm_gaussian_obj_site: mmCIF['ihm_gaussian_obj_site']
 }

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

@@ -169,8 +169,8 @@ function getFormatData(format: mmCIF_Format): FormatData {
     }
 }
 
-function createStandardModel(format: mmCIF_Format, atom_site: AtomSite, entities: Entities, formatData: FormatData, previous?: Model): Model {
-    const atomic = getAtomicHierarchyAndConformation(format, atom_site, entities, formatData, previous);
+function createStandardModel(format: mmCIF_Format, atom_site: AtomSite, sourceIndex: Column<number>, entities: Entities, formatData: FormatData, previous?: Model): Model {
+    const atomic = getAtomicHierarchyAndConformation(atom_site, sourceIndex, entities, formatData, previous);
     if (previous && atomic.sameAsPrevious) {
         return {
             ...previous,
@@ -209,7 +209,7 @@ function createStandardModel(format: mmCIF_Format, atom_site: AtomSite, entities
 }
 
 function createModelIHM(format: mmCIF_Format, data: IHMData, formatData: FormatData): Model {
-    const atomic = getAtomicHierarchyAndConformation(format, data.atom_site, data.entities, formatData);
+    const atomic = getAtomicHierarchyAndConformation(data.atom_site, data.atom_site_sourceIndex, data.entities, formatData);
     const coarse = getIHMCoarse(data, formatData);
 
     return {
@@ -255,8 +255,8 @@ async function readStandard(ctx: RuntimeContext, format: mmCIF_Format, formatDat
     let modelStart = 0;
     while (modelStart < atomCount) {
         const modelEnd = findModelEnd(format.data.atom_site.pdbx_PDB_model_num, modelStart);
-        const atom_site = await sortAtomSite(ctx, format.data.atom_site, modelStart, modelEnd);
-        const model = createStandardModel(format, atom_site, entities, formatData, models.length > 0 ? models[models.length - 1] : void 0);
+        const { atom_site, sourceIndex } = await sortAtomSite(ctx, format.data.atom_site, modelStart, modelEnd);
+        const model = createStandardModel(format, atom_site, sourceIndex, entities, formatData, models.length > 0 ? models[models.length - 1] : void 0);
         attachProps(model);
         models.push(model);
         modelStart = modelEnd;
@@ -265,14 +265,17 @@ async function readStandard(ctx: RuntimeContext, format: mmCIF_Format, formatDat
 }
 
 function splitTable<T extends Table<any>>(table: T, col: Column<number>) {
-    const ret = new Map<number, T>()
+    const ret = new Map<number, { table: T, start: number, end: number }>()
     const rowCount = table._rowCount;
     let modelStart = 0;
     while (modelStart < rowCount) {
         const modelEnd = findModelEnd(col, modelStart);
         const id = col.value(modelStart);
-        const window = Table.window(table, table._schema, modelStart, modelEnd) as T;
-        ret.set(id, window);
+        ret.set(id, {
+            table: Table.window(table, table._schema, modelStart, modelEnd) as T,
+            start: modelStart,
+            end: modelEnd
+        });
         modelStart = modelEnd;
     }
     return ret;
@@ -286,8 +289,9 @@ async function readIHM(ctx: RuntimeContext, format: mmCIF_Format, formatData: Fo
         throw new Error('expected _atom_site.ihm_model_id to be defined')
     }
 
-    // TODO: will IHM require sorting or will we trust it?
     const atom_sites = splitTable(format.data.atom_site, format.data.atom_site.ihm_model_id);
+    // TODO: will coarse IHM records require sorting or will we trust it?
+    // ==> Probably implement a sort as as well and store the sourceIndex same as with atomSite
     const sphere_sites = splitTable(format.data.ihm_sphere_obj_site, format.data.ihm_sphere_obj_site.model_id);
     const gauss_sites = splitTable(format.data.ihm_gaussian_obj_site, format.data.ihm_gaussian_obj_site.model_id);
 
@@ -296,13 +300,26 @@ async function readIHM(ctx: RuntimeContext, format: mmCIF_Format, formatData: Fo
     const { model_id, model_name } = ihm_model_list;
     for (let i = 0; i < ihm_model_list._rowCount; i++) {
         const id = model_id.value(i);
+
+        let atom_site, atom_site_sourceIndex;
+        if (atom_sites.has(id)) {
+            const e = atom_sites.get(id)!;
+            const { atom_site: sorted, sourceIndex } = await sortAtomSite(ctx, e.table, e.start, e.end);
+            atom_site = sorted;
+            atom_site_sourceIndex = sourceIndex;
+        } else {
+            atom_site = Table.window(format.data.atom_site, format.data.atom_site._schema, 0, 0);
+            atom_site_sourceIndex = Column.ofIntArray([]);
+        }
+
         const data: IHMData = {
             model_id: id,
             model_name: model_name.value(i),
             entities: entities,
-            atom_site: atom_sites.has(id) ? atom_sites.get(id)! : Table.window(format.data.atom_site, format.data.atom_site._schema, 0, 0),
-            ihm_sphere_obj_site: sphere_sites.has(id) ? sphere_sites.get(id)! : Table.window(format.data.ihm_sphere_obj_site, format.data.ihm_sphere_obj_site._schema, 0, 0),
-            ihm_gaussian_obj_site: gauss_sites.has(id) ? gauss_sites.get(id)! : Table.window(format.data.ihm_gaussian_obj_site, format.data.ihm_gaussian_obj_site._schema, 0, 0)
+            atom_site,
+            atom_site_sourceIndex,
+            ihm_sphere_obj_site: sphere_sites.has(id) ? sphere_sites.get(id)!.table : Table.window(format.data.ihm_sphere_obj_site, format.data.ihm_sphere_obj_site._schema, 0, 0),
+            ihm_gaussian_obj_site: gauss_sites.has(id) ? gauss_sites.get(id)!.table : Table.window(format.data.ihm_gaussian_obj_site, format.data.ihm_gaussian_obj_site._schema, 0, 0)
         };
         const model = createModelIHM(format, data, formatData);
         attachProps(model);

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

@@ -9,6 +9,8 @@ import { createRangeArray, makeBuckets } from 'mol-data/util';
 import { Column, Table } from 'mol-data/db';
 import { RuntimeContext } from 'mol-task';
 
+export type SortedAtomSite = mmCIF_Database['atom_site'] & { sourceIndex: Column<number> }
+
 function isIdentity(xs: ArrayLike<number>) {
     for (let i = 0, _i = xs.length; i < _i; i++) {
         if (xs[i] !== i) return false;
@@ -36,8 +38,11 @@ export async function sortAtomSite(ctx: RuntimeContext, atom_site: mmCIF_Databas
     }
 
     if (isIdentity(indices) && indices.length === atom_site._rowCount) {
-        return atom_site;
+        return { atom_site, sourceIndex: Column.ofIntArray(indices) };
     }
 
-    return Table.view(atom_site, atom_site._schema, indices) as mmCIF_Database['atom_site'];
+    return {
+        atom_site: Table.view(atom_site, atom_site._schema, indices) as mmCIF_Database['atom_site'],
+        sourceIndex: Column.ofIntArray(indices)
+    };
 }
\ No newline at end of file

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

@@ -38,7 +38,14 @@ export const AtomsSchema = {
      * The net integer charge assigned to this atom.
      * This is the formal charge assignment normally found in chemical diagrams.
      */
-    pdbx_formal_charge: mmCIF.atom_site.pdbx_formal_charge
+    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
 };