Merge pull request #285 from BradyAJohnston/dev-overhaul-cartoon

feat: improve cartoon styles

MolecularNodes/load.py

@@ -367,7 +367,8 @@ def create_molecule(mol_array,
                     file = None,
                     calculate_ss = False,
                     del_solvent = False, 
-                    include_bonds = False, 
+                    include_bonds = False,
+                    starting_style = 0,
                     collection = None
                     ):
     import biotite.structure as struc
@@ -410,6 +411,7 @@ def create_molecule(mol_array,
         locations = locations, 
         bonds = bond_idx
         )
+
 
     # 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

MolecularNodes/nodes.py

@@ -19,14 +19,14 @@
         'IMAGE'     : 'NodeSocketImage'
     }
 
-def mol_append_node(node_name):
+def mol_append_node(node_name, link = True):
     node = bpy.data.node_groups.get(node_name)
-    if not node:
+    if not node or link:
         bpy.ops.wm.append(
             directory = os.path.join(
                     pkg.ADDON_DIR, 'assets', 'node_append_file.blend' + r'/NodeTree'), 
                     filename = node_name, 
-                    link = False
+                    link = link
                 )
 
     return bpy.data.node_groups[node_name]
@@ -68,7 +68,7 @@ def gn_new_group_empty(name = "Geometry Nodes"):
     group.links.new(output_node.inputs[0], input_node.outputs[0])
     return group
 
-def add_custom_node_group(parent_group, node_name, location = [0,0], width = 200):
+def add_custom_node_group(parent_group, node_name, location = [0,0], width = 200, show_options = True):
 
     mol_append_node(node_name)
 
