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

stcrpy/__init__.py

@@ -2,4 +2,4 @@ from .tcr_processing.TCRParser import TCRParser
 from .tcr_processing.TCRIO import TCRIO
 from .tcr_geometry.TCRDock import TCRDock
 from .tcr_geometry.TCRGeom import TCRGeom
-from .tcr_methods.tcr_methods import load_TCRs, fetch_TCR, yield_TCRs, load_TCR
+from .tcr_methods.tcr_methods import load_TCRs, fetch_TCRs, yield_TCRs, load_TCR

stcrpy/tcr_formats/tcr_formats.py

@@ -99,7 +99,7 @@ def get_sequences(
             for chain in entity.get_chains()
         }
     except AttributeError as e:
-        if entity.level == "C":
+        if entity.level == "C" or entity.level == "F":  # covers chains and fragments
             sequences = {
                 entity.id: seq1(
                     "".join(
@@ -112,3 +112,22 @@ def get_sequences(
     if amino_acids_only:
         sequences = {k: seq.replace("X", "") for k, seq in sequences.items()}
     return sequences
+
+
+def merge_chains(chains, new_chain_id=None):
+    from Bio import PDB
+
+    if new_chain_id is None:
+        new_chain_id = f"{chains[0].id}_{chains[1].id}"
+    new_chain = PDB.Chain.Chain(new_chain_id)
+    new_res_id = 1
+
+    for chain in chains:
+        for residue in chain.get_residues():
+            new_residue = residue.copy()
+            new_residue.id = (" ", new_res_id, " ")
+
+            new_chain.add(new_residue)
+            new_res_id += 1
+
+    return new_chain

stcrpy/tcr_geometry/TCRAngle.py

@@ -0,0 +1,177 @@
+import os
+import math
+import numpy as np
+
+from Bio.SVDSuperimposer import SVDSuperimposer
+
+
+def superimpose(coordsA, coordsB):
+    sup = SVDSuperimposer()
+    sup.set(coordsA, coordsB)
+    sup.run()
+    return sup
+
+
+class TCRAngle:
+
+    def __init__(self, tcr_type="abTCR"):
+        """Unified class for calculating TCR orientation angles for abTCRs and gdTCRs."""
+
+        self.dat_path = os.path.join(os.path.dirname(__file__), "reference_data")
+
+        from stcrpy.tcr_processing.TCRParser import TCRParser  # avoids circular import
+
+        self.sparser = TCRParser(QUIET=True)
+
+        self.chain_configs = {
+            "abTCR": {
+                "chains": ("A", "B"),
+                "consensus": ("consensus_A.pdb", "consensus_B.pdb"),
+                "pc": ("pcA.txt", "pcB.txt"),
+                "coreset": ("Acoreset.txt", "Bcoreset.txt"),
+                "angle_labels": ["BA", "BC1", "AC1", "BC2", "AC2", "dc"],
+            },
+            "gdTCR": {
+                "chains": ("G", "D"),
+                "consensus": ("consensus_G.pdb", "consensus_D.pdb"),
+                "pc": ("pcA.txt", "pcB.txt"),
+                "coreset": ("Acoreset.txt", "Bcoreset.txt"),
+                "angle_labels": ["DG", "DC1", "GC1", "DC2", "GC2", "dc"],
+            },
+        }
+
+        self._init_tcr_type_specific_reference_data(tcr_type)
+
+    def _init_tcr_type_specific_reference_data(self, tcr_type):
+        if tcr_type not in self.chain_configs:
+            raise ValueError("tcr_type must be 'abTCR' or 'gdTCR'")
+        self.tcr_type = tcr_type
+        self.cfg = self.chain_configs[tcr_type]
+        self._read_consensus()
+        self._read_pc()
+        self._read_coreset()
+
+    def _normalise(self, v):
+        a = np.array(v)
+        return a / np.linalg.norm(a)
+
+    def _read_consensus(self):
+        c1, c2 = self.cfg["chains"]
+        f1, f2 = self.cfg["consensus"]
+        self.consensus_1 = self.sparser.get_tcr_structure(
+            c1, os.path.join(self.dat_path, f1)
+        )
+        self.consensus_2 = self.sparser.get_tcr_structure(
+            c2, os.path.join(self.dat_path, f2)
+        )
+
+        self.consensus_1_atoms = sorted(
+            list(self.consensus_1.get_atoms()), key=lambda x: x.parent.id[1]
+        )
+        self.consensus_2_atoms = sorted(
+            list(self.consensus_2.get_atoms()), key=lambda x: x.parent.id[1]
+        )
+
+    def _read_pc(self):
+        f1, f2 = self.cfg["pc"]
+        self.pos1 = [
+            list(map(float, x.split())) for x in open(os.path.join(self.dat_path, f1))
+        ]
+        self.pos2 = [
+            list(map(float, x.split())) for x in open(os.path.join(self.dat_path, f2))
+        ]
+        self.c1 = [
+            6 * 0.5 * self.pos1[0][i] - 2 * 0.5 * self.pos1[1][i] + self.pos1[2][i]
+            for i in range(3)
+        ]
+        self.c2 = [
+            -10 * 0.5 * self.pos2[0][i] + 1 * 0.5 * self.pos2[1][i] + self.pos2[2][i]
+            for i in range(3)
+        ]
+        self.p1a = [self.c1[i] + self.pos1[0][i] for i in range(3)]
+        self.p1b = [self.c1[i] + self.pos1[1][i] for i in range(3)]
+        self.p2a = [self.c2[i] + self.pos2[0][i] for i in range(3)]
+        self.p2b = [self.c2[i] + self.pos2[1][i] for i in range(3)]
+
+    def _read_coreset(self):
+        f1, f2 = self.cfg["coreset"]
+        self.coreset1 = [
+            int(l.strip()[1:]) for l in open(os.path.join(self.dat_path, f1))
+        ]
+        self.coreset2 = [
+            int(l.strip()[1:]) for l in open(os.path.join(self.dat_path, f2))
+        ]
+
+    def calculate_angles(self, tcr):
+        if self.tcr_type not in str(type(tcr)):
+            self.tcr_type = str(type(tcr)).split(".")[-1].split("'>")[0]
+            self._init_tcr_type_specific_reference_data(self.tcr_type)
+
+        c1, c2 = self.cfg["chains"]
+        tcr_chain_1_coreset_atoms = [
+            res["CA"]
+            for res in tcr[tcr.get_domain_assignment()["V" + c1]].get_residues()
+            if res.id[1] in self.coreset1
+        ]
+        tcr_chain_2_coreset_atoms = [
+            res["CA"]
+            for res in tcr[tcr.get_domain_assignment()["V" + c2]].get_residues()
+            if res.id[1] in self.coreset2
+        ]
+
+        tcr_chain_1_coreset_atoms = sorted(
+            tcr_chain_1_coreset_atoms, key=lambda x: x.parent.id[1]
+        )
+        tcr_chain_2_coreset_atoms = sorted(
+            tcr_chain_2_coreset_atoms, key=lambda x: x.parent.id[1]
+        )
+
+        rot1, tran1 = superimpose(
+            np.asarray([a.coord for a in tcr_chain_1_coreset_atoms]),
+            np.asarray([a.coord for a in self.consensus_1_atoms]),
+        ).get_rotran()
+        rot2, tran2 = superimpose(
+            np.asarray([a.coord for a in tcr_chain_2_coreset_atoms]),
+            np.asarray([a.coord for a in self.consensus_2_atoms]),
+        ).get_rotran()
+
+        points1 = [
+            np.dot(self.c1, rot1) + tran1,
+            np.dot(self.p1a, rot1) + tran1,
+            np.dot(self.p1b, rot1) + tran1,
+        ]
+        points2 = [
+            np.dot(self.c2, rot2) + tran2,
+            np.dot(self.p2a, rot2) + tran2,
+            np.dot(self.p2b, rot2) + tran2,
+        ]
+
+        C = self._normalise([points2[0][i] - points1[0][i] for i in range(3)])
+        Cminus = [-x for x in C]
+
+        v1a = self._normalise([points1[1][i] - points1[0][i] for i in range(3)])
+        v1b = self._normalise([points1[2][i] - points1[0][i] for i in range(3)])
+        v2a = self._normalise([points2[1][i] - points2[0][i] for i in range(3)])
+        v2b = self._normalise([points2[2][i] - points2[0][i] for i in range(3)])
+
+        dc = np.linalg.norm(points2[0] - points1[0])
+
+        n_x = np.cross(v1a, C)
+        n_y = np.cross(C, n_x)
+        tmp1 = self._normalise([0, np.dot(v1a, n_x), np.dot(v1a, n_y)])
+        tmp2 = self._normalise([0, np.dot(v2a, n_x), np.dot(v2a, n_y)])
+
+        angle = math.degrees(math.acos(np.dot(tmp1, tmp2)))
+        if np.dot(np.cross(tmp1, tmp2), [1, 0, 0]) < 0:
+            angle = -angle
+
+        results = [
+            angle,
+            math.degrees(math.acos(np.dot(v1a, C))),
+            math.degrees(math.acos(np.dot(v2a, Cminus))),
+            math.degrees(math.acos(np.dot(v1b, C))),
+            math.degrees(math.acos(np.dot(v2b, Cminus))),
+            dc,
+        ]
+
+        return dict(zip(self.cfg["angle_labels"], results))

stcrpy/tcr_geometry/reference_data/Acoreset.txt

@@ -0,0 +1,30 @@
+A20
+A22
+A45
+A86
+A52
+A18
+A88
+A94
+A43
+A44
+A89
+A90
+A85
+A16
+A99
+A104
+A17
+A100
+A42
+A41
+A40
+A98
+A97
+A101
+A39
+A38
+A96
+A95
+A103
+A102

stcrpy/tcr_geometry/reference_data/Bcoreset.txt

@@ -0,0 +1,30 @@
+B22
+B23
+B104
+B16
+B24
+B44
+B43
+B86
+B6
+B42
+B5
+B15
+B41
+B53
+B8
+B52
+B51
+B9
+B10
+B12
+B11
+B103
+B14
+B102
+B40
+B39
+B38
+B101
+B100
+B99

stcrpy/tcr_geometry/reference_data/consensus_A.pdb

@@ -0,0 +1,31 @@
+ATOM      1  CA  GLY A  16     -12.766  -6.146   4.154  0.00  0.00           C  
+ATOM      2  CA  GLU A  17     -12.984  -3.096   1.985  0.00  0.00           C  
+ATOM      3  CA  ASN A  18     -10.777  -2.216  -0.940  0.00  0.00           C  
+ATOM      4  CA  THR A  20      -6.622   2.515  -3.475  0.00  0.00           C  
+ATOM      5  CA  ASN A  22      -2.375   6.893  -5.615  0.00  0.00           C  
+ATOM      6  CA  ASN A  38      10.871   1.660 -12.510  0.00  0.00           C  
+ATOM      7  CA  LEU A  39       8.338   2.153  -9.768  0.00  0.00           C  
+ATOM      8  CA  GLN A  40       8.338   0.459  -6.382  0.00  0.00           C  
+ATOM      9  CA  TRP A  41       5.923   0.253  -3.477  0.00  0.00           C  
+ATOM     10  CA  TYR A  42       7.139  -0.195   0.076  0.00  0.00           C  
+ATOM     11  CA  ARG A  43       5.258  -0.682   3.297  0.00  0.00           C  
+ATOM     12  CA  GLN A  44       6.468   0.441   6.681  0.00  0.00           C  
+ATOM     13  CA  ASN A  45       4.955  -0.914   9.806  0.00  0.00           C  
+ATOM     14  CA  HIS A  52       9.234  -4.199  -0.523  0.00  0.00           C  
+ATOM     15  CA  LYS A  85       2.672   9.293 -15.022  0.00  0.00           C  
+ATOM     16  CA  SER A  86       0.797   7.515 -12.294  0.00  0.00           C  
+ATOM     17  CA  SER A  88      -1.761   3.077  -8.014  0.00  0.00           C  
+ATOM     18  CA  LEU A  89      -1.639   0.408  -5.358  0.00  0.00           C  
+ATOM     19  CA  LEU A  90      -4.860  -1.486  -4.764  0.00  0.00           C  
+ATOM     20  CA  SER A  94      -8.077  -5.060   3.718  0.00  0.00           C  
+ATOM     21  CA  ARG A  95      -7.018  -6.656   6.969  0.00  0.00           C  
+ATOM     22  CA  ALA A  96      -5.916  -5.018  10.209  0.00  0.00           C  
+ATOM     23  CA  ALA A  97      -2.373  -6.175   9.542  0.00  0.00           C  
+ATOM     24  CA  ASP A  98      -2.346  -3.975   6.457  0.00  0.00           C  
+ATOM     25  CA  THR A  99      -2.284  -0.874   8.638  0.00  0.00           C  
+ATOM     26  CA  ALA A 100       0.968   0.902   7.897  0.00  0.00           C  
+ATOM     27  CA  SER A 101       2.587   3.693   5.962  0.00  0.00           C  
+ATOM     28  CA  TYR A 102       2.875   3.032   2.230  0.00  0.00           C  
+ATOM     29  CA  PHE A 103       5.639   4.573   0.152  0.00  0.00           C  
+ATOM     30  CA  CYS A 104       5.899   4.893  -3.569  0.00  0.00           C  
+TER

stcrpy/tcr_geometry/reference_data/consensus_B.pdb

@@ -0,0 +1,31 @@
+ATOM      1  CA  THR B   5       7.680   4.957   6.318  0.00  0.00           C  
+ATOM      2  CA  GLN B   6       4.157   4.013   5.323  0.00  0.00           C  
+ATOM      3  CA  PRO B   8      -1.597   4.163   7.655  0.00  0.00           C  
+ATOM      4  CA  LYS B   9      -0.563   0.753   8.980  0.00  0.00           C  
+ATOM      5  CA  PHE B  10      -4.125  -0.534   8.872  0.00  0.00           C  
+ATOM      6  CA  GLN B  11      -7.228   0.591   7.033  0.00  0.00           C  
+ATOM      7  CA  VAL B  12     -10.787  -0.634   6.761  0.00  0.00           C  
+ATOM      8  CA  LYS B  14     -16.366   0.265   2.522  0.00  0.00           C  
+ATOM      9  CA  THR B  15     -16.732  -0.673  -1.121  0.00  0.00           C  
+ATOM     10  CA  GLY B  16     -16.733   2.365  -3.364  0.00  0.00           C  
+ATOM     11  CA  GLN B  22       1.267   7.341   1.834  0.00  0.00           C  
+ATOM     12  CA  CYS B  23       4.911   6.380   1.563  0.00  0.00           C  
+ATOM     13  CA  ALA B  24       7.861   8.498   2.594  0.00  0.00           C  
+ATOM     14  CA  TYR B  38      13.247   5.051  -7.303  0.00  0.00           C  
+ATOM     15  CA  MET B  39      10.086   4.125  -5.461  0.00  0.00           C  
+ATOM     16  CA  SER B  40       7.830   1.118  -5.919  0.00  0.00           C  
+ATOM     17  CA  TRP B  41       4.421  -0.118  -4.828  0.00  0.00           C  
+ATOM     18  CA  TYR B  42       3.741  -3.827  -4.381  0.00  0.00           C  
+ATOM     19  CA  ARG B  43       0.720  -5.770  -3.219  0.00  0.00           C  
+ATOM     20  CA  GLN B  44       0.883  -9.078  -1.429  0.00  0.00           C  
+ATOM     21  CA  ARG B  51       3.180  -8.202  -7.285  0.00  0.00           C  
+ATOM     22  CA  LEU B  52       4.411  -4.907  -8.681  0.00  0.00           C  
+ATOM     23  CA  ILE B  53       1.744  -2.252  -9.110  0.00  0.00           C  
+ATOM     24  CA  ASP B  86       6.783  11.037  -1.067  0.00  0.00           C  
+ATOM     25  CA  THR B  99      -8.080  -6.335   0.175  0.00  0.00           C  
+ATOM     26  CA  SER B 100      -4.378  -6.994   0.577  0.00  0.00           C  
+ATOM     27  CA  VAL B 101      -1.115  -5.893   2.125  0.00  0.00           C  
+ATOM     28  CA  TYR B 102       0.511  -2.993   0.305  0.00  0.00           C  
+ATOM     29  CA  PHE B 103       4.257  -2.495   0.489  0.00  0.00           C  
+ATOM     30  CA  CYS B 104       6.223   0.500  -0.633  0.00  0.00           C  
+TER

stcrpy/tcr_geometry/reference_data/consensus_D.pdb

@@ -0,0 +1,31 @@
+ATOM      1  CA  GLY D  16     -12.766  -6.146   4.154  0.00  0.00           C  
+ATOM      2  CA  GLU D  17     -12.984  -3.096   1.985  0.00  0.00           C  
+ATOM      3  CA  ASN D  18     -10.777  -2.216  -0.940  0.00  0.00           C  
+ATOM      4  CA  THR D  20      -6.622   2.515  -3.475  0.00  0.00           C  
+ATOM      5  CA  ASN D  22      -2.375   6.893  -5.615  0.00  0.00           C  
+ATOM      6  CA  ASN D  38      10.871   1.660 -12.510  0.00  0.00           C  
+ATOM      7  CA  LEU D  39       8.338   2.153  -9.768  0.00  0.00           C  
+ATOM      8  CA  GLN D  40       8.338   0.459  -6.382  0.00  0.00           C  
+ATOM      9  CA  TRP D  41       5.923   0.253  -3.477  0.00  0.00           C  
+ATOM     10  CA  TYR D  42       7.139  -0.195   0.076  0.00  0.00           C  
+ATOM     11  CA  ARG D  43       5.258  -0.682   3.297  0.00  0.00           C  
+ATOM     12  CA  GLN D  44       6.468   0.441   6.681  0.00  0.00           C  
+ATOM     13  CA  ASN D  45       4.955  -0.914   9.806  0.00  0.00           C  
+ATOM     14  CA  HIS D  52       9.234  -4.199  -0.523  0.00  0.00           C  
+ATOM     15  CA  LYS D  85       2.672   9.293 -15.022  0.00  0.00           C  
+ATOM     16  CA  SER D  86       0.797   7.515 -12.294  0.00  0.00           C  
+ATOM     17  CA  SER D  88      -1.761   3.077  -8.014  0.00  0.00           C  
+ATOM     18  CA  LEU D  89      -1.639   0.408  -5.358  0.00  0.00           C  
+ATOM     19  CA  LEU D  90      -4.860  -1.486  -4.764  0.00  0.00           C  
+ATOM     20  CA  SER D  94      -8.077  -5.060   3.718  0.00  0.00           C  
+ATOM     21  CA  ARG D  95      -7.018  -6.656   6.969  0.00  0.00           C  
+ATOM     22  CA  ALA D  96      -5.916  -5.018  10.209  0.00  0.00           C  
+ATOM     23  CA  ALA D  97      -2.373  -6.175   9.542  0.00  0.00           C  
+ATOM     24  CA  ASP D  98      -2.346  -3.975   6.457  0.00  0.00           C  
+ATOM     25  CA  THR D  99      -2.284  -0.874   8.638  0.00  0.00           C  
+ATOM     26  CA  ALA D 100       0.968   0.902   7.897  0.00  0.00           C  
+ATOM     27  CA  SER D 101       2.587   3.693   5.962  0.00  0.00           C  
+ATOM     28  CA  TYR D 102       2.875   3.032   2.230  0.00  0.00           C  
+ATOM     29  CA  PHE D 103       5.639   4.573   0.152  0.00  0.00           C  
+ATOM     30  CA  CYS D 104       5.899   4.893  -3.569  0.00  0.00           C  
+TER

stcrpy/tcr_geometry/reference_data/consensus_G.pdb

@@ -0,0 +1,31 @@
+ATOM      1  CA  THR G   5       7.680   4.957   6.318  0.00  0.00           C  
+ATOM      2  CA  GLN G   6       4.157   4.013   5.323  0.00  0.00           C  
+ATOM      3  CA  PRO G   8      -1.597   4.163   7.655  0.00  0.00           C  
+ATOM      4  CA  LYS G   9      -0.563   0.753   8.980  0.00  0.00           C  
+ATOM      5  CA  PHE G  10      -4.125  -0.534   8.872  0.00  0.00           C  
+ATOM      6  CA  GLN G  11      -7.228   0.591   7.033  0.00  0.00           C  
+ATOM      7  CA  VAL G  12     -10.787  -0.634   6.761  0.00  0.00           C  
+ATOM      8  CA  LYS G  14     -16.366   0.265   2.522  0.00  0.00           C  
+ATOM      9  CA  THR G  15     -16.732  -0.673  -1.121  0.00  0.00           C  
+ATOM     10  CA  GLY G  16     -16.733   2.365  -3.364  0.00  0.00           C  
+ATOM     11  CA  GLN G  22       1.267   7.341   1.834  0.00  0.00           C  
+ATOM     12  CA  CYS G  23       4.911   6.380   1.563  0.00  0.00           C  
+ATOM     13  CA  ALA G  24       7.861   8.498   2.594  0.00  0.00           C  
+ATOM     14  CA  TYR G  38      13.247   5.051  -7.303  0.00  0.00           C  
+ATOM     15  CA  MET G  39      10.086   4.125  -5.461  0.00  0.00           C  
+ATOM     16  CA  SER G  40       7.830   1.118  -5.919  0.00  0.00           C  
+ATOM     17  CA  TRP G  41       4.421  -0.118  -4.828  0.00  0.00           C  
+ATOM     18  CA  TYR G  42       3.741  -3.827  -4.381  0.00  0.00           C  
+ATOM     19  CA  ARG G  43       0.720  -5.770  -3.219  0.00  0.00           C  
+ATOM     20  CA  GLN G  44       0.883  -9.078  -1.429  0.00  0.00           C  
+ATOM     21  CA  ARG G  51       3.180  -8.202  -7.285  0.00  0.00           C  
+ATOM     22  CA  LEU G  52       4.411  -4.907  -8.681  0.00  0.00           C  
+ATOM     23  CA  ILE G  53       1.744  -2.252  -9.110  0.00  0.00           C  
+ATOM     24  CA  ASP G  86       6.783  11.037  -1.067  0.00  0.00           C  
+ATOM     25  CA  THR G  99      -8.080  -6.335   0.175  0.00  0.00           C  
+ATOM     26  CA  SER G 100      -4.378  -6.994   0.577  0.00  0.00           C  
+ATOM     27  CA  VAL G 101      -1.115  -5.893   2.125  0.00  0.00           C  
+ATOM     28  CA  TYR G 102       0.511  -2.993   0.305  0.00  0.00           C  
+ATOM     29  CA  PHE G 103       4.257  -2.495   0.489  0.00  0.00           C  
+ATOM     30  CA  CYS G 104       6.223   0.500  -0.633  0.00  0.00           C  
+TER

stcrpy/tcr_geometry/reference_data/pcA.txt

@@ -0,0 +1,3 @@
+0.170468    0.120233    -0.978000
+-0.460937   0.886969    0.028699
+5.223556    2.001083    1.418880

stcrpy/tcr_geometry/reference_data/pcB.txt

@@ -0,0 +1,3 @@
+0.595917    0.758482    -0.263794
+-0.056280   0.367130    0.928465
+2.455201    -3.709282   -1.446316

stcrpy/tcr_interactions/TCRInteractionProfiler.py

@@ -9,6 +9,7 @@ try:
     import plip
     from plip.basic.remote import VisualizerData
     from plip.visualization.visualize import visualize_in_pymol
+
 except ModuleNotFoundError as e:
     if "pymol" in str(e):
         warnings.warn(
@@ -413,7 +414,8 @@ class TCRInteractionProfiler:
             heatmap_b[pair[0][1], ligand_number_mapping[int(pair[1][1])]] = (
                 heatmap_b[pair[0][1], ligand_number_mapping[int(pair[1][1])]] + 1
             )
-        ax_beta.imshow(heatmap_b.T, cmap="PuRd")
+
+        im_beta = ax_beta.imshow(heatmap_b.T, cmap="PuRd")
         for i in plot_index:
             ax_beta.axhline(y=i - 0.5, color="blue", linewidth=1)
             ax_beta.axhline(y=i + 0.5, color="blue", linewidth=1)
@@ -427,6 +429,11 @@ class TCRInteractionProfiler:
             ax_beta.set_xticks([], [], rotation=90)
             ax_beta.set_yticks([], [])
 
+        cbar = fig.colorbar(
+            im_beta, ax=[ax_alpha, ax_beta], orientation="vertical", shrink=0.8
+        )
+        cbar.set_label("Interaction count", rotation=270, labelpad=15)
+
         if save_as is not None:
             fig.savefig(save_as, bbox_inches="tight", dpi=200)
 

stcrpy/tcr_methods/tcr_batch_operations.py

@@ -19,20 +19,24 @@ class TCRBatchOperator:
     def _load_geometry_filter(self):
         self._geometry_filter = DockingGeometryFilter()
 
-    def tcrs_from_file_list(self, file_list):
+    def tcrs_from_file_list(self, file_list, **kwargs):
         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():
+                for tcr in self._tcr_parser.get_tcr_structure(
+                    tcr_id, file, **kwargs
+                ).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):
+    def tcrs_from_file_dict(self, file_dict, **kwargs):
         for tcr_id, file in file_dict.items():
             try:
-                for tcr in self._tcr_parser.get_tcr_structure(tcr_id, file).get_TCRs():
+                for tcr in self._tcr_parser.get_tcr_structure(
+                    tcr_id, file, **kwargs
+                ).get_TCRs():
                     yield tcr_id, tcr
             except Exception as e:
                 warnings.warn(f"Loading {tcr_id}: {file} failed with error {str(e)}")
@@ -167,18 +171,18 @@ class TCRBatchOperator:
         return germlines_and_alleles_df, geometries_df, interactions_df
 
 
-def batch_load_TCRs(tcr_files):
+def batch_load_TCRs(tcr_files, **kwargs):
     if isinstance(tcr_files, dict):
-        return dict(TCRBatchOperator().tcrs_from_file_dict(tcr_files))
+        return dict(TCRBatchOperator().tcrs_from_file_dict(tcr_files), **kwargs)
     else:
-        return list(TCRBatchOperator().tcrs_from_file_list(tcr_files))
+        return list(TCRBatchOperator().tcrs_from_file_list(tcr_files), **kwargs)
 
 
-def batch_yield_TCRs(tcr_files):
+def batch_yield_TCRs(tcr_files, **kwargs):
     if isinstance(tcr_files, dict):
-        return TCRBatchOperator().tcrs_from_file_dict(tcr_files)
+        return TCRBatchOperator().tcrs_from_file_dict(tcr_files, **kwargs)
     else:
-        return TCRBatchOperator().tcrs_from_file_list(tcr_files)
+        return TCRBatchOperator().tcrs_from_file_list(tcr_files, **kwargs)
 
 
 def get_TCR_interactions(tcr_files, renumber=True, save_as_csv=None):

stcrpy/tcr_methods/tcr_methods.py

@@ -2,23 +2,24 @@ import warnings
 import requests
 import os
 
+from ..tcr_processing import TCR
 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):
+def load_TCR(tcr_structure_file, tcr_id=None, **kwargs):
     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()
+        tcr_parser.get_tcr_structure(tcr_id, tcr_structure_file, **kwargs).get_TCRs()
     )
     if len(tcr_structure) == 1:
         return tcr_structure[0]
     return tcr_structure
 
 
-def load_TCRs(tcr_structure_files, tcr_ids=None):
+def load_TCRs(tcr_structure_files, tcr_ids=None, **kwargs):
     tcr_parser = TCRParser()
     if isinstance(tcr_structure_files, str):  # loading single file
         tcr_id = tcr_structure_files.split("/")[-1].split(".")[
@@ -29,7 +30,9 @@ def load_TCRs(tcr_structure_files, tcr_ids=None):
                 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)
+        tcr_structure = tcr_parser.get_tcr_structure(
+            tcr_id, tcr_structure_files, **kwargs
+        )
         return list(tcr_structure.get_TCRs())
 
     if len(tcr_structure_files) > 10:
@@ -39,15 +42,15 @@ def load_TCRs(tcr_structure_files, tcr_ids=None):
 
     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)))
