1use bon::bon;
2
3use polars::{
4 frame::DataFrame,
5 prelude::{col, IntoLazy},
6};
7
8use crate::{
9 components::{Axis, FacetConfig, Legend, Line as LineStyle, Mode, Rgb, Shape, Text},
10 ir::data::ColumnData,
11 ir::layout::LayoutIR,
12 ir::line::LineIR,
13 ir::marker::MarkerIR,
14 ir::trace::{TimeSeriesPlotIR, TraceIR},
15};
16
17#[derive(Clone)]
169#[allow(dead_code)]
170pub struct TimeSeriesPlot {
171 traces: Vec<TraceIR>,
172 layout: LayoutIR,
173}
174
175#[bon]
176impl TimeSeriesPlot {
177 #[builder(on(String, into), on(Text, into))]
178 pub fn new(
179 data: &DataFrame,
180 x: &str,
181 y: &str,
182 additional_series: Option<Vec<&str>>,
183 facet: Option<&str>,
184 facet_config: Option<&FacetConfig>,
185 size: Option<usize>,
186 color: Option<Rgb>,
187 colors: Option<Vec<Rgb>>,
188 shape: Option<Shape>,
189 shapes: Option<Vec<Shape>>,
190 width: Option<f64>,
191 line: Option<LineStyle>,
192 lines: Option<Vec<LineStyle>>,
193 plot_title: Option<Text>,
194 x_title: Option<Text>,
195 y_title: Option<Text>,
196 y2_title: Option<Text>,
197 legend_title: Option<Text>,
198 x_axis: Option<&Axis>,
199 y_axis: Option<&Axis>,
200 y2_axis: Option<&Axis>,
201 legend: Option<&Legend>,
202 ) -> Self {
203 let grid = facet.map(|facet_column| {
204 let config = facet_config.cloned().unwrap_or_default();
205 let facet_categories =
206 crate::data::get_unique_groups(data, facet_column, config.sorter);
207 let n_facets = facet_categories.len();
208 let (ncols, nrows) =
209 crate::faceting::calculate_grid_dimensions(n_facets, config.cols, config.rows);
210 crate::ir::facet::GridSpec {
211 kind: crate::ir::facet::FacetKind::Axis,
212 rows: nrows,
213 cols: ncols,
214 h_gap: config.h_gap,
215 v_gap: config.v_gap,
216 scales: config.scales.clone(),
217 n_facets,
218 facet_categories,
219 title_style: config.title_style.clone(),
220 x_title: x_title.clone(),
221 y_title: y_title.clone(),
222 x_axis: x_axis.cloned(),
223 y_axis: y_axis.cloned(),
224 legend_title: legend_title.clone(),
225 legend: legend.cloned(),
226 }
227 });
228
229 let layout = LayoutIR {
230 title: plot_title.clone(),
231 x_title: if grid.is_some() {
232 None
233 } else {
234 x_title.clone()
235 },
236 y_title: if grid.is_some() {
237 None
238 } else {
239 y_title.clone()
240 },
241 y2_title: if grid.is_some() {
242 None
243 } else {
244 y2_title.clone()
245 },
246 z_title: None,
247 legend_title: if grid.is_some() {
248 None
249 } else {
250 legend_title.clone()
251 },
252 legend: if grid.is_some() {
253 None
254 } else {
255 legend.cloned()
256 },
257 dimensions: None,
258 bar_mode: None,
259 box_mode: None,
260 box_gap: None,
261 margin_bottom: None,
262 axes_2d: if grid.is_some() {
263 None
264 } else {
265 Some(crate::ir::layout::Axes2dIR {
266 x_axis: x_axis.cloned(),
267 y_axis: y_axis.cloned(),
268 y2_axis: y2_axis.cloned(),
269 })
270 },
271 scene_3d: None,
272 polar: None,
273 mapbox: None,
274 grid,
275 annotations: vec![],
276 };
277
278 let traces = match facet {
279 Some(facet_column) => {
280 let config = facet_config.cloned().unwrap_or_default();
281 Self::create_ir_traces_faceted(
282 data,
283 x,
284 y,
285 additional_series,
286 facet_column,
287 &config,
288 size,
289 color,
290 colors,
291 shape,
292 shapes,
293 width,
294 line,
295 lines,
296 )
297 }
298 None => Self::create_ir_traces(
299 data,
300 x,
301 y,
302 additional_series,
303 y2_axis.is_some(),
304 size,
305 color,
306 colors,
307 shape,
308 shapes,
309 width,
310 line,
311 lines,
312 ),
313 };
314
315 Self { traces, layout }
316 }
317}
318
319#[bon]
320impl TimeSeriesPlot {
321 #[builder(
322 start_fn = try_builder,
323 finish_fn = try_build,
324 builder_type = TimeSeriesPlotTryBuilder,
325 on(String, into),
326 on(Text, into),
327 )]
328 pub fn try_new(
329 data: &DataFrame,
330 x: &str,
331 y: &str,
332 additional_series: Option<Vec<&str>>,
333 facet: Option<&str>,
334 facet_config: Option<&FacetConfig>,
335 size: Option<usize>,
336 color: Option<Rgb>,
337 colors: Option<Vec<Rgb>>,
338 shape: Option<Shape>,
339 shapes: Option<Vec<Shape>>,
340 width: Option<f64>,
341 line: Option<LineStyle>,
342 lines: Option<Vec<LineStyle>>,
343 plot_title: Option<Text>,
344 x_title: Option<Text>,
345 y_title: Option<Text>,
346 y2_title: Option<Text>,
347 legend_title: Option<Text>,
348 x_axis: Option<&Axis>,
349 y_axis: Option<&Axis>,
350 y2_axis: Option<&Axis>,
351 legend: Option<&Legend>,
352 ) -> Result<Self, crate::io::PlotlarsError> {
353 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
354 Self::__orig_new(
355 data,
356 x,
357 y,
358 additional_series,
359 facet,
360 facet_config,
361 size,
362 color,
363 colors,
364 shape,
365 shapes,
366 width,
367 line,
368 lines,
369 plot_title,
370 x_title,
371 y_title,
372 y2_title,
373 legend_title,
374 x_axis,
375 y_axis,
376 y2_axis,
377 legend,
378 )
379 }))
380 .map_err(|panic| {
381 let msg = panic
382 .downcast_ref::<String>()
383 .cloned()
384 .or_else(|| panic.downcast_ref::<&str>().map(|s| s.to_string()))
385 .unwrap_or_else(|| "unknown error".to_string());
386 crate::io::PlotlarsError::PlotBuild { message: msg }
387 })
388 }
389}
390
391impl TimeSeriesPlot {
392 #[allow(clippy::too_many_arguments)]
393 fn create_ir_traces(
394 data: &DataFrame,
395 x_col: &str,
396 y_col: &str,
397 additional_series: Option<Vec<&str>>,
398 has_y2_axis: bool,
399 size: Option<usize>,
400 color: Option<Rgb>,
401 colors: Option<Vec<Rgb>>,
402 shape: Option<Shape>,
403 shapes: Option<Vec<Shape>>,
404 width: Option<f64>,
405 style: Option<LineStyle>,
406 styles: Option<Vec<LineStyle>>,
407 ) -> Vec<TraceIR> {
408 let mut traces = Vec::new();
409
410 let mode = Self::resolve_mode(shape, shapes.as_ref());
411
412 let marker_ir = MarkerIR {
413 opacity: None,
414 size,
415 color: Self::resolve_color(0, color, colors.clone()),
416 shape: Self::resolve_shape(0, shape, shapes.clone()),
417 };
418
419 let line_ir = Self::resolve_line_ir(0, width, style, styles.clone());
420
421 traces.push(TraceIR::TimeSeriesPlot(TimeSeriesPlotIR {
422 x: ColumnData::String(crate::data::get_string_column(data, x_col)),
423 y: ColumnData::Numeric(crate::data::get_numeric_column(data, y_col)),
424 name: Some(y_col.to_string()),
425 marker: Some(marker_ir),
426 line: Some(line_ir),
427 mode,
428 show_legend: None,
429 legend_group: None,
430 y_axis_ref: Some(String::new()),
431 subplot_ref: None,
432 }));
433
434 if let Some(additional_series) = additional_series {
435 let mut y_axis_ref = String::new();
436
437 for (i, series) in additional_series.into_iter().enumerate() {
438 let subset = data
439 .clone()
440 .lazy()
441 .select([col(x_col), col(series)])
442 .collect()
443 .unwrap();
444
445 let marker_ir = MarkerIR {
446 opacity: None,
447 size,
448 color: Self::resolve_color(i + 1, color, colors.clone()),
449 shape: Self::resolve_shape(i + 1, shape, shapes.clone()),
450 };
451
452 let line_ir = Self::resolve_line_ir(i + 1, width, style, styles.clone());
453
454 if has_y2_axis {
455 y_axis_ref = "y2".to_string();
456 }
457
458 traces.push(TraceIR::TimeSeriesPlot(TimeSeriesPlotIR {
459 x: ColumnData::String(crate::data::get_string_column(&subset, x_col)),
460 y: ColumnData::Numeric(crate::data::get_numeric_column(&subset, series)),
461 name: Some(series.to_string()),
462 marker: Some(marker_ir),
463 line: Some(line_ir),
464 mode,
465 show_legend: None,
466 legend_group: None,
467 y_axis_ref: Some(y_axis_ref.clone()),
468 subplot_ref: None,
469 }));
470 }
471 }
472
473 traces
474 }
475
476 #[allow(clippy::too_many_arguments)]
477 fn create_ir_traces_faceted(
478 data: &DataFrame,
479 x: &str,
480 y: &str,
481 additional_series: Option<Vec<&str>>,
482 facet_column: &str,
483 config: &FacetConfig,
484 size: Option<usize>,
485 color: Option<Rgb>,
486 colors: Option<Vec<Rgb>>,
487 shape: Option<Shape>,
488 shapes: Option<Vec<Shape>>,
489 width: Option<f64>,
490 style: Option<LineStyle>,
491 styles: Option<Vec<LineStyle>>,
492 ) -> Vec<TraceIR> {
493 const MAX_FACETS: usize = 8;
494
495 let facet_categories = crate::data::get_unique_groups(data, facet_column, config.sorter);
496
497 if facet_categories.len() > MAX_FACETS {
498 panic!(
499 "Facet column '{}' has {} unique values, but plotly.rs supports maximum {} subplots",
500 facet_column,
501 facet_categories.len(),
502 MAX_FACETS
503 );
504 }
505
506 let all_y_cols = if let Some(ref add_series) = additional_series {
507 let mut cols = vec![y];
508 cols.extend(add_series.iter().copied());
509 cols
510 } else {
511 vec![y]
512 };
513
514 if let Some(ref color_vec) = colors {
515 if additional_series.is_none() {
516 let color_count = color_vec.len();
517 let facet_count = facet_categories.len();
518 if color_count != facet_count {
519 panic!(
520 "When using colors with facet (without additional_series), colors.len() must equal number of facets. \
521 Expected {} colors for {} facets, but got {} colors. \
522 Each facet must be assigned exactly one color.",
523 facet_count, facet_count, color_count
524 );
525 }
526 } else {
527 let color_count = color_vec.len();
528 let series_count = all_y_cols.len();
529 if color_count < series_count {
530 panic!(
531 "When using colors with additional_series, colors.len() must be >= number of series. \
532 Need at least {} colors for {} series, but got {} colors",
533 series_count, series_count, color_count
534 );
535 }
536 }
537 }
538
539 let mut traces = Vec::new();
540
541 let mode_for_facet = Self::resolve_mode(shape, shapes.as_ref());
542
543 if config.highlight_facet {
544 for (facet_idx, facet_value) in facet_categories.iter().enumerate() {
545 let subplot_ref = format!(
546 "{}{}",
547 crate::faceting::get_axis_reference(facet_idx, "x"),
548 crate::faceting::get_axis_reference(facet_idx, "y")
549 );
550
551 for other_facet_value in facet_categories.iter() {
552 if other_facet_value != facet_value {
553 let other_data = crate::data::filter_data_by_group(
554 data,
555 facet_column,
556 other_facet_value,
557 );
558
559 let grey_color = config.unhighlighted_color.unwrap_or(Rgb(200, 200, 200));
560
561 for (series_idx, y_col) in all_y_cols.iter().enumerate() {
562 let marker_ir = MarkerIR {
563 opacity: None,
564 size,
565 color: Some(grey_color),
566 shape: Self::resolve_shape(series_idx, shape, None),
567 };
568
569 let line_ir =
570 Self::resolve_line_ir(series_idx, width, style, styles.clone());
571
572 traces.push(TraceIR::TimeSeriesPlot(TimeSeriesPlotIR {
573 x: ColumnData::String(crate::data::get_string_column(
574 &other_data,
575 x,
576 )),
577 y: ColumnData::Numeric(crate::data::get_numeric_column(
578 &other_data,
579 y_col,
580 )),
581 name: None,
582 marker: Some(marker_ir),
583 line: Some(line_ir),
584 mode: mode_for_facet,
585 show_legend: Some(false),
586 legend_group: None,
587 y_axis_ref: None,
588 subplot_ref: Some(subplot_ref.clone()),
589 }));
590 }
591 }
592 }
593
594 let facet_data = crate::data::filter_data_by_group(data, facet_column, facet_value);
595
596 for (series_idx, y_col) in all_y_cols.iter().enumerate() {
597 let color_index = if additional_series.is_none() {
598 facet_idx
599 } else {
600 series_idx
601 };
602
603 let marker_ir = MarkerIR {
604 opacity: None,
605 size,
606 color: Self::resolve_color(color_index, color, colors.clone()),
607 shape: Self::resolve_shape(color_index, shape, shapes.clone()),
608 };
609
610 let line_ir = Self::resolve_line_ir(series_idx, width, style, styles.clone());
611
612 let show_legend = facet_idx == 0;
613 let name = if show_legend {
614 Some(y_col.to_string())
615 } else {
616 None
617 };
618
619 traces.push(TraceIR::TimeSeriesPlot(TimeSeriesPlotIR {
620 x: ColumnData::String(crate::data::get_string_column(&facet_data, x)),
621 y: ColumnData::Numeric(crate::data::get_numeric_column(&facet_data, y_col)),
622 name,
623 marker: Some(marker_ir),
624 line: Some(line_ir),
625 mode: mode_for_facet,
626 show_legend: Some(show_legend),
627 legend_group: Some(y_col.to_string()),
628 y_axis_ref: None,
629 subplot_ref: Some(subplot_ref.clone()),
630 }));
631 }
632 }
633 } else {
634 for (facet_idx, facet_value) in facet_categories.iter().enumerate() {
635 let facet_data = crate::data::filter_data_by_group(data, facet_column, facet_value);
636
637 let subplot_ref = format!(
638 "{}{}",
639 crate::faceting::get_axis_reference(facet_idx, "x"),
640 crate::faceting::get_axis_reference(facet_idx, "y")
641 );
642
643 for (series_idx, y_col) in all_y_cols.iter().enumerate() {
644 let color_index = if additional_series.is_none() {
645 facet_idx
646 } else {
647 series_idx
648 };
649
650 let marker_ir = MarkerIR {
651 opacity: None,
652 size,
653 color: Self::resolve_color(color_index, color, colors.clone()),
654 shape: Self::resolve_shape(color_index, shape, shapes.clone()),
655 };
656
657 let line_ir = Self::resolve_line_ir(series_idx, width, style, styles.clone());
658
659 let show_legend = facet_idx == 0;
660 let name = if show_legend {
661 Some(y_col.to_string())
662 } else {
663 None
664 };
665
666 traces.push(TraceIR::TimeSeriesPlot(TimeSeriesPlotIR {
667 x: ColumnData::String(crate::data::get_string_column(&facet_data, x)),
668 y: ColumnData::Numeric(crate::data::get_numeric_column(&facet_data, y_col)),
669 name,
670 marker: Some(marker_ir),
671 line: Some(line_ir),
672 mode: mode_for_facet,
673 show_legend: Some(show_legend),
674 legend_group: Some(y_col.to_string()),
675 y_axis_ref: None,
676 subplot_ref: Some(subplot_ref.clone()),
677 }));
678 }
679 }
680 }
681
682 traces
683 }
684
685 fn resolve_color(index: usize, color: Option<Rgb>, colors: Option<Vec<Rgb>>) -> Option<Rgb> {
686 if let Some(c) = color {
687 return Some(c);
688 }
689 if let Some(ref cs) = colors {
690 return cs.get(index).copied();
691 }
692 None
693 }
694
695 fn resolve_shape(
696 index: usize,
697 shape: Option<Shape>,
698 shapes: Option<Vec<Shape>>,
699 ) -> Option<Shape> {
700 if let Some(s) = shape {
701 return Some(s);
702 }
703 if let Some(ref ss) = shapes {
704 return ss.get(index).cloned();
705 }
706 None
707 }
708
709 fn resolve_line_ir(
710 index: usize,
711 width: Option<f64>,
712 style: Option<LineStyle>,
713 styles: Option<Vec<LineStyle>>,
714 ) -> LineIR {
715 let resolved_style = if style.is_some() {
716 style
717 } else {
718 styles.and_then(|ss| ss.get(index).cloned())
719 };
720
721 LineIR {
722 width,
723 style: resolved_style,
724 color: None,
725 }
726 }
727
728 fn resolve_mode(shape: Option<Shape>, shapes: Option<&Vec<Shape>>) -> Option<Mode> {
729 if shape.is_some() || shapes.is_some() {
730 Some(Mode::LinesMarkers)
731 } else {
732 Some(Mode::Lines)
733 }
734 }
735}
736
737impl crate::Plot for TimeSeriesPlot {
738 fn ir_traces(&self) -> &[TraceIR] {
739 &self.traces
740 }
741
742 fn ir_layout(&self) -> &LayoutIR {
743 &self.layout
744 }
745}
746
747#[cfg(test)]
748mod tests {
749 use super::*;
750 use crate::Plot;
751 use polars::prelude::*;
752
753 #[test]
754 fn test_basic_one_trace() {
755 let df = df![
756 "x" => ["2024-01", "2024-02", "2024-03"],
757 "y" => [1.0, 2.0, 3.0]
758 ]
759 .unwrap();
760 let plot = TimeSeriesPlot::builder().data(&df).x("x").y("y").build();
761 assert_eq!(plot.ir_traces().len(), 1);
762 assert!(matches!(plot.ir_traces()[0], TraceIR::TimeSeriesPlot(_)));
763 }
764
765 #[test]
766 fn test_with_additional_series() {
767 let df = df![
768 "x" => ["2024-01", "2024-02"],
769 "y" => [1.0, 2.0],
770 "y2" => [3.0, 4.0]
771 ]
772 .unwrap();
773 let plot = TimeSeriesPlot::builder()
774 .data(&df)
775 .x("x")
776 .y("y")
777 .additional_series(vec!["y2"])
778 .build();
779 assert_eq!(plot.ir_traces().len(), 2);
780 }
781
782 #[test]
783 fn test_layout_titles() {
784 let df = df![
785 "x" => ["2024-01", "2024-02"],
786 "y" => [1.0, 2.0]
787 ]
788 .unwrap();
789 let plot = TimeSeriesPlot::builder()
790 .data(&df)
791 .x("x")
792 .y("y")
793 .plot_title("My Title")
794 .x_title("X")
795 .y_title("Y")
796 .build();
797 let layout = plot.ir_layout();
798 assert!(layout.title.is_some());
799 assert!(layout.x_title.is_some());
800 assert!(layout.y_title.is_some());
801 }
802
803 #[test]
804 fn test_faceted_trace_count() {
805 let df = df![
806 "x" => ["2024-01", "2024-02", "2024-01", "2024-02"],
807 "y" => [1.0, 2.0, 3.0, 4.0],
808 "facet_col" => ["a", "a", "b", "b"]
809 ]
810 .unwrap();
811 let plot = TimeSeriesPlot::builder()
812 .data(&df)
813 .x("x")
814 .y("y")
815 .facet("facet_col")
816 .build();
817 assert_eq!(plot.ir_traces().len(), 2);
819 }
820
821 #[test]
822 fn test_resolve_color_both_none() {
823 let result = TimeSeriesPlot::resolve_color(0, None, None);
824 assert!(result.is_none());
825 }
826}