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

stcrpy/tcr_processing/MHC.py

@@ -7,7 +7,13 @@ The MHC class. This is similar to the Fab class.
 
 """
 
+import warnings
+
+from Bio import BiopythonWarning
+
 from .Entity import Entity
+from .MHCchain import MHCchain
+from .utils.region_definitions import IMGT_MH1_ABD, IMGT_MH2_ABD
 
 
 class MHC(Entity):
@@ -38,6 +44,23 @@ class MHC(Entity):
     def _add_tcr(self, tcr=None):
         self.tcr.append(tcr)
 
+    def copy(self, copy_siblings = True ):
+        """
+        Return a copy of the MHC object. This returns a shallow copy of the MHC object.
+        If the copy_siblings flag is set to True, the antigen and TCR objects will also be copied. Warning - if the copy_siblings flag is set to False, the antigen and TCR objects will not be copied, and the reference will still point to the same TCR and antigen objects as the original.
+
+        copy_siblings: Whether to copy sibling entities (ie. TCR and Antigen objects). Default True. 
+
+        """
+        shallow = super().copy()
+        if copy_siblings:
+            shallow.antigen = [a.copy(copy_siblings=False) for a in self.get_antigen()]
+            shallow.tcr = [t.copy(copy_siblings=False) for t in self.get_TCR()]
+            for t in shallow.tcr:
+                t.MHC = [shallow if m.id == shallow.id else m.copy(copy_siblings=False) for m in t.get_MHC()]
+                t.antigen = [ag.copy(copy_siblings=False) 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 t.antigen]
+        return shallow
+
     def get_TCR(self):
         return self.tcr
 
@@ -78,6 +101,14 @@ class MHC(Entity):
     def get_allele_assignments(self):
         return {c.id: c.get_allele_assignments() for c in self.get_chains()}
 
+    def crop(self, *args, **kwargs):
+        """Raises NotImplementedError."""
+        raise NotImplementedError()
+
+    def standardise_chain_names():
+        """Raises NotImplementedError."""
+        raise NotImplementedError()
+
 
 class MH1(MHC):
     """
@@ -89,15 +120,15 @@ class MH1(MHC):
         if self.MHC_type == "MH1":
             return "<%s %s%s GA1/GA2=%s; B2M=%s>" % (
                 self.MHC_type,
-                self.MH1,
-                self.B2M,
+                self.MH1 if self.MH1 else '',
+                self.B2M if self.B2M else '',
                 self.MH1,
                 self.B2M,
             )
         else:
             return "<GA1/GA2 %s%s GA1=%s; GA2=%s>" % (
-                self.GA1,
-                self.GA2,
+                self.GA1 if self.GA1 else '',
+                self.GA2 if self.GA1 else '',
                 self.GA1,
                 self.GA2,
             )
@@ -152,6 +183,60 @@ class MH1(MHC):
         if hasattr(self, "B2M"):
             return self.child_dict[self.B2M]
 
+    def crop(self, *, remove_het_atoms: bool = True) -> None:
+        """Crop to antigen binding domain.
+
+        This method mutates the MH1 object.
+
+        Args:
+            remove_het_atoms: remove het atoms from structure as well
+
+        """
+        alpha_chain = self.get_alpha()
+        new_alpha_chain = MHCchain(alpha_chain.id)
+
+        for residue in alpha_chain:
+            if residue.id[1] in IMGT_MH1_ABD or (not remove_het_atoms and residue.id[0] != ' '):
+                new_alpha_chain.add(residue.copy())
+
+        new_alpha_chain.analyse(alpha_chain.chain_type)
+
+        del self[alpha_chain.id]
+        self.add(new_alpha_chain)
+
+        if hasattr(self, 'B2M'):
+            del self[self.B2M]
+            self.B2M = None
+
+    def standardise_chain_names(self) -> None:
+        """Standardise MHC chain name to A and B2M chain name to B."""
+        new_id = []
+        new_child_dict = {}
+
+        for MH1_domain in set(['MH1', 'GA1', 'GA2']):
+            if hasattr(self, MH1_domain):
+                new_child_dict['A'] = self.child_dict[getattr(self, MH1_domain)]
+                setattr(self, MH1_domain, 'A')
+                new_id.append('A')
+                break
+
+        if hasattr(self, 'B2M'):
+            new_child_dict['B'] = self.child_dict[self.B2M]
+            self.B2M = 'B'
+            new_id.append('B')
+
+        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
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = ''.join(new_id)
+
 
 class MH2(MHC):
     """
@@ -194,6 +279,60 @@ class MH2(MHC):
         if hasattr(self, "GB"):
             return self.child_dict[self.GB]
 
+    def crop(self, *, remove_het_atoms: bool = True) -> None:
+        """Crop to antigen binding domain.
+
+        This method mutates the MH2 object.
+
+        Args:
+            remove_het_atoms: remove het atoms from structure as well
+
+        """
+        new_child_dict = {}
+
+        for chain in self:
+            new_chain = MHCchain(chain.id)
+
+            for residue in chain:
+                if residue.id[1] in IMGT_MH2_ABD or (not remove_het_atoms and residue.id[0] != ' '):
+                    new_chain.add(residue.copy())
+
+            new_chain.analyse(chain.chain_type)
+            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)
+
+    def standardise_chain_names(self) -> None:
+        """Standardise MHC chain 1 name to A and MHC chain 2 name to B."""
+        new_id = []
+        new_child_dict = {}
+
+        if hasattr(self, 'GA'):
+            new_child_dict['A'] = self.child_dict[self.GA]
+            self.GA = 'A'
+            new_id.append('A')
+
+        if hasattr(self, 'GB'):
+            new_child_dict['B'] = self.child_dict[self.GB]
+            self.GB = 'B'
+            new_id.append('B')
+
+        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
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = ''.join(new_id)
+
 
 class CD1(MHC):
     """
