stcrpy 1.0.0__py3-none-any.whl → 1.0.5__py3-none-any.whl

stcrpy/tcr_processing/TCR.py

@@ -1,13 +1,17 @@
 """
 Created on 3rd April 2024
-@author: Nele Quast based on work by Dunbar and Leem 
+Nele Quast based on work by Dunbar and Leem
 The TCR class.
 """
 
 import sys
 import warnings
 
+from Bio import BiopythonWarning
+
 from .Entity import Entity
+from .TCRchain import TCRchain
+from .utils.region_definitions import IMGT_VARIABLE_DOMAIN
 
 try:
     from .. import tcr_interactions
@@ -18,17 +22,34 @@ except ImportError as e:
     print(e)
 
 
+class TCRError(Exception):
+    """Error raised when there is an issue with the TCR."""
+
+
 class TCR(Entity):
     """
-    TCR class. This is generic which is inherited later.
-    Holds paired TCR chains.
+    TCR class. Inherits from PDB.Entity.
+    This is a base class for TCR strucutres, enabling antigen and MHC association.
+    abTCR and gdTCR are the instantiated subclasses of this class.
     """
 
     def _add_antigen(self, antigen=None):
+        """
+        Append associated antigen to TCR antigen field.
+
+        Args:
+            antigen (Antigen, optional): Antigen to associate with TCR. Defaults to None.
+        """
         if antigen not in self.antigen:
             self.antigen.append(antigen)
 
     def _add_mhc(self, mhc=None):
+        """
+        Append associated MHC to TCR MHC field. If antigen are associted with MHC but not TCR, add them to TCR antigen.
+
+        Args:
+            mhc (MHC, optional): MHC to associate with TCR. Defaults to None.
+        """
         if mhc not in self.MHC:
             self.MHC.append(mhc)
             # If there are any het antigens that are in the MHC but not in close proximity of the TCR
@@ -36,19 +57,41 @@ class TCR(Entity):
             if set(mhc.antigen) - set(self.antigen):
                 self.antigen.extend(mhc.antigen)
 
+    def copy(self, copy_siblings = True ):
+        """
+        Return a copy of the TCR object. This returns a shallow copy of the TCR object.
+        If the copy_siblings flag is set to True, the antigen and MHC objects will also be copied. Warning - if the copy_siblings flag is set to False, the antigen and MHC objects will not be copied, and the reference will still point to the same MHC and antigen objects as the original.
+
+        copy_siblings: Whether to copy sibling entities (ie. MHC and Antigen objects). Default True. 
+
+        """
+        shallow = super().copy()
+        if copy_siblings:
+            shallow.antigen = [a.copy() for a in self.get_antigen()]
+            shallow.MHC = [m.copy(copy_siblings=False) for m in self.get_MHC()]
+            for m in shallow.MHC:
+                m.tcr = [t.copy(copy_siblings=False) if t.id != shallow.id else shallow for t in m.tcr]
+                m.antigen = [ag.copy() if ag.id not in [a.id for a in shallow.antigen] else [a for a in shallow.antigen if a.id==ag.id][0] for ag in m.antigen]
+
+        return shallow
+
     def get_antigen(self):
         """
-        Return a list of bound antigens.
+        Return a list of TCR associated antigens.
         """
         return self.antigen
 
     def get_MHC(self):
-        """ """
+        """
+        Return a list of TCR associated MHCs.
+        """
         return self.MHC
 
     def is_bound(self):
-        """
-        Check whether this TCR is bound to an antigen
+        """True or False if the TCR is associated with an antigen.
+
+        Returns:
+            bool: Whether TCR is associated with an antigen.
         """
         if self.get_antigen():
             return True
@@ -56,22 +99,40 @@ class TCR(Entity):
             return False
 
     def get_chains(self):
+        """Returns generator of TCR chains.
+
+        Yields:
+            Chain: TCR chain
+        """
         for c in self:
             yield c
 
     def get_residues(self):
+        """Returns generator of TCR residues.
+
+        Yields:
+            Residue: TCR residue
+        """
         for c in self.get_chains():
             for r in c:
                 yield r
 
     def get_atoms(self):
+        """Returns generator of TCR atoms.
+
+        Yields:
+            Atom: TCR atoms
+        """
         for r in self.get_residues():
             for a in r:
                 yield a
 
     def get_frameworks(self):
         """
