fix(timsseek): correct obs_mobility unit + refit supersimpleprediction

rust/timsseek/src/fragment_mass/elution_group_converter.rs

@@ -7,30 +7,30 @@ use rustyms::prelude::{
 
 /// Super simple 1/k0 prediction.
 ///
-/// This is a simple prediction of the retention time based on the m/z and charge.
-/// On my data it gets MAPE 1.82802 so, this prediction + 10% error is a pretty solid way
-/// to set an extraction window for mobility if you dont know anything for the peptide.
+/// Refit on `hela_iccoff_gt20peps` (34k IDs, holdout MAPE 1.36%) after the
+/// mobility unit-bug fix. Use `scripts/refit_mobility.py` against a fresh
+/// `results.parquet` to refit on a different dataset.
 ///
 /// Example:
 /// ```
 /// use timsseek::fragment_mass::elution_group_converter::supersimpleprediction;
 /// let mass = 1810.917339999999;
 /// let charge = 2;
 /// let out = supersimpleprediction(mass / charge as f64, charge);
-/// assert!((out - 1.105151).abs() < 0.001 );
+/// assert!((out - 1.144405).abs() < 0.001);
 /// ```
 pub fn supersimpleprediction(mz: f64, charge: i32) -> f64 {
-    let intercept_ = -1.660e+00;
+    let intercept_ = -1.319388e+00;
     let log1p_mz = (mz + 1.).ln();
     let sq_mz_over_charge = mz.powi(2) / charge as f64;
     let log1p_sq_mz_over_charge = (sq_mz_over_charge + 1.).ln();
 
     intercept_
-        + (-3.798e-01 * log1p_mz)
-        + (-2.389e-04 * mz)
-        + (3.957e-01 * log1p_sq_mz_over_charge)
-        + (4.157e-07 * sq_mz_over_charge)
-        + (1.417e-01 * charge as f64)
+        + (-2.954677e-01 * log1p_mz)
+        + (-9.277763e-05 * mz)
+        + (3.219103e-01 * log1p_sq_mz_over_charge)
+        + (4.005229e-07 * sq_mz_over_charge)
+        + (1.176651e-01 * charge as f64)
 }
 
 fn count_carbon_sulphur(form: &MolecularFormula) -> (u16, u16) {

rust/timsseek/src/scoring/offsets.rs

@@ -14,6 +14,10 @@ const PRECURSOR_TOP_N: usize = 3;
 /// Balances statistical power with outlier resistance.
 const FRAGMENT_TOP_N: usize = 7;
 
+/// Number of top fragments contributing to the obs-mobility estimate.
+/// Subset of FRAGMENT_TOP_N to bias toward the highest-confidence ions.
+const FRAGMENT_OBS_MOB_TOP_N: usize = 3;
+
 /// Container for measured m/z and mobility offsets from top ions.
 ///
 /// Tracks the highest-intensity precursors and fragments to calculate
@@ -118,25 +122,31 @@ impl MzMobilityOffsets {
     /// Used by Phase-2 calibration to estimate population-level offsets;
     /// rescoring uses the per-ion arrays directly.
     pub fn weighted_ms1(&self) -> Option<(f32, f32)> {
-        let (mut w, mut mz, mut mob) = (0.0f64, 0.0f64, 0.0f64);
-        for v in self.ms1.get_values() {
-            if v.intensity <= 0.0 || v.mz_error_ppm.is_nan() {
+        let (mut w_mz, mut mz) = (0.0f64, 0.0f64);
+        let (mut w_mob, mut mob) = (0.0f64, 0.0f64);
+        for v in self.ms1.get_values_sorted() {
+            if v.intensity <= 0.0 {
                 continue;
             }
-            w += v.intensity;
-            mz += v.intensity * v.mz_error_ppm as f64;
-            mob += v.intensity * v.mobility_error_pct as f64;
+            if !v.mz_error_ppm.is_nan() {
+                w_mz += v.intensity;
+                mz += v.intensity * v.mz_error_ppm as f64;
+            }
+            if !v.mobility_error_pct.is_nan() {
+                w_mob += v.intensity;
+                mob += v.intensity * v.mobility_error_pct as f64;
+            }
         }
-        if w > 0.0 {
-            Some(((mz / w) as f32, (mob / w) as f32))
+        if w_mz > 0.0 && w_mob > 0.0 {
+            Some(((mz / w_mz) as f32, (mob / w_mob) as f32))
         } else {
             None
         }
     }
 
     pub fn ms1_mz_errors(&self) -> [f32; PRECURSOR_TOP_N] {
         let mut out = [0.0; PRECURSOR_TOP_N];
-        let vals = self.ms1.get_values();
+        let vals = self.ms1.get_values_sorted();
         for i in 0..PRECURSOR_TOP_N {
             out[i] = vals[i].mz_error_ppm;
         }
@@ -145,7 +155,7 @@ impl MzMobilityOffsets {
 
     pub fn ms2_mz_errors(&self) -> [f32; FRAGMENT_TOP_N] {
         let mut out = [0.0; FRAGMENT_TOP_N];
-        let vals = self.ms2.get_values();
+        let vals = self.ms2.get_values_sorted();
         for i in 0..FRAGMENT_TOP_N {
             out[i] = vals[i].mz_error_ppm;
         }
@@ -154,7 +164,7 @@ impl MzMobilityOffsets {
 
     pub fn ms1_mobility_errors(&self) -> [f32; PRECURSOR_TOP_N] {
         let mut out = [0.0; PRECURSOR_TOP_N];
-        let vals = self.ms1.get_values();
+        let vals = self.ms1.get_values_sorted();
         for i in 0..PRECURSOR_TOP_N {
             out[i] = vals[i].mobility_error_pct;
         }
@@ -163,37 +173,41 @@ impl MzMobilityOffsets {
 
     pub fn ms2_mobility_errors(&self) -> [f32; FRAGMENT_TOP_N] {
         let mut out = [0.0; FRAGMENT_TOP_N];
-        let vals = self.ms2.get_values();
+        let vals = self.ms2.get_values_sorted();
         for i in 0..FRAGMENT_TOP_N {
             out[i] = vals[i].mobility_error_pct;
         }
         out
     }
 
+    /// Intensity-weighted absolute mobility deltas (1/k0 units) for MS1 and MS2.
+    /// Converts the stored percent error back to absolute via `ref_mobility`.
+    /// The collector's "mz" slot carries the ppm error (unused by current
+    /// consumers); only `mean_mobility()` is meaningful here.
     pub fn avg_delta_mobs(&self) -> (MzMobilityStatsCollector, MzMobilityStatsCollector) {
-        let mut ms2 = MzMobilityStatsCollector::default();
         let mut ms1 = MzMobilityStatsCollector::default();
-        let vals = self.ms2.get_values();
-        for v in vals.iter().take(3) {
+        let mut ms2 = MzMobilityStatsCollector::default();
+        let pct_to_abs = self.ref_mobility / 100.0;
+
+        for v in self.ms2.get_values_sorted().iter().take(FRAGMENT_OBS_MOB_TOP_N) {
             if v.mobility_error_pct.is_nan() {
                 continue;
             }
             ms2.add(
                 v.intensity,
                 v.mz_error_ppm as f64,
-                v.mobility_error_pct as f64,
+                v.mobility_error_pct as f64 * pct_to_abs,
             );
         }
 
-        let vals = self.ms1.get_values();
-        for v in vals.iter() {
+        for v in self.ms1.get_values_sorted().iter() {
             if v.mobility_error_pct.is_nan() {
                 continue;
             }
             ms1.add(
                 v.intensity,
                 v.mz_error_ppm as f64,
-                v.mobility_error_pct as f64,
+                v.mobility_error_pct as f64 * pct_to_abs,
             );
         }
 

rust/timsseek/src/scoring/results.rs

@@ -299,11 +299,11 @@ impl ScoredCandidateBuilder {
         self.ms2_mz_errors = SetField::Some(offsets.ms2_mz_errors());
         self.ms2_mobility_errors = SetField::Some(offsets.ms2_mobility_errors());
 
-        let mob_errors = offsets.avg_delta_mobs();
-        let cum_err = mob_errors.0 + mob_errors.1;
+        let (ms1_err, ms2_err) = offsets.avg_delta_mobs();
+        let cum_err = ms1_err + ms2_err;
         let obs_mob = (offsets.ref_mobility + cum_err.mean_mobility().unwrap_or(f64::NAN)) as f32;
-        let d_err = match (mob_errors.0.mean_mobility(), mob_errors.1.mean_mobility()) {
-            (Ok(mz), Ok(mob)) => mz - mob,
+        let d_err = match (ms1_err.mean_mobility(), ms2_err.mean_mobility()) {
+            (Ok(ms1_mob), Ok(ms2_mob)) => ms1_mob - ms2_mob,
             _ => f64::NAN,
         };
         self.delta_ms1_ms2_mobility = SetField::Some(d_err as f32);
@@ -404,8 +404,8 @@ impl ScoredCandidateBuilder {
         let sq_delta_ms1_ms2_mobility = delta_ms1_ms2_mobility * delta_ms1_ms2_mobility;
 
         let relints = expect_some!(relative_intensities);
-        let ms1_intensity_ratios = relints.ms1.get_values();
-        let ms2_intensity_ratios = relints.ms2.get_values();
+        let ms1_intensity_ratios = relints.ms1.get_values_sorted();
+        let ms2_intensity_ratios = relints.ms2.get_values_sorted();
 
         let scoring = ScoringFields {
             // Identity

rust/timsseek/src/utils/top_n_array.rs

@@ -16,7 +16,7 @@
 /// top_3.push(3.0);
 /// top_3.push(12.0);
 ///
-/// assert_eq!(top_3.get_values(), [12.0, 3.0, 3.0]);
+/// assert_eq!(top_3.get_values_sorted(), [12.0, 3.0, 3.0]);
 /// assert_eq!(top_3.max_val(), 12.0);
 /// assert_eq!(top_3.min_val(), 3.0);
 /// ```
@@ -60,7 +60,7 @@ impl<const N: usize, T: Clone + Copy + PartialOrd + Default> TopNArray<N, T> {
         self.len == 0
     }
 
-    pub fn get_values(&self) -> [T; N] {
+    pub fn get_values_sorted(&self) -> [T; N] {
         self.array
     }
 

scripts/refit_mobility.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env -S uv run --script
+# /// script
+# requires-python = ">=3.12"
+# dependencies = ["polars", "numpy", "scikit-learn"]
+# ///
+"""Refit `supersimpleprediction` 1/k0 model from a timsseek results.parquet.
+
+Mirrors the feature set in
+rust/timsseek/src/fragment_mass/elution_group_converter.rs:supersimpleprediction.
+Prints the new intercept + coefs ready to paste back into the Rust source.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import numpy as np
+import polars as pl
+from sklearn.linear_model import HuberRegressor, LinearRegression
+
+
+def build_features(mz: np.ndarray, z: np.ndarray) -> np.ndarray:
+    sq_mz_over_z = mz**2 / z
+    return np.column_stack([
+        np.log1p(mz),
+        mz,
+        np.log1p(sq_mz_over_z),
+        sq_mz_over_z,
+        z,
+    ])
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("parquet", type=Path, help="path to results.parquet")
+    p.add_argument(
+        "--min-main-score",
+        type=float,
+        default=0.0,
+        help="filter rows with main_score < threshold (default: 0.0)",
+    )
+    p.add_argument(
+        "--robust",
+        action="store_true",
+        help="use HuberRegressor instead of OLS",
+    )
+    p.add_argument(
+        "--holdout",
+        type=float,
+        default=0.2,
+        help="fraction of rows reserved for holdout MAPE (default: 0.2)",
+    )
+    p.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+    )
+    args = p.parse_args()
+
+    if not args.parquet.exists():
+        print(f"parquet not found: {args.parquet}", file=sys.stderr)
+        return 1
+
+    df = pl.read_parquet(args.parquet).filter(
+        pl.col("obs_mobility").is_finite()
+        & pl.col("precursor_mz").is_finite()
+        & (pl.col("main_score") > args.min_main_score)
+    )
+    print(f"rows after filter: {df.height}")
+    if df.height < 100:
+        print("not enough rows to fit", file=sys.stderr)
+        return 1
+
+    mz = df["precursor_mz"].to_numpy().astype(np.float64)
+    z = df["precursor_charge"].cast(pl.Float64).to_numpy()
+    y = df["obs_mobility"].to_numpy().astype(np.float64)
+
+    X = build_features(mz, z)
+
+    rng = np.random.default_rng(args.seed)
+    idx = rng.permutation(len(y))
+    cut = int(len(y) * (1 - args.holdout))
+    tr, ho = idx[:cut], idx[cut:]
+
+    model = HuberRegressor(max_iter=500) if args.robust else LinearRegression()
+    model.fit(X[tr], y[tr])
+
+    def mape(a: np.ndarray, b: np.ndarray) -> float:
+        return float(np.mean(np.abs((a - b) / b)) * 100)
+
+    tr_mape = mape(model.predict(X[tr]), y[tr])
+    ho_mape = mape(model.predict(X[ho]), y[ho])
+
+    feats = [
+        "log1p_mz",
+        "mz",
+        "log1p_sq_mz_over_charge",
+        "sq_mz_over_charge",
+        "charge",
+    ]
+    print()
+    print(f"intercept: {model.intercept_:.6e}")
+    for name, coef in zip(feats, model.coef_):
+        print(f"  {name:>24s}: {coef:+.6e}")
+    print()
+    print(f"train MAPE: {tr_mape:.4f}%")
+    print(f"holdout MAPE: {ho_mape:.4f}%")
+
+    print()
+    print("--- Rust paste block ---")
+    print(f"    let intercept_ = {model.intercept_:.3e};")
+    print("    let log1p_mz = (mz + 1.).ln();")
+    print("    let sq_mz_over_charge = mz.powi(2) / charge as f64;")
+    print("    let log1p_sq_mz_over_charge = (sq_mz_over_charge + 1.).ln();")
+    print()
+    print("    intercept_")
+    c = model.coef_
+    print(f"        + ({c[0]:+.3e} * log1p_mz)")
+    print(f"        + ({c[1]:+.3e} * mz)")
+    print(f"        + ({c[2]:+.3e} * log1p_sq_mz_over_charge)")
+    print(f"        + ({c[3]:+.3e} * sq_mz_over_charge)")
+    print(f"        + ({c[4]:+.3e} * charge as f64)")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())