feat: port RNA-annotation + fusion-prep scripts from feature/rna-integration

Ports two self-contained additions (no changes to existing scripts) from the
modules repo feature/rna-integration branch:
- annotate_rna.py: annotate neoantigens with RNA data (reads rna_* columns from
  the MAF, kallisto abundance + GTF for expression/VAF) -> console annotate_rna.py
- prepare_fusion_fasta.py: build fusion neoantigen FASTAs from AGFusion output
  -> console prepare_fusion_fasta.py
Both with their pytest suites (imports repointed to the package). pytest 101 passed.

Author: johnoooh

pyproject.toml

@@ -35,6 +35,8 @@ dev = ["pytest"]
 "convertannotjson.py" = "neoantigen_utils.convertannotjson:main"
 "format_netmhcpan_output.py" = "neoantigen_utils.format_netmhcpan_output:console"
 "generateHLAString.sh" = "neoantigen_utils.hla_string:console"
+"annotate_rna.py" = "neoantigen_utils.annotate_rna:main"
+"prepare_fusion_fasta.py" = "neoantigen_utils.prepare_fusion_fasta:main"
 
 [tool.hatch.build.targets.wheel]
 packages = ["src/neoantigen_utils"]

src/neoantigen_utils/annotate_rna.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+"""Annotate neoantigen TSV with RNA-seq data from FORTE outputs."""
+
+import argparse
+import gzip
+import re
+import sys
+
+import pandas as pd
+
+VERSION = "1.0.0"
+
+
+def build_mutation_id(row):
+    """Reconstruct mutation_id from MAF columns.
+
+    Matches the format used by generate_input.py:
+    - SNP/DNP/TNP: chr_pos_ref_alt
+    - DEL: chr_pos_ref_D
+    - INS: chr_pos_I_alt
+    """
+    chrom = str(row["Chromosome"])
+    pos = str(row["Start_Position"])
+    ref = str(row["Reference_Allele"])
+    alt = str(row["Tumor_Seq_Allele2"])
+    vtype = str(row["Variant_Type"])
+
+    if vtype == "DEL":
+        return f"{chrom}_{pos}_{ref}_D"
+    elif vtype == "INS":
+        return f"{chrom}_{pos}_I_{alt}"
+    else:
+        return f"{chrom}_{pos}_{ref}_{alt}"
+
+
+def extract_rna_columns_from_maf(maf_df):
+    """Extract RNA VAF columns from MAF if present.
+
+    Returns DataFrame with mutation_id, rna_alt_count, rna_ref_count, rna_vaf.
+    If rna columns are absent, returns DataFrame with NaN values.
+    """
+    maf_df = maf_df.copy()
+    maf_df["mutation_id"] = maf_df.apply(build_mutation_id, axis=1)
+
+    has_rna = "rna_t_alt_count" in maf_df.columns
+
+    result = pd.DataFrame({"mutation_id": maf_df["mutation_id"]})
+
+    if has_rna:
+        result["rna_alt_count"] = maf_df["rna_t_alt_count"].values
+        result["rna_ref_count"] = maf_df["rna_t_ref_count"].values
+        result["rna_vaf"] = maf_df["rna_t_variant_frequency"].values
+    else:
+        result["rna_alt_count"] = pd.NA
+        result["rna_ref_count"] = pd.NA
+        result["rna_vaf"] = pd.NA
+
+    return result
+
+
+def parse_gtf_gene_map(gtf_path):
+    """Parse GTF to build transcript_id -> gene_name mapping."""
+    tx2gene = {}
+    opener = gzip.open if gtf_path.endswith(".gz") else open
+    with opener(gtf_path, "rt") as f:
+        for line in f:
+            if line.startswith("#"):
+                continue
+            fields = line.strip().split("\t")
+            if len(fields) < 9:
+                continue
+            attrs = fields[8]
+            tx_match = re.search(r'transcript_id "([^"]+)"', attrs)
+            gene_match = re.search(r'gene_name "([^"]+)"', attrs)
+            if tx_match and gene_match:
+                tx_id = tx_match.group(1).split(".")[0]
+                gene_name = gene_match.group(1)
+                tx2gene[tx_id] = gene_name
+    return tx2gene
+
+
+def annotate_expression(abundance_path, gtf_path):
+    """Parse Kallisto abundance and map to gene-level TPM.
+
+    Returns DataFrame with Gene, rna_tpm columns.
+    Returns None if abundance_path is None.
+    """
+    if abundance_path is None:
+        return None
+
+    abundance = pd.read_csv(abundance_path, sep="\t")
+    tx2gene = parse_gtf_gene_map(gtf_path)
+
+    abundance["transcript_id"] = abundance["target_id"].str.split(".").str[0]
+    abundance["Gene"] = abundance["transcript_id"].map(tx2gene)
+
+    gene_tpm = (
+        abundance.dropna(subset=["Gene"])
+        .groupby("Gene")["tpm"]
+        .sum()
+        .reset_index()
+        .rename(columns={"tpm": "rna_tpm"})
+    )
+
+    return gene_tpm
+
+
+def merge_annotations(neo_tsv, rna_vaf_df, expression_df, tpm_threshold=1.0):
+    """Merge RNA annotations onto neoantigen TSV.
+
+    Returns annotated DataFrame with rna_tpm, rna_vaf, rna_alt_count,
+    rna_ref_count, rna_expressed columns.
+    """
+    result = neo_tsv.copy()
+
+    if rna_vaf_df is not None:
+        result = result.merge(rna_vaf_df, on="mutation_id", how="left")
+    else:
+        result["rna_alt_count"] = pd.NA
+        result["rna_ref_count"] = pd.NA
+        result["rna_vaf"] = pd.NA
+
+    if expression_df is not None:
+        result = result.merge(expression_df, on="Gene", how="left")
+    else:
+        result["rna_tpm"] = pd.NA
+
+    tpm_ok = result["rna_tpm"].notna() & (result["rna_tpm"] > tpm_threshold)
+    vaf_ok = result["rna_alt_count"].isna() | (result["rna_alt_count"] > 0)
+    result["rna_expressed"] = tpm_ok & vaf_ok
+
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Annotate neoantigens with RNA data")
+    parser.add_argument("--neoantigen_tsv", required=True, help="Neoantigen TSV from convertannotjson")
+    parser.add_argument("--maf", required=True, help="Input MAF (may contain rna_* columns)")
+    parser.add_argument("--kallisto_abundance", default=None, help="Kallisto abundance.tsv (optional)")
+    parser.add_argument("--gtf", default=None, help="GTF file for transcript-to-gene mapping")
+    parser.add_argument("--tpm_threshold", type=float, default=1.0, help="TPM threshold for expressed")
+    parser.add_argument("--output_annotated", required=True, help="Output annotated TSV")
+    parser.add_argument("--output_report", required=True, help="Output detailed RNA report")
+    parser.add_argument("-v", "--version", action="version", version=f"%(prog)s {VERSION}")
+    args = parser.parse_args()
+
+    neo_tsv = pd.read_csv(args.neoantigen_tsv, sep="\t")
+    maf_df = pd.read_csv(args.maf, sep="\t", comment="#")
+
+    rna_vaf_df = extract_rna_columns_from_maf(maf_df)
+    expression_df = annotate_expression(args.kallisto_abundance, args.gtf)
+
+    annotated = merge_annotations(neo_tsv, rna_vaf_df, expression_df, args.tpm_threshold)
+    annotated.to_csv(args.output_annotated, sep="\t", index=False)
+
+    report_cols = ["mutation_id", "Gene", "rna_tpm", "rna_vaf", "rna_alt_count", "rna_ref_count", "rna_expressed"]
+    available_cols = [c for c in report_cols if c in annotated.columns]
+    report = annotated[available_cols].drop_duplicates()
+    report.to_csv(args.output_report, sep="\t", index=False)
+
+
+if __name__ == "__main__":
+    main()

