Temp decoy

frustratometer/__init__.py

@@ -13,6 +13,7 @@
 from . import align
 from . import frustration
 from . import optimization
+from . import numba_util
 
 # Handle versioneer
 from ._version import get_versions

frustratometer/classes/DCA.py

@@ -78,6 +78,8 @@ class DCA(Frustratometer):
     #     self._decoy_fluctuation = {}
     #     return self
 
+    alphabet = '-ACDEFGHIKLMNPQRSTVWY'
+
     @classmethod
     def from_potts_model_file(cls,pdb_structure: object,
                               potts_model_file: Union[Path,str] = None,
@@ -146,8 +148,8 @@ def from_potts_model_file(cls,pdb_structure: object,
             self.potts_model["J"]= self.potts_model["familycouplings"].reshape(int(len(self.filtered_aligned_sequence)),21,int(len(self.filtered_aligned_sequence)),21).transpose(0,2,1,3)
 
         if self.filtered_aligned_sequence is not None:
-            self.aa_freq = frustration.compute_aa_freq(self.sequence)
-            self.contact_freq = frustration.compute_contact_freq(self.sequence)
+            self.aa_freq = frustration.compute_aa_freq(self.sequence, self.alphabet)
+            self.contact_freq = frustration.compute_contact_freq(self.sequence, self.alphabet)
         else:
             self.aa_freq = None
             self.contact_freq = None   
@@ -222,8 +224,8 @@ def from_pottsmodel(cls,pdb_structure : object,
             self.potts_model["J"]= self.potts_model["familycouplings"].reshape(int(len(self.filtered_aligned_sequence)),21,int(len(self.filtered_aligned_sequence)),21).transpose(0,2,1,3)
 
         if self.filtered_aligned_sequence is not None:
-            self.aa_freq = frustration.compute_aa_freq(self.sequence)
-            self.contact_freq = frustration.compute_contact_freq(self.sequence)
+            self.aa_freq = frustration.compute_aa_freq(self.sequence, self.alphabet)
+            self.contact_freq = frustration.compute_contact_freq(self.sequence, self.alphabet)
         else:
             self.aa_freq = None
             self.contact_freq = None   

frustratometer/classes/Frustratometer.py

@@ -61,9 +61,27 @@ def native_energy(self,sequence:str = None,ignore_couplings_of_gaps:bool=False,i
         if sequence is None:
             sequence=self.sequence
         else:
-            return frustration.compute_native_energy(sequence, self.potts_model, self.mask,ignore_couplings_of_gaps,ignore_fields_of_gaps)
+            return frustration.compute_native_energy(sequence, self.potts_model, self.mask, self.alphabet, 
+                                                     ignore_couplings_of_gaps, ignore_fields_of_gaps)
         if not self._native_energy:
-            self._native_energy=frustration.compute_native_energy(sequence, self.potts_model, self.mask,ignore_couplings_of_gaps,ignore_fields_of_gaps)
+            self._native_energy=frustration.compute_native_energy(sequence, self.potts_model, self.mask, self.alphabet, 
+                                                                  ignore_couplings_of_gaps, ignore_fields_of_gaps)
+        else:
+            new = frustration.compute_native_energy(
+                        sequence, self.potts_model, self.mask, self.alphabet, 
+                        ignore_couplings_of_gaps, ignore_fields_of_gaps)
+            if not (self._native_energy == new):
+                raise AssertionError(f"""
+                                        It seems that you have changed parameters of an object such that
+                                        the native energy of your system is now different from what it was
+                                        originally computed to be. Our code probably should prevent this 
+                                        from happening, but you can prevent it too by not changing the alphabet
+                                        or any other parameters after initializing your DCA or AWSEM-family
+                                        class (anything that inherits from _AWSEMBase).
+
+                                        Previous value of {self.__class__}._native_energy: {self._native_energy}
+                                        New value of {self.__class__}._native_energy: {new}""")
+
         energy_value=self._native_energy
         return energy_value
 
@@ -89,7 +107,7 @@ def sequences_energies(self, sequences:np.array, split_couplings_and_fields:bool
         output (if split_couplings_and_fields==True): np.array
             Array containing computed fields and couplings energies of the protein sequences. 
         """
-        output=frustration.compute_sequences_energy(sequences, self.potts_model, self.mask, split_couplings_and_fields)
+        output=frustration.compute_sequences_energy(sequences, self.potts_model, self.mask, self.alphabet, split_couplings_and_fields)
         return output
 
     def fields_energy(self, sequence:str = None, ignore_fields_of_gaps:bool = False) -> float:
@@ -114,7 +132,7 @@ def fields_energy(self, sequence:str = None, ignore_fields_of_gaps:bool = False)
         """
         if sequence is None:
             sequence=self.sequence
-        fields_energy=frustration.compute_fields_energy(sequence, self.potts_model,ignore_fields_of_gaps)
+        fields_energy=frustration.compute_fields_energy(sequence, self.potts_model, self.alphabet, ignore_fields_of_gaps)
         return fields_energy
 
     def couplings_energy(self, sequence:str = None,ignore_couplings_of_gaps:bool = False) -> float:
@@ -139,7 +157,7 @@ def couplings_energy(self, sequence:str = None,ignore_couplings_of_gaps:bool = F
         """
         if sequence is None:
             sequence=self.sequence
-        couplings_energy=frustration.compute_couplings_energy(sequence, self.potts_model, self.mask,ignore_couplings_of_gaps)
+        couplings_energy=frustration.compute_couplings_energy(sequence, self.potts_model, self.mask, self.alphabet, ignore_couplings_of_gaps)
         return couplings_energy
 
     def decoy_fluctuation(self, sequence:str = None,kind:str = 'singleresidue',mask:np.array = None) -> np.array:
@@ -167,13 +185,13 @@ def decoy_fluctuation(self, sequence:str = None,kind:str = 'singleresidue',mask:
         if not isinstance(mask, np.ndarray):
             mask=self.mask
         if kind == 'singleresidue':
-            fluctuation = frustration.compute_singleresidue_decoy_energy_fluctuation(sequence, self.potts_model, mask)
+            fluctuation = frustration.compute_singleresidue_decoy_energy_fluctuation(sequence, self.potts_model, mask, self.alphabet)
         elif kind == 'mutational':
-            fluctuation = frustration.compute_mutational_decoy_energy_fluctuation(sequence, self.potts_model, mask)
+            fluctuation = frustration.compute_mutational_decoy_energy_fluctuation(sequence, self.potts_model, mask, self.alphabet)
         elif kind == 'configurational':
-            fluctuation = frustration.compute_configurational_decoy_energy_fluctuation(sequence, self.potts_model, mask)
+            fluctuation = frustration.compute_configurational_decoy_energy_fluctuation(sequence, self.potts_model, mask, self.alphabet)
         elif kind == 'contact':
-            fluctuation = frustration.compute_contact_decoy_energy_fluctuation(sequence, self.potts_model, mask)
+            fluctuation = frustration.compute_contact_decoy_energy_fluctuation(sequence, self.potts_model, mask, self.alphabet)
         else:
             raise Exception("Wrong kind of decoy generation selected")
         self._decoy_fluctuation[kind] = fluctuation
@@ -211,7 +229,8 @@ def scores(self):
         """
         return frustration.compute_scores(self.potts_model)
 
-    def frustration(self, sequence:str = None, kind:str = 'singleresidue', mask:np.array = None, aa_freq:np.array = None, correction:int = 0) -> np.array:
+    def frustration(self, sequence:str = None, kind:str = 'singleresidue', mask:np.array = None, aa_freq:np.array = None, 
+                          correction:int = 0, n_decoys:int = 4000) -> np.array:
         """
         Calculates frustration index values.
 
@@ -242,9 +261,11 @@ def frustration(self, sequence:str = None, kind:str = 'singleresidue', mask:np.a
             frustration_values=frustration.compute_single_frustration(decoy_fluctuation, aa_freq, correction)
             return frustration_values
         elif kind in ['mutational', 'configurational', 'contact']:
-            if kind == 'configurational' and 'configurational_frustration' in dir(self):
-                #TODO: Correct this function for different aa_freq than WT
-                return self.configurational_frustration(None, correction)
+            if kind == 'configurational':
+                if 'configurational_frustration' in dir(self):
+                    return self.configurational_frustration(aa_freq=aa_freq, correction=correction, n_decoys=n_decoys)
+                else:
+                    raise ValueError("kind='configurational' may only be used on objects implementing self.configurational_frustration")
             if aa_freq is None:
                 aa_freq = self.contact_freq
             frustration_values=frustration.compute_pair_frustration(decoy_fluctuation, aa_freq, correction)
@@ -268,7 +289,7 @@ def plot_decoy_energy(self, sequence:str = None, kind:str = 'singleresidue', met
         native_energy = self.native_energy(sequence=sequence)
         decoy_energy = self.decoy_energy(kind=kind,sequence=sequence)
         if kind == 'singleresidue':
-            g = frustration.plot_singleresidue_decoy_energy(decoy_energy, native_energy, method)
+            g = frustration.plot_singleresidue_decoy_energy(decoy_energy, native_energy, method, self.alphabet)
             return g
 
     def roc(self):
@@ -292,6 +313,7 @@ def auc(self):
         return frustration.compute_auc(self.roc())
 
     def vmd(self, sequence: str = None, single:Union[str,np.array] = 'singleresidue', pair:Union[str,np.array] = 'mutational',
+             tcl_script:str = 'frustration.tcl', call_vmd:bool=True,
              aa_freq:np.array = None, correction:int = 0, max_connections:Union[int,None] = None, movie_name=None, still_image_name=None):
         """
         Calculates frustration indices and superimposes frustration patterns onto PDB structure using the VMD software.
@@ -317,12 +339,14 @@ def vmd(self, sequence: str = None, single:Union[str,np.array] = 'singleresidue'
                     from the sequence that was passed to this vmd function. Proceeding further may not\n\
                     perform the computation that you intend to perform.")
 
-
+        #breakpoint()
         tcl_script = frustration.write_tcl_script(self.pdb_file, self.chain, self.mask, self.distance_matrix, self.distance_cutoff,
                                       -self.frustration(kind=single, sequence=sequence, aa_freq=aa_freq),
                                       -self.frustration(kind=pair, sequence=sequence, aa_freq=aa_freq),
-                                      max_connections=max_connections, movie_name=movie_name, still_image_name=still_image_name)
-        frustration.call_vmd(self.pdb_file, tcl_script)
+                                      max_connections=max_connections, movie_name=movie_name, still_image_name=still_image_name,
+                                      tcl_script=tcl_script,)
+        if call_vmd:
+            frustration.call_vmd(self.pdb_file, tcl_script)
 
     def view_pair_frustration(self, sequence:str = None, pair:str = 'mutational', aa_freq:np.array = None):
         """

frustratometer/classes/Gamma.py

@@ -34,6 +34,10 @@ def __init__(self, data, segment_definition=None, description=None, alphabet=Non
 
         self._validate_segments()
 
+    @property
+    def q(self):
+        return len(self.alphabet)
+
     def _init_from_array(self, gamma_array):
         self.gamma_array = gamma_array
 
@@ -399,7 +403,7 @@ def correlate_segments(self, other):
         return correlations
 
     # Plotting
-    def plot_gamma(self, new_order=None):
+    def plot_gamma(self, new_order=None, scale=[-5,5]):
         import matplotlib.pyplot as plt
         import seaborn as sns
         if new_order:
@@ -408,16 +412,21 @@ def plot_gamma(self, new_order=None):
 
         # Plot setup
         f, axes = plt.subplots(2, 2, figsize=(18, 16))
-        titles = ['Burial Gammas', 'Direct Gammas', 'Water Gammas', 'Protein Gammas']
-
+        f.subplots_adjust(hspace=50) # fix overlap between axis ticks of upper subplots and titles of lower subplots
+        titles = ['Burial Gammas', 'Direct Gammas', 'Protein Gammas', 'Water Gammas']
         for i, (title, name) in enumerate(zip(titles, segments)):
             ax = axes[i // 2, i % 2]
-            sns.heatmap(segments[name].reshape(-1, 20), ax=ax, cmap='RdBu_r', center=0)
+            foo = sns.heatmap(segments[name].reshape(-1, 20), ax=ax, cmap='RdBu_r', center=0, vmin=scale[0], vmax=scale[1])
+            foo.collections[0].colorbar.ax.tick_params(labelsize=16)
             ax.set_title(title)
             ax.set_xticks(np.arange(len(self.alphabet)) + 0.5)
-            ax.set_xticklabels(self.alphabet)
-            ax.set_yticks(np.arange(segments[name].shape[0] // 20) + 0.5)
-            ax.set_yticklabels(range(segments[name].shape[0] // 20))
+            ax.set_xticklabels(self.alphabet, size=16)
+            if i==0: # burial
+                ax.set_yticks([0.5,1.5,2.5])
+                ax.set_yticklabels(['low','medium','high'], rotation=45, size=16)
+            else: # direct, prot, or wat
+                ax.set_yticks(np.arange(len(self.alphabet)) + 0.5)
+                ax.set_yticklabels(self.alphabet, rotation=0, fontsize=16)
 
         plt.tight_layout()
         plt.show()
@@ -648,4 +657,4 @@ class O():
 
     self.gamma1 = Gamma(np.arange(0,1260,1))
     self.gamma2 = Gamma(np.arange(0,1260,1)*5+10)
-    self.gamma3 = Gamma(np.arange(1260,0,-1)*2-4)
+    self.gamma3 = Gamma(np.arange(1260,0,-1)*2-4)