add trr format support

CHANGELOG.md

@@ -13,6 +13,7 @@ Note that since we don't clearly distinguish between a public and private interf
 - Fix loading of some compressed files within sessions
 - Fix wrong element assignment for atoms with Charmm ion names
 - Fix handling of empty symmetry cell data
+- Add support for ``trr`` coordinates files
 
 ## [v3.3.1] - 2022-02-27
 

src/mol-io/reader/trr/parser.ts

@@ -0,0 +1,157 @@
+/**
+ * Copyright (c) 2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * Adapted from NGL.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { Task } from '../../../mol-task';
+import { ReaderResult as Result } from '../result';
+
+export interface TrrFile {
+    frames: { count: number, x: Float32Array, y: Float32Array, z: Float32Array }[],
+    boxes: number[][],
+    times: number[],
+    timeOffset: number,
+    deltaTime: number
+}
+
+async function parseInternal(data: Uint8Array) {
+    // https://github.com/gromacs/gromacs/blob/master/src/gromacs/fileio/trrio.cpp
+
+    const dv = new DataView(data.buffer);
+
+    const f: TrrFile = {
+        frames: [],
+        boxes: [],
+        times: [],
+        timeOffset: 0,
+        deltaTime: 0
+    };
+    const coordinates = f.frames;
+    const boxes = f.boxes;
+    const times = f.times;
+
+    let offset = 0;
+
+    while (true) {
+        // const magicnum = dv.getInt32(offset)
+        // const i1 = dv.getFloat32(offset + 4)
+        offset += 8;
+
+        const versionSize = dv.getInt32(offset);
+        offset += 4;
+        offset += versionSize;
+
+        // const irSize = dv.getInt32(offset)
+        // const eSize = dv.getInt32(offset + 4)
+        const boxSize = dv.getInt32(offset + 8);
+        const virSize = dv.getInt32(offset + 12);
+        const presSize = dv.getInt32(offset + 16);
+        // const topSize = dv.getInt32(offset + 20)
+        // const symSize = dv.getInt32(offset + 24)
+        const coordSize = dv.getInt32(offset + 28);
+        const velocitySize = dv.getInt32(offset + 32);
+        const forceSize = dv.getInt32(offset + 36);
+        const natoms = dv.getInt32(offset + 40);
+        // const step = dv.getInt32(offset + 44)
+        // const nre = dv.getInt32(offset + 48)
+        offset += 52;
+
+        const floatSize = boxSize / 9;
+        const natoms3 = natoms * 3;
+
+        // let lambda
+        if (floatSize === 8) {
+            times.push(dv.getFloat64(offset));
+            // lambda = dv.getFloat64(offset + 8)
+        } else {
+            times.push(dv.getFloat32(offset));
+            // lambda = dv.getFloat32(offset + 4)
+        }
+        offset += 2 * floatSize;
+
+        if (boxSize) {
+            const box = new Float32Array(9);
+            if (floatSize === 8) {
+                for (let i = 0; i < 9; ++i) {
+                    box[i] = dv.getFloat64(offset) * 10;
+                    offset += 8;
+                }
+            } else {
+                for (let i = 0; i < 9; ++i) {
+                    box[i] = dv.getFloat32(offset) * 10;
+                    offset += 4;
+                }
+            }
+            boxes.push(box as unknown as number[]);
+        }
+
+        // ignore, unused
+        offset += virSize;
+
+        // ignore, unused
+        offset += presSize;
+
+        if (coordSize) {
+            const x = new Float32Array(natoms);
+            const y = new Float32Array(natoms);
+            const z = new Float32Array(natoms);
+            if (floatSize === 8) {
+                for (let i = 0; i < natoms; ++i) {
+                    x[i] = dv.getFloat64(offset) * 10;
+                    y[i] = dv.getFloat64(offset + 8) * 10;
+                    z[i] = dv.getFloat64(offset + 16) * 10;
+                    offset += 24;
+                }
+            } else {
+                const tmp = new Uint32Array(data.buffer, offset, natoms3);
+                for (let i = 0; i < natoms3; ++i) {
+                    const value = tmp[i];
+                    tmp[i] = (
+                        ((value & 0xFF) << 24) | ((value & 0xFF00) << 8) |
+                        ((value >> 8) & 0xFF00) | ((value >> 24) & 0xFF)
+                    );
+                }
+                const frameCoords = new Float32Array(data.buffer, offset, natoms3);
+                for (let i = 0; i < natoms; ++i) {
+                    x[i] = frameCoords[i * 3] * 10;
+                    y[i] = frameCoords[i * 3 + 1] * 10;
+                    z[i] = frameCoords[i * 3 + 2] * 10;
+                    offset += 12;
+                }
+            }
+            coordinates.push({ count: natoms, x, y, z });
+        }
+
+        // ignore, unused
+        offset += velocitySize;
+
+        // ignore, unused
+        offset += forceSize;
+
+        if (offset >= data.byteLength) break;
+    }
+
+    if (times.length >= 1) {
+        f.timeOffset = times[0];
+    }
+    if (times.length >= 2) {
+        f.deltaTime = times[1] - times[0];
+    }
+
+    return f;
+}
+
+export function parseTrr(data: Uint8Array) {
+    return Task.create<Result<TrrFile>>('Parse TRR', async ctx => {
+        try {
+            ctx.update({ canAbort: true, message: 'Parsing trajectory...' });
+            const file = await parseInternal(data);
+            return Result.success(file);
+        } catch (e) {
+            return Result.error('' + e);
+        }
+    });
+}

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

@@ -0,0 +1,39 @@
+/**
+ * Copyright (c) 2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { Task } from '../../mol-task';
+import { TrrFile } from '../../mol-io/reader/trr/parser';
+import { Coordinates, Frame, Time } from '../../mol-model/structure/coordinates';
+import { Cell } from '../../mol-math/geometry/spacegroup/cell';
+import { Vec3 } from '../../mol-math/linear-algebra';
+
+export function coordinatesFromTrr(file: TrrFile): Task<Coordinates> {
+    return Task.create('Parse TRR', async ctx => {
+        await ctx.update('Converting to coordinates');
+
+        const deltaTime = Time(file.deltaTime, 'step');
+        const offsetTime = Time(file.timeOffset, deltaTime.unit);
+
+        const frames: Frame[] = [];
+        for (let i = 0, il = file.frames.length; i < il; ++i) {
+            const box = file.boxes[i];
+            const x = Vec3.fromArray(Vec3(), box, 0);
+            const y = Vec3.fromArray(Vec3(), box, 3);
+            const z = Vec3.fromArray(Vec3(), box, 6);
+            frames.push({
+                elementCount: file.frames[i].count,
+                cell: Cell.fromBasis(x, y, z),
+                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)
+            });
+        }
+
+        return Coordinates.create(frames, deltaTime, offsetTime);
+    });
+}

src/mol-plugin-state/formats/coordinates.ts

@@ -42,11 +42,28 @@ const XtcProvider = DataFormatProvider({
 });
 type XtcProvider = typeof XtcProvider;
 
-export type CoordinatesProvider = DcdProvider | XtcProvider;
+export { TrrProvider };
+const TrrProvider = DataFormatProvider({
+    label: 'TRR',
+    description: 'TRR',
+    category: CoordinatesFormatCategory,
+    binaryExtensions: ['trr'],
+    parse: (plugin, data) => {
+        const coordinates = plugin.state.data.build()
+            .to(data)
+            .apply(StateTransforms.Model.CoordinatesFromTrr);
+
+        return coordinates.commit();
+    }
+});
+type TrrProvider = typeof TrrProvider;
+
+export type CoordinatesProvider = DcdProvider | XtcProvider | TrrProvider;
 
 export const BuiltInCoordinatesFormats = [
     ['dcd', DcdProvider] as const,
     ['xtc', XtcProvider] as const,
+    ['trr', TrrProvider] as const,
 ] as const;
 
 export type BuiltInCoordinatesFormat = (typeof BuiltInCoordinatesFormats)[number][0]