pymol view MVP

choppa/IO/input.py

@@ -1,11 +1,9 @@
-
 import pandas as pd
 from typing import Optional
 from pathlib import Path
 import numpy as np 
 import logging, sys
 from collections import OrderedDict
-
 logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 logger = logging.getLogger()
 
@@ -137,6 +135,7 @@ def get_fitness_basedict(self):
         return fitness_basedict
 
 from Bio.PDB import PDBParser
+from rdkit import Chem
 
 class ComplexFactory():
     """
@@ -170,16 +169,12 @@ def load_pdb(self):
         self.check_validity(complex)
         return complex
 
-    def load_pdb_string(self):
+    def load_pdb_rdkit(self):
         """
-        Loads an input PDB file to a string
-
-        TODO: implement validity checks here like with `self.load_pdb()`
+        Loads an input PDB file to an RDKit object for easier string retrieval
         """
-        with open(self.path_to_pdb_file, 'r') as file:
-            pdb_string = file.read()
-
-        return pdb_string
+        return Chem.MolFromPDBFile(self.path_to_pdb_file)
+
 
 if __name__ == "__main__":
     from choppa.data.toy_data.resources import TOY_COMPLEX, TOY_FITNESS_DATA_COMPLETE, TOY_FITNESS_DATA_COMPLETE_NOCONF

choppa/render/render.py

@@ -1,6 +1,9 @@
 import logging, sys
 import pymol2
-from choppa.render.utils import show_contacts, get_ligand_resnames
+from io import StringIO
+from rdkit import Chem
+
+from choppa.render.utils import show_contacts, get_ligand_resnames_from_pdb_str
 
 logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 logger = logging.getLogger()
@@ -12,16 +15,19 @@ class PYMOL():
     settings in the GUI application themselves, but a combination of pre-set `ray` settings is provided
     in the `.pse` file. 
     """
-    def __init__(self, filled_aligned_fitness_dict, complex, complex_pdb_str, fitness_threshold):
+    def __init__(self, filled_aligned_fitness_dict, complex, complex_rdkit, fitness_threshold):
         self.fitness_dict = filled_aligned_fitness_dict
         self.complex = complex
-        self.complex_pdb_str = complex_pdb_str
         self.fitness_threshold = fitness_threshold
 
+        # get the PDB file as a string from RDKit
+        self.complex_pdb_str = Chem.MolToPDBBlock(complex_rdkit)
+
     def pymol_start_session(self):
         """
         Boots up a session with the publicly available PyMOL API.
         """
+        logger.info("Starting PyMOL session")
         p = pymol2.PyMOL()
         p.start()
         return p
@@ -31,6 +37,7 @@ def pymol_setup_system(self, p, remove_solvents=True):
         Sets up the protein (and ligands if present) without changing any of the looks. Also removes
         some stuff we're not interested in.
         """
+        logger.info("PyMOL session: setting up system")
         p.cmd.read_pdbstr(self.complex_pdb_str, 'complex')
 
         if remove_solvents:
@@ -61,7 +68,6 @@ def pymol_color_coder(self):
         - residue_mutability_levels : `{residue index : number of fit mutations, ..}`
         - mutability_color_dict : `{number of fit mutations : color, ..}`
         """
