stcrpy 1.0.0__py3-none-any.whl

stcrpy/tcr_interactions/utils.py

@@ -0,0 +1,170 @@
+import warnings
+
+try:
+    from plip.structure.preparation import PDBComplex
+except ModuleNotFoundError:
+    warnings.warn(
+        """\n\nPLIP package not found. \nProfiling interactions will not be possible \nTo enable interaction profiling, install PLIP with:
+        \npip install plip --no-deps\n\n"""
+    )
+from rdkit import Chem
+
+
+def return_interactions(
+    protein_file=None, ligand_file=None, complex_file=None, pymol_visualization=False
+):
+    with open(protein_file, "r") as f:
+        protein = f.read()
+    protein = [line for line in protein.split("\n") if line.startswith("ATOM")]
+    ligand = Chem.MolFromMolFile(ligand_file)
+    ligand_pdb_block = Chem.MolToPDBBlock(ligand)
+    complex_pdb_block = "\n".join(protein) + "\n" + ligand_pdb_block
+    # return complex_pdb_block, ligand_pdb_block, protein
+    my_mol = PDBComplex()
+    my_mol.load_pdb(complex_pdb_block, as_string=True)
+    my_mol.analyze()
+    return my_mol
+
+
+class Interaction:
+
+    def __init__(
+        self,
+        type,
+        protein_atom,
+        protein_chain,
+        protein_residue,
+        protein_number,
+        ligand_atom,
+        distance,
+        angle,
+        plip_id,
+    ) -> None:
+        self.type = type
+        self.protein_atom = protein_atom
+        self.protein_chain = protein_chain
+        self.protein_residue = protein_residue
+        self.protein_number = protein_number
+        self.ligand_atom = ligand_atom
+        self.distance = distance
+        self.angle = angle
+        self.plip_id = plip_id
+
+    def to_tuple(self):
+        return (
+            self.type,
+            self.protein_atom,
+            self.protein_chain,
+            self.protein_residue,
+            self.protein_number,
+            self.ligand_atom,
+            self.distance,
+            self.angle,
+            self.plip_id,
+        )
+
+
+def parse_interaction(interaction) -> Interaction:
+    if "saltbridge" in str(type(interaction)):
+        return Interaction("saltbridge", *process_saltbridge(interaction))
+    elif "hydroph" in str(type(interaction)):
+        return Interaction("hydrophobic", *process_hydrophobic(interaction))
+    elif "hbond" in str(type(interaction)):
+        return Interaction("hbond", *process_hbond(interaction))
+    elif "pistack" in str(type(interaction)):
+        return Interaction("pistack", *process_pi_stack(interaction))
+    else:
+        raise NotImplementedError(f"Parsing not implemented for {type(interaction)}")
+
+
+def process_pi_stack(interaction):
+    protein_ring_atoms = [
+        (j.coords, j.atomicnum) for j in interaction.proteinring.atoms
+    ]
+    protein_chain = interaction.reschain
+    protein_residue = interaction.restype
+    protein_number = interaction.resnr
+    ligand_ring_atoms = [(j.coords, j.atomicnum) for j in interaction.ligandring.atoms]
+    distance = interaction.distance
+    angle = interaction.angle
+    plip_id = None
+    return (
+        protein_ring_atoms,
+        protein_chain,
+        protein_residue,
+        protein_number,
+        ligand_ring_atoms,
+        distance,
+        angle,
+        plip_id,
+    )
+
+
+def process_hydrophobic(interaction):
+    protein_atom = [(interaction.bsatom.coords, interaction.bsatom.atomicnum)]
+    protein_chain = interaction.reschain
+    protein_residue = interaction.restype
+    protein_number = interaction.resnr
+    ligand_atom = [(interaction.ligatom.coords, interaction.ligatom.atomicnum)]
+    distance = interaction.distance
+    plip_id = None
+    return (
+        protein_atom,
+        protein_chain,
+        protein_residue,
+        protein_number,
+        ligand_atom,
+        distance,
+        None,
+        plip_id,
+    )
+
+
+def process_hbond(interaction):
+    if interaction.protisdon:
+        protein_atom = [(interaction.d.coords, interaction.d.atomicnum)]
+        ligand_atom = [(interaction.a.coords, interaction.a.atomicnum)]
+    else:
+        protein_atom = [(interaction.a.coords, interaction.a.atomicnum)]
+        ligand_atom = [(interaction.d.coords, interaction.d.atomicnum)]
+
+    protein_chain = interaction.reschain
+    protein_residue = interaction.restype
+    protein_number = interaction.resnr
+    distance = interaction.distance_ad
+    angle = interaction.angle
+    plip_id = None
+    return (
+        protein_atom,
+        protein_chain,
+        protein_residue,
+        protein_number,
+        ligand_atom,
+        distance,
+        angle,
+        plip_id,
+    )
+
+
+def process_saltbridge(interaction):
+    if interaction.protispos:
+        protein_atom = [(a.coords, a.atomicnum) for a in interaction.positive.atoms]
+        ligand_atom = [(a.coords, a.atomicnum) for a in interaction.negative.atoms]
+    else:
+        protein_atom = [(a.coords, a.atomicnum) for a in interaction.negative.atoms]
+        ligand_atom = [(a.coords, a.atomicnum) for a in interaction.positive.atoms]
+    protein_chain = interaction.reschain
+    protein_residue = interaction.restype
+    protein_number = interaction.resnr
+    distance = interaction.distance
+    plip_id = None
+    return (
+        protein_atom,
+        protein_chain,
+        protein_residue,
+        protein_number,
+        ligand_atom,
+        distance,
+        None,
+        plip_id,
+    )

