plotlars/plots/
heatmap.rs

1use bon::bon;
2
3use plotly::{HeatMap as HeatMapPlotly, Layout as LayoutPlotly, Trace};
4
5use polars::frame::DataFrame;
6use serde::Serialize;
7
8use crate::{
9    common::{Layout, Marker, PlotHelper, Polar},
10    components::{Axis, ColorBar, Palette, Text},
11};
12
13/// A structure representing a heat map.
14///
15/// The `HeatMap` struct enables the creation of heat map visualizations with options for color scaling,
16/// axis customization, legend adjustments, and data value formatting. Users can customize the color
17/// scale, adjust the color bar, and set titles for the plot and axes, as well as format ticks and scales
18/// for improved data readability.
19///
20/// # Arguments
21///
22/// * `data` - A reference to the `DataFrame` containing the data to be plotted.
23/// * `x` - A string slice specifying the column name for x-axis values.
24/// * `y` - A string slice specifying the column name for y-axis values.
25/// * `z` - A string slice specifying the column name for z-axis values, which are represented by the color intensity.
26/// * `auto_color_scale` - An optional boolean for enabling automatic color scaling based on data.
27/// * `color_bar` - An optional reference to a `ColorBar` struct for customizing the color bar appearance.
28/// * `color_scale` - An optional `Palette` enum for specifying the color scale (e.g., Viridis).
29/// * `reverse_scale` - An optional boolean to reverse the color scale direction.
30/// * `show_scale` - An optional boolean to display the color scale on the plot.
31/// * `plot_title` - An optional `Text` struct for setting the title of the plot.
32/// * `x_title` - An optional `Text` struct for labeling the x-axis.
33/// * `y_title` - An optional `Text` struct for labeling the y-axis.
34/// * `x_axis` - An optional reference to an `Axis` struct for customizing x-axis appearance.
35/// * `y_axis` - An optional reference to an `Axis` struct for customizing y-axis appearance.
36///
37/// # Example
38///
39/// ```rust
40/// use polars::prelude::*;
41/// use plotlars::{ColorBar, HeatMap, Palette, Plot, Text, ValueExponent};
42///
43/// let dataset = LazyCsvReader::new(PlPath::new("data/heatmap.csv"))
44///     .finish()
45///     .unwrap()
46///     .collect()
47///     .unwrap();
48///
49/// HeatMap::builder()
50///     .data(&dataset)
51///     .x("x")
52///     .y("y")
53///     .z("z")
54///     .color_bar(
55///         &ColorBar::new()
56///             .length(290)
57///             .value_exponent(ValueExponent::None)
58///             .separate_thousands(true)
59///             .tick_length(5)
60///             .tick_step(2500.0)
61///     )
62///     .plot_title(
63///         Text::from("Heat Map")
64///             .font("Arial")
65///             .size(18)
66///     )
67///     .color_scale(Palette::Viridis)
68///     .build()
69///     .plot();
70/// ```
71///
72/// ![Example](https://imgur.com/5uFih4M.png)
73#[derive(Clone, Serialize)]
74pub struct HeatMap {
75    pub traces: Vec<Box<dyn Trace + 'static>>,
76    pub layout: LayoutPlotly,
77}
78
79#[bon]
80impl HeatMap {
81    #[builder(on(String, into), on(Text, into))]
82    pub fn new(
83        data: &DataFrame,
84        x: &str,
85        y: &str,
86        z: &str,
87        auto_color_scale: Option<bool>,
88        color_bar: Option<&ColorBar>,
89        color_scale: Option<Palette>,
90        reverse_scale: Option<bool>,
91        show_scale: Option<bool>,
92        plot_title: Option<Text>,
93        x_title: Option<Text>,
94        y_title: Option<Text>,
95        x_axis: Option<&Axis>,
96        y_axis: Option<&Axis>,
97    ) -> Self {
98        let legend = None;
99        let legend_title = None;
100        let z_title = None;
101        let z_axis = None;
102        let y2_title = None;
103        let y2_axis = None;
104
105        let layout = Self::create_layout(
106            plot_title,
107            x_title,
108            y_title,
109            y2_title,
110            z_title,
111            legend_title,
112            x_axis,
113            y_axis,
114            y2_axis,
115            z_axis,
116            legend,
117        );
118
119        let traces = Self::create_traces(
120            data,
121            x,
122            y,
123            z,
124            auto_color_scale,
125            color_bar,
126            color_scale,
127            reverse_scale,
128            show_scale,
129        );
130
131        Self { traces, layout }
132    }
133
134    #[allow(clippy::too_many_arguments)]
135    fn create_traces(
136        data: &DataFrame,
137        x: &str,
138        y: &str,
139        z: &str,
140        auto_color_scale: Option<bool>,
141        color_bar: Option<&ColorBar>,
142        color_scale: Option<Palette>,
143        reverse_scale: Option<bool>,
144        show_scale: Option<bool>,
145    ) -> Vec<Box<dyn Trace + 'static>> {
146        let mut traces: Vec<Box<dyn Trace + 'static>> = Vec::new();
147        let trace = Self::create_trace(
148            data,
149            x,
150            y,
151            z,
152            auto_color_scale,
153            color_bar,
154            color_scale,
155            reverse_scale,
156            show_scale,
157        );
158
159        traces.push(trace);
160        traces
161    }
162
163    #[allow(clippy::too_many_arguments)]
164    fn create_trace(
165        data: &DataFrame,
166        x: &str,
167        y: &str,
168        z: &str,
169        auto_color_scale: Option<bool>,
170        color_bar: Option<&ColorBar>,
171        color_scale: Option<Palette>,
172        reverse_scale: Option<bool>,
173        show_scale: Option<bool>,
174    ) -> Box<dyn Trace + 'static> {
175        let x = Self::get_string_column(data, x);
176        let y = Self::get_string_column(data, y);
177        let z = Self::get_numeric_column(data, z);
178
179        let mut trace = HeatMapPlotly::default().x(x).y(y).z(z);
180
181        trace = Self::set_auto_color_scale(trace, auto_color_scale);
182        trace = Self::set_color_bar(trace, color_bar);
183        trace = Self::set_color_scale(trace, color_scale);
184        trace = Self::set_reverse_scale(trace, reverse_scale);
185        trace = Self::set_show_scale(trace, show_scale);
186        trace
187    }
188
189    fn set_auto_color_scale<X, Y, Z>(
190        mut trace: Box<HeatMapPlotly<X, Y, Z>>,
191        auto_color_scale: Option<bool>,
192    ) -> Box<HeatMapPlotly<X, Y, Z>>
193    where
194        X: Serialize + Clone,
195        Y: Serialize + Clone,
196        Z: Serialize + Clone,
197    {
198        if let Some(auto_color_scale) = auto_color_scale {
199            trace = trace.auto_color_scale(auto_color_scale);
200        }
201
202        trace
203    }
204
205    fn set_color_bar<X, Y, Z>(
206        mut trace: Box<HeatMapPlotly<X, Y, Z>>,
207        color_bar: Option<&ColorBar>,
208    ) -> Box<HeatMapPlotly<X, Y, Z>>
209    where
210        X: Serialize + Clone,
211        Y: Serialize + Clone,
212        Z: Serialize + Clone,
213    {
214        if let Some(color_bar) = color_bar {
215            trace = trace.color_bar(color_bar.to_plotly())
216        }
217
218        trace
219    }
220
221    fn set_color_scale<X, Y, Z>(
222        mut trace: Box<HeatMapPlotly<X, Y, Z>>,
223        color_scale: Option<Palette>,
224    ) -> Box<HeatMapPlotly<X, Y, Z>>
225    where
226        X: Serialize + Clone,
227        Y: Serialize + Clone,
228        Z: Serialize + Clone,
229    {
230        if let Some(color_scale) = color_scale {
231            trace = trace.color_scale(color_scale.to_plotly());
232        }
233
234        trace
235    }
236
237    fn set_reverse_scale<X, Y, Z>(
238        mut trace: Box<HeatMapPlotly<X, Y, Z>>,
239        reverse_scale: Option<bool>,
240    ) -> Box<HeatMapPlotly<X, Y, Z>>
241    where
242        X: Serialize + Clone,
243        Y: Serialize + Clone,
244        Z: Serialize + Clone,
245    {
246        if let Some(reverse_scale) = reverse_scale {
247            trace = trace.reverse_scale(reverse_scale);
248        }
249        trace
250    }
251
252    fn set_show_scale<X, Y, Z>(
253        mut trace: Box<HeatMapPlotly<X, Y, Z>>,
254        show_scale: Option<bool>,
255    ) -> Box<HeatMapPlotly<X, Y, Z>>
256    where
257        X: Serialize + Clone,
258        Y: Serialize + Clone,
259        Z: Serialize + Clone,
260    {
261        if let Some(show_scale) = show_scale {
262            trace = trace.show_scale(show_scale);
263        }
264        trace
265    }
266}
267
268impl Layout for HeatMap {}
269impl Marker for HeatMap {}
270impl Polar for HeatMap {}
271
272impl PlotHelper for HeatMap {
273    fn get_layout(&self) -> &LayoutPlotly {
274        &self.layout
275    }
276
277    fn get_traces(&self) -> &Vec<Box<dyn Trace + 'static>> {
278        &self.traces
279    }
280}
plotlars/plots/heatmap.rs

plotlars/plots/
heatmap.rs
