Skip to main content

ggplot_rs/geom/
density2d.rs

1use crate::aes::Aesthetic;
2use crate::coord::Coord;
3use crate::data::DataFrame;
4use crate::position::identity::PositionIdentity;
5use crate::position::Position;
6use crate::render::backend::{DrawBackend, LineStyle, Linetype};
7use crate::render::RenderError;
8use crate::scale::ScaleSet;
9use crate::stat::density2d::StatDensity2d;
10use crate::stat::Stat;
11use crate::theme::Theme;
12
13use super::{Geom, GeomParams};
14
15/// 2D density contour geometry — draws density contour lines from point cloud data.
16pub struct GeomDensity2d {
17    pub color: (u8, u8, u8),
18    pub alpha: f64,
19    pub width: f64,
20    pub n_grid: usize,
21    pub n_levels: usize,
22}
23
24impl Default for GeomDensity2d {
25    fn default() -> Self {
26        GeomDensity2d {
27            color: (50, 50, 50),
28            alpha: 1.0,
29            width: 0.8,
30            n_grid: 50,
31            n_levels: 10,
32        }
33    }
34}
35
36impl GeomDensity2d {
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    pub fn with_color(mut self, color: (u8, u8, u8)) -> Self {
42        self.color = color;
43        self
44    }
45
46    pub fn with_alpha(mut self, alpha: f64) -> Self {
47        self.alpha = alpha;
48        self
49    }
50
51    pub fn with_levels(mut self, n: usize) -> Self {
52        self.n_levels = n;
53        self
54    }
55}
56
57impl Geom for GeomDensity2d {
58    fn draw(
59        &self,
60        data: &DataFrame,
61        coord: &dyn Coord,
62        scales: &ScaleSet,
63        _theme: &Theme,
64        backend: &mut dyn DrawBackend,
65    ) -> Result<(), RenderError> {
66        let x_col = match data.column("x") {
67            Some(c) => c,
68            None => return Ok(()),
69        };
70        let y_col = match data.column("y") {
71            Some(c) => c,
72            None => return Ok(()),
73        };
74        let group_col = data.column("group");
75        let level_col = data.column("level");
76
77        let plot_area = backend.plot_area();
78        let x_scale = scales.get(&Aesthetic::X);
79        let y_scale = scales.get(&Aesthetic::Y);
80
81        // Draw segment pairs grouped by group column (same as GeomContour)
82        let nrows = data.nrows();
83        let mut i = 0;
84        while i + 1 < nrows {
85            let same_group = match group_col {
86                Some(gc) => gc[i].to_group_key() == gc[i + 1].to_group_key(),
87                None => true,
88            };
89
90            if same_group {
91                let nx0 = x_scale.map(|s| s.map(&x_col[i])).unwrap_or(0.0);
92                let ny0 = y_scale.map(|s| s.map(&y_col[i])).unwrap_or(0.0);
93                let nx1 = x_scale.map(|s| s.map(&x_col[i + 1])).unwrap_or(0.0);
94                let ny1 = y_scale.map(|s| s.map(&y_col[i + 1])).unwrap_or(0.0);
95
96                let p0 = coord.transform((nx0, ny0), &plot_area);
97                let p1 = coord.transform((nx1, ny1), &plot_area);
98
99                let color = if let Some(lc) = level_col {
100                    scales
101                        .map_color(&Aesthetic::Color, &lc[i])
102                        .unwrap_or(self.color)
103                } else {
104                    self.color
105                };
106
107                backend.draw_line(
108                    &[p0, p1],
109                    &LineStyle {
110                        color,
111                        alpha: self.alpha,
112                        width: self.width,
113                        linetype: Linetype::Solid,
114                    },
115                )?;
116
117                i += 2;
118            } else {
119                i += 1;
120            }
121        }
122
123        Ok(())
124    }
125
126    fn required_aes(&self) -> Vec<Aesthetic> {
127        vec![Aesthetic::X, Aesthetic::Y]
128    }
129
130    fn default_stat(&self) -> Box<dyn Stat> {
131        Box::new(StatDensity2d {
132            n_grid: self.n_grid,
133            n_levels: self.n_levels,
134        })
135    }
136
137    fn default_position(&self) -> Box<dyn Position> {
138        Box::new(PositionIdentity)
139    }
140
141    fn default_params(&self) -> GeomParams {
142        GeomParams::default()
143    }
144
145    fn name(&self) -> &str {
146        "density_2d"
147    }
148}