updates for chain-wise calculations and fixes multichain processing

frustratometer/classes/Frustratometer.py

@@ -211,7 +211,10 @@ 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=4000,
+                          ) -> np.array:
         """
         Calculates frustration index values.
 
@@ -225,6 +228,8 @@ def frustration(self, sequence:str = None, kind:str = 'singleresidue', mask:np.a
             A 2D Boolean array that determines which residue pairs should be considered in the energy computation. The mask should have dimensions (L, L), where L is the length of the sequence.
         aa_freq: np.array
             Array of frequencies of all 21 possible amino acids within sequence
+        n_decoys: int
+            Number of decoys
 
         Returns
         -------
@@ -243,8 +248,10 @@ def frustration(self, sequence:str = None, kind:str = 'singleresidue', mask:np.a
             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)
+                #TODO: Correct this function for different aa_freq than WT --- think this has been done already?
+                frustration_values = self.configurational_frustration(aa_freq=aa_freq, correction=correction, n_decoys=n_decoys,)
+                assert np.all(frustration_values==frustration_values.T), f"configurational frustration matrix was not symmetric! Max difference: {np.max(frustration_values-frustration_values.T)}"
+                return frustration_values
             if aa_freq is None:
                 aa_freq = self.contact_freq
             frustration_values=frustration.compute_pair_frustration(decoy_fluctuation, aa_freq, correction)
@@ -321,7 +328,8 @@ def vmd(self, sequence: str = None, single:Union[str,np.array] = 'singleresidue'
         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)
+                                      max_connections=max_connections, movie_name=movie_name, still_image_name=still_image_name,
+                                                                                                               min_contact_distance=self.min_contact_distance)
         frustration.call_vmd(self.pdb_file, tcl_script)
 
     def view_pair_frustration(self, sequence:str = None, pair:str = 'mutational', aa_freq:np.array = None):
@@ -338,6 +346,8 @@ def view_pair_frustration(self, sequence:str = None, pair:str = 'mutational', aa
         aa_freq: np.array
             Array of frequencies of all 21 possible amino acids within sequence
         """
+        if type(self.chain) == list:
+            raise NotImplementedError("This function not supported for multichain systems")
         import py3Dmol
 
         if sequence is None:

frustratometer/classes/Structure.py

