Merge pull request #164 from BradyAJohnston/density

Adds support for importing `.map` files via [`mrcfile`](https://github.com/ccpem/mrcfile) as another python package which can be installed and [`pyopenvdb`](https://pypi.org/project/pyopenvdb/) which ships with Blender 3.5.

Enables import and some starting nodes for working with volumetric data.

MolecularNodes/__init__.py

@@ -16,8 +16,8 @@
     "name"        : "MolecularNodes",
     "author"      : "Brady Johnston", 
     "description" : "Importer and nodes for working with structural biology data in Blender.",
-    "blender"     : (3, 4, 0),
-    "version"     : (2, 4, 3),
+    "blender"     : (3, 5, 0),
+    "version"     : (2, 5, 2),
     "location"    : "Scene Properties -> MolecularNodes",
     "warning"     : "",
     "doc_url"     : "https://bradyajohnston.github.io/MolecularNodes/", 
@@ -67,6 +67,16 @@ def register():
         description = "Delete the solvent from the structure on import",
         default = True
         )
+    bpy.types.Scene.mol_import_map_nodes = bpy.props.BoolProperty(
+        name = "mol_import_map_nodes", 
+        description = "Creating starting node tree for imported map.",
+        default = True
+        )
+    bpy.types.Scene.mol_import_map_invert = bpy.props.BoolProperty(
+        name = "mol_import_map_invert", 
+        description = "Invert the values in the map. Low becomes high, high becomes low.",
+        default = False
+        )
     bpy.types.Scene.mol_import_include_bonds = bpy.props.BoolProperty(
         name = "mol_import_include_bonds", 
         description = "Include bonds in the imported structure.",
@@ -102,6 +112,14 @@ def register():
         subtype = 'FILE_PATH', 
         maxlen = 0
         )
