Merge pull request #520 from BradyAJohnston/fix-plddt

Add AlphaFold Download
BradyAJohnston · Jun 5, 2024 · 304b70e · 304b70e
2 parents ae4db8c + 8f7c588
commit 304b70e
Show file tree

Hide file tree

Showing 16 changed files with 383 additions and 262 deletions.
diff --git a/molecularnodes/assets/MN_data_file.blend b/molecularnodes/assets/MN_data_file.blend
diff --git a/molecularnodes/blender/nodes.py b/molecularnodes/blender/nodes.py
@@ -158,6 +158,7 @@ def format_node_name(name):
         .title()
         .replace("Dna", "DNA")
         .replace("Topo ", "Topology ")
+        .replace("Plddt", "pLDDT")
     )
 
 
@@ -488,7 +489,7 @@ def create_starting_nodes_density(object, threshold=0.8, style="density_surface"
 
 
 def create_starting_node_tree(
-    object, coll_frames=None, style="spheres", name=None, set_color=True
+    object, coll_frames=None, style="spheres", name=None, color="common"
 ):
     """
     Create a starting node tree for the inputted object.
@@ -506,9 +507,9 @@ def create_starting_node_tree(
     name : str, optional
         Name of the node tree. If None, a default name will be generated.
         The default is None.
-    set_color : bool, optional
-        Whether to set up nodes for generating colors in the node tree.
-        The default is True.
+    color : str, optional
+        None doesn't add ay set_color nodes, 'common' adds the color by common elements
+        and 'plddt' adds color by pLDDT score.
     """
     # ensure there is a geometry nodes modifier called 'MolecularNodes' that is created and applied to the object
     mod = get_mod(object)
@@ -533,16 +534,28 @@ def create_starting_node_tree(
     link(node_input.outputs[0], node_style.inputs[0])
 
     # if requested, setup the nodes for generating colors in the 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])
-
-        link(node_input.outputs["Geometry"], node_color_set.inputs[0])
-        link(node_random_color.outputs["Color"], node_color_common.inputs["Carbon"])
-        link(node_color_common.outputs[0], node_color_set.inputs["Color"])
-        link(node_color_set.outputs[0], node_style.inputs[0])
-        to_animate = node_color_set
+    if color is not None:
+        if color == "common":
+            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]
+            )
+
+            link(node_input.outputs["Geometry"], node_color_set.inputs[0])
+            link(node_random_color.outputs["Color"], node_color_common.inputs["Carbon"])
+            link(node_color_common.outputs[0], node_color_set.inputs["Color"])
+            link(node_color_set.outputs[0], node_style.inputs[0])
+            to_animate = node_color_set
+        elif color.lower() == "plddt":
+            node_color_set = add_custom(group, "MN_color_set", [200, 0])
+            node_color_plddt = add_custom(group, "MN_color_pLDDT", [-50, -150])
+
+            link(node_input.outputs["Geometry"], node_color_set.inputs["Atoms"])
+            link(node_color_plddt.outputs[0], node_color_set.inputs["Color"])
+            link(node_color_set.outputs["Atoms"], node_style.inputs["Atoms"])
+        else:
+            to_animate = node_input
     else:
         to_animate = node_input
 

diff --git a/molecularnodes/color.py b/molecularnodes/color.py
@@ -1,17 +1,34 @@
 import random
 import colorsys
 import numpy as np
+import numpy.typing as npt
 
 
 def random_rgb(seed=None):
-    """Random Pastel RGB values
-    """
+    """Random Pastel RGB values"""
     if seed:
         random.seed(seed)
     r, g, b = colorsys.hls_to_rgb(random.random(), 0.6, 0.6)
     return np.array((r, g, b, 1))
 
 
+def plddt(b_factor: np.ndarray) -> npt.NDArray[np.float32]:
+    colors = np.zeros((len(b_factor), 4), float)
+
+    for i, value in enumerate(b_factor):
+        if value > 90:
+            color = np.array([0.000000, 0.086496, 0.672395, 1.000000])
+        elif value > 70:
+            color = np.array([0.130157, 0.597176, 0.896205, 1.000000])
+        elif value > 50:
+            color = np.array([1.000169, 0.708345, 0.006512, 1.000000])
+        else:
+            color = np.array([1.000169, 0.205070, 0.059507, 1.000000])
+        colors[i, :] = color
+
+    return colors
+
+
 def color_from_atomic_number(atomic_number: int):
     r, g, b = list(iupac_colors_rgb.values())[int(atomic_number - 1)]
     return np.array((r, g, b, 1))
