1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7#[derive(Clone, Debug, Default)]
9#[non_exhaustive]
10pub enum SmoothMethod {
11 #[default]
13 Lm,
14 Loess { span: f64 },
16 #[cfg(feature = "regression")]
18 Glm { family: SmoothFamily },
19 #[cfg(feature = "regression")]
21 Rlm,
22 Gam,
28}
29
30#[cfg(feature = "regression")]
45#[derive(Clone, Copy, Debug, Default, PartialEq)]
46pub enum SmoothFamily {
47 #[default]
49 Gaussian,
50 Poisson,
52 Binomial(SmoothBinomialLink),
55 Gamma(SmoothGammaLink),
58 NegativeBinomial,
61}
62
63#[cfg(feature = "regression")]
64impl SmoothFamily {
65 pub fn binomial() -> Self {
67 SmoothFamily::Binomial(SmoothBinomialLink::Logit)
68 }
69
70 pub fn gamma() -> Self {
72 SmoothFamily::Gamma(SmoothGammaLink::Inverse)
73 }
74}
75
76#[cfg(feature = "regression")]
78#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
79pub enum SmoothBinomialLink {
80 #[default]
82 Logit,
83 Probit,
85 Cloglog,
87}
88
89#[cfg(feature = "regression")]
91#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
92pub enum SmoothGammaLink {
93 #[default]
95 Inverse,
96 Log,
98}
99
100pub struct StatSmooth {
102 pub n_points: usize,
104 pub se: bool,
106 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 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 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 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 #[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 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 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 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 #[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 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 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 #[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 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#[cfg(feature = "regression")]
499const GLM_TOL: f64 = 1e-8;
500#[cfg(feature = "regression")]
501const GLM_MAX_ITER: usize = 100;
502
503#[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 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#[cfg(all(test, feature = "regression"))]
614mod gam_backend_parity_tests {
615 use super::*;
616
617 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 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 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 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}