Merge pull request #523 from BradyAJohnston/fix-mda-load

Fix: MDA Loading only first frame
BradyAJohnston · Jun 11, 2024 · 9f6fb58 · 9f6fb58
2 parents f834c00 + da09aa0
commit 9f6fb58
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 58 deletions.
diff --git a/molecularnodes/io/md.py b/molecularnodes/io/md.py
@@ -82,7 +82,9 @@ def load(
     custom_selections: dict = {},
     in_memory: bool = False,
 ):
-    universe = mda.Universe(topology=path_resolve(top), coordinate=path_resolve(traj))
+    top = path_resolve(top)
+    traj = path_resolve(traj)
+    universe = mda.Universe(top, traj)
 
     if in_memory:
         universe.transfer_to_memory(start=start, step=step, stop=stop)

diff --git a/tests/__snapshots__/test_ops.ambr b/tests/__snapshots__/test_ops.ambr
@@ -2970,12 +2970,12 @@
 # ---
 # name: test_op_api_mda.14
   [[0.5 0.5 0.3]
-   [1.2 0.3 0.8]
-   [0.9 0.3 0.5]
+   [1.2 0.2 0.7]
+   [0.9 0.2 0.5]
    [0.6 0.3 0.5]
    [0.4 0.3 0.5]
    [0.7 0.5 0.6]
-   [1.3 0.4 0.7]
+   [1.3 0.3 0.7]
    [1.0 0.6 0.5]
    [1.0 0.4 0.4]
    [0.4 0.2 0.4]
@@ -2988,9 +2988,9 @@
    [1.0 0.3 0.5]
    [0.6 0.5 0.4]
    [1.1 0.6 0.7]
-   [1.0 0.4 0.4]
+   [1.0 0.5 0.3]
    [0.7 0.5 0.5]
-   [1.1 0.6 0.5]
+   [1.1 0.6 0.4]
    [1.2 0.5 0.6]
    [0.9 0.2 0.6]
    [1.3 0.5 0.7]
@@ -3000,14 +3000,14 @@
    [0.8 0.5 0.4]
    [0.6 0.3 0.3]
    [0.7 0.3 0.7]
-   [1.2 0.6 0.4]
-   [1.2 0.3 0.7]
+   [1.1 0.6 0.4]
+   [1.1 0.3 0.7]
    [0.7 0.4 0.8]
    [1.0 0.4 0.5]
    [0.5 0.4 0.5]
    [0.5 0.3 0.4]
-   [1.1 0.6 0.4]
-   [1.0 0.2 0.6]
+   [1.1 0.5 0.4]
+   [0.9 0.2 0.6]
    [1.1 0.5 0.7]
    [1.0 0.3 0.4]
    [1.2 0.5 0.4]
@@ -3016,58 +3016,58 @@
    [1.4 0.3 0.6]
    [0.7 0.4 0.7]
    [0.6 0.5 0.5]
-   [1.0 0.3 0.5]
+   [1.0 0.2 0.5]
    [0.9 0.3 0.3]
    [0.7 0.4 0.8]
    [1.1 0.5 0.5]
    [1.2 0.2 0.7]
-   [0.8 0.3 0.8]
+   [0.8 0.3 0.7]
    [0.6 0.6 0.6]
    [0.9 0.4 0.7]
    [0.6 0.2 0.5]
    [0.8 0.4 0.7]
    [0.6 0.5 0.7]
    [0.5 0.2 0.5]
    [1.1 0.4 0.4]
-   [1.0 0.3 0.3]
-   [1.0 0.5 0.6]
-   [1.1 0.6 0.5]
+   [1.0 0.3 0.4]
+   [1.0 0.5 0.5]
+   [1.1 0.5 0.5]
    [1.2 0.6 0.5]
-   [0.9 0.3 0.4]
+   [0.9 0.4 0.4]
    [0.8 0.6 0.5]
    [1.2 0.3 0.6]
    [0.4 0.4 0.3]
    [0.9 0.5 0.7]
    [1.2 0.5 0.7]
-   [1.1 0.6 0.6]
+   [1.1 0.7 0.6]
    [0.7 0.5 0.4]
-   [1.3 0.4 0.7]
-   [1.0 0.2 0.6]
+   [1.3 0.5 0.7]
+   [0.9 0.2 0.6]
    [0.6 0.5 0.4]
    [1.3 0.4 0.7]
-   [1.3 0.4 0.8]
-   [1.1 0.5 0.5]
+   [1.3 0.5 0.8]
+   [1.1 0.6 0.5]
    [0.6 0.5 0.4]
    [1.1 0.5 0.6]
-   [1.1 0.6 0.8]
-   [1.3 0.4 0.7]
+   [1.1 0.5 0.8]
+   [1.4 0.3 0.7]
    [1.0 0.3 0.4]
    [1.2 0.5 0.6]
-   [0.6 0.5 0.5]
+   [0.5 0.5 0.5]
    [0.8 0.4 0.8]
-   [1.0 0.3 0.5]
-   [1.2 0.4 0.8]
+   [0.9 0.3 0.5]
+   [1.2 0.4 0.7]
    [1.2 0.3 0.7]
    [0.7 0.4 0.7]
    [0.5 0.2 0.5]
    [0.4 0.2 0.4]
    [0.7 0.6 0.5]
    [0.9 0.2 0.6]
-   [1.1 0.5 0.4]
+   [1.1 0.6 0.4]
    [0.8 0.4 0.8]
    [0.8 0.2 0.6]
    [0.5 0.4 0.5]
-   [1.2 0.4 0.8]
+   [1.2 0.4 0.7]
    [1.1 0.4 0.4]]
 # ---
 # name: test_op_api_mda.15
