Skip to main content

ggplot_rs/geom/
density.rs

1use crate::aes::Aesthetic;
2use crate::coord::Coord;
3use crate::data::{DataFrame, Value};
4use crate::position::identity::PositionIdentity;
5use crate::position::Position;
6use crate::render::backend::{DrawBackend, LineStyle, Linetype, RectStyle};
7use crate::render::RenderError;
8use crate::scale::ScaleSet;
9use crate::stat::density::StatDensity;
10use crate::stat::Stat;
11use crate::theme::Theme;
12
13use super::{Geom, GeomParams};
14
15/// Density curve geometry — filled area under a kernel density estimate.
16pub struct GeomDensity {
17    pub fill: (u8, u8, u8),
18    pub color: (u8, u8, u8),
19    pub alpha: f64,
20    pub line_width: f64,
21}
22
23impl Default for GeomDensity {
24    fn default() -> Self {
25        GeomDensity {
26            fill: (97, 156, 255),
27            color: (50, 50, 50),
28            alpha: 0.3,
29            line_width: 1.0,
30        }
31    }
32}
33
34impl Geom for GeomDensity {
35    fn draw(
36        &self,
37        data: &DataFrame,
38        coord: &dyn Coord,
39        scales: &ScaleSet,
40        _theme: &Theme,
41        backend: &mut dyn DrawBackend,
42    ) -> Result<(), RenderError> {
43        let x_col = data
44            .column("x")
45            .ok_or(RenderError::MissingAesthetic("x".into()))?;
46        let y_col = data
47            .column("y")
48            .ok_or(RenderError::MissingAesthetic("y".into()))?;
49        let fill_col = data.column("fill");
50        let color_col = data.column("color");
51
52        let plot_area = backend.plot_area();
53        let x_scale = scales.get(&Aesthetic::X);
54        let y_scale = scales.get(&Aesthetic::Y);
55
56        let base_ny = y_scale.map(|s| s.map(&Value::Float(0.0))).unwrap_or(0.0);
57
58        // Determine groups from color or fill column
59        let group_col = color_col.or(fill_col);
60        let groups: Vec<(String, Vec<usize>)> = if let Some(gc) = group_col {
61            let mut groups: Vec<(String, Vec<usize>)> = Vec::new();
62            for (i, v) in gc.iter().enumerate() {
63                let key = v.to_group_key();
64                if let Some(entry) = groups.iter_mut().find(|(k, _)| k == &key) {
65                    entry.1.push(i);
66                } else {
67                    groups.push((key, vec![i]));
68                }
69            }
70            groups
71        } else {
72            // Single group with all indices
73            vec![("".to_string(), (0..data.nrows()).collect())]
74        };
75
76        for (_, indices) in &groups {
77            if indices.is_empty() {
78                continue;
79            }
80            let first_idx = indices[0];
81
82            // Resolve fill: use fill scale if fill column exists, else use color scale, else default
83            let fill_color = fill_col
84                .and_then(|fc| scales.map_color(&Aesthetic::Fill, &fc[first_idx]))
85                .or_else(|| {
86                    color_col.and_then(|cc| scales.map_color(&Aesthetic::Color, &cc[first_idx]))
87                })
88                .unwrap_or(self.fill);
89
90            // Resolve line: use color scale if color column exists, else use fill scale, else default
91            let line_color = color_col
92                .and_then(|cc| scales.map_color(&Aesthetic::Color, &cc[first_idx]))
93                .or_else(|| {
94                    fill_col.and_then(|fc| scales.map_color(&Aesthetic::Fill, &fc[first_idx]))
95                })
96                .unwrap_or(self.color);
97
98            let mut upper: Vec<(f64, f64)> = Vec::new();
99            let mut lower: Vec<(f64, f64)> = Vec::new();
100
101            for &i in indices {
102                let nx = x_scale.map(|s| s.map(&x_col[i])).unwrap_or(0.0);
103                let ny = y_scale.map(|s| s.map(&y_col[i])).unwrap_or(0.0);
104                upper.push(coord.transform((nx, ny), &plot_area));
105                lower.push(coord.transform((nx, base_ny), &plot_area));
106            }
107
108            // Filled polygon
109            let mut polygon = upper.clone();
110            lower.reverse();
111            polygon.extend(lower);
112
113            if polygon.len() >= 3 {
114                backend.draw_polygon(
115                    &polygon,
116                    &RectStyle {
117                        fill: Some(fill_color),
118                        stroke: None,
119                        stroke_width: 0.0,
120                        alpha: self.alpha,
121                        clip: true,
122                    },
123                )?;
124            }
125
126            // Top line
127            if upper.len() >= 2 {
128                backend.draw_line(
129                    &upper,
130                    &LineStyle {
131                        color: line_color,
132                        alpha: 1.0,
133                        width: self.line_width,
134                        linetype: Linetype::Solid,
135                    },
136                )?;
137            }
138        }
139
140        Ok(())
141    }
142
143    fn required_aes(&self) -> Vec<Aesthetic> {
144        vec![Aesthetic::X]
145    }
146
147    fn default_stat(&self) -> Box<dyn Stat> {
148        Box::new(StatDensity::default())
149    }
150
151    fn default_position(&self) -> Box<dyn Position> {
152        Box::new(PositionIdentity)
153    }
154
155    fn default_params(&self) -> GeomParams {
156        GeomParams::default()
157    }
158
159    fn name(&self) -> &str {
160        "density"
161    }
162
163    fn set_series_color(&mut self, color: (u8, u8, u8)) {
164        self.fill = color;
165    }
166}