plotlars/plots/
timeseriesplot.rs

1use bon::bon;
2
3use plotly::{
4    Layout as LayoutPlotly, Scatter, Trace,
5    common::{Line as LinePlotly, Marker as MarkerPlotly, Mode},
6};
7
8use polars::{
9    frame::DataFrame,
10    prelude::{IntoLazy, col},
11};
12use serde::Serialize;
13
14use crate::{
15    common::{Layout, Line, Marker, PlotHelper, Polar},
16    components::{Axis, Legend, Line as LineStyle, Rgb, Shape, Text},
17};
18
19/// A structure representing a time series plot.
20///
21/// The `TimeSeriesPlot` struct facilitates the creation and customization of time series plots with various options
22/// for data selection, grouping, layout configuration, and aesthetic adjustments. It supports the addition of multiple
23/// series, customization of marker shapes, colors, sizes, opacity settings, and comprehensive layout customization
24/// including titles, axes, and legends.
25///
26/// # Arguments
27///
28/// * `data` - A reference to the `DataFrame` containing the data to be plotted.
29/// * `x` - A string slice specifying the column name to be used for the x-axis, typically representing time or dates.
30/// * `y` - A string slice specifying the column name to be used for the y-axis, typically representing the primary metric.
31/// * `additional_series` - An optional vector of string slices specifying additional y-axis columns to be plotted as series.
32/// * `size` - An optional `usize` specifying the size of the markers or line thickness.
33/// * `color` - An optional `Rgb` value specifying the color of the markers. This is used when `group` is not specified.
34/// * `colors` - An optional vector of `Rgb` values specifying the colors for the markers. This is used when `group` is specified to differentiate between groups.
35/// * `with_shape` - An optional `bool` indicating whether to use shapes for markers in the plot.
36/// * `shape` - An optional `Shape` specifying the shape of the markers. This is used when `group` is not specified.
37/// * `shapes` - An optional vector of `Shape` values specifying multiple shapes for the markers when plotting multiple groups.
38/// * `width` - An optional `f64` specifying the width of the plotted lines.
39/// * `line` - An optional `LineStyle` specifying the style of the line. This is used when `additional_series` is not specified.
40/// * `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.
41/// * `plot_title` - An optional `Text` struct specifying the title of the plot.
42/// * `x_title` - An optional `Text` struct specifying the title of the x-axis.
43/// * `y_title` - An optional `Text` struct specifying the title of the y-axis.
44/// * `legend_title` - An optional `Text` struct specifying the title of the legend.
45/// * `x_axis` - An optional reference to an `Axis` struct for customizing the x-axis.
46/// * `y_axis` - An optional reference to an `Axis` struct for customizing the y-axis.
47/// * `y_axis2` - An optional reference to an `Axis` struct for customizing the y-axis2.
48/// * `legend` - An optional reference to a `Legend` struct for customizing the legend of the plot (e.g., positioning, font, etc.).
49///
50/// # Examples
51///
52/// ```rust
53/// use polars::prelude::*;
54/// use plotlars::{Axis, Legend, Line, Plot, Rgb, Shape, Text, TimeSeriesPlot};
55///
56/// let dataset = LazyCsvReader::new(PlPath::new("data/revenue_and_cost.csv"))
57///     .finish()
58///     .unwrap()
59///     .select([
60///         col("Date").cast(DataType::String),
61///         col("Revenue").cast(DataType::Int32),
62///         col("Cost").cast(DataType::Int32),
63///     ])
64///     .collect()
65///     .unwrap();
66///
67/// TimeSeriesPlot::builder()
68///     .data(&dataset)
69///     .x("Date")
70///     .y("Revenue")
71///     .additional_series(vec!["Cost"])
72///     .size(8)
73///     .colors(vec![
74///         Rgb(0, 0, 255),
75///         Rgb(255, 0, 0),
76///     ])
77///     .lines(vec![Line::Dash, Line::Solid])
78///     .with_shape(true)
79///     .shapes(vec![Shape::Circle, Shape::Square])
80///     .plot_title(
81///         Text::from("Time Series Plot")
82///             .font("Arial")
83///             .size(18)
84///     )
85///     .legend(
86///         &Legend::new()
87///             .x(0.05)
88///             .y(0.9)
89///     )
90///     .x_title("x")
91///     .y_title(
92///         Text::from("y")
93///             .color(Rgb(0, 0, 255))
94///     )
95///     .y_title2(
96///         Text::from("y2")
97///             .color(Rgb(255, 0, 0))
98///     )
99///     .y_axis(
100///         &Axis::new()
101///             .value_color(Rgb(0, 0, 255))
102///             .show_grid(false)
103///             .zero_line_color(Rgb(0, 0, 0))
104///     )
105///     .y_axis2(
106///         &Axis::new()
107///             .axis_side(plotlars::AxisSide::Right)
108///             .value_color(Rgb(255, 0, 0))
109///             .show_grid(false)
110///     )
111///     .build()
112///     .plot();
113/// ```
114///
115/// ![Example1](https://imgur.com/hL27Xcn.png)
116///
117/// ```rust
118/// use polars::prelude::*;
119/// use plotlars::{Plot, TimeSeriesPlot, Rgb, Line};
120///
121/// let dataset = LazyCsvReader::new(PlPath::new("data/debilt_2023_temps.csv"))
122///     .with_has_header(true)
123///     .with_try_parse_dates(true)
124///     .finish()
125///     .unwrap()
126///     .with_columns(vec![
127///         (col("tavg") / lit(10)).alias("tavg"),
128///         (col("tmin") / lit(10)).alias("tmin"),
129///         (col("tmax") / lit(10)).alias("tmax"),
130///     ])
131///     .collect()
132///     .unwrap();
133///
134///     TimeSeriesPlot::builder()
135///     .data(&dataset)
136///     .x("date")
137///     .y("tavg")
138///     .additional_series(vec!["tmin", "tmax"])
139///     .colors(vec![
140///         Rgb(128, 128, 128),
141///         Rgb(0, 122, 255),
142///         Rgb(255, 128, 0),
143///     ])
144///     .lines(vec![
145///         Line::Solid,
146///         Line::Dot,
147///         Line::Dot,
148///     ])
149///     .plot_title("Temperature at De Bilt (2023)")
150///     .legend_title("Legend")
151///     .build()
152///     .plot();
153/// ```
154///
155/// ![Example2](https://imgur.com/NBioox6.png)
156#[derive(Clone, Serialize)]
157pub struct TimeSeriesPlot {
158    traces: Vec<Box<dyn Trace + 'static>>,
159    layout: LayoutPlotly,
160}
161
162#[bon]
163impl TimeSeriesPlot {
164    #[builder(on(String, into), on(Text, into))]
165    pub fn new(
166        data: &DataFrame,
167        x: &str,
168        y: &str,
169        additional_series: Option<Vec<&str>>,
170        size: Option<usize>,
171        color: Option<Rgb>,
172        colors: Option<Vec<Rgb>>,
173        with_shape: Option<bool>,
174        shape: Option<Shape>,
175        shapes: Option<Vec<Shape>>,
176        width: Option<f64>,
177        line: Option<LineStyle>,
178        lines: Option<Vec<LineStyle>>,
179        plot_title: Option<Text>,
180        x_title: Option<Text>,
181        y_title: Option<Text>,
182        y_title2: Option<Text>,
183        legend_title: Option<Text>,
184        x_axis: Option<&Axis>,
185        y_axis: Option<&Axis>,
186        y_axis2: Option<&Axis>,
187        legend: Option<&Legend>,
188    ) -> Self {
189        let z_title = None;
190        let z_axis = None;
191        let mut has_y_axis2 = false;
192
193        if y_axis2.is_some() {
194            has_y_axis2 = true;
195        }
196
197        let layout = Self::create_layout(
198            plot_title,
199            x_title,
200            y_title,
201            y_title2,
202            z_title,
203            legend_title,
204            x_axis,
205            y_axis,
206            y_axis2,
207            z_axis,
208            legend,
209        );
210
211        let traces = Self::create_traces(
212            data,
213            x,
214            y,
215            additional_series,
216            has_y_axis2,
217            size,
218            color,
219            colors,
220            with_shape,
221            shape,
222            shapes,
223            width,
224            line,
225            lines,
226        );
227
228        Self { traces, layout }
229    }
230
231    #[allow(clippy::too_many_arguments)]
232    fn create_traces(
233        data: &DataFrame,
234        x_col: &str,
235        y_col: &str,
236        additional_series: Option<Vec<&str>>,
237        has_y_axis2: bool,
238        size: Option<usize>,
239        color: Option<Rgb>,
240        colors: Option<Vec<Rgb>>,
241        with_shape: Option<bool>,
242        shape: Option<Shape>,
243        shapes: Option<Vec<Shape>>,
244        width: Option<f64>,
245        style: Option<LineStyle>,
246        styles: Option<Vec<LineStyle>>,
247    ) -> Vec<Box<dyn Trace + 'static>> {
248        let mut traces: Vec<Box<dyn Trace + 'static>> = Vec::new();
249
250        let opacity = None;
251
252        let marker = Self::create_marker(
253            0,
254            opacity,
255            size,
256            color,
257            colors.clone(),
258            shape,
259            shapes.clone(),
260        );
261
262        let line = Self::create_line(0, width, style, styles.clone());
263
264        let name = Some(y_col);
265        let mut y_axis_index = "";
266
267        let trace = Self::create_trace(
268            data,
269            x_col,
270            y_col,
271            name,
272            with_shape,
273            marker,
274            line,
275            y_axis_index,
276        );
277
278        traces.push(trace);
279
280        if let Some(additional_series) = additional_series {
281            let additional_series = additional_series.into_iter();
282
283            for (i, series) in additional_series.enumerate() {
284                let marker = Self::create_marker(
285                    i + 1,
286                    opacity,
287                    size,
288                    color,
289                    colors.clone(),
290                    shape,
291                    shapes.clone(),
292                );
293
294                let line = Self::create_line(i + 1, width, style, styles.clone());
295
296                let subset = data
297                    .clone()
298                    .lazy()
299                    .select([col(x_col), col(series)])
300                    .collect()
301                    .unwrap();
302
303                let name = Some(series);
304
305                if has_y_axis2 {
306                    y_axis_index = "y2";
307                }
308
309                let trace = Self::create_trace(
310                    &subset,
311                    x_col,
312                    series,
313                    name,
314                    with_shape,
315                    marker,
316                    line,
317                    y_axis_index,
318                );
319
320                traces.push(trace);
321            }
322        }
323
324        traces
325    }
326
327    #[allow(clippy::too_many_arguments)]
328    fn create_trace(
329        data: &DataFrame,
330        x_col: &str,
331        y_col: &str,
332        name: Option<&str>,
333        with_shape: Option<bool>,
334        marker: MarkerPlotly,
335        line: LinePlotly,
336        index: &str,
337    ) -> Box<dyn Trace + 'static> {
338        let x_data = Self::get_string_column(data, x_col);
339        let y_data = Self::get_numeric_column(data, y_col);
340
341        let mut trace = Scatter::default().x(x_data).y(y_data);
342
343        if let Some(with_shape) = with_shape {
344            if with_shape {
345                trace = trace.mode(Mode::LinesMarkers);
346            } else {
347                trace = trace.mode(Mode::Lines);
348            }
349        }
350
351        trace = trace.marker(marker);
352        trace = trace.line(line);
353
354        if let Some(name) = name {
355            trace = trace.name(name);
356        }
357
358        trace.y_axis(index)
359    }
360}
361
362impl Layout for TimeSeriesPlot {}
363impl Line for TimeSeriesPlot {}
364impl Marker for TimeSeriesPlot {}
365impl Polar for TimeSeriesPlot {}
366
367impl PlotHelper for TimeSeriesPlot {
368    fn get_layout(&self) -> &LayoutPlotly {
369        &self.layout
370    }
371
372    fn get_traces(&self) -> &Vec<Box<dyn Trace + 'static>> {
373        &self.traces
374    }
375}