+    bpy.types.Scene.mol_import_map = bpy.props.StringProperty(
+        name = 'path_map', 
+        description = 'File path for the map file.', 
+        options = {'TEXTEDIT_UPDATE'}, 
+        default = '', 
+        subtype = 'FILE_PATH', 
+        maxlen = 0
+        )
     bpy.types.Scene.mol_import_local_name = bpy.props.StringProperty(
         name = 'mol_name', 
         description = 'Name of the molecule on import', 
@@ -163,6 +181,7 @@ def register():
     bpy.utils.register_class(MOL_MT_Add_Node_Menu_Properties)
     bpy.utils.register_class(MOL_MT_Add_Node_Menu_Styling)
     bpy.utils.register_class(MOL_MT_Add_Node_Menu_Color)
+    bpy.utils.register_class(MOL_MT_Add_Density_Menu)
     bpy.utils.register_class(MOL_MT_Add_Node_Menu_Bonds)
     bpy.utils.register_class(MOL_MT_Add_Node_Menu_Selections)
     bpy.utils.register_class(MOL_MT_Add_Node_Menu_Membranes)
@@ -179,6 +198,7 @@ def register():
     bpy.utils.register_class(MOL_OT_Import_Method_Selection)
     bpy.utils.register_class(MOL_OT_Import_Protein_Local)
     bpy.utils.register_class(MOL_OT_Import_Protein_MD)
+    bpy.utils.register_class(MOL_OT_Import_Map)
     bpy.utils.register_class(MOL_OT_Assembly_Bio)
     bpy.utils.register_class(MOL_OT_Default_Style)
     bpy.utils.register_class(MOL_OT_Color_Chain)
@@ -198,10 +218,13 @@ def unregister():
     del bpy.types.Scene.mol_import_center
     del bpy.types.Scene.mol_import_del_solvent
     del bpy.types.Scene.mol_import_include_bonds
+    del bpy.types.Scene.mol_import_map_nodes
+    del bpy.types.Scene.mol_import_map_invert
     del bpy.types.Scene.mol_import_panel_selection
     del bpy.types.Scene.mol_import_local_path
     del bpy.types.Scene.mol_import_md_topology
     del bpy.types.Scene.mol_import_md_trajectory
+    del bpy.types.Scene.mol_import_map
     del bpy.types.Scene.mol_import_local_name
     del bpy.types.Scene.mol_import_md_name
     del bpy.types.Scene.mol_import_md_frame_start
@@ -225,6 +248,7 @@ def unregister():
     bpy.utils.unregister_class(MOL_MT_Add_Node_Menu_Styling)
     bpy.utils.unregister_class(MOL_MT_Add_Node_Menu_Color)
     bpy.utils.unregister_class(MOL_MT_Add_Node_Menu_Bonds)
+    bpy.utils.unregister_class(MOL_MT_Add_Density_Menu)
     bpy.utils.unregister_class(MOL_MT_Add_Node_Menu_Selections)
     bpy.utils.unregister_class(MOL_MT_Add_Node_Menu_Membranes)
     bpy.utils.unregister_class(MOL_MT_Add_Node_Menu_DNA)
@@ -239,6 +263,7 @@ def unregister():
     bpy.utils.unregister_class(MOL_OT_Import_Method_Selection)
     bpy.utils.unregister_class(MOL_OT_Import_Protein_Local)
     bpy.utils.unregister_class(MOL_OT_Import_Protein_MD)
+    bpy.utils.unregister_class(MOL_OT_Import_Map)
     bpy.utils.unregister_class(MOL_OT_Assembly_Bio)
     bpy.utils.unregister_class(MOL_OT_Default_Style)
     bpy.utils.unregister_class(MOL_OT_Color_Chain)

MolecularNodes/density.py

@@ -0,0 +1,143 @@
+import bpy
+import pyopenvdb as vdb
+import numpy as np
+import os
+
+def map_to_grid(file: str, invert: bool = False) -> vdb.FloatGrid:
+    """Reads an MRC file and converts it into a pyopenvdb FloatGrid object.
+
+    This function reads a file in MRC format, and converts it into a pyopenvdb FloatGrid object,
+    which can be used to represent volumetric data in Blender.
+
+    Args:
+        file (str): The path to the MRC file.
+        invert (bool): Whether to invert the data from the grid, defaulting to False. Some file types
+        such as EM tomograms have inverted values, where a high value == low density.
+
+    Returns:
+        pyopenvdb.FloatGrid: A pyopenvdb FloatGrid object containing the density data.
+    """
+    import mrcfile
+    volume = mrcfile.read(file)
+
+    dataType = volume.dtype
+
+    # enables different grid types
+
+    if dataType == "float32" or dataType == "float64":
+        grid = vdb.FloatGrid()
+    elif dataType == "int8" or dataType == "int16" or dataType == "int32":
+        volume = volume.astype('int32')
+        grid = vdb.Int32Grid()
+    elif dataType == "int64":
+        grid = vdb.Int64Grid()
+
+    if invert:
+        volume = np.max(volume) - volume
+
+    try:
+        grid.copyFromArray(volume)
+    except ValueError:
+        print(f"Grid data type '{volume.dtype}' is an unsupported type.")
+
+    grid.gridClass = vdb.GridClass.FOG_VOLUME
+    grid.name = 'density'
+    return grid
+
+def path_to_vdb(file: str):
+    # Set up file paths
+    folder_path = os.path.dirname(file)
+    name = os.path.basename(file).split(".")[0]
+    file_name = name + '.vdb'
+    file_path = os.path.join(folder_path, file_name)
+    return file_path
+
+
+def map_to_vdb(file: str, invert: bool = False, world_scale=0.01, overwrite=False) -> str:
+    """
+    Converts an MRC file to a .vdb file using pyopenvdb.
+
+    Args:
+        file (str): The path to the input MRC file.
+        invert (bool): Whether to invert the data from the grid, defaulting to False. Some file types
+        such as EM tomograms have inverted values, where a high value == low density.
+        world_scale (float, optional): The scaling factor to apply to the voxel size of the input file. Defaults to 0.01.
+        overwrite (bool, optional): If True, the .vdb file will be overwritten if it already exists. Defaults to False.
+
+    Returns:
+        str: The path to the converted .vdb file.
+    """
+    import mrcfile
+    file_path = path_to_vdb(file)
+
+    # If the map has already been converted to a .vdb and overwrite is False, return that instead
+    if os.path.exists(file_path) and not overwrite:
+        return file_path
+
+    # Read in the MRC file and convert it to a pyopenvdb grid
+    grid = map_to_grid(file, invert = invert)
+
+    # Read the voxel size from the MRC file and convert it to a numpy array
+    with mrcfile.open(file) as mrc:
+        voxel_size = np.array([mrc.voxel_size.x, mrc.voxel_size.y, mrc.voxel_size.z])
+
+    # Rotate and scale the grid for import into Blender
+    grid.transform.rotate(np.pi / 2, vdb.Axis(1))
+    grid.transform.scale(np.array((-1, 1, 1)) * world_scale * voxel_size)
+
+    # Write the grid to a .vdb file
+    vdb.write(file_path, grid)
+
+    # Return the path to the output file
+    return file_path
+
+def vdb_to_volume(file: str) -> bpy.types.Object:
+    """Imports a VDB file as a Blender volume object.
+
+    Args:
+        file (str): Path to the VDB file.
+
+    Returns:
+        bpy.types.Object: A Blender object containing the imported volume data.
+    """
+    # extract name of file for object name
+    name = os.path.basename(file).split('.')[0]
+
+    # import the volume object
+    bpy.ops.object.volume_import(
+        filepath = file, 
+        files = [], 
+        scale = [1, 1, 1], 
+        rotation = [0, 0, 0]
+    )
+
+    # get reference to imported object and return
+    vol = bpy.context.scene.objects[name]
+    return vol
+
+
+def load(file: str, name: str = None, invert: bool = False, world_scale: float = 0.01) -> bpy.types.Object:
+    """
+    Loads an MRC file into Blender as a volumetric object.
+
+    Args:
+        file (str): Path to the MRC file.
+        name (str, optional): If not None, renames the object with the new name.
+        invert (bool): Whether to invert the data from the grid, defaulting to False. Some file types
+        such as EM tomograms have inverted values, where a high value == low density.
+        world_scale (float, optional): Scale of the object in the world. Defaults to 0.01.
+
+    Returns:
+        bpy.types.Object: The loaded volumetric object.
+    """
+    # Convert MRC file to VDB format
+    vdb_file = map_to_vdb(file, invert = invert, world_scale = world_scale)
+
+    # Import VDB file into Blender
+    vol_object = vdb_to_volume(vdb_file)
+
+    if name:
+        # Rename object to specified name
+        vol_object.name = name
+
+    return vol_object

MolecularNodes/nodes.py

@@ -97,6 +97,37 @@ def add_custom_node_group_to_node(parent_group, node_name, location = [0,0], wid
 
     return node
 
+def create_starting_nodes_density(obj):
+    # ensure there is a geometry nodes modifier called 'MolecularNodes' that is created and applied to the object
+    node_mod = obj.modifiers.get('MolecularNodes')
+    if not node_mod:
+        node_mod = obj.modifiers.new("MolecularNodes", "NODES")
+    obj.modifiers.active = node_mod
+    # create a new GN node group, specific to this particular molecule
+    node_group = gn_new_group_empty(f"MOL_density_{str(obj.name)}")
+    node_mod.node_group = node_group
+    # move the input and output nodes for the group
+    node_input = node_mod.node_group.nodes[bpy.app.translations.pgettext_data("Group Input",)]
+    node_input.location = [0, 0]
+    node_output = node_mod.node_group.nodes[bpy.app.translations.pgettext_data("Group Output",)]
+    node_output.location = [800, 0]
+
+    node_density = add_custom_node_group(node_mod, 'MOL_style_density_surface', [400, 0])
+    node_density.inputs['Material'].default_value = mol_base_material()
+
+
+    link = node_group.links.new
+    link(
+        node_input.outputs[0], 
+        node_density.inputs[0]
+    )
+    link(
+        node_density.outputs[0], 
+        node_output.inputs[0]
+    )
+
+
+
 def create_starting_node_tree(obj, coll_frames, starting_style = "atoms"):
 
     # ensure there is a geometry nodes modifier called 'MolecularNodes' that is created and applied to the object

MolecularNodes/pkg.py

@@ -268,6 +268,7 @@ def install_all_packages(pypi_mirror_provider: str='Default') -> list:
     pkgs = get_pkgs()
     results = []
     for pkg in pkgs.items():
+
         try:
             result = install_package(package=f"{pkg.get('name')}=={pkg.get('version')}", 
                                      pypi_mirror_provider=mirror_url)

MolecularNodes/requirements.txt

@@ -1,2 +1,3 @@
 biotite==0.36.1   # General parsing of structural files.
-MDAnalysis==2.2.0 # Reading of molecular dynamics trajectories.
+MDAnalysis==2.2.0 # Reading of molecular dynamics trajectories.
+mrcfile==1.4.3    # Importing EM density files.