Skip to content

Commit

Permalink
Merge pull request BradyAJohnston#338 from BradyAJohnston/color
Browse files Browse the repository at this point in the history 
addition of initial colors on import
  • Loading branch information
@BradyAJohnston
BradyAJohnston authored Oct 31, 2023
2 parents 26dc40a + c17ab42 commit 67cca0f
Show file tree
Hide file tree
Showing 110 changed files with 1,866 additions and 27 deletions.
160 changes: 158 additions & 2 deletions molecularnodes/color.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,164 @@
import colorsys
import numpy as np

def random_rgb():
def random_rgb(seed=None):
"""Random Pastel RGB values
"""
random.seed(seed)
r, g, b = colorsys.hls_to_rgb(random.random(), 0.6, 0.6)
return np.array((r, g, b, 1))
return np.array((r, g, b, 1))


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))


def colors_from_elements(atomic_numbers):
colors = np.array(list(map(color_from_atomic_number, atomic_numbers)))
return colors


def equidistant_colors(some_list):
u = np.unique(some_list)
num_colors = len(u)

hues = [i / (num_colors + 1) for i in range(num_colors)]

# Convert HSL to RGB
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]

return dict(zip(u, colors))


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
)))

colors[mask] = chain_colors[mask]

return colors / 255


iupac_colors_rgb = {
"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
"Ne": (179, 227, 245), # Neon
"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
"Ar": (128, 209, 227), # Argon
"Ca": (61, 255, 0), # Calcium
"Sc": (230, 230, 230), # Scandium
"Ti": (191, 194, 199), # Titanium
"V": (166, 166, 171), # Vanadium
"Cr": (138, 153, 199), # Chromium
"Mn": (156, 122, 199), # Manganese
"Fe": (224, 102, 51), # Iron
"Ni": (199, 138, 138), # Nickel
"Co": (255, 217, 143), # Cobalt
"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
"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
"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
"Pt": (208, 208, 224), # Platinum
"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
}
9 changes: 7 additions & 2 deletions molecularnodes/load.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from . import coll
import warnings
from . import data
from . import color
from . import assembly
from . import nodes
from . import pkg
Expand Down Expand Up @@ -530,6 +531,9 @@ def att_charge():
)))
return charge

def att_color():
return color.color_chains(att_atomic_number(), att_chain_id()).reshape(-1)

def att_is_alpha():
return np.isin(MN_array.atom_name, 'CA')

Expand Down Expand Up @@ -610,6 +614,7 @@ def att_sec_struct():
{'name': 'atom_name', 'value': att_atom_name, 'type': 'INT', 'domain': 'POINT'},
{'name': 'lipophobicity', 'value': att_lipophobicity, 'type': 'FLOAT', 'domain': 'POINT'},
{'name': 'charge', 'value': att_charge, 'type': 'FLOAT', 'domain': 'POINT'},
{'name': 'Color', 'value': att_color, 'type': 'FLOAT_COLOR', 'domain': 'POINT'},

{'name': 'is_backbone', 'value': att_is_backbone, 'type': 'BOOLEAN', 'domain': 'POINT'},
{'name': 'is_alpha_carbon', 'value': att_is_alpha, 'type': 'BOOLEAN', 'domain': 'POINT'},
Expand All @@ -627,8 +632,8 @@ def att_sec_struct():
try:
obj.add_attribute(MN_object, att['name'], att['value'](), att['type'], att['domain'])
print(f'Added {att["name"]} after {time.process_time() - start} s')
except:
# warnings.warn(f"Unable to add attribute: {att['name']}")
except :
warnings.warn(f"Unable to add attribute: {att['name']}")
print(f'Failed adding {att["name"]} after {time.process_time() - start} s')