@@ -3318,12 +3318,12 @@
 # ---
 # name: test_op_api_mda.38
   [[0.5 0.5 0.3]
-   [1.2 0.3 0.8]
-   [0.9 0.3 0.5]
+   [1.2 0.2 0.7]
+   [0.9 0.2 0.5]
    [0.6 0.3 0.5]
    [0.4 0.3 0.5]
    [0.7 0.5 0.6]
-   [1.3 0.4 0.7]
+   [1.3 0.3 0.7]
    [1.0 0.6 0.5]
    [1.0 0.4 0.4]
    [0.4 0.2 0.4]
@@ -3336,9 +3336,9 @@
    [1.0 0.3 0.5]
    [0.6 0.5 0.4]
    [1.1 0.6 0.7]
-   [1.0 0.4 0.4]
+   [1.0 0.5 0.3]
    [0.7 0.5 0.5]
-   [1.1 0.6 0.5]
+   [1.1 0.6 0.4]
    [1.2 0.5 0.6]
    [0.9 0.2 0.6]
    [1.3 0.5 0.7]
@@ -3348,14 +3348,14 @@
    [0.8 0.5 0.4]
    [0.6 0.3 0.3]
    [0.7 0.3 0.7]
-   [1.2 0.6 0.4]
-   [1.2 0.3 0.7]
+   [1.1 0.6 0.4]
+   [1.1 0.3 0.7]
    [0.7 0.4 0.8]
    [1.0 0.4 0.5]
    [0.5 0.4 0.5]
    [0.5 0.3 0.4]
-   [1.1 0.6 0.4]
-   [1.0 0.2 0.6]
+   [1.1 0.5 0.4]
+   [0.9 0.2 0.6]
    [1.1 0.5 0.7]
    [1.0 0.3 0.4]
    [1.2 0.5 0.4]
@@ -3364,58 +3364,58 @@
    [1.4 0.3 0.6]
    [0.7 0.4 0.7]
    [0.6 0.5 0.5]
-   [1.0 0.3 0.5]
+   [1.0 0.2 0.5]
    [0.9 0.3 0.3]
    [0.7 0.4 0.8]
    [1.1 0.5 0.5]
    [1.2 0.2 0.7]
-   [0.8 0.3 0.8]
+   [0.8 0.3 0.7]
    [0.6 0.6 0.6]
    [0.9 0.4 0.7]
    [0.6 0.2 0.5]
    [0.8 0.4 0.7]
    [0.6 0.5 0.7]
    [0.5 0.2 0.5]
    [1.1 0.4 0.4]
-   [1.0 0.3 0.3]
-   [1.0 0.5 0.6]
-   [1.1 0.6 0.5]
+   [1.0 0.3 0.4]
+   [1.0 0.5 0.5]
+   [1.1 0.5 0.5]
    [1.2 0.6 0.5]
-   [0.9 0.3 0.4]
+   [0.9 0.4 0.4]
    [0.8 0.6 0.5]
    [1.2 0.3 0.6]
    [0.4 0.4 0.3]
    [0.9 0.5 0.7]
    [1.2 0.5 0.7]
-   [1.1 0.6 0.6]
+   [1.1 0.7 0.6]
    [0.7 0.5 0.4]
-   [1.3 0.4 0.7]
-   [1.0 0.2 0.6]
+   [1.3 0.5 0.7]
+   [0.9 0.2 0.6]
    [0.6 0.5 0.4]
    [1.3 0.4 0.7]
-   [1.3 0.4 0.8]
-   [1.1 0.5 0.5]
+   [1.3 0.5 0.8]
+   [1.1 0.6 0.5]
    [0.6 0.5 0.4]
    [1.1 0.5 0.6]
-   [1.1 0.6 0.8]
-   [1.3 0.4 0.7]
+   [1.1 0.5 0.8]
+   [1.4 0.3 0.7]
    [1.0 0.3 0.4]
    [1.2 0.5 0.6]
-   [0.6 0.5 0.5]
+   [0.5 0.5 0.5]
    [0.8 0.4 0.8]
-   [1.0 0.3 0.5]
-   [1.2 0.4 0.8]
+   [0.9 0.3 0.5]
+   [1.2 0.4 0.7]
    [1.2 0.3 0.7]
    [0.7 0.4 0.7]
    [0.5 0.2 0.5]
    [0.4 0.2 0.4]
    [0.7 0.6 0.5]
    [0.9 0.2 0.6]
-   [1.1 0.5 0.4]
+   [1.1 0.6 0.4]
    [0.8 0.4 0.8]
    [0.8 0.2 0.6]
    [0.5 0.4 0.5]
-   [1.2 0.4 0.8]
+   [1.2 0.4 0.7]
    [1.1 0.4 0.4]]
 # ---
 # name: test_op_api_mda.39

