Skip to main content

plotlars_core/plots/
timeseriesplot.rs

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/// A structure representing a time series plot.
18///
19/// The `TimeSeriesPlot` struct facilitates the creation and customization of time series plots with various options
20/// for data selection, grouping, layout configuration, and aesthetic adjustments. It supports the addition of multiple
21/// series, customization of marker shapes, colors, sizes, opacity settings, and comprehensive layout customization
22/// including titles, axes, and legends.
23///
24/// # Backend Support
25///
26/// | Backend | Supported |
27/// |---------|-----------|
28/// | Plotly  | Yes       |
29/// | Plotters| Yes       |
30///
31/// # Arguments
32///
33/// * `data` - A reference to the `DataFrame` containing the data to be plotted.
34/// * `x` - A string slice specifying the column name to be used for the x-axis, typically representing time or dates.
35/// * `y` - A string slice specifying the column name to be used for the y-axis, typically representing the primary metric.
36/// * `additional_series` - An optional vector of string slices specifying additional y-axis columns to be plotted as series.
37/// * `facet` - An optional string slice specifying the column name to be used for faceting (creating multiple subplots).
38/// * `facet_config` - An optional reference to a `FacetConfig` struct for customizing facet behavior (grid dimensions, scales, gaps, etc.).
39/// * `size` - An optional `usize` specifying the size of the markers or line thickness.
40/// * `color` - An optional `Rgb` value specifying the color of the markers. This is used when `group` is not specified.
41/// * `colors` - An optional vector of `Rgb` values specifying the colors for the markers. This is used when `group` is specified to differentiate between groups.
42/// * `shape` - An optional `Shape` specifying the shape of the markers. When set, markers are automatically displayed on the plot. This is used when `group` is not specified.
43/// * `shapes` - An optional vector of `Shape` values specifying multiple shapes for the markers when plotting multiple groups. When set, markers are automatically displayed on the plot.
44/// * `width` - An optional `f64` specifying the width of the plotted lines.
45/// * `line` - An optional `LineStyle` specifying the style of the line. This is used when `additional_series` is not specified.
46/// * `lines` - An optional vector of `LineStyle` enums specifying the styles of lines for each plotted series. This is used when `additional_series` is specified to differentiate between multiple series.
47/// * `plot_title` - An optional `Text` struct specifying the title of the plot.
48/// * `x_title` - An optional `Text` struct specifying the title of the x-axis.
49/// * `y_title` - An optional `Text` struct specifying the title of the y-axis.
50/// * `y2_title` - An optional `Text` struct specifying the title of the secondary y-axis.
51/// * `legend_title` - An optional `Text` struct specifying the title of the legend.
52/// * `x_axis` - An optional reference to an `Axis` struct for customizing the x-axis.
53/// * `y_axis` - An optional reference to an `Axis` struct for customizing the y-axis.
54/// * `y2_axis` - An optional reference to an `Axis` struct for customizing the secondary y-axis.
55/// * `legend` - An optional reference to a `Legend` struct for customizing the legend of the plot (e.g., positioning, font, etc.).
56///
57/// # Examples
58///
59/// ```rust
60/// use polars::prelude::*;
61/// use plotlars::{Axis, Legend, Line, Plot, Rgb, Shape, Text, TimeSeriesPlot};
62///
63/// let dataset = LazyCsvReader::new(PlRefPath::new("data/revenue_and_cost.csv"))
64///     .finish()
65///     .unwrap()
66///     .select([
67///         col("Date").cast(DataType::String),
68///         col("Revenue").cast(DataType::Int32),
69///         col("Cost").cast(DataType::Int32),
70///     ])
71///     .collect()
72///     .unwrap();
73///
74/// TimeSeriesPlot::builder()
75///     .data(&dataset)
76///     .x("Date")
77///     .y("Revenue")
78///     .additional_series(vec!["Cost"])
79///     .size(8)
80///     .colors(vec![
81///         Rgb(0, 0, 255),
82///         Rgb(255, 0, 0),
83///     ])
84///     .lines(vec![
85///         Line::Dash,
86///         Line::Solid,
87///     ])
88///     .shapes(vec![
89///         Shape::Circle,
90///         Shape::Square,
91///     ])
92///     .plot_title(
93///         Text::from("Time Series Plot")
94///             .font("Arial")
95///             .size(18)
96///     )
97///     .legend(
98///         &Legend::new()
99///             .x(0.05)
100///             .y(0.9)
101///     )
102///     .x_title("x")
103///     .y_title(
104///         Text::from("y")
105///             .color(Rgb(0, 0, 255))
106///     )
107///     .y2_title(
108///         Text::from("y2")
109///             .color(Rgb(255, 0, 0))
110///     )
111///     .y_axis(
112///         &Axis::new()
113///             .value_color(Rgb(0, 0, 255))
114///             .show_grid(false)
115///             .zero_line_color(Rgb(0, 0, 0))
116///     )
117///     .y2_axis(
118///         &Axis::new()
119///             .axis_side(plotlars::AxisSide::Right)
120///             .value_color(Rgb(255, 0, 0))
121///             .show_grid(false)
122///     )
123///     .build()
124///     .plot();
125/// ```
126///
127/// ![Example1](https://imgur.com/hL27Xcn.png)
128///
129/// ```rust
130/// use polars::prelude::*;
131/// use plotlars::{Plot, TimeSeriesPlot, Rgb, Line};
132///
133/// let dataset = LazyCsvReader::new(PlRefPath::new("data/debilt_2023_temps.csv"))
134///     .with_has_header(true)
135///     .with_try_parse_dates(true)
136///     .finish()
137///     .unwrap()
138///     .with_columns(vec![
139///         (col("tavg") / lit(10)).alias("tavg"),
140///         (col("tmin") / lit(10)).alias("tmin"),
141///         (col("tmax") / lit(10)).alias("tmax"),
142///     ])
143///     .collect()
144///     .unwrap();
145///
146///     TimeSeriesPlot::builder()
147///     .data(&dataset)
148///     .x("date")
149///     .y("tavg")
150///     .additional_series(vec!["tmin", "tmax"])
151///     .colors(vec![
152///         Rgb(128, 128, 128),
153///         Rgb(0, 122, 255),
154///         Rgb(255, 128, 0),
155///     ])
156///     .lines(vec![
157///         Line::Solid,
158///         Line::Dot,
159///         Line::Dot,
160///     ])
161///     .plot_title("Temperature at De Bilt (2023)")
162///     .legend_title("Legend")
163///     .build()
164///     .plot();
165/// ```
166///
167/// ![Example2](https://imgur.com/NBioox6.png)
168#[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        // 2 facets, 1 series each = 2 traces
818        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}