From 341dd33328ae709227268b81ae7c236e79d7874d Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Sat, 2 Nov 2024 00:52:16 +0800 Subject: [PATCH] cleanup collections code --- molecularnodes/blender/bpyd/collection.py | 23 ++++++ molecularnodes/blender/coll.py | 80 ++++++-------------- molecularnodes/blender/mesh.py | 7 +- molecularnodes/entities/molecule/molecule.py | 8 +- molecularnodes/entities/trajectory/dna.py | 2 +- tests/test_coll.py | 8 +- 6 files changed, 56 insertions(+), 72 deletions(-) create mode 100644 molecularnodes/blender/bpyd/collection.py diff --git a/molecularnodes/blender/bpyd/collection.py b/molecularnodes/blender/bpyd/collection.py new file mode 100644 index 00000000..3c3b14c3 --- /dev/null +++ b/molecularnodes/blender/bpyd/collection.py @@ -0,0 +1,23 @@ +import bpy +from bpy.types import Collection + + +def create_collection( + name: str = "NewCollection", parent: Collection | str | None = None +) -> Collection: + if isinstance(parent, str): + try: + parent = bpy.data.collections[name] + except KeyError: + parent = bpy.data.collections.new(name) + bpy.context.scene.collection.children.linke(parent) + try: + coll = bpy.data.collections[name] + except KeyError: + coll = bpy.data.collections.new(name) + if parent is None: + bpy.context.scene.collection.children.link(coll) + else: + parent.children.link(coll) + + return coll diff --git a/molecularnodes/blender/coll.py b/molecularnodes/blender/coll.py index f5359bba..048c3e46 100644 --- a/molecularnodes/blender/coll.py +++ b/molecularnodes/blender/coll.py @@ -1,70 +1,34 @@ import bpy +from bpy.types import Collection +from .bpyd.collection import create_collection -def mn(): - """Return the MolecularNodes Collection +def mn() -> Collection: + "Return the 'MolecularNodes' collection, creating it first if required" + return create_collection("MolecularNodes") - The collection called 'MolecularNodes' inside the Blender scene is returned. If the - collection does not exist first, it is created. - """ - coll = bpy.data.collections.get('MolecularNodes') - if not coll: - coll = bpy.data.collections.new('MolecularNodes') - bpy.context.scene.collection.children.link(coll) - return coll +def data() -> Collection: + "Return the MolecularNodes/data collection and disable it" + name = ".MN_data" -def armature(name='MN_armature'): - coll = bpy.data.collections.get(name) - if not coll: - coll = bpy.data.collections.new(name) - mn().children.link(coll) - return coll + try: + return bpy.data.collections[name] + except KeyError: + collection = create_collection(name=name, parent=mn()) + bpy.context.view_layer.layer_collection.children["MolecularNodes"].children[ + collection.name + ].exclude = True + return collection -def data(suffix=""): - """A collection for storing MN related data objects. - """ - name = f"MN_data{suffix}" - collection = bpy.data.collections.get(name) - if not collection: - collection = bpy.data.collections.new(name) - mn().children.link(collection) +def frames(name: str = "") -> Collection: + "Return a collection for storing the objects that are the frames of a trajectory" + return create_collection(f".data_{name}_frames", parent=data()) - # disable the view of the data collection - bpy.context.view_layer.layer_collection.children['MolecularNodes'].children[name].exclude = True - return collection - -def frames(name="", parent=None, suffix="_frames"): - """Create a Collection for Frames of a Trajectory - - Args: - name (str, optional): Name of the collection for the frames. Defaults to "". - parent (_type_, optional): A blender collection which will become the parent - collection. Defaults to the MolecularNodes collection if None. - """ - coll_frames = bpy.data.collections.new(name + suffix) - if not parent: - mn().children.link(coll_frames) - else: - parent.children.link(coll_frames) - - return coll_frames - - -def cellpack(name="", parent=None, fallback=False): +def cellpack(name: str = "") -> Collection: + "Return a collection for storing the instances for a CellPack Ensemble" full_name = f"cellpack_{name}" - coll = bpy.data.collections.get(full_name) - if coll and fallback: - return coll - - coll = bpy.data.collections.new(full_name) - - if parent: - parent.children.link(coll) - else: - data().children.link(coll) - - return coll + return create_collection(full_name, parent=data()) diff --git a/molecularnodes/blender/mesh.py b/molecularnodes/blender/mesh.py index 0157ae72..3060ee00 100644 --- a/molecularnodes/blender/mesh.py +++ b/molecularnodes/blender/mesh.py @@ -76,7 +76,12 @@ def evaluate_using_mesh(obj): return bob.evaluate().object -def create_data_object(array, collection=None, name="DataObject", world_scale=0.01): +def create_data_object( + array: np.ndarray, + name: str = "DataObject", + collection: str | bpy.types.Collection | None = None, + world_scale: float = 0.01, +) -> bpy.types.Object: # still requires a unique call TODO: figure out why # I think this has to do with the bcif instancing extraction array = np.unique(array) diff --git a/molecularnodes/entities/molecule/molecule.py b/molecularnodes/entities/molecule/molecule.py index 919255d0..8b045574 100644 --- a/molecularnodes/entities/molecule/molecule.py +++ b/molecularnodes/entities/molecule/molecule.py @@ -647,19 +647,13 @@ def att_sec_struct(): coll_frames = None if frames: - coll_frames = bl.coll.frames(bob.name, parent=bl.coll.data()) + coll_frames = bl.coll.frames(bob.name) for i, frame in enumerate(frames): frame = bl.mesh.create_object( name=bob.name + "_frame_" + str(i), collection=coll_frames, vertices=frame.coord * world_scale, - # vertices=frame.coord * world_scale - centroid ) - # TODO if update_attribute - # bl.mesh.store_named_attribute(attribute) - - # this has started to throw errors for me. I'm not sure why. - # mol.mn['molcule_type'] = 'pdb' # add custom properties to the actual blender object, such as number of chains, biological assemblies etc # currently biological assemblies can be problematic to holding off on doing that diff --git a/molecularnodes/entities/trajectory/dna.py b/molecularnodes/entities/trajectory/dna.py index 36f5f56b..4dde6182 100644 --- a/molecularnodes/entities/trajectory/dna.py +++ b/molecularnodes/entities/trajectory/dna.py @@ -288,7 +288,7 @@ def load(top, traj, name="oxDNA", setup_nodes=True, world_scale=0.01): # create a collection to store all of the frame objects that are part of the trajectory # they will contain all of the possible attributes which can be interpolated betewen # frames such as position, base_vector, base_normal, velocity, angular_velocity - collection = coll.frames(f"{name}_frames", parent=coll.data()) + collection = coll.frames(name) for i, frame in enumerate(trajectory): fill_n = int(np.ceil(np.log10(n_frames))) frame_name = f"{name}_frame_{str(i).zfill(fill_n)}" diff --git a/tests/test_coll.py b/tests/test_coll.py index 07abbacc..c01236ac 100644 --- a/tests/test_coll.py +++ b/tests/test_coll.py @@ -4,9 +4,7 @@ def test_coll(): assert mn.blender.coll.mn().name == "MolecularNodes" assert mn.blender.coll.mn().name == "MolecularNodes" - assert mn.blender.coll.data().name == "MN_data" + assert mn.blender.coll.data().name == ".MN_data" assert mn.blender.coll.cellpack().name == "cellpack_" - assert mn.blender.coll.cellpack(fallback=True).name == "cellpack_" - assert mn.blender.coll.cellpack().name == "cellpack_.001" - assert mn.blender.coll.frames().name == "_frames" - assert mn.blender.coll.frames("4OZS").name == "4OZS_frames" + assert mn.blender.coll.frames().name == ".data__frames" + assert mn.blender.coll.frames("4OZS").name == ".data_4OZS_frames"