@@ -14,7 +14,7 @@
 class Structure:
 
     def __init__(self, pdb_file: Union[Path,str], chain: Union[str,None]=None, seq_selection: str = None, aligned_sequence: str = None, filtered_aligned_sequence: str = None,
-                distance_matrix_method:str = 'CB', pdb_directory: Path = Path.cwd(), repair_pdb:bool = True)->object:
+                distance_matrix_method:str = 'CB', pdb_directory: Path = Path.cwd(), repair_pdb:bool = True, return_distance_midpoints:bool = False)->object:
 
         """
         Generates structure object. Both PDB and CIF format files are accepted as input.
@@ -55,6 +55,11 @@ def __init__(self, pdb_file: Union[Path,str], chain: Union[str,None]=None, seq_s
             If True, provided pdb file will be repaired with missing residues inserted and heteroatoms removed.
             Note that a pdb file will be produced, regardless of input file format.
 
+        return_distance_midpoints: bool
+            Whether to return a matrix of the same shape as distance_matrix representing the same contacts as distance_matrix
+            that indicates the absolute coordinates of the midpoint between the pair of atoms. This helps us compute the pair distribution
+            functions of the different classes of contacts. So this matrix isn't really a matrix because each "element" has 3 channels: x, y, and z
+
         Returns
         -------
         Structure object
@@ -77,7 +82,7 @@ def __init__(self, pdb_file: Union[Path,str], chain: Union[str,None]=None, seq_s
 
         self.pdbID=pdb_file.stem
         self.pdb_file=pdb_file
-        self.chain=chain
+        self.chain=chain # will be None if no chain supplied
         self.distance_matrix_method=distance_matrix_method
         self.filtered_aligned_sequence=filtered_aligned_sequence
         self.aligned_sequence=aligned_sequence
@@ -88,18 +93,18 @@ def __init__(self, pdb_file: Union[Path,str], chain: Union[str,None]=None, seq_s
             self.init_index_shift=0
 
             if repair_pdb:
-                fixer=pdb.repair_pdb(pdb_file, chain, pdb_directory)
+                fixer=pdb.repair_pdb(pdb_file, chain, pdb_directory) # for this function, chain can be str or list (or None)
                 self.pdb_file=str(pdb_directory/f"{self.pdbID}_cleaned.pdb")
 
             if ".pdb" in str(pdb_file) or repair_pdb==True:
-                self.structure = prody.parsePDB(str(self.pdb_file), chain=self.chain).select(f"protein")
+                self.structure = prody.parsePDB(str(self.pdb_file), chain=self.chain).select(f"protein") # for this function, chain should be a string containing the chain ids like "AB" or "A B"
             else:
-                self.structure=prody.parseMMCIF(str(self.pdb_file),chain=self.chain).select(f"protein")
+                self.structure=prody.parseMMCIF(str(self.pdb_file),chain=self.chain).select(f"protein")  # for this function, chain should be a string containing the chain ids like "AB" or "A B"
         else:
             assert len(self.seq_selection.replace("to"," to ").replace(":"," : ").split())>=4, "Please correctly input your residue selection"
 
             if self.chain==None:
-                raise ValueError("Please provide a chain name")
+                raise ValueError("self.chain==None. Please provide chain name(s)")
 
             self.init_index=int(self.seq_selection.replace("to"," to ").replace(":"," : ").split()[1].replace("`",""))
             self.fin_index=int(self.seq_selection.replace("to"," to ").replace(":"," : ").split()[3].replace("`",""))
@@ -116,7 +121,7 @@ def __init__(self, pdb_file: Union[Path,str], chain: Union[str,None]=None, seq_s
 
             with open(pdb_file,"r") as f:
                 for line in f:
-                    if line.split()[0]=="ATOM" and line.split()[4+shift]==self.chain:
+                    if line.split()[0]=="ATOM" and (line.split()[4+shift] in self.chain):
                         try:
                             res_index=''.join(i for i in line.split()[5+index_shift] if i.isdigit())
                             next_res_index=''.join(i for i in next(f).split()[5+index_shift] if i.isdigit())
@@ -133,27 +138,35 @@ def __init__(self, pdb_file: Union[Path,str], chain: Union[str,None]=None, seq_s
                 self.init_index_shift=self.init_index-self.pdb_init_index
                 self.fin_index_shift=self.fin_index-self.pdb_init_index+1
                 if repair_pdb:
-                    fixer=pdb.repair_pdb(pdb_file, chain, pdb_directory)
+                    fixer=pdb.repair_pdb(pdb_file, chain, pdb_directory) # for this function, chain can be str or list (or None)
                     self.pdb_file=f"{pdb_directory}/{self.pdbID}_cleaned.pdb"
                     self.select_gap_indices=[i for i in gap_indices if self.init_index<=i<=self.fin_index]
                     self.fin_index_shift-=len(self.select_gap_indices)
-                    self.seq_selection=f"resnum `{self.init_index_shift+1}to{self.fin_index_shift}`"
+                    #self.seq_selection=f"resnum `{self.init_index_shift+1}to{self.fin_index_shift}`" # WE'RE KEEPING IDs NOW SO DON'T WANT TO DO THIS!!!!
             elif "resindex" in self.seq_selection:
                 self.init_index_shift=self.init_index
                 self.fin_index_shift=self.fin_index+1
                 if repair_pdb:
-                    fixer=pdb.repair_pdb(pdb_file, chain, pdb_directory)
+                    fixer=pdb.repair_pdb(pdb_file, chain, pdb_directory) # for this function, chain can be str or list (or None)
                     self.pdb_file=f"{pdb_directory}/{self.pdbID}_cleaned.pdb"
-                    self.chain="A"
+                    # self.chain="A" I don't know why we would want to change the chain ID here
 
             if ".pdb" in str(pdb_file) or repair_pdb==True:
                 self.structure = prody.parsePDB(str(self.pdb_file), chain=self.chain).select(f"protein and {self.seq_selection}")
             else:
                 self.structure=prody.parseMMCIF(str(self.pdb_file),chain=self.chain).select(f"protein and {self.seq_selection}")
 
-        self.sequence=pdb.get_sequence(self.pdb_file,self.chain)
-        self.distance_matrix=pdb.get_distance_matrix(pdb_file=self.pdb_file,chain=self.chain,
-                                                     method=self.distance_matrix_method)
+        self.sequence, self.start_mask = pdb.get_sequence(self.pdb_file,self.chain,return_start_mask=True) # this function can now accept chain as list or string
+        if return_distance_midpoints:
+            self.distance_matrix, self.midpoint_matrix = pdb.get_distance_matrix(pdb_file=self.pdb_file,chain=self.chain, # for this function, chain should be a string containing the chain ids                                     
+                                                                                 method=self.distance_matrix_method,      # separated by a space, like "A B"
+                                                                                 return_distance_midpoints=True)
+        else:
+            self.distance_matrix=pdb.get_distance_matrix(pdb_file=self.pdb_file,chain=self.chain, # for this function, chain should be a string containing the chain ids                                     
+                                                     method=self.distance_matrix_method,      # separated by a space, like "A B"
+                                                    return_distance_midpoints=False)            
+            self.midpoint_matrix = None
+
         self.full_pdb_distance_matrix=self.distance_matrix
 
         self.z_coordinates=self.structure.select('((name CB) or (resname GLY and name CA))').getCoords()
@@ -180,7 +193,7 @@ def __init__(self, pdb_file: Union[Path,str], chain: Union[str,None]=None, seq_s
             else:
                 self.full_to_aligned_index_dict=dict(zip(range(self.init_index_shift,self.fin_index_shift+1), range(len(self.sequence))))
                 self.mapped_distance_matrix=self.distance_matrix
-
+       
     @classmethod
     def full_pdb(cls,pdb_file: Union[Path,str], chain: Union[str,None]=None, aligned_sequence: str = None, filtered_aligned_sequence: str = None,
                 distance_matrix_method:str = 'CB', pdb_directory: Path = Path.cwd(), repair_pdb:bool = True):