peptide-property-ng: unified Depthcharge peptide property predictor

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/

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 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 |
+
+## 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/*'
+```

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"]

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"

packages/peptide-property-ng/src/peptide_property_ng/data/__init__.py

@@ -0,0 +1 @@
+"""Data — Sage-parquet dataset, fragment-intensity targets, splits, collation."""

packages/peptide-property-ng/src/peptide_property_ng/data/ccs_pretrain.py

@@ -0,0 +1,137 @@
+"""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+\]")
+
+# 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], 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:
+    """ionmob ``sequence-tokenized`` list -> a ProForma string, or ``None`` to skip.
+
+    A ``<START>`` / ``<END>`` marker carrying a terminal modification (e.g.
+    ``<START>[UNIMOD:1]``) flags a terminal mod -> skip, consistent with the
+    other loaders.
+    """
+    if len(tokens) < 5 or tokens[0] != "<START>" or tokens[-1] != "<END>":
+        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
+            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": 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": normalize_ccs(ccs_val),
+                    "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))

packages/peptide-property-ng/src/peptide_property_ng/data/chronologer_rt.py

@@ -0,0 +1,106 @@
+"""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:
+    """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(
+    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 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:
+        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_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)
+    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
+            rt = _normalise_hi(hi[start + j])
+            if not np.isfinite(rt):
+                continue  # non-finite HI -> useless as an RT example
+            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": rt,
+                    "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))