@@ -77,10 +77,12 @@ def add_custom_node_group(parent_group, node_name, location = [0,0], width = 200
 
     node.location = location
     node.width = 200 # unsure if width will work TODO check
+    node.show_options = show_options
+    node.name = node_name
 
     return node
 
-def add_custom_node_group_to_node(parent_group, node_name, location = [0,0], width = 200):
+def add_custom_node_group_to_node(parent_group, node_name, location = [0,0], width = 200, show_options = False):
 
     mol_append_node(node_name)
 
@@ -89,6 +91,8 @@ def add_custom_node_group_to_node(parent_group, node_name, location = [0,0], wid
 
     node.location = location
     node.width = 200 # unsure if width will work TODO check
+    node.show_options = show_options
+    node.name = node_name
 
     return node
 
@@ -274,23 +278,25 @@ def create_starting_node_tree(obj, coll_frames, starting_style = "atoms"):
     # node_properties = add_custom_node_group(node_group, 'MOL_prop_setup', [0, 0])
     node_colour = add_custom_node_group(node_mod, 'MOL_color_set_common', [200, 0])
 
-    node_random_colour = node_group.nodes.new("FunctionNodeRandomValue")
-    node_random_colour.data_type = 'FLOAT_VECTOR'
-    node_random_colour.location = [-60, -200]
 
-    node_chain_id = node_group.nodes.new("GeometryNodeInputNamedAttribute")
-    node_chain_id.location = [-250, -450]
-    node_chain_id.data_type = "INT"
-    node_chain_id.inputs['Name'].default_value = "chain_id"
+    node_random_color = add_custom_node_group(node_mod, 'MOL_color_random', [-60, -200])
+    # node_random_colour = node_group.nodes.new("FunctionNodeRandomValue")
+    # node_random_colour.data_type = 'FLOAT_VECTOR'
+    # node_random_colour.location = [-60, -200]
+
+    # node_chain_id = node_group.nodes.new("GeometryNodeInputNamedAttribute")
+    # node_chain_id.location = [-250, -450]
+    # node_chain_id.data_type = "INT"
+    # node_chain_id.inputs['Name'].default_value = "chain_id"
 
 
 
     # create the links between the the nodes that have been established
     link = node_group.links.new
-    link(node_input.outputs['Geometry'], node_colour.inputs['Atoms'])
-    link(node_colour.outputs['Atoms'], node_output.inputs['Geometry'])
-    link(node_random_colour.outputs['Value'], node_colour.inputs['Carbon'])
-    link(node_chain_id.outputs[4], node_random_colour.inputs['ID'])
+    link(node_input.outputs['Geometry'], node_colour.inputs[0])
+    link(node_colour.outputs[0], node_output.inputs['Geometry'])
+    link(node_random_color.outputs['Color'], node_colour.inputs['Carbon'])
+    # link(node_chain_id.outputs[4], node_random_colour.inputs['ID'])
 
     styles = [
         'MOL_style_atoms_cycles', 
@@ -626,7 +632,7 @@ def chain_selection(node_name, input_list, attribute = 'chain_id', starting_valu
     counter = 0
     for chain_name in input_list:
         current_node = chain_group.nodes.new("GeometryNodeGroup")
-        current_node.node_tree = mol_append_node('MOL_utils_bool_chain')
+        current_node.node_tree = mol_append_node('.utils_bool_chain')
         current_node.location = [counter * node_sep_dis, 200]
         current_node.inputs["number_matched"].default_value = counter + starting_value
         group_link = chain_group.links.new

MolecularNodes/ui.py

@@ -296,7 +296,7 @@ def draw(self, context):
 
         MOL_PT_panel_ui(self.layout, bpy.context.scene)
 
-def mol_add_node(node_name):
+def mol_add_node(node_name, label: str = '', show_options = True):
     prev_context = bpy.context.area.type
     bpy.context.area.type = 'NODE_EDITOR'
     # actually invoke the operator to add a node to the current node tree
@@ -308,8 +308,12 @@ def mol_add_node(node_name):
         use_transform=True
         )
     bpy.context.area.type = prev_context
-    bpy.context.active_node.node_tree = bpy.data.node_groups[node_name]
-    bpy.context.active_node.width = 200.0
+    node = bpy.context.active_node
+    node.node_tree = bpy.data.node_groups[node_name]
+    node.width = 200.0
+    node.show_options = show_options
+    if label != '':
+        node.label = label
 
     # if added node has a 'Material' input, set it to the default MN material
     input_mat = bpy.context.active_node.inputs.get('Material')
@@ -328,6 +332,7 @@ class MOL_OT_Add_Custom_Node_Group(bpy.types.Operator):
         subtype = 'NONE', 
         maxlen = 0
     )
+    node_label: bpy.props.StringProperty(name = 'node_label', default = '')
     node_description: bpy.props.StringProperty(
         name = "node_description", 
         description="", 
@@ -346,7 +351,7 @@ def description(cls, context, properties):
     def execute(self, context):
         try:
             nodes.mol_append_node(self.node_name)
-            mol_add_node(self.node_name)
+            mol_add_node(self.node_name, label=self.node_label)
         except RuntimeError:
             self.report({'ERROR'}, 
                         message='Failed to add node. Ensure you are not in edit mode.')
@@ -361,6 +366,7 @@ def menu_item_interface(layout_function,
                         node_description='Add custom MolecularNodes node group.'):
     op=layout_function.operator('mol.add_custom_node_group', 
                                 text = label, emboss = True, depress=False)
+    op.node_label = label
     op.node_name = node_name
     op.node_description = node_description
 
@@ -623,12 +629,18 @@ def draw(self, context):
         menu_item_interface(layout, 'Set Color Common', 'MOL_color_set_common', 
                             "Choose a color for the most common elements in PDB \
                             structures")
+        menu_item_interface(layout, 'Common Elements', 'MOL_color_element_common', 
+                            "Choose a color for the most common elements in PDB \
+                            structures")
         layout.separator()
         menu_item_interface(layout, 'Goodsell Colors', 'MOL_color_goodsell', 
                             "Adjusts the given colors to copy the 'Goodsell Style'.\n \
                             Darkens the non-carbon atoms and keeps the carbon atoms \
                             the same color. Highlights differences without being too \
                             visually busy")
+        # menu_item_interface(layout, 'Color by B Factor', 'MOL_color_map_attribute')
+        menu_item_interface(layout, 'Color by Attribute', 'MOL_color_map_attribute')
+        menu_item_interface(layout, 'Random Color', 'MOL_color_random')
         layout.separator()
         menu_item_interface(layout, 'Color by SS', 'MOL_color_sec_struct', 
                             "Specify colors based on the secondary structure")
@@ -898,7 +910,7 @@ def poll(cls, context):
     def draw(self, context):
         layout = self.layout
         layout.operator_context = "INVOKE_DEFAULT"
-        menu_item_interface(layout, 'Booelean Chain', 'MOL_utils_bool_chain')
+        menu_item_interface(layout, 'Booelean Chain', '.utils_bool_chain')
         menu_item_interface(layout, 'Rotation Matrix', 'MOL_utils_rotation_matrix')
         menu_item_interface(layout, 'Curve Resample', 'MOL_utils_curve_resample')
         menu_item_interface(layout, 'Determine Secondary Structure', 'MOL_utils_dssp')