stcrpy/tcr_methods/tcr_batch_operations.py

@@ -0,0 +1,223 @@
+import warnings
+import os
+import pandas as pd
+
+from ..tcr_processing.TCRParser import TCRParser
+from ..tcr_interactions.TCRInteractionProfiler import TCRInteractionProfiler
+from ..tcr_geometry.TCRGeom import TCRGeom
+from ..tcr_geometry.TCRGeomFiltering import DockingGeometryFilter
+from ..tcr_formats.tcr_formats import get_sequences
+
+
+class TCRBatchOperator:
+    def __init__(self):
+        self._tcr_parser = TCRParser()
+
+    def _load_geometry_calculator(self):
+        self._geometry_calculator = TCRGeom()
+
+    def _load_geometry_filter(self):
+        self._geometry_filter = DockingGeometryFilter()
+
+    def tcrs_from_file_list(self, file_list):
+        for file in file_list:
+            tcr_id = file.split("/")[-1].split(".")[0]
+            try:
+                for tcr in self._tcr_parser.get_tcr_structure(tcr_id, file).get_TCRs():
+                    yield tcr
+            except Exception as e:
+                warnings.warn(f"Loading {file} failed with error {str(e)}")
+                yield None
+
+    def tcrs_from_file_dict(self, file_dict):
+        for tcr_id, file in file_dict.items():
+            try:
+                for tcr in self._tcr_parser.get_tcr_structure(tcr_id, file).get_TCRs():
+                    yield tcr_id, tcr
+            except Exception as e:
+                warnings.warn(f"Loading {tcr_id}: {file} failed with error {str(e)}")
+                yield None
+
+    def get_TCR_pMHC_interactions(self, tcr_generator, renumber=True, save_as_csv=None):
+        interaction_analysis_dict = {}
+        for tcr in tcr_generator:
+            if tcr is None:  # handles case where file could not be parsed in generator
+                continue
+            tcr_id = f"{tcr.parent.parent.id}_{tcr.id}"
+            if isinstance(
+                tcr, tuple
+            ):  # handle case where tcr is passed as (key, value)
+                tcr_id, tcr = tcr
+            try:
+                interaction_analysis_dict[tcr_id] = tcr.profile_peptide_interactions()
+            except Exception as e:
+                warnings.warn(
+                    f"Interactions profile failed for {tcr} with error {str(e)}"
+                )
+        interactions_df = pd.concat(
+            interaction_analysis_dict.values(),
+            keys=interaction_analysis_dict.keys(),
+            axis=0,
+        )
+
+        if save_as_csv is not None:
+            interactions_df.to_csv(save_as_csv)
+
+        return interactions_df
+
+    def get_TCR_geometry(self, tcr_generator, mode="rudolph", save_as_csv=None):
+        geometries_dict = {}
+        for tcr in tcr_generator:
+            if tcr is None:  # handles case where file could not be parsed in generator
+                continue
+
+            if isinstance(
+                tcr, tuple
+            ):  # handle case where tcr is passed as (key, value)
+                tcr_id, tcr = tcr
+            else:
+                tcr_id = f"{tcr.parent.parent.id}_{tcr.id}"
+            try:
+                geometries_dict[tcr_id] = tcr.calculate_docking_geometry(
+                    mode=mode, as_df=True
+                )
+            except Exception as e:
+                warnings.warn(
+                    f"Geometry calculation failed for {tcr} with error {str(e)}"
+                )
+        geometries_df = pd.concat(geometries_dict).droplevel(1)
+
+        if save_as_csv is not None:
+            geometries_df.to_csv(save_as_csv)
+
+        return geometries_df
+
+    def get_germlines_and_alleles(self, tcr_generator, save_as_csv=None):
+        germlines_and_alleles_dict = {}
+        for tcr in tcr_generator:
+            if tcr is None:  # handles case where file could not be parsed in generator
+                continue
+            tcr_id = f"{tcr.parent.parent.id}_{tcr.id}"
+            if isinstance(
+                tcr, tuple
+            ):  # handle case where tcr is passed as (key, value)
+                tcr_id, tcr = tcr
+            germlines_and_alleles_dict[tcr_id] = tcr.get_germlines_and_alleles()
+
+        germlines_and_alleles_df = pd.DataFrame(germlines_and_alleles_dict).T
+
+        if save_as_csv is not None:
+            germlines_and_alleles_df.to_csv(save_as_csv)
+
+        return germlines_and_alleles_df
+
+    def full_analysis(self, tcr_generator, geometry_mode="rudolph", save_dir=None):
+        from tqdm import tqdm
+
+        germlines_and_alleles_dict = {}
+        geometries_dict = {}
+        interaction_analysis_dict = {}
+
+        for tcr in tqdm(tcr_generator):
+            if tcr is None:  # handles case where file could not be parsed in generator
+                continue
+            if isinstance(
+                tcr, tuple
+            ):  # handle case where tcr is passed as (key, value)
+                tcr_id, tcr = tcr
+            else:
+                tcr_id = f"{tcr.parent.parent.id}_{tcr.id}"
+            try:
+                germlines_and_alleles_dict[tcr_id] = tcr.get_germlines_and_alleles()
+            except Exception as e:
+                warnings.warn(
+                    f"Germline and allele retrieval failed for {tcr} with error {str(e)}"
+                )
+            try:
+                geometries_dict[tcr_id] = tcr.calculate_docking_geometry(
+                    mode=geometry_mode, as_df=True
+                )
+            except Exception as e:
+                warnings.warn(
+                    f"Geometry calculation failed for {tcr} with error {str(e)}"
+                )
+            try:
+                interaction_analysis_dict[tcr_id] = tcr.profile_peptide_interactions()
+            except Exception as e:
+                warnings.warn(
+                    f"Interaction profiling failed for {tcr} with error {str(e)}"
+                )
+        germlines_and_alleles_df = pd.DataFrame(germlines_and_alleles_dict).T
+
+        geometries_df = pd.concat(geometries_dict).droplevel(1)
+
+        interactions_df = pd.concat(
+            interaction_analysis_dict.values(),
+            keys=interaction_analysis_dict.keys(),
+            axis=0,
+        )
+
+        if save_dir is not None:
+            geometries_df.to_csv(os.path.join(save_dir, "geometries.csv"))
+            germlines_and_alleles_df.to_csv(
+                os.path.join(save_dir, "germlines_and_alleles.csv")
+            )
+            interactions_df.to_csv(os.path.join(save_dir, "interactions.csv"))
+
+        return germlines_and_alleles_df, geometries_df, interactions_df
+
+
+def batch_load_TCRs(tcr_files):
+    if isinstance(tcr_files, dict):
+        return dict(TCRBatchOperator().tcrs_from_file_dict(tcr_files))
+    else:
+        return list(TCRBatchOperator().tcrs_from_file_list(tcr_files))
+
+
+def batch_yield_TCRs(tcr_files):
+    if isinstance(tcr_files, dict):
+        return TCRBatchOperator().tcrs_from_file_dict(tcr_files)
+    else:
+        return TCRBatchOperator().tcrs_from_file_list(tcr_files)
+
+
+def get_TCR_interactions(tcr_files, renumber=True, save_as_csv=None):
+    batch_ops = TCRBatchOperator()
+    if isinstance(tcr_files, list):
+        tcr_generator = batch_ops.tcrs_from_file_list(tcr_files)
+    if isinstance(tcr_files, dict):
+        tcr_generator = batch_ops.tcrs_from_file_dict(tcr_files)
+
+    return batch_ops.get_TCR_pMHC_interactions(
+        tcr_generator, renumber=renumber, save_as_csv=save_as_csv
+    )
+
+
+def get_TCR_geometry(tcr_files, mode="rudolph", save_as_csv=None):
+    batch_ops = TCRBatchOperator()
+    if isinstance(tcr_files, list):
+        tcr_generator = batch_ops.tcrs_from_file_list(tcr_files)
+    if isinstance(tcr_files, dict):
+        tcr_generator = batch_ops.tcrs_from_file_dict(tcr_files)
+
+    return batch_ops.get_TCR_geometry(tcr_generator, mode=mode, save_as_csv=save_as_csv)
+
+
+def get_germlines_and_alleles(tcr_files, save_as_csv=None):
+    batch_ops = TCRBatchOperator()
+    if isinstance(tcr_files, list):
+        tcr_generator = batch_ops.tcrs_from_file_list(tcr_files)
+    if isinstance(tcr_files, dict):
+        tcr_generator = batch_ops.tcrs_from_file_dict(tcr_files)
+
+    return batch_ops.get_germlines_and_alleles(tcr_generator, save_as_csv=save_as_csv)
+
+
+def analyse_tcrs(tcr_files, save_dir=None):
+    batch_ops = TCRBatchOperator()
+    if isinstance(tcr_files, list):
+        tcr_generator = batch_ops.tcrs_from_file_list(tcr_files)
+    if isinstance(tcr_files, dict):
+        tcr_generator = batch_ops.tcrs_from_file_dict(tcr_files)
+
+    return batch_ops.full_analysis(tcr_generator, save_dir=save_dir)

