From 1e4c2873a79ae2bc89eb576ee2ef8bbbc6fe6838 Mon Sep 17 00:00:00 2001 From: Brady Johnston <36021261+BradyAJohnston@users.noreply.github.com> Date: Sun, 22 Sep 2024 15:51:19 +0800 Subject: [PATCH] Add trajectory reloading (#616) * Operator to reload trajectory * fixed problem with pickling while saving --- molecularnodes/__init__.py | 4 +- molecularnodes/blender/__init__.py | 7 ++++ molecularnodes/entities/__init__.py | 3 +- molecularnodes/entities/entity.py | 15 ++++--- .../entities/trajectory/trajectory.py | 10 ++++- molecularnodes/entities/trajectory/ui.py | 40 ++++++++++++++----- molecularnodes/props.py | 12 ++++++ molecularnodes/session.py | 13 +++--- molecularnodes/ui/panel.py | 10 +++++ tests/test_ops.py | 2 +- 10 files changed, 89 insertions(+), 27 deletions(-) diff --git a/molecularnodes/__init__.py b/molecularnodes/__init__.py index 2dfb1d1a..a6558be7 100644 --- a/molecularnodes/__init__.py +++ b/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) diff --git a/molecularnodes/blender/__init__.py b/molecularnodes/blender/__init__.py index b69d77ea..94a1d5c5 100644 --- a/molecularnodes/blender/__init__.py +++ b/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 diff --git a/molecularnodes/entities/__init__.py b/molecularnodes/entities/__init__.py index 58128dfb..b2cec060 100644 --- a/molecularnodes/entities/__init__.py +++ b/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 diff --git a/molecularnodes/entities/entity.py b/molecularnodes/entities/entity.py index c91e4408..1e63d261 100644 --- a/molecularnodes/entities/entity.py +++ b/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: """ diff --git a/molecularnodes/entities/trajectory/trajectory.py b/molecularnodes/entities/trajectory/trajectory.py index 9eff302a..928b5b9d 100644 --- a/molecularnodes/entities/trajectory/trajectory.py +++ b/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}") diff --git a/molecularnodes/entities/trajectory/ui.py b/molecularnodes/entities/trajectory/ui.py index 6400d2f9..b719096e 100644 --- a/molecularnodes/entities/trajectory/ui.py +++ b/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] diff --git a/molecularnodes/props.py b/molecularnodes/props.py index 72d64e12..e22758db 100644 --- a/molecularnodes/props.py +++ b/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="", + ) diff --git a/molecularnodes/session.py b/molecularnodes/session.py index 8ef9d7ca..87415006 100644 --- a/molecularnodes/session.py +++ b/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 diff --git a/molecularnodes/ui/panel.py b/molecularnodes/ui/panel.py index d4819649..e7dc8c74 100644 --- a/molecularnodes/ui/panel.py +++ b/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() diff --git a/tests/test_ops.py b/tests/test_ops.py index 9e6002e8..100473fd 100644 --- a/tests/test_ops.py +++ b/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