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