stcrpy 1.0.0__py3-none-any.whl

stcrpy/tcr_metrics/tcr_interface_rmsd.py

@@ -0,0 +1,237 @@
+import warnings
+import Bio
+from Bio.PDB.Superimposer import Superimposer
+import numpy as np
+
+class InterfaceRMSD:
+    def __init__(self):
+        return
+
+    def get_interface_rmsd(self, dock: "abTCR", reference: "abTCR") -> float:
+        """
+        Calculates the root-mean-square deviation (RMSD) between the interface residues of a
+        docked TCR structure and a reference TCR structure.
+
+        Args:
+            dock (abTCR): The docked TCR structure.
+            reference (abTCR): The reference TCR structure.
+
+        Returns:
+            float: The RMSD value between the interface residues of the docked TCR structure and
+                    the reference TCR structure.
+        """
+
+        # check all residues in dock can be mapped to reference
+        dock_to_ref_chain_mapping = self.check_residue_mapping(dock, reference)
+
+        # get the interface of the reference
+        tcr_interface, antigen_interface = self.get_interface_residues(reference)
+
+        # align the dock by MHC
+        self.align_by_mhc(dock, reference, dock_to_ref_chain_mapping)
+
+        # get the docked residues found in the reference interface
+        ref_to_dock_chain_mapping = {v: k for k, v in dock_to_ref_chain_mapping.items()}
+
+        try:
+            docked_tcr_interface = [
+                dock[ref_to_dock_chain_mapping[res.parent.id]][res.id]
+                for res in tcr_interface
+            ]
+            docked_antigen_interface = [
+                dock.parent[ref_to_dock_chain_mapping[res.parent.id]][res.id]
+                for res in antigen_interface
+            ]
+        except KeyError as e:
+            warnings.warn(
+                f"""Key error {str(e)} matching dock chains to reference chains for dock: {
+                    ' '.join([str(i) for i in dock.full_id])
+                } to reference {
+                    ' '.join([str(i) for i in reference.full_id])
+                }. Interface RMSD could not be calculated."""
+            )
+            return None
+
+        # extract coordinates from interfaces
+        reference_coordinates = np.asarray(
+            [
+                atom.get_coord()
+                for res in (tcr_interface + antigen_interface)
+                for atom in res
+                if atom.element in ["N", "O", "C", "S"]
+                and (
+                    atom.id
+                    in [
+                        a.id
+                        for a in dock.parent[ref_to_dock_chain_mapping[res.parent.id]][
+                            res.id
+                        ].get_atoms()
+                    ]
+                )
+            ]
+        )
+
+        docked_coordinates = np.asarray(
+            [
+                atom.get_coord()
+                for res in (docked_tcr_interface + docked_antigen_interface)
+                for atom in res
+                if atom.element in ["N", "O", "C", "S"]
+                and (
+                    atom.id
+                    in [
+                        a.id
+                        for a in reference.parent[
+                            dock_to_ref_chain_mapping[res.parent.id]
+                        ][res.id].get_atoms()
+                    ]
+                )
+            ]
+        )
+
+        # calculate rmsd
+        rmsd = np.sqrt(
+            ((docked_coordinates - reference_coordinates) ** 2).sum()
+            / len(docked_coordinates)
+        )
+        return rmsd
+
+    def check_residue_mapping(self, dock: "abTCR", reference: "abTCR") -> dict:
+        chain_mapping = {}
+
+        for i, tcr_chain in enumerate(dock.get_chains()):
+            try:
+                for j, res in enumerate(tcr_chain.get_residues()):
+                    if j > 3 and j < len(tcr_chain) - 2:
+                        assert (
+                            res.resname
+                            == list(reference.get_chains())[i][res.id].resname
+                        ), f"""
+    TCR chain mapping {tcr_chain.id} -> {list(reference.get_chains())[i]} failed. Trying chain swap."""
+                        chain_mapping[tcr_chain.id] = list(reference.get_chains())[i].id
+            except (AssertionError, KeyError):
+                for j, res in enumerate(tcr_chain.get_residues()):
+                    if (
+                        j > 3 and j < len(tcr_chain) - 2
+                    ):  # avoids small mismatches at beginnings and ends of sequences
+                        assert (
+                            res.resname
+                            == list(reference.get_chains())[1 - i][res.id].resname
+                        ), f"""
+    TCR chain mapping {tcr_chain.id} -> {list(reference.get_chains())[1-i]} failed. Residue mapping failed"""
+                        chain_mapping[tcr_chain.id] = list(reference.get_chains())[
+                            1 - i
+                        ].id
+
+        for i, antigen_chain in enumerate(dock.get_antigen()):
+            try:
+                for res in antigen_chain.get_residues():
+                    assert (
+                        res.resname == reference.get_antigen()[i][res.id].resname
+                    ), f"""
+    antigen chain mapping {antigen_chain.id} -> {reference.get_antigen()[i]} failed. Trying chain swap."""
+                    chain_mapping[antigen_chain.id] = reference.get_antigen()[i].id
+            except (AssertionError, KeyError):
+                for res in antigen_chain.get_residues():
+                    assert (
+                        res.resname == reference.get_antigen()[1 - i][res.id].resname
+                    ), f"""
+    antigen chain mapping {antigen_chain.id} -> {reference.get_antigen()[1-i]} failed. Residue mapping failed"""
+                    chain_mapping[antigen_chain.id] = reference.get_antigen()[1 - i].id
+        return chain_mapping
+
+    def get_interface_residues(
+        self, tcr: "abTCR", angstrom_cutoff: float = 8.0
+    ) -> list:
+        """
+        Retrieves the interface residues between a TCR and its antigen based on a distance cutoff.
+
+        Args:
+            tcr (abTCR): The TCR object.
+            angstrom_cutoff (float, optional): The distance cutoff in angstroms. Defaults to 8.0.
+
+        Returns:
+            tuple: A tuple containing two lists: the interface residues of the TCR and the
+                    interface residues of the antigen.
+        """
+        tcr_c_alphas = [atom for atom in tcr.get_atoms() if atom.id == "CA"]
+        antigen_c_alphas = [
+            atom
+            for chain in tcr.get_antigen()
+            for atom in chain.get_atoms()
+            if atom.id == "CA"
+        ]
+
+        tcr_c_coords = np.asarray([[x.get_coord()] for x in tcr_c_alphas])
+        antigen_c_coords = np.asarray([[x.get_coord() for x in antigen_c_alphas]])
+
+        tcr_c_coords = np.broadcast_to(
+            tcr_c_coords, (tcr_c_coords.shape[0], antigen_c_coords.shape[1], 3)
+        )
+
+        antigen_c_coords = np.broadcast_to(
+            antigen_c_coords, (tcr_c_coords.shape[0], antigen_c_coords.shape[1], 3)
+        )
+
+        pairwise_distances = np.sqrt(((tcr_c_coords - antigen_c_coords) ** 2).sum(-1))
+        contacts = np.argwhere(pairwise_distances <= angstrom_cutoff)
+        tcr_interface_idx = set(contacts[:, 0])
+        antigen_interface_idx = set(contacts[:, 1])
+        tcr_interface = [tcr_c_alphas[idx].parent for idx in tcr_interface_idx]
+        antigen_interface = [
+            antigen_c_alphas[idx].parent for idx in antigen_interface_idx
+        ]
+
+        return tcr_interface, antigen_interface
+
+    def align_by_mhc(
+        self, dock: "abTCR", reference: "abTCR", chain_mapping: dict
+    ) -> None:
+        """
+        Aligns the docked TCR structure to the reference TCR structure by aligning the MHC.
+
+        Args:
+            dock (abTCR): The docked TCR structure.
+            reference (abTCR): The reference TCR structure.
+            chain_mapping (dict): A dictionary mapping the chain IDs of the docked TCR structure to the chain IDs of the
+                                    reference TCR structure.
+
+        Returns:
+            None
+        """
+        mhc_chain = dock.get_MHC()
+        assert len(mhc_chain) >= 1, ValueError("No MHC chains found")
+        if hasattr(mhc_chain[0], "get_MH1"):
+            mhc_chain = mhc_chain[
+                0
+            ].get_MH1()  # This will only work for class I MHC, ie. single chain helices.
+            reference_mhc_chain = reference.get_MHC()[0].get_MH1()
+        else:
+            # For Class II MHC try creating new entity with GA and GB chains
+            class_II_mhc_chain = Bio.PDB.Entity.Entity()
+            class_II_mhc_chain.add(mhc_chain[0].get_GA())
+            class_II_mhc_chain.add(mhc_chain[0].get_GB())
+            mhc_chain = class_II_mhc_chain
+            reference_mhc_chain = Bio.PDB.Entity.Entity()
+            reference_mhc_chain.add(reference.get_MHC()[0].get_GA())
+            reference_mhc_chain.add(reference.get_MHC()[0].get_GB())
+
+        mutual_residue_ids = set(
+            [r.id for r in reference_mhc_chain.get_residues()]
+        ).intersection(set([r.id for r in mhc_chain.get_residues()]))
+        reference_atoms = [
+            a
+            for res in mutual_residue_ids
+            for a in reference_mhc_chain[res].get_atoms()
+            if a.id in ["N", "C", "O", "CA"]
+        ]
+        docked_atoms = [
+            a
+            for res in mutual_residue_ids
+            for a in mhc_chain[res].get_atoms()
+            if a.id in ["N", "C", "O", "CA"]
+        ]
+
+        superimposer = Superimposer()
+        superimposer.set_atoms(reference_atoms, docked_atoms)
+        superimposer.apply(dock.parent.get_atoms())

