diff --git a/MolecularNodes/obj.py b/MolecularNodes/obj.py index 1780e4c9..f3d6e377 100644 --- a/MolecularNodes/obj.py +++ b/MolecularNodes/obj.py @@ -1,10 +1,45 @@ import bpy import numpy as np -def create_object(name, collection, locations, bonds=[]): +def create_object(name: str, collection: bpy.types.Collection, locations, bonds=[]) -> bpy.types.Object: """ - Creates a mesh with the given name in the given collection, from the supplied - values for the locations of vertices, and if supplied, bonds as edges. + Create a mesh with the given name in the given collection, using the supplied + vertex locations and, if provided, bonds as edges. + + Parameters + ---------- + name : str + The name of the mesh object to be created. + collection : bpy.types.Collection + The collection to which the mesh object will be added. + locations : array-like + The list of vertex locations for the mesh, an nx3 np array of locations. + Each element in the list represents a 3D point (x, y, z) for a vertex. + bonds : list of tuples, optional + The list of vertex index pairs representing bonds as edges for the mesh. + Each tuple should contain two vertex indices (e.g., (index1, index2)). + + Returns + ------- + bpy.types.Object + The newly created mesh object. + + Notes + ----- + - The 'name' should be a unique identifier for the created mesh object. + - The 'locations' list should contain at least three 3D points to define a 3D triangle. + - If 'bonds' are not provided, the mesh will have no edges. + - If 'bonds' are provided, they should be valid vertex indices within the 'locations' list. + + Example + ------- + ```python + # Create a mesh object named "MyMesh" in the collection "MyCollection" + # with vertex locations and bond edges. + 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)] + my_object = create_object("MyMesh", bpy.data.collections['Collection'], locations, bonds) + ``` """ # create a new mesh mol_mesh = bpy.data.meshes.new(name) @@ -13,31 +48,93 @@ def create_object(name, collection, locations, bonds=[]): collection.objects.link(mol_object) return mol_object -def add_attribute(object, name, data, type = "FLOAT", domain = "POINT", add = True): - if not add: - return None - attribute = object.data.attributes.new(name, type, domain) - attribute.data.foreach_set('value', data) -def get_attribute(obj: bpy.types.Object, - att_name = 'position') -> np.array: - """Retrieve Attribute from Object as Numpy Array +def add_attribute(object: bpy.types.Object, name: str, data, type="FLOAT", domain="POINT"): + """ + Add an attribute to the given object's geometry on the given domain. + + Parameters + ---------- + object : bpy.types.Object + The object to which the attribute will be added. + name : str + The name of the attribute. + data : array-like + The data to be assigned to the attribute. For "FLOAT_VECTOR" attributes, it should be a 1D array + representing the vector data. + type : str, optional, default: "FLOAT" + The data type of the attribute. Possible values are "FLOAT", "FLOAT_VECTOR", "INT", or "BOOLEAN". + domain : str, optional, default: "POINT" + The domain to which the attribute is added. Possible values are "POINT" or other domains supported + by the object. + + Returns + ------- + Any + The newly created attribute, which can be further manipulated or used in the 3D environment. + + Notes + ----- + - The function supports adding both scalar and vector attributes. + - The "FLOAT_VECTOR" attribute requires the input data to be a 1D array, and it will be reshaped internally + to represent vectors with 3 components (x, y, z). + """ + + if type == "FLOAT_VECTOR": + att = object.data.attributes.new(name, type, domain) + # currently vectors have to be added as a 1d array. may change in the future + # but currently must be reshaped then added as a 'vector' but supplying a 1d array + vec_1d = data.reshape(len(data) * 3) + att.data.foreach_set('vector', vec_1d) + else: + att = object.data.attributes.new(name, type, domain) + att.data.foreach_set('value', data) + + return att + +def get_attribute(obj: bpy.types.Object, att_name='position') -> np.array: + """ + Retrieve an attribute from the object as a NumPy array. + + Parameters + ---------- + obj : bpy.types.Object + The Blender object from which the attribute will be retrieved. + att_name : str, optional + The name of the attribute to retrieve. Default is 'position'. + + Returns + ------- + np.array + The attribute data as a NumPy array. + + Notes + ----- + - This function retrieves the specified attribute from the object and returns it as a NumPy array. + - The function assumes that the attribute data type is one of ['INT', 'FLOAT', 'BOOLEAN', 'FLOAT_VECTOR']. + + Example + ------- + ```python + # Assuming 'my_object' is a Blender object with an attribute named 'my_attribute' + attribute_data = get_attribute(my_object, 'my_attribute') + ``` """ + + # Get the attribute from the object's mesh att = obj.to_mesh().attributes[att_name] + + # Map attribute values to a NumPy array based on the attribute data type if att.data_type in ['INT', 'FLOAT', 'BOOLEAN']: - d_type = { - 'INT': int, - 'FLOAT': float, - 'BOOLEAN': bool - } - att_array = np.array(list(map( - lambda x: x.value, - att.data.values() - )), dtype = d_type.get(att.data_type)) + # Define the mapping of Blender data types to NumPy data types + d_type = {'INT': int, 'FLOAT': float, 'BOOLEAN': bool} + # Convert attribute values to a NumPy array with the appropriate data type + att_array = np.array(list(map(lambda x: x.value, att.data.values())), dtype=d_type.get(att.data_type)) elif att.data_type == "FLOAT_VECTOR": - att_array = np.array(list(map( - lambda x: x.vector, - att.data.values() - ))) - - return att_array \ No newline at end of file + # Convert attribute vectors to a NumPy array + att_array = np.array(list(map(lambda x: x.vector, att.data.values()))) + else: + # Unsupported data type, return an empty NumPy array + att_array = np.array([]) + + return att_array diff --git a/MolecularNodes/star.py b/MolecularNodes/star.py index 9fd91e97..118bdffd 100644 --- a/MolecularNodes/star.py +++ b/MolecularNodes/star.py @@ -3,6 +3,8 @@ from . import coll from . import nodes from .obj import create_object +from .obj import add_attribute + bpy.types.Scene.mol_import_star_file_path = bpy.props.StringProperty( @@ -87,28 +89,24 @@ def load_star_file( target_meta=target_metadata)) obj = create_object(obj_name, coll.mn(), xyz * world_scale) - - # vectors have to be added as a 1D array currently - rotations = eulers.reshape(len(eulers) * 3) + # create the attribute and add the data for the rotations - attribute = obj.data.attributes.new('MOLRotation', 'FLOAT_VECTOR', 'POINT') - attribute.data.foreach_set('vector', rotations) + add_attribute(obj, 'MOLRotation', eulers, 'FLOAT_VECTOR', 'POINT') # create the attribute and add the data for the image id - attribute_imgid = obj.data.attributes.new('MOLImageId', 'INT', 'POINT') - attribute_imgid.data.foreach_set('value', image_id) + add_attribute(obj, 'MOLIMageId', image_id, 'INT', 'POINT') + # create attribute for every column in the STAR file for col in df.columns: - col_type = df[col].dtype + col_type = df[col].dtype # If col_type is numeric directly add if np.issubdtype(col_type, np.number): - attribute = obj.data.attributes.new(col, 'FLOAT', 'POINT') - attribute.data.foreach_set('value', df[col].to_numpy().reshape(-1)) + add_attribute(obj, col, df[col].to_numpy().reshape(-1), 'FLOAT', 'POINT') + # If col_type is object, convert to category and add integer values elif col_type == object: - attribute = obj.data.attributes.new(col, 'INT', 'POINT') - codes = df[col].astype('category').cat.codes - attribute.data.foreach_set('value', codes.to_numpy().reshape(-1)) + codes = df[col].astype('category').cat.codes.to_numpy().reshape(-1) + add_attribute(obj, col, codes, 'INT', 'POINT') # Add the category names as a property to the blender object obj[col + '_categories'] = list(df[col].astype('category').cat.categories) diff --git a/tests/test_load.py b/tests/test_load.py index 5ba03e79..cd24b71e 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -2,87 +2,7 @@ import os import pytest import MolecularNodes as mn - -# ensure we can successfully install all of the required pacakges -# def test_install_packages(): - # mn.pkg.install_all_packages() - # assert mn.pkg.is_current('biotite') == True - -def apply_mods(obj): - """ - Applies the modifiers on the modifier stack - - This will realise the computations inside of any Geometry Nodes modifiers, ensuring - that the result of the node trees can be compared by looking at the resulting - vertices of the object. - """ - bpy.context.view_layer.objects.active = obj - for modifier in obj.modifiers: - bpy.ops.object.modifier_apply(modifier = modifier.name) - -def get_verts(obj, float_decimals=4, n_verts=100, apply_modifiers=True, seed=42): - """ - Randomly samples a specified number of vertices from an object. - - Parameters - ---------- - obj : object - Object from which to sample vertices. - float_decimals : int, optional - Number of decimal places to round the vertex coordinates, defaults to 4. - n_verts : int, optional - Number of vertices to sample, defaults to 100. - apply_modifiers : bool, optional - Whether to apply all modifiers on the object before sampling vertices, defaults to True. - seed : int, optional - Seed for the random number generator, defaults to 42. - - Returns - ------- - str - String representation of the randomly selected vertices. - - Notes - ----- - This function randomly samples a specified number of vertices from the given object. - By default, it applies all modifiers on the object before sampling vertices. The - random seed can be set externally for reproducibility. - - If the number of vertices to sample (`n_verts`) exceeds the number of vertices - available in the object, all available vertices will be sampled. - - The vertex coordinates are rounded to the specified number of decimal places - (`float_decimals`) before being included in the output string. - - Examples - -------- - >>> obj = mn.load.molecule_rcsb('6n2y', starting_style=2) - >>> get_verts(obj, float_decimals=3, n_verts=50, apply_modifiers=True, seed=42) - '1.234,2.345,3.456\n4.567,5.678,6.789\n...' - """ - - import random - - random.seed(seed) - - if apply_modifiers: - apply_mods(obj) - - vert_list = [(v.co.x, v.co.y, v.co.z) for v in obj.data.vertices] - - if n_verts > len(vert_list): - n_verts = len(vert_list) - - random_verts = random.sample(vert_list, n_verts) - - verts_string = "" - for i, vert in enumerate(random_verts): - if i < n_verts: - rounded = [round(x, float_decimals) for x in vert] - verts_string += "{},{},{}\n".format(rounded[0], rounded[1], rounded[2]) - - return verts_string - +from .utils import get_verts, apply_mods def test_open_rcsb(snapshot): mn.load.open_structure_rcsb('4ozs') diff --git a/tests/test_obj.py b/tests/test_obj.py new file mode 100644 index 00000000..90526b4c --- /dev/null +++ b/tests/test_obj.py @@ -0,0 +1,15 @@ +import bpy +import MolecularNodes as mn +from .utils import apply_mods, get_verts + +def test_creat_obj(): + # Create a mesh object named "MyMesh" in the collection "MyCollection" + # with vertex locations and bond edges. + 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 = mn.obj.create_object(name, bpy.data.collections['Collection'], locations, bonds) + + assert len(my_object.data.vertices) == 3 + assert my_object.name == name + assert my_object.name != "name" \ No newline at end of file diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 00000000..56fb5e89 --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,76 @@ +import bpy + +def apply_mods(obj): + """ + Applies the modifiers on the modifier stack + + This will realise the computations inside of any Geometry Nodes modifiers, ensuring + that the result of the node trees can be compared by looking at the resulting + vertices of the object. + """ + bpy.context.view_layer.objects.active = obj + for modifier in obj.modifiers: + bpy.ops.object.modifier_apply(modifier = modifier.name) + +def get_verts(obj, float_decimals=4, n_verts=100, apply_modifiers=True, seed=42): + """ + Randomly samples a specified number of vertices from an object. + + Parameters + ---------- + obj : object + Object from which to sample vertices. + float_decimals : int, optional + Number of decimal places to round the vertex coordinates, defaults to 4. + n_verts : int, optional + Number of vertices to sample, defaults to 100. + apply_modifiers : bool, optional + Whether to apply all modifiers on the object before sampling vertices, defaults to True. + seed : int, optional + Seed for the random number generator, defaults to 42. + + Returns + ------- + str + String representation of the randomly selected vertices. + + Notes + ----- + This function randomly samples a specified number of vertices from the given object. + By default, it applies all modifiers on the object before sampling vertices. The + random seed can be set externally for reproducibility. + + If the number of vertices to sample (`n_verts`) exceeds the number of vertices + available in the object, all available vertices will be sampled. + + The vertex coordinates are rounded to the specified number of decimal places + (`float_decimals`) before being included in the output string. + + Examples + -------- + >>> obj = mn.load.molecule_rcsb('6n2y', starting_style=2) + >>> get_verts(obj, float_decimals=3, n_verts=50, apply_modifiers=True, seed=42) + '1.234,2.345,3.456\n4.567,5.678,6.789\n...' + """ + + import random + + random.seed(seed) + + if apply_modifiers: + apply_mods(obj) + + vert_list = [(v.co.x, v.co.y, v.co.z) for v in obj.data.vertices] + + if n_verts > len(vert_list): + n_verts = len(vert_list) + + random_verts = random.sample(vert_list, n_verts) + + verts_string = "" + for i, vert in enumerate(random_verts): + if i < n_verts: + rounded = [round(x, float_decimals) for x in vert] + verts_string += "{},{},{}\n".format(rounded[0], rounded[1], rounded[2]) + + return verts_string \ No newline at end of file