-
         # instantiate a dict to populate.
         residue_mutability_levels = {}
         for k in [
@@ -82,10 +88,11 @@ def pymol_color_coder(self):
                 residue_mutability_levels[f"n_fit_{self.count_fit_residues(fitness_data)}"].append(str(i))
             else:
                 residue_mutability_levels["no_fitness_data"].append(str(i))
-        
+
         # make the index values separated by '+' so that PyMOL can read it
         for k,v in residue_mutability_levels.items():
             residue_mutability_levels[k] = "+".join(v)
+        logger.info(f"PyMOL session: fitness degree per residue found using threshold {self.fitness_threshold}:\n{residue_mutability_levels}\n")
 
         mutability_color_dict = { # TODO: convert color coding to match 3DMol color specs
             "n_fit_0" : "white",
@@ -103,6 +110,8 @@ def pymol_color_by_fitness(self, p):
         With a pymol session set up with a system using `self.pymol_setup_system()`,
         integrates `fitness` data by coloring residues by mutability degree.
         """
+        logger.info("PyMOL session: coloring system surface with fitness data..")
+
         # first get the fitness degree per residue and what colors to make them
         residue_mutability_levels, mutability_color_dict = self.pymol_color_coder()
 
@@ -121,26 +130,78 @@ def pymol_color_by_fitness(self, p):
 
         for fitness_degree_name, color in mutability_color_dict.items():
             p.cmd.set("surface_color", color, f"({fitness_degree_name})")
+
+        return mutability_color_dict
+
+    def pymol_select_components(self, p):
+        """
+        Makes selections in PyMOL for ligand, protein, binding site. Returns whether there
+        is/are (a) ligand(s) present in the system.
+        """
+        ligands = get_ligand_resnames_from_pdb_str(self.complex_pdb_str)
+        if ligands:
+            # need to use the PyMOL DSL for selection. Construct the string first.
+            ligand_selector = f"resn {ligands[0]}"
+            for lig in ligands[1:]: # add additional ligand entries, if there are any
+                ligand_selector += f" and resn {lig}"
 
-    def pymol_prettify_system(self, p):
+        # make the selections
+        p.cmd.select("ligand", ligand_selector)
+        p.cmd.select( # PyMOL will have made only the protein cartoon, so can just select that way
+            "receptor", "rep cartoon"
+        )  
+
+        return ligands
+
+    def pymol_prettify_system(self, p, ligands_in_system):
         """
         With a pymol session set up with a system using `self.pymol_setup_system()`,
-        makes the session pretty.
+        makes the session pretty. This code isn't pretty though, that's just because
+        of how the PyMOL API is constructed.
         """
-
-        # tmp
-        p.cmd.show("surface")
+        logger.info("PyMOL session: prettifying view")
+
+        # reset the view
+        p.cmd.select("None") 
+        p.cmd.set("bg_rgb", "white")
+        p.cmd.bg_color("white")
+        p.cmd.hide("everything")
+
+        # show the protein surface
+        p.cmd.show("surface", "receptor")
+        p.cmd.set("surface_mode", 3)
+
+        # set some variables to improve the image when the user `ray`s the session once loaded
+        # p.cmd.set("surface_quality", 2) # turn on after dev
+        p.cmd.set("antialias", 2)
 
-    def pymol_add_interactions(self, p):
+        if ligands_in_system:
+            # select the ligand and subpocket residues, show them as sticks w/o nonpolar Hs
+            # TODO: set ligand stick color that conforms to HTML view
+            p.cmd.show("sticks", "ligand")
+            p.cmd.show("spheres", "ligand")
+            p.cmd.set("stick_radius", "0.15")
+            p.cmd.set("sphere_scale", "0.15")
+
+    def pymol_add_interactions(self, p, mutability_color_dict):
         """
         Adds interactions to pymol session if a ligand is present. Interactions are colored by
         fitness of contacted residues, not by interaction type.
         """
+        logger.info("PyMOL session: adding ligand-protein interactions (contacts) colored by fitness degree")
+
+        # for each degree of fitness (0, 1, .. 5, no_data), show ligand-protein contacts and color them
+        # in the same way that we colored the residue surface
+        for fitness_degree_selection, color in mutability_color_dict.items():
+            show_contacts(p, fitness_degree_selection, "ligand", contact_color=color)
+
+        #TODO: color backbone interactions green instead. How can we retrieve that data from pymol?
 
     def pymol_write_session(self, p, out_filename):
         """
         Writes out a pymol session to a `.pse` file.
         """
+        logger.info("PyMOL session: writing session file to {out_filename}")
         p.cmd.save(out_filename)
 
     def render(self):
@@ -154,10 +215,17 @@ def render(self):
         self.pymol_setup_system(p)
 
         # color it by fitness
-        self.pymol_color_by_fitness(p)
+        mutability_color_dict = self.pymol_color_by_fitness(p)
+
+        # make selections and figure out whether there is a ligand present
+        ligands_in_system = self.pymol_select_components(p)
 
         # make the view pretty
-        self.pymol_prettify_system(p)
+        self.pymol_prettify_system(p, ligands_in_system)
+
+        if ligands_in_system:
+            # add interactions
+            self.pymol_add_interactions(p, mutability_color_dict)
 
         # finally write to a session file (`.pse`)
         self.pymol_write_session(p, "test_sess.pse")
@@ -179,10 +247,12 @@ class HTML():
                                     # confidence_colname="confidence"
                                     ).get_fitness_basedict()
     complex = ComplexFactory(TOY_COMPLEX).load_pdb()
+    complex_rdkit = ComplexFactory(TOY_COMPLEX).load_pdb_rdkit()
 
     from choppa.align.align import AlignFactory
     filled_aligned_fitness_dict = AlignFactory(fitness_dict, complex).align_fitness()
 
     PYMOL(filled_aligned_fitness_dict, 
-          complex, ComplexFactory(TOY_COMPLEX).load_pdb_string(), 
+          complex,
+          complex_rdkit,
           fitness_threshold=0.7).render()