Skip to main content

ggplot_rs/stat/
ecdf.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// Empirical cumulative distribution function.
8/// Sorts x values and assigns y = rank / n.
9pub struct StatEcdf;
10
11impl Default for StatEcdf {
12    fn default() -> Self {
13        StatEcdf
14    }
15}
16
17impl Stat for StatEcdf {
18    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
19        let x_col = match data.column("x") {
20            Some(c) => c,
21            None => return DataFrame::new(),
22        };
23
24        let mut values: Vec<f64> = x_col
25            .iter()
26            .filter_map(|v| v.as_f64())
27            .filter(|v| v.is_finite())
28            .collect();
29        if values.is_empty() {
30            return DataFrame::new();
31        }
32
33        values.sort_by(|a, b| a.total_cmp(b));
34        let n = values.len() as f64;
35
36        let mut x_vals = Vec::with_capacity(values.len() + 2);
37        let mut y_vals = Vec::with_capacity(values.len() + 2);
38
39        // ggplot2 pads the step to ±Inf (y = 0 before the first point, y = 1
40        // after the last) so it spans the panel. Scales ignore non-finite values
41        // when training, and geom_step clamps the ±Inf segments to the panel edge.
42        x_vals.push(Value::Float(f64::NEG_INFINITY));
43        y_vals.push(Value::Float(0.0));
44        for (i, &x) in values.iter().enumerate() {
45            x_vals.push(Value::Float(x));
46            y_vals.push(Value::Float((i + 1) as f64 / n));
47        }
48        x_vals.push(Value::Float(f64::INFINITY));
49        y_vals.push(Value::Float(1.0));
50
51        let mut result = DataFrame::new();
52        result.add_column("x".to_string(), x_vals);
53        result.add_column("y".to_string(), y_vals);
54
55        // Carry over grouping columns
56        let nrows = values.len() + 2;
57        for col_name in &["color", "fill", "group"] {
58            if let Some(col) = data.column(col_name) {
59                if let Some(first) = col.first() {
60                    result.add_column(col_name.to_string(), vec![first.clone(); nrows]);
61                }
62            }
63        }
64
65        result
66    }
67
68    fn required_aes(&self) -> Vec<Aesthetic> {
69        vec![Aesthetic::X]
70    }
71
72    fn name(&self) -> &str {
73        "ecdf"
74    }
75}
76
77/// ECDF with a simultaneous Dvoretzky–Kiefer–Wolfowitz confidence band:
78/// `F̂(x) ± ε`, `ε = √(ln(2 / (1 − level)) / (2n))`, clamped to `[0, 1]`.
79/// Output `x` (±Inf padded, as [`StatEcdf`]), `y`, `ymin`, `ymax` — draw it
80/// with `geom_stepribbon` (see `GGPlot::stat_ecdf_band`).
81#[derive(Clone, Debug)]
82pub struct StatEcdfBand {
83    /// Confidence level (default 0.95).
84    pub level: f64,
85}
86
87impl Default for StatEcdfBand {
88    fn default() -> Self {
89        StatEcdfBand { level: 0.95 }
90    }
91}
92
93impl StatEcdfBand {
94    pub fn new(level: f64) -> Self {
95        StatEcdfBand { level }
96    }
97
98    /// The DKW half-width `ε` for `n` observations at this level.
99    pub fn epsilon(&self, n: usize) -> f64 {
100        ((2.0 / (1.0 - self.level)).ln() / (2.0 * n as f64)).sqrt()
101    }
102}
103
104impl Stat for StatEcdfBand {
105    fn compute_group(&self, data: &DataFrame, scales: &ScaleSet) -> DataFrame {
106        if !(self.level > 0.0 && self.level < 1.0) {
107            return DataFrame::new();
108        }
109        // Only finite x count towards n (±Inf / NaN input rows are dropped).
110        let finite: Vec<Value> = data
111            .column("x")
112            .map(|c| {
113                c.iter()
114                    .filter(|v| v.as_f64().is_some_and(f64::is_finite))
115                    .cloned()
116                    .collect()
117            })
118            .unwrap_or_default();
119        let n = finite.len();
120        if n == 0 {
121            return DataFrame::new();
122        }
123        let mut input = DataFrame::new();
124        input.add_column("x".into(), finite);
125        for col in ["color", "fill", "group"] {
126            if let Some(first) = data.column(col).and_then(|c| c.first()) {
127                input.add_column(col.into(), vec![first.clone(); n]);
128            }
129        }
130        let mut out = StatEcdf.compute_group(&input, scales);
131        let eps = self.epsilon(n);
132        let y: Vec<f64> = out
133            .column("y")
134            .map(|c| c.iter().filter_map(|v| v.as_f64()).collect())
135            .unwrap_or_default();
136        out.add_column(
137            "ymin".into(),
138            y.iter()
139                .map(|&v| Value::Float((v - eps).max(0.0)))
140                .collect(),
141        );
142        out.add_column(
143            "ymax".into(),
144            y.iter()
145                .map(|&v| Value::Float((v + eps).min(1.0)))
146                .collect(),
147        );
148        out
149    }
150
151    fn required_aes(&self) -> Vec<Aesthetic> {
152        vec![Aesthetic::X]
153    }
154
155    fn name(&self) -> &str {
156        "ecdf_band"
157    }
158}