Skip to main content

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}
12
13impl StatFunction {
14    pub fn new(fun: impl Fn(f64) -> f64 + Send + Sync + 'static) -> Self {
15        StatFunction {
16            fun: Box::new(fun),
17            n_points: 101,
18        }
19    }
20
21    pub fn with_n_points(mut self, n: usize) -> Self {
22        self.n_points = n;
23        self
24    }
25}
26
27impl Stat for StatFunction {
28    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
29        // Determine x range from data
30        let x_col = match data.column("x") {
31            Some(c) => c,
32            None => return DataFrame::new(),
33        };
34
35        let values: Vec<f64> = x_col.iter().filter_map(|v| v.as_f64()).collect();
36        if values.is_empty() {
37            return DataFrame::new();
38        }
39
40        let x_min = values.iter().cloned().fold(f64::INFINITY, f64::min);
41        let x_max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
42
43        if (x_max - x_min).abs() < f64::EPSILON {
44            return DataFrame::new();
45        }
46
47        let step = (x_max - x_min) / (self.n_points - 1).max(1) as f64;
48
49        let mut x_vals = Vec::with_capacity(self.n_points);
50        let mut y_vals = Vec::with_capacity(self.n_points);
51
52        for i in 0..self.n_points {
53            let x = x_min + i as f64 * step;
54            let y = (self.fun)(x);
55            x_vals.push(Value::Float(x));
56            y_vals.push(Value::Float(y));
57        }
58
59        let mut result = DataFrame::new();
60        result.add_column("x".to_string(), x_vals);
61        result.add_column("y".to_string(), y_vals);
62        result
63    }
64
65    fn required_aes(&self) -> Vec<Aesthetic> {
66        vec![Aesthetic::X]
67    }
68
69    fn name(&self) -> &str {
70        "function"
71    }
72}