-        Obtain framework regions from a TCR structure object.
+        Obtain framework regions from a TCR structure object as generator.
+
+        Yields:
+            Fragment: TCR framework regions
         """
         for f in self.get_fragments():
             if "fw" in f.id:
@@ -79,15 +140,20 @@ class TCR(Entity):
 
     def get_CDRs(self):
         """
-        Obtain CDR loops from a TCR structure object
+        Obtain complementarity determining regions (CDRs) from a TCR structure object as generator.
+
+        Yields:
+            Fragment: TCR CDR regions
         """
         for f in self.get_fragments():
             if "cdr" in f.id:
                 yield f
 
     def get_TCR_type(self):
-        """
-        Get the TCR type
+        """Get TCR type according to variable region assignments.
+
+        Returns:
+            str: TCR type (abTCR, gdTCR, dbTCR)
         """
         if hasattr(self, "tcr_type"):
             return self.tcr_type
@@ -102,9 +168,22 @@ class TCR(Entity):
             return self.tcr_type
 
     def get_germline_assignments(self):
+        """Retrive germline assignments for all TCR chains.
+        This is a dictionary with the chain ID as key and the germline assignments as value.
+
+        Returns:
+            dict: dict with TCR chain ID as key and germline assignments as value
+        """
         return {c.id: c.get_germline_assignments() for c in self.get_chains()}
 
     def get_MHC_allele_assignments(self):
