plotlars_core/plots/
ohlc.rs

1use bon::bon;
2
3use polars::frame::DataFrame;
4
5use crate::{
6    components::{Axis, Text},
7    ir::data::ColumnData,
8    ir::layout::LayoutIR,
9    ir::trace::{OhlcPlotIR, TraceIR},
10};
11
12/// A structure representing an OHLC (Open-High-Low-Close) financial chart.
13///
14/// The `OhlcPlot` struct facilitates the creation and customization of OHLC charts commonly used
15/// for visualizing financial data such as stock prices. It supports multiple OHLC series, custom
16/// styling for increasing/decreasing values, hover information, and comprehensive layout customization
17/// including range selectors and sliders for interactive time navigation.
18///
19/// # Backend Support
20///
21/// | Backend | Supported |
22/// |---------|-----------|
23/// | Plotly  | Yes       |
24/// | Plotters| --        |
25///
26/// # Arguments
27///
28/// * `data` - A reference to the `DataFrame` containing the data to be plotted.
29/// * `dates` - A string slice specifying the column name for dates/timestamps.
30/// * `open` - A string slice specifying the column name for opening values.
31/// * `high` - A string slice specifying the column name for high values.
32/// * `low` - A string slice specifying the column name for low values.
33/// * `close` - A string slice specifying the column name for closing values.
34/// * `tick_width` - An optional `f64` specifying the width of the open/close ticks (0-1 range).
35/// * `plot_title` - An optional `Text` struct specifying the title of the plot.
36/// * `x_title` - An optional `Text` struct specifying the title of the x-axis.
37/// * `y_title` - An optional `Text` struct specifying the title of the y-axis.
38/// * `x_axis` - An optional reference to an `Axis` struct for customizing the x-axis.
39/// * `y_axis` - An optional reference to an `Axis` struct for customizing the y-axis.
40///
41/// # Examples
42///
43/// ```rust
44/// use plotlars::{Axis, OhlcPlot, Plot};
45/// use polars::prelude::*;
46///
47/// let stock_data = LazyCsvReader::new(PlRefPath::new("data/stock_prices.csv"))
48///     .finish()
49///     .unwrap()
50///     .collect()
51///     .unwrap();
52///
53/// OhlcPlot::builder()
54///     .data(&stock_data)
55///     .dates("date")
56///     .open("open")
57///     .high("high")
58///     .low("low")
59///     .close("close")
60///     .plot_title("OHLC Plot")
61///     .y_title("Price ($)")
62///     .y_axis(
63///         &Axis::new()
64///             .show_axis(true)
65///     )
66///     .build()
67///     .plot();
68/// ```
69/// ![Exmple](https://imgur.com/Sv8r9VN.png)
70#[derive(Clone)]
71#[allow(dead_code)]
72pub struct OhlcPlot {
73    traces: Vec<TraceIR>,
74    layout: LayoutIR,
75}
76
77#[bon]
78impl OhlcPlot {
79    #[builder(on(String, into), on(Text, into))]
80    pub fn new(
81        data: &DataFrame,
82        dates: &str,
83        open: &str,
84        high: &str,
85        low: &str,
86        close: &str,
87        tick_width: Option<f64>,
88        plot_title: Option<Text>,
89        x_title: Option<Text>,
90        y_title: Option<Text>,
91        x_axis: Option<&Axis>,
92        y_axis: Option<&Axis>,
93    ) -> Self {
94        // Build IR
95        let ir_trace = TraceIR::OhlcPlot(OhlcPlotIR {
96            dates: ColumnData::String(crate::data::get_string_column(data, dates)),
97            open: ColumnData::Numeric(crate::data::get_numeric_column(data, open)),
98            high: ColumnData::Numeric(crate::data::get_numeric_column(data, high)),
99            low: ColumnData::Numeric(crate::data::get_numeric_column(data, low)),
100            close: ColumnData::Numeric(crate::data::get_numeric_column(data, close)),
101            tick_width,
102        });
103        let traces = vec![ir_trace];
104        let layout = LayoutIR {
105            title: plot_title.clone(),
106            x_title: x_title.clone(),
107            y_title: y_title.clone(),
108            y2_title: None,
109            z_title: None,
110            legend_title: None,
111            legend: None,
112            dimensions: None,
113            bar_mode: None,
114            box_mode: None,
115            box_gap: None,
116            margin_bottom: None,
117            axes_2d: Some(crate::ir::layout::Axes2dIR {
118                x_axis: x_axis.cloned(),
119                y_axis: y_axis.cloned(),
120                y2_axis: None,
121            }),
122            scene_3d: None,
123            polar: None,
124            mapbox: None,
125            grid: None,
126            annotations: vec![],
127        };
128        Self { traces, layout }
129    }
130}
131
132#[bon]
133impl OhlcPlot {
134    #[builder(
135        start_fn = try_builder,
136        finish_fn = try_build,
137        builder_type = OhlcPlotTryBuilder,
138        on(String, into),
139        on(Text, into),
140    )]
141    pub fn try_new(
142        data: &DataFrame,
143        dates: &str,
144        open: &str,
145        high: &str,
146        low: &str,
147        close: &str,
148        tick_width: Option<f64>,
149        plot_title: Option<Text>,
150        x_title: Option<Text>,
151        y_title: Option<Text>,
152        x_axis: Option<&Axis>,
153        y_axis: Option<&Axis>,
154    ) -> Result<Self, crate::io::PlotlarsError> {
155        std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
156            Self::__orig_new(
157                data, dates, open, high, low, close, tick_width, plot_title, x_title, y_title,
158                x_axis, y_axis,
159            )
160        }))
161        .map_err(|panic| {
162            let msg = panic
163                .downcast_ref::<String>()
164                .cloned()
165                .or_else(|| panic.downcast_ref::<&str>().map(|s| s.to_string()))
166                .unwrap_or_else(|| "unknown error".to_string());
167            crate::io::PlotlarsError::PlotBuild { message: msg }
168        })
169    }
170}
171
172impl crate::Plot for OhlcPlot {
173    fn ir_traces(&self) -> &[TraceIR] {
174        &self.traces
175    }
176
177    fn ir_layout(&self) -> &LayoutIR {
178        &self.layout
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185    use crate::Plot;
186    use polars::prelude::*;
187
188    fn sample_df() -> DataFrame {
189        df![
190            "date" => ["2024-01-01", "2024-01-02", "2024-01-03"],
191            "open" => [100.0, 102.0, 101.0],
192            "high" => [105.0, 106.0, 104.0],
193            "low" => [98.0, 100.0, 99.0],
194            "close" => [103.0, 101.0, 102.0]
195        ]
196        .unwrap()
197    }
198
199    #[test]
200    fn test_basic_one_trace() {
201        let df = sample_df();
202        let plot = OhlcPlot::builder()
203            .data(&df)
204            .dates("date")
205            .open("open")
206            .high("high")
207            .low("low")
208            .close("close")
209            .build();
210        assert_eq!(plot.ir_traces().len(), 1);
211    }
212
213    #[test]
214    fn test_trace_variant() {
215        let df = sample_df();
216        let plot = OhlcPlot::builder()
217            .data(&df)
218            .dates("date")
219            .open("open")
220            .high("high")
221            .low("low")
222            .close("close")
223            .build();
224        assert!(matches!(plot.ir_traces()[0], TraceIR::OhlcPlot(_)));
225    }
226
227    #[test]
228    fn test_layout_has_axes() {
229        let df = sample_df();
230        let plot = OhlcPlot::builder()
231            .data(&df)
232            .dates("date")
233            .open("open")
234            .high("high")
235            .low("low")
236            .close("close")
237            .build();
238        assert!(plot.ir_layout().axes_2d.is_some());
239    }
240}
plotlars_core/plots/ohlc.rs

plotlars_core/plots/
ohlc.rs
