plotlars/plots/
ohlc.rs

1use bon::bon;
2
3use plotly::{Layout as LayoutPlotly, Ohlc as OhlcPlotly, Trace};
4
5use polars::frame::DataFrame;
6use serde::Serialize;
7
8use crate::{
9    common::{Layout, PlotHelper, Polar},
10    components::{Axis, Text},
11};
12
13/// A structure representing an OHLC (Open-High-Low-Close) financial chart.
14///
15/// The `OhlcPlot` struct facilitates the creation and customization of OHLC charts commonly used
16/// for visualizing financial data such as stock prices. It supports multiple OHLC series, custom
17/// styling for increasing/decreasing values, hover information, and comprehensive layout customization
18/// including range selectors and sliders for interactive time navigation.
19///
20/// # Arguments
21///
22/// * `data` - A reference to the `DataFrame` containing the data to be plotted.
23/// * `dates` - A string slice specifying the column name for dates/timestamps.
24/// * `open` - A string slice specifying the column name for opening values.
25/// * `high` - A string slice specifying the column name for high values.
26/// * `low` - A string slice specifying the column name for low values.
27/// * `close` - A string slice specifying the column name for closing values.
28/// * `tick_width` - An optional `f64` specifying the width of the open/close ticks (0-1 range).
29/// * `plot_title` - An optional `Text` struct specifying the title of the plot.
30/// * `x_title` - An optional `Text` struct specifying the title of the x-axis.
31/// * `y_title` - An optional `Text` struct specifying the title of the y-axis.
32/// * `x_axis` - An optional reference to an `Axis` struct for customizing the x-axis.
33/// * `y_axis` - An optional reference to an `Axis` struct for customizing the y-axis.
34///
35/// # Examples
36///
37/// ```rust
38/// use plotlars::{Axis, OhlcPlot, Plot};
39/// use polars::prelude::*;
40///
41/// let stock_data = LazyCsvReader::new(PlPath::new("data/stock_prices.csv"))
42///     .finish()
43///     .unwrap()
44///     .collect()
45///     .unwrap();
46///
47/// OhlcPlot::builder()
48///     .data(&stock_data)
49///     .dates("date")
50///     .open("open")
51///     .high("high")
52///     .low("low")
53///     .close("close")
54///     .plot_title("OHLC Plot")
55///     .y_title("Price ($)")
56///     .y_axis(
57///         &Axis::new()
58///             .show_axis(true)
59///     )
60///     .build()
61///     .plot();
62/// ```
63/// ![Exmple](https://imgur.com/Sv8r9VN.png)
64#[derive(Clone, Serialize)]
65pub struct OhlcPlot {
66    traces: Vec<Box<dyn Trace + 'static>>,
67    layout: LayoutPlotly,
68}
69
70#[bon]
71impl OhlcPlot {
72    #[builder(on(String, into), on(Text, into))]
73    pub fn new(
74        data: &DataFrame,
75        dates: &str,
76        open: &str,
77        high: &str,
78        low: &str,
79        close: &str,
80        tick_width: Option<f64>,
81        plot_title: Option<Text>,
82        x_title: Option<Text>,
83        y_title: Option<Text>,
84        x_axis: Option<&Axis>,
85        y_axis: Option<&Axis>,
86    ) -> Self {
87        let z_title = None;
88        let y_title2 = None;
89        let z_axis = None;
90        let y_axis2 = None;
91        let legend_title = None;
92        let legend = None;
93
94        let layout = Self::create_layout(
95            plot_title,
96            x_title,
97            y_title,
98            y_title2,
99            z_title,
100            legend_title,
101            x_axis,
102            y_axis,
103            y_axis2,
104            z_axis,
105            legend,
106            None,
107        );
108
109        let traces = Self::create_traces(data, dates, open, high, low, close, tick_width);
110
111        Self { traces, layout }
112    }
113
114    fn create_traces(
115        data: &DataFrame,
116        dates_col: &str,
117        open_col: &str,
118        high_col: &str,
119        low_col: &str,
120        close_col: &str,
121        tick_width: Option<f64>,
122    ) -> Vec<Box<dyn Trace + 'static>> {
123        let mut traces: Vec<Box<dyn Trace + 'static>> = Vec::new();
124
125        let trace = Self::create_trace(
126            data, dates_col, open_col, high_col, low_col, close_col, tick_width,
127        );
128
129        traces.push(trace);
130        traces
131    }
132
133    fn create_trace(
134        data: &DataFrame,
135        dates_col: &str,
136        open_col: &str,
137        high_col: &str,
138        low_col: &str,
139        close_col: &str,
140        tick_width: Option<f64>,
141    ) -> Box<dyn Trace + 'static> {
142        let dates_data = Self::get_string_column(data, dates_col);
143        let open_data = Self::get_numeric_column(data, open_col);
144        let high_data = Self::get_numeric_column(data, high_col);
145        let low_data = Self::get_numeric_column(data, low_col);
146        let close_data = Self::get_numeric_column(data, close_col);
147
148        // Convert Option<f32> to f32 for OHLC trace
149        let open_values: Vec<f32> = open_data.into_iter().map(|v| v.unwrap_or(0.0)).collect();
150        let high_values: Vec<f32> = high_data.into_iter().map(|v| v.unwrap_or(0.0)).collect();
151        let low_values: Vec<f32> = low_data.into_iter().map(|v| v.unwrap_or(0.0)).collect();
152        let close_values: Vec<f32> = close_data.into_iter().map(|v| v.unwrap_or(0.0)).collect();
153        let dates_values: Vec<String> = dates_data
154            .into_iter()
155            .map(|v| v.unwrap_or_default())
156            .collect();
157
158        let mut trace = *OhlcPlotly::new(
159            dates_values,
160            open_values,
161            high_values,
162            low_values,
163            close_values,
164        );
165
166        // Set tick width
167        if let Some(tick_w) = tick_width {
168            trace = trace.tick_width(tick_w);
169        }
170
171        // Return trace as Box
172        Box::new(trace)
173    }
174}
175
176impl Layout for OhlcPlot {}
177impl Polar for OhlcPlot {}
178
179impl PlotHelper for OhlcPlot {
180    fn get_layout(&self) -> &LayoutPlotly {
181        &self.layout
182    }
183
184    fn get_traces(&self) -> &Vec<Box<dyn Trace + 'static>> {
185        &self.traces
186    }
187}