ggplot_rs/stat/
function.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// Evaluate a function over the x range and produce (x, y) pairs.
8pub struct StatFunction {
9    pub fun: Box<dyn Fn(f64) -> f64 + Send + Sync>,
10    pub n_points: usize,
11    /// Evaluate over this x range instead of the data's (ggplot2's `xlim`).
12    pub xlim: Option<(f64, f64)>,
13}
14
15impl StatFunction {
16    pub fn new(fun: impl Fn(f64) -> f64 + Send + Sync + 'static) -> Self {
17        StatFunction {
18            fun: Box::new(fun),
19            n_points: 101,
20            xlim: None,
21        }
22    }
23
24    /// Evaluate over `[lo, hi]` regardless of the layer's data.
25    pub fn with_xlim(mut self, lo: f64, hi: f64) -> Self {
26        self.xlim = Some((lo, hi));
27        self
28    }
29
30    pub fn with_n_points(mut self, n: usize) -> Self {
31        self.n_points = n;
32        self
33    }
34}
35
36impl Stat for StatFunction {
37    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
38        // Determine x range from data
39        let x_col = match data.column("x") {
40            Some(c) => c,
41            None => return DataFrame::new(),
42        };
43
44        let values: Vec<f64> = x_col.iter().filter_map(|v| v.as_f64()).collect();
45        let (x_min, x_max) = match self.xlim {
46            Some(lim) => lim,
47            None if values.is_empty() => return DataFrame::new(),
48            None => (
49                values.iter().cloned().fold(f64::INFINITY, f64::min),
50                values.iter().cloned().fold(f64::NEG_INFINITY, f64::max),
51            ),
52        };
53
54        if (x_max - x_min).abs() < f64::EPSILON {
55            return DataFrame::new();
56        }
57
58        let step = (x_max - x_min) / (self.n_points - 1).max(1) as f64;
59
60        let mut x_vals = Vec::with_capacity(self.n_points);
61        let mut y_vals = Vec::with_capacity(self.n_points);
62
63        for i in 0..self.n_points {
64            let x = x_min + i as f64 * step;
65            let y = (self.fun)(x);
66            x_vals.push(Value::Float(x));
67            y_vals.push(Value::Float(y));
68        }
69
70        let mut result = DataFrame::new();
71        result.add_column("x".to_string(), x_vals);
72        result.add_column("y".to_string(), y_vals);
73        result
74    }
75
76    fn required_aes(&self) -> Vec<Aesthetic> {
77        vec![Aesthetic::X]
78    }
79
80    fn name(&self) -> &str {
81        "function"
82    }
83}
ggplot_rs/stat/function.rs

ggplot_rs/stat/
function.rs
