immunebuilder/ImmuneBuilder/TCRBuilder2.py

import torch
import numpy as np
import os
import sys
import argparse
from ImmuneBuilder.models import StructureModule
from ImmuneBuilder.util import get_encoding, to_pdb, find_alignment_transform, download_file, sequence_dict_from_fasta, add_errors_as_bfactors, are_weights_ready
from ImmuneBuilder.refine import refine
from ImmuneBuilder.sequence_checks import number_sequences

embed_dim = {
    "tcr2_plus_model_1":256,
    "tcr2_plus_model_2":256,
    "tcr2_plus_model_3":128,
    "tcr2_plus_model_4":128,
    "tcr2_model_1":128,
    "tcr2_model_2":128,
    "tcr2_model_3":256,
    "tcr2_model_4":256,
}

model_urls = {
    "tcr2_plus_model_1": "https://zenodo.org/records/10892159/files/tcr_model_1?download=1",
    "tcr2_plus_model_2": "https://zenodo.org/record/10892159/files/tcr_model_2?download=1",
    "tcr2_plus_model_3": "https://zenodo.org/record/10892159/files/tcr_model_3?download=1",
    "tcr2_plus_model_4": "https://zenodo.org/record/10892159/files/tcr_model_4?download=1",
    "tcr2_model_1": "https://zenodo.org/record/7258553/files/tcr_model_1?download=1",
    "tcr2_model_2": "https://zenodo.org/record/7258553/files/tcr_model_2?download=1",
    "tcr2_model_3": "https://zenodo.org/record/7258553/files/tcr_model_3?download=1",
    "tcr2_model_4": "https://zenodo.org/record/7258553/files/tcr_model_4?download=1",
}

class TCR:
    def __init__(self, numbered_sequences, predictions, header=None):
        self.numbered_sequences = numbered_sequences
        self.atoms = [x[0] for x in predictions]
        self.encodings = [x[1] for x in predictions]
        if header is not None:
            self.header = header
        else:
            self.header = "REMARK  TCR STRUCTURE MODELLED USING TCRBUILDER2+                               \n"

        with torch.no_grad():
            traces = torch.stack([x[:,0] for x in self.atoms])
            self.R,self.t = find_alignment_transform(traces)
            self.aligned_traces = (traces-self.t) @ self.R
            self.error_estimates = (self.aligned_traces - self.aligned_traces.mean(0)).square().sum(-1)
            self.ranking = [x.item() for x in self.error_estimates.mean(-1).argsort()]
        

    def save_single_unrefined(self, filename, index=0):
        atoms = (self.atoms[index] - self.t[index]) @ self.R[index]
        unrefined = to_pdb(self.numbered_sequences, atoms, chain_ids="BA")

        with open(filename, "w+") as file:
            file.write(unrefined)


    def save_all(self, dirname=None, filename=None, check_for_strained_bonds=True, n_threads=-1):
        if dirname is None:
            dirname="TCRBuilder2_output"
        if filename is None:
            filename="final_model.pdb"
        os.makedirs(dirname, exist_ok = True)

        for i in range(len(self.atoms)):

            unrefined_filename = os.path.join(dirname,f"rank{i}_unrefined.pdb")
            self.save_single_unrefined(unrefined_filename, index=self.ranking[i])

        np.save(os.path.join(dirname,"error_estimates"), self.error_estimates.mean(0).cpu().numpy())
        final_filename = os.path.join(dirname, filename)
        refine(os.path.join(dirname,"rank0_unrefined.pdb"), final_filename, check_for_strained_bonds=check_for_strained_bonds, n_threads=n_threads)
        add_errors_as_bfactors(final_filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), header=[self.header])


    def save(self, filename=None, check_for_strained_bonds=True, n_threads=-1):
        if filename is None:
            filename = "TCRBuilder2_output.pdb"

        for i in range(len(self.atoms)):
            self.save_single_unrefined(filename, index=self.ranking[i])
            success = refine(filename, filename, check_for_strained_bonds=check_for_strained_bonds, n_threads=n_threads)
            if success:
                break
            else:
                self.save_single_unrefined(filename, index=self.ranking[i])
                success = refine(filename, filename, check_for_strained_bonds=check_for_strained_bonds, n_threads=n_threads)
                if success:
                    break

        if not success:
            print(f"FAILED TO REFINE {filename}.\nSaving anyways.", flush=True)
        add_errors_as_bfactors(filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), header=[self.header])


