stcrpy 1.0.0__py3-none-any.whl

stcrpy/tcr_datasets/utils.py

@@ -0,0 +1,350 @@
+CHAIN_TYPE_ONE_HOT_ENCODING = {
+    "A": 0,  # alpha TCR
+    "D": 0,  # delta TCR
+    "B": 1,  # beta TCR
+    "G": 1,  # gamma TCR
+    "MHC": 2,
+    "antigen": 3,
+}
+
+MHC_CHAIN_TYPES = [
+    "GA1",
+    "GA1L",
+    "GA2",
+    "GA2L",
+    "GA",
+    "GB",
+    "CD1",
+    "CD1L",
+    "MR1",
+    "MR1",
+    "MHC",
+    "MH1",
+    "MH2",
+]
+
+TCR_REGION_ONE_HOT_ENCODING = {
+    "NOT_CDR": 0,
+    "CDRA1": 1,
+    "CDRA2": 2,
+    "CDRA3": 3,
+    "CDRB1": 4,
+    "CDRB2": 5,
+    "CDRB3": 6,
+}
+
+AMINO_ACID_ONEHOT_ENCODING = {
+    "ALA": 0,
+    "ARG": 1,
+    "ASN": 2,
+    "ASP": 3,
+    "CYS": 4,
+    "GLU": 5,
+    "GLN": 6,
+    "GLY": 7,
+    "HIS": 8,
+    "ILE": 9,
+    "LEU": 10,
+    "LYS": 11,
+    "MET": 12,
+    "PHE": 13,
+    "PRO": 14,
+    "SER": 15,
+    "THR": 16,
+    "TRP": 17,
+    "TYR": 18,
+    "VAL": 19,
+    "UNK": 20,
+}
+
+
+ATOM14_ATOM_NAMES = {
+    "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""],
+    "ARG": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD",
+        "NE",
+        "CZ",
+        "NH1",
+        "NH2",
+        "",
+        "",
+        "",
+    ],
+    "ASN": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "OD1",
+        "ND2",
+        "",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "ASP": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "OD1",
+        "OD2",
+        "",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""],
+    "GLN": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD",
+        "OE1",
+        "NE2",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "GLU": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD",
+        "OE1",
+        "OE2",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""],
+    "HIS": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "ND1",
+        "CD2",
+        "CE1",
+        "NE2",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "ILE": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG1",
+        "CG2",
+        "CD1",
+        "",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "LEU": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "LYS": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD",
+        "CE",
+        "NZ",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "MET": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "SD",
+        "CE",
+        "",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "PHE": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "CE1",
+        "CE2",
+        "CZ",
+        "",
+        "",
+        "",
+    ],
+    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""],
+    "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""],
+    "THR": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "OG1",
+        "CG2",
+        "",
+        "",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "TRP": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "NE1",
+        "CE2",
+        "CE3",
+        "CZ2",
+        "CZ3",
+        "CH2",
+    ],
+    "TYR": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "CE1",
+        "CE2",
+        "CZ",
+        "OH",
+        "",
+        "",
+    ],
+    "VAL": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG1",
+        "CG2",
+        "",
+        "",
+        "",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "UNK": ["", "", "", "", "", "", "", "", "", "", "", "", "", ""],
+}
+
+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",
+]
+
+ATOM37_ATOM_ONEHOT_ENCODING = {atom_type: i for i, atom_type in enumerate(ATOM_TYPES)}

stcrpy/tcr_formats/tcr_formats.py

@@ -0,0 +1,114 @@
+from Bio.SeqUtils import seq1
+import json
+import os
+
+
+def to_AF3_json(
+    tcr: "TCR",
+    tcr_only: bool = True,
+    save: bool = True,
+    save_dir: str = "",
+    name: str = None,
+    V_domain_only: bool = False,
+) -> dict:
+    """Converts TCR object to dict in Alphafold 3 JSON input format, ie. amino acid sequences.
+    Eg:
+    {
+        "name": Job name,
+        "modelSeeds": [],
+        "sequences": [
+            {"proteinChain": {"sequence": AAAAAAAAAAAAAA, "count": 1}},
+            {"proteinChain": {"sequence": AAAAAAAAAAAAAA, "count": 1}},
+            {"proteinChain": {"sequence": AAAAAAAAAAAAAA, "count": 1}},
+        ],
+    }
+
+    Args:
+        tcr (TCR): TCR structure object
+        tcr_only (bool, optional): Whether to include TCR sequence only, excluding antigen and MHC. Defaults to True.
+        save (bool, optional): Whether to save dict as JSON file. Defaults to True.
+        save_dir (str, optional): Directory to save JSON files to. Defaults to "".
+        name (str, optional): TCR ID to use as name for AF3 job. Defaults to None.
+        V_domain_only (bool, optional): Include full TCR sequence or only the variable domain (1-128 IMGT numbering). Defaults to False.
+
+    Returns:
+        dict: Nested dictionary of AF3 sequence inputs.
+    """
+    if V_domain_only:
+        residue_nrs = list(range(128))
+    else:
+        residue_nrs = None
+    tcr_sequences = get_sequences(tcr, residues_to_include=residue_nrs)
+    if not tcr_only:
+        if len(tcr.get_MHC()) > 0:
+            mhc_sequences = get_sequences(tcr.get_MHC()[0])
+            tcr_sequences.update(mhc_sequences)
+
+        if len(tcr.get_antigen()) > 0:
+            antigen_sequence = get_sequences(tcr.get_antigen()[0])
+            tcr_sequences.update(antigen_sequence)
+    name = name if name is not None else f"{tcr.parent.parent.id}_{tcr.id}"
+    tcr_json = {
+        "name": name,
+        "modelSeeds": [],
+        "sequences": [
+            {"proteinChain": {"sequence": seq, "count": 1}}
+            for _, seq in tcr_sequences.items()
+        ],
+    }
+    if save:
+        path = os.path.join(save_dir, f"{name}.json")
+        with open(path, "w") as f:
+            json.dump(tcr_json, f)
+    return tcr_json
+
+
+def get_sequences(
+    entity: "Bio.PDB.Entity",
+    amino_acids_only: bool = True,
+    residues_to_include: list = None,
+) -> dict:
+    """Extract seqeunces from strcuture objects as dictionary.
+
+    Args:
+        entity (Bio.PDB.Entity): Stucture object
+        amino_acids_only (bool, optional): Whether to remove non-amino acid 'X' from sequences. Defaults to True.
+        residues_to_include (list, optional): List of residue IDs to include in sequence. Defaults to None.
+
+    Raises:
+        e: AttributeError if entity has no attribute .get_chains(). The assuems entity is chain level and returns single sequence
+
+    Returns:
+        dict: Dictionary of amino acid sequences, keyed by chain ID in strcuctre entity.
+    """
+
+    if residues_to_include is None:
+
+        def residue_filter(res):
+            return True
+
+    else:
+
+        def residue_filter(res):
+            return res.id[1] in residues_to_include
+    try:
+        sequences = {
+            chain.id: seq1(
+                "".join(residue.resname for residue in chain if residue_filter(residue))
+            )
+            for chain in entity.get_chains()
+        }
+    except AttributeError as e:
+        if entity.level == "C":
+            sequences = {
+                entity.id: seq1(
+                    "".join(
+                        residue.resname for residue in entity if residue_filter(residue)
+                    )
+                )
+            }
+        else:
+            raise e
+    if amino_acids_only:
+        sequences = {k: seq.replace("X", "") for k, seq in sequences.items()}
+    return sequences