diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..ecdabf8 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,10 @@ +* text=auto +*.py text eol=lf +*.sh text eol=lf +*.sbatch text eol=lf +*.yml text eol=lf +*.yaml text eol=lf +*.md text eol=lf +*.tsv text eol=lf whitespace=-blank-at-eol +*.csv text eol=lf +*.json text eol=lf diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..114808f --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,18 @@ +name: CPU tests +on: + push: + pull_request: +jobs: + tests: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + cache: pip + - run: python -m pip install torch==2.4.1 --index-url https://download.pytorch.org/whl/cpu + - run: python -m pip install -r requirements-dev.txt + - run: python -m compileall -q irrm_codec benchmark scripts tests + - run: python -m pytest -q + - run: find scripts slurm benchmark/slurm -type f \( -name '*.sh' -o -name '*.sbatch' \) -print0 | xargs -0 -n1 bash -n diff --git a/.gitignore b/.gitignore index d579468..1130398 100644 --- a/.gitignore +++ b/.gitignore @@ -47,3 +47,16 @@ data/ # Reports reports/ + +# Generated experiments and local publication scratch +.publish-*/ +.tmp*/ +tmp/ +output/ +manuscript/ +*.pt +*.pth +*.npy +*.npz +experiments/**/wandb/ +experiments/**/predictions/ diff --git a/README.md b/README.md index ead456a..e7106a6 100644 --- a/README.md +++ b/README.md @@ -7,11 +7,51 @@ working with TCR CDR3 amino-acid sequences, TCRemP embeddings, and generation pr - forward model: predicts a TCRemP embedding from CDR3 sequence input - inverse model: reconstructs a CDR3 sequence from a TCRemP embedding -- joint Transformer: maps char- or WordPiece-tokenized CDR3s through one 320-D - bottleneck and jointly predicts TCRemP, `log10(pgen)`, and the original CDR3 +- joint Transformer: maps char- or WordPiece-tokenized CDR3s through a configurable latent + bottleneck (320 dimensions by default; 128 in the DATA-ANCHOR study) and jointly predicts TCRemP, `log10(pgen)`, and the original CDR3 The repository is organized as a small training package with Python entrypoints, Slurm launchers, and analysis notebooks for both single-run and multi-chain workflows. +## Canonical repository and publication workflows + +The maintained repository is [VEK239/irrm-codec](https://github.com/VEK239/irrm-codec). +It combines the original Antigenomics implementation, student contributions, and +Elizaveta Vlasova's joint RTP-CODEC and evaluation work. Original contributor +commits remain in the Git history; see [integration provenance](docs/integration_provenance.md). + +The joint codec uses **R** (sequence reconstruction), **T** (TCRemP prediction), +and **P** (generation-probability prediction) supervision. Input tokenizers include +characters, WordPiece, and sequence-only terminal anchors. Reconstruction always +uses character tokens. Latent width is configurable and is independent of the +Transformer hidden width. + +- Joint codec and matched objective/tokenizer studies: `irrm_codec/train_multitask.py`, + `benchmark/prepare_trb_joint.py`, `experiments/trb/`, and [tokenizer study](docs/anchored_tokenizer_study.md). +- Student reconstruction/Pgen benchmarks: `benchmark/prepare_splits.py`, + `benchmark/train_decoder.py`, `benchmark/train_pgen_arm.py`, and [historical report](REPORT.md). +- Author external-model, residue-PCA, latent, and VDJdb evaluations: `benchmark/` + and their matching launchers in `slurm/`. +- Historical result tables: `results/` (student runs) and `results/author/` (author runs). They describe the original runs, not a + new rerun of this integrated revision. + +For a CPU development setup without the Pgen-generation dependency: + +```bash +python -m venv .venv +# Linux/macOS: source .venv/bin/activate +# PowerShell: .\.venv\Scripts\Activate.ps1 +python -m pip install torch==2.4.1 --index-url https://download.pytorch.org/whl/cpu +python -m pip install -r requirements-dev.txt +python -m pytest -q +python -m irrm_codec.train_multitask --help +``` + +Install `requirements.txt` for OLGA/Pgen generation through mirpy. Install +`requirements-benchmark.txt` only when running the optional ESM/TCR-BERT/SCEPTR +encoders; their model weights and external datasets are obtained separately. +See [reproduction and compatibility notes](docs/publication.md) before reusing +historical checkpoints or submitting cluster examples. + ## Repository layout - `irrm_codec/`: package with data loading, tokenization, datasets, models, losses, utilities and training entrypoints @@ -48,7 +88,7 @@ Register the environment as a Jupyter kernel: python -m ipykernel install --user --name irrm-codec --display-name "Python (irrm-codec)" ``` -Project dependencies, including notebook packages, are installed from [requirements.txt](/c:/Users/lizzka239/projects/irrm-codec/requirements.txt). +Project dependencies, including notebook packages, are installed from [requirements.txt](requirements.txt). ## Input data diff --git a/REPORT.md b/REPORT.md new file mode 100644 index 0000000..05ad27b --- /dev/null +++ b/REPORT.md @@ -0,0 +1,390 @@ +# IRRM-CODEC baseline benchmark + +Comparison of IRRM-CODEC against a small set of baselines on CDR3 sequence +reconstruction and Pgen prediction, using human TRB data from +[Zenodo record 19520535](https://zenodo.org/records/19520535). + +One pretrained checkpoint per model. The pretrained encoders are frozen — only the +task heads are trained. + +## Status + +| issue | scope | state | +| --- | --- | --- | +| 1 | train/validation/test splits | done | +| 2 | sequence reconstruction | done | +| 3 | Pgen prediction | done | + +## 1. Dataset + +Source files are `trb_background_100k.tsv` and `trb_background_embeddings.parquet`, +both inside `redcea_bg.gz` on the Zenodo record. The embeddings parquet is the only +place TCRemP vectors are published, and it holds exactly 100,000 rows per chain rather +than a large background pool. + +The AIRR table arrived clean: no duplicate CDR3 amino-acid sequences, no non-standard +residues, no missing V/J calls, and nothing longer than the model's 40-residue limit. +Two filters still removed rows: + +| filter | rows | reason | +| --- | ---: | --- | +| non-functional V gene | 335 | `TRBV21-1`, `TRBV23-1`, `TRBV6-7` are pseudogenes; SCEPTR rejects them | +| Pgen = 0 | 2 | `CASSFL`, `CASTTL` give `log10(0) = -inf`, which has no finite gradient | + +Both were dropped for every representation and both targets at once, so all models see +one identical row set. Final dataset: **99,663 unique CDR3 sequences**. + +The embeddings parquet has no `clone_id`, so it is aligned to the AIRR table by row +order. Because a silent misalignment would corrupt every downstream result invisibly, +the pipeline asserts it: mean TCRemP CDR3-distance correlates with sequence length at +**r = 0.525**, against **r = 0.008** under a shuffled control. The run aborts below +r = 0.3. + +Split: one seeded permutation over de-duplicated sequences, so a CDR3 can appear in only +one split. Nested training subsets are prefixes of the same permutation. + +| split | rows | +| --- | ---: | +| train | 79,730 | +| validation | 9,966 | +| test | 9,967 | +| subsets | 1,000 ⊂ 10,000 ⊂ 79,730 | + +Both targets are cached: `log10(Pgen)` and `log10(Pgen_1mm)`, computed with mirpy's OLGA +wrapper (1 h 43 min on 8 CPU cores). Verified: no CDR3 overlap between any pair of +splits, and target means agree across splits to within 0.01 for both targets. + +Reproduce with `python -m benchmark.prepare_splits`. + +## 2. Sequence reconstruction + +Every representation is reduced to the same **64-dimensional bottleneck** by a +`StandardScaler` + PCA fitted on the training split only, then decoded by the same +IRRM-CODEC decoder (`InverseModel`, 14.48 M parameters at this input width) with +identical settings: AdamW, lr 1e-3, weight decay 1e-4, batch 256, up to 60 epochs, +early stopping with patience 8. Three seeds per representation (42, 43, 44), 15 runs +total, each on one GTX 1080 Ti. + +Decoder inputs are standardized with train-split statistics. Without it a shared +learning rate would act differently on each representation — the projections arrive with +standard deviations from 1.0 to 12.2 — making "identical training settings" true only on +paper. The repo's own forward/inverse trainers normalize their embeddings the same way. + +### Results at a 64-dimensional bottleneck + +Mean ± standard deviation over three seeds, on the held-out test split. + +| representation | source dim | expl. var | exact match | token acc | norm. Levenshtein | within ED ≤ 1 | mean ED | train s | seq/s | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| aligned one-hot | 1000 | 0.380 | **0.9144** ± 0.0006 | 0.9878 | 0.0106 | 0.9643 | 0.171 | 1792 | 6807 | +| SCEPTR | 64 | 1.000 | **0.9005** ± 0.0081 | 0.9866 | 0.0112 | 0.9586 | 0.176 | 1322 | 6781 | +| TCRemP | 9000 | 0.950 | 0.8017 ± 0.0038 | 0.9784 | 0.0197 | 0.9406 | 0.295 | 1796 | 6968 | +| TCR-BERT | 768 | 0.910 | 0.7700 ± 0.0056 | 0.9644 | 0.0291 | 0.8879 | 0.451 | 1174 | 6864 | +| ESM-2 8M | 320 | 0.974 | 0.6510 ± 0.0014 | 0.9326 | 0.0560 | 0.7692 | 0.849 | 1435 | 6720 | + +Seed variance is small — at most ±0.008 on exact match — so gaps larger than about 0.02 +are real rather than noise. + +### Explained variance does not predict reconstruction quality + +The relationship is inverted: aligned one-hot retains the least variance (0.380) and +reconstructs best, while ESM-2 retains almost all of it (0.974) and reconstructs worst. + +Explained variance measures how much statistical spread survives the projection, not how +much residue-level information survives. One-hot has no correlations for PCA to exploit, +so its variance is spread evenly and the retained fraction looks low — but the 64 +components it keeps describe the dominant axes of sequence variation, which is exactly +what the decoder needs. + +The consequence is that explained variance cannot be used to normalize or excuse the +comparison. It is reported here because it is informative about each representation's +redundancy, not because it explains the ranking. + +### One-hot is a ceiling, not a peer + +The bottleneck equalizes dimensionality but not provenance. Aligned one-hot is a +**reversible encoding of the target itself**, so reconstructing from its PCA projection is +an autoencoding task: compress a sequence, decompress it. The other four invert a lossy +transformation of the sequence into a biological feature space. These are different +problems. + +One-hot's 0.9144 should therefore be read as an approximate ceiling — what 64 numbers can +carry about a CDR3 when they encode the string directly. Its value is diagnostic: it shows +the decoder's capacity is not the limiting factor, so the spread among the other four +reflects their representations rather than the measuring instrument. + +### Mean pooling handicaps the protein language models + +ESM-2 and TCR-BERT emit one vector per residue, and the decoder needs one vector per +sequence. Both are mean-pooled over real residues, excluding BOS/EOS and padding — the +same rule for both, chosen for comparability. + +Averaging over positions discards positional information, which is precisely what +sequence reconstruction requires. Their low scores partly measure this pooling choice +rather than the representations themselves. A CLS token, or per-position embeddings +concatenated before PCA, would likely score higher. Under the issue's constraint of one +frozen checkpoint producing one vector per receptor, mean pooling is the standard choice, +but the caveat belongs with the numbers. + +### Where IRRM-CODEC wins and loses + +**Loses: TCRemP is not the strongest compact representation.** SCEPTR reaches 0.9005 +against TCRemP's 0.8017 — about 10 points — and does so while losing nothing to the +projection (explained variance 1.000, since it is natively 64-dimensional) where TCRemP +loses 5%. SCEPTR lands within 1.4 points of the one-hot ceiling, meaning its 64 numbers +are nearly as informative about the sequence as a direct compression of the sequence. +This is a genuine negative result for the claim that TCRemP is a good compact receptor +representation. + +**Wins: TCRemP survives aggressive compression, and inversion improves under it.** The +committed multi-chain notebook reports exact match 0.7371 for the inverse model on raw +9000-dimensional TRB embeddings. At 64 dimensions — a 140× reduction — this benchmark +reaches **0.8017**. Compression did not cost accuracy; it appears to have removed noise. +That supports the project's premise that the TCRemP space is invertible and usable as a +design space, and it means the practical form of the model can be far smaller than the +published one. + +**Wins: it beats both pretrained protein language models.** TCRemP is ahead of TCR-BERT +by 3 points and ESM-2 8M by 15, though the mean-pooling caveat above applies to both. + +**Neutral: inference speed does not separate the representations.** All five run at +6,700–7,000 sequences/second, because the cost is dominated by the shared decoder, not by +the input. Encoder cost is what differs and is paid once (measured separately: one-hot +0.013 ms/sequence, SCEPTR 1.30, ESM-2 2.07, TCR-BERT 10.11 on 8 CPU threads). + +Reconstruction quality is not what TCRemP was designed for — it encodes similarity to a +reference panel, for clustering and specificity work. Sequence recoverability is a +property IRRM-CODEC adds on top. The reconstruction benchmark therefore bounds one claim +about the representation; the Pgen benchmark tests the project's speed argument, where +the comparison is against OLGA rather than against other encoders. + +Reproduce with `benchmark/slurm/reconstruction_array.sbatch`, then +`python -m benchmark.collect_reconstruction`. + +## 3. Pgen prediction + +Six arms, both targets, three training sizes. The three frozen-embedding arms share one +regression head (512 → 256 → 1) and one training configuration, so only the representation +differs between them. The two IRRM arms differ only in initialization. Ridge is +deterministic and selects its penalty on validation, which is that arm's equivalent of +early stopping. + +Targets are standardized on train statistics and mapped back before any metric, so a +freshly initialized network does not spend its first epochs locating an offset near −6.5 +that Ridge fits for free. Reported errors are in log10 units. + +Neural arms ran ten seeds on every cell: 180 runs. Ten rather than three because the two +IRRM arms spread widely across seeds, and the effect they exist to measure is small by +construction — the arms are the same architecture on the same data, differing only in +their starting weights. + +The whole sweep was then run **twice**, under two learning-rate configurations: a constant +1e-3, and a cosine schedule decaying to zero over the epoch budget (with early-stopping +patience widened from 8 to 12 so the schedule reaches its tail). The second configuration +was not a tuning pass but a robustness check, and it changed the headline conclusion. All +tables below report the cosine runs, which converge better and vary less; the constant-rate +numbers are kept in `results/pgen_constlr/` for comparison. + +### Results on the full training split + +Mean ± standard deviation over seeds, on the held-out test split. + +**`log10_pgen_1mm`, train = 79,730** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_pretrained | **0.2181** ± 0.0322 | 0.1248 | 0.9813 | 0.9935 | 0.9921 | +0.0338 | 187 | 1584 | +| irrm_scratch | 0.2435 ± 0.0248 | 0.1583 | 0.9770 | 0.9916 | 0.9937 | +0.0466 | 105 | 1583 | +| tcr_bert_mlp | 0.2766 ± 0.0006 | 0.1671 | 0.9706 | 0.9854 | 0.9908 | +0.0175 | 71 | 1668 | +| sceptr_mlp | 0.4465 ± 0.0060 | 0.2510 | 0.9233 | 0.9621 | 0.9783 | +0.0642 | 58 | 1193 | +| esm2_8m_mlp | 0.4539 ± 0.0015 | 0.2655 | 0.9208 | 0.9607 | 0.9762 | +0.0549 | 68 | 1295 | +| tfidf_ridge | 0.5542 | 0.3671 | 0.8819 | 0.9391 | 0.9566 | +0.0035 | 9 | 717 | + +**`log10_pgen`, train = 79,730** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_pretrained | **0.3102** ± 0.0388 | 0.1925 | 0.9722 | 0.9912 | 0.9883 | +0.1177 | 159 | 1593 | +| irrm_scratch | 0.3114 ± 0.0333 | 0.2027 | 0.9721 | 0.9921 | 0.9928 | +0.0987 | 131 | 1580 | +| tcr_bert_mlp | 0.3758 ± 0.0019 | 0.2333 | 0.9598 | 0.9800 | 0.9851 | +0.0238 | 74 | 1666 | +| esm2_8m_mlp | 0.5728 ± 0.0005 | 0.3494 | 0.9065 | 0.9535 | 0.9673 | +0.0746 | 63 | 1296 | +| sceptr_mlp | 0.5795 ± 0.0046 | 0.3442 | 0.9043 | 0.9532 | 0.9668 | +0.0976 | 61 | 1193 | +| tfidf_ridge | 0.6210 | 0.4145 | 0.8901 | 0.9435 | 0.9563 | +0.0087 | 8 | 714 | + +The ordering is identical on both targets and at every training size: the two IRRM arms +lead, TCR-BERT is the best frozen representation, then SCEPTR and ESM-2, with Ridge last. +The gaps are large — IRRM halves the error of SCEPTR, ESM-2 and Ridge — and far exceed +seed spread. Per-size tables are in `results/pgen/pgen_summary.md`. + +For reference, the committed pgen notebook reports RMSE 0.5914, R² 0.9437 and Pearson +0.9874 on `log10_pgen_1mm`. The best arm here reaches RMSE 0.2181, R² 0.9813 and Pearson +0.9935 on a different split of the same data. + +### Does IRRM pretraining improve Pgen prediction? + +**No — not consistently.** Welch's t-test on RMSE, ten seeds per arm, under both +learning-rate configurations. A positive delta means pretraining won. + +| target | train | Δ RMSE (constant lr) | p | Δ RMSE (cosine) | p | +| --- | --- | ---: | ---: | ---: | ---: | +| log10_pgen | 1k | −0.017 | 0.527 | +0.045 | 0.113 | +| log10_pgen | 10k | −0.037 | 0.157 | −0.030 | **0.006** | +| log10_pgen | all | +0.045 | 0.118 | +0.001 | 0.940 | +| log10_pgen_1mm | 1k | −0.025 | 0.307 | +0.068 | **0.044** | +| log10_pgen_1mm | 10k | −0.013 | 0.500 | −0.006 | 0.763 | +| log10_pgen_1mm | all | +0.056 | **0.012** | +0.025 | 0.065 | + +Under the constant rate, one cell was significant: pretraining cut RMSE by 20.9% on +`log10_pgen_1mm` at full data, p = 0.012, and it replicated across an accidental rerun. +That looked like a real, narrow effect. + +It did not survive the schedule change. The same cell falls to Δ = +0.025 at p = 0.065, +while two different cells become significant — with **opposite signs**: pretraining +"helps" at 1k on one target and "hurts" at 10k on the other. No cell is significant under +both configurations, and three cells flip sign between them. That is the signature of +noise being read as effect, not of a small effect being measured precisely. + +The mechanism is visible in the absolute numbers. Cosine scheduling improved almost every +arm, but not the pretrained one: + +| arm, `log10_pgen_1mm` all | constant lr | cosine | +| --- | ---: | ---: | +| irrm_scratch | 0.2653 | **0.2435** | +| irrm_pretrained | 0.2098 | 0.2181 | +| tcr_bert_mlp | 0.3115 | **0.2766** | + +The from-scratch arm gained what the pretrained arm did not. The transferred encoder was +supplying a good starting point that compensated for a poorly annealed learning rate; once +the optimizer no longer needed compensating, the advantage disappeared. Pretraining was +substituting for training quality rather than adding information. + +Seed spread also fell where the comparison lives — 0.071 → 0.033 and 0.055 → 0.017 on the +from-scratch arm at `all` and 10k — so the cosine numbers are the more trustworthy ones, +not merely the more recent. + +With six comparisons, a Bonferroni threshold is p < 0.008. Only the 10k "hurts" result +meets it, and it contradicts the 1k "helps" result on the other target. The honest +reading is that the sequence-to-TCRemP initialization has no reliable effect on Pgen +prediction at any training size. + +### Is the advantage larger in the low-data regime? + +The question presupposes an advantage that the data does not support. Across both +configurations and all six cells there is no consistent direction: the sign flips between +configurations at 1k on both targets, and the two nominally significant results point +opposite ways. + +One asymmetry does survive both runs. The pretrained arm carries a larger positive bias at +small data — +0.32 and +0.24 at 1k against +0.23 for the from-scratch arm — so the +transferred features do push predictions systematically, and a small training set corrects +that less well. But this shows up as bias rather than as a reliable RMSE penalty. + +### Quality against speed + +End-to-end inference: string in, log10 Pgen out, including tokenization or encoder passes, +measured on one device at a time over 5,000 test sequences. + +| arm | GPU (GTX 1080 Ti) | CPU (Xeon Gold, 8 threads) | RMSE (1mm, all) | +| --- | ---: | ---: | ---: | +| tfidf_ridge | 86,661 seq/s | 82,842 seq/s | 0.5542 | +| **irrm** | **64,814 seq/s** | **3,777 seq/s** | **0.2098** | +| esm2_8m_mlp | 9,393 seq/s | 813 seq/s | 0.4794 | +| tcr_bert_mlp | 4,039 seq/s | 216 seq/s | 0.3115 | +| sceptr_mlp | 2,495 seq/s | 1,265 seq/s | 0.4560 | + +**IRRM dominates all three pretrained encoders on both axes at once** — 4.6× to 17.5× +faster on CPU, 6.9× to 26× on GPU, and lower RMSE than any of them. There is no trade-off +to make against the large frozen encoders; they are simply worse in both respects, because +their per-sequence encoding cost dwarfs the regression head that follows it. + +The real trade-off is against Ridge, which is 1.3× faster on GPU and 22× on CPU while +carrying 2.6× the error. Ridge is the right choice when throughput matters more than +precision; IRRM when it does not. + +Caveat on the SCEPTR row: it is the slowest arm on GPU yet mid-pack on CPU, which suggests +its library evaluates on CPU unless `sceptr.enable_hardware_acceleration()` is called, +which this benchmark does not do. Its GPU figure should be read as a CPU computation plus +transfer overhead, not as a GPU measurement. + +### Against OLGA + +The practical argument for a neural Pgen predictor is speed against the exact calculation. +Measured on the same 8-core laptop CPU, so the two numbers are directly comparable: + +| | sequences/second | IRRM speedup | +| --- | ---: | ---: | +| IRRM | 1,969 | — | +| OLGA exact | 42.7 | **46×** | +| OLGA 1-mismatch | 2.7 | **729×** | + +On the cluster GPU IRRM reaches 64,814 sequences/second, but OLGA was not benchmarked on +that node, so the laptop comparison is the honest one to quote. + +### Where IRRM-CODEC wins and loses + +**Wins: accuracy at every training size, on both targets.** IRRM leads all four baselines +in all six cells. On the full split it reaches R² 0.981 and Pearson 0.994 for +`log10_pgen_1mm`, and halves the error of SCEPTR, ESM-2 and Ridge. These margins are ten +times the seed spread, so they are not in question. + +**Wins: the quality–speed trade-off against pretrained encoders.** Faster and more accurate +than frozen ESM-2, TCR-BERT and SCEPTR simultaneously. For this task the large pretrained +encoders contribute nothing on either axis. + +**Wins: 46× to 729× over OLGA** on identical hardware, which is the argument the model +exists to make. + +**Loses: the pretraining step buys nothing.** The sequence-to-TCRemP initialization shows no +reliable effect once the learning rate is properly annealed, and what looked like a 20.9% +gain under a constant rate was the pretrained encoder compensating for an untuned +optimizer. The practical implication is direct: train the Pgen model from scratch and skip +the pretraining stage, which costs an extra training run for no measurable return. + +**Loses: seed stability.** Even with cosine scheduling the IRRM arms vary by 0.017–0.092 +RMSE across seeds where the frozen arms vary by 0.0005–0.006 — still an order of magnitude +wider, because these arms fit a full encoder rather than a head over fixed features. Any +comparison between IRRM variants needs many seeds; three is not enough, as the first sweep +demonstrated. + +**Loses: throughput against a trivial baseline.** TF-IDF + Ridge is 22× faster on CPU. If +an application needs coarse Pgen estimates at maximum rate, k-mers remain the right tool. + +Reproduce with `benchmark/slurm/pgen_array.sbatch` and +`benchmark/slurm/pgen_seeds_array.sbatch`, then `python -m benchmark.collect_pgen` and +`python -m benchmark.benchmark_speed`. + +## Limitations + +- **Mean pooling** for ESM-2 and TCR-BERT discards positional information, which penalizes + them most on reconstruction. A CLS token or per-position features would likely score + higher. +- **Aligned one-hot is not a peer baseline** in section 2: it encodes the target + reversibly, so it bounds the task rather than competing in it. +- **GPU training is not bit-reproducible.** Re-running a configuration with the same seed + gives slightly different results, because cuDNN selects algorithms at runtime and some + backward kernels accumulate non-deterministically. Conclusions here rest on ten-seed + means rather than single runs. +- **One chain, one species.** Everything is human TRB. Nothing here shows the ordering + holds for TRA or for other loci. +- **One bottleneck width.** Reconstruction was compared at 64 dimensions. Projections at + 32 and 128 are cached in `data/benchmark/trb/bottleneck/` but were not trained. +- **Six statistical comparisons** were made for the pretraining question, under each of two + learning-rate configurations. No cell is significant under both, and three flip sign + between them, which is why the conclusion is stated as absence of a reliable effect + rather than as a measured null. +- **Sections 2 and 3 used different learning-rate schedules** — constant for reconstruction, + cosine for Pgen. The issues require identical settings across the arms compared within + each experiment, which holds; they are not comparable across sections. Reconstruction was + not rerun because its seed spread (±0.0006 to ±0.008) leaves nothing for a schedule to + fix. + +## Environment + +Sections 2 and 3 were produced on the ctlab Slurm cluster: GTX 1080 Ti, torch 2.6.0+cu124, +Python 3.11. Section 1 timings and the OLGA comparison are from an 8-core CPU +(Ryzen 7 8845HS); the CPU speed column in section 3 is from a cluster Xeon Gold node. +Peak memory is peak resident set size via `getrusage`, not +`torch.cuda.max_memory_allocated`, and the device is recorded with every measurement. + +Encoder checkpoints: `esm2_t6_8M_UR50D` (fair-esm), `wukevin/tcr-bert` (HuggingFace), +SCEPTR default variant (`sceptr` 1.2.0). diff --git a/benchmark/EXTERNAL_RESIDUE_BENCHMARK.md b/benchmark/EXTERNAL_RESIDUE_BENCHMARK.md new file mode 100644 index 0000000..b23e518 --- /dev/null +++ b/benchmark/EXTERNAL_RESIDUE_BENCHMARK.md @@ -0,0 +1,57 @@ +# Frozen residue-state benchmark preregistration + +## Question + +Does the poor reconstruction of public encoders in the common 64D sequence-vector +benchmark arise primarily from global pooling? The corrective analysis exposes frozen +per-residue states to a trainable, position-aware task head. + +## Locked inputs + +- Cohort: `/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/cohort` +- Split: the existing 79,544 / 9,943 / 9,943 train/validation/test manifests. +- Models: ESM2-8M, `wukevin/tcr-bert@ef65ddc`, and SCEPTR 1.2.0 default. +- Encoder parameters are frozen. Validation loss alone chooses checkpoints. +- `log10_pgen_1mm` is normalized using the training split only. + +## Representation interfaces + +The published pooled/vector results remain unchanged and are the compact-vector panel. +The new sensitivity panel uses: + +- ESM2: final-layer amino-acid states, excluding BOS/EOS/padding. +- TCR-BERT: final-layer amino-acid states, excluding CLS/SEP/padding. +- SCEPTR: default model penultimate-layer CDR3B residue states selected with compartment + mask 6. The existing default SCEPTR vector is a learned CLS representation, not a + mean-pooled vector. + +States are zero-padded to 40 residues and stored as float16; task heads cast to float32. +No amino-acid states are averaged before reaching the trainable head. + +## Task heads + +Each representation gets its own freshly initialized decoder/probe for seeds 42, 43, +and 44. Reconstruction uses the same learned output queries and three-layer +cross-attention decoder for every representation. Pgen uses the same two-layer +position-aware attention probe for every representation. Only the unavoidable input +projection width changes with the frozen encoder dimension; parameter counts are +reported. All heads use AdamW at `3e-4`, matching the locked RTP training rate. A +discarded `1e-3` pilot was stopped after its TCR-BERT validation loss diverged; no +pilot checkpoint or test result enters this comparison. + +## Endpoints + +- Reconstruction primary: exact CDR3 match. Secondary: residue accuracy, normalized + Levenshtein distance, edit distance <=1, and mean edit distance. +- Pgen primary: raw-unit RMSE. Secondary: MAE, R2, Pearson, Spearman, and bias. +- The prespecified comparison is within encoder: residue-state head versus its existing + pooled/vector head. RTP remains the fixed compact-vector reference. + +## Interpretation boundary + +This is a best-available-frozen-interface sensitivity analysis, not a matched compact +representation comparison. A residue tensor contains more numbers than a 64D or 128D +vector, and its task head is position-aware. If it closes the gap, the pooled benchmark +was interface-limited. If RTP remains better, that is stronger evidence for accessible +task information, but still not universal representation superiority because RTP was +trained on the benchmark reconstruction and pgen objectives. diff --git a/benchmark/EXTERNAL_RESIDUE_PCA256_BENCHMARK.md b/benchmark/EXTERNAL_RESIDUE_PCA256_BENCHMARK.md new file mode 100644 index 0000000..a82bb51 --- /dev/null +++ b/benchmark/EXTERNAL_RESIDUE_PCA256_BENCHMARK.md @@ -0,0 +1,35 @@ +# Frozen residue-state PCA-256 benchmark preregistration + +## Question and correction + +The full `[40,D]` residue tensors showed that ESM2 and TCR-BERT retain near-exact +sequence identity, but give each clonotype far more information than a compact IRRM +embedding. That diagnostic array was stopped and is excluded from comparative +results. The corrected comparison maps every public encoder to exactly one 256D +point before any supervised head is fitted. + +## Leakage-safe representation + +For ESM2-8M, TCR-BERT and SCEPTR, the zero-padded residue states are flattened in +sequence order. Per-coordinate mean and scale and a seeded randomized 256-component +PCA are fitted on the locked 79,544-row training manifest only. The fixed transform +is then applied to train, validation and test. Reports record source/projection +hashes, the train-manifest and row-index hashes, solver seed, explained variance, +shapes and finiteness. Validation and test rows are never used to fit scaling or PCA. + +## Downstream heads + +Each encoder receives its own freshly initialized reconstruction decoder and pgen MLP +for seeds 42, 43 and 44. The architecture, optimizer, early stopping and training +budget are the established compact-vector benchmark implementations; only the 256D +input values differ. Checkpoints are selected on validation loss. Reconstruction is +scored by exact match, residue accuracy and edit-distance measures. Pgen prediction +is scored in raw `log10_pgen_1mm` units by RMSE, MAE, R2, Pearson, Spearman and bias. + +## Interpretation boundary + +This asks how much useful information can be retained in a common 256D linear +compression of each frozen residue representation. It does not prove that the public +model's native pooling is poor in every application, and 256D remains twice the +selected RTP latent width of 128. Existing native-vector/PCA-64 results and the new +residue-PCA-256 results are reported as distinct interfaces. diff --git a/benchmark/LATENT_BENCHMARK_PREREGISTRATION.md b/benchmark/LATENT_BENCHMARK_PREREGISTRATION.md new file mode 100644 index 0000000..b63584b --- /dev/null +++ b/benchmark/LATENT_BENCHMARK_PREREGISTRATION.md @@ -0,0 +1,83 @@ +# DATA-ANCHOR latent benchmark preregistration + +Status: locked before inspecting any result from the evaluations defined here. + +## Claim under test + +At the same 128-dimensional bottleneck and encoder architecture, the frozen +R+T+P representation is a better *balanced, multipurpose representation* than +any frozen single-objective representation. This is not a claim that R+T+P must +win every specialist task or every unrelated downstream endpoint. + +The matched seed-42 DATA-ANCHOR models are R, T, P, RT, RP, TP, and RTP. All +seven use the locked TRB benchmark, train-only normalizer, tokenizer, model +architecture, and validation-selected checkpoint. No encoder is tuned using an +evaluation result. + +## Primary evaluation: frozen latent sufficiency + +Identical fresh readouts are fitted from each frozen latent representation using +the locked training split, selected using validation only, and evaluated once on +test. Readouts measure sequence content, standardized TCRemP target recovery, +and `log10_pgen_1mm` recovery. The original task heads are not reused as probes. + +Primary model-level endpoints, in order, are: + +1. worst-task normalized regret across R, T, and P targets (lower is better); +2. mean normalized regret across the same targets (lower is better); +3. Pareto hypervolume with a fixed null-reference point (higher is better); +4. number of targets within 5% of the corresponding best specialist. + +Task losses are oriented so lower is better and normalized using fixed trivial +train-derived nulls and the best validation-eligible specialist. The test set +does not set signs, weights, thresholds, readout capacity, or normalization. +Paired test-clonotype bootstrap intervals are reported for RTP minus every +comparator. These intervals do not measure training-seed variability. + +## Controlled one-amino-acid mutation challenge + +Valid single-substitution variants are grouped by parent clonotype. A parent's +variants cannot cross fitting/evaluation folds. TCRemP and pgen effects must be +computed by an existing validated scorer; absence of such a scorer blocks that +endpoint rather than allowing a proxy. + +Primary endpoints are balanced prediction of delta-TCRemP and delta-log-pgen, +followed by pgen direction accuracy and mutant-effect rank correlation. Mutation +position/substitution recovery and triplet consistency are secondary. Since all +pairs have edit distance one, this analysis tests organization beyond raw edit +distance. Parent-clustered bootstrap intervals are required. + +## Secondary evaluations + +These are reported separately and cannot replace a failed primary endpoint: + +- nested low-label readout curves; +- fixed subgroup performance by CDR3 length, V/J family, and pgen quantile; +- joint sequence/TCRemP/pgen retrieval using train-defined scaling; +- masking and conservative/nonconservative substitution robustness; +- synthetic-versus-observed discrimination only if a provenance-verified, + benchmark-deduplicated synthetic source already exists. + +## Baselines and efficiency + +Every feasible track includes the seven frozen models plus R|T|P concatenation +(384 dimensions), train-only PCA of that concatenation to 128 dimensions, and +sequence composition/length controls (with V/J reported separately). Runtime, +encoder count, representation dimension, and parameter count are reported. The +384-dimensional concatenation is an expensive upper bound, not a size-matched +single-encoder baseline. + +## Decision rule + +The primary claim is supported only if RTP improves the preregistered balanced +endpoint over each single-objective model with a paired 95% interval excluding +no improvement, while remaining competitive on every constituent task. Wins on +an isolated secondary metric, subgroup, perturbation type, or retrospectively +chosen endpoint do not establish the claim. All negative results are retained. + +## Execution controls + +All remote computation runs through CPU Slurm using +`/home/evlasova/.conda/envs/irrm-codec/bin/python`. No computation runs on an +Aldan frontend. Inputs, model checkpoints, configs, scripts, and outputs receive +SHA-256 manifests; outputs are isolated by study and Slurm job identifier. diff --git a/benchmark/__init__.py b/benchmark/__init__.py index 6e25b91..69b6918 100644 --- a/benchmark/__init__.py +++ b/benchmark/__init__.py @@ -1 +1 @@ -"""Baseline benchmark for IRRM-CODEC on reconstruction and pgen prediction.""" +"""Dataset preparation utilities for reproducible IRRM-CODEC benchmarks.""" diff --git a/benchmark/audit_external_residue_alignment.py b/benchmark/audit_external_residue_alignment.py new file mode 100644 index 0000000..e9eb9c8 --- /dev/null +++ b/benchmark/audit_external_residue_alignment.py @@ -0,0 +1,154 @@ +"""Audit row/mask provenance and direct residue-token retention for frozen encoders.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd + + +NAMES = ("esm2_8m", "tcr_bert", "sceptr") +AA = "ACDEFGHIKLMNPQRSTVWY" +AA_TO_ID = {aa: i for i, aa in enumerate(AA)} + + +def sequence_sha256(sequences) -> str: + return hashlib.sha256("\n".join(sequences).encode("ascii")).hexdigest() + + +def length_matched_permutation(lengths: np.ndarray, seed: int) -> np.ndarray: + rng = np.random.default_rng(seed) + result = np.arange(len(lengths)) + for length in np.unique(lengths): + positions = np.flatnonzero(lengths == length) + if len(positions) > 1: + order = rng.permutation(positions) + result[order] = np.roll(order, 1) + return result + + +def labels_for_rows(sequences: np.ndarray, rows: np.ndarray, max_len: int = 40): + labels = np.full((len(rows), max_len), -1, dtype=np.int16) + mask = np.zeros((len(rows), max_len), dtype=bool) + for out_row, source_row in enumerate(rows): + sequence = sequences[source_row] + encoded = [AA_TO_ID[aa] for aa in sequence] + labels[out_row, : len(encoded)] = encoded + mask[out_row, : len(encoded)] = True + return labels, mask + + +def centroid_probe(matrix, train_rows, test_rows, train_labels, test_labels, train_mask, test_mask): + dim = matrix.shape[-1] + sums = np.zeros((len(AA), dim), dtype=np.float64) + counts = np.zeros(len(AA), dtype=np.int64) + for local_start in range(0, len(train_rows), 1024): + local_stop = min(local_start + 1024, len(train_rows)) + block = np.asarray(matrix[train_rows[local_start:local_stop]], dtype=np.float32) + labels = train_labels[local_start:local_stop] + mask = train_mask[local_start:local_stop] + for aa_id in range(len(AA)): + selected = block[labels == aa_id] + if len(selected): + sums[aa_id] += selected.sum(axis=0, dtype=np.float64) + counts[aa_id] += len(selected) + if (counts == 0).any(): + raise ValueError("Missing amino acid in training-centroid audit") + centroids = (sums / counts[:, None]).astype(np.float32) + centroid_norm = np.square(centroids).sum(axis=1) + correct = total = 0 + predictions_by_row = np.full_like(test_labels, -1) + for local_start in range(0, len(test_rows), 512): + local_stop = min(local_start + 512, len(test_rows)) + block = np.asarray(matrix[test_rows[local_start:local_stop]], dtype=np.float32) + flat = block.reshape(-1, dim) + distances = ( + np.square(flat).sum(axis=1, keepdims=True) + - 2 * flat @ centroids.T + + centroid_norm[None, :] + ) + predicted = distances.argmin(axis=1).reshape(block.shape[:2]) + predictions_by_row[local_start:local_stop] = predicted + labels = test_labels[local_start:local_stop] + mask = test_mask[local_start:local_stop] + correct += int((predicted[mask] == labels[mask]).sum()) + total += int(mask.sum()) + return correct / total, predictions_by_row + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-dir", required=True) + parser.add_argument("--residue-dir", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def main(): + args = parse_args() + dataset_dir, residue_dir = Path(args.dataset_dir), Path(args.residue_dir) + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + sequences = frame["junction_aa"].astype(str).to_numpy() + train_rows = pd.read_csv(dataset_dir / "manifests/train.tsv", sep="\t")["row_index"].to_numpy(np.int64) + test_rows = pd.read_csv(dataset_dir / "manifests/test.tsv", sep="\t")["row_index"].to_numpy(np.int64) + train_labels, train_mask = labels_for_rows(sequences, train_rows) + test_labels, test_mask = labels_for_rows(sequences, test_rows) + test_lengths = test_mask.sum(axis=1) + perm = length_matched_permutation(test_lengths, args.seed) + manifest = json.loads((residue_dir / "residue_representations.json").read_text(encoding="utf-8")) + expected_sequence_sha = sequence_sha256(sequences) + report = { + "status": "accepted", + "sequence_sha256": expected_sequence_sha, + "train_rows": len(train_rows), + "test_rows": len(test_rows), + "models": {}, + "interpretation": "Centroid accuracy measures direct amino-acid identity retained by raw residue states; it is not an epitope or biological endpoint.", + } + for name in NAMES: + entry = manifest[name] + if entry["sequence_sha256"] != expected_sequence_sha or entry["pooling"] != "none": + raise ValueError(f"Sequence hash/pooling provenance mismatch for {name}") + matrix = np.load(entry["features"], mmap_mode="r") + mask = np.load(entry["mask"], mmap_mode="r") + expected_shape = tuple(entry["shape"]) + if tuple(matrix.shape) != expected_shape or mask.shape != matrix.shape[:2]: + raise ValueError(f"Shape mismatch for {name}") + for start in range(0, len(matrix), 4096): + if not np.isfinite(matrix[start:start + 4096]).all(): + raise ValueError(f"Non-finite residue states for {name}") + expected_mask = np.zeros_like(mask, dtype=bool) + for row, sequence in enumerate(sequences): + expected_mask[row, : len(sequence)] = True + if not np.array_equal(mask, expected_mask): + raise ValueError(f"Residue mask/sequence alignment mismatch for {name}") + accuracy, predicted = centroid_probe( + matrix, train_rows, test_rows, train_labels, test_labels, train_mask, test_mask + ) + permuted_labels = test_labels[perm] + permuted_mask = test_mask[perm] + permuted_accuracy = float( + (predicted[permuted_mask] == permuted_labels[permuted_mask]).mean() + ) + report["models"][name] = { + "model": entry["model"], + "shape": entry["shape"], + "features_sha256": entry["features_sha256"], + "mask_sha256": entry["mask_sha256"], + "mask_matches_sequence_lengths": True, + "train_centroid_test_residue_accuracy": accuracy, + "length_matched_permuted_row_accuracy": permuted_accuracy, + } + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2), encoding="utf-8") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/benchmark_speed.py b/benchmark/benchmark_speed.py new file mode 100644 index 0000000..3b44436 --- /dev/null +++ b/benchmark/benchmark_speed.py @@ -0,0 +1,200 @@ +"""End-to-end inference speed for every Pgen arm, measured on one device. + +The per-arm speeds recorded during training cover only the trained head. For the frozen +arms that omits the encoder, which is the dominant cost and the reason those numbers +looked orders of magnitude faster than they are in practice. Encoder timings were also +collected on a different machine than the heads, so the two cannot be combined. + +This script times the full path a new sequence actually travels — string in, log10 Pgen +out — for every arm on whichever device is requested, so the quality-versus-speed +comparison rests on one set of measurements. +""" + +import argparse +import json +import time +import warnings +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from irrm_codec.pgen_model import PgenModel +from irrm_codec.tokenization import PAD_ID, encode +from irrm_codec.utils import setup_logging + +ARMS = ("irrm", "tfidf_ridge", "sceptr_mlp", "tcr_bert_mlp", "esm2_8m_mlp") + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument("--output-dir", default="results/pgen") + p.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda")) + p.add_argument("--n-sequences", type=int, default=5000) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--repeats", type=int, default=3) + p.add_argument("--threads", type=int, default=8) + p.add_argument("--arms", nargs="+", default=list(ARMS), choices=list(ARMS)) + return p.parse_args() + + +def timed(fn, sequences, repeats): + """Fastest of several passes, so a stray scheduling hiccup does not set the number.""" + fn(sequences[: min(64, len(sequences))]) # warm up lazy initialization + best = float("inf") + for _ in range(repeats): + started = time.perf_counter() + fn(sequences) + best = min(best, time.perf_counter() - started) + return { + "sequences_per_second": len(sequences) / best, + "ms_per_sequence": best / len(sequences) * 1000, + "seconds_total": best, + } + + +def make_irrm(args, device): + """Tokenize and run the sequence encoder: the whole pipeline for this arm.""" + model = PgenModel(max_len=40).to(device).eval() + + @torch.no_grad() + def run(sequences): + for start in range(0, len(sequences), args.batch_size): + batch = sequences[start : start + args.batch_size] + tokens = torch.tensor([encode(s, max_len=40) for s in batch], dtype=torch.long).to(device) + model(tokens, tokens.ne(PAD_ID)) + if device.type == "cuda": + torch.cuda.synchronize() + + return run + + +def make_tfidf_ridge(args, device, train_sequences, train_targets): + from sklearn.feature_extraction.text import TfidfVectorizer + from sklearn.linear_model import Ridge + + vectorizer = TfidfVectorizer(analyzer="char", ngram_range=(1, 3)) + features = vectorizer.fit_transform(train_sequences) + model = Ridge(alpha=1.0).fit(features, train_targets) + + def run(sequences): + model.predict(vectorizer.transform(sequences)) + + return run + + +def make_frozen(args, device, representation): + """Encoder plus the regression head, which is what predicting a new sequence costs.""" + from benchmark.train_pgen_arm import RegressionHead + + if representation == "sceptr_mlp": + import sceptr + + head = RegressionHead(64).to(device).eval() + + def encode_batch(batch): + frame = pd.DataFrame({"TRBV": ["TRBV20-1"] * len(batch), "CDR3B": batch, "TRBJ": ["TRBJ2-7"] * len(batch)}) + return torch.from_numpy(sceptr.calc_vector_representations(frame)) + + elif representation == "esm2_8m_mlp": + import esm + + encoder, alphabet = esm.pretrained.esm2_t6_8M_UR50D() + encoder = encoder.to(device).eval() + converter = alphabet.get_batch_converter() + head = RegressionHead(encoder.embed_dim).to(device).eval() + last = encoder.num_layers + + def encode_batch(batch): + _, _, tokens = converter([(str(i), s) for i, s in enumerate(batch)]) + tokens = tokens.to(device) + states = encoder(tokens, repr_layers=[last])["representations"][last] + keep = ( + (tokens != alphabet.padding_idx) + & (tokens != alphabet.cls_idx) + & (tokens != alphabet.eos_idx) + ).unsqueeze(-1).to(states.dtype) + return (states * keep).sum(1) / keep.sum(1).clamp_min(1e-9) + + else: + from transformers import AutoModel, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained("wukevin/tcr-bert") + encoder = AutoModel.from_pretrained("wukevin/tcr-bert").to(device).eval() + head = RegressionHead(encoder.config.hidden_size).to(device).eval() + special = set(tokenizer.all_special_ids) + + def encode_batch(batch): + enc = tokenizer([" ".join(s) for s in batch], return_tensors="pt", padding=True) + enc = {k: v.to(device) for k, v in enc.items()} + states = encoder(**enc).last_hidden_state + keep = enc["attention_mask"].bool() + for token_id in special: + keep &= enc["input_ids"] != token_id + keep = keep.unsqueeze(-1).to(states.dtype) + return (states * keep).sum(1) / keep.sum(1).clamp_min(1e-9) + + @torch.no_grad() + def run(sequences): + for start in range(0, len(sequences), args.batch_size): + head(encode_batch(sequences[start : start + args.batch_size]).to(device)) + if device.type == "cuda": + torch.cuda.synchronize() + + return run + + +def main(): + args = parse_args() + warnings.filterwarnings("ignore") + torch.set_num_threads(args.threads) + + device = torch.device( + ("cuda" if torch.cuda.is_available() else "cpu") if args.device == "auto" else args.device + ) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / f"benchmark_speed_{device.type}.log") + + frame = pd.read_parquet(Path(args.dataset_dir) / "dataset.parquet") + test_rows = pd.read_csv(Path(args.dataset_dir) / "manifests" / "test.tsv", sep="\t")["row_index"] + sequences = frame.junction_aa.iloc[test_rows].tolist()[: args.n_sequences] + train_rows = pd.read_csv(Path(args.dataset_dir) / "manifests" / "train_10k.tsv", sep="\t")["row_index"] + log.info("device=%s sequences=%d batch=%d", device, len(sequences), args.batch_size) + + builders = { + "irrm": lambda: make_irrm(args, device), + "tfidf_ridge": lambda: make_tfidf_ridge( + args, device, + frame.junction_aa.iloc[train_rows].tolist(), + frame.log10_pgen_1mm.to_numpy()[train_rows], + ), + "sceptr_mlp": lambda: make_frozen(args, device, "sceptr_mlp"), + "tcr_bert_mlp": lambda: make_frozen(args, device, "tcr_bert_mlp"), + "esm2_8m_mlp": lambda: make_frozen(args, device, "esm2_8m_mlp"), + } + + results = [] + for name in args.arms: + log.info("timing %s", name) + measurement = timed(builders[name](), sequences, args.repeats) + results.append({"arm": name, "device": device.type, **measurement}) + log.info( + "%-13s %9.1f seq/s %7.3f ms/seq", + name, measurement["sequences_per_second"], measurement["ms_per_sequence"], + ) + + table = pd.DataFrame(results).sort_values("sequences_per_second", ascending=False) + path = output_dir / f"speed_end_to_end_{device.type}.csv" + table.to_csv(path, index=False) + + log.info("=" * 56) + for _, row in table.iterrows(): + log.info("%-13s %9.1f seq/s %7.3f ms/seq", row.arm, row.sequences_per_second, row.ms_per_sequence) + log.info("wrote %s", path) + + +if __name__ == "__main__": + main() diff --git a/benchmark/benchmark_vdjdb_irrm_runtime.py b/benchmark/benchmark_vdjdb_irrm_runtime.py new file mode 100644 index 0000000..1a728ff --- /dev/null +++ b/benchmark/benchmark_vdjdb_irrm_runtime.py @@ -0,0 +1,144 @@ +"""Measure frozen IRRM encoder runtime on the VDJdb uniform-100 cohort. + +The measured interval starts after the checkpoint, tokenizer and cohort are in +memory and excludes warm-up. It reports tokenization separately from the +sequence-to-latent encoder pass; the latter includes host-to-device transfer. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import platform +import time +from dataclasses import asdict +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from irrm_codec.multitask_data import resolve_encoder_tokenizer +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import PAD_ID + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def synchronize(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def make_batch(encoded: list[list[int]], rows: list[int], device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + width = max(len(encoded[row]) for row in rows) + tokens = torch.full((len(rows), width), PAD_ID, dtype=torch.long) + for local_index, row in enumerate(rows): + ids = encoded[row] + tokens[local_index, : len(ids)] = torch.tensor(ids, dtype=torch.long) + tokens = tokens.to(device, non_blocking=device.type == "cuda") + return tokens, tokens.ne(PAD_ID) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--sequence-column", default="cdr3") + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--device", choices=("cpu", "cuda"), required=True) + parser.add_argument("--batch-size", type=int, required=True) + parser.add_argument("--warmup-batches", type=int, default=3) + args = parser.parse_args() + + if args.device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable") + if args.batch_size < 1 or args.warmup_batches < 0: + raise ValueError("batch size must be positive and warmup batches non-negative") + + device = torch.device(args.device) + torch.set_grad_enabled(False) + cohort = pd.read_csv(args.cohort, sep="\t") + sequences = cohort[args.sequence_column].astype(str).tolist() + if not sequences or any(not sequence for sequence in sequences): + raise ValueError("Cohort has an empty sequence") + + payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + model_config = IRRMCodecConfig(**payload["model_config"]) + model = IRRMCodecTransformer(model_config) + model.load_state_dict(payload["model_state"], strict=True) + model.to(device).eval() + tokenizer = resolve_encoder_tokenizer("data_anchor", str(args.tokenizer)) + if tokenizer.vocab_size != model_config.input_vocab_size: + raise ValueError("Tokenizer vocabulary does not match checkpoint") + + tokenization_started = time.perf_counter() + encoded = [tokenizer.encode(sequence, model_config.max_sequence_len) for sequence in sequences] + tokenization_seconds = time.perf_counter() - tokenization_started + rows = [list(range(start, min(start + args.batch_size, len(encoded)))) for start in range(0, len(encoded), args.batch_size)] + + with torch.inference_mode(): + for batch_rows in rows[: args.warmup_batches]: + tokens, mask = make_batch(encoded, batch_rows, device) + latent = model.encode(tokens, mask) + if latent.shape != (len(batch_rows), model_config.latent_dim): + raise ValueError(f"Unexpected latent shape: {tuple(latent.shape)}") + synchronize(device) + + encoded_rows = 0 + latent_sum = 0.0 + inference_started = time.perf_counter() + for batch_rows in rows: + tokens, mask = make_batch(encoded, batch_rows, device) + latent = model.encode(tokens, mask) + latent_sum += float(latent.sum()) + encoded_rows += len(batch_rows) + synchronize(device) + inference_seconds = time.perf_counter() - inference_started + + if encoded_rows != len(sequences) or not np.isfinite(latent_sum): + raise RuntimeError("Encoder output validation failed") + result = { + "status": "complete", + "cohort": { + "path": str(args.cohort.resolve()), + "sha256": sha256_file(args.cohort), + "rows": len(sequences), + "sequence_column": args.sequence_column, + }, + "checkpoint": {"path": str(args.checkpoint.resolve()), "sha256": sha256_file(args.checkpoint)}, + "tokenizer": {"path": str(args.tokenizer.resolve()), "sha256": sha256_file(args.tokenizer)}, + "model": asdict(model_config), + "runtime": { + "device": str(device), + "batch_size": args.batch_size, + "warmup_batches": args.warmup_batches, + "tokenization_seconds": tokenization_seconds, + "encoder_seconds": inference_seconds, + "end_to_end_seconds": tokenization_seconds + inference_seconds, + "encoder_sequences_per_second": len(sequences) / inference_seconds, + "end_to_end_sequences_per_second": len(sequences) / (tokenization_seconds + inference_seconds), + "ms_per_sequence_encoder": inference_seconds / len(sequences) * 1000, + "ms_per_sequence_end_to_end": (tokenization_seconds + inference_seconds) / len(sequences) * 1000, + }, + "validation": {"latent_dim": model_config.latent_dim, "latent_sum": latent_sum}, + "environment": {"python": platform.python_version(), "torch": torch.__version__}, + } + if device.type == "cuda": + result["environment"]["cuda_device"] = torch.cuda.get_device_name(device) + result["environment"]["cuda_runtime"] = torch.version.cuda + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(result["runtime"], sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/cache_embeddings.py b/benchmark/cache_embeddings.py new file mode 100644 index 0000000..5b0c982 --- /dev/null +++ b/benchmark/cache_embeddings.py @@ -0,0 +1,207 @@ +"""Issue 2/3: generate and cache frozen representations for the benchmark dataset. + +One pretrained checkpoint per encoder, no fine-tuning. Every matrix is written in the +row order of ``dataset.parquet`` so that ``row_index`` identifies the same sequence in +every representation, including the TCRemP matrix produced by ``prepare_splits``. + +Encoders: + onehot aligned one-hot from the repo tokenizer (gap-padded to 40 x vocab) + esm2_8m ESM-2 8M (esm2_t6_8M_UR50D), mean-pooled over residues + tcr_bert TCR-BERT (wukevin/tcr-bert), mean-pooled over residues + sceptr SCEPTR default variant, already one vector per receptor + +Transformer encoders are mean-pooled over real residues only, excluding BOS/EOS and +padding, so the pooling rule is identical across models. +""" + +import argparse +import json +import time +import warnings +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from irrm_codec.tokenization import AA_VOCAB, encode +from irrm_codec.utils import setup_logging + +ENCODERS = ("onehot", "esm2_8m", "tcr_bert", "sceptr") + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument("--output-dir", default="data/benchmark/trb/representations") + p.add_argument("--encoders", nargs="+", default=list(ENCODERS), choices=list(ENCODERS)) + p.add_argument("--batch-size", type=int, default=64) + p.add_argument("--max-len", type=int, default=40) + p.add_argument("--threads", type=int, default=8) + p.add_argument("--overwrite", action="store_true", help="Recompute even if the cache exists.") + return p.parse_args() + + +def encode_onehot(sequences, args, log): + """Aligned one-hot: the repo's gap padding gives every sequence the same layout.""" + vocab_size = len(AA_VOCAB) + out = np.zeros((len(sequences), args.max_len * vocab_size), dtype=np.float32) + positions = np.arange(args.max_len) + for row, seq in enumerate(sequences): + block = np.zeros((args.max_len, vocab_size), dtype=np.float32) + block[positions, encode(seq, max_len=args.max_len)] = 1.0 + out[row] = block.ravel() + log.info("onehot done dim=%d", out.shape[1]) + return out + + +def _mean_pool(states, keep): + """Average hidden states over the positions flagged in ``keep``.""" + mask = keep.unsqueeze(-1).to(states.dtype) + return (states * mask).sum(1) / mask.sum(1).clamp_min(1e-9) + + +def encode_esm2(sequences, args, log): + import esm + + model, alphabet = esm.pretrained.esm2_t6_8M_UR50D() + model.eval() + batch_converter = alphabet.get_batch_converter() + last_layer = model.num_layers + chunks = [] + with torch.no_grad(): + for start in range(0, len(sequences), args.batch_size): + batch = sequences[start : start + args.batch_size] + _, _, tokens = batch_converter([(str(i), s) for i, s in enumerate(batch)]) + states = model(tokens, repr_layers=[last_layer])["representations"][last_layer] + keep = ( + (tokens != alphabet.padding_idx) + & (tokens != alphabet.cls_idx) + & (tokens != alphabet.eos_idx) + ) + chunks.append(_mean_pool(states, keep).numpy().astype(np.float32)) + if start % (args.batch_size * 100) == 0: + log.info("esm2_8m %d/%d", start, len(sequences)) + return np.concatenate(chunks) + + +def encode_tcr_bert(sequences, args, log): + from transformers import AutoModel, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained("wukevin/tcr-bert") + model = AutoModel.from_pretrained("wukevin/tcr-bert") + model.eval() + special = set(tokenizer.all_special_ids) + chunks = [] + with torch.no_grad(): + for start in range(0, len(sequences), args.batch_size): + batch = sequences[start : start + args.batch_size] + # TCR-BERT expects residues separated by spaces. + enc = tokenizer([" ".join(s) for s in batch], return_tensors="pt", padding=True) + states = model(**enc).last_hidden_state + keep = enc["attention_mask"].bool() + for token_id in special: + keep &= enc["input_ids"] != token_id + chunks.append(_mean_pool(states, keep).numpy().astype(np.float32)) + if start % (args.batch_size * 100) == 0: + log.info("tcr_bert %d/%d", start, len(sequences)) + return np.concatenate(chunks) + + +def encode_sceptr(sequences, args, log, frame=None): + import sceptr + + # SCEPTR reads V and J gene calls in addition to the CDR3, so it needs the full table. + instances = pd.DataFrame( + { + "TRBV": frame["v_call"].to_numpy(), + "CDR3B": frame["junction_aa"].to_numpy(), + "TRBJ": frame["j_call"].to_numpy(), + } + ) + chunks = [] + stride = max(args.batch_size, 1024) + for start in range(0, len(instances), stride): + chunks.append( + sceptr.calc_vector_representations(instances.iloc[start : start + stride]).astype(np.float32) + ) + if start % (stride * 10) == 0: + log.info("sceptr %d/%d", start, len(instances)) + return np.concatenate(chunks) + + +def main(): + args = parse_args() + warnings.filterwarnings("ignore") + torch.set_num_threads(args.threads) + + dataset_dir = Path(args.dataset_dir) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "cache_embeddings.log") + + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + sequences = frame["junction_aa"].tolist() + log.info("dataset rows=%d from %s", len(frame), dataset_dir / "dataset.parquet") + + builders = { + "onehot": lambda: encode_onehot(sequences, args, log), + "esm2_8m": lambda: encode_esm2(sequences, args, log), + "tcr_bert": lambda: encode_tcr_bert(sequences, args, log), + "sceptr": lambda: encode_sceptr(sequences, args, log, frame=frame), + } + + meta_path = output_dir / "representations.json" + meta = json.loads(meta_path.read_text(encoding="utf-8")) if meta_path.exists() else {} + + for name in args.encoders: + path = output_dir / f"{name}.npy" + if path.exists() and not args.overwrite: + cached = np.load(path, mmap_mode="r") + if cached.shape[0] == len(frame): + log.info("%s cached dim=%d, skipping", name, cached.shape[1]) + continue + log.warning("%s cache has %d rows, expected %d; recomputing", name, cached.shape[0], len(frame)) + + log.info("encoding %s", name) + started = time.perf_counter() + matrix = builders[name]() + elapsed = time.perf_counter() - started + + if matrix.shape[0] != len(frame): + raise ValueError(f"{name} produced {matrix.shape[0]} rows, expected {len(frame)}.") + if not np.isfinite(matrix).all(): + raise ValueError(f"{name} produced non-finite values.") + + np.save(path, matrix) + meta[name] = { + "dim": int(matrix.shape[1]), + "rows": int(matrix.shape[0]), + "seconds": round(elapsed, 2), + "ms_per_sequence": round(elapsed / len(frame) * 1000, 4), + "path": str(path.resolve()), + } + meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8") + log.info("%s done dim=%d in %.1f s", name, matrix.shape[1], elapsed) + + # TCRemP comes from prepare_splits and needs no encoder pass. + tcremp = dataset_dir / "embeddings.npy" + if tcremp.exists(): + matrix = np.load(tcremp, mmap_mode="r") + meta["tcremp"] = { + "dim": int(matrix.shape[1]), + "rows": int(matrix.shape[0]), + "seconds": None, + "ms_per_sequence": None, + "path": str(tcremp.resolve()), + } + meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8") + + log.info("=" * 60) + for name, info in meta.items(): + log.info("%-10s dim=%-5d rows=%d", name, info["dim"], info["rows"]) + log.info("wrote %s", meta_path) + + +if __name__ == "__main__": + main() diff --git a/benchmark/cache_external_model_embeddings.py b/benchmark/cache_external_model_embeddings.py new file mode 100644 index 0000000..5388153 --- /dev/null +++ b/benchmark/cache_external_model_embeddings.py @@ -0,0 +1,329 @@ +"""Issue 2/3: generate and cache frozen representations for the benchmark dataset. + +One pretrained checkpoint per encoder, no fine-tuning. Every matrix is written in the +row order of ``dataset.parquet`` so that ``row_index`` identifies the same sequence in +every representation, including the TCRemP matrix produced by ``prepare_splits``. + +Encoders: + onehot aligned one-hot from the repo tokenizer (gap-padded to 40 x vocab) + esm2_8m ESM-2 8M (esm2_t6_8M_UR50D), mean-pooled over residues + tcr_bert TCR-BERT (wukevin/tcr-bert), mean-pooled over residues + sceptr SCEPTR default variant, already one vector per receptor + +Transformer encoders are mean-pooled over real residues only, excluding BOS/EOS and +padding, so the pooling rule is identical across models. +""" + +import argparse +import json +import time +import warnings +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from irrm_codec.tokenization import AA_VOCAB, encode +from irrm_codec.utils import setup_logging + +ENCODERS = ("onehot", "esm2_8m", "esm2_35m", "tcr_bert", "sceptr", "sceptr_cdr3") +TCR_BERT_REVISION = "ef65ddc" + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument( + "--dataset-file", + type=Path, + help="Optional parquet/TSV cohort. Sequence column may be junction_aa or cdr3.", + ) + p.add_argument("--output-dir", default="data/benchmark/trb/representations") + p.add_argument("--encoders", nargs="+", default=list(ENCODERS), choices=list(ENCODERS)) + p.add_argument("--batch-size", type=int, default=64) + p.add_argument("--max-len", type=int, default=40) + p.add_argument("--threads", type=int, default=8) + p.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda")) + p.add_argument("--rtp-embeddings", type=Path) + p.add_argument("--overwrite", action="store_true", help="Recompute even if the cache exists.") + return p.parse_args() + + +def encode_onehot(sequences, args, log): + """Aligned one-hot: the repo's gap padding gives every sequence the same layout.""" + vocab_size = len(AA_VOCAB) + out = np.zeros((len(sequences), args.max_len * vocab_size), dtype=np.float32) + positions = np.arange(args.max_len) + for row, seq in enumerate(sequences): + block = np.zeros((args.max_len, vocab_size), dtype=np.float32) + block[positions, encode(seq, max_len=args.max_len)] = 1.0 + out[row] = block.ravel() + log.info("onehot done dim=%d", out.shape[1]) + return out + + +def _mean_pool(states, keep): + """Average hidden states over the positions flagged in ``keep``.""" + mask = keep.unsqueeze(-1).to(states.dtype) + return (states * mask).sum(1) / mask.sum(1).clamp_min(1e-9) + + +def encode_esm2(sequences, args, log, model_name="esm2_8m"): + import esm + + factories = { + "esm2_8m": esm.pretrained.esm2_t6_8M_UR50D, + "esm2_35m": esm.pretrained.esm2_t12_35M_UR50D, + } + model, alphabet = factories[model_name]() + device = choose_device(args.device) + model = model.eval().to(device) + batch_converter = alphabet.get_batch_converter() + last_layer = model.num_layers + chunks = [] + with torch.no_grad(): + for start in range(0, len(sequences), args.batch_size): + batch = sequences[start : start + args.batch_size] + _, _, tokens = batch_converter([(str(i), s) for i, s in enumerate(batch)]) + tokens = tokens.to(device) + states = model(tokens, repr_layers=[last_layer])["representations"][last_layer] + keep = ( + (tokens != alphabet.padding_idx) + & (tokens != alphabet.cls_idx) + & (tokens != alphabet.eos_idx) + ) + chunks.append(_mean_pool(states, keep).cpu().numpy().astype(np.float32)) + if start % (args.batch_size * 100) == 0: + log.info("%s %d/%d", model_name, start, len(sequences)) + return np.concatenate(chunks) + + +def encode_tcr_bert(sequences, args, log): + from transformers import AutoModel, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained("wukevin/tcr-bert", revision=TCR_BERT_REVISION) + model = AutoModel.from_pretrained("wukevin/tcr-bert", revision=TCR_BERT_REVISION) + device = choose_device(args.device) + model = model.eval().to(device) + special = set(tokenizer.all_special_ids) + chunks = [] + with torch.no_grad(): + for start in range(0, len(sequences), args.batch_size): + batch = sequences[start : start + args.batch_size] + # TCR-BERT expects residues separated by spaces. + enc = tokenizer([" ".join(s) for s in batch], return_tensors="pt", padding=True) + enc = {name: value.to(device) for name, value in enc.items()} + states = model(**enc).last_hidden_state + keep = enc["attention_mask"].bool() + for token_id in special: + keep &= enc["input_ids"] != token_id + chunks.append(_mean_pool(states, keep).cpu().numpy().astype(np.float32)) + if start % (args.batch_size * 100) == 0: + log.info("tcr_bert %d/%d", start, len(sequences)) + return np.concatenate(chunks) + + +def encode_sceptr(sequences, args, log, frame=None, cdr3_only=False): + import logging + import sceptr + + if choose_device(args.device).type == "cuda" and hasattr(sceptr, "enable_hardware_acceleration"): + sceptr.enable_hardware_acceleration() + if cdr3_only: + model = sceptr.variant.cdr3_only() + instances = pd.DataFrame({"CDR3B": frame["junction_aa"].to_numpy()}) + else: + model = sceptr + # The default author model derives beta CDR1/2 from V/J and embeds them + # together with CDR3B. Keep this annotation-aware modality explicit. + instances = pd.DataFrame( + { + "TRBV": frame["v_call"].to_numpy(), + "CDR3B": frame["junction_aa"].to_numpy(), + "TRBJ": frame["j_call"].to_numpy(), + } + ) + chunks = [] + stride = max(args.batch_size, 1024) + # libtcrlm logs an INFO line for every accepted ambiguous terminal residue; on a + # 100k-row cohort this can dominate runtime and create multi-megabyte Slurm logs. + # Suppression changes no token/model output and is limited to the SCEPTR call. + root_logger = logging.getLogger() + previous_level = root_logger.level + try: + root_logger.setLevel(logging.WARNING) + for start in range(0, len(instances), stride): + chunks.append( + model.calc_vector_representations(instances.iloc[start : start + stride]).astype(np.float32) + ) + finally: + root_logger.setLevel(previous_level) + log.info("sceptr encoded %d rows in %d chunks", len(instances), len(chunks)) + return np.concatenate(chunks) + + +def choose_device(requested): + if requested == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but no GPU is visible") + if requested == "auto": + requested = "cuda" if torch.cuda.is_available() else "cpu" + return torch.device(requested) + + +def sha256_file(path, chunk_size=1024 * 1024): + import hashlib + + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + while chunk := handle.read(chunk_size): + digest.update(chunk) + return digest.hexdigest() + + +def update_catalog(meta_path, name, entry): + """Merge one completed encoder atomically; Slurm array tasks may finish together.""" + import fcntl + + lock_path = meta_path.with_suffix(meta_path.suffix + ".lock") + lock_path.touch(exist_ok=True) + with lock_path.open("r+") as lock: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX) + current = json.loads(meta_path.read_text(encoding="utf-8")) if meta_path.exists() else {} + current[name] = entry + temporary = meta_path.with_suffix(meta_path.suffix + f".{name}.tmp") + temporary.write_text(json.dumps(current, indent=2, sort_keys=True), encoding="utf-8") + temporary.replace(meta_path) + fcntl.flock(lock.fileno(), fcntl.LOCK_UN) + return current + + +def main(): + args = parse_args() + warnings.filterwarnings("ignore") + torch.set_num_threads(args.threads) + + dataset_dir = Path(args.dataset_dir) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "cache_embeddings.log") + + dataset_file = args.dataset_file or dataset_dir / "dataset.parquet" + if dataset_file.suffix.lower() in {".tsv", ".txt"}: + frame = pd.read_csv(dataset_file, sep="\t") + else: + frame = pd.read_parquet(dataset_file) + if "junction_aa" not in frame and "cdr3" in frame: + frame = frame.rename(columns={"cdr3": "junction_aa"}) + required = {"junction_aa"} + if any(name == "sceptr" for name in args.encoders): + required.update({"v_call", "j_call"}) + missing = required.difference(frame.columns) + if missing: + raise ValueError(f"Input cohort is missing columns: {sorted(missing)}") + sequences = frame["junction_aa"].tolist() + log.info("dataset rows=%d from %s", len(frame), dataset_file) + + builders = { + "onehot": lambda: encode_onehot(sequences, args, log), + "esm2_8m": lambda: encode_esm2(sequences, args, log, "esm2_8m"), + "esm2_35m": lambda: encode_esm2(sequences, args, log, "esm2_35m"), + "tcr_bert": lambda: encode_tcr_bert(sequences, args, log), + "sceptr": lambda: encode_sceptr(sequences, args, log, frame=frame), + "sceptr_cdr3": lambda: encode_sceptr( + sequences, args, log, frame=frame, cdr3_only=True + ), + } + + meta_path = output_dir / "representations.json" + meta = json.loads(meta_path.read_text(encoding="utf-8")) if meta_path.exists() else {} + + for name in args.encoders: + path = output_dir / f"{name}.npy" + if path.exists() and not args.overwrite: + cached = np.load(path, mmap_mode="r") + if cached.shape[0] == len(frame): + log.info("%s cached dim=%d, skipping", name, cached.shape[1]) + continue + log.warning("%s cache has %d rows, expected %d; recomputing", name, cached.shape[0], len(frame)) + + log.info("encoding %s", name) + started = time.perf_counter() + matrix = builders[name]() + elapsed = time.perf_counter() - started + + if matrix.shape[0] != len(frame): + raise ValueError(f"{name} produced {matrix.shape[0]} rows, expected {len(frame)}.") + if not np.isfinite(matrix).all(): + raise ValueError(f"{name} produced non-finite values.") + + np.save(path, matrix) + entry = { + "dim": int(matrix.shape[1]), + "rows": int(matrix.shape[0]), + "seconds": round(elapsed, 2), + "ms_per_sequence": round(elapsed / len(frame) * 1000, 4), + "path": str(path.resolve()), + "sha256": sha256_file(path), + "sequence_sha256": __import__("hashlib").sha256( + "\n".join(sequences).encode("utf-8") + ).hexdigest(), + } + entry["model"] = { + "esm2_8m": {"id": "esm2_t6_8M_UR50D", "library": "fair-esm==2.0.0"}, + "esm2_35m": {"id": "esm2_t12_35M_UR50D", "library": "fair-esm==2.0.0"}, + "tcr_bert": {"id": "wukevin/tcr-bert", "revision": TCR_BERT_REVISION}, + "sceptr": { + "id": "SCEPTR default native CLS64", + "library": "sceptr==1.2.0", + "modality": "CDR3+V+J", + "pooling": "author-native CLS", + }, + "sceptr_cdr3": { + "id": "SCEPTR cdr3_only native CLS64", + "library": "sceptr==1.2.0", + "modality": "CDR3 sequence only", + "pooling": "author-native CLS", + }, + "onehot": {"id": "aligned repository one-hot", "modality": "sequence"}, + }[name] + meta = update_catalog(meta_path, name, entry) + log.info("%s done dim=%d in %.1f s", name, matrix.shape[1], elapsed) + + # TCRemP comes from prepare_splits and needs no encoder pass. + tcremp = dataset_dir / "embeddings.npy" + if tcremp.exists(): + matrix = np.load(tcremp, mmap_mode="r") + entry = { + "dim": int(matrix.shape[1]), + "rows": int(matrix.shape[0]), + "seconds": None, + "ms_per_sequence": None, + "path": str(tcremp.resolve()), + "sha256": sha256_file(tcremp), + } + meta = update_catalog(meta_path, "tcremp", entry) + + if args.rtp_embeddings: + matrix = np.load(args.rtp_embeddings, mmap_mode="r") + if matrix.shape[0] != len(frame) or not np.isfinite(matrix).all(): + raise ValueError(f"RTP matrix failed shape/finiteness gate: {matrix.shape}") + entry = { + "dim": int(matrix.shape[1]), + "rows": int(matrix.shape[0]), + "seconds": 0.0, + "ms_per_sequence": 0.0, + "path": str(args.rtp_embeddings.resolve()), + "sha256": sha256_file(args.rtp_embeddings), + "checkpoint_role": "DATA-ANCHOR R+T+P seed42 validation-selected frozen encoder", + } + meta = update_catalog(meta_path, "rtp", entry) + + log.info("=" * 60) + for name, info in meta.items(): + log.info("%-10s dim=%-5d rows=%d", name, info["dim"], info["rows"]) + log.info("wrote %s", meta_path) + + +if __name__ == "__main__": + main() diff --git a/benchmark/cache_external_residue_embeddings.py b/benchmark/cache_external_residue_embeddings.py new file mode 100644 index 0000000..5266955 --- /dev/null +++ b/benchmark/cache_external_residue_embeddings.py @@ -0,0 +1,250 @@ +"""Cache frozen per-residue states without global pooling. + +Outputs are float16 ``[N, 40, D]`` arrays plus boolean masks. The float16 +storage is an I/O choice only; heads cast each batch to float32 before use. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from irrm_codec.utils import setup_logging + + +ENCODERS = ("esm2_8m", "tcr_bert", "sceptr") +DIMENSIONS = {"esm2_8m": 320, "tcr_bert": 768, "sceptr": 64} +TCR_BERT_REVISION = "ef65ddc" + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-dir", default="data/benchmark/trb") + parser.add_argument("--output-dir", required=True) + parser.add_argument("--encoders", nargs="+", choices=ENCODERS, default=list(ENCODERS)) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--max-len", type=int, default=40) + parser.add_argument("--threads", type=int, default=8) + parser.add_argument("--device", choices=("cpu", "cuda", "auto"), default="auto") + parser.add_argument("--limit", type=int) + parser.add_argument("--overwrite", action="store_true") + return parser.parse_args() + + +def sha256_file(path, chunk_size=1024 * 1024): + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + while block := handle.read(chunk_size): + digest.update(block) + return digest.hexdigest() + + +def choose_device(requested): + if requested == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but unavailable") + if requested == "auto": + requested = "cuda" if torch.cuda.is_available() else "cpu" + return torch.device(requested) + + +def open_outputs(output_dir, name, n_rows, max_len, dim): + feature_tmp = output_dir / f"{name}_residue.tmp.npy" + mask_tmp = output_dir / f"{name}_residue_mask.tmp.npy" + features = np.lib.format.open_memmap( + feature_tmp, mode="w+", dtype=np.float16, shape=(n_rows, max_len, dim) + ) + mask = np.lib.format.open_memmap( + mask_tmp, mode="w+", dtype=np.bool_, shape=(n_rows, max_len) + ) + features[:] = 0 + mask[:] = False + return features, mask, feature_tmp, mask_tmp + + +def assign_rows(features, mask, start, states, keep, sequences): + for offset, sequence in enumerate(sequences): + selected = states[offset][keep[offset]].detach().cpu().numpy() + if len(selected) != len(sequence): + raise ValueError( + f"Residue-state count {len(selected)} != sequence length {len(sequence)}" + ) + features[start + offset, : len(sequence)] = selected.astype(np.float16) + mask[start + offset, : len(sequence)] = True + + +def encode_esm2(frame, args, features, mask, log): + import esm + + model, alphabet = esm.pretrained.esm2_t6_8M_UR50D() + device = choose_device(args.device) + model = model.eval().to(device) + converter = alphabet.get_batch_converter() + sequences = frame.junction_aa.tolist() + with torch.no_grad(): + for start in range(0, len(sequences), args.batch_size): + batch = sequences[start : start + args.batch_size] + _, _, tokens = converter([(str(i), seq) for i, seq in enumerate(batch)]) + tokens = tokens.to(device) + states = model(tokens, repr_layers=[model.num_layers])["representations"][model.num_layers] + keep = ( + (tokens != alphabet.padding_idx) + & (tokens != alphabet.cls_idx) + & (tokens != alphabet.eos_idx) + ) + assign_rows(features, mask, start, states, keep, batch) + if start % (args.batch_size * 100) == 0: + log.info("esm2_8m residue states %d/%d", start, len(sequences)) + + +def encode_tcr_bert(frame, args, features, mask, log): + from transformers import AutoModel, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained("wukevin/tcr-bert", revision=TCR_BERT_REVISION) + model = AutoModel.from_pretrained("wukevin/tcr-bert", revision=TCR_BERT_REVISION) + device = choose_device(args.device) + model = model.eval().to(device) + specials = set(tokenizer.all_special_ids) + sequences = frame.junction_aa.tolist() + with torch.no_grad(): + for start in range(0, len(sequences), args.batch_size): + batch = sequences[start : start + args.batch_size] + encoded = tokenizer([" ".join(seq) for seq in batch], return_tensors="pt", padding=True) + encoded = {key: value.to(device) for key, value in encoded.items()} + states = model(**encoded).last_hidden_state + keep = encoded["attention_mask"].bool() + for token_id in specials: + keep &= encoded["input_ids"] != token_id + assign_rows(features, mask, start, states, keep, batch) + if start % (args.batch_size * 100) == 0: + log.info("tcr_bert residue states %d/%d", start, len(sequences)) + + +def encode_sceptr(frame, args, features, mask, log): + import logging + import sceptr + + if choose_device(args.device).type == "cuda": + sceptr.enable_hardware_acceleration() + else: + sceptr.disable_hardware_acceleration() + instances = pd.DataFrame( + { + "TRBV": frame.v_call.to_numpy(), + "CDR3B": frame.junction_aa.to_numpy(), + "TRBJ": frame.j_call.to_numpy(), + } + ) + sequences = frame.junction_aa.tolist() + # libtcrlm emits one INFO message for many accepted ambiguous terminal residues. + # Suppress only that internal chatter; it does not change tokenisation or outputs. + root_logger = logging.getLogger() + previous_level = root_logger.level + try: + root_logger.setLevel(logging.WARNING) + for start in range(0, len(instances), args.batch_size): + batch = instances.iloc[start : start + args.batch_size] + result = sceptr.calc_residue_representations(batch) + for offset, sequence in enumerate(sequences[start : start + len(batch)]): + selected = result.representation_array[offset][result.compartment_mask[offset] == 6] + if len(selected) != len(sequence): + raise ValueError( + f"SCEPTR CDR3B state count {len(selected)} != sequence length {len(sequence)}" + ) + features[start + offset, : len(sequence)] = selected.astype(np.float16) + mask[start + offset, : len(sequence)] = True + if start % (args.batch_size * 100) == 0: + log.warning("sceptr residue states %d/%d", start, len(instances)) + finally: + root_logger.setLevel(previous_level) + + +def update_catalog(path, name, entry): + import fcntl + + lock_path = path.with_suffix(path.suffix + ".lock") + lock_path.touch(exist_ok=True) + with lock_path.open("r+") as lock: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX) + current = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {} + current[name] = entry + temporary = path.with_suffix(path.suffix + f".{name}.tmp") + temporary.write_text(json.dumps(current, indent=2, sort_keys=True), encoding="utf-8") + temporary.replace(path) + fcntl.flock(lock.fileno(), fcntl.LOCK_UN) + + +def main(): + args = parse_args() + torch.set_num_threads(args.threads) + dataset_dir = Path(args.dataset_dir) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "cache_residue_embeddings.log") + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + if args.limit is not None: + frame = frame.iloc[: args.limit].copy() + if frame.junction_aa.str.len().max() > args.max_len: + raise ValueError("A sequence exceeds max_len") + sequence_sha = hashlib.sha256("\n".join(frame.junction_aa).encode()).hexdigest() + builders = {"esm2_8m": encode_esm2, "tcr_bert": encode_tcr_bert, "sceptr": encode_sceptr} + catalog_path = output_dir / "residue_representations.json" + + for name in args.encoders: + feature_path = output_dir / f"{name}_residue.npy" + mask_path = output_dir / f"{name}_residue_mask.npy" + if feature_path.exists() and mask_path.exists() and not args.overwrite: + cached = np.load(feature_path, mmap_mode="r") + if cached.shape == (len(frame), args.max_len, DIMENSIONS[name]): + log.info("%s residue cache already complete", name) + continue + started = time.perf_counter() + features, mask, feature_tmp, mask_tmp = open_outputs( + output_dir, name, len(frame), args.max_len, DIMENSIONS[name] + ) + builders[name](frame, args, features, mask, log) + features.flush() + mask.flush() + del features, mask + feature_tmp.replace(feature_path) + mask_tmp.replace(mask_path) + features = np.load(feature_path, mmap_mode="r") + mask = np.load(mask_path, mmap_mode="r") + lengths = frame.junction_aa.str.len().to_numpy() + if not np.array_equal(mask.sum(axis=1), lengths): + raise ValueError(f"{name} mask counts do not equal CDR3 lengths") + if not np.isfinite(features).all(): + raise ValueError(f"{name} residue cache contains non-finite values") + update_catalog( + catalog_path, + name, + { + "model": { + "esm2_8m": "esm2_t6_8M_UR50D", + "tcr_bert": f"wukevin/tcr-bert@{TCR_BERT_REVISION}", + "sceptr": "sceptr==1.2.0 default penultimate residue states; CDR3B mask=6", + }[name], + "rows": len(frame), + "shape": list(features.shape), + "dtype": str(features.dtype), + "mask_true": int(mask.sum()), + "sequence_sha256": sequence_sha, + "features": str(feature_path.resolve()), + "mask": str(mask_path.resolve()), + "features_sha256": sha256_file(feature_path), + "mask_sha256": sha256_file(mask_path), + "seconds": round(time.perf_counter() - started, 2), + "pooling": "none", + }, + ) + log.info("%s accepted shape=%s", name, features.shape) + + +if __name__ == "__main__": + main() diff --git a/benchmark/collect_author_global_comparison.py b/benchmark/collect_author_global_comparison.py new file mode 100644 index 0000000..daeed52 --- /dev/null +++ b/benchmark/collect_author_global_comparison.py @@ -0,0 +1,245 @@ +"""Assemble the preregistered author-native global embedding comparison table.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd + +METHODS = [ + ("rtp", "DATA-ANCHOR RTP", "sequence-only", 128), + ("tcremp", "TCRemP prototype distances", "sequence-derived alignment", 9000), + ("tcr_bert", "TCR-BERT final residue mean", "sequence-only", 768), + ("esm2_35m", "ESM2-35M final residue mean", "sequence-only", 480), + ("esm2_8m", "ESM2-8M final residue mean", "sequence-only sensitivity", 320), + ("sceptr_cdr3", "SCEPTR CDR3-only native CLS", "sequence-only", 64), + ("sceptr", "SCEPTR default native CLS", "annotation-aware V+CDR3+J", 64), + ("tfidf", "CDR3 TF-IDF ridge", "sequence control", None), +] + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def load_json(path: Path) -> dict: + return json.loads(path.read_text(encoding="utf-8")) + + +def aggregate_runs(root: Path, key: str, expected: set[str]) -> pd.DataFrame: + rows = [] + for path in sorted(root.glob("*/metrics.json")): + data = load_json(path) + name = data[key] + if name not in expected: + continue + row = {"name": name, "seed": int(data["seed"])} + row.update({f"test_{metric}": value for metric, value in data["test"].items()}) + rows.append(row) + frame = pd.DataFrame(rows) + for name in expected: + seeds = sorted(frame.loc[frame["name"] == name, "seed"].tolist()) + if seeds != [42, 43, 44]: + raise ValueError(f"{name} seeds are {seeds}, expected [42, 43, 44]") + return frame + + +def format_cell(mean, std, digits=4): + if pd.isna(mean): + return "NA" + if pd.isna(std): + return f"{mean:.{digits}f}" + return f"{mean:.{digits}f} +/- {std:.{digits}f}" + + +def markdown_table(frame: pd.DataFrame) -> str: + columns = [ + "method", "modality", "reconstruction_exact", "pgen_rmse", + "vdjdb_p_at_1", "vdjdb_p_at_5", "vdjdb_p_at_10", + "vdjdb_map", "vdjdb_within_between_ratio", + ] + labels = [ + "Method", "Modality", "Reconstruction exact", "pgen RMSE", + "VDJdb P@1", "P@5", "P@10", "MAP", "Within/between", + ] + lines = ["| " + " | ".join(labels) + " |", "|" + "|".join(["---"] * len(labels)) + "|"] + for _, row in frame.iterrows(): + values = [str(row[col]) if not pd.isna(row[col]) else "NA" for col in columns] + lines.append("| " + " | ".join(values) + " |") + return "\n".join(lines) + "\n" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--old-reconstruction-summary", type=Path, required=True) + parser.add_argument("--old-pgen-summary", type=Path, required=True) + parser.add_argument("--new-reconstruction-dir", type=Path, required=True) + parser.add_argument("--new-pgen-dir", type=Path, required=True) + parser.add_argument("--control-pgen-dir", type=Path, required=True) + parser.add_argument("--reference-vdjdb-metrics", type=Path, required=True) + parser.add_argument("--external-vdjdb-summary", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.output_dir.exists(): + raise FileExistsError(args.output_dir) + args.output_dir.mkdir(parents=True) + + reconstruction = pd.read_csv(args.old_reconstruction_summary) + new_recon = aggregate_runs( + args.new_reconstruction_dir, "representation", {"esm2_35m", "sceptr_cdr3"} + ) + recon_new_rows = [] + for name, block in new_recon.groupby("name"): + recon_new_rows.append({ + "representation": name, + "seeds": len(block), + "exact_match_mean": block["test_exact_match"].mean(), + "exact_match_std": block["test_exact_match"].std(ddof=1), + "token_accuracy_mean": block["test_token_accuracy"].mean(), + "token_accuracy_std": block["test_token_accuracy"].std(ddof=1), + }) + reconstruction = pd.concat([reconstruction, pd.DataFrame(recon_new_rows)], ignore_index=True) + + pgen = pd.read_csv(args.old_pgen_summary) + primary_pgen = aggregate_runs( + args.new_pgen_dir, + "arm", + {"esm2_35m_mlp", "sceptr_cdr3_mlp", "tcremp_mlp"}, + ) + control_pgen = aggregate_runs(args.control_pgen_dir, "arm", {"onehot_mlp"}) + new_pgen = pd.concat([primary_pgen, control_pgen], ignore_index=True) + pgen_new_rows = [] + for arm, block in new_pgen.groupby("name"): + pgen_new_rows.append({ + "arm": arm, + "seeds": len(block), + "rmse_mean": block["test_rmse"].mean(), + "rmse_std": block["test_rmse"].std(ddof=1), + "r2_mean": block["test_r2"].mean(), + "r2_std": block["test_r2"].std(ddof=1), + "spearman_rho_mean": block["test_spearman_rho"].mean(), + "spearman_rho_std": block["test_spearman_rho"].std(ddof=1), + }) + pgen = pd.concat([pgen, pd.DataFrame(pgen_new_rows)], ignore_index=True) + + external = pd.read_csv(args.external_vdjdb_summary, sep="\t").set_index("model") + reference = load_json(args.reference_vdjdb_metrics)["rtp"]["cosine"] + rows = [] + for key, label, modality, native_dim in METHODS: + recon = reconstruction[reconstruction["representation"] == key] + pgen_arm = "tfidf_ridge" if key == "tfidf" else f"{key}_mlp" + pgen_row = pgen[pgen["arm"] == pgen_arm] + row = { + "condition": key, + "method": label, + "modality": modality, + "native_global_dim": native_dim, + "reconstruction_common_dim": 64 if len(recon) else np.nan, + "reconstruction_exact_mean": recon["exact_match_mean"].iloc[0] if len(recon) else np.nan, + "reconstruction_exact_std": recon["exact_match_std"].iloc[0] if len(recon) else np.nan, + "reconstruction_token_accuracy_mean": recon["token_accuracy_mean"].iloc[0] if len(recon) else np.nan, + "pgen_rmse_mean": pgen_row["rmse_mean"].iloc[0] if len(pgen_row) else np.nan, + "pgen_rmse_std": pgen_row["rmse_std"].iloc[0] if len(pgen_row) else np.nan, + "pgen_r2_mean": pgen_row["r2_mean"].iloc[0] if len(pgen_row) else np.nan, + "pgen_spearman_mean": pgen_row["spearman_rho_mean"].iloc[0] if len(pgen_row) else np.nan, + } + if key == "rtp": + macro = reference["macro_epitope"] + knn = reference["knn_classification"] + elif key in external.index: + macro = { + metric: external.loc[key, f"cosine_{metric}"] + for metric in ("precision_at_1", "precision_at_5", "precision_at_10", + "average_precision", "within_between_ratio") + } + knn = None + else: + macro = knn = None + for source, target in ( + ("precision_at_1", "vdjdb_cosine_precision_at_1"), + ("precision_at_5", "vdjdb_cosine_precision_at_5"), + ("precision_at_10", "vdjdb_cosine_precision_at_10"), + ("average_precision", "vdjdb_cosine_map"), + ("within_between_ratio", "vdjdb_cosine_within_between_ratio"), + ): + row[target] = ( + float(macro[source]) + if macro is not None + else np.nan + ) + row["vdjdb_cosine_knn_macro_f1_at_5"] = ( + float(knn["k_5"]["macro_f1"]) if knn is not None + else ( + float(external.loc[key, "cosine_knn_macro_f1_at_5"]) + if key in external.index else np.nan + ) + ) + rows.append(row) + + final = pd.DataFrame(rows) + final["reconstruction_exact"] = [ + format_cell(mean, std) for mean, std in zip( + final["reconstruction_exact_mean"], final["reconstruction_exact_std"] + ) + ] + final["pgen_rmse"] = [ + format_cell(mean, std) for mean, std in zip(final["pgen_rmse_mean"], final["pgen_rmse_std"]) + ] + for src, dst in ( + ("vdjdb_cosine_precision_at_1", "vdjdb_p_at_1"), + ("vdjdb_cosine_precision_at_5", "vdjdb_p_at_5"), + ("vdjdb_cosine_precision_at_10", "vdjdb_p_at_10"), + ("vdjdb_cosine_map", "vdjdb_map"), + ("vdjdb_cosine_within_between_ratio", "vdjdb_within_between_ratio"), + ): + final[dst] = final[src].map(lambda value: "NA" if pd.isna(value) else f"{value:.4f}") + + table_path = args.output_dir / "final_author_global_comparison.tsv" + final.to_csv(table_path, sep="\t", index=False, lineterminator="\n") + md_path = args.output_dir / "final_author_global_comparison.md" + md_path.write_text(markdown_table(final), encoding="utf-8") + manifest = { + "status": "complete", + "protocol": { + "reconstruction": "train-only PCA64 plus identical learned decoder; three probe seeds", + "pgen": "native author-global vector plus identical MLP; three probe seeds", + "vdjdb": "native author-global vector on the fixed REDCEA gallery; macro over 87 epitopes", + }, + "inputs": { + str(path): sha256(path) + for path in ( + args.old_reconstruction_summary, + args.old_pgen_summary, + args.reference_vdjdb_metrics, + args.external_vdjdb_summary, + ) + }, + "outputs": {table_path.name: sha256(table_path), md_path.name: sha256(md_path)}, + "limitations": [ + "SCEPTR-default is annotation-aware and is not equivalent to sequence-only methods.", + "SCEPTR-default was not run on full REDCEA because unsupported V/J rows would require method-specific filtering.", + "REDCEA membership is sequence/TCRemP-cluster-conditioned, so retrieval is motif-conditioned.", + "Aligned one-hot is retained as an internal high-dimensional sequence control but omitted from the manuscript-facing comparison.", + "Probe mean/std reflects probe seeds, not pretrained encoder training-seed uncertainty.", + ], + } + (args.output_dir / "MANIFEST.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + print(md_path.read_text(encoding="utf-8")) + + +if __name__ == "__main__": + main() diff --git a/benchmark/collect_external_pgen.py b/benchmark/collect_external_pgen.py new file mode 100644 index 0000000..2879daa --- /dev/null +++ b/benchmark/collect_external_pgen.py @@ -0,0 +1,176 @@ +"""Issue 3: aggregate Pgen runs into the per-run CSV and the summary tables. + +Reads every ``metrics.json`` written by ``train_external_model_pgen`` and produces: + pgen_runs.csv one row per run (arm x target x training size x seed) + pgen_summary.csv mean and std across seeds + pgen_summary.md Markdown tables for REPORT.md, one per target + +The Markdown groups by training size so the low-data question is readable directly: +whether the pretrained IRRM arm's advantage over the from-scratch arm grows as the +training set shrinks. +""" + +import argparse +import json +from pathlib import Path + +import pandas as pd +from scipy import stats + +METRICS = ("rmse", "mae", "r2", "pearson_r", "spearman_rho", "bias") +SUBSET_ORDER = {"1k": 0, "10k": 1, "all": 2} + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--runs-dir", default="artifacts/benchmark/pgen") + p.add_argument("--output-dir", default="results/pgen") + return p.parse_args() + + +def load_runs(runs_dir): + rows = [] + for path in sorted(Path(runs_dir).glob("*/metrics.json")): + run = json.loads(path.read_text(encoding="utf-8")) + memory = run.get("memory") or {} + rows.append( + { + "arm": run["arm"], + "target": run["target"], + "train_subset": run["train_subset"], + "train_size": run["train_size"], + "seed": run["seed"], + "training_seconds": run["training_seconds"], + **{name: run["test"][name] for name in METRICS}, + "inference_seq_per_second": run["inference"]["sequences_per_second"], + "inference_ms_per_sequence": run["inference"]["ms_per_sequence"], + "device": run["inference"]["device"], + "peak_rss_mb": memory.get("peak_rss_mb"), + "peak_gpu_mb": memory.get("peak_gpu_mb"), + "epochs_ran": (run.get("arm_details") or {}).get("epochs_ran"), + } + ) + if not rows: + raise SystemExit(f"No metrics.json found under {runs_dir}.") + frame = pd.DataFrame(rows) + frame["subset_order"] = frame.train_subset.map(SUBSET_ORDER).fillna(99) + return frame.sort_values(["target", "subset_order", "arm", "seed"]) + + +def summarize(runs): + aggregated = ( + runs.groupby(["target", "train_subset", "subset_order", "arm"]) + .agg( + seeds=("seed", "count"), + train_size=("train_size", "first"), + **{f"{name}_{stat}": (name, stat) for name in METRICS for stat in ("mean", "std")}, + training_seconds_mean=("training_seconds", "mean"), + inference_seq_per_second_mean=("inference_seq_per_second", "mean"), + peak_rss_mb_mean=("peak_rss_mb", "mean"), + peak_gpu_mb_mean=("peak_gpu_mb", "mean"), + ) + .reset_index() + ) + return aggregated.sort_values(["target", "subset_order", "rmse_mean"]) + + +def to_markdown(summary): + blocks = [] + for target in summary.target.unique(): + blocks.append(f"### {target}\n") + for subset in summary[summary.target == target].train_subset.unique(): + rows = summary[(summary.target == target) & (summary.train_subset == subset)] + size = int(rows.train_size.iloc[0]) + blocks.append(f"**train = {subset} ({size:,} sequences)**\n") + blocks.append( + "| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB |\n" + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |" + ) + for _, row in rows.iterrows(): + std = row["rmse_std"] + rmse = f"{row['rmse_mean']:.4f}" + (f" ± {std:.4f}" if pd.notna(std) else "") + cells = [ + row["arm"], + rmse, + f"{row['mae_mean']:.4f}", + f"{row['r2_mean']:.4f}", + f"{row['pearson_r_mean']:.4f}", + f"{row['spearman_rho_mean']:.4f}", + f"{row['bias_mean']:+.4f}", + f"{row['training_seconds_mean']:.0f}", + f"{row['inference_seq_per_second_mean']:.0f}", + f"{row['peak_rss_mb_mean']:.0f}" if pd.notna(row["peak_rss_mb_mean"]) else "-", + ] + blocks.append("| " + " | ".join(cells) + " |") + blocks.append("") + return "\n".join(blocks) + "\n" + + +def pretraining_delta(runs): + """Does IRRM pretraining help? RMSE difference with a significance test. + + Seed spread for these two arms is comparable to the effect being measured, so the + difference of means alone cannot answer the question. Welch's t-test does not assume + the two arms have equal variance, which they do not. + """ + rows = [] + for (target, subset, order), group in runs.groupby(["target", "train_subset", "subset_order"]): + scratch = group.loc[group.arm == "irrm_scratch", "rmse"].to_numpy() + pretrained = group.loc[group.arm == "irrm_pretrained", "rmse"].to_numpy() + if len(scratch) < 2 or len(pretrained) < 2: + continue + # Lower RMSE is better, so a positive delta means pretraining won. + delta = scratch.mean() - pretrained.mean() + _, p_value = stats.ttest_ind(scratch, pretrained, equal_var=False) + rows.append( + { + "target": target, + "train_subset": subset, + "subset_order": order, + "n_scratch": len(scratch), + "n_pretrained": len(pretrained), + "irrm_scratch": scratch.mean(), + "irrm_pretrained": pretrained.mean(), + "rmse_delta": delta, + "relative_improvement": delta / scratch.mean(), + "p_value": p_value, + "verdict": ( + "no detectable effect" + if p_value >= 0.05 + else ("pretraining helps" if delta > 0 else "pretraining hurts") + ), + } + ) + if not rows: + return None + return pd.DataFrame(rows).sort_values(["target", "subset_order"]) + + +def main(): + args = parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + runs = load_runs(args.runs_dir) + runs.drop(columns=["subset_order"]).to_csv(output_dir / "pgen_runs.csv", index=False) + + summary = summarize(runs) + summary.drop(columns=["subset_order"]).to_csv(output_dir / "pgen_summary.csv", index=False) + (output_dir / "pgen_summary.md").write_text(to_markdown(summary), encoding="utf-8") + + print(f"runs: {len(runs)} arms: {runs.arm.nunique()} targets: {runs.target.nunique()}") + delta = pretraining_delta(runs) + if delta is not None: + delta.drop(columns=["subset_order"]).to_csv(output_dir / "pretraining_delta.csv", index=False) + print("\nDoes IRRM pretraining help? (positive rmse_delta means pretraining wins)") + print( + delta[ + ["target", "train_subset", "n_scratch", "n_pretrained", "irrm_scratch", + "irrm_pretrained", "rmse_delta", "p_value", "verdict"] + ].to_string(index=False) + ) + print(f"\nwrote {output_dir}/pgen_runs.csv, pgen_summary.csv, pgen_summary.md") + + +if __name__ == "__main__": + main() diff --git a/benchmark/collect_external_reconstruction.py b/benchmark/collect_external_reconstruction.py new file mode 100644 index 0000000..0c41543 --- /dev/null +++ b/benchmark/collect_external_reconstruction.py @@ -0,0 +1,148 @@ +"""Issue 2: aggregate reconstruction runs into the per-run CSV and the summary table. + +Reads every ``metrics.json`` written by ``train_decoder`` and produces: + reconstruction_runs.csv one row per run (representation x seed) + reconstruction_summary.csv mean and std across seeds, per representation + reconstruction_summary.md the same table in Markdown, for REPORT.md + +Explained variance from the PCA fit is joined in, because it is the main explanation +for differences at a fixed bottleneck width: a representation that survives the +projection intact starts from a different place than one that loses most of its variance. +""" + +import argparse +import json +from pathlib import Path + +import pandas as pd + +METRICS = ( + "exact_match", + "token_accuracy", + "normalized_levenshtein", + "within_edit_distance_1", + "mean_edit_distance", +) + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--runs-dir", default="artifacts/benchmark/reconstruction") + p.add_argument("--bottleneck-dir", default="data/benchmark/trb/bottleneck") + p.add_argument("--output-dir", default="results/reconstruction") + p.add_argument("--bottleneck-dim", type=int, default=64) + return p.parse_args() + + +def load_runs(runs_dir, bottleneck_dim): + rows = [] + for path in sorted(Path(runs_dir).glob("*/metrics.json")): + run = json.loads(path.read_text(encoding="utf-8")) + if run["bottleneck_dim"] != bottleneck_dim: + continue + rows.append( + { + "representation": run["representation"], + "seed": run["seed"], + "bottleneck_dim": run["bottleneck_dim"], + "train_subset": run["train_subset"], + "train_size": run["train_size"], + "epochs_ran": run["epochs_ran"], + "best_epoch": run["best_epoch"], + "best_val_loss": run["best_val_loss"], + "training_seconds": run["training_seconds"], + "test_loss": run["test"]["loss"], + **{name: run["test"][name] for name in METRICS}, + "inference_seq_per_second": run["inference"]["sequences_per_second"], + "inference_ms_per_sequence": run["inference"]["ms_per_sequence"], + "device": run["inference"]["device"], + } + ) + if not rows: + raise SystemExit(f"No metrics.json found under {runs_dir} for dim={bottleneck_dim}.") + return pd.DataFrame(rows) + + +def summarize(runs, bottleneck_dir, bottleneck_dim): + aggregated = runs.groupby("representation").agg( + seeds=("seed", "count"), + **{f"{name}_{stat}": (name, stat) for name in METRICS for stat in ("mean", "std")}, + training_seconds_mean=("training_seconds", "mean"), + epochs_mean=("epochs_ran", "mean"), + inference_seq_per_second_mean=("inference_seq_per_second", "mean"), + ) + + summary_path = Path(bottleneck_dir) / "bottleneck_summary.json" + if summary_path.exists(): + bottleneck = json.loads(summary_path.read_text(encoding="utf-8")) + aggregated["source_dim"] = [ + bottleneck.get(name, {}).get("source_dim") for name in aggregated.index + ] + aggregated["explained_variance"] = [ + bottleneck.get(name, {}).get("explained_variance", {}).get(str(bottleneck_dim)) + for name in aggregated.index + ] + + return aggregated.sort_values("exact_match_mean", ascending=False).reset_index() + + +def to_markdown(summary, bottleneck_dim): + header = ( + f"### Reconstruction at a {bottleneck_dim}-dimensional bottleneck\n\n" + "| representation | source dim | expl. var | exact match | token acc | " + "norm. Levenshtein | within ED 1 | train s | seq/s |\n" + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |\n" + ) + + def cell(row, name): + std = row[f"{name}_std"] + spread = f" ± {std:.4f}" if pd.notna(std) else "" + return f"{row[f'{name}_mean']:.4f}{spread}" + + lines = [] + for _, row in summary.iterrows(): + expl = row.get("explained_variance") + source_dim = row.get("source_dim") + cells = [ + row["representation"], + f"{source_dim:.0f}" if pd.notna(source_dim) else "-", + f"{expl:.3f}" if pd.notna(expl) else "-", + cell(row, "exact_match"), + cell(row, "token_accuracy"), + cell(row, "normalized_levenshtein"), + cell(row, "within_edit_distance_1"), + f"{row['training_seconds_mean']:.0f}", + f"{row['inference_seq_per_second_mean']:.0f}", + ] + lines.append("| " + " | ".join(cells) + " |") + return header + "\n".join(lines) + "\n" + + +def main(): + args = parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + runs = load_runs(args.runs_dir, args.bottleneck_dim) + runs = runs.sort_values(["representation", "seed"]) + runs.to_csv(output_dir / "reconstruction_runs.csv", index=False) + + summary = summarize(runs, args.bottleneck_dir, args.bottleneck_dim) + summary.to_csv(output_dir / "reconstruction_summary.csv", index=False) + (output_dir / "reconstruction_summary.md").write_text( + to_markdown(summary, args.bottleneck_dim), encoding="utf-8" + ) + + print(f"runs: {len(runs)} representations: {runs.representation.nunique()}") + print( + summary[ + ["representation", "seeds", "exact_match_mean", "exact_match_std", "token_accuracy_mean"] + ].to_string(index=False) + ) + print(f"\nwrote {output_dir}/reconstruction_runs.csv") + print(f"wrote {output_dir}/reconstruction_summary.csv") + print(f"wrote {output_dir}/reconstruction_summary.md") + + +if __name__ == "__main__": + main() diff --git a/benchmark/collect_external_residue_benchmark.py b/benchmark/collect_external_residue_benchmark.py new file mode 100644 index 0000000..cd916f0 --- /dev/null +++ b/benchmark/collect_external_residue_benchmark.py @@ -0,0 +1,77 @@ +"""Collect unpooled reconstruction and pgen runs into auditable tables.""" + +import argparse +import json +from pathlib import Path + +import pandas as pd + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reconstruction-dir", required=True) + parser.add_argument("--pgen-dir", required=True) + parser.add_argument("--output-dir", required=True) + return parser.parse_args() + + +def collect(root, task): + rows = [] + for path in sorted(Path(root).glob("*/metrics.json")): + payload = json.loads(path.read_text(encoding="utf-8")) + base = { + "representation": payload["representation"], + "seed": payload["seed"], + "interface": payload["interface"], + "parameters": payload["parameters"], + "best_epoch": payload["best_epoch"], + "best_val_loss": payload["best_val_loss"], + "epochs_ran": payload["epochs_ran"], + "training_seconds": payload["training_seconds"], + "metrics_path": str(path), + } + if task == "reconstruction": + names = ("loss", "token_accuracy", "exact_match", "normalized_levenshtein", "within_edit_distance_1", "mean_edit_distance") + else: + names = ("rmse", "mae", "r2", "pearson_r", "spearman_rho", "bias") + rows.append({**base, **{name: payload["test"][name] for name in names}}) + if not rows: + raise ValueError(f"No {task} metrics found under {root}") + return pd.DataFrame(rows) + + +def summarize(frame, task): + if task == "reconstruction": + names = ("token_accuracy", "exact_match", "normalized_levenshtein", "within_edit_distance_1", "mean_edit_distance") + ascending = False + key = "exact_match_mean" + else: + names = ("rmse", "mae", "r2", "pearson_r", "spearman_rho", "bias") + ascending = True + key = "rmse_mean" + result = frame.groupby("representation").agg( + seeds=("seed", "count"), + parameters=("parameters", "first"), + training_seconds_mean=("training_seconds", "mean"), + **{f"{name}_{stat}": (name, stat) for name in names for stat in ("mean", "std")}, + ).reset_index() + return result.sort_values(key, ascending=ascending) + + +def main(): + args = parse_args() + output = Path(args.output_dir) + output.mkdir(parents=True, exist_ok=True) + reconstruction = collect(args.reconstruction_dir, "reconstruction") + pgen = collect(args.pgen_dir, "pgen") + reconstruction.to_csv(output / "reconstruction_residue_runs.tsv", sep="\t", index=False) + pgen.to_csv(output / "pgen_residue_runs.tsv", sep="\t", index=False) + summarize(reconstruction, "reconstruction").to_csv( + output / "reconstruction_residue_summary.tsv", sep="\t", index=False + ) + summarize(pgen, "pgen").to_csv(output / "pgen_residue_summary.tsv", sep="\t", index=False) + print(f"status=complete reconstruction_runs={len(reconstruction)} pgen_runs={len(pgen)}") + + +if __name__ == "__main__": + main() diff --git a/benchmark/collect_pgen.py b/benchmark/collect_pgen.py new file mode 100644 index 0000000..d2f1db5 --- /dev/null +++ b/benchmark/collect_pgen.py @@ -0,0 +1,176 @@ +"""Issue 3: aggregate Pgen runs into the per-run CSV and the summary tables. + +Reads every ``metrics.json`` written by ``train_pgen_arm`` and produces: + pgen_runs.csv one row per run (arm x target x training size x seed) + pgen_summary.csv mean and std across seeds + pgen_summary.md Markdown tables for REPORT.md, one per target + +The Markdown groups by training size so the low-data question is readable directly: +whether the pretrained IRRM arm's advantage over the from-scratch arm grows as the +training set shrinks. +""" + +import argparse +import json +from pathlib import Path + +import pandas as pd +from scipy import stats + +METRICS = ("rmse", "mae", "r2", "pearson_r", "spearman_rho", "bias") +SUBSET_ORDER = {"1k": 0, "10k": 1, "all": 2} + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--runs-dir", default="artifacts/benchmark/pgen") + p.add_argument("--output-dir", default="results/pgen") + return p.parse_args() + + +def load_runs(runs_dir): + rows = [] + for path in sorted(Path(runs_dir).glob("*/metrics.json")): + run = json.loads(path.read_text(encoding="utf-8")) + memory = run.get("memory") or {} + rows.append( + { + "arm": run["arm"], + "target": run["target"], + "train_subset": run["train_subset"], + "train_size": run["train_size"], + "seed": run["seed"], + "training_seconds": run["training_seconds"], + **{name: run["test"][name] for name in METRICS}, + "inference_seq_per_second": run["inference"]["sequences_per_second"], + "inference_ms_per_sequence": run["inference"]["ms_per_sequence"], + "device": run["inference"]["device"], + "peak_rss_mb": memory.get("peak_rss_mb"), + "peak_gpu_mb": memory.get("peak_gpu_mb"), + "epochs_ran": (run.get("arm_details") or {}).get("epochs_ran"), + } + ) + if not rows: + raise SystemExit(f"No metrics.json found under {runs_dir}.") + frame = pd.DataFrame(rows) + frame["subset_order"] = frame.train_subset.map(SUBSET_ORDER).fillna(99) + return frame.sort_values(["target", "subset_order", "arm", "seed"]) + + +def summarize(runs): + aggregated = ( + runs.groupby(["target", "train_subset", "subset_order", "arm"]) + .agg( + seeds=("seed", "count"), + train_size=("train_size", "first"), + **{f"{name}_{stat}": (name, stat) for name in METRICS for stat in ("mean", "std")}, + training_seconds_mean=("training_seconds", "mean"), + inference_seq_per_second_mean=("inference_seq_per_second", "mean"), + peak_rss_mb_mean=("peak_rss_mb", "mean"), + peak_gpu_mb_mean=("peak_gpu_mb", "mean"), + ) + .reset_index() + ) + return aggregated.sort_values(["target", "subset_order", "rmse_mean"]) + + +def to_markdown(summary): + blocks = [] + for target in summary.target.unique(): + blocks.append(f"### {target}\n") + for subset in summary[summary.target == target].train_subset.unique(): + rows = summary[(summary.target == target) & (summary.train_subset == subset)] + size = int(rows.train_size.iloc[0]) + blocks.append(f"**train = {subset} ({size:,} sequences)**\n") + blocks.append( + "| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB |\n" + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |" + ) + for _, row in rows.iterrows(): + std = row["rmse_std"] + rmse = f"{row['rmse_mean']:.4f}" + (f" ± {std:.4f}" if pd.notna(std) else "") + cells = [ + row["arm"], + rmse, + f"{row['mae_mean']:.4f}", + f"{row['r2_mean']:.4f}", + f"{row['pearson_r_mean']:.4f}", + f"{row['spearman_rho_mean']:.4f}", + f"{row['bias_mean']:+.4f}", + f"{row['training_seconds_mean']:.0f}", + f"{row['inference_seq_per_second_mean']:.0f}", + f"{row['peak_rss_mb_mean']:.0f}" if pd.notna(row["peak_rss_mb_mean"]) else "-", + ] + blocks.append("| " + " | ".join(cells) + " |") + blocks.append("") + return "\n".join(blocks) + "\n" + + +def pretraining_delta(runs): + """Does IRRM pretraining help? RMSE difference with a significance test. + + Seed spread for these two arms is comparable to the effect being measured, so the + difference of means alone cannot answer the question. Welch's t-test does not assume + the two arms have equal variance, which they do not. + """ + rows = [] + for (target, subset, order), group in runs.groupby(["target", "train_subset", "subset_order"]): + scratch = group.loc[group.arm == "irrm_scratch", "rmse"].to_numpy() + pretrained = group.loc[group.arm == "irrm_pretrained", "rmse"].to_numpy() + if len(scratch) < 2 or len(pretrained) < 2: + continue + # Lower RMSE is better, so a positive delta means pretraining won. + delta = scratch.mean() - pretrained.mean() + _, p_value = stats.ttest_ind(scratch, pretrained, equal_var=False) + rows.append( + { + "target": target, + "train_subset": subset, + "subset_order": order, + "n_scratch": len(scratch), + "n_pretrained": len(pretrained), + "irrm_scratch": scratch.mean(), + "irrm_pretrained": pretrained.mean(), + "rmse_delta": delta, + "relative_improvement": delta / scratch.mean(), + "p_value": p_value, + "verdict": ( + "no detectable effect" + if p_value >= 0.05 + else ("pretraining helps" if delta > 0 else "pretraining hurts") + ), + } + ) + if not rows: + return None + return pd.DataFrame(rows).sort_values(["target", "subset_order"]) + + +def main(): + args = parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + runs = load_runs(args.runs_dir) + runs.drop(columns=["subset_order"]).to_csv(output_dir / "pgen_runs.csv", index=False) + + summary = summarize(runs) + summary.drop(columns=["subset_order"]).to_csv(output_dir / "pgen_summary.csv", index=False) + (output_dir / "pgen_summary.md").write_text(to_markdown(summary), encoding="utf-8") + + print(f"runs: {len(runs)} arms: {runs.arm.nunique()} targets: {runs.target.nunique()}") + delta = pretraining_delta(runs) + if delta is not None: + delta.drop(columns=["subset_order"]).to_csv(output_dir / "pretraining_delta.csv", index=False) + print("\nDoes IRRM pretraining help? (positive rmse_delta means pretraining wins)") + print( + delta[ + ["target", "train_subset", "n_scratch", "n_pretrained", "irrm_scratch", + "irrm_pretrained", "rmse_delta", "p_value", "verdict"] + ].to_string(index=False) + ) + print(f"\nwrote {output_dir}/pgen_runs.csv, pgen_summary.csv, pgen_summary.md") + + +if __name__ == "__main__": + main() diff --git a/benchmark/collect_reconstruction.py b/benchmark/collect_reconstruction.py new file mode 100644 index 0000000..0c41543 --- /dev/null +++ b/benchmark/collect_reconstruction.py @@ -0,0 +1,148 @@ +"""Issue 2: aggregate reconstruction runs into the per-run CSV and the summary table. + +Reads every ``metrics.json`` written by ``train_decoder`` and produces: + reconstruction_runs.csv one row per run (representation x seed) + reconstruction_summary.csv mean and std across seeds, per representation + reconstruction_summary.md the same table in Markdown, for REPORT.md + +Explained variance from the PCA fit is joined in, because it is the main explanation +for differences at a fixed bottleneck width: a representation that survives the +projection intact starts from a different place than one that loses most of its variance. +""" + +import argparse +import json +from pathlib import Path + +import pandas as pd + +METRICS = ( + "exact_match", + "token_accuracy", + "normalized_levenshtein", + "within_edit_distance_1", + "mean_edit_distance", +) + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--runs-dir", default="artifacts/benchmark/reconstruction") + p.add_argument("--bottleneck-dir", default="data/benchmark/trb/bottleneck") + p.add_argument("--output-dir", default="results/reconstruction") + p.add_argument("--bottleneck-dim", type=int, default=64) + return p.parse_args() + + +def load_runs(runs_dir, bottleneck_dim): + rows = [] + for path in sorted(Path(runs_dir).glob("*/metrics.json")): + run = json.loads(path.read_text(encoding="utf-8")) + if run["bottleneck_dim"] != bottleneck_dim: + continue + rows.append( + { + "representation": run["representation"], + "seed": run["seed"], + "bottleneck_dim": run["bottleneck_dim"], + "train_subset": run["train_subset"], + "train_size": run["train_size"], + "epochs_ran": run["epochs_ran"], + "best_epoch": run["best_epoch"], + "best_val_loss": run["best_val_loss"], + "training_seconds": run["training_seconds"], + "test_loss": run["test"]["loss"], + **{name: run["test"][name] for name in METRICS}, + "inference_seq_per_second": run["inference"]["sequences_per_second"], + "inference_ms_per_sequence": run["inference"]["ms_per_sequence"], + "device": run["inference"]["device"], + } + ) + if not rows: + raise SystemExit(f"No metrics.json found under {runs_dir} for dim={bottleneck_dim}.") + return pd.DataFrame(rows) + + +def summarize(runs, bottleneck_dir, bottleneck_dim): + aggregated = runs.groupby("representation").agg( + seeds=("seed", "count"), + **{f"{name}_{stat}": (name, stat) for name in METRICS for stat in ("mean", "std")}, + training_seconds_mean=("training_seconds", "mean"), + epochs_mean=("epochs_ran", "mean"), + inference_seq_per_second_mean=("inference_seq_per_second", "mean"), + ) + + summary_path = Path(bottleneck_dir) / "bottleneck_summary.json" + if summary_path.exists(): + bottleneck = json.loads(summary_path.read_text(encoding="utf-8")) + aggregated["source_dim"] = [ + bottleneck.get(name, {}).get("source_dim") for name in aggregated.index + ] + aggregated["explained_variance"] = [ + bottleneck.get(name, {}).get("explained_variance", {}).get(str(bottleneck_dim)) + for name in aggregated.index + ] + + return aggregated.sort_values("exact_match_mean", ascending=False).reset_index() + + +def to_markdown(summary, bottleneck_dim): + header = ( + f"### Reconstruction at a {bottleneck_dim}-dimensional bottleneck\n\n" + "| representation | source dim | expl. var | exact match | token acc | " + "norm. Levenshtein | within ED 1 | train s | seq/s |\n" + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |\n" + ) + + def cell(row, name): + std = row[f"{name}_std"] + spread = f" ± {std:.4f}" if pd.notna(std) else "" + return f"{row[f'{name}_mean']:.4f}{spread}" + + lines = [] + for _, row in summary.iterrows(): + expl = row.get("explained_variance") + source_dim = row.get("source_dim") + cells = [ + row["representation"], + f"{source_dim:.0f}" if pd.notna(source_dim) else "-", + f"{expl:.3f}" if pd.notna(expl) else "-", + cell(row, "exact_match"), + cell(row, "token_accuracy"), + cell(row, "normalized_levenshtein"), + cell(row, "within_edit_distance_1"), + f"{row['training_seconds_mean']:.0f}", + f"{row['inference_seq_per_second_mean']:.0f}", + ] + lines.append("| " + " | ".join(cells) + " |") + return header + "\n".join(lines) + "\n" + + +def main(): + args = parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + runs = load_runs(args.runs_dir, args.bottleneck_dim) + runs = runs.sort_values(["representation", "seed"]) + runs.to_csv(output_dir / "reconstruction_runs.csv", index=False) + + summary = summarize(runs, args.bottleneck_dir, args.bottleneck_dim) + summary.to_csv(output_dir / "reconstruction_summary.csv", index=False) + (output_dir / "reconstruction_summary.md").write_text( + to_markdown(summary, args.bottleneck_dim), encoding="utf-8" + ) + + print(f"runs: {len(runs)} representations: {runs.representation.nunique()}") + print( + summary[ + ["representation", "seeds", "exact_match_mean", "exact_match_std", "token_accuracy_mean"] + ].to_string(index=False) + ) + print(f"\nwrote {output_dir}/reconstruction_runs.csv") + print(f"wrote {output_dir}/reconstruction_summary.csv") + print(f"wrote {output_dir}/reconstruction_summary.md") + + +if __name__ == "__main__": + main() diff --git a/benchmark/combine_external_pca128_results.py b/benchmark/combine_external_pca128_results.py new file mode 100644 index 0000000..8b38a28 --- /dev/null +++ b/benchmark/combine_external_pca128_results.py @@ -0,0 +1,49 @@ +"""Combine the all-model PCA-128 sensitivity tables without recomputing metrics.""" + +import argparse +import hashlib +import json +from pathlib import Path + +import pandas as pd + + +def sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--main-root", required=True) + parser.add_argument("--sceptr-root", required=True) + return parser.parse_args() + + +def main(): + args = parse_args() + main_root, sceptr_root = Path(args.main_root), Path(args.sceptr_root) + output = main_root / "results" / "all_baselines" + output.mkdir(parents=True, exist_ok=True) + sources = {} + for task, filename in (("reconstruction", "reconstruction_summary.csv"), ("pgen", "pgen_summary.csv")): + paths = [ + main_root / "results" / task / filename, + sceptr_root / "results" / task / filename, + ] + frame = pd.concat([pd.read_csv(path) for path in paths], ignore_index=True) + identity_column = "representation" if task == "reconstruction" else "arm" + if len(frame) != 3 or frame[identity_column].nunique() != 3: + raise ValueError(f"Expected three distinct PCA-128 baseline rows for {task}") + destination = output / filename + frame.to_csv(destination, index=False) + sources[task] = { + "inputs": [{"path": str(path), "sha256": sha256(path)} for path in paths], + "output": str(destination), + "output_sha256": sha256(destination), + } + (output / "MANIFEST.json").write_text(json.dumps(sources, indent=2), encoding="utf-8") + print("status=accepted baselines=3 tasks=2") + + +if __name__ == "__main__": + main() diff --git a/benchmark/create_vdjdb_tcremp_umap.py b/benchmark/create_vdjdb_tcremp_umap.py new file mode 100644 index 0000000..f947633 --- /dev/null +++ b/benchmark/create_vdjdb_tcremp_umap.py @@ -0,0 +1,115 @@ +"""Fit and render a native-TCRemP UMAP for the full VDJdb cohort.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D +import numpy as np +import pandas as pd +import umap + + +HIGHLIGHTS = { + "SSYRRPVGI": "#4477AA", + "STPESANL": "#EE6677", + "SLLMWITQV": "#228833", + "KYNKANVFL": "#AA3377", + "SIINFEKL": "#CCBB44", +} +SEED = 1729 +MAX_BACKGROUND = 5_000 +MAX_PER_HIGHLIGHT = 1_800 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--tcremp", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + return parser.parse_args() + + +def cap(frame: pd.DataFrame, maximum: int, rng: np.random.Generator) -> pd.DataFrame: + if len(frame) <= maximum: + return frame + return frame.iloc[rng.choice(len(frame), size=maximum, replace=False)] + + +def main() -> None: + args = parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + + cohort = pd.read_csv(args.cohort, sep="\t") + labels = cohort.label.astype(str).to_numpy() + matrix = np.load(args.tcremp, mmap_mode="r") + if matrix.shape[0] != len(cohort): + raise ValueError(f"TCRemP rows ({matrix.shape[0]}) do not match cohort ({len(cohort)})") + if not np.isfinite(matrix).all(): + raise ValueError("TCRemP matrix contains non-finite values") + + coordinates = umap.UMAP( + n_neighbors=30, + min_dist=0.15, + metric="cosine", + n_components=2, + random_state=SEED, + transform_seed=SEED, + ).fit_transform(np.asarray(matrix, dtype=np.float32)) + frame = pd.DataFrame({"index": np.arange(len(cohort)), "epitope": labels, + "umap_1": coordinates[:, 0], "umap_2": coordinates[:, 1]}) + frame.to_parquet(args.output_dir / "umap_coordinates.parquet", index=False) + + plt.rcParams.update({"font.family": "STIXGeneral", "font.size": 7.5, + "pdf.fonttype": 42, "ps.fonttype": 42, "svg.fonttype": "none"}) + rng = np.random.default_rng(SEED) + fig, ax = plt.subplots(figsize=(3.38, 3.38)) + other = cap(frame[~frame.epitope.isin(HIGHLIGHTS)], MAX_BACKGROUND, rng) + ax.scatter(other.umap_1, other.umap_2, s=1.1, c="#9e9e9e", alpha=.20, + linewidths=0, rasterized=True) + for epitope, color in HIGHLIGHTS.items(): + selected = cap(frame[frame.epitope == epitope], MAX_PER_HIGHLIGHT, rng) + ax.scatter(selected.umap_1, selected.umap_2, s=8.0, c=color, alpha=.76, + linewidths=0, rasterized=True) + ax.set_xticks([]) + ax.set_yticks([]) + for spine in ax.spines.values(): + spine.set_visible(False) + ax.set_title("TCRemP prototype embedding", fontsize=9, fontweight="bold", pad=5) + handles = [Line2D([], [], marker="o", linestyle="", markerfacecolor="#9e9e9e", + markeredgewidth=0, markersize=3.9, alpha=.7, label="Other epitopes")] + handles += [Line2D([], [], marker="o", linestyle="", markerfacecolor=color, + markeredgewidth=0, markersize=4.5, label=epitope) + for epitope, color in HIGHLIGHTS.items()] + fig.legend(handles=handles, loc="lower center", ncol=3, frameon=False, fontsize=6.8, + bbox_to_anchor=(.5, -.045), handletextpad=.25, columnspacing=.7, borderaxespad=0) + fig.subplots_adjust(left=.012, right=.995, top=.89, bottom=.105) + for suffix, kwargs in ((".png", {"dpi": 600}), (".pdf", {}), (".svg", {})): + fig.savefig((args.output_dir / "vdjdb_whole_cohort_tcremp_umap").with_suffix(suffix), + bbox_inches="tight", pad_inches=.01, **kwargs) + plt.close(fig) + + metadata = { + "cohort_size": int(len(cohort)), + "source_embeddings": {"TCRemP": str(args.tcremp)}, + "matrix_shape": [int(dim) for dim in matrix.shape], + "highlights": {epitope: {"n_full_cohort": int((labels == epitope).sum()), "color": color} + for epitope, color in HIGHLIGHTS.items()}, + "umap": {"metric": "cosine", "n_neighbors": 30, "min_dist": 0.15, + "n_components": 2, "random_state": SEED, "transform_seed": SEED}, + "display_sampling": {"seed": SEED, "max_background": MAX_BACKGROUND, + "max_per_highlight": MAX_PER_HIGHLIGHT}, + "caveat": "The native 9,000-dimensional TCRemP representation is projected separately; compare only qualitative within-panel structure.", + } + (args.output_dir / "README.md").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/benchmark/create_vdjdb_umap_candidates.py b/benchmark/create_vdjdb_umap_candidates.py new file mode 100644 index 0000000..0fb5064 --- /dev/null +++ b/benchmark/create_vdjdb_umap_candidates.py @@ -0,0 +1,89 @@ +"""Create matched, exploratory RP/RTP UMAPs from saved VDJdb embeddings only.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import umap + + +TARGETS = ["SLLMWITQV", "GLCTLVAML", "VEALYLVCG", "LLLDRLNQL"] +SEED = 1729 +N_TARGET = 1500 +N_BACKGROUND = 1500 + + +def args() -> argparse.Namespace: + p = argparse.ArgumentParser() + p.add_argument("--cohort", type=Path, required=True) + p.add_argument("--native-evaluation", type=Path, required=True) + p.add_argument("--output-dir", type=Path, required=True) + return p.parse_args() + + +def sample_indices(labels: np.ndarray, target: str, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]: + positive = np.flatnonzero(labels == target) + background = np.flatnonzero(labels != target) + return (rng.choice(positive, size=min(N_TARGET, len(positive)), replace=False), + rng.choice(background, size=min(N_BACKGROUND, len(background)), replace=False)) + + +def project(values: np.ndarray) -> np.ndarray: + # Projection is intentionally fit separately in each representation space. + # This preserves native neighborhoods but does not make coordinate systems comparable. + return umap.UMAP(n_neighbors=30, min_dist=0.15, metric="cosine", random_state=SEED, + n_components=2, transform_seed=SEED).fit_transform(values) + + +def main() -> None: + a = args() + if a.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {a.output_dir}") + a.output_dir.mkdir(parents=True) + cohort = pd.read_csv(a.cohort, sep="\t") + labels = cohort.label.astype(str).to_numpy() + rng = np.random.default_rng(SEED) + models = {"RP": np.load(a.native_evaluation / "embeddings_rp.npy", mmap_mode="r"), + "RTP": np.load(a.native_evaluation / "embeddings_rtp.npy", mmap_mode="r")} + if any(len(matrix) != len(cohort) for matrix in models.values()): + raise ValueError("Embedding rows do not match cohort rows") + fig, axes = plt.subplots(len(TARGETS), 2, figsize=(7.0, 10.0), constrained_layout=True) + all_rows = [] + for row, target in enumerate(TARGETS): + positive, background = sample_indices(labels, target, rng) + indices = np.concatenate([background, positive]) + is_target = np.concatenate([np.zeros(len(background), dtype=bool), np.ones(len(positive), dtype=bool)]) + for column, (name, matrix) in enumerate(models.items()): + coords = project(np.asarray(matrix[indices], dtype=np.float32)) + frame = pd.DataFrame({"index": indices, "epitope": labels[indices], "selected_epitope": target, + "is_selected_epitope": is_target, "model": name, + "umap_1": coords[:, 0], "umap_2": coords[:, 1]}) + all_rows.append(frame) + ax = axes[row, column] + ax.scatter(coords[~is_target, 0], coords[~is_target, 1], s=3.2, c="#bdbdbd", alpha=.33, linewidths=0) + ax.scatter(coords[is_target, 0], coords[is_target, 1], s=6.5, c="#0072B2", alpha=.76, linewidths=0) + ax.set_title(f"{target} (n={len(positive):,}) — {name}", fontsize=9) + ax.set_xticks([]); ax.set_yticks([]) + ax.text(.02, .03, "separate UMAP fit", transform=ax.transAxes, fontsize=6.5, color="#555555") + for spine in ax.spines.values(): spine.set_linewidth(.5) + fig.text(.5, .002, "Same sequence subset in each RP/RTP pair; UMAP: cosine, n_neighbors=30, min_dist=0.15, seed=1729. Coordinates are not directly comparable across separate fits.", ha="center", fontsize=7) + fig.savefig(a.output_dir / "vdjdb_umap_contact_sheet.png", dpi=450, bbox_inches="tight") + fig.savefig(a.output_dir / "vdjdb_umap_contact_sheet.pdf", bbox_inches="tight") + plt.close(fig) + pd.concat(all_rows, ignore_index=True).to_csv(a.output_dir / "umap_coordinates.tsv", sep="\t", index=False) + (a.output_dir / "README.md").write_text(json.dumps({ + "source_embeddings": {name: str(a.native_evaluation / f"embeddings_{name.lower()}.npy") for name in models}, + "cohort": str(a.cohort), "targets": TARGETS, "sampling": {"target_max": N_TARGET, "background": N_BACKGROUND, "seed": SEED}, + "umap": {"metric": "cosine", "n_neighbors": 30, "min_dist": .15, "random_state": SEED}, + "caveat": "RP and RTP use matched sequence subsets but separate UMAP fits; compare qualitative within-panel structure only." + }, indent=2) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/benchmark/create_vdjdb_whole_umap.py b/benchmark/create_vdjdb_whole_umap.py new file mode 100644 index 0000000..22900c9 --- /dev/null +++ b/benchmark/create_vdjdb_whole_umap.py @@ -0,0 +1,82 @@ +"""Fit one UMAP per saved RP/RTP embedding set across the complete VDJdb cohort.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D +import numpy as np +import pandas as pd +import umap + + +HIGHLIGHTS = { + "SLLMWITQV": ("SLLMWITQV", "#0072B2"), + "GLCTLVAML": ("GLCTLVAML", "#D55E00"), + "VEALYLVCG": ("VEALYLVCG", "#009E73"), + "LLLDRLNQL": ("LLLDRLNQL", "#CC79A7"), +} +SEED = 1729 + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser() + p.add_argument("--cohort", type=Path, required=True) + p.add_argument("--native-evaluation", type=Path, required=True) + p.add_argument("--output-dir", type=Path, required=True) + return p.parse_args() + + +def fit(matrix: np.ndarray) -> np.ndarray: + return umap.UMAP(n_neighbors=30, min_dist=0.15, metric="cosine", n_components=2, + random_state=SEED, transform_seed=SEED).fit_transform(matrix) + + +def main() -> None: + a = parse_args() + if a.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {a.output_dir}") + a.output_dir.mkdir(parents=True) + cohort = pd.read_csv(a.cohort, sep="\t") + labels = cohort.label.astype(str).to_numpy() + matrices = {"RP": np.load(a.native_evaluation / "embeddings_rp.npy", mmap_mode="r"), + "RTP": np.load(a.native_evaluation / "embeddings_rtp.npy", mmap_mode="r")} + if any(len(x) != len(cohort) for x in matrices.values()): + raise ValueError("Embedding rows do not match the full cohort") + fig, axes = plt.subplots(1, 2, figsize=(12, 5.4), constrained_layout=True) + frames = [] + highlighted = np.isin(labels, list(HIGHLIGHTS)) + for ax, (name, matrix) in zip(axes, matrices.items()): + coords = fit(np.asarray(matrix, dtype=np.float32)) + ax.scatter(coords[~highlighted, 0], coords[~highlighted, 1], s=.45, c="#a9a9a9", alpha=.16, linewidths=0, rasterized=True) + for epitope, (display, color) in HIGHLIGHTS.items(): + mask = labels == epitope + ax.scatter(coords[mask, 0], coords[mask, 1], s=3.8, c=color, alpha=.85, linewidths=0, rasterized=True, label=f"{display} (n={mask.sum():,})") + ax.set_title(name, fontsize=12, weight="bold") + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + ax.tick_params(length=2, labelsize=7) + frames.append(pd.DataFrame({"index": np.arange(len(cohort)), "epitope": labels, "model": name, + "umap_1": coords[:, 0], "umap_2": coords[:, 1], + "highlighted": highlighted})) + handles = [Line2D([], [], marker="o", linestyle="", markerfacecolor="#a9a9a9", markeredgewidth=0, markersize=4, alpha=.55, label="Other VDJdb epitopes")] + handles += [Line2D([], [], marker="o", linestyle="", markerfacecolor=color, markeredgewidth=0, markersize=5, label=f"{display} (n={(labels == ep).sum():,})") for ep, (display, color) in HIGHLIGHTS.items()] + fig.legend(handles=handles, loc="lower center", ncol=3, frameon=False, fontsize=8, bbox_to_anchor=(.5, -.085)) + fig.text(.5, .005, "Full 65,756-sequence VDJdb cohort. UMAP: cosine, n_neighbors=30, min_dist=0.15, seed=1729. RP and RTP are separate fits; coordinates are not directly comparable.", ha="center", fontsize=7.5) + fig.savefig(a.output_dir / "vdjdb_whole_cohort_rp_rtp_umap.png", dpi=450, bbox_inches="tight") + fig.savefig(a.output_dir / "vdjdb_whole_cohort_rp_rtp_umap.pdf", bbox_inches="tight") + plt.close(fig) + pd.concat(frames, ignore_index=True).to_parquet(a.output_dir / "umap_coordinates.parquet", index=False) + (a.output_dir / "README.md").write_text(json.dumps({ + "cohort_size": int(len(cohort)), "source_embeddings": {name: str(a.native_evaluation / f"embeddings_{name.lower()}.npy") for name in matrices}, + "highlights": {ep: {"n": int((labels == ep).sum()), "color": color} for ep, (_, color) in HIGHLIGHTS.items()}, + "umap": {"metric": "cosine", "n_neighbors": 30, "min_dist": .15, "random_state": SEED}, + "caveat": "Both panels use the full identical cohort; RP/RTP are separately fitted UMAPs, so compare qualitative within-panel structure only." + }, indent=2) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/benchmark/diagnose_vdjdb_decoded_tcremp_geometry.py b/benchmark/diagnose_vdjdb_decoded_tcremp_geometry.py new file mode 100644 index 0000000..a65b977 --- /dev/null +++ b/benchmark/diagnose_vdjdb_decoded_tcremp_geometry.py @@ -0,0 +1,31 @@ +"""Diagnose raw and target-standardized TCRemP geometry on locked VDJdb pairs.""" +import json +from pathlib import Path +import numpy as np +import pandas as pd +from evaluate_vdjdb_uniform100_separation import separation_metrics +from irrm_codec.multitask_data import TargetStandardizer + +ROOT=Path('/home/evlasova/irrm-codec') +out=ROOT/'experiments/trb/vdjdb_decoded_tcremp/1429942/geometry_diagnostic' +out.mkdir(exist_ok=False) +cohort=pd.read_csv(ROOT/'experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv',sep='\t') +pairs=np.load(ROOT/'experiments/trb/vdjdb-uniform100-v1/evaluation-1427813/uniform_pairs.npz') +original,within,between=pairs['original_indices'],pairs['within'],pairs['between'] +labels=cohort.iloc[original].label.astype(str).to_numpy() +native=np.load(ROOT/'experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/tcremp.npy',mmap_mode='r') +decoded=np.load(ROOT/'experiments/trb/vdjdb_decoded_tcremp/1429942/decoded_tcremp.npy',mmap_mode='r') +s=TargetStandardizer.load(ROOT/'experiments/trb/anchor-data-rtp-seed42/target_standardizer.npz') +variants={'native_raw':native[original], 'decoded_raw':decoded[original], + 'native_standardized':(native[original]-s.tcremp_mean)/s.tcremp_std, + 'decoded_standardized':(decoded[original]-s.tcremp_mean)/s.tcremp_std, + 'native_centered':native[original]-s.tcremp_mean, + 'decoded_centered':decoded[original]-s.tcremp_mean} +rows=[] +for name,matrix in variants.items(): + query,epitope=separation_metrics(matrix,within,between,labels) + query.insert(0,'representation',name) + query.to_parquet(out/f'per_query_{name}.parquet',index=False) + rows.append({'representation':name,'macro_cohen_d':float(epitope.cohen_d.mean()),'macro_auroc':float(epitope.auroc.mean()),'median_cohen_d':float(epitope.cohen_d.median())}) +pd.DataFrame(rows).to_csv(out/'summary.tsv',sep='\t',index=False) +(out/'RESULTS.json').write_text(json.dumps({'status':'complete','note':'All variants use unchanged separation_metrics and identical fixed uniform-100 pairs.'},indent=2)+'\n') \ No newline at end of file diff --git a/benchmark/evaluate_latent_secondary.py b/benchmark/evaluate_latent_secondary.py new file mode 100644 index 0000000..7f0f303 --- /dev/null +++ b/benchmark/evaluate_latent_secondary.py @@ -0,0 +1,659 @@ +"""Preregistered secondary tests of matched frozen DATA-ANCHOR representations. + +This benchmark intentionally asks complementary *representation sufficiency* +questions on the locked TRB train/validation/test split. It does not search for +a task on which RTP wins. Probe hyperparameters are selected on validation; +test labels are used once for reporting and never for representation/model +selection. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +from dataclasses import asdict +from pathlib import Path + +import numpy as np +import pandas as pd +from sklearn.decomposition import PCA + +MODELS = ("r", "t", "p", "rt", "rp", "tp", "rtp") +AA = "ACDEFGHIKLMNPQRSTVWY" +CONSERVATIVE_GROUPS = ("AILMV", "FWY", "STNQ", "KRH", "DE", "CGP") +FRACTIONS = (0.01, 0.05, 0.10, 0.25, 0.50, 1.00) +ALPHAS = (0.1, 10.0, 1000.0) +JOINT_QUANTILES = (0.20, 0.25, 0.30) + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(1024 * 1024): + digest.update(block) + return digest.hexdigest() + + +def stable_json(path: Path, value: object) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def assert_finite_tree(value: object, where: str = "root") -> None: + if isinstance(value, dict): + for key, child in value.items(): + assert_finite_tree(child, f"{where}.{key}") + elif isinstance(value, list): + for index, child in enumerate(value): + assert_finite_tree(child, f"{where}[{index}]") + elif isinstance(value, (float, np.floating)) and not math.isfinite(float(value)): + raise FloatingPointError(f"Non-finite value at {where}: {value}") + + +def preregistration() -> dict: + return { + "status": "preregistered_before_test_evaluation", + "models": list(MODELS), + "primary_comparison": "RTP versus every other matched encoder; all conditions reported", + "no_selection_rule": "No track, subgroup, representation, or test metric may be omitted based on RTP rank.", + "tracks": { + "low_label": { + "nested_train_fractions": list(FRACTIONS), + "subset_rule": "prefixes of one seed-42 permutation of the locked train manifest", + "probe": "ridge with intercept; representation standardized on each train subset", + "alpha_grid_validation_only": list(ALPHAS), + "targets": { + "sequence_content": "20 AA fractions plus standardized CDR3 length", + "tcremp": "fixed seed-42 random-sign 128D sketch of locked train-standardized 9000D TCRemP vector", + "pgen": "locked train-standardized log10_pgen_1mm", + }, + "metrics": ["sequence RMSE/cosine and length MAE", "TCRemP MSE/cosine", "pgen RMSE/MAE/R2"], + }, + "strata": { + "fit": "full locked train; alpha selected on validation globally, never within test subgroup", + "length_bins": ["<=12", "13-15", "16-18", ">=19"], + "gene_families": "TRBV/TRBJ prefix before '-' and '*'; report test groups n>=100", + "pgen_bins": "quartiles defined from locked train targets", + }, + "joint_retrieval": { + "queries": "all locked test rows", + "gallery": "256 deterministic seed-42 test candidates per query, excluding self", + "relevance": ( + "simultaneously below the query-specific q-quantile for sequence 2-mer/length distance, " + "TCRemP cosine distance, and absolute standardized pgen difference" + ), + "primary_q": 0.25, + "sensitivity_q": list(JOINT_QUANTILES), + "metrics": ["P@1", "P@5", "P@10", "MAP", "eligible query count"], + "caveat": "Target-defined relevance measures joint target sufficiency, not independent biology.", + }, + "robustness": { + "rows": "deterministic 5000-row test subset", + "perturbations": [ + "one within-biochemical-group AA substitution", + "one outside-group AA substitution", + "one central encoder token replaced by UNK id 3", + ], + "no_indels": True, + "metrics": ["cosine drift", "Euclidean drift", "nonconservative greater than conservative fraction"], + }, + "synthetic_discrimination": { + "status": "disabled unless an explicit provenance manifest is supplied", + "required_manifest": ["source_path", "source_sha256", "generator", "generator_sha256", "benchmark_overlap=0"], + "reason": "No nearby file is accepted as synthetic merely from its filename.", + }, + }, + "representations": [ + *MODELS, + "concat_r_t_p_384d", + "pca128_of_concat_fit_train_only", + "aa_composition_length_vj_train_schema", + ], + "caveats": [ + "All neural models are one training seed; sampling intervals do not represent training-seed variation.", + "The low-cost sequence probe tests composition/length, not exact autoregressive reconstruction.", + "TCRemP probing/relevance uses one preregistered random projection for CPU feasibility; it is not the full 9000D target.", + "Concat/PCA use three encoders and are efficiency upper bounds, not single-encoder peers.", + ], + } + + +def nested_subsets(indices: np.ndarray, fractions: tuple[float, ...], seed: int) -> dict[float, np.ndarray]: + order = np.random.default_rng(seed).permutation(np.asarray(indices, dtype=np.int64)) + result = {} + previous: set[int] = set() + for fraction in fractions: + size = max(100, int(round(len(order) * fraction))) + subset = order[: min(size, len(order))] + if not previous.issubset(set(map(int, subset))): + raise AssertionError("Nested subset construction failed.") + previous = set(map(int, subset)) + result[fraction] = subset + return result + + +def sequence_features(sequences: pd.Series) -> np.ndarray: + output = np.zeros((len(sequences), len(AA) + 1), dtype=np.float32) + for row, sequence in enumerate(sequences.astype(str)): + for column, aa in enumerate(AA): + output[row, column] = sequence.count(aa) / len(sequence) + output[row, -1] = len(sequence) + return output + + +def bigram_features(sequences: pd.Series) -> np.ndarray: + lookup = {a + b: i for i, (a, b) in enumerate((a, b) for a in AA for b in AA)} + output = np.zeros((len(sequences), len(lookup)), dtype=np.float32) + for row, sequence in enumerate(sequences.astype(str)): + for left, right in zip(sequence, sequence[1:]): + output[row, lookup[left + right]] += 1 + norm = np.linalg.norm(output[row]) + if norm: + output[row] /= norm + return output + + +def gene_family(value: object) -> str: + text = str(value).upper().split("*", 1)[0] + return text.split("-", 1)[0] + + +def control_features(table: pd.DataFrame, train: np.ndarray) -> np.ndarray: + basic = sequence_features(table["junction_aa"]) + mean = basic[train].mean(axis=0) + std = basic[train].std(axis=0) + basic = (basic - mean) / np.where(std < 1e-8, 1.0, std) + categories = {} + for column in ("v_call", "j_call"): + values = table[column].map(gene_family).to_numpy() + categories[column] = {value: i for i, value in enumerate(sorted(set(values[train])))} + width = sum(len(value) for value in categories.values()) + genes = np.zeros((len(table), width), dtype=np.float32) + offset = 0 + for column in ("v_call", "j_call"): + values = table[column].map(gene_family).to_numpy() + for row, value in enumerate(values): + index = categories[column].get(value) + if index is not None: + genes[row, offset + index] = 1 + offset += len(categories[column]) + return np.concatenate((basic, genes), axis=1) + + +def standardize_targets( + table: pd.DataFrame, tcremp: np.ndarray, standardizer: object, +) -> tuple[np.ndarray, dict[str, slice]]: + sequence = sequence_features(table["junction_aa"]).astype(np.float64) + sequence[:, -1] = (sequence[:, -1] - table.loc[table["split"] == "train", "junction_aa"].str.len().mean()) / max( + float(table.loc[table["split"] == "train", "junction_aa"].str.len().std(ddof=0)), 1e-8 + ) + t = np.asarray(tcremp, dtype=np.float64) + p = ((table[standardizer.pgen_target].to_numpy(np.float64) - standardizer.pgen_mean) / standardizer.pgen_std)[:, None] + y = np.concatenate((sequence, t, p), axis=1) + slices = { + "sequence": slice(0, sequence.shape[1]), + "tcremp": slice(sequence.shape[1], sequence.shape[1] + t.shape[1]), + "pgen": slice(sequence.shape[1] + t.shape[1], y.shape[1]), + } + return y, slices + + +def project_tcremp( + tcremp: np.ndarray, standardizer: object, dimension: int = 128, seed: int = 42, + chunk_size: int = 512, +) -> tuple[np.ndarray, dict]: + """Apply one fixed Johnson-Lindenstrauss sign sketch after locked standardization.""" + source_dim = int(tcremp.shape[1]) + rng = np.random.default_rng(seed) + projection = rng.choice(np.asarray([-1.0, 1.0], dtype=np.float32), size=(source_dim, dimension)) + projection /= np.sqrt(dimension) + output = np.empty((len(tcremp), dimension), dtype=np.float32) + for start in range(0, len(tcremp), chunk_size): + stop = min(start + chunk_size, len(tcremp)) + standardized = ( + np.asarray(tcremp[start:stop], dtype=np.float32) - standardizer.tcremp_mean + ) / standardizer.tcremp_std + output[start:stop] = standardized @ projection + if not np.isfinite(output).all(): + raise ValueError("TCRemP random projection emitted non-finite values.") + return output, { + "method": "dense Rademacher random projection after locked train standardization", + "seed": seed, "source_dimension": source_dim, "output_dimension": dimension, + "scale": "1/sqrt(output_dimension)", + "projection_sha256": hashlib.sha256(projection.tobytes()).hexdigest(), + } + + +def ridge_coefficients(x: np.ndarray, y: np.ndarray, indices: np.ndarray, alpha: float) -> dict: + xx = np.asarray(x[indices], dtype=np.float64) + yy = np.asarray(y[indices], dtype=np.float64) + mean_x, scale_x = xx.mean(axis=0), xx.std(axis=0) + scale_x = np.where(scale_x < 1e-8, 1.0, scale_x) + mean_y = yy.mean(axis=0) + z = (xx - mean_x) / scale_x + gram = z.T @ z + beta = np.linalg.solve(gram + np.eye(gram.shape[0]) * alpha, z.T @ (yy - mean_y)) + return {"mean_x": mean_x, "scale_x": scale_x, "mean_y": mean_y, "beta": beta, "alpha": alpha} + + +def ridge_predict(fit: dict, x: np.ndarray, indices: np.ndarray) -> np.ndarray: + z = (np.asarray(x[indices], dtype=np.float64) - fit["mean_x"]) / fit["scale_x"] + return z @ fit["beta"] + fit["mean_y"] + + +def cosine_mean(a: np.ndarray, b: np.ndarray) -> float: + denominator = np.linalg.norm(a, axis=1) * np.linalg.norm(b, axis=1) + valid = denominator > 0 + return float(np.mean(np.sum(a[valid] * b[valid], axis=1) / denominator[valid])) + + +def target_metrics(truth: np.ndarray, predicted: np.ndarray, slices: dict[str, slice]) -> dict: + sequence_true, sequence_pred = truth[:, slices["sequence"]], predicted[:, slices["sequence"]] + t_true, t_pred = truth[:, slices["tcremp"]], predicted[:, slices["tcremp"]] + p_true, p_pred = truth[:, slices["pgen"]].ravel(), predicted[:, slices["pgen"]].ravel() + residual = p_true - p_pred + return { + "sequence": { + "aa_composition_rmse": float(np.sqrt(np.mean((sequence_true[:, :-1] - sequence_pred[:, :-1]) ** 2))), + "aa_composition_cosine": cosine_mean(sequence_true[:, :-1], sequence_pred[:, :-1]), + "standardized_length_mae": float(np.mean(np.abs(sequence_true[:, -1] - sequence_pred[:, -1]))), + }, + "tcremp": { + "standardized_mse": float(np.mean((t_true - t_pred) ** 2)), + "cosine": cosine_mean(t_true, t_pred), + }, + "pgen": { + "standardized_rmse": float(np.sqrt(np.mean(residual ** 2))), + "standardized_mae": float(np.mean(np.abs(residual))), + "r2": float(1 - np.sum(residual ** 2) / np.sum((p_true - p_true.mean()) ** 2)), + }, + } + + +def selection_losses(metrics: dict) -> dict[str, float]: + return { + "sequence": metrics["sequence"]["aa_composition_rmse"] + metrics["sequence"]["standardized_length_mae"], + "tcremp": metrics["tcremp"]["standardized_mse"], + "pgen": metrics["pgen"]["standardized_rmse"], + } + + +def choose_task_fits( + x: np.ndarray, y: np.ndarray, train: np.ndarray, val: np.ndarray, slices: dict[str, slice], +) -> tuple[dict[str, dict], dict]: + fits = {alpha: ridge_coefficients(x, y, train, alpha) for alpha in ALPHAS} + validations = {alpha: target_metrics(y[val], ridge_predict(fit, x, val), slices) for alpha, fit in fits.items()} + selected = {} + report = {} + for task in ("sequence", "tcremp", "pgen"): + alpha = min(ALPHAS, key=lambda value: (selection_losses(validations[value])[task], value)) + selected[task] = fits[alpha] + report[task] = {"alpha": alpha, "validation": validations[alpha][task]} + return selected, report + + +def merged_predictions(fits: dict[str, dict], x: np.ndarray, indices: np.ndarray, slices: dict[str, slice]) -> np.ndarray: + width = max(value.stop for value in slices.values()) + result = np.empty((len(indices), width), dtype=np.float64) + for task, section in slices.items(): + result[:, section] = ridge_predict(fits[task], x, indices)[:, section] + return result + + +def length_strata(table: pd.DataFrame, test: np.ndarray) -> dict[str, np.ndarray]: + lengths = table.iloc[test]["junction_aa"].str.len().to_numpy() + rules = { + "le12": lengths <= 12, + "13_15": (lengths >= 13) & (lengths <= 15), + "16_18": (lengths >= 16) & (lengths <= 18), + "ge19": lengths >= 19, + } + return {name: np.flatnonzero(mask) for name, mask in rules.items() if mask.sum() >= 100} + + +def categorical_strata(table: pd.DataFrame, test: np.ndarray, column: str) -> dict[str, np.ndarray]: + values = table.iloc[test][column].map(gene_family).to_numpy() + return {str(value): np.flatnonzero(values == value) for value in sorted(set(values)) if np.sum(values == value) >= 100} + + +def pgen_strata(table: pd.DataFrame, train: np.ndarray, test: np.ndarray, target: str) -> dict[str, np.ndarray]: + edges = np.quantile(table.iloc[train][target].to_numpy(float), [0.25, 0.5, 0.75]) + groups = np.digitize(table.iloc[test][target].to_numpy(float), edges, right=True) + return {f"q{value + 1}": np.flatnonzero(groups == value) for value in range(4)} + + +def make_candidate_gallery(n: int, width: int, seed: int) -> np.ndarray: + if n <= width: + raise ValueError("Test set must be larger than the fixed candidate gallery.") + gallery = np.empty((n, width), dtype=np.int32) + for query in range(n): + rng = np.random.default_rng(seed + query * 1_000_003) + selected = rng.choice(n - 1, width, replace=False) + gallery[query] = selected + (selected >= query) + return gallery + + +def cosine_candidate_distance(features: np.ndarray, gallery: np.ndarray, chunk: int = 1024) -> np.ndarray: + values = np.asarray(features, dtype=np.float32) + norms = np.linalg.norm(values, axis=1, keepdims=True) + unit = values / np.where(norms == 0, 1.0, norms) + result = np.empty(gallery.shape, dtype=np.float32) + for start in range(0, len(values), chunk): + stop = min(start + chunk, len(values)) + result[start:stop] = 1 - np.sum(unit[start:stop, None, :] * unit[gallery[start:stop]], axis=2) + return np.maximum(result, 0) + + +def joint_relevance( + sequences: pd.Series, tcremp: np.ndarray, pgen: np.ndarray, gallery: np.ndarray, quantile: float, +) -> np.ndarray: + bigrams = bigram_features(sequences) + seq = cosine_candidate_distance(bigrams, gallery) + lengths = sequences.str.len().to_numpy(float) + seq += 0.2 * np.abs(lengths[:, None] - lengths[gallery]) / 40.0 + t_distance = cosine_candidate_distance(tcremp, gallery) + p_distance = np.abs(pgen[:, None] - pgen[gallery]) + thresholds = [np.quantile(value, quantile, axis=1, keepdims=True) for value in (seq, t_distance, p_distance)] + return (seq <= thresholds[0]) & (t_distance <= thresholds[1]) & (p_distance <= thresholds[2]) + + +def retrieval_metrics(distance: np.ndarray, relevant: np.ndarray) -> dict: + rows = [] + for query in range(len(distance)): + positives = int(relevant[query].sum()) + if positives == 0: + continue + order = np.argsort(distance[query], kind="mergesort") + hits = relevant[query, order].astype(float) + precision = np.cumsum(hits) / np.arange(1, len(hits) + 1) + rows.append([ + hits[0], hits[:5].mean(), hits[:10].mean(), float((precision * hits).sum() / positives), positives + ]) + values = np.asarray(rows, dtype=float) + return { + "eligible_queries": int(len(values)), + "precision_at_1": float(values[:, 0].mean()), + "precision_at_5": float(values[:, 1].mean()), + "precision_at_10": float(values[:, 2].mean()), + "map": float(values[:, 3].mean()), + "mean_relevant_candidates": float(values[:, 4].mean()), + } + + +def deterministic_substitutions(sequence: str, seed: int) -> tuple[str, str]: + if len(sequence) < 3: + raise ValueError("Robustness substitution requires a sequence of length at least three.") + rng = np.random.default_rng(seed) + position = int(rng.integers(1, len(sequence) - 1)) + source = sequence[position] + group = next(group for group in CONSERVATIVE_GROUPS if source in group) + conservative_pool = [aa for aa in group if aa != source] + if not conservative_pool: + conservative_pool = [aa for aa in AA if aa != source] + nonconservative_pool = [aa for aa in AA if aa not in group] + conservative = sequence[:position] + conservative_pool[int(rng.integers(len(conservative_pool)))] + sequence[position + 1:] + nonconservative = sequence[:position] + nonconservative_pool[int(rng.integers(len(nonconservative_pool)))] + sequence[position + 1:] + return conservative, nonconservative + + +def embedding_drift(original: np.ndarray, altered: np.ndarray) -> dict: + euclidean = np.linalg.norm(original - altered, axis=1) + cosine = 1 - np.sum(original * altered, axis=1) / ( + np.linalg.norm(original, axis=1) * np.linalg.norm(altered, axis=1) + ) + return { + "euclidean_mean": float(euclidean.mean()), "euclidean_median": float(np.median(euclidean)), + "cosine_mean": float(cosine.mean()), "cosine_median": float(np.median(cosine)), + } + + +def extract_altered_embeddings( + token_sets: dict[str, list[list[int]]], models: dict, config, batch_size: int, +) -> dict[str, dict[str, np.ndarray]]: + import torch + from irrm_codec.multitask_transformer import IRRMCodecTransformer + from irrm_codec.tokenization import PAD_ID + + result = {kind: {} for kind in token_sets} + for name, (checkpoint_path, _) in models.items(): + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + if checkpoint["model_config"] != asdict(config): + raise ValueError(f"{name} altered-input extraction found config mismatch.") + model = IRRMCodecTransformer(config) + model.load_state_dict(checkpoint["model_state"], strict=True) + model.eval() + with torch.no_grad(): + for kind, rows in token_sets.items(): + blocks = [] + for start in range(0, len(rows), batch_size): + block = rows[start:start + batch_size] + width = max(map(len, block)) + tokens = torch.full((len(block), width), PAD_ID, dtype=torch.long) + for index, ids in enumerate(block): + tokens[index, :len(ids)] = torch.tensor(ids) + blocks.append(model.encode(tokens, tokens.ne(PAD_ID)).numpy().astype(np.float32)) + result[kind][name] = np.concatenate(blocks) + del checkpoint, model + return result + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--model", action="append", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--preflight-only", action="store_true") + parser.add_argument("--accepted-preflight", type=Path) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def validate_inputs(args: argparse.Namespace) -> tuple[pd.DataFrame, np.ndarray, dict, object, dict]: + from benchmark.evaluate_vdjdb_frozen import parse_models, validate_run_configs + from irrm_codec.multitask_data import ( + TargetStandardizer, + load_prepared_benchmark, + resolve_encoder_tokenizer, + select_split_indices, + ) + + table, tcremp = load_prepared_benchmark(args.data_dir) + indices = {split: select_split_indices(table, args.data_dir, split) for split in ("train", "val", "test")} + if any(np.intersect1d(indices[a], indices[b]).size for a, b in (("train", "val"), ("train", "test"), ("val", "test"))): + raise ValueError("Locked split manifests overlap.") + models = parse_models(args.model) + if set(models) != set(MODELS): + raise ValueError("Full seven-condition factorial required.") + config, model_report = validate_run_configs(models) + tokenizer = resolve_encoder_tokenizer("data_anchor", str(args.tokenizer)) + if tokenizer.vocab_size != config.input_vocab_size: + raise ValueError("Tokenizer/checkpoint vocabulary mismatch.") + standardizer = TargetStandardizer.load(args.data_dir / "target_standardizer.npz") + if standardizer.train_rows != len(indices["train"]) or standardizer.pgen_target != "log10_pgen_1mm": + raise ValueError("Locked train-only normalizer mismatch.") + return table, tcremp, indices, config, {"models": models, "report": model_report, "tokenizer": tokenizer, "standardizer": standardizer} + + +def main() -> None: + from benchmark.evaluate_vdjdb_frozen import extract_embeddings + + args = parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + table, tcremp, indices, config, runtime = validate_inputs(args) + registration = preregistration() + stable_json(args.output_dir / "PREREGISTRATION.json", registration) + input_hashes = { + "dataset": sha256(args.data_dir / "dataset.parquet"), + "tcremp": sha256(args.data_dir / "embeddings.npy"), + "normalizer": sha256(args.data_dir / "target_standardizer.npz"), + "tokenizer": sha256(args.tokenizer), + "splits": {split: sha256(args.data_dir / "manifests" / f"{split}.tsv") for split in indices}, + } + preflight = { + "status": "accepted_for_secondary_frozen_evaluation", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "rows": len(table), "splits": {key: len(value) for key, value in indices.items()}, + "input_hashes": input_hashes, "model_report": runtime["report"], + "preregistration_sha256": sha256(args.output_dir / "PREREGISTRATION.json"), + "synthetic_track": "blocked_no_validated_source_manifest", + } + stable_json(args.output_dir / "PREFLIGHT.json", preflight) + if args.preflight_only: + print(json.dumps(preflight, indent=2, sort_keys=True)) + return + if args.accepted_preflight is None: + raise ValueError("--accepted-preflight is required for evaluation.") + accepted = json.loads(args.accepted_preflight.read_text(encoding="utf-8")) + if accepted["status"] != "accepted_for_secondary_frozen_evaluation" or accepted["input_hashes"] != input_hashes: + raise ValueError("Accepted preflight is absent, stale, or refers to different inputs.") + + cohort = pd.DataFrame({"cdr3": table["junction_aa"].astype(str)}) + embeddings, extraction = extract_embeddings( + cohort, runtime["models"], config, args.tokenizer, args.batch_size, args.output_dir + ) + representations = dict(embeddings) + concat = np.concatenate([embeddings[name] for name in ("r", "t", "p")], axis=1) + representations["concat_r_t_p"] = concat + pca = PCA(n_components=128, svd_solver="randomized", random_state=args.seed) + pca.fit(concat[indices["train"]]) + representations["pca128_concat_r_t_p"] = pca.transform(concat).astype(np.float32) + representations["sequence_length_vj_control"] = control_features(table, indices["train"]) + np.save(args.output_dir / "pca_components.npy", pca.components_, allow_pickle=False) + + tcremp_sketch, tcremp_projection_report = project_tcremp( + tcremp, runtime["standardizer"], dimension=128, seed=args.seed + ) + np.save(args.output_dir / "tcremp_random_projection.npy", tcremp_sketch, allow_pickle=False) + y, slices = standardize_targets(table, tcremp_sketch, runtime["standardizer"]) + subsets = nested_subsets(indices["train"], FRACTIONS, args.seed) + efficiency, final_fits = {}, {} + rows = [] + for representation, features in representations.items(): + efficiency[representation] = {} + for fraction, subset in subsets.items(): + fits, validation = choose_task_fits(features, y, subset, indices["val"], slices) + prediction = merged_predictions(fits, features, indices["test"], slices) + metrics = target_metrics(y[indices["test"]], prediction, slices) + efficiency[representation][str(fraction)] = { + "train_rows": len(subset), "validation_selection": validation, "test": metrics, + } + for task in metrics: + rows.append({"representation": representation, "fraction": fraction, "train_rows": len(subset), + "task": task, **metrics[task]}) + if fraction == 1.0: + final_fits[representation] = fits + pd.DataFrame(rows).to_csv(args.output_dir / "low_label_curves.tsv", sep="\t", index=False) + stable_json(args.output_dir / "low_label.json", efficiency) + + strata_definitions = { + "length": length_strata(table, indices["test"]), + "v_family": categorical_strata(table, indices["test"], "v_call"), + "j_family": categorical_strata(table, indices["test"], "j_call"), + "pgen_train_quartile": pgen_strata(table, indices["train"], indices["test"], runtime["standardizer"].pgen_target), + } + strata_rows = [] + for representation, features in representations.items(): + prediction = merged_predictions(final_fits[representation], features, indices["test"], slices) + test_truth = y[indices["test"]] + for stratum_type, groups in strata_definitions.items(): + for stratum, positions in groups.items(): + metrics = target_metrics(test_truth[positions], prediction[positions], slices) + for task, values in metrics.items(): + strata_rows.append({"representation": representation, "stratum_type": stratum_type, + "stratum": stratum, "rows": len(positions), "task": task, **values}) + pd.DataFrame(strata_rows).to_csv(args.output_dir / "stratified_test_metrics.tsv", sep="\t", index=False) + + test = indices["test"] + gallery = make_candidate_gallery(len(test), 256, args.seed) + np.save(args.output_dir / "joint_gallery.npy", gallery, allow_pickle=False) + test_sequences = table.iloc[test]["junction_aa"].reset_index(drop=True) + test_tcremp = tcremp_sketch[test] + test_pgen = y[test, slices["pgen"]].ravel() + retrieval = {} + for quantile in JOINT_QUANTILES: + relevant = joint_relevance(test_sequences, test_tcremp, test_pgen, gallery, quantile) + retrieval[str(quantile)] = { + name: retrieval_metrics(cosine_candidate_distance(features[test], gallery), relevant) + for name, features in representations.items() + } + retrieval[str(quantile)]["relevance"] = { + "eligible_queries": int(np.sum(relevant.sum(axis=1) > 0)), + "mean_relevant": float(relevant.sum(axis=1).mean()), + } + stable_json(args.output_dir / "joint_retrieval.json", retrieval) + + robust_positions = np.random.default_rng(args.seed).choice(test, size=min(5000, len(test)), replace=False) + robust_sequences = table.iloc[robust_positions]["junction_aa"].astype(str).tolist() + conservative, nonconservative = zip(*[ + deterministic_substitutions(sequence, args.seed + row * 17) + for row, sequence in enumerate(robust_sequences) + ]) + tokenizer = runtime["tokenizer"] + token_sets = { + "original": [tokenizer.encode(value, config.max_sequence_len) for value in robust_sequences], + "conservative": [tokenizer.encode(value, config.max_sequence_len) for value in conservative], + "nonconservative": [tokenizer.encode(value, config.max_sequence_len) for value in nonconservative], + } + masked = [] + for row, tokens in enumerate(token_sets["original"]): + copy = list(tokens) + position = int(np.random.default_rng(args.seed + row * 31).integers(len(copy))) + copy[position] = 3 + masked.append(copy) + token_sets["token_mask"] = masked + altered = extract_altered_embeddings(token_sets, runtime["models"], config, args.batch_size) + for kind in altered: + altered[kind]["concat_r_t_p"] = np.concatenate( + [altered[kind][name] for name in ("r", "t", "p")], axis=1 + ) + altered[kind]["pca128_concat_r_t_p"] = pca.transform( + altered[kind]["concat_r_t_p"] + ).astype(np.float32) + robustness = {} + for name in (*MODELS, "concat_r_t_p", "pca128_concat_r_t_p"): + original = altered["original"][name] + conservative_drift = np.linalg.norm(original - altered["conservative"][name], axis=1) + nonconservative_drift = np.linalg.norm(original - altered["nonconservative"][name], axis=1) + robustness[name] = { + kind: embedding_drift(original, altered[kind][name]) + for kind in ("conservative", "nonconservative", "token_mask") + } + robustness[name]["nonconservative_euclidean_gt_conservative_fraction"] = float( + np.mean(nonconservative_drift > conservative_drift) + ) + stable_json(args.output_dir / "robustness.json", robustness) + + report = { + "status": "complete", "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "study": "secondary frozen-representation benchmark on locked TRB splits", + "input_hashes": input_hashes, "models": runtime["report"], "extraction": extraction, + "representations": {name: list(value.shape) for name, value in representations.items()}, + "pca": {"fit_split": "train", "explained_variance_sum": float(pca.explained_variance_ratio_.sum())}, + "tcremp_projection": tcremp_projection_report, + "tracks": { + "low_label": "low_label.json", "stratified": "stratified_test_metrics.tsv", + "joint_retrieval": "joint_retrieval.json", "robustness": "robustness.json", + "synthetic_discrimination": "blocked_no_validated_source_manifest", + }, + "preregistration": {"path": str(args.output_dir / "PREREGISTRATION.json"), + "sha256": sha256(args.output_dir / "PREREGISTRATION.json")}, + "caveats": registration["caveats"], "nonfinite_scan": "passed", + } + assert_finite_tree(efficiency, "efficiency") + assert_finite_tree(retrieval, "retrieval") + assert_finite_tree(robustness, "robustness") + report["artifact_sha256"] = { + path.name: sha256(path) for path in sorted(args.output_dir.iterdir()) if path.is_file() + } + stable_json(args.output_dir / "RESULTS.json", report) + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_latent_sufficiency.py b/benchmark/evaluate_latent_sufficiency.py new file mode 100644 index 0000000..f8ba4e6 --- /dev/null +++ b/benchmark/evaluate_latent_sufficiency.py @@ -0,0 +1,517 @@ +"""Frozen linear-probe benchmark for balanced DATA-ANCHOR latent sufficiency. + +The seven encoders are never updated. Fresh ridge readouts are fitted on the +locked training manifest, their regularization is selected on validation, and +the locked test split is evaluated once. A fixed-position sequence probe is +used instead of any checkpoint decoder so reconstruction-supervised encoders +receive no privileged task head. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import time +from dataclasses import asdict +from pathlib import Path + +import numpy as np +import pandas as pd +from sklearn.decomposition import PCA + +AA = "ACDEFGHIKLMNPQRSTVWY" +AA_TO_CLASS = {aa: index + 1 for index, aa in enumerate(AA)} +PAD_CLASS = 0 +EOS_CLASS = len(AA) + 1 +N_CLASSES = len(AA) + 2 +TASKS = ("sequence", "tcremp", "pgen") +SPECIALISTS = {"sequence": "r", "tcremp": "t", "pgen": "p"} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", type=Path, required=True) + parser.add_argument("--normalizer", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--model", action="append", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=256) + parser.add_argument("--alphas", default="0.01,1,100") + parser.add_argument("--bootstrap", type=int, default=2000) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def write_json(path: Path, payload: object) -> None: + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def sequence_targets(sequences: list[str], max_length: int) -> np.ndarray: + """Return fixed targets: PAD=0, amino acids=1..20, EOS=21.""" + targets = np.full((len(sequences), max_length + 1), PAD_CLASS, dtype=np.uint8) + for row, sequence in enumerate(sequences): + if not sequence or len(sequence) > max_length or any(aa not in AA_TO_CLASS for aa in sequence): + raise ValueError(f"Invalid benchmark sequence at row {row}: {sequence!r}") + targets[row, : len(sequence)] = [AA_TO_CLASS[aa] for aa in sequence] + targets[row, len(sequence)] = EOS_CLASS + return targets + + +def one_hot_flat(targets: np.ndarray) -> np.ndarray: + eye = np.eye(N_CLASSES, dtype=np.float32) + return eye[targets].reshape(len(targets), -1) + + +def sequence_metrics(scores: np.ndarray, targets: np.ndarray) -> tuple[dict[str, float], np.ndarray]: + predicted = scores.reshape(len(targets), targets.shape[1], N_CLASSES).argmax(axis=2) + active = np.arange(targets.shape[1])[None, :] <= (targets == EOS_CLASS).argmax(axis=1)[:, None] + correct = predicted == targets + row_error = 1.0 - (correct & active).sum(axis=1) / active.sum(axis=1) + exact = np.all(correct | ~active, axis=1) + predicted_eos = np.where((predicted == EOS_CLASS).any(axis=1), + (predicted == EOS_CLASS).argmax(axis=1), targets.shape[1]) + true_length = (targets == EOS_CLASS).argmax(axis=1) + return { + "token_accuracy": float(1.0 - row_error.mean()), + "exact_accuracy": float(exact.mean()), + "length_accuracy": float((predicted_eos == true_length).mean()), + "task_loss": float(row_error.mean()), + }, row_error.astype(np.float64) + + +def tcremp_metrics(predicted: np.ndarray, target: np.ndarray) -> tuple[dict[str, float], np.ndarray]: + residual = predicted - target + row_mse = np.mean(residual * residual, axis=1) + pred_norm = np.linalg.norm(predicted, axis=1) + target_norm = np.linalg.norm(target, axis=1) + safe = (pred_norm > 0) & (target_norm > 0) + cosine = np.zeros(len(target), dtype=np.float64) + cosine[safe] = np.sum(predicted[safe] * target[safe], axis=1) / ( + pred_norm[safe] * target_norm[safe] + ) + return { + "standardized_mse": float(row_mse.mean()), + "standardized_rmse": float(np.sqrt(row_mse.mean())), + "mean_cosine": float(cosine.mean()), + "task_loss": float(row_mse.mean()), + }, row_mse.astype(np.float64) + + +def pgen_metrics(predicted: np.ndarray, target: np.ndarray, std: float) -> tuple[dict[str, float], np.ndarray]: + residual = predicted.reshape(-1) - target.reshape(-1) + squared = residual * residual + raw = residual * std + return { + "standardized_mse": float(squared.mean()), + "standardized_rmse": float(np.sqrt(squared.mean())), + "raw_rmse_log10_pgen_1mm": float(np.sqrt(np.mean(raw * raw))), + "raw_mae_log10_pgen_1mm": float(np.mean(np.abs(raw))), + "task_loss": float(squared.mean()), + }, squared.astype(np.float64) + + +def standardize_representation( + array: np.ndarray, train_indices: np.ndarray +) -> tuple[np.ndarray, dict[str, object]]: + mean = np.asarray(array[train_indices], dtype=np.float64).mean(axis=0) + std = np.asarray(array[train_indices], dtype=np.float64).std(axis=0) + zero = std < 1e-8 + std[zero] = 1.0 + result = (np.asarray(array, dtype=np.float32) - mean.astype(np.float32)) / std.astype(np.float32) + if not np.isfinite(result).all(): + raise ValueError("Non-finite standardized representation.") + return result.astype(np.float32), { + "train_only": True, + "zero_variance_features": int(zero.sum()), + "dimension": int(result.shape[1]), + } + + +def composition_features( + table: pd.DataFrame, + include_vj: bool, + fit_indices: np.ndarray | None = None, +) -> tuple[np.ndarray, dict[str, object]]: + sequences = table["junction_aa"].astype(str).tolist() + values = np.zeros((len(table), len(AA) + 1), dtype=np.float32) + for row, sequence in enumerate(sequences): + values[row, : len(AA)] = [sequence.count(aa) / len(sequence) for aa in AA] + values[row, -1] = len(sequence) / 40.0 + columns = [f"fraction_{aa}" for aa in AA] + ["length_over_40"] + if include_vj: + if not {"v_call", "j_call"}.issubset(table.columns): + raise ValueError("V/J control requested but benchmark lacks v_call/j_call.") + if fit_indices is None: + raise ValueError("V/J vocabulary must be fitted from explicit train indices.") + categoricals = table[["v_call", "j_call"]].fillna("").astype(str) + encoded_parts = [] + for field, prefix in (("v_call", "v"), ("j_call", "j")): + vocabulary = sorted(set(categoricals.iloc[fit_indices][field])) + mapping = {value: index for index, value in enumerate(vocabulary)} + encoded = np.zeros((len(table), len(vocabulary)), dtype=np.float32) + for row, value in enumerate(categoricals[field]): + index = mapping.get(value) + if index is not None: + encoded[row, index] = 1.0 + encoded_parts.append(encoded) + columns.extend([f"{prefix}_{value}" for value in vocabulary]) + values = np.concatenate([values, *encoded_parts], axis=1) + return values, {"columns": columns, "dimension": len(columns), "sequence_only": not include_vj} + + +def append_intercept(x: np.ndarray) -> np.ndarray: + return np.concatenate([x, np.ones((len(x), 1), dtype=x.dtype)], axis=1) + + +def ridge_weights(x: np.ndarray, y: np.ndarray, alpha: float) -> np.ndarray: + xa = append_intercept(np.asarray(x, dtype=np.float64)) + gram = xa.T @ xa + penalty = np.eye(gram.shape[0], dtype=np.float64) * alpha + penalty[-1, -1] = 0.0 + return np.linalg.solve(gram + penalty, xa.T @ np.asarray(y, dtype=np.float64)) + + +def ridge_predict(x: np.ndarray, weights: np.ndarray) -> np.ndarray: + return append_intercept(np.asarray(x, dtype=np.float64)) @ weights + + +def fit_probe( + x: np.ndarray, + train: np.ndarray, + val: np.ndarray, + test: np.ndarray, + target: np.ndarray, + alphas: list[float], + score, +) -> tuple[float, np.ndarray, np.ndarray, dict[str, float]]: + chosen: tuple[float, float, np.ndarray] | None = None + trace = {} + xa = append_intercept(np.asarray(x[train], dtype=np.float64)) + gram = xa.T @ xa + cross = xa.T @ np.asarray(target[train], dtype=np.float64) + for alpha in alphas: + penalty = np.eye(gram.shape[0], dtype=np.float64) * alpha + penalty[-1, -1] = 0.0 + weights = np.linalg.solve(gram + penalty, cross) + val_pred = ridge_predict(x[val], weights) + loss = float(score(val_pred, val)) + trace[str(alpha)] = loss + if chosen is None or (loss, alpha) < (chosen[0], chosen[1]): + chosen = (loss, alpha, weights) + assert chosen is not None + return chosen[1], ridge_predict(x[val], chosen[2]), ridge_predict(x[test], chosen[2]), trace + + +def aggregate_losses( + losses: dict[str, float], specialist: dict[str, float], null: dict[str, float] +) -> dict[str, object]: + regret = {} + skill = {} + within = {} + for task in TASKS: + gap = null[task] - specialist[task] + if not np.isfinite(gap) or gap <= 0: + raise ValueError(f"Invalid validation specialist/null gap for {task}: {gap}") + regret[task] = float((losses[task] - specialist[task]) / gap) + skill[task] = float(np.clip((null[task] - losses[task]) / gap, 0.0, 1.0)) + within[task] = bool(losses[task] <= 1.05 * specialist[task]) + return { + "normalized_regret_by_task": regret, + "mean_normalized_regret": float(np.mean(list(regret.values()))), + "worst_task_normalized_regret": float(max(regret.values())), + "fixed_reference_hypervolume": float(np.prod(list(skill.values()))), + "tasks_within_5pct_of_specialist_loss": int(sum(within.values())), + "within_5pct_by_task": within, + } + + +def bootstrap_aggregate_deltas( + per_row: dict[str, dict[str, np.ndarray]], + specialist: dict[str, float], + null: dict[str, float], + replicates: int, + seed: int, +) -> dict[str, object]: + rng = np.random.default_rng(seed) + n = len(next(iter(next(iter(per_row.values())).values()))) + output = {} + for comparator in sorted(set(per_row) - {"rtp"}): + mean_delta = np.empty(replicates, dtype=np.float64) + worst_delta = np.empty(replicates, dtype=np.float64) + task_delta = {task: np.empty(replicates, dtype=np.float64) for task in TASKS} + for replicate in range(replicates): + sample = rng.integers(0, n, n) + aggregates = {} + for name in ("rtp", comparator): + losses = {task: float(per_row[name][task][sample].mean()) for task in TASKS} + aggregates[name] = aggregate_losses(losses, specialist, null) + mean_delta[replicate] = (aggregates["rtp"]["mean_normalized_regret"] + - aggregates[comparator]["mean_normalized_regret"]) + worst_delta[replicate] = (aggregates["rtp"]["worst_task_normalized_regret"] + - aggregates[comparator]["worst_task_normalized_regret"]) + for task in TASKS: + task_delta[task][replicate] = ( + per_row["rtp"][task][sample].mean() - per_row[comparator][task][sample].mean() + ) + output[comparator] = { + "orientation": "RTP minus comparator; negative favors RTP for all regret/loss endpoints", + "mean_normalized_regret": ci(mean_delta), + "worst_task_normalized_regret": ci(worst_delta), + "raw_task_loss": {task: ci(values) for task, values in task_delta.items()}, + } + return output + + +def ci(values: np.ndarray) -> dict[str, float]: + return { + "mean": float(values.mean()), + "lower_95": float(np.quantile(values, 0.025)), + "upper_95": float(np.quantile(values, 0.975)), + } + + +def file_manifest(paths: list[Path]) -> dict[str, dict[str, object]]: + return {str(path): {"sha256": sha256(path), "bytes": path.stat().st_size} for path in paths} + + +def main() -> None: + # Keep statistical helpers importable in lightweight local test environments + # that do not have the cluster's PyTorch runtime. + from benchmark.evaluate_vdjdb_frozen import ( + extract_embeddings, + parse_models, + validate_run_configs, + ) + from irrm_codec.multitask_data import ( + TargetStandardizer, + load_prepared_benchmark, + select_split_indices, + ) + + args = parse_args() + started = time.time() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + alphas = sorted({float(item) for item in args.alphas.split(",")}) + if not alphas or any(alpha <= 0 for alpha in alphas): + raise ValueError("All ridge alphas must be positive.") + + table, tcremp_raw = load_prepared_benchmark(args.data_dir) + train = select_split_indices(table, args.data_dir, "train") + val = select_split_indices(table, args.data_dir, "val") + test = select_split_indices(table, args.data_dir, "test") + if set(train) & set(val) or set(train) & set(test) or set(val) & set(test): + raise ValueError("Locked manifests overlap.") + normalizer = TargetStandardizer.load(args.normalizer) + if normalizer.pgen_target != "log10_pgen_1mm" or normalizer.train_rows != len(train): + raise ValueError("Locked train-only normalizer does not match requested target/split.") + tcremp_target = ((np.asarray(tcremp_raw, dtype=np.float32) - normalizer.tcremp_mean) + / normalizer.tcremp_std).astype(np.float32) + pgen_target = ((table["log10_pgen_1mm"].to_numpy(dtype=np.float32) - normalizer.pgen_mean) + / normalizer.pgen_std).astype(np.float32)[:, None] + seq_target = sequence_targets(table["junction_aa"].astype(str).tolist(), 40) + seq_one_hot = one_hot_flat(seq_target) + + models = parse_models(args.model) + model_config, model_report = validate_run_configs(models) + if model_config.latent_dim != 128 or model_config.max_sequence_len != 40: + raise ValueError("This preregistration requires the matched 128-D, length-40 models.") + cohort = table[["junction_aa"]].rename(columns={"junction_aa": "cdr3"}) + embeddings, extraction = extract_embeddings( + cohort, models, model_config, args.tokenizer, args.batch_size, args.output_dir, + ) + representations: dict[str, np.ndarray] = dict(embeddings) + concat = np.concatenate([embeddings[name] for name in ("r", "t", "p")], axis=1) + representations["concat_r_t_p"] = concat + pca = PCA(n_components=128, svd_solver="randomized", random_state=args.seed) + pca.fit(concat[train]) + representations["pca128_concat_r_t_p"] = pca.transform(concat).astype(np.float32) + comp, comp_meta = composition_features(table, include_vj=False) + comp_vj, comp_vj_meta = composition_features(table, include_vj=True, fit_indices=train) + representations["composition_length"] = comp + representations["composition_length_vj"] = comp_vj + + representation_meta = {} + results = {} + per_row = {} + validation_losses = {} + test_losses = {} + for name, raw in representations.items(): + x, scaling = standardize_representation(raw, train) + sequence_alpha, sequence_val_pred, sequence_test_pred, sequence_trace = fit_probe( + x, train, val, test, seq_one_hot, alphas, + lambda pred, indices: sequence_metrics(pred, seq_target[indices])[0]["task_loss"], + ) + t_alpha, t_val_pred, t_test_pred, t_trace = fit_probe( + x, train, val, test, tcremp_target, alphas, + lambda pred, indices: tcremp_metrics(pred, tcremp_target[indices])[0]["task_loss"], + ) + p_alpha, p_val_pred, p_test_pred, p_trace = fit_probe( + x, train, val, test, pgen_target, alphas, + lambda pred, indices: pgen_metrics(pred, pgen_target[indices], normalizer.pgen_std)[0]["task_loss"], + ) + seq_val, _ = sequence_metrics(sequence_val_pred, seq_target[val]) + seq_test, seq_rows = sequence_metrics(sequence_test_pred, seq_target[test]) + t_val, _ = tcremp_metrics(t_val_pred, tcremp_target[val]) + t_test, t_rows = tcremp_metrics(t_test_pred, tcremp_target[test]) + p_val, _ = pgen_metrics(p_val_pred, pgen_target[val], normalizer.pgen_std) + p_test, p_rows = pgen_metrics(p_test_pred, pgen_target[test], normalizer.pgen_std) + validation_losses[name] = { + "sequence": seq_val["task_loss"], "tcremp": t_val["task_loss"], "pgen": p_val["task_loss"], + } + test_losses[name] = { + "sequence": seq_test["task_loss"], "tcremp": t_test["task_loss"], "pgen": p_test["task_loss"], + } + per_row[name] = {"sequence": seq_rows, "tcremp": t_rows, "pgen": p_rows} + dim = x.shape[1] + results[name] = { + "validation": {"sequence": seq_val, "tcremp": t_val, "pgen": p_val}, + "test": {"sequence": seq_test, "tcremp": t_test, "pgen": p_test}, + "selected_alpha_validation_only": { + "sequence": sequence_alpha, "tcremp": t_alpha, "pgen": p_alpha, + }, + "validation_alpha_trace": { + "sequence": sequence_trace, "tcremp": t_trace, "pgen": p_trace, + }, + "probe_parameter_counts": { + "sequence": int((dim + 1) * seq_one_hot.shape[1]), + "tcremp": int((dim + 1) * tcremp_target.shape[1]), + "pgen": int(dim + 1), + }, + } + representation_meta[name] = {**scaling} + del x + + # Train-derived constant probes are the fixed null representation. + majority = np.zeros((len(val), seq_target.shape[1] * N_CLASSES), dtype=np.float32) + majority_test = np.zeros((len(test), seq_target.shape[1] * N_CLASSES), dtype=np.float32) + train_majority = np.apply_along_axis(lambda x: np.bincount(x, minlength=N_CLASSES).argmax(), 0, + seq_target[train]) + for position, cls in enumerate(train_majority): + majority[:, position * N_CLASSES + cls] = 1 + majority_test[:, position * N_CLASSES + cls] = 1 + null_val_seq, _ = sequence_metrics(majority, seq_target[val]) + null_test_seq, null_seq_rows = sequence_metrics(majority_test, seq_target[test]) + zero_t_val = np.zeros_like(tcremp_target[val]) + zero_t_test = np.zeros_like(tcremp_target[test]) + null_val_t, _ = tcremp_metrics(zero_t_val, tcremp_target[val]) + null_test_t, null_t_rows = tcremp_metrics(zero_t_test, tcremp_target[test]) + zero_p_val = np.zeros_like(pgen_target[val]) + zero_p_test = np.zeros_like(pgen_target[test]) + null_val_p, _ = pgen_metrics(zero_p_val, pgen_target[val], normalizer.pgen_std) + null_test_p, null_p_rows = pgen_metrics(zero_p_test, pgen_target[test], normalizer.pgen_std) + validation_null = {"sequence": null_val_seq["task_loss"], "tcremp": null_val_t["task_loss"], + "pgen": null_val_p["task_loss"]} + test_null = {"sequence": null_test_seq["task_loss"], "tcremp": null_test_t["task_loss"], + "pgen": null_test_p["task_loss"]} + per_row["constant_null"] = {"sequence": null_seq_rows, "tcremp": null_t_rows, "pgen": null_p_rows} + validation_specialist = {task: validation_losses[name][task] for task, name in SPECIALISTS.items()} + for name in results: + results[name]["validation_aggregate"] = aggregate_losses( + validation_losses[name], validation_specialist, validation_null, + ) + # Validation-defined reference and scale are intentionally reused on test. + results[name]["test_aggregate_validation_fixed_scale"] = aggregate_losses( + test_losses[name], validation_specialist, validation_null, + ) + + parameter_counts = {name: int(report["parameter_count"]) for name, report in model_report.items()} + for name in model_report: + representation_meta[name].update({ + "encoder_count": 1, "encoder_parameter_count": parameter_counts[name], + "representation_dimension": 128, + }) + representation_meta["concat_r_t_p"].update({ + "encoder_count": 3, + "encoder_parameter_count": sum(parameter_counts[name] for name in ("r", "t", "p")), + "representation_dimension": 384, + }) + representation_meta["pca128_concat_r_t_p"].update({ + "encoder_count": 3, + "encoder_parameter_count": sum(parameter_counts[name] for name in ("r", "t", "p")), + "representation_dimension": 128, + "train_only_pca_parameters": int(pca.components_.size + pca.mean_.size), + "train_explained_variance_ratio": float(pca.explained_variance_ratio_.sum()), + }) + representation_meta["composition_length"].update({ + "encoder_count": 0, "encoder_parameter_count": 0, **comp_meta, + }) + representation_meta["composition_length_vj"].update({ + "encoder_count": 0, "encoder_parameter_count": 0, **comp_vj_meta, + "annotation_aware_modality": True, + }) + + bootstrap = bootstrap_aggregate_deltas( + {name: per_row[name] for name in results}, validation_specialist, validation_null, + args.bootstrap, args.seed + 1000, + ) + input_paths = [ + args.data_dir / "dataset.parquet", args.data_dir / "embeddings.npy", + args.data_dir / "manifests" / "train.tsv", args.data_dir / "manifests" / "val.tsv", + args.data_dir / "manifests" / "test.tsv", args.normalizer, args.tokenizer, + *[path for pair in models.values() for path in pair], + ] + payload = { + "status": "complete", + "study": "frozen DATA-ANCHOR latent sufficiency with fresh linear probes", + "selection": "ridge alpha selected on validation; test evaluated once; all models reported", + "fixed_specialists": SPECIALISTS, + "task_loss_orientation": "lower is better: sequence token error, TCRemP standardized MSE, pgen standardized MSE", + "aggregate_definition": { + "regret": "(loss - fixed validation specialist loss)/(validation constant-null loss - validation specialist loss)", + "hypervolume": "product over tasks of clipped (validation-null - loss)/(validation-null - validation-specialist)", + "test_scale": "validation specialist and null values reused unchanged", + }, + "split_rows": {"train": len(train), "val": len(val), "test": len(test)}, + "split_overlap": 0, + "alphas": alphas, + "models": model_report, + "embedding_extraction": extraction, + "representations": representation_meta, + "validation_null_losses": validation_null, + "test_null_losses_descriptive": test_null, + "validation_specialist_losses": validation_specialist, + "results": results, + "paired_test_clonotype_bootstrap": bootstrap, + "bootstrap_replicates": args.bootstrap, + "limitations": [ + "All encoders are single seed 42; bootstrap intervals do not quantify training-seed variability.", + "Linear-probe sufficiency does not establish causal biological usefulness.", + "The V/J control is annotation-aware and is not an equivalent sequence-only inference condition.", + ], + "runtime_seconds": float(time.time() - started), + "slurm": {"job_id": os.environ.get("SLURM_JOB_ID"), "node": os.environ.get("SLURMD_NODENAME")}, + "inputs": file_manifest(input_paths), + "model_config": asdict(model_config), + } + results_path = args.output_dir / "RESULTS.json" + write_json(results_path, payload) + summary_rows = [] + for name, result in results.items(): + row = {"representation": name} + for task in TASKS: + row[f"val_{task}_loss"] = result["validation"][task]["task_loss"] + row[f"test_{task}_loss"] = result["test"][task]["task_loss"] + row.update({f"test_{key}": value for key, value in + result["test_aggregate_validation_fixed_scale"].items() + if isinstance(value, (int, float))}) + summary_rows.append(row) + summary_path = args.output_dir / "summary.tsv" + pd.DataFrame(summary_rows).to_csv(summary_path, sep="\t", index=False) + outputs = file_manifest([path for path in args.output_dir.iterdir() if path.is_file()]) + write_json(args.output_dir / "OUTPUT_HASHES.json", outputs) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_reconstruction_substitutions.py b/benchmark/evaluate_reconstruction_substitutions.py new file mode 100644 index 0000000..15a6f6d --- /dev/null +++ b/benchmark/evaluate_reconstruction_substitutions.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python +"""Evaluate structured reconstruction substitutions for one frozen IRRM-CODEC run. + +The analysis follows the earlier inverse-model protocol, but runs autoregressive +reconstruction from a specified final multi-task checkpoint on its locked test split. +It counts only wrong amino-acid positions among target/prediction pairs of equal +length and summarizes substitutions within predefined physicochemical groups. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from dataclasses import asdict +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import torch + +from irrm_codec.multitask_data import ( + MultiTaskBenchmarkDataset, + build_multitask_dataloader, + load_prepared_benchmark, + resolve_encoder_tokenizer, + select_split_indices, +) +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import decode + + +AA_ORDER = list("CSTAGPDEQNHRKMILVWYF") +AA_GROUPS = { + "C": set("C"), + "STAGP": set("STAGP"), + "DEQN": set("DEQN"), + "HRK": set("HRK"), + "MILV": set("MILV"), + "WYF": set("WYF"), +} +AA_TO_GROUP = {aa: name for name, aas in AA_GROUPS.items() for aa in aas} +GROUP_BOUNDARIES = np.cumsum([len(group) for group in AA_GROUPS.values()])[:-1] + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-dir", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--num-workers", type=int, default=8) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + return parser.parse_args() + + +def save_figure(matrix: pd.DataFrame, output_path: Path) -> None: + fig, ax = plt.subplots(figsize=(9, 7.5)) + image = ax.imshow(matrix.to_numpy(), cmap="Greys") + fig.colorbar(image, ax=ax, label="Wrong aligned positions") + ax.set_xticks(np.arange(len(AA_ORDER)), AA_ORDER) + ax.set_yticks(np.arange(len(AA_ORDER)), AA_ORDER) + for boundary in GROUP_BOUNDARIES: + ax.axhline(boundary - 0.5, color="black", linewidth=1.5) + ax.axvline(boundary - 0.5, color="black", linewidth=1.5) + ax.set_xlabel("Predicted amino acid") + ax.set_ylabel("Target amino acid") + ax.set_title("RTP / DATA-ANCHOR: substitution errors on locked TRB test split") + fig.tight_layout() + fig.savefig(output_path.with_suffix(".png"), dpi=240, bbox_inches="tight") + fig.savefig(output_path.with_suffix(".svg"), bbox_inches="tight") + fig.savefig(output_path.with_suffix(".pdf"), bbox_inches="tight") + plt.close(fig) + + +def main() -> None: + args = parse_args() + if args.batch_size < 1: + raise ValueError("batch-size must be positive") + if args.device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but unavailable") + for path in (args.data_dir / "READY.json", args.tokenizer, args.checkpoint): + if not path.is_file(): + raise FileNotFoundError(path) + + ready = json.loads((args.data_dir / "READY.json").read_text(encoding="utf-8")) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Benchmark READY.json is not accepting") + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + config = IRRMCodecConfig(**checkpoint["model_config"]) + tokenizer_info = checkpoint["tokenizer"] + if tokenizer_info["type"] != "data_anchor": + raise ValueError(f"Expected a DATA-ANCHOR checkpoint, got {tokenizer_info['type']!r}") + tokenizer = resolve_encoder_tokenizer("data_anchor", str(args.tokenizer)) + if tokenizer.vocab_size != config.input_vocab_size: + raise ValueError("Tokenizer vocabulary differs from checkpoint input embedding") + + table, embeddings = load_prepared_benchmark(args.data_dir) + test_indices = select_split_indices(table, args.data_dir, "test") + dataset = MultiTaskBenchmarkDataset( + table, embeddings, test_indices, tokenizer, "log10_pgen_1mm", config.max_sequence_len + ) + device = torch.device(args.device) + loader = build_multitask_dataloader( + dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=device.type == "cuda" + ) + model = IRRMCodecTransformer(config).to(device) + model.load_state_dict(checkpoint["model_state"], strict=True) + model.eval() + + matrix = pd.DataFrame(0, index=AA_ORDER, columns=AA_ORDER, dtype=np.int64) + total_sequences = 0 + exact_sequences = 0 + length_matched_sequences = 0 + length_mismatched_sequences = 0 + wrong_positions = 0 + within_group = 0 + between_group = 0 + unexpected_residues = 0 + + with torch.inference_mode(): + for batch in loader: + tokens = batch["encoder_tokens"].to(device, non_blocking=True) + mask = batch["encoder_mask"].to(device, non_blocking=True) + generated = model.reconstruct(tokens, mask).cpu().numpy() + predictions = [decode(row.tolist()) for row in generated] + for target, prediction in zip(batch["sequence"], predictions): + total_sequences += 1 + if target == prediction: + exact_sequences += 1 + if len(target) != len(prediction): + length_mismatched_sequences += 1 + continue + length_matched_sequences += 1 + for target_aa, predicted_aa in zip(target, prediction): + if target_aa == predicted_aa: + continue + if target_aa not in AA_TO_GROUP or predicted_aa not in AA_TO_GROUP: + unexpected_residues += 1 + continue + matrix.loc[target_aa, predicted_aa] += 1 + wrong_positions += 1 + if AA_TO_GROUP[target_aa] == AA_TO_GROUP[predicted_aa]: + within_group += 1 + else: + between_group += 1 + + if total_sequences != len(test_indices): + raise RuntimeError(f"Observed {total_sequences} sequences, expected {len(test_indices)}") + if wrong_positions != within_group + between_group: + raise RuntimeError("Substitution counts are inconsistent") + args.output_dir.mkdir(parents=True, exist_ok=False) + matrix.to_csv(args.output_dir / "substitution_matrix.tsv", sep="\t") + summary = { + "status": "complete", + "analysis": "locked_test_autoregressive_substitution_errors", + "checkpoint": str(args.checkpoint), + "checkpoint_sha256": sha256(args.checkpoint), + "data_dir": str(args.data_dir), + "benchmark_ready_sha256": sha256(args.data_dir / "READY.json"), + "tokenizer": str(args.tokenizer), + "tokenizer_sha256": sha256(args.tokenizer), + "model_config": asdict(config), + "test_sequences": total_sequences, + "exact_sequences": exact_sequences, + "exact_match": exact_sequences / total_sequences, + "length_matched_sequences": length_matched_sequences, + "length_mismatched_sequences": length_mismatched_sequences, + "wrong_aligned_positions": wrong_positions, + "within_group_errors": within_group, + "between_group_errors": between_group, + "within_group_fraction": within_group / wrong_positions if wrong_positions else None, + "between_group_fraction": between_group / wrong_positions if wrong_positions else None, + "unexpected_residue_errors": unexpected_residues, + "groups": {name: "".join(sorted(aas)) for name, aas in AA_GROUPS.items()}, + "note": "Only incorrect aligned positions from equal-length target/prediction pairs enter the matrix.", + } + (args.output_dir / "RESULTS.json").write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8") + save_figure(matrix, args.output_dir / "substitution_confusion") + print(json.dumps(summary, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_rtp_error_analysis.py b/benchmark/evaluate_rtp_error_analysis.py new file mode 100644 index 0000000..d8e9eef --- /dev/null +++ b/benchmark/evaluate_rtp_error_analysis.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python +"""Locked-test error analysis for the DATA-ANCHOR RTP checkpoint. + +Produces reconstruction and pgen error stratified by CDR3 length, relative +position profiles, a grouped amino-acid substitution matrix, and an empirical +BLOSUM62 comparison. Edit distance is used for the primary residue error rate +so insertions/deletions are not silently omitted; substitution analyses retain +the explicitly labelled equal-length, aligned-position definition. +""" +from __future__ import annotations + +import argparse, hashlib, json +from collections import Counter, defaultdict +from dataclasses import asdict +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import torch +from Bio.Align import substitution_matrices + +from irrm_codec.multitask_data import (MultiTaskBenchmarkDataset, TargetStandardizer, + build_multitask_dataloader, load_prepared_benchmark, resolve_encoder_tokenizer, select_split_indices) +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import decode + +AA_ORDER = list("CSTAGPDEQNHRKMILVWYF") +GROUPS = ["C", "STAGP", "DEQN", "HRK", "MILV", "WYF"] +BOUNDARIES = np.cumsum([len(x) for x in GROUPS])[:-1] + +def sha256(path): + h = hashlib.sha256() + with Path(path).open("rb") as f: + for b in iter(lambda: f.read(1048576), b""): h.update(b) + return h.hexdigest() + +def edit_distance(a, b): + prev = list(range(len(b) + 1)) + for i, x in enumerate(a, 1): + cur = [i] + for j, y in enumerate(b, 1): cur.append(min(prev[j] + 1, cur[-1] + 1, prev[j-1] + (x != y))) + prev = cur + return prev[-1] + +def args(): + p = argparse.ArgumentParser(description=__doc__) + for n in ("data-dir", "tokenizer", "checkpoint", "output-dir"): p.add_argument("--" + n, type=Path, required=True) + p.add_argument("--batch-size", type=int, default=64); p.add_argument("--num-workers", type=int, default=8) + p.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + return p.parse_args() + +def rate_rows(items): + rows = [] + for length, x in sorted(items.items()): + n=x["n"] + rows.append({"cdr3_length": length, "n_sequences": n, "exact_sequences": x["exact"], + "sequence_error_rate": 1-x["exact"]/n, "total_target_residues": x["res"], + "edit_errors": x["edit"], "residue_error_rate": x["edit"]/x["res"], + "mean_abs_log10_pgen_error": np.mean(x["pabs"]), "median_abs_log10_pgen_error": np.median(x["pabs"]), + "mean_signed_log10_pgen_error": np.mean(x["psigned"]), "mean_true_log10_pgen": np.mean(x["truep"])}) + return pd.DataFrame(rows) + +def save_figures(out, matrix, length, pos, exact): + plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False}) + fig, ax = plt.subplots(figsize=(8.2, 6.8)); im=ax.imshow(matrix, cmap="Greys"); fig.colorbar(im, ax=ax, label="Wrong aligned positions") + ax.set(xticks=range(20), xticklabels=AA_ORDER, yticks=range(20), yticklabels=AA_ORDER, xlabel="Predicted amino acid", ylabel="True amino acid", title="RTP reconstruction substitution errors") + for x in BOUNDARIES: ax.axhline(x-.5,color="black",lw=1.25); ax.axvline(x-.5,color="black",lw=1.25) + fig.tight_layout(); [fig.savefig((out/"figure_error_substitution_matrix").with_suffix("."+s), dpi=300 if s=="png" else None, bbox_inches="tight") for s in ("png","svg","pdf")]; plt.close(fig) + fig, ax = plt.subplots(2,1,figsize=(8,6.4),sharex=True); x=length.cdr3_length + ax[0].plot(x,100*length.sequence_error_rate,"o-",label="Sequence error",color="#332288"); ax[0].plot(x,100*length.residue_error_rate,"o-",label="Residue error (edit-aware)",color="#EE6677"); ax[0].set(ylabel="Error rate (%)",title="Reconstruction error increases with CDR3 length"); ax[0].legend(); ax[1].bar(x,length.n_sequences,color="#88CCEE"); ax[1].set(xlabel="CDR3 length (aa)",ylabel="Test sequences (N)") + for a in ax: a.grid(axis="y",color="#ddd"); + fig.tight_layout(); [fig.savefig((out/"figure_error_reconstruction_by_length").with_suffix("."+s),dpi=300 if s=="png" else None,bbox_inches="tight") for s in ("png","svg","pdf")]; plt.close(fig) + fig, ax = plt.subplots(figsize=(8,4.5)); + for g, z in pos.groupby("length_group", sort=False): ax.plot(z.relative_bin+1,100*z.error_rate,"o-",label=g) + ax.set(xticks=range(1,11),xlabel="Relative CDR3 position decile (N → C)",ylabel="Aligned residue error (%)",title="Central reconstruction errors persist within length strata"); ax.grid(axis="y",color="#ddd"); ax.legend(title="CDR3 length"); fig.tight_layout(); [fig.savefig((out/"figure_error_relative_position").with_suffix("."+s),dpi=300 if s=="png" else None,bbox_inches="tight") for s in ("png","svg","pdf")]; plt.close(fig) + fig, ax = plt.subplots(2,1,figsize=(8,6.4),sharex=True); ax[0].plot(x,length.mean_abs_log10_pgen_error,"o-",color="#117733"); ax[0].set(ylabel="Mean absolute log10(pgen) error",title="pgen error by CDR3 length"); ax[1].bar(x,length.n_sequences,color="#88CCEE"); ax[1].set(xlabel="CDR3 length (aa)",ylabel="Test sequences (N)") + for a in ax: a.grid(axis="y",color="#ddd") + fig.tight_layout(); [fig.savefig((out/"figure_error_pgen_by_length").with_suffix("."+s),dpi=300 if s=="png" else None,bbox_inches="tight") for s in ("png","svg","pdf")]; plt.close(fig) + fig, ax=plt.subplots(figsize=(5.5,4.3)); ax.boxplot([exact.loc[exact.exact_reconstruction,"abs_log10_pgen_error"],exact.loc[~exact.exact_reconstruction,"abs_log10_pgen_error"]],tick_labels=["Exact","Non-exact"],showfliers=False); ax.set(ylabel="Absolute log10(pgen) error",title="pgen error and reconstruction outcome"); ax.grid(axis="y",color="#ddd"); fig.tight_layout(); [fig.savefig((out/"figure_error_pgen_by_reconstruction").with_suffix("."+s),dpi=300 if s=="png" else None,bbox_inches="tight") for s in ("png","svg","pdf")]; plt.close(fig) + # IEEE-ready single-column view: 3.5-in wide with an offset count axis. + plt.rcParams.update({"font.family": "Arial", "font.size": 8, "pdf.fonttype": 42, "ps.fonttype": 42}) + plot = length.loc[length.cdr3_length.between(9, 25)].copy() + x = plot.cdr3_length.to_numpy() + residue_error = 100 * plot.residue_error_rate.to_numpy() + pgen_error = plot.mean_abs_log10_pgen_error.to_numpy() + n_sequences = plot.n_sequences.to_numpy() + fig, ax1 = plt.subplots(figsize=(3.5, 3.15)) + ax3 = ax1.twinx() + ax3.spines["right"].set_position(("axes", 1.30)) + bars = ax3.bar(x, n_sequences, width=0.68, color="#BDBDBD", edgecolor="none", alpha=0.75, label="Test sequences") + ax3.set_ylim(0, n_sequences.max() * 1.16) + ax3.set_yticks([]) + ax3.spines["right"].set_visible(False) + line1, = ax1.plot(x, residue_error, marker="o", markersize=3.8, linewidth=1.35, color="#0072B2", label="Residue error") + ax1.set(xlabel="CDR3 length (aa)", ylabel="Residue (%)", xlim=(8.5, 25.5), xticks=np.arange(9, 26, 2), ylim=(0, max(15, residue_error.max() * 1.05))) + ax1.grid(axis="y", color="#D9D9D9", linewidth=0.5) + ax2 = ax1.twinx() + line2, = ax2.plot(x, pgen_error, marker="s", markersize=3.6, linewidth=1.35, linestyle="--", color="#D55E00", label=r"Mean absolute $\log_{10}(P_{gen})$ error") + ax2.set_ylabel(r"MAE $\log_{10}(P_{gen})$", labelpad=1) + ax2.set_ylim(0, max(0.62, pgen_error.max() * 1.05)) + for xi, ni in zip(x, n_sequences): + if ni < 50: + ax3.annotate(f"n={ni}", (xi, ni), textcoords="offset points", xytext=(0, 2), ha="center", fontsize=6.5, color="#4D4D4D") + ax1.legend([bars, line1, line2], ["Test sequences", "Residue error", r"Mean $|\Delta\log_{10}P_{gen}|$"], loc="upper left", frameon=False, fontsize=7, handlelength=1.8, borderpad=0.15, labelspacing=0.25) + for axis in (ax1, ax2, ax3): + axis.tick_params(direction="out", length=2.5, width=0.7, labelsize=7.5) + fig.subplots_adjust(left=0.16, right=0.82, bottom=0.18, top=0.97) + [fig.savefig((out/"figure_error_length_combined_ieee").with_suffix("."+s), dpi=600 if s=="png" else None, bbox_inches="tight") for s in ("png", "svg", "pdf")] + plt.close(fig) + +def main(): + a=args(); out=a.output_dir + if a.device=="cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable") + for path in (a.data_dir/"READY.json",a.data_dir/"target_standardizer.npz",a.tokenizer,a.checkpoint): + if not path.is_file(): raise FileNotFoundError(path) + ready=json.loads((a.data_dir/"READY.json").read_text()); + if ready.get("status")!="ready" or not all(ready.get("checks",{}).values()): raise ValueError("Benchmark is not accepting") + payload=torch.load(a.checkpoint,map_location="cpu",weights_only=False); config=IRRMCodecConfig(**payload["model_config"]) + if payload["tokenizer"]["type"]!="data_anchor": raise ValueError("Expected DATA-ANCHOR checkpoint") + tok=resolve_encoder_tokenizer("data_anchor",str(a.tokenizer)); table, embeddings=load_prepared_benchmark(a.data_dir); idx=select_split_indices(table,a.data_dir,"test") + ds=MultiTaskBenchmarkDataset(table,embeddings,idx,tok,"log10_pgen_1mm",config.max_sequence_len); dev=torch.device(a.device) + loader=build_multitask_dataloader(ds,a.batch_size,False,a.num_workers,pin_memory=dev.type=="cuda") + model=IRRMCodecTransformer(config).to(dev); model.load_state_dict(payload["model_state"],strict=True); model.eval(); norm=TargetStandardizer.load(a.data_dir/"target_standardizer.npz") + bylen=defaultdict(lambda:{"n":0,"exact":0,"res":0,"edit":0,"pabs":[],"psigned":[],"truep":[]}); mat=np.zeros((20,20),dtype=int); pos=defaultdict(lambda:[0,0]); region=defaultdict(lambda:[0,0]); records=[]; scores=[]; target_errors=[]; pred_errors=[] + groups=(("≤16",lambda n:n<=16),("17–18",lambda n:17<=n<=18),("19–20",lambda n:19<=n<=20),("≥21",lambda n:n>=21)) + with torch.inference_mode(): + for batch in loader: + t=batch["encoder_tokens"].to(dev); m=batch["encoder_mask"].to(dev); generated=[decode(r.tolist()) for r in model.reconstruct(t,m).cpu().numpy()] + pred=(model(t,m)["pgen_standardized"].squeeze(-1).cpu().numpy()*norm.pgen_std+norm.pgen_mean) + for target, guess, truep, pp in zip(batch["sequence"],generated,batch["pgen_target"].numpy(),pred): + n=len(target); exact=target==guess; ed=edit_distance(target,guess); signed=float(pp-truep); item=bylen[n]; item["n"]+=1; item["exact"]+=exact; item["res"]+=n; item["edit"]+=ed; item["pabs"].append(abs(signed)); item["psigned"].append(signed); item["truep"].append(float(truep)); records.append({"cdr3_length":n,"exact_reconstruction":exact,"edit_errors":ed,"residue_error_rate":ed/n,"true_log10_pgen":float(truep),"pred_log10_pgen":float(pp),"abs_log10_pgen_error":abs(signed),"signed_log10_pgen_error":signed}) + if len(target)==len(guess): + for j,(x,y) in enumerate(zip(target,guess)): + if x!=y: + mat[AA_ORDER.index(x),AA_ORDER.index(y)]+=1; scores.append(float(substitution_matrices.load("BLOSUM62")[x,y])); target_errors.append(x); pred_errors.append(y) + for label, cond in groups: + if cond(n): + key=(label,min(9,int(j*10/n))); pos[key][0]+=1; pos[key][1]+=x!=y + r="N_edge_1_3" if j<3 else "C_edge_1_3" if j>=n-3 else "middle"; region[r][0]+=1; region[r][1]+=x!=y + rows=[] + for (g,b),(eligible,wrong) in pos.items(): rows.append({"length_group":g,"relative_bin":b,"eligible_positions":eligible,"wrong_positions":wrong,"error_rate":wrong/eligible}) + position=pd.DataFrame(rows).sort_values(["length_group","relative_bin"]); length=rate_rows(bylen); rec=pd.DataFrame(records); reg=pd.DataFrame([{"region":k,"eligible_positions":v[0],"wrong_positions":v[1],"error_rate":v[1]/v[0]} for k,v in region.items()]) + rng=np.random.default_rng(42); pred_counts=Counter(pred_errors); choices=np.array(list(pred_counts)); weights=np.array([pred_counts[x] for x in choices],float); weights/=weights.sum(); baseline=[] + for x in target_errors: + keep=choices!=x; w=weights[keep]; w/=w.sum(); baseline.append(float(substitution_matrices.load("BLOSUM62")[x,rng.choice(choices[keep],p=w)])) + bl=pd.DataFrame({"observed_blosum62":scores,"frequency_matched_random_blosum62":baseline}) + summary={"status":"complete","analysis":"locked_test_rtp_error_analysis","test_sequences":len(rec),"checkpoint_sha256":sha256(a.checkpoint),"benchmark_ready_sha256":sha256(a.data_dir/"READY.json"),"normalizer_sha256":sha256(a.data_dir/"target_standardizer.npz"),"model_config":asdict(config),"blosum62":{"observed_mean":float(np.mean(scores)),"observed_median":float(np.median(scores)),"positive_fraction":float(np.mean(np.array(scores)>0)),"baseline_mean":float(np.mean(baseline)),"baseline_median":float(np.median(baseline)),"mean_difference_observed_minus_baseline":float(np.mean(scores)-np.mean(baseline))},"pgen_abs_error":{"exact_mean":float(rec[rec.exact_reconstruction].abs_log10_pgen_error.mean()),"nonexact_mean":float(rec[~rec.exact_reconstruction].abs_log10_pgen_error.mean()),"spearman_residue_error":float(rec.residue_error_rate.corr(rec.abs_log10_pgen_error,method="spearman")),"spearman_length":float(rec.cdr3_length.corr(rec.abs_log10_pgen_error,method="spearman"))},"note":"Substitutions and position profiles include aligned equal-length predictions only. Primary residue error is Levenshtein edit distance divided by target length."} + out.mkdir(parents=True,exist_ok=False); pd.DataFrame(mat,index=AA_ORDER,columns=AA_ORDER).to_csv(out/"substitution_matrix.tsv",sep="\t"); length.to_csv(out/"summary_by_cdr3_length.tsv",sep="\t",index=False); position.to_csv(out/"relative_position_error_rates.tsv",sep="\t",index=False); reg.to_csv(out/"edge_middle_error_rates.tsv",sep="\t",index=False); rec.to_csv(out/"per_sequence_errors.tsv",sep="\t",index=False); bl.to_csv(out/"blosum62_observed_vs_frequency_matched_random.tsv",sep="\t",index=False); (out/"RESULTS.json").write_text(json.dumps(summary,indent=2,sort_keys=True)+"\n") + save_figures(out,mat,length,position,rec); print(json.dumps(summary,indent=2,sort_keys=True)) + +if __name__ == "__main__": main() diff --git a/benchmark/evaluate_rtp_reconstruction_by_length.py b/benchmark/evaluate_rtp_reconstruction_by_length.py new file mode 100644 index 0000000..a38323d --- /dev/null +++ b/benchmark/evaluate_rtp_reconstruction_by_length.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +"""Quantify frozen RTP/DATA-ANCHOR reconstruction errors by target CDR3 length.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from dataclasses import asdict +from collections import defaultdict +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import torch + +from irrm_codec.multitask_data import MultiTaskBenchmarkDataset, build_multitask_dataloader, load_prepared_benchmark, resolve_encoder_tokenizer, select_split_indices +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import decode + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-dir", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--num-workers", type=int, default=8) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.batch_size < 1: + raise ValueError("batch-size must be positive") + if args.device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but unavailable") + for path in (args.data_dir / "READY.json", args.tokenizer, args.checkpoint): + if not path.is_file(): + raise FileNotFoundError(path) + ready = json.loads((args.data_dir / "READY.json").read_text(encoding="utf-8")) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Benchmark READY.json is not accepting") + + payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + config = IRRMCodecConfig(**payload["model_config"]) + if payload["tokenizer"]["type"] != "data_anchor": + raise ValueError("This evaluation is registered for DATA-ANCHOR checkpoints only") + tokenizer = resolve_encoder_tokenizer("data_anchor", str(args.tokenizer)) + if tokenizer.vocab_size != config.input_vocab_size: + raise ValueError("Tokenizer vocabulary differs from checkpoint input embedding") + table, embeddings = load_prepared_benchmark(args.data_dir) + indices = select_split_indices(table, args.data_dir, "test") + dataset = MultiTaskBenchmarkDataset(table, embeddings, indices, tokenizer, "log10_pgen_1mm", config.max_sequence_len) + device = torch.device(args.device) + loader = build_multitask_dataloader(dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=device.type == "cuda") + model = IRRMCodecTransformer(config).to(device) + model.load_state_dict(payload["model_state"], strict=True) + model.eval() + + counts = defaultdict(lambda: {"sequences": 0, "exact": 0, "length_mismatch": 0, "aligned_positions": 0, "wrong_aligned_positions": 0}) + with torch.inference_mode(): + for batch in loader: + tokens = batch["encoder_tokens"].to(device, non_blocking=True) + mask = batch["encoder_mask"].to(device, non_blocking=True) + predictions = [decode(row.tolist()) for row in model.reconstruct(tokens, mask).cpu().numpy()] + for target, prediction in zip(batch["sequence"], predictions): + item = counts[len(target)] + item["sequences"] += 1 + item["exact"] += int(target == prediction) + if len(target) != len(prediction): + item["length_mismatch"] += 1 + continue + item["aligned_positions"] += len(target) + item["wrong_aligned_positions"] += sum(a != b for a, b in zip(target, prediction)) + + rows = [] + for length, item in sorted(counts.items()): + sequences = item["sequences"] + rows.append({ + "target_length": length, + **item, + "sequence_error_rate": 1 - item["exact"] / sequences, + "aligned_position_error_rate": item["wrong_aligned_positions"] / item["aligned_positions"] if item["aligned_positions"] else np.nan, + }) + result = pd.DataFrame(rows) + if int(result["sequences"].sum()) != len(indices): + raise RuntimeError("Processed sequence count differs from locked test split") + + args.output_dir.mkdir(parents=True, exist_ok=False) + result.to_csv(args.output_dir / "error_rates_by_target_length.tsv", sep="\t", index=False) + fig, axes = plt.subplots(1, 2, figsize=(12, 4.2)) + axes[0].plot(result["target_length"], result["sequence_error_rate"] * 100, marker="o", color="#332288") + axes[0].set_xlabel("Target CDR3 length (aa)") + axes[0].set_ylabel("Sequences not reconstructed exactly (%)") + axes[0].set_title("Sequence-level reconstruction error") + axes[1].plot(result["target_length"], result["aligned_position_error_rate"] * 100, marker="o", color="#EE6677") + axes[1].set_xlabel("Target CDR3 length (aa)") + axes[1].set_ylabel("Wrong aligned positions (%)") + axes[1].set_title("Position-level reconstruction error") + for axis in axes: + axis.grid(axis="y", color="#DDDDDD", linewidth=0.7) + axis.spines[["top", "right"]].set_visible(False) + fig.suptitle("RTP / DATA-ANCHOR reconstruction error by target CDR3 length", weight="bold") + fig.tight_layout(rect=(0, 0, 1, 0.93)) + for suffix in ("png", "svg", "pdf"): + fig.savefig((args.output_dir / "reconstruction_error_by_length").with_suffix(f".{suffix}"), dpi=240 if suffix == "png" else None, bbox_inches="tight") + plt.close(fig) + summary = { + "status": "complete", + "analysis": "locked_test_autoregressive_error_by_target_cdr3_length", + "checkpoint": str(args.checkpoint), + "checkpoint_sha256": sha256(args.checkpoint), + "tokenizer_sha256": sha256(args.tokenizer), + "benchmark_ready_sha256": sha256(args.data_dir / "READY.json"), + "model_config": asdict(config), + "test_sequences": len(indices), + "length_range": [int(result["target_length"].min()), int(result["target_length"].max())], + } + (args.output_dir / "RESULTS.json").write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(summary, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_rtp_reconstruction_by_position.py b/benchmark/evaluate_rtp_reconstruction_by_position.py new file mode 100644 index 0000000..ccf5ef7 --- /dev/null +++ b/benchmark/evaluate_rtp_reconstruction_by_position.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python +"""Map autoregressive RTP/DATA-ANCHOR reconstruction errors along CDR3 position.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from dataclasses import asdict +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import torch + +from irrm_codec.multitask_data import ( + MultiTaskBenchmarkDataset, + build_multitask_dataloader, + load_prepared_benchmark, + resolve_encoder_tokenizer, + select_split_indices, +) +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import decode + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-dir", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--num-workers", type=int, default=8) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + return parser.parse_args() + + +def rate_table(counts: dict[int, list[int]], label: str) -> pd.DataFrame: + rows = [] + for position, (eligible, wrong) in sorted(counts.items()): + rows.append( + { + label: position, + "eligible_positions": eligible, + "wrong_positions": wrong, + "error_rate": wrong / eligible if eligible else np.nan, + } + ) + return pd.DataFrame(rows) + + +def save_figure(absolute: pd.DataFrame, edge: pd.DataFrame, regions: pd.DataFrame, stem: Path) -> None: + fig, axes = plt.subplots(1, 3, figsize=(15, 4.4)) + axes[0].plot(absolute["position_1based"], absolute["error_rate"] * 100, marker="o", color="#332288") + axes[0].set_xlabel("Position from N terminus") + axes[0].set_ylabel("Autoregressive error rate (%)") + axes[0].set_title("Absolute CDR3 position") + + axes[1].plot(edge["distance_from_nearest_edge"], edge["error_rate"] * 100, marker="o", color="#EE6677") + axes[1].set_xlabel("Distance from nearest CDR3 edge") + axes[1].set_ylabel("Autoregressive error rate (%)") + axes[1].set_title("Distance from either edge") + + axes[2].bar(regions["region"], regions["error_rate"] * 100, color=["#4477AA", "#CCBB44", "#66CCEE"]) + axes[2].set_ylabel("Autoregressive error rate (%)") + axes[2].set_title("First/last 3 residues versus middle") + for idx, row in regions.iterrows(): + axes[2].text(idx, row["error_rate"] * 100, f"{row['wrong_positions']}/{row['eligible_positions']}", ha="center", va="bottom", fontsize=8) + + for ax in axes: + ax.grid(axis="y", color="#DDDDDD", linewidth=0.7) + ax.spines[["top", "right"]].set_visible(False) + fig.suptitle("RTP / DATA-ANCHOR reconstruction errors along locked-test CDR3 sequences", weight="bold") + fig.tight_layout(rect=(0, 0, 1, 0.93)) + fig.savefig(stem.with_suffix(".png"), dpi=240, bbox_inches="tight") + fig.savefig(stem.with_suffix(".svg"), bbox_inches="tight") + fig.savefig(stem.with_suffix(".pdf"), bbox_inches="tight") + plt.close(fig) + + +def main() -> None: + args = parse_args() + if args.batch_size < 1: + raise ValueError("batch-size must be positive") + if args.device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but unavailable") + for path in (args.data_dir / "READY.json", args.tokenizer, args.checkpoint): + if not path.is_file(): + raise FileNotFoundError(path) + + ready = json.loads((args.data_dir / "READY.json").read_text(encoding="utf-8")) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Benchmark READY.json is not accepting") + payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + config = IRRMCodecConfig(**payload["model_config"]) + if payload["tokenizer"]["type"] != "data_anchor": + raise ValueError("This evaluation is registered for DATA-ANCHOR checkpoints only") + tokenizer = resolve_encoder_tokenizer("data_anchor", str(args.tokenizer)) + if tokenizer.vocab_size != config.input_vocab_size: + raise ValueError("Tokenizer vocabulary differs from checkpoint input embedding") + + table, embeddings = load_prepared_benchmark(args.data_dir) + indices = select_split_indices(table, args.data_dir, "test") + dataset = MultiTaskBenchmarkDataset(table, embeddings, indices, tokenizer, "log10_pgen_1mm", config.max_sequence_len) + device = torch.device(args.device) + loader = build_multitask_dataloader(dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=device.type == "cuda") + model = IRRMCodecTransformer(config).to(device) + model.load_state_dict(payload["model_state"], strict=True) + model.eval() + + absolute: dict[int, list[int]] = {} + edge: dict[int, list[int]] = {} + region = {name: [0, 0] for name in ("N_edge_1_3", "middle", "C_edge_1_3")} + matched_sequences = 0 + mismatched_sequences = 0 + + with torch.inference_mode(): + for batch in loader: + tokens = batch["encoder_tokens"].to(device, non_blocking=True) + mask = batch["encoder_mask"].to(device, non_blocking=True) + predictions = [decode(row.tolist()) for row in model.reconstruct(tokens, mask).cpu().numpy()] + for target, prediction in zip(batch["sequence"], predictions): + if len(target) != len(prediction): + mismatched_sequences += 1 + continue + matched_sequences += 1 + for position, (target_aa, prediction_aa) in enumerate(zip(target, prediction)): + wrong = int(target_aa != prediction_aa) + absolute.setdefault(position + 1, [0, 0]) + absolute[position + 1][0] += 1 + absolute[position + 1][1] += wrong + distance = min(position, len(target) - position - 1) + edge.setdefault(distance, [0, 0]) + edge[distance][0] += 1 + edge[distance][1] += wrong + region_name = "N_edge_1_3" if position < 3 else "C_edge_1_3" if position >= len(target) - 3 else "middle" + region[region_name][0] += 1 + region[region_name][1] += wrong + + if matched_sequences + mismatched_sequences != len(indices): + raise RuntimeError("Processed sequence count differs from locked test split") + absolute_table = rate_table(absolute, "position_1based") + edge_table = rate_table(edge, "distance_from_nearest_edge") + region_table = pd.DataFrame( + [{"region": name, "eligible_positions": values[0], "wrong_positions": values[1], "error_rate": values[1] / values[0]} for name, values in region.items()] + ) + if int(region_table["wrong_positions"].sum()) != int(absolute_table["wrong_positions"].sum()): + raise RuntimeError("Position and region error totals differ") + + args.output_dir.mkdir(parents=True, exist_ok=False) + absolute_table.to_csv(args.output_dir / "absolute_position_error_rates.tsv", sep="\t", index=False) + edge_table.to_csv(args.output_dir / "edge_distance_error_rates.tsv", sep="\t", index=False) + region_table.to_csv(args.output_dir / "region_error_rates.tsv", sep="\t", index=False) + save_figure(absolute_table, edge_table, region_table, args.output_dir / "reconstruction_error_by_position") + result = { + "status": "complete", + "analysis": "locked_test_autoregressive_error_by_cdr3_position", + "checkpoint": str(args.checkpoint), + "checkpoint_sha256": sha256(args.checkpoint), + "tokenizer_sha256": sha256(args.tokenizer), + "benchmark_ready_sha256": sha256(args.data_dir / "READY.json"), + "model_config": asdict(config), + "matched_length_sequences": matched_sequences, + "mismatched_length_sequences": mismatched_sequences, + "regions": region_table.to_dict(orient="records"), + "note": "Position analysis includes only target/prediction pairs with equal sequence length.", + } + (args.output_dir / "RESULTS.json").write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_single_aa_challenge.py b/benchmark/evaluate_single_aa_challenge.py new file mode 100644 index 0000000..2cc5e28 --- /dev/null +++ b/benchmark/evaluate_single_aa_challenge.py @@ -0,0 +1,339 @@ +"""Frozen-encoder delta benchmark on controlled one-AA substitutions.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy.stats import pearsonr, spearmanr +from sklearn.decomposition import PCA +from sklearn.metrics import balanced_accuracy_score, f1_score, mean_absolute_error, mean_squared_error, r2_score + + +ALPHAS = (0.01, 0.1, 1.0, 10.0, 100.0) +AA = "ACDEFGHIKLMNPQRSTVWY" + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def ridge_fit(x: np.ndarray, y: np.ndarray, alpha: float) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + mean_x, std_x = x.mean(0), x.std(0) + std_x[std_x < 1e-8] = 1.0 + z = (x - mean_x) / std_x + mean_y = y.mean(0) + centered = y - mean_y + coef = np.linalg.solve(z.T @ z + alpha * np.eye(z.shape[1]), z.T @ centered) + return coef, mean_x, std_x, mean_y + + +def ridge_predict(model: tuple[np.ndarray, ...], x: np.ndarray) -> np.ndarray: + coef, mean_x, std_x, mean_y = model + return (x - mean_x) / std_x @ coef + mean_y + + +def finite_correlation(x: np.ndarray, y: np.ndarray, kind: str) -> float: + if len(x) < 2 or np.std(x) < 1e-12 or np.std(y) < 1e-12: + return 0.0 + return float((spearmanr if kind == "spearman" else pearsonr)(x, y).statistic) + + +def choose_alpha_regression(x: np.ndarray, y: np.ndarray, train: np.ndarray, val: np.ndarray) -> float: + scores = [] + scale = y[train].std(0) + scale = np.where(scale < 1e-8, 1.0, scale) + for alpha in ALPHAS: + pred = ridge_predict(ridge_fit(x[train], y[train], alpha), x[val]) + scores.append((float(np.mean(((pred - y[val]) / scale) ** 2)), alpha)) + return min(scores)[1] + + +def choose_alpha_classifier(x: np.ndarray, labels: np.ndarray, train: np.ndarray, val: np.ndarray) -> tuple[float, np.ndarray]: + # Keep the complete declared label space. A test-only rare position remains + # unlearnable (correctly scoring as an error) rather than crashing encoding. + classes = np.array(sorted(set(labels))) + mapping = {value: i for i, value in enumerate(classes)} + onehot = np.eye(len(classes))[np.array([mapping[value] for value in labels])] + scores = [] + for alpha in ALPHAS: + pred = classes[np.argmax(ridge_predict(ridge_fit(x[train], onehot[train], alpha), x[val]), axis=1)] + scores.append((-f1_score(labels[val], pred, average="macro", zero_division=0), alpha)) + return min(scores)[1], classes + + +def regression_metrics(truth: np.ndarray, pred: np.ndarray, train_truth: np.ndarray) -> dict: + scale = float(np.std(train_truth)) or 1.0 + return { + "rmse": float(np.sqrt(mean_squared_error(truth, pred))), + "nrmse_train_sd": float(np.sqrt(mean_squared_error(truth, pred)) / scale), + "mae": float(mean_absolute_error(truth, pred)), + "r2": float(r2_score(truth, pred)), + "pearson": finite_correlation(truth, pred, "pearson"), + "spearman": finite_correlation(truth, pred, "spearman"), + } + + +def tcremp_metrics(truth: np.ndarray, pred: np.ndarray, train_truth: np.ndarray) -> dict: + denom = np.linalg.norm(truth, axis=1) * np.linalg.norm(pred, axis=1) + cosine = np.divide(np.sum(truth * pred, axis=1), denom, out=np.zeros(len(truth)), where=denom > 1e-12) + scale = train_truth.std(0) + active = scale >= 1e-8 + return { + "mean_delta_cosine": float(cosine.mean()), + "median_delta_cosine": float(np.median(cosine)), + "standardized_rmse": float(np.sqrt(np.mean(((truth[:, active] - pred[:, active]) / scale[active]) ** 2))), + "raw_rmse": float(np.sqrt(np.mean((truth - pred) ** 2))), + "active_dimensions": int(active.sum()), + } + + +def rank_summary(frame: pd.DataFrame, true_pgen: np.ndarray, pred_pgen: np.ndarray, + true_t: np.ndarray, pred_t: np.ndarray, mask: np.ndarray) -> dict: + rows = np.flatnonzero(mask) + pgen_r, tcremp_r, pgen_pairs, tcremp_pairs = [], [], [], [] + for parent in sorted(frame.iloc[rows]["parent_id"].unique()): + idx = rows[frame.iloc[rows]["parent_id"].to_numpy() == parent] + if len(idx) < 2: + continue + true_te = np.linalg.norm(true_t[idx], axis=1) + pred_te = np.linalg.norm(pred_t[idx], axis=1) + pgen_r.append(finite_correlation(true_pgen[idx], pred_pgen[idx], "spearman")) + tcremp_r.append(finite_correlation(true_te, pred_te, "spearman")) + for a in range(len(idx)): + for b in range(a + 1, len(idx)): + pgen_pairs.append(np.sign(true_pgen[idx[a]] - true_pgen[idx[b]]) == np.sign(pred_pgen[idx[a]] - pred_pgen[idx[b]])) + tcremp_pairs.append(np.sign(true_te[a] - true_te[b]) == np.sign(pred_te[a] - pred_te[b])) + return { + "parent_mean_pgen_spearman": float(np.mean(pgen_r)), + "parent_mean_tcremp_effect_spearman": float(np.mean(tcremp_r)), + "pgen_pair_order_accuracy": float(np.mean(pgen_pairs)), + "tcremp_effect_pair_order_accuracy": float(np.mean(tcremp_pairs)), + "parents": len(pgen_r), + } + + +def identity_features(frame: pd.DataFrame, include_position: bool) -> np.ndarray: + result = np.zeros((len(frame), 20 + (40 if include_position else 0) + (40 if include_position else 0)), np.float64) + lookup = {aa: i for i, aa in enumerate(AA)} + for i, row in enumerate(frame.itertuples(index=False)): + result[i, lookup[row.to_aa]] += 1 + result[i, lookup[row.from_aa]] -= 1 + if include_position: + result[i, 20 + min(int(row.mutation_position_zero_based), 39)] = 1 + result[i, 60 + min(int(row.length), 39)] = 1 + return result + + +def evaluate_representation(name: str, x: np.ndarray, frame: pd.DataFrame, delta_pgen: np.ndarray, + delta_t: np.ndarray, splits: dict[str, np.ndarray]) -> tuple[dict, dict[str, np.ndarray]]: + train, val, test = (splits[key] for key in ("train", "val", "test")) + p_alpha = choose_alpha_regression(x, delta_pgen[:, None], train, val) + p_model = ridge_fit(x[train], delta_pgen[train, None], p_alpha) + p_pred = ridge_predict(p_model, x).ravel() + t_alpha = choose_alpha_regression(x, delta_t, train, val) + t_model = ridge_fit(x[train], delta_t[train], t_alpha) + t_pred = ridge_predict(t_model, x) + + classifiers = {} + class_predictions = {} + for label_name, labels in { + "mutation_position": frame["mutation_position_zero_based"].to_numpy(int), + "target_amino_acid": frame["to_aa"].astype(str).to_numpy(), + }.items(): + alpha, classes = choose_alpha_classifier(x, labels, train, val) + class_to_index = {value: i for i, value in enumerate(classes)} + onehot = np.eye(len(classes))[np.array([class_to_index[value] for value in labels])] + pred = classes[np.argmax(ridge_predict(ridge_fit(x[train], onehot[train], alpha), x), axis=1)] + classifiers[label_name] = { + "alpha": alpha, + "macro_f1": float(f1_score(labels[test], pred[test], average="macro", zero_division=0)), + "balanced_accuracy": float(balanced_accuracy_score(labels[test], pred[test])), + } + class_predictions[label_name] = pred + direction_truth = (delta_pgen > 0).astype(int) + direction_alpha, direction_classes = choose_alpha_classifier(x, direction_truth, train, val) + direction_onehot = np.eye(len(direction_classes))[direction_truth] + direction_pred = direction_classes[np.argmax(ridge_predict( + ridge_fit(x[train], direction_onehot[train], direction_alpha), x), axis=1)] + metrics = { + "representation": name, + "dimension": int(x.shape[1]), + "pgen": {"alpha": p_alpha, **regression_metrics(delta_pgen[test], p_pred[test], delta_pgen[train])}, + "tcremp": {"alpha": t_alpha, **tcremp_metrics(delta_t[test], t_pred[test], delta_t[train])}, + "mutation_identity": classifiers, + "pgen_direction": { + "alpha": direction_alpha, + "macro_f1": float(f1_score(direction_truth[test], direction_pred[test], average="macro", zero_division=0)), + "balanced_accuracy": float(balanced_accuracy_score(direction_truth[test], direction_pred[test])), + }, + "within_parent_ranking": rank_summary(frame, delta_pgen, p_pred, delta_t, t_pred, test), + } + predictions = {"pgen": p_pred, "tcremp": t_pred, + "position": class_predictions["mutation_position"], + "to_aa": class_predictions["target_amino_acid"]} + return metrics, predictions + + +def parent_bootstrap(frame: pd.DataFrame, truth_p: np.ndarray, truth_t: np.ndarray, + predictions: dict[str, dict[str, np.ndarray]], test: np.ndarray, + repeats: int, seed: int) -> dict: + parent_rows = {parent: np.flatnonzero((frame["parent_id"].to_numpy() == parent) & test) + for parent in sorted(frame.loc[test, "parent_id"].unique())} + parents = np.array(list(parent_rows), dtype=object) + rng = np.random.default_rng(seed) + result = {} + for comparator in sorted(set(predictions) - {"rtp"}): + samples = {"pgen_nrmse": [], "tcremp_one_minus_cosine": []} + for _ in range(repeats): + chosen = rng.choice(parents, size=len(parents), replace=True) + idx = np.concatenate([parent_rows[parent] for parent in chosen]) + scale = np.std(truth_p[np.flatnonzero(frame["probe_split"].to_numpy() == "train")]) or 1.0 + def score(name: str) -> tuple[float, float]: + p = float(np.sqrt(np.mean((predictions[name]["pgen"][idx] - truth_p[idx]) ** 2)) / scale) + pred, true = predictions[name]["tcremp"][idx], truth_t[idx] + denom = np.linalg.norm(pred, axis=1) * np.linalg.norm(true, axis=1) + cos = np.divide(np.sum(pred * true, axis=1), denom, out=np.zeros(len(idx)), where=denom > 1e-12) + return p, float(1 - cos.mean()) + rtp, other = score("rtp"), score(comparator) + samples["pgen_nrmse"].append(rtp[0] - other[0]) + samples["tcremp_one_minus_cosine"].append(rtp[1] - other[1]) + result[comparator] = { + metric: {"rtp_minus_comparator_mean": float(np.mean(values)), + "ci95": [float(x) for x in np.quantile(values, [0.025, 0.975])]} + for metric, values in samples.items() + } + return result + + +def bootstrap_or_disabled(frame: pd.DataFrame, truth_p: np.ndarray, truth_t: np.ndarray, + predictions: dict[str, dict[str, np.ndarray]], test: np.ndarray, + repeats: int, seed: int) -> dict: + if repeats < 0: + raise ValueError("bootstrap repeats must be non-negative") + if repeats == 0: + return { + "status": "disabled_by_user", + "repeats": 0, + "uncertainty_available": False, + "significance_claim_permitted": False, + } + return { + "status": "enabled", + "repeats": repeats, + "uncertainty_available": True, + "significance_claim_permitted": True, + "comparisons": parent_bootstrap(frame, truth_p, truth_t, predictions, test, repeats, seed), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--prepared", type=Path, required=True) + parser.add_argument("--targets", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--model", action="append", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--bootstrap", type=int, default=0) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + parents = pd.read_csv(args.prepared / "parents.tsv", sep="\t") + frame = pd.read_csv(args.targets / "targets.tsv", sep="\t") + target_manifest = json.loads((args.targets / "TARGETS.json").read_text()) + if target_manifest["status"] != "accepted_for_frozen_delta_evaluation": + raise ValueError("Target gate did not accept this challenge.") + parent_t = np.load(args.prepared / "parent_tcremp.npy", mmap_mode="r") + mutant_t = np.load(args.targets / "mutant_tcremp.npy", mmap_mode="r") + parent_row = {identifier: i for i, identifier in enumerate(parents["parent_id"].astype(str))} + pidx = np.array([parent_row[value] for value in frame["parent_id"].astype(str)]) + delta_t = np.asarray(mutant_t, np.float64) - np.asarray(parent_t[pidx], np.float64) + parent_pgen = parents.set_index("parent_id")["log10_pgen_1mm"] + delta_pgen = frame["log10_pgen_1mm"].to_numpy(float) - parent_pgen.loc[frame["parent_id"]].to_numpy(float) + if not np.isfinite(delta_t).all() or not np.isfinite(delta_pgen).all(): + raise ValueError("Nonfinite delta targets.") + + from benchmark.evaluate_vdjdb_frozen import extract_embeddings, parse_models, validate_run_configs + models = parse_models(args.model) + config, parity = validate_run_configs(models) + sequences = pd.DataFrame({"cdr3": pd.concat((parents["junction_aa"], frame["junction_aa"]), ignore_index=True)}) + embedded, extraction = extract_embeddings(sequences, models, config, args.tokenizer, + args.batch_size, args.output_dir) + n_parent = len(parents) + representations = {name: array[n_parent:] - array[:n_parent][pidx] for name, array in embedded.items()} + representations["concat_rtp"] = np.concatenate([representations[name] for name in ("r", "t", "p")], axis=1) + split_masks = {name: frame["probe_split"].to_numpy() == name for name in ("train", "val", "test")} + pca = PCA(n_components=128, random_state=args.seed).fit(representations["concat_rtp"][split_masks["train"]]) + representations["pca_concat_rtp_128"] = pca.transform(representations["concat_rtp"]) + representations["delta_composition_control"] = identity_features(frame, False) + representations["mutation_identity_control"] = identity_features(frame, True) + + metrics, predictions = {}, {} + for name, array in representations.items(): + metrics[name], predictions[name] = evaluate_representation( + name, np.asarray(array, np.float64), frame, delta_pgen, delta_t, split_masks + ) + seven = ["r", "t", "p", "rt", "rp", "tp", "rtp"] + errors = { + name: { + "sequence": float(1 - np.mean([metrics[name]["mutation_identity"][key]["macro_f1"] + for key in ("mutation_position", "target_amino_acid")])), + "tcremp": float(1 - metrics[name]["tcremp"]["mean_delta_cosine"]), + "pgen": float(metrics[name]["pgen"]["nrmse_train_sd"]), + } for name in seven + } + best = {task: min(errors[name][task] for name in seven) for task in ("sequence", "tcremp", "pgen")} + aggregate = {} + for name in seven: + regrets = {task: (errors[name][task] - best[task]) / max(best[task], 1e-8) for task in best} + aggregate[name] = {"task_errors": errors[name], "relative_regrets": regrets, + "mean_relative_regret": float(np.mean(list(regrets.values()))), + "worst_task_relative_regret": float(max(regrets.values()))} + bootstrap = bootstrap_or_disabled(frame, delta_pgen, delta_t, predictions, split_masks["test"], + args.bootstrap, args.seed) + payload = { + "status": "complete", + "study": "parent-disjoint controlled one-amino-acid frozen-delta benchmark", + "all_pairs_edit_distance": 1, + "probe_split_unit": "parent clonotype", + "rows": len(frame), + "parents": int(frame["parent_id"].nunique()), + "split_parent_counts": { + str(key): int(value) + for key, value in frame.groupby("probe_split")["parent_id"].nunique().items() + }, + "model_parity": parity, + "extraction": extraction, + "metrics": metrics, + "preregistered_aggregate": {"definition": "relative error regret against best of seven per task; lower is better", + "best_task_errors": best, "models": aggregate}, + "parent_clustered_bootstrap": bootstrap, + "controls": ["delta_composition_control", "mutation_identity_control"], + "ensemble_baselines": ["concat_rtp", "pca_concat_rtp_128"], + "limitations": [ + "fixed seed42 encoders", + "linear probes only", + "mutants are in-silico and are not observed repertoire clonotypes", + "bootstrap uncertainty was disabled by user direction; point estimates only and no significance claim", + ], + } + result_path = args.output_dir / "RESULTS.json" + result_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + hashes = {path.name: sha256(path) for path in args.output_dir.iterdir() if path.is_file()} + (args.output_dir / "HASHES.json").write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n") + print(json.dumps({"status": "complete", "results": str(result_path)}, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_vdjdb_author_global.py b/benchmark/evaluate_vdjdb_author_global.py new file mode 100644 index 0000000..326c85a --- /dev/null +++ b/benchmark/evaluate_vdjdb_author_global.py @@ -0,0 +1,162 @@ +"""Evaluate author-native frozen global embeddings on the fixed REDCEA gallery.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd + +from benchmark.evaluate_vdjdb_redcea_clonotypes import ( + METRIC_COLUMNS, + candidate_distances, + summarize_candidates, +) + +EXTERNAL = ("esm2_35m", "tcr_bert", "sceptr_cdr3") +MODALITY = { + "esm2_35m": "sequence-only", + "tcr_bert": "sequence-only", + "sceptr_cdr3": "sequence-only", + "sceptr": "annotation-aware V+CDR3+J", +} + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--preflight", type=Path, required=True) + parser.add_argument("--representations-dir", type=Path, required=True) + parser.add_argument("--reference-evaluation", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.output_dir.exists(): + raise FileExistsError(args.output_dir) + args.output_dir.mkdir(parents=True) + gate = json.loads(args.preflight.read_text(encoding="utf-8")) + if gate["status"] != "accepted": + raise ValueError("Author-global preflight did not accept.") + + cohort = pd.read_csv(args.cohort, sep="\t") + if sha256(args.cohort) != gate["input_hashes"]["redcea_cohort"]: + raise ValueError("Cohort hash differs from accepted preflight.") + gallery_file = args.reference_evaluation / "candidate_gallery.npz" + gallery = np.load(gallery_file) + candidates = gallery["candidates"] + relevant = gallery["relevant"] + if candidates.shape[0] != len(cohort): + raise ValueError("Candidate gallery row count mismatch.") + + labels_text = cohort["label"].astype(str) + label_names = sorted(labels_text.unique()) + label_to_id = {label: index for index, label in enumerate(label_names)} + label_ids = np.array([label_to_id[value] for value in labels_text], dtype=np.int32) + sequence_sha = hashlib.sha256( + "\n".join(cohort["cdr3"].astype(str)).encode("utf-8") + ).hexdigest() + catalog = json.loads( + (args.representations_dir / "representations.json").read_text(encoding="utf-8") + ) + + summaries = {} + per_epitope = [] + rows = [] + for name in EXTERNAL: + entry = catalog[name] + if entry["rows"] != len(cohort) or entry["sequence_sha256"] != sequence_sha: + raise ValueError(f"{name} catalog provenance mismatch") + source = Path(entry["path"]) + if sha256(source) != entry["sha256"]: + raise ValueError(f"{name} representation hash mismatch") + matrix = np.load(source, mmap_mode="r") + if matrix.shape != (len(cohort), entry["dim"]) or not np.isfinite(matrix).all(): + raise ValueError(f"{name} shape/finiteness mismatch: {matrix.shape}") + summaries[name] = {} + row = {"model": name, "modality": MODALITY[name], "native_dim": int(matrix.shape[1])} + for distance in ("cosine", "euclidean"): + distances, norms = candidate_distances(matrix, candidates, distance) + summary, _, per_label = summarize_candidates( + distances, candidates, relevant, label_ids, label_names + ) + summary["embedding_norms"] = norms + summaries[name][distance] = summary + per_label.insert(0, "model", name) + per_label.insert(1, "modality", MODALITY[name]) + per_label.insert(2, "distance", distance) + per_epitope.append(per_label) + for metric in METRIC_COLUMNS: + row[f"{distance}_{metric}"] = summary["macro_epitope"][metric] + for k in (1, 5, 10): + row[f"{distance}_knn_macro_f1_at_{k}"] = ( + summary["knn_classification"][f"k_{k}"]["macro_f1"] + ) + rows.append(row) + + pd.DataFrame(rows).to_csv( + args.output_dir / "external_retrieval_summary.tsv", + sep="\t", index=False, lineterminator="\n", + ) + pd.concat(per_epitope, ignore_index=True).to_csv( + args.output_dir / "external_per_epitope.tsv", + sep="\t", index=False, lineterminator="\n", + ) + (args.output_dir / "external_metrics.json").write_text( + json.dumps(summaries, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + result = { + "status": "complete", + "cohort": { + "sha256": sha256(args.cohort), + "clonotypes": len(cohort), + "epitopes": len(label_names), + }, + "candidate_gallery": { + "path": str(gallery_file), + "sha256": sha256(gallery_file), + "shape": list(candidates.shape), + }, + "representations": { + name: { + "sha256": catalog[name]["sha256"], + "dim": catalog[name]["dim"], + "modality": MODALITY[name], + "model": catalog[name]["model"], + } + for name in EXTERNAL + }, + "selection": "No VDJdb labels were used to fit, pool, tune, or select an encoder.", + "unavailable": { + "sceptr": ( + "Annotation-aware SCEPTR-default rejects unsupported/nonfunctional " + "VDJdb V/J alleles; the fixed cohort was not filtered per method." + ) + }, + "uncertainty": "Per-epitope distributions retained; no bootstrap requested.", + "circularity_caveat": ( + "REDCEA membership is sequence/TCRemP-cluster-conditioned; this is " + "motif-conditioned retrieval geometry, not independent binding evidence." + ), + } + (args.output_dir / "RESULTS.json").write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + print(json.dumps(result, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_vdjdb_decoded_tcremp.py b/benchmark/evaluate_vdjdb_decoded_tcremp.py new file mode 100644 index 0000000..8191388 --- /dev/null +++ b/benchmark/evaluate_vdjdb_decoded_tcremp.py @@ -0,0 +1,196 @@ +"""Decode saved VDJdb RTP latents through the frozen TCRemP head. + +This intentionally reuses the frozen uniform-100 VDJdb cohort, candidate pairs, +and ``separation_metrics`` implementation. It does not encode sequences or +train a model: its sole intervention is ``saved RTP -> frozen T head``. +""" +from __future__ import annotations + +import argparse +import hashlib +import json +from dataclasses import asdict +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import torch +import umap + +from evaluate_vdjdb_uniform100_separation import separation_metrics +from irrm_codec.multitask_data import TargetStandardizer +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer + + +SEED = 1729 +HIGHLIGHTS = { + "SSYRRPVGI": "#4477AA", "STPESANL": "#EE6677", "SLLMWITQV": "#228833", + "KYNKANVFL": "#AA3377", "SIINFEKL": "#CCBB44", +} + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--cohort", type=Path, required=True) + p.add_argument("--rtp", type=Path, required=True, + help="Existing full-cohort RTP matrix used by the VDJdb analysis.") + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--standardizer", type=Path, required=True) + p.add_argument("--uniform-evaluation", type=Path, required=True, + help="Directory containing the locked uniform_pairs.npz.") + p.add_argument("--existing-separation", type=Path, required=True, + help="Completed RP/RTP separation directory; values are reused, not recomputed.") + p.add_argument("--native-tcremp", type=Path, + help="Optional existing aligned native TCRemP matrix for reconstruction diagnostics.") + p.add_argument("--output-dir", type=Path, required=True) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--device", choices=("cpu", "cuda"), default="cpu") + return p.parse_args() + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def decode(a: argparse.Namespace, rtp: np.ndarray, out: Path) -> tuple[Path, dict]: + payload = torch.load(a.checkpoint, map_location="cpu", weights_only=False) + config = IRRMCodecConfig(**payload["model_config"]) + if rtp.ndim != 2 or rtp.shape[1] != config.latent_dim: + raise ValueError(f"RTP shape {rtp.shape} is incompatible with latent_dim={config.latent_dim}") + standardizer = TargetStandardizer.load(a.standardizer) + if len(standardizer.tcremp_mean) != config.tcremp_dim: + raise ValueError("TCRemP standardizer and checkpoint output dimensions differ") + device = torch.device(a.device) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("--device cuda requested but CUDA is unavailable") + model = IRRMCodecTransformer(config).to(device) + model.load_state_dict(payload["model_state"], strict=True) + model.eval() + mean = torch.from_numpy(standardizer.tcremp_mean).to(device) + std = torch.from_numpy(standardizer.tcremp_std).to(device) + decoded_path = out / "decoded_tcremp.npy" + decoded = np.lib.format.open_memmap(decoded_path, mode="w+", dtype=np.float32, + shape=(len(rtp), config.tcremp_dim)) + with torch.inference_mode(): + for start in range(0, len(rtp), a.batch_size): + stop = min(start + a.batch_size, len(rtp)) + z = torch.from_numpy(np.asarray(rtp[start:stop], dtype=np.float32)).to(device) + decoded[start:stop] = (model.tcremp_head(z) * std + mean).cpu().numpy() + decoded.flush() + return decoded_path, { + "checkpoint_epoch_zero_based": int(payload["epoch"]), + "checkpoint_best_validation_loss": float(payload["best_val_loss"]), + "model_config": asdict(config), "standardizer_train_rows": standardizer.train_rows, + } + + +def diagnostics(decoded: np.ndarray, native: np.ndarray | None) -> dict: + result = {"decoded": {"shape": [int(x) for x in decoded.shape], + "nan_count": int(np.isnan(decoded).sum()), "inf_count": int(np.isinf(decoded).sum()), + "min": float(np.min(decoded)), "max": float(np.max(decoded)), + "mean": float(np.mean(decoded)), "std": float(np.std(decoded))}} + if native is not None: + if native.shape != decoded.shape: + raise ValueError(f"Native TCRemP shape {native.shape} does not match decoded {decoded.shape}") + # Chunked global sums avoid a second multi-GB temporary array. + cosine, squared, dot, sx, sy, sxx, syy, n = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0 + for start in range(0, len(decoded), 128): + x, y = np.asarray(decoded[start:start+128], np.float64), np.asarray(native[start:start+128], np.float64) + cosine += float(np.sum(np.sum(x*y, axis=1) / (np.linalg.norm(x, axis=1)*np.linalg.norm(y, axis=1)))) + squared += float(np.square(x-y).sum()); dot += float((x*y).sum()); sx += float(x.sum()); sy += float(y.sum()) + sxx += float(np.square(x).sum()); syy += float(np.square(y).sum()); n += x.size + pearson = (n*dot-sx*sy) / np.sqrt((n*sxx-sx*sx)*(n*syy-sy*sy)) + result["native_tcremp"] = {"available": True, "global_mean_cosine": cosine / len(decoded), + "global_pearson": float(pearson), "global_mse": squared / n} + else: + result["native_tcremp"] = {"available": False} + return result + + +def plot_umap(matrix: np.ndarray, labels: np.ndarray, out: Path) -> None: + coords = umap.UMAP(n_neighbors=30, min_dist=.15, metric="cosine", n_components=2, + random_state=SEED, transform_seed=SEED).fit_transform(np.asarray(matrix, np.float32)) + frame = pd.DataFrame({"index": np.arange(len(labels)), "epitope": labels, + "umap_1": coords[:, 0], "umap_2": coords[:, 1]}) + frame.to_parquet(out / "umap_coordinates.parquet", index=False) + plt.rcParams.update({"font.family": "STIXGeneral", "font.size": 7.5, "pdf.fonttype": 42, "svg.fonttype": "none"}) + fig, ax = plt.subplots(figsize=(3.38, 3.38)) + other = ~np.isin(labels, list(HIGHLIGHTS)) + ax.scatter(coords[other, 0], coords[other, 1], s=1.1, c="#9e9e9e", alpha=.20, linewidths=0, rasterized=True) + for epitope, color in HIGHLIGHTS.items(): + keep = labels == epitope + ax.scatter(coords[keep, 0], coords[keep, 1], s=8, c=color, alpha=.76, linewidths=0, rasterized=True, label=epitope) + ax.set(xticks=[], yticks=[], title="Decoded TCRemP-like representation") + for spine in ax.spines.values(): spine.set_visible(False) + ax.legend(loc="lower center", bbox_to_anchor=(.5, -.16), ncol=3, frameon=False, fontsize=6.8) + fig.subplots_adjust(left=.012, right=.995, top=.89, bottom=.14) + for suffix, kwargs in ((".png", {"dpi": 600}), (".pdf", {}), (".svg", {})): + fig.savefig((out / "vdjdb_decoded_tcremp_umap").with_suffix(suffix), bbox_inches="tight", pad_inches=.01, **kwargs) + plt.close(fig) + + +def plot_effects(table: pd.DataFrame, out: Path) -> None: + long = table.melt(id_vars=["epitope", "N"], value_vars=["d_RP", "d_RTP", "d_decoded_TCRemP"], + var_name="representation", value_name="cohen_d") + order = table.sort_values("d_RTP", ascending=False).epitope + wide = long.pivot(index="representation", columns="epitope", values="cohen_d")[order] + fig, ax = plt.subplots(figsize=(11, 2.2)) + image = ax.imshow(wide, aspect="auto", cmap="viridis") + ax.set(yticks=range(len(wide.index)), yticklabels=wide.index, xticks=range(len(order)), xticklabels=order) + ax.tick_params(axis="x", labelrotation=90, labelsize=5); ax.tick_params(axis="y", labelsize=7) + fig.colorbar(image, ax=ax, label="Cohen's d", shrink=.8) + fig.tight_layout() + for suffix in (".png", ".pdf", ".svg"): fig.savefig((out / "cohen_d_comparison").with_suffix(suffix), dpi=600 if suffix==".png" else None) + plt.close(fig) + + +def main() -> None: + a = parse_args() + if a.output_dir.exists(): raise FileExistsError(f"Refusing to overwrite {a.output_dir}") + a.output_dir.mkdir(parents=True) + cohort = pd.read_csv(a.cohort, sep="\t") + rtp = np.load(a.rtp, mmap_mode="r") + if len(rtp) != len(cohort): raise ValueError("Saved RTP row count does not match cohort") + cohort.insert(0, "cohort_index", np.arange(len(cohort))) + cohort.to_csv(a.output_dir / "receptor_metadata.tsv", sep="\t", index=False) + decoded_path, decoder_meta = decode(a, rtp, a.output_dir) + decoded = np.load(decoded_path, mmap_mode="r") + native = np.load(a.native_tcremp, mmap_mode="r") if a.native_tcremp else None + check = diagnostics(decoded, native) + check["ordering"] = {"cohort_rows": len(cohort), "rtp_rows": len(rtp), "decoded_rows": len(decoded), + "exact_row_order_preserved": True, "identifier_columns": [c for c in ("cid", "cdr3") if c in cohort]} + (a.output_dir / "sanity_checks.json").write_text(json.dumps(check, indent=2, sort_keys=True)+"\n") + pairs = np.load(a.uniform_evaluation / "uniform_pairs.npz") + original, within, between = pairs["original_indices"], pairs["within"], pairs["between"] + labels = cohort.iloc[original].label.astype(str).to_numpy() + query, decoded_epitope = separation_metrics(decoded[original], within, between, labels) + query.insert(0, "model", "decoded_TCRemP") + query.to_parquet(a.output_dir / "per_query_decoded_tcremp.parquet", index=False) + old = pd.read_csv(a.existing_separation / "per_epitope.tsv", sep="\t") + old = old[old.model.isin(["RP", "RTP"])].pivot(index=["epitope", "queries"], columns="model", values="cohen_d").reset_index() + result = old.merge(decoded_epitope, left_on=["epitope", "queries"], right_on=["epitope", "queries"], validate="one_to_one") + result = result.rename(columns={"queries": "N", "RP": "d_RP", "RTP": "d_RTP", "cohen_d": "d_decoded_TCRemP"}) + result["decoded_minus_RP"] = result.d_decoded_TCRemP-result.d_RP + result["decoded_minus_RTP"] = result.d_decoded_TCRemP-result.d_RTP + result.sort_values("epitope").to_csv(a.output_dir / "per_epitope_cohen_d.tsv", sep="\t", index=False) + summary = {"median_cohen_d": {k: float(result[k].median()) for k in ("d_RP", "d_RTP", "d_decoded_TCRemP")}, + "fraction_decoded_gt_RP": float((result.d_decoded_TCRemP > result.d_RP).mean()), + "fraction_decoded_gt_RTP": float((result.d_decoded_TCRemP > result.d_RTP).mean()), + "pearson_d_RTP_vs_decoded": float(result.d_RTP.corr(result.d_decoded_TCRemP)), + "top_decoded": result.nlargest(10, "d_decoded_TCRemP")[["epitope", "N", "d_decoded_TCRemP"]].to_dict("records"), + "top_gains_vs_RP": result.nlargest(10, "decoded_minus_RP")[["epitope", "decoded_minus_RP"]].to_dict("records"), + "largest_losses_vs_RTP": result.nsmallest(10, "decoded_minus_RTP")[["epitope", "decoded_minus_RTP"]].to_dict("records")} + (a.output_dir / "aggregate_statistics.json").write_text(json.dumps(summary, indent=2, sort_keys=True)+"\n") + plot_effects(result, a.output_dir); plot_umap(decoded, cohort.label.astype(str).to_numpy(), a.output_dir) + manifest = {"status": "complete", "decoder": decoder_meta, "inputs": {name: {"path": str(path), "sha256": sha256(path)} for name, path in {"cohort":a.cohort,"rtp":a.rtp,"checkpoint":a.checkpoint,"standardizer":a.standardizer,"pairs":a.uniform_evaluation/"uniform_pairs.npz"}.items()}, "protocol": "Saved RTP latents decoded only through frozen TCRemP head; Cohen's d is imported from evaluate_vdjdb_uniform100_separation.separation_metrics without modification."} + (a.output_dir / "RESULTS.json").write_text(json.dumps(manifest, indent=2, sort_keys=True)+"\n") + + +if __name__ == "__main__": main() diff --git a/benchmark/evaluate_vdjdb_epitope_prediction.py b/benchmark/evaluate_vdjdb_epitope_prediction.py new file mode 100644 index 0000000..bea57b4 --- /dev/null +++ b/benchmark/evaluate_vdjdb_epitope_prediction.py @@ -0,0 +1,627 @@ +"""Frozen-embedding multiclass epitope prediction on the manuscript VDJdb cohort. + +The selected rows are exactly ``original_indices`` from the published uniform-100 +separation evaluation. No encoder is loaded or fine-tuned here. Every dense +representation uses the same cached folds, training-only scaling, and the same +small, training-only C selection procedure. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import warnings +from dataclasses import dataclass +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from joblib import Parallel, delayed +from scipy import sparse +from scipy.stats import wilcoxon +from sklearn.exceptions import ConvergenceWarning +from sklearn.feature_extraction.text import TfidfVectorizer +from sklearn.linear_model import LogisticRegression +from sklearn.metrics import ( + accuracy_score, + balanced_accuracy_score, + f1_score, + roc_auc_score, + top_k_accuracy_score, +) +from sklearn.model_selection import StratifiedGroupKFold, StratifiedKFold, StratifiedShuffleSplit +from sklearn.preprocessing import StandardScaler + + +NATIVE_NAMES = ("r", "t", "p", "rt", "rp", "tp", "rtp") +EXTERNAL_FILES = { + "esm2_35m": "esm2_35m.npy", + "tcr_bert": "tcr_bert.npy", + "sceptr_cdr3": "sceptr_cdr3.npy", +} +DISPLAY = { + "r": "R", "t": "T", "p": "P", "rt": "RT", "rp": "RP", "tp": "TP", + "rtp": "RTP", "esm2_35m": "ESM2-35M", "tcr_bert": "TCR-BERT", + "sceptr_cdr3": "SCEPTR (CDR3-only)", "kmer_tfidf": "3-mer TF-IDF", +} +PRIMARY_ORDER = ( + "rtp", "rp", "r", "t", "p", "rt", "tp", "tcr_bert", "sceptr_cdr3", + "esm2_35m", "kmer_tfidf", +) + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def stable_json(path: Path, value: object) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--uniform-evaluation", type=Path, required=True) + parser.add_argument("--native-evaluation", type=Path, required=True) + parser.add_argument("--external-representations", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--seeds", type=int, nargs="+", default=[42, 43, 44]) + parser.add_argument("--folds", type=int, default=5) + parser.add_argument("--c-grid", type=float, nargs="+", default=[0.1, 1.0]) + parser.add_argument("--max-iter", type=int, default=400) + parser.add_argument("--jobs", type=int, default=1) + parser.add_argument("--representations", nargs="+", default=list(PRIMARY_ORDER)) + parser.add_argument("--skip-similarity-grouped", action="store_true") + return parser.parse_args() + + +def selected_cohort(cohort_path: Path, uniform_dir: Path) -> tuple[pd.DataFrame, np.ndarray]: + cohort = pd.read_csv(cohort_path, sep="\t") + pairs = np.load(uniform_dir / "uniform_pairs.npz") + original = np.asarray(pairs["original_indices"], dtype=np.int64) + if len(original) != 65_756 or len(np.unique(original)) != len(original): + raise ValueError("Uniform evaluation does not contain the locked 65,756 unique rows.") + selected = cohort.iloc[original].copy().reset_index(drop=True) + selected.insert(0, "prediction_row_index", np.arange(len(selected), dtype=np.int64)) + selected.insert(1, "source_cohort_row_index", original) + if selected["label"].nunique() != 40 or selected.groupby("label").size().min() < 101: + raise ValueError("Selected cohort differs from the locked 40-epitope support gate.") + return selected, original + + +def audit_dataset(frame: pd.DataFrame) -> tuple[dict, pd.DataFrame]: + class_sizes = frame.groupby("label", sort=True).size().rename("sequences").reset_index() + duplicate_rows = int(frame["cdr3"].duplicated(keep=False).sum()) + labels_per_sequence = frame.groupby("cdr3")["label"].nunique() + conflicting = int((labels_per_sequence > 1).sum()) + sizes = class_sizes["sequences"].to_numpy() + report = { + "sequences": int(len(frame)), + "epitopes": int(len(class_sizes)), + "class_size_min": int(sizes.min()), + "class_size_median": float(np.median(sizes)), + "class_size_max": int(sizes.max()), + "duplicate_cdr3_rows": duplicate_rows, + "duplicate_cdr3_sequences_exist": bool(duplicate_rows), + "identical_cdr3_with_different_labels": conflicting, + } + return report, class_sizes + + +class UnionFind: + def __init__(self, n: int): + self.parent = np.arange(n, dtype=np.int64) + + def find(self, value: int) -> int: + while self.parent[value] != value: + self.parent[value] = self.parent[self.parent[value]] + value = int(self.parent[value]) + return value + + def union(self, left: int, right: int) -> None: + a, b = self.find(left), self.find(right) + if a != b: + self.parent[max(a, b)] = min(a, b) + + +def redcea_component_groups(frame: pd.DataFrame) -> np.ndarray: + """Connect rows sharing any retained REDCEA cluster identifier.""" + union = UnionFind(len(frame)) + owner: dict[str, int] = {} + for index, raw in enumerate(frame["source_cids"].astype(str)): + identifiers = [value for value in raw.split("|") if value] + if not identifiers: + identifiers = [str(frame.iloc[index]["cid"])] + for identifier in identifiers: + if identifier in owner: + union.union(index, owner[identifier]) + else: + owner[identifier] = index + roots = np.array([union.find(i) for i in range(len(frame))], dtype=np.int64) + _, groups = np.unique(roots, return_inverse=True) + return groups.astype(np.int64) + + +def _valid_folds(folds: list[tuple[np.ndarray, np.ndarray]], y: np.ndarray) -> bool: + expected = set(np.unique(y).tolist()) + return all(set(np.unique(y[train])) == expected and set(np.unique(y[test])) == expected + for train, test in folds) + + +def make_splits( + y: np.ndarray, exact_groups: np.ndarray, similarity_groups: np.ndarray, + seeds: list[int], requested_folds: int, include_similarity: bool, +) -> tuple[dict[str, list[tuple[int, int, np.ndarray, np.ndarray]]], dict]: + protocols: dict[str, list[tuple[int, int, np.ndarray, np.ndarray]]] = { + "sequence_stratified": [], "exact_cdr3_grouped": [] + } + report: dict[str, object] = {} + for repeat, seed in enumerate(seeds): + splitter = StratifiedKFold(requested_folds, shuffle=True, random_state=seed) + folds = list(splitter.split(np.zeros(len(y)), y)) + for fold, (train, test) in enumerate(folds): + item = (repeat, fold, train.astype(np.int64), test.astype(np.int64)) + protocols["sequence_stratified"].append(item) + protocols["exact_cdr3_grouped"].append(item) + if len(np.unique(exact_groups)) == len(exact_groups): + report["exact_cdr3_grouped"] = { + "alias_of": "sequence_stratified", + "reason": "the locked cohort has one row per unique exact CDR3", + } + else: + raise ValueError("Unexpected duplicate exact CDR3s in the locked deduplicated cohort.") + + if include_similarity: + per_label_groups = pd.DataFrame({"y": y, "group": similarity_groups}).groupby("y")[ + "group" + ].nunique() + maximum = min(requested_folds, int(per_label_groups.min())) + chosen: tuple[int, list[list[tuple[np.ndarray, np.ndarray]]]] | None = None + for n_splits in range(maximum, 1, -1): + repeated = [] + for seed in seeds: + splitter = StratifiedGroupKFold(n_splits, shuffle=True, random_state=seed) + folds = list(splitter.split(np.zeros(len(y)), y, similarity_groups)) + repeated.append(folds) + if all(_valid_folds(folds, y) for folds in repeated): + chosen = n_splits, repeated + break + if chosen is None: + raise ValueError("Could not construct class-complete REDCEA-grouped folds.") + n_splits, repeated = chosen + protocols["redcea_similarity_grouped"] = [] + for repeat, folds in enumerate(repeated): + for fold, (train, test) in enumerate(folds): + if np.intersect1d(similarity_groups[train], similarity_groups[test]).size: + raise AssertionError("REDCEA group leakage detected.") + protocols["redcea_similarity_grouped"].append( + (repeat, fold, train.astype(np.int64), test.astype(np.int64)) + ) + report["redcea_similarity_grouped"] = { + "folds": n_splits, + "groups": int(len(np.unique(similarity_groups))), + "minimum_groups_per_epitope": int(per_label_groups.min()), + "definition": "connected components of rows sharing any retained REDCEA source_cid", + } + return protocols, report + + +def write_split_assignments( + path: Path, protocols: dict[str, list[tuple[int, int, np.ndarray, np.ndarray]]], + frame: pd.DataFrame, +) -> None: + blocks = [] + for protocol, folds in protocols.items(): + by_repeat: dict[int, np.ndarray] = {} + for repeat, fold, _, test in folds: + by_repeat.setdefault(repeat, np.full(len(frame), -1, dtype=np.int16))[test] = fold + for repeat, assignments in by_repeat.items(): + if np.any(assignments < 0): + raise AssertionError(f"Incomplete assignments for {protocol} repeat {repeat}") + blocks.append(pd.DataFrame({ + "protocol": protocol, + "repeat": repeat, + "prediction_row_index": frame["prediction_row_index"], + "source_cohort_row_index": frame["source_cohort_row_index"], + "cdr3": frame["cdr3"], + "label": frame["label"], + "test_fold": assignments, + })) + pd.concat(blocks, ignore_index=True).to_csv(path, index=False, lineterminator="\n") + + +def load_dense_matrix( + name: str, original: np.ndarray, native_dir: Path, external_dir: Path, source_rows: int, +) -> np.ndarray: + path = (native_dir / f"embeddings_{name}.npy") if name in NATIVE_NAMES else ( + external_dir / EXTERNAL_FILES[name] + ) + matrix = np.load(path, mmap_mode="r") + if matrix.ndim != 2 or matrix.shape[0] != source_rows: + raise ValueError(f"{name} shape {matrix.shape} does not align to {source_rows} source rows") + selected = np.asarray(matrix[original], dtype=np.float32) + if not np.isfinite(selected).all(): + raise ValueError(f"{name} contains non-finite values") + return selected + + +def inner_split( + train: np.ndarray, y: np.ndarray, groups: np.ndarray | None, seed: int +) -> tuple[np.ndarray, np.ndarray]: + if groups is None: + split = StratifiedShuffleSplit(n_splits=1, test_size=0.15, random_state=seed) + inner_train, inner_val = next(split.split(np.zeros(len(train)), y[train])) + else: + sub_y, sub_groups = y[train], groups[train] + per_label = pd.DataFrame({"y": sub_y, "g": sub_groups}).groupby("y")["g"].nunique() + n_splits = min(3, int(per_label.min())) + if n_splits < 2: + raise ValueError("Outer training data cannot support group-aware C selection.") + split = StratifiedGroupKFold(n_splits, shuffle=True, random_state=seed) + candidates = list(split.split(np.zeros(len(train)), sub_y, sub_groups)) + valid = [item for item in candidates if _valid_folds([item], sub_y)] + if not valid: + raise ValueError("No class-complete group-aware inner split.") + inner_train, inner_val = valid[0] + return train[inner_train], train[inner_val] + + +def classifier(c_value: float, max_iter: int, seed: int) -> LogisticRegression: + return LogisticRegression( + C=c_value, + penalty="l2", + solver="lbfgs", + class_weight="balanced", + max_iter=max_iter, + tol=1e-4, + random_state=seed, + ) + + +def transform_dense( + matrix: np.ndarray, fit_rows: np.ndarray, *row_sets: np.ndarray +) -> tuple[list[np.ndarray], StandardScaler]: + scaler = StandardScaler().fit(matrix[fit_rows]) + return [scaler.transform(matrix[rows]) for rows in row_sets], scaler + + +@dataclass(frozen=True) +class FoldSpec: + protocol: str + repeat: int + fold: int + train: np.ndarray + test: np.ndarray + seed: int + + +def score_predictions(y_true: np.ndarray, probabilities: np.ndarray, classes: np.ndarray) -> dict: + predicted = classes[np.argmax(probabilities, axis=1)] + return { + "macro_f1": float(f1_score(y_true, predicted, average="macro")), + "balanced_accuracy": float(balanced_accuracy_score(y_true, predicted)), + "accuracy": float(accuracy_score(y_true, predicted)), + "macro_auroc": float(roc_auc_score( + y_true, probabilities, labels=classes, multi_class="ovr", average="macro" + )), + "top3_accuracy": float(top_k_accuracy_score( + y_true, probabilities, k=3, labels=classes + )), + } + + +def evaluate_fold( + spec: FoldSpec, name: str, matrix: np.ndarray | None, sequences: np.ndarray, + y: np.ndarray, classes: np.ndarray, groups: np.ndarray | None, + c_grid: tuple[float, ...], max_iter: int, +) -> tuple[dict, pd.DataFrame]: + tune_train, tune_val = inner_split(spec.train, y, groups, spec.seed + spec.fold * 1009) + best: tuple[float, float] | None = None + if name == "kmer_tfidf": + vectorizer = TfidfVectorizer(analyzer="char", ngram_range=(3, 3), lowercase=False) + inner_train_x = vectorizer.fit_transform(sequences[tune_train]) + inner_val_x = vectorizer.transform(sequences[tune_val]) + else: + (inner_train_x, inner_val_x), _ = transform_dense( + matrix, tune_train, tune_train, tune_val + ) + for c_value in c_grid: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ConvergenceWarning) + model = classifier(c_value, max_iter, spec.seed).fit(inner_train_x, y[tune_train]) + value = float(f1_score(y[tune_val], model.predict(inner_val_x), average="macro")) + candidate = (value, -c_value) + if best is None or candidate > best: + best = candidate + chosen_c = -best[1] + + if name == "kmer_tfidf": + vectorizer = TfidfVectorizer(analyzer="char", ngram_range=(3, 3), lowercase=False) + train_x = vectorizer.fit_transform(sequences[spec.train]) + test_x = vectorizer.transform(sequences[spec.test]) + dimension = int(train_x.shape[1]) + else: + (train_x, test_x), _ = transform_dense(matrix, spec.train, spec.train, spec.test) + dimension = int(matrix.shape[1]) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ConvergenceWarning) + model = classifier(chosen_c, max_iter, spec.seed).fit(train_x, y[spec.train]) + probabilities = model.predict_proba(test_x) + if not np.array_equal(model.classes_, classes): + raise ValueError("A fold lacks one or more classes despite split validation.") + metrics = score_predictions(y[spec.test], probabilities, classes) + row = { + "protocol": spec.protocol, "repeat": spec.repeat, "fold": spec.fold, + "seed": spec.seed, "representation": name, "display_name": DISPLAY[name], + "dim": dimension, "train_rows": len(spec.train), "test_rows": len(spec.test), + "selected_c": chosen_c, **metrics, + } + predicted = classes[np.argmax(probabilities, axis=1)] + top3 = classes[np.argpartition(probabilities, -3, axis=1)[:, -3:]] + prediction = pd.DataFrame({ + "protocol": spec.protocol, + "repeat": spec.repeat, + "fold": spec.fold, + "representation": name, + "prediction_row_index": spec.test, + "true_label": y[spec.test], + "predicted_label": predicted, + "correct": predicted == y[spec.test], + "top3_labels": ["|".join(map(str, sorted(values))) for values in top3], + "true_label_probability": probabilities[ + np.arange(len(spec.test)), np.searchsorted(classes, y[spec.test]) + ], + }) + return row, prediction + + +def aggregate_metrics(raw: pd.DataFrame) -> pd.DataFrame: + metrics = ["macro_f1", "balanced_accuracy", "accuracy", "macro_auroc", "top3_accuracy"] + grouped = raw.groupby(["protocol", "representation", "display_name"], sort=False) + rows = [] + for keys, frame in grouped: + row = dict(zip(("protocol", "representation", "display_name"), keys)) + row["dim"] = int(round(float(frame["dim"].median()))) + row["dim_min"] = int(frame["dim"].min()) + row["dim_max"] = int(frame["dim"].max()) + row["observations"] = len(frame) + for metric in metrics: + row[f"{metric}_mean"] = float(frame[metric].mean()) + row[f"{metric}_std"] = float(frame[metric].std(ddof=1)) + rows.append(row) + return pd.DataFrame(rows) + + +def paired_statistics(raw: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]: + comparisons = ("rp", "tcr_bert", "sceptr_cdr3", "esm2_35m") + differences, tests = [], [] + for protocol in raw["protocol"].unique(): + subset = raw[raw.protocol == protocol] + wide = subset.pivot(index=["repeat", "fold", "seed"], columns="representation", + values="macro_f1") + if "rtp" not in wide: + continue + for comparator in comparisons: + if comparator not in wide: + continue + paired = wide[["rtp", comparator]].dropna().copy() + paired["difference"] = paired["rtp"] - paired[comparator] + for index, values in paired.iterrows(): + differences.append({ + "protocol": protocol, "repeat": index[0], "fold": index[1], + "seed": index[2], "comparison": f"rtp_minus_{comparator}", + "rtp_macro_f1": values["rtp"], "comparator_macro_f1": values[comparator], + "difference": values["difference"], + }) + delta = paired["difference"].to_numpy() + try: + statistic, p_value = wilcoxon(delta, alternative="two-sided", method="auto") + except ValueError: + statistic, p_value = 0.0, 1.0 + tests.append({ + "protocol": protocol, "comparison": f"rtp_vs_{comparator}", + "observations": len(delta), "mean_paired_difference": float(delta.mean()), + "median_paired_difference": float(np.median(delta)), + "wins": int((delta > 0).sum()), "ties": int((delta == 0).sum()), + "losses": int((delta < 0).sum()), "wilcoxon_statistic": float(statistic), + "wilcoxon_two_sided_p": float(p_value), + }) + return pd.DataFrame(differences), pd.DataFrame(tests) + + +def latex_table(aggregate: pd.DataFrame, protocol: str) -> str: + data = aggregate[aggregate.protocol == protocol].set_index("representation") + lines = [ + r"\begin{tabular}{lrrrrr}", r"\toprule", + r"Representation & Dim & Macro F1 & Balanced accuracy & Accuracy & Macro AUROC \\", + r"\midrule", + ] + for name in PRIMARY_ORDER: + if name not in data.index: + continue + row = data.loc[name] + cell = lambda metric: f"{row[f'{metric}_mean']:.3f} $\\pm$ {row[f'{metric}_std']:.3f}" + lines.append( + f"{DISPLAY[name]} & {int(row['dim'])} & {cell('macro_f1')} & " + f"{cell('balanced_accuracy')} & {cell('accuracy')} & {cell('macro_auroc')} \\\\" + ) + lines.extend([r"\bottomrule", r"\end{tabular}", ""]) + return "\n".join(lines) + + +def create_figure(raw: pd.DataFrame, aggregate: pd.DataFrame, output: Path) -> None: + standard = aggregate[aggregate.protocol == "sequence_stratified"].set_index("representation") + order = [name for name in PRIMARY_ORDER if name in standard.index] + fig, axes = plt.subplots(1, 3, figsize=(11.2, 3.65), gridspec_kw={"width_ratios": [1.55, 1, 1]}) + ax = axes[0] + positions = np.arange(len(order)) + means = standard.loc[order, "macro_f1_mean"].to_numpy() + errors = standard.loc[order, "macro_f1_std"].to_numpy() + colors = ["#d55e00" if name == "rtp" else "#4477aa" if name in NATIVE_NAMES else "#228833" + for name in order] + ax.errorbar(positions, means, yerr=errors, fmt="none", ecolor="#333333", capsize=2, lw=.9) + ax.scatter(positions, means, c=colors, s=34, edgecolor="white", linewidth=.5, zorder=3) + ax.set_xticks(positions, [DISPLAY[name] for name in order], rotation=50, ha="right") + ax.set_ylabel("Macro F1") + ax.set_title("A Sequence-stratified performance", loc="left", fontweight="bold") + ax.grid(axis="y", color="#dddddd", linewidth=.6) + + def paired_panel(panel, protocol: str, title: str) -> None: + subset = raw[(raw.protocol == protocol) & raw.representation.isin(["rp", "rtp"])] + wide = subset.pivot(index=["repeat", "fold"], columns="representation", values="macro_f1").dropna() + for _, row in wide.iterrows(): + panel.plot([0, 1], [row.rp, row.rtp], color="#999999", alpha=.65, lw=.7) + panel.scatter([0, 1], [row.rp, row.rtp], c=["#4477aa", "#d55e00"], s=18, zorder=3) + delta = (wide.rtp - wide.rp).mean() + panel.set_xticks([0, 1], ["RP", "RTP"]) + panel.set_ylabel("Macro F1") + panel.set_title(title, loc="left", fontweight="bold") + panel.text(.5, .02, f"mean Δ={delta:+.3f}", transform=panel.transAxes, + ha="center", va="bottom", fontsize=8) + panel.grid(axis="y", color="#dddddd", linewidth=.6) + + paired_panel(axes[1], "sequence_stratified", "B TCRemP effect") + if "redcea_similarity_grouped" in set(raw.protocol): + paired_panel(axes[2], "redcea_similarity_grouped", "C Cluster-grouped") + else: + axes[2].axis("off") + fig.suptitle("Frozen TCR representations for multiclass epitope prediction", y=1.01, fontsize=12) + fig.tight_layout() + fig.savefig(output / "epitope_prediction.png", dpi=600, bbox_inches="tight") + fig.savefig(output / "epitope_prediction.pdf", bbox_inches="tight") + plt.close(fig) + + +def main() -> None: + args = parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + if len(set(args.seeds)) != len(args.seeds) or args.folds < 2: + raise ValueError("Seeds must be unique and at least two folds are required.") + if any(value <= 0 for value in args.c_grid): + raise ValueError("Every C value must be positive.") + unknown = set(args.representations).difference(PRIMARY_ORDER) + if unknown: + raise ValueError(f"Unknown representations: {sorted(unknown)}") + + frame, original = selected_cohort(args.cohort, args.uniform_evaluation) + audit, class_sizes = audit_dataset(frame) + frame.to_csv(args.output_dir / "samples.tsv", sep="\t", index=False, lineterminator="\n") + class_sizes.to_csv(args.output_dir / "class_sizes.csv", index=False, lineterminator="\n") + exact_groups = pd.factorize(frame["cdr3"], sort=True)[0] + similarity_groups = redcea_component_groups(frame) + label_names = np.array(sorted(frame["label"].astype(str).unique())) + label_to_id = {label: index for index, label in enumerate(label_names)} + y_id = np.array([label_to_id[label] for label in frame["label"].astype(str)], dtype=np.int32) + protocols, split_report = make_splits( + y_id, exact_groups, similarity_groups, args.seeds, args.folds, + not args.skip_similarity_grouped, + ) + write_split_assignments(args.output_dir / "split_assignments.csv", protocols, frame) + + specs = [] + for protocol, folds in protocols.items(): + if protocol == "exact_cdr3_grouped": + continue # exact groups are singletons; metrics are copied below with an explicit flag + group_values = similarity_groups if protocol == "redcea_similarity_grouped" else None + for repeat, fold, train, test in folds: + specs.append((FoldSpec(protocol, repeat, fold, train, test, args.seeds[repeat]), group_values)) + + all_metrics, all_predictions, rep_manifest = [], [], {} + sequences = frame["cdr3"].astype(str).to_numpy() + source_rows = int(pd.read_csv(args.cohort, sep="\t", usecols=["cdr3"]).shape[0]) + for name in args.representations: + matrix = None if name == "kmer_tfidf" else load_dense_matrix( + name, original, args.native_evaluation, args.external_representations, source_rows + ) + rep_manifest[name] = { + "display_name": DISPLAY[name], + "dim": "fold-specific vocabulary" if matrix is None else int(matrix.shape[1]), + } + tasks = Parallel(n_jobs=args.jobs, verbose=10)( + delayed(evaluate_fold)( + spec, name, matrix, sequences, y_id, np.arange(len(label_names)), group_values, + tuple(sorted(set(args.c_grid))), args.max_iter, + ) + for spec, group_values in specs + ) + for metrics, predictions in tasks: + predictions["true_label"] = predictions["true_label"].map(dict(enumerate(label_names))) + predictions["predicted_label"] = predictions["predicted_label"].map(dict(enumerate(label_names))) + # top3 values are encoded class IDs; preserve IDs compactly and map is in label_classes.tsv. + all_metrics.append(metrics) + all_predictions.append(predictions) + + raw = pd.DataFrame(all_metrics) + exact = raw[raw.protocol == "sequence_stratified"].copy() + exact["protocol"] = "exact_cdr3_grouped" + exact["evaluation_reused"] = True + raw["evaluation_reused"] = False + raw = pd.concat([raw, exact], ignore_index=True) + raw.to_csv(args.output_dir / "fold_metrics.csv", index=False, lineterminator="\n") + pd.concat(all_predictions, ignore_index=True).to_parquet( + args.output_dir / "predictions.parquet", index=False, compression="zstd" + ) + pd.DataFrame({"class_id": np.arange(len(label_names)), "label": label_names}).to_csv( + args.output_dir / "label_classes.tsv", sep="\t", index=False, lineterminator="\n" + ) + aggregate = aggregate_metrics(raw) + aggregate.to_csv(args.output_dir / "aggregated_metrics.csv", index=False, lineterminator="\n") + differences, tests = paired_statistics(raw) + differences.to_csv(args.output_dir / "paired_macro_f1_differences.csv", index=False, + lineterminator="\n") + tests.to_csv(args.output_dir / "paired_tests.csv", index=False, lineterminator="\n") + for protocol in aggregate.protocol.unique(): + (args.output_dir / f"table_{protocol}.tex").write_text( + latex_table(aggregate, protocol), encoding="utf-8" + ) + create_figure(raw, aggregate, args.output_dir) + + input_files = { + "cohort": args.cohort, + "uniform_pairs": args.uniform_evaluation / "uniform_pairs.npz", + } + manifest = { + "status": "complete", + "dataset": audit, + "protocols": { + "outer": { + "requested_folds": args.folds, "seeds": args.seeds, + "standard": "repeated shuffled StratifiedKFold", + "exact_cdr3_grouped": split_report["exact_cdr3_grouped"], + "similarity_grouped": split_report.get("redcea_similarity_grouped"), + }, + "classifier": "L2 multinomial logistic regression; balanced class weights", + "standardization": "training-fold statistics only", + "hyperparameter_selection": { + "c_grid": sorted(set(args.c_grid)), + "criterion": "inner-training macro F1; smaller C wins ties", + "inner_split": "stratified 15% validation; group-aware for grouped protocol", + }, + "tfidf": "character 3-mer TF-IDF fit on training rows only", + "exact_group_evaluation_reused": True, + }, + "representations": rep_manifest, + "inputs": {name: {"path": str(path), "sha256": sha256(path)} + for name, path in input_files.items()}, + "outputs": {name: sha256(args.output_dir / name) for name in ( + "samples.tsv", "class_sizes.csv", "split_assignments.csv", "fold_metrics.csv", + "aggregated_metrics.csv", "paired_macro_f1_differences.csv", "paired_tests.csv", + "predictions.parquet", "epitope_prediction.png", "epitope_prediction.pdf", + )}, + "software": {"python": os.sys.version, "slurm_job_id": os.environ.get("SLURM_JOB_ID")}, + } + stable_json(args.output_dir / "RESULTS.json", manifest) + print(json.dumps({"status": "complete", "output": str(args.output_dir), **audit}, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_vdjdb_frozen.py b/benchmark/evaluate_vdjdb_frozen.py new file mode 100644 index 0000000..bbccf54 --- /dev/null +++ b/benchmark/evaluate_vdjdb_frozen.py @@ -0,0 +1,497 @@ +"""Evaluate matched frozen IRRM-CODEC encoders on a locked VDJdb cohort.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import math +import os +from collections import defaultdict +from dataclasses import asdict +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from sklearn.metrics import average_precision_score, roc_auc_score + +from irrm_codec.multitask_data import resolve_encoder_tokenizer +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import PAD_ID + + +AA = "ACDEFGHIKLMNPQRSTVWY" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--cohort-preflight", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument( + "--model", action="append", required=True, + help="NAME:CHECKPOINT:RUN_CONFIG; legacy R/P/RP/RTP or the full seven-condition factorial", + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--max-pairs-per-label", type=int, default=2000) + parser.add_argument("--bootstrap", type=int, default=2000) + parser.add_argument("--permutations", type=int, default=1000) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def stable_json(path: Path, value: object) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def parse_models(values: list[str]) -> dict[str, tuple[Path, Path]]: + result = {} + for value in values: + fields = value.split(":", 2) + if len(fields) != 3: + raise ValueError(f"Invalid --model value: {value!r}") + name, checkpoint, config = fields + if name in result: + raise ValueError(f"Duplicate model name: {name}") + result[name] = (Path(checkpoint), Path(config)) + legacy = {"r", "p", "rp", "rtp"} + factorial = {"r", "t", "p", "rt", "rp", "tp", "rtp"} + if set(result) not in (legacy, factorial): + raise ValueError("Models must be legacy r/p/rp/rtp or the full seven-condition factorial.") + return result + + +def validate_run_configs(models: dict[str, tuple[Path, Path]]) -> tuple[IRRMCodecConfig, dict]: + expected_weights = { + "r": (1.0, 0.0, 0.0), + "t": (0.0, 1.0, 0.0), + "p": (0.0, 0.0, 1.0), + "rt": (1.0, 1.0, 0.0), + "rp": (1.0, 0.0, 1.0), + "tp": (0.0, 1.0, 1.0), + "rtp": (1.0, 1.0, 1.0), + } + configs = {name: json.loads(path.read_text(encoding="utf-8")) for name, (_, path) in models.items()} + reference = configs["rtp"] + ignored = {"output_dir", "reconstruction_loss_weight", "tcremp_loss_weight", "pgen_loss_weight"} + ref_training = {key: value for key, value in reference["training"].items() if key not in ignored} + report = {} + for name, payload in configs.items(): + if payload["model"] != reference["model"]: + raise ValueError(f"{name} architecture differs from full DATA-ANCHOR reference.") + comparable = {key: value for key, value in payload["training"].items() if key not in ignored} + if comparable != ref_training: + differing = sorted(key for key in set(comparable) | set(ref_training) + if comparable.get(key) != ref_training.get(key)) + raise ValueError(f"{name} has non-weight training differences: {differing}") + weights = ( + float(payload["training"]["reconstruction_loss_weight"]), + float(payload["training"]["tcremp_loss_weight"]), + float(payload["training"]["pgen_loss_weight"]), + ) + if weights != expected_weights[name]: + raise ValueError(f"{name} has weights {weights}, expected {expected_weights[name]}") + report[name] = { + "checkpoint": str(models[name][0]), + "checkpoint_sha256": sha256(models[name][0]), + "run_config": str(models[name][1]), + "run_config_sha256": sha256(models[name][1]), + "loss_weights_reconstruction_tcremp_pgen": list(weights), + "parameter_count": payload["parameter_count"], + } + return IRRMCodecConfig(**reference["model"]), report + + +def extract_embeddings( + cohort: pd.DataFrame, + models: dict[str, tuple[Path, Path]], + model_config: IRRMCodecConfig, + tokenizer_path: Path, + batch_size: int, + output: Path, +) -> tuple[dict[str, np.ndarray], dict]: + tokenizer = resolve_encoder_tokenizer("data_anchor", str(tokenizer_path)) + if tokenizer.vocab_size != model_config.input_vocab_size: + raise ValueError("Tokenizer vocabulary differs from matched checkpoint architecture.") + encoded = [tokenizer.encode(sequence, model_config.max_sequence_len) + for sequence in cohort["cdr3"].astype(str)] + results = {} + metadata = {} + for name, (checkpoint_path, _) in models.items(): + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + if checkpoint["model_config"] != asdict(model_config): + raise ValueError(f"{name} checkpoint model_config differs from run_config.") + model = IRRMCodecTransformer(model_config) + model.load_state_dict(checkpoint["model_state"], strict=True) + model.eval() + batches = [] + with torch.no_grad(): + for start in range(0, len(encoded), batch_size): + rows = encoded[start:start + batch_size] + width = max(map(len, rows)) + tokens = torch.full((len(rows), width), PAD_ID, dtype=torch.long) + for index, ids in enumerate(rows): + tokens[index, :len(ids)] = torch.tensor(ids, dtype=torch.long) + batches.append(model.encode(tokens, tokens.ne(PAD_ID)).numpy().astype(np.float32)) + array = np.concatenate(batches) + if array.shape != (len(cohort), model_config.latent_dim) or not np.isfinite(array).all(): + raise ValueError(f"{name} emitted invalid latent array {array.shape}.") + path = output / f"embeddings_{name}.npy" + np.save(path, array, allow_pickle=False) + results[name] = array + metadata[name] = { + "best_epoch_one_based": int(checkpoint["epoch"]), + "best_validation_loss": float(checkpoint["best_val_loss"]), + "shape": list(array.shape), + "finite": True, + "sha256": sha256(path), + } + del checkpoint, model + return results, metadata + + +def distance_matrix(array: np.ndarray, metric: str) -> np.ndarray: + x = array.astype(np.float64) + if metric == "cosine": + norms = np.linalg.norm(x, axis=1, keepdims=True) + if np.any(norms == 0): + raise ValueError("Zero-norm latent prevents cosine distance.") + unit = x / norms + return np.maximum(0.0, 1.0 - unit @ unit.T) + squared = np.sum(x * x, axis=1, keepdims=True) + return np.sqrt(np.maximum(0.0, squared + squared.T - 2.0 * (x @ x.T))) + + +def build_pair_manifest(cohort: pd.DataFrame, limit: int, seed: int) -> pd.DataFrame: + labels = cohort["label"].astype(str).to_numpy() + donors = cohort["donor_key"].astype(str).to_numpy() + lengths = cohort["length"].to_numpy(dtype=int) + by_label = {label: np.flatnonzero(labels == label) for label in sorted(set(labels))} + by_length = {length: np.flatnonzero(lengths == length) for length in sorted(set(lengths))} + rows = [] + for label_index, (label, indices) in enumerate(by_label.items()): + positives = [(int(indices[i]), int(indices[j])) for i in range(len(indices)) + for j in range(i + 1, len(indices)) if donors[indices[i]] != donors[indices[j]]] + rng = np.random.default_rng(seed + label_index) + if len(positives) > limit: + chosen = rng.choice(len(positives), size=limit, replace=False) + positives = [positives[index] for index in sorted(chosen)] + for pair_index, (query, positive) in enumerate(positives): + if rng.integers(2): + query, positive = positive, query + candidates = by_length.get(int(lengths[positive]), np.empty(0, dtype=int)) + candidates = candidates[(labels[candidates] != label) & (donors[candidates] != donors[query])] + match_type = "exact_length" + if len(candidates) == 0: + allowed = np.flatnonzero((labels != label) & (donors != donors[query])) + difference = np.abs(lengths[allowed] - lengths[positive]) + candidates = allowed[difference == difference.min()] + match_type = "nearest_length" + negative = int(candidates[rng.integers(len(candidates))]) + rows.append({ + "label": label, "query_index": query, "positive_index": positive, + "negative_index": negative, "query_donor": donors[query], + "positive_donor": donors[positive], "negative_donor": donors[negative], + "query_length": int(lengths[query]), "positive_length": int(lengths[positive]), + "negative_length": int(lengths[negative]), "match_type": match_type, + "pair_index_within_label": pair_index, + }) + result = pd.DataFrame(rows) + if result.empty or (result["query_donor"] == result["positive_donor"]).any(): + raise ValueError("Could not construct donor-disjoint positive pairs.") + return result + + +def pair_metrics(matrix: np.ndarray, pairs: pd.DataFrame) -> tuple[dict, pd.DataFrame]: + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + within = matrix[q, p] + between = matrix[q, n] + truth = np.concatenate((np.ones(len(within), dtype=int), np.zeros(len(between), dtype=int))) + score = -np.concatenate((within, between)) + per_label = [] + label_values = pairs["label"].astype(str).to_numpy() + for label in sorted(set(label_values)): + mask = label_values == label + pooled = math.sqrt((float(within[mask].var()) + float(between[mask].var())) / 2.0) + per_label.append({ + "label": label, + "pairs": int(mask.sum()), + "within_mean": float(within[mask].mean()), + "between_mean": float(between[mask].mean()), + "within_between_ratio": float(within[mask].mean() / between[mask].mean()), + "standardized_separation": float((between[mask].mean() - within[mask].mean()) / pooled) + if pooled > 0 else float("nan"), + }) + table = pd.DataFrame(per_label) + return { + "pairs_per_class": len(within), + "same_label_auroc": float(roc_auc_score(truth, score)), + "same_label_auprc_balanced": float(average_precision_score(truth, score)), + "macro_within_between_ratio": float(table["within_between_ratio"].mean()), + "macro_standardized_separation": float(table["standardized_separation"].mean()), + "within_expected_lower_fraction": float((table["within_mean"] < table["between_mean"]).mean()), + }, table + + +def retrieval_metrics(matrix: np.ndarray, cohort: pd.DataFrame) -> tuple[dict, pd.DataFrame]: + labels = cohort["label"].astype(str).to_numpy() + donors = cohort["donor_key"].astype(str).to_numpy() + rows = [] + for query in range(len(cohort)): + allowed = np.flatnonzero((donors != donors[query]) & (np.arange(len(cohort)) != query)) + relevant = labels[allowed] == labels[query] + total_relevant = int(relevant.sum()) + if total_relevant == 0: + continue + order = np.argsort(matrix[query, allowed], kind="mergesort") + ranked = relevant[order] + values = {"query_index": query, "label": labels[query], "query_donor": donors[query], + "relevant_donor_excluded": total_relevant} + for k in (1, 5, 10): + values[f"precision_at_{k}"] = float(ranked[:k].sum() / k) + r = total_relevant + hits = ranked[:r].astype(float) + precision = np.cumsum(hits) / np.arange(1, r + 1) + values["ap_at_r"] = float((precision * hits).sum() / r) + rows.append(values) + table = pd.DataFrame(rows) + return { + "eligible_queries": len(table), + "precision_at_1": float(table["precision_at_1"].mean()), + "precision_at_5": float(table["precision_at_5"].mean()), + "precision_at_10": float(table["precision_at_10"].mean()), + "map_at_r": float(table["ap_at_r"].mean()), + }, table + + +def bootstrap_primary( + matrices: dict[str, np.ndarray], pairs: pd.DataFrame, + retrieval: dict[str, pd.DataFrame], replicates: int, seed: int, +) -> dict: + donors = sorted( + set(pairs["query_donor"].astype(str)).union( + *(set(table["query_donor"].astype(str)) for table in retrieval.values()) + ) + ) + donor_index = {donor: index for index, donor in enumerate(donors)} + pair_donor = pairs["query_donor"].map(donor_index).to_numpy(int) + labels = sorted(set(pairs["label"].astype(str))) + label_index = {label: index for index, label in enumerate(labels)} + pair_label = pairs["label"].map(label_index).to_numpy(int) + rng = np.random.default_rng(seed) + weights = rng.multinomial(len(donors), np.full(len(donors), 1 / len(donors)), size=replicates) + model_boot = {} + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + for name, matrix in matrices.items(): + within = matrix[q, p] + between = matrix[q, n] + donor_label_within = np.zeros((len(donors), len(labels))) + donor_label_between = np.zeros_like(donor_label_within) + donor_label_count = np.zeros_like(donor_label_within) + np.add.at(donor_label_within, (pair_donor, pair_label), within) + np.add.at(donor_label_between, (pair_donor, pair_label), between) + np.add.at(donor_label_count, (pair_donor, pair_label), 1) + summed_within = weights @ donor_label_within + summed_between = weights @ donor_label_between + counts = weights @ donor_label_count + mean_within = np.full_like(summed_within, np.nan) + mean_between = np.full_like(summed_between, np.nan) + np.divide(summed_within, counts, out=mean_within, where=counts > 0) + np.divide(summed_between, counts, out=mean_between, where=counts > 0) + ratios = np.full_like(mean_within, np.nan) + np.divide(mean_within, mean_between, out=ratios, where=mean_between > 0) + macro_ratio = np.nanmean(ratios, axis=1) + + query_table = retrieval[name] + query_donor_index = query_table["query_donor"].map(donor_index).to_numpy(int) + donor_ap_sum = np.zeros(len(donors)) + donor_query_count = np.zeros(len(donors)) + np.add.at(donor_ap_sum, query_donor_index, query_table["ap_at_r"].to_numpy(float)) + np.add.at(donor_query_count, query_donor_index, 1) + boot_map = (weights @ donor_ap_sum) / np.maximum(weights @ donor_query_count, 1) + model_boot[name] = {"macro_ratio": macro_ratio, "map_at_r": boot_map} + + def summarize(values: np.ndarray) -> dict: + return {"ci95": [float(np.nanpercentile(values, 2.5)), float(np.nanpercentile(values, 97.5))]} + + report = {"replicates": replicates, "cluster": "query donor", "models": {}} + for name, values in model_boot.items(): + report["models"][name] = { + "macro_within_between_ratio": summarize(values["macro_ratio"]), + "map_at_r": summarize(values["map_at_r"]), + } + report["paired_differences_rtp_minus_other"] = {} + for other in ("r", "p", "rp"): + report["paired_differences_rtp_minus_other"][other] = { + "macro_ratio_lower_is_better": summarize( + model_boot["rtp"]["macro_ratio"] - model_boot[other]["macro_ratio"] + ), + "map_at_r_higher_is_better": summarize( + model_boot["rtp"]["map_at_r"] - model_boot[other]["map_at_r"] + ), + } + return report + + +def assert_finite_tree(value: object, path: str = "root") -> None: + if isinstance(value, dict): + for key, child in value.items(): + assert_finite_tree(child, f"{path}.{key}") + elif isinstance(value, list): + for index, child in enumerate(value): + assert_finite_tree(child, f"{path}[{index}]") + elif isinstance(value, (float, np.floating)) and not math.isfinite(float(value)): + raise FloatingPointError(f"Non-finite result at {path}: {value}") + + +def controls( + cohort: pd.DataFrame, pairs: pd.DataFrame, matrices: dict[str, np.ndarray], + permutations: int, seed: int, +) -> dict: + lengths = cohort["length"].to_numpy(float) + length_matrix = np.abs(lengths[:, None] - lengths[None, :]) + length_metrics, _ = pair_metrics(length_matrix, pairs) + composition = np.zeros((len(cohort), len(AA)), dtype=float) + for index, sequence in enumerate(cohort["cdr3"].astype(str)): + for aa_index, aa in enumerate(AA): + composition[index, aa_index] = sequence.count(aa) / len(sequence) + composition_metrics = {} + for metric in ("cosine", "euclidean"): + value, _ = pair_metrics(distance_matrix(composition, metric), pairs) + composition_metrics[metric] = value + + rng = np.random.default_rng(seed + 10000) + all_i, all_j = np.triu_indices(len(cohort), k=1) + donor = cohort["donor_key"].astype(str).to_numpy() + same_length = lengths[all_i] == lengths[all_j] + donor_disjoint = donor[all_i] != donor[all_j] + eligible = np.flatnonzero(same_length & donor_disjoint) + if len(eligible) > 100000: + eligible = rng.choice(eligible, size=100000, replace=False) + pair_i, pair_j = all_i[eligible], all_j[eligible] + labels = cohort["label"].astype(str).to_numpy() + actual_same = labels[pair_i] == labels[pair_j] + model_distance = {name: matrix[pair_i, pair_j] for name, matrix in matrices.items()} + actual = { + name: float(distance[~actual_same].mean() - distance[actual_same].mean()) + for name, distance in model_distance.items() + } + null = {name: np.empty(permutations, dtype=float) for name in matrices} + groups = [np.flatnonzero(lengths == value) for value in sorted(set(lengths))] + for replicate in range(permutations): + shuffled = labels.copy() + for group in groups: + shuffled[group] = rng.permutation(shuffled[group]) + same = shuffled[pair_i] == shuffled[pair_j] + for name, distance in model_distance.items(): + null[name][replicate] = float(distance[~same].mean() - distance[same].mean()) + return { + "length_only_on_primary_pairs": length_metrics, + "amino_acid_composition": composition_metrics, + "shuffled_label_null": { + "pairs": len(pair_i), "permutations": permutations, + "permutation_strata": "exact CDR3 length", + "models": { + name: { + "observed_between_minus_within": actual[name], + "null_ci95": [float(np.percentile(null[name], 2.5)), float(np.percentile(null[name], 97.5))], + "one_sided_p": float((1 + np.sum(null[name] >= actual[name])) / (permutations + 1)), + } for name in matrices + }, + }, + } + + +def main() -> None: + args = parse_args() + models = parse_models(args.model) + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + preflight = json.loads(args.cohort_preflight.read_text(encoding="utf-8")) + if preflight["status"] != "ready_for_cohort_review_not_embedding_extraction": + raise ValueError("Cohort preflight status is not accepting.") + if sha256(args.cohort) != preflight["cohort"]["cohort_sha256"]: + raise ValueError("Cohort hash differs from accepted CPU preflight.") + cohort = pd.read_csv(args.cohort, sep="\t") + if cohort["cdr3"].duplicated().any() or len(cohort) != preflight["cohort"]["unique_cdr3"]: + raise ValueError("Cohort is not the expected unique-clonotype table.") + + config, model_report = validate_run_configs(models) + embeddings, extraction = extract_embeddings( + cohort, models, config, args.tokenizer, args.batch_size, args.output_dir + ) + pairs = build_pair_manifest(cohort, args.max_pairs_per_label, args.seed) + pairs.to_csv(args.output_dir / "pair_manifest.tsv", sep="\t", index=False) + matrices_by_metric = {} + metrics = {} + retrieval_tables = {} + per_label_tables = [] + for metric in ("cosine", "euclidean"): + matrices = {name: distance_matrix(array, metric) for name, array in embeddings.items()} + matrices_by_metric[metric] = matrices + metrics[metric] = {} + retrieval_tables[metric] = {} + for name, matrix in matrices.items(): + pair_value, per_label = pair_metrics(matrix, pairs) + retrieval_value, per_query = retrieval_metrics(matrix, cohort) + metrics[metric][name] = {"pairwise_and_contrast": pair_value, "retrieval": retrieval_value} + retrieval_tables[metric][name] = per_query + per_label.insert(0, "model", name) + per_label.insert(0, "distance", metric) + per_label_tables.append(per_label) + per_query.to_csv(args.output_dir / f"retrieval_queries_{metric}_{name}.tsv", sep="\t", index=False) + stable_json( + args.output_dir / f"bootstrap_{metric}.json", + bootstrap_primary(matrices, pairs, retrieval_tables[metric], args.bootstrap, args.seed), + ) + pd.concat(per_label_tables, ignore_index=True).to_csv( + args.output_dir / "per_label_metrics.tsv", sep="\t", index=False + ) + control_report = controls( + cohort, pairs, matrices_by_metric["cosine"], args.permutations, args.seed + ) + assert_finite_tree(metrics, "metrics") + assert_finite_tree(control_report, "controls") + stable_json(args.output_dir / "metrics.json", metrics) + stable_json(args.output_dir / "controls.json", control_report) + + artifact_hashes = { + path.name: sha256(path) for path in sorted(args.output_dir.iterdir()) if path.is_file() + } + manifest = { + "status": "complete", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "study": "matched DATA-ANCHOR frozen-encoder VDJdb epitope separation", + "cohort": {"path": str(args.cohort), "sha256": sha256(args.cohort), "rows": len(cohort), + "labels": int(cohort["label"].nunique()), "donors": int(cohort["donor_key"].nunique())}, + "models": model_report, + "extraction": extraction, + "pairs": {"rows": len(pairs), "sha256": sha256(args.output_dir / "pair_manifest.tsv"), + "nearest_length_fallbacks": int((pairs["match_type"] != "exact_length").sum())}, + "selection_statement": "No checkpoint or model was selected using VDJdb outcomes; all are reported.", + "inactive_head_statement": "Only frozen encoder latents are evaluated; inactive-head predictions are ignored.", + "artifact_sha256": artifact_hashes, + "nonfinite_scan": "passed", + } + stable_json(args.output_dir / "RESULTS.json", manifest) + print(json.dumps(manifest, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_vdjdb_redcea_clonotypes.py b/benchmark/evaluate_vdjdb_redcea_clonotypes.py new file mode 100644 index 0000000..df9d78b --- /dev/null +++ b/benchmark/evaluate_vdjdb_redcea_clonotypes.py @@ -0,0 +1,506 @@ +"""Evaluate frozen DATA-ANCHOR latents on REDCEA motif-member clonotypes. + +Every retained clonotype is a query. Its deterministic candidate gallery has +up to 20 same-epitope clonotypes and four times as many other-epitope controls, +selected without model embeddings. This keeps the all-clonotype evaluation +tractable while fixing positive prevalence and using identical comparisons for +all seven frozen encoders. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +from pathlib import Path + +import numpy as np +import pandas as pd +from sklearn.cluster import MiniBatchKMeans +from sklearn.metrics import ( + adjusted_mutual_info_score, + adjusted_rand_score, + balanced_accuracy_score, + calinski_harabasz_score, + completeness_score, + davies_bouldin_score, + f1_score, + homogeneity_score, + precision_recall_fscore_support, + silhouette_score, + v_measure_score, +) + +PRIMARY_MODELS = ("r", "t", "p", "rt", "rp", "tp", "rtp") +METRIC_COLUMNS = ( + "precision_at_1", "precision_at_5", "precision_at_10", "average_precision", + "within_between_ratio", "between_minus_within", "nearest_margin", +) + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def stable_json(path: Path, value: object) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def assert_finite_tree(value: object, path: str = "root") -> None: + if isinstance(value, dict): + for key, child in value.items(): + assert_finite_tree(child, f"{path}.{key}") + elif isinstance(value, list): + for index, child in enumerate(value): + assert_finite_tree(child, f"{path}[{index}]") + elif isinstance(value, (float, np.floating)) and not math.isfinite(float(value)): + raise FloatingPointError(f"Non-finite result at {path}: {value}") + + +def build_candidate_gallery( + cohort: pd.DataFrame, max_positives: int, negative_ratio: int, seed: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict]: + labels_text = cohort["label"].astype(str).to_numpy() + label_names = np.array(sorted(set(labels_text)), dtype=object) + label_to_id = {label: index for index, label in enumerate(label_names)} + labels = np.array([label_to_id[value] for value in labels_text], dtype=np.int32) + species = cohort["species"].fillna("").astype(str).to_numpy() + lengths = cohort["length"].to_numpy(dtype=np.int16) + v_call = cohort["v_call"].fillna("").astype(str).to_numpy() + j_call = cohort["j_call"].fillna("").astype(str).to_numpy() + by_label = {label: np.flatnonzero(labels == label) for label in range(len(label_names))} + by_label_species_length: dict[tuple[int, str, int], np.ndarray] = {} + by_label_species_length_vj: dict[tuple[int, str, int, str, str], np.ndarray] = {} + by_label_species_length_v: dict[tuple[int, str, int, str], np.ndarray] = {} + by_label_species_length_j: dict[tuple[int, str, int, str], np.ndarray] = {} + by_label_species: dict[tuple[int, str], dict[int, np.ndarray]] = {} + by_label_length: dict[tuple[int, int], np.ndarray] = {} + for label, indices in by_label.items(): + for sp in sorted(set(species[indices])): + sp_indices = indices[species[indices] == sp] + lengths_map = {} + for length in sorted(set(lengths[sp_indices])): + values = sp_indices[lengths[sp_indices] == length] + by_label_species_length[(label, sp, int(length))] = values + for v_value in sorted(set(v_call[values])): + by_label_species_length_v[(label, sp, int(length), v_value)] = values[ + v_call[values] == v_value + ] + for j_value in sorted(set(j_call[values])): + by_label_species_length_j[(label, sp, int(length), j_value)] = values[ + j_call[values] == j_value + ] + for v_value, j_value in sorted(set(zip(v_call[values], j_call[values]))): + by_label_species_length_vj[( + label, sp, int(length), v_value, j_value + )] = values[(v_call[values] == v_value) & (j_call[values] == j_value)] + lengths_map[int(length)] = values + by_label_species[(label, sp)] = lengths_map + for length in sorted(set(lengths[indices])): + by_label_length[(label, int(length))] = indices[lengths[indices] == length] + + width = max_positives * (negative_ratio + 1) + candidates = np.full((len(cohort), width), -1, dtype=np.int32) + relevant = np.zeros((len(cohort), width), dtype=bool) + tiers = np.full((len(cohort), width), -1, dtype=np.int8) + tier_counts: dict[str, int] = {} + + def choose_group(query: int, negative_label: int) -> tuple[np.ndarray, int]: + exact_vj = by_label_species_length_vj.get(( + negative_label, species[query], int(lengths[query]), v_call[query], j_call[query] + )) + if exact_vj is not None: + return exact_vj, 0 + one_gene = [ + group for group in ( + by_label_species_length_v.get(( + negative_label, species[query], int(lengths[query]), v_call[query] + )), + by_label_species_length_j.get(( + negative_label, species[query], int(lengths[query]), j_call[query] + )), + ) if group is not None + ] + if one_gene: + return min(one_gene, key=len), 1 + exact = by_label_species_length.get( + (negative_label, species[query], int(lengths[query])) + ) + if exact is not None: + return exact, 2 + same_species = by_label_species.get((negative_label, species[query])) + if same_species: + nearest_length = min( + same_species, key=lambda value: (abs(value - int(lengths[query])), value) + ) + return same_species[nearest_length], 3 + available_lengths = sorted( + length for (label, length) in by_label_length if label == negative_label + ) + nearest_length = min( + available_lengths, key=lambda value: (abs(value - int(lengths[query])), value) + ) + return by_label_length[(negative_label, nearest_length)], 4 + + for query in range(len(cohort)): + rng = np.random.default_rng(seed + query * 1_000_003) + positive_pool = by_label[int(labels[query])] + positive_pool = positive_pool[positive_pool != query] + n_positive = min(max_positives, len(positive_pool)) + if n_positive < 1: + raise ValueError(f"Query {query} has no same-epitope candidate.") + positives = ( + positive_pool if len(positive_pool) == n_positive + else positive_pool[np.sort(rng.choice(len(positive_pool), n_positive, replace=False))] + ) + candidates[query, :n_positive] = positives + relevant[query, :n_positive] = True + tiers[query, :n_positive] = 0 + + quota = negative_ratio * n_positive + other_labels = np.delete(np.arange(len(label_names), dtype=np.int32), labels[query]) + selected: list[int] = [] + selected_tiers: list[int] = [] + if quota > len(other_labels): + raise ValueError("Negative quota exceeds the number of other epitope labels.") + for negative_label in rng.permutation(other_labels)[:quota]: + group, tier = choose_group(query, int(negative_label)) + selected.append(int(group[rng.integers(len(group))])) + selected_tiers.append(tier) + start = n_positive + stop = start + quota + candidates[query, start:stop] = selected + tiers[query, start:stop] = selected_tiers + for tier in selected_tiers: + name = { + 0: "same_species_length_vj", 1: "same_species_length_one_gene", + 2: "same_species_length", 3: "same_species_nearest_length", + 4: "nearest_length_any_species", + }[tier] + tier_counts[name] = tier_counts.get(name, 0) + 1 + return candidates, relevant, tiers, { + "queries": len(cohort), "width": width, "max_positives": max_positives, + "negative_ratio": negative_ratio, "selection_seed": seed, + "negative_match_tiers": dict(sorted(tier_counts.items())), + "label_names": label_names.tolist(), + } + + +def candidate_distances( + features: np.ndarray, candidates: np.ndarray, metric: str, chunk_size: int = 4096, +) -> tuple[np.ndarray, dict]: + values = np.asarray(features, dtype=np.float32) + norms = np.linalg.norm(values, axis=1) + if not np.isfinite(values).all() or np.any(norms == 0): + raise ValueError("Features are non-finite or contain a zero-norm row.") + output = np.full(candidates.shape, np.inf, dtype=np.float32) + for start in range(0, len(values), chunk_size): + stop = min(start + chunk_size, len(values)) + ids = candidates[start:stop] + valid = ids >= 0 + safe = np.maximum(ids, 0) + query = values[start:stop, None, :] + gallery = values[safe] + if metric == "cosine": + query = query / norms[start:stop, None, None] + gallery = gallery / norms[safe, None] + block = 1.0 - np.sum(query * gallery, axis=2) + elif metric == "euclidean": + block = np.linalg.norm(query - gallery, axis=2) + else: + raise ValueError(metric) + output[start:stop][valid] = np.maximum(block[valid], 0.0) + return output, { + "latent_norm_min": float(norms.min()), "latent_norm_median": float(np.median(norms)), + "latent_norm_max": float(norms.max()), + "euclidean_scale_note": "raw latent Euclidean distance; compare retrieval/effect, not absolute scale", + } + + +def _majority_label(labels: np.ndarray, distances: np.ndarray) -> int: + unique, counts = np.unique(labels, return_counts=True) + winners = unique[counts == counts.max()] + if len(winners) == 1: + return int(winners[0]) + first_rank = {int(label): int(np.flatnonzero(labels == label)[0]) for label in winners} + return min((int(label) for label in winners), key=lambda value: (first_rank[value], value)) + + +def summarize_candidates( + distances: np.ndarray, + candidates: np.ndarray, + relevant: np.ndarray, + label_ids: np.ndarray, + label_names: list[str], +) -> tuple[dict, pd.DataFrame, pd.DataFrame]: + rows = [] + predictions = {1: np.empty(len(distances), dtype=np.int32), + 5: np.empty(len(distances), dtype=np.int32), + 10: np.empty(len(distances), dtype=np.int32)} + for query in range(len(distances)): + valid = candidates[query] >= 0 + order = np.argsort(distances[query, valid], kind="mergesort") + hits = relevant[query, valid][order] + ranked_candidates = candidates[query, valid][order] + ranked_labels = label_ids[ranked_candidates] + positive_count = int(hits.sum()) + precision_curve = np.cumsum(hits) / np.arange(1, len(hits) + 1) + within = distances[query, relevant[query]] + between = distances[query, valid & ~relevant[query]] + row = { + "query_index": query, "label": label_names[int(label_ids[query])], + "positive_candidates": positive_count, + "negative_candidates": int(len(hits) - positive_count), + "precision_at_1": float(hits[:1].mean()), + "precision_at_5": float(hits[:5].mean()), + "precision_at_10": float(hits[:10].mean()), + "average_precision": float((precision_curve * hits).sum() / positive_count), + "within_mean": float(within.mean()), "between_mean": float(between.mean()), + "within_between_ratio": float(within.mean() / between.mean()), + "between_minus_within": float(between.mean() - within.mean()), + "nearest_same": float(within.min()), "nearest_other": float(between.min()), + "nearest_margin": float(between.min() - within.min()), + "nearest_same_is_closer": float(within.min() < between.min()), + } + rows.append(row) + for k in predictions: + predictions[k][query] = _majority_label(ranked_labels[:k], distances[query, valid][order][:k]) + per_query = pd.DataFrame(rows) + per_label = per_query.groupby("label", sort=True)[list(METRIC_COLUMNS) + [ + "nearest_same_is_closer", "within_mean", "between_mean" + ]].mean().reset_index() + summary = { + "queries": len(per_query), + "macro_epitope": {column: float(per_label[column].mean()) for column in METRIC_COLUMNS}, + "micro_query": {column: float(per_query[column].mean()) for column in METRIC_COLUMNS}, + "nearest_same_is_closer_macro": float(per_label["nearest_same_is_closer"].mean()), + "knn_classification": {}, + } + for k, prediction in predictions.items(): + precision, recall, f1, support = precision_recall_fscore_support( + label_ids, prediction, labels=np.arange(len(label_names)), zero_division=0 + ) + summary["knn_classification"][f"k_{k}"] = { + "macro_f1": float(f1_score(label_ids, prediction, average="macro")), + "weighted_f1": float(f1_score(label_ids, prediction, average="weighted")), + "balanced_accuracy": float(balanced_accuracy_score(label_ids, prediction)), + } + f1_lookup = {label_names[index]: float(f1[index]) for index in range(len(label_names))} + per_label[f"knn_f1_at_{k}"] = per_label["label"].map(f1_lookup) + return summary, per_query, per_label + + +def clustering_metrics(features: np.ndarray, label_ids: np.ndarray, seed: int, sample: int) -> dict: + values = np.asarray(features, dtype=np.float64) + unit = values / np.linalg.norm(values, axis=1, keepdims=True) + n_labels = len(np.unique(label_ids)) + kmeans = MiniBatchKMeans( + n_clusters=n_labels, random_state=seed, n_init=10, batch_size=4096, + max_iter=300, reassignment_ratio=0.01, + ).fit(unit) + predicted = kmeans.labels_ + sample_size = min(sample, len(unit)) + return { + "true_label_geometry": { + "silhouette_cosine": float(silhouette_score( + unit, label_ids, metric="cosine", sample_size=sample_size, random_state=seed + )), + "silhouette_euclidean_unit": float(silhouette_score( + unit, label_ids, metric="euclidean", sample_size=sample_size, random_state=seed + )), + "davies_bouldin_unit_lower_is_better": float(davies_bouldin_score(unit, label_ids)), + "calinski_harabasz_unit_higher_is_better": float(calinski_harabasz_score(unit, label_ids)), + }, + "label_free_minibatch_kmeans_k_equals_epitopes": { + "adjusted_rand": float(adjusted_rand_score(label_ids, predicted)), + "adjusted_mutual_information": float(adjusted_mutual_info_score(label_ids, predicted)), + "v_measure": float(v_measure_score(label_ids, predicted)), + "homogeneity": float(homogeneity_score(label_ids, predicted)), + "completeness": float(completeness_score(label_ids, predicted)), + "inertia": float(kmeans.inertia_), "n_init": 10, "seed": seed, + }, + "interpretation": ( + "Secondary: epitopes can contain multiple REDCEA motifs, so a one-cluster-per-epitope " + "K-means geometry is not assumed by the primary retrieval/contrast analysis." + ), + } + + +def baseline_features(cohort: pd.DataFrame) -> dict[str, np.ndarray]: + amino = "ACDEFGHIKLMNPQRSTVWY" + composition = np.zeros((len(cohort), len(amino)), dtype=np.float32) + for row, sequence in enumerate(cohort["cdr3"].astype(str)): + for column, aa in enumerate(amino): + composition[row, column] = sequence.count(aa) / len(sequence) + length = cohort["length"].to_numpy(np.float32)[:, None] + length = (length - length.mean()) / max(float(length.std()), 1e-8) + basic = np.concatenate((composition, length), axis=1) + genes = pd.get_dummies( + cohort[["v_call", "j_call"]].fillna("").astype(str), dtype=np.float32 + ).to_numpy(np.float32) + return {"aa_composition_length": basic, "aa_composition_length_vj": np.concatenate((basic, genes), axis=1)} + + +def bootstrap_by_epitope(per_label: pd.DataFrame, replicates: int, seed: int) -> dict: + labels = sorted(per_label["label"].unique()) + rng = np.random.default_rng(seed) + samples = rng.integers(0, len(labels), size=(replicates, len(labels))) + report = {"replicates": replicates, "unit": "epitope label", "paired_rtp_minus_other": {}} + for distance in ("cosine", "euclidean"): + block = per_label[per_label["distance"] == distance] + arrays = { + model: block[block["model"] == model].set_index("label").loc[labels] + for model in PRIMARY_MODELS + } + report["paired_rtp_minus_other"][distance] = {} + for other in PRIMARY_MODELS[:-1]: + metrics = {} + for column in METRIC_COLUMNS: + delta = arrays["rtp"][column].to_numpy() - arrays[other][column].to_numpy() + boot = delta[samples].mean(axis=1) + metrics[column] = { + "estimate": float(delta.mean()), + "ci95": [float(np.percentile(boot, 2.5)), float(np.percentile(boot, 97.5))], + "direction": "lower is better" if column == "within_between_ratio" else "higher is better", + } + report["paired_rtp_minus_other"][distance][other] = metrics + return report + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--preflight", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--model", action="append", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--max-positive-candidates", type=int, default=20) + parser.add_argument("--negative-ratio", type=int, default=4) + parser.add_argument("--bootstrap", type=int, default=2000) + parser.add_argument("--silhouette-sample", type=int, default=10000) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def main() -> None: + from benchmark.evaluate_vdjdb_frozen import ( + extract_embeddings, parse_models, validate_run_configs, + ) + + args = parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + preflight = json.loads(args.preflight.read_text(encoding="utf-8")) + if preflight["status"] != "accepted_for_frozen_embedding_evaluation": + raise ValueError("Cohort preflight is not accepting.") + if sha256(args.cohort) != preflight["cohort"]["sha256"]: + raise ValueError("Cohort differs from accepted preflight.") + cohort = pd.read_csv(args.cohort, sep="\t") + if cohort["cdr3"].duplicated().any() or (cohort.groupby("label").size() < 20).any(): + raise ValueError("Cohort violates uniqueness or minimum-support rules.") + models = parse_models(args.model) + if tuple(sorted(models)) != tuple(sorted(PRIMARY_MODELS)): + raise ValueError("The full seven-condition factorial is required.") + config, model_report = validate_run_configs(models) + embeddings, extraction = extract_embeddings( + cohort, models, config, args.tokenizer, args.batch_size, args.output_dir + ) + candidates, relevant, tiers, gallery_report = build_candidate_gallery( + cohort, args.max_positive_candidates, args.negative_ratio, args.seed + ) + gallery_path = args.output_dir / "candidate_gallery.npz" + np.savez_compressed(gallery_path, candidates=candidates, relevant=relevant, match_tier=tiers) + label_names = gallery_report["label_names"] + label_to_id = {label: index for index, label in enumerate(label_names)} + label_ids = np.array([label_to_id[value] for value in cohort["label"].astype(str)], dtype=np.int32) + + summaries = {} + clustering = {} + per_labels = [] + for model in PRIMARY_MODELS: + summaries[model] = {} + for metric in ("cosine", "euclidean"): + distances, norm_report = candidate_distances(embeddings[model], candidates, metric) + summary, per_query, per_label = summarize_candidates( + distances, candidates, relevant, label_ids, label_names + ) + summary["embedding_norms"] = norm_report + summaries[model][metric] = summary + per_query.to_parquet(args.output_dir / f"per_query_{model}_{metric}.parquet", index=False) + per_label.insert(0, "model", model) + per_label.insert(1, "distance", metric) + per_labels.append(per_label) + clustering[model] = clustering_metrics( + embeddings[model], label_ids, args.seed, args.silhouette_sample + ) + + controls = {} + for name, features in baseline_features(cohort).items(): + controls[name] = {"retrieval": {}, "clustering": clustering_metrics( + features, label_ids, args.seed, args.silhouette_sample + )} + for metric in ("cosine", "euclidean"): + distances, _ = candidate_distances(features, candidates, metric) + summary, _, _ = summarize_candidates( + distances, candidates, relevant, label_ids, label_names + ) + controls[name]["retrieval"][metric] = summary + + per_label_table = pd.concat(per_labels, ignore_index=True) + per_label_path = args.output_dir / "per_epitope_metrics.tsv" + per_label_table.to_csv(per_label_path, sep="\t", index=False, lineterminator="\n") + bootstrap = bootstrap_by_epitope(per_label_table, args.bootstrap, args.seed) + chance = { + "candidate_positive_fraction": float(relevant.sum() / (candidates >= 0).sum()), + "expected_precision_at_k_under_random_ranking": float( + relevant.sum() / (candidates >= 0).sum() + ), + "note": "Candidate prevalence is fixed at 1:4 positive:negative; MAP chance varies slightly with gallery size.", + } + assert_finite_tree(summaries, "summaries") + assert_finite_tree(clustering, "clustering") + assert_finite_tree(controls, "controls") + assert_finite_tree(bootstrap, "bootstrap") + stable_json(args.output_dir / "metrics.json", summaries) + stable_json(args.output_dir / "clustering.json", clustering) + stable_json(args.output_dir / "controls.json", controls) + stable_json(args.output_dir / "bootstrap.json", bootstrap) + artifact_hashes = { + path.name: sha256(path) for path in sorted(args.output_dir.iterdir()) if path.is_file() + } + result = { + "status": "complete", "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "study": "REDCEA motif-member per-clonotype geometry across seven matched DATA-ANCHOR encoders", + "cohort": {"path": str(args.cohort), "sha256": sha256(args.cohort), + "clonotypes": len(cohort), "epitopes": int(cohort["label"].nunique())}, + "models": model_report, "extraction": extraction, + "candidate_gallery": {**gallery_report, "path": str(gallery_path), "sha256": sha256(gallery_path)}, + "chance_control": chance, + "primary_metrics": ["macro P@1/5/10", "macro MAP over fixed candidate gallery", + "within/between distance ratio", "nearest-same margin", "kNN macro-F1"], + "secondary_metrics": ["silhouette", "Davies-Bouldin", "Calinski-Harabasz", + "MiniBatchKMeans ARI/AMI/V-measure"], + "selection_statement": "No VDJdb outcome selected a checkpoint, model, gallery member, or support threshold.", + "circularity_caveat": ( + "REDCEA membership is sequence/TCRemP-cluster-conditioned. This measures organization of " + "REDCEA-supported epitope labels, not independent binding prediction beyond sequence motifs." + ), + "single_seed_caveat": ( + "Paired intervals quantify cohort/epitope uncertainty for fixed seed-42 models, not training-seed variability." + ), + "artifact_sha256": artifact_hashes, "nonfinite_scan": "passed", + } + stable_json(args.output_dir / "RESULTS.json", result) + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_vdjdb_uniform100.py b/benchmark/evaluate_vdjdb_uniform100.py new file mode 100644 index 0000000..dc0c9d8 --- /dev/null +++ b/benchmark/evaluate_vdjdb_uniform100.py @@ -0,0 +1,163 @@ +"""Uniform 100-within / 100-between cosine-distance evaluation for Section 7. + +Only epitopes with at least 101 unique clonotypes are retained. For every +retained query, 100 distinct same-epitope comparators and 100 distinct +different-epitope comparators are drawn without replacement. The candidate +pairs are deterministic and shared by all representations. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd + + +N_WITHIN = 100 +N_BETWEEN = 100 +SEED = 42 + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--native-evaluation", type=Path, required=True) + parser.add_argument("--external-representations", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--seed", type=int, default=SEED) + return parser.parse_args() + + +def build_pairs(labels: np.ndarray, seed: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + counts = pd.Series(labels).value_counts() + eligible_labels = np.array(sorted(counts[counts >= N_WITHIN + 1].index), dtype=object) + keep = np.isin(labels, eligible_labels) + original_indices = np.flatnonzero(keep) + retained_labels = labels[keep] + groups = {label: np.flatnonzero(retained_labels == label) for label in eligible_labels} + all_indices = np.arange(len(retained_labels), dtype=np.int32) + within = np.empty((len(retained_labels), N_WITHIN), dtype=np.int32) + between = np.empty((len(retained_labels), N_BETWEEN), dtype=np.int32) + + for query in range(len(retained_labels)): + rng = np.random.default_rng(seed + int(original_indices[query]) * 1_000_003) + label = retained_labels[query] + same = groups[label] + same = same[same != query] + within[query] = rng.choice(same, size=N_WITHIN, replace=False) + other = all_indices[retained_labels != label] + between[query] = rng.choice(other, size=N_BETWEEN, replace=False) + return original_indices, within, between + + +def cosine_summary( + matrix: np.ndarray, within: np.ndarray, between: np.ndarray, + labels: np.ndarray, batch_size: int = 256, +) -> tuple[pd.DataFrame, pd.DataFrame]: + values = np.asarray(matrix, dtype=np.float32) + norms = np.linalg.norm(values, axis=1, keepdims=True) + if np.any(norms == 0) or not np.isfinite(values).all(): + raise ValueError("Embedding matrix contains a zero-norm or non-finite vector.") + unit = values / norms + within_mean = np.empty(len(unit), dtype=np.float32) + between_mean = np.empty(len(unit), dtype=np.float32) + for start in range(0, len(unit), batch_size): + stop = min(start + batch_size, len(unit)) + query = unit[start:stop] + within_mean[start:stop] = (1.0 - np.sum(query[:, None] * unit[within[start:stop]], axis=2)).mean(axis=1) + between_mean[start:stop] = (1.0 - np.sum(query[:, None] * unit[between[start:stop]], axis=2)).mean(axis=1) + per_query = pd.DataFrame({ + "epitope": labels, + "within_mean": within_mean, + "between_mean": between_mean, + "within_between_ratio": within_mean / between_mean, + "between_minus_within": between_mean - within_mean, + }) + per_epitope = per_query.groupby("epitope", sort=True).agg( + queries=("epitope", "size"), + within_mean=("within_mean", "mean"), + between_mean=("between_mean", "mean"), + within_between_ratio=("within_between_ratio", "mean"), + between_minus_within=("between_minus_within", "mean"), + ).reset_index() + return per_query, per_epitope + + +def main() -> None: + args = parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + cohort = pd.read_csv(args.cohort, sep="\t") + labels = cohort["label"].astype(str).to_numpy() + original_indices, within, between = build_pairs(labels, args.seed) + retained = cohort.iloc[original_indices].reset_index(drop=True) + np.savez_compressed(args.output_dir / "uniform_pairs.npz", within=within, between=between, original_indices=original_indices) + + native = { + "RTP": args.native_evaluation / "embeddings_rtp.npy", + "RP": args.native_evaluation / "embeddings_rp.npy", + "T": args.native_evaluation / "embeddings_t.npy", + } + catalog = json.loads((args.external_representations / "representations.json").read_text(encoding="utf-8")) + external = { + "ESM2-35M": Path(catalog["esm2_35m"]["path"]), + "TCR-BERT": Path(catalog["tcr_bert"]["path"]), + "SCEPTR (CDR3-only)": Path(catalog["sceptr_cdr3"]["path"]), + } + summary_rows = [] + all_epitopes = [] + for name, path in {**native, **external}.items(): + matrix = np.load(path, mmap_mode="r")[original_indices] + per_query, per_epitope = cosine_summary(matrix, within, between, retained["label"].astype(str).to_numpy()) + per_query.insert(0, "model", name) + per_query.to_parquet(args.output_dir / f"per_query_{name.lower().replace(' ', '_').replace('(', '').replace(')', '')}.parquet", index=False) + per_epitope.insert(0, "model", name) + all_epitopes.append(per_epitope) + summary_rows.append({ + "model": name, + "macro_within_between_ratio": float(per_epitope["within_between_ratio"].mean()), + "macro_between_minus_within": float(per_epitope["between_minus_within"].mean()), + "micro_within_between_ratio": float(per_query["within_mean"].mean() / per_query["between_mean"].mean()), + "queries": len(per_query), + "epitopes": len(per_epitope), + }) + pd.DataFrame(summary_rows).to_csv(args.output_dir / "summary.tsv", sep="\t", index=False, lineterminator="\n") + pd.concat(all_epitopes, ignore_index=True).to_csv(args.output_dir / "per_epitope.tsv", sep="\t", index=False, lineterminator="\n") + result = { + "status": "complete", + "protocol": { + "metric": "cosine distance", + "within_candidates_per_query": N_WITHIN, + "between_candidates_per_query": N_BETWEEN, + "within_sampling": "uniform without replacement from same epitope, excluding query", + "between_sampling": "uniform without replacement from all different-epitope clonotypes", + "pair_seed": args.seed, + "query_seed": "seed + original_cohort_index * 1,000,003", + }, + "cohort": { + "source": str(args.cohort), "source_sha256": sha256(args.cohort), + "retained_clonotypes": len(retained), "retained_epitopes": int(retained["label"].nunique()), + "minimum_epitope_size": N_WITHIN + 1, + }, + "models": list(native) + list(external), + "artifacts": {name: sha256(args.output_dir / name) for name in ("uniform_pairs.npz", "summary.tsv", "per_epitope.tsv")}, + } + (args.output_dir / "RESULTS.json").write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(result, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_vdjdb_uniform100_separation.py b/benchmark/evaluate_vdjdb_uniform100_separation.py new file mode 100644 index 0000000..c2e1705 --- /dev/null +++ b/benchmark/evaluate_vdjdb_uniform100_separation.py @@ -0,0 +1,93 @@ +"""AUROC and standardized distance separation on the fixed uniform-100 pairs.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import pandas as pd + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--native-evaluation", type=Path, required=True) + parser.add_argument("--external-representations", type=Path, required=True) + parser.add_argument("--uniform-evaluation", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + return parser.parse_args() + + +def separation_metrics( + matrix: np.ndarray, within: np.ndarray, between: np.ndarray, labels: np.ndarray, + batch_size: int = 256, +) -> tuple[pd.DataFrame, pd.DataFrame]: + values = np.asarray(matrix, dtype=np.float32) + norms = np.linalg.norm(values, axis=1, keepdims=True) + if np.any(norms == 0) or not np.isfinite(values).all(): + raise ValueError("Embedding matrix contains a zero-norm or non-finite vector.") + unit = values / norms + n_within, n_between = within.shape[1], between.shape[1] + auroc = np.empty(len(unit), dtype=np.float32) + cohen_d = np.empty(len(unit), dtype=np.float32) + for start in range(0, len(unit), batch_size): + stop = min(start + batch_size, len(unit)) + query = unit[start:stop] + d_within = 1.0 - np.sum(query[:, None] * unit[within[start:stop]], axis=2) + d_between = 1.0 - np.sum(query[:, None] * unit[between[start:stop]], axis=2) + ordered = np.argsort(np.concatenate([d_within, d_between], axis=1), axis=1) + ranks = np.empty_like(ordered) + ranks[np.arange(stop-start)[:, None], ordered] = np.arange(1, n_within+n_between+1) + positive_rank_sum = ranks[:, :n_within].sum(axis=1) + auroc[start:stop] = ( + n_within * n_between - (positive_rank_sum - n_within * (n_within + 1) / 2) + ) / (n_within * n_between) + pooled_sd = np.sqrt((d_within.var(axis=1, ddof=1) + d_between.var(axis=1, ddof=1)) / 2) + cohen_d[start:stop] = (d_between.mean(axis=1) - d_within.mean(axis=1)) / pooled_sd + per_query = pd.DataFrame({"epitope": labels, "auroc": auroc, "cohen_d": cohen_d}) + per_epitope = per_query.groupby("epitope", sort=True).agg( + queries=("epitope", "size"), auroc=("auroc", "mean"), cohen_d=("cohen_d", "mean"), + ).reset_index() + return per_query, per_epitope + + +def main() -> None: + args = parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + cohort = pd.read_csv(args.cohort, sep="\t") + pairs = np.load(args.uniform_evaluation / "uniform_pairs.npz") + original = pairs["original_indices"] + within, between = pairs["within"], pairs["between"] + labels = cohort.iloc[original]["label"].astype(str).to_numpy() + catalog = json.loads((args.external_representations / "representations.json").read_text(encoding="utf-8")) + paths = { + "RTP": args.native_evaluation / "embeddings_rtp.npy", + "RP": args.native_evaluation / "embeddings_rp.npy", + "T": args.native_evaluation / "embeddings_t.npy", + "ESM2-35M": Path(catalog["esm2_35m"]["path"]), + "TCR-BERT": Path(catalog["tcr_bert"]["path"]), + "SCEPTR (CDR3-only)": Path(catalog["sceptr_cdr3"]["path"]), + } + summary, per_epitope = [], [] + for name, path in paths.items(): + query, epitope = separation_metrics(np.load(path, mmap_mode="r")[original], within, between, labels) + query.insert(0, "model", name) + query.to_parquet(args.output_dir / f"per_query_{name.lower().replace(' ', '_').replace('(', '').replace(')', '')}.parquet", index=False) + epitope.insert(0, "model", name) + per_epitope.append(epitope) + summary.append({"model": name, "macro_auroc": float(epitope["auroc"].mean()), "macro_cohen_d": float(epitope["cohen_d"].mean()), "queries":len(query), "epitopes":len(epitope)}) + pd.DataFrame(summary).to_csv(args.output_dir / "summary.tsv", sep="\t", index=False, lineterminator="\n") + pd.concat(per_epitope, ignore_index=True).to_csv(args.output_dir / "per_epitope.tsv", sep="\t", index=False, lineterminator="\n") + (args.output_dir / "RESULTS.json").write_text(json.dumps({ + "status":"complete", "pairs":str(args.uniform_evaluation / "uniform_pairs.npz"), + "metrics":{"auroc":"P(within distance < between distance)", "cohen_d":"(mean between - mean within) / pooled within/between distance SD"}, + "queries":len(labels), "epitopes":int(pd.Series(labels).nunique()), "models":list(paths), + }, indent=2, sort_keys=True)+"\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_vdjdb_viral_motif_top10.py b/benchmark/evaluate_vdjdb_viral_motif_top10.py new file mode 100644 index 0000000..ad869d0 --- /dev/null +++ b/benchmark/evaluate_vdjdb_viral_motif_top10.py @@ -0,0 +1,425 @@ +"""Evaluate seven frozen DATA-ANCHOR head combinations on ten viral epitopes.""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import numpy as np +import pandas as pd + +from benchmark.evaluate_vdjdb_frozen import ( + assert_finite_tree, + extract_embeddings, + sha256, + stable_json, +) +from benchmark.evaluate_vdjdb_ylq_motif_pooled import ( + bootstrap_clonotypes, + composition_features, + pair_metrics, + pooled_retrieval, + primary_distances, + shuffled_label_control, + vector_distance, + vj_pair_values, +) +from irrm_codec.multitask_transformer import IRRMCodecConfig + + +MODEL_NAMES = ("r", "t", "p", "rt", "rp", "tp", "rtp") +EXPECTED_WEIGHTS = { + "r": (1.0, 0.0, 0.0), + "t": (0.0, 1.0, 0.0), + "p": (0.0, 0.0, 1.0), + "rt": (1.0, 1.0, 0.0), + "rp": (1.0, 0.0, 1.0), + "tp": (0.0, 1.0, 1.0), + "rtp": (1.0, 1.0, 1.0), +} +PARAMETER_COUNT = 21_535_913 + + +def parse_seven_models(values: list[str]) -> dict[str, tuple[Path, Path]]: + result = {} + for value in values: + fields = value.split(":", 2) + if len(fields) != 3: + raise ValueError(f"Invalid --model value: {value!r}") + name, checkpoint, config = fields + if name in result: + raise ValueError(f"Duplicate model name: {name}") + result[name] = (Path(checkpoint), Path(config)) + if set(result) != set(MODEL_NAMES): + raise ValueError(f"Exact required model names are {', '.join(MODEL_NAMES)}.") + return {name: result[name] for name in MODEL_NAMES} + + +def validate_seven_configs( + models: dict[str, tuple[Path, Path]], +) -> tuple[IRRMCodecConfig, dict]: + configs = { + name: json.loads(path.read_text(encoding="utf-8")) + for name, (_, path) in models.items() + } + reference = configs["rtp"] + ignored = {"output_dir", "reconstruction_loss_weight", "tcremp_loss_weight", "pgen_loss_weight"} + reference_training = { + key: value for key, value in reference["training"].items() if key not in ignored + } + report = {} + for name in MODEL_NAMES: + payload = configs[name] + if payload["model"] != reference["model"]: + raise ValueError(f"{name} architecture differs from RTP.") + comparable = { + key: value for key, value in payload["training"].items() if key not in ignored + } + if comparable != reference_training: + differing = sorted( + key for key in set(comparable) | set(reference_training) + if comparable.get(key) != reference_training.get(key) + ) + raise ValueError(f"{name} has non-weight training differences: {differing}") + weights = ( + float(payload["training"]["reconstruction_loss_weight"]), + float(payload["training"]["tcremp_loss_weight"]), + float(payload["training"]["pgen_loss_weight"]), + ) + if weights != EXPECTED_WEIGHTS[name]: + raise ValueError(f"{name} has weights {weights}; expected {EXPECTED_WEIGHTS[name]}.") + if int(payload["parameter_count"]) != PARAMETER_COUNT: + raise ValueError(f"{name} parameter count is not {PARAMETER_COUNT}.") + report[name] = { + "checkpoint": str(models[name][0]), + "checkpoint_sha256": sha256(models[name][0]), + "run_config": str(models[name][1]), + "run_config_sha256": sha256(models[name][1]), + "loss_weights_reconstruction_tcremp_pgen": list(weights), + "parameter_count": int(payload["parameter_count"]), + } + return IRRMCodecConfig(**reference["model"]), report + + +def validate_epitope_tables(directory: Path, report: dict) -> tuple[pd.DataFrame, pd.DataFrame]: + candidates_path = directory / "candidates.tsv" + pairs_path = directory / "pair_manifest.tsv" + if sha256(candidates_path) != report["artifacts"]["candidates_sha256"] \ + or sha256(pairs_path) != report["artifacts"]["pairs_sha256"]: + raise ValueError(f"Preflight artifacts changed in {directory}.") + candidates = pd.read_csv(candidates_path, sep="\t") + pairs = pd.read_csv(pairs_path, sep="\t") + if candidates["cdr3"].duplicated().any(): + raise ValueError(f"Duplicate CDR3 in {directory}.") + if set(candidates["mhc_a"]) != {"HLA-A*02:01"} \ + or set(candidates["mhc_b"]) != {"B2M"} \ + or set(candidates["mhc_class"]) != {"MHCI"}: + raise ValueError(f"Context mismatch in {directory}.") + if candidates["motif_cids"].fillna("").eq("").any(): + raise ValueError(f"Missing official motif membership in {directory}.") + labels = candidates["is_target"].to_numpy(int) + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + if not (np.all(labels[q] == 1) and np.all(labels[p] == 1) and np.all(labels[n] == 0)): + raise ValueError(f"Invalid target/negative triplets in {directory}.") + if not np.all(candidates["length"].to_numpy(int)[p] + == candidates["length"].to_numpy(int)[n]): + raise ValueError(f"Non-exact-length negative comparator in {directory}.") + return candidates, pairs + + +def summarize_distribution(rows: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + summary_rows = [] + win_rows = [] + for (distance, metric, model), group in rows.groupby(["distance", "metric", "model"]): + values = group["value"].to_numpy(float) + summary_rows.append({ + "distance": distance, + "metric": metric, + "model": model, + "epitopes": len(values), + "median": float(np.median(values)), + "q1": float(np.percentile(values, 25)), + "q3": float(np.percentile(values, 75)), + }) + lower_better = {"within_between_ratio"} + for (distance, metric, epitope), group in rows.groupby(["distance", "metric", "epitope"]): + target = group["value"].min() if metric in lower_better else group["value"].max() + winners = sorted(group.loc[np.isclose(group["value"], target), "model"].tolist()) + for model in winners: + win_rows.append({ + "distance": distance, + "metric": metric, + "epitope": epitope, + "model": model, + "tied_winners": len(winners), + "fractional_win": 1.0 / len(winners), + }) + wins = pd.DataFrame(win_rows) + win_summary = wins.groupby(["distance", "metric", "model"], as_index=False).agg( + win_count=("fractional_win", "sum"), + epitopes_with_best_or_tied=("epitope", "count"), + ) + complete = pd.MultiIndex.from_product( + [sorted(rows["distance"].unique()), sorted(rows["metric"].unique()), MODEL_NAMES], + names=["distance", "metric", "model"], + ).to_frame(index=False) + win_summary = complete.merge( + win_summary, on=["distance", "metric", "model"], how="left" + ).fillna({"win_count": 0.0, "epitopes_with_best_or_tied": 0}) + win_summary["epitopes_with_best_or_tied"] = win_summary[ + "epitopes_with_best_or_tied" + ].astype(int) + return pd.DataFrame(summary_rows), wins, win_summary + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--preflight-dir", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--model", action="append", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--bootstrap", type=int, default=2000) + parser.add_argument("--permutations", type=int, default=1000) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + + preflight = json.loads((args.preflight_dir / "PREFLIGHT.json").read_text(encoding="utf-8")) + expected_epitopes = int(preflight["support_gate"]["selected_epitopes"]) + if preflight["status"] != "ready" or len(preflight["selected"]) != expected_epitopes \ + or expected_epitopes < 2 or expected_epitopes > 10: + raise ValueError("Up-to-top-10 viral motif preflight is not accepting.") + master_path = args.preflight_dir / "master_candidates.tsv" + if sha256(master_path) != preflight["artifacts"]["master_candidates_sha256"]: + raise ValueError("Master candidate hash changed after CPU gate.") + master = pd.read_csv(master_path, sep="\t") + if master["cdr3"].duplicated().any(): + raise ValueError("Master candidates are not unique by exact CDR3.") + + models = parse_seven_models(args.model) + config, model_report = validate_seven_configs(models) + embeddings, extraction = extract_embeddings( + master, models, config, args.tokenizer, args.batch_size, args.output_dir + ) + master_index = {cdr3: index for index, cdr3 in enumerate(master["cdr3"].astype(str))} + metric_rows = [] + control_rows = [] + epitope_results = [] + paired_effect_rows = [] + for selected in preflight["selected"]: + rank = int(selected["rank"]) + epitope = selected["epitope"] + source_dir = args.preflight_dir / f"epitope-{rank:02d}-{epitope}" + target_dir = args.output_dir / f"epitope-{rank:02d}-{epitope}" + target_dir.mkdir() + candidates, pairs = validate_epitope_tables(source_dir, selected) + candidates = candidates.copy() + candidates["is_ylq"] = candidates["is_target"] + indices = np.asarray([master_index[value] for value in candidates["cdr3"].astype(str)]) + subset_embeddings = {name: array[indices] for name, array in embeddings.items()} + metrics = {"cosine": {}, "euclidean": {}} + bootstraps = {} + cosine_values = {} + for distance in ("cosine", "euclidean"): + pair_values = {} + retrieval_tables = {} + for name, array in subset_embeddings.items(): + within, between = primary_distances(array, pairs, distance) + pair_values[name] = (within, between) + retrieval, retrieval_table = pooled_retrieval(array, candidates, distance) + retrieval_tables[name] = retrieval_table + pair_result = pair_metrics(within, between) + metrics[distance][name] = { + "pairwise_and_contrast": pair_result, + "retrieval": retrieval, + } + values = { + "within_between_ratio": pair_result["within_between_ratio"], + "between_minus_within": pair_result["between_minus_within"], + "same_label_auroc": pair_result["same_label_auroc"], + "same_label_auprc_balanced": pair_result["same_label_auprc_balanced"], + "precision_at_1": retrieval["precision_at_1"], + "map_at_r": retrieval["map_at_r"], + } + for metric, value in values.items(): + metric_rows.append({ + "rank": rank, "epitope": epitope, "distance": distance, + "model": name, "metric": metric, "value": value, + }) + if int(selected["positive_cdr3"]) < 3: + bootstraps[distance] = { + "status": "not_estimable_with_fewer_than_three_target_clonotypes", + "requested_replicates": args.bootstrap, + } + else: + try: + bootstraps[distance] = bootstrap_clonotypes( + pair_values, pairs, retrieval_tables, args.bootstrap, + args.seed + rank * 1000, + ) + bootstraps[distance]["small_n_caveat"] = ( + "descriptive/unstable: fewer than 5 target clonotypes" + if int(selected["positive_cdr3"]) < 5 else "none" + ) + difference = bootstraps[distance][ + "paired_differences_rtp_minus_other" + ]["rp"] + for metric, payload in difference.items(): + paired_effect_rows.append({ + "rank": rank, "epitope": epitope, "distance": distance, + "metric": metric, "rtp_minus_rp_ci95_low": payload["ci95"][0], + "rtp_minus_rp_ci95_high": payload["ci95"][1], + "direction": payload["direction"], + }) + except (ValueError, ZeroDivisionError) as error: + bootstraps[distance] = { + "status": "not_estimable_at_available_target_support", + "reason": str(error), + "requested_replicates": args.bootstrap, + } + if distance == "cosine": + cosine_values = pair_values + + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + lengths = candidates["length"].to_numpy(float) + exact_length_result = pair_metrics( + np.abs(lengths[q] - lengths[p]), np.abs(lengths[q] - lengths[n]) + ) + controls = { + "exact_length": exact_length_result, + "amino_acid_composition": {}, + } + for metric in ("within_between_ratio", "between_minus_within", + "same_label_auroc", "same_label_auprc_balanced"): + control_rows.append({ + "rank": rank, "epitope": epitope, "distance": "absolute", + "control": "exact_cdr3_length", "metric": metric, + "value": exact_length_result[metric], + }) + composition = composition_features(candidates) + for distance in ("cosine", "euclidean"): + result = pair_metrics( + vector_distance(composition[q], composition[p], distance), + vector_distance(composition[q], composition[n], distance), + ) + controls["amino_acid_composition"][distance] = result + for metric in ("within_between_ratio", "between_minus_within", + "same_label_auroc", "same_label_auprc_balanced"): + control_rows.append({ + "rank": rank, "epitope": epitope, "distance": distance, + "control": "amino_acid_composition", "metric": metric, + "value": result[metric], + }) + vj_within, vj_between = vj_pair_values(candidates, pairs) + vj_result = pair_metrics(vj_within, vj_between) + controls["v_j_identity"] = vj_result + for metric in ("within_between_ratio", "between_minus_within", + "same_label_auroc", "same_label_auprc_balanced"): + control_rows.append({ + "rank": rank, "epitope": epitope, "distance": "mismatch_fraction", + "control": "v_j_identity", "metric": metric, "value": vj_result[metric], + }) + try: + controls["shuffled_labels_cosine"] = shuffled_label_control( + cosine_values, candidates, pairs, args.permutations, args.seed + rank * 1000 + ) + except ValueError as error: + controls["shuffled_labels_cosine"] = { + "status": "not_estimable_at_available_target_support", + "reason": str(error), + "requested_permutations": args.permutations, + } + assert_finite_tree(metrics, f"metrics_{epitope}") + assert_finite_tree(bootstraps, f"bootstrap_{epitope}") + assert_finite_tree(controls, f"controls_{epitope}") + stable_json(target_dir / "metrics.json", metrics) + stable_json(target_dir / "bootstrap.json", bootstraps) + stable_json(target_dir / "controls.json", controls) + candidates.drop(columns=["is_ylq"]).to_csv( + target_dir / "candidates.tsv", sep="\t", index=False + ) + pairs.to_csv(target_dir / "pair_manifest.tsv", sep="\t", index=False) + epitope_results.append({ + "rank": rank, + "epitope": epitope, + "positive_cdr3": int(selected["positive_cdr3"]), + "negative_cdr3": int(selected["selected_unique_negative_cdr3"]), + "pairs": len(pairs), + "artifact_sha256": { + path.name: sha256(path) for path in sorted(target_dir.iterdir()) if path.is_file() + }, + }) + + metric_table = pd.DataFrame(metric_rows) + metric_table.to_csv(args.output_dir / "per_epitope_metrics.tsv", sep="\t", index=False) + pd.DataFrame(control_rows).to_csv( + args.output_dir / "per_epitope_controls.tsv", sep="\t", index=False + ) + pd.DataFrame(paired_effect_rows).to_csv( + args.output_dir / "per_epitope_rtp_minus_rp_bootstrap.tsv", sep="\t", index=False + ) + distribution, wins, win_summary = summarize_distribution(metric_table) + distribution.to_csv(args.output_dir / "distribution_summary.tsv", sep="\t", index=False) + wins.to_csv(args.output_dir / "win_assignments.tsv", sep="\t", index=False) + win_summary.to_csv(args.output_dir / "win_summary.tsv", sep="\t", index=False) + point_effects = metric_table.pivot_table( + index=["rank", "epitope", "distance", "metric"], columns="model", values="value" + ).reset_index() + point_effects["rtp_minus_rp"] = point_effects["rtp"] - point_effects["rp"] + point_effects.to_csv(args.output_dir / "rtp_minus_rp_effects.tsv", sep="\t", index=False) + effect_summary_rows = [] + for (distance, metric), group in point_effects.groupby(["distance", "metric"]): + values = group["rtp_minus_rp"].to_numpy(float) + negative_favors_rtp = metric == "within_between_ratio" + effect_summary_rows.append({ + "distance": distance, + "metric": metric, + "epitopes": len(values), + "median_rtp_minus_rp": float(np.median(values)), + "q1_rtp_minus_rp": float(np.percentile(values, 25)), + "q3_rtp_minus_rp": float(np.percentile(values, 75)), + "rtp_win_count": int(np.sum(values < 0 if negative_favors_rtp else values > 0)), + "direction": "negative_favors_rtp" if negative_favors_rtp + else "positive_favors_rtp", + }) + pd.DataFrame(effect_summary_rows).to_csv( + args.output_dir / "rtp_minus_rp_distribution.tsv", sep="\t", index=False + ) + + result = { + "status": "complete", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "models": model_report, + "extraction": extraction, + "epitopes": epitope_results, + "epitope_weighting": "each epitope contributes one point to distribution summaries", + "pooled_pair_estimate_reported": False, + "selection_statement": ( + "Up to ten selected only by post-filter motif-member support; no model metric used. " + f"Actual evaluable count={expected_epitopes}, shortfall=" + f"{preflight['support_gate']['evaluable_epitope_shortfall']}." + ), + "claim_scope": "exploratory within-VDJdb motif/epitope organization; pooled donors", + "circularity_caveat": preflight["circularity_caveat"], + "single_seed": True, + "models_trained_during_evaluation": False, + } + hashes = { + path.name: sha256(path) for path in sorted(args.output_dir.iterdir()) if path.is_file() + } + result["artifact_sha256"] = hashes + stable_json(args.output_dir / "RESULTS.json", result) + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_vdjdb_ylq_frozen.py b/benchmark/evaluate_vdjdb_ylq_frozen.py new file mode 100644 index 0000000..c89026c --- /dev/null +++ b/benchmark/evaluate_vdjdb_ylq_frozen.py @@ -0,0 +1,310 @@ +"""Frozen-encoder evaluation for exact YLQPRTFLL in one dominant MHC context.""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import numpy as np +import pandas as pd +from sklearn.metrics import average_precision_score, roc_auc_score + +from benchmark.evaluate_vdjdb_frozen import ( + AA, + assert_finite_tree, + distance_matrix, + extract_embeddings, + pair_metrics, + parse_models, + sha256, + stable_json, + validate_run_configs, +) + + +def binary_retrieval(matrix: np.ndarray, candidates: pd.DataFrame) -> tuple[dict, pd.DataFrame]: + positive = candidates["is_ylq"].to_numpy(int).astype(bool) + donors = candidates["donor_key"].astype(str).to_numpy() + rows = [] + for query in np.flatnonzero(positive): + allowed = np.flatnonzero((donors != donors[query]) & (np.arange(len(candidates)) != query)) + relevant = positive[allowed] + r = int(relevant.sum()) + if r == 0: + continue + ranked = relevant[np.argsort(matrix[query, allowed], kind="mergesort")] + hits = ranked[:r].astype(float) + precision = np.cumsum(hits) / np.arange(1, r + 1) + rows.append({ + "query_index": int(query), "label": "YLQPRTFLL", "query_donor": donors[query], + "relevant_donor_excluded": r, + "precision_at_1": float(ranked[:1].sum()), + "precision_at_5": float(ranked[:5].sum() / 5), + "precision_at_10": float(ranked[:10].sum() / 10), + "ap_at_r": float((precision * hits).sum() / r), + }) + table = pd.DataFrame(rows) + return { + "eligible_positive_queries": len(table), + "precision_at_1": float(table["precision_at_1"].mean()), + "precision_at_5": float(table["precision_at_5"].mean()), + "precision_at_10": float(table["precision_at_10"].mean()), + "map_at_r": float(table["ap_at_r"].mean()), + }, table + + +def composition_matrix(candidates: pd.DataFrame) -> np.ndarray: + result = np.zeros((len(candidates), len(AA)), dtype=float) + for row, sequence in enumerate(candidates["cdr3"].astype(str)): + for column, aa in enumerate(AA): + result[row, column] = sequence.count(aa) / len(sequence) + return result + + +def vj_distance(candidates: pd.DataFrame) -> np.ndarray: + v = candidates["v_call"].fillna("").astype(str).to_numpy() + j = candidates["j_call"].fillna("").astype(str).to_numpy() + v_comparable = (v[:, None] != "") & (v[None, :] != "") + j_comparable = (j[:, None] != "") & (j[None, :] != "") + v_mismatch = (~v_comparable | (v[:, None] != v[None, :])).astype(float) + j_mismatch = (~j_comparable | (j[:, None] != j[None, :])).astype(float) + return (v_mismatch + j_mismatch) / 2.0 + + +def bootstrap_ylq( + matrices: dict[str, np.ndarray], pairs: pd.DataFrame, + retrieval: dict[str, pd.DataFrame], replicates: int, seed: int, +) -> dict: + """Paired query-donor bootstrap for all prespecified primary metrics.""" + donors = sorted(set(pairs["query_donor"].astype(str)).union( + *(set(table["query_donor"].astype(str)) for table in retrieval.values()))) + donor_index = {donor: index for index, donor in enumerate(donors)} + pair_donor = pairs["query_donor"].map(donor_index).to_numpy(int) + rng = np.random.default_rng(seed) + weights = rng.multinomial( + len(donors), np.full(len(donors), 1 / len(donors)), size=replicates + ) + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + metric_values: dict[str, dict[str, np.ndarray]] = {} + for name, matrix in matrices.items(): + within = matrix[q, p] + between = matrix[q, n] + result = {key: np.empty(replicates) for key in ( + "within_between_ratio", "between_minus_within", "auroc", "auprc", + "precision_at_1", "map_at_r", + )} + query_table = retrieval[name] + query_donor = query_table["query_donor"].map(donor_index).to_numpy(int) + query_p1 = query_table["precision_at_1"].to_numpy(float) + query_ap = query_table["ap_at_r"].to_numpy(float) + for replicate, donor_weight in enumerate(weights): + pair_weight = donor_weight[pair_donor].astype(float) + within_mean = float(np.average(within, weights=pair_weight)) + between_mean = float(np.average(between, weights=pair_weight)) + result["within_between_ratio"][replicate] = within_mean / between_mean + result["between_minus_within"][replicate] = between_mean - within_mean + truth = np.concatenate((np.ones(len(within), dtype=int), np.zeros(len(between), dtype=int))) + score = -np.concatenate((within, between)) + sample_weight = np.concatenate((pair_weight, pair_weight)) + result["auroc"][replicate] = roc_auc_score( + truth, score, sample_weight=sample_weight + ) + result["auprc"][replicate] = average_precision_score( + truth, score, sample_weight=sample_weight + ) + query_weight = donor_weight[query_donor].astype(float) + result["precision_at_1"][replicate] = np.average(query_p1, weights=query_weight) + result["map_at_r"][replicate] = np.average(query_ap, weights=query_weight) + metric_values[name] = result + + def interval(values: np.ndarray) -> list[float]: + return [float(np.percentile(values, 2.5)), float(np.percentile(values, 97.5))] + + report = { + "replicates": replicates, + "cluster": "query donor", + "paired_across_frozen_models": True, + "models": { + name: {metric: {"ci95": interval(values)} for metric, values in result.items()} + for name, result in metric_values.items() + }, + "paired_differences_rtp_minus_other": {}, + } + for other in ("r", "p", "rp"): + report["paired_differences_rtp_minus_other"][other] = { + metric: { + "ci95": interval(metric_values["rtp"][metric] - metric_values[other][metric]), + "direction": "negative_favors_rtp" if metric == "within_between_ratio" + else "positive_favors_rtp", + } + for metric in metric_values["rtp"] + } + return report + + +def shuffle_null( + matrices: dict[str, np.ndarray], candidates: pd.DataFrame, + permutations: int, seed: int, +) -> dict: + rng = np.random.default_rng(seed + 20000) + labels = candidates["is_ylq"].to_numpy(int).astype(bool) + lengths = candidates["length"].to_numpy(int) + donors = candidates["donor_key"].astype(str).to_numpy() + left, right = np.triu_indices(len(candidates), k=1) + eligible = (lengths[left] == lengths[right]) & (donors[left] != donors[right]) + left, right = left[eligible], right[eligible] + if len(left) > 100000: + selected = rng.choice(len(left), size=100000, replace=False) + left, right = left[selected], right[selected] + + def contrast(distance: np.ndarray, value: np.ndarray) -> float: + positive_pair = value[left] & value[right] + negative_pair = value[left] ^ value[right] + if not positive_pair.any() or not negative_pair.any(): + return float("nan") + return float(distance[left[negative_pair], right[negative_pair]].mean() - + distance[left[positive_pair], right[positive_pair]].mean()) + + observed = {name: contrast(matrix, labels) for name, matrix in matrices.items()} + null = {name: np.empty(permutations) for name in matrices} + groups = [np.flatnonzero(lengths == length) for length in sorted(set(lengths))] + for replicate in range(permutations): + shuffled = labels.copy() + for group in groups: + shuffled[group] = rng.permutation(shuffled[group]) + for name, matrix in matrices.items(): + null[name][replicate] = contrast(matrix, shuffled) + return { + "pairs": len(left), "permutations": permutations, + "strata": "exact CDR3 length", "models": { + name: { + "observed_between_minus_within": observed[name], + "null_ci95": [float(np.nanpercentile(null[name], 2.5)), + float(np.nanpercentile(null[name], 97.5))], + "one_sided_p": float((1 + np.nansum(null[name] >= observed[name])) / + (1 + np.isfinite(null[name]).sum())), + } for name in matrices + }, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--preflight-dir", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--model", action="append", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--bootstrap", type=int, default=2000) + parser.add_argument("--permutations", type=int, default=1000) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + preflight = json.loads((args.preflight_dir / "PREFLIGHT.json").read_text()) + if preflight["status"] != "ready": + raise ValueError("YLQ preflight is not accepting.") + candidate_path = args.preflight_dir / "candidates.tsv" + pair_path = args.preflight_dir / "pair_manifest.tsv" + if sha256(candidate_path) != preflight["artifacts"]["candidates_sha256"]: + raise ValueError("Candidate table hash changed.") + if sha256(pair_path) != preflight["artifacts"]["pairs_sha256"]: + raise ValueError("Pair manifest hash changed.") + candidates = pd.read_csv(candidate_path, sep="\t") + pairs = pd.read_csv(pair_path, sep="\t") + if candidates["cdr3"].duplicated().any(): + raise ValueError("YLQ candidate table is not exact-CDR3 deduplicated.") + ylq = candidates["is_ylq"].to_numpy(int) + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + donors = candidates["donor_key"].astype(str).to_numpy() + lengths = candidates["length"].to_numpy(int) + if not (np.all(ylq[q] == 1) and np.all(ylq[p] == 1) and np.all(ylq[n] == 0)): + raise ValueError("Pair manifest does not contain YLQ/YLQ/non-YLQ triplets.") + if not (np.all(donors[q] != donors[p]) and np.all(donors[q] != donors[n]) + and np.all(donors[p] != donors[n])): + raise ValueError("Pair manifest is not fully donor-disjoint within triplets.") + if not np.all(lengths[p] == lengths[n]): + raise ValueError("Negative comparator is not exact-length matched to the positive comparator.") + pairs["label"] = "YLQPRTFLL" + models = parse_models(args.model) + config, model_report = validate_run_configs(models) + embeddings, extraction = extract_embeddings( + candidates.rename(columns={"candidate_index": "row_index"}), models, config, + args.tokenizer, 128, args.output_dir, + ) + + metrics = {} + bootstraps = {} + matrices_by_metric = {} + for metric in ("cosine", "euclidean"): + matrices = {name: distance_matrix(value, metric) for name, value in embeddings.items()} + matrices_by_metric[metric] = matrices + metrics[metric] = {} + retrieval_tables = {} + for name, matrix in matrices.items(): + pair_value, per_pair_label = pair_metrics(matrix, pairs) + pair_value.update({ + "within_mean": float(per_pair_label.loc[0, "within_mean"]), + "between_mean": float(per_pair_label.loc[0, "between_mean"]), + "between_minus_within": float( + per_pair_label.loc[0, "between_mean"] - per_pair_label.loc[0, "within_mean"] + ), + }) + retrieval_value, per_query = binary_retrieval(matrix, candidates) + metrics[metric][name] = {"pairwise_and_contrast": pair_value, "retrieval": retrieval_value} + retrieval_tables[name] = per_query + per_query.to_csv(args.output_dir / f"retrieval_{metric}_{name}.tsv", sep="\t", index=False) + bootstraps[metric] = bootstrap_ylq( + matrices, pairs, retrieval_tables, args.bootstrap, args.seed + ) + + lengths = candidates["length"].to_numpy(float) + length_control, _ = pair_metrics(np.abs(lengths[:, None] - lengths[None, :]), pairs) + comp = composition_matrix(candidates) + composition_control = {} + for metric in ("cosine", "euclidean"): + composition_control[metric], _ = pair_metrics(distance_matrix(comp, metric), pairs) + vj_control, _ = pair_metrics(vj_distance(candidates), pairs) + controls = { + "exact_length": length_control, + "amino_acid_composition": composition_control, + "v_j_identity": vj_control, + "shuffled_labels": shuffle_null( + matrices_by_metric["cosine"], candidates, args.permutations, args.seed + ), + } + assert_finite_tree(metrics, "metrics") + assert_finite_tree(bootstraps, "bootstraps") + assert_finite_tree(controls, "controls") + stable_json(args.output_dir / "metrics.json", metrics) + stable_json(args.output_dir / "bootstrap.json", bootstraps) + stable_json(args.output_dir / "controls.json", controls) + pairs.to_csv(args.output_dir / "pair_manifest.tsv", sep="\t", index=False) + candidates.to_csv(args.output_dir / "candidates.tsv", sep="\t", index=False) + artifact_hashes = { + path.name: sha256(path) for path in sorted(args.output_dir.iterdir()) if path.is_file() + } + result = { + "status": "complete", "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "target": "exact YLQPRTFLL in dominant exact MHC context", + "primary_context": preflight["primary_context"], + "positive_cdr3": preflight["positive_cdr3"], + "positive_donors": preflight["positive_donors"], + "negative_cdr3": preflight["same_context_non_ylq_negative_cdr3"], + "pairs": len(pairs), "models": model_report, "extraction": extraction, + "motif_expansion_performed": False, + "selection_statement": "All four prespecified frozen checkpoints are reported; YLQ outcomes selected none.", + "artifact_sha256": artifact_hashes, + } + stable_json(args.output_dir / "RESULTS.json", result) + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/evaluate_vdjdb_ylq_motif_pooled.py b/benchmark/evaluate_vdjdb_ylq_motif_pooled.py new file mode 100644 index 0000000..fc1bed1 --- /dev/null +++ b/benchmark/evaluate_vdjdb_ylq_motif_pooled.py @@ -0,0 +1,391 @@ +"""CPU-only frozen evaluation on pooled official VDJdb motif-member clonotypes.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +from pathlib import Path + +import numpy as np +import pandas as pd +from sklearn.metrics import average_precision_score, roc_auc_score + +from benchmark.evaluate_vdjdb_frozen import ( + AA, + assert_finite_tree, + extract_embeddings, + parse_models, + sha256, + stable_json, + validate_run_configs, +) + + +def vector_distance(left: np.ndarray, right: np.ndarray, metric: str) -> np.ndarray: + left = left.astype(np.float64, copy=False) + right = right.astype(np.float64, copy=False) + if metric == "cosine": + denominator = np.linalg.norm(left, axis=1) * np.linalg.norm(right, axis=1) + if np.any(denominator == 0): + raise ValueError("Zero-norm vector prevents cosine distance.") + return np.maximum(0.0, 1.0 - np.sum(left * right, axis=1) / denominator) + if metric == "euclidean": + return np.linalg.norm(left - right, axis=1) + raise ValueError(metric) + + +def primary_distances(array: np.ndarray, pairs: pd.DataFrame, metric: str) -> tuple[np.ndarray, np.ndarray]: + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + return vector_distance(array[q], array[p], metric), vector_distance(array[q], array[n], metric) + + +def pair_metrics(within: np.ndarray, between: np.ndarray) -> dict: + truth = np.concatenate((np.ones(len(within), dtype=int), np.zeros(len(between), dtype=int))) + score = -np.concatenate((within, between)) + within_mean = float(within.mean()) + between_mean = float(between.mean()) + if between_mean == 0: + ratio = 1.0 if within_mean == 0 else float("inf") + else: + ratio = within_mean / between_mean + pooled = math.sqrt((float(within.var()) + float(between.var())) / 2.0) + separation = (between_mean - within_mean) / pooled if pooled > 0 else 0.0 + return { + "pairs_per_class": len(within), + "within_mean": within_mean, + "between_mean": between_mean, + "between_minus_within": between_mean - within_mean, + "within_between_ratio": ratio, + "standardized_separation": separation, + "same_label_auroc": float(roc_auc_score(truth, score)), + "same_label_auprc_balanced": float(average_precision_score(truth, score)), + } + + +def pooled_retrieval(array: np.ndarray, candidates: pd.DataFrame, metric: str) -> tuple[dict, pd.DataFrame]: + positive = candidates["is_ylq"].to_numpy(int).astype(bool) + cdr3 = candidates["cdr3"].astype(str).to_numpy() + if metric == "cosine": + norms = np.linalg.norm(array, axis=1, keepdims=True) + if np.any(norms == 0): + raise ValueError("Zero-norm latent prevents cosine retrieval.") + prepared = array / norms + else: + prepared = array.astype(np.float64, copy=False) + rows = [] + all_indices = np.arange(len(candidates)) + for query in np.flatnonzero(positive): + if metric == "cosine": + distance = np.maximum(0.0, 1.0 - prepared @ prepared[query]) + else: + distance = np.linalg.norm(prepared - prepared[query], axis=1) + allowed = all_indices[all_indices != query] + relevant = positive[allowed] + total_relevant = int(relevant.sum()) + order = np.argsort(distance[allowed], kind="mergesort") + ranked = relevant[order] + hits = ranked[:total_relevant].astype(float) + precision = np.cumsum(hits) / np.arange(1, total_relevant + 1) + rows.append({ + "query_index": int(query), + "query_cdr3": cdr3[query], + "relevant_pooled": total_relevant, + "precision_at_1": float(ranked[0]), + "precision_at_5": float(ranked[:5].mean()), + "precision_at_10": float(ranked[:10].mean()), + "ap_at_r": float((precision * hits).sum() / total_relevant), + }) + table = pd.DataFrame(rows) + return { + "eligible_positive_queries": len(table), + "gallery_excludes_query_only": True, + "donor_filter": "none_pooled_by_user_protocol", + "precision_at_1": float(table["precision_at_1"].mean()), + "precision_at_5": float(table["precision_at_5"].mean()), + "precision_at_10": float(table["precision_at_10"].mean()), + "map_at_r": float(table["ap_at_r"].mean()), + }, table + + +def bootstrap_clonotypes( + pair_values: dict[str, tuple[np.ndarray, np.ndarray]], + pairs: pd.DataFrame, + retrieval: dict[str, pd.DataFrame], + replicates: int, + seed: int, +) -> dict: + clusters = sorted(set(pairs["query_cdr3"].astype(str)).union( + *(set(table["query_cdr3"].astype(str)) for table in retrieval.values()))) + cluster_index = {value: index for index, value in enumerate(clusters)} + pair_cluster = pairs["query_cdr3"].map(cluster_index).to_numpy(int) + rng = np.random.default_rng(seed) + weights = rng.multinomial( + len(clusters), np.full(len(clusters), 1 / len(clusters)), size=replicates + ) + values_by_model: dict[str, dict[str, np.ndarray]] = {} + for name, (within, between) in pair_values.items(): + result = {metric: np.empty(replicates) for metric in ( + "within_between_ratio", "between_minus_within", "auroc", "auprc", + "precision_at_1", "map_at_r", + )} + query_table = retrieval[name] + query_cluster = query_table["query_cdr3"].map(cluster_index).to_numpy(int) + query_p1 = query_table["precision_at_1"].to_numpy(float) + query_ap = query_table["ap_at_r"].to_numpy(float) + truth = np.concatenate((np.ones(len(within), dtype=int), np.zeros(len(between), dtype=int))) + score = -np.concatenate((within, between)) + for replicate, cluster_weight in enumerate(weights): + pair_weight = cluster_weight[pair_cluster].astype(float) + within_mean = float(np.average(within, weights=pair_weight)) + between_mean = float(np.average(between, weights=pair_weight)) + result["within_between_ratio"][replicate] = within_mean / between_mean + result["between_minus_within"][replicate] = between_mean - within_mean + sample_weight = np.concatenate((pair_weight, pair_weight)) + result["auroc"][replicate] = roc_auc_score(truth, score, sample_weight=sample_weight) + result["auprc"][replicate] = average_precision_score( + truth, score, sample_weight=sample_weight + ) + query_weight = cluster_weight[query_cluster].astype(float) + result["precision_at_1"][replicate] = np.average(query_p1, weights=query_weight) + result["map_at_r"][replicate] = np.average(query_ap, weights=query_weight) + values_by_model[name] = result + + def interval(values: np.ndarray) -> list[float]: + return [float(np.percentile(values, 2.5)), float(np.percentile(values, 97.5))] + + report = { + "replicates": replicates, + "cluster": "query exact CDR3 clonotype", + "paired_across_frozen_models": True, + "models": { + name: {metric: {"ci95": interval(values)} for metric, values in model_values.items()} + for name, model_values in values_by_model.items() + }, + "paired_differences_rtp_minus_other": {}, + } + for other in ("r", "p", "rp"): + report["paired_differences_rtp_minus_other"][other] = { + metric: { + "ci95": interval(values_by_model["rtp"][metric] - values_by_model[other][metric]), + "direction": "negative_favors_rtp" if metric == "within_between_ratio" + else "positive_favors_rtp", + } + for metric in values_by_model["rtp"] + } + return report + + +def composition_features(candidates: pd.DataFrame) -> np.ndarray: + result = np.zeros((len(candidates), len(AA)), dtype=float) + for row, sequence in enumerate(candidates["cdr3"].astype(str)): + for column, aa in enumerate(AA): + result[row, column] = sequence.count(aa) / len(sequence) + return result + + +def vj_pair_values(candidates: pd.DataFrame, pairs: pd.DataFrame) -> tuple[np.ndarray, np.ndarray]: + v = candidates["v_call"].fillna("").astype(str).to_numpy() + j = candidates["j_call"].fillna("").astype(str).to_numpy() + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + + def value(left: np.ndarray, right: np.ndarray) -> np.ndarray: + v_mismatch = (v[left] == "") | (v[right] == "") | (v[left] != v[right]) + j_mismatch = (j[left] == "") | (j[right] == "") | (j[left] != j[right]) + return (v_mismatch.astype(float) + j_mismatch.astype(float)) / 2.0 + + return value(q, p), value(q, n) + + +def shuffled_label_control( + pair_values: dict[str, tuple[np.ndarray, np.ndarray]], + candidates: pd.DataFrame, + pairs: pd.DataFrame, + permutations: int, + seed: int, +) -> dict: + rng = np.random.default_rng(seed + 30000) + labels = candidates["is_ylq"].to_numpy(int).astype(bool) + lengths = candidates["length"].to_numpy(int) + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + groups = [np.flatnonzero(lengths == length) for length in sorted(set(lengths))] + observed = { + name: float(between.mean() - within.mean()) + for name, (within, between) in pair_values.items() + } + null = {name: np.full(permutations, np.nan) for name in pair_values} + valid_pair_counts = [] + for replicate in range(permutations): + shuffled = labels.copy() + for group in groups: + shuffled[group] = rng.permutation(shuffled[group]) + same_mask = shuffled[q] & shuffled[p] + different_mask = shuffled[q] ^ shuffled[n] + valid_pair_counts.append((int(same_mask.sum()), int(different_mask.sum()))) + if not same_mask.any() or not different_mask.any(): + continue + for name, (within, between) in pair_values.items(): + null[name][replicate] = float( + between[different_mask].mean() - within[same_mask].mean() + ) + models = {} + for name, values in null.items(): + finite = values[np.isfinite(values)] + if len(finite) < permutations * 0.95: + raise ValueError(f"Too few valid shuffled-label replicates for {name}: {len(finite)}") + models[name] = { + "observed_between_minus_within": observed[name], + "null_ci95": [float(np.percentile(finite, 2.5)), float(np.percentile(finite, 97.5))], + "one_sided_p": float((1 + np.sum(finite >= observed[name])) / (1 + len(finite))), + } + return { + "permutations": permutations, + "strata": "exact CDR3 length", + "fixed_edge_design": "primary query-positive and query-negative edges; labels permuted", + "valid_same_pair_range": [min(value[0] for value in valid_pair_counts), + max(value[0] for value in valid_pair_counts)], + "valid_different_pair_range": [min(value[1] for value in valid_pair_counts), + max(value[1] for value in valid_pair_counts)], + "models": models, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--preflight-dir", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--model", action="append", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--bootstrap", type=int, default=2000) + parser.add_argument("--permutations", type=int, default=1000) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + preflight = json.loads((args.preflight_dir / "PREFLIGHT.json").read_text(encoding="utf-8")) + if preflight["status"] != "ready": + raise ValueError("Pooled official-motif preflight is not accepting.") + candidate_path = args.preflight_dir / "candidates.tsv" + pair_path = args.preflight_dir / "pair_manifest.tsv" + if sha256(candidate_path) != preflight["artifacts"]["candidates_sha256"]: + raise ValueError("Candidate hash changed after CPU gate.") + if sha256(pair_path) != preflight["artifacts"]["pairs_sha256"]: + raise ValueError("Pair hash changed after CPU gate.") + candidates = pd.read_csv(candidate_path, sep="\t") + pairs = pd.read_csv(pair_path, sep="\t") + if candidates["cdr3"].duplicated().any(): + raise ValueError("Candidate table is not exact-CDR3 deduplicated.") + if set(candidates["mhc_a"]) != {"HLA-A*02:01"} or set(candidates["mhc_b"]) != {"B2M"} \ + or set(candidates["mhc_class"]) != {"MHCI"}: + raise ValueError("Candidate table is not restricted to exact A*02:01/B2M/MHCI.") + if candidates["motif_cids"].fillna("").eq("").any(): + raise ValueError("A candidate lacks explicit official motif membership.") + labels = candidates["is_ylq"].to_numpy(int) + q = pairs["query_index"].to_numpy(int) + p = pairs["positive_index"].to_numpy(int) + n = pairs["negative_index"].to_numpy(int) + if not (np.all(labels[q] == 1) and np.all(labels[p] == 1) and np.all(labels[n] == 0)): + raise ValueError("Pair manifest does not contain YLQ/YLQ/non-YLQ triplets.") + if not (np.all(q != p) and np.all(q != n) and np.all(p != n)): + raise ValueError("Pair manifest contains repeated clonotypes within a triplet.") + if not np.all(candidates["length"].to_numpy(int)[p] == candidates["length"].to_numpy(int)[n]): + raise ValueError("Negative comparator is not exact-length matched.") + + models = parse_models(args.model) + config, model_report = validate_run_configs(models) + embeddings, extraction = extract_embeddings( + candidates, models, config, args.tokenizer, args.batch_size, args.output_dir + ) + metrics: dict[str, dict] = {} + bootstraps: dict[str, dict] = {} + cosine_pair_values: dict[str, tuple[np.ndarray, np.ndarray]] = {} + for metric in ("cosine", "euclidean"): + metrics[metric] = {} + pair_values = {} + retrieval_tables = {} + for name, array in embeddings.items(): + within, between = primary_distances(array, pairs, metric) + pair_values[name] = (within, between) + np.savez_compressed( + args.output_dir / f"primary_distances_{metric}_{name}.npz", + within=within, between=between, + ) + retrieval_value, retrieval_table = pooled_retrieval(array, candidates, metric) + retrieval_table.to_csv( + args.output_dir / f"retrieval_{metric}_{name}.tsv", sep="\t", index=False + ) + retrieval_tables[name] = retrieval_table + metrics[metric][name] = { + "pairwise_and_contrast": pair_metrics(within, between), + "retrieval": retrieval_value, + } + bootstraps[metric] = bootstrap_clonotypes( + pair_values, pairs, retrieval_tables, args.bootstrap, args.seed + ) + if metric == "cosine": + cosine_pair_values = pair_values + + lengths = candidates["length"].to_numpy(float) + length_within = np.abs(lengths[q] - lengths[p]) + length_between = np.abs(lengths[q] - lengths[n]) + composition = composition_features(candidates) + composition_control = {} + for metric in ("cosine", "euclidean"): + within = vector_distance(composition[q], composition[p], metric) + between = vector_distance(composition[q], composition[n], metric) + composition_control[metric] = pair_metrics(within, between) + vj_within, vj_between = vj_pair_values(candidates, pairs) + controls = { + "exact_length": pair_metrics(length_within, length_between), + "amino_acid_composition": composition_control, + "v_j_identity": pair_metrics(vj_within, vj_between), + "shuffled_labels_cosine": shuffled_label_control( + cosine_pair_values, candidates, pairs, args.permutations, args.seed + ), + } + assert_finite_tree(metrics, "metrics") + assert_finite_tree(bootstraps, "bootstraps") + assert_finite_tree(controls, "controls") + stable_json(args.output_dir / "metrics.json", metrics) + stable_json(args.output_dir / "bootstrap.json", bootstraps) + stable_json(args.output_dir / "controls.json", controls) + candidates.to_csv(args.output_dir / "candidates.tsv", sep="\t", index=False) + pairs.to_csv(args.output_dir / "pair_manifest.tsv", sep="\t", index=False) + hashes = {path.name: sha256(path) for path in sorted(args.output_dir.iterdir()) if path.is_file()} + result = { + "status": "complete", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "target": "exact YLQPRTFLL, official motif-member, pooled HLA-A*02:01/B2M/MHCI", + "positive_cdr3": preflight["positive_cdr3"], + "negative_cdr3": preflight["negative_cdr3"], + "candidate_cdr3": preflight["candidate_cdr3"], + "pairs": len(pairs), + "pooling": preflight["pooling"], + "models": model_report, + "extraction": extraction, + "artifact_sha256": hashes, + "selection_statement": ( + "All seven nonempty matched R/T/P factorial checkpoints reported; " + "no YLQ-based selection." + ), + "circularity_caveat": preflight["circularity_caveat"], + "claim_scope": "within-VDJdb motif/epitope organization; not cross-donor or independent motif evidence", + "exploratory_n40_protocol": preflight["exploratory_n40_protocol"], + "compact_control_selection": preflight["compact_control_selection"], + "models_trained": False, + } + stable_json(args.output_dir / "RESULTS.json", result) + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/external_residue_models.py b/benchmark/external_residue_models.py new file mode 100644 index 0000000..35810be --- /dev/null +++ b/benchmark/external_residue_models.py @@ -0,0 +1,152 @@ +"""Shared models and datasets for the frozen residue-state benchmark. + +The public encoders remain frozen. Only these task heads are trained. Every +residue-level condition uses the same head topology; the sole shape-dependent +component is the initial linear projection from the public encoder width. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +from torch.utils.data import Dataset + + +class ResidueArrayDataset(Dataset): + """Index a memory-mapped frozen residue matrix without loading it eagerly.""" + + def __init__(self, features_path, mask_path, rows, targets): + self.features = np.load(Path(features_path), mmap_mode="r") + self.mask = np.load(Path(mask_path), mmap_mode="r") + self.rows = np.asarray(rows, dtype=np.int64) + self.targets = np.asarray(targets) + if self.features.shape[:2] != self.mask.shape: + raise ValueError("Frozen feature and mask shapes do not agree.") + if len(self.rows) != len(self.targets): + raise ValueError("Rows and targets must have identical lengths.") + + def __len__(self): + return len(self.rows) + + def __getitem__(self, index): + row = self.rows[index] + return ( + torch.from_numpy(np.asarray(self.features[row], dtype=np.float32)), + torch.from_numpy(np.asarray(self.mask[row], dtype=np.bool_)), + torch.as_tensor(self.targets[index]), + ) + + +class FrozenResidueDecoder(nn.Module): + """Position-aware decoder over frozen per-residue states. + + Learned output queries cross-attend to every residue state. Unlike global + mean pooling, this interface never averages away residue order before the + task head sees it. + """ + + def __init__( + self, + input_dim, + vocab_size=25, + max_len=40, + hidden_dim=512, + num_layers=3, + nhead=8, + ff_mult=4, + dropout=0.2, + ): + super().__init__() + self.max_len = max_len + self.input_norm = nn.LayerNorm(input_dim) + self.input_projection = nn.Linear(input_dim, hidden_dim) + self.source_position = nn.Parameter(torch.randn(max_len, hidden_dim) * 0.02) + self.output_queries = nn.Parameter(torch.randn(max_len, hidden_dim) * 0.02) + layer = nn.TransformerDecoderLayer( + d_model=hidden_dim, + nhead=nhead, + dim_feedforward=hidden_dim * ff_mult, + dropout=dropout, + activation="gelu", + batch_first=True, + norm_first=True, + ) + self.decoder = nn.TransformerDecoder(layer, num_layers=num_layers) + self.output = nn.Linear(hidden_dim, vocab_size) + + def forward(self, features, residue_mask): + if features.ndim != 3 or residue_mask.shape != features.shape[:2]: + raise ValueError("Expected features [B,L,D] and mask [B,L].") + memory = self.input_projection(self.input_norm(features)) + memory = memory + self.source_position[: memory.shape[1]].unsqueeze(0) + queries = self.output_queries.unsqueeze(0).expand(features.shape[0], -1, -1) + decoded = self.decoder( + queries, + memory, + memory_key_padding_mask=~residue_mask.bool(), + ) + return self.output(decoded) + + +class FrozenResiduePgenHead(nn.Module): + """Trainable position-aware pgen probe over frozen residue states.""" + + def __init__( + self, + input_dim, + max_len=40, + hidden_dim=256, + num_layers=2, + nhead=8, + ff_mult=4, + dropout=0.2, + ): + super().__init__() + self.input_norm = nn.LayerNorm(input_dim) + self.input_projection = nn.Linear(input_dim, hidden_dim) + self.source_position = nn.Parameter(torch.randn(max_len, hidden_dim) * 0.02) + layer = nn.TransformerEncoderLayer( + d_model=hidden_dim, + nhead=nhead, + dim_feedforward=hidden_dim * ff_mult, + dropout=dropout, + activation="gelu", + batch_first=True, + norm_first=True, + ) + self.context = nn.TransformerEncoder(layer, num_layers=num_layers) + self.query = nn.Parameter(torch.randn(1, 1, hidden_dim) * 0.02) + self.pool = nn.MultiheadAttention(hidden_dim, nhead, dropout=dropout, batch_first=True) + self.regressor = nn.Sequential( + nn.LayerNorm(hidden_dim), + nn.Linear(hidden_dim, 512), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(512, 256), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(256, 1), + ) + + def forward(self, features, residue_mask): + if features.ndim != 3 or residue_mask.shape != features.shape[:2]: + raise ValueError("Expected features [B,L,D] and mask [B,L].") + hidden = self.input_projection(self.input_norm(features)) + hidden = hidden + self.source_position[: hidden.shape[1]].unsqueeze(0) + hidden = self.context(hidden, src_key_padding_mask=~residue_mask.bool()) + query = self.query.expand(features.shape[0], -1, -1) + pooled, _ = self.pool( + query, + hidden, + hidden, + key_padding_mask=~residue_mask.bool(), + need_weights=False, + ) + return self.regressor(pooled[:, 0]).squeeze(-1) + + +def parameter_count(model): + return int(sum(parameter.numel() for parameter in model.parameters())) diff --git a/benchmark/extract_vdjdb_latent.py b/benchmark/extract_vdjdb_latent.py new file mode 100644 index 0000000..ba8c5de --- /dev/null +++ b/benchmark/extract_vdjdb_latent.py @@ -0,0 +1,79 @@ +"""Extract one frozen IRRM-CODEC latent matrix for an ordered VDJdb cohort.""" +from __future__ import annotations + +import argparse +import hashlib +import json +from dataclasses import asdict +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from irrm_codec.multitask_data import resolve_encoder_tokenizer +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import PAD_ID + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=512) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + args = parser.parse_args() + if args.output.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output}") + + cohort = pd.read_csv(args.cohort, sep="\t") + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + config = IRRMCodecConfig(**checkpoint["model_config"]) + tokenizer = resolve_encoder_tokenizer("data_anchor", str(args.tokenizer)) + if tokenizer.vocab_size != config.input_vocab_size: + raise ValueError("Tokenizer vocabulary differs from checkpoint architecture") + encoded = [tokenizer.encode(sequence, config.max_sequence_len) + for sequence in cohort["cdr3"].astype(str)] + device = torch.device(args.device) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but unavailable") + model = IRRMCodecTransformer(config).to(device) + model.load_state_dict(checkpoint["model_state"], strict=True) + model.eval() + matrix = np.lib.format.open_memmap( + args.output, mode="w+", dtype=np.float32, + shape=(len(cohort), config.latent_dim), + ) + with torch.inference_mode(): + for start in range(0, len(encoded), args.batch_size): + rows = encoded[start:start + args.batch_size] + width = max(map(len, rows)) + tokens = torch.full((len(rows), width), PAD_ID, dtype=torch.long, device=device) + for index, ids in enumerate(rows): + tokens[index, :len(ids)] = torch.tensor(ids, dtype=torch.long, device=device) + matrix[start:start + len(rows)] = model.encode(tokens, tokens.ne(PAD_ID)).float().cpu().numpy() + matrix.flush() + if not np.isfinite(matrix).all(): + raise ValueError("Extracted latent contains non-finite values") + metadata = { + "status": "complete", "rows": len(cohort), "shape": list(matrix.shape), + "checkpoint": str(args.checkpoint), "checkpoint_sha256": sha256(args.checkpoint), + "checkpoint_epoch_zero_based": int(checkpoint["epoch"]), + "model_config": asdict(config), "cohort_sha256": sha256(args.cohort), + "latent_sha256": sha256(args.output), "ordering_preserved": True, + } + args.output.with_suffix(".json").write_text(json.dumps(metadata, indent=2, sort_keys=True)+"\n") + + +if __name__ == "__main__": + main() diff --git a/benchmark/finalize_external_residue_pca256.py b/benchmark/finalize_external_residue_pca256.py new file mode 100644 index 0000000..67dd22e --- /dev/null +++ b/benchmark/finalize_external_residue_pca256.py @@ -0,0 +1,57 @@ +"""Fail-closed catalog construction for train-only residue PCA vectors.""" + +import argparse +import json +from pathlib import Path + +import numpy as np + + +NAMES = ("esm2_8m_residue", "tcr_bert_residue", "sceptr_residue") + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--rows", type=int, default=99430) + parser.add_argument("--components", type=int, default=256) + return parser.parse_args() + + +def main(): + args = parse_args() + output = Path(args.output_dir) + catalog, summary = {}, {} + for name in NAMES: + report_path = output / f"{name}_report.json" + report = json.loads(report_path.read_text(encoding="utf-8")) + array_path = output / f"{name}_{args.components}.npy" + matrix = np.load(array_path, mmap_mode="r") + if report.get("status") != "accepted" or report.get("validation_or_test_used_for_fit"): + raise ValueError(f"Rejected or leaky projection report for {name}") + if tuple(matrix.shape) != (args.rows, args.components): + raise ValueError(f"Wrong projected shape for {name}: {matrix.shape}") + if not np.isfinite(matrix).all(): + raise ValueError(f"Non-finite projection for {name}") + catalog[name] = { + "path": str(array_path), + "shape": list(matrix.shape), + "sha256": report["output_sha256"], + "interface": f"flattened_residue_states_train_only_pca{args.components}", + } + summary[name] = { + "source_dim": report["flattened_dim"], + "components_fitted": args.components, + "explained_variance": {str(args.components): report["explained_variance_ratio"]}, + "solver": report["solver"], + "fit_split": "train", + "fit_rows_sha256": report["fit_rows_sha256"], + "capped_by_source_dim": False, + } + (output / "representations.json").write_text(json.dumps(catalog, indent=2), encoding="utf-8") + (output / "bottleneck_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + print(f"status=accepted representations={len(catalog)} rows={args.rows} dim={args.components}") + + +if __name__ == "__main__": + main() diff --git a/benchmark/finalize_external_residue_pca_single.py b/benchmark/finalize_external_residue_pca_single.py new file mode 100644 index 0000000..c8a250f --- /dev/null +++ b/benchmark/finalize_external_residue_pca_single.py @@ -0,0 +1,71 @@ +"""Fail-closed catalog construction for one train-only residue PCA vector.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--representation", required=True) + parser.add_argument("--rows", type=int, default=99430) + parser.add_argument("--components", type=int, required=True) + return parser.parse_args() + + +def main(): + args = parse_args() + output = Path(args.output_dir) + name = f"{args.representation}_residue" + report_path = output / f"{name}_report.json" + report = json.loads(report_path.read_text(encoding="utf-8")) + array_path = output / f"{name}_{args.components}.npy" + matrix = np.load(array_path, mmap_mode="r") + if report.get("status") != "accepted" or report.get("validation_or_test_used_for_fit"): + raise ValueError("Rejected or leaky projection report") + if report.get("fit_split") != "train" or report.get("components") != args.components: + raise ValueError("Projection provenance or dimensionality mismatch") + if tuple(matrix.shape) != (args.rows, args.components): + raise ValueError(f"Wrong projected shape: {matrix.shape}") + if not np.isfinite(matrix).all(): + raise ValueError("Non-finite projection") + catalog = { + name: { + "path": str(array_path), + "shape": list(matrix.shape), + "sha256": report["output_sha256"], + "interface": f"flattened_residue_states_train_only_pca{args.components}", + } + } + summary = { + name: { + "source_dim": report["flattened_dim"], + "components_fitted": args.components, + "explained_variance": { + str(args.components): report["explained_variance_ratio"] + }, + "solver": report["solver"], + "fit_split": "train", + "fit_rows_sha256": report["fit_rows_sha256"], + "capped_by_source_dim": False, + } + } + (output / "representations.json").write_text( + json.dumps(catalog, indent=2), encoding="utf-8" + ) + (output / "bottleneck_summary.json").write_text( + json.dumps(summary, indent=2), encoding="utf-8" + ) + print( + f"status=accepted representation={name} rows={args.rows} " + f"dim={args.components} explained={report['explained_variance_ratio']:.6f}" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/finalize_vdjdb_epitope_prediction.py b/benchmark/finalize_vdjdb_epitope_prediction.py new file mode 100644 index 0000000..eb74689 --- /dev/null +++ b/benchmark/finalize_vdjdb_epitope_prediction.py @@ -0,0 +1,117 @@ +"""Finalize an already-scored VDJdb epitope-prediction output directory.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import pandas as pd + +from benchmark.evaluate_vdjdb_epitope_prediction import ( + aggregate_metrics, + create_figure, + latex_table, + paired_statistics, + sha256, + stable_json, +) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + output = args.output_dir + raw_path = output / "fold_metrics.csv" + predictions = output / "predictions.parquet" + samples_path = output / "samples.tsv" + splits_path = output / "split_assignments.csv" + for path in (raw_path, predictions, samples_path, splits_path): + if not path.exists(): + raise FileNotFoundError(path) + + raw = pd.read_csv(raw_path) + expected_representations = { + "r", "t", "p", "rt", "rp", "tp", "rtp", "tcr_bert", "sceptr_cdr3", + "esm2_35m", "kmer_tfidf", + } + if set(raw["representation"]) != expected_representations: + raise ValueError("Fold metrics do not contain all eleven planned representations.") + observed = raw.groupby(["protocol", "representation"]).size() + expected_observations = { + "sequence_stratified": 15, + "exact_cdr3_grouped": 15, + "redcea_similarity_grouped": 12, + } + for (protocol, _), count in observed.items(): + if int(count) != expected_observations[protocol]: + raise ValueError(f"Incomplete {protocol} representation: {count}") + + aggregate = aggregate_metrics(raw) + aggregate.to_csv(output / "aggregated_metrics.csv", index=False, lineterminator="\n") + differences, tests = paired_statistics(raw) + differences.to_csv(output / "paired_macro_f1_differences.csv", index=False, + lineterminator="\n") + tests.to_csv(output / "paired_tests.csv", index=False, lineterminator="\n") + for protocol in aggregate["protocol"].unique(): + (output / f"table_{protocol}.tex").write_text( + latex_table(aggregate, protocol), encoding="utf-8" + ) + create_figure(raw, aggregate, output) + + samples = pd.read_csv(samples_path, sep="\t") + sizes = samples.groupby("label").size() + label_counts = samples.groupby("cdr3")["label"].nunique() + audit = { + "sequences": int(len(samples)), + "epitopes": int(samples["label"].nunique()), + "class_size_min": int(sizes.min()), + "class_size_median": float(np.median(sizes)), + "class_size_max": int(sizes.max()), + "duplicate_cdr3_rows": int(samples["cdr3"].duplicated(keep=False).sum()), + "duplicate_cdr3_sequences_exist": bool(samples["cdr3"].duplicated().any()), + "identical_cdr3_with_different_labels": int((label_counts > 1).sum()), + } + outputs = ( + "samples.tsv", "class_sizes.csv", "split_assignments.csv", "fold_metrics.csv", + "aggregated_metrics.csv", "paired_macro_f1_differences.csv", "paired_tests.csv", + "predictions.parquet", "epitope_prediction.png", "epitope_prediction.pdf", + "table_sequence_stratified.tex", "table_exact_cdr3_grouped.tex", + "table_redcea_similarity_grouped.tex", + ) + representation_summary = ( + raw.groupby(["representation", "display_name"], sort=False) + .agg(dim=("dim", "median"), dim_min=("dim", "min"), dim_max=("dim", "max")) + .reset_index() + ) + for column in ("dim", "dim_min", "dim_max"): + representation_summary[column] = representation_summary[column].round().astype(int) + + result = { + "status": "complete", + "recovered_from_completed_scoring": True, + "dataset": audit, + "protocols": { + "sequence_stratified": "5 folds x 3 seeds", + "exact_cdr3_grouped": "alias of sequence-stratified because every CDR3 is unique", + "redcea_similarity_grouped": "4 folds x 3 seeds; connected source_cid components", + "inner_tuning": "training-only C in {0.1, 1.0}; macro F1 selection", + "classifier": "balanced L2 multinomial logistic regression", + }, + "representations": representation_summary.to_dict(orient="records"), + + "convergence": { + "fold_fits": int(len(raw)), + "converged_before_max_iter": int(raw["converged_before_max_iter"].sum()) + if "converged_before_max_iter" in raw else None, + }, + "outputs": {name: sha256(output / name) for name in outputs}, + } + stable_json(output / "RESULTS.json", result) + print(json.dumps({"status": "complete", "output": str(output), **audit}, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/fit_bottleneck.py b/benchmark/fit_bottleneck.py new file mode 100644 index 0000000..20f7d22 --- /dev/null +++ b/benchmark/fit_bottleneck.py @@ -0,0 +1,160 @@ +"""Issue 2: reduce every representation to a common bottleneck dimension with PCA. + +The scaler and the PCA are fitted on the training split only and then applied to +validation and test, so no information from the held-out rows leaks into the +projection. Representations arrive on very different scales (TCRemP holds alignment +distances in the hundreds, transformer activations sit near zero), so each one is +standardized before the PCA. + +One PCA is fitted at the largest requested width. Principal components are ordered by +explained variance, so the narrower bottlenecks are prefixes of the same fit and cost +nothing extra. +""" + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +from joblib import dump +from sklearn.decomposition import PCA, IncrementalPCA +from sklearn.preprocessing import StandardScaler + +from irrm_codec.utils import setup_logging + +# Above this input width the full SVD needs more memory than it is worth, so the +# incremental solver streams the training split in batches instead. +INCREMENTAL_PCA_MIN_DIM = 2000 + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument("--representations-dir", default="data/benchmark/trb/representations") + p.add_argument("--output-dir", default="data/benchmark/trb/bottleneck") + p.add_argument("--dims", type=int, nargs="+", default=[32, 64, 128]) + p.add_argument("--representations", nargs="+") + p.add_argument("--batch-size", type=int, default=4096, help="IncrementalPCA batch size.") + p.add_argument("--seed", type=int, default=42) + p.add_argument("--overwrite", action="store_true") + return p.parse_args() + + +def fit_projection(matrix, train_rows, n_components, batch_size, seed, log): + """Standardize and fit a PCA on the training rows only.""" + scaler = StandardScaler() + incremental = matrix.shape[1] >= INCREMENTAL_PCA_MIN_DIM + + if incremental: + # Two streamed passes: one to learn the scaling, one to fit the components. + for start in range(0, len(train_rows), batch_size): + scaler.partial_fit(np.asarray(matrix[train_rows[start : start + batch_size]])) + pca = IncrementalPCA(n_components=n_components, batch_size=batch_size) + for start in range(0, len(train_rows), batch_size): + block = np.asarray(matrix[train_rows[start : start + batch_size]], dtype=np.float64) + if len(block) < n_components: + break # A trailing batch smaller than the component count cannot be fitted. + pca.partial_fit(scaler.transform(block)) + else: + train = np.asarray(matrix[train_rows], dtype=np.float64) + scaler.fit(train) + pca = PCA(n_components=n_components, random_state=seed) + pca.fit(scaler.transform(train)) + + log.info( + "fitted %s pca n_components=%d", + "incremental" if incremental else "full", + n_components, + ) + return scaler, pca + + +def transform_all(matrix, scaler, pca, batch_size): + """Project every row, streaming so the full float64 copy never materializes.""" + out = np.empty((matrix.shape[0], pca.n_components_), dtype=np.float32) + for start in range(0, matrix.shape[0], batch_size): + block = np.asarray(matrix[start : start + batch_size], dtype=np.float64) + out[start : start + batch_size] = pca.transform(scaler.transform(block)).astype(np.float32) + return out + + +def main(): + args = parse_args() + dataset_dir = Path(args.dataset_dir) + reps_dir = Path(args.representations_dir) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "fit_bottleneck.log") + + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + train_rows = pd.read_csv(dataset_dir / "manifests" / "train.tsv", sep="\t")["row_index"].to_numpy() + log.info("dataset rows=%d train rows=%d", len(frame), len(train_rows)) + + catalog = json.loads((reps_dir / "representations.json").read_text(encoding="utf-8")) + names = args.representations or list(catalog) + max_dim = max(args.dims) + + summary_path = output_dir / "bottleneck_summary.json" + summary = json.loads(summary_path.read_text(encoding="utf-8")) if summary_path.exists() else {} + + for name in names: + source = Path(catalog[name]["path"]) + matrix = np.load(source, mmap_mode="r") + if matrix.shape[0] != len(frame): + raise ValueError(f"{name} has {matrix.shape[0]} rows, expected {len(frame)}.") + + # A representation narrower than the bottleneck cannot be reduced to it; SCEPTR + # is already 64-dimensional, so its widest projection is capped at its own width. + n_components = min(max_dim, matrix.shape[1], len(train_rows)) + if name in summary and not args.overwrite: + log.info("%s already projected, skipping", name) + continue + + log.info("fitting %s dim=%d -> %d", name, matrix.shape[1], n_components) + started = time.perf_counter() + scaler, pca = fit_projection( + matrix, train_rows, n_components, args.batch_size, args.seed, log + ) + projected = transform_all(matrix, scaler, pca, args.batch_size) + elapsed = time.perf_counter() - started + + dump({"scaler": scaler, "pca": pca}, output_dir / f"{name}_projection.joblib") + + ratios = np.asarray(pca.explained_variance_ratio_, dtype=np.float64) + explained = {} + for dim in sorted(args.dims): + if dim > n_components: + explained[str(dim)] = None + continue + np.save(output_dir / f"{name}_{dim}.npy", projected[:, :dim]) + explained[str(dim)] = round(float(ratios[:dim].sum()), 6) + + summary[name] = { + "source_dim": int(matrix.shape[1]), + "components_fitted": int(n_components), + "explained_variance": explained, + "seconds": round(elapsed, 2), + "solver": "incremental" if matrix.shape[1] >= INCREMENTAL_PCA_MIN_DIM else "full", + "capped_by_source_dim": bool(n_components < max_dim), + } + summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") + log.info("%s done in %.1f s explained=%s", name, elapsed, explained) + + log.info("=" * 68) + header = " ".join(f"{d:>8d}" for d in sorted(args.dims)) + log.info("%-10s %6s %s", "repr", "dim", header) + for name, info in summary.items(): + cells = " ".join( + f"{info['explained_variance'][str(d)]:8.4f}" + if info["explained_variance"].get(str(d)) is not None + else " -" + for d in sorted(args.dims) + ) + log.info("%-10s %6d %s", name, info["source_dim"], cells) + log.info("wrote %s", summary_path) + + +if __name__ == "__main__": + main() diff --git a/benchmark/fit_external_model_bottleneck.py b/benchmark/fit_external_model_bottleneck.py new file mode 100644 index 0000000..20f7d22 --- /dev/null +++ b/benchmark/fit_external_model_bottleneck.py @@ -0,0 +1,160 @@ +"""Issue 2: reduce every representation to a common bottleneck dimension with PCA. + +The scaler and the PCA are fitted on the training split only and then applied to +validation and test, so no information from the held-out rows leaks into the +projection. Representations arrive on very different scales (TCRemP holds alignment +distances in the hundreds, transformer activations sit near zero), so each one is +standardized before the PCA. + +One PCA is fitted at the largest requested width. Principal components are ordered by +explained variance, so the narrower bottlenecks are prefixes of the same fit and cost +nothing extra. +""" + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +from joblib import dump +from sklearn.decomposition import PCA, IncrementalPCA +from sklearn.preprocessing import StandardScaler + +from irrm_codec.utils import setup_logging + +# Above this input width the full SVD needs more memory than it is worth, so the +# incremental solver streams the training split in batches instead. +INCREMENTAL_PCA_MIN_DIM = 2000 + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument("--representations-dir", default="data/benchmark/trb/representations") + p.add_argument("--output-dir", default="data/benchmark/trb/bottleneck") + p.add_argument("--dims", type=int, nargs="+", default=[32, 64, 128]) + p.add_argument("--representations", nargs="+") + p.add_argument("--batch-size", type=int, default=4096, help="IncrementalPCA batch size.") + p.add_argument("--seed", type=int, default=42) + p.add_argument("--overwrite", action="store_true") + return p.parse_args() + + +def fit_projection(matrix, train_rows, n_components, batch_size, seed, log): + """Standardize and fit a PCA on the training rows only.""" + scaler = StandardScaler() + incremental = matrix.shape[1] >= INCREMENTAL_PCA_MIN_DIM + + if incremental: + # Two streamed passes: one to learn the scaling, one to fit the components. + for start in range(0, len(train_rows), batch_size): + scaler.partial_fit(np.asarray(matrix[train_rows[start : start + batch_size]])) + pca = IncrementalPCA(n_components=n_components, batch_size=batch_size) + for start in range(0, len(train_rows), batch_size): + block = np.asarray(matrix[train_rows[start : start + batch_size]], dtype=np.float64) + if len(block) < n_components: + break # A trailing batch smaller than the component count cannot be fitted. + pca.partial_fit(scaler.transform(block)) + else: + train = np.asarray(matrix[train_rows], dtype=np.float64) + scaler.fit(train) + pca = PCA(n_components=n_components, random_state=seed) + pca.fit(scaler.transform(train)) + + log.info( + "fitted %s pca n_components=%d", + "incremental" if incremental else "full", + n_components, + ) + return scaler, pca + + +def transform_all(matrix, scaler, pca, batch_size): + """Project every row, streaming so the full float64 copy never materializes.""" + out = np.empty((matrix.shape[0], pca.n_components_), dtype=np.float32) + for start in range(0, matrix.shape[0], batch_size): + block = np.asarray(matrix[start : start + batch_size], dtype=np.float64) + out[start : start + batch_size] = pca.transform(scaler.transform(block)).astype(np.float32) + return out + + +def main(): + args = parse_args() + dataset_dir = Path(args.dataset_dir) + reps_dir = Path(args.representations_dir) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "fit_bottleneck.log") + + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + train_rows = pd.read_csv(dataset_dir / "manifests" / "train.tsv", sep="\t")["row_index"].to_numpy() + log.info("dataset rows=%d train rows=%d", len(frame), len(train_rows)) + + catalog = json.loads((reps_dir / "representations.json").read_text(encoding="utf-8")) + names = args.representations or list(catalog) + max_dim = max(args.dims) + + summary_path = output_dir / "bottleneck_summary.json" + summary = json.loads(summary_path.read_text(encoding="utf-8")) if summary_path.exists() else {} + + for name in names: + source = Path(catalog[name]["path"]) + matrix = np.load(source, mmap_mode="r") + if matrix.shape[0] != len(frame): + raise ValueError(f"{name} has {matrix.shape[0]} rows, expected {len(frame)}.") + + # A representation narrower than the bottleneck cannot be reduced to it; SCEPTR + # is already 64-dimensional, so its widest projection is capped at its own width. + n_components = min(max_dim, matrix.shape[1], len(train_rows)) + if name in summary and not args.overwrite: + log.info("%s already projected, skipping", name) + continue + + log.info("fitting %s dim=%d -> %d", name, matrix.shape[1], n_components) + started = time.perf_counter() + scaler, pca = fit_projection( + matrix, train_rows, n_components, args.batch_size, args.seed, log + ) + projected = transform_all(matrix, scaler, pca, args.batch_size) + elapsed = time.perf_counter() - started + + dump({"scaler": scaler, "pca": pca}, output_dir / f"{name}_projection.joblib") + + ratios = np.asarray(pca.explained_variance_ratio_, dtype=np.float64) + explained = {} + for dim in sorted(args.dims): + if dim > n_components: + explained[str(dim)] = None + continue + np.save(output_dir / f"{name}_{dim}.npy", projected[:, :dim]) + explained[str(dim)] = round(float(ratios[:dim].sum()), 6) + + summary[name] = { + "source_dim": int(matrix.shape[1]), + "components_fitted": int(n_components), + "explained_variance": explained, + "seconds": round(elapsed, 2), + "solver": "incremental" if matrix.shape[1] >= INCREMENTAL_PCA_MIN_DIM else "full", + "capped_by_source_dim": bool(n_components < max_dim), + } + summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") + log.info("%s done in %.1f s explained=%s", name, elapsed, explained) + + log.info("=" * 68) + header = " ".join(f"{d:>8d}" for d in sorted(args.dims)) + log.info("%-10s %6s %s", "repr", "dim", header) + for name, info in summary.items(): + cells = " ".join( + f"{info['explained_variance'][str(d)]:8.4f}" + if info["explained_variance"].get(str(d)) is not None + else " -" + for d in sorted(args.dims) + ) + log.info("%-10s %6d %s", name, info["source_dim"], cells) + log.info("wrote %s", summary_path) + + +if __name__ == "__main__": + main() diff --git a/benchmark/fit_external_residue_pca256.py b/benchmark/fit_external_residue_pca256.py new file mode 100644 index 0000000..5ae99f6 --- /dev/null +++ b/benchmark/fit_external_residue_pca256.py @@ -0,0 +1,176 @@ +"""Fit a train-only PCA from flattened residue states to one compact clonotype vector. + +The public encoder states are already frozen. We flatten the zero-padded [40,D] +tensor in residue order, standardize each flattened coordinate using the locked +training split only, fit a seeded randomized PCA, and apply that fixed projection to +all rows. Validation and test rows never participate in fitting. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + + +REPRESENTATIONS = ("esm2_8m", "tcr_bert", "sceptr") + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(8 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def rows_sha256(rows: np.ndarray) -> str: + return hashlib.sha256(np.asarray(rows, dtype=" 0 else 0.0 + return ( + mean.cpu().numpy().astype(np.float32), + scale.cpu().numpy().astype(np.float32), + components.cpu().numpy().astype(np.float32), + explained, + ) + + +def transform_to_memmap(matrix, mean, scale, components, output, batch_size, device): + projected = np.lib.format.open_memmap( + output, mode="w+", dtype=np.float32, shape=(matrix.shape[0], components.shape[1]) + ) + mean_t = torch.from_numpy(mean).to(device) + scale_t = torch.from_numpy(scale).to(device) + components_t = torch.from_numpy(components).to(device) + for start in range(0, matrix.shape[0], batch_size): + stop = min(start + batch_size, matrix.shape[0]) + block = np.asarray(matrix[start:stop], dtype=np.float32).reshape(stop - start, -1) + x = torch.from_numpy(block).to(device) + x.sub_(mean_t).div_(scale_t) + projected[start:stop] = (x @ components_t).cpu().numpy() + projected.flush() + return projected + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-dir", required=True) + parser.add_argument("--residue-dir", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--representation", choices=REPRESENTATIONS, required=True) + parser.add_argument("--components", type=int, default=256) + parser.add_argument("--oversample", type=int, default=32) + parser.add_argument("--power-iterations", type=int, default=4) + parser.add_argument("--batch-size", type=int, default=2048) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def main(): + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("Residue PCA fitting must run in a GPU Slurm allocation.") + device = torch.device("cuda") + dataset_dir, residue_dir, output_dir = map( + Path, (args.dataset_dir, args.residue_dir, args.output_dir) + ) + output_dir.mkdir(parents=True, exist_ok=True) + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + train_manifest = dataset_dir / "manifests" / "train.tsv" + train_rows = pd.read_csv(train_manifest, sep="\t")["row_index"].to_numpy(np.int64) + source = residue_dir / f"{args.representation}_residue.npy" + source_mask = residue_dir / f"{args.representation}_residue_mask.npy" + matrix = np.load(source, mmap_mode="r") + mask = np.load(source_mask, mmap_mode="r") + if matrix.ndim != 3 or matrix.shape[:2] != mask.shape or matrix.shape[0] != len(frame): + raise ValueError(f"Unexpected residue/mask shapes: {matrix.shape}, {mask.shape}") + started = time.perf_counter() + mean, scale, components, explained = fit_randomized_pca( + matrix, + train_rows, + args.components, + args.oversample, + args.power_iterations, + args.seed, + device, + ) + output = output_dir / f"{args.representation}_residue_{args.components}.npy" + projected = transform_to_memmap( + matrix, mean, scale, components, output, args.batch_size, device + ) + for start in range(0, len(projected), args.batch_size): + if not np.isfinite(projected[start : start + args.batch_size]).all(): + raise ValueError("Non-finite projected values.") + projection_path = output_dir / f"{args.representation}_residue_projection.npz" + np.savez( + projection_path, + mean=mean, + scale=scale, + components=components, + fit_rows=train_rows, + ) + report = { + "status": "accepted", + "representation": f"{args.representation}_residue", + "source": str(source), + "source_sha256": sha256(source), + "source_shape": list(matrix.shape), + "flattened_dim": int(matrix.shape[1] * matrix.shape[2]), + "output": str(output), + "output_sha256": sha256(output), + "output_shape": list(projected.shape), + "projection": str(projection_path), + "projection_sha256": sha256(projection_path), + "components": args.components, + "explained_variance_ratio": explained, + "fit_split": "train", + "fit_rows": int(len(train_rows)), + "fit_rows_sha256": rows_sha256(train_rows), + "train_manifest": str(train_manifest), + "train_manifest_sha256": sha256(train_manifest), + "validation_or_test_used_for_fit": False, + "standardization": "per_flattened_coordinate_train_only", + "solver": "torch_randomized_pca_lowrank", + "oversample": args.oversample, + "power_iterations": args.power_iterations, + "seed": args.seed, + "finite": True, + "seconds": round(time.perf_counter() - started, 2), + } + report_path = output_dir / f"{args.representation}_residue_report.json" + report_path.write_text(json.dumps(report, indent=2), encoding="utf-8") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/preflight_author_global_embeddings.py b/benchmark/preflight_author_global_embeddings.py new file mode 100644 index 0000000..a6cfb71 --- /dev/null +++ b/benchmark/preflight_author_global_embeddings.py @@ -0,0 +1,134 @@ +"""Fail-closed CPU preflight for author-native global embedding benchmarks.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import logging +from pathlib import Path + +import numpy as np +import pandas as pd + +from benchmark.cache_external_model_embeddings import encode_esm2, encode_sceptr + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def read_table(path: Path) -> pd.DataFrame: + if path.suffix.lower() == ".parquet": + return pd.read_parquet(path) + return pd.read_csv(path, sep="\t") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--locked-dataset", type=Path, required=True) + parser.add_argument("--redcea-cohort", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--locked-sha256", required=True) + parser.add_argument("--redcea-sha256", required=True) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + args.output_dir.mkdir(parents=True, exist_ok=False) + + actual_hashes = { + "locked_dataset": sha256_file(args.locked_dataset), + "redcea_cohort": sha256_file(args.redcea_cohort), + } + expected_hashes = { + "locked_dataset": args.locked_sha256, + "redcea_cohort": args.redcea_sha256, + } + if actual_hashes != expected_hashes: + raise RuntimeError( + f"Input hash mismatch: expected={expected_hashes}, actual={actual_hashes}" + ) + + locked = read_table(args.locked_dataset) + redcea = read_table(args.redcea_cohort) + for name, frame in (("locked", locked), ("redcea", redcea)): + if "junction_aa" not in frame and "cdr3" not in frame: + raise ValueError(f"{name} input lacks junction_aa/cdr3") + seq_col = "junction_aa" if "junction_aa" in frame else "cdr3" + if frame[seq_col].astype(str).duplicated().any() and name == "redcea": + raise ValueError("REDCEA cohort must contain unique clonotypes") + + locked_sequences = locked["junction_aa"].astype(str).head(2).tolist() + redcea_seq_col = "junction_aa" if "junction_aa" in redcea else "cdr3" + redcea_sequences = redcea[redcea_seq_col].astype(str).head(2).tolist() + + class SmokeArgs: + device = "cpu" + batch_size = 2 + + log = logging.getLogger("author_global_preflight") + esm35 = encode_esm2(locked_sequences, SmokeArgs(), log, model_name="esm2_35m") + if esm35.shape != (2, 480) or not np.isfinite(esm35).all(): + raise RuntimeError(f"ESM2-35M smoke failed: shape={esm35.shape}") + + sceptr_cdr3 = encode_sceptr( + locked_sequences, + SmokeArgs(), + log, + frame=locked.head(2).copy(), + cdr3_only=True, + ) + if sceptr_cdr3.shape != (2, 64) or not np.isfinite(sceptr_cdr3).all(): + raise RuntimeError(f"SCEPTR-CDR3 smoke failed: shape={sceptr_cdr3.shape}") + + if not {"v_call", "j_call"}.issubset(redcea.columns): + raise ValueError("REDCEA cohort lacks v_call/j_call for SCEPTR-default") + sceptr_default = encode_sceptr( + redcea_sequences, + SmokeArgs(), + log, + frame=redcea.head(2).rename(columns={"cdr3": "junction_aa"}).copy(), + cdr3_only=False, + ) + if sceptr_default.shape != (2, 64) or not np.isfinite(sceptr_default).all(): + raise RuntimeError(f"SCEPTR-default smoke failed: shape={sceptr_default.shape}") + + report = { + "status": "accepted", + "input_hashes": actual_hashes, + "row_counts": {"locked": len(locked), "redcea": len(redcea)}, + "representations": { + "esm2_35m": { + "model": "esm2_t12_35M_UR50D", + "pooling": "final-layer residue mean excluding special/pad tokens", + "modality": "sequence-only", + "smoke_shape": list(esm35.shape), + }, + "sceptr_cdr3": { + "model": "SCEPTR variant.cdr3_only()", + "pooling": "native contextualized CLS64", + "modality": "sequence-only", + "smoke_shape": list(sceptr_cdr3.shape), + }, + "sceptr": { + "model": "SCEPTR default", + "pooling": "native contextualized CLS64", + "modality": "annotation-aware TRBV+CDR3B+TRBJ", + "smoke_shape": list(sceptr_default.shape), + }, + }, + } + output = args.output_dir / "PREFLIGHT.json" + output.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8") + print(json.dumps(report, sort_keys=True)) + + + +if __name__ == "__main__": + main() diff --git a/benchmark/preflight_external_model_benchmark.py b/benchmark/preflight_external_model_benchmark.py new file mode 100644 index 0000000..9a84772 --- /dev/null +++ b/benchmark/preflight_external_model_benchmark.py @@ -0,0 +1,217 @@ +"""Fail-closed preflight for the locked external-model benchmark. + +This script performs metadata/import/model-cache checks on a Slurm CPU node. It does +not fit a model or produce benchmark scores. Its JSON output is the gate used before +external representation extraction. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib +import importlib.util +import json +import os +import platform +import sys +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd + + +REQUIRED_COLUMNS = { + "junction_aa", + "v_call", + "j_call", + "log10_pgen_1mm", +} +SPECIAL_REPRESENTATIONS = ("onehot", "tcremp", "rtp") +EXTERNAL_REPRESENTATIONS = ("esm2_8m", "tcr_bert", "sceptr") + + +def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(chunk_size): + digest.update(chunk) + return digest.hexdigest() + + +def read_manifest_rows(path: Path) -> np.ndarray: + frame = pd.read_csv(path, sep="\t") + if "row_index" not in frame: + raise ValueError(f"{path} has no row_index column") + rows = frame["row_index"].to_numpy(dtype=np.int64) + if len(rows) != len(np.unique(rows)): + raise ValueError(f"{path} contains duplicate row_index values") + return rows + + +def validate_splits(dataset_dir: Path, n_rows: int) -> dict[str, Any]: + manifests = dataset_dir / "manifests" + split_rows = { + split: read_manifest_rows(manifests / f"{split}.tsv") + for split in ("train", "val", "test") + } + for split, rows in split_rows.items(): + if np.any(rows < 0) or np.any(rows >= n_rows): + raise ValueError(f"{split} has out-of-range row_index") + sets = {name: set(rows.tolist()) for name, rows in split_rows.items()} + overlaps = { + "train_val": len(sets["train"] & sets["val"]), + "train_test": len(sets["train"] & sets["test"]), + "val_test": len(sets["val"] & sets["test"]), + } + if any(overlaps.values()): + raise ValueError(f"split overlap detected: {overlaps}") + if len(set.union(*sets.values())) != n_rows: + raise ValueError("train/val/test manifests do not cover the dataset exactly") + return { + "counts": {name: int(len(rows)) for name, rows in split_rows.items()}, + "overlaps": overlaps, + "manifest_sha256": { + name: sha256_file(manifests / f"{name}.tsv") for name in split_rows + }, + } + + +def module_report(name: str) -> dict[str, Any]: + spec = importlib.util.find_spec(name) + report: dict[str, Any] = {"available": spec is not None} + if spec is None: + return report + try: + module = importlib.import_module(name) + report.update( + { + "import_ok": True, + "version": getattr(module, "__version__", None), + "path": getattr(module, "__file__", None), + } + ) + except Exception as exc: # preflight records the exact import blocker + report.update({"import_ok": False, "error": f"{type(exc).__name__}: {exc}"}) + return report + + +def inspect_hf_cache() -> dict[str, Any]: + cache_root = Path( + os.environ.get( + "HF_HOME", + os.environ.get("HUGGINGFACE_HUB_CACHE", str(Path.home() / ".cache" / "huggingface")), + ) + ) + hub_root = cache_root if cache_root.name == "hub" else cache_root / "hub" + candidates = { + "tcr_bert": hub_root / "models--wukevin--tcr-bert", + } + return { + "root": str(hub_root), + "models": { + name: { + "path": str(path), + "exists": path.exists(), + "snapshots": sorted(p.name for p in (path / "snapshots").glob("*")) + if (path / "snapshots").exists() + else [], + } + for name, path in candidates.items() + }, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-dir", type=Path, required=True) + parser.add_argument("--rtp-embeddings", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--expected-rows", type=int, default=99_430) + parser.add_argument("--expected-latent-dim", type=int, default=128) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + args.output.parent.mkdir(parents=True, exist_ok=True) + dataset_path = args.dataset_dir / "dataset.parquet" + ready_path = args.dataset_dir / "READY.json" + if not dataset_path.exists() or not ready_path.exists(): + raise FileNotFoundError("locked dataset.parquet/READY.json is missing") + + frame = pd.read_parquet(dataset_path) + missing = sorted(REQUIRED_COLUMNS - set(frame.columns)) + if missing: + raise ValueError(f"dataset missing required columns: {missing}") + if len(frame) != args.expected_rows: + raise ValueError(f"dataset has {len(frame)} rows, expected {args.expected_rows}") + if frame["junction_aa"].duplicated().any(): + raise ValueError("locked dataset contains duplicate CDR3 amino-acid sequences") + if not np.isfinite(frame["log10_pgen_1mm"].to_numpy(dtype=float)).all(): + raise ValueError("log10_pgen_1mm contains non-finite targets") + + split_report = validate_splits(args.dataset_dir, len(frame)) + rtp = np.load(args.rtp_embeddings, mmap_mode="r") + if rtp.shape != (args.expected_rows, args.expected_latent_dim): + raise ValueError( + f"RTP embedding shape is {rtp.shape}, expected " + f"({args.expected_rows}, {args.expected_latent_dim})" + ) + # Streamed finiteness check avoids making an unnecessary 50 MB copy. + for start in range(0, len(rtp), 8192): + if not np.isfinite(np.asarray(rtp[start : start + 8192])).all(): + raise ValueError(f"RTP embeddings non-finite at block beginning {start}") + + v = frame["v_call"].fillna("").astype(str).str.strip() + j = frame["j_call"].fillna("").astype(str).str.strip() + sceptr_candidate = v.ne("") & j.ne("") + report = { + "status": "accepted_metadata_gate", + "note": "Imports/caches are inventoried; encoder execution is a separate Slurm gate.", + "python": { + "executable": sys.executable, + "version": sys.version, + "platform": platform.platform(), + }, + "dataset": { + "path": str(dataset_path), + "rows": int(len(frame)), + "dataset_sha256": sha256_file(dataset_path), + "ready_sha256": sha256_file(ready_path), + "unique_cdr3": int(frame["junction_aa"].nunique()), + "finite_log10_pgen_1mm": True, + "split": split_report, + }, + "rtp": { + "path": str(args.rtp_embeddings), + "shape": list(rtp.shape), + "dtype": str(rtp.dtype), + "finite": True, + "sha256": sha256_file(args.rtp_embeddings), + }, + "sceptr_metadata": { + "candidate_rows_with_v_and_j": int(sceptr_candidate.sum()), + "missing_v": int(v.eq("").sum()), + "missing_j": int(j.eq("").sum()), + "candidate_by_split": {}, + "warning": "SCEPTR uses V/J calls and is annotation-aware, unlike the sequence-only encoders.", + }, + "modules": { + name: module_report(name) + for name in ("torch", "sklearn", "transformers", "esm", "sceptr") + }, + "hf_cache": inspect_hf_cache(), + "planned_representations": [*SPECIAL_REPRESENTATIONS, *EXTERNAL_REPRESENTATIONS], + } + for split in ("train", "val", "test"): + rows = read_manifest_rows(args.dataset_dir / "manifests" / f"{split}.tsv") + report["sceptr_metadata"]["candidate_by_split"][split] = int(sceptr_candidate.iloc[rows].sum()) + + args.output.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/preflight_external_residue_benchmark.py b/benchmark/preflight_external_residue_benchmark.py new file mode 100644 index 0000000..ce61e04 --- /dev/null +++ b/benchmark/preflight_external_residue_benchmark.py @@ -0,0 +1,137 @@ +"""Fail-closed gate for the unpooled frozen-residue benchmark.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from benchmark.external_residue_models import ( + FrozenResidueDecoder, + FrozenResiduePgenHead, + parameter_count, +) +from irrm_codec.losses import inverse_loss, pgen_loss +from irrm_codec.tokenization import encode + + +REPRESENTATIONS = ("esm2_8m", "tcr_bert", "sceptr") + + +def sha256_file(path, chunk_size=1024 * 1024): + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + while block := handle.read(chunk_size): + digest.update(block) + return digest.hexdigest() + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cohort-dir", required=True) + parser.add_argument("--smoke-dir", required=True) + parser.add_argument("--output", required=True) + return parser.parse_args() + + +def load_rows(cohort, split): + path = cohort / "manifests" / f"{split}.tsv" + rows = pd.read_csv(path, sep="\t")["row_index"].to_numpy(dtype=np.int64) + if len(np.unique(rows)) != len(rows): + raise ValueError(f"Duplicate row identity in {split}") + return rows, sha256_file(path) + + +def main(): + args = parse_args() + cohort, smoke = Path(args.cohort_dir), Path(args.smoke_dir) + cohort_report = json.loads((cohort / "COHORT.json").read_text(encoding="utf-8")) + if cohort_report.get("status") != "accepted" or cohort_report.get("split_overlap") != 0: + raise ValueError("Locked cohort is not accepted/non-overlapping") + frame = pd.read_parquet(cohort / "dataset.parquet") + expected_counts = {"train": 79544, "val": 9943, "test": 9943} + split_rows, manifest_hashes = {}, {} + for split, expected in expected_counts.items(): + split_rows[split], manifest_hashes[split] = load_rows(cohort, split) + if len(split_rows[split]) != expected: + raise ValueError(f"{split} count changed") + if len(set(split_rows["train"]) & set(split_rows["val"])): + raise ValueError("train/val overlap") + if len(set(split_rows["train"]) & set(split_rows["test"])): + raise ValueError("train/test overlap") + if len(set(split_rows["val"]) & set(split_rows["test"])): + raise ValueError("val/test overlap") + + catalog = json.loads((smoke / "residue_representations.json").read_text(encoding="utf-8")) + targets = torch.tensor( + [encode(sequence, max_len=40) for sequence in frame.junction_aa.iloc[:2]], dtype=torch.long + ) + pgen = torch.tensor(frame.log10_pgen_1mm.iloc[:2].to_numpy(), dtype=torch.float32) + checks = {} + for name in REPRESENTATIONS: + entry = catalog[name] + features = np.load(entry["features"]) + mask = np.load(entry["mask"]) + if features.shape[0] < 2 or features.shape[1] != 40: + raise ValueError(f"{name} smoke shape invalid: {features.shape}") + lengths = frame.junction_aa.iloc[: len(features)].str.len().to_numpy() + if not np.array_equal(mask.sum(axis=1), lengths): + raise ValueError(f"{name} mask/length mismatch") + if not np.isfinite(features).all(): + raise ValueError(f"{name} smoke features non-finite") + x = torch.from_numpy(features[:2].astype(np.float32)) + m = torch.from_numpy(mask[:2]) + decoder = FrozenResidueDecoder(features.shape[-1]).eval() + regressor = FrozenResiduePgenHead(features.shape[-1]).eval() + with torch.no_grad(): + logits = decoder(x, m) + reconstruction_smoke = inverse_loss(logits, targets) + prediction = regressor(x, m) + pgen_smoke = pgen_loss(prediction, pgen) + if not torch.isfinite(reconstruction_smoke) or not torch.isfinite(pgen_smoke): + raise ValueError(f"{name} smoke loss is non-finite") + checks[name] = { + "shape": list(features.shape), + "dtype": str(features.dtype), + "mask_counts_equal_lengths": True, + "features_sha256": sha256_file(entry["features"]), + "mask_sha256": sha256_file(entry["mask"]), + "decoder_parameters": parameter_count(decoder), + "pgen_head_parameters": parameter_count(regressor), + "finite_reconstruction_smoke_loss": float(reconstruction_smoke), + "finite_pgen_smoke_loss": float(pgen_smoke), + } + report = { + "status": "accepted", + "protocol": "frozen per-residue states; no global pooling before trainable task head", + "cohort": { + "rows": len(frame), + "dataset_sha256": sha256_file(cohort / "dataset.parquet"), + "split_counts": expected_counts, + "manifest_sha256": manifest_hashes, + "overlap": 0, + }, + "representations": checks, + "controls": { + "encoder_frozen": True, + "decoder_trained_separately_per_representation": True, + "pgen_head_trained_separately_per_representation": True, + "checkpoint_selection": "validation loss only", + "test_used_for_selection": False, + "pgen_normalization": "train-only", + "important_scope": "best-available frozen interface, not a matched compact-vector comparison", + }, + } + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2), encoding="utf-8") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/preflight_external_sceptr_pca128.py b/benchmark/preflight_external_sceptr_pca128.py new file mode 100644 index 0000000..9dc8c78 --- /dev/null +++ b/benchmark/preflight_external_sceptr_pca128.py @@ -0,0 +1,105 @@ +"""Standalone fail-closed smoke gate for the SCEPTR residue-PCA 128D study.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from benchmark.fit_external_residue_pca256 import fit_randomized_pca +from benchmark.train_external_model_pgen import FROZEN_ARMS, RegressionHead +from irrm_codec.inverse_model import InverseModel + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(8 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-dir", required=True) + parser.add_argument("--residue-dir", required=True) + parser.add_argument("--output", required=True) + return parser.parse_args() + + +def main(): + args = parse_args() + dataset_dir = Path(args.dataset_dir) + residue_dir = Path(args.residue_dir) + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + splits = { + name: pd.read_csv(dataset_dir / "manifests" / f"{name}.tsv", sep="\t")[ + "row_index" + ].to_numpy(np.int64) + for name in ("train", "val", "test") + } + if tuple(map(len, splits.values())) != (79544, 9943, 9943): + raise ValueError("Locked split counts changed") + split_sets = {name: set(rows.tolist()) for name, rows in splits.items()} + if any( + split_sets[a] & split_sets[b] + for a, b in (("train", "val"), ("train", "test"), ("val", "test")) + ): + raise ValueError("Split row overlap") + source = residue_dir / "sceptr_residue.npy" + matrix = np.load(source, mmap_mode="r") + if tuple(matrix.shape) != (len(frame), 40, 64): + raise ValueError(f"Unexpected SCEPTR residue shape {matrix.shape}") + for start in range(0, len(matrix), 4096): + if not np.isfinite(matrix[start : start + 4096]).all(): + raise ValueError("Non-finite SCEPTR residue states") + + rng = np.random.default_rng(7) + toy = rng.normal(size=(30, 4, 5)).astype(np.float32) + mean, scale, components, explained = fit_randomized_pca( + toy, + np.arange(20), + n_components=6, + oversample=2, + niter=2, + seed=42, + device=torch.device("cpu"), + ) + scores = ((toy.reshape(30, -1) - mean) / scale) @ components + if scores.shape != (30, 6) or not np.isfinite(scores).all() or not 0 < explained <= 1: + raise ValueError("Randomized PCA smoke failed") + if FROZEN_ARMS.get("sceptr_residue_pca128_mlp") != "sceptr_residue": + raise ValueError("Pgen arm mapping is missing or incorrect") + with torch.no_grad(): + features = torch.randn(4, 128) + prediction = RegressionHead(128)(features) + logits = InverseModel(embedding_dim=128, max_len=40)(features) + if prediction.shape != (4,) or logits.shape != (4, 40, 25): + raise ValueError("Head shape smoke failed") + if not torch.isfinite(prediction).all() or not torch.isfinite(logits).all(): + raise ValueError("Non-finite head smoke") + + report = { + "status": "accepted", + "rows": len(frame), + "split_counts": {name: len(rows) for name, rows in splits.items()}, + "split_overlap": 0, + "source_shape": list(matrix.shape), + "source_sha256": sha256(source), + "target_dim": 128, + "pgen_arm": "sceptr_residue_pca128_mlp", + "finite_smoke": True, + } + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2), encoding="utf-8") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/prepare_external_model_cohort.py b/benchmark/prepare_external_model_cohort.py new file mode 100644 index 0000000..e42ede1 --- /dev/null +++ b/benchmark/prepare_external_model_cohort.py @@ -0,0 +1,188 @@ +"""Build the locked, common-row cohort used by every external representation. + +The source benchmark and its split labels remain immutable. Rows with V/J calls that +SCEPTR cannot standardize as functional human genes are removed *before* model scoring; +the same filtered rows are then used for every representation. Original row identity is +retained in ``source_row_index`` and all emitted matrices are checked by that mapping. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd + + +def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(chunk_size): + digest.update(chunk) + return digest.hexdigest() + + +def manifest_rows(path: Path) -> np.ndarray: + return pd.read_csv(path, sep="\t")["row_index"].to_numpy(dtype=np.int64) + + +def deterministic_rank(sequence: str, seed: int) -> str: + return hashlib.sha256(f"{seed}|{sequence}".encode("utf-8")).hexdigest() + + +def nested_training_subsets( + frame: pd.DataFrame, train_rows: np.ndarray, sizes=(1_000, 10_000), seed=42 +) -> dict[str, np.ndarray]: + ranked = sorted( + train_rows.tolist(), + key=lambda row: (deterministic_rank(str(frame.at[row, "junction_aa"]), seed), row), + ) + subsets: dict[str, np.ndarray] = {} + previous: list[int] = [] + for size in sizes: + if len(ranked) < size: + raise ValueError(f"only {len(ranked)} filtered train rows, cannot emit {size}") + previous_set = set(previous) + chosen = previous + [row for row in ranked if row not in previous_set][: size - len(previous)] + subsets[str(size // 1000) + "k"] = np.asarray(chosen, dtype=np.int64) + previous = chosen + subsets["all"] = np.asarray(ranked, dtype=np.int64) + return subsets + + +def subset_matrix(source: Path, source_rows: np.ndarray, output: Path) -> dict[str, object]: + matrix = np.load(source, mmap_mode="r") + if np.any(source_rows < 0) or np.any(source_rows >= matrix.shape[0]): + raise ValueError(f"{source}: row mapping exceeds matrix shape {matrix.shape}") + target = np.lib.format.open_memmap( + output, mode="w+", dtype=np.float32, shape=(len(source_rows), matrix.shape[1]) + ) + for start in range(0, len(source_rows), 4096): + block = np.asarray(matrix[source_rows[start : start + 4096]], dtype=np.float32) + if not np.isfinite(block).all(): + raise ValueError(f"{source}: non-finite block at filtered row {start}") + target[start : start + len(block)] = block + target.flush() + del target + return { + "source": str(source), + "source_sha256": sha256_file(source), + "output": str(output), + "output_sha256": sha256_file(output), + "shape": [int(len(source_rows)), int(matrix.shape[1])], + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-dataset-dir", type=Path, required=True) + parser.add_argument("--rtp-embeddings", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + import tidytcells as tt + + source_dataset = args.source_dataset_dir / "dataset.parquet" + source_frame = pd.read_parquet(source_dataset).reset_index(drop=True) + standardized: dict[str, pd.Series] = {} + usable = np.ones(len(source_frame), dtype=bool) + dropped: dict[str, int] = {} + for column, gene_type in (("v_call", "v"), ("j_call", "j")): + values = source_frame[column].map( + lambda value: tt.tr.standardize( + value, species="homosapiens", enforce_functional=True, log_failures=False + ) + ) + standardized[column] = values + invalid = values.isna().to_numpy() + dropped[f"nonfunctional_or_invalid_{gene_type}"] = int(invalid.sum()) + usable &= ~invalid + + source_rows = np.flatnonzero(usable).astype(np.int64) + output = args.output_dir + manifests = output / "manifests" + reps = output / "representations" + manifests.mkdir(parents=True, exist_ok=True) + reps.mkdir(parents=True, exist_ok=True) + + filtered = source_frame.iloc[source_rows].copy() + filtered["source_row_index"] = source_rows + filtered["v_call"] = standardized["v_call"].iloc[source_rows].to_numpy() + filtered["j_call"] = standardized["j_call"].iloc[source_rows].to_numpy() + filtered["row_index"] = np.arange(len(filtered), dtype=np.int64) + filtered.to_parquet(output / "dataset.parquet", index=False) + + source_to_new = {int(source): int(new) for new, source in enumerate(source_rows)} + split_new: dict[str, np.ndarray] = {} + source_manifest_hashes: dict[str, str] = {} + for split in ("train", "val", "test"): + path = args.source_dataset_dir / "manifests" / f"{split}.tsv" + original = manifest_rows(path) + kept = [source_to_new[int(row)] for row in original if int(row) in source_to_new] + split_new[split] = np.asarray(kept, dtype=np.int64) + filtered.iloc[split_new[split]].to_csv(manifests / f"{split}.tsv", sep="\t", index=False) + source_manifest_hashes[split] = sha256_file(path) + + sets = {name: set(rows.tolist()) for name, rows in split_new.items()} + if any(sets[a] & sets[b] for a, b in (("train", "val"), ("train", "test"), ("val", "test"))): + raise ValueError("filtered split overlap") + if len(set.union(*sets.values())) != len(filtered): + raise ValueError("filtered splits do not cover the common cohort") + + subsets = nested_training_subsets(filtered, split_new["train"], seed=args.seed) + for name, rows in subsets.items(): + filtered.iloc[rows].to_csv(manifests / f"train_{name}.tsv", sep="\t", index=False) + + matrix_reports = { + "tcremp": subset_matrix( + args.source_dataset_dir / "embeddings.npy", source_rows, reps / "tcremp.npy" + ), + "rtp": subset_matrix(args.rtp_embeddings, source_rows, reps / "rtp.npy"), + } + catalog = { + name: { + "path": report["output"], + "rows": report["shape"][0], + "dim": report["shape"][1], + "sha256": report["output_sha256"], + "source": report["source"], + "source_sha256": report["source_sha256"], + } + for name, report in matrix_reports.items() + } + (reps / "representations.json").write_text( + json.dumps(catalog, indent=2, sort_keys=True), encoding="utf-8" + ) + + report = { + "status": "accepted", + "definition": "common functional-human-V/J intersection fixed before model evaluation", + "source_dataset": str(source_dataset), + "source_dataset_sha256": sha256_file(source_dataset), + "source_rows": int(len(source_frame)), + "common_rows": int(len(filtered)), + "dropped_total": int(len(source_frame) - len(filtered)), + "dropped": dropped, + "split_counts": {name: int(len(rows)) for name, rows in split_new.items()}, + "split_overlap": 0, + "source_manifest_sha256": source_manifest_hashes, + "sequence_sha256": hashlib.sha256( + "\n".join(filtered["junction_aa"].astype(str)).encode("utf-8") + ).hexdigest(), + "nested_train_counts": {name: int(len(rows)) for name, rows in subsets.items()}, + "matrices": matrix_reports, + "sceptr_modality": "sequence plus V/J annotations; reported separately from sequence-only models", + } + (output / "COHORT.json").write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/prepare_single_aa_challenge.py b/benchmark/prepare_single_aa_challenge.py new file mode 100644 index 0000000..7281314 --- /dev/null +++ b/benchmark/prepare_single_aa_challenge.py @@ -0,0 +1,203 @@ +"""Prepare a deterministic, parent-disjoint single-amino-acid challenge. + +Only locked benchmark *test* clonotypes are eligible parents. The generated +mutants are scoring inputs, never additional training observations for an +IRRM encoder. Probe folds are assigned at the parent level. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd + + +AA = "ACDEFGHIKLMNPQRSTVWY" + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def stable_json(path: Path, payload: object) -> None: + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def parent_folds(parent_ids: np.ndarray, seed: int) -> dict[str, str]: + ids = np.asarray(sorted(map(str, parent_ids)), dtype=object) + rng = np.random.default_rng(seed) + shuffled = ids[rng.permutation(len(ids))] + n_train = int(np.floor(0.6 * len(ids))) + n_val = int(np.floor(0.2 * len(ids))) + if min(n_train, n_val, len(ids) - n_train - n_val) < 1: + raise ValueError("At least five parents are required for 60/20/20 probe folds.") + return { + str(identifier): ("train" if i < n_train else "val" if i < n_train + n_val else "test") + for i, identifier in enumerate(shuffled) + } + + +def choose_parents(table: pd.DataFrame, indices: np.ndarray, limit: int, seed: int) -> pd.DataFrame: + eligible = table.iloc[indices].copy() + eligible = eligible[eligible["junction_aa"].astype(str).str.len() >= 4] + eligible["length"] = eligible["junction_aa"].astype(str).str.len() + if eligible.empty: + raise ValueError("No held-out parent has a mutable internal residue.") + # Proportional, deterministic length-stratified sampling prevents the most + # common CDR3 length from monopolizing a bounded challenge. + if len(eligible) > limit: + rng = np.random.default_rng(seed) + selected: list[int] = [] + groups = list(eligible.groupby("length", sort=True)) + raw = {length: limit * len(group) / len(eligible) for length, group in groups} + quotas = {length: int(np.floor(value)) for length, value in raw.items()} + remainder_order = sorted( + groups, key=lambda item: (-(raw[item[0]] - quotas[item[0]]), -len(item[1]), item[0]) + ) + for length, _group in remainder_order[: limit - sum(quotas.values())]: + quotas[length] += 1 + for length, group in groups: + take = min(quotas[length], len(group)) + if take: + selected.extend(rng.choice(group.index.to_numpy(), size=take, replace=False).tolist()) + eligible = eligible.loc[sorted(selected)] + if len(eligible) != limit: + raise AssertionError(f"Length-stratified sampler selected {len(eligible)} parents, expected {limit}.") + return eligible.sort_values("row_index", kind="mergesort").reset_index(drop=True) + + +def generate_mutants( + parents: pd.DataFrame, + benchmark_sequences: set[str], + variants_per_parent: int, + seed: int, +) -> pd.DataFrame: + if variants_per_parent < 2: + raise ValueError("At least two mutants per parent are required for within-parent ranking.") + used = set(benchmark_sequences) + rows: list[dict] = [] + for row in parents.to_dict(orient="records"): + parent_id = f"parent-{int(row['row_index']):06d}" + sequence = str(row["junction_aa"]) + candidates = [(pos, aa) for pos in range(1, len(sequence) - 1) for aa in AA if aa != sequence[pos]] + rng = np.random.default_rng(seed + int(row["row_index"]) * 104729) + accepted = 0 + for candidate_index in rng.permutation(len(candidates)): + position, new_aa = candidates[int(candidate_index)] + mutant = sequence[:position] + new_aa + sequence[position + 1:] + if mutant in used: + continue + used.add(mutant) + mutant_id = f"{parent_id}-m{accepted + 1:02d}" + rows.append({ + "clone_id": mutant_id, + "parent_id": parent_id, + "parent_row_index": int(row["row_index"]), + "parent_sequence": sequence, + "junction_aa": mutant, + "mutation_position_zero_based": int(position), + "mutation_position_fraction": float(position / (len(sequence) - 1)), + "from_aa": sequence[position], + "to_aa": new_aa, + "length": len(sequence), + "v_call": str(row.get("v_call", "")), + "j_call": str(row.get("j_call", "")), + "locus": "TRB", + }) + accepted += 1 + if accepted == variants_per_parent: + break + if accepted != variants_per_parent: + raise ValueError(f"Could not generate {variants_per_parent} unique mutants for {parent_id}.") + result = pd.DataFrame(rows) + if result["junction_aa"].duplicated().any() or set(result["junction_aa"]) & benchmark_sequences: + raise AssertionError("Generated mutants are not unique and benchmark-disjoint.") + mismatch = [ + sum(a != b for a, b in zip(row.parent_sequence, row.junction_aa)) + for row in result.itertuples(index=False) + ] + if set(mismatch) != {1}: + raise AssertionError("A generated variant is not exactly one substitution from its parent.") + return result + + +def prepare(data_dir: Path, output_dir: Path, max_parents: int, variants: int, seed: int) -> dict: + # Keep deterministic generation utilities importable in lightweight local + # test environments that do not install the model's PyTorch dependency. + from irrm_codec.multitask_data import load_prepared_benchmark, select_split_indices + if output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {output_dir}") + output_dir.mkdir(parents=True) + table, embeddings = load_prepared_benchmark(data_dir) + test_indices = select_split_indices(table, data_dir, "test") + parents = choose_parents(table, test_indices, max_parents, seed) + folds = parent_folds(np.array([f"parent-{int(i):06d}" for i in parents["row_index"]]), seed) + parents["parent_id"] = [f"parent-{int(i):06d}" for i in parents["row_index"]] + parents["probe_split"] = parents["parent_id"].map(folds) + mutants = generate_mutants(parents, set(table["junction_aa"].astype(str)), variants, seed) + mutants["probe_split"] = mutants["parent_id"].map(folds) + if mutants["probe_split"].isna().any(): + raise AssertionError("A mutant lacks a parent-derived probe split.") + + parent_target_path = output_dir / "parent_tcremp.npy" + np.save(parent_target_path, np.asarray(embeddings[parents["row_index"].to_numpy(int)], dtype=np.float32)) + parents_path = output_dir / "parents.tsv" + mutants_path = output_dir / "mutants.tsv" + scoring_path = output_dir / "mutants_airr.tsv" + parents.to_csv(parents_path, sep="\t", index=False, lineterminator="\n") + mutants.to_csv(mutants_path, sep="\t", index=False, lineterminator="\n") + mutants[["clone_id", "junction_aa", "v_call", "j_call", "locus"]].to_csv( + scoring_path, sep="\t", index=False, lineterminator="\n" + ) + split_counts = mutants.groupby("probe_split")["parent_id"].nunique().to_dict() + report = { + "status": "ready_for_independent_target_scoring", + "seed": seed, + "source_split": "locked benchmark test only", + "parents": int(len(parents)), + "mutants": int(len(mutants)), + "variants_per_parent": variants, + "all_pairs_edit_distance": 1, + "parent_disjoint_probe_folds": True, + "probe_parent_counts": {str(key): int(value) for key, value in split_counts.items()}, + "benchmark_mutant_overlap": 0, + "tcremp_shape": [int(len(parents)), int(embeddings.shape[1])], + "sources": { + "dataset": {"path": str(data_dir / "dataset.parquet"), "sha256": sha256(data_dir / "dataset.parquet")}, + "embeddings": {"path": str(data_dir / "embeddings.npy"), "sha256": sha256(data_dir / "embeddings.npy")}, + "test_manifest": {"path": str(data_dir / "manifests/test.tsv"), "sha256": sha256(data_dir / "manifests/test.tsv")}, + }, + "artifacts": { + path.name: sha256(path) for path in (parents_path, mutants_path, scoring_path, parent_target_path) + }, + } + stable_json(output_dir / "PREPARED.json", report) + return report + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--max-parents", type=int, default=2000) + parser.add_argument("--variants-per-parent", type=int, default=3) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + report = prepare(args.data_dir, args.output_dir, args.max_parents, args.variants_per_parent, args.seed) + print(json.dumps(report, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/prepare_splits.py b/benchmark/prepare_splits.py index 40ea2e9..aa5983b 100644 --- a/benchmark/prepare_splits.py +++ b/benchmark/prepare_splits.py @@ -17,6 +17,7 @@ import numpy as np import pandas as pd +import tidytcells as tt import pyarrow.parquet as pq from irrm_codec.dataio import iter_embedding_batches, normalize_locus_name @@ -73,6 +74,19 @@ def clean_sequences(df, cdr3_col, locus, min_len, max_len, log): drop |= mask df = df[~drop].copy() + # SCEPTR consumes V/J gene calls and rejects anything non-standard or + # non-functional, while the other representations only read the CDR3. Resolving the + # gene names here keeps a single row set that every representation can encode. + for column, gene_type in (("v_call", "V"), ("j_call", "J")): + df[column] = df[column].map( + lambda value: tt.tr.standardize( + value, species="homosapiens", enforce_functional=True, log_failures=False + ) + ) + unusable = df[column].isna() + report[f"dropped_non_functional_{gene_type.lower()}_call"] = int(unusable.sum()) + df = df[~unusable].copy() + duplicated = df[cdr3_col].duplicated(keep="first") report["dropped_duplicate_cdr3"] = int(duplicated.sum()) df = df[~duplicated].copy() @@ -266,21 +280,35 @@ def check_alignment(emb, sequences, min_corr, log): return {"length_corr": corr, "shuffled_length_corr": shuffled, "min_corr_threshold": min_corr} +def _pgen_lookup(cache_path, needed): + """Return per-sequence pgen values from the cache, or None if it cannot serve them. + + Values are keyed on the CDR3 sequence rather than on row position, so tightening + the cleaning filters later reuses the cache instead of recomputing it. + """ + if not cache_path.exists(): + return None + cached = pd.read_csv(cache_path, sep="\t") + if not all(c in cached.columns for c in ("junction_aa", *PGEN_COLUMNS)): + return None + cached = cached.drop_duplicates("junction_aa").set_index("junction_aa") + if not needed.isin(cached.index).all(): + return None + return cached.loc[needed, list(PGEN_COLUMNS)].reset_index(drop=True) + + + def ensure_pgen(df, args, output_dir, log): cache_path = output_dir / "pgen.tsv" clean_airr_path = output_dir / "cleaned_airr.tsv" df.to_csv(clean_airr_path, sep="\t", index=False) + hit = _pgen_lookup(cache_path, df["junction_aa"]) + if hit is not None: + log.info("reusing cached pgen values for all %d sequences from %s", len(hit), cache_path) + return hit, {"source": "cache", "path": str(cache_path)} if cache_path.exists(): - cached = pd.read_csv(cache_path, sep="\t") - same_rows = len(cached) == len(df) - same_seqs = same_rows and np.array_equal( - cached["junction_aa"].astype(str).to_numpy(), df["junction_aa"].to_numpy() - ) - if same_seqs and all(c in cached.columns for c in PGEN_COLUMNS): - log.info("reusing cached pgen table %s", cache_path) - return cached, {"source": "cache", "path": str(cache_path)} - log.warning("cached pgen table does not match the cleaned data; recomputing") + log.warning("cached pgen table does not cover the cleaned data; recomputing") if args.skip_pgen: raise FileNotFoundError( @@ -300,10 +328,10 @@ def ensure_pgen(df, args, output_dir, log): log.info("computing pgen: %s", " ".join(cmd)) subprocess.run(cmd, check=True, cwd=Path.cwd()) - table = pd.read_csv(cache_path, sep="\t") - if not np.array_equal(table["junction_aa"].astype(str).to_numpy(), df["junction_aa"].to_numpy()): - raise ValueError("pgen output rows do not line up with the cleaned AIRR rows.") - return table, {"source": "computed", "path": str(cache_path)} + hit = _pgen_lookup(cache_path, df["junction_aa"]) + if hit is None: + raise ValueError("pgen output does not cover every cleaned sequence.") + return hit, {"source": "computed", "path": str(cache_path)} def make_splits(n, train_fraction, val_fraction, seed): diff --git a/benchmark/prepare_trb_joint.py b/benchmark/prepare_trb_joint.py new file mode 100644 index 0000000..d958624 --- /dev/null +++ b/benchmark/prepare_trb_joint.py @@ -0,0 +1,648 @@ +"""Prepare the identity-aligned TRB benchmark used by joint codec training. + +The TCRemP Parquet intentionally contains only numeric embedding columns. Its +authoritative row identities live in the representations TSV written beside it +from the same ``analysis_repertoire`` by TCRemP. This preparation therefore +uses that sidecar's ``clone_id`` values, verifies every AIRR and pgen record by +biological identity, and refuses a bare row-count-only alignment. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import logging +import os +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq + + +VALID_AA = set("ACDEFGHIKLMNPQRSTVWY") +IDENTITY_COLUMNS = ["junction_aa", "v_identity", "j_identity", "locus_identity"] +PGEN_COLUMNS = ["log10_pgen", "log10_pgen_1mm"] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--airr-path", required=True) + parser.add_argument("--embeddings-path", required=True) + parser.add_argument("--representations-path", required=True) + parser.add_argument("--pgen-path", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--tcremp-writer-source") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--train-fraction", type=float, default=0.8) + parser.add_argument("--val-fraction", type=float, default=0.1) + parser.add_argument("--min-len", type=int, default=1) + parser.add_argument("--max-len", type=int, default=40) + parser.add_argument("--embedding-batch-size", type=int, default=256) + parser.add_argument("--standardizer-chunk-size", type=int, default=1024) + return parser.parse_args() + + +def setup_logging(path: Path) -> logging.Logger: + path.parent.mkdir(parents=True, exist_ok=True) + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s | %(levelname)s | %(message)s", + handlers=[logging.StreamHandler(), logging.FileHandler(path, encoding="utf-8")], + force=True, + ) + return logging.getLogger("prepare_trb_joint") + + +def normalize_gene(value: object) -> str: + value = str(value).strip().upper() + return value.split("*", 1)[0] + + +def normalize_locus(value: object) -> str: + normalized = str(value).strip().lower() + aliases = {"trb": "beta", "b": "beta"} + return aliases.get(normalized, normalized) + + +def add_identity_columns( + table: pd.DataFrame, + *, + cdr3_col: str, + v_col: str, + j_col: str, + locus_col: str | None = None, + locus_value: str = "beta", +) -> pd.DataFrame: + required = {cdr3_col, v_col, j_col} + if locus_col is not None: + required.add(locus_col) + missing = required.difference(table.columns) + if missing: + raise ValueError(f"Source table is missing identity columns: {sorted(missing)}") + result = table.copy() + result["junction_aa"] = result[cdr3_col].astype(str).str.strip().str.upper() + result["v_identity"] = result[v_col].map(normalize_gene) + result["j_identity"] = result[j_col].map(normalize_gene) + if locus_col is None: + result["locus_identity"] = normalize_locus(locus_value) + else: + result["locus_identity"] = result[locus_col].map(normalize_locus) + return result + + +def identity_strings(table: pd.DataFrame) -> pd.Series: + return table[IDENTITY_COLUMNS].astype(str).agg("\x1f".join, axis=1) + + +def require_unique_identity(table: pd.DataFrame, label: str) -> None: + duplicated = table.duplicated(IDENTITY_COLUMNS, keep=False) + if duplicated.any(): + examples = table.loc[duplicated, IDENTITY_COLUMNS].head(5).to_dict("records") + raise ValueError(f"{label} contains duplicate biological identities: {examples}") + + +def validate_writer_contract(path: str | None) -> dict: + if not path: + raise ValueError( + "--tcremp-writer-source is required: the clone-id sidecar contract must be auditable." + ) + source_path = Path(path) + source = source_path.read_text(encoding="utf-8") + required_fragments = [ + "clone_representations = get_representations_df(analysis_repertoire", + "embeddings = run_tcremp_embedding(\n analysis_repertoire", + "embeddings.to_parquet", + "index=False", + ] + missing = [fragment for fragment in required_fragments if fragment not in source] + if missing: + raise ValueError(f"TCRemP writer contract is not recognizable; missing {missing}") + return { + "path": str(source_path), + "sha256": hashlib.sha256(source.encode("utf-8")).hexdigest(), + "evidence": ( + "representations clone_id and embedding rows are emitted from the same " + "analysis_repertoire before index-free Parquet serialization" + ), + } + + +def prepare_identity_table( + airr: pd.DataFrame, + pgen: pd.DataFrame, + representations: pd.DataFrame, + *, + min_len: int, + max_len: int, +) -> tuple[pd.DataFrame, dict]: + airr = add_identity_columns( + airr, + cdr3_col="junction_aa", + v_col="v_call", + j_col="j_call", + locus_col="locus", + ) + pgen = add_identity_columns( + pgen, + cdr3_col="junction_aa", + v_col="v_call", + j_col="j_call", + locus_col="locus", + ) + representations = add_identity_columns( + representations, + cdr3_col="cdr3aa_TRB", + v_col="v_TRB", + j_col="j_TRB", + locus_value="beta", + ) + if "clone_id" not in representations.columns: + raise ValueError("TCRemP representations must contain clone_id.") + if representations["clone_id"].isna().any() or representations["clone_id"].duplicated().any(): + raise ValueError("TCRemP clone_id values must be complete and unique.") + for column in PGEN_COLUMNS: + if column not in pgen.columns: + raise ValueError(f"pgen table is missing required target {column!r}.") + + require_unique_identity(airr, "AIRR") + require_unique_identity(pgen, "pgen") + require_unique_identity(representations, "TCRemP representations") + + representation_identity = representations[["clone_id", *IDENTITY_COLUMNS]].copy() + joined = airr.merge( + representation_identity, + on=IDENTITY_COLUMNS, + how="inner", + validate="one_to_one", + ) + if len(joined) != len(airr) or len(joined) != len(representations): + raise ValueError( + "AIRR and TCRemP representations do not match one-to-one by CDR3/V/J/locus: " + f"airr={len(airr)} representations={len(representations)} matched={len(joined)}" + ) + joined = joined.merge( + pgen[[*IDENTITY_COLUMNS, *PGEN_COLUMNS]], + on=IDENTITY_COLUMNS, + how="inner", + validate="one_to_one", + ) + if len(joined) != len(airr) or len(joined) != len(pgen): + raise ValueError( + "AIRR and pgen targets do not match one-to-one by CDR3/V/J/locus: " + f"airr={len(airr)} pgen={len(pgen)} matched={len(joined)}" + ) + + report = { + "rows_input": int(len(joined)), + "rows_after_locus_filter": 0, + "dropped_missing": 0, + "dropped_invalid_chars": 0, + "dropped_too_short": 0, + "dropped_too_long": 0, + "dropped_duplicate_cdr3": 0, + "dropped_non_finite_pgen": 0, + } + keep = pd.Series(True, index=joined.index) + locus_ok = joined["locus_identity"].eq("beta") + report["rows_after_locus_filter"] = int(locus_ok.sum()) + keep &= locus_ok + + missing = joined["junction_aa"].isin(["", "NAN", "NONE"]) + report["dropped_missing"] = int((keep & missing).sum()) + keep &= ~missing + invalid = ~joined["junction_aa"].map(lambda sequence: set(sequence) <= VALID_AA) + report["dropped_invalid_chars"] = int((keep & invalid).sum()) + keep &= ~invalid + too_short = joined["junction_aa"].str.len().lt(min_len) + report["dropped_too_short"] = int((keep & too_short).sum()) + keep &= ~too_short + too_long = joined["junction_aa"].str.len().gt(max_len) + report["dropped_too_long"] = int((keep & too_long).sum()) + keep &= ~too_long + joined = joined.loc[keep].copy() + + duplicated = joined["junction_aa"].duplicated(keep="first") + report["dropped_duplicate_cdr3"] = int(duplicated.sum()) + joined = joined.loc[~duplicated].copy() + + target_values = joined[PGEN_COLUMNS].apply(pd.to_numeric, errors="coerce") + finite = np.isfinite(target_values.to_numpy(dtype=np.float64)).all(axis=1) + report["dropped_non_finite_pgen"] = int((~finite).sum()) + joined = joined.loc[finite].copy() + joined[PGEN_COLUMNS] = target_values.loc[finite].to_numpy(dtype=np.float64) + if joined.empty: + raise ValueError("No rows remain after benchmark validation.") + + record_ids = identity_strings(joined).map( + lambda value: hashlib.sha256(value.encode("utf-8")).hexdigest() + ) + joined.insert(0, "record_id", record_ids.to_numpy()) + joined = joined.reset_index(drop=True) + joined.insert(0, "row_index", np.arange(len(joined), dtype=np.int64)) + report["rows_kept"] = int(len(joined)) + return joined, report + + +def make_splits( + table: pd.DataFrame, + *, + train_fraction: float, + val_fraction: float, + seed: int, +) -> tuple[dict[str, np.ndarray], dict]: + if not 0 < train_fraction < 1 or not 0 <= val_fraction < 1: + raise ValueError("Split fractions are invalid.") + if train_fraction + val_fraction >= 1: + raise ValueError("train_fraction + val_fraction must be less than one.") + order = np.random.default_rng(seed).permutation(len(table)) + train_end = int(len(table) * train_fraction) + val_end = train_end + int(len(table) * val_fraction) + splits = { + "train": order[:train_end], + "val": order[train_end:val_end], + "test": order[val_end:], + } + table["split"] = "" + for name, indices in splits.items(): + table.loc[indices, "split"] = name + sets = { + name: set(table.iloc[indices]["junction_aa"].tolist()) + for name, indices in splits.items() + } + overlaps = { + "train_val": len(sets["train"] & sets["val"]), + "train_test": len(sets["train"] & sets["test"]), + "val_test": len(sets["val"] & sets["test"]), + } + if any(overlaps.values()) or (table["split"] == "").any(): + raise ValueError(f"Split construction failed: overlaps={overlaps}") + return splits, overlaps + + +def parquet_embedding_columns(parquet: pq.ParquetFile) -> list[str]: + columns = [ + field.name + for field in parquet.schema_arrow + if pa.types.is_integer(field.type) or pa.types.is_floating(field.type) + ] + if not columns: + raise ValueError("TCRemP Parquet has no numeric embedding columns.") + if "clone_id" in parquet.schema_arrow.names: + raise ValueError( + "This preparation expects identity in the authoritative sidecar; " + "a physical clone_id column requires a dedicated direct-id reader." + ) + return columns + + +def materialize_embeddings( + parquet_path: str, + representations: pd.DataFrame, + table: pd.DataFrame, + output_path: Path, + *, + batch_size: int, +) -> tuple[np.memmap, dict]: + parquet = pq.ParquetFile(parquet_path) + columns = parquet_embedding_columns(parquet) + if parquet.metadata.num_rows != len(representations): + raise ValueError( + "TCRemP Parquet and authoritative clone-id sidecar row counts differ: " + f"parquet={parquet.metadata.num_rows} sidecar={len(representations)}" + ) + source_clone_ids = representations["clone_id"].astype(str).to_numpy() + destination_by_clone = { + str(clone_id): int(row_index) + for clone_id, row_index in zip(table["clone_id"], table["row_index"]) + } + temporary_path = output_path.with_suffix(".npy.tmp") + temporary_path.unlink(missing_ok=True) + output = np.lib.format.open_memmap( + temporary_path, + mode="w+", + dtype=np.float32, + shape=(len(table), len(columns)), + ) + written = np.zeros(len(table), dtype=bool) + source_offset = 0 + try: + for record_batch in parquet.iter_batches(batch_size=batch_size, columns=columns): + matrix = record_batch.to_pandas().to_numpy(dtype=np.float32, copy=False) + if not np.isfinite(matrix).all(): + raise ValueError("TCRemP source embeddings contain NaN or infinite values.") + batch_ids = source_clone_ids[source_offset : source_offset + len(matrix)] + for source_row, clone_id in enumerate(batch_ids): + destination = destination_by_clone.get(clone_id) + if destination is not None: + if written[destination]: + raise ValueError(f"Duplicate embedding clone_id {clone_id!r}.") + output[destination] = matrix[source_row] + written[destination] = True + source_offset += len(matrix) + if source_offset != len(representations) or not written.all(): + missing = table.loc[~written, "clone_id"].head(10).tolist() + raise ValueError( + "Embedding identity materialization was incomplete: " + f"source={source_offset}/{len(representations)} " + f"written={int(written.sum())}/{len(table)} missing={missing}" + ) + output.flush() + output._mmap.close() + temporary_path.replace(output_path) + except Exception: + output._mmap.close() + temporary_path.unlink(missing_ok=True) + raise + matrix = np.load(output_path, mmap_mode="r") + return matrix, { + "source_rows": int(parquet.metadata.num_rows), + "selected_rows": int(len(table)), + "dimension": int(len(columns)), + "dtype": str(matrix.dtype), + "all_finite": True, + "mapping": "TCRemP representations clone_id sidecar writer contract", + } + + +def combine_moments( + count: int, + mean: np.ndarray, + m2: np.ndarray, + batch: np.ndarray, +) -> tuple[int, np.ndarray, np.ndarray]: + batch_count = len(batch) + batch_mean = batch.mean(axis=0, dtype=np.float64) + batch_m2 = batch.var(axis=0, dtype=np.float64) * batch_count + if count == 0: + return batch_count, batch_mean, batch_m2 + delta = batch_mean - mean + total = count + batch_count + return ( + total, + mean + delta * (batch_count / total), + m2 + batch_m2 + delta * delta * (count * batch_count / total), + ) + + +def write_standardizer( + path: Path, + table: pd.DataFrame, + embeddings: np.ndarray, + train_indices: np.ndarray, + *, + chunk_size: int, +) -> dict: + count = 0 + mean = np.zeros(embeddings.shape[1], dtype=np.float64) + m2 = np.zeros(embeddings.shape[1], dtype=np.float64) + for start in range(0, len(train_indices), chunk_size): + indices = train_indices[start : start + chunk_size] + batch = np.asarray(embeddings[indices], dtype=np.float32) + if not np.isfinite(batch).all(): + raise ValueError("TCRemP train targets contain non-finite values.") + count, mean, m2 = combine_moments(count, mean, m2, batch) + raw_std = np.sqrt(m2 / count) + zero_variance = raw_std < 1e-8 + std = np.where(zero_variance, 1.0, raw_std) + + pgen_stats = {} + for column in PGEN_COLUMNS: + values = table.iloc[train_indices][column].to_numpy(dtype=np.float64) + if not np.isfinite(values).all(): + raise ValueError(f"Training target {column} contains non-finite values.") + target_std = float(values.std()) + if target_std < 1e-8: + target_std = 1.0 + pgen_stats[column] = {"mean": float(values.mean()), "std": target_std} + default = pgen_stats["log10_pgen_1mm"] + np.savez( + path, + tcremp_mean=mean.astype(np.float32), + tcremp_std=std.astype(np.float32), + pgen_mean=np.float32(default["mean"]), + pgen_std=np.float32(default["std"]), + pgen_target=np.asarray("log10_pgen_1mm"), + log10_pgen_mean=np.float32(pgen_stats["log10_pgen"]["mean"]), + log10_pgen_std=np.float32(pgen_stats["log10_pgen"]["std"]), + log10_pgen_1mm_mean=np.float32(default["mean"]), + log10_pgen_1mm_std=np.float32(default["std"]), + train_rows=np.int64(count), + ) + return { + "fit_split": "train", + "train_rows": int(count), + "embedding_dimension": int(embeddings.shape[1]), + "zero_variance_embedding_dimensions": int(zero_variance.sum()), + "pgen": pgen_stats, + "default_pgen_target": "log10_pgen_1mm", + } + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def source_info(path: str, *, hash_file: bool) -> dict: + source = Path(path) + stat = source.stat() + result = {"path": str(source), "size_bytes": int(stat.st_size), "mtime_ns": int(stat.st_mtime_ns)} + if hash_file: + result["sha256"] = sha256_file(source) + return result + + +def summarize_target(table: pd.DataFrame, indices: np.ndarray, column: str) -> dict: + values = table.iloc[indices][column].to_numpy(dtype=np.float64) + return { + "rows": int(len(values)), + "finite": int(np.isfinite(values).sum()), + "mean": float(values.mean()), + "std": float(values.std()), + "min": float(values.min()), + "max": float(values.max()), + } + + +def main() -> None: + args = parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + ready_path = output_dir / "READY.json" + if ready_path.exists(): + raise FileExistsError(f"Refusing to overwrite an already ready benchmark: {ready_path}") + log = setup_logging(output_dir / "prepare.log") + log.info("starting identity-aligned TRB benchmark preparation") + writer_contract = validate_writer_contract(args.tcremp_writer_source) + + airr = pd.read_csv(args.airr_path, sep="\t") + pgen = pd.read_csv(args.pgen_path, sep="\t") + representations = pd.read_csv(args.representations_path, sep="\t") + table, cleaning = prepare_identity_table( + airr, + pgen, + representations, + min_len=args.min_len, + max_len=args.max_len, + ) + splits, overlaps = make_splits( + table, + train_fraction=args.train_fraction, + val_fraction=args.val_fraction, + seed=args.seed, + ) + embeddings, embedding_summary = materialize_embeddings( + args.embeddings_path, + representations, + table, + output_dir / "embeddings.npy", + batch_size=args.embedding_batch_size, + ) + normalizer = write_standardizer( + output_dir / "target_standardizer.npz", + table, + embeddings, + splits["train"], + chunk_size=args.standardizer_chunk_size, + ) + + output_columns = [ + "row_index", + "record_id", + "clone_id", + "junction_aa", + "v_call", + "j_call", + "locus", + *PGEN_COLUMNS, + "split", + ] + table[output_columns].to_parquet(output_dir / "dataset.parquet", index=False) + table[output_columns].to_csv(output_dir / "dataset.tsv", sep="\t", index=False) + table[["row_index", "record_id", "clone_id"]].to_csv( + output_dir / "embedding_row_id.tsv", sep="\t", index=False + ) + + manifest_dir = output_dir / "manifests" + manifest_dir.mkdir(exist_ok=True) + manifest_columns = [ + "row_index", + "record_id", + "clone_id", + "junction_aa", + "v_call", + "j_call", + *PGEN_COLUMNS, + ] + for name, indices in splits.items(): + table.iloc[np.sort(indices)][manifest_columns].to_csv( + manifest_dir / f"{name}.tsv", sep="\t", index=False + ) + train = splits["train"] + for name, indices in { + "1k": train[: min(1_000, len(train))], + "10k": train[: min(10_000, len(train))], + "all": train, + }.items(): + pd.DataFrame({"row_index": np.sort(indices)}).to_csv( + manifest_dir / f"train_{name}.tsv", sep="\t", index=False + ) + + summary = { + "slurm": { + "job_id": os.environ.get("SLURM_JOB_ID"), + "node": os.environ.get("SLURMD_NODENAME"), + "cpus_per_task": os.environ.get("SLURM_CPUS_PER_TASK"), + }, + "status": "ready", + "benchmark": "TRB joint IRRM-CODEC", + "seed": args.seed, + "max_len": args.max_len, + "sources": { + "airr": source_info(args.airr_path, hash_file=True), + "tcremp_embeddings": source_info(args.embeddings_path, hash_file=False), + "tcremp_representations": source_info(args.representations_path, hash_file=True), + "pgen": source_info(args.pgen_path, hash_file=True), + "tcremp_writer": writer_contract, + }, + "alignment": { + "method": ( + "TCRemP authoritative representations clone_id sidecar; AIRR and pgen " + "joined one-to-one by normalized CDR3/V/J/locus identity" + ), + "row_order_only": False, + "airr_rows": int(len(airr)), + "representations_rows": int(len(representations)), + "pgen_rows": int(len(pgen)), + "identity_matches_before_cleaning": int(cleaning["rows_input"]), + "selected_embedding_rows": int(embedding_summary["selected_rows"]), + }, + "cleaning": cleaning, + "dataset": { + "rows": int(len(table)), + "unique_record_ids": int(table["record_id"].nunique()), + "unique_clone_ids": int(table["clone_id"].nunique()), + "unique_cdr3": int(table["junction_aa"].nunique()), + "cdr3_length_min": int(table["junction_aa"].str.len().min()), + "cdr3_length_max": int(table["junction_aa"].str.len().max()), + }, + "embeddings": embedding_summary, + "splits": { + "sizes": {name: int(len(indices)) for name, indices in splits.items()}, + "fractions": { + "train": args.train_fraction, + "val": args.val_fraction, + "test": 1.0 - args.train_fraction - args.val_fraction, + }, + "cdr3_overlap": overlaps, + }, + "targets": { + column: { + name: summarize_target(table, indices, column) + for name, indices in splits.items() + } + for column in PGEN_COLUMNS + }, + "standardizer": normalizer, + "checks": { + "identity_alignment": True, + "valid_amino_acids": True, + "length_at_most_40": bool(table["junction_aa"].str.len().le(args.max_len).all()), + "unique_cdr3": bool(table["junction_aa"].is_unique), + "finite_pgen_targets": bool( + np.isfinite(table[PGEN_COLUMNS].to_numpy(dtype=np.float64)).all() + ), + "finite_embeddings": bool(embedding_summary["all_finite"]), + "non_overlapping_splits": not any(overlaps.values()), + "train_only_standardization": normalizer["fit_split"] == "train", + }, + "outputs": { + "dataset_parquet": str(output_dir / "dataset.parquet"), + "dataset_tsv": str(output_dir / "dataset.tsv"), + "embeddings_npy": str(output_dir / "embeddings.npy"), + "embedding_row_ids": str(output_dir / "embedding_row_id.tsv"), + "manifests": str(manifest_dir), + "standardizer": str(output_dir / "target_standardizer.npz"), + "log": str(output_dir / "prepare.log"), + }, + } + if not all(summary["checks"].values()): + raise ValueError(f"Final readiness checks failed: {summary['checks']}") + (output_dir / "dataset_summary.json").write_text( + json.dumps(summary, indent=2), encoding="utf-8" + ) + ready_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") + log.info("benchmark ready at %s", output_dir) + log.info("rows=%d embedding_shape=%s splits=%s", len(table), embeddings.shape, summary["splits"]) + if isinstance(embeddings, np.memmap): + embeddings._mmap.close() + + +if __name__ == "__main__": + main() diff --git a/benchmark/prepare_vdjdb_epitope_cohort.py b/benchmark/prepare_vdjdb_epitope_cohort.py new file mode 100644 index 0000000..0b7a572 --- /dev/null +++ b/benchmark/prepare_vdjdb_epitope_cohort.py @@ -0,0 +1,420 @@ +"""Prepare an auditable, leakage-controlled VDJdb TRB cohort. + +This command intentionally stops before full embedding extraction or evaluation. +It filters and hashes the source, removes exact benchmark overlaps, locks an +objective label rule, and smoke-loads one frozen encoder checkpoint on CPU. +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import math +import os +import re +from collections import Counter, defaultdict +from dataclasses import asdict +from pathlib import Path +from typing import Iterable + +VALID_AA = frozenset("ACDEFGHIKLMNPQRSTVWY") +MISSING = frozenset({"", "NA", "N/A", "NONE", "NULL", "UNKNOWN"}) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--vdjdb-tsv", type=Path, required=True) + parser.add_argument("--benchmark-dir", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--run-config", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--min-support", type=int, default=50) + parser.add_argument("--min-donors", type=int, default=2) + parser.add_argument("--sensitivity-thresholds", type=int, nargs="+", default=[25, 50, 100]) + parser.add_argument("--min-vdjdb-score", type=float, default=1.0) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(chunk_size): + digest.update(chunk) + return digest.hexdigest() + + +def stable_json(path: Path, payload: object) -> None: + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def clean_text(value: object) -> str: + return re.sub(r"\s+", " ", str(value or "").strip()) + + +def normalized_context(value: object) -> str: + text = clean_text(value).upper() + return "" if text in MISSING else text + + +def parse_json_object(raw: str) -> dict: + try: + value = json.loads(raw or "{}") + except json.JSONDecodeError: + return {} + return value if isinstance(value, dict) else {} + + +def normalize_vdjdb_row(row: dict[str, str], min_score: float) -> tuple[dict | None, str]: + if normalized_context(row.get("gene")) != "TRB": + return None, "not_trb" + if normalized_context(row.get("species")) not in {"HOMOSAPIENS", "HOMO SAPIENS"}: + return None, "not_human" + + sequence = normalized_context(row.get("cdr3")) + if not sequence or len(sequence) > 40 or set(sequence).difference(VALID_AA): + return None, "invalid_cdr3" + + fix = parse_json_object(row.get("cdr3fix", "")) + if fix.get("good") is not True: + return None, "cdr3fix_not_good" + fixed_sequence = normalized_context(fix.get("cdr3")) + if fixed_sequence and fixed_sequence != sequence: + return None, "cdr3fix_changes_sequence" + + epitope = normalized_context(row.get("antigen.epitope")) + mhc_a = normalized_context(row.get("mhc.a")) + mhc_b = normalized_context(row.get("mhc.b")) + mhc_class = normalized_context(row.get("mhc.class")) + if not epitope or not mhc_a or not mhc_class: + return None, "missing_epitope_or_mhc" + + try: + score = float(row.get("vdjdb.score", "nan")) + except ValueError: + score = math.nan + if not math.isfinite(score) or score < min_score: + return None, "low_or_invalid_score" + + meta = parse_json_object(row.get("meta", "")) + subject = clean_text(meta.get("subject.id") or meta.get("subject_id")) + reference = clean_text(row.get("reference.id")) + study = clean_text(meta.get("study.id") or meta.get("study_id") or reference) + if not subject: + return None, "missing_donor" + donor_key = f"{study}|{subject}" + label = "|".join((epitope, mhc_a, mhc_b or "-", mhc_class)) + return { + "cdr3": sequence, + "label": label, + "epitope": epitope, + "mhc_a": mhc_a, + "mhc_b": mhc_b, + "mhc_class": mhc_class, + "donor_key": donor_key, + "reference_id": reference, + "v_call": normalized_context(row.get("v.segm")), + "j_call": normalized_context(row.get("j.segm")), + "vdjdb_score": score, + "length": len(sequence), + }, "accepted_row" + + +def read_benchmark_sequences(benchmark_dir: Path) -> tuple[dict[str, set[str]], dict[str, str]]: + split_sequences: dict[str, set[str]] = {} + hashes: dict[str, str] = {} + for split in ("train", "val", "test"): + path = benchmark_dir / "manifests" / f"{split}.tsv" + hashes[str(path)] = sha256_file(path) + sequences: set[str] = set() + with path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle, delimiter="\t") + if "junction_aa" not in (reader.fieldnames or []): + raise ValueError(f"Missing junction_aa in {path}") + for row in reader: + sequences.add(normalized_context(row["junction_aa"])) + split_sequences[split] = sequences + return split_sequences, hashes + + +def prepare_records( + vdjdb_tsv: Path, + benchmark_sequences: dict[str, set[str]], + min_score: float, +) -> tuple[list[dict], dict]: + counts: Counter[str] = Counter() + by_sequence: dict[str, list[dict]] = defaultdict(list) + with vdjdb_tsv.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle, delimiter="\t") + required = { + "gene", "cdr3", "species", "mhc.a", "mhc.b", "mhc.class", + "antigen.epitope", "reference.id", "vdjdb.score", "meta", "cdr3fix", + "v.segm", "j.segm", + } + missing = required.difference(reader.fieldnames or []) + if missing: + raise ValueError(f"VDJdb source is missing columns: {sorted(missing)}") + for row in reader: + counts["source_rows"] += 1 + normalized, reason = normalize_vdjdb_row(row, min_score=min_score) + counts[reason] += 1 + if normalized is not None: + by_sequence[normalized["cdr3"]].append(normalized) + + retained: list[dict] = [] + overlap_by_split: Counter[str] = Counter() + for sequence, rows in sorted(by_sequence.items()): + labels = {row["label"] for row in rows} + if len(labels) != 1: + counts["ambiguous_multilabel_clonotypes"] += 1 + continue + donors = {row["donor_key"] for row in rows} + if len(donors) != 1: + counts["multi_donor_clonotypes"] += 1 + continue + overlapping_splits = [ + split for split, sequences in benchmark_sequences.items() if sequence in sequences + ] + if overlapping_splits: + counts["benchmark_overlap_clonotypes"] += 1 + for split in overlapping_splits: + overlap_by_split[split] += 1 + continue + representative = min( + rows, + key=lambda item: ( + item["label"], item["donor_key"], item["reference_id"], -item["vdjdb_score"] + ), + ).copy() + representative["source_rows"] = len(rows) + retained.append(representative) + + counts["deduplicated_nonoverlap_clonotypes"] = len(retained) + report = { + "filter_counts": dict(sorted(counts.items())), + "benchmark_overlap_by_split": dict(sorted(overlap_by_split.items())), + } + return retained, report + + +def label_summary(records: Iterable[dict]) -> list[dict]: + groups: dict[str, list[dict]] = defaultdict(list) + for record in records: + groups[record["label"]].append(record) + summaries = [] + for label, rows in sorted(groups.items()): + lengths = sorted(row["length"] for row in rows) + summaries.append({ + "label": label, + "epitope": rows[0]["epitope"], + "mhc_a": rows[0]["mhc_a"], + "mhc_b": rows[0]["mhc_b"], + "mhc_class": rows[0]["mhc_class"], + "unique_cdr3": len(rows), + "donors": len({row["donor_key"] for row in rows}), + "length_min": lengths[0], + "length_median": lengths[len(lengths) // 2], + "length_max": lengths[-1], + }) + return summaries + + +def write_tsv(path: Path, rows: list[dict], fields: list[str]) -> None: + with path.open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n") + writer.writeheader() + writer.writerows({field: row.get(field, "") for field in fields} for row in rows) + + +def smoke_checkpoint( + checkpoint_path: Path, + run_config_path: Path, + tokenizer_path: Path, + sequences: list[str], +) -> dict: + import torch + + from irrm_codec.multitask_data import resolve_encoder_tokenizer + from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer + from irrm_codec.tokenization import PAD_ID + + run_config = json.loads(run_config_path.read_text(encoding="utf-8")) + training = run_config["training"] + expected = { + "tokenizer_type": "data_anchor", + "latent_dim": 128, + "encoder_layers": 4, + "decoder_layers": 4, + "tcremp_loss_weight": 1.0, + "pgen_loss_weight": 1.0, + "reconstruction_loss_weight": 1.0, + "seed": 42, + } + observed = {key: training[key] for key in expected} + if observed != expected: + raise ValueError(f"Frozen checkpoint run is not the locked DATA-ANCHOR R+T+P model: {observed}") + + payload = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + model_config = IRRMCodecConfig(**payload["model_config"]) + if asdict(model_config) != run_config["model"]: + raise ValueError("Checkpoint model_config differs from run_config.json") + model = IRRMCodecTransformer(model_config) + model.load_state_dict(payload["model_state"], strict=True) + model.eval() + tokenizer = resolve_encoder_tokenizer("data_anchor", str(tokenizer_path)) + if tokenizer.vocab_size != model_config.input_vocab_size: + raise ValueError("Tokenizer vocabulary does not match checkpoint input embedding") + encoded = [tokenizer.encode(sequence, model_config.max_sequence_len) for sequence in sequences[:2]] + width = max(map(len, encoded)) + tokens = torch.full((len(encoded), width), PAD_ID, dtype=torch.long) + for index, ids in enumerate(encoded): + tokens[index, : len(ids)] = torch.tensor(ids) + with torch.no_grad(): + latent = model.encode(tokens, tokens.ne(PAD_ID)) + if latent.shape != (len(encoded), model_config.latent_dim) or not torch.isfinite(latent).all(): + raise ValueError(f"Frozen encoder smoke output is invalid: {tuple(latent.shape)}") + return { + "checkpoint_epoch_zero_based": int(payload["epoch"]), + "checkpoint_best_val_loss": float(payload["best_val_loss"]), + "model_config": asdict(model_config), + "tokenizer_type": tokenizer.name, + "tokenizer_vocab_size": tokenizer.vocab_size, + "smoke_sequences": sequences[:2], + "smoke_token_lengths": [len(ids) for ids in encoded], + "smoke_latent_shape": list(latent.shape), + "smoke_latent_finite": True, + } + + +def main() -> None: + args = parse_args() + if args.min_support < 2 or args.min_donors < 2: + raise ValueError("Primary support and donor thresholds must both be at least two.") + output = args.output_dir + if output.exists(): + raise FileExistsError(f"Refusing to overwrite existing preflight directory: {output}") + output.mkdir(parents=True) + + benchmark_sequences, benchmark_hashes = read_benchmark_sequences(args.benchmark_dir) + records, inventory = prepare_records( + args.vdjdb_tsv, benchmark_sequences, min_score=args.min_vdjdb_score + ) + summaries = label_summary(records) + eligible_labels = { + row["label"] for row in summaries + if row["unique_cdr3"] >= args.min_support and row["donors"] >= args.min_donors + } + cohort = [row for row in records if row["label"] in eligible_labels] + if len(eligible_labels) < 2: + raise ValueError( + f"Primary rule yielded {len(eligible_labels)} eligible labels; at least two are required." + ) + + fields = [ + "cdr3", "label", "epitope", "mhc_a", "mhc_b", "mhc_class", "donor_key", + "reference_id", "vdjdb_score", "length", "source_rows", + ] + write_tsv(output / "cohort.tsv", cohort, fields) + write_tsv( + output / "label_summary.tsv", + summaries, + ["label", "epitope", "mhc_a", "mhc_b", "mhc_class", "unique_cdr3", "donors", + "length_min", "length_median", "length_max"], + ) + sensitivity = [] + for threshold in sorted(set(args.sensitivity_thresholds + [args.min_support])): + selected = [ + row for row in summaries + if row["unique_cdr3"] >= threshold and row["donors"] >= args.min_donors + ] + sensitivity.append({ + "min_unique_cdr3": threshold, + "min_donors": args.min_donors, + "eligible_labels": len(selected), + "eligible_clonotypes": sum(row["unique_cdr3"] for row in selected), + }) + write_tsv( + output / "threshold_sensitivity.tsv", sensitivity, + ["min_unique_cdr3", "min_donors", "eligible_labels", "eligible_clonotypes"], + ) + + checkpoint = smoke_checkpoint( + args.checkpoint, args.run_config, args.tokenizer, + [row["cdr3"] for row in cohort], + ) + protocol = { + "study_type": "single-cohort preflight; model comparison unresolved", + "primary_label": "antigen.epitope|mhc.a|mhc.b-or-dash|mhc.class", + "primary_selection": { + "all_labels_meeting_rule": True, + "min_unique_cdr3": args.min_support, + "min_donors": args.min_donors, + "vdjdb_score_min": args.min_vdjdb_score, + "exact_cdr3_deduplication": True, + "ambiguous_multilabel_cdr3_removed": True, + "multi_donor_cdr3_removed_for_clustered_inference": True, + "all_locked_benchmark_split_overlaps_removed": True, + }, + "future_metrics_after_model_protocol_acceptance": { + "primary": "per-label within/between cosine-distance ratio; lower is better", + "secondary_distance": "Euclidean", + "pairwise": ["same-label AUROC", "same-label AUPRC with sampled prevalence reported"], + "retrieval": ["donor-excluded precision@1/5/10", "donor-excluded MAP@R"], + "negative_control": "between-label pairs matched on ordered CDR3 lengths", + "nulls": [ + "absolute length-difference score", + "1000 label permutations within length bins and study/donor strata", + ], + "uncertainty": "2000 donor-cluster bootstrap replicates, fixed seed", + "visualization": "descriptive only; no UMAP/scaler/metric learning fit on evaluation data", + }, + "seed": args.seed, + "embedding_extraction_authorized": False, + "blocking_gate": "accept cohort and resolve matched model-comparison protocol", + } + stable_json(output / "evaluation_protocol.json", protocol) + stable_json(output / "inventory.json", inventory) + + manifest = { + "status": "ready_for_cohort_review_not_embedding_extraction", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "source": { + "path": str(args.vdjdb_tsv), + "sha256": sha256_file(args.vdjdb_tsv), + "bytes": args.vdjdb_tsv.stat().st_size, + }, + "benchmark": { + "path": str(args.benchmark_dir), + "manifest_sha256": benchmark_hashes, + "unique_sequences_by_split": { + split: len(values) for split, values in benchmark_sequences.items() + }, + }, + "frozen_checkpoint": { + "path": str(args.checkpoint), + "sha256": sha256_file(args.checkpoint), + "run_config_path": str(args.run_config), + "run_config_sha256": sha256_file(args.run_config), + "tokenizer_path": str(args.tokenizer), + "tokenizer_sha256": sha256_file(args.tokenizer), + **checkpoint, + }, + "cohort": { + "eligible_labels": len(eligible_labels), + "unique_cdr3": len(cohort), + "donors": len({row["donor_key"] for row in cohort}), + "cohort_sha256": sha256_file(output / "cohort.tsv"), + "label_summary_sha256": sha256_file(output / "label_summary.tsv"), + }, + "embedding_extraction_performed": False, + "evaluation_performed": False, + } + stable_json(output / "PREFLIGHT.json", manifest) + print(json.dumps(manifest, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/prepare_vdjdb_redcea_clonotypes.py b/benchmark/prepare_vdjdb_redcea_clonotypes.py new file mode 100644 index 0000000..6580093 --- /dev/null +++ b/benchmark/prepare_vdjdb_redcea_clonotypes.py @@ -0,0 +1,209 @@ +"""Prepare a leakage-controlled, per-clonotype REDCEA VDJdb benchmark. + +The REDCEA member table is used only as a membership/label source. Its UMAP +coordinates and cluster identifiers are retained for provenance but never used +as model features or evaluation targets. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from collections import Counter, defaultdict +from pathlib import Path + +import pandas as pd + +from benchmark.prepare_vdjdb_epitope_cohort import read_benchmark_sequences + + +AA = frozenset("ACDEFGHIKLMNPQRSTVWY") +REQUIRED_COLUMNS = { + "species", "antigen.epitope", "antigen.gene", "antigen.species", + "mhc.a", "mhc.b", "mhc.class", "gene", "cdr3aa", "cid", "csz", + "v.segm", "j.segm", +} + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def stable_json(path: Path, value: object) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def canonical_aa(value: object, max_length: int) -> str | None: + sequence = str(value).strip().upper() + if not sequence or len(sequence) > max_length or any(aa not in AA for aa in sequence): + return None + return sequence + + +def prepare_cohort( + members: pd.DataFrame, + benchmark_sequences: dict[str, set[str]], + min_clonotypes: int, + max_length: int, +) -> tuple[pd.DataFrame, pd.DataFrame, dict]: + missing = REQUIRED_COLUMNS.difference(members.columns) + if missing: + raise ValueError(f"REDCEA member table is missing columns: {sorted(missing)}") + counts: Counter[str] = Counter(source_rows=len(members)) + rows_by_sequence: dict[str, list[dict]] = defaultdict(list) + for raw in members.to_dict(orient="records"): + if str(raw["gene"]).strip().upper() != "TRB": + counts["non_trb_rows"] += 1 + continue + sequence = canonical_aa(raw["cdr3aa"], max_length=max_length) + if sequence is None: + counts["invalid_sequence_rows"] += 1 + continue + label = str(raw["antigen.epitope"]).strip().upper() + if not label: + counts["empty_epitope_rows"] += 1 + continue + rows_by_sequence[sequence].append({ + "cdr3": sequence, + "label": label, + "species": str(raw["species"]).strip(), + "antigen_gene": str(raw["antigen.gene"]).strip(), + "antigen_species": str(raw["antigen.species"]).strip(), + "mhc_a": str(raw["mhc.a"]).strip(), + "mhc_b": str(raw["mhc.b"]).strip(), + "mhc_class": str(raw["mhc.class"]).strip(), + "v_call": str(raw["v.segm"]).strip(), + "j_call": str(raw["j.segm"]).strip(), + "cid": str(raw["cid"]).strip(), + "csz": int(float(raw["csz"])), + "length": len(sequence), + }) + + retained: list[dict] = [] + overlap_by_split: Counter[str] = Counter() + for sequence, source_rows in sorted(rows_by_sequence.items()): + labels = {row["label"] for row in source_rows} + if len(labels) != 1: + counts["ambiguous_multilabel_cdr3"] += 1 + continue + overlapping = [ + split for split, sequences in benchmark_sequences.items() if sequence in sequences + ] + if overlapping: + counts["benchmark_overlap_cdr3"] += 1 + overlap_by_split.update(overlapping) + continue + representative = min( + source_rows, + key=lambda row: ( + row["label"], row["species"], row["cid"], row["v_call"], row["j_call"] + ), + ).copy() + representative["source_rows"] = len(source_rows) + representative["source_cids"] = "|".join(sorted({row["cid"] for row in source_rows})) + retained.append(representative) + + provisional = pd.DataFrame(retained) + support = provisional.groupby("label")["cdr3"].nunique().sort_index() + eligible_labels = set(support[support >= min_clonotypes].index) + cohort = provisional[provisional["label"].isin(eligible_labels)].copy() + cohort.sort_values(["label", "cdr3"], inplace=True, kind="mergesort") + cohort.reset_index(drop=True, inplace=True) + if cohort.empty or cohort["cdr3"].duplicated().any(): + raise ValueError("Prepared cohort is empty or contains duplicate CDR3 sequences.") + if (cohort.groupby("label").size() < min_clonotypes).any(): + raise AssertionError("A retained label is below the preregistered support threshold.") + + summary = ( + cohort.groupby("label", sort=True) + .agg( + unique_cdr3=("cdr3", "size"), + redcea_clusters=("cid", "nunique"), + species=("species", lambda x: "|".join(sorted(set(x)))), + antigen_species=("antigen_species", lambda x: "|".join(sorted(set(x)))), + length_min=("length", "min"), + length_median=("length", "median"), + length_max=("length", "max"), + ) + .reset_index() + ) + report = { + "filter_counts": dict(sorted(counts.items())), + "benchmark_overlap_by_split": dict(sorted(overlap_by_split.items())), + "labels_before_support_gate": int(len(support)), + "labels_after_support_gate": int(cohort["label"].nunique()), + "clonotypes_after_support_gate": int(len(cohort)), + "min_clonotypes_per_epitope": int(min_clonotypes), + "support_before_gate": {str(k): int(v) for k, v in support.items()}, + } + return cohort, summary, report + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--cluster-members", type=Path, required=True) + parser.add_argument("--benchmark-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--min-clonotypes", type=int, default=20) + parser.add_argument("--max-length", type=int, default=40) + parser.add_argument("--source-git-commit", required=True) + parser.add_argument("--source-git-blob", required=True) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.min_clonotypes < 11: + raise ValueError("At least 11 clonotypes are required for precision@10.") + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + benchmark_sequences, benchmark_hashes = read_benchmark_sequences(args.benchmark_dir) + members = pd.read_csv(args.cluster_members, sep="\t", dtype=str, keep_default_na=False) + cohort, summary, report = prepare_cohort( + members, benchmark_sequences, args.min_clonotypes, args.max_length + ) + cohort_path = args.output_dir / "cohort.tsv" + summary_path = args.output_dir / "epitope_support.tsv" + cohort.to_csv(cohort_path, sep="\t", index=False, lineterminator="\n") + summary.to_csv(summary_path, sep="\t", index=False, lineterminator="\n") + manifest = { + "status": "accepted_for_frozen_embedding_evaluation", + "study": "REDCEA motif-member per-clonotype epitope geometry", + "source": { + "path": str(args.cluster_members), + "sha256": sha256(args.cluster_members), + "git_commit": args.source_git_commit, + "git_blob": args.source_git_blob, + "expected_repository_path": "results/redcea/cluster_members_TRB.txt", + }, + "cohort": { + "path": str(cohort_path), "sha256": sha256(cohort_path), + "rows": len(cohort), "labels": int(cohort["label"].nunique()), + "species": sorted(cohort["species"].unique().tolist()), + "minimum_label_support": args.min_clonotypes, + }, + "epitope_support": {"path": str(summary_path), "sha256": sha256(summary_path)}, + "benchmark_manifest_sha256": benchmark_hashes, + "filters": report, + "protocol": { + "unit": "unique exact TRB CDR3 amino-acid clonotype", + "label": "exact antigen.epitope from REDCEA member table", + "donors_and_publications": "pooled; not used for filtering or splitting", + "ambiguous_multilabel_cdr3": "removed entirely", + "benchmark_train_val_test_overlap": "removed entirely", + "redcea_cid": "provenance only; not an evaluation unit or feature", + "redcea_x_y": "excluded; UMAP visualization coordinates are never features", + }, + } + stable_json(args.output_dir / "PREFLIGHT.json", manifest) + print(json.dumps(manifest, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/prepare_vdjdb_viral_motif_top10.py b/benchmark/prepare_vdjdb_viral_motif_top10.py new file mode 100644 index 0000000..f65f903 --- /dev/null +++ b/benchmark/prepare_vdjdb_viral_motif_top10.py @@ -0,0 +1,453 @@ +"""Prepare an exploratory top-10 viral official-motif VDJdb benchmark.""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import os +from collections import Counter, defaultdict +from pathlib import Path + +from benchmark.prepare_vdjdb_epitope_cohort import ( + VALID_AA, + clean_text, + normalized_context, + parse_json_object, + read_benchmark_sequences, + sha256_file, + stable_json, + write_tsv, +) +from benchmark.prepare_vdjdb_ylq_motif_pooled import ( + TARGET_CONTEXT, + build_pairs, + select_compact_negatives, +) + + +def read_species_list(path: Path) -> list[str]: + values = [] + for line in path.read_text(encoding="utf-8").splitlines(): + value = line.strip() + if value and not value.startswith("#"): + values.append(value) + if not values or len(values) != len(set(values)): + raise ValueError("Viral antigen-species list must be nonempty and unique.") + return values + + +def read_official_members(path: Path) -> tuple[dict[tuple[str, str, str], set[str]], dict]: + counts: Counter[str] = Counter() + membership: dict[tuple[str, str, str], set[str]] = defaultdict(set) + with path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle, delimiter="\t") + required = { + "species", "antigen.epitope", "antigen.species", "mhc.a", "mhc.b", + "mhc.class", "gene", "cdr3aa", "cid", "v.segm", "j.segm", + } + missing = required.difference(reader.fieldnames or []) + if missing: + raise ValueError(f"Official cluster-members file lacks columns: {sorted(missing)}") + for row in reader: + counts["source_rows"] += 1 + if normalized_context(row["species"]) != "HOMOSAPIENS": + continue + if normalized_context(row["gene"]) != "TRB": + continue + context = ( + normalized_context(row["mhc.a"]), normalized_context(row["mhc.b"]), + normalized_context(row["mhc.class"]), + ) + if context != TARGET_CONTEXT: + continue + counts["exact_context_trb_rows"] += 1 + sequence = normalized_context(row["cdr3aa"]) + peptide = normalized_context(row["antigen.epitope"]) + antigen_species = clean_text(row["antigen.species"]) + cid = clean_text(row["cid"]) + if not sequence or len(sequence) > 40 or set(sequence).difference(VALID_AA): + counts["invalid_member_cdr3"] += 1 + continue + if not peptide or not antigen_species or not cid: + counts["missing_member_label_species_or_cid"] += 1 + continue + membership[(sequence, peptide, antigen_species)].add(cid) + counts["valid_exact_context_membership_rows"] += 1 + return membership, { + "counts": dict(sorted(counts.items())), + "unique_sequence_peptide_species_memberships": len(membership), + } + + +def read_joined_records( + vdjdb_path: Path, + membership: dict[tuple[str, str, str], set[str]], + benchmark: dict[str, set[str]], + min_score: float, +) -> tuple[list[dict], dict]: + counts: Counter[str] = Counter() + overlaps: Counter[str] = Counter() + grouped: dict[str, list[dict]] = defaultdict(list) + with vdjdb_path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle, delimiter="\t") + required = { + "gene", "cdr3", "v.segm", "j.segm", "species", "mhc.a", "mhc.b", + "mhc.class", "antigen.epitope", "antigen.species", "reference.id", + "vdjdb.score", "meta", "cdr3fix", + } + missing = required.difference(reader.fieldnames or []) + if missing: + raise ValueError(f"Pinned VDJdb table lacks columns: {sorted(missing)}") + for row in reader: + counts["source_rows"] += 1 + if normalized_context(row["gene"]) != "TRB" \ + or normalized_context(row["species"]) != "HOMOSAPIENS": + continue + context = ( + normalized_context(row["mhc.a"]), normalized_context(row["mhc.b"]), + normalized_context(row["mhc.class"]), + ) + if context != TARGET_CONTEXT: + continue + counts["exact_context_trb_rows"] += 1 + sequence = normalized_context(row["cdr3"]) + if not sequence or len(sequence) > 40 or set(sequence).difference(VALID_AA): + counts["invalid_cdr3"] += 1 + continue + fix = parse_json_object(row.get("cdr3fix", "")) + if fix.get("good") is not True: + counts["cdr3fix_not_good"] += 1 + continue + fixed = normalized_context(fix.get("cdr3")) + if fixed and fixed != sequence: + counts["cdr3fix_changes_sequence"] += 1 + continue + try: + score = float(row["vdjdb.score"]) + except ValueError: + score = math.nan + if not math.isfinite(score) or score < min_score: + counts["low_or_invalid_score"] += 1 + continue + peptide = normalized_context(row["antigen.epitope"]) + antigen_species = clean_text(row["antigen.species"]) + key = (sequence, peptide, antigen_species) + if key not in membership: + counts["not_exact_official_motif_member"] += 1 + continue + counts["exact_official_membership_join_rows"] += 1 + overlapping = [split for split, sequences in benchmark.items() if sequence in sequences] + if overlapping: + counts["benchmark_overlap_rows"] += 1 + for split in overlapping: + overlaps[split] += 1 + continue + meta = parse_json_object(row.get("meta", "")) + subject = clean_text(meta.get("subject.id") or meta.get("subject_id")) + study = clean_text(meta.get("study.id") or meta.get("study_id") or row["reference.id"]) + grouped[sequence].append({ + "cdr3": sequence, + "epitope": peptide, + "antigen_species": antigen_species, + "mhc_a": TARGET_CONTEXT[0], + "mhc_b": TARGET_CONTEXT[1], + "mhc_class": TARGET_CONTEXT[2], + "v_call": normalized_context(row["v.segm"]), + "j_call": normalized_context(row["j.segm"]), + "reference_id": clean_text(row["reference.id"]), + "vdjdb_score": score, + "donor_key": f"{study}|{subject}" if subject else "", + "motif_cids": membership[key], + "length": len(sequence), + }) + + retained = [] + for sequence, rows in sorted(grouped.items()): + labels = {(row["epitope"], row["antigen_species"]) for row in rows} + if len(labels) != 1: + counts["ambiguous_multilabel_or_species_cdr3"] += 1 + continue + representative = min( + rows, + key=lambda row: ( + -row["vdjdb_score"], row["epitope"], row["antigen_species"], + row["v_call"], row["j_call"], row["reference_id"], + ), + ).copy() + representative["source_rows"] = len(rows) + representative["donor_count_descriptive_only"] = len({ + row["donor_key"] for row in rows if row["donor_key"] + }) + representative["motif_cids"] = ",".join(sorted({ + cid for row in rows for cid in row["motif_cids"] + })) + retained.append(representative) + counts["pooled_exact_cdr3_after_global_dedup"] = len(retained) + return retained, { + "filter_counts": dict(sorted(counts.items())), + "benchmark_overlap_rows_by_split": dict(sorted(overlaps.items())), + } + + +def rank_epitopes(records: list[dict], viral_species: set[str]) -> tuple[list[dict], list[dict]]: + species_counts = Counter(record["antigen_species"] for record in records) + viral = [record for record in records if record["antigen_species"] in viral_species] + by_epitope: dict[str, list[dict]] = defaultdict(list) + for record in viral: + by_epitope[record["epitope"]].append(record) + ranking = [] + for peptide, rows in by_epitope.items(): + ranking.append({ + "epitope": peptide, + "unique_cdr3": len(rows), + "antigen_species": ",".join(sorted({row["antigen_species"] for row in rows})), + "species_count": len({row["antigen_species"] for row in rows}), + }) + ranking.sort(key=lambda row: (-row["unique_cdr3"], row["epitope"])) + inventory = [ + { + "antigen_species": species, + "unique_cdr3": count, + "classified_viral_by_exact_allowlist": int(species in viral_species), + } + for species, count in sorted(species_counts.items(), key=lambda item: (-item[1], item[0])) + ] + return ranking, inventory + + +def choose_top10(ranking: list[dict], requested_threshold: int) -> tuple[list[dict], dict]: + requested_epitopes = 10 + evaluable = [row for row in ranking if row["unique_cdr3"] >= 2] + if not evaluable: + raise ValueError("No viral epitope has two distinct post-filter CDR3s.") + passing = [row for row in evaluable if row["unique_cdr3"] >= requested_threshold] + if len(passing) >= requested_epitopes: + threshold = requested_threshold + relaxed = False + selected_count = requested_epitopes + elif len(evaluable) >= requested_epitopes: + threshold = int(evaluable[requested_epitopes - 1]["unique_cdr3"]) + relaxed = True + selected_count = requested_epitopes + else: + threshold = int(evaluable[-1]["unique_cdr3"]) + relaxed = True + selected_count = len(evaluable) + selected = [row.copy() for row in evaluable if row["unique_cdr3"] >= threshold][ + :selected_count + ] + for index, row in enumerate(selected, start=1): + row["rank"] = index + return selected, { + "requested_minimum_unique_cdr3": requested_threshold, + "epitopes_meeting_requested_threshold": len(passing), + "requested_epitopes": requested_epitopes, + "selected_epitopes": len(selected), + "evaluable_epitope_shortfall": requested_epitopes - len(selected), + "all_post_filter_viral_epitopes": len(ranking), + "viral_epitopes_with_two_distinct_cdr3": len(evaluable), + "viral_epitopes_below_distinct_pair_minimum": len(ranking) - len(evaluable), + "effective_minimum_unique_cdr3": threshold, + "threshold_transparently_lowered": relaxed, + "selection_rule": ( + "post-filter unique CDR3 descending, then exact peptide lexical; at least two " + "distinct target CDR3s required for within-target pairs" + ), + "model_metrics_used_for_selection": False, + } + + +def write_epitope_dataset( + output_dir: Path, + rank: int, + epitope: str, + positives: list[dict], + eligible_negatives: list[dict], + controls_per_positive: int, + max_pairs: int, + seed: int, +) -> dict: + directory = output_dir / f"epitope-{rank:02d}-{epitope}" + directory.mkdir() + negatives, assignments, strata, control_report = select_compact_negatives( + positives, eligible_negatives, controls_per_positive, seed + rank * 1000 + ) + candidates = [] + fields = [ + "candidate_index", "is_target", "cdr3", "epitope", "antigen_species", + "mhc_a", "mhc_b", "mhc_class", "v_call", "j_call", "reference_id", + "vdjdb_score", "length", "motif_cids", "source_rows", + "donor_count_descriptive_only", + ] + for index, record in enumerate([*positives, *negatives]): + candidate = dict(record) + candidate["candidate_index"] = index + candidate["is_target"] = int(index < len(positives)) + candidates.append(candidate) + if len({row["cdr3"] for row in candidates}) != len(candidates): + raise ValueError(f"{epitope} candidate set is not exact-CDR3 deduplicated.") + write_tsv(directory / "candidates.tsv", candidates, fields) + write_tsv(directory / "control_selection.tsv", assignments, list(assignments[0])) + write_tsv(directory / "control_length_strata.tsv", strata, list(strata[0])) + pairs, matching = build_pairs(positives, negatives, seed + rank * 1000) + pairs_before_cap = len(pairs) + if len(pairs) > max_pairs: + pairs = pairs[:max_pairs] + matching["pairs_before_deterministic_cap"] = pairs_before_cap + matching["primary_pair_cap"] = max_pairs + matching["pairs_after_deterministic_cap"] = len(pairs) + matching["positive_pairs_used_for_evaluation"] = len(pairs) + matching["pair_cap_rule"] = "first seeded-shuffled distinct positive pairs" + write_tsv(directory / "pair_manifest.tsv", pairs, list(pairs[0])) + report = { + "rank": rank, + "epitope": epitope, + "positive_cdr3": len(positives), + "selected_unique_negative_cdr3": len(negatives), + "eligible_viral_negative_cdr3_before_compaction": len(eligible_negatives), + "candidate_cdr3": len(candidates), + "matching": matching, + "compact_control_selection": control_report, + "artifacts": { + "candidates_sha256": sha256_file(directory / "candidates.tsv"), + "pairs_sha256": sha256_file(directory / "pair_manifest.tsv"), + "control_selection_sha256": sha256_file(directory / "control_selection.tsv"), + "control_length_strata_sha256": sha256_file(directory / "control_length_strata.tsv"), + }, + } + stable_json(directory / "PREFLIGHT.json", report) + return {"report": report, "candidates": candidates} + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--vdjdb-tsv", type=Path, required=True) + parser.add_argument("--cluster-members", type=Path, required=True) + parser.add_argument("--motif-pwms", type=Path, required=True) + parser.add_argument("--benchmark-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--viral-species-list", type=Path) + parser.add_argument("--inventory-only", action="store_true") + parser.add_argument("--min-score", type=float, default=1.0) + parser.add_argument("--requested-min-positives", type=int, default=40) + parser.add_argument("--controls-per-positive", type=int, default=4) + parser.add_argument("--max-pairs-per-epitope", type=int, default=5000) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + + benchmark, benchmark_hashes = read_benchmark_sequences(args.benchmark_dir) + membership, motif_report = read_official_members(args.cluster_members) + records, join_report = read_joined_records( + args.vdjdb_tsv, membership, benchmark, args.min_score + ) + all_species = Counter(record["antigen_species"] for record in records) + write_tsv( + args.output_dir / "antigen_species_inventory.tsv", + [{"antigen_species": key, "unique_cdr3": value} + for key, value in sorted(all_species.items(), key=lambda item: (-item[1], item[0]))], + ["antigen_species", "unique_cdr3"], + ) + if args.inventory_only: + report = { + "status": "inventory_complete", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "exact_context": list(TARGET_CONTEXT), + "unique_antigen_species": len(all_species), + "retained_unique_cdr3": len(records), + "motif_inventory": motif_report, + "vdjdb_join_and_quality": join_report, + "artifact_sha256": sha256_file(args.output_dir / "antigen_species_inventory.tsv"), + } + stable_json(args.output_dir / "INVENTORY.json", report) + print(json.dumps(report, indent=2, sort_keys=True)) + return + if args.viral_species_list is None: + raise ValueError("--viral-species-list is required outside inventory-only mode.") + viral_species_values = read_species_list(args.viral_species_list) + ranking, inventory = rank_epitopes(records, set(viral_species_values)) + write_tsv(args.output_dir / "antigen_species_classification.tsv", inventory, list(inventory[0])) + write_tsv(args.output_dir / "viral_epitope_candidates.tsv", ranking, list(ranking[0])) + selected, gate = choose_top10(ranking, args.requested_min_positives) + write_tsv(args.output_dir / "selected_top10.tsv", selected, list(selected[0])) + + viral_records = [record for record in records + if record["antigen_species"] in set(viral_species_values)] + datasets = [] + master: dict[str, dict] = {} + for selected_row in selected: + epitope = selected_row["epitope"] + positives = [record for record in viral_records if record["epitope"] == epitope] + negatives = [record for record in viral_records if record["epitope"] != epitope] + dataset = write_epitope_dataset( + args.output_dir, selected_row["rank"], epitope, positives, negatives, + args.controls_per_positive, args.max_pairs_per_epitope, args.seed, + ) + datasets.append(dataset["report"]) + for record in dataset["candidates"]: + master.setdefault(record["cdr3"], record) + master_rows = [] + for index, cdr3 in enumerate(sorted(master)): + row = dict(master[cdr3]) + row["master_index"] = index + master_rows.append(row) + write_tsv( + args.output_dir / "master_candidates.tsv", master_rows, + ["master_index", "cdr3", "epitope", "antigen_species", "mhc_a", "mhc_b", + "mhc_class", "v_call", "j_call", "reference_id", "vdjdb_score", "length", + "motif_cids", "source_rows", "donor_count_descriptive_only"], + ) + manifest = { + "status": "ready", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "exact_context": {"mhc_a": TARGET_CONTEXT[0], "mhc_b": TARGET_CONTEXT[1], + "mhc_class": TARGET_CONTEXT[2]}, + "pooling": "all records pooled; donor metadata descriptive only", + "viral_species_rule": ( + "Exact, case-sensitive antigen.species membership in the pinned allowlist; no " + "substring, epitope, antigen-gene, or model-score inference." + ), + "viral_antigen_species_exact": viral_species_values, + "support_gate": gate, + "selected": datasets, + "master_candidate_cdr3": len(master_rows), + "motif_inventory": motif_report, + "vdjdb_join_and_quality": join_report, + "sources": { + "vdjdb": {"path": str(args.vdjdb_tsv), "sha256": sha256_file(args.vdjdb_tsv)}, + "cluster_members": {"path": str(args.cluster_members), + "sha256": sha256_file(args.cluster_members)}, + "motif_pwms": {"path": str(args.motif_pwms), + "sha256": sha256_file(args.motif_pwms)}, + "viral_species_list": {"path": str(args.viral_species_list), + "sha256": sha256_file(args.viral_species_list)}, + "benchmark_manifests": benchmark_hashes, + }, + "artifacts": { + "selected_top10_sha256": sha256_file(args.output_dir / "selected_top10.tsv"), + "candidate_ranking_sha256": sha256_file( + args.output_dir / "viral_epitope_candidates.tsv" + ), + "master_candidates_sha256": sha256_file(args.output_dir / "master_candidates.tsv"), + }, + "official_motif_rule": ( + "Exact sequence+peptide+antigen.species membership in official cluster_members.txt " + "after exact HomoSapiens/TRB/HLA-A*02:01/B2M/MHCI filtering." + ), + "circularity_caveat": ( + "Official motif membership is sequence-cluster-conditioned. This is an exploratory " + "within-VDJdb motif/epitope-organization analysis, not evidence independent of " + "sequence motif or composition." + ), + "embedding_extraction_performed": False, + } + stable_json(args.output_dir / "PREFLIGHT.json", manifest) + print(json.dumps(manifest, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/prepare_vdjdb_ylq.py b/benchmark/prepare_vdjdb_ylq.py new file mode 100644 index 0000000..1f8ca75 --- /dev/null +++ b/benchmark/prepare_vdjdb_ylq.py @@ -0,0 +1,216 @@ +"""Prepare an exact-peptide, dominant-MHC YLQ VDJdb evaluation cohort.""" + +from __future__ import annotations + +import argparse +import csv +import json +import os +from collections import Counter, defaultdict +from pathlib import Path + +import numpy as np + +from benchmark.prepare_vdjdb_epitope_cohort import ( + prepare_records, + read_benchmark_sequences, + sha256_file, + stable_json, + write_tsv, +) + + +TARGET_PEPTIDE = "YLQPRTFLL" +AA = "ACDEFGHIKLMNPQRSTVWY" + + +def composition(sequence: str) -> np.ndarray: + return np.asarray([sequence.count(aa) / len(sequence) for aa in AA], dtype=float) + + +def context(record: dict) -> tuple[str, str, str]: + return record["mhc_a"], record["mhc_b"] or "-", record["mhc_class"] + + +def build_pairs(positives: list[dict], negatives: list[dict], seed: int) -> tuple[list[dict], dict]: + rng = np.random.default_rng(seed) + positive_pairs = [ + (left, right) + for left in range(len(positives)) + for right in range(left + 1, len(positives)) + if positives[left]["donor_key"] != positives[right]["donor_key"] + ] + rng.shuffle(positive_pairs) + negative_by_length: dict[int, list[int]] = defaultdict(list) + negative_composition = [] + for index, record in enumerate(negatives): + negative_by_length[int(record["length"])].append(index) + negative_composition.append(composition(record["cdr3"])) + negative_composition = np.asarray(negative_composition) + usage: Counter[int] = Counter() + rows = [] + for pair_number, (query, positive) in enumerate(positive_pairs): + if rng.integers(2): + query, positive = positive, query + query_record = positives[query] + positive_record = positives[positive] + candidates = [ + index for index in negative_by_length[int(positive_record["length"])] + if negatives[index]["donor_key"] not in { + query_record["donor_key"], positive_record["donor_key"] + } + ] + if not candidates: + raise ValueError( + f"No exact-length donor-disjoint negative for positive length {positive_record['length']}." + ) + target_composition = composition(positive_record["cdr3"]) + def rank(index: int) -> tuple: + candidate = negatives[index] + v_missing = not positive_record["v_call"] or not candidate["v_call"] + j_missing = not positive_record["j_call"] or not candidate["j_call"] + v_mismatch = int(v_missing or candidate["v_call"] != positive_record["v_call"]) + j_mismatch = int(j_missing or candidate["j_call"] != positive_record["j_call"]) + comp_l1 = float(np.abs(negative_composition[index] - target_composition).sum()) + return v_mismatch + j_mismatch, usage[index], comp_l1, candidate["cdr3"] + negative = min(candidates, key=rank) + usage[negative] += 1 + negative_record = negatives[negative] + rows.append({ + "pair_id": pair_number, + "query_index": query, + "positive_index": positive, + "negative_index": len(positives) + negative, + "query_donor": query_record["donor_key"], + "positive_donor": positive_record["donor_key"], + "negative_donor": negative_record["donor_key"], + "query_length": query_record["length"], + "positive_length": positive_record["length"], + "negative_length": negative_record["length"], + "v_exact_match": int( + bool(positive_record["v_call"]) + and positive_record["v_call"] == negative_record["v_call"] + ), + "j_exact_match": int( + bool(positive_record["j_call"]) + and positive_record["j_call"] == negative_record["j_call"] + ), + "vj_exact_match": int( + bool(positive_record["v_call"] and positive_record["j_call"]) + and positive_record["v_call"] == negative_record["v_call"] + and positive_record["j_call"] == negative_record["j_call"] + ), + "composition_l1": float( + np.abs(negative_composition[negative] - target_composition).sum() + ), + }) + if not rows: + raise ValueError("No donor-disjoint positive YLQ pairs are available.") + report = { + "positive_pairs": len(rows), + "exact_length_fraction": 1.0, + "v_exact_match_fraction": float(np.mean([row["v_exact_match"] for row in rows])), + "j_exact_match_fraction": float(np.mean([row["j_exact_match"] for row in rows])), + "vj_exact_match_fraction": float(np.mean([row["vj_exact_match"] for row in rows])), + "mean_composition_l1": float(np.mean([row["composition_l1"] for row in rows])), + "unique_negatives_used": len(usage), + "max_negative_reuse": max(usage.values()), + } + return rows, report + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--vdjdb-tsv", type=Path, required=True) + parser.add_argument("--benchmark-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--min-positive-cdr3", type=int, default=50) + parser.add_argument("--min-positive-donors", type=int, default=2) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + + benchmark, benchmark_hashes = read_benchmark_sequences(args.benchmark_dir) + records, inventory = prepare_records(args.vdjdb_tsv, benchmark, min_score=1.0) + exact_peptide = [record for record in records if record["epitope"] == TARGET_PEPTIDE] + distribution = [] + by_context: dict[tuple[str, str, str], list[dict]] = defaultdict(list) + for record in exact_peptide: + by_context[context(record)].append(record) + for key, rows in sorted(by_context.items(), key=lambda item: (-len(item[1]), item[0])): + distribution.append({ + "mhc_a": key[0], "mhc_b": key[1], "mhc_class": key[2], + "unique_cdr3": len(rows), "donors": len({row["donor_key"] for row in rows}), + }) + if not distribution: + raise ValueError("No exact YLQPRTFLL records remain after leakage filtering.") + dominant = distribution[0] + dominant_context = (dominant["mhc_a"], dominant["mhc_b"], dominant["mhc_class"]) + if dominant["unique_cdr3"] < args.min_positive_cdr3 or dominant["donors"] < args.min_positive_donors: + stable_json(args.output_dir / "BLOCKED.json", { + "status": "blocked_insufficient_support", "mhc_distribution": distribution, + }) + raise ValueError("Dominant exact YLQ peptide+MHC context has insufficient support.") + + positives = [ + record for record in exact_peptide if context(record) == dominant_context + ] + negatives = [ + record for record in records + if record["epitope"] != TARGET_PEPTIDE and context(record) == dominant_context + ] + if len(negatives) < 50: + raise ValueError(f"Only {len(negatives)} same-context non-YLQ negatives remain.") + candidates = [] + candidate_fields = [ + "candidate_index", "is_ylq", "cdr3", "epitope", "mhc_a", "mhc_b", "mhc_class", + "donor_key", "reference_id", "v_call", "j_call", "vdjdb_score", "length", "source_rows", + ] + for index, record in enumerate([*positives, *negatives]): + candidate = dict(record) + candidate["candidate_index"] = index + candidate["is_ylq"] = int(index < len(positives)) + candidates.append(candidate) + write_tsv(args.output_dir / "candidates.tsv", candidates, candidate_fields) + write_tsv( + args.output_dir / "mhc_distribution.tsv", distribution, + ["mhc_a", "mhc_b", "mhc_class", "unique_cdr3", "donors"], + ) + pairs, match_report = build_pairs(positives, negatives, seed=args.seed) + pair_fields = list(pairs[0]) + write_tsv(args.output_dir / "pair_manifest.tsv", pairs, pair_fields) + stable_json(args.output_dir / "inventory.json", inventory) + manifest = { + "status": "ready", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "target_peptide_definition": "exact equality to YLQPRTFLL", + "primary_context_rule": "dominant exact peptide+MHC context after all filters", + "primary_context": { + **dominant, + "label": "|".join(dominant_context), + }, + "mhc_distribution": distribution, + "positive_cdr3": len(positives), + "positive_donors": len({row["donor_key"] for row in positives}), + "same_context_non_ylq_negative_cdr3": len(negatives), + "same_context_non_ylq_negative_donors": len({row["donor_key"] for row in negatives}), + "matching": match_report, + "leakage": inventory, + "source": {"path": str(args.vdjdb_tsv), "sha256": sha256_file(args.vdjdb_tsv)}, + "benchmark_manifest_sha256": benchmark_hashes, + "artifacts": { + "candidates_sha256": sha256_file(args.output_dir / "candidates.tsv"), + "pairs_sha256": sha256_file(args.output_dir / "pair_manifest.tsv"), + "mhc_distribution_sha256": sha256_file(args.output_dir / "mhc_distribution.tsv"), + }, + "motif_expansion_performed": False, + "embedding_extraction_performed": False, + } + stable_json(args.output_dir / "PREFLIGHT.json", manifest) + print(json.dumps(manifest, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/prepare_vdjdb_ylq_motif_pooled.py b/benchmark/prepare_vdjdb_ylq_motif_pooled.py new file mode 100644 index 0000000..701d984 --- /dev/null +++ b/benchmark/prepare_vdjdb_ylq_motif_pooled.py @@ -0,0 +1,458 @@ +"""Prepare pooled exact-YLQ evaluation data from official VDJdb motif members.""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import os +from collections import Counter, defaultdict +from pathlib import Path + +import numpy as np + +from benchmark.prepare_vdjdb_epitope_cohort import ( + VALID_AA, + clean_text, + normalized_context, + parse_json_object, + read_benchmark_sequences, + sha256_file, + stable_json, + write_tsv, +) + + +TARGET_PEPTIDE = "YLQPRTFLL" +TARGET_CONTEXT = ("HLA-A*02:01", "B2M", "MHCI") +AA = "ACDEFGHIKLMNPQRSTVWY" + + +def composition(sequence: str) -> np.ndarray: + return np.asarray([sequence.count(aa) / len(sequence) for aa in AA], dtype=float) + + +def read_official_members(path: Path) -> tuple[dict[tuple[str, str], set[str]], dict]: + counts: Counter[str] = Counter() + membership: dict[tuple[str, str], set[str]] = defaultdict(set) + with path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle, delimiter="\t") + required = { + "species", "antigen.epitope", "mhc.a", "mhc.b", "mhc.class", "gene", + "cdr3aa", "cid", "csz", "v.segm", "j.segm", + } + missing = required.difference(reader.fieldnames or []) + if missing: + raise ValueError(f"Official cluster-members file lacks columns: {sorted(missing)}") + for row in reader: + counts["source_rows"] += 1 + if normalized_context(row["species"]) != "HOMOSAPIENS": + continue + if normalized_context(row["gene"]) != "TRB": + continue + context = ( + normalized_context(row["mhc.a"]), normalized_context(row["mhc.b"]), + normalized_context(row["mhc.class"]), + ) + if context != TARGET_CONTEXT: + continue + counts["exact_context_trb_rows"] += 1 + sequence = normalized_context(row["cdr3aa"]) + peptide = normalized_context(row["antigen.epitope"]) + cid = clean_text(row["cid"]) + if not sequence or len(sequence) > 40 or set(sequence).difference(VALID_AA): + counts["invalid_member_cdr3"] += 1 + continue + if not peptide or not cid: + counts["missing_member_label_or_cid"] += 1 + continue + membership[(sequence, peptide)].add(cid) + counts["valid_exact_context_membership_rows"] += 1 + if peptide == TARGET_PEPTIDE: + counts["exact_ylq_membership_rows"] += 1 + report = { + "counts": dict(sorted(counts.items())), + "unique_sequence_peptide_memberships": len(membership), + "unique_ylq_member_cdr3": len({key[0] for key in membership if key[1] == TARGET_PEPTIDE}), + } + return membership, report + + +def read_joined_records( + vdjdb_path: Path, + membership: dict[tuple[str, str], set[str]], + benchmark: dict[str, set[str]], + min_score: float, +) -> tuple[list[dict], dict]: + counts: Counter[str] = Counter() + overlap: Counter[str] = Counter() + grouped: dict[str, list[dict]] = defaultdict(list) + with vdjdb_path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle, delimiter="\t") + required = { + "gene", "cdr3", "v.segm", "j.segm", "species", "mhc.a", "mhc.b", + "mhc.class", "antigen.epitope", "reference.id", "vdjdb.score", "meta", "cdr3fix", + } + missing = required.difference(reader.fieldnames or []) + if missing: + raise ValueError(f"Pinned VDJdb table lacks columns: {sorted(missing)}") + for row in reader: + counts["source_rows"] += 1 + if normalized_context(row["gene"]) != "TRB": + continue + if normalized_context(row["species"]) != "HOMOSAPIENS": + continue + context = ( + normalized_context(row["mhc.a"]), normalized_context(row["mhc.b"]), + normalized_context(row["mhc.class"]), + ) + if context != TARGET_CONTEXT: + continue + counts["exact_context_trb_rows"] += 1 + sequence = normalized_context(row["cdr3"]) + if not sequence or len(sequence) > 40 or set(sequence).difference(VALID_AA): + counts["invalid_cdr3"] += 1 + continue + fix = parse_json_object(row.get("cdr3fix", "")) + if fix.get("good") is not True: + counts["cdr3fix_not_good"] += 1 + continue + fixed = normalized_context(fix.get("cdr3")) + if fixed and fixed != sequence: + counts["cdr3fix_changes_sequence"] += 1 + continue + try: + score = float(row["vdjdb.score"]) + except ValueError: + score = math.nan + if not math.isfinite(score) or score < min_score: + counts["low_or_invalid_score"] += 1 + continue + peptide = normalized_context(row["antigen.epitope"]) + key = (sequence, peptide) + if key not in membership: + counts["not_official_motif_member"] += 1 + continue + counts["exact_official_membership_join_rows"] += 1 + overlapping = [split for split, sequences in benchmark.items() if sequence in sequences] + if overlapping: + counts["benchmark_overlap_rows"] += 1 + for split in overlapping: + overlap[split] += 1 + continue + meta = parse_json_object(row.get("meta", "")) + subject = clean_text(meta.get("subject.id") or meta.get("subject_id")) + study = clean_text(meta.get("study.id") or meta.get("study_id") or row["reference.id"]) + grouped[sequence].append({ + "cdr3": sequence, + "epitope": peptide, + "mhc_a": TARGET_CONTEXT[0], + "mhc_b": TARGET_CONTEXT[1], + "mhc_class": TARGET_CONTEXT[2], + "v_call": normalized_context(row["v.segm"]), + "j_call": normalized_context(row["j.segm"]), + "reference_id": clean_text(row["reference.id"]), + "vdjdb_score": score, + "donor_key": f"{study}|{subject}" if subject else "", + "motif_cids": membership[key], + "length": len(sequence), + }) + + retained = [] + for sequence, rows in sorted(grouped.items()): + epitopes = {row["epitope"] for row in rows} + if len(epitopes) != 1: + counts["ambiguous_multilabel_cdr3"] += 1 + continue + representative = min( + rows, + key=lambda row: ( + -row["vdjdb_score"], row["epitope"], row["v_call"], row["j_call"], + row["reference_id"], + ), + ).copy() + representative["source_rows"] = len(rows) + representative["donor_count_descriptive_only"] = len({ + row["donor_key"] for row in rows if row["donor_key"] + }) + representative["motif_cids"] = ",".join(sorted({ + cid for row in rows for cid in row["motif_cids"] + })) + retained.append(representative) + counts["pooled_exact_cdr3_after_dedup"] = len(retained) + return retained, { + "filter_counts": dict(sorted(counts.items())), + "benchmark_overlap_rows_by_split": dict(sorted(overlap.items())), + } + + +def build_pairs(positives: list[dict], negatives: list[dict], seed: int) -> tuple[list[dict], dict]: + rng = np.random.default_rng(seed) + positive_pairs = [(left, right) for left in range(len(positives)) + for right in range(left + 1, len(positives))] + rng.shuffle(positive_pairs) + negative_by_length: dict[int, list[int]] = defaultdict(list) + negative_composition = np.asarray([composition(record["cdr3"]) for record in negatives]) + for index, record in enumerate(negatives): + negative_by_length[int(record["length"])].append(index) + usage: Counter[int] = Counter() + rows = [] + excluded_no_control_pairs = 0 + forced_orientation_pairs = 0 + for pair_id, (query, positive) in enumerate(positive_pairs): + query_supported = bool(negative_by_length.get(int(positives[query]["length"]), [])) + positive_supported = bool(negative_by_length.get(int(positives[positive]["length"]), [])) + if not query_supported and not positive_supported: + excluded_no_control_pairs += 1 + continue + if query_supported and not positive_supported: + query, positive = positive, query + forced_orientation_pairs += 1 + elif not query_supported and positive_supported: + forced_orientation_pairs += 1 + elif rng.integers(2): + query, positive = positive, query + qrow, prow = positives[query], positives[positive] + candidates = negative_by_length.get(int(prow["length"]), []) + if not candidates: + raise ValueError(f"No exact-length official-motif negative for length {prow['length']}.") + target = composition(prow["cdr3"]) + + def rank(index: int) -> tuple: + candidate = negatives[index] + v_mismatch = int(not prow["v_call"] or not candidate["v_call"] + or prow["v_call"] != candidate["v_call"]) + j_mismatch = int(not prow["j_call"] or not candidate["j_call"] + or prow["j_call"] != candidate["j_call"]) + comp_l1 = float(np.abs(negative_composition[index] - target).sum()) + return comp_l1, v_mismatch + j_mismatch, usage[index], candidate["cdr3"] + + negative = min(candidates, key=rank) + usage[negative] += 1 + nrow = negatives[negative] + rows.append({ + "pair_id": pair_id, + "query_index": query, + "positive_index": positive, + "negative_index": len(positives) + negative, + "query_cdr3": qrow["cdr3"], + "positive_cdr3": prow["cdr3"], + "negative_cdr3": nrow["cdr3"], + "query_length": qrow["length"], + "positive_length": prow["length"], + "negative_length": nrow["length"], + "v_exact_match": int(bool(prow["v_call"]) + and prow["v_call"] == nrow["v_call"]), + "j_exact_match": int(bool(prow["j_call"]) + and prow["j_call"] == nrow["j_call"]), + "composition_l1": float(np.abs(negative_composition[negative] - target).sum()), + }) + if not rows: + raise ValueError("No distinct pooled positive CDR3 pairs are available.") + return rows, { + "positive_pairs": len(rows), + "all_possible_positive_pairs": len(positive_pairs), + "excluded_pairs_both_endpoints_lack_exact_length_controls": excluded_no_control_pairs, + "forced_orientation_to_supported_comparator": forced_orientation_pairs, + "donor_constraint": "none_pooled_by_user_protocol", + "distinct_cdr3_triplets": True, + "exact_length_fraction": 1.0, + "v_exact_match_fraction": float(np.mean([row["v_exact_match"] for row in rows])), + "j_exact_match_fraction": float(np.mean([row["j_exact_match"] for row in rows])), + "mean_composition_l1": float(np.mean([row["composition_l1"] for row in rows])), + "unique_negatives_used": len(usage), + "max_negative_reuse": max(usage.values()), + } + + +def select_compact_negatives( + positives: list[dict], negatives: list[dict], per_positive: int, seed: int, +) -> tuple[list[dict], list[dict], list[dict], dict]: + """Select globally unique same-length controls in seeded per-stratum rounds.""" + if per_positive < 1: + raise ValueError("Controls per positive must be positive.") + negative_by_length: dict[int, list[int]] = defaultdict(list) + negative_composition = np.asarray([composition(record["cdr3"]) for record in negatives]) + for index, record in enumerate(negatives): + negative_by_length[int(record["length"])].append(index) + positive_by_length: dict[int, list[int]] = defaultdict(list) + for index, record in enumerate(positives): + positive_by_length[int(record["length"])].append(index) + + selected: list[int] = [] + used: set[int] = set() + assignments = [] + strata = [] + for length in sorted(positive_by_length): + positive_indices = sorted(positive_by_length[length], key=lambda index: positives[index]["cdr3"]) + rng = np.random.default_rng(seed + length) + positive_indices = [positive_indices[index] for index in rng.permutation(len(positive_indices))] + available_in_stratum = negative_by_length.get(length, []) + for slot in range(1, per_positive + 1): + for positive_index in positive_indices: + candidates = [index for index in available_in_stratum if index not in used] + if not candidates: + break + positive = positives[positive_index] + target = composition(positive["cdr3"]) + + def rank(index: int) -> tuple: + candidate = negatives[index] + comp_l1 = float(np.abs(negative_composition[index] - target).sum()) + v_mismatch = int(not positive["v_call"] or not candidate["v_call"] + or positive["v_call"] != candidate["v_call"]) + j_mismatch = int(not positive["j_call"] or not candidate["j_call"] + or positive["j_call"] != candidate["j_call"]) + return comp_l1, v_mismatch + j_mismatch, candidate["cdr3"] + + chosen = min(candidates, key=rank) + used.add(chosen) + selected.append(chosen) + control = negatives[chosen] + assignments.append({ + "selection_order": len(assignments), + "length": length, + "slot_for_positive": slot, + "positive_cdr3": positive["cdr3"], + "control_cdr3": control["cdr3"], + "composition_l1": float(np.abs(negative_composition[chosen] - target).sum()), + "v_exact_match": int(bool(positive["v_call"]) + and positive["v_call"] == control["v_call"]), + "j_exact_match": int(bool(positive["j_call"]) + and positive["j_call"] == control["j_call"]), + }) + selected_count = sum(1 for row in assignments if row["length"] == length) + target_count = per_positive * len(positive_indices) + strata.append({ + "length": length, + "positive_cdr3": len(positive_indices), + "eligible_negative_cdr3": len(available_in_stratum), + "nominal_target": target_count, + "selected_unique_controls": selected_count, + "shortfall": target_count - selected_count, + "all_available_retained_if_short": int(len(available_in_stratum) < target_count), + }) + selected_records = [negatives[index] for index in selected] + if len({record["cdr3"] for record in selected_records}) != len(selected_records): + raise ValueError("Compact control selector reused a CDR3.") + report = { + "selection_seed": seed, + "controls_per_positive_cap": per_positive, + "nominal_global_cap": per_positive * len(positives), + "eligible_controls_before_compaction": len(negatives), + "selected_unique_controls": len(selected_records), + "selection_without_replacement": True, + "gallery_control_reuse": 0, + "selection_priority": "minimum AA-composition L1, then V/J mismatch count, then CDR3", + "length_strata": strata, + } + return selected_records, assignments, strata, report + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--vdjdb-tsv", type=Path, required=True) + parser.add_argument("--cluster-members", type=Path, required=True) + parser.add_argument("--motif-pwms", type=Path, required=True) + parser.add_argument("--benchmark-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--min-score", type=float, default=1.0) + parser.add_argument("--min-positives", type=int, default=50) + parser.add_argument("--min-negatives", type=int, default=100) + parser.add_argument("--controls-per-positive", type=int, default=4) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + args.output_dir.mkdir(parents=True) + + benchmark, benchmark_hashes = read_benchmark_sequences(args.benchmark_dir) + membership, motif_report = read_official_members(args.cluster_members) + records, join_report = read_joined_records( + args.vdjdb_tsv, membership, benchmark, args.min_score + ) + positives = [record for record in records if record["epitope"] == TARGET_PEPTIDE] + eligible_negatives = [record for record in records if record["epitope"] != TARGET_PEPTIDE] + if len(positives) < args.min_positives or len(eligible_negatives) < args.min_negatives: + blocked = { + "status": "blocked_insufficient_official_motif_support", + "positive_cdr3": len(positives), "negative_cdr3": len(eligible_negatives), + "motif": motif_report, "join": join_report, + } + stable_json(args.output_dir / "BLOCKED.json", blocked) + raise ValueError( + f"Insufficient support after join: {len(positives)} positives, " + f"{len(eligible_negatives)} negatives" + ) + + negatives, control_assignments, length_strata, control_report = select_compact_negatives( + positives, eligible_negatives, args.controls_per_positive, args.seed + ) + + candidates = [] + fields = [ + "candidate_index", "is_ylq", "cdr3", "epitope", "mhc_a", "mhc_b", "mhc_class", + "v_call", "j_call", "reference_id", "vdjdb_score", "length", "motif_cids", + "source_rows", "donor_count_descriptive_only", + ] + for index, record in enumerate([*positives, *negatives]): + candidate = dict(record) + candidate["candidate_index"] = index + candidate["is_ylq"] = int(index < len(positives)) + candidates.append(candidate) + if len({row["cdr3"] for row in candidates}) != len(candidates): + raise ValueError("Candidate set is not exact-CDR3 deduplicated.") + write_tsv(args.output_dir / "candidates.tsv", candidates, fields) + write_tsv( + args.output_dir / "control_selection.tsv", control_assignments, + list(control_assignments[0]), + ) + write_tsv(args.output_dir / "control_length_strata.tsv", length_strata, list(length_strata[0])) + pairs, matching = build_pairs(positives, negatives, args.seed) + write_tsv(args.output_dir / "pair_manifest.tsv", pairs, list(pairs[0])) + manifest = { + "status": "ready", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "target_peptide": TARGET_PEPTIDE, + "exact_context": {"mhc_a": TARGET_CONTEXT[0], "mhc_b": TARGET_CONTEXT[1], + "mhc_class": TARGET_CONTEXT[2]}, + "pooling": "all records pooled; donor metadata is descriptive only and never constrains pairs/retrieval", + "positive_cdr3": len(positives), + "negative_cdr3": len(negatives), + "eligible_negative_cdr3_before_compaction": len(eligible_negatives), + "candidate_cdr3": len(candidates), + "matching": matching, + "compact_control_selection": control_report, + "motif_inventory": motif_report, + "vdjdb_join_and_quality": join_report, + "sources": { + "vdjdb": {"path": str(args.vdjdb_tsv), "sha256": sha256_file(args.vdjdb_tsv)}, + "cluster_members": {"path": str(args.cluster_members), + "sha256": sha256_file(args.cluster_members)}, + "motif_pwms": {"path": str(args.motif_pwms), "sha256": sha256_file(args.motif_pwms)}, + "benchmark_manifests": benchmark_hashes, + }, + "artifacts": { + "candidates_sha256": sha256_file(args.output_dir / "candidates.tsv"), + "pairs_sha256": sha256_file(args.output_dir / "pair_manifest.tsv"), + "control_selection_sha256": sha256_file(args.output_dir / "control_selection.tsv"), + "control_length_strata_sha256": sha256_file(args.output_dir / "control_length_strata.tsv"), + }, + "official_motif_rule": ( + "Exact sequence+peptide membership in official cluster_members.txt after exact " + "HomoSapiens/TRB/HLA-A*02:01/B2M/MHCI filtering; no PWM scanning or motif mining." + ), + "circularity_caveat": ( + "Official motif membership is sequence-cluster-conditioned. This evaluates within-VDJdb " + "motif/epitope organization, not independent epitope evidence beyond sequence motif/composition." + ), + "exploratory_n40_protocol": True, + "embedding_extraction_performed": False, + } + stable_json(args.output_dir / "PREFLIGHT.json", manifest) + print(json.dumps(manifest, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/slurm/forward_encoder.sbatch b/benchmark/slurm/forward_encoder.sbatch new file mode 100644 index 0000000..570c931 --- /dev/null +++ b/benchmark/slurm/forward_encoder.sbatch @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +# Pretrain the sequence-to-TCRemP encoder that the irrm_pretrained Pgen arm starts from. +# +# Trained on the benchmark's own train split, so nothing the Pgen arms are evaluated on +# was seen during pretraining. Must finish before the Pgen arrays start. + +#SBATCH --job-name=irrm-fwd +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=48G +#SBATCH --time=02:00:00 +#SBATCH --output=/mnt/tank/scratch/mpodsytnik/irrm-codec/logs/%x_%j.out +#SBATCH --error=/mnt/tank/scratch/mpodsytnik/irrm-codec/logs/%x_%j.err + +set -euo pipefail + +SCRATCH="${SCRATCH:-/mnt/tank/scratch/mpodsytnik/irrm-codec}" +REPO="${REPO:-/nfs/home/mpodsytnik/irrm-codec}" +CONDA_ROOT="${CONDA_ROOT:-/mnt/tank/scratch/mpodsytnik/miniconda3}" +CONDA_ENV="${CONDA_ENV:-irrm-codec}" + +DATASET_DIR="${DATASET_DIR:-$SCRATCH/data/benchmark/trb}" +OUTPUT_DIR="${OUTPUT_DIR:-$SCRATCH/artifacts/benchmark/forward_encoder}" +EPOCHS="${EPOCHS:-40}" +PATIENCE="${PATIENCE:-12}" + +source "$CONDA_ROOT/bin/activate" +conda activate "$CONDA_ENV" + +echo "job=${SLURM_JOB_ID:-local} node=${SLURMD_NODENAME:-?}" +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || true + +cd "$REPO" +export PYTHONUNBUFFERED=1 +python -m benchmark.train_forward_encoder \ + --dataset-dir "$DATASET_DIR" \ + --output-dir "$OUTPUT_DIR" \ + --train-subset all \ + --epochs "$EPOCHS" \ + --patience "$PATIENCE" \ + --threads "${SLURM_CPUS_PER_TASK:-4}" diff --git a/benchmark/slurm/pgen_array.sbatch b/benchmark/slurm/pgen_array.sbatch new file mode 100644 index 0000000..4f97f9d --- /dev/null +++ b/benchmark/slurm/pgen_array.sbatch @@ -0,0 +1,91 @@ +#!/usr/bin/env bash +# Issue 3: every Pgen arm across both targets, three training sizes and three seeds. +# +# 5 neural arms x 2 targets x 3 subsets x 3 seeds = 90 runs, plus tfidf_ridge which is +# deterministic and needs one run per target and subset = 6. Total 96 tasks, throttled to +# 2 concurrent because the cluster allows at most 2 GPUs per user. +# +# The log directory must exist before submitting: Slurm opens the --output file before the +# script body runs, so a missing directory kills the job immediately. +# +# mkdir -p /mnt/tank/scratch/mpodsytnik/irrm-codec/logs +# sbatch benchmark/slurm/pgen_array.sbatch +# +# tfidf_ridge does not need a GPU, but it finishes in seconds, so keeping one array is +# simpler than splitting it onto the CPU partition. + +#SBATCH --job-name=irrm-pgen +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=24G +#SBATCH --time=02:00:00 +#SBATCH --array=0-95%2 +#SBATCH --output=/mnt/tank/scratch/mpodsytnik/irrm-codec/logs/%x_%A_%a.out +#SBATCH --error=/mnt/tank/scratch/mpodsytnik/irrm-codec/logs/%x_%A_%a.err + +set -euo pipefail + +SCRATCH="${SCRATCH:-/mnt/tank/scratch/mpodsytnik/irrm-codec}" +REPO="${REPO:-/nfs/home/mpodsytnik/irrm-codec}" +CONDA_ROOT="${CONDA_ROOT:-/mnt/tank/scratch/mpodsytnik/miniconda3}" +CONDA_ENV="${CONDA_ENV:-irrm-codec}" + +DATASET_DIR="${DATASET_DIR:-$SCRATCH/data/benchmark/trb}" +REPRESENTATIONS_DIR="${REPRESENTATIONS_DIR:-$SCRATCH/data/benchmark/trb/representations}" +PRETRAINED_CHECKPOINT="${PRETRAINED_CHECKPOINT:-$SCRATCH/artifacts/benchmark/forward_encoder/best.pt}" +OUTPUT_DIR="${OUTPUT_DIR:-$SCRATCH/artifacts/benchmark/pgen}" + +EPOCHS="${EPOCHS:-60}" +BATCH_SIZE="${BATCH_SIZE:-256}" +# Left to the trainer's own default so the cosine schedule is not cut short. +PATIENCE="${PATIENCE:-12}" + +# Build the task list in a fixed order so a task id always maps to the same run. +CONFIGS=() +for target in log10_pgen log10_pgen_1mm; do + for subset in 1k 10k all; do + for arm in irrm_scratch irrm_pretrained sceptr_mlp tcr_bert_mlp esm2_8m_mlp; do + for seed in 42 43 44; do + CONFIGS+=("$arm $target $subset $seed") + done + done + CONFIGS+=("tfidf_ridge $target $subset 42") + done +done + +if [ "${SLURM_ARRAY_TASK_ID:-0}" -ge "${#CONFIGS[@]}" ]; then + echo "task ${SLURM_ARRAY_TASK_ID} is past the end of the ${#CONFIGS[@]}-run list" >&2 + exit 1 +fi +read -r ARM TARGET SUBSET SEED <<< "${CONFIGS[${SLURM_ARRAY_TASK_ID:-0}]}" + +mkdir -p "$OUTPUT_DIR" + +# conda init is deliberately not used: it would edit the shared account's ~/.bashrc. +source "$CONDA_ROOT/bin/activate" +conda activate "$CONDA_ENV" + +echo "job=${SLURM_JOB_ID:-local} task=${SLURM_ARRAY_TASK_ID:-0} node=${SLURMD_NODENAME:-?}" +echo "arm=$ARM target=$TARGET subset=$SUBSET seed=$SEED (${#CONFIGS[@]} runs total)" +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || true + +cd "$REPO" +# Slurm redirects stdout to a file, where Python block-buffers by default. +export PYTHONUNBUFFERED=1 +# Called directly rather than through srun: the batch task already holds its own +# allocation, and an inner step has to request the GPU again, which fails outright when +# the per-user GPU quota is busy elsewhere. +python -m benchmark.train_pgen_arm \ + --arm "$ARM" \ + --target "$TARGET" \ + --train-subset "$SUBSET" \ + --seed "$SEED" \ + --dataset-dir "$DATASET_DIR" \ + --representations-dir "$REPRESENTATIONS_DIR" \ + --pretrained-checkpoint "$PRETRAINED_CHECKPOINT" \ + --output-dir "$OUTPUT_DIR" \ + --epochs "$EPOCHS" \ + --batch-size "$BATCH_SIZE" \ + --patience "$PATIENCE" \ + --threads "${SLURM_CPUS_PER_TASK:-4}" diff --git a/benchmark/slurm/pgen_seeds_array.sbatch b/benchmark/slurm/pgen_seeds_array.sbatch new file mode 100644 index 0000000..33d7972 --- /dev/null +++ b/benchmark/slurm/pgen_seeds_array.sbatch @@ -0,0 +1,95 @@ +#!/usr/bin/env bash +# Issue 3, follow-up: extra seeds for the two IRRM arms on the full training split. +# +# The main array ran three seeds per arm. For the frozen-embedding arms that is plenty +# (their RMSE spread across seeds is about 0.005), but the IRRM arms train a full encoder +# and spread by 0.03-0.06 — larger than the scratch-vs-pretrained difference they are +# meant to measure. Seven more seeds take the standard error from 0.033 to 0.018 and make +# the pretraining question answerable. +# +# All three training sizes, so the low-data question gets the same statistical treatment +# as the full-data one: 2 arms x 2 targets x 3 subsets x 7 seeds. Results land in the same +# output directory, so collect_pgen picks them up with no changes. +# +# mkdir -p /mnt/tank/scratch/mpodsytnik/irrm-codec/logs +# sbatch benchmark/slurm/pgen_seeds_array.sbatch +# +# To repeat only some of them, override SUBSETS: +# sbatch --export=ALL,SUBSETS="1k 10k" benchmark/slurm/pgen_seeds_array.sbatch + +#SBATCH --job-name=irrm-pgen-seeds +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=24G +#SBATCH --time=01:00:00 +#SBATCH --array=0-55%2 +#SBATCH --output=/mnt/tank/scratch/mpodsytnik/irrm-codec/logs/%x_%A_%a.out +#SBATCH --error=/mnt/tank/scratch/mpodsytnik/irrm-codec/logs/%x_%A_%a.err + +set -euo pipefail + +SCRATCH="${SCRATCH:-/mnt/tank/scratch/mpodsytnik/irrm-codec}" +REPO="${REPO:-/nfs/home/mpodsytnik/irrm-codec}" +CONDA_ROOT="${CONDA_ROOT:-/mnt/tank/scratch/mpodsytnik/miniconda3}" +CONDA_ENV="${CONDA_ENV:-irrm-codec}" + +DATASET_DIR="${DATASET_DIR:-$SCRATCH/data/benchmark/trb}" +REPRESENTATIONS_DIR="${REPRESENTATIONS_DIR:-$SCRATCH/data/benchmark/trb/representations}" +PRETRAINED_CHECKPOINT="${PRETRAINED_CHECKPOINT:-$SCRATCH/artifacts/benchmark/forward_encoder/best.pt}" +OUTPUT_DIR="${OUTPUT_DIR:-$SCRATCH/artifacts/benchmark/pgen}" + +# "all" is excluded by default because it was covered by an earlier submission of this +# same file; add it back through SUBSETS and widen --array to match if it is ever needed. +SUBSETS="${SUBSETS:-1k 10k}" +EPOCHS="${EPOCHS:-60}" +BATCH_SIZE="${BATCH_SIZE:-256}" +# Left to the trainer's own default so the cosine schedule is not cut short. +PATIENCE="${PATIENCE:-12}" + +# Seeds 42-44 already ran in the main array; these continue the sequence. +CONFIGS=() +for target in log10_pgen log10_pgen_1mm; do + for subset in $SUBSETS; do + for arm in irrm_scratch irrm_pretrained; do + for seed in 45 46 47 48 49 50 51; do + CONFIGS+=("$arm $target $subset $seed") + done + done + done +done + +# Exit cleanly rather than failing, so an array wider than the config list leaves no +# spurious FAILED tasks in the accounting record. +if [ "${SLURM_ARRAY_TASK_ID:-0}" -ge "${#CONFIGS[@]}" ]; then + echo "task ${SLURM_ARRAY_TASK_ID} is past the end of the ${#CONFIGS[@]}-run list; nothing to do" + exit 0 +fi +read -r ARM TARGET TRAIN_SUBSET SEED <<< "${CONFIGS[${SLURM_ARRAY_TASK_ID:-0}]}" + +mkdir -p "$OUTPUT_DIR" + +source "$CONDA_ROOT/bin/activate" +conda activate "$CONDA_ENV" + +echo "job=${SLURM_JOB_ID:-local} task=${SLURM_ARRAY_TASK_ID:-0} node=${SLURMD_NODENAME:-?}" +echo "arm=$ARM target=$TARGET subset=$TRAIN_SUBSET seed=$SEED (${#CONFIGS[@]} runs total)" + +cd "$REPO" +export PYTHONUNBUFFERED=1 +# Called directly rather than through srun: the batch task already holds its own +# allocation, and an inner step has to request the GPU again, which fails outright when +# the per-user GPU quota is busy elsewhere. +python -m benchmark.train_pgen_arm \ + --arm "$ARM" \ + --target "$TARGET" \ + --train-subset "$TRAIN_SUBSET" \ + --seed "$SEED" \ + --dataset-dir "$DATASET_DIR" \ + --representations-dir "$REPRESENTATIONS_DIR" \ + --pretrained-checkpoint "$PRETRAINED_CHECKPOINT" \ + --output-dir "$OUTPUT_DIR" \ + --epochs "$EPOCHS" \ + --batch-size "$BATCH_SIZE" \ + --patience "$PATIENCE" \ + --threads "${SLURM_CPUS_PER_TASK:-4}" diff --git a/benchmark/slurm/reconstruction_array.sbatch b/benchmark/slurm/reconstruction_array.sbatch new file mode 100644 index 0000000..80ccb54 --- /dev/null +++ b/benchmark/slurm/reconstruction_array.sbatch @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +# Issue 2: train the IRRM-CODEC decoder for every representation and seed. +# +# 5 representations x 3 seeds = 15 tasks. The array is throttled to 2 concurrent +# tasks because the cluster allows at most 2 GPUs per user by default. +# +# The log directory must exist before submitting: Slurm opens the --output file before +# the script body runs, so a missing directory kills the job immediately. +# +# Submit from the repository root: +# mkdir -p /mnt/tank/scratch/mpodsytnik/irrm-codec/logs +# sbatch benchmark/slurm/reconstruction_array.sbatch +# +# Override anything without editing the file: +# sbatch --export=ALL,BOTTLENECK_DIM=32,TRAIN_SUBSET=10k \ +# benchmark/slurm/reconstruction_array.sbatch + +#SBATCH --job-name=irrm-recon +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=16G +# A 60-epoch run takes ~40 min on a GTX 1080 Ti (39 s/epoch at train=79730), so this +# leaves ample headroom while keeping the job small enough to schedule quickly. +#SBATCH --time=02:00:00 +#SBATCH --array=0-14%2 +#SBATCH --output=/mnt/tank/scratch/mpodsytnik/irrm-codec/logs/%x_%A_%a.out +#SBATCH --error=/mnt/tank/scratch/mpodsytnik/irrm-codec/logs/%x_%A_%a.err + +set -euo pipefail + +SCRATCH="${SCRATCH:-/mnt/tank/scratch/mpodsytnik/irrm-codec}" +REPO="${REPO:-/nfs/home/mpodsytnik/irrm-codec}" +CONDA_ROOT="${CONDA_ROOT:-/mnt/tank/scratch/mpodsytnik/miniconda3}" +CONDA_ENV="${CONDA_ENV:-irrm-codec}" + +DATASET_DIR="${DATASET_DIR:-$SCRATCH/data/benchmark/trb}" +BOTTLENECK_DIR="${BOTTLENECK_DIR:-$SCRATCH/data/benchmark/trb/bottleneck}" +OUTPUT_DIR="${OUTPUT_DIR:-$SCRATCH/artifacts/benchmark/reconstruction}" + +BOTTLENECK_DIM="${BOTTLENECK_DIM:-64}" +TRAIN_SUBSET="${TRAIN_SUBSET:-all}" +EPOCHS="${EPOCHS:-60}" +BATCH_SIZE="${BATCH_SIZE:-256}" +PATIENCE="${PATIENCE:-8}" + +# Task id maps to (representation, seed): ids 0-2 are the first representation, and so on. +REPRESENTATIONS=(tcremp onehot tcr_bert sceptr esm2_8m) +SEEDS=(42 43 44) +REPRESENTATION="${REPRESENTATIONS[$((SLURM_ARRAY_TASK_ID / ${#SEEDS[@]}))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % ${#SEEDS[@]}))]}" + +mkdir -p "$OUTPUT_DIR" + +# conda init is deliberately not used: it would edit the shared account's ~/.bashrc. +source "$CONDA_ROOT/bin/activate" +conda activate "$CONDA_ENV" + +echo "job=${SLURM_JOB_ID:-local} task=${SLURM_ARRAY_TASK_ID:-0} node=${SLURMD_NODENAME:-?}" +echo "representation=$REPRESENTATION seed=$SEED dim=$BOTTLENECK_DIM subset=$TRAIN_SUBSET" +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || true + +cd "$REPO" +# Slurm redirects stdout to a file, where Python block-buffers by default and the log +# appears only in bursts. Unbuffered output keeps the .out file usable for monitoring. +export PYTHONUNBUFFERED=1 +# Called directly rather than through srun: the batch task already holds its own +# allocation, and an inner step has to request the GPU again, which fails outright when +# the per-user GPU quota is busy elsewhere. +python -m benchmark.train_decoder \ + --dataset-dir "$DATASET_DIR" \ + --bottleneck-dir "$BOTTLENECK_DIR" \ + --output-dir "$OUTPUT_DIR" \ + --representation "$REPRESENTATION" \ + --bottleneck-dim "$BOTTLENECK_DIM" \ + --train-subset "$TRAIN_SUBSET" \ + --seed "$SEED" \ + --epochs "$EPOCHS" \ + --batch-size "$BATCH_SIZE" \ + --patience "$PATIENCE" \ + --threads "${SLURM_CPUS_PER_TASK:-4}" diff --git a/benchmark/slurm/submit_pgen_all.sh b/benchmark/slurm/submit_pgen_all.sh new file mode 100644 index 0000000..dd8c214 --- /dev/null +++ b/benchmark/slurm/submit_pgen_all.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +# Submit the whole Pgen benchmark as one dependency chain, so it can be started and left. +# +# bash benchmark/slurm/submit_pgen_all.sh +# +# Stages, each waiting for the previous one to succeed: +# 1. pretrain the sequence-to-TCRemP encoder (~10 min) +# 2. main sweep: 6 arms x 2 targets x 3 sizes x 3 seeds (96 runs) +# 3. extra seeds for the two IRRM arms, 7 more each (84 runs) +# +# Existing results are archived rather than left in place: run names are identical +# between submissions, so mixing a finished constant-learning-rate sweep with a new +# cosine-schedule one would silently average two different configurations together. + +set -euo pipefail + +SCRATCH="${SCRATCH:-/mnt/tank/scratch/mpodsytnik/irrm-codec}" +REPO="${REPO:-/nfs/home/mpodsytnik/irrm-codec}" +PGEN_DIR="$SCRATCH/artifacts/benchmark/pgen" +FORWARD_DIR="$SCRATCH/artifacts/benchmark/forward_encoder" + +mkdir -p "$SCRATCH/logs" + +stamp="$(date +%Y%m%d_%H%M%S)" +for dir in "$PGEN_DIR" "$FORWARD_DIR"; do + if [ -d "$dir" ]; then + mv "$dir" "${dir}_constlr_${stamp}" + echo "archived $(basename "$dir") -> $(basename "${dir}")_constlr_${stamp}" + fi +done + +cd "$REPO" + +forward_id=$(sbatch --parsable benchmark/slurm/forward_encoder.sbatch) +echo "1/3 forward encoder : $forward_id" + +main_id=$(sbatch --parsable --dependency=afterok:"$forward_id" benchmark/slurm/pgen_array.sbatch) +echo "2/3 main sweep : $main_id (after $forward_id)" + +# The extra seeds only cover 1k and 10k here; the main sweep already supplies three seeds +# per cell, and SUBSETS controls which sizes get topped up to ten. +seeds_id=$(sbatch --parsable --dependency=afterok:"$main_id" \ + --export=ALL,SUBSETS="1k 10k all" --array=0-83%2 benchmark/slurm/pgen_seeds_array.sbatch) +echo "3/3 extra seeds : $seeds_id (after $main_id)" + +echo +echo "watch with: squeue -u \$USER" +echo "count with: ls $PGEN_DIR/*/metrics.json 2>/dev/null | wc -l # expect 180" diff --git a/benchmark/train_decoder.py b/benchmark/train_decoder.py new file mode 100644 index 0000000..6ea2366 --- /dev/null +++ b/benchmark/train_decoder.py @@ -0,0 +1,259 @@ +"""Issue 2: train the IRRM-CODEC decoder on one representation and score reconstruction. + +The decoder architecture, optimizer, schedule and early stopping are identical for every +representation, so the only thing that varies between runs is the bottleneck vector the +decoder reads. Inputs are standardized with train-split statistics: the projections +arrive on very different scales, and a shared learning rate on unequal scales would make +"identical training settings" true on paper only. The repo's forward/inverse trainers +normalize their embeddings the same way. + +Reported metrics: exact sequence match, amino-acid token accuracy, normalized +Levenshtein distance, fraction of predictions within edit distance 1, inference speed. +""" + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from rapidfuzz.distance import Levenshtein +from torch.utils.data import DataLoader, TensorDataset + +from irrm_codec.inverse_model import InverseModel +from irrm_codec.losses import inverse_loss +from irrm_codec.tokenization import AA_VOCAB, ID2AA, encode + +GAP_ID = AA_VOCAB["-"] +from irrm_codec.utils import choose_device, save_checkpoint, set_seed, setup_logging + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument("--bottleneck-dir", default="data/benchmark/trb/bottleneck") + p.add_argument("--representation", required=True) + p.add_argument("--bottleneck-dim", type=int, default=64) + p.add_argument("--output-dir", default="artifacts/benchmark/reconstruction") + p.add_argument("--seed", type=int, default=42) + p.add_argument("--epochs", type=int, default=60) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--weight-decay", type=float, default=1e-4) + p.add_argument("--patience", type=int, default=8, help="Early-stopping patience in epochs.") + p.add_argument("--max-len", type=int, default=40) + p.add_argument("--train-subset", default="all", help="Training subset manifest: 1k, 10k or all.") + p.add_argument("--threads", type=int, default=8) + p.add_argument("--num-workers", type=int, default=0) + return p.parse_args() + + +def load_split_rows(dataset_dir, name): + return pd.read_csv(Path(dataset_dir) / "manifests" / f"{name}.tsv", sep="\t")["row_index"].to_numpy() + + +def decode_tokens(token_ids): + """Turn predicted token ids back into a gap-free amino-acid string.""" + return "".join(ID2AA[int(t)] for t in token_ids if int(t) > GAP_ID) + + +def build_tensors(matrix, rows, sequences, max_len): + features = torch.from_numpy(np.asarray(matrix[rows], dtype=np.float32)) + targets = torch.tensor( + [encode(sequences[i], max_len=max_len) for i in rows], dtype=torch.long + ) + return features, targets + + +@torch.no_grad() +def evaluate(model, loader, device, sequences, rows, max_len, collect_predictions=False): + """Run the decoder over a split and score it at the sequence level.""" + model.eval() + loss_sum, steps = 0.0, 0 + token_hits, token_total = 0, 0 + predictions = [] + + for features, targets in loader: + features, targets = features.to(device), targets.to(device) + logits = model(features) + loss_sum += inverse_loss(logits, targets).item() + steps += 1 + + predicted = logits.argmax(dim=-1) + # encode() gap-pads to a fixed width, so every position is filled and most of + # them are gaps. Scoring those would let a model inflate token accuracy by + # predicting padding, so accuracy counts real residues only. + residues = targets.gt(GAP_ID) + token_hits += predicted.eq(targets).logical_and(residues).sum().item() + token_total += residues.sum().item() + predictions.extend(decode_tokens(row) for row in predicted.cpu().numpy()) + + truth = [sequences[i] for i in rows] + distances = np.array([Levenshtein.distance(p, t) for p, t in zip(predictions, truth)]) + lengths = np.array([max(len(t), 1) for t in truth]) + + metrics = { + "loss": loss_sum / max(steps, 1), + "token_accuracy": token_hits / max(token_total, 1), + "exact_match": float(np.mean(distances == 0)), + "normalized_levenshtein": float(np.mean(distances / lengths)), + "within_edit_distance_1": float(np.mean(distances <= 1)), + "mean_edit_distance": float(distances.mean()), + } + return (metrics, predictions, truth) if collect_predictions else (metrics, None, None) + + +@torch.no_grad() +def measure_inference_speed(model, features, device, batch_size): + """Throughput on the test split, excluding metric computation.""" + model.eval() + warmup = features[: min(batch_size, len(features))].to(device) + model(warmup) + + started = time.perf_counter() + for start in range(0, len(features), batch_size): + model(features[start : start + batch_size].to(device)) + elapsed = time.perf_counter() - started + return { + "sequences_per_second": len(features) / elapsed, + "ms_per_sequence": elapsed / len(features) * 1000, + "batch_size": batch_size, + "device": str(device), + } + + +def main(): + args = parse_args() + torch.set_num_threads(args.threads) + set_seed(args.seed) + device = choose_device() + + run_name = f"{args.representation}_d{args.bottleneck_dim}_{args.train_subset}_seed{args.seed}" + output_dir = Path(args.output_dir) / run_name + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "train.log") + + dataset_dir = Path(args.dataset_dir) + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + sequences = frame["junction_aa"].tolist() + + source = Path(args.bottleneck_dir) / f"{args.representation}_{args.bottleneck_dim}.npy" + matrix = np.load(source, mmap_mode="r") + log.info("run=%s device=%s features=%s", run_name, device, matrix.shape) + + train_rows = load_split_rows(dataset_dir, f"train_{args.train_subset}") + val_rows = load_split_rows(dataset_dir, "val") + test_rows = load_split_rows(dataset_dir, "test") + + train_x, train_y = build_tensors(matrix, train_rows, sequences, args.max_len) + val_x, val_y = build_tensors(matrix, val_rows, sequences, args.max_len) + test_x, test_y = build_tensors(matrix, test_rows, sequences, args.max_len) + + # Train-split statistics only, so no held-out information reaches the projection. + mean = train_x.mean(dim=0, keepdim=True) + std = train_x.std(dim=0, keepdim=True).clamp_min(1e-6) + train_x, val_x, test_x = ((t - mean) / std for t in (train_x, val_x, test_x)) + np.save(output_dir / "mean.npy", mean.numpy()) + np.save(output_dir / "std.npy", std.numpy()) + + loaders = { + name: DataLoader( + TensorDataset(x, y), + batch_size=args.batch_size, + shuffle=(name == "train"), + num_workers=args.num_workers, + ) + for name, (x, y) in ( + ("train", (train_x, train_y)), + ("val", (val_x, val_y)), + ("test", (test_x, test_y)), + ) + } + + model = InverseModel(embedding_dim=args.bottleneck_dim, max_len=args.max_len).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + log.info("model params=%.2fM train=%d val=%d test=%d", + sum(p.numel() for p in model.parameters()) / 1e6, + len(train_rows), len(val_rows), len(test_rows)) + + best_val = float("inf") + best_epoch = 0 + history = [] + training_started = time.perf_counter() + + for epoch in range(1, args.epochs + 1): + model.train() + epoch_loss, steps = 0.0, 0 + for features, targets in loaders["train"]: + features, targets = features.to(device), targets.to(device) + logits = model(features) + loss = inverse_loss(logits, targets) + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + epoch_loss += loss.item() + steps += 1 + + train_loss = epoch_loss / max(steps, 1) + val_metrics, _, _ = evaluate(model, loaders["val"], device, sequences, val_rows, args.max_len) + history.append({"epoch": epoch, "train_loss": train_loss, **{f"val_{k}": v for k, v in val_metrics.items()}}) + log.info( + "epoch=%d train_loss=%.4f val_loss=%.4f val_exact=%.4f val_tok=%.4f", + epoch, train_loss, val_metrics["loss"], val_metrics["exact_match"], val_metrics["token_accuracy"], + ) + + if val_metrics["loss"] < best_val - 1e-5: + best_val, best_epoch = val_metrics["loss"], epoch + save_checkpoint(output_dir / "best.pt", model, optimizer, epoch, val_metrics) + elif epoch - best_epoch >= args.patience: + log.info("early stopping at epoch=%d (best epoch=%d)", epoch, best_epoch) + break + + training_seconds = time.perf_counter() - training_started + save_checkpoint(output_dir / "last.pt", model, optimizer, len(history), history[-1]) + + model.load_state_dict(torch.load(output_dir / "best.pt", map_location=device)["model_state"]) + test_metrics, predictions, truth = evaluate( + model, loaders["test"], device, sequences, test_rows, args.max_len, collect_predictions=True + ) + speed = measure_inference_speed(model, test_x, device, args.batch_size) + + pd.DataFrame( + { + "row_index": test_rows, + "true_cdr3": truth, + "predicted_cdr3": predictions, + "edit_distance": [Levenshtein.distance(p, t) for p, t in zip(predictions, truth)], + } + ).to_csv(output_dir / "test_predictions.csv", index=False) + + pd.DataFrame(history).to_csv(output_dir / "history.csv", index=False) + result = { + "run_name": run_name, + "representation": args.representation, + "bottleneck_dim": args.bottleneck_dim, + "train_subset": args.train_subset, + "train_size": int(len(train_rows)), + "seed": args.seed, + "epochs_ran": len(history), + "best_epoch": best_epoch, + "best_val_loss": best_val, + "training_seconds": round(training_seconds, 2), + "test": test_metrics, + "inference": speed, + "config": vars(args), + } + (output_dir / "metrics.json").write_text(json.dumps(result, indent=2), encoding="utf-8") + + log.info("=" * 60) + log.info("test exact_match=%.4f token_accuracy=%.4f", test_metrics["exact_match"], test_metrics["token_accuracy"]) + log.info("test norm_levenshtein=%.4f within_ed1=%.4f", test_metrics["normalized_levenshtein"], test_metrics["within_edit_distance_1"]) + log.info("training_seconds=%.1f inference=%.0f seq/s", training_seconds, speed["sequences_per_second"]) + log.info("wrote %s", output_dir / "metrics.json") + + +if __name__ == "__main__": + main() diff --git a/benchmark/train_external_model_decoder.py b/benchmark/train_external_model_decoder.py new file mode 100644 index 0000000..5bed724 --- /dev/null +++ b/benchmark/train_external_model_decoder.py @@ -0,0 +1,271 @@ +"""Issue 2: train the IRRM-CODEC decoder on one representation and score reconstruction. + +The decoder architecture, optimizer, schedule and early stopping are identical for every +representation, so the only thing that varies between runs is the bottleneck vector the +decoder reads. Inputs are standardized with train-split statistics: the projections +arrive on very different scales, and a shared learning rate on unequal scales would make +"identical training settings" true on paper only. The repo's forward/inverse trainers +normalize their embeddings the same way. + +Reported metrics: exact sequence match, amino-acid token accuracy, normalized +Levenshtein distance, fraction of predictions within edit distance 1, inference speed. +""" + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import DataLoader, TensorDataset + +from irrm_codec.inverse_model import InverseModel +from irrm_codec.losses import inverse_loss +from irrm_codec.tokenization import AA_VOCAB, ID2AA, encode + +GAP_ID = AA_VOCAB["-"] +from irrm_codec.utils import choose_device, save_checkpoint, set_seed, setup_logging + + +def levenshtein_distance(left: str, right: str) -> int: + if len(left) > len(right): + left, right = right, left + previous = list(range(len(left) + 1)) + for right_index, right_char in enumerate(right, start=1): + current = [right_index] + for left_index, left_char in enumerate(left, start=1): + current.append(min(current[-1] + 1, previous[left_index] + 1, + previous[left_index - 1] + (left_char != right_char))) + previous = current + return previous[-1] + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument("--bottleneck-dir", default="data/benchmark/trb/bottleneck") + p.add_argument("--representation", required=True) + p.add_argument("--bottleneck-dim", type=int, default=64) + p.add_argument("--output-dir", default="artifacts/benchmark/reconstruction") + p.add_argument("--seed", type=int, default=42) + p.add_argument("--epochs", type=int, default=60) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--weight-decay", type=float, default=1e-4) + p.add_argument("--patience", type=int, default=8, help="Early-stopping patience in epochs.") + p.add_argument("--max-len", type=int, default=40) + p.add_argument("--train-subset", default="all", help="Training subset manifest: 1k, 10k or all.") + p.add_argument("--threads", type=int, default=8) + p.add_argument("--num-workers", type=int, default=0) + return p.parse_args() + + +def load_split_rows(dataset_dir, name): + return pd.read_csv(Path(dataset_dir) / "manifests" / f"{name}.tsv", sep="\t")["row_index"].to_numpy() + + +def decode_tokens(token_ids): + """Turn predicted token ids back into a gap-free amino-acid string.""" + return "".join(ID2AA[int(t)] for t in token_ids if int(t) > GAP_ID) + + +def build_tensors(matrix, rows, sequences, max_len): + features = torch.from_numpy(np.asarray(matrix[rows], dtype=np.float32)) + targets = torch.tensor( + [encode(sequences[i], max_len=max_len) for i in rows], dtype=torch.long + ) + return features, targets + + +@torch.no_grad() +def evaluate(model, loader, device, sequences, rows, max_len, collect_predictions=False): + """Run the decoder over a split and score it at the sequence level.""" + model.eval() + loss_sum, steps = 0.0, 0 + token_hits, token_total = 0, 0 + predictions = [] + + for features, targets in loader: + features, targets = features.to(device), targets.to(device) + logits = model(features) + loss_sum += inverse_loss(logits, targets).item() + steps += 1 + + predicted = logits.argmax(dim=-1) + # encode() gap-pads to a fixed width, so every position is filled and most of + # them are gaps. Scoring those would let a model inflate token accuracy by + # predicting padding, so accuracy counts real residues only. + residues = targets.gt(GAP_ID) + token_hits += predicted.eq(targets).logical_and(residues).sum().item() + token_total += residues.sum().item() + predictions.extend(decode_tokens(row) for row in predicted.cpu().numpy()) + + truth = [sequences[i] for i in rows] + distances = np.array([levenshtein_distance(p, t) for p, t in zip(predictions, truth)]) + lengths = np.array([max(len(t), 1) for t in truth]) + + metrics = { + "loss": loss_sum / max(steps, 1), + "token_accuracy": token_hits / max(token_total, 1), + "exact_match": float(np.mean(distances == 0)), + "normalized_levenshtein": float(np.mean(distances / lengths)), + "within_edit_distance_1": float(np.mean(distances <= 1)), + "mean_edit_distance": float(distances.mean()), + } + return (metrics, predictions, truth) if collect_predictions else (metrics, None, None) + + +@torch.no_grad() +def measure_inference_speed(model, features, device, batch_size): + """Throughput on the test split, excluding metric computation.""" + model.eval() + warmup = features[: min(batch_size, len(features))].to(device) + model(warmup) + + started = time.perf_counter() + for start in range(0, len(features), batch_size): + model(features[start : start + batch_size].to(device)) + elapsed = time.perf_counter() - started + return { + "sequences_per_second": len(features) / elapsed, + "ms_per_sequence": elapsed / len(features) * 1000, + "batch_size": batch_size, + "device": str(device), + } + + +def main(): + args = parse_args() + torch.set_num_threads(args.threads) + set_seed(args.seed) + device = choose_device() + + run_name = f"{args.representation}_d{args.bottleneck_dim}_{args.train_subset}_seed{args.seed}" + output_dir = Path(args.output_dir) / run_name + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "train.log") + + dataset_dir = Path(args.dataset_dir) + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + sequences = frame["junction_aa"].tolist() + + source = Path(args.bottleneck_dir) / f"{args.representation}_{args.bottleneck_dim}.npy" + matrix = np.load(source, mmap_mode="r") + log.info("run=%s device=%s features=%s", run_name, device, matrix.shape) + + train_rows = load_split_rows(dataset_dir, f"train_{args.train_subset}") + val_rows = load_split_rows(dataset_dir, "val") + test_rows = load_split_rows(dataset_dir, "test") + + train_x, train_y = build_tensors(matrix, train_rows, sequences, args.max_len) + val_x, val_y = build_tensors(matrix, val_rows, sequences, args.max_len) + test_x, test_y = build_tensors(matrix, test_rows, sequences, args.max_len) + + # Train-split statistics only, so no held-out information reaches the projection. + mean = train_x.mean(dim=0, keepdim=True) + std = train_x.std(dim=0, keepdim=True).clamp_min(1e-6) + train_x, val_x, test_x = ((t - mean) / std for t in (train_x, val_x, test_x)) + np.save(output_dir / "mean.npy", mean.numpy()) + np.save(output_dir / "std.npy", std.numpy()) + + loaders = { + name: DataLoader( + TensorDataset(x, y), + batch_size=args.batch_size, + shuffle=(name == "train"), + num_workers=args.num_workers, + ) + for name, (x, y) in ( + ("train", (train_x, train_y)), + ("val", (val_x, val_y)), + ("test", (test_x, test_y)), + ) + } + + model = InverseModel(embedding_dim=args.bottleneck_dim, max_len=args.max_len).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + log.info("model params=%.2fM train=%d val=%d test=%d", + sum(p.numel() for p in model.parameters()) / 1e6, + len(train_rows), len(val_rows), len(test_rows)) + + best_val = float("inf") + best_epoch = 0 + history = [] + training_started = time.perf_counter() + + for epoch in range(1, args.epochs + 1): + model.train() + epoch_loss, steps = 0.0, 0 + for features, targets in loaders["train"]: + features, targets = features.to(device), targets.to(device) + logits = model(features) + loss = inverse_loss(logits, targets) + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + epoch_loss += loss.item() + steps += 1 + + train_loss = epoch_loss / max(steps, 1) + val_metrics, _, _ = evaluate(model, loaders["val"], device, sequences, val_rows, args.max_len) + history.append({"epoch": epoch, "train_loss": train_loss, **{f"val_{k}": v for k, v in val_metrics.items()}}) + log.info( + "epoch=%d train_loss=%.4f val_loss=%.4f val_exact=%.4f val_tok=%.4f", + epoch, train_loss, val_metrics["loss"], val_metrics["exact_match"], val_metrics["token_accuracy"], + ) + + if val_metrics["loss"] < best_val - 1e-5: + best_val, best_epoch = val_metrics["loss"], epoch + save_checkpoint(output_dir / "best.pt", model, optimizer, epoch, val_metrics) + elif epoch - best_epoch >= args.patience: + log.info("early stopping at epoch=%d (best epoch=%d)", epoch, best_epoch) + break + + training_seconds = time.perf_counter() - training_started + save_checkpoint(output_dir / "last.pt", model, optimizer, len(history), history[-1]) + + model.load_state_dict(torch.load(output_dir / "best.pt", map_location=device)["model_state"]) + test_metrics, predictions, truth = evaluate( + model, loaders["test"], device, sequences, test_rows, args.max_len, collect_predictions=True + ) + speed = measure_inference_speed(model, test_x, device, args.batch_size) + + pd.DataFrame( + { + "row_index": test_rows, + "true_cdr3": truth, + "predicted_cdr3": predictions, + "edit_distance": [levenshtein_distance(p, t) for p, t in zip(predictions, truth)], + } + ).to_csv(output_dir / "test_predictions.csv", index=False) + + pd.DataFrame(history).to_csv(output_dir / "history.csv", index=False) + result = { + "run_name": run_name, + "representation": args.representation, + "bottleneck_dim": args.bottleneck_dim, + "train_subset": args.train_subset, + "train_size": int(len(train_rows)), + "seed": args.seed, + "epochs_ran": len(history), + "best_epoch": best_epoch, + "best_val_loss": best_val, + "training_seconds": round(training_seconds, 2), + "test": test_metrics, + "inference": speed, + "config": vars(args), + } + (output_dir / "metrics.json").write_text(json.dumps(result, indent=2), encoding="utf-8") + + log.info("=" * 60) + log.info("test exact_match=%.4f token_accuracy=%.4f", test_metrics["exact_match"], test_metrics["token_accuracy"]) + log.info("test norm_levenshtein=%.4f within_ed1=%.4f", test_metrics["normalized_levenshtein"], test_metrics["within_edit_distance_1"]) + log.info("training_seconds=%.1f inference=%.0f seq/s", training_seconds, speed["sequences_per_second"]) + log.info("wrote %s", output_dir / "metrics.json") + + +if __name__ == "__main__": + main() diff --git a/benchmark/train_external_model_pgen.py b/benchmark/train_external_model_pgen.py new file mode 100644 index 0000000..66002ab --- /dev/null +++ b/benchmark/train_external_model_pgen.py @@ -0,0 +1,425 @@ +"""Issue 3: train one Pgen-prediction arm and score it on the benchmark test split. + +Arms: + irrm_scratch PgenModel from a random initialization + irrm_pretrained PgenModel initialized from the sequence-to-TCRemP encoder + tfidf_ridge 1-3-mer character TF-IDF + Ridge regression + sceptr_mlp frozen SCEPTR embeddings + the shared regression head + tcr_bert_mlp frozen TCR-BERT embeddings + the same head + esm2_8m_mlp frozen ESM-2 8M embeddings + the same head + rtp_mlp frozen DATA-ANCHOR R+T+P embeddings + the same head + sceptr_residue_pca128_mlp frozen train-only residue-PCA 128D SCEPTR vector + +The four frozen-embedding arms share one head and one training configuration, so the +only thing that varies between them is the representation. The two IRRM arms differ only +in initialization, which is the comparison the issue asks about. + +Targets stay in raw log10 units, matching the repo's pgen trainer, so RMSE, MAE and bias +are directly interpretable as log10 error. +""" + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +from scipy import stats +from torch.utils.data import DataLoader, TensorDataset + +from irrm_codec.losses import pgen_loss +from irrm_codec.pgen_model import PgenModel +from irrm_codec.tokenization import PAD_ID, encode +from irrm_codec.utils import choose_device, save_checkpoint, set_seed, setup_logging + +SEQUENCE_ARMS = ("irrm_scratch", "irrm_pretrained") +FROZEN_ARMS = { + "sceptr_mlp": "sceptr", + "tcr_bert_mlp": "tcr_bert", + "esm2_8m_mlp": "esm2_8m", + "onehot_mlp": "onehot", + "tcremp_mlp": "tcremp", + "esm2_35m_mlp": "esm2_35m", + "rtp_mlp": "rtp", + "sceptr_cdr3_mlp": "sceptr_cdr3", + "sceptr_residue_pca256_mlp": "sceptr_residue", + "tcr_bert_residue_pca256_mlp": "tcr_bert_residue", + "esm2_8m_residue_pca256_mlp": "esm2_8m_residue", + "sceptr_residue_pca128_mlp": "sceptr_residue", + "tcr_bert_residue_pca128_mlp": "tcr_bert_residue", + "esm2_8m_residue_pca128_mlp": "esm2_8m_residue", +} +ARMS = (*SEQUENCE_ARMS, "tfidf_ridge", *FROZEN_ARMS) + +# Ridge has no training loop, so the penalty is selected on validation instead. This is +# the baseline's equivalent of early stopping; fixing it arbitrarily would handicap it. +RIDGE_ALPHAS = (0.01, 0.1, 1.0, 10.0, 100.0) + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--arm", required=True, choices=ARMS) + p.add_argument("--target", default="log10_pgen_1mm", choices=("log10_pgen", "log10_pgen_1mm")) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument("--representations-dir", default="data/benchmark/trb/representations") + p.add_argument("--pretrained-checkpoint", default="artifacts/benchmark/forward_encoder/best.pt") + p.add_argument("--output-dir", default="artifacts/benchmark/pgen") + p.add_argument("--train-subset", default="all") + p.add_argument("--seed", type=int, default=42) + p.add_argument("--epochs", type=int, default=60) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--weight-decay", type=float, default=1e-4) + # Wider than the reconstruction runs use: the cosine schedule only pays off if training + # reaches the low-learning-rate tail, and stopping at epoch 17 of 60 would decay the + # rate by under a fifth. + p.add_argument("--patience", type=int, default=12) + p.add_argument("--max-len", type=int, default=40) + p.add_argument("--threads", type=int, default=8) + return p.parse_args() + + +class RegressionHead(nn.Module): + """One head shared by every frozen representation, so only the input differs.""" + + def __init__(self, input_dim, hidden_dims=(512, 256), dropout=0.2): + super().__init__() + layers = [] + previous = input_dim + for width in hidden_dims: + layers += [nn.Linear(previous, width), nn.GELU(), nn.Dropout(dropout)] + previous = width + layers.append(nn.Linear(previous, 1)) + self.net = nn.Sequential(*layers) + + def forward(self, features): + return self.net(features).squeeze(-1) + + +def load_rows(dataset_dir, name): + return pd.read_csv(Path(dataset_dir) / "manifests" / f"{name}.tsv", sep="\t")["row_index"].to_numpy() + + +def peak_memory_mb(): + """True peak resident set size, plus GPU peak when a device was used.""" + result = {} + try: + import resource + + # ru_maxrss is in kilobytes on Linux. + result["peak_rss_mb"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 + except (ImportError, AttributeError): + try: + import psutil + + result["peak_rss_mb"] = psutil.Process().memory_info().rss / 1024**2 + except ImportError: + result["peak_rss_mb"] = None + if torch.cuda.is_available(): + result["peak_gpu_mb"] = torch.cuda.max_memory_allocated() / 1024**2 + return result + + +def regression_metrics(predictions, truth): + residual = predictions - truth + total_variance = np.var(truth) + return { + "rmse": float(np.sqrt(np.mean(residual**2))), + "mae": float(np.mean(np.abs(residual))), + "bias": float(np.mean(residual)), + "r2": float(1.0 - np.mean(residual**2) / total_variance) if total_variance > 0 else float("nan"), + "pearson_r": float(stats.pearsonr(predictions, truth)[0]), + "spearman_rho": float(stats.spearmanr(predictions, truth)[0]), + } + + +def build_features(args, frame, rows_by_split, log): + """Return per-split model inputs for the requested arm.""" + if args.arm in SEQUENCE_ARMS: + tokens = { + name: torch.tensor( + [encode(frame.junction_aa.iloc[i], max_len=args.max_len) for i in rows], + dtype=torch.long, + ) + for name, rows in rows_by_split.items() + } + return tokens, None + + representation = FROZEN_ARMS[args.arm] + representations_dir = Path(args.representations_dir) + catalog_path = representations_dir / "representations.json" + catalog = json.loads(catalog_path.read_text(encoding="utf-8")) + source = Path(catalog[representation]["path"]) + matrix = np.load(source, mmap_mode="r") + if matrix.shape[0] != len(frame): + raise ValueError(f"{representation} rows={matrix.shape[0]} dataset rows={len(frame)}") + features = { + name: torch.from_numpy(np.asarray(matrix[rows], dtype=np.float32)) + for name, rows in rows_by_split.items() + } + # Train-split statistics only: a shared learning rate across representations with + # different scales would not be the identical setting the issue requires. + mean = features["train"].mean(dim=0, keepdim=True) + raw_std = features["train"].std(dim=0, keepdim=True) + zero_variance = raw_std < 1e-6 + std = torch.where(zero_variance, torch.ones_like(raw_std), raw_std) + for name in features: + features[name] = (features[name] - mean) / std + log.info("%s features dim=%d zero_variance_columns=%d", representation, + features["train"].shape[1], int(zero_variance.sum())) + return features, (mean, std) + + +def build_model(args, input_dim, log): + if args.arm in SEQUENCE_ARMS: + model = PgenModel(max_len=args.max_len) + if args.arm == "irrm_pretrained": + state = torch.load(args.pretrained_checkpoint, map_location="cpu")["model_state"] + own = model.state_dict() + # Everything but the output layer matches: PgenModel is ForwardModel with + # output_dim=1, so only mlp. differs in shape. + transferable = {k: v for k, v in state.items() if k in own and own[k].shape == v.shape} + model.load_state_dict(transferable, strict=False) + log.info( + "transferred %d/%d tensors from %s", + len(transferable), len(own), args.pretrained_checkpoint, + ) + return model + return RegressionHead(input_dim) + + +def run_epoch(model, loader, device, is_sequence, optimizer=None): + is_train = optimizer is not None + model.train(mode=is_train) + loss_sum, steps = 0.0, 0 + for batch in loader: + inputs, target = batch[0].to(device), batch[-1].to(device) + with torch.set_grad_enabled(is_train): + prediction = model(inputs, inputs.ne(PAD_ID)) if is_sequence else model(inputs) + loss = pgen_loss(prediction, target) + if is_train: + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + loss_sum += loss.item() + steps += 1 + return loss_sum / max(steps, 1) + + +@torch.no_grad() +def predict_all(model, inputs, device, is_sequence, batch_size): + model.eval() + chunks = [] + for start in range(0, len(inputs), batch_size): + block = inputs[start : start + batch_size].to(device) + prediction = model(block, block.ne(PAD_ID)) if is_sequence else model(block) + chunks.append(prediction.cpu().numpy()) + return np.concatenate(chunks) + + +def fit_ridge(args, frame, rows_by_split, targets, log): + """1-3-mer character TF-IDF into Ridge, with the penalty chosen on validation.""" + from sklearn.feature_extraction.text import TfidfVectorizer + from sklearn.linear_model import Ridge + + started = time.perf_counter() + sequences = { + name: frame.junction_aa.iloc[rows].tolist() for name, rows in rows_by_split.items() + } + vectorizer = TfidfVectorizer(analyzer="char", ngram_range=(1, 3)) + train_features = vectorizer.fit_transform(sequences["train"]) + val_features = vectorizer.transform(sequences["val"]) + + best = (None, float("inf"), None) + for alpha in RIDGE_ALPHAS: + model = Ridge(alpha=alpha).fit(train_features, targets["train"]) + rmse = float(np.sqrt(np.mean((model.predict(val_features) - targets["val"]) ** 2))) + log.info("ridge alpha=%.3g val_rmse=%.4f", alpha, rmse) + if rmse < best[1]: + best = (model, rmse, alpha) + model, val_rmse, alpha = best + training_seconds = time.perf_counter() - started + log.info("selected alpha=%.3g val_rmse=%.4f features=%d", alpha, val_rmse, train_features.shape[1]) + + # Vectorization is part of this arm's inference path, so it belongs inside the timed + # section; excluding it would report the Ridge solve alone and overstate throughput + # by orders of magnitude. + inference_started = time.perf_counter() + test_features = vectorizer.transform(sequences["test"]) + predictions = model.predict(test_features) + elapsed = time.perf_counter() - inference_started + + return ( + predictions, + training_seconds, + { + "sequences_per_second": len(predictions) / elapsed, + "ms_per_sequence": elapsed / len(predictions) * 1000, + "batch_size": len(predictions), + "device": "cpu", + }, + {"alpha": alpha, "val_rmse": val_rmse, "n_features": int(train_features.shape[1])}, + [], + ) + + +def fit_neural(args, features, targets, log): + device = choose_device() + is_sequence = args.arm in SEQUENCE_ARMS + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + loaders = { + name: DataLoader( + TensorDataset(features[name], torch.from_numpy(targets[name].astype(np.float32))), + batch_size=args.batch_size, + shuffle=(name == "train"), + ) + for name in ("train", "val") + } + + model = build_model(args, features["train"].shape[-1], log).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + # A constant learning rate leaves the model bouncing around the minimum: validation + # loss oscillates, early stopping fires at an arbitrary epoch, and the IRRM arms end up + # spreading roughly tenfold wider across seeds than the frozen-embedding arms. + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) + log.info("device=%s params=%.2fM", device, sum(p.numel() for p in model.parameters()) / 1e6) + + best_val, best_epoch = float("inf"), 0 + best_state = None + history = [] + started = time.perf_counter() + + for epoch in range(1, args.epochs + 1): + train_loss = run_epoch(model, loaders["train"], device, is_sequence, optimizer) + val_loss = run_epoch(model, loaders["val"], device, is_sequence) + scheduler.step() + history.append({"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss}) + if epoch % 5 == 0 or epoch == 1: + log.info("epoch=%d train_loss=%.4f val_loss=%.4f", epoch, train_loss, val_loss) + + if val_loss < best_val - 1e-6: + best_val, best_epoch = val_loss, epoch + best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} + elif epoch - best_epoch >= args.patience: + log.info("early stopping at epoch=%d (best epoch=%d)", epoch, best_epoch) + break + + training_seconds = time.perf_counter() - started + if best_state is not None: + model.load_state_dict(best_state) + + inference_started = time.perf_counter() + predictions = predict_all(model, features["test"], device, is_sequence, args.batch_size) + elapsed = time.perf_counter() - inference_started + + return ( + predictions, + training_seconds, + { + "sequences_per_second": len(predictions) / elapsed, + "ms_per_sequence": elapsed / len(predictions) * 1000, + "batch_size": args.batch_size, + "device": str(device), + }, + {"best_val_loss": best_val, "best_epoch": best_epoch, "epochs_ran": len(history)}, + history, + model, + optimizer, + ) + + +def main(): + args = parse_args() + torch.set_num_threads(args.threads) + set_seed(args.seed) + + run_name = f"{args.arm}_{args.target}_{args.train_subset}_seed{args.seed}" + output_dir = Path(args.output_dir) / run_name + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "train.log") + + dataset_dir = Path(args.dataset_dir) + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + rows_by_split = { + "train": load_rows(dataset_dir, f"train_{args.train_subset}"), + "val": load_rows(dataset_dir, "val"), + "test": load_rows(dataset_dir, "test"), + } + targets = { + name: frame[args.target].to_numpy()[rows] for name, rows in rows_by_split.items() + } + # log10 Pgen sits near -6.5, so a freshly initialized network spends its first epochs + # just finding the offset while Ridge fits an intercept for free. Standardizing the + # target on train statistics removes that asymmetry; predictions are mapped back + # before any metric is computed, so the reported errors stay in log10 units. + target_mean = float(targets["train"].mean()) + target_std = max(float(targets["train"].std()), 1e-6) + scaled_targets = {name: (values - target_mean) / target_std for name, values in targets.items()} + log.info( + "run=%s train=%d val=%d test=%d target_mean=%.3f", + run_name, *(len(r) for r in rows_by_split.values()), targets["train"].mean(), + ) + + history, model, optimizer = [], None, None + if args.arm == "tfidf_ridge": + scaled_predictions, training_seconds, speed, extra, history = fit_ridge( + args, frame, rows_by_split, scaled_targets, log + ) + else: + features, _ = build_features(args, frame, rows_by_split, log) + scaled_predictions, training_seconds, speed, extra, history, model, optimizer = fit_neural( + args, features, scaled_targets, log + ) + save_checkpoint(output_dir / "best.pt", model, optimizer, extra["best_epoch"], extra) + + predictions = scaled_predictions * target_std + target_mean + metrics = regression_metrics(predictions, targets["test"]) + pd.DataFrame( + { + "row_index": rows_by_split["test"], + "cdr3": frame.junction_aa.iloc[rows_by_split["test"]].to_numpy(), + "true": targets["test"], + "predicted": predictions, + } + ).to_csv(output_dir / "test_predictions.csv", index=False) + if history: + pd.DataFrame(history).to_csv(output_dir / "history.csv", index=False) + + result = { + "run_name": run_name, + "arm": args.arm, + "target": args.target, + "train_subset": args.train_subset, + "train_size": int(len(rows_by_split["train"])), + "seed": args.seed, + "training_seconds": round(training_seconds, 2), + "test": metrics, + "inference": speed, + "memory": peak_memory_mb(), + "target_scaling": {"mean": target_mean, "std": target_std}, + "arm_details": extra, + "config": vars(args), + } + (output_dir / "metrics.json").write_text(json.dumps(result, indent=2), encoding="utf-8") + + log.info("=" * 60) + log.info( + "test rmse=%.4f mae=%.4f r2=%.4f pearson=%.4f spearman=%.4f bias=%.4f", + metrics["rmse"], metrics["mae"], metrics["r2"], + metrics["pearson_r"], metrics["spearman_rho"], metrics["bias"], + ) + log.info( + "training=%.0f s inference=%.0f seq/s peak_rss=%.0f MB", + training_seconds, speed["sequences_per_second"], result["memory"].get("peak_rss_mb") or 0, + ) + log.info("wrote %s", output_dir / "metrics.json") + + +if __name__ == "__main__": + main() diff --git a/benchmark/train_external_residue_decoder.py b/benchmark/train_external_residue_decoder.py new file mode 100644 index 0000000..0529ca6 --- /dev/null +++ b/benchmark/train_external_residue_decoder.py @@ -0,0 +1,222 @@ +"""Train one decoder on one frozen, unpooled residue-state representation.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import DataLoader + +from benchmark.external_residue_models import ( + FrozenResidueDecoder, + ResidueArrayDataset, + parameter_count, +) +from irrm_codec.losses import inverse_loss +from irrm_codec.tokenization import AA_VOCAB, ID2AA, encode +from irrm_codec.utils import choose_device, save_checkpoint, set_seed, setup_logging + + +GAP_ID = AA_VOCAB["-"] +REPRESENTATIONS = ("esm2_8m", "tcr_bert", "sceptr") + + +def levenshtein_distance(left: str, right: str) -> int: + """Return exact edit distance without an optional third-party dependency.""" + if len(left) > len(right): + left, right = right, left + previous = list(range(len(left) + 1)) + for right_index, right_char in enumerate(right, start=1): + current = [right_index] + for left_index, left_char in enumerate(left, start=1): + current.append( + min( + current[-1] + 1, + previous[left_index] + 1, + previous[left_index - 1] + (left_char != right_char), + ) + ) + previous = current + return previous[-1] + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-dir", required=True) + parser.add_argument("--residue-dir", required=True) + parser.add_argument("--representation", choices=REPRESENTATIONS, required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--epochs", type=int, default=60) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--weight-decay", type=float, default=1e-4) + parser.add_argument("--patience", type=int, default=8) + parser.add_argument("--max-len", type=int, default=40) + parser.add_argument("--threads", type=int, default=4) + parser.add_argument("--num-workers", type=int, default=0) + return parser.parse_args() + + +def load_rows(dataset_dir, split): + return pd.read_csv(Path(dataset_dir) / "manifests" / f"{split}.tsv", sep="\t")[ + "row_index" + ].to_numpy() + + +def decode_tokens(token_ids): + return "".join(ID2AA[int(token)] for token in token_ids if int(token) > GAP_ID) + + +def make_loader(features_path, mask_path, rows, targets, args, shuffle): + dataset = ResidueArrayDataset(features_path, mask_path, rows, targets[rows]) + return DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=shuffle, + num_workers=args.num_workers, + pin_memory=torch.cuda.is_available(), + ) + + +def train_epoch(model, loader, device, optimizer): + model.train() + total, steps = 0.0, 0 + for features, mask, target in loader: + features, mask, target = features.to(device), mask.to(device), target.long().to(device) + loss = inverse_loss(model(features, mask), target) + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + total += loss.item() + steps += 1 + return total / max(steps, 1) + + +@torch.no_grad() +def evaluate(model, loader, device, sequences, rows, collect=False): + model.eval() + total, steps, token_hits, token_total = 0.0, 0, 0, 0 + predictions = [] + for features, mask, target in loader: + features, mask, target = features.to(device), mask.to(device), target.long().to(device) + logits = model(features, mask) + total += inverse_loss(logits, target).item() + steps += 1 + predicted = logits.argmax(dim=-1) + residues = target.gt(GAP_ID) + token_hits += predicted.eq(target).logical_and(residues).sum().item() + token_total += residues.sum().item() + predictions.extend(decode_tokens(row) for row in predicted.cpu().numpy()) + truth = [sequences[index] for index in rows] + distances = np.asarray([levenshtein_distance(pred, true) for pred, true in zip(predictions, truth)]) + lengths = np.asarray([max(len(sequence), 1) for sequence in truth]) + metrics = { + "loss": total / max(steps, 1), + "token_accuracy": token_hits / max(token_total, 1), + "exact_match": float(np.mean(distances == 0)), + "normalized_levenshtein": float(np.mean(distances / lengths)), + "within_edit_distance_1": float(np.mean(distances <= 1)), + "mean_edit_distance": float(distances.mean()), + } + return metrics, (predictions if collect else None), (truth if collect else None) + + +def main(): + args = parse_args() + torch.set_num_threads(args.threads) + set_seed(args.seed) + device = choose_device() + dataset_dir, residue_dir = Path(args.dataset_dir), Path(args.residue_dir) + run_name = f"{args.representation}_residue_seed{args.seed}" + output_dir = Path(args.output_dir) / run_name + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "train.log") + + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + sequences = frame.junction_aa.tolist() + targets = np.asarray([encode(sequence, max_len=args.max_len) for sequence in sequences], dtype=np.int64) + rows = {split: load_rows(dataset_dir, split) for split in ("train", "val", "test")} + feature_path = residue_dir / f"{args.representation}_residue.npy" + mask_path = residue_dir / f"{args.representation}_residue_mask.npy" + feature_matrix = np.load(feature_path, mmap_mode="r") + if feature_matrix.shape[0] != len(frame) or feature_matrix.shape[1] != args.max_len: + raise ValueError(f"Frozen feature shape mismatch: {feature_matrix.shape}") + loaders = { + split: make_loader(feature_path, mask_path, rows[split], targets, args, split == "train") + for split in rows + } + model = FrozenResidueDecoder( + input_dim=feature_matrix.shape[-1], max_len=args.max_len + ).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + log.info( + "run=%s device=%s frozen_shape=%s params=%d train=%d val=%d test=%d", + run_name, + device, + feature_matrix.shape, + parameter_count(model), + *(len(rows[split]) for split in ("train", "val", "test")), + ) + + best_val, best_epoch, history = float("inf"), 0, [] + started = time.perf_counter() + for epoch in range(1, args.epochs + 1): + train_loss = train_epoch(model, loaders["train"], device, optimizer) + val_metrics, _, _ = evaluate(model, loaders["val"], device, sequences, rows["val"]) + history.append({"epoch": epoch, "train_loss": train_loss, **{f"val_{k}": v for k, v in val_metrics.items()}}) + log.info( + "epoch=%d train_loss=%.4f val_loss=%.4f val_exact=%.4f val_token=%.4f", + epoch, + train_loss, + val_metrics["loss"], + val_metrics["exact_match"], + val_metrics["token_accuracy"], + ) + if val_metrics["loss"] < best_val - 1e-5: + best_val, best_epoch = val_metrics["loss"], epoch + save_checkpoint(output_dir / "best.pt", model, optimizer, epoch, val_metrics) + elif epoch - best_epoch >= args.patience: + break + training_seconds = time.perf_counter() - started + checkpoint = torch.load(output_dir / "best.pt", map_location=device) + model.load_state_dict(checkpoint["model_state"]) + test, predictions, truth = evaluate( + model, loaders["test"], device, sequences, rows["test"], collect=True + ) + pd.DataFrame( + { + "row_index": rows["test"], + "true_cdr3": truth, + "predicted_cdr3": predictions, + "edit_distance": [levenshtein_distance(pred, true) for pred, true in zip(predictions, truth)], + } + ).to_csv(output_dir / "test_predictions.csv", index=False) + pd.DataFrame(history).to_csv(output_dir / "history.csv", index=False) + result = { + "run_name": run_name, + "representation": args.representation, + "interface": "frozen_residue_states_no_global_pooling", + "seed": args.seed, + "train_size": int(len(rows["train"])), + "best_epoch": best_epoch, + "best_val_loss": best_val, + "epochs_ran": len(history), + "training_seconds": round(training_seconds, 2), + "parameters": parameter_count(model), + "frozen_shape": list(feature_matrix.shape), + "test": test, + "config": vars(args), + } + (output_dir / "metrics.json").write_text(json.dumps(result, indent=2), encoding="utf-8") + log.info("status=complete test=%s", test) + + +if __name__ == "__main__": + main() diff --git a/benchmark/train_external_residue_pgen.py b/benchmark/train_external_residue_pgen.py new file mode 100644 index 0000000..d01dcd3 --- /dev/null +++ b/benchmark/train_external_residue_pgen.py @@ -0,0 +1,182 @@ +"""Train a position-aware pgen probe on frozen, unpooled residue states.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from scipy import stats +from torch.utils.data import DataLoader + +from benchmark.external_residue_models import ( + FrozenResiduePgenHead, + ResidueArrayDataset, + parameter_count, +) +from irrm_codec.losses import pgen_loss +from irrm_codec.utils import choose_device, save_checkpoint, set_seed, setup_logging + + +REPRESENTATIONS = ("esm2_8m", "tcr_bert", "sceptr") + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-dir", required=True) + parser.add_argument("--residue-dir", required=True) + parser.add_argument("--representation", choices=REPRESENTATIONS, required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--target", default="log10_pgen_1mm") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--epochs", type=int, default=60) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--weight-decay", type=float, default=1e-4) + parser.add_argument("--patience", type=int, default=12) + parser.add_argument("--threads", type=int, default=4) + parser.add_argument("--num-workers", type=int, default=0) + return parser.parse_args() + + +def load_rows(dataset_dir, split): + return pd.read_csv(Path(dataset_dir) / "manifests" / f"{split}.tsv", sep="\t")[ + "row_index" + ].to_numpy() + + +def regression_metrics(prediction, truth): + residual = prediction - truth + variance = np.var(truth) + return { + "rmse": float(np.sqrt(np.mean(residual**2))), + "mae": float(np.mean(np.abs(residual))), + "bias": float(np.mean(residual)), + "r2": float(1 - np.mean(residual**2) / variance), + "pearson_r": float(stats.pearsonr(prediction, truth)[0]), + "spearman_rho": float(stats.spearmanr(prediction, truth)[0]), + } + + +def run_epoch(model, loader, device, optimizer=None): + training = optimizer is not None + model.train(training) + total, steps = 0.0, 0 + for features, mask, target in loader: + features, mask, target = features.to(device), mask.to(device), target.float().to(device) + with torch.set_grad_enabled(training): + loss = pgen_loss(model(features, mask), target) + if training: + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + total += loss.item() + steps += 1 + return total / max(steps, 1) + + +@torch.no_grad() +def predict(model, loader, device): + model.eval() + output = [] + for features, mask, _ in loader: + output.append(model(features.to(device), mask.to(device)).cpu().numpy()) + return np.concatenate(output) + + +def main(): + args = parse_args() + torch.set_num_threads(args.threads) + set_seed(args.seed) + device = choose_device() + dataset_dir, residue_dir = Path(args.dataset_dir), Path(args.residue_dir) + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + rows = {split: load_rows(dataset_dir, split) for split in ("train", "val", "test")} + raw_target = frame[args.target].to_numpy(dtype=np.float64) + target_mean = float(raw_target[rows["train"]].mean()) + target_std = float(raw_target[rows["train"]].std()) + scaled = ((raw_target - target_mean) / target_std).astype(np.float32) + feature_path = residue_dir / f"{args.representation}_residue.npy" + mask_path = residue_dir / f"{args.representation}_residue_mask.npy" + feature_matrix = np.load(feature_path, mmap_mode="r") + loaders = { + split: DataLoader( + ResidueArrayDataset(feature_path, mask_path, rows[split], scaled[rows[split]]), + batch_size=args.batch_size, + shuffle=(split == "train"), + num_workers=args.num_workers, + pin_memory=torch.cuda.is_available(), + ) + for split in rows + } + model = FrozenResiduePgenHead(input_dim=feature_matrix.shape[-1]).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) + run_name = f"{args.representation}_residue_{args.target}_seed{args.seed}" + output_dir = Path(args.output_dir) / run_name + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "train.log") + log.info( + "run=%s device=%s frozen_shape=%s params=%d", + run_name, + device, + feature_matrix.shape, + parameter_count(model), + ) + best_val, best_epoch, history = float("inf"), 0, [] + started = time.perf_counter() + for epoch in range(1, args.epochs + 1): + train_loss = run_epoch(model, loaders["train"], device, optimizer) + val_loss = run_epoch(model, loaders["val"], device) + scheduler.step() + history.append({"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss}) + if epoch % 5 == 0 or epoch == 1: + log.info("epoch=%d train_loss=%.5f val_loss=%.5f", epoch, train_loss, val_loss) + if val_loss < best_val - 1e-6: + best_val, best_epoch = val_loss, epoch + save_checkpoint(output_dir / "best.pt", model, optimizer, epoch, {"val_loss": val_loss}) + elif epoch - best_epoch >= args.patience: + break + training_seconds = time.perf_counter() - started + model.load_state_dict(torch.load(output_dir / "best.pt", map_location=device)["model_state"]) + scaled_prediction = predict(model, loaders["test"], device) + prediction = scaled_prediction * target_std + target_mean + truth = raw_target[rows["test"]] + metrics = regression_metrics(prediction, truth) + pd.DataFrame(history).to_csv(output_dir / "history.csv", index=False) + pd.DataFrame( + { + "row_index": rows["test"], + "cdr3": frame.junction_aa.iloc[rows["test"]].to_numpy(), + "true": truth, + "predicted": prediction, + } + ).to_csv(output_dir / "test_predictions.csv", index=False) + result = { + "run_name": run_name, + "representation": args.representation, + "interface": "frozen_residue_states_no_global_pooling", + "target": args.target, + "seed": args.seed, + "train_size": int(len(rows["train"])), + "best_epoch": best_epoch, + "best_val_loss": best_val, + "epochs_ran": len(history), + "training_seconds": round(training_seconds, 2), + "parameters": parameter_count(model), + "frozen_shape": list(feature_matrix.shape), + "target_scaling": {"fit_split": "train", "mean": target_mean, "std": target_std}, + "test": metrics, + "config": vars(args), + } + (output_dir / "metrics.json").write_text(json.dumps(result, indent=2), encoding="utf-8") + log.info("status=complete test=%s", metrics) + + +if __name__ == "__main__": + main() diff --git a/benchmark/train_forward_encoder.py b/benchmark/train_forward_encoder.py new file mode 100644 index 0000000..8589958 --- /dev/null +++ b/benchmark/train_forward_encoder.py @@ -0,0 +1,206 @@ +"""Issue 3, step 1: pretrain the sequence-to-TCRemP encoder on the benchmark split. + +One of the Pgen arms is initialized from a pretrained sequence-to-TCRemP encoder. The +published checkpoints were trained with their own random split, so reusing one would leak +benchmark test sequences into pretraining and inflate exactly the comparison the issue +asks about. This trains the encoder on the benchmark training split only. + +PgenModel is ForwardModel with output_dim=1, so everything except the final output layer +transfers directly into the Pgen arm. +""" + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import DataLoader, Dataset + +from irrm_codec.forward_model import ForwardModel +from irrm_codec.losses import forward_loss, forward_metrics +from irrm_codec.tokenization import PAD_ID, encode +from irrm_codec.utils import choose_device, save_checkpoint, set_seed, setup_logging + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument("--output-dir", default="artifacts/benchmark/forward_encoder") + p.add_argument("--train-subset", default="all") + p.add_argument("--seed", type=int, default=42) + p.add_argument("--epochs", type=int, default=40) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--weight-decay", type=float, default=1e-4) + p.add_argument("--patience", type=int, default=6) + p.add_argument("--max-len", type=int, default=40) + p.add_argument("--threads", type=int, default=8) + return p.parse_args() + + +class SequenceToEmbedding(Dataset): + """Tokenized CDR3 in, normalized TCRemP vector out. + + Shuffled training reads the 9000-wide rows in random order, which is slow against a + memory-mapped file on network storage. The split is materialized once instead: the + full training slice is about 2.9 GB resident, which is cheaper than re-reading it + from NFS every epoch. + """ + + def __init__(self, matrix, rows, sequences, mean, std, max_len): + self.tokens = torch.tensor( + [encode(sequences[i], max_len=max_len) for i in rows], dtype=torch.long + ) + targets = np.empty((len(rows), matrix.shape[1]), dtype=np.float32) + for start in range(0, len(rows), 4096): + block = rows[start : start + 4096] + targets[start : start + len(block)] = np.asarray(matrix[block], dtype=np.float32) + self.targets = torch.from_numpy((targets - mean) / std) + + def __len__(self): + return len(self.tokens) + + def __getitem__(self, index): + tokens = self.tokens[index] + return tokens, tokens.ne(PAD_ID), self.targets[index] + + +def load_rows(dataset_dir, name): + return pd.read_csv(Path(dataset_dir) / "manifests" / f"{name}.tsv", sep="\t")["row_index"].to_numpy() + + +def train_statistics(matrix, rows, batch_size=4096): + """Streaming mean and standard deviation over the training rows only.""" + total = np.zeros(matrix.shape[1], dtype=np.float64) + total_sq = np.zeros(matrix.shape[1], dtype=np.float64) + for start in range(0, len(rows), batch_size): + block = np.asarray(matrix[rows[start : start + batch_size]], dtype=np.float64) + total += block.sum(axis=0) + total_sq += (block ** 2).sum(axis=0) + mean = total / len(rows) + variance = np.maximum(total_sq / len(rows) - mean ** 2, 0.0) + return mean.astype(np.float32), np.sqrt(variance).clip(1e-6).astype(np.float32) + + +def run_epoch(model, loader, device, optimizer=None): + is_train = optimizer is not None + model.train(mode=is_train) + sums = {"loss": 0.0, "mse": 0.0, "cosine": 0.0} + steps = 0 + + for tokens, mask, target in loader: + tokens, mask, target = tokens.to(device), mask.to(device), target.to(device) + with torch.set_grad_enabled(is_train): + prediction = model(tokens, mask) + loss = forward_loss(prediction, target) + if is_train: + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + + metrics = forward_metrics(prediction.detach(), target) + sums["loss"] += loss.item() + sums["mse"] += metrics["mse"] + sums["cosine"] += metrics["cosine"] + steps += 1 + + return {name: value / max(steps, 1) for name, value in sums.items()} + + +def main(): + args = parse_args() + torch.set_num_threads(args.threads) + set_seed(args.seed) + device = choose_device() + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "train.log") + + dataset_dir = Path(args.dataset_dir) + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + sequences = frame["junction_aa"].tolist() + matrix = np.load(dataset_dir / "embeddings.npy", mmap_mode="r") + + train_rows = load_rows(dataset_dir, f"train_{args.train_subset}") + val_rows = load_rows(dataset_dir, "val") + log.info("device=%s embeddings=%s train=%d val=%d", device, matrix.shape, len(train_rows), len(val_rows)) + + mean, std = train_statistics(matrix, train_rows) + np.save(output_dir / "mean.npy", mean) + np.save(output_dir / "std.npy", std) + + loaders = { + name: DataLoader( + SequenceToEmbedding(matrix, rows, sequences, mean, std, args.max_len), + batch_size=args.batch_size, + shuffle=(name == "train"), + ) + for name, rows in (("train", train_rows), ("val", val_rows)) + } + + model = ForwardModel(output_dim=matrix.shape[1], max_len=args.max_len).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + # A constant learning rate leaves the model bouncing around the minimum, which shows + # up as an oscillating validation loss and, downstream, as large seed-to-seed spread. + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) + log.info("model params=%.2fM output_dim=%d", sum(p.numel() for p in model.parameters()) / 1e6, matrix.shape[1]) + + best_val = float("inf") + best_epoch = 0 + history = [] + started = time.perf_counter() + + for epoch in range(1, args.epochs + 1): + train_metrics = run_epoch(model, loaders["train"], device, optimizer) + val_metrics = run_epoch(model, loaders["val"], device) + scheduler.step() + history.append( + { + "epoch": epoch, + **{f"train_{k}": v for k, v in train_metrics.items()}, + **{f"val_{k}": v for k, v in val_metrics.items()}, + } + ) + log.info( + "epoch=%d train_loss=%.4f val_loss=%.4f val_cosine=%.4f", + epoch, train_metrics["loss"], val_metrics["loss"], val_metrics["cosine"], + ) + + if val_metrics["loss"] < best_val - 1e-5: + best_val, best_epoch = val_metrics["loss"], epoch + save_checkpoint(output_dir / "best.pt", model, optimizer, epoch, val_metrics) + elif epoch - best_epoch >= args.patience: + log.info("early stopping at epoch=%d (best epoch=%d)", epoch, best_epoch) + break + + elapsed = time.perf_counter() - started + pd.DataFrame(history).to_csv(output_dir / "history.csv", index=False) + (output_dir / "metrics.json").write_text( + json.dumps( + { + "best_val_loss": best_val, + "best_epoch": best_epoch, + "epochs_ran": len(history), + "training_seconds": round(elapsed, 2), + "embedding_dim": int(matrix.shape[1]), + "train_size": int(len(train_rows)), + "final": history[-1], + "config": vars(args), + }, + indent=2, + ), + encoding="utf-8", + ) + + log.info("=" * 60) + log.info("best_val_loss=%.4f at epoch=%d in %.0f s", best_val, best_epoch, elapsed) + log.info("checkpoint for Pgen transfer: %s", output_dir / "best.pt") + + +if __name__ == "__main__": + main() diff --git a/benchmark/train_pgen_arm.py b/benchmark/train_pgen_arm.py new file mode 100644 index 0000000..6a7bd25 --- /dev/null +++ b/benchmark/train_pgen_arm.py @@ -0,0 +1,399 @@ +"""Issue 3: train one Pgen-prediction arm and score it on the benchmark test split. + +Arms: + irrm_scratch PgenModel from a random initialization + irrm_pretrained PgenModel initialized from the sequence-to-TCRemP encoder + tfidf_ridge 1-3-mer character TF-IDF + Ridge regression + sceptr_mlp frozen SCEPTR embeddings + the shared regression head + tcr_bert_mlp frozen TCR-BERT embeddings + the same head + esm2_8m_mlp frozen ESM-2 8M embeddings + the same head + +The three frozen-embedding arms share one head and one training configuration, so the +only thing that varies between them is the representation. The two IRRM arms differ only +in initialization, which is the comparison the issue asks about. + +Targets stay in raw log10 units, matching the repo's pgen trainer, so RMSE, MAE and bias +are directly interpretable as log10 error. +""" + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +from scipy import stats +from torch.utils.data import DataLoader, TensorDataset + +from irrm_codec.losses import pgen_loss +from irrm_codec.pgen_model import PgenModel +from irrm_codec.tokenization import PAD_ID, encode +from irrm_codec.utils import choose_device, save_checkpoint, set_seed, setup_logging + +SEQUENCE_ARMS = ("irrm_scratch", "irrm_pretrained") +FROZEN_ARMS = {"sceptr_mlp": "sceptr", "tcr_bert_mlp": "tcr_bert", "esm2_8m_mlp": "esm2_8m"} +ARMS = (*SEQUENCE_ARMS, "tfidf_ridge", *FROZEN_ARMS) + +# Ridge has no training loop, so the penalty is selected on validation instead. This is +# the baseline's equivalent of early stopping; fixing it arbitrarily would handicap it. +RIDGE_ALPHAS = (0.01, 0.1, 1.0, 10.0, 100.0) + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--arm", required=True, choices=ARMS) + p.add_argument("--target", default="log10_pgen_1mm", choices=("log10_pgen", "log10_pgen_1mm")) + p.add_argument("--dataset-dir", default="data/benchmark/trb") + p.add_argument("--representations-dir", default="data/benchmark/trb/representations") + p.add_argument("--pretrained-checkpoint", default="artifacts/benchmark/forward_encoder/best.pt") + p.add_argument("--output-dir", default="artifacts/benchmark/pgen") + p.add_argument("--train-subset", default="all") + p.add_argument("--seed", type=int, default=42) + p.add_argument("--epochs", type=int, default=60) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--weight-decay", type=float, default=1e-4) + # Wider than the reconstruction runs use: the cosine schedule only pays off if training + # reaches the low-learning-rate tail, and stopping at epoch 17 of 60 would decay the + # rate by under a fifth. + p.add_argument("--patience", type=int, default=12) + p.add_argument("--max-len", type=int, default=40) + p.add_argument("--threads", type=int, default=8) + return p.parse_args() + + +class RegressionHead(nn.Module): + """One head shared by every frozen representation, so only the input differs.""" + + def __init__(self, input_dim, hidden_dims=(512, 256), dropout=0.2): + super().__init__() + layers = [] + previous = input_dim + for width in hidden_dims: + layers += [nn.Linear(previous, width), nn.GELU(), nn.Dropout(dropout)] + previous = width + layers.append(nn.Linear(previous, 1)) + self.net = nn.Sequential(*layers) + + def forward(self, features): + return self.net(features).squeeze(-1) + + +def load_rows(dataset_dir, name): + return pd.read_csv(Path(dataset_dir) / "manifests" / f"{name}.tsv", sep="\t")["row_index"].to_numpy() + + +def peak_memory_mb(): + """True peak resident set size, plus GPU peak when a device was used.""" + result = {} + try: + import resource + + # ru_maxrss is in kilobytes on Linux. + result["peak_rss_mb"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 + except (ImportError, AttributeError): + try: + import psutil + + result["peak_rss_mb"] = psutil.Process().memory_info().rss / 1024**2 + except ImportError: + result["peak_rss_mb"] = None + if torch.cuda.is_available(): + result["peak_gpu_mb"] = torch.cuda.max_memory_allocated() / 1024**2 + return result + + +def regression_metrics(predictions, truth): + residual = predictions - truth + total_variance = np.var(truth) + return { + "rmse": float(np.sqrt(np.mean(residual**2))), + "mae": float(np.mean(np.abs(residual))), + "bias": float(np.mean(residual)), + "r2": float(1.0 - np.mean(residual**2) / total_variance) if total_variance > 0 else float("nan"), + "pearson_r": float(stats.pearsonr(predictions, truth)[0]), + "spearman_rho": float(stats.spearmanr(predictions, truth)[0]), + } + + +def build_features(args, frame, rows_by_split, log): + """Return per-split model inputs for the requested arm.""" + if args.arm in SEQUENCE_ARMS: + tokens = { + name: torch.tensor( + [encode(frame.junction_aa.iloc[i], max_len=args.max_len) for i in rows], + dtype=torch.long, + ) + for name, rows in rows_by_split.items() + } + return tokens, None + + representation = FROZEN_ARMS[args.arm] + matrix = np.load(Path(args.representations_dir) / f"{representation}.npy", mmap_mode="r") + features = { + name: torch.from_numpy(np.asarray(matrix[rows], dtype=np.float32)) + for name, rows in rows_by_split.items() + } + # Train-split statistics only: a shared learning rate across representations with + # different scales would not be the identical setting the issue requires. + mean = features["train"].mean(dim=0, keepdim=True) + std = features["train"].std(dim=0, keepdim=True).clamp_min(1e-6) + for name in features: + features[name] = (features[name] - mean) / std + log.info("%s features dim=%d", representation, features["train"].shape[1]) + return features, (mean, std) + + +def build_model(args, input_dim, log): + if args.arm in SEQUENCE_ARMS: + model = PgenModel(max_len=args.max_len) + if args.arm == "irrm_pretrained": + state = torch.load(args.pretrained_checkpoint, map_location="cpu")["model_state"] + own = model.state_dict() + # Everything but the output layer matches: PgenModel is ForwardModel with + # output_dim=1, so only mlp. differs in shape. + transferable = {k: v for k, v in state.items() if k in own and own[k].shape == v.shape} + model.load_state_dict(transferable, strict=False) + log.info( + "transferred %d/%d tensors from %s", + len(transferable), len(own), args.pretrained_checkpoint, + ) + return model + return RegressionHead(input_dim) + + +def run_epoch(model, loader, device, is_sequence, optimizer=None): + is_train = optimizer is not None + model.train(mode=is_train) + loss_sum, steps = 0.0, 0 + for batch in loader: + inputs, target = batch[0].to(device), batch[-1].to(device) + with torch.set_grad_enabled(is_train): + prediction = model(inputs, inputs.ne(PAD_ID)) if is_sequence else model(inputs) + loss = pgen_loss(prediction, target) + if is_train: + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + loss_sum += loss.item() + steps += 1 + return loss_sum / max(steps, 1) + + +@torch.no_grad() +def predict_all(model, inputs, device, is_sequence, batch_size): + model.eval() + chunks = [] + for start in range(0, len(inputs), batch_size): + block = inputs[start : start + batch_size].to(device) + prediction = model(block, block.ne(PAD_ID)) if is_sequence else model(block) + chunks.append(prediction.cpu().numpy()) + return np.concatenate(chunks) + + +def fit_ridge(args, frame, rows_by_split, targets, log): + """1-3-mer character TF-IDF into Ridge, with the penalty chosen on validation.""" + from sklearn.feature_extraction.text import TfidfVectorizer + from sklearn.linear_model import Ridge + + started = time.perf_counter() + sequences = { + name: frame.junction_aa.iloc[rows].tolist() for name, rows in rows_by_split.items() + } + vectorizer = TfidfVectorizer(analyzer="char", ngram_range=(1, 3)) + train_features = vectorizer.fit_transform(sequences["train"]) + val_features = vectorizer.transform(sequences["val"]) + + best = (None, float("inf"), None) + for alpha in RIDGE_ALPHAS: + model = Ridge(alpha=alpha).fit(train_features, targets["train"]) + rmse = float(np.sqrt(np.mean((model.predict(val_features) - targets["val"]) ** 2))) + log.info("ridge alpha=%.3g val_rmse=%.4f", alpha, rmse) + if rmse < best[1]: + best = (model, rmse, alpha) + model, val_rmse, alpha = best + training_seconds = time.perf_counter() - started + log.info("selected alpha=%.3g val_rmse=%.4f features=%d", alpha, val_rmse, train_features.shape[1]) + + # Vectorization is part of this arm's inference path, so it belongs inside the timed + # section; excluding it would report the Ridge solve alone and overstate throughput + # by orders of magnitude. + inference_started = time.perf_counter() + test_features = vectorizer.transform(sequences["test"]) + predictions = model.predict(test_features) + elapsed = time.perf_counter() - inference_started + + return ( + predictions, + training_seconds, + { + "sequences_per_second": len(predictions) / elapsed, + "ms_per_sequence": elapsed / len(predictions) * 1000, + "batch_size": len(predictions), + "device": "cpu", + }, + {"alpha": alpha, "val_rmse": val_rmse, "n_features": int(train_features.shape[1])}, + [], + ) + + +def fit_neural(args, features, targets, log): + device = choose_device() + is_sequence = args.arm in SEQUENCE_ARMS + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + loaders = { + name: DataLoader( + TensorDataset(features[name], torch.from_numpy(targets[name].astype(np.float32))), + batch_size=args.batch_size, + shuffle=(name == "train"), + ) + for name in ("train", "val") + } + + model = build_model(args, features["train"].shape[-1], log).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + # A constant learning rate leaves the model bouncing around the minimum: validation + # loss oscillates, early stopping fires at an arbitrary epoch, and the IRRM arms end up + # spreading roughly tenfold wider across seeds than the frozen-embedding arms. + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) + log.info("device=%s params=%.2fM", device, sum(p.numel() for p in model.parameters()) / 1e6) + + best_val, best_epoch = float("inf"), 0 + best_state = None + history = [] + started = time.perf_counter() + + for epoch in range(1, args.epochs + 1): + train_loss = run_epoch(model, loaders["train"], device, is_sequence, optimizer) + val_loss = run_epoch(model, loaders["val"], device, is_sequence) + scheduler.step() + history.append({"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss}) + if epoch % 5 == 0 or epoch == 1: + log.info("epoch=%d train_loss=%.4f val_loss=%.4f", epoch, train_loss, val_loss) + + if val_loss < best_val - 1e-6: + best_val, best_epoch = val_loss, epoch + best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} + elif epoch - best_epoch >= args.patience: + log.info("early stopping at epoch=%d (best epoch=%d)", epoch, best_epoch) + break + + training_seconds = time.perf_counter() - started + if best_state is not None: + model.load_state_dict(best_state) + + inference_started = time.perf_counter() + predictions = predict_all(model, features["test"], device, is_sequence, args.batch_size) + elapsed = time.perf_counter() - inference_started + + return ( + predictions, + training_seconds, + { + "sequences_per_second": len(predictions) / elapsed, + "ms_per_sequence": elapsed / len(predictions) * 1000, + "batch_size": args.batch_size, + "device": str(device), + }, + {"best_val_loss": best_val, "best_epoch": best_epoch, "epochs_ran": len(history)}, + history, + model, + optimizer, + ) + + +def main(): + args = parse_args() + torch.set_num_threads(args.threads) + set_seed(args.seed) + + run_name = f"{args.arm}_{args.target}_{args.train_subset}_seed{args.seed}" + output_dir = Path(args.output_dir) / run_name + output_dir.mkdir(parents=True, exist_ok=True) + log = setup_logging(output_dir / "train.log") + + dataset_dir = Path(args.dataset_dir) + frame = pd.read_parquet(dataset_dir / "dataset.parquet") + rows_by_split = { + "train": load_rows(dataset_dir, f"train_{args.train_subset}"), + "val": load_rows(dataset_dir, "val"), + "test": load_rows(dataset_dir, "test"), + } + targets = { + name: frame[args.target].to_numpy()[rows] for name, rows in rows_by_split.items() + } + # log10 Pgen sits near -6.5, so a freshly initialized network spends its first epochs + # just finding the offset while Ridge fits an intercept for free. Standardizing the + # target on train statistics removes that asymmetry; predictions are mapped back + # before any metric is computed, so the reported errors stay in log10 units. + target_mean = float(targets["train"].mean()) + target_std = max(float(targets["train"].std()), 1e-6) + scaled_targets = {name: (values - target_mean) / target_std for name, values in targets.items()} + log.info( + "run=%s train=%d val=%d test=%d target_mean=%.3f", + run_name, *(len(r) for r in rows_by_split.values()), targets["train"].mean(), + ) + + history, model, optimizer = [], None, None + if args.arm == "tfidf_ridge": + scaled_predictions, training_seconds, speed, extra, history = fit_ridge( + args, frame, rows_by_split, scaled_targets, log + ) + else: + features, _ = build_features(args, frame, rows_by_split, log) + scaled_predictions, training_seconds, speed, extra, history, model, optimizer = fit_neural( + args, features, scaled_targets, log + ) + save_checkpoint(output_dir / "best.pt", model, optimizer, extra["best_epoch"], extra) + + predictions = scaled_predictions * target_std + target_mean + metrics = regression_metrics(predictions, targets["test"]) + pd.DataFrame( + { + "row_index": rows_by_split["test"], + "cdr3": frame.junction_aa.iloc[rows_by_split["test"]].to_numpy(), + "true": targets["test"], + "predicted": predictions, + } + ).to_csv(output_dir / "test_predictions.csv", index=False) + if history: + pd.DataFrame(history).to_csv(output_dir / "history.csv", index=False) + + result = { + "run_name": run_name, + "arm": args.arm, + "target": args.target, + "train_subset": args.train_subset, + "train_size": int(len(rows_by_split["train"])), + "seed": args.seed, + "training_seconds": round(training_seconds, 2), + "test": metrics, + "inference": speed, + "memory": peak_memory_mb(), + "target_scaling": {"mean": target_mean, "std": target_std}, + "arm_details": extra, + "config": vars(args), + } + (output_dir / "metrics.json").write_text(json.dumps(result, indent=2), encoding="utf-8") + + log.info("=" * 60) + log.info( + "test rmse=%.4f mae=%.4f r2=%.4f pearson=%.4f spearman=%.4f bias=%.4f", + metrics["rmse"], metrics["mae"], metrics["r2"], + metrics["pearson_r"], metrics["spearman_rho"], metrics["bias"], + ) + log.info( + "training=%.0f s inference=%.0f seq/s peak_rss=%.0f MB", + training_seconds, speed["sequences_per_second"], result["memory"].get("peak_rss_mb") or 0, + ) + log.info("wrote %s", output_dir / "metrics.json") + + +if __name__ == "__main__": + main() diff --git a/benchmark/validate_single_aa_targets.py b/benchmark/validate_single_aa_targets.py new file mode 100644 index 0000000..68c1998 --- /dev/null +++ b/benchmark/validate_single_aa_targets.py @@ -0,0 +1,99 @@ +"""Identity-align independently scored mutation-challenge targets.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def align_targets(prepared: Path, pgen_path: Path, representations_path: Path, + tcremp_path: Path, output_dir: Path) -> dict: + mutants = pd.read_csv(prepared / "mutants.tsv", sep="\t", dtype={"clone_id": str}) + pgen = pd.read_csv(pgen_path, sep="\t", dtype={"clone_id": str}) + representations = pd.read_csv(representations_path, sep="\t", dtype={"clone_id": str}) + tcremp = pd.read_parquet(tcremp_path) + for name, frame in (("mutants", mutants), ("pgen", pgen), + ("representations", representations), ("tcremp", tcremp)): + if "clone_id" not in frame or frame["clone_id"].isna().any() or frame["clone_id"].duplicated().any(): + raise ValueError(f"{name} lacks complete unique clone_id values.") + expected = mutants["clone_id"].astype(str).tolist() + expected_set = set(expected) + for name, frame in (("pgen", pgen), ("representations", representations), ("tcremp", tcremp)): + observed = set(frame["clone_id"].astype(str)) + if observed != expected_set: + raise ValueError(f"{name} clone IDs differ: missing={len(expected_set-observed)} extra={len(observed-expected_set)}") + pgen = pgen.set_index("clone_id").loc[expected] + representations = representations.set_index("clone_id").loc[expected] + tcremp = tcremp.set_index("clone_id").loc[expected] + sidecar_sequence_column = next( + (column for column in ("cdr3aa_TRB", "junction_aa") if column in representations), None + ) + if sidecar_sequence_column is None: + raise ValueError(f"TCRemP sidecar lacks a recognized TRB sequence column: {list(representations.columns)}") + if not np.array_equal(representations[sidecar_sequence_column].astype(str).to_numpy(), mutants["junction_aa"].astype(str).to_numpy()): + raise ValueError("TCRemP sidecar sequences do not align to generated mutants by clone_id.") + if not np.array_equal(pgen["junction_aa"].astype(str).to_numpy(), mutants["junction_aa"].astype(str).to_numpy()): + raise ValueError("Pgen sequences do not align to generated mutants by clone_id.") + exact_values = pd.to_numeric(pgen["log10_pgen"], errors="coerce").to_numpy(np.float64) + target_values = pd.to_numeric(pgen["log10_pgen_1mm"], errors="coerce").to_numpy(np.float64) + if not np.isfinite(target_values).all(): + raise ValueError( + f"Nonfinite declared log10_pgen_1mm targets: {int((~np.isfinite(target_values)).sum())}" + ) + matrix = tcremp.to_numpy(np.float32) + if matrix.ndim != 2 or matrix.shape[1] != 9000 or not np.isfinite(matrix).all(): + raise ValueError(f"Invalid TCRemP matrix {matrix.shape}.") + output_dir.mkdir(parents=True, exist_ok=False) + matrix_path = output_dir / "mutant_tcremp.npy" + target_path = output_dir / "targets.tsv" + np.save(matrix_path, matrix, allow_pickle=False) + result = mutants.copy() + # Exact pgen can legitimately be zero for a valid mutant. It is not an + # evaluation target; retain its nonfinite count as a diagnostic but keep + # only the finite 1-mismatch target in the accepted table. + result["log10_pgen_1mm"] = target_values + result.to_csv(target_path, sep="\t", index=False, lineterminator="\n") + report = { + "status": "accepted_for_frozen_delta_evaluation", + "alignment": "unique clone_id plus exact junction_aa agreement", + "tcremp_sidecar_sequence_column": sidecar_sequence_column, + "rows": len(result), + "tcremp_shape": list(matrix.shape), + "pgen_target": "log10_pgen_1mm", + "finite_pgen_target": True, + "nonfinite_exact_log10_pgen_diagnostic_rows": int((~np.isfinite(exact_values)).sum()), + "finite_tcremp": True, + "sources": {str(path): sha256(path) for path in (pgen_path, representations_path, tcremp_path)}, + "artifacts": {matrix_path.name: sha256(matrix_path), target_path.name: sha256(target_path)}, + } + (output_dir / "TARGETS.json").write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + return report + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--prepared", type=Path, required=True) + parser.add_argument("--pgen", type=Path, required=True) + parser.add_argument("--representations", type=Path, required=True) + parser.add_argument("--tcremp", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + print(json.dumps(align_targets(args.prepared, args.pgen, args.representations, + args.tcremp, args.output_dir), sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/viral_antigen_species_2025_07_30.txt b/benchmark/viral_antigen_species_2025_07_30.txt new file mode 100644 index 0000000..cf35dc8 --- /dev/null +++ b/benchmark/viral_antigen_species_2025_07_30.txt @@ -0,0 +1,7 @@ +# Exact antigen.species values in the pinned VDJdb 2025-07-30 snapshot that +# denote viral antigens after the preregistered A*02:01/B2M/MHCI/TRB filters. +# Matching is exact and case-sensitive; HomoSapiens is intentionally excluded. +InfluenzaA +EBV +SARS-CoV-2 +HIV-1 diff --git a/data/benchmark/trb/dataset_summary.json b/data/benchmark/trb/dataset_summary.json index 1c46115..c45b3af 100644 --- a/data/benchmark/trb/dataset_summary.json +++ b/data/benchmark/trb/dataset_summary.json @@ -13,14 +13,16 @@ "dropped_invalid_chars": 0, "dropped_too_short": 0, "dropped_too_long": 0, + "dropped_non_functional_v_call": 335, + "dropped_non_functional_j_call": 0, "dropped_duplicate_cdr3": 0, - "rows_kept": 99998, + "rows_kept": 99663, "dropped_non_finite_pgen": 2 }, "embedding_dim": 9000, "embedding_alignment": { - "length_corr": 0.5274442408242208, - "shuffled_length_corr": 0.0008568142944854723, + "length_corr": 0.5251705966285264, + "shuffled_length_corr": 0.008426662808806782, "min_corr_threshold": 0.3 }, "alignment_mode": "row_order", @@ -29,9 +31,9 @@ "path": "data\\benchmark\\trb\\pgen.tsv" }, "split_sizes": { - "train": 79998, - "val": 9999, - "test": 10001 + "train": 79730, + "val": 9966, + "test": 9967 }, "split_fractions": { "train": 0.8, @@ -41,7 +43,7 @@ "train_subsets": { "1k": 1000, "10k": 10000, - "all": 79998 + "all": 79730 }, "cdr3_overlap_between_splits": { "train_val": 0, @@ -51,54 +53,54 @@ "targets": { "log10_pgen": { "train": { - "n_finite": 79998, + "n_finite": 79730, "n_non_finite": 0, - "mean": -8.34572029024153, - "std": 1.8845455361342092, + "mean": -8.342815174291603, + "std": 1.8788593865723895, "min": -28.06169119699202, "max": -5.092073042426336 }, "val": { - "n_finite": 9999, + "n_finite": 9966, "n_non_finite": 0, - "mean": -8.351515526786393, - "std": 1.909456406125435, - "min": -23.09528585214532, - "max": -5.365645796309452 + "mean": -8.330849724502892, + "std": 1.8649352353948874, + "min": -23.57673660375092, + "max": -5.252935754631411 }, "test": { - "n_finite": 10001, + "n_finite": 9967, "n_non_finite": 0, - "mean": -8.34374768827681, - "std": 1.8479154152116826, - "min": -21.193822220414155, - "max": -5.3048644505016 + "mean": -8.33931896402931, + "std": 1.8736677602653675, + "min": -22.300802363627707, + "max": -5.272577113539358 } }, "log10_pgen_1mm": { "train": { - "n_finite": 79998, + "n_finite": 79730, "n_non_finite": 0, - "mean": -6.52977416306943, - "std": 1.6300446772117814, + "mean": -6.526844812213512, + "std": 1.6264507836142756, "min": -25.268624790379345, "max": -3.967721318105829 }, "val": { - "n_finite": 9999, + "n_finite": 9966, "n_non_finite": 0, - "mean": -6.532824795313865, - "std": 1.6590462485672821, - "min": -20.17601369470008, - "max": -4.055332526062316 + "mean": -6.516624042597638, + "std": 1.611059243463444, + "min": -19.405392492212744, + "max": -3.999265344121248 }, "test": { - "n_finite": 10001, + "n_finite": 9967, "n_non_finite": 0, - "mean": -6.522161324076908, - "std": 1.5951257276034596, - "min": -18.745888024917427, - "max": -4.127677056652921 + "mean": -6.521180768534297, + "std": 1.6127152970829823, + "min": -19.90410715459764, + "max": -4.018798001189545 } } }, diff --git a/docs/anchored_tokenizer_study.md b/docs/anchored_tokenizer_study.md new file mode 100644 index 0000000..bc72e4c --- /dev/null +++ b/docs/anchored_tokenizer_study.md @@ -0,0 +1,70 @@ +# Leakage-safe sequence-only anchored tokenizer study + +## Formal mapping + +For a validated amino-acid CDR3 string `s`, each candidate implements + +`T(s) = [a_N(s), WP(m(s)), a_C(s)]`, + +where `a_N(s)` and `a_C(s)` are literal substrings of `s`, `m(s)` is exactly +the intervening substring, and `WP` is a WordPiece model fitted only to middle +spans from a benchmark-disjoint independently generated TRB corpus. Concatenating +the decoded token surfaces must reproduce `s` byte-for-byte. The encoder accepts +only `s`; neither `v_call` nor `j_call` is an inference input. The reconstruction +decoder remains the unchanged 25-token character decoder. + +All tokenizers reserve PAD/BOS/EOS/UNK IDs 0/1/2/3. Normal validated inputs never +emit UNK. Input tokens have exactly one of three boundary types: `N_ANCHOR`, +`MIDDLE_WORDPIECE`, or `C_ANCHOR`. + +## Candidates + +1. **EDGE-3 (`edge_k`)** protects the first and last three residues as one token + each. Its vocabulary exhaustively enumerates every 3-mer over the 20 + standard amino acids (8,000 tokens per side). This makes every valid literal + edge representable without observing evaluation sequences. K=3 is + the largest feasible non-overlapping choice because the locked benchmark's + recorded minimum CDR3 length is seven; K=4 would require eight residues and + failed the CPU coverage gate. Inputs shorter than six or containing invalid + amino acids fail closed. +2. **DATA-ANCHOR (`data_anchor`)** mines literal prefixes and suffixes of length + 2–8 from locked train rows only. A motif is retained when it has at least 50 + train occurrences and at least 25 occurrences within one annotated V (N side) + or J (C side) gene. V/J calls restrict the train-only library construction but + are discarded before encoding. All 20 one-residue identity fallbacks are + reserved constants (not learned motifs), guaranteeing sequence-only coverage + without held-out fitting. +3. **GERMLINE-ANCHOR (`germline_anchor`)** translates the functional V and J + CDR3-compatible segments in the documented OLGA human `TRB_orig` model using + its V/J anchor-index CSVs. Literal prefixes/suffixes of length 1–8 are retained + only if at least 20 locked train sequences support them. The exact source-file + hashes are recorded. All 20 one-residue identity fallbacks are reserved + constants; multi-residue germline motifs remain reference-derived and + train-supported. Encoding still operates solely by literal sequence match. + +For variable anchors, all matching prefix/suffix pairs that do not overlap are +considered. Selection maximizes total protected length, then N length, then C +length, then uses lexical order. Distinct same-length literal matches cannot both +match the same edge; collision counts and overlap-driven fallback counts are +nonetheless audited. If no identity fallback pair exists, encoding fails closed. + +## Corpus and leakage controls + +The preferred central corpus is a deterministic 500,000-row subset streamed from +the independently generated OLGA human TRB corpus +`/projects/immunestatus/vdjrearm/sample/TRB_1e7.tsv.gz`. Its generation script, +OLGA model files, and corpus are hashed. Every locked train, validation, and test +CDR3 is removed before fitting, duplicates and invalid rows are removed, and the +recorded final overlap must be zero. The external corpus influences central +WordPiece only; DATA and GERMLINE anchor retention uses locked train evidence +only. Validation and test strings are used solely for post-fit audits and metrics. + +## Matched experiment rule + +All three candidates use full R+T+P supervision, seed 42, d-model 320, latent 128, +four encoder and four decoder layers, batch 64, 40 epochs, AdamW at 3e-4, and the +locked train-only normalizer. Apart from input tokenizer path/type and the implied +input embedding vocabulary size/parameter count, configurations match the selected +character baseline. Selection uses validation metrics only; test metrics are +descriptive. Any annotation-aware V/J-input condition would be a separate modality +and is outside this study. diff --git a/docs/integration_provenance.md b/docs/integration_provenance.md new file mode 100644 index 0000000..800833c --- /dev/null +++ b/docs/integration_provenance.md @@ -0,0 +1,31 @@ +# Integration provenance + +Canonical destination: https://github.com/VEK239/irrm-codec + +The publication integration starts at Antigenomics main `31ccbad1c83a765fdb91a3686cbe31d4fcf6165c`. +Real merge commits preserve the original authors, commit IDs, and parent history: + +1. Maria Demidova's `demidovamaria/main` at `7b0e7f1b82c64490528163210d154fddcb97a290` (45 unique commits). +2. `Consoooomer/benchmark/issue-3-pgen` at `b96f06eb42d9c92e3982103bb903a1dc67d59138` (19 unique commits), including the splits and reconstruction branches. +3. VEK239 main at `aa2589d21428fd59fd34b55d675fc2cb62a6980b`. +4. Antigenomics `agent/multitask-tokenizer-evaluation` at `9abb43c99de96108e178b8bbb5a0ed5d5cf414c6`, authored by Elizaveta Vlasova. +5. Previously uncommitted local author source, test, Slurm, and experiment configuration files. + +Published history is not rewritten and student work is not squashed. The integration +retains both the published VEK239 main and student tips as ancestors. Contributor +identity follows the original Git commits; no author fields were reassigned. + +Resolution decisions: + +- Keep the working single-task training fixes and W&B integration from VEK239 main. +- Combine memory-mapped/clone-ID-aware embedding preparation with the student's V/J + standardization and sequence-keyed Pgen cache. +- Retain student padded/anchored WordPiece APIs together with validated unpadded + WordPiece input and the author's anchored tokenizer implementation. +- Preserve both known WordPiece special-token layouts without renumbering stored vocabularies. +- Retain author gradient-accumulation and resume fixes alongside the experimental + locked standardizer and local TCRemP geometry-loss options. +- Retain historical reports and result tables with their original provenance. + +The original GPL-3.0 license remains in LICENSE. Historical cluster paths and input +hashes in experiment manifests are provenance records and are not rewritten. diff --git a/docs/publication.md b/docs/publication.md new file mode 100644 index 0000000..fc5ff6d --- /dev/null +++ b/docs/publication.md @@ -0,0 +1,48 @@ +# Reproduction and compatibility + +## Data and results + +The repository contains code, tests, experiment configurations, and historical +student result tables. Large TCRemP embeddings, pretrained weights, checkpoints, +raw repertoires, and cluster output are external runtime artifacts. CPU tests use +synthetic fixtures; they do not regenerate the published research results. + +The author's joint preparation (`benchmark.prepare_trb_joint`) validates biological +identity using AIRR, Pgen, and the TCRemP clone-ID sidecar and writer contract. The +student preparation (`benchmark.prepare_splits`) retains its distinct benchmark +workflow and V/J cleaning. Do not interchange split manifests between cohorts. +Normalize targets and fit tokenizer/PCA models using the prescribed training data. + +## Token IDs and historical checkpoints + +The author's character and anchored-tokenizer convention is PAD=0, BOS=1, EOS=2, +UNK=3. Student WordPiece files and the earlier VEK239 main may use PAD=0, UNK=1, +BOS=2, EOS=3. The integrated WordPiece loader accepts both known layouts, preserves +all stored IDs, and checks UNK/EOS using each tokenizer's own vocabulary. The +reconstruction decoder uses the author's character convention. + +A character reconstruction checkpoint trained under the earlier VEK239 special-ID +order is not automatically interchangeable with the author's decoder convention. +Use its original revision or explicitly migrate special-token embedding/output +rows before evaluation. Do not relabel token IDs in existing tokenizer files. +Author DATA-ANCHOR configurations remain compatible with the author convention. + +## Cluster examples + +Slurm examples retain the partitions, accounts, data paths, and environments used +for their historical experiments. Inspect these before submission; #SBATCH paths +are literal and do not expand shell variables. Adapt them to your cluster and use +the Python module entrypoints with explicit CLI paths. Scripts that expose REPO, +SCRATCH, CONDA_ROOT, PYTHON_BIN, or IRRM_ROOT can be configured through those variables. +No cluster jobs are submitted by installing this repository or running its CPU tests. + +## Validation scope + +Run `python -m pytest -q` from the repository root. Tests cover reconstruction, +objective combinations, tokenizers, data identity/alignment, standardizers, metrics, +external benchmark bookkeeping, and VDJdb cohort rules. Optional external encoders +need requirements-benchmark.txt plus model downloads; complete GPU training and +external-data experiments require separately provisioned resources. + +Tab-separated result files preserve empty terminal fields. Trailing tabs in these +files represent missing values and are retained for table-schema fidelity. diff --git a/docs/vdjdb_epitope_validation.md b/docs/vdjdb_epitope_validation.md new file mode 100644 index 0000000..a182f1e --- /dev/null +++ b/docs/vdjdb_epitope_validation.md @@ -0,0 +1,59 @@ +# VDJdb epitope-separation validation gate + +This study asks whether frozen IRRM-CODEC latent vectors separate epitope/MHC +labels: within-label distances must be lower than matched between-label +distances. It is a quantitative held-out validation, not a visualization study. + +## Cohort rule (locked before embeddings) + +- Source: official VDJdb `2025-07-30` `vdjdb.txt`; the CPU preflight records its + byte size and SHA-256. +- Retain human TRB rows with canonical CDR3 amino acids (length at most 40), + `cdr3fix.good=true` without a sequence-changing fix, VDJdb score at least 1, + a peptide, MHC-A and MHC class, and a subject identifier. +- The label is the exact normalized tuple `(epitope, mhc.a, mhc.b, mhc.class)`. + This avoids pooling the same peptide across distinct MHC contexts. +- Donors are keyed by `(study-or-reference, subject.id)` so generic identifiers + such as `donor1` do not collide across studies. +- Deduplicate exact CDR3s. Remove CDR3s assigned to multiple labels, CDR3s + observed under multiple donors (for unambiguous donor-cluster inference), and + overlaps with any locked benchmark train, validation, or test split. +- Primary cohort: every label with at least 50 unique remaining CDR3s and at + least two donors. Report fixed threshold sensitivities at 25, 50, and 100. + +The CPU job writes the immutable candidate table, per-label counts, overlap and +filter counts, threshold sensitivity, hashes, and a finite two-sequence CPU +smoke encoding. It explicitly does not extract the full embedding matrix. + +## Model-comparison gate + +The completed DATA-ANCHOR run is R+T+P only. The required primary matrix is +matched DATA-ANCHOR R-only, P-only, R+P, and R+T+P encoders, with the same train +data, tokenizer bundle, seed, architecture, optimizer, batch size, 40-epoch +budget, and validation-only checkpoint selection. R+P versus R+T+P is the direct +test of TCRemP's incremental value conditional on reconstruction and pgen. +T-only and R+T are useful optional attribution conditions, but are not required +for this primary comparison. The current multi-task implementation always +constructs the shared encoder, decoder, TCRemP head, and pgen head. Therefore +zeroing unused loss weights yields checkpoint-compatible P-only encoders; its +decoder and TCRemP head remain present but intentionally untrained, and their +metrics are not interpretable. + +A separate, narrower auxiliary-head analysis uses R-only, R+T, R+P, and R+T+P. +The existing character-tokenizer quartet can answer that question for the char +condition, but cannot be mixed with DATA-ANCHOR R+T+P: tokenizer vocabulary, +input embeddings, latent width, and objective would be confounded. + +No embeddings or outcome metrics are produced until the cohort and matched +model-comparison matrix are accepted. + +## Future metrics (preregistered) + +Primary: per-label within/between cosine-distance ratio (lower is better), with +length-matched between-label pairs. Secondary: Euclidean contrast, same-label +pair AUROC/AUPRC, donor-excluded precision@1/5/10, and MAP@R. Controls are an +absolute length-difference score and 1,000 label permutations within length and +study/donor strata. Uncertainty uses 2,000 donor-cluster bootstrap replicates. +All compared checkpoints use the same cohort, pair manifest, and resamples. +UMAP is not a primary endpoint and no scaler or metric learner is fit on these +held-out records. diff --git a/experiments/trb/anchored-tokenizer-study/stage_manifest.json b/experiments/trb/anchored-tokenizer-study/stage_manifest.json new file mode 100644 index 0000000..9aedc0b --- /dev/null +++ b/experiments/trb/anchored-tokenizer-study/stage_manifest.json @@ -0,0 +1,73 @@ +{ + "study": "leakage-safe sequence-only TRB terminal-anchor tokenizers", + "status": "germline_submitted", + "base_git_revision": "3c873ff7b05e876a80cd226c8099fc1f8b282e91", + "uncommitted_remote_code_is_hash_locked": true, + "cpu_preflight_job_id": "1426217", + "cpu_preflight_sha256": "08d1d582842c2bdab0bafcd5d429f3b3042d02cb715b50e2638b5da652435a38", + "benchmark": { + "path": "/home/evlasova/irrm-codec/data/benchmark/trb", + "ready_sha256": "ce08945c0c0a975c2b05129695c26399456f185fe79b89914c0ad52e4fdcf167", + "target_standardizer_sha256": "ed0052ad510b7e0299628796162037cbb1a3c7e5e8d661ca4a668eb49a830b24", + "rows": 99430, + "train_rows": 79544, + "val_rows": 9943, + "test_rows": 9943 + }, + "matched_training": { + "inference_modality": "sequence_only", + "decoder_vocab_size": 25, + "objectives": ["reconstruction", "tcremp", "log10_pgen_1mm"], + "loss_weights": {"reconstruction": 1.0, "tcremp": 1.0, "pgen": 1.0}, + "seed": 42, + "d_model": 320, + "latent_dim": 128, + "encoder_layers": 4, + "decoder_layers": 4, + "nhead": 8, + "ff_dim": 1280, + "batch_size": 64, + "epochs": 40, + "optimizer": "AdamW", + "learning_rate": 0.0003, + "weight_decay": 0.0001, + "normalizer_fit": "locked train only", + "selection": "validation joint objective only" + }, + "gpu_batch_script_sha256": "dcf3ef80f59e42ccde5b48cb97925603ec6062000b2b2cde12b378268effd452", + "runs": { + "edge_k": { + "condition": "EDGE-3", + "bundle_sha256": "8333622c01ad10bd3d176f71131705ad68feb0a1eed5039ed6ee352081888a24", + "input_vocab_size": 16256, + "parameter_count": 26344873, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/anchor-edge3-rtp-seed42", + "job_id": "1426221", + "terminal": true, + "best_epoch": 31, + "best_validation_joint_loss": 0.2154474285803342, + "test_loss_descriptive": 0.2144757355439235 + }, + "data_anchor": { + "condition": "DATA-ANCHOR", + "bundle_sha256": "8ef9aeddf8c069d4db898369f01f97e80160ba2ceb3c2354f18547467beeed4d", + "input_vocab_size": 1228, + "parameter_count": 21535913, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/anchor-data-rtp-seed42", + "job_id": "1426222", + "terminal": true, + "best_epoch": 25, + "best_validation_joint_loss": 0.20576476722275314, + "test_loss_descriptive": 0.20348657863661399 + }, + "germline_anchor": { + "condition": "GERMLINE-ANCHOR", + "bundle_sha256": "6e48150050325d916a2d09b3a7a0db4d74ad7233319dd49e01ba904d93d435b1", + "input_vocab_size": 387, + "parameter_count": 21266793, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/anchor-germline-rtp-seed42", + "job_id": "1426317", + "submission_gate": "satisfied: EDGE-3 and DATA-ANCHOR terminal before submission" + } + } +} diff --git a/experiments/trb/phase1-rq1-char-seed42/experiment_manifest.json b/experiments/trb/phase1-rq1-char-seed42/experiment_manifest.json new file mode 100644 index 0000000..3377a8e --- /dev/null +++ b/experiments/trb/phase1-rq1-char-seed42/experiment_manifest.json @@ -0,0 +1,35 @@ +{ + "phase": "Phase 1 / RQ1", + "tokenizer": "char", + "seed": 42, + "gpu_concurrency_cap": 2, + "submission_rule": "Exactly the two registered baselines; no further GPU jobs until both are terminal and summarized.", + "base_git_revision": "3c873ff7b05e876a80cd226c8099fc1f8b282e91", + "saved_pr_revision": "330dfdfb48bd867a195599043405e4d89999673e", + "benchmark_path": "/home/evlasova/irrm-codec/data/benchmark/trb", + "benchmark_preflight": "/home/evlasova/irrm-codec/experiments/trb/phase1-rq1-char-seed42/preflight.json", + "training_code_sha256": { + "irrm_codec/tokenization.py": "649d157fa9c87f304a314f5009fed304d511562fa0f56df57dc4c0f06aa11936", + "irrm_codec/multitask_data.py": "5c8f1b1f1e21d2044028291d7e49e54090cce475d5e124cc3798ff2876b5f52f", + "irrm_codec/multitask_losses.py": "1d46e24c335fee02ac7344c6278de5b112d9f3bc57915a25df31db444d361649", + "irrm_codec/multitask_metrics.py": "891e52658e9470f8d213f517125670fbedd752efeb636493ca2a41a99e309e64", + "irrm_codec/multitask_transformer.py": "deab3ff9edc851a37683f1c2cd314e262c6b09171d407db7621962f58ed67cfa", + "irrm_codec/train_multitask.py": "0eae9a7d4dd233e52c386a4e086d262b6236d4fb8505550ea77898f5d30ce4de" + }, + "experiments": [ + { + "experiment_id": "rq1-char-r-seed42", + "condition": "R", + "loss_weights": {"reconstruction": 1.0, "tcremp": 0.0, "pgen": 0.0}, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq1-char-r-seed42", + "job_id": "1425896" + }, + { + "experiment_id": "rq1-char-rtp-seed42", + "condition": "R+T+P", + "loss_weights": {"reconstruction": 1.0, "tcremp": 1.0, "pgen": 1.0}, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq1-char-rtp-seed42", + "job_id": "1425897" + } + ] +} diff --git a/experiments/trb/phase2-rq23-char-seed42/experiment_manifest.json b/experiments/trb/phase2-rq23-char-seed42/experiment_manifest.json new file mode 100644 index 0000000..730e96b --- /dev/null +++ b/experiments/trb/phase2-rq23-char-seed42/experiment_manifest.json @@ -0,0 +1,36 @@ +{ + "phase": "Phase 2 / RQ2-RQ3", + "tokenizer": "char", + "seed": 42, + "gpu_concurrency_cap": 2, + "submission_rule": "Exactly the two registered head ablations; do not launch further variants.", + "base_git_revision": "3c873ff7b05e876a80cd226c8099fc1f8b282e91", + "saved_pr_revision": "330dfdfb48bd867a195599043405e4d89999673e", + "benchmark_path": "/home/evlasova/irrm-codec/data/benchmark/trb", + "reused_benchmark_preflight": "/home/evlasova/irrm-codec/experiments/trb/phase1-rq1-char-seed42/preflight.json", + "benchmark_ready_sha256": "ce08945c0c0a975c2b05129695c26399456f185fe79b89914c0ad52e4fdcf167", + "training_code_sha256": { + "irrm_codec/tokenization.py": "649d157fa9c87f304a314f5009fed304d511562fa0f56df57dc4c0f06aa11936", + "irrm_codec/multitask_data.py": "5c8f1b1f1e21d2044028291d7e49e54090cce475d5e124cc3798ff2876b5f52f", + "irrm_codec/multitask_losses.py": "1d46e24c335fee02ac7344c6278de5b112d9f3bc57915a25df31db444d361649", + "irrm_codec/multitask_metrics.py": "891e52658e9470f8d213f517125670fbedd752efeb636493ca2a41a99e309e64", + "irrm_codec/multitask_transformer.py": "deab3ff9edc851a37683f1c2cd314e262c6b09171d407db7621962f58ed67cfa", + "irrm_codec/train_multitask.py": "0eae9a7d4dd233e52c386a4e086d262b6236d4fb8505550ea77898f5d30ce4de" + }, + "experiments": [ + { + "experiment_id": "rq2-char-rt-seed42", + "condition": "R+T", + "loss_weights": {"reconstruction": 1.0, "tcremp": 1.0, "pgen": 0.0}, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq2-char-rt-seed42", + "job_id": "1425901" + }, + { + "experiment_id": "rq2-char-rp-seed42", + "condition": "R+P", + "loss_weights": {"reconstruction": 1.0, "tcremp": 0.0, "pgen": 1.0}, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq2-char-rp-seed42", + "job_id": "1425902" + } + ] +} diff --git a/experiments/trb/rq1-char-r-seed42/experiment_config.json b/experiments/trb/rq1-char-r-seed42/experiment_config.json new file mode 100644 index 0000000..9176424 --- /dev/null +++ b/experiments/trb/rq1-char-r-seed42/experiment_config.json @@ -0,0 +1,56 @@ +{ + "experiment_id": "rq1-char-r-seed42", + "research_question": "Phase 1 / RQ1 character-tokenizer matched baseline", + "condition": "reconstruction_only_R", + "base_git_revision": "3c873ff7b05e876a80cd226c8099fc1f8b282e91", + "saved_pr_revision": "330dfdfb48bd867a195599043405e4d89999673e", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq1-char-r-seed42", + "tokenizer": "char", + "seed": 42, + "train_subset": "all", + "pgen_target": "log10_pgen_1mm", + "model": { + "max_sequence_len": 40, + "d_model": 320, + "latent_dim": 320, + "nhead": 8, + "encoder_layers": 4, + "decoder_layers": 4, + "ff_dim": 1280, + "dropout": 0.1, + "tcremp_head_dim": 1024, + "pgen_head_dim": 256, + "decoder_memory_tokens": 4 + }, + "loss": { + "tcremp": 0.0, + "pgen": 0.0, + "reconstruction": 1.0, + "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, + "label_smoothing": 0.0 + }, + "training": { + "batch_size": 64, + "epochs": 40, + "learning_rate": 0.0003, + "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, + "max_grad_norm": 1.0, + "early_stopping_patience": 0, + "scheduler_factor": 0.5, + "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, + "amp": true + }, + "slurm": { + "partition": "medium", + "constraint": "gpu", + "gpus": 1, + "cpus": 8, + "memory": "48G", + "time": "12:00:00", + "job_id": "1425896" + } +} diff --git a/experiments/trb/rq1-char-rtp-seed42/experiment_config.json b/experiments/trb/rq1-char-rtp-seed42/experiment_config.json new file mode 100644 index 0000000..9bba372 --- /dev/null +++ b/experiments/trb/rq1-char-rtp-seed42/experiment_config.json @@ -0,0 +1,56 @@ +{ + "experiment_id": "rq1-char-rtp-seed42", + "research_question": "Phase 1 / RQ1 character-tokenizer matched baseline", + "condition": "full_multitask_R_plus_T_plus_P", + "base_git_revision": "3c873ff7b05e876a80cd226c8099fc1f8b282e91", + "saved_pr_revision": "330dfdfb48bd867a195599043405e4d89999673e", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq1-char-rtp-seed42", + "tokenizer": "char", + "seed": 42, + "train_subset": "all", + "pgen_target": "log10_pgen_1mm", + "model": { + "max_sequence_len": 40, + "d_model": 320, + "latent_dim": 320, + "nhead": 8, + "encoder_layers": 4, + "decoder_layers": 4, + "ff_dim": 1280, + "dropout": 0.1, + "tcremp_head_dim": 1024, + "pgen_head_dim": 256, + "decoder_memory_tokens": 4 + }, + "loss": { + "tcremp": 1.0, + "pgen": 1.0, + "reconstruction": 1.0, + "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, + "label_smoothing": 0.0 + }, + "training": { + "batch_size": 64, + "epochs": 40, + "learning_rate": 0.0003, + "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, + "max_grad_norm": 1.0, + "early_stopping_patience": 0, + "scheduler_factor": 0.5, + "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, + "amp": true + }, + "slurm": { + "partition": "medium", + "constraint": "gpu", + "gpus": 1, + "cpus": 8, + "memory": "48G", + "time": "12:00:00", + "job_id": "1425897" + } +} diff --git a/experiments/trb/rq2-char-rp-seed42/experiment_config.json b/experiments/trb/rq2-char-rp-seed42/experiment_config.json new file mode 100644 index 0000000..dcc9475 --- /dev/null +++ b/experiments/trb/rq2-char-rp-seed42/experiment_config.json @@ -0,0 +1,56 @@ +{ + "experiment_id": "rq2-char-rp-seed42", + "research_question": "Phase 2 / RQ2-RQ3 character-tokenizer head ablation", + "condition": "reconstruction_plus_pgen_R_plus_P", + "base_git_revision": "3c873ff7b05e876a80cd226c8099fc1f8b282e91", + "saved_pr_revision": "330dfdfb48bd867a195599043405e4d89999673e", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq2-char-rp-seed42", + "tokenizer": "char", + "seed": 42, + "train_subset": "all", + "pgen_target": "log10_pgen_1mm", + "model": { + "max_sequence_len": 40, + "d_model": 320, + "latent_dim": 320, + "nhead": 8, + "encoder_layers": 4, + "decoder_layers": 4, + "ff_dim": 1280, + "dropout": 0.1, + "tcremp_head_dim": 1024, + "pgen_head_dim": 256, + "decoder_memory_tokens": 4 + }, + "loss": { + "tcremp": 0.0, + "pgen": 1.0, + "reconstruction": 1.0, + "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, + "label_smoothing": 0.0 + }, + "training": { + "batch_size": 64, + "epochs": 40, + "learning_rate": 0.0003, + "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, + "max_grad_norm": 1.0, + "early_stopping_patience": 0, + "scheduler_factor": 0.5, + "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, + "amp": true + }, + "slurm": { + "partition": "medium", + "constraint": "gpu", + "gpus": 1, + "cpus": 8, + "memory": "48G", + "time": "12:00:00", + "job_id": "1425902" + } +} diff --git a/experiments/trb/rq2-char-rt-seed42/experiment_config.json b/experiments/trb/rq2-char-rt-seed42/experiment_config.json new file mode 100644 index 0000000..8fea052 --- /dev/null +++ b/experiments/trb/rq2-char-rt-seed42/experiment_config.json @@ -0,0 +1,56 @@ +{ + "experiment_id": "rq2-char-rt-seed42", + "research_question": "Phase 2 / RQ2-RQ3 character-tokenizer head ablation", + "condition": "reconstruction_plus_tcremp_R_plus_T", + "base_git_revision": "3c873ff7b05e876a80cd226c8099fc1f8b282e91", + "saved_pr_revision": "330dfdfb48bd867a195599043405e4d89999673e", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq2-char-rt-seed42", + "tokenizer": "char", + "seed": 42, + "train_subset": "all", + "pgen_target": "log10_pgen_1mm", + "model": { + "max_sequence_len": 40, + "d_model": 320, + "latent_dim": 320, + "nhead": 8, + "encoder_layers": 4, + "decoder_layers": 4, + "ff_dim": 1280, + "dropout": 0.1, + "tcremp_head_dim": 1024, + "pgen_head_dim": 256, + "decoder_memory_tokens": 4 + }, + "loss": { + "tcremp": 1.0, + "pgen": 0.0, + "reconstruction": 1.0, + "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, + "label_smoothing": 0.0 + }, + "training": { + "batch_size": 64, + "epochs": 40, + "learning_rate": 0.0003, + "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, + "max_grad_norm": 1.0, + "early_stopping_patience": 0, + "scheduler_factor": 0.5, + "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, + "amp": true + }, + "slurm": { + "partition": "medium", + "constraint": "gpu", + "gpus": 1, + "cpus": 8, + "memory": "48G", + "time": "12:00:00", + "job_id": "1425901" + } +} diff --git a/experiments/trb/rq4-wordpiece-r-seed42/experiment_config.json b/experiments/trb/rq4-wordpiece-r-seed42/experiment_config.json new file mode 100644 index 0000000..e006c7a --- /dev/null +++ b/experiments/trb/rq4-wordpiece-r-seed42/experiment_config.json @@ -0,0 +1,33 @@ +{ + "experiment_id": "rq4-wordpiece-r-seed42", + "research_question": "RQ4 WordPiece tokenizer study matched character baseline", + "condition": "wordpiece_reconstruction_only_R", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq4-wordpiece-r-seed42", + "tokenizer": { + "type": "wordpiece", + "path": "/home/evlasova/irrm-codec/experiments/trb/tokenizers/wordpiece-seed42/tokenizer.json", + "sha256": "064c1b6a60fec3a14913e821954c7eaf1a70b2335b77ef79fe4e2ed4fac62634", + "vocab_size": 256, + "train_manifest_sha256": "6399cbea503191a652e7ecef0b8b859793f4ece16322fd49ce2ed03364484b01" + }, + "seed": 42, + "model": { + "max_sequence_len": 40, "d_model": 320, "latent_dim": 320, "nhead": 8, + "encoder_layers": 4, "decoder_layers": 4, "ff_dim": 1280, "dropout": 0.1, + "tcremp_head_dim": 1024, "pgen_head_dim": 256, "decoder_memory_tokens": 4, + "input_vocab_size": 256, "output_vocab_size": 25, + "share_input_output_embeddings": false, + "expected_parameter_count": 21779177, + "embedding_adjustment": "WordPiece requires a separate 256x320 encoder embedding; the 25x320 character decoder embedding remains unchanged. This adds 81,920 parameters versus the shared character baseline." + }, + "loss": {"reconstruction": 1.0, "tcremp": 0.0, "pgen": 0.0}, + "training": { + "batch_size": 64, "epochs": 40, "learning_rate": 0.0003, + "weight_decay": 0.0001, "gradient_accumulation_steps": 1, + "max_grad_norm": 1.0, "early_stopping_patience": 0, + "scheduler_factor": 0.5, "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, "amp": true + }, + "slurm": {"partition": "medium", "gpus": 1, "cpus": 8, "memory": "48G", "time": "12:00:00", "job_id": "1425944"} +} diff --git a/experiments/trb/rq4-wordpiece-rtp-seed42/experiment_config.json b/experiments/trb/rq4-wordpiece-rtp-seed42/experiment_config.json new file mode 100644 index 0000000..6caa21f --- /dev/null +++ b/experiments/trb/rq4-wordpiece-rtp-seed42/experiment_config.json @@ -0,0 +1,33 @@ +{ + "experiment_id": "rq4-wordpiece-rtp-seed42", + "research_question": "RQ4 WordPiece tokenizer study matched character baseline", + "condition": "wordpiece_full_R_plus_T_plus_P", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq4-wordpiece-rtp-seed42", + "tokenizer": { + "type": "wordpiece", + "path": "/home/evlasova/irrm-codec/experiments/trb/tokenizers/wordpiece-seed42/tokenizer.json", + "sha256": "064c1b6a60fec3a14913e821954c7eaf1a70b2335b77ef79fe4e2ed4fac62634", + "vocab_size": 256, + "train_manifest_sha256": "6399cbea503191a652e7ecef0b8b859793f4ece16322fd49ce2ed03364484b01" + }, + "seed": 42, + "model": { + "max_sequence_len": 40, "d_model": 320, "latent_dim": 320, "nhead": 8, + "encoder_layers": 4, "decoder_layers": 4, "ff_dim": 1280, "dropout": 0.1, + "tcremp_head_dim": 1024, "pgen_head_dim": 256, "decoder_memory_tokens": 4, + "input_vocab_size": 256, "output_vocab_size": 25, + "share_input_output_embeddings": false, + "expected_parameter_count": 21779177, + "embedding_adjustment": "WordPiece requires a separate 256x320 encoder embedding; the 25x320 character decoder embedding remains unchanged. This adds 81,920 parameters versus the shared character baseline." + }, + "loss": {"reconstruction": 1.0, "tcremp": 1.0, "pgen": 1.0}, + "training": { + "batch_size": 64, "epochs": 40, "learning_rate": 0.0003, + "weight_decay": 0.0001, "gradient_accumulation_steps": 1, + "max_grad_norm": 1.0, "early_stopping_patience": 0, + "scheduler_factor": 0.5, "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, "amp": true + }, + "slurm": {"partition": "medium", "gpus": 1, "cpus": 8, "memory": "48G", "time": "12:00:00", "job_id": "1425945"} +} diff --git a/experiments/trb/rq4-wordpiece-seed42/experiment_manifest.json b/experiments/trb/rq4-wordpiece-seed42/experiment_manifest.json new file mode 100644 index 0000000..9c3745b --- /dev/null +++ b/experiments/trb/rq4-wordpiece-seed42/experiment_manifest.json @@ -0,0 +1,25 @@ +{ + "phase": "RQ4 tokenizer study", + "tokenizer_dir": "/home/evlasova/irrm-codec/experiments/trb/tokenizers/wordpiece-seed42", + "tokenizer_sha256": "064c1b6a60fec3a14913e821954c7eaf1a70b2335b77ef79fe4e2ed4fac62634", + "train_manifest_sha256": "6399cbea503191a652e7ecef0b8b859793f4ece16322fd49ce2ed03364484b01", + "preflight_job_id": "1425943", + "preflight_status": "ready", + "gpu_concurrency_cap": 2, + "submission_rule": "Exactly the registered WordPiece R and R+T+P baselines; no other GPU variants.", + "seed": 42, + "epochs": 40, + "architecture_adjustment": { + "reason": "WordPiece input vocabulary differs from the character reconstruction vocabulary.", + "input_vocab_size": 256, + "output_vocab_size": 25, + "shared_embeddings": false, + "character_baseline_parameters": 21697257, + "wordpiece_parameters": 21779177, + "parameter_increase": 81920 + }, + "experiments": [ + {"experiment_id": "rq4-wordpiece-r-seed42", "weights_r_t_p": [1, 0, 0], "job_id": "1425944"}, + {"experiment_id": "rq4-wordpiece-rtp-seed42", "weights_r_t_p": [1, 1, 1], "job_id": "1425945"} + ] +} diff --git a/experiments/trb/rq5-char-rtp-d128-depths-seed42/experiment_manifest.json b/experiments/trb/rq5-char-rtp-d128-depths-seed42/experiment_manifest.json new file mode 100644 index 0000000..d4cbd28 --- /dev/null +++ b/experiments/trb/rq5-char-rtp-d128-depths-seed42/experiment_manifest.json @@ -0,0 +1,52 @@ +{ + "phase": "RQ5 Stage 2 encoder-depth sweep", + "tokenizer": "char", + "supervision": "R+T+P", + "latent_dim": 128, + "decoder_layers": 4, + "seed": 42, + "epochs": 40, + "gpu_concurrency_cap": 2, + "submission_status": "completed", + "submission_rule": "Only encoder depths 2 and 6 after successful CPU preflight and explicit continuation.", + "selection": { + "criterion": "minimum best-checkpoint validation loss only", + "validation_losses": { + "d128": 0.19693734171098476, + "d320": 0.20322377112032147, + "d512": 0.2087959912276347 + }, + "selected_latent_dim": 128, + "test_metrics_used_for_selection": false, + "selected_baseline_job_id": "1425949" + }, + "benchmark_path": "/home/evlasova/irrm-codec/data/benchmark/trb", + "benchmark_ready_sha256": "ce08945c0c0a975c2b05129695c26399456f185fe79b89914c0ad52e4fdcf167", + "preflight": { + "job_id": "1425954", + "status": "ready", + "report_path": "/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-depths-seed42/preflight.json", + "runtime": { + "python": "3.11.15", + "torch": "2.4.1+cu121" + } + }, + "candidates": [ + { + "encoder_layers": 2, + "decoder_layers": 4, + "parameter_count": 18677033, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-e2d4-seed42", + "job_id": "1425955", + "status": "completed" + }, + { + "encoder_layers": 6, + "decoder_layers": 4, + "parameter_count": 23608873, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-e6d4-seed42", + "job_id": "1425956", + "status": "completed" + } + ] +} diff --git a/experiments/trb/rq5-char-rtp-d128-e2d4-seed42/experiment_config.json b/experiments/trb/rq5-char-rtp-d128-e2d4-seed42/experiment_config.json new file mode 100644 index 0000000..878cbfb --- /dev/null +++ b/experiments/trb/rq5-char-rtp-d128-e2d4-seed42/experiment_config.json @@ -0,0 +1,59 @@ +{ + "experiment_id": "rq5-char-rtp-d128-e2d4-seed42", + "research_question": "RQ5 Stage 2 encoder-depth sweep at validation-selected latent_dim=128", + "condition": "char_R_plus_T_plus_P_latent_d128_encoder2_decoder4", + "selection_basis": "latent_dim=128 selected solely by minimum Stage-1 validation loss", + "baseline_job_id": "1425949", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-e2d4-seed42", + "tokenizer": "char", + "seed": 42, + "model": { + "input_vocab_size": 25, + "output_vocab_size": 25, + "share_input_output_embeddings": true, + "max_sequence_len": 40, + "d_model": 320, + "latent_dim": 128, + "nhead": 8, + "encoder_layers": 2, + "decoder_layers": 4, + "ff_dim": 1280, + "dropout": 0.1, + "tcremp_dim": 9000, + "tcremp_head_dim": 1024, + "pgen_head_dim": 256, + "decoder_memory_tokens": 4, + "parameter_count": 18677033 + }, + "loss": { + "reconstruction": 1, + "tcremp": 1, + "pgen": 1, + "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, + "label_smoothing": 0 + }, + "training": { + "batch_size": 64, + "epochs": 40, + "learning_rate": 0.0003, + "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, + "max_grad_norm": 1, + "early_stopping_patience": 0, + "scheduler_factor": 0.5, + "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, + "amp": true + }, + "slurm": { + "partition": "medium", + "gpus": 1, + "cpus": 8, + "memory": "48G", + "time": "12:00:00", + "status": "completed", + "job_id": "1425955" + } +} diff --git a/experiments/trb/rq5-char-rtp-d128-e6d4-seed42/experiment_config.json b/experiments/trb/rq5-char-rtp-d128-e6d4-seed42/experiment_config.json new file mode 100644 index 0000000..0fd1126 --- /dev/null +++ b/experiments/trb/rq5-char-rtp-d128-e6d4-seed42/experiment_config.json @@ -0,0 +1,59 @@ +{ + "experiment_id": "rq5-char-rtp-d128-e6d4-seed42", + "research_question": "RQ5 Stage 2 encoder-depth sweep at validation-selected latent_dim=128", + "condition": "char_R_plus_T_plus_P_latent_d128_encoder6_decoder4", + "selection_basis": "latent_dim=128 selected solely by minimum Stage-1 validation loss", + "baseline_job_id": "1425949", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-e6d4-seed42", + "tokenizer": "char", + "seed": 42, + "model": { + "input_vocab_size": 25, + "output_vocab_size": 25, + "share_input_output_embeddings": true, + "max_sequence_len": 40, + "d_model": 320, + "latent_dim": 128, + "nhead": 8, + "encoder_layers": 6, + "decoder_layers": 4, + "ff_dim": 1280, + "dropout": 0.1, + "tcremp_dim": 9000, + "tcremp_head_dim": 1024, + "pgen_head_dim": 256, + "decoder_memory_tokens": 4, + "parameter_count": 23608873 + }, + "loss": { + "reconstruction": 1, + "tcremp": 1, + "pgen": 1, + "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, + "label_smoothing": 0 + }, + "training": { + "batch_size": 64, + "epochs": 40, + "learning_rate": 0.0003, + "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, + "max_grad_norm": 1, + "early_stopping_patience": 0, + "scheduler_factor": 0.5, + "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, + "amp": true + }, + "slurm": { + "partition": "medium", + "gpus": 1, + "cpus": 8, + "memory": "48G", + "time": "12:00:00", + "status": "completed", + "job_id": "1425956" + } +} diff --git a/experiments/trb/rq5-char-rtp-d128-seed42/experiment_config.json b/experiments/trb/rq5-char-rtp-d128-seed42/experiment_config.json new file mode 100644 index 0000000..54bb4c4 --- /dev/null +++ b/experiments/trb/rq5-char-rtp-d128-seed42/experiment_config.json @@ -0,0 +1,57 @@ +{ + "experiment_id": "rq5-char-rtp-d128-seed42", + "research_question": "RQ5 Stage 1 character full-multitask latent bottleneck width sweep", + "condition": "char_R_plus_T_plus_P_latent_d128", + "baseline_job_id": "1425897", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-seed42", + "tokenizer": "char", + "seed": 42, + "model": { + "input_vocab_size": 25, + "output_vocab_size": 25, + "share_input_output_embeddings": true, + "max_sequence_len": 40, + "d_model": 320, + "latent_dim": 128, + "nhead": 8, + "encoder_layers": 4, + "decoder_layers": 4, + "ff_dim": 1280, + "dropout": 0.1, + "tcremp_dim": 9000, + "tcremp_head_dim": 1024, + "pgen_head_dim": 256, + "decoder_memory_tokens": 4, + "parameter_count": 21142953 + }, + "loss": { + "reconstruction": 1, + "tcremp": 1, + "pgen": 1, + "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, + "label_smoothing": 0 + }, + "training": { + "batch_size": 64, + "epochs": 40, + "learning_rate": 0.0003, + "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, + "max_grad_norm": 1, + "early_stopping_patience": 0, + "scheduler_factor": 0.5, + "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, + "amp": true + }, + "slurm": { + "partition": "medium", + "gpus": 1, + "cpus": 8, + "memory": "48G", + "time": "12:00:00", + "job_id": "1425949" + } +} diff --git a/experiments/trb/rq5-char-rtp-d512-seed42/experiment_config.json b/experiments/trb/rq5-char-rtp-d512-seed42/experiment_config.json new file mode 100644 index 0000000..f34fe2e --- /dev/null +++ b/experiments/trb/rq5-char-rtp-d512-seed42/experiment_config.json @@ -0,0 +1,57 @@ +{ + "experiment_id": "rq5-char-rtp-d512-seed42", + "research_question": "RQ5 Stage 1 character full-multitask latent bottleneck width sweep", + "condition": "char_R_plus_T_plus_P_latent_d512", + "baseline_job_id": "1425897", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d512-seed42", + "tokenizer": "char", + "seed": 42, + "model": { + "input_vocab_size": 25, + "output_vocab_size": 25, + "share_input_output_embeddings": true, + "max_sequence_len": 40, + "d_model": 320, + "latent_dim": 512, + "nhead": 8, + "encoder_layers": 4, + "decoder_layers": 4, + "ff_dim": 1280, + "dropout": 0.1, + "tcremp_dim": 9000, + "tcremp_head_dim": 1024, + "pgen_head_dim": 256, + "decoder_memory_tokens": 4, + "parameter_count": 22251561 + }, + "loss": { + "reconstruction": 1, + "tcremp": 1, + "pgen": 1, + "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, + "label_smoothing": 0 + }, + "training": { + "batch_size": 64, + "epochs": 40, + "learning_rate": 0.0003, + "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, + "max_grad_norm": 1, + "early_stopping_patience": 0, + "scheduler_factor": 0.5, + "scheduler_patience": 2, + "scheduler_min_lr": 0.000001, + "amp": true + }, + "slurm": { + "partition": "medium", + "gpus": 1, + "cpus": 8, + "memory": "48G", + "time": "12:00:00", + "job_id": "1425950" + } +} diff --git a/experiments/trb/rq5-char-rtp-widths-seed42/experiment_manifest.json b/experiments/trb/rq5-char-rtp-widths-seed42/experiment_manifest.json new file mode 100644 index 0000000..17d0f79 --- /dev/null +++ b/experiments/trb/rq5-char-rtp-widths-seed42/experiment_manifest.json @@ -0,0 +1,34 @@ +{ + "phase": "RQ5 Stage 1 bottleneck-width sweep", + "tokenizer": "char", + "supervision": "R+T+P", + "seed": 42, + "epochs": 40, + "gpu_concurrency_cap": 2, + "submission_rule": "Exactly d128 and d512; d320 is reused from completed job 1425897.", + "benchmark_path": "/home/evlasova/irrm-codec/data/benchmark/trb", + "benchmark_ready_sha256": "ce08945c0c0a975c2b05129695c26399456f185fe79b89914c0ad52e4fdcf167", + "preflight_job_id": "1425948", + "preflight_status": "ready", + "width_definition": "latent_dim only; d_model=320 and nhead=8 fixed", + "baselines": [ + { + "latent_dim": 128, + "parameter_count": 21142953, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-seed42", + "job_id": "1425949" + }, + { + "latent_dim": 320, + "parameter_count": 21697257, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq1-char-rtp-seed42", + "job_id": "1425897" + }, + { + "latent_dim": 512, + "parameter_count": 22251561, + "output_dir": "/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d512-seed42", + "job_id": "1425950" + } + ] +} diff --git a/experiments/trb/wordpiece-sweep-seed42/experiment_manifest.json b/experiments/trb/wordpiece-sweep-seed42/experiment_manifest.json new file mode 100644 index 0000000..60da434 --- /dev/null +++ b/experiments/trb/wordpiece-sweep-seed42/experiment_manifest.json @@ -0,0 +1,159 @@ +{ + "phase": "WordPiece vocabulary-size sweep Stage A", + "git_revision": "3c873ff7b05e876a80cd226c8099fc1f8b282e91", + "implementation_state": "uncommitted focused sweep files preserved in the existing dirty worktree", + "implementation_sha256": { + "scripts/prepare_wordpiece_sweep.py": "3695f7a58aaf735bf56e6de960ce0eb11193cd4a929f77b1ea3e54155b4ecec5", + "scripts/preflight_wordpiece_sweep.py": "379f31dd2bccf1f51b99f148ca15aea86d2618a5be816eeae208bdaebb9a67ec", + "slurm/prepare_wordpiece_sweep.sbatch": "ace4cd922c31de4ab5671774f947e1972ed5c8010757b1017a9623c6727510fd", + "slurm/train_wordpiece_sweep.sbatch": "68171f1923559a94cbe6c8e3782fff2abf65bc519ad281fa5753cdd0e4c3699f" + }, + "tokenizer_type": "wordpiece", + "vocab_sizes": [44, 64, 128, 256, 512, 1024, 2048, 4096], + "tokenizer_root": "/home/evlasova/irrm-codec/experiments/trb/tokenizers/wordpiece-sweep-seed42", + "fit_split": "locked train manifest only", + "fit_rows": 79544, + "special_token_ids": {"pad": 0, "bos": 1, "eos": 2, "unk": 3}, + "selected_char_baseline": { + "job_id": "1425949", + "experiment_id": "rq5-char-rtp-d128-seed42", + "selection_basis": "minimum validation objective", + "latent_dim": 128, + "encoder_layers": 4, + "decoder_layers": 4 + }, + "model_invariants": { + "character_decoder_vocab_size": 25, + "separate_encoder_decoder_embeddings": true, + "full_rtp_weights": [1, 1, 1], + "seed": 42, + "epochs": 40, + "batch_size": 64, + "optimizer": "AdamW", + "learning_rate": 0.0003, + "locked_train_only_normalizer": true + }, + "selection_rule": { + "primary": "minimum full validation objective", + "secondary": ["validation pgen RMSE", "validation reconstruction", "validation TCRemP"], + "test_metrics_used_for_selection": false + }, + "gpu_concurrency_cap": 2, + "submission_order": [[44, 64], [128, 256], [512, 1024], [2048, 4096]], + "cpu_gate": { + "job_id": "1426071", + "status": "accepted", + "benchmark_ready_sha256": "ce08945c0c0a975c2b05129695c26399456f185fe79b89914c0ad52e4fdcf167", + "train_manifest_sha256": "6399cbea503191a652e7ecef0b8b859793f4ece16322fd49ce2ed03364484b01", + "train_corpus_sha256": "2478f48a82470e66ec78b0ec9fa06841868c6eed4697d387735f275a9c65597c" + }, + "completed_conditions": { + "44": { + "job_id": "1426073", + "best_epoch": 30, + "best_validation_loss": 0.19666322000361444, + "validation_pgen_rmse": 0.3436580250763355, + "validation_reconstruction_loss": 0.004775686473650997, + "validation_tcremp_cosine": 0.9922292748291728, + "walltime_seconds_from_log": 2890.775, + "artifacts_verified": true, + "error_scan_clean": true, + "best_checkpoint_sha256": "5942370dbc2bbf36c3ad059077517c1c0cc41028c002ca4615e5e2ceda856f56" + }, + "64": { + "job_id": "1426074", + "best_epoch": 28, + "best_validation_loss": 0.20685211541870166, + "validation_pgen_rmse": 0.3242212150811544, + "validation_reconstruction_loss": 0.007868499768547392, + "validation_tcremp_cosine": 0.9918381557562432, + "walltime_seconds_from_log": 3059.946, + "artifacts_verified": true, + "error_scan_clean": true, + "best_checkpoint_sha256": "36aa6aefd1745e0c11a2453829dc585ca09cc77a15680c4d05b70deb0b2dacef" + }, + "128": { + "job_id": "1426122", + "best_epoch": 26, + "best_validation_loss": 0.20901615753579544, + "validation_pgen_rmse": 0.2838516566076911, + "validation_reconstruction_loss": 0.009563876309816215, + "validation_tcremp_cosine": 0.9917293340642124, + "walltime_seconds_from_log": 2814.008, + "artifacts_verified": true, + "error_scan_clean": true, + "best_checkpoint_sha256": "545d0485a07e60bbc549a22a9e8940d8a5c392420a9ccd9b1604bf95360826ef" + }, + "256": { + "job_id": "1426123", + "best_epoch": 23, + "best_validation_loss": 0.20649890744307267, + "validation_pgen_rmse": 0.27295139267261015, + "validation_reconstruction_loss": 0.009048107677470638, + "validation_tcremp_cosine": 0.9917872234284248, + "walltime_seconds_from_log": 2953.805, + "artifacts_verified": true, + "error_scan_clean": true, + "best_checkpoint_sha256": "096de58b79b41b69693d40e6b4bfd57f31fae9d2aab1c5e9d727b3154b0db251" + }, + "512": { + "job_id": "1426176", + "best_epoch": 25, + "best_validation_loss": 0.2104906694103532, + "validation_pgen_rmse": 0.30843847231625027, + "validation_reconstruction_loss": 0.008969121246607934, + "validation_tcremp_cosine": 0.9916971205322241, + "walltime_seconds_from_log": 2827.815, + "artifacts_verified": true, + "error_scan_clean": true, + "best_checkpoint_sha256": "740c521273e8ee34589ff928b262fd8c60aa34bbd88d567bba0daae73223e03a" + }, + "1024": { + "job_id": "1426177", + "best_epoch": 28, + "best_validation_loss": 0.21072169335227994, + "validation_pgen_rmse": 0.3169206069630301, + "validation_reconstruction_loss": 0.006642488566201312, + "validation_tcremp_cosine": 0.9915933634905607, + "walltime_seconds_from_log": 2946.908, + "artifacts_verified": true, + "error_scan_clean": true, + "best_checkpoint_sha256": "64d8967b0633b4c645daf025f189fb209cc824297820e52e1def790817b7c338" + }, + "2048": { + "job_id": "1426180", + "best_epoch": 24, + "best_validation_loss": 0.21218146349443082, + "validation_pgen_rmse": 0.3269114392204356, + "validation_reconstruction_loss": 0.011945437021755085, + "validation_tcremp_cosine": 0.9917859220653428, + "walltime_seconds_from_log": 2821.125, + "artifacts_verified": true, + "error_scan_clean": true, + "best_checkpoint_sha256": "4351751bc12f601785eb6efcd4e5d1b984919c0a0992a3214d6c219a998a225b" + }, + "4096": { + "job_id": "1426181", + "best_epoch": 25, + "best_validation_loss": 0.21656744164125932, + "validation_pgen_rmse": 0.4111111972200827, + "validation_reconstruction_loss": 0.012061045652630643, + "validation_tcremp_cosine": 0.9918066222555235, + "walltime_seconds_from_log": 2929.382, + "artifacts_verified": true, + "error_scan_clean": true, + "best_checkpoint_sha256": "932cd6e4135862beb4fba3781a5593d7915e3471e40670d4a2c47b2a3276d6d6" + } + }, + "validation_ranking": [44, 256, 64, 128, 512, 1024, 2048, 4096], + "selected_condition": { + "vocab_size": 44, + "selection_basis": "minimum full validation objective", + "best_validation_loss": 0.19666322000361444, + "character_baseline_validation_loss": 0.19693734171098476, + "replication_required_before_confirmatory_claim": true + }, + "active_conditions": {}, + "submission_status": "stage_a_complete", + "replication_status": "not_authorized" +} diff --git a/experiments/trb/wordpiece-sweep-seed42/validation_results.md b/experiments/trb/wordpiece-sweep-seed42/validation_results.md new file mode 100644 index 0000000..4c947c3 --- /dev/null +++ b/experiments/trb/wordpiece-sweep-seed42/validation_results.md @@ -0,0 +1,80 @@ +# Leakage-safe TRB WordPiece vocabulary sweep (Stage A) + +Status: complete. Selection used only the full validation objective. Test metrics were +recorded after checkpoint selection and were not used to rank conditions. + +## Locked controls + +- Train-only tokenizer fit: 79,544 sequences; train manifest SHA-256 + `6399cbea503191a652e7ecef0b8b859793f4ece16322fd49ce2ed03364484b01`. +- Benchmark READY SHA-256: + `ce08945c0c0a975c2b05129695c26399456f185fe79b89914c0ad52e4fdcf167`. +- Special IDs: PAD/BOS/EOS/UNK = 0/1/2/3; zero UNK and exact roundtrip on train, + validation, and test for every vocabulary. +- Matched model: full R+T+P, seed 42, latent 128, encoder/decoder 4+4, 40 epochs, + batch 64, AdamW 3e-4, locked train-only normalizer, separate WordPiece encoder and + unchanged 25-token character decoder. + +## Validation-only ranking + +| Rank | WordPiece vocab | Job | Best epoch | Full val objective | Val pgen RMSE | Val recon loss | Val token acc | Val TCRemP cosine | Val TCRemP std MSE | Params | Walltime | +|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 1 | 44 | 1426073 | 30 | 0.196663 | 0.343658 | 0.004776 | 0.998500 | 0.992229 | 0.258091 | 21,157,033 | 48m11s | +| 2 | 256 | 1426123 | 23 | 0.206499 | 0.272951 | 0.009048 | 0.996913 | 0.991787 | 0.270087 | 21,224,873 | 49m14s | +| 3 | 64 | 1426074 | 28 | 0.206852 | 0.324221 | 0.007868 | 0.997243 | 0.991838 | 0.269134 | 21,163,433 | 51m00s | +| 4 | 128 | 1426122 | 26 | 0.209016 | 0.283852 | 0.009564 | 0.996720 | 0.991729 | 0.272229 | 21,183,913 | 46m54s | +| 5 | 512 | 1426176 | 25 | 0.210491 | 0.308438 | 0.008969 | 0.997031 | 0.991697 | 0.273659 | 21,306,793 | 47m08s | +| 6 | 1024 | 1426177 | 28 | 0.210722 | 0.316921 | 0.006642 | 0.997884 | 0.991593 | 0.276770 | 21,470,633 | 49m07s | +| 7 | 2048 | 1426180 | 24 | 0.212181 | 0.326911 | 0.011945 | 0.996184 | 0.991786 | 0.270806 | 21,798,313 | 47m01s | +| 8 | 4096 | 1426181 | 25 | 0.216567 | 0.411111 | 0.012061 | 0.995960 | 0.991807 | 0.269719 | 22,453,673 | 48m49s | + +The preregistered selection is WordPiece-44. Its validation objective is 0.000274 +(0.139%) lower than the matched character baseline (0.196937), a small single-seed +difference that requires replication before a confirmatory claim. + +## Held-out test metrics (descriptive only) + +| Vocab | Total | Pgen RMSE | Recon loss | Exact match | Length accuracy | Mean edit | TCRemP cosine | TCRemP std MSE | +|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 44 | 0.195894 | 0.341319 | 0.005795 | 0.978980 | 0.990747 | 0.047068 | 0.992311 | 0.255408 | +| 64 | 0.205969 | 0.325752 | 0.009292 | 0.964397 | 0.985015 | 0.081867 | 0.991931 | 0.265650 | +| 128 | 0.207857 | 0.289369 | 0.010379 | 0.959871 | 0.977371 | 0.087298 | 0.991808 | 0.269085 | +| 256 | 0.203443 | 0.270520 | 0.009503 | 0.963693 | 0.982199 | 0.084783 | 0.991930 | 0.265257 | +| 512 | 0.207351 | 0.315414 | 0.009045 | 0.966610 | 0.981897 | 0.075128 | 0.991832 | 0.268639 | +| 1024 | 0.209171 | 0.315975 | 0.007251 | 0.972543 | 0.983405 | 0.059841 | 0.991686 | 0.273637 | +| 2048 | 0.210649 | 0.329018 | 0.012147 | 0.955245 | 0.974052 | 0.096047 | 0.991866 | 0.268131 | +| 4096 | 0.216030 | 0.419660 | 0.013340 | 0.951725 | 0.970331 | 0.107010 | 0.991912 | 0.266441 | + +## Token interpretation + +The winning 44-token tokenizer is effectively the smallest valid standard WordPiece +system: `C` is the start token in all 79,544 training sequences, and most continuation +pieces are single amino acids. `##F` is strongly end-enriched (mean normalized position +0.934; end log2 enrichment 3.378). Frequent pieces include `##S` (166,534 occurrences, +94.0% of sequences) and `##A` (118,254, 97.7%). + +Train-only descriptive presence associations show lower `log10_pgen_1mm` for sequences +containing `##V` (mean difference -1.211; high-vs-low-quartile log2 enrichment -1.483), +`##W` (-1.267; -1.852), and internal `##C` (-1.342; -1.909). These are descriptive, +unadjusted associations, not biological or causal claims. + +At vocabulary 64, multi-residue boundary pieces emerge (`CA`, `CAS`, `##YEQYF`, +`##TQYF`, `##EQFF`). At 256, the same boundary families remain common. At 4096, +the vocabulary fragments into much rarer context-specific pieces: `##YEQYF` occurs in +2.49% of train sequences and is positively associated with higher pgen +(mean difference +0.665; quartile log2 enrichment +1.005), while rare single-residue +fallbacks such as `##F` show strong negative associations. This pattern is consistent +with reduced statistical sharing at very large vocabularies and accompanies worse joint +validation objectives, but it does not establish a biological mechanism. + +## Artifacts + +- Tokenizers and audits: + `/home/evlasova/irrm-codec/experiments/trb/tokenizers/wordpiece-sweep-seed42` +- Model runs: + `/home/evlasova/irrm-codec/experiments/trb/wp-sweep-v-rtp-seed42` +- Each run contains `best.pt`, `last.pt`, `history.json`, `run_config.json`, + `test_metrics.json`, `target_standardizer.npz`, and logs. All terminal artifact and + failure scans passed. + +No replication or additional architecture jobs were launched. diff --git a/irrm_codec/anchored_tokenization.py b/irrm_codec/anchored_tokenization.py new file mode 100644 index 0000000..bab2431 --- /dev/null +++ b/irrm_codec/anchored_tokenization.py @@ -0,0 +1,186 @@ +"""Sequence-only terminal-anchor tokenizers with a WordPiece middle. + +The serialized bundle keeps protected literal N/C anchors in one input token and +uses an independently trained WordPiece model only for the remaining middle. +No V/J annotation is accepted by the runtime encoder. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +from tokenizers import Tokenizer + +from irrm_codec.tokenization import BOS_ID, EOS_ID, PAD_ID, UNK_ID, VALID_AA + + +SPECIAL_TOKENS = ["[PAD]", "[BOS]", "[EOS]", "[UNK]"] +EXPECTED_SPECIAL_IDS = { + "[PAD]": PAD_ID, + "[BOS]": BOS_ID, + "[EOS]": EOS_ID, + "[UNK]": UNK_ID, +} +SUPPORTED_KINDS = {"edge_k", "data_anchor", "germline_anchor"} + + +@dataclass(frozen=True) +class AnchoredEncoding: + ids: list[int] + tokens: list[str] + boundary_types: list[str] + n_anchor: str + middle: str + c_anchor: str + + +def _normalize(sequence: str, max_len: int) -> str: + sequence = "" if sequence is None else str(sequence).strip().upper() + if not sequence: + raise ValueError("Sequence must not be empty.") + invalid = sorted(set(sequence).difference(VALID_AA)) + if invalid: + raise ValueError(f"Sequence contains unsupported amino acids: {invalid}") + if len(sequence) > max_len: + raise ValueError( + f"Sequence length {len(sequence)} exceeds max_len={max_len}." + ) + return sequence + + +class AnchoredTokenizer: + """Runtime encoder for an immutable anchored-tokenizer bundle.""" + + def __init__(self, bundle_path: str | Path): + self.bundle_path = Path(bundle_path) + payload = json.loads(self.bundle_path.read_text(encoding="utf-8")) + kind = payload.get("tokenizer_type") + if kind not in SUPPORTED_KINDS: + raise ValueError(f"Unsupported anchored tokenizer type: {kind!r}.") + self.kind = str(kind) + self.k = int(payload.get("k") or 0) + self.n_anchors = tuple(payload["n_anchors"]) + self.c_anchors = tuple(payload["c_anchors"]) + self.n_anchor_set = set(self.n_anchors) + self.c_anchor_set = set(self.c_anchors) + self.vocab = {str(token): int(identifier) for token, identifier in payload["vocab"].items()} + self.id_to_token = {identifier: token for token, identifier in self.vocab.items()} + if len(self.vocab) != len(self.id_to_token): + raise ValueError("Anchored vocabulary IDs must be unique.") + for token, expected in EXPECTED_SPECIAL_IDS.items(): + if self.vocab.get(token) != expected: + raise ValueError( + f"Anchored tokenizer has {token}={self.vocab.get(token)}, expected {expected}." + ) + expected_ids = list(range(len(self.vocab))) + if sorted(self.id_to_token) != expected_ids: + raise ValueError("Anchored vocabulary IDs must be contiguous from zero.") + central_path = self.bundle_path.parent / payload["central_tokenizer_file"] + self.central = Tokenizer.from_file(str(central_path)) + for token, expected in EXPECTED_SPECIAL_IDS.items(): + if self.central.token_to_id(token) != expected: + raise ValueError( + f"Central WordPiece has {token}={self.central.token_to_id(token)}, expected {expected}." + ) + self.central_to_input = { + int(key): int(value) for key, value in payload["central_to_input"].items() + } + self.input_to_central = {value: key for key, value in self.central_to_input.items()} + + @property + def vocab_size(self) -> int: + return len(self.vocab) + + def _select_anchors(self, sequence: str) -> tuple[str, str]: + if self.kind == "edge_k": + if len(sequence) < 2 * self.k: + raise ValueError( + f"EDGE-{self.k} requires sequence length >= {2 * self.k}, got {len(sequence)}." + ) + n_anchor = sequence[: self.k] + c_anchor = sequence[-self.k :] + if n_anchor not in self.n_anchor_set or c_anchor not in self.c_anchor_set: + raise ValueError( + f"EDGE-{self.k} literal is absent from the exhaustive anchor inventory: {sequence!r}." + ) + return n_anchor, c_anchor + + n_matches = [anchor for anchor in self.n_anchors if sequence.startswith(anchor)] + c_matches = [anchor for anchor in self.c_anchors if sequence.endswith(anchor)] + feasible = [ + (n_anchor, c_anchor) + for n_anchor in n_matches + for c_anchor in c_matches + if len(n_anchor) + len(c_anchor) <= len(sequence) + ] + if not feasible: + raise ValueError(f"No non-overlapping protected anchor pair matches {sequence!r}.") + # Longest combined literal coverage, then longer N, then longer C, then lexical. + return min( + feasible, + key=lambda pair: ( + -(len(pair[0]) + len(pair[1])), + -len(pair[0]), + -len(pair[1]), + pair[0], + pair[1], + ), + ) + + def encode_with_boundaries(self, sequence: str, max_len: int) -> AnchoredEncoding: + sequence = _normalize(sequence, max_len=max_len) + n_anchor, c_anchor = self._select_anchors(sequence) + middle = sequence[len(n_anchor) : len(sequence) - len(c_anchor)] + central_encoding = self.central.encode(middle) if middle else None + central_ids = [] if central_encoding is None else central_encoding.ids + if UNK_ID in central_ids: + raise ValueError(f"Central WordPiece emitted [UNK] for middle {middle!r}.") + try: + input_middle_ids = [self.central_to_input[identifier] for identifier in central_ids] + n_id = self.vocab[f"[N:{n_anchor}]"] + c_id = self.vocab[f"[C:{c_anchor}]"] + except KeyError as error: + raise ValueError(f"Anchored tokenizer bundle is missing token {error.args[0]!r}.") from error + ids = [n_id, *input_middle_ids, c_id] + if len(ids) > max_len: + raise ValueError( + f"Sequence {sequence!r} encodes to {len(ids)} tokens, exceeding max_len={max_len}." + ) + tokens = [self.id_to_token[identifier] for identifier in ids] + return AnchoredEncoding( + ids=ids, + tokens=tokens, + boundary_types=["N_ANCHOR", *(["MIDDLE_WORDPIECE"] * len(input_middle_ids)), "C_ANCHOR"], + n_anchor=n_anchor, + middle=middle, + c_anchor=c_anchor, + ) + + def encode(self, sequence: str, max_len: int) -> list[int]: + return self.encode_with_boundaries(sequence, max_len).ids + + def decode(self, ids: list[int]) -> str: + tokens = [self.id_to_token[int(identifier)] for identifier in ids] + if len(tokens) < 2 or not tokens[0].startswith("[N:") or not tokens[-1].startswith("[C:"): + raise ValueError("Anchored encoding must start with N_ANCHOR and end with C_ANCHOR.") + n_anchor = tokens[0][3:-1] + c_anchor = tokens[-1][3:-1] + central_ids = [self.input_to_central[int(identifier)] for identifier in ids[1:-1]] + middle = self.central.decode(central_ids, skip_special_tokens=True) if central_ids else "" + return f"{n_anchor}{middle}{c_anchor}" + + +class AnchoredUnpaddedEncodeFn: + """Picklable callable used by multi-worker benchmark DataLoaders.""" + + def __init__(self, tokenizer: AnchoredTokenizer): + self.tokenizer = tokenizer + + def __call__(self, sequence: str, max_len: int) -> list[int]: + return self.tokenizer.encode(sequence, max_len) + + +def load_anchored_tokenizer(path: str | Path) -> AnchoredTokenizer: + return AnchoredTokenizer(path) diff --git a/irrm_codec/multitask_data.py b/irrm_codec/multitask_data.py index 55b6072..1550b74 100644 --- a/irrm_codec/multitask_data.py +++ b/irrm_codec/multitask_data.py @@ -15,13 +15,6 @@ encode_raw, encode_reconstruction_pair, ) -from irrm_codec.wordpiece_tokenization import ( - WordpieceUnpaddedEncodeFn, - load_wordpiece_tokenizer, - wordpiece_vocab_size, -) - - REQUIRED_DATASET_COLUMNS = { "row_index", "junction_aa", @@ -95,11 +88,36 @@ def resolve_encoder_tokenizer( path=None, encode=encode_raw, ) + if tokenizer_type in {"edge_k", "data_anchor", "germline_anchor"}: + if not tokenizer_path: + raise ValueError("--tokenizer-path is required for anchored input.") + from irrm_codec.anchored_tokenization import ( + AnchoredUnpaddedEncodeFn, + load_anchored_tokenizer, + ) + + tokenizer = load_anchored_tokenizer(tokenizer_path) + if tokenizer.kind != tokenizer_type: + raise ValueError( + f"Tokenizer bundle kind {tokenizer.kind!r} does not match {tokenizer_type!r}." + ) + return EncoderTokenizer( + name=tokenizer.kind, + vocab_size=tokenizer.vocab_size, + path=str(Path(tokenizer_path)), + encode=AnchoredUnpaddedEncodeFn(tokenizer), + ) if tokenizer_type != "wordpiece": raise ValueError(f"Unsupported tokenizer type: {tokenizer_type!r}.") if not tokenizer_path: raise ValueError("--tokenizer-path is required for WordPiece input.") + from irrm_codec.wordpiece_tokenization import ( + WordpieceUnpaddedEncodeFn, + load_wordpiece_tokenizer, + wordpiece_vocab_size, + ) + tokenizer = load_wordpiece_tokenizer(tokenizer_path) return EncoderTokenizer( diff --git a/irrm_codec/multitask_losses.py b/irrm_codec/multitask_losses.py index e454026..b110f26 100644 --- a/irrm_codec/multitask_losses.py +++ b/irrm_codec/multitask_losses.py @@ -17,28 +17,53 @@ class MultiTaskLossWeights: class IRRMCodecMultiTaskLoss(nn.Module): - """Combine normalized losses without mixing incompatible target scales. - - The TCRemP and pgen heads predict standardized train-target values. TCRemP - cosine loss is nevertheless computed after de-standardization, so the model - is optimized and evaluated in the original embedding geometry as well. - """ + """Combine normalized losses without mixing incompatible target scales.""" def __init__( self, weights: MultiTaskLossWeights | None = None, tcremp_mse_fraction: float = 0.7, + tcremp_centered_cosine_weight: float = 0.0, + tcremp_pairwise_log_distance_weight: float = 0.0, pgen_huber_delta: float = 0.5, label_smoothing: float = 0.0, ): super().__init__() if not 0.0 <= tcremp_mse_fraction <= 1.0: raise ValueError("tcremp_mse_fraction must lie in [0, 1].") + for name, value in ( + ("tcremp_centered_cosine_weight", tcremp_centered_cosine_weight), + ("tcremp_pairwise_log_distance_weight", tcremp_pairwise_log_distance_weight), + ): + if not 0.0 <= value <= 1.0: + raise ValueError(f"{name} must lie in [0, 1].") + if tcremp_centered_cosine_weight + tcremp_pairwise_log_distance_weight > 1.0: + raise ValueError("TCRemP geometry weights must sum to at most 1.") self.weights = weights or MultiTaskLossWeights() self.tcremp_mse_fraction = tcremp_mse_fraction + self.tcremp_centered_cosine_weight = tcremp_centered_cosine_weight + self.tcremp_pairwise_log_distance_weight = tcremp_pairwise_log_distance_weight self.pgen_huber_delta = pgen_huber_delta self.label_smoothing = label_smoothing + @staticmethod + def _pairwise_log_cosine_distance( + prediction: torch.Tensor, + target: torch.Tensor, + ) -> torch.Tensor: + """Match off-diagonal raw-space cosine distances on a log scale.""" + if prediction.size(0) < 2: + return prediction.new_zeros(()) + pred_unit = F.normalize(prediction.float(), dim=-1) + target_unit = F.normalize(target.float(), dim=-1) + pred_distance = (1.0 - pred_unit @ pred_unit.T).clamp_min(1e-7) + target_distance = (1.0 - target_unit @ target_unit.T).clamp_min(1e-7) + mask = ~torch.eye(prediction.size(0), dtype=torch.bool, device=prediction.device) + return F.smooth_l1_loss( + pred_distance[mask].log(), + target_distance[mask].log(), + ) + def forward( self, outputs: dict[str, torch.Tensor | None], @@ -66,10 +91,34 @@ def forward( tcremp_target, dim=-1, ).mean() - tcremp_loss = ( + tcremp_legacy = ( self.tcremp_mse_fraction * tcremp_mse + (1.0 - self.tcremp_mse_fraction) * tcremp_cosine ) + if self.tcremp_centered_cosine_weight > 0: + tcremp_centered_cosine = 1.0 - F.cosine_similarity( + pred_tcremp_std.float(), + target_tcremp_std.float(), + dim=-1, + ).mean() + else: + tcremp_centered_cosine = pred_tcremp_std.new_zeros(()) + if self.tcremp_pairwise_log_distance_weight > 0: + tcremp_pairwise = self._pairwise_log_cosine_distance( + pred_tcremp_raw, + tcremp_target, + ) + else: + tcremp_pairwise = pred_tcremp_std.new_zeros(()) + geometry_weight = ( + self.tcremp_centered_cosine_weight + + self.tcremp_pairwise_log_distance_weight + ) + tcremp_loss = ( + (1.0 - geometry_weight) * tcremp_legacy + + self.tcremp_centered_cosine_weight * tcremp_centered_cosine + + self.tcremp_pairwise_log_distance_weight * tcremp_pairwise + ) safe_pgen_std = pgen_std.clamp_min(1e-8) target_pgen_std = (pgen_target - pgen_mean) / safe_pgen_std @@ -78,14 +127,12 @@ def forward( target_pgen_std, delta=self.pgen_huber_delta, ) - reconstruction_loss = F.cross_entropy( reconstruction_logits.reshape(-1, reconstruction_logits.size(-1)), reconstruction_target.reshape(-1), ignore_index=PAD_ID, label_smoothing=self.label_smoothing, ) - total = ( self.weights.tcremp * tcremp_loss + self.weights.pgen * pgen_loss @@ -96,6 +143,8 @@ def forward( "tcremp_loss": tcremp_loss, "tcremp_mse_standardized": tcremp_mse, "tcremp_cosine_loss_raw": tcremp_cosine, + "tcremp_centered_cosine_loss": tcremp_centered_cosine, + "tcremp_pairwise_log_distance_loss": tcremp_pairwise, "pgen_loss": pgen_loss, "reconstruction_loss": reconstruction_loss, } diff --git a/irrm_codec/multitask_metrics.py b/irrm_codec/multitask_metrics.py index e89bd33..2f82276 100644 --- a/irrm_codec/multitask_metrics.py +++ b/irrm_codec/multitask_metrics.py @@ -188,6 +188,8 @@ def compute(self, include_spearman: bool = True) -> dict[str, float | int]: covariance = self.pgen_sum_cross - self.pgen_sum_pred * self.pgen_sum_target / n pearson_denom = math.sqrt(max(pred_var, 0.0) * max(target_var, 0.0)) pearson = covariance / pearson_denom if pearson_denom > 0 else float("nan") + if math.isfinite(pearson): + pearson = max(-1.0, min(1.0, pearson)) metrics.update( { diff --git a/irrm_codec/tokenization.py b/irrm_codec/tokenization.py index f400d97..cadb514 100644 --- a/irrm_codec/tokenization.py +++ b/irrm_codec/tokenization.py @@ -1,8 +1,8 @@ AA_VOCAB = { "": 0, - "": 1, - "": 2, - "": 3, + "": 1, + "": 2, + "": 3, "-": 4, "A": 5, "C": 6, diff --git a/irrm_codec/train_multitask.py b/irrm_codec/train_multitask.py index 2f8f08a..8631ae2 100644 --- a/irrm_codec/train_multitask.py +++ b/irrm_codec/train_multitask.py @@ -37,7 +37,11 @@ def parse_args() -> argparse.Namespace: choices=["log10_pgen", "log10_pgen_1mm"], default="log10_pgen_1mm", ) - parser.add_argument("--tokenizer-type", choices=["char", "wordpiece"], default="char") + parser.add_argument( + "--tokenizer-type", + choices=["char", "wordpiece", "edge_k", "data_anchor", "germline_anchor"], + default="char", + ) parser.add_argument("--tokenizer-path") parser.add_argument("--max-sequence-len", type=int, default=40) @@ -56,6 +60,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--pgen-loss-weight", type=float, default=1.0) parser.add_argument("--reconstruction-loss-weight", type=float, default=1.0) parser.add_argument("--tcremp-mse-fraction", type=float, default=0.7) + parser.add_argument("--tcremp-centered-cosine-weight", type=float, default=0.0) + parser.add_argument("--tcremp-pairwise-log-distance-weight", type=float, default=0.0) parser.add_argument("--pgen-huber-delta", type=float, default=0.5) parser.add_argument("--label-smoothing", type=float, default=0.0) @@ -70,6 +76,13 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--scheduler-patience", type=int, default=2) parser.add_argument("--scheduler-min-lr", type=float, default=1e-6) parser.add_argument("--standardizer-chunk-size", type=int, default=1024) + parser.add_argument( + "--standardizer-path", + help=( + "Locked train-only target standardizer. Defaults to " + "/target_standardizer.npz for train-subset=all." + ), + ) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--num-workers", type=int, default=0) @@ -130,6 +143,16 @@ def validate_args(args: argparse.Namespace) -> None: == 0 ): raise ValueError("At least one task loss weight must be positive.") + geometry_weight = ( + args.tcremp_centered_cosine_weight + + args.tcremp_pairwise_log_distance_weight + ) + if ( + args.tcremp_centered_cosine_weight < 0 + or args.tcremp_pairwise_log_distance_weight < 0 + or geometry_weight > 1 + ): + raise ValueError("TCRemP geometry weights must be non-negative and sum to at most 1.") def choose_device(requested: str) -> torch.device: @@ -180,6 +203,7 @@ def _resolve_best_checkpoint(output_dir: Path, resume: str | None) -> Path: ) + def run_epoch( *, model: IRRMCodecTransformer, @@ -247,9 +271,7 @@ def run_epoch( if is_train: accumulation_group_size = _gradient_accumulation_group_size( - step, - limit, - gradient_accumulation_steps, + step, limit, gradient_accumulation_steps ) scaled_loss = losses["loss"] / accumulation_group_size scaler.scale(scaled_loss).backward() @@ -383,18 +405,42 @@ def main() -> None: test_indices = select_split_indices(table, args.data_dir, "test") tokenizer = resolve_encoder_tokenizer(args.tokenizer_type, args.tokenizer_path) - standardizer_path = output_dir / "target_standardizer.npz" + output_standardizer_path = output_dir / "target_standardizer.npz" + locked_standardizer_path = ( + Path(args.standardizer_path) + if args.standardizer_path + else Path(args.data_dir) / "target_standardizer.npz" + ) resume_standardizer_path = ( - Path(args.resume).parent / "target_standardizer.npz" - if args.resume - else None + Path(args.resume).parent / "target_standardizer.npz" if args.resume else None ) if resume_standardizer_path is not None and resume_standardizer_path.exists(): standardizer = TargetStandardizer.load(resume_standardizer_path) - if resume_standardizer_path.resolve() != standardizer_path.resolve(): - standardizer.save(standardizer_path) - elif args.resume and standardizer_path.exists(): - standardizer = TargetStandardizer.load(standardizer_path) + standardizer.save(output_standardizer_path) + if standardizer.pgen_target != args.pgen_target: + raise ValueError("Resume checkpoint and requested pgen target do not match.") + if standardizer.train_rows != len(train_indices): + raise ValueError("Resume standardizer and requested train subset do not match.") + if len(standardizer.tcremp_mean) != embeddings.shape[1]: + raise ValueError("Resume standardizer and TCRemP dimension do not match.") + elif args.resume and output_standardizer_path.exists(): + standardizer = TargetStandardizer.load(output_standardizer_path) + if standardizer.pgen_target != args.pgen_target: + raise ValueError("Resume checkpoint and requested pgen target do not match.") + if standardizer.train_rows != len(train_indices): + raise ValueError("Resume standardizer and requested train subset do not match.") + if len(standardizer.tcremp_mean) != embeddings.shape[1]: + raise ValueError("Resume standardizer and TCRemP dimension do not match.") + elif args.train_subset == "all" and locked_standardizer_path.exists(): + logger.info("loading locked benchmark standardizer from %s", locked_standardizer_path) + standardizer = TargetStandardizer.load(locked_standardizer_path) + if standardizer.pgen_target != args.pgen_target: + raise ValueError("Locked benchmark standardizer pgen target does not match.") + if standardizer.train_rows != len(train_indices): + raise ValueError("Locked benchmark standardizer train row count does not match.") + if len(standardizer.tcremp_mean) != embeddings.shape[1]: + raise ValueError("Locked benchmark standardizer TCRemP dimension does not match.") + standardizer.save(output_standardizer_path) else: logger.info( "computing train-only target statistics rows=%d embedding_dim=%d", @@ -408,14 +454,7 @@ def main() -> None: args.pgen_target, chunk_size=args.standardizer_chunk_size, ) - standardizer.save(standardizer_path) - if args.resume: - if standardizer.pgen_target != args.pgen_target: - raise ValueError("Resume checkpoint and requested pgen target do not match.") - if standardizer.train_rows != len(train_indices): - raise ValueError("Resume standardizer and requested train subset do not match.") - if len(standardizer.tcremp_mean) != embeddings.shape[1]: - raise ValueError("Resume standardizer and TCRemP dimension do not match.") + standardizer.save(output_standardizer_path) standardizer_tensors = standardizer.as_torch(device) datasets = { @@ -480,6 +519,8 @@ def main() -> None: reconstruction=args.reconstruction_loss_weight, ), tcremp_mse_fraction=args.tcremp_mse_fraction, + tcremp_centered_cosine_weight=args.tcremp_centered_cosine_weight, + tcremp_pairwise_log_distance_weight=args.tcremp_pairwise_log_distance_weight, pgen_huber_delta=args.pgen_huber_delta, label_smoothing=args.label_smoothing, ) @@ -525,7 +566,12 @@ def main() -> None: "test_rows": len(test_indices), "tcremp_dim": int(embeddings.shape[1]), "pgen_target": args.pgen_target, - "standardizer": str(standardizer_path), + "standardizer": str(output_standardizer_path), + "locked_standardizer_source": ( + str(locked_standardizer_path) + if args.train_subset == "all" and locked_standardizer_path.exists() + else None + ), }, ) logger.info( @@ -660,11 +706,6 @@ def main() -> None: break best_checkpoint_path = _resolve_best_checkpoint(output_dir, args.resume) - if best_checkpoint_path != output_dir / "best.pt": - logger.warning( - "No best checkpoint was created in this output directory; evaluating %s instead.", - best_checkpoint_path, - ) best_checkpoint = torch.load(best_checkpoint_path, map_location=device) model.load_state_dict(best_checkpoint["model_state"]) test_metrics = run_epoch( diff --git a/irrm_codec/wordpiece_tokenization.py b/irrm_codec/wordpiece_tokenization.py index 0a77c57..351a3fa 100644 --- a/irrm_codec/wordpiece_tokenization.py +++ b/irrm_codec/wordpiece_tokenization.py @@ -2,27 +2,29 @@ Requires: pip install tokenizers Requires a tokenizer.json with special tokens at [PAD]=0, [UNK]=1, [BOS]=2, [EOS]=3 — -same ids irrm_codec.tokenization uses. A mismatched file raises ValueError. +The author layout [PAD]/[BOS]/[EOS]/[UNK]=0/1/2/3 is also accepted. IDs are never remapped. """ from pathlib import Path from tokenizers import Tokenizer from tokenizers.models import WordPiece -from irrm_codec.tokenization import BOS_ID, EOS_ID, PAD_ID, UNK_ID +from irrm_codec.tokenization import BOS_ID, EOS_ID, PAD_ID, UNK_ID, VALID_AA _EXPECTED_SPECIAL_IDS = {"[PAD]": PAD_ID, "[UNK]": UNK_ID, "[BOS]": BOS_ID, "[EOS]": EOS_ID} -def _validate_special_ids(tokenizer, tokenizer_path): - """Check PAD/UNK/BOS/EOS ids match, since ForwardModel/InverseModel assume padding_idx=0.""" - for token, expected_id in _EXPECTED_SPECIAL_IDS.items(): - actual_id = tokenizer.token_to_id(token) - if actual_id != expected_id: - raise ValueError( - f"Tokenizer at {tokenizer_path} has {token}={actual_id}, expected {expected_id}." - ) +def validate_wordpiece_tokenizer(tokenizer, source): + """Accept the published author and student special-token layouts without remapping IDs.""" + observed = {token: tokenizer.token_to_id(token) for token in _EXPECTED_SPECIAL_IDS} + author = {"[PAD]": 0, "[BOS]": 1, "[EOS]": 2, "[UNK]": 3} + student = {"[PAD]": 0, "[UNK]": 1, "[BOS]": 2, "[EOS]": 3} + if observed not in (author, student): + raise ValueError(f"Tokenizer at {source} has unsupported special IDs: {observed}.") + +def _validate_special_ids(tokenizer, tokenizer_path): + validate_wordpiece_tokenizer(tokenizer, tokenizer_path) def load_wordpiece_tokenizer(path) -> Tokenizer: """Load and validate a tokenizer.json. Returns the Tokenizer itself — it's picklable, @@ -48,23 +50,28 @@ def encode_wordpiece(seq, tokenizer, max_len): return ids + [PAD_ID] * (max_len - len(ids)) -def encode_wordpiece_unpadded(seq, tokenizer, max_len): - """Encode a CDR3 without padding; batching supplies PAD tokens later. - - This is the preferred interface for the joint Transformer because it keeps - sequence length explicit and avoids passing a fixed-width buffer through the - dataset for every example. - """ - seq = "" if seq is None else str(seq).strip().upper() - if not seq: +def encode_wordpiece_unpadded( + sequence: str, + tokenizer: Tokenizer, + max_len: int, +) -> list[int]: + sequence = "" if sequence is None else str(sequence).strip().upper() + if not sequence: raise ValueError("Sequence must not be empty.") - ids = tokenizer.encode(seq).ids - if len(ids) > max_len: + invalid = sorted(set(sequence).difference(VALID_AA)) + if invalid: + raise ValueError(f"Sequence contains unsupported amino acids: {invalid}") + token_ids = tokenizer.encode(sequence).ids + if not token_ids: + raise ValueError("WordPiece encoding must contain at least one token.") + if tokenizer.token_to_id("[UNK]") in token_ids: + raise ValueError(f"WordPiece encoding contains [UNK] for sequence {sequence!r}.") + if len(token_ids) > max_len: raise ValueError( - f"Sequence '{seq}' encodes to {len(ids)} WordPiece tokens, " - f"exceeds max_len={max_len}." + f"Sequence {sequence!r} encodes to {len(token_ids)} tokens, " + f"exceeding max_len={max_len}." ) - return ids + return token_ids class WordpieceUnpaddedEncodeFn: @@ -116,8 +123,9 @@ def encode_wordpiece_anchored(seq, tokenizer, max_len, left_anchor=1, right_anch def decode_wordpiece(token_ids, tokenizer, stop_at_eos=True): """Decode ids back to an amino-acid string, stopping at the first EOS by default.""" ids = list(token_ids) - if stop_at_eos and EOS_ID in ids: - ids = ids[: ids.index(EOS_ID)] + eos_id = tokenizer.token_to_id("[EOS]") + if stop_at_eos and eos_id in ids: + ids = ids[: ids.index(eos_id)] return tokenizer.decode(ids, skip_special_tokens=True) diff --git a/requirements-benchmark.txt b/requirements-benchmark.txt new file mode 100644 index 0000000..6be1846 --- /dev/null +++ b/requirements-benchmark.txt @@ -0,0 +1,4 @@ +# Optional external representation encoders; install after requirements.txt. +fair-esm +transformers +sceptr diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..aa395c5 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,2 @@ +-r requirements_no_mirpy.txt +pytest>=8,<9 diff --git a/requirements.txt b/requirements.txt index 394cc7c..e7a17c9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,3 +15,9 @@ tokenizers>=0.20,<1 wandb torch==2.4.1 mirpy-lib @ git+https://github.com/antigenomics/mirpy.git@11a5fd1ac12edd22e12a2bb167ab9dfb55ccefd4 + +scipy +joblib +tidytcells>=2,<3 +rapidfuzz>=3,<4 +biopython diff --git a/requirements_no_mirpy.txt b/requirements_no_mirpy.txt index 7df91da..5ec5a69 100644 --- a/requirements_no_mirpy.txt +++ b/requirements_no_mirpy.txt @@ -13,3 +13,9 @@ ipywidgets tqdm tokenizers>=0.20,<1 torch==2.4.1 + +scipy +joblib +tidytcells>=2,<3 +rapidfuzz>=3,<4 +biopython diff --git a/results/author/README.md b/results/author/README.md new file mode 100644 index 0000000..c26c2c8 --- /dev/null +++ b/results/author/README.md @@ -0,0 +1,8 @@ +# Author result archive + +These tables were recovered from the local manuscript/tables directory during +publication integration. They record the original author experiments, including +R/T/P factorial, tokenization/architecture comparisons, external models, latent +sufficiency, reconstruction errors, and downstream evaluation. They were not +regenerated by the integration CPU tests. Preserve the recorded hashes and run +provenance when reproducing or comparing these results. diff --git a/results/author/architecture_results.tsv b/results/author/architecture_results.tsv new file mode 100644 index 0000000..ed1f2ad --- /dev/null +++ b/results/author/architecture_results.tsv @@ -0,0 +1,6 @@ +condition latent_dim encoder_layers decoder_layers best_epoch best_validation_joint_loss test_reconstruction_exact_match test_tcremp_cosine test_pgen_rmse parameters remote_output +d128-e4d4 128 4 4 30 0.19693734171098476 0.9815950920245399 0.992289113360584 0.343395136972041 21142953 /home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-seed42 +d320-e4d4 320 4 4 28 0.20322377112032147 0.971236045459117 0.9922739834801287 0.3678781484793622 21697257 /home/evlasova/irrm-codec/experiments/trb/rq1-char-rtp-seed42 +d512-e4d4 512 4 4 28 0.2087959912276347 0.9676154078246002 0.9919841033189718 0.3447591167246714 22251561 /home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d512-seed42 +d128-e2d4 128 2 4 38 0.2047392800008221 0.9614804384994469 0.9922445351528616 0.3915365453431801 18677033 /home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-e2d4-seed42 +d128-e6d4 128 6 4 28 0.20005685247478622 0.9690234335713568 0.992220890590342 0.3328977050320967 23608873 /home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-e6d4-seed42 diff --git a/results/author/downstream_distribution_summary.tsv b/results/author/downstream_distribution_summary.tsv new file mode 100644 index 0000000..2aec099 --- /dev/null +++ b/results/author/downstream_distribution_summary.tsv @@ -0,0 +1,85 @@ +distance metric model epitopes median q1 q3 +cosine between_minus_within p 6 0.0145443152895451 -0.0001492005485736 0.0223180058645856 +cosine between_minus_within r 6 0.1578570451716967 0.0660611825275395 0.2059822047681393 +cosine between_minus_within rp 6 0.1482288745391226 0.0703371563373106 0.1812976137030069 +cosine between_minus_within rt 6 0.1485969787030767 0.1105945336383283 0.1936926422338269 +cosine between_minus_within rtp 6 0.1181646810375563 0.1091203728571297 0.1666706381533964 +cosine between_minus_within t 6 0.1445365267388816 0.0487515996728354 0.1554817927078322 +cosine between_minus_within tp 6 0.0558434676762929 0.0177303851753551 0.1254538239660743 +cosine map_at_r p 6 0.1072723899090958 0.0499999999999999 0.3327528171366436 +cosine map_at_r r 6 0.567181302109368 0.2961775566854647 0.9219392854417802 +cosine map_at_r rp 6 0.5712887961220562 0.3008174984726075 0.8470631754505331 +cosine map_at_r rt 6 0.5785180537443047 0.394076685303714 0.8092542051437623 +cosine map_at_r rtp 6 0.5720908539934587 0.3674453785075067 0.8205274195292898 +cosine map_at_r t 6 0.4102288717786875 0.241837828946378 0.5799432736096324 +cosine map_at_r tp 6 0.3849056822155335 0.2746976756621377 0.617755165445629 +cosine precision_at_1 p 6 0.15 0.0 0.6293478260869565 +cosine precision_at_1 r 6 0.997835497835498 0.982613401091662 1.0 +cosine precision_at_1 rp 6 0.9967532467532468 0.9820722755505364 1.0 +cosine precision_at_1 rt 6 0.9869659326181066 0.9758152173913044 0.998917748917749 +cosine precision_at_1 rtp 6 0.9826369282891022 0.9758152173913044 0.9967532467532468 +cosine precision_at_1 t 6 0.8125 0.7625 0.9073369565217392 +cosine precision_at_1 tp 6 0.879891304347826 0.80625 0.97557876905703 +cosine same_label_auprc_balanced p 6 0.539248199520719 0.5053728183074511 0.5879886551902822 +cosine same_label_auprc_balanced r 6 0.8475517022311078 0.6490237688365724 0.9912878787878788 +cosine same_label_auprc_balanced rp 6 0.8210715594987712 0.6474908672137286 0.9811625874125874 +cosine same_label_auprc_balanced rt 6 0.8034531209606439 0.7060939707370587 0.974431818181818 +cosine same_label_auprc_balanced rtp 6 0.7985890749138258 0.7036019824914448 0.9706818181818182 +cosine same_label_auprc_balanced t 6 0.7141583539785815 0.6082069732618179 0.920682486972982 +cosine same_label_auprc_balanced tp 6 0.7098644847055979 0.6231296867497326 0.8973546314400598 +cosine same_label_auroc p 6 0.504926571111111 0.4405555555555555 0.5609014279785969 +cosine same_label_auroc r 6 0.8457943800000001 0.6161970710939357 0.99 +cosine same_label_auroc rp 6 0.81789524 0.604472590273484 0.9775 +cosine same_label_auroc rt 6 0.78903498 0.6648428768234964 0.9675 +cosine same_label_auroc rtp 6 0.77995546 0.665116706405517 0.96 +cosine same_label_auroc t 6 0.7309987030273752 0.5775910405350774 0.9163325281803544 +cosine same_label_auroc tp 6 0.7320116917146258 0.5748606921593343 0.8970451933333334 +cosine within_between_ratio p 6 0.8179821629283468 0.7973325641955162 1.0103488838501615 +cosine within_between_ratio r 6 0.7461820844136712 0.5318130153584926 0.896319378478054 +cosine within_between_ratio rp 6 0.7360651990271578 0.5343375419702201 0.8872807430525629 +cosine within_between_ratio rt 6 0.7598338444758671 0.604724903085157 0.8425857353682646 +cosine within_between_ratio rtp 6 0.7715636576640352 0.6536983166194029 0.8369475654044015 +cosine within_between_ratio t 6 0.7699551601043487 0.7099430371505517 0.9071517773056192 +cosine within_between_ratio tp 6 0.8493645602381839 0.7302541305366014 0.9526920339881136 +euclidean between_minus_within p 6 0.2757123701977926 -0.059618771315719 0.6019926425050299 +euclidean between_minus_within r 6 1.6116113735004554 0.7905967767468112 2.6223420450288786 +euclidean between_minus_within rp 6 1.5069047169069432 0.8281081652673254 2.4094358188863936 +euclidean between_minus_within rt 6 1.712747171445943 1.1221761363938518 2.3251597072080803 +euclidean between_minus_within rtp 6 1.5569426077436344 1.1079599161060276 1.7965944639688614 +euclidean between_minus_within t 6 1.7479573007303433 0.738132089438619 2.701660348818735 +euclidean between_minus_within tp 6 1.0889405252505675 0.3118886026511105 1.9733571716059028 +euclidean map_at_r p 6 0.1069136915349483 0.0499999999999999 0.3323742775309038 +euclidean map_at_r r 6 0.5686769709072258 0.296089041743911 0.9225953504042612 +euclidean map_at_r rp 6 0.5738011794219053 0.301259424271951 0.8472519876707403 +euclidean map_at_r rt 6 0.5801873236682918 0.4032949812936193 0.8079024354126133 +euclidean map_at_r rtp 6 0.5800994521465858 0.3816460500084527 0.8178226775033324 +euclidean map_at_r t 6 0.4067968767490306 0.2407373145960313 0.5793277904451506 +euclidean map_at_r tp 6 0.3811436671924311 0.2747485823983004 0.6158232511979151 +euclidean precision_at_1 p 6 0.1375 0.0 0.6230978260869565 +euclidean precision_at_1 r 6 0.997835497835498 0.982613401091662 1.0 +euclidean precision_at_1 rp 6 0.997835497835498 0.982613401091662 1.0 +euclidean precision_at_1 rt 6 0.9869659326181066 0.9758152173913044 0.998917748917749 +euclidean precision_at_1 rtp 6 0.9826369282891022 0.9758152173913044 0.9967532467532468 +euclidean precision_at_1 t 6 0.8125 0.7625 0.9073369565217392 +euclidean precision_at_1 tp 6 0.879891304347826 0.80625 0.97557876905703 +euclidean same_label_auprc_balanced p 6 0.5193031996501787 0.5037457912457912 0.5462770451630575 +euclidean same_label_auprc_balanced r 6 0.850156124431406 0.6546985430978252 0.9912878787878788 +euclidean same_label_auprc_balanced rp 6 0.8228910796395241 0.6550536252738018 0.9811625874125874 +euclidean same_label_auprc_balanced rt 6 0.8085321068617215 0.7101054409871914 0.9740311771561772 +euclidean same_label_auprc_balanced rtp 6 0.8005937366322124 0.712261035132749 0.9706818181818182 +euclidean same_label_auprc_balanced t 6 0.7141135767918775 0.606869109378777 0.9199743118632536 +euclidean same_label_auprc_balanced tp 6 0.708716250339988 0.6232321985930138 0.8963979330461858 +euclidean same_label_auroc p 6 0.4924432822222222 0.4336111111111111 0.5602215882639715 +euclidean same_label_auroc r 6 0.84832026 0.6218091192984543 0.99 +euclidean same_label_auroc rp 6 0.8188651 0.6075487212247325 0.9775 +euclidean same_label_auroc rt 6 0.78825486 0.6805495752075895 0.9675 +euclidean same_label_auroc rtp 6 0.7754911 0.6878897774032066 0.96 +euclidean same_label_auroc t 6 0.7301146150906672 0.5756820298335315 0.9150484842120004 +euclidean same_label_auroc tp 6 0.730726850510397 0.5733397310340332 0.8974585733333333 +euclidean within_between_ratio p 6 0.9349746031223288 0.8766179889749843 1.0221335500292326 +euclidean within_between_ratio r 6 0.8656682254017987 0.7232751133175406 0.9347943555106916 +euclidean within_between_ratio rp 6 0.858484774291782 0.7266690133104422 0.9261933907546518 +euclidean within_between_ratio rt 6 0.8595267247842686 0.7694552586796968 0.9104275665656826 +euclidean within_between_ratio rtp 6 0.8676658082373783 0.8053474848489822 0.9064168776507302 +euclidean within_between_ratio t 6 0.8669592497804925 0.8033823364614108 0.936983570542994 +euclidean within_between_ratio tp 6 0.8990771644752427 0.8353543721278662 0.9696957552483164 diff --git a/results/author/downstream_per_epitope_controls.tsv b/results/author/downstream_per_epitope_controls.tsv new file mode 100644 index 0000000..7122789 --- /dev/null +++ b/results/author/downstream_per_epitope_controls.tsv @@ -0,0 +1,97 @@ +rank epitope distance control metric value +1 GILGFVFTL absolute exact_cdr3_length within_between_ratio 1.0 +1 GILGFVFTL absolute exact_cdr3_length between_minus_within 0.0 +1 GILGFVFTL absolute exact_cdr3_length same_label_auroc 0.5 +1 GILGFVFTL absolute exact_cdr3_length same_label_auprc_balanced 0.5 +1 GILGFVFTL cosine amino_acid_composition within_between_ratio 0.917498976188238 +1 GILGFVFTL cosine amino_acid_composition between_minus_within 0.0227998023803722 +1 GILGFVFTL cosine amino_acid_composition same_label_auroc 0.55205354 +1 GILGFVFTL cosine amino_acid_composition same_label_auprc_balanced 0.5690443818486688 +1 GILGFVFTL euclidean amino_acid_composition within_between_ratio 0.963622217536067 +1 GILGFVFTL euclidean amino_acid_composition between_minus_within 0.009632094464311 +1 GILGFVFTL euclidean amino_acid_composition same_label_auroc 0.54093074 +1 GILGFVFTL euclidean amino_acid_composition same_label_auprc_balanced 0.5540953848468919 +1 GILGFVFTL mismatch_fraction v_j_identity within_between_ratio 0.5486555697823303 +1 GILGFVFTL mismatch_fraction v_j_identity between_minus_within 0.423 +1 GILGFVFTL mismatch_fraction v_j_identity same_label_auroc 0.8715852800000001 +1 GILGFVFTL mismatch_fraction v_j_identity same_label_auprc_balanced 0.8325820334830544 +2 GLCTLVAML absolute exact_cdr3_length within_between_ratio 1.0 +2 GLCTLVAML absolute exact_cdr3_length between_minus_within 0.0 +2 GLCTLVAML absolute exact_cdr3_length same_label_auroc 0.5 +2 GLCTLVAML absolute exact_cdr3_length same_label_auprc_balanced 0.5 +2 GLCTLVAML cosine amino_acid_composition within_between_ratio 0.954935606190231 +2 GLCTLVAML cosine amino_acid_composition between_minus_within 0.0128782543677601 +2 GLCTLVAML cosine amino_acid_composition same_label_auroc 0.5206604588205093 +2 GLCTLVAML cosine amino_acid_composition same_label_auprc_balanced 0.5860336166083757 +2 GLCTLVAML euclidean amino_acid_composition within_between_ratio 0.9648202050628965 +2 GLCTLVAML euclidean amino_acid_composition between_minus_within 0.0089334411397177 +2 GLCTLVAML euclidean amino_acid_composition same_label_auroc 0.5192853975588696 +2 GLCTLVAML euclidean amino_acid_composition same_label_auprc_balanced 0.5944189527018383 +2 GLCTLVAML mismatch_fraction v_j_identity within_between_ratio 0.8186555497655029 +2 GLCTLVAML mismatch_fraction v_j_identity between_minus_within 0.1681159420289855 +2 GLCTLVAML mismatch_fraction v_j_identity same_label_auroc 0.6016854535695115 +2 GLCTLVAML mismatch_fraction v_j_identity same_label_auprc_balanced 0.6104768679055156 +3 YLQPRTFLL absolute exact_cdr3_length within_between_ratio 1.0 +3 YLQPRTFLL absolute exact_cdr3_length between_minus_within 0.0 +3 YLQPRTFLL absolute exact_cdr3_length same_label_auroc 0.4999999999999999 +3 YLQPRTFLL absolute exact_cdr3_length same_label_auprc_balanced 0.5 +3 YLQPRTFLL cosine amino_acid_composition within_between_ratio 0.9531587680806616 +3 YLQPRTFLL cosine amino_acid_composition between_minus_within 0.0114973732487685 +3 YLQPRTFLL cosine amino_acid_composition same_label_auroc 0.5639429250891795 +3 YLQPRTFLL cosine amino_acid_composition same_label_auprc_balanced 0.596767984185325 +3 YLQPRTFLL euclidean amino_acid_composition within_between_ratio 0.9119303437002227 +3 YLQPRTFLL euclidean amino_acid_composition between_minus_within 0.0221158269223845 +3 YLQPRTFLL euclidean amino_acid_composition same_label_auroc 0.6133992865636146 +3 YLQPRTFLL euclidean amino_acid_composition same_label_auprc_balanced 0.6361629629451064 +3 YLQPRTFLL mismatch_fraction v_j_identity within_between_ratio 0.8848996832101372 +3 YLQPRTFLL mismatch_fraction v_j_identity between_minus_within 0.0751724137931034 +3 YLQPRTFLL mismatch_fraction v_j_identity same_label_auroc 0.5652775267538646 +3 YLQPRTFLL mismatch_fraction v_j_identity same_label_auprc_balanced 0.5440484630600354 +4 SLYNTVATL absolute exact_cdr3_length within_between_ratio 1.0 +4 SLYNTVATL absolute exact_cdr3_length between_minus_within 0.0 +4 SLYNTVATL absolute exact_cdr3_length same_label_auroc 0.5 +4 SLYNTVATL absolute exact_cdr3_length same_label_auprc_balanced 0.5 +4 SLYNTVATL cosine amino_acid_composition within_between_ratio 0.4896836516369041 +4 SLYNTVATL cosine amino_acid_composition between_minus_within 0.0630725596060895 +4 SLYNTVATL cosine amino_acid_composition same_label_auroc 0.88 +4 SLYNTVATL cosine amino_acid_composition same_label_auprc_balanced 0.9102272727272728 +4 SLYNTVATL euclidean amino_acid_composition within_between_ratio 0.7248228937768533 +4 SLYNTVATL euclidean amino_acid_composition between_minus_within 0.0564259005060354 +4 SLYNTVATL euclidean amino_acid_composition same_label_auroc 0.84 +4 SLYNTVATL euclidean amino_acid_composition same_label_auprc_balanced 0.8726984126984126 +4 SLYNTVATL mismatch_fraction v_j_identity within_between_ratio 0.4615384615384615 +4 SLYNTVATL mismatch_fraction v_j_identity between_minus_within 0.35 +4 SLYNTVATL mismatch_fraction v_j_identity same_label_auroc 0.79 +4 SLYNTVATL mismatch_fraction v_j_identity same_label_auprc_balanced 0.7529411764705882 +5 ALSKGVHFV absolute exact_cdr3_length within_between_ratio 1.0 +5 ALSKGVHFV absolute exact_cdr3_length between_minus_within 0.0 +5 ALSKGVHFV absolute exact_cdr3_length same_label_auroc 0.5 +5 ALSKGVHFV absolute exact_cdr3_length same_label_auprc_balanced 0.5 +5 ALSKGVHFV cosine amino_acid_composition within_between_ratio 0.43903884813344 +5 ALSKGVHFV cosine amino_acid_composition between_minus_within 0.0618176669520792 +5 ALSKGVHFV cosine amino_acid_composition same_label_auroc 0.7916666666666667 +5 ALSKGVHFV cosine amino_acid_composition same_label_auprc_balanced 0.7500000000000001 +5 ALSKGVHFV euclidean amino_acid_composition within_between_ratio 0.6709093005374359 +5 ALSKGVHFV euclidean amino_acid_composition between_minus_within 0.0494387749198131 +5 ALSKGVHFV euclidean amino_acid_composition same_label_auroc 0.8194444444444445 +5 ALSKGVHFV euclidean amino_acid_composition same_label_auprc_balanced 0.763888888888889 +5 ALSKGVHFV mismatch_fraction v_j_identity within_between_ratio 0.0 +5 ALSKGVHFV mismatch_fraction v_j_identity between_minus_within 0.5 +5 ALSKGVHFV mismatch_fraction v_j_identity same_label_auroc 1.0 +5 ALSKGVHFV mismatch_fraction v_j_identity same_label_auprc_balanced 1.0 +6 FLYALALLL absolute exact_cdr3_length within_between_ratio 1.0 +6 FLYALALLL absolute exact_cdr3_length between_minus_within 0.0 +6 FLYALALLL absolute exact_cdr3_length same_label_auroc 0.5 +6 FLYALALLL absolute exact_cdr3_length same_label_auprc_balanced 0.5 +6 FLYALALLL cosine amino_acid_composition within_between_ratio 0.8643636499681533 +6 FLYALALLL cosine amino_acid_composition between_minus_within 0.0276565312034742 +6 FLYALALLL cosine amino_acid_composition same_label_auroc 1.0 +6 FLYALALLL cosine amino_acid_composition same_label_auprc_balanced 1.0 +6 FLYALALLL euclidean amino_acid_composition within_between_ratio 0.983532756656882 +6 FLYALALLL euclidean amino_acid_composition between_minus_within 0.0036101150869386 +6 FLYALALLL euclidean amino_acid_composition same_label_auroc 1.0 +6 FLYALALLL euclidean amino_acid_composition same_label_auprc_balanced 1.0 +6 FLYALALLL mismatch_fraction v_j_identity within_between_ratio 1.0 +6 FLYALALLL mismatch_fraction v_j_identity between_minus_within 0.0 +6 FLYALALLL mismatch_fraction v_j_identity same_label_auroc 0.5 +6 FLYALALLL mismatch_fraction v_j_identity same_label_auprc_balanced 0.5 diff --git a/results/author/downstream_per_epitope_metrics.tsv b/results/author/downstream_per_epitope_metrics.tsv new file mode 100644 index 0000000..c657e1d --- /dev/null +++ b/results/author/downstream_per_epitope_metrics.tsv @@ -0,0 +1,505 @@ +rank epitope distance model metric value +1 GILGFVFTL cosine r within_between_ratio 0.814834956105159 +1 GILGFVFTL cosine r between_minus_within 0.1233428152129985 +1 GILGFVFTL cosine r same_label_auroc 0.73158876 +1 GILGFVFTL cosine r same_label_auprc_balanced 0.7299518893107003 +1 GILGFVFTL cosine r precision_at_1 0.9956709956709956 +1 GILGFVFTL cosine r map_at_r 0.8377571417671207 +1 GILGFVFTL cosine t within_between_ratio 0.7684862592239499 +1 GILGFVFTL cosine t between_minus_within 0.1538348010838437 +1 GILGFVFTL cosine t same_label_auroc 0.71333396 +1 GILGFVFTL cosine t same_label_auprc_balanced 0.674158188636663 +1 GILGFVFTL cosine t precision_at_1 0.9913419913419912 +1 GILGFVFTL cosine t map_at_r 0.8391364170281714 +1 GILGFVFTL cosine p within_between_ratio 0.8230982847313358 +1 GILGFVFTL cosine p between_minus_within 0.0206568423984859 +1 GILGFVFTL cosine p same_label_auroc 0.53763092 +1 GILGFVFTL cosine p same_label_auprc_balanced 0.521491273229804 +1 GILGFVFTL cosine p precision_at_1 0.8982683982683982 +1 GILGFVFTL cosine p map_at_r 0.7469039220366203 +1 GILGFVFTL cosine rt within_between_ratio 0.8470954060378486 +1 GILGFVFTL cosine rt between_minus_within 0.1070347350523491 +1 GILGFVFTL cosine rt same_label_auroc 0.70806996 +1 GILGFVFTL cosine rt same_label_auprc_balanced 0.7050656379180732 +1 GILGFVFTL cosine rt precision_at_1 0.9956709956709956 +1 GILGFVFTL cosine rt map_at_r 0.8290056068583498 +1 GILGFVFTL cosine rp within_between_ratio 0.8013631708826814 +1 GILGFVFTL cosine rp between_minus_within 0.1202215255443032 +1 GILGFVFTL cosine rp same_label_auroc 0.7257904800000001 +1 GILGFVFTL cosine rp same_label_auprc_balanced 0.7174927693471926 +1 GILGFVFTL cosine rp precision_at_1 0.9935064935064936 +1 GILGFVFTL cosine rp map_at_r 0.8382527018021324 +1 GILGFVFTL cosine tp within_between_ratio 0.7243638016662785 +1 GILGFVFTL cosine tp between_minus_within 0.137600465097264 +1 GILGFVFTL cosine tp same_label_auroc 0.75484744 +1 GILGFVFTL cosine tp same_label_auprc_balanced 0.7053104436509567 +1 GILGFVFTL cosine tp precision_at_1 0.9891774891774892 +1 GILGFVFTL cosine tp map_at_r 0.8556084509363918 +1 GILGFVFTL cosine rtp within_between_ratio 0.829036865860044 +1 GILGFVFTL cosine rtp between_minus_within 0.110323732416114 +1 GILGFVFTL cosine rtp same_label_auroc 0.71991092 +1 GILGFVFTL cosine rtp same_label_auprc_balanced 0.7144508771003788 +1 GILGFVFTL cosine rtp precision_at_1 0.987012987012987 +1 GILGFVFTL cosine rtp map_at_r 0.835703226039053 +1 GILGFVFTL euclidean r within_between_ratio 0.8980290222030745 +1 GILGFVFTL euclidean r between_minus_within 1.2324917821683776 +1 GILGFVFTL euclidean r same_label_auroc 0.73664052 +1 GILGFVFTL euclidean r same_label_auprc_balanced 0.7351607337112969 +1 GILGFVFTL euclidean r precision_at_1 0.9956709956709956 +1 GILGFVFTL euclidean r map_at_r 0.8403814016170456 +1 GILGFVFTL euclidean t within_between_ratio 0.8627014070154413 +1 GILGFVFTL euclidean t between_minus_within 1.7771783413246125 +1 GILGFVFTL euclidean t same_label_auroc 0.71670196 +1 GILGFVFTL euclidean t same_label_auprc_balanced 0.6769013347021693 +1 GILGFVFTL euclidean t precision_at_1 0.9913419913419912 +1 GILGFVFTL euclidean t map_at_r 0.8439129645545131 +1 GILGFVFTL euclidean p within_between_ratio 0.8973830858007207 +1 GILGFVFTL euclidean p between_minus_within 0.4762456869122298 +1 GILGFVFTL euclidean p same_label_auroc 0.54044212 +1 GILGFVFTL euclidean p same_label_auprc_balanced 0.5236232343171925 +1 GILGFVFTL euclidean p precision_at_1 0.8982683982683982 +1 GILGFVFTL euclidean p map_at_r 0.7478342545035729 +1 GILGFVFTL euclidean rt within_between_ratio 0.9181802458439324 +1 GILGFVFTL euclidean rt between_minus_within 1.0255226173407426 +1 GILGFVFTL euclidean rt same_label_auroc 0.70650972 +1 GILGFVFTL euclidean rt same_label_auprc_balanced 0.7064940862833436 +1 GILGFVFTL euclidean rt precision_at_1 0.9956709956709956 +1 GILGFVFTL euclidean rt map_at_r 0.8272032472168177 +1 GILGFVFTL euclidean rp within_between_ratio 0.8919376258759846 +1 GILGFVFTL euclidean rp between_minus_within 1.1943521196618183 +1 GILGFVFTL euclidean rp same_label_auroc 0.7277302 +1 GILGFVFTL euclidean rp same_label_auprc_balanced 0.7211318096286985 +1 GILGFVFTL euclidean rp precision_at_1 0.9956709956709956 +1 GILGFVFTL euclidean rp map_at_r 0.8390079506829614 +1 GILGFVFTL euclidean tp within_between_ratio 0.8430098768730618 +1 GILGFVFTL euclidean tp between_minus_within 1.764684817162717 +1 GILGFVFTL euclidean tp same_label_auroc 0.75650096 +1 GILGFVFTL euclidean tp same_label_auprc_balanced 0.7068407684952329 +1 GILGFVFTL euclidean tp precision_at_1 0.9891774891774892 +1 GILGFVFTL euclidean tp map_at_r 0.8574148540779141 +1 GILGFVFTL euclidean rtp within_between_ratio 0.9120983001562664 +1 GILGFVFTL euclidean rtp between_minus_within 1.0340184856414627 +1 GILGFVFTL euclidean rtp same_label_auroc 0.7109821999999999 +1 GILGFVFTL euclidean rtp same_label_auprc_balanced 0.7101946466646145 +1 GILGFVFTL euclidean rtp precision_at_1 0.987012987012987 +1 GILGFVFTL euclidean rtp map_at_r 0.8320969033377764 +2 GLCTLVAML cosine r within_between_ratio 0.9304252886506998 +2 GLCTLVAML cosine r between_minus_within 0.0469673049657198 +2 GLCTLVAML cosine r same_label_auroc 0.5325370487059209 +2 GLCTLVAML cosine r same_label_auprc_balanced 0.6220477286785298 +2 GLCTLVAML cosine r precision_at_1 0.9782608695652174 +2 GLCTLVAML cosine r map_at_r 0.2960349214300812 +2 GLCTLVAML cosine t within_between_ratio 0.6904286297927523 +2 GLCTLVAML cosine t between_minus_within 0.2327473467424893 +2 GLCTLVAML cosine t same_label_auroc 0.7486634460547503 +2 GLCTLVAML cosine t same_label_auprc_balanced 0.7541585193205 +2 GLCTLVAML cosine t precision_at_1 0.9347826086956522 +2 GLCTLVAML cosine t map_at_r 0.5281064277718629 +2 GLCTLVAML cosine p within_between_ratio 0.2016262709941349 +2 GLCTLVAML cosine p between_minus_within 0.1346310777629427 +2 GLCTLVAML cosine p same_label_auroc 0.8133053280123224 +2 GLCTLVAML cosine p same_label_auprc_balanced 0.764362746527017 +2 GLCTLVAML cosine p precision_at_1 0.7391304347826086 +2 GLCTLVAML cosine p map_at_r 0.3971554962427943 +2 GLCTLVAML cosine rt within_between_ratio 0.8290567233595127 +2 GLCTLVAML cosine rt between_minus_within 0.1212739293962659 +2 GLCTLVAML cosine rt same_label_auroc 0.6504338490979953 +2 GLCTLVAML cosine rt same_label_auprc_balanced 0.7091789691940154 +2 GLCTLVAML cosine rt precision_at_1 0.9782608695652174 +2 GLCTLVAML cosine rt map_at_r 0.3897568779087489 +2 GLCTLVAML cosine rp within_between_ratio 0.9159199337758568 +2 GLCTLVAML cosine rp between_minus_within 0.0537090332683131 +2 GLCTLVAML cosine rp same_label_auroc 0.5408042194683657 +2 GLCTLVAML cosine rp same_label_auprc_balanced 0.6241568998359073 +2 GLCTLVAML cosine rp precision_at_1 0.9782608695652174 +2 GLCTLVAML cosine rp map_at_r 0.30432489044198 +2 GLCTLVAML cosine tp within_between_ratio 0.7479251171475703 +2 GLCTLVAML cosine tp between_minus_within 0.1461652305207085 +2 GLCTLVAML cosine tp same_label_auroc 0.7091759434292515 +2 GLCTLVAML cosine tp same_label_auprc_balanced 0.7144185257602391 +2 GLCTLVAML cosine tp precision_at_1 0.9347826086956522 +2 GLCTLVAML cosine tp map_at_r 0.4710206617825163 +2 GLCTLVAML cosine rtp within_between_ratio 0.8395844652525207 +2 GLCTLVAML cosine rtp between_minus_within 0.1087192530041349 +2 GLCTLVAML cosine rtp same_label_auroc 0.6468519685406894 +2 GLCTLVAML cosine rtp same_label_auprc_balanced 0.6999856842884669 +2 GLCTLVAML cosine rtp precision_at_1 0.9782608695652174 +2 GLCTLVAML cosine rtp map_at_r 0.3668666020143697 +2 GLCTLVAML euclidean r within_between_ratio 0.9470494666132306 +2 GLCTLVAML euclidean r between_minus_within 0.643298441606289 +2 GLCTLVAML euclidean r same_label_auroc 0.543551541459544 +2 GLCTLVAML euclidean r same_label_auprc_balanced 0.6278778128933347 +2 GLCTLVAML euclidean r precision_at_1 0.9782608695652174 +2 GLCTLVAML euclidean r map_at_r 0.2969725401974061 +2 GLCTLVAML euclidean t within_between_ratio 0.7836093129434006 +2 GLCTLVAML euclidean t between_minus_within 3.009821017983443 +2 GLCTLVAML euclidean t same_label_auroc 0.7435272701813344 +2 GLCTLVAML euclidean t same_label_auprc_balanced 0.7513258188815857 +2 GLCTLVAML euclidean t precision_at_1 0.9347826086956522 +2 GLCTLVAML euclidean t map_at_r 0.525644495113936 +2 GLCTLVAML euclidean p within_between_ratio 0.4021732315244646 +2 GLCTLVAML euclidean p between_minus_within 3.49534116615842 +2 GLCTLVAML euclidean p same_label_auroc 0.810628952834372 +2 GLCTLVAML euclidean p same_label_auprc_balanced 0.7625697321880942 +2 GLCTLVAML euclidean p precision_at_1 0.7391304347826086 +2 GLCTLVAML euclidean p map_at_r 0.3968899090179063 +2 GLCTLVAML euclidean rt within_between_ratio 0.8871695287309327 +2 GLCTLVAML euclidean rt between_minus_within 1.4121366935531796 +2 GLCTLVAML euclidean rt same_label_auroc 0.6718961936101193 +2 GLCTLVAML euclidean rt same_label_auprc_balanced 0.7209395050987346 +2 GLCTLVAML euclidean rt precision_at_1 0.9782608695652174 +2 GLCTLVAML euclidean rt map_at_r 0.4009350926126311 +2 GLCTLVAML euclidean rp within_between_ratio 0.9376119790475408 +2 GLCTLVAML euclidean rp between_minus_within 0.7060268471358278 +2 GLCTLVAML euclidean rp same_label_auroc 0.5572181381129082 +2 GLCTLVAML euclidean rp same_label_auprc_balanced 0.6330275638221696 +2 GLCTLVAML euclidean rp precision_at_1 0.9782608695652174 +2 GLCTLVAML euclidean rp map_at_r 0.3085944081608493 +2 GLCTLVAML euclidean tp within_between_ratio 0.832802537212801 +2 GLCTLVAML euclidean tp between_minus_within 2.042914623086965 +2 GLCTLVAML euclidean tp same_label_auroc 0.7049527410207939 +2 GLCTLVAML euclidean tp same_label_auprc_balanced 0.710591732184743 +2 GLCTLVAML euclidean tp precision_at_1 0.9347826086956522 +2 GLCTLVAML euclidean tp map_at_r 0.4632930047916605 +2 GLCTLVAML euclidean rtp within_between_ratio 0.8893726101341216 +2 GLCTLVAML euclidean rtp between_minus_within 1.3297842074997224 +2 GLCTLVAML euclidean rtp same_label_auroc 0.6801923032042755 +2 GLCTLVAML euclidean rtp same_label_auprc_balanced 0.7184602005371521 +2 GLCTLVAML euclidean rtp precision_at_1 0.9782608695652174 +2 GLCTLVAML euclidean rtp map_at_r 0.3851989042931715 +3 YLQPRTFLL cosine r within_between_ratio 0.9234808526023524 +3 YLQPRTFLL cosine r between_minus_within 0.0446161337209015 +3 YLQPRTFLL cosine r same_label_auroc 0.5777331747919143 +3 YLQPRTFLL cosine r same_label_auprc_balanced 0.6199786587907083 +3 YLQPRTFLL cosine r precision_at_1 1.0 +3 YLQPRTFLL cosine r map_at_r 0.2966054624516153 +3 YLQPRTFLL cosine t within_between_ratio 0.9523943494125764 +3 YLQPRTFLL cosine t between_minus_within 0.019922715432474 +3 YLQPRTFLL cosine t same_label_auroc 0.5323434007134364 +3 YLQPRTFLL cosine t same_label_auprc_balanced 0.5862232348035362 +3 YLQPRTFLL cosine t precision_at_1 0.825 +3 YLQPRTFLL cosine t map_at_r 0.292351315785512 +3 YLQPRTFLL cosine p within_between_ratio 0.8128660411253578 +3 YLQPRTFLL cosine p between_minus_within 0.0228717270199522 +3 YLQPRTFLL cosine p same_label_auroc 0.5686582639714625 +3 YLQPRTFLL cosine p same_label_auprc_balanced 0.557005125811634 +3 YLQPRTFLL cosine p precision_at_1 0.3 +3 YLQPRTFLL cosine p map_at_r 0.1395447798181916 +3 YLQPRTFLL cosine rt within_between_ratio 0.8632042895855496 +3 YLQPRTFLL cosine rt between_minus_within 0.0844520133839246 +3 YLQPRTFLL cosine rt same_label_auroc 0.6346615933412604 +3 YLQPRTFLL cosine rt same_label_auprc_balanced 0.6779547121190657 +3 YLQPRTFLL cosine rt precision_at_1 0.975 +3 YLQPRTFLL cosine rt map_at_r 0.4070361074886092 +3 YLQPRTFLL cosine rp within_between_ratio 0.9310489264658908 +3 YLQPRTFLL cosine rp between_minus_within 0.0366156392001849 +3 YLQPRTFLL cosine rp same_label_auroc 0.5640332936979786 +3 YLQPRTFLL cosine rp same_label_auprc_balanced 0.5943170885404121 +3 YLQPRTFLL cosine rp precision_at_1 1.0 +3 YLQPRTFLL cosine rp map_at_r 0.29964836781615 +3 YLQPRTFLL cosine tp within_between_ratio 0.9533213775412188 +3 YLQPRTFLL cosine tp between_minus_within 0.0160828353071133 +3 YLQPRTFLL cosine tp same_label_auroc 0.5300889417360286 +3 YLQPRTFLL cosine tp same_label_auprc_balanced 0.5957361011159912 +3 YLQPRTFLL cosine tp precision_at_1 0.825 +3 YLQPRTFLL cosine tp map_at_r 0.2987907026485508 +3 YLQPRTFLL cosine rtp within_between_ratio 0.9120239325822764 +3 YLQPRTFLL cosine rtp between_minus_within 0.0486971483198341 +3 YLQPRTFLL cosine rtp same_label_auroc 0.5801588585017836 +3 YLQPRTFLL cosine rtp same_label_auprc_balanced 0.6137564655276165 +3 YLQPRTFLL cosine rtp precision_at_1 0.975 +3 YLQPRTFLL cosine rtp map_at_r 0.3691817079869173 +3 YLQPRTFLL euclidean r within_between_ratio 0.9530694427510062 +3 YLQPRTFLL euclidean r between_minus_within 0.5335972980226611 +3 YLQPRTFLL euclidean r same_label_auroc 0.5835319857312723 +3 YLQPRTFLL euclidean r same_label_auprc_balanced 0.6232151785513477 +3 YLQPRTFLL euclidean r precision_at_1 1.0 +3 YLQPRTFLL euclidean r map_at_r 0.2957945422594127 +3 YLQPRTFLL euclidean t within_between_ratio 0.9589057298754774 +3 YLQPRTFLL euclidean t between_minus_within 0.4112640325394672 +3 YLQPRTFLL euclidean t same_label_auroc 0.5286753864447087 +3 YLQPRTFLL euclidean t same_label_auprc_balanced 0.5835250342709797 +3 YLQPRTFLL euclidean t precision_at_1 0.825 +3 YLQPRTFLL euclidean t map_at_r 0.2879492583841253 +3 YLQPRTFLL euclidean p within_between_ratio 0.8696962900330723 +3 YLQPRTFLL euclidean p between_minus_within 0.6439082943692966 +3 YLQPRTFLL euclidean p same_label_auroc 0.566814744351962 +3 YLQPRTFLL euclidean p same_label_auprc_balanced 0.5538283154450125 +3 YLQPRTFLL euclidean p precision_at_1 0.275 +3 YLQPRTFLL euclidean p map_at_r 0.1388273830698967 +3 YLQPRTFLL euclidean rt within_between_ratio 0.9200626838923368 +3 YLQPRTFLL euclidean rt between_minus_within 0.9411507264122356 +3 YLQPRTFLL euclidean rt same_label_auroc 0.6455286563614744 +3 YLQPRTFLL euclidean rt same_label_auprc_balanced 0.6898694962560761 +3 YLQPRTFLL euclidean rt precision_at_1 0.975 +3 YLQPRTFLL euclidean rt map_at_r 0.4103746473365836 +3 YLQPRTFLL euclidean rp within_between_ratio 0.9583005918155642 +3 YLQPRTFLL euclidean rp between_minus_within 0.4338869883892666 +3 YLQPRTFLL euclidean rp same_label_auroc 0.5674882282996433 +3 YLQPRTFLL euclidean rp same_label_auprc_balanced 0.5968576844192361 +3 YLQPRTFLL euclidean rp precision_at_1 1.0 +3 YLQPRTFLL euclidean rp map_at_r 0.2988144296423183 +3 YLQPRTFLL euclidean tp within_between_ratio 0.9551444520774234 +3 YLQPRTFLL euclidean tp between_minus_within 0.413196233338418 +3 YLQPRTFLL euclidean tp same_label_auroc 0.529468727705113 +3 YLQPRTFLL euclidean tp same_label_auprc_balanced 0.595362675292274 +3 YLQPRTFLL euclidean tp precision_at_1 0.825 +3 YLQPRTFLL euclidean tp map_at_r 0.2989943295932016 +3 YLQPRTFLL euclidean rtp within_between_ratio 0.9433479558101976 +3 YLQPRTFLL euclidean rtp between_minus_within 0.6197374939630098 +3 YLQPRTFLL euclidean rtp same_label_auroc 0.596221640903686 +3 YLQPRTFLL euclidean rtp same_label_auprc_balanced 0.6267377767138391 +3 YLQPRTFLL euclidean rtp precision_at_1 0.975 +3 YLQPRTFLL euclidean rtp map_at_r 0.3804617652468797 +4 SLYNTVATL cosine r within_between_ratio 0.4832409495705956 +4 SLYNTVATL cosine r between_minus_within 0.1923712751303949 +4 SLYNTVATL cosine r same_label_auroc 0.96 +4 SLYNTVATL cosine r same_label_auprc_balanced 0.9651515151515152 +4 SLYNTVATL cosine r precision_at_1 1.0 +4 SLYNTVATL cosine r map_at_r 0.95 +4 SLYNTVATL cosine t within_between_ratio 1.1742857931789124 +4 SLYNTVATL cosine t between_minus_within -0.0485228913031524 +4 SLYNTVATL cosine t same_label_auroc 0.42 +4 SLYNTVATL cosine t same_label_auprc_balanced 0.5704667116509221 +4 SLYNTVATL cosine t precision_at_1 0.8 +4 SLYNTVATL cosine t map_at_r 0.225 +4 SLYNTVATL cosine p within_between_ratio 1.0727657502231038 +4 SLYNTVATL cosine p between_minus_within -0.0030095301249663 +4 SLYNTVATL cosine p same_label_auroc 0.43 +4 SLYNTVATL cosine p same_label_auprc_balanced 0.4691727278569384 +4 SLYNTVATL cosine p precision_at_1 0.0 +4 SLYNTVATL cosine p map_at_r 0.0749999999999999 +4 SLYNTVATL cosine rt within_between_ratio 0.5760962155828022 +4 SLYNTVATL cosine rt between_minus_within 0.1759200280098875 +4 SLYNTVATL cosine rt same_label_auroc 0.8700000000000001 +4 SLYNTVATL cosine rt same_label_auprc_balanced 0.8977272727272726 +4 SLYNTVATL cosine rt precision_at_1 1.0 +4 SLYNTVATL cosine rt map_at_r 0.75 +4 SLYNTVATL cosine rp within_between_ratio 0.4888609802364154 +4 SLYNTVATL cosine rp between_minus_within 0.1829847437593619 +4 SLYNTVATL cosine rp same_label_auroc 0.91 +4 SLYNTVATL cosine rp same_label_auprc_balanced 0.9246503496503496 +4 SLYNTVATL cosine rp precision_at_1 1.0 +4 SLYNTVATL cosine rp map_at_r 0.85 +4 SLYNTVATL cosine tp within_between_ratio 1.030978839230196 +4 SLYNTVATL cosine tp between_minus_within -0.0053958569484331 +4 SLYNTVATL cosine tp same_label_auroc 0.44 +4 SLYNTVATL cosine tp same_label_auprc_balanced 0.5771453108295213 +4 SLYNTVATL cosine tp precision_at_1 0.8 +4 SLYNTVATL cosine tp map_at_r 0.2666666666666666 +4 SLYNTVATL cosine rtp within_between_ratio 0.6335676056698618 +4 SLYNTVATL cosine rtp between_minus_within 0.1260056296589986 +4 SLYNTVATL cosine rtp same_label_auroc 0.8400000000000001 +4 SLYNTVATL cosine rtp same_label_auprc_balanced 0.8827272727272727 +4 SLYNTVATL cosine rtp precision_at_1 1.0 +4 SLYNTVATL cosine rtp map_at_r 0.775 +4 SLYNTVATL euclidean r within_between_ratio 0.6865976748898798 +4 SLYNTVATL euclidean r between_minus_within 2.8328790717609937 +4 SLYNTVATL euclidean r same_label_auroc 0.96 +4 SLYNTVATL euclidean r same_label_auprc_balanced 0.9651515151515152 +4 SLYNTVATL euclidean r precision_at_1 1.0 +4 SLYNTVATL euclidean r map_at_r 0.95 +4 SLYNTVATL euclidean t within_between_ratio 1.0462181651429443 +4 SLYNTVATL euclidean t between_minus_within -0.3722083127557116 +4 SLYNTVATL euclidean t same_label_auroc 0.42 +4 SLYNTVATL euclidean t same_label_auprc_balanced 0.5704667116509221 +4 SLYNTVATL euclidean t precision_at_1 0.8 +4 SLYNTVATL euclidean t map_at_r 0.225 +4 SLYNTVATL euclidean p within_between_ratio 1.0386560265576643 +4 SLYNTVATL euclidean p between_minus_within -0.1045513795820771 +4 SLYNTVATL euclidean p same_label_auroc 0.43 +4 SLYNTVATL euclidean p same_label_auprc_balanced 0.4691727278569384 +4 SLYNTVATL euclidean p precision_at_1 0.0 +4 SLYNTVATL euclidean p map_at_r 0.0749999999999999 +4 SLYNTVATL euclidean rt within_between_ratio 0.7486457046270609 +4 SLYNTVATL euclidean rt between_minus_within 2.429093726497872 +4 SLYNTVATL euclidean rt same_label_auroc 0.87 +4 SLYNTVATL euclidean rt same_label_auprc_balanced 0.8961247086247085 +4 SLYNTVATL euclidean rt precision_at_1 1.0 +4 SLYNTVATL euclidean rt map_at_r 0.75 +4 SLYNTVATL euclidean rp within_between_ratio 0.6938813768447298 +4 SLYNTVATL euclidean rp between_minus_within 2.606095320464502 +4 SLYNTVATL euclidean rp same_label_auroc 0.91 +4 SLYNTVATL euclidean rp same_label_auprc_balanced 0.9246503496503496 +4 SLYNTVATL euclidean rp precision_at_1 1.0 +4 SLYNTVATL euclidean rp map_at_r 0.85 +4 SLYNTVATL euclidean tp within_between_ratio 0.9867719221564634 +4 SLYNTVATL euclidean tp between_minus_within 0.0847202375826423 +4 SLYNTVATL euclidean tp same_label_auroc 0.44 +4 SLYNTVATL euclidean tp same_label_auprc_balanced 0.5771453108295213 +4 SLYNTVATL euclidean tp precision_at_1 0.8 +4 SLYNTVATL euclidean tp map_at_r 0.2666666666666666 +4 SLYNTVATL euclidean rtp within_between_ratio 0.7918103110184312 +4 SLYNTVATL euclidean rtp between_minus_within 1.7841010079875463 +4 SLYNTVATL euclidean rtp same_label_auroc 0.8400000000000001 +4 SLYNTVATL euclidean rtp same_label_auprc_balanced 0.8827272727272727 +4 SLYNTVATL euclidean rtp precision_at_1 1.0 +4 SLYNTVATL euclidean rtp map_at_r 0.775 +5 ALSKGVHFV cosine r within_between_ratio 0.3866865701359755 +5 ALSKGVHFV cosine r between_minus_within 0.2910739948539576 +5 ALSKGVHFV cosine r same_label_auroc 1.0 +5 ALSKGVHFV cosine r same_label_auprc_balanced 1.0 +5 ALSKGVHFV cosine r precision_at_1 1.0 +5 ALSKGVHFV cosine r map_at_r 1.0 +5 ALSKGVHFV cosine t within_between_ratio 0.2201952872084881 +5 ALSKGVHFV cosine t between_minus_within 0.1352382523939194 +5 ALSKGVHFV cosine t same_label_auroc 0.9722222222222224 +5 ALSKGVHFV cosine t same_label_auprc_balanced 0.976190476190476 +5 ALSKGVHFV cosine t precision_at_1 0.75 +5 ALSKGVHFV cosine t map_at_r 0.5972222222222222 +5 ALSKGVHFV cosine p within_between_ratio 0.7921547385522357 +5 ALSKGVHFV cosine p between_minus_within 0.0084317881806044 +5 ALSKGVHFV cosine p same_label_auroc 0.4722222222222222 +5 ALSKGVHFV cosine p same_label_auprc_balanced 0.5983164983164984 +5 ALSKGVHFV cosine p precision_at_1 0.0 +5 ALSKGVHFV cosine p map_at_r 0.0416666666666666 +5 ALSKGVHFV cosine rt within_between_ratio 0.3962952790798194 +5 ALSKGVHFV cosine rt between_minus_within 0.2712548608840007 +5 ALSKGVHFV cosine rt same_label_auroc 1.0 +5 ALSKGVHFV cosine rt same_label_auprc_balanced 1.0 +5 ALSKGVHFV cosine rt precision_at_1 1.0 +5 ALSKGVHFV cosine rt map_at_r 1.0 +5 ALSKGVHFV cosine rp within_between_ratio 0.3960752750678162 +5 ALSKGVHFV cosine rp between_minus_within 0.2488982920455894 +5 ALSKGVHFV cosine rp same_label_auroc 1.0 +5 ALSKGVHFV cosine rp same_label_auprc_balanced 1.0 +5 ALSKGVHFV cosine rp precision_at_1 1.0 +5 ALSKGVHFV cosine rp map_at_r 1.0 +5 ALSKGVHFV cosine tp within_between_ratio 0.3950118209766391 +5 ALSKGVHFV cosine tp between_minus_within 0.0890139005725051 +5 ALSKGVHFV cosine tp same_label_auroc 0.9444444444444444 +5 ALSKGVHFV cosine tp same_label_auprc_balanced 0.9583333333333331 +5 ALSKGVHFV cosine tp precision_at_1 1.0 +5 ALSKGVHFV cosine tp map_at_r 0.6666666666666666 +5 ALSKGVHFV cosine rtp within_between_ratio 0.4186267993922287 +5 ALSKGVHFV cosine rtp between_minus_within 0.2138412621092893 +5 ALSKGVHFV cosine rtp same_label_auroc 1.0 +5 ALSKGVHFV cosine rtp same_label_auprc_balanced 1.0 +5 ALSKGVHFV cosine rtp precision_at_1 1.0 +5 ALSKGVHFV cosine rtp map_at_r 1.0 +5 ALSKGVHFV euclidean r within_between_ratio 0.6251015024335775 +5 ALSKGVHFV euclidean r between_minus_within 3.838027858236756 +5 ALSKGVHFV euclidean r same_label_auroc 1.0 +5 ALSKGVHFV euclidean r same_label_auprc_balanced 1.0 +5 ALSKGVHFV euclidean r precision_at_1 1.0 +5 ALSKGVHFV euclidean r map_at_r 1.0 +5 ALSKGVHFV euclidean t within_between_ratio 0.4450798133529459 +5 ALSKGVHFV euclidean t between_minus_within 3.637365859970788 +5 ALSKGVHFV euclidean t same_label_auroc 0.9722222222222224 +5 ALSKGVHFV euclidean t same_label_auprc_balanced 0.976190476190476 +5 ALSKGVHFV euclidean t precision_at_1 0.75 +5 ALSKGVHFV euclidean t map_at_r 0.5972222222222222 +5 ALSKGVHFV euclidean p within_between_ratio 0.9725661204439372 +5 ALSKGVHFV euclidean p between_minus_within 0.0751790534833554 +5 ALSKGVHFV euclidean p same_label_auroc 0.4444444444444444 +5 ALSKGVHFV euclidean p same_label_auprc_balanced 0.514983164983165 +5 ALSKGVHFV euclidean p precision_at_1 0.0 +5 ALSKGVHFV euclidean p map_at_r 0.0416666666666666 +5 ALSKGVHFV euclidean rt within_between_ratio 0.6210177573640201 +5 ALSKGVHFV euclidean rt between_minus_within 3.774737004163465 +5 ALSKGVHFV euclidean rt same_label_auroc 1.0 +5 ALSKGVHFV euclidean rt same_label_auprc_balanced 1.0 +5 ALSKGVHFV euclidean rt precision_at_1 1.0 +5 ALSKGVHFV euclidean rt map_at_r 1.0 +5 ALSKGVHFV euclidean rp within_between_ratio 0.6226201452868303 +5 ALSKGVHFV euclidean rp between_minus_within 3.4502983786053063 +5 ALSKGVHFV euclidean rp same_label_auroc 1.0 +5 ALSKGVHFV euclidean rp same_label_auprc_balanced 1.0 +5 ALSKGVHFV euclidean rp precision_at_1 1.0 +5 ALSKGVHFV euclidean rp map_at_r 1.0 +5 ALSKGVHFV euclidean tp within_between_ratio 0.6149599833028031 +5 ALSKGVHFV euclidean tp between_minus_within 2.342975754521881 +5 ALSKGVHFV euclidean tp same_label_auroc 0.9444444444444444 +5 ALSKGVHFV euclidean tp same_label_auprc_balanced 0.9583333333333331 +5 ALSKGVHFV euclidean tp precision_at_1 1.0 +5 ALSKGVHFV euclidean tp map_at_r 0.6666666666666666 +5 ALSKGVHFV euclidean rtp within_between_ratio 0.642568503402784 +5 ALSKGVHFV euclidean rtp between_minus_within 3.1288563750265324 +5 ALSKGVHFV euclidean rtp same_label_auroc 1.0 +5 ALSKGVHFV euclidean rtp same_label_auprc_balanced 1.0 +5 ALSKGVHFV euclidean rtp precision_at_1 1.0 +5 ALSKGVHFV euclidean rtp map_at_r 1.0 +6 FLYALALLL cosine r within_between_ratio 0.6775292127221834 +6 FLYALALLL cosine r between_minus_within 0.2105191813140541 +6 FLYALALLL cosine r same_label_auroc 1.0 +6 FLYALALLL cosine r same_label_auprc_balanced 1.0 +6 FLYALALLL cosine r precision_at_1 0.0 +6 FLYALALLL cosine r map_at_r 0.0 +6 FLYALALLL cosine t within_between_ratio 0.7714240609847476 +6 FLYALALLL cosine t between_minus_within 0.1560307899158284 +6 FLYALALLL cosine t same_label_auroc 1.0 +6 FLYALALLL cosine t same_label_auprc_balanced 1.0 +6 FLYALALLL cosine t precision_at_1 0.0 +6 FLYALALLL cosine t map_at_r 0.0 +6 FLYALALLL cosine p within_between_ratio 1.7646471937707764 +6 FLYALALLL cosine p between_minus_within -0.0382589357909233 +6 FLYALALLL cosine p same_label_auroc 0.0 +6 FLYALALLL cosine p same_label_auprc_balanced 0.5 +6 FLYALALLL cosine p precision_at_1 0.0 +6 FLYALALLL cosine p map_at_r 0.0 +6 FLYALALLL cosine rt within_between_ratio 0.6906109655922215 +6 FLYALALLL cosine rt between_minus_within 0.1996168469751401 +6 FLYALALLL cosine rt same_label_auroc 1.0 +6 FLYALALLL cosine rt same_label_auprc_balanced 1.0 +6 FLYALALLL cosine rt precision_at_1 0.0 +6 FLYALALLL cosine rt map_at_r 0.0 +6 FLYALALLL cosine rp within_between_ratio 0.6707672271716343 +6 FLYALALLL cosine rp between_minus_within 0.1762362235339419 +6 FLYALALLL cosine rp same_label_auroc 1.0 +6 FLYALALLL cosine rp same_label_auprc_balanced 1.0 +6 FLYALALLL cosine rp precision_at_1 0.0 +6 FLYALALLL cosine rp map_at_r 0.0 +6 FLYALALLL cosine tp within_between_ratio 0.9508040033287974 +6 FLYALALLL cosine tp between_minus_within 0.0226730347800806 +6 FLYALALLL cosine tp same_label_auroc 1.0 +6 FLYALALLL cosine tp same_label_auprc_balanced 1.0 +6 FLYALALLL cosine tp precision_at_1 0.0 +6 FLYALALLL cosine tp map_at_r 0.0 +6 FLYALALLL cosine rtp within_between_ratio 0.7140904494680264 +6 FLYALALLL cosine rtp between_minus_within 0.1802256409848623 +6 FLYALALLL cosine rtp same_label_auroc 1.0 +6 FLYALALLL cosine rtp same_label_auprc_balanced 1.0 +6 FLYALALLL cosine rtp precision_at_1 0.0 +6 FLYALALLL cosine rtp map_at_r 0.0 +6 FLYALALLL euclidean r within_between_ratio 0.833307428600523 +6 FLYALALLL euclidean r between_minus_within 1.990730964832533 +6 FLYALALLL euclidean r same_label_auroc 1.0 +6 FLYALALLL euclidean r same_label_auprc_balanced 1.0 +6 FLYALALLL euclidean r precision_at_1 0.0 +6 FLYALALLL euclidean r map_at_r 0.0 +6 FLYALALLL euclidean t within_between_ratio 0.8712170925455438 +6 FLYALALLL euclidean t between_minus_within 1.7187362601360745 +6 FLYALALLL euclidean t same_label_auroc 1.0 +6 FLYALALLL euclidean t same_label_auprc_balanced 1.0 +6 FLYALALLL euclidean t precision_at_1 0.0 +6 FLYALALLL euclidean t map_at_r 0.0 +6 FLYALALLL euclidean p within_between_ratio 1.3238794339624518 +6 FLYALALLL euclidean p between_minus_within -1.1806467313926985 +6 FLYALALLL euclidean p same_label_auroc 0.0 +6 FLYALALLL euclidean p same_label_auprc_balanced 0.5 +6 FLYALALLL euclidean p precision_at_1 0.0 +6 FLYALALLL euclidean p map_at_r 0.0 +6 FLYALALLL euclidean rt within_between_ratio 0.8318839208376045 +6 FLYALALLL euclidean rt between_minus_within 2.0133576493387064 +6 FLYALALLL euclidean rt same_label_auroc 1.0 +6 FLYALALLL euclidean rt same_label_auprc_balanced 1.0 +6 FLYALALLL euclidean rt precision_at_1 0.0 +6 FLYALALLL euclidean rt map_at_r 0.0 +6 FLYALALLL euclidean rp within_between_ratio 0.8250319227075794 +6 FLYALALLL euclidean rp between_minus_within 1.819457314152068 +6 FLYALALLL euclidean rp same_label_auroc 1.0 +6 FLYALALLL euclidean rp same_label_auprc_balanced 1.0 +6 FLYALALLL euclidean rp precision_at_1 0.0 +6 FLYALALLL euclidean rp map_at_r 0.0 +6 FLYALALLL euclidean tp within_between_ratio 0.9745461896386142 +6 FLYALALLL euclidean tp between_minus_within 0.278119392422008 +6 FLYALALLL euclidean tp same_label_auroc 1.0 +6 FLYALALLL euclidean tp same_label_auprc_balanced 1.0 +6 FLYALALLL euclidean tp precision_at_1 0.0 +6 FLYALALLL euclidean tp map_at_r 0.0 +6 FLYALALLL euclidean rtp within_between_ratio 0.845959006340635 +6 FLYALALLL euclidean rtp between_minus_within 1.8007589492959664 +6 FLYALALLL euclidean rtp same_label_auroc 1.0 +6 FLYALALLL euclidean rtp same_label_auprc_balanced 1.0 +6 FLYALALLL euclidean rtp precision_at_1 0.0 +6 FLYALALLL euclidean rtp map_at_r 0.0 diff --git a/results/author/downstream_rtp_minus_rp_effects.tsv b/results/author/downstream_rtp_minus_rp_effects.tsv new file mode 100644 index 0000000..af9ab26 --- /dev/null +++ b/results/author/downstream_rtp_minus_rp_effects.tsv @@ -0,0 +1,73 @@ +rank epitope distance metric p r rp rt rtp t tp rtp_minus_rp +1 GILGFVFTL cosine between_minus_within 0.0206568423984859 0.1233428152129985 0.1202215255443032 0.1070347350523491 0.110323732416114 0.1538348010838437 0.137600465097264 -0.0098977931281891 +1 GILGFVFTL cosine map_at_r 0.7469039220366203 0.8377571417671207 0.8382527018021324 0.8290056068583498 0.835703226039053 0.8391364170281714 0.8556084509363918 -0.0025494757630794 +1 GILGFVFTL cosine precision_at_1 0.8982683982683982 0.9956709956709956 0.9935064935064936 0.9956709956709956 0.987012987012987 0.9913419913419912 0.9891774891774892 -0.0064935064935065 +1 GILGFVFTL cosine same_label_auprc_balanced 0.521491273229804 0.7299518893107003 0.7174927693471926 0.7050656379180732 0.7144508771003788 0.674158188636663 0.7053104436509567 -0.0030418922468138 +1 GILGFVFTL cosine same_label_auroc 0.53763092 0.73158876 0.7257904800000001 0.70806996 0.71991092 0.71333396 0.75484744 -0.0058795600000001 +1 GILGFVFTL cosine within_between_ratio 0.8230982847313358 0.814834956105159 0.8013631708826814 0.8470954060378486 0.829036865860044 0.7684862592239499 0.7243638016662785 0.0276736949773626 +1 GILGFVFTL euclidean between_minus_within 0.4762456869122298 1.2324917821683776 1.1943521196618183 1.0255226173407426 1.0340184856414627 1.7771783413246125 1.764684817162717 -0.1603336340203558 +1 GILGFVFTL euclidean map_at_r 0.7478342545035729 0.8403814016170456 0.8390079506829614 0.8272032472168177 0.8320969033377764 0.8439129645545131 0.8574148540779141 -0.0069110473451849 +1 GILGFVFTL euclidean precision_at_1 0.8982683982683982 0.9956709956709956 0.9956709956709956 0.9956709956709956 0.987012987012987 0.9913419913419912 0.9891774891774892 -0.0086580086580086 +1 GILGFVFTL euclidean same_label_auprc_balanced 0.5236232343171925 0.7351607337112969 0.7211318096286985 0.7064940862833436 0.7101946466646145 0.6769013347021693 0.7068407684952329 -0.0109371629640839 +1 GILGFVFTL euclidean same_label_auroc 0.54044212 0.73664052 0.7277302 0.70650972 0.7109821999999999 0.71670196 0.75650096 -0.016748 +1 GILGFVFTL euclidean within_between_ratio 0.8973830858007207 0.8980290222030745 0.8919376258759846 0.9181802458439324 0.9120983001562664 0.8627014070154413 0.8430098768730618 0.0201606742802817 +2 GLCTLVAML cosine between_minus_within 0.1346310777629427 0.0469673049657198 0.0537090332683131 0.1212739293962659 0.1087192530041349 0.2327473467424893 0.1461652305207085 0.0550102197358218 +2 GLCTLVAML cosine map_at_r 0.3971554962427943 0.2960349214300812 0.30432489044198 0.3897568779087489 0.3668666020143697 0.5281064277718629 0.4710206617825163 0.0625417115723897 +2 GLCTLVAML cosine precision_at_1 0.7391304347826086 0.9782608695652174 0.9782608695652174 0.9782608695652174 0.9782608695652174 0.9347826086956522 0.9347826086956522 0.0 +2 GLCTLVAML cosine same_label_auprc_balanced 0.764362746527017 0.6220477286785298 0.6241568998359073 0.7091789691940154 0.6999856842884669 0.7541585193205 0.7144185257602391 0.0758287844525595 +2 GLCTLVAML cosine same_label_auroc 0.8133053280123224 0.5325370487059209 0.5408042194683657 0.6504338490979953 0.6468519685406894 0.7486634460547503 0.7091759434292515 0.1060477490723237 +2 GLCTLVAML cosine within_between_ratio 0.2016262709941349 0.9304252886506998 0.9159199337758568 0.8290567233595127 0.8395844652525207 0.6904286297927523 0.7479251171475703 -0.076335468523336 +2 GLCTLVAML euclidean between_minus_within 3.49534116615842 0.643298441606289 0.7060268471358278 1.4121366935531796 1.3297842074997224 3.009821017983443 2.042914623086965 0.6237573603638946 +2 GLCTLVAML euclidean map_at_r 0.3968899090179063 0.2969725401974061 0.3085944081608493 0.4009350926126311 0.3851989042931715 0.525644495113936 0.4632930047916605 0.0766044961323221 +2 GLCTLVAML euclidean precision_at_1 0.7391304347826086 0.9782608695652174 0.9782608695652174 0.9782608695652174 0.9782608695652174 0.9347826086956522 0.9347826086956522 0.0 +2 GLCTLVAML euclidean same_label_auprc_balanced 0.7625697321880942 0.6278778128933347 0.6330275638221696 0.7209395050987346 0.7184602005371521 0.7513258188815857 0.710591732184743 0.0854326367149824 +2 GLCTLVAML euclidean same_label_auroc 0.810628952834372 0.543551541459544 0.5572181381129082 0.6718961936101193 0.6801923032042755 0.7435272701813344 0.7049527410207939 0.1229741650913672 +2 GLCTLVAML euclidean within_between_ratio 0.4021732315244646 0.9470494666132306 0.9376119790475408 0.8871695287309327 0.8893726101341216 0.7836093129434006 0.832802537212801 -0.0482393689134192 +3 YLQPRTFLL cosine between_minus_within 0.0228717270199522 0.0446161337209015 0.0366156392001849 0.0844520133839246 0.0486971483198341 0.019922715432474 0.0160828353071133 0.0120815091196492 +3 YLQPRTFLL cosine map_at_r 0.1395447798181916 0.2966054624516153 0.29964836781615 0.4070361074886092 0.3691817079869173 0.292351315785512 0.2987907026485508 0.0695333401707673 +3 YLQPRTFLL cosine precision_at_1 0.3 1.0 1.0 0.975 0.975 0.825 0.825 -0.025 +3 YLQPRTFLL cosine same_label_auprc_balanced 0.557005125811634 0.6199786587907083 0.5943170885404121 0.6779547121190657 0.6137564655276165 0.5862232348035362 0.5957361011159912 0.0194393769872044 +3 YLQPRTFLL cosine same_label_auroc 0.5686582639714625 0.5777331747919143 0.5640332936979786 0.6346615933412604 0.5801588585017836 0.5323434007134364 0.5300889417360286 0.016125564803805 +3 YLQPRTFLL cosine within_between_ratio 0.8128660411253578 0.9234808526023524 0.9310489264658908 0.8632042895855496 0.9120239325822764 0.9523943494125764 0.9533213775412188 -0.0190249938836142 +3 YLQPRTFLL euclidean between_minus_within 0.6439082943692966 0.5335972980226611 0.4338869883892666 0.9411507264122356 0.6197374939630098 0.4112640325394672 0.413196233338418 0.1858505055737431 +3 YLQPRTFLL euclidean map_at_r 0.1388273830698967 0.2957945422594127 0.2988144296423183 0.4103746473365836 0.3804617652468797 0.2879492583841253 0.2989943295932016 0.0816473356045614 +3 YLQPRTFLL euclidean precision_at_1 0.275 1.0 1.0 0.975 0.975 0.825 0.825 -0.025 +3 YLQPRTFLL euclidean same_label_auprc_balanced 0.5538283154450125 0.6232151785513477 0.5968576844192361 0.6898694962560761 0.6267377767138391 0.5835250342709797 0.595362675292274 0.0298800922946029 +3 YLQPRTFLL euclidean same_label_auroc 0.566814744351962 0.5835319857312723 0.5674882282996433 0.6455286563614744 0.596221640903686 0.5286753864447087 0.529468727705113 0.0287334126040427 +3 YLQPRTFLL euclidean within_between_ratio 0.8696962900330723 0.9530694427510062 0.9583005918155642 0.9200626838923368 0.9433479558101976 0.9589057298754774 0.9551444520774234 -0.0149526360053664 +4 SLYNTVATL cosine between_minus_within -0.0030095301249663 0.1923712751303949 0.1829847437593619 0.1759200280098875 0.1260056296589986 -0.0485228913031524 -0.0053958569484331 -0.0569791141003633 +4 SLYNTVATL cosine map_at_r 0.0749999999999999 0.95 0.85 0.75 0.775 0.225 0.2666666666666666 -0.0749999999999999 +4 SLYNTVATL cosine precision_at_1 0.0 1.0 1.0 1.0 1.0 0.8 0.8 0.0 +4 SLYNTVATL cosine same_label_auprc_balanced 0.4691727278569384 0.9651515151515152 0.9246503496503496 0.8977272727272726 0.8827272727272727 0.5704667116509221 0.5771453108295213 -0.0419230769230769 +4 SLYNTVATL cosine same_label_auroc 0.43 0.96 0.91 0.8700000000000001 0.8400000000000001 0.42 0.44 -0.0699999999999998 +4 SLYNTVATL cosine within_between_ratio 1.0727657502231038 0.4832409495705956 0.4888609802364154 0.5760962155828022 0.6335676056698618 1.1742857931789124 1.030978839230196 0.1447066254334464 +4 SLYNTVATL euclidean between_minus_within -0.1045513795820771 2.8328790717609937 2.606095320464502 2.429093726497872 1.7841010079875463 -0.3722083127557116 0.0847202375826423 -0.8219943124769555 +4 SLYNTVATL euclidean map_at_r 0.0749999999999999 0.95 0.85 0.75 0.775 0.225 0.2666666666666666 -0.0749999999999999 +4 SLYNTVATL euclidean precision_at_1 0.0 1.0 1.0 1.0 1.0 0.8 0.8 0.0 +4 SLYNTVATL euclidean same_label_auprc_balanced 0.4691727278569384 0.9651515151515152 0.9246503496503496 0.8961247086247085 0.8827272727272727 0.5704667116509221 0.5771453108295213 -0.0419230769230769 +4 SLYNTVATL euclidean same_label_auroc 0.43 0.96 0.91 0.87 0.8400000000000001 0.42 0.44 -0.0699999999999998 +4 SLYNTVATL euclidean within_between_ratio 1.0386560265576643 0.6865976748898798 0.6938813768447298 0.7486457046270609 0.7918103110184312 1.0462181651429443 0.9867719221564634 0.0979289341737014 +5 ALSKGVHFV cosine between_minus_within 0.0084317881806044 0.2910739948539576 0.2488982920455894 0.2712548608840007 0.2138412621092893 0.1352382523939194 0.0890139005725051 -0.0350570299363 +5 ALSKGVHFV cosine map_at_r 0.0416666666666666 1.0 1.0 1.0 1.0 0.5972222222222222 0.6666666666666666 0.0 +5 ALSKGVHFV cosine precision_at_1 0.0 1.0 1.0 1.0 1.0 0.75 1.0 0.0 +5 ALSKGVHFV cosine same_label_auprc_balanced 0.5983164983164984 1.0 1.0 1.0 1.0 0.976190476190476 0.9583333333333331 0.0 +5 ALSKGVHFV cosine same_label_auroc 0.4722222222222222 1.0 1.0 1.0 1.0 0.9722222222222224 0.9444444444444444 0.0 +5 ALSKGVHFV cosine within_between_ratio 0.7921547385522357 0.3866865701359755 0.3960752750678162 0.3962952790798194 0.4186267993922287 0.2201952872084881 0.3950118209766391 0.0225515243244124 +5 ALSKGVHFV euclidean between_minus_within 0.0751790534833554 3.838027858236756 3.4502983786053063 3.774737004163465 3.1288563750265324 3.637365859970788 2.342975754521881 -0.3214420035787739 +5 ALSKGVHFV euclidean map_at_r 0.0416666666666666 1.0 1.0 1.0 1.0 0.5972222222222222 0.6666666666666666 0.0 +5 ALSKGVHFV euclidean precision_at_1 0.0 1.0 1.0 1.0 1.0 0.75 1.0 0.0 +5 ALSKGVHFV euclidean same_label_auprc_balanced 0.514983164983165 1.0 1.0 1.0 1.0 0.976190476190476 0.9583333333333331 0.0 +5 ALSKGVHFV euclidean same_label_auroc 0.4444444444444444 1.0 1.0 1.0 1.0 0.9722222222222224 0.9444444444444444 0.0 +5 ALSKGVHFV euclidean within_between_ratio 0.9725661204439372 0.6251015024335775 0.6226201452868303 0.6210177573640201 0.642568503402784 0.4450798133529459 0.6149599833028031 0.0199483581159537 +6 FLYALALLL cosine between_minus_within -0.0382589357909233 0.2105191813140541 0.1762362235339419 0.1996168469751401 0.1802256409848623 0.1560307899158284 0.0226730347800806 0.0039894174509204 +6 FLYALALLL cosine map_at_r 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 +6 FLYALALLL cosine precision_at_1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 +6 FLYALALLL cosine same_label_auprc_balanced 0.5 1.0 1.0 1.0 1.0 1.0 1.0 0.0 +6 FLYALALLL cosine same_label_auroc 0.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 +6 FLYALALLL cosine within_between_ratio 1.7646471937707764 0.6775292127221834 0.6707672271716343 0.6906109655922215 0.7140904494680264 0.7714240609847476 0.9508040033287974 0.043323222296392 +6 FLYALALLL euclidean between_minus_within -1.1806467313926985 1.990730964832533 1.819457314152068 2.0133576493387064 1.8007589492959664 1.7187362601360745 0.278119392422008 -0.0186983648561014 +6 FLYALALLL euclidean map_at_r 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 +6 FLYALALLL euclidean precision_at_1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 +6 FLYALALLL euclidean same_label_auprc_balanced 0.5 1.0 1.0 1.0 1.0 1.0 1.0 0.0 +6 FLYALALLL euclidean same_label_auroc 0.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 +6 FLYALALLL euclidean within_between_ratio 1.3238794339624518 0.833307428600523 0.8250319227075794 0.8318839208376045 0.845959006340635 0.8712170925455438 0.9745461896386142 0.0209270836330556 diff --git a/results/author/external_author_global_comparison.md b/results/author/external_author_global_comparison.md new file mode 100644 index 0000000..e23a9fe --- /dev/null +++ b/results/author/external_author_global_comparison.md @@ -0,0 +1,10 @@ +| Method | Modality | Reconstruction exact | pgen RMSE | VDJdb P@1 | P@5 | P@10 | MAP | Within/between | +|---|---|---|---|---|---|---|---|---| +| DATA-ANCHOR RTP | sequence-only | 0.8831 +/- 0.0058 | 0.2375 +/- 0.0008 | 0.3879 | 0.3272 | 0.2940 | 0.3136 | 0.9526 | +| TCRemP prototype distances | sequence-derived alignment | 0.5181 +/- 0.0024 | 0.6038 +/- 0.0019 | NA | NA | NA | NA | NA | +| TCR-BERT final residue mean | sequence-only | 0.4778 +/- 0.0179 | 0.4299 +/- 0.0022 | 0.3590 | 0.3053 | 0.2735 | 0.2999 | 0.9722 | +| ESM2-35M final residue mean | sequence-only | 0.2042 +/- 0.0132 | 0.7086 +/- 0.0033 | 0.3597 | 0.3081 | 0.2782 | 0.2968 | 0.9410 | +| ESM2-8M final residue mean | sequence-only sensitivity | 0.2179 +/- 0.0075 | 0.7439 +/- 0.0012 | NA | NA | NA | NA | NA | +| SCEPTR CDR3-only native CLS | sequence-only | 0.8213 +/- 0.0043 | 0.6637 +/- 0.0062 | 0.3553 | 0.2778 | 0.2366 | 0.2651 | 1.0469 | +| SCEPTR default native CLS | annotation-aware V+CDR3+J | 0.7263 +/- 0.0125 | 0.8073 +/- 0.0084 | NA | NA | NA | NA | NA | +| CDR3 TF-IDF ridge | sequence control | NA | 0.7860 | NA | NA | NA | NA | NA | diff --git a/results/author/external_author_global_comparison.tsv b/results/author/external_author_global_comparison.tsv new file mode 100644 index 0000000..418179b --- /dev/null +++ b/results/author/external_author_global_comparison.tsv @@ -0,0 +1,9 @@ +condition method modality native_global_dim reconstruction_common_dim reconstruction_exact_mean reconstruction_exact_std reconstruction_token_accuracy_mean pgen_rmse_mean pgen_rmse_std pgen_r2_mean pgen_spearman_mean vdjdb_cosine_precision_at_1 vdjdb_cosine_precision_at_5 vdjdb_cosine_precision_at_10 vdjdb_cosine_map vdjdb_cosine_within_between_ratio vdjdb_cosine_knn_macro_f1_at_5 reconstruction_exact pgen_rmse vdjdb_p_at_1 vdjdb_p_at_5 vdjdb_p_at_10 vdjdb_map vdjdb_within_between_ratio +rtp DATA-ANCHOR RTP sequence-only 128.0 64.0 0.8831338630192095 0.0058150137887855 0.9810685896331252 0.2375036700301933 0.0007641427238683 0.9910726990213594 0.9958299890234856 0.38789317660994344 0.32720557444226145 0.2939801617089443 0.31364862922705045 0.9525868556453406 0.2744565109481644 0.8831 +/- 0.0058 0.2375 +/- 0.0008 0.3879 0.3272 0.2940 0.3136 0.9526 +tcremp TCRemP prototype distances sequence-derived alignment 9000.0 64.0 0.518052901538771 0.0024158527908004 0.9371434757441994 0.6038112324768586 0.0019115151879035455 0.9422992036297707 0.9650743522262878 0.5181 +/- 0.0024 0.6038 +/- 0.0019 NA NA NA NA NA +tcr_bert TCR-BERT final residue mean sequence-only 768.0 64.0 0.4777565456435013 0.0178585766284086 0.90975931989201 0.4298724150079793 0.0022051496033906 0.9707542136303974 0.9863067900009792 0.358981305173744 0.3053478420625738 0.2734693644304071 0.2998904226736852 0.9721619846959844 0.2582267395006297 0.4778 +/- 0.0179 0.4299 +/- 0.0022 0.3590 0.3053 0.2735 0.2999 0.9722 +esm2_35m ESM2-35M final residue mean sequence-only 480.0 64.0 0.20416373327969425 0.013152045845910268 0.8082159095428093 0.7086315695015749 0.0033364098056348564 0.9205262070492487 0.9632367899599094 0.3597203532617486 0.3080627550986812 0.2781556804093438 0.2968139690814502 0.9410002050511648 0.2406928227557169 0.2042 +/- 0.0132 0.7086 +/- 0.0033 0.3597 0.3081 0.2782 0.2968 0.9410 +esm2_8m ESM2-8M final residue mean sequence-only sensitivity 320.0 64.0 0.2179087465218411 0.0075197155491379 0.8092609039533044 0.7439029042798556 0.0012426103217311 0.9124190107532284 0.9588768535120026 0.2179 +/- 0.0075 0.7439 +/- 0.0012 NA NA NA NA NA +sceptr_cdr3 SCEPTR CDR3-only native CLS sequence-only 64.0 64.0 0.8213148278520902 0.004335551696797592 0.9707998073498677 0.6636995360955414 0.006248417396838993 0.9302819519415871 0.9732935578951564 0.355322581407025 0.2777975134354664 0.2366489373866297 0.2651021244243266 1.0469384621629017 0.217804653528936 0.8213 +/- 0.0043 0.6637 +/- 0.0062 0.3553 0.2778 0.2366 0.2651 1.0469 +sceptr SCEPTR default native CLS annotation-aware V+CDR3+J 64.0 64.0 0.7263401387911094 0.0125148065817577 0.9556794231100624 0.8072807829525773 0.0083542304525543 0.8968529612267799 0.95929616602314 0.7263 +/- 0.0125 0.8073 +/- 0.0084 NA NA NA NA NA +tfidf CDR3 TF-IDF ridge sequence control 0.785974434572495 0.9022327566653896 0.9653913662774836 NA 0.7860 NA NA NA NA NA diff --git a/results/author/external_author_global_comparison_manifest.json b/results/author/external_author_global_comparison_manifest.json new file mode 100644 index 0000000..8f43089 --- /dev/null +++ b/results/author/external_author_global_comparison_manifest.json @@ -0,0 +1,25 @@ +{ + "inputs": { + "/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/vdjdb/evaluation-1427581/external_retrieval_summary.tsv": "58ec7b647ea881689495770f1a2bbe94b90184bd8e67f143aca96ada4909393a", + "/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/results/pgen/pgen_summary.csv": "1be5e255e91c0b55f2c7429aeb48b75ce332af524fee17779f43070a88b8a362", + "/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/results/reconstruction/reconstruction_summary.csv": "5f15ae270e153ef179cddb8651d8e7f75d5ceb3b77aa8c66c48c66fc5d0739d4", + "/home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728/metrics.json": "8a460d8f870f6dd5733c6cbf41c7a1614ba1e3729235169da739ae9a20492d56" + }, + "limitations": [ + "SCEPTR-default is annotation-aware and is not equivalent to sequence-only methods.", + "SCEPTR-default was not run on full REDCEA because unsupported V/J rows would require method-specific filtering.", + "REDCEA membership is sequence/TCRemP-cluster-conditioned, so retrieval is motif-conditioned.", + "Aligned one-hot is retained as an internal high-dimensional sequence control but omitted from the manuscript-facing comparison.", + "Probe mean/std reflects probe seeds, not pretrained encoder training-seed uncertainty." + ], + "outputs": { + "final_author_global_comparison.md": "b1d4019d103b5d545f3bc46f99212dadd99f73482dc9c8ec6df456d3c0382f97", + "final_author_global_comparison.tsv": "d15807f7d1401f39b10b3aa3a4364d70d9db9c87c0fe238b92c6ee2f0e9cf3b3" + }, + "protocol": { + "pgen": "native author-global vector plus identical MLP; three probe seeds", + "reconstruction": "train-only PCA64 plus identical learned decoder; three probe seeds", + "vdjdb": "native author-global vector on the fixed REDCEA gallery; macro over 87 epitopes" + }, + "status": "complete" +} diff --git a/results/author/external_pgen_results.tsv b/results/author/external_pgen_results.tsv new file mode 100644 index 0000000..482eb3c --- /dev/null +++ b/results/author/external_pgen_results.tsv @@ -0,0 +1,6 @@ +target train_subset arm seeds train_size rmse_mean rmse_std mae_mean mae_std r2_mean r2_std pearson_r_mean pearson_r_std spearman_rho_mean spearman_rho_std bias_mean bias_std training_seconds_mean inference_seq_per_second_mean peak_rss_mb_mean peak_gpu_mb_mean +log10_pgen_1mm all rtp_mlp 3 79544 0.2375036700301933 0.0007641427238683 0.1643160829299884 0.0004204611887859 0.9910726990213594 5.749766422037862e-05 0.9955328156356806 2.5301945949134647e-05 0.9958299890234856 3.869003023257598e-05 -0.0063614919328252 0.0031365612188622 74.25666666666666 828631.1210576348 1053.1158854166667 21.77001953125 +log10_pgen_1mm all tcr_bert_mlp 3 79544 0.4298724150079793 0.0022051496033906 0.3040288464942088 0.0013097794566628 0.9707542136303974 0.0003003689953682 0.9853613777329778 0.0001596107640185 0.9863067900009792 8.707651289475979e-05 0.0143588582656017 0.003591457594319 81.21 347583.1589531299 1549.37890625 28.14501953125 +log10_pgen_1mm all esm2_8m_mlp 3 79544 0.7439029042798556 0.0012426103217311 0.5208345937732529 0.0009216602240305 0.9124190107532284 0.000292707843767 0.9559005092380672 0.0001831596672659 0.9588768535120026 0.0002171379827709 0.0663262720525637 0.0016679294577322 72.43 633788.4195212495 1127.73828125 23.45751953125 +log10_pgen_1mm all tfidf_ridge 1 79544 0.785974434572495 0.5622312803788474 0.9022327566653896 0.9498971636908724 0.9653913662774836 0.0069710814573569 9.0 47475.47340859601 554.30859375 +log10_pgen_1mm all sceptr_mlp 3 79544 0.8072807829525773 0.0083542304525543 0.536926930754249 0.00316616759605 0.8968529612267799 0.0021294897691397 0.948402732471954 0.0008946298748116 0.95929616602314 0.0004665126697379 0.1042935665880658 0.0134872722784602 70.33999999999999 899745.0778283556 1022.0872395833334 21.20751953125 diff --git a/results/author/external_pgen_runs.tsv b/results/author/external_pgen_runs.tsv new file mode 100644 index 0000000..3cae4a5 --- /dev/null +++ b/results/author/external_pgen_runs.tsv @@ -0,0 +1,14 @@ +arm target train_subset train_size seed training_seconds rmse mae r2 pearson_r spearman_rho bias inference_seq_per_second inference_ms_per_sequence device peak_rss_mb peak_gpu_mb epochs_ran +esm2_8m_mlp log10_pgen_1mm all 79544 42 68.4 0.7434317223716531 0.520374913637156 0.9125300844628657 0.955902023500962 0.9589726094925388 0.066005454145396 636372.7860640743 0.0015714059775952 cuda 1131.9296875 23.45751953125 60.0 +esm2_8m_mlp log10_pgen_1mm all 79544 43 80.73 0.7429647954773283 0.5202331828575915 0.9126399243552248 0.9560829070791672 0.9586282966795376 0.0681313072794392 632408.853215861 0.0015812555357422 cuda 1126.29296875 23.45751953125 60.0 +esm2_8m_mlp log10_pgen_1mm all 79544 44 68.16 0.7453121949905853 0.521895684825011 0.9120870234415944 0.955716597134072 0.9590296543639312 0.0648420547328558 632583.6192838132 0.0015808186767974 cuda 1124.9921875 23.45751953125 60.0 +rtp_mlp log10_pgen_1mm all 79544 42 69.83 0.2370257379741502 0.1640886523530489 0.9911086533036424 0.9955494464625272 0.9958659210528849 -0.0074270037056943 905496.9714494308 0.0011043659244925 cuda 1050.56640625 21.77001953125 60.0 +rtp_mlp log10_pgen_1mm all 79544 43 66.05 0.2371002967282686 0.1640583226327462 0.991103058704482 0.9955453028703484 0.9958350139936892 -0.0028309819243721 828891.7263762074 0.0012064301864513 cuda 1054.671875 21.77001953125 60.0 +rtp_mlp log10_pgen_1mm all 79544 44 86.89 0.2383849753881611 0.1648012738041702 0.9910063850559538 0.995503697574166 0.9957890320238824 -0.0088264901684091 751504.6653472661 0.0013306637285317 cuda 1054.109375 21.77001953125 60.0 +sceptr_mlp log10_pgen_1mm all 79544 42 65.91 0.8103093724079922 0.5394118196202454 0.896084997028523 0.9479077517246606 0.9587721075749124 0.1036191435472859 1020112.022406086 0.0009802844962471 cuda 1021.7265625 21.20751953125 60.0 +sceptr_mlp log10_pgen_1mm all 79544 43 79.47 0.8136983165694129 0.5380069016025932 0.8952139752293935 0.9478649788339956 0.9594502316971228 0.1181053979058701 652900.7605952359 0.0015316263364256 cuda 1023.546875 21.20751953125 60.0 +sceptr_mlp log10_pgen_1mm all 79544 44 65.64 0.7978346598803266 0.5333620710399083 0.8992599114224236 0.949435466857206 0.9596661587973848 0.0911561583110414 1026222.4504837452 0.0009744475961606 cuda 1020.98828125 21.20751953125 60.0 +tcr_bert_mlp log10_pgen_1mm all 79544 42 86.02 0.4323397416215389 0.3047337966650835 0.970418046685067 0.985203459333992 0.9863763092964248 0.017726926052555 339208.0863098805 0.0029480429280994 cuda 1553.72265625 28.14501953125 60.0 +tcr_bert_mlp log10_pgen_1mm all 79544 43 72.09 0.4280938578492894 0.3025175762726348 0.9709962252575708 0.9855226286722562 0.986334939969123 0.0147702138797483 377791.1592998166 0.0026469650635905 cuda 1541.89453125 28.14501953125 60.0 +tcr_bert_mlp log10_pgen_1mm all 79544 44 85.52 0.4291836455531097 0.3048351665449082 0.9708483689485546 0.9853580451926848 0.98620912073739 0.0105794348645018 325750.2312496926 0.0030698366541863 cuda 1552.51953125 28.14501953125 60.0 +tfidf_ridge log10_pgen_1mm all 79544 42 9.0 0.785974434572495 0.5622312803788474 0.9022327566653896 0.9498971636908724 0.9653913662774836 0.0069710814573569 47475.47340859601 0.0210635077062536 cpu 554.30859375 diff --git a/results/author/external_reconstruction_results.tsv b/results/author/external_reconstruction_results.tsv new file mode 100644 index 0000000..b71574b --- /dev/null +++ b/results/author/external_reconstruction_results.tsv @@ -0,0 +1,7 @@ +representation seeds exact_match_mean exact_match_std token_accuracy_mean token_accuracy_std normalized_levenshtein_mean normalized_levenshtein_std within_edit_distance_1_mean within_edit_distance_1_std mean_edit_distance_mean mean_edit_distance_std training_seconds_mean epochs_mean inference_seq_per_second_mean source_dim explained_variance +rtp 3 0.8831338630192095 0.0058150137887855 0.9810685896331252 0.000560607194438 0.0130685588662131 0.0006119496602856 0.9404606255657244 0.0010878662200937 0.2455999195413859 0.0102899773206707 1254.5 59.0 13499.942816795909 128 0.968979 +onehot 3 0.8165878842736934 0.0037756299507998 0.9703579492481784 0.0004862497339854 0.0238099449743504 0.0004189465126061 0.9079084112776158 0.0018279266325006 0.4476516142009454 0.0077166637856656 998.82 47.0 13627.791460900138 1000 0.315236 +sceptr 3 0.7263401387911094 0.0125148065817577 0.9556794231100624 0.0024136647944061 0.0338323573330202 0.0020863791671621 0.8548057259713701 0.0092747603719289 0.5916725334406115 0.0353314332915674 657.6666666666666 31.0 13631.894064641703 64 1.0 +tcremp 3 0.518052901538771 0.0024158527908004 0.9371434757441994 0.0011054314818466 0.057005945958024 0.0008528920928819 0.7730061349693251 0.0042000629624502 0.934929095846324 0.0142856529454408 944.4033333333332 44.333333333333336 13506.954848081185 9000 0.959459 +tcr_bert 3 0.4777565456435013 0.0178585766284086 0.90975931989201 0.0028108844032423 0.0747954350264354 0.0028478342464577 0.6928827650943713 0.0079810659894325 1.270374467799792 0.0452585664172693 794.2933333333334 37.333333333333336 13529.272169955244 768 0.902213 +esm2_8m 3 0.2179087465218411 0.0075197155491379 0.8092609039533044 0.002076273541677 0.1568623343259775 0.001567719017009 0.3763786918770324 0.0073475867775242 2.5583827818565825 0.0221288616149074 790.4200000000001 37.333333333333336 13643.431387205686 320 0.970557 diff --git a/results/author/external_reconstruction_runs.tsv b/results/author/external_reconstruction_runs.tsv new file mode 100644 index 0000000..0e12620 --- /dev/null +++ b/results/author/external_reconstruction_runs.tsv @@ -0,0 +1,19 @@ +representation seed bottleneck_dim train_subset train_size epochs_ran best_epoch best_val_loss training_seconds test_loss exact_match token_accuracy normalized_levenshtein within_edit_distance_1 mean_edit_distance inference_seq_per_second inference_ms_per_sequence device +esm2_8m 42 64 all 79544 34 26 0.2073100912265288 715.4 0.2092864666229639 0.213617620436488 0.8070692877689258 0.1582913051676092 0.3692044654530825 2.5780951423111738 13737.439343699409 0.0727937700018776 cuda +esm2_8m 43 64 all 79544 38 30 0.206404773852764 809.74 0.2055744937597177 0.2135170471688625 0.8095149723617757 0.15711028054773 0.3760434476516142 2.562606859096852 13462.130148638444 0.0742824492824517 cuda +esm2_8m 44 64 all 79544 40 32 0.2068329342664816 846.12 0.2104361947530355 0.226591571960173 0.8111984517292117 0.1551854172625933 0.3838881625264004 2.5344463441617218 13730.72466927921 0.0728293680112438 cuda +onehot 42 64 all 79544 52 44 0.0400615889005936 1109.88 0.0400808783582387 0.8158503469777733 0.9707578308302072 0.0235220490955847 0.9084783264608268 0.4418183646786684 13437.93176766472 0.0744162135430892 cuda +onehot 43 64 all 79544 42 34 0.0394708287830536 891.93 0.0394849172100806 0.8132354420195113 0.9698166730736092 0.0242905725131555 0.9058634215025646 0.4564014884843608 13716.154190149347 0.0729067336322435 cuda +onehot 44 64 all 79544 47 39 0.0393366367102433 994.65 0.0403386892225497 0.8206778638237956 0.970499343840719 0.023617213314311 0.909383485869456 0.4447349894398069 13729.288424886348 0.0728369868162543 cuda +rtp 42 64 all 79544 60 52 0.0221906138154176 1276.43 0.0217311147791452 0.8838378758925878 0.9806996381182148 0.0133091188279378 0.9395554661570954 0.251030875993161 13431.362790330624 0.0744526088387628 cuda +rtp 43 64 all 79544 57 49 0.0216901884533655 1207.86 0.0220113017906745 0.8769988936940561 0.9807924283195696 0.0135236748927534 0.9401589057628482 0.2520366086694157 13726.01962618779 0.0728543326640817 cuda +rtp 44 64 all 79544 60 58 0.0212099805522041 1279.21 0.0209383756543199 0.8885648194709846 0.9817137024615916 0.012372882877948 0.9416675047772302 0.233732273961581 13342.446033869304 0.0749487760686112 cuda +sceptr 42 64 all 79544 36 28 0.0558294934722093 762.97 0.0567849224003461 0.7353917328774012 0.9575020877795304 0.0322961628513188 0.8625163431559891 0.5667303630694961 13706.086211692991 0.0729602881927647 cuda +sceptr 43 64 all 79544 31 23 0.0555792820568268 659.06 0.056470837826148 0.7315699487076335 0.9565940693805592 0.0329932651279869 0.8573871065070904 0.5761842502262898 13432.203122547407 0.0744479510082297 cuda +sceptr 44 64 all 79544 26 18 0.0571236062126281 550.97 0.0566988510963244 0.7120587347882933 0.9529421121700976 0.0362076440197549 0.8445137282510309 0.6321029870260485 13757.392859684716 0.0726881910111359 cuda +tcr_bert 42 64 all 79544 32 24 0.1105848010151814 681.6 0.1102362450880882 0.4596198330483757 0.9069844510134016 0.0776677400683502 0.6845016594589158 1.3154983405410845 13456.478415255033 0.0743136479798713 cuda +tcr_bert 43 64 all 79544 37 29 0.1082922487686842 785.5 0.1106430744895568 0.4783264608267122 0.9096886225957396 0.0747458457766419 0.6937544000804586 1.2706426631801266 13714.924236424891 0.0729132719044952 cuda +tcr_bert 44 64 all 79544 43 35 0.1096065732148977 915.78 0.1109284074642719 0.4953233430554158 0.9126048860668884 0.0719727192343142 0.7003922357437393 1.2249823996781657 13416.413858185811 0.0745355659545241 cuda +tcremp 42 64 all 79544 42 34 0.0778516114522249 895.52 0.0781389696475787 0.5153374233128835 0.9358753429923514 0.0579907698564299 0.7681786181233028 0.9514231117369004 13352.933211938987 0.0748899125104505 cuda +tcremp 43 64 all 79544 46 38 0.0764745114705501 975.58 0.0765142872547492 0.5199637936236549 0.9379034716791048 0.0565094918624861 0.7750176003218344 0.9268832344362868 13694.12270506786 0.0730240280109308 cuda +tcremp 44 64 all 79544 45 37 0.0775704320806723 962.11 0.0771636510124573 0.5188574876797747 0.9376516125611422 0.0565175761551559 0.7758221864628382 0.926480941365785 13473.808627236707 0.0742180646664777 cuda diff --git a/results/author/factorial_results.tsv b/results/author/factorial_results.tsv new file mode 100644 index 0000000..de1a927 --- /dev/null +++ b/results/author/factorial_results.tsv @@ -0,0 +1,8 @@ +condition reconstruction_weight tcremp_weight pgen_weight best_epoch best_validation_joint_loss test_reconstruction_loss_active test_reconstruction_token_accuracy_active test_reconstruction_exact_match_active test_tcremp_cosine_active test_tcremp_standardized_mse_active test_pgen_rmse_active test_pgen_r2_active parameters checkpoint_sha256 run_config_sha256 remote_output +R 1.0 0.0 0.0 40 0.005638612275925364 0.0062070361559643066 0.9979169387082533 0.9765664286432666 NA NA NA NA 21535913 5b3c29358a905eabc1eae262115b75db225562f072aa627488bb5764d502b23e 78002783c3a29f27c3921e31233fd2e3e4c364d8d2772cc00aea7e537dbe814c /home/evlasova/irrm-codec/experiments/trb/data-anchor-r-seed42 +T 0.0 1.0 0.0 7 0.19805858843759142 NA NA NA 0.9914940227196632 0.2756815801099142 NA NA 21535913 754277ecf372ecb412b8d588ecfd5801aa4be6e1bbd87a6d82917c58b9643ec9 f899ff4f1f4b9503d9872cfab22b97625ee3e58495a8d87d366bcb284fe493fe /home/evlasova/irrm-codec/experiments/trb/data-anchor-t-seed42 +P 0.0 0.0 1.0 38 0.005830517954523791 NA NA NA NA NA 0.2752929237653586 0.9880059307820565 21535913 6b9ea5fd0f23c21e3f7a2a4d0e983857134f11bbfe3bb7b7df7327c43bc9e5c4 936d97c4f2b5fb7324e4086c87f2a61525a7d8c54c8ae326bb0565fb924c6cf9 /home/evlasova/irrm-codec/experiments/trb/data-anchor-p-seed42 +RT 1.0 1.0 0.0 25 0.20361182044314435 0.009265304249247219 0.9968038999881856 0.9631901840490797 0.991784158396095 0.2699985514131456 NA NA 21535913 e0c6f100dce2788afd5015aff4b866af9a5fc73b77a70abc9b0d9813f5208bda ad9ad21656776121c77d969dc1432c537fde71661bf2c17985416299fdbb22d9 /home/evlasova/irrm-codec/experiments/trb/data-anchor-rt-seed42 +RP 1.0 0.0 1.0 39 0.00867614229220961 0.004723803990094098 0.99837707762046 0.9812933722216635 NA NA 0.22639804111600034 0.9918881150217478 21535913 997ce62b620ce18fbc56d88fe7414949f0a3ad108d0788aa99ba607d7df6d95b d93acfeda0bc251b80ef49bafc9b842656de2de979b1d36a7d49a0d36bdd8202 /home/evlasova/irrm-codec/experiments/trb/data-anchor-rp-seed42 +TP 0.0 1.0 1.0 12 0.20485042515300259 NA NA NA 0.9916881372136608 0.27117700439744896 0.34955449828735125 0.9806622364855879 21535913 6df6885f0652404e597e5a354c5ed44dfb63abc4a7c7e00cd6ce2550487a8e06 2e7f1edfab65c9c5882654ae0f295734d3d55858c4450a2635991680a3b27bf6 /home/evlasova/irrm-codec/experiments/trb/data-anchor-tp-seed42 +RTP 1.0 1.0 1.0 25 0.20576476722275314 0.008851642950913693 0.9967355009606954 0.9635924771195816 0.9918976390216716 0.2662354936682937 0.270233460057513 0.9884427449327019 21535913 5ecd99c754eef189a68e72be8628dd7c3858235b651fb970ad21ad5a69dbf222 a58b8a7e85a9a1aa328a5ae0ac196c6c079a55fcf1d24088f25dd2e43fc81e9f /home/evlasova/irrm-codec/experiments/trb/anchor-data-rtp-seed42 diff --git a/results/author/latent_sufficiency_summary.tsv b/results/author/latent_sufficiency_summary.tsv new file mode 100644 index 0000000..2db4f1a --- /dev/null +++ b/results/author/latent_sufficiency_summary.tsv @@ -0,0 +1,10 @@ +representation val_sequence_loss test_sequence_loss val_tcremp_loss test_tcremp_loss val_pgen_loss test_pgen_loss test_mean_normalized_regret test_worst_task_normalized_regret test_fixed_reference_hypervolume test_tasks_within_5pct_of_specialist_loss +r 0.4935558985376048 0.49224470019186 0.548148507589488 0.5506662230820216 0.0614951493528322 0.0653795911789573 0.1400491045037798 0.3722305268570544 0.5932814628050733 1 +t 0.4509414655450394 0.4490276663440188 0.2865061488114791 0.2821712378823272 0.117128964292562 0.1173834178493009 -0.0455000984472581 0.1080270920473844 0.8919729079526156 2 +p 0.6331614164344623 0.6350980117693602 0.8335210218559798 0.837555952241106 0.0115658875926671 0.0121837848696961 0.511661387724634 0.7764896312264542 0.0540858276076921 0 +rt 0.4260993086890938 0.4239572890785302 0.3990870048344427 0.3952955325072009 0.05701791602458 0.0586903250223235 -0.0570831929098857 0.1532961774080955 0.8059702254687741 1 +rp 0.4843092262983698 0.4828165204583172 0.5390500380432909 0.5401505206608808 0.0122720687155209 0.0126924586574998 0.1003535424521918 0.3574127484198617 0.6418482152942699 1 +tp 0.4384337059212025 0.4379317115096868 0.295442866793231 0.2910054249698389 0.019219655658664 0.0200339606698381 -0.0942818672224091 0.0086448937860575 0.985069940870266 2 +rtp 0.42446096833766 0.42360938070699 0.3902434068917082 0.3884632290116502 0.0208702339405904 0.0210434258818411 -0.0737242868597041 0.1436687121796011 0.8480459079597709 1 +concat_r_t_p 0.3800976034848598 0.3797286195409372 0.2779766205636249 0.2734238594102645 0.01125705180847 0.0119070518391921 -0.2091851165085149 0.0003482880518308 0.9996517119481692 3 +pca128_concat_r_t_p 0.4163787686308496 0.4150748132636582 0.2840515017505456 0.2797396175536809 0.0131192125953123 0.013556514821328 -0.1425721751878963 0.0020321932513549 0.9979678067486452 2 diff --git a/results/author/low_label_sequence_summary.tsv b/results/author/low_label_sequence_summary.tsv new file mode 100644 index 0000000..9eed89f --- /dev/null +++ b/results/author/low_label_sequence_summary.tsv @@ -0,0 +1,43 @@ +model train_fraction train_rows test_sequence_composition_rmse +r 0.01 795 0.01835094344455989 +r 0.05 3977 0.017523018379142774 +r 0.1 7954 0.017369375370154497 +r 0.25 19886 0.017267594094227467 +r 0.5 39772 0.017234372857679845 +r 1.0 79544 0.017218643463211603 +t 0.01 795 0.03147869204585527 +t 0.05 3977 0.029924893812797497 +t 0.1 7954 0.02970307989219788 +t 0.25 19886 0.02955591518184327 +t 0.5 39772 0.029504023885818596 +t 1.0 79544 0.029479354201959598 +p 0.01 795 0.04532366426571271 +p 0.05 3977 0.044574953547430206 +p 0.1 7954 0.04428606909530054 +p 0.25 19886 0.04387545458540854 +p 0.5 39772 0.04377336751769786 +p 1.0 79544 0.043726839108308116 +rt 0.01 795 0.016244031377328154 +rt 0.05 3977 0.015384379770039164 +rt 0.1 7954 0.015258002249398772 +rt 0.25 19886 0.015178035629387987 +rt 0.5 39772 0.015153765261059585 +rt 1.0 79544 0.015137956375535723 +rp 0.01 795 0.02006764040046825 +rp 0.05 3977 0.019064061738322383 +rp 0.1 7954 0.018926092173410757 +rp 0.25 19886 0.018820184904729538 +rp 0.5 39772 0.018784321535024042 +rp 1.0 79544 0.01876798572396856 +tp 0.01 795 0.03074244468361117 +tp 0.05 3977 0.02869302085614458 +tp 0.1 7954 0.028474638592191966 +tp 0.25 19886 0.028330022822164024 +tp 0.5 39772 0.028274905585994548 +tp 1.0 79544 0.02825358338746815 +rtp 0.01 795 0.016079114347764864 +rtp 0.05 3977 0.015141720406547479 +rtp 0.1 7954 0.015008292222499603 +rtp 0.25 19886 0.014917465005251243 +rtp 0.5 39772 0.014892085094678839 +rtp 1.0 79544 0.014880154198137008 diff --git a/results/author/provenance_and_hashes.tsv b/results/author/provenance_and_hashes.tsv new file mode 100644 index 0000000..dbb79c3 --- /dev/null +++ b/results/author/provenance_and_hashes.tsv @@ -0,0 +1,16 @@ +kind name path sha256 +training_source airr /projects/immunestatus/vdjrearm/airr_format/trb_background_100k.tsv 403e1db78a1a4b3006191d722ef30f533f194b27b7141e3af0e01c1e819cfaf3 +training_source tcremp_embeddings /projects/immunestatus/vdjrearm/tcremp/trb_background_100k_tcremp.parquet not recorded for large binary +training_source tcremp_representations /projects/immunestatus/vdjrearm/tcremp/trb_background_100k_tcremp_representations.tsv 52ae2d30b446b2afd5609a9da987b95facacb8bd1a2670851e98afe21fb5b2c1 +training_source pgen /projects/immunestatus/vdjrearm/pgen_1mm_background_100k/trb_background_100k_pgen.tsv 25de4450e1593c7805acaf1819439b563bf42357e24525f49d75ba93afe36caa +training_source tcremp_writer /home/evlasova/tcremp-git/tcremp/tcremp/tcremp_run.py 0052c8021fe23f542ff8d75c4e9246d78c0cd691fa93094fba176cbfec4c4252 +benchmark READY /home/evlasova/irrm-codec/data/benchmark/trb/READY.json ce08945c0c0a975c2b05129695c26399456f185fe79b89914c0ad52e4fdcf167 +benchmark train-only normalizer /home/evlasova/irrm-codec/data/benchmark/trb/target_standardizer.npz ed0052ad510b7e0299628796162037cbb1a3c7e5e8d661ca4a668eb49a830b24 +tokenizer DATA-ANCHOR /home/evlasova/irrm-codec/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json 8ef9aeddf8c069d4db898369f01f97e80160ba2ceb3c2354f18547467beeed4d +downstream VDJdb 2025-07-30 vdjdb.txt pinned official snapshot 191290d17ed3f489ac32b476eda6b9e3db93fab05c0968fb5da567187a27a94a +downstream VDJdb cluster_members.txt pinned official motif membership b5d572d912c180ff0221f8966cdacedbcd978467f5367f947dbee71bdd12b1c4 +result viral multi-epitope evaluation /home/evlasova/irrm-codec/experiments/trb/vdjdb-viral-motif-top10/evaluation-1426501 712a79c9d4a2661c7da3db9d8caa0e9c04dbf59e790edfa820889bc07309ed76 +student_pr reconstruction benchmark PR #3 https://github.com/antigenomics/irrm-codec/pull/3 git-commit:603ea761f7439a2faffe314525c5ac67def28357 +student_pr pgen benchmark PR #4 https://github.com/antigenomics/irrm-codec/pull/4 git-commit:b96f06eb42d9c92e3982103bb903a1dc67d59138 +result external reconstruction locked-cohort rerun /home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/results/reconstruction/reconstruction_summary.csv 5f15ae270e153ef179cddb8651d8e7f75d5ceb3b77aa8c66c48c66fc5d0739d4 +result external pgen locked-cohort rerun /home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/results/pgen/pgen_summary.csv 1be5e255e91c0b55f2c7429aeb48b75ce332af524fee17779f43070a88b8a362 diff --git a/results/author/rtp_data_anchor_error_analysis_by_length.tsv b/results/author/rtp_data_anchor_error_analysis_by_length.tsv new file mode 100644 index 0000000..f488d55 --- /dev/null +++ b/results/author/rtp_data_anchor_error_analysis_by_length.tsv @@ -0,0 +1,23 @@ +cdr3_length n_sequences exact_sequences sequence_error_rate total_target_residues edit_errors residue_error_rate mean_abs_log10_pgen_error median_abs_log10_pgen_error mean_signed_log10_pgen_error mean_true_log10_pgen +7 3 3 0.0 21 0 0.0 0.22788333892822266 0.20193147659301758 0.09326235453287761 -6.067805767059326 +8 18 17 0.05555555555555558 144 2 0.013888888888888888 0.19520963562859428 0.1939537525177002 0.10459881358676487 -6.952210797203912 +9 58 58 0.0 522 0 0.0 0.13554670070779734 0.09454631805419922 0.012654896440177128 -6.39963698387146 +10 161 161 0.0 1610 0 0.0 0.1254880768912179 0.08246088027954102 0.026373279761083376 -6.248149824438628 +11 368 368 0.0 4048 0 0.0 0.1158735207889391 0.08518195152282715 0.027172174142754597 -6.210726919381515 +12 793 790 0.003783102143757877 9516 4 0.0004203446826397646 0.13495038015695118 0.09746122360229492 0.030675953303429697 -6.365015834156532 +13 1195 1187 0.006694560669456062 15535 8 0.0005149662053427744 0.14315945134502076 0.10540390014648438 0.07770905754057433 -6.662306459179483 +14 1538 1534 0.002600780234070177 21532 5 0.00023221252089912687 0.1622169594460873 0.1293473243713379 0.08179769993448444 -7.131347616255206 +15 1632 1623 0.005514705882352922 24480 11 0.00044934640522875815 0.17407382820166795 0.1297016143798828 0.07869276112201166 -7.785865174204695 +16 1358 1343 0.011045655375552244 21728 18 0.0008284241531664212 0.19903086597684205 0.15855979919433594 0.0986646977369845 -8.657309157858892 +17 1077 1045 0.02971216341689875 18309 57 0.003113223005079469 0.22972634775060796 0.18281316757202148 0.09377350802762484 -9.486332648975935 +18 750 707 0.05733333333333335 13500 86 0.00637037037037037 0.24847226905822753 0.20470094680786133 0.0733973954518636 -10.422224502563477 +19 447 390 0.12751677852348997 8493 120 0.01412928293889085 0.3059646789926277 0.24474716186523438 0.07611795506488024 -11.592920788715883 +20 284 229 0.19366197183098588 5680 111 0.019542253521126762 0.3081234794267466 0.2438054084777832 0.08418968193967577 -12.651549712033338 +21 123 83 0.3252032520325203 2583 102 0.03948896631823461 0.3743546493654329 0.32733917236328125 0.03476118072261655 -13.97880065731886 +22 64 30 0.53125 1408 102 0.07244318181818182 0.4636581093072891 0.41320085525512695 0.09101755917072296 -15.086607545614243 +23 36 9 0.75 828 80 0.0966183574879227 0.45497918128967285 0.24609375 0.09665173954433864 -16.783052576912773 +24 18 2 0.8888888888888888 432 49 0.11342592592592593 0.4684235254923503 0.43729591369628906 0.2559358808729384 -17.646304077572292 +25 15 1 0.9333333333333333 375 54 0.144 0.5853490193684896 0.41939735412597656 0.23636004130045574 -18.517236709594727 +26 2 0 1.0 52 8 0.15384615384615385 0.9295778274536133 0.9295778274536133 0.9295778274536133 -19.613176345825195 +27 2 1 0.5 54 5 0.09259259259259259 1.296828269958496 1.296828269958496 1.296828269958496 -21.37657928466797 +28 1 0 1.0 28 8 0.2857142857142857 1.1843910217285156 1.1843910217285156 1.1843910217285156 -25.132247924804688 diff --git a/results/author/rtp_data_anchor_error_rates_by_target_length.tsv b/results/author/rtp_data_anchor_error_rates_by_target_length.tsv new file mode 100644 index 0000000..2808a96 --- /dev/null +++ b/results/author/rtp_data_anchor_error_rates_by_target_length.tsv @@ -0,0 +1,23 @@ +target_length sequences exact length_mismatch aligned_positions wrong_aligned_positions sequence_error_rate aligned_position_error_rate +7 3 3 0 21 0 0.0 0.0 +8 18 17 0 144 2 0.05555555555555558 0.013888888888888888 +9 58 58 0 522 0 0.0 0.0 +10 161 161 0 1610 0 0.0 0.0 +11 368 368 0 4048 0 0.0 0.0 +12 793 790 2 9492 2 0.003783102143757877 0.00021070375052675939 +13 1195 1187 8 15431 0 0.006694560669456062 0.0 +14 1538 1534 3 21490 2 0.002600780234070177 9.306654257794323e-05 +15 1632 1623 9 24345 0 0.005514705882352922 0.0 +16 1358 1343 12 21536 6 0.011045655375552244 0.0002786032689450223 +17 1077 1045 16 18037 38 0.02971216341689875 0.0021067805067361534 +18 750 707 21 13122 55 0.05733333333333335 0.004191434232586496 +19 447 390 21 8094 95 0.12751677852348997 0.011737089201877934 +20 284 229 25 5180 75 0.19366197183098588 0.01447876447876448 +21 123 83 12 2331 97 0.3252032520325203 0.041613041613041614 +22 64 30 11 1166 92 0.53125 0.07890222984562607 +23 36 9 14 506 43 0.75 0.08498023715415019 +24 18 2 8 240 35 0.8888888888888888 0.14583333333333334 +25 15 1 5 250 46 0.9333333333333333 0.184 +26 2 0 2 0 0 1.0 +27 2 1 1 27 0 0.5 0.0 +28 1 0 1 0 0 1.0 diff --git a/results/author/rtp_data_anchor_region_error_rates.tsv b/results/author/rtp_data_anchor_region_error_rates.tsv new file mode 100644 index 0000000..07523e5 --- /dev/null +++ b/results/author/rtp_data_anchor_region_error_rates.tsv @@ -0,0 +1,4 @@ +region eligible_positions wrong_positions error_rate +N_edge_1_3 29316 3 0.00010233319688907081 +middle 88960 584 0.006564748201438849 +C_edge_1_3 29316 1 3.411106562969027e-05 diff --git a/results/author/rtp_data_anchor_substitution_matrix.tsv b/results/author/rtp_data_anchor_substitution_matrix.tsv new file mode 100644 index 0000000..6b729e6 --- /dev/null +++ b/results/author/rtp_data_anchor_substitution_matrix.tsv @@ -0,0 +1,21 @@ + C S T A G P D E Q N H R K M I L V W Y F +C 0 1 0 0 2 1 0 0 0 0 0 1 0 0 1 1 0 0 0 0 +S 0 0 1 2 10 5 2 0 2 0 2 5 1 0 1 2 1 2 0 1 +T 0 4 0 0 5 5 0 1 2 0 0 2 2 2 1 3 0 0 1 0 +A 0 3 0 0 9 3 1 1 1 1 0 2 1 0 2 1 4 1 1 1 +G 3 8 6 6 0 7 2 2 3 2 0 15 2 2 2 7 7 0 3 1 +P 1 3 4 2 6 0 3 1 2 2 2 10 1 0 2 4 2 2 0 3 +D 0 3 0 3 1 4 0 0 0 4 1 6 1 0 0 2 3 1 0 0 +E 0 0 1 3 1 2 0 0 1 0 0 3 1 1 1 1 4 0 1 0 +Q 0 2 3 1 5 2 0 1 0 1 0 1 1 0 0 1 2 0 1 2 +N 1 0 0 1 2 3 3 0 1 0 0 5 2 0 0 2 1 1 1 1 +H 0 1 0 0 1 2 2 1 0 0 0 1 3 0 1 1 0 0 0 1 +R 1 4 2 2 11 6 6 5 2 6 2 0 3 0 2 7 6 1 0 1 +K 0 1 2 1 2 1 2 1 0 2 3 1 0 0 0 4 0 3 0 1 +M 0 0 2 0 2 0 0 1 0 0 0 0 1 0 0 0 0 0 0 0 +I 0 1 0 2 2 3 0 1 0 0 1 2 0 0 0 2 1 0 0 0 +L 1 3 4 2 5 4 4 0 0 1 1 10 1 0 0 0 3 0 5 4 +V 0 1 0 4 6 3 3 3 3 2 0 4 1 0 0 4 0 2 0 3 +W 0 2 0 1 3 1 1 1 0 1 0 1 0 0 0 0 1 0 0 0 +Y 0 0 1 1 2 0 0 0 2 0 0 0 0 0 1 5 1 0 0 1 +F 0 1 1 0 3 3 0 0 2 1 1 0 2 0 1 2 3 0 1 0 diff --git a/results/author/rtp_data_anchor_substitution_results.json b/results/author/rtp_data_anchor_substitution_results.json new file mode 100644 index 0000000..417bcc8 --- /dev/null +++ b/results/author/rtp_data_anchor_substitution_results.json @@ -0,0 +1,47 @@ +{ + "analysis": "locked_test_autoregressive_substitution_errors", + "benchmark_ready_sha256": "ce08945c0c0a975c2b05129695c26399456f185fe79b89914c0ad52e4fdcf167", + "between_group_errors": 463, + "between_group_fraction": 0.7874149659863946, + "checkpoint": "/home/evlasova/irrm-codec/experiments/trb/anchor-data-rtp-seed42/best.pt", + "checkpoint_sha256": "5ecd99c754eef189a68e72be8628dd7c3858235b651fb970ad21ad5a69dbf222", + "data_dir": "/home/evlasova/irrm-codec/data/benchmark/trb", + "exact_match": 0.9635924771195816, + "exact_sequences": 9581, + "groups": { + "C": "C", + "DEQN": "DENQ", + "HRK": "HKR", + "MILV": "ILMV", + "STAGP": "AGPST", + "WYF": "FWY" + }, + "length_matched_sequences": 9772, + "length_mismatched_sequences": 171, + "model_config": { + "d_model": 320, + "decoder_layers": 4, + "decoder_memory_tokens": 4, + "dropout": 0.1, + "encoder_layers": 4, + "ff_dim": 1280, + "input_vocab_size": 1228, + "latent_dim": 128, + "max_sequence_len": 40, + "nhead": 8, + "output_vocab_size": 25, + "pgen_head_dim": 256, + "share_input_output_embeddings": false, + "tcremp_dim": 9000, + "tcremp_head_dim": 1024 + }, + "note": "Only incorrect aligned positions from equal-length target/prediction pairs enter the matrix.", + "status": "complete", + "test_sequences": 9943, + "tokenizer": "/home/evlasova/irrm-codec/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json", + "tokenizer_sha256": "8ef9aeddf8c069d4db898369f01f97e80160ba2ceb3c2354f18547467beeed4d", + "unexpected_residue_errors": 0, + "within_group_errors": 125, + "within_group_fraction": 0.21258503401360543, + "wrong_aligned_positions": 588 +} diff --git a/results/author/tokenizer_results.tsv b/results/author/tokenizer_results.tsv new file mode 100644 index 0000000..6fd6452 --- /dev/null +++ b/results/author/tokenizer_results.tsv @@ -0,0 +1,13 @@ +tokenizer family requested_vocab actual_input_vocab best_epoch best_validation_joint_loss test_reconstruction_exact_match test_reconstruction_token_accuracy test_tcremp_cosine test_pgen_rmse parameters remote_output +Character character 25.0 25 30 0.19693734171098476 0.9815950920245399 0.9984206042743174 0.992289113360584 0.343395136972041 21142953 /home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-seed42 +WordPiece-44 wordpiece 44.0 44 30 0.19666322000361444 0.9789801870662778 0.9981159176973157 0.9923110324276845 0.34131871315836837 21157033 /home/evlasova/irrm-codec/experiments/trb/wp-sweep-v44-rtp-seed42 +WordPiece-64 wordpiece 64.0 64 28 0.20685211541870166 0.9643970632605854 0.9968474266420431 0.991931193049712 0.32575211456733877 21163433 /home/evlasova/irrm-codec/experiments/trb/wp-sweep-v64-rtp-seed42 +WordPiece-128 wordpiece 128.0 128 26 0.20901615753579544 0.9598712662174395 0.996505431504592 0.9918078695230294 0.2893693995404677 21183913 /home/evlasova/irrm-codec/experiments/trb/wp-sweep-v128-rtp-seed42 +WordPiece-256 wordpiece 256.0 256 23 0.20649890744307267 0.9636930503872071 0.996660883839797 0.9919296334857589 0.2705196099398586 21224873 /home/evlasova/irrm-codec/experiments/trb/wp-sweep-v256-rtp-seed42 +WordPiece-512 wordpiece 512.0 512 25 0.2104906694103532 0.9666096751483456 0.9969966608838398 0.9918318263819961 0.31541390221471494 21306793 /home/evlasova/irrm-codec/experiments/trb/wp-sweep-v512-rtp-seed42 +WordPiece-1024 wordpiece 1024.0 1024 28 0.21072169335227994 0.972543497938248 0.9975003264499039 0.991686036310289 0.3159752198188979 21470633 /home/evlasova/irrm-codec/experiments/trb/wp-sweep-v1024-rtp-seed42 +WordPiece-2048 wordpiece 2048.0 2048 24 0.21218146349443082 0.9552448959066681 0.9959893297517115 0.9918661757301039 0.3290175757219649 21798313 /home/evlasova/irrm-codec/experiments/trb/wp-sweep-v2048-rtp-seed42 +WordPiece-4096 wordpiece 4096.0 4096 25 0.21656744164125932 0.9517248315397767 0.9954297013449735 0.9919115816768741 0.4196596483780269 22453673 /home/evlasova/irrm-codec/experiments/trb/wp-sweep-v4096-rtp-seed42 +EDGE-3 edge_anchor NA 16256 31 0.2154474285803342 0.9647993563310872 0.9967292828672872 0.9917393089431783 0.3389353724618627 26344873 /home/evlasova/irrm-codec/experiments/trb/anchor-edge3-rtp-seed42 +DATA-ANCHOR data_anchor NA 1228 25 0.20576476722275314 0.9635924771195816 0.9967355009606954 0.9918976390216716 0.270233460057513 21535913 /home/evlasova/irrm-codec/experiments/trb/anchor-data-rtp-seed42 +GERMLINE-ANCHOR germline_anchor NA 387 28 0.20924612771849702 0.9716383385296188 0.9972826931806169 0.9917630262086177 0.25866407306127065 21266793 /home/evlasova/irrm-codec/experiments/trb/anchor-germline-rtp-seed42 diff --git a/results/pgen/pgen_runs.csv b/results/pgen/pgen_runs.csv new file mode 100644 index 0000000..7fae920 --- /dev/null +++ b/results/pgen/pgen_runs.csv @@ -0,0 +1,181 @@ +arm,target,train_subset,train_size,seed,training_seconds,rmse,mae,r2,pearson_r,spearman_rho,bias,inference_seq_per_second,inference_ms_per_sequence,device,peak_rss_mb,peak_gpu_mb,epochs_ran +esm2_8m_mlp,log10_pgen,1k,1000,42,5.34,1.0771408535945268,0.7260378855007953,0.669508854815493,0.8216067160495758,0.8331470316606933,0.13865496676993277,705960.7233955928,0.0014165093989791824,cuda,1197.78125,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen,1k,1000,43,5.1,1.0751237520233003,0.7263235609554174,0.6707454805022155,0.8223345359775261,0.8340828674416486,0.13843278495773942,747093.4212828245,0.00133852068765766,cuda,1194.6875,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen,1k,1000,44,4.11,1.0695933554075805,0.7286857503677006,0.6741241143295938,0.8228841204325575,0.8323480349886503,0.10114906892531643,735015.6226967543,0.001360515299431357,cuda,1195.64453125,23.45751953125,46.0 +irrm_pretrained,log10_pgen,1k,1000,42,6.24,0.8649319159558528,0.573572627859008,0.7869023415752135,0.9225699791418592,0.9119520006525227,0.4028323986191931,80863.50773405924,0.0123665177039903,cuda,1552.5234375,255.98876953125,25.0 +irrm_pretrained,log10_pgen,1k,1000,43,5.9,0.8313225259801558,0.5495318621784235,0.8031416099725512,0.9236803569147729,0.9135004792110046,0.36259704153653866,78945.20760898925,0.012667013366446999,cuda,1550.91015625,255.98876953125,23.0 +irrm_pretrained,log10_pgen,1k,1000,44,6.16,0.8039642867965016,0.5216304005308618,0.8158853501548303,0.929212427359479,0.922879820889854,0.3472389288297001,80992.25445801619,0.012346859668147256,cuda,1545.44140625,255.98876953125,25.0 +irrm_pretrained,log10_pgen,1k,1000,45,6.53,0.8174569543838605,0.5322813395104757,0.8096536218416155,0.9271785368544058,0.9203442313044025,0.3599058488751396,80060.79152636931,0.012490508536511706,cuda,1544.26171875,255.98876953125,26.0 +irrm_pretrained,log10_pgen,1k,1000,46,7.85,0.8237597476667486,0.5406323265712351,0.8067070717851613,0.9297969710103117,0.9226493652846909,0.3918390707428683,79906.00704535747,0.012514703674685743,cuda,1563.98046875,255.98876953125,30.0 +irrm_pretrained,log10_pgen,1k,1000,47,6.61,0.795733847970004,0.5171656860759076,0.8196357352293764,0.9285150090131103,0.921082252441927,0.3313119003076507,77841.70608175163,0.012846583795963706,cuda,1548.0859375,255.98876953125,25.0 +irrm_pretrained,log10_pgen,1k,1000,48,5.72,0.8021244491460586,0.5302790450689656,0.816727062835779,0.9220099669223497,0.9179323947568008,0.31359157081286243,80136.58582260921,0.012478694840002361,cuda,1566.8984375,255.98876953125,22.0 +irrm_pretrained,log10_pgen,1k,1000,49,4.68,0.8318569691142761,0.5448034584308773,0.8028884147796086,0.907612296788354,0.9008832577185555,0.2114075241262159,80186.96882282455,0.012470854238293123,cuda,1556.875,255.98876953125,18.0 +irrm_pretrained,log10_pgen,1k,1000,50,5.84,0.7830557427649305,0.5141476876670233,0.8253372917973597,0.922166417284953,0.9141818947508201,0.24003536004227982,80261.9705837884,0.012459200698991854,cuda,1547.0625,255.98876953125,23.0 +irrm_pretrained,log10_pgen,1k,1000,51,4.96,0.7998190142313116,0.5398007120038264,0.8177790607593768,0.9173102792460404,0.9079210168206249,0.28461443290058946,79910.51413303308,0.012513997824306627,cuda,1562.07421875,255.98876953125,19.0 +irrm_scratch,log10_pgen,1k,1000,42,7.69,0.8441119514619323,0.5319142441170579,0.7970379080578323,0.9051526732558525,0.906339826203102,0.25311759515793814,78739.84026758223,0.01270005116344778,cuda,1565.42578125,255.98876953125,30.0 +irrm_scratch,log10_pgen,1k,1000,43,10.56,0.7437902178581257,0.4714709206837621,0.8424146805105919,0.9204954050970439,0.9219248220744659,0.1288298033379222,80851.88208602746,0.012368295878826752,cuda,1546.63671875,255.98876953125,44.0 +irrm_scratch,log10_pgen,1k,1000,44,6.41,0.9439559508218183,0.6187347933622375,0.7461844128303036,0.8696195540239993,0.8685103108517291,0.18257265629955535,78595.5292329371,0.012723370015567363,cuda,1561.90234375,255.98876953125,25.0 +irrm_scratch,log10_pgen,1k,1000,45,6.98,0.885964311952314,0.5553881292737052,0.7764126193240595,0.8968896378559235,0.8945789702406476,0.26261517141442153,80666.26935827773,0.012396755272746265,cuda,1546.95703125,255.98876953125,28.0 +irrm_scratch,log10_pgen,1k,1000,46,5.84,0.9763104933844673,0.6321716080463298,0.7284869263676981,0.870469034070754,0.861312281330798,0.31175582583126027,80848.42252814662,0.012368825126450171,cuda,1545.22265625,255.98876953125,23.0 +irrm_scratch,log10_pgen,1k,1000,47,8.73,0.8783602742491172,0.5487071160322827,0.780234152019625,0.9088537530042312,0.9072049580677213,0.3422900177627635,76363.0089395081,0.013095345690112373,cuda,1551.3046875,255.98876953125,34.0 +irrm_scratch,log10_pgen,1k,1000,48,10.64,0.919467769293943,0.58948479091107,0.7591826060141407,0.9116245478710833,0.9121262001095223,0.4535502971643727,79905.32213663701,0.012514810944507722,cuda,1558.87890625,255.98876953125,42.0 +irrm_scratch,log10_pgen,1k,1000,49,6.74,0.8736720995109767,0.5701165193450962,0.7825738551116111,0.8916065482310117,0.8876233950864687,0.20014242571748278,80802.20622960075,0.012375899702026502,cuda,1573.88671875,255.98876953125,27.0 +irrm_scratch,log10_pgen,1k,1000,50,10.54,0.7922520694606673,0.49711589769390385,0.8212106701729678,0.9145798789121482,0.9180335296945293,0.1337827529050968,80315.75408298723,0.012450857386792856,cuda,1540.3359375,255.98876953125,43.0 +irrm_scratch,log10_pgen,1k,1000,51,11.01,0.7476454702404467,0.48343691228369096,0.8407768378569647,0.9201273521939399,0.9226283285508657,0.0296063515905427,80351.0919746068,0.012445381580078938,cuda,1551.796875,255.98876953125,45.0 +sceptr_mlp,log10_pgen,1k,1000,42,4.74,1.0981920422083575,0.7610458447203725,0.6564646628435225,0.8121566102601996,0.8074398847933616,0.10412556109970554,1062266.7426666566,0.0009413831383722458,cuda,1172.00390625,21.20751953125,60.0 +sceptr_mlp,log10_pgen,1k,1000,43,3.42,1.0999811361432124,0.7624915073977926,0.655344425928786,0.8120102287560248,0.8046603984441668,0.11799728837458752,1061074.718902842,0.00094244069921297,cuda,1168.94140625,21.20751953125,43.0 +sceptr_mlp,log10_pgen,1k,1000,44,3.2,1.0940586433868646,0.758922927361791,0.6590458075780972,0.8133347744266023,0.8086149192965398,0.09307519652395542,1071874.0294636576,0.0009329454511556532,cuda,1174.91796875,21.20751953125,39.0 +tcr_bert_mlp,log10_pgen,1k,1000,42,4.65,0.8673881393172919,0.5780485788243145,0.7856903186825444,0.8898886415389792,0.8930102409604314,0.10531888878673068,476370.05146904866,0.002099208371550984,cuda,1242.5546875,28.14501953125,48.0 +tcr_bert_mlp,log10_pgen,1k,1000,43,5.87,0.858948772460653,0.5726932261716688,0.7898403393004456,0.892423247305388,0.8953746242397653,0.08094347219936113,483633.3470054774,0.002067682069881493,cuda,1246.63671875,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen,1k,1000,44,5.51,0.8662151855680355,0.5750684743816339,0.7862695412187646,0.8906440030475861,0.8943119327291345,0.10118385512174216,483340.43726194254,0.0020689351084814326,cuda,1247.06640625,28.14501953125,57.0 +tfidf_ridge,log10_pgen,1k,1000,42,0.18,1.2769573384805994,0.8800498060420275,0.5355193690361973,0.7459276280012653,0.7755959954308803,0.24748512995339594,82046.7449062024,0.012188173962821092,cpu,647.671875,0.0, +esm2_8m_mlp,log10_pgen,10k,10000,42,13.12,0.7892901399554304,0.5117788675519823,0.8225450219460526,0.9077396313574956,0.922136015327215,0.06535016127457582,751323.1196541691,0.0013309852629855125,cuda,1218.4375,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen,10k,10000,43,11.67,0.7910012590527175,0.5090540726294859,0.821774770988962,0.908184530983501,0.9233618174289024,0.09305254228575181,707665.0616763182,0.0014130978822540683,cuda,1217.71484375,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen,10k,10000,44,12.8,0.7937840810408098,0.5114952986728927,0.8205185365266188,0.9069227656983063,0.9224128719321808,0.07858451337605844,691845.9266758637,0.0014454085244163175,cuda,1217.4921875,23.45751953125,60.0 +irrm_pretrained,log10_pgen,10k,10000,42,22.1,0.49633429136522883,0.3364633920574051,0.929828074355773,0.9788465234302741,0.9769528791960383,0.2582476449257332,79422.11814450966,0.012590951026771734,cuda,1614.30859375,255.36376953125,31.0 +irrm_pretrained,log10_pgen,10k,10000,43,25.12,0.5193359156545152,0.36831604542598495,0.9231734116100119,0.978126930098433,0.9718469313258652,0.3012286896888966,78484.32285378146,0.012741398073383758,cuda,1594.49609375,255.36376953125,34.0 +irrm_pretrained,log10_pgen,10k,10000,44,32.11,0.5252833859981841,0.3584885153591458,0.9214036891479529,0.9735898112520421,0.9643449436380355,0.2684473076598781,78905.69137201828,0.012673357049560336,cuda,1570.359375,255.36376953125,45.0 +irrm_pretrained,log10_pgen,10k,10000,45,22.37,0.5275998646962152,0.39018114650797886,0.9207089474596369,0.9729629104504584,0.958600728526982,0.29486428735038855,78146.06342119222,0.01279654989925971,cuda,1576.4296875,255.36376953125,31.0 +irrm_pretrained,log10_pgen,10k,10000,46,31.75,0.5158049435613763,0.37157946936378206,0.9242145502587643,0.9783526876739728,0.9718356785668307,0.3058773639971277,78655.62426811836,0.012713649014992712,cuda,1608.67578125,255.36376953125,43.0 +irrm_pretrained,log10_pgen,10k,10000,47,18.23,0.5332813946641329,0.35700741491474547,0.9189920399063798,0.97604323639183,0.9748606762478106,0.2938373550204833,78264.75742759714,0.012777143031780323,cuda,1581.08203125,255.36376953125,25.0 +irrm_pretrained,log10_pgen,10k,10000,48,24.91,0.4579817124860741,0.3023822172194899,0.9402536876166383,0.9782543855947416,0.9757298823937622,0.20419803463546618,78031.82970995203,0.01281528324681155,cuda,1594.7578125,255.36376953125,34.0 +irrm_pretrained,log10_pgen,10k,10000,49,25.92,0.4976124581595938,0.3438377496106781,0.9294661936068322,0.9784862915435858,0.9759692349627562,0.27500778441402923,78638.12822918102,0.012716477649183417,cuda,1580.44921875,255.36376953125,35.0 +irrm_pretrained,log10_pgen,10k,10000,50,19.67,0.5425341893413741,0.39877899006847456,0.9161565664369719,0.9734359952461339,0.9633476928509621,0.321624207014219,77982.10842038045,0.01282345425452298,cuda,1586.1640625,255.36376953125,27.0 +irrm_pretrained,log10_pgen,10k,10000,51,13.38,0.5341332976578411,0.37650608358281984,0.9187330170116145,0.9667227466615107,0.9633437546517095,0.23337067889230595,78704.76482654782,0.012705711048166311,cuda,1613.80078125,255.36376953125,18.0 +irrm_scratch,log10_pgen,10k,10000,42,20.95,0.498671630386942,0.32079206399879684,0.9291656104710792,0.9657117315186489,0.9694134769384374,0.09424418102101653,78508.39615413654,0.012737491134536582,cuda,1583.9453125,255.36376953125,28.0 +irrm_scratch,log10_pgen,10k,10000,43,33.65,0.4706407993114887,0.29323152256924806,0.936905140469252,0.9740480810857403,0.9781106763867039,0.16780619885200848,79078.70122623147,0.0126456300431536,cuda,1602.7890625,255.36376953125,46.0 +irrm_scratch,log10_pgen,10k,10000,44,33.05,0.5079075377997477,0.32965094824733154,0.9265174619388699,0.9712447393219146,0.9755254688138308,0.22044626558539696,78956.04437864256,0.012665274810429812,cuda,1590.59765625,255.36376953125,45.0 +irrm_scratch,log10_pgen,10k,10000,45,27.82,0.47713739863574445,0.2912506286196337,0.9351512291587034,0.9718732546431008,0.9750787340415884,0.11394230996095928,79112.17619742312,0.012640279259977842,cuda,1582.5078125,255.36376953125,39.0 +irrm_scratch,log10_pgen,10k,10000,46,29.27,0.4579175921407287,0.29244273401005727,0.9402704161710751,0.9708733181295428,0.9743106241169605,-0.0422780559037735,78999.96939361731,0.012658232752186277,cuda,1574.63671875,255.36376953125,40.0 +irrm_scratch,log10_pgen,10k,10000,47,36.55,0.465288916422154,0.3107330220285725,0.9383319456243314,0.9721828697208292,0.9783745998124969,0.1545819731891662,78634.00965443617,0.012717143693862043,cuda,1587.2578125,255.36376953125,51.0 +irrm_scratch,log10_pgen,10k,10000,48,24.69,0.4868968973441002,0.32282410811390244,0.9324712289538608,0.9666661177257057,0.9732218700912197,0.08242033868522627,79166.7916318792,0.012631559008352135,cuda,1600.56640625,255.36376953125,34.0 +irrm_scratch,log10_pgen,10k,10000,49,30.38,0.49955472550048546,0.3238743178975122,0.9289145077923109,0.9669152473849376,0.9734628057385909,0.14312938242896417,78494.10538506562,0.012739810143632277,cuda,1578.953125,255.36376953125,41.0 +irrm_scratch,log10_pgen,10k,10000,50,23.59,0.5021787474881457,0.3100156924156718,0.9281657618398458,0.9687455318772152,0.9699728610796904,0.16583821028234638,78789.2636408248,0.012692084603794266,cuda,1612.04296875,255.36376953125,32.0 +irrm_scratch,log10_pgen,10k,10000,51,25.4,0.4806922198139846,0.28437546648930967,0.9341813427954785,0.971277292615862,0.9743939805805218,0.09846404782942433,78746.8403341449,0.012698922213979886,cuda,1590.25,255.36376953125,35.0 +sceptr_mlp,log10_pgen,10k,10000,42,10.85,0.8197409601112305,0.5322377064815003,0.8085884658779793,0.900190469155707,0.9106309096825668,0.07351937350297826,1045992.739596697,0.0009560295804592005,cuda,1184.78515625,21.20751953125,60.0 +sceptr_mlp,log10_pgen,10k,10000,43,10.9,0.8232791649384097,0.5329595503445926,0.8069325407912433,0.8993982334295626,0.910789330836775,0.07795665485235405,1002912.2882333912,0.0009970961685607412,cuda,1178.8515625,21.20751953125,60.0 +sceptr_mlp,log10_pgen,10k,10000,44,11.99,0.8213314982545481,0.5383096758524992,0.8078449561119915,0.8998291633130819,0.9092193480549974,0.06639299892483241,1015685.420632926,0.0009845568122626475,cuda,1184.00390625,21.20751953125,60.0 +tcr_bert_mlp,log10_pgen,10k,10000,42,13.04,0.5491397489252545,0.3589531191073769,0.9141024862729251,0.9565472815665261,0.9603263904726105,0.028682170324545696,483779.5195345584,0.002067057325953142,cuda,1284.14453125,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen,10k,10000,43,13.05,0.5525621465801277,0.360994787679031,0.9130284736191095,0.9561370066319027,0.9611197559421756,0.022896591890857958,485116.2109191465,0.0020613617469210244,cuda,1278.78515625,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen,10k,10000,44,12.88,0.5531349877060555,0.3610178326203211,0.9128480534001411,0.9561711841901395,0.9604234621650158,0.03684545612776406,471155.4772059833,0.0021224416320704525,cuda,1266.859375,28.14501953125,60.0 +tfidf_ridge,log10_pgen,10k,10000,42,0.58,0.7659173935252817,0.5084553048733631,0.8328991356685383,0.9129315175825418,0.9404240652599085,0.03423890204930655,79360.64652403284,0.01260070379715429,cpu,652.96484375,0.0, +esm2_8m_mlp,log10_pgen,all,79730,42,62.4,0.5731677057093365,0.3493723214550509,0.9064210307249039,0.9536062412202215,0.9675506650848761,0.07753757589264149,737700.7938537622,0.0013555631339041703,cuda,1292.87109375,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen,all,79730,43,63.05,0.5722816128637928,0.3497724788661413,0.9067101452692726,0.9532051754785794,0.9667819816787405,0.06594397495518078,744636.9475580577,0.0013429363171937319,cuda,1295.58203125,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen,all,79730,44,63.38,0.5729095378508662,0.34892790104491134,0.9065053119604066,0.9536574004133471,0.9674232426731361,0.08018154737074212,736760.647026297,0.0013572929065038773,cuda,1300.19921875,23.45751953125,60.0 +irrm_pretrained,log10_pgen,all,79730,42,123.85,0.3057602609523895,0.2083267044257563,0.9733696476549023,0.9917018462735949,0.9894025191718101,0.1588558915729577,77911.56266318262,0.012835065371786638,cuda,1596.3984375,254.86376953125,27.0 +irrm_pretrained,log10_pgen,all,79730,43,118.85,0.3227698346701869,0.19583600915331453,0.9703243178058237,0.9886645939613051,0.9832962906240956,0.11649408246044561,77761.84919341622,0.012859776489017266,cuda,1579.3515625,254.86376953125,26.0 +irrm_pretrained,log10_pgen,all,79730,44,211.48,0.28060503644475776,0.14739548233582567,0.9775712146156854,0.9929846069686562,0.9923312548184898,0.08320273340059216,76331.09878614836,0.013100820188657731,cuda,1584.953125,254.86376953125,47.0 +irrm_pretrained,log10_pgen,all,79730,45,215.6,0.2906337822506723,0.1697879711770245,0.9759393686285257,0.9930494157956323,0.9915791451404403,0.11848370121901147,76213.84869992676,0.013120974954791398,cuda,1603.67578125,254.86376953125,48.0 +irrm_pretrained,log10_pgen,all,79730,46,144.15,0.33693704512417527,0.2029260203209646,0.9676620595011232,0.989456937144701,0.985854278301268,0.07777895175197753,77506.95307456442,0.012902068270416523,cuda,1622.1328125,254.86376953125,32.0 +irrm_pretrained,log10_pgen,all,79730,47,198.11,0.28260784579697307,0.16704077865556358,0.977249902558756,0.9926284885670357,0.9930547129437717,0.12034238878009869,76224.02236325688,0.013119223690851051,cuda,1597.26171875,254.86376953125,44.0 +irrm_pretrained,log10_pgen,all,79730,48,157.15,0.2956240286589957,0.16988515750641128,0.9751060224184863,0.9918395106629686,0.9882366619574364,0.037514705427660146,77747.24938970283,0.012862191368180339,cuda,1585.90625,254.86376953125,35.0 +irrm_pretrained,log10_pgen,all,79730,49,202.71,0.2568489903064575,0.15425740537613158,0.9812081058492581,0.9950565011110328,0.9960072133460977,0.11595127294142848,76347.31414166577,0.013098037714129096,cuda,1580.203125,254.86376953125,45.0 +irrm_pretrained,log10_pgen,all,79730,50,99.96,0.39284684013561044,0.2914359988728974,0.9560396278441484,0.9858889109334719,0.9763490870794563,0.21786421304129522,77578.69993954332,0.012890136091211825,cuda,1600.61328125,254.86376953125,21.0 +irrm_pretrained,log10_pgen,all,79730,51,115.95,0.33713538320078806,0.21844640571312485,0.9676239768220666,0.9906641206786925,0.9865010824105642,0.13035446438857884,77558.41612855578,0.012893507241592779,cuda,1579.63671875,254.86376953125,25.0 +irrm_scratch,log10_pgen,all,79730,42,114.17,0.2719535505097029,0.16653379054223957,0.9789329222438513,0.9919265890391352,0.9926972842928662,0.09609930151292388,77975.04366490629,0.012824616095086117,cuda,1585.13671875,254.86376953125,26.0 +irrm_scratch,log10_pgen,all,79730,43,148.87,0.2806572979150532,0.1858581395148315,0.977562859309801,0.9949507777304938,0.995145779501445,0.05192284273870285,78143.844898836,0.012796913196357142,cuda,1572.2421875,254.86376953125,33.0 +irrm_scratch,log10_pgen,all,79730,44,182.86,0.2834810367182595,0.16944260740206932,0.9771091006088122,0.9950294500699801,0.9956005429373097,0.10711986412509826,76530.48611350135,0.013066688202096523,cuda,1584.03515625,254.86376953125,41.0 +irrm_scratch,log10_pgen,all,79730,45,93.79,0.3284325426622101,0.22519455615620482,0.9692739171693161,0.988786459724294,0.9899692333708238,0.13219256068879753,77059.23870669039,0.012977029319044889,cuda,1574.84765625,254.86376953125,21.0 +irrm_scratch,log10_pgen,all,79730,46,114.35,0.27585077654973195,0.18547613126024184,0.9783247930033384,0.9915338350103832,0.9910657074195383,0.025506009252100864,77201.34301823228,0.012953142534888735,cuda,1576.3515625,254.86376953125,25.0 +irrm_scratch,log10_pgen,all,79730,47,198.35,0.33604797728519187,0.21124316666040283,0.967832493067254,0.9958645289807903,0.9960726435747845,0.10184083611487582,76554.00202823328,0.013062674367189815,cuda,1596.35546875,254.86376953125,44.0 +irrm_scratch,log10_pgen,all,79730,48,95.99,0.32095970522074435,0.2228937454828952,0.9706562335892933,0.9904940076818866,0.9904438010414041,0.17860603614144913,77715.69472656505,0.012867413763955924,cuda,1575.47265625,254.86376953125,21.0 +irrm_scratch,log10_pgen,all,79730,49,142.4,0.324897288809792,0.19739056002251965,0.969931829346724,0.9933873020382896,0.9936995466077865,0.12007436457931936,77726.9541308783,0.012865549810638133,cuda,1571.75390625,254.86376953125,31.0 +irrm_scratch,log10_pgen,all,79730,50,89.87,0.3152489172546755,0.21929226072923863,0.9716911622598429,0.9860163867885153,0.9893431929980233,0.04317832128811549,76554.0589797288,0.01306266464936622,cuda,1584.25390625,254.86376953125,20.0 +irrm_scratch,log10_pgen,all,79730,51,131.13,0.3766005567885428,0.24336180453560238,0.9596004295556182,0.9933946378049742,0.9936023254985136,0.13065965375684127,77555.69803371314,0.012893959120389892,cuda,1576.484375,254.86376953125,29.0 +sceptr_mlp,log10_pgen,all,79730,42,57.69,0.5766064944696535,0.3428484145267614,0.905294785689099,0.9535790459641675,0.9663776368101202,0.08926576441701496,1035103.8993361475,0.0009660865934727316,cuda,1194.8359375,21.20751953125,60.0 +sceptr_mlp,log10_pgen,all,79730,43,64.55,0.584833564731118,0.3453063517799877,0.9025729817429096,0.95246948351646,0.9672135688288871,0.10801653991967043,1061797.5338926413,0.0009417991359745523,cuda,1194.54296875,21.20751953125,60.0 +sceptr_mlp,log10_pgen,all,79730,44,59.51,0.5770612206641812,0.3445521396506703,0.9051453530234606,0.9536648620712153,0.9669085617186784,0.09536860715000833,991191.4022597936,0.0010088868786796617,cuda,1190.8828125,21.20751953125,60.0 +tcr_bert_mlp,log10_pgen,all,79730,42,79.88,0.3777564977751135,0.23458992981807736,0.9593520433627332,0.979747865553209,0.9847126718927225,0.024266606714189013,478528.72105608066,0.0020897387262212123,cuda,1670.765625,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen,all,79730,43,70.85,0.375505857040682,0.23313485418218585,0.9598349545543144,0.9801372299422181,0.9855585447736372,0.024686716803757173,481572.8058755746,0.0020765292138575885,cuda,1662.76171875,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen,all,79730,44,71.6,0.3740375957016166,0.23232463675184953,0.9601484382877309,0.9802273073279212,0.9850856333747358,0.02237037464962718,486510.9071794783,0.0020554523757698425,cuda,1665.94921875,28.14501953125,60.0 +tfidf_ridge,log10_pgen,all,79730,42,8.4,0.6210225099572112,0.41454671189516185,0.8901425494415403,0.9434950444517985,0.9563135820470278,0.008684176700797186,68718.25314155618,0.014552174339182484,cpu,714.09765625,0.0, +esm2_8m_mlp,log10_pgen_1mm,1k,1000,42,5.11,0.8638028049753916,0.568791218194971,0.7131110578070958,0.8469102887490368,0.8667108625359236,0.10372949114561908,734687.1376026621,0.0013611235978120881,cuda,1199.7890625,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen_1mm,1k,1000,43,5.23,0.8619774244542309,0.5692703601559694,0.7143222790817589,0.8478099524307426,0.8669547470373846,0.1071127415864139,727775.8588480044,0.0013740494244792606,cuda,1200.953125,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen_1mm,1k,1000,44,5.08,0.8590033542150199,0.5676105701335532,0.7162902189773253,0.8479008164271059,0.8685698320651293,0.08222814137429608,745106.2393934651,0.00134209049277862,cuda,1192.59765625,23.45751953125,60.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,42,12.53,0.6529562585935635,0.40768538159968476,0.8360721408566971,0.9442740681192724,0.949476176152167,0.3004492179255113,80240.86719978109,0.012462477474355214,cuda,1550.8359375,255.98876953125,52.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,43,5.16,0.6400113298790213,0.40304623811145285,0.8425074866906714,0.9312456988056497,0.9373496156964188,0.2231928880169614,80682.86935446602,0.012394204717815323,cuda,1549.69921875,255.98876953125,20.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,44,6.26,0.6200875499689317,0.3844985011699642,0.8521604565547631,0.9459296723810654,0.9503241412598608,0.2694864386133546,80083.91447115551,0.012486902102671045,cuda,1553.51953125,255.98876953125,25.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,45,6.54,0.6310536518942342,0.3986318635819051,0.8468852047631974,0.9443783978709075,0.9489821784491898,0.2896516474185923,77981.26403690618,0.01282359310727164,cuda,1549.2734375,255.98876953125,26.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,46,6.59,0.6195555852385934,0.3962891510289275,0.8524140068427615,0.9443500979728823,0.9469694023446114,0.27819833731219673,78697.9342770863,0.012706813834262993,cuda,1562.43359375,255.98876953125,26.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,47,6.33,0.6266547317455818,0.3904620047620022,0.849012416025269,0.943407263239133,0.9469707393444504,0.263707425816376,78310.37693883583,0.012769699739551095,cuda,1565.94921875,255.98876953125,25.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,48,5.67,0.6236619614132954,0.39930113323908234,0.8504511417494165,0.9389343391116015,0.9458162046101952,0.25761830099041527,78693.53284131692,0.012707524543553905,cuda,1543.70703125,255.98876953125,22.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,49,4.71,0.6499044803178063,0.40785618003239943,0.8376008876652727,0.9264759291777597,0.9328795005102539,0.17716616132361881,78385.95400081249,0.012757387630820119,cuda,1555.32421875,255.98876953125,18.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,50,5.9,0.5994330346497347,0.380280920901402,0.8618452136538101,0.939269318647871,0.9421541893901585,0.18829553681659392,79014.97934627304,0.012655828151490463,cuda,1548.5,255.98876953125,23.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,51,4.96,0.5994008301275314,0.3870256531801975,0.8618600579787317,0.9348727785466787,0.9388998800257852,0.17224279149582014,78674.18185650515,0.012710650132007891,cuda,1560.15234375,255.98876953125,19.0 +irrm_scratch,log10_pgen_1mm,1k,1000,42,6.48,0.7756210635607437,0.488775886202463,0.768695661554632,0.9158074553176181,0.9222819561703773,0.3852233068305189,79599.61677741556,0.012562874552478065,cuda,1574.03125,255.98876953125,25.0 +irrm_scratch,log10_pgen_1mm,1k,1000,43,11.15,0.5616567093437265,0.34603153059101177,0.8787095823268849,0.9386583217064812,0.9474188542893907,0.07016255632983405,79666.25790176584,0.01255236566066742,cuda,1544.94921875,255.98876953125,44.0 +irrm_scratch,log10_pgen_1mm,1k,1000,44,6.88,0.7058828649367268,0.43409787223543767,0.8084201324866083,0.9140735719556143,0.9168842625543453,0.21586295803326302,79344.34843686518,0.012603292102092269,cuda,1563.3046875,255.98876953125,27.0 +irrm_scratch,log10_pgen_1mm,1k,1000,45,7.91,0.6740001737872768,0.4130208410802216,0.8253355155724036,0.9192456087003898,0.926908100052709,0.20133236035721191,77125.38840246688,0.012965899047167906,cuda,1583.46875,255.98876953125,31.0 +irrm_scratch,log10_pgen_1mm,1k,1000,46,7.63,0.682312413436616,0.42114535487571414,0.8210007817215697,0.9154231043861375,0.9224241754873767,0.19040918863685738,77714.31885101968,0.01286764157216671,cuda,1544.8125,255.98876953125,30.0 +irrm_scratch,log10_pgen_1mm,1k,1000,47,9.5,0.65108646567971,0.3912826224971895,0.8370096379269262,0.9316193931638769,0.9398473653548969,0.21678677651645972,78072.18619042147,0.012808658868101327,cuda,1571.74609375,255.98876953125,38.0 +irrm_scratch,log10_pgen_1mm,1k,1000,48,6.33,0.8344811555092708,0.5589508598595957,0.7322572888207102,0.9108633601503565,0.91322708473897,0.4732333279414907,78272.52281380334,0.012775875416444992,cuda,1548.8046875,255.98876953125,25.0 +irrm_scratch,log10_pgen_1mm,1k,1000,49,6.92,0.82629497434108,0.5637715445233729,0.7374845841235782,0.912489113508546,0.9179172506572928,0.48450341909887745,77540.4721325073,0.01289649098719854,cuda,1551.09375,255.98876953125,27.0 +irrm_scratch,log10_pgen_1mm,1k,1000,50,8.19,0.6103351399149818,0.3876196098409824,0.8567741727273026,0.9259979083199648,0.9311154743696028,0.019025652453153675,75476.72065839717,0.013249118288086949,cuda,1547.109375,255.98876953125,31.0 +irrm_scratch,log10_pgen_1mm,1k,1000,51,8.32,0.6251580173699457,0.3920347913893575,0.8497327980864278,0.9250595921373767,0.9335045905874384,0.0355205504947733,78642.68759238251,0.012715740402758845,cuda,1561.4453125,255.98876953125,33.0 +sceptr_mlp,log10_pgen_1mm,1k,1000,42,4.88,0.8729629962973278,0.585332278608628,0.7069941717233896,0.8422360133559548,0.855733573185971,0.07803339644954202,1052124.6577599368,0.0009504577167966822,cuda,1173.1171875,21.20751953125,60.0 +sceptr_mlp,log10_pgen_1mm,1k,1000,43,3.81,0.872707884195867,0.5878114142571446,0.7071654009980505,0.8424537096287541,0.8531245639028813,0.08160378789314612,1025225.5230945263,0.0009753951471883124,cuda,1171.46875,21.20751953125,43.0 +sceptr_mlp,log10_pgen_1mm,1k,1000,44,4.35,0.875039007473917,0.5886985468686378,0.7055989083216988,0.8411393439083864,0.8544929176996288,0.07046295775930289,1009110.5557346264,0.0009909716971219333,cuda,1169.109375,21.20751953125,50.0 +tcr_bert_mlp,log10_pgen_1mm,1k,1000,42,5.29,0.6852645656848846,0.43437482691831775,0.8194484836357487,0.9093048373927972,0.9256513310486022,0.08790968486708374,469312.9234898688,0.0021307744789209647,cuda,1243.875,28.14501953125,47.0 +tcr_bert_mlp,log10_pgen_1mm,1k,1000,43,4.25,0.6829331715832121,0.434071752462488,0.8206749316672025,0.9093094169413265,0.9264527464334696,0.0736007711157669,468664.7860406887,0.0021337212220445805,cuda,1238.6015625,28.14501953125,38.0 +tcr_bert_mlp,log10_pgen_1mm,1k,1000,44,5.38,0.6872351224425811,0.4356053796900238,0.8184085974905176,0.9083054888112506,0.9259385459713246,0.076839266100189,484858.9313030664,0.002062455562719003,cuda,1242.58203125,28.14501953125,57.0 +tfidf_ridge,log10_pgen_1mm,1k,1000,42,0.19,1.085296409794027,0.73535530121258,0.5471218978219647,0.750774560158103,0.795403603438914,0.18915308819686613,80950.28937014837,0.012353260350033597,cpu,643.671875,0.0, +esm2_8m_mlp,log10_pgen_1mm,10k,10000,42,13.41,0.6260295155704774,0.3933335705938847,0.8493135476799875,0.9221821094460882,0.9420381756814984,0.04854568221982546,739495.3955814017,0.0013522734637364253,cuda,1219.70703125,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen_1mm,10k,10000,43,11.81,0.6223697320332455,0.3882320646994455,0.8510702310367887,0.9238662517056982,0.9434740284740887,0.0649370734917326,735547.6744995526,0.0013595311829112556,cuda,1225.6953125,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen_1mm,10k,10000,44,13.37,0.6259706667490265,0.3916924130701659,0.8493418763868927,0.9223776550508551,0.9421140639783299,0.053312992388272695,731884.5499928056,0.0013663357151204106,cuda,1219.30859375,23.45751953125,60.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,42,23.04,0.3522035785638962,0.23605209785853865,0.9523050807498069,0.983930424800023,0.9843882615672284,0.18039631196874276,79187.71578948178,0.012628221309710091,cuda,1592.3515625,255.36376953125,32.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,43,29.11,0.38753899215660315,0.27109869129044467,0.9422548651039686,0.9811142357156362,0.9703975743994855,0.2089216769812675,79484.49264258641,0.012581070429632676,cuda,1595.453125,255.36376953125,39.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,44,32.26,0.3822850747742425,0.25218843297052235,0.9438099686538676,0.9813273427294593,0.9755213089198129,0.18395725029628748,78665.73866989763,0.012712014364935492,cuda,1596.4609375,255.36376953125,45.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,45,17.65,0.4041191447823732,0.27990819980915715,0.9372081267064315,0.979920025727262,0.9765428583968161,0.2257876246112694,78729.95469238408,0.012701645820923287,cuda,1589.9296875,255.36376953125,24.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,46,31.18,0.3911848445918753,0.26226686425507795,0.9411632560109777,0.9829105055368497,0.9776255144294108,0.20629808755274887,77487.35419934198,0.012905331590331831,cuda,1591.9296875,255.36376953125,43.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,47,23.39,0.38348089009514835,0.2531676301319391,0.9434578858919235,0.9826728115818364,0.9819501282337215,0.19949019843057192,78626.90485599045,0.012718292826502017,cuda,1600.2109375,255.36376953125,32.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,48,25.34,0.31816722084810306,0.20541534208333775,0.9610779722310352,0.9830906838143463,0.9823737812284563,0.10997886851599012,78570.88446529236,0.01272736086408365,cuda,1581.47265625,255.36376953125,34.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,49,29.76,0.3814527636735218,0.2392016173657349,0.9440543761849206,0.9826898812367412,0.9814819849844005,0.17505504218945128,77532.6625809653,0.012897789998579327,cuda,1575.05078125,255.36376953125,41.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,50,22.5,0.40909544792676755,0.261909274021916,0.9356521733237381,0.983110380156732,0.9833059108264491,0.22079621710108355,78582.57884899958,0.012725466822888961,cuda,1583.88671875,255.36376953125,30.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,51,13.44,0.4038788122079963,0.2812596679686494,0.9372827900596038,0.9746190105546972,0.9714999790318952,0.16937099385437487,77728.24803117815,0.012865335644756622,cuda,1569.32421875,255.36376953125,18.0 +irrm_scratch,log10_pgen_1mm,10k,10000,42,20.33,0.36921385599324186,0.23028669602259502,0.9475868125935152,0.9744331202623125,0.9801626054392585,0.03312593973357653,79122.29371026848,0.012638662924280468,cuda,1585.8984375,255.36376953125,28.0 +irrm_scratch,log10_pgen_1mm,10k,10000,43,21.57,0.36219569504086635,0.22513798777308272,0.9495604553291446,0.9756419521768758,0.9802466698170175,0.04881839728880603,76747.8627017076,0.013029678805345422,cuda,1586.53515625,255.36376953125,29.0 +irrm_scratch,log10_pgen_1mm,10k,10000,44,18.83,0.3831162822600322,0.23419808420311247,0.9435653535535624,0.9728618066319854,0.9764946080845648,0.07214146335570987,80006.72679311575,0.012498949026946643,cuda,1576.73046875,255.36376953125,25.0 +irrm_scratch,log10_pgen_1mm,10k,10000,45,28.99,0.35224366238484983,0.22376578294680224,0.9522942239413621,0.9760330880069608,0.9834552471268139,-0.012863246857717952,76926.3501805443,0.012999446843026143,cuda,1571.3359375,255.36376953125,39.0 +irrm_scratch,log10_pgen_1mm,10k,10000,46,39.09,0.33900270671157107,0.20915864891427285,0.9558133658824247,0.9800842419254892,0.9855767889778858,0.10221936446662258,78725.08869905554,0.012702430908940939,cuda,1602.68359375,255.36376953125,53.0 +irrm_scratch,log10_pgen_1mm,10k,10000,47,35.52,0.33813072800923394,0.21959607994470112,0.9560403861991751,0.9801886364299571,0.9853709016134247,0.10500751758902857,77535.64982109178,0.01289729308140748,cuda,1579.80859375,255.36376953125,49.0 +irrm_scratch,log10_pgen_1mm,10k,10000,48,32.9,0.3378822660799348,0.2061434139007981,0.9561049664141406,0.9787600364711352,0.9840258078648648,0.06777958021349939,78439.5066056244,0.012748677844542899,cuda,1582.52734375,255.36376953125,45.0 +irrm_scratch,log10_pgen_1mm,10k,10000,49,17.22,0.5211322195049435,0.34843281043792573,0.8955807814829972,0.9674054163214658,0.970678072425513,0.30572248198059976,77709.23790181913,0.012868482911432472,cuda,1572.69921875,255.36376953125,23.0 +irrm_scratch,log10_pgen_1mm,10k,10000,50,19.9,0.37348189977801466,0.23958918079854716,0.9463680351794995,0.9731939619808125,0.9802889403153295,0.04181760784932039,78409.30288658597,0.012753588709319811,cuda,1582.83203125,255.36376953125,27.0 +irrm_scratch,log10_pgen_1mm,10k,10000,51,25.7,0.37806560313552356,0.22143062541018307,0.9450435182028492,0.9757677766800191,0.9812830800850978,0.05084201869418621,77880.5342036482,0.012840179002690461,cuda,1570.671875,255.36376953125,35.0 +sceptr_mlp,log10_pgen_1mm,10k,10000,42,11.75,0.6371332141671232,0.39614468171809036,0.8439207819168444,0.9190377199718023,0.9394482016893153,0.04125073329136141,1057753.3728966473,0.0009453999633785235,cuda,1184.921875,21.20751953125,60.0 +sceptr_mlp,log10_pgen_1mm,10k,10000,43,10.84,0.6407606075461617,0.39765060478869485,0.842138509788058,0.9183379902101502,0.9390568683069896,0.04882956748748253,1009009.5155340733,0.0009910709310513245,cuda,1186.05078125,21.20751953125,60.0 +sceptr_mlp,log10_pgen_1mm,10k,10000,44,12.35,0.637657661090663,0.3986599607309693,0.8436637274958344,0.9191095715446563,0.9392093880388864,0.04484058808243563,1051063.8584379375,0.0009514169781141299,cuda,1186.1640625,21.20751953125,60.0 +tcr_bert_mlp,log10_pgen_1mm,10k,10000,42,13.18,0.4111318934993246,0.25845667002793843,0.9350099417710966,0.967297966003536,0.9756017679174287,0.008384003869639601,479165.2697996234,0.0020869626056541587,cuda,1280.30859375,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen_1mm,10k,10000,43,14.65,0.41350585090839187,0.2565953029194565,0.9342572438418231,0.966898297425101,0.97593350933347,0.01531694211326789,485958.66084314627,0.0020577882041756053,cuda,1279.98046875,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen_1mm,10k,10000,44,13.26,0.41106930870267516,0.25502984257628103,0.9350297265655259,0.9675431906442596,0.97598224972199,0.026138567755919037,470986.4438187892,0.002123203359935233,cuda,1285.44921875,28.14501953125,60.0 +tfidf_ridge,log10_pgen_1mm,10k,10000,42,0.55,0.679848528075636,0.44560910401677106,0.8222912089245368,0.9069106930892485,0.9418412049600424,0.018927686789477982,79617.02327020983,0.012560127959144239,cpu,651.0703125,0.0, +esm2_8m_mlp,log10_pgen_1mm,all,79730,42,68.38,0.45471149866490485,0.2655386843892589,0.9205019524469311,0.9605820531819961,0.9763708989995882,0.05194563645716011,732619.9504161674,0.0013649641938251157,cuda,1291.36328125,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen_1mm,all,79730,43,65.72,0.45216768528691753,0.26522786955099303,0.9213889435621597,0.9608687465987668,0.9762787240690396,0.05328308914404441,715461.1194133183,0.001397699990769596,cuda,1297.3359375,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen_1mm,all,79730,44,69.37,0.4549245239649274,0.26568223603381097,0.9204274477889403,0.9606880913298007,0.9760459387204964,0.05959356687875681,732726.0886193139,0.0013647664734912253,cuda,1296.890625,23.45751953125,60.0 +irrm_pretrained,log10_pgen_1mm,all,79730,42,197.32,0.2010198735884292,0.11098549405704658,0.9844631640433253,0.9957314596428587,0.9967040399178354,0.061738928489644775,77704.89332565277,0.012869202404140863,cuda,1599.6015625,254.86376953125,43.0 +irrm_pretrained,log10_pgen_1mm,all,79730,43,201.51,0.19844419233978014,0.10745138967990057,0.9848587623496916,0.994074687121133,0.9932460677425246,0.00047019676069362895,76702.7717764188,0.013037338506030836,cuda,1580.54296875,254.86376953125,45.0 +irrm_pretrained,log10_pgen_1mm,all,79730,44,211.1,0.22832235688634206,0.12134789785717585,0.9799561350952133,0.9932992236706577,0.9903972598369296,0.008108554940664215,77555.84641114557,0.012893934452068462,cuda,1590.78515625,254.86376953125,45.0 +irrm_pretrained,log10_pgen_1mm,all,79730,45,218.6,0.214663543875197,0.10917544386411444,0.9822825515055111,0.9931988998785725,0.9902501268322028,0.01889345747215228,77644.95169252451,0.012879137383715803,cuda,1587.48828125,254.86376953125,49.0 +irrm_pretrained,log10_pgen_1mm,all,79730,46,202.2,0.19382356300065293,0.10396384714341159,0.9855556589259883,0.9949288895266005,0.996159075626305,0.03274914680978691,76330.32244284997,0.013100953434969701,cuda,1582.69140625,254.86376953125,45.0 +irrm_pretrained,log10_pgen_1mm,all,79730,47,264.54,0.20821800322930945,0.10973956171749896,0.9833305548661806,0.9948324953928825,0.9952488381802871,0.008130821087864773,77603.79416967419,0.01288596789241495,cuda,1589.7421875,254.86376953125,56.0 +irrm_pretrained,log10_pgen_1mm,all,79730,48,157.77,0.21873573252501982,0.12205762993538553,0.9816039721229748,0.9933163416010036,0.9946624243808212,0.010566345166154303,76365.73780152647,0.013094877739530084,cuda,1580.36328125,254.86376953125,35.0 +irrm_pretrained,log10_pgen_1mm,all,79730,49,140.76,0.21032706256829567,0.1248383391297992,0.9829911519147718,0.9953674218542565,0.995268712887437,0.06345402711592327,76480.93193590423,0.013075154482140229,cuda,1578.29296875,254.86376953125,31.0 +irrm_pretrained,log10_pgen_1mm,all,79730,50,98.67,0.3050127360048115,0.2223549543074473,0.9642298684623126,0.9866466817661931,0.9767291951976588,0.1524269310931368,77514.2341875952,0.012900856345685635,cuda,1577.84375,254.86376953125,21.0 +irrm_pretrained,log10_pgen_1mm,all,79730,51,174.39,0.2025477173651247,0.11567008877527189,0.9842260922848269,0.9934815527437457,0.9921715207108387,-0.018485427401458462,77031.87929353648,0.012981638370646724,cuda,1575.97265625,254.86376953125,39.0 +irrm_scratch,log10_pgen_1mm,all,79730,42,100.9,0.19633344317316612,0.12458027456075794,0.9851791485170951,0.9931292366678238,0.9947710352522292,0.050957250625054906,77995.24701890766,0.0128212940944668,cuda,1589.96484375,254.86376953125,23.0 +irrm_scratch,log10_pgen_1mm,all,79730,43,117.63,0.24310322146472474,0.15923510689506007,0.9772769817622599,0.9934988890065156,0.9949599529712726,0.05080266396051272,78482.72969605918,0.012741656716996334,cuda,1590.64453125,254.86376953125,25.0 +irrm_scratch,log10_pgen_1mm,all,79730,44,93.54,0.2137935159265193,0.13860440255569756,0.9824258775437755,0.9917096398136698,0.9936772681741826,-0.005426926835017065,77045.82669410027,0.012979288339268017,cuda,1585.46484375,254.86376953125,21.0 +irrm_scratch,log10_pgen_1mm,all,79730,45,88.95,0.24961439845675001,0.16239639861369026,0.976043473157776,0.9894739139539525,0.9908986908663686,0.050949834173454445,76718.78796334841,0.013034616767899663,cuda,1587.9609375,254.86376953125,19.0 +irrm_scratch,log10_pgen_1mm,all,79730,46,114.28,0.24344421065623903,0.1654815839553585,0.9772131920875918,0.9932718413750967,0.9946481437848445,0.000411995501209572,76869.21650122834,0.01300910878913434,cuda,1570.91796875,254.86376953125,25.0 +irrm_scratch,log10_pgen_1mm,all,79730,47,162.25,0.2623964374412026,0.15666522343761038,0.9735271646889943,0.9947957802738262,0.9963339789925373,0.07214033209339996,77821.99642085207,0.01284983739805542,cuda,1580.6953125,254.86376953125,36.0 +irrm_scratch,log10_pgen_1mm,all,79730,48,94.34,0.2756881202965431,0.1874225016620911,0.9707772761679718,0.9914597073267497,0.9936345681852634,0.13766648666405745,76817.75721149794,0.013017823434323356,cuda,1566.234375,254.86376953125,21.0 +irrm_scratch,log10_pgen_1mm,all,79730,49,94.61,0.25040927656190526,0.16263100405986253,0.9758906547425884,0.9906799443983614,0.9925507522500705,0.11173074335760001,76616.09987191304,0.013052086985265526,cuda,1590.24609375,254.86376953125,21.0 +irrm_scratch,log10_pgen_1mm,all,79730,50,101.62,0.22903029253875162,0.14824262140358074,0.979831646498065,0.9914505658360079,0.9938290217871487,0.042517517343130445,77764.52751135893,0.012859333580518852,cuda,1592.57421875,254.86376953125,22.0 +irrm_scratch,log10_pgen_1mm,all,79730,51,80.82,0.27079802850363044,0.17755006597799897,0.9718047736339673,0.9866817276800025,0.9915138366409755,-0.04579108726461479,76768.9593331389,0.013026098161113529,cuda,1578.109375,254.86376953125,18.0 +sceptr_mlp,log10_pgen_1mm,all,79730,42,58.13,0.4422856412715362,0.24850351471532656,0.9247874575108979,0.9628673381017279,0.9787701463480561,0.06165277549235565,1026396.787782811,0.0009742820826243693,cuda,1195.98828125,21.20751953125,60.0 +sceptr_mlp,log10_pgen_1mm,all,79730,43,58.16,0.4532838787285646,0.2528874805480784,0.9210003556570947,0.9610954079350825,0.9781999681797215,0.07523289538804193,1036979.2601728663,0.0009643394409192893,cuda,1188.16796875,21.20751953125,60.0 +sceptr_mlp,log10_pgen_1mm,all,79730,44,58.97,0.4438398228288394,0.251492624705254,0.9242579384996292,0.9624833450691401,0.9780063631797343,0.05573852857717545,1067511.1234785551,0.0009367583887476829,cuda,1195.93359375,21.20751953125,60.0 +tcr_bert_mlp,log10_pgen_1mm,all,79730,42,70.1,0.27590123549840345,0.1670392069257,0.9707320786177651,0.9854458261580342,0.9908209029750301,0.015594443637952673,486114.4564435013,0.0020571287003397836,cuda,1669.57421875,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen_1mm,all,79730,43,72.57,0.27667355629747115,0.1669020757783433,0.9705679919147793,0.985455078383879,0.9908143589668412,0.01800416090842117,481155.3265214974,0.002078330935728135,cuda,1670.77734375,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen_1mm,all,79730,44,70.66,0.277176078142094,0.16742615472770725,0.9704609801791605,0.9854239740579062,0.9906712868241694,0.01896631708478492,471095.082362874,0.0021227137311310806,cuda,1664.90234375,28.14501953125,60.0 +tfidf_ridge,log10_pgen_1mm,all,79730,42,8.65,0.5542270263983469,0.36706299488768934,0.8818972557236601,0.9391297435576847,0.9565962719889497,0.0035393838152091624,67434.86203217156,0.01482912502323982,cpu,717.2109375,0.0, diff --git a/results/pgen/pgen_summary.csv b/results/pgen/pgen_summary.csv new file mode 100644 index 0000000..28f519b --- /dev/null +++ b/results/pgen/pgen_summary.csv @@ -0,0 +1,37 @@ +target,train_subset,arm,seeds,train_size,rmse_mean,rmse_std,mae_mean,mae_std,r2_mean,r2_std,pearson_r_mean,pearson_r_std,spearman_rho_mean,spearman_rho_std,bias_mean,bias_std,training_seconds_mean,inference_seq_per_second_mean,peak_rss_mb_mean,peak_gpu_mb_mean +log10_pgen,1k,irrm_pretrained,10,1000,0.81540254540097,0.023651584767435244,0.5363845145896604,0.01761324618463643,0.8104657560730872,0.011097509580039815,0.9230052240535637,0.006701123055612074,0.9153326713831202,0.007112225627182163,0.32453740767930384,0.06286931463118653,6.049,79910.55138167983,1553.811328125,255.98876953125 +log10_pgen,1k,irrm_scratch,10,1000,0.8605530608233808,0.07907712609580614,0.5498540931749136,0.05496680052186431,0.7874514668265794,0.038402525773423256,0.9009418384515987,0.01869080920210034,0.900028262220985,0.021674349507537095,0.22982628971813562,0.12186389434072437,8.514,79743.93268363111,1554.234765625,255.98876953125 +log10_pgen,1k,tcr_bert_mlp,3,1000,0.8641840324486602,0.004571642536440778,0.5752700931258724,0.0026833632063272307,0.7872667330672515,0.0022475456545458564,0.8909852972973177,0.0013013139211785886,0.8942322659764437,0.0011842031778845218,0.09581540536927798,0.013044363649120726,5.343333333333334,481114.6119121562,1245.4192708333333,28.14501953125 +log10_pgen,1k,esm2_8m_mlp,3,1000,1.0739526536751358,0.003907657319455539,0.7270157322746377,0.0014533144669035686,0.6714594832157674,0.0023890385099123545,0.8222751241532199,0.000640771264232261,0.8331926446969974,0.0008683152203000417,0.1260789402176629,0.021590187659655407,4.8500000000000005,729356.5891250572,1196.0377604166667,23.45751953125 +log10_pgen,1k,sceptr_mlp,3,1000,1.097410607246145,0.0030375912434647237,0.7608200931599853,0.0017949689951640149,0.6569516321168019,0.0018981335130340972,0.8125005378142754,0.000726167976573551,0.806905067511356,0.0020307833308386533,0.1050660153327495,0.012487634130457757,3.7866666666666666,1065071.8303443855,1171.9544270833333,21.20751953125 +log10_pgen,1k,tfidf_ridge,1,1000,1.2769573384805994,,0.8800498060420275,,0.5355193690361973,,0.7459276280012653,,0.7755959954308803,,0.24748512995339594,,0.18,82046.7449062024,647.671875,0.0 +log10_pgen,10k,irrm_scratch,10,10000,0.48468864648435217,0.017093841387999915,0.3079190504390036,0.01637888879825529,0.9330074645214808,0.004709659869646582,0.9699538184023497,0.0027724153081702158,0.9741865097600041,0.0029359356023705156,0.1198594851930735,0.0708025137861426,28.534999999999997,78848.62979964017,1590.3546875,255.36376953125 +log10_pgen,10k,irrm_pretrained,10,10000,0.5149901453584536,0.025042621888465344,0.3603541024110505,0.027959074461221787,0.9242930177410577,0.007147733530583638,0.9754821518342982,0.0038646581191170513,0.9696832402360751,0.0066390374950387595,0.27567033535985275,0.035946894046752415,23.556,78523.54086732784,1592.05234375,255.36376953125 +log10_pgen,10k,tcr_bert_mlp,3,10000,0.5516122944038125,0.0021603582102378664,0.3603219131355763,0.001185466400222133,0.9133263377640586,0.0006781907878428965,0.956285157462856,0.00022764843655370855,0.9606232028599339,0.00043275795888619406,0.02947473944772257,0.0070081258362111026,12.99,480017.0692198961,1276.5963541666667,28.14501953125 +log10_pgen,10k,tfidf_ridge,1,10000,0.7659173935252817,,0.5084553048733631,,0.8328991356685383,,0.9129315175825418,,0.9404240652599085,,0.03423890204930655,,0.58,79360.64652403284,652.96484375,0.0 +log10_pgen,10k,esm2_8m_mlp,3,10000,0.7913584933496525,0.0022681686173861636,0.5107760796181203,0.0014980266616373488,0.8216127764872111,0.0010229088193688608,0.9076156426797676,0.0006399553131931383,0.9226369015627661,0.0006428760803300136,0.07899573897879536,0.013855768051405546,12.530000000000001,716944.7026687836,1217.8815104166667,23.45751953125 +log10_pgen,10k,sceptr_mlp,3,10000,0.8214505411013961,0.0017721037636965911,0.534502310892864,0.0033169693385023216,0.8077886542604048,0.0008293970081266804,0.8998059552994505,0.0003966274326964802,0.9102131961914464,0.000864334949118416,0.07262300909338824,0.005833706923441675,11.246666666666668,1021530.1494876714,1182.546875,21.20751953125 +log10_pgen,all,irrm_pretrained,10,79730,0.31017690475410065,0.03876912606322715,0.1925337933537014,0.0421692419704994,0.9722094243698776,0.0071918253654449016,0.9911934932097092,0.002627633642411745,0.988261224579343,0.005649934632632095,0.11768424049840458,0.04843363342634476,158.781,77118.10143799629,1593.01328125,254.86376953125 +log10_pgen,all,irrm_scratch,10,79730,0.3114129649713904,0.03330243911503736,0.2026686762306246,0.025686145645484902,0.972091574015385,0.006053100947354877,0.9921383974868743,0.0030753230018360387,0.9927640057242495,0.0024501388116768847,0.09871997901982245,0.046786331452667265,131.178,77301.63643012849,1579.693359375,254.86376953125 +log10_pgen,all,tcr_bert_mlp,3,79730,0.37576665017247074,0.0018731171749935922,0.23334980691737092,0.001147842183293328,0.959778478734926,0.0004011899273065553,0.9800374676077827,0.0002548146313769847,0.9851189500136984,0.0004239194870380052,0.023774566055857788,0.0012340738433489955,74.11,482204.1447037112,1666.4921875,28.14501953125 +log10_pgen,all,esm2_8m_mlp,3,79730,0.5727862854746651,0.0004557230451824632,0.34935756712203453,0.0004224821792519374,0.906545495984861,0.00014868716086941896,0.9534896057040494,0.00024764840122839165,0.9672519631455843,0.0004119721497089627,0.0745543660728548,0.007573095317883042,62.94333333333333,739699.4628127055,1296.2174479166667,23.45751953125 +log10_pgen,all,sceptr_mlp,3,79730,0.5795004266216509,0.004624225940520761,0.3442356353191398,0.0012591644567361206,0.9043377068184898,0.0015301220456800072,0.9532377971839475,0.000666761216473622,0.9668332557858953,0.00042302343036070976,0.09755030382889791,0.009563877131517034,60.583333333333336,1029364.2784961942,1193.4205729166667,21.20751953125 +log10_pgen,all,tfidf_ridge,1,79730,0.6210225099572112,,0.41454671189516185,,0.8901425494415403,,0.9434950444517985,,0.9563135820470278,,0.008684176700797186,,8.4,68718.25314155618,714.09765625,0.0 +log10_pgen_1mm,1k,irrm_pretrained,10,1000,0.6262719413828294,0.01830265138298122,0.3955077027607018,0.00961961574615435,0.8490809012780591,0.00880504553440322,0.9393137563872822,0.006548271545908237,0.9439822027783091,0.00589778494371969,0.24200087457294406,0.04804549899851409,6.465000000000001,79076.58743231386,1553.939453125,255.98876953125 +log10_pgen_1mm,1k,tcr_bert_mlp,3,1000,0.685144286570226,0.002153496133772639,0.43468398635694316,0.0008122116639941274,0.8195106709311563,0.0011344461621800414,0.9089732477151248,0.0005783007075420176,0.9260142078177989,0.0004060297908214042,0.07944990736101322,0.007503188590872171,4.973333333333334,474278.88027787465,1241.6861979166667,28.14501953125 +log10_pgen_1mm,1k,irrm_scratch,10,1000,0.6946828977880078,0.09169235422773352,0.4396730913095346,0.07394649046737808,0.8115420155347044,0.05006709040749404,0.9209237429346361,0.009130494064272706,0.92715291142624,0.01084600102257937,0.229206009669244,0.16953090969254042,7.931,78145.45197570449,1559.0765625,255.98876953125 +log10_pgen_1mm,1k,esm2_8m_mlp,3,1000,0.8615945278815476,0.0024225274486241732,0.5685573828281646,0.0008542452596914314,0.7145745186220601,0.0016045201855514647,0.8475403525356285,0.0005475393573428439,0.8674118138794791,0.0010102596207092706,0.09769012470210969,0.013496899355350991,5.140000000000001,735856.4119480439,1197.7799479166667,23.45751953125 +log10_pgen_1mm,1k,sceptr_mlp,3,1000,0.873569962655704,0.001278608640091633,0.5872807465781368,0.0017447483024049287,0.7065861603477129,0.000859261187461126,0.8419430222976985,0.0007044658486279333,0.8544503515961605,0.001305025388521403,0.07670004736733034,0.005688838971975157,4.346666666666667,1028820.2455296965,1171.2317708333333,21.20751953125 +log10_pgen_1mm,1k,tfidf_ridge,1,1000,1.085296409794027,,0.73535530121258,,0.5471218978219647,,0.750774560158103,,0.795403603438914,,0.18915308819686613,,0.19,80950.28937014837,643.671875,0.0 +log10_pgen_1mm,10k,irrm_scratch,10,10000,0.37544649188982115,0.05393243068687214,0.23577393103520206,0.04089320136394839,0.9447957898778669,0.01791525764663305,0.9754370036887015,0.0038427589703087235,0.9807582721749771,0.00451351584966454,0.08146111243136314,0.08586850228496051,26.005000000000003,78150.25535034612,1581.172265625,255.36376953125 +log10_pgen_1mm,10k,irrm_pretrained,10,10000,0.38134067696205276,0.027442508022011,0.2542467817755318,0.02295537004741706,0.9438266494916274,0.0076933826463340274,0.9815385301853583,0.0027043074180465344,0.9785087302017675,0.004963049737352929,0.18800522715017878,0.03347324976757551,24.767,78459.65347761178,1587.60703125,255.36376953125 +log10_pgen_1mm,10k,tcr_bert_mlp,3,10000,0.41190235103679723,0.0013890241512824646,0.25669393850789196,0.0017155417017192197,0.9347656373928152,0.00044039284910229917,0.9672464846909655,0.00032551430104860114,0.9758391756576296,0.00020704041320170286,0.01661317124627551,0.00894797698767368,13.696666666666667,478703.45815385296,1281.9127604166667,28.14501953125 +log10_pgen_1mm,10k,esm2_8m_mlp,3,10000,0.6247899714509164,0.0020961953451629095,0.39108601612116534,0.002604252012107772,0.8499085517012229,0.0010061435225250147,0.9228086720675472,0.0009210947503702813,0.9425420893779722,0.0008079743916434745,0.05559858269994359,0.008431332591102741,12.863333333333335,735642.5400245866,1221.5703125,23.45751953125 +log10_pgen_1mm,10k,sceptr_mlp,3,10000,0.6385171609346493,0.0019604975562168936,0.39748508241258484,0.0012657825132708074,0.8432410064002456,0.0009634019559633482,0.9188284272422029,0.0004262476087169893,0.9392381526783972,0.00019724605799668942,0.044973629620426524,0.0037911682871111067,11.646666666666667,1039275.5822895527,1185.7122395833333,21.20751953125 +log10_pgen_1mm,10k,tfidf_ridge,1,10000,0.679848528075636,,0.44560910401677106,,0.8222912089245368,,0.9069106930892485,,0.9418412049600424,,0.018927686789477982,,0.55,79617.02327020983,651.0703125,0.0 +log10_pgen_1mm,all,irrm_pretrained,10,79730,0.21811147813829623,0.032210945564203126,0.12475846464670519,0.034986743339859155,0.9813497911570795,0.006241034205243758,0.9934877653197903,0.002576512029083435,0.992083726131284,0.005850870340073854,0.033805298153456245,0.049056661532578,186.68599999999998,77093.5363036828,1584.332421875,254.86376953125 +log10_pgen_1mm,all,irrm_scratch,10,79730,0.24346109450194323,0.024810691892533376,0.1582809183121708,0.018059866352145078,0.9769970188800086,0.004526437952595065,0.9916151246332007,0.002322840744846866,0.9936817248904892,0.001651159132919006,0.04659588096187876,0.05440170415855433,104.894,77290.11482224047,1583.28125,254.86376953125 +log10_pgen_1mm,all,tcr_bert_mlp,3,79730,0.2765836233126562,0.0006421618938097269,0.16712247914391684,0.000271781866644313,0.9705870169039016,0.00013654689125638776,0.9854416261999398,1.5971834883716513e-05,0.9907688495886803,8.455516409449843e-05,0.017521640543719588,0.001736951945831234,71.11,479454.95510929095,1668.41796875,28.14501953125 +log10_pgen_1mm,all,sceptr_mlp,3,79730,0.44646978094298007,0.005952127061135091,0.2509612066562197,0.00223977520363394,0.9233485838892074,0.002050787523950637,0.9621486970353169,0.0009321620614320843,0.9783254925691707,0.0003970622889583402,0.06420806648585768,0.00999523478389865,58.419999999999995,1043629.0571447442,1193.36328125,21.20751953125 +log10_pgen_1mm,all,esm2_8m_mlp,3,79730,0.45393456930558324,0.00153386906096142,0.2654829299913543,0.00023225770136071967,0.9207727812660104,0.0005349109434407152,0.9607129637035211,0.00014495604518536024,0.976231853929708,0.00016747358961218396,0.05494076415998711,0.0040845593124909495,67.82333333333334,726935.7194829332,1295.1966145833333,23.45751953125 +log10_pgen_1mm,all,tfidf_ridge,1,79730,0.5542270263983469,,0.36706299488768934,,0.8818972557236601,,0.9391297435576847,,0.9565962719889497,,0.0035393838152091624,,8.65,67434.86203217156,717.2109375,0.0 diff --git a/results/pgen/pgen_summary.md b/results/pgen/pgen_summary.md new file mode 100644 index 0000000..a0f87e6 --- /dev/null +++ b/results/pgen/pgen_summary.md @@ -0,0 +1,69 @@ +### log10_pgen + +**train = 1k (1,000 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_pretrained | 0.8154 ± 0.0237 | 0.5364 | 0.8105 | 0.9230 | 0.9153 | +0.3245 | 6 | 79911 | 1554 | +| irrm_scratch | 0.8606 ± 0.0791 | 0.5499 | 0.7875 | 0.9009 | 0.9000 | +0.2298 | 9 | 79744 | 1554 | +| tcr_bert_mlp | 0.8642 ± 0.0046 | 0.5753 | 0.7873 | 0.8910 | 0.8942 | +0.0958 | 5 | 481115 | 1245 | +| esm2_8m_mlp | 1.0740 ± 0.0039 | 0.7270 | 0.6715 | 0.8223 | 0.8332 | +0.1261 | 5 | 729357 | 1196 | +| sceptr_mlp | 1.0974 ± 0.0030 | 0.7608 | 0.6570 | 0.8125 | 0.8069 | +0.1051 | 4 | 1065072 | 1172 | +| tfidf_ridge | 1.2770 | 0.8800 | 0.5355 | 0.7459 | 0.7756 | +0.2475 | 0 | 82047 | 648 | + +**train = 10k (10,000 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_scratch | 0.4847 ± 0.0171 | 0.3079 | 0.9330 | 0.9700 | 0.9742 | +0.1199 | 29 | 78849 | 1590 | +| irrm_pretrained | 0.5150 ± 0.0250 | 0.3604 | 0.9243 | 0.9755 | 0.9697 | +0.2757 | 24 | 78524 | 1592 | +| tcr_bert_mlp | 0.5516 ± 0.0022 | 0.3603 | 0.9133 | 0.9563 | 0.9606 | +0.0295 | 13 | 480017 | 1277 | +| tfidf_ridge | 0.7659 | 0.5085 | 0.8329 | 0.9129 | 0.9404 | +0.0342 | 1 | 79361 | 653 | +| esm2_8m_mlp | 0.7914 ± 0.0023 | 0.5108 | 0.8216 | 0.9076 | 0.9226 | +0.0790 | 13 | 716945 | 1218 | +| sceptr_mlp | 0.8215 ± 0.0018 | 0.5345 | 0.8078 | 0.8998 | 0.9102 | +0.0726 | 11 | 1021530 | 1183 | + +**train = all (79,730 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_pretrained | 0.3102 ± 0.0388 | 0.1925 | 0.9722 | 0.9912 | 0.9883 | +0.1177 | 159 | 77118 | 1593 | +| irrm_scratch | 0.3114 ± 0.0333 | 0.2027 | 0.9721 | 0.9921 | 0.9928 | +0.0987 | 131 | 77302 | 1580 | +| tcr_bert_mlp | 0.3758 ± 0.0019 | 0.2333 | 0.9598 | 0.9800 | 0.9851 | +0.0238 | 74 | 482204 | 1666 | +| esm2_8m_mlp | 0.5728 ± 0.0005 | 0.3494 | 0.9065 | 0.9535 | 0.9673 | +0.0746 | 63 | 739699 | 1296 | +| sceptr_mlp | 0.5795 ± 0.0046 | 0.3442 | 0.9043 | 0.9532 | 0.9668 | +0.0976 | 61 | 1029364 | 1193 | +| tfidf_ridge | 0.6210 | 0.4145 | 0.8901 | 0.9435 | 0.9563 | +0.0087 | 8 | 68718 | 714 | + +### log10_pgen_1mm + +**train = 1k (1,000 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_pretrained | 0.6263 ± 0.0183 | 0.3955 | 0.8491 | 0.9393 | 0.9440 | +0.2420 | 6 | 79077 | 1554 | +| tcr_bert_mlp | 0.6851 ± 0.0022 | 0.4347 | 0.8195 | 0.9090 | 0.9260 | +0.0794 | 5 | 474279 | 1242 | +| irrm_scratch | 0.6947 ± 0.0917 | 0.4397 | 0.8115 | 0.9209 | 0.9272 | +0.2292 | 8 | 78145 | 1559 | +| esm2_8m_mlp | 0.8616 ± 0.0024 | 0.5686 | 0.7146 | 0.8475 | 0.8674 | +0.0977 | 5 | 735856 | 1198 | +| sceptr_mlp | 0.8736 ± 0.0013 | 0.5873 | 0.7066 | 0.8419 | 0.8545 | +0.0767 | 4 | 1028820 | 1171 | +| tfidf_ridge | 1.0853 | 0.7354 | 0.5471 | 0.7508 | 0.7954 | +0.1892 | 0 | 80950 | 644 | + +**train = 10k (10,000 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_scratch | 0.3754 ± 0.0539 | 0.2358 | 0.9448 | 0.9754 | 0.9808 | +0.0815 | 26 | 78150 | 1581 | +| irrm_pretrained | 0.3813 ± 0.0274 | 0.2542 | 0.9438 | 0.9815 | 0.9785 | +0.1880 | 25 | 78460 | 1588 | +| tcr_bert_mlp | 0.4119 ± 0.0014 | 0.2567 | 0.9348 | 0.9672 | 0.9758 | +0.0166 | 14 | 478703 | 1282 | +| esm2_8m_mlp | 0.6248 ± 0.0021 | 0.3911 | 0.8499 | 0.9228 | 0.9425 | +0.0556 | 13 | 735643 | 1222 | +| sceptr_mlp | 0.6385 ± 0.0020 | 0.3975 | 0.8432 | 0.9188 | 0.9392 | +0.0450 | 12 | 1039276 | 1186 | +| tfidf_ridge | 0.6798 | 0.4456 | 0.8223 | 0.9069 | 0.9418 | +0.0189 | 1 | 79617 | 651 | + +**train = all (79,730 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_pretrained | 0.2181 ± 0.0322 | 0.1248 | 0.9813 | 0.9935 | 0.9921 | +0.0338 | 187 | 77094 | 1584 | +| irrm_scratch | 0.2435 ± 0.0248 | 0.1583 | 0.9770 | 0.9916 | 0.9937 | +0.0466 | 105 | 77290 | 1583 | +| tcr_bert_mlp | 0.2766 ± 0.0006 | 0.1671 | 0.9706 | 0.9854 | 0.9908 | +0.0175 | 71 | 479455 | 1668 | +| sceptr_mlp | 0.4465 ± 0.0060 | 0.2510 | 0.9233 | 0.9621 | 0.9783 | +0.0642 | 58 | 1043629 | 1193 | +| esm2_8m_mlp | 0.4539 ± 0.0015 | 0.2655 | 0.9208 | 0.9607 | 0.9762 | +0.0549 | 68 | 726936 | 1295 | +| tfidf_ridge | 0.5542 | 0.3671 | 0.8819 | 0.9391 | 0.9566 | +0.0035 | 9 | 67435 | 717 | diff --git a/results/pgen/pretraining_delta.csv b/results/pgen/pretraining_delta.csv new file mode 100644 index 0000000..9ca3d57 --- /dev/null +++ b/results/pgen/pretraining_delta.csv @@ -0,0 +1,7 @@ +target,train_subset,n_scratch,n_pretrained,irrm_scratch,irrm_pretrained,rmse_delta,relative_improvement,p_value,verdict +log10_pgen,1k,10,10,0.8605530608233808,0.81540254540097,0.045150515422410775,0.05246685820768631,0.11262996695012906,no detectable effect +log10_pgen,10k,10,10,0.48468864648435217,0.5149901453584536,-0.030301498874101385,-0.06251745134508664,0.006104018174065397,pretraining hurts +log10_pgen,all,10,10,0.31141296497139037,0.31017690475410065,0.001236060217289725,0.003969199604143913,0.9399003420219125,no detectable effect +log10_pgen_1mm,1k,10,10,0.6946828977880077,0.6262719413828294,0.06841095640517836,0.09847796256826079,0.04394678471543129,pretraining helps +log10_pgen_1mm,10k,10,10,0.37544649188982115,0.38134067696205276,-0.005894185072231606,-0.01569913476235468,0.7628154827592442,no detectable effect +log10_pgen_1mm,all,10,10,0.24346109450194325,0.21811147813829623,0.02534961636364702,0.10412183686065161,0.06525425563284348,no detectable effect diff --git a/results/pgen/speed_end_to_end_cpu.csv b/results/pgen/speed_end_to_end_cpu.csv new file mode 100644 index 0000000..cc1d905 --- /dev/null +++ b/results/pgen/speed_end_to_end_cpu.csv @@ -0,0 +1,6 @@ +arm,device,sequences_per_second,ms_per_sequence,seconds_total +tfidf_ridge,cpu,82842.20607184026,0.012071141600608825,0.06035570800304413 +irrm,cpu,3776.528412201475,0.26479345336556437,1.3239672668278217 +sceptr_mlp,cpu,1264.9458164759842,0.7905476953834295,3.9527384769171476 +esm2_8m_mlp,cpu,812.9761726854197,1.2300483502447606,6.150241751223803 +tcr_bert_mlp,cpu,215.6417930730961,4.637320000678301,23.186600003391504 diff --git a/results/pgen/speed_end_to_end_cuda.csv b/results/pgen/speed_end_to_end_cuda.csv new file mode 100644 index 0000000..5a644ee --- /dev/null +++ b/results/pgen/speed_end_to_end_cuda.csv @@ -0,0 +1,6 @@ +arm,device,sequences_per_second,ms_per_sequence,seconds_total +tfidf_ridge,cuda,86660.47446799547,0.011539286002516747,0.05769643001258373 +irrm,cuda,64813.57391703495,0.01542886681854725,0.07714433409273624 +esm2_8m_mlp,cuda,9392.62324529416,0.10646652951836587,0.5323326475918293 +tcr_bert_mlp,cuda,4038.584267802087,0.24761152267456055,1.2380576133728027 +sceptr_mlp,cuda,2495.308062377455,0.4007521215826273,2.0037606079131365 diff --git a/results/pgen_constlr/pgen_runs.csv b/results/pgen_constlr/pgen_runs.csv new file mode 100644 index 0000000..5eaca29 --- /dev/null +++ b/results/pgen_constlr/pgen_runs.csv @@ -0,0 +1,181 @@ +arm,target,train_subset,train_size,seed,training_seconds,rmse,mae,r2,pearson_r,spearman_rho,bias,inference_seq_per_second,inference_ms_per_sequence,device,peak_rss_mb,peak_gpu_mb,epochs_ran +esm2_8m_mlp,log10_pgen,1k,1000,42,3.94,1.0915541051919524,0.7402943732362133,0.660605057411438,0.8169628083446634,0.8234968097446708,0.15321451674286654,516076.35393841757,0.0019376977696585731,cuda,1199.19140625,23.45751953125,36.0 +esm2_8m_mlp,log10_pgen,1k,1000,43,3.62,1.0795015732119884,0.7359788250816464,0.6680586231435096,0.8196909145226835,0.825906573915183,0.11602056096385072,521901.14098656137,0.0019160716876565505,cuda,1202.671875,23.45751953125,33.0 +esm2_8m_mlp,log10_pgen,1k,1000,44,4.23,1.0792830312542607,0.7377109537684872,0.6681930106739722,0.8189650446270591,0.8286560440522147,0.09293785484005393,533565.0251171781,0.0018741858122735583,cuda,1198.1796875,23.45751953125,39.0 +irrm_pretrained,log10_pgen,1k,1000,42,9.58,0.8726550760693804,0.5738058147665613,0.783079762948996,0.9159395638398967,0.904600743389775,0.3580331865933481,51294.28117321873,0.01949535069266377,cuda,1545.14453125,255.98876953125,24.0 +irrm_pretrained,log10_pgen,1k,1000,43,9.41,0.8380269458087515,0.5430632993028036,0.799953573378763,0.9165763115138791,0.9104608804983463,0.29427279298237435,51134.894099149045,0.01955611755176474,cuda,1530.4296875,255.98876953125,24.0 +irrm_pretrained,log10_pgen,1k,1000,44,10.64,0.8284073138270375,0.5369289222086728,0.8045198421966144,0.9202394479936183,0.9145521115122686,0.31079326245117705,50260.79320984828,0.019896223997596132,cuda,1533.0625,255.98876953125,27.0 +irrm_pretrained,log10_pgen,1k,1000,45,6.98,0.8730714403127777,0.5729750874305046,0.7828727180827383,0.9150187388529426,0.9086835384224735,0.3783469468940029,80348.14928043708,0.012445837383381736,cuda,1551.4609375,255.98876953125,28.0 +irrm_pretrained,log10_pgen,1k,1000,46,10.14,0.8955693146946709,0.5783578312292941,0.7715383856109055,0.9135258281103794,0.9083706530804575,0.3708126548954926,50534.95524079303,0.019788283085146102,cuda,1527.92578125,255.98876953125,26.0 +irrm_pretrained,log10_pgen,1k,1000,47,5.36,0.8396945597007993,0.5545821248403418,0.7991566249701977,0.9113736381289581,0.9056997342503031,0.28816276817374503,80618.55931828934,0.012404091668915962,cuda,1551.01953125,255.98876953125,21.0 +irrm_pretrained,log10_pgen,1k,1000,48,7.28,0.8150718432048386,0.5475463423471961,0.8107627565877321,0.9143665674978133,0.905040473254371,0.1834927141464866,50143.35650795583,0.019942821335491138,cuda,1533.8359375,255.98876953125,18.0 +irrm_pretrained,log10_pgen,1k,1000,49,8.05,0.8704777136313411,0.5613098881755785,0.7841608882495538,0.9197475352791084,0.9121646002518446,0.3501417096914958,80054.3104401068,0.012491519750809137,cuda,1556.7109375,255.98876953125,32.0 +irrm_pretrained,log10_pgen,1k,1000,50,9.15,0.8728682173569736,0.5556686607249849,0.7829737868166594,0.9069782520813081,0.9077522395634259,0.26813947600043,49952.29768064458,0.020019099149216475,cuda,1524.78125,255.98876953125,23.0 +irrm_pretrained,log10_pgen,1k,1000,51,5.82,0.8444090542591324,0.550895488311041,0.7968950094359974,0.9150288894370522,0.9074025004235194,0.29766531813792796,80316.62807986789,0.01245072189790621,cuda,1566.3046875,255.98876953125,23.0 +irrm_scratch,log10_pgen,1k,1000,42,13.39,0.8273421285536087,0.5551784896127165,0.8050222248124117,0.909914579519032,0.9121780503326388,0.008624446700217612,50720.321822123,0.019715963228841813,cuda,1529.171875,255.98876953125,29.0 +irrm_scratch,log10_pgen,1k,1000,43,16.43,0.7659800942598978,0.4832207026282427,0.8328717756010071,0.9169360837216642,0.9190804948555896,0.13152816635376696,51292.88083902965,0.019495882930386756,cuda,1526.18359375,255.98876953125,37.0 +irrm_scratch,log10_pgen,1k,1000,44,14.37,0.7734791074501303,0.48697228424591754,0.8295833567191366,0.9177807024004445,0.9177315561830938,0.134461495555832,50365.65252164455,0.019854800840121183,cuda,1539.015625,255.98876953125,36.0 +irrm_scratch,log10_pgen,1k,1000,45,6.0,0.8539224789496141,0.5486367348054839,0.7922927172235745,0.8934616865743961,0.8922835537384012,0.14120462825912458,81472.33337984262,0.012274105312999588,cuda,1563.94921875,255.98876953125,24.0 +irrm_scratch,log10_pgen,1k,1000,46,10.49,0.821753873046916,0.5321979550634945,0.8076472714601192,0.9006659274404937,0.8979426733955749,0.10031384699661687,50611.856025323556,0.019758216325827915,cuda,1526.07421875,255.98876953125,27.0 +irrm_scratch,log10_pgen,1k,1000,47,7.49,0.8155386677194593,0.5207589655369478,0.8105459269103279,0.9062289998614832,0.9034322766161501,0.1744060220632375,81164.71916927739,0.012320624160780953,cuda,1555.8046875,255.98876953125,30.0 +irrm_scratch,log10_pgen,1k,1000,48,10.13,1.0387836221780697,0.6855247208672894,0.6926274929312484,0.894016387851271,0.8914120357400651,0.5572090034998788,50667.35209156709,0.019736575106446837,cuda,1527.1015625,255.98876953125,26.0 +irrm_scratch,log10_pgen,1k,1000,49,5.84,0.8624344844562729,0.5699924510092282,0.7881311746289452,0.8912662667482427,0.8891159235445192,0.07916370451294726,80878.741597902,0.01236418841642734,cuda,1541.0234375,255.98876953125,23.0 +irrm_scratch,log10_pgen,1k,1000,50,11.18,0.8471909549401387,0.5283402923875888,0.7955545485946303,0.9049882896587529,0.9049888220927975,0.2020672699252005,50459.30508681347,0.019817950292409597,cuda,1538.03125,255.98876953125,29.0 +irrm_scratch,log10_pgen,1k,1000,51,7.99,0.7730128518213162,0.5031581391390016,0.8297887501673299,0.9138801307741508,0.9151692724245305,0.02280371236558214,80576.88220665505,0.012410507488181363,cuda,1577.8984375,255.98876953125,32.0 +sceptr_mlp,log10_pgen,1k,1000,42,4.47,1.1015078458851626,0.7649751967259993,0.6543870382698407,0.8100436681628357,0.804978245151527,0.07793970883382847,778354.4681485145,0.0012847616875364225,cuda,1170.66796875,21.20751953125,43.0 +sceptr_mlp,log10_pgen,1k,1000,43,3.26,1.1016636574753431,0.7619187374556721,0.6542892553733989,0.811423946447547,0.8051246776840334,0.12022796173619192,882398.3426362096,0.0011332750206810786,cuda,1166.8515625,21.20751953125,32.0 +sceptr_mlp,log10_pgen,1k,1000,44,3.67,1.0896400265903015,0.761396253152809,0.6617942958022931,0.8142103565529715,0.8080201289452699,0.04627952788417081,908732.0562786325,0.0011004343833705182,cuda,1171.33203125,21.20751953125,35.0 +tcr_bert_mlp,log10_pgen,1k,1000,42,5.05,0.8760345003106056,0.5819093645526308,0.7813964290543859,0.8881443789227367,0.8913580831616918,0.12819454572911565,318430.06197725417,0.0031404070136802323,cuda,1246.765625,28.14501953125,40.0 +tcr_bert_mlp,log10_pgen,1k,1000,43,4.17,0.8632398237103549,0.5830586150560361,0.7877353046815565,0.8891474748834872,0.8893524914152939,0.055539820699031614,309981.80367169366,0.0032259958105770458,cuda,1244.59375,28.14501953125,35.0 +tcr_bert_mlp,log10_pgen,1k,1000,44,6.12,0.8580475606372164,0.579303658373242,0.7902811083390895,0.8897587965843224,0.8932709859909886,0.02707135115754949,291458.1496005449,0.003431024321572549,cuda,1249.66015625,28.14501953125,52.0 +tfidf_ridge,log10_pgen,1k,1000,42,0.44,1.276955755696797,0.8800488403651113,0.5355205204794624,0.7459285398734663,0.7755971181289645,0.24748485510618118,46905.79064077635,0.021319329369339647,cpu,655.55078125,0.0, +esm2_8m_mlp,log10_pgen,10k,10000,42,14.47,0.7966986650954415,0.5174849735727736,0.8191980913374267,0.9071021327654052,0.9216778640250258,0.11214545445353795,520928.62497661123,0.001919648781145206,cuda,1217.65234375,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen,10k,10000,43,12.23,0.8015686329227331,0.5239774840812171,0.8169809655284199,0.904688112817364,0.9187355591684678,0.06898057147978498,543791.4727871665,0.0018389402005047405,cuda,1217.609375,23.45751953125,52.0 +esm2_8m_mlp,log10_pgen,10k,10000,44,11.7,0.8130470777845348,0.5269840306024154,0.8117017783783829,0.9031960836644991,0.9162947786091444,0.11342834681753776,519796.4815531242,0.001923829874746465,cuda,1216.08203125,23.45751953125,48.0 +irrm_pretrained,log10_pgen,10k,10000,42,38.49,0.5019841051682237,0.35627390506351453,0.9282214363308795,0.975111086691799,0.9682088622524755,0.2690681770530914,49387.35804129182,0.020248096672106232,cuda,1536.140625,255.36376953125,28.0 +irrm_pretrained,log10_pgen,10k,10000,43,29.14,0.6072004539701912,0.427645045166025,0.8949783089310082,0.9724442690808554,0.9664332662967745,0.36938994676744685,49579.00388832995,0.020169828386475165,cuda,1534.71484375,255.36376953125,21.0 +irrm_pretrained,log10_pgen,10k,10000,44,54.5,0.4370812238168363,0.29203450514300294,0.9455824314863231,0.9769737912781561,0.9738792892972243,0.16793524652596767,48117.34292600225,0.020782527446244488,cuda,1558.37890625,255.36376953125,39.0 +irrm_pretrained,log10_pgen,10k,10000,45,19.05,0.5500174080896365,0.36935044114925275,0.91382769653112,0.9726686230820053,0.9645339276878059,0.27234471256262516,78345.31960212644,0.0127640043474002,cuda,1600.17578125,255.36376953125,26.0 +irrm_pretrained,log10_pgen,10k,10000,46,13.76,0.5911542166154463,0.406977227846606,0.900455695805151,0.9714829073068255,0.9660735651292331,0.3354204946989577,78635.7118076855,0.012716868417820625,cuda,1586.33203125,255.36376953125,18.0 +irrm_pretrained,log10_pgen,10k,10000,47,15.36,0.4888588784078921,0.3293540748698081,0.9319259097726058,0.9732700643861482,0.9664655902116925,0.21819266897231107,78277.66738632193,0.012775035759110239,cuda,1579.4765625,255.36376953125,21.0 +irrm_pretrained,log10_pgen,10k,10000,48,16.01,0.5354970143071539,0.3902579098488256,0.9183175154346566,0.974014324072708,0.9671961798973695,0.31852752852926525,78617.67365585294,0.01271978619435468,cuda,1596.01171875,255.36376953125,21.0 +irrm_pretrained,log10_pgen,10k,10000,49,18.22,0.5581728783451227,0.40858714159675347,0.9112532837722653,0.9728205507714903,0.960817491415762,0.3282748804806309,77998.74070560814,0.012820719808468646,cuda,1599.359375,255.36376953125,25.0 +irrm_pretrained,log10_pgen,10k,10000,50,19.25,0.5270736550875357,0.38627417251271373,0.9208670328180113,0.9744715976870983,0.9647418996062248,0.3086164301141052,78419.26439319062,0.01275196863599798,cuda,1584.53515625,255.36376953125,26.0 +irrm_pretrained,log10_pgen,10k,10000,51,12.61,0.6251439984244976,0.45201107165020654,0.8886795471851842,0.9676513537306013,0.9604122315429587,0.3910038787630061,78106.44570462324,0.012803040657895516,cuda,1590.1171875,255.36376953125,17.0 +irrm_scratch,log10_pgen,10k,10000,42,19.37,0.5751521553443715,0.36240860395121677,0.9057719214879167,0.9520192928529396,0.9524838673585823,0.04430426241683526,50632.00661867529,0.019750352924609478,cuda,1542.84375,255.36376953125,14.0 +irrm_scratch,log10_pgen,10k,10000,43,31.69,0.5007316165717091,0.3300130508894094,0.928579175452831,0.9642083396331315,0.9665581193770536,-0.06257023294812729,49718.20607856898,0.020113356431640233,cuda,1534.55078125,255.36376953125,23.0 +irrm_scratch,log10_pgen,10k,10000,44,61.01,0.4469771841868142,0.27836988726873746,0.9430903987777541,0.9741013193473115,0.9754319376831495,0.13597858457402748,49778.46625232864,0.02008900786398215,cuda,1544.23046875,255.36376953125,45.0 +irrm_scratch,log10_pgen,10k,10000,45,38.13,0.4631258469189825,0.3008167975501209,0.938903986870447,0.9704180906902137,0.9711355946769996,0.08284721902999719,49424.02072992871,0.020233076654454583,cuda,1549.046875,255.36376953125,28.0 +irrm_scratch,log10_pgen,10k,10000,46,15.11,0.513747792852385,0.3283076451165438,0.9248178449988129,0.9619094348416142,0.9659900784041277,0.03270629272861952,79673.31215256036,0.012551254278034482,cuda,1578.6796875,255.36376953125,20.0 +irrm_scratch,log10_pgen,10k,10000,47,23.61,0.45577350966876207,0.29101608587140165,0.9408284438149535,0.9699673282824827,0.9704972919439828,0.005078207319413182,79138.31108404836,0.012636104894100608,cuda,1592.99609375,255.36376953125,32.0 +irrm_scratch,log10_pgen,10k,10000,48,41.17,0.4734539790509061,0.3158020615534417,0.9361486074139419,0.9680498037965779,0.9701016241272683,-0.05792117853843921,49208.22809005352,0.02032180468213466,cuda,1542.109375,255.36376953125,30.0 +irrm_scratch,log10_pgen,10k,10000,49,14.14,0.6207776331020315,0.42591710844635566,0.890229168663141,0.9608567893677247,0.9620175525332944,0.338611312952432,78777.80804498789,0.012693930242752208,cuda,1598.8828125,255.36376953125,19.0 +irrm_scratch,log10_pgen,10k,10000,50,20.27,0.5104941019143432,0.3157260706700292,0.92576712354424,0.966965276583171,0.9666370862800718,0.10494240099662171,78415.86048844548,0.012752522178180384,cuda,1574.81640625,255.36376953125,28.0 +irrm_scratch,log10_pgen,10k,10000,51,18.5,0.4912820150441888,0.3233667366122514,0.9312493888304615,0.9660599780498389,0.9695236673192581,-0.05855530574608454,78525.09469403228,0.012734782478091016,cuda,1578.3046875,255.36376953125,25.0 +sceptr_mlp,log10_pgen,10k,10000,42,12.43,0.8241140085650038,0.5328266240608573,0.8065407833721275,0.8997464229402171,0.9116146522108373,0.10096275534285942,933740.0276133372,0.0010709619063413453,cuda,1180.078125,21.20751953125,56.0 +sceptr_mlp,log10_pgen,10k,10000,43,8.72,0.8398266277213052,0.5446345861976845,0.7990934422983463,0.8950853150574498,0.9065704590796785,0.08551405539283576,928490.7344967334,0.0010770166710839895,cuda,1183.64453125,21.20751953125,39.0 +sceptr_mlp,log10_pgen,10k,10000,44,9.32,0.8345979343986126,0.543566797366414,0.8015873110179307,0.8967334069477051,0.9075547226147163,0.09449457219799369,897752.4414176691,0.0011138928215231234,cuda,1182.296875,21.20751953125,42.0 +tcr_bert_mlp,log10_pgen,10k,10000,42,14.7,0.5716066534740867,0.37969244209306996,0.906930071018644,0.952757972654319,0.9536264682586161,0.048066586677211275,317680.7349723767,0.003147814424714653,cuda,1273.0234375,28.14501953125,58.0 +tcr_bert_mlp,log10_pgen,10k,10000,43,15.03,0.5636854793599815,0.3738462210991635,0.9094916751782824,0.9538177223571953,0.9569626148875244,0.011505552497695236,310494.78990068275,0.0032206659580982593,cuda,1272.17578125,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen,10k,10000,44,11.24,0.5689464613822552,0.3818325411493399,0.9077943288924409,0.9531500150003065,0.9539839680301143,0.019939592151627557,319234.63163778995,0.003132492220125479,cuda,1271.4765625,28.14501953125,44.0 +tfidf_ridge,log10_pgen,10k,10000,42,0.8,0.76592457502429,0.5084579081617813,0.83289600206607,0.9129297738264299,0.9404228939259158,0.03423942511518709,59776.55467359597,0.016728966824207284,cpu,649.6875,0.0, +esm2_8m_mlp,log10_pgen,all,79730,42,74.59,0.5951112784062191,0.3778985089314805,0.8991185783382184,0.9484219706765102,0.961556478055205,0.008018448818370851,512009.60178724234,0.00195308837277535,cuda,1302.80859375,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen,all,79730,43,50.88,0.617624643001898,0.3918682055753462,0.8913414103811155,0.9452014036779254,0.9567233374476869,0.056252967770279834,537971.2513126577,0.0018588353886197177,cuda,1297.15625,23.45751953125,41.0 +esm2_8m_mlp,log10_pgen,all,79730,44,72.62,0.5966358832752859,0.37357570394098827,0.8986010236337659,0.9487803111146647,0.9612983902630552,0.06572506115436895,530540.8837711269,0.0018848688773840019,cuda,1300.69140625,23.45751953125,59.0 +irrm_pretrained,log10_pgen,all,79730,42,139.98,0.3713165663531919,0.2728079802272566,0.9607261494233188,0.988563084501092,0.9869164526749541,0.23043004494876027,44512.322638950456,0.022465689065727414,cuda,1571.984375,254.86376953125,14.0 +irrm_pretrained,log10_pgen,all,79730,43,347.53,0.2732877421635648,0.16781113564368397,0.9787257069575539,0.991832386281639,0.9888689143888211,0.11093694782506916,47394.69275674851,0.02109940885433022,cuda,1590.50390625,254.86376953125,35.0 +irrm_pretrained,log10_pgen,all,79730,44,205.57,0.3504203019450061,0.19649729263615834,0.9650221306775171,0.9851718707287127,0.978683330148102,0.1121013101544567,48625.990574486685,0.020565133752250688,cuda,1557.7890625,254.86376953125,22.0 +irrm_pretrained,log10_pgen,all,79730,45,442.88,0.22182566736977571,0.15680747819213534,0.9859835373059486,0.9942586166515042,0.9937141588261896,0.09475363882593475,47269.34832899524,0.02115535828926576,cuda,1563.43359375,254.86376953125,46.0 +irrm_pretrained,log10_pgen,all,79730,46,152.94,0.25087874409297584,0.17141298237291647,0.9820715573743454,0.9930255075484793,0.9897789108642574,0.11165818089101294,76249.21253309137,0.01311488954152819,cuda,1579.53125,254.86376953125,34.0 +irrm_pretrained,log10_pgen,all,79730,47,90.82,0.3081060827680913,0.18067044300141025,0.9729594589703192,0.98795043358461,0.982963120190587,0.08818773965778262,77044.89818934872,0.012979444758851634,cuda,1571.203125,254.86376953125,20.0 +irrm_pretrained,log10_pgen,all,79730,48,135.12,0.2986887215728427,0.18298794555902867,0.9745872022580503,0.9892383244365203,0.9876915622839695,0.10148357882720346,77142.23328055044,0.01296306779664008,cuda,1605.953125,254.86376953125,30.0 +irrm_pretrained,log10_pgen,all,79730,49,197.31,0.2363387040497491,0.14275755332583007,0.9840894742264049,0.9932642509249219,0.9919847270618742,0.077734183311473,76508.9197560553,0.01307037144411982,cuda,1597.3515625,254.86376953125,44.0 +irrm_pretrained,log10_pgen,all,79730,50,230.84,0.3051675125435428,0.204067754933738,0.9734727990480465,0.9907285747007517,0.9885527241317941,0.151946996887396,45097.88213776874,0.022173990276197833,cuda,1599.87109375,254.86376953125,24.0 +irrm_pretrained,log10_pgen,all,79730,51,103.4,0.29346090454675655,0.17002071923826675,0.9754689953051937,0.9899495870083614,0.9870642278998465,0.07495573131594488,76568.00698079141,0.013060285090754284,cuda,1580.66796875,254.86376953125,23.0 +irrm_scratch,log10_pgen,all,79730,42,180.5,0.279814286287663,0.1880315391386849,0.9776974459635999,0.9893302583812692,0.9909131558161663,0.014503944017790777,44026.03046998928,0.02271383518624643,cuda,1556.59765625,254.86376953125,19.0 +irrm_scratch,log10_pgen,all,79730,43,163.51,0.32091193596616924,0.20097612255800545,0.9706649675550438,0.9903812553054268,0.9903384271183792,0.10042179829816784,48413.8793434837,0.02065523386186973,cuda,1554.2265625,254.86376953125,18.0 +irrm_scratch,log10_pgen,all,79730,44,158.27,0.33497006170309135,0.22447315337168178,0.9680385246398328,0.9903139887992916,0.9907750466785762,0.15124634206940737,48154.72149260951,0.020766395672197457,cuda,1565.8125,254.86376953125,17.0 +irrm_scratch,log10_pgen,all,79730,45,69.33,0.38364183086153864,0.24361249078632374,0.9580756110250899,0.9868639987947267,0.9873375371802209,0.15954624188589653,78605.74365544904,0.012721716677387846,cuda,1566.375,254.86376953125,15.0 +irrm_scratch,log10_pgen,all,79730,46,256.28,0.2765557228636666,0.19643691692277543,0.9782138679476076,0.9945744989211008,0.9943793188551895,-0.03322250443971604,47319.816123156845,0.021132795558574266,cuda,1559.6171875,254.86376953125,28.0 +irrm_scratch,log10_pgen,all,79730,47,49.79,0.4767214417749651,0.3107282061051032,0.935264247058595,0.9741276964314304,0.9767865361194387,0.20746135859338838,78385.25126887314,0.01275750200212857,cuda,1592.8828125,254.86376953125,11.0 +irrm_scratch,log10_pgen,all,79730,48,138.08,0.28539163367483134,0.20092957071021314,0.976799502012103,0.9939909339821749,0.99446663064038,0.17717354120223236,77375.05021257492,0.012924062695309,cuda,1574.5078125,254.86376953125,31.0 +irrm_scratch,log10_pgen,all,79730,49,219.8,0.27950952345275093,0.17243624843163155,0.9777460016551317,0.9928358034944527,0.9934787961501889,0.10517735134727471,47216.97727532152,0.02117882290873925,cuda,1556.86328125,254.86376953125,23.0 +irrm_scratch,log10_pgen,all,79730,50,126.62,0.29019055088587026,0.2017640000268431,0.9760127000527028,0.9935500344917078,0.9938903933807892,-0.011590887368701517,76962.32866136862,0.012993369839417967,cuda,1586.01953125,254.86376953125,28.0 +irrm_scratch,log10_pgen,all,79730,51,94.65,0.4301566653976132,0.28585552544683723,0.9472930184547584,0.9920160751800176,0.9918333428144999,0.12769040293382197,76544.8788356785,0.013064231274658285,cuda,1580.16015625,254.86376953125,21.0 +sceptr_mlp,log10_pgen,all,79730,42,68.73,0.5795817078970179,0.35118457974334066,0.9043149314162111,0.9530617534275462,0.9646389496705083,0.08103153479383794,903096.3845214696,0.0011073015207893643,cuda,1205.3125,21.20751953125,60.0 +sceptr_mlp,log10_pgen,all,79730,43,63.41,0.5864017681293087,0.35538681299615843,0.9020497893897336,0.9512580809778278,0.9650385869823047,0.0895127140611525,932787.822787842,0.0010720551614956547,cuda,1197.5,21.20751953125,54.0 +sceptr_mlp,log10_pgen,all,79730,44,69.09,0.5839895730378692,0.35415758598602276,0.9028539788212233,0.9529103660424704,0.9650201847137837,0.09078357713498891,840967.2787517394,0.0011891069073273759,cuda,1197.9453125,21.20751953125,60.0 +tcr_bert_mlp,log10_pgen,all,79730,42,51.11,0.41770773759434954,0.2717698329057919,0.9502996013491293,0.9748411659041659,0.9781232625528591,0.005214108014531167,319254.11607227504,0.0031323010406343918,cuda,1670.78125,28.14501953125,39.0 +tcr_bert_mlp,log10_pgen,all,79730,43,79.57,0.4082671293067538,0.2585217788211574,0.9525207705490689,0.9769878226010591,0.980975805469901,0.061827468664735444,319375.7427781945,0.0031311081777882456,cuda,1669.08203125,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen,all,79730,44,62.81,0.4052985618099394,0.26193266943809845,0.953208716605654,0.9764250140488404,0.9789578701970347,0.02276455072057831,313163.7880490887,0.003193217217832508,cuda,1666.2890625,28.14501953125,47.0 +tfidf_ridge,log10_pgen,all,79730,42,9.48,0.6210237131878172,0.4145452767389449,0.8901421237436864,0.9434948178107029,0.9563123982662174,0.008686388197035454,46069.398916552265,0.02170638262095298,cpu,719.30078125,0.0, +esm2_8m_mlp,log10_pgen_1mm,1k,1000,42,3.72,0.8734384426546405,0.5811203823236456,0.7066749222466395,0.8428611811931889,0.8595136817400962,0.09854411271895464,531093.36010315,0.0018829081195927174,cuda,1199.21875,23.45751953125,34.0 +esm2_8m_mlp,log10_pgen_1mm,1k,1000,43,3.68,0.8676287772120702,0.5797603585575622,0.7105640414240096,0.8447384362193011,0.8588298556977834,0.08767572292374373,535898.101289993,0.0018660263911979516,cuda,1193.8125,23.45751953125,33.0 +esm2_8m_mlp,log10_pgen_1mm,1k,1000,44,4.15,0.8644075419017213,0.5799084298984422,0.7127092228837868,0.8447698749869166,0.8628554857306527,0.0476420760920801,537734.4054442797,0.0018596541152575005,cuda,1200.45703125,23.45751953125,39.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,42,9.55,0.6800570702951657,0.42365609350148575,0.8221821685480444,0.933677407276072,0.9388725488205081,0.2802916052170561,50542.42011624055,0.019785360449700248,cuda,1565.16015625,255.98876953125,24.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,43,10.39,0.6584753322447946,0.4057328196643852,0.8332892484227971,0.9367753943134008,0.945499893493824,0.2825335584781099,49021.009200498906,0.02039941682779192,cuda,1525.6953125,255.98876953125,26.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,44,10.59,0.6350050580507642,0.3884640395355252,0.8449617139927675,0.937801711477871,0.9459502912847946,0.22436193183759795,50416.68569021295,0.019834703259641744,cuda,1531.1953125,255.98876953125,27.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,45,7.0,0.6751810614818567,0.4172362618894598,0.8247229346342868,0.9339740537359036,0.9430059456408061,0.28812560326747977,78496.26553820826,0.012739459554458005,cuda,1545.91015625,255.98876953125,28.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,46,6.15,0.6685058569575997,0.40767457030235643,0.828171569821408,0.9354996251517697,0.94440611285962,0.25846381532024154,78715.08871845975,0.012704044628301188,cuda,1562.8828125,255.98876953125,24.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,47,5.45,0.642218084734104,0.4031555004716163,0.8414195480163037,0.9303131458716969,0.9390571523791333,0.2127156710366962,78601.48299372749,0.012722406269100565,cuda,1568.2421875,255.98876953125,21.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,48,9.68,0.6342088584075344,0.38472686559228664,0.8453502590541507,0.9380904108313092,0.9460257613442419,0.18461700954641433,78680.31644771349,0.012709659100882542,cuda,1557.07421875,255.98876953125,39.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,49,6.84,0.6778525531966003,0.4063589834535896,0.8233331515954379,0.9308549053584529,0.9425814646334755,0.2311877650974568,79128.20936664913,0.01263771805281719,cuda,1549.296875,255.98876953125,27.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,50,4.99,0.6604886606165908,0.40637006388958447,0.8322682333755885,0.9292828110884466,0.9392455130418712,0.21391503595003497,79443.85923264487,0.01258750531078282,cuda,1549.6328125,255.98876953125,19.0 +irrm_pretrained,log10_pgen_1mm,1k,1000,51,5.89,0.6372390774341838,0.39241801627750306,0.8438689107279675,0.9333474225263396,0.941621721540056,0.19241422764986238,78580.33201102511,0.012725830680630063,cuda,1567.53515625,255.98876953125,23.0 +irrm_scratch,log10_pgen_1mm,1k,1000,42,23.13,0.5675127906661855,0.35095961425392186,0.8761671416561154,0.937109291811129,0.9487226634897591,0.07096433712849722,49901.41476917611,0.020039511998319045,cuda,1547.078125,255.98876953125,60.0 +irrm_scratch,log10_pgen_1mm,1k,1000,43,15.61,0.5487943112163485,0.33568773672051616,0.8842012714556997,0.9412726757779909,0.949084309670342,0.06602086182336614,49383.01988887116,0.020249875407586358,cuda,1530.03125,255.98876953125,40.0 +irrm_scratch,log10_pgen_1mm,1k,1000,44,9.52,0.6648318977548506,0.42654361325296536,0.8300550415051868,0.9127595437394332,0.9188339077448422,0.0633631568418735,49671.77068032257,0.020132159299007017,cuda,1536.828125,255.98876953125,23.0 +irrm_scratch,log10_pgen_1mm,1k,1000,45,8.8,0.5754026215571453,0.3537011322377345,0.8727000416146815,0.9359365018938964,0.9437906456738564,0.08697872636101138,78318.26257008441,0.012768413996737137,cuda,1567.0390625,255.98876953125,35.0 +irrm_scratch,log10_pgen_1mm,1k,1000,46,10.78,0.5914339264050594,0.3559291734268365,0.8655078129659881,0.9347964630335666,0.9440643595280251,0.08504480157971957,78297.12279984003,0.012771861394657061,cuda,1540.38671875,255.98876953125,44.0 +irrm_scratch,log10_pgen_1mm,1k,1000,47,10.54,0.6312892435409316,0.37127948106885506,0.8467708585419784,0.938358781754315,0.947146302229228,0.20864989401004486,78701.99357101899,0.012706158441814075,cuda,1542.140625,255.98876953125,42.0 +irrm_scratch,log10_pgen_1mm,1k,1000,48,5.42,0.7789986520475849,0.4847890187497359,0.766676758356806,0.9101877824346927,0.9124383703561733,0.33916751335536144,78214.11359544039,0.012785416263520711,cuda,1536.45703125,255.98876953125,21.0 +irrm_scratch,log10_pgen_1mm,1k,1000,49,6.06,0.7225787829148815,0.4803692043319228,0.799250256201727,0.9158206089433159,0.9218188258377523,0.31839776090918176,77636.35061136573,0.012880564221853094,cuda,1561.94921875,255.98876953125,23.0 +irrm_scratch,log10_pgen_1mm,1k,1000,50,9.09,0.6073520635952829,0.3868976601348002,0.8581708134341803,0.9344983541620274,0.9446860349277294,0.19716863613669644,73919.43957497159,0.013528241092598736,cuda,1545.10546875,255.98876953125,35.0 +irrm_scratch,log10_pgen_1mm,1k,1000,51,8.61,0.6263301451334247,0.37914429651656706,0.8491687887564309,0.9280868284134893,0.9410847460057685,0.07778436074450072,79116.83397937867,0.01263953509894802,cuda,1553.6875,255.98876953125,34.0 +sceptr_mlp,log10_pgen_1mm,1k,1000,42,5.15,0.8810064957993323,0.5850634018423648,0.7015697722686547,0.8403622848328763,0.8545658323343029,0.10923790320890234,882821.5452160854,0.0011327317569659373,cuda,1177.9609375,21.20751953125,51.0 +sceptr_mlp,log10_pgen_1mm,1k,1000,43,3.26,0.8746353820476067,0.5889435925023606,0.7058704398980913,0.8414885522865576,0.8527221642098379,0.07441716460277809,914632.7917439189,0.0010933349525915341,cuda,1174.5703125,21.20751953125,32.0 +sceptr_mlp,log10_pgen_1mm,1k,1000,44,3.58,0.8782573629605056,0.5871160035537236,0.70342933697857,0.8410865411091808,0.8535996617833912,0.09703216727003491,890798.2918995839,0.001122588591708622,cuda,1170.9921875,21.20751953125,34.0 +tcr_bert_mlp,log10_pgen_1mm,1k,1000,42,6.18,0.6816212251140086,0.4366693685684581,0.82136325351949,0.9088467547401823,0.924966468087286,0.06101595375050499,314954.2306226356,0.0031750645102403986,cuda,1248.34765625,28.14501953125,52.0 +tcr_bert_mlp,log10_pgen_1mm,1k,1000,43,4.14,0.6793982191316419,0.4378072850036896,0.8225265476865682,0.9087161823299943,0.9238383313662871,0.053667050221566405,317582.12648712896,0.0031487918135107277,cuda,1246.90234375,28.14501953125,35.0 +tcr_bert_mlp,log10_pgen_1mm,1k,1000,44,4.65,0.6878677483806852,0.4429540581453104,0.8180741200953934,0.9072606006109809,0.9226372171396086,0.04333846563685214,309389.4351772415,0.003232172421877059,cuda,1249.19140625,28.14501953125,39.0 +tfidf_ridge,log10_pgen_1mm,1k,1000,42,0.32,1.0852944231152912,0.7353522020723512,0.5471235558436147,0.750775753234674,0.7954033678341129,0.18915192191042499,47177.071821432546,0.021196737342772085,cpu,654.21875,0.0, +esm2_8m_mlp,log10_pgen_1mm,10k,10000,42,11.56,0.6430747866190542,0.40745519085125576,0.8409961815941281,0.9185196137541594,0.9363519119678613,0.08332422824857887,538169.2784562818,0.0018581514033436883,cuda,1216.796875,23.45751953125,50.0 +esm2_8m_mlp,log10_pgen_1mm,10k,10000,43,13.99,0.6306860686151857,0.40229598440644404,0.8470635288924342,0.9211499809443683,0.939748810891317,0.038799049758178106,527803.3593734449,0.0018946450079194257,cuda,1215.671875,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen_1mm,10k,10000,44,8.66,0.6611521713168533,0.42557124266900453,0.8319310656720827,0.9124578749759534,0.9301727900061972,0.039136087227110615,517421.681469499,0.0019326596387688251,cuda,1215.515625,23.45751953125,36.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,42,41.65,0.35014197467911184,0.24343691858470326,0.9528618056553566,0.9816822892074405,0.9790687929250388,0.1665855199568644,48129.36627137599,0.020777335698989468,cuda,1548.6015625,255.36376953125,30.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,43,29.39,0.4313657428770761,0.2741951117819419,0.9284555591077652,0.9790466822962967,0.9782558105357633,0.2068312263721413,48760.62302536395,0.020508351574585647,cuda,1543.375,255.36376953125,21.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,44,29.98,0.37998173937666613,0.23754270438177846,0.9444850386158077,0.9784090079761774,0.9786407582301763,0.15045958642347596,45442.359880088385,0.02200589939956382,cuda,1536.1328125,255.36376953125,21.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,45,19.5,0.38967208225694244,0.24806408147764908,0.9416174347071837,0.9797404689589103,0.9770233552637545,0.16339216413745158,78750.25101444994,0.012698372222540718,cuda,1575.671875,255.36376953125,26.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,46,13.32,0.45681643268603583,0.2994579549083242,0.9197642298987043,0.9786137525002476,0.9770029606653783,0.24480818914323121,78082.57722053208,0.01280695432446672,cuda,1598.0390625,255.36376953125,18.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,47,15.65,0.3302367401562472,0.2121627214782346,0.9580689858485731,0.9807385058013642,0.9785634446932512,0.08867936156471229,78239.31749460826,0.012781297588247921,cuda,1575.87890625,255.36376953125,21.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,48,9.01,0.4601046655746993,0.3112046844265132,0.918604974508381,0.9710594170731376,0.9674161539957851,0.23073410752599088,78125.415720281,0.012799931888751595,cuda,1584.9296875,255.36376953125,12.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,49,13.28,0.42220643993618356,0.2868667920655672,0.9314615472701736,0.9776291443929536,0.9737151019226785,0.22247335337541346,78233.47222931472,0.012782252551297241,cuda,1577.703125,255.36376953125,18.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,50,11.29,0.40000050278216176,0.2671247144719602,0.9384815104664057,0.9757313885303056,0.9723141763195362,0.166277798141714,78714.28512446466,0.012704174323869923,cuda,1581.85546875,255.36376953125,15.0 +irrm_pretrained,log10_pgen_1mm,10k,10000,51,13.43,0.46614344171337685,0.3209338557302386,0.916454368508382,0.9769192225657825,0.9705874373403192,0.26804017515069967,78940.1169968257,0.01266783022427255,cuda,1577.5625,255.36376953125,18.0 +irrm_scratch,log10_pgen_1mm,10k,10000,42,19.45,0.44597789242184976,0.27716016426362267,0.923526450047828,0.9619298081954418,0.9655188940132988,0.045043227895579865,49678.4435373698,0.0201294551277108,cuda,1547.33984375,255.36376953125,14.0 +irrm_scratch,log10_pgen_1mm,10k,10000,43,29.16,0.41747837743011523,0.2587127286235791,0.9329880025986563,0.9706742053550433,0.9756462190319197,0.058843259232818426,50159.388972695255,0.019936447003857237,cuda,1569.26171875,255.36376953125,21.0 +irrm_scratch,log10_pgen_1mm,10k,10000,44,44.48,0.3686214976297636,0.24951271291904806,0.9477548587464316,0.9758187390687287,0.9826749820055413,0.08981302069059369,49087.11666474849,0.02037194416672963,cuda,1539.15625,255.36376953125,32.0 +irrm_scratch,log10_pgen_1mm,10k,10000,45,21.04,0.37169950839382454,0.24491287615278978,0.9468787161492357,0.9752518625507858,0.9805370387548386,0.09421506924227235,76535.86619263089,0.01306576968088568,cuda,1584.4375,255.36376953125,28.0 +irrm_scratch,log10_pgen_1mm,10k,10000,46,14.89,0.38276351460773933,0.2471445237431523,0.9436692340359019,0.971782066968998,0.9755962370095889,0.01284569682245751,78559.09322070763,0.012729271163944735,cuda,1589.38671875,255.36376953125,20.0 +irrm_scratch,log10_pgen_1mm,10k,10000,47,11.18,0.45773894183737657,0.27476089458618436,0.9194398415263564,0.9648299509760926,0.9713430094180486,0.14051059713170738,78499.43459207573,0.012738945257332424,cuda,1589.20703125,255.36376953125,15.0 +irrm_scratch,log10_pgen_1mm,10k,10000,48,23.87,0.3456903560926719,0.20971918710266105,0.9540527929814415,0.9771313469305555,0.9817604989495604,0.04363672955178386,78526.85782180715,0.012734496549819887,cuda,1594.03515625,255.36376953125,32.0 +irrm_scratch,log10_pgen_1mm,10k,10000,49,14.1,0.4073696215734271,0.2520798791207299,0.9361939487404316,0.9702873629167561,0.9727032762915658,0.11533500071674543,78575.93218125278,0.012726543258733205,cuda,1607.69140625,255.36376953125,19.0 +irrm_scratch,log10_pgen_1mm,10k,10000,50,17.33,0.3638889125453413,0.22850331503421772,0.9490877564538704,0.9744755401401348,0.980159241404828,0.03594200359354233,78979.94634026023,0.012661441876546926,cuda,1608.0625,255.36376953125,23.0 +irrm_scratch,log10_pgen_1mm,10k,10000,51,21.79,0.3959283377167753,0.23382869758806193,0.939727700302192,0.9764227136087295,0.9823862851349912,0.13576840298631876,78324.60890559829,0.012767379422286838,cuda,1599.27734375,255.36376953125,30.0 +sceptr_mlp,log10_pgen_1mm,10k,10000,42,9.63,0.645827637038086,0.4068036199738689,0.8396319527003634,0.9164770048749564,0.9360747724957823,0.027719664424225597,916432.6280458241,0.0010911876873397367,cuda,1181.9296875,21.20751953125,42.0 +sceptr_mlp,log10_pgen_1mm,10k,10000,43,9.19,0.6532333481759839,0.409299590819611,0.8359329819489905,0.9145515186929961,0.9359316492492474,0.032295013070945074,887704.5331397445,0.001126500950111266,cuda,1179.8046875,21.20751953125,39.0 +sceptr_mlp,log10_pgen_1mm,10k,10000,44,9.67,0.6495695749089844,0.4056266461442478,0.8377682178782393,0.9160401203643372,0.9367067372055377,0.05914198352740034,915879.2358882766,0.0010918470042943357,cuda,1178.87890625,21.20751953125,42.0 +tcr_bert_mlp,log10_pgen_1mm,10k,10000,42,8.95,0.43260012130778047,0.27748984096282453,0.9280455160181926,0.9633803919392315,0.9706797264555369,0.01180787024418463,310603.97321174777,0.003219533831649574,cuda,1276.2109375,28.14501953125,34.0 +tcr_bert_mlp,log10_pgen_1mm,10k,10000,43,15.22,0.421858887434182,0.26746360618628684,0.9315743399900089,0.9652323883296584,0.9732567446625279,0.009886096964405733,317514.9456476062,0.0031494580450705757,cuda,1272.53515625,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen_1mm,10k,10000,44,8.96,0.42981001830337606,0.27754906861007,0.9289706799220001,0.9639476620845376,0.9706564063195163,0.010106714404347067,317600.07131345535,0.003148613902586471,cuda,1274.28125,28.14501953125,35.0 +tfidf_ridge,log10_pgen_1mm,10k,10000,42,0.79,0.6798483740334681,0.4456090287409803,0.8222912894561352,0.9069107278125044,0.9418410114465724,0.018927879241947294,58098.109898958115,0.01721226390564443,cpu,648.9140625,0.0, +esm2_8m_mlp,log10_pgen_1mm,all,79730,42,74.3,0.46899616986313375,0.28366711611461526,0.9154286659694916,0.9569846240109816,0.9730241066544465,0.018638176980045787,547451.8876378721,0.0018266445373213852,cuda,1299.3828125,23.45751953125,60.0 +esm2_8m_mlp,log10_pgen_1mm,all,79730,43,59.95,0.4863207128254932,0.29147505686690234,0.9090651984987221,0.9550695754445274,0.9709942057450834,0.0882283553982555,529579.7460452183,0.0018882897381702637,cuda,1304.2421875,23.45751953125,49.0 +esm2_8m_mlp,log10_pgen_1mm,all,79730,44,71.01,0.48283508427727606,0.28513132216226056,0.91036404937305,0.9563536575939271,0.9721687047226008,0.08814281761993255,554802.0910919767,0.0018024445402355506,cuda,1299.171875,23.45751953125,59.0 +irrm_pretrained,log10_pgen_1mm,all,79730,42,140.33,0.29083007096135444,0.2098718946239174,0.9674790511927934,0.9907727648593888,0.9874548324427458,0.17782213587315893,44069.23582390069,0.022691566606599874,cuda,1558.08203125,254.86376953125,14.0 +irrm_pretrained,log10_pgen_1mm,all,79730,43,309.62,0.19178641612678832,0.11479367564442367,0.9858576924817093,0.994653716168127,0.9941628093444557,0.06551083671283722,47695.29669604378,0.020966427913697155,cuda,1578.21875,254.86376953125,33.0 +irrm_pretrained,log10_pgen_1mm,all,79730,44,220.33,0.23311122370676388,0.12960197515861094,0.9791065115378612,0.9896635201545613,0.9850500673303398,0.02886887512971913,43989.910272901674,0.02273248555853528,cuda,1558.9921875,254.86376953125,22.0 +irrm_pretrained,log10_pgen_1mm,all,79730,45,269.72,0.15868507534782078,0.08505078419565594,0.9903181855762656,0.9953405868926234,0.9970612987096276,-0.018548766321781516,76408.93078313439,0.013087475374288683,cuda,1586.76953125,254.86376953125,60.0 +irrm_pretrained,log10_pgen_1mm,all,79730,46,327.58,0.17893124003444805,0.11649784998203631,0.9876900317543066,0.9944873753403464,0.9934121864733891,0.05748933470935837,48046.353383919544,0.02081323408687008,cuda,1561.046875,254.86376953125,34.0 +irrm_pretrained,log10_pgen_1mm,all,79730,47,91.12,0.22684790947864455,0.12761343549003248,0.9802141755269492,0.9913132664242993,0.9879244544819463,0.0428836753199869,76291.03266716524,0.013107700407762183,cuda,1583.94921875,254.86376953125,20.0 +irrm_pretrained,log10_pgen_1mm,all,79730,48,91.39,0.2672039101747195,0.15850310311757398,0.9725482391013425,0.989855109751394,0.985813195353334,0.10367822122398167,76762.82293730929,0.013027139463287856,cuda,1583.31640625,254.86376953125,20.0 +irrm_pretrained,log10_pgen_1mm,all,79730,49,492.85,0.15331362972511323,0.09336122804377556,0.9909625455636024,0.9958130660577037,0.995566207785459,0.03562421480685179,47207.914573014874,0.021182888696625945,cuda,1568.7890625,254.86376953125,51.0 +irrm_pretrained,log10_pgen_1mm,all,79730,50,268.43,0.15444098240428322,0.09742089487884842,0.9908291476734726,0.9961100967107266,0.9968357584571697,0.056485239300779894,76450.86408964357,0.01308029689275239,cuda,1590.7421875,254.86376953125,60.0 +irrm_pretrained,log10_pgen_1mm,all,79730,51,104.31,0.24259668004354265,0.12878927401599008,0.9773715766288749,0.9901045476909557,0.9860099208508144,0.056035867666271395,76305.9339249849,0.013105140695651316,cuda,1581.4140625,254.86376953125,23.0 +irrm_scratch,log10_pgen_1mm,all,79730,42,182.2,0.24742771296521945,0.157050990621428,0.9764613651971807,0.9929284722679887,0.9944522727001026,0.09369640877203453,44012.93829160942,0.022720591689981286,cuda,1557.33984375,254.86376953125,19.0 +irrm_scratch,log10_pgen_1mm,all,79730,43,164.46,0.21579927855743838,0.14238205001424453,0.9820945777897901,0.9920916216231709,0.9937298104111683,-0.02161755598062506,48196.633862531206,0.0207483369658605,cuda,1573.13671875,254.86376953125,18.0 +irrm_scratch,log10_pgen_1mm,all,79730,44,121.45,0.32828928131963364,0.20981698564201592,0.9585620754190175,0.9815517547820619,0.9873401004095298,0.05558387703943304,47504.06554863723,0.021050829827946958,cuda,1552.19921875,254.86376953125,13.0 +irrm_scratch,log10_pgen_1mm,all,79730,45,211.91,0.2848072379190708,0.20553148194253587,0.9688120640790565,0.9922746098718975,0.9950765364472423,0.042364935671912546,47078.90490162043,0.021240935873289193,cuda,1564.90625,254.86376953125,22.0 +irrm_scratch,log10_pgen_1mm,all,79730,46,54.58,0.32282093426343655,0.1940359213846715,0.9599310493194535,0.9819233271806727,0.9836484176549345,0.09733460792348787,76869.10938834613,0.013009126916612967,cuda,1576.37109375,254.86376953125,12.0 +irrm_scratch,log10_pgen_1mm,all,79730,47,147.26,0.2559991761785086,0.15067211437522288,0.9748022522085213,0.9940023026524442,0.9961921259351151,0.11158205967586128,76798.71471554741,0.013021051246806303,cuda,1617.33984375,254.86376953125,32.0 +irrm_scratch,log10_pgen_1mm,all,79730,48,182.77,0.2335171541479004,0.15523083215682953,0.9790336820331119,0.9911588526121001,0.9941198058114114,0.07813050404736843,48080.641854246656,0.0207983912326178,cuda,1572.85546875,254.86376953125,19.0 +irrm_scratch,log10_pgen_1mm,all,79730,49,77.54,0.27385095790334146,0.19382559150444986,0.9711654539881946,0.9909897033406608,0.9936805074797037,0.16402027947482764,76598.80083274224,0.013055034662795256,cuda,1589.9375,254.86376953125,17.0 +irrm_scratch,log10_pgen_1mm,all,79730,50,71.91,0.23204268378956183,0.15422337004000328,0.9792976165218135,0.9896301157879785,0.9930708539631528,0.003519611339806289,76968.537959956,0.01299232162263839,cuda,1580.3046875,254.86376953125,16.0 +irrm_scratch,log10_pgen_1mm,all,79730,51,134.45,0.25863070214883893,0.16719389261637624,0.9742815526056321,0.9873098753313578,0.991357120638237,-0.015157397736740554,45549.57022164899,0.02195410396045221,cuda,1559.40625,254.86376953125,14.0 +sceptr_mlp,log10_pgen_1mm,all,79730,42,47.72,0.4590071571918785,0.26491759622764105,0.918992821822185,0.959962995648993,0.9752358006705864,0.066629467263402,908790.3950185104,0.0011003637422682397,cuda,1199.44140625,21.20751953125,40.0 +sceptr_mlp,log10_pgen_1mm,all,79730,43,54.04,0.46201987164360353,0.2647853562151094,0.9179259433906491,0.958619255812622,0.975345369198441,0.04454950212876287,920832.7705944345,0.0010859735143379616,cuda,1194.40234375,21.20751953125,47.0 +sceptr_mlp,log10_pgen_1mm,all,79730,44,57.35,0.4469669894594107,0.2609785939350046,0.9231868653260547,0.9610926944199375,0.975631702210054,0.033990439569041256,909364.3039765669,0.0010996692916437245,cuda,1198.41015625,21.20751953125,49.0 +tcr_bert_mlp,log10_pgen_1mm,all,79730,42,58.15,0.3098943987334235,0.19844359774924528,0.9630757194284447,0.9815324691892897,0.9850971523921663,0.013213505453879129,318647.4263871291,0.003138264794221455,cuda,1670.19921875,28.14501953125,44.0 +tcr_bert_mlp,log10_pgen_1mm,all,79730,43,79.21,0.30509861900551083,0.19526537661917165,0.9642097219020543,0.9820934038858983,0.9869364844825658,-0.0119739816780719,316693.254885937,0.0031576296134256747,cuda,1670.44921875,28.14501953125,60.0 +tcr_bert_mlp,log10_pgen_1mm,all,79730,44,48.95,0.3193936611517518,0.20682718145954618,0.960777328144496,0.9802463373113919,0.983617638847924,0.007206928689406769,320776.8613929061,0.00311743183612967,cuda,1673.2109375,28.14501953125,37.0 +tfidf_ridge,log10_pgen_1mm,all,79730,42,9.66,0.5542224821064449,0.3670566172834339,0.8818991924431581,0.9391308258396079,0.9565974661198886,0.0035328855355764607,44726.53155722081,0.022358094070421077,cpu,716.12109375,0.0, diff --git a/results/pgen_constlr/pgen_summary.csv b/results/pgen_constlr/pgen_summary.csv new file mode 100644 index 0000000..70f0a0d --- /dev/null +++ b/results/pgen_constlr/pgen_summary.csv @@ -0,0 +1,37 @@ +target,train_subset,arm,seeds,train_size,rmse_mean,rmse_std,mae_mean,mae_std,r2_mean,r2_std,pearson_r_mean,pearson_r_std,spearman_rho_mean,spearman_rho_std,bias_mean,bias_std,training_seconds_mean,inference_seq_per_second_mean,peak_rss_mb_mean,peak_gpu_mb_mean +log10_pgen,1k,irrm_scratch,10,1000,0.8379438263375423,0.07867249592710443,0.5413980735295911,0.0580446314810705,0.798406523904873,0.040498992051886576,0.9049139054549931,0.009840024528434311,0.904333465892336,0.011361519326849367,0.1551782296232404,0.15390117173821022,10.331,62821.00447401784,1542.425390625,255.98876953125 +log10_pgen,1k,irrm_pretrained,10,1000,0.8550251478865704,0.025339719180604656,0.5575133459336978,0.013908713333604477,0.7915913348278159,0.01233799418292649,0.9148794772734956,0.0038483435506742065,0.9084727474646785,0.0031749550782257897,0.30998608299664804,0.05861494528750305,8.241,62465.82250303106,1542.067578125,255.98876953125 +log10_pgen,1k,tcr_bert_mlp,3,1000,0.8657739615527257,0.009257369532250346,0.581423879327303,0.0019239794308733246,0.7864709473583439,0.004575295736734588,0.8890168834635155,0.0008150930534225373,0.8913271868559915,0.001959429986301301,0.07026857252856558,0.05214573058201129,5.113333333333333,306623.3384164976,1247.0065104166667,28.14501953125 +log10_pgen,1k,esm2_8m_mlp,3,1000,1.083446236552734,0.007022470402211829,0.7379947173621156,0.002171722887800877,0.6656188970763066,0.004342632397618175,0.8185395891648021,0.001412940229596168,0.8260198092373562,0.0025814804548312297,0.12072431084892372,0.030412381356488497,3.93,523847.50668071903,1200.0143229166667,23.45751953125 +log10_pgen,1k,sceptr_mlp,3,1000,1.0976038433169357,0.006897307587346891,0.7627633957781601,0.0019332084823762091,0.6568235298151776,0.0043050872520435145,0.8118926570544515,0.0021225197902238176,0.8060410172602768,0.0017155240942351325,0.0814823994847304,0.03710128987080741,3.7999999999999994,856494.9556877855,1169.6171875,21.20751953125 +log10_pgen,1k,tfidf_ridge,1,1000,1.276955755696797,,0.8800488403651113,,0.5355205204794624,,0.7459285398734663,,0.7755971181289645,,0.24748485510618118,,0.44,46905.79064077635,655.55078125,0.0 +log10_pgen,10k,irrm_scratch,10,10000,0.5051515834654494,0.054951652825074,0.3271744047929508,0.041713085823000436,0.9265386059854499,0.016676524110579356,0.9654555653445005,0.006190666031767679,0.9670376819703789,0.006265392894722844,0.05654215627852953,0.12147756529555402,28.3,64329.13142336295,1563.64609375,255.36376953125 +log10_pgen,10k,irrm_pretrained,10,10000,0.5422183832232536,0.057352525306723294,0.38187654948467087,0.04724556949707446,0.9154108858067206,0.017514758663127572,0.9730908568087688,0.0024674093272134758,0.9658762303337521,0.0038073550349844287,0.2978773964467407,0.06780555517243432,23.639,69548.45281110329,1576.52421875,255.36376953125 +log10_pgen,10k,tcr_bert_mlp,3,10000,0.5680795314054411,0.004031119636653914,0.37845706811385776,0.004133997900017408,0.9080720250297891,0.0013031846861411404,0.9532419033372737,0.0005358170938291702,0.9548576837254182,0.0018316667053695878,0.026503910442178022,0.01914405941888179,13.656666666666666,315803.3855036165,1272.2252604166667,28.14501953125 +log10_pgen,10k,tfidf_ridge,1,10000,0.76592457502429,,0.5084579081617813,,0.83289600206607,,0.9129297738264299,,0.9404228939259158,,0.03423942511518709,,0.8,59776.55467359597,649.6875,0.0 +log10_pgen,10k,esm2_8m_mlp,3,10000,0.8037714586009032,0.008393865622556233,0.5228154960854687,0.00485496479799742,0.8159602784147432,0.003850977620423332,0.9049954430824227,0.0019710768157375443,0.9189027339342127,0.0026954336757552023,0.09818479091695358,0.025299728810900795,12.799999999999999,528172.193105634,1217.1145833333333,23.45751953125 +log10_pgen,10k,sceptr_mlp,3,10000,0.8328461902283072,0.00800144086501869,0.5403426692083185,0.00653094512617079,0.8024071788961349,0.003790759788035956,0.8971883816484573,0.002363627118951327,0.9085799446350773,0.002673813985561092,0.09365712764456295,0.007758322471483018,10.156666666666666,919994.4011759133,1182.0065104166667,21.20751953125 +log10_pgen,all,irrm_pretrained,10,79730,0.29094909474054964,0.04741822316545527,0.18458412851304246,0.035762733378958134,0.9753107011546698,0.00800254329667357,0.9903982636366593,0.0027969455400128004,0.9876218128470395,0.00429003780174096,0.11541883526450339,0.04597981539220331,204.63899999999998,61641.35071767869,1581.82890625,254.86376953125 +log10_pgen,all,irrm_scratch,10,79730,0.33578636528681594,0.07126425216797283,0.22252437734980995,0.0446965625642714,0.9665805886364465,0.014984588581583783,0.9897984543781598,0.0059964201920581885,0.9904199184753828,0.0052830720928398155,0.09984075885395624,0.08298558555562141,145.683,62300.467733850506,1569.30625,254.86376953125 +log10_pgen,all,tcr_bert_mlp,3,79730,0.4104244762370142,0.006479777792889647,0.26407476038834926,0.006878891588527168,0.9520096961679507,0.001520406413171397,0.9760846675180218,0.0011130635940318988,0.9793523127399316,0.0014666079786796635,0.029935375799948306,0.028979884407407768,64.49666666666667,317264.54896651936,1668.7174479166667,28.14501953125 +log10_pgen,all,sceptr_mlp,3,79730,0.5833243496880652,0.0034583517709830796,0.3535763262418406,0.0021605759059627762,0.9030728998757227,0.0011483300575644375,0.9524100668159482,0.0010005163984225006,0.9648992404555322,0.0002256061400935981,0.08710927532999313,0.005301695083341565,67.07666666666667,892283.828687017,1200.2526041666667,21.20751953125 +log10_pgen,all,esm2_8m_mlp,3,79730,0.6031239348944677,0.012581097198351444,0.381114139482605,0.009560811361581226,0.8963536707843666,0.004348451618740226,0.9474678951563668,0.0019709996661059505,0.9598594019219823,0.0027189754719963667,0.04333215924767322,0.030947112086310898,66.03,526840.578957009,1300.21875,23.45751953125 +log10_pgen,all,tfidf_ridge,1,79730,0.6210237131878172,,0.4145452767389449,,0.8901421237436864,,0.9434948178107029,,0.9563123982662174,,0.008686388197035454,,9.48,46069.398916552265,719.30078125,0.0 +log10_pgen_1mm,1k,irrm_scratch,10,1000,0.6314524434831694,0.07261686713463565,0.3925300930693855,0.05359654264664398,0.8448668784488793,0.037029600737626185,0.9288826831963857,0.01158522271877412,0.9371670165463476,0.013830910702485056,0.1513540048890253,0.1074965335519036,10.756,69316.03220404696,1546.0703125,255.98876953125 +log10_pgen_1mm,1k,irrm_pretrained,10,1000,0.6569231613419195,0.018426714329123007,0.4035793214577793,0.012179916879442062,0.8339567738188751,0.009299913259190176,0.9339616887631262,0.0031160599077837098,0.942626640503833,0.002856669189694695,0.236862622340095,0.0381077679541767,7.6530000000000005,70162.56693153805,1552.2625,255.98876953125 +log10_pgen_1mm,1k,tcr_bert_mlp,3,1000,0.682962397542112,0.004391160329175946,0.43914357057248604,0.0033486676619782895,0.8206546404338172,0.0023092481985602415,0.9082745125603858,0.0008804972285221423,0.9238140055310605,0.0011648159957729678,0.052673823202974514,0.008880499503683912,4.989999999999999,313975.2640956687,1248.1471354166667,28.14501953125 +log10_pgen_1mm,1k,esm2_8m_mlp,3,1000,0.868491587256144,0.004576857325871695,0.58026305692655,0.0007461477118348201,0.7099827288514787,0.0030588624151971274,0.8441231641331356,0.0010930223257392714,0.8603996743895107,0.002154103414763927,0.0779539705782595,0.02680743711530656,3.85,534908.622279141,1197.8294270833333,23.45751953125 +log10_pgen_1mm,1k,sceptr_mlp,3,1000,0.8779664136024815,0.003195506416212866,0.5870409992994831,0.0019411824020671429,0.7036231830484386,0.0021568768483211037,0.840979126076205,0.0005707653513373905,0.8536292194425107,0.0009221893947024148,0.09356241169390511,0.01766777750843526,3.9966666666666666,896084.2096198626,1174.5078125,21.20751953125 +log10_pgen_1mm,1k,tfidf_ridge,1,1000,1.0852944231152912,,0.7353522020723512,,0.5471235558436147,,0.750775753234674,,0.7954033678341129,,0.18915192191042499,,0.32,47177.071821432546,654.21875,0.0 +log10_pgen_1mm,10k,irrm_scratch,10,10000,0.3957156960248885,0.036479454715929324,0.24763349791340467,0.020427223218532808,0.9393319301582345,0.011325238425838581,0.9718603596711267,0.005104883477754595,0.9768325682014181,0.005697693481619573,0.07719530078638195,0.044325297668403874,21.729,69692.66884291462,1582.785546875,255.36376953125 +log10_pgen_1mm,10k,irrm_pretrained,10,10000,0.4086669762038501,0.046905976308239894,0.2700989539306911,0.03504841161795715,0.9350255454586733,0.014519371549026574,0.9779569879302615,0.002987619383779735,0.9752587991891681,0.004033122939199344,0.19082814817916946,0.0534713821299041,19.65,69141.77849773047,1569.975,255.36376953125 +log10_pgen_1mm,10k,tcr_bert_mlp,3,10000,0.4280896756817795,0.005573437927670511,0.2741675052530605,0.00580582242250533,0.9295301786434006,0.0018297346174931833,0.9641868141178093,0.0009488771937410975,0.9715309591458604,0.001494619581813537,0.010600227204312479,0.0010516507488302893,11.043333333333335,315239.66339093645,1274.3424479166667,28.14501953125 +log10_pgen_1mm,10k,esm2_8m_mlp,3,10000,0.6449710088503644,0.015321311873335977,0.41177413930890144,0.012223927299852232,0.839996925386215,0.0076155594348420825,0.9173758232248271,0.004457507052465646,0.9354245042884585,0.0048549055347065495,0.053753121744622524,0.025609883901828074,11.403333333333334,527798.1064330753,1215.9947916666667,23.45751953125 +log10_pgen_1mm,10k,sceptr_mlp,3,10000,0.6495435200410181,0.0037029243182388977,0.4072432856459092,0.001875529278483193,0.8377777175091977,0.0018495036731835572,0.91568954797743,0.0010094800681917497,0.9362377196501891,0.000412436863463324,0.03971888700752367,0.016975746305393863,9.496666666666668,906672.1323579485,1180.2044270833333,21.20751953125 +log10_pgen_1mm,10k,tfidf_ridge,1,10000,0.6798483740334681,,0.4456090287409803,,0.8222912894561352,,0.9069107278125044,,0.9418410114465724,,0.018927879241947294,,0.79,58098.109898958115,648.9140625,0.0 +log10_pgen_1mm,all,irrm_pretrained,10,79730,0.20977471380034785,0.04934704474272933,0.1261504115150865,0.036345323942476096,0.9822377157037178,0.008214782973867617,0.9928114050050126,0.0026846786292086143,0.9909290731229282,0.004906346197013534,0.06058496344211638,0.05150767433523792,231.56800000000004,61322.829515201796,1575.13203125,254.86376953125 +log10_pgen_1mm,all,irrm_scratch,10,79730,0.265318511919295,0.03763402497421702,0.17299632302977777,0.025127194527394783,0.9724441689161771,0.007977149306655412,0.9893860635450332,0.0044306439958006795,0.9922667551450598,0.003882647933197412,0.060945733022736606,0.059777932774134894,134.853,58765.79175768857,1574.3796875,254.86376953125 +log10_pgen_1mm,all,tcr_bert_mlp,3,79730,0.31146222629689535,0.007275343286903232,0.20017871860932104,0.005973007290388119,0.9626875898249984,0.0017488039251256668,0.9812907367955267,0.0009469633811877807,0.9852170919075519,0.0016626705031840912,0.0028154841550713326,0.013155453843526574,62.10333333333333,318705.8475553241,1671.2864583333333,28.14501953125 +log10_pgen_1mm,all,sceptr_mlp,3,79730,0.45599800609829755,0.00796483259339435,0.2635605154592517,0.0022369870190182945,0.9200352101796296,0.0027810528614914155,0.9598916486271841,0.0012382618601008899,0.9754042906930271,0.00020442190227909016,0.048389802987068714,0.016654953173757057,53.03666666666667,912995.8231965039,1197.41796875,21.20751953125 +log10_pgen_1mm,all,esm2_8m_mlp,3,79730,0.4793839889886344,0.009163377719603393,0.28675783171459274,0.00415031742047246,0.9116193046137546,0.003362317482388092,0.9561359523498121,0.0009759095413796944,0.9720623390407103,0.0010191220089720955,0.06500311666607794,0.04015323839050756,68.42,543944.5749250223,1300.9322916666667,23.45751953125 +log10_pgen_1mm,all,tfidf_ridge,1,79730,0.5542224821064449,,0.3670566172834339,,0.8818991924431581,,0.9391308258396079,,0.9565974661198886,,0.0035328855355764607,,9.66,44726.53155722081,716.12109375,0.0 diff --git a/results/pgen_constlr/pgen_summary.md b/results/pgen_constlr/pgen_summary.md new file mode 100644 index 0000000..b8227f3 --- /dev/null +++ b/results/pgen_constlr/pgen_summary.md @@ -0,0 +1,69 @@ +### log10_pgen + +**train = 1k (1,000 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_scratch | 0.8379 ± 0.0787 | 0.5414 | 0.7984 | 0.9049 | 0.9043 | +0.1552 | 10 | 62821 | 1542 | +| irrm_pretrained | 0.8550 ± 0.0253 | 0.5575 | 0.7916 | 0.9149 | 0.9085 | +0.3100 | 8 | 62466 | 1542 | +| tcr_bert_mlp | 0.8658 ± 0.0093 | 0.5814 | 0.7865 | 0.8890 | 0.8913 | +0.0703 | 5 | 306623 | 1247 | +| esm2_8m_mlp | 1.0834 ± 0.0070 | 0.7380 | 0.6656 | 0.8185 | 0.8260 | +0.1207 | 4 | 523848 | 1200 | +| sceptr_mlp | 1.0976 ± 0.0069 | 0.7628 | 0.6568 | 0.8119 | 0.8060 | +0.0815 | 4 | 856495 | 1170 | +| tfidf_ridge | 1.2770 | 0.8800 | 0.5355 | 0.7459 | 0.7756 | +0.2475 | 0 | 46906 | 656 | + +**train = 10k (10,000 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_scratch | 0.5052 ± 0.0550 | 0.3272 | 0.9265 | 0.9655 | 0.9670 | +0.0565 | 28 | 64329 | 1564 | +| irrm_pretrained | 0.5422 ± 0.0574 | 0.3819 | 0.9154 | 0.9731 | 0.9659 | +0.2979 | 24 | 69548 | 1577 | +| tcr_bert_mlp | 0.5681 ± 0.0040 | 0.3785 | 0.9081 | 0.9532 | 0.9549 | +0.0265 | 14 | 315803 | 1272 | +| tfidf_ridge | 0.7659 | 0.5085 | 0.8329 | 0.9129 | 0.9404 | +0.0342 | 1 | 59777 | 650 | +| esm2_8m_mlp | 0.8038 ± 0.0084 | 0.5228 | 0.8160 | 0.9050 | 0.9189 | +0.0982 | 13 | 528172 | 1217 | +| sceptr_mlp | 0.8328 ± 0.0080 | 0.5403 | 0.8024 | 0.8972 | 0.9086 | +0.0937 | 10 | 919994 | 1182 | + +**train = all (79,730 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_pretrained | 0.2909 ± 0.0474 | 0.1846 | 0.9753 | 0.9904 | 0.9876 | +0.1154 | 205 | 61641 | 1582 | +| irrm_scratch | 0.3358 ± 0.0713 | 0.2225 | 0.9666 | 0.9898 | 0.9904 | +0.0998 | 146 | 62300 | 1569 | +| tcr_bert_mlp | 0.4104 ± 0.0065 | 0.2641 | 0.9520 | 0.9761 | 0.9794 | +0.0299 | 64 | 317265 | 1669 | +| sceptr_mlp | 0.5833 ± 0.0035 | 0.3536 | 0.9031 | 0.9524 | 0.9649 | +0.0871 | 67 | 892284 | 1200 | +| esm2_8m_mlp | 0.6031 ± 0.0126 | 0.3811 | 0.8964 | 0.9475 | 0.9599 | +0.0433 | 66 | 526841 | 1300 | +| tfidf_ridge | 0.6210 | 0.4145 | 0.8901 | 0.9435 | 0.9563 | +0.0087 | 9 | 46069 | 719 | + +### log10_pgen_1mm + +**train = 1k (1,000 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_scratch | 0.6315 ± 0.0726 | 0.3925 | 0.8449 | 0.9289 | 0.9372 | +0.1514 | 11 | 69316 | 1546 | +| irrm_pretrained | 0.6569 ± 0.0184 | 0.4036 | 0.8340 | 0.9340 | 0.9426 | +0.2369 | 8 | 70163 | 1552 | +| tcr_bert_mlp | 0.6830 ± 0.0044 | 0.4391 | 0.8207 | 0.9083 | 0.9238 | +0.0527 | 5 | 313975 | 1248 | +| esm2_8m_mlp | 0.8685 ± 0.0046 | 0.5803 | 0.7100 | 0.8441 | 0.8604 | +0.0780 | 4 | 534909 | 1198 | +| sceptr_mlp | 0.8780 ± 0.0032 | 0.5870 | 0.7036 | 0.8410 | 0.8536 | +0.0936 | 4 | 896084 | 1175 | +| tfidf_ridge | 1.0853 | 0.7354 | 0.5471 | 0.7508 | 0.7954 | +0.1892 | 0 | 47177 | 654 | + +**train = 10k (10,000 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_scratch | 0.3957 ± 0.0365 | 0.2476 | 0.9393 | 0.9719 | 0.9768 | +0.0772 | 22 | 69693 | 1583 | +| irrm_pretrained | 0.4087 ± 0.0469 | 0.2701 | 0.9350 | 0.9780 | 0.9753 | +0.1908 | 20 | 69142 | 1570 | +| tcr_bert_mlp | 0.4281 ± 0.0056 | 0.2742 | 0.9295 | 0.9642 | 0.9715 | +0.0106 | 11 | 315240 | 1274 | +| esm2_8m_mlp | 0.6450 ± 0.0153 | 0.4118 | 0.8400 | 0.9174 | 0.9354 | +0.0538 | 11 | 527798 | 1216 | +| sceptr_mlp | 0.6495 ± 0.0037 | 0.4072 | 0.8378 | 0.9157 | 0.9362 | +0.0397 | 9 | 906672 | 1180 | +| tfidf_ridge | 0.6798 | 0.4456 | 0.8223 | 0.9069 | 0.9418 | +0.0189 | 1 | 58098 | 649 | + +**train = all (79,730 sequences)** + +| arm | RMSE | MAE | R² | Pearson r | Spearman ρ | bias | train s | seq/s | peak RSS MB | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| irrm_pretrained | 0.2098 ± 0.0493 | 0.1262 | 0.9822 | 0.9928 | 0.9909 | +0.0606 | 232 | 61323 | 1575 | +| irrm_scratch | 0.2653 ± 0.0376 | 0.1730 | 0.9724 | 0.9894 | 0.9923 | +0.0609 | 135 | 58766 | 1574 | +| tcr_bert_mlp | 0.3115 ± 0.0073 | 0.2002 | 0.9627 | 0.9813 | 0.9852 | +0.0028 | 62 | 318706 | 1671 | +| sceptr_mlp | 0.4560 ± 0.0080 | 0.2636 | 0.9200 | 0.9599 | 0.9754 | +0.0484 | 53 | 912996 | 1197 | +| esm2_8m_mlp | 0.4794 ± 0.0092 | 0.2868 | 0.9116 | 0.9561 | 0.9721 | +0.0650 | 68 | 543945 | 1301 | +| tfidf_ridge | 0.5542 | 0.3671 | 0.8819 | 0.9391 | 0.9566 | +0.0035 | 10 | 44727 | 716 | diff --git a/results/pgen_constlr/pretraining_delta.csv b/results/pgen_constlr/pretraining_delta.csv new file mode 100644 index 0000000..5584101 --- /dev/null +++ b/results/pgen_constlr/pretraining_delta.csv @@ -0,0 +1,7 @@ +target,train_subset,n_scratch,n_pretrained,irrm_scratch,irrm_pretrained,rmse_delta,relative_improvement,p_value,verdict +log10_pgen,1k,10,10,0.8379438263375423,0.8550251478865704,-0.017081321549028083,-0.020384805057501967,0.5270257428158892,no detectable effect +log10_pgen,10k,10,10,0.5051515834654494,0.5422183832232537,-0.03706679975780425,-0.0733775780796686,0.15732847985243456,no detectable effect +log10_pgen,all,10,10,0.335786365286816,0.29094909474054964,0.04483727054626635,0.13352915776662955,0.117526241878387,no detectable effect +log10_pgen_1mm,1k,10,10,0.6314524434831694,0.6569231613419194,-0.02547071785874999,-0.04033671596589345,0.3072070679790528,no detectable effect +log10_pgen_1mm,10k,10,10,0.3957156960248885,0.4086669762038501,-0.012951280178961633,-0.0327287502342264,0.49998985357929854,no detectable effect +log10_pgen_1mm,all,10,10,0.26531851191929495,0.2097747138003479,0.05554379811894705,0.20934761663310725,0.011632512578723934,pretraining helps diff --git a/results/reconstruction/reconstruction_runs.csv b/results/reconstruction/reconstruction_runs.csv new file mode 100644 index 0000000..61092dc --- /dev/null +++ b/results/reconstruction/reconstruction_runs.csv @@ -0,0 +1,16 @@ +representation,seed,bottleneck_dim,train_subset,train_size,epochs_ran,best_epoch,best_val_loss,training_seconds,test_loss,exact_match,token_accuracy,normalized_levenshtein,within_edit_distance_1,mean_edit_distance,inference_seq_per_second,inference_ms_per_sequence,device +esm2_8m,42,64,all,79730,36,28,0.0733299881984026,1449.82,0.072866829064412,0.6525534263068125,0.9325163975146336,0.0563025242603035,0.7704424601183907,0.855824219925755,6925.270491542189,0.14439869189532695,cuda +esm2_8m,43,64,all,79730,35,27,0.07321660067790593,1481.52,0.07167374705656981,0.6505468044547005,0.9328907864327216,0.05586594799605164,0.7669308718771948,0.845590448479984,6302.629250780989,0.15866394169958278,cuda +esm2_8m,44,64,all,79730,34,26,0.07201942428946495,1374.37,0.07244877650951728,0.6499448178990669,0.9323724017769074,0.05579528646580723,0.7703421290257851,0.8460921039430119,6931.001781474648,0.14427928768866063,cuda +onehot,42,64,all,79730,46,38,0.015626185442297123,1947.5,0.01686641280181133,0.9147185712852413,0.988012354834297,0.01044870014041422,0.9650847797732517,0.16925855322564462,6740.392889160309,0.14835930433790723,cuda +onehot,43,64,all,79730,44,36,0.015526399291001068,1773.28,0.0167713965026614,0.9147185712852413,0.9878251603752529,0.01055498694685954,0.9642821310324069,0.171666499448179,6959.059296335155,0.143697582879719,cuda +onehot,44,64,all,79730,39,31,0.01566104817753419,1654.49,0.016842290782966673,0.9137152603591853,0.9875299691129142,0.010679486103341483,0.9635798133841678,0.17397411457810777,6721.714112326583,0.14877157571550906,cuda +sceptr,42,64,all,79730,34,26,0.015631044331269387,1438.6,0.016878187322081663,0.9052874485803151,0.9872923761456661,0.010762903291556672,0.9606702116986054,0.1700612019664894,6681.0765368225,0.14967647720970384,cuda +sceptr,43,64,all,79730,25,17,0.015448446170641826,1005.99,0.01793186505062458,0.8911407645229257,0.9852404368830683,0.012352635473980028,0.9537473663088191,0.19353867763619945,6942.691494033216,0.1440363583574805,cuda +sceptr,44,64,all,79730,36,28,0.014794133317012053,1520.72,0.016622064921718378,0.9049864553024982,0.9872923761456661,0.01057906079226048,0.9614728604394502,0.16584729607705428,6718.565058588532,0.14884130633246928,cuda +tcr_bert,42,64,all,79730,29,21,0.04164245581397644,1170.33,0.040557941994988,0.7707434533962075,0.9647642429784078,0.02883887644221048,0.8893348048560249,0.4475770041135748,6948.343259792289,0.14391919952870805,cuda +tcr_bert,43,64,all,79730,26,18,0.041497998273907565,1101.03,0.04137205002972713,0.7640212701916324,0.9634538817650997,0.030204758605446042,0.8806059997993378,0.4672418982642721,6700.394915657967,0.14924493445350925,cuda +tcr_bert,44,64,all,79730,31,23,0.041648901043794095,1249.9,0.04166440746914118,0.7751580214708538,0.9649730367981108,0.028198279038403613,0.8936490418380656,0.4376442259456205,6943.732958313603,0.14401475488810647,cuda +tcremp,42,64,all,79730,39,31,0.0247877309910762,1566.19,0.025743332906411245,0.7975318551219023,0.9779254534065792,0.020130054818765956,0.9388983646031905,0.3027992374836962,7139.259510791705,0.14007054912185218,cuda +tcremp,43,64,all,79730,46,38,0.024846020345695507,1927.35,0.025916891745649852,0.8049563559747166,0.9786814310296416,0.01929511392118879,0.942008628473964,0.2894552021671516,6815.03971278562,0.14673428800772959,cuda +tcremp,44,64,all,79730,47,39,0.024844986697038014,1893.08,0.026344755951028604,0.8025484097521822,0.9785662344394606,0.019611650793997584,0.9408046553626969,0.29266579713053076,6948.998411166585,0.14390563083063387,cuda diff --git a/results/reconstruction/reconstruction_summary.csv b/results/reconstruction/reconstruction_summary.csv new file mode 100644 index 0000000..3e1d5dc --- /dev/null +++ b/results/reconstruction/reconstruction_summary.csv @@ -0,0 +1,6 @@ +representation,seeds,exact_match_mean,exact_match_std,token_accuracy_mean,token_accuracy_std,normalized_levenshtein_mean,normalized_levenshtein_std,within_edit_distance_1_mean,within_edit_distance_1_std,mean_edit_distance_mean,mean_edit_distance_std,training_seconds_mean,epochs_mean,inference_seq_per_second_mean,source_dim,explained_variance +onehot,3,0.9143841343098892,0.000579261833239334,0.9877891614408214,0.00024319938015680297,0.01056105773020508,0.00011551268752343819,0.9643155747299422,0.0007530403832110864,0.1716330557506438,0.002357958561529363,1791.756666666667,43.0,6807.05543260735,1000,0.379862 +sceptr,3,0.9004715561352463,0.008082103892102458,0.9866083963914668,0.0011846876856216142,0.011231533185932395,0.0009752447238186415,0.9586301461489583,0.004247612915008504,0.17648239189324774,0.014920687424517944,1321.77,31.666666666666668,6780.777696481416,64,1.0 +tcremp,3,0.8016788736162671,0.003787858636103861,0.9783910396252272,0.00040730266245565675,0.019678939844650776,0.00042151801606229864,0.9405705494799506,0.001568291908323191,0.29497341226045953,0.006964885880127865,1795.54,44.0,6967.76587824797,9000,0.949757 +tcr_bert,3,0.7699742483528978,0.005608080292114219,0.9643970538472061,0.0008234554615370438,0.02908063802868671,0.00102485441003165,0.8878632821644761,0.006644867995664911,0.45082104277448914,0.015063146809206184,1173.7533333333333,28.666666666666668,6864.157044587952,768,0.910277 +esm2_8m,3,0.6510150162201933,0.0013658792586644186,0.9325931952414209,0.0002675893887576123,0.05598791957405413,0.0002747368639656366,0.7692384870071235,0.00199908285794834,0.8491689241162503,0.005769110528054075,1435.2366666666667,35.0,6719.6338412659425,320,0.974346 diff --git a/results/reconstruction/reconstruction_summary.md b/results/reconstruction/reconstruction_summary.md new file mode 100644 index 0000000..c492a6e --- /dev/null +++ b/results/reconstruction/reconstruction_summary.md @@ -0,0 +1,9 @@ +### Reconstruction at a 64-dimensional bottleneck + +| representation | source dim | expl. var | exact match | token acc | norm. Levenshtein | within ED 1 | train s | seq/s | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| onehot | 1000 | 0.380 | 0.9144 ± 0.0006 | 0.9878 ± 0.0002 | 0.0106 ± 0.0001 | 0.9643 ± 0.0008 | 1792 | 6807 | +| sceptr | 64 | 1.000 | 0.9005 ± 0.0081 | 0.9866 ± 0.0012 | 0.0112 ± 0.0010 | 0.9586 ± 0.0042 | 1322 | 6781 | +| tcremp | 9000 | 0.950 | 0.8017 ± 0.0038 | 0.9784 ± 0.0004 | 0.0197 ± 0.0004 | 0.9406 ± 0.0016 | 1796 | 6968 | +| tcr_bert | 768 | 0.910 | 0.7700 ± 0.0056 | 0.9644 ± 0.0008 | 0.0291 ± 0.0010 | 0.8879 ± 0.0066 | 1174 | 6864 | +| esm2_8m | 320 | 0.974 | 0.6510 ± 0.0014 | 0.9326 ± 0.0003 | 0.0560 ± 0.0003 | 0.7692 ± 0.0020 | 1435 | 6720 | diff --git a/scripts/preflight_anchored_tokenizers.py b/scripts/preflight_anchored_tokenizers.py new file mode 100644 index 0000000..2313df4 --- /dev/null +++ b/scripts/preflight_anchored_tokenizers.py @@ -0,0 +1,203 @@ +"""CPU gate for the three leakage-safe anchored tokenizers.""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import platform +from dataclasses import asdict +from pathlib import Path + +import torch + +from irrm_codec.anchored_tokenization import AnchoredTokenizer +from irrm_codec.multitask_data import ( + MultiTaskBenchmarkDataset, + TargetStandardizer, + collate_multitask, + load_prepared_benchmark, + resolve_encoder_tokenizer, + select_split_indices, +) +from irrm_codec.multitask_losses import IRRMCodecMultiTaskLoss, MultiTaskLossWeights +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import AA_VOCAB + + +CANDIDATES = ("edge_k", "data_anchor", "germline_anchor") + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", required=True) + parser.add_argument("--tokenizer-root", required=True) + parser.add_argument("--baseline-run-config", required=True) + parser.add_argument("--output-path", required=True) + args = parser.parse_args() + + data_dir = Path(args.data_dir) + root = Path(args.tokenizer_root) + baseline_path = Path(args.baseline_run_config) + baseline = json.loads(baseline_path.read_text(encoding="utf-8")) + if ( + baseline["model"]["latent_dim"] != 128 + or baseline["model"]["encoder_layers"] != 4 + or baseline["model"]["decoder_layers"] != 4 + ): + raise ValueError("Baseline is not the validation-selected char d128 4+4 model.") + ready = json.loads((data_dir / "READY.json").read_text(encoding="utf-8")) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Locked benchmark is not ready.") + preparation = json.loads((root / "preparation_manifest.json").read_text(encoding="utf-8")) + if preparation["external_corpus"]["overlap_after_filter"] != 0: + raise ValueError("External tokenizer corpus overlaps the locked benchmark.") + if preparation["inference_modality"] != "sequence_only": + raise ValueError("Anchored study is not sequence-only.") + + table, embeddings = load_prepared_benchmark(data_dir) + split_indices = {split: select_split_indices(table, data_dir, split) for split in ("train", "val", "test")} + standardizer = TargetStandardizer.load(data_dir / "target_standardizer.npz") + if standardizer.train_rows != len(split_indices["train"]) or standardizer.pgen_target != "log10_pgen_1mm": + raise ValueError("Locked train-only normalizer does not match the benchmark.") + + expected_training = { + "train_subset":"all","pgen_target":"log10_pgen_1mm","tokenizer_type":"char", + "tokenizer_path":None,"max_sequence_len":40,"d_model":320,"latent_dim":128,"nhead":8, + "encoder_layers":4,"decoder_layers":4,"ff_dim":1280,"dropout":0.1, + "tcremp_head_dim":1024,"pgen_head_dim":256,"decoder_memory_tokens":4, + "tcremp_loss_weight":1.0,"pgen_loss_weight":1.0,"reconstruction_loss_weight":1.0, + "tcremp_mse_fraction":0.7,"pgen_huber_delta":0.5,"label_smoothing":0.0, + "batch_size":64,"epochs":40,"lr":0.0003,"weight_decay":0.0001, + "gradient_accumulation_steps":1,"max_grad_norm":1.0,"early_stopping_patience":0, + "scheduler_factor":0.5,"scheduler_patience":2,"scheduler_min_lr":1e-6, + "seed":42,"num_workers":8,"amp":True,"skip_generation_metrics":False, + "save_test_predictions":False,"val_generation_every":0,"max_train_batches":0, + "max_eval_batches":0, + } + mismatches = { + key: {"expected": value, "actual": baseline["training"].get(key)} + for key, value in expected_training.items() if baseline["training"].get(key) != value + } + if mismatches: + raise ValueError(f"Selected baseline differs from locked settings: {mismatches}") + + standardizer_tensors = standardizer.as_torch(torch.device("cpu")) + candidate_reports = {} + for kind in CANDIDATES: + directory = root / kind + bundle_path = directory / "anchored_tokenizer.json" + config = json.loads((directory / "tokenizer_config.json").read_text(encoding="utf-8")) + validation = json.loads((directory / "validation_report.json").read_text(encoding="utf-8")) + if config["bundle_sha256"] != sha256(bundle_path): + raise ValueError(f"{kind} bundle checksum changed.") + if config["inference_modality"] != "sequence_only" or config["accepts_v_call_or_j_call"]: + raise ValueError(f"{kind} violates sequence-only inference.") + if not all(validation["checks"].values()): + raise ValueError(f"{kind} preparation checks are not all accepting.") + runtime = AnchoredTokenizer(bundle_path) + roundtrip_rows = 0 + max_encoded = 0 + for split, indices in split_indices.items(): + for sequence in table.iloc[indices]["junction_aa"].astype(str): + encoded = runtime.encode_with_boundaries(sequence, 40) + if runtime.decode(encoded.ids) != sequence: + raise ValueError(f"{kind} roundtrip failed in {split} for {sequence!r}.") + if encoded.n_anchor != sequence[:len(encoded.n_anchor)] or encoded.c_anchor != sequence[-len(encoded.c_anchor):]: + raise ValueError(f"{kind} protected token does not match a literal edge.") + if encoded.boundary_types[0] != "N_ANCHOR" or encoded.boundary_types[-1] != "C_ANCHOR": + raise ValueError(f"{kind} boundary typing failed.") + max_encoded = max(max_encoded, len(encoded.ids)) + roundtrip_rows += 1 + if max_encoded > 40: + raise ValueError(f"{kind} exceeds max input length.") + + tokenizer = resolve_encoder_tokenizer(kind, str(bundle_path)) + dataset = MultiTaskBenchmarkDataset( + table, embeddings, split_indices["train"][:2], tokenizer, "log10_pgen_1mm", 40 + ) + batch = collate_multitask([dataset[0], dataset[1]]) + model_payload = dict(baseline["model"]) + model_payload["input_vocab_size"] = tokenizer.vocab_size + model_payload["output_vocab_size"] = len(AA_VOCAB) + model_payload["share_input_output_embeddings"] = False + model_config = IRRMCodecConfig(**model_payload) + model_config.validate() + changed_model_fields = sorted( + key for key, value in asdict(model_config).items() if value != baseline["model"][key] + ) + if changed_model_fields != ["input_vocab_size", "share_input_output_embeddings"]: + raise ValueError(f"{kind} changed unexpected model fields: {changed_model_fields}") + torch.manual_seed(42) + model = IRRMCodecTransformer(model_config) + outputs = model(batch["encoder_tokens"], batch["encoder_mask"], batch["decoder_input"]) + losses = IRRMCodecMultiTaskLoss(weights=MultiTaskLossWeights(1.0, 1.0, 1.0))( + outputs, + tcremp_target=batch["tcremp_target"], + pgen_target=batch["pgen_target"], + reconstruction_target=batch["reconstruction_target"], + **standardizer_tensors, + ) + losses["loss"].backward() + if not torch.isfinite(losses["loss"]): + raise FloatingPointError(f"{kind} finite smoke gate failed.") + candidate_reports[kind] = { + "tokenizer_type": kind, + "bundle_path": str(bundle_path), + "bundle_sha256": sha256(bundle_path), + "input_vocab_size": tokenizer.vocab_size, + "roundtrip_rows": roundtrip_rows, + "max_encoded_length": max_encoded, + "parameter_count": sum(parameter.numel() for parameter in model.parameters()), + "parameter_delta_vs_char": sum(parameter.numel() for parameter in model.parameters()) - baseline["parameter_count"], + "changed_model_fields_vs_char": changed_model_fields, + "changed_training_fields_vs_char": ["tokenizer_type", "tokenizer_path"], + "finite_full_multitask_smoke_loss": True, + "smoke_loss": float(losses["loss"].detach()), + } + del runtime, tokenizer, dataset, batch, model, outputs, losses + gc.collect() + + report = { + "status": "ready", + "selection_rule": { + "primary": "minimum best-checkpoint full validation objective", + "secondary": ["validation pgen RMSE", "validation reconstruction", "validation TCRemP"], + "test_metrics_used_for_selection": False, + }, + "study_constraints": { + "sequence_only_inference": True, + "v_j_calls_are_model_inputs": False, + "char_decoder_vocab_size": len(AA_VOCAB), + "configs_differ_only_in_tokenizer_and_implied_input_embedding": True, + }, + "baseline": { + "job_id": "1425949", + "run_config_sha256": sha256(baseline_path), + "parameter_count": baseline["parameter_count"], + }, + "benchmark": { + "rows": len(table), + "split_rows": {split: len(value) for split, value in split_indices.items()}, + "normalizer_fit_split": "train", + "normalizer_train_rows": standardizer.train_rows, + }, + "candidates": candidate_reports, + "runtime": {"python": platform.python_version(), "torch": torch.__version__}, + } + output = Path(args.output_path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/scripts/preflight_data_anchor_objectives.py b/scripts/preflight_data_anchor_objectives.py new file mode 100644 index 0000000..4266c31 --- /dev/null +++ b/scripts/preflight_data_anchor_objectives.py @@ -0,0 +1,228 @@ +"""CPU gate for matched DATA-ANCHOR objective-ablation training runs.""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import os +from dataclasses import asdict +from pathlib import Path + +import torch + +from irrm_codec.multitask_data import ( + MultiTaskBenchmarkDataset, + TargetStandardizer, + collate_multitask, + load_prepared_benchmark, + resolve_encoder_tokenizer, + select_split_indices, +) +from irrm_codec.multitask_losses import IRRMCodecMultiTaskLoss, MultiTaskLossWeights +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer + + +CONDITIONS = { + "r": {"reconstruction": 1.0, "tcremp": 0.0, "pgen": 0.0}, + "t": {"reconstruction": 0.0, "tcremp": 1.0, "pgen": 0.0}, + "p": {"reconstruction": 0.0, "tcremp": 0.0, "pgen": 1.0}, + "rt": {"reconstruction": 1.0, "tcremp": 1.0, "pgen": 0.0}, + "rp": {"reconstruction": 1.0, "tcremp": 0.0, "pgen": 1.0}, + "tp": {"reconstruction": 0.0, "tcremp": 1.0, "pgen": 1.0}, +} + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def stable_json(path: Path, value: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--reference-run-config", type=Path, required=True) + parser.add_argument("--output-root", type=Path, required=True) + args = parser.parse_args() + + ready = json.loads((args.data_dir / "READY.json").read_text(encoding="utf-8")) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Locked benchmark READY.json is not accepting.") + reference = json.loads(args.reference_run_config.read_text(encoding="utf-8")) + training = reference["training"] + required = { + "data_dir": str(args.data_dir), + "train_subset": "all", + "pgen_target": "log10_pgen_1mm", + "tokenizer_type": "data_anchor", + "tokenizer_path": str(args.tokenizer), + "max_sequence_len": 40, + "d_model": 320, + "latent_dim": 128, + "nhead": 8, + "encoder_layers": 4, + "decoder_layers": 4, + "ff_dim": 1280, + "dropout": 0.1, + "tcremp_head_dim": 1024, + "pgen_head_dim": 256, + "decoder_memory_tokens": 4, + "tcremp_loss_weight": 1.0, + "pgen_loss_weight": 1.0, + "reconstruction_loss_weight": 1.0, + "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, + "label_smoothing": 0.0, + "batch_size": 64, + "epochs": 40, + "lr": 0.0003, + "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, + "max_grad_norm": 1.0, + "early_stopping_patience": 0, + "scheduler_factor": 0.5, + "scheduler_patience": 2, + "scheduler_min_lr": 1e-6, + "seed": 42, + "num_workers": 8, + "amp": True, + "skip_generation_metrics": False, + "save_test_predictions": False, + "val_generation_every": 0, + "max_train_batches": 0, + "max_eval_batches": 0, + } + mismatches = { + key: {"expected": value, "actual": training.get(key)} + for key, value in required.items() if training.get(key) != value + } + if mismatches: + raise ValueError(f"Completed DATA-ANCHOR reference is not the locked run: {mismatches}") + + table, embeddings = load_prepared_benchmark(args.data_dir) + indices = { + split: select_split_indices(table, args.data_dir, split) + for split in ("train", "val", "test") + } + normalizer = TargetStandardizer.load(args.data_dir / "target_standardizer.npz") + if normalizer.train_rows != len(indices["train"]) or normalizer.pgen_target != "log10_pgen_1mm": + raise ValueError("Train-only normalizer does not match the locked train manifest.") + tokenizer = resolve_encoder_tokenizer("data_anchor", str(args.tokenizer)) + config = IRRMCodecConfig(**reference["model"]) + if tokenizer.vocab_size != config.input_vocab_size: + raise ValueError("Tokenizer vocabulary does not match reference input embedding.") + dataset = MultiTaskBenchmarkDataset( + table, embeddings, indices["train"][:2], tokenizer, "log10_pgen_1mm", 40 + ) + batch = collate_multitask([dataset[0], dataset[1]]) + standardized = normalizer.as_torch(torch.device("cpu")) + + reports = {} + configs_dir = args.output_root / "configs" + for name, weights in CONDITIONS.items(): + torch.manual_seed(42) + model = IRRMCodecTransformer(config) + outputs = model(batch["encoder_tokens"], batch["encoder_mask"], batch["decoder_input"]) + criterion = IRRMCodecMultiTaskLoss( + weights=MultiTaskLossWeights( + tcremp=weights["tcremp"], + pgen=weights["pgen"], + reconstruction=weights["reconstruction"], + ) + ) + losses = criterion( + outputs, + tcremp_target=batch["tcremp_target"], + pgen_target=batch["pgen_target"], + reconstruction_target=batch["reconstruction_target"], + **standardized, + ) + losses["loss"].backward() + encoder_grads = [ + parameter.grad for parameter in model.encoder.parameters() if parameter.grad is not None + ] + if not torch.isfinite(losses["loss"]) or not encoder_grads or not all( + torch.isfinite(gradient).all() for gradient in encoder_grads + ): + raise FloatingPointError(f"{name} has non-finite loss or encoder gradient.") + condition_config = { + "condition": name, + "experiment_id": f"data-anchor-{name}-seed42", + "loss_weights": weights, + "model": asdict(config), + "training_reference": str(args.reference_run_config), + "training_reference_sha256": sha256(args.reference_run_config), + "benchmark_ready_sha256": sha256(args.data_dir / "READY.json"), + "normalizer_sha256": sha256(args.data_dir / "target_standardizer.npz"), + "tokenizer_sha256": sha256(args.tokenizer), + "seed": 42, + "epochs": 40, + "batch_size": 64, + "optimizer": "AdamW", + "learning_rate": 0.0003, + "inactive_head_metrics_interpretable": False, + } + config_path = configs_dir / f"{name}.json" + stable_json(config_path, condition_config) + reports[name] = { + "config_path": str(config_path), + "config_sha256": sha256(config_path), + "smoke_loss": float(losses["loss"].detach()), + "finite_encoder_gradient": True, + "parameter_count": sum(parameter.numel() for parameter in model.parameters()), + "architecture_equal_to_full_reference": asdict(config) == reference["model"], + } + del model, outputs, losses, criterion + gc.collect() + + report = { + "status": "ready", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "required_conditions": ["r", "t", "p", "rt", "rp", "tp"], + "newly_required_conditions": ["t", "rt", "tp"], + "fixed_reference": { + "condition": "rtp", + "job_id": "1426222", + "run_config": str(args.reference_run_config), + "run_config_sha256": sha256(args.reference_run_config), + "parameter_count": reference["parameter_count"], + }, + "checks": { + "locked_benchmark": True, + "locked_train_only_normalizer": True, + "same_data_splits": True, + "same_data_anchor_tokenizer": True, + "same_model_architecture": True, + "same_seed_optimizer_budget": True, + "only_loss_weights_differ": True, + "finite_smoke_losses_and_encoder_gradients": True, + "output_directories_isolated": True, + }, + "hashes": { + "benchmark_ready": sha256(args.data_dir / "READY.json"), + "normalizer": sha256(args.data_dir / "target_standardizer.npz"), + "tokenizer": sha256(args.tokenizer), + }, + "splits": {split: len(value) for split, value in indices.items()}, + "conditions": reports, + "note": ( + "Every condition retains the identical decoder and all auxiliary heads. Metrics from " + "heads with zero loss weight are non-interpretable." + ), + } + stable_json(args.output_root / "PREFLIGHT.json", report) + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/scripts/preflight_phase1_rq1.py b/scripts/preflight_phase1_rq1.py new file mode 100644 index 0000000..ba5a3c4 --- /dev/null +++ b/scripts/preflight_phase1_rq1.py @@ -0,0 +1,132 @@ +"""Validate the locked TRB benchmark and Phase-1 runtime on a Slurm CPU node.""" + +import argparse +import hashlib +import json +import platform +import subprocess +from pathlib import Path + +import numpy as np +import torch + +from irrm_codec.multitask_data import ( + TargetStandardizer, + load_prepared_benchmark, + select_split_indices, +) +from irrm_codec.multitask_losses import IRRMCodecMultiTaskLoss +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import AA_VOCAB + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", required=True) + parser.add_argument("--output-path", required=True) + args = parser.parse_args() + + data_dir = Path(args.data_dir).resolve() + ready = json.loads((data_dir / "READY.json").read_text(encoding="utf-8")) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Locked benchmark READY checks are not all true.") + + table, embeddings = load_prepared_benchmark(data_dir) + split_indices = { + split: select_split_indices(table, data_dir, split) + for split in ("train", "val", "test") + } + standardizer = TargetStandardizer.load(data_dir / "target_standardizer.npz") + if standardizer.pgen_target != "log10_pgen_1mm": + raise ValueError("Unexpected locked pgen target.") + if standardizer.train_rows != len(split_indices["train"]): + raise ValueError("Locked standardizer was not fit on the full train split.") + if len(standardizer.tcremp_mean) != embeddings.shape[1]: + raise ValueError("Locked TCRemP standardizer dimension mismatch.") + + torch.manual_seed(42) + model_config = IRRMCodecConfig( + input_vocab_size=len(AA_VOCAB), + output_vocab_size=len(AA_VOCAB), + max_sequence_len=8, + d_model=32, + latent_dim=24, + nhead=4, + encoder_layers=1, + decoder_layers=1, + ff_dim=64, + dropout=0.0, + tcremp_dim=12, + tcremp_head_dim=24, + pgen_head_dim=12, + decoder_memory_tokens=2, + ) + model = IRRMCodecTransformer(model_config) + tokens = torch.randint(5, len(AA_VOCAB), (2, 6)) + decoder_input = torch.randint(5, len(AA_VOCAB), (2, 7)) + outputs = model(tokens, tokens.ne(0), decoder_input) + losses = IRRMCodecMultiTaskLoss()( + outputs, + tcremp_target=torch.randn(2, 12), + pgen_target=torch.randn(2), + reconstruction_target=torch.randint(5, len(AA_VOCAB), (2, 7)), + tcremp_mean=torch.zeros(12), + tcremp_std=torch.ones(12), + pgen_mean=torch.tensor(0.0), + pgen_std=torch.tensor(1.0), + ) + losses["loss"].backward() + if not torch.isfinite(losses["loss"]): + raise FloatingPointError("Runtime smoke-test loss is non-finite.") + + artifact_paths = [ + data_dir / "READY.json", + data_dir / "dataset.parquet", + data_dir / "embeddings.npy", + data_dir / "target_standardizer.npz", + *(data_dir / "manifests" / f"{name}.tsv" for name in ("train", "val", "test")), + ] + revision = subprocess.check_output( + ["git", "rev-parse", "HEAD"], text=True + ).strip() + report = { + "status": "ready", + "git_revision": revision, + "benchmark_path": str(data_dir), + "benchmark_ready_summary": ready, + "prepared_artifact_sha256": { + str(path.relative_to(data_dir)): sha256(path) for path in artifact_paths + }, + "dataset_rows": len(table), + "embedding_shape": list(embeddings.shape), + "embedding_dtype": str(embeddings.dtype), + "split_rows": {name: len(values) for name, values in split_indices.items()}, + "standardizer": { + "pgen_target": standardizer.pgen_target, + "train_rows": standardizer.train_rows, + "tcremp_dim": len(standardizer.tcremp_mean), + }, + "runtime": { + "python": platform.python_version(), + "torch": torch.__version__, + "torch_cuda_build": torch.version.cuda, + "smoke_loss_finite": True, + }, + } + output_path = Path(args.output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + temporary = output_path.with_suffix(output_path.suffix + ".tmp") + temporary.write_text(json.dumps(report, indent=2), encoding="utf-8") + temporary.replace(output_path) + + +if __name__ == "__main__": + main() diff --git a/scripts/preflight_rq5_depths.py b/scripts/preflight_rq5_depths.py new file mode 100644 index 0000000..285ea8e --- /dev/null +++ b/scripts/preflight_rq5_depths.py @@ -0,0 +1,164 @@ +"""Validate RQ5 Stage-2 encoder-depth candidates against the selected d128 run.""" + +import argparse +import gc +import hashlib +import json +import platform +from dataclasses import asdict +from pathlib import Path + +import torch + +from irrm_codec.multitask_data import TargetStandardizer, load_prepared_benchmark, select_split_indices +from irrm_codec.multitask_losses import IRRMCodecMultiTaskLoss, MultiTaskLossWeights +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import AA_VOCAB + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--baseline-run-config", required=True) + parser.add_argument("--d128-test-metrics", required=True) + parser.add_argument("--d320-test-metrics", required=True) + parser.add_argument("--d512-test-metrics", required=True) + parser.add_argument("--data-dir", required=True) + parser.add_argument("--output-path", required=True) + args = parser.parse_args() + + baseline_path = Path(args.baseline_run_config) + baseline = json.loads(baseline_path.read_text()) + width_metrics = { + "128": json.loads(Path(args.d128_test_metrics).read_text()), + "320": json.loads(Path(args.d320_test_metrics).read_text()), + "512": json.loads(Path(args.d512_test_metrics).read_text()), + } + validation_losses = { + width: metrics["best_checkpoint_val_loss"] for width, metrics in width_metrics.items() + } + if min(validation_losses, key=validation_losses.get) != "128": + raise ValueError(f"latent_dim=128 was not selected by validation: {validation_losses}") + + data_dir = Path(args.data_dir) + ready = json.loads((data_dir / "READY.json").read_text()) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Locked benchmark READY checks failed.") + table, embeddings = load_prepared_benchmark(data_dir) + splits = {name: select_split_indices(table, data_dir, name) for name in ("train", "val", "test")} + standardizer = TargetStandardizer.load(data_dir / "target_standardizer.npz") + if standardizer.train_rows != len(splits["train"]) or standardizer.pgen_target != "log10_pgen_1mm": + raise ValueError("Locked train-only normalizer does not match the selected d128 run.") + + expected_training = { + "train_subset": "all", "pgen_target": "log10_pgen_1mm", "tokenizer_type": "char", + "tokenizer_path": None, "max_sequence_len": 40, "d_model": 320, "latent_dim": 128, + "nhead": 8, "encoder_layers": 4, "decoder_layers": 4, "ff_dim": 1280, + "dropout": 0.1, "tcremp_head_dim": 1024, "pgen_head_dim": 256, + "decoder_memory_tokens": 4, "tcremp_loss_weight": 1.0, "pgen_loss_weight": 1.0, + "reconstruction_loss_weight": 1.0, "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, "label_smoothing": 0.0, "batch_size": 64, + "epochs": 40, "lr": 0.0003, "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, "max_grad_norm": 1.0, + "early_stopping_patience": 0, "scheduler_factor": 0.5, "scheduler_patience": 2, + "scheduler_min_lr": 1e-6, "seed": 42, "num_workers": 8, "amp": True, + "skip_generation_metrics": False, "save_test_predictions": False, + "val_generation_every": 0, "max_train_batches": 0, "max_eval_batches": 0, + } + mismatches = { + key: {"expected": expected, "actual": baseline["training"].get(key)} + for key, expected in expected_training.items() + if baseline["training"].get(key) != expected + } + if mismatches: + raise ValueError(f"Selected d128 run differs from locked RQ5 settings: {mismatches}") + if baseline["model"]["latent_dim"] != 128 or baseline["model"]["encoder_layers"] != 4: + raise ValueError("Selected baseline is not the completed d128, 4+4-layer run.") + if baseline["parameter_count"] != width_metrics["128"]["parameter_count"]: + raise ValueError("Selected d128 run/metric parameter counts disagree.") + + results = {} + for depth in (2, 6): + torch.manual_seed(42) + model_payload = dict(baseline["model"]) + model_payload["encoder_layers"] = depth + config = IRRMCodecConfig(**model_payload) + config.validate() + changed_model_fields = sorted( + key for key, value in asdict(config).items() if value != baseline["model"][key] + ) + if changed_model_fields != ["encoder_layers"]: + raise ValueError(f"encoder-depth {depth} changed unexpected model fields: {changed_model_fields}") + + model = IRRMCodecTransformer(config) + parameter_count = sum(parameter.numel() for parameter in model.parameters()) + tokens = torch.randint(5, len(AA_VOCAB), (2, 12)) + decoder_input = torch.randint(5, len(AA_VOCAB), (2, 13)) + outputs = model(tokens, tokens.ne(0), decoder_input) + losses = IRRMCodecMultiTaskLoss(weights=MultiTaskLossWeights(1.0, 1.0, 1.0))( + outputs, + tcremp_target=torch.randn(2, embeddings.shape[1]), + pgen_target=torch.randn(2), + reconstruction_target=torch.randint(5, len(AA_VOCAB), (2, 13)), + tcremp_mean=torch.zeros(embeddings.shape[1]), + tcremp_std=torch.ones(embeddings.shape[1]), + pgen_mean=torch.tensor(0.0), + pgen_std=torch.tensor(1.0), + ) + losses["loss"].backward() + finite = bool(torch.isfinite(losses["loss"])) + if not finite: + raise FloatingPointError(f"encoder-depth {depth} smoke loss is non-finite.") + results[str(depth)] = { + "encoder_layers": depth, + "decoder_layers": config.decoder_layers, + "latent_dim": config.latent_dim, + "parameter_count": parameter_count, + "parameter_delta_vs_e4": parameter_count - baseline["parameter_count"], + "changed_model_fields_vs_e4": changed_model_fields, + "changed_training_fields_vs_e4": [], + "finite_full_multitask_smoke_loss": finite, + "smoke_loss": float(losses["loss"].detach()), + "model_config": asdict(config), + } + del model, outputs, losses + gc.collect() + + report = { + "status": "ready", + "interpretation": "RQ5 Stage 2 varies encoder_layers only; latent_dim=128 and decoder_layers=4 remain fixed.", + "selection": { + "criterion": "minimum best-checkpoint validation loss only", + "validation_losses": validation_losses, + "selected_latent_dim": 128, + "test_metrics_used_for_selection": False, + }, + "baseline": { + "job_id": "1425949", "run_config_path": str(baseline_path.resolve()), + "run_config_sha256": sha256(baseline_path), "latent_dim": 128, + "encoder_layers": 4, "decoder_layers": 4, + "parameter_count": baseline["parameter_count"], "non_depth_training_fields_match": True, + }, + "benchmark": { + "path": str(data_dir.resolve()), "ready_sha256": sha256(data_dir / "READY.json"), + "rows": len(table), "embedding_shape": list(embeddings.shape), + "split_rows": {key: len(value) for key, value in splits.items()}, + "normalizer_fit_split": "train", "normalizer_train_rows": standardizer.train_rows, + }, + "candidates": results, + "runtime": {"python": platform.python_version(), "torch": torch.__version__}, + } + output = Path(args.output_path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/scripts/preflight_rq5_widths.py b/scripts/preflight_rq5_widths.py new file mode 100644 index 0000000..f6a5b75 --- /dev/null +++ b/scripts/preflight_rq5_widths.py @@ -0,0 +1,140 @@ +"""Validate RQ5 latent-width candidates against the completed d320 baseline.""" + +import argparse +import hashlib +import json +import platform +from dataclasses import asdict +from pathlib import Path + +import torch + +from irrm_codec.multitask_data import TargetStandardizer, load_prepared_benchmark, select_split_indices +from irrm_codec.multitask_losses import IRRMCodecMultiTaskLoss, MultiTaskLossWeights +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import AA_VOCAB + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--baseline-run-config", required=True) + parser.add_argument("--baseline-test-metrics", required=True) + parser.add_argument("--data-dir", required=True) + parser.add_argument("--output-path", required=True) + args = parser.parse_args() + + baseline_path = Path(args.baseline_run_config) + baseline = json.loads(baseline_path.read_text()) + baseline_metrics = json.loads(Path(args.baseline_test_metrics).read_text()) + data_dir = Path(args.data_dir) + ready = json.loads((data_dir / "READY.json").read_text()) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Locked benchmark READY checks failed.") + table, embeddings = load_prepared_benchmark(data_dir) + splits = {name: select_split_indices(table, data_dir, name) for name in ("train", "val", "test")} + standardizer = TargetStandardizer.load(data_dir / "target_standardizer.npz") + if standardizer.train_rows != len(splits["train"]) or standardizer.pgen_target != "log10_pgen_1mm": + raise ValueError("Locked train-only normalizer does not match baseline.") + + expected_training = { + "train_subset": "all", "pgen_target": "log10_pgen_1mm", "tokenizer_type": "char", + "tokenizer_path": None, "max_sequence_len": 40, "d_model": 320, "nhead": 8, + "encoder_layers": 4, "decoder_layers": 4, "ff_dim": 1280, "dropout": 0.1, + "tcremp_head_dim": 1024, "pgen_head_dim": 256, "decoder_memory_tokens": 4, + "tcremp_loss_weight": 1.0, "pgen_loss_weight": 1.0, "reconstruction_loss_weight": 1.0, + "tcremp_mse_fraction": 0.7, "pgen_huber_delta": 0.5, "label_smoothing": 0.0, + "batch_size": 64, "epochs": 40, "lr": 0.0003, "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, "max_grad_norm": 1.0, + "early_stopping_patience": 0, "scheduler_factor": 0.5, "scheduler_patience": 2, + "scheduler_min_lr": 1e-6, "seed": 42, "num_workers": 8, "amp": True, + "skip_generation_metrics": False, "save_test_predictions": False, + "val_generation_every": 0, "max_train_batches": 0, "max_eval_batches": 0, + } + mismatches = { + key: {"expected": expected, "actual": baseline["training"].get(key)} + for key, expected in expected_training.items() + if baseline["training"].get(key) != expected + } + if mismatches: + raise ValueError(f"Completed d320 baseline differs from locked RQ5 settings: {mismatches}") + if baseline["model"]["latent_dim"] != 320 or baseline["parameter_count"] != baseline_metrics["parameter_count"]: + raise ValueError("d320 baseline model/metric metadata is inconsistent.") + + results = {} + torch.manual_seed(42) + for width in (128, 512): + model_payload = dict(baseline["model"]) + model_payload["latent_dim"] = width + config = IRRMCodecConfig(**model_payload) + config.validate() + changed_model_fields = sorted( + key for key, value in asdict(config).items() if value != baseline["model"][key] + ) + if changed_model_fields != ["latent_dim"]: + raise ValueError(f"d{width} changed unexpected model fields: {changed_model_fields}") + model = IRRMCodecTransformer(config) + parameter_count = sum(parameter.numel() for parameter in model.parameters()) + tokens = torch.randint(5, len(AA_VOCAB), (2, 12)) + decoder_input = torch.randint(5, len(AA_VOCAB), (2, 13)) + outputs = model(tokens, tokens.ne(0), decoder_input) + losses = IRRMCodecMultiTaskLoss(weights=MultiTaskLossWeights(1.0, 1.0, 1.0))( + outputs, + tcremp_target=torch.randn(2, embeddings.shape[1]), + pgen_target=torch.randn(2), + reconstruction_target=torch.randint(5, len(AA_VOCAB), (2, 13)), + tcremp_mean=torch.zeros(embeddings.shape[1]), + tcremp_std=torch.ones(embeddings.shape[1]), + pgen_mean=torch.tensor(0.0), + pgen_std=torch.tensor(1.0), + ) + losses["loss"].backward() + finite = bool(torch.isfinite(losses["loss"])) + if not finite: + raise FloatingPointError(f"d{width} smoke loss is non-finite.") + results[str(width)] = { + "latent_dim": width, + "parameter_count": parameter_count, + "parameter_delta_vs_d320": parameter_count - baseline["parameter_count"], + "changed_model_fields_vs_d320": changed_model_fields, + "changed_training_fields_vs_d320": [], + "finite_full_multitask_smoke_loss": finite, + "smoke_loss": float(losses["loss"].detach()), + "model_config": asdict(config), + } + del model, outputs, losses + + report = { + "status": "ready", + "interpretation": "RQ5 width is latent_dim only; d_model=320 and nhead=8 remain fixed.", + "baseline": { + "job_id": "1425897", + "run_config_path": str(baseline_path.resolve()), + "run_config_sha256": sha256(baseline_path), + "latent_dim": 320, + "parameter_count": baseline["parameter_count"], + "non_width_training_fields_match": True, + }, + "benchmark": { + "path": str(data_dir.resolve()), "ready_sha256": sha256(data_dir / "READY.json"), + "rows": len(table), "embedding_shape": list(embeddings.shape), + "split_rows": {key: len(value) for key, value in splits.items()}, + "normalizer_fit_split": "train", "normalizer_train_rows": standardizer.train_rows, + }, + "candidates": results, + "runtime": {"python": platform.python_version(), "torch": torch.__version__}, + } + output = Path(args.output_path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/scripts/preflight_vdjdb_evaluation.py b/scripts/preflight_vdjdb_evaluation.py new file mode 100644 index 0000000..8d4b96c --- /dev/null +++ b/scripts/preflight_vdjdb_evaluation.py @@ -0,0 +1,80 @@ +"""Small CPU-only functional gate for frozen VDJdb evaluation statistics.""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import numpy as np +import pandas as pd + +from benchmark.evaluate_vdjdb_frozen import ( + assert_finite_tree, + bootstrap_primary, + build_pair_manifest, + controls, + distance_matrix, + pair_metrics, + retrieval_metrics, + sha256, +) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--cohort", type=Path, required=True) + parser.add_argument("--cohort-preflight", type=Path, required=True) + parser.add_argument("--evaluation-script", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + preflight = json.loads(args.cohort_preflight.read_text(encoding="utf-8")) + if sha256(args.cohort) != preflight["cohort"]["cohort_sha256"]: + raise ValueError("Cohort hash changed after preparation.") + full = pd.read_csv(args.cohort, sep="\t") + chosen = [] + for _, group in full.groupby("label", sort=True): + if group["donor_key"].nunique() >= 2: + chosen.append(group.head(20)) + if len(chosen) == 4: + break + cohort = pd.concat(chosen, ignore_index=True) + pairs = build_pair_manifest(cohort, limit=50, seed=42) + rng = np.random.default_rng(42) + arrays = {name: rng.normal(size=(len(cohort), 16)) for name in ("r", "p", "rp", "rtp")} + matrices = {name: distance_matrix(value, "cosine") for name, value in arrays.items()} + retrieval = {} + report = {} + for name, matrix in matrices.items(): + pair_report, _ = pair_metrics(matrix, pairs) + retrieval_report, retrieval[name] = retrieval_metrics(matrix, cohort) + report[name] = {"pairs": pair_report, "retrieval": retrieval_report} + report["bootstrap"] = bootstrap_primary(matrices, pairs, retrieval, replicates=20, seed=42) + report["controls"] = controls(cohort, pairs, matrices, permutations=20, seed=42) + assert_finite_tree(report) + payload = { + "status": "ready", + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "cohort_sha256": sha256(args.cohort), + "evaluation_script_sha256": sha256(args.evaluation_script), + "smoke_rows": len(cohort), + "smoke_pairs": len(pairs), + "checks": { + "pair_contrast": True, + "pairwise_auroc_auprc": True, + "donor_excluded_retrieval": True, + "cosine_distance": True, + "euclidean_distance": bool(np.isfinite(distance_matrix(arrays["r"], "euclidean")).all()), + "donor_cluster_bootstrap": True, + "length_composition_and_shuffle_controls": True, + "finite_results": True, + }, + } + args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(payload, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/scripts/preflight_wordpiece_rq4.py b/scripts/preflight_wordpiece_rq4.py new file mode 100644 index 0000000..1cc30e8 --- /dev/null +++ b/scripts/preflight_wordpiece_rq4.py @@ -0,0 +1,140 @@ +"""CPU preflight for the locked RQ4 WordPiece benchmark runs.""" + +import argparse +import hashlib +import json +import platform +from pathlib import Path + +import torch + +from irrm_codec.multitask_data import ( + MultiTaskBenchmarkDataset, + TargetStandardizer, + collate_multitask, + load_prepared_benchmark, + resolve_encoder_tokenizer, + select_split_indices, +) +from irrm_codec.multitask_losses import IRRMCodecMultiTaskLoss, MultiTaskLossWeights +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import AA_VOCAB + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", required=True) + parser.add_argument("--tokenizer-dir", required=True) + parser.add_argument("--output-path", required=True) + args = parser.parse_args() + + data_dir = Path(args.data_dir) + tokenizer_dir = Path(args.tokenizer_dir) + tokenizer_path = tokenizer_dir / "tokenizer.json" + config = json.loads((tokenizer_dir / "tokenizer_config.json").read_text()) + validation = json.loads((tokenizer_dir / "validation_report.json").read_text()) + if config["training_split"] != "train" or not validation["train_only_fit"]: + raise ValueError("Tokenizer provenance does not prove train-only fitting.") + if not all(validation["checks"].values()): + raise ValueError("Tokenizer validation checks are not all true.") + if sha256(tokenizer_path) != config["tokenizer_sha256"]: + raise ValueError("Tokenizer checksum does not match immutable config.") + train_manifest = data_dir / "manifests" / "train.tsv" + if sha256(train_manifest) != config["train_manifest_sha256"]: + raise ValueError("Locked train-manifest checksum changed.") + + ready = json.loads((data_dir / "READY.json").read_text()) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Locked benchmark is not ready.") + table, embeddings = load_prepared_benchmark(data_dir) + train_indices = select_split_indices(table, data_dir, "train") + val_indices = select_split_indices(table, data_dir, "val") + test_indices = select_split_indices(table, data_dir, "test") + standardizer = TargetStandardizer.load(data_dir / "target_standardizer.npz") + if standardizer.train_rows != len(train_indices) or standardizer.pgen_target != "log10_pgen_1mm": + raise ValueError("Locked train-only target standardizer does not match RQ4.") + if len(standardizer.tcremp_mean) != embeddings.shape[1]: + raise ValueError("Locked TCRemP standardizer has the wrong dimension.") + + tokenizer = resolve_encoder_tokenizer("wordpiece", str(tokenizer_path)) + dataset = MultiTaskBenchmarkDataset( + table, embeddings, train_indices[:2], tokenizer, "log10_pgen_1mm", 40 + ) + batch = collate_multitask([dataset[0], dataset[1]]) + torch.manual_seed(42) + model_config = IRRMCodecConfig( + input_vocab_size=tokenizer.vocab_size, + output_vocab_size=len(AA_VOCAB), + share_input_output_embeddings=False, + max_sequence_len=40, + d_model=32, + latent_dim=24, + nhead=4, + encoder_layers=1, + decoder_layers=1, + ff_dim=64, + dropout=0.0, + tcremp_dim=embeddings.shape[1], + tcremp_head_dim=24, + pgen_head_dim=12, + decoder_memory_tokens=2, + ) + model = IRRMCodecTransformer(model_config) + outputs = model(batch["encoder_tokens"], batch["encoder_mask"], batch["decoder_input"]) + standardizer_tensors = standardizer.as_torch(torch.device("cpu")) + losses = IRRMCodecMultiTaskLoss(weights=MultiTaskLossWeights(1.0, 1.0, 1.0))( + outputs, + tcremp_target=batch["tcremp_target"], + pgen_target=batch["pgen_target"], + reconstruction_target=batch["reconstruction_target"], + **standardizer_tensors, + ) + losses["loss"].backward() + if not torch.isfinite(losses["loss"]): + raise FloatingPointError("WordPiece smoke loss is non-finite.") + + report = { + "status": "ready", + "tokenizer": { + "path": str(tokenizer_path.resolve()), + "sha256": sha256(tokenizer_path), + "vocab_size": tokenizer.vocab_size, + "special_token_ids": config["special_token_ids"], + "train_manifest_sha256": config["train_manifest_sha256"], + "train_only_fit": True, + }, + "benchmark": { + "path": str(data_dir.resolve()), + "ready_sha256": sha256(data_dir / "READY.json"), + "rows": len(table), + "embedding_shape": list(embeddings.shape), + "split_rows": { + "train": len(train_indices), "val": len(val_indices), "test": len(test_indices) + }, + "normalizer_fit_split": "train", + "normalizer_train_rows": standardizer.train_rows, + }, + "model_smoke": { + "input_vocab_size": tokenizer.vocab_size, + "output_vocab_size": len(AA_VOCAB), + "shared_input_output_embeddings": False, + "finite_loss": True, + "loss": float(losses["loss"].detach()), + }, + "runtime": {"python": platform.python_version(), "torch": torch.__version__}, + } + output_path = Path(args.output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(report, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/scripts/preflight_wordpiece_sweep.py b/scripts/preflight_wordpiece_sweep.py new file mode 100644 index 0000000..9c106a2 --- /dev/null +++ b/scripts/preflight_wordpiece_sweep.py @@ -0,0 +1,189 @@ +"""CPU gate for the leakage-safe TRB WordPiece vocabulary sweep.""" + +import argparse +import gc +import hashlib +import json +import platform +from dataclasses import asdict +from pathlib import Path + +import torch + +from irrm_codec.multitask_data import ( + MultiTaskBenchmarkDataset, + TargetStandardizer, + collate_multitask, + load_prepared_benchmark, + resolve_encoder_tokenizer, + select_split_indices, +) +from irrm_codec.multitask_losses import IRRMCodecMultiTaskLoss, MultiTaskLossWeights +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import AA_VOCAB + + +VOCAB_SIZES = (44, 64, 128, 256, 512, 1024, 2048, 4096) + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", required=True) + parser.add_argument("--tokenizer-root", required=True) + parser.add_argument("--baseline-run-config", required=True) + parser.add_argument("--output-path", required=True) + args = parser.parse_args() + + data_dir = Path(args.data_dir) + tokenizer_root = Path(args.tokenizer_root) + baseline_path = Path(args.baseline_run_config) + baseline = json.loads(baseline_path.read_text()) + if ( + baseline["model"]["latent_dim"] != 128 + or baseline["model"]["encoder_layers"] != 4 + or baseline["model"]["decoder_layers"] != 4 + ): + raise ValueError("Baseline is not the validation-selected char d128 4+4 model.") + + ready = json.loads((data_dir / "READY.json").read_text()) + if ready.get("status") != "ready" or not all(ready.get("checks", {}).values()): + raise ValueError("Locked benchmark is not ready.") + table, embeddings = load_prepared_benchmark(data_dir) + split_indices = { + split: select_split_indices(table, data_dir, split) for split in ("train", "val", "test") + } + standardizer = TargetStandardizer.load(data_dir / "target_standardizer.npz") + if standardizer.train_rows != len(split_indices["train"]): + raise ValueError("Locked standardizer is not train-only for the locked manifest.") + if standardizer.pgen_target != "log10_pgen_1mm": + raise ValueError("Locked standardizer uses the wrong pgen target.") + + expected_training = { + "train_subset": "all", "pgen_target": "log10_pgen_1mm", "tokenizer_type": "char", + "tokenizer_path": None, "max_sequence_len": 40, "d_model": 320, "latent_dim": 128, + "nhead": 8, "encoder_layers": 4, "decoder_layers": 4, "ff_dim": 1280, + "dropout": 0.1, "tcremp_head_dim": 1024, "pgen_head_dim": 256, + "decoder_memory_tokens": 4, "tcremp_loss_weight": 1.0, "pgen_loss_weight": 1.0, + "reconstruction_loss_weight": 1.0, "tcremp_mse_fraction": 0.7, + "pgen_huber_delta": 0.5, "label_smoothing": 0.0, "batch_size": 64, + "epochs": 40, "lr": 0.0003, "weight_decay": 0.0001, + "gradient_accumulation_steps": 1, "max_grad_norm": 1.0, + "early_stopping_patience": 0, "scheduler_factor": 0.5, "scheduler_patience": 2, + "scheduler_min_lr": 1e-6, "seed": 42, "num_workers": 8, "amp": True, + "skip_generation_metrics": False, "save_test_predictions": False, + "val_generation_every": 0, "max_train_batches": 0, "max_eval_batches": 0, + } + mismatches = { + key: {"expected": expected, "actual": baseline["training"].get(key)} + for key, expected in expected_training.items() + if baseline["training"].get(key) != expected + } + if mismatches: + raise ValueError(f"Selected char baseline differs from locked Stage-A settings: {mismatches}") + + train_manifest = data_dir / "manifests" / "train.tsv" + standardizer_tensors = standardizer.as_torch(torch.device("cpu")) + candidates = {} + for vocab_size in VOCAB_SIZES: + token_dir = tokenizer_root / f"v{vocab_size}-seed42" + config = json.loads((token_dir / "tokenizer_config.json").read_text()) + validation = json.loads((token_dir / "validation_report.json").read_text()) + tokenizer_path = token_dir / "tokenizer.json" + if config["actual_vocab_size"] != vocab_size or config["requested_vocab_size"] != vocab_size: + raise ValueError(f"WordPiece v{vocab_size} actual/requested size mismatch.") + if config["training_split"] != "train" or config["training_rows"] != 79_544: + raise ValueError(f"WordPiece v{vocab_size} lacks train-only provenance.") + if config["train_manifest_sha256"] != sha256(train_manifest): + raise ValueError(f"WordPiece v{vocab_size} train manifest checksum changed.") + if config["tokenizer_sha256"] != sha256(tokenizer_path): + raise ValueError(f"WordPiece v{vocab_size} tokenizer checksum changed.") + if not all(validation["checks"].values()): + raise ValueError(f"WordPiece v{vocab_size} validation checks are not all true.") + + tokenizer = resolve_encoder_tokenizer("wordpiece", str(tokenizer_path)) + dataset = MultiTaskBenchmarkDataset( + table, embeddings, split_indices["train"][:2], tokenizer, "log10_pgen_1mm", 40 + ) + batch = collate_multitask([dataset[0], dataset[1]]) + model_payload = dict(baseline["model"]) + model_payload["input_vocab_size"] = vocab_size + model_payload["output_vocab_size"] = len(AA_VOCAB) + model_payload["share_input_output_embeddings"] = False + model_config = IRRMCodecConfig(**model_payload) + model_config.validate() + changed_model_fields = sorted( + key for key, value in asdict(model_config).items() if value != baseline["model"][key] + ) + if changed_model_fields != ["input_vocab_size", "share_input_output_embeddings"]: + raise ValueError(f"WordPiece v{vocab_size} changed unexpected model fields: {changed_model_fields}") + torch.manual_seed(42) + model = IRRMCodecTransformer(model_config) + outputs = model(batch["encoder_tokens"], batch["encoder_mask"], batch["decoder_input"]) + losses = IRRMCodecMultiTaskLoss(weights=MultiTaskLossWeights(1.0, 1.0, 1.0))( + outputs, + tcremp_target=batch["tcremp_target"], + pgen_target=batch["pgen_target"], + reconstruction_target=batch["reconstruction_target"], + **standardizer_tensors, + ) + losses["loss"].backward() + if not torch.isfinite(losses["loss"]): + raise FloatingPointError(f"WordPiece v{vocab_size} smoke loss is non-finite.") + candidates[str(vocab_size)] = { + "requested_vocab_size": vocab_size, + "actual_vocab_size": tokenizer.vocab_size, + "tokenizer_sha256": config["tokenizer_sha256"], + "vocab_sha256": config["vocab_sha256"], + "token_statistics_sha256": config["token_statistics_sha256"], + "parameter_count": sum(parameter.numel() for parameter in model.parameters()), + "parameter_delta_vs_char": sum(parameter.numel() for parameter in model.parameters()) + - baseline["parameter_count"], + "changed_model_fields_vs_char": changed_model_fields, + "changed_training_fields_vs_char": ["tokenizer_type", "tokenizer_path"], + "separate_wordpiece_encoder_char_decoder": True, + "finite_full_multitask_smoke_loss": True, + "smoke_loss": float(losses["loss"].detach()), + } + del model, outputs, losses, dataset, batch + gc.collect() + + report = { + "status": "ready", + "selection_rule": { + "primary": "minimum best-checkpoint full validation objective", + "secondary": ["validation pgen RMSE", "validation reconstruction", "validation TCRemP"], + "test_metrics_used_for_selection": False, + }, + "baseline": { + "job_id": "1425949", + "run_config_sha256": sha256(baseline_path), + "parameter_count": baseline["parameter_count"], + "non_tokenizer_training_fields_match": True, + }, + "benchmark": { + "path": str(data_dir.resolve()), + "ready_sha256": sha256(data_dir / "READY.json"), + "rows": len(table), + "embedding_shape": list(embeddings.shape), + "split_rows": {split: len(indices) for split, indices in split_indices.items()}, + "normalizer_fit_split": "train", + "normalizer_train_rows": standardizer.train_rows, + }, + "candidates": candidates, + "runtime": {"python": platform.python_version(), "torch": torch.__version__}, + } + output = Path(args.output_path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_anchored_tokenizers.py b/scripts/prepare_anchored_tokenizers.py new file mode 100644 index 0000000..2e575ab --- /dev/null +++ b/scripts/prepare_anchored_tokenizers.py @@ -0,0 +1,553 @@ +"""Prepare three leakage-safe sequence-only terminal-anchor tokenizers. + +This script is intentionally a CPU batch workload: it streams an independently +generated OLGA TRB corpus, removes every locked benchmark sequence, mines anchor +libraries from train-only evidence, fits central WordPiece models, and audits +exact round trips on every locked split. +""" + +from __future__ import annotations + +import argparse +import csv +import gzip +import hashlib +import itertools +import json +import math +import shutil +import tempfile +from collections import Counter, defaultdict +from pathlib import Path + +import numpy as np +from tokenizers import Tokenizer +from tokenizers.decoders import WordPiece as WordPieceDecoder +from tokenizers.models import WordPiece +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.trainers import WordPieceTrainer + +from irrm_codec.anchored_tokenization import AnchoredTokenizer, SPECIAL_TOKENS +from irrm_codec.multitask_data import load_prepared_benchmark, select_split_indices +from irrm_codec.tokenization import BOS_ID, EOS_ID, PAD_ID, UNK_ID, VALID_AA +from irrm_codec.wordpiece_tokenization import validate_wordpiece_tokenizer + + +ALPHABET = "ACDEFGHIKLMNPQRSTVWY" +CENTRAL_VOCAB_SIZE = 256 +EDGE_K = 3 +MAX_ANCHOR_LEN = 8 +MIN_DATA_TOTAL_SUPPORT = 50 +MIN_DATA_GENE_SUPPORT = 25 +MIN_GERMLINE_SUPPORT = 20 +SPECIAL_IDS = {"pad": PAD_ID, "bos": BOS_ID, "eos": EOS_ID, "unk": UNK_ID} +CODONS = { + "TTT":"F","TTC":"F","TTA":"L","TTG":"L","TCT":"S","TCC":"S","TCA":"S","TCG":"S", + "TAT":"Y","TAC":"Y","TAA":"*","TAG":"*","TGT":"C","TGC":"C","TGA":"*","TGG":"W", + "CTT":"L","CTC":"L","CTA":"L","CTG":"L","CCT":"P","CCC":"P","CCA":"P","CCG":"P", + "CAT":"H","CAC":"H","CAA":"Q","CAG":"Q","CGT":"R","CGC":"R","CGA":"R","CGG":"R", + "ATT":"I","ATC":"I","ATA":"I","ATG":"M","ACT":"T","ACC":"T","ACA":"T","ACG":"T", + "AAT":"N","AAC":"N","AAA":"K","AAG":"K","AGT":"S","AGC":"S","AGA":"R","AGG":"R", + "GTT":"V","GTC":"V","GTA":"V","GTG":"V","GCT":"A","GCC":"A","GCA":"A","GCG":"A", + "GAT":"D","GAC":"D","GAA":"E","GAG":"E","GGT":"G","GGC":"G","GGA":"G","GGG":"G", +} + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def translate(nt: str) -> str: + return "".join(CODONS.get(nt[index:index + 3].upper(), "X") for index in range(0, len(nt) - 2, 3)) + + +def canonical_call(value: object) -> str: + return str(value).split(",", 1)[0].split("*", 1)[0].strip() + + +def select_pair(sequence: str, n_anchors: tuple[str, ...], c_anchors: tuple[str, ...]) -> tuple[str, str, bool]: + n_matches = [anchor for anchor in n_anchors if sequence.startswith(anchor)] + c_matches = [anchor for anchor in c_anchors if sequence.endswith(anchor)] + feasible = [(n, c) for n in n_matches for c in c_matches if len(n) + len(c) <= len(sequence)] + if not feasible: + raise ValueError(f"No feasible anchor pair for {sequence!r}.") + pair = min(feasible, key=lambda value: (-(len(value[0]) + len(value[1])), -len(value[0]), -len(value[1]), value)) + raw_longest_overlap = bool(n_matches and c_matches and len(n_matches[0]) + len(c_matches[0]) > len(sequence)) + return pair[0], pair[1], raw_longest_overlap + + +def build_data_anchors(train) -> tuple[tuple[str, ...], tuple[str, ...], dict]: + side_payload = {} + libraries = [] + for side, gene_column in (("N", "v_call"), ("C", "j_call")): + total = Counter() + by_gene: dict[str, Counter] = defaultdict(Counter) + for sequence, raw_gene in zip(train["junction_aa"], train[gene_column], strict=True): + gene = canonical_call(raw_gene) + for length in range(1, min(MAX_ANCHOR_LEN, len(sequence)) + 1): + motif = sequence[:length] if side == "N" else sequence[-length:] + total[motif] += 1 + by_gene[motif][gene] += 1 + retained = set(ALPHABET) + evidence = {} + for motif, count in total.items(): + top_gene, top_count = by_gene[motif].most_common(1)[0] + if len(motif) >= 2 and count >= MIN_DATA_TOTAL_SUPPORT and top_count >= MIN_DATA_GENE_SUPPORT: + retained.add(motif) + evidence[motif] = { + "train_count": count, + "top_gene": top_gene, + "top_gene_count": top_count, + "top_gene_fraction": top_count / count, + } + library = tuple(sorted(retained, key=lambda value: (-len(value), value))) + libraries.append(library) + side_payload[side] = { + "candidate_rule": ( + f"train-only literal {side} flanks length 2..{MAX_ANCHOR_LEN}; total support >= " + f"{MIN_DATA_TOTAL_SUPPORT} and support within at least one annotated gene >= " + f"{MIN_DATA_GENE_SUPPORT}; exhaustive one-residue identity fallbacks" + ), + "retained": len(library), + "evidence": evidence, + } + return libraries[0], libraries[1], side_payload + + +def read_olga_genes(model_params: Path, section: str) -> dict[str, str]: + genes = {} + active = False + marker = f"#GeneChoice;{section}_gene;" + for raw_line in model_params.read_text(encoding="utf-8").splitlines(): + if raw_line.startswith("#"): + active = raw_line.startswith(marker) + continue + if active and raw_line.startswith("%"): + fields = raw_line[1:].strip().split(";") + if len(fields) >= 2: + genes[fields[0].strip()] = fields[1].strip().upper() + if not genes: + raise ValueError(f"No {section} genes found in {model_params}.") + return genes + + +def read_anchor_csv(path: Path) -> dict[str, int]: + with path.open(encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + return { + row["gene"]: int(row["anchor_index"]) + for row in rows + if row["function"].strip("()") == "F" + } + + +def build_germline_anchors(train, model_dir: Path) -> tuple[tuple[str, ...], tuple[str, ...], dict]: + params = model_dir / "model_params.txt" + v_anchor_path = model_dir / "V_gene_CDR3_anchors.csv" + j_anchor_path = model_dir / "J_gene_CDR3_anchors.csv" + v_genes = read_olga_genes(params, "V") + j_genes = read_olga_genes(params, "J") + v_indices = read_anchor_csv(v_anchor_path) + j_indices = read_anchor_csv(j_anchor_path) + reference = {"N": defaultdict(set), "C": defaultdict(set)} + for gene, index in v_indices.items(): + if gene not in v_genes: + continue + amino = translate(v_genes[gene][index:]) + for length in range(1, min(MAX_ANCHOR_LEN, len(amino)) + 1): + motif = amino[:length] + if set(motif) <= VALID_AA: + reference["N"][motif].add(gene) + for gene, index in j_indices.items(): + if gene not in j_genes: + continue + frame_start = index % 3 + amino = translate(j_genes[gene][frame_start:index + 3]) + for length in range(1, min(MAX_ANCHOR_LEN, len(amino)) + 1): + motif = amino[-length:] + if set(motif) <= VALID_AA: + reference["C"][motif].add(gene) + + payload = {} + libraries = [] + for side in ("N", "C"): + support = Counter() + for sequence in train["junction_aa"]: + for motif in reference[side]: + if (sequence.startswith(motif) if side == "N" else sequence.endswith(motif)): + support[motif] += 1 + retained = { + motif for motif, count in support.items() + if count >= MIN_GERMLINE_SUPPORT and len(motif) >= 2 + } + # Exhaustive one-residue identity fallbacks are constants, not learned motifs. + retained.update(ALPHABET) + library = tuple(sorted(retained, key=lambda value: (-len(value), value))) + libraries.append(library) + payload[side] = { + "retained": len(library), + "minimum_train_support": MIN_GERMLINE_SUPPORT, + "motifs": { + motif: {"train_support": support[motif], "source_genes": sorted(reference[side][motif])} + for motif in library + }, + } + payload["source_files"] = { + path.name: sha256(path) for path in (params, v_anchor_path, j_anchor_path) + } + payload["derivation"] = ( + "Functional OLGA TRB_orig V sequences translated from each V CDR3 anchor; functional J " + "sequences translated in the frame ending at each J conserved-F anchor; literal prefixes/" + "suffixes length 1..8 retained only with train support >=20." + ) + return libraries[0], libraries[1], payload + + +def stream_external_corpus(source: Path, output: Path, benchmark_sequences: set[str], rows: int) -> dict: + accepted = 0 + source_rows = 0 + rejected_overlap = 0 + rejected_invalid = 0 + rejected_duplicate = 0 + seen = set() + with gzip.open(source, "rt", encoding="utf-8") as input_handle, output.open("w", encoding="utf-8") as out: + header = input_handle.readline().rstrip("\n").split("\t") + aa_index = header.index("junction_aa") + for line in input_handle: + source_rows += 1 + fields = line.rstrip("\n").split("\t") + if len(fields) <= aa_index: + rejected_invalid += 1 + continue + sequence = fields[aa_index].strip().upper() + if not (8 <= len(sequence) <= 40 and set(sequence) <= VALID_AA and sequence.startswith("C") and sequence.endswith("F")): + rejected_invalid += 1 + continue + if sequence in benchmark_sequences: + rejected_overlap += 1 + continue + if sequence in seen: + rejected_duplicate += 1 + continue + seen.add(sequence) + out.write(f"{sequence}\n") + accepted += 1 + if accepted == rows: + break + if accepted != rows: + raise ValueError(f"Requested {rows} disjoint external rows but obtained {accepted}.") + return { + "source_path": str(source), + "source_sha256": sha256(source), + "source_rows_scanned": source_rows, + "accepted_unique_rows": accepted, + "benchmark_overlap_rejected": rejected_overlap, + "invalid_rejected": rejected_invalid, + "duplicate_rejected": rejected_duplicate, + "output_sha256": sha256(output), + "overlap_after_filter": len(seen.intersection(benchmark_sequences)), + "selection": "first valid unique rows in source order after locked-benchmark exclusion", + } + + +def fit_central_wordpiece(middle_path: Path, output_path: Path) -> Tokenizer: + tokenizer = Tokenizer(WordPiece(unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.decoder = WordPieceDecoder(prefix="##") + tokenizer.train( + [str(middle_path)], + WordPieceTrainer( + vocab_size=CENTRAL_VOCAB_SIZE, + min_frequency=2, + special_tokens=SPECIAL_TOKENS, + continuing_subword_prefix="##", + show_progress=False, + ), + ) + validate_wordpiece_tokenizer(tokenizer, middle_path) + if tokenizer.get_vocab_size() != CENTRAL_VOCAB_SIZE: + raise ValueError( + f"Central WordPiece requested {CENTRAL_VOCAB_SIZE}, got {tokenizer.get_vocab_size()}." + ) + tokenizer.save(str(output_path)) + return tokenizer + + +def write_bundle(directory: Path, kind: str, n_anchors: tuple[str, ...], c_anchors: tuple[str, ...], central: Tokenizer, provenance: dict) -> Path: + vocab = {token: identifier for identifier, token in enumerate(SPECIAL_TOKENS)} + for anchor in sorted(n_anchors): + vocab[f"[N:{anchor}]"] = len(vocab) + for anchor in sorted(c_anchors): + vocab[f"[C:{anchor}]"] = len(vocab) + central_by_id = {identifier: token for token, identifier in central.get_vocab().items()} + central_to_input = {} + for central_id in range(len(SPECIAL_TOKENS), central.get_vocab_size()): + token = f"[M:{central_by_id[central_id]}]" + central_to_input[str(central_id)] = len(vocab) + vocab[token] = len(vocab) + bundle = { + "format_version": 1, + "tokenizer_type": kind, + "inference_modality": "sequence_only", + "accepts_v_call_or_j_call": False, + "k": EDGE_K if kind == "edge_k" else None, + "special_token_ids": SPECIAL_IDS, + "central_tokenizer_file": "central_tokenizer.json", + "central_to_input": central_to_input, + "n_anchors": list(n_anchors), + "c_anchors": list(c_anchors), + "vocab": vocab, + "vocab_size": len(vocab), + "boundary_types": ["N_ANCHOR", "MIDDLE_WORDPIECE", "C_ANCHOR"], + "selection_rule": ( + f"literal first/last {EDGE_K} residues" if kind == "edge_k" else + "longest non-overlapping prefix/suffix pair; maximum total length, then N length, then C length, then lexical" + ), + "unknown_behavior": ( + f"fail closed only for invalid amino acids or length <{2 * EDGE_K}; every standard-amino-acid EDGE-{EDGE_K} literal is exhaustively enumerated" + if kind == "edge_k" else + "exhaustive one-residue identity fallbacks guarantee valid-amino-acid coverage without heldout fitting; otherwise fail closed" + ), + "provenance": provenance, + } + path = directory / "anchored_tokenizer.json" + path.write_text(json.dumps(bundle, indent=2), encoding="utf-8") + return path + + +def audit(kind: str, bundle_path: Path, sequences: dict[str, list[str]], train_pgen: np.ndarray) -> tuple[dict, list[dict]]: + tokenizer = AnchoredTokenizer(bundle_path) + counts = np.zeros(tokenizer.vocab_size, dtype=np.int64) + presence = np.zeros(tokenizer.vocab_size, dtype=np.int64) + position_sum = np.zeros(tokenizer.vocab_size, dtype=np.float64) + pgen_sum = np.zeros(tokenizer.vocab_size, dtype=np.float64) + reports = {} + for split, values in sequences.items(): + lengths, n_lengths, c_lengths = [], [], [] + mismatches = boundary_failures = overlap_adjustments = 0 + n_used, c_used = Counter(), Counter() + for row_number, sequence in enumerate(values): + encoded = tokenizer.encode_with_boundaries(sequence, 40) + mismatches += int(tokenizer.decode(encoded.ids) != sequence) + boundary_failures += int( + encoded.boundary_types[0] != "N_ANCHOR" + or encoded.boundary_types[-1] != "C_ANCHOR" + or encoded.n_anchor != sequence[:len(encoded.n_anchor)] + or encoded.c_anchor != sequence[-len(encoded.c_anchor):] + ) + lengths.append(len(encoded.ids)) + n_lengths.append(len(encoded.n_anchor)) + c_lengths.append(len(encoded.c_anchor)) + n_used[encoded.n_anchor] += 1 + c_used[encoded.c_anchor] += 1 + if kind != "edge_k": + _n, _c, adjusted = select_pair(sequence, tokenizer.n_anchors, tokenizer.c_anchors) + overlap_adjustments += int(adjusted) + if split == "train": + unique = np.unique(encoded.ids) + counts += np.bincount(encoded.ids, minlength=tokenizer.vocab_size) + presence[unique] += 1 + denominator = max(len(encoded.ids) - 1, 1) + for position, identifier in enumerate(encoded.ids): + position_sum[identifier] += position / denominator + pgen_sum[unique] += float(train_pgen[row_number]) + reports[split] = { + "rows": len(values), + "roundtrip_mismatches": mismatches, + "indel_or_substitution_failures": mismatches, + "boundary_failures": boundary_failures, + "token_length_mean": float(np.mean(lengths)), + "token_length_p95": float(np.quantile(lengths, 0.95)), + "token_length_max": int(max(lengths)), + "mean_n_anchor_length": float(np.mean(n_lengths)), + "mean_c_anchor_length": float(np.mean(c_lengths)), + "n_anchor_length_gt1_coverage": float(np.mean(np.asarray(n_lengths) > 1)), + "c_anchor_length_gt1_coverage": float(np.mean(np.asarray(c_lengths) > 1)), + "unique_n_anchors_used": len(n_used), + "unique_c_anchors_used": len(c_used), + "longest_pair_overlap_adjustments": overlap_adjustments, + "same_length_literal_collisions": 0, + } + global_mean = float(train_pgen.mean()) + global_std = float(train_pgen.std()) + rows = [] + for identifier in range(tokenizer.vocab_size): + token = tokenizer.id_to_token[identifier] + present = int(presence[identifier]) + absent = len(sequences["train"]) - present + present_mean = float(pgen_sum[identifier] / present) if present else math.nan + absent_mean = float((train_pgen.sum() - pgen_sum[identifier]) / absent) if absent else math.nan + fraction = present / len(sequences["train"]) + association = ( + (present_mean - global_mean) * math.sqrt(fraction * (1 - fraction)) / global_std + if present and absent and global_std > 0 else math.nan + ) + boundary = "SPECIAL" if identifier < 4 else ( + "N_ANCHOR" if token.startswith("[N:") else "C_ANCHOR" if token.startswith("[C:") else "MIDDLE_WORDPIECE" + ) + rows.append({ + "token_id": identifier, + "token": token, + "boundary_type": boundary, + "train_occurrence_count": int(counts[identifier]), + "train_sequence_count": present, + "train_mean_normalized_position": ( + float(position_sum[identifier] / counts[identifier]) if counts[identifier] else "" + ), + "train_log10_pgen_1mm_mean_present": present_mean if present else "", + "train_log10_pgen_1mm_mean_absent": absent_mean if absent else "", + "train_pgen_presence_correlation": association if math.isfinite(association) else "", + }) + return reports, rows + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", required=True) + parser.add_argument("--output-root", required=True) + parser.add_argument("--external-corpus", required=True) + parser.add_argument("--generation-script", required=True) + parser.add_argument("--olga-model-dir", required=True) + parser.add_argument("--external-rows", type=int, default=500_000) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + if args.seed != 42 or args.external_rows < 100_000: + raise ValueError("Anchored study requires seed 42 and a genuinely larger external corpus.") + + data_dir = Path(args.data_dir) + output_root = Path(args.output_root) + if (output_root / "READY.json").exists(): + raise FileExistsError(f"Refusing to overwrite immutable tokenizer root {output_root}.") + output_root.mkdir(parents=True, exist_ok=True) + table, _ = load_prepared_benchmark(data_dir) + indices = {split: select_split_indices(table, data_dir, split) for split in ("train", "val", "test")} + sequences = {split: table.iloc[value]["junction_aa"].astype(str).tolist() for split, value in indices.items()} + train = table.iloc[indices["train"]].reset_index(drop=True) + benchmark_sequences = set(table["junction_aa"].astype(str)) + train_manifest = data_dir / "manifests" / "train.tsv" + + corpus_path = output_root / "olga_trb_disjoint_500k.txt" + corpus_report = stream_external_corpus( + Path(args.external_corpus), corpus_path, benchmark_sequences, args.external_rows + ) + corpus_report.update({ + "generation_script": str(Path(args.generation_script)), + "generation_script_sha256": sha256(Path(args.generation_script)), + "declared_generator": "olga-generate_sequences --VDJ_model_folder=../model/Homo+sapiens/TRB_orig/", + "fit_usage": "central WordPiece only; never anchors or evaluation metrics", + }) + + edge_n = tuple("".join(value) for value in itertools.product(ALPHABET, repeat=EDGE_K)) + edge_c = tuple("".join(value) for value in itertools.product(ALPHABET, repeat=EDGE_K)) + data_n, data_c, data_provenance = build_data_anchors(train) + germ_n, germ_c, germ_provenance = build_germline_anchors(train, Path(args.olga_model_dir)) + candidates = { + "edge_k": (edge_n, edge_c, { + "anchor_fit": "exhaustive standard-amino-acid K-mer literal space on each side; no observed sequences used", + "k": EDGE_K, + "canonical_anchor_count_per_side": len(edge_n), + }), + "data_anchor": (data_n, data_c, { + "anchor_fit": "locked train rows only; V/J calls restrict candidate support but are not runtime inputs", + "train_manifest_sha256": sha256(train_manifest), + "mining": data_provenance, + }), + "germline_anchor": (germ_n, germ_c, { + "anchor_fit": "OLGA TRB_orig germline-compatible motifs retained by locked train support only", + "train_manifest_sha256": sha256(train_manifest), + "germline": germ_provenance, + }), + } + + top_manifest = { + "status": "preparation_complete", + "inference_modality": "sequence_only", + "forbidden_model_inputs": ["v_call", "j_call"], + "seed": args.seed, + "locked_split_rows": {split: len(value) for split, value in indices.items()}, + "train_manifest_sha256": sha256(train_manifest), + "external_corpus": corpus_report, + "central_wordpiece_vocab_size": CENTRAL_VOCAB_SIZE, + "candidates": {}, + } + train_pgen = train["log10_pgen_1mm"].to_numpy(np.float64) + for kind, (n_anchors, c_anchors, provenance) in candidates.items(): + final_dir = output_root / kind + if final_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {final_dir}.") + with tempfile.TemporaryDirectory(prefix=f".{kind}-", dir=output_root) as temporary: + candidate_dir = Path(temporary) + middle_path = candidate_dir / "central_corpus.txt" + with corpus_path.open(encoding="utf-8") as input_handle, middle_path.open("w", encoding="utf-8") as out: + for line in input_handle: + sequence = line.strip() + if kind == "edge_k": + n_anchor, c_anchor = sequence[:EDGE_K], sequence[-EDGE_K:] + else: + n_anchor, c_anchor, _ = select_pair(sequence, n_anchors, c_anchors) + middle = sequence[len(n_anchor):len(sequence) - len(c_anchor)] + if middle: + out.write(f"{middle}\n") + central_path = candidate_dir / "central_tokenizer.json" + central = fit_central_wordpiece(middle_path, central_path) + provenance = dict(provenance) + provenance.update({ + "central_corpus_sha256": sha256(middle_path), + "external_raw_corpus_sha256": sha256(corpus_path), + "central_fit_rows_source": args.external_rows, + "heldout_raw_sequences_used_for_fit": False, + }) + bundle_path = write_bundle(candidate_dir, kind, n_anchors, c_anchors, central, provenance) + reports, token_rows = audit(kind, bundle_path, sequences, train_pgen) + if any( + report[check] + for report in reports.values() + for check in ("roundtrip_mismatches", "boundary_failures") + ): + raise ValueError(f"{kind} failed exact round-trip/boundary audit: {reports}") + if max(report["token_length_max"] for report in reports.values()) > 40: + raise ValueError(f"{kind} exceeds model max input length: {reports}") + with (candidate_dir / "token_statistics.tsv").open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(token_rows[0])) + writer.writeheader() + writer.writerows(token_rows) + validation = { + "status": "prepared", + "checks": { + "sequence_only_inference": True, + "special_ids_0_1_2_3": True, + "train_only_anchor_fit": True, + "no_val_test_vocab_or_anchor_fit": True, + "external_corpus_benchmark_overlap_zero": corpus_report["overlap_after_filter"] == 0, + "exact_roundtrip_all_splits": True, + "no_indels_substitutions_all_splits": True, + "literal_protected_boundaries_all_splits": True, + "max_encoded_length_at_most_40": True, + }, + "encoding": reports, + } + (candidate_dir / "validation_report.json").write_text(json.dumps(validation, indent=2), encoding="utf-8") + config = json.loads(bundle_path.read_text(encoding="utf-8")) + config["bundle_sha256"] = sha256(bundle_path) + config["central_tokenizer_sha256"] = sha256(central_path) + config["token_statistics_sha256"] = sha256(candidate_dir / "token_statistics.tsv") + (candidate_dir / "tokenizer_config.json").write_text(json.dumps(config, indent=2), encoding="utf-8") + candidate_dir.rename(final_dir) + top_manifest["candidates"][kind] = { + "path": str(final_dir), + "input_vocab_size": config["vocab_size"], + "n_anchor_count": len(n_anchors), + "c_anchor_count": len(c_anchors), + "bundle_sha256": config["bundle_sha256"], + "central_tokenizer_sha256": config["central_tokenizer_sha256"], + "encoding": reports, + } + (output_root / "preparation_manifest.json").write_text(json.dumps(top_manifest, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_wordpiece_rq4.py b/scripts/prepare_wordpiece_rq4.py new file mode 100644 index 0000000..ce21df5 --- /dev/null +++ b/scripts/prepare_wordpiece_rq4.py @@ -0,0 +1,177 @@ +"""Train and validate the immutable RQ4 WordPiece tokenizer on TRB train only.""" + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +from tokenizers import Tokenizer +from tokenizers.decoders import WordPiece as WordPieceDecoder +from tokenizers.models import WordPiece +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.trainers import WordPieceTrainer + +from irrm_codec.multitask_data import load_prepared_benchmark, select_split_indices +from irrm_codec.tokenization import BOS_ID, EOS_ID, PAD_ID, UNK_ID +from irrm_codec.wordpiece_tokenization import validate_wordpiece_tokenizer + + +SPECIAL_TOKENS = ["[PAD]", "[BOS]", "[EOS]", "[UNK]"] +PREFIX = "##" + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def encode_report(tokenizer: Tokenizer, sequences: list[str]) -> dict: + token_lengths = [] + unknown_sequences = 0 + roundtrip_mismatches = 0 + for sequence in sequences: + encoding = tokenizer.encode(sequence) + token_lengths.append(len(encoding.ids)) + unknown_sequences += int(UNK_ID in encoding.ids) + decoded = tokenizer.decode(encoding.ids, skip_special_tokens=True) + roundtrip_mismatches += int(decoded != sequence) + values = np.asarray(token_lengths, dtype=np.int64) + return { + "rows": len(sequences), + "unknown_sequences": unknown_sequences, + "roundtrip_mismatches": roundtrip_mismatches, + "token_length_min": int(values.min()), + "token_length_max": int(values.max()), + "token_length_mean": float(values.mean()), + "token_length_p50": float(np.quantile(values, 0.5)), + "token_length_p95": float(np.quantile(values, 0.95)), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--vocab-size", type=int, default=256) + parser.add_argument("--min-frequency", type=int, default=2) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + if args.seed != 42: + raise ValueError("RQ4 tokenizer seed must remain fixed at 42.") + if args.vocab_size != 256: + raise ValueError("RQ4 uses the preregistered 256-token WordPiece vocabulary.") + output_dir = Path(args.output_dir) + tokenizer_path = output_dir / "tokenizer.json" + if tokenizer_path.exists(): + raise FileExistsError(f"Refusing to overwrite immutable tokenizer: {tokenizer_path}") + output_dir.mkdir(parents=True, exist_ok=True) + + data_dir = Path(args.data_dir) + table, _embeddings = load_prepared_benchmark(data_dir) + split_indices = { + split: select_split_indices(table, data_dir, split) + for split in ("train", "val", "test") + } + train_manifest = data_dir / "manifests" / "train.tsv" + sequences = { + split: table.iloc[indices]["junction_aa"].astype(str).tolist() + for split, indices in split_indices.items() + } + + corpus_path = output_dir / "train_cdr3.txt" + corpus_path.write_text( + "".join(f"{sequence}\n" for sequence in sequences["train"]), + encoding="utf-8", + ) + tokenizer = Tokenizer(WordPiece(unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.decoder = WordPieceDecoder(prefix=PREFIX) + tokenizer.train( + [str(corpus_path)], + WordPieceTrainer( + vocab_size=args.vocab_size, + min_frequency=args.min_frequency, + special_tokens=SPECIAL_TOKENS, + continuing_subword_prefix=PREFIX, + show_progress=False, + ), + ) + validate_wordpiece_tokenizer(tokenizer, "newly trained RQ4 tokenizer") + actual_vocab_size = tokenizer.get_vocab_size() + if actual_vocab_size != args.vocab_size: + raise ValueError( + f"Requested vocabulary {args.vocab_size}, got {actual_vocab_size}." + ) + + reports = {split: encode_report(tokenizer, values) for split, values in sequences.items()} + for split, report in reports.items(): + if report["unknown_sequences"] or report["roundtrip_mismatches"]: + raise ValueError(f"Tokenizer failed {split} coverage/roundtrip validation: {report}") + if report["token_length_max"] > 40: + raise ValueError(f"Tokenizer exceeds model width on {split}: {report}") + + tokenizer.save(str(tokenizer_path)) + tokenizer.model.save(str(output_dir)) + special_ids = { + "pad": PAD_ID, + "bos": BOS_ID, + "eos": EOS_ID, + "unk": UNK_ID, + } + config = { + "status": "immutable", + "tokenizer_type": "wordpiece", + "seed": args.seed, + "requested_vocab_size": args.vocab_size, + "actual_vocab_size": actual_vocab_size, + "min_frequency": args.min_frequency, + "continuing_subword_prefix": PREFIX, + "special_tokens": SPECIAL_TOKENS, + "special_token_ids": special_ids, + "training_split": "train", + "training_rows": len(split_indices["train"]), + "train_manifest_path": str(train_manifest.resolve()), + "train_manifest_sha256": sha256(train_manifest), + "corpus_path": str(corpus_path.resolve()), + "corpus_sha256": sha256(corpus_path), + "tokenizer_path": str(tokenizer_path.resolve()), + } + (output_dir / "tokenizer_config.json").write_text( + json.dumps(config, indent=2), encoding="utf-8" + ) + validation = { + "status": "ready", + "train_only_fit": True, + "split_rows": {split: len(indices) for split, indices in split_indices.items()}, + "split_index_overlap": { + "train_val": int(np.intersect1d(split_indices["train"], split_indices["val"]).size), + "train_test": int(np.intersect1d(split_indices["train"], split_indices["test"]).size), + "val_test": int(np.intersect1d(split_indices["val"], split_indices["test"]).size), + }, + "encoding": reports, + "checks": { + "special_ids_match_model": True, + "train_manifest_matches_locked_split": True, + "no_split_overlap": True, + "no_unknown_tokens": True, + "roundtrip_exact": True, + "all_encoded_lengths_at_most_40": True, + }, + } + (output_dir / "validation_report.json").write_text( + json.dumps(validation, indent=2), encoding="utf-8" + ) + config["tokenizer_sha256"] = sha256(tokenizer_path) + config["vocab_sha256"] = sha256(output_dir / "vocab.txt") + (output_dir / "tokenizer_config.json").write_text( + json.dumps(config, indent=2), encoding="utf-8" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_wordpiece_sweep.py b/scripts/prepare_wordpiece_sweep.py new file mode 100644 index 0000000..9d6680c --- /dev/null +++ b/scripts/prepare_wordpiece_sweep.py @@ -0,0 +1,354 @@ +"""Train and audit the leakage-safe TRB WordPiece vocabulary sweep.""" + +import argparse +import csv +import hashlib +import json +import math +import tempfile +from pathlib import Path + +import numpy as np +from tokenizers import Tokenizer +from tokenizers.decoders import WordPiece as WordPieceDecoder +from tokenizers.models import WordPiece +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.trainers import WordPieceTrainer + +from irrm_codec.multitask_data import load_prepared_benchmark, select_split_indices +from irrm_codec.tokenization import BOS_ID, EOS_ID, PAD_ID, UNK_ID +from irrm_codec.wordpiece_tokenization import validate_wordpiece_tokenizer + + +SPECIAL_TOKENS = ["[PAD]", "[BOS]", "[EOS]", "[UNK]"] +PREFIX = "##" +DEFAULT_VOCAB_SIZES = (44, 64, 128, 256, 512, 1024, 2048, 4096) + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _safe_float(value: float) -> float | None: + return float(value) if math.isfinite(value) else None + + +def audit_tokenizer( + tokenizer: Tokenizer, + sequences: dict[str, list[str]], + train_pgen: np.ndarray, +) -> tuple[dict, list[dict]]: + vocab_size = tokenizer.get_vocab_size() + token_by_id = {identifier: token for token, identifier in tokenizer.get_vocab().items()} + split_counts = {split: np.zeros(vocab_size, dtype=np.int64) for split in sequences} + split_presence = {split: np.zeros(vocab_size, dtype=np.int64) for split in sequences} + reports = {} + train_start = np.zeros(vocab_size, dtype=np.int64) + train_end = np.zeros(vocab_size, dtype=np.int64) + train_internal = np.zeros(vocab_size, dtype=np.int64) + train_position_sum = np.zeros(vocab_size, dtype=np.float64) + train_pgen_present_sum = np.zeros(vocab_size, dtype=np.float64) + train_low_presence = np.zeros(vocab_size, dtype=np.int64) + train_high_presence = np.zeros(vocab_size, dtype=np.int64) + low_threshold, high_threshold = np.quantile(train_pgen, [0.25, 0.75]) + + for split, values in sequences.items(): + token_lengths = [] + unknown_sequences = 0 + roundtrip_mismatches = 0 + for row_number, (sequence, encoding) in enumerate( + zip(values, tokenizer.encode_batch(values), strict=True) + ): + ids = encoding.ids + token_lengths.append(len(ids)) + unknown_sequences += int(UNK_ID in ids) + roundtrip_mismatches += int( + tokenizer.decode(ids, skip_special_tokens=True) != sequence + ) + split_counts[split] += np.bincount(ids, minlength=vocab_size) + unique_ids = np.unique(ids) + split_presence[split][unique_ids] += 1 + if split == "train": + train_start[ids[0]] += 1 + train_end[ids[-1]] += 1 + if len(ids) > 2: + train_internal += np.bincount(ids[1:-1], minlength=vocab_size) + denominator = max(len(ids) - 1, 1) + for position, token_id in enumerate(ids): + train_position_sum[token_id] += position / denominator + pgen = float(train_pgen[row_number]) + train_pgen_present_sum[unique_ids] += pgen + if pgen <= low_threshold: + train_low_presence[unique_ids] += 1 + if pgen >= high_threshold: + train_high_presence[unique_ids] += 1 + lengths = np.asarray(token_lengths, dtype=np.int64) + reports[split] = { + "rows": len(values), + "unknown_sequences": unknown_sequences, + "roundtrip_mismatches": roundtrip_mismatches, + "token_length_min": int(lengths.min()), + "token_length_max": int(lengths.max()), + "token_length_mean": float(lengths.mean()), + "token_length_p50": float(np.quantile(lengths, 0.5)), + "token_length_p95": float(np.quantile(lengths, 0.95)), + } + + train_rows = len(sequences["train"]) + train_global_mean = float(train_pgen.mean()) + train_global_std = float(train_pgen.std()) + low_rows = int(np.count_nonzero(train_pgen <= low_threshold)) + high_rows = int(np.count_nonzero(train_pgen >= high_threshold)) + total_train_occurrences = int(split_counts["train"].sum()) + global_start_fraction = train_rows / total_train_occurrences + global_end_fraction = train_rows / total_train_occurrences + rows = [] + for token_id in range(vocab_size): + count = int(split_counts["train"][token_id]) + present = int(split_presence["train"][token_id]) + absent = train_rows - present + mean_present = ( + train_pgen_present_sum[token_id] / present if present else math.nan + ) + mean_absent = ( + (float(train_pgen.sum()) - train_pgen_present_sum[token_id]) / absent + if absent + else math.nan + ) + presence_fraction = present / train_rows + association = ( + (mean_present - train_global_mean) + * math.sqrt(presence_fraction * (1.0 - presence_fraction)) + / train_global_std + if present and absent and train_global_std > 0 + else math.nan + ) + high_rate = (train_high_presence[token_id] + 0.5) / (high_rows + 1.0) + low_rate = (train_low_presence[token_id] + 0.5) / (low_rows + 1.0) + start_fraction = train_start[token_id] / count if count else math.nan + end_fraction = train_end[token_id] / count if count else math.nan + rows.append( + { + "token_id": token_id, + "token": token_by_id[token_id], + "surface": token_by_id[token_id].removeprefix(PREFIX), + "is_special": token_id < len(SPECIAL_TOKENS), + "train_count": count, + "val_count": int(split_counts["val"][token_id]), + "test_count": int(split_counts["test"][token_id]), + "train_sequence_count": present, + "train_sequence_fraction": presence_fraction, + "train_start_count": int(train_start[token_id]), + "train_internal_count": int(train_internal[token_id]), + "train_end_count": int(train_end[token_id]), + "train_mean_normalized_position": _safe_float( + train_position_sum[token_id] / count if count else math.nan + ), + "start_log2_enrichment": _safe_float( + math.log2(start_fraction / global_start_fraction) + if count and start_fraction > 0 + else math.nan + ), + "end_log2_enrichment": _safe_float( + math.log2(end_fraction / global_end_fraction) + if count and end_fraction > 0 + else math.nan + ), + "train_pgen_mean_present": _safe_float(mean_present), + "train_pgen_mean_absent": _safe_float(mean_absent), + "train_pgen_mean_difference": _safe_float(mean_present - mean_absent), + "train_pgen_presence_correlation": _safe_float(association), + "high_vs_low_pgen_log2_enrichment": _safe_float(math.log2(high_rate / low_rate)), + } + ) + + eligible = [row for row in rows if not row["is_special"] and row["train_sequence_count"] >= 50] + summary = { + "encoding": reports, + "pgen_association": { + "split": "train", + "target": "log10_pgen_1mm", + "global_mean": train_global_mean, + "global_std": train_global_std, + "low_quartile_threshold": float(low_threshold), + "high_quartile_threshold": float(high_threshold), + "interpretation": "Descriptive token-presence associations only; not biological or causal claims.", + }, + "top_tokens": { + "frequency": sorted(eligible, key=lambda row: row["train_count"], reverse=True)[:25], + "absolute_pgen_association": sorted( + eligible, + key=lambda row: abs(row["train_pgen_presence_correlation"] or 0.0), + reverse=True, + )[:25], + "start_enrichment": sorted( + eligible, key=lambda row: row["start_log2_enrichment"] or -math.inf, reverse=True + )[:25], + "end_enrichment": sorted( + eligible, key=lambda row: row["end_log2_enrichment"] or -math.inf, reverse=True + )[:25], + }, + } + return summary, rows + + +def write_token_statistics(path: Path, rows: list[dict]) -> None: + with path.open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", required=True) + parser.add_argument("--output-root", required=True) + parser.add_argument("--vocab-sizes", nargs="+", type=int, default=list(DEFAULT_VOCAB_SIZES)) + parser.add_argument("--min-frequency", type=int, default=2) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + vocab_sizes = tuple(args.vocab_sizes) + if args.seed != 42 or vocab_sizes != DEFAULT_VOCAB_SIZES: + raise ValueError("WordPiece Stage A requires seed 42 and the locked eight-size sweep.") + + data_dir = Path(args.data_dir) + output_root = Path(args.output_root) + output_root.mkdir(parents=True, exist_ok=True) + table, _embeddings = load_prepared_benchmark(data_dir) + split_indices = { + split: select_split_indices(table, data_dir, split) for split in ("train", "val", "test") + } + sequences = { + split: table.iloc[indices]["junction_aa"].astype(str).tolist() + for split, indices in split_indices.items() + } + train_pgen = table.iloc[split_indices["train"]]["log10_pgen_1mm"].to_numpy(np.float64) + train_manifest = data_dir / "manifests" / "train.tsv" + corpus_path = output_root / "train_cdr3.txt" + corpus_text = "".join(f"{sequence}\n" for sequence in sequences["train"]) + if corpus_path.exists() and corpus_path.read_text(encoding="utf-8") != corpus_text: + raise ValueError("Existing WordPiece sweep corpus differs from the locked train split.") + corpus_path.write_text(corpus_text, encoding="utf-8") + + sweep = { + "status": "ready", + "tokenizer_type": "wordpiece", + "vocab_sizes": list(vocab_sizes), + "seed": args.seed, + "training_split": "train", + "training_rows": len(split_indices["train"]), + "train_manifest_sha256": sha256(train_manifest), + "corpus_sha256": sha256(corpus_path), + "tokenizers": [], + } + for vocab_size in vocab_sizes: + final_dir = output_root / f"v{vocab_size}-seed42" + if final_dir.exists(): + raise FileExistsError(f"Refusing to overwrite immutable tokenizer directory: {final_dir}") + with tempfile.TemporaryDirectory(prefix=f".v{vocab_size}-", dir=output_root) as temporary: + token_dir = Path(temporary) + tokenizer = Tokenizer(WordPiece(unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.decoder = WordPieceDecoder(prefix=PREFIX) + tokenizer.train( + [str(corpus_path)], + WordPieceTrainer( + vocab_size=vocab_size, + min_frequency=args.min_frequency, + special_tokens=SPECIAL_TOKENS, + continuing_subword_prefix=PREFIX, + show_progress=False, + ), + ) + validate_wordpiece_tokenizer(tokenizer, f"WordPiece v{vocab_size}") + actual_vocab_size = tokenizer.get_vocab_size() + if actual_vocab_size != vocab_size: + raise ValueError( + f"Requested WordPiece vocabulary {vocab_size}, got {actual_vocab_size}; size is infeasible." + ) + audit, rows = audit_tokenizer(tokenizer, sequences, train_pgen) + for split, report in audit["encoding"].items(): + if report["unknown_sequences"] or report["roundtrip_mismatches"]: + raise ValueError(f"WordPiece v{vocab_size} failed {split}: {report}") + if report["token_length_max"] > 40: + raise ValueError(f"WordPiece v{vocab_size} exceeds max length on {split}: {report}") + tokenizer_path = token_dir / "tokenizer.json" + tokenizer.save(str(tokenizer_path)) + tokenizer.model.save(str(token_dir)) + write_token_statistics(token_dir / "token_statistics.tsv", rows) + config = { + "status": "immutable", + "tokenizer_type": "wordpiece", + "seed": args.seed, + "requested_vocab_size": vocab_size, + "actual_vocab_size": actual_vocab_size, + "min_frequency": args.min_frequency, + "continuing_subword_prefix": PREFIX, + "special_tokens": SPECIAL_TOKENS, + "special_token_ids": {"pad": PAD_ID, "bos": BOS_ID, "eos": EOS_ID, "unk": UNK_ID}, + "training_split": "train", + "training_rows": len(split_indices["train"]), + "train_manifest_path": str(train_manifest.resolve()), + "train_manifest_sha256": sha256(train_manifest), + "corpus_path": str(corpus_path.resolve()), + "corpus_sha256": sha256(corpus_path), + } + (token_dir / "audit_report.json").write_text(json.dumps(audit, indent=2), encoding="utf-8") + validation = { + "status": "ready", + "train_only_fit": True, + "split_rows": {split: len(indices) for split, indices in split_indices.items()}, + "split_index_overlap": { + "train_val": int(np.intersect1d(split_indices["train"], split_indices["val"]).size), + "train_test": int(np.intersect1d(split_indices["train"], split_indices["test"]).size), + "val_test": int(np.intersect1d(split_indices["val"], split_indices["test"]).size), + }, + "encoding": audit["encoding"], + "checks": { + "special_ids_match_model": True, + "requested_vocab_size_feasible": True, + "train_manifest_matches_locked_split": True, + "no_split_overlap": True, + "no_unknown_tokens": True, + "roundtrip_exact": True, + "all_encoded_lengths_at_most_40": True, + "pgen_association_train_only": True, + }, + } + (token_dir / "validation_report.json").write_text( + json.dumps(validation, indent=2), encoding="utf-8" + ) + config.update( + { + "tokenizer_sha256": sha256(tokenizer_path), + "vocab_sha256": sha256(token_dir / "vocab.txt"), + "token_statistics_sha256": sha256(token_dir / "token_statistics.tsv"), + "audit_report_sha256": sha256(token_dir / "audit_report.json"), + } + ) + (token_dir / "tokenizer_config.json").write_text( + json.dumps(config, indent=2), encoding="utf-8" + ) + token_dir.rename(final_dir) + sweep["tokenizers"].append( + { + "requested_vocab_size": vocab_size, + "actual_vocab_size": actual_vocab_size, + "directory": str(final_dir.resolve()), + "tokenizer_sha256": config["tokenizer_sha256"], + "vocab_sha256": config["vocab_sha256"], + "token_statistics_sha256": config["token_statistics_sha256"], + "max_encoded_lengths": { + split: audit["encoding"][split]["token_length_max"] for split in sequences + }, + } + ) + (output_root / "sweep_manifest.json").write_text(json.dumps(sweep, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_vdjdb_tcremp_umap.sbatch b/scripts/run_vdjdb_tcremp_umap.sbatch new file mode 100644 index 0000000..0c64396 --- /dev/null +++ b/scripts/run_vdjdb_tcremp_umap.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-tcremp-umap +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=08:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/tcremp-umap-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/tcremp-umap-%j.err + +set -euo pipefail +repo=/home/evlasova/irrm-codec +runtime=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/umap-runtime-v1 +out=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/tcremp-umap-${SLURM_JOB_ID} +cd "$repo" +"$runtime/bin/python" benchmark/create_vdjdb_tcremp_umap.py \ + --cohort /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv \ + --tcremp /home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/tcremp.npy \ + --output-dir "$out" diff --git a/scripts/run_vdjdb_umap_candidates.sbatch b/scripts/run_vdjdb_umap_candidates.sbatch new file mode 100644 index 0000000..c400b96 --- /dev/null +++ b/scripts/run_vdjdb_umap_candidates.sbatch @@ -0,0 +1,18 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-umap +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/umap-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/umap-%j.err + +set -euo pipefail +repo=/home/evlasova/irrm-codec +out=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/umap-${SLURM_JOB_ID} +cd "$repo" +/home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/create_vdjdb_umap_candidates.py \ + --cohort /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv \ + --native-evaluation /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728 \ + --output-dir "$out" diff --git a/scripts/run_vdjdb_umap_candidates_runtime.sbatch b/scripts/run_vdjdb_umap_candidates_runtime.sbatch new file mode 100644 index 0000000..6941960 --- /dev/null +++ b/scripts/run_vdjdb_umap_candidates_runtime.sbatch @@ -0,0 +1,25 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-umap +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/umap-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/umap-%j.err + +set -euo pipefail +repo=/home/evlasova/irrm-codec +runtime=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/umap-runtime-v1 +out=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/umap-${SLURM_JOB_ID} +base_python=/home/evlasova/.conda/envs/irrm-codec/bin/python +cd "$repo" +if [[ ! -x "$runtime/bin/python" ]]; then + "$base_python" -m venv --system-site-packages "$runtime" + "$runtime/bin/python" -m pip install --upgrade pip + "$runtime/bin/python" -m pip install umap-learn +fi +"$runtime/bin/python" benchmark/create_vdjdb_umap_candidates.py \ + --cohort /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv \ + --native-evaluation /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728 \ + --output-dir "$out" diff --git a/scripts/run_vdjdb_uniform100.sbatch b/scripts/run_vdjdb_uniform100.sbatch new file mode 100644 index 0000000..238fb25 --- /dev/null +++ b/scripts/run_vdjdb_uniform100.sbatch @@ -0,0 +1,22 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-uniform100 +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --time=04:00:00 +#SBATCH --output=slurm-%j.out +#SBATCH --error=slurm-%j.err + +set -euo pipefail + +repo=/home/evlasova/irrm-codec +out=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/evaluation-${SLURM_JOB_ID} +mkdir -p /home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1 + +cd "$repo" +python benchmark/evaluate_vdjdb_uniform100.py \ + --cohort /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv \ + --native-evaluation /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728 \ + --external-representations /home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/vdjdb/representations \ + --output-dir "$out" diff --git a/scripts/run_vdjdb_uniform100_direct.sbatch b/scripts/run_vdjdb_uniform100_direct.sbatch new file mode 100644 index 0000000..da2db2a --- /dev/null +++ b/scripts/run_vdjdb_uniform100_direct.sbatch @@ -0,0 +1,21 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-uniform100 +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --time=04:00:00 +#SBATCH --output=slurm-%j.out +#SBATCH --error=slurm-%j.err + +set -euo pipefail + +repo=/home/evlasova/irrm-codec +out=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/evaluation-${SLURM_JOB_ID} + +cd "$repo" +/home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/evaluate_vdjdb_uniform100.py \ + --cohort /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv \ + --native-evaluation /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728 \ + --external-representations /home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/vdjdb/representations \ + --output-dir "$out" diff --git a/scripts/run_vdjdb_uniform100_separation.sbatch b/scripts/run_vdjdb_uniform100_separation.sbatch new file mode 100644 index 0000000..32f3198 --- /dev/null +++ b/scripts/run_vdjdb_uniform100_separation.sbatch @@ -0,0 +1,20 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-u100-sep +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --time=04:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/separation-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/separation-%j.err + +set -euo pipefail +repo=/home/evlasova/irrm-codec +out=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/separation-${SLURM_JOB_ID} +cd "$repo" +/home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/evaluate_vdjdb_uniform100_separation.py \ + --cohort /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv \ + --native-evaluation /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728 \ + --external-representations /home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/vdjdb/representations \ + --uniform-evaluation /home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/evaluation-1427813 \ + --output-dir "$out" diff --git a/scripts/run_vdjdb_whole_umap.sbatch b/scripts/run_vdjdb_whole_umap.sbatch new file mode 100644 index 0000000..18902a0 --- /dev/null +++ b/scripts/run_vdjdb_whole_umap.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-whole-umap +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=08:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/whole-umap-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/whole-umap-%j.err + +set -euo pipefail +repo=/home/evlasova/irrm-codec +runtime=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/umap-runtime-v1 +out=/home/evlasova/irrm-codec/experiments/trb/vdjdb-uniform100-v1/whole-umap-${SLURM_JOB_ID} +cd "$repo" +"$runtime/bin/python" benchmark/create_vdjdb_whole_umap.py \ + --cohort /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv \ + --native-evaluation /home/evlasova/irrm-codec/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728 \ + --output-dir "$out" diff --git a/slurm/benchmark_vdjdb_irrm_cpu.sbatch b/slurm/benchmark_vdjdb_irrm_cpu.sbatch new file mode 100644 index 0000000..4ffa90e --- /dev/null +++ b/slurm/benchmark_vdjdb_irrm_cpu.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-irrm-cpu +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=16 +#SBATCH --mem=32G +#SBATCH --time=2:00:00 + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/vdjdb-runtime-benchmark/irrm-cpu-${SLURM_JOB_ID}" +mkdir -p "$OUT" +cd "$ROOT" +/home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/benchmark_vdjdb_irrm_runtime.py \ + --cohort "$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/cohort.airr.tsv" \ + --sequence-column junction_aa \ + --checkpoint "$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt" \ + --tokenizer "$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" \ + --output "$OUT/RESULTS.json" --device cpu --batch-size 256 diff --git a/slurm/benchmark_vdjdb_irrm_cpu_fixed.sbatch b/slurm/benchmark_vdjdb_irrm_cpu_fixed.sbatch new file mode 100644 index 0000000..57bb086 --- /dev/null +++ b/slurm/benchmark_vdjdb_irrm_cpu_fixed.sbatch @@ -0,0 +1,13 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-irrm-cpu +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=16 +#SBATCH --mem=32G +#SBATCH --time=2:00:00 +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/vdjdb-runtime-benchmark/irrm-cpu-${SLURM_JOB_ID}" +mkdir -p "$OUT" +cd "$ROOT" +PYTHONPATH="$ROOT" /home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/benchmark_vdjdb_irrm_runtime.py --cohort "$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/cohort.airr.tsv" --sequence-column junction_aa --checkpoint "$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt" --tokenizer "$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" --output "$OUT/RESULTS.json" --device cpu --batch-size 256 diff --git a/slurm/benchmark_vdjdb_irrm_gpu.sbatch b/slurm/benchmark_vdjdb_irrm_gpu.sbatch new file mode 100644 index 0000000..76931a9 --- /dev/null +++ b/slurm/benchmark_vdjdb_irrm_gpu.sbatch @@ -0,0 +1,20 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-irrm-gpu +#SBATCH --partition=short +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=2:00:00 + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/vdjdb-runtime-benchmark/irrm-gpu-${SLURM_JOB_ID}" +mkdir -p "$OUT" +cd "$ROOT" +/home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/benchmark_vdjdb_irrm_runtime.py \ + --cohort "$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/cohort.airr.tsv" \ + --sequence-column junction_aa \ + --checkpoint "$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt" \ + --tokenizer "$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" \ + --output "$OUT/RESULTS.json" --device cuda --batch-size 512 diff --git a/slurm/benchmark_vdjdb_irrm_gpu_fixed.sbatch b/slurm/benchmark_vdjdb_irrm_gpu_fixed.sbatch new file mode 100644 index 0000000..2836eba --- /dev/null +++ b/slurm/benchmark_vdjdb_irrm_gpu_fixed.sbatch @@ -0,0 +1,14 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-irrm-gpu +#SBATCH --partition=short +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=2:00:00 +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/vdjdb-runtime-benchmark/irrm-gpu-${SLURM_JOB_ID}" +mkdir -p "$OUT" +cd "$ROOT" +PYTHONPATH="$ROOT" /home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/benchmark_vdjdb_irrm_runtime.py --cohort "$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/cohort.airr.tsv" --sequence-column junction_aa --checkpoint "$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt" --tokenizer "$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" --output "$OUT/RESULTS.json" --device cuda --batch-size 512 diff --git a/slurm/benchmark_vdjdb_tcremp_runtime.sbatch b/slurm/benchmark_vdjdb_tcremp_runtime.sbatch new file mode 100644 index 0000000..abe1cff --- /dev/null +++ b/slurm/benchmark_vdjdb_tcremp_runtime.sbatch @@ -0,0 +1,30 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-tcremp-runtime +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --time=2:00:00 + +set -euo pipefail +THREADS="${1:?Pass the requested TCRemP worker count (1 or 16)}" +case "$THREADS" in 1|16) ;; *) echo "threads must be 1 or 16" >&2; exit 2;; esac +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/vdjdb-runtime-benchmark/tcremp-${THREADS}cpu-${SLURM_JOB_ID}" +TCREMP_ROOT=/home/evlasova/tcremp-git/tcremp +PYTHON=/home/evlasova/.conda/envs/vdjdb-redcea/bin/python +mkdir -p "$OUT/raw_tcremp" +cd "$ROOT" +STARTED="$(date -Iseconds)" +SECONDS=0 +PYTHONPATH="$TCREMP_ROOT:$ROOT" "$PYTHON" -m tcremp.tcremp_run \ + -i "$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/cohort.airr.tsv" \ + -o "$OUT/raw_tcremp" -e cohort -c TRB -s HomoSapiens -x clone_id \ + -p "$TCREMP_ROOT/tcremp/resources/tcremp_prototypes_TRB.tsv" \ + -np "$THREADS" -llen 4 -hlen 41 --skip-clustering 2>&1 | tee "$OUT/tcremp.log" +python - "$OUT/RESULTS.json" "$THREADS" "$SECONDS" "$STARTED" <<'PY' +import json, sys +from pathlib import Path +path, threads, seconds, started = map(str, sys.argv[1:]) +Path(path).write_text(json.dumps({"status": "complete", "threads": int(threads), "wall_seconds_including_load_and_write": int(seconds), "started": started}, indent=2) + "\n") +PY diff --git a/slurm/cache_author_global_embeddings.sbatch b/slurm/cache_author_global_embeddings.sbatch new file mode 100644 index 0000000..075ff14 --- /dev/null +++ b/slurm/cache_author_global_embeddings.sbatch @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=irrm-author-global-cache +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=08:00:00 +#SBATCH --array=0-1%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/cache-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/cache-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-global-author-v1" +RUNTIME="$REPO/experiments/trb/external-model-benchmark/runtime-v2" +LOCKED="$REPO/experiments/trb/external-model-benchmark/cohort" +LOCKED_REPS="$LOCKED/representations" +REDCEA="$REPO/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" +REDCEA_REPS="$BASE/vdjdb/representations" + +export PYTHONPATH="$RUNTIME/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export HF_HOME="$RUNTIME/huggingface" +export TORCH_HOME="$RUNTIME/torch" +export PYTHONUNBUFFERED=1 +export TOKENIZERS_PARALLELISM=false +export HF_HUB_DISABLE_TELEMETRY=1 +mkdir -p "$BASE/logs" "$REDCEA_REPS" +cd "$REPO" + +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader +if [[ "$SLURM_ARRAY_TASK_ID" == 0 ]]; then + "$PYTHON" benchmark/cache_external_model_embeddings.py --dataset-dir "$LOCKED" --output-dir "$LOCKED_REPS" --encoders esm2_35m --device cuda --batch-size 256 --threads "${SLURM_CPUS_PER_TASK:-8}" + "$PYTHON" benchmark/cache_external_model_embeddings.py --dataset-dir "$LOCKED" --dataset-file "$REDCEA" --output-dir "$REDCEA_REPS" --encoders esm2_35m tcr_bert --device cuda --batch-size 256 --threads "${SLURM_CPUS_PER_TASK:-8}" + sha256sum "$LOCKED_REPS/esm2_35m.npy" "$REDCEA_REPS/esm2_35m.npy" "$REDCEA_REPS/tcr_bert.npy" +else + "$PYTHON" benchmark/cache_external_model_embeddings.py --dataset-dir "$LOCKED" --output-dir "$LOCKED_REPS" --encoders sceptr_cdr3 --device cuda --batch-size 256 --threads "${SLURM_CPUS_PER_TASK:-8}" + "$PYTHON" benchmark/cache_external_model_embeddings.py --dataset-dir "$LOCKED" --dataset-file "$REDCEA" --output-dir "$REDCEA_REPS" --encoders sceptr sceptr_cdr3 --device cuda --batch-size 256 --threads "${SLURM_CPUS_PER_TASK:-8}" + sha256sum "$LOCKED_REPS/sceptr_cdr3.npy" "$REDCEA_REPS/sceptr.npy" "$REDCEA_REPS/sceptr_cdr3.npy" +fi +echo "status=accepted task=$SLURM_ARRAY_TASK_ID" diff --git a/slurm/cache_author_global_sceptr_cdr3.sbatch b/slurm/cache_author_global_sceptr_cdr3.sbatch new file mode 100644 index 0000000..0b19c3b --- /dev/null +++ b/slurm/cache_author_global_sceptr_cdr3.sbatch @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=irrm-sceptr-cdr3-redcea +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=03:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/sceptr-cdr3-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/sceptr-cdr3-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-global-author-v1" +RUNTIME="$REPO/experiments/trb/external-model-benchmark/runtime-v2" +LOCKED="$REPO/experiments/trb/external-model-benchmark/cohort" +REDCEA="$REPO/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" +OUT="$BASE/vdjdb/representations" + +export PYTHONPATH="$RUNTIME/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export HF_HOME="$RUNTIME/huggingface" +export TORCH_HOME="$RUNTIME/torch" +export PYTHONUNBUFFERED=1 +export TOKENIZERS_PARALLELISM=false +mkdir -p "$BASE/logs" "$OUT" +cd "$REPO" + +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader +"$PYTHON" benchmark/cache_external_model_embeddings.py --dataset-dir "$LOCKED" --dataset-file "$REDCEA" --output-dir "$OUT" --encoders sceptr_cdr3 --device cuda --batch-size 256 --threads "${SLURM_CPUS_PER_TASK:-8}" +sha256sum "$OUT/sceptr_cdr3.npy" +echo "status=accepted" diff --git a/slurm/cache_external_model_embeddings.sbatch b/slurm/cache_external_model_embeddings.sbatch new file mode 100644 index 0000000..9ef52ae --- /dev/null +++ b/slurm/cache_external_model_embeddings.sbatch @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-cache +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=06:00:00 +#SBATCH --array=0-2%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/cache-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/cache-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +RUNTIME="$REPO/experiments/trb/external-model-benchmark/runtime-v2" +COHORT="$REPO/experiments/trb/external-model-benchmark/cohort" +ENCODERS=(esm2_8m tcr_bert sceptr) +ENCODER="${ENCODERS[$SLURM_ARRAY_TASK_ID]}" + +export PYTHONPATH="$RUNTIME/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export HF_HOME="$RUNTIME/huggingface" +export TORCH_HOME="$RUNTIME/torch" +export PYTHONUNBUFFERED=1 +export TOKENIZERS_PARALLELISM=false +export HF_HUB_DISABLE_TELEMETRY=1 +mkdir -p "$REPO/experiments/trb/external-model-benchmark/logs" +cd "$REPO" + +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader +echo "encoder=$ENCODER python=$PYTHON" +"$PYTHON" benchmark/cache_external_model_embeddings.py \ + --dataset-dir "$COHORT" \ + --output-dir "$COHORT/representations" \ + --encoders "$ENCODER" \ + --rtp-embeddings "$COHORT/representations/rtp.npy" \ + --device cuda \ + --batch-size 256 \ + --threads "${SLURM_CPUS_PER_TASK:-8}" + +sha256sum "$COHORT/representations/$ENCODER.npy" +echo "status=accepted encoder=$ENCODER" diff --git a/slurm/cache_external_residue_embeddings.sbatch b/slurm/cache_external_residue_embeddings.sbatch new file mode 100644 index 0000000..3180f62 --- /dev/null +++ b/slurm/cache_external_residue_embeddings.sbatch @@ -0,0 +1,37 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-cache +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=06:00:00 +#SBATCH --array=0-2%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/cache-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/cache-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-v1" +RUNTIME="$BASE/runtime-v2" +ENCODERS=(esm2_8m tcr_bert sceptr) +ENCODER="${ENCODERS[$SLURM_ARRAY_TASK_ID]}" +export PYTHONPATH="$RUNTIME/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export HF_HOME="$RUNTIME/huggingface" +export TORCH_HOME="$RUNTIME/torch" +export TOKENIZERS_PARALLELISM=false +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/representations" +cd "$REPO" +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader +"$PYTHON" benchmark/cache_external_residue_embeddings.py \ + --dataset-dir "$BASE/cohort" \ + --output-dir "$ROOT/representations" \ + --encoders "$ENCODER" \ + --device cuda \ + --batch-size 128 \ + --threads "${SLURM_CPUS_PER_TASK:-8}" +sha256sum "$ROOT/representations/${ENCODER}_residue.npy" "$ROOT/representations/${ENCODER}_residue_mask.npy" +echo "status=accepted encoder=$ENCODER" diff --git a/slurm/collect_author_global_comparison.sbatch b/slurm/collect_author_global_comparison.sbatch new file mode 100644 index 0000000..b79e414 --- /dev/null +++ b/slurm/collect_author_global_comparison.sbatch @@ -0,0 +1,25 @@ +#!/bin/bash +#SBATCH --job-name=irrm-author-global-final +#SBATCH --partition=medium +#SBATCH --cpus-per-task=4 +#SBATCH --mem=24G +#SBATCH --time=01:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/final-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/final-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +OLD="$REPO/experiments/trb/external-model-benchmark" +BASE="$REPO/experiments/trb/external-global-author-v1" +VDJDB_EVAL="${1:?pass completed VDJdb evaluation path}" +OUT="$BASE/final-$SLURM_JOB_ID" + +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$BASE/logs" +cd "$REPO" + +"$PYTHON" benchmark/collect_author_global_comparison.py --old-reconstruction-summary "$OLD/results/reconstruction/reconstruction_summary.csv" --old-pgen-summary "$OLD/results/pgen/pgen_summary.csv" --new-reconstruction-dir "$BASE/reconstruction" --new-pgen-dir "$BASE/pgen" --control-pgen-dir "$BASE/pgen-onehot-fixed" --reference-vdjdb-metrics "$REPO/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728/metrics.json" --external-vdjdb-summary "$VDJDB_EVAL/external_retrieval_summary.tsv" --output-dir "$OUT" +sha256sum "$OUT/final_author_global_comparison.tsv" "$OUT/final_author_global_comparison.md" "$OUT/MANIFEST.json" > "$OUT/HASHES.sha256" +echo "status=complete output=$OUT" diff --git a/slurm/collect_external_model_benchmark.sbatch b/slurm/collect_external_model_benchmark.sbatch new file mode 100644 index 0000000..06a1ded --- /dev/null +++ b/slurm/collect_external_model_benchmark.sbatch @@ -0,0 +1,29 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-collect +#SBATCH --partition=medium +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=01:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/collect-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/collect-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/results/reconstruction" "$ROOT/results/pgen" +cd "$REPO" + +"$PYTHON" benchmark/collect_external_reconstruction.py \ + --runs-dir "$ROOT/reconstruction" \ + --bottleneck-dir "$ROOT/bottleneck" \ + --bottleneck-dim 64 \ + --output-dir "$ROOT/results/reconstruction" +"$PYTHON" benchmark/collect_external_pgen.py \ + --runs-dir "$ROOT/pgen" \ + --output-dir "$ROOT/results/pgen" +find "$ROOT/results" -type f ! -name HASHES.sha256 ! -name HASHES.sha256.tmp -print0 | sort -z | xargs -0 sha256sum > "$ROOT/results/HASHES.sha256.tmp" +mv "$ROOT/results/HASHES.sha256.tmp" "$ROOT/results/HASHES.sha256" +echo "status=accepted" diff --git a/slurm/collect_external_pca256_pgen.sbatch b/slurm/collect_external_pca256_pgen.sbatch new file mode 100644 index 0000000..b1e8596 --- /dev/null +++ b/slurm/collect_external_pca256_pgen.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-pca256-pgen-sum +#SBATCH --partition=short +#SBATCH --cpus-per-task=2 +#SBATCH --mem=4G +#SBATCH --time=00:20:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/pgen-summary-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/pgen-summary-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark/residue-pca256-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +mkdir -p "$ROOT/results/pgen" +cd "$REPO" +"$PYTHON" benchmark/collect_external_pgen.py --runs-dir "$ROOT/pgen" --output-dir "$ROOT/results/pgen" +sha256sum "$ROOT/results/pgen"/* > "$ROOT/results/pgen/HASHES.sha256" +echo status=accepted diff --git a/slurm/collect_external_pca256_reconstruction.sbatch b/slurm/collect_external_pca256_reconstruction.sbatch new file mode 100644 index 0000000..d4c9cdb --- /dev/null +++ b/slurm/collect_external_pca256_reconstruction.sbatch @@ -0,0 +1,21 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-pca256-recon-sum +#SBATCH --partition=short +#SBATCH --cpus-per-task=2 +#SBATCH --mem=4G +#SBATCH --time=00:20:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/recon-summary-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/recon-summary-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark/residue-pca256-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +mkdir -p "$ROOT/results/reconstruction" +cd "$REPO" +"$PYTHON" benchmark/collect_external_reconstruction.py \ + --runs-dir "$ROOT/reconstruction" --output-dir "$ROOT/results/reconstruction" \ + --bottleneck-dir "$ROOT/bottleneck" --bottleneck-dim 256 +sha256sum "$ROOT/results/reconstruction"/* > "$ROOT/results/reconstruction/HASHES.sha256" +echo status=accepted diff --git a/slurm/collect_external_residue_benchmark.sbatch b/slurm/collect_external_residue_benchmark.sbatch new file mode 100644 index 0000000..a0c4636 --- /dev/null +++ b/slurm/collect_external_residue_benchmark.sbatch @@ -0,0 +1,25 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-collect +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=2 +#SBATCH --mem=8G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/collect-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/collect-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark/residue-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/results" +cd "$REPO" +"$PYTHON" benchmark/collect_external_residue_benchmark.py \ + --reconstruction-dir "$ROOT/reconstruction" \ + --pgen-dir "$ROOT/pgen" \ + --output-dir "$ROOT/results" +find "$ROOT/results" -maxdepth 1 -type f ! -name HASHES.sha256 -print0 | sort -z | xargs -0 sha256sum > "$ROOT/results/HASHES.sha256" +cat "$ROOT/results/HASHES.sha256" +echo "status=complete" diff --git a/slurm/collect_external_residue_pca128_all.sbatch b/slurm/collect_external_residue_pca128_all.sbatch new file mode 100644 index 0000000..6fbc049 --- /dev/null +++ b/slurm/collect_external_residue_pca128_all.sbatch @@ -0,0 +1,28 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-pca128-sum +#SBATCH --partition=short +#SBATCH --cpus-per-task=2 +#SBATCH --mem=4G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/summary-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/summary-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca128-v1" +SCEPTR_ROOT="$BASE/residue-pca128-sceptr-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +mkdir -p "$ROOT/results/reconstruction" "$ROOT/results/pgen" +cd "$REPO" +"$PYTHON" benchmark/collect_external_reconstruction.py \ + --runs-dir "$ROOT/reconstruction" --output-dir "$ROOT/results/reconstruction" \ + --bottleneck-dir "$ROOT/bottleneck" --bottleneck-dim 128 +"$PYTHON" benchmark/collect_external_pgen.py \ + --runs-dir "$ROOT/pgen" --output-dir "$ROOT/results/pgen" +"$PYTHON" benchmark/combine_external_pca128_results.py \ + --main-root "$ROOT" --sceptr-root "$SCEPTR_ROOT" +sha256sum "$ROOT/results/reconstruction"/* "$ROOT/results/pgen"/* \ + "$ROOT/results/all_baselines"/* > "$ROOT/results/HASHES.sha256" +echo status=accepted diff --git a/slurm/collect_external_rtp128_reconstruction.sbatch b/slurm/collect_external_rtp128_reconstruction.sbatch new file mode 100644 index 0000000..c739dcf --- /dev/null +++ b/slurm/collect_external_rtp128_reconstruction.sbatch @@ -0,0 +1,72 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rtp128-sum +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=2 +#SBATCH --mem=8G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/native-rtp128-v1/logs/summary-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/native-rtp128-v1/logs/summary-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/native-rtp128-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +mkdir -p "$ROOT/results/reconstruction" +cd "$REPO" + +"$PYTHON" benchmark/collect_external_reconstruction.py \ + --runs-dir "$ROOT/reconstruction" \ + --output-dir "$ROOT/results/reconstruction" \ + --bottleneck-dir "$ROOT/bottleneck" \ + --bottleneck-dim 128 + +"$PYTHON" - "$ROOT" <<'PY' +import hashlib +import json +import sys +from pathlib import Path + +import pandas as pd + +root = Path(sys.argv[1]) +runs = pd.read_csv(root / "results/reconstruction/reconstruction_runs.csv") +summary = pd.read_csv(root / "results/reconstruction/reconstruction_summary.csv") +assert len(runs) == 3 and set(runs.seed) == {42, 43, 44} +assert runs.representation.eq("rtp").all() and runs.bottleneck_dim.eq(128).all() +assert len(summary) == 1 and summary.loc[0, "representation"] == "rtp" +assert int(summary.loc[0, "seeds"]) == 3 + +def sha256(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + +artifacts = {} +for path in sorted((root / "results/reconstruction").glob("*")): + if path.is_file(): + artifacts[str(path)] = sha256(path) +for path in sorted((root / "reconstruction").glob("*/metrics.json")): + artifacts[str(path)] = sha256(path) +report = { + "status": "accepted", + "interface": "native_rtp_128d", + "seeds": [42, 43, 44], + "pgen_reference": { + "note": "not rerun; previously completed native RTP-128 external MLP", + "rmse_mean": 0.2375036700, + "rmse_std": 0.0007641427, + }, + "reconstruction": summary.iloc[0].to_dict(), + "artifacts": artifacts, +} +(root / "RESULTS.json").write_text(json.dumps(report, indent=2), encoding="utf-8") +print(json.dumps(report, indent=2)) +PY + +sha256sum "$ROOT/RESULTS.json" "$ROOT/results/reconstruction"/* > "$ROOT/RESULT_HASHES.sha256" +if grep -R -E -i 'traceback|out of memory|nan|inf' "$ROOT/logs"/*.err; then + echo "error/nonfinite signature found" >&2 + exit 1 +fi +echo status=accepted diff --git a/slurm/collect_external_sceptr_pca128.sbatch b/slurm/collect_external_sceptr_pca128.sbatch new file mode 100644 index 0000000..6f9c995 --- /dev/null +++ b/slurm/collect_external_sceptr_pca128.sbatch @@ -0,0 +1,24 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-sceptr-pca128-sum +#SBATCH --partition=short +#SBATCH --cpus-per-task=2 +#SBATCH --mem=4G +#SBATCH --time=00:20:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/summary-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/summary-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +mkdir -p "$ROOT/results/reconstruction" "$ROOT/results/pgen" +cd "$REPO" +"$PYTHON" benchmark/collect_external_reconstruction.py \ + --runs-dir "$ROOT/reconstruction" --output-dir "$ROOT/results/reconstruction" \ + --bottleneck-dir "$ROOT/bottleneck" --bottleneck-dim 128 +"$PYTHON" benchmark/collect_external_pgen.py \ + --runs-dir "$ROOT/pgen" --output-dir "$ROOT/results/pgen" +sha256sum "$ROOT/results/reconstruction"/* "$ROOT/results/pgen"/* \ + > "$ROOT/results/HASHES.sha256" +echo status=accepted diff --git a/slurm/diagnose_vdjdb_decoded_tcremp_geometry.sbatch b/slurm/diagnose_vdjdb_decoded_tcremp_geometry.sbatch new file mode 100644 index 0000000..037e830 --- /dev/null +++ b/slurm/diagnose_vdjdb_decoded_tcremp_geometry.sbatch @@ -0,0 +1,10 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-tcremp-geom +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --time=04:00:00 +set -euo pipefail +cd /home/evlasova/irrm-codec +PYTHONPATH=/home/evlasova/irrm-codec /home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/diagnose_vdjdb_decoded_tcremp_geometry.py \ No newline at end of file diff --git a/slurm/evaluate_latent_secondary.sbatch b/slurm/evaluate_latent_secondary.sbatch new file mode 100644 index 0000000..cbed559 --- /dev/null +++ b/slurm/evaluate_latent_secondary.sbatch @@ -0,0 +1,64 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-latent-secondary +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=32 +#SBATCH --mem=128G +#SBATCH --time=16:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +MODE="${1:?mode must be preflight or evaluate}" +PREFLIGHT_ROOT="${2:-}" +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +BASE="$ROOT/experiments/trb/latent-secondary" +OUTPUT="$BASE/${MODE}-${SLURM_JOB_ID}" +DATA="$ROOT/data/benchmark/trb" +TOKENIZER="$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +export OMP_NUM_THREADS="${SLURM_CPUS_PER_TASK}" +export MKL_NUM_THREADS="${SLURM_CPUS_PER_TASK}" + +test -x "$PYTHON" +test -f "$DATA/READY.json" +test -f "$TOKENIZER" +test ! -e "$OUTPUT" +mkdir -p "$BASE" +cd "$ROOT" + +COMMON=( + --data-dir "$DATA" + --tokenizer "$TOKENIZER" + --output-dir "$OUTPUT" + --model "r:$ROOT/experiments/trb/data-anchor-r-seed42/best.pt:$ROOT/experiments/trb/data-anchor-r-seed42/run_config.json" + --model "t:$ROOT/experiments/trb/data-anchor-t-seed42/best.pt:$ROOT/experiments/trb/data-anchor-t-seed42/run_config.json" + --model "p:$ROOT/experiments/trb/data-anchor-p-seed42/best.pt:$ROOT/experiments/trb/data-anchor-p-seed42/run_config.json" + --model "rt:$ROOT/experiments/trb/data-anchor-rt-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rt-seed42/run_config.json" + --model "rp:$ROOT/experiments/trb/data-anchor-rp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rp-seed42/run_config.json" + --model "tp:$ROOT/experiments/trb/data-anchor-tp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-tp-seed42/run_config.json" + --model "rtp:$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt:$ROOT/experiments/trb/anchor-data-rtp-seed42/run_config.json" + --batch-size 128 + --seed 42 +) + +if [[ "$MODE" == "preflight" ]]; then + "$PYTHON" -m compileall -q benchmark/evaluate_latent_secondary.py tests/test_latent_secondary.py + "$PYTHON" - <<'PY' +import runpy +namespace = runpy.run_path('tests/test_latent_secondary.py') +for name, value in sorted(namespace.items()): + if name.startswith('test_') and callable(value): + value() +print('direct tests passed') +PY + "$PYTHON" -m benchmark.evaluate_latent_secondary "${COMMON[@]}" --preflight-only +elif [[ "$MODE" == "evaluate" ]]; then + test -n "$PREFLIGHT_ROOT" + test -f "$PREFLIGHT_ROOT/PREFLIGHT.json" + "$PYTHON" -m benchmark.evaluate_latent_secondary \ + "${COMMON[@]}" --accepted-preflight "$PREFLIGHT_ROOT/PREFLIGHT.json" +else + echo "unsupported mode: $MODE" >&2 + exit 2 +fi diff --git a/slurm/evaluate_latent_sufficiency.sbatch b/slurm/evaluate_latent_sufficiency.sbatch new file mode 100644 index 0000000..2fff4b4 --- /dev/null +++ b/slurm/evaluate_latent_sufficiency.sbatch @@ -0,0 +1,55 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-latent-suff +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=32 +#SBATCH --mem=192G +#SBATCH --time=16:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +BASE="$ROOT/experiments/trb/latent-sufficiency" +OUTPUT="$BASE/evaluation-${SLURM_JOB_ID}" +DATA="$ROOT/data/benchmark/trb" +TOKENIZER="$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +export OMP_NUM_THREADS="${SLURM_CPUS_PER_TASK}" +export MKL_NUM_THREADS="${SLURM_CPUS_PER_TASK}" + +test -x "$PYTHON" +test -f "$DATA/READY.json" +test -f "$DATA/target_standardizer.npz" +test -f "$TOKENIZER" +test ! -e "$OUTPUT" +mkdir -p "$BASE" +cd "$ROOT" + +"$PYTHON" -m compileall -q benchmark/evaluate_latent_sufficiency.py tests/test_latent_sufficiency.py +"$PYTHON" - <<'PY' +import runpy +scope = runpy.run_path("tests/test_latent_sufficiency.py") +for name in sorted(scope): + if name.startswith("test_"): + scope[name]() +print("direct latent-sufficiency tests passed") +PY + +"$PYTHON" -m benchmark.evaluate_latent_sufficiency \ + --data-dir "$DATA" \ + --normalizer "$DATA/target_standardizer.npz" \ + --tokenizer "$TOKENIZER" \ + --output-dir "$OUTPUT" \ + --model "r:$ROOT/experiments/trb/data-anchor-r-seed42/best.pt:$ROOT/experiments/trb/data-anchor-r-seed42/run_config.json" \ + --model "t:$ROOT/experiments/trb/data-anchor-t-seed42/best.pt:$ROOT/experiments/trb/data-anchor-t-seed42/run_config.json" \ + --model "p:$ROOT/experiments/trb/data-anchor-p-seed42/best.pt:$ROOT/experiments/trb/data-anchor-p-seed42/run_config.json" \ + --model "rt:$ROOT/experiments/trb/data-anchor-rt-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rt-seed42/run_config.json" \ + --model "rp:$ROOT/experiments/trb/data-anchor-rp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rp-seed42/run_config.json" \ + --model "tp:$ROOT/experiments/trb/data-anchor-tp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-tp-seed42/run_config.json" \ + --model "rtp:$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt:$ROOT/experiments/trb/anchor-data-rtp-seed42/run_config.json" \ + --batch-size 256 --alphas 0.01,1,100 --bootstrap 2000 --seed 42 + +test -s "$OUTPUT/RESULTS.json" +test -s "$OUTPUT/summary.tsv" +test -s "$OUTPUT/OUTPUT_HASHES.json" diff --git a/slurm/evaluate_rtp_data_anchor_length_errors.sbatch b/slurm/evaluate_rtp_data_anchor_length_errors.sbatch new file mode 100644 index 0000000..94d1e13 --- /dev/null +++ b/slurm/evaluate_rtp_data_anchor_length_errors.sbatch @@ -0,0 +1,28 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rtp-lenerr +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +CHECKPOINT="$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/best.pt" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/rtp-data-anchor-length-errors-$SLURM_JOB_ID" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER" +test -f "$CHECKPOINT" +cd "$ROOT_DIR" +export PYTHONPATH="$ROOT_DIR${PYTHONPATH:+:$PYTHONPATH}" +"$PYTHON_BIN" benchmark/evaluate_rtp_reconstruction_by_length.py \ + --data-dir "$DATA_DIR" --tokenizer "$TOKENIZER" --checkpoint "$CHECKPOINT" \ + --output-dir "$OUTPUT_DIR" --batch-size 64 --num-workers 8 --device cuda diff --git a/slurm/evaluate_rtp_data_anchor_position_errors.sbatch b/slurm/evaluate_rtp_data_anchor_position_errors.sbatch new file mode 100644 index 0000000..c42323c --- /dev/null +++ b/slurm/evaluate_rtp_data_anchor_position_errors.sbatch @@ -0,0 +1,28 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rtp-poserr +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +CHECKPOINT="$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/best.pt" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/rtp-data-anchor-position-errors-$SLURM_JOB_ID" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER" +test -f "$CHECKPOINT" +cd "$ROOT_DIR" +export PYTHONPATH="$ROOT_DIR${PYTHONPATH:+:$PYTHONPATH}" +"$PYTHON_BIN" benchmark/evaluate_rtp_reconstruction_by_position.py \ + --data-dir "$DATA_DIR" --tokenizer "$TOKENIZER" --checkpoint "$CHECKPOINT" \ + --output-dir "$OUTPUT_DIR" --batch-size 64 --num-workers 8 --device cuda diff --git a/slurm/evaluate_rtp_data_anchor_substitutions.sbatch b/slurm/evaluate_rtp_data_anchor_substitutions.sbatch new file mode 100644 index 0000000..670c803 --- /dev/null +++ b/slurm/evaluate_rtp_data_anchor_substitutions.sbatch @@ -0,0 +1,33 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rtp-subst +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +CHECKPOINT="$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/best.pt" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/rtp-data-anchor-substitutions-$SLURM_JOB_ID" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER" +test -f "$CHECKPOINT" +cd "$ROOT_DIR" +"$PYTHON_BIN" benchmark/evaluate_reconstruction_substitutions.py \ + --data-dir "$DATA_DIR" \ + --tokenizer "$TOKENIZER" \ + --checkpoint "$CHECKPOINT" \ + --output-dir "$OUTPUT_DIR" \ + --batch-size 64 \ + --num-workers 8 \ + --device cuda diff --git a/slurm/evaluate_rtp_data_anchor_substitutions_v2.sbatch b/slurm/evaluate_rtp_data_anchor_substitutions_v2.sbatch new file mode 100644 index 0000000..8959f3a --- /dev/null +++ b/slurm/evaluate_rtp_data_anchor_substitutions_v2.sbatch @@ -0,0 +1,34 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rtp-subst +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +CHECKPOINT="$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/best.pt" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/rtp-data-anchor-substitutions-$SLURM_JOB_ID" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER" +test -f "$CHECKPOINT" +cd "$ROOT_DIR" +export PYTHONPATH="$ROOT_DIR${PYTHONPATH:+:$PYTHONPATH}" +"$PYTHON_BIN" benchmark/evaluate_reconstruction_substitutions.py \ + --data-dir "$DATA_DIR" \ + --tokenizer "$TOKENIZER" \ + --checkpoint "$CHECKPOINT" \ + --output-dir "$OUTPUT_DIR" \ + --batch-size 64 \ + --num-workers 8 \ + --device cuda diff --git a/slurm/evaluate_rtp_error_analysis.sbatch b/slurm/evaluate_rtp_error_analysis.sbatch new file mode 100644 index 0000000..762a559 --- /dev/null +++ b/slurm/evaluate_rtp_error_analysis.sbatch @@ -0,0 +1,14 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rtp-errana +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +cd "$ROOT_DIR"; export PYTHONPATH="$ROOT_DIR${PYTHONPATH:+:$PYTHONPATH}" +/home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/evaluate_rtp_error_analysis.py --data-dir "$ROOT_DIR/data/benchmark/trb" --tokenizer "$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" --checkpoint "$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/best.pt" --output-dir "$ROOT_DIR/experiments/trb/rtp-data-anchor-error-analysis-$SLURM_JOB_ID" --batch-size 64 --num-workers 8 --device cuda diff --git a/slurm/evaluate_single_aa_challenge.sbatch b/slurm/evaluate_single_aa_challenge.sbatch new file mode 100644 index 0000000..10d440d --- /dev/null +++ b/slurm/evaluate_single_aa_challenge.sbatch @@ -0,0 +1,32 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-aa1-eval +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=32 +#SBATCH --mem=192G +#SBATCH --time=16:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +TARGET_ROOT="${1:?Validated target job directory required}" +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +BASE="$ROOT/experiments/trb/single-aa-challenge" +OUT="$BASE/evaluation-${SLURM_JOB_ID}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +TOKENIZER="$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +test -f "$TARGET_ROOT/targets/TARGETS.json" +test ! -e "$OUT" +cd "$ROOT" +export OMP_NUM_THREADS="$SLURM_CPUS_PER_TASK" +export MKL_NUM_THREADS="$SLURM_CPUS_PER_TASK" +"$PYTHON" -m benchmark.evaluate_single_aa_challenge \ + --prepared "$TARGET_ROOT/prepared" --targets "$TARGET_ROOT/targets" \ + --tokenizer "$TOKENIZER" --output-dir "$OUT" --bootstrap 0 --batch-size 128 --seed 42 \ + --model "r:$ROOT/experiments/trb/data-anchor-r-seed42/best.pt:$ROOT/experiments/trb/data-anchor-r-seed42/run_config.json" \ + --model "t:$ROOT/experiments/trb/data-anchor-t-seed42/best.pt:$ROOT/experiments/trb/data-anchor-t-seed42/run_config.json" \ + --model "p:$ROOT/experiments/trb/data-anchor-p-seed42/best.pt:$ROOT/experiments/trb/data-anchor-p-seed42/run_config.json" \ + --model "rt:$ROOT/experiments/trb/data-anchor-rt-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rt-seed42/run_config.json" \ + --model "rp:$ROOT/experiments/trb/data-anchor-rp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rp-seed42/run_config.json" \ + --model "tp:$ROOT/experiments/trb/data-anchor-tp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-tp-seed42/run_config.json" \ + --model "rtp:$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt:$ROOT/experiments/trb/anchor-data-rtp-seed42/run_config.json" diff --git a/slurm/evaluate_tcremp_geometry_biology.sbatch b/slurm/evaluate_tcremp_geometry_biology.sbatch new file mode 100644 index 0000000..f8f1cf8 --- /dev/null +++ b/slurm/evaluate_tcremp_geometry_biology.sbatch @@ -0,0 +1,32 @@ +#!/bin/bash -l +#SBATCH --job-name=tcremp-geom-bio +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --time=16:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err +set -euo pipefail +VARIANT="${1:?Variant centered or relational required}" +case "$VARIANT" in centered|relational) ;; *) exit 2 ;; esac +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +MODEL="$ROOT/experiments/trb/tcremp-geometry-t-${VARIANT}-seed42" +OUT="$ROOT/experiments/trb/vdjdb-tcremp-geometry-${VARIANT}-${SLURM_JOB_ID}" +COHORT="$ROOT/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" +TOKENIZER="$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +test -f "$MODEL/best.pt" +mkdir -p "$OUT" +cd "$ROOT" +PYTHONPATH="$ROOT" "$PYTHON" benchmark/extract_vdjdb_latent.py \ + --cohort "$COHORT" --checkpoint "$MODEL/best.pt" --tokenizer "$TOKENIZER" \ + --output "$OUT/embeddings_t.npy" --batch-size 512 --device cuda +PYTHONPATH="$ROOT" "$PYTHON" benchmark/evaluate_vdjdb_decoded_tcremp.py \ + --cohort "$COHORT" --rtp "$OUT/embeddings_t.npy" \ + --checkpoint "$MODEL/best.pt" --standardizer "$MODEL/target_standardizer.npz" \ + --uniform-evaluation "$ROOT/experiments/trb/vdjdb-uniform100-v1/evaluation-1427813" \ + --existing-separation "$ROOT/experiments/trb/vdjdb-uniform100-v1/separation-1428106" \ + --native-tcremp "$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/tcremp.npy" \ + --output-dir "$OUT/decoded_eval" --batch-size 256 --device cuda diff --git a/slurm/evaluate_vdjdb_author_global.sbatch b/slurm/evaluate_vdjdb_author_global.sbatch new file mode 100644 index 0000000..6df8a01 --- /dev/null +++ b/slurm/evaluate_vdjdb_author_global.sbatch @@ -0,0 +1,26 @@ +#!/bin/bash +#SBATCH --job-name=irrm-author-global-vdjdb +#SBATCH --partition=medium +#SBATCH --cpus-per-task=16 +#SBATCH --mem=96G +#SBATCH --time=06:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/vdjdb-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/vdjdb-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-global-author-v1" +COHORT="$REPO/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" +REFERENCE="$REPO/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728" +GATE="${AUTHOR_GLOBAL_PREFLIGHT:-$BASE/preflight-1427529/PREFLIGHT.json}" +OUT="$BASE/vdjdb/evaluation-$SLURM_JOB_ID" + +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$BASE/logs" "$BASE/vdjdb" +cd "$REPO" + +"$PYTHON" benchmark/evaluate_vdjdb_author_global.py --cohort "$COHORT" --preflight "$GATE" --representations-dir "$BASE/vdjdb/representations" --reference-evaluation "$REFERENCE" --output-dir "$OUT" +sha256sum "$OUT/RESULTS.json" "$OUT/external_metrics.json" "$OUT/external_retrieval_summary.tsv" "$OUT/external_per_epitope.tsv" > "$OUT/HASHES.sha256" +echo "status=complete output=$OUT" diff --git a/slurm/evaluate_vdjdb_decoded_tcremp.sbatch b/slurm/evaluate_vdjdb_decoded_tcremp.sbatch new file mode 100644 index 0000000..a612a96 --- /dev/null +++ b/slurm/evaluate_vdjdb_decoded_tcremp.sbatch @@ -0,0 +1,21 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-decoded-tcremp +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --time=16:00:00 + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/vdjdb_decoded_tcremp/${SLURM_JOB_ID}" +cd "$ROOT" +PYTHONPATH="$ROOT" /home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/evaluate_vdjdb_decoded_tcremp.py \ + --cohort "$ROOT/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" \ + --rtp "$ROOT/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728/embeddings_rtp.npy" \ + --checkpoint "$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt" \ + --standardizer "$ROOT/experiments/trb/anchor-data-rtp-seed42/target_standardizer.npz" \ + --uniform-evaluation "$ROOT/experiments/trb/vdjdb-uniform100-v1/evaluation-1427813" \ + --existing-separation "$ROOT/experiments/trb/vdjdb-uniform100-v1/separation-1428106" \ + --native-tcremp "$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/tcremp.npy" \ + --output-dir "$OUT" --batch-size 256 diff --git a/slurm/evaluate_vdjdb_epitope_prediction.sbatch b/slurm/evaluate_vdjdb_epitope_prediction.sbatch new file mode 100644 index 0000000..3e6f023 --- /dev/null +++ b/slurm/evaluate_vdjdb_epitope_prediction.sbatch @@ -0,0 +1,42 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-epitope-prediction +#SBATCH --partition=medium +#SBATCH --constraint=hpc +#SBATCH --cpus-per-task=20 +#SBATCH --mem=160G +#SBATCH --time=16:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/epitope-prediction-v1/logs/run-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/epitope-prediction-v1/logs/run-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/epitope-prediction-v1" +COHORT="$REPO/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" +UNIFORM="$REPO/experiments/trb/vdjdb-uniform100-v1/evaluation-1427813" +NATIVE="$REPO/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728" +EXTERNAL="$REPO/experiments/trb/external-global-author-v1/vdjdb/representations" +OUT="${EPITOPE_PREDICTION_OUT:-$BASE/evaluation-$SLURM_JOB_ID}" + +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +mkdir -p "$BASE/logs" +cd "$REPO" + +"$PYTHON" benchmark/evaluate_vdjdb_epitope_prediction.py \ + --cohort "$COHORT" \ + --uniform-evaluation "$UNIFORM" \ + --native-evaluation "$NATIVE" \ + --external-representations "$EXTERNAL" \ + --output-dir "$OUT" \ + --seeds 42 43 44 \ + --folds 5 \ + --c-grid 0.1 1.0 \ + --max-iter 400 \ + --jobs 16 + +sha256sum "$OUT"/* > "$OUT/HASHES.sha256" +echo "status=complete output=$OUT" diff --git a/slurm/evaluate_vdjdb_frozen.sbatch b/slurm/evaluate_vdjdb_frozen.sbatch new file mode 100644 index 0000000..743a604 --- /dev/null +++ b/slurm/evaluate_vdjdb_frozen.sbatch @@ -0,0 +1,34 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-vdjdb-eval +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=08:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +VALIDATION_ROOT="$ROOT_DIR/experiments/trb/vdjdb-data-anchor-validation" +COHORT_DIR="$VALIDATION_ROOT/cohort-preflight-1426341" +OUTPUT_DIR="$VALIDATION_ROOT/evaluation-${SLURM_JOB_ID}" +TOKENIZER="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +for condition in r p rp; do + test -f "$ROOT_DIR/experiments/trb/data-anchor-$condition-seed42/best.pt" + test -f "$ROOT_DIR/experiments/trb/data-anchor-$condition-seed42/run_config.json" +done +test -f "$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/best.pt" +test -f "$COHORT_DIR/PREFLIGHT.json" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m benchmark.evaluate_vdjdb_frozen \ + --cohort "$COHORT_DIR/cohort.tsv" --cohort-preflight "$COHORT_DIR/PREFLIGHT.json" \ + --tokenizer "$TOKENIZER" --output-dir "$OUTPUT_DIR" \ + --model "r:$ROOT_DIR/experiments/trb/data-anchor-r-seed42/best.pt:$ROOT_DIR/experiments/trb/data-anchor-r-seed42/run_config.json" \ + --model "p:$ROOT_DIR/experiments/trb/data-anchor-p-seed42/best.pt:$ROOT_DIR/experiments/trb/data-anchor-p-seed42/run_config.json" \ + --model "rp:$ROOT_DIR/experiments/trb/data-anchor-rp-seed42/best.pt:$ROOT_DIR/experiments/trb/data-anchor-rp-seed42/run_config.json" \ + --model "rtp:$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/best.pt:$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/run_config.json" \ + --batch-size 128 --max-pairs-per-label 2000 --bootstrap 2000 --permutations 1000 --seed 42 diff --git a/slurm/evaluate_vdjdb_redcea_clonotypes.sbatch b/slurm/evaluate_vdjdb_redcea_clonotypes.sbatch new file mode 100644 index 0000000..8313c8e --- /dev/null +++ b/slurm/evaluate_vdjdb_redcea_clonotypes.sbatch @@ -0,0 +1,44 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-redcea-eval +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=32 +#SBATCH --mem=128G +#SBATCH --time=16:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +PREFLIGHT_ROOT="${1:?Preflight job directory is required}" +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +BASE="$ROOT/experiments/trb/vdjdb-redcea-clonotypes" +OUTPUT="$BASE/evaluation-${SLURM_JOB_ID}" +TOKENIZER="$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +export OMP_NUM_THREADS="${SLURM_CPUS_PER_TASK}" +export MKL_NUM_THREADS="${SLURM_CPUS_PER_TASK}" + +test -x "$PYTHON" +test -f "$PREFLIGHT_ROOT/cohort/PREFLIGHT.json" +test -f "$PREFLIGHT_ROOT/cohort/cohort.tsv" +test -f "$TOKENIZER" +test ! -e "$OUTPUT" +cd "$ROOT" +"$PYTHON" -m benchmark.evaluate_vdjdb_redcea_clonotypes \ + --cohort "$PREFLIGHT_ROOT/cohort/cohort.tsv" \ + --preflight "$PREFLIGHT_ROOT/cohort/PREFLIGHT.json" \ + --tokenizer "$TOKENIZER" \ + --output-dir "$OUTPUT" \ + --model "r:$ROOT/experiments/trb/data-anchor-r-seed42/best.pt:$ROOT/experiments/trb/data-anchor-r-seed42/run_config.json" \ + --model "t:$ROOT/experiments/trb/data-anchor-t-seed42/best.pt:$ROOT/experiments/trb/data-anchor-t-seed42/run_config.json" \ + --model "p:$ROOT/experiments/trb/data-anchor-p-seed42/best.pt:$ROOT/experiments/trb/data-anchor-p-seed42/run_config.json" \ + --model "rt:$ROOT/experiments/trb/data-anchor-rt-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rt-seed42/run_config.json" \ + --model "rp:$ROOT/experiments/trb/data-anchor-rp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rp-seed42/run_config.json" \ + --model "tp:$ROOT/experiments/trb/data-anchor-tp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-tp-seed42/run_config.json" \ + --model "rtp:$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt:$ROOT/experiments/trb/anchor-data-rtp-seed42/run_config.json" \ + --batch-size 128 \ + --max-positive-candidates 20 \ + --negative-ratio 4 \ + --bootstrap 2000 \ + --silhouette-sample 10000 \ + --seed 42 diff --git a/slurm/evaluate_vdjdb_tonly_decoded_tcremp.sbatch b/slurm/evaluate_vdjdb_tonly_decoded_tcremp.sbatch new file mode 100644 index 0000000..c04c729 --- /dev/null +++ b/slurm/evaluate_vdjdb_tonly_decoded_tcremp.sbatch @@ -0,0 +1,20 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-tonly-tcremp +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --time=16:00:00 +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/vdjdb_tonly_decoded_tcremp/${SLURM_JOB_ID}" +cd "$ROOT" +PYTHONPATH="$ROOT" /home/evlasova/.conda/envs/irrm-codec/bin/python benchmark/evaluate_vdjdb_decoded_tcremp.py \ + --cohort "$ROOT/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" \ + --rtp "$ROOT/experiments/trb/vdjdb-redcea-clonotypes/evaluation-1426728/embeddings_t.npy" \ + --checkpoint "$ROOT/experiments/trb/data-anchor-t-seed42/best.pt" \ + --standardizer "$ROOT/experiments/trb/data-anchor-t-seed42/target_standardizer.npz" \ + --uniform-evaluation "$ROOT/experiments/trb/vdjdb-uniform100-v1/evaluation-1427813" \ + --existing-separation "$ROOT/experiments/trb/vdjdb-uniform100-v1/separation-1428106" \ + --native-tcremp "$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/tcremp.npy" \ + --output-dir "$OUT" --batch-size 256 \ No newline at end of file diff --git a/slurm/evaluate_vdjdb_uniform100_tcremp.sbatch b/slurm/evaluate_vdjdb_uniform100_tcremp.sbatch new file mode 100644 index 0000000..b7993da --- /dev/null +++ b/slurm/evaluate_vdjdb_uniform100_tcremp.sbatch @@ -0,0 +1,45 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-tcremp +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=32 +#SBATCH --mem=96G +#SBATCH --time=16:00:00 + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-${SLURM_JOB_ID}" +COHORT="$ROOT/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" +PAIRS="$ROOT/experiments/trb/vdjdb-uniform100-v1/evaluation-1427813/uniform_pairs.npz" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +TCREMP_PYTHON=/home/evlasova/.conda/envs/vdjdb-redcea/bin/python +TCREMP_ROOT=/home/evlasova/tcremp-git/tcremp +mkdir -p "$OUT/raw_tcremp" +cd "$ROOT" +"$PYTHON" - "$COHORT" "$OUT/cohort.airr.tsv" <<'PY' +import sys +import pandas as pd +x = pd.read_csv(sys.argv[1], sep="\t") +pd.DataFrame({"clone_id": range(len(x)), "junction_aa": x.cdr3, "v_call": x.v_call, "j_call": x.j_call, "locus": "TRB"}).to_csv(sys.argv[2], sep="\t", index=False) +PY +PYTHONPATH="$TCREMP_ROOT:$ROOT" "$TCREMP_PYTHON" -m tcremp.tcremp_run -i "$OUT/cohort.airr.tsv" -o "$OUT/raw_tcremp" -e cohort -c TRB -s HomoSapiens -x clone_id -p "$TCREMP_ROOT/tcremp/resources/tcremp_prototypes_TRB.tsv" -np "$SLURM_CPUS_PER_TASK" -llen 4 -hlen 41 --skip-clustering +"$PYTHON" - "$COHORT" "$PAIRS" "$OUT/raw_tcremp/cohort_tcremp.parquet" "$OUT" <<'PY' +import json, sys +from pathlib import Path +import numpy as np +import pandas as pd +from benchmark.evaluate_vdjdb_uniform100_separation import separation_metrics +cohort, pairs, raw, out = map(Path, sys.argv[1:]) +x, p, r = pd.read_csv(cohort, sep="\t"), np.load(pairs), pd.read_parquet(raw) +v = r.drop(columns="clone_id").select_dtypes(include=[np.number]); v.index = r.clone_id.astype(int) +if v.shape[1] != 9000 or not v.index.is_unique or set(v.index) != set(range(len(x))): raise ValueError("Invalid TCRemP alignment") +m = v.loc[np.arange(len(x))].to_numpy(np.float32) +if not np.isfinite(m).all(): raise ValueError("Non-finite TCRemP") +np.save(out / "tcremp.npy", m, allow_pickle=False) +q, e = separation_metrics(m[p["original_indices"]], p["within"], p["between"], x.iloc[p["original_indices"]].label.astype(str).to_numpy()) +e.insert(0, "model", "TCRemP prototypes") +e.to_csv(out / "per_epitope.tsv", sep="\t", index=False) +s = {"model": "TCRemP prototypes", "macro_auroc": float(e.auroc.mean()), "macro_cohen_d": float(e.cohen_d.mean()), "queries": len(q), "epitopes": len(e)} +pd.DataFrame([s]).to_csv(out / "summary.tsv", sep="\t", index=False) +(out / "RESULTS.json").write_text(json.dumps({"status": "complete", "summary": s, "matrix_shape": list(m.shape), "protocol": "cosine distance on Fig. 2 fixed uniform-100 pairs"}, indent=2) + "\n") +PY diff --git a/slurm/evaluate_vdjdb_uniform100_tcremp_pca128.sbatch b/slurm/evaluate_vdjdb_uniform100_tcremp_pca128.sbatch new file mode 100644 index 0000000..2df132b --- /dev/null +++ b/slurm/evaluate_vdjdb_uniform100_tcremp_pca128.sbatch @@ -0,0 +1,61 @@ +#!/bin/bash -l +#SBATCH --job-name=vdjdb-tcremp128 +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=32 +#SBATCH --mem=192G +#SBATCH --time=16:00:00 + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-pca128-${SLURM_JOB_ID}" +TRAIN_DATA="$ROOT/experiments/trb/external-model-benchmark/cohort" +TRAIN_TCREMP="$TRAIN_DATA/representations/tcremp.npy" +VDJDB_TCREMP="$ROOT/experiments/trb/vdjdb-uniform100-v1/tcremp-1429926/tcremp.npy" +COHORT="$ROOT/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" +PAIRS="$ROOT/experiments/trb/vdjdb-uniform100-v1/evaluation-1427813/uniform_pairs.npz" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +mkdir -p "$OUT" +cd "$ROOT" +"$PYTHON" - "$TRAIN_TCREMP" "$TRAIN_DATA/manifests/train.tsv" "$VDJDB_TCREMP" "$COHORT" "$PAIRS" "$OUT" <<'PY' +import hashlib +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +from joblib import dump +from benchmark.evaluate_vdjdb_uniform100_separation import separation_metrics +from benchmark.fit_external_model_bottleneck import fit_projection, transform_all + +train_matrix_path, train_rows_path, vdjdb_matrix_path, cohort_path, pairs_path, out = map(Path, sys.argv[1:]) +out.mkdir(parents=True, exist_ok=True) +train_matrix = np.load(train_matrix_path, mmap_mode="r") +train_rows = pd.read_csv(train_rows_path, sep="\t")["row_index"].to_numpy(dtype=int) +vdjdb = np.load(vdjdb_matrix_path, mmap_mode="r") +if train_matrix.shape[1] != 9000 or vdjdb.shape != (68302, 9000): + raise ValueError(f"Unexpected TCRemP shapes: train={train_matrix.shape}, VDJdb={vdjdb.shape}") +class Log: + def info(self, *args): print(*args, flush=True) +scaler, pca = fit_projection(train_matrix, train_rows, 128, 4096, 42, Log()) +dump({"scaler": scaler, "pca": pca, "fit_rows": len(train_rows), "source": str(train_matrix_path)}, out / "tcremp_pca128.joblib") +projected = transform_all(vdjdb, scaler, pca, 4096) +np.save(out / "tcremp_pca128.npy", projected, allow_pickle=False) +cohort = pd.read_csv(cohort_path, sep="\t") +pairs = np.load(pairs_path) +original = pairs["original_indices"] +query, epitope = separation_metrics(projected[original], pairs["within"], pairs["between"], cohort.iloc[original].label.astype(str).to_numpy()) +query.insert(0, "model", "TCRemP prototypes PCA-128") +epitope.insert(0, "model", "TCRemP prototypes PCA-128") +query.to_parquet(out / "per_query.parquet", index=False) +epitope.to_csv(out / "per_epitope.tsv", sep="\t", index=False, lineterminator="\n") +summary = {"model": "TCRemP prototypes PCA-128", "macro_auroc": float(epitope.auroc.mean()), "macro_cohen_d": float(epitope.cohen_d.mean()), "queries": len(query), "epitopes": len(epitope)} +pd.DataFrame([summary]).to_csv(out / "summary.tsv", sep="\t", index=False, lineterminator="\n") +def sha256(path): + h = hashlib.sha256() + with open(path, "rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): h.update(block) + return h.hexdigest() +(out / "RESULTS.json").write_text(json.dumps({"status": "complete", "summary": summary, "projection": {"method": "StandardScaler plus IncrementalPCA", "output_dim": 128, "fit_data": str(train_matrix_path), "fit_rows": len(train_rows), "fit_split": str(train_rows_path), "seed": 42}, "source_tcremp": {"path": str(vdjdb_matrix_path), "sha256": sha256(vdjdb_matrix_path), "shape": list(vdjdb.shape)}, "protocol": "cosine distance on the identical Fig. 2 fixed uniform-100 pairs"}, indent=2, sort_keys=True) + "\n") +PY diff --git a/slurm/evaluate_vdjdb_viral_motif_top10.sbatch b/slurm/evaluate_vdjdb_viral_motif_top10.sbatch new file mode 100644 index 0000000..51de588 --- /dev/null +++ b/slurm/evaluate_vdjdb_viral_motif_top10.sbatch @@ -0,0 +1,28 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-viral10-eval +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +PREFLIGHT_DIR="${1:?Top-10 preflight directory is required}" +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUTPUT="$ROOT/experiments/trb/vdjdb-viral-motif-top10/evaluation-${SLURM_JOB_ID}" +TOKENIZER="$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +test -f "$PREFLIGHT_DIR/PREFLIGHT.json" +cd "$ROOT" +"$PYTHON" -m benchmark.evaluate_vdjdb_viral_motif_top10 \ + --preflight-dir "$PREFLIGHT_DIR" --tokenizer "$TOKENIZER" --output-dir "$OUTPUT" \ + --model "r:$ROOT/experiments/trb/data-anchor-r-seed42/best.pt:$ROOT/experiments/trb/data-anchor-r-seed42/run_config.json" \ + --model "t:$ROOT/experiments/trb/data-anchor-t-seed42/best.pt:$ROOT/experiments/trb/data-anchor-t-seed42/run_config.json" \ + --model "p:$ROOT/experiments/trb/data-anchor-p-seed42/best.pt:$ROOT/experiments/trb/data-anchor-p-seed42/run_config.json" \ + --model "rt:$ROOT/experiments/trb/data-anchor-rt-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rt-seed42/run_config.json" \ + --model "rp:$ROOT/experiments/trb/data-anchor-rp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rp-seed42/run_config.json" \ + --model "tp:$ROOT/experiments/trb/data-anchor-tp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-tp-seed42/run_config.json" \ + --model "rtp:$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt:$ROOT/experiments/trb/anchor-data-rtp-seed42/run_config.json" \ + --batch-size 128 --bootstrap 2000 --permutations 1000 --seed 42 diff --git a/slurm/evaluate_vdjdb_ylq_frozen.sbatch b/slurm/evaluate_vdjdb_ylq_frozen.sbatch new file mode 100644 index 0000000..49fa542 --- /dev/null +++ b/slurm/evaluate_vdjdb_ylq_frozen.sbatch @@ -0,0 +1,25 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-ylq-eval +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=08:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +PREFLIGHT_DIR="${1:?YLQ preflight directory is required}" +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUTPUT="$ROOT/experiments/trb/vdjdb-ylq-data-anchor-validation/evaluation-${SLURM_JOB_ID}" +TOKENIZER="$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +test -f "$PREFLIGHT_DIR/PREFLIGHT.json" +cd "$ROOT" +"$PYTHON" -m benchmark.evaluate_vdjdb_ylq_frozen \ + --preflight-dir "$PREFLIGHT_DIR" --tokenizer "$TOKENIZER" --output-dir "$OUTPUT" \ + --model "r:$ROOT/experiments/trb/data-anchor-r-seed42/best.pt:$ROOT/experiments/trb/data-anchor-r-seed42/run_config.json" \ + --model "p:$ROOT/experiments/trb/data-anchor-p-seed42/best.pt:$ROOT/experiments/trb/data-anchor-p-seed42/run_config.json" \ + --model "rp:$ROOT/experiments/trb/data-anchor-rp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rp-seed42/run_config.json" \ + --model "rtp:$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt:$ROOT/experiments/trb/anchor-data-rtp-seed42/run_config.json" \ + --bootstrap 2000 --permutations 1000 --seed 42 diff --git a/slurm/evaluate_vdjdb_ylq_motif_pooled.sbatch b/slurm/evaluate_vdjdb_ylq_motif_pooled.sbatch new file mode 100644 index 0000000..eab1504 --- /dev/null +++ b/slurm/evaluate_vdjdb_ylq_motif_pooled.sbatch @@ -0,0 +1,28 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-ylq-motif-eval +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=08:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +PREFLIGHT_DIR="${1:?Official-motif pooled preflight directory is required}" +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUTPUT="$ROOT/experiments/trb/vdjdb-ylq-official-motif-pooled/evaluation-${SLURM_JOB_ID}" +TOKENIZER="$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +test -f "$PREFLIGHT_DIR/PREFLIGHT.json" +cd "$ROOT" +"$PYTHON" -m benchmark.evaluate_vdjdb_ylq_motif_pooled \ + --preflight-dir "$PREFLIGHT_DIR" --tokenizer "$TOKENIZER" --output-dir "$OUTPUT" \ + --model "r:$ROOT/experiments/trb/data-anchor-r-seed42/best.pt:$ROOT/experiments/trb/data-anchor-r-seed42/run_config.json" \ + --model "t:$ROOT/experiments/trb/data-anchor-t-seed42/best.pt:$ROOT/experiments/trb/data-anchor-t-seed42/run_config.json" \ + --model "p:$ROOT/experiments/trb/data-anchor-p-seed42/best.pt:$ROOT/experiments/trb/data-anchor-p-seed42/run_config.json" \ + --model "rt:$ROOT/experiments/trb/data-anchor-rt-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rt-seed42/run_config.json" \ + --model "rp:$ROOT/experiments/trb/data-anchor-rp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-rp-seed42/run_config.json" \ + --model "tp:$ROOT/experiments/trb/data-anchor-tp-seed42/best.pt:$ROOT/experiments/trb/data-anchor-tp-seed42/run_config.json" \ + --model "rtp:$ROOT/experiments/trb/anchor-data-rtp-seed42/best.pt:$ROOT/experiments/trb/anchor-data-rtp-seed42/run_config.json" \ + --batch-size 128 --bootstrap 2000 --permutations 1000 --seed 42 diff --git a/slurm/finalize_external_residue_pca128_all.sbatch b/slurm/finalize_external_residue_pca128_all.sbatch new file mode 100644 index 0000000..538ccc4 --- /dev/null +++ b/slurm/finalize_external_residue_pca128_all.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-pca128-gate +#SBATCH --partition=short +#SBATCH --cpus-per-task=2 +#SBATCH --mem=8G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/gate-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/gate-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark/residue-pca128-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +cd "$REPO" +"$PYTHON" benchmark/finalize_external_residue_pca256.py \ + --output-dir "$ROOT/bottleneck" --rows 99430 --components 128 +sha256sum "$ROOT/bottleneck"/* > "$ROOT/bottleneck/HASHES.sha256" +echo status=accepted diff --git a/slurm/finalize_external_residue_pca256.sbatch b/slurm/finalize_external_residue_pca256.sbatch new file mode 100644 index 0000000..c49abf6 --- /dev/null +++ b/slurm/finalize_external_residue_pca256.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-pca-gate +#SBATCH --partition=short +#SBATCH --cpus-per-task=2 +#SBATCH --mem=8G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/finalize-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/finalize-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark/residue-pca256-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +cd "$REPO" +"$PYTHON" benchmark/finalize_external_residue_pca256.py --output-dir "$ROOT/bottleneck" +sha256sum "$ROOT/bottleneck"/*_256.npy "$ROOT/bottleneck/representations.json" \ + "$ROOT/bottleneck/bottleneck_summary.json" > "$ROOT/bottleneck/HASHES.sha256" +echo status=accepted diff --git a/slurm/finalize_external_sceptr_pca128.sbatch b/slurm/finalize_external_sceptr_pca128.sbatch new file mode 100644 index 0000000..6c1830e --- /dev/null +++ b/slurm/finalize_external_sceptr_pca128.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-sceptr-pca128-gate +#SBATCH --partition=short +#SBATCH --cpus-per-task=2 +#SBATCH --mem=8G +#SBATCH --time=00:20:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/gate-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/gate-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +cd "$REPO" +"$PYTHON" benchmark/finalize_external_residue_pca_single.py \ + --output-dir "$ROOT/bottleneck" --representation sceptr \ + --rows 99430 --components 128 +sha256sum "$ROOT/bottleneck"/* > "$ROOT/bottleneck/HASHES.sha256" diff --git a/slurm/finalize_vdjdb_epitope_prediction.sbatch b/slurm/finalize_vdjdb_epitope_prediction.sbatch new file mode 100644 index 0000000..d6d55a1 --- /dev/null +++ b/slurm/finalize_vdjdb_epitope_prediction.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=finalize-epitope-prediction +#SBATCH --partition=short +#SBATCH --constraint=hpc +#SBATCH --cpus-per-task=2 +#SBATCH --mem=16G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/epitope-prediction-v1/logs/finalize-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/epitope-prediction-v1/logs/finalize-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +OUT="${EPITOPE_PREDICTION_OUT:-$REPO/experiments/trb/epitope-prediction-v1/evaluation-1429982}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export MPLBACKEND=Agg +cd "$REPO" +"$PYTHON" benchmark/finalize_vdjdb_epitope_prediction.py --output-dir "$OUT" +sha256sum "$OUT"/* > "$OUT/HASHES.sha256" diff --git a/slurm/fit_author_global_bottleneck.sbatch b/slurm/fit_author_global_bottleneck.sbatch new file mode 100644 index 0000000..982211b --- /dev/null +++ b/slurm/fit_author_global_bottleneck.sbatch @@ -0,0 +1,23 @@ +#!/bin/bash +#SBATCH --job-name=irrm-author-global-pca +#SBATCH --partition=medium +#SBATCH --cpus-per-task=16 +#SBATCH --mem=96G +#SBATCH --time=03:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/pca-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/pca-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +OLD="$REPO/experiments/trb/external-model-benchmark" +BASE="$REPO/experiments/trb/external-global-author-v1" +OUT="$BASE/bottleneck" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$BASE/logs" "$OUT" +cd "$REPO" + +"$PYTHON" benchmark/fit_external_model_bottleneck.py --dataset-dir "$OLD/cohort" --representations-dir "$OLD/cohort/representations" --output-dir "$OUT" --dims 64 --representations esm2_35m sceptr_cdr3 --batch-size 4096 --seed 42 +sha256sum "$OUT"/*_64.npy "$OUT/bottleneck_summary.json" > "$OUT/HASHES.sha256" +echo "status=accepted output=$OUT" diff --git a/slurm/fit_external_model_bottleneck.sbatch b/slurm/fit_external_model_bottleneck.sbatch new file mode 100644 index 0000000..629b999 --- /dev/null +++ b/slurm/fit_external_model_bottleneck.sbatch @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-pca +#SBATCH --partition=medium +#SBATCH --cpus-per-task=32 +#SBATCH --mem=192G +#SBATCH --time=08:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/pca-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/pca-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark" +COHORT="$ROOT/cohort" +OUT="$ROOT/bottleneck" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$OUT" +cd "$REPO" + +"$PYTHON" benchmark/fit_external_model_bottleneck.py \ + --dataset-dir "$COHORT" \ + --representations-dir "$COHORT/representations" \ + --output-dir "$OUT" \ + --dims 64 \ + --representations onehot tcremp esm2_8m tcr_bert sceptr rtp \ + --batch-size 4096 \ + --seed 42 + +sha256sum "$OUT"/*_64.npy "$OUT/bottleneck_summary.json" > "$OUT/HASHES.sha256" +echo "status=accepted" diff --git a/slurm/fit_external_residue_pca128_all.sbatch b/slurm/fit_external_residue_pca128_all.sbatch new file mode 100644 index 0000000..c29f50f --- /dev/null +++ b/slurm/fit_external_residue_pca128_all.sbatch @@ -0,0 +1,28 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-pca128 +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64G +#SBATCH --time=04:00:00 +#SBATCH --array=0-2%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/pca-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/pca-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca128-v1" +REPRESENTATIONS=(esm2_8m tcr_bert sceptr) +REPRESENTATION="${REPRESENTATIONS[$SLURM_ARRAY_TASK_ID]}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/bottleneck" +cd "$REPO" +"$PYTHON" benchmark/fit_external_residue_pca256.py \ + --dataset-dir "$BASE/cohort" --residue-dir "$BASE/residue-v1/representations" \ + --output-dir "$ROOT/bottleneck" --representation "$REPRESENTATION" \ + --components 128 --oversample 32 --power-iterations 4 --batch-size 2048 --seed 42 +echo "status=accepted representation=$REPRESENTATION" diff --git a/slurm/fit_external_residue_pca256.sbatch b/slurm/fit_external_residue_pca256.sbatch new file mode 100644 index 0000000..dccd9a0 --- /dev/null +++ b/slurm/fit_external_residue_pca256.sbatch @@ -0,0 +1,30 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-pca256 +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64G +#SBATCH --time=04:00:00 +#SBATCH --array=0-2%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/pca-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/pca-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca256-v1" +REPRESENTATIONS=(esm2_8m tcr_bert sceptr) +REPRESENTATION="${REPRESENTATIONS[$SLURM_ARRAY_TASK_ID]}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/bottleneck" +cd "$REPO" +"$PYTHON" benchmark/fit_external_residue_pca256.py \ + --dataset-dir "$BASE/cohort" \ + --residue-dir "$BASE/residue-v1/representations" \ + --output-dir "$ROOT/bottleneck" \ + --representation "$REPRESENTATION" \ + --components 256 --oversample 32 --power-iterations 4 --batch-size 2048 --seed 42 +echo "status=accepted representation=$REPRESENTATION" diff --git a/slurm/fit_external_sceptr_pca128.sbatch b/slurm/fit_external_sceptr_pca128.sbatch new file mode 100644 index 0000000..1bfc27e --- /dev/null +++ b/slurm/fit_external_sceptr_pca128.sbatch @@ -0,0 +1,27 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-sceptr-pca128 +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64G +#SBATCH --time=04:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/pca-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/pca-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca128-sceptr-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/bottleneck" +cd "$REPO" +"$PYTHON" benchmark/fit_external_residue_pca256.py \ + --dataset-dir "$BASE/cohort" \ + --residue-dir "$BASE/residue-v1/representations" \ + --output-dir "$ROOT/bottleneck" --representation sceptr \ + --components 128 --oversample 32 --power-iterations 4 \ + --batch-size 2048 --seed 42 +echo status=accepted diff --git a/slurm/install_external_benchmark_runtime.sbatch b/slurm/install_external_benchmark_runtime.sbatch new file mode 100644 index 0000000..c949951 --- /dev/null +++ b/slurm/install_external_benchmark_runtime.sbatch @@ -0,0 +1,109 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-runtime +#SBATCH --partition=medium +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/runtime-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/runtime-%j.err + +set -euo pipefail + +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark/runtime-v2" +SITE="$ROOT/site-packages" +HF_HOME="$ROOT/huggingface" +TORCH_HOME="$ROOT/torch" +REPORT="$ROOT/RUNTIME.json" + +mkdir -p "$REPO/experiments/trb/external-model-benchmark/logs" "$SITE" "$HF_HOME" "$TORCH_HOME" + +# The student report named the model/checkpoint versions but did not pin the +# transformers/fair-esm library builds. Pin compatible library versions here and record +# the complete installed distribution inventory below. --target keeps the shared conda +# environment unchanged; its Python executable remains the runtime. +"$PYTHON" -m pip install \ + --disable-pip-version-check \ + --target "$SITE" \ + --upgrade \ + --no-deps \ + "sceptr==1.2.0" \ + "fair-esm==2.0.0" \ + "transformers==4.44.2" \ + "libtcrlm==1.1.3" \ + "tidytcells==2.2.1" \ + "blosum==2.2.0" \ + "tokenizers==0.19.1" \ + "huggingface-hub==0.36.2" \ + "safetensors==0.4.5" \ + "regex==2024.11.6" \ + "requests==2.32.3" \ + "filelock==3.16.1" \ + "packaging==24.2" \ + "PyYAML==6.0.2" \ + "tqdm==4.67.1" \ + "fsspec==2024.10.0" \ + "urllib3==2.2.3" \ + "charset-normalizer==3.4.0" \ + "chardet==5.2.0" \ + "idna==3.10" \ + "certifi==2024.8.30" \ + "rapidfuzz==3.10.1" + +export PYTHONPATH="$SITE${PYTHONPATH:+:$PYTHONPATH}" +export HF_HOME TORCH_HOME + +"$PYTHON" - <<'PY' +import hashlib +import importlib +import importlib.metadata +import json +import os +import pathlib +import sys + +root = pathlib.Path(os.environ["HF_HOME"]).parent +packages = [ + "sceptr", "fair-esm", "transformers", "tokenizers", "huggingface-hub", + "safetensors", "rapidfuzz", +] +modules = {} +for name in ("torch", "sceptr", "esm", "transformers"): + module = importlib.import_module(name) + modules[name] = { + "version": getattr(module, "__version__", None), + "path": str(pathlib.Path(module.__file__).resolve()), + } + +distributions = {} +for package in packages: + try: + distributions[package] = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + distributions[package] = None + +report = { + "status": "accepted", + "python": sys.executable, + "site_packages": str(root / "site-packages"), + "hf_home": os.environ["HF_HOME"], + "torch_home": os.environ["TORCH_HOME"], + "modules": modules, + "distributions": distributions, + "provenance_note": ( + "Student report pins SCEPTR 1.2.0, esm2_t6_8M_UR50D and " + "wukevin/tcr-bert but not fair-esm/transformers library versions; " + "this rerun pins compatible library builds explicitly." + ), +} +if modules["torch"]["version"] != "2.4.1+cu121" or "/.conda/envs/irrm-codec/" not in modules["torch"]["path"]: + raise RuntimeError(f"isolated packages shadowed validated conda torch: {modules['torch']}") +(root / "RUNTIME.json").write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8") +print(json.dumps(report, indent=2, sort_keys=True)) +PY + +cd "$REPO" +sha256sum slurm/install_external_benchmark_runtime.sbatch "$REPORT" > "$ROOT/HASHES.sha256" +echo "status=accepted" +echo "runtime=$ROOT" diff --git a/slurm/install_wordpiece_dependency.sbatch b/slurm/install_wordpiece_dependency.sbatch new file mode 100644 index 0000000..16029c5 --- /dev/null +++ b/slurm/install_wordpiece_dependency.sbatch @@ -0,0 +1,14 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-wp-dependency +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=1 +#SBATCH --mem=4G +#SBATCH --time=00:20:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" +"$PYTHON_BIN" -m pip install --disable-pip-version-check "tokenizers==0.23.1" +"$PYTHON_BIN" -c 'import tokenizers; assert tokenizers.__version__ == "0.23.1"; print(tokenizers.__version__)' diff --git a/slurm/preflight_author_global_embeddings.sbatch b/slurm/preflight_author_global_embeddings.sbatch new file mode 100644 index 0000000..6fb6631 --- /dev/null +++ b/slurm/preflight_author_global_embeddings.sbatch @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=irrm-author-global-gate +#SBATCH --partition=medium +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=01:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/preflight-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/preflight-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-global-author-v1" +RUNTIME="$REPO/experiments/trb/external-model-benchmark/runtime-v2" +LOCKED="$REPO/experiments/trb/external-model-benchmark/cohort/dataset.parquet" +REDCEA="$REPO/experiments/trb/vdjdb-redcea-clonotypes/preflight-1426719/cohort/cohort.tsv" +OUT="$BASE/preflight-$SLURM_JOB_ID" + +export PYTHONPATH="$RUNTIME/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export HF_HOME="$RUNTIME/huggingface" +export TORCH_HOME="$RUNTIME/torch" +export PYTHONUNBUFFERED=1 +export TOKENIZERS_PARALLELISM=false +export HF_HUB_DISABLE_TELEMETRY=1 +mkdir -p "$BASE/logs" +cd "$REPO" + +"$PYTHON" -m py_compile benchmark/cache_external_model_embeddings.py benchmark/preflight_author_global_embeddings.py benchmark/train_external_model_pgen.py +"$PYTHON" benchmark/preflight_author_global_embeddings.py --locked-dataset "$LOCKED" --redcea-cohort "$REDCEA" --locked-sha256 44aa9f8f73c23c0e42ef2bc113a5e4ee3f49fcc335e3c84f0d2eacd59021c191 --redcea-sha256 b3e5852e6952b8f4f529a60287aa76d4308dae37151b7f2becfb4ff87bd98575 --output-dir "$OUT" +sha256sum "$OUT/PREFLIGHT.json" > "$OUT/HASHES.sha256" +echo "status=accepted output=$OUT" diff --git a/slurm/preflight_data_anchor_factorial.sbatch b/slurm/preflight_data_anchor_factorial.sbatch new file mode 100644 index 0000000..644bf26 --- /dev/null +++ b/slurm/preflight_data_anchor_factorial.sbatch @@ -0,0 +1,30 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-data-factorial-pre +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +REFERENCE="$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/run_config.json" +OUTPUT_ROOT="$ROOT_DIR/experiments/trb/data-anchor-full-factorial-gate" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER" +test -f "$REFERENCE" +test ! -e "$OUTPUT_ROOT/PREFLIGHT.json" +mkdir -p "$OUTPUT_ROOT" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m scripts.preflight_data_anchor_objectives \ + --data-dir "$DATA_DIR" \ + --tokenizer "$TOKENIZER" \ + --reference-run-config "$REFERENCE" \ + --output-root "$OUTPUT_ROOT" diff --git a/slurm/preflight_data_anchor_objectives.sbatch b/slurm/preflight_data_anchor_objectives.sbatch new file mode 100644 index 0000000..3e02ac1 --- /dev/null +++ b/slurm/preflight_data_anchor_objectives.sbatch @@ -0,0 +1,30 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-data-obj-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +REFERENCE="$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42/run_config.json" +OUTPUT_ROOT="$ROOT_DIR/experiments/trb/data-anchor-objective-matrix" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER" +test -f "$REFERENCE" +test ! -e "$OUTPUT_ROOT/PREFLIGHT.json" +mkdir -p "$OUTPUT_ROOT" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m scripts.preflight_data_anchor_objectives \ + --data-dir "$DATA_DIR" \ + --tokenizer "$TOKENIZER" \ + --reference-run-config "$REFERENCE" \ + --output-root "$OUTPUT_ROOT" diff --git a/slurm/preflight_external_model_benchmark.sbatch b/slurm/preflight_external_model_benchmark.sbatch new file mode 100644 index 0000000..12e36f4 --- /dev/null +++ b/slurm/preflight_external_model_benchmark.sbatch @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-preflight +#SBATCH --partition=medium +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=01:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/preflight-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/preflight-%j.err + +set -euo pipefail + +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +DATASET="$REPO/data/benchmark/trb" +RTP="$REPO/experiments/trb/latent-sufficiency/evaluation-1426737/embeddings_rtp.npy" +OUT="$REPO/experiments/trb/external-model-benchmark/preflight-${SLURM_JOB_ID}" + +mkdir -p "$REPO/experiments/trb/external-model-benchmark/logs" "$OUT" +cd "$REPO" + +"$PYTHON" benchmark/preflight_external_model_benchmark.py \ + --dataset-dir "$DATASET" \ + --rtp-embeddings "$RTP" \ + --output "$OUT/PREFLIGHT.json" + +sha256sum \ + benchmark/preflight_external_model_benchmark.py \ + slurm/preflight_external_model_benchmark.sbatch \ + "$OUT/PREFLIGHT.json" > "$OUT/HASHES.sha256" + +echo "status=accepted" +echo "output=$OUT" diff --git a/slurm/preflight_external_residue_benchmark.sbatch b/slurm/preflight_external_residue_benchmark.sbatch new file mode 100644 index 0000000..673fe1a --- /dev/null +++ b/slurm/preflight_external_residue_benchmark.sbatch @@ -0,0 +1,39 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/preflight-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/preflight-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-v1" +RUNTIME="$BASE/runtime-v2" +SMOKE="$ROOT/preflight-$SLURM_JOB_ID/smoke" +export PYTHONPATH="$RUNTIME/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export HF_HOME="$RUNTIME/huggingface" +export TORCH_HOME="$RUNTIME/torch" +export TOKENIZERS_PARALLELISM=false +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$SMOKE" +cd "$REPO" + +"$PYTHON" benchmark/cache_external_residue_embeddings.py \ + --dataset-dir "$BASE/cohort" \ + --output-dir "$SMOKE" \ + --encoders esm2_8m tcr_bert sceptr \ + --device cpu \ + --batch-size 8 \ + --threads "${SLURM_CPUS_PER_TASK:-8}" \ + --limit 8 +"$PYTHON" benchmark/preflight_external_residue_benchmark.py \ + --cohort-dir "$BASE/cohort" \ + --smoke-dir "$SMOKE" \ + --output "$ROOT/preflight-$SLURM_JOB_ID/PREFLIGHT.json" +sha256sum "$ROOT/preflight-$SLURM_JOB_ID/PREFLIGHT.json" +echo "status=accepted job=$SLURM_JOB_ID" diff --git a/slurm/preflight_external_residue_pca128_all.sbatch b/slurm/preflight_external_residue_pca128_all.sbatch new file mode 100644 index 0000000..8a822c2 --- /dev/null +++ b/slurm/preflight_external_residue_pca128_all.sbatch @@ -0,0 +1,32 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-pca128-pre +#SBATCH --partition=short +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=01:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/preflight-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/preflight-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca128-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +mkdir -p "$ROOT/logs" +cd "$REPO" +"$PYTHON" -m py_compile \ + benchmark/audit_external_residue_alignment.py \ + benchmark/fit_external_residue_pca256.py \ + benchmark/finalize_external_residue_pca256.py \ + benchmark/train_external_model_decoder.py \ + benchmark/train_external_model_pgen.py \ + benchmark/combine_external_pca128_results.py +"$PYTHON" benchmark/audit_external_residue_alignment.py \ + --dataset-dir "$BASE/cohort" --residue-dir "$BASE/residue-v1/representations" \ + --output "$ROOT/RESIDUE_AUDIT.json" --seed 42 +bash -n slurm/fit_external_residue_pca128_all.sbatch \ + slurm/finalize_external_residue_pca128_all.sbatch \ + slurm/train_external_residue_pca128_remaining.sbatch \ + slurm/collect_external_residue_pca128_all.sbatch +echo status=accepted diff --git a/slurm/preflight_external_rtp128_reconstruction.sbatch b/slurm/preflight_external_rtp128_reconstruction.sbatch new file mode 100644 index 0000000..6b84b62 --- /dev/null +++ b/slurm/preflight_external_rtp128_reconstruction.sbatch @@ -0,0 +1,100 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rtp128-pre +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=4 +#SBATCH --mem=16G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/native-rtp128-v1/logs/preflight-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/native-rtp128-v1/logs/preflight-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/native-rtp128-v1" +SOURCE="$BASE/cohort/representations/rtp.npy" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/bottleneck" "$ROOT/reconstruction" "$ROOT/results" +cd "$REPO" + +"$PYTHON" -m py_compile \ + benchmark/train_external_model_decoder.py \ + benchmark/collect_external_reconstruction.py +bash -n \ + slurm/train_external_rtp128_reconstruction.sbatch \ + slurm/collect_external_rtp128_reconstruction.sbatch + +"$PYTHON" - "$SOURCE" "$BASE/cohort" "$ROOT" <<'PY' +import hashlib +import json +import os +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from irrm_codec.inverse_model import InverseModel + +source, cohort, root = map(Path, sys.argv[1:]) +matrix = np.load(source, mmap_mode="r") +assert matrix.shape == (99_430, 128), matrix.shape +for start in range(0, len(matrix), 8192): + assert np.isfinite(np.asarray(matrix[start:start + 8192])).all() + +splits = {} +sets = {} +for name in ("train", "val", "test"): + rows = pd.read_csv(cohort / "manifests" / f"{name}.tsv", sep="\t")["row_index"].to_numpy() + splits[name] = int(len(rows)) + sets[name] = set(map(int, rows)) +assert splits == {"train": 79_544, "val": 9_943, "test": 9_943}, splits +assert not (sets["train"] & sets["val"] or sets["train"] & sets["test"] or sets["val"] & sets["test"]) +assert len(sets["train"] | sets["val"] | sets["test"]) == 99_430 + +smoke = InverseModel(embedding_dim=128, max_len=40)(torch.zeros(2, 128)) +assert smoke.shape == (2, 40, 25) and torch.isfinite(smoke).all() + +target = root / "bottleneck" / "rtp_128.npy" +if target.is_symlink() or target.exists(): + target.unlink() +target.symlink_to(source) +assert target.resolve() == source.resolve() + +def sha256(path): + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + +report = { + "status": "accepted", + "representation": "rtp", + "interface": "native_128d_no_projection", + "source": str(source), + "source_sha256": sha256(source), + "shape": list(matrix.shape), + "dtype": str(matrix.dtype), + "finite": True, + "splits": splits, + "split_overlap": 0, + "decoder": "independently_trained_InverseModel", + "seeds": [42, 43, 44], + "environment_python": sys.executable, +} +(root / "PREFLIGHT.json").write_text(json.dumps(report, indent=2), encoding="utf-8") +print(json.dumps(report, indent=2)) +PY + +sha256sum \ + benchmark/train_external_model_decoder.py \ + benchmark/collect_external_reconstruction.py \ + slurm/preflight_external_rtp128_reconstruction.sbatch \ + slurm/train_external_rtp128_reconstruction.sbatch \ + slurm/collect_external_rtp128_reconstruction.sbatch \ + "$ROOT/PREFLIGHT.json" > "$ROOT/PREFLIGHT_HASHES.sha256" +echo status=accepted diff --git a/slurm/preflight_external_sceptr_pca128.sbatch b/slurm/preflight_external_sceptr_pca128.sbatch new file mode 100644 index 0000000..160a8a4 --- /dev/null +++ b/slurm/preflight_external_sceptr_pca128.sbatch @@ -0,0 +1,36 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-sceptr-pca128-pre +#SBATCH --partition=short +#SBATCH --cpus-per-task=2 +#SBATCH --mem=8G +#SBATCH --time=00:20:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/preflight-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/preflight-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca128-sceptr-v1" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +mkdir -p "$ROOT/logs" +cd "$REPO" +test -f "$BASE/cohort/dataset.parquet" +test -f "$BASE/cohort/manifests/train.tsv" +test -f "$BASE/residue-v1/representations/sceptr_residue.npy" +"$PYTHON" -m py_compile \ + benchmark/fit_external_residue_pca256.py \ + benchmark/finalize_external_residue_pca_single.py \ + benchmark/preflight_external_sceptr_pca128.py \ + benchmark/train_external_model_decoder.py \ + benchmark/train_external_model_pgen.py +"$PYTHON" benchmark/preflight_external_sceptr_pca128.py \ + --dataset-dir "$BASE/cohort" \ + --residue-dir "$BASE/residue-v1/representations" \ + --output "$ROOT/PREFLIGHT.json" +bash -n \ + slurm/fit_external_sceptr_pca128.sbatch \ + slurm/finalize_external_sceptr_pca128.sbatch \ + slurm/train_external_sceptr_pca128_heads.sbatch \ + slurm/collect_external_sceptr_pca128.sbatch +echo status=accepted diff --git a/slurm/preflight_phase1_rq1.sbatch b/slurm/preflight_phase1_rq1.sbatch new file mode 100644 index 0000000..7277147 --- /dev/null +++ b/slurm/preflight_phase1_rq1.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rq1-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=4 +#SBATCH --mem=8G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/phase1-rq1-char-seed42/preflight-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/phase1-rq1-char-seed42/preflight-%j.err + +set -euo pipefail + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +cd "$ROOT_DIR" +export PYTHONPATH="$ROOT_DIR${PYTHONPATH:+:$PYTHONPATH}" + +/home/evlasova/.conda/envs/irrm-codec/bin/python scripts/preflight_phase1_rq1.py \ + --data-dir "$ROOT_DIR/data/benchmark/trb" \ + --output-path "$ROOT_DIR/experiments/trb/phase1-rq1-char-seed42/preflight.json" diff --git a/slurm/preflight_rq5_depths.sbatch b/slurm/preflight_rq5_depths.sbatch new file mode 100644 index 0000000..76eff19 --- /dev/null +++ b/slurm/preflight_rq5_depths.sbatch @@ -0,0 +1,21 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rq5-depth-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=4 +#SBATCH --mem=16G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-depths-seed42/preflight-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-d128-depths-seed42/preflight-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m scripts.preflight_rq5_depths \ + --baseline-run-config "$ROOT_DIR/experiments/trb/rq5-char-rtp-d128-seed42/run_config.json" \ + --d128-test-metrics "$ROOT_DIR/experiments/trb/rq5-char-rtp-d128-seed42/test_metrics.json" \ + --d320-test-metrics "$ROOT_DIR/experiments/trb/rq1-char-rtp-seed42/test_metrics.json" \ + --d512-test-metrics "$ROOT_DIR/experiments/trb/rq5-char-rtp-d512-seed42/test_metrics.json" \ + --data-dir "$ROOT_DIR/data/benchmark/trb" \ + --output-path "$ROOT_DIR/experiments/trb/rq5-char-rtp-d128-depths-seed42/preflight.json" diff --git a/slurm/preflight_rq5_widths.sbatch b/slurm/preflight_rq5_widths.sbatch new file mode 100644 index 0000000..b878e2a --- /dev/null +++ b/slurm/preflight_rq5_widths.sbatch @@ -0,0 +1,19 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rq5-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=4 +#SBATCH --mem=16G +#SBATCH --time=00:30:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-widths-seed42/preflight-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/rq5-char-rtp-widths-seed42/preflight-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m scripts.preflight_rq5_widths \ + --baseline-run-config "$ROOT_DIR/experiments/trb/rq1-char-rtp-seed42/run_config.json" \ + --baseline-test-metrics "$ROOT_DIR/experiments/trb/rq1-char-rtp-seed42/test_metrics.json" \ + --data-dir "$ROOT_DIR/data/benchmark/trb" \ + --output-path "$ROOT_DIR/experiments/trb/rq5-char-rtp-widths-seed42/preflight.json" diff --git a/slurm/preflight_single_aa_challenge.sbatch b/slurm/preflight_single_aa_challenge.sbatch new file mode 100644 index 0000000..c7051e3 --- /dev/null +++ b/slurm/preflight_single_aa_challenge.sbatch @@ -0,0 +1,52 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-aa1-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +BASE="$ROOT/experiments/trb/single-aa-challenge" +OUT="$BASE/preflight-${SLURM_JOB_ID}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +TCREMP_PYTHON=/home/evlasova/.conda/envs/vdjdb-redcea/bin/python +TCREMP_ROOT=/home/evlasova/tcremp-git/tcremp +PROTOTYPES="$TCREMP_ROOT/tcremp/resources/tcremp_prototypes_TRB.tsv" +MIRPY_PROVENANCE=/home/evlasova/.conda/envs/vdjdb-redcea/lib/python3.11/site-packages/mir-1.0.0.dist-info/direct_url.json +mkdir -p "$BASE" +test ! -e "$OUT" +mkdir -p "$OUT/raw_tcremp" +test "$(git -C "$TCREMP_ROOT" rev-parse HEAD)" = 372bbb64aa9572e1e2ba04a0060a9c03cd7796c6 +test "$(sha256sum "$PROTOTYPES" | cut -d' ' -f1)" = 38cbb5e75d900ef61476bf3af2ed2ff0150a05988bfe568946cd3de5e167f36b +grep -q 468ddae2107da212bb0db9497a3dc9cc44b00eaf "$MIRPY_PROVENANCE" +cd "$ROOT" +"$PYTHON" -m compileall -q benchmark/prepare_single_aa_challenge.py benchmark/validate_single_aa_targets.py benchmark/evaluate_single_aa_challenge.py +"$PYTHON" -c "import runpy; ns=runpy.run_path('tests/test_single_aa_challenge.py'); ns['test_mutants_are_unique_one_substitution_and_parent_split_is_disjoint']()" +PYTHONPATH="$TCREMP_ROOT:$ROOT" "$TCREMP_PYTHON" -c "from tcremp.tcremp_run import main; from mir.embedding.prototype_embedding import PrototypeEmbedding; print('tcremp_import_ok')" +"$PYTHON" -m benchmark.prepare_single_aa_challenge \ + --data-dir "$ROOT/data/benchmark/trb" --output-dir "$OUT/prepared" \ + --max-parents 10 --variants-per-parent 3 --seed 42 +PYTHONPATH="$TCREMP_ROOT:$ROOT" "$TCREMP_PYTHON" -m tcremp.tcremp_run \ + -i "$OUT/prepared/mutants_airr.tsv" -o "$OUT/raw_tcremp" -e mutants \ + -c TRB -s HomoSapiens -x clone_id -p "$PROTOTYPES" -np 8 \ + -llen 4 -hlen 41 --skip-clustering +"$PYTHON" -m irrm_codec.calc_pgen_1mm \ + --airr-path "$OUT/prepared/mutants_airr.tsv" --output-path "$OUT/mutants_pgen.tsv" \ + --chain TRB --species human --locus beta --clone-id-col clone_id \ + --cdr3-col junction_aa --threads 8 --chunk-size 10 +"$PYTHON" -m benchmark.validate_single_aa_targets \ + --prepared "$OUT/prepared" --pgen "$OUT/mutants_pgen.tsv" \ + --representations "$OUT/raw_tcremp/mutants_tcremp_representations.tsv" \ + --tcremp "$OUT/raw_tcremp/mutants_tcremp.parquet" --output-dir "$OUT/targets" +sha256sum "$OUT/prepared/PREPARED.json" "$OUT/targets/TARGETS.json" > "$OUT/SHA256SUMS" +cp "$MIRPY_PROVENANCE" "$OUT/mirpy_direct_url.json" +sha256sum "$TCREMP_PYTHON" "$MIRPY_PROVENANCE" "$PROTOTYPES" >> "$OUT/SHA256SUMS" +printf '{"status":"accepted","job_id":"%s","project_python":"%s","tcremp_python":"%s","tcremp_commit":"%s","mirpy_commit":"%s","prototype_sha256":"%s"}\n' \ + "$SLURM_JOB_ID" "$PYTHON" "$TCREMP_PYTHON" \ + 372bbb64aa9572e1e2ba04a0060a9c03cd7796c6 \ + 468ddae2107da212bb0db9497a3dc9cc44b00eaf \ + 38cbb5e75d900ef61476bf3af2ed2ff0150a05988bfe568946cd3de5e167f36b > "$OUT/PREFLIGHT.json" diff --git a/slurm/preflight_vdjdb_epitope.sbatch b/slurm/preflight_vdjdb_epitope.sbatch new file mode 100644 index 0000000..560de05 --- /dev/null +++ b/slurm/preflight_vdjdb_epitope.sbatch @@ -0,0 +1,38 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-vdjdb-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUTPUT_ROOT="$ROOT_DIR/experiments/trb/vdjdb-data-anchor-validation" +OUTPUT_DIR="$OUTPUT_ROOT/cohort-preflight-${SLURM_JOB_ID}" +VDJDB_SOURCE="/projects/immunestatus/vdjdb/vdjdb-2025-07-30/vdjdb.txt" +BENCHMARK_DIR="$ROOT_DIR/data/benchmark/trb" +EXPERIMENT_DIR="$ROOT_DIR/experiments/trb/anchor-data-rtp-seed42" +TOKENIZER="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +test -f "$VDJDB_SOURCE" +test -f "$BENCHMARK_DIR/READY.json" +test -f "$EXPERIMENT_DIR/best.pt" +test -f "$EXPERIMENT_DIR/run_config.json" +test -f "$TOKENIZER" +mkdir -p "$OUTPUT_ROOT" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m benchmark.prepare_vdjdb_epitope_cohort \ + --vdjdb-tsv "$VDJDB_SOURCE" \ + --benchmark-dir "$BENCHMARK_DIR" \ + --checkpoint "$EXPERIMENT_DIR/best.pt" \ + --run-config "$EXPERIMENT_DIR/run_config.json" \ + --tokenizer "$TOKENIZER" \ + --output-dir "$OUTPUT_DIR" \ + --min-support 50 --min-donors 2 \ + --sensitivity-thresholds 25 50 100 \ + --min-vdjdb-score 1 --seed 42 diff --git a/slurm/preflight_vdjdb_evaluation.sbatch b/slurm/preflight_vdjdb_evaluation.sbatch new file mode 100644 index 0000000..77762bf --- /dev/null +++ b/slurm/preflight_vdjdb_evaluation.sbatch @@ -0,0 +1,25 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-vdjdb-eval-gate +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=4 +#SBATCH --mem=16G +#SBATCH --time=01:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +VALIDATION_ROOT="$ROOT_DIR/experiments/trb/vdjdb-data-anchor-validation" +COHORT_DIR="$VALIDATION_ROOT/cohort-preflight-1426341" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" +OUTPUT="$VALIDATION_ROOT/evaluation-code-preflight-${SLURM_JOB_ID}.json" + +test -x "$PYTHON_BIN" +test -f "$COHORT_DIR/PREFLIGHT.json" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m scripts.preflight_vdjdb_evaluation \ + --cohort "$COHORT_DIR/cohort.tsv" \ + --cohort-preflight "$COHORT_DIR/PREFLIGHT.json" \ + --evaluation-script "$ROOT_DIR/benchmark/evaluate_vdjdb_frozen.py" \ + --output "$OUTPUT" diff --git a/slurm/prepare_anchored_tokenizers.sbatch b/slurm/prepare_anchored_tokenizers.sbatch new file mode 100644 index 0000000..29f3983 --- /dev/null +++ b/slurm/prepare_anchored_tokenizers.sbatch @@ -0,0 +1,35 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-anchor-preflight +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=08:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER_ROOT="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42" +SOURCE_ROOT="/projects/immunestatus/vdjrearm" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +test -f "$DATA_DIR/READY.json" +test ! -e "$TOKENIZER_ROOT/READY.json" +mkdir -p "$TOKENIZER_ROOT" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m scripts.prepare_anchored_tokenizers \ + --data-dir "$DATA_DIR" \ + --output-root "$TOKENIZER_ROOT" \ + --external-corpus "$SOURCE_ROOT/sample/TRB_1e7.tsv.gz" \ + --generation-script "$SOURCE_ROOT/scripts/sample_murugan.sh" \ + --olga-model-dir "$SOURCE_ROOT/model/Homo+sapiens/TRB_orig" \ + --external-rows 500000 --seed 42 +"$PYTHON_BIN" -m scripts.preflight_anchored_tokenizers \ + --data-dir "$DATA_DIR" \ + --tokenizer-root "$TOKENIZER_ROOT" \ + --baseline-run-config "$ROOT_DIR/experiments/trb/rq5-char-rtp-d128-seed42/run_config.json" \ + --output-path "$TOKENIZER_ROOT/preflight.json" +"$PYTHON_BIN" -c 'import json,pathlib; p=pathlib.Path("/home/evlasova/irrm-codec/experiments/trb/tokenizers/anchored-seed42"); r=json.loads((p/"preflight.json").read_text()); assert r["status"]=="ready"; (p/"READY.json").write_text(json.dumps({"status":"ready","preflight_job_id":__import__("os").environ["SLURM_JOB_ID"]},indent=2))' diff --git a/slurm/prepare_external_model_cohort.sbatch b/slurm/prepare_external_model_cohort.sbatch new file mode 100644 index 0000000..178de9d --- /dev/null +++ b/slurm/prepare_external_model_cohort.sbatch @@ -0,0 +1,44 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-cohort +#SBATCH --partition=medium +#SBATCH --cpus-per-task=16 +#SBATCH --mem=96G +#SBATCH --time=04:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/cohort-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/cohort-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +RUNTIME="$REPO/experiments/trb/external-model-benchmark/runtime-v2" +SOURCE="$REPO/data/benchmark/trb" +RTP="$REPO/experiments/trb/latent-sufficiency/evaluation-1426737/embeddings_rtp.npy" +OUT="$REPO/experiments/trb/external-model-benchmark/cohort" + +export PYTHONPATH="$RUNTIME/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$REPO/experiments/trb/external-model-benchmark/logs" "$OUT" +cd "$REPO" + +"$PYTHON" benchmark/prepare_external_model_cohort.py \ + --source-dataset-dir "$SOURCE" \ + --rtp-embeddings "$RTP" \ + --output-dir "$OUT" \ + --seed 42 + +"$PYTHON" benchmark/cache_external_model_embeddings.py \ + --dataset-dir "$OUT" \ + --output-dir "$OUT/representations" \ + --encoders onehot \ + --rtp-embeddings "$OUT/representations/rtp.npy" \ + --device cpu \ + --threads "${SLURM_CPUS_PER_TASK:-16}" + +"$PYTHON" -m pytest tests/test_external_model_benchmark.py -q +sha256sum \ + benchmark/prepare_external_model_cohort.py \ + benchmark/cache_external_model_embeddings.py \ + "$OUT/COHORT.json" \ + "$OUT/representations/representations.json" > "$OUT/PREPARATION_HASHES.sha256" +echo "status=accepted" +echo "cohort=$OUT" diff --git a/slurm/prepare_single_aa_challenge.sbatch b/slurm/prepare_single_aa_challenge.sbatch new file mode 100644 index 0000000..5a30f71 --- /dev/null +++ b/slurm/prepare_single_aa_challenge.sbatch @@ -0,0 +1,45 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-aa1-targets +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=32 +#SBATCH --mem=96G +#SBATCH --time=16:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +PREFLIGHT="${1:?Accepted preflight directory required}" +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +BASE="$ROOT/experiments/trb/single-aa-challenge" +OUT="$BASE/targets-${SLURM_JOB_ID}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +TCREMP_PYTHON=/home/evlasova/.conda/envs/vdjdb-redcea/bin/python +TCREMP_ROOT=/home/evlasova/tcremp-git/tcremp +PROTOTYPES="$TCREMP_ROOT/tcremp/resources/tcremp_prototypes_TRB.tsv" +MIRPY_PROVENANCE=/home/evlasova/.conda/envs/vdjdb-redcea/lib/python3.11/site-packages/mir-1.0.0.dist-info/direct_url.json +test "$("$PYTHON" -c "import json; print(json.load(open('$PREFLIGHT/PREFLIGHT.json'))['status'])")" = accepted +test ! -e "$OUT" +mkdir -p "$OUT/raw_tcremp" +test "$(git -C "$TCREMP_ROOT" rev-parse HEAD)" = 372bbb64aa9572e1e2ba04a0060a9c03cd7796c6 +test "$(sha256sum "$PROTOTYPES" | cut -d' ' -f1)" = 38cbb5e75d900ef61476bf3af2ed2ff0150a05988bfe568946cd3de5e167f36b +grep -q 468ddae2107da212bb0db9497a3dc9cc44b00eaf "$MIRPY_PROVENANCE" +cd "$ROOT" +"$PYTHON" -m benchmark.prepare_single_aa_challenge \ + --data-dir "$ROOT/data/benchmark/trb" --output-dir "$OUT/prepared" \ + --max-parents 2000 --variants-per-parent 3 --seed 42 +PYTHONPATH="$TCREMP_ROOT:$ROOT" "$TCREMP_PYTHON" -m tcremp.tcremp_run \ + -i "$OUT/prepared/mutants_airr.tsv" -o "$OUT/raw_tcremp" -e mutants \ + -c TRB -s HomoSapiens -x clone_id -p "$PROTOTYPES" -np 32 \ + -llen 4 -hlen 41 --skip-clustering +"$PYTHON" -m irrm_codec.calc_pgen_1mm \ + --airr-path "$OUT/prepared/mutants_airr.tsv" --output-path "$OUT/mutants_pgen.tsv" \ + --chain TRB --species human --locus beta --clone-id-col clone_id \ + --cdr3-col junction_aa --threads 32 --chunk-size 50 +"$PYTHON" -m benchmark.validate_single_aa_targets \ + --prepared "$OUT/prepared" --pgen "$OUT/mutants_pgen.tsv" \ + --representations "$OUT/raw_tcremp/mutants_tcremp_representations.tsv" \ + --tcremp "$OUT/raw_tcremp/mutants_tcremp.parquet" --output-dir "$OUT/targets" +sha256sum "$OUT/prepared/PREPARED.json" "$OUT/targets/TARGETS.json" > "$OUT/SHA256SUMS" +cp "$MIRPY_PROVENANCE" "$OUT/mirpy_direct_url.json" +sha256sum "$TCREMP_PYTHON" "$MIRPY_PROVENANCE" "$PROTOTYPES" >> "$OUT/SHA256SUMS" diff --git a/slurm/prepare_trb_joint.sbatch b/slurm/prepare_trb_joint.sbatch new file mode 100644 index 0000000..9292356 --- /dev/null +++ b/slurm/prepare_trb_joint.sbatch @@ -0,0 +1,29 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-trb-prepare +#SBATCH --partition=medium +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/data/benchmark/trb-prepare-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/data/benchmark/trb-prepare-%j.err + +set -euo pipefail + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUTPUT_DIR="$ROOT_DIR/data/benchmark/trb" +cd "$ROOT_DIR" + +/home/evlasova/.conda/envs/irrm-codec/bin/python -m benchmark.prepare_trb_joint \ + --airr-path /projects/immunestatus/vdjrearm/airr_format/trb_background_100k.tsv \ + --embeddings-path /projects/immunestatus/vdjrearm/tcremp/trb_background_100k_tcremp.parquet \ + --representations-path /projects/immunestatus/vdjrearm/tcremp/trb_background_100k_tcremp_representations.tsv \ + --pgen-path /projects/immunestatus/vdjrearm/pgen_1mm_background_100k/trb_background_100k_pgen.tsv \ + --tcremp-writer-source /home/evlasova/tcremp-git/tcremp/tcremp/tcremp_run.py \ + --output-dir "$OUTPUT_DIR" \ + --seed 42 \ + --train-fraction 0.8 \ + --val-fraction 0.1 \ + --max-len 40 \ + --embedding-batch-size 256 \ + --standardizer-chunk-size 1024 diff --git a/slurm/prepare_vdjdb_redcea_clonotypes.sbatch b/slurm/prepare_vdjdb_redcea_clonotypes.sbatch new file mode 100644 index 0000000..4d498cf --- /dev/null +++ b/slurm/prepare_vdjdb_redcea_clonotypes.sbatch @@ -0,0 +1,41 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-redcea-pre +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=4 +#SBATCH --mem=16G +#SBATCH --time=02:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +BASE="$ROOT/experiments/trb/vdjdb-redcea-clonotypes" +OUTPUT="$BASE/preflight-${SLURM_JOB_ID}" +SOURCE="$OUTPUT/cluster_members_TRB.txt" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +COMMIT="77961e2bec255311176d5168b4b8534fdbee24e4" +BLOB="904320e903c646dc266639eb570461893d29aa2d" +URL="https://raw.githubusercontent.com/antigenomics/vdjdb-motifs/${COMMIT}/results/redcea/cluster_members_TRB.txt" + +test -x "$PYTHON" +test -f "$ROOT/data/benchmark/trb/READY.json" +test ! -e "$OUTPUT" +mkdir -p "$OUTPUT" +cd "$ROOT" +"$PYTHON" -m compileall -q \ + benchmark/prepare_vdjdb_redcea_clonotypes.py \ + benchmark/evaluate_vdjdb_redcea_clonotypes.py +"$PYTHON" -c "import runpy; m=runpy.run_path('tests/test_vdjdb_redcea_clonotypes.py'); m['test_preparation_deduplicates_removes_ambiguity_overlap_and_applies_support']()" +"$PYTHON" -c "import runpy; m=runpy.run_path('tests/test_vdjdb_redcea_evaluation.py'); m['test_perfect_epitope_geometry_scores_perfect_retrieval']()" +curl --fail --location --retry 3 --output "$SOURCE" "$URL" +ACTUAL_BLOB="$(git hash-object "$SOURCE")" +test "$ACTUAL_BLOB" = "$BLOB" +"$PYTHON" -m benchmark.prepare_vdjdb_redcea_clonotypes \ + --cluster-members "$SOURCE" \ + --benchmark-dir "$ROOT/data/benchmark/trb" \ + --output-dir "$OUTPUT/cohort" \ + --min-clonotypes 20 \ + --max-length 40 \ + --source-git-commit "$COMMIT" \ + --source-git-blob "$BLOB" diff --git a/slurm/prepare_vdjdb_viral_motif_top10.sbatch b/slurm/prepare_vdjdb_viral_motif_top10.sbatch new file mode 100644 index 0000000..10065a7 --- /dev/null +++ b/slurm/prepare_vdjdb_viral_motif_top10.sbatch @@ -0,0 +1,33 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-viral10-pre +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +MODE="${1:-prepare}" +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +SNAPSHOT="/projects/immunestatus/vdjdb/vdjdb-2025-07-30" +OUTPUT="$ROOT/experiments/trb/vdjdb-viral-motif-top10/${MODE}-${SLURM_JOB_ID}" +SPECIES="$ROOT/benchmark/viral_antigen_species_2025_07_30.txt" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +mkdir -p "$ROOT/experiments/trb/vdjdb-viral-motif-top10" +cd "$ROOT" +if test "$MODE" = "inventory"; then + "$PYTHON" -m benchmark.prepare_vdjdb_viral_motif_top10 \ + --vdjdb-tsv "$SNAPSHOT/vdjdb.txt" --cluster-members "$SNAPSHOT/cluster_members.txt" \ + --motif-pwms "$SNAPSHOT/motif_pwms.txt" --benchmark-dir "$ROOT/data/benchmark/trb" \ + --output-dir "$OUTPUT" --inventory-only --min-score 1 --seed 42 +else + test -f "$SPECIES" + "$PYTHON" -m benchmark.prepare_vdjdb_viral_motif_top10 \ + --vdjdb-tsv "$SNAPSHOT/vdjdb.txt" --cluster-members "$SNAPSHOT/cluster_members.txt" \ + --motif-pwms "$SNAPSHOT/motif_pwms.txt" --benchmark-dir "$ROOT/data/benchmark/trb" \ + --output-dir "$OUTPUT" --viral-species-list "$SPECIES" --min-score 1 \ + --requested-min-positives 40 --controls-per-positive 4 \ + --max-pairs-per-epitope 5000 --seed 42 +fi diff --git a/slurm/prepare_vdjdb_ylq.sbatch b/slurm/prepare_vdjdb_ylq.sbatch new file mode 100644 index 0000000..f01d7fa --- /dev/null +++ b/slurm/prepare_vdjdb_ylq.sbatch @@ -0,0 +1,20 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-ylq-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUTPUT="$ROOT/experiments/trb/vdjdb-ylq-data-anchor-validation/preflight-${SLURM_JOB_ID}" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +mkdir -p "$ROOT/experiments/trb/vdjdb-ylq-data-anchor-validation" +cd "$ROOT" +"$PYTHON" -m benchmark.prepare_vdjdb_ylq \ + --vdjdb-tsv /projects/immunestatus/vdjdb/vdjdb-2025-07-30/vdjdb.txt \ + --benchmark-dir "$ROOT/data/benchmark/trb" --output-dir "$OUTPUT" \ + --min-positive-cdr3 50 --min-positive-donors 2 --seed 42 diff --git a/slurm/prepare_vdjdb_ylq_motif_pooled.sbatch b/slurm/prepare_vdjdb_ylq_motif_pooled.sbatch new file mode 100644 index 0000000..43dcde5 --- /dev/null +++ b/slurm/prepare_vdjdb_ylq_motif_pooled.sbatch @@ -0,0 +1,24 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-ylq-motif-pre +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +SNAPSHOT="/projects/immunestatus/vdjdb/vdjdb-2025-07-30" +OUTPUT="$ROOT/experiments/trb/vdjdb-ylq-official-motif-pooled/preflight-${SLURM_JOB_ID}" +PYTHON="/home/evlasova/.conda/envs/irrm-codec/bin/python" +mkdir -p "$ROOT/experiments/trb/vdjdb-ylq-official-motif-pooled" +cd "$ROOT" +"$PYTHON" -m benchmark.prepare_vdjdb_ylq_motif_pooled \ + --vdjdb-tsv "$SNAPSHOT/vdjdb.txt" \ + --cluster-members "$SNAPSHOT/cluster_members.txt" \ + --motif-pwms "$SNAPSHOT/motif_pwms.txt" \ + --benchmark-dir "$ROOT/data/benchmark/trb" \ + --output-dir "$OUTPUT" --min-score 1 --min-positives 40 --min-negatives 100 \ + --controls-per-positive 4 --seed 42 diff --git a/slurm/prepare_wordpiece_rq4.sbatch b/slurm/prepare_wordpiece_rq4.sbatch new file mode 100644 index 0000000..dcc3959 --- /dev/null +++ b/slurm/prepare_wordpiece_rq4.sbatch @@ -0,0 +1,29 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rq4-wp-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=01:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER_DIR="$ROOT_DIR/experiments/trb/tokenizers/wordpiece-seed42" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -f "$DATA_DIR/READY.json" +test ! -e "$TOKENIZER_DIR/tokenizer.json" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m scripts.prepare_wordpiece_rq4 \ + --data-dir "$DATA_DIR" \ + --output-dir "$TOKENIZER_DIR" \ + --vocab-size 256 \ + --min-frequency 2 \ + --seed 42 +"$PYTHON_BIN" -m scripts.preflight_wordpiece_rq4 \ + --data-dir "$DATA_DIR" \ + --tokenizer-dir "$TOKENIZER_DIR" \ + --output-path "$TOKENIZER_DIR/preflight.json" diff --git a/slurm/prepare_wordpiece_sweep.sbatch b/slurm/prepare_wordpiece_sweep.sbatch new file mode 100644 index 0000000..613bfeb --- /dev/null +++ b/slurm/prepare_wordpiece_sweep.sbatch @@ -0,0 +1,30 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-wp-sweep-preflight +#SBATCH --partition=short +#SBATCH --constraint=compute +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=02:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER_ROOT="$ROOT_DIR/experiments/trb/tokenizers/wordpiece-sweep-seed42" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -f "$DATA_DIR/READY.json" +test ! -e "$TOKENIZER_ROOT/sweep_manifest.json" +cd "$ROOT_DIR" +"$PYTHON_BIN" -m scripts.prepare_wordpiece_sweep \ + --data-dir "$DATA_DIR" \ + --output-root "$TOKENIZER_ROOT" \ + --vocab-sizes 44 64 128 256 512 1024 2048 4096 \ + --min-frequency 2 \ + --seed 42 +"$PYTHON_BIN" -m scripts.preflight_wordpiece_sweep \ + --data-dir "$DATA_DIR" \ + --tokenizer-root "$TOKENIZER_ROOT" \ + --baseline-run-config "$ROOT_DIR/experiments/trb/rq5-char-rtp-d128-seed42/run_config.json" \ + --output-path "$TOKENIZER_ROOT/preflight.json" diff --git a/slurm/smoke_tcremp_geometry_loss.sbatch b/slurm/smoke_tcremp_geometry_loss.sbatch new file mode 100644 index 0000000..0b7ef9b --- /dev/null +++ b/slurm/smoke_tcremp_geometry_loss.sbatch @@ -0,0 +1,31 @@ +#!/bin/bash -l +#SBATCH --job-name=tcremp-geom-smoke +#SBATCH --partition=short +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=00:30:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err +set -euo pipefail +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/tcremp-geometry-smoke-${SLURM_JOB_ID}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +cd "$ROOT" +PYTHONPATH="$ROOT" "$PYTHON" tests/test_tcremp_geometry_loss.py +PYTHONPATH="$ROOT" "$PYTHON" -m irrm_codec.train_multitask \ + --data-dir "$ROOT/data/benchmark/trb" --standardizer-path "$ROOT/data/benchmark/trb/target_standardizer.npz" \ + --output-dir "$OUT" --train-subset all --pgen-target log10_pgen_1mm \ + --tokenizer-type data_anchor --tokenizer-path "$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" \ + --max-sequence-len 40 --d-model 320 --latent-dim 128 --nhead 8 \ + --encoder-layers 4 --decoder-layers 4 --ff-dim 1280 --dropout 0.1 \ + --tcremp-head-dim 1024 --pgen-head-dim 256 --decoder-memory-tokens 4 \ + --tcremp-loss-weight 1 --pgen-loss-weight 0 --reconstruction-loss-weight 0 \ + --tcremp-mse-fraction 1 --tcremp-centered-cosine-weight 0.15 \ + --tcremp-pairwise-log-distance-weight 0.35 \ + --batch-size 4 --epochs 1 --lr 0.0003 --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 --max-grad-norm 1 --early-stopping-patience 0 \ + --scheduler-factor 0.5 --scheduler-patience 2 --scheduler-min-lr 0.000001 \ + --seed 42 --num-workers 0 --device cuda --amp --no-progress \ + --max-train-batches 2 --max-eval-batches 2 --skip-generation-metrics --no-save-test-predictions \ No newline at end of file diff --git a/slurm/train_anchored_tokenizer.sbatch b/slurm/train_anchored_tokenizer.sbatch new file mode 100644 index 0000000..177c06e --- /dev/null +++ b/slurm/train_anchored_tokenizer.sbatch @@ -0,0 +1,56 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-anchor-rtp +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +KIND="${1:?Tokenizer kind is required}" +case "$KIND" in + edge_k|data_anchor|germline_anchor) ;; + *) echo "Refusing unregistered anchored tokenizer" >&2; exit 2 ;; +esac + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER_ROOT="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42" +TOKENIZER_PATH="$TOKENIZER_ROOT/$KIND/anchored_tokenizer.json" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" +case "$KIND" in + edge_k) EXPERIMENT_ID="anchor-edge3-rtp-seed42" ;; + data_anchor) EXPERIMENT_ID="anchor-data-rtp-seed42" ;; + germline_anchor) EXPERIMENT_ID="anchor-germline-rtp-seed42" ;; +esac +OUTPUT_DIR="$ROOT_DIR/experiments/trb/$EXPERIMENT_ID" + +test -x "$PYTHON_BIN" +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER_ROOT/READY.json" +test -f "$TOKENIZER_ROOT/preflight.json" +test -f "$TOKENIZER_PATH" +mkdir -p "$OUTPUT_DIR" +if test -e "$OUTPUT_DIR/best.pt" || test -e "$OUTPUT_DIR/test_metrics.json"; then + echo "Refusing to overwrite completed experiment: $OUTPUT_DIR" >&2; exit 3 +fi +cd "$ROOT_DIR" +printf '{"job_id":"%s","experiment_id":"%s","tokenizer_type":"%s","node":"%s","cuda_visible_devices":"%s"}\n' \ + "$SLURM_JOB_ID" "$EXPERIMENT_ID" "$KIND" "${SLURMD_NODENAME:-unknown}" \ + "${CUDA_VISIBLE_DEVICES:-unknown}" > "$OUTPUT_DIR/slurm_runtime.json" +"$PYTHON_BIN" -m irrm_codec.train_multitask \ + --data-dir "$DATA_DIR" --standardizer-path "$DATA_DIR/target_standardizer.npz" \ + --output-dir "$OUTPUT_DIR" --train-subset all --pgen-target log10_pgen_1mm \ + --tokenizer-type "$KIND" --tokenizer-path "$TOKENIZER_PATH" \ + --max-sequence-len 40 --d-model 320 --latent-dim 128 --nhead 8 \ + --encoder-layers 4 --decoder-layers 4 --ff-dim 1280 --dropout 0.1 \ + --tcremp-head-dim 1024 --pgen-head-dim 256 --decoder-memory-tokens 4 \ + --tcremp-loss-weight 1 --pgen-loss-weight 1 --reconstruction-loss-weight 1 \ + --tcremp-mse-fraction 0.7 --pgen-huber-delta 0.5 --label-smoothing 0 \ + --batch-size 64 --epochs 40 --lr 0.0003 --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 --max-grad-norm 1 --early-stopping-patience 0 \ + --scheduler-factor 0.5 --scheduler-patience 2 --scheduler-min-lr 0.000001 \ + --seed 42 --num-workers 8 --device cuda --amp --no-progress --no-save-test-predictions diff --git a/slurm/train_author_global_onehot_fixed.sbatch b/slurm/train_author_global_onehot_fixed.sbatch new file mode 100644 index 0000000..feb9f9c --- /dev/null +++ b/slurm/train_author_global_onehot_fixed.sbatch @@ -0,0 +1,27 @@ +#!/bin/bash +#SBATCH --job-name=irrm-onehot-pgen-fixed +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=03:00:00 +#SBATCH --array=0-2%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/pgen-onehot-fixed-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/pgen-onehot-fixed-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +OLD="$REPO/experiments/trb/external-model-benchmark" +BASE="$REPO/experiments/trb/external-global-author-v1" +SEEDS=(42 43 44) +SEED="${SEEDS[$SLURM_ARRAY_TASK_ID]}" + +export PYTHONPATH="$OLD/runtime-v2/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$BASE/logs" "$BASE/pgen-onehot-fixed" +cd "$REPO" + +"$PYTHON" benchmark/train_external_model_pgen.py --dataset-dir "$OLD/cohort" --representations-dir "$OLD/cohort/representations" --output-dir "$BASE/pgen-onehot-fixed" --arm onehot_mlp --target log10_pgen_1mm --train-subset all --seed "$SEED" --epochs 60 --batch-size 256 --lr 0.001 --weight-decay 0.0001 --patience 12 --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted seed=$SEED" diff --git a/slurm/train_author_global_pgen.sbatch b/slurm/train_author_global_pgen.sbatch new file mode 100644 index 0000000..9d08180 --- /dev/null +++ b/slurm/train_author_global_pgen.sbatch @@ -0,0 +1,28 @@ +#!/bin/bash +#SBATCH --job-name=irrm-author-global-pgen +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=03:00:00 +#SBATCH --array=0-5%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/pgen-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/pgen-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +OLD="$REPO/experiments/trb/external-model-benchmark" +BASE="$REPO/experiments/trb/external-global-author-v1" +ARMS=(esm2_35m_mlp sceptr_cdr3_mlp) +SEEDS=(42 43 44) +ARM="${ARMS[$((SLURM_ARRAY_TASK_ID / 3))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % 3))]}" +export PYTHONPATH="$OLD/runtime-v2/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$BASE/logs" "$BASE/pgen" +cd "$REPO" + +"$PYTHON" benchmark/train_external_model_pgen.py --dataset-dir "$OLD/cohort" --representations-dir "$OLD/cohort/representations" --output-dir "$BASE/pgen" --arm "$ARM" --target log10_pgen_1mm --train-subset all --seed "$SEED" --epochs 60 --batch-size 256 --lr 0.001 --weight-decay 0.0001 --patience 12 --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted arm=$ARM seed=$SEED" diff --git a/slurm/train_author_global_pgen_controls.sbatch b/slurm/train_author_global_pgen_controls.sbatch new file mode 100644 index 0000000..8ccf296 --- /dev/null +++ b/slurm/train_author_global_pgen_controls.sbatch @@ -0,0 +1,29 @@ +#!/bin/bash +#SBATCH --job-name=irrm-author-global-controls +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64G +#SBATCH --time=03:00:00 +#SBATCH --array=0-5%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/pgen-controls-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/pgen-controls-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +OLD="$REPO/experiments/trb/external-model-benchmark" +BASE="$REPO/experiments/trb/external-global-author-v1" +ARMS=(onehot_mlp tcremp_mlp) +SEEDS=(42 43 44) +ARM="${ARMS[$((SLURM_ARRAY_TASK_ID / 3))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % 3))]}" + +export PYTHONPATH="$OLD/runtime-v2/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$BASE/logs" "$BASE/pgen" +cd "$REPO" + +"$PYTHON" benchmark/train_external_model_pgen.py --dataset-dir "$OLD/cohort" --representations-dir "$OLD/cohort/representations" --output-dir "$BASE/pgen" --arm "$ARM" --target log10_pgen_1mm --train-subset all --seed "$SEED" --epochs 60 --batch-size 256 --lr 0.001 --weight-decay 0.0001 --patience 12 --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted arm=$ARM seed=$SEED" diff --git a/slurm/train_author_global_reconstruction.sbatch b/slurm/train_author_global_reconstruction.sbatch new file mode 100644 index 0000000..58c02fa --- /dev/null +++ b/slurm/train_author_global_reconstruction.sbatch @@ -0,0 +1,28 @@ +#!/bin/bash +#SBATCH --job-name=irrm-author-global-recon +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=24G +#SBATCH --time=03:00:00 +#SBATCH --array=0-5%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/recon-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-global-author-v1/logs/recon-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +OLD="$REPO/experiments/trb/external-model-benchmark" +BASE="$REPO/experiments/trb/external-global-author-v1" +REPS=(esm2_35m sceptr_cdr3) +SEEDS=(42 43 44) +REP="${REPS[$((SLURM_ARRAY_TASK_ID / 3))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % 3))]}" +export PYTHONPATH="$OLD/runtime-v2/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$BASE/logs" "$BASE/reconstruction" +cd "$REPO" + +"$PYTHON" benchmark/train_external_model_decoder.py --dataset-dir "$OLD/cohort" --bottleneck-dir "$BASE/bottleneck" --output-dir "$BASE/reconstruction" --representation "$REP" --bottleneck-dim 64 --train-subset all --seed "$SEED" --epochs 60 --batch-size 256 --lr 0.001 --weight-decay 0.0001 --patience 8 --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted representation=$REP seed=$SEED" diff --git a/slurm/train_data_anchor_objective.sbatch b/slurm/train_data_anchor_objective.sbatch new file mode 100644 index 0000000..980f30a --- /dev/null +++ b/slurm/train_data_anchor_objective.sbatch @@ -0,0 +1,57 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-data-obj +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +CONDITION="${1:?Condition t, rt, or tp is required}" +case "$CONDITION" in + t) RECON=0; TCREMP=1; PGEN=0 ;; + rt) RECON=1; TCREMP=1; PGEN=0 ;; + tp) RECON=0; TCREMP=1; PGEN=1 ;; + *) echo "Refusing unregistered DATA-ANCHOR objective condition" >&2; exit 2 ;; +esac + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER="$ROOT_DIR/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" +MATRIX_ROOT="$ROOT_DIR/experiments/trb/data-anchor-full-factorial-gate" +PREFLIGHT="$MATRIX_ROOT/PREFLIGHT.json" +CONDITION_CONFIG="$MATRIX_ROOT/configs/$CONDITION.json" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/data-anchor-$CONDITION-seed42" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -x "$PYTHON_BIN" +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER" +test -f "$PREFLIGHT" +test -f "$CONDITION_CONFIG" +"$PYTHON_BIN" -c 'import json,sys; p=json.load(open(sys.argv[1])); assert p["status"]=="ready" and all(p["checks"].values()); c=json.load(open(sys.argv[2])); assert c["condition"]==sys.argv[3]' "$PREFLIGHT" "$CONDITION_CONFIG" "$CONDITION" +mkdir -p "$OUTPUT_DIR" +if test -e "$OUTPUT_DIR/best.pt" || test -e "$OUTPUT_DIR/test_metrics.json"; then + echo "Refusing to overwrite completed experiment: $OUTPUT_DIR" >&2; exit 3 +fi +cd "$ROOT_DIR" +printf '{"job_id":"%s","condition":"%s","node":"%s","cuda_visible_devices":"%s","preflight_sha256":"%s","condition_config_sha256":"%s"}\n' \ + "$SLURM_JOB_ID" "$CONDITION" "${SLURMD_NODENAME:-unknown}" "${CUDA_VISIBLE_DEVICES:-unknown}" \ + "$(sha256sum "$PREFLIGHT" | cut -d' ' -f1)" "$(sha256sum "$CONDITION_CONFIG" | cut -d' ' -f1)" \ + > "$OUTPUT_DIR/slurm_runtime.json" +"$PYTHON_BIN" -m irrm_codec.train_multitask \ + --data-dir "$DATA_DIR" --standardizer-path "$DATA_DIR/target_standardizer.npz" \ + --output-dir "$OUTPUT_DIR" --train-subset all --pgen-target log10_pgen_1mm \ + --tokenizer-type data_anchor --tokenizer-path "$TOKENIZER" \ + --max-sequence-len 40 --d-model 320 --latent-dim 128 --nhead 8 \ + --encoder-layers 4 --decoder-layers 4 --ff-dim 1280 --dropout 0.1 \ + --tcremp-head-dim 1024 --pgen-head-dim 256 --decoder-memory-tokens 4 \ + --tcremp-loss-weight "$TCREMP" --pgen-loss-weight "$PGEN" --reconstruction-loss-weight "$RECON" \ + --tcremp-mse-fraction 0.7 --pgen-huber-delta 0.5 --label-smoothing 0 \ + --batch-size 64 --epochs 40 --lr 0.0003 --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 --max-grad-norm 1 --early-stopping-patience 0 \ + --scheduler-factor 0.5 --scheduler-patience 2 --scheduler-min-lr 0.000001 \ + --seed 42 --num-workers 8 --device cuda --amp --no-progress --no-save-test-predictions diff --git a/slurm/train_external_pca256_pgen.sbatch b/slurm/train_external_pca256_pgen.sbatch new file mode 100644 index 0000000..061059b --- /dev/null +++ b/slurm/train_external_pca256_pgen.sbatch @@ -0,0 +1,32 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-pca256-pgen +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=03:00:00 +#SBATCH --array=0-8%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/pgen-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/pgen-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca256-v1" +ARMS=(esm2_8m_residue_pca256_mlp tcr_bert_residue_pca256_mlp sceptr_residue_pca256_mlp) +SEEDS=(42 43 44) +ARM="${ARMS[$((SLURM_ARRAY_TASK_ID / 3))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % 3))]}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/pgen" +cd "$REPO" +"$PYTHON" benchmark/train_external_model_pgen.py \ + --dataset-dir "$BASE/cohort" --representations-dir "$ROOT/bottleneck" \ + --output-dir "$ROOT/pgen" --arm "$ARM" --target log10_pgen_1mm \ + --train-subset all --seed "$SEED" --epochs 60 --batch-size 256 \ + --lr 0.001 --weight-decay 0.0001 --patience 12 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted arm=$ARM seed=$SEED" diff --git a/slurm/train_external_pca256_reconstruction.sbatch b/slurm/train_external_pca256_reconstruction.sbatch new file mode 100644 index 0000000..7a78ecb --- /dev/null +++ b/slurm/train_external_pca256_reconstruction.sbatch @@ -0,0 +1,32 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-pca256-recon +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --array=0-8%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/recon-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/recon-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca256-v1" +REPRESENTATIONS=(esm2_8m_residue tcr_bert_residue sceptr_residue) +SEEDS=(42 43 44) +REPRESENTATION="${REPRESENTATIONS[$((SLURM_ARRAY_TASK_ID / 3))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % 3))]}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/reconstruction" +cd "$REPO" +"$PYTHON" benchmark/train_external_model_decoder.py \ + --dataset-dir "$BASE/cohort" --bottleneck-dir "$ROOT/bottleneck" \ + --output-dir "$ROOT/reconstruction" --representation "$REPRESENTATION" \ + --bottleneck-dim 256 --train-subset all --seed "$SEED" --epochs 60 \ + --batch-size 256 --lr 0.001 --weight-decay 0.0001 --patience 8 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted representation=$REPRESENTATION seed=$SEED" diff --git a/slurm/train_external_pca256_sceptr_reconstruction.sbatch b/slurm/train_external_pca256_sceptr_reconstruction.sbatch new file mode 100644 index 0000000..c16c479 --- /dev/null +++ b/slurm/train_external_pca256_sceptr_reconstruction.sbatch @@ -0,0 +1,30 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-pca256-sceptr-recon +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --array=0-2%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/recon-sceptr-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca256-v1/logs/recon-sceptr-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca256-v1" +SEEDS=(42 43 44) +SEED="${SEEDS[$SLURM_ARRAY_TASK_ID]}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/reconstruction" +cd "$REPO" +"$PYTHON" benchmark/train_external_model_decoder.py \ + --dataset-dir "$BASE/cohort" --bottleneck-dir "$ROOT/bottleneck" \ + --output-dir "$ROOT/reconstruction" --representation sceptr_residue \ + --bottleneck-dim 256 --train-subset all --seed "$SEED" --epochs 60 \ + --batch-size 256 --lr 0.001 --weight-decay 0.0001 --patience 8 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted representation=sceptr_residue seed=$SEED" diff --git a/slurm/train_external_pgen.sbatch b/slurm/train_external_pgen.sbatch new file mode 100644 index 0000000..52d8eb9 --- /dev/null +++ b/slurm/train_external_pgen.sbatch @@ -0,0 +1,40 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-pgen +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=03:00:00 +#SBATCH --array=0-11%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/pgen-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/pgen-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark" +ARMS=(sceptr_mlp tcr_bert_mlp esm2_8m_mlp rtp_mlp) +SEEDS=(42 43 44) +ARM="${ARMS[$((SLURM_ARRAY_TASK_ID / 3))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % 3))]}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/pgen" +cd "$REPO" + +"$PYTHON" benchmark/train_external_model_pgen.py \ + --dataset-dir "$ROOT/cohort" \ + --representations-dir "$ROOT/cohort/representations" \ + --output-dir "$ROOT/pgen" \ + --arm "$ARM" \ + --target log10_pgen_1mm \ + --train-subset all \ + --seed "$SEED" \ + --epochs 60 \ + --batch-size 256 \ + --lr 0.001 \ + --weight-decay 0.0001 \ + --patience 12 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted arm=$ARM seed=$SEED" diff --git a/slurm/train_external_pgen_tfidf.sbatch b/slurm/train_external_pgen_tfidf.sbatch new file mode 100644 index 0000000..37fdfc0 --- /dev/null +++ b/slurm/train_external_pgen_tfidf.sbatch @@ -0,0 +1,28 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-tfidf +#SBATCH --partition=medium +#SBATCH --cpus-per-task=16 +#SBATCH --mem=96G +#SBATCH --time=02:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/tfidf-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/tfidf-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/pgen" +cd "$REPO" + +"$PYTHON" benchmark/train_external_model_pgen.py \ + --dataset-dir "$ROOT/cohort" \ + --representations-dir "$ROOT/cohort/representations" \ + --output-dir "$ROOT/pgen" \ + --arm tfidf_ridge \ + --target log10_pgen_1mm \ + --train-subset all \ + --seed 42 \ + --threads "${SLURM_CPUS_PER_TASK:-16}" +echo "status=accepted arm=tfidf_ridge" diff --git a/slurm/train_external_reconstruction.sbatch b/slurm/train_external_reconstruction.sbatch new file mode 100644 index 0000000..8d6f093 --- /dev/null +++ b/slurm/train_external_reconstruction.sbatch @@ -0,0 +1,41 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-recon +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=24G +#SBATCH --time=03:00:00 +#SBATCH --array=0-17%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/recon-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/recon-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +ROOT="$REPO/experiments/trb/external-model-benchmark" +RUNTIME="$ROOT/runtime-v2" +REPRESENTATIONS=(onehot tcremp esm2_8m tcr_bert sceptr rtp) +SEEDS=(42 43 44) +REPRESENTATION="${REPRESENTATIONS[$((SLURM_ARRAY_TASK_ID / 3))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % 3))]}" +export PYTHONPATH="$RUNTIME/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/reconstruction" +cd "$REPO" + +"$PYTHON" benchmark/train_external_model_decoder.py \ + --dataset-dir "$ROOT/cohort" \ + --bottleneck-dir "$ROOT/bottleneck" \ + --output-dir "$ROOT/reconstruction" \ + --representation "$REPRESENTATION" \ + --bottleneck-dim 64 \ + --train-subset all \ + --seed "$SEED" \ + --epochs 60 \ + --batch-size 256 \ + --lr 0.001 \ + --weight-decay 0.0001 \ + --patience 8 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted representation=$REPRESENTATION seed=$SEED" diff --git a/slurm/train_external_residue_pca128_remaining.sbatch b/slurm/train_external_residue_pca128_remaining.sbatch new file mode 100644 index 0000000..2ee199d --- /dev/null +++ b/slurm/train_external_residue_pca128_remaining.sbatch @@ -0,0 +1,51 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-pca128-heads +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --array=0-11%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/head-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-v1/logs/head-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca128-v1" +REPRESENTATIONS=(esm2_8m_residue tcr_bert_residue) +ARMS=(esm2_8m_residue_pca128_mlp tcr_bert_residue_pca128_mlp) +SEEDS=(42 43 44) +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/reconstruction" "$ROOT/pgen" +cd "$REPO" + +if (( SLURM_ARRAY_TASK_ID < 6 )); then + REP_INDEX=$((SLURM_ARRAY_TASK_ID / 3)) + SEED_INDEX=$((SLURM_ARRAY_TASK_ID % 3)) + ARM="${ARMS[$REP_INDEX]}" + SEED="${SEEDS[$SEED_INDEX]}" + "$PYTHON" benchmark/train_external_model_pgen.py \ + --dataset-dir "$BASE/cohort" --representations-dir "$ROOT/bottleneck" \ + --output-dir "$ROOT/pgen" --arm "$ARM" --target log10_pgen_1mm \ + --train-subset all --seed "$SEED" --epochs 60 --batch-size 256 \ + --lr 0.001 --weight-decay 0.0001 --patience 12 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" + echo "status=accepted task=pgen arm=$ARM seed=$SEED" +else + OFFSET=$((SLURM_ARRAY_TASK_ID - 6)) + REP_INDEX=$((OFFSET / 3)) + SEED_INDEX=$((OFFSET % 3)) + REPRESENTATION="${REPRESENTATIONS[$REP_INDEX]}" + SEED="${SEEDS[$SEED_INDEX]}" + "$PYTHON" benchmark/train_external_model_decoder.py \ + --dataset-dir "$BASE/cohort" --bottleneck-dir "$ROOT/bottleneck" \ + --output-dir "$ROOT/reconstruction" --representation "$REPRESENTATION" \ + --bottleneck-dim 128 --train-subset all --seed "$SEED" --epochs 60 \ + --batch-size 256 --lr 0.001 --weight-decay 0.0001 --patience 8 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" + echo "status=accepted task=reconstruction representation=$REPRESENTATION seed=$SEED" +fi diff --git a/slurm/train_external_residue_pgen.sbatch b/slurm/train_external_residue_pgen.sbatch new file mode 100644 index 0000000..e1207eb --- /dev/null +++ b/slurm/train_external_residue_pgen.sbatch @@ -0,0 +1,39 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-pgen +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64G +#SBATCH --time=06:00:00 +#SBATCH --array=0-8%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/pgen-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/pgen-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-v1" +REPRESENTATIONS=(esm2_8m tcr_bert sceptr) +SEEDS=(42 43 44) +REPRESENTATION="${REPRESENTATIONS[$((SLURM_ARRAY_TASK_ID / 3))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % 3))]}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/pgen" +cd "$REPO" +"$PYTHON" benchmark/train_external_residue_pgen.py \ + --dataset-dir "$BASE/cohort" \ + --residue-dir "$ROOT/representations" \ + --output-dir "$ROOT/pgen" \ + --representation "$REPRESENTATION" \ + --target log10_pgen_1mm \ + --seed "$SEED" \ + --epochs 60 \ + --batch-size 128 \ + --lr 0.0003 \ + --weight-decay 0.0001 \ + --patience 12 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted representation=$REPRESENTATION seed=$SEED" diff --git a/slurm/train_external_residue_reconstruction.sbatch b/slurm/train_external_residue_reconstruction.sbatch new file mode 100644 index 0000000..ede23d8 --- /dev/null +++ b/slurm/train_external_residue_reconstruction.sbatch @@ -0,0 +1,38 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-res-recon +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64G +#SBATCH --time=06:00:00 +#SBATCH --array=0-8%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/recon-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-v1/logs/recon-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-v1" +REPRESENTATIONS=(esm2_8m tcr_bert sceptr) +SEEDS=(42 43 44) +REPRESENTATION="${REPRESENTATIONS[$((SLURM_ARRAY_TASK_ID / 3))]}" +SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID % 3))]}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/reconstruction" +cd "$REPO" +"$PYTHON" benchmark/train_external_residue_decoder.py \ + --dataset-dir "$BASE/cohort" \ + --residue-dir "$ROOT/representations" \ + --output-dir "$ROOT/reconstruction" \ + --representation "$REPRESENTATION" \ + --seed "$SEED" \ + --epochs 60 \ + --batch-size 64 \ + --lr 0.0003 \ + --weight-decay 0.0001 \ + --patience 8 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted representation=$REPRESENTATION seed=$SEED" diff --git a/slurm/train_external_rtp128_reconstruction.sbatch b/slurm/train_external_rtp128_reconstruction.sbatch new file mode 100644 index 0000000..d2fbe32 --- /dev/null +++ b/slurm/train_external_rtp128_reconstruction.sbatch @@ -0,0 +1,40 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rtp128-recon +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --array=0-2%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/native-rtp128-v1/logs/recon-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/native-rtp128-v1/logs/recon-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/native-rtp128-v1" +SEEDS=(42 43 44) +SEED="${SEEDS[$SLURM_ARRAY_TASK_ID]}" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +cd "$REPO" + +test "$("$PYTHON" -c 'import json; print(json.load(open("experiments/trb/external-model-benchmark/native-rtp128-v1/PREFLIGHT.json"))["status"])')" = accepted + +"$PYTHON" benchmark/train_external_model_decoder.py \ + --dataset-dir "$BASE/cohort" \ + --bottleneck-dir "$ROOT/bottleneck" \ + --output-dir "$ROOT/reconstruction" \ + --representation rtp \ + --bottleneck-dim 128 \ + --train-subset all \ + --seed "$SEED" \ + --epochs 60 \ + --batch-size 256 \ + --lr 0.001 \ + --weight-decay 0.0001 \ + --patience 8 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" +echo "status=accepted representation=rtp native_dim=128 seed=$SEED" diff --git a/slurm/train_external_sceptr_pca128_heads.sbatch b/slurm/train_external_sceptr_pca128_heads.sbatch new file mode 100644 index 0000000..e21be11 --- /dev/null +++ b/slurm/train_external_sceptr_pca128_heads.sbatch @@ -0,0 +1,42 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-sceptr-pca128-heads +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=32G +#SBATCH --time=04:00:00 +#SBATCH --array=0-5%2 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/head-%A-%a.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/residue-pca128-sceptr-v1/logs/head-%A-%a.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +BASE="$REPO/experiments/trb/external-model-benchmark" +ROOT="$BASE/residue-pca128-sceptr-v1" +SEEDS=(42 43 44) +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +mkdir -p "$ROOT/logs" "$ROOT/reconstruction" "$ROOT/pgen" +cd "$REPO" + +if (( SLURM_ARRAY_TASK_ID < 3 )); then + SEED="${SEEDS[$SLURM_ARRAY_TASK_ID]}" + "$PYTHON" benchmark/train_external_model_decoder.py \ + --dataset-dir "$BASE/cohort" --bottleneck-dir "$ROOT/bottleneck" \ + --output-dir "$ROOT/reconstruction" --representation sceptr_residue \ + --bottleneck-dim 128 --train-subset all --seed "$SEED" --epochs 60 \ + --batch-size 256 --lr 0.001 --weight-decay 0.0001 --patience 8 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" + echo "status=accepted task=reconstruction seed=$SEED" +else + SEED="${SEEDS[$((SLURM_ARRAY_TASK_ID - 3))]}" + "$PYTHON" benchmark/train_external_model_pgen.py \ + --dataset-dir "$BASE/cohort" --representations-dir "$ROOT/bottleneck" \ + --output-dir "$ROOT/pgen" --arm sceptr_residue_pca128_mlp \ + --target log10_pgen_1mm --train-subset all --seed "$SEED" --epochs 60 \ + --batch-size 256 --lr 0.001 --weight-decay 0.0001 --patience 12 \ + --threads "${SLURM_CPUS_PER_TASK:-4}" + echo "status=accepted task=pgen seed=$SEED" +fi diff --git a/slurm/train_phase1_rq1_char.sbatch b/slurm/train_phase1_rq1_char.sbatch new file mode 100644 index 0000000..1a5c336 --- /dev/null +++ b/slurm/train_phase1_rq1_char.sbatch @@ -0,0 +1,82 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rq1-char +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail + +EXPERIMENT_ID="${1:?experiment id is required}" +TCREMP_WEIGHT="${2:?TCRemP loss weight is required}" +PGEN_WEIGHT="${3:?pgen loss weight is required}" +RECONSTRUCTION_WEIGHT="${4:?reconstruction loss weight is required}" + +case "$EXPERIMENT_ID:$TCREMP_WEIGHT:$PGEN_WEIGHT:$RECONSTRUCTION_WEIGHT" in + rq1-char-r-seed42:0:0:1|rq1-char-rtp-seed42:1:1:1) ;; + *) echo "Refusing unregistered Phase-1 configuration" >&2; exit 2 ;; +esac + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/$EXPERIMENT_ID" +CONFIG_PATH="$OUTPUT_DIR/experiment_config.json" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -f "$DATA_DIR/READY.json" +test -f "$DATA_DIR/target_standardizer.npz" +test -f "$CONFIG_PATH" +if test -e "$OUTPUT_DIR/best.pt" || test -e "$OUTPUT_DIR/test_metrics.json"; then + echo "Refusing to overwrite an existing completed experiment: $OUTPUT_DIR" >&2 + exit 3 +fi + +cd "$ROOT_DIR" +printf '{"job_id":"%s","experiment_id":"%s","node":"%s","cuda_visible_devices":"%s"}\n' \ + "$SLURM_JOB_ID" "$EXPERIMENT_ID" "${SLURMD_NODENAME:-unknown}" "${CUDA_VISIBLE_DEVICES:-unknown}" \ + > "$OUTPUT_DIR/slurm_runtime.json" + +"$PYTHON_BIN" -m irrm_codec.train_multitask \ + --data-dir "$DATA_DIR" \ + --standardizer-path "$DATA_DIR/target_standardizer.npz" \ + --output-dir "$OUTPUT_DIR" \ + --train-subset all \ + --pgen-target log10_pgen_1mm \ + --tokenizer-type char \ + --max-sequence-len 40 \ + --d-model 320 \ + --latent-dim 320 \ + --nhead 8 \ + --encoder-layers 4 \ + --decoder-layers 4 \ + --ff-dim 1280 \ + --dropout 0.1 \ + --tcremp-head-dim 1024 \ + --pgen-head-dim 256 \ + --decoder-memory-tokens 4 \ + --tcremp-loss-weight "$TCREMP_WEIGHT" \ + --pgen-loss-weight "$PGEN_WEIGHT" \ + --reconstruction-loss-weight "$RECONSTRUCTION_WEIGHT" \ + --tcremp-mse-fraction 0.7 \ + --pgen-huber-delta 0.5 \ + --label-smoothing 0 \ + --batch-size 64 \ + --epochs 40 \ + --lr 0.0003 \ + --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 \ + --max-grad-norm 1 \ + --early-stopping-patience 0 \ + --scheduler-factor 0.5 \ + --scheduler-patience 2 \ + --scheduler-min-lr 0.000001 \ + --seed 42 \ + --num-workers 8 \ + --device cuda \ + --amp \ + --no-progress \ + --no-save-test-predictions diff --git a/slurm/train_phase2_rq23_char.sbatch b/slurm/train_phase2_rq23_char.sbatch new file mode 100644 index 0000000..aa8eb16 --- /dev/null +++ b/slurm/train_phase2_rq23_char.sbatch @@ -0,0 +1,82 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rq23-char +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail + +EXPERIMENT_ID="${1:?experiment id is required}" +TCREMP_WEIGHT="${2:?TCRemP loss weight is required}" +PGEN_WEIGHT="${3:?pgen loss weight is required}" +RECONSTRUCTION_WEIGHT="${4:?reconstruction loss weight is required}" + +case "$EXPERIMENT_ID:$TCREMP_WEIGHT:$PGEN_WEIGHT:$RECONSTRUCTION_WEIGHT" in + rq2-char-rt-seed42:1:0:1|rq2-char-rp-seed42:0:1:1) ;; + *) echo "Refusing unregistered Phase-2 configuration" >&2; exit 2 ;; +esac + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/$EXPERIMENT_ID" +CONFIG_PATH="$OUTPUT_DIR/experiment_config.json" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -f "$DATA_DIR/READY.json" +test -f "$DATA_DIR/target_standardizer.npz" +test -f "$CONFIG_PATH" +if test -e "$OUTPUT_DIR/best.pt" || test -e "$OUTPUT_DIR/test_metrics.json"; then + echo "Refusing to overwrite an existing completed experiment: $OUTPUT_DIR" >&2 + exit 3 +fi + +cd "$ROOT_DIR" +printf '{"job_id":"%s","experiment_id":"%s","node":"%s","cuda_visible_devices":"%s"}\n' \ + "$SLURM_JOB_ID" "$EXPERIMENT_ID" "${SLURMD_NODENAME:-unknown}" "${CUDA_VISIBLE_DEVICES:-unknown}" \ + > "$OUTPUT_DIR/slurm_runtime.json" + +"$PYTHON_BIN" -m irrm_codec.train_multitask \ + --data-dir "$DATA_DIR" \ + --standardizer-path "$DATA_DIR/target_standardizer.npz" \ + --output-dir "$OUTPUT_DIR" \ + --train-subset all \ + --pgen-target log10_pgen_1mm \ + --tokenizer-type char \ + --max-sequence-len 40 \ + --d-model 320 \ + --latent-dim 320 \ + --nhead 8 \ + --encoder-layers 4 \ + --decoder-layers 4 \ + --ff-dim 1280 \ + --dropout 0.1 \ + --tcremp-head-dim 1024 \ + --pgen-head-dim 256 \ + --decoder-memory-tokens 4 \ + --tcremp-loss-weight "$TCREMP_WEIGHT" \ + --pgen-loss-weight "$PGEN_WEIGHT" \ + --reconstruction-loss-weight "$RECONSTRUCTION_WEIGHT" \ + --tcremp-mse-fraction 0.7 \ + --pgen-huber-delta 0.5 \ + --label-smoothing 0 \ + --batch-size 64 \ + --epochs 40 \ + --lr 0.0003 \ + --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 \ + --max-grad-norm 1 \ + --early-stopping-patience 0 \ + --scheduler-factor 0.5 \ + --scheduler-patience 2 \ + --scheduler-min-lr 0.000001 \ + --seed 42 \ + --num-workers 8 \ + --device cuda \ + --amp \ + --no-progress \ + --no-save-test-predictions diff --git a/slurm/train_rq4_wordpiece.sbatch b/slurm/train_rq4_wordpiece.sbatch new file mode 100644 index 0000000..c67242e --- /dev/null +++ b/slurm/train_rq4_wordpiece.sbatch @@ -0,0 +1,84 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rq4-wp +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail + +EXPERIMENT_ID="${1:?experiment id is required}" +TCREMP_WEIGHT="${2:?TCRemP loss weight is required}" +PGEN_WEIGHT="${3:?pgen loss weight is required}" +RECONSTRUCTION_WEIGHT="${4:?reconstruction loss weight is required}" +case "$EXPERIMENT_ID:$TCREMP_WEIGHT:$PGEN_WEIGHT:$RECONSTRUCTION_WEIGHT" in + rq4-wordpiece-r-seed42:0:0:1|rq4-wordpiece-rtp-seed42:1:1:1) ;; + *) echo "Refusing unregistered RQ4 configuration" >&2; exit 2 ;; +esac + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER_DIR="$ROOT_DIR/experiments/trb/tokenizers/wordpiece-seed42" +TOKENIZER_PATH="$TOKENIZER_DIR/tokenizer.json" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/$EXPERIMENT_ID" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER_DIR/preflight.json" +test -f "$TOKENIZER_PATH" +test -f "$OUTPUT_DIR/experiment_config.json" +if test -e "$OUTPUT_DIR/best.pt" || test -e "$OUTPUT_DIR/test_metrics.json"; then + echo "Refusing to overwrite completed experiment: $OUTPUT_DIR" >&2 + exit 3 +fi + +cd "$ROOT_DIR" +printf '{"job_id":"%s","experiment_id":"%s","node":"%s","cuda_visible_devices":"%s"}\n' \ + "$SLURM_JOB_ID" "$EXPERIMENT_ID" "${SLURMD_NODENAME:-unknown}" "${CUDA_VISIBLE_DEVICES:-unknown}" \ + > "$OUTPUT_DIR/slurm_runtime.json" + +"$PYTHON_BIN" -m irrm_codec.train_multitask \ + --data-dir "$DATA_DIR" \ + --standardizer-path "$DATA_DIR/target_standardizer.npz" \ + --output-dir "$OUTPUT_DIR" \ + --train-subset all \ + --pgen-target log10_pgen_1mm \ + --tokenizer-type wordpiece \ + --tokenizer-path "$TOKENIZER_PATH" \ + --max-sequence-len 40 \ + --d-model 320 \ + --latent-dim 320 \ + --nhead 8 \ + --encoder-layers 4 \ + --decoder-layers 4 \ + --ff-dim 1280 \ + --dropout 0.1 \ + --tcremp-head-dim 1024 \ + --pgen-head-dim 256 \ + --decoder-memory-tokens 4 \ + --tcremp-loss-weight "$TCREMP_WEIGHT" \ + --pgen-loss-weight "$PGEN_WEIGHT" \ + --reconstruction-loss-weight "$RECONSTRUCTION_WEIGHT" \ + --tcremp-mse-fraction 0.7 \ + --pgen-huber-delta 0.5 \ + --label-smoothing 0 \ + --batch-size 64 \ + --epochs 40 \ + --lr 0.0003 \ + --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 \ + --max-grad-norm 1 \ + --early-stopping-patience 0 \ + --scheduler-factor 0.5 \ + --scheduler-patience 2 \ + --scheduler-min-lr 0.000001 \ + --seed 42 \ + --num-workers 8 \ + --device cuda \ + --amp \ + --no-progress \ + --no-save-test-predictions diff --git a/slurm/train_rq5_depth.sbatch b/slurm/train_rq5_depth.sbatch new file mode 100644 index 0000000..8f6671e --- /dev/null +++ b/slurm/train_rq5_depth.sbatch @@ -0,0 +1,46 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rq5-depth +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +EXPERIMENT_ID="${1:?experiment id is required}" +ENCODER_LAYERS="${2:?encoder layer count is required}" +case "$EXPERIMENT_ID:$ENCODER_LAYERS" in + rq5-char-rtp-d128-e2d4-seed42:2|rq5-char-rtp-d128-e6d4-seed42:6) ;; + *) echo "Refusing unregistered RQ5 depth" >&2; exit 2 ;; +esac + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/$EXPERIMENT_ID" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" +test -f "$ROOT_DIR/experiments/trb/rq5-char-rtp-d128-depths-seed42/preflight.json" +test -f "$DATA_DIR/READY.json" +test -f "$OUTPUT_DIR/experiment_config.json" +if test -e "$OUTPUT_DIR/best.pt" || test -e "$OUTPUT_DIR/test_metrics.json"; then + echo "Refusing to overwrite completed experiment: $OUTPUT_DIR" >&2; exit 3 +fi + +cd "$ROOT_DIR" +printf '{"job_id":"%s","experiment_id":"%s","node":"%s","cuda_visible_devices":"%s"}\n' \ + "$SLURM_JOB_ID" "$EXPERIMENT_ID" "${SLURMD_NODENAME:-unknown}" "${CUDA_VISIBLE_DEVICES:-unknown}" \ + > "$OUTPUT_DIR/slurm_runtime.json" +"$PYTHON_BIN" -m irrm_codec.train_multitask \ + --data-dir "$DATA_DIR" --standardizer-path "$DATA_DIR/target_standardizer.npz" \ + --output-dir "$OUTPUT_DIR" --train-subset all --pgen-target log10_pgen_1mm \ + --tokenizer-type char --max-sequence-len 40 --d-model 320 --latent-dim 128 \ + --nhead 8 --encoder-layers "$ENCODER_LAYERS" --decoder-layers 4 --ff-dim 1280 --dropout 0.1 \ + --tcremp-head-dim 1024 --pgen-head-dim 256 --decoder-memory-tokens 4 \ + --tcremp-loss-weight 1 --pgen-loss-weight 1 --reconstruction-loss-weight 1 \ + --tcremp-mse-fraction 0.7 --pgen-huber-delta 0.5 --label-smoothing 0 \ + --batch-size 64 --epochs 40 --lr 0.0003 --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 --max-grad-norm 1 --early-stopping-patience 0 \ + --scheduler-factor 0.5 --scheduler-patience 2 --scheduler-min-lr 0.000001 \ + --seed 42 --num-workers 8 --device cuda --amp --no-progress --no-save-test-predictions diff --git a/slurm/train_rq5_width.sbatch b/slurm/train_rq5_width.sbatch new file mode 100644 index 0000000..83d37e7 --- /dev/null +++ b/slurm/train_rq5_width.sbatch @@ -0,0 +1,46 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-rq5-width +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +EXPERIMENT_ID="${1:?experiment id is required}" +LATENT_DIM="${2:?latent dimension is required}" +case "$EXPERIMENT_ID:$LATENT_DIM" in + rq5-char-rtp-d128-seed42:128|rq5-char-rtp-d512-seed42:512) ;; + *) echo "Refusing unregistered RQ5 width" >&2; exit 2 ;; +esac + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/$EXPERIMENT_ID" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" +test -f "$ROOT_DIR/experiments/trb/rq5-char-rtp-widths-seed42/preflight.json" +test -f "$DATA_DIR/READY.json" +test -f "$OUTPUT_DIR/experiment_config.json" +if test -e "$OUTPUT_DIR/best.pt" || test -e "$OUTPUT_DIR/test_metrics.json"; then + echo "Refusing to overwrite completed experiment: $OUTPUT_DIR" >&2; exit 3 +fi + +cd "$ROOT_DIR" +printf '{"job_id":"%s","experiment_id":"%s","node":"%s","cuda_visible_devices":"%s"}\n' \ + "$SLURM_JOB_ID" "$EXPERIMENT_ID" "${SLURMD_NODENAME:-unknown}" "${CUDA_VISIBLE_DEVICES:-unknown}" \ + > "$OUTPUT_DIR/slurm_runtime.json" +"$PYTHON_BIN" -m irrm_codec.train_multitask \ + --data-dir "$DATA_DIR" --standardizer-path "$DATA_DIR/target_standardizer.npz" \ + --output-dir "$OUTPUT_DIR" --train-subset all --pgen-target log10_pgen_1mm \ + --tokenizer-type char --max-sequence-len 40 --d-model 320 --latent-dim "$LATENT_DIM" \ + --nhead 8 --encoder-layers 4 --decoder-layers 4 --ff-dim 1280 --dropout 0.1 \ + --tcremp-head-dim 1024 --pgen-head-dim 256 --decoder-memory-tokens 4 \ + --tcremp-loss-weight 1 --pgen-loss-weight 1 --reconstruction-loss-weight 1 \ + --tcremp-mse-fraction 0.7 --pgen-huber-delta 0.5 --label-smoothing 0 \ + --batch-size 64 --epochs 40 --lr 0.0003 --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 --max-grad-norm 1 --early-stopping-patience 0 \ + --scheduler-factor 0.5 --scheduler-patience 2 --scheduler-min-lr 0.000001 \ + --seed 42 --num-workers 8 --device cuda --amp --no-progress --no-save-test-predictions diff --git a/slurm/train_tcremp_geometry_ablation.sbatch b/slurm/train_tcremp_geometry_ablation.sbatch new file mode 100644 index 0000000..9d32e91 --- /dev/null +++ b/slurm/train_tcremp_geometry_ablation.sbatch @@ -0,0 +1,36 @@ +#!/bin/bash -l +#SBATCH --job-name=tcremp-geom-t +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=16:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err +set -euo pipefail +VARIANT="${1:?Variant centered or relational required}" +case "$VARIANT" in + centered) C=0.30; P=0.00 ;; + relational) C=0.15; P=0.35 ;; + *) echo "Unknown variant: $VARIANT" >&2; exit 2 ;; +esac +ROOT="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +OUT="$ROOT/experiments/trb/tcremp-geometry-t-${VARIANT}-seed42" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +test ! -e "$OUT/best.pt" +cd "$ROOT" +PYTHONPATH="$ROOT" "$PYTHON" -m irrm_codec.train_multitask \ + --data-dir "$ROOT/data/benchmark/trb" --standardizer-path "$ROOT/data/benchmark/trb/target_standardizer.npz" \ + --output-dir "$OUT" --train-subset all --pgen-target log10_pgen_1mm \ + --tokenizer-type data_anchor --tokenizer-path "$ROOT/experiments/trb/tokenizers/anchored-seed42/data_anchor/anchored_tokenizer.json" \ + --max-sequence-len 40 --d-model 320 --latent-dim 128 --nhead 8 \ + --encoder-layers 4 --decoder-layers 4 --ff-dim 1280 --dropout 0.1 \ + --tcremp-head-dim 1024 --pgen-head-dim 256 --decoder-memory-tokens 4 \ + --tcremp-loss-weight 1 --pgen-loss-weight 0 --reconstruction-loss-weight 0 \ + --tcremp-mse-fraction 1 --tcremp-centered-cosine-weight "$C" \ + --tcremp-pairwise-log-distance-weight "$P" \ + --batch-size 64 --epochs 40 --lr 0.0003 --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 --max-grad-norm 1 --early-stopping-patience 0 \ + --scheduler-factor 0.5 --scheduler-patience 2 --scheduler-min-lr 0.000001 \ + --seed 42 --num-workers 8 --device cuda --amp --no-progress --no-save-test-predictions \ No newline at end of file diff --git a/slurm/train_wordpiece_sweep.sbatch b/slurm/train_wordpiece_sweep.sbatch new file mode 100644 index 0000000..cd87db7 --- /dev/null +++ b/slurm/train_wordpiece_sweep.sbatch @@ -0,0 +1,52 @@ +#!/bin/bash -l +#SBATCH --job-name=irrm-wp-sweep +#SBATCH --partition=medium +#SBATCH --constraint=gpu +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=12:00:00 +#SBATCH --output=slurm-%x-%j.out +#SBATCH --error=slurm-%x-%j.err + +set -euo pipefail +VOCAB_SIZE="${1:?WordPiece vocabulary size is required}" +case "$VOCAB_SIZE" in + 44|64|128|256|512|1024|2048|4096) ;; + *) echo "Refusing unregistered WordPiece vocabulary size" >&2; exit 2 ;; +esac + +ROOT_DIR="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +DATA_DIR="$ROOT_DIR/data/benchmark/trb" +TOKENIZER_ROOT="$ROOT_DIR/experiments/trb/tokenizers/wordpiece-sweep-seed42" +TOKENIZER_DIR="$TOKENIZER_ROOT/v${VOCAB_SIZE}-seed42" +TOKENIZER_PATH="$TOKENIZER_DIR/tokenizer.json" +EXPERIMENT_ID="wp-sweep-v${VOCAB_SIZE}-rtp-seed42" +OUTPUT_DIR="$ROOT_DIR/experiments/trb/$EXPERIMENT_ID" +PYTHON_BIN="${IRRM_PYTHON:-$(command -v python)}" + +test -f "$DATA_DIR/READY.json" +test -f "$TOKENIZER_ROOT/preflight.json" +test -f "$TOKENIZER_PATH" +mkdir -p "$OUTPUT_DIR" +if test -e "$OUTPUT_DIR/best.pt" || test -e "$OUTPUT_DIR/test_metrics.json"; then + echo "Refusing to overwrite completed experiment: $OUTPUT_DIR" >&2; exit 3 +fi + +cd "$ROOT_DIR" +printf '{"job_id":"%s","experiment_id":"%s","vocab_size":%s,"node":"%s","cuda_visible_devices":"%s"}\n' \ + "$SLURM_JOB_ID" "$EXPERIMENT_ID" "$VOCAB_SIZE" "${SLURMD_NODENAME:-unknown}" \ + "${CUDA_VISIBLE_DEVICES:-unknown}" > "$OUTPUT_DIR/slurm_runtime.json" +"$PYTHON_BIN" -m irrm_codec.train_multitask \ + --data-dir "$DATA_DIR" --standardizer-path "$DATA_DIR/target_standardizer.npz" \ + --output-dir "$OUTPUT_DIR" --train-subset all --pgen-target log10_pgen_1mm \ + --tokenizer-type wordpiece --tokenizer-path "$TOKENIZER_PATH" \ + --max-sequence-len 40 --d-model 320 --latent-dim 128 --nhead 8 \ + --encoder-layers 4 --decoder-layers 4 --ff-dim 1280 --dropout 0.1 \ + --tcremp-head-dim 1024 --pgen-head-dim 256 --decoder-memory-tokens 4 \ + --tcremp-loss-weight 1 --pgen-loss-weight 1 --reconstruction-loss-weight 1 \ + --tcremp-mse-fraction 0.7 --pgen-huber-delta 0.5 --label-smoothing 0 \ + --batch-size 64 --epochs 40 --lr 0.0003 --weight-decay 0.0001 \ + --gradient-accumulation-steps 1 --max-grad-norm 1 --early-stopping-patience 0 \ + --scheduler-factor 0.5 --scheduler-patience 2 --scheduler-min-lr 0.000001 \ + --seed 42 --num-workers 8 --device cuda --amp --no-progress --no-save-test-predictions diff --git a/slurm/validate_external_model_cohort.sbatch b/slurm/validate_external_model_cohort.sbatch new file mode 100644 index 0000000..a9a6df2 --- /dev/null +++ b/slurm/validate_external_model_cohort.sbatch @@ -0,0 +1,63 @@ +#!/bin/bash +#SBATCH --job-name=irrm-ext-validate +#SBATCH --partition=medium +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=01:00:00 +#SBATCH --output=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/validate-%j.out +#SBATCH --error=/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/logs/validate-%j.err + +set -euo pipefail +REPO="${IRRM_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}" +PYTHON=/home/evlasova/.conda/envs/irrm-codec/bin/python +RUNTIME="$REPO/experiments/trb/external-model-benchmark/runtime-v2" +ROOT="$REPO/experiments/trb/external-model-benchmark" +COHORT="$ROOT/cohort" +export PYTHONPATH="$RUNTIME/site-packages:$REPO${PYTHONPATH:+:$PYTHONPATH}" +export PYTHONUNBUFFERED=1 +cd "$REPO" + +"$PYTHON" -m py_compile \ + benchmark/prepare_external_model_cohort.py \ + benchmark/cache_external_model_embeddings.py \ + benchmark/fit_external_model_bottleneck.py \ + benchmark/train_external_model_decoder.py \ + benchmark/train_external_model_pgen.py + +"$PYTHON" - <<'PY' +import json +from pathlib import Path +import numpy as np +import pandas as pd +import torch +import esm +import sceptr +import transformers +import benchmark.cache_external_model_embeddings +import benchmark.fit_external_model_bottleneck +import benchmark.train_external_model_decoder +import benchmark.train_external_model_pgen + +root = Path("/home/evlasova/irrm-codec/experiments/trb/external-model-benchmark/cohort") +report = json.loads((root / "COHORT.json").read_text()) +catalog = json.loads((root / "representations/representations.json").read_text()) +assert report["status"] == "accepted" and report["common_rows"] == 99430 +assert report["dropped_total"] == 0 and report["split_overlap"] == 0 +expected = {"onehot": (99430, 1000), "tcremp": (99430, 9000), "rtp": (99430, 128)} +for name, shape in expected.items(): + matrix = np.load(catalog[name]["path"], mmap_mode="r") + assert matrix.shape == shape, (name, matrix.shape, shape) +frame = pd.read_parquet(root / "dataset.parquet") +assert len(frame) == 99430 and frame.junction_aa.nunique() == 99430 +print(json.dumps({ + "status": "accepted", + "rows": len(frame), + "split_counts": report["split_counts"], + "representations": {name: list(shape) for name, shape in expected.items()}, + "torch": torch.__version__, + "esm": esm.__version__, + "transformers": transformers.__version__, + "sceptr_import": True, +}, indent=2, sort_keys=True)) +PY +echo "status=accepted" diff --git a/tests/test_anchored_tokenization.py b/tests/test_anchored_tokenization.py new file mode 100644 index 0000000..1fa4413 --- /dev/null +++ b/tests/test_anchored_tokenization.py @@ -0,0 +1,90 @@ +import json +import tempfile +import unittest +from pathlib import Path + +from tokenizers import Tokenizer +from tokenizers.decoders import WordPiece as WordPieceDecoder +from tokenizers.models import WordPiece + +from irrm_codec.anchored_tokenization import AnchoredTokenizer, SPECIAL_TOKENS +from irrm_codec.multitask_data import resolve_encoder_tokenizer + + +class AnchoredTokenizerTest(unittest.TestCase): + def make_bundle(self, root: Path, kind: str, n_anchors: list[str], c_anchors: list[str]): + central_vocab = {token: identifier for identifier, token in enumerate(SPECIAL_TOKENS)} + for amino_acid in "ACDEFGHIKLMNPQRSTVWY": + central_vocab[amino_acid] = len(central_vocab) + for amino_acid in "ACDEFGHIKLMNPQRSTVWY": + central_vocab[f"##{amino_acid}"] = len(central_vocab) + central = Tokenizer(WordPiece(vocab=central_vocab, unk_token="[UNK]")) + central.decoder = WordPieceDecoder(prefix="##") + central.save(str(root / "central_tokenizer.json")) + + vocab = {token: identifier for identifier, token in enumerate(SPECIAL_TOKENS)} + for anchor in sorted(n_anchors): + vocab[f"[N:{anchor}]"] = len(vocab) + for anchor in sorted(c_anchors): + vocab[f"[C:{anchor}]"] = len(vocab) + central_to_input = {} + by_id = {identifier: token for token, identifier in central_vocab.items()} + for identifier in range(4, len(central_vocab)): + central_to_input[str(identifier)] = len(vocab) + vocab[f"[M:{by_id[identifier]}]"] = len(vocab) + payload = { + "tokenizer_type": kind, + "k": 3 if kind == "edge_k" else None, + "central_tokenizer_file": "central_tokenizer.json", + "n_anchors": n_anchors, + "c_anchors": c_anchors, + "central_to_input": central_to_input, + "vocab": vocab, + } + path = root / "anchored_tokenizer.json" + path.write_text(json.dumps(payload), encoding="utf-8") + return path + + def test_edge_roundtrip_and_boundary_types(self): + with tempfile.TemporaryDirectory() as temporary: + path = self.make_bundle(Path(temporary), "edge_k", ["CAS"], ["QYF"]) + tokenizer = AnchoredTokenizer(path) + encoding = tokenizer.encode_with_boundaries("CASSLGQETQYF", 40) + self.assertEqual(tokenizer.decode(encoding.ids), "CASSLGQETQYF") + self.assertEqual(encoding.n_anchor, "CAS") + self.assertEqual(encoding.middle, "SLGQET") + self.assertEqual(encoding.c_anchor, "QYF") + self.assertEqual(encoding.boundary_types[0], "N_ANCHOR") + self.assertEqual(encoding.boundary_types[-1], "C_ANCHOR") + self.assertNotIn(3, encoding.ids) + + def test_variable_longest_nonoverlapping_selection(self): + with tempfile.TemporaryDirectory() as temporary: + path = self.make_bundle( + Path(temporary), "data_anchor", ["CASSL", "CASS", "C"], ["ETQYF", "TQYF", "F"] + ) + tokenizer = AnchoredTokenizer(path) + encoding = tokenizer.encode_with_boundaries("CASSLGQETQYF", 40) + self.assertEqual((encoding.n_anchor, encoding.c_anchor), ("CASSL", "ETQYF")) + self.assertEqual(encoding.middle, "GQ") + self.assertEqual(tokenizer.decode(encoding.ids), "CASSLGQETQYF") + + def test_fail_closed_unknown_anchor(self): + with tempfile.TemporaryDirectory() as temporary: + path = self.make_bundle(Path(temporary), "edge_k", ["CAS"], ["QYF"]) + tokenizer = AnchoredTokenizer(path) + with self.assertRaisesRegex(ValueError, "absent from the exhaustive"): + tokenizer.encode("CARSLGQETQYF", 40) + + def test_resolver_keeps_sequence_only_interface(self): + with tempfile.TemporaryDirectory() as temporary: + path = self.make_bundle(Path(temporary), "germline_anchor", ["CASS", "C"], ["TQYF", "F"]) + resolved = resolve_encoder_tokenizer("germline_anchor", str(path)) + self.assertEqual(resolved.name, "germline_anchor") + self.assertEqual(resolved.path, str(path)) + self.assertGreater(resolved.vocab_size, 4) + self.assertTrue(resolved.encode("CASSLGQETQYF", 40)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_author_global_embeddings.py b/tests/test_author_global_embeddings.py new file mode 100644 index 0000000..201bf38 --- /dev/null +++ b/tests/test_author_global_embeddings.py @@ -0,0 +1,60 @@ +import sys +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +torch = pytest.importorskip("torch") + +from benchmark.cache_external_model_embeddings import ENCODERS, _mean_pool, encode_sceptr + + +class _Log: + def info(self, *args, **kwargs): + pass + + +class _FakeModel: + def __init__(self, seen): + self.seen = seen + + def calc_vector_representations(self, frame): + self.seen.append(tuple(frame.columns)) + return np.ones((len(frame), 64), dtype=np.float32) + + +def test_author_global_encoder_catalog_contains_preregistered_conditions(): + assert "esm2_35m" in ENCODERS + assert "sceptr" in ENCODERS + assert "sceptr_cdr3" in ENCODERS + + +def test_mean_pool_uses_only_real_residue_positions(): + states = torch.tensor([[[1.0, 2.0], [3.0, 6.0], [99.0, 99.0]]]) + keep = torch.tensor([[True, True, False]]) + assert torch.allclose(_mean_pool(states, keep), torch.tensor([[2.0, 4.0]])) + + +def test_sceptr_default_and_cdr3_only_use_native_vector_api(monkeypatch): + default_seen, cdr3_seen = [], [] + fake = SimpleNamespace( + enable_hardware_acceleration=lambda: None, + calc_vector_representations=_FakeModel(default_seen).calc_vector_representations, + variant=SimpleNamespace(cdr3_only=lambda: _FakeModel(cdr3_seen)), + ) + monkeypatch.setitem(sys.modules, "sceptr", fake) + frame = pd.DataFrame( + { + "junction_aa": ["CASSF", "CASRF"], + "v_call": ["TRBV1", "TRBV2"], + "j_call": ["TRBJ1", "TRBJ2"], + } + ) + args = SimpleNamespace(device="cpu", batch_size=2) + default = encode_sceptr(frame["junction_aa"].tolist(), args, _Log(), frame=frame) + cdr3 = encode_sceptr( + frame["junction_aa"].tolist(), args, _Log(), frame=frame, cdr3_only=True + ) + assert default.shape == cdr3.shape == (2, 64) + assert default_seen == [("TRBV", "CDR3B", "TRBJ")] + assert cdr3_seen == [("CDR3B",)] diff --git a/tests/test_data_anchor_factorial.py b/tests/test_data_anchor_factorial.py new file mode 100644 index 0000000..ed0ace5 --- /dev/null +++ b/tests/test_data_anchor_factorial.py @@ -0,0 +1,16 @@ +from scripts.preflight_data_anchor_objectives import CONDITIONS + + +def test_full_nonempty_three_head_factorial_is_registered() -> None: + observed = { + name: (weights["reconstruction"], weights["tcremp"], weights["pgen"]) + for name, weights in CONDITIONS.items() + } + assert observed == { + "r": (1.0, 0.0, 0.0), + "t": (0.0, 1.0, 0.0), + "p": (0.0, 0.0, 1.0), + "rt": (1.0, 1.0, 0.0), + "rp": (1.0, 0.0, 1.0), + "tp": (0.0, 1.0, 1.0), + } diff --git a/tests/test_external_model_benchmark.py b/tests/test_external_model_benchmark.py new file mode 100644 index 0000000..266db47 --- /dev/null +++ b/tests/test_external_model_benchmark.py @@ -0,0 +1,52 @@ +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +from benchmark.preflight_external_model_benchmark import read_manifest_rows, validate_splits +from benchmark.prepare_external_model_cohort import nested_training_subsets + + +def _write_manifest(root: Path, name: str, rows: list[int]) -> None: + pd.DataFrame({"row_index": rows}).to_csv(root / f"{name}.tsv", sep="\t", index=False) + + +def test_validate_splits_requires_exact_nonoverlapping_cover(tmp_path: Path) -> None: + manifests = tmp_path / "manifests" + manifests.mkdir() + _write_manifest(manifests, "train", [0, 2, 4]) + _write_manifest(manifests, "val", [1]) + _write_manifest(manifests, "test", [3, 5]) + report = validate_splits(tmp_path, 6) + assert report["counts"] == {"train": 3, "val": 1, "test": 2} + assert report["overlaps"] == {"train_val": 0, "train_test": 0, "val_test": 0} + + +def test_validate_splits_rejects_overlap(tmp_path: Path) -> None: + manifests = tmp_path / "manifests" + manifests.mkdir() + _write_manifest(manifests, "train", [0, 1]) + _write_manifest(manifests, "val", [1, 2]) + _write_manifest(manifests, "test", [3]) + with pytest.raises(ValueError, match="overlap"): + validate_splits(tmp_path, 4) + + +def test_read_manifest_rejects_duplicate_identity(tmp_path: Path) -> None: + path = tmp_path / "rows.tsv" + pd.DataFrame({"row_index": np.array([2, 2])}).to_csv(path, sep="\t", index=False) + with pytest.raises(ValueError, match="duplicate"): + read_manifest_rows(path) + + +def test_nested_training_subsets_are_exact_deterministic_and_nested() -> None: + frame = pd.DataFrame({"junction_aa": [f"CASS{i:05d}F" for i in range(12_000)]}) + train = np.arange(len(frame), dtype=np.int64) + first = nested_training_subsets(frame, train, seed=42) + second = nested_training_subsets(frame, train[::-1], seed=42) + assert len(first["1k"]) == 1_000 + assert len(first["10k"]) == 10_000 + assert set(first["1k"]) < set(first["10k"]) + assert np.array_equal(first["1k"], second["1k"]) + assert np.array_equal(first["10k"], second["10k"]) diff --git a/tests/test_external_residue_benchmark.py b/tests/test_external_residue_benchmark.py new file mode 100644 index 0000000..a6ef820 --- /dev/null +++ b/tests/test_external_residue_benchmark.py @@ -0,0 +1,78 @@ +import numpy as np +import pytest +import torch + +from benchmark.external_residue_models import ( + FrozenResidueDecoder, + FrozenResiduePgenHead, + ResidueArrayDataset, +) +from benchmark.train_external_residue_decoder import levenshtein_distance +from benchmark.fit_external_residue_pca256 import fit_randomized_pca + + +@pytest.mark.parametrize( + ("left", "right", "expected"), + [("", "", 0), ("CASS", "CASS", 0), ("CASS", "CAS", 1), ("CASS", "CATS", 1), ("CA", "CASS", 2)], +) +def test_exact_levenshtein_distance(left, right, expected): + assert levenshtein_distance(left, right) == expected + assert levenshtein_distance(right, left) == expected + + +def test_train_only_flattened_pca_is_finite_and_256_style(): + rng = np.random.default_rng(7) + matrix = rng.normal(size=(30, 4, 5)).astype(np.float32) + train_rows = np.arange(20) + mean, scale, components, explained = fit_randomized_pca( + matrix, train_rows, n_components=6, oversample=2, niter=2, seed=42, + device=torch.device("cpu"), + ) + scores = ((matrix.reshape(30, -1) - mean) / scale) @ components + assert mean.shape == scale.shape == (20,) + assert components.shape == (20, 6) + assert scores.shape == (30, 6) + assert np.isfinite(scores).all() + assert 0 < explained <= 1 + + +def test_residue_dataset_requires_matching_shapes(tmp_path): + features = np.zeros((4, 5, 3), dtype=np.float16) + mask = np.zeros((4, 4), dtype=bool) + np.save(tmp_path / "features.npy", features) + np.save(tmp_path / "mask.npy", mask) + with pytest.raises(ValueError, match="shapes"): + ResidueArrayDataset(tmp_path / "features.npy", tmp_path / "mask.npy", [0], [1]) + + +def test_position_aware_decoder_is_finite_and_order_sensitive(): + torch.manual_seed(7) + model = FrozenResidueDecoder(input_dim=8, max_len=6, hidden_dim=32, num_layers=1, nhead=4).eval() + features = torch.randn(2, 6, 8) + mask = torch.tensor([[1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 0]], dtype=torch.bool) + with torch.no_grad(): + original = model(features, mask) + reversed_states = model(features.flip(1), mask.flip(1)) + assert original.shape == (2, 6, 25) + assert torch.isfinite(original).all() + assert not torch.allclose(original, reversed_states) + + +def test_position_aware_pgen_head_is_finite(): + torch.manual_seed(9) + model = FrozenResiduePgenHead(input_dim=12, max_len=7, hidden_dim=32, num_layers=1, nhead=4).eval() + features = torch.randn(3, 7, 12) + mask = torch.tensor( + [[1, 1, 1, 1, 0, 0, 0], [1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 0]], + dtype=torch.bool, + ) + with torch.no_grad(): + prediction = model(features, mask) + assert prediction.shape == (3,) + assert torch.isfinite(prediction).all() + + +def test_models_reject_mask_shape_mismatch(): + model = FrozenResidueDecoder(input_dim=4, max_len=5, hidden_dim=16, num_layers=1, nhead=4) + with pytest.raises(ValueError, match="mask"): + model(torch.randn(2, 5, 4), torch.ones(2, 4, dtype=torch.bool)) diff --git a/tests/test_integration_compatibility.py b/tests/test_integration_compatibility.py new file mode 100644 index 0000000..588f619 --- /dev/null +++ b/tests/test_integration_compatibility.py @@ -0,0 +1,51 @@ +"""Regression coverage for formats brought together by the publication merge.""" +import tempfile +from pathlib import Path + +import pytest +from tokenizers import Tokenizer +from tokenizers.decoders import WordPiece as WordPieceDecoder +from tokenizers.models import WordPiece + +from irrm_codec.wordpiece_tokenization import ( + decode_wordpiece, encode_wordpiece_unpadded, filter_vocab_by_max_token_length, + load_wordpiece_tokenizer, +) + + +@pytest.mark.parametrize("specials", [ + ["[PAD]", "[BOS]", "[EOS]", "[UNK]"], + ["[PAD]", "[UNK]", "[BOS]", "[EOS]"], +]) +def test_author_and_student_vocabularies_preserve_ids_and_decode(specials): + vocab = {token: i for i, token in enumerate(specials + ["C", "##A", "CASS"])} + tokenizer = Tokenizer(WordPiece(vocab=vocab, unk_token="[UNK]")) + tokenizer.decoder = WordPieceDecoder(prefix="##") + tokenizer.add_special_tokens(specials) + with tempfile.TemporaryDirectory() as temporary: + path = Path(temporary) / "tokenizer.json" + tokenizer.save(str(path)) + restored = load_wordpiece_tokenizer(path) + assert restored.get_vocab() == vocab + ids = encode_wordpiece_unpadded("CA", restored, 40) + assert ids == [vocab["C"], vocab["##A"]] + assert decode_wordpiece(ids + [vocab["[EOS]"], vocab["C"]], restored) == "CA" + filtered = filter_vocab_by_max_token_length(restored, 1) + assert {s: filtered.token_to_id(s) for s in specials} == {s: vocab[s] for s in specials} + assert decode_wordpiece(encode_wordpiece_unpadded("CA", filtered, 40), filtered) == "CA" + with pytest.raises(ValueError, match="UNK"): + encode_wordpiece_unpadded("CY", restored, 40) + + +def test_sequence_keyed_pgen_cache_survives_row_filtering(tmp_path): + import pandas as pd + from benchmark.prepare_splits import _pgen_lookup + path = tmp_path / "pgen.tsv" + pd.DataFrame({ + "junction_aa": ["CASSF", "CASRF", "CATGF"], + "log10_pgen": [-5., -6., -7.], + "log10_pgen_1mm": [-4., -5., -6.], + }).to_csv(path, sep="\t", index=False) + result = _pgen_lookup(path, pd.Series(["CATGF", "CASSF"])) + assert result.log10_pgen.tolist() == [-7., -5.] + assert _pgen_lookup(path, pd.Series(["CQQF"])) is None diff --git a/tests/test_latent_secondary.py b/tests/test_latent_secondary.py new file mode 100644 index 0000000..da473ab --- /dev/null +++ b/tests/test_latent_secondary.py @@ -0,0 +1,82 @@ +import numpy as np +import pandas as pd + +from benchmark.evaluate_latent_secondary import ( + deterministic_substitutions, + joint_relevance, + make_candidate_gallery, + nested_subsets, + project_tcremp, + preregistration, + retrieval_metrics, + ridge_coefficients, + ridge_predict, +) + + +def test_preregistration_covers_tracks_and_does_not_assume_synthetic_source() -> None: + report = preregistration() + assert set(report["models"]) == {"r", "t", "p", "rt", "rp", "tp", "rtp"} + assert set(report["tracks"]) == { + "low_label", "strata", "joint_retrieval", "robustness", "synthetic_discrimination" + } + assert report["tracks"]["synthetic_discrimination"]["status"].startswith("disabled") + assert "No track" in report["no_selection_rule"] + + +def test_nested_subsets_are_deterministic_and_strictly_nested() -> None: + indices = np.arange(1000) + first = nested_subsets(indices, (0.1, 0.25, 1.0), 42) + second = nested_subsets(indices, (0.1, 0.25, 1.0), 42) + assert all(np.array_equal(first[key], second[key]) for key in first) + assert set(first[0.1]).issubset(first[0.25]) + assert set(first[0.25]).issubset(first[1.0]) + + +def test_ridge_fit_uses_only_selected_rows() -> None: + rng = np.random.default_rng(2) + x = rng.normal(size=(40, 3)) + y = x @ np.array([[1.0, 0.0], [0.0, 2.0], [1.0, -1.0]]) + train = np.arange(30) + fit = ridge_coefficients(x, y, train, 0.1) + prediction = ridge_predict(fit, x, np.arange(30, 40)) + assert np.mean((prediction - y[30:]) ** 2) < 0.01 + + +def test_tcremp_projection_is_deterministic_and_finite() -> None: + class Standardizer: + tcremp_mean = np.asarray([1.0, 2.0, 3.0], dtype=np.float32) + tcremp_std = np.asarray([1.0, 2.0, 4.0], dtype=np.float32) + + values = np.arange(18, dtype=np.float32).reshape(6, 3) + first, report = project_tcremp(values, Standardizer(), dimension=2, seed=42, chunk_size=2) + second, _ = project_tcremp(values, Standardizer(), dimension=2, seed=42, chunk_size=3) + assert np.array_equal(first, second) + assert np.isfinite(first).all() + assert report["source_dimension"] == 3 + + +def test_joint_gallery_excludes_query_and_retrieval_is_perfect_when_ranked() -> None: + gallery = make_candidate_gallery(20, 5, 42) + assert gallery.shape == (20, 5) + assert all(query not in gallery[query] for query in range(20)) + relevant = np.zeros_like(gallery, dtype=bool) + relevant[:, 0] = True + distance = np.ones_like(gallery, dtype=float) + distance[:, 0] = 0 + metrics = retrieval_metrics(distance, relevant) + assert metrics["precision_at_1"] == 1.0 + assert metrics["map"] == 1.0 + + +def test_joint_relevance_and_substitutions_are_valid_and_no_indels() -> None: + sequences = pd.Series(["CASSLGQETQYF", "CASSLGQETQFF", "CASSPGQETQYF", "CARDRGNEQFF"]) + tcremp = np.asarray([[1, 0], [0.9, 0.1], [0.8, 0.2], [0, 1]], dtype=np.float32) + pgen = np.asarray([0.0, 0.1, 0.2, 2.0]) + gallery = np.asarray([[1, 2, 3], [0, 2, 3], [0, 1, 3], [0, 1, 2]], dtype=np.int32) + relevant = joint_relevance(sequences, tcremp, pgen, gallery, 0.5) + assert relevant.shape == gallery.shape + conservative, nonconservative = deterministic_substitutions(sequences.iloc[0], 42) + assert len(conservative) == len(nonconservative) == len(sequences.iloc[0]) + assert sum(a != b for a, b in zip(conservative, sequences.iloc[0])) == 1 + assert sum(a != b for a, b in zip(nonconservative, sequences.iloc[0])) == 1 diff --git a/tests/test_latent_sufficiency.py b/tests/test_latent_sufficiency.py new file mode 100644 index 0000000..d57de21 --- /dev/null +++ b/tests/test_latent_sufficiency.py @@ -0,0 +1,59 @@ +import numpy as np +import pandas as pd + +from benchmark.evaluate_latent_sufficiency import ( + EOS_CLASS, + aggregate_losses, + composition_features, + fit_probe, + one_hot_flat, + sequence_metrics, + sequence_targets, +) + + +def test_sequence_target_roundtrip_and_linear_probe_selection(): + sequences = ["CASS", "CASR", "CATS", "CATR"] * 3 + targets = sequence_targets(sequences, 4) + assert np.all(targets[:, 4] == EOS_CLASS) + one_hot = one_hot_flat(targets) + metrics, rows = sequence_metrics(one_hot, targets) + assert metrics["token_accuracy"] == 1.0 + assert metrics["exact_accuracy"] == 1.0 + assert np.all(rows == 0) + + x = np.arange(len(sequences), dtype=np.float32)[:, None] + train = np.arange(0, 8) + val = np.arange(8, 10) + test = np.arange(10, 12) + alpha, val_pred, test_pred, trace = fit_probe( + x, train, val, test, np.arange(len(sequences), dtype=np.float32)[:, None], + [0.01, 1.0], lambda pred, index: np.mean((pred[:, 0] - index) ** 2), + ) + assert alpha in {0.01, 1.0} + assert val_pred.shape == (2, 1) + assert test_pred.shape == (2, 1) + assert set(trace) == {"0.01", "1.0"} + + +def test_controls_and_aggregate_orientation(): + table = pd.DataFrame({ + "junction_aa": ["CASS", "CATS"], + "v_call": ["TRBV1", "TRBV2"], + "j_call": ["TRBJ1", "TRBJ1"], + }) + sequence_only, meta = composition_features(table, include_vj=False) + annotated, annotated_meta = composition_features(table, include_vj=True, fit_indices=np.array([0])) + assert sequence_only.shape == (2, 21) + assert annotated.shape[1] > sequence_only.shape[1] + assert meta["sequence_only"] is True + assert annotated_meta["sequence_only"] is False + + aggregate = aggregate_losses( + {"sequence": 0.2, "tcremp": 0.5, "pgen": 0.4}, + {"sequence": 0.2, "tcremp": 0.4, "pgen": 0.3}, + {"sequence": 0.8, "tcremp": 1.0, "pgen": 1.0}, + ) + assert aggregate["normalized_regret_by_task"]["sequence"] == 0 + assert aggregate["worst_task_normalized_regret"] > 0 + assert 0 <= aggregate["fixed_reference_hypervolume"] <= 1 diff --git a/tests/test_multitask_baselines.py b/tests/test_multitask_baselines.py new file mode 100644 index 0000000..f5dc463 --- /dev/null +++ b/tests/test_multitask_baselines.py @@ -0,0 +1,191 @@ +import json +import logging +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import pandas as pd +import torch + +from irrm_codec.multitask_data import TargetStandardizer, resolve_encoder_tokenizer +from irrm_codec.multitask_losses import IRRMCodecMultiTaskLoss, MultiTaskLossWeights +from irrm_codec.multitask_transformer import IRRMCodecConfig, IRRMCodecTransformer +from irrm_codec.tokenization import AA_VOCAB, EOS_ID, encode_raw, encode_reconstruction_pair + + +class CharacterTokenizerTest(unittest.TestCase): + def test_character_reconstruction_pair(self): + encoded = encode_raw("CASSLGQETQYF", max_len=40) + decoder_input, target = encode_reconstruction_pair("CASSLGQETQYF", max_len=40) + self.assertEqual(decoder_input[1:], encoded) + self.assertEqual(target[:-1], encoded) + self.assertEqual(target[-1], EOS_ID) + self.assertEqual(resolve_encoder_tokenizer("char").vocab_size, len(AA_VOCAB)) + + def test_character_tokenizer_rejects_overflow(self): + with self.assertRaises(ValueError): + encode_raw("A" * 41, max_len=40) + + +class LockedStandardizerTest(unittest.TestCase): + def test_loader_accepts_benchmark_normalizer_with_extra_keys(self): + with tempfile.TemporaryDirectory() as temporary: + path = Path(temporary) / "target_standardizer.npz" + np.savez( + path, + tcremp_mean=np.zeros(12, dtype=np.float32), + tcremp_std=np.ones(12, dtype=np.float32), + pgen_mean=np.float32(-8.3), + pgen_std=np.float32(2.5), + pgen_target=np.asarray("log10_pgen_1mm"), + train_rows=np.int64(79_544), + log10_pgen_mean=np.float32(-10.4), + ) + loaded = TargetStandardizer.load(path) + self.assertEqual(loaded.pgen_target, "log10_pgen_1mm") + self.assertEqual(loaded.train_rows, 79_544) + self.assertEqual(loaded.tcremp_mean.shape, (12,)) + + +class MatchedBaselineLossTest(unittest.TestCase): + def test_tiny_locked_benchmark_trains_and_writes_terminal_artifacts(self): + from irrm_codec.train_multitask import main + + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + data_dir = root / "data" + manifests = data_dir / "manifests" + output_dir = root / "output" + manifests.mkdir(parents=True) + table = pd.DataFrame( + { + "row_index": range(6), + "junction_aa": ["CASSF", "CASRF", "CAT", "CQQ", "CAR", "CAS"], + "split": ["train", "train", "train", "val", "test", "test"], + "log10_pgen": [-5.0, -6.0, -7.0, -6.0, -7.0, -8.0], + "log10_pgen_1mm": [-4.0, -5.0, -6.0, -5.0, -6.0, -7.0], + } + ) + table.to_parquet(data_dir / "dataset.parquet", index=False) + embeddings = np.arange(72, dtype=np.float32).reshape(6, 12) + np.save(data_dir / "embeddings.npy", embeddings) + for name, rows in (("train", [0, 1, 2]), ("val", [3]), ("test", [4, 5])): + pd.DataFrame({"row_index": rows}).to_csv( + manifests / f"{name}.tsv", sep="\t", index=False + ) + np.savez( + data_dir / "target_standardizer.npz", + tcremp_mean=embeddings[:3].mean(axis=0), + tcremp_std=embeddings[:3].std(axis=0), + pgen_mean=np.float32(-5.0), + pgen_std=np.float32(np.std([-4.0, -5.0, -6.0])), + pgen_target=np.asarray("log10_pgen_1mm"), + train_rows=np.int64(3), + ) + argv = [ + "train_multitask", + "--data-dir", str(data_dir), + "--output-dir", str(output_dir), + "--standardizer-path", str(data_dir / "target_standardizer.npz"), + "--d-model", "32", + "--latent-dim", "24", + "--nhead", "4", + "--encoder-layers", "1", + "--decoder-layers", "1", + "--ff-dim", "64", + "--tcremp-head-dim", "24", + "--pgen-head-dim", "12", + "--decoder-memory-tokens", "2", + "--batch-size", "2", + "--epochs", "1", + "--early-stopping-patience", "0", + "--num-workers", "0", + "--device", "cpu", + "--no-amp", + "--no-progress", + "--skip-generation-metrics", + "--no-save-test-predictions", + ] + with patch("sys.argv", argv): + main() + self.assertTrue((output_dir / "best.pt").is_file()) + self.assertTrue((output_dir / "last.pt").is_file()) + metrics = json.loads((output_dir / "test_metrics.json").read_text()) + self.assertEqual(metrics["samples"], 2) + self.assertEqual(metrics["tokenizer_type"], "char") + logging.shutdown() + + def test_phase1_configs_are_matched_except_condition_and_weights(self): + root = Path(__file__).resolve().parents[1] + r = json.loads( + (root / "experiments/trb/rq1-char-r-seed42/experiment_config.json").read_text() + ) + rtp = json.loads( + (root / "experiments/trb/rq1-char-rtp-seed42/experiment_config.json").read_text() + ) + for key in ( + "base_git_revision", + "saved_pr_revision", + "data_dir", + "tokenizer", + "seed", + "train_subset", + "pgen_target", + "model", + "training", + ): + self.assertEqual(r[key], rtp[key], key) + self.assertEqual(r["loss"]["reconstruction"], 1.0) + self.assertEqual(rtp["loss"]["reconstruction"], 1.0) + self.assertEqual((r["loss"]["tcremp"], r["loss"]["pgen"]), (0.0, 0.0)) + self.assertEqual((rtp["loss"]["tcremp"], rtp["loss"]["pgen"]), (1.0, 1.0)) + + def test_reconstruction_and_multitask_configs_share_architecture(self): + torch.manual_seed(42) + config = IRRMCodecConfig( + input_vocab_size=len(AA_VOCAB), + output_vocab_size=len(AA_VOCAB), + max_sequence_len=8, + d_model=32, + latent_dim=24, + nhead=4, + encoder_layers=1, + decoder_layers=1, + ff_dim=64, + dropout=0.0, + tcremp_dim=12, + tcremp_head_dim=24, + pgen_head_dim=12, + decoder_memory_tokens=2, + ) + model = IRRMCodecTransformer(config) + tokens = torch.randint(5, len(AA_VOCAB), (3, 6)) + decoder_input = torch.randint(5, len(AA_VOCAB), (3, 7)) + target = torch.randint(5, len(AA_VOCAB), (3, 7)) + outputs = model(tokens, tokens.ne(0), decoder_input) + common = dict( + outputs=outputs, + tcremp_target=torch.randn(3, 12), + pgen_target=torch.randn(3), + reconstruction_target=target, + tcremp_mean=torch.zeros(12), + tcremp_std=torch.ones(12), + pgen_mean=torch.tensor(0.0), + pgen_std=torch.tensor(1.0), + ) + r_losses = IRRMCodecMultiTaskLoss( + MultiTaskLossWeights(tcremp=0.0, pgen=0.0, reconstruction=1.0) + )(**common) + rtp_losses = IRRMCodecMultiTaskLoss( + MultiTaskLossWeights(tcremp=1.0, pgen=1.0, reconstruction=1.0) + )(**common) + self.assertTrue(torch.isfinite(r_losses["loss"])) + self.assertTrue(torch.isfinite(rtp_losses["loss"])) + self.assertTrue(torch.allclose(r_losses["loss"], r_losses["reconstruction_loss"])) + self.assertGreater(float(rtp_losses["loss"]), float(r_losses["loss"])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_prepare_trb_joint.py b/tests/test_prepare_trb_joint.py new file mode 100644 index 0000000..b4ef90f --- /dev/null +++ b/tests/test_prepare_trb_joint.py @@ -0,0 +1,104 @@ +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import pandas as pd + +from benchmark.prepare_trb_joint import ( + add_identity_columns, + make_splits, + prepare_identity_table, + write_standardizer, +) + + +class JointTRBPreparationTest(unittest.TestCase): + def sources(self): + airr = pd.DataFrame( + { + "junction_aa": ["CASSF", "CASRF", "CAT", "CAA"], + "v_call": ["TRBV1", "TRBV2", "TRBV3", "TRBV4"], + "j_call": ["TRBJ1", "TRBJ2", "TRBJ1", "TRBJ2"], + "locus": ["beta"] * 4, + } + ) + pgen = airr.copy() + pgen["log10_pgen"] = [-5.0, -6.0, -7.0, -8.0] + pgen["log10_pgen_1mm"] = [-4.0, -5.0, -6.0, -7.0] + representations = pd.DataFrame( + { + "clone_id": [12, 10, 13, 11], + "cdr3aa_TRB": ["CAT", "CASSF", "CAA", "CASRF"], + "v_TRB": ["TRBV3*01", "TRBV1*01", "TRBV4*01", "TRBV2*01"], + "j_TRB": ["TRBJ1*01", "TRBJ1*01", "TRBJ2*01", "TRBJ2*01"], + } + ) + return airr, pgen, representations + + def test_identity_join_does_not_depend_on_row_order(self): + airr, pgen, representations = self.sources() + table, report = prepare_identity_table( + airr, + pgen.sample(frac=1.0, random_state=2), + representations, + min_len=1, + max_len=40, + ) + self.assertEqual(report["rows_kept"], 4) + observed = dict(zip(table["junction_aa"], table["clone_id"])) + self.assertEqual(observed, {"CASSF": 10, "CASRF": 11, "CAT": 12, "CAA": 13}) + self.assertTrue(table["record_id"].is_unique) + + def test_non_finite_targets_are_removed_before_fixed_split(self): + airr, pgen, representations = self.sources() + pgen.loc[1, "log10_pgen_1mm"] = -np.inf + table, report = prepare_identity_table( + airr, pgen, representations, min_len=1, max_len=40 + ) + self.assertEqual(report["dropped_non_finite_pgen"], 1) + splits, overlaps = make_splits( + table, train_fraction=0.5, val_fraction=0.25, seed=42 + ) + self.assertEqual({name: len(rows) for name, rows in splits.items()}, {"train": 1, "val": 0, "test": 2}) + self.assertEqual(overlaps, {"train_val": 0, "train_test": 0, "val_test": 0}) + + def test_standardizer_uses_only_train_rows(self): + airr, pgen, representations = self.sources() + table, _ = prepare_identity_table( + airr, pgen, representations, min_len=1, max_len=40 + ) + embeddings = np.arange(20, dtype=np.float32).reshape(4, 5) + with tempfile.TemporaryDirectory() as temporary: + output = Path(temporary) / "target_standardizer.npz" + summary = write_standardizer( + output, + table, + embeddings, + np.array([0, 2]), + chunk_size=1, + ) + with np.load(output, allow_pickle=False) as standardizer: + np.testing.assert_allclose( + standardizer["tcremp_mean"], embeddings[[0, 2]].mean(axis=0) + ) + self.assertEqual(int(standardizer["train_rows"]), 2) + self.assertEqual(str(standardizer["pgen_target"]), "log10_pgen_1mm") + self.assertEqual(summary["fit_split"], "train") + + def test_identity_normalizes_gene_alleles(self): + table = add_identity_columns( + pd.DataFrame( + {"cdr3": [" cassf "], "v": ["TRBV7-6*01"], "j": ["TRBJ2-7*01"]} + ), + cdr3_col="cdr3", + v_col="v", + j_col="j", + ) + self.assertEqual(table.loc[0, "junction_aa"], "CASSF") + self.assertEqual(table.loc[0, "v_identity"], "TRBV7-6") + self.assertEqual(table.loc[0, "j_identity"], "TRBJ2-7") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_single_aa_challenge.py b/tests/test_single_aa_challenge.py new file mode 100644 index 0000000..91f3894 --- /dev/null +++ b/tests/test_single_aa_challenge.py @@ -0,0 +1,69 @@ +import numpy as np +import pandas as pd + +from benchmark.prepare_single_aa_challenge import choose_parents, generate_mutants, parent_folds +from benchmark.validate_single_aa_targets import align_targets + + +def test_mutants_are_unique_one_substitution_and_parent_split_is_disjoint(): + parents = pd.DataFrame({ + "row_index": [10, 20, 30, 40, 50], + "junction_aa": ["CASSF", "CASRF", "CATQF", "CARGF", "CSAWF"], + "v_call": ["TRBV1"] * 5, + "j_call": ["TRBJ1-1"] * 5, + }) + benchmark = set(parents["junction_aa"]) + mutants = generate_mutants(parents, benchmark, variants_per_parent=3, seed=42) + assert len(mutants) == 15 + assert mutants["junction_aa"].nunique() == len(mutants) + assert not (set(mutants["junction_aa"]) & benchmark) + for row in mutants.itertuples(index=False): + assert len(row.parent_sequence) == len(row.junction_aa) + assert sum(a != b for a, b in zip(row.parent_sequence, row.junction_aa)) == 1 + assert row.mutation_position_zero_based not in (0, len(row.junction_aa) - 1) + folds = parent_folds(mutants["parent_id"].unique(), seed=42) + mutants["fold"] = mutants["parent_id"].map(folds) + by_fold = {fold: set(group["parent_id"]) for fold, group in mutants.groupby("fold")} + assert by_fold["train"].isdisjoint(by_fold["val"]) + assert by_fold["train"].isdisjoint(by_fold["test"]) + assert by_fold["val"].isdisjoint(by_fold["test"]) + + +def test_length_stratified_parent_limit_is_exact_even_with_more_strata_than_limit(): + table = pd.DataFrame({ + "row_index": np.arange(20), + "junction_aa": ["C" + "A" * length + "F" for length in range(2, 22)], + }) + chosen = choose_parents(table, np.arange(20), limit=5, seed=42) + assert len(chosen) == 5 + assert chosen["row_index"].nunique() == 5 + + +def test_validator_accepts_exact_zero_but_requires_finite_1mm_target(tmp_path): + prepared = tmp_path / "prepared" + prepared.mkdir() + mutants = pd.DataFrame({ + "clone_id": ["a", "b"], "junction_aa": ["CASSF", "CASRF"], + "parent_id": ["p1", "p2"], "parent_sequence": ["CASAF", "CASSF"], + }) + mutants.to_csv(prepared / "mutants.tsv", sep="\t", index=False) + pgen = mutants[["clone_id", "junction_aa"]].copy() + pgen["log10_pgen"] = [-np.inf, -5.0] + pgen["log10_pgen_1mm"] = [-4.0, -4.5] + pgen_path = tmp_path / "pgen.tsv" + pgen.to_csv(pgen_path, sep="\t", index=False) + representations = pd.DataFrame({ + "clone_id": ["a", "b"], "cdr3aa_TRB": ["CASSF", "CASRF"], + }) + reps_path = tmp_path / "representations.tsv" + representations.to_csv(reps_path, sep="\t", index=False) + tcremp = pd.DataFrame(np.zeros((2, 9000), dtype=np.float32)) + tcremp.insert(0, "clone_id", ["a", "b"]) + tcremp_path = tmp_path / "tcremp.parquet" + tcremp.to_parquet(tcremp_path, index=False) + report = align_targets(prepared, pgen_path, reps_path, tcremp_path, tmp_path / "targets") + assert report["finite_pgen_target"] is True + assert report["nonfinite_exact_log10_pgen_diagnostic_rows"] == 1 + accepted = pd.read_csv(tmp_path / "targets" / "targets.tsv", sep="\t") + assert "log10_pgen" not in accepted + assert np.isfinite(accepted["log10_pgen_1mm"]).all() diff --git a/tests/test_single_aa_evaluation.py b/tests/test_single_aa_evaluation.py new file mode 100644 index 0000000..ca90733 --- /dev/null +++ b/tests/test_single_aa_evaluation.py @@ -0,0 +1,62 @@ +import numpy as np +import pandas as pd + +from benchmark.evaluate_single_aa_challenge import ( + bootstrap_or_disabled, + evaluate_representation, + identity_features, + ridge_fit, + ridge_predict, +) + + +def test_zero_bootstrap_is_explicitly_disabled_without_resampling(): + frame = pd.DataFrame({"parent_id": ["p1"], "probe_split": ["test"]}) + disabled = bootstrap_or_disabled( + frame, + np.array([0.0]), + np.zeros((1, 2)), + {}, + np.array([True]), + repeats=0, + seed=42, + ) + assert disabled == { + "status": "disabled_by_user", + "repeats": 0, + "uncertainty_available": False, + "significance_claim_permitted": False, + } + + +def test_ridge_and_delta_evaluator_are_finite(): + rng = np.random.default_rng(42) + parents = np.repeat([f"p{i}" for i in range(15)], 3) + split = np.repeat(["train"] * 9 + ["val"] * 3 + ["test"] * 3, 3) + n = len(parents) + frame = pd.DataFrame({ + "parent_id": parents, + "probe_split": split, + "mutation_position_zero_based": np.tile([1, 2, 3], 15), + "from_aa": np.tile(["A", "C", "D"], 15), + "to_aa": np.tile(["G", "H", "I"], 15), + "length": 8, + }) + x = rng.normal(size=(n, 8)) + pgen = x[:, 0] - 0.5 * x[:, 1] + tcremp = x @ rng.normal(size=(8, 12)) + masks = {name: split == name for name in ("train", "val", "test")} + metrics, predictions = evaluate_representation("toy", x, frame, pgen, tcremp, masks) + assert metrics["pgen"]["r2"] > 0.9 + assert metrics["tcremp"]["mean_delta_cosine"] > 0.9 + assert np.isfinite(predictions["tcremp"]).all() + assert identity_features(frame, False).shape == (n, 20) + assert identity_features(frame, True).shape == (n, 100) + + +def test_ridge_preserves_multioutput_shape(): + x = np.arange(40, dtype=float).reshape(10, 4) + y = np.column_stack((x[:, 0], x[:, 1] - x[:, 2])) + pred = ridge_predict(ridge_fit(x, y, 1.0), x) + assert pred.shape == y.shape + assert np.isfinite(pred).all() diff --git a/tests/test_tcremp_geometry_loss.py b/tests/test_tcremp_geometry_loss.py new file mode 100644 index 0000000..e05e005 --- /dev/null +++ b/tests/test_tcremp_geometry_loss.py @@ -0,0 +1,29 @@ +import unittest + +import torch + +from irrm_codec.multitask_losses import IRRMCodecMultiTaskLoss + + +class TCRemPGeometryLossTest(unittest.TestCase): + def test_pairwise_loss_is_zero_for_exact_prediction(self): + target = torch.tensor([[1.0, 2.0, 3.0], [2.0, 1.0, 4.0], [3.0, 4.0, 1.0]]) + loss = IRRMCodecMultiTaskLoss._pairwise_log_cosine_distance(target, target) + self.assertAlmostEqual(float(loss), 0.0, places=7) + + def test_pairwise_loss_detects_collapsed_geometry(self): + target = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]) + prediction = torch.ones_like(target) + loss = IRRMCodecMultiTaskLoss._pairwise_log_cosine_distance(prediction, target) + self.assertGreater(float(loss), 1.0) + + def test_geometry_weights_are_validated(self): + with self.assertRaises(ValueError): + IRRMCodecMultiTaskLoss( + tcremp_centered_cosine_weight=0.6, + tcremp_pairwise_log_distance_weight=0.5, + ) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tests/test_vdjdb_cohort.py b/tests/test_vdjdb_cohort.py new file mode 100644 index 0000000..49ca16e --- /dev/null +++ b/tests/test_vdjdb_cohort.py @@ -0,0 +1,64 @@ +import json +import unittest +from pathlib import Path + +from benchmark.evaluate_vdjdb_frozen import parse_models +from benchmark.prepare_vdjdb_epitope_cohort import normalize_vdjdb_row + + +def row(**changes): + value = { + "gene": "TRB", + "cdr3": "CASSLGQETQYF", + "species": "HomoSapiens", + "mhc.a": "HLA-A*02:01", + "mhc.b": "B2M", + "mhc.class": "MHCI", + "antigen.epitope": "GILGFVFTL", + "reference.id": "PMID:1", + "vdjdb.score": "2", + "meta": json.dumps({"subject.id": "donor1", "study.id": "study1"}), + "cdr3fix": json.dumps({"good": True, "cdr3": "CASSLGQETQYF"}), + } + value.update(changes) + return value + + +class NormalizeVDJdbRowTests(unittest.TestCase): + def test_accepts_unmodified_human_trb_and_builds_scoped_donor(self): + normalized, reason = normalize_vdjdb_row(row(), min_score=1) + self.assertEqual(reason, "accepted_row") + self.assertEqual(normalized["donor_key"], "study1|donor1") + self.assertEqual( + normalized["label"], "GILGFVFTL|HLA-A*02:01|B2M|MHCI" + ) + + def test_rejects_sequence_changing_fix(self): + normalized, reason = normalize_vdjdb_row( + row(cdr3fix=json.dumps({"good": True, "cdr3": "CASSIRSSYEQYF"})), + min_score=1, + ) + self.assertIsNone(normalized) + self.assertEqual(reason, "cdr3fix_changes_sequence") + + def test_rejects_missing_donor_and_low_score(self): + normalized, reason = normalize_vdjdb_row( + row(meta="{}"), min_score=1 + ) + self.assertIsNone(normalized) + self.assertEqual(reason, "missing_donor") + normalized, reason = normalize_vdjdb_row(row(**{"vdjdb.score": "0"}), min_score=1) + self.assertIsNone(normalized) + self.assertEqual(reason, "low_or_invalid_score") + + def test_frozen_evaluator_accepts_full_seven_condition_factorial(self): + values = [f"{name}:{name}.pt:{name}.json" for name in ( + "r", "t", "p", "rt", "rp", "tp", "rtp" + )] + parsed = parse_models(values) + self.assertEqual(set(parsed), {"r", "t", "p", "rt", "rp", "tp", "rtp"}) + self.assertEqual(parsed["tp"], (Path("tp.pt"), Path("tp.json"))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_vdjdb_epitope_prediction.py b/tests/test_vdjdb_epitope_prediction.py new file mode 100644 index 0000000..af6b476 --- /dev/null +++ b/tests/test_vdjdb_epitope_prediction.py @@ -0,0 +1,47 @@ +import numpy as np +import pandas as pd + +from benchmark.evaluate_vdjdb_epitope_prediction import ( + audit_dataset, + make_splits, + redcea_component_groups, + score_predictions, +) + + +def test_dataset_audit_detects_duplicates_and_label_conflicts(): + frame = pd.DataFrame({"cdr3": ["CASSF", "CASSF", "CASRF"], "label": ["A", "B", "B"]}) + report, sizes = audit_dataset(frame) + assert report["sequences"] == 3 + assert report["epitopes"] == 2 + assert report["duplicate_cdr3_rows"] == 2 + assert report["identical_cdr3_with_different_labels"] == 1 + assert sizes.set_index("label").loc["B", "sequences"] == 2 + + +def test_redcea_components_merge_rows_sharing_any_source_cluster(): + frame = pd.DataFrame({ + "source_cids": ["a|b", "b|c", "d", "e"], + "cid": ["a", "b", "d", "e"], + }) + groups = redcea_component_groups(frame) + assert groups[0] == groups[1] + assert len(np.unique(groups)) == 3 + + +def test_grouped_splits_have_no_group_leakage_and_shared_exact_alias(): + y = np.repeat(np.arange(2), 12) + exact = np.arange(len(y)) + similarity = np.repeat(np.arange(8), 3) + protocols, report = make_splits(y, exact, similarity, [42], 3, True) + assert protocols["sequence_stratified"] == protocols["exact_cdr3_grouped"] + assert report["exact_cdr3_grouped"]["alias_of"] == "sequence_stratified" + for _, _, train, test in protocols["redcea_similarity_grouped"]: + assert not set(similarity[train]).intersection(similarity[test]) + + +def test_multiclass_metrics_are_perfect_for_perfect_probabilities(): + y = np.array([0, 1, 2, 0, 1, 2]) + probabilities = np.eye(3)[y] + result = score_predictions(y, probabilities, np.arange(3)) + assert all(value == 1.0 for value in result.values()) diff --git a/tests/test_vdjdb_redcea_clonotypes.py b/tests/test_vdjdb_redcea_clonotypes.py new file mode 100644 index 0000000..f035cd3 --- /dev/null +++ b/tests/test_vdjdb_redcea_clonotypes.py @@ -0,0 +1,37 @@ +import pandas as pd + +from benchmark.prepare_vdjdb_redcea_clonotypes import prepare_cohort + + +def member(sequence: str, label: str, cid: str, species: str = "HomoSapiens") -> dict: + return { + "species": species, "antigen.epitope": label, "antigen.gene": "g", + "antigen.species": "virus", "mhc.a": "", "mhc.b": "", "mhc.class": "", + "gene": "TRB", "cdr3aa": sequence, "cid": cid, "csz": "5", + "v.segm": "TRBV1", "j.segm": "TRBJ1", + } + + +def test_preparation_deduplicates_removes_ambiguity_overlap_and_applies_support() -> None: + rows = [ + member("CASSAAAFF", "E1", "E1.1"), + member("CASSAAAFF", "E1", "E1.2"), + member("CASSQQQFF", "E1", "E1.1"), + member("CASSCCCFF", "E2", "E2.1"), + member("CASSCCCFF", "E3", "E3.1"), + member("CASSDDDFF", "E2", "E2.1"), + member("CASSEEEXF", "E2", "E2.1"), # invalid AA X + member("CASSFFFYF", "E2", "E2.1"), + ] + benchmark = {"train": {"CASSDDDFF"}, "val": set(), "test": set()} + cohort, summary, report = prepare_cohort(pd.DataFrame(rows), benchmark, 2, 40) + assert cohort[["label", "cdr3"]].to_dict("records") == [ + {"label": "E1", "cdr3": "CASSAAAFF"}, + {"label": "E1", "cdr3": "CASSQQQFF"}, + ] + assert summary.loc[0, "unique_cdr3"] == 2 + assert cohort.loc[0, "source_rows"] == 2 + assert cohort.loc[0, "source_cids"] == "E1.1|E1.2" + assert report["filter_counts"]["ambiguous_multilabel_cdr3"] == 1 + assert report["filter_counts"]["benchmark_overlap_cdr3"] == 1 + assert report["filter_counts"]["invalid_sequence_rows"] == 1 diff --git a/tests/test_vdjdb_redcea_evaluation.py b/tests/test_vdjdb_redcea_evaluation.py new file mode 100644 index 0000000..c51404d --- /dev/null +++ b/tests/test_vdjdb_redcea_evaluation.py @@ -0,0 +1,36 @@ +import numpy as np +import pandas as pd + +from benchmark.evaluate_vdjdb_redcea_clonotypes import ( + build_candidate_gallery, + candidate_distances, + summarize_candidates, +) + + +def test_perfect_epitope_geometry_scores_perfect_retrieval() -> None: + rows = [] + features = [] + for label_id, label in enumerate(("E1", "E2", "E3")): + for index in range(4): + rows.append({ + "label": label, "species": "HomoSapiens", "length": 10, + "v_call": "TRBV1", "j_call": "TRBJ1", "cdr3": f"C{label}{index}F", + }) + vector = np.zeros(3, dtype=np.float32) + vector[label_id] = 10.0 + vector += index * 1e-3 + features.append(vector) + cohort = pd.DataFrame(rows) + candidates, relevant, _, report = build_candidate_gallery(cohort, 2, 1, 42) + distances, _ = candidate_distances(np.asarray(features), candidates, "cosine") + names = report["label_names"] + label_ids = np.array([names.index(value) for value in cohort["label"]], dtype=np.int32) + summary, per_query, _ = summarize_candidates( + distances, candidates, relevant, label_ids, names + ) + assert summary["macro_epitope"]["precision_at_1"] == 1.0 + assert summary["macro_epitope"]["precision_at_5"] > 0.0 + assert summary["macro_epitope"]["within_between_ratio"] < 0.01 + assert summary["knn_classification"]["k_1"]["macro_f1"] == 1.0 + assert per_query["nearest_same_is_closer"].all() diff --git a/tests/test_vdjdb_viral_motif_top10.py b/tests/test_vdjdb_viral_motif_top10.py new file mode 100644 index 0000000..7e8c09c --- /dev/null +++ b/tests/test_vdjdb_viral_motif_top10.py @@ -0,0 +1,79 @@ +import tempfile +from pathlib import Path + +import pandas as pd + +from benchmark.evaluate_vdjdb_viral_motif_top10 import ( + EXPECTED_WEIGHTS, + MODEL_NAMES, + summarize_distribution, +) +from benchmark.prepare_vdjdb_viral_motif_top10 import ( + choose_top10, + rank_epitopes, + read_species_list, +) + + +def test_species_allowlist_is_exact_and_rejects_duplicates() -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "viral.txt" + path.write_text("# pinned exact values\nInfluenzaA\nSARS-CoV-2\n", encoding="utf-8") + assert read_species_list(path) == ["InfluenzaA", "SARS-CoV-2"] + path.write_text("InfluenzaA\nInfluenzaA\n", encoding="utf-8") + try: + read_species_list(path) + except ValueError: + pass + else: + raise AssertionError("Duplicate species values must be rejected.") + + +def test_top10_is_ranked_by_support_and_relaxation_is_explicit() -> None: + records = [] + for index, count in enumerate(range(50, 38, -1)): + for row in range(count): + records.append({ + "cdr3": f"C{index:02d}{row:03d}", + "epitope": f"PEP{index:02d}", + "antigen_species": "VirusExact", + }) + records.append({"cdr3": "OTHER", "epitope": "PEPX", "antigen_species": "Tumor"}) + ranking, inventory = rank_epitopes(records, {"VirusExact"}) + selected, gate = choose_top10(ranking, requested_threshold=45) + assert [row["epitope"] for row in selected] == [f"PEP{i:02d}" for i in range(10)] + assert gate["threshold_transparently_lowered"] is True + assert gate["effective_minimum_unique_cdr3"] == 41 + assert next(row for row in inventory if row["antigen_species"] == "Tumor") \ + ["classified_viral_by_exact_allowlist"] == 0 + + +def test_epitope_count_shortfall_excludes_singleton_targets() -> None: + ranking = [ + {"epitope": f"E{index}", "unique_cdr3": count} + for index, count in enumerate((50, 46, 40, 5, 4, 2, 1, 1)) + ] + selected, gate = choose_top10(ranking, requested_threshold=40) + assert len(selected) == 6 + assert gate["selected_epitopes"] == 6 + assert gate["evaluable_epitope_shortfall"] == 4 + assert gate["effective_minimum_unique_cdr3"] == 2 + assert gate["viral_epitopes_below_distinct_pair_minimum"] == 2 + + +def test_full_factorial_model_labels_and_distribution_summary() -> None: + assert tuple(EXPECTED_WEIGHTS) == MODEL_NAMES + rows = [] + for epitope_index in range(10): + for model_index, model in enumerate(MODEL_NAMES): + rows.append({ + "epitope": f"E{epitope_index}", + "distance": "cosine", + "metric": "same_label_auroc", + "model": model, + "value": 0.5 + model_index / 100, + }) + summary, wins, win_summary = summarize_distribution(pd.DataFrame(rows)) + assert len(summary) == 7 + assert set(wins["model"]) == {"rtp"} + assert win_summary.loc[win_summary["model"] == "rtp", "win_count"].item() == 10 diff --git a/tests/test_vdjdb_ylq.py b/tests/test_vdjdb_ylq.py new file mode 100644 index 0000000..d7ef3f8 --- /dev/null +++ b/tests/test_vdjdb_ylq.py @@ -0,0 +1,45 @@ +import unittest + +from benchmark.prepare_vdjdb_ylq import build_pairs + + +def record(cdr3: str, donor: str, v_call: str = "TRBV1", j_call: str = "TRBJ1") -> dict: + return { + "cdr3": cdr3, + "donor_key": donor, + "length": len(cdr3), + "v_call": v_call, + "j_call": j_call, + } + + +class YLQPairConstructionTests(unittest.TestCase): + def test_pairs_are_exact_length_and_fully_donor_disjoint(self) -> None: + positives = [ + record("CASS", "study|p1"), + record("CASR", "study|p2"), + record("CAST", "study|p3"), + ] + negatives = [ + record("CATG", "study|n1"), + record("CAVG", "study|n2", v_call="TRBV2"), + ] + pairs, report = build_pairs(positives, negatives, seed=42) + self.assertEqual(len(pairs), 3) + self.assertEqual(report["exact_length_fraction"], 1.0) + for pair in pairs: + self.assertEqual(pair["positive_length"], pair["negative_length"]) + self.assertEqual( + len({pair["query_donor"], pair["positive_donor"], pair["negative_donor"]}), + 3, + ) + + def test_missing_exact_length_negative_is_blocking(self) -> None: + positives = [record("CASS", "p1"), record("CASR", "p2")] + negatives = [record("CASSL", "n1")] + with self.assertRaisesRegex(ValueError, "No exact-length donor-disjoint negative"): + build_pairs(positives, negatives, seed=42) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_vdjdb_ylq_motif_pooled.py b/tests/test_vdjdb_ylq_motif_pooled.py new file mode 100644 index 0000000..2c6df16 --- /dev/null +++ b/tests/test_vdjdb_ylq_motif_pooled.py @@ -0,0 +1,129 @@ +import csv +import json +import tempfile +import unittest +from pathlib import Path + +from benchmark.prepare_vdjdb_ylq_motif_pooled import ( + build_pairs, + read_joined_records, + read_official_members, + select_compact_negatives, +) + + +class OfficialMotifPooledTests(unittest.TestCase): + def test_official_members_require_exact_context_and_chain(self) -> None: + fields = [ + "species", "antigen.epitope", "mhc.a", "mhc.b", "mhc.class", "gene", + "cdr3aa", "cid", "csz", "v.segm", "j.segm", + ] + rows = [ + ["HomoSapiens", "YLQPRTFLL", "HLA-A*02:01", "B2M", "MHCI", "TRB", + "CASST", "H.B.YLQPRTFLL.1", "5", "TRBV1", "TRBJ1"], + ["HomoSapiens", "YLQPRTFLL", "HLA-A*02", "B2M", "MHCI", "TRB", + "CASSA", "wrong_mhc", "5", "TRBV1", "TRBJ1"], + ["HomoSapiens", "YLQPRTFLL", "HLA-A*02:01", "B2M", "MHCI", "TRA", + "CAVAA", "wrong_chain", "5", "TRAV1", "TRAJ1"], + ] + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "cluster_members.txt" + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.writer(handle, delimiter="\t") + writer.writerow(fields) + writer.writerows(rows) + membership, report = read_official_members(path) + self.assertEqual(membership, {("CASST", "YLQPRTFLL"): {"H.B.YLQPRTFLL.1"}}) + self.assertEqual(report["counts"]["exact_context_trb_rows"], 1) + + def test_join_pools_donors_and_removes_ambiguous_cdr3(self) -> None: + fields = [ + "gene", "cdr3", "v.segm", "j.segm", "species", "mhc.a", "mhc.b", + "mhc.class", "antigen.epitope", "reference.id", "vdjdb.score", "meta", "cdr3fix", + ] + + def row(cdr3: str, epitope: str, donor: str) -> list[str]: + return [ + "TRB", cdr3, "TRBV1", "TRBJ1", "HomoSapiens", "HLA-A*02:01", "B2M", + "MHCI", epitope, "PMID:1", "1", + json.dumps({"study.id": "s", "subject.id": donor}), + json.dumps({"good": True, "cdr3": cdr3}), + ] + + rows = [ + row("CASST", "YLQPRTFLL", "d1"), + row("CASST", "YLQPRTFLL", "d2"), + row("CASSA", "YLQPRTFLL", "d3"), + row("CASSA", "OTHER", "d3"), + row("CASSG", "OTHER", "d4"), + ] + membership = { + ("CASST", "YLQPRTFLL"): {"ylq"}, + ("CASSA", "YLQPRTFLL"): {"ylq"}, + ("CASSA", "OTHER"): {"other"}, + ("CASSG", "OTHER"): {"other"}, + } + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "vdjdb.txt" + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.writer(handle, delimiter="\t") + writer.writerow(fields) + writer.writerows(rows) + records, report = read_joined_records( + path, membership, {"train": set(), "val": set(), "test": set()}, 1.0 + ) + self.assertEqual({record["cdr3"] for record in records}, {"CASST", "CASSG"}) + pooled = next(record for record in records if record["cdr3"] == "CASST") + self.assertEqual(pooled["source_rows"], 2) + self.assertEqual(pooled["donor_count_descriptive_only"], 2) + self.assertEqual(report["filter_counts"]["ambiguous_multilabel_cdr3"], 1) + + def test_pair_builder_has_no_donor_dependency_and_exact_length_negative(self) -> None: + def record(cdr3: str) -> dict: + return {"cdr3": cdr3, "length": len(cdr3), "v_call": "TRBV1", "j_call": "TRBJ1"} + + positives = [record("CASS"), record("CASR"), record("CAST")] + negatives = [record("CAVG"), record("CATG")] + pairs, report = build_pairs(positives, negatives, seed=42) + self.assertEqual(len(pairs), 3) + self.assertEqual(report["donor_constraint"], "none_pooled_by_user_protocol") + for pair in pairs: + self.assertEqual(pair["positive_length"], pair["negative_length"]) + self.assertEqual(len({pair["query_cdr3"], pair["positive_cdr3"], pair["negative_cdr3"]}), 3) + + def test_compact_controls_are_unique_capped_and_same_length(self) -> None: + def record(cdr3: str, v: str = "TRBV1", j: str = "TRBJ1") -> dict: + return {"cdr3": cdr3, "length": len(cdr3), "v_call": v, "j_call": j} + + positives = [record("CASS"), record("CASR"), record("CASST")] + negatives = [ + record("CAVG"), record("CATG"), record("CAGG"), record("CAAG"), + record("CGGG"), record("CTTG"), record("CAVGT"), + ] + selected, assignments, strata, report = select_compact_negatives( + positives, negatives, per_positive=2, seed=42 + ) + self.assertEqual(len(selected), 5) + self.assertEqual(len({row["cdr3"] for row in selected}), 5) + self.assertLessEqual(len(selected), 2 * len(positives)) + self.assertTrue(all(len(row["positive_cdr3"]) == len(row["control_cdr3"]) + for row in assignments)) + length_five = next(row for row in strata if row["length"] == 5) + self.assertEqual(length_five["shortfall"], 1) + self.assertEqual(report["gallery_control_reuse"], 0) + + def test_pairs_orient_away_from_unsupported_control_length(self) -> None: + def record(cdr3: str) -> dict: + return {"cdr3": cdr3, "length": len(cdr3), "v_call": "TRBV1", "j_call": "TRBJ1"} + + positives = [record("CASS"), record("CASRT"), record("CASST")] + negatives = [record("CAVG")] + pairs, report = build_pairs(positives, negatives, seed=42) + self.assertEqual(len(pairs), 2) + self.assertTrue(all(pair["positive_length"] == 4 for pair in pairs)) + self.assertEqual(report["excluded_pairs_both_endpoints_lack_exact_length_controls"], 1) + self.assertEqual(report["forced_orientation_to_supported_comparator"], 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_wordpiece_sweep.py b/tests/test_wordpiece_sweep.py new file mode 100644 index 0000000..d768955 --- /dev/null +++ b/tests/test_wordpiece_sweep.py @@ -0,0 +1,71 @@ +import tempfile +import unittest +from pathlib import Path + +import numpy as np +from tokenizers import Tokenizer +from tokenizers.decoders import WordPiece as WordPieceDecoder +from tokenizers.models import WordPiece +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.trainers import WordPieceTrainer + +from irrm_codec.multitask_data import resolve_encoder_tokenizer +from irrm_codec.tokenization import AA_VOCAB +from irrm_codec.wordpiece_tokenization import validate_wordpiece_tokenizer +from scripts.prepare_wordpiece_sweep import SPECIAL_TOKENS, audit_tokenizer + + +class WordPieceSweepTest(unittest.TestCase): + def _tokenizer(self, sequences, vocab_size=44): + tokenizer = Tokenizer(WordPiece(unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.decoder = WordPieceDecoder(prefix="##") + tokenizer.train_from_iterator( + sequences, + WordPieceTrainer( + vocab_size=vocab_size, + min_frequency=1, + special_tokens=SPECIAL_TOKENS, + continuing_subword_prefix="##", + show_progress=False, + ), + ) + return tokenizer + + def test_special_ids_roundtrip_and_separate_encoder_vocab(self): + sequences = ["CASSLGQETQYF", "CASSPGQETQYF", "CASSIRSSYEQYF"] + tokenizer = self._tokenizer(sequences) + validate_wordpiece_tokenizer(tokenizer, "test") + for sequence in sequences: + encoding = tokenizer.encode(sequence) + self.assertNotIn(3, encoding.ids) + self.assertEqual(tokenizer.decode(encoding.ids), sequence) + with tempfile.TemporaryDirectory() as temporary: + path = Path(temporary) / "tokenizer.json" + tokenizer.save(str(path)) + resolved = resolve_encoder_tokenizer("wordpiece", str(path)) + self.assertEqual(resolved.name, "wordpiece") + self.assertEqual(resolved.vocab_size, tokenizer.get_vocab_size()) + self.assertNotEqual(resolved.vocab_size, len(AA_VOCAB)) + + def test_audit_uses_train_only_for_pgen_association(self): + sequences = { + "train": ["CASSF", "CASRF", "CATGF", "CQQYF"], + "val": ["CASSF"], + "test": ["CASRF"], + } + tokenizer = self._tokenizer(sum(sequences.values(), [])) + summary, rows = audit_tokenizer( + tokenizer, + sequences, + np.asarray([-8.0, -7.0, -6.0, -5.0], dtype=np.float64), + ) + self.assertEqual(summary["pgen_association"]["split"], "train") + self.assertEqual(summary["pgen_association"]["target"], "log10_pgen_1mm") + self.assertTrue(rows) + self.assertEqual(summary["encoding"]["val"]["unknown_sequences"], 0) + self.assertEqual(summary["encoding"]["test"]["roundtrip_mismatches"], 0) + + +if __name__ == "__main__": + unittest.main()