class TCRBuilder2:
    def __init__(self, model_ids = [1,2,3,4], weights_dir=None, numbering_scheme='imgt', use_TCRBuilder2_PLUS_weights=True):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.scheme = numbering_scheme
        self.use_TCRBuilder2_PLUS_weights = use_TCRBuilder2_PLUS_weights
        if weights_dir is None:
            weights_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "trained_model")

        self.models = {}
        for id in model_ids:
            if use_TCRBuilder2_PLUS_weights:
                model_file = f"tcr_model_{id}"
                model_key = f'tcr2_plus_model_{id}'
            else:
                model_file = f"tcr2_model_{id}"
                model_key = f'tcr2_model_{id}'
            model = StructureModule(rel_pos_dim=64, embed_dim=embed_dim[model_key]).to(self.device)
            weights_path = os.path.join(weights_dir, model_file)

            try:
                if not are_weights_ready(weights_path):
                    print(f"Downloading weights for {model_file}...", flush=True)
                    download_file(model_urls[model_key], weights_path)

                model.load_state_dict(torch.load(weights_path, map_location=torch.device(self.device)))
            except Exception as e:
                print(f"ERROR: {model_file} not downloaded or corrupted.", flush=True)
                raise e

            model.eval()
            model.to(torch.get_default_dtype())

            self.models[model_file] = model


    def predict(self, sequence_dict):
        numbered_sequences = number_sequences(sequence_dict, scheme=self.scheme)
        sequence_dict = {chain: "".join([x[1] for x in numbered_sequences[chain]]) for chain in numbered_sequences}

        with torch.no_grad():
            sequence_dict = {"H":sequence_dict["B"], "L":sequence_dict["A"]} 
            encoding = torch.tensor(get_encoding(sequence_dict), device = self.device, dtype=torch.get_default_dtype())
            full_seq = sequence_dict["H"] + sequence_dict["L"]
            outputs = []

            for model_file in self.models:
                pred = self.models[model_file](encoding, full_seq)
                outputs.append(pred)

        if self.use_TCRBuilder2_PLUS_weights:
            header = "REMARK  TCR STRUCTURE MODELLED USING TCRBUILDER2+                               \n"
        else:
            header = "REMARK  TCR STRUCTURE MODELLED USING TCRBUILDER2                                \n"

        return TCR(numbered_sequences, outputs, header)


def command_line_interface():
    description="""
        TCRBuilder2                                                 || ||
        A Method for T-Cell Receptor Structure Prediction          |VB|VA| 
        Author: Brennan Abanades Kenyon, Nele Quast                |CB|CA|
        Supervisor: Charlotte Deane                             -------------

        By default TCRBuilder2 will use TCRBuilder2+ weights.
        To use the original weights, see options.
    """

    parser = argparse.ArgumentParser(prog="TCRBuilder2", description=description, formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("-B", "--beta_sequence", help="Beta chain amino acid sequence", default=None)
    parser.add_argument("-A", "--alpha_sequence", help="Alpha chain amino acid sequence", default=None)
    parser.add_argument("-f", "--fasta_file", help="Fasta file containing a beta and alpha chain named B and A", default=None)

    parser.add_argument("-o", "--output", help="Path to where the output model should be saved. Defaults to the same directory as input file.", default=None)
    parser.add_argument("--to_directory", help="Save all unrefined models and the top ranked refined model to a directory. " 
    "If this flag is set the output argument will be assumed to be a directory", default=False, action="store_true")
    parser.add_argument("--n_threads", help="The number of CPU threads to be used. If this option is set, refinement will be performed on CPU instead of GPU. By default, all available cores will be used.", type=int, default=-1)
    parser.add_argument("-u", "--no_sidechain_bond_check", help="Don't check for strained bonds. This is a bit faster but will rarely generate unphysical side chains", default=False, action="store_true")
    parser.add_argument("-v", "--verbose", help="Verbose output", default=False, action="store_true")
    parser.add_argument("-og", "--original_weights", help="use original TCRBuilder2 weights instead of TCRBuilder2+", default=False, action="store_true")

    args = parser.parse_args()

    if (args.beta_sequence is not None) and (args.alpha_sequence is not None):
        seqs = {"B":args.beta_sequence, "A":args.alpha_sequence}
    elif args.fasta_file is not None:
        seqs = sequence_dict_from_fasta(args.fasta_file)
    else:
        raise ValueError("Missing input sequences")

    check_for_strained_bonds = not args.no_sidechain_bond_check
    
    if args.n_threads > 0:
        torch.set_num_threads(args.n_threads)    

    if args.verbose:
        print(description, flush=True)
        print(f"Sequences loaded succesfully.\nAlpha and Beta chains are:", flush=True)
        [print(f"{chain}: {seqs[chain]}", flush=True) for chain in "AB"]
        print("Running sequences through deep learning model...", flush=True)

    try:
        tcr = TCRBuilder2(use_TCRBuilder2_PLUS_weights=not args.original_weights).predict(seqs)
    except AssertionError as e:
        print(e, flush=True)
        sys.exit(1)
    
    if args.verbose:
        print("TCR modelled succesfully, starting refinement.", flush=True)

    if args.to_directory:
        tcr.save_all(args.output,check_for_strained_bonds=check_for_strained_bonds, n_threads=args.n_threads)
        if args.verbose:
            print("Refinement finished. Saving all outputs to directory", flush=True)
    else:
        tcr.save(args.output,check_for_strained_bonds=check_for_strained_bonds, n_threads=args.n_threads)
        if args.verbose:
            outfile = "TCRBuilder2_output.pdb" if args.output is None else args.output
            print(f"Refinement finished. Saving final structure to {outfile}", flush=True)