feat: build leakage-aware MitoInteract recovery pipeline

AUDIT.md

@@ -13,10 +13,11 @@ The published checkpoint is a failed proof of concept, not a scientifically usab
 | Training set was only 32 pairs | `config.json`: `max_train_samples: 32` | Cannot support the claim that the model learned from the 1.9M-pair dataset |
 | Validation set was only 16 pairs | `max_val_samples: 16` | Correlation metrics are unstable and not representative |
 | Reported correlation is negative | Checkpoint: Pearson `-0.9107416`, Spearman `-0.5` | Predictions are directionally wrong on the tiny validation set |
-| Regression error is enormous | Checkpoint: RMSE `7.0257`, MAE `6.8183` | Error is roughly the scale of the entire useful pKd range |
+| Regression error is enormous | Checkpoint: RMSE `7.0257`, MAE `6.8183` | Error is roughly the scale of the entire reported target range |
 | Training barely began | 32 samples / batch 4 × 2 epochs = 16 steps, but warmup is 500 steps | The run completed only 3.2% of warmup; approximate final LR was `3.2e-05`, not the configured `1e-3` |
 | Model card contradicts checkpoint | README says ESM-2 650M; config/checkpoint use ESM-2 35M | Architecture and reproducibility claims are incorrect |
 | Target semantics are undocumented | Dataset has `affinity_uM`, `neg_log10_affinity_M`, and normalized `affinity`; config names none | Published inference labels raw output as pKd without proving it was trained on pKd |
+| Source target conflates measurement types | The pinned BindingDB notebook combines available `Ki`, `Kd`, `IC50`, and `EC50` values and averages them in linear µM | The dataset target is generic pAffinity, not pKd; assay provenance must be retained in any scientific successor |
 | Evaluation split is undocumented | No training/evaluation code is published | Protein/ligand leakage cannot be ruled out |
 | Reproduction assets are missing | No training script, requirements, seed, split manifest, predictions, or metric calculation | Results cannot be independently reproduced |
 | Serialization is unsafe and oversized | `full_model.pt` is a 315.7 MB pickle checkpoint loaded with `weights_only=False` | Avoidable code-execution risk and duplicated frozen encoders |
@@ -27,7 +28,7 @@ The separate `mitointeract_weights.pt` file contains 1,278,980 trainable head pa
 
 MitoInteract v2 does not earn a model-card performance claim until all gates pass:
 
-1. **Unit semantics:** pKd conversion tests pass and the exact target column is fixed to `neg_log10_affinity_M`.
+1. **Unit semantics:** pAffinity conversion tests pass, the exact target column is fixed to `neg_log10_affinity_M`, and the model card does not relabel mixed measurements as pKd.
 2. **Overfit sanity check:** the trainable head can overfit a fixed 32-example set to near-zero training error.
 3. **Real validation:** at least 1,000 validation examples, with a published split manifest.
 4. **Leakage-aware evaluation:** report random-pair, cold-protein, and cold-ligand/scaffold results separately.
@@ -40,7 +41,7 @@ MitoInteract v2 does not earn a model-card performance claim until all gates pas
 For the first bounded experiment, use ESM-2 35M plus ChemBERTa with frozen encoders and a simpler regression head. On a fixed 10,000-pair development sample:
 
 - validation/test size: at least 1,000 each
-- primary target: pKd (`neg_log10_affinity_M`)
+- primary target: generic pAffinity (`neg_log10_affinity_M`), not pKd
 - primary split: cold-protein
 - RMSE: lower than training-mean baseline
 - Pearson: `> 0.30`
@@ -48,3 +49,21 @@ For the first bounded experiment, use ESM-2 35M plus ChemBERTa with frozen encod
 - no NaNs or unit-conversion failures
 
 These are recovery gates, not publishable state-of-the-art claims.
+
+## Recovery evidence as of 2026-07-19
+
+The deterministic 10,000-pair development sample is now backed by the immutable dataset revision `11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b`. The original Parquet file was verified against SHA-256 `d50f4cae794298f32c39f75a9e8603528d20e50e4fa757025def412f5b2a78d8`.
+
+| Gate | Result |
+|---|---|
+| pAffinity/µM numeric conversion | passed, while preserving the mixed-measurement warning |
+| Python test suite | passed, 33 tests with the optional model dependencies installed |
+| 32-pair synthetic-embedding overfit | passed |
+| 32-pair pinned-encoder overfit | passed, RMSE below `1e-6` pAffinity |
+| Random-pair combined ridge | RMSE `1.3358`, Pearson `0.5371`, Spearman `0.5187` |
+| Exact protein-disjoint combined ridge | RMSE `1.3462`, Pearson `0.4384`, Spearman `0.4283` |
+| MMseqs2 50% similarity-disjoint combined ridge | RMSE `1.4349`, Pearson `0.3625`, Spearman `0.3652` |
+| Scaffold-disjoint combined ridge | RMSE `1.3649`, Pearson `0.5198`, Spearman `0.5354` |
+| Full 10,000-pair neural head | pending GPU embedding run |
+
+The MMseqs2 split clustered 2,255 unique proteins into 1,421 groups at 50% sequence identity and 80% bidirectional coverage, then held whole clusters out. The combined ridge model beat its train-mean RMSE of `1.5336`, but the ligand-only ridge slightly outperformed the combined model on this split. This is evidence of dataset shortcut risk, not evidence that the recovery model has learned general protein-ligand interaction biology.

README.md

@@ -4,91 +4,138 @@ tags:
 - biology
 - protein
 - drug-target-interaction
-- mitochondria
-- apoptosis
 - binding-affinity
 - esm2
 - chemberta
 datasets:
 - jglaser/binding_affinity
-metrics:
-- pearsonr
-- spearmanr
-- rmse
-- mae
 ---
 
-# MitoInteract: Protein-Molecule Binding Affinity Prediction for Mitochondrial Apoptosis Research
+# MitoInteract
 
-## Overview
+> **Research recovery in progress. The published v1 checkpoint is deprecated and must not be used to rank compounds or guide experiments.**
 
-MitoInteract is a dual-encoder model that predicts **binding affinity (pKd)** between any protein and any molecule. 
-It combines:
-- **ESM-2 650M** (protein encoder) for protein sequence understanding
-- **ChemBERTa** (molecule encoder) for SMILES-based molecular representation
-- **Bidirectional cross-attention** fusion layer
-- **4-layer MLP** regression head
+MitoInteract is an experimental protein-ligand binding-affinity project. The recovery target is **pAffinity**, defined as `-log10(the dataset's affinity value in molar units)`. Higher values indicate a numerically stronger reported affinity.
 
-## Intended Use
+This target is **not pKd**. The pinned dataset's BindingDB preprocessing combines available `Ki`, `Kd`, `IC50`, and `EC50` values and, when more than one exists, averages them in linear micromolar units. Those measurements are not scientifically interchangeable. The generic `pAffinity` name preserves the source semantics without pretending every row is a dissociation constant.
 
-This model is designed for **mitochondrial apoptosis research**, enabling researchers to:
-- Predict how ceramides interact with mitochondrial membrane proteins (VDAC1, VDAC2)
-- Screen BCL-2 family protein interactions with BH3 mimetic drugs (venetoclax, navitoclax, ABT-737)
-- Explore protein-lipid interactions in the apoptosis pathway
-- Run in-silico binding experiments before wet-lab validation
+The first published checkpoint did not complete a meaningful training run. Its validation correlations were negative, its evaluation set contained only 16 pairs, and its model card overstated the architecture and intended use. The weights remain in this repository for provenance only. See [`AUDIT.md`](AUDIT.md) for the artifact-level evidence.
 
-## Quick Start
+## Current status
 
