diff --git a/docs/build_node_docs.py b/docs/build_node_docs.py index cbb10f35..d8a7c493 100644 --- a/docs/build_node_docs.py +++ b/docs/build_node_docs.py @@ -1,12 +1,14 @@ +import os +import pathlib +import sys + import bpy +import griffe from quartodoc import MdRenderer + import molecularnodes as mn -import griffe -import os -import sys -import pathlib -sys.path.insert(0, os.path.abspath('..')) +sys.path.insert(0, os.path.abspath("..")) folder = pathlib.Path(__file__).resolve().parent file_output_qmd = os.path.join(folder, "nodes/index.qmd") @@ -28,12 +30,12 @@ def get_values(sockets): default = None if dtype == "Float": default = round(socket.default_value, 2) - elif dtype in ['Geometry', 'Collection', 'Object']: + elif dtype in ["Geometry", "Collection", "Object"]: default = None elif dtype == "Vector": default = [round(x, 2) for x in socket.default_value] elif dtype == "Material": - default = '`MN Default`' + default = "`MN Default`" elif dtype == "Color": default = col_to_rgb_str(socket.default_value) else: @@ -44,13 +46,13 @@ def get_values(sockets): name=socket.name, annotation=dtype, value=default, - description=socket.description + description=socket.description, ) ) return param_list -cat = '' +cat = "" text = griffe.docstrings.dataclasses.DocstringSectionText params = griffe.docstrings.dataclasses.DocstringSectionParameters @@ -58,7 +60,8 @@ def get_values(sockets): for category, node_list in mn.ui.node_info.menu_items.items(): objects = [] objects.append( - [text(title=None, value=f"## {mn.blender.nodes.format_node_name(category)}")]) + [text(title=None, value=f"## {mn.blender.nodes.format_node_name(category)}")] + ) for item in node_list: if isinstance(item, str): @@ -66,24 +69,26 @@ def get_values(sockets): iter_list = [item] - if item['label'] == "custom": - iter_list = item['values'] + if item["label"] == "custom": + iter_list = item["values"] for entry in iter_list: - name = entry['name'] + name = entry["name"] if name.startswith("mn."): - name = entry['backup'] + name = entry["backup"] entry_list = [] - desc = entry.get('description') - urls = entry.get('video_url') + desc = entry.get("description") + urls = entry.get("video_url") - inputs = params(get_values( - mn.blender.nodes.inputs(bpy.data.node_groups[name]))) - outputs = params(get_values( - mn.blender.nodes.outputs(bpy.data.node_groups[name]))) + inputs = params( + get_values(mn.blender.nodes.inputs(bpy.data.node_groups[name])) + ) + outputs = params( + get_values(mn.blender.nodes.outputs(bpy.data.node_groups[name])) + ) - title = mn.blender.nodes.format_node_name(entry.get('label')) + title = mn.blender.nodes.format_node_name(entry.get("label")) entry_list.append(text(title=None, value=f"### {title}")) if desc: entry_list.append(text(title=None, value=desc)) @@ -91,15 +96,14 @@ def get_values(sockets): if not isinstance(urls, list): urls = [urls] [ - entry_list.append( - text(title=None, value=f"![]({url}.mp4)") - ) for url in urls + entry_list.append(text(title=None, value=f"![]({url}.mp4)")) + for url in urls ] - if len(inputs.as_dict()['value']) > 0: + if len(inputs.as_dict()["value"]) > 0: entry_list.append(text(value="\n#### Inputs")) entry_list.append(inputs) - if len(outputs.as_dict()['value']) > 0: + if len(outputs.as_dict()["value"]) > 0: entry_list.append(text(value="\n#### Outputs")) entry_list.append(outputs) @@ -116,10 +120,10 @@ def get_values(sockets): """ for category, object in categories.items(): - with open(os.path.join(folder, f'nodes/{category}.qmd'), 'w') as file: + with open(os.path.join(folder, f"nodes/{category}.qmd"), "w") as file: file.write(header) for doc in object: - section = '' + section = "" for sec in doc: file.write(ren.render(sec)) file.write("\n\n") diff --git a/docs/install.py b/docs/install.py index e3edf92e..3bf5bd38 100644 --- a/docs/install.py +++ b/docs/install.py @@ -4,16 +4,12 @@ def main(): - python = os.path.realpath(sys.executable) - commands = [ - f'{python} -m pip install .', - f'{python} -m pip install quartodoc' - ] + commands = [f"{python} -m pip install .", f"{python} -m pip install quartodoc"] for command in commands: - subprocess.run(command.split(' ')) + subprocess.run(command.split(" ")) if __name__ == "__main__": diff --git a/molecularnodes/__init__.py b/molecularnodes/__init__.py index fbe07ee6..c5360371 100644 --- a/molecularnodes/__init__.py +++ b/molecularnodes/__init__.py @@ -29,7 +29,7 @@ "warning": "", "doc_url": "https://bradyajohnston.github.io/MolecularNodes/", "tracker_url": "https://github.com/BradyAJohnston/MolecularNodes/issues", - "category": "Import" + "category": "Import", } auto_load.init() @@ -40,8 +40,7 @@ def register(): auto_load.register() bpy.types.NODE_MT_add.append(MN_add_node_menu) - bpy.types.Object.mn = bpy.props.PointerProperty( - type=MolecularNodesObjectProperties) + bpy.types.Object.mn = bpy.props.PointerProperty(type=MolecularNodesObjectProperties) for func in universe_funcs: try: bpy.app.handlers.load_post.append(func) diff --git a/molecularnodes/auto_load.py b/molecularnodes/auto_load.py index 72254c06..b90f277a 100644 --- a/molecularnodes/auto_load.py +++ b/molecularnodes/auto_load.py @@ -1,6 +1,4 @@ -import os import bpy -import sys import typing import inspect import pkgutil @@ -18,6 +16,7 @@ modules = None ordered_classes = None + def init(): global modules global ordered_classes @@ -25,6 +24,7 @@ def init(): modules = get_all_submodules(Path(__file__).parent) ordered_classes = get_ordered_classes_to_register(modules) + def register(): for cls in ordered_classes: bpy.utils.register_class(cls) @@ -35,6 +35,7 @@ def register(): if hasattr(module, "register"): module.register() + def unregister(): for cls in reversed(ordered_classes): bpy.utils.unregister_class(cls) @@ -49,13 +50,16 @@ def unregister(): # Import modules ################################################# + def get_all_submodules(directory): return list(iter_submodules(directory, directory.name)) + def iter_submodules(path, package_name): for name in sorted(iter_submodule_names(path)): yield importlib.import_module("." + name, package_name) + def iter_submodule_names(path, root=""): for _, module_name, is_package in pkgutil.iter_modules([str(path)]): if is_package: @@ -69,22 +73,30 @@ def iter_submodule_names(path, root=""): # Find classes to register ################################################# + def get_ordered_classes_to_register(modules): return toposort(get_register_deps_dict(modules)) + def get_register_deps_dict(modules): my_classes = set(iter_my_classes(modules)) - my_classes_by_idname = {cls.bl_idname : cls for cls in my_classes if hasattr(cls, "bl_idname")} + my_classes_by_idname = { + cls.bl_idname: cls for cls in my_classes if hasattr(cls, "bl_idname") + } deps_dict = {} for cls in my_classes: - deps_dict[cls] = set(iter_my_register_deps(cls, my_classes, my_classes_by_idname)) + deps_dict[cls] = set( + iter_my_register_deps(cls, my_classes, my_classes_by_idname) + ) return deps_dict + def iter_my_register_deps(cls, my_classes, my_classes_by_idname): yield from iter_my_deps_from_annotations(cls, my_classes) yield from iter_my_deps_from_parent_id(cls, my_classes_by_idname) + def iter_my_deps_from_annotations(cls, my_classes): for value in typing.get_type_hints(cls, {}, {}).values(): dependency = get_dependency_from_annotation(value) @@ -92,6 +104,7 @@ def iter_my_deps_from_annotations(cls, my_classes): if dependency in my_classes: yield dependency + def get_dependency_from_annotation(value): if blender_version >= (2, 93): if isinstance(value, bpy.props._PropertyDeferred): @@ -102,6 +115,7 @@ def get_dependency_from_annotation(value): return value[1]["type"] return None + def iter_my_deps_from_parent_id(cls, my_classes_by_idname): if bpy.types.Panel in cls.__bases__: parent_idname = getattr(cls, "bl_parent_id", None) @@ -110,6 +124,7 @@ def iter_my_deps_from_parent_id(cls, my_classes_by_idname): if parent_cls is not None: yield parent_cls + def iter_my_classes(modules): base_types = get_register_base_types() for cls in get_classes_in_modules(modules): @@ -117,6 +132,7 @@ def iter_my_classes(modules): if not getattr(cls, "is_registered", False): yield cls + def get_classes_in_modules(modules): classes = set() for module in modules: @@ -124,24 +140,38 @@ def get_classes_in_modules(modules): classes.add(cls) return classes + def iter_classes_in_module(module): for value in module.__dict__.values(): if inspect.isclass(value): yield value + def get_register_base_types(): - return set(getattr(bpy.types, name) for name in [ - "Panel", "Operator", "PropertyGroup", - "AddonPreferences", "Header", "Menu", - "Node", "NodeSocket", "NodeTree", - "UIList", "RenderEngine", - "Gizmo", "GizmoGroup", - ]) + return set( + getattr(bpy.types, name) + for name in [ + "Panel", + "Operator", + "PropertyGroup", + "AddonPreferences", + "Header", + "Menu", + "Node", + "NodeSocket", + "NodeTree", + "UIList", + "RenderEngine", + "Gizmo", + "GizmoGroup", + ] + ) # Find order to register to solve dependencies ################################################# + def toposort(deps_dict): sorted_list = [] sorted_values = set() @@ -153,5 +183,5 @@ def toposort(deps_dict): sorted_values.add(value) else: unsorted.append(value) - deps_dict = {value : deps_dict[value] - sorted_values for value in unsorted} + deps_dict = {value: deps_dict[value] - sorted_values for value in unsorted} return sorted_list diff --git a/molecularnodes/io/parse/star.py b/molecularnodes/io/parse/star.py index b6a273a5..73088568 100644 --- a/molecularnodes/io/parse/star.py +++ b/molecularnodes/io/parse/star.py @@ -3,21 +3,22 @@ from .ensemble import Ensemble from ... import blender as bl + @bpy.app.handlers.persistent def _rehydrate_ensembles(scene): for obj in bpy.data.objects: - if hasattr(obj, 'mn') and 'molecule_type' in obj.mn.keys(): - if obj.mn['molecule_type'] == 'star': + if hasattr(obj, "mn") and "molecule_type" in obj.mn.keys(): + if obj.mn["molecule_type"] == "star": ensemble = StarFile.from_blender_object(obj) - if not hasattr(bpy.types.Scene, 'MN_starfile_ensembles'): + if not hasattr(bpy.types.Scene, "MN_starfile_ensembles"): bpy.types.Scene.MN_starfile_ensembles = [] bpy.types.Scene.MN_starfile_ensembles.append(ensemble) + class StarFile(Ensemble): def __init__(self, file_path): super().__init__(file_path) - - + @classmethod def from_starfile(cls, file_path): self = cls(file_path) @@ -28,10 +29,11 @@ def from_starfile(cls, file_path): self._create_mn_columns() self.n_images = self._n_images() return self - + @classmethod def from_blender_object(cls, blender_object): import bpy + self = cls(blender_object["starfile_path"]) self.object = blender_object self.star_node = bl.nodes.get_star_node(self.object) @@ -44,10 +46,10 @@ def from_blender_object(cls, blender_object): self.n_images = self._n_images() bpy.app.handlers.depsgraph_update_post.append(self._update_micrograph_texture) return self - def _read(self): import starfile + star = starfile.read(self.file_path) return star @@ -58,88 +60,117 @@ def _n_images(self): def _create_mn_columns(self): # only RELION 3.1 and cisTEM STAR files are currently supported, fail gracefully - if isinstance(self.data, dict) and 'particles' in self.data and 'optics' in self.data: - self.star_type = 'relion' + if ( + isinstance(self.data, dict) + and "particles" in self.data + and "optics" in self.data + ): + self.star_type = "relion" elif "cisTEMAnglePsi" in self.data: - self.star_type = 'cistem' + self.star_type = "cistem" else: raise ValueError( - 'File is not a valid RELION>=3.1 or cisTEM STAR file, other formats are not currently supported.' + "File is not a valid RELION>=3.1 or cisTEM STAR file, other formats are not currently supported." ) # Get absolute position and orientations - if self.star_type == 'relion': - df = self.data['particles'].merge(self.data['optics'], on='rlnOpticsGroup') + if self.star_type == "relion": + df = self.data["particles"].merge(self.data["optics"], on="rlnOpticsGroup") # get necessary info from dataframes # Standard cryoEM starfile don't have rlnCoordinateZ. If this column is not present # Set it to "0" if "rlnCoordinateZ" not in df: - df['rlnCoordinateZ'] = 0 + df["rlnCoordinateZ"] = 0 - self.positions = df[['rlnCoordinateX', 'rlnCoordinateY', - 'rlnCoordinateZ']].to_numpy() - pixel_size = df['rlnImagePixelSize'].to_numpy().reshape((-1, 1)) + self.positions = df[ + ["rlnCoordinateX", "rlnCoordinateY", "rlnCoordinateZ"] + ].to_numpy() + pixel_size = df["rlnImagePixelSize"].to_numpy().reshape((-1, 1)) self.positions = self.positions * pixel_size - shift_column_names = ['rlnOriginXAngst', - 'rlnOriginYAngst', 'rlnOriginZAngst'] + shift_column_names = [ + "rlnOriginXAngst", + "rlnOriginYAngst", + "rlnOriginZAngst", + ] if all([col in df.columns for col in shift_column_names]): shifts_ang = df[shift_column_names].to_numpy() self.positions -= shifts_ang - df['MNAnglePhi'] = df['rlnAngleRot'] - df['MNAngleTheta'] = df['rlnAngleTilt'] - df['MNAnglePsi'] = df['rlnAnglePsi'] - df['MNPixelSize'] = df['rlnImagePixelSize'] + df["MNAnglePhi"] = df["rlnAngleRot"] + df["MNAngleTheta"] = df["rlnAngleTilt"] + df["MNAnglePsi"] = df["rlnAnglePsi"] + df["MNPixelSize"] = df["rlnImagePixelSize"] try: - df['MNImageId'] = df['rlnMicrographName'].astype( - 'category').cat.codes.to_numpy() + df["MNImageId"] = ( + df["rlnMicrographName"].astype("category").cat.codes.to_numpy() + ) except KeyError: try: - df['MNImageId'] = df['rlnTomoName'].astype( - 'category').cat.codes.to_numpy() + df["MNImageId"] = ( + df["rlnTomoName"].astype("category").cat.codes.to_numpy() + ) except KeyError: - df['MNImageId'] = 0.0 - + df["MNImageId"] = 0.0 + self.data = df - elif self.star_type == 'cistem': + elif self.star_type == "cistem": df = self.data - df['cisTEMZFromDefocus'] = ( - df['cisTEMDefocus1'] + df['cisTEMDefocus2']) / 2 - df['cisTEMZFromDefocus'] = df['cisTEMZFromDefocus'] - \ - df['cisTEMZFromDefocus'].median() - self.positions = df[['cisTEMOriginalXPosition', - 'cisTEMOriginalYPosition', 'cisTEMZFromDefocus']].to_numpy() - df['MNAnglePhi'] = df['cisTEMAnglePhi'] - df['MNAngleTheta'] = df['cisTEMAngleTheta'] - df['MNAnglePsi'] = df['cisTEMAnglePsi'] - df['MNPixelSize'] = df['cisTEMPixelSize'] - df['MNImageId'] = df['cisTEMOriginalImageFilename'].astype( - 'category').cat.codes.to_numpy() - + df["cisTEMZFromDefocus"] = (df["cisTEMDefocus1"] + df["cisTEMDefocus2"]) / 2 + df["cisTEMZFromDefocus"] = ( + df["cisTEMZFromDefocus"] - df["cisTEMZFromDefocus"].median() + ) + self.positions = df[ + [ + "cisTEMOriginalXPosition", + "cisTEMOriginalYPosition", + "cisTEMZFromDefocus", + ] + ].to_numpy() + df["MNAnglePhi"] = df["cisTEMAnglePhi"] + df["MNAngleTheta"] = df["cisTEMAngleTheta"] + df["MNAnglePsi"] = df["cisTEMAnglePsi"] + df["MNPixelSize"] = df["cisTEMPixelSize"] + df["MNImageId"] = ( + df["cisTEMOriginalImageFilename"] + .astype("category") + .cat.codes.to_numpy() + ) + def _convert_mrc_to_tiff(self): import mrcfile from pathlib import Path - if self.star_type == 'relion': - micrograph_path = self.object['rlnMicrographName_categories'][self.star_node.inputs['Image'].default_value - 1] - elif self.star_type == 'cistem': - micrograph_path = self.object['cisTEMOriginalImageFilename_categories'][self.star_node.inputs['Image'].default_value - 1].strip("'") + + if self.star_type == "relion": + micrograph_path = self.object["rlnMicrographName_categories"][ + self.star_node.inputs["Image"].default_value - 1 + ] + elif self.star_type == "cistem": + micrograph_path = self.object["cisTEMOriginalImageFilename_categories"][ + self.star_node.inputs["Image"].default_value - 1 + ].strip("'") else: return False - + # This could be more elegant if not Path(micrograph_path).exists(): pot_micrograph_path = Path(self.file_path).parent / micrograph_path if not pot_micrograph_path.exists(): - if self.star_type == 'relion': - pot_micrograph_path = Path(self.file_path).parent.parent.parent / micrograph_path + if self.star_type == "relion": + pot_micrograph_path = ( + Path(self.file_path).parent.parent.parent / micrograph_path + ) if not pot_micrograph_path.exists(): - raise FileNotFoundError(f"Micrograph file {micrograph_path} not found") + raise FileNotFoundError( + f"Micrograph file {micrograph_path} not found" + ) else: - raise FileNotFoundError(f"Micrograph file {micrograph_path} not found") + raise FileNotFoundError( + f"Micrograph file {micrograph_path} not found" + ) micrograph_path = pot_micrograph_path - tiff_path = Path(micrograph_path).with_suffix('.tiff') + tiff_path = Path(micrograph_path).with_suffix(".tiff") if not tiff_path.exists(): with mrcfile.open(micrograph_path) as mrc: micrograph_data = mrc.data.copy() @@ -148,26 +179,31 @@ def _convert_mrc_to_tiff(self): if micrograph_data.ndim == 3: micrograph_data = np.sum(micrograph_data, axis=0) # Normalize the data to 0-1 - micrograph_data = (micrograph_data - micrograph_data.min()) / (micrograph_data.max() - micrograph_data.min()) - + micrograph_data = (micrograph_data - micrograph_data.min()) / ( + micrograph_data.max() - micrograph_data.min() + ) + if micrograph_data.dtype != np.float32: micrograph_data = micrograph_data.astype(np.float32) from PIL import Image + # Need to invert in Y to generate the correct tiff - Image.fromarray(micrograph_data[::-1,:]).save(tiff_path) + Image.fromarray(micrograph_data[::-1, :]).save(tiff_path) return tiff_path - + def _update_micrograph_texture(self, *_): try: - show_micrograph = self.star_node.inputs['Show Micrograph'] - _ = self.object['mn'] + show_micrograph = self.star_node.inputs["Show Micrograph"] + _ = self.object["mn"] except ReferenceError: - bpy.app.handlers.depsgraph_update_post.remove(self._update_micrograph_texture) + bpy.app.handlers.depsgraph_update_post.remove( + self._update_micrograph_texture + ) return - if self.star_node.inputs['Image'].default_value == self.current_image: + if self.star_node.inputs["Image"].default_value == self.current_image: return else: - self.current_image = self.star_node.inputs['Image'].default_value + self.current_image = self.star_node.inputs["Image"].default_value if not show_micrograph: return tiff_path = self._convert_mrc_to_tiff() @@ -176,40 +212,48 @@ def _update_micrograph_texture(self, *_): image_obj = bpy.data.images[tiff_path.name] except KeyError: image_obj = bpy.data.images.load(str(tiff_path)) - image_obj.colorspace_settings.name = 'Non-Color' - self.micrograph_material.node_tree.nodes['Image Texture'].image = image_obj - self.star_node.inputs['Micrograph'].default_value = image_obj - - + image_obj.colorspace_settings.name = "Non-Color" + self.micrograph_material.node_tree.nodes["Image Texture"].image = image_obj + self.star_node.inputs["Micrograph"].default_value = image_obj - def create_model(self, name='StarFileObject', node_setup=True, world_scale=0.01): + def create_model(self, name="StarFileObject", node_setup=True, world_scale=0.01): from molecularnodes.blender.nodes import get_star_node, MN_micrograph_material + blender_object = bl.obj.create_object( - self.positions * world_scale, collection=bl.coll.mn(), name=name) + self.positions * world_scale, collection=bl.coll.mn(), name=name + ) + + blender_object.mn["molecule_type"] = "star" - blender_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): bl.obj.set_attribute( - blender_object, col, self.data[col].to_numpy().reshape(-1), 'FLOAT', 'POINT') + blender_object, + col, + self.data[col].to_numpy().reshape(-1), + "FLOAT", + "POINT", + ) # If col_type is object, convert to category and add integer values elif col_type == object: - codes = self.data[col].astype( - 'category').cat.codes.to_numpy().reshape(-1) - bl.obj.set_attribute(blender_object, col, codes, 'INT', 'POINT') + codes = ( + self.data[col].astype("category").cat.codes.to_numpy().reshape(-1) + ) + bl.obj.set_attribute(blender_object, col, codes, "INT", "POINT") # Add the category names as a property to the blender object - blender_object[f'{col}_categories'] = list( - self.data[col].astype('category').cat.categories) + blender_object[f"{col}_categories"] = list( + self.data[col].astype("category").cat.categories + ) if node_setup: bl.nodes.create_starting_nodes_starfile( - blender_object, n_images=self.n_images) - self.node_group = blender_object.modifiers['MolecularNodes'].node_group + blender_object, n_images=self.n_images + ) + self.node_group = blender_object.modifiers["MolecularNodes"].node_group blender_object["starfile_path"] = str(self.file_path) self.object = blender_object