@@ -32,8 +49,7 @@ def equidistant_colors(some_list):
     colors = [colorsys.hls_to_rgb(hue, 0.6, 0.6) for hue in hues]
 
     # Convert RGB to 8-bit integer values
-    colors = [(int(r * 255), int(g * 255), int(b * 255), 1)
-              for (r, g, b) in colors]
+    colors = [(int(r * 255), int(g * 255), int(b * 255), 1) for (r, g, b) in colors]
 
     return dict(zip(u, colors))
 
@@ -48,127 +64,124 @@ def color_chains(atomic_numbers, chain_ids):
     mask = atomic_numbers == 6
     colors = colors_from_elements(atomic_numbers)
     chain_color_dict = equidistant_colors(chain_ids)
-    chain_colors = np.array(list(map(
-        lambda x: chain_color_dict.get(x),
-        chain_ids
-    )))
+    chain_colors = np.array(list(map(lambda x: chain_color_dict.get(x), chain_ids)))
 
     colors[mask] = chain_colors[mask]
 
     return colors / 255
 
 
 iupac_colors_rgb = {
-    "H": (255, 255, 255),   # Hydrogen
+    "H": (255, 255, 255),  # Hydrogen
     "He": (217, 255, 255),  # Helium
     "Li": (204, 128, 255),  # Lithium
-    "Be": (194, 255, 0),    # Beryllium
-    "B": (255, 181, 181),   # Boron
-    "C": (144, 144, 144),   # Carbon
-    "N": (48, 80, 248),     # Nitrogen
-    "O": (255, 13, 13),     # Oxygen
-    "F": (144, 224, 80),    # Fluorine
+    "Be": (194, 255, 0),  # Beryllium
+    "B": (255, 181, 181),  # Boron
+    "C": (144, 144, 144),  # Carbon
+    "N": (48, 80, 248),  # Nitrogen
+    "O": (255, 13, 13),  # Oxygen
+    "F": (144, 224, 80),  # Fluorine
     "Ne": (179, 227, 245),  # Neon
-    "Na": (171, 92, 242),   # Sodium
-    "Mg": (138, 255, 0),    # Magnesium
+    "Na": (171, 92, 242),  # Sodium
+    "Mg": (138, 255, 0),  # Magnesium
     "Al": (191, 166, 166),  # Aluminum
     "Si": (240, 200, 160),  # Silicon
-    "P": (255, 128, 0),     # Phosphorus
-    "S": (255, 255, 48),    # Sulfur
-    "Cl": (31, 240, 31),    # Chlorine
-    "K": (143, 64, 212),    # Potassium
+    "P": (255, 128, 0),  # Phosphorus
+    "S": (255, 255, 48),  # Sulfur
+    "Cl": (31, 240, 31),  # Chlorine
+    "K": (143, 64, 212),  # Potassium
     "Ar": (128, 209, 227),  # Argon
-    "Ca": (61, 255, 0),     # Calcium
+    "Ca": (61, 255, 0),  # Calcium
     "Sc": (230, 230, 230),  # Scandium
     "Ti": (191, 194, 199),  # Titanium
-    "V": (166, 166, 171),   # Vanadium
+    "V": (166, 166, 171),  # Vanadium
     "Cr": (138, 153, 199),  # Chromium
     "Mn": (156, 122, 199),  # Manganese
-    "Fe": (224, 102, 51),   # Iron
+    "Fe": (224, 102, 51),  # Iron
     "Ni": (199, 138, 138),  # Nickel
     "Co": (255, 217, 143),  # Cobalt
-    "Cu": (200, 128, 51),   # Copper
+    "Cu": (200, 128, 51),  # Copper
     "Zn": (125, 128, 176),  # Zinc
     "Ga": (194, 143, 143),  # Gallium
     "Ge": (102, 143, 143),  # Germanium
     "As": (189, 128, 227),  # Arsenic
-    "Se": (255, 161, 0),    # Selenium
-    "Br": (166, 41, 41),    # Bromine
-    "Kr": (92, 184, 209),   # Krypton
-    "Rb": (112, 46, 176),   # Rubidium
-    "Sr": (0, 255, 0),      # Strontium
-    "Y": (148, 255, 255),   # Yttrium
-    "Zr": (148, 224, 224),   # Zirconium
-    "Nb": (115, 194, 201),   # Niobium
-    "Mo": (84, 181, 181),   # Molybdenum
-    "Tc": (59, 158, 158),   # Technetium
-    "Ru": (36, 125, 125),   # Ruthenium
-    "Rh": (10, 125, 140),   # Rhodium
-    "Pd": (0, 105, 133),    # Palladium
+    "Se": (255, 161, 0),  # Selenium
+    "Br": (166, 41, 41),  # Bromine
+    "Kr": (92, 184, 209),  # Krypton
+    "Rb": (112, 46, 176),  # Rubidium
+    "Sr": (0, 255, 0),  # Strontium
+    "Y": (148, 255, 255),  # Yttrium
+    "Zr": (148, 224, 224),  # Zirconium
+    "Nb": (115, 194, 201),  # Niobium
+    "Mo": (84, 181, 181),  # Molybdenum
+    "Tc": (59, 158, 158),  # Technetium
+    "Ru": (36, 125, 125),  # Ruthenium
+    "Rh": (10, 125, 140),  # Rhodium
+    "Pd": (0, 105, 133),  # Palladium
     "Ag": (192, 192, 192),  # Silver
     "Cd": (255, 217, 143),  # Cadmium
     "In": (166, 117, 115),  # Indium
     "Sn": (102, 128, 128),  # Tin
-    "Sb": (158, 99, 181),   # Antimony
-    "Te": (212, 122, 0),    # Tellurium
-    "I": (148, 0, 148),     # Iodine
-    "Xe": (66, 158, 176),   # Xenon
-    "Cs": (87, 23, 143),    # Cesium
-    "Ba": (0, 201, 0),      # Barium
+    "Sb": (158, 99, 181),  # Antimony
+    "Te": (212, 122, 0),  # Tellurium
+    "I": (148, 0, 148),  # Iodine
+    "Xe": (66, 158, 176),  # Xenon
+    "Cs": (87, 23, 143),  # Cesium
+    "Ba": (0, 201, 0),  # Barium
     "La": (112, 212, 255),  # Lanthanum
     "Ce": (255, 255, 199),  # Cerium
     "Pr": (217, 255, 199),  # Praseodymium
     "Nd": (199, 255, 199),  # Neodymium
     "Pm": (163, 255, 199),  # Promethium
     "Sm": (143, 255, 199),  # Samarium
-    "Eu": (97, 255, 199),   # Europium
-    "Gd": (69, 255, 199),   # Gadolinium
-    "Tb": (48, 255, 199),   # Terbium
-    "Dy": (31, 255, 199),   # Dysprosium
-    "Ho": (0, 255, 156),    # Holmium
-    "Er": (0, 230, 117),    # Erbium
-    "Tm": (0, 212, 82),     # Thulium
-    "Yb": (0, 191, 56),     # Ytterbium
-    "Lu": (0, 171, 36),     # Lutetium
-    "Hf": (77, 194, 255),   # Hafnium
-    "Ta": (77, 166, 255),   # Tantalum
-    "W": (33, 148, 214),    # Tungsten
-    "Re": (38, 125, 171),    # Rhenium
-    "Os": (38, 102, 150),    # Osmium
-    "Ir": (23, 84, 135),    # Iridium
+    "Eu": (97, 255, 199),  # Europium
+    "Gd": (69, 255, 199),  # Gadolinium
+    "Tb": (48, 255, 199),  # Terbium
+    "Dy": (31, 255, 199),  # Dysprosium
+    "Ho": (0, 255, 156),  # Holmium
+    "Er": (0, 230, 117),  # Erbium
+    "Tm": (0, 212, 82),  # Thulium
+    "Yb": (0, 191, 56),  # Ytterbium
+    "Lu": (0, 171, 36),  # Lutetium
+    "Hf": (77, 194, 255),  # Hafnium
+    "Ta": (77, 166, 255),  # Tantalum
+    "W": (33, 148, 214),  # Tungsten
+    "Re": (38, 125, 171),  # Rhenium
+    "Os": (38, 102, 150),  # Osmium
+    "Ir": (23, 84, 135),  # Iridium
     "Pt": (208, 208, 224),  # Platinum
-    "Au": (255, 209, 35),   # Gold
+    "Au": (255, 209, 35),  # Gold
     "Hg": (184, 184, 208),  # Mercury
-    "Tl": (166, 84, 77),    # Thallium
-    "Pb": (87, 89, 97),     # Lead
-    "Bi": (158, 79, 181),    # Bismuth
-    "Th": (255, 161, 0),   # Thorium
-    "Pa": (255, 161, 0),   # Protactinium
-    "U": (255, 161, 0),    # Uranium
-    "Np": (255, 161, 0),   # Neptunium
-    "Pu": (255, 161, 0),   # Plutonium
-    "Am": (255, 161, 0),   # Americium
-    "Cm": (255, 161, 0),   # Curium
-    "Bk": (255, 161, 0),   # Berkelium
-    "Cf": (255, 161, 0),   # Californium
-    "Es": (255, 161, 0),   # Einsteinium
-    "Fm": (255, 161, 0),   # Fermium
-    "Md": (255, 161, 0),   # Mendelevium
-    "No": (255, 161, 0),   # Nobelium
-    "Lr": (255, 161, 0),   # Lawrencium
-    "Rf": (204, 0, 89),    # Rutherfordium
-    "Db": (209, 0, 79),    # Dubnium
-    "Sg": (217, 0, 69),    # Seaborgium
-    "Bh": (224, 0, 56),    # Bohrium
-    "Hs": (230, 0, 46),    # Hassium
-    "Mt": (235, 0, 38),    # Meitnerium
-    "Ds": (240, 0, 33),    # Darmstadtium
-    "Rg": (241, 0, 30),    # Roentgenium
-    "Cn": (242, 0, 26),    # Copernicium
-    "Nh": (242, 0, 26),    # Nihonium
-    "Fl": (242, 0, 26),    # Flerovium
-    "Mc": (242, 0, 26),    # Moscovium
-    "Lv": (242, 0, 26),    # Livermorium
-    "Ts": (242, 0, 26),    # Tennessine
-    "Og": (242, 0, 26)     # Oganesson
+    "Tl": (166, 84, 77),  # Thallium
+    "Pb": (87, 89, 97),  # Lead
+    "Bi": (158, 79, 181),  # Bismuth
+    "Th": (255, 161, 0),  # Thorium
+    "Pa": (255, 161, 0),  # Protactinium
+    "U": (255, 161, 0),  # Uranium
+    "Np": (255, 161, 0),  # Neptunium
+    "Pu": (255, 161, 0),  # Plutonium
+    "Am": (255, 161, 0),  # Americium
+    "Cm": (255, 161, 0),  # Curium
+    "Bk": (255, 161, 0),  # Berkelium
+    "Cf": (255, 161, 0),  # Californium
+    "Es": (255, 161, 0),  # Einsteinium
+    "Fm": (255, 161, 0),  # Fermium
+    "Md": (255, 161, 0),  # Mendelevium
+    "No": (255, 161, 0),  # Nobelium
+    "Lr": (255, 161, 0),  # Lawrencium
+    "Rf": (204, 0, 89),  # Rutherfordium
+    "Db": (209, 0, 79),  # Dubnium
+    "Sg": (217, 0, 69),  # Seaborgium
+    "Bh": (224, 0, 56),  # Bohrium
+    "Hs": (230, 0, 46),  # Hassium
+    "Mt": (235, 0, 38),  # Meitnerium
+    "Ds": (240, 0, 33),  # Darmstadtium
+    "Rg": (241, 0, 30),  # Roentgenium
+    "Cn": (242, 0, 26),  # Copernicium
+    "Nh": (242, 0, 26),  # Nihonium
+    "Fl": (242, 0, 26),  # Flerovium
+    "Mc": (242, 0, 26),  # Moscovium
+    "Lv": (242, 0, 26),  # Livermorium
+    "Ts": (242, 0, 26),  # Tennessine
+    "Og": (242, 0, 26),  # Oganesson
 }