Skip to main content

ggplot_rs/stat/
smooth.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// Smoothing method selection.
8#[derive(Clone, Debug, Default)]
9pub enum SmoothMethod {
10    /// Linear regression (y = mx + b).
11    #[default]
12    Lm,
13    /// LOESS with configurable span.
14    Loess { span: f64 },
15    /// Generalized linear model via anofox-regression (Gaussian or Poisson).
16    #[cfg(feature = "regression")]
17    Glm { family: SmoothFamily },
18    /// Robust linear regression (Huber M-estimator) via anofox-regression.
19    #[cfg(feature = "regression")]
20    Rlm,
21    /// Penalized B-spline (P-spline) GAM smoother — ggplot2's
22    /// `method = "gam"`. λ is chosen by GCV. Backed by anofox-regression's
23    /// `PSplineRegressor` when the `regression` feature is enabled, or by a
24    /// small dependency-free (wasm-safe) P-spline fit otherwise — see
25    /// [`StatSmooth::compute_gam`].
26    Gam,
27}
28
29/// GLM family for regression-backed smoothing (`SmoothMethod::Glm`), R's
30/// `glm(family = …)`.
31///
32/// For every non-Gaussian family the confidence band is computed as ggplot2's
33/// `predictdf.glm` does: the linear predictor and its standard error come from
34/// `predict(type = "link", se.fit = TRUE)`, the interval
35/// `η ± qnorm(0.975)·se(η)` is formed on the link scale and both ends are
36/// mapped through the inverse link — so the band respects the response range
37/// (probabilities stay in `[0, 1]`, means stay positive). Dispersion follows
38/// R: fixed at 1 for binomial / Poisson / negative binomial, Pearson-estimated
39/// for Gamma.
40///
41/// `Gaussian` keeps the ordinary-least-squares fit with its `t`-based interval
42/// (identical to `method = "lm"`).
43#[cfg(feature = "regression")]
44#[derive(Clone, Copy, Debug, Default, PartialEq)]
45pub enum SmoothFamily {
46    /// Ordinary least squares (identity link).
47    #[default]
48    Gaussian,
49    /// Poisson regression (log link) for count responses.
50    Poisson,
51    /// Binomial regression for 0/1 (or proportion) responses —
52    /// `binomial(link)`, R's default link is logit.
53    Binomial(SmoothBinomialLink),
54    /// Gamma regression for positive continuous responses — `Gamma(link)`;
55    /// R's default link is the inverse.
56    Gamma(SmoothGammaLink),
57    /// Negative-binomial regression (log link, θ estimated by maximum
58    /// likelihood) for over-dispersed counts — R's `MASS::glm.nb`.
59    NegativeBinomial,
60}
61
62#[cfg(feature = "regression")]
63impl SmoothFamily {
64    /// `binomial("logit")` — logistic regression.
65    pub fn binomial() -> Self {
66        SmoothFamily::Binomial(SmoothBinomialLink::Logit)
67    }
68
69    /// `Gamma("inverse")` — R's canonical Gamma link.
70    pub fn gamma() -> Self {
71        SmoothFamily::Gamma(SmoothGammaLink::Inverse)
72    }
73}
74
75/// Link function for [`SmoothFamily::Binomial`].
76#[cfg(feature = "regression")]
77#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
78pub enum SmoothBinomialLink {
79    /// `log(μ / (1 − μ))` (canonical).
80    #[default]
81    Logit,
82    /// `Φ⁻¹(μ)`.
83    Probit,
84    /// `log(−log(1 − μ))`.
85    Cloglog,
86}
87
88/// Link function for [`SmoothFamily::Gamma`].
89#[cfg(feature = "regression")]
90#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
91pub enum SmoothGammaLink {
92    /// `1 / μ` (R's canonical Gamma link).
93    #[default]
94    Inverse,
95    /// `log(μ)`.
96    Log,
97}
98
99/// Smoothing statistic — supports both linear regression and LOESS.
100pub struct StatSmooth {
101    /// Number of points to generate for the fitted line.
102    pub n_points: usize,
103    /// Whether to compute confidence interval.
104    pub se: bool,
105    /// Smoothing method.
106    pub method: SmoothMethod,
107}
108
109impl Default for StatSmooth {
110    fn default() -> Self {
111        StatSmooth {
112            n_points: 80,
113            se: true,
114            method: SmoothMethod::Lm,
115        }
116    }
117}
118
119impl Stat for StatSmooth {
120    fn compute_group(&self, data: &DataFrame, scales: &ScaleSet) -> DataFrame {
121        match &self.method {
122            SmoothMethod::Lm => self.compute_lm(data),
123            SmoothMethod::Loess { span } => {
124                let loess = super::loess::StatLoess {
125                    span: *span,
126                    n_points: self.n_points,
127                    se: self.se,
128                };
129                loess.compute_group(data, scales)
130            }
131            #[cfg(feature = "regression")]
132            SmoothMethod::Glm { family } => self.compute_glm(data, Some(*family)),
133            #[cfg(feature = "regression")]
134            SmoothMethod::Rlm => self.compute_glm(data, None),
135            SmoothMethod::Gam => self.compute_gam(data),
136        }
137    }
138
139    fn required_aes(&self) -> Vec<Aesthetic> {
140        vec![Aesthetic::X, Aesthetic::Y]
141    }
142
143    fn name(&self) -> &str {
144        "smooth"
145    }
146}
147
148impl StatSmooth {
149    fn compute_lm(&self, data: &DataFrame) -> DataFrame {
150        let x_col = match data.column("x") {
151            Some(c) => c,
152            None => return DataFrame::new(),
153        };
154        let y_col = match data.column("y") {
155            Some(c) => c,
156            None => return DataFrame::new(),
157        };
158
159        let pairs: Vec<(f64, f64)> = x_col
160            .iter()
161            .zip(y_col.iter())
162            .filter_map(|(x, y)| Some((x.as_f64()?, y.as_f64()?)))
163            .collect();
164
165        if pairs.len() < 2 {
166            return DataFrame::new();
167        }
168
169        let n = pairs.len() as f64;
170        let sum_x: f64 = pairs.iter().map(|(x, _)| x).sum();
171        let sum_y: f64 = pairs.iter().map(|(_, y)| y).sum();
172        let sum_xy: f64 = pairs.iter().map(|(x, y)| x * y).sum();
173        let sum_xx: f64 = pairs.iter().map(|(x, _)| x * x).sum();
174
175        let mean_x = sum_x / n;
176        let mean_y = sum_y / n;
177
178        let denom = sum_xx - sum_x * sum_x / n;
179        let (slope, intercept) = if denom.abs() < f64::EPSILON {
180            (0.0, mean_y)
181        } else {
182            let m = (sum_xy - sum_x * sum_y / n) / denom;
183            let b = mean_y - m * mean_x;
184            (m, b)
185        };
186
187        // Generate fitted values across x range
188        let x_min = pairs.iter().map(|(x, _)| *x).fold(f64::INFINITY, f64::min);
189        let x_max = pairs
190            .iter()
191            .map(|(x, _)| *x)
192            .fold(f64::NEG_INFINITY, f64::max);
193
194        let step = (x_max - x_min) / (self.n_points - 1).max(1) as f64;
195
196        // Compute standard error of prediction if requested
197        let se_values = if self.se && pairs.len() > 2 {
198            let residuals: Vec<f64> = pairs
199                .iter()
200                .map(|(x, y)| y - (slope * x + intercept))
201                .collect();
202            let sse: f64 = residuals.iter().map(|r| r * r).sum();
203            let mse = sse / (n - 2.0);
204            Some((mse, sum_xx, mean_x, n))
205        } else {
206            None
207        };
208
209        let mut x_vals = Vec::with_capacity(self.n_points);
210        let mut y_vals = Vec::with_capacity(self.n_points);
211        let mut ymin_vals = Vec::with_capacity(self.n_points);
212        let mut ymax_vals = Vec::with_capacity(self.n_points);
213
214        for i in 0..self.n_points {
215            let x = x_min + i as f64 * step;
216            let y = slope * x + intercept;
217            x_vals.push(Value::Float(x));
218            y_vals.push(Value::Float(y));
219
220            if let Some((mse, sum_xx, mean_x, n)) = se_values {
221                let se_pred = (mse
222                    * (1.0 / n + (x - mean_x).powi(2) / (sum_xx - n * mean_x * mean_x)))
223                    .sqrt();
224                // 95% confidence interval for the mean response: R's
225                // qt(0.975, n − 2), not the large-sample normal 1.96.
226                let t_val = crate::stat::dist::qt(0.975, (n - 2.0).max(1.0));
227                ymin_vals.push(Value::Float(y - t_val * se_pred));
228                ymax_vals.push(Value::Float(y + t_val * se_pred));
229            }
230        }
231
232        let mut result = DataFrame::new();
233        result.add_column("x".to_string(), x_vals);
234        result.add_column("y".to_string(), y_vals);
235        if !ymin_vals.is_empty() {
236            result.add_column("ymin".to_string(), ymin_vals);
237            result.add_column("ymax".to_string(), ymax_vals);
238        }
239        result
240    }
241
242    /// GLM / robust-linear smoothing backed by anofox-regression. `family = None`
243    /// selects the robust (Huber) fit; `Some(..)` selects a GLM family. A
244    /// confidence-interval ribbon (ymin/ymax) is emitted when `self.se` is set.
245    #[cfg(feature = "regression")]
246    fn compute_glm(&self, data: &DataFrame, family: Option<SmoothFamily>) -> DataFrame {
247        use anofox_regression::solvers::{
248            FittedRegressor, HuberRegressor, OlsRegressor, Regressor,
249        };
250        use anofox_regression::{IntervalType, RegressionOptions};
251        use faer::{Col, Mat};
252
253        let (x_col, y_col) = match (data.column("x"), data.column("y")) {
254            (Some(x), Some(y)) => (x, y),
255            _ => return DataFrame::new(),
256        };
257        let pairs: Vec<(f64, f64)> = x_col
258            .iter()
259            .zip(y_col.iter())
260            .filter_map(|(x, y)| Some((x.as_f64()?, y.as_f64()?)))
261            .collect();
262        if pairs.len() < 2 {
263            return DataFrame::new();
264        }
265
266        let n = pairs.len();
267        let x = Mat::from_fn(n, 1, |i, _| pairs[i].0);
268        let y = Col::from_fn(n, |i| pairs[i].1);
269        let x_min = pairs.iter().map(|p| p.0).fold(f64::INFINITY, f64::min);
270        let x_max = pairs.iter().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max);
271        let steps = self.n_points.max(2);
272        let grid = Mat::from_fn(steps, 1, |k, _| {
273            x_min + (x_max - x_min) * k as f64 / (steps - 1) as f64
274        });
275        let interval = if self.se {
276            Some(IntervalType::Confidence)
277        } else {
278            None
279        };
280
281        // Fit the requested model and predict (with interval) over the grid.
282        let pred = match family {
283            None => match HuberRegressor::new().fit(&x, &y) {
284                Ok(f) => f.predict_with_interval(&grid, interval, 0.95),
285                Err(_) => return DataFrame::new(),
286            },
287            Some(SmoothFamily::Gaussian) => {
288                match OlsRegressor::new(RegressionOptions::default()).fit(&x, &y) {
289                    Ok(f) => f.predict_with_interval(&grid, interval, 0.95),
290                    Err(_) => return DataFrame::new(),
291                }
292            }
293            Some(family) => match glm_link_prediction(family, &x, &y, &grid, self.se) {
294                Some(p) => p,
295                None => return DataFrame::new(),
296            },
297        };
298
299        let mut x_vals = Vec::with_capacity(steps);
300        let mut y_vals = Vec::with_capacity(steps);
301        let mut ymin_vals = Vec::with_capacity(steps);
302        let mut ymax_vals = Vec::with_capacity(steps);
303        for k in 0..steps {
304            x_vals.push(Value::Float(grid[(k, 0)]));
305            y_vals.push(Value::Float(pred.fit[k]));
306            if self.se {
307                // An inverse link can swap the ends (e.g. Gamma's 1/η).
308                let (a, b) = (pred.lower[k], pred.upper[k]);
309                ymin_vals.push(Value::Float(a.min(b)));
310                ymax_vals.push(Value::Float(a.max(b)));
311            }
312        }
313
314        let mut result = DataFrame::new();
315        result.add_column("x".to_string(), x_vals);
316        result.add_column("y".to_string(), y_vals);
317        if self.se {
318            result.add_column("ymin".to_string(), ymin_vals);
319            result.add_column("ymax".to_string(), ymax_vals);
320        }
321        for col_name in &["color", "fill", "group"] {
322            if let Some(col) = data.column(col_name) {
323                if let Some(first) = col.first() {
324                    result.add_column(col_name.to_string(), vec![first.clone(); steps]);
325                }
326            }
327        }
328        result
329    }
330
331    /// `method = "gam"` dispatcher: anofox-regression's P-spline when the
332    /// `regression` feature is enabled, the dependency-free (wasm-safe)
333    /// P-spline otherwise. Both always compile when `regression` is on
334    /// (see [`Self::compute_gam_builtin`]'s doc comment), which is what lets
335    /// `tests/gam_smooth.rs` and the parity test below compare them.
336    fn compute_gam(&self, data: &DataFrame) -> DataFrame {
337        #[cfg(feature = "regression")]
338        return self.compute_gam_regression(data);
339        #[cfg(not(feature = "regression"))]
340        return self.compute_gam_builtin(data);
341    }
342
343    /// GAM smoothing via anofox-regression's penalized B-spline (P-spline) with
344    /// GCV-selected smoothing — ggplot2's `method = "gam"`. Emits a t-based
345    /// confidence ribbon (ymin/ymax) when `self.se` is set.
346    #[cfg(feature = "regression")]
347    fn compute_gam_regression(&self, data: &DataFrame) -> DataFrame {
348        use anofox_regression::solvers::{FittedRegressor, PSplineRegressor, Regressor};
349        use anofox_regression::IntervalType;
350        use faer::{Col, Mat};
351
352        let (x_col, y_col) = match (data.column("x"), data.column("y")) {
353            (Some(x), Some(y)) => (x, y),
354            _ => return DataFrame::new(),
355        };
356        let pairs: Vec<(f64, f64)> = x_col
357            .iter()
358            .zip(y_col.iter())
359            .filter_map(|(x, y)| Some((x.as_f64()?, y.as_f64()?)))
360            .collect();
361        // The P-spline needs enough distinct support to build a basis; fall back
362        // to a straight line for tiny groups.
363        if pairs.len() < 6 {
364            return self.compute_lm(data);
365        }
366
367        let n = pairs.len();
368        let x = Mat::from_fn(n, 1, |i, _| pairs[i].0);
369        let y = Col::from_fn(n, |i| pairs[i].1);
370        let x_min = pairs.iter().map(|p| p.0).fold(f64::INFINITY, f64::min);
371        let x_max = pairs.iter().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max);
372        let steps = self.n_points.max(2);
373        let grid = Mat::from_fn(steps, 1, |k, _| {
374            x_min + (x_max - x_min) * k as f64 / (steps - 1) as f64
375        });
376        let interval = if self.se {
377            Some(IntervalType::Confidence)
378        } else {
379            None
380        };
381
382        let pred = match PSplineRegressor::new().fit(&x, &y) {
383            Ok(f) => f.predict_with_interval(&grid, interval, 0.95),
384            // Degenerate input (e.g. collinear x) — fall back to a line.
385            Err(_) => return self.compute_lm(data),
386        };
387
388        let mut x_vals = Vec::with_capacity(steps);
389        let mut y_vals = Vec::with_capacity(steps);
390        let mut ymin_vals = Vec::with_capacity(steps);
391        let mut ymax_vals = Vec::with_capacity(steps);
392        for k in 0..steps {
393            x_vals.push(Value::Float(grid[(k, 0)]));
394            y_vals.push(Value::Float(pred.fit[k]));
395            if self.se {
396                ymin_vals.push(Value::Float(pred.lower[k]));
397                ymax_vals.push(Value::Float(pred.upper[k]));
398            }
399        }
400
401        let mut result = DataFrame::new();
402        result.add_column("x".to_string(), x_vals);
403        result.add_column("y".to_string(), y_vals);
404        if self.se {
405            result.add_column("ymin".to_string(), ymin_vals);
406            result.add_column("ymax".to_string(), ymax_vals);
407        }
408        for col_name in &["color", "fill", "group"] {
409            if let Some(col) = data.column(col_name) {
410                if let Some(first) = col.first() {
411                    result.add_column(col_name.to_string(), vec![first.clone(); steps]);
412                }
413            }
414        }
415        result
416    }
417
418    /// GAM smoothing via a small, dependency-free penalized B-spline
419    /// (P-spline) fit — ggplot2's `method = "gam"` without the `regression`
420    /// feature's anofox-regression/faer/nalgebra/statrs/rand dependency tree,
421    /// so this builds for `wasm32-unknown-unknown` too (#174). Same cubic
422    /// B-spline basis, GCV-selected λ and confidence-band construction as
423    /// [`gam::fit`](super::gam::fit) documents. Emits a t-based confidence
424    /// ribbon (ymin/ymax) when `self.se` is set.
425    ///
426    /// This function has no `#[cfg]` of its own — it always compiles — so
427    /// that with the `regression` feature on, both backends exist side by
428    /// side and can be compared directly (see the parity test below and
429    /// `tests/gam_smooth.rs`); only [`Self::compute_gam`]'s dispatch is
430    /// feature-gated. With `regression` on, its only non-test caller is
431    /// gone (the dispatcher picks `compute_gam_regression` instead), hence
432    /// the `allow` below for that configuration.
433    #[cfg_attr(feature = "regression", allow(dead_code))]
434    fn compute_gam_builtin(&self, data: &DataFrame) -> DataFrame {
435        let (x_col, y_col) = match (data.column("x"), data.column("y")) {
436            (Some(x), Some(y)) => (x, y),
437            _ => return DataFrame::new(),
438        };
439        let pairs: Vec<(f64, f64)> = x_col
440            .iter()
441            .zip(y_col.iter())
442            .filter_map(|(x, y)| Some((x.as_f64()?, y.as_f64()?)))
443            .collect();
444        // The P-spline needs enough distinct support to build a basis; fall
445        // back to a straight line for tiny groups or a singular/degenerate
446        // fit — matching the regression backend's fallback behaviour.
447        if pairs.len() < 6 {
448            return self.compute_lm(data);
449        }
450        let Some(fit) = super::gam::fit(&pairs) else {
451            return self.compute_lm(data);
452        };
453
454        let (x_min, x_max) = fit.domain();
455        let steps = self.n_points.max(2);
456        let t_val = self
457            .se
458            .then(|| crate::stat::dist::qt(0.975, (fit.n() - fit.edf()).max(1.0)));
459
460        let mut x_vals = Vec::with_capacity(steps);
461        let mut y_vals = Vec::with_capacity(steps);
462        let mut ymin_vals = Vec::with_capacity(steps);
463        let mut ymax_vals = Vec::with_capacity(steps);
464        for k in 0..steps {
465            let x = x_min + (x_max - x_min) * k as f64 / (steps - 1) as f64;
466            let y = fit.predict(x);
467            x_vals.push(Value::Float(x));
468            y_vals.push(Value::Float(y));
469            if let Some(t) = t_val {
470                let se = fit.se(x);
471                ymin_vals.push(Value::Float(y - t * se));
472                ymax_vals.push(Value::Float(y + t * se));
473            }
474        }
475
476        let mut result = DataFrame::new();
477        result.add_column("x".to_string(), x_vals);
478        result.add_column("y".to_string(), y_vals);
479        if self.se {
480            result.add_column("ymin".to_string(), ymin_vals);
481            result.add_column("ymax".to_string(), ymax_vals);
482        }
483        for col_name in &["color", "fill", "group"] {
484            if let Some(col) = data.column(col_name) {
485                if let Some(first) = col.first() {
486                    result.add_column(col_name.to_string(), vec![first.clone(); steps]);
487                }
488            }
489        }
490        result
491    }
492}
493
494/// IRLS convergence settings shared with the anofox-statistics DuckDB
495/// extension's GLM aggregates (R's `glm.control(epsilon = 1e-8)`), so a fit in
496/// SQL and the plotted smooth agree numerically.
497#[cfg(feature = "regression")]
498const GLM_TOL: f64 = 1e-8;
499#[cfg(feature = "regression")]
500const GLM_MAX_ITER: usize = 100;
501
502/// Fitted mean and (optionally) a 95% confidence band for a non-Gaussian GLM
503/// family over `grid`, computed like ggplot2's `predictdf.glm`: the band is
504/// `linkinv(η ± qnorm(0.975)·se(η))`. `None` when the fit fails (degenerate
505/// or out-of-range responses, e.g. a negative count).
506#[cfg(feature = "regression")]
507fn glm_link_prediction(
508    family: SmoothFamily,
509    x: &faer::Mat<f64>,
510    y: &faer::Col<f64>,
511    grid: &faer::Mat<f64>,
512    se: bool,
513) -> Option<anofox_regression::PredictionResult> {
514    use anofox_regression::core::PredictionType;
515    use anofox_regression::solvers::{
516        BinomialRegressor, GammaRegressor, NegativeBinomialRegressor, PoissonRegressor, Regressor,
517    };
518    use anofox_regression::{BinomialLink, PredictionResult};
519    use faer::Col;
520
521    // Link-scale prediction (η and se(η)) plus the inverse link.
522    let (link_pred, linkinv): (PredictionResult, Box<dyn Fn(f64) -> f64>) = match family {
523        SmoothFamily::Gaussian => return None,
524        SmoothFamily::Poisson => {
525            let f = PoissonRegressor::log()
526                .tolerance(GLM_TOL)
527                .max_iterations(GLM_MAX_ITER)
528                .build()
529                .fit(x, y)
530                .ok()?;
531            (
532                f.predict_with_se(grid, PredictionType::Link, None, 0.95),
533                Box::new(f64::exp),
534            )
535        }
536        SmoothFamily::NegativeBinomial => {
537            let f = NegativeBinomialRegressor::builder()
538                .tolerance(GLM_TOL)
539                .max_iterations(GLM_MAX_ITER)
540                .build()
541                .fit(x, y)
542                .ok()?;
543            (
544                f.predict_with_se(grid, PredictionType::Link, None, 0.95),
545                Box::new(f64::exp),
546            )
547        }
548        SmoothFamily::Binomial(link) => {
549            let link = match link {
550                SmoothBinomialLink::Logit => BinomialLink::Logit,
551                SmoothBinomialLink::Probit => BinomialLink::Probit,
552                SmoothBinomialLink::Cloglog => BinomialLink::Cloglog,
553            };
554            let f = BinomialRegressor::builder()
555                .link(link)
556                .tolerance(GLM_TOL)
557                .max_iterations(GLM_MAX_ITER)
558                .build()
559                .fit(x, y)
560                .ok()?;
561            (
562                f.predict_with_se(grid, PredictionType::Link, None, 0.95),
563                Box::new(move |eta| link.link_inverse(eta)),
564            )
565        }
566        SmoothFamily::Gamma(link) => {
567            let (power, inv): (f64, Box<dyn Fn(f64) -> f64>) = match link {
568                SmoothGammaLink::Inverse => (-1.0, Box::new(|eta: f64| 1.0 / eta)),
569                SmoothGammaLink::Log => (0.0, Box::new(f64::exp)),
570            };
571            let f = GammaRegressor::builder()
572                .link_power(power)
573                .tolerance(GLM_TOL)
574                .max_iterations(GLM_MAX_ITER)
575                .build()
576                .fit(x, y)
577                .ok()?;
578            (
579                f.inner()
580                    .predict_with_se(grid, PredictionType::Link, None, 0.95),
581                inv,
582            )
583        }
584    };
585
586    let n = grid.nrows();
587    let eta = &link_pred.fit;
588    let fit = Col::from_fn(n, |k| linkinv(eta[k]));
589    if !se {
590        return Some(PredictionResult::with_intervals(
591            fit,
592            Col::zeros(n),
593            Col::zeros(n),
594            Col::zeros(n),
595        ));
596    }
597    let z = crate::stat::dist::qnorm(0.975);
598    let lower = Col::from_fn(n, |k| linkinv(eta[k] - z * link_pred.se[k]));
599    let upper = Col::from_fn(n, |k| linkinv(eta[k] + z * link_pred.se[k]));
600    Some(PredictionResult::with_intervals(
601        fit,
602        lower,
603        upper,
604        link_pred.se.clone(),
605    ))
606}
607
608// With the `regression` feature on, both `compute_gam_regression` and
609// `compute_gam_builtin` compile side by side (only `compute_gam`'s dispatch
610// is feature-gated) — exercise that to check the two backends agree, which
611// is only possible in this configuration.
612#[cfg(all(test, feature = "regression"))]
613mod gam_backend_parity_tests {
614    use super::*;
615
616    /// A noisy sine wave — the same shape of fixture used across the other
617    /// GAM tests (`tests/gam_smooth.rs`, `src/stat/gam.rs`).
618    fn sine_frame() -> DataFrame {
619        let mut df = DataFrame::new();
620        let x: Vec<Value> = (0..80)
621            .map(|i| Value::Float(i as f64 / 79.0 * std::f64::consts::TAU))
622            .collect();
623        let y: Vec<Value> = (0..80)
624            .map(|i| {
625                let xv = i as f64 / 79.0 * std::f64::consts::TAU;
626                Value::Float(xv.sin() + 0.05 * ((i * 7 % 5) as f64 - 2.0))
627            })
628            .collect();
629        df.add_column("x".into(), x);
630        df.add_column("y".into(), y);
631        df
632    }
633
634    #[test]
635    fn builtin_matches_regression_backend() {
636        let data = sine_frame();
637        let stat = StatSmooth {
638            n_points: 50,
639            se: false,
640            method: SmoothMethod::Gam,
641        };
642        let reg = stat.compute_gam_regression(&data);
643        let builtin = stat.compute_gam_builtin(&data);
644        assert_eq!(reg.nrows(), builtin.nrows());
645
646        let reg_x: Vec<f64> = reg
647            .column("x")
648            .unwrap()
649            .iter()
650            .filter_map(Value::as_f64)
651            .collect();
652        let reg_y: Vec<f64> = reg
653            .column("y")
654            .unwrap()
655            .iter()
656            .filter_map(Value::as_f64)
657            .collect();
658        let builtin_x: Vec<f64> = builtin
659            .column("x")
660            .unwrap()
661            .iter()
662            .filter_map(Value::as_f64)
663            .collect();
664        let builtin_y: Vec<f64> = builtin
665            .column("y")
666            .unwrap()
667            .iter()
668            .filter_map(Value::as_f64)
669            .collect();
670
671        // Both backends grid the same [x_min, x_max] with the same formula,
672        // so the x values should match closely.
673        let max_x_diff = reg_x
674            .iter()
675            .zip(&builtin_x)
676            .map(|(a, b)| (a - b).abs())
677            .fold(0.0, f64::max);
678        assert!(max_x_diff < 1e-9, "x grids disagree: {max_x_diff}");
679
680        // The two P-spline fits (anofox-regression's PSplineRegressor vs. the
681        // dependency-free GCV fit here) should track the same curve closely —
682        // this is the issue #174 "Done when" criterion.
683        let max_y_diff = reg_y
684            .iter()
685            .zip(&builtin_y)
686            .map(|(a, b)| (a - b).abs())
687            .fold(0.0, f64::max);
688        // Measured max_y_diff on this fixture is ~0.093; the bound below is
689        // about 2x that — tight enough to catch a real backend divergence,
690        // loose enough for GCV's lambda grid to land on a neighbouring step
691        // between the two independent implementations.
692        assert!(
693            max_y_diff < 0.19,
694            "fitted curves disagree by up to {max_y_diff}"
695        );
696    }
697
698    #[test]
699    fn builtin_falls_back_to_lm_for_tiny_groups_like_regression() {
700        let mut df = DataFrame::new();
701        df.add_column("x".into(), (0..4).map(|i| Value::Float(i as f64)).collect());
702        df.add_column(
703            "y".into(),
704            (0..4).map(|i| Value::Float(2.0 * i as f64 + 1.0)).collect(),
705        );
706        let stat = StatSmooth {
707            n_points: 20,
708            se: false,
709            method: SmoothMethod::Gam,
710        };
711        let reg = stat.compute_gam_regression(&df);
712        let builtin = stat.compute_gam_builtin(&df);
713        let lm = stat.compute_lm(&df);
714        assert_eq!(reg.nrows(), lm.nrows());
715        assert_eq!(builtin.nrows(), lm.nrows());
716    }
717}