diff --git a/rust/timsseek/src/fragment_mass/elution_group_converter.rs b/rust/timsseek/src/fragment_mass/elution_group_converter.rs index cd428283..02bf67ba 100644 --- a/rust/timsseek/src/fragment_mass/elution_group_converter.rs +++ b/rust/timsseek/src/fragment_mass/elution_group_converter.rs @@ -7,9 +7,9 @@ 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: /// ``` @@ -17,20 +17,20 @@ use rustyms::prelude::{ /// 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) { diff --git a/rust/timsseek/src/scoring/offsets.rs b/rust/timsseek/src/scoring/offsets.rs index adf3f2e8..c3269ae8 100644 --- a/rust/timsseek/src/scoring/offsets.rs +++ b/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,17 +122,23 @@ 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 } @@ -136,7 +146,7 @@ impl MzMobilityOffsets { 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, ); } diff --git a/rust/timsseek/src/scoring/results.rs b/rust/timsseek/src/scoring/results.rs index 4a0a4064..b6ad1843 100644 --- a/rust/timsseek/src/scoring/results.rs +++ b/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 diff --git a/rust/timsseek/src/utils/top_n_array.rs b/rust/timsseek/src/utils/top_n_array.rs index cf7ad5b5..4db82445 100644 --- a/rust/timsseek/src/utils/top_n_array.rs +++ b/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 TopNArray { self.len == 0 } - pub fn get_values(&self) -> [T; N] { + pub fn get_values_sorted(&self) -> [T; N] { self.array } diff --git a/scripts/refit_mobility.py b/scripts/refit_mobility.py new file mode 100755 index 00000000..9e49f2d7 --- /dev/null +++ b/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())