From d21386997378d5014ca842ac431c4084bc7aa47c Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 17:54:29 +0200 Subject: [PATCH 01/19] feat(peptide-property-ng): unified peptide property predictor New research-track package: a clean-slate, Depthcharge-based neural network predicting fragment intensity, ion mobility/CCS, retention time and precursor charge from one shared encoder. Key design points: - Hybrid modification encoding: each residue carries both a learnable per-UNIMOD token embedding and an atomic-composition chemistry feature (from sagepy's UNIMOD tables), so rare/unseen modifications generalise and the scheme stays open-vocab (no re-mint when UNIMOD grows). - Fragment-indexed intensity head: one prediction per cleavage site, removing Prosit's fixed 174-vector 30-residue cap. - Instrument / acquisition-mode conditioning via the encoder global token; charge / m-z / collision energy conditioned at the heads so the charge head never sees its own label. - CCS head ports the imspy sqrt(m/z) per-charge physics prior. Trains on Sage search results. See docs/nn-architecture-exploration in the claudius-proteomics project for the design rationale. --- packages/peptide-property-ng/.gitignore | 11 + packages/peptide-property-ng/README.md | 40 ++++ packages/peptide-property-ng/pyproject.toml | 38 ++++ .../src/peptide_property_ng/__init__.py | 8 + .../src/peptide_property_ng/data/__init__.py | 1 + .../src/peptide_property_ng/data/collate.py | 61 ++++++ .../data/fragment_targets.py | 64 ++++++ .../peptide_property_ng/data/sage_dataset.py | 206 ++++++++++++++++++ .../src/peptide_property_ng/data/splits.py | 48 ++++ .../src/peptide_property_ng/eval/__init__.py | 1 + .../src/peptide_property_ng/eval/evaluate.py | 52 +++++ .../src/peptide_property_ng/eval/metrics.py | 68 ++++++ .../src/peptide_property_ng/losses.py | 72 ++++++ .../src/peptide_property_ng/model/__init__.py | 1 + .../src/peptide_property_ng/model/config.py | 105 +++++++++ .../peptide_property_ng/model/embedding.py | 93 ++++++++ .../src/peptide_property_ng/model/encoder.py | 107 +++++++++ .../model/heads/__init__.py | 1 + .../peptide_property_ng/model/heads/ccs.py | 64 ++++++ .../model/heads/intensity.py | 59 +++++ .../peptide_property_ng/model/heads/scalar.py | 47 ++++ .../peptide_property_ng/model/multitask.py | 75 +++++++ .../modifications/__init__.py | 1 + .../modifications/build_table.py | 95 ++++++++ .../modifications/composition.py | 86 ++++++++ .../modifications/data/.gitkeep | 1 + .../data/mod_composition_table.npz | Bin 0 -> 15340 bytes .../src/peptide_property_ng/train/__init__.py | 1 + .../src/peptide_property_ng/train/train.py | 145 ++++++++++++ .../peptide-property-ng/tests/__init__.py | 0 .../tests/test_composition.py | 59 +++++ .../peptide-property-ng/tests/test_encoder.py | 69 ++++++ .../tests/test_multitask.py | 69 ++++++ 33 files changed, 1748 insertions(+) create mode 100644 packages/peptide-property-ng/.gitignore create mode 100644 packages/peptide-property-ng/README.md create mode 100644 packages/peptide-property-ng/pyproject.toml create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/__init__.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/__init__.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/collate.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/fragment_targets.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/splits.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/eval/__init__.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/losses.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/__init__.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/config.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/embedding.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/heads/__init__.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/heads/ccs.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/heads/intensity.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/heads/scalar.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/modifications/__init__.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/modifications/build_table.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/modifications/composition.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/modifications/data/.gitkeep create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/modifications/data/mod_composition_table.npz create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/train/__init__.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/train/train.py create mode 100644 packages/peptide-property-ng/tests/__init__.py create mode 100644 packages/peptide-property-ng/tests/test_composition.py create mode 100644 packages/peptide-property-ng/tests/test_encoder.py create mode 100644 packages/peptide-property-ng/tests/test_multitask.py diff --git a/packages/peptide-property-ng/.gitignore b/packages/peptide-property-ng/.gitignore new file mode 100644 index 000000000..015345982 --- /dev/null +++ b/packages/peptide-property-ng/.gitignore @@ -0,0 +1,11 @@ +# training / evaluation outputs +runs/ + +# python build / cache +__pycache__/ +*.py[cod] +*.egg-info/ +build/ +dist/ +.pytest_cache/ +.venv/ diff --git a/packages/peptide-property-ng/README.md b/packages/peptide-property-ng/README.md new file mode 100644 index 000000000..62fe1097d --- /dev/null +++ b/packages/peptide-property-ng/README.md @@ -0,0 +1,40 @@ +# peptide-property-ng + +A next-generation, clean-slate neural network for **peptide property prediction** +on timsTOF proteomics data — fragment intensity, ion mobility / CCS, retention +time and precursor charge from **one shared encoder**. + +It is a research-track sibling of `imspy-predictors`, not a replacement for the +production model. See `docs/nn-architecture-exploration/` in the +`claudius-proteomics` project for the design rationale. + +## What is different from the production `UnifiedPeptideModel` + +| | production `imspy_predictors` | this model | +|---|---|---| +| Modifications | opaque `[UNIMOD:N]` tokens (no chemistry) | **hybrid**: per-mod token **+** atomic-composition feature | +| Unseen mods | embedding stays ~random | generalize via chemistry; open-vocab | +| Conditioning | dormant instrument embedding, never used | instrument / acq-mode / charge / m-z / CE inside the encoder | +| Intensity output | fixed 174-vector → 30-residue cap | **fragment-indexed** → no length cap | +| Backbone | bespoke transformer | built on **Depthcharge** primitives | + +## Layout + +``` +src/peptide_property_ng/ + modifications/ UNIMOD atomic-composition table (from sagepy) + model/ hybrid embedding, Depthcharge-based encoder, task heads + data/ Sage-parquet dataset, fragment targets, splits, collate + train/ multi-task training + eval/ per-task + OOD + unseen-modification evaluation +``` + +## Quick start + +```bash +# 1. build the modification-composition table (one-off) +python -m peptide_property_ng.modifications.build_table + +# 2. train the first-prototype model on processed Sage results +python -m peptide_property_ng.train.train --datasets-glob '/scratch/claudius-proteomics/*' +``` diff --git a/packages/peptide-property-ng/pyproject.toml b/packages/peptide-property-ng/pyproject.toml new file mode 100644 index 000000000..2c77fac91 --- /dev/null +++ b/packages/peptide-property-ng/pyproject.toml @@ -0,0 +1,38 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "peptide-property-ng" +version = "0.0.1" +description = "Next-generation unified neural predictor for peptide properties (fragment intensity, ion mobility/CCS, retention time, charge) on timsTOF proteomics data." +readme = "README.md" +requires-python = ">=3.11" +license = "MIT" + +# Core, pip-installable dependencies. +dependencies = [ + "torch>=2.0", + "depthcharge-ms>=0.4.9", + "numpy>=1.24", + "pyarrow>=15", + "scipy>=1.10", +] + +# NOTE: two further dependencies are Rust-built monorepo artifacts, not on PyPI, +# and must be provided by the environment: +# - imspy_connector -> the Rust-backed `ProformaTokenizer` +# - sagepy / sagepy_connector -> the UNIMOD atomic-composition tables +# During development we run against the claudius-proteomics `.venv`, which has both. + +[project.optional-dependencies] +dev = ["pytest>=7"] + +[tool.hatch.build.targets.wheel] +packages = ["src/peptide_property_ng"] + +[tool.hatch.build.targets.wheel.force-include] +"src/peptide_property_ng/modifications/data" = "peptide_property_ng/modifications/data" + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/packages/peptide-property-ng/src/peptide_property_ng/__init__.py b/packages/peptide-property-ng/src/peptide_property_ng/__init__.py new file mode 100644 index 000000000..c23002225 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/__init__.py @@ -0,0 +1,8 @@ +"""peptide-property-ng — unified neural peptide property predictor. + +A clean-slate, Depthcharge-based model that predicts fragment intensity, ion +mobility / CCS, retention time and precursor charge from one shared encoder, +with a hybrid (token + atomic-composition) modification representation. +""" + +__version__ = "0.0.1" diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/__init__.py b/packages/peptide-property-ng/src/peptide_property_ng/data/__init__.py new file mode 100644 index 000000000..6e5ae62bb --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/__init__.py @@ -0,0 +1 @@ +"""Data — Sage-parquet dataset, fragment-intensity targets, splits, collation.""" diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/collate.py b/packages/peptide-property-ng/src/peptide_property_ng/data/collate.py new file mode 100644 index 000000000..2dc3fd511 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/collate.py @@ -0,0 +1,61 @@ +"""Variable-length collation of prepared examples into model batches.""" +from __future__ import annotations + +import numpy as np +import torch + +from peptide_property_ng.data.fragment_targets import N_ION_CHANNELS + + +def collate_examples( + examples: list[dict], + *, + pad_token_id: int = 61, + max_charge: int = 8, +) -> dict[str, torch.Tensor]: + """Collate a list of example dicts into a padded batch. + + Token sequences pad to the batch-max length with ``pad_token_id``; intensity + targets pad to ``(batch_max_len - 1, n_ion_channels)`` with ``-1`` (the + masked-out marker the spectral-angle loss ignores). + """ + batch_size = len(examples) + max_len = max(len(e["tokens"]) for e in examples) + + tokens = np.full((batch_size, max_len), pad_token_id, dtype=np.int64) + intensity = np.full((batch_size, max_len - 1, N_ION_CHANNELS), -1.0, dtype=np.float32) + for i, e in enumerate(examples): + length = len(e["tokens"]) + tokens[i, :length] = e["tokens"] + tgt = e["intensity_target"] + intensity[i, : tgt.shape[0], :] = tgt + + def _col(key: str, dtype) -> torch.Tensor: + return torch.tensor([e[key] for e in examples], dtype=dtype) + + charge = _col("charge", torch.long) + return { + "tokens": torch.from_numpy(tokens), + "charge": charge, + "precursor_mz": _col("precursor_mz", torch.float32), + "collision_energy": _col("collision_energy", torch.float32), + "instrument": _col("instrument", torch.long), + "acq_mode": _col("acq_mode", torch.long), + # targets + "intensity_target": torch.from_numpy(intensity), + "ccs_target": _col("ccs_target", torch.float32), + "ccs_valid": _col("ccs_valid", torch.bool), + "rt_target": _col("rt_target", torch.float32), + "rt_valid": _col("rt_valid", torch.bool), + # charge head target: class j == charge j+1 + "charge_target": charge.clamp(1, max_charge) - 1, + } + + +def make_collate_fn(pad_token_id: int = 61, max_charge: int = 8): + """Return a ``collate_fn`` bound to the given padding / charge settings.""" + + def _fn(examples: list[dict]) -> dict[str, torch.Tensor]: + return collate_examples(examples, pad_token_id=pad_token_id, max_charge=max_charge) + + return _fn diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/fragment_targets.py b/packages/peptide-property-ng/src/peptide_property_ng/data/fragment_targets.py new file mode 100644 index 000000000..aa4a6aa1b --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/fragment_targets.py @@ -0,0 +1,64 @@ +"""Fragment-indexed intensity targets from Sage ``matched_fragments``. + +Each peptide of length ``L`` yields an ``(L-1, 6)`` target — one row per +cleavage site, six channels for b/y ions at fragment charges 1-3. This is the +variable-length successor to Prosit's fixed 174-vector (no 30-residue cap). + +Convention (Prosit / imspy compatible): + -1 physically impossible channel (fragment charge > precursor charge) + 0 possible but unobserved fragment + 0..1 observed, base-peak-normalised intensity +The masked spectral-angle loss ignores -1 entries. +""" +from __future__ import annotations + +import numpy as np + +# Channel order: y ions at charge 1/2/3, then b ions at charge 1/2/3. +ION_CHANNELS: tuple[str, ...] = ("y+1", "y+2", "y+3", "b+1", "b+2", "b+3") +N_ION_CHANNELS = len(ION_CHANNELS) +MAX_FRAGMENT_CHARGE = 3 + + +def build_intensity_target( + peptide_len: int, + precursor_charge: int, + frag_type: np.ndarray, # ('b'|'y') per matched fragment + frag_ordinal: np.ndarray, # int + frag_charge: np.ndarray, # int + frag_intensity: np.ndarray, # float +) -> np.ndarray: + """Return the ``(peptide_len - 1, 6)`` intensity target for one PSM.""" + n_sites = peptide_len - 1 + target = np.zeros((n_sites, N_ION_CHANNELS), dtype=np.float32) + + # Mark fragment charges above the precursor charge as impossible. + for fc in range(1, MAX_FRAGMENT_CHARGE + 1): + if fc > precursor_charge: + target[:, fc - 1] = -1.0 # y channel + target[:, 3 + fc - 1] = -1.0 # b channel + + for ftype, ordi, fch, inten in zip( + frag_type, frag_ordinal, frag_charge, frag_intensity + ): + fch = int(fch) + if fch < 1 or fch > MAX_FRAGMENT_CHARGE: + continue + ordi = int(ordi) + if ftype == "b": + site, channel = ordi - 1, 3 + fch - 1 + elif ftype == "y": + # cleavage site i produces y_{n_sites - i + ...}: y_o sits at site n_sites - o + site, channel = n_sites - ordi, fch - 1 + else: + continue + if 0 <= site < n_sites and target[site, channel] >= 0.0: + target[site, channel] = max(target[site, channel], float(inten)) + + # Base-peak normalisation over the observed (positive) entries. + observed = target > 0.0 + if observed.any(): + peak = float(target[observed].max()) + if peak > 0.0: + target[observed] /= peak + return target diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py new file mode 100644 index 000000000..e5b503676 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py @@ -0,0 +1,206 @@ +"""Sage-parquet -> multi-task training examples. + +Reads a dataset's ``results.sage.parquet`` + ``matched_fragments.sage.parquet``, +filters to confident target PSMs, converts the Sage delta-mass peptide strings +to UNIMOD form (via ``sagepy_rescore``'s residue-specificity-aware converter), +tokenises, and builds per-task targets. +""" +from __future__ import annotations + +import glob +from collections import defaultdict +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq +from torch.utils.data import Dataset + +from peptide_property_ng.data.fragment_targets import build_intensity_target +from peptide_property_ng.data.splits import peptide_split + +_PROTON = 1.007276 + +_RES_COLS = [ + "psm_id", "peptide", "stripped_peptide", "charge", "calcmass", + "aligned_rt", "ion_mobility", "spectrum_q", "is_decoy", "rank", + "matched_peaks", "peptide_len", +] +_FRAG_COLS = [ + "psm_id", "fragment_type", "fragment_ordinals", "fragment_charge", "fragment_intensity", +] + + +def discover_sage_dirs(glob_pattern: str) -> list[Path]: + """Find non-empty ``.../processed/sage`` directories under a dataset-root glob. + + Datasets whose ``results.sage.parquet`` has zero rows (e.g. the legacy + ``TimsCompressionType=1`` deposits Sage cannot read) are skipped. + """ + out: list[Path] = [] + for root in sorted(glob.glob(glob_pattern)): + res = Path(root) / "processed" / "sage" / "results.sage.parquet" + if not res.exists(): + continue + try: + if pq.read_metadata(res).num_rows == 0: + continue + except Exception: + continue + out.append(res.parent) + return out + + +def _accession(sage_dir: Path) -> str: + # ...//processed/sage -> + return sage_dir.parts[-3] if len(sage_dir.parts) >= 3 else sage_dir.name + + +def prepare_examples( + sage_dir: str | Path, + *, + cap: int = 8000, + q_max: float = 0.01, + min_peaks: int = 6, + max_len: int = 64, + seed: int = 0, + instrument: int = 0, + acq_mode: int = 0, +) -> list[dict]: + """Load and prepare one dataset's PSMs into a list of example dicts.""" + from imspy_predictors.utilities.tokenizers import ProformaTokenizer + from sagepy_rescore.sage_loader import _parse_sage_peptide + + sage_dir = Path(sage_dir) + res_path = sage_dir / "results.sage.parquet" + frag_path = sage_dir / "matched_fragments.sage.parquet" + if not res_path.exists(): + return [] + res = pq.read_table(res_path, columns=_RES_COLS) + if res.num_rows == 0: + return [] + + keep = pc.and_( + pc.and_(pc.equal(res["is_decoy"], False), pc.equal(res["rank"], 1)), + pc.and_( + pc.less_equal(res["spectrum_q"], q_max), + pc.and_( + pc.greater_equal(res["matched_peaks"], min_peaks), + pc.less_equal(res["peptide_len"], max_len), + ), + ), + ) + res = res.filter(keep) + if res.num_rows == 0: + return [] + if res.num_rows > cap: + idx = np.sort(np.random.RandomState(seed).choice(res.num_rows, cap, replace=False)) + res = res.take(pa.array(idx)) + df = res.to_pydict() + n = len(df["psm_id"]) + + # Fragments for the kept PSMs, grouped by psm_id. + frag = pq.read_table(frag_path, columns=_FRAG_COLS) + frag = frag.filter(pc.is_in(frag["psm_id"], value_set=pa.array(list(set(df["psm_id"]))))) + fd = frag.to_pydict() + by_psm: dict[int, tuple[list, list, list, list]] = defaultdict(lambda: ([], [], [], [])) + for pid, ft, fo, fc, fi in zip( + fd["psm_id"], fd["fragment_type"], fd["fragment_ordinals"], + fd["fragment_charge"], fd["fragment_intensity"], + ): + cols = by_psm[pid] + cols[0].append(ft); cols[1].append(fo); cols[2].append(fc); cols[3].append(fi) + + # Sage emits delta-mass peptides; convert to UNIMOD, then tokenise in one batch. + # Batch tokenisation right-pads to the batch-max length, so the attention + # mask is needed to recover each peptide's true token count. + parsed = [_parse_sage_peptide(p) for p in df["peptide"]] + tok = ProformaTokenizer.with_defaults() + _encoded = tok([m for _, m in parsed]) + token_ids, attn_mask = _encoded["input_ids"], _encoded["attention_mask"] + + accession = _accession(sage_dir) + examples: list[dict] = [] + for k in range(n): + stripped, modseq = parsed[k] + if "[+" in modseq or "[-" in modseq: + continue # an unconverted delta mass — skip rather than train on a dropped mod + n_real = int(sum(attn_mask[k])) # real tokens incl. [CLS]/[SEP], excl. padding + residue_ids = token_ids[k][1 : n_real - 1] # strip [CLS]/[SEP] and right-padding + length = len(residue_ids) + # Terminal mods become standalone tokens -> token count != residue count; + # skip those rather than mis-align the fragment indexing. + if length != int(df["peptide_len"][k]) or length < 3 or length > max_len: + continue + + charge = int(df["charge"][k]) + mz = (float(df["calcmass"][k]) + charge * _PROTON) / charge + cols = by_psm.get(df["psm_id"][k]) + if not cols or not cols[0]: + continue # no matched fragments -> no intensity target + target = build_intensity_target( + length, charge, + np.asarray(cols[0]), + np.asarray(cols[1], dtype=np.int32), + np.asarray(cols[2], dtype=np.int32), + np.asarray(cols[3], dtype=np.float32), + ) + + im = df["ion_mobility"][k] + rt = df["aligned_rt"][k] + im = float(im) if im is not None else float("nan") + rt = float(rt) if rt is not None else float("nan") + examples.append( + { + "accession": accession, + "stripped": stripped, + "modseq": modseq, + "tokens": np.asarray(residue_ids, dtype=np.int64), + "charge": charge, + "precursor_mz": mz, + "collision_energy": 0.0, # not in Sage output for timsTOF diaPASEF + "instrument": int(instrument), + "acq_mode": int(acq_mode), + "intensity_target": target, + "ccs_target": im, + "ccs_valid": bool(np.isfinite(im) and im > 0.0), + "rt_target": rt, + "rt_valid": bool(np.isfinite(rt)), + } + ) + return examples + + +class SagePropertyDataset(Dataset): + """A thin wrapper over a list of prepared example dicts.""" + + def __init__(self, examples: list[dict]): + self.examples = examples + + def __len__(self) -> int: + return len(self.examples) + + def __getitem__(self, i: int) -> dict: + return self.examples[i] + + +def build_split_datasets( + sage_dirs: list[str | Path], + *, + cap: int = 8000, + seed: int = 0, + val_frac: float = 0.1, + test_frac: float = 0.1, + **prepare_kwargs, +) -> dict[str, SagePropertyDataset]: + """Prepare every dataset and split peptide-level into train / val / test.""" + examples: list[dict] = [] + for sage_dir in sage_dirs: + examples.extend(prepare_examples(sage_dir, cap=cap, seed=seed, **prepare_kwargs)) + + buckets: dict[str, list[dict]] = {"train": [], "val": [], "test": []} + for ex in examples: + split = peptide_split(ex["stripped"], val_frac=val_frac, test_frac=test_frac, seed=seed) + buckets[split].append(ex) + return {name: SagePropertyDataset(rows) for name, rows in buckets.items()} diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/splits.py b/packages/peptide-property-ng/src/peptide_property_ng/data/splits.py new file mode 100644 index 000000000..c85b05b9f --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/splits.py @@ -0,0 +1,48 @@ +"""Leakage-free dataset splits. + +Peptide-level split — a peptide (its *stripped*, modification-free sequence) is +hashed to a fixed bucket, so the same peptide never appears in two splits +regardless of which dataset, charge or run it came from. This is essential: +peptides recur heavily across the deposits. + +Also provides the unseen-modification split used by the key generalisation +probe — hold every PSM bearing a chosen modification out of train/val. +""" +from __future__ import annotations + +import hashlib +import re + +_UNIMOD_RE = re.compile(r"\[UNIMOD:(\d+)\]") + + +def _bucket(key: str, seed: int) -> float: + """Stable hash of ``key`` into ``[0, 1)``.""" + digest = hashlib.sha1(f"{seed}:{key}".encode()).hexdigest() + return (int(digest[:8], 16) % 1_000_000) / 1_000_000.0 + + +def peptide_split( + stripped_peptide: str, + *, + val_frac: float = 0.1, + test_frac: float = 0.1, + seed: int = 0, +) -> str: + """Assign a peptide to ``"train"`` / ``"val"`` / ``"test"`` by stable hash.""" + f = _bucket(stripped_peptide, seed) + if f < test_frac: + return "test" + if f < test_frac + val_frac: + return "val" + return "train" + + +def modifications_in(modified_peptide: str) -> set[int]: + """UNIMOD ids present in a UNIMOD-format peptide string.""" + return {int(x) for x in _UNIMOD_RE.findall(modified_peptide)} + + +def has_modification(modified_peptide: str, unimod_id: int) -> bool: + """True if the peptide carries the given UNIMOD modification.""" + return unimod_id in modifications_in(modified_peptide) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/eval/__init__.py b/packages/peptide-property-ng/src/peptide_property_ng/eval/__init__.py new file mode 100644 index 000000000..a1e61af4e --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/eval/__init__.py @@ -0,0 +1 @@ +"""Evaluation — per-task metrics, OOD probes, the unseen-modification split.""" diff --git a/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate.py b/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate.py new file mode 100644 index 000000000..0b32fa81f --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate.py @@ -0,0 +1,52 @@ +"""Evaluate a trained checkpoint on the held-out test split. + + python -m peptide_property_ng.eval.evaluate --checkpoint runs/prototype/best.pt +""" +from __future__ import annotations + +import argparse +import json + +import torch +from torch.utils.data import DataLoader + +from peptide_property_ng.data.collate import make_collate_fn +from peptide_property_ng.data.sage_dataset import build_split_datasets, discover_sage_dirs +from peptide_property_ng.eval.metrics import evaluate_split +from peptide_property_ng.model.config import PRESETS +from peptide_property_ng.model.multitask import UnifiedPeptidePropertyModel +from peptide_property_ng.modifications.composition import CompositionTable + + +def main() -> None: + ap = argparse.ArgumentParser(description="Evaluate a trained checkpoint.") + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--datasets-glob", default="/scratch/claudius-proteomics/*") + ap.add_argument("--cap", type=int, default=4000) + ap.add_argument("--max-datasets", type=int, default=0) + ap.add_argument("--batch-size", type=int, default=64) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--split", default="test", choices=["train", "val", "test"]) + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + args = ap.parse_args() + + ckpt = torch.load(args.checkpoint, map_location=args.device) + cfg = PRESETS[ckpt.get("preset", "small")] + model = UnifiedPeptidePropertyModel(cfg, CompositionTable.load()).to(args.device) + model.load_state_dict(ckpt["model_state_dict"]) + + sage_dirs = discover_sage_dirs(args.datasets_glob) + if args.max_datasets: + sage_dirs = sage_dirs[: args.max_datasets] + splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed) + loader = DataLoader( + splits[args.split], + batch_size=args.batch_size, + collate_fn=make_collate_fn(cfg.pad_token_id, cfg.max_charge), + ) + metrics = evaluate_split(model, loader, args.device) + print(json.dumps({"split": args.split, "checkpoint": args.checkpoint, "metrics": metrics}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py b/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py new file mode 100644 index 000000000..94803418f --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py @@ -0,0 +1,68 @@ +"""Per-task evaluation metrics.""" +from __future__ import annotations + +import torch + +from peptide_property_ng.losses import masked_spectral_angle + + +def _pearson(x: torch.Tensor, y: torch.Tensor) -> float: + if x.numel() < 2: + return float("nan") + x = x - x.mean() + y = y - y.mean() + denom = x.norm() * y.norm() + return float((x * y).sum() / denom) if denom > 0 else float("nan") + + +@torch.no_grad() +def evaluate_split(model, loader, device: str = "cpu") -> dict[str, float]: + """Run the model over a DataLoader and return per-task metrics. + + Returns: intensity spectral angle (similarity, higher better), CCS / RT + median absolute error, RT Pearson r, charge accuracy. + """ + model.eval() + sa_sum = sa_n = 0.0 + ccs_abs: list[torch.Tensor] = [] + rt_pred: list[torch.Tensor] = [] + rt_true: list[torch.Tensor] = [] + charge_correct = charge_n = 0.0 + + for batch in loader: + batch = {k: v.to(device) for k, v in batch.items()} + out = model(batch) + + if "intensity" in out: + sa = 1.0 - masked_spectral_angle(out["intensity"], batch["intensity_target"]) + sa_sum += float(sa.sum()) + sa_n += sa.numel() + + if "ccs" in out: + mean, _ = out["ccs"] + v = batch["ccs_valid"] + if v.any(): + ccs_abs.append((mean[v] - batch["ccs_target"][v]).abs().cpu()) + + if "rt" in out: + v = batch["rt_valid"] + if v.any(): + rt_pred.append(out["rt"][v].cpu()) + rt_true.append(batch["rt_target"][v].cpu()) + + if "charge" in out: + pred = out["charge"].argmax(dim=-1) + charge_correct += float((pred == batch["charge_target"]).sum()) + charge_n += pred.numel() + + ccs = torch.cat(ccs_abs) if ccs_abs else torch.tensor([]) + rtp = torch.cat(rt_pred) if rt_pred else torch.tensor([]) + rtt = torch.cat(rt_true) if rt_true else torch.tensor([]) + return { + "intensity_sa": sa_sum / sa_n if sa_n else float("nan"), + "ccs_mae": float(ccs.mean()) if ccs.numel() else float("nan"), + "ccs_median_ae": float(ccs.median()) if ccs.numel() else float("nan"), + "rt_mae": float((rtp - rtt).abs().mean()) if rtp.numel() else float("nan"), + "rt_pearson": _pearson(rtp, rtt), + "charge_acc": charge_correct / charge_n if charge_n else float("nan"), + } diff --git a/packages/peptide-property-ng/src/peptide_property_ng/losses.py b/packages/peptide-property-ng/src/peptide_property_ng/losses.py new file mode 100644 index 000000000..933ebaf60 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/losses.py @@ -0,0 +1,72 @@ +"""Multi-task losses for the unified peptide property model.""" +from __future__ import annotations + +import math + +import torch +from torch.nn import functional as F + + +def masked_spectral_angle(pred: torch.Tensor, target: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: + """Per-sample spectral-angle distance in ``[0, 1]`` (0 = identical spectra). + + ``pred`` / ``target`` are ``(B, sites, channels)``. Entries where + ``target < 0`` are masked out (impossible / padded fragment channels). + """ + mask = (target >= 0).float() + p = (pred * mask).reshape(pred.shape[0], -1) + t = (target.clamp(min=0.0) * mask).reshape(target.shape[0], -1) + p = F.normalize(p, dim=1, eps=eps) + t = F.normalize(t, dim=1, eps=eps) + cos = (p * t).sum(dim=1).clamp(-1.0 + eps, 1.0 - eps) + return 2.0 * torch.arccos(cos) / math.pi + + +class MultiTaskLoss: + """Weighted sum of per-task losses; tasks with no valid targets are skipped. + + - intensity : masked spectral-angle distance + - ccs : Gaussian negative log-likelihood (mean, std) vs 1/K0 + - rt : L1 on normalised retention time + - charge : cross-entropy + """ + + DEFAULT_WEIGHTS = {"intensity": 1.0, "ccs": 1.0, "rt": 0.5, "charge": 0.3} + + def __init__(self, weights: dict[str, float] | None = None): + self.weights = dict(self.DEFAULT_WEIGHTS if weights is None else weights) + + def __call__( + self, + outputs: dict[str, object], + batch: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, float]]: + parts: dict[str, torch.Tensor] = {} + + if "intensity" in outputs: + parts["intensity"] = masked_spectral_angle( + outputs["intensity"], batch["intensity_target"] + ).mean() + + if "ccs" in outputs: + mean, _std = outputs["ccs"] + valid = batch["ccs_valid"] + if valid.any(): + # L1 on the mean — stable for the prototype. Gaussian-NLL on + # (mean, std) is the eventual target but is unbounded-below and + # destabilises the fixed-weight multi-task sum; revisit together + # with uncertainty-based task weighting. + parts["ccs"] = F.l1_loss(mean[valid], batch["ccs_target"][valid]) + + if "rt" in outputs: + valid = batch["rt_valid"] + if valid.any(): + parts["rt"] = F.l1_loss(outputs["rt"][valid], batch["rt_target"][valid]) + + if "charge" in outputs: + parts["charge"] = F.cross_entropy(outputs["charge"], batch["charge_target"]) + + if not parts: + raise ValueError("no task produced a loss for this batch") + total = sum(self.weights.get(name, 1.0) * value for name, value in parts.items()) + return total, {name: float(value.detach()) for name, value in parts.items()} diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/__init__.py b/packages/peptide-property-ng/src/peptide_property_ng/model/__init__.py new file mode 100644 index 000000000..d3956ded4 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/__init__.py @@ -0,0 +1 @@ +"""Model — hybrid residue embedding, Depthcharge-based encoder, multi-task heads.""" diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py new file mode 100644 index 000000000..a53d435fa --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py @@ -0,0 +1,105 @@ +"""Model configuration — presets, instrument / acquisition vocabularies. + +Conditioning policy (deliberate, see the charge-leakage note): + - Instrument model and acquisition mode are conditioned *inside* the encoder + via the prepended global token — they cannot leak any task label. + - Charge, precursor m/z and collision energy are conditioned at the *heads* + that legitimately use them (intensity, CCS). They are deliberately kept out + of the encoder so the charge head — which reads the shared encoder output — + never sees its own label. One encoder pass, no leakage. +""" +from __future__ import annotations + +from dataclasses import dataclass, field + +# timsTOF-focused instrument vocabulary (from timstof_catalog.tsv instrument_model). +INSTRUMENTS: tuple[str, ...] = ( + "unknown", + "timsTOF", + "timsTOF Pro", + "timsTOF Pro 2", + "timsTOF HT", + "timsTOF SCP", + "timsTOF Ultra", + "timsTOF fleX", +) +ACQUISITION_MODES: tuple[str, ...] = ( + "unknown", + "DDA", + "DDA-PASEF", + "DIA", + "diaPASEF", + "PRM", + "MALDI", +) +INSTRUMENT_TO_ID: dict[str, int] = {n: i for i, n in enumerate(INSTRUMENTS)} +ACQ_TO_ID: dict[str, int] = {n: i for i, n in enumerate(ACQUISITION_MODES)} + + +def instrument_id(name: str | None) -> int: + """Map an instrument-model string to an id, tolerating unknowns/aliases.""" + if not name: + return 0 + n = " ".join(str(name).split()) # normalise whitespace + if n in INSTRUMENT_TO_ID: + return INSTRUMENT_TO_ID[n] + low = n.lower() + for cand, idx in INSTRUMENT_TO_ID.items(): + if cand.lower() == low: + return idx + return 0 # "unknown" + + +def acquisition_id(name: str | None) -> int: + """Map an acquisition-mode string to an id, tolerating unknowns.""" + if not name: + return 0 + n = str(name).strip() + for cand, idx in ACQ_TO_ID.items(): + if cand.lower() == n.lower(): + return idx + return 0 + + +@dataclass +class ModelConfig: + """Architecture + vocabulary configuration for the unified predictor.""" + + # encoder + d_model: int = 256 + n_layers: int = 6 + n_heads: int = 8 + dim_feedforward: int = 1024 + dropout: float = 0.1 + max_seq_len: int = 64 # peptide-length cap — far above the old 174-vector's 30 + + # tokenizer / modification vocabulary + vocab_size: int = 2175 + pad_token_id: int = 61 + n_elements: int = 31 # atomic-composition channels (see modifications.composition) + + # hybrid embedding + comp_fusion: str = "add" # "add" | "gate" + + # conditioning + max_charge: int = 8 + n_instruments: int = len(INSTRUMENTS) + n_acq_modes: int = len(ACQUISITION_MODES) + + # heads + n_ion_channels: int = 6 # b/y ions x fragment charges 1-3 (Prosit-compatible) + tasks: tuple[str, ...] = ("intensity", "ccs", "rt", "charge") + + def __post_init__(self) -> None: + if self.d_model % self.n_heads != 0: + raise ValueError(f"d_model {self.d_model} not divisible by n_heads {self.n_heads}") + + +# First-prototype preset — roughly on par with the production BASE encoder so +# iteration is fast on the ~19 datasets available now. +SMALL = ModelConfig(d_model=256, n_layers=6, n_heads=8, dim_feedforward=1024) + +# Research preset — larger; for the full ~138-dataset corpus. +RESEARCH = ModelConfig(d_model=384, n_layers=8, n_heads=8, dim_feedforward=1536) + +PRESETS: dict[str, ModelConfig] = {"small": SMALL, "research": RESEARCH} diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/embedding.py b/packages/peptide-property-ng/src/peptide_property_ng/model/embedding.py new file mode 100644 index 000000000..3fdfc79e6 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/embedding.py @@ -0,0 +1,93 @@ +"""Hybrid residue embedding — learnable token embedding + atomic-composition. + +The central design decision: every residue is represented by BOTH + (a) a learnable per-token embedding (great for common modifications, which + have plenty of training data), and + (b) a chemistry feature derived from the modification's atomic composition + (lets rare / unseen modifications generalise — they inherit meaning from + chemically similar mods, and the scheme is open-vocab). + +A bare residue or special token has an all-zero composition delta; with the +bias-free CompositionEncoder this maps to an exact zero contribution, so +unmodified residues fall back to the pure token embedding. +""" +from __future__ import annotations + +import torch +from torch import nn + +from peptide_property_ng.modifications.composition import CompositionTable + + +class CompositionEncoder(nn.Module): + """Maps a signed atomic-composition vector to a d_model chemistry embedding. + + Bias-free so that an all-zero composition (a bare residue) maps to an exact + zero vector — the additive-fusion semantics then mean "unmodified residue = + pure token embedding". + """ + + def __init__(self, n_elements: int, d_model: int, hidden: int | None = None): + super().__init__() + hidden = hidden or d_model // 2 + self.net = nn.Sequential( + nn.Linear(n_elements, hidden, bias=False), + nn.GELU(), + nn.Linear(hidden, d_model, bias=False), + ) + + def forward(self, comp: torch.Tensor) -> torch.Tensor: # (..., n_elements) -> (..., d_model) + return self.net(comp) + + +class HybridResidueEmbedding(nn.Module): + """Token embedding + composition embedding, fused per residue. + + Parameters + ---------- + vocab_size, d_model, pad_token_id + Tokenizer / model dimensions. + composition_table + Per-token atomic-composition deltas (see modifications.composition). + fusion + ``"add"`` — ``token + composition`` (default). + ``"gate"`` — ``token + sigmoid(g) * composition`` with a learnable + per-channel gate ``g`` (init 0 → gate 0.5); the recommended upgrade if + the unseen-modification probe shows the composition branch is ignored. + """ + + def __init__( + self, + vocab_size: int, + d_model: int, + composition_table: CompositionTable, + pad_token_id: int, + fusion: str = "add", + ): + super().__init__() + if fusion not in ("add", "gate"): + raise ValueError(f"unknown fusion '{fusion}'") + self.d_model = d_model + self.pad_token_id = pad_token_id + self.fusion = fusion + + # +1 row to match depthcharge's nn.Embedding(n_tokens + 1, ...) convention. + self.token_emb = nn.Embedding(vocab_size + 1, d_model, padding_idx=pad_token_id) + self.comp_encoder = CompositionEncoder(composition_table.n_elements, d_model) + + # Composition lookup as a frozen buffer, indexed by token id; signed-log1p + # normalised. One extra zero row to align with token_emb's vocab_size + 1. + comp = composition_table.as_tensor(normalize=True) # (vocab_size, n_elements) + comp = torch.cat([comp, torch.zeros(1, comp.shape[1])], dim=0) + self.register_buffer("composition", comp, persistent=False) + + if fusion == "gate": + self.gate = nn.Parameter(torch.zeros(d_model)) + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + """tokens: (batch, seq_len) long -> (batch, seq_len, d_model).""" + token_vec = self.token_emb(tokens) + comp_vec = self.comp_encoder(self.composition[tokens]) + if self.fusion == "gate": + comp_vec = torch.sigmoid(self.gate) * comp_vec + return token_vec + comp_vec diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py b/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py new file mode 100644 index 000000000..2b80dec94 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py @@ -0,0 +1,107 @@ +"""Depthcharge-based shared encoder for the unified peptide property predictor.""" +from __future__ import annotations + +import torch +from depthcharge.transformers import AnalyteTransformerEncoder +from torch import nn + +from peptide_property_ng.model.config import ModelConfig +from peptide_property_ng.model.embedding import HybridResidueEmbedding +from peptide_property_ng.modifications.composition import CompositionTable + + +class PeptidePropertyEncoder(AnalyteTransformerEncoder): + """Shared transformer encoder. + + Built on Depthcharge's ``AnalyteTransformerEncoder``, but: + - the plain token embedding is replaced with the hybrid (token + + atomic-composition) residue embedding, and + - conditioning on instrument model + acquisition mode happens via the + prepended global token. Charge / precursor m-z / collision energy are + deliberately *not* fed here — they are conditioned at the heads that use + them, so the charge head (which reads this shared output) never sees its + own label. One encoder pass, no label leakage. + + ``forward(tokens, instrument, acq_mode) -> (latent, padding_mask)`` + - ``latent`` ``(batch, 1 + seq_len, d_model)`` — index 0 is the global token + - ``padding_mask`` ``(batch, 1 + seq_len)`` bool — ``True`` where padded + """ + + def __init__(self, cfg: ModelConfig, composition_table: CompositionTable): + super().__init__( + n_tokens=cfg.vocab_size, + d_model=cfg.d_model, + nhead=cfg.n_heads, + dim_feedforward=cfg.dim_feedforward, + n_layers=cfg.n_layers, + dropout=cfg.dropout, + positional_encoder=True, + padding_int=cfg.pad_token_id, + ) + self.cfg = cfg + + # The base class built a plain nn.Embedding token encoder; drop it and + # use the hybrid residue embedding instead. + del self.token_encoder + self.hybrid_embedding = HybridResidueEmbedding( + vocab_size=cfg.vocab_size, + d_model=cfg.d_model, + composition_table=composition_table, + pad_token_id=cfg.pad_token_id, + fusion=cfg.comp_fusion, + ) + + # Leak-free conditioning (instrument + acquisition mode only). + self.instrument_emb = nn.Embedding(cfg.n_instruments, cfg.d_model) + self.acq_emb = nn.Embedding(cfg.n_acq_modes, cfg.d_model) + + def global_token_hook( # noqa: D102 (documented on the class) + self, + tokens: torch.Tensor, + *, + instrument: torch.Tensor | None = None, + acq_mode: torch.Tensor | None = None, + **_: object, + ) -> torch.Tensor: + g = torch.zeros( + tokens.shape[0], + self.d_model, + device=tokens.device, + dtype=self.instrument_emb.weight.dtype, + ) + if instrument is not None: + g = g + self.instrument_emb(instrument) + if acq_mode is not None: + g = g + self.acq_emb(acq_mode) + return g + + def forward( # type: ignore[override] + self, + tokens: torch.Tensor, + *, + instrument: torch.Tensor | None = None, + acq_mode: torch.Tensor | None = None, + attn_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + encoded = self.hybrid_embedding(tokens) # (B, L, d) + global_token = self.global_token_hook( + tokens, instrument=instrument, acq_mode=acq_mode + ) # (B, d) + encoded = torch.cat([global_token[:, None, :], encoded], dim=1) # (B, 1+L, d) + + # Explicit padding mask from the token ids — robust, and the global + # token at position 0 is never padding. + pad = tokens == self.cfg.pad_token_id # (B, L) + padding_mask = torch.cat( + [ + torch.zeros(tokens.shape[0], 1, dtype=torch.bool, device=tokens.device), + pad, + ], + dim=1, + ) # (B, 1+L) + + encoded = self.positional_encoder(encoded) + latent = self.transformer_encoder( + encoded, mask=attn_mask, src_key_padding_mask=padding_mask + ) + return latent, padding_mask diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/heads/__init__.py b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/__init__.py new file mode 100644 index 000000000..627540428 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/__init__.py @@ -0,0 +1 @@ +"""Task heads — fragment-indexed intensity, CCS/IM, retention time, charge.""" diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/heads/ccs.py b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/ccs.py new file mode 100644 index 000000000..cfbd2f6a2 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/ccs.py @@ -0,0 +1,64 @@ +"""Ion-mobility head — a per-charge physics prior plus a learned correction. + +Predicts inverse reduced ion mobility (1/K0), the raw quantity Sage reports in +its ``ion_mobility`` column — this avoids baking a Mason-Schamp CCS conversion +(with its gas/temperature constants) into the prototype. + +The physics prior ``1/K0 ~ slope[z]*sqrt(m/z) + intercept[z]`` is ported from +the production ``SquareRootProjectionLayer`` (imspy_predictors); slopes and +intercepts are learnable. A Gaussian-NLL ``(mean, std)`` output keeps a +calibrated uncertainty, useful downstream for the pep-centric SNR work. +""" +from __future__ import annotations + +import torch +from torch import nn +from torch.nn import functional as F + +from peptide_property_ng.model.config import ModelConfig + + +class SqrtMzProjection(nn.Module): + """Per-charge physics prior: ``property ~ slope[z]*sqrt(m/z) + intercept[z]``.""" + + def __init__(self, max_charge: int, slope_init: float = 0.03, intercept_init: float = 0.60): + super().__init__() + # Index by charge state directly (row 0 unused). Init at a 1/K0 ballpark. + self.slopes = nn.Parameter(torch.full((max_charge + 1,), float(slope_init))) + self.intercepts = nn.Parameter(torch.full((max_charge + 1,), float(intercept_init))) + + def forward(self, mz: torch.Tensor, charge: torch.Tensor) -> torch.Tensor: + z = charge.clamp(0, self.slopes.shape[0] - 1) + return self.slopes[z] * torch.sqrt(mz.clamp(min=0.0)) + self.intercepts[z] + + +class IonMobilityHead(nn.Module): + """Predict inverse ion mobility (1/K0) as ``(mean, std)``.""" + + def __init__(self, cfg: ModelConfig): + super().__init__() + d = cfg.d_model + self.max_charge = cfg.max_charge + self.physics = SqrtMzProjection(cfg.max_charge) + self.charge_emb = nn.Embedding(cfg.max_charge + 1, d) + self.net = nn.Sequential( + nn.Linear(2 * d, d), + nn.LayerNorm(d), + nn.GELU(), + nn.Linear(d, d // 2), + nn.GELU(), + nn.Linear(d // 2, 2), # correction to the mean, log-std + ) + + def forward( + self, + latent: torch.Tensor, # (B, 1+L, d) + charge: torch.Tensor, # (B,) long + precursor_mz: torch.Tensor, # (B,) float + ) -> tuple[torch.Tensor, torch.Tensor]: + glob = latent[:, 0] # (B, d) + h = torch.cat([glob, self.charge_emb(charge.clamp(0, self.max_charge))], dim=-1) + correction, log_std = self.net(h).unbind(dim=-1) # (B,), (B,) + mean = self.physics(precursor_mz, charge) + correction + std = F.softplus(log_std) + 1e-4 + return mean, std diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/heads/intensity.py b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/intensity.py new file mode 100644 index 000000000..484329a0a --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/intensity.py @@ -0,0 +1,59 @@ +"""Fragment-indexed intensity head. + +Unlike Prosit's fixed 174-value vector (29 positions x 6 channels -> a hard +30-residue cap), this head predicts one intensity per *cleavage site*. A peptide +of length L has L-1 sites; the site representation is built from the two +flanking residue outputs, so the output is ``(batch, L-1, n_ion_channels)`` and +the peptide length is bounded only by the encoder's ``max_seq_len``. + +Charge and collision energy are conditioned here (head-level) via FiLM — they +are deliberately kept out of the shared encoder (see config.py). +""" +from __future__ import annotations + +import torch +from torch import nn + +from peptide_property_ng.model.config import ModelConfig + + +class IntensityHead(nn.Module): + """Predict per-cleavage-site fragment intensities in ``[0, 1]``.""" + + def __init__(self, cfg: ModelConfig): + super().__init__() + d = cfg.d_model + self.charge_emb = nn.Embedding(cfg.max_charge + 1, d) + self.ce_proj = nn.Linear(1, d) + + self.site_in = nn.Linear(2 * d, d) + # FiLM modulation from (charge, CE); zero-init -> starts as identity. + self.film = nn.Linear(d, 2 * d) + nn.init.zeros_(self.film.weight) + nn.init.zeros_(self.film.bias) + + self.mlp = nn.Sequential( + nn.LayerNorm(d), + nn.GELU(), + nn.Linear(d, d), + nn.GELU(), + nn.Linear(d, cfg.n_ion_channels), + ) + self.max_charge = cfg.max_charge + + def forward( + self, + latent: torch.Tensor, # (B, 1+L, d) — index 0 is the global token + charge: torch.Tensor, # (B,) long + collision_energy: torch.Tensor, # (B,) float (0.0 when unknown) + ) -> torch.Tensor: + residues = latent[:, 1:] # (B, L, d) + sites = torch.cat([residues[:, :-1], residues[:, 1:]], dim=-1) # (B, L-1, 2d) + h = self.site_in(sites) # (B, L-1, d) + + cond = self.charge_emb(charge.clamp(0, self.max_charge)) + cond = cond + self.ce_proj(collision_energy.unsqueeze(-1).float()) # (B, d) + scale, shift = self.film(cond).chunk(2, dim=-1) # (B, d) each + h = h * (1.0 + scale[:, None, :]) + shift[:, None, :] + + return torch.sigmoid(self.mlp(h)) # (B, L-1, n_ion_channels) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/heads/scalar.py b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/scalar.py new file mode 100644 index 000000000..018392f36 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/scalar.py @@ -0,0 +1,47 @@ +"""Scalar heads — retention time and precursor charge. + +Both read only the encoder's global token (index 0). That token aggregates the +peptide sequence plus instrument / acquisition-mode conditioning — but *not* +charge, m/z or collision energy (those are kept out of the encoder). So the +charge head never sees its own label: no leakage, one encoder pass. +""" +from __future__ import annotations + +import torch +from torch import nn + +from peptide_property_ng.model.config import ModelConfig + + +def _mlp(d_in: int, d_hidden: int, d_out: int) -> nn.Sequential: + return nn.Sequential( + nn.Linear(d_in, d_hidden), + nn.LayerNorm(d_hidden), + nn.GELU(), + nn.Linear(d_hidden, d_hidden // 2), + nn.GELU(), + nn.Linear(d_hidden // 2, d_out), + ) + + +class RetentionTimeHead(nn.Module): + """Predict (per-dataset min-max normalised) retention time.""" + + def __init__(self, cfg: ModelConfig): + super().__init__() + self.net = _mlp(cfg.d_model, cfg.d_model, 1) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + return self.net(latent[:, 0]).squeeze(-1) # (B,) + + +class ChargeHead(nn.Module): + """Predict precursor-charge logits — class ``j`` corresponds to charge ``j + 1``.""" + + def __init__(self, cfg: ModelConfig): + super().__init__() + self.max_charge = cfg.max_charge + self.net = _mlp(cfg.d_model, cfg.d_model, cfg.max_charge) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + return self.net(latent[:, 0]) # (B, max_charge) logits diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py b/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py new file mode 100644 index 000000000..382da6360 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py @@ -0,0 +1,75 @@ +"""Unified multi-task peptide property model — one encoder, four heads.""" +from __future__ import annotations + +import torch +from torch import nn + +from peptide_property_ng.model.config import SMALL, ModelConfig +from peptide_property_ng.model.encoder import PeptidePropertyEncoder +from peptide_property_ng.model.heads.ccs import IonMobilityHead +from peptide_property_ng.model.heads.intensity import IntensityHead +from peptide_property_ng.model.heads.scalar import ChargeHead, RetentionTimeHead +from peptide_property_ng.modifications.composition import CompositionTable + + +class UnifiedPeptidePropertyModel(nn.Module): + """Shared Depthcharge-based encoder driving all property heads. + + ``forward(batch)`` returns a dict over the active tasks: + - ``intensity`` : ``(B, L-1, n_ion_channels)`` in ``[0, 1]`` + - ``ccs`` : ``(mean (B,), std (B,))`` — inverse ion mobility, 1/K0 + - ``rt`` : ``(B,)`` — normalised retention time + - ``charge`` : ``(B, max_charge)`` logits (class j -> charge j+1) + + The input ``batch`` dict carries: ``tokens`` (B, L) long, ``charge`` (B,) + long, ``precursor_mz`` (B,) float, ``collision_energy`` (B,) float, and + optionally ``instrument`` / ``acq_mode`` (B,) long. + """ + + HEAD_TYPES: dict[str, type[nn.Module]] = { + "intensity": IntensityHead, + "ccs": IonMobilityHead, + "rt": RetentionTimeHead, + "charge": ChargeHead, + } + + def __init__( + self, + cfg: ModelConfig = SMALL, + composition_table: CompositionTable | None = None, + ): + super().__init__() + self.cfg = cfg + if composition_table is None: + composition_table = CompositionTable.load() + self.encoder = PeptidePropertyEncoder(cfg, composition_table) + self.heads = nn.ModuleDict( + {t: self.HEAD_TYPES[t](cfg) for t in cfg.tasks} + ) + + def forward( + self, + batch: dict[str, torch.Tensor], + tasks: list[str] | None = None, + ) -> dict[str, object]: + active = list(self.heads) if tasks is None else [t for t in tasks if t in self.heads] + latent, _ = self.encoder( + batch["tokens"], + instrument=batch.get("instrument"), + acq_mode=batch.get("acq_mode"), + ) + out: dict[str, object] = {} + if "intensity" in active: + out["intensity"] = self.heads["intensity"]( + latent, batch["charge"], batch["collision_energy"] + ) + if "ccs" in active: + out["ccs"] = self.heads["ccs"](latent, batch["charge"], batch["precursor_mz"]) + if "rt" in active: + out["rt"] = self.heads["rt"](latent) + if "charge" in active: + out["charge"] = self.heads["charge"](latent) + return out + + def num_parameters(self) -> int: + return sum(p.numel() for p in self.parameters() if p.requires_grad) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/modifications/__init__.py b/packages/peptide-property-ng/src/peptide_property_ng/modifications/__init__.py new file mode 100644 index 000000000..8892dff88 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/modifications/__init__.py @@ -0,0 +1 @@ +"""UNIMOD modification chemistry — atomic-composition features for the hybrid encoding.""" diff --git a/packages/peptide-property-ng/src/peptide_property_ng/modifications/build_table.py b/packages/peptide-property-ng/src/peptide_property_ng/modifications/build_table.py new file mode 100644 index 000000000..9509ab81e --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/modifications/build_table.py @@ -0,0 +1,95 @@ +"""Build the UNIMOD atomic-composition lookup table for the hybrid encoding. + +Maps every ProformaTokenizer vocabulary token to a signed element-count vector +and writes ``data/mod_composition_table.npz``: + - ``counts`` (vocab_size, n_elements) int16 — signed element deltas + - ``elements`` (n_elements,) str — element order (== composition.ELEMENTS) + - ``tokens`` (vocab_size,) str — id-indexed token strings (provenance) + +Run once; re-run when UNIMOD or the tokenizer vocabulary changes — no model +vocab re-mint is needed, only a fresh table: + + python -m peptide_property_ng.modifications.build_table +""" +from __future__ import annotations + +import re +from pathlib import Path + +import numpy as np + +from peptide_property_ng.modifications.composition import DATA_DIR, DEFAULT_TABLE, ELEMENTS + +# Matches every UNIMOD reference inside a token, e.g. "C[UNIMOD:4]" -> "4". +_UNIMOD_RE = re.compile(r"\[UNIMOD:(\d+)\]") + + +def build(out_path: Path = DEFAULT_TABLE) -> dict: + """Build the composition table and write it to ``out_path``. Returns a summary.""" + from imspy_predictors.utilities.tokenizers import ProformaTokenizer + from sagepy.core.unimod import modification_atomic_composition + + tok = ProformaTokenizer.with_defaults() + vocab = list(tok.get_vocab()) # id-indexed list of token strings + comp = modification_atomic_composition() # {'[UNIMOD:N]': {element: count}} + + # Every element sagepy reports must be in our canonical alphabet. + seen = {e for c in comp.values() for e in c} + unknown = seen - set(ELEMENTS) + if unknown: + raise RuntimeError( + f"UNIMOD composition uses elements absent from ELEMENTS: {sorted(unknown)} " + "— add them to composition.ELEMENTS and rebuild." + ) + + elem_index = {e: i for i, e in enumerate(ELEMENTS)} + counts = np.zeros((len(vocab), len(ELEMENTS)), dtype=np.int16) + + n_mod_tokens = n_resolved = 0 + unresolved: set[str] = set() + for tid, token in enumerate(vocab): + ids = _UNIMOD_RE.findall(token) + if not ids: + continue # bare residue or special token -> zero delta + n_mod_tokens += 1 + resolved_any = False + for uid in ids: # sum all mods on the token (terminal + residue, etc.) + key = f"[UNIMOD:{uid}]" + mod = comp.get(key) + if mod is None: + unresolved.add(key) + continue + resolved_any = True + for elem, c in mod.items(): + counts[tid, elem_index[elem]] += c + n_resolved += int(resolved_any) + + DATA_DIR.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + out_path, + counts=counts, + elements=np.array(ELEMENTS), + tokens=np.array(vocab), + ) + + summary = { + "vocab_size": len(vocab), + "mod_tokens": n_mod_tokens, + "resolved": n_resolved, + "unresolved": len(unresolved), + "n_elements": len(ELEMENTS), + "nonzero_rows": int((counts != 0).any(axis=1).sum()), + "path": str(out_path), + } + return summary + + +if __name__ == "__main__": + s = build() + print(f"vocab size : {s['vocab_size']}") + print(f"modification tokens : {s['mod_tokens']}") + print(f" with composition : {s['resolved']}") + print(f" unresolved UNIMOD : {s['unresolved']} (id absent from sagepy table)") + print(f"elements : {s['n_elements']}") + print(f"nonzero rows : {s['nonzero_rows']}") + print(f"written : {s['path']}") diff --git a/packages/peptide-property-ng/src/peptide_property_ng/modifications/composition.py b/packages/peptide-property-ng/src/peptide_property_ng/modifications/composition.py new file mode 100644 index 000000000..4b39bad0f --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/modifications/composition.py @@ -0,0 +1,86 @@ +"""UNIMOD atomic-composition features for the hybrid modification encoding. + +Each vocabulary token gets a fixed vector of element-count *deltas* relative to +the bare residue: a modified-residue token (e.g. ``C[UNIMOD:4]``) carries the +modification's atomic composition; a bare residue or a special token carries +zeros. Deltas are *signed* — UNIMOD neutral-loss modifications have negative +counts. + +Source: ``sagepy.core.unimod.modification_atomic_composition()`` — the same +UNIMOD view the Sage search engine uses, so it is consistent with the training +labels produced by Sage. +""" +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch + +# Canonical element order. Verified against sagepy's UNIMOD composition tables +# (31 elements incl. stable isotopes and metals); ``build_table.py`` asserts the +# live data is a subset of this set and fails loudly if UNIMOD ever adds one. +ELEMENTS: tuple[str, ...] = ( + "C", "H", "N", "O", "S", "P", "Se", # common organic + "2H", "13C", "15N", "18O", # stable isotopes (labels) + "Ag", "Al", "As", "B", "Br", "Ca", "Cl", "Cu", # metals / halogens / ... + "F", "Fe", "Hg", "I", "K", "Li", "Mg", "Mo", + "Na", "Ni", "Si", "Zn", +) +N_ELEMENTS = len(ELEMENTS) + +DATA_DIR = Path(__file__).parent / "data" +DEFAULT_TABLE = DATA_DIR / "mod_composition_table.npz" + + +def signed_log1p(x): + """Sign-preserving log compression: ``sign(x) * log1p(|x|)``. + + Plain ``log1p`` is invalid for the ~430 UNIMOD modifications with negative + (neutral-loss) atom counts; this keeps the sign and compresses magnitude so + a 120-atom glycan does not dwarf a single-atom methylation. + """ + if isinstance(x, torch.Tensor): + return torch.sign(x) * torch.log1p(torch.abs(x)) + x = np.asarray(x, dtype=np.float32) + return np.sign(x) * np.log1p(np.abs(x)) + + +class CompositionTable: + """Per-token atomic-composition delta vectors, indexed by tokenizer id.""" + + def __init__(self, counts: np.ndarray, elements: tuple[str, ...]): + if counts.ndim != 2 or counts.shape[1] != len(elements): + raise ValueError( + f"counts shape {counts.shape} inconsistent with {len(elements)} elements" + ) + self.counts = counts.astype(np.float32) # (vocab_size, n_elements), signed + self.elements = tuple(elements) + + @classmethod + def load(cls, path: str | Path = DEFAULT_TABLE) -> "CompositionTable": + path = Path(path) + if not path.exists(): + raise FileNotFoundError( + f"composition table not found at {path}; run " + "`python -m peptide_property_ng.modifications.build_table` first." + ) + with np.load(path, allow_pickle=False) as z: + return cls(z["counts"], tuple(str(e) for e in z["elements"])) + + @property + def vocab_size(self) -> int: + return self.counts.shape[0] + + @property + def n_elements(self) -> int: + return self.counts.shape[1] + + def as_tensor(self, normalize: bool = True) -> torch.Tensor: + """Return the ``(vocab_size, n_elements)`` table as a float tensor. + + With ``normalize`` (default) the signed-log1p compression is applied — + this is the form the model's ``CompositionEncoder`` consumes. + """ + t = torch.from_numpy(self.counts).float() + return signed_log1p(t) if normalize else t diff --git a/packages/peptide-property-ng/src/peptide_property_ng/modifications/data/.gitkeep b/packages/peptide-property-ng/src/peptide_property_ng/modifications/data/.gitkeep new file mode 100644 index 000000000..1758240c3 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/modifications/data/.gitkeep @@ -0,0 +1 @@ +# Generated artifacts (mod_composition_table.npz) land here — see build_table.py diff --git a/packages/peptide-property-ng/src/peptide_property_ng/modifications/data/mod_composition_table.npz b/packages/peptide-property-ng/src/peptide_property_ng/modifications/data/mod_composition_table.npz new file mode 100644 index 0000000000000000000000000000000000000000..6f31aea443c5533cd156ad39ff21cb9db1d8a33f GIT binary patch literal 15340 zcmb7L2|SeD_it4yBvDygiAV_9Q=!HB+R|9Ej_ga;VHB0@*>@9?o$R|QVX_W|F_;-U zWoGO%7|Z{e_TIPpeg66M%<|lO?m6e4^F7ObM*Y^#T>{&-ZQBd{Kd}vGPHVmO_20g2 zhqsy9yV%`#K5ysXv3=XYZC4m}0Hd~Ti#-1MWgEfZUL-rhvE$>%j}j!=SC?22JBFHk zk6$0!M(aCz!q@+LV1#-1iGkwlK|RJdqOP?{3ixx3FpzWyM{pysajw)0k7QmZN2BwT zIk(je&I-H~E!sF(bVA)du~%D9TTd@uZ)>&JzMfoWFFGE(P_wKiyD?e=E5zyP!6|R8 zXJU2w;N%ZTpZQe;%zG*m=Cg2bW!BzneN0fccWbkWL3R_JxklcPD!|rz+?!bI%@oXE zhwi?jcHkc0PRnRPiu(DyzIa<9)+ISlDE#X`_o`X3i&2d?n(s;$jgzD(# zU68BfR1Q~7+i@rS<~Avf{dotZMLey`=Du`j@8f=Z(Ky}KpTtF3Lrea|3Qr_ zW{^$>mX~uZ-&o@u$g&L84oE3DU+JHxl5ZW5-^^GUc1CPKq6TmC(E`e=Y^2)i@kaN> znX0UWC!Tvy<5lgbN3X*8;2Gx{SM;We2-nUYD!VxE-Hy)7+#6}wd7eWtG7gicR&qlR zRE0`og>N4VnIL9eDIB4wBXrCT)4G7ecS_x;?%8Wr6O)ejig|7WJ9D&#$N7;93@KEK ziYq$TjQJjwCUc6Lh|r!& zJU651A+%>`t82Eq*Z%sPi9VTu1?Yg7sdD>ZHfoX<*xhBYPQ`4nU7^A5K3?K zS9*k*!i_~R+$U#iks3A`$~|2EJ$FJ2RE!(vM7k}w*{)t|Vk^&Px}TJP{%pSyk-XzpZzP+c^KRb~s2|qLeSlKuvR2 z)TYs^R>Qzll}1VcaWGG-7FltlrU~62OAZO?jEE{FZC>{x=OuMT)0I^xJW6>Y>+?)_ zj9_!WaR-|bsyraZK7&l_`YIs~AE8o>obP@3u}n`w(b6jDcI?7@Gpa{m^zcxoYQ*Y* zF9`y>J5|o3bBSqlz|c#=i+h?6$9(VQp0vrX+N;Xt^eAA0fWd zC^Wk|*ZU@)(+g_HgyR@RIjGp034lH6Y$e#-U-~6UXj9w`n#y zG#VZ&jN9pJYsFxzE>6~0bhznB6ECawRED3+-iyJkuPiBqBwkX~%NCUh=qt5pHVSD1#x&jY(-qB64Dr@YzJ0p{b%9ul(ST=l`)~ zs$=Q&9eia2PA{&ibvsWUI#r}@uxxoYs_wK>@z&!~*^Ew{hDD(e`OMjqtvbpjiI7S5 zKqm zYm!wLeugY!d_6ns;iH69{c#Ex?<{4f=@{dU!Cf_Kdi{xyn#R$MmnxH=n%bMyJ(=fH zFER6z)UZ`|#&37=#Xi5Eua>mFSSh2m#)7tHdxSrh%UlXl{$NK~cqg@9^H!}CRgty# zgFp^37*Uv7EZ&j@pJS@fj9&{JY$V)%#a#2K3B~cg$KS=ZVjkc6!h>6gJyyFdqhSE5 z-^qxN@#baSj%a#TCNRa3zbP!BdD3B*{!`b+MYs%PeQv@B7Ou(S7y8DU?s;{yf`_X% zD8V6m7yVnZ^n~Z#mBW&D(>e;`87?}g;BYU=2hpEGWHZ_n)%Lw9Ne6n zJ_$|~apjJ+Ouw&fXPmhMQrb`UP|odT3hNFm_8W7OhkJ73D?*kSKC(F0h~3bV3E@q= zcLR(j^5oQ+jSc)0oPK`rNr4qxfwiX7=@M-O-CJ?%iM(b)s-|VtvP-|cx{mdh{R0qG zw|{Hhvr;>VWoHVaGO9!Q=Kdm4bGT%x)m72F1QdOZzQPRcAYGR-KRC9n6t^7&H$OOb zNmB!l!!5kxzHghCTlCL-wfc72-UVJvi-~(oN81&1w}i#D%xcMDb`sab{M>mLtv(kS z#S-a3Va*jtRiptTFNNLQ^vK{OZ|jyf`T#~Xm8IX(c4B;+YSG!jT$2mrE^XtcQ74Vc z$I_wTG}f^vnG-4on&3AQV&`~9uD7VyESsGeFo~LVKdsXu)t~*eUF_pbHof&ip7zbx z1K1*|{f9p|DLX*&t)VaH2lwaDI6LEtj~%iDDdJ`W^7pKMx-bBVyoZ4^j0GJWV!K<^ z*+qlD8pcl(GXHukE=qwH#O1K3$!n{7jbp%UcMsuVOe_YG`1phVD)wCH@(TXtqSKuQ z)`uf;`dzU7$0?=yCAf#i+Qd&t&Eu4`+Eh8@L;TdZ0moS8<#I@1*9Pv~tyDaH3o!r% zK@?wX81O}SXdQ8zEuzozDOBi_gGKnc$1_G^Sk}#-m`kMDT&&MfsSEBeGH~!WAgY8G zU>;8fXN`M(WcEtXXz89gE5&*@c4kABW$NJxR+HGlt;|wndW85zskrt}>!!A3Wu=Qj zgZ7^4D(3TdhoJ(csYw+!ro)rTRS(aGRg5bQ*n=;VWL6C=yn{J?c(-19FCTgb1wk_s zT$#)=^=vJIH|p1eTA|ycHd5F9q2%1-gR^v9t0O@yO&7eBw3{If z=H;aEvqCp*A#B;K>>6z;5V_q-L4lu=FSPBHp{TL4d8-Nc99voqIliLFHuvd5e)SV2 zt4ol5?6VSO&TMmTQfzWQg}E;~@f!?W0lP8Yrz$5pfen?3CrBj~Z;V9_rf(oxXVR#5q1TFg(Amxkb@C zA$CG0Is|eJr-?Rxc{~AXpk~&oIBdPYI?N!_{w!t3yZQn=_s{~{!b2!o@?AZzkHWip zZl6@wNe>6}g_!iZNdYA|CB`erB3%HU51u|6Lr4f55D+g#oiu>Gl9pG%h=xrEeO4V8g42g*a&4(3*n@JGwJ}Zvx{8|;MXxV7+n z+eA&H8>Dc=do!HG+DZaqdh9q8o6rH{n9|{S;NvM$L@~&rZKI~!&{;)+&?A*25U$T? zqvEXxysv6Orhk_ z_N{`-Vy&BdKCm3v#AFrJnR|D7iw`&%X*#OVi7ZLFr>{LJmyQ^Gj_2{HMU(1T#s3?L&boGpRNu# zvEj?y#<_S^=ooKgm9HF75bdR~M-#lEi4h?M%3zy4*T#@_6G%S+A0<%A>#>c3b*ApA z`tsUsh01U-F^h`Yb8w!Ey~|0xEKwXf+&u$}vx=h2d8FH?mE_CGtU1}U_u2+uy@YCS z#&<+Xo@b^yT zk|IAkYoRW&e=ua4S?pHgUjx`%evu6Co+n9^?`kE9GGjK>+O$n-)#AZ+0TCges;v1W zO;Sn_9kZsd@Wb&}mD-VKvdWXX0ev$%qo)T(t@#SFjT=M#@XRtt8ebJ8>&VzQ42*h* z7VEa_jG3i-6N?>PFU3ATHJQ6J#~SOTRc_7KI_aPT3Y{5{f2W8@e0y&(^zrq!?%7Q> zH1EQ$cdiB?n!!3g!9(6eXZ&e?%I@e8Xhj8k@|gl_p5Yk(Jdca#W8WsUh7KDpmhJA1 zSzI`SLCoFRXWZFznyl%%tB`(5MYB5w78|R@$J$hv(4uJj@>O}#b~I-l zu#agQ)LsC`27nxrF;W$=evdX47y~&3v>!1@(oGotpWzypZ7sG;A zcDF@cMhuqiPCld7W?lU<%^RvgWIQ!l+|9T+FKck)NoFBESRQ?3qOH#fFt2HMh#_4o zR1teY?8&VK79nh>BBX=D-qBih+Ma3kT5KuGq!M+>UF;Cw;1qdZmH@Ivs+uWvj38jN z9c_rA1rq1<#lfk5i8+?|Hm;gH9LHN+U5<6oZHxA4YY}O1yuGE&`7JxUCn-A-ln8US zf`|CFHH=XBJo)?^R|CjCn^o;qRv{DFtGc}C+~+p*1=#>Yot&~*c~IY7*DjX(wSc?Z zytS4z5y}%tNe?KNa4$?v@KnH}V$yBs#m5t?cc$m2g&mbdcr~g6Rws!7+~8I%>XJ%R zNdP<*fY%ycu7kOL;Mg=9`XeqphuRvh@Q1O*KtjR7K6mwcj@nIoa-s<`J7A_WQ7P#* zTXq*2B+WF&JZ3u;mD&-Mk2m|Y=qp@8%kCO1dfnC6yjnPd9>Z}9lO%?EVi4Ch$wo8c z=|J6a2%I#icvuX(yBu~@@)p7AQWBJU+l!j`uZ z03C*lc^*-O`1Z4dMvg<-Y@h;*d>?T`(C8@>Kxz#VGX~xyj60Z%5RLx8!j0o93>+z= z=`vijF^Te4a`R(Jf}|X{XC>iYwyhL(sgf|qJ{`RjlfJsiF3oeKx2Hzu)~=2;SFx0; zP46oxX|AK@&U3+7D>K~t=o|d-U3$z$v9OzLeKES@nkNPO5EQ=XJs~L86~=BSDFDO+4jche*9lumbBnTXn3y7S=uEVoJ%$d;4} zVzw`}xzbp^bDFn20hM>>MRs;o7+Tr3@ zi9O!gSlgi_15&+Zr>oD>?jjF~TdySA2ZJ1^6P=|5%T0^lLYz6C$(hDyWc7;B>RzUe zt2jWXMY^azc;aqEs|W0{Gu@8RL5gX*cd9G`HEyJhGu(w1iBHiTQ6hed6YWqe7(Jcv zaz&g%Tg_UkjO;3Rmg?R}HQ41n@k%j-Cx|kM*;mpybuT|#lr!wQuC|a-_q6VEs9t5| ztow@ggwj|yj(KoDJZs9+*Ii!f+VkbTVtXn&a>HfrEfen?mukV9H4mi3VSH1pENmOQ zV#rGu`@2}~7AN*Rov_g=M|SZoWn}ooC8mfKy>wZ={{VKN+A}3+_yZ*>`BoQGZ*y0| z?T(hb=FM=kOjc!GNV#xo`jC(qF*#Aoj(G~qa%0l8pt<@@^ ztpz~gj$+Ovh8YPv53L|2P!$CvFpwiqSy1ZY@;ZMp>)@$>SngGuGrzkq>zzQqMr3rA zQ-Dj04sZ1N^A&1`Nm%iiyh*n5Ef_h)(mxa!Ogy9Fq(E;|ijwT%?K4%*sT%q8(L)t~ zG2m{iYEhQ1?Rz`Vm#?9TCH;%M>ZH6Rbpv^GK|Pfq3H9=lZq$X)JN&|}q`+m%ra_Sy zx<>6We&Qnvtk94vKDRdrb+%qb%g7*F>g2WM-q(uOSPxsJ+GgXd!b7^%g!4~J-W#`V! z)c}bGN+;csyb9@3-aN`$hKt^PMRT?V8K4^tT=CHc2MTlXG)W~>f)7S#J2n~&xtfOK z51kz>x!4tJhYQCYT=kuOSEw!43Pn0GcLTMNwB}~7IdU&YiThk+lb?6;X|4@QCO_qX zj9%sAkTsC+mj3YG1>TbI?nnisUECdqB@kj=Kg(V)SR7h>$s+tVQ1vhA|I;C(Mt7j- z-pyTDr`0;S*2&>bYb{qb3z~~%cvhyqvZM71&>I2;_aeidE16I5-(gt{MP`J~cq}}n zfu`Zv)0bToqLiXco#TWOv+FC)`}gBBlC$9cMTsopHVW6$Und!aG`lD5cE&kfdNY2R zlA9?dgB2*d;yUg=LQ|7e)?NgTGEiyU8c|SE<6gN}aFhQZo*VQ5q@xAda^YLuP#-Qc zpR-#T3r4%q{>%b=4Ht#EC9#|!zhWl7&^B{jCQyl4y#2JwdH=$;u;sRHRzAaHJL-JY z^5AYIq!f1PPle@31nmgP6AVQLckQS?Yd^i7v~B6)1=SXhAdxqvwbt|2?me5=SWbv| z-d{Q}^g!1?G{kF7SS6B|LzqnnEE$2Dn&YC- ztxSevjW*Z2F77AKg(xepzLF?BZYa+ETyJ%?rm}zlXApmajkL*a04-7 zI$PfCg*$Mo!A{??6;{Qrn0(*zsBTI}>PE$o#OCI~8hsfp)S*(u5&m;~LNe|=Pdnpf zD}QP#SE0+Xwo_hrj%*@vt%rkof2XQ>|Im5&Yn>n2#thTV%x1Nzk>YnD7x1GqokWi`MEux!87TlW>ca+}j-<1(qJc<#d(cuqNM>-KIgeuhn0AmAl~U zmd|=!^jz_j5J$aK4#w7>Cx2&fj7y1IcoFW|%fydjRp&#-RC_A!)J?mQvnAgol{FO^ zw9W5BV0RoIyaprlLukp4yIaO?pL?EJBd@Rbn7f`^RMs$t9?n}^HZ4`olrR^>nK^F~zR^LHX9U)R}AFB2JxWhP#X4VBlG`hDUi z8{!jc&ueSjg=OC~S(&-AJ;*`E3L&;E={rb2<}RE~j0t$FCjrYBJWAQoWO$~Zy%UWY z!Y)#}OS(!BQW-j#Z+y@^hA-yCX}B-DMwy>8IX6aIWvMN60cqCeB;9p@_~Mb#+(=(l z3?aS4NC?ZV<&DHpBj?3bf@J?zy{WvJr47fSh>lR61U zLs=c3W2eC5JKw+)gBX`92D;nDQ7lV!ZUyC|yQ^%(Kqc4k*iJYo>P&QK^8wFb+B--; zKNx(z?vclQ?;GsWbLh@>^Q$`00&7Z9A}7yoTMtA|v+fwzjJk{Up_1OY9P!A8H9w8^ zJCCoGB4bBZA|5D6)uqRkX!379;fYl6?j*vyBb%a#uic~1oX5h=qg~0N@$O4oxFRXO z`zp3tI`%vjo?XWzd6f1h@?Imj6@$VM#w%oci61_!4*?U;Dft$ zLq4cEXCY7Z+_r2GEn*Rf@H}J&>1DgGLQt>Qr`I&qYddv!fK)VZ197oDFoEywRfb7l ziM~8kp?l`!hZx89$nmanO**2`v89SAOn4}O1syrk*0V7>q62;&`D1ZoPW{6(cSLVT zG5cb~ZrQVw>rYw0(o0NL&J(0I*Nu}h`?XD;czr0k4bNDCkUoXba)`k7uB{2AjmB3l z-cIy962FtYBcMdXr_DEoe~p)K5Zzn5PvcC?MDT!QsvZAqZlKbq*@dj_r7Yyqq>~nw zD1MdF@bD?AVuL>=@wgh52QdcjcCQiN4_194-EZD^!omni(@g}0_u(l?m(P#i2 zu3tcRuUet0LfBHB=k7W)eQbxvX4?|#tx=^oWxnox@jO<(3bNjqmP#`b1u}LEYDeeF zf%BcDswryIhQhal@{+Lk36C)xA-o9aO*_U#h&J(1baD-eC{baZ?i?I+XvFH)tHsEG z6G2(-d(K$bwUb>CFUMAOBWyq%k|(ImRZG3hN>ssw5IL*#XyX)T+M9e_Hj1%c?RC)=?BBT zw+YIR>DWomVazvFKvCL02hoiyACt398_cgbRw*YDUMmp%pFVeNk0%@KGSW*Z_Y+|{{!~;I%4xRG`d@z?a6&YX6!q(`X$<2mTu(&UvBBOSQ%h}?H zfR$BMS)Il}jkANz`7o=*_c7dM+~rwy$H|##Hg=-8ll;9sIm-^BqDPGw?J7zM1y8Ic zJWsaahVz3*q7rR^kS6D>k;x}q5_GJ@a1(pC2gddNVLI_*m!tPKOUgQ(z6y@qVZ+0qLjzL7ER(_RPLdJft=KdGP^URlOZnPH> zmbPoy7b%+G!~P(Z0vUk_hs zmSiI=c0>hqR29bTsfXBeUe1t>g3EUvqR|&1*yi$vK*^Kahw9R+C5`6;!^ViqQd1KL z9O+YHOLj@|)*ieK#`N~ny*$ZRPTs+8Y8vuwV$Kf269a^sx?@by9X^|6VcThHuQH6Z zAxy)_To{L{d0XFyw#0fVvFwj%Gk12xKdTX&7U3oDj=YN37|qiEU=QQ6ycJ(FJmz#2 zz29dI>tsWZ$$P^4Bt_w$ur-yp@@UKQRNdom&uQ*bQxJ`B=O;O5u2e$j$+XQyb}AvP z6eTw#-`g63h@*w=Dx6meG}?}P|DZqdEHux56>ckY<-J6!9-cT_koRrFj^zFi^?>ds-e>0qakfrv1N zjpYzIh?-71wa`x^cepCo1jd4{(=48D`;-LsnYFwr0TOz(rzE@;b3=@v<+!4l*F4Q= zc=kG94zsj%u241G(h|d_g>|$gBeSwzqh3hky1i(MbOLiz@2(5XUB1blAF#zqwLd%oFD zQsnMt*%L=xw6EcEX+1l&h0a+t$FSARH7T@#x8CLI0q1{`!J=J=C|<+~*bxl0gz zq?SPL!h_={agVG9_!R{3OpSGKqmrLH+Vb}G9G&o9N`s2qd91M2U6OjI2=ASpPzM#= zAVMbY-@l#f3d2D@4sV)(ILV+_@8Z@ggzPnBs!=4lvDQ}3^$y1v+ox~bjB3kG`B&H! zJh0X2_#<9tx9q}w`m0H}1-O$TF%-kGaOqSx$8&jiEe4*xafN$r8qd77vAQI^XLT?7 zPHqn1S5Btoc(@eT916yj>G>8X>OpUAR%kR3fN)u!_)_?cXKN@;{{;VXe(5E6t0Fs# zkszCG?Nw18vfnnc>pnjP15Jx>xi7oTh2Fm zPY-DxBE^9rYVM>Yr%~uE4ZeSW|AZ|XT0kQ{EG6~0f_GiSukO~g-a}N>ekp(GY-s*4 z!@jj!bD{5x<$e{|XWPQyH-V|oZ``-HHn;o!{Koa6FY;bE_C?+d3@-ExO7HUgXiVl@ zvV80Hdl0MU&`R3V4-eWbATM8L34F?U)1^=!cQ-E* zab*cJ|JDt`GaCwrx_QlF=9AqJduBryG4tIuV~b-IMru`?8};>&eiZB^8H{@XSy-%- zfRe#pt#vZcnoJl&B?FE_5HU-f(IlF#$Qp6Mw>G_Jt=)*5Gn985+p6T0h=R{3C!Bl& zhQjK%lU-&Zq4~?s!?N4#B-SQ#mlwB_yXR|A8JR=GWhT1~ecWc?Y^_ggzVn7Bz3fPz zA%1FDlX%Szo+2}W-BIBardB0hDQItBkAh3$XCm!iw@}EBK#+wu&Co#HR3?L=y$>?1 zaI?y5V};G37$)}yMQQOv)eOeMMrSsb-ppVO_P`u4@LCQhH|m;D@6}_(HYDe!7o}`- zYQ<9yfubBRnzesM0gZ3cR~dQjR3$CW&M)^_=R>eAbzA6p%GA^rB7b(K#>i35^WAwO z!pITFR*&^ep7m)`t1PTh^rAR<(GTr+;yO+~;Zztd6%O$m0*%XFA)1aC(92F3(l6pC zE+9!$qqICJ`P$N;Tz6OR$vCG`J@Vm(05HX6lNTj|URow0V}JnyiqfDY_j0*C%9rNz z;$s{}(W8we-rTEEl(`GY1T-+h+dU3Uk=b~}gd~9LHP*N1c zISy=3zet@}CZjAdkMA%}sgr0g@$RIKpc3+$<0X?WAhDd2RPIxF3@iM9Y_ReQA)q5C`o!^9@s+6kY&wBVGJHDFdZ__=S8UMi0Z&>-kx1hhv z7q)*hs@A4mvmk!32RQ?Kd)V$17|i zbM0!<3hzd$^&@xg4XP$LshaemsvQ^%e$PchM$uM~{*nG>fC1nsr z-hoP2gi1G_N>>@6`;B6kOzyY0|0WJjqUHJRvY%JvPjdK=GkmwIE)oC# zdchvMk01bm0Hy|jUpY(-04wU7pB49S&m?X_%YGOiUgX2y zB`xw~r!2yGq{v&Lw~7jZ!lCv(f-v$uR&B`&Xs?3TF+YzehPLqa)IJzeT{G?orCz zsNo3<1fm}R3Lt9%M0n)Sgj1cb|4D99x4&U8tou6*|EkIU*^k^w_a6~I{tv%|ZRZ|p z10ogse0USEP7;FMp{3|UHr z2RUceUoGcLSose5Yv^Bt(0A!y!S?@++iKUq{~+=gcizfN2HbfqFPRGRKP>9s=nof9 zbP)9SAUVbT(9mc9B!Bl${94LijU27^r_Ij0`?t+5IWzw`^9Z139zf&-0)dUt=x01* zqT-nW+3c^=t$%^d9}30e#sJ^?^7%JGK8fY`akc?UM+LxY8r83n@%Kur^0D|FFR54t zgv#$`MOB*7*HG}kG?t?OCyj{={f+gM)o<7jAlbLL`8jd>l6d?N?C&`ICm2IM=K;Ly zPZqj#=TBaVp=!*@`46eZR|Rg|5gYlM{{mt0OY-{z!8RmCs97LH_HTFsIrE#nrtU4C z-1j&4@MQUHSzi;df8}JhQ-e9I2xlOYT)vI*p=BLHZK&HUk30O5Llf z>q(^0b39O0h|*B08}PmJfK)-!iUff}_v@_`6$m1l3Hh?_lpXb!yBPJBbRu=u_giDF z0pJfq;Q>u>wB>^?(_}OQH0strpfUl@=;bq+yYG0;&=9kKs&9}uJKVbiK z>kD6%_h$b)>nE1K3i_SkPq)4ZY6?qKskCU&lq2gcJc?W$mc5nYu&`c_jNnNcGMX(J zr{#h9P?X&p0$^^o1X{N#k|wo$Y5)-LH&TFv=|5c|3sJvW{4-OLaF{TVK*Q}{LDc4LV6rwGITN2Wr3!@OuHOhpsMnnWn7~q*n58n|reF2ZMP2d+pn>Gm zfC5wx@j!3#1Te{EL{bZx8`WTDU=UR-DXLf;Qe>-l)bFVU&S^KKqMF1<4r-kP(ANa$ zukijvzY-DyZCK`)`E$?(aB&KQpWx;BY1U`1ek|DWiMlMc|IboZ73Eb?b52pK=s;ENXaDFBq1eDho&;1U2E{LfDP*w%vH|5jksDC<*6vGHmB zxwQKgEU%D#{}4YaA+GMNHHQY8H;YG(@P>g(iA(h%ZapeZElx_6JP=U#0Trjrv@CFXi$6QioGF z>q~K6gKpGR{i(eEF^8%&t^{-H^~Z3)x4sM>u_LWUfgrqJ>notogRey=W`3!(F#LF4 zK$9}LvJ>AI`abi7;3{f1Oxq$5hBFvKkoQb`$Y9O^8knh=CQ6seb{vH-h{Tg5R!hf zsQ;SU8nJ}L%kvXv|A2Pt*05j_0bTtrB#aUbH-A6zi)G1B{%B-6bH6pRul=GgDG5o& z<8PvuT>oaX|LRB$brisyerw#HO|;AUPh52Z$h71?l=RmsYms<|6!jOhk);4S`>6#9 z?8cW4j)Pm_(%&W$^Z$_1JCONFf72F_umSDK&)WKPC$xm~PiWLZ18D2tjnnU3cl|m2 zEh(JYE4L=MnOsx2wOjqx_8kH{fBm&A^|!Fp-^K3xdW1xMdGPs|<;QPl|2>a-EJZzX zq8>jT_+=jQ__kjTss25adcH$Fy`i4u{JP@n6We||?NPtAn}$kwC-9d8yz)y?Km8x8 C6HK81 literal 0 HcmV?d00001 diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/__init__.py b/packages/peptide-property-ng/src/peptide_property_ng/train/__init__.py new file mode 100644 index 000000000..99224b64a --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/__init__.py @@ -0,0 +1 @@ +"""Training — multi-task optimisation.""" diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py new file mode 100644 index 000000000..1c8cab2f6 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py @@ -0,0 +1,145 @@ +"""Train the unified peptide property model — first prototype. + +Example: + python -m peptide_property_ng.train.train \\ + --datasets-glob '/scratch/claudius-proteomics/*' --cap 4000 --epochs 12 +""" +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch +from torch.utils.data import DataLoader + +from peptide_property_ng.data.collate import make_collate_fn +from peptide_property_ng.data.sage_dataset import build_split_datasets, discover_sage_dirs +from peptide_property_ng.eval.metrics import evaluate_split +from peptide_property_ng.losses import MultiTaskLoss +from peptide_property_ng.model.config import PRESETS +from peptide_property_ng.model.multitask import UnifiedPeptidePropertyModel +from peptide_property_ng.modifications.composition import CompositionTable + + +def _to_device(batch: dict, device: str) -> dict: + return {k: v.to(device) for k, v in batch.items()} + + +def train_epoch(model, loader, loss_fn, optimizer, device) -> dict[str, float]: + model.train() + totals: dict[str, float] = {} + n_batches = 0 + for batch in loader: + batch = _to_device(batch, device) + optimizer.zero_grad() + loss, parts = loss_fn(model(batch), batch) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + for name, value in parts.items(): + totals[name] = totals.get(name, 0.0) + value + totals["_total"] = totals.get("_total", 0.0) + float(loss.detach()) + n_batches += 1 + return {k: v / max(n_batches, 1) for k, v in totals.items()} + + +def main() -> None: + ap = argparse.ArgumentParser(description="Train the unified peptide property model.") + ap.add_argument("--datasets-glob", default="/scratch/claudius-proteomics/*") + ap.add_argument("--preset", default="small", choices=sorted(PRESETS)) + ap.add_argument("--cap", type=int, default=4000, help="max PSMs sampled per dataset") + ap.add_argument("--max-datasets", type=int, default=0, help="limit datasets (0 = all)") + ap.add_argument("--epochs", type=int, default=12) + ap.add_argument("--batch-size", type=int, default=64) + ap.add_argument("--lr", type=float, default=3e-4) + ap.add_argument("--weight-decay", type=float, default=0.01) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--patience", type=int, default=4) + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + ap.add_argument("--out", default="runs/prototype") + args = ap.parse_args() + + torch.manual_seed(args.seed) + out_dir = Path(args.out) + out_dir.mkdir(parents=True, exist_ok=True) + + print(f"[{time.strftime('%H:%M:%S')}] discovering datasets ...", flush=True) + sage_dirs = discover_sage_dirs(args.datasets_glob) + if args.max_datasets: + sage_dirs = sage_dirs[: args.max_datasets] + print(f" {len(sage_dirs)} datasets") + + print(f"[{time.strftime('%H:%M:%S')}] preparing examples (cap {args.cap}/dataset) ...", flush=True) + t0 = time.time() + splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed) + for name, ds in splits.items(): + print(f" {name}: {len(ds):,} examples") + print(f" prepared in {time.time() - t0:.0f}s") + if len(splits["train"]) == 0: + raise SystemExit("no training examples — check --datasets-glob") + + cfg = PRESETS[args.preset] + collate = make_collate_fn(cfg.pad_token_id, cfg.max_charge) + loaders = { + name: DataLoader( + ds, batch_size=args.batch_size, shuffle=(name == "train"), + collate_fn=collate, drop_last=(name == "train"), + ) + for name, ds in splits.items() + } + + device = args.device + model = UnifiedPeptidePropertyModel(cfg, CompositionTable.load()).to(device) + print(f"model: '{args.preset}' preset, {model.num_parameters():,} parameters, device={device}") + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + loss_fn = MultiTaskLoss() + + best_sa, best_epoch, bad = -1.0, -1, 0 + history: list[dict] = [] + for epoch in range(1, args.epochs + 1): + t0 = time.time() + tr = train_epoch(model, loaders["train"], loss_fn, optimizer, device) + val = evaluate_split(model, loaders["val"], device) + history.append({"epoch": epoch, "train": tr, "val": val}) + print( + f"[{time.strftime('%H:%M:%S')}] epoch {epoch:2d} " + f"train_loss={tr.get('_total', 0):.4f} " + f"(int={tr.get('intensity', 0):.3f} ccs={tr.get('ccs', 0):.3f} " + f"rt={tr.get('rt', 0):.3f} chg={tr.get('charge', 0):.3f}) " + f"val: SA={val['intensity_sa']:.4f} ccs_mae={val['ccs_mae']:.4f} " + f"rt_mae={val['rt_mae']:.4f} chg_acc={val['charge_acc']:.3f} " + f"[{time.time() - t0:.0f}s]", + flush=True, + ) + if val["intensity_sa"] > best_sa: + best_sa, best_epoch, bad = val["intensity_sa"], epoch, 0 + torch.save( + {"model_state_dict": model.state_dict(), "preset": args.preset, + "epoch": epoch, "val": val}, + out_dir / "best.pt", + ) + else: + bad += 1 + if bad >= args.patience: + print(f"early stop (no val SA gain for {args.patience} epochs)") + break + + ckpt = torch.load(out_dir / "best.pt", map_location=device) + model.load_state_dict(ckpt["model_state_dict"]) + test = evaluate_split(model, loaders["test"], device) + print(f"\n=== test (best epoch {best_epoch}) ===") + for k, v in test.items(): + print(f" {k:16s} {v:.4f}") + (out_dir / "metrics.json").write_text( + json.dumps( + {"args": vars(args), "best_epoch": best_epoch, "history": history, "test": test}, + indent=2, + ) + ) + print(f"\nsaved -> {out_dir}/best.pt , {out_dir}/metrics.json") + + +if __name__ == "__main__": + main() diff --git a/packages/peptide-property-ng/tests/__init__.py b/packages/peptide-property-ng/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/peptide-property-ng/tests/test_composition.py b/packages/peptide-property-ng/tests/test_composition.py new file mode 100644 index 000000000..f0e98a67c --- /dev/null +++ b/packages/peptide-property-ng/tests/test_composition.py @@ -0,0 +1,59 @@ +"""Tests for the UNIMOD atomic-composition table.""" +import numpy as np +import torch + +from peptide_property_ng.modifications.composition import ( + DEFAULT_TABLE, + ELEMENTS, + CompositionTable, + signed_log1p, +) + + +def _tokens() -> list[str]: + """Id-indexed token strings, read straight from the built table.""" + with np.load(DEFAULT_TABLE, allow_pickle=False) as z: + return [str(t) for t in z["tokens"]] + + +def test_table_loads(): + t = CompositionTable.load() + assert t.vocab_size == 2175 + assert t.n_elements == len(ELEMENTS) == 31 + assert t.elements == ELEMENTS + + +def test_known_modifications(): + """Carbamidomethyl (UNIMOD:4, +C2H3NO) and Oxidation (UNIMOD:35, +O).""" + t = CompositionTable.load() + vocab = _tokens() + ei = {e: i for i, e in enumerate(t.elements)} + + cam = next(i for i, v in enumerate(vocab) if "[UNIMOD:4]" in v) + row = t.counts[cam] + assert (row[ei["C"]], row[ei["H"]], row[ei["N"]], row[ei["O"]]) == (2, 3, 1, 1) + + ox = next(i for i, v in enumerate(vocab) if "[UNIMOD:35]" in v) + assert t.counts[ox][ei["O"]] == 1 + + +def test_bare_residue_is_zero(): + """A bare amino-acid token carries an all-zero composition delta.""" + t = CompositionTable.load() + ala = _tokens().index("A") + assert not t.counts[ala].any() + + +def test_signed_log1p_handles_negative_losses(): + """Neutral-loss modifications have negative counts; signed_log1p must stay finite.""" + x = torch.tensor([-40.0, -1.0, 0.0, 3.0, 120.0]) + out = signed_log1p(x) + assert torch.isfinite(out).all() + assert out[0] < 0 and out[2] == 0 and out[4] > 0 + # numpy path agrees + assert np.allclose(signed_log1p(x.numpy()), out.numpy()) + + +def test_table_has_negative_entries(): + """The table must preserve neutral-loss (negative) compositions.""" + assert (CompositionTable.load().counts < 0).any() diff --git a/packages/peptide-property-ng/tests/test_encoder.py b/packages/peptide-property-ng/tests/test_encoder.py new file mode 100644 index 000000000..f8c181430 --- /dev/null +++ b/packages/peptide-property-ng/tests/test_encoder.py @@ -0,0 +1,69 @@ +"""Tests for the hybrid embedding and the Depthcharge-based encoder.""" +import torch + +from peptide_property_ng.model.config import SMALL +from peptide_property_ng.model.encoder import PeptidePropertyEncoder +from peptide_property_ng.modifications.composition import CompositionTable + +# Random token ids in 1..59 — safely valid and never the pad id (61). +_LO, _HI = 1, 60 + + +def _encoder() -> PeptidePropertyEncoder: + return PeptidePropertyEncoder(SMALL, CompositionTable.load()) + + +def test_forward_shape(): + enc = _encoder() + b, length = 4, 12 + tokens = torch.randint(_LO, _HI, (b, length)) + latent, mask = enc(tokens) + assert latent.shape == (b, 1 + length, SMALL.d_model) + assert mask.shape == (b, 1 + length) + assert not mask[:, 0].any() # global token is never masked + + +def test_padding_mask(): + enc = _encoder() + tokens = torch.randint(_LO, _HI, (2, 8)) + tokens[0, 5:] = SMALL.pad_token_id # pad the last 3 residues of sample 0 + _, mask = enc(tokens) + assert mask[0, 6:].all() # +1 offset for the global token + assert not mask[0, :6].any() + assert not mask[1].any() + + +def test_instrument_conditioning_changes_output(): + enc = _encoder().eval() + tokens = torch.randint(_LO, _HI, (3, 10)) + base, _ = enc(tokens) + cond, _ = enc(tokens, instrument=torch.tensor([2, 2, 2])) + assert not torch.allclose(base, cond) + + +def test_long_peptide_runs(): + """A 50-residue peptide must encode fine — no fixed-length cap.""" + enc = _encoder() + tokens = torch.randint(_LO, _HI, (1, 50)) + latent, _ = enc(tokens) + assert latent.shape == (1, 51, SMALL.d_model) + + +def test_unmodified_residue_is_pure_token_embedding(): + """Bias-free composition encoder => a zero-composition token gets no chemistry term.""" + enc = _encoder() + he = enc.hybrid_embedding + zero_rows = (~he.composition.bool().any(dim=1)).nonzero().flatten() + tid = int(zero_rows[zero_rows > 0][0]) + tokens = torch.tensor([[tid]]) + assert torch.allclose(he(tokens), he.token_emb(tokens), atol=1e-6) + + +def test_modified_residue_differs_from_token_only(): + """A modified-residue token's hybrid embedding includes a chemistry term.""" + enc = _encoder() + he = enc.hybrid_embedding + mod_rows = he.composition.bool().any(dim=1).nonzero().flatten() + tid = int(mod_rows[0]) + tokens = torch.tensor([[tid]]) + assert not torch.allclose(he(tokens), he.token_emb(tokens), atol=1e-6) diff --git a/packages/peptide-property-ng/tests/test_multitask.py b/packages/peptide-property-ng/tests/test_multitask.py new file mode 100644 index 000000000..0dd3b3f4f --- /dev/null +++ b/packages/peptide-property-ng/tests/test_multitask.py @@ -0,0 +1,69 @@ +"""Tests for the unified multi-task model — output shapes and the leakage guard.""" +import torch + +from peptide_property_ng.model.config import SMALL +from peptide_property_ng.model.multitask import UnifiedPeptidePropertyModel + + +def _batch(b: int = 5, length: int = 14) -> dict[str, torch.Tensor]: + return { + "tokens": torch.randint(1, 60, (b, length)), + "charge": torch.randint(1, 5, (b,)), + "precursor_mz": torch.rand(b) * 800 + 300, + "collision_energy": torch.zeros(b), + "instrument": torch.randint(0, SMALL.n_instruments, (b,)), + "acq_mode": torch.randint(0, SMALL.n_acq_modes, (b,)), + } + + +def test_forward_all_heads(): + model = UnifiedPeptidePropertyModel(SMALL).eval() + b, length = 5, 14 + out = model(_batch(b, length)) + + assert set(out) == {"intensity", "ccs", "rt", "charge"} + assert out["intensity"].shape == (b, length - 1, SMALL.n_ion_channels) + assert (out["intensity"] >= 0).all() and (out["intensity"] <= 1).all() + + mean, std = out["ccs"] + assert mean.shape == (b,) and std.shape == (b,) + assert (std > 0).all() + + assert out["rt"].shape == (b,) + assert out["charge"].shape == (b, SMALL.max_charge) + + +def test_task_subset(): + model = UnifiedPeptidePropertyModel(SMALL).eval() + out = model(_batch(), tasks=["intensity"]) + assert set(out) == {"intensity"} + + +def test_charge_head_does_not_see_charge(): + """The charge head must be invariant to the precursor charge input + (charge is conditioned only at the intensity/CCS heads, never the encoder).""" + model = UnifiedPeptidePropertyModel(SMALL).eval() + batch = _batch() + with torch.no_grad(): + logits_a = model(batch, tasks=["charge"])["charge"] + batch2 = dict(batch, charge=torch.full_like(batch["charge"], 5)) + logits_b = model(batch2, tasks=["charge"])["charge"] + assert torch.allclose(logits_a, logits_b), "charge leaked into the charge head" + + +def test_backward_runs(): + model = UnifiedPeptidePropertyModel(SMALL) + out = model(_batch()) + loss = ( + out["intensity"].mean() + + out["ccs"][0].mean() + + out["rt"].mean() + + out["charge"].mean() + ) + loss.backward() + assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in model.parameters()) + + +def test_parameter_count_reported(): + model = UnifiedPeptidePropertyModel(SMALL) + assert model.num_parameters() > 1_000_000 From 098019dde88363e70cf2d24a99efa5004fd13c10 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 18:30:38 +0200 Subject: [PATCH 02/19] fix(peptide-property-ng): address review findings - losses/metrics: skip degenerate intensity targets (no observed peak) from the spectral-angle loss and metric -- they were a constant, zero-gradient term diluting both. - encoder: assert composition-table dimensions match the model config, rather than silently mis-indexing chemistry vectors. - train: save the final model as best.pt if no epoch improved val SA, so the post-training reload cannot crash. - embedding: add token_only / composition_only fusion modes for the modification-encoding ablation; expose via `train --comp-fusion`. - tests: synthetic fragment-target tests + loss / degenerate-handling tests (24 total). --- .../src/peptide_property_ng/eval/metrics.py | 7 +-- .../src/peptide_property_ng/losses.py | 17 ++++-- .../src/peptide_property_ng/model/config.py | 5 +- .../peptide_property_ng/model/embedding.py | 15 ++++-- .../src/peptide_property_ng/model/encoder.py | 13 +++++ .../src/peptide_property_ng/train/train.py | 14 +++++ .../tests/test_fragment_targets.py | 53 +++++++++++++++++++ .../peptide-property-ng/tests/test_losses.py | 33 ++++++++++++ 8 files changed, 146 insertions(+), 11 deletions(-) create mode 100644 packages/peptide-property-ng/tests/test_fragment_targets.py create mode 100644 packages/peptide-property-ng/tests/test_losses.py diff --git a/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py b/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py index 94803418f..cc10f6a25 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py @@ -3,7 +3,7 @@ import torch -from peptide_property_ng.losses import masked_spectral_angle +from peptide_property_ng.losses import intensity_signal_mask, masked_spectral_angle def _pearson(x: torch.Tensor, y: torch.Tensor) -> float: @@ -35,8 +35,9 @@ def evaluate_split(model, loader, device: str = "cpu") -> dict[str, float]: if "intensity" in out: sa = 1.0 - masked_spectral_angle(out["intensity"], batch["intensity_target"]) - sa_sum += float(sa.sum()) - sa_n += sa.numel() + signal = intensity_signal_mask(batch["intensity_target"]) + sa_sum += float(sa[signal].sum()) + sa_n += int(signal.sum()) if "ccs" in out: mean, _ = out["ccs"] diff --git a/packages/peptide-property-ng/src/peptide_property_ng/losses.py b/packages/peptide-property-ng/src/peptide_property_ng/losses.py index 933ebaf60..3d9ba871a 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/losses.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/losses.py @@ -22,6 +22,16 @@ def masked_spectral_angle(pred: torch.Tensor, target: torch.Tensor, eps: float = return 2.0 * torch.arccos(cos) / math.pi +def intensity_signal_mask(target: torch.Tensor) -> torch.Tensor: + """``(B,)`` bool — True where a target has at least one observed (positive) peak. + + A target with no positive b/y charge-1-3 peak is degenerate: the spectral + angle against it is a constant with zero gradient, so such samples are + excluded from the intensity loss and metric rather than diluting them. + """ + return (target > 0).flatten(1).any(dim=1) + + class MultiTaskLoss: """Weighted sum of per-task losses; tasks with no valid targets are skipped. @@ -44,9 +54,10 @@ def __call__( parts: dict[str, torch.Tensor] = {} if "intensity" in outputs: - parts["intensity"] = masked_spectral_angle( - outputs["intensity"], batch["intensity_target"] - ).mean() + sa = masked_spectral_angle(outputs["intensity"], batch["intensity_target"]) + signal = intensity_signal_mask(batch["intensity_target"]) + if signal.any(): + parts["intensity"] = sa[signal].mean() if "ccs" in outputs: mean, _std = outputs["ccs"] diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py index a53d435fa..a69516333 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py @@ -78,8 +78,9 @@ class ModelConfig: pad_token_id: int = 61 n_elements: int = 31 # atomic-composition channels (see modifications.composition) - # hybrid embedding - comp_fusion: str = "add" # "add" | "gate" + # hybrid embedding fusion: "add" | "gate" | "token_only" | "composition_only" + # ("token_only" / "composition_only" are ablation modes) + comp_fusion: str = "add" # conditioning max_charge: int = 8 diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/embedding.py b/packages/peptide-property-ng/src/peptide_property_ng/model/embedding.py index 3fdfc79e6..38e28b6a5 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/embedding.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/embedding.py @@ -54,8 +54,13 @@ class HybridResidueEmbedding(nn.Module): ``"gate"`` — ``token + sigmoid(g) * composition`` with a learnable per-channel gate ``g`` (init 0 → gate 0.5); the recommended upgrade if the unseen-modification probe shows the composition branch is ignored. + ``"token_only"`` / ``"composition_only"`` — ablation modes that drop one + branch, for measuring each branch's contribution (e.g. the + unseen-modification probe's control is ``token_only``). """ + _FUSIONS = ("add", "gate", "token_only", "composition_only") + def __init__( self, vocab_size: int, @@ -65,8 +70,8 @@ def __init__( fusion: str = "add", ): super().__init__() - if fusion not in ("add", "gate"): - raise ValueError(f"unknown fusion '{fusion}'") + if fusion not in self._FUSIONS: + raise ValueError(f"unknown fusion '{fusion}'; expected one of {self._FUSIONS}") self.d_model = d_model self.pad_token_id = pad_token_id self.fusion = fusion @@ -86,8 +91,12 @@ def __init__( def forward(self, tokens: torch.Tensor) -> torch.Tensor: """tokens: (batch, seq_len) long -> (batch, seq_len, d_model).""" - token_vec = self.token_emb(tokens) + if self.fusion == "token_only": + return self.token_emb(tokens) comp_vec = self.comp_encoder(self.composition[tokens]) + if self.fusion == "composition_only": + return comp_vec + token_vec = self.token_emb(tokens) if self.fusion == "gate": comp_vec = torch.sigmoid(self.gate) * comp_vec return token_vec + comp_vec diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py b/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py index 2b80dec94..2287f9522 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py @@ -40,6 +40,19 @@ def __init__(self, cfg: ModelConfig, composition_table: CompositionTable): ) self.cfg = cfg + # Fail loudly on a composition-table / config mismatch rather than + # silently indexing the wrong chemistry vectors. + if composition_table.vocab_size != cfg.vocab_size: + raise ValueError( + f"composition table vocab {composition_table.vocab_size} != " + f"cfg.vocab_size {cfg.vocab_size} — rebuild the table or fix the config" + ) + if composition_table.n_elements != cfg.n_elements: + raise ValueError( + f"composition table has {composition_table.n_elements} elements != " + f"cfg.n_elements {cfg.n_elements}" + ) + # The base class built a plain nn.Embedding token encoder; drop it and # use the hybrid residue embedding instead. del self.token_encoder diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py index 1c8cab2f6..2581c8ebe 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py @@ -7,6 +7,7 @@ from __future__ import annotations import argparse +import dataclasses import json import time from pathlib import Path @@ -49,6 +50,9 @@ def main() -> None: ap = argparse.ArgumentParser(description="Train the unified peptide property model.") ap.add_argument("--datasets-glob", default="/scratch/claudius-proteomics/*") ap.add_argument("--preset", default="small", choices=sorted(PRESETS)) + ap.add_argument("--comp-fusion", default=None, + choices=["add", "gate", "token_only", "composition_only"], + help="modification-encoding fusion / ablation (default: preset value)") ap.add_argument("--cap", type=int, default=4000, help="max PSMs sampled per dataset") ap.add_argument("--max-datasets", type=int, default=0, help="limit datasets (0 = all)") ap.add_argument("--epochs", type=int, default=12) @@ -81,6 +85,8 @@ def main() -> None: raise SystemExit("no training examples — check --datasets-glob") cfg = PRESETS[args.preset] + if args.comp_fusion: + cfg = dataclasses.replace(cfg, comp_fusion=args.comp_fusion) collate = make_collate_fn(cfg.pad_token_id, cfg.max_charge) loaders = { name: DataLoader( @@ -126,6 +132,14 @@ def main() -> None: print(f"early stop (no val SA gain for {args.patience} epochs)") break + if best_epoch == -1: + print("warning: no epoch improved val SA; saving the final model as best.pt") + torch.save( + {"model_state_dict": model.state_dict(), "preset": args.preset, + "epoch": epoch, "val": val}, + out_dir / "best.pt", + ) + best_epoch = epoch ckpt = torch.load(out_dir / "best.pt", map_location=device) model.load_state_dict(ckpt["model_state_dict"]) test = evaluate_split(model, loaders["test"], device) diff --git a/packages/peptide-property-ng/tests/test_fragment_targets.py b/packages/peptide-property-ng/tests/test_fragment_targets.py new file mode 100644 index 000000000..3fd5f3a01 --- /dev/null +++ b/packages/peptide-property-ng/tests/test_fragment_targets.py @@ -0,0 +1,53 @@ +"""Tests for fragment-indexed intensity-target construction. + +Channel layout: y+1=0, y+2=1, y+3=2, b+1=3, b+2=4, b+3=5. +""" +import numpy as np + +from peptide_property_ng.data.fragment_targets import N_ION_CHANNELS, build_intensity_target + + +def test_by_ion_site_placement(): + """b_o -> site o-1 ; y_o -> site (L-1)-o ; base-peak normalised.""" + length = 5 # 4 cleavage sites + ftype = np.array(["b", "b", "y", "y"]) + ford = np.array([1, 4, 1, 4], dtype=np.int32) + fch = np.array([1, 1, 1, 1], dtype=np.int32) + finten = np.array([10.0, 40.0, 20.0, 80.0], dtype=np.float32) + + t = build_intensity_target(length, 2, ftype, ford, fch, finten) + assert t.shape == (length - 1, N_ION_CHANNELS) + assert t[0, 3] == 10.0 / 80.0 # b1 -> site 0, channel b+1 + assert t[3, 3] == 40.0 / 80.0 # b4 -> site 3, channel b+1 + assert t[3, 0] == 20.0 / 80.0 # y1 -> site 3, channel y+1 + assert t[0, 0] == 80.0 / 80.0 # y4 -> site 0, channel y+1 (base peak) + + +def test_impossible_fragment_charges_masked(): + """Precursor charge 1 -> fragment charges 2 and 3 are impossible (-1).""" + t = build_intensity_target( + 5, 1, np.array(["b"]), np.array([1], np.int32), + np.array([1], np.int32), np.array([5.0], np.float32), + ) + assert (t[:, 1] == -1).all() and (t[:, 2] == -1).all() # y+2, y+3 + assert (t[:, 4] == -1).all() and (t[:, 5] == -1).all() # b+2, b+3 + assert (t[:, 0] >= 0).all() and (t[:, 3] >= 0).all() # charge-1 channels valid + + +def test_unobserved_possible_fragment_is_zero(): + """An observable-but-unmatched fragment is 0 (a real zero peak), not masked.""" + t = build_intensity_target( + 5, 2, np.array(["b"]), np.array([1], np.int32), + np.array([1], np.int32), np.array([5.0], np.float32), + ) + assert t[0, 3] == 1.0 # b1 observed -> base peak + assert t[1, 3] == 0.0 and t[2, 3] == 0.0 # unobserved but possible -> 0 + + +def test_out_of_range_ordinal_ignored(): + """A fragment ordinal beyond the peptide is dropped, not crashed on.""" + t = build_intensity_target( + 5, 2, np.array(["b"]), np.array([99], np.int32), + np.array([1], np.int32), np.array([5.0], np.float32), + ) + assert (t[t >= 0] == 0).all() # nothing placed diff --git a/packages/peptide-property-ng/tests/test_losses.py b/packages/peptide-property-ng/tests/test_losses.py new file mode 100644 index 000000000..3fa536551 --- /dev/null +++ b/packages/peptide-property-ng/tests/test_losses.py @@ -0,0 +1,33 @@ +"""Tests for the multi-task losses and degenerate-target handling.""" +import torch + +from peptide_property_ng.losses import MultiTaskLoss, intensity_signal_mask, masked_spectral_angle + + +def test_spectral_angle_identical_is_zero(): + x = torch.rand(4, 6, 6) + sa = masked_spectral_angle(x, x.clone()) + assert torch.allclose(sa, torch.zeros(4), atol=1e-3) + + +def test_spectral_angle_fully_masked_is_finite(): + """A target with every entry masked (-1) must not produce NaN.""" + sa = masked_spectral_angle(torch.rand(2, 5, 6), torch.full((2, 5, 6), -1.0)) + assert torch.isfinite(sa).all() + + +def test_signal_mask(): + target = torch.zeros(3, 4, 6) + target[0, 0, 0] = 0.5 # sample 0 has an observed peak + target[2] = -1.0 # sample 2 fully masked -> degenerate + assert intensity_signal_mask(target).tolist() == [True, False, False] + + +def test_loss_skips_degenerate_intensity_samples(): + """A batch mixing real and degenerate intensity targets yields a finite loss.""" + target = torch.zeros(4, 5, 6) + target[0, 1, 0] = 1.0 # only sample 0 carries signal + total, parts = MultiTaskLoss()({"intensity": torch.rand(4, 5, 6)}, + {"intensity_target": target}) + assert "intensity" in parts + assert torch.isfinite(torch.tensor(parts["intensity"])) From 34ffa869be7d722ec08cbbf96d2a325ced6db557 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 18:39:36 +0200 Subject: [PATCH 03/19] docs(peptide-property-ng): fix conditioning description in README charge / m-z / CE are conditioned at the heads, not in the encoder (only instrument / acq-mode go through the encoder global token, which keeps the charge head leak-free). The README table said otherwise. --- packages/peptide-property-ng/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/peptide-property-ng/README.md b/packages/peptide-property-ng/README.md index 62fe1097d..162a6b118 100644 --- a/packages/peptide-property-ng/README.md +++ b/packages/peptide-property-ng/README.md @@ -14,7 +14,7 @@ production model. See `docs/nn-architecture-exploration/` in the |---|---|---| | Modifications | opaque `[UNIMOD:N]` tokens (no chemistry) | **hybrid**: per-mod token **+** atomic-composition feature | | Unseen mods | embedding stays ~random | generalize via chemistry; open-vocab | -| Conditioning | dormant instrument embedding, never used | instrument / acq-mode / charge / m-z / CE inside the encoder | +| Conditioning | dormant instrument embedding, never used | instrument / acq-mode in the encoder; charge / m-z / CE at the heads (keeps the charge head leak-free) | | Intensity output | fixed 174-vector → 30-residue cap | **fragment-indexed** → no length cap | | Backbone | bespoke transformer | built on **Depthcharge** primitives | From c73528fc305eed658104b4016a2b6fa358492e47 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 19:48:58 +0200 Subject: [PATCH 04/19] refactor(peptide-property-ng): unify intensity encoding via the Prosit 174-vector Both data sources now reach the site-indexed intensity target through one proven encoder + one conversion, replacing the hand-rolled site encoder: - Sage fragments -> 174-vector via imspy's observed_fragments_to_intensity_target (the encoding the rescoring pipeline already uses), then prosit174_to_sites. - Wilhelmlab HF MS2 datasets store the 174-vector natively -> prosit174_to_sites. prosit174_to_sites is the single ordinal->site remap; the hand-rolled build_intensity_target is removed. Adds data/hf_intensity.py for the timsTOF-ms2 / prospect-ptms-ms2 / Prosit-2025 intensity-pretraining datasets. Intensity peptides are capped at 30 aa (the 174-vector limit). --- .../data/fragment_targets.py | 101 ++++++++-------- .../peptide_property_ng/data/hf_intensity.py | 114 ++++++++++++++++++ .../peptide_property_ng/data/sage_dataset.py | 19 +-- .../tests/test_fragment_targets.py | 99 ++++++++------- 4 files changed, 231 insertions(+), 102 deletions(-) create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/hf_intensity.py diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/fragment_targets.py b/packages/peptide-property-ng/src/peptide_property_ng/data/fragment_targets.py index aa4a6aa1b..dc4f58177 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/fragment_targets.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/fragment_targets.py @@ -1,64 +1,67 @@ -"""Fragment-indexed intensity targets from Sage ``matched_fragments``. +"""Fragment-intensity targets. -Each peptide of length ``L`` yields an ``(L-1, 6)`` target — one row per -cleavage site, six channels for b/y ions at fragment charges 1-3. This is the -variable-length successor to Prosit's fixed 174-vector (no 30-residue cap). +The MS2 intensity target the model trains against is **cleavage-site-indexed**: +``(L-1, 6)`` — one row per cleavage site, channels ``[y+1, y+2, y+3, b+1, b+2, b+3]``. -Convention (Prosit / imspy compatible): - -1 physically impossible channel (fragment charge > precursor charge) - 0 possible but unobserved fragment - 0..1 observed, base-peak-normalised intensity -The masked spectral-angle loss ignores -1 entries. +Both data sources reach it through the *same single conversion*, +``prosit174_to_sites``: + - Public Wilhelmlab HF datasets store intensities natively as the Prosit + 174-vector. + - Sage search results are turned into that 174-vector by the proven + ``imspy_predictors`` encoder ``observed_fragments_to_intensity_target`` + (see ``data/sage_dataset.py``) — the fragment→vector encoding is *not* + re-implemented here. + +So there is exactly one place where the ordinal→site remap happens, and it is +this file. + +Convention (Prosit / imspy compatible): ``-1`` = impossible/masked channel, +``0`` = possible but unobserved, ``0..1`` = observed (base-peak-normalised). """ from __future__ import annotations import numpy as np -# Channel order: y ions at charge 1/2/3, then b ions at charge 1/2/3. +# Channel order — verified against imspy_predictors/intensity/utility.py:211 +# ("Ion types in order: y+1, y+2, y+3, b+1, b+2, b+3"). ION_CHANNELS: tuple[str, ...] = ("y+1", "y+2", "y+3", "b+1", "b+2", "b+3") N_ION_CHANNELS = len(ION_CHANNELS) -MAX_FRAGMENT_CHARGE = 3 +# The Prosit 174-vector covers fragment ordinals 1..29 (peptides up to 30 aa). +PROSIT_MAX_ORDINAL = 29 +PROSIT_VECTOR_LEN = PROSIT_MAX_ORDINAL * N_ION_CHANNELS # 174 -def build_intensity_target( - peptide_len: int, - precursor_charge: int, - frag_type: np.ndarray, # ('b'|'y') per matched fragment - frag_ordinal: np.ndarray, # int - frag_charge: np.ndarray, # int - frag_intensity: np.ndarray, # float -) -> np.ndarray: - """Return the ``(peptide_len - 1, 6)`` intensity target for one PSM.""" - n_sites = peptide_len - 1 - target = np.zeros((n_sites, N_ION_CHANNELS), dtype=np.float32) - # Mark fragment charges above the precursor charge as impossible. - for fc in range(1, MAX_FRAGMENT_CHARGE + 1): - if fc > precursor_charge: - target[:, fc - 1] = -1.0 # y channel - target[:, 3 + fc - 1] = -1.0 # b channel +def prosit174_to_sites(intensities_174, peptide_len: int) -> np.ndarray: + """Convert a Prosit 174-element intensity vector to our site-indexed target. + + The Prosit vector is **ordinal-indexed**: reshaped to (29, 6), row ``r`` holds + the ions of fragment ordinal ``r+1`` — ``y_{r+1}`` in columns 0:3, ``b_{r+1}`` + in columns 3:6. - for ftype, ordi, fch, inten in zip( - frag_type, frag_ordinal, frag_charge, frag_intensity - ): - fch = int(fch) - if fch < 1 or fch > MAX_FRAGMENT_CHARGE: - continue - ordi = int(ordi) - if ftype == "b": - site, channel = ordi - 1, 3 + fch - 1 - elif ftype == "y": - # cleavage site i produces y_{n_sites - i + ...}: y_o sits at site n_sites - o - site, channel = n_sites - ordi, fch - 1 - else: - continue - if 0 <= site < n_sites and target[site, channel] >= 0.0: - target[site, channel] = max(target[site, channel], float(inten)) + Our ``(L-1, 6)`` target is **cleavage-site-indexed**: site ``i`` holds the two + ions produced by cleaving bond ``i`` — the b ion ``b_{i+1}`` and its + *complementary* y ion ``y_{L-1-i}``. The b and y of one cleavage have + different ordinals (``b_k`` pairs with ``y_{L-k}``), so this is a genuine + remap, not a reshape. + + Mapping (``L`` = peptide_len, ``n_sites`` = L-1): + site i, b channels (3:6) <- ordinal i+1 -> Prosit row i + site i, y channels (0:3) <- ordinal L-1-i -> Prosit row (L-1-i)-1 = L-2-i + + The ``-1`` masked-channel markers in the Prosit vector carry through. + """ + arr = np.asarray(intensities_174, dtype=np.float32).reshape(-1) + if arr.size != PROSIT_VECTOR_LEN: + raise ValueError(f"expected {PROSIT_VECTOR_LEN} intensity values, got {arr.size}") + prosit = arr.reshape(PROSIT_MAX_ORDINAL, N_ION_CHANNELS) # row r -> fragment ordinal r+1 + + n_sites = peptide_len - 1 + if not 1 <= n_sites <= PROSIT_MAX_ORDINAL: + raise ValueError(f"peptide_len {peptide_len} outside the 174-vector range [2, 30]") - # Base-peak normalisation over the observed (positive) entries. - observed = target > 0.0 - if observed.any(): - peak = float(target[observed].max()) - if peak > 0.0: - target[observed] /= peak + target = np.full((n_sites, N_ION_CHANNELS), -1.0, dtype=np.float32) + for i in range(n_sites): + target[i, 3:6] = prosit[i, 3:6] # b_{i+1} + target[i, 0:3] = prosit[n_sites - 1 - i, 0:3] # y_{L-1-i} (Prosit row L-2-i) return target diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/hf_intensity.py b/packages/peptide-property-ng/src/peptide_property_ng/data/hf_intensity.py new file mode 100644 index 000000000..719fdef3f --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/hf_intensity.py @@ -0,0 +1,114 @@ +"""Wilhelmlab HuggingFace MS2 datasets -> intensity-pretraining examples. + +The three Wilhelmlab MS2 datasets — `timsTOF-ms2`, `prospect-ptms-ms2`, +`Prosit-2025-lac-ms2` — share one schema: + - `intensities_raw` : the Prosit 174-vector (`-1` = masked) + - `modified_sequence` : `[]-...-[]` terminal placeholders + inline `[UNIMOD:N]` + - `precursor_charge_onehot` : 6-way one-hot charge + - `collision_energy_aligned_normed` : normalised collision energy + +Each row becomes an example dict identical in shape to the Sage loader's, so the +same collate function and model consume it. CCS / RT are absent here (masked). +The intensity target is produced by the *same* `prosit174_to_sites` conversion +the Sage path uses. +""" +from __future__ import annotations + +import itertools +import re + +import numpy as np + +from peptide_property_ng.data.fragment_targets import prosit174_to_sites +from peptide_property_ng.data.sage_dataset import SagePropertyDataset + +_UNIMOD_RE = re.compile(r"\[UNIMOD:\d+\]") + + +def _strip_terminal_placeholders(modseq: str) -> str: + """Drop empty `[]-` / `-[]` terminal-modification placeholders. + + A non-empty terminal bracket (a real terminal mod) is left in place; such a + peptide then fails the token-count check and is skipped — same policy as the + Sage loader. + """ + modseq = re.sub(r"^\[\]-", "", modseq) + modseq = re.sub(r"-\[\]$", "", modseq) + return modseq + + +def prepare_hf_intensity_examples( + dataset_name: str, + *, + split: str = "train", + cap: int | None = None, + instrument: int = 0, + acq_mode: int = 0, + max_len: int = 30, + batch: int = 2048, +) -> list[dict]: + """Stream a Wilhelmlab HF MS2 dataset into intensity-pretraining examples.""" + from datasets import load_dataset + from imspy_predictors.utilities.tokenizers import ProformaTokenizer + + tok = ProformaTokenizer.with_defaults() + rows = load_dataset(dataset_name, split=split, streaming=True) + if cap is not None: + rows = itertools.islice(rows, cap) + + examples: list[dict] = [] + buf: list[dict] = [] + + def _flush() -> None: + if not buf: + return + modseqs = [_strip_terminal_placeholders(r["modified_sequence"]) for r in buf] + enc = tok(modseqs) + ids, attn = enc["input_ids"], enc["attention_mask"] + for row, modseq, tok_ids, am in zip(buf, modseqs, ids, attn): + if "-" in modseq or "[+" in modseq: + continue # an unresolved terminal mod / delta mass — skip + n_real = int(sum(am)) + residue_ids = tok_ids[1 : n_real - 1] # strip [CLS]/[SEP] + right-padding + length = len(residue_ids) + stripped = _UNIMOD_RE.sub("", modseq) + if length != len(stripped) or not 3 <= length <= max_len: + continue + charge = int(np.argmax(row["precursor_charge_onehot"])) + 1 + try: + target = prosit174_to_sites(row["intensities_raw"], length) + except ValueError: + continue + examples.append( + { + "accession": dataset_name, + "stripped": stripped, + "modseq": modseq, + "tokens": np.asarray(residue_ids, dtype=np.int64), + "charge": charge, + "precursor_mz": 0.0, # absent here; unused by the intensity task + "collision_energy": float(row["collision_energy_aligned_normed"]), + "instrument": int(instrument), + "acq_mode": int(acq_mode), + "intensity_target": target, + "ccs_target": float("nan"), "ccs_valid": False, + "rt_target": float("nan"), "rt_valid": False, + } + ) + buf.clear() + + for row in rows: + buf.append(row) + if len(buf) >= batch: + _flush() + _flush() + return examples + + +def build_hf_intensity_dataset( + dataset_name: str, *, split: str = "train", cap: int | None = None, **kwargs +) -> SagePropertyDataset: + """Prepared HF intensity examples wrapped as a Dataset.""" + return SagePropertyDataset( + prepare_hf_intensity_examples(dataset_name, split=split, cap=cap, **kwargs) + ) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py index e5b503676..663adba95 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py @@ -10,6 +10,7 @@ import glob from collections import defaultdict from pathlib import Path +from types import SimpleNamespace import numpy as np import pyarrow as pa @@ -17,7 +18,7 @@ import pyarrow.parquet as pq from torch.utils.data import Dataset -from peptide_property_ng.data.fragment_targets import build_intensity_target +from peptide_property_ng.data.fragment_targets import prosit174_to_sites from peptide_property_ng.data.splits import peptide_split _PROTON = 1.007276 @@ -63,12 +64,13 @@ def prepare_examples( cap: int = 8000, q_max: float = 0.01, min_peaks: int = 6, - max_len: int = 64, + max_len: int = 30, # the Prosit 174-vector intensity encoding caps peptides at 30 aa seed: int = 0, instrument: int = 0, acq_mode: int = 0, ) -> list[dict]: """Load and prepare one dataset's PSMs into a list of example dicts.""" + from imspy_predictors.intensity.predictors import observed_fragments_to_intensity_target from imspy_predictors.utilities.tokenizers import ProformaTokenizer from sagepy_rescore.sage_loader import _parse_sage_peptide @@ -139,13 +141,14 @@ def prepare_examples( cols = by_psm.get(df["psm_id"][k]) if not cols or not cols[0]: continue # no matched fragments -> no intensity target - target = build_intensity_target( - length, charge, - np.asarray(cols[0]), - np.asarray(cols[1], dtype=np.int32), - np.asarray(cols[2], dtype=np.int32), - np.asarray(cols[3], dtype=np.float32), + # Encode fragments with the proven imspy encoder (Sage fragments -> + # Prosit 174-vector), then the single shared 174 -> site conversion. + frag = SimpleNamespace( + ion_types=cols[0], fragment_ordinals=cols[1], + charges=cols[2], intensities=cols[3], ) + prosit174 = observed_fragments_to_intensity_target(stripped, charge, frag) + target = prosit174_to_sites(prosit174, length) im = df["ion_mobility"][k] rt = df["aligned_rt"][k] diff --git a/packages/peptide-property-ng/tests/test_fragment_targets.py b/packages/peptide-property-ng/tests/test_fragment_targets.py index 3fd5f3a01..5713bef6c 100644 --- a/packages/peptide-property-ng/tests/test_fragment_targets.py +++ b/packages/peptide-property-ng/tests/test_fragment_targets.py @@ -1,53 +1,62 @@ -"""Tests for fragment-indexed intensity-target construction. +"""Tests for the Prosit-174 -> site-indexed intensity-target conversion. -Channel layout: y+1=0, y+2=1, y+3=2, b+1=3, b+2=4, b+3=5. +This conversion is the single place the ordinal->site remap happens, and it is +easy to get wrong — hence the explicit b/y placement tests below. """ import numpy as np +import pytest -from peptide_property_ng.data.fragment_targets import N_ION_CHANNELS, build_intensity_target +from peptide_property_ng.data.fragment_targets import ( + N_ION_CHANNELS, + PROSIT_MAX_ORDINAL, + PROSIT_VECTOR_LEN, + prosit174_to_sites, +) -def test_by_ion_site_placement(): - """b_o -> site o-1 ; y_o -> site (L-1)-o ; base-peak normalised.""" - length = 5 # 4 cleavage sites - ftype = np.array(["b", "b", "y", "y"]) - ford = np.array([1, 4, 1, 4], dtype=np.int32) - fch = np.array([1, 1, 1, 1], dtype=np.int32) - finten = np.array([10.0, 40.0, 20.0, 80.0], dtype=np.float32) +def _prosit_vector(entries: dict[tuple[int, int], float]) -> np.ndarray: + """Build a 174-vector; ``entries`` maps (fragment_ordinal, channel) -> value.""" + v = np.zeros((PROSIT_MAX_ORDINAL, N_ION_CHANNELS), np.float32) + for (ordinal, channel), value in entries.items(): + v[ordinal - 1, channel] = value + return v.reshape(-1) + - t = build_intensity_target(length, 2, ftype, ford, fch, finten) +def test_b_ion_site_mapping(): + """b_k (Prosit ordinal k) lands at site k-1, b channels (3:6).""" + length = 5 # 4 cleavage sites + t = prosit174_to_sites(_prosit_vector({(1, 3): 0.11, (4, 3): 0.44}), length) assert t.shape == (length - 1, N_ION_CHANNELS) - assert t[0, 3] == 10.0 / 80.0 # b1 -> site 0, channel b+1 - assert t[3, 3] == 40.0 / 80.0 # b4 -> site 3, channel b+1 - assert t[3, 0] == 20.0 / 80.0 # y1 -> site 3, channel y+1 - assert t[0, 0] == 80.0 / 80.0 # y4 -> site 0, channel y+1 (base peak) - - -def test_impossible_fragment_charges_masked(): - """Precursor charge 1 -> fragment charges 2 and 3 are impossible (-1).""" - t = build_intensity_target( - 5, 1, np.array(["b"]), np.array([1], np.int32), - np.array([1], np.int32), np.array([5.0], np.float32), - ) - assert (t[:, 1] == -1).all() and (t[:, 2] == -1).all() # y+2, y+3 - assert (t[:, 4] == -1).all() and (t[:, 5] == -1).all() # b+2, b+3 - assert (t[:, 0] >= 0).all() and (t[:, 3] >= 0).all() # charge-1 channels valid - - -def test_unobserved_possible_fragment_is_zero(): - """An observable-but-unmatched fragment is 0 (a real zero peak), not masked.""" - t = build_intensity_target( - 5, 2, np.array(["b"]), np.array([1], np.int32), - np.array([1], np.int32), np.array([5.0], np.float32), - ) - assert t[0, 3] == 1.0 # b1 observed -> base peak - assert t[1, 3] == 0.0 and t[2, 3] == 0.0 # unobserved but possible -> 0 - - -def test_out_of_range_ordinal_ignored(): - """A fragment ordinal beyond the peptide is dropped, not crashed on.""" - t = build_intensity_target( - 5, 2, np.array(["b"]), np.array([99], np.int32), - np.array([1], np.int32), np.array([5.0], np.float32), - ) - assert (t[t >= 0] == 0).all() # nothing placed + assert t[0, 3] == 0.11 # b1 -> site 0 + assert t[3, 3] == 0.44 # b4 -> site 3 + + +def test_y_ion_site_mapping(): + """y_o (Prosit ordinal o) lands at site L-1-o, y channels (0:3).""" + length = 5 + t = prosit174_to_sites(_prosit_vector({(1, 0): 0.81, (4, 0): 0.84}), length) + assert t[3, 0] == 0.81 # y1 -> site L-1-1 = 3 + assert t[0, 0] == 0.84 # y4 -> site L-1-4 = 0 + + +def test_complementary_ions_share_a_site(): + """Site i must hold b_{i+1} and y_{L-1-i} — the complementary pair of one cleavage.""" + length = 6 # site 2 should carry b3 and y3 + t = prosit174_to_sites(_prosit_vector({(3, 3): 0.3, (3, 0): 0.9}), length) + assert t[2, 3] == 0.3 # b3 at site 2 + assert t[2, 0] == 0.9 # y3 at site 2 + + +def test_minus_one_mask_carries_through(): + t = prosit174_to_sites(np.full(PROSIT_VECTOR_LEN, -1.0, np.float32), 10) + assert (t == -1.0).all() + + +def test_shape_and_bounds(): + v = np.zeros(PROSIT_VECTOR_LEN, np.float32) + assert prosit174_to_sites(v, 30).shape == (29, N_ION_CHANNELS) # max peptide length + assert prosit174_to_sites(v, 3).shape == (2, N_ION_CHANNELS) + with pytest.raises(ValueError): + prosit174_to_sites(v, 31) # peptide longer than the 174-vector + with pytest.raises(ValueError): + prosit174_to_sites(np.zeros(100), 10) # wrong vector length From d62293a734b93adcc13d1d1a4a0e30ae7747222b Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 20:22:22 +0200 Subject: [PATCH 05/19] test(peptide-property-ng): pin all six intensity channels in the 174->site test Codex's conversion review noted the tests pinned y-channel-0 / b-channel-3 but not the charge order within each y/b triplet. This test sets distinct values across all six channels of one ordinal so a charge permutation would fail. --- .../tests/test_fragment_targets.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/packages/peptide-property-ng/tests/test_fragment_targets.py b/packages/peptide-property-ng/tests/test_fragment_targets.py index 5713bef6c..2cb063e90 100644 --- a/packages/peptide-property-ng/tests/test_fragment_targets.py +++ b/packages/peptide-property-ng/tests/test_fragment_targets.py @@ -47,6 +47,18 @@ def test_complementary_ions_share_a_site(): assert t[2, 0] == 0.9 # y3 at site 2 +def test_all_six_channels_pinned(): + """Pin every channel — a charge permutation within a y/b triplet must fail.""" + length = 5 # 4 sites; ordinal 2 -> b2 at site 1, y2 at site L-1-2 = 2 + v = _prosit_vector({ + (2, 0): 0.1, (2, 1): 0.2, (2, 2): 0.3, # y2 charges 1,2,3 + (2, 3): 0.4, (2, 4): 0.5, (2, 5): 0.6, # b2 charges 1,2,3 + }) + t = prosit174_to_sites(v, length) + assert list(t[1, 3:6]) == [0.4, 0.5, 0.6] # b2 -> site 1, b+1/b+2/b+3 + assert list(t[2, 0:3]) == [0.1, 0.2, 0.3] # y2 -> site 2, y+1/y+2/y+3 + + def test_minus_one_mask_carries_through(): t = prosit174_to_sites(np.full(PROSIT_VECTOR_LEN, -1.0, np.float32), 10) assert (t == -1.0).all() From 249d73577f3d274166eafdcf55c03723077d5ca0 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 20:47:54 +0200 Subject: [PATCH 06/19] feat(peptide-property-ng): staged pretraining on public corpora Adds the pretrain-then-fine-tune pipeline: - data/chronologer_rt.py -- Chronologer DB (2.64M peptide-RT, harmonised HI) - data/ccs_pretrain.py -- ionmob CCS parquets (UNIMOD-annotated) - train/pretrain.py -- staged single-task curriculum: Orbitrap intensity -> timsTOF intensity -> CCS -> RT, accumulating in the shared encoder; saves a checkpoint train.py picks up via --init-from. - config: add an `orbitrap` instrument id so the embedding can bridge the Orbitrap-HCD -> timsTOF-PASEF domain shift across pretraining stages. CCS/RT pretrain on their native units (CCS, harmonised HI); the unit shift to Sage 1/K0 / aligned_rt is absorbed during the campaign fine-tune -- no error-prone physics conversion. Intensity uses the one proven 174-vector encoding throughout. --- .../peptide_property_ng/data/ccs_pretrain.py | 122 +++++++++++++ .../data/chronologer_rt.py | 97 +++++++++++ .../src/peptide_property_ng/model/config.py | 5 +- .../src/peptide_property_ng/train/pretrain.py | 160 ++++++++++++++++++ .../src/peptide_property_ng/train/train.py | 7 + 5 files changed, 390 insertions(+), 1 deletion(-) create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py b/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py new file mode 100644 index 000000000..4502c3202 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py @@ -0,0 +1,122 @@ +"""ionmob CCS datasets -> ion-mobility-pretraining examples. + +Uses the deduplicated, UNIMOD-annotated ionmob CCS parquets +(`*_unique_unimod*.parquet`, Meier / Tenzer / Chang / ...): schema +``mz, charge, sequence-tokenized, ccs, name``. + +Target-unit note: these provide **CCS** (Angstrom^2); the Sage fine-tuning data +provides **inverse ion mobility** (1/K0). They are deliberately *not* converted +into each other — that physics conversion is error-prone, and it is unnecessary: +pretraining is single-task per stage, so the CCS head simply warms on CCS here +and the unit/scale shift to 1/K0 is absorbed during the campaign fine-tune — +exactly as the RT head pretrains on Chronologer ``HI`` and fine-tunes on Sage +``aligned_rt``. +""" +from __future__ import annotations + +import glob +import re + +import numpy as np +import pyarrow.parquet as pq + +from peptide_property_ng.data.sage_dataset import SagePropertyDataset + +DEFAULT_CCS_GLOB = ( + "/home/administrator/Documents/promotion/rust/rustims-submission/" + "re-scoring/models/unimod/**/*_unique_unimod*.parquet" +) +_UNIMOD_RE = re.compile(r"\[UNIMOD:\d+\]") + + +def _tokens_to_proforma(tokens: list[str]) -> str | None: + """ionmob ``sequence-tokenized`` list -> a ProForma string, or ``None`` to skip. + + A ```` / ```` marker carrying a terminal modification (e.g. + ``[UNIMOD:1]``) flags a terminal mod -> skip, consistent with the + other loaders. + """ + if len(tokens) < 5 or tokens[0] != "" or tokens[-1] != "": + return None + return "".join(tokens[1:-1]) + + +def prepare_ccs_examples( + data_glob: str = DEFAULT_CCS_GLOB, + *, + cap: int | None = None, + instrument: int = 0, + acq_mode: int = 0, + max_len: int = 64, + seed: int = 0, +) -> list[dict]: + """Load the ionmob CCS parquets into ion-mobility-pretraining example dicts.""" + from imspy_predictors.utilities.tokenizers import ProformaTokenizer + + files = sorted(set(glob.glob(data_glob, recursive=True))) + if not files: + return [] + seqs: list[list[str]] = [] + charge: list[int] = [] + mz: list[float] = [] + ccs: list[float] = [] + for path in files: + d = pq.read_table(path, columns=["mz", "charge", "sequence-tokenized", "ccs"]).to_pydict() + seqs += d["sequence-tokenized"] + charge += d["charge"] + mz += d["mz"] + ccs += d["ccs"] + + idx = np.arange(len(seqs)) + if cap is not None and len(idx) > cap: + idx = np.sort(np.random.RandomState(seed).choice(len(idx), cap, replace=False)) + + tok = ProformaTokenizer.with_defaults() + proformas: list[str] = [] + src: list[int] = [] + for i in idx: + pf = _tokens_to_proforma(seqs[i]) + if pf is not None: + proformas.append(pf) + src.append(int(i)) + + examples: list[dict] = [] + chunk = 20000 # tokenise in chunks (batch tokenisation right-pads to the chunk max) + for start in range(0, len(proformas), chunk): + block = proformas[start : start + chunk] + block_src = src[start : start + chunk] + enc = tok(block) + ids, attn = enc["input_ids"], enc["attention_mask"] + for j, pf in enumerate(block): + i = block_src[j] + n_real = int(sum(attn[j])) + residue_ids = ids[j][1 : n_real - 1] + length = len(residue_ids) + stripped = _UNIMOD_RE.sub("", pf) + if length != len(stripped) or not 3 <= length <= max_len: + continue + examples.append( + { + "accession": "ionmob", + "stripped": stripped, + "modseq": pf, + "tokens": np.asarray(residue_ids, dtype=np.int64), + "charge": int(charge[i]), + "precursor_mz": float(mz[i]), + "collision_energy": 0.0, + "instrument": int(instrument), + "acq_mode": int(acq_mode), + # masked placeholder — CCS pretraining runs only the CCS task + "intensity_target": np.full((length - 1, 6), -1.0, dtype=np.float32), + "ccs_target": float(ccs[i]), + "ccs_valid": True, + "rt_target": float("nan"), + "rt_valid": False, + } + ) + return examples + + +def build_ccs_dataset(data_glob: str = DEFAULT_CCS_GLOB, *, cap: int | None = None, **kwargs): + """Prepared ionmob CCS examples wrapped as a Dataset.""" + return SagePropertyDataset(prepare_ccs_examples(data_glob, cap=cap, **kwargs)) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py b/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py new file mode 100644 index 000000000..ae2c677c5 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py @@ -0,0 +1,97 @@ +"""Chronologer retention-time database -> RT-pretraining examples. + +The Chronologer DB (``Chronologer_DB_220308.gz``, ~2.64 M peptide-RT records +harmonised across 11 community datasets) is the RT pretraining corpus. Its +``HI`` (hydrophobicity index) column is the cross-dataset-harmonised RT target +-- so it sidesteps the per-run RT-scale problem -- and is min-max normalised to +``[0, 1]`` to match the model's RT-head output range and the Sage ``aligned_rt`` +fine-tuning target. + +Peptides use Sage-style ``[+mass]`` delta notation; they are converted to UNIMOD +form with the same ``sagepy_rescore`` converter the Sage loader uses. +""" +from __future__ import annotations + +import re +from pathlib import Path + +import numpy as np + +from peptide_property_ng.data.sage_dataset import SagePropertyDataset + +DEFAULT_CHRONOLOGER_DB = Path( + "/home/administrator/Documents/promotion/chronologer/data/Chronologer_DB_220308.gz" +) +# HI range across the full DB — fixed so normalisation is reproducible for any cap. +HI_MIN, HI_MAX = -1.01, 31.07 +_UNIMOD_RE = re.compile(r"\[UNIMOD:\d+\]") + + +def _normalise_hi(hi: float) -> float: + return (float(hi) - HI_MIN) / (HI_MAX - HI_MIN) + + +def prepare_chronologer_examples( + db_path: str | Path = DEFAULT_CHRONOLOGER_DB, + *, + cap: int | None = None, + instrument: int = 0, + acq_mode: int = 0, + max_len: int = 64, + seed: int = 0, +) -> list[dict]: + """Load the Chronologer DB into RT-pretraining example dicts.""" + import pandas as pd + from imspy_predictors.utilities.tokenizers import ProformaTokenizer + from sagepy_rescore.sage_loader import _parse_sage_peptide + + df = pd.read_csv(db_path, sep="\t", compression="gzip", usecols=["PeptideModSeq", "HI"]) + if cap is not None and len(df) > cap: + df = df.sample(cap, random_state=seed) + peptides = df["PeptideModSeq"].tolist() + hi = df["HI"].to_numpy(dtype=np.float32) + + tok = ProformaTokenizer.with_defaults() + parsed = [_parse_sage_peptide(p) for p in peptides] # [+mass] -> UNIMOD + + examples: list[dict] = [] + chunk = 20000 # tokenise in chunks (batch tokenisation right-pads to the chunk max) + for start in range(0, len(parsed), chunk): + block = parsed[start : start + chunk] + enc = tok([m for _, m in block]) + ids, attn = enc["input_ids"], enc["attention_mask"] + for j, (stripped, modseq) in enumerate(block): + if "[+" in modseq or "[-" in modseq: + continue # an unconverted delta mass — skip + n_real = int(sum(attn[j])) + residue_ids = ids[j][1 : n_real - 1] # strip [CLS]/[SEP] + right-padding + length = len(residue_ids) + if length != len(stripped) or not 3 <= length <= max_len: + continue + examples.append( + { + "accession": "Chronologer", + "stripped": stripped, + "modseq": modseq, + "tokens": np.asarray(residue_ids, dtype=np.int64), + "charge": 2, # dummy — RT is charge-independent; the RT head ignores charge + "precursor_mz": 0.0, + "collision_energy": 0.0, + "instrument": int(instrument), + "acq_mode": int(acq_mode), + # masked placeholder — RT pretraining runs only the RT task + "intensity_target": np.full((length - 1, 6), -1.0, dtype=np.float32), + "ccs_target": float("nan"), + "ccs_valid": False, + "rt_target": _normalise_hi(hi[start + j]), + "rt_valid": True, + } + ) + return examples + + +def build_chronologer_dataset( + db_path: str | Path = DEFAULT_CHRONOLOGER_DB, *, cap: int | None = None, **kwargs +) -> SagePropertyDataset: + """Prepared Chronologer RT examples wrapped as a Dataset.""" + return SagePropertyDataset(prepare_chronologer_examples(db_path, cap=cap, **kwargs)) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py index a69516333..229bf1a9c 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py @@ -12,7 +12,9 @@ from dataclasses import dataclass, field -# timsTOF-focused instrument vocabulary (from timstof_catalog.tsv instrument_model). +# Instrument vocabulary — timsTOF-focused (the target domain), plus `orbitrap` +# so Orbitrap-HCD pretraining corpora (PROSPECT) carry a distinct id and the +# instrument embedding can learn the HCD<->PASEF domain shift. INSTRUMENTS: tuple[str, ...] = ( "unknown", "timsTOF", @@ -22,6 +24,7 @@ "timsTOF SCP", "timsTOF Ultra", "timsTOF fleX", + "orbitrap", ) ACQUISITION_MODES: tuple[str, ...] = ( "unknown", diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py new file mode 100644 index 000000000..197747121 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py @@ -0,0 +1,160 @@ +"""Staged pretraining — warm the shared encoder + heads on large public corpora, +then hand off to ``train.py`` for the multi-task campaign fine-tune. + +Default curriculum (each stage is single-task; the shared encoder accumulates): + 1. intensity — Wilhelmlab/prospect-ptms-ms2 (Orbitrap-HCD, broad) + 2. intensity — Wilhelmlab/timsTOF-ms2 (timsTOF — instrument domain) + 3. ccs — ionmob CCS parquets (timsTOF) + 4. rt — Chronologer DB (harmonised retention time) + + python -m peptide_property_ng.train.pretrain --cap 50000 --epochs 3 + python -m peptide_property_ng.train.train --init-from runs/pretrain/pretrained.pt \\ + --datasets-glob '/scratch/claudius-proteomics/*' +""" +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch +from torch.utils.data import DataLoader + +from peptide_property_ng.data.ccs_pretrain import prepare_ccs_examples +from peptide_property_ng.data.chronologer_rt import prepare_chronologer_examples +from peptide_property_ng.data.collate import make_collate_fn +from peptide_property_ng.data.hf_intensity import prepare_hf_intensity_examples +from peptide_property_ng.data.sage_dataset import SagePropertyDataset +from peptide_property_ng.eval.metrics import evaluate_split +from peptide_property_ng.losses import MultiTaskLoss +from peptide_property_ng.model.config import PRESETS, instrument_id +from peptide_property_ng.model.multitask import UnifiedPeptidePropertyModel +from peptide_property_ng.modifications.composition import CompositionTable + +_STAGE_METRIC = {"intensity": "intensity_sa", "ccs": "ccs_mae", "rt": "rt_mae"} + + +def _default_stages() -> list[dict]: + """Curriculum: each stage = (name, task, builder taking a per-stage cap).""" + return [ + { + "name": "intensity-orbitrap", "task": "intensity", + "build": lambda cap: prepare_hf_intensity_examples( + "Wilhelmlab/prospect-ptms-ms2", split="train", cap=cap, + instrument=instrument_id("orbitrap")), + }, + { + "name": "intensity-timstof", "task": "intensity", + "build": lambda cap: prepare_hf_intensity_examples( + "Wilhelmlab/timsTOF-ms2", split="train", cap=cap, + instrument=instrument_id("timsTOF Pro")), + }, + { + "name": "ccs-ionmob", "task": "ccs", + "build": lambda cap: prepare_ccs_examples( + cap=cap, instrument=instrument_id("timsTOF Pro")), + }, + { + "name": "rt-chronologer", "task": "rt", + "build": lambda cap: prepare_chronologer_examples(cap=cap), + }, + ] + + +def _train_task_epoch(model, loader, task, loss_fn, optimizer, device) -> float: + model.train() + total, n_batches = 0.0, 0 + for batch in loader: + batch = {k: v.to(device) for k, v in batch.items()} + optimizer.zero_grad() + loss, _ = loss_fn(model(batch, tasks=[task]), batch) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + total += float(loss.detach()) + n_batches += 1 + return total / max(n_batches, 1) + + +def main() -> None: + ap = argparse.ArgumentParser(description="Staged pretraining for the unified model.") + ap.add_argument("--preset", default="small", choices=sorted(PRESETS)) + ap.add_argument("--cap", type=int, default=50000, help="max examples per stage (0 = all)") + ap.add_argument("--epochs", type=int, default=3, help="epochs per stage") + ap.add_argument("--batch-size", type=int, default=128) + ap.add_argument("--lr", type=float, default=3e-4) + ap.add_argument("--val-frac", type=float, default=0.03) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--stages", default="all", help="comma-separated stage names, or 'all'") + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + ap.add_argument("--out", default="runs/pretrain") + args = ap.parse_args() + + torch.manual_seed(args.seed) + out_dir = Path(args.out) + out_dir.mkdir(parents=True, exist_ok=True) + cfg = PRESETS[args.preset] + device = args.device + cap = None if args.cap == 0 else args.cap + + model = UnifiedPeptidePropertyModel(cfg, CompositionTable.load()).to(device) + loss_fn = MultiTaskLoss() + collate = make_collate_fn(cfg.pad_token_id, cfg.max_charge) + print(f"model: '{args.preset}' preset, {model.num_parameters():,} parameters, device={device}") + + stages = _default_stages() + if args.stages != "all": + wanted = {s.strip() for s in args.stages.split(",")} + stages = [s for s in stages if s["name"] in wanted] + + history: list[dict] = [] + for stage in stages: + name, task = stage["name"], stage["task"] + t0 = time.time() + print(f"\n=== stage '{name}' (task={task}) ===", flush=True) + examples = stage["build"](cap) + if len(examples) < args.batch_size * 2: + print(f" only {len(examples)} examples — skipping", flush=True) + continue + + # random hold-out for a stage metric + gen = torch.Generator().manual_seed(args.seed) + perm = torch.randperm(len(examples), generator=gen).tolist() + n_val = max(args.batch_size, int(len(examples) * args.val_frac)) + val_ex = [examples[i] for i in perm[:n_val]] + train_ex = [examples[i] for i in perm[n_val:]] + print(f" {len(train_ex):,} train / {len(val_ex):,} val [prepared in {time.time()-t0:.0f}s]", + flush=True) + + train_loader = DataLoader(SagePropertyDataset(train_ex), batch_size=args.batch_size, + shuffle=True, collate_fn=collate, drop_last=True) + val_loader = DataLoader(SagePropertyDataset(val_ex), batch_size=args.batch_size, + collate_fn=collate) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01) + metric = _STAGE_METRIC[task] + for epoch in range(1, args.epochs + 1): + te = time.time() + train_loss = _train_task_epoch(model, train_loader, task, loss_fn, optimizer, device) + val = evaluate_split(model, val_loader, device) + print(f" epoch {epoch}: train_loss={train_loss:.4f} " + f"val[{metric}]={val[metric]:.4f} [{time.time()-te:.0f}s]", flush=True) + history.append({"stage": name, "epoch": epoch, + "train_loss": train_loss, "val": val}) + + torch.save( + {"model_state_dict": model.state_dict(), "preset": args.preset, "stage": name}, + out_dir / "pretrained.pt", + ) + print(f" saved -> {out_dir}/pretrained.pt", flush=True) + + (out_dir / "pretrain_history.json").write_text(json.dumps(history, indent=2)) + print( + f"\npretraining complete. fine-tune on the campaign data with:\n" + f" python -m peptide_property_ng.train.train " + f"--init-from {out_dir}/pretrained.pt --datasets-glob '/scratch/claudius-proteomics/*'" + ) + + +if __name__ == "__main__": + main() diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py index 2581c8ebe..25996e97f 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py @@ -61,6 +61,8 @@ def main() -> None: ap.add_argument("--weight-decay", type=float, default=0.01) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--patience", type=int, default=4) + ap.add_argument("--init-from", default=None, + help="checkpoint to initialise weights from (e.g. a pretrained.pt)") ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") ap.add_argument("--out", default="runs/prototype") args = ap.parse_args() @@ -98,6 +100,11 @@ def main() -> None: device = args.device model = UnifiedPeptidePropertyModel(cfg, CompositionTable.load()).to(device) + if args.init_from: + ckpt = torch.load(args.init_from, map_location=device) + missing, unexpected = model.load_state_dict(ckpt["model_state_dict"], strict=False) + print(f"initialised from {args.init_from} " + f"(missing={len(missing)}, unexpected={len(unexpected)})") print(f"model: '{args.preset}' preset, {model.num_parameters():,} parameters, device={device}") optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) loss_fn = MultiTaskLoss() From bcd48e9314f95f18694c866146a6d3abad7d3f74 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 21:03:43 +0200 Subject: [PATCH 07/19] fix(peptide-property-ng): address codex pretraining-build review - ccs / rt loaders: guard finite, positive targets -- a NaN no longer produces a NaN loss (codex A1/A2). - train --init-from: reject a preset mismatch with a clear error instead of a shape crash (codex A3). - normalise the CCS / ion-mobility target to ~[0,1] with fixed bounds in both the ionmob (CCS) and Sage (1/K0) loaders, and re-init the physics prior for that range -- the CCS head now sees one consistent target scale across pretraining and fine-tune, removing the negative transfer codex flagged (B1); still no physics conversion. - pretrain: reorder the curriculum (RT, CCS, then intensity last) so the shared encoder ends tuned for the priority task, and save per-stage checkpoints so the fine-tune handoff is selectable (B2); add the Prosit-2025 intensity stage (B3). - hf_intensity: validate the precursor-charge one-hot before argmax (C3). - pretrain: add --chronologer-db / --ccs-glob so the data paths are configurable (for running on monster3). --- .../peptide_property_ng/data/ccs_pretrain.py | 21 ++++++- .../data/chronologer_rt.py | 12 +++- .../peptide_property_ng/data/hf_intensity.py | 5 +- .../peptide_property_ng/data/sage_dataset.py | 12 +++- .../peptide_property_ng/model/heads/ccs.py | 13 +++-- .../src/peptide_property_ng/train/pretrain.py | 57 +++++++++++++------ .../src/peptide_property_ng/train/train.py | 6 ++ 7 files changed, 96 insertions(+), 30 deletions(-) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py b/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py index 4502c3202..3c40e031c 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py @@ -28,6 +28,17 @@ ) _UNIMOD_RE = re.compile(r"\[UNIMOD:\d+\]") +# CCS (Angstrom^2) is min-max normalised to ~[0,1] with fixed bounds so the CCS +# head sees the *same target scale* in pretraining (CCS) and fine-tuning (Sage +# 1/K0, normalised the same way in sage_dataset.py) — this avoids the negative +# transfer of feeding a ~400-scale target into a [0,1]-scale head. +CCS_MIN, CCS_MAX = 200.0, 1200.0 + + +def normalize_ccs(ccs: float) -> float: + """Map a CCS value (Angstrom^2) to ~[0,1].""" + return (float(ccs) - CCS_MIN) / (CCS_MAX - CCS_MIN) + def _tokens_to_proforma(tokens: list[str]) -> str | None: """ionmob ``sequence-tokenized`` list -> a ProForma string, or ``None`` to skip. @@ -95,20 +106,24 @@ def prepare_ccs_examples( stripped = _UNIMOD_RE.sub("", pf) if length != len(stripped) or not 3 <= length <= max_len: continue + ccs_val, mz_val, z = float(ccs[i]), float(mz[i]), int(charge[i]) + if not (np.isfinite(ccs_val) and ccs_val > 0 + and np.isfinite(mz_val) and mz_val > 0 and z >= 1): + continue # skip malformed rows rather than emit a NaN target examples.append( { "accession": "ionmob", "stripped": stripped, "modseq": pf, "tokens": np.asarray(residue_ids, dtype=np.int64), - "charge": int(charge[i]), - "precursor_mz": float(mz[i]), + "charge": z, + "precursor_mz": mz_val, "collision_energy": 0.0, "instrument": int(instrument), "acq_mode": int(acq_mode), # masked placeholder — CCS pretraining runs only the CCS task "intensity_target": np.full((length - 1, 6), -1.0, dtype=np.float32), - "ccs_target": float(ccs[i]), + "ccs_target": normalize_ccs(ccs_val), "ccs_valid": True, "rt_target": float("nan"), "rt_valid": False, diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py b/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py index ae2c677c5..b3efbfeb8 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py @@ -28,7 +28,12 @@ def _normalise_hi(hi: float) -> float: - return (float(hi) - HI_MIN) / (HI_MAX - HI_MIN) + """HI -> ~[0,1], clipped (a few records sit just outside the fixed bounds). + + A non-finite HI propagates as NaN through ``np.clip`` and is filtered by the + caller. + """ + return float(np.clip((float(hi) - HI_MIN) / (HI_MAX - HI_MIN), 0.0, 1.0)) def prepare_chronologer_examples( @@ -68,6 +73,9 @@ def prepare_chronologer_examples( length = len(residue_ids) if length != len(stripped) or not 3 <= length <= max_len: continue + rt = _normalise_hi(hi[start + j]) + if not np.isfinite(rt): + continue # non-finite HI -> useless as an RT example examples.append( { "accession": "Chronologer", @@ -83,7 +91,7 @@ def prepare_chronologer_examples( "intensity_target": np.full((length - 1, 6), -1.0, dtype=np.float32), "ccs_target": float("nan"), "ccs_valid": False, - "rt_target": _normalise_hi(hi[start + j]), + "rt_target": rt, "rt_valid": True, } ) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/hf_intensity.py b/packages/peptide-property-ng/src/peptide_property_ng/data/hf_intensity.py index 719fdef3f..8f380a969 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/hf_intensity.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/hf_intensity.py @@ -74,7 +74,10 @@ def _flush() -> None: stripped = _UNIMOD_RE.sub("", modseq) if length != len(stripped) or not 3 <= length <= max_len: continue - charge = int(np.argmax(row["precursor_charge_onehot"])) + 1 + onehot = row["precursor_charge_onehot"] + if sum(onehot) != 1: + continue # malformed / all-zero one-hot -> skip rather than default to charge 1 + charge = int(np.argmax(onehot)) + 1 try: target = prosit174_to_sites(row["intensities_raw"], length) except ValueError: diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py index 663adba95..e331af641 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py @@ -23,6 +23,16 @@ _PROTON = 1.007276 +# Inverse ion mobility (1/K0) is min-max normalised to ~[0,1] with fixed bounds, +# matching the CCS normalisation in ccs_pretrain.py so the CCS/IM head sees one +# consistent target scale across pretraining (CCS) and fine-tuning (1/K0). +IM_MIN, IM_MAX = 0.5, 1.9 + + +def normalize_inverse_mobility(one_over_k0: float) -> float: + """Map an inverse ion mobility (1/K0) value to ~[0,1].""" + return (float(one_over_k0) - IM_MIN) / (IM_MAX - IM_MIN) + _RES_COLS = [ "psm_id", "peptide", "stripped_peptide", "charge", "calcmass", "aligned_rt", "ion_mobility", "spectrum_q", "is_decoy", "rank", @@ -166,7 +176,7 @@ def prepare_examples( "instrument": int(instrument), "acq_mode": int(acq_mode), "intensity_target": target, - "ccs_target": im, + "ccs_target": normalize_inverse_mobility(im), "ccs_valid": bool(np.isfinite(im) and im > 0.0), "rt_target": rt, "rt_valid": bool(np.isfinite(rt)), diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/heads/ccs.py b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/ccs.py index cfbd2f6a2..97c6002fa 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/heads/ccs.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/ccs.py @@ -1,8 +1,10 @@ """Ion-mobility head — a per-charge physics prior plus a learned correction. -Predicts inverse reduced ion mobility (1/K0), the raw quantity Sage reports in -its ``ion_mobility`` column — this avoids baking a Mason-Schamp CCS conversion -(with its gas/temperature constants) into the prototype. +Predicts a min-max-normalised ion-mobility target in ~[0,1]. Both data sources +are normalised to that same scale with fixed bounds (campaign data = Sage +``ion_mobility`` 1/K0; pretraining data = ionmob CCS — see the data loaders), +so the head sees one consistent target scale and no Mason-Schamp CCS<->1/K0 +conversion is needed. The physics prior ``1/K0 ~ slope[z]*sqrt(m/z) + intercept[z]`` is ported from the production ``SquareRootProjectionLayer`` (imspy_predictors); slopes and @@ -21,9 +23,10 @@ class SqrtMzProjection(nn.Module): """Per-charge physics prior: ``property ~ slope[z]*sqrt(m/z) + intercept[z]``.""" - def __init__(self, max_charge: int, slope_init: float = 0.03, intercept_init: float = 0.60): + def __init__(self, max_charge: int, slope_init: float = 0.0, intercept_init: float = 0.5): super().__init__() - # Index by charge state directly (row 0 unused). Init at a 1/K0 ballpark. + # Index by charge state directly (row 0 unused). Init at the centre of the + # normalised [0,1] target range; the per-charge slopes are learnable. self.slopes = nn.Parameter(torch.full((max_charge + 1,), float(slope_init))) self.intercepts = nn.Parameter(torch.full((max_charge + 1,), float(intercept_init))) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py index 197747121..c8e4a1b76 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py @@ -35,30 +35,46 @@ _STAGE_METRIC = {"intensity": "intensity_sa", "ccs": "ccs_mae", "rt": "rt_mae"} -def _default_stages() -> list[dict]: - """Curriculum: each stage = (name, task, builder taking a per-stage cap).""" +def _default_stages( + chronologer_db: str | None = None, ccs_glob: str | None = None +) -> list[dict]: + """Curriculum: each stage = name, task, and a builder taking a per-stage cap. + + The order is deliberate: the auxiliary tasks (RT, CCS) run first and + intensity runs last, so the shared encoder ends tuned for the priority task + rather than left drifted toward RT — a catastrophic-forgetting mitigation. + Within intensity, the broad Orbitrap corpora precede the timsTOF one. + """ + rt_kw = {} if chronologer_db is None else {"db_path": chronologer_db} + ccs_kw = {} if ccs_glob is None else {"data_glob": ccs_glob} return [ { - "name": "intensity-orbitrap", "task": "intensity", + "name": "rt-chronologer", "task": "rt", + "build": lambda cap: prepare_chronologer_examples(cap=cap, **rt_kw), + }, + { + "name": "ccs-ionmob", "task": "ccs", + "build": lambda cap: prepare_ccs_examples( + cap=cap, instrument=instrument_id("timsTOF Pro"), **ccs_kw), + }, + { + "name": "intensity-prospect", "task": "intensity", "build": lambda cap: prepare_hf_intensity_examples( "Wilhelmlab/prospect-ptms-ms2", split="train", cap=cap, instrument=instrument_id("orbitrap")), }, + { + "name": "intensity-prosit2025", "task": "intensity", + "build": lambda cap: prepare_hf_intensity_examples( + "Wilhelmlab/Prosit-2025-lac-ms2", split="train", cap=cap, + instrument=instrument_id("orbitrap")), + }, { "name": "intensity-timstof", "task": "intensity", "build": lambda cap: prepare_hf_intensity_examples( "Wilhelmlab/timsTOF-ms2", split="train", cap=cap, instrument=instrument_id("timsTOF Pro")), }, - { - "name": "ccs-ionmob", "task": "ccs", - "build": lambda cap: prepare_ccs_examples( - cap=cap, instrument=instrument_id("timsTOF Pro")), - }, - { - "name": "rt-chronologer", "task": "rt", - "build": lambda cap: prepare_chronologer_examples(cap=cap), - }, ] @@ -87,6 +103,10 @@ def main() -> None: ap.add_argument("--val-frac", type=float, default=0.03) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--stages", default="all", help="comma-separated stage names, or 'all'") + ap.add_argument("--chronologer-db", default=None, + help="path to Chronologer_DB_*.gz (default: the package's known local path)") + ap.add_argument("--ccs-glob", default=None, + help="glob for ionmob *_unique_unimod*.parquet (default: known local path)") ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") ap.add_argument("--out", default="runs/pretrain") args = ap.parse_args() @@ -103,7 +123,7 @@ def main() -> None: collate = make_collate_fn(cfg.pad_token_id, cfg.max_charge) print(f"model: '{args.preset}' preset, {model.num_parameters():,} parameters, device={device}") - stages = _default_stages() + stages = _default_stages(args.chronologer_db, args.ccs_glob) if args.stages != "all": wanted = {s.strip() for s in args.stages.split(",")} stages = [s for s in stages if s["name"] in wanted] @@ -142,11 +162,12 @@ def main() -> None: history.append({"stage": name, "epoch": epoch, "train_loss": train_loss, "val": val}) - torch.save( - {"model_state_dict": model.state_dict(), "preset": args.preset, "stage": name}, - out_dir / "pretrained.pt", - ) - print(f" saved -> {out_dir}/pretrained.pt", flush=True) + # Save a per-stage checkpoint (so the campaign fine-tune can pick the + # handoff point, not just the post-RT state) and the running pretrained.pt. + ckpt = {"model_state_dict": model.state_dict(), "preset": args.preset, "stage": name} + torch.save(ckpt, out_dir / f"after-{name}.pt") + torch.save(ckpt, out_dir / "pretrained.pt") + print(f" saved -> {out_dir}/after-{name}.pt (+ pretrained.pt)", flush=True) (out_dir / "pretrain_history.json").write_text(json.dumps(history, indent=2)) print( diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py index 25996e97f..bf84d8eb9 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py @@ -102,6 +102,12 @@ def main() -> None: model = UnifiedPeptidePropertyModel(cfg, CompositionTable.load()).to(device) if args.init_from: ckpt = torch.load(args.init_from, map_location=device) + ckpt_preset = ckpt.get("preset") + if ckpt_preset is not None and ckpt_preset != args.preset: + raise SystemExit( + f"--init-from checkpoint is preset '{ckpt_preset}', but --preset is " + f"'{args.preset}'; they must match (the architectures differ)." + ) missing, unexpected = model.load_state_dict(ckpt["model_state_dict"], strict=False) print(f"initialised from {args.init_from} " f"(missing={len(missing)}, unexpected={len(unexpected)})") From 9bc83f11e49351f2941aa706d6418c2ed9ccefee Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 21:11:28 +0200 Subject: [PATCH 08/19] fix(peptide-property-ng): final-review polish - clip the normalised CCS / ion-mobility targets to [0,1] (codex: values beyond the fixed bounds would otherwise leave the [0,1] range). - pretrain: peptide-level stage hold-out (a peptide no longer crosses train/val within a stage) instead of a random row split. Final codex review: build verified sound for capped runs. Remaining note -- streaming/sharded stage loading for an uncapped full-scale run -- is a follow-up; a finite --cap avoids it. --- .../src/peptide_property_ng/data/ccs_pretrain.py | 4 ++-- .../src/peptide_property_ng/data/sage_dataset.py | 4 ++-- .../src/peptide_property_ng/train/pretrain.py | 14 ++++++++------ 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py b/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py index 3c40e031c..329914934 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py @@ -36,8 +36,8 @@ def normalize_ccs(ccs: float) -> float: - """Map a CCS value (Angstrom^2) to ~[0,1].""" - return (float(ccs) - CCS_MIN) / (CCS_MAX - CCS_MIN) + """Map a CCS value (Angstrom^2) to [0,1], clipped at the fixed bounds.""" + return float(np.clip((float(ccs) - CCS_MIN) / (CCS_MAX - CCS_MIN), 0.0, 1.0)) def _tokens_to_proforma(tokens: list[str]) -> str | None: diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py index e331af641..b85b7127f 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py @@ -30,8 +30,8 @@ def normalize_inverse_mobility(one_over_k0: float) -> float: - """Map an inverse ion mobility (1/K0) value to ~[0,1].""" - return (float(one_over_k0) - IM_MIN) / (IM_MAX - IM_MIN) + """Map an inverse ion mobility (1/K0) value to [0,1], clipped at the fixed bounds.""" + return float(np.clip((float(one_over_k0) - IM_MIN) / (IM_MAX - IM_MIN), 0.0, 1.0)) _RES_COLS = [ "psm_id", "peptide", "stripped_peptide", "charge", "calcmass", diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py index c8e4a1b76..556cfb444 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py @@ -26,6 +26,7 @@ from peptide_property_ng.data.collate import make_collate_fn from peptide_property_ng.data.hf_intensity import prepare_hf_intensity_examples from peptide_property_ng.data.sage_dataset import SagePropertyDataset +from peptide_property_ng.data.splits import peptide_split from peptide_property_ng.eval.metrics import evaluate_split from peptide_property_ng.losses import MultiTaskLoss from peptide_property_ng.model.config import PRESETS, instrument_id @@ -138,12 +139,13 @@ def main() -> None: print(f" only {len(examples)} examples — skipping", flush=True) continue - # random hold-out for a stage metric - gen = torch.Generator().manual_seed(args.seed) - perm = torch.randperm(len(examples), generator=gen).tolist() - n_val = max(args.batch_size, int(len(examples) * args.val_frac)) - val_ex = [examples[i] for i in perm[:n_val]] - train_ex = [examples[i] for i in perm[n_val:]] + # peptide-level hold-out — a peptide never crosses train/val within a + # stage (stage metrics are pretraining diagnostics, not the final eval). + train_ex, val_ex = [], [] + for ex in examples: + bucket = peptide_split(ex["stripped"], val_frac=args.val_frac, + test_frac=0.0, seed=args.seed) + (val_ex if bucket == "val" else train_ex).append(ex) print(f" {len(train_ex):,} train / {len(val_ex):,} val [prepared in {time.time()-t0:.0f}s]", flush=True) From d3bef364578fb3503c29a41027d18b9814b06a77 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 21:42:28 +0200 Subject: [PATCH 09/19] refactor(peptide-property-ng): vendor the delta-mass to UNIMOD conversion The Chronologer RT loader and the Sage loader converted [+mass] peptide notation via sagepy_rescore._parse_sage_peptide, which imports sagepy_rescore.pipeline -- a heavy module that pulls the whole search/fdr stack and is version-fragile (its sagepy.core.fdr imports break against older sagepy installs, e.g. monster3's sagepy 0.4.4). Vendor just the conversion into data/mass_to_unimod.py (_BRACKET_MOD, _mass_to_unimod + tables, parse_delta_mass_peptide) -- copied verbatim from sagepy_rescore; it needs only sagepy.core.unimod and sagepy_connector.py_unimod. Verified byte-identical to sagepy_rescore._parse_sage_peptide. Drops the sagepy_rescore dependency. --- .../data/chronologer_rt.py | 5 +- .../data/mass_to_unimod.py | 94 +++++++++++++++++++ .../peptide_property_ng/data/sage_dataset.py | 5 +- 3 files changed, 100 insertions(+), 4 deletions(-) create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/mass_to_unimod.py diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py b/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py index b3efbfeb8..038ed8fe2 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py @@ -48,7 +48,8 @@ def prepare_chronologer_examples( """Load the Chronologer DB into RT-pretraining example dicts.""" import pandas as pd from imspy_predictors.utilities.tokenizers import ProformaTokenizer - from sagepy_rescore.sage_loader import _parse_sage_peptide + + from peptide_property_ng.data.mass_to_unimod import parse_delta_mass_peptide df = pd.read_csv(db_path, sep="\t", compression="gzip", usecols=["PeptideModSeq", "HI"]) if cap is not None and len(df) > cap: @@ -57,7 +58,7 @@ def prepare_chronologer_examples( hi = df["HI"].to_numpy(dtype=np.float32) tok = ProformaTokenizer.with_defaults() - parsed = [_parse_sage_peptide(p) for p in peptides] # [+mass] -> UNIMOD + parsed = [parse_delta_mass_peptide(p) for p in peptides] # [+mass] -> UNIMOD examples: list[dict] = [] chunk = 20000 # tokenise in chunks (batch tokenisation right-pads to the chunk max) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/mass_to_unimod.py b/packages/peptide-property-ng/src/peptide_property_ng/data/mass_to_unimod.py new file mode 100644 index 000000000..c06167d26 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/mass_to_unimod.py @@ -0,0 +1,94 @@ +"""Delta-mass peptide -> UNIMOD conversion. + +Sage and the Chronologer DB write peptides with delta-mass brackets +(``MC[+57.0216]PEPTIDE``); the tokenizer needs UNIMOD form +(``MC[UNIMOD:4]PEPTIDE``). The mass->UNIMOD step is residue-specificity aware +(the residue immediately before the bracket), with a small hardcoded table for +the common search modifications and the quantized sagepy table as the general +case. + +This logic is **vendored verbatim** from ``sagepy_rescore`` (``sage_loader. +_parse_sage_peptide`` + ``pipeline._mass_to_unimod`` and its tables) so this +package does not import ``sagepy_rescore.pipeline`` — a heavy module whose +unrelated ``sagepy.core.fdr`` imports are version-fragile across sagepy +releases. The conversion itself is unchanged; it needs only ``sagepy.core.unimod`` +and ``sagepy_connector.py_unimod``. +""" +from __future__ import annotations + +import re + +from sagepy.core.unimod import unimod_to_mass +from sagepy_connector import py_unimod + +# Sage / Chronologer write only positive delta masses, e.g. "[+57.0216]". +_BRACKET_MOD = re.compile(r"\[\+([0-9.]+)\]") + + +def _build_mass_to_unimod() -> dict[float, str]: + out: dict[float, str] = {} + for unimod_id, mass in unimod_to_mass().items(): + out[round(float(mass), 4)] = unimod_id + return out + + +_MASS_TO_UNIMOD = _build_mass_to_unimod() + +# Residue-specificity-disambiguated entries for the common search modifications. +_COMMON_SAGE_MASS_TO_UNIMOD = { + ("C", py_unimod.quanzied_mass(57.0216)): "[UNIMOD:4]", + ("M", py_unimod.quanzied_mass(15.9949)): "[UNIMOD:35]", + ("[", py_unimod.quanzied_mass(42.0)): "[UNIMOD:1]", +} + + +def _mass_to_unimod(value, specificity: str | None = None): + """Resolve a delta mass (+ optional residue specificity) to a UNIMOD token.""" + if isinstance(value, list): + return [_mass_to_unimod(v, specificity=specificity) for v in value] + if isinstance(value, str): + return value + if isinstance(value, int): + return f"[UNIMOD:{value}]" + + quantized = py_unimod.quanzied_mass(float(value)) + common = _COMMON_SAGE_MASS_TO_UNIMOD.get((specificity, quantized)) + if common is not None: + return common + + candidates = py_unimod.quantized_mass_to_unimod_candidates().get(quantized, []) + if candidates: + return candidates[0] + + key = round(float(value), 4) + if key in _MASS_TO_UNIMOD: + return _MASS_TO_UNIMOD[key] + raise ValueError( + f"Mass {value!r} for specificity {specificity!r} has no UNIMOD entry " + "in sagepy's quantized table." + ) + + +def parse_delta_mass_peptide(peptide: str) -> tuple[str, str]: + """Convert a delta-mass peptide string to ``(unmodified, UNIMOD-modified)``. + + Example: ``IIPGFMC[+57.0216]QGGDFTR`` -> + ``("IIPGFMCQGGDFTR", "IIPGFMC[UNIMOD:4]QGGDFTR")``. A delta mass with no + UNIMOD match is left as the raw ``[+mass]`` bracket (the caller skips it). + """ + out_parts: list[str] = [] + last_end = 0 + for m in _BRACKET_MOD.finditer(peptide): + start = m.start() + out_parts.append(peptide[last_end:start]) + specificity = "[" if start == 0 else peptide[start - 1] + try: + unimod = _mass_to_unimod(float(m.group(1)), specificity=specificity) + except ValueError: + unimod = m.group(0) # unknown mass -> leave the raw bracket + out_parts.append(unimod) + last_end = m.end() + out_parts.append(peptide[last_end:]) + modified = "".join(out_parts).replace("-", "") + unmodified = _BRACKET_MOD.sub("", peptide).replace("-", "") + return unmodified, modified diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py index b85b7127f..d533a2f8d 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py @@ -82,7 +82,8 @@ def prepare_examples( """Load and prepare one dataset's PSMs into a list of example dicts.""" from imspy_predictors.intensity.predictors import observed_fragments_to_intensity_target from imspy_predictors.utilities.tokenizers import ProformaTokenizer - from sagepy_rescore.sage_loader import _parse_sage_peptide + + from peptide_property_ng.data.mass_to_unimod import parse_delta_mass_peptide sage_dir = Path(sage_dir) res_path = sage_dir / "results.sage.parquet" @@ -127,7 +128,7 @@ def prepare_examples( # Sage emits delta-mass peptides; convert to UNIMOD, then tokenise in one batch. # Batch tokenisation right-pads to the batch-max length, so the attention # mask is needed to recover each peptide's true token count. - parsed = [_parse_sage_peptide(p) for p in df["peptide"]] + parsed = [parse_delta_mass_peptide(p) for p in df["peptide"]] tok = ProformaTokenizer.with_defaults() _encoded = tok([m for _, m in parsed]) token_ids, attn_mask = _encoded["input_ids"], _encoded["attention_mask"] From b330850adff51dfa68da0d8a3996dfb3bc1e3d42 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 23:20:48 +0200 Subject: [PATCH 10/19] feat(peptide-property-ng): reserve embedding headroom for new instruments/modes The instrument and acquisition-mode embeddings were sized exactly to the current vocabularies (9 / 7). Adding a new instrument later would change the embedding shape and break load_state_dict on every checkpoint. Size the embeddings to a fixed capacity instead -- 32 instrument slots (9 used), 16 acquisition-mode slots (7 used). A new entry appended to INSTRUMENTS / ACQUISITION_MODES takes the next free id with the embedding shape unchanged: existing checkpoints stay loadable and the new id starts from an untrained row, ready to fine-tune. __post_init__ asserts the vocabularies stay within capacity. --- .../src/peptide_property_ng/model/config.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py index 229bf1a9c..429ef80c0 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py @@ -87,8 +87,13 @@ class ModelConfig: # conditioning max_charge: int = 8 - n_instruments: int = len(INSTRUMENTS) - n_acq_modes: int = len(ACQUISITION_MODES) + # Embedding *capacity* for instrument / acquisition-mode ids, deliberately + # larger than the current vocabularies. A new instrument or mode appended to + # INSTRUMENTS / ACQUISITION_MODES later gets the next free id without + # resizing the embedding -- so existing checkpoints stay loadable and the + # new id simply starts from an untrained row, ready to fine-tune. + n_instruments: int = 32 # 9 ids in use (see INSTRUMENTS); the rest reserved + n_acq_modes: int = 16 # 7 ids in use (see ACQUISITION_MODES); the rest reserved # heads n_ion_channels: int = 6 # b/y ions x fragment charges 1-3 (Prosit-compatible) @@ -97,6 +102,14 @@ class ModelConfig: def __post_init__(self) -> None: if self.d_model % self.n_heads != 0: raise ValueError(f"d_model {self.d_model} not divisible by n_heads {self.n_heads}") + if len(INSTRUMENTS) > self.n_instruments: + raise ValueError( + f"{len(INSTRUMENTS)} instruments exceed embedding capacity {self.n_instruments}" + ) + if len(ACQUISITION_MODES) > self.n_acq_modes: + raise ValueError( + f"{len(ACQUISITION_MODES)} acquisition modes exceed capacity {self.n_acq_modes}" + ) # First-prototype preset — roughly on par with the production BASE encoder so From 92921876db4191e0679519090aa35870a0eb63e2 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 22 May 2026 23:32:03 +0200 Subject: [PATCH 11/19] feat(peptide-property-ng): bump embedding headroom to 64 instruments / 32 acq modes Plenty of room to add new instrument models or acquisition modes later without resizing the embedding (so existing checkpoints keep loading). --- .../src/peptide_property_ng/model/config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py index 429ef80c0..8947f96f9 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py @@ -92,8 +92,8 @@ class ModelConfig: # INSTRUMENTS / ACQUISITION_MODES later gets the next free id without # resizing the embedding -- so existing checkpoints stay loadable and the # new id simply starts from an untrained row, ready to fine-tune. - n_instruments: int = 32 # 9 ids in use (see INSTRUMENTS); the rest reserved - n_acq_modes: int = 16 # 7 ids in use (see ACQUISITION_MODES); the rest reserved + n_instruments: int = 64 # 9 ids in use (see INSTRUMENTS); the rest reserved + n_acq_modes: int = 32 # 7 ids in use (see ACQUISITION_MODES); the rest reserved # heads n_ion_channels: int = 6 # b/y ions x fragment charges 1-3 (Prosit-compatible) From 05391cffebdef9f19d1b02c307a8c1159279188e Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Sat, 23 May 2026 08:22:42 +0200 Subject: [PATCH 12/19] feat(peptide-property-ng): wire per-dataset instrument/acq from timstof_catalog The Sage loader passed `instrument=unknown` for every campaign sample, which silently bypassed the encoder's metadata conditioning -- the very mechanism the multi-instrument pretraining curriculum was built to exercise. Wire per-dataset `instrument_model` and `acquisition_mode` through a new `load_catalog_metadata(timstof_catalog.tsv)` -> per-accession dispatch in `build_split_datasets`; expose `--catalog` on `train.py`. Also extend ACQUISITION_MODES with `PASEF` and `single-cell` (the catalog's dominant values that the previous enum did not cover). All 25 currently-processed campaign datasets resolve to real (instrument, acq) ids -- spanning Pro / Pro 2 / HT / SCP / fleX / plain timsTOF. --- .../peptide_property_ng/data/sage_dataset.py | 44 ++++++++++++++++++- .../src/peptide_property_ng/model/config.py | 2 + .../src/peptide_property_ng/train/train.py | 5 ++- 3 files changed, 48 insertions(+), 3 deletions(-) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py index d533a2f8d..2b56914b2 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py @@ -7,6 +7,7 @@ """ from __future__ import annotations +import csv import glob from collections import defaultdict from pathlib import Path @@ -199,6 +200,27 @@ def __getitem__(self, i: int) -> dict: return self.examples[i] +def load_catalog_metadata(catalog_path: str | Path) -> dict[str, tuple[int, int]]: + """Read ``timstof_catalog.tsv`` -> ``{accession: (instrument_id, acq_id)}``. + + Unknown / missing values map to id 0 via ``instrument_id`` / ``acquisition_id``. + """ + from peptide_property_ng.model.config import acquisition_id, instrument_id + + out: dict[str, tuple[int, int]] = {} + with open(catalog_path) as f: + reader = csv.DictReader(f, delimiter="\t") + for row in reader: + acc = (row.get("accession") or "").strip() + if not acc: + continue + out[acc] = ( + instrument_id(row.get("instrument_model")), + acquisition_id(row.get("acquisition_mode")), + ) + return out + + def build_split_datasets( sage_dirs: list[str | Path], *, @@ -206,12 +228,30 @@ def build_split_datasets( seed: int = 0, val_frac: float = 0.1, test_frac: float = 0.1, + catalog_path: str | Path | None = None, **prepare_kwargs, ) -> dict[str, SagePropertyDataset]: - """Prepare every dataset and split peptide-level into train / val / test.""" + """Prepare every dataset and split peptide-level into train / val / test. + + If ``catalog_path`` is given (a ``timstof_catalog.tsv``), per-dataset + ``instrument`` and ``acq_mode`` ids are looked up by accession and passed to + ``prepare_examples`` — so the encoder's metadata conditioning actually sees + the right instrument instead of falling back to ``unknown``. Otherwise the + loader's defaults (0/unknown) are used. + """ + metadata: dict[str, tuple[int, int]] = ( + load_catalog_metadata(catalog_path) if catalog_path else {} + ) examples: list[dict] = [] for sage_dir in sage_dirs: - examples.extend(prepare_examples(sage_dir, cap=cap, seed=seed, **prepare_kwargs)) + sage_dir_p = Path(sage_dir) + accession = _accession(sage_dir_p) + kw = dict(prepare_kwargs) + if accession in metadata: + instr, acq = metadata[accession] + kw.setdefault("instrument", instr) + kw.setdefault("acq_mode", acq) + examples.extend(prepare_examples(sage_dir_p, cap=cap, seed=seed, **kw)) buckets: dict[str, list[dict]] = {"train": [], "val": [], "test": []} for ex in examples: diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py index 8947f96f9..d1429663f 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py @@ -34,6 +34,8 @@ "diaPASEF", "PRM", "MALDI", + "PASEF", # catalog's dominant value (716 deposits) + "single-cell", # catalog: 70 deposits ) INSTRUMENT_TO_ID: dict[str, int] = {n: i for i, n in enumerate(INSTRUMENTS)} ACQ_TO_ID: dict[str, int] = {n: i for i, n in enumerate(ACQUISITION_MODES)} diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py index bf84d8eb9..a2e620882 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py @@ -49,6 +49,8 @@ def train_epoch(model, loader, loss_fn, optimizer, device) -> dict[str, float]: def main() -> None: ap = argparse.ArgumentParser(description="Train the unified peptide property model.") ap.add_argument("--datasets-glob", default="/scratch/claudius-proteomics/*") + ap.add_argument("--catalog", default=None, + help="path to timstof_catalog.tsv for per-dataset instrument/acq conditioning") ap.add_argument("--preset", default="small", choices=sorted(PRESETS)) ap.add_argument("--comp-fusion", default=None, choices=["add", "gate", "token_only", "composition_only"], @@ -79,7 +81,8 @@ def main() -> None: print(f"[{time.strftime('%H:%M:%S')}] preparing examples (cap {args.cap}/dataset) ...", flush=True) t0 = time.time() - splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed) + splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed, + catalog_path=args.catalog) for name, ds in splits.items(): print(f" {name}: {len(ds):,} examples") print(f" prepared in {time.time() - t0:.0f}s") From 503ce6797adce444d74bf83656067655779b13c9 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Sat, 23 May 2026 08:35:31 +0200 Subject: [PATCH 13/19] experiment(peptide-property-ng): CE default 0.26 (timsTOF-ms2 median) for Sage examples The fine-tune intensity SA capped at 0.66 because every campaign example was fed CE=0 -- squarely out-of-distribution for the head's CE FiLM, which was pretrained around the timsTOF-ms2 median normalised CE (~0.26). Default the Sage loader's CE to that in-distribution value; expose --default-ce on train.py so we can sweep it later. --- .../src/peptide_property_ng/data/sage_dataset.py | 6 +++++- .../src/peptide_property_ng/train/train.py | 5 ++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py index 2b56914b2..0fb2da44c 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py @@ -79,6 +79,10 @@ def prepare_examples( seed: int = 0, instrument: int = 0, acq_mode: int = 0, + # Sage doesn't surface per-PSM CE for timsTOF diaPASEF; the intensity head's CE + # FiLM was tuned by pretraining around the timsTOF-ms2 median (~0.26 normalised). + # Default to that here so the head is fed an in-distribution CE rather than 0. + default_ce: float = 0.26, ) -> list[dict]: """Load and prepare one dataset's PSMs into a list of example dicts.""" from imspy_predictors.intensity.predictors import observed_fragments_to_intensity_target @@ -174,7 +178,7 @@ def prepare_examples( "tokens": np.asarray(residue_ids, dtype=np.int64), "charge": charge, "precursor_mz": mz, - "collision_energy": 0.0, # not in Sage output for timsTOF diaPASEF + "collision_energy": float(default_ce), "instrument": int(instrument), "acq_mode": int(acq_mode), "intensity_target": target, diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py index a2e620882..7c361df09 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py @@ -51,6 +51,9 @@ def main() -> None: ap.add_argument("--datasets-glob", default="/scratch/claudius-proteomics/*") ap.add_argument("--catalog", default=None, help="path to timstof_catalog.tsv for per-dataset instrument/acq conditioning") + ap.add_argument("--default-ce", type=float, default=0.26, + help="default normalized collision energy for Sage PSMs that lack one " + "(~0.26 = timsTOF-ms2 median, in-distribution for the pretrained CE FiLM)") ap.add_argument("--preset", default="small", choices=sorted(PRESETS)) ap.add_argument("--comp-fusion", default=None, choices=["add", "gate", "token_only", "composition_only"], @@ -82,7 +85,7 @@ def main() -> None: print(f"[{time.strftime('%H:%M:%S')}] preparing examples (cap {args.cap}/dataset) ...", flush=True) t0 = time.time() splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed, - catalog_path=args.catalog) + catalog_path=args.catalog, default_ce=args.default_ce) for name, ds in splits.items(): print(f" {name}: {len(ds):,} examples") print(f" prepared in {time.time() - t0:.0f}s") From 6b8809ee8384f3b29f5ad896ff47823b1ca326a6 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Sat, 23 May 2026 09:13:01 +0200 Subject: [PATCH 14/19] fix(peptide-property-ng): masked SA -> target > 0 (canonical Prosit/v4 convention) The intensity loss / eval masked target < 0 (impossible-channel only), which treated target == 0 as a real no-peak label. That is correct for densely annotated PROSPECT targets but wrong for Sage matched_fragments, where 0 means unmatched / unknown, not real zero. v4's finetune_dia_pasef_intensity.py canonical-Prosit loss masks target > 0; this aligns the package to that. This explains a sizable chunk of the campaign fine-tune intensity-SA gap: re-evaluating the existing best.pt with the v4 metric already lifts SA from 0.673 to 0.733 with no retraining. The remaining ~0.10 gap to v4's 0.83 will need a retrain with the corrected training loss. --- .../src/peptide_property_ng/losses.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/losses.py b/packages/peptide-property-ng/src/peptide_property_ng/losses.py index 3d9ba871a..fb522c8e2 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/losses.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/losses.py @@ -10,12 +10,19 @@ def masked_spectral_angle(pred: torch.Tensor, target: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: """Per-sample spectral-angle distance in ``[0, 1]`` (0 = identical spectra). - ``pred`` / ``target`` are ``(B, sites, channels)``. Entries where - ``target < 0`` are masked out (impossible / padded fragment channels). + Canonical Prosit / Sage-fine-tune loss: mask ``target > 0`` (only observed + peaks contribute). Zeros are treated as "unmatched / unknown" rather than + real "no peak" labels — which is the correct interpretation for Sage + ``matched_fragments`` (Sage only reports peaks it matched) and matches the + production v4 / imspy_predictors `masked_spectral_distance` convention. + + Pretraining on PROSPECT-style densely-annotated targets is unaffected in + practice — there the observable b/y positions almost all carry positive + intensity, so `> 0` and `>= 0` masks coincide. """ - mask = (target >= 0).float() + mask = (target > 0).float() p = (pred * mask).reshape(pred.shape[0], -1) - t = (target.clamp(min=0.0) * mask).reshape(target.shape[0], -1) + t = (target * mask).reshape(target.shape[0], -1) p = F.normalize(p, dim=1, eps=eps) t = F.normalize(t, dim=1, eps=eps) cos = (p * t).sum(dim=1).clamp(-1.0 + eps, 1.0 - eps) From cafc9aa8809c5e75c827056345c8fe973cc1aa9f Mon Sep 17 00:00:00 2001 From: David Teschner Date: Sun, 24 May 2026 08:58:06 +0200 Subject: [PATCH 15/19] feat(peptide-property-ng): pooled intensity head + per-PSM CE calibration Two-prong fine-tune-time levers explored on the SMALL preset campaign data. * Pooled (v4-style) intensity head: attention-pool residue latents, concat with charge/CE/instrument side-embeddings, MLP to Prosit-174 grid. Sits behind cfg.intensity_head = 'site' (default, local FiLM) or 'pooled' (new). Plumbed through train.py and pretrain.py via --intensity-head. * Per-PSM CE calibration: new calibrate_ce.py runs an inference-time CE grid sweep against a frozen checkpoint and writes per-PSM optimal CEs to parquet keyed by (accession, psm_id). train.py and evaluate_optimal_nce.py accept --ce-calibration to use those calibrated CEs instead of the flat default; added load_ce_calibration to sage_dataset. * Loss / metric: masked_spectral_angle now slices pred to the target's site dimension when pred is longer (pooled head outputs fixed (B,29,n_ion)). * train.py: --tasks and --freeze-encoder for specialization sweeps; the eval helper also takes a tasks list. evaluate_split short-circuits skipped tasks. Campaign results on the 25 timsTOF datasets: canonical baseline test SA 0.7597 canonical + inference-time optimal-CE 0.7945 round-1 calibrated fine-tune 0.7867 round-1 + inference-time optimal-CE 0.8058 round-2 calibrated fine-tune 0.8090 round-2 + inference-time optimal-CE 0.8170 pooled head + round-2 calibration (random 0.7510 (head-init disadvantage) pooled-head init, encoder warm-started) --- .../peptide_property_ng/data/sage_dataset.py | 34 +++- .../eval/evaluate_optimal_nce.py | 158 +++++++++++++++ .../src/peptide_property_ng/eval/metrics.py | 8 +- .../src/peptide_property_ng/losses.py | 5 + .../src/peptide_property_ng/model/config.py | 3 + .../model/heads/intensity_pooled.py | 101 ++++++++++ .../peptide_property_ng/model/multitask.py | 28 ++- .../peptide_property_ng/train/calibrate_ce.py | 183 ++++++++++++++++++ .../src/peptide_property_ng/train/pretrain.py | 7 + .../src/peptide_property_ng/train/train.py | 50 ++++- 10 files changed, 557 insertions(+), 20 deletions(-) create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate_optimal_nce.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/model/heads/intensity_pooled.py create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/train/calibrate_ce.py diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py index 0fb2da44c..c1540a0db 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/sage_dataset.py @@ -83,6 +83,7 @@ def prepare_examples( # FiLM was tuned by pretraining around the timsTOF-ms2 median (~0.26 normalised). # Default to that here so the head is fed an in-distribution CE rather than 0. default_ce: float = 0.26, + ce_calibration: dict[tuple[str, int], float] | None = None, ) -> list[dict]: """Load and prepare one dataset's PSMs into a list of example dicts.""" from imspy_predictors.intensity.predictors import observed_fragments_to_intensity_target @@ -154,7 +155,8 @@ def prepare_examples( charge = int(df["charge"][k]) mz = (float(df["calcmass"][k]) + charge * _PROTON) / charge - cols = by_psm.get(df["psm_id"][k]) + psm_id = int(df["psm_id"][k]) + cols = by_psm.get(psm_id) if not cols or not cols[0]: continue # no matched fragments -> no intensity target # Encode fragments with the proven imspy encoder (Sage fragments -> @@ -170,15 +172,21 @@ def prepare_examples( rt = df["aligned_rt"][k] im = float(im) if im is not None else float("nan") rt = float(rt) if rt is not None else float("nan") + ce_val = float(default_ce) + if ce_calibration is not None: + cal = ce_calibration.get((accession, psm_id)) + if cal is not None: + ce_val = float(cal) examples.append( { "accession": accession, + "psm_id": psm_id, "stripped": stripped, "modseq": modseq, "tokens": np.asarray(residue_ids, dtype=np.int64), "charge": charge, "precursor_mz": mz, - "collision_energy": float(default_ce), + "collision_energy": ce_val, "instrument": int(instrument), "acq_mode": int(acq_mode), "intensity_target": target, @@ -225,6 +233,20 @@ def load_catalog_metadata(catalog_path: str | Path) -> dict[str, tuple[int, int] return out +def load_ce_calibration(parquet_path: str | Path) -> dict[tuple[str, int], float]: + """Read a CE calibration parquet -> ``{(accession, psm_id): calibrated_ce}``. + + The parquet must have columns ``accession`` (string), ``psm_id`` (int64), + ``calibrated_ce`` (float). Missing keys at lookup time fall back to the + dataset's ``default_ce`` — so partial calibrations are safe. + """ + table = pq.read_table(parquet_path) + accs = table["accession"].to_pylist() + pids = table["psm_id"].to_pylist() + ces = table["calibrated_ce"].to_pylist() + return {(a, int(p)): float(c) for a, p, c in zip(accs, pids, ces)} + + def build_split_datasets( sage_dirs: list[str | Path], *, @@ -233,6 +255,7 @@ def build_split_datasets( val_frac: float = 0.1, test_frac: float = 0.1, catalog_path: str | Path | None = None, + ce_calibration: dict[tuple[str, int], float] | None = None, **prepare_kwargs, ) -> dict[str, SagePropertyDataset]: """Prepare every dataset and split peptide-level into train / val / test. @@ -242,6 +265,10 @@ def build_split_datasets( ``prepare_examples`` — so the encoder's metadata conditioning actually sees the right instrument instead of falling back to ``unknown``. Otherwise the loader's defaults (0/unknown) are used. + + If ``ce_calibration`` is given, each PSM's collision energy is overridden by + the calibrated value for that ``(accession, psm_id)``; PSMs not in the map + fall back to ``default_ce``. """ metadata: dict[str, tuple[int, int]] = ( load_catalog_metadata(catalog_path) if catalog_path else {} @@ -255,7 +282,8 @@ def build_split_datasets( instr, acq = metadata[accession] kw.setdefault("instrument", instr) kw.setdefault("acq_mode", acq) - examples.extend(prepare_examples(sage_dir_p, cap=cap, seed=seed, **kw)) + examples.extend(prepare_examples(sage_dir_p, cap=cap, seed=seed, + ce_calibration=ce_calibration, **kw)) buckets: dict[str, list[dict]] = {"train": [], "val": [], "test": []} for ex in examples: diff --git a/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate_optimal_nce.py b/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate_optimal_nce.py new file mode 100644 index 000000000..cfdb57694 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate_optimal_nce.py @@ -0,0 +1,158 @@ +"""Per-PSM optimal NCE calibration — v4's inference recipe applied to ppng. + +For each PSM, sweep the model over a grid of collision energies and pick the CE +that maximises spectral-angle similarity to the observed spectrum. Reports the +mean optimal-CE SA, the mean fixed-CE SA at the original (per-PSM dataset/catalog) +value, and the optimal-CE distribution. No retraining — this isolates "how good +is the model when the right CE is used per spectrum" from "how good is the +flat CE we used at fine-tune time." + +Example: + python -m peptide_property_ng.eval.evaluate_optimal_nce \\ + --checkpoint runs/finetune-canonicalSA/best.pt \\ + --datasets-glob '/scratch/claudius-proteomics/*' \\ + --catalog .../timstof_catalog.tsv +""" +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from torch.utils.data import DataLoader + +from peptide_property_ng.data.collate import make_collate_fn +from peptide_property_ng.data.sage_dataset import ( + build_split_datasets, discover_sage_dirs, load_ce_calibration, +) +from peptide_property_ng.losses import intensity_signal_mask, masked_spectral_angle +from peptide_property_ng.model.config import PRESETS +from peptide_property_ng.model.multitask import UnifiedPeptidePropertyModel +from peptide_property_ng.modifications.composition import CompositionTable + + +@torch.no_grad() +def evaluate_optimal_nce( + model, + loader, + device: str, + ce_grid: torch.Tensor, +): + """Return per-PSM (opt_sa, opt_ce, fixed_sa) tensors over the signal subset. + + ``fixed_sa`` uses the loader's original (per-PSM dataset/catalog) CE — i.e. + the same CE the model saw during fine-tune; ``opt_sa`` is the per-PSM max + over the supplied grid. + """ + model.eval() + ce_grid = ce_grid.to(device) + n_ce = ce_grid.numel() + + opt_sa_list: list[torch.Tensor] = [] + opt_ce_list: list[torch.Tensor] = [] + fixed_sa_list: list[torch.Tensor] = [] + + for batch in loader: + batch = {k: v.to(device) for k, v in batch.items()} + target = batch["intensity_target"] + signal = intensity_signal_mask(target) + if not signal.any(): + continue + + B = target.shape[0] + ce_orig = batch["collision_energy"].clone() + + # Per-CE SA similarity (B, n_ce). + sa_matrix = torch.zeros(B, n_ce, device=device) + for j, ce in enumerate(ce_grid): + batch["collision_energy"] = ce.expand(B).float() + pred = model(batch, tasks=["intensity"])["intensity"] + sa_matrix[:, j] = 1.0 - masked_spectral_angle(pred, target) + + opt_sa, opt_idx = sa_matrix.max(dim=1) + opt_ce = ce_grid[opt_idx] + + # Fixed reference: model's original CE (reproduces the training test SA). + batch["collision_energy"] = ce_orig + pred = model(batch, tasks=["intensity"])["intensity"] + fixed_sa = 1.0 - masked_spectral_angle(pred, target) + + opt_sa_list.append(opt_sa[signal].cpu()) + opt_ce_list.append(opt_ce[signal].cpu()) + fixed_sa_list.append(fixed_sa[signal].cpu()) + + return ( + torch.cat(opt_sa_list), + torch.cat(opt_ce_list), + torch.cat(fixed_sa_list), + ) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--checkpoint", required=True, help="path to a best.pt") + ap.add_argument("--datasets-glob", default="/scratch/claudius-proteomics/*") + ap.add_argument("--catalog", default=None) + ap.add_argument("--default-ce", type=float, default=0.26) + ap.add_argument("--ce-calibration", default=None, + help="parquet of per-PSM calibrated CEs; if supplied, the 'fixed' baseline " + "uses these (matches a fine-tune that was trained with the same calibration)") + ap.add_argument("--cap", type=int, default=4000) + ap.add_argument("--max-datasets", type=int, default=0) + ap.add_argument("--batch-size", type=int, default=64) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--ce-min", type=float, default=0.10) + ap.add_argument("--ce-max", type=float, default=0.40) + ap.add_argument("--ce-step", type=float, default=0.02) + ap.add_argument("--split", default="test", choices=["train", "val", "test"]) + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + args = ap.parse_args() + + print(f"checkpoint: {args.checkpoint}") + sage_dirs = discover_sage_dirs(args.datasets_glob) + if args.max_datasets: + sage_dirs = sage_dirs[: args.max_datasets] + print(f" {len(sage_dirs)} datasets, cap {args.cap}/dataset") + ce_cal = load_ce_calibration(args.ce_calibration) if args.ce_calibration else None + splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed, + catalog_path=args.catalog, default_ce=args.default_ce, + ce_calibration=ce_cal) + ds = splits[args.split] + + ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=False) + preset = ckpt.get("preset", "small") + cfg = PRESETS[preset] + collate = make_collate_fn(cfg.pad_token_id, cfg.max_charge) + loader = DataLoader(ds, batch_size=args.batch_size, shuffle=False, collate_fn=collate) + model = UnifiedPeptidePropertyModel(cfg, CompositionTable.load()).to(args.device) + model.load_state_dict(ckpt["model_state_dict"]) + + n_pts = int(round((args.ce_max - args.ce_min) / args.ce_step)) + 1 + ce_grid = torch.linspace(args.ce_min, args.ce_max, n_pts) + print(f"split={args.split}, |dataset|={len(ds):,}, preset={preset}") + print(f"CE grid ({n_pts} pts): {args.ce_min:.3f} .. {args.ce_max:.3f} step {args.ce_step:.3f}") + print() + + opt_sa, opt_ce, fixed_sa = evaluate_optimal_nce(model, loader, args.device, ce_grid) + n = opt_sa.numel() + lift = float(opt_sa.mean() - fixed_sa.mean()) + print(f"PSMs scored: {n:,}") + print(f"fixed (original per-PSM CE) SA: {float(fixed_sa.mean()):.4f}") + print(f"per-PSM optimal CE SA: {float(opt_sa.mean()):.4f} (lift +{lift:.4f})") + print() + print("optimal CE distribution (per-PSM):") + for q in [0.05, 0.25, 0.5, 0.75, 0.95]: + print(f" q{int(q*100):02d} {float(opt_ce.quantile(q)):.3f}") + print(f" mean {float(opt_ce.mean()):.3f}") + # Mode by histogram (CE grid is discrete, so this is exact). + vals, counts = torch.unique(opt_ce, return_counts=True) + print(f" mode {float(vals[counts.argmax()]):.3f} ({int(counts.max())} / {n} PSMs)") + print() + print("share of PSMs at each CE grid point:") + for v, c in zip(vals.tolist(), counts.tolist()): + bar = "#" * int(40 * c / n) + print(f" {v:.3f} {c:>5d} {c/n:5.1%} {bar}") + + +if __name__ == "__main__": + main() diff --git a/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py b/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py index cc10f6a25..fc2f0b96c 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/eval/metrics.py @@ -16,9 +16,13 @@ def _pearson(x: torch.Tensor, y: torch.Tensor) -> float: @torch.no_grad() -def evaluate_split(model, loader, device: str = "cpu") -> dict[str, float]: +def evaluate_split(model, loader, device: str = "cpu", + tasks: list[str] | None = None) -> dict[str, float]: """Run the model over a DataLoader and return per-task metrics. + ``tasks`` restricts the forward to the named tasks (e.g. ``["intensity"]``); + metrics for skipped tasks come back as ``nan``. + Returns: intensity spectral angle (similarity, higher better), CCS / RT median absolute error, RT Pearson r, charge accuracy. """ @@ -31,7 +35,7 @@ def evaluate_split(model, loader, device: str = "cpu") -> dict[str, float]: for batch in loader: batch = {k: v.to(device) for k, v in batch.items()} - out = model(batch) + out = model(batch, tasks=tasks) if "intensity" in out: sa = 1.0 - masked_spectral_angle(out["intensity"], batch["intensity_target"]) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/losses.py b/packages/peptide-property-ng/src/peptide_property_ng/losses.py index fb522c8e2..073843fae 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/losses.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/losses.py @@ -20,6 +20,11 @@ def masked_spectral_angle(pred: torch.Tensor, target: torch.Tensor, eps: float = practice — there the observable b/y positions almost all carry positive intensity, so `> 0` and `>= 0` masks coincide. """ + # Pooled head outputs a fixed (B, 29, n_ion) grid; site head outputs + # (B, L-1, n_ion). Slice pred to the batch's target sites so both shapes + # work transparently here. + if pred.shape[1] > target.shape[1]: + pred = pred[:, : target.shape[1], :] mask = (target > 0).float() p = (pred * mask).reshape(pred.shape[0], -1) t = (target * mask).reshape(target.shape[0], -1) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py index d1429663f..1582cde9c 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/config.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/config.py @@ -100,6 +100,9 @@ class ModelConfig: # heads n_ion_channels: int = 6 # b/y ions x fragment charges 1-3 (Prosit-compatible) tasks: tuple[str, ...] = ("intensity", "ccs", "rt", "charge") + # intensity head topology: "site" (local FiLM, per-cleavage-site) or + # "pooled" (v4-style attention-pooled Prosit-174); see heads/intensity_pooled.py + intensity_head: str = "site" def __post_init__(self) -> None: if self.d_model % self.n_heads != 0: diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/heads/intensity_pooled.py b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/intensity_pooled.py new file mode 100644 index 000000000..71d376733 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/intensity_pooled.py @@ -0,0 +1,101 @@ +"""V4-compatible pooled intensity head — Prosit-174 ablation. + +Mirrors ``imspy_predictors.models.heads.IntensityHead``: attention-pool the +residue latents, concat with charge / CE / instrument side-embeddings, MLP → +sigmoid → reshape to ``(B, max_sites, n_ion_channels)``. + +The hypothesis (codex review): a global head that predicts the *whole* Prosit-174 +vector can model cross-fragment competition ("this peptide's base peak should +dominate, suppress the rest"), which the per-cleavage-site FiLM head in +``intensity.py`` cannot express except through encoder context. If swapping +heads on the same pretrained encoder closes the remaining SA gap to v4 (~0.83), +the gap is head topology, not capacity. +""" +from __future__ import annotations + +import torch +from torch import nn +from torch.nn import functional as F + +from peptide_property_ng.model.config import ModelConfig + + +# Prosit-174 ceiling — fragments per peptide cap (29 = 30 aa - 1 cleavage site). +# This is the public-corpus ground truth shape; even though ``cfg.max_seq_len`` +# allows longer peptides in the encoder, intensity supervision is bounded here. +_MAX_SITES = 29 + + +class PooledIntensityHead(nn.Module): + """Predict ``(B, 29, n_ion_channels)`` via attention-pool + side-conditioned MLP.""" + + def __init__(self, cfg: ModelConfig): + super().__init__() + d = cfg.d_model + self.max_sites = _MAX_SITES + self.n_ion = cfg.n_ion_channels + self.num_outputs = self.max_sites * self.n_ion + self.max_charge = cfg.max_charge + + # Single-head self-attention over the residue latents. + self.attention = nn.MultiheadAttention( + embed_dim=d, num_heads=1, batch_first=True, dropout=cfg.dropout, + ) + + # Side embeddings — match v4's d_model // 4 sizing. + side = max(d // 4, 16) + self.charge_embed = nn.Linear(cfg.max_charge + 1, side) + self.ce_embed = nn.Linear(1, side) + # Instrument conditioning lives in the head (v4 parity) on top of the + # encoder's global-token conditioning, not as a replacement. + self.instrument_embed = nn.Embedding(cfg.n_instruments, side) + + in_dim = d + 3 * side + hidden = max(d * 2, 512) + self.fc = nn.Sequential( + nn.Linear(in_dim, hidden), + nn.LayerNorm(hidden), + nn.GELU(), + nn.Dropout(cfg.dropout), + nn.Linear(hidden, hidden), + nn.LayerNorm(hidden), + nn.GELU(), + nn.Dropout(cfg.dropout), + nn.Linear(hidden, self.num_outputs), + ) + + def forward( + self, + latent: torch.Tensor, # (B, 1+L, d) — index 0 is the global token + charge: torch.Tensor, # (B,) long + collision_energy: torch.Tensor, # (B,) float + instrument: torch.Tensor | None = None, # (B,) long + padding_mask: torch.Tensor | None = None, # (B, L) bool, True at pad + ) -> torch.Tensor: + residues = latent[:, 1:] # (B, L, d) + B = residues.size(0) + device = residues.device + + # Self-attention pool, then mean over valid residues. + pooled, _ = self.attention( + residues, residues, residues, key_padding_mask=padding_mask, + ) + if padding_mask is not None: + valid = (~padding_mask).unsqueeze(-1).float() + denom = valid.sum(dim=1).clamp(min=1.0) + pooled = (pooled * valid).sum(dim=1) / denom # (B, d) + else: + pooled = pooled.mean(dim=1) + + # Side-info embeddings (v4 layout). + chg = charge.view(-1).long().clamp(0, self.max_charge) + chg_onehot = F.one_hot(chg, num_classes=self.max_charge + 1).float() + chg_emb = self.charge_embed(chg_onehot) + ce_emb = self.ce_embed(collision_energy.view(-1, 1).float()) + if instrument is None: + instrument = torch.zeros(B, dtype=torch.long, device=device) + inst_emb = self.instrument_embed(instrument.view(-1).long()) + + features = torch.cat([pooled, chg_emb, ce_emb, inst_emb], dim=-1) + out = torch.sigmoid(self.fc(features)) # (B, max_sites * n_ion) + return out.view(B, self.max_sites, self.n_ion) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py b/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py index 382da6360..88d14ddc8 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py @@ -8,6 +8,7 @@ from peptide_property_ng.model.encoder import PeptidePropertyEncoder from peptide_property_ng.model.heads.ccs import IonMobilityHead from peptide_property_ng.model.heads.intensity import IntensityHead +from peptide_property_ng.model.heads.intensity_pooled import PooledIntensityHead from peptide_property_ng.model.heads.scalar import ChargeHead, RetentionTimeHead from peptide_property_ng.modifications.composition import CompositionTable @@ -43,9 +44,15 @@ def __init__( if composition_table is None: composition_table = CompositionTable.load() self.encoder = PeptidePropertyEncoder(cfg, composition_table) - self.heads = nn.ModuleDict( - {t: self.HEAD_TYPES[t](cfg) for t in cfg.tasks} - ) + # Intensity head is dispatched on cfg.intensity_head ("site" | "pooled"). + # Other heads always use the canonical class. Stored under one key + # ("intensity") regardless, so checkpoint code keys are stable. + head_types: dict[str, type[nn.Module]] = dict(self.HEAD_TYPES) + if cfg.intensity_head == "pooled": + head_types["intensity"] = PooledIntensityHead + elif cfg.intensity_head != "site": + raise ValueError(f"unknown intensity_head '{cfg.intensity_head}'") + self.heads = nn.ModuleDict({t: head_types[t](cfg) for t in cfg.tasks}) def forward( self, @@ -60,9 +67,18 @@ def forward( ) out: dict[str, object] = {} if "intensity" in active: - out["intensity"] = self.heads["intensity"]( - latent, batch["charge"], batch["collision_energy"] - ) + head = self.heads["intensity"] + if isinstance(head, PooledIntensityHead): + pad_mask = batch["tokens"] == self.cfg.pad_token_id + out["intensity"] = head( + latent, batch["charge"], batch["collision_energy"], + instrument=batch.get("instrument"), + padding_mask=pad_mask, + ) + else: + out["intensity"] = head( + latent, batch["charge"], batch["collision_energy"] + ) if "ccs" in active: out["ccs"] = self.heads["ccs"](latent, batch["charge"], batch["precursor_mz"]) if "rt" in active: diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/calibrate_ce.py b/packages/peptide-property-ng/src/peptide_property_ng/train/calibrate_ce.py new file mode 100644 index 000000000..130851b42 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/calibrate_ce.py @@ -0,0 +1,183 @@ +"""Compute per-PSM optimal CE with a frozen checkpoint and write it to a parquet. + +For each prepared PSM, sweep collision energy over a grid and pick the value +that maximises spectral-angle similarity to the observed spectrum. The output +is a parquet of ``(accession, psm_id, calibrated_ce)`` that ``train.train`` +can load via ``--ce-calibration`` to fine-tune with per-PSM CE instead of a +flat default. + +Recipe: calibrate the *pretrained* (or other freeze-point) checkpoint once +over the full prepared corpus; then fine-tune from that same checkpoint with +the calibration applied. v4's effective recipe. + +Example: + python -m peptide_property_ng.train.calibrate_ce \\ + --checkpoint /scratch/.../small-pretrained.pt \\ + --datasets-glob '/scratch/claudius-proteomics/*' \\ + --catalog .../timstof_catalog.tsv \\ + --out /scratch/.../ce_calibration_pretrained_small.parquet +""" +from __future__ import annotations + +import argparse +import time +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import torch +from torch.utils.data import DataLoader + +from peptide_property_ng.data.collate import make_collate_fn +from peptide_property_ng.data.sage_dataset import ( + SagePropertyDataset, + build_split_datasets, + discover_sage_dirs, +) +from peptide_property_ng.losses import intensity_signal_mask, masked_spectral_angle +from peptide_property_ng.model.config import PRESETS +from peptide_property_ng.model.multitask import UnifiedPeptidePropertyModel +from peptide_property_ng.modifications.composition import CompositionTable + + +@torch.no_grad() +def calibrate(model, loader, device: str, ce_grid: torch.Tensor) -> tuple[list, list, list, list]: + """Sweep the grid per batch, return per-example ``(acc, psm, ce, fixed_sa, opt_sa)``.""" + model.eval() + ce_grid = ce_grid.to(device) + n_ce = ce_grid.numel() + + accs: list[str] = [] + pids: list[int] = [] + opt_ces: list[float] = [] + fixed_sas: list[float] = [] + opt_sas: list[float] = [] + + for batch in loader: + meta_acc = batch.pop("_accession") + meta_pid = batch.pop("_psm_id") + batch = {k: v.to(device) for k, v in batch.items()} + target = batch["intensity_target"] + signal = intensity_signal_mask(target) + B = target.shape[0] + + ce_orig = batch["collision_energy"].clone() + sa_matrix = torch.full((B, n_ce), float("nan"), device=device) + for j, ce in enumerate(ce_grid): + batch["collision_energy"] = ce.expand(B).float() + pred = model(batch, tasks=["intensity"])["intensity"] + sa_matrix[:, j] = 1.0 - masked_spectral_angle(pred, target) + + # PSMs with no signal: keep their CE at the loader's default (no calibration); + # also mark fixed_sa / opt_sa as NaN so the parquet records "uncalibrated". + opt_sa, opt_idx = sa_matrix.max(dim=1) + opt_ce = ce_grid[opt_idx] + opt_ce = torch.where(signal, opt_ce, ce_orig) + + batch["collision_energy"] = ce_orig + pred = model(batch, tasks=["intensity"])["intensity"] + fixed_sa = 1.0 - masked_spectral_angle(pred, target) + + for i in range(B): + accs.append(meta_acc[i]) + pids.append(int(meta_pid[i].item())) + opt_ces.append(float(opt_ce[i].item())) + if bool(signal[i]): + fixed_sas.append(float(fixed_sa[i].item())) + opt_sas.append(float(opt_sa[i].item())) + else: + fixed_sas.append(float("nan")) + opt_sas.append(float("nan")) + return accs, pids, opt_ces, fixed_sas, opt_sas + + +def _collate_with_meta(pad_token_id: int, max_charge: int): + """Same as ``make_collate_fn`` but also returns per-example accession + psm_id.""" + inner = make_collate_fn(pad_token_id, max_charge) + + def _fn(examples: list[dict]) -> dict: + out = inner(examples) + out["_accession"] = [e["accession"] for e in examples] + out["_psm_id"] = torch.tensor([e["psm_id"] for e in examples], dtype=torch.int64) + return out + + return _fn + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--checkpoint", required=True, help="checkpoint to calibrate against") + ap.add_argument("--datasets-glob", default="/scratch/claudius-proteomics/*") + ap.add_argument("--catalog", default=None) + ap.add_argument("--default-ce", type=float, default=0.26) + ap.add_argument("--cap", type=int, default=4000) + ap.add_argument("--max-datasets", type=int, default=0) + ap.add_argument("--batch-size", type=int, default=64) + ap.add_argument("--seed", type=int, default=0, + help="MUST match what train.train uses (controls per-dataset PSM sampling)") + ap.add_argument("--ce-min", type=float, default=0.00) + ap.add_argument("--ce-max", type=float, default=0.60) + ap.add_argument("--ce-step", type=float, default=0.02) + ap.add_argument("--out", required=True, help="output parquet path") + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + args = ap.parse_args() + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + + sage_dirs = discover_sage_dirs(args.datasets_glob) + if args.max_datasets: + sage_dirs = sage_dirs[: args.max_datasets] + print(f"[{time.strftime('%H:%M:%S')}] {len(sage_dirs)} datasets (cap {args.cap}/dataset)") + splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed, + catalog_path=args.catalog, default_ce=args.default_ce) + examples = splits["train"].examples + splits["val"].examples + splits["test"].examples + print(f" total {len(examples):,} prepared examples to calibrate " + f"({len(splits['train']):,} train + {len(splits['val']):,} val + {len(splits['test']):,} test)") + + ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=False) + preset = ckpt.get("preset", "small") + cfg = PRESETS[preset] + model = UnifiedPeptidePropertyModel(cfg, CompositionTable.load()).to(args.device) + model.load_state_dict(ckpt["model_state_dict"]) + print(f" checkpoint preset={preset}, {model.num_parameters():,} parameters") + + n_pts = int(round((args.ce_max - args.ce_min) / args.ce_step)) + 1 + ce_grid = torch.linspace(args.ce_min, args.ce_max, n_pts) + print(f" CE grid ({n_pts} pts): {args.ce_min:.3f} .. {args.ce_max:.3f} step {args.ce_step:.3f}") + + loader = DataLoader( + SagePropertyDataset(examples), batch_size=args.batch_size, shuffle=False, + collate_fn=_collate_with_meta(cfg.pad_token_id, cfg.max_charge), + ) + t0 = time.time() + accs, pids, opt_ces, fixed_sas, opt_sas = calibrate(model, loader, args.device, ce_grid) + print(f"[{time.strftime('%H:%M:%S')}] calibrated {len(accs):,} PSMs in {time.time() - t0:.0f}s") + + table = pa.table({ + "accession": pa.array(accs, type=pa.string()), + "psm_id": pa.array(pids, type=pa.int64()), + "calibrated_ce": pa.array(opt_ces, type=pa.float32()), + "fixed_sa": pa.array(fixed_sas, type=pa.float32()), + "opt_sa": pa.array(opt_sas, type=pa.float32()), + }) + pq.write_table(table, out_path) + print(f" wrote {out_path} ({out_path.stat().st_size / 1e6:.2f} MB)") + + # Quick summary of the calibration so the user can see the lift. + import numpy as np + fa = np.array(fixed_sas, dtype=np.float64) + oa = np.array(opt_sas, dtype=np.float64) + valid = ~np.isnan(fa) & ~np.isnan(oa) + print(f"\nover the {int(valid.sum()):,} signal PSMs:") + print(f" fixed (per-PSM default/catalog CE) mean SA {fa[valid].mean():.4f}") + print(f" per-PSM optimal CE mean SA {oa[valid].mean():.4f} " + f"(lift +{(oa[valid] - fa[valid]).mean():.4f})") + ce_arr = np.array(opt_ces, dtype=np.float64)[valid] + print(f" CE: mean {ce_arr.mean():.3f} median {np.median(ce_arr):.3f} " + f"q05 {np.quantile(ce_arr, 0.05):.3f} q95 {np.quantile(ce_arr, 0.95):.3f}") + + +if __name__ == "__main__": + main() diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py index 556cfb444..9853bf6fc 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py @@ -14,6 +14,7 @@ from __future__ import annotations import argparse +import dataclasses import json import time from pathlib import Path @@ -97,6 +98,10 @@ def _train_task_epoch(model, loader, task, loss_fn, optimizer, device) -> float: def main() -> None: ap = argparse.ArgumentParser(description="Staged pretraining for the unified model.") ap.add_argument("--preset", default="small", choices=sorted(PRESETS)) + ap.add_argument("--intensity-head", default=None, + choices=["site", "pooled"], + help="intensity head topology (default: preset value); " + "use 'pooled' to pretrain the v4-style Prosit-174 head") ap.add_argument("--cap", type=int, default=50000, help="max examples per stage (0 = all)") ap.add_argument("--epochs", type=int, default=3, help="epochs per stage") ap.add_argument("--batch-size", type=int, default=128) @@ -116,6 +121,8 @@ def main() -> None: out_dir = Path(args.out) out_dir.mkdir(parents=True, exist_ok=True) cfg = PRESETS[args.preset] + if args.intensity_head: + cfg = dataclasses.replace(cfg, intensity_head=args.intensity_head) device = args.device cap = None if args.cap == 0 else args.cap diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py index 7c361df09..8bd521b02 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py @@ -16,7 +16,9 @@ from torch.utils.data import DataLoader from peptide_property_ng.data.collate import make_collate_fn -from peptide_property_ng.data.sage_dataset import build_split_datasets, discover_sage_dirs +from peptide_property_ng.data.sage_dataset import ( + build_split_datasets, discover_sage_dirs, load_ce_calibration, +) from peptide_property_ng.eval.metrics import evaluate_split from peptide_property_ng.losses import MultiTaskLoss from peptide_property_ng.model.config import PRESETS @@ -28,14 +30,14 @@ def _to_device(batch: dict, device: str) -> dict: return {k: v.to(device) for k, v in batch.items()} -def train_epoch(model, loader, loss_fn, optimizer, device) -> dict[str, float]: +def train_epoch(model, loader, loss_fn, optimizer, device, tasks=None) -> dict[str, float]: model.train() totals: dict[str, float] = {} n_batches = 0 for batch in loader: batch = _to_device(batch, device) optimizer.zero_grad() - loss, parts = loss_fn(model(batch), batch) + loss, parts = loss_fn(model(batch, tasks=tasks), batch) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() @@ -54,10 +56,17 @@ def main() -> None: ap.add_argument("--default-ce", type=float, default=0.26, help="default normalized collision energy for Sage PSMs that lack one " "(~0.26 = timsTOF-ms2 median, in-distribution for the pretrained CE FiLM)") + ap.add_argument("--ce-calibration", default=None, + help="parquet of per-PSM calibrated CEs (from calibrate_ce.py); " + "overrides default_ce per (accession, psm_id) where present") ap.add_argument("--preset", default="small", choices=sorted(PRESETS)) ap.add_argument("--comp-fusion", default=None, choices=["add", "gate", "token_only", "composition_only"], help="modification-encoding fusion / ablation (default: preset value)") + ap.add_argument("--intensity-head", default=None, + choices=["site", "pooled"], + help="intensity head topology: 'site' (local FiLM, default) or " + "'pooled' (v4-style Prosit-174 attention-pooled ablation)") ap.add_argument("--cap", type=int, default=4000, help="max PSMs sampled per dataset") ap.add_argument("--max-datasets", type=int, default=0, help="limit datasets (0 = all)") ap.add_argument("--epochs", type=int, default=12) @@ -68,6 +77,12 @@ def main() -> None: ap.add_argument("--patience", type=int, default=4) ap.add_argument("--init-from", default=None, help="checkpoint to initialise weights from (e.g. a pretrained.pt)") + ap.add_argument("--tasks", default="intensity,ccs,rt,charge", + help="comma-separated tasks to train + evaluate; restricts forward and loss " + "(e.g. 'intensity' for a specialized fine-tune)") + ap.add_argument("--freeze-encoder", action="store_true", + help="freeze the shared encoder (only the task heads train) -- linear-probe / " + "feature-extraction mode on the pretrained encoder") ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") ap.add_argument("--out", default="runs/prototype") args = ap.parse_args() @@ -83,9 +98,13 @@ def main() -> None: print(f" {len(sage_dirs)} datasets") print(f"[{time.strftime('%H:%M:%S')}] preparing examples (cap {args.cap}/dataset) ...", flush=True) + ce_cal = load_ce_calibration(args.ce_calibration) if args.ce_calibration else None + if ce_cal is not None: + print(f" CE calibration loaded: {len(ce_cal):,} keys from {args.ce_calibration}") t0 = time.time() splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed, - catalog_path=args.catalog, default_ce=args.default_ce) + catalog_path=args.catalog, default_ce=args.default_ce, + ce_calibration=ce_cal) for name, ds in splits.items(): print(f" {name}: {len(ds):,} examples") print(f" prepared in {time.time() - t0:.0f}s") @@ -95,6 +114,8 @@ def main() -> None: cfg = PRESETS[args.preset] if args.comp_fusion: cfg = dataclasses.replace(cfg, comp_fusion=args.comp_fusion) + if args.intensity_head: + cfg = dataclasses.replace(cfg, intensity_head=args.intensity_head) collate = make_collate_fn(cfg.pad_token_id, cfg.max_charge) loaders = { name: DataLoader( @@ -117,16 +138,27 @@ def main() -> None: missing, unexpected = model.load_state_dict(ckpt["model_state_dict"], strict=False) print(f"initialised from {args.init_from} " f"(missing={len(missing)}, unexpected={len(unexpected)})") - print(f"model: '{args.preset}' preset, {model.num_parameters():,} parameters, device={device}") - optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + + task_list = [t.strip() for t in args.tasks.split(",") if t.strip()] + if args.freeze_encoder: + for p in model.encoder.parameters(): + p.requires_grad = False + n_frozen = sum(p.numel() for p in model.encoder.parameters()) + print(f"encoder frozen: {n_frozen:,} params no longer training") + + n_train = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"model: '{args.preset}' preset, {model.num_parameters():,} total parameters, " + f"{n_train:,} trainable, device={device}, tasks={task_list}") + optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], + lr=args.lr, weight_decay=args.weight_decay) loss_fn = MultiTaskLoss() best_sa, best_epoch, bad = -1.0, -1, 0 history: list[dict] = [] for epoch in range(1, args.epochs + 1): t0 = time.time() - tr = train_epoch(model, loaders["train"], loss_fn, optimizer, device) - val = evaluate_split(model, loaders["val"], device) + tr = train_epoch(model, loaders["train"], loss_fn, optimizer, device, tasks=task_list) + val = evaluate_split(model, loaders["val"], device, tasks=task_list) history.append({"epoch": epoch, "train": tr, "val": val}) print( f"[{time.strftime('%H:%M:%S')}] epoch {epoch:2d} " @@ -161,7 +193,7 @@ def main() -> None: best_epoch = epoch ckpt = torch.load(out_dir / "best.pt", map_location=device) model.load_state_dict(ckpt["model_state_dict"]) - test = evaluate_split(model, loaders["test"], device) + test = evaluate_split(model, loaders["test"], device, tasks=task_list) print(f"\n=== test (best epoch {best_epoch}) ===") for k, v in test.items(): print(f" {k:16s} {v:.4f}") From f845c2e422d59edf4918a51648885fc026fea676 Mon Sep 17 00:00:00 2001 From: David Teschner Date: Mon, 25 May 2026 22:13:53 +0200 Subject: [PATCH 16/19] fix(peptide-property-ng): clamp_eps in masked_spectral_angle + checkpoint metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The real bug: the arccos clamp_eps was 1e-8, which let arccos's near-boundary derivative blow up to ~7000 on sparse intensity targets (Sage matched_fragments with 1 positive position normalize to one-hot vectors whose cos=1 exactly). Even with clamp absorbing the boundary in forward, the gradient chain through F.normalize on tiny-magnitude masked vectors compounded into NaN gradients. This was the silent failure mode underneath three earlier negative results: * RESEARCH preset pretraining at lr=3e-4 NaN'd intensity stages from epoch 1 * RESEARCH preset pretraining at lr=1e-4 NaN'd intensity (the cause, not lr) * Pooled intensity head pretraining NaN'd intensity (same root cause) Decoupled the two eps in masked_spectral_angle: - norm_eps (1e-8): F.normalize denominator floor — accurate for normal mags - clamp_eps (1e-4): arccos boundary buffer — derivative bounded at ~70 instead of ~7000 The fix unblocks scaling experiments. Empirically, with the fix in place: * RESEARCH preset now pretrains cleanly through all 5 stages, with each stage reaching higher val SA than SMALL (timsTOF 0.8574 vs 0.8484). * Pooled-head pretraining likewise runs to completion. Recipe / infra additions: * pretrain.py: --intensity-head flag; per-stage eval restricted to current stage's task (so the pooled head's fixed (B,29,n_ion) output doesn't shape- clash with Chronologer's long-peptide intensity placeholders during RT eval) * train.py + pretrain.py: save intensity_head in checkpoint metadata * calibrate_ce.py + evaluate_optimal_nce.py: auto-detect intensity_head from checkpoint (back-compat for legacy checkpoints inferring from state-dict keys) * evaluate_optimal_nce.py: --ce-calibration to use the calibration parquet as the 'fixed' baseline (matches what a fine-tune that was trained with the same calibration actually saw) * encoder.py: disable nested-tensor fast path on TransformerEncoder — its SDPA backward had a real NaN risk on certain shape combinations Campaign fine-tune results, ordered by test SA on the 7,724-PSM hold-out: config test SA --------------------------------------------- ------- canonical baseline (flat CE 0.26) 0.7597 + inference-time optimal-CE search 0.7945 round-1 per-PSM CE calibration 0.7867 round-1 calibration + inference search 0.8058 round-2 calibration 0.8090 round-2 calibration + inference search 0.8170 per-(accession, charge) CE labels 0.7616 ABLATION pooled head, pretrained + r1 calibration 0.7673 ABLATION RESEARCH preset + r1 calibration 0.7868 ABLATION RESEARCH preset + r2 calibration 0.8073 ABLATION RESEARCH preset + r2 + inference search 0.8137 ABLATION The +0.057 SA lift from baseline came from iterative per-PSM CE calibration + inference-time CE search alone. Pooled-head topology, RESEARCH capacity, and per-(run, charge) label smoothing each underperformed the SMALL site-head per-PSM recipe — the per-PSM CE specificity is doing real work, not noise. --- .../eval/evaluate_optimal_nce.py | 6 +++++ .../src/peptide_property_ng/losses.py | 26 ++++++++++++++----- .../src/peptide_property_ng/model/encoder.py | 8 ++++++ .../peptide_property_ng/train/calibrate_ce.py | 13 +++++++++- .../src/peptide_property_ng/train/pretrain.py | 13 ++++++++-- .../src/peptide_property_ng/train/train.py | 1 + 6 files changed, 58 insertions(+), 9 deletions(-) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate_optimal_nce.py b/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate_optimal_nce.py index cfdb57694..a38b72daa 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate_optimal_nce.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/eval/evaluate_optimal_nce.py @@ -119,9 +119,15 @@ def main() -> None: ce_calibration=ce_cal) ds = splits[args.split] + import dataclasses ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=False) preset = ckpt.get("preset", "small") cfg = PRESETS[preset] + intensity_head = ckpt.get("intensity_head") + if intensity_head is None: + keys = ckpt["model_state_dict"].keys() + intensity_head = "pooled" if any("heads.intensity.attention." in k for k in keys) else "site" + cfg = dataclasses.replace(cfg, intensity_head=intensity_head) collate = make_collate_fn(cfg.pad_token_id, cfg.max_charge) loader = DataLoader(ds, batch_size=args.batch_size, shuffle=False, collate_fn=collate) model = UnifiedPeptidePropertyModel(cfg, CompositionTable.load()).to(args.device) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/losses.py b/packages/peptide-property-ng/src/peptide_property_ng/losses.py index 073843fae..ce66b7d26 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/losses.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/losses.py @@ -7,7 +7,12 @@ from torch.nn import functional as F -def masked_spectral_angle(pred: torch.Tensor, target: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: +def masked_spectral_angle( + pred: torch.Tensor, + target: torch.Tensor, + norm_eps: float = 1e-8, + clamp_eps: float = 1e-4, +) -> torch.Tensor: """Per-sample spectral-angle distance in ``[0, 1]`` (0 = identical spectra). Canonical Prosit / Sage-fine-tune loss: mask ``target > 0`` (only observed @@ -16,9 +21,18 @@ def masked_spectral_angle(pred: torch.Tensor, target: torch.Tensor, eps: float = ``matched_fragments`` (Sage only reports peaks it matched) and matches the production v4 / imspy_predictors `masked_spectral_distance` convention. + The two eps are deliberately decoupled: + - ``norm_eps`` (1e-8) keeps ``F.normalize`` accurate for non-degenerate + magnitudes; + - ``clamp_eps`` (1e-4) bounds ``arccos``'s near-boundary derivative + (≈ -1/sqrt(2·clamp_eps), so ~70 instead of ~7000 at 1e-8). With 1e-8 + the SA-loss backward NaN-s on sparse one-positive-fragment targets + whose normalize-to-one-hot forces cos exactly to ±1; with 1e-4 those + same samples are absorbed by the clamp (grad 0 above 1-clamp_eps). + Pretraining on PROSPECT-style densely-annotated targets is unaffected in - practice — there the observable b/y positions almost all carry positive - intensity, so `> 0` and `>= 0` masks coincide. + practice — observable b/y positions almost all carry positive intensity, + so the > 0 vs >= 0 masks coincide and cos rarely hits the boundary. """ # Pooled head outputs a fixed (B, 29, n_ion) grid; site head outputs # (B, L-1, n_ion). Slice pred to the batch's target sites so both shapes @@ -28,9 +42,9 @@ def masked_spectral_angle(pred: torch.Tensor, target: torch.Tensor, eps: float = mask = (target > 0).float() p = (pred * mask).reshape(pred.shape[0], -1) t = (target * mask).reshape(target.shape[0], -1) - p = F.normalize(p, dim=1, eps=eps) - t = F.normalize(t, dim=1, eps=eps) - cos = (p * t).sum(dim=1).clamp(-1.0 + eps, 1.0 - eps) + p = F.normalize(p, dim=1, eps=norm_eps) + t = F.normalize(t, dim=1, eps=norm_eps) + cos = (p * t).sum(dim=1).clamp(-1.0 + clamp_eps, 1.0 - clamp_eps) return 2.0 * torch.arccos(cos) / math.pi diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py b/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py index 2287f9522..227870bc9 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/encoder.py @@ -68,6 +68,14 @@ def __init__(self, cfg: ModelConfig, composition_table: CompositionTable): self.instrument_emb = nn.Embedding(cfg.n_instruments, cfg.d_model) self.acq_emb = nn.Embedding(cfg.n_acq_modes, cfg.d_model) + # PyTorch's TransformerEncoder converts padded batches to a NestedTensor + # for the SDPA fast path; the backward pass on that path returns NaN + # gradients on certain shape combinations (silently observed as the + # `_nested_tensor_from_mask` warning we kept seeing — actually the smoke). + # Disable the fast path: slightly slower forward, finite gradients. + if hasattr(self.transformer_encoder, "enable_nested_tensor"): + self.transformer_encoder.enable_nested_tensor = False + def global_token_hook( # noqa: D102 (documented on the class) self, tokens: torch.Tensor, diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/calibrate_ce.py b/packages/peptide-property-ng/src/peptide_property_ng/train/calibrate_ce.py index 130851b42..06ca203c5 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/calibrate_ce.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/calibrate_ce.py @@ -136,12 +136,23 @@ def main() -> None: print(f" total {len(examples):,} prepared examples to calibrate " f"({len(splits['train']):,} train + {len(splits['val']):,} val + {len(splits['test']):,} test)") + import dataclasses ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=False) preset = ckpt.get("preset", "small") cfg = PRESETS[preset] + intensity_head = ckpt.get("intensity_head") + if intensity_head is None: + # Legacy checkpoint without metadata — infer from the state-dict keys. + keys = ckpt["model_state_dict"].keys() + if any("heads.intensity.attention." in k for k in keys): + intensity_head = "pooled" + else: + intensity_head = "site" + cfg = dataclasses.replace(cfg, intensity_head=intensity_head) model = UnifiedPeptidePropertyModel(cfg, CompositionTable.load()).to(args.device) model.load_state_dict(ckpt["model_state_dict"]) - print(f" checkpoint preset={preset}, {model.num_parameters():,} parameters") + print(f" checkpoint preset={preset}, intensity_head={intensity_head}, " + f"{model.num_parameters():,} parameters") n_pts = int(round((args.ce_max - args.ce_min) / args.ce_step)) + 1 ce_grid = torch.linspace(args.ce_min, args.ce_max, n_pts) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py index 9853bf6fc..c9d106964 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/pretrain.py @@ -165,7 +165,11 @@ def main() -> None: for epoch in range(1, args.epochs + 1): te = time.time() train_loss = _train_task_epoch(model, train_loader, task, loss_fn, optimizer, device) - val = evaluate_split(model, val_loader, device) + # Restrict eval to the current stage's task: other heads' outputs are + # at shapes the current corpus's targets/placeholders don't share + # (eg. pooled-head intensity (B,29,6) vs Chronologer's (B,L-1,6) for + # L up to max_seq_len), and per-stage metrics are diagnostic anyway. + val = evaluate_split(model, val_loader, device, tasks=[task]) print(f" epoch {epoch}: train_loss={train_loss:.4f} " f"val[{metric}]={val[metric]:.4f} [{time.time()-te:.0f}s]", flush=True) history.append({"stage": name, "epoch": epoch, @@ -173,7 +177,12 @@ def main() -> None: # Save a per-stage checkpoint (so the campaign fine-tune can pick the # handoff point, not just the post-RT state) and the running pretrained.pt. - ckpt = {"model_state_dict": model.state_dict(), "preset": args.preset, "stage": name} + ckpt = { + "model_state_dict": model.state_dict(), + "preset": args.preset, + "intensity_head": cfg.intensity_head, + "stage": name, + } torch.save(ckpt, out_dir / f"after-{name}.pt") torch.save(ckpt, out_dir / "pretrained.pt") print(f" saved -> {out_dir}/after-{name}.pt (+ pretrained.pt)", flush=True) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py index 8bd521b02..1ca5e1655 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py @@ -174,6 +174,7 @@ def main() -> None: best_sa, best_epoch, bad = val["intensity_sa"], epoch, 0 torch.save( {"model_state_dict": model.state_dict(), "preset": args.preset, + "intensity_head": cfg.intensity_head, "epoch": epoch, "val": val}, out_dir / "best.pt", ) From 7ad044b7887791e13d94f92b79c21645e7dede83 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Mon, 8 Jun 2026 16:41:03 +0200 Subject: [PATCH 17/19] peptide-property-ng: RT-sigma head + IM-sigma loss for timsim peak shapes Add per-peptide peak-shape supervision so the predictor can configure timsim's elution/mobility peak widths (replacing its sampled/fixed defaults): - RTSigmaHead: new head predicting the gradient-normalised RT EMG sigma from the encoder global token (apex stays Chronologer; lambda left to timsim sampling, as it is unidentifiable at timsTOF MS1 frame sampling). - IM sigma: supervise the existing IonMobilityHead (mean, std) std output with an L1 term vs the measured Gaussian mobilogram width -- no model change. - MultiTaskLoss: rt_sigma + im_sigma terms, both masked to peptides carrying a valid label and skipped when absent, so existing shape-free training is unchanged. Weights 0.5 each. - tests/test_shape_heads.py covers the new head output, masked loss terms, and the backward-compatible no-target path. --- .../src/peptide_property_ng/losses.py | 28 ++++++- .../peptide_property_ng/model/heads/scalar.py | 20 +++++ .../peptide_property_ng/model/multitask.py | 6 +- .../tests/test_shape_heads.py | 77 +++++++++++++++++++ 4 files changed, 127 insertions(+), 4 deletions(-) create mode 100644 packages/peptide-property-ng/tests/test_shape_heads.py diff --git a/packages/peptide-property-ng/src/peptide_property_ng/losses.py b/packages/peptide-property-ng/src/peptide_property_ng/losses.py index ce66b7d26..b2389b04e 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/losses.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/losses.py @@ -62,12 +62,20 @@ class MultiTaskLoss: """Weighted sum of per-task losses; tasks with no valid targets are skipped. - intensity : masked spectral-angle distance - - ccs : Gaussian negative log-likelihood (mean, std) vs 1/K0 + - ccs : L1 on the mean 1/K0 + - im_sigma : L1 on the ccs head's std vs the measured IM Gaussian width + (timsim shape target; masked to peptides with a valid label) - rt : L1 on normalised retention time + - rt_sigma : L1 on the gradient-normalised RT EMG peak-width (timsim shape) - charge : cross-entropy + + The two shape terms (``im_sigma``, ``rt_sigma``) are *additive and optional*: + they contribute only when the batch carries the corresponding target+mask, + so existing (shape-free) training is unchanged. """ - DEFAULT_WEIGHTS = {"intensity": 1.0, "ccs": 1.0, "rt": 0.5, "charge": 0.3} + DEFAULT_WEIGHTS = {"intensity": 1.0, "ccs": 1.0, "im_sigma": 0.5, + "rt": 0.5, "rt_sigma": 0.5, "charge": 0.3} def __init__(self, weights: dict[str, float] | None = None): self.weights = dict(self.DEFAULT_WEIGHTS if weights is None else weights) @@ -86,7 +94,7 @@ def __call__( parts["intensity"] = sa[signal].mean() if "ccs" in outputs: - mean, _std = outputs["ccs"] + mean, std = outputs["ccs"] valid = batch["ccs_valid"] if valid.any(): # L1 on the mean — stable for the prototype. Gaussian-NLL on @@ -94,12 +102,26 @@ def __call__( # destabilises the fixed-weight multi-task sum; revisit together # with uncertainty-based task weighting. parts["ccs"] = F.l1_loss(mean[valid], batch["ccs_target"][valid]) + # IM peak-width (timsim Gaussian sigma): supervise the head's std on + # the subset of peptides that carry a measured IM-sigma label. L1 + # (not Gaussian-NLL) per the design — the std is the timsim shape param. + if "im_sigma_target" in batch and "im_sigma_valid" in batch: + sv = batch["im_sigma_valid"] + if sv.any(): + parts["im_sigma"] = F.l1_loss(std[sv], batch["im_sigma_target"][sv]) if "rt" in outputs: valid = batch["rt_valid"] if valid.any(): parts["rt"] = F.l1_loss(outputs["rt"][valid], batch["rt_target"][valid]) + if "rt_sigma" in outputs and "rt_sigma_target" in batch: + sv = batch["rt_sigma_valid"] + if sv.any(): + parts["rt_sigma"] = F.l1_loss( + outputs["rt_sigma"][sv], batch["rt_sigma_target"][sv] + ) + if "charge" in outputs: parts["charge"] = F.cross_entropy(outputs["charge"], batch["charge_target"]) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/heads/scalar.py b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/scalar.py index 018392f36..3e735cc7f 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/heads/scalar.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/heads/scalar.py @@ -35,6 +35,26 @@ def forward(self, latent: torch.Tensor) -> torch.Tensor: return self.net(latent[:, 0]).squeeze(-1) # (B,) +class RTSigmaHead(nn.Module): + """Predict the RT elution-peak *shape* width (gradient-normalised EMG sigma). + + A positive scalar per peptide, read from the encoder global token only (same + leakage-free footing as the RT-apex head). Output is the gradient-normalised + sigma; timsim rescales by the target gradient at simulation time (see the + SIM shape-predictor design). ``lambda`` (EMG tail) is deliberately *not* + predicted — it is unidentifiable at timsTOF MS1 frame sampling, so it stays + sampled by timsim. softplus keeps sigma strictly positive. + """ + + def __init__(self, cfg: ModelConfig): + super().__init__() + self.net = _mlp(cfg.d_model, cfg.d_model, 1) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + raw = self.net(latent[:, 0]).squeeze(-1) # (B,) + return nn.functional.softplus(raw) + 1e-4 + + class ChargeHead(nn.Module): """Predict precursor-charge logits — class ``j`` corresponds to charge ``j + 1``.""" diff --git a/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py b/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py index 88d14ddc8..b70155d3e 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/model/multitask.py @@ -9,7 +9,7 @@ from peptide_property_ng.model.heads.ccs import IonMobilityHead from peptide_property_ng.model.heads.intensity import IntensityHead from peptide_property_ng.model.heads.intensity_pooled import PooledIntensityHead -from peptide_property_ng.model.heads.scalar import ChargeHead, RetentionTimeHead +from peptide_property_ng.model.heads.scalar import ChargeHead, RetentionTimeHead, RTSigmaHead from peptide_property_ng.modifications.composition import CompositionTable @@ -20,6 +20,7 @@ class UnifiedPeptidePropertyModel(nn.Module): - ``intensity`` : ``(B, L-1, n_ion_channels)`` in ``[0, 1]`` - ``ccs`` : ``(mean (B,), std (B,))`` — inverse ion mobility, 1/K0 - ``rt`` : ``(B,)`` — normalised retention time + - ``rt_sigma`` : ``(B,)`` — gradient-normalised RT EMG peak-width (timsim shape) - ``charge`` : ``(B, max_charge)`` logits (class j -> charge j+1) The input ``batch`` dict carries: ``tokens`` (B, L) long, ``charge`` (B,) @@ -31,6 +32,7 @@ class UnifiedPeptidePropertyModel(nn.Module): "intensity": IntensityHead, "ccs": IonMobilityHead, "rt": RetentionTimeHead, + "rt_sigma": RTSigmaHead, "charge": ChargeHead, } @@ -83,6 +85,8 @@ def forward( out["ccs"] = self.heads["ccs"](latent, batch["charge"], batch["precursor_mz"]) if "rt" in active: out["rt"] = self.heads["rt"](latent) + if "rt_sigma" in active: + out["rt_sigma"] = self.heads["rt_sigma"](latent) if "charge" in active: out["charge"] = self.heads["charge"](latent) return out diff --git a/packages/peptide-property-ng/tests/test_shape_heads.py b/packages/peptide-property-ng/tests/test_shape_heads.py new file mode 100644 index 000000000..9f3e41e4a --- /dev/null +++ b/packages/peptide-property-ng/tests/test_shape_heads.py @@ -0,0 +1,77 @@ +"""Tests for the timsim peak-shape supervision: RT-sigma head + IM-sigma loss. + +Covers the additive wiring: a new ``rt_sigma`` head, and an ``im_sigma`` loss +term riding the existing ``ccs`` head's ``std``. Both must be optional — absent +targets ⇒ no term, existing training unchanged. +""" +import dataclasses + +import torch + +from peptide_property_ng.model.config import SMALL +from peptide_property_ng.model.multitask import UnifiedPeptidePropertyModel +from peptide_property_ng.losses import MultiTaskLoss + + +def _cfg_with_shape(): + return dataclasses.replace(SMALL, tasks=tuple(SMALL.tasks) + ("rt_sigma",)) + + +def _batch(b: int = 6, length: int = 14) -> dict[str, torch.Tensor]: + return { + "tokens": torch.randint(1, 60, (b, length)), + "charge": torch.randint(1, 5, (b,)), + "precursor_mz": torch.rand(b) * 800 + 300, + "collision_energy": torch.zeros(b), + "instrument": torch.randint(0, SMALL.n_instruments, (b,)), + "acq_mode": torch.randint(0, SMALL.n_acq_modes, (b,)), + } + + +def test_rt_sigma_head_emits_positive_scalar(): + model = UnifiedPeptidePropertyModel(_cfg_with_shape()).eval() + out = model(_batch(6)) + assert "rt_sigma" in out + assert out["rt_sigma"].shape == (6,) + assert (out["rt_sigma"] > 0).all() # softplus keeps sigma strictly positive + + +def test_shape_loss_terms_present_and_masked(): + b = 6 + model = UnifiedPeptidePropertyModel(_cfg_with_shape()) + out = model(_batch(b)) + # only a subset of peptides carry a measured shape label + rt_valid = torch.tensor([True, True, False, True, False, False]) + im_valid = torch.tensor([True, False, True, False, False, True]) + batch = { + "rt_sigma_target": torch.rand(b), + "rt_sigma_valid": rt_valid, + "im_sigma_target": torch.rand(b) * 0.02, + "im_sigma_valid": im_valid, + # CCS-mean term inactive here (no ccs_valid set true) + "ccs_valid": torch.zeros(b, dtype=torch.bool), + "rt_valid": torch.zeros(b, dtype=torch.bool), + "charge_target": torch.randint(0, SMALL.max_charge, (b,)), + "intensity_target": torch.zeros(b, 13, SMALL.n_ion_channels), + } + total, logd = MultiTaskLoss()(out, batch) + assert "rt_sigma" in logd and "im_sigma" in logd + total.backward() + assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in model.parameters()) + + +def test_shape_terms_absent_when_no_target(): + """No shape targets in the batch ⇒ no shape loss terms (backward-compatible).""" + b = 5 + model = UnifiedPeptidePropertyModel(SMALL) # default tasks, no rt_sigma head + out = model(_batch(b)) + batch = { + "ccs_valid": torch.ones(b, dtype=torch.bool), + "ccs_target": torch.rand(b), + "rt_valid": torch.ones(b, dtype=torch.bool), + "rt_target": torch.rand(b), + "charge_target": torch.randint(0, SMALL.max_charge, (b,)), + "intensity_target": torch.zeros(b, 13, SMALL.n_ion_channels), + } + _total, logd = MultiTaskLoss()(out, batch) + assert "rt_sigma" not in logd and "im_sigma" not in logd From 4511e3b5b4b01dc8f3895b5a0f0a634dd9aeb171 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Thu, 11 Jun 2026 16:07:32 +0200 Subject: [PATCH 18/19] peptide-property-ng: HF-corpus data loader + train wiring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validate the published HF timsTOF corpus by training on it directly. - data/hf_corpus_dataset.py: streams Tier-1 + Tier-3 parquets (pyarrow.dataset pushdown; the 391M-row Tier-3 can't be materialised) into the same example dicts the Sage path emits. Real per-PSM CE (volts/100); prefers Sage aligned_rt via --hf-rt-lookup, else per-accession-normalised rt_seconds. - train.py: --hf-corpus / --hf-max-datasets / --hf-rt-lookup. Validation (15 ds, cap 4k, from cap0-pretrained.pt): intensity SA 0.734, ccs_mae 0.0226, RT Pearson 0.880 (aligned), charge 0.804 — in-neighborhood of the 56-dataset baseline; gap is data volume. --- .../data/hf_corpus_dataset.py | 183 ++++++++++++++++++ .../src/peptide_property_ng/train/train.py | 47 +++-- 2 files changed, 218 insertions(+), 12 deletions(-) create mode 100644 packages/peptide-property-ng/src/peptide_property_ng/data/hf_corpus_dataset.py diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/hf_corpus_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/hf_corpus_dataset.py new file mode 100644 index 000000000..af2536f17 --- /dev/null +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/hf_corpus_dataset.py @@ -0,0 +1,183 @@ +"""HF-corpus -> multi-task training examples (validation of the published corpus). + +Mirrors ``sage_dataset.prepare_examples`` but sources from the aggregated HF +corpus parquets (TO_HF_CORPUS.md): Tier-1 (per-precursor RT/IM/charge/CE) + +Tier-3 (per-precursor matched b/y fragments). Produces the SAME example dicts +``SagePropertyDataset`` wraps, so it drops into the existing collate + train loop. + +Differences vs the Sage-parquet path, by design: +- RT: Tier-1 ships ``rt_seconds`` (raw apex), not Sage ``aligned_rt`` -> we + **per-accession min-max normalise** RT so the head sees a comparable scale. +- CE: Tier-1 ships real per-precursor ``collision_energy_mean_v`` (volts) -> + ``CE = clip(volts/100, 0, 1)`` (the production convention), not a flat 0.26. +- The split is already baked into the parquet (``split`` column). +""" +from __future__ import annotations + +from collections import defaultdict +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pyarrow.dataset as pds + +from peptide_property_ng.data.fragment_targets import prosit174_to_sites +from peptide_property_ng.data.sage_dataset import (SagePropertyDataset, + normalize_inverse_mobility) + +_PROTON = 1.007276 + +_T1_COLS = ["accession", "raw_file", "precursor_id", "sage_psm_id", + "modified_sequence", "sequence", "charge", "mz", "rt_seconds", + "mobility", "collision_energy_mean_v"] + + +def load_rt_aligned_lookup(path): + """{(accession, int sage_psm_id): aligned_rt} from a lookup parquet.""" + import pyarrow.parquet as pq + t = pq.read_table(path, columns=["accession", "psm_id", "aligned_rt"]) + accs = t["accession"].to_pylist() + pids = t["psm_id"].to_pylist() + rts = t["aligned_rt"].to_pylist() + return {(a, int(p)): r for a, p, r in zip(accs, pids, rts) if p is not None} +_T3_COLS = ["accession", "raw_file", "precursor_id", "fragment_type", + "fragment_ordinal", "fragment_charge", "fragment_intensity"] + + +def _key(acc, raw, pid): + return (acc, raw, int(pid)) + + +def prepare_examples_hf(corpus_dir, split, *, cap=4000, max_len=30, seed=0, + accessions=None, rt_lookup=None): + """Build example dicts from ``corpus_dir/{tier1,tier3}/{split}.parquet``. + + rt_lookup: optional {(accession, int sage_psm_id): aligned_rt} (or a path) + — when a precursor's aligned RT is found it is used (cross-run consistent, + ~[0,1]); otherwise we fall back to per-accession-normalised raw rt_seconds. + """ + from imspy_predictors.intensity.predictors import observed_fragments_to_intensity_target + from imspy_predictors.utilities.tokenizers import ProformaTokenizer + + from peptide_property_ng.data.mass_to_unimod import parse_delta_mass_peptide + + if isinstance(rt_lookup, (str, Path)): + rt_lookup = load_rt_aligned_lookup(rt_lookup) + + corpus_dir = Path(corpus_dir) + afilter = (pds.field("accession").isin(list(accessions)) + if accessions else None) + + # Tier-3 is ~391M rows -> NEVER materialise it. Pass 1: stream Tier-1 + # (pushdown + batches), pick up to `cap` precursors/accession + props + + # per-accession RT range. Pass 2: stream Tier-3, keep only the selected + # precursors' fragments. Memory stays bounded by the selection (~50k prec). + t1set = pds.dataset(corpus_dir / "tier1" / f"{split}.parquet") + sel: dict = {} # key -> props + per_acc_count: dict = defaultdict(int) + rt_lo, rt_hi = {}, {} + for rb in t1set.scanner(columns=_T1_COLS, filter=afilter, + batch_size=131072).to_batches(): + d = rb.to_pydict() + for i in range(len(d["precursor_id"])): + acc = d["accession"][i] + rt = d["rt_seconds"][i] + if rt is not None and np.isfinite(rt): + rt_lo[acc] = rt if acc not in rt_lo else min(rt_lo[acc], rt) + rt_hi[acc] = rt if acc not in rt_hi else max(rt_hi[acc], rt) + if per_acc_count[acc] >= cap or not d["modified_sequence"][i]: + continue + key = _key(acc, d["raw_file"][i], d["precursor_id"][i]) + sel[key] = {"acc": acc, "modseq_raw": d["modified_sequence"][i], + "charge": d["charge"][i], "mz": d["mz"][i], + "rt": rt, "mobility": d["mobility"][i], + "ce_v": d["collision_energy_mean_v"][i], + "sage_psm_id": d["sage_psm_id"][i]} + per_acc_count[acc] += 1 + + by_prec: dict = defaultdict(lambda: ([], [], [], [])) + t3set = pds.dataset(corpus_dir / "tier3" / f"{split}.parquet") + for rb in t3set.scanner(columns=_T3_COLS, filter=afilter, + batch_size=262144).to_batches(): + fd = rb.to_pydict() + for acc, raw, pid, ft, fo, fc, fi in zip( + fd["accession"], fd["raw_file"], fd["precursor_id"], + fd["fragment_type"], fd["fragment_ordinal"], + fd["fragment_charge"], fd["fragment_intensity"]): + key = _key(acc, raw, pid) + if key not in sel: + continue + cols = by_prec[key] + cols[0].append(ft); cols[1].append(int(fo)) + cols[2].append(int(fc)); cols[3].append(float(fi)) + + tok = ProformaTokenizer.with_defaults() + examples: list[dict] = [] + for key, p in sel.items(): + acc = p["acc"] + cols = by_prec.get(key) + if not cols or not cols[0]: + continue # no matched fragments -> no intensity target + modseq_raw = p["modseq_raw"] + stripped, modseq = parse_delta_mass_peptide(modseq_raw) + if "[+" in modseq or "[-" in modseq: + continue # unconverted delta mass + enc = tok([modseq]) + ids, mask = enc["input_ids"][0], enc["attention_mask"][0] + n_real = int(sum(mask)) + residue_ids = ids[1:n_real - 1] + length = len(residue_ids) + if length < 3 or length > max_len: + continue + + charge = int(p["charge"]) + mz = float(p["mz"]) if p["mz"] is not None else 0.0 + frag = SimpleNamespace(ion_types=cols[0], fragment_ordinals=cols[1], + charges=cols[2], intensities=cols[3]) + prosit174 = observed_fragments_to_intensity_target(stripped, charge, frag) + target = prosit174_to_sites(prosit174, length) + + im = p["mobility"] + im = float(im) if im is not None else float("nan") + # RT target: prefer Sage aligned_rt (cross-run consistent, ~[0,1]); + # else fall back to per-accession-normalised raw rt_seconds. + al = None + if rt_lookup is not None and p["sage_psm_id"] is not None: + al = rt_lookup.get((acc, int(p["sage_psm_id"]))) + if al is not None and np.isfinite(al): + rt_norm = float(min(1.0, max(0.0, al))) + rt_valid = True + else: + rt = float(p["rt"]) if p["rt"] is not None else float("nan") + lo, hi = rt_lo.get(acc), rt_hi.get(acc) + if np.isfinite(rt) and lo is not None and hi is not None and hi > lo: + rt_norm = (rt - lo) / (hi - lo) + rt_valid = True + else: + rt_norm, rt_valid = float("nan"), False + + ce = float(np.clip((p["ce_v"] or 26.0) / 100.0, 0.0, 1.0)) + + examples.append({ + "accession": acc, "psm_id": int(key[2]), + "stripped": stripped, "modseq": modseq, + "tokens": np.asarray(residue_ids, dtype=np.int64), + "charge": charge, "precursor_mz": mz, "collision_energy": ce, + "instrument": 0, "acq_mode": 0, + "intensity_target": target, + "ccs_target": normalize_inverse_mobility(im), + "ccs_valid": bool(np.isfinite(im) and im > 0.0), + "rt_target": rt_norm, "rt_valid": rt_valid, + }) + return examples + + +def build_split_datasets_hf(corpus_dir, *, cap=4000, seed=0, accessions=None, + rt_lookup=None): + """{split: SagePropertyDataset} from the HF corpus parquets.""" + if isinstance(rt_lookup, (str, Path)): + rt_lookup = load_rt_aligned_lookup(rt_lookup) # load once, reuse per split + return {sp: SagePropertyDataset( + prepare_examples_hf(corpus_dir, sp, cap=cap, seed=seed, + accessions=accessions, rt_lookup=rt_lookup)) + for sp in ("train", "val", "test")} diff --git a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py index 1ca5e1655..d1ec08c6d 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/train/train.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/train/train.py @@ -51,6 +51,14 @@ def train_epoch(model, loader, loss_fn, optimizer, device, tasks=None) -> dict[s def main() -> None: ap = argparse.ArgumentParser(description="Train the unified peptide property model.") ap.add_argument("--datasets-glob", default="/scratch/claudius-proteomics/*") + ap.add_argument("--hf-corpus", default=None, + help="path to aggregated HF corpus dir (tier1/ + tier3/ with " + "{train,val,test}.parquet); uses the HF loader instead of Sage") + ap.add_argument("--hf-max-datasets", type=int, default=0, + help="limit the HF run to the first N accessions (0 = all)") + ap.add_argument("--hf-rt-lookup", default=None, + help="aligned_rt lookup parquet (accession,psm_id,aligned_rt); " + "uses Sage cross-run-aligned RT instead of raw rt_seconds") ap.add_argument("--catalog", default=None, help="path to timstof_catalog.tsv for per-dataset instrument/acq conditioning") ap.add_argument("--default-ce", type=float, default=0.26, @@ -92,19 +100,34 @@ def main() -> None: out_dir.mkdir(parents=True, exist_ok=True) print(f"[{time.strftime('%H:%M:%S')}] discovering datasets ...", flush=True) - sage_dirs = discover_sage_dirs(args.datasets_glob) - if args.max_datasets: - sage_dirs = sage_dirs[: args.max_datasets] - print(f" {len(sage_dirs)} datasets") - - print(f"[{time.strftime('%H:%M:%S')}] preparing examples (cap {args.cap}/dataset) ...", flush=True) - ce_cal = load_ce_calibration(args.ce_calibration) if args.ce_calibration else None - if ce_cal is not None: - print(f" CE calibration loaded: {len(ce_cal):,} keys from {args.ce_calibration}") t0 = time.time() - splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed, - catalog_path=args.catalog, default_ce=args.default_ce, - ce_calibration=ce_cal) + if args.hf_corpus: + import pyarrow.compute as _pc + import pyarrow.dataset as _pds + from peptide_property_ng.data.hf_corpus_dataset import build_split_datasets_hf + accs = None + if args.hf_max_datasets: + all_accs = sorted(_pc.unique(_pds.dataset( + f"{args.hf_corpus}/tier1/val.parquet").to_table( + columns=["accession"])["accession"]).to_pylist()) + accs = all_accs[: args.hf_max_datasets] + print(f"[{time.strftime('%H:%M:%S')}] HF corpus {args.hf_corpus} " + f"(cap {args.cap}/dataset, {len(accs) if accs else 'all'} datasets) ...", flush=True) + splits = build_split_datasets_hf(args.hf_corpus, cap=args.cap, + seed=args.seed, accessions=accs, + rt_lookup=args.hf_rt_lookup) + else: + sage_dirs = discover_sage_dirs(args.datasets_glob) + if args.max_datasets: + sage_dirs = sage_dirs[: args.max_datasets] + print(f" {len(sage_dirs)} datasets") + print(f"[{time.strftime('%H:%M:%S')}] preparing examples (cap {args.cap}/dataset) ...", flush=True) + ce_cal = load_ce_calibration(args.ce_calibration) if args.ce_calibration else None + if ce_cal is not None: + print(f" CE calibration loaded: {len(ce_cal):,} keys from {args.ce_calibration}") + splits = build_split_datasets(sage_dirs, cap=args.cap, seed=args.seed, + catalog_path=args.catalog, default_ce=args.default_ce, + ce_calibration=ce_cal) for name, ds in splits.items(): print(f" {name}: {len(ds):,} examples") print(f" prepared in {time.time() - t0:.0f}s") From 88b93e7a980bb1997802efdc92da3c09288d0715 Mon Sep 17 00:00:00 2001 From: theGreatHerrLebert Date: Fri, 12 Jun 2026 13:09:52 +0200 Subject: [PATCH 19/19] peptide-property-ng: HF loader prefers native tier1.rt_aligned (v0.2); schema-aware (falls back for v0.1) --- .../data/hf_corpus_dataset.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/packages/peptide-property-ng/src/peptide_property_ng/data/hf_corpus_dataset.py b/packages/peptide-property-ng/src/peptide_property_ng/data/hf_corpus_dataset.py index af2536f17..1cc18b852 100644 --- a/packages/peptide-property-ng/src/peptide_property_ng/data/hf_corpus_dataset.py +++ b/packages/peptide-property-ng/src/peptide_property_ng/data/hf_corpus_dataset.py @@ -29,7 +29,7 @@ _T1_COLS = ["accession", "raw_file", "precursor_id", "sage_psm_id", "modified_sequence", "sequence", "charge", "mz", "rt_seconds", - "mobility", "collision_energy_mean_v"] + "rt_aligned", "mobility", "collision_energy_mean_v"] def load_rt_aligned_lookup(path): @@ -73,10 +73,11 @@ def prepare_examples_hf(corpus_dir, split, *, cap=4000, max_len=30, seed=0, # per-accession RT range. Pass 2: stream Tier-3, keep only the selected # precursors' fragments. Memory stays bounded by the selection (~50k prec). t1set = pds.dataset(corpus_dir / "tier1" / f"{split}.parquet") + cols = [c for c in _T1_COLS if c in t1set.schema.names] # rt_aligned only in v0.2+ sel: dict = {} # key -> props per_acc_count: dict = defaultdict(int) rt_lo, rt_hi = {}, {} - for rb in t1set.scanner(columns=_T1_COLS, filter=afilter, + for rb in t1set.scanner(columns=cols, filter=afilter, batch_size=131072).to_batches(): d = rb.to_pydict() for i in range(len(d["precursor_id"])): @@ -92,7 +93,8 @@ def prepare_examples_hf(corpus_dir, split, *, cap=4000, max_len=30, seed=0, "charge": d["charge"][i], "mz": d["mz"][i], "rt": rt, "mobility": d["mobility"][i], "ce_v": d["collision_energy_mean_v"][i], - "sage_psm_id": d["sage_psm_id"][i]} + "sage_psm_id": d["sage_psm_id"][i], + "rt_aligned": d["rt_aligned"][i] if "rt_aligned" in d else None} per_acc_count[acc] += 1 by_prec: dict = defaultdict(lambda: ([], [], [], [])) @@ -139,10 +141,13 @@ def prepare_examples_hf(corpus_dir, split, *, cap=4000, max_len=30, seed=0, im = p["mobility"] im = float(im) if im is not None else float("nan") - # RT target: prefer Sage aligned_rt (cross-run consistent, ~[0,1]); - # else fall back to per-accession-normalised raw rt_seconds. + # RT target: prefer the native rt_aligned column (v0.2+), else an + # external aligned_rt lookup, else per-accession-normalised rt_seconds. al = None - if rt_lookup is not None and p["sage_psm_id"] is not None: + ra = p.get("rt_aligned") + if ra is not None and np.isfinite(ra): + al = ra + elif rt_lookup is not None and p["sage_psm_id"] is not None: al = rt_lookup.get((acc, int(p["sage_psm_id"]))) if al is not None and np.isfinite(al): rt_norm = float(min(1.0, max(0.0, al)))