src/neoantigen_utils/prepare_fusion_fasta.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+"""Convert AGFusion output to pipeline-compatible FASTA format."""
+
+import argparse
+import os
+
+VERSION = "1.0.0"
+
+
+def parse_agfusion_dir(agfusion_dir):
+    """Parse AGFusion output directory for fusion protein sequences.
+
+    AGFusion creates subdirectories per fusion event, each containing
+    *_protein.fa files with fusion protein sequences.
+    """
+    fusions = []
+    for entry in sorted(os.listdir(agfusion_dir)):
+        subdir = os.path.join(agfusion_dir, entry)
+        if not os.path.isdir(subdir):
+            continue
+        gene_pair = entry
+        for fname in sorted(os.listdir(subdir)):
+            if fname.endswith("_protein.fa"):
+                fpath = os.path.join(subdir, fname)
+                sequences = read_fasta(fpath)
+                for seq_id, seq in sequences:
+                    fusions.append({
+                        "gene_pair": gene_pair,
+                        "transcript_pair": seq_id,
+                        "sequence": seq,
+                    })
+    return fusions
+
+
+def read_fasta(path):
+    """Read FASTA file, return list of (header, sequence) tuples."""
+    sequences = []
+    current_header = None
+    current_seq = []
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if line.startswith(">"):
+                if current_header:
+                    sequences.append((current_header, "".join(current_seq)))
+                current_header = line[1:]
+                current_seq = []
+            elif line:
+                current_seq.append(line)
+    if current_header:
+        sequences.append((current_header, "".join(current_seq)))
+    return sequences
+
+
+def extract_junction_peptides(fusion_seq, junction_pos, peptide_lengths=None):
+    """Extract peptide windows spanning the fusion junction.
+
+    For each peptide length, slide a window across the junction point.
+    Each peptide must include at least 1 AA from each fusion partner.
+    """
+    if peptide_lengths is None:
+        peptide_lengths = [9, 10, 11]
+    peptides = []
+    for plen in peptide_lengths:
+        if len(fusion_seq) < plen:
+            continue
+        start_min = max(0, junction_pos - plen + 1)
+        start_max = min(junction_pos, len(fusion_seq) - plen)
+        for start in range(start_min, start_max + 1):
+            pep = fusion_seq[start:start + plen]
+            if len(pep) == plen:
+                peptides.append(pep)
+    return peptides
+
+
+def infer_junction_position(fusion_seq, five_prime_len=None):
+    """Infer junction position from AGFusion sequence.
+
+    AGFusion marks the junction with '*' in some outputs.
+    If not present, use five_prime_len if provided, else midpoint.
+    """
+    if "*" in fusion_seq:
+        return fusion_seq.index("*")
+    if five_prime_len is not None:
+        return five_prime_len
+    return len(fusion_seq) // 2
+
+
+def write_fasta_pair(peptides, mut_path, wt_path):
+    """Write MUT and WT FASTA files for pipeline compatibility."""
+    with open(mut_path, "w") as mut_f, open(wt_path, "w") as wt_f:
+        for pep in peptides:
+            mut_f.write(f">{pep['id']}_M\n{pep['mut_seq']}\n")
+            wt_f.write(f">{pep['id']}_W\n{pep['wt_seq']}\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare fusion FASTAs for neoantigen pipeline")
+    parser.add_argument("--agfusion_dir", required=True, help="AGFusion output directory")
+    parser.add_argument("--output_prefix", required=True, help="Output file prefix")
+    parser.add_argument("--peptide_lengths", default="9,10,11", help="Comma-separated peptide lengths")
+    parser.add_argument("-v", "--version", action="version", version=f"%(prog)s {VERSION}")
+    args = parser.parse_args()
+
+    peptide_lengths = [int(x) for x in args.peptide_lengths.split(",")]
+    fusions = parse_agfusion_dir(args.agfusion_dir)
+
+    all_peptides = []
+    pep_counter = 0
+    for fusion in fusions:
+        seq = fusion["sequence"].replace("*", "")
+        junction = infer_junction_position(fusion["sequence"])
+        junction_peps = extract_junction_peptides(seq, junction, peptide_lengths)
+        for pep in junction_peps:
+            pep_counter += 1
+            all_peptides.append({
+                "id": f"{fusion['gene_pair']}_{pep_counter}",
+                "mut_seq": pep,
+                "wt_seq": pep,
+            })
+
+    mut_path = f"{args.output_prefix}.SV.MUT.fa"
+    wt_path = f"{args.output_prefix}.SV.WT.fa"
+    write_fasta_pair(all_peptides, mut_path, wt_path)
+
+
+if __name__ == "__main__":
+    main()