Merge pull request #279 from BradyAJohnston/dev-add_attribute

Improves the 'add_attribute()` function

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
+        # 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

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)
 

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')

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"

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