koth_ms/output/
align_report.rs

1//! JSON diagnostics report for the `koth_align` stage.
2//!
3//! Captures per-anchor data + fitted models so downstream plotting scripts can
4//! visualise RT warp, mass-ppm drift, IM drift, and per-run quality without
5//! re-running anything. Consensus feature data is intentionally NOT duplicated
6//! here — it already lives in `consensus_features.tsv` / `intensity_matrix.tsv`.
7
8use std::fs::File;
9use std::io::BufWriter;
10use std::path::Path;
11use std::time::Instant;
12
13use rand::seq::SliceRandom;
14use rand::SeedableRng;
15use rayon::prelude::*;
16use serde::Serialize;
17
18/// Max anchor records emitted per run in the JSON report.
19/// `residual_stats` still aggregates over every anchor; only the per-anchor
20/// detail array gets capped. With 20 runs and ~70k anchors each, the full
21/// list would balloon the JSON past 800 MB and turn the writer into the new
22/// bottleneck — see lfq/mod.rs for the matching loop-hoisting fix.
23const MAX_ANCHORS_PER_RUN: usize = 500;
24
25use crate::alignment::{drift::DriftFit, warp::RtWarp, AlignmentResult, RunAlignment, RunInput};
26use crate::config::AlignConfig;
27use crate::error::KothError;
28use crate::lfq::IntensityMatrix;
29
30// ── helpers ──────────────────────────────────────────────────────────────────
31
32fn median_sorted(sorted: &[f64]) -> f64 {
33    let n = sorted.len();
34    if n == 0 {
35        return f64::NAN;
36    }
37    if n % 2 == 0 {
38        (sorted[n / 2 - 1] + sorted[n / 2]) / 2.0
39    } else {
40        sorted[n / 2]
41    }
42}
43
44fn percentile_sorted(sorted: &[f64], p: f64) -> f64 {
45    let n = sorted.len();
46    if n == 0 {
47        return f64::NAN;
48    }
49    let idx = ((p * (n as f64 - 1.0)).round() as usize).min(n - 1);
50    sorted[idx]
51}
52
53fn summary(values: &[f64]) -> ResidualSummary {
54    if values.is_empty() {
55        return ResidualSummary {
56            n: 0,
57            median: f64::NAN,
58            mad: f64::NAN,
59            std: f64::NAN,
60        };
61    }
62    let mut sorted: Vec<f64> = values.iter().copied().filter(|v| v.is_finite()).collect();
63    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
64    let med = median_sorted(&sorted);
65    let mut abs_dev: Vec<f64> = sorted.iter().map(|x| (x - med).abs()).collect();
66    abs_dev.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
67    let mad = median_sorted(&abs_dev);
68    ResidualSummary {
69        n: sorted.len(),
70        median: med,
71        mad,
72        std: crate::stats::std_dev(&sorted),
73    }
74}
75
76// ── serialisable types ───────────────────────────────────────────────────────
77
78#[derive(Serialize)]
79pub struct ResidualSummary {
80    pub n: usize,
81    pub median: f64,
82    pub mad: f64,
83    pub std: f64,
84}
85
86#[derive(Serialize)]
87pub struct AnchorRecord {
88    pub ref_rt_norm: f64,
89    pub run_rt_norm: f64,
90    pub ref_rt: f64,
91    pub run_rt: f64,
92    pub ref_mz: f64,
93    pub run_mz: f64,
94    pub ppm_error: f64,
95    pub ref_im: f64,
96    pub run_im: f64,
97    pub im_delta: f64,
98    pub rt_residual_after_warp: f64,
99    pub ppm_residual_after_fit: f64,
100    pub im_residual_after_fit: f64,
101    pub rt_active: bool,
102    pub mass_active: bool,
103    pub im_active: bool,
104}
105
106#[derive(Serialize)]
107pub struct ResidualStats {
108    pub ppm_before: ResidualSummary,
109    pub ppm_after: ResidualSummary,
110    pub im_before: ResidualSummary,
111    pub im_after: ResidualSummary,
112    pub rt_norm_before: ResidualSummary,
113    pub rt_norm_after: ResidualSummary,
114}
115
116#[derive(Serialize)]
117pub struct AlignmentDiagnostics {
118    pub n_anchors: usize,
119    pub n_anchors_rt_active: usize,
120    pub n_anchors_mass_active: usize,
121    pub n_anchors_im_active: usize,
122    pub rt_warp: RtWarp,
123    pub mass_drift: DriftFit,
124    pub im_drift: DriftFit,
125    pub anchors: Vec<AnchorRecord>,
126    pub residual_stats: ResidualStats,
127}
128
129#[derive(Serialize)]
130pub struct LfqRunSummary {
131    pub n_detected: usize,
132    pub detection_rate: f64,
133    pub qvalue_percentiles: QvaluePercentiles,
134}
135
136#[derive(Serialize)]
137pub struct QvaluePercentiles {
138    pub p25: f64,
139    pub p50: f64,
140    pub p75: f64,
141    pub p95: f64,
142}
143
144#[derive(Serialize)]
145pub struct RunAlignReport {
146    pub name: String,
147    pub is_reference: bool,
148    pub rt_range: [f64; 2],
149    pub n_features: usize,
150    pub n_features_high_score: usize,
151    pub n_hills: usize,
152    /// Absent for the reference run.
153    #[serde(skip_serializing_if = "Option::is_none")]
154    pub alignment: Option<AlignmentDiagnostics>,
155    pub lfq: LfqRunSummary,
156}
157
158#[derive(Serialize)]
159pub struct Timing {
160    pub alignment_sec: f64,
161    pub lfq_sec: f64,
162    pub total_sec: f64,
163}
164
165#[derive(Serialize)]
166pub struct AlignReport {
167    pub schema_version: u32,
168    pub reference_run: String,
169    pub n_runs: usize,
170    pub n_consensus_features: usize,
171    pub config: AlignConfig,
172    pub runs: Vec<RunAlignReport>,
173    pub timing: Timing,
174}
175
176// ── builder ──────────────────────────────────────────────────────────────────
177
178fn build_anchors_block(al: &RunAlignment) -> AlignmentDiagnostics {
179    let (run_lo, run_hi) = al.run_rt_range;
180    let (ref_lo, ref_hi) = al.ref_rt_range;
181    let run_span = run_hi - run_lo;
182    let ref_span = ref_hi - ref_lo;
183
184    let mut records: Vec<AnchorRecord> = Vec::with_capacity(al.anchors.len());
185
186    let mut ppm_before: Vec<f64> = Vec::with_capacity(al.anchors.len());
187    let mut ppm_after: Vec<f64> = Vec::with_capacity(al.anchors.len());
188    let mut im_before: Vec<f64> = Vec::new();
189    let mut im_after: Vec<f64> = Vec::new();
190    let mut rt_before: Vec<f64> = Vec::with_capacity(al.anchors.len());
191    let mut rt_after: Vec<f64> = Vec::with_capacity(al.anchors.len());
192
193    for (i, a) in al.anchors.iter().enumerate() {
194        let ppm = a.ppm_error();
195        let im_d = a.im_delta();
196        let rt_norm_delta = a.ref_rt_norm - a.run_rt_norm;
197
198        // Apply the fitted models to compute residuals.
199        let warp_predicted_delta = al.rt_warp.delta_at_run_norm(a.run_rt_norm);
200        let rt_resid = rt_norm_delta - warp_predicted_delta;
201
202        let ppm_resid = ppm - al.mass_drift.predict(a.ref_rt_norm);
203
204        // IM is only meaningful when both sides have IM.
205        let im_resid = if a.ref_im != 0.0 && a.run_im != 0.0 {
206            im_d - al.im_drift.predict(a.ref_rt_norm)
207        } else {
208            f64::NAN
209        };
210
211        // Absolute RTs derived from the per-run RT span.
212        let run_rt_abs = run_lo + a.run_rt_norm * run_span;
213        let ref_rt_abs = ref_lo + a.ref_rt_norm * ref_span;
214
215        let rt_active = *al.rt_active.get(i).unwrap_or(&false);
216        let mass_active = *al.mass_active.get(i).unwrap_or(&false);
217        let im_active = *al.im_active.get(i).unwrap_or(&false);
218
219        records.push(AnchorRecord {
220            ref_rt_norm: a.ref_rt_norm,
221            run_rt_norm: a.run_rt_norm,
222            ref_rt: ref_rt_abs,
223            run_rt: run_rt_abs,
224            ref_mz: a.ref_mz,
225            run_mz: a.run_mz,
226            ppm_error: ppm,
227            ref_im: a.ref_im,
228            run_im: a.run_im,
229            im_delta: im_d,
230            rt_residual_after_warp: rt_resid,
231            ppm_residual_after_fit: ppm_resid,
232            im_residual_after_fit: im_resid,
233            rt_active,
234            mass_active,
235            im_active,
236        });
237
238        ppm_before.push(ppm);
239        ppm_after.push(ppm_resid);
240        if a.ref_im != 0.0 && a.run_im != 0.0 {
241            im_before.push(im_d);
242            im_after.push(im_resid);
243        }
244        rt_before.push(rt_norm_delta);
245        rt_after.push(rt_resid);
246    }
247
248    // Subsample the per-anchor detail to keep the JSON report tractable.
249    // Residual stats above were computed over the full anchor set, so the
250    // summary numbers stay exact — only the scatter-plot detail is capped.
251    // Deterministic seed so re-running on the same input produces the same
252    // sample (helpful for diffing reports).
253    let n_total = records.len();
254    let sampled_records = if n_total > MAX_ANCHORS_PER_RUN {
255        let mut rng = rand::rngs::StdRng::seed_from_u64(0x0a11_c0de);
256        records.shuffle(&mut rng);
257        records.truncate(MAX_ANCHORS_PER_RUN);
258        records
259    } else {
260        records
261    };
262
263    AlignmentDiagnostics {
264        n_anchors: al.n_anchors,
265        n_anchors_rt_active: al.rt_active.iter().filter(|&&v| v).count(),
266        n_anchors_mass_active: al.mass_active.iter().filter(|&&v| v).count(),
267        n_anchors_im_active: al.im_active.iter().filter(|&&v| v).count(),
268        rt_warp: al.rt_warp.clone(),
269        mass_drift: al.mass_drift.clone(),
270        im_drift: al.im_drift.clone(),
271        anchors: sampled_records,
272        residual_stats: ResidualStats {
273            ppm_before: summary(&ppm_before),
274            ppm_after: summary(&ppm_after),
275            im_before: summary(&im_before),
276            im_after: summary(&im_after),
277            rt_norm_before: summary(&rt_before),
278            rt_norm_after: summary(&rt_after),
279        },
280    }
281}
282
283fn build_lfq_summary(matrix: &IntensityMatrix, run_idx: usize) -> LfqRunSummary {
284    let mut detected_q: Vec<f64> = Vec::new();
285    for feat in 0..matrix.n_features {
286        let intensity = matrix.intensity(feat, run_idx);
287        if intensity > 0.0 {
288            detected_q.push(matrix.q_value(feat, run_idx));
289        }
290    }
291    let n_detected = detected_q.len();
292    detected_q.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
293
294    let qp = QvaluePercentiles {
295        p25: percentile_sorted(&detected_q, 0.25),
296        p50: percentile_sorted(&detected_q, 0.50),
297        p75: percentile_sorted(&detected_q, 0.75),
298        p95: percentile_sorted(&detected_q, 0.95),
299    };
300
301    let detection_rate = if matrix.n_features == 0 {
302        0.0
303    } else {
304        n_detected as f64 / matrix.n_features as f64
305    };
306
307    LfqRunSummary {
308        n_detected,
309        detection_rate,
310        qvalue_percentiles: qp,
311    }
312}
313
314pub fn build_align_report(
315    runs: &[RunInput],
316    alignment: &AlignmentResult,
317    matrix: &IntensityMatrix,
318    config: &AlignConfig,
319    timing: Timing,
320) -> AlignReport {
321    let t_outer = Instant::now();
322    let min_combined = config.alignment.min_anchor_combined_score;
323
324    // Each run's block is independent — parallelise across runs. The big
325    // per-run cost is `build_anchors_block`, which does six sorts over the
326    // ~70k anchors (for residual stats) plus the per-anchor record build.
327    // Precompute `rt_range` for every run up front: `RunInput::rt_range`
328    // falls back to a full feature sweep when `scan_times` is empty (which
329    // it is in the align pipeline), so calling it inside the per-run closure
330    // would be the same trap as before — only this time the alignment table
331    // already stores it, so reuse that and fall back for the reference run.
332    let rt_ranges: Vec<(f64, f64)> = runs
333        .par_iter()
334        .enumerate()
335        .map(|(run_idx, ri)| {
336            if run_idx == alignment.reference_idx {
337                ri.rt_range()
338            } else {
339                alignment
340                    .alignments
341                    .get(&ri.name)
342                    .map(|al| al.run_rt_range)
343                    .unwrap_or_else(|| ri.rt_range())
344            }
345        })
346        .collect();
347
348    let runs_out: Vec<RunAlignReport> = runs
349        .par_iter()
350        .enumerate()
351        .map(|(run_idx, ri)| {
352            let is_reference = run_idx == alignment.reference_idx;
353            let rt_range = rt_ranges[run_idx];
354            let n_features = ri.features.len();
355            let n_features_high_score = ri
356                .features
357                .iter()
358                .filter(|f| f.combined_score >= min_combined)
359                .count();
360            // Hills are streamed per-run during LFQ (not held on `ri`), so the
361            // count comes from the matrix, which records it as each run loads.
362            let n_hills = matrix.hills_per_run.get(run_idx).copied().unwrap_or(0);
363
364            let alignment_block = if is_reference {
365                None
366            } else {
367                alignment.alignments.get(&ri.name).map(build_anchors_block)
368            };
369
370            let lfq = build_lfq_summary(matrix, run_idx);
371
372            RunAlignReport {
373                name: ri.name.clone(),
374                is_reference,
375                rt_range: [rt_range.0, rt_range.1],
376                n_features,
377                n_features_high_score,
378                n_hills,
379                alignment: alignment_block,
380                lfq,
381            }
382        })
383        .collect();
384
385    let report = AlignReport {
386        schema_version: 1,
387        reference_run: alignment.reference_name.clone(),
388        n_runs: runs.len(),
389        n_consensus_features: matrix.n_features,
390        config: config.clone(),
391        runs: runs_out,
392        timing,
393    };
394    log::info!(
395        "[timing] build_align_report inner: {:.2?}",
396        t_outer.elapsed()
397    );
398    report
399}
400
401// ── writer ───────────────────────────────────────────────────────────────────
402
403pub fn write_align_report(report: &AlignReport, path: &Path) -> Result<(), KothError> {
404    // Stream directly to disk via a buffered writer to avoid materialising the
405    // full JSON in memory (~hundreds of MB on multi-run datasets) and to batch
406    // the file I/O. Compact (non-pretty) output cuts the file ~3× by dropping
407    // whitespace — the schema is consumed by `scripts/plot_align_report.py`
408    // via `json.load`, which doesn't care about formatting.
409    let file = File::create(path)?;
410    let mut writer = BufWriter::with_capacity(1 << 20, file); // 1 MiB buffer
411    serde_json::to_writer(&mut writer, report)?;
412    use std::io::Write;
413    writer.flush()?;
414    log::info!("Wrote align report to {}", path.display());
415    Ok(())
416}
koth_ms/output/align_report.rs

koth_ms/output/
align_report.rs
