Add trajectory reloading (#616)

* Operator to reload trajectory

* fixed problem with pickling while saving

molecularnodes/__init__.py

@@ -45,8 +45,8 @@ def register():
     for op in all_classes:
         try:
             bpy.utils.register_class(op)
-        except Exception:
-            # print(e)
+        except Exception as e:
+            print(e)
             pass
 
     bpy.types.NODE_MT_add.append(MN_add_node_menu)

molecularnodes/blender/__init__.py

@@ -10,3 +10,10 @@ def path_resolve(path: Union[str, Path]) -> Path:
         return Path(bpy.path.abspath(str(path)))
     else:
         raise ValueError(f"Unable to resolve path: {path}")
+
+
+def active_object(context: bpy.types.Context = None) -> bpy.types.Object:
+    if context is None:
+        return bpy.context.active_object
+
+    return context.active_object

molecularnodes/entities/__init__.py

@@ -7,7 +7,7 @@
 from .molecule.pdb import PDB
 from .molecule.pdbx import BCIF, CIF
 from .molecule.sdf import SDF
-from .molecule.ui import MN_OT_Import_wwPDB, fetch, load_local
+from .molecule.ui import fetch, load_local
 from .trajectory.trajectory import Trajectory
 
 CLASSES = (
@@ -16,7 +16,6 @@
         MN_OT_Import_Map,
         MN_OT_Import_OxDNA_Trajectory,
         MN_OT_Import_Star_File,
-        MN_OT_Import_wwPDB,
     ]
     + trajectory.CLASSES
     + molecule.CLASSES

molecularnodes/entities/entity.py

@@ -15,7 +15,7 @@ def __init__(self, message):
 class MolecularEntity(metaclass=ABCMeta):
     def __init__(self) -> None:
         self.uuid: str = str(uuid1())
-        self.object_ref: bpy.types.Object | None
+        self._object: bpy.types.Object | None
         self.type: str = ""
 
     @property
@@ -45,24 +45,27 @@ def object(self) -> bpy.types.Object | None:
             # if the connection is broken then trying to the name will raise a connection
             # error. If we are loading from a saved session then the object_ref will be
             # None and get an AttributeError
-            self.object_ref.name
-            return self.object_ref
+            self._object.name
+            return self._object
         except (ReferenceError, AttributeError):
             for obj in bpy.data.objects:
                 if obj.mn.uuid == self.uuid:
                     print(
                         Warning(
-                            f"Lost connection to object: {self.object_ref}, now connected to {obj}"
+                            f"Lost connection to object: {self._object}, now connected to {obj}"
                         )
                     )
-                    self.object_ref = obj
+                    self._object = obj
                     return obj
 
             return None
 
     @object.setter
     def object(self, value):
-        self.object_ref = value
+        if isinstance(value, bpy.types.Object) or value is None:
+            self._object = value
+        else:
+            raise TypeError(f"The `object` must be a Blender object, not {value=}")
 
     def named_attribute(self, name="position", evaluate=False) -> np.ndarray | None:
         """

molecularnodes/entities/trajectory/trajectory.py

@@ -7,7 +7,7 @@
 
 from ... import data
 from ..entity import MolecularEntity, ObjectMissingError
-from ...blender import coll, mesh, nodes
+from ...blender import coll, mesh, nodes, path_resolve
 from ...utils import lerp, correct_periodic_positions
 from .selections import Selection, TrajectorySelectionItem
 
@@ -416,6 +416,13 @@ def _attributes_2_blender(self):
             },
         }
 
+    def save_filepaths_on_object(self) -> None:
+        obj = self.object
+        obj.mn.filepath_topology = str(path_resolve(self.universe.filename))
+        obj.mn.filepath_trajectory = str(
+            path_resolve(self.universe.trajectory.filename)
+        )
+
     def create_object(
         self,
         style: str = "vdw",
@@ -447,6 +454,7 @@ def create_object(
         obj["atom_type_unique"] = self.atom_type_unique
         self.subframes = subframes
         obj.mn.molecule_type = "md"
+        self.save_filepaths_on_object()
 
         if style is not None:
             nodes.create_starting_node_tree(obj, style=style, name=f"MN_{obj.name}")

molecularnodes/entities/trajectory/ui.py

@@ -8,7 +8,8 @@
 import bpy
 import MDAnalysis as mda
 
-from ...blender import path_resolve
+from ... import blender as bl
+from ...session import get_session
 from .trajectory import Trajectory
 from bpy.props import StringProperty
 
@@ -39,8 +40,8 @@ def load(
     style="spheres",
     subframes: int = 0,
 ):
-    top = path_resolve(top)
-    traj = path_resolve(traj)
+    top = bl.path_resolve(top)
+    traj = bl.path_resolve(traj)
 
     universe = mda.Universe(top, traj)
 
@@ -51,8 +52,31 @@ def load(
     return traj
 
 
-class MN_OT_Import_Protein_MD(bpy.types.Operator):
-    bl_idname = "mn.import_protein_md"
+class MN_OT_Reload_Trajectory(bpy.types.Operator):
+    bl_idname = "mn.reload_trajectory"
+    bl_label = "Reload Trajectory"
+    bl_description = (
+        "Reload the `mda.UNiverse` of the current Object to renable updating"
+    )
+    bl_options = {"REGISTER", "UNDO"}
+
+    @classmethod
+    def poll(cls, context):
+        obj = bl.active_object(context)
+        traj = get_session(context).trajectories.get(obj.mn.uuid)
+        return not traj
+
+    def execute(self, context):
+        obj = bl.active_object(context)
+        universe = mda.Universe(obj.mn.filepath_topology, obj.mn.filepath_trajectory)
+        traj = Trajectory(universe)
+        traj.object = obj
+        obj.mn.uuid = traj.uuid
+        return {"FINISHED"}
+
+
+class MN_OT_Import_Trajectory(bpy.types.Operator):
+    bl_idname = "mn.import_trajectory"
     bl_label = "Import Protein MD"
     bl_description = "Load molecular dynamics trajectory"
     bl_options = {"REGISTER", "UNDO"}
@@ -93,7 +117,7 @@ def panel(layout, scene):
     col = layout.column(align=True)
     row_import = col.row()
     row_import.prop(scene, "MN_import_md_name")
-    row_import.operator("mn.import_protein_md", text="Load")
+    row_import.operator("mn.import_trajectory", text="Load")
     col.separator()
     col.prop(scene, "MN_import_md_topology")
     col.prop(scene, "MN_import_md_trajectory")
@@ -107,6 +131,4 @@ def panel(layout, scene):
     col.enabled = scene.MN_import_node_setup
 
 
-CLASSES = [
-    MN_OT_Import_Protein_MD,
-]
+CLASSES = [MN_OT_Import_Trajectory, MN_OT_Reload_Trajectory]

molecularnodes/props.py

@@ -94,3 +94,15 @@ class MolecularNodesObjectProperties(bpy.types.PropertyGroup):
         default=True,
         update=_update_trajectories,
     )
+    filepath_trajectory: StringProperty(  # type: ignore
+        name="Trajectory",
+        description="Filepath for the `trajectory` of the Object",
+        subtype="FILE_PATH",
+        default="",
+    )
+    filepath_topology: StringProperty(  # type: ignore
+        name="Topology",
+        description="Filepath for the Topology of the Object",
+        subtype="FILE_PATH",
+        default="",
+    )

molecularnodes/session.py

@@ -21,9 +21,7 @@ def trim(dictionary: dict):
         if hasattr(item, "calculations"):
             item.calculations = {}
         try:
-            if isinstance(item.object, bpy.types.Object):
-                item.name = item.object.name
-                item.object = None
+            item.object = None
             if hasattr(item, "frames"):
                 if isinstance(item.frames, bpy.types.Collection):
                     item.frames_name = item.frames.name
@@ -45,6 +43,7 @@ def trim(dictionary: dict):
 def make_paths_relative(trajectories: Dict[str, Trajectory]) -> None:
     for key, traj in trajectories.items():
         traj.universe.load_new(make_path_relative(traj.universe.trajectory.filename))
+        traj.save_filepaths_on_object()
 
 
 def trim_root_folder(filename):
@@ -54,7 +53,10 @@ def trim_root_folder(filename):
 
 def make_path_relative(filepath):
     "Take a path and make it relative, in an actually usable way"
-    filepath = os.path.relpath(filepath)
+    try:
+        filepath = os.path.relpath(filepath)
+    except ValueError:
+        return filepath
 
     # count the number of "../../../" there are to remove
     n_to_remove = int(filepath.count("..") - 2)
@@ -127,12 +129,11 @@ def __repr__(self) -> str:
     def pickle(self, filepath) -> None:
         pickle_path = self.stashpath(filepath)
 
+        make_paths_relative(self.trajectories)
         self.molecules = trim(self.molecules)
         self.trajectories = trim(self.trajectories)
         self.ensembles = trim(self.ensembles)
 
-        make_paths_relative(self.trajectories)
-
         # don't save anything if there is nothing to save
         if self.n_items == 0:
             return None

molecularnodes/ui/panel.py

@@ -137,6 +137,16 @@ def panel_md_properties(layout, context):
     obj = context.active_object
     session = get_session()
     universe = session.trajectories.get(obj.mn.uuid)
+    trajectory_is_linked = bool(universe)
+    col = layout.column()
+    col.enabled = False
+    if not trajectory_is_linked:
+        col.enabled = True
+        col.label(text="Object not linked to a trajectory, please reload one")
+        col.prop(obj.mn, "filepath_topology")
+        col.prop(obj.mn, "filepath_trajectory")
+        col.operator("mn.reload_trajectory")
+        return None
 
     layout.label(text="Trajectory Playback", icon="OPTIONS")
     row = layout.row()

tests/test_ops.py

@@ -84,7 +84,7 @@ def test_op_api_mda(snapshot_custom: NumpySnapshotExtension):
     bpy.context.scene.MN_import_style = "ribbon"
 
     with ObjectTracker() as o:
-        bpy.ops.mn.import_protein_md()
+        bpy.ops.mn.import_trajectory()
         obj_1 = o.latest()
 
     assert obj_1.name == name