1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7#[derive(Clone, Debug, Default)]
9pub enum SmoothMethod {
10 #[default]
12 Lm,
13 Loess { span: f64 },
15 #[cfg(feature = "regression")]
17 Glm { family: SmoothFamily },
18 #[cfg(feature = "regression")]
20 Rlm,
21 Gam,
27}
28
29#[cfg(feature = "regression")]
44#[derive(Clone, Copy, Debug, Default, PartialEq)]
45pub enum SmoothFamily {
46 #[default]
48 Gaussian,
49 Poisson,
51 Binomial(SmoothBinomialLink),
54 Gamma(SmoothGammaLink),
57 NegativeBinomial,
60}
61
62#[cfg(feature = "regression")]
63impl SmoothFamily {
64 pub fn binomial() -> Self {
66 SmoothFamily::Binomial(SmoothBinomialLink::Logit)
67 }
68
69 pub fn gamma() -> Self {
71 SmoothFamily::Gamma(SmoothGammaLink::Inverse)
72 }
73}
74
75#[cfg(feature = "regression")]
77#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
78pub enum SmoothBinomialLink {
79 #[default]
81 Logit,
82 Probit,
84 Cloglog,
86}
87
88#[cfg(feature = "regression")]
90#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
91pub enum SmoothGammaLink {
92 #[default]
94 Inverse,
95 Log,
97}
98
99pub struct StatSmooth {
101 pub n_points: usize,
103 pub se: bool,
105 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 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 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 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 #[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 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 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 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 #[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 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 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 #[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 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#[cfg(feature = "regression")]
498const GLM_TOL: f64 = 1e-8;
499#[cfg(feature = "regression")]
500const GLM_MAX_ITER: usize = 100;
501
502#[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 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#[cfg(all(test, feature = "regression"))]
613mod gam_backend_parity_tests {
614 use super::*;
615
616 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 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 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 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}