-```python
-from model import load_model, predict_binding
+The recovery branch now provides:
 
-# Load model
-model, config = load_model("full_model.pt", device="cuda")
+- explicit and tested pAffinity/µM conversions;
+- a dataset revision pinned to `11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b`;
+- deterministic sampling, pair deduplication, and split manifests;
+- Bemis-Murcko scaffold grouping with an exact-ligand fallback for acyclic molecules;
+- MMseqs2 protein clustering at 50% identity and 80% bidirectional coverage;
+- train-mean, protein-only, ligand-only, and combined character n-gram ridge baselines;
+- a corrected LayerNorm-based v2 head with standardized targets;
+- pinned ESM-2 and ChemBERTa revisions, long-protein chunking, and cached frozen embeddings;
+- safe `safetensors` output for new trainable weights;
+- unit tests and a successful 32-pair end-to-end overfit sanity check.
 
-# Predict ceramide C16 binding to VDAC1
-result = predict_binding(
-    model,
-    protein_seq="MPPYLTFGLKAGALLPLTLPYVRAEAVTKLKLTLNAFEGASK...",  # VDAC1
-    smiles="CCCCCCCCCCCCCCCC(=O)N[C@@H](CO)[C@H](O)/C=C/CCCCCCCCCCCCC",  # Ceramide C16
-    device="cuda"
-)
-print(f"Predicted pKd: {result['pKd']:.3f}")
-print(f"Predicted Kd: {result['Kd_uM']:.3f} µM")
+No v2 neural checkpoint has earned a performance claim yet.
+
+## Reproducible development baseline
+
+These are **development results**, not an external benchmark. They use a deterministic 10,000-pair sample from `jglaser/binding_affinity`. Each test split contains 1,000 pairs.
+
+| Split | Model | Test RMSE ↓ | Pearson ↑ | Spearman ↑ |
+|---|---|---:|---:|---:|
+| Random pair | train mean | 1.5650 | undefined | undefined |
+| Random pair | combined character ridge | **1.3358** | **0.5371** | **0.5187** |
+| Exact protein-disjoint | train mean | 1.4951 | undefined | undefined |
+| Exact protein-disjoint | combined character ridge | **1.3462** | **0.4384** | **0.4283** |
+| Protein-similarity-disjoint, MMseqs2 50% | train mean | 1.5336 | undefined | undefined |
+| Protein-similarity-disjoint, MMseqs2 50% | combined character ridge | **1.4349** | **0.3625** | **0.3652** |
+| Ligand scaffold-disjoint | train mean | 1.5834 | undefined | undefined |
+| Ligand scaffold-disjoint | combined character ridge | **1.3649** | **0.5198** | **0.5354** |
+
+The similarity-disjoint result is the most important current recovery signal. It passes the initial gate of beating the train mean with positive Pearson and Spearman correlation. It also shows that ligand features carry more signal than protein character features on this sample, which is a warning against interpreting random-split performance as learned protein-ligand interaction biology.
+
+## Run the recovery checks
+
+```bash
+cd recovery
+uv sync --dev
+uv run pytest -q
+
+uv run python scripts/dataset_smoke.py \
+  --rows 100 \
+  --output artifacts/dataset-smoke.json
 ```
 
-## Key Apoptosis Targets
-
-| Protein | Role in Apoptosis |
-|---------|-------------------|
-| BCL-2 | Anti-apoptotic, prevents MOMP |
-| BCL-XL | Anti-apoptotic, sequesters BAX/BAK |
-| BAX | Pro-apoptotic, forms pores in outer membrane |
-| BAK | Pro-apoptotic, oligomerizes in membrane |
-| VDAC1 | Voltage-dependent anion channel, ceramide target |
-| Cytochrome c | Released during MOMP, activates caspase cascade |
-
-## Key Molecules
-
-| Molecule | Role |
-|----------|------|
-| Ceramide C16 | Lipid mediator, promotes MOMP via VDAC |
-| Ceramide C2 | Short-chain ceramide analog |
-| Venetoclax | BCL-2 inhibitor (FDA-approved) |
-| Navitoclax | BCL-2/BCL-XL dual inhibitor |
-| ABT-737 | BCL-2/BCL-XL/BCL-w inhibitor |
-| Cardiolipin | Mitochondrial inner membrane lipid |
-
-## Training Details
-
-- **Dataset**: jglaser/binding_affinity (1.9M protein-ligand pairs)
-- **Architecture**: ESM-2 650M (frozen) + ChemBERTa (frozen) + Cross-Attention + MLP
-- **Training**: AdamW, lr=1e-3, cosine schedule, early stopping
-- **Best Validation Pearson R**: -0.9107
-
-## Citation
-
-Based on:
-- BAPULM (arxiv:2411.04150) - frozen encoder + MLP pattern
-- SSM-DTA (arxiv:2206.09818) - CLS cross-attention fusion
-- ESM-2 (arxiv:2202.03555) - protein language model
-- ChemBERTa (arxiv:2010.09885) - molecular language model
+To rebuild the 10,000-pair development data from the converted Parquet file:
+
+```bash
+uv run python scripts/prepare_sample.py \
+  --parquet /path/to/train.parquet \
+  --sample-size 10000 \
+  --seed 42 \
+  --output-dir artifacts/dev-10k
+
+uv run python scripts/cluster_proteins.py \
+  --data-dir artifacts/dev-10k \
+  --min-seq-id 0.5 \
+  --coverage 0.8
+
+uv run python scripts/run_baselines.py \
+  --data-dir artifacts/dev-10k \
+  --output artifacts/dev-10k/baselines.json
+```
+
+MMseqs2 is required for the sequence-similarity split. Generated data, embeddings, predictions, and weights are intentionally excluded from Git.
+
+## v2 architecture under evaluation
+
+The bounded neural recovery experiment uses frozen, revision-pinned encoders:
+
+- protein: `facebook/esm2_t12_35M_UR50D`;
+- ligand: `DeepChem/ChemBERTa-77M-MLM`;
+- fusion: projected protein and ligand vectors plus elementwise product and absolute difference;
+- regression: LayerNorm-based MLP over a standardized pAffinity target;
+- serialization: trainable head in `safetensors`, with encoder names, revisions, and target scaler stored separately.
+
+Proteins longer than ESM-2's context are split into residue chunks and pooled with length weighting instead of silently truncating the sequence. This is a reproducible baseline strategy, not proof that full-sequence pooling captures a binding site.
+
+## Intended use
+
+Current acceptable use:
+
+- reproducing the failure analysis;
+- testing protein-ligand affinity data pipelines;
+- comparing leakage-aware splitting strategies;
+- developing and auditing computational baselines.
+
+Current unacceptable use:
+
+- ranking compounds for laboratory work;
+- estimating clinical activity, safety, dosage, or efficacy;
+- claiming specificity for mitochondrial apoptosis proteins;
+- treating a predicted pAffinity as experimental evidence;
+- loading `full_model.pt` from an untrusted source.
+
+The source dataset combines BindingDB, PDBbind-cn, BioLIP, and BindingMOAD. Its mixed assays, measurement types, source-specific preprocessing, duplicates, and possible censoring limit what a single regression target can support. See [`recovery/DATA_PROVENANCE.md`](recovery/DATA_PROVENANCE.md) for the pinned preprocessing evidence. A publishable model requires source-aware deduplication, stronger external evaluation, uncertainty estimates, and wet-lab validation.
+
+## Legacy v1 facts
+
+| Item | Recovered value |
+|---|---:|
+| Training pairs | 32 |
+| Validation pairs | 16 |
+| Optimizer steps | approximately 16 |
+| Configured warmup steps | 500 |
+| Pearson | -0.9107 |
+| Spearman | -0.5000 |
+| RMSE | 7.0257 in the undocumented legacy target scale |
+| MAE | 6.8183 in the undocumented legacy target scale |
+
+The original README claimed ESM-2 650M, while the checkpoint used ESM-2 35M. The original `full_model.pt` is a Python pickle loaded with `weights_only=False`; it is retained only as a legacy artifact.
+
+## References
+
+- Lin et al., [Language models of protein sequences at the scale of evolution](https://doi.org/10.1101/2022.07.20.500902)
+- Chithrananda et al., [ChemBERTa](https://arxiv.org/abs/2010.09885)
+- Steinegger and Söding, [MMseqs2](https://doi.org/10.1038/nbt.3988)

model.py

@@ -1,10 +1,15 @@
+"""Deprecated MitoInteract v1 architecture retained for artifact forensics.
+
+The v1 checkpoint failed its evaluation gates and uses unsafe pickle
+serialization. New work lives under ``recovery/``. Loading and inference are
+disabled by default so this module cannot silently present an undocumented
+legacy score as pKd.
 """
-MitoInteract Model Class Definition
-Copy this file to load the model for inference.
-"""
+
 import torch
 import torch.nn as nn
-from transformers import EsmModel, EsmTokenizer, AutoModel, AutoTokenizer
+from transformers import EsmModel, AutoModel, AutoTokenizer
+
 
 class MitoInteract(nn.Module):
     def __init__(
@@ -25,23 +30,47 @@ class MitoInteract(nn.Module):
         self.mol_encoder = AutoModel.from_pretrained(mol_model_name)
         self.mol_dim = mol_dim
         if freeze_encoders:
-            for p in self.esm.parameters(): p.requires_grad = False
-            for p in self.mol_encoder.parameters(): p.requires_grad = False
+            for p in self.esm.parameters():
+                p.requires_grad = False
+            for p in self.mol_encoder.parameters():
+                p.requires_grad = False
         self.prot_proj = nn.Sequential(
-            nn.Linear(protein_dim, proj_dim), nn.LayerNorm(proj_dim), nn.ReLU(), nn.Dropout(dropout))
+            nn.Linear(protein_dim, proj_dim),
+            nn.LayerNorm(proj_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+        )
         self.mol_proj = nn.Sequential(
-            nn.Linear(mol_dim, proj_dim), nn.LayerNorm(proj_dim), nn.ReLU(), nn.Dropout(dropout))
-        self.cross_attn_mol2prot = nn.MultiheadAttention(proj_dim, n_heads, dropout=dropout, batch_first=True)
-        self.cross_attn_prot2mol = nn.MultiheadAttention(proj_dim, n_heads, dropout=dropout, batch_first=True)
+            nn.Linear(mol_dim, proj_dim),
+            nn.LayerNorm(proj_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+        )
+        self.cross_attn_mol2prot = nn.MultiheadAttention(
+            proj_dim, n_heads, dropout=dropout, batch_first=True
+        )
+        self.cross_attn_prot2mol = nn.MultiheadAttention(
+            proj_dim, n_heads, dropout=dropout, batch_first=True
+        )
         self.ln_mol2prot = nn.LayerNorm(proj_dim)
         self.ln_prot2mol = nn.LayerNorm(proj_dim)
         fused_dim = proj_dim * 2
         self.mlp = nn.Sequential(
-            nn.Linear(fused_dim, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(dropout),
-            nn.Linear(512, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(dropout),
-            nn.Linear(256, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(dropout),
-            nn.Linear(128, 1))
-    
+            nn.Linear(fused_dim, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, 128),
+            nn.BatchNorm1d(128),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(128, 1),
+        )
+
     def encode_protein(self, input_ids, attention_mask):
         ctx = torch.no_grad() if self.freeze_encoders else torch.enable_grad()
         with ctx:
@@ -49,32 +78,44 @@ class MitoInteract(nn.Module):
         mask = attention_mask.unsqueeze(-1).float()
         pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
         return pooled, out.last_hidden_state
-    
+
     def encode_molecule(self, input_ids, attention_mask):
         ctx = torch.no_grad() if self.freeze_encoders else torch.enable_grad()
         with ctx:
             out = self.mol_encoder(input_ids=input_ids, attention_mask=attention_mask)
         return out.pooler_output, out.last_hidden_state
-    
-    def forward(self, prot_input_ids, prot_attention_mask, mol_input_ids, mol_attention_mask):
+
+    def forward(
+        self, prot_input_ids, prot_attention_mask, mol_input_ids, mol_attention_mask
+    ):
         prot_pooled, prot_seq = self.encode_protein(prot_input_ids, prot_attention_mask)
         mol_pooled, mol_seq = self.encode_molecule(mol_input_ids, mol_attention_mask)
         prot_seq_proj = self.prot_proj(prot_seq)
         mol_seq_proj = self.mol_proj(mol_seq)
         prot_q = self.prot_proj(prot_pooled).unsqueeze(1)
         mol_q = self.mol_proj(mol_pooled).unsqueeze(1)
-        prot_pad_mask = (prot_attention_mask == 0)
-        mol_pad_mask = (mol_attention_mask == 0)
-        h_prot2mol, _ = self.cross_attn_prot2mol(prot_q, mol_seq_proj, mol_seq_proj, key_padding_mask=mol_pad_mask)
-        h_mol2prot, _ = self.cross_attn_mol2prot(mol_q, prot_seq_proj, prot_seq_proj, key_padding_mask=prot_pad_mask)
+        prot_pad_mask = prot_attention_mask == 0
+        mol_pad_mask = mol_attention_mask == 0
+        h_prot2mol, _ = self.cross_attn_prot2mol(
+            prot_q, mol_seq_proj, mol_seq_proj, key_padding_mask=mol_pad_mask
+        )
+        h_mol2prot, _ = self.cross_attn_mol2prot(
+            mol_q, prot_seq_proj, prot_seq_proj, key_padding_mask=prot_pad_mask
+        )
         h_prot2mol = self.ln_prot2mol(h_prot2mol.squeeze(1))
         h_mol2prot = self.ln_mol2prot(h_mol2prot.squeeze(1))
         fused = torch.cat([h_prot2mol, h_mol2prot], dim=-1)
         return self.mlp(fused).squeeze(-1)
 
 
-def load_model(checkpoint_path, device="cpu"):
-    """Load trained MitoInteract model."""
+def load_model(checkpoint_path, device="cpu", *, allow_unsafe_legacy=False):
+    """Load the failed v1 pickle only after explicit risk acknowledgement."""
+    if not allow_unsafe_legacy:
+        raise RuntimeError(
+            "MitoInteract v1 is deprecated and full_model.pt requires unsafe pickle "
+            "deserialization. See AUDIT.md. Pass allow_unsafe_legacy=True only for "
+            "controlled artifact forensics, never for untrusted files."
+        )
     checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
     config = checkpoint["config"]
     model = MitoInteract(
@@ -92,21 +133,35 @@ def load_model(checkpoint_path, device="cpu"):
     return model, config
 
 
-def predict_binding(model, protein_seq, smiles, device="cpu"):
-    """Predict binding affinity (pKd) for a protein-molecule pair."""
-    prot_tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
-    mol_tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
-    
-    prot_enc = prot_tokenizer(protein_seq, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    mol_enc = mol_tokenizer(smiles, return_tensors="pt", padding=True, truncation=True, max_length=200)
-    
+def predict_binding(
+    model, protein_seq, smiles, device="cpu", *, allow_invalid_legacy_output=False
+):
+    """Return the undocumented v1 score only after explicit acknowledgement."""
+    if not allow_invalid_legacy_output:
+        raise RuntimeError(
+            "v1 output has undocumented target semantics and failed validation. "
+            "It cannot be interpreted as pKd or converted to Kd. See AUDIT.md."
+        )
+    prot_tokenizer = AutoTokenizer.from_pretrained(model.esm.name_or_path)
+    mol_tokenizer = AutoTokenizer.from_pretrained(model.mol_encoder.name_or_path)
+
+    prot_enc = prot_tokenizer(
+        protein_seq, return_tensors="pt", padding=True, truncation=True, max_length=512
+    )
+    mol_enc = mol_tokenizer(
+        smiles, return_tensors="pt", padding=True, truncation=True, max_length=200
+    )
+
     model = model.to(device)
     with torch.no_grad():
-        pKd = model(
-            prot_enc["input_ids"].to(device), prot_enc["attention_mask"].to(device),
-            mol_enc["input_ids"].to(device), mol_enc["attention_mask"].to(device),
+        legacy_score = model(
+            prot_enc["input_ids"].to(device),
+            prot_enc["attention_mask"].to(device),
+            mol_enc["input_ids"].to(device),
+            mol_enc["attention_mask"].to(device),
         )
-    
-    pKd_val = pKd.item()
-    Kd_uM = 10 ** (-pKd_val) * 1e6
-    return {"pKd": pKd_val, "Kd_uM": Kd_uM}
+
+    return {
+        "legacy_score": legacy_score.item(),
+        "warning": "Do not interpret this value as pKd, Kd, or validated affinity.",
+    }

recovery/DATA_PROVENANCE.md

@@ -0,0 +1,61 @@
+# Dataset target provenance
+
+## Pinned source
+
+- Dataset: `jglaser/binding_affinity`
+- Revision: `11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b`
+- File: `data/all.parquet`
+- SHA-256: `d50f4cae794298f32c39f75a9e8603528d20e50e4fa757025def412f5b2a78d8`
+
+## Why the target is pAffinity, not pKd
+
+The combined Parquet exposes `affinity_uM` and `neg_log10_affinity_M`, but it does not retain the measurement type or source row.
+
+The pinned `bindingdb_single.ipynb` selects:
+
+```python
+['IC50 (nM)', 'Ki (nM)', 'Kd (nM)', 'EC50 (nM)']
+```
+
+Its `to_uM` function converts every available field to micromolar, appends the values to one list, and returns:
+
+```python
+np.mean(vals[~np.isnan(vals)])
+```
+
+It stores the result in one column:
+
+```python
+df_affinity['affinity_uM'] = df_affinity[
+    ['IC50 (nM)', 'Ki (nM)', 'Kd (nM)', 'EC50 (nM)']
+].parallel_apply(to_uM, axis=1)
+```
+
+The pinned `combine_dbs.ipynb` then keeps only `seq`, `smiles`, and `affinity_uM` from each source and computes:
+
+```python
+df_nr['neg_log10_affinity_M'] = 6 - np.log(df_nr['affinity_uM']) / np.log(10)
+```
+
+Therefore:
+
+- the numeric transform is `-log10(affinity in M)`;
+- not every row is a dissociation constant;
+- some BindingDB rows may average heterogeneous measurement types before the logarithm;
+- the final file cannot separate `Ki`, `Kd`, `IC50`, or `EC50`;
+- the final file cannot reconstruct assay conditions, censoring operators, or original source provenance.
+
+MitoInteract recovery calls this value **pAffinity**. It must not be presented as pKd or converted to a physical Kd.
+
+## Consequence for v2
+
+The current 10,000-pair experiment is useful for pipeline recovery and leakage analysis. It is not sufficient for a scientifically defensible affinity predictor. A successor dataset should retain at least:
+
+- measurement type;
+- numeric value, unit, and relation operator (`=`, `<`, `>`, etc.);
+- assay and source identifiers;
+- protein and ligand identifiers;
+- replicate measurements before aggregation;
+- source database and version.
+
+A clean v2 training target should either model one measurement type at a time or use a multi-task/censoring-aware objective that keeps these distinctions explicit.

recovery/README.md

@@ -1,24 +1,171 @@
-# MitoInteract v2 Recovery Spike
+# MitoInteract v2 Recovery
 
-This directory contains the smallest reproducible checks needed before retraining the model.
+This package contains the reproducible data, baseline, embedding, and training pipeline used to recover MitoInteract from the failed v1 checkpoint.
 
-## Run locally
+## 1. Install and test
 
 ```bash
-cd /root/repos/MitoInteract/recovery
+cd recovery
 uv sync --dev
 uv run pytest -q
