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/// 
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/// 
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}