stcrpy/tcr_metrics/tcr_rmsd.py

@@ -0,0 +1,179 @@
+import warnings
+import numpy as np
+import pandas as pd
+
+from Bio.PDB.Superimposer import Superimposer
+
+from . import constants
+
+
+class RMSD:
+    def __init__(self):
+        return
+
+    @staticmethod
+    def _retrieve_chain(tcr, chain_type):
+        assert chain_type in ["A", "B", "G", "D"], ValueError(
+            "TCR chain type not recognised"
+        )
+        try:
+            return tcr[tcr.get_domain_assignment()[f"V{chain_type}"]]
+        except KeyError:
+            # map chain type A to G and B to D
+            chain_type = {"A": "G", "B": "D"}[chain_type]
+            return tcr[tcr.get_domain_assignment()[f"V{chain_type}"]]
+
+    @staticmethod
+    def _rmsd(x1, x2):
+        assert x1.shape == x2.shape
+        assert x1.shape[-1] == 3
+        return np.sqrt(np.mean(3 * (x1 - x2) ** 2))
+
+    def calculate_rmsd(self, tcr_to_align, tcr_ref, save_alignment=False):
+        rmsds = {}
+        for chain_type in ["A", "B"]:
+            chain_to_align = self._retrieve_chain(tcr_to_align, chain_type)
+            ref_chain = self._retrieve_chain(tcr_ref, chain_type)
+
+            ref_residue_numbering = [
+                x.id
+                for x in ref_chain.get_residues()
+                if all([not a.is_disordered() for a in x.get_atoms()])
+            ]
+            residue_numbering_intersection = [
+                x.id
+                for x in chain_to_align.get_residues()
+                if x.id in ref_residue_numbering
+                and all([not a.is_disordered() for a in x.get_atoms()])
+            ]
+
+            # Get residues to align
+            ref_residues = [
+                ref_chain[x]
+                for x in residue_numbering_intersection
+                if (x in chain_to_align) and (x in ref_chain)
+            ]
+            to_align_residues = [
+                chain_to_align[x]
+                for x in residue_numbering_intersection
+                if (x in chain_to_align) and (x in ref_chain)
+            ]
+
+            # Get backbone atoms to align
+            fixed = []
+            moved = []
+            for i in range(len(to_align_residues)):
+                fixed += [
+                    ref_residues[i][atom]
+                    for atom in constants.ATOM_TYPES[:4]
+                    if (atom in to_align_residues[i]) and (atom in ref_residues[i])
+                ]
+                moved += [
+                    to_align_residues[i][atom]
+                    for atom in constants.ATOM_TYPES[:4]
+                    if (atom in to_align_residues[i]) and (atom in ref_residues[i])
+                ]
+
+            # Calculate superimposer and align
+            imposer = Superimposer()
+            imposer.set_atoms(fixed, moved)
+            imposer.apply(tcr_to_align.get_atoms())
+
+            rmsds[chain_type] = (
+                imposer.rms
+            )  # whole chain RMSD after alignment calculated across all atoms
+
+            if save_alignment:
+                tcr_ref.save(
+                    save_as=f"{tcr_ref.parent.parent.id}_RMSD_reference_alignment_{chain_type}.pdb",
+                    tcr_only=True,
+                )
+                tcr_to_align.save(
+                    save_as=f"{tcr_to_align.parent.parent.id}_RMSD_aligned_to_{tcr_ref.id}_{chain_type}.pdb",
+                    tcr_only=True,
+                )
+
+            # calculate CDR RMSD
+            for CDR_loop_nr, (ref_CDR, aligned_CDR) in enumerate(
+                zip(ref_chain.get_CDRs(), chain_to_align.get_CDRs())
+            ):
+                ref_CDR_atom_coords = np.asarray(
+                    [
+                        r[a].get_coord()
+                        for r in ref_CDR
+                        for a in constants.ATOM_TYPES[:4]
+                        if r.id in residue_numbering_intersection
+                        and a in aligned_CDR[r.id]
+                    ]
+                )
+                aligned_CDR_atom_coords = np.asarray(
+                    [
+                        r[a].get_coord()
+                        for r in aligned_CDR
+                        for a in constants.ATOM_TYPES[:4]
+                        if r.id in residue_numbering_intersection and a in ref_CDR[r.id]
+                    ]
+                )
+
+                rmsds[f"CDR{chain_type}{CDR_loop_nr + 1}"] = self._rmsd(
+                    ref_CDR_atom_coords, aligned_CDR_atom_coords
+                )
+
+            # calculate framework RMSD
+            ref_framework_residues = {
+                r.id: r
+                for fw in ref_chain.get_frameworks()
+                for r in fw.get_residues()
+                if r.id in residue_numbering_intersection
+            }
+            aligned_framework_residues = {
+                r.id: r
+                for fw in chain_to_align.get_frameworks()
+                for r in fw.get_residues()
+                if r.id in residue_numbering_intersection
+            }
+
+            ref_framework_atom_coords = np.asarray(
+                [
+                    r[a].get_coord()
+                    for r_id, r in ref_framework_residues.items()
+                    for a in constants.ATOM_TYPES[:4]
+                    if r_id in residue_numbering_intersection
+                    and a in aligned_framework_residues[r_id]
+                ]
+            )
+            aligned_framework_atom_coords = np.asarray(
+                [
+                    r[a].get_coord()
+                    for r_id, r in aligned_framework_residues.items()
+                    for a in constants.ATOM_TYPES[:4]
+                    if r_id in residue_numbering_intersection
+                    and a in ref_framework_residues[r_id]
+                ]
+            )
+
+            rmsds[f"FW{chain_type}"] = self._rmsd(
+                ref_framework_atom_coords, aligned_framework_atom_coords
+            )
+        return rmsds
+
+    def rmsd_from_files(self, pred_and_target_files: list) -> pd.DataFrame:
+        """Calculates the RMSD between TCR structures from a list of files.
+
+        Args:
+            pred_and_target_files (list of tuples): List of tuples, where each tuple contains
+                the path to the predicticted PDB at index 0 and the path to the target PDB at index 1.
+
+        Returns:
+            pandas.Dataframe: Pandas dataframe indexed by the TCR ID of the predicted structure, with columns
+                containing the RMSD of the whole alpha and beta chain, and all framework and CDR regions.
+        """
+        from ..tcr_methods.tcr_methods import load_TCRs
+
+        all_rmsds = {}
+        for pred_tcr_file, target_tcr_file in pred_and_target_files:
+            pred_tcr, target_tcr = load_TCRs([pred_tcr_file, target_tcr_file])
+            all_rmsds[pred_tcr.parent.parent.id] = self.calculate_rmsd(
+                pred_tcr, target_tcr
+            )
+        return pd.DataFrame(all_rmsds).transpose()

