PSF parser fix, correct topology/coordinates ordering (wip)

src/mol-data/db/column.ts

@@ -220,7 +220,13 @@ namespace Column {
         } else {
             for (let i = 0, _i = c.rowCount; i < _i; i++) array[offset + i] = c.value(i);
         }
+    }
 
+    export function isIdentity<T extends number>(c: Column<T>) {
+        for (let i = 0, _i = c.rowCount; i < _i; i++) {
+            if (i !== c.value(i)) return false;
+        }
+        return true;
     }
 }
 

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

@@ -8,19 +8,13 @@ import { createRangeArray, makeBuckets } from '../../../mol-data/util';
 import { Column, Table } from '../../../mol-data/db';
 import { RuntimeContext } from '../../../mol-task';
 import { AtomSite } from './schema';
+import { arrayIsIdentity } from '../../../mol-util/array';
 
 export type SortedAtomSite = {
     atom_site: AtomSite
     sourceIndex: Column<number>
 }
 
-function isIdentity(xs: ArrayLike<number>) {
-    for (let i = 0, _i = xs.length; i < _i; i++) {
-        if (xs[i] !== i) return false;
-    }
-    return true;
-}
-
 export async function sortAtomSite(ctx: RuntimeContext, atom_site: AtomSite, start: number, end: number): Promise<SortedAtomSite> {
     const indices = createRangeArray(start, end - 1);
 
@@ -40,7 +34,7 @@ export async function sortAtomSite(ctx: RuntimeContext, atom_site: AtomSite, sta
         if (ctx.shouldUpdate) await ctx.update();
     }
 
-    if (isIdentity(indices) && indices.length === atom_site._rowCount) {
+    if (arrayIsIdentity(indices) && indices.length === atom_site._rowCount) {
         return { atom_site, sourceIndex: Column.ofIntArray(indices) };
     }
 

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

@@ -4,17 +4,17 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { PsfFile } from '../../mol-io/reader/psf/parser';
 import { Column, Table } from '../../mol-data/db';
-import { EntityBuilder } from './common/entity';
+import { PsfFile } from '../../mol-io/reader/psf/parser';
-import { ComponentBuilder } from './common/component';
+import { getMoleculeType, MoleculeType } from '../../mol-model/structure/model/types';
-import { guessElementSymbolString } from './util';
+import { Topology } from '../../mol-model/structure/topology/topology';
-import { MoleculeType, getMoleculeType } from '../../mol-model/structure/model/types';
-import { getChainId } from './common/util';
 import { Task } from '../../mol-task';
 import { ModelFormat } from '../format';
-import { Topology } from '../../mol-model/structure/topology/topology';
+import { BasicSchema, createBasic } from './basic/schema';
-import { createBasic, BasicSchema } from './basic/schema';
+import { ComponentBuilder } from './common/component';
+import { EntityBuilder } from './common/entity';
+import { getChainId } from './common/util';
+import { guessElementSymbolString } from './util';
 
 function getBasic(atoms: PsfFile['atoms']) {
     const auth_atom_id = atoms.atomName;
@@ -32,16 +32,28 @@ function getBasic(atoms: PsfFile['atoms']) {
     let currentAsymIndex = 0;
     let currentAsymId = '';
     let currentSeqId = 0;
+    let currentSegmentName = atoms.segmentName.value(0), segmentChanged = false;
     let prevMoleculeType = MoleculeType.Unknown;
     let prevResidueNumber = -1;
 
     for (let i = 0, il = atoms.count; i < il; ++i) {
         const residueNumber = atoms.residueId.value(i);
-        if (residueNumber !== prevResidueNumber) {
+
+        if (currentSegmentName !== atoms.segmentName.value(i)) {
+            currentAsymId = getChainId(currentAsymIndex);
+            currentAsymIndex += 1;
+            currentSeqId = 0;
+            segmentChanged = true;
+            currentSegmentName = atoms.segmentName.value(i);
+        } else {
+            segmentChanged = false;
+        }
+
+        if (segmentChanged || residueNumber !== prevResidueNumber) {
             const compId = atoms.residueName.value(i);
             const moleculeType = getMoleculeType(componentBuilder.add(compId, i).type, compId);
 
-            if (moleculeType !== prevMoleculeType || residueNumber !== prevResidueNumber + 1) {
+            if (!segmentChanged && (moleculeType !== prevMoleculeType || residueNumber !== prevResidueNumber + 1)) {
                 currentAsymId = getChainId(currentAsymIndex);
                 currentAsymIndex += 1;
                 currentSeqId = 0;

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

@@ -101,4 +101,25 @@ namespace Coordinates {
             z: frame.z,
         };
     }
+
+    function reorderCoords(xs: ArrayLike<number>, index: ArrayLike<number>) {
+        const ret = new Float32Array(xs.length);
+        for (let i = 0, _i = xs.length; i < _i; i++) {
+            ret[i] = xs[index[i]];
+        }
+        return ret;
+    }
+
+    export function getAtomicConformationReordered(frame: Frame, atomId: Column<number>, srcIndex: ArrayLike<number>): AtomicConformation {
+        return {
+            id: UUID.create22(),
+            atomId,
+            occupancy: Column.ofConst(1, frame.elementCount, Column.Schema.int),
+            B_iso_or_equiv: Column.ofConst(0, frame.elementCount, Column.Schema.float),
+            xyzDefined: true,
+            x: reorderCoords(frame.x, srcIndex),
+            y: reorderCoords(frame.y, srcIndex),
+            z: reorderCoords(frame.z, srcIndex)
+        };
+    }
 }

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

@@ -88,16 +88,23 @@ export interface Model extends Readonly<{
 export namespace Model {
     export type Trajectory = ReadonlyArray<Model>
 
-    export function trajectoryFromModelAndCoordinates(model: Model, coordinates: Coordinates): Trajectory {
+    function _trajectoryFromModelAndCoordinates(model: Model, coordinates: Coordinates) {
         const trajectory: Mutable<Model.Trajectory> = [];
         const { frames } = coordinates;
+
+        const srcIndex = model.atomicHierarchy.atoms.sourceIndex;
+        const isIdentity = Column.isIdentity(srcIndex);
+        const srcIndexArray = isIdentity ? void 0 : srcIndex.toArray({ array: Int32Array });
+
         for (let i = 0, il = frames.length; i < il; ++i) {
             const f = frames[i];
             const m = {
                 ...model,
                 id: UUID.create22(),
                 modelNum: i,
-                atomicConformation: Coordinates.getAtomicConformation(f, model.atomicConformation.atomId),
+                atomicConformation: isIdentity
+                    ? Coordinates.getAtomicConformation(f, model.atomicConformation.atomId)
+                    : Coordinates.getAtomicConformationReordered(f, model.atomicConformation.atomId, srcIndexArray!),
                 // TODO: add support for supplying sphere and gaussian coordinates in addition to atomic coordinates?
                 // coarseConformation: coarse.conformation,
                 customProperties: new CustomProperties(),
@@ -106,15 +113,28 @@ export namespace Model {
             };
             trajectory.push(m);
         }
-        return trajectory;
+        return { trajectory, srcIndexArray };
+    }
+
+    export function trajectoryFromModelAndCoordinates(model: Model, coordinates: Coordinates): Trajectory {
+        return _trajectoryFromModelAndCoordinates(model, coordinates).trajectory;
     }
 
     export function trajectoryFromTopologyAndCoordinates(topology: Topology, coordinates: Coordinates): Task<Trajectory> {
         return Task.create('Create Trajectory', async ctx => {
             const model = (await createModels(topology.basic, topology.sourceData, ctx))[0];
             if (!model) throw new Error('found no model');
-            const trajectory = trajectoryFromModelAndCoordinates(model, coordinates);
+            const { trajectory, srcIndexArray } = _trajectoryFromModelAndCoordinates(model, coordinates);
-            const bondData = { pairs: topology.bonds, count: model.atomicHierarchy.atoms._rowCount };
+
+            const pairs = srcIndexArray
+                ? {
+                    indexA: Column.ofIntArray(Column.mapToArray(topology.bonds.indexA, i => srcIndexArray[i], Int32Array)),
+                    indexB: Column.ofIntArray(Column.mapToArray(topology.bonds.indexB, i => srcIndexArray[i], Int32Array)),
+                    order: topology.bonds.order
+                }
+                : topology.bonds;
+
+            const bondData = { pairs, count: model.atomicHierarchy.atoms._rowCount };
             const indexPairBonds = IndexPairBonds.fromData(bondData);
 
             let index = 0;

src/mol-util/array.ts

@@ -115,4 +115,11 @@ export function arrayEqual<T>(xs?: ArrayLike<T>, ys?: ArrayLike<T>) {
         if (xs[i] !== ys[i]) return false;
     }
     return true;
+}
+
+export function arrayIsIdentity(xs: ArrayLike<number>) {
+    for (let i = 0, _i = xs.length; i < _i; i++) {
+        if (xs[i] !== i) return false;
+    }
+    return true;
 }