In-place reordering support for Frame.x/y/z

CHANGELOG.md

@@ -6,6 +6,7 @@ Note that since we don't clearly distinguish between a public and private interf
 ## [Unreleased]
 
 - Add ability to specify ``volumeIndex`` in ``Viewer.loadVolumeFromUrl`` to better support Volume Server inputs.
+- Support in-place reordering for trajectory ``Frame.x/y/z`` arrays for better memory efficiency.
 
 ## [v2.0.6] - 2021-06-01
 

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

@@ -38,6 +38,8 @@ export function coordinatesFromDcd(dcdFile: DcdFile): Task<Coordinates> {
                 x: dcdFrame.x,
                 y: dcdFrame.y,
                 z: dcdFrame.z,
+
+                xyzOrdering: { isIdentity: true }
             };
 
             if (dcdFrame.cell) {

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

@@ -30,6 +30,7 @@ export function coordinatesFromXtc(file: XtcFile): Task<Coordinates> {
                 x: file.frames[i].x,
                 y: file.frames[i].y,
                 z: file.frames[i].z,
+                xyzOrdering: { isIdentity: true },
                 time: Time(offsetTime.value + deltaTime.value * i, deltaTime.unit)
             });
         }

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

@@ -1,10 +1,11 @@
 /**
- * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2019-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { UUID } from '../../../mol-util';
+import { arrayEqual, UUID } from '../../../mol-util';
 import { Cell } from '../../../mol-math/geometry/spacegroup/cell';
 import { AtomicConformation } from '../model/properties/atomic';
 import { Column } from '../../../mol-data/db';
@@ -34,6 +35,12 @@ export interface Frame {
         readonly fy: ArrayLike<number>
         readonly fz: ArrayLike<number>
     }
+
+    readonly xyzOrdering: {
+        isIdentity: boolean,
+        frozen?: boolean,
+        index?: ArrayLike<number>,
+    }
 }
 
 //
@@ -85,41 +92,106 @@ namespace Coordinates {
             hasVelocities,
             hasForces,
             deltaTime,
-            timeOffset
+            timeOffset,
         };
     }
 
-    export function getAtomicConformation(frame: Frame, atomId: Column<number>): AtomicConformation {
+    /**
+     * Only use ordering if it's not identity.
+     */
+    export function getAtomicConformation(frame: Frame, atomId: Column<number>, ordering?: ArrayLike<number>): AtomicConformation {
+        let { x, y, z } = frame;
+
+        if (frame.xyzOrdering.frozen) {
+            if (ordering) {
+                if (frame.xyzOrdering.isIdentity) {
+                    // simple list reordering
+                    x = getOrderedCoords(x, ordering);
+                    y = getOrderedCoords(y, ordering);
+                    z = getOrderedCoords(z, ordering);
+                } else if (!arrayEqual(frame.xyzOrdering.index! as any, ordering as any)) {
+                    x = getSourceOrderedCoords(x, frame.xyzOrdering.index!, ordering);
+                    y = getSourceOrderedCoords(y, frame.xyzOrdering.index!, ordering);
+                    z = getSourceOrderedCoords(z, frame.xyzOrdering.index!, ordering);
+                }
+            } else if (!frame.xyzOrdering.isIdentity) {
+                x = getInvertedCoords(x, frame.xyzOrdering.index!);
+                y = getInvertedCoords(y, frame.xyzOrdering.index!);
+                z = getInvertedCoords(z, frame.xyzOrdering.index!);
+            }
+        } else if (ordering) {
+            if (frame.xyzOrdering.isIdentity) {
+                frame.xyzOrdering.isIdentity = false;
+                frame.xyzOrdering.index = ordering;
+                reorderCoordsInPlace(x as unknown as number[], ordering);
+                reorderCoordsInPlace(y as unknown as number[], ordering);
+                reorderCoordsInPlace(z as unknown as number[], ordering);
+            } else {
+                // is current ordering is not the same as requested?
+                //   => copy the conformations into a new array
+                if (!arrayEqual(frame.xyzOrdering.index! as any, ordering as any)) {
+                    x = getSourceOrderedCoords(x, frame.xyzOrdering.index!, ordering);
+                    y = getSourceOrderedCoords(y, frame.xyzOrdering.index!, ordering);
+                    z = getSourceOrderedCoords(z, frame.xyzOrdering.index!, ordering);
+                }
+            }
+        }
+
+        // once the conformation has been accessed at least once, freeze it.
+        //   => any other request to the frame with different ordering will result in a copy.
+        frame.xyzOrdering.frozen = true;
+
         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: frame.x,
-            y: frame.y,
-            z: frame.z,
+            x,
+            y,
+            z,
         };
     }
 
-    function reorderCoords(xs: ArrayLike<number>, index: ArrayLike<number>) {
+    const _reorderBuffer = [0.123];
+    function reorderCoordsInPlace(xs: number[], index: ArrayLike<number>) {
+        const buffer = _reorderBuffer;
+
+        for (let i = 0, _i = xs.length; i < _i; i++) {
+            buffer[i] = xs[index[i]];
+        }
+        for (let i = 0, _i = xs.length; i < _i; i++) {
+            xs[i] = buffer[i];
+        }
+    }
+
+    function getSourceOrderedCoords(xs: ArrayLike<number>, srcIndex: ArrayLike<number>, index: ArrayLike<number>) {
         const ret = new Float32Array(xs.length);
+
+        for (let i = 0, _i = xs.length; i < _i; i++) {
+            ret[i] = xs[srcIndex[index[i]]];
+        }
+
+        return ret;
+    }
+
+    function getOrderedCoords(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)
-        };
+    function getInvertedCoords(xs: ArrayLike<number>, index: ArrayLike<number>) {
+        const ret = new Float32Array(xs.length);
+
+        for (let i = 0, _i = xs.length; i < _i; i++) {
+            ret[index[i]] = xs[i];
+        }
+
+        return ret;
     }
 }

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

@@ -105,9 +105,7 @@ export namespace Model {
                 ...model,
                 id: UUID.create22(),
                 modelNum: i,
-                atomicConformation: isIdentity
-                    ? Coordinates.getAtomicConformation(f, model.atomicConformation.atomId)
-                    : Coordinates.getAtomicConformationReordered(f, model.atomicConformation.atomId, srcIndexArray!),
+                atomicConformation: Coordinates.getAtomicConformation(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(),