From 0ebd08e65e899a3e0b7867c535b6213e54f13602 Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Wed, 30 Oct 2024 16:12:16 +0800 Subject: [PATCH 01/12] initial refactoring --- molecularnodes/blender/bones.py | 115 --------- molecularnodes/blender/databpy/__init__.py | 0 molecularnodes/blender/mesh.py | 218 +++++++++++------- molecularnodes/entities/ensemble/cellpack.py | 5 +- molecularnodes/entities/ensemble/star.py | 18 +- molecularnodes/entities/entity.py | 4 +- molecularnodes/entities/molecule/molecule.py | 2 +- molecularnodes/entities/trajectory/dna.py | 16 +- .../entities/trajectory/trajectory.py | 8 +- tests/test_obj.py | 6 +- tests/test_select.py | 6 +- 11 files changed, 178 insertions(+), 220 deletions(-) delete mode 100644 molecularnodes/blender/bones.py create mode 100644 molecularnodes/blender/databpy/__init__.py diff --git a/molecularnodes/blender/bones.py b/molecularnodes/blender/bones.py deleted file mode 100644 index d52148ea..00000000 --- a/molecularnodes/blender/bones.py +++ /dev/null @@ -1,115 +0,0 @@ -import bpy -import numpy as np -from . import mesh, coll - - -def clear_armature(object): - for mod in object.modifiers: - if mod.type == "ARMATURE": - if mod.object: - bpy.data.objects.remove(mod.object) - object.modifiers.remove(mod) - - -def add_bones(object, name="armature"): - # creates bones and assigns correct weights - - clear_armature(object) - - bone_positions, bone_weights, chain_ids = get_bone_positions(object) - - armature = create_bones(bone_positions, chain_ids) - for i in range(bone_weights.shape[1]): - group = object.vertex_groups.new(name=f"mn_armature_{i}") - vertex_indices = np.where(bone_weights[:, i] == 1)[0] - group.add(vertex_indices.tolist(), 1, "ADD") - - object.select_set(True) - armature.select_set(True) - bpy.context.view_layer.objects.active = armature - bpy.ops.object.parent_set(type="ARMATURE") - - bpy.context.view_layer.objects.active = object - bpy.ops.object.modifier_move_to_index("EXEC_DEFAULT", modifier="Armature", index=0) - - return armature - - -def get_bone_positions(object): - positions, atom_name, chain_id, res_id, sec_struct = [ - mesh.named_attribute(object, att) - for att in ["position", "atom_name", "chain_id", "res_id", "sec_struct"] - ] - - is_alpha_carbon = atom_name == 2 - idx = np.where(is_alpha_carbon)[0] - bone_positions = positions[idx, :] - bone_positions = np.vstack((bone_positions, positions[-1])) - group_ids = np.cumsum(is_alpha_carbon) - groups = np.unique(group_ids) - bone_weights = np.zeros((len(group_ids), len(groups))) - - for i, unique_id in enumerate(groups): - bone_weights[:, i] = ((group_ids - 1) == unique_id).astype(int) - - print("get_bone_positions") - return bone_positions, bone_weights, chain_id[idx] - - -def get_bone_weights(object): - print("hello world") - - -def create_bones(positions, chain_ids, name="armature"): - bpy.ops.object.add(type="ARMATURE", enter_editmode=True) - object = bpy.context.active_object - object.name = name - coll.armature().objects.link(object) - armature = object.data - armature.name = f"{name}_frame" - arm_name = armature.name - bones = [] - # add bones - for i, position in enumerate(positions): - try: - pos_a = position - pos_b = positions[i + 1, :] - except: - continue - - bone_name = f"mn_armature_{i}" - bone = armature.edit_bones.new(bone_name) - bone.head = pos_a - bone.tail = pos_b - bones.append(bone.name) - - armature = bpy.data.armatures[arm_name] - bones_a = bones.copy() - bones_b = bones.copy() - bones_b.pop(0) - bones = zip(bones_a, bones_b) - - for bone_a, bone_b in bones: - armature.edit_bones.active = armature.edit_bones[bone_a] - for bone in [bone_a, bone_b]: - armature.edit_bones[bone].select = True - bpy.ops.armature.parent_set(type="CONNECTED") - for bone in [bone_a, bone_b]: - armature.edit_bones[bone].select = False - bpy.ops.object.editmode_toggle() - - return object - - -class MN_MT_Add_Armature(bpy.types.Operator): - bl_idname = "mn.add_armature" - bl_label = "Add Armature" - bl_description = ( - "Automatically add armature for each amino acid of the structure " - ) - - def execute(self, context): - object = context.active_object - add_bones(bpy.data.objects[object.name], name=object.name) - - return {"FINISHED"} diff --git a/molecularnodes/blender/databpy/__init__.py b/molecularnodes/blender/databpy/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/molecularnodes/blender/mesh.py b/molecularnodes/blender/mesh.py index f241ef7b..78c13049 100644 --- a/molecularnodes/blender/mesh.py +++ b/molecularnodes/blender/mesh.py @@ -1,7 +1,7 @@ import bpy import numpy as np -from typing import Optional +from typing import Optional, Type from enum import Enum from . import coll @@ -16,19 +16,9 @@ class AttributeTypeInfo: width: int -TYPES = { - key: AttributeTypeInfo(*values) - for key, values in { - "FLOAT_VECTOR": ("vector", float, [3]), - "FLOAT_COLOR": ("color", float, [4]), - "QUATERNION": ("value", float, [4]), - "INT": ("value", int, [1]), - "FLOAT": ("value", float, [1]), - "INT32_2D": ("value", int, [2]), - "FLOAT4X4": ("value", float, [4, 4]), - "BOOLEAN": ("value", bool, [1]), - }.items() -} +@dataclass +class DomainInfo: + name: str class AttributeMismatchError(Exception): @@ -37,6 +27,73 @@ def __init__(self, message): super().__init__(self.message) +# https://docs.blender.org/api/current/bpy_types_enum_items/attribute_domain_items.html#rna-enum-attribute-domain-items +class Domains: + POINT = DomainInfo(name="POINT") + EDGE = DomainInfo(name="EDGE") + FACE = DomainInfo(name="FACE") + CORNER = DomainInfo(name="CORNER") + CURVE = DomainInfo(name="CURVE") + INSTANCE = DomainInfo(name="INSTNANCE") + LAYER = DomainInfo(name="LAYER") + + +@dataclass +class AttributeInfo: + type_name: str + value_name: str + dtype: Type + dimensions: tuple + + +# https://docs.blender.org/api/current/bpy_types_enum_items/attribute_type_items.html#rna-enum-attribute-type-items +class AttributeTypes(Enum): + # https://docs.blender.org/api/current/bpy.types.FloatAttribute.html#bpy.types.FloatAttribute + FLOAT = AttributeInfo( + type_name="FLOAT", value_name="value", dtype=float, dimensions=tuple + ) + # https://docs.blender.org/api/current/bpy.types.FloatVectorAttribute.html#bpy.types.FloatVectorAttribute + FLOAT_VECTOR = AttributeInfo( + type_name="FLOAT_VECTOR", value_name="vector", dtype=float, dimensions=(3,) + ) + # https://docs.blender.org/api/current/bpy.types.Float2Attribute.html#bpy.types.Float2Attribute + FLOAT2 = AttributeInfo( + type_name="FLOAT2", value_name="vector", dtype=float, dimensions=(2,) + ) + # alternatively use color_srgb to get the color info in sRGB color space, otherwise linear color space + # https://docs.blender.org/api/current/bpy.types.FloatColorAttributeValue.html#bpy.types.FloatColorAttributeValue + FLOAT_COLOR = AttributeInfo( + type_name="FLOAT_COLOR", value_name="color", dtype=float, dimensions=(4,) + ) + # https://docs.blender.org/api/current/bpy.types.ByteColorAttribute.html#bpy.types.ByteColorAttribute + # TODO unsure about this, int values are stored but float values are returned + BYTE_COLOR = AttributeInfo( + type_name="BYTE_COLOR", value_name="color", dtype=int, dimensions=(4,) + ) + # https://docs.blender.org/api/current/bpy.types.QuaternionAttribute.html#bpy.types.QuaternionAttribute + QUATERNION = AttributeInfo( + type_name="QUATERNION", value_name="value", dtype=float, dimensions=(4,) + ) + # https://docs.blender.org/api/current/bpy.types.IntAttribute.html#bpy.types.IntAttribute + INT = AttributeInfo(type_name="INT", value_name="value", dtype=int, dimensions=(1,)) + # https://docs.blender.org/api/current/bpy.types.ByteIntAttributeValue.html#bpy.types.ByteIntAttributeValue + INT8 = AttributeInfo( + type_name="INT8", value_name="value", dtype=int, dimensions=(1,) + ) + # https://docs.blender.org/api/current/bpy.types.Int2Attribute.html#bpy.types.Int2Attribute + INT32_2D = AttributeInfo( + type_name="INT32_2D", value_name="value", dtype=int, dimensions=(2,) + ) + # https://docs.blender.org/api/current/bpy.types.Float4x4Attribute.html#bpy.types.Float4x4Attribute + FLOAT4X4 = AttributeInfo( + type_name="FLOAT4X4", value_name="value", dtype=float, dimensions=(4, 4) + ) + # https://docs.blender.org/api/current/bpy.types.BoolAttribute.html#bpy.types.BoolAttribute + BOOLEAN = AttributeInfo( + type_name="BOOLEAN", value_name="value", dtype=bool, dimensions=(1,) + ) + + def centre(array: np.array): return np.mean(array, axis=0) @@ -153,20 +210,37 @@ def create_object( return object -class AttributeDataType(Enum): - FLOAT_VECTOR = "FLOAT_VECTOR" - FLOAT_COLOR = "FLOAT_COLOR" - QUATERNION = "QUATERNION" - FLOAT = "FLOAT" - INT = "INT" - BOOLEAN = "BOOLEAN" - FLOAT4X4 = "FLOAT4X4" +def guess_attribute_from_array(array: np.ndarray) -> AttributeInfo: + if not isinstance(array, np.ndarray): + raise ValueError(f"`array` must be a numpy array, not {type(array)=}") + + dtype = array.dtype + shape = array.shape + if len(array) == 1: + if np.issubdtype(dtype, np.int_): + return AttributeTypes.INT + elif np.issubdtype(dtype, np.float_): + return AttributeTypes.FLOAT + elif np.issubdtype(dtype, np.bool_): + AttributeTypes.BOOLEAN + elif len(shape) == 3 and shape[1:] == (4, 4): + return AttributeTypes.FLOAT4X4 + else: + if shape[1] == 3: + return AttributeTypes.FLOAT_VECTOR + elif shape[1] == 4: + return AttributeTypes.FLOAT_COLOR + else: + return AttributeTypes.FLOAT + + # if we didn't match against anything return float + return AttributeTypes.FLOAT def store_named_attribute( obj: bpy.types.Object, - name: str, data: np.ndarray, + name: str, data_type: Optional[str] = None, domain: str = "POINT", overwrite: bool = True, @@ -197,39 +271,23 @@ def store_named_attribute( The added attribute. """ - # if the datatype isn't specified, try to guess the datatype based on the - # datatype of the ndarray. This should work but ultimately won't guess between - # the quaternion and color datatype, so will just default to color - if data_type is None: - dtype = data.dtype - shape = data.shape - - if len(shape) == 1: - if np.issubdtype(dtype, np.int_): - data_type = "INT" - elif np.issubdtype(dtype, np.float_): - data_type = "FLOAT" - elif np.issubdtype(dtype, np.bool_): - data_type = "BOOLEAN" - elif len(shape) == 3 and shape[1:] == (4, 4): - data_type = "FLOAT4X4" - else: - if shape[1] == 3: - data_type = "FLOAT_VECTOR" - elif shape[1] == 4: - data_type = "FLOAT_COLOR" - else: - data_type = "FLOAT" - # catch if the data_type still wasn't determined and report info about the data + if isinstance(data_type, str): + try: + attr_info = AttributeTypes[data_type].value + except KeyError: + raise ValueError( + f"Given data type {data_type=} does not match any of the possible attribute types: {list(AttributeTypes)=}" + ) + if data_type is None: - data_type = "FLOAT" - # raise ValueError( - # f"Unable to determine data type for {data}, {shape=}, {dtype=}" - # ) + attr_info = guess_attribute_from_array(data) + + print(f"{data_type=}") + print(f"{attr_info=}") attribute = obj.data.attributes.get(name) # type: ignore if not attribute or not overwrite: - attribute = obj.data.attributes.new(name, data_type, domain) # type: ignore + attribute = obj.data.attributes.new(name, attr_info.type_name, domain) if len(data) != len(attribute.data): raise AttributeMismatchError( @@ -238,7 +296,7 @@ def store_named_attribute( # the 'foreach_set' requires a 1D array, regardless of the shape of the attribute # it also requires the order to be 'c' or blender might crash!! - attribute.data.foreach_set(TYPES[data_type].dname, data.reshape(-1)) + attribute.data.foreach_set(attr_info.value_name, data.reshape(-1)) # The updating of data doesn't work 100% of the time (see: # https://projects.blender.org/blender/blender/issues/118507) so this resetting of a @@ -255,7 +313,7 @@ def store_named_attribute( def named_attribute( - object: bpy.types.Object, name="position", evaluate=False + obj: bpy.types.Object, name="position", evaluate=False ) -> np.ndarray: """ Get the attribute data from the object. @@ -268,8 +326,8 @@ def named_attribute( np.ndarray: The attribute data as a numpy array. """ if evaluate: - object = evaluated(object) - attribute_names = object.data.attributes.keys() + obj = evaluate(obj) + attribute_names = obj.data.attributes.keys() verbose = False if name not in attribute_names: if verbose: @@ -283,10 +341,10 @@ def named_attribute( ) # Get the attribute and some metadata about it from the object - att = object.data.attributes[name] + att = obj.data.attributes[name] n_att = len(att.data) - data_type = TYPES[att.data_type] - dim = data_type.width + attr_info = AttributeTypes[att.data_type].value + dim = attr_info.dimensions n_values = n_att for dimension in dim: n_values *= dimension @@ -295,11 +353,11 @@ def named_attribute( # we have the initialise the array first with the appropriate length, then we can # fill it with the given data using the 'foreach_get' method which is super fast C++ # internal method - array = np.zeros(n_values, dtype=data_type.dtype) + array = np.zeros(n_values, dtype=attr_info.dtype) # it is currently not really consistent, but to get the values you need to use one of # the 'value', 'vector', 'color' etc from the types dict. This I could only figure # out through trial and error. I assume this might be changed / improved in the future - att.data.foreach_get(data_type.dname, array) + att.data.foreach_get(attr_info.value_name, array) if dim == [1]: return array @@ -326,25 +384,25 @@ def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types. # import the volume object with ObjectTracker() as o: bpy.ops.object.volume_import(filepath=file, files=[]) - object = o.latest() + obj = o.latest() if collection: # Move the object to the MolecularNodes collection - initial_collection = object.users_collection[0] - initial_collection.objects.unlink(object) + initial_collection = obj.users_collection[0] + initial_collection.objects.unlink(obj) collection = coll.mn() - collection.objects.link(object) + collection.objects.link(obj) - return object + return obj -def evaluated(object): +def evaluate(obj): "Return an object which has the modifiers evaluated." - object.update_tag() - return object.evaluated_get(bpy.context.evaluated_depsgraph_get()) + obj.update_tag() + return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) -def evaluate_using_mesh(object): +def evaluate_using_mesh(obj): """ Evaluate the object using a debug object. Some objects can't currently have their Geometry Node trees evaluated (such as volumes), so we source the geometry they create @@ -365,18 +423,16 @@ def evaluate_using_mesh(object): """ # create an empty mesh object. It's modifiers can be evaluated but some other # object types can't be currently through the API - debug = create_object() - mod = nodes.get_mod(debug) + debug_obj = create_object() + mod = nodes.get_mod(debug_obj) mod.node_group = nodes.create_debug_group() - mod.node_group.nodes["Object Info"].inputs["Object"].default_value = object + mod.node_group.nodes["Object Info"].inputs["Object"].default_value = obj # need to use 'evaluate' otherwise the modifiers won't be taken into account - return evaluated(debug) + return evaluate(debug_obj) -def create_data_object( - array, collection=None, name="DataObject", world_scale=0.01, fallback=False -): +def create_data_object(array, collection=None, name="DataObject", world_scale=0.01): # still requires a unique call TODO: figure out why # I think this has to do with the bcif instancing extraction array = np.unique(array) @@ -385,7 +441,7 @@ def create_data_object( if not collection: collection = coll.data() - object = create_object(locations, collection=collection, name=name) + obj = create_object(locations, collection=collection, name=name) attributes = [ ("rotation", "QUATERNION"), @@ -404,8 +460,6 @@ def create_data_object( if np.issubdtype(data.dtype, str): data = np.unique(data, return_inverse=True)[1] - store_named_attribute( - object, name=column, data=data, data_type=type, domain="POINT" - ) + store_named_attribute(obj=obj, data=data, name=column, data_type=type) - return object + return obj diff --git a/molecularnodes/entities/ensemble/cellpack.py b/molecularnodes/entities/ensemble/cellpack.py index 6a56afdb..8ebc28f4 100644 --- a/molecularnodes/entities/ensemble/cellpack.py +++ b/molecularnodes/entities/ensemble/cellpack.py @@ -69,11 +69,10 @@ def _create_object_instances( colors = np.tile(color.random_rgb(i), (len(chain_atoms), 1)) bl.mesh.store_named_attribute( - obj, - name="Color", + obj=obj, data=colors, + name="Color", data_type="FLOAT_COLOR", - overwrite=True, ) if node_setup: diff --git a/molecularnodes/entities/ensemble/star.py b/molecularnodes/entities/ensemble/star.py index d76eedcb..c7180fb2 100644 --- a/molecularnodes/entities/ensemble/star.py +++ b/molecularnodes/entities/ensemble/star.py @@ -223,20 +223,24 @@ def create_object(self, name="StarFileObject", node_setup=True, world_scale=0.01 # If col_type is numeric directly add if np.issubdtype(col_type, np.number): bl.mesh.store_named_attribute( - blender_object, - col, - self.data[col].to_numpy().reshape(-1), - "FLOAT", - "POINT", + obj=blender_object, + name=col, + data=self.data[col].to_numpy().reshape(-1), + data_type="FLOAT", + domain="POINT", ) # If col_type is object, convert to category and add integer values - elif col_type == object: + elif isinstance(col_type, object): codes = ( self.data[col].astype("category").cat.codes.to_numpy().reshape(-1) ) bl.mesh.store_named_attribute( - blender_object, col, codes, "INT", "POINT" + obj=blender_object, + data=codes, + name=col, + data_type="INT", + domain="POINT", ) # Add the category names as a property to the blender object blender_object[f"{col}_categories"] = list( diff --git a/molecularnodes/entities/entity.py b/molecularnodes/entities/entity.py index 1e63d261..d5d71ab7 100644 --- a/molecularnodes/entities/entity.py +++ b/molecularnodes/entities/entity.py @@ -150,9 +150,9 @@ def store_named_attribute( ) return None bl.mesh.store_named_attribute( - self.object, - name=name, + obj=self.object, data=data, + name=name, data_type=data_type, domain=domain, overwrite=overwrite, diff --git a/molecularnodes/entities/molecule/molecule.py b/molecularnodes/entities/molecule/molecule.py index e9c76f33..8952b25f 100644 --- a/molecularnodes/entities/molecule/molecule.py +++ b/molecularnodes/entities/molecule/molecule.py @@ -625,8 +625,8 @@ def att_sec_struct(): try: bl.mesh.store_named_attribute( obj, - name=att["name"], data=att["value"](), + name=att["name"], data_type=att["type"], domain=att["domain"], ) diff --git a/molecularnodes/entities/trajectory/dna.py b/molecularnodes/entities/trajectory/dna.py index 6c823fa4..e0b9cdab 100644 --- a/molecularnodes/entities/trajectory/dna.py +++ b/molecularnodes/entities/trajectory/dna.py @@ -190,7 +190,9 @@ def store_named_attributes_to_dna_mol(mol, frame, scale_dna=0.1): if att != "angular_velocity": data *= scale_dna - mesh.store_named_attribute(mol, att, data, data_type="FLOAT_VECTOR") + mesh.store_named_attribute( + obj=mol, data=data, name=att, data_type="FLOAT_VECTOR" + ) def toplogy_to_bond_idx_pairs(topology: np.ndarray): @@ -260,12 +262,16 @@ def load(top, traj, name="oxDNA", setup_nodes=True, world_scale=0.01): ) # adding additional toplogy information from the topology and frames objects - mesh.store_named_attribute(obj, "res_name", topology[:, 1], "INT") - mesh.store_named_attribute(obj, "chain_id", topology[:, 0], "INT") mesh.store_named_attribute( - obj, - "Color", + obj=obj, data=topology[:, 1], name="res_name", data_type="INT" + ) + mesh.store_named_attribute( + obj=obj, data=topology[:, 0], name="chain_id", data_type="INT" + ) + mesh.store_named_attribute( + obj=obj, data=color.color_chains_equidistant(topology[:, 0]), + name="Color", data_type="FLOAT_COLOR", ) store_named_attributes_to_dna_mol(obj, trajectory[0], scale_dna=scale_dna) diff --git a/molecularnodes/entities/trajectory/trajectory.py b/molecularnodes/entities/trajectory/trajectory.py index 928b5b9d..11135425 100644 --- a/molecularnodes/entities/trajectory/trajectory.py +++ b/molecularnodes/entities/trajectory/trajectory.py @@ -438,7 +438,11 @@ def create_object( for att_name, att in self._attributes_2_blender.items(): try: mesh.store_named_attribute( - obj, att_name, att["value"], att["type"], att["domain"] + obj=obj, + data=att["value"], + name=att_name, + data_type=att["type"], + domain=att["domain"], ) except Exception as e: print(e) @@ -467,7 +471,7 @@ def create_object( def _update_calculations(self): for name, func in self.calculations.items(): try: - self.store_named_attribute(name=name, data=func(self.universe)) + self.store_named_attribute(data=func(self.universe), name=name) except Exception as e: print(e) diff --git a/tests/test_obj.py b/tests/test_obj.py index f340e2e1..0129b624 100644 --- a/tests/test_obj.py +++ b/tests/test_obj.py @@ -50,11 +50,13 @@ def test_matrix_read_write(): arr = np.array((5, 4, 4), float) arr = np.random.rand(5, 4, 4) - mesh.store_named_attribute(obj, "test_matrix", arr, "FLOAT4X4") + mesh.store_named_attribute( + obj=obj, data=arr, name="test_matrix", data_type="FLOAT4X4" + ) assert np.allclose(mesh.named_attribute(obj, "test_matrix"), arr) arr2 = np.random.rand(5, 4, 4) - mesh.store_named_attribute(obj, "test_matrix2", arr2) + mesh.store_named_attribute(obj=obj, data=arr2, name="test_matrix2") assert not np.allclose(mesh.named_attribute(obj, "test_matrix2"), arr) diff --git a/tests/test_select.py b/tests/test_select.py index 56469b75..40d399ac 100644 --- a/tests/test_select.py +++ b/tests/test_select.py @@ -29,7 +29,11 @@ def evaluate(object): def test_select_multiple_residues(selection): n_atoms = 100 object = mn.blender.mesh.create_object(np.zeros((n_atoms, 3))) - mn.blender.mesh.store_named_attribute(object, "res_id", np.arange(n_atoms) + 1) + mn.blender.mesh.store_named_attribute( + obj=object, + data=np.arange(n_atoms) + 1, + name="res_id", + ) mod = nodes.get_mod(object) group = nodes.new_group(fallback=False) From 78b6f136674ad915032c7462b64399d4e18c6d70 Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Wed, 30 Oct 2024 17:01:16 +0800 Subject: [PATCH 02/12] minor tweaks --- molecularnodes/blender/mesh.py | 26 +++++++++++++------------- tests/utils.py | 4 ++-- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/molecularnodes/blender/mesh.py b/molecularnodes/blender/mesh.py index 78c13049..078d4339 100644 --- a/molecularnodes/blender/mesh.py +++ b/molecularnodes/blender/mesh.py @@ -39,7 +39,7 @@ class Domains: @dataclass -class AttributeInfo: +class AttributeType: type_name: str value_name: str dtype: Type @@ -49,47 +49,47 @@ class AttributeInfo: # https://docs.blender.org/api/current/bpy_types_enum_items/attribute_type_items.html#rna-enum-attribute-type-items class AttributeTypes(Enum): # https://docs.blender.org/api/current/bpy.types.FloatAttribute.html#bpy.types.FloatAttribute - FLOAT = AttributeInfo( + FLOAT = AttributeType( type_name="FLOAT", value_name="value", dtype=float, dimensions=tuple ) # https://docs.blender.org/api/current/bpy.types.FloatVectorAttribute.html#bpy.types.FloatVectorAttribute - FLOAT_VECTOR = AttributeInfo( + FLOAT_VECTOR = AttributeType( type_name="FLOAT_VECTOR", value_name="vector", dtype=float, dimensions=(3,) ) # https://docs.blender.org/api/current/bpy.types.Float2Attribute.html#bpy.types.Float2Attribute - FLOAT2 = AttributeInfo( + FLOAT2 = AttributeType( type_name="FLOAT2", value_name="vector", dtype=float, dimensions=(2,) ) # alternatively use color_srgb to get the color info in sRGB color space, otherwise linear color space # https://docs.blender.org/api/current/bpy.types.FloatColorAttributeValue.html#bpy.types.FloatColorAttributeValue - FLOAT_COLOR = AttributeInfo( + FLOAT_COLOR = AttributeType( type_name="FLOAT_COLOR", value_name="color", dtype=float, dimensions=(4,) ) # https://docs.blender.org/api/current/bpy.types.ByteColorAttribute.html#bpy.types.ByteColorAttribute # TODO unsure about this, int values are stored but float values are returned - BYTE_COLOR = AttributeInfo( + BYTE_COLOR = AttributeType( type_name="BYTE_COLOR", value_name="color", dtype=int, dimensions=(4,) ) # https://docs.blender.org/api/current/bpy.types.QuaternionAttribute.html#bpy.types.QuaternionAttribute - QUATERNION = AttributeInfo( + QUATERNION = AttributeType( type_name="QUATERNION", value_name="value", dtype=float, dimensions=(4,) ) # https://docs.blender.org/api/current/bpy.types.IntAttribute.html#bpy.types.IntAttribute - INT = AttributeInfo(type_name="INT", value_name="value", dtype=int, dimensions=(1,)) + INT = AttributeType(type_name="INT", value_name="value", dtype=int, dimensions=(1,)) # https://docs.blender.org/api/current/bpy.types.ByteIntAttributeValue.html#bpy.types.ByteIntAttributeValue - INT8 = AttributeInfo( + INT8 = AttributeType( type_name="INT8", value_name="value", dtype=int, dimensions=(1,) ) # https://docs.blender.org/api/current/bpy.types.Int2Attribute.html#bpy.types.Int2Attribute - INT32_2D = AttributeInfo( + INT32_2D = AttributeType( type_name="INT32_2D", value_name="value", dtype=int, dimensions=(2,) ) # https://docs.blender.org/api/current/bpy.types.Float4x4Attribute.html#bpy.types.Float4x4Attribute - FLOAT4X4 = AttributeInfo( + FLOAT4X4 = AttributeType( type_name="FLOAT4X4", value_name="value", dtype=float, dimensions=(4, 4) ) # https://docs.blender.org/api/current/bpy.types.BoolAttribute.html#bpy.types.BoolAttribute - BOOLEAN = AttributeInfo( + BOOLEAN = AttributeType( type_name="BOOLEAN", value_name="value", dtype=bool, dimensions=(1,) ) @@ -210,7 +210,7 @@ def create_object( return object -def guess_attribute_from_array(array: np.ndarray) -> AttributeInfo: +def guess_attribute_from_array(array: np.ndarray) -> AttributeType: if not isinstance(array, np.ndarray): raise ValueError(f"`array` must be a numpy array, not {type(array)=}") diff --git a/tests/utils.py b/tests/utils.py index 1f2e2a85..4e20b9f3 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -65,7 +65,7 @@ def sample_attribute( random.seed(seed) if error: attribute = mn.blender.mesh.named_attribute( - object, attribute, evaluate=evaluate + obj=object, naame=attribute, evaluate=evaluate ) length = len(attribute) @@ -81,7 +81,7 @@ def sample_attribute( else: try: attribute = mn.blender.mesh.named_attribute( - object=object, name=attribute, evaluate=evaluate + obj=object, name=attribute, evaluate=evaluate ) length = len(attribute) From d84b8d363c4c4e3d2166a71dcd94d29dc6df7f7c Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 31 Oct 2024 07:30:18 +0800 Subject: [PATCH 03/12] working refactor and tests --- molecularnodes/blender/attribute.py | 178 +++++++++++++++ molecularnodes/blender/mesh.py | 285 ++----------------------- molecularnodes/blender/object.py | 110 ++++++++++ molecularnodes/entities/entity.py | 2 +- molecularnodes/entities/molecule/ui.py | 18 +- tests/test_attributes.py | 7 +- tests/test_nodes.py | 4 +- tests/test_select.py | 8 +- tests/utils.py | 5 +- 9 files changed, 328 insertions(+), 289 deletions(-) create mode 100644 molecularnodes/blender/attribute.py create mode 100644 molecularnodes/blender/object.py diff --git a/molecularnodes/blender/attribute.py b/molecularnodes/blender/attribute.py new file mode 100644 index 00000000..e32fb05f --- /dev/null +++ b/molecularnodes/blender/attribute.py @@ -0,0 +1,178 @@ +from dataclasses import dataclass +from enum import Enum +from typing import Type +from functools import reduce + +import bpy +import numpy as np +from numpy.ma.core import prod + + +@dataclass +class AttributeTypeInfo: + dname: str + dtype: type + width: int + + +@dataclass +class DomainInfo: + name: str + + +class AttributeMismatchError(Exception): + def __init__(self, message): + self.message = message + super().__init__(self.message) + + +# https://docs.blender.org/api/current/bpy_types_enum_items/attribute_domain_items.html#rna-enum-attribute-domain-items +class Domains: + POINT = DomainInfo(name="POINT") + EDGE = DomainInfo(name="EDGE") + FACE = DomainInfo(name="FACE") + CORNER = DomainInfo(name="CORNER") + CURVE = DomainInfo(name="CURVE") + INSTANCE = DomainInfo(name="INSTNANCE") + LAYER = DomainInfo(name="LAYER") + + +@dataclass +class AttributeType: + type_name: str + value_name: str + dtype: Type + dimensions: tuple + + +# https://docs.blender.org/api/current/bpy_types_enum_items/attribute_type_items.html#rna-enum-attribute-type-items +class AttributeTypes(Enum): + # https://docs.blender.org/api/current/bpy.types.FloatAttribute.html#bpy.types.FloatAttribute + FLOAT = AttributeType( + type_name="FLOAT", value_name="value", dtype=float, dimensions=(1,) + ) + # https://docs.blender.org/api/current/bpy.types.FloatVectorAttribute.html#bpy.types.FloatVectorAttribute + FLOAT_VECTOR = AttributeType( + type_name="FLOAT_VECTOR", value_name="vector", dtype=float, dimensions=(3,) + ) + # https://docs.blender.org/api/current/bpy.types.Float2Attribute.html#bpy.types.Float2Attribute + FLOAT2 = AttributeType( + type_name="FLOAT2", value_name="vector", dtype=float, dimensions=(2,) + ) + # alternatively use color_srgb to get the color info in sRGB color space, otherwise linear color space + # https://docs.blender.org/api/current/bpy.types.FloatColorAttributeValue.html#bpy.types.FloatColorAttributeValue + FLOAT_COLOR = AttributeType( + type_name="FLOAT_COLOR", value_name="color", dtype=float, dimensions=(4,) + ) + # https://docs.blender.org/api/current/bpy.types.ByteColorAttribute.html#bpy.types.ByteColorAttribute + # TODO unsure about this, int values are stored but float values are returned + BYTE_COLOR = AttributeType( + type_name="BYTE_COLOR", value_name="color", dtype=int, dimensions=(4,) + ) + # https://docs.blender.org/api/current/bpy.types.QuaternionAttribute.html#bpy.types.QuaternionAttribute + QUATERNION = AttributeType( + type_name="QUATERNION", value_name="value", dtype=float, dimensions=(4,) + ) + # https://docs.blender.org/api/current/bpy.types.IntAttribute.html#bpy.types.IntAttribute + INT = AttributeType(type_name="INT", value_name="value", dtype=int, dimensions=(1,)) + # https://docs.blender.org/api/current/bpy.types.ByteIntAttributeValue.html#bpy.types.ByteIntAttributeValue + INT8 = AttributeType( + type_name="INT8", value_name="value", dtype=int, dimensions=(1,) + ) + # https://docs.blender.org/api/current/bpy.types.Int2Attribute.html#bpy.types.Int2Attribute + INT32_2D = AttributeType( + type_name="INT32_2D", value_name="value", dtype=int, dimensions=(2,) + ) + # https://docs.blender.org/api/current/bpy.types.Float4x4Attribute.html#bpy.types.Float4x4Attribute + FLOAT4X4 = AttributeType( + type_name="FLOAT4X4", value_name="value", dtype=float, dimensions=(4, 4) + ) + # https://docs.blender.org/api/current/bpy.types.BoolAttribute.html#bpy.types.BoolAttribute + BOOLEAN = AttributeType( + type_name="BOOLEAN", value_name="value", dtype=bool, dimensions=(1,) + ) + +def guess_atype_from_array(array: np.ndarray) -> AttributeType: + if not isinstance(array, np.ndarray): + raise ValueError(f"`array` must be a numpy array, not {type(array)=}") + + dtype = array.dtype + shape = array.shape + n_row = shape[0] + + # for 1D arrays we we use the float, int of boolean attribute types + if shape == (n_row, 1) or shape == (n_row,): + if np.issubdtype(dtype, np.int_): + return AttributeTypes.INT.value + elif np.issubdtype(dtype, np.float_): + return AttributeTypes.FLOAT.value + elif np.issubdtype(dtype, np.bool_): + return AttributeTypes.BOOLEAN.value + + # for 2D arrays we use the float_vector, float_color, float4x4 attribute types + elif shape == (n_row, 4, 4): + return AttributeTypes.FLOAT4X4.value + elif shape == (n_row, 3): + return AttributeTypes.FLOAT_VECTOR.value + elif shape == (n_row, 4): + return AttributeTypes.FLOAT_COLOR.value + + # if we didn't match against anything return float + return AttributeTypes.FLOAT.value + +class Attribute: + def __init__(self, attribute: bpy.types.Attribute): + self.attribute = attribute + self.n_attr = len(attribute.data) + + @property + def atype(self): + try: + atype = AttributeTypes[self.attribute.data_type].value + except KeyError: + raise ValueError(f"Unknown attribute type: {self.attribute.data_type}") + + return atype + + @property + def value_name(self): + return self.atype.value_name + + @property + def is_1d(self): + return self.atype.dimensions == (1,) + + @property + def type_name(self): + return self.atype.type_name + + @property + def shape(self): + return (self.n_attr, *self.atype.dimensions) + + @property + def dtype(self) -> Type: + return self.atype.dtype + + @property + def n_values(self) -> int: + return np.prod(self.shape, dtype=int) + + + def as_array(self) -> np.ndarray: + """ + Returns the attribute data as a numpy array + """ + # initialize empty 1D array that is needed to then be filled with values + # from the Blender attribute + array = np.zeros(self.n_values, dtype=self.dtype) + self.attribute.data.foreach_get(self.value_name, array) + + # if the attribute has more than one dimension reshape the array before returning + if self.is_1d: + return array + else: + return array.reshape(self.shape) + + def __str__(self): + return "Attribute: {}, type: {}, size: {}".format(self.attribute.name, self.type_name, self.shape) diff --git a/molecularnodes/blender/mesh.py b/molecularnodes/blender/mesh.py index 078d4339..5fa369a8 100644 --- a/molecularnodes/blender/mesh.py +++ b/molecularnodes/blender/mesh.py @@ -1,242 +1,25 @@ +from typing import Optional + import bpy import numpy as np -from typing import Optional, Type -from enum import Enum - -from . import coll -from . import nodes -from dataclasses import dataclass - - -@dataclass -class AttributeTypeInfo: - dname: str - dtype: type - width: int - - -@dataclass -class DomainInfo: - name: str - - -class AttributeMismatchError(Exception): - def __init__(self, message): - self.message = message - super().__init__(self.message) - - -# https://docs.blender.org/api/current/bpy_types_enum_items/attribute_domain_items.html#rna-enum-attribute-domain-items -class Domains: - POINT = DomainInfo(name="POINT") - EDGE = DomainInfo(name="EDGE") - FACE = DomainInfo(name="FACE") - CORNER = DomainInfo(name="CORNER") - CURVE = DomainInfo(name="CURVE") - INSTANCE = DomainInfo(name="INSTNANCE") - LAYER = DomainInfo(name="LAYER") - - -@dataclass -class AttributeType: - type_name: str - value_name: str - dtype: Type - dimensions: tuple - - -# https://docs.blender.org/api/current/bpy_types_enum_items/attribute_type_items.html#rna-enum-attribute-type-items -class AttributeTypes(Enum): - # https://docs.blender.org/api/current/bpy.types.FloatAttribute.html#bpy.types.FloatAttribute - FLOAT = AttributeType( - type_name="FLOAT", value_name="value", dtype=float, dimensions=tuple - ) - # https://docs.blender.org/api/current/bpy.types.FloatVectorAttribute.html#bpy.types.FloatVectorAttribute - FLOAT_VECTOR = AttributeType( - type_name="FLOAT_VECTOR", value_name="vector", dtype=float, dimensions=(3,) - ) - # https://docs.blender.org/api/current/bpy.types.Float2Attribute.html#bpy.types.Float2Attribute - FLOAT2 = AttributeType( - type_name="FLOAT2", value_name="vector", dtype=float, dimensions=(2,) - ) - # alternatively use color_srgb to get the color info in sRGB color space, otherwise linear color space - # https://docs.blender.org/api/current/bpy.types.FloatColorAttributeValue.html#bpy.types.FloatColorAttributeValue - FLOAT_COLOR = AttributeType( - type_name="FLOAT_COLOR", value_name="color", dtype=float, dimensions=(4,) - ) - # https://docs.blender.org/api/current/bpy.types.ByteColorAttribute.html#bpy.types.ByteColorAttribute - # TODO unsure about this, int values are stored but float values are returned - BYTE_COLOR = AttributeType( - type_name="BYTE_COLOR", value_name="color", dtype=int, dimensions=(4,) - ) - # https://docs.blender.org/api/current/bpy.types.QuaternionAttribute.html#bpy.types.QuaternionAttribute - QUATERNION = AttributeType( - type_name="QUATERNION", value_name="value", dtype=float, dimensions=(4,) - ) - # https://docs.blender.org/api/current/bpy.types.IntAttribute.html#bpy.types.IntAttribute - INT = AttributeType(type_name="INT", value_name="value", dtype=int, dimensions=(1,)) - # https://docs.blender.org/api/current/bpy.types.ByteIntAttributeValue.html#bpy.types.ByteIntAttributeValue - INT8 = AttributeType( - type_name="INT8", value_name="value", dtype=int, dimensions=(1,) - ) - # https://docs.blender.org/api/current/bpy.types.Int2Attribute.html#bpy.types.Int2Attribute - INT32_2D = AttributeType( - type_name="INT32_2D", value_name="value", dtype=int, dimensions=(2,) - ) - # https://docs.blender.org/api/current/bpy.types.Float4x4Attribute.html#bpy.types.Float4x4Attribute - FLOAT4X4 = AttributeType( - type_name="FLOAT4X4", value_name="value", dtype=float, dimensions=(4, 4) - ) - # https://docs.blender.org/api/current/bpy.types.BoolAttribute.html#bpy.types.BoolAttribute - BOOLEAN = AttributeType( - type_name="BOOLEAN", value_name="value", dtype=bool, dimensions=(1,) - ) - - -def centre(array: np.array): - return np.mean(array, axis=0) +from . import coll, nodes +from .attribute import ( + Attribute, + AttributeMismatchError, + AttributeTypes, + guess_atype_from_array, +) +from .object import ObjectTracker, create_object +def centre(array: np.ndarray): + return np.mean(array, axis=0) + def centre_weighted(array: np.ndarray, weight: np.ndarray): return np.sum(array * weight.reshape((len(array), 1)), axis=0) / np.sum(weight) -class ObjectTracker: - """ - A context manager for tracking new objects in Blender. - - This class provides a way to track new objects that are added to Blender's bpy.data.objects collection. - It stores the current objects when entering the context and provides a method to find new objects that were added when exiting the context. - - Methods - ------- - new_objects(): - Returns a list of new objects that were added to bpy.data.objects while in the context. - """ - - def __enter__(self): - """ - Store the current objects and their names when entering the context. - - Returns - ------- - self - The instance of the class. - """ - self.objects = list(bpy.context.scene.objects) - return self - - def __exit__(self, type, value, traceback): - pass - - def new_objects(self): - """ - Find new objects that were added to bpy.data.objects while in the context. - - Use new_objects()[-1] to get the most recently added object. - - Returns - ------- - list - A list of new objects. - """ - obj_names = list([o.name for o in self.objects]) - current_objects = bpy.context.scene.objects - new_objects = [] - for obj in current_objects: - if obj.name not in obj_names: - new_objects.append(obj) - return new_objects - - def latest(self): - """ - Get the most recently added object. - - This method returns the most recently added object to bpy.data.objects while in the context. - - Returns - ------- - bpy.types.Object - The most recently added object. - """ - return self.new_objects()[-1] - - -def create_object( - vertices: np.ndarray = [], - edges: np.ndarray = [], - faces: np.ndarray = [], - name: str = "NewObject", - collection: bpy.types.Collection = None, -) -> bpy.types.Object: - """ - Create a new Blender object, initialised with locations for each vertex. - - If edges and faces are supplied then these are also created on the mesh. - - Parameters - ---------- - vertices : np.ndarray, optional - The vertices of the vertices as a numpy array. Defaults to None. - edges : np.ndarray, optional - The edges of the object as a numpy array. Defaults to None. - faces : np.ndarray, optional - The faces of the object as a numpy array. Defaults to None. - name : str, optional - The name of the object. Defaults to 'NewObject'. - collection : bpy.types.Collection, optional - The collection to link the object to. Defaults to None. - - Returns - ------- - bpy.types.Object - The created object. - """ - mesh = bpy.data.meshes.new(name) - - mesh.from_pydata(vertices=vertices, edges=edges, faces=faces) - - object = bpy.data.objects.new(name, mesh) - - if not collection: - # Add the object to the scene if no collection is specified - collection = bpy.data.collections["Collection"] - - collection.objects.link(object) - - object["type"] = "molecule" - - return object - - -def guess_attribute_from_array(array: np.ndarray) -> AttributeType: - if not isinstance(array, np.ndarray): - raise ValueError(f"`array` must be a numpy array, not {type(array)=}") - - dtype = array.dtype - shape = array.shape - if len(array) == 1: - if np.issubdtype(dtype, np.int_): - return AttributeTypes.INT - elif np.issubdtype(dtype, np.float_): - return AttributeTypes.FLOAT - elif np.issubdtype(dtype, np.bool_): - AttributeTypes.BOOLEAN - elif len(shape) == 3 and shape[1:] == (4, 4): - return AttributeTypes.FLOAT4X4 - else: - if shape[1] == 3: - return AttributeTypes.FLOAT_VECTOR - elif shape[1] == 4: - return AttributeTypes.FLOAT_COLOR - else: - return AttributeTypes.FLOAT - - # if we didn't match against anything return float - return AttributeTypes.FLOAT - - def store_named_attribute( obj: bpy.types.Object, data: np.ndarray, @@ -273,21 +56,18 @@ def store_named_attribute( if isinstance(data_type, str): try: - attr_info = AttributeTypes[data_type].value + atype = AttributeTypes[data_type].value except KeyError: raise ValueError( f"Given data type {data_type=} does not match any of the possible attribute types: {list(AttributeTypes)=}" ) if data_type is None: - attr_info = guess_attribute_from_array(data) - - print(f"{data_type=}") - print(f"{attr_info=}") + atype = guess_atype_from_array(data) attribute = obj.data.attributes.get(name) # type: ignore if not attribute or not overwrite: - attribute = obj.data.attributes.new(name, attr_info.type_name, domain) + attribute = obj.data.attributes.new(name, atype.type_name, domain) if len(data) != len(attribute.data): raise AttributeMismatchError( @@ -296,7 +76,7 @@ def store_named_attribute( # the 'foreach_set' requires a 1D array, regardless of the shape of the attribute # it also requires the order to be 'c' or blender might crash!! - attribute.data.foreach_set(attr_info.value_name, data.reshape(-1)) + attribute.data.foreach_set(atype.value_name, data.reshape(-1)) # The updating of data doesn't work 100% of the time (see: # https://projects.blender.org/blender/blender/issues/118507) so this resetting of a @@ -326,7 +106,7 @@ def named_attribute( np.ndarray: The attribute data as a numpy array. """ if evaluate: - obj = evaluate(obj) + obj = evaluate_object(obj) attribute_names = obj.data.attributes.keys() verbose = False if name not in attribute_names: @@ -340,30 +120,7 @@ def named_attribute( f"The selected attribute '{name}' does not exist on the mesh." ) - # Get the attribute and some metadata about it from the object - att = obj.data.attributes[name] - n_att = len(att.data) - attr_info = AttributeTypes[att.data_type].value - dim = attr_info.dimensions - n_values = n_att - for dimension in dim: - n_values *= dimension - - # data to and from attributes has to be given and taken as a 1D array - # we have the initialise the array first with the appropriate length, then we can - # fill it with the given data using the 'foreach_get' method which is super fast C++ - # internal method - array = np.zeros(n_values, dtype=attr_info.dtype) - # it is currently not really consistent, but to get the values you need to use one of - # the 'value', 'vector', 'color' etc from the types dict. This I could only figure - # out through trial and error. I assume this might be changed / improved in the future - att.data.foreach_get(attr_info.value_name, array) - - if dim == [1]: - return array - else: - # return an array with one row per item, even if a 1D attribute. Does this make sense? - return array.reshape((n_att, *dim)) + return Attribute(obj.data.attributes[name]).as_array() def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types.Object: @@ -396,7 +153,7 @@ def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types. return obj -def evaluate(obj): +def evaluate_object(obj): "Return an object which has the modifiers evaluated." obj.update_tag() return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) @@ -429,7 +186,7 @@ def evaluate_using_mesh(obj): mod.node_group.nodes["Object Info"].inputs["Object"].default_value = obj # need to use 'evaluate' otherwise the modifiers won't be taken into account - return evaluate(debug_obj) + return evaluate_object(debug_obj) def create_data_object(array, collection=None, name="DataObject", world_scale=0.01): diff --git a/molecularnodes/blender/object.py b/molecularnodes/blender/object.py new file mode 100644 index 00000000..b4737254 --- /dev/null +++ b/molecularnodes/blender/object.py @@ -0,0 +1,110 @@ +import bpy +import numpy as np + + +class ObjectTracker: + """ + A context manager for tracking new objects in Blender. + + This class provides a way to track new objects that are added to Blender's bpy.data.objects collection. + It stores the current objects when entering the context and provides a method to find new objects that were added when exiting the context. + + Methods + ------- + new_objects(): + Returns a list of new objects that were added to bpy.data.objects while in the context. + """ + + def __enter__(self): + """ + Store the current objects and their names when entering the context. + + Returns + ------- + self + The instance of the class. + """ + self.objects = list(bpy.context.scene.objects) + return self + + def __exit__(self, type, value, traceback): + pass + + def new_objects(self): + """ + Find new objects that were added to bpy.data.objects while in the context. + + Use new_objects()[-1] to get the most recently added object. + + Returns + ------- + list + A list of new objects. + """ + obj_names = list([o.name for o in self.objects]) + current_objects = bpy.context.scene.objects + new_objects = [] + for obj in current_objects: + if obj.name not in obj_names: + new_objects.append(obj) + return new_objects + + def latest(self): + """ + Get the most recently added object. + + This method returns the most recently added object to bpy.data.objects while in the context. + + Returns + ------- + bpy.types.Object + The most recently added object. + """ + return self.new_objects()[-1] + + +def create_object( + vertices: np.ndarray = [], + edges: np.ndarray = [], + faces: np.ndarray = [], + name: str = "NewObject", + collection: bpy.types.Collection = None, +) -> bpy.types.Object: + """ + Create a new Blender object, initialised with locations for each vertex. + + If edges and faces are supplied then these are also created on the mesh. + + Parameters + ---------- + vertices : np.ndarray, optional + The vertices of the vertices as a numpy array. Defaults to None. + edges : np.ndarray, optional + The edges of the object as a numpy array. Defaults to None. + faces : np.ndarray, optional + The faces of the object as a numpy array. Defaults to None. + name : str, optional + The name of the object. Defaults to 'NewObject'. + collection : bpy.types.Collection, optional + The collection to link the object to. Defaults to None. + + Returns + ------- + bpy.types.Object + The created object. + """ + mesh = bpy.data.meshes.new(name) + + mesh.from_pydata(vertices=vertices, edges=edges, faces=faces) + + object = bpy.data.objects.new(name, mesh) + + if not collection: + # Add the object to the scene if no collection is specified + collection = bpy.data.collections["Collection"] + + collection.objects.link(object) + + object["type"] = "molecule" + + return object diff --git a/molecularnodes/entities/entity.py b/molecularnodes/entities/entity.py index d5d71ab7..41011db1 100644 --- a/molecularnodes/entities/entity.py +++ b/molecularnodes/entities/entity.py @@ -178,6 +178,6 @@ def list_attributes(cls, evaluate=False) -> list | None: warnings.warn("No object created") return None if evaluate: - return list(bl.mesh.evaluated(cls.object).data.attributes.keys()) + return list(bl.mesh.evaluate_object(cls.object).data.attributes.keys()) return list(cls.object.data.attributes.keys()) diff --git a/molecularnodes/entities/molecule/ui.py b/molecularnodes/entities/molecule/ui.py index f574d7df..74c4113a 100644 --- a/molecularnodes/entities/molecule/ui.py +++ b/molecularnodes/entities/molecule/ui.py @@ -44,16 +44,16 @@ def parse(filepath) -> Molecule: def fetch( - pdb_code, - style="spheres", - centre="", - del_solvent=True, - del_hydrogen=False, - cache_dir=None, - build_assembly=False, + pdb_code: str, + style: str | None="spheres", + centre: str="", + del_solvent: bool=True, + del_hydrogen: bool=False, + cache_dir: str | None=None, + build_assembly: bool=False, database: str = "rcsb", - format="bcif", - color="common", + format: str="bcif", + color: str="common", ) -> Molecule: if build_assembly: centre = "" diff --git a/tests/test_attributes.py b/tests/test_attributes.py index 8f461d49..903cdbdf 100644 --- a/tests/test_attributes.py +++ b/tests/test_attributes.py @@ -1,11 +1,12 @@ -import molecularnodes as mn -import pytest import itertools + import numpy as np +import pytest +import molecularnodes as mn +from .constants import attributes, codes, data_dir from .utils import sample_attribute -from .constants import codes, attributes, data_dir mn._test_register() diff --git a/tests/test_nodes.py b/tests/test_nodes.py index b36c52c2..6bf0f61b 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -284,7 +284,7 @@ def test_topo_bonds(): # compare the number of edges before and after deleting them with bonds = mol.data.edges - no_bonds = mn.blender.mesh.evaluated(mol).data.edges + no_bonds = mn.blender.mesh.evaluate_object(mol).data.edges assert len(bonds) > len(no_bonds) assert len(no_bonds) == 0 @@ -292,5 +292,5 @@ def test_topo_bonds(): # are the same (other attributes will be different, but for now this is good) node_find = nodes.add_custom(group, "Topology Find Bonds") nodes.insert_last_node(group, node=node_find) - bonds_new = mn.blender.mesh.evaluated(mol).data.edges + bonds_new = mn.blender.mesh.evaluate_object(mol).data.edges assert len(bonds) == len(bonds_new) diff --git a/tests/test_select.py b/tests/test_select.py index 40d399ac..467b6d41 100644 --- a/tests/test_select.py +++ b/tests/test_select.py @@ -11,10 +11,6 @@ def create_debug_group(name="MolecularNodesDebugGroup"): group.links.new(info.outputs["Geometry"], group.nodes["Group Output"].inputs[0]) return group - -def evaluate(object): - object.update_tag() - dg = bpy.context.evaluated_depsgraph_get() return object.evaluated_get(dg) @@ -45,9 +41,9 @@ def test_select_multiple_residues(selection): node_sel = nodes.add_custom(group, node_sel_group.name) group.links.new(node_sel.outputs["Selection"], sep.inputs["Selection"]) - vertices_count = len(mn.blender.mesh.evaluated(object).data.vertices) + vertices_count = len(mn.blender.mesh.evaluate_object(object).data.vertices) assert vertices_count == len(selection[1]) assert ( - mn.blender.mesh.named_attribute(mn.blender.mesh.evaluated(object), "res_id") + mn.blender.mesh.named_attribute(mn.blender.mesh.evaluate_object(object), "res_id") == selection[1] ).all() diff --git a/tests/utils.py b/tests/utils.py index 4e20b9f3..923b16ab 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,11 +1,8 @@ import bpy - -# from .conftest import molecularnodes as mn import molecularnodes as mn import numpy as np import random -# import pathlib from syrupy.extensions.amber import AmberSnapshotExtension @@ -65,7 +62,7 @@ def sample_attribute( random.seed(seed) if error: attribute = mn.blender.mesh.named_attribute( - obj=object, naame=attribute, evaluate=evaluate + obj=object, name=attribute, evaluate=evaluate ) length = len(attribute) From 71b60d4541cd792f61132d61939fc5b76410ef62 Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 31 Oct 2024 07:43:28 +0800 Subject: [PATCH 04/12] more refactoring --- molecularnodes/blender/attribute.py | 16 ++++++++++++++-- molecularnodes/blender/mesh.py | 27 ++++++++++----------------- molecularnodes/blender/object.py | 5 +++++ 3 files changed, 29 insertions(+), 19 deletions(-) diff --git a/molecularnodes/blender/attribute.py b/molecularnodes/blender/attribute.py index e32fb05f..a0b0991e 100644 --- a/molecularnodes/blender/attribute.py +++ b/molecularnodes/blender/attribute.py @@ -1,11 +1,9 @@ from dataclasses import dataclass from enum import Enum from typing import Type -from functools import reduce import bpy import numpy as np -from numpy.ma.core import prod @dataclass @@ -121,6 +119,9 @@ def guess_atype_from_array(array: np.ndarray) -> AttributeType: return AttributeTypes.FLOAT.value class Attribute: + """ + Wrapper around a Blender attribute to provide a more convenient interface with numpy arrays + """ def __init__(self, attribute: bpy.types.Attribute): self.attribute = attribute self.n_attr = len(attribute.data) @@ -158,6 +159,17 @@ def dtype(self) -> Type: def n_values(self) -> int: return np.prod(self.shape, dtype=int) + def from_array(self, array: np.ndarray) -> None: + """ + Set the attribute data from a numpy array + """ + if array.shape != self.shape: + raise ValueError( + f"Array shape {array.shape} does not match attribute shape {self.shape}" + ) + + self.attribute.data.foreach_set(self.value_name, array.reshape(-1)) + def as_array(self) -> np.ndarray: """ diff --git a/molecularnodes/blender/mesh.py b/molecularnodes/blender/mesh.py index 5fa369a8..1e6a7fde 100644 --- a/molecularnodes/blender/mesh.py +++ b/molecularnodes/blender/mesh.py @@ -10,7 +10,7 @@ AttributeTypes, guess_atype_from_array, ) -from .object import ObjectTracker, create_object +from .object import ObjectTracker, create_object, evaluate_object def centre(array: np.ndarray): @@ -96,7 +96,7 @@ def named_attribute( obj: bpy.types.Object, name="position", evaluate=False ) -> np.ndarray: """ - Get the attribute data from the object. + Get the named attribute data from the object, optionally evaluating modifiers first. Parameters: object (bpy.types.Object): The Blender object. @@ -107,20 +107,17 @@ def named_attribute( """ if evaluate: obj = evaluate_object(obj) - attribute_names = obj.data.attributes.keys() verbose = False - if name not in attribute_names: + try: + attr = Attribute(obj.data.attributes[name]) + except KeyError: + message = f"The selected attribute '{name}' does not exist on the mesh." if verbose: - raise AttributeError( - f"The selected attribute '{name}' does not exist on the mesh. \ - Possible attributes are: {attribute_names=}" - ) - else: - raise AttributeError( - f"The selected attribute '{name}' does not exist on the mesh." - ) + message += f"Possible attributes are: {obj.data.attributes.keys()}" + + raise AttributeError(message) - return Attribute(obj.data.attributes[name]).as_array() + return attr.as_array() def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types.Object: @@ -153,10 +150,6 @@ def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types. return obj -def evaluate_object(obj): - "Return an object which has the modifiers evaluated." - obj.update_tag() - return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) def evaluate_using_mesh(obj): diff --git a/molecularnodes/blender/object.py b/molecularnodes/blender/object.py index b4737254..568261e1 100644 --- a/molecularnodes/blender/object.py +++ b/molecularnodes/blender/object.py @@ -63,6 +63,11 @@ def latest(self): return self.new_objects()[-1] +def evaluate_object(obj): + "Return an object which has the modifiers evaluated." + obj.update_tag() + return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) + def create_object( vertices: np.ndarray = [], edges: np.ndarray = [], From 7230a9b01932c2a9773d8ffa33679489668ca392 Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 31 Oct 2024 10:33:35 +0800 Subject: [PATCH 05/12] cleanup --- molecularnodes/blender/attribute.py | 44 +++++++++---------- .../entities/trajectory/trajectory.py | 6 +-- molecularnodes/utils.py | 8 +++- 3 files changed, 29 insertions(+), 29 deletions(-) diff --git a/molecularnodes/blender/attribute.py b/molecularnodes/blender/attribute.py index a0b0991e..82a1feb3 100644 --- a/molecularnodes/blender/attribute.py +++ b/molecularnodes/blender/attribute.py @@ -90,6 +90,7 @@ class AttributeTypes(Enum): type_name="BOOLEAN", value_name="value", dtype=bool, dimensions=(1,) ) + def guess_atype_from_array(array: np.ndarray) -> AttributeType: if not isinstance(array, np.ndarray): raise ValueError(f"`array` must be a numpy array, not {type(array)=}") @@ -118,22 +119,16 @@ def guess_atype_from_array(array: np.ndarray) -> AttributeType: # if we didn't match against anything return float return AttributeTypes.FLOAT.value + class Attribute: """ Wrapper around a Blender attribute to provide a more convenient interface with numpy arrays """ + def __init__(self, attribute: bpy.types.Attribute): self.attribute = attribute self.n_attr = len(attribute.data) - - @property - def atype(self): - try: - atype = AttributeTypes[self.attribute.data_type].value - except KeyError: - raise ValueError(f"Unknown attribute type: {self.attribute.data_type}") - - return atype + self.atype = AttributeTypes[self.attribute.data_type].value @property def value_name(self): @@ -170,21 +165,22 @@ def from_array(self, array: np.ndarray) -> None: self.attribute.data.foreach_set(self.value_name, array.reshape(-1)) - def as_array(self) -> np.ndarray: - """ - Returns the attribute data as a numpy array - """ - # initialize empty 1D array that is needed to then be filled with values - # from the Blender attribute - array = np.zeros(self.n_values, dtype=self.dtype) - self.attribute.data.foreach_get(self.value_name, array) - - # if the attribute has more than one dimension reshape the array before returning - if self.is_1d: - return array - else: - return array.reshape(self.shape) + """ + Returns the attribute data as a numpy array + """ + # initialize empty 1D array that is needed to then be filled with values + # from the Blender attribute + array = np.zeros(self.n_values, dtype=self.dtype) + self.attribute.data.foreach_get(self.value_name, array) + + # if the attribute has more than one dimension reshape the array before returning + if self.is_1d: + return array + else: + return array.reshape(self.shape) def __str__(self): - return "Attribute: {}, type: {}, size: {}".format(self.attribute.name, self.type_name, self.shape) + return "Attribute: {}, type: {}, size: {}".format( + self.attribute.name, self.type_name, self.shape + ) diff --git a/molecularnodes/entities/trajectory/trajectory.py b/molecularnodes/entities/trajectory/trajectory.py index 11135425..b78c08db 100644 --- a/molecularnodes/entities/trajectory/trajectory.py +++ b/molecularnodes/entities/trajectory/trajectory.py @@ -13,7 +13,7 @@ class Trajectory(MolecularEntity): - def __init__(self, universe: mda.Universe, world_scale=0.01): + def __init__(self, universe: mda.Universe, world_scale: float=0.01): super().__init__() self.universe: mda.Universe = universe self.selections: Dict[str, Selection] = {} @@ -144,7 +144,7 @@ def positions(self) -> np.ndarray: return self.atoms.positions * self.world_scale @property - def bonds(self) -> List[List[int]]: + def bonds(self) -> np.ndarray: if hasattr(self.atoms, "bonds"): bond_indices = self.atoms.bonds.indices atm_indices = self.atoms.indices @@ -157,7 +157,7 @@ def bonds(self) -> List[List[int]]: bonds = [[index_map[bond[0]], index_map[bond[1]]] for bond in bond_indices] else: bonds = [] - return bonds + return np.array(bonds) @property def elements(self) -> List[str]: diff --git a/molecularnodes/utils.py b/molecularnodes/utils.py index eb2ec6b6..c03f8b69 100644 --- a/molecularnodes/utils.py +++ b/molecularnodes/utils.py @@ -8,7 +8,9 @@ MN_DATA_FILE = os.path.join(ADDON_DIR, "assets", "MN_data_file_4.2.blend") -def correct_periodic_1d(value1, value2, boundary): +def correct_periodic_1d( + value1: np.ndarray, value2: np.ndarray, boundary: float +) -> np.ndarray: diff = value2 - value1 half = boundary / 2 value2[diff > half] -= boundary @@ -16,7 +18,9 @@ def correct_periodic_1d(value1, value2, boundary): return value2 -def correct_periodic_positions(positions_1, positions_2, dimensions): +def correct_periodic_positions( + positions_1: np.ndarray, positions_2: np.ndarray, dimensions: np.ndarray +) -> np.ndarray: if not np.allclose(dimensions[3:], 90.0): raise ValueError( f"Only works with orthorhombic unitcells, and not dimensions={dimensions}" From 4a2384c8c0faeb40e72c8fbc7de24c41ff3f60df Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 31 Oct 2024 12:23:31 +0800 Subject: [PATCH 06/12] more refactoring --- molecularnodes/blender/databpy/__init__.py | 11 ++ .../blender/{ => databpy}/attribute.py | 139 ++++++++++++++++-- .../blender/{ => databpy}/object.py | 100 ++++++++++--- molecularnodes/blender/databpy/utils.py | 7 + molecularnodes/blender/mesh.py | 128 ++-------------- molecularnodes/blender/node.py | 37 ----- molecularnodes/entities/ensemble/cellpack.py | 5 +- molecularnodes/entities/ensemble/star.py | 14 +- molecularnodes/entities/entity.py | 38 ++--- molecularnodes/entities/molecule/molecule.py | 31 ++-- molecularnodes/entities/trajectory/dna.py | 19 +-- .../entities/trajectory/trajectory.py | 8 +- tests/test_obj.py | 19 ++- tests/test_ops.py | 2 +- 14 files changed, 305 insertions(+), 253 deletions(-) rename molecularnodes/blender/{ => databpy}/attribute.py (58%) rename molecularnodes/blender/{ => databpy}/object.py (50%) create mode 100644 molecularnodes/blender/databpy/utils.py delete mode 100644 molecularnodes/blender/node.py diff --git a/molecularnodes/blender/databpy/__init__.py b/molecularnodes/blender/databpy/__init__.py index e69de29b..185bd143 100644 --- a/molecularnodes/blender/databpy/__init__.py +++ b/molecularnodes/blender/databpy/__init__.py @@ -0,0 +1,11 @@ +from .object import ObjectTracker, BlenderObject, bob, create_object +from .attribute import ( + named_attribute, + store_named_attribute, + Attribute, + AttributeType, + AttributeTypeInfo, + AttributeTypes, + Domains, + DomainType, +) diff --git a/molecularnodes/blender/attribute.py b/molecularnodes/blender/databpy/attribute.py similarity index 58% rename from molecularnodes/blender/attribute.py rename to molecularnodes/blender/databpy/attribute.py index 82a1feb3..c9a93fe1 100644 --- a/molecularnodes/blender/attribute.py +++ b/molecularnodes/blender/databpy/attribute.py @@ -1,6 +1,7 @@ from dataclasses import dataclass from enum import Enum from typing import Type +from .utils import evaluate_object import bpy import numpy as np @@ -14,9 +15,12 @@ class AttributeTypeInfo: @dataclass -class DomainInfo: +class DomainType: name: str + def __str__(self): + return self.name + class AttributeMismatchError(Exception): def __init__(self, message): @@ -26,13 +30,13 @@ def __init__(self, message): # https://docs.blender.org/api/current/bpy_types_enum_items/attribute_domain_items.html#rna-enum-attribute-domain-items class Domains: - POINT = DomainInfo(name="POINT") - EDGE = DomainInfo(name="EDGE") - FACE = DomainInfo(name="FACE") - CORNER = DomainInfo(name="CORNER") - CURVE = DomainInfo(name="CURVE") - INSTANCE = DomainInfo(name="INSTNANCE") - LAYER = DomainInfo(name="LAYER") + POINT = DomainType(name="POINT") + EDGE = DomainType(name="EDGE") + FACE = DomainType(name="FACE") + CORNER = DomainType(name="CORNER") + CURVE = DomainType(name="CURVE") + INSTANCE = DomainType(name="INSTNANCE") + LAYER = DomainType(name="LAYER") @dataclass @@ -42,6 +46,9 @@ class AttributeType: dtype: Type dimensions: tuple + def __str__(self) -> str: + return self.type_name + # https://docs.blender.org/api/current/bpy_types_enum_items/attribute_type_items.html#rna-enum-attribute-type-items class AttributeTypes(Enum): @@ -102,22 +109,22 @@ def guess_atype_from_array(array: np.ndarray) -> AttributeType: # for 1D arrays we we use the float, int of boolean attribute types if shape == (n_row, 1) or shape == (n_row,): if np.issubdtype(dtype, np.int_): - return AttributeTypes.INT.value + return AttributeTypes.INT elif np.issubdtype(dtype, np.float_): - return AttributeTypes.FLOAT.value + return AttributeTypes.FLOAT elif np.issubdtype(dtype, np.bool_): - return AttributeTypes.BOOLEAN.value + return AttributeTypes.BOOLEAN # for 2D arrays we use the float_vector, float_color, float4x4 attribute types elif shape == (n_row, 4, 4): - return AttributeTypes.FLOAT4X4.value + return AttributeTypes.FLOAT4X4 elif shape == (n_row, 3): - return AttributeTypes.FLOAT_VECTOR.value + return AttributeTypes.FLOAT_VECTOR elif shape == (n_row, 4): - return AttributeTypes.FLOAT_COLOR.value + return AttributeTypes.FLOAT_COLOR # if we didn't match against anything return float - return AttributeTypes.FLOAT.value + return AttributeTypes.FLOAT class Attribute: @@ -184,3 +191,105 @@ def __str__(self): return "Attribute: {}, type: {}, size: {}".format( self.attribute.name, self.type_name, self.shape ) + + +def store_named_attribute( + obj: bpy.types.Object, + data: np.ndarray, + name: str, + atype: str | AttributeType | None = None, + domain: str | DomainType = Domains.POINT, + overwrite: bool = True, +) -> bpy.types.Attribute: + """ + Adds and sets the values of an attribute on the object. + + Parameters + ---------- + obj : bpy.types.Object + The Blender object. + name : str + The name of the attribute. + data : np.ndarray + The attribute data as a numpy array. + atype : str, AttributeType, optional + The attribute type to store the data as. One of the AttributeType enums or a string + of the same name. + 'FLOAT_VECTOR', 'FLOAT_COLOR', 'FLOAT4X4', 'QUATERNION', 'FLOAT', 'INT', 'BOOLEAN' + domain : str, optional + The domain of the attribute. Defaults to 'POINT'. Currently, only 'POINT', 'EDGE', + and 'FACE' have been tested. + overwrite : bool + Setting to false will create a new attribute if the given name is already an + attribute on the mesh. + + Returns + ------- + bpy.types.Attribute + The added attribute. + """ + + if isinstance(atype, str): + try: + atype = AttributeTypes[atype] + except KeyError: + raise ValueError( + f"Given data type {atype=} does not match any of the possible attribute types: {list(AttributeTypes)=}" + ) + + if atype is None: + atype = guess_atype_from_array(data) + + attribute = obj.data.attributes.get(name) # type: ignore + if not attribute or not overwrite: + attribute = obj.data.attributes.new(name, atype.value.type_name, str(domain)) + + if len(data) != len(attribute.data): + raise AttributeMismatchError( + f"Data length {len(data)}, dimensions {data.shape} does not equal the size of the target domain {domain}, len={len(attribute.data)=}" + ) + + # the 'foreach_set' requires a 1D array, regardless of the shape of the attribute + # it also requires the order to be 'c' or blender might crash!! + attribute.data.foreach_set(atype.value.value_name, data.reshape(-1)) + + # The updating of data doesn't work 100% of the time (see: + # https://projects.blender.org/blender/blender/issues/118507) so this resetting of a + # single vertex is the current fix. Not great as I can see it breaking when we are + # missing a vertex - but for now we shouldn't be dealing with any situations where this + # is the case For now we will set a single vert to it's own position, which triggers a + # proper refresh of the object data. + try: + obj.data.vertices[0].co = obj.data.vertices[0].co # type: ignore + except AttributeError: + obj.data.update() # type: ignore + + return attribute + + +def named_attribute( + obj: bpy.types.Object, name="position", evaluate=False +) -> np.ndarray: + """ + Get the named attribute data from the object, optionally evaluating modifiers first. + + Parameters: + object (bpy.types.Object): The Blender object. + name (str, optional): The name of the attribute. Defaults to 'position'. + + Returns: + np.ndarray: The attribute data as a numpy array. + """ + if evaluate: + obj = evaluate_object(obj) + verbose = False + try: + attr = Attribute(obj.data.attributes[name]) + except KeyError: + message = f"The selected attribute '{name}' does not exist on the mesh." + if verbose: + message += f"Possible attributes are: {obj.data.attributes.keys()}" + + raise AttributeError(message) + + return attr.as_array() diff --git a/molecularnodes/blender/object.py b/molecularnodes/blender/databpy/object.py similarity index 50% rename from molecularnodes/blender/object.py rename to molecularnodes/blender/databpy/object.py index 568261e1..f81f02c2 100644 --- a/molecularnodes/blender/object.py +++ b/molecularnodes/blender/databpy/object.py @@ -1,5 +1,13 @@ import bpy import numpy as np +from typing import Union, Optional, Type +from .attribute import ( + store_named_attribute, + named_attribute, + Domains, + DomainType, +) +from mathutils import Matrix class ObjectTracker: @@ -63,22 +71,18 @@ def latest(self): return self.new_objects()[-1] -def evaluate_object(obj): - "Return an object which has the modifiers evaluated." - obj.update_tag() - return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) - def create_object( - vertices: np.ndarray = [], - edges: np.ndarray = [], - faces: np.ndarray = [], + vertices: np.ndarray | None = None, + edges: np.ndarray | None = None, + faces: np.ndarray | None = None, name: str = "NewObject", - collection: bpy.types.Collection = None, + collection: bpy.types.Collection | None = None, ) -> bpy.types.Object: """ - Create a new Blender object, initialised with locations for each vertex. + Create a new Blender object and corresponding mesh. - If edges and faces are supplied then these are also created on the mesh. + Vertices are created for each row in the vertices array. If edges and / or faces are created then they are also + initialized but default to None. Parameters ---------- @@ -99,17 +103,75 @@ def create_object( The created object. """ mesh = bpy.data.meshes.new(name) - + if vertices is None: + vertices = [] + if edges is None: + edges = [] + if faces is None: + faces = [] mesh.from_pydata(vertices=vertices, edges=edges, faces=faces) - - object = bpy.data.objects.new(name, mesh) - + obj = bpy.data.objects.new(name, mesh) if not collection: - # Add the object to the scene if no collection is specified collection = bpy.data.collections["Collection"] + collection.objects.link(obj) + return obj + + +class BlenderObject: + """ + A convenience class for working with Blender objects + """ + + def __init__(self, object: bpy.types.Object): + self.object = object + + def store_named_attribute( + self, + data: np.ndarray, + name: str, + dtype: Type | None = None, + domain: str | DomainType = Domains.POINT, + ) -> None: + store_named_attribute(self.object, data=data, name=name, domain=domain) - collection.objects.link(object) + def named_attribute(self, name: str, evaluate: bool = False) -> np.ndarray: + return named_attribute(self.object, name=name, evaluate=evaluate) - object["type"] = "molecule" + def transform_origin(self, matrix: Matrix) -> None: + self.object.matrix_local = matrix * self.object.matrix_world - return object + def transform_points(self, matrix: Matrix) -> None: + self.position = self.position * matrix + + @property + def selected(self) -> np.ndarray: + return named_attribute(self.object, ".select_vert") + + @property + def position(self) -> np.ndarray: + return named_attribute(self.object, name="position", evaluate=False) + + @position.setter + def position(self, value: np.ndarray) -> None: + store_named_attribute(self.object, "position", value) + + def selected_positions(self, mask: Optional[np.ndarray] = None) -> np.ndarray: + if mask is not None: + return self.position[np.logical_and(self.selected, mask)] + + return self.position[self.selected] + + +def bob(object: Union[bpy.types.Object, BlenderObject]) -> BlenderObject: + """ + Convenience function to convert a Blender object to a BlenderObject + """ + if isinstance(object, BlenderObject): + return object + elif isinstance(object, bpy.types.Object): + return BlenderObject(object) + else: + raise ValueError( + f"Unknown object type: {object=}" + "Expected bpy.types.Object or BlenderObject" + ) diff --git a/molecularnodes/blender/databpy/utils.py b/molecularnodes/blender/databpy/utils.py new file mode 100644 index 00000000..5b969c2d --- /dev/null +++ b/molecularnodes/blender/databpy/utils.py @@ -0,0 +1,7 @@ +import bpy + + +def evaluate_object(obj: bpy.types.Object): + "Return an object which has the modifiers evaluated." + obj.update_tag() + return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) diff --git a/molecularnodes/blender/mesh.py b/molecularnodes/blender/mesh.py index 1e6a7fde..6059fc7c 100644 --- a/molecularnodes/blender/mesh.py +++ b/molecularnodes/blender/mesh.py @@ -4,120 +4,26 @@ import numpy as np from . import coll, nodes -from .attribute import ( +from .databpy.attribute import ( Attribute, AttributeMismatchError, AttributeTypes, guess_atype_from_array, + store_named_attribute, + named_attribute, ) -from .object import ObjectTracker, create_object, evaluate_object +from .databpy.object import ObjectTracker, create_object +from .databpy.utils import evaluate_object -def centre(array: np.ndarray): - return np.mean(array, axis=0) +def centre(position: np.ndarray): + "Calculate the centroid of the vectors" + return np.mean(position, axis=0) -def centre_weighted(array: np.ndarray, weight: np.ndarray): - return np.sum(array * weight.reshape((len(array), 1)), axis=0) / np.sum(weight) - -def store_named_attribute( - obj: bpy.types.Object, - data: np.ndarray, - name: str, - data_type: Optional[str] = None, - domain: str = "POINT", - overwrite: bool = True, -) -> bpy.types.Attribute: - """ - Adds and sets the values of an attribute on the object. - - Parameters - ---------- - obj : bpy.types.Object - The Blender object. - name : str - The name of the attribute. - data : np.ndarray - The attribute data as a numpy array. - type : str, optional - The data type of the attribute. Defaults to None. Possible values are: - 'FLOAT_VECTOR', 'FLOAT_COLOR', 'FLOAT4X4', 'QUATERNION', 'FLOAT', 'INT', 'BOOLEAN' - domain : str, optional - The domain of the attribute. Defaults to 'POINT'. Currently, only 'POINT', 'EDGE', - and 'FACE' have been tested. - overwrite : bool, optional - Whether to overwrite an existing attribute with the same name. Defaults to True. - - Returns - ------- - bpy.types.Attribute - The added attribute. - """ - - if isinstance(data_type, str): - try: - atype = AttributeTypes[data_type].value - except KeyError: - raise ValueError( - f"Given data type {data_type=} does not match any of the possible attribute types: {list(AttributeTypes)=}" - ) - - if data_type is None: - atype = guess_atype_from_array(data) - - attribute = obj.data.attributes.get(name) # type: ignore - if not attribute or not overwrite: - attribute = obj.data.attributes.new(name, atype.type_name, domain) - - if len(data) != len(attribute.data): - raise AttributeMismatchError( - f"Data length {len(data)}, dimensions {data.shape} does not equal the size of the target domain {domain}, len={len(attribute.data)=}" - ) - - # the 'foreach_set' requires a 1D array, regardless of the shape of the attribute - # it also requires the order to be 'c' or blender might crash!! - attribute.data.foreach_set(atype.value_name, data.reshape(-1)) - - # The updating of data doesn't work 100% of the time (see: - # https://projects.blender.org/blender/blender/issues/118507) so this resetting of a - # single vertex is the current fix. Not great as I can see it breaking when we are - # missing a vertex - but for now we shouldn't be dealing with any situations where this - # is the case For now we will set a single vert to it's own position, which triggers a - # proper refresh of the object data. - try: - obj.data.vertices[0].co = obj.data.vertices[0].co # type: ignore - except AttributeError: - obj.data.update() # type: ignore - - return attribute - - -def named_attribute( - obj: bpy.types.Object, name="position", evaluate=False -) -> np.ndarray: - """ - Get the named attribute data from the object, optionally evaluating modifiers first. - - Parameters: - object (bpy.types.Object): The Blender object. - name (str, optional): The name of the attribute. Defaults to 'position'. - - Returns: - np.ndarray: The attribute data as a numpy array. - """ - if evaluate: - obj = evaluate_object(obj) - verbose = False - try: - attr = Attribute(obj.data.attributes[name]) - except KeyError: - message = f"The selected attribute '{name}' does not exist on the mesh." - if verbose: - message += f"Possible attributes are: {obj.data.attributes.keys()}" - - raise AttributeError(message) - - return attr.as_array() +def centre_weighted(position: np.ndarray, weight: np.ndarray): + "Calculate the weighted centroid of the vectors" + return np.sum(position * weight.reshape((-1, 1)), axis=0) / np.sum(weight) def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types.Object: @@ -150,8 +56,6 @@ def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types. return obj - - def evaluate_using_mesh(obj): """ Evaluate the object using a debug object. Some objects can't currently have their @@ -194,10 +98,10 @@ def create_data_object(array, collection=None, name="DataObject", world_scale=0. obj = create_object(locations, collection=collection, name=name) attributes = [ - ("rotation", "QUATERNION"), - ("assembly_id", "INT"), - ("chain_id", "INT"), - ("transform_id", "INT"), + ("rotation", AttributeTypes.QUATERNION), + ("assembly_id", AttributeTypes.INT), + ("chain_id", AttributeTypes.INT), + ("transform_id", AttributeTypes.INT), ] for column, type in attributes: @@ -210,6 +114,6 @@ def create_data_object(array, collection=None, name="DataObject", world_scale=0. if np.issubdtype(data.dtype, str): data = np.unique(data, return_inverse=True)[1] - store_named_attribute(obj=obj, data=data, name=column, data_type=type) + store_named_attribute(obj=obj, data=data, name=column, atype=type) return obj diff --git a/molecularnodes/blender/node.py b/molecularnodes/blender/node.py deleted file mode 100644 index eaf0ed07..00000000 --- a/molecularnodes/blender/node.py +++ /dev/null @@ -1,37 +0,0 @@ -from abc import ABCMeta -from typing import Optional, Any -import warnings -import time -import numpy as np -import bpy - - -class Node(metaclass=ABCMeta): - def __init__(self, node: bpy.types.Node, chain=[]): - - self.node = node - self.group = node.id_data - self.chain = chain - - @property - def location(self): - return np.array(self.node.location) - - def new(self, name): - "Add a new node to the node group." - try: - return self.group.nodes.new(f'GeometryNode{name}') - except RuntimeError: - return self.group.nodes.new(f'ShaderNode{name}') - - def link(self, name, linkto=0, linkfrom=0): - "Create a new node along in the chain and create a link to it. Return the new node." - new_node = self.new(name) - new_node.location = self.location + np.array((200, 0)) - - self.group.links.new( - self.node.outputs[linkfrom], - new_node.inputs[linkto] - ) - - return Node(new_node, chain=self.chain + [self]) diff --git a/molecularnodes/entities/ensemble/cellpack.py b/molecularnodes/entities/ensemble/cellpack.py index 8ebc28f4..32bd7a11 100644 --- a/molecularnodes/entities/ensemble/cellpack.py +++ b/molecularnodes/entities/ensemble/cellpack.py @@ -8,6 +8,7 @@ from .cif import OldCIF from ..molecule import molecule from ... import blender as bl +from ...blender.databpy import store_named_attribute, AttributeTypes from ... import color @@ -68,11 +69,11 @@ def _create_object_instances( ) colors = np.tile(color.random_rgb(i), (len(chain_atoms), 1)) - bl.mesh.store_named_attribute( + store_named_attribute( obj=obj, data=colors, name="Color", - data_type="FLOAT_COLOR", + atype=AttributeTypes.FLOAT_COLOR, ) if node_setup: diff --git a/molecularnodes/entities/ensemble/star.py b/molecularnodes/entities/ensemble/star.py index c7180fb2..65209efa 100644 --- a/molecularnodes/entities/ensemble/star.py +++ b/molecularnodes/entities/ensemble/star.py @@ -7,6 +7,7 @@ from PIL import Image from ... import blender as bl +from ...blender.databpy import AttributeTypes, store_named_attribute from .ensemble import Ensemble @@ -222,12 +223,11 @@ def create_object(self, name="StarFileObject", node_setup=True, world_scale=0.01 col_type = self.data[col].dtype # If col_type is numeric directly add if np.issubdtype(col_type, np.number): - bl.mesh.store_named_attribute( + store_named_attribute( obj=blender_object, name=col, data=self.data[col].to_numpy().reshape(-1), - data_type="FLOAT", - domain="POINT", + atype=AttributeTypes.FLOAT, ) # If col_type is object, convert to category and add integer values @@ -235,12 +235,8 @@ def create_object(self, name="StarFileObject", node_setup=True, world_scale=0.01 codes = ( self.data[col].astype("category").cat.codes.to_numpy().reshape(-1) ) - bl.mesh.store_named_attribute( - obj=blender_object, - data=codes, - name=col, - data_type="INT", - domain="POINT", + store_named_attribute( + obj=blender_object, data=codes, name=col, atype=AttributeTypes.INT ) # Add the category names as a property to the blender object blender_object[f"{col}_categories"] = list( diff --git a/molecularnodes/entities/entity.py b/molecularnodes/entities/entity.py index 41011db1..2badc132 100644 --- a/molecularnodes/entities/entity.py +++ b/molecularnodes/entities/entity.py @@ -2,6 +2,8 @@ import bpy from uuid import uuid1 from .. import blender as bl +from ..blender import databpy as db +from ..blender.databpy import AttributeTypes, AttributeType, Domains, DomainType import warnings import numpy as np @@ -90,38 +92,24 @@ def named_attribute(self, name="position", evaluate=False) -> np.ndarray | None: "No object yet created. Use `create_object()` to create a corresponding object." ) return None - return bl.mesh.named_attribute(self.object, name=name, evaluate=evaluate) + return db.named_attribute(self.object, name=name, evaluate=evaluate) def set_position(self, positions: np.ndarray) -> None: "A slightly optimised way to set the positions of the object's mesh" - obj = self.object - attribute = obj.data.attributes["position"] - n_points = len(attribute.data) - if positions.shape != (n_points, 3): - raise AttributeError( - f"Expected an array of dimension {(n_points, 3)} to set the position" - / f"but got {positions.shape=}" - ) - - # actually set the data for the positions - attribute.data.foreach_set("vector", positions.reshape(-1)) - # trigger a depsgraph update. The second method is better but bugs out sometimes - # so we try the first method initially - try: - obj.data.vertices[0].co = obj.data.vertices[0].co # type: ignore - except AttributeError: - obj.data.update() # type: ignore + self.store_named_attribute( + data=positions, name="position", atype=AttributeTypes.FLOAT_VECTOR + ) def set_boolean(self, boolean: np.ndarray, name="boolean") -> None: - self.store_named_attribute(boolean, name=name, data_type="BOOLEAN") + self.store_named_attribute(boolean, name=name, atype=AttributeTypes.BOOLEAN) def store_named_attribute( self, data: np.ndarray, - name="NewAttribute", - data_type=None, - domain="POINT", - overwrite=True, + name: str = "NewAttribute", + atype: str | AttributeType | None = None, + domain: str | DomainType = Domains.POINT, + overwrite: bool = True, ): """ Set an attribute for the molecule. @@ -149,11 +137,11 @@ def store_named_attribute( "No object yet created. Use `create_object()` to create a corresponding object." ) return None - bl.mesh.store_named_attribute( + db.store_named_attribute( obj=self.object, data=data, name=name, - data_type=data_type, + atype=atype, domain=domain, overwrite=overwrite, ) diff --git a/molecularnodes/entities/molecule/molecule.py b/molecularnodes/entities/molecule/molecule.py index 8952b25f..a95a8bf5 100644 --- a/molecularnodes/entities/molecule/molecule.py +++ b/molecularnodes/entities/molecule/molecule.py @@ -13,6 +13,12 @@ from ... import blender as bl from ... import color, data, utils +from ...blender.databpy import ( + Domains, + AttributeTypes, + store_named_attribute, + named_attribute, +) from ..entity import MolecularEntity @@ -112,13 +118,13 @@ def centre(self, centre_type: str = "centroid") -> np.ndarray: :return: np.ndarray of shape (3,) user-defined centroid of all atoms in the Molecule object """ - positions = self.named_attribute(name="position") + position = self.named_attribute(name="position") if centre_type == "centroid": - return bl.mesh.centre(positions) + return bl.mesh.centre(position) elif centre_type == "mass": mass = self.named_attribute(name="mass") - return bl.mesh.centre_weighted(positions, mass) + return bl.mesh.centre_weighted(position, mass) else: raise ValueError( f"`{centre_type}` not a supported selection of ['centroid', 'mass']" @@ -296,7 +302,7 @@ def centre_array(atom_array, centre): atom_array.coord -= bl.mesh.centre(atom_array.coord) elif centre == "mass": atom_array.coord -= bl.mesh.centre_weighted( - array=atom_array.coord, weight=atom_array.mass + position=atom_array.coord, weight=atom_array.mass ) if centre in ["mass", "centroid"]: @@ -336,8 +342,12 @@ def centre_array(atom_array, centre): # 'AROMATIC_SINGLE' = 5, 'AROMATIC_DOUBLE' = 6, 'AROMATIC_TRIPLE' = 7 # https://www.biotite-python.org/apidoc/biotite.structure.BondType.html#biotite.structure.BondType if array.bonds: - bl.mesh.store_named_attribute( - obj, name="bond_type", data=bond_types, data_type="INT", domain="EDGE" + store_named_attribute( + obj=obj, + data=bond_types, + name="bond_type", + atype=AttributeTypes.INT, + domain=Domains.EDGE, ) # The attributes for the model are initially defined as single-use functions. This allows @@ -623,17 +633,18 @@ def att_sec_struct(): if verbose: start = time.process_time() try: - bl.mesh.store_named_attribute( - obj, + store_named_attribute( + obj=obj, data=att["value"](), name=att["name"], - data_type=att["type"], + atype=att["type"], domain=att["domain"], ) if verbose: print(f'Added {att["name"]} after {time.process_time() - start} s') - except: + except Exception as e: if verbose: + print(e) warnings.warn(f"Unable to add attribute: {att['name']}") print( f'Failed adding {att["name"]} after {time.process_time() - start} s' diff --git a/molecularnodes/entities/trajectory/dna.py b/molecularnodes/entities/trajectory/dna.py index e0b9cdab..e5873f71 100644 --- a/molecularnodes/entities/trajectory/dna.py +++ b/molecularnodes/entities/trajectory/dna.py @@ -2,6 +2,7 @@ import bpy from ... import color from ...blender import mesh, coll, nodes +from ...blender.databpy import store_named_attribute, AttributeTypes bpy.types.Scene.MN_import_oxdna_topology = bpy.props.StringProperty( name="Toplogy", @@ -176,7 +177,7 @@ def read_trajectory(filepath): return np.stack(frames) -def store_named_attributes_to_dna_mol(mol, frame, scale_dna=0.1): +def store_named_attributes_to_dna_mol(obj, frame, scale_dna=0.1): attributes = ("base_vector", "base_normal", "velocity", "angular_velocity") for i, att in enumerate(attributes): col_idx = np.array([3, 4, 5]) + i * 3 @@ -190,8 +191,8 @@ def store_named_attributes_to_dna_mol(mol, frame, scale_dna=0.1): if att != "angular_velocity": data *= scale_dna - mesh.store_named_attribute( - obj=mol, data=data, name=att, data_type="FLOAT_VECTOR" + store_named_attribute( + obj=obj, data=data, name=att, atype=AttributeTypes.FLOAT_VECTOR ) @@ -262,17 +263,17 @@ def load(top, traj, name="oxDNA", setup_nodes=True, world_scale=0.01): ) # adding additional toplogy information from the topology and frames objects - mesh.store_named_attribute( - obj=obj, data=topology[:, 1], name="res_name", data_type="INT" + store_named_attribute( + obj=obj, data=topology[:, 1], name="res_name", atype=AttributeTypes.INT ) - mesh.store_named_attribute( - obj=obj, data=topology[:, 0], name="chain_id", data_type="INT" + store_named_attribute( + obj=obj, data=topology[:, 0], name="chain_id", atype=AttributeTypes.INT ) - mesh.store_named_attribute( + store_named_attribute( obj=obj, data=color.color_chains_equidistant(topology[:, 0]), name="Color", - data_type="FLOAT_COLOR", + atype=AttributeTypes.FLOAT_COLOR, ) store_named_attributes_to_dna_mol(obj, trajectory[0], scale_dna=scale_dna) diff --git a/molecularnodes/entities/trajectory/trajectory.py b/molecularnodes/entities/trajectory/trajectory.py index b78c08db..274eced0 100644 --- a/molecularnodes/entities/trajectory/trajectory.py +++ b/molecularnodes/entities/trajectory/trajectory.py @@ -8,12 +8,13 @@ from ... import data from ..entity import MolecularEntity, ObjectMissingError from ...blender import coll, mesh, nodes, path_resolve +from ...blender import databpy as db from ...utils import lerp, correct_periodic_positions from .selections import Selection, TrajectorySelectionItem class Trajectory(MolecularEntity): - def __init__(self, universe: mda.Universe, world_scale: float=0.01): + def __init__(self, universe: mda.Universe, world_scale: float = 0.01): super().__init__() self.universe: mda.Universe = universe self.selections: Dict[str, Selection] = {} @@ -437,11 +438,10 @@ def create_object( for att_name, att in self._attributes_2_blender.items(): try: - mesh.store_named_attribute( - obj=obj, + self.store_named_attribute( data=att["value"], name=att_name, - data_type=att["type"], + atype=att["type"], domain=att["domain"], ) except Exception as e: diff --git a/tests/test_obj.py b/tests/test_obj.py index 0129b624..261b5f3a 100644 --- a/tests/test_obj.py +++ b/tests/test_obj.py @@ -2,6 +2,7 @@ import numpy as np import molecularnodes as mn from molecularnodes.blender import mesh +from molecularnodes.blender import databpy as db from .constants import data_dir mn.register() @@ -13,7 +14,7 @@ def test_creat_obj(): locations = [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]] bonds = [(0, 1), (1, 2), (2, 0)] name = "MyMesh" - my_object = mesh.create_object(locations, bonds, name=name) + my_object = db.create_object(locations, bonds, name=name) assert len(my_object.data.vertices) == 3 assert my_object.name == name @@ -25,9 +26,7 @@ def test_set_position(): pos_a = mol.named_attribute("position") - mol.store_named_attribute( - data=mol.named_attribute("position") + 10, name="position" - ) + mol.store_named_attribute(mol.named_attribute("position") + 10, name="position") pos_b = mol.named_attribute("position") print(f"{pos_a=}") @@ -38,25 +37,25 @@ def test_set_position(): def test_eval_mesh(): - a = mesh.create_object(np.zeros((3, 3))) + a = db.create_object(np.zeros((3, 3))) assert len(a.data.vertices) == 3 - b = mesh.create_object(np.zeros((5, 3))) + b = db.create_object(np.zeros((5, 3))) assert len(b.data.vertices) == 5 assert len(mesh.evaluate_using_mesh(b).data.vertices) == 5 def test_matrix_read_write(): - obj = mesh.create_object(np.zeros((5, 3))) + obj = db.create_object(np.zeros((5, 3))) arr = np.array((5, 4, 4), float) arr = np.random.rand(5, 4, 4) - mesh.store_named_attribute( - obj=obj, data=arr, name="test_matrix", data_type="FLOAT4X4" + db.store_named_attribute( + obj=obj, data=arr, name="test_matrix", atype=db.AttributeTypes.FLOAT4X4 ) assert np.allclose(mesh.named_attribute(obj, "test_matrix"), arr) arr2 = np.random.rand(5, 4, 4) - mesh.store_named_attribute(obj=obj, data=arr2, name="test_matrix2") + db.store_named_attribute(obj=obj, data=arr2, name="test_matrix2") assert not np.allclose(mesh.named_attribute(obj, "test_matrix2"), arr) diff --git a/tests/test_ops.py b/tests/test_ops.py index 100473fd..ed59ab3e 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -3,7 +3,7 @@ import numpy as np import molecularnodes as mn -from molecularnodes.blender.mesh import ObjectTracker, named_attribute +from molecularnodes.blender.databpy import ObjectTracker, named_attribute from .utils import sample_attribute, NumpySnapshotExtension from .constants import data_dir, codes, attributes From e30e248193d4394ef733cf52f52f8e487ff8a47a Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 31 Oct 2024 13:12:00 +0800 Subject: [PATCH 07/12] BlenderObject is subclassed by Molecule --- molecularnodes/blender/databpy/__init__.py | 2 +- molecularnodes/blender/databpy/object.py | 80 +++++++++----- molecularnodes/entities/ensemble/star.py | 31 +++--- molecularnodes/entities/entity.py | 102 ++---------------- molecularnodes/entities/molecule/molecule.py | 37 +++---- .../entities/trajectory/trajectory.py | 9 +- 6 files changed, 99 insertions(+), 162 deletions(-) diff --git a/molecularnodes/blender/databpy/__init__.py b/molecularnodes/blender/databpy/__init__.py index 185bd143..0752a2c0 100644 --- a/molecularnodes/blender/databpy/__init__.py +++ b/molecularnodes/blender/databpy/__init__.py @@ -1,4 +1,4 @@ -from .object import ObjectTracker, BlenderObject, bob, create_object +from .object import ObjectTracker, BlenderObject, create_object from .attribute import ( named_attribute, store_named_attribute, diff --git a/molecularnodes/blender/databpy/object.py b/molecularnodes/blender/databpy/object.py index f81f02c2..94e0569d 100644 --- a/molecularnodes/blender/databpy/object.py +++ b/molecularnodes/blender/databpy/object.py @@ -1,15 +1,22 @@ import bpy import numpy as np -from typing import Union, Optional, Type +from typing import Union, Optional from .attribute import ( - store_named_attribute, - named_attribute, + AttributeTypes, + AttributeType, Domains, DomainType, ) +from . import attribute from mathutils import Matrix +class ObjectMissingError(Exception): + def __init__(self, message): + self.message = message + super().__init__(self.message) + + class ObjectTracker: """ A context manager for tracking new objects in Blender. @@ -122,20 +129,36 @@ class BlenderObject: A convenience class for working with Blender objects """ - def __init__(self, object: bpy.types.Object): - self.object = object + def __init__(self, object: bpy.types.Object | None): + self._object = object def store_named_attribute( self, data: np.ndarray, name: str, - dtype: Type | None = None, + atype: str | AttributeType | None = None, domain: str | DomainType = Domains.POINT, ) -> None: - store_named_attribute(self.object, data=data, name=name, domain=domain) + attribute.store_named_attribute( + self.object, data=data, name=name, atype=atype, domain=domain + ) + return self + + @property + def object(self) -> bpy.types.Object: + obj = self._object + if obj is None: + raise ObjectMissingError( + "Object is deleted and unable to establish a connection with a new Blender Object." + ) + return obj + + @object.setter + def object(self, value: bpy.types.Object) -> None: + self._object = value def named_attribute(self, name: str, evaluate: bool = False) -> np.ndarray: - return named_attribute(self.object, name=name, evaluate=evaluate) + return attribute.named_attribute(self.object, name=name, evaluate=evaluate) def transform_origin(self, matrix: Matrix) -> None: self.object.matrix_local = matrix * self.object.matrix_world @@ -145,33 +168,38 @@ def transform_points(self, matrix: Matrix) -> None: @property def selected(self) -> np.ndarray: - return named_attribute(self.object, ".select_vert") + return self.named_attribute(".select_vert") + + @property + def name(self) -> str: + obj = self.object + if obj is None: + return None + + return obj.name + + @name.setter + def name(self, value: str) -> None: + obj = self.object + if obj is None: + raise ObjectMissingError + obj.name = value @property def position(self) -> np.ndarray: - return named_attribute(self.object, name="position", evaluate=False) + return self.named_attribute("position") @position.setter def position(self, value: np.ndarray) -> None: - store_named_attribute(self.object, "position", value) + self.store_named_attribute( + value, + name="position", + atype=AttributeTypes.FLOAT_VECTOR, + domain=Domains.POINT, + ) def selected_positions(self, mask: Optional[np.ndarray] = None) -> np.ndarray: if mask is not None: return self.position[np.logical_and(self.selected, mask)] return self.position[self.selected] - - -def bob(object: Union[bpy.types.Object, BlenderObject]) -> BlenderObject: - """ - Convenience function to convert a Blender object to a BlenderObject - """ - if isinstance(object, BlenderObject): - return object - elif isinstance(object, bpy.types.Object): - return BlenderObject(object) - else: - raise ValueError( - f"Unknown object type: {object=}" - "Expected bpy.types.Object or BlenderObject" - ) diff --git a/molecularnodes/entities/ensemble/star.py b/molecularnodes/entities/ensemble/star.py index 65209efa..f546efc6 100644 --- a/molecularnodes/entities/ensemble/star.py +++ b/molecularnodes/entities/ensemble/star.py @@ -7,7 +7,7 @@ from PIL import Image from ... import blender as bl -from ...blender.databpy import AttributeTypes, store_named_attribute +from ...blender.databpy import AttributeTypes, store_named_attribute, BlenderObject from .ensemble import Ensemble @@ -212,19 +212,20 @@ def _update_micrograph_texture(self, *_): self.star_node.inputs["Micrograph"].default_value = image_obj def create_object(self, name="StarFileObject", node_setup=True, world_scale=0.01): - blender_object = bl.mesh.create_object( - self.positions * world_scale, collection=bl.coll.mn(), name=name + bob = BlenderObject( + bl.mesh.create_object( + self.positions * world_scale, collection=bl.coll.mn(), name=name + ) ) - blender_object.mn["molecule_type"] = "star" + bob.object.mn["molecule_type"] = "star" # create attribute for every column in the STAR file for col in self.data.columns: col_type = self.data[col].dtype # If col_type is numeric directly add if np.issubdtype(col_type, np.number): - store_named_attribute( - obj=blender_object, + bob.store_named_attribute( name=col, data=self.data[col].to_numpy().reshape(-1), atype=AttributeTypes.FLOAT, @@ -235,21 +236,19 @@ def create_object(self, name="StarFileObject", node_setup=True, world_scale=0.01 codes = ( self.data[col].astype("category").cat.codes.to_numpy().reshape(-1) ) - store_named_attribute( - obj=blender_object, data=codes, name=col, atype=AttributeTypes.INT + bob.store_named_attribute( + data=codes, name=col, atype=AttributeTypes.INT ) # Add the category names as a property to the blender object - blender_object[f"{col}_categories"] = list( + bob.object[f"{col}_categories"] = list( self.data[col].astype("category").cat.categories ) - blender_object.mn.uuid = self.uuid + bob.object.mn.uuid = self.uuid if node_setup: - bl.nodes.create_starting_nodes_starfile( - blender_object, n_images=self.n_images - ) + bl.nodes.create_starting_nodes_starfile(bob.object, n_images=self.n_images) - blender_object["starfile_path"] = str(self.file_path) - self.object = blender_object + bob.object["starfile_path"] = str(self.file_path) + self.object = bob.object bpy.app.handlers.depsgraph_update_post.append(self._update_micrograph_texture) - return blender_object + return self.object diff --git a/molecularnodes/entities/entity.py b/molecularnodes/entities/entity.py index 2badc132..0464c6e3 100644 --- a/molecularnodes/entities/entity.py +++ b/molecularnodes/entities/entity.py @@ -2,39 +2,23 @@ import bpy from uuid import uuid1 from .. import blender as bl -from ..blender import databpy as db -from ..blender.databpy import AttributeTypes, AttributeType, Domains, DomainType +from ..blender.databpy import ( + AttributeTypes, + BlenderObject, +) import warnings import numpy as np -class ObjectMissingError(Exception): - def __init__(self, message): - self.message = message - super().__init__(self.message) - - -class MolecularEntity(metaclass=ABCMeta): +class MolecularEntity( + BlenderObject, + metaclass=ABCMeta, +): def __init__(self) -> None: self.uuid: str = str(uuid1()) self._object: bpy.types.Object | None self.type: str = "" - @property - def name(self) -> str: - obj = self.object - if obj is None: - return None - - return obj.name - - @name.setter - def name(self, value: str) -> None: - obj = self.object - if obj is None: - raise ObjectMissingError - obj.name = value - @property def object(self) -> bpy.types.Object | None: # If we don't have connection to an object, attempt to re-stablish to a new @@ -69,30 +53,9 @@ def object(self, 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: - """ - Get the value of an object for the data molecule. - - Parameters - ---------- - name : str, optional - The name of the attribute. Default is 'position'. - evaluate : bool, optional - Whether to first evaluate all node trees before getting the requsted attribute. - False (default) will sample the underlying atomic geometry, while True will - sample the geometry that is created through the Geometry Nodes tree. - - Returns - ------- - np.ndarray - The value of the attribute. - """ - if self.object is None: - warnings.warn( - "No object yet created. Use `create_object()` to create a corresponding object." - ) - return None - return db.named_attribute(self.object, name=name, evaluate=evaluate) + @property + def bob(self) -> BlenderObject: + return BlenderObject(self.object) def set_position(self, positions: np.ndarray) -> None: "A slightly optimised way to set the positions of the object's mesh" @@ -103,49 +66,6 @@ def set_position(self, positions: np.ndarray) -> None: def set_boolean(self, boolean: np.ndarray, name="boolean") -> None: self.store_named_attribute(boolean, name=name, atype=AttributeTypes.BOOLEAN) - def store_named_attribute( - self, - data: np.ndarray, - name: str = "NewAttribute", - atype: str | AttributeType | None = None, - domain: str | DomainType = Domains.POINT, - overwrite: bool = True, - ): - """ - Set an attribute for the molecule. - - Parameters - ---------- - data : np.ndarray - The data to be set as the attribute. Must be of length equal to the length - of the domain. - name : str, optional - The name of the new attribute. Default is 'NewAttribute'. - type : str, optional - If value is None (Default), the data type is inferred. The data type of the - attribute. Possbible values are ('FLOAT_VECTOR', 'FLOAT_COLOR", 'QUATERNION', - 'FLOAT', 'INT', 'BOOLEAN'). - domain : str, optional - The domain of the attribute. Default is 'POINT'. Possible values are - currently ['POINT', 'EDGE', 'FACE', 'SPLINE'] - overwrite : bool, optional - Whether to overwrite an existing attribute with the same name, or create a - new attribute with always a unique name. Default is True. - """ - if not self.object: - warnings.warn( - "No object yet created. Use `create_object()` to create a corresponding object." - ) - return None - db.store_named_attribute( - obj=self.object, - data=data, - name=name, - atype=atype, - domain=domain, - overwrite=overwrite, - ) - @classmethod def list_attributes(cls, evaluate=False) -> list | None: """ diff --git a/molecularnodes/entities/molecule/molecule.py b/molecularnodes/entities/molecule/molecule.py index a95a8bf5..7b03a7aa 100644 --- a/molecularnodes/entities/molecule/molecule.py +++ b/molecularnodes/entities/molecule/molecule.py @@ -13,12 +13,7 @@ from ... import blender as bl from ... import color, data, utils -from ...blender.databpy import ( - Domains, - AttributeTypes, - store_named_attribute, - named_attribute, -) +from ...blender.databpy import Domains, AttributeTypes, BlenderObject from ..entity import MolecularEntity @@ -330,11 +325,13 @@ def centre_array(atom_array, centre): bond_types = bonds_array[:, 2].copy(order="C") # creating the blender object and meshes and everything - obj = bl.mesh.create_object( - name=name, - collection=collection, - vertices=array.coord * world_scale, - edges=bond_idx, + bob = BlenderObject( + bl.mesh.create_object( + name=name, + collection=collection, + vertices=array.coord * world_scale, + edges=bond_idx, + ) ) # Add information about the bond types to the model on the edge domain @@ -342,8 +339,7 @@ def centre_array(atom_array, centre): # 'AROMATIC_SINGLE' = 5, 'AROMATIC_DOUBLE' = 6, 'AROMATIC_TRIPLE' = 7 # https://www.biotite-python.org/apidoc/biotite.structure.BondType.html#biotite.structure.BondType if array.bonds: - store_named_attribute( - obj=obj, + bob.store_named_attribute( data=bond_types, name="bond_type", atype=AttributeTypes.INT, @@ -403,7 +399,7 @@ def att_res_name(): res_nums.append(res_num) counter += 1 - obj["ligands"] = np.unique(other_res) + bob.object["ligands"] = np.unique(other_res) return np.array(res_nums) def att_chain_id(): @@ -633,8 +629,7 @@ def att_sec_struct(): if verbose: start = time.process_time() try: - store_named_attribute( - obj=obj, + bob.store_named_attribute( data=att["value"](), name=att["name"], atype=att["type"], @@ -652,10 +647,10 @@ def att_sec_struct(): coll_frames = None if frames: - coll_frames = bl.coll.frames(obj.name, parent=bl.coll.data()) + coll_frames = bl.coll.frames(bob.name, parent=bl.coll.data()) for i, frame in enumerate(frames): frame = bl.mesh.create_object( - name=obj.name + "_frame_" + str(i), + name=bob.name + "_frame_" + str(i), collection=coll_frames, vertices=frame.coord * world_scale, # vertices=frame.coord * world_scale - centroid @@ -669,9 +664,9 @@ def att_sec_struct(): # 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 try: - obj["chain_ids"] = list(np.unique(array.chain_id)) + bob.object["chain_ids"] = list(np.unique(array.chain_id)) except AttributeError: - obj["chain_ids"] = None + bob.object["chain_ids"] = None warnings.warn("No chain information detected.") - return obj, coll_frames + return bob.object, coll_frames diff --git a/molecularnodes/entities/trajectory/trajectory.py b/molecularnodes/entities/trajectory/trajectory.py index 274eced0..4d5df0df 100644 --- a/molecularnodes/entities/trajectory/trajectory.py +++ b/molecularnodes/entities/trajectory/trajectory.py @@ -6,7 +6,7 @@ import numpy.typing as npt from ... import data -from ..entity import MolecularEntity, ObjectMissingError +from ..entity import MolecularEntity from ...blender import coll, mesh, nodes, path_resolve from ...blender import databpy as db from ...utils import lerp, correct_periodic_positions @@ -44,8 +44,6 @@ def add_selection( ) -> TrajectorySelectionItem: "Adds a new selection with the given name, selection string and selection parameters." obj = self.object - # if obj is None: - # raise ObjectMissingError("Universe contains no object to add seleciton to") obj.mn_trajectory_selections.add() sel = obj.mn_trajectory_selections[-1] @@ -535,10 +533,7 @@ def _update_positions(self, frame): universe = self.universe frame_mapping = self.frame_mapping obj = self.object - if obj is None: - raise ObjectMissingError( - "Object is deleted and unable to establish a connection with a new Blender Object." - ) + subframes: int = obj.mn.subframes interpolate: bool = obj.mn.interpolate offset: int = obj.mn.offset From ecaf08caa9a1924481546a42c816baa557f78e4d Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 31 Oct 2024 13:40:10 +0800 Subject: [PATCH 08/12] cleanup some trajectory --- molecularnodes/entities/trajectory/trajectory.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/molecularnodes/entities/trajectory/trajectory.py b/molecularnodes/entities/trajectory/trajectory.py index 4d5df0df..3cb92dc3 100644 --- a/molecularnodes/entities/trajectory/trajectory.py +++ b/molecularnodes/entities/trajectory/trajectory.py @@ -144,19 +144,12 @@ def positions(self) -> np.ndarray: @property def bonds(self) -> np.ndarray: + # the code to remap indices for a selection was removed as we don't subset the trajectory anymore + # when importing it, everything is imported and the selections just update if hasattr(self.atoms, "bonds"): - bond_indices = self.atoms.bonds.indices - atm_indices = self.atoms.indices - bond_filtering = np.all(np.isin(bond_indices, atm_indices), axis=1) - bond_indices = bond_indices[bond_filtering] - index_map = { - index: i for i, index in enumerate(self.universe.atoms.indices) - } - - bonds = [[index_map[bond[0]], index_map[bond[1]]] for bond in bond_indices] + return self.atoms.bonds.indices else: - bonds = [] - return np.array(bonds) + return None @property def elements(self) -> List[str]: From 499e5398066bfaa242d6a0bb8c7549f9ac6b1e3b Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 31 Oct 2024 13:50:12 +0800 Subject: [PATCH 09/12] cleanup --- .../entities/trajectory/trajectory.py | 51 +++++++++---------- 1 file changed, 23 insertions(+), 28 deletions(-) diff --git a/molecularnodes/entities/trajectory/trajectory.py b/molecularnodes/entities/trajectory/trajectory.py index 3cb92dc3..21603326 100644 --- a/molecularnodes/entities/trajectory/trajectory.py +++ b/molecularnodes/entities/trajectory/trajectory.py @@ -152,21 +152,21 @@ def bonds(self) -> np.ndarray: return None @property - def elements(self) -> List[str]: + def elements(self) -> np.ndarray: + if hasattr(self.atoms, "elements"): + return self.atoms.elements + try: - elements = self.atoms.elements.tolist() - except Exception: - try: - elements = [ - x - if x in data.elements.keys() - else mda.topology.guessers.guess_atom_element(x) - for x in self.atoms.names - ] + guessed_elements = [ + x + if x in data.elements.keys() + else mda.topology.guessers.guess_atom_element(x) + for x in self.atoms.names + ] + return np.array(guessed_elements) - except Exception: - elements = ["X"] * self.atoms.n_atoms - return elements + except Exception: + return np.repeat("X", self.n_atoms) @property def atomic_number(self) -> np.ndarray: @@ -179,7 +179,6 @@ def atomic_number(self) -> np.ndarray: @property def vdw_radii(self) -> np.ndarray: - # pm to Angstrom return ( np.array( [ @@ -187,25 +186,21 @@ def vdw_radii(self) -> np.ndarray: for element in self.elements ] ) - * 0.01 - * self.world_scale + * 0.01 # pm to Angstrom + * self.world_scale # Angstrom to world scale ) @property def mass(self) -> np.ndarray: # units: daltons - try: - masses = np.array([x.mass for x in self.atoms]) - except mda.exceptions.NoDataError: - masses = np.array( - [ - data.elements.get(element, {"standard_mass": 0}).get( - "standard_mass" - ) - for element in self.elements - ] - ) - return masses + if hasattr(self.atoms, "masses"): + return np.array([x.mass for x in self.atoms]) + else: + masses = [ + data.elements.get(element, {"standard_mass": 0}).get("standard_mass") + for element in self.elements + ] + return np.array(masses) @property def res_id(self) -> np.ndarray: From 6ae137fc3ac7a94aa3142b7358ff6157b66e66cd Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 31 Oct 2024 17:45:12 +0800 Subject: [PATCH 10/12] more refactoring --- molecularnodes/blender/__init__.py | 20 +----- .../blender/{databpy => bpyd}/__init__.py | 0 .../blender/{databpy => bpyd}/attribute.py | 36 ++++++++++- .../blender/{databpy => bpyd}/object.py | 64 ++++++++++++++----- molecularnodes/blender/databpy/utils.py | 7 -- molecularnodes/blender/mesh.py | 20 ++---- molecularnodes/blender/nodes.py | 3 +- molecularnodes/blender/utils.py | 11 ++++ molecularnodes/blender_manifest.toml | 23 ++++--- molecularnodes/entities/ensemble/cellpack.py | 2 +- molecularnodes/entities/ensemble/star.py | 2 +- molecularnodes/entities/entity.py | 2 +- molecularnodes/entities/molecule/molecule.py | 2 +- molecularnodes/entities/trajectory/dna.py | 2 +- .../entities/trajectory/trajectory.py | 2 +- molecularnodes/load.py | 0 molecularnodes/logger.py | 35 ---------- tests/test_density.py | 6 +- tests/test_load.py | 2 +- tests/test_nodes.py | 2 +- tests/test_obj.py | 21 +++--- tests/test_ops.py | 2 +- tests/test_select.py | 18 ++---- tests/test_trajectory.py | 3 +- tests/utils.py | 4 +- 25 files changed, 143 insertions(+), 146 deletions(-) rename molecularnodes/blender/{databpy => bpyd}/__init__.py (100%) rename molecularnodes/blender/{databpy => bpyd}/attribute.py (91%) rename molecularnodes/blender/{databpy => bpyd}/object.py (81%) delete mode 100644 molecularnodes/blender/databpy/utils.py create mode 100644 molecularnodes/blender/utils.py delete mode 100644 molecularnodes/load.py delete mode 100644 molecularnodes/logger.py diff --git a/molecularnodes/blender/__init__.py b/molecularnodes/blender/__init__.py index 94a1d5c5..a9e8ba98 100644 --- a/molecularnodes/blender/__init__.py +++ b/molecularnodes/blender/__init__.py @@ -1,19 +1 @@ -from pathlib import Path -from typing import Union -import bpy - - -def path_resolve(path: Union[str, Path]) -> Path: - if isinstance(path, str): - return Path(bpy.path.abspath(path)) - elif isinstance(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 +from .utils import path_resolve diff --git a/molecularnodes/blender/databpy/__init__.py b/molecularnodes/blender/bpyd/__init__.py similarity index 100% rename from molecularnodes/blender/databpy/__init__.py rename to molecularnodes/blender/bpyd/__init__.py diff --git a/molecularnodes/blender/databpy/attribute.py b/molecularnodes/blender/bpyd/attribute.py similarity index 91% rename from molecularnodes/blender/databpy/attribute.py rename to molecularnodes/blender/bpyd/attribute.py index c9a93fe1..41de672a 100644 --- a/molecularnodes/blender/databpy/attribute.py +++ b/molecularnodes/blender/bpyd/attribute.py @@ -1,11 +1,26 @@ from dataclasses import dataclass from enum import Enum from typing import Type -from .utils import evaluate_object - import bpy import numpy as np +from pathlib import Path + + +def evaluate_object(obj: bpy.types.Object): + "Return an object which has the modifiers evaluated." + obj.update_tag() + return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) + + +def path_resolve(path: str | Path) -> Path: + if isinstance(path, str): + return Path(bpy.path.abspath(path)) + elif isinstance(path, Path): + return Path(bpy.path.abspath(str(path))) + else: + raise ValueError(f"Unable to resolve path: {path}") + @dataclass class AttributeTypeInfo: @@ -161,6 +176,21 @@ def dtype(self) -> Type: def n_values(self) -> int: return np.prod(self.shape, dtype=int) + @classmethod + def from_object( + cls, + obj: bpy.types.Object, + name: str, + atype: AttributeType, + domain: DomainType, + ): + att = obj.data.get(name) + if att is None: + att = obj.data.attributes.new( + name=name, type=atype.value.type_name, domain=domain.value.name + ) + return Attribute(att) + def from_array(self, array: np.ndarray) -> None: """ Set the attribute data from a numpy array @@ -250,7 +280,7 @@ def store_named_attribute( ) # the 'foreach_set' requires a 1D array, regardless of the shape of the attribute - # it also requires the order to be 'c' or blender might crash!! + # so we have to flatten it first attribute.data.foreach_set(atype.value.value_name, data.reshape(-1)) # The updating of data doesn't work 100% of the time (see: diff --git a/molecularnodes/blender/databpy/object.py b/molecularnodes/blender/bpyd/object.py similarity index 81% rename from molecularnodes/blender/databpy/object.py rename to molecularnodes/blender/bpyd/object.py index 94e0569d..36e88876 100644 --- a/molecularnodes/blender/databpy/object.py +++ b/molecularnodes/blender/bpyd/object.py @@ -1,7 +1,8 @@ import bpy import numpy as np -from typing import Union, Optional +from typing import Optional from .attribute import ( + evaluate_object, AttributeTypes, AttributeType, Domains, @@ -124,25 +125,22 @@ def create_object( return obj +def active_object(context: bpy.types.Context = None) -> bpy.types.Object: + if context is None: + return bpy.context.active_object + + return context.active_object + + class BlenderObject: """ A convenience class for working with Blender objects """ - def __init__(self, object: bpy.types.Object | None): - self._object = object - - def store_named_attribute( - self, - data: np.ndarray, - name: str, - atype: str | AttributeType | None = None, - domain: str | DomainType = Domains.POINT, - ) -> None: - attribute.store_named_attribute( - self.object, data=data, name=name, atype=atype, domain=domain - ) - return self + def __init__(self, obj: bpy.types.Object | None): + if not isinstance(obj, bpy.types.Object): + raise ValueError(f"{obj} must be a Blender object of type bpy.types.Object") + self._object = obj @property def object(self) -> bpy.types.Object: @@ -157,6 +155,39 @@ def object(self) -> bpy.types.Object: def object(self, value: bpy.types.Object) -> None: self._object = value + def store_named_attribute( + self, + data: np.ndarray, + name: str, + atype: str | AttributeType | None = None, + domain: str | DomainType = Domains.POINT, + ) -> None: + """ + Parameters + ---------- + data : np.ndarray + The data to be stored as an attribute. + name : str + The name for the attribute. Will overwrite an already existing attribute. + atype : str or AttributeType or None, optional + The attribute type to store the data as. Either string or selection from the + AttributeTypes enum. None will attempt to infer the attribute type from the + input array. + domain : str or DomainType, optional + The domain to store the attribute on. Defaults to Domains.POINT. + + Returns + ------- + self + """ + attribute.store_named_attribute( + self.object, data=data, name=name, atype=atype, domain=domain + ) + return self + + def evaluate(self): + return BlenderObject(evaluate_object(self.object)) + def named_attribute(self, name: str, evaluate: bool = False) -> np.ndarray: return attribute.named_attribute(self.object, name=name, evaluate=evaluate) @@ -203,3 +234,6 @@ def selected_positions(self, mask: Optional[np.ndarray] = None) -> np.ndarray: return self.position[np.logical_and(self.selected, mask)] return self.position[self.selected] + + def __len__(self) -> int: + return len(self.object.data.vertices) diff --git a/molecularnodes/blender/databpy/utils.py b/molecularnodes/blender/databpy/utils.py deleted file mode 100644 index 5b969c2d..00000000 --- a/molecularnodes/blender/databpy/utils.py +++ /dev/null @@ -1,7 +0,0 @@ -import bpy - - -def evaluate_object(obj: bpy.types.Object): - "Return an object which has the modifiers evaluated." - obj.update_tag() - return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) diff --git a/molecularnodes/blender/mesh.py b/molecularnodes/blender/mesh.py index 6059fc7c..d3aa47a1 100644 --- a/molecularnodes/blender/mesh.py +++ b/molecularnodes/blender/mesh.py @@ -1,19 +1,9 @@ -from typing import Optional - import bpy import numpy as np from . import coll, nodes -from .databpy.attribute import ( - Attribute, - AttributeMismatchError, - AttributeTypes, - guess_atype_from_array, - store_named_attribute, - named_attribute, -) -from .databpy.object import ObjectTracker, create_object -from .databpy.utils import evaluate_object +from .bpyd.attribute import AttributeTypes, evaluate_object +from .bpyd.object import ObjectTracker, create_object, BlenderObject def centre(position: np.ndarray): @@ -95,7 +85,7 @@ def create_data_object(array, collection=None, name="DataObject", world_scale=0. if not collection: collection = coll.data() - obj = create_object(locations, collection=collection, name=name) + bob = BlenderObject(create_object(locations, collection=collection, name=name)) attributes = [ ("rotation", AttributeTypes.QUATERNION), @@ -114,6 +104,6 @@ def create_data_object(array, collection=None, name="DataObject", world_scale=0. if np.issubdtype(data.dtype, str): data = np.unique(data, return_inverse=True)[1] - store_named_attribute(obj=obj, data=data, name=column, atype=type) + bob.store_named_attribute(data=data, name=column, atype=type) - return obj + return bob.object diff --git a/molecularnodes/blender/nodes.py b/molecularnodes/blender/nodes.py index 60ad89dd..00dc1148 100644 --- a/molecularnodes/blender/nodes.py +++ b/molecularnodes/blender/nodes.py @@ -13,6 +13,7 @@ from .. import color, utils from . import mesh +from . import bpyd import re NODE_WIDTH = 180 @@ -725,7 +726,7 @@ def create_assembly_node_tree( "name": "assembly_id", "type": "NodeSocketInt", "min": 1, - "max": max(mesh.named_attribute(data_object, "assembly_id")), + "max": max(bpyd.named_attribute(data_object, "assembly_id")), "default": 1, }, ) diff --git a/molecularnodes/blender/utils.py b/molecularnodes/blender/utils.py new file mode 100644 index 00000000..99dd22c4 --- /dev/null +++ b/molecularnodes/blender/utils.py @@ -0,0 +1,11 @@ +import bpy +from pathlib import Path + + +def path_resolve(path: str | Path) -> Path: + if isinstance(path, str): + return Path(bpy.path.abspath(path)) + elif isinstance(path, Path): + return Path(bpy.path.abspath(str(path))) + else: + raise ValueError(f"Unable to resolve path: {path}") diff --git a/molecularnodes/blender_manifest.toml b/molecularnodes/blender_manifest.toml index 61ae37dc..7d6751bf 100644 --- a/molecularnodes/blender_manifest.toml +++ b/molecularnodes/blender_manifest.toml @@ -35,17 +35,16 @@ wheels = [ "./wheels/biotite-0.41.2-cp311-cp311-macosx_11_0_arm64.whl", "./wheels/biotite-0.41.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", "./wheels/biotite-0.41.2-cp311-cp311-win_amd64.whl", - "./wheels/colorama-0.4.6-py2.py3-none-any.whl", "./wheels/contourpy-1.3.0-cp311-cp311-macosx_10_9_x86_64.whl", "./wheels/contourpy-1.3.0-cp311-cp311-macosx_11_0_arm64.whl", "./wheels/contourpy-1.3.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", "./wheels/contourpy-1.3.0-cp311-cp311-win_amd64.whl", "./wheels/cycler-0.12.1-py3-none-any.whl", "./wheels/fasteners-0.19-py3-none-any.whl", - "./wheels/fonttools-4.54.0-cp311-cp311-macosx_10_9_universal2.whl", - "./wheels/fonttools-4.54.0-cp311-cp311-macosx_11_0_arm64.whl", - "./wheels/fonttools-4.54.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", - "./wheels/fonttools-4.54.0-cp311-cp311-win_amd64.whl", + "./wheels/fonttools-4.54.1-cp311-cp311-macosx_10_9_universal2.whl", + "./wheels/fonttools-4.54.1-cp311-cp311-macosx_11_0_arm64.whl", + "./wheels/fonttools-4.54.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", + "./wheels/fonttools-4.54.1-cp311-cp311-win_amd64.whl", "./wheels/joblib-1.4.2-py3-none-any.whl", "./wheels/kiwisolver-1.4.7-cp311-cp311-macosx_10_9_x86_64.whl", "./wheels/kiwisolver-1.4.7-cp311-cp311-macosx_11_0_arm64.whl", @@ -62,17 +61,17 @@ wheels = [ "./wheels/msgpack-1.1.0-cp311-cp311-macosx_11_0_arm64.whl", "./wheels/msgpack-1.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", "./wheels/msgpack-1.1.0-cp311-cp311-win_amd64.whl", - "./wheels/networkx-3.3-py3-none-any.whl", + "./wheels/networkx-3.4.2-py3-none-any.whl", "./wheels/packaging-24.1-py3-none-any.whl", "./wheels/pandas-2.2.3-cp311-cp311-macosx_10_9_x86_64.whl", "./wheels/pandas-2.2.3-cp311-cp311-macosx_11_0_arm64.whl", "./wheels/pandas-2.2.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", "./wheels/pandas-2.2.3-cp311-cp311-win_amd64.whl", - "./wheels/pillow-10.4.0-cp311-cp311-macosx_10_10_x86_64.whl", - "./wheels/pillow-10.4.0-cp311-cp311-macosx_11_0_arm64.whl", - "./wheels/pillow-10.4.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", - "./wheels/pillow-10.4.0-cp311-cp311-win_amd64.whl", - "./wheels/pyparsing-3.1.4-py3-none-any.whl", + "./wheels/pillow-11.0.0-cp311-cp311-macosx_10_10_x86_64.whl", + "./wheels/pillow-11.0.0-cp311-cp311-macosx_11_0_arm64.whl", + "./wheels/pillow-11.0.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", + "./wheels/pillow-11.0.0-cp311-cp311-win_amd64.whl", + "./wheels/pyparsing-3.2.0-py3-none-any.whl", "./wheels/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", "./wheels/pytz-2024.2-py2.py3-none-any.whl", "./wheels/scipy-1.14.1-cp311-cp311-macosx_10_13_x86_64.whl", @@ -82,7 +81,7 @@ wheels = [ "./wheels/six-1.16.0-py2.py3-none-any.whl", "./wheels/starfile-0.5.6-py3-none-any.whl", "./wheels/threadpoolctl-3.5.0-py3-none-any.whl", - "./wheels/tqdm-4.66.5-py3-none-any.whl", + "./wheels/tqdm-4.66.6-py3-none-any.whl", "./wheels/typing_extensions-4.12.2-py3-none-any.whl", "./wheels/tzdata-2024.2-py2.py3-none-any.whl", ] diff --git a/molecularnodes/entities/ensemble/cellpack.py b/molecularnodes/entities/ensemble/cellpack.py index 32bd7a11..dc274099 100644 --- a/molecularnodes/entities/ensemble/cellpack.py +++ b/molecularnodes/entities/ensemble/cellpack.py @@ -8,7 +8,7 @@ from .cif import OldCIF from ..molecule import molecule from ... import blender as bl -from ...blender.databpy import store_named_attribute, AttributeTypes +from ...blender.bpyd import store_named_attribute, AttributeTypes from ... import color diff --git a/molecularnodes/entities/ensemble/star.py b/molecularnodes/entities/ensemble/star.py index f546efc6..75e694ff 100644 --- a/molecularnodes/entities/ensemble/star.py +++ b/molecularnodes/entities/ensemble/star.py @@ -7,7 +7,7 @@ from PIL import Image from ... import blender as bl -from ...blender.databpy import AttributeTypes, store_named_attribute, BlenderObject +from ...blender.bpyd import AttributeTypes, store_named_attribute, BlenderObject from .ensemble import Ensemble diff --git a/molecularnodes/entities/entity.py b/molecularnodes/entities/entity.py index 0464c6e3..bb6ec38e 100644 --- a/molecularnodes/entities/entity.py +++ b/molecularnodes/entities/entity.py @@ -2,7 +2,7 @@ import bpy from uuid import uuid1 from .. import blender as bl -from ..blender.databpy import ( +from ..blender.bpyd import ( AttributeTypes, BlenderObject, ) diff --git a/molecularnodes/entities/molecule/molecule.py b/molecularnodes/entities/molecule/molecule.py index 7b03a7aa..b9229edd 100644 --- a/molecularnodes/entities/molecule/molecule.py +++ b/molecularnodes/entities/molecule/molecule.py @@ -13,7 +13,7 @@ from ... import blender as bl from ... import color, data, utils -from ...blender.databpy import Domains, AttributeTypes, BlenderObject +from ...blender.bpyd import Domains, AttributeTypes, BlenderObject from ..entity import MolecularEntity diff --git a/molecularnodes/entities/trajectory/dna.py b/molecularnodes/entities/trajectory/dna.py index e5873f71..36f5f56b 100644 --- a/molecularnodes/entities/trajectory/dna.py +++ b/molecularnodes/entities/trajectory/dna.py @@ -2,7 +2,7 @@ import bpy from ... import color from ...blender import mesh, coll, nodes -from ...blender.databpy import store_named_attribute, AttributeTypes +from ...blender.bpyd import store_named_attribute, AttributeTypes bpy.types.Scene.MN_import_oxdna_topology = bpy.props.StringProperty( name="Toplogy", diff --git a/molecularnodes/entities/trajectory/trajectory.py b/molecularnodes/entities/trajectory/trajectory.py index 21603326..4169c52b 100644 --- a/molecularnodes/entities/trajectory/trajectory.py +++ b/molecularnodes/entities/trajectory/trajectory.py @@ -8,7 +8,7 @@ from ... import data from ..entity import MolecularEntity from ...blender import coll, mesh, nodes, path_resolve -from ...blender import databpy as db +from ...blender import bpyd as db from ...utils import lerp, correct_periodic_positions from .selections import Selection, TrajectorySelectionItem diff --git a/molecularnodes/load.py b/molecularnodes/load.py deleted file mode 100644 index e69de29b..00000000 diff --git a/molecularnodes/logger.py b/molecularnodes/logger.py deleted file mode 100644 index 06bae942..00000000 --- a/molecularnodes/logger.py +++ /dev/null @@ -1,35 +0,0 @@ -import os -import logging -from .utils import ADDON_DIR - - -def start_logging(logfile_name: str = "side-packages-install") -> logging.Logger: - """ - Configure and start logging to a file. - - Parameters - ---------- - logfile_name : str, optional - The name of the log file. Defaults to 'side-packages-install'. - - Returns - ------- - logging.Logger - A Logger object that can be used to write log messages. - - This function sets up a logging configuration with a specified log file name and logging level. - The log file will be created in the `ADDON_DIR/logs` directory. If the directory - does not exist, it will be created. The function returns a Logger object that can be used to - write log messages. - - """ - # Create the logs directory if it doesn't exist - logs_dir = os.path.join(os.path.abspath(ADDON_DIR), "logs") - os.makedirs(logs_dir, exist_ok=True) - - # Set up logging configuration - logfile_path = os.path.join(logs_dir, f"{logfile_name}.log") - logging.basicConfig(filename=logfile_path, level=logging.INFO) - - # Return logger object - return logging.getLogger() diff --git a/tests/test_density.py b/tests/test_density.py index 6e470dc5..20b4c6dd 100644 --- a/tests/test_density.py +++ b/tests/test_density.py @@ -29,7 +29,7 @@ def density_file(): def test_density_load(density_file): obj = mn.entities.density.load(density_file).object evaluated = mn.blender.mesh.evaluate_using_mesh(obj) - pos = mn.blender.mesh.named_attribute(evaluated, "position") + pos = mn.blender.bpyd.named_attribute(evaluated, "position") assert len(pos) > 1000 @@ -50,7 +50,7 @@ def test_density_centered(density_file): obj = mn.entities.density.load(density_file, center=True, overwrite=True).object evaluated = mn.blender.mesh.evaluate_using_mesh(obj) - pos = mn.blender.mesh.named_attribute(evaluated, "position") + pos = mn.blender.bpyd.named_attribute(evaluated, "position") assert len(pos) > 1000 @@ -69,7 +69,7 @@ def test_density_invert(density_file): style_node.inputs["Threshold"].default_value = 0.01 evaluated = mn.blender.mesh.evaluate_using_mesh(obj) - pos = mn.blender.mesh.named_attribute(evaluated, "position") + pos = mn.blender.bpyd.named_attribute(evaluated, "position") # At this threshold after inverting we should have a cube the size of the volume assert pos[:, 0].max() > 2.0 assert pos[:, 1].max() > 2.0 diff --git a/tests/test_load.py b/tests/test_load.py index 8e17d178..c4afd7fc 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -64,7 +64,7 @@ def test_download_format(code, format): mol2 = o def verts(object): - return mn.blender.mesh.named_attribute(object, "position") + return mn.blender.bpyd.named_attribute(object, "position") assert np.isclose(verts(mol), verts(mol2)).all() diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 6bf0f61b..3b8b559d 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -266,7 +266,7 @@ def test_node_topology(snapshot_custom: NumpySnapshotExtension, code, node_name) group.links.new(output, input) - assert snapshot_custom == mn.blender.mesh.named_attribute( + assert snapshot_custom == mn.blender.bpyd.named_attribute( mol.object, "test_attribute", evaluate=True ) diff --git a/tests/test_obj.py b/tests/test_obj.py index 261b5f3a..d0379a1c 100644 --- a/tests/test_obj.py +++ b/tests/test_obj.py @@ -1,8 +1,7 @@ -import bpy import numpy as np import molecularnodes as mn from molecularnodes.blender import mesh -from molecularnodes.blender import databpy as db +from molecularnodes.blender import bpyd from .constants import data_dir mn.register() @@ -14,7 +13,7 @@ def test_creat_obj(): locations = [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]] bonds = [(0, 1), (1, 2), (2, 0)] name = "MyMesh" - my_object = db.create_object(locations, bonds, name=name) + my_object = bpyd.create_object(locations, bonds, name=name) assert len(my_object.data.vertices) == 3 assert my_object.name == name @@ -37,25 +36,25 @@ def test_set_position(): def test_eval_mesh(): - a = db.create_object(np.zeros((3, 3))) + a = bpyd.create_object(np.zeros((3, 3))) assert len(a.data.vertices) == 3 - b = db.create_object(np.zeros((5, 3))) + b = bpyd.create_object(np.zeros((5, 3))) assert len(b.data.vertices) == 5 assert len(mesh.evaluate_using_mesh(b).data.vertices) == 5 def test_matrix_read_write(): - obj = db.create_object(np.zeros((5, 3))) + obj = bpyd.create_object(np.zeros((5, 3))) arr = np.array((5, 4, 4), float) arr = np.random.rand(5, 4, 4) - db.store_named_attribute( - obj=obj, data=arr, name="test_matrix", atype=db.AttributeTypes.FLOAT4X4 + bpyd.store_named_attribute( + obj=obj, data=arr, name="test_matrix", atype=bpyd.AttributeTypes.FLOAT4X4 ) - assert np.allclose(mesh.named_attribute(obj, "test_matrix"), arr) + assert np.allclose(bpyd.named_attribute(obj, "test_matrix"), arr) arr2 = np.random.rand(5, 4, 4) - db.store_named_attribute(obj=obj, data=arr2, name="test_matrix2") + bpyd.store_named_attribute(obj=obj, data=arr2, name="test_matrix2") - assert not np.allclose(mesh.named_attribute(obj, "test_matrix2"), arr) + assert not np.allclose(bpyd.named_attribute(obj, "test_matrix2"), arr) diff --git a/tests/test_ops.py b/tests/test_ops.py index ed59ab3e..763c4d39 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -3,7 +3,7 @@ import numpy as np import molecularnodes as mn -from molecularnodes.blender.databpy import ObjectTracker, named_attribute +from molecularnodes.blender.bpyd import ObjectTracker, named_attribute from .utils import sample_attribute, NumpySnapshotExtension from .constants import data_dir, codes, attributes diff --git a/tests/test_select.py b/tests/test_select.py index 467b6d41..c896d295 100644 --- a/tests/test_select.py +++ b/tests/test_select.py @@ -1,6 +1,6 @@ import molecularnodes as mn +from molecularnodes.blender import bpyd from molecularnodes.blender import nodes -import bpy import numpy as np import pytest @@ -11,8 +11,6 @@ def create_debug_group(name="MolecularNodesDebugGroup"): group.links.new(info.outputs["Geometry"], group.nodes["Group Output"].inputs[0]) return group - return object.evaluated_get(dg) - custom_selections = [ ("1, 3, 5-7", np.array((1, 3, 5, 6, 7))), @@ -24,14 +22,13 @@ def create_debug_group(name="MolecularNodesDebugGroup"): @pytest.mark.parametrize("selection", custom_selections) def test_select_multiple_residues(selection): n_atoms = 100 - object = mn.blender.mesh.create_object(np.zeros((n_atoms, 3))) - mn.blender.mesh.store_named_attribute( - obj=object, + bob = bpyd.BlenderObject(bpyd.create_object(np.zeros((n_atoms, 3)))) + bob.store_named_attribute( data=np.arange(n_atoms) + 1, name="res_id", ) - mod = nodes.get_mod(object) + mod = nodes.get_mod(bob.object) group = nodes.new_group(fallback=False) mod.node_group = group sep = group.nodes.new("GeometryNodeSeparateGeometry") @@ -41,9 +38,6 @@ def test_select_multiple_residues(selection): node_sel = nodes.add_custom(group, node_sel_group.name) group.links.new(node_sel.outputs["Selection"], sep.inputs["Selection"]) - vertices_count = len(mn.blender.mesh.evaluate_object(object).data.vertices) + vertices_count = len(bob.evaluate()) assert vertices_count == len(selection[1]) - assert ( - mn.blender.mesh.named_attribute(mn.blender.mesh.evaluate_object(object), "res_id") - == selection[1] - ).all() + assert (bob.evaluate().named_attribute("res_id") == selection[1]).all() diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py index f76eecad..0a06acf8 100644 --- a/tests/test_trajectory.py +++ b/tests/test_trajectory.py @@ -2,12 +2,11 @@ import os import pytest import molecularnodes as mn -from molecularnodes.blender.mesh import named_attribute import MDAnalysis as mda import numpy as np from .constants import data_dir -from .utils import sample_attribute, NumpySnapshotExtension +from .utils import NumpySnapshotExtension mn._test_register() diff --git a/tests/utils.py b/tests/utils.py index 923b16ab..f8b11e7f 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -61,7 +61,7 @@ def sample_attribute( random.seed(seed) if error: - attribute = mn.blender.mesh.named_attribute( + attribute = mn.blender.bpyd.named_attribute( obj=object, name=attribute, evaluate=evaluate ) length = len(attribute) @@ -77,7 +77,7 @@ def sample_attribute( return attribute[idx, :] else: try: - attribute = mn.blender.mesh.named_attribute( + attribute = mn.blender.bpyd.named_attribute( obj=object, name=attribute, evaluate=evaluate ) length = len(attribute) From 2048c8bbfb59a944806526bf7bb5b039106d77aa Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 31 Oct 2024 18:42:11 +0800 Subject: [PATCH 11/12] tinkering --- molecularnodes/blender/bpyd/attribute.py | 12 ++++---- molecularnodes/blender/bpyd/object.py | 38 +++++++++++++++++++++--- molecularnodes/blender/mesh.py | 8 ++--- molecularnodes/entities/entity.py | 4 +-- tests/test_nodes.py | 12 ++++---- 5 files changed, 51 insertions(+), 23 deletions(-) diff --git a/molecularnodes/blender/bpyd/attribute.py b/molecularnodes/blender/bpyd/attribute.py index 41de672a..13f272ea 100644 --- a/molecularnodes/blender/bpyd/attribute.py +++ b/molecularnodes/blender/bpyd/attribute.py @@ -7,12 +7,6 @@ from pathlib import Path -def evaluate_object(obj: bpy.types.Object): - "Return an object which has the modifiers evaluated." - obj.update_tag() - return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) - - def path_resolve(path: str | Path) -> Path: if isinstance(path, str): return Path(bpy.path.abspath(path)) @@ -297,6 +291,12 @@ def store_named_attribute( return attribute +def evaluate_object(obj: bpy.types.Object): + "Return an object which has the modifiers evaluated." + obj.update_tag() + return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) + + def named_attribute( obj: bpy.types.Object, name="position", evaluate=False ) -> np.ndarray: diff --git a/molecularnodes/blender/bpyd/object.py b/molecularnodes/blender/bpyd/object.py index 36e88876..be8fa840 100644 --- a/molecularnodes/blender/bpyd/object.py +++ b/molecularnodes/blender/bpyd/object.py @@ -2,7 +2,6 @@ import numpy as np from typing import Optional from .attribute import ( - evaluate_object, AttributeTypes, AttributeType, Domains, @@ -185,12 +184,43 @@ def store_named_attribute( ) return self - def evaluate(self): - return BlenderObject(evaluate_object(self.object)) - def named_attribute(self, name: str, evaluate: bool = False) -> np.ndarray: + """ + Retrieve a named attribute from the object. + + Optionally, evaluate the object before reading the named attribute + + Parameters + ---------- + name : str + Name of the attribute to get. + evaluate : bool, optional + Whether to evaluate the object before reading the attribute (default is False). + Returns + ------- + np.ndarray + The attribute read from the mesh as a numpy array. + """ return attribute.named_attribute(self.object, name=name, evaluate=evaluate) + def evaluate(self): + obj = self.object + obj.update_tag() + evluated_obj = obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) + return BlenderObject(evluated_obj) + + @property + def attributes(self): + return self.object.data.attributes + + @property + def vertices(self): + return self.object.data.vertices + + @property + def edges(self): + return self.object.data.edges + def transform_origin(self, matrix: Matrix) -> None: self.object.matrix_local = matrix * self.object.matrix_world diff --git a/molecularnodes/blender/mesh.py b/molecularnodes/blender/mesh.py index d3aa47a1..0157ae72 100644 --- a/molecularnodes/blender/mesh.py +++ b/molecularnodes/blender/mesh.py @@ -2,7 +2,7 @@ import numpy as np from . import coll, nodes -from .bpyd.attribute import AttributeTypes, evaluate_object +from .bpyd.attribute import AttributeTypes from .bpyd.object import ObjectTracker, create_object, BlenderObject @@ -67,13 +67,13 @@ def evaluate_using_mesh(obj): """ # create an empty mesh object. It's modifiers can be evaluated but some other # object types can't be currently through the API - debug_obj = create_object() - mod = nodes.get_mod(debug_obj) + bob = BlenderObject(create_object()) + mod = nodes.get_mod(bob.object) mod.node_group = nodes.create_debug_group() mod.node_group.nodes["Object Info"].inputs["Object"].default_value = obj # need to use 'evaluate' otherwise the modifiers won't be taken into account - return evaluate_object(debug_obj) + return bob.evaluate().object def create_data_object(array, collection=None, name="DataObject", world_scale=0.01): diff --git a/molecularnodes/entities/entity.py b/molecularnodes/entities/entity.py index bb6ec38e..4755407c 100644 --- a/molecularnodes/entities/entity.py +++ b/molecularnodes/entities/entity.py @@ -86,6 +86,6 @@ def list_attributes(cls, evaluate=False) -> list | None: warnings.warn("No object created") return None if evaluate: - return list(bl.mesh.evaluate_object(cls.object).data.attributes.keys()) + return list(cls.bob.evaluate().attributes.keys()) - return list(cls.object.data.attributes.keys()) + return list(cls.bob.attributes.keys()) diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 3b8b559d..18be2eb3 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -272,10 +272,8 @@ def test_node_topology(snapshot_custom: NumpySnapshotExtension, code, node_name) def test_topo_bonds(): - mol = mn.entities.fetch( - "1BNA", del_solvent=True, style=None, cache_dir=data_dir - ).object - group = nodes.get_mod(mol).node_group = nodes.new_group() + mol = mn.entities.fetch("1BNA", del_solvent=True, style=None, cache_dir=data_dir) + group = nodes.get_mod(mol.object).node_group = nodes.new_group() # add the node that will break bonds, set the cutoff to 0 node_break = nodes.add_custom(group, "Topology Break Bonds") @@ -283,8 +281,8 @@ def test_topo_bonds(): node_break.inputs["Cutoff"].default_value = 0 # compare the number of edges before and after deleting them with - bonds = mol.data.edges - no_bonds = mn.blender.mesh.evaluate_object(mol).data.edges + bonds = mol.object.data.edges + no_bonds = mol.evaluate().object.data.edges assert len(bonds) > len(no_bonds) assert len(no_bonds) == 0 @@ -292,5 +290,5 @@ def test_topo_bonds(): # are the same (other attributes will be different, but for now this is good) node_find = nodes.add_custom(group, "Topology Find Bonds") nodes.insert_last_node(group, node=node_find) - bonds_new = mn.blender.mesh.evaluate_object(mol).data.edges + bonds_new = mol.evaluate().edges assert len(bonds) == len(bonds_new) From af70fa716c0a5f96afbfe9c1667eb414e78ba535 Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Fri, 1 Nov 2024 08:03:09 +0800 Subject: [PATCH 12/12] cleanup --- molecularnodes/blender/bpyd/object.py | 9 +-------- molecularnodes/entities/trajectory/ui.py | 4 ++-- molecularnodes/ui/menu.py | 1 + molecularnodes/ui/panel.py | 8 +++++--- 4 files changed, 9 insertions(+), 13 deletions(-) diff --git a/molecularnodes/blender/bpyd/object.py b/molecularnodes/blender/bpyd/object.py index be8fa840..e5f66404 100644 --- a/molecularnodes/blender/bpyd/object.py +++ b/molecularnodes/blender/bpyd/object.py @@ -109,13 +109,13 @@ def create_object( bpy.types.Object The created object. """ - mesh = bpy.data.meshes.new(name) if vertices is None: vertices = [] if edges is None: edges = [] if faces is None: faces = [] + mesh = bpy.data.meshes.new(name) mesh.from_pydata(vertices=vertices, edges=edges, faces=faces) obj = bpy.data.objects.new(name, mesh) if not collection: @@ -124,13 +124,6 @@ def create_object( return obj -def active_object(context: bpy.types.Context = None) -> bpy.types.Object: - if context is None: - return bpy.context.active_object - - return context.active_object - - class BlenderObject: """ A convenience class for working with Blender objects diff --git a/molecularnodes/entities/trajectory/ui.py b/molecularnodes/entities/trajectory/ui.py index b719096e..bc69e1ba 100644 --- a/molecularnodes/entities/trajectory/ui.py +++ b/molecularnodes/entities/trajectory/ui.py @@ -62,12 +62,12 @@ class MN_OT_Reload_Trajectory(bpy.types.Operator): @classmethod def poll(cls, context): - obj = bl.active_object(context) + obj = context.active_object traj = get_session(context).trajectories.get(obj.mn.uuid) return not traj def execute(self, context): - obj = bl.active_object(context) + obj = context.active_object universe = mda.Universe(obj.mn.filepath_topology, obj.mn.filepath_trajectory) traj = Trajectory(universe) traj.object = obj diff --git a/molecularnodes/ui/menu.py b/molecularnodes/ui/menu.py index 8fc4c8b5..fe6b6578 100644 --- a/molecularnodes/ui/menu.py +++ b/molecularnodes/ui/menu.py @@ -106,6 +106,7 @@ def menu( op.description = f"Choose custom selections for {self.label}" else: raise ValueError(f"Data type currently not supported: {self.dtype}") + # test if the object has the currently tested property to enable operator row.enabled = bool(context.active_object.get(self.property_id)) diff --git a/molecularnodes/ui/panel.py b/molecularnodes/ui/panel.py index e7dc8c74..a04db68f 100644 --- a/molecularnodes/ui/panel.py +++ b/molecularnodes/ui/panel.py @@ -110,13 +110,15 @@ def panel_import(layout, context): chosen_panel[selection](col, scene) -def ui_from_node(layout, node): +def ui_from_node( + layout: bpy.types.UILayout, node: bpy.types.NodeGroup, context: bpy.types.Context +): """ Generate the UI for a particular node, which displays the relevant node inputs for user control in a panel, rather than through the node editor. """ col = layout.column(align=True) - ntree = bpy.context.active_object.modifiers["MolecularNodes"].node_group + ntree = context.active_object.modifiers["MolecularNodes"].node_group tree = node.node_tree.interface.items_tree @@ -213,7 +215,7 @@ def panel_object(layout, context): if mol_type == "star": layout.label(text="Ensemble") box = layout.box() - ui_from_node(box, nodes.get_star_node(object)) + ui_from_node(box, nodes.get_star_node(object), context=context) return None