stcrpy/tcr_methods/tcr_methods.py

@@ -0,0 +1,150 @@
+import warnings
+import requests
+import os
+
+from ..tcr_processing.TCRParser import TCRParser
+from .tcr_batch_operations import batch_load_TCRs, batch_yield_TCRs
+
+
+def load_TCR(tcr_structure_file, tcr_id=None):
+    tcr_parser = TCRParser()
+    if tcr_id is None:
+        tcr_id = tcr_structure_file.split("/")[-1].split(".")[0]
+    tcr_structure = list(
+        tcr_parser.get_tcr_structure(tcr_id, tcr_structure_file).get_TCRs()
+    )
+    if len(tcr_structure) == 1:
+        return tcr_structure[0]
+    return tcr_structure
+
+
+def load_TCRs(tcr_structure_files, tcr_ids=None):
+    tcr_parser = TCRParser()
+    if isinstance(tcr_structure_files, str):  # loading single file
+        tcr_id = tcr_structure_files.split("/")[-1].split(".")[
+            0
+        ]  # set tcr_id to file name without extension
+        if tcr_ids is not None:
+            if not isinstance(tcr_ids, str):
+                warnings.warn(f"TCR ID: {tcr_ids} for a single TCR should be type str.")
+            tcr_id = tcr_ids
+
+        tcr_structure = tcr_parser.get_tcr_structure(tcr_id, tcr_structure_files)
+        return list(tcr_structure.get_TCRs())
+
+    if len(tcr_structure_files) > 10:
+        warnings.warn(
+            "Loading more than 10 TCR structure objects into memory. Consider applying generator methods to reduce memory load."
+        )
+
+    if tcr_ids is not None:
+        if len(tcr_structure_files) == len(tcr_ids):
+            return batch_load_TCRs(dict(zip(tcr_ids, tcr_structure_files)))
+        else:
+            warnings.warn(
+                f"Length of TCR IDs {len(tcr_ids)} does not match length of files {len(tcr_structure_files)}. TCR IDs reverted to default."
+            )
+    return batch_load_TCRs(tcr_structure_files)
+
+
+def yield_TCRs(tcr_structure_files, tcr_ids=None):
+    tcr_parser = TCRParser()
+    if isinstance(tcr_structure_files, str):  # loading single file
+        tcr_id = tcr_structure_files.split("/")[-1].split(".")[
+            0
+        ]  # set tcr_id to file name without extension
+        if tcr_ids is not None:
+            if not isinstance(tcr_ids, str):
+                warnings.warn(f"TCR ID: {tcr_ids} for a single TCR should be type str.")
+            tcr_id = tcr_ids
+
+        tcr_structure = tcr_parser.get_tcr_structure(tcr_id, tcr_structure_files)
+        return list(tcr_structure.get_TCRs())
+
+    if tcr_ids is not None:
+        if len(tcr_structure_files) == len(tcr_ids):
+            return batch_yield_TCRs(dict(zip(tcr_ids, tcr_structure_files)))
+        else:
+            warnings.warn(
+                f"Length of TCR IDs {len(tcr_ids)} does not match length of files {len(tcr_structure_files)}. TCR IDs reverted to default."
+            )
+    return batch_yield_TCRs(tcr_structure_files)
+
+
+def fetch_TCR(pdb_id: str):
+    """
+    Fetches and parses a T-cell receptor (TCR) structure from the STCRDab or RCSB PDB databases.
+
+    The function first attempts to download a PDB file from the STCRDab database.
+    If the PDB file is not found, it falls back to downloading a CIF file from RCSB PDB.
+    The downloaded file is then parsed using `TCRParser` to extract TCR structures.
+
+    Parameters:
+        pdb_id (str): The PDB identifier of the structure to be fetched.
+
+    Returns:
+        - A single TCR structure if exactly one is found.
+        - A list of TCR structures if multiple are found.
+        - None if no TCRs are identified (with a `UserWarning` issued).
+
+    Raises:
+        - A warning if no TCR structures are found in the downloaded file.
+        - Prints an error message if the file cannot be downloaded.
+
+    Notes:
+        - STCRDab returns an error message if the requested PDB ID does not exist.
+        - The function temporarily saves the downloaded file and deletes it after parsing.
+
+    Example:
+        tcr = fetch_TCR("6eqa")
+
+    """
+
+    stcrdab_base_url = "https://opig.stats.ox.ac.uk/webapps/stcrdab-stcrpred/pdb/"
+    pdb_base_url = "https://files.rcsb.org/download/"
+
+    filename = f"{pdb_id.upper()}.pdb"
+
+    url = stcrdab_base_url + pdb_id.lower()
+    TCR_FOUND = False
+
+    try:
+        response = requests.get(url, stream=True, timeout=10)
+        if response.status_code == 200:
+            with open(filename, "wb") as file:
+                for chunk in response.iter_content(chunk_size=1024):
+                    file.write(chunk)
+                if (
+                    not b"does not exist" in chunk
+                ):  # STCRDab returns '$PDB does not exist for downloading' if PDB code not found in database
+                    TCR_FOUND = True
+
+    except requests.exceptions.Timeout:
+        warnings.warn(f"Request to STCRDab ({url}) timed out. Trying RCSB.")
+
+    if not TCR_FOUND:
+        if os.path.exists(filename):
+            os.remove(filename)  # remove the file written with response from STCRDab
+
+        # Request from RCSB data base
+        filename = f"{pdb_id.upper()}.cif"
+        url = pdb_base_url + filename
+        response = requests.get(url, stream=True, timeout=10)
+
+        if response.status_code == 200:
+            with open(filename, "wb") as file:
+                for chunk in response.iter_content(chunk_size=1024):
+                    file.write(chunk)
+        else:
+            print("Failed to download file")
+
+    tcr_parser = TCRParser()
+    tcr = list(tcr_parser.get_tcr_structure(pdb_id, filename).get_TCRs())
+    os.remove(filename)
+    if len(tcr) == 1:
+        return tcr[0]
+    elif len(tcr) == 0:
+        warnings.warn(f"No TCRs identified in {pdb_id}")
+        return None
+    else:
+        return tcr

