Skip to main content

plotlars_core/plots/
contourplot.rs

1use bon::bon;
2
3use crate::{
4    components::{Axis, ColorBar, Coloring, FacetConfig, FacetScales, Legend, Palette, Text},
5    ir::data::ColumnData,
6    ir::layout::LayoutIR,
7    ir::trace::{ContourPlotIR, TraceIR},
8};
9use polars::frame::DataFrame;
10
11/// A structure representing a contour plot.
12///
13/// The `ContourPlot` struct enables the creation of contour visualizations that display level
14/// curves of a three‑dimensional surface on a two‑dimensional plane. It offers extensive
15/// configuration options for contour styling, color scaling, axis appearance, legends, and
16/// annotations. Users can fine‑tune the contour interval, choose from predefined color palettes,
17/// reverse or hide the color scale, and set custom titles for both the plot and its axes in
18/// order to improve the readability of complex surfaces.
19///
20/// # Backend Support
21///
22/// | Backend | Supported |
23/// |---------|-----------|
24/// | Plotly  | Yes       |
25/// | Plotters| --        |
26///
27/// # Arguments
28///
29/// * `data` - A reference to the `DataFrame` containing the data to be plotted.
30/// * `x` - A string slice specifying the column name for x‑axis values.
31/// * `y` - A string slice specifying the column name for y‑axis values.
32/// * `z` - A string slice specifying the column name for z‑axis values whose magnitude
33///   determines each contour line.
34/// * `facet` - An optional string slice specifying the column name to be used for faceting (creating multiple subplots).
35/// * `facet_config` - An optional reference to a `FacetConfig` struct for customizing facet behavior (grid dimensions, scales, gaps, etc.).
36/// * `color_bar` - An optional reference to a `ColorBar` struct for customizing the color bar
37///   appearance.
38/// * `color_scale` - An optional `Palette` enum for specifying the color palette (e.g.,
39///   `Palette::Viridis`).
40/// * `reverse_scale` - An optional boolean to reverse the color scale direction.
41/// * `show_scale` - An optional boolean to display the color scale on the plot.
42/// * `contours` - An optional reference to a `Contours` struct for configuring the contour
43///   interval, size, and coloring.
44/// * `plot_title` - An optional `Text` struct for setting the title of the plot.
45/// * `x_title` - An optional `Text` struct for labeling the x‑axis.
46/// * `y_title` - An optional `Text` struct for labeling the y‑axis.
47/// * `x_axis` - An optional reference to an `Axis` struct for customizing x‑axis appearance.
48/// * `y_axis` - An optional reference to an `Axis` struct for customizing y‑axis appearance.
49///
50/// # Example
51///
52/// ```rust
53/// use plotlars::{Coloring, ContourPlot, Palette, Plot, Text};
54/// use polars::prelude::*;
55///
56/// let dataset = LazyCsvReader::new(PlRefPath::new("data/contour_surface.csv"))
57///     .finish()
58///     .unwrap()
59///     .collect()
60///     .unwrap();
61///
62/// ContourPlot::builder()
63///     .data(&dataset)
64///     .x("x")
65///     .y("y")
66///     .z("z")
67///     .color_scale(Palette::Viridis)
68///     .reverse_scale(true)
69///     .coloring(Coloring::Fill)
70///     .show_lines(false)
71///     .plot_title(
72///         Text::from("Contour Plot")
73///             .font("Arial")
74///             .size(18)
75///     )
76///     .build()
77///     .plot();
78/// ```
79///
80/// ![Example](https://imgur.com/VWgxHC8.png)
81#[derive(Clone)]
82#[allow(dead_code)]
83pub struct ContourPlot {
84    traces: Vec<TraceIR>,
85    layout: LayoutIR,
86}
87
88#[bon]
89impl ContourPlot {
90    #[builder(on(String, into), on(Text, into))]
91    pub fn new(
92        data: &DataFrame,
93        x: &str,
94        y: &str,
95        z: &str,
96        facet: Option<&str>,
97        facet_config: Option<&FacetConfig>,
98        color_bar: Option<&ColorBar>,
99        color_scale: Option<Palette>,
100        reverse_scale: Option<bool>,
101        show_scale: Option<bool>,
102        show_lines: Option<bool>,
103        coloring: Option<Coloring>,
104        plot_title: Option<Text>,
105        x_title: Option<Text>,
106        y_title: Option<Text>,
107        x_axis: Option<&Axis>,
108        y_axis: Option<&Axis>,
109        legend: Option<&Legend>,
110    ) -> Self {
111        let grid = facet.map(|facet_column| {
112            let config = facet_config.cloned().unwrap_or_default();
113            let facet_categories =
114                crate::data::get_unique_groups(data, facet_column, config.sorter);
115            let n_facets = facet_categories.len();
116            let (ncols, nrows) =
117                crate::faceting::calculate_grid_dimensions(n_facets, config.cols, config.rows);
118            crate::ir::facet::GridSpec {
119                kind: crate::ir::facet::FacetKind::Axis,
120                rows: nrows,
121                cols: ncols,
122                h_gap: config.h_gap,
123                v_gap: config.v_gap,
124                scales: config.scales.clone(),
125                n_facets,
126                facet_categories,
127                title_style: config.title_style.clone(),
128                x_title: x_title.clone(),
129                y_title: y_title.clone(),
130                x_axis: x_axis.cloned(),
131                y_axis: y_axis.cloned(),
132                legend_title: None,
133                legend: legend.cloned(),
134            }
135        });
136
137        let layout = LayoutIR {
138            title: plot_title.clone(),
139            x_title: if grid.is_some() {
140                None
141            } else {
142                x_title.clone()
143            },
144            y_title: if grid.is_some() {
145                None
146            } else {
147                y_title.clone()
148            },
149            y2_title: None,
150            z_title: None,
151            legend_title: None,
152            legend: if grid.is_some() {
153                None
154            } else {
155                legend.cloned()
156            },
157            dimensions: None,
158            bar_mode: None,
159            box_mode: None,
160            box_gap: None,
161            margin_bottom: None,
162            axes_2d: if grid.is_some() {
163                None
164            } else {
165                Some(crate::ir::layout::Axes2dIR {
166                    x_axis: x_axis.cloned(),
167                    y_axis: y_axis.cloned(),
168                    y2_axis: None,
169                })
170            },
171            scene_3d: None,
172            polar: None,
173            mapbox: None,
174            grid,
175            annotations: vec![],
176        };
177
178        let traces = match facet {
179            Some(facet_column) => {
180                let config = facet_config.cloned().unwrap_or_default();
181                Self::create_ir_traces_faceted(
182                    data,
183                    x,
184                    y,
185                    z,
186                    facet_column,
187                    &config,
188                    color_bar,
189                    color_scale,
190                    reverse_scale,
191                    show_scale,
192                    show_lines,
193                    coloring,
194                )
195            }
196            None => Self::create_ir_traces(
197                data,
198                x,
199                y,
200                z,
201                color_bar,
202                color_scale,
203                reverse_scale,
204                show_scale,
205                show_lines,
206                coloring,
207            ),
208        };
209
210        Self { traces, layout }
211    }
212}
213
214#[bon]
215impl ContourPlot {
216    #[builder(
217        start_fn = try_builder,
218        finish_fn = try_build,
219        builder_type = ContourPlotTryBuilder,
220        on(String, into),
221        on(Text, into),
222    )]
223    pub fn try_new(
224        data: &DataFrame,
225        x: &str,
226        y: &str,
227        z: &str,
228        facet: Option<&str>,
229        facet_config: Option<&FacetConfig>,
230        color_bar: Option<&ColorBar>,
231        color_scale: Option<Palette>,
232        reverse_scale: Option<bool>,
233        show_scale: Option<bool>,
234        show_lines: Option<bool>,
235        coloring: Option<Coloring>,
236        plot_title: Option<Text>,
237        x_title: Option<Text>,
238        y_title: Option<Text>,
239        x_axis: Option<&Axis>,
240        y_axis: Option<&Axis>,
241        legend: Option<&Legend>,
242    ) -> Result<Self, crate::io::PlotlarsError> {
243        std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
244            Self::__orig_new(
245                data,
246                x,
247                y,
248                z,
249                facet,
250                facet_config,
251                color_bar,
252                color_scale,
253                reverse_scale,
254                show_scale,
255                show_lines,
256                coloring,
257                plot_title,
258                x_title,
259                y_title,
260                x_axis,
261                y_axis,
262                legend,
263            )
264        }))
265        .map_err(|panic| {
266            let msg = panic
267                .downcast_ref::<String>()
268                .cloned()
269                .or_else(|| panic.downcast_ref::<&str>().map(|s| s.to_string()))
270                .unwrap_or_else(|| "unknown error".to_string());
271            crate::io::PlotlarsError::PlotBuild { message: msg }
272        })
273    }
274}
275
276impl ContourPlot {
277    #[allow(clippy::too_many_arguments)]
278    fn create_ir_traces(
279        data: &DataFrame,
280        x: &str,
281        y: &str,
282        z: &str,
283        color_bar: Option<&ColorBar>,
284        color_scale: Option<Palette>,
285        reverse_scale: Option<bool>,
286        show_scale: Option<bool>,
287        show_lines: Option<bool>,
288        coloring: Option<Coloring>,
289    ) -> Vec<TraceIR> {
290        vec![TraceIR::ContourPlot(ContourPlotIR {
291            x: ColumnData::Numeric(crate::data::get_numeric_column(data, x)),
292            y: ColumnData::Numeric(crate::data::get_numeric_column(data, y)),
293            z: ColumnData::Numeric(crate::data::get_numeric_column(data, z)),
294            color_scale,
295            color_bar: color_bar.cloned(),
296            coloring,
297            show_lines,
298            show_labels: None,
299            n_contours: None,
300            reverse_scale,
301            show_scale,
302            z_min: None,
303            z_max: None,
304            subplot_ref: None,
305        })]
306    }
307
308    #[allow(clippy::too_many_arguments)]
309    fn create_ir_traces_faceted(
310        data: &DataFrame,
311        x: &str,
312        y: &str,
313        z: &str,
314        facet_column: &str,
315        config: &FacetConfig,
316        color_bar: Option<&ColorBar>,
317        color_scale: Option<Palette>,
318        reverse_scale: Option<bool>,
319        show_scale: Option<bool>,
320        show_lines: Option<bool>,
321        coloring: Option<Coloring>,
322    ) -> Vec<TraceIR> {
323        const MAX_FACETS: usize = 8;
324
325        let facet_categories = crate::data::get_unique_groups(data, facet_column, config.sorter);
326
327        if facet_categories.len() > MAX_FACETS {
328            panic!(
329                "Facet column '{}' has {} unique values, but plotly.rs supports maximum {} subplots",
330                facet_column,
331                facet_categories.len(),
332                MAX_FACETS
333            );
334        }
335
336        let use_global_z = !matches!(config.scales, FacetScales::Free);
337        let z_range = if use_global_z {
338            Self::calculate_global_z_range(data, z)
339        } else {
340            None
341        };
342
343        let mut traces = Vec::new();
344
345        for (facet_idx, facet_value) in facet_categories.iter().enumerate() {
346            let facet_data = crate::data::filter_data_by_group(data, facet_column, facet_value);
347
348            let subplot_ref = format!(
349                "{}{}",
350                crate::faceting::get_axis_reference(facet_idx, "x"),
351                crate::faceting::get_axis_reference(facet_idx, "y")
352            );
353
354            let show_scale_for_facet = if facet_idx == 0 {
355                show_scale
356            } else {
357                Some(false)
358            };
359
360            let (z_min, z_max) = match z_range {
361                Some((zmin, zmax)) => (Some(zmin), Some(zmax)),
362                None => (None, None),
363            };
364
365            traces.push(TraceIR::ContourPlot(ContourPlotIR {
366                x: ColumnData::Numeric(crate::data::get_numeric_column(&facet_data, x)),
367                y: ColumnData::Numeric(crate::data::get_numeric_column(&facet_data, y)),
368                z: ColumnData::Numeric(crate::data::get_numeric_column(&facet_data, z)),
369                color_scale,
370                color_bar: color_bar.cloned(),
371                coloring,
372                show_lines,
373                show_labels: None,
374                n_contours: None,
375                reverse_scale,
376                show_scale: show_scale_for_facet,
377                z_min,
378                z_max,
379                subplot_ref: Some(subplot_ref),
380            }));
381        }
382
383        traces
384    }
385
386    fn calculate_global_z_range(data: &DataFrame, z: &str) -> Option<(f64, f64)> {
387        let z_data = crate::data::get_numeric_column(data, z);
388
389        let mut z_min = f64::INFINITY;
390        let mut z_max = f64::NEG_INFINITY;
391        let mut found_valid = false;
392
393        for val in z_data.iter().flatten() {
394            let val_f64 = *val as f64;
395            if !val_f64.is_nan() {
396                z_min = z_min.min(val_f64);
397                z_max = z_max.max(val_f64);
398                found_valid = true;
399            }
400        }
401
402        if found_valid {
403            Some((z_min, z_max))
404        } else {
405            None
406        }
407    }
408}
409
410impl crate::Plot for ContourPlot {
411    fn ir_traces(&self) -> &[TraceIR] {
412        &self.traces
413    }
414
415    fn ir_layout(&self) -> &LayoutIR {
416        &self.layout
417    }
418}
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423    use crate::Plot;
424    use polars::prelude::*;
425
426    #[test]
427    fn test_basic_one_trace() {
428        let df = df![
429            "x" => [1.0, 2.0, 3.0],
430            "y" => [4.0, 5.0, 6.0],
431            "z" => [7.0, 8.0, 9.0]
432        ]
433        .unwrap();
434        let plot = ContourPlot::builder()
435            .data(&df)
436            .x("x")
437            .y("y")
438            .z("z")
439            .build();
440        assert_eq!(plot.ir_traces().len(), 1);
441        assert!(matches!(plot.ir_traces()[0], TraceIR::ContourPlot(_)));
442    }
443
444    #[test]
445    fn test_layout_has_axes() {
446        let df = df![
447            "x" => [1.0, 2.0],
448            "y" => [3.0, 4.0],
449            "z" => [5.0, 6.0]
450        ]
451        .unwrap();
452        let plot = ContourPlot::builder()
453            .data(&df)
454            .x("x")
455            .y("y")
456            .z("z")
457            .build();
458        assert!(plot.ir_layout().axes_2d.is_some());
459    }
460
461    #[test]
462    fn test_layout_title() {
463        let df = df![
464            "x" => [1.0],
465            "y" => [2.0],
466            "z" => [3.0]
467        ]
468        .unwrap();
469        let plot = ContourPlot::builder()
470            .data(&df)
471            .x("x")
472            .y("y")
473            .z("z")
474            .plot_title("Contour")
475            .build();
476        assert!(plot.ir_layout().title.is_some());
477    }
478}