Metadata-Version: 2.4
Name: tcr-pmhc-analyzer
Version: 0.2.1
Summary: Analyzer for TCR-pMHC binding predictor outputs
Project-URL: Homepage, https://github.com/qbic-pipelines/tcr-pmhc-analyzer
Project-URL: Repository, https://github.com/qbic-pipelines/tcr-pmhc-analyzer.git
Project-URL: Issues, https://github.com/qbic-pipelines/tcr-pmhc-analyzer/issues
Author: Mark Polster
License-Expression: MIT
License-File: LICENSE
Keywords: analysis,peptide,tcr,tcr-pMHC,tcr-peptide,tcr-pmhc
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Requires-Python: >=3.10
Requires-Dist: matplotlib>=3.8.0
Requires-Dist: pandas>=3.0.0
Requires-Dist: rich-click>=1.9.6
Requires-Dist: scikit-learn>=1.4.0
Provides-Extra: dev
Requires-Dist: ruff==0.14.14; extra == 'dev'
Description-Content-Type: text/markdown

# tcr-pmhc-analyzer

[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![PyPI version](https://img.shields.io/pypi/v/tcr-pmhc-analyzer.svg)](https://pypi.org/project/tcr-pmhc-analyzer/)

Analyzer for TCR-pMHC binding predictor outputs. Merges predictions from multiple models into a unified table, detects data leakage against bundled training sets, identifies seen/unseen peptides, and benchmarks model performance via ROC curves.

## Installation

```bash
pip install tcr-pmhc-analyzer
```

For development:

```bash
git clone https://github.com/qbic-pipelines/tcr-pmhc-analyzer.git
cd tcr-pmhc-analyzer
pip install -e ".[dev]"
```

## Input format

Both commands accept a TSV configuration file with two required columns (`model`, `file`) and one optional column (`training_data`):

| Column          | Required | Description                                                                                  |
|-----------------|----------|----------------------------------------------------------------------------------------------|
| `model`         | Yes      | Name of the prediction model                                                                 |
| `file`          | Yes      | Path to the model's prediction output                                                        |
| `training_data` | No       | Path to a training-data file. Only needed for **custom** (non-built-in) models, or to override a built-in model's bundled training data for leakage detection. Omit the column entirely when using only built-in models. |

Example `input.tsv`:

```tsv
model	file
ergo2	results/ergo2_predictions.csv
mixtcrpred	results/mixtcrpred_predictions.csv
t2pmhc-gcn	results/t2pmhc_gcn_predictions.csv
```

Each prediction file must contain the following columns:

| Column         | Description                                      |
|----------------|--------------------------------------------------|
| `identifier`   | Unique sample identifier (used for merging)       |
| `binding_score` | Model's predicted binding score                  |
| `binder`       | Ground truth label (0/1), required for benchmarking |
| `peptide`      | Peptide sequence                                  |
| `cdr3a`        | CDR3 alpha chain sequence                         |
| `cdr3b`        | CDR3 beta chain sequence                          |
| `va`, `vb`     | V gene alpha/beta                                 |
| `ja`, `jb`     | J gene alpha/beta                                 |
| `mhc`          | MHC allele                                        |
| `organism`     | Source organism                                   |
| `mhc_class`    | MHC class                                         |

## Commands

### `create-analyzer-table`

Merges predictions from multiple models into a single table with rank-normalized scores, data leakage annotations, and seen-peptide flags.

```bash
tcr-pmhc-analyzer create-analyzer-table [OPTIONS]
```

| Option              | Short | Required | Description                                                 |
|---------------------|-------|----------|-------------------------------------------------------------|
| `--input PATH`      | `-i`  | Yes      | Path to TSV config file with `model` and `file` columns     |
| `--output PATH`     | `-o`  | Yes      | Output file path (`.csv` or `.tsv`)                         |
| `--ergo-version`    |       | If ergo2 | ERGO training data version: `vdjdb`, `mcpas`, or `both`     |

**Example:**

```bash
tcr-pmhc-analyzer create-analyzer-table \
  -i input.tsv \
  -o analyzer_table.csv \
  --ergo-version vdjdb
```

**Output columns added:**
- `binding_score_{model}` — raw binding score per model
- `rank_score_{model}` — rank-normalized score in [0, 1] (1 = highest)
- `sample_in_train_{model}` — `True` if the sample appears in the model's training data (data leakage)
- `seen_in_{model}` — `True` if the peptide was seen in the model's training data

### `benchmark`

Generates ROC curve plots comparing model performance, split by seen vs unseen peptides. Data leakage samples are automatically removed before analysis.

```bash
tcr-pmhc-analyzer benchmark [OPTIONS]
```

| Option              | Short | Required | Description                                                        |
|---------------------|-------|----------|--------------------------------------------------------------------|
| `--input PATH`      | `-i`  | *        | Path to TSV config file with `model` and `file` columns            |
| `--table PATH`      |       | *        | Path to a pre-created analyzer table (alternative to `--input`)    |
| `--output PATH`     | `-o`  | Yes      | Output directory for ROC curve plots                               |
| `--ergo-version`    |       | If ergo2 | ERGO training data version: `vdjdb`, `mcpas`, or `both`             |
| `--models`          | `-m`  | No       | Space-separated list of models to benchmark (default: all available)|

\* Either `--input` or `--table` must be provided.

**Examples:**

```bash
# Benchmark from raw predictions
tcr-pmhc-analyzer benchmark -i input.tsv -o results/

# Benchmark from a pre-created analyzer table
tcr-pmhc-analyzer benchmark --table analyzer_table.csv -o results/

# Benchmark specific models only
tcr-pmhc-analyzer benchmark -i input.tsv -o results/ -m "ergo2 mixtcrpred tabr-bert"
```

**Output files:**
- `roc_curve_unseen.png` — ROC curves for peptides unseen by all selected models
- `roc_curve_seen.png` — ROC curves for peptides seen by all selected models

## Supported models

| Model          | Training data                          |
|----------------|----------------------------------------|
| `ergo2`        | mcpas or vdjdb (specify with `--ergo-version vdjdb\|mcpas`) |
| `ergo2-vdjdb`  | vdjdb (use with `--ergo-version both`, alongside `ergo2-mcpas`) |
| `ergo2-mcpas`  | mcpas (use with `--ergo-version both`, alongside `ergo2-vdjdb`) |
| `mixtcrpred`   | 146 pMHC training set                  |
| `t2pmhc-gcn`   | t2pmhc core training set               |
| `t2pmhc-gat`   | t2pmhc core training set               |
| `tabr-bert`    | TCR-pMHC training set                  |
| `tulip-tcr`    | TULIP training set                     |
| `atm-tcr`      | ATM-TCR training set                   |
| `nettcr_2_2`   | NetTCR_2_2 training set                |
| `tcren`        | TCRen PDB-based structural training set |
| `stag-llm`     | STAG-LLM structural training set (`final_dataset_modeled.csv`) |

> **Note on score direction.** Most models output `binding_score` such that
> higher = more likely binder. `tcren` is the exception: lower (more negative)
> score = more likely binder. The analyzer is aware of this and inverts
> direction internally for `rank_score_tcren` and for ROC/PR benchmarking, so
> `rank_score_{model}` (1.0 = strongest binder) and AUC/AUPRC are directly
> comparable across all models.

> **Note on the `stag-llm` training set.** STAG-LLM trains in two stages: a
> self-supervised ESM-2 masked-language-model pretraining on a large
> sequence-only corpus (`full_seq_df_new.csv`, binders only, no labels), then
> the supervised binding classifier — the deployed checkpoint — trained with
> labels only on the structurally-modeled set (`final_dataset_modeled.csv`,
> 46,201 pairs). The bundled `stag-llm_train.csv` is that **supervised
> structural** set, since binding-benchmark leakage means the labeled pairs the
> model actually learned; the MLM pretraining corpus is treated like ESM-2's own
> pretraining and is **not** counted as label leakage. To experiment with the
> full sequence corpus later, point `stag-llm`'s `training_data` column in the
> input samplesheet at an alternate file — no code change needed.

### ERGO-II "both" variant mode

The upstream [`tcrpmhcbinding`](https://github.com/qbic-pipelines/tcrpmhcbinding) Nextflow pipeline can run ERGO-II against both training-database variants (`vdjdb` and `mcpas`) in a single run by passing `--ergo2_variant both`. In that mode the pipeline emits two prediction columns labelled `ergo2-vdjdb` and `ergo2-mcpas` instead of the single `ergo2` column.

When you pass that predictions TSV to this tool, use `--ergo-version both`:

```bash
tcr-pmhc-analyzer create-analyzer-table \
  -i input.tsv \
  -o analyzer_table.csv \
  --ergo-version both
```

Each ERGO column is then checked against its respective training set:
- `sample_in_train_ergo2-vdjdb` and `seen_in_ergo2-vdjdb` are computed against the vdjdb training data.
- `sample_in_train_ergo2-mcpas` and `seen_in_ergo2-mcpas` are computed against the mcpas training data.

`--ergo-version both` requires the input config to contain both `ergo2-vdjdb` and `ergo2-mcpas` model rows. Single-variant inputs continue to use `--ergo-version vdjdb` or `--ergo-version mcpas` with the legacy `ergo2` model name; that path is unchanged.

## How it works

1. **Merge**: Prediction outputs from multiple models are merged on the `identifier` column into a single DataFrame.
2. **Rank normalization**: Each model's `binding_score` is rank-normalized to [0, 1] using descending order with average tie-breaking. NaN values are preserved.
3. **Data leakage detection**: Each sample is checked against bundled training data to flag samples that appear in a model's training set. The match rule is **per-model** because each training file exposes the TCR-pMHC tuple at a different granularity — see [Data leakage detection](#data-leakage-detection) below for the per-model rules.
4. **Seen peptide detection**: Each peptide is checked against training data to identify whether it was seen during model training.
5. **Benchmarking**: ROC curves are generated after removing leaked samples, separately for seen and unseen peptides.

### Data leakage detection

`sample_in_train_{model}` is computed by a per-model matcher in
`Analyzer._find_leakage_<model>`. Each function does its own normalization
(V/J `*NN` allele stripping, MHC field truncation, etc.) and its own
`pd.merge`. The rules are derived from what each bundled training file
actually ships (column presence, NaN rates, MHC notation), not from the
underlying model paper, so the leakage check reflects the granularity the
training file was actually curated at.

| model                          | match keys                                                                                   | notes                                                                                                                                                                  |
|--------------------------------|----------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `t2pmhc-gcn`, `t2pmhc-gat`     | `(peptide, cdr3a, cdr3b, va, vb, ja, jb, mhc)` with 2-field MHC                              | Default rule. Training is complete and clean (100% non-NaN, V/J carry the `*` allele).                                                                                 |
| `mixtcrpred`                   | as above, **with `*NN` stripped from V/J** on both sides                                     | Training file is internally mixed on the V/J allele suffix (~95.5% stripped, ~4.5% retained). Mirrors `benchmarker.translate_mixtcrpred`.                              |
| `ergo2_vdjdb`                  | full-chain **OR** TRB-only `(peptide, cdr3b, vb, jb, mhc)`; MHC at first colon field          | 82.7% of vdjdb training rows have NaN α-side fields. The TRB-only union catches them. MHC first-field handles the mix of `HLA-X*NN` and `HLA-X*NN:NN`.                  |
| `ergo2_mcpas`                  | as `ergo2_vdjdb`, **with `*NN` stripped from V/J**                                           | mcpas stores V/J at gene level (<3% have the allele suffix) and uses short-form MHC (`HLA-A2`). Without the strip the merge produces zero matches.                     |
| `tabr-bert`                    | `(peptide, cdr3b, mhc)` with 2-field MHC                                                     | β-only training. Default rule auto-drops the absent α-side / V/J columns.                                                                                              |
| `tulip-tcr`                    | `(peptide, cdr3b)` **OR** `(peptide, cdr3a)`. **MHC excluded.**                              | Training `mhc` column is ~40% non-allele text (cell lines, free text, malformed alleles); MHC-conditioned matching would over-reject. Lossy on MHC, high recall on TCR. |
| `atm-tcr`                      | `(peptide, cdr3b)`                                                                            | atm-tcr training has no MHC column and no α-side. Default rule auto-drops everything except these two.                                                                 |
| `tcren`                        | `(peptide, cdr3a, cdr3b)` **OR** `(peptide, cdr3b)`. **MHC excluded.**                       | Training stores only `mhc_class` ("MHCI"/"MHCII"), too coarse to discriminate. The fallback group handles the 1.4% of training rows with NaN cdr3a.                    |
| `stag-llm`                     | `(peptide, cdr3a, cdr3b)`. **MHC excluded.**                                                 | Default rule. Training ships only those three columns (source MHC is a raw ~180 aa sequence, dropped upstream), so the 8 default keys auto-drop to the triple. Equivalent to the tcren rule; no `(peptide, cdr3b)` fallback is needed (no NaN cdr3a). |
| `nettcr_peptide`, `nettcr_pan` | `(peptide, mhc)` only — known limitation                                                     | NetTCR-2.2 training file uses `anarcii_*` CDR sequences, no V/J. Test tables don't carry `anarcii_*`. The resulting flag is advisory only.                              |

Several of these rules mirror the reference implementations in `benchmarker.find_leakage_*` (from the t2pmhc revisions benchmark pipeline at `t2pmhc_revisions/benchmark_m3/benchmarker.py`); the others are derived directly from the training-file profile.

**Regeneration note.** Per-model rules changed in `[dev]` — see [CHANGELOG.md](CHANGELOG.md). Analyzer tables generated before that change have stale `sample_in_train_*` columns for `mixtcrpred`, `ergo2_vdjdb`, `ergo2_mcpas`, `tulip-tcr`, `tcren`, `nettcr_peptide`, and `nettcr_pan` and should be regenerated.

## Citations

If you use tcr-pmhc-analyzer in your research, please cite the underlying prediction models:

**ATM-TCR**
> Cai, M. et al. (2022). ATM-TCR: TCR-Epitope Binding Affinity Prediction Using a Multi-Head Self-Attention Model. *Frontiers in Immunology*, 13, 893247. https://doi.org/10.3389/fimmu.2022.893247

**ERGO-II**
> Springer, I. et al. (2021). Contribution of T Cell Receptor Alpha and Beta CDR3, MHC Typing, V and J Genes to Peptide Binding Prediction. *Frontiers in Immunology*, 12, 664514. https://doi.org/10.3389/fimmu.2021.664514

**MIXTCRpred**
> Croce, G. et al. (2024). Deep learning predictions of TCR-epitope interactions reveal epitope-specific chains in dual alpha T cells. *Nature Communications*, 15, 3211. https://doi.org/10.1038/s41467-024-47461-8

**t2pmhc**
> Polster, M. et al. (2026). t2pmhc: A Structure-Informed Graph Neural Network to Predict TCR-pMHC Binding. *bioRxiv*. https://doi.org/10.64898/2026.02.27.708137

**TABR-BERT**
> Zhang, J. et al. (2024). Accurate TCR-pMHC interaction prediction using a BERT-based transfer learning method. *Briefings in Bioinformatics*, 25(1), bbad436. https://doi.org/10.1093/bib/bbad436

**TULIP**
> Meynard-Piganeau, B. et al. (2024). TULIP — a Transformer-based Unsupervised Language model for Interacting Peptides and T-cell receptors that generalizes to unseen epitopes. *Proceedings of the National Academy of Sciences*, 121(13). https://doi.org/10.1073/pnas.2316401121

**NetTCR_2_2**
> Jensen, M. F., & Nielsen, M. (2023). NetTCR 2.2-Improved TCR specificity predictions by combining pan-and peptide-specific training strategies, loss-scaling and integration of sequence similarity. *bioRxiv*, 2023-10. https://www.biorxiv.org/content/10.1101/2023.10.12.562001v2.abstract

**TCRen**
> Karnaukhov, V. K. et al. (2024). Structure-based prediction of T cell receptor recognition of unseen epitopes using TCRen. *Nature Computational Science*, 4(7), 510–521. https://doi.org/10.1038/s43588-024-00653-0

## License

[MIT](LICENSE)