rsomics_quantile_transform/
transform.rs

1//! Per-column forward transform — sklearn `QuantileTransformer._transform_col`.
2//!
3//! Forward map:
4//!   1. Double-interp averaging (handles ties):
5//!      `0.5 * (interp(x, quantiles, refs) - interp(-x, -quantiles[::-1], -refs[::-1]))`
6//!   2. Force boundary values to exactly 0/1.
7//!   3. For normal output: apply `ndtri` then clip to `[CLIP_MIN, CLIP_MAX]`.
8//!
9//! Uniform output is purely linear interpolation → BIT-EXACT vs sklearn (0 ULP).
10//! Normal output adds the Cephes `ndtri` transcendental; cross-arch last bits can
11//! differ ≤1 ULP, so compat tolerance is ≤1e-12 relative.
12
13use crate::ndtri::{CLIP_MAX, CLIP_MIN, ndtri};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum OutputDistribution {
17    Uniform,
18    Normal,
19}
20
21/// Linear interpolation matching `np.interp(x, xp, fp)`.
22///
23/// numpy's C implementation computes `slope = (fp[j+1]-fp[j])/(xp[j+1]-xp[j])` then
24/// returns `fp[j] + slope * (x - xp[j])` — which is `fp[j].mul_add(slope, dx)` in FP.
25/// Using the same `fp[i].mul_add(slope, dx)` → `fp[i] + slope*(x-xp[i])` order
26/// via `f64::mul_add` (FMA) matches numpy bit-for-bit including in the subsample regime.
27fn np_interp(x: f64, xp: &[f64], fp: &[f64]) -> f64 {
28    debug_assert_eq!(xp.len(), fp.len());
29    let n = xp.len();
30    if x <= xp[0] {
31        return fp[0];
32    }
33    if x >= xp[n - 1] {
34        return fp[n - 1];
35    }
36    let idx = xp.partition_point(|&v| v <= x);
37    let i = idx - 1;
38    let slope = (fp[i + 1] - fp[i]) / (xp[i + 1] - xp[i]);
39    // FMA: slope * (x - xp[i]) + fp[i]
40    slope.mul_add(x - xp[i], fp[i])
41}
42
43/// Transform a single column in place, matching `_transform_col(inverse=False)`.
44pub fn transform_col(
45    col: &mut [f64],
46    quantiles: &[f64],
47    references: &[f64],
48    dist: OutputDistribution,
49) {
50    let lower_bound_x = quantiles[0];
51    let upper_bound_x = quantiles[quantiles.len() - 1];
52
53    // Build reversed views for the second interp.
54    let q_rev: Vec<f64> = quantiles.iter().rev().map(|&v| -v).collect();
55    let r_rev: Vec<f64> = references.iter().rev().map(|&v| -v).collect();
56
57    for v in col.iter_mut() {
58        if v.is_nan() {
59            continue;
60        }
61        let x = *v;
62
63        // sklearn uses `== lower/upper_bound_x` for uniform boundary detection.
64        let at_lower = x == lower_bound_x;
65        let at_upper = x == upper_bound_x;
66
67        // Double-interp averaging (ties handling, per sklearn comment).
68        let fwd = np_interp(x, quantiles, references);
69        let rev = -np_interp(-x, &q_rev, &r_rev);
70        let mut y = 0.5 * (fwd + rev);
71
72        // Exact boundary values must land at exactly 0 / 1 (sklearn sets them after interp).
73        if at_upper {
74            y = 1.0;
75        }
76        if at_lower {
77            y = 0.0;
78        }
79
80        if dist == OutputDistribution::Normal {
81            y = ndtri(y).clamp(CLIP_MIN, CLIP_MAX);
82        }
83
84        *v = y;
85    }
86}
87
88/// Transform every column of the matrix (row-major `data`, `n_rows × n_cols`).
89pub fn transform_matrix(
90    data: &mut [f64],
91    n_rows: usize,
92    n_cols: usize,
93    quantiles_per_col: &[Vec<f64>],
94    references: &[f64],
95    dist: OutputDistribution,
96) {
97    // Work column by column; extract → transform → scatter back.
98    for j in 0..n_cols {
99        let mut col: Vec<f64> = (0..n_rows).map(|i| data[i * n_cols + j]).collect();
100        transform_col(&mut col, &quantiles_per_col[j], references, dist);
101        for (i, v) in col.into_iter().enumerate() {
102            data[i * n_cols + j] = v;
103        }
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    fn close(a: f64, b: f64) {
112        assert!(
113            (a - b).abs() < 1e-12,
114            "got={a} want={b} diff={}",
115            (a - b).abs()
116        );
117    }
118
119    #[test]
120    fn np_interp_basic() {
121        let xp = [0.0, 1.0, 2.0];
122        let fp = [0.0, 0.5, 1.0];
123        close(np_interp(0.5, &xp, &fp), 0.25);
124        close(np_interp(0.0, &xp, &fp), 0.0);
125        close(np_interp(2.0, &xp, &fp), 1.0);
126        close(np_interp(-1.0, &xp, &fp), 0.0); // below → fp[0]
127        close(np_interp(3.0, &xp, &fp), 1.0); // above → fp[-1]
128    }
129
130    #[test]
131    fn uniform_ties_average() {
132        // quantiles = [1,2,2,2,3], refs = [0,.25,.5,.75,1]
133        // x=2 → fwd=interp(2,[1,2,2,2,3],[0,.25,.5,.75,1])=0.25
134        //        rev=-interp(-2,[-3,-2,-2,-2,-1],[−1,−.75,−.5,−.25,0])
135        //        fwd for rev interp: x=-2 in [-3,-2,-2,-2,-1] → 0.75
136        //        rev term = -(−0.75) but wait — refs_rev = [−1,−.75,−.5,−.25,0]
137        //        -interp(-2,[-3,-2,-2,-2,-1],[-1,-.75,-.5,-.25,0]) = -(-0.75) = 0.75
138        //        average = 0.5*(0.25+0.75) = 0.5
139        let quantiles = [1.0, 2.0, 2.0, 2.0, 3.0];
140        let refs = [0.0, 0.25, 0.5, 0.75, 1.0];
141        let mut col = [2.0];
142        transform_col(&mut col, &quantiles, &refs, OutputDistribution::Uniform);
143        close(col[0], 0.5);
144    }
145
146    #[test]
147    fn boundary_forced_to_exact() {
148        let quantiles = [1.0, 2.0, 3.0];
149        let refs = [0.0, 0.5, 1.0];
150        let mut col = [1.0, 3.0];
151        transform_col(&mut col, &quantiles, &refs, OutputDistribution::Uniform);
152        assert_eq!(col[0], 0.0);
153        assert_eq!(col[1], 1.0);
154    }
155}
rsomics_quantile_transform/transform.rs

rsomics_quantile_transform/
transform.rs