diff --git a/tests/test_ops.py b/tests/test_ops.py
@@ -2,7 +2,7 @@
 import pytest
 import numpy as np
 import molecularnodes as mn
-from molecularnodes.blender.obj import ObjectTracker
+from molecularnodes.blender.obj import ObjectTracker, get_attribute
 
 from .utils import sample_attribute, NumpySnapshotExtension
 from .constants import data_dir, codes, attributes
@@ -90,10 +90,21 @@ def test_op_api_mda(snapshot_custom: NumpySnapshotExtension):
     name = "NewTrajectoryInMemory"
 
     obj_2, universe = mn.io.md.load(topo, traj, name="test", style="ribbon")
-    frames_coll = bpy.data.collections.get(f"{obj_2.name}_frames")
 
-    assert not frames_coll
+    # test the 'frames' collection doesn't exist, as it should only be created when reading
+    # into memory
+    assert not bpy.data.collections.get(f"{obj_2.name}_frames")
 
     for mol in [obj_1, obj_2]:
         for att in attributes:
             assert snapshot_custom == sample_attribute(mol, att)
+
+    # capture positions, change the frame number and test that the positions have updated
+    # and cahnged
+    pos_1, pos_2 = [get_attribute(x, "position") for x in [obj_1, obj_2]]
+    bpy.context.scene.frame_set(4)
+
+    assert not np.allclose(get_attribute(obj_1, "position"), pos_1)
+    assert np.allclose(
+        get_attribute(obj_1, "position"), get_attribute(obj_2, "position")
+    )