diff --git a/gui/user_interface.py b/gui/user_interface.py index 1b534d7..36b7832 100644 --- a/gui/user_interface.py +++ b/gui/user_interface.py @@ -37,10 +37,10 @@ def __init__(self, protein, window, pdb_renderer, embedding_renderer): self.res_layout = None self.color_mode = DropDown(bounding_box=[16, -80, 280, 32], title="Coloring Mode", - options=["Functional Similarity", "Amino Acid Order", "Atom Type"], + options=["Functional Similarity", "Amino Acid Order", "Atom Type", "Residue Type"], window=window, batch=self.batch, bg_batch=self.bg_batch) self.color_palette = DropDown(bounding_box=[312, -80, 280, 32], title="Color Palette", - options=["Rainbow", "Monocolor", "Nature", "Grape", "Penguin", "Lemon"], + options=["Rainbow", "Monocolor", "Nature", "Penguin", "Grape", "Lemon", "Mulberry"], window=window, batch=self.batch, bg_batch=self.bg_batch) go_titles = [] for i in range(len(self.protein.go_ids)): diff --git a/protein.py b/protein.py index bc6b3f0..ee64db0 100644 --- a/protein.py +++ b/protein.py @@ -48,14 +48,16 @@ class Protein: CLUSTER_INDEX = 1 # Color residues by their associated cluster in the embedding space RESIDUE_INDEX = 2 # Color residues by their index in the amino acid sequence ATOM_TYPE = 3 # Color residues by their atoms using CPK coloring + RESIDUE_TYPE = 4 # Colors by identity of each amino acid # Color palettes (for color modes 1 and 2) RAINBOW = 6 MONOCOLOR = 7 POISSON = 8 - GRAPE = 9 - PENGUIN = 10 + PENGUIN = 9 + GRAPE = 10 LEMON = 11 + MULBERRY = 12 POISSON_PALETTE = [(187, 176, 148), (128, 118, 101), (89, 82, 70), (51, 51, 51), (25, 31, 34), (47, 68, 67), (59, 94, 88), (90, 140, 108), (139, 180, 141), (192, 208, 165), (247, 239, 199), @@ -65,6 +67,8 @@ class Protein: GRAPE_PALETTE = [(3, 6, 55), (60, 7, 83), (114, 4, 85), (145, 10, 103), (194, 48, 131)] + MULBERRY_PALETTE = [(197, 235, 195), (183, 200, 181), (167, 144, 165), (135, 92, 116), (84, 65, 78), (177, 133, 167), (195, 162, 158), (232, 218, 197), (255, 244, 233)] + MONOCOLOR_PALETTE = [(59, 212, 59)] PENGUIN_PALETTE = [(43, 48, 58), (146, 220, 229), (238, 229, 233), (124, 124, 124), (214, 73, 51)] @@ -136,6 +140,7 @@ def __init__(self, protein_path, prompt_for_chain=True, chain_id=None, verbose=F # Parse the .pdb file for structure data self.residues = [] self.atoms = [] + self.residue_map = {} atom_index = 0 res_index = 0 @@ -153,7 +158,8 @@ def __init__(self, protein_path, prompt_for_chain=True, chain_id=None, verbose=F atom_index += 1 self.residues.append(residue) self.atoms.extend(residue_atoms) - + if not bio_residue.get_resname() in self.residue_map: + self.residue_map[bio_residue.get_resname()] = len(self.residue_map) res_index += 1 if not isfile(f"data/{protein_name}_data.json"): @@ -250,6 +256,8 @@ def get_color_from_palette(palette): new_color.rgb = get_color_from_palette(self.PENGUIN_PALETTE) case self.LEMON: new_color.rgb = get_color_from_palette(self.LEMON_PALETTE) + case self.MULBERRY: + new_color.rgb = get_color_from_palette(self.MULBERRY_PALETTE) new_color.set_luminance(new_color.get_luminance() * luminance) if residue.highlighted: new_color.set_luminance(min(1.0, new_color.get_luminance() + self.HIGHLIGHT_LUMINANCE)) @@ -286,6 +294,11 @@ def get_color_from_palette(palette): if residue.highlighted: color.set_luminance(min(1.0, color.get_luminance() + self.HIGHLIGHT_LUMINANCE)) atom.color = [int(c * 255) for c in color.rgb] + case self.RESIDUE_TYPE: + color = get_color(self.residue_map[residue.bio_residue.get_resname()] / len(self.residue_map)) + for atom in residue.atoms: + atom.color = color + residue.color = color # Adapted from https://github.com/tbepler/prose @staticmethod