From c57ae2c86a4aa17d6ab941e7c79fefdffffa4acc Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Sat, 25 May 2024 14:09:35 +0200 Subject: [PATCH] cleanup fixing ruff --- build.py | 4 +- docs/build_node_docs.py | 17 +--- molecularnodes/auto_load.py | 8 +- molecularnodes/blender/bones.py | 7 +- molecularnodes/blender/nodes.py | 139 +++++++--------------------- molecularnodes/blender/obj.py | 27 ++---- molecularnodes/color.py | 6 +- molecularnodes/io/cellpack.py | 4 +- molecularnodes/io/density.py | 4 +- molecularnodes/io/dna.py | 12 +-- molecularnodes/io/md.py | 19 +--- molecularnodes/io/parse/bcif.py | 78 ++++------------ molecularnodes/io/parse/cellpack.py | 16 +--- molecularnodes/io/parse/cif.py | 29 ++---- molecularnodes/io/parse/ensemble.py | 4 +- molecularnodes/io/parse/mda.py | 102 +++++--------------- molecularnodes/io/parse/molecule.py | 73 ++++----------- molecularnodes/io/parse/mrc.py | 16 +--- molecularnodes/io/parse/pdb.py | 52 +++-------- molecularnodes/io/parse/pdbx.py | 71 +++----------- molecularnodes/io/parse/star.py | 60 +++--------- molecularnodes/io/retrieve.py | 8 +- molecularnodes/io/star.py | 4 +- molecularnodes/io/wwpdb.py | 5 +- molecularnodes/pkg.py | 20 +--- molecularnodes/props.py | 4 +- molecularnodes/ui/func.py | 4 +- molecularnodes/ui/node_info.py | 4 +- molecularnodes/ui/ops.py | 4 +- molecularnodes/ui/panel.py | 4 +- molecularnodes/ui/pref.py | 4 +- molecularnodes/utils.py | 24 ++--- pyproject.toml | 12 +-- tests/install.py | 5 +- tests/test_assembly.py | 4 +- tests/test_cellpack.py | 8 +- tests/test_density.py | 12 +-- tests/test_download.py | 15 +-- tests/test_load.py | 47 +++------- tests/test_mda.py | 24 ++--- tests/test_nodes.py | 109 +++++----------------- tests/test_ops.py | 8 +- tests/test_pdbx.py | 4 +- tests/test_select.py | 5 +- tests/test_star.py | 27 ++---- tests/utils.py | 20 +--- 46 files changed, 271 insertions(+), 862 deletions(-) diff --git a/build.py b/build.py index 82005f29..f7e37867 100644 --- a/build.py +++ b/build.py @@ -18,9 +18,7 @@ def zip_template() -> None: # Add the file to the zip file zipf.write( file_path, - arcname=os.path.relpath( - file_path, start="molecularnodes/assets/template/" - ), + arcname=os.path.relpath(file_path, start="molecularnodes/assets/template/"), ) diff --git a/docs/build_node_docs.py b/docs/build_node_docs.py index d8a7c493..d7adf22a 100644 --- a/docs/build_node_docs.py +++ b/docs/build_node_docs.py @@ -59,9 +59,7 @@ def get_values(sockets): categories = {} 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)}")] - ) + objects.append([text(title=None, value=f"## {mn.blender.nodes.format_node_name(category)}")]) for item in node_list: if isinstance(item, str): @@ -81,12 +79,8 @@ def get_values(sockets): 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")) entry_list.append(text(title=None, value=f"### {title}")) @@ -95,10 +89,7 @@ def get_values(sockets): if urls: 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: entry_list.append(text(value="\n#### Inputs")) diff --git a/molecularnodes/auto_load.py b/molecularnodes/auto_load.py index b90f277a..f7ef1088 100644 --- a/molecularnodes/auto_load.py +++ b/molecularnodes/auto_load.py @@ -80,15 +80,11 @@ def get_ordered_classes_to_register(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 diff --git a/molecularnodes/blender/bones.py b/molecularnodes/blender/bones.py index 6b7dab67..065190d4 100644 --- a/molecularnodes/blender/bones.py +++ b/molecularnodes/blender/bones.py @@ -37,8 +37,7 @@ def add_bones(object, name="armature"): def get_bone_positions(object): positions, atom_name, chain_id, res_id, sec_struct = [ - obj.get_attribute(object, att) - for att in ["position", "atom_name", "chain_id", "res_id", "sec_struct"] + obj.get_attribute(object, att) for att in ["position", "atom_name", "chain_id", "res_id", "sec_struct"] ] is_alpha_carbon = atom_name == 2 @@ -104,9 +103,7 @@ def create_bones(positions, chain_ids, name="armature"): 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 " - ) + bl_description = "Automatically add armature for each amino acid of the structure " def execute(self, context): object = context.active_object diff --git a/molecularnodes/blender/nodes.py b/molecularnodes/blender/nodes.py index aaa818b8..62e69153 100644 --- a/molecularnodes/blender/nodes.py +++ b/molecularnodes/blender/nodes.py @@ -116,9 +116,7 @@ def node_tree_debug( ) -> bpy.types.GeometryNodeTree: group = new_tree(name=name, fallback=False) info = group.nodes.new("GeometryNodeObjectInfo") - group.links.new( - info.outputs["Geometry"], group.nodes["Group Output"].inputs[0] - ) + group.links.new(info.outputs["Geometry"], group.nodes["Group Output"].inputs[0]) return group @@ -159,9 +157,7 @@ def get_input(group: bpy.types.GeometryNodeGroup) -> bpy.types.Node: ] -def get_mod( - object: bpy.types.Object, name: str = "MolecularNodes" -) -> bpy.types.Modifier: +def get_mod(object: bpy.types.Object, name: str = "MolecularNodes") -> bpy.types.Modifier: node_mod = object.modifiers.get(name) if not node_mod: node_mod = object.modifiers.new(name, "NODES") @@ -171,13 +167,7 @@ def get_mod( def format_node_name(name: str) -> str: "Formats a node's name for nicer printing." - return ( - name.strip("MN_") - .replace("_", " ") - .title() - .replace("Dna", "DNA") - .replace("Topo ", "Topology ") - ) + return name.strip("MN_").replace("_", " ").title().replace("Dna", "DNA").replace("Topo ", "Topology ") def get_nodes_last_output( @@ -272,12 +262,7 @@ def append(node_name: str, link: bool = False) -> bpy.types.GeometryNodeTree: if not node or link: node_name_components = node_name.split("_") if node_name_components[0] == "MN": - data_file = ( - MN_DATA_FILE[:-6] - + "_" - + node_name_components[1] - + ".blend" - ) + data_file = MN_DATA_FILE[:-6] + "_" + node_name_components[1] + ".blend" bpy.ops.wm.append( "EXEC_DEFAULT", directory=os.path.join(data_file, "NodeTree"), @@ -319,9 +304,7 @@ def MN_micrograph_material() -> bpy.types.Material: return bpy.data.materials[mat_name] -def new_tree( - name: str = "Geometry Nodes", geometry: bool = True, fallback: bool = True -) -> bpy.types.GeometryNodeTree: +def new_tree(name: str = "Geometry Nodes", geometry: bool = True, fallback: bool = True) -> bpy.types.GeometryNodeTree: group = bpy.data.node_groups.get(name) # if the group already exists, return it and don't create a new one if group and fallback: @@ -334,19 +317,13 @@ def new_tree( input_node.location.x = -200 - input_node.width output_node.location.x = 200 if geometry: - group.interface.new_socket( - "Geometry", in_out="INPUT", socket_type="NodeSocketGeometry" - ) - group.interface.new_socket( - "Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry" - ) + group.interface.new_socket("Geometry", in_out="INPUT", socket_type="NodeSocketGeometry") + group.interface.new_socket("Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry") group.links.new(output_node.inputs[0], input_node.outputs[0]) return group -def assign_material( - node: bpy.types.GeometryNode, material: str = "default" -) -> None: +def assign_material(node: bpy.types.GeometryNode, material: str = "default") -> None: material_socket = node.inputs.get("Material") if material_socket: if not material: @@ -370,9 +347,7 @@ def add_node( # actually invoke the operator to add a node to the current node tree # use_transform=True ensures it appears where the user's mouse is and is currently # being moved so the user can place it where they wish - bpy.ops.node.add_node( - "INVOKE_DEFAULT", type="GeometryNodeGroup", use_transform=True - ) + bpy.ops.node.add_node("INVOKE_DEFAULT", type="GeometryNodeGroup", use_transform=True) bpy.context.area.type = prev_context node = bpy.context.active_node node.node_tree = bpy.data.node_groups[node_name] @@ -462,9 +437,7 @@ def change_style_node(object: bpy.types.Object, style: str) -> None: pass -def create_starting_nodes_starfile( - object: bpy.types.Object, n_images: int = 1 -) -> None: +def create_starting_nodes_starfile(object: bpy.types.Object, n_images: int = 1) -> None: # ensure there is a geometry nodes modifier called 'MolecularNodes' that is created and applied to the object node_mod = get_mod(object) @@ -570,9 +543,7 @@ def create_starting_node_tree( if set_color: node_color_set = add_custom(group, "MN_color_set", [200, 0]) node_color_common = add_custom(group, "MN_color_common", [-50, -150]) - node_random_color = add_custom( - group, "MN_color_attribute_random", [-300, -150] - ) + node_random_color = add_custom(group, "MN_color_attribute_random", [-300, -150]) link(node_input.outputs["Geometry"], node_color_set.inputs[0]) link( @@ -594,9 +565,7 @@ def create_starting_node_tree( node_animate = add_custom(group, "MN_animate_value", [500, -300]) node_animate_frames.inputs["Frames"].default_value = coll_frames - node_animate.inputs["To Max"].default_value = ( - len(coll_frames.objects) - 1 - ) + node_animate.inputs["To Max"].default_value = len(coll_frames.objects) - 1 link(to_animate.outputs[0], node_animate_frames.inputs[0]) link(node_animate_frames.outputs[0], node_style.inputs[0]) @@ -665,15 +634,9 @@ def assembly_initialise(mol: bpy.types.Object) -> bpy.types.GeometryNodeTree: Setup the required data object and nodes for building an assembly. """ - transforms = utils.array_quaternions_from_dict( - mol["biological_assemblies"] - ) - data_object = obj.create_data_object( - array=transforms, name=f"data_assembly_{mol.name}" - ) - tree_assembly = create_assembly_node_tree( - name=mol.name, iter_list=mol["chain_ids"], data_object=data_object - ) + transforms = utils.array_quaternions_from_dict(mol["biological_assemblies"]) + data_object = obj.create_data_object(array=transforms, name=f"data_assembly_{mol.name}") + tree_assembly = create_assembly_node_tree(name=mol.name, iter_list=mol["chain_ids"], data_object=data_object) return tree_assembly @@ -695,9 +658,7 @@ def create_assembly_node_tree( tree = new_tree(name=node_group_name) link = tree.links.new - n_assemblies = len( - np.unique(obj.get_attribute(data_object, "assembly_id")) - ) + n_assemblies = len(np.unique(obj.get_attribute(data_object, "assembly_id"))) node_group_instances = split_geometry_to_instances( name=f".MN_utils_split_{name}", @@ -707,9 +668,7 @@ def create_assembly_node_tree( node_group_assembly_instance = append(".MN_assembly_instance_chains") node_instances = add_custom(tree, node_group_instances.name, [0, 0]) - node_assembly = add_custom( - tree, node_group_assembly_instance.name, [200, 0] - ) + node_assembly = add_custom(tree, node_group_assembly_instance.name, [200, 0]) node_assembly.inputs["data_object"].default_value = data_object out_sockets = outputs(tree) @@ -742,9 +701,7 @@ def create_assembly_node_tree( for info in socket_info: socket = tree.interface.items_tree.get(info["name"]) if not socket: - socket = tree.interface.new_socket( - info["name"], in_out="INPUT", socket_type=info["type"] - ) + socket = tree.interface.new_socket(info["name"], in_out="INPUT", socket_type=info["type"]) socket.default_value = info["default"] socket.min_value = info["min"] socket.max_value = info["max"] @@ -765,18 +722,14 @@ def create_assembly_node_tree( def add_inverse_selection(group: bpy.types.GeometryNodeTree) -> None: output = get_output(group) if "Inverted" not in output.inputs.keys(): - group.interface.new_socket( - "Inverted", in_out="OUTPUT", socket_type="NodeSocketBool" - ) + group.interface.new_socket("Inverted", in_out="OUTPUT", socket_type="NodeSocketBool") loc = output.location bool_math = group.nodes.new("FunctionNodeBooleanMath") bool_math.location = [loc[0], -100] bool_math.operation = "NOT" - group.links.new( - output.inputs["Selection"].links[0].from_socket, bool_math.inputs[0] - ) + group.links.new(output.inputs["Selection"].links[0].from_socket, bool_math.inputs[0]) group.links.new(bool_math.outputs[0], output.inputs["Inverted"]) @@ -865,13 +818,9 @@ def custom_iswitch( # is colors, then generate a random pastel color for each value default_lookup = None if default_values: - default_lookup = dict( - zip(iter_list, itertools.cycle(default_values)) - ) + default_lookup = dict(zip(iter_list, itertools.cycle(default_values))) elif dtype == "RGBA": - default_lookup = dict( - zip(iter_list, [color.random_rgb() for i in iter_list]) - ) + default_lookup = dict(zip(iter_list, [color.random_rgb() for i in iter_list])) # for each item in the iter_list, we create a new socket on the interface for this # node group, and link it to the interface on the index switch. The index switch @@ -881,9 +830,7 @@ def custom_iswitch( if i > 1: node_iswitch.index_switch_items.new() - socket = group.interface.new_socket( - name=f"{prefix}{item}", in_out="INPUT", socket_type=socket_type - ) + socket = group.interface.new_socket(name=f"{prefix}{item}", in_out="INPUT", socket_type=socket_type) # if a set of default values was given, then use it for setting # the defaults on the created sockets of the node group if default_lookup: @@ -894,9 +841,7 @@ def custom_iswitch( node_iswitch.inputs[str(i)], ) - socket_out = group.interface.new_socket( - name="Color", in_out="OUTPUT", socket_type=socket_type - ) + socket_out = group.interface.new_socket(name="Color", in_out="OUTPUT", socket_type=socket_type) link( node_iswitch.outputs["Output"], node_output.inputs[socket_out.identifier], @@ -908,14 +853,10 @@ def custom_iswitch( except Exception as e: node_name = group.name bpy.data.node_groups.remove(group) - raise NodeGroupCreationError( - f"Unable to make node group: {node_name}.\nError: {e}" - ) + raise NodeGroupCreationError(f"Unable to make node group: {node_name}.\nError: {e}") -def resid_multiple_selection( - node_name: str, input_resid_string: str -) -> bpy.types.GeometryNodeTree: +def resid_multiple_selection(node_name: str, input_resid_string: str) -> bpy.types.GeometryNodeTree: """ Returns a node group that takes an integer input and creates a boolean tick box for each item in the input list. Outputs are the selected @@ -966,28 +907,18 @@ def resid_multiple_selection( # set two new inputs current_node.node_tree = append("MN_select_res_id_range") [resid_start, resid_end] = residue_id.split("-")[:2] - socket_1 = tree.interface.new_socket( - "res_id: Min", in_out="INPUT", socket_type="NodeSocketInt" - ) + socket_1 = tree.interface.new_socket("res_id: Min", in_out="INPUT", socket_type="NodeSocketInt") socket_1.default_value = int(resid_start) - socket_2 = tree.interface.new_socket( - "res_id: Max", in_out="INPUT", socket_type="NodeSocketInt" - ) + socket_2 = tree.interface.new_socket("res_id: Max", in_out="INPUT", socket_type="NodeSocketInt") socket_2.default_value = int(resid_end) # a residue range - link( - node_input.outputs[socket_1.identifier], current_node.inputs[0] - ) - link( - node_input.outputs[socket_2.identifier], current_node.inputs[1] - ) + link(node_input.outputs[socket_1.identifier], current_node.inputs[0]) + link(node_input.outputs[socket_2.identifier], current_node.inputs[1]) else: # create a node current_node.node_tree = append("MN_select_res_id_single") - socket = tree.interface.new_socket( - "res_id", in_out="INPUT", socket_type="NodeSocketInt" - ) + socket = tree.interface.new_socket("res_id", in_out="INPUT", socket_type="NodeSocketInt") socket.default_value = int(residue_id) link(node_input.outputs[socket.identifier], current_node.inputs[0]) @@ -1011,12 +942,8 @@ def resid_multiple_selection( 800, (residue_id_index + 1) / 2 * node_sep_dis, ] - tree.interface.new_socket( - "Selection", in_out="OUTPUT", socket_type="NodeSocketBool" - ) - tree.interface.new_socket( - "Inverted", in_out="OUTPUT", socket_type="NodeSocketBool" - ) + tree.interface.new_socket("Selection", in_out="OUTPUT", socket_type="NodeSocketBool") + tree.interface.new_socket("Inverted", in_out="OUTPUT", socket_type="NodeSocketBool") link(prev.outputs[0], residue_id_group_out.inputs["Selection"]) invert_bool_math = new_node("FunctionNodeBooleanMath") invert_bool_math.location = [ diff --git a/molecularnodes/blender/obj.py b/molecularnodes/blender/obj.py index c0f54a34..d5a9cb2b 100644 --- a/molecularnodes/blender/obj.py +++ b/molecularnodes/blender/obj.py @@ -1,6 +1,6 @@ from dataclasses import dataclass from typing import Union, List, Optional -from typing import Optional, Type +from typing import Type from types import TracebackType import bpy import numpy as np @@ -40,10 +40,7 @@ def centre(array: np.array) -> np.ndarray: def centre_weighted(array: np.ndarray, weight: np.ndarray) -> np.ndarray: - return np.array( - np.sum(array * weight.reshape((len(array), 1)), axis=0) - / np.sum(weight) - ) + return np.array(np.sum(array * weight.reshape((len(array), 1)), axis=0) / np.sum(weight)) class ObjectTracker: @@ -240,9 +237,7 @@ def set_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[type].dname, data.reshape(-1).copy(order="C") - ) + attribute.data.foreach_set(TYPES[type].dname, data.reshape(-1).copy(order="C")) # 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 @@ -258,9 +253,7 @@ def set_attribute( return attribute -def get_attribute( - object: bpy.types.Object, name: str = "position", evaluate: bool = False -) -> np.ndarray: +def get_attribute(object: bpy.types.Object, name: str = "position", evaluate: bool = False) -> np.ndarray: """ Get the attribute data from the object. @@ -282,9 +275,7 @@ def get_attribute( Possible attributes are: {attribute_names=}" ) else: - raise AttributeError( - f"The selected attribute '{name}' does not exist on the mesh." - ) + raise AttributeError(f"The selected attribute '{name}' does not exist on the mesh.") # Get the attribute and some metadata about it from the object att = object.data.attributes[name] @@ -309,9 +300,7 @@ def get_attribute( return array -def import_vdb( - file: str, collection: bpy.types.Collection = None -) -> bpy.types.Object: +def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types.Object: """ Imports a VDB file as a Blender volume object, in the MolecularNodes collection. @@ -411,8 +400,6 @@ def create_data_object( if np.issubdtype(data.dtype, str): data = np.unique(data, return_inverse=True)[1] - set_attribute( - object, name=column, data=data, type=type, domain="POINT" - ) + set_attribute(object, name=column, data=data, type=type, domain="POINT") return object diff --git a/molecularnodes/color.py b/molecularnodes/color.py index b8965eaa..8ff433fa 100644 --- a/molecularnodes/color.py +++ b/molecularnodes/color.py @@ -2,7 +2,7 @@ import colorsys import numpy as np from numpy.typing import NDArray -from typing import List, Optional, Any, Dict, Tuple +from typing import List, Dict, Tuple def random_rgb(seed: int = 6) -> NDArray[np.float64]: @@ -52,9 +52,7 @@ def color_chains_equidistant( return chain_colors / 255 -def color_chains( - atomic_numbers: NDArray[np.int32], chain_ids: NDArray[np.character] -) -> NDArray[np.float32]: +def color_chains(atomic_numbers: NDArray[np.int32], chain_ids: NDArray[np.character]) -> NDArray[np.float32]: mask = atomic_numbers == 6 colors = colors_from_elements(atomic_numbers) chain_color_dict = equidistant_colors(chain_ids) diff --git a/molecularnodes/io/cellpack.py b/molecularnodes/io/cellpack.py index 6fe350d5..71cce99a 100644 --- a/molecularnodes/io/cellpack.py +++ b/molecularnodes/io/cellpack.py @@ -24,9 +24,7 @@ def load( fraction: float = 1, ): ensemble = parse.CellPack(file_path) - model = ensemble.create_model( - name=name, node_setup=node_setup, world_scale=world_scale, fraction=fraction - ) + model = ensemble.create_model(name=name, node_setup=node_setup, world_scale=world_scale, fraction=fraction) return model diff --git a/molecularnodes/io/density.py b/molecularnodes/io/density.py index dd18976a..c6f84075 100644 --- a/molecularnodes/io/density.py +++ b/molecularnodes/io/density.py @@ -52,9 +52,7 @@ def load( center: bool = False, overwrite: bool = False, ): - density = parse.MRC( - file_path=file_path, center=center, invert=invert, overwrite=overwrite - ) + density = parse.MRC(file_path=file_path, center=center, invert=invert, overwrite=overwrite) density.create_model(name=name, setup_nodes=setup_nodes, style=style) return density diff --git a/molecularnodes/io/dna.py b/molecularnodes/io/dna.py index 96c79810..3100bf2e 100644 --- a/molecularnodes/io/dna.py +++ b/molecularnodes/io/dna.py @@ -124,9 +124,7 @@ def read_topology_old(filepath): # convert the columns to numeric array_int = np.zeros(array_str.shape, dtype=int) - array_int[:, (0, 2, 3)] = array_str[:, (0, 2, 3)].astype( - int - ) # easy convert numeric columns to int + array_int[:, (0, 2, 3)] = array_str[:, (0, 2, 3)].astype(int) # easy convert numeric columns to int # convert bases (A, C, G, T) to (30, 31, 32, 33) array_int[:, 1] = base_to_int(array_str[:, 1]) @@ -286,15 +284,11 @@ def load(top, traj, name="oxDNA", setup_nodes=True, world_scale=0.01): for i, frame in enumerate(trajectory): fill_n = int(np.ceil(np.log10(n_frames))) frame_name = f"{name}_frame_{str(i).zfill(fill_n)}" - frame_mol = obj.create_object( - frame[:, 0:3] * scale_dna, name=frame_name, collection=collection - ) + frame_mol = obj.create_object(frame[:, 0:3] * scale_dna, name=frame_name, collection=collection) set_attributes_to_dna_mol(frame_mol, frame, scale_dna) if setup_nodes: - nodes.create_starting_node_tree( - mol, coll_frames=collection, style="oxdna", set_color=False - ) + nodes.create_starting_node_tree(mol, coll_frames=collection, style="oxdna", set_color=False) return mol, collection diff --git a/molecularnodes/io/md.py b/molecularnodes/io/md.py index 65924956..c3a596ca 100644 --- a/molecularnodes/io/md.py +++ b/molecularnodes/io/md.py @@ -62,9 +62,7 @@ class MockUniverse: description="True will load all of the requested frames into the scene and into memory. False will stream the trajectory from a live MDAnalysis session", default=False, ) -bpy.types.Scene.list_index = bpy.props.IntProperty( - name="Index for trajectory selection list.", default=0 -) +bpy.types.Scene.list_index = bpy.props.IntProperty(name="Index for trajectory selection list.", default=0) def load( @@ -117,8 +115,7 @@ def execute(self, context): if not pkg.is_current("MDAnalysis"): self.report( {"ERROR"}, - message="MDAnalysis is not installed. " - "Please install it to use this feature.", + message="MDAnalysis is not installed. " "Please install it to use this feature.", ) return {"CANCELLED"} top = scene.MN_import_md_topology @@ -158,9 +155,7 @@ class TrajectorySelectionItem(bpy.types.PropertyGroup): bl_idname = "testing" - name: bpy.props.StringProperty( - name="Attribute Name", description="Attribute", default="custom_selection" - ) + name: bpy.props.StringProperty(name="Attribute Name", description="Attribute", default="custom_selection") selection: bpy.props.StringProperty( name="Selection String", @@ -171,17 +166,13 @@ class TrajectorySelectionItem(bpy.types.PropertyGroup): # have to manually register this class otherwise the PropertyGroup registration fails bpy.utils.register_class(TrajectorySelectionItem) -bpy.types.Scene.trajectory_selection_list = bpy.props.CollectionProperty( - type=TrajectorySelectionItem -) +bpy.types.Scene.trajectory_selection_list = bpy.props.CollectionProperty(type=TrajectorySelectionItem) class MN_UL_TrajectorySelectionListUI(bpy.types.UIList): """UI List""" - def draw_item( - self, context, layout, data, item, icon, active_data, active_propname, index - ): + def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index): custom_icon = "VIS_SEL_11" if self.layout_type in {"DEFAULT", "COMPACT"}: diff --git a/molecularnodes/io/parse/bcif.py b/molecularnodes/io/parse/bcif.py index 233b5c8e..362c6ae0 100644 --- a/molecularnodes/io/parse/bcif.py +++ b/molecularnodes/io/parse/bcif.py @@ -52,14 +52,7 @@ def _atom_array_from_bcif(open_bcif): # other fields should be single self-contained fields mol = AtomArray(n_atoms) coord_field_names = [f"Cartn_{axis}" for axis in "xyz"] - mol.coord = np.hstack( - list( - [ - np.array(atom_site[column]).reshape((n_atoms, 1)) - for column in coord_field_names - ] - ) - ) + mol.coord = np.hstack(list([np.array(atom_site[column]).reshape((n_atoms, 1)) for column in coord_field_names])) # the list of current atom_site_lookup = { @@ -118,9 +111,7 @@ def _get_ops_from_bcif(open_bcif): is_petworld = False cats = open_bcif.data_blocks[0] assembly_gen = cats["pdbx_struct_assembly_gen"] - gen_arr = np.column_stack( - list([assembly_gen[name] for name in assembly_gen.field_names]) - ) + gen_arr = np.column_stack(list([assembly_gen[name] for name in assembly_gen.field_names])) dtype = [ ("assembly_id", int), ("chain_id", "U10"), @@ -148,12 +139,8 @@ def _get_ops_from_bcif(open_bcif): print("PetWorld!") is_petworld = True op_ids = np.array(ops["id"]) - struct_ops = np.column_stack( - list([np.array(ops[name]).reshape((ops.row_count, 1)) for name in ok_names]) - ) - rotations = np.array( - list([rotation_from_matrix(x[0:9].reshape((3, 3))) for x in struct_ops]) - ) + struct_ops = np.column_stack(list([np.array(ops[name]).reshape((ops.row_count, 1)) for name in ok_names])) + rotations = np.array(list([rotation_from_matrix(x[0:9].reshape((3, 3))) for x in struct_ops])) translations = struct_ops[:, 9:12] gen_list = [] @@ -325,19 +312,13 @@ def _decode_fixed_point(data: np.ndarray, encoding: FixedPointEncoding) -> np.nd return np.array(data, dtype=_get_dtype(encoding["srcType"])) / encoding["factor"] -def _decode_interval_quantization( - data: np.ndarray, encoding: IntervalQuantizationEncoding -) -> np.ndarray: +def _decode_interval_quantization(data: np.ndarray, encoding: IntervalQuantizationEncoding) -> np.ndarray: delta = (encoding["max"] - encoding["min"]) / (encoding["numSteps"] - 1) - return ( - np.array(data, dtype=_get_dtype(encoding["srcType"])) * delta + encoding["min"] - ) + return np.array(data, dtype=_get_dtype(encoding["srcType"])) * delta + encoding["min"] def _decode_run_length(data: np.ndarray, encoding: RunLengthEncoding) -> np.ndarray: - return np.repeat( - np.array(data[::2], dtype=_get_dtype(encoding["srcType"])), repeats=data[1::2] - ) + return np.repeat(np.array(data[::2], dtype=_get_dtype(encoding["srcType"])), repeats=data[1::2]) def _decode_delta(data: np.ndarray, encoding: DeltaEncoding) -> np.ndarray: @@ -347,9 +328,7 @@ def _decode_delta(data: np.ndarray, encoding: DeltaEncoding) -> np.ndarray: return np.cumsum(result, out=result) -def _decode_integer_packing_signed( - data: np.ndarray, encoding: IntegerPackingEncoding -) -> np.ndarray: +def _decode_integer_packing_signed(data: np.ndarray, encoding: IntegerPackingEncoding) -> np.ndarray: upper_limit = 0x7F if encoding["byteCount"] == 1 else 0x7FFF lower_limit = -upper_limit - 1 n = len(data) @@ -370,9 +349,7 @@ def _decode_integer_packing_signed( return output -def _decode_integer_packing_unsigned( - data: np.ndarray, encoding: IntegerPackingEncoding -) -> np.ndarray: +def _decode_integer_packing_unsigned(data: np.ndarray, encoding: IntegerPackingEncoding) -> np.ndarray: upper_limit = 0xFF if encoding["byteCount"] == 1 else 0xFFFF n = len(data) output = np.zeros(encoding["srcSize"], dtype="i4") @@ -392,9 +369,7 @@ def _decode_integer_packing_unsigned( return output -def _decode_integer_packing( - data: np.ndarray, encoding: IntegerPackingEncoding -) -> np.ndarray: +def _decode_integer_packing(data: np.ndarray, encoding: IntegerPackingEncoding) -> np.ndarray: if len(data) == encoding["srcSize"]: return data if encoding["isUnsigned"]: @@ -404,9 +379,7 @@ def _decode_integer_packing( def _decode_string_array(data: bytes, encoding: StringArrayEncoding) -> List[str]: - offsets = _decode( - EncodedData(encoding=encoding["offsetEncoding"], data=encoding["offsets"]) - ) + offsets = _decode(EncodedData(encoding=encoding["offsetEncoding"], data=encoding["offsets"])) indices = _decode(EncodedData(encoding=encoding["dataEncoding"], data=data)) str = encoding["stringData"] @@ -496,12 +469,8 @@ def __contains__(self, key: str): def __init__(self, category: EncodedCategory, lazy: bool): self.field_names = [c["name"] for c in category["columns"]] - self._field_cache = { - c["name"]: None if lazy else _decode_column(c) for c in category["columns"] - } - self._columns: Dict[str, EncodedColumn] = { - c["name"]: c for c in category["columns"] - } + self._field_cache = {c["name"]: None if lazy else _decode_column(c) for c in category["columns"]} + self._columns: Dict[str, EncodedColumn] = {c["name"]: c for c in category["columns"]} self.row_count = category["rowCount"] self.name = category["name"][1:] @@ -527,17 +496,9 @@ def __getitem__(self, index_or_name: Union[int, str]): Access a data block by index or header (case sensitive) """ if isinstance(index_or_name, str): - return ( - self._block_map[index_or_name] - if index_or_name in self._block_map - else None - ) + return self._block_map[index_or_name] if index_or_name in self._block_map else None else: - return ( - self.data_blocks[index_or_name] - if index_or_name < len(self.data_blocks) - else None - ) + return self.data_blocks[index_or_name] if index_or_name < len(self.data_blocks) else None def __len__(self): return len(self.data_blocks) @@ -566,17 +527,12 @@ def loads(data: Union[bytes, EncodedFile], lazy=True) -> CifFile: """ import msgpack - file: EncodedFile = ( - data if isinstance(data, dict) and "dataBlocks" in data else msgpack.loads(data) - ) # type: ignore + file: EncodedFile = data if isinstance(data, dict) and "dataBlocks" in data else msgpack.loads(data) # type: ignore data_blocks = [ CifDataBlock( header=block["header"], - categories={ - cat["name"][1:]: CifCategory(category=cat, lazy=lazy) - for cat in block["categories"] - }, + categories={cat["name"][1:]: CifCategory(category=cat, lazy=lazy) for cat in block["categories"]}, ) for block in file["dataBlocks"] ] diff --git a/molecularnodes/io/parse/cellpack.py b/molecularnodes/io/parse/cellpack.py index f8bf0ee1..02ff3e55 100644 --- a/molecularnodes/io/parse/cellpack.py +++ b/molecularnodes/io/parse/cellpack.py @@ -50,9 +50,7 @@ def _read(self, file_path): return data - def _create_object_instances( - self, name: str = "CellPack", node_setup: bool = True - ) -> bpy.types.Collection: + def _create_object_instances(self, name: str = "CellPack", node_setup: bool = True) -> bpy.types.Collection: collection = bl.coll.cellpack(name) if self.file_type == "cif": @@ -77,18 +75,14 @@ def _create_object_instances( ) if node_setup: - bl.nodes.create_starting_node_tree( - model, name=f"MN_pack_instance_{name}", set_color=False - ) + bl.nodes.create_starting_node_tree(model, name=f"MN_pack_instance_{name}", set_color=False) self.data_collection = collection return collection def _create_data_object(self, name="DataObject"): - data_object = bl.obj.create_data_object( - self.transformations, name=name, collection=bl.coll.mn() - ) + data_object = bl.obj.create_data_object(self.transformations, name=name, collection=bl.coll.mn()) data_object["chain_ids"] = self.chain_ids @@ -100,9 +94,7 @@ def _setup_node_tree(self, name="CellPack", fraction=1.0, as_points=False): group = bl.nodes.new_tree(name=f"MN_ensemble_{name}", fallback=False) mod.node_group = group - node_pack = bl.nodes.add_custom( - group, "MN_pack_instances", location=[-100, 0] - ) + node_pack = bl.nodes.add_custom(group, "MN_pack_instances", location=[-100, 0]) node_pack.inputs["Collection"].default_value = self.data_collection node_pack.inputs["Fraction"].default_value = fraction node_pack.inputs["As Points"].default_value = as_points diff --git a/molecularnodes/io/parse/cif.py b/molecularnodes/io/parse/cif.py index 4f085c56..e1c8d835 100644 --- a/molecularnodes/io/parse/cif.py +++ b/molecularnodes/io/parse/cif.py @@ -11,9 +11,7 @@ def __init__(self, file_path, extra_fields=None, sec_struct=True): super().__init__() self.file_path = file_path self.file = self._read() - self.array = self._get_structure( - extra_fields=extra_fields, sec_struct=sec_struct - ) + self.array = self._get_structure(extra_fields=extra_fields, sec_struct=sec_struct) self.n_atoms = self.array.array_length() def _read(self): @@ -35,9 +33,7 @@ def _get_structure(self, extra_fields: str = None, sec_struct=True, bonds=True): try: array = pdbx.get_structure(self.file, extra_fields=extra_fields) try: - array.set_annotation( - "sec_struct", _get_secondary_structure(array, self.file) - ) + array.set_annotation("sec_struct", _get_secondary_structure(array, self.file)) except KeyError: warnings.warn("No secondary structure information.") try: @@ -51,9 +47,7 @@ def _get_structure(self, extra_fields: str = None, sec_struct=True, bonds=True): # pdbx files don't seem to have bond information defined, so connect them based # on their residue names if not array.bonds and bonds: - array.bonds = struc.bonds.connect_via_residue_names( - array, inter_residue=True - ) + array.bonds = struc.bonds.connect_via_residue_names(array, inter_residue=True) return array @@ -184,9 +178,7 @@ def _get_entity_id(array, file): idx.append(i) entity_lookup = dict(zip(chains, idx)) - chain_id_int = np.array( - [entity_lookup.get(chain, -1) for chain in array.chain_id], int - ) + chain_id_int = np.array([entity_lookup.get(chain, -1) for chain in array.chain_id], int) return chain_id_int @@ -280,14 +272,9 @@ def _get_transformations(struct_oper): transformation_dict = {} for index, id in enumerate(struct_oper["id"]): rotation_matrix = np.array( - [ - [float(struct_oper[f"matrix[{i}][{j}]"][index]) for j in (1, 2, 3)] - for i in (1, 2, 3) - ] - ) - translation_vector = np.array( - [float(struct_oper[f"vector[{i}]"][index]) for i in (1, 2, 3)] + [[float(struct_oper[f"matrix[{i}][{j}]"][index]) for j in (1, 2, 3)] for i in (1, 2, 3)] ) + translation_vector = np.array([float(struct_oper[f"vector[{i}]"][index]) for i in (1, 2, 3)]) transformation_dict[id] = (rotation_matrix, translation_vector) return transformation_dict @@ -313,9 +300,7 @@ def _parse_operation_expression(expression): for gexpr in expr.split(","): if "-" in gexpr: first, last = gexpr.split("-") - operations.append( - [str(id) for id in range(int(first), int(last) + 1)] - ) + operations.append([str(id) for id in range(int(first), int(last) + 1)]) else: operations.append([gexpr]) else: diff --git a/molecularnodes/io/parse/ensemble.py b/molecularnodes/io/parse/ensemble.py index cdd00e58..dd1ea37c 100644 --- a/molecularnodes/io/parse/ensemble.py +++ b/molecularnodes/io/parse/ensemble.py @@ -79,8 +79,6 @@ def get_attribute(self, name="position", evaluate=False) -> np.ndarray | None: The value of the attribute. """ if not self.object: - warnings.warn( - "No object yet created. Use `create_model()` to create a corresponding object." - ) + warnings.warn("No object yet created. Use `create_model()` to create a corresponding object.") return None return bl.obj.get_attribute(self.object, name=name, evaluate=evaluate) diff --git a/molecularnodes/io/parse/mda.py b/molecularnodes/io/parse/mda.py index c8df734c..0ab7f7e8 100644 --- a/molecularnodes/io/parse/mda.py +++ b/molecularnodes/io/parse/mda.py @@ -33,9 +33,7 @@ class MockUniverse: class AtomGroupInBlender: - def __init__( - self, ag: mda.AtomGroup, style: str = "vdw", world_scale: float = 0.01 - ): + def __init__(self, ag: mda.AtomGroup, style: str = "vdw", world_scale: float = 0.01): """ AtomGroup in Blender. It will be dynamically updated when the frame changes or @@ -153,9 +151,7 @@ def elements(self) -> List[str]: except AttributeError: # If 'elements' attribute doesn't exist try: elements = [ - x - if x in data.elements.keys() - else mda.topology.guessers.guess_atom_element(x) + x if x in data.elements.keys() else mda.topology.guessers.guess_atom_element(x) for x in self.ag.atoms.names ] except ( @@ -168,22 +164,14 @@ def elements(self) -> List[str]: @property def atomic_number(self) -> np.ndarray: return np.array( - [ - data.elements.get(element, data.elements.get("X")).get("atomic_number") - for element in self.elements - ] + [data.elements.get(element, data.elements.get("X")).get("atomic_number") for element in self.elements] ) @property def vdw_radii(self) -> np.ndarray: # pm to Angstrom return ( - np.array( - [ - data.elements.get(element, {}).get("vdw_radii", 100) - for element in self.elements - ] - ) + np.array([data.elements.get(element, {}).get("vdw_radii", 100) for element in self.elements]) * 0.01 * self.world_scale ) @@ -195,12 +183,7 @@ def mass(self) -> np.ndarray: masses = np.array([x.mass for x in self.ag.atoms]) except mda.exceptions.NoDataError: masses = np.array( - [ - data.elements.get(element, {"standard_mass": 0}).get( - "standard_mass" - ) - for element in self.elements - ] + [data.elements.get(element, {"standard_mass": 0}).get("standard_mass") for element in self.elements] ) return masses @@ -215,12 +198,7 @@ def res_name(self) -> np.ndarray: @property def res_num(self) -> np.ndarray: return np.array( - [ - data.residues.get(res_name, data.residues.get("UNK")).get( - "res_name_num" - ) - for res_name in self.res_name - ] + [data.residues.get(res_name, data.residues.get("UNK")).get("res_name_num") for res_name in self.res_name] ) @property @@ -256,9 +234,7 @@ def atom_type_unique(self) -> np.ndarray: @property def atom_type_num(self) -> np.ndarray: - atom_type_unique, atom_type_index = np.unique( - self.atom_type, return_inverse=True - ) + atom_type_unique, atom_type_index = np.unique(self.atom_type, return_inverse=True) return atom_type_index @property @@ -271,9 +247,7 @@ def atom_name(self) -> np.ndarray: @property def atom_name_num(self) -> np.ndarray: if hasattr(self.ag, "names"): - return np.array( - list(map(lambda x: data.atom_names.get(x, -1), self.atom_name)) - ) + return np.array(list(map(lambda x: data.atom_names.get(x, -1), self.atom_name))) else: return np.repeat(-1, self.ag.n_atoms) @@ -299,9 +273,7 @@ def is_alpha_carbon(self) -> np.ndarray: @property def is_solvent(self) -> np.ndarray: - return self.bool_selection( - self.ag, "name OW or name HW1 or name HW2 or resname W or resname PW" - ) + return self.bool_selection(self.ag, "name OW or name HW1 or name HW2 or resname W or resname PW") @property def _attributes_2_blender(self): @@ -466,12 +438,8 @@ def __init__(self, world_scale: float = 0.01, in_memory: bool = False): if in_memory: return bpy.types.Scene.mda_session = self - bpy.app.handlers.frame_change_post.append( - self._update_trajectory_handler_wrapper() - ) - bpy.app.handlers.depsgraph_update_pre.append( - self._update_style_handler_wrapper() - ) + bpy.app.handlers.frame_change_post.append(self._update_trajectory_handler_wrapper()) + bpy.app.handlers.depsgraph_update_pre.append(self._update_style_handler_wrapper()) log.info("MDAnalysis session is initialized.") @property @@ -548,9 +516,7 @@ def show( custom_selections=custom_selections, ) if frame_mapping is not None: - warnings.warn( - "Custom frame_mapping not supported" "when in_memory is on." - ) + warnings.warn("Custom frame_mapping not supported" "when in_memory is on.") if subframes != 0: warnings.warn("Custom subframes not supported" "when in_memory is on.") log.info(f"{atoms} is loaded in memory.") @@ -562,9 +528,7 @@ def show( # if any frame_mapping is out of range, then raise an error if frame_mapping and (len(frame_mapping) > universe.trajectory.n_frames): - raise ValueError( - "one or more mapping values are" "out of range for the trajectory" - ) + raise ValueError("one or more mapping values are" "out of range for the trajectory") mol_object = self._process_atomgroup( ag=atoms, @@ -677,15 +641,11 @@ def in_memory( print("KeyError: 'occupancy' not found in ts.data") add_occupancy = False except TypeError: - print( - "TypeError: ts.data is not a dictionary or similar mapping type" - ) + print("TypeError: ts.data is not a dictionary or similar mapping type") add_occupancy = False # disable the frames collection from the viewer - bpy.context.view_layer.layer_collection.children[coll.mn().name].children[ - coll_frames.name - ].exclude = True + bpy.context.view_layer.layer_collection.children[coll.mn().name].children[coll_frames.name].exclude = True if node_setup: nodes.create_starting_node_tree( @@ -698,9 +658,7 @@ def in_memory( return mol_object - def transfer_to_memory( - self, start=None, stop=None, step=None, verbose=False, **kwargs - ): + def transfer_to_memory(self, start=None, stop=None, step=None, verbose=False, **kwargs): """ Transfer the trajectories in the session to memory. This is an alternative way to make sure the blender session is @@ -732,9 +690,7 @@ def transfer_to_memory( for rep_name in self.rep_names: universe = self.universe_reps[rep_name]["universe"] - universe.transfer_to_memory( - start=start, stop=stop, step=step, verbose=verbose, **kwargs - ) + universe.transfer_to_memory(start=start, stop=stop, step=step, verbose=verbose, **kwargs) log.info("The trajectories in this session is transferred to memory.") def _process_atomgroup( @@ -767,9 +723,7 @@ def _process_atomgroup( return_object : bool Whether to return the blender object or not. Default: False """ - ag_blender = AtomGroupInBlender( - ag=ag, style=style, world_scale=self.world_scale - ) + ag_blender = AtomGroupInBlender(ag=ag, style=style, world_scale=self.world_scale) # create the initial model mol_object = obj.create_object( name=name, @@ -780,9 +734,7 @@ def _process_atomgroup( # add the attributes for the model in blender for att_name, att in ag_blender._attributes_2_blender.items(): - obj.set_attribute( - mol_object, att_name, att["value"], att["type"], att["domain"] - ) + obj.set_attribute(mol_object, att_name, att["value"], att["type"], att["domain"]) mol_object["chain_ids"] = ag_blender.chain_ids mol_object["atom_type_unique"] = ag_blender.atom_type_unique mol_object.mn["subframes"] = subframes @@ -795,9 +747,7 @@ def _process_atomgroup( # instead, the name generated by blender is used. if mol_object.name != name: warnings.warn( - "The name of the object is changed to {} because {} is already used.".format( - mol_object.name, name - ) + "The name of the object is changed to {} because {} is already used.".format(mol_object.name, name) ) self.atom_reps[mol_object.name] = ag_blender @@ -879,9 +829,7 @@ def _update_trajectory(self, frame): mol_object.data.clear_geometry() mol_object.data.from_pydata(ag_rep.positions, ag_rep.bonds, faces=[]) for att_name, att in ag_rep._attributes_2_blender.items(): - obj.set_attribute( - mol_object, att_name, att["value"], att["type"], att["domain"] - ) + obj.set_attribute(mol_object, att_name, att["value"], att["type"], att["domain"]) mol_object["chain_id"] = ag_rep.chain_ids mol_object["atom_type_unique"] = ag_rep.atom_type_unique mol_object.mn["subframes"] = subframes @@ -954,12 +902,8 @@ def _rejuvenate(cls, mol_objects): cls = pickle.load(f) except FileNotFoundError: return None - bpy.app.handlers.frame_change_post.append( - cls._update_trajectory_handler_wrapper() - ) - bpy.app.handlers.depsgraph_update_pre.append( - cls._update_style_handler_wrapper() - ) + bpy.app.handlers.frame_change_post.append(cls._update_trajectory_handler_wrapper()) + bpy.app.handlers.depsgraph_update_pre.append(cls._update_style_handler_wrapper()) log.info("MDAnalysis session is loaded from {}".format(blend_file_name)) return cls diff --git a/molecularnodes/io/parse/molecule.py b/molecularnodes/io/parse/molecule.py index 2737b7ea..7a6436f2 100644 --- a/molecularnodes/io/parse/molecule.py +++ b/molecularnodes/io/parse/molecule.py @@ -1,5 +1,5 @@ from abc import ABCMeta -from typing import Optional, Any, Tuple, Union, List +from typing import Optional, Any, Tuple, Union, List, Dict import biotite.structure from numpy.typing import NDArray from pathlib import Path @@ -132,9 +132,7 @@ def set_attribute( overwrite=overwrite, ) - def get_attribute( - self, name: str = "position", evaluate: bool = False - ) -> np.ndarray: + def get_attribute(self, name: str = "position", evaluate: bool = False) -> np.ndarray: """ Get the value of an attribute for the associated object. @@ -192,9 +190,7 @@ def centre(self, centre_type: str = "centroid") -> np.ndarray: mass = self.get_attribute(name="mass") return bl.obj.centre_weighted(positions, mass) else: - raise ValueError( - f"`{centre_type}` not a supported selection of ['centroid', 'mass']" - ) + raise ValueError(f"`{centre_type}` not a supported selection of ['centroid', 'mass']") def create_model( self, @@ -290,9 +286,7 @@ def create_model( return model - def assemblies( - self, as_array: bool = False - ) -> Dict[str, List[float]] | None: + def assemblies(self, as_array: bool = False) -> Dict[str, List[float]] | None: """ Get the biological assemblies of the molecule. @@ -348,25 +342,16 @@ def _create_model( array = array[mask] try: - mass = np.array( - [ - data.elements.get(x, {}).get("standard_mass", 0.0) - for x in np.char.title(array.element) - ] - ) + mass = np.array([data.elements.get(x, {}).get("standard_mass", 0.0) for x in np.char.title(array.element)]) array.set_annotation("mass", mass) except AttributeError: pass - def centre_array( - atom_array: biotite.structure.AtomArray, centre: str - ) -> None: + def centre_array(atom_array: biotite.structure.AtomArray, centre: str) -> None: if centre == "centroid": atom_array.coord -= bl.obj.centre(atom_array.coord) elif centre == "mass": - atom_array.coord -= bl.obj.centre_weighted( - array=atom_array.coord, weight=atom_array.mass - ) + atom_array.coord -= bl.obj.centre_weighted(array=atom_array.coord, weight=atom_array.mass) if centre in ["mass", "centroid"]: if is_stack: @@ -405,9 +390,7 @@ def centre_array( # '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.obj.set_attribute( - mol, name="bond_type", data=bond_types, type="INT", domain="EDGE" - ) + bl.obj.set_attribute(mol, name="bond_type", data=bond_types, type="INT", domain="EDGE") # The attributes for the model are initially defined as single-use functions. This allows # for a loop that attempts to add each attibute by calling the function. Only during this @@ -420,12 +403,7 @@ def centre_array( def att_atomic_number() -> NDArray[np.int32]: atomic_number = np.array( - [ - data.elements.get(x, {"atomic_number": -1}).get( - "atomic_number" - ) - for x in np.char.title(array.element) - ] + [data.elements.get(x, {"atomic_number": -1}).get("atomic_number") for x in np.char.title(array.element)] ) return atomic_number @@ -444,26 +422,16 @@ def att_res_name() -> NDArray[np.int32]: res_nums = [] for name in res_names: - res_num = data.residues.get(name, {"res_name_num": -1}).get( - "res_name_num" - ) + res_num = data.residues.get(name, {"res_name_num": -1}).get("res_name_num") if res_num == 9999: - if ( - res_names[counter - 1] != name - or res_ids[counter] != res_ids[counter - 1] - ): + if res_names[counter - 1] != name or res_ids[counter] != res_ids[counter - 1]: id_counter += 1 unique_res_name = str(id_counter + 100) + "_" + str(name) other_res.append(unique_res_name) - num = ( - np.where(np.isin(np.unique(other_res), unique_res_name))[ - 0 - ][0] - + 100 - ) + num = np.where(np.isin(np.unique(other_res), unique_res_name))[0][0] + 100 res_nums.append(num) else: res_nums.append(res_num) @@ -490,8 +458,7 @@ def att_vdw_radii() -> NDArray[np.float64]: map( # divide by 100 to convert from picometres to angstroms which is # what all of coordinates are in - lambda x: data.elements.get(x, {}).get("vdw_radii", 100.0) - / 100, + lambda x: data.elements.get(x, {}).get("vdw_radii", 100.0) / 100, np.char.title(array.element), ) ) @@ -502,9 +469,7 @@ def att_mass() -> NDArray[np.float64]: return array.mass def att_atom_name() -> NDArray[np.int32]: - atom_name = np.array( - [data.atom_names.get(x, -1) for x in array.atom_name] - ) + atom_name = np.array([data.atom_names.get(x, -1) for x in array.atom_name]) return atom_name @@ -746,15 +711,11 @@ def att_sec_struct() -> NDArray[np.int32]: domain=att["domain"], ) if verbose: - print( - f'Added {att["name"]} after {time.process_time() - start} s' - ) - except: + print(f'Added {att["name"]} after {time.process_time() - start} s') + except ValueError: if verbose: warnings.warn(f"Unable to add attribute: {att['name']}") - print( - f'Failed adding {att["name"]} after {time.process_time() - start} s' - ) + print(f'Failed adding {att["name"]} after {time.process_time() - start} s') coll_frames = None if frames: diff --git a/molecularnodes/io/parse/mrc.py b/molecularnodes/io/parse/mrc.py index ed1f2f7c..facd8e16 100644 --- a/molecularnodes/io/parse/mrc.py +++ b/molecularnodes/io/parse/mrc.py @@ -19,13 +19,9 @@ def __init__(self, file_path, center=False, invert=False, overwrite=False): super().__init__(self) self.file_path = file_path self.grid = self.map_to_grid(self.file_path, center=center) - self.file_vdb = self.map_to_vdb( - self.file_path, center=center, invert=invert, overwrite=overwrite - ) + self.file_vdb = self.map_to_vdb(self.file_path, center=center, invert=invert, overwrite=overwrite) - def create_model( - self, name="NewDensity", style="density_surface", setup_nodes=True - ) -> bpy.types.Object: + def create_model(self, name="NewDensity", style="density_surface", setup_nodes=True) -> bpy.types.Object: """ Loads an MRC file into Blender as a volumetric object. @@ -53,9 +49,7 @@ def create_model( object.name = name if setup_nodes: - nodes.create_starting_nodes_density( - object=object, style=style, threshold=self.threshold - ) + nodes.create_starting_nodes_density(object=object, style=style, threshold=self.threshold) return object @@ -181,9 +175,7 @@ def map_to_grid(self, file: str, invert: bool = False, center: bool = False): try: grid.copyFromArray(volume.astype(float)) except Exception as e: - print( - f"Grid data type '{volume.dtype}' is an unsupported type.\nError: {e}" - ) + print(f"Grid data type '{volume.dtype}' is an unsupported type.\nError: {e}") grid.gridClass = vdb.GridClass.FOG_VOLUME grid.name = "density" diff --git a/molecularnodes/io/parse/pdb.py b/molecularnodes/io/parse/pdb.py index 82c8b12a..104c8d69 100644 --- a/molecularnodes/io/parse/pdb.py +++ b/molecularnodes/io/parse/pdb.py @@ -1,6 +1,6 @@ import numpy as np -from typing import List, Union, AnyStr +from typing import Union, AnyStr from pathlib import Path from .assembly import AssemblyParser @@ -51,9 +51,7 @@ def _get_sec_struct(file, array): lines_helix = lines[np.char.startswith(lines, "HELIX")] lines_sheet = lines[np.char.startswith(lines, "SHEET")] if len(lines_helix) == 0 and len(lines_sheet) == 0: - raise struc.BadStructureError( - "No secondary structure information detected." - ) + raise struc.BadStructureError("No secondary structure information detected.") sec_struct = np.zeros(array.array_length(), int) @@ -80,9 +78,7 @@ def _get_mask(line, start1, end1, start2, end2, chainid): # create a mask for the array based on these values mask = np.logical_and( - np.logical_and( - array.chain_id == chain_id, array.res_id >= start_num - ), + np.logical_and(array.chain_id == chain_id, array.res_id >= start_num), array.res_id <= end_num, ) @@ -95,9 +91,7 @@ def _get_mask(line, start1, end1, start2, end2, chainid): # assign remaining AA atoms to 3 (loop), while all other remaining # atoms will be 0 (not relevant) - mask = np.logical_and( - sec_struct == 0, struc.filter_canonical_amino_acids(array) - ) + mask = np.logical_and(sec_struct == 0, struc.filter_canonical_amino_acids(array)) sec_struct[mask] = 3 @@ -122,9 +116,7 @@ def _comp_secondary_structure(array): conv_sse_char_int = {"a": 1, "b": 2, "c": 3, "": 0} char_sse = annotate_sse(array) - int_sse = np.array( - [conv_sse_char_int[char] for char in char_sse], dtype=int - ) + int_sse = np.array([conv_sse_char_int[char] for char in char_sse], dtype=int) atom_sse = spread_residue_wise(array, int_sse) return atom_sse @@ -145,9 +137,7 @@ def get_transformations(self, assembly_id): # Get lines containing transformations for assemblies remark_lines = self._file.get_remark(350) if remark_lines is None: - raise biotite.InvalidFileError( - "File does not contain assembly information (REMARK 350)" - ) + raise biotite.InvalidFileError("File does not contain assembly information (REMARK 350)") # Get lines corresponding to selected assembly ID assembly_start_i = None assembly_stop_i = None @@ -171,9 +161,7 @@ def get_transformations(self, assembly_id): # Get transformations for a sets of chains transformations = [] chain_set_start_indices = [ - i - for i, line in enumerate(assembly_lines) - if line.startswith("APPLY THE FOLLOWING TO CHAINS") + i for i, line in enumerate(assembly_lines) if line.startswith("APPLY THE FOLLOWING TO CHAINS") ] # Add exclusive stop at end of records chain_set_start_indices.append(len(assembly_lines)) @@ -184,12 +172,10 @@ def get_transformations(self, assembly_id): affected_chain_ids = [] transform_start = None for j, line in enumerate(assembly_lines[start:stop]): - if line.startswith( - "APPLY THE FOLLOWING TO CHAINS:" - ) or line.startswith(" AND CHAINS:"): - affected_chain_ids += [ - chain_id.strip() for chain_id in line[30:].split(",") - ] + if line.startswith("APPLY THE FOLLOWING TO CHAINS:") or line.startswith( + " AND CHAINS:" + ): + affected_chain_ids += [chain_id.strip() for chain_id in line[30:].split(",")] else: # Chain specification has finished # BIOMT lines start directly after chain specification @@ -197,13 +183,9 @@ def get_transformations(self, assembly_id): break # Parse transformations from BIOMT lines if transform_start is None: - raise biotite.InvalidFileError( - "No 'BIOMT' records found for chosen assembly" - ) + raise biotite.InvalidFileError("No 'BIOMT' records found for chosen assembly") - matrices = _parse_transformations( - assembly_lines[transform_start:stop] - ) + matrices = _parse_transformations(assembly_lines[transform_start:stop]) for matrix in matrices: transformations.append((affected_chain_ids, matrix.tolist())) @@ -228,9 +210,7 @@ def _parse_transformations(lines): # Each transformation requires 3 lines for the (x,y,z) components if len(lines) % 3 != 0: - raise biotite.InvalidFileError( - "Invalid number of transformation vectors" - ) + raise biotite.InvalidFileError("Invalid number of transformation vectors") n_transformations = len(lines) // 3 matrices = np.tile(np.identity(4), (n_transformations, 1, 1)) @@ -243,9 +223,7 @@ def _parse_transformations(lines): transformations = [float(e) for e in line.split()[2:]] if len(transformations) != 4: - raise biotite.InvalidFileError( - "Invalid number of transformation vector elements" - ) + raise biotite.InvalidFileError("Invalid number of transformation vector elements") matrices[transformation_i, component_i, :] = transformations component_i += 1 diff --git a/molecularnodes/io/parse/pdbx.py b/molecularnodes/io/parse/pdbx.py index 87ada18f..f722e600 100644 --- a/molecularnodes/io/parse/pdbx.py +++ b/molecularnodes/io/parse/pdbx.py @@ -1,7 +1,6 @@ import numpy as np import warnings import itertools -from typing import List from .molecule import Molecule @@ -13,12 +12,7 @@ def __init__(self, file_path): @property def entity_ids(self): - return ( - self.file.block.get("entity") - .get("pdbx_description") - .as_array() - .tolist() - ) + return self.file.block.get("entity").get("pdbx_description").as_array().tolist() def _get_entity_id(self, array, file): chain_ids = file.block["entity_poly"]["pdbx_strand_id"].as_array() @@ -35,14 +29,10 @@ def _get_entity_id(self, array, file): idx.append(i) entity_lookup = dict(zip(chains, idx)) - chain_id_int = np.array( - [entity_lookup.get(chain, -1) for chain in array.chain_id], int - ) + chain_id_int = np.array([entity_lookup.get(chain, -1) for chain in array.chain_id], int) return chain_id_int - def get_structure( - self, extra_fields=["b_factor", "occupancy", "atom_id"], bonds=True - ): + def get_structure(self, extra_fields=["b_factor", "occupancy", "atom_id"], bonds=True): import biotite.structure.io.pdbx as pdbx import biotite.structure as struc @@ -55,16 +45,12 @@ def get_structure( except KeyError: warnings.warn("No secondary structure information.") try: - array.set_annotation( - "entity_id", self._get_entity_id(array, self.file) - ) + array.set_annotation("entity_id", self._get_entity_id(array, self.file)) except KeyError: warnings.warn("No entity ID information") if not array.bonds and bonds: - array.bonds = struc.bonds.connect_via_residue_names( - array, inter_residue=True - ) + array.bonds = struc.bonds.connect_via_residue_names(array, inter_residue=True) return array @@ -87,9 +73,7 @@ def _extract_matrices(self, category): "vector[3]", ] - columns = [ - category[name].as_array().astype(float) for name in matrix_columns - ] + columns = [category[name].as_array().astype(float) for name in matrix_columns] matrices = np.empty((len(columns[0]), 4, 4), float) col_mask = np.tile((0, 1, 2, 3), 3) @@ -148,15 +132,9 @@ def _get_secondary_structure(self, file, array): # as normalquit sheet = file.block.get("struct_sheet_range") if sheet: - starts = np.append( - starts, sheet["beg_auth_seq_id"].as_array().astype(int) - ) - ends = np.append( - ends, sheet["end_auth_seq_id"].as_array().astype(int) - ) - chains = np.append( - chains, sheet["end_auth_asym_id"].as_array().astype(str) - ) + starts = np.append(starts, sheet["beg_auth_seq_id"].as_array().astype(int)) + ends = np.append(ends, sheet["end_auth_seq_id"].as_array().astype(int)) + chains = np.append(chains, sheet["end_auth_asym_id"].as_array().astype(str)) id_label = np.append(id_label, np.repeat("STRN", len(sheet["id"]))) if not conf and not sheet: @@ -252,9 +230,7 @@ def get_transformations(self, assembly_id): struct_oper_category = self._file.block["pdbx_struct_oper_list"] - if assembly_id not in assembly_gen_category["assembly_id"].as_array( - str - ): + if assembly_id not in assembly_gen_category["assembly_id"].as_array(str): raise KeyError(f"File has no Assembly ID '{assembly_id}'") # Extract all possible transformations indexed by operation ID @@ -308,9 +284,7 @@ def _extract_matrices(category, scale=True): "vector[3]", ] - columns = [ - category[name].as_array().astype(float) for name in matrix_columns - ] + columns = [category[name].as_array().astype(float) for name in matrix_columns] n = 4 if scale else 3 matrices = np.empty((len(columns[0]), n, 4), float) @@ -354,17 +328,9 @@ def _get_transformations(struct_oper): transformation_dict = {} for index, id in enumerate(struct_oper["id"].as_array()): rotation_matrix = np.array( - [ - [ - float(struct_oper[f"matrix[{i}][{j}]"][index]) - for j in (1, 2, 3) - ] - for i in (1, 2, 3) - ] - ) - translation_vector = np.array( - [float(struct_oper[f"vector[{i}]"][index]) for i in (1, 2, 3)] + [[float(struct_oper[f"matrix[{i}][{j}]"][index]) for j in (1, 2, 3)] for i in (1, 2, 3)] ) + translation_vector = np.array([float(struct_oper[f"vector[{i}]"][index]) for i in (1, 2, 3)]) transformation_dict[id] = (rotation_matrix, translation_vector) return transformation_dict @@ -390,20 +356,13 @@ def _parse_operation_expression(expression): for gexpr in expr.split(","): if "-" in gexpr: first, last = gexpr.split("-") - operations.append( - [ - str(id) - for id in range(int(first), int(last) + 1) - ] - ) + operations.append([str(id) for id in range(int(first), int(last) + 1)]) else: operations.append([gexpr]) else: # Range of operation IDs, they must be integers first, last = expr.split("-") - operations.append( - [str(id) for id in range(int(first), int(last) + 1)] - ) + operations.append([str(id) for id in range(int(first), int(last) + 1)]) elif "," in expr: # List of operation IDs operations.append(expr.split(",")) diff --git a/molecularnodes/io/parse/star.py b/molecularnodes/io/parse/star.py index 73088568..5e337cf6 100644 --- a/molecularnodes/io/parse/star.py +++ b/molecularnodes/io/parse/star.py @@ -60,11 +60,7 @@ 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 - ): + 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" @@ -83,9 +79,7 @@ def _create_mn_columns(self): if "rlnCoordinateZ" not in df: df["rlnCoordinateZ"] = 0 - self.positions = df[ - ["rlnCoordinateX", "rlnCoordinateY", "rlnCoordinateZ"] - ].to_numpy() + 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 = [ @@ -101,14 +95,10 @@ def _create_mn_columns(self): 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 @@ -117,9 +107,7 @@ def _create_mn_columns(self): elif self.star_type == "cistem": df = self.data df["cisTEMZFromDefocus"] = (df["cisTEMDefocus1"] + df["cisTEMDefocus2"]) / 2 - df["cisTEMZFromDefocus"] = ( - df["cisTEMZFromDefocus"] - df["cisTEMZFromDefocus"].median() - ) + df["cisTEMZFromDefocus"] = df["cisTEMZFromDefocus"] - df["cisTEMZFromDefocus"].median() self.positions = df[ [ "cisTEMOriginalXPosition", @@ -131,11 +119,7 @@ def _create_mn_columns(self): df["MNAngleTheta"] = df["cisTEMAngleTheta"] df["MNAnglePsi"] = df["cisTEMAnglePsi"] df["MNPixelSize"] = df["cisTEMPixelSize"] - df["MNImageId"] = ( - df["cisTEMOriginalImageFilename"] - .astype("category") - .cat.codes.to_numpy() - ) + df["MNImageId"] = df["cisTEMOriginalImageFilename"].astype("category").cat.codes.to_numpy() def _convert_mrc_to_tiff(self): import mrcfile @@ -157,17 +141,11 @@ def _convert_mrc_to_tiff(self): 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 - ) + 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") @@ -196,9 +174,7 @@ def _update_micrograph_texture(self, *_): 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: return @@ -219,9 +195,7 @@ def _update_micrograph_texture(self, *_): 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 - ) + blender_object = bl.obj.create_object(self.positions * world_scale, collection=bl.coll.mn(), name=name) blender_object.mn["molecule_type"] = "star" @@ -240,19 +214,13 @@ def create_model(self, name="StarFileObject", node_setup=True, world_scale=0.01) # 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) - ) + 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 - ) + bl.nodes.create_starting_nodes_starfile(blender_object, n_images=self.n_images) self.node_group = blender_object.modifiers["MolecularNodes"].node_group blender_object["starfile_path"] = str(self.file_path) diff --git a/molecularnodes/io/retrieve.py b/molecularnodes/io/retrieve.py index 3c99cd2a..80de2b9d 100644 --- a/molecularnodes/io/retrieve.py +++ b/molecularnodes/io/retrieve.py @@ -55,9 +55,7 @@ def download( """ supported_formats = ["cif", "pdb", "bcif"] if format not in supported_formats: - raise ValueError( - f"File format '{format}' not in: {supported_formats=}" - ) + raise ValueError(f"File format '{format}' not in: {supported_formats=}") _is_binary = format in ["bcif"] filename = f"{code}.{format}" @@ -112,9 +110,7 @@ def _url(code: str, format: str, database: str = "rcsb") -> str: def get_alphafold_url(code: str, format: str) -> str: if format not in ["pdb", "cif", "bcif"]: - ValueError( - f"Format {format} not currently supported from AlphaFold databse." - ) + ValueError(f"Format {format} not currently supported from AlphaFold databse.") # we have to first query the database, then they'll return some JSON with a list # of metadata, some items of which will be the URLs for the computed models diff --git a/molecularnodes/io/star.py b/molecularnodes/io/star.py index e11793ac..d574434a 100644 --- a/molecularnodes/io/star.py +++ b/molecularnodes/io/star.py @@ -25,9 +25,7 @@ def load(file_path, name="NewStarInstances", node_setup=True, world_scale=0.01): class MN_OT_Import_Star_File(bpy.types.Operator): bl_idname = "mn.import_star_file" bl_label = "Load" - bl_description = ( - "Will import the given file, setting up the points to instance an object." - ) + bl_description = "Will import the given file, setting up the points to instance an object." bl_options = {"REGISTER"} @classmethod diff --git a/molecularnodes/io/wwpdb.py b/molecularnodes/io/wwpdb.py index 8a5e4cfb..d8418e21 100644 --- a/molecularnodes/io/wwpdb.py +++ b/molecularnodes/io/wwpdb.py @@ -1,6 +1,5 @@ from pathlib import Path from typing import Optional, Union, Set -from pathlib import Path import bpy @@ -134,9 +133,7 @@ def execute(self, context: bpy.types.Context) -> Set[str]: # the UI for the panel, which will display the operator and the properties -def panel( - layout: bpy.types.UILayout, scene: bpy.types.Scene -) -> bpy.types.UILayout: +def panel(layout: bpy.types.UILayout, scene: bpy.types.Scene) -> bpy.types.UILayout: layout.label(text="Download from PDB", icon="IMPORT") layout.separator() diff --git a/molecularnodes/pkg.py b/molecularnodes/pkg.py index 1b61cad4..b695bdde 100644 --- a/molecularnodes/pkg.py +++ b/molecularnodes/pkg.py @@ -246,9 +246,7 @@ def run_python(cmd_list: list = None, mirror_url: str = "", timeout: int = 600): log.info(f"Running Pip: '{cmd_list}'") # run the command and capture the output - result = subprocess.run( - cmd_list, timeout=timeout, stdout=subprocess.PIPE, stderr=subprocess.PIPE - ) + result = subprocess.run(cmd_list, timeout=timeout, stdout=subprocess.PIPE, stderr=subprocess.PIPE) if result.returncode != 0: log.error("Command failed: %s", cmd_list) @@ -365,9 +363,7 @@ def install_all_packages(pypi_mirror_provider: str = "Default") -> list: ) results.append(result) except InstallationError as e: - raise InstallationError( - f"Error installing package {pkg.get('name')}: {str(e)}" - ) + raise InstallationError(f"Error installing package {pkg.get('name')}: {str(e)}") return results @@ -380,13 +376,9 @@ class MN_OT_Install_Package(bpy.types.Operator): description="Python Package to Install", default="biotite", ) # type: ignore - version: bpy.props.StringProperty( - name="Python Package", description="Python Package to Install", default="0.36.1" - ) # type: ignore + version: bpy.props.StringProperty(name="Python Package", description="Python Package to Install", default="0.36.1") # type: ignore - description: bpy.props.StringProperty( - name="Operator description", default="Install specified python package." - ) # type: ignore + description: bpy.props.StringProperty(name="Operator description", default="Install specified python package.") # type: ignore @classmethod def description(cls, context, properties): @@ -399,9 +391,7 @@ def execute(self, context): pypi_mirror_provider=bpy.context.scene.pypi_mirror_provider, ) if result.returncode == 0 and is_current(self.package): - self.report( - {"INFO"}, f"Successfully installed {self.package} v{self.version}" - ) + self.report({"INFO"}, f"Successfully installed {self.package} v{self.version}") else: log_dir = os.path.join(os.path.abspath(ADDON_DIR), "logs") self.report( diff --git a/molecularnodes/props.py b/molecularnodes/props.py index f3a267a5..37ed26c4 100644 --- a/molecularnodes/props.py +++ b/molecularnodes/props.py @@ -37,9 +37,7 @@ subtype="NONE", default=0, ) -bpy.types.Scene.MN_import_build_assembly = BoolProperty( - name="Build Assembly", default=False -) +bpy.types.Scene.MN_import_build_assembly = BoolProperty(name="Build Assembly", default=False) bpy.types.Scene.MN_import_node_setup = BoolProperty( name="Setup Nodes", default=True, diff --git a/molecularnodes/ui/func.py b/molecularnodes/ui/func.py index 6b3bbdab..08e6ed9a 100644 --- a/molecularnodes/ui/func.py +++ b/molecularnodes/ui/func.py @@ -45,9 +45,7 @@ def button_custom_color(layout, label, field, prefix, property_id, starting_valu op.description = f"Choose individual colors for each {label}" -def button_custom_selection( - layout, label, field, prefix, property_id, starting_value=0 -): +def button_custom_selection(layout, label, field, prefix, property_id, starting_value=0): op = layout.operator("mn.selection_custom", text=label) op.field = field op.prefix = prefix diff --git a/molecularnodes/ui/node_info.py b/molecularnodes/ui/node_info.py index 8f0489dc..86c218f5 100644 --- a/molecularnodes/ui/node_info.py +++ b/molecularnodes/ui/node_info.py @@ -495,9 +495,7 @@ "description": "A sphere atom representation, visible in EEVEE and Cycles. Based on mesh instancing which slows down viewport performance", }, ], - "cellpack": [ - {"label": "Pack Instances", "name": "MN_pack_instances", "description": ""} - ], + "cellpack": [{"label": "Pack Instances", "name": "MN_pack_instances", "description": ""}], "density": [ { "label": "Style Surface", diff --git a/molecularnodes/ui/ops.py b/molecularnodes/ui/ops.py index 6f941f85..b1accdbe 100644 --- a/molecularnodes/ui/ops.py +++ b/molecularnodes/ui/ops.py @@ -7,9 +7,7 @@ class MN_OT_Add_Custom_Node_Group(bpy.types.Operator): bl_label = "Add Custom Node Group" # bl_description = "Add Molecular Nodes custom node group." bl_options = {"REGISTER", "UNDO"} - node_name: bpy.props.StringProperty( - name="node_name", description="", default="", subtype="NONE", maxlen=0 - ) + node_name: bpy.props.StringProperty(name="node_name", description="", default="", subtype="NONE", maxlen=0) node_label: bpy.props.StringProperty(name="node_label", default="") node_description: bpy.props.StringProperty( name="node_description", diff --git a/molecularnodes/ui/panel.py b/molecularnodes/ui/panel.py index 41f87f0b..be617268 100644 --- a/molecularnodes/ui/panel.py +++ b/molecularnodes/ui/panel.py @@ -126,9 +126,7 @@ def panel_object(layout, context): row = layout.row(align=True) row.label(text="Style") - current_style = nodes.format_node_name( - nodes.get_style_node(object).node_tree.name - ).replace("Style ", "") + current_style = nodes.format_node_name(nodes.get_style_node(object).node_tree.name).replace("Style ", "") row.operator_menu_enum("mn.style_change", "style", text=current_style) box = layout.box() ui_from_node(box, nodes.get_style_node(object)) diff --git a/molecularnodes/ui/pref.py b/molecularnodes/ui/pref.py index 843f88bf..4bd4f022 100644 --- a/molecularnodes/ui/pref.py +++ b/molecularnodes/ui/pref.py @@ -30,9 +30,7 @@ def draw(self, context): col_main = layout.column(heading="", align=False) row_import = col_main.row() - row_import.prop( - bpy.context.scene, "pypi_mirror_provider", text="Set PyPI Mirror" - ) + row_import.prop(bpy.context.scene, "pypi_mirror_provider", text="Set PyPI Mirror") pkgs = pkg.get_pkgs() for package in pkgs.values(): diff --git a/molecularnodes/utils.py b/molecularnodes/utils.py index 76af938c..3ffe203c 100644 --- a/molecularnodes/utils.py +++ b/molecularnodes/utils.py @@ -88,9 +88,7 @@ def _zipfile_root_namelist(file_to_extract): def template_install(): print(os.path.abspath(ADDON_DIR)) - template = os.path.join( - os.path.abspath(ADDON_DIR), "assets", "template", "Molecular Nodes.zip" - ) + template = os.path.join(os.path.abspath(ADDON_DIR), "assets", "template", "Molecular Nodes.zip") _install_template(template) bpy.utils.refresh_script_paths() @@ -118,9 +116,7 @@ def _install_template(filepath, subfolder="", overwrite=True): try: os.makedirs(path_app_templates, exist_ok=True) except PermissionError: - print( - "Permission denied: You do not have the necessary permissions to create the directory." - ) + print("Permission denied: You do not have the necessary permissions to create the directory.") traceback.print_exc() except OSError as e: print(f"OS error: {e}") @@ -136,15 +132,11 @@ def _install_template(filepath, subfolder="", overwrite=True): traceback.print_exc() return {"CANCELLED"} except PermissionError: - print( - "Permission denied: You do not have the necessary permissions to open the file." - ) + print("Permission denied: You do not have the necessary permissions to open the file.") traceback.print_exc() return {"CANCELLED"} except zipfile.BadZipFile: - print( - "Bad zip file: The file is not a zip file or it is corrupted." - ) + print("Bad zip file: The file is not a zip file or it is corrupted.") traceback.print_exc() return {"CANCELLED"} @@ -154,9 +146,7 @@ def _install_template(filepath, subfolder="", overwrite=True): _module_filesystem_remove(path_app_templates, f) else: for f in file_to_extract_root: - path_dest = os.path.join( - path_app_templates, os.path.basename(f) - ) + path_dest = os.path.join(path_app_templates, os.path.basename(f)) if os.path.exists(path_dest): # self.report({'WARNING'}, tip_("File already installed to %r\n") % path_dest) return {"CANCELLED"} @@ -164,9 +154,7 @@ def _install_template(filepath, subfolder="", overwrite=True): try: # extract the file to "bl_app_templates_user" file_to_extract.extractall(path_app_templates) except PermissionError: - print( - "Permission denied: You do not have the necessary permissions to write to the directory." - ) + print("Permission denied: You do not have the necessary permissions to write to the directory.") traceback.print_exc() return {"CANCELLED"} except OSError as e: diff --git a/pyproject.toml b/pyproject.toml index 73235a08..3834a45e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,15 +41,15 @@ ignore_missing_imports = true [tool.ruff] # Enable pycodestyle (`E`) and Pyflakes (`F`) codes by default. -select = ["E", "F"] -ignore = [] +lint.select = ["E", "F"] +lint.ignore = ["E501"] # Allow autofix for all enabled rules (when `--fix`) is provided. -fixable = ["A", "B", "C", "D", "E", "F", "G", "I", "N", "Q", "S", "T", "W", "ANN", "ARG", "BLE", "COM", "DJ", "DTZ", "EM", "ERA", "EXE", "FBT", "ICN", "INP", "ISC", "NPY", "PD", "PGH", "PIE", "PL", "PT", "PTH", "PYI", "RET", "RSE", "RUF", "SIM", "SLF", "TCH", "TID", "TRY", "UP", "YTT"] -unfixable = [] +lint.fixable = ["A", "B", "C", "D", "E", "F", "G", "I", "N", "Q", "S", "T", "W", "ANN", "ARG", "BLE", "COM", "DJ", "DTZ", "EM", "ERA", "EXE", "FBT", "ICN", "INP", "ISC", "NPY", "PD", "PGH", "PIE", "PL", "PT", "PTH", "PYI", "RET", "RSE", "RUF", "SIM", "SLF", "TCH", "TID", "TRY", "UP", "YTT"] +lint.unfixable = [] # Exclude a variety of commonly ignored directories. -exclude = [ +lint.exclude = [ ".bzr", ".direnv", ".eggs", @@ -71,5 +71,5 @@ exclude = [ "node_modules", "venv"] -line-length = 79 +line-length = 120 diff --git a/tests/install.py b/tests/install.py index d8b84475..a75df146 100644 --- a/tests/install.py +++ b/tests/install.py @@ -3,10 +3,7 @@ import os import pathlib -REQUIREMENTS = ( - pathlib.Path(pathlib.Path(__file__).resolve().parent.parent) - / "molecularnodes/requirements.txt" -) +REQUIREMENTS = pathlib.Path(pathlib.Path(__file__).resolve().parent.parent) / "molecularnodes/requirements.txt" def main(): diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 5749a482..08033d5e 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -11,9 +11,7 @@ DATA_DIR = join(dirname(realpath(__file__)), "data") -@pytest.mark.parametrize( - "pdb_id, format", itertools.product(["1f2n", "5zng"], ["pdb", "cif"]) -) +@pytest.mark.parametrize("pdb_id, format", itertools.product(["1f2n", "5zng"], ["pdb", "cif"])) def test_get_transformations(pdb_id, format): """ Compare an assembly built from transformation information in diff --git a/tests/test_cellpack.py b/tests/test_cellpack.py index aa76f886..0d8d8a9a 100644 --- a/tests/test_cellpack.py +++ b/tests/test_cellpack.py @@ -12,18 +12,14 @@ def test_load_cellpack(snapshot_custom: NumpySnapshotExtension, format): bpy.ops.wm.read_homefile(app_template="") name = f"Cellpack_{format}" - ens = mn.io.cellpack.load( - data_dir / f"square1.{format}", name=name, node_setup=False, fraction=0.1 - ) + ens = mn.io.cellpack.load(data_dir / f"square1.{format}", name=name, node_setup=False, fraction=0.1) coll = bpy.data.collections[f"cellpack_{name}"] instance_names = [object.name for object in coll.objects] assert snapshot_custom == "\n".join(instance_names) assert ens.name == name - ens.modifiers["MolecularNodes"].node_group.nodes["MN_pack_instances"].inputs[ - "As Points" - ].default_value = False + ens.modifiers["MolecularNodes"].node_group.nodes["MN_pack_instances"].inputs["As Points"].default_value = False mn.blender.nodes.realize_instances(ens) for attribute in ens.data.attributes.keys(): assert snapshot_custom == sample_attribute(ens, attribute, evaluate=True) diff --git a/tests/test_density.py b/tests/test_density.py index 357232af..0862dba6 100644 --- a/tests/test_density.py +++ b/tests/test_density.py @@ -117,12 +117,8 @@ def test_density_naming_api(density_file, name): assert object.name == object_name -@pytest.mark.parametrize( - "invert,node_setup,center", list(itertools.product([True, False], repeat=3)) -) -def test_density_operator( - snapshot_custom: NumpySnapshotExtension, density_file, invert, node_setup, center -): +@pytest.mark.parametrize("invert,node_setup,center", list(itertools.product([True, False], repeat=3))) +def test_density_operator(snapshot_custom: NumpySnapshotExtension, density_file, invert, node_setup, center): scene = bpy.context.scene scene.MN_import_density = str(density_file) scene.MN_import_density_invert = invert @@ -134,6 +130,4 @@ def test_density_operator( for bob in bpy.data.objects: if bob.name not in bobs: new_bob = bob - assert snapshot_custom == sample_attribute( - mn.blender.obj.evaluate_using_mesh(new_bob), "position" - ) + assert snapshot_custom == sample_attribute(mn.blender.obj.evaluate_using_mesh(new_bob), "position") diff --git a/tests/test_download.py b/tests/test_download.py index 916089f0..d99a5172 100644 --- a/tests/test_download.py +++ b/tests/test_download.py @@ -23,10 +23,7 @@ def _filestart(format): def test_download_raises_error_on_invalid_format(): with pytest.raises(ValueError) as excinfo: download("1abc", "invalid_format") - assert ( - "File format 'invalid_format' not in: supported_formats=['cif', 'pdb', 'bcif']" - in str(excinfo.value) - ) + assert "File format 'invalid_format' not in: supported_formats=['cif', 'pdb', 'bcif']" in str(excinfo.value) def test_fail_download_pdb_large_structure_raises(): @@ -41,14 +38,8 @@ def test_fail_download_pdb_large_structure_raises(): @pytest.mark.parametrize("format", ["cif", "bcif", "pdb"]) def test_compare_biotite(format): - struc_download = load_structure( - mn.io.download("4ozs", format=format, cache=tempfile.TemporaryDirectory().name) - ) - struc_biotite = load_structure( - rcsb.fetch( - "4ozs", format=format, target_path=tempfile.TemporaryDirectory().name - ) - ) + struc_download = load_structure(mn.io.download("4ozs", format=format, cache=tempfile.TemporaryDirectory().name)) + struc_biotite = load_structure(rcsb.fetch("4ozs", format=format, target_path=tempfile.TemporaryDirectory().name)) assert struc_download == struc_biotite diff --git a/tests/test_load.py b/tests/test_load.py index 805632b2..e6deb3a2 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -15,9 +15,7 @@ def useful_function(snapshot_custom, style, code, assembly, cache_dir=None): - obj = mn.io.fetch( - code, style=style, build_assembly=assembly, cache_dir=cache_dir - ).object + obj = mn.io.fetch(code, style=style, build_assembly=assembly, cache_dir=cache_dir).object node = mn.blender.nodes.get_style_node(obj) if "EEVEE" in node.inputs.keys(): @@ -29,9 +27,7 @@ def useful_function(snapshot_custom, style, code, assembly, cache_dir=None): assert snapshot_custom == sample_attribute(obj, att, evaluate=dont_realise) -@pytest.mark.parametrize( - "assembly, code, style", itertools.product([False], codes, styles) -) +@pytest.mark.parametrize("assembly, code, style", itertools.product([False], codes, styles)) def test_style_1(snapshot_custom: NumpySnapshotExtension, assembly, code, style): useful_function(snapshot_custom, style, code, assembly, cache_dir=data_dir) @@ -48,9 +44,7 @@ def test_style_2(snapshot_custom: NumpySnapshotExtension, assembly, code, style) useful_function(snapshot_custom, style, code, assembly, cache_dir=data_dir) -@pytest.mark.parametrize( - "code, format", itertools.product(codes, ["bcif", "cif", "pdb"]) -) +@pytest.mark.parametrize("code, format", itertools.product(codes, ["bcif", "cif", "pdb"])) def test_download_format(code, format): mol = mn.io.fetch(code, format=format, style=None).object scene = bpy.context.scene @@ -76,9 +70,7 @@ def test_style_positions(snapshot_custom: NumpySnapshotExtension, code, style): assert snapshot_custom == sample_attribute(mol, "position") -@pytest.mark.parametrize( - "code, centre_method", itertools.product(codes, centre_methods) -) +@pytest.mark.parametrize("code, centre_method", itertools.product(codes, centre_methods)) def test_centring(snapshot_custom: NumpySnapshotExtension, code, centre_method): """fetch a pdb structure using code and translate the model using the centre_method. Check the CoG and CoM values against the snapshot file. @@ -103,27 +95,16 @@ def test_centring_different(code): each by a different centring method. Check that their centroids and positions are in fact different. """ - mols = [ - mn.io.fetch(code, centre=method, cache_dir=data_dir) - for method in centre_methods - ] + mols = [mn.io.fetch(code, centre=method, cache_dir=data_dir) for method in centre_methods] for mol1, mol2 in itertools.combinations(mols, 2): - assert not np.allclose( - mol1.centre(centre_type="centroid"), mol2.centre(centre_type="centroid") - ) - assert not np.allclose( - mol1.centre(centre_type="mass"), mol2.centre(centre_type="mass") - ) - assert not np.allclose( - mol1.get_attribute("position"), mol2.get_attribute("position") - ) + assert not np.allclose(mol1.centre(centre_type="centroid"), mol2.centre(centre_type="centroid")) + assert not np.allclose(mol1.centre(centre_type="mass"), mol2.centre(centre_type="mass")) + assert not np.allclose(mol1.get_attribute("position"), mol2.get_attribute("position")) # THESE TEST FUNCTIONS ARE NOT RUN def test_local_pdb(snapshot_custom): - molecules = [ - mn.io.load(data_dir / f"1l58.{ext}", style="spheres") for ext in ("cif", "pdb") - ] + molecules = [mn.io.load(data_dir / f"1l58.{ext}", style="spheres") for ext in ("cif", "pdb")] molecules.append(mn.io.fetch("1l58", format="bcif")) for att in ["position"]: for mol in molecules: @@ -134,11 +115,7 @@ def test_rcsb_nmr(snapshot_custom): mol = mn.io.fetch("2M6Q", style="cartoon") assert len(mol.frames.objects) == 10 assert ( - mol.object.modifiers["MolecularNodes"] - .node_group.nodes["MN_animate_value"] - .inputs["To Max"] - .default_value - == 9 + mol.object.modifiers["MolecularNodes"].node_group.nodes["MN_animate_value"].inputs["To Max"].default_value == 9 ) assert snapshot_custom == sample_attribute(mol, "position", evaluate=True) @@ -172,6 +149,4 @@ def test_rcsb_cache(snapshot_custom): assert os.path.exists(file) obj_2 = mn.io.fetch("6BQN", style="cartoon", cache_dir=test_cache) - assert ( - sample_attribute(obj_1, "position") == sample_attribute(obj_2, "position") - ).all() + assert (sample_attribute(obj_1, "position") == sample_attribute(obj_2, "position")).all() diff --git a/tests/test_mda.py b/tests/test_mda.py index ce304f19..bf8d57b2 100644 --- a/tests/test_mda.py +++ b/tests/test_mda.py @@ -50,9 +50,7 @@ def reload_mda_session(self, mda_session): mn.mda.create_session() @pytest.mark.parametrize("in_memory", [False, True]) - def test_show_universe( - self, snapshot_custom: NumpySnapshotExtension, in_memory, mda_session, universe - ): + def test_show_universe(self, snapshot_custom: NumpySnapshotExtension, in_memory, mda_session, universe): remove_all_molecule_objects(mda_session) mda_session.show(universe, in_memory=in_memory) bob = bpy.data.objects["atoms"] @@ -74,9 +72,7 @@ def test_same_name_atoms(self, in_memory, mda_session, universe): assert (verts_1 == verts_2).all() @pytest.mark.parametrize("in_memory", [False, True]) - def test_show_multiple_selection( - self, snapshot_custom: NumpySnapshotExtension, in_memory, mda_session, universe - ): + def test_show_multiple_selection(self, snapshot_custom: NumpySnapshotExtension, in_memory, mda_session, universe): remove_all_molecule_objects(mda_session) custom_selections = {"name_ca": "name CA"} mda_session.show( @@ -132,9 +128,7 @@ def test_attributes_added(self, in_memory, mda_session, universe): assert att in attributes @pytest.mark.parametrize("in_memory", [False, True]) - def test_trajectory_update( - self, snapshot_custom: NumpySnapshotExtension, in_memory, universe - ): + def test_trajectory_update(self, snapshot_custom: NumpySnapshotExtension, in_memory, universe): # remove_all_molecule_objects(mda_session) mda_session = mn.io.MDAnalysisSession() @@ -169,9 +163,7 @@ def test_trajectory_update( assert not np.isclose(pos_a, pos_b).all() @pytest.mark.parametrize("in_memory", [False, True]) - def test_show_updated_atoms( - self, snapshot_custom: NumpySnapshotExtension, in_memory, mda_session, universe - ): + def test_show_updated_atoms(self, snapshot_custom: NumpySnapshotExtension, in_memory, mda_session, universe): remove_all_molecule_objects(mda_session) updating_ag = universe.select_atoms("around 5 resid 1", updating=True) mda_session.show(updating_ag, in_memory=in_memory, style="vdw") @@ -322,9 +314,7 @@ def test_subframes(self, mda_session, universe): # now using subframes, there should be a difference assert not np.isclose(verts_b, verts_c).all() - assert np.isclose( - verts_c, mn.utils.lerp(verts_a, verts_b, t=fraction) - ).all() + assert np.isclose(verts_c, mn.utils.lerp(verts_a, verts_b, t=fraction)).all() def test_subframe_mapping(self, mda_session, universe): remove_all_molecule_objects(mda_session) @@ -354,9 +344,7 @@ def test_subframe_mapping(self, mda_session, universe): def test_martini(snapshot_custom: NumpySnapshotExtension, toplogy): session = mn.io.MDAnalysisSession() remove_all_molecule_objects(session) - universe = mda.Universe( - data_dir / "martini" / toplogy, data_dir / "martini/pent/PENT2_100frames.xtc" - ) + universe = mda.Universe(data_dir / "martini" / toplogy, data_dir / "martini/pent/PENT2_100frames.xtc") mol = session.show(universe, style="ribbon") diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 6353267f..0f1748bf 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -15,53 +15,28 @@ def test_node_name_format(): - assert ( - mn.blender.nodes.format_node_name("MN_style_cartoon") == "Style Cartoon" - ) - assert ( - mn.blender.nodes.format_node_name("MN_dna_double_helix") - == "DNA Double Helix" - ) - assert ( - mn.blender.nodes.format_node_name("MN_topo_vector_angle") - == "Topology Vector Angle" - ) + assert mn.blender.nodes.format_node_name("MN_style_cartoon") == "Style Cartoon" + assert mn.blender.nodes.format_node_name("MN_dna_double_helix") == "DNA Double Helix" + assert mn.blender.nodes.format_node_name("MN_topo_vector_angle") == "Topology Vector Angle" def test_get_nodes(): bob = mn.io.fetch("4ozs", style="spheres").object - assert ( - nodes.get_nodes_last_output(bob.modifiers["MolecularNodes"].node_group)[ - 0 - ].name - == "MN_style_spheres" - ) + assert nodes.get_nodes_last_output(bob.modifiers["MolecularNodes"].node_group)[0].name == "MN_style_spheres" nodes.realize_instances(bob) - assert ( - nodes.get_nodes_last_output(bob.modifiers["MolecularNodes"].node_group)[ - 0 - ].name - == "Realize Instances" - ) + assert nodes.get_nodes_last_output(bob.modifiers["MolecularNodes"].node_group)[0].name == "Realize Instances" assert nodes.get_style_node(bob).name == "MN_style_spheres" bob2 = mn.io.fetch("1cd3", style="cartoon", build_assembly=True).object - assert ( - nodes.get_nodes_last_output( - bob2.modifiers["MolecularNodes"].node_group - )[0].name - == "MN_assembly_1cd3" - ) + assert nodes.get_nodes_last_output(bob2.modifiers["MolecularNodes"].node_group)[0].name == "MN_assembly_1cd3" assert nodes.get_style_node(bob2).name == "MN_style_cartoon" def test_selection(): chain_ids = [let for let in "ABCDEFG123456"] - node = nodes.custom_iswitch( - "test_node", chain_ids, prefix="Chain ", dtype="BOOLEAN" - ) + node = nodes.custom_iswitch("test_node", chain_ids, prefix="Chain ", dtype="BOOLEAN") input_sockets = nodes.inputs(node) for letter, socket in zip(chain_ids, input_sockets.values()): @@ -71,14 +46,10 @@ def test_selection(): @pytest.mark.parametrize("code", codes) @pytest.mark.parametrize("attribute", ["chain_id", "entity_id"]) -def test_selection_working( - snapshot_custom: NumpySnapshotExtension, attribute, code -): +def test_selection_working(snapshot_custom: NumpySnapshotExtension, attribute, code): mol = mn.io.fetch(code, style="ribbon", cache_dir=data_dir).object group = mol.modifiers["MolecularNodes"].node_group - node_sel = nodes.add_selection( - group, mol.name, mol[f"{attribute}s"], attribute - ) + node_sel = nodes.add_selection(group, mol.name, mol[f"{attribute}s"], attribute) n = len(node_sel.inputs) @@ -103,17 +74,13 @@ def test_color_custom(snapshot_custom: NumpySnapshotExtension, code, attribute): ) group = mol.modifiers["MolecularNodes"].node_group node_col = mn.blender.nodes.add_custom(group, group_col.name, [0, -200]) - group.links.new( - node_col.outputs[0], group.nodes["MN_color_set"].inputs["Color"] - ) + group.links.new(node_col.outputs[0], group.nodes["MN_color_set"].inputs["Color"]) assert snapshot_custom == sample_attribute(mol, "Color", n=50) def test_custom_resid_selection(): - node = mn.blender.nodes.resid_multiple_selection( - "new_node", "1, 5, 10-20, 40-100" - ) + node = mn.blender.nodes.resid_multiple_selection("new_node", "1, 5, 10-20, 40-100") numbers = [1, 5, 10, 20, 40, 100] assert len(nodes.outputs(node)) == 2 counter = 0 @@ -152,9 +119,7 @@ def test_color_lookup_supplied(): for item in nodes.inputs(node).values(): assert np.allclose(np.array(item.default_value), col) - node = mn.blender.nodes.custom_iswitch( - name="test2", iter_list=range(10, 20), dtype="RGBA", start=10 - ) + node = mn.blender.nodes.custom_iswitch(name="test2", iter_list=range(10, 20), dtype="RGBA", start=10) for item in nodes.inputs(node).values(): assert not np.allclose(np.array(item.default_value), col) @@ -168,9 +133,7 @@ def test_color_chain(snapshot_custom: NumpySnapshotExtension): ) group = mol.modifiers["MolecularNodes"].node_group node_col = mn.blender.nodes.add_custom(group, group_col.name, [0, -200]) - group.links.new( - node_col.outputs[0], group.nodes["MN_color_set"].inputs["Color"] - ) + group.links.new(node_col.outputs[0], group.nodes["MN_color_set"].inputs["Color"]) assert snapshot_custom == sample_attribute(mol, "Color") @@ -185,9 +148,7 @@ def test_color_entity(snapshot_custom: NumpySnapshotExtension): ) group = mol.modifiers["MolecularNodes"].node_group node_col = mn.blender.nodes.add_custom(group, group_col.name, [0, -200]) - group.links.new( - node_col.outputs[0], group.nodes["MN_color_set"].inputs["Color"] - ) + group.links.new(node_col.outputs[0], group.nodes["MN_color_set"].inputs["Color"]) assert snapshot_custom == sample_attribute(mol, "Color") @@ -208,21 +169,13 @@ def test_change_style(): for style in ["ribbon", "cartoon", "presets", "ball_and_stick", "surface"]: style_node_1 = nodes.get_style_node(model) - links_in_1 = [ - link.from_socket.name for link in get_links(style_node_1.inputs) - ] - links_out_1 = [ - link.from_socket.name for link in get_links(style_node_1.outputs) - ] + links_in_1 = [link.from_socket.name for link in get_links(style_node_1.inputs)] + links_out_1 = [link.from_socket.name for link in get_links(style_node_1.outputs)] nodes.change_style_node(model, style) style_node_2 = nodes.get_style_node(model) - links_in_2 = [ - link.from_socket.name for link in get_links(style_node_2.inputs) - ] - links_out_2 = [ - link.from_socket.name for link in get_links(style_node_2.outputs) - ] + links_in_2 = [link.from_socket.name for link in get_links(style_node_2.inputs)] + links_out_2 = [link.from_socket.name for link in get_links(style_node_2.outputs)] assert len(links_in_1) == len(links_in_2) assert len(links_out_1) == len(links_out_2) @@ -240,19 +193,13 @@ def test_node_topology(snapshot_custom: NumpySnapshotExtension): node_att = group.nodes.new("GeometryNodeStoreNamedAttribute") node_att.inputs[2].default_value = "test_attribute" nodes.insert_last_node(group, node_att) - node_names = [ - node["name"] - for node in mn.ui.node_info.menu_items["topology"] - if not node == "break" - ] + node_names = [node["name"] for node in mn.ui.node_info.menu_items["topology"] if not node == "break"] for node_name in node_names: # exclude these particular nodes, as they aren't field nodes and so we shouldn't # be testing them here. Will create their own particular tests later if "backbone" in node_name or "bonds" in node_name: continue - node_topo = nodes.add_custom( - group, node_name, location=[x - 300 for x in node_att.location] - ) + node_topo = nodes.add_custom(group, node_name, location=[x - 300 for x in node_att.location]) if node_name == "MN_topo_point_mask": node_topo.inputs["atom_name"].default_value = 61 @@ -275,9 +222,7 @@ def test_node_topology(snapshot_custom: NumpySnapshotExtension): group.links.new(output, input) - assert snapshot_custom == mn.blender.obj.get_attribute( - mol, "test_attribute", evaluate=True - ) + assert snapshot_custom == mn.blender.obj.get_attribute(mol, "test_attribute", evaluate=True) def test_compute_backbone(snapshot_custom: NumpySnapshotExtension): @@ -296,9 +241,7 @@ def test_compute_backbone(snapshot_custom: NumpySnapshotExtension): nodes.insert_last_node(group, node_att) node_names = ["MN_topo_backbone"] for node_name in node_names: - node_topo = nodes.add_custom( - group, node_name, location=[x - 300 for x in node_att.location] - ) + node_topo = nodes.add_custom(group, node_name, location=[x - 300 for x in node_att.location]) if node_name == "MN_topo_point_mask": node_topo.inputs["atom_name"].default_value = 61 @@ -321,9 +264,7 @@ def test_compute_backbone(snapshot_custom: NumpySnapshotExtension): group.links.new(output, input) - assert snapshot_custom == mn.blender.obj.get_attribute( - mol, "test_attribute", evaluate=True - ) + assert snapshot_custom == mn.blender.obj.get_attribute(mol, "test_attribute", evaluate=True) for angle in ["Phi", "Psi"]: output = node_backbone.outputs[angle] @@ -335,9 +276,7 @@ def test_compute_backbone(snapshot_custom: NumpySnapshotExtension): group.links.new(output, input) - assert snapshot_custom == mn.blender.obj.get_attribute( - mol, "test_attribute", evaluate=True - ) + assert snapshot_custom == mn.blender.obj.get_attribute(mol, "test_attribute", evaluate=True) def test_topo_bonds(): diff --git a/tests/test_ops.py b/tests/test_ops.py index 66a33c87..0c9d477d 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -14,9 +14,7 @@ @pytest.mark.parametrize("code", codes) -def test_op_api_cartoon( - snapshot_custom: NumpySnapshotExtension, code, style="ribbon", format="bcif" -): +def test_op_api_cartoon(snapshot_custom: NumpySnapshotExtension, code, style="ribbon", format="bcif"): scene = bpy.context.scene scene.MN_import_node_setup = True scene.MN_pdb_code = code @@ -63,9 +61,7 @@ def test_op_local(snapshot_custom, code, file_format): bpy.ops.mn.import_protein_local() bob_centred = o.latest() - bob_pos, bob_centred_pos = [ - sample_attribute(x, "position", evaluate=False) for x in [bob, bob_centred] - ] + bob_pos, bob_centred_pos = [sample_attribute(x, "position", evaluate=False) for x in [bob, bob_centred]] assert snapshot_custom == bob_pos assert snapshot_custom == bob_centred_pos diff --git a/tests/test_pdbx.py b/tests/test_pdbx.py index 3223ccd9..0ae5c45a 100644 --- a/tests/test_pdbx.py +++ b/tests/test_pdbx.py @@ -23,6 +23,4 @@ def test_get_ss_from_mmcif(snapshot_custom): def test_secondary_structure_no_helix(snapshot_custom): m = mn.io.fetch("7ZL4", cache_dir=data_dir) - assert snapshot_custom == sample_attribute( - m.object, "sec_struct", n=500, evaluate=False - ) + assert snapshot_custom == sample_attribute(m.object, "sec_struct", n=500, evaluate=False) diff --git a/tests/test_select.py b/tests/test_select.py index 4114aff9..71b7fed5 100644 --- a/tests/test_select.py +++ b/tests/test_select.py @@ -36,7 +36,4 @@ def test_select_multiple_residues(selection): vertices_count = len(mn.blender.obj.evaluated(object).data.vertices) assert vertices_count == len(selection[1]) - assert ( - mn.blender.obj.get_attribute(mn.blender.obj.evaluated(object), "res_id") - == selection[1] - ).all() + assert (mn.blender.obj.get_attribute(mn.blender.obj.evaluated(object), "res_id") == selection[1]).all() diff --git a/tests/test_star.py b/tests/test_star.py index 424c13dd..e3a74724 100644 --- a/tests/test_star.py +++ b/tests/test_star.py @@ -25,17 +25,13 @@ def test_starfile_attributes(type): elif type == "cistem": df = star - euler_angles = df[ - ["cisTEMAnglePhi", "cisTEMAngleTheta", "cisTEMAnglePsi"] - ].to_numpy() + euler_angles = df[["cisTEMAnglePhi", "cisTEMAngleTheta", "cisTEMAnglePsi"]].to_numpy() # Calculate Scipy rotation from the euler angles rot_from_euler = R.from_euler(seq="ZYZ", angles=euler_angles, degrees=True).inv() # Activate the rotation debug mode in the nodetreee and get the quaternion attribute - debugnode = mn.blender.nodes.star_node(ensemble.node_group).node_tree.nodes[ - "Switch.001" - ] + debugnode = mn.blender.nodes.star_node(ensemble.node_group).node_tree.nodes["Switch.001"] debugnode.inputs["Switch"].default_value = True quat_attribute = ensemble.get_attribute("MNDEBUGEuler", evaluate=True) @@ -75,30 +71,19 @@ def test_micrograph_loading(): assert tiff_path.exists() # Ensure montage get only loaded once assert sum(1 for image in bpy.data.images.keys() if "montage" in image) == 1 - assert ( - ensemble.micrograph_material.node_tree.nodes["Image Texture"].image.name - == "montage.tiff" - ) + assert ensemble.micrograph_material.node_tree.nodes["Image Texture"].image.name == "montage.tiff" assert ensemble.star_node.inputs["Micrograph"].default_value.name == "montage.tiff" -@pytest.mark.skipif( - importlib.util.find_spec("pyopenvdb"), reason="Test may segfault on GHA" -) +@pytest.mark.skipif(importlib.util.find_spec("pyopenvdb"), reason="Test may segfault on GHA") def test_rehydration(tmp_path): bpy.ops.wm.read_homefile() ensemble = mn.io.star.load(data_dir / "cistem.star") bpy.ops.wm.save_as_mainfile(filepath=str(tmp_path / "test.blend")) assert ensemble._update_micrograph_texture in bpy.app.handlers.depsgraph_update_post bpy.ops.wm.read_homefile() - assert ( - ensemble._update_micrograph_texture - not in bpy.app.handlers.depsgraph_update_post - ) + assert ensemble._update_micrograph_texture not in bpy.app.handlers.depsgraph_update_post bpy.ops.wm.open_mainfile(filepath=str(tmp_path / "test.blend")) new_ensemble = bpy.types.Scene.MN_starfile_ensembles[0] - assert ( - new_ensemble._update_micrograph_texture - in bpy.app.handlers.depsgraph_update_post - ) + assert new_ensemble._update_micrograph_texture in bpy.app.handlers.depsgraph_update_post assert new_ensemble.data.equals(ensemble.data) diff --git a/tests/utils.py b/tests/utils.py index d28eeb3c..0f6318aa 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -13,15 +13,11 @@ class NumpySnapshotExtension(AmberSnapshotExtension): def serialize(self, data, **kwargs): if isinstance(data, np.ndarray): - return np.array2string( - data, precision=1, threshold=1e3, floatmode="maxprec_equal" - ) + return np.array2string(data, precision=1, threshold=1e3, floatmode="maxprec_equal") return super().serialize(data, **kwargs) -def sample_attribute( - object, attribute, n=100, evaluate=True, error: bool = False, seed=6 -): +def sample_attribute(object, attribute, n=100, evaluate=True, error: bool = False, seed=6): if isinstance(object, mn.io.parse.molecule.Molecule): object = object.object @@ -41,9 +37,7 @@ def sample_attribute( return attribute[idx, :] else: try: - attribute = mn.blender.obj.get_attribute( - object=object, name=attribute, evaluate=evaluate - ) + attribute = mn.blender.obj.get_attribute(object=object, name=attribute, evaluate=evaluate) length = len(attribute) if n > length: @@ -59,15 +53,11 @@ def sample_attribute( return np.array(e) -def sample_attribute_to_string( - object, attribute, n=100, evaluate=True, precision=3, seed=6 -): +def sample_attribute_to_string(object, attribute, n=100, evaluate=True, precision=3, seed=6): if isinstance(object, mn.io.parse.molecule.Molecule): object = object.object try: - array = sample_attribute( - object, attribute=attribute, n=n, evaluate=evaluate, seed=seed - ) + array = sample_attribute(object, attribute=attribute, n=n, evaluate=evaluate, seed=seed) except AttributeError as e: print(f"Error {e}, unable to sample attribute {attribute} from {object}") return str(e)