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#[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}