+            return batch_load_TCRs(dict(zip(tcr_ids, tcr_structure_files)), **kwargs)
         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)
+    return batch_load_TCRs(tcr_structure_files, **kwargs)
 
 
-def yield_TCRs(tcr_structure_files, tcr_ids=None):
+def yield_TCRs(tcr_structure_files, tcr_ids=None, **kwargs):
     tcr_parser = TCRParser()
     if isinstance(tcr_structure_files, str):  # loading single file
         tcr_id = tcr_structure_files.split("/")[-1].split(".")[
@@ -58,20 +61,22 @@ def yield_TCRs(tcr_structure_files, tcr_ids=None):
                 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)
+        tcr_structure = tcr_parser.get_tcr_structure(
+            tcr_id, tcr_structure_files, **kwargs
+        )
         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)))
+            return batch_yield_TCRs(dict(zip(tcr_ids, tcr_structure_files)), **kwargs)
         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)
+    return batch_yield_TCRs(tcr_structure_files, **kwargs)
 
 
-def fetch_TCR(pdb_id: str):
+def fetch_TCRs(pdb_id: str, **kwargs) -> list[TCR]:
     """
     Fetches and parses a T-cell receptor (TCR) structure from the STCRDab or RCSB PDB databases.
 
@@ -83,9 +88,7 @@ def fetch_TCR(pdb_id: str):
         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).
+        A list of the TCR structures found in the fetched PDB file
 
     Raises:
         - A warning if no TCR structures are found in the downloaded file.
@@ -96,7 +99,7 @@ def fetch_TCR(pdb_id: str):
         - The function temporarily saves the downloaded file and deletes it after parsing.
 
     Example:
-        tcr = fetch_TCR("6eqa")
+        tcr = fetch_TCRs("6eqa")
 
     """
 
@@ -139,12 +142,10 @@ def fetch_TCR(pdb_id: str):
             print("Failed to download file")
 
     tcr_parser = TCRParser()
-    tcr = list(tcr_parser.get_tcr_structure(pdb_id, filename).get_TCRs())
+    tcrs = list(tcr_parser.get_tcr_structure(pdb_id, filename, **kwargs).get_TCRs())
     os.remove(filename)
-    if len(tcr) == 1:
-        return tcr[0]
-    elif len(tcr) == 0:
+
+    if len(tcrs) == 0:
         warnings.warn(f"No TCRs identified in {pdb_id}")
-        return None
-    else:
-        return tcr
+
+    return tcrs