+        """
+        Retrieve MHC allele assignments for all TCR associated MHCs.
+        This is a list of dictionaries with the MHC ID as key and the allele assignments as value.
+
+        Returns:
+            dict: dict with MHC chain ID as key and allele assignments as value
+        """
         return [
             (
                 mhc.get_allele_assignments()
@@ -116,6 +195,11 @@ class TCR(Entity):
         ]
 
     def get_germlines_and_alleles(self):
+        """Get all germline and allele assignments for TCR and MHC chains as a dictionary with the chain ID as key and the germline assignments as value.
+
+        Returns:
+            dict: Dictionary of TCR germline and MHC allele assignemnts with amino acid sequences.
+        """
         from ..tcr_formats.tcr_formats import get_sequences
 
         germlines_and_alleles = {}
@@ -166,23 +250,79 @@ class TCR(Entity):
 
         return germlines_and_alleles
 
+    def get_chain_mapping(self):
+        """Get a dictionary of chain IDs to chain types.
+
+        Returns:
+            dict: Dictionary of chain IDs to chain types
+        """
+        tcr_chain_mapping = {v: k for k, v in self.get_domain_assignment().items()}
+        antigen_chain_mapping = {c.id: "Ag" for c in self.get_antigen()}
+        mhc_chain_mapping = {
+            c.id: c.chain_type for m in self.get_MHC() for c in m.get_chains()
+        }
+        chain_mapping = {
+            **tcr_chain_mapping,
+            **antigen_chain_mapping,
+            **mhc_chain_mapping,
+        }
+
+        return chain_mapping
+
     def save(self, save_as=None, tcr_only: bool = False, format: str = "pdb"):
+        """Save TCR object as PDB or MMCIF file.
+
+        Args:
+            save_as (str, optional): File path to save TCR to. Defaults to None.
+            tcr_only (bool, optional): Whether to save TCR only or to include MHC and antigen. Defaults to False.
+            format (str, optional): Whether to save as PDB or MMCIF. Defaults to "pdb".
+        """
         from . import TCRIO
 
         tcrio = TCRIO.TCRIO()
         tcrio.save(self, save_as=save_as, tcr_only=tcr_only, format=format)
 
     def get_scanning_angle(self, mode="rudolph"):
+        """
+        Returns TCR:pMHC complex scanning (aka crossing, incident) angle of TCR to MHC.
+        See paper for details.
+
+        Args:
+            mode (str, optional): Mode for calculating the scanning angle. Options "rudolph", "cys", "com". Defaults to "rudolph".
+
+        Returns:
+            float: Scanning angle of TCR to MHC in degrees
+        """
         if not hasattr(self, "geometry") or self.geometry.mode != mode:
             self.calculate_docking_geometry(mode=mode)
         return self.geometry.get_scanning_angle()
 
-    def get_pitch_angle(self, mode="rudolph"):
+    def get_pitch_angle(self, mode="cys"):
+        """
+        Returns TCR:pMHC complex pitch angle of TCR to MHC.
+        See paper for details.
+
+        Args:
+            mode (str, optional): Mode for calculating the scanning angle. Options "rudolph", "cys", "com". Defaults to "cys".
+
+        Returns:
+            float: Pitch angle of TCR to MHC in degrees
+        """
         if not hasattr(self, "geometry") or self.geometry.mode != mode:
             self.calculate_docking_geometry(mode=mode)
         return self.geometry.get_pitch_angle()
 
     def calculate_docking_geometry(self, mode="rudolph", as_df=False):
+        """Calculate docking geometry of TCR to MHC.
+        This is a wrapper function for the TCRGeom class.
+
+        Args:
+            mode (str, optional): Mode for calculating the geometry. Options "rudolph", "cys", "com". Defaults to "rudolph".
+            as_df (bool, optional): Whether to return as dictionary or dataframe. Defaults to False.
+
+        Returns:
+            [dict, DataFrame]: TCR to MHC geometry.
+        """
         if len(self.get_MHC()) == 0:
             warnings.warn(
                 f"No MHC found for TCR {self}. Docking geometry cannot be calcuated"
@@ -203,20 +343,39 @@ class TCR(Entity):
         return self.geometry.to_dict()
 
     def score_docking_geometry(self, **kwargs):
+        """
+        Score docking geometry of TCR to MHC.
+        This is a wrapper function for the TCRGeomFiltering class.
+        The score is calculated as the negative log of the TCR:pMHC complex geometry feature probabilities based on the distributions fit by maximum likelihood estimation of TCR to Class I MHC strucutres from STCRDab.
+        Please see the paper methods for details.
+
+        Returns:
+            float: TCR:pMHC complex score as negative log of TCR:pMHC complex geometry feature probabilities
+        """
         from ..tcr_geometry.TCRGeomFiltering import DockingGeometryFilter
 
         geom_filter = DockingGeometryFilter()
         if not hasattr(self, "geometry"):
             self.calculate_docking_geometry(mode="com")
         return geom_filter.score_docking_geometry(
-            self.geometry.get_scanning_angle(),
-            self.geometry.get_pitch_angle(),
+            self.geometry.get_scanning_angle(mode="com"),
+            self.geometry.get_pitch_angle(mode="com"),
             self.geometry.tcr_com[-1],  # z component of TCR centre of mass
         )
 
     def profile_peptide_interactions(
         self, renumber: bool = True, save_to: str = None, **kwargs
     ) -> "pd.DataFrame":
+        """
+        Profile the interactions of the peptide to the TCR and the MHC.
+
+        Args:
+            renumber (bool, optional): Whether to renumber the interacting residues. Defaults to True.
+            save_to (str, optional): Path to save intraction data to as csv. Defaults to None.
+
+        Returns:
+            pd.DataFrame: Dataframe of peptide interactions
+        """
         if len(self.get_antigen()) == 0:
             warnings.warn(
                 f"No peptide antigen found for TCR {self}. Peptide interactions cannot be profiled"
@@ -238,6 +397,25 @@ class TCR(Entity):
         return interactions
 
     def get_interaction_heatmap(self, plotting_kwargs={}, **interaction_kwargs):
+        """
+        Get interaction heatmap of TCR to MHC and peptide.
+        Generates heatmap image.
+        Plotting kwargs are passed to heatmap function.
+
+        Args:
+            plotting_kwargs (dict, optional):
+                save_as: path to save heatmap image to
+                interaction_type: type of interaction (eg. saltbridge, h_bond) to plot. All interactions are plotted by default.
+                antigen_name: name of antigen for plot title
+                mutation_index: index of antigen residues to highlight in plot
+                Defaults to {
+                    save_as:None,
+                    interaction_type:None,
+                    antigen_name:None,
+                    mutation_index:None
+                    }.
+            interaction_kwargs: kwargs for TCRInteractionProfiler class. See TCRInteractionProfiler for details.
+        """
         from ..tcr_interactions import TCRInteractionProfiler
 
         interaction_profiler = TCRInteractionProfiler.TCRInteractionProfiler(
@@ -251,6 +429,44 @@ class TCR(Entity):
     def profile_MHC_interactions(self):
         raise NotImplementedError
 
+    def get_TCR_angles(self):
+        from ..tcr_geometry.TCRAngle import TCRAngle
+
+        return TCRAngle().calculate_angles(self)
+
+    def crop(self, *, crop_mhc: bool = True, remove_het_atoms: bool = True) -> None:
+        """Crop TCR to variable domain and optionally crop MHC to antigen binding domain.
+
+        This method mutates the TCR object.
+
+        Args:
+            crop_mhc: crop mhc to antigen binding domain
+            remove_het_atoms: remove het atoms from structure as well
+
+        """
+        new_child_dict = {}
+        for chain in self:
+            new_chain = TCRchain(chain.id)
+
+            for residue in chain:
+                if residue.id[1] in IMGT_VARIABLE_DOMAIN or (not remove_het_atoms and residue.id[0] != ' '):
+                    new_chain.add(residue.copy())
+
+            new_chain.analyse(chain.chain_type)
+            new_chain.set_engineered(chain.engineered)
+            new_chain.xtra.update(chain.xtra)
+            new_child_dict[new_chain.id] = new_chain
+
+        for chain_id in new_child_dict:
+            del self[chain_id]
+
+        for new_chain in new_child_dict.values():
+            self.add(new_chain)
+
+        if crop_mhc:
+            for mhc in self.get_MHC():
+                mhc.crop(remove_het_atoms=remove_het_atoms)
+
     def _create_interaction_visualiser(self):
         """Function called during TCR initialisation checks if pymol is installed and assigns a visualisation method accordingly.
         If pymol is installed, method to generate interaction visualisations is returned.
@@ -266,6 +482,15 @@ class TCR(Entity):
             def visualise_interactions(
                 save_as=None, antigen_residues_to_highlight=None, **interaction_kwargs
             ):
+                """Visualise peptide interactions in pymol.
+
+                Args:
+                    save_as (str, optional): path to save pymol session to. Defaults to None.
+                    antigen_residues_to_highlight (list[int], optional): antigen residues to highlight red in pymol session. Defaults to None.
+                    **interaction_kwargs: kwargs for TCRInteractionProfiler class. See TCRInteractionProfiler for details.
+                Returns:
+                    str: path to saved pymol session
+                """
                 from ..tcr_interactions import TCRInteractionProfiler
 
                 interaction_profiler = TCRInteractionProfiler.TCRInteractionProfiler(
@@ -304,9 +529,39 @@ class TCR(Entity):
 
             return visualise_interactions
 
+    def standardise_chain_names(self):
+        """Raises NotImplementedError."""
+        raise NotImplementedError()
+
+    def _validate_chain_standardising(self) -> None:
+        if (hasattr(self, 'antigen') and len(self.antigen) > 1) or (hasattr(self, 'MHC') and len(self.MHC) > 1):
+            msg = 'More than one antigen or MHC molecule is not currently supported for standardising.'
+            raise TCRError(msg)
+
+    def _standardise_antigen_chain_names(self) -> None:
+        """Will give the antigen the chain id C. Does not support multiple antigens."""
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.antigen[0].id = 'C'
+
+    def _standardise_mhc_chain_names(self) -> None:
+        """Will give the MHC first chain id A and second chain B. Does not support more than one MHC molecule."""
+        self.MHC[0].standardise_chain_names()
+
 
 class abTCR(TCR):
+    """
+    abTCR class. Inherits from TCR.
+    This is a subclass of TCR for TCRs with alpha and beta chains.
+    """
     def __init__(self, c1, c2):
+        """
+        Initialise abTCR object. This is a subclass of TCR for TCRs with alpha and beta chains.
+
+        Args:
+            c1 (TCRchain): alpha or beta type TCR chain
+            c2 (TCRchain): alpha or beta type TCR chain
+        """
 
         if c1.chain_type == "B":
             Entity.__init__(self, c1.id + c2.id)
@@ -328,9 +583,22 @@ class abTCR(TCR):
         self.visualise_interactions = self._create_interaction_visualiser()
 
     def __repr__(self):
+        """
+        String representation of the abTCR object.
+
+        Returns:
+            str: String representation of the abTCR objec
+        """
         return "<TCR %s%s beta=%s; alpha=%s>" % (self.VB, self.VA, self.VB, self.VA)
 
     def _add_domain(self, chain):
+        """
+        Add a variable alpha or variable beta domain to the TCR object.
+        Links the domain to the chain ID.
+
+        Args:
+            chain (TCRchain): TCR chain whose domain is being added.
+        """
         if chain.chain_type == "B":
             self.VB = chain.id
         elif chain.chain_type == "A" or chain.chain_type == "D":
@@ -340,14 +608,32 @@ class abTCR(TCR):
         self.add(chain)
 
     def get_VB(self):
+        """
+        Retrieve the variable beta chain of the TCR
+
+        Returns:
+            TCRchain: VB chain
+        """
         if hasattr(self, "VB"):
             return self.child_dict[self.VB]
 
     def get_VA(self):
+        """
+        Retrieve the variable alpha chain of the TCR
+
+        Returns:
+            TCRchain: VA chain
+        """
         if hasattr(self, "VA"):
             return self.child_dict[self.VA]
 
     def get_domain_assignment(self):
+        """
+        Retrieve the domain assignment of the TCR as a dict with variable domain type as key and chain ID as value.
+
+        Returns:
+            dict: domain assignment from domain to chain ID, e.g. {"VA": "D", "VB": "E"}
+        """
         try:
             return {"VA": self.VA, "VB": self.VB}
         except AttributeError:
@@ -358,6 +644,12 @@ class abTCR(TCR):
         return None
 
     def is_engineered(self):
+        """
+        Flag for engineered TCRs.
+
+        Returns:
+            bool: Flag for engineered TCRs
+        """
         if self.engineered:
             return True
         else:
@@ -371,6 +663,12 @@ class abTCR(TCR):
             return False
 
     def get_fragments(self):
+        """
+        Retrieve the fragments, ie FW and CDR loops of the TCR as a generator.
+
+        Yields:
+            Fragment: fragment of TCR chain.
+        """
         vb, va = self.get_VB(), self.get_VA()
 
         # If a variable domain exists
@@ -379,6 +677,56 @@ class abTCR(TCR):
                 for frag in var_domain.get_fragments():
                     yield frag
 
+    def standardise_chain_names(self) -> None:
+        """
+        Standardise the TCR, antigen, and MHC chain names to the following convention.
+
+        Convention:
+            - A - MHC chain 1
+            - B - MHC chain 2 (eg B2M)
+            - C - antigen chain
+            - D - TCR alpha chain
+            - E - TCR beta chain
+
+        Note, this mutates the original object.
+
+        Raises:
+            TCRError: if there is more than one antigen or MHC molecules attached to the TCR.
+
+        """
+        self._validate_chain_standardising()
+
+        new_id = []
+        new_child_dict = {}
+
+        if hasattr(self, 'VB'):
+            new_child_dict['E'] = self.child_dict[self.VB]
+            self.VB = 'E'
+            new_id.append('E')
+
+        if hasattr(self, 'VA'):
+            new_child_dict['D'] = self.child_dict[self.VA]
+            self.VA = 'D'
+            new_id.append('D')
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+
+            for chain_id, chain in new_child_dict.items():
+                chain.id = chain_id
+
+        self.child_dict = new_child_dict
+
+        if hasattr(self, 'antigen') and self.antigen:
+            self._standardise_antigen_chain_names()
+
+        if hasattr(self, 'MHC') and self.MHC:
+            self._standardise_mhc_chain_names()
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = ''.join(new_id)
+
 
 class gdTCR(TCR):
 
@@ -455,6 +803,56 @@ class gdTCR(TCR):
                 for frag in var_domain.get_fragments():
                     yield frag
 
+    def standardise_chain_names(self) -> None:
+        """
+        Standardise the TCR, antigen, and MHC chain names to the following convention.
+
+        Convention:
+            - A - MHC chain 1
+            - B - MHC chain 2 (eg B2M)
+            - C - antigen chain
+            - D - TCR delta chain
+            - E - TCR gamma chain
+
+        Note, this mutates the original object.
+
+        Raises:
+            TCRError: if there is more than one antigen or MHC molecules attached to the TCR.
+
+        """
+        self._validate_chain_standardising()
+
+        new_id = []
+        new_child_dict = {}
+
+        if hasattr(self, 'VG'):
+            new_child_dict['E'] = self.child_dict[self.VG]
+            self.VG = 'E'
+            new_id.append('E')
+
+        if hasattr(self, 'VD'):
+            new_child_dict['D'] = self.child_dict[self.VD]
+            self.VD = 'D'
+            new_id.append('D')
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+
+            for chain_id, chain in new_child_dict.items():
+                chain.id = chain_id
+
+        self.child_dict = new_child_dict
+
+        if hasattr(self, 'antigen') and self.antigen:
+            self._standardise_antigen_chain_names()
+
+        if hasattr(self, 'MHC') and self.MHC:
+            self._standardise_mhc_chain_names()
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = ''.join(new_id)
+
 
 class dbTCR(TCR):
     def __init__(self, c1, c2):
@@ -530,3 +928,53 @@ class dbTCR(TCR):
             if var_domain:
                 for frag in var_domain.get_fragments():
                     yield frag
+
+    def standardise_chain_names(self) -> None:
+        """
+        Standardise the TCR, antigen, and MHC chain names to the following convention.
+
+        Convention:
+            - A - MHC chain 1
+            - B - MHC chain 2 (eg B2M)
+            - C - antigen chain
+            - D - TCR delta chain
+            - E - TCR beta chain
+
+        Note, this mutates the original object.
+
+        Raises:
+            TCRError: if there is more than one antigen or MHC molecules attached to the TCR.
+
+        """
+        self._validate_chain_standardising()
+
+        new_id = []
+        new_child_dict = {}
+
+        if hasattr(self, 'VB'):
+            new_child_dict['E'] = self.child_dict[self.VB]
+            self.VB = 'E'
+            new_id.append('E')
+
+        if hasattr(self, 'VD'):
+            new_child_dict['D'] = self.child_dict[self.VD]
+            self.VD = 'D'
+            new_id.append('D')
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+
+            for chain_id, chain in new_child_dict.items():
+                chain.id = chain_id
+
+        self.child_dict = new_child_dict
+
+        if hasattr(self, 'antigen') and self.antigen:
+            self._standardise_antigen_chain_names()
+
+        if hasattr(self, 'MHC') and self.MHC:
+            self._standardise_mhc_chain_names()
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = ''.join(new_id)