Add sequence reconstruction

[giulioisac]

Add something along this to sonia:

from sonnia import Sonia
from sonnia.utils import nt2aa,gene_to_num_str
import numpy as np
def reconstruct_sequence(sonia_model:Sonia,
                         ntjunction:str,
                         V:str,
                         J:str)->str:
    """Return full sequence and CDR3 anchor position"""
    try:
        V = sonia_model.pgen_model.V_mask_mapping[gene_to_num_str(V, "V")][0]
        J = sonia_model.pgen_model.J_mask_mapping[gene_to_num_str(J, "J")][0]
    except:
        return np.nan
    fullV_gene = sonia_model.genomic_data.genV[V][2]
    endV = -len(sonia_model.genomic_data.genV[V][1])
    begin = fullV_gene[: endV]
    fullJ_gene = sonia_model.genomic_data.genJ[J][2]
    beginJ = len(sonia_model.genomic_data.genJ[J][1])
    end = fullJ_gene[beginJ:]
    ntseq = begin + ntjunction + end
    return ntseq

# generate sequences using the model
model=Sonia(ppost_model='human_B_kappa')
seqs=model.generate_sequences_post(int(10),nucleotide=True)

# reconstruct sequences
reconstructed_seqs=[reconstruct_sequence(model,seq[3],seq[1],seq[2]) for seq in seqs]
reconstructed_seqs_aa=[nt2aa(seq) for seq in reconstructed_seqs]

[ZhaochenYe999]

Hi, is there a command line implementation of reconstructing the full sequence? Really appreciate your help!

[giulioisac]

Hello, this script should work for you:

"""
Standalone script: given a CDR3 junction (nucleotide *or* amino-acid), a V gene,
a J gene, and a sonnia/sonia model, return the full reconstructed sequence by
stitching the germline V/J flanks around the junction.

Example — nucleotide junction (prints both full nt and full aa):

    python reconstruct_sequence_demo.py \
        --model human_B_kappa \
        --junction TGTCAGCAGGATTATAACTTACTCTCCTTCACTTTC \
        --v IGKV3-7*04 --j IGKJ4*01

    # nt: GAAATTGTAATGACACAGTCT...TTCACTTTCGGCGGAGGGACCAAGGTGGAGATCAAAC
    # aa: EIVMTQSPPTLSLSPGERVTLSC...QQDYNLLSFTFGGGTKVEIK

Example — amino-acid junction (prints full aa only):

    python reconstruct_sequence_demo.py \
        --model human_B_kappa \
        --junction CQQYGSSPLTF --v IGKV3-20*01 --j IGKJ2*01

    # aa: EIVLTQSPGTLSLSPGERATLSC...QQYGSSPLTFGQGTKLEIK

Any of the default sonnia models works for --model (e.g. human_B_heavy,
human_T_beta, mouse_T_alpha, ...).
"""

import argparse

from olga.utils import nt2aa

from sonnia.utils import define_pgen_model, gene_to_num_str


def _resolve_gene(pgen_model, gene, gene_type):
    mapping = (pgen_model.V_mask_mapping if gene_type == "V"
               else pgen_model.J_mask_mapping)
    try:
        return mapping[gene_to_num_str(gene, gene_type)][0]
    except KeyError as err:
        raise ValueError(f"{gene_type} gene {gene!r} not found in model") from err


def _v_flank_nt(genomic_data, V):
    full = genomic_data.genV[V][2]
    cdr3_overlap = genomic_data.genV[V][1]
    return full[: -len(cdr3_overlap)] if cdr3_overlap else full


def _j_flank_nt(genomic_data, J):
    full = genomic_data.genJ[J][2]
    cdr3_overlap = genomic_data.genJ[J][1]
    return full[len(cdr3_overlap):]


def reconstruct_sequence(genomic_data, pgen_model, junction: str, V: str, J: str):
    """Return (full_nt, full_aa). full_nt is None when the junction is AA-only."""
    V = _resolve_gene(pgen_model, V, "V")
    J = _resolve_gene(pgen_model, J, "J")

    v_nt = _v_flank_nt(genomic_data, V)
    j_nt = _j_flank_nt(genomic_data, J)

    is_nt = set(junction.upper()) <= set("ACGTN")

    if is_nt:
        full_nt = v_nt + junction.upper() + j_nt
        usable = full_nt[: len(full_nt) - len(full_nt) % 3]
        return full_nt, nt2aa(usable)

    v_aa = nt2aa(v_nt[: len(v_nt) - len(v_nt) % 3])
    j_aa = nt2aa(j_nt[: len(j_nt) - len(j_nt) % 3])
    return None, v_aa + junction + j_aa


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True,
                        help="sonnia model name, e.g. human_B_kappa")
    parser.add_argument("--junction", required=True,
                        help="CDR3 junction (nucleotide or amino-acid)")
    parser.add_argument("--v", required=True, help="V gene, e.g. IGKV1-5*01")
    parser.add_argument("--j", required=True, help="J gene, e.g. IGKJ5*01")
    args = parser.parse_args()

    genomic_data, _gm, pgen_model, _sg, _norm = define_pgen_model(
        args.model, compute_norm=False)
    full_nt, full_aa = reconstruct_sequence(
        genomic_data, pgen_model, args.junction, args.v, args.j)

    if full_nt is not None:
        print(f"nt: {full_nt}")
    print(f"aa: {full_aa}")


if __name__ == "__main__":
    main()

[ZhaochenYe999]

Thank you very much!