@@ -253,6 +392,60 @@ class CD1(MHC):
         if hasattr(self, "B2M"):
             return self.child_dict[self.B2M]
 
+    def crop(self, *, remove_het_atoms: bool = True) -> None:
+        """Crop to antigen binding domain.
+
+        This method mutates the CD1 object.
+
+        Args:
+            remove_het_atoms: remove het atoms from structure as well
+
+        """
+        if hasattr(self, 'CD1'):
+            alpha_chain = self.child_dict[self.CD1]
+
+            new_alpha_chain = MHCchain(alpha_chain.id)
+
+            for residue in alpha_chain:
+                if residue.id[1] in IMGT_MH1_ABD or (not remove_het_atoms and residue.id[0] != ' '):
+                    new_alpha_chain.add(residue.copy())
+
+            new_alpha_chain.analyse(alpha_chain.chain_type)
+
+            del self[alpha_chain.id]
+            self.add(new_alpha_chain)
+
+        if hasattr(self, 'B2M'):
+            del self[self.B2M]
+            self.B2M = None
+
+    def standardise_chain_names(self) -> None:
+        """Standardise CD1 chain name to A and B2M chain name to B."""
+        new_id = []
+        new_child_dict = {}
+
+        if hasattr(self, 'CD1'):
+            new_child_dict['A'] = self.child_dict[self.CD1]
+            self.CD1 = 'A'
+            new_id.append('A')
+
+        if hasattr(self, 'B2M'):
+            new_child_dict['B'] = self.child_dict[self.B2M]
+            self.B2M = 'B'
+            new_id.append('B')
+
+        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
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = ''.join(new_id)
+
 
 class MR1(MHC):
     """
@@ -309,6 +502,60 @@ class MR1(MHC):
         if hasattr(self, "B2M"):
             return self.child_dict[self.B2M]
 
+    def crop(self, *, remove_het_atoms: bool = True) -> None:
+        """Crop to antigen binding domain.
+
+        This method mutates the MR1 object.
+
+        Args:
+            remove_het_atoms: remove het atoms from structure as well
+
+        """
+        if hasattr(self, 'MR1'):
+            alpha_chain = self.child_dict[self.CD1]
+
+            new_alpha_chain = MHCchain(alpha_chain.id)
+
+            for residue in alpha_chain:
+                if residue.id[1] in IMGT_MH1_ABD or (not remove_het_atoms and residue.id[0] != ' '):
+                    new_alpha_chain.add(residue.copy())
+
+            new_alpha_chain.analyse(alpha_chain.chain_type)
+
+            del self[alpha_chain.id]
+            self.add(new_alpha_chain)
+
+        if hasattr(self, 'B2M'):
+            del self[self.B2M]
+            self.B2M = None
+
+    def standardise_chain_names(self) -> None:
+        """Standardise MR1 chain name to A and B2M chain name to B."""
+        new_id = []
+        new_child_dict = {}
+
+        if hasattr(self, 'MR1'):
+            new_child_dict['A'] = self.child_dict[self.MR1]
+            self.MR1 = 'A'
+            new_id.append('A')
+
+        if hasattr(self, 'B2M'):
+            new_child_dict['B'] = self.child_dict[self.B2M]
+            self.B2M = 'B'
+            new_id.append('B')
+
+        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
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = ''.join(new_id)
+
 
 class scMH1(MHC):
     """