stcrpy/tcr_methods/tcr_reformatting.py

@@ -0,0 +1,18 @@
+def tcrs_to_AF3_json(tcrs, path=None, **kwargs):
+    from ..tcr_formats.tcr_formats import to_AF3_json
+    import json
+
+    if isinstance(tcrs[0], str):
+        from .tcr_methods import load_TCRs
+
+        tcrs = load_TCRs(tcrs)
+    else:
+        from ..tcr_processing.TCR import TCR
+
+        assert isinstance(tcrs[0], TCR)
+    multiple_job_json = [to_AF3_json(tcr, save=False, **kwargs) for tcr in tcrs]
+    path = path if path is not None else "stcrpy_AF3_TCRs.json"
+    with open(path, "w") as f:
+        json.dump(multiple_job_json, f)
+    print(f"{len(tcrs)} saved as AF3 json job: {path}")
+    return multiple_job_json

stcrpy/tcr_metrics/__init__.py

@@ -0,0 +1,2 @@
+from .tcr_rmsd import RMSD
+from .tcr_interface_rmsd import InterfaceRMSD

stcrpy/tcr_metrics/constants.py

@@ -0,0 +1,39 @@
+ATOM_TYPES = [
+    "N",
+    "CA",
+    "C",
+    "CB",
+    "O",
+    "CG",
+    "CG1",
+    "CG2",
+    "OG",
+    "OG1",
+    "SG",
+    "CD",
+    "CD1",
+    "CD2",
+    "ND1",
+    "ND2",
+    "OD1",
+    "OD2",
+    "SD",
+    "CE",
+    "CE1",
+    "CE2",
+    "CE3",
+    "NE",
+    "NE1",
+    "NE2",
+    "OE1",
+    "OE2",
+    "CH2",
+    "NH1",
+    "NH2",
+    "OH",
+    "CZ",
+    "CZ2",
+    "CZ3",
+    "NZ",
+    "OXT",
+]