stcrpy 1.0.0__py3-none-any.whl

stcrpy/tcr_geometry/TCRCoM.py

@@ -0,0 +1,350 @@
+import os
+import Bio
+from typing import Union
+import warnings
+import numpy as np
+
+from ..tcr_processing.TCRParser import TCRParser
+from ..tcr_processing.TCRIO import TCRIO
+from ..tcr_processing import abTCR, MHCchain
+
+
+# Some of this code is adapted and refactored from https://github.com/EsamTolba/TCR-CoM/
+# Please see the TCRCoM_LICENSE for the license that applies to those code sections.
+
+class TCRCoM:
+    def __init__(self):
+        """Abstract class for calculating TCR centre of mass after aligning TCR:pMHC complex to reference MHC structure."""
+        self.set_reffile()
+
+        tcr_parser = TCRParser()
+        self.ref_model = list(
+            tcr_parser.get_tcr_structure("reference", self.reffile).get_TCRs()
+        )[0]
+        self.set_reference_residues()
+
+    def set_reffile(self):
+        """Super method for setting MHC reference file.
+
+        Raises:
+            NotImplementedError
+        """
+        raise NotImplementedError(
+            "TCRCom cannot be insantiated directly, instantiate its subclass"
+        )
+
+    def set_reference_residues(self):
+        """Set TCR and MHC reference residues."""
+        self.set_tcr_reference()
+        self.set_mhc_reference()
+
+    def set_tcr_reference(self):
+        """Set TCR variable domain residues in reference model as residues numbered 1 to 121 for VA and 1 to 126 for VB."""
+        self.reference_VA_residues = [
+            r
+            for r in self.ref_model.get_VA().get_residues()
+            if r.get_id()[1] >= 1 and r.get_id()[1] <= 121
+        ]
+        self.reference_VB_residues = [
+            r
+            for r in self.ref_model.get_VB().get_residues()
+            if r.get_id()[1] >= 1 and r.get_id()[1] <= 126
+        ]
+
+    def set_mhc_reference(self):
+        """Super method for setting MHC reference residues for superposition. Overwritten by MHC class sepcific methods.
+
+        Raises:
+            NotImplementedError
+        """
+        raise NotImplementedError(
+            "TCRCom cannot be insantiated directly, instantiate its subclass"
+        )
+
+    def get_filtered_TCR_residues(self, tcr: "TCR") -> "tuple[list[Bio.PDB.Residue]]":
+        """Get variable domain residues of query TCR and filter out those without a counterpart in the reference TCR.
+
+        Args:
+            tcr (TCR): TCR structure object
+
+        Returns:
+            tuple[list[Bio.PDB.Residue]]: VA_residues, VB_residues.
+        """
+        tcr_A_residues = [
+            tcr.get_VA()[r.get_id()]
+            for r in self.reference_VA_residues
+            if r.get_id() in tcr.get_VA()
+        ]
+        tcr_B_residues = [
+            tcr.get_VB()[r.get_id()]
+            for r in self.reference_VB_residues
+            if r.get_id() in tcr.get_VB()
+        ]
+        return tcr_A_residues, tcr_B_residues
+
+    def center_of_mass(
+        self,
+        entity: "Union[Bio.PDB.Entity.Entity, list[Bio.PDB.Atom.Atom]]",
+        geometric: bool = False,
+    ) -> np.array:
+        """Calculate the mass weighted or purely geometric centre of mass of an entity or a list of atoms.
+
+        Args:
+            entity (Union[Bio.PDB.Entity.Entity, list[Bio.PDB.Atom.Atom]]): Structure entity whose centre of mass will be calculated
+            geometric (bool, optional): Whether to calculate geometric mean. Defaults to False.
+
+        Raises:
+            ValueError: Checks input type to ensure atoms can be retrieved.
+            ValueError: Unknown atoms.
+
+        Returns:
+            np.array: centre of mass
+        """
+        # Structure, Model, Chain, Residue
+        if isinstance(entity, Bio.PDB.Entity.Entity):
+            atom_list = entity.get_atoms()
+        # List of Atoms
+        elif hasattr(entity, "__iter__") and [x for x in entity if x.level == "A"]:
+            atom_list = entity
+        else:
+            raise ValueError(
+                "Center of Mass can only be calculated from the following objects:\n"
+                "Structure, Model, Chain, Residue, list of Atoms."
+            )
+
+        masses, positions = zip(*[(atom.mass, atom.coord) for atom in atom_list])
+        positions = np.asarray(positions)
+        if "ukn" in set(masses) and not geometric:
+            raise ValueError(
+                "Some Atoms don't have an element assigned.\n"
+                "Try adding them manually or calculate the geometrical center of mass instead"
+            )
+
+        if geometric:
+            return positions.sum(axis=0) / len(atom_list)
+        else:
+            return np.matmul(np.asarray(masses), positions) / len(atom_list)
+
+    def add_com(
+        self,
+        mhc_com: np.array,
+        tcr_com: np.array,
+        VA_com: np.array,
+        VB_com: np.array,
+        tcr: "TCR",
+    ) -> "TCR":
+        """
+        Function to add pseudoatoms at MHC-CoM, TCR-CoM, and XYZ axis to the output PDB file
+
+        Args:
+            mhc_com (np.array): MHC centre of mass
+            tcr_com (np.array): TCR centre of mass
+            VA_com (np.array): Alpha chain centre of mass
+            VB_com (np.array): Beta chain entre of mass
+            tcr (TCR): TCR structure object
+
+        Returns:
+            TCR: Copy of the original TCR strucutre object with added pseudo-atoms.
+        """
+        new_structure = tcr.copy()
+
+        # mhc com
+        mhc_com_chain = "X"
+        new_structure.add(Bio.PDB.Chain.Chain(mhc_com_chain))
+        res_id = (" ", 1, " ")
+        new_residue = Bio.PDB.Residue.Residue(res_id, "MCM", " ")
+        new_atom = Bio.PDB.Atom.Atom("C", mhc_com, 0, 0.0, " ", "C", 1, "C")
+        new_residue.add(new_atom)
+        new_structure.child_dict[mhc_com_chain].add(new_residue)
+
+        # tcr com
+        tcr_com_chain = "Y"
+        new_structure.add(Bio.PDB.Chain.Chain(tcr_com_chain))
+        pseudo_atom_ids = ["TCM", "ACM", "BCM"]
+        tcr_com_list = [tcr_com, VA_com, VB_com]
+        for i in range(3):
+            res_id = (" ", i + 1, " ")
+            new_residue = Bio.PDB.Residue.Residue(res_id, pseudo_atom_ids[i], " ")
+            new_atom = Bio.PDB.Atom.Atom("C", tcr_com_list[i], 0, 0.0, " ", "C", 1, "C")
+            new_residue.add(new_atom)
+            new_structure.child_dict[tcr_com_chain].add(new_residue)
+
+        # X,Y,Z atoms
+        pos = [[50, 0, 0], [0, 50, 0], [0, 0, 50]]
+        resn = ["X", "Y", "Z"]
+        xyz_chain = "Z"
+        new_structure.add(Bio.PDB.Chain.Chain(xyz_chain))
+        for i in [0, 1, 2]:
+            res_id = (" ", i + 1, " ")
+            new_residue = Bio.PDB.Residue.Residue(res_id, resn[i], " ")
+            new_atom = Bio.PDB.Atom.Atom("O", pos[i], 0, 0.0, " ", "O", 1, "O")
+            new_residue.add(new_atom)
+            new_structure.child_dict[xyz_chain].add(new_residue)
+
+        return new_structure
+
+    def calculate_centres_of_mass(
+        self,
+        tcr: abTCR,
+        save_aligned_as: str = None,
+    ) -> tuple[np.array]:
+        """Calculate the TCR and MHC centres of mass of an stcrpy TCR structure object.
+
+        Args:
+            tcr (abTCR): TCR structure object
+            save_aligned_as (str): Path to same alignment to. If None or False alignment is not saved. Defaults to None.
+
+        Raises:
+            NotImplementedError: Alpha Beta TCR compatible only, Gamma Delta TCRs not implemented.
+
+        Returns:
+            tuple[np.array]: tcr_com, mhc_com, tcr_VA_com, tcr_VB_com
+        """
+
+        assert len(tcr.get_MHC()) > 0, "No MHC associated with TCR"
+        if not isinstance(tcr, abTCR):
+            raise NotImplementedError(
+                f"TCR MHC geometry only implemented for abTCR types, not {type(tcr)}"
+            )
+
+        mhc_atoms = [res["CA"] for res in self.get_filtered_MHC_residues(tcr)]
+        ref_mhc_atoms = [
+            res["CA"]
+            for res in self.reference_MHC_residues
+            if (res.parent.chain_type, res.get_id())
+            in [(a.parent.parent.chain_type, a.parent.get_id()) for a in mhc_atoms]
+        ]
+
+        superimposer = Bio.PDB.Superimposer()
+        superimposer.set_atoms(ref_mhc_atoms, mhc_atoms)
+        self.mhc_alignment_transform = (x.astype("f") for x in superimposer.rotran)
+        superimposer.apply(tcr.parent.get_atoms())
+
+        mhc_com = self.center_of_mass(mhc_atoms, geometric=True)
+
+        tcr_VA_residues, tcr_VB_residues = self.get_filtered_TCR_residues(tcr)
+        tcr_VA_atoms = [res["CA"] for res in tcr_VA_residues]
+        tcr_VA_com = self.center_of_mass(tcr_VA_atoms, geometric=True)
+
+        tcr_VB_atoms = [res["CA"] for res in tcr_VB_residues]
+        tcr_VB_com = self.center_of_mass(tcr_VB_atoms, geometric=True)
+
+        tcr_com = self.center_of_mass(tcr_VA_atoms + tcr_VB_atoms, geometric=True)
+
+        if save_aligned_as:
+            if not isinstance(save_aligned_as, str):
+                save_aligned_as = os.path.join(
+                    os.getcwd(), f"{tcr.parent.parent.id}_{tcr.id}.pdb"
+                )
+            aligned_tcr = self.add_com(mhc_com, tcr_com, tcr_VA_com, tcr_VB_com, tcr)
+            io = TCRIO()
+            io.save(aligned_tcr, save_as=save_aligned_as)
+
+        # com_distance = tcr_com - mhc_com
+        # r = np.sqrt(np.sum(np.square(com_distance)))
+        # theta = np.degrees(np.arctan2(com_distance[1], com_distance[0]))
+        # phi = np.degrees(np.arccos(com_distance[2] / r))
+
+        # return r, theta, phi
+
+        return tcr_com, mhc_com, tcr_VA_com, tcr_VB_com
+
+
+class MHCI_TCRCoM(TCRCoM):
+    def __init__(self):
+        """TCRCoM module for MHC Class I complexes."""
+        super().__init__()
+
+    def set_reffile(self):
+        """Sets reference file for MHC class I structures"""
+        self.reffile = os.path.join(
+            os.path.dirname(__file__),
+            "reference_data/dock_reference_1_imgt_numbered.pdb",
+        )
+
+    def set_mhc_reference(self):
+        """Set class I reference MHC residues"""
+        mhc = self.ref_model.get_MHC()[0].get_MH1()
+        self.reference_MHC_residues = [
+            r for r in mhc.get_residues() if r.get_id()[1] >= 1 and r.get_id()[1] <= 179
+        ]
+
+    def get_filtered_MHC_residues(self, tcr: "TCR") -> list[Bio.PDB.Residue]:
+        """Retrieve MHC residues from query TCR and filter out those whose counterpart is not found in reference.
+
+        Args:
+            tcr (TCR): TCR structure object associated with MHC Class I.
+
+        Returns:
+            list[Bio.PDB.Residue]: filtered_MHC_residues
+        """
+        mhc = tcr.get_MHC()[0]
+        if not isinstance(mhc, MHCchain):  # handle single MHC chain case
+            mhc = mhc.get_MH1()
+        filtered_MHC_residues = [
+            mhc[ref_res.get_id()]
+            for ref_res in self.reference_MHC_residues
+            if ref_res.get_id() in mhc
+        ]
+        return filtered_MHC_residues
+
+
+class MHCII_TCRCoM(TCRCoM):
+    def __init__(self):
+        """TCRCoM module for MHC Class II complexes."""
+        super().__init__()
+
+    def set_reffile(self):
+        """Sets reference file for MHC class II structures"""
+        self.reffile = os.path.join(
+            os.path.dirname(__file__),
+            "reference_data/dock_reference_2_imgt_numbered.pdb",
+        )
+
+    def set_mhc_reference(self):
+        """Set class II reference MHC residues"""
+        mhc = self.ref_model.get_MHC()[0]
+        self.reference_MHC_residues = [
+            r
+            for r in mhc.get_GA().get_residues()
+            if r.get_id()[1] >= 1 and r.get_id()[1] <= 88
+        ]
+        self.reference_MHC_residues.extend(
+            [
+                r
+                for r in mhc.get_GB().get_residues()
+                if r.get_id()[1] >= 1 and r.get_id()[1] <= 87
+            ]
+        )
+
+    def get_filtered_MHC_residues(self, tcr: "TCR") -> list[Bio.PDB.Residue]:
+        """Retrieve MHC residues from query TCR and filter out those whose counterpart is not found in reference.
+
+        Args:
+            tcr (TCR): TCR structure object associated with MHC Class II.
+
+        Returns:
+            list[Bio.PDB.Residue]: filtered_MHC_residues
+        """
+        mhc = tcr.get_MHC()[0]
+        if hasattr(mhc, "get_GA"):
+            filtered_MHC_residues = [
+                mhc.get_GA()[ref_res.get_id()]
+                for ref_res in self.reference_MHC_residues
+                if ref_res.parent.chain_type == "GA"
+                and ref_res.get_id() in mhc.get_GA()
+            ]
+        else:
+            warnings.warn(f"No GA chain found for MHC class II: {mhc}")
+        if hasattr(mhc, "get_GA"):
+            filtered_MHC_residues.extend(
+                [
+                    mhc.get_GB()[ref_res.get_id()]
+                    for ref_res in self.reference_MHC_residues
+                    if ref_res.parent.chain_type == "GB"
+                    and ref_res.get_id() in mhc.get_GB()
+                ]
+            )
+        else:
+            warnings.warn(f"No GB chain found for MHC class II: {mhc}")
+        return filtered_MHC_residues