-uv run python scripts/dataset_smoke.py --rows 100 --output artifacts/dataset-smoke.json
+uv run python scripts/dataset_smoke.py \
+  --rows 100 \
+  --output artifacts/dataset-smoke.json
 ```
 
-The smoke script reads public rows from the Hugging Face Dataset Viewer API, verifies that `neg_log10_affinity_M` agrees with `affinity_uM`, and reports a deterministic train-mean baseline. It is a plumbing check, not the final scientific evaluation.
+The Dataset Viewer smoke test verifies that `neg_log10_affinity_M` agrees numerically with `affinity_uM`. It is a plumbing check, not a performance result.
 
-## Next experiment
+## Target semantics
 
-Use `config.json` to build a fixed 10,000-pair development sample and report three splits separately:
+The target is generic **pAffinity**, not pKd:
 
-1. random pair
-2. cold protein
-3. cold ligand/scaffold
+```text
+pAffinity = -log10(dataset affinity in molar units)
+```
+
+The pinned source preprocessing combines heterogeneous measurements. In particular, its BindingDB notebook collects `Ki`, `Kd`, `IC50`, and `EC50` and averages the available values in linear micromolar units. These quantities are not interchangeable. The recovery pipeline preserves the dataset's numeric target so the failed run can be audited, but it does not relabel the target as a dissociation constant. See [`DATA_PROVENANCE.md`](DATA_PROVENANCE.md) for the pinned code evidence. A scientifically stronger successor dataset must retain measurement type, assay provenance, relation operators/censoring, and source identifiers.
+
+## 2. Download the pinned source data
+
+The recovery configuration pins `jglaser/binding_affinity` to commit:
+
+```text
+11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b
+```
+
+Download `data/all.parquet` at that revision and verify its SHA-256:
+
+```text
+d50f4cae794298f32c39f75a9e8603528d20e50e4fa757025def412f5b2a78d8
+```
+
+For example:
+
+```bash
+hf download jglaser/binding_affinity data/all.parquet \
+  --repo-type dataset \
+  --revision 11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b \
+  --local-dir .cache/binding-affinity
+
+sha256sum .cache/binding-affinity/data/all.parquet
+```
+
+## 3. Build the fixed development sample
+
+```bash
+uv run python scripts/prepare_sample.py \
+  --parquet .cache/binding-affinity/data/all.parquet \
+  --dataset-revision 11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b \
+  --sample-size 10000 \
+  --candidate-multiplier 3 \
+  --seed 42 \
+  --output-dir artifacts/dev-10k
+```
+
+The script:
+
+- samples source row indexes deterministically;
+- validates the pAffinity/µM conversion;
+- canonicalizes SMILES;
+- aggregates duplicate protein-ligand measurements by median pAffinity;
+- emits pair IDs and entity IDs derived from content hashes;
+- writes 8,000/1,000/1,000 manifests for random-pair, exact-protein, and scaffold-disjoint splits;
+- records row counts, target distributions, sequence lengths, and overlap checks.
+
+Generated data stays under `artifacts/` and is ignored by Git.
+
+## 4. Add the protein-similarity split
+
+Install MMseqs2 through the operating system package manager, then run:
+
+```bash
+uv run python scripts/cluster_proteins.py \
+  --data-dir artifacts/dev-10k \
+  --min-seq-id 0.5 \
+  --coverage 0.8 \
+  --seed 42
+```
+
+The resulting split keeps MMseqs2 clusters disjoint across train, validation, and test. The clustering requires at least 50% actual sequence identity over 80% of both sequences.
+
+## 5. Run non-neural baselines
+
+```bash
+uv run python scripts/run_baselines.py \
+  --data-dir artifacts/dev-10k \
+  --output artifacts/dev-10k/baselines.json
+```
+
+The benchmark reports train-mean, protein-only, ligand-only, and combined character n-gram ridge models. It selects ridge regularization on validation data and reports RMSE, MAE, Pearson, Spearman, and R² on the untouched test split.
+
+These shortcut-sensitive baselines must be run before a transformer model. A neural model that cannot beat them has not justified its cost.
+
+## 6. Verify the v2 trainable head
+
+The default environment keeps PyTorch optional. Install a CPU build for local checks:
+
+```bash
+uv pip install --python .venv/bin/python torch \
+  --index-url https://download.pytorch.org/whl/cpu
+uv pip install --python .venv/bin/python transformers safetensors
+
+.venv/bin/python -m pytest -q
+.venv/bin/python scripts/overfit_sanity.py \
+  --sample artifacts/dev-10k/sample.jsonl \
+  --examples 32 \
+  --max-rmse 0.05 \
+  --output artifacts/overfit-sanity.json
+```
+
+The overfit check isolates the trainable fusion/head path. Passing it does not establish generalization.
+
+## 7. Cache pinned encoder embeddings
+
+```bash
+.venv/bin/python scripts/embed_sample.py \
+  --sample artifacts/dev-10k/sample.jsonl \
+  --batch-size 8 \
+  --device auto \
+  --output artifacts/embeddings-10k.npz
+```
+
+Defaults are pinned to:
+
+- `facebook/esm2_t12_35M_UR50D@6fbf070e65b0b7291e7bbcd451118c216cff79d8`
+- `DeepChem/ChemBERTa-77M-MLM@ed8a5374f2024ec8da53760af91a33fb8f6a15ff`
+
+Protein sequences longer than 1,022 residues are chunked, mean-pooled per chunk, and combined with length weighting. Embeddings are cached so the frozen encoders do not run again for each split or head experiment.
+
+## 8. Train the v2 head
+
+```bash
+.venv/bin/python scripts/train_head.py \
+  --embeddings artifacts/embeddings-10k.npz \
+  --manifest artifacts/dev-10k/split-protein_similarity_50.jsonl \
+  --output-dir artifacts/v2-head-protein-similarity \
+  --epochs 100 \
+  --batch-size 256 \
+  --patience 10 \
+  --device auto
+```
+
+The trainer:
+
+- standardizes pAffinity using training data only;
+- uses AdamW, gradient clipping, 5% warmup, and cosine decay;
+- selects the checkpoint by validation RMSE;
+- saves only trainable weights in `model.safetensors`;
+- records target scaling, encoder revisions, split counts, optimizer settings, history, and test metrics.
+
+## Acceptance gates
+
+A v2 checkpoint is not releasable unless all of these pass:
 
-The first model must beat the train-mean baseline and produce positive Pearson/Spearman correlation before any larger GPU run.
+1. Tests and pAffinity conversion checks pass.
+2. The 32-pair overfit sanity reaches RMSE ≤ 0.05 pAffinity.
+3. Validation and test each contain at least 1,000 pairs.
+4. The published primary result uses the MMseqs2 protein-similarity split, not only a random split.
+5. Test RMSE beats the train mean and both Pearson and Spearman are positive.
+6. The neural model beats or materially complements the combined character ridge baseline.
+7. Dataset, encoder, split, seed, predictions, metrics, and weights are reproducible.
+8. Model-card language remains computational and does not imply experimental or clinical validity.

recovery/cloud/modal_embed.py

@@ -0,0 +1,92 @@
+"""Run the 10k frozen-encoder pass on a temporary Modal L4.
+
+This script is intentionally not run by tests. Build the sample locally first,
+review the current Modal price, then invoke it with:
+
+    modal run cloud/modal_embed.py
+
+Retrieve the result with the `modal volume get` command printed on completion.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+from pathlib import Path
+
+import modal
+
+app = modal.App("mitointeract-v2-embeddings")
+output_volume = modal.Volume.from_name(
+    "mitointeract-v2-artifacts", create_if_missing=True
+)
+cache_volume = modal.Volume.from_name("mitointeract-hf-cache", create_if_missing=True)
+
+recovery_dir = Path(__file__).parents[1]
+image = (
+    modal.Image.debian_slim(python_version="3.12")
+    .pip_install(
+        "numpy>=2.0",
+        "safetensors>=0.5",
+        "torch>=2.4",
+        "transformers>=4.45",
+    )
+    .add_local_dir(recovery_dir / "src", remote_path="/app/src")
+    .add_local_file(
+        recovery_dir / "scripts/embed_sample.py", remote_path="/app/embed_sample.py"
+    )
+    .add_local_file(
+        recovery_dir / "artifacts/dev-10k/sample.jsonl",
+        remote_path="/input/sample.jsonl",
+    )
+)
+
+
+@app.function(
+    image=image,
+    gpu="L4",
+    cpu=4,
+    memory=16_384,
+    timeout=3_600,
+    volumes={"/output": output_volume, "/cache": cache_volume},
+)
+def embed() -> dict:
+    env = os.environ.copy()
+    env["PYTHONPATH"] = "/app/src"
+    env["HF_HOME"] = "/cache/huggingface"
+    subprocess.run(
+        [
+            "python",
+            "/app/embed_sample.py",
+            "--sample",
+            "/input/sample.jsonl",
+            "--batch-size",
+            "32",
+            "--device",
+            "cuda",
+            "--output",
+            "/output/embeddings-10k.npz",
+        ],
+        check=True,
+        env=env,
+    )
+    output_volume.commit()
+    cache_volume.commit()
+    metadata = json.loads(Path("/output/embeddings-10k.json").read_text())
+    metadata["output_bytes"] = Path("/output/embeddings-10k.npz").stat().st_size
+    return metadata
+
+
+@app.local_entrypoint()
+def main() -> None:
+    result = embed.remote()
+    print(json.dumps(result, indent=2))
+    print(
+        "modal volume get mitointeract-v2-artifacts "
+        "embeddings-10k.npz artifacts/embeddings-10k.npz"
+    )
+    print(
+        "modal volume get mitointeract-v2-artifacts "
+        "embeddings-10k.json artifacts/embeddings-10k.json"
+    )

recovery/config.json

@@ -1,8 +1,8 @@
 {
   "dataset": "jglaser/binding_affinity",
-  "dataset_revision": "main",
+  "dataset_revision": "11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b",
   "target_column": "neg_log10_affinity_M",
-  "target_unit": "pKd",
+  "target_unit": "pAffinity",
   "sample_size": 10000,
   "seed": 42,
   "splits": ["random_pair", "cold_protein", "cold_ligand_scaffold"],
@@ -11,8 +11,12 @@
   "minimum_test_rows": 1000,
   "model": {
     "protein_encoder": "facebook/esm2_t12_35M_UR50D",
-    "molecule_encoder": "seyonec/ChemBERTa-zinc-base-v1",
-    "freeze_encoders": true
+    "protein_revision": "6fbf070e65b0b7291e7bbcd451118c216cff79d8",
+    "molecule_encoder": "DeepChem/ChemBERTa-77M-MLM",
+    "molecule_revision": "ed8a5374f2024ec8da53760af91a33fb8f6a15ff",
+    "freeze_encoders": true,
+    "protein_long_sequence_strategy": "length_weighted_chunk_mean",
+    "target": "standardized_pAffinity"
   },
   "acceptance": {
     "beat_train_mean_rmse": true,

recovery/pyproject.toml

@@ -3,7 +3,20 @@ name = "mitointeract-recovery"
 version = "0.1.0"
 description = "Reproducibility and target-semantics checks for MitoInteract v2"
 requires-python = ">=3.11"
-dependencies = []
+dependencies = [
+  "numpy>=2.0",
+  "pyarrow>=18",
+  "rdkit>=2025.3",
+  "scikit-learn>=1.6",
+  "scipy>=1.14",
+]
+
+[project.optional-dependencies]
+model = [
+  "safetensors>=0.5",
+  "torch>=2.4",
+  "transformers>=4.45",
+]
 
 [build-system]
 requires = ["setuptools>=64"]
@@ -16,7 +29,10 @@ where = ["src"]
 package = true
 
 [dependency-groups]
-dev = ["pytest>=9"]
+dev = [
+  "pytest>=9",
+  "ruff>=0.14",
+]
 
 [tool.pytest.ini_options]
 pythonpath = ["src"]

recovery/reports/README.md

@@ -0,0 +1,9 @@
+# Checked recovery reports
+
+These JSON files are small, reviewable outputs from the deterministic recovery pipeline. Large source data, split manifests, embeddings, predictions, and model weights remain under the ignored `artifacts/` directory.
+
+- `dev-10k-audit.json`: immutable source revision and hash, sample statistics, split audits, and hashes for reproducing the sample and manifests.
+- `dev-10k-baselines.json`: validation selection and untouched test metrics for every cheap baseline and split.
+- `overfit-actual-embeddings.json`: 32-pair optimization sanity check using the pinned frozen encoders.
+
+The checked reports support model-card statements but are not an external benchmark or a releasable neural checkpoint.

recovery/reports/dev-10k-audit.json

@@ -0,0 +1,149 @@
+{
+  "source_parquet": "data/all.parquet",
+  "source_sha256": "d50f4cae794298f32c39f75a9e8603528d20e50e4fa757025def412f5b2a78d8",
+  "dataset": "jglaser/binding_affinity",
+  "dataset_revision": "11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b",
+  "source_rows": 1836729,
+  "candidate_rows": 30000,
+  "sample_rows": 10000,
+  "seed": 42,
+  "target": {
+    "column": "neg_log10_affinity_M",
+    "unit": "pAffinity"
+  },
+  "rejected": {},
+  "unique_proteins": 2255,
+  "unique_ligands": 9897,
+  "unique_scaffold_groups": 8463,
+  "duplicate_measurements_in_sample": 51,
+  "protein_length": {
+    "min": 19,
+    "median": 479.0,
+    "max": 7096,
+    "over_512": 4318
+  },
+  "splits": {
+    "random_pair": {
+      "train": {
+        "rows": 8000,
+        "mean_paffinity": 6.5244491582587365,
+        "std_paffinity": 1.5626149739941089,
+        "min_paffinity": 8.881784197001252e-16,
+        "max_paffinity": 14.0
+      },
+      "validation": {
+        "rows": 1000,
+        "mean_paffinity": 6.507777023136616,
+        "std_paffinity": 1.5282972208815573,
+        "min_paffinity": 0.7708500027656555,
+        "max_paffinity": 10.74472713470459
+      },
+      "test": {
+        "rows": 1000,
+        "mean_paffinity": 6.420703814268112,
+        "std_paffinity": 1.5623837036360122,
+        "min_paffinity": 1.2039999961853027,
+        "max_paffinity": 10.95860767364502
+      }
+    },
+    "cold_protein_exact": {
+      "train": {
+        "rows": 8000,
+        "mean_paffinity": 6.476698521107435,
+        "std_paffinity": 1.5651385137888398,
+        "min_paffinity": 0.6000000238418579,
+        "max_paffinity": 14.0
+      },
+      "validation": {
+        "rows": 1000,
+        "mean_paffinity": 6.65395347070694,
+        "std_paffinity": 1.5705924306056336,
+        "min_paffinity": 2.3372421264648438,
+        "max_paffinity": 13.0
+      },
+      "test": {
+        "rows": 1000,
+        "mean_paffinity": 6.656532463908196,
+        "std_paffinity": 1.4850308092775097,
+        "min_paffinity": 8.881784197001252e-16,
+        "max_paffinity": 12.698969841003418
+      }
+    },
+    "cold_ligand_scaffold": {
+      "train": {
+        "rows": 8000,
+        "mean_paffinity": 6.537638315707445,
+        "std_paffinity": 1.5519537903490064,
+        "min_paffinity": 0.6000000238418579,
+        "max_paffinity": 13.214670181274414
+      },
+      "validation": {
+        "rows": 1000,
+        "mean_paffinity": 6.312392803311348,
+        "std_paffinity": 1.5805480574681217,
+        "min_paffinity": 8.881784197001252e-16,
+        "max_paffinity": 13.0
+      },
+      "test": {
+        "rows": 1000,
+        "mean_paffinity": 6.5105747745037075,
+        "std_paffinity": 1.583959954584063,
+        "min_paffinity": 2.0315170288085938,
+        "max_paffinity": 14.0
+      }
+    }
+  },
+  "limitations": [
+    "the source target can combine Ki, Kd, IC50, and EC50; pAffinity is not pKd",
+    "assay type, relation operator, and source provenance are absent from the combined Parquet",
+    "cold_protein_exact is entity-disjoint but not sequence-similarity-clustered",
+    "scaffold groups use Bemis-Murcko with exact-molecule fallback for acyclic ligands",
+    "this development sample is not an external biological benchmark"
+  ],
+  "split_manifests": {
+    "protein_similarity_50": "split-protein_similarity_50.jsonl"
+  },
+  "external_split_audits": {
+    "protein_similarity_50": {
+      "tool": "MMseqs2",
+      "command": [
+        "mmseqs",
+        "easy-cluster",
+        "artifacts/dev-10k/mmseqs-identity-50-coverage-80/proteins.fasta",
+        "artifacts/dev-10k/mmseqs-identity-50-coverage-80/clusters",
+        "artifacts/dev-10k/mmseqs-identity-50-coverage-80/tmp",
+        "--min-seq-id",
+        "0.5",
+        "-c",
+        "0.8",
+        "--cov-mode",
+        "0",
+        "--alignment-mode",
+        "3",
+        "--threads",
+        "8"
+      ],
+      "min_sequence_identity": 0.5,
+      "minimum_alignment_coverage": 0.8,
+      "coverage_mode": "both query and target",
+      "alignment_mode": "actual identical residues / aligned columns",
+      "unique_proteins": 2255,
+      "clusters": 1421,
+      "largest_cluster_proteins": 15,
+      "split_rows": {
+        "validation": 1000,
+        "train": 8000,
+        "test": 1000
+      },
+      "manifest": "split-protein_similarity_50.jsonl",
+      "protein_cluster_overlap": 0
+    }
+  },
+  "sample_sha256": "b3f10e03a3882a6b57fc38de5384a639e88494a738dcca45a04196f87790186b",
+  "manifest_sha256": {
+    "split-random_pair.jsonl": "b316172efa0832e9a9677e33c91f999c32b0c2496c1d1ebcf65887a8ff919635",
+    "split-cold_protein_exact.jsonl": "5c71713f1fe2fdf7091ac57a1d5f4027b8e1c90f50af379a45074467a3f1ba53",
+    "split-protein_similarity_50.jsonl": "42edbf243966f7f9a1c6a16d86d394a94a5386ee4a96defbb9be7f364deaab78",
+    "split-cold_ligand_scaffold.jsonl": "b38a49046902eba950c91f97a5781ace3e95cc9f6252295f026cc895f38557e1"
+  }
+}

recovery/reports/dev-10k-baselines.json

@@ -0,0 +1,864 @@
+{
+  "sample_rows": 10000,
+  "target": "pAffinity",
+  "packages": {
+    "numpy": "2.5.1",
+    "scipy": "1.18.0",
+    "scikit-learn": "1.9.0"
+  },
+  "splits": {
+    "cold_ligand_scaffold": {
+      "rows": {
+        "train": 8000,
+        "validation": 1000,
+        "test": 1000
+      },
+      "mean_baseline": {
+        "prediction": 6.537638315707445,
+        "validation": {
+          "rmse": 1.5957348058680887,
+          "mae": 1.2895477378728986,
+          "pearson_r": null,
+          "spearman_r": null,
+          "r2": -0.020329719922339562
+        },
+        "test": {
+          "rmse": 1.5833990791492547,
+          "mae": 1.3026160468125345,
+          "pearson_r": null,
+          "spearman_r": null,
+          "r2": -0.00029222365027892216
+        }
+      },
+      "feature_dimensions": {
+        "protein": 4096,
+        "ligand": 4096
+      },
+      "models": [
+        {
+          "name": "protein_char3_ridge",
+          "selected_alpha": 1.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.527804068122877,
+                "mae": 1.1779963380098344,
+                "pearson_r": 0.3896740635545639,
+                "spearman_r": 0.3792641691445764,
+                "r2": 0.06469248187168464
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.4578139871351854,
+                "mae": 1.139968386530876,
+                "pearson_r": 0.4151059016140626,
+                "spearman_r": 0.3969308305062462,
+                "r2": 0.14842416968739125
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.4723908858399237,
+                "mae": 1.1793289355039596,
+                "pearson_r": 0.39965520038704133,
+                "spearman_r": 0.39853528582342707,
+                "r2": 0.1313089600423143
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.554020183734794,
+                "mae": 1.257076021552086,
+                "pearson_r": 0.3325547399758985,
+                "spearman_r": 0.3505718847248242,
+                "r2": 0.032318560087709836
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.4578139871351854,
+            "mae": 1.139968386530876,
+            "pearson_r": 0.4151059016140626,
+            "spearman_r": 0.3969308305062462,
+            "r2": 0.14842416968739125
+          },
+          "test": {
+            "rmse": 1.39215070131138,
+            "mae": 1.102342973947525,
+            "pearson_r": 0.47756240225675844,
+            "spearman_r": 0.4740036884171099,
+            "r2": 0.22675233445655052
+          },
+          "fit_and_eval_seconds": 0.9593746850005118
+        },
+        {
+          "name": "ligand_char2_5_ridge",
+          "selected_alpha": 10.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.5653168764143357,
+                "mae": 1.2302509154081345,
+                "pearson_r": 0.35459474446304934,
+                "spearman_r": 0.32893305055591643,
+                "r2": 0.018198624347302683
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.4623783005609468,
+                "mae": 1.1591383439302445,
+                "pearson_r": 0.4035976657359315,
+                "spearman_r": 0.37928879379063246,
+                "r2": 0.1430833734386795
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.4539723529408441,
+                "mae": 1.157127993941307,
+                "pearson_r": 0.4061425224793911,
+                "spearman_r": 0.39141890069825885,
+                "r2": 0.15290640459726268
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.5244228028633535,
+                "mae": 1.22633742749691,
+                "pearson_r": 0.3536892569060556,
+                "spearman_r": 0.3496830979478171,
+                "r2": 0.06882785931877777
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.4539723529408441,
+            "mae": 1.157127993941307,
+            "pearson_r": 0.4061425224793911,
+            "spearman_r": 0.39141890069825885,
+            "r2": 0.15290640459726268
+          },
+          "test": {
+            "rmse": 1.4494043209545224,
+            "mae": 1.1682950427532197,
+            "pearson_r": 0.4147866313943454,
+            "spearman_r": 0.42730966425629635,
+            "r2": 0.16184330209141884
+          },
+          "fit_and_eval_seconds": 0.33016813500034914
+        },
+        {
+          "name": "combined_char_ridge",
+          "selected_alpha": 10.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.6081591147564207,
+                "mae": 1.2413834615945816,
+                "pearson_r": 0.4169937363739051,
+                "spearman_r": 0.40429283856812304,
+                "r2": -0.03628004246158212
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.4107695549582286,
+                "mae": 1.1013368602991105,
+                "pearson_r": 0.4874112220852569,
+                "spearman_r": 0.47383417558330937,
+                "r2": 0.20249895881043067
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.383549482454833,
+                "mae": 1.0967753337621688,
+                "pearson_r": 0.49859659893827324,
+                "spearman_r": 0.4945176557261553,
+                "r2": 0.23297681168774853
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.4929408447274506,
+                "mae": 1.200959615111351,
+                "pearson_r": 0.43016081195245404,
+                "spearman_r": 0.4306665031943897,
+                "r2": 0.10689133934020467
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.383549482454833,
+            "mae": 1.0967753337621688,
+            "pearson_r": 0.49859659893827324,
+            "spearman_r": 0.4945176557261553,
+            "r2": 0.23297681168774853
+          },
+          "test": {
+            "rmse": 1.3649470834344057,
+            "mae": 1.094587422132492,
+            "pearson_r": 0.5198256551098647,
+            "spearman_r": 0.5354444265178329,
+            "r2": 0.2566767000232979
+          },
+          "fit_and_eval_seconds": 1.7034703729987086
+        }
+      ]
+    },
+    "cold_protein_exact": {
+      "rows": {
+        "train": 8000,
+        "validation": 1000,
+        "test": 1000
+      },
+      "mean_baseline": {
+        "prediction": 6.476698521107435,
+        "validation": {
+          "rmse": 1.579782624176554,
+          "mae": 1.3174895944134593,
+          "pearson_r": null,
+          "spearman_r": null,
+          "r2": -0.01274982588145357
+        },
+        "test": {
+          "rmse": 1.4951426135931578,
+          "mae": 1.2163382358753083,
+          "pearson_r": null,
+          "spearman_r": null,
+          "r2": -0.014679353070982115
+        }
+      },
+      "feature_dimensions": {
+        "protein": 4096,
+        "ligand": 4096
+      },
+      "models": [
+        {
+          "name": "protein_char3_ridge",
+          "selected_alpha": 1.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.5913619001081973,
+                "mae": 1.2860890839099883,
+                "pearson_r": 0.24087607131254302,
+                "spearman_r": 0.24239051227885716,
+                "r2": -0.02765046684019734
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.4859462663648586,
+                "mae": 1.2179750628471375,
+                "pearson_r": 0.3474792859155311,
+                "spearman_r": 0.33699637680490147,
+                "r2": 0.10398821505374234
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.5008105727933332,
+                "mae": 1.2442660155296326,
+                "pearson_r": 0.3831242209268426,
+                "spearman_r": 0.3726678732274564,
+                "r2": 0.08597247826157572
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.553986133423702,
+                "mae": 1.2952312564849853,
+                "pearson_r": 0.3282098020500765,
+                "spearman_r": 0.30711899038156937,
+                "r2": 0.02005479994650128
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.4859462663648586,
+            "mae": 1.2179750628471375,
+            "pearson_r": 0.3474792859155311,
+            "spearman_r": 0.33699637680490147,
+            "r2": 0.10398821505374234
+          },
+          "test": {
+            "rmse": 1.4585936156876136,
+            "mae": 1.1720177594423293,
+            "pearson_r": 0.2545678403741694,
+            "spearman_r": 0.2484803560755494,
+            "r2": 0.03432230415157966
+          },
+          "fit_and_eval_seconds": 0.9749435680005263
+        },
+        {
+          "name": "ligand_char2_5_ridge",
+          "selected_alpha": 10.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.6018811256320094,
+                "mae": 1.2999265880584716,
+                "pearson_r": 0.30216596804115387,
+                "spearman_r": 0.29339634813312904,
+                "r2": -0.04128132661572703
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.4768323814132662,
+                "mae": 1.2094637866020204,
+                "pearson_r": 0.37716980445129944,
+                "spearman_r": 0.36773099907324625,
+                "r2": 0.11494568428343488
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.4613319935586893,
+                "mae": 1.2048628535270691,
+                "pearson_r": 0.3911678449191863,
+                "spearman_r": 0.3873732410005544,
+                "r2": 0.1334267140116593
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.5197247823028681,
+                "mae": 1.2624011125564576,
+                "pearson_r": 0.3423555617155642,
+                "spearman_r": 0.3437308713625526,
+                "r2": 0.06278894470085006
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.4613319935586893,
+            "mae": 1.2048628535270691,
+            "pearson_r": 0.3911678449191863,
+            "spearman_r": 0.3873732410005544,
+            "r2": 0.1334267140116593
+          },
+          "test": {
+            "rmse": 1.3526643485350685,
+            "mae": 1.0842266837358474,
+            "pearson_r": 0.4362727683300783,
+            "spearman_r": 0.4191402824529675,
+            "r2": 0.16949228447261422
+          },
+          "fit_and_eval_seconds": 0.37831321699923137
+        },
+        {
+          "name": "combined_char_ridge",
+          "selected_alpha": 1.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.609541967480026,
+                "mae": 1.2925864708423616,
+                "pearson_r": 0.3470317009211976,
+                "spearman_r": 0.3383021776125508,
+                "r2": -0.05126479701152342
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.4096703358753107,
+                "mae": 1.149479679107666,
+                "pearson_r": 0.4577114457735785,
+                "spearman_r": 0.4511058044444574,
+                "r2": 0.1936146421708813
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.4111861303751425,
+                "mae": 1.165886058807373,
+                "pearson_r": 0.48119186255022445,
+                "spearman_r": 0.481658449066694,
+                "r2": 0.19187952492255234
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.500300406119371,
+                "mae": 1.2459029321670532,
+                "pearson_r": 0.41484717418396205,
+                "spearman_r": 0.4187622144178708,
+                "r2": 0.08659377868944229
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.4096703358753107,
+            "mae": 1.149479679107666,
+            "pearson_r": 0.4577114457735785,
+            "spearman_r": 0.4511058044444574,
+            "r2": 0.1936146421708813
+          },
+          "test": {
+            "rmse": 1.3462287870584813,
+            "mae": 1.0676061400175094,
+            "pearson_r": 0.4384364045433516,
+            "spearman_r": 0.42825367206750203,
+            "r2": 0.17737608661000348
+          },
+          "fit_and_eval_seconds": 1.6212121779990412
+        }
+      ]
+    },
+    "protein_similarity_50": {
+      "rows": {
+        "train": 8000,
+        "validation": 1000,
+        "test": 1000
+      },
+      "mean_baseline": {
+        "prediction": 6.4936330486238,
+        "validation": {
+          "rmse": 1.5513375597083523,
+          "mae": 1.2790881541401746,
+          "pearson_r": null,
+          "spearman_r": null,
+          "r2": -0.013098447070553743
+        },
+        "test": {
+          "rmse": 1.533614785804936,
+          "mae": 1.2540174638183117,
+          "pearson_r": null,
+          "spearman_r": null,
+          "r2": -5.474558686380249e-05
+        }
+      },
+      "feature_dimensions": {
+        "protein": 4096,
+        "ligand": 4096
+      },
+      "models": [
+        {
+          "name": "protein_char3_ridge",
+          "selected_alpha": 10.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.7552403417237736,
+                "mae": 1.401314944624901,
+                "pearson_r": 0.18824801515316072,
+                "spearman_r": 0.18074336545364095,
+                "r2": -0.29691710691928597
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.5601924885212413,
+                "mae": 1.2610043984651567,
+                "pearson_r": 0.2319965337290784,
+                "spearman_r": 0.23655822614618757,
+                "r2": -0.024696847762769814
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.5356435415785854,
+                "mae": 1.2523401461839676,
+                "pearson_r": 0.21912642622037123,
+                "spearman_r": 0.2117505257188302,
+                "r2": 0.0072957755423069814
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.5432243885780599,
+                "mae": 1.2689961529970168,
+                "pearson_r": 0.19756629969633466,
+                "spearman_r": 0.19839572987049425,
+                "r2": -0.0025295698557341595
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.5356435415785854,
+            "mae": 1.2523401461839676,
+            "pearson_r": 0.21912642622037123,
+            "spearman_r": 0.2117505257188302,
+            "r2": 0.0072957755423069814
+          },
+          "test": {
+            "rmse": 1.5253163396597564,
+            "mae": 1.2421429891586304,
+            "pearson_r": 0.12107755752116321,
+            "spearman_r": 0.11902515580955182,
+            "r2": 0.010738639684262585
+          },
+          "fit_and_eval_seconds": 0.8420703150004556
+        },
+        {
+          "name": "ligand_char2_5_ridge",
+          "selected_alpha": 10.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.5799239713380486,
+                "mae": 1.271127613902092,
+                "pearson_r": 0.29265036816534756,
+                "spearman_r": 0.2854553015351282,
+                "r2": -0.050779065362796194
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.4859846732161874,
+                "mae": 1.2077869259119034,
+                "pearson_r": 0.33129654600169606,
+                "spearman_r": 0.3199643083980807,
+                "r2": 0.07046083081253718
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.4720109593236679,
+                "mae": 1.1972612363100053,
+                "pearson_r": 0.3227538249245248,
+                "spearman_r": 0.316034036664011,
+                "r2": 0.08786079723640017
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.5116912475349742,
+                "mae": 1.23919963991642,
+                "pearson_r": 0.25043786181937333,
+                "spearman_r": 0.24831507291908664,
+                "r2": 0.03802179626202029
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.4720109593236679,
+            "mae": 1.1972612363100053,
+            "pearson_r": 0.3227538249245248,
+            "spearman_r": 0.316034036664011,
+            "r2": 0.08786079723640017
+          },
+          "test": {
+            "rmse": 1.429188581726333,
+            "mae": 1.1390189518928528,
+            "pearson_r": 0.36960662009993567,
+            "spearman_r": 0.3651127131105154,
+            "r2": 0.1314990922237853
+          },
+          "fit_and_eval_seconds": 0.33130267600063235
+        },
+        {
+          "name": "combined_char_ridge",
+          "selected_alpha": 10.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.695446275393881,
+                "mae": 1.3751192058324815,
+                "pearson_r": 0.28274147950808365,
+                "spearman_r": 0.2880691599027914,
+                "r2": -0.21006054325174417
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.4883475906405312,
+                "mae": 1.2132670460939408,
+                "pearson_r": 0.35310402171271504,
+                "spearman_r": 0.36051834925201537,
+                "r2": 0.06750229342577019
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.4627022031596522,
+                "mae": 1.1911575428247452,
+                "pearson_r": 0.36076986449064835,
+                "spearman_r": 0.3723077063029598,
+                "r2": 0.09936075775091613
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.5056113027663196,
+                "mae": 1.2321353279352187,
+                "pearson_r": 0.2885589685976343,
+                "spearman_r": 0.29454935455844794,
+                "r2": 0.04574428935903241
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.4627022031596522,
+            "mae": 1.1911575428247452,
+            "pearson_r": 0.36076986449064835,
+            "spearman_r": 0.3723077063029598,
+            "r2": 0.09936075775091613
+          },
+          "test": {
+            "rmse": 1.4348684777145742,
+            "mae": 1.1439314622879029,
+            "pearson_r": 0.36247186743462856,
+            "spearman_r": 0.3652137611435173,
+            "r2": 0.12458216430470337
+          },
+          "fit_and_eval_seconds": 1.6895546520008793
+        }
+      ]
+    },
+    "random_pair": {
+      "rows": {
+        "train": 8000,
+        "validation": 1000,
+        "test": 1000
+      },
+      "mean_baseline": {
+        "prediction": 6.5244491582587365,
+        "validation": {
+          "rmse": 1.5276238617697966,
+          "mae": 1.2404643696633129,
+          "pearson_r": null,
+          "spearman_r": null,
+          "r2": -0.00011912455955709333
+        },
+        "test": {
+          "rmse": 1.565044692955996,
+          "mae": 1.2979176334988474,
+          "pearson_r": null,
+          "spearman_r": null,
+          "r2": -0.004413634248849885
+        }
+      },
+      "feature_dimensions": {
+        "protein": 4096,
+        "ligand": 4096
+      },
+      "models": [
+        {
+          "name": "protein_char3_ridge",
+          "selected_alpha": 1.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.4410565980112502,
+                "mae": 1.1288913173079491,
+                "pearson_r": 0.4079221480849496,
+                "spearman_r": 0.3893575689504746,
+                "r2": 0.1100185770832387
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.3745768065170116,
+                "mae": 1.100080373942852,
+                "pearson_r": 0.44118951305358517,
+                "spearman_r": 0.41362525923337845,
+                "r2": 0.1902389400564115
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.3955357347301767,
+                "mae": 1.1366457569003106,
+                "pearson_r": 0.4230883826225867,
+                "spearman_r": 0.4074744496313428,
+                "r2": 0.16535693645157734
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.4814185114679912,
+                "mae": 1.2066325885653495,
+                "pearson_r": 0.3582189512625387,
+                "spearman_r": 0.36664589962084587,
+                "r2": 0.059466217343855376
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.3745768065170116,
+            "mae": 1.100080373942852,
+            "pearson_r": 0.44118951305358517,
+            "spearman_r": 0.41362525923337845,
+            "r2": 0.1902389400564115
+          },
+          "test": {
+            "rmse": 1.3880651061666627,
+            "mae": 1.1175376932621002,
+            "pearson_r": 0.4640112366064133,
+            "spearman_r": 0.4357397752396537,
+            "r2": 0.20990596198875733
+          },
+          "fit_and_eval_seconds": 0.894843939000566
+        },
+        {
+          "name": "ligand_char2_5_ridge",
+          "selected_alpha": 10.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.4856668026573707,
+                "mae": 1.1865731849074364,
+                "pearson_r": 0.3688529817391914,
+                "spearman_r": 0.37857178194968666,
+                "r2": 0.05406411042117554
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.3916760241468447,
+                "mae": 1.1180368124842643,
+                "pearson_r": 0.42160711722938904,
+                "spearman_r": 0.42128867139651427,
+                "r2": 0.16996738915693665
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.3915522045386022,
+                "mae": 1.1223492309451104,
+                "pearson_r": 0.42089896113032044,
+                "spearman_r": 0.41150350807895536,
+                "r2": 0.17011508120440122
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.462418544451856,
+                "mae": 1.1855577809214592,
+                "pearson_r": 0.35716066077001357,
+                "spearman_r": 0.34302021119123416,
+                "r2": 0.08343718024568059
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.3915522045386022,
+            "mae": 1.1223492309451104,
+            "pearson_r": 0.42089896113032044,
+            "spearman_r": 0.41150350807895536,
+            "r2": 0.17011508120440122
+          },
+          "test": {
+            "rmse": 1.4268869133537447,
+            "mae": 1.175576719045639,
+            "pearson_r": 0.4211078957570381,
+            "spearman_r": 0.40335385177971433,
+            "r2": 0.1650927724032577
+          },
+          "fit_and_eval_seconds": 0.4427177519992256
+        },
+        {
+          "name": "combined_char_ridge",
+          "selected_alpha": 10.0,
+          "validation_trials": [
+            {
+              "alpha": 0.1,
+              "metrics": {
+                "rmse": 1.4956017228690686,
+                "mae": 1.1701721274256707,
+                "pearson_r": 0.4361198131829029,
+                "spearman_r": 0.4281298523300441,
+                "r2": 0.0413705240127229
+              }
+            },
+            {
+              "alpha": 1.0,
+              "metrics": {
+                "rmse": 1.3140696311066697,
+                "mae": 1.0519584881663322,
+                "pearson_r": 0.5201378315237909,
+                "spearman_r": 0.5029682145871107,
+                "r2": 0.25995927103340644
+              }
+            },
+            {
+              "alpha": 10.0,
+              "metrics": {
+                "rmse": 1.307333278135575,
+                "mae": 1.0559683529734611,
+                "pearson_r": 0.5300812334907434,
+                "spearman_r": 0.5194865614899057,
+                "r2": 0.26752720749752146
+              }
+            },
+            {
+              "alpha": 100.0,
+              "metrics": {
+                "rmse": 1.4257502567100055,
+                "mae": 1.15831251257658,
+                "pearson_r": 0.45176846530905573,
+                "spearman_r": 0.44875210828870205,
+                "r2": 0.1288242419468516
+              }
+            }
+          ],
+          "validation": {
+            "rmse": 1.307333278135575,
+            "mae": 1.0559683529734611,
+            "pearson_r": 0.5300812334907434,
+            "spearman_r": 0.5194865614899057,
+            "r2": 0.26752720749752146
+          },
+          "test": {
+            "rmse": 1.335818862102712,
+            "mae": 1.0964913694858551,
+            "pearson_r": 0.5371481422264178,
+            "spearman_r": 0.5186748222826298,
+            "r2": 0.2682642831079074
+          },
+          "fit_and_eval_seconds": 1.8038349809994543
+        }
+      ]
+    }
+  },
+  "total_seconds": 26.8222462800004
+}

recovery/reports/overfit-actual-embeddings.json

@@ -0,0 +1,60 @@
+{
+  "examples": 32,
+  "steps": 1000,
+  "seed": 42,
+  "rmse_paffinity": 2.6320698458933376e-07,
+  "mae_paffinity": 9.685754776000977e-08,
+  "acceptance_max_rmse_paffinity": 0.05,
+  "passed": true,
+  "target_scaler": {
+    "mean": 6.446505546569824,
+    "std": 1.7902863025665283
+  },
+  "history": [
+    {
+      "step": 1,
+      "standardized_mse": 1.0336081981658936
+    },
+    {
+      "step": 100,
+      "standardized_mse": 2.772125753836008e-06
+    },
+    {
+      "step": 200,
+      "standardized_mse": 1.03808101004077e-10
+    },
+    {
+      "step": 300,
+      "standardized_mse": 1.9242752194188056e-14
+    },
+    {
+      "step": 400,
+      "standardized_mse": 1.2282174246831118e-14
+    },
+    {
+      "step": 500,
+      "standardized_mse": 1.1517161193925346e-14
+    },
+    {
+      "step": 600,
+      "standardized_mse": 5.5680982768535614e-05
+    },
+    {
+      "step": 700,
+      "standardized_mse": 3.629010425498791e-09
+    },
+    {
+      "step": 800,
+      "standardized_mse": 6.048471863336416e-14
+    },
+    {
+      "step": 900,
+      "standardized_mse": 1.406742154404575e-14
+    },
+    {
+      "step": 1000,
+      "standardized_mse": 1.5145552184365335e-14
+    }
+  ],
+  "scope": "trainable head over 32 embeddings from the pinned frozen encoders recorded in recovery/config.json"
+}

recovery/scripts/cluster_proteins.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+"""Create a sequence-similarity-disjoint protein split using MMseqs2."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+import subprocess
+from collections import Counter
+from pathlib import Path
+
+from mitointeract_recovery.splits import assign_grouped_splits, assert_group_disjoint
+
+
+def read_rows(path: Path) -> list[dict]:
+    with path.open() as handle:
+        return [json.loads(line) for line in handle if line.strip()]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-dir", type=Path, default=Path("artifacts/dev-10k"))
+    parser.add_argument("--min-seq-id", type=float, default=0.5)
+    parser.add_argument("--coverage", type=float, default=0.8)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--threads", type=int, default=8)
+    args = parser.parse_args()
+
+    if not shutil.which("mmseqs"):
+        raise RuntimeError("MMseqs2 is required; install the mmseqs2 package")
+    rows = read_rows(args.data_dir / "sample.jsonl")
+    proteins = sorted({row["protein_id"]: row["sequence"] for row in rows}.items())
+    suffix = (
+        f"identity-{int(args.min_seq_id * 100)}-coverage-{int(args.coverage * 100)}"
+    )
+    work = args.data_dir / f"mmseqs-{suffix}"
+    work.mkdir(parents=True, exist_ok=True)
+    fasta = work / "proteins.fasta"
+    with fasta.open("w") as handle:
+        for protein_id, sequence in proteins:
+            handle.write(f">{protein_id}\n{sequence}\n")
+
+    prefix = work / "clusters"
+    tmp = work / "tmp"
+    command = [
+        "mmseqs",
+        "easy-cluster",
+        str(fasta),
+        str(prefix),
+        str(tmp),
+        "--min-seq-id",
+        str(args.min_seq_id),
+        "-c",
+        str(args.coverage),
+        "--cov-mode",
+        "0",
+        "--alignment-mode",
+        "3",
+        "--threads",
+        str(args.threads),
+    ]
+    cluster_tsv = prefix.with_name(prefix.name + "_cluster.tsv")
+    if not cluster_tsv.exists():
+        subprocess.run(command, check=True)
+    protein_to_cluster: dict[str, str] = {}
+    with cluster_tsv.open() as handle:
+        for line in handle:
+            representative, member = line.rstrip("\n").split("\t")
+            protein_to_cluster[member] = representative
+    missing = {protein_id for protein_id, _ in proteins} - protein_to_cluster.keys()
+    if missing:
+        raise AssertionError(f"MMseqs2 omitted {len(missing)} proteins")
+
+    enriched = [
+        dict(row, protein_similarity_cluster=protein_to_cluster[row["protein_id"]])
+        for row in rows
+    ]
+    assignments = assign_grouped_splits(
+        enriched,
+        group_key="protein_similarity_cluster",
+        seed=args.seed,
+    )
+    assert_group_disjoint(enriched, assignments, group_key="protein_similarity_cluster")
+    manifest_name = f"split-protein_similarity_{int(args.min_seq_id * 100)}.jsonl"
+    manifest = args.data_dir / manifest_name
+    with manifest.open("w") as handle:
+        for row in rows:
+            handle.write(
+                json.dumps(
+                    {"pair_id": row["pair_id"], "split": assignments[row["pair_id"]]}
+                )
+                + "\n"
+            )
+
+    cluster_sizes = Counter(protein_to_cluster.values())
+    row_counts = Counter(assignments.values())
+    report = {
+        "tool": "MMseqs2",
+        "command": command,
+        "min_sequence_identity": args.min_seq_id,
+        "minimum_alignment_coverage": args.coverage,
+        "coverage_mode": "both query and target",
+        "alignment_mode": "actual identical residues / aligned columns",
+        "unique_proteins": len(proteins),
+        "clusters": len(cluster_sizes),
+        "largest_cluster_proteins": max(cluster_sizes.values()),
+        "split_rows": dict(row_counts),
+        "manifest": manifest_name,
+        "protein_cluster_overlap": 0,
+    }
+    (args.data_dir / f"protein-clusters-{suffix}.json").write_text(
+        json.dumps(report, indent=2) + "\n"
+    )
+    audit_path = args.data_dir / "audit.json"
+    audit = json.loads(audit_path.read_text())
+    audit.setdefault("split_manifests", {})[
+        f"protein_similarity_{int(args.min_seq_id * 100)}"
+    ] = manifest_name
+    audit.setdefault("external_split_audits", {})[
+        f"protein_similarity_{int(args.min_seq_id * 100)}"
+    ] = report
+    audit_path.write_text(json.dumps(audit, indent=2) + "\n")
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

recovery/scripts/dataset_smoke.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 """Validate dataset target semantics and compute a deterministic mean baseline."""
+
 from __future__ import annotations
 
 import argparse
@@ -11,14 +12,20 @@ import urllib.parse
 import urllib.request
 from pathlib import Path
 
-from mitointeract_recovery import micromolar_to_pkd
+from mitointeract_recovery import micromolar_to_paffinity
 
 
 def fetch_rows(dataset: str, rows: int) -> list[dict]:
     if not 2 <= rows <= 100:
         raise ValueError("Dataset Viewer smoke checks support 2-100 rows")
     query = urllib.parse.urlencode(
-        {"dataset": dataset, "config": "default", "split": "train", "offset": 0, "length": rows}
+        {
+            "dataset": dataset,
+            "config": "default",
+            "split": "train",
+            "offset": 0,
+            "length": rows,
+        }
     )
     url = f"https://datasets-server.huggingface.co/rows?{query}"
     with urllib.request.urlopen(url, timeout=60) as response:
@@ -31,7 +38,9 @@ def pearson(xs: list[float], ys: list[float]) -> float | None:
         return None
     xm, ym = statistics.mean(xs), statistics.mean(ys)
     numerator = sum((x - xm) * (y - ym) for x, y in zip(xs, ys))
-    denominator = math.sqrt(sum((x - xm) ** 2 for x in xs) * sum((y - ym) ** 2 for y in ys))
+    denominator = math.sqrt(
+        sum((x - xm) ** 2 for x in xs) * sum((y - ym) ** 2 for y in ys)
+    )
     return numerator / denominator if denominator else None
 
 
@@ -40,17 +49,26 @@ def main() -> None:
     parser.add_argument("--dataset", default="jglaser/binding_affinity")
     parser.add_argument("--rows", type=int, default=100)
     parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--output", type=Path, default=Path("artifacts/dataset-smoke.json"))
+    parser.add_argument(
+        "--output", type=Path, default=Path("artifacts/dataset-smoke.json")
+    )
     args = parser.parse_args()
 
     rows = fetch_rows(args.dataset, args.rows)
     errors = []
     for index, row in enumerate(rows):
-        calculated = micromolar_to_pkd(float(row["affinity_uM"]))
+        calculated = micromolar_to_paffinity(float(row["affinity_uM"]))
         published = float(row["neg_log10_affinity_M"])
         delta = abs(calculated - published)
         if delta > 1e-5:
-            errors.append({"row": index, "calculated": calculated, "published": published, "delta": delta})
+            errors.append(
+                {
+                    "row": index,
+                    "calculated": calculated,
+                    "published": published,
+                    "delta": delta,
+                }
+            )
 
     shuffled = list(rows)
     random.Random(args.seed).shuffle(shuffled)
@@ -60,16 +78,27 @@ def main() -> None:
     test_targets = [float(row["neg_log10_affinity_M"]) for row in test]
     mean_prediction = statistics.mean(train_targets)
     predictions = [mean_prediction] * len(test_targets)
-    rmse = math.sqrt(statistics.mean((pred - target) ** 2 for pred, target in zip(predictions, test_targets)))
-    mae = statistics.mean(abs(pred - target) for pred, target in zip(predictions, test_targets))
+    rmse = math.sqrt(
+        statistics.mean(
+            (pred - target) ** 2 for pred, target in zip(predictions, test_targets)
+        )
+    )
+    mae = statistics.mean(
+        abs(pred - target) for pred, target in zip(predictions, test_targets)
+    )
 
     report = {
         "dataset": args.dataset,
         "rows": len(rows),
         "seed": args.seed,
         "target_column": "neg_log10_affinity_M",
+        "target_unit": "pAffinity (-log10 of the dataset's mixed affinity value in M)",
         "unit_check": {"passed": not errors, "errors": errors},
-        "split": {"train_rows": len(train), "test_rows": len(test), "kind": "seeded smoke only"},
+        "split": {
+            "train_rows": len(train),
+            "test_rows": len(test),
+            "kind": "seeded smoke only",
+        },
         "train_mean_baseline": {
             "prediction": mean_prediction,
             "rmse": rmse,

recovery/scripts/embed_sample.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+"""Cache pinned frozen-encoder embeddings for a MitoInteract sample."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+DEFAULT_PROTEIN_MODEL = "facebook/esm2_t12_35M_UR50D"
+DEFAULT_PROTEIN_REVISION = "6fbf070e65b0b7291e7bbcd451118c216cff79d8"
+DEFAULT_LIGAND_MODEL = "DeepChem/ChemBERTa-77M-MLM"
+DEFAULT_LIGAND_REVISION = "ed8a5374f2024ec8da53760af91a33fb8f6a15ff"
+
+
+def read_rows(path: Path, limit: int | None) -> list[dict]:
+    rows = []
+    with path.open() as handle:
+        for line in handle:
+            if line.strip():
+                rows.append(json.loads(line))
+            if limit and len(rows) >= limit:
+                break
+    return rows
+
+
+def masked_mean(
+    last_hidden: torch.Tensor, attention: torch.Tensor, special: torch.Tensor
+) -> torch.Tensor:
+    mask = attention.bool() & ~special.bool()
+    weights = mask.unsqueeze(-1).to(last_hidden.dtype)
+    return (last_hidden * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1)
+
+
+def encode_texts(
+    model, tokenizer, texts: list[str], batch_size: int, device: torch.device
+) -> np.ndarray:
+    outputs = []
+    for start in range(0, len(texts), batch_size):
+        batch = texts[start : start + batch_size]
+        encoded = tokenizer(
+            batch,
+            padding=True,
+            truncation=True,
+            max_length=min(getattr(tokenizer, "model_max_length", 512), 512),
+            return_special_tokens_mask=True,
+            return_tensors="pt",
+        )
+        special = encoded.pop("special_tokens_mask")
+        encoded = {key: value.to(device) for key, value in encoded.items()}
+        with torch.inference_mode():
+            hidden = model(**encoded).last_hidden_state
+        pooled = masked_mean(hidden, encoded["attention_mask"], special.to(device))
+        outputs.append(pooled.float().cpu().numpy())
+    return np.concatenate(outputs, axis=0)
+
+
+def encode_proteins(
+    model,
+    tokenizer,
+    entities: list[tuple[str, str]],
+    batch_size: int,
+    device: torch.device,
+    chunk_residues: int,
+) -> dict[str, np.ndarray]:
+    chunks: list[str] = []
+    owners: list[str] = []
+    weights: list[int] = []
+    for entity_id, sequence in entities:
+        for start in range(0, len(sequence), chunk_residues):
+            chunk = sequence[start : start + chunk_residues]
+            chunks.append(chunk)
+            owners.append(entity_id)
+            weights.append(len(chunk))
+
+    chunk_embeddings = []
+    for start in range(0, len(chunks), batch_size):
+        batch = chunks[start : start + batch_size]
+        encoded = tokenizer(
+            batch,
+            padding=True,
+            truncation=True,
+            max_length=chunk_residues + 2,
+            return_special_tokens_mask=True,
+            return_tensors="pt",
+        )
+        special = encoded.pop("special_tokens_mask")
+        encoded = {key: value.to(device) for key, value in encoded.items()}
+        with torch.inference_mode():
+            hidden = model(**encoded).last_hidden_state
+        pooled = masked_mean(hidden, encoded["attention_mask"], special.to(device))
+        chunk_embeddings.extend(pooled.float().cpu().numpy())
+
+    accum: dict[str, list[tuple[np.ndarray, int]]] = defaultdict(list)
+    for owner, embedding, weight in zip(owners, chunk_embeddings, weights):
+        accum[owner].append((embedding, weight))
+    result = {}
+    for owner, values in accum.items():
+        matrix = np.stack([value for value, _ in values])
+        entity_weights = np.asarray([weight for _, weight in values], dtype=np.float32)
+        result[owner] = np.average(matrix, axis=0, weights=entity_weights)
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--sample", type=Path, default=Path("artifacts/dev-10k/sample.jsonl")
+    )
+    parser.add_argument("--limit", type=int)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--protein-chunk-residues", type=int, default=1022)
+    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
+    parser.add_argument("--protein-model", default=DEFAULT_PROTEIN_MODEL)
+    parser.add_argument("--protein-revision", default=DEFAULT_PROTEIN_REVISION)
+    parser.add_argument("--ligand-model", default=DEFAULT_LIGAND_MODEL)
+    parser.add_argument("--ligand-revision", default=DEFAULT_LIGAND_REVISION)
+    parser.add_argument("--output", type=Path, default=Path("artifacts/embeddings.npz"))
+    args = parser.parse_args()
+
+    device_name = (
+        "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device
+    )
+    if device_name == "auto":
+        device_name = "cpu"
+    device = torch.device(device_name)
+    rows = read_rows(args.sample, args.limit)
+    proteins = sorted({row["protein_id"]: row["sequence"] for row in rows}.items())
+    ligands = sorted({row["ligand_id"]: row["smiles"] for row in rows}.items())
+
+    started = time.monotonic()
+    protein_tokenizer = AutoTokenizer.from_pretrained(
+        args.protein_model, revision=args.protein_revision
+    )
+    protein_model = (
+        AutoModel.from_pretrained(args.protein_model, revision=args.protein_revision)
+        .eval()
+        .to(device)
+    )
+    protein_embeddings = encode_proteins(
+        protein_model,
+        protein_tokenizer,
+        proteins,
+        args.batch_size,
+        device,
+        args.protein_chunk_residues,
+    )
+    del protein_model
+
+    ligand_tokenizer = AutoTokenizer.from_pretrained(
+        args.ligand_model, revision=args.ligand_revision
+    )
+    ligand_model = (
+        AutoModel.from_pretrained(args.ligand_model, revision=args.ligand_revision)
+        .eval()
+        .to(device)
+    )
+    ligand_matrix = encode_texts(
+        ligand_model,
+        ligand_tokenizer,
+        [smiles for _, smiles in ligands],
+        args.batch_size,
+        device,
+    )
+    ligand_embeddings = {
+        entity_id: embedding
+        for (entity_id, _), embedding in zip(ligands, ligand_matrix)
+    }
+
+    pair_protein = np.stack(
+        [protein_embeddings[row["protein_id"]] for row in rows]
+    ).astype(np.float32)
+    pair_ligand = np.stack(
+        [ligand_embeddings[row["ligand_id"]] for row in rows]
+    ).astype(np.float32)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(
+        args.output,
+        pair_ids=np.asarray([row["pair_id"] for row in rows]),
+        paffinity=np.asarray([row["paffinity"] for row in rows], dtype=np.float32),
+        protein=pair_protein,
+        ligand=pair_ligand,
+    )
+    metadata = {
+        "rows": len(rows),
+        "unique_proteins": len(proteins),
+        "unique_ligands": len(ligands),
+        "protein_dim": int(pair_protein.shape[1]),
+        "ligand_dim": int(pair_ligand.shape[1]),
+        "protein_model": args.protein_model,
+        "protein_revision": args.protein_revision,
+        "ligand_model": args.ligand_model,
+        "ligand_revision": args.ligand_revision,
+        "protein_chunk_residues": args.protein_chunk_residues,
+        "device": str(device),
+        "seconds": time.monotonic() - started,
+    }
+    args.output.with_suffix(".json").write_text(json.dumps(metadata, indent=2) + "\n")
+    print(json.dumps(metadata, indent=2))
+
+
+if __name__ == "__main__":
+    main()

recovery/scripts/overfit_sanity.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+"""Prove the v2 trainable head can memorize a fixed 32-example set."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from mitointeract_recovery.model import MitoInteractHead, TargetScaler
+
+
+def read_targets(path: Path, count: int) -> torch.Tensor:
+    values = []
+    with path.open() as handle:
+        for line in handle:
+            if line.strip():
+                values.append(float(json.loads(line)["paffinity"]))
+            if len(values) == count:
+                break
+    if len(values) != count:
+        raise ValueError(f"needed {count} examples, found {len(values)}")
+    return torch.tensor(values, dtype=torch.float32)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--sample", type=Path, default=Path("artifacts/dev-10k/sample.jsonl")
+    )
+    parser.add_argument(
+        "--embeddings", type=Path, help="Optional NPZ from embed_sample.py"
+    )
+    parser.add_argument("--examples", type=int, default=32)
+    parser.add_argument("--steps", type=int, default=2000)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--max-rmse", type=float, default=0.05)
+    parser.add_argument(
+        "--output", type=Path, default=Path("artifacts/overfit-sanity.json")
+    )
+    args = parser.parse_args()
+
+    torch.manual_seed(args.seed)
+    torch.set_num_threads(min(torch.get_num_threads(), 8))
+    embedding_scope = "synthetic fixed embeddings"
+    if args.embeddings:
+        arrays = np.load(args.embeddings)
+        available = len(arrays["paffinity"])
+        if available < args.examples:
+            raise ValueError(f"needed {args.examples} embeddings, found {available}")
+        targets = torch.from_numpy(arrays["paffinity"][: args.examples]).float()
+        protein = torch.from_numpy(arrays["protein"][: args.examples]).float()
+        ligand = torch.from_numpy(arrays["ligand"][: args.examples]).float()
+        embedding_scope = f"actual pinned frozen encoders from {args.embeddings}"
+    else:
+        targets = read_targets(args.sample, args.examples)
+        # Fixed synthetic embeddings isolate the trainable fusion/head path from
+        # encoder downloads and prove optimization can fit a tiny set exactly.
+        protein = torch.randn(args.examples, 32)
+        ligand = torch.randn(args.examples, 24)
+
+    scaler = TargetScaler.fit(targets)
+    encoded_targets = scaler.encode(targets)
+    model = MitoInteractHead(
+        protein.shape[1],
+        ligand.shape[1],
+        projection_dim=64,
+        hidden_dim=128,
+        dropout=0.0,
+    )
+    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.0)
+
+    history = []
+    for step in range(1, args.steps + 1):
+        model.train()
+        optimizer.zero_grad(set_to_none=True)
+        prediction = model(protein, ligand)
+        loss = torch.nn.functional.mse_loss(prediction, encoded_targets)
+        loss.backward()
+        optimizer.step()
+        if step == 1 or step % 100 == 0 or step == args.steps:
+            history.append({"step": step, "standardized_mse": float(loss.detach())})
+
+    model.eval()
+    with torch.no_grad():
+        decoded = scaler.decode(model(protein, ligand))
+        rmse = math.sqrt(float(torch.mean((decoded - targets) ** 2)))
+        mae = float(torch.mean(torch.abs(decoded - targets)))
+    report = {
+        "examples": args.examples,
+        "steps": args.steps,
+        "seed": args.seed,
+        "rmse_paffinity": rmse,
+        "mae_paffinity": mae,
+        "acceptance_max_rmse_paffinity": args.max_rmse,
+        "passed": rmse <= args.max_rmse,
+        "target_scaler": {"mean": scaler.mean, "std": scaler.std},
+        "history": history,
+        "scope": f"trainable head over {embedding_scope}",
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(report, indent=2) + "\n")
+    print(json.dumps(report, indent=2))
+    if not report["passed"]:
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

recovery/scripts/prepare_sample.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+"""Create a deterministic, deduplicated DTA development sample and split manifests."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+import random
+import statistics
+from collections import Counter, defaultdict
+from pathlib import Path
+
+import pyarrow.parquet as pq
+
+from mitointeract_recovery import micromolar_to_paffinity
+from mitointeract_recovery.chemistry import canonicalize_smiles, scaffold_id, stable_id
+from mitointeract_recovery.splits import (
+    assign_grouped_splits,
+    assign_random_splits,
+    assert_group_disjoint,
+)
+
+COLUMNS = ["seq", "smiles", "smiles_can", "affinity_uM", "neg_log10_affinity_M"]
+
+
+def stable_pair_id(sequence: str, smiles: str) -> str:
+    return stable_id("pair", f"{sequence}\0{smiles}")
+
+
+def sample_raw_rows(
+    parquet_path: Path, candidate_count: int, seed: int
+) -> tuple[list[dict], int]:
+    parquet = pq.ParquetFile(parquet_path)
+    total_rows = parquet.metadata.num_rows
+    if candidate_count > total_rows:
+        candidate_count = total_rows
+    wanted = sorted(random.Random(seed).sample(range(total_rows), candidate_count))
+    selected: list[dict] = []
+    pointer = 0
+    offset = 0
+
+    for batch in parquet.iter_batches(batch_size=8192, columns=COLUMNS):
+        batch_end = offset + batch.num_rows
+        local_indices: list[int] = []
+        while pointer < len(wanted) and wanted[pointer] < batch_end:
+            local_indices.append(wanted[pointer] - offset)
+            pointer += 1
+        if local_indices:
+            data = batch.to_pydict()
+            for local in local_indices:
+                selected.append({column: data[column][local] for column in COLUMNS})
+        offset = batch_end
+        if pointer == len(wanted):
+            break
+
+    if len(selected) != candidate_count:
+        raise RuntimeError(
+            f"requested {candidate_count} rows but selected {len(selected)}"
+        )
+    return selected, total_rows
+
+
+def normalize_and_deduplicate(
+    raw_rows: list[dict], sample_size: int, seed: int
+) -> tuple[list[dict], Counter]:
+    rejected = Counter()
+    pairs: dict[str, dict] = {}
+    targets: dict[str, list[float]] = defaultdict(list)
+
+    for raw in raw_rows:
+        sequence = "".join(str(raw.get("seq") or "").split()).upper()
+        source_smiles = str(raw.get("smiles_can") or raw.get("smiles") or "").strip()
+        if not sequence:
+            rejected["empty_sequence"] += 1
+            continue
+        try:
+            canonical = canonicalize_smiles(source_smiles)
+        except ValueError:
+            rejected["invalid_smiles"] += 1
+            continue
+        try:
+            affinity_um = float(raw["affinity_uM"])
+            published_paffinity = float(raw["neg_log10_affinity_M"])
+            calculated_paffinity = micromolar_to_paffinity(affinity_um)
+        except (TypeError, ValueError, OverflowError):
+            rejected["invalid_affinity"] += 1
+            continue
+        if (
+            not math.isfinite(published_paffinity)
+            or abs(calculated_paffinity - published_paffinity) > 1e-4
+        ):
+            rejected["unit_mismatch"] += 1
+            continue
+
+        pair_id = stable_pair_id(sequence, canonical)
+        if pair_id not in pairs:
+            pairs[pair_id] = {
+                "pair_id": pair_id,
+                "protein_id": stable_id("protein", sequence),
+                "ligand_id": stable_id("ligand", canonical),
+                "scaffold_id": scaffold_id(canonical),
+                "sequence": sequence,
+                "smiles": canonical,
+                "protein_length": len(sequence),
+                "smiles_length": len(canonical),
+            }
+        targets[pair_id].append(published_paffinity)
+
+    normalized: list[dict] = []
+    for pair_id, row in pairs.items():
+        values = targets[pair_id]
+        normalized.append(
+            {
+                **row,
+                "paffinity": statistics.median(values),
+                "replicate_count": len(values),
+                "replicate_paffinity_range": max(values) - min(values),
+            }
+        )
+
+    normalized.sort(
+        key=lambda row: hashlib.sha256(f"{seed}:{row['pair_id']}".encode()).hexdigest()
+    )
+    if len(normalized) < sample_size:
+        raise RuntimeError(
+            f"only {len(normalized)} valid unique pairs remained; increase --candidate-multiplier"
+        )
+    return normalized[:sample_size], rejected
+
+
+def write_jsonl(path: Path, rows: list[dict]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w") as handle:
+        for row in rows:
+            handle.write(json.dumps(row, sort_keys=True) + "\n")
+
+
+def target_stats(rows: list[dict], assignments: dict[str, str]) -> dict:
+    result = {}
+    for split in ("train", "validation", "test"):
+        values = [
+            row["paffinity"] for row in rows if assignments[row["pair_id"]] == split
+        ]
+        result[split] = {
+            "rows": len(values),
+            "mean_paffinity": statistics.mean(values),
+            "std_paffinity": statistics.stdev(values) if len(values) > 1 else 0.0,
+            "min_paffinity": min(values),
+            "max_paffinity": max(values),
+        }
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--parquet", type=Path, required=True)
+    parser.add_argument(
+        "--dataset-revision",
+        default="11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b",
+    )
+    parser.add_argument("--sample-size", type=int, default=10_000)
+    parser.add_argument("--candidate-multiplier", type=int, default=3)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/dev-10k"))
+    args = parser.parse_args()
+    if args.sample_size < 100:
+        raise ValueError("sample size must be at least 100")
+
+    raw_rows, source_rows = sample_raw_rows(
+        args.parquet,
+        args.sample_size * args.candidate_multiplier,
+        args.seed,
+    )
+    rows, rejected = normalize_and_deduplicate(raw_rows, args.sample_size, args.seed)
+
+    manifests = {
+        "random_pair": assign_random_splits(rows, seed=args.seed),
+        "cold_protein_exact": assign_grouped_splits(rows, "protein_id", seed=args.seed),
+        "cold_ligand_scaffold": assign_grouped_splits(
+            rows, "scaffold_id", seed=args.seed
+        ),
+    }
+    assert_group_disjoint(rows, manifests["cold_protein_exact"], "protein_id")
+    assert_group_disjoint(rows, manifests["cold_ligand_scaffold"], "scaffold_id")
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    write_jsonl(args.output_dir / "sample.jsonl", rows)
+    for name, assignments in manifests.items():
+        write_jsonl(
+            args.output_dir / f"split-{name}.jsonl",
+            [
+                {"pair_id": pair_id, "split": split}
+                for pair_id, split in sorted(assignments.items())
+            ],
+        )
+
+    with args.parquet.open("rb") as handle:
+        source_sha256 = hashlib.file_digest(handle, "sha256").hexdigest()
+    audit = {
+        "source_parquet": str(args.parquet),
+        "source_sha256": source_sha256,
+        "dataset": "jglaser/binding_affinity",
+        "dataset_revision": args.dataset_revision,
+        "source_rows": source_rows,
+        "candidate_rows": len(raw_rows),
+        "sample_rows": len(rows),
+        "seed": args.seed,
+        "target": {"column": "neg_log10_affinity_M", "unit": "pAffinity"},
+        "rejected": dict(rejected),
+        "unique_proteins": len({row["protein_id"] for row in rows}),
+        "unique_ligands": len({row["ligand_id"] for row in rows}),
+        "unique_scaffold_groups": len({row["scaffold_id"] for row in rows}),
+        "duplicate_measurements_in_sample": sum(
+            row["replicate_count"] - 1 for row in rows
+        ),
+        "protein_length": {
+            "min": min(row["protein_length"] for row in rows),
+            "median": statistics.median(row["protein_length"] for row in rows),
+            "max": max(row["protein_length"] for row in rows),
+            "over_512": sum(row["protein_length"] > 512 for row in rows),
+        },
+        "splits": {
+            name: target_stats(rows, assignments)
+            for name, assignments in manifests.items()
+        },
+        "limitations": [
+            "the source target can combine Ki, Kd, IC50, and EC50; pAffinity is not pKd",
+            "assay type, relation operator, and source provenance are absent from the combined Parquet",
+            "cold_protein_exact is entity-disjoint but not sequence-similarity-clustered",
+            "scaffold groups use Bemis-Murcko with exact-molecule fallback for acyclic ligands",
+            "this development sample is not an external biological benchmark",
+        ],
+    }
+    (args.output_dir / "audit.json").write_text(json.dumps(audit, indent=2) + "\n")
+    print(json.dumps(audit, indent=2))
+
+
+if __name__ == "__main__":
+    main()

recovery/scripts/run_baselines.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+"""Benchmark leakage-auditing character n-gram ridge baselines."""
+
+from __future__ import annotations
+
+import argparse
+import importlib.metadata
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+from scipy.sparse import hstack
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import Ridge
+
+from mitointeract_recovery.metrics import regression_metrics
+
+ALPHAS = (0.1, 1.0, 10.0, 100.0)
+
+
+def read_jsonl(path: Path) -> list[dict]:
+    with path.open() as handle:
+        return [json.loads(line) for line in handle if line.strip()]
+
+
+def read_manifest(path: Path) -> dict[str, str]:
+    return {row["pair_id"]: row["split"] for row in read_jsonl(path)}
+
+
+def partition(rows: list[dict], manifest: dict[str, str]) -> dict[str, list[dict]]:
+    result = {"train": [], "validation": [], "test": []}
+    for row in rows:
+        result[manifest[row["pair_id"]]].append(row)
+    return result
+
+
+def targets(rows: list[dict]) -> np.ndarray:
+    return np.asarray([row["paffinity"] for row in rows], dtype=np.float64)
+
+
+def select_ridge(
+    train_x,
+    train_y: np.ndarray,
+    validation_x,
+    validation_y: np.ndarray,
+) -> tuple[Ridge, float, list[dict]]:
+    trials = []
+    best = None
+    for alpha in ALPHAS:
+        model = Ridge(alpha=alpha, solver="lsqr", tol=1e-4)
+        model.fit(train_x, train_y)
+        predictions = model.predict(validation_x)
+        metrics = regression_metrics(validation_y, predictions)
+        trials.append({"alpha": alpha, "metrics": metrics})
+        if best is None or metrics["rmse"] < best[0]:
+            best = (metrics["rmse"], model, alpha)
+    return best[1], best[2], trials
+
+
+def evaluate_feature_set(
+    name: str,
+    train_x,
+    validation_x,
+    test_x,
+    train_y: np.ndarray,
+    validation_y: np.ndarray,
+    test_y: np.ndarray,
+) -> dict:
+    started = time.monotonic()
+    model, alpha, trials = select_ridge(train_x, train_y, validation_x, validation_y)
+    return {
+        "name": name,
+        "selected_alpha": alpha,
+        "validation_trials": trials,
+        "validation": regression_metrics(validation_y, model.predict(validation_x)),
+        "test": regression_metrics(test_y, model.predict(test_x)),
+        "fit_and_eval_seconds": time.monotonic() - started,
+    }
+
+
+def benchmark_split(rows: list[dict], manifest_path: Path) -> dict:
+    manifest = read_manifest(manifest_path)
+    splits = partition(rows, manifest)
+    train_y = targets(splits["train"])
+    validation_y = targets(splits["validation"])
+    test_y = targets(splits["test"])
+
+    mean = float(train_y.mean())
+    result = {
+        "rows": {name: len(values) for name, values in splits.items()},
+        "mean_baseline": {
+            "prediction": mean,
+            "validation": regression_metrics(
+                validation_y, np.full_like(validation_y, mean)
+            ),
+            "test": regression_metrics(test_y, np.full_like(test_y, mean)),
+        },
+    }
+
+    protein_vectorizer = TfidfVectorizer(
+        analyzer="char",
+        ngram_range=(3, 3),
+        lowercase=False,
+        min_df=2,
+        max_features=4096,
+        sublinear_tf=True,
+        dtype=np.float32,
+    )
+    ligand_vectorizer = TfidfVectorizer(
+        analyzer="char",
+        ngram_range=(2, 5),
+        lowercase=False,
+        min_df=2,
+        max_features=4096,
+        sublinear_tf=True,
+        dtype=np.float32,
+    )
+
+    protein_train = protein_vectorizer.fit_transform(
+        [row["sequence"] for row in splits["train"]]
+    )
+    protein_validation = protein_vectorizer.transform(
+        [row["sequence"] for row in splits["validation"]]
+    )
+    protein_test = protein_vectorizer.transform(
+        [row["sequence"] for row in splits["test"]]
+    )
+    ligand_train = ligand_vectorizer.fit_transform(
+        [row["smiles"] for row in splits["train"]]
+    )
+    ligand_validation = ligand_vectorizer.transform(
+        [row["smiles"] for row in splits["validation"]]
+    )
+    ligand_test = ligand_vectorizer.transform([row["smiles"] for row in splits["test"]])
+
+    result["feature_dimensions"] = {
+        "protein": protein_train.shape[1],
+        "ligand": ligand_train.shape[1],
+    }
+    result["models"] = [
+        evaluate_feature_set(
+            "protein_char3_ridge",
+            protein_train,
+            protein_validation,
+            protein_test,
+            train_y,
+            validation_y,
+            test_y,
+        ),
+        evaluate_feature_set(
+            "ligand_char2_5_ridge",
+            ligand_train,
+            ligand_validation,
+            ligand_test,
+            train_y,
+            validation_y,
+            test_y,
+        ),
+        evaluate_feature_set(
+            "combined_char_ridge",
+            hstack([protein_train, ligand_train], format="csr"),
+            hstack([protein_validation, ligand_validation], format="csr"),
+            hstack([protein_test, ligand_test], format="csr"),
+            train_y,
+            validation_y,
+            test_y,
+        ),
+    ]
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-dir", type=Path, default=Path("artifacts/dev-10k"))
+    parser.add_argument(
+        "--output", type=Path, default=Path("artifacts/dev-10k/baselines.json")
+    )
+    args = parser.parse_args()
+    rows = read_jsonl(args.data_dir / "sample.jsonl")
+
+    started = time.monotonic()
+    report = {
+        "sample_rows": len(rows),
+        "target": "pAffinity",
+        "packages": {
+            package: importlib.metadata.version(package)
+            for package in ("numpy", "scipy", "scikit-learn")
+        },
+        "splits": {},
+    }
+    for manifest_path in sorted(args.data_dir.glob("split-*.jsonl")):
+        split_name = manifest_path.stem.removeprefix("split-")
+        report["splits"][split_name] = benchmark_split(rows, manifest_path)
+    report["total_seconds"] = time.monotonic() - started
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(report, indent=2) + "\n")
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

recovery/scripts/train_head.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env python3
+"""Train and evaluate the MitoInteract v2 head over cached frozen embeddings."""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import math
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+from safetensors.torch import save_file
+from torch.utils.data import DataLoader, TensorDataset
+
+from mitointeract_recovery.metrics import regression_metrics
+from mitointeract_recovery.model import MitoInteractHead, TargetScaler
+
+
+def read_manifest(path: Path) -> dict[str, str]:
+    with path.open() as handle:
+        return {
+            row["pair_id"]: row["split"]
+            for row in (json.loads(line) for line in handle if line.strip())
+        }
+
+
+def evaluate(
+    model: MitoInteractHead,
+    protein: torch.Tensor,
+    ligand: torch.Tensor,
+    targets: torch.Tensor,
+    indices: np.ndarray,
+    scaler: TargetScaler,
+) -> dict:
+    model.eval()
+    with torch.inference_mode():
+        prediction = (
+            scaler.decode(model(protein[indices], ligand[indices])).cpu().numpy()
+        )
+    return regression_metrics(targets[indices].cpu().numpy(), prediction)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--embeddings", type=Path, required=True)
+    parser.add_argument("--manifest", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/v2-head"))
+    parser.add_argument("--epochs", type=int, default=100)
+    parser.add_argument("--batch-size", type=int, default=256)
+    parser.add_argument("--learning-rate", type=float, default=1e-3)
+    parser.add_argument("--weight-decay", type=float, default=1e-4)
+    parser.add_argument("--warmup-ratio", type=float, default=0.05)
+    parser.add_argument("--patience", type=int, default=10)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--projection-dim", type=int, default=256)
+    parser.add_argument("--hidden-dim", type=int, default=512)
+    parser.add_argument("--dropout", type=float, default=0.1)
+    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
+    args = parser.parse_args()
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    device_name = (
+        "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device
+    )
+    if device_name == "auto":
+        device_name = "cpu"
+    device = torch.device(device_name)
+
+    arrays = np.load(args.embeddings)
+    pair_ids = arrays["pair_ids"].astype(str)
+    protein = torch.from_numpy(arrays["protein"]).float().to(device)
+    ligand = torch.from_numpy(arrays["ligand"]).float().to(device)
+    targets = torch.from_numpy(arrays["paffinity"]).float().to(device)
+    manifest = read_manifest(args.manifest)
+    split_indices = {
+        split: np.asarray(
+            [
+                index
+                for index, pair_id in enumerate(pair_ids)
+                if manifest.get(pair_id) == split
+            ]
+        )
+        for split in ("train", "validation", "test")
+    }
+    if any(not len(indices) for indices in split_indices.values()):
+        raise ValueError(
+            f"all splits must be non-empty: { {k: len(v) for k, v in split_indices.items()} }"
+        )
+
+    scaler = TargetScaler.fit(targets[split_indices["train"]])
+    encoded_targets = scaler.encode(targets)
+    train_dataset = TensorDataset(
+        protein[split_indices["train"]],
+        ligand[split_indices["train"]],
+        encoded_targets[split_indices["train"]],
+    )
+    generator = torch.Generator().manual_seed(args.seed)
+    loader = DataLoader(
+        train_dataset,
+        batch_size=min(args.batch_size, len(train_dataset)),
+        shuffle=True,
+        generator=generator,
+    )
+
+    model = MitoInteractHead(
+        protein.shape[1],
+        ligand.shape[1],
+        projection_dim=args.projection_dim,
+        hidden_dim=args.hidden_dim,
+        dropout=args.dropout,
+    ).to(device)
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=args.learning_rate,
+        weight_decay=args.weight_decay,
+    )
+    total_steps = max(1, args.epochs * len(loader))
+    warmup_steps = max(1, round(total_steps * args.warmup_ratio))
+
+    def lr_factor(step: int) -> float:
+        if step < warmup_steps:
+            return (step + 1) / warmup_steps
+        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+        return 0.5 * (1 + math.cos(math.pi * progress))
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_factor)
+    best_state = None
+    best_validation = math.inf
+    best_epoch = 0
+    remaining_patience = args.patience
+    history = []
+
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        losses = []
+        for batch_protein, batch_ligand, batch_target in loader:
+            optimizer.zero_grad(set_to_none=True)
+            prediction = model(batch_protein, batch_ligand)
+            loss = torch.nn.functional.mse_loss(prediction, batch_target)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            scheduler.step()
+            losses.append(float(loss.detach()))
+
+        validation = evaluate(
+            model,
+            protein,
+            ligand,
+            targets,
+            split_indices["validation"],
+            scaler,
+        )
+        history.append(
+            {
+                "epoch": epoch,
+                "train_standardized_mse": float(np.mean(losses)),
+                "validation": validation,
+                "learning_rate": optimizer.param_groups[0]["lr"],
+            }
+        )
+        if validation["rmse"] < best_validation:
+            best_validation = validation["rmse"]
+            best_epoch = epoch
+            best_state = copy.deepcopy(
+                {key: value.detach().cpu() for key, value in model.state_dict().items()}
+            )
+            remaining_patience = args.patience
+        else:
+            remaining_patience -= 1
+            if remaining_patience <= 0:
+                break
+
+    model.load_state_dict(best_state)
+    test_metrics = evaluate(
+        model, protein, ligand, targets, split_indices["test"], scaler
+    )
+    validation_metrics = evaluate(
+        model, protein, ligand, targets, split_indices["validation"], scaler
+    )
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    save_file(best_state, args.output_dir / "model.safetensors")
+
+    encoder_metadata_path = args.embeddings.with_suffix(".json")
+    encoder_metadata = (
+        json.loads(encoder_metadata_path.read_text())
+        if encoder_metadata_path.exists()
+        else None
+    )
+    report = {
+        "embedding_file": str(args.embeddings),
+        "manifest": str(args.manifest),
+        "rows": {name: len(indices) for name, indices in split_indices.items()},
+        "seed": args.seed,
+        "device": str(device),
+        "best_epoch": best_epoch,
+        "target_scaler": {"mean": scaler.mean, "std": scaler.std},
+        "head": {
+            "protein_dim": int(protein.shape[1]),
+            "ligand_dim": int(ligand.shape[1]),
+            "projection_dim": args.projection_dim,
+            "hidden_dim": args.hidden_dim,
+            "dropout": args.dropout,
+        },
+        "optimizer": {
+            "name": "AdamW",
+            "learning_rate": args.learning_rate,
+            "weight_decay": args.weight_decay,
+            "warmup_ratio": args.warmup_ratio,
+            "warmup_steps": warmup_steps,
+            "total_steps": total_steps,
+        },
+        "validation": validation_metrics,
+        "test": test_metrics,
+        "encoders": encoder_metadata,
+        "history": history,
+    }
+    (args.output_dir / "report.json").write_text(json.dumps(report, indent=2) + "\n")
+    (args.output_dir / "config.json").write_text(
+        json.dumps(
+            {
+                "format": "MitoInteract-v2-head",
+                "target": "pAffinity",
+                "target_scaler": report["target_scaler"],
+                "head": report["head"],
+                "encoders": encoder_metadata,
+            },
+            indent=2,
+        )
+        + "\n"
+    )
+    print(
+        json.dumps(
+            {key: value for key, value in report.items() if key != "history"}, indent=2
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

recovery/src/mitointeract_recovery/__init__.py

@@ -1,3 +1,3 @@
-from .affinity import micromolar_to_pkd, pkd_to_micromolar
+from .affinity import micromolar_to_paffinity, paffinity_to_micromolar
 
-__all__ = ["micromolar_to_pkd", "pkd_to_micromolar"]
+__all__ = ["micromolar_to_paffinity", "paffinity_to_micromolar"]

recovery/src/mitointeract_recovery/affinity.py

@@ -1,18 +1,17 @@
-"""Binding-affinity unit conversions used by MitoInteract."""
-from __future__ import annotations
+"""Conversions for the dataset's generic molar affinity scale."""
 
 import math
 
 
-def micromolar_to_pkd(affinity_um: float) -> float:
-    """Convert dissociation constant in micromolar to pKd."""
+def micromolar_to_paffinity(affinity_um: float) -> float:
+    """Convert a positive micromolar affinity value to -log10(value in M)."""
     if not math.isfinite(affinity_um) or affinity_um <= 0:
-        raise ValueError("affinity_um must be a positive finite number")
+        raise ValueError("affinity_um must be finite and greater than zero")
     return -math.log10(affinity_um * 1e-6)
 
 
-def pkd_to_micromolar(pkd: float) -> float:
-    """Convert pKd to dissociation constant in micromolar."""
-    if not math.isfinite(pkd):
-        raise ValueError("pkd must be finite")
-    return (10 ** (-pkd)) * 1e6
+def paffinity_to_micromolar(paffinity: float) -> float:
+    """Convert generic pAffinity to a micromolar affinity value."""
+    if not math.isfinite(paffinity):
+        raise ValueError("paffinity must be finite")
+    return (10 ** (-paffinity)) * 1e6

recovery/src/mitointeract_recovery/chemistry.py

@@ -0,0 +1,34 @@
+"""Ligand canonicalization and grouping helpers."""
+
+from __future__ import annotations
+
+import hashlib
+
+from rdkit import Chem
+from rdkit.Chem.Scaffolds import MurckoScaffold
+
+
+def canonicalize_smiles(smiles: str) -> str:
+    """Return isomeric canonical SMILES or raise for an invalid molecule."""
+    molecule = Chem.MolFromSmiles(smiles)
+    if molecule is None:
+        raise ValueError(f"invalid SMILES: {smiles!r}")
+    return Chem.MolToSmiles(molecule, canonical=True, isomericSmiles=True)
+
+
+def stable_id(namespace: str, value: str) -> str:
+    digest = hashlib.sha256(value.encode("utf-8")).hexdigest()[:20]
+    return f"{namespace}_{digest}"
+
+
+def scaffold_id(smiles: str) -> str:
+    """Return a stable Bemis-Murcko scaffold group, preserving acyclic diversity."""
+    canonical = canonicalize_smiles(smiles)
+    molecule = Chem.MolFromSmiles(canonical)
+    scaffold = MurckoScaffold.GetScaffoldForMol(molecule)
+    scaffold_smiles = Chem.MolToSmiles(scaffold, canonical=True, isomericSmiles=True)
+    if scaffold_smiles:
+        return stable_id("scaffold", scaffold_smiles)
+    # Empty Murcko scaffolds are common for acyclic molecules. Treat each
+    # canonical molecule as its own group instead of collapsing all of them.
+    return stable_id("acyclic", canonical)

recovery/src/mitointeract_recovery/metrics.py

@@ -0,0 +1,37 @@
+"""Regression metrics with JSON-safe handling for undefined correlations."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+import numpy as np
+from scipy.stats import pearsonr, spearmanr
+from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
+
+
+def _correlation(function, y_true: np.ndarray, y_pred: np.ndarray) -> float | None:
+    if len(y_true) < 2 or np.ptp(y_true) == 0 or np.ptp(y_pred) == 0:
+        return None
+    value = float(function(y_true, y_pred).statistic)
+    return value if math.isfinite(value) else None
+
+
+def regression_metrics(
+    y_true: Sequence[float], y_pred: Sequence[float]
+) -> dict[str, float | None]:
+    truth = np.asarray(y_true, dtype=np.float64)
+    predicted = np.asarray(y_pred, dtype=np.float64)
+    if truth.shape != predicted.shape or truth.ndim != 1 or not len(truth):
+        raise ValueError(
+            "y_true and y_pred must be non-empty one-dimensional sequences of equal length"
+        )
+    if not np.isfinite(truth).all() or not np.isfinite(predicted).all():
+        raise ValueError("metrics inputs must be finite")
+    return {
+        "rmse": float(math.sqrt(mean_squared_error(truth, predicted))),
+        "mae": float(mean_absolute_error(truth, predicted)),
+        "pearson_r": _correlation(pearsonr, truth, predicted),
+        "spearman_r": _correlation(spearmanr, truth, predicted),
+        "r2": float(r2_score(truth, predicted)) if len(truth) >= 2 else None,
+    }

recovery/src/mitointeract_recovery/model.py

@@ -0,0 +1,84 @@
+"""Corrected trainable MitoInteract v2 regression head.
+
+The frozen protein and molecule encoders are intentionally kept outside this
+module. Their pooled embeddings can be cached, versioned, and reused across
+head experiments without serializing hundreds of megabytes of frozen weights.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+
+@dataclass(frozen=True)
+class TargetScaler:
+    mean: float
+    std: float
+
+    @classmethod
+    def fit(cls, values: torch.Tensor) -> "TargetScaler":
+        mean = float(values.mean())
+        std = float(values.std(unbiased=False))
+        if not std > 0:
+            raise ValueError("target standard deviation must be positive")
+        return cls(mean=mean, std=std)
+
+    def encode(self, values: torch.Tensor) -> torch.Tensor:
+        return (values - self.mean) / self.std
+
+    def decode(self, values: torch.Tensor) -> torch.Tensor:
+        return values * self.std + self.mean
+
+
+class MitoInteractHead(nn.Module):
+    """Fuse pooled protein/ligand embeddings and regress standardized pAffinity."""
+
+    def __init__(
+        self,
+        protein_dim: int,
+        ligand_dim: int,
+        *,
+        projection_dim: int = 256,
+        hidden_dim: int = 512,
+        dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.protein_projection = nn.Sequential(
+            nn.LayerNorm(protein_dim),
+            nn.Linear(protein_dim, projection_dim),
+            nn.GELU(),
+        )
+        self.ligand_projection = nn.Sequential(
+            nn.LayerNorm(ligand_dim),
+            nn.Linear(ligand_dim, projection_dim),
+            nn.GELU(),
+        )
+        interaction_dim = projection_dim * 4
+        self.regressor = nn.Sequential(
+            nn.LayerNorm(interaction_dim),
+            nn.Linear(interaction_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim // 2, 1),
+        )
+
+    def forward(
+        self, protein_embedding: torch.Tensor, ligand_embedding: torch.Tensor
+    ) -> torch.Tensor:
+        if protein_embedding.ndim != 2 or ligand_embedding.ndim != 2:
+            raise ValueError("encoder embeddings must be rank-2 [batch, features]")
+        if protein_embedding.shape[0] != ligand_embedding.shape[0]:
+            raise ValueError("protein and ligand batch sizes must match")
+        protein = self.protein_projection(protein_embedding)
+        ligand = self.ligand_projection(ligand_embedding)
+        interaction = torch.cat(
+            [protein, ligand, protein * ligand, torch.abs(protein - ligand)],
+            dim=-1,
+        )
+        return self.regressor(interaction).squeeze(-1)

recovery/src/mitointeract_recovery/splits.py

@@ -0,0 +1,108 @@
+"""Deterministic entity-disjoint split assignment."""
+
+from __future__ import annotations
+
+import hashlib
+from collections import defaultdict
+from collections.abc import Mapping, Sequence
+
+SPLIT_RATIOS = {"train": 0.8, "validation": 0.1, "test": 0.1}
+
+
+def _tie_breaker(seed: int, group: str) -> str:
+    return hashlib.sha256(f"{seed}:{group}".encode()).hexdigest()
+
+
+def assign_random_splits(
+    rows: Sequence[Mapping[str, object]],
+    *,
+    id_key: str = "pair_id",
+    seed: int = 42,
+    ratios: Mapping[str, float] = SPLIT_RATIOS,
+) -> dict[str, str]:
+    """Assign individual rows deterministically with exact approximate counts."""
+    if abs(sum(ratios.values()) - 1.0) > 1e-9 or any(
+        value <= 0 for value in ratios.values()
+    ):
+        raise ValueError("split ratios must be positive and sum to 1")
+    ordered = sorted(
+        (str(row[id_key]) for row in rows),
+        key=lambda pair_id: _tie_breaker(seed, pair_id),
+    )
+    assignments: dict[str, str] = {}
+    cursor = 0
+    names = list(ratios)
+    for index, name in enumerate(names):
+        if index == len(names) - 1:
+            stop = len(ordered)
+        else:
+            stop = cursor + round(len(ordered) * ratios[name])
+        for pair_id in ordered[cursor:stop]:
+            assignments[pair_id] = name
+        cursor = stop
+    return assignments
+
+
+def assign_grouped_splits(
+    rows: Sequence[Mapping[str, object]],
+    group_key: str,
+    *,
+    id_key: str = "pair_id",
+    seed: int = 42,
+    ratios: Mapping[str, float] = SPLIT_RATIOS,
+) -> dict[str, str]:
+    """Assign whole groups to train/validation/test with approximate row balance."""
+    if not rows:
+        return {}
+    if abs(sum(ratios.values()) - 1.0) > 1e-9 or any(
+        value <= 0 for value in ratios.values()
+    ):
+        raise ValueError("split ratios must be positive and sum to 1")
+
+    groups: dict[str, list[str]] = defaultdict(list)
+    for row in rows:
+        pair_id = str(row[id_key])
+        group = str(row[group_key])
+        groups[group].append(pair_id)
+
+    targets = {name: len(rows) * ratio for name, ratio in ratios.items()}
+    counts = {name: 0 for name in ratios}
+    assignments: dict[str, str] = {}
+    ordered = sorted(
+        groups.items(), key=lambda item: (-len(item[1]), _tie_breaker(seed, item[0]))
+    )
+
+    for _, pair_ids in ordered:
+        # Choose the split with the largest normalized deficit. Once all are
+        # over target, choose the least-filled split relative to its target.
+        def priority(name: str) -> tuple[float, float, str]:
+            deficit = (targets[name] - counts[name]) / targets[name]
+            fill = counts[name] / targets[name]
+            return deficit, -fill, name
+
+        selected = max(ratios, key=priority)
+        for pair_id in pair_ids:
+            assignments[pair_id] = selected
+        counts[selected] += len(pair_ids)
+
+    return assignments
+
+
+def assert_group_disjoint(
+    rows: Sequence[Mapping[str, object]],
+    assignments: Mapping[str, str],
+    group_key: str,
+    *,
+    id_key: str = "pair_id",
+) -> None:
+    """Raise when an entity group occurs in more than one split."""
+    seen: dict[str, str] = {}
+    for row in rows:
+        pair_id = str(row[id_key])
+        group = str(row[group_key])
+        split = assignments[pair_id]
+        previous = seen.setdefault(group, split)
+        if previous != split:
+            raise AssertionError(
+                f"group {group!r} appears in both {previous!r} and {split!r}"
+            )

recovery/tests/test_affinity.py

@@ -1,28 +1,31 @@
 import math
+
 import pytest
 
-from mitointeract_recovery import micromolar_to_pkd, pkd_to_micromolar
+from mitointeract_recovery import micromolar_to_paffinity, paffinity_to_micromolar
 
 
-def test_one_micromolar_is_pkd_six():
-    assert micromolar_to_pkd(1.0) == pytest.approx(6.0)
+def test_one_micromolar_is_paffinity_six():
+    assert micromolar_to_paffinity(1.0) == pytest.approx(6.0)
 
 
 def test_dataset_example_500_micromolar():
-    assert micromolar_to_pkd(500.0) == pytest.approx(3.3010299957)
+    assert micromolar_to_paffinity(500.0) == pytest.approx(3.3010299957)
 
 
-@pytest.mark.parametrize("pkd", [2.0, 3.3010299957, 6.0, 7.5850267, 9.0])
-def test_round_trip(pkd):
-    assert micromolar_to_pkd(pkd_to_micromolar(pkd)) == pytest.approx(pkd)
+@pytest.mark.parametrize("paffinity", [2.0, 3.3010299957, 6.0, 7.5850267, 9.0])
+def test_round_trip(paffinity):
+    assert micromolar_to_paffinity(paffinity_to_micromolar(paffinity)) == pytest.approx(
+        paffinity
+    )
 
 
-@pytest.mark.parametrize("bad", [0.0, -1.0, math.inf, -math.inf, math.nan])
-def test_invalid_micromolar_rejected(bad):
+@pytest.mark.parametrize("bad", [0.0, -1.0, math.inf, math.nan])
+def test_invalid_affinity_rejected(bad):
     with pytest.raises(ValueError):
-        micromolar_to_pkd(bad)
+        micromolar_to_paffinity(bad)
 
 
-def test_nonfinite_pkd_rejected():
+def test_nonfinite_paffinity_rejected():
     with pytest.raises(ValueError):
-        pkd_to_micromolar(math.nan)
+        paffinity_to_micromolar(math.nan)

recovery/tests/test_checked_reports.py

@@ -0,0 +1,55 @@
+import json
+from pathlib import Path
+
+
+REPORTS = Path(__file__).parents[1] / "reports"
+
+
+def load(name: str) -> dict:
+    return json.loads((REPORTS / name).read_text())
+
+
+def model(report: dict, split: str, name: str) -> dict:
+    return next(
+        item for item in report["splits"][split]["models"] if item["name"] == name
+    )
+
+
+def test_checked_sample_and_manifests():
+    audit = load("dev-10k-audit.json")
+    assert audit["dataset_revision"] == "11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b"
+    assert (
+        audit["source_sha256"]
+        == "d50f4cae794298f32c39f75a9e8603528d20e50e4fa757025def412f5b2a78d8"
+    )
+    assert audit["sample_rows"] == 10_000
+    assert (
+        audit["external_split_audits"]["protein_similarity_50"][
+            "protein_cluster_overlap"
+        ]
+        == 0
+    )
+
+
+def test_combined_baseline_beats_train_mean_on_every_checked_split():
+    report = load("dev-10k-baselines.json")
+    assert report["target"] == "pAffinity"
+    for split_name, split in report["splits"].items():
+        combined = model(report, split_name, "combined_char_ridge")["test"]
+        mean = split["mean_baseline"]["test"]
+        assert combined["rmse"] < mean["rmse"]
+        assert combined["pearson_r"] > 0
+        assert combined["spearman_r"] > 0
+
+
+def test_similarity_split_passes_initial_recovery_gate():
+    report = load("dev-10k-baselines.json")
+    result = model(report, "protein_similarity_50", "combined_char_ridge")["test"]
+    assert result["pearson_r"] > 0.30
+    assert result["spearman_r"] > 0.30
+
+
+def test_actual_encoder_overfit_sanity_passed():
+    report = load("overfit-actual-embeddings.json")
+    assert report["passed"] is True
+    assert report["rmse_paffinity"] <= report["acceptance_max_rmse_paffinity"]

recovery/tests/test_chemistry.py

@@ -0,0 +1,20 @@
+import pytest
+
+from mitointeract_recovery.chemistry import canonicalize_smiles, scaffold_id
+
+
+def test_canonical_smiles_normalizes_equivalent_strings():
+    assert canonicalize_smiles("C(C)O") == canonicalize_smiles("CCO")
+
+
+def test_bemis_murcko_scaffold_groups_ring_substitutions():
+    assert scaffold_id("Cc1ccccc1") == scaffold_id("Oc1ccccc1")
+
+
+def test_acyclic_scaffolds_do_not_collapse_into_one_group():
+    assert scaffold_id("CCO") != scaffold_id("CCCC")
+
+
+def test_invalid_smiles_rejected():
+    with pytest.raises(ValueError):
+        canonicalize_smiles("not-a-smiles")

recovery/tests/test_legacy_guard.py

@@ -0,0 +1,29 @@
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+pytest.importorskip("torch")
+pytest.importorskip("transformers")
+
+
+def load_legacy_module():
+    path = Path(__file__).parents[2] / "model.py"
+    spec = importlib.util.spec_from_file_location("mitointeract_legacy_model", path)
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_legacy_pickle_loading_is_blocked_by_default():
+    legacy = load_legacy_module()
+    with pytest.raises(RuntimeError, match="unsafe pickle"):
+        legacy.load_model("not-opened.pt")
+
+
+def test_legacy_output_cannot_be_called_pkd_by_default():
+    legacy = load_legacy_module()
+    with pytest.raises(RuntimeError, match="cannot be interpreted as pKd"):
+        legacy.predict_binding(None, "PROTEIN", "CC")

recovery/tests/test_metrics.py

@@ -0,0 +1,18 @@
+import pytest
+
+from mitointeract_recovery.metrics import regression_metrics
+
+
+def test_perfect_predictions():
+    metrics = regression_metrics([1.0, 2.0, 3.0], [1.0, 2.0, 3.0])
+    assert metrics["rmse"] == pytest.approx(0.0)
+    assert metrics["mae"] == pytest.approx(0.0)
+    assert metrics["pearson_r"] == pytest.approx(1.0)
+    assert metrics["spearman_r"] == pytest.approx(1.0)
+    assert metrics["r2"] == pytest.approx(1.0)
+
+
+def test_constant_predictions_have_null_correlations():
+    metrics = regression_metrics([1.0, 2.0, 3.0], [2.0, 2.0, 2.0])
+    assert metrics["pearson_r"] is None
+    assert metrics["spearman_r"] is None

recovery/tests/test_model.py

@@ -0,0 +1,32 @@
+import pytest
+
+torch = pytest.importorskip("torch")
+from torch import nn  # noqa: E402
+
+from mitointeract_recovery.model import MitoInteractHead, TargetScaler  # noqa: E402
+
+
+def test_head_output_shape():
+    model = MitoInteractHead(16, 12, projection_dim=8, hidden_dim=16, dropout=0.0)
+    output = model(torch.randn(4, 16), torch.randn(4, 12))
+    assert output.shape == (4,)
+
+
+def test_head_uses_no_batch_norm():
+    model = MitoInteractHead(16, 12, projection_dim=8, hidden_dim=16)
+    assert not any(
+        isinstance(module, nn.modules.batchnorm._BatchNorm)
+        for module in model.modules()
+    )
+
+
+def test_target_scaler_round_trip():
+    values = torch.tensor([3.0, 5.0, 7.0, 9.0])
+    scaler = TargetScaler.fit(values)
+    assert torch.allclose(scaler.decode(scaler.encode(values)), values)
+
+
+def test_mismatched_batches_rejected():
+    model = MitoInteractHead(16, 12, projection_dim=8, hidden_dim=16)
+    with pytest.raises(ValueError):
+        model(torch.randn(3, 16), torch.randn(4, 12))

recovery/tests/test_splits.py

@@ -0,0 +1,45 @@
+from mitointeract_recovery.splits import (
+    assign_grouped_splits,
+    assign_random_splits,
+    assert_group_disjoint,
+)
+
+
+ROWS = [
+    {"pair_id": "a1", "protein_id": "p1", "scaffold_id": "s1"},
+    {"pair_id": "a2", "protein_id": "p1", "scaffold_id": "s2"},
+    {"pair_id": "b1", "protein_id": "p2", "scaffold_id": "s1"},
+    {"pair_id": "c1", "protein_id": "p3", "scaffold_id": "s3"},
+    {"pair_id": "d1", "protein_id": "p4", "scaffold_id": "s4"},
+    {"pair_id": "e1", "protein_id": "p5", "scaffold_id": "s5"},
+]
+
+
+def test_random_split_has_expected_counts():
+    rows = [{"pair_id": str(index)} for index in range(100)]
+    assignments = assign_random_splits(rows, seed=42)
+    assert list(assignments.values()).count("train") == 80
+    assert list(assignments.values()).count("validation") == 10
+    assert list(assignments.values()).count("test") == 10
+
+
+def test_grouped_split_is_deterministic():
+    first = assign_grouped_splits(ROWS, "protein_id", seed=42)
+    second = assign_grouped_splits(ROWS, "protein_id", seed=42)
+    assert first == second
+
+
+def test_cold_protein_groups_are_disjoint():
+    assignments = assign_grouped_splits(ROWS, "protein_id", seed=42)
+    assert_group_disjoint(ROWS, assignments, "protein_id")
+
+
+def test_cold_scaffold_groups_are_disjoint():
+    assignments = assign_grouped_splits(ROWS, "scaffold_id", seed=7)
+    assert_group_disjoint(ROWS, assignments, "scaffold_id")
+
+
+def test_every_pair_is_assigned_once():
+    assignments = assign_grouped_splits(ROWS, "protein_id", seed=42)
+    assert set(assignments) == {row["pair_id"] for row in ROWS}
+    assert set(assignments.values()) <= {"train", "validation", "test"}