stcrpy/tcr_ml/geometry_predictor.py

@@ -0,0 +1,3 @@
+import torch_geometric
+
+# from torch.

stcrpy/tcr_processing/AGchain.py

@@ -0,0 +1,89 @@
+"""
+Created on 10 May 2017
+@author: leem
+
+Based on the AGchain class from ABDB.
+
+"""
+
+import sys
+from Bio.PDB.Chain import Chain
+from .Chemical_components import get_res_type
+
+
+class AGchain(Chain):
+    """
+    Non-TCR and non-MHC (peptide) chains are described using this class.
+    """
+
+    def __init__(self, identifier):
+        Chain.__init__(self, identifier)
+        self.level = "C"
+        self.type = ""
+        self.engineered = False
+
+    def set_type(self):
+        """
+        Use the type check to check the residue name from the chemical component dictionary
+        For ease of use I have binned these into four types
+        peptide
+        nucleic-acid
+        saccharide (carbohydrate)
+        non-polymer
+        """
+        # Most structures are going to be proteins.
+        # Check the composition of the chain.
+        composition = {
+            "peptide": 0,
+            "nucleic-acid": 0,
+            "non-polymer": 0,
+            "saccharide": 0,
+            None: 0,
+        }
+
+        for r in self.child_list:
+            composition[get_res_type(r)] += 1
+
+        if (
+            composition["nucleic-acid"]
+            or composition["peptide"]
+            or composition["saccharide"]
+        ):
+            composition["non-polymer"] = 0
+            composition[None] = 0
+
+        chain_comp_type = max(composition, key=lambda x: composition[x])
+
+        if chain_comp_type == "peptide":
+            if composition["peptide"] < 30:  # peptide
+                self.type = "peptide"
+            else:
+                self.type = "protein"
+        elif chain_comp_type == "nucleic-acid":
+            self.type = "nucleic-acid"
+        elif chain_comp_type == "saccharide":
+            self.type = "carbohydrate"
+        elif chain_comp_type == "non-polymer":
+            self.type = "non-polymer"
+        elif chain_comp_type is None:
+            print(
+                "Warning: Unknown antigen type for chain %s" % self.id, file=sys.stderr
+            )
+            self.type = "unknown"
+        else:
+            print(
+                "Warning: Unknown antigen type for chain %s" % self.id, file=sys.stderr
+            )
+            self.type = "unknown"
+
+    def get_type(self):
+        return self.type
+
+    def set_engineered(self, engineered):
+        if engineered:
+            self.engineered = True
+        else:
+            self.engineered = False
+
+    def is_engineered(self):
+        return self.engineered