stcrpy/tcr_geometry/TCRCoM_LICENCE

@@ -0,0 +1,168 @@
+This license applies to the modified TCRCoM code located in TCRCoM.py within the stcrpy package. Original license can be found here: https://github.com/EsamTolba/TCR-CoM/blob/master/LICENSE
+
+
+                    GNU LESSER GENERAL PUBLIC LICENSE
+                       Version 3, 29 June 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+
+  This version of the GNU Lesser General Public License incorporates
+the terms and conditions of version 3 of the GNU General Public
+License, supplemented by the additional permissions listed below.
+
+  0. Additional Definitions.
+
+  As used herein, "this License" refers to version 3 of the GNU Lesser
+General Public License, and the "GNU GPL" refers to version 3 of the GNU
+General Public License.
+
+  "The Library" refers to a covered work governed by this License,
+other than an Application or a Combined Work as defined below.
+
+  An "Application" is any work that makes use of an interface provided
+by the Library, but which is not otherwise based on the Library.
+Defining a subclass of a class defined by the Library is deemed a mode
+of using an interface provided by the Library.
+
+  A "Combined Work" is a work produced by combining or linking an
+Application with the Library.  The particular version of the Library
+with which the Combined Work was made is also called the "Linked
+Version".
+
+  The "Minimal Corresponding Source" for a Combined Work means the
+Corresponding Source for the Combined Work, excluding any source code
+for portions of the Combined Work that, considered in isolation, are
+based on the Application, and not on the Linked Version.
+
+  The "Corresponding Application Code" for a Combined Work means the
+object code and/or source code for the Application, including any data
+and utility programs needed for reproducing the Combined Work from the
+Application, but excluding the System Libraries of the Combined Work.
+
+  1. Exception to Section 3 of the GNU GPL.
+
+  You may convey a covered work under sections 3 and 4 of this License
+without being bound by section 3 of the GNU GPL.
+
+  2. Conveying Modified Versions.
+
+  If you modify a copy of the Library, and, in your modifications, a
+facility refers to a function or data to be supplied by an Application
+that uses the facility (other than as an argument passed when the
+facility is invoked), then you may convey a copy of the modified
+version:
+
+   a) under this License, provided that you make a good faith effort to
+   ensure that, in the event an Application does not supply the
+   function or data, the facility still operates, and performs
+   whatever part of its purpose remains meaningful, or
+
+   b) under the GNU GPL, with none of the additional permissions of
+   this License applicable to that copy.
+
+  3. Object Code Incorporating Material from Library Header Files.
+
+  The object code form of an Application may incorporate material from
+a header file that is part of the Library.  You may convey such object
+code under terms of your choice, provided that, if the incorporated
+material is not limited to numerical parameters, data structure
+layouts and accessors, or small macros, inline functions and templates
+(ten or fewer lines in length), you do both of the following:
+
+   a) Give prominent notice with each copy of the object code that the
+   Library is used in it and that the Library and its use are
+   covered by this License.
+
+   b) Accompany the object code with a copy of the GNU GPL and this license
+   document.
+
+  4. Combined Works.
+
+  You may convey a Combined Work under terms of your choice that,
+taken together, effectively do not restrict modification of the
+portions of the Library contained in the Combined Work and reverse
+engineering for debugging such modifications, if you also do each of
+the following:
+
+   a) Give prominent notice with each copy of the Combined Work that
+   the Library is used in it and that the Library and its use are
+   covered by this License.
+
+   b) Accompany the Combined Work with a copy of the GNU GPL and this license
+   document.
+
+   c) For a Combined Work that displays copyright notices during
+   execution, include the copyright notice for the Library among
+   these notices, as well as a reference directing the user to the
+   copies of the GNU GPL and this license document.
+
+   d) Do one of the following:
+
+       0) Convey the Minimal Corresponding Source under the terms of this
+       License, and the Corresponding Application Code in a form
+       suitable for, and under terms that permit, the user to
+       recombine or relink the Application with a modified version of
+       the Linked Version to produce a modified Combined Work, in the
+       manner specified by section 6 of the GNU GPL for conveying
+       Corresponding Source.
+
+       1) Use a suitable shared library mechanism for linking with the
+       Library.  A suitable mechanism is one that (a) uses at run time
+       a copy of the Library already present on the user's computer
+       system, and (b) will operate properly with a modified version
+       of the Library that is interface-compatible with the Linked
+       Version.
+
+   e) Provide Installation Information, but only if you would otherwise
+   be required to provide such information under section 6 of the
+   GNU GPL, and only to the extent that such information is
+   necessary to install and execute a modified version of the
+   Combined Work produced by recombining or relinking the
+   Application with a modified version of the Linked Version. (If
+   you use option 4d0, the Installation Information must accompany
+   the Minimal Corresponding Source and Corresponding Application
+   Code. If you use option 4d1, you must provide the Installation
+   Information in the manner specified by section 6 of the GNU GPL
+   for conveying Corresponding Source.)
+
+  5. Combined Libraries.
+
+  You may place library facilities that are a work based on the
+Library side by side in a single library together with other library
+facilities that are not Applications and are not covered by this
+License, and convey such a combined library under terms of your
+choice, if you do both of the following:
+
+   a) Accompany the combined library with a copy of the same work based
+   on the Library, uncombined with any other library facilities,
+   conveyed under the terms of this License.
+
+   b) Give prominent notice with the combined library that part of it
+   is a work based on the Library, and explaining where to find the
+   accompanying uncombined form of the same work.
+
+  6. Revised Versions of the GNU Lesser General Public License.
+
+  The Free Software Foundation may publish revised and/or new versions
+of the GNU Lesser General Public License from time to time. Such new
+versions will be similar in spirit to the present version, but may
+differ in detail to address new problems or concerns.
+
+  Each version is given a distinguishing version number. If the
+Library as you received it specifies that a certain numbered version
+of the GNU Lesser General Public License "or any later version"
+applies to it, you have the option of following the terms and
+conditions either of that published version or of any later version
+published by the Free Software Foundation. If the Library as you
+received it does not specify a version number of the GNU Lesser
+General Public License, you may choose any version of the GNU Lesser
+General Public License ever published by the Free Software Foundation.
+
+  If the Library as you received it specifies that a proxy can decide
+whether future versions of the GNU Lesser General Public License shall
+apply, that proxy's public statement of acceptance of any version is
+permanent authorization for you to choose that version for the
+Library.