if MN_frames:
Expand Down
12 changes: 7 additions & 5 deletions molecularnodes/obj.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,14 +81,14 @@ def add_attribute(object: bpy.types.Object, name: str, data, type="FLOAT", domai
to represent vectors with 3 components (x, y, z).
"""

if type == "FLOAT_VECTOR":
att = object.data.attributes.new(name, type, domain)
att = object.data.attributes.new(name, type, domain)
if type == "FLOAT_VECTOR" :
# currently vectors have to be added as a 1d array. may change in the future
# but currently must be reshaped then added as a 'vector' but supplying a 1d array
vec_1d = data.reshape(len(data) * 3).copy(order = 'c')
att.data.foreach_set('vector', vec_1d)
att.data.foreach_set('vector', data.reshape(-1))
elif type == "FLOAT_COLOR":
att.data.foreach_set('color', data.reshape(-1))
else:
att = object.data.attributes.new(name, type, domain)
att.data.foreach_set('value', data.copy(order = 'c'))

return att
Expand Down Expand Up @@ -134,6 +134,8 @@ def get_attribute(obj: bpy.types.Object, att_name='position') -> np.array:
elif att.data_type == "FLOAT_VECTOR":
# Convert attribute vectors to a NumPy array
att_array = np.array(list(map(lambda x: x.vector, att.data.values())))
elif att.data_type == "FLOAT_COLOR":
att_array = np.array(list(map(lambda x: x.color, att.data.values())))
else:
# Unsupported data type, return an empty NumPy array
att_array = np.array([])
Expand Down
2 changes: 2 additions & 0 deletions tests/constants.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
from pathlib import Path
import os

codes = ['4ozs', '8H1B', '1BNA', '8U8W']
attributes = ['b_factor', 'vdw_radii', 'lipophobicity', 'charge', 'res_id', 'res_name', 'atomic_number', 'chain_id', 'entity_id', 'atom_name', 'sec_struct', 'Color', 'position', '.select_vert', 'is_backbone', 'is_alpha_carbon', 'is_solvent', 'is_nucleic', 'is_peptide', 'is_hetero', 'is_carb', 'bond_type', '.edge_verts', '.select_edge', 'sharp_face']
test_data_directory = Path(os.path.abspath(Path(__file__).parent / 'data'))
1 change: 1 addition & 0 deletions tests/snapshots/test_attributes/test_attribute/1BNA/att_.edge_verts_values.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
[]
9 changes: 9 additions & 0 deletions tests/snapshots/test_attributes/test_attribute/1BNA/att_.select_edge_values.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
[ True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True]
9 changes: 9 additions & 0 deletions tests/snapshots/test_attributes/test_attribute/1BNA/att_.select_vert_values.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
[ True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True True True True True True True True True
True True True True]
100 changes: 100 additions & 0 deletions tests/snapshots/test_attributes/test_attribute/1BNA/att_Color_values.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
[[1. 0.051 0.051 0.004]
[0.357 0.839 0.357 0.004]
[0.188 0.314 0.973 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[0.839 0.357 0.357 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[0.188 0.314 0.973 0.004]
[0.188 0.314 0.973 0.004]
[1. 0.502 0. 0.004]
[0.188 0.314 0.973 0.004]
[0.357 0.839 0.357 0.004]
[0.357 0.839 0.357 0.004]
[0.188 0.314 0.973 0.004]
[1. 0.051 0.051 0.004]
[0.357 0.839 0.357 0.004]
[0.188 0.314 0.973 0.004]
[0.839 0.357 0.357 0.004]
[0.357 0.839 0.357 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[0.357 0.839 0.357 0.004]
[1. 0.051 0.051 0.004]
[0.839 0.357 0.357 0.004]
[1. 0.051 0.051 0.004]
[0.188 0.314 0.973 0.004]
[1. 0.051 0.051 0.004]
[0.839 0.357 0.357 0.004]
[0.188 0.314 0.973 0.004]
[0.357 0.839 0.357 0.004]
[1. 0.051 0.051 0.004]
[0.357 0.839 0.357 0.004]
[1. 0.051 0.051 0.004]
[0.357 0.839 0.357 0.004]
[0.188 0.314 0.973 0.004]
[0.188 0.314 0.973 0.004]
[0.357 0.839 0.357 0.004]
[0.357 0.839 0.357 0.004]
[1. 0.051 0.051 0.004]
[0.839 0.357 0.357 0.004]
[0.357 0.839 0.357 0.004]
[0.839 0.357 0.357 0.004]
[0.839 0.357 0.357 0.004]
[0.839 0.357 0.357 0.004]
[0.839 0.357 0.357 0.004]
[0.357 0.839 0.357 0.004]
[0.839 0.357 0.357 0.004]
[1. 0.051 0.051 0.004]
[0.839 0.357 0.357 0.004]
[1. 0.051 0.051 0.004]
[0.357 0.839 0.357 0.004]
[0.188 0.314 0.973 0.004]
[1. 0.051 0.051 0.004]
[1. 0.502 0. 0.004]
[0.357 0.839 0.357 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[0.188 0.314 0.973 0.004]
[0.839 0.357 0.357 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[0.188 0.314 0.973 0.004]
[0.839 0.357 0.357 0.004]
[0.357 0.839 0.357 0.004]
[0.839 0.357 0.357 0.004]
[0.357 0.839 0.357 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[0.357 0.839 0.357 0.004]
[0.188 0.314 0.973 0.004]
[0.357 0.839 0.357 0.004]
[0.839 0.357 0.357 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[0.839 0.357 0.357 0.004]
[1. 0.051 0.051 0.004]
[0.839 0.357 0.357 0.004]
[0.188 0.314 0.973 0.004]
[1. 0.051 0.051 0.004]
[1. 0.502 0. 0.004]
[0.839 0.357 0.357 0.004]
[0.357 0.839 0.357 0.004]
[0.188 0.314 0.973 0.004]
[0.357 0.839 0.357 0.004]
[0.357 0.839 0.357 0.004]
[0.357 0.839 0.357 0.004]
[1. 0.051 0.051 0.004]
[0.188 0.314 0.973 0.004]
[1. 0.051 0.051 0.004]
[1. 0.051 0.051 0.004]
[0.839 0.357 0.357 0.004]
[0.357 0.839 0.357 0.004]
[0.357 0.839 0.357 0.004]
[0.839 0.357 0.357 0.004]
[0.357 0.839 0.357 0.004]
[0.839 0.357 0.357 0.004]
[0.188 0.314 0.973 0.004]]
5 changes: 5 additions & 0 deletions tests/snapshots/test_attributes/test_attribute/1BNA/att_atom_name_values.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
[52 70 74 53 58 56 70 52 53 62 62 50 73 54 70 63 52 61 64 77 61 56 51 59
51 59 72 62 56 61 64 57 52 68 58 67 64 63 57 65 58 57 59 71 61 59 55 55
57 56 57 56 57 62 52 50 61 53 53 58 62 57 51 52 66 54 59 71 54 75 58 67
66 67 67 52 72 54 58 68 64 52 50 54 71 74 54 59 68 52 63 53 76 77 54 57
67 57 67 66]
5 changes: 5 additions & 0 deletions tests/snapshots/test_attributes/test_attribute/1BNA/att_atomic_number_values.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
[ 8 6 7 8 8 8 6 8 8 7 7 15 7 6 6 7 8 6 7 6 6 8 8 6
8 6 8 7 8 6 7 6 8 6 8 6 7 7 6 6 8 6 6 6 6 6 6 6
6 8 6 8 6 7 8 15 6 8 8 8 7 6 8 8 7 6 6 6 6 8 8 6
7 6 6 8 8 6 8 6 7 8 15 6 6 7 6 6 6 8 7 8 8 6 6 6
6 6 6 7]
9 changes: 9 additions & 0 deletions tests/snapshots/test_attributes/test_attribute/1BNA/att_b_factor_values.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
[65.53 32.12 23.72 35.42 42.11 48.38 28.01 78.33 64.35 27.94 35.01 42.39
40.53 50.41 31.71 45.86 51.53 39.56 20.55 36.33 38.4 36.88 42.55 39.28
45.14 51.34 44.39 19.52 44.25 22.18 29.16 43.98 40.43 17.93 79.04 27.74
27.94 30.25 53.8 17.96 55.65 40.14 22. 15.82 41.08 53.79 58.69 40.86
29.82 50.51 23.76 32. 46.45 23.28 53.39 43.68 24.85 39.51 50.67 52.04
25.44 28.58 60.41 54.2 27.6 71.73 29.08 15.08 34.84 30.04 78.59 21.66
39.18 56.67 30.78 46.81 31.85 34.88 53.45 35.5 25.12 45.71 56.5 37.89
23.69 54.91 40.72 40.32 24.58 37.36 17.75 59.7 32.54 24.98 55.04 38.51
38.99 34.34 19.88 28.01]
3 changes: 3 additions & 0 deletions tests/snapshots/test_attributes/test_attribute/1BNA/att_bond_type_values.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
[1 1 1 2 1 1 1 2 1 1 1 1 1 1 1 1 2 1 2 1 1 1 1 1 1 1 1 2 1 1 1 2 2 1 2 1 1
1 2 1 1 1 1 2 2 1 1 1 1 2 1 1 1 1 1 2 1 2 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
1 1 1 1 1 1 1 2 1 1 1 1 1 1 1 2 1 1 1 2 1 1 1 1 1 2]
Loading

0 comments on commit 67cca0f

Please sign in to comment.