@@ -385,6 +632,48 @@ class scMH1(MHC):
     def get_B2M(self):
         return None
 
+    def crop(self, *, remove_het_atoms: bool = True) -> None:
+        """Crop to antigen binding domain.
+
+        This method mutates the scMH1 object.
+
+        Args:
+            remove_het_atoms: remove het atoms from structure as well
+
+        """
+        alpha_chain = self.get_alpha()
+        new_alpha_chain = MHCchain(alpha_chain.id)
+
+        for residue in alpha_chain:
+            if residue.id[1] in IMGT_MH1_ABD or (not remove_het_atoms and residue.id[0] != ' '):
+                new_alpha_chain.add(residue.copy())
+
+        new_alpha_chain.analyse(alpha_chain.chain_type)
+
+        del self[alpha_chain.id]
+        self.add(new_alpha_chain)
+
+    def standardise_chain_names(self) -> None:
+        """Standardise MHC chain name to A."""
+        new_child_dict = {}
+        for MH1_domain in set(['MH1', 'GA1', 'GA2']):
+            if hasattr(self, MH1_domain):
+                new_child_dict['A'] = self.child_dict[getattr(self, MH1_domain)]
+                setattr(self, MH1_domain, 'A')
+                break
+
+        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
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = 'A'
+
 
 class scCD1(MHC):
     """
@@ -448,6 +737,48 @@ class scCD1(MHC):
     def get_B2M(self):
         return None
 
+    def crop(self, *, remove_het_atoms: bool = True) -> None:
+        """Crop to antigen binding domain.
+
+        This method mutates the CD1 object.
+
+        Args:
+            remove_het_atoms: remove het atoms from structure as well
+
+        """
+        if hasattr(self, 'CD1'):
+            alpha_chain = self.child_dict[self.CD1]
+
+            new_alpha_chain = MHCchain(alpha_chain.id)
+
+            for residue in alpha_chain:
+                if residue.id[1] in IMGT_MH1_ABD or (not remove_het_atoms and residue.id[0] != ' '):
+                    new_alpha_chain.add(residue.copy())
+
+            new_alpha_chain.analyse(alpha_chain.chain_type)
+
+            del self[alpha_chain.id]
+            self.add(new_alpha_chain)
+
+    def standardise_chain_names(self) -> None:
+        """Standardise CD1 chain name to A."""
+        new_child_dict = {}
+        if hasattr(self, 'CD1'):
+            new_child_dict['A'] = self.child_dict[self.CD1]
+            self.CD1 = 'A'
+
+        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
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = 'A'
+
 
 class scMH2(MHC):
     """
@@ -514,3 +845,57 @@ class scMH2(MHC):
     def get_GB(self):
         if hasattr(self, "GB"):
             return self.child_dict[self.GB]
+
+    def crop(self, *, remove_het_atoms: bool = True) -> None:
+        """Crop to antigen binding domain.
+
+        This method mutates the scMH2 object.
+
+        Args:
+            remove_het_atoms: remove het atoms from structure as well
+
+        """
+        new_child_dict = {}
+
+        for chain in self:
+            new_chain = MHCchain(chain.id)
+
+            for residue in chain:
+                if residue.id[1] in IMGT_MH2_ABD or (not remove_het_atoms and residue.id[0] != ' '):
+                    new_chain.add(residue.copy())
+
+            new_chain.analyse(chain.chain_type)
+            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)
+
+    def standardise_chain_names(self) -> None:
+        """Standardise MHC chain 1 name to A or MHC chain 2 name to B."""
+        new_id = []
+        new_child_dict = {}
+
+        if hasattr(self, 'GA'):
+            new_child_dict['A'] = self.child_dict[self.GA]
+            self.GA = 'A'
+            new_id.append('A')
+
+        if hasattr(self, 'GB'):
+            new_child_dict['B'] = self.child_dict[self.GB]
+            self.GB = 'B'
+            new_id.append('B')
+
+        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
+
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', BiopythonWarning)
+            self.id = ''.join(new_id)