plotlars_core/plots/
candlestick.rs

1use bon::bon;
2
3use polars::frame::DataFrame;
4
5use crate::{
6    components::{Axis, Direction, Text},
7    ir::data::ColumnData,
8    ir::layout::LayoutIR,
9    ir::trace::{CandlestickPlotIR, TraceIR},
10};
11
12/// A structure representing a Candlestick financial chart.
13///
14/// The `CandlestickPlot` struct facilitates the creation and customization of candlestick charts commonly used
15/// for visualizing financial data such as stock prices. It supports custom styling for increasing/decreasing
16/// values, whisker width configuration, 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| Yes       |
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/// * `increasing` - An optional reference to a `Direction` struct for customizing increasing candlesticks.
35/// * `decreasing` - An optional reference to a `Direction` struct for customizing decreasing candlesticks.
36/// * `whisker_width` - An optional `f64` specifying the width of the whiskers (0-1 range).
37/// * `plot_title` - An optional `Text` struct specifying the title of the plot.
38/// * `x_title` - An optional `Text` struct specifying the title of the x-axis.
39/// * `y_title` - An optional `Text` struct specifying the title of the y-axis.
40/// * `x_axis` - An optional reference to an `Axis` struct for customizing the x-axis.
41/// * `y_axis` - An optional reference to an `Axis` struct for customizing the y-axis.
42///
43/// # Examples
44///
45/// ```rust
46/// use plotlars::{Axis, CandlestickPlot, Direction, Plot, Rgb};
47/// use polars::prelude::*;
48///
49/// let stock_data = LazyCsvReader::new(PlRefPath::new("data/stock_prices.csv"))
50///     .finish()
51///     .unwrap()
52///     .collect()
53///     .unwrap();
54///
55/// let increasing = Direction::new()
56///     .line_color(Rgb(0, 200, 100))
57///     .line_width(0.5);
58///
59/// let decreasing = Direction::new()
60///     .line_color(Rgb(200, 50, 50))
61///     .line_width(0.5);
62///
63/// CandlestickPlot::builder()
64///     .data(&stock_data)
65///     .dates("date")
66///     .open("open")
67///     .high("high")
68///     .low("low")
69///     .close("close")
70///     .increasing(&increasing)
71///     .decreasing(&decreasing)
72///     .whisker_width(0.1)
73///     .plot_title("Candlestick Plot")
74///     .y_title("Price ($)")
75///     .y_axis(
76///         &Axis::new()
77///             .show_axis(true)
78///             .show_grid(true)
79///     )
80///     .build()
81///     .plot();
82/// ```
83///
84/// ![Example](https://imgur.com/fNDRLDX.png)
85#[derive(Clone)]
86#[allow(dead_code)]
87pub struct CandlestickPlot {
88    traces: Vec<TraceIR>,
89    layout: LayoutIR,
90}
91
92#[bon]
93impl CandlestickPlot {
94    #[builder(on(String, into), on(Text, into))]
95    pub fn new(
96        data: &DataFrame,
97        dates: &str,
98        open: &str,
99        high: &str,
100        low: &str,
101        close: &str,
102        increasing: Option<&Direction>,
103        decreasing: Option<&Direction>,
104        whisker_width: Option<f64>,
105        plot_title: Option<Text>,
106        x_title: Option<Text>,
107        y_title: Option<Text>,
108        x_axis: Option<&Axis>,
109        y_axis: Option<&Axis>,
110    ) -> Self {
111        // Build IR
112        let ir_trace = TraceIR::CandlestickPlot(CandlestickPlotIR {
113            dates: ColumnData::String(crate::data::get_string_column(data, dates)),
114            open: ColumnData::Numeric(crate::data::get_numeric_column(data, open)),
115            high: ColumnData::Numeric(crate::data::get_numeric_column(data, high)),
116            low: ColumnData::Numeric(crate::data::get_numeric_column(data, low)),
117            close: ColumnData::Numeric(crate::data::get_numeric_column(data, close)),
118            increasing: increasing.cloned(),
119            decreasing: decreasing.cloned(),
120            whisker_width,
121        });
122        let traces = vec![ir_trace];
123        let layout = LayoutIR {
124            title: plot_title.clone(),
125            x_title: x_title.clone(),
126            y_title: y_title.clone(),
127            y2_title: None,
128            z_title: None,
129            legend_title: None,
130            legend: None,
131            dimensions: None,
132            bar_mode: None,
133            box_mode: None,
134            box_gap: None,
135            margin_bottom: None,
136            axes_2d: Some(crate::ir::layout::Axes2dIR {
137                x_axis: x_axis.cloned(),
138                y_axis: y_axis.cloned(),
139                y2_axis: None,
140            }),
141            scene_3d: None,
142            polar: None,
143            mapbox: None,
144            grid: None,
145            annotations: vec![],
146        };
147        Self { traces, layout }
148    }
149}
150
151#[bon]
152impl CandlestickPlot {
153    #[builder(
154        start_fn = try_builder,
155        finish_fn = try_build,
156        builder_type = CandlestickPlotTryBuilder,
157        on(String, into),
158        on(Text, into),
159    )]
160    pub fn try_new(
161        data: &DataFrame,
162        dates: &str,
163        open: &str,
164        high: &str,
165        low: &str,
166        close: &str,
167        increasing: Option<&Direction>,
168        decreasing: Option<&Direction>,
169        whisker_width: Option<f64>,
170        plot_title: Option<Text>,
171        x_title: Option<Text>,
172        y_title: Option<Text>,
173        x_axis: Option<&Axis>,
174        y_axis: Option<&Axis>,
175    ) -> Result<Self, crate::io::PlotlarsError> {
176        std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
177            Self::__orig_new(
178                data,
179                dates,
180                open,
181                high,
182                low,
183                close,
184                increasing,
185                decreasing,
186                whisker_width,
187                plot_title,
188                x_title,
189                y_title,
190                x_axis,
191                y_axis,
192            )
193        }))
194        .map_err(|panic| {
195            let msg = panic
196                .downcast_ref::<String>()
197                .cloned()
198                .or_else(|| panic.downcast_ref::<&str>().map(|s| s.to_string()))
199                .unwrap_or_else(|| "unknown error".to_string());
200            crate::io::PlotlarsError::PlotBuild { message: msg }
201        })
202    }
203}
204
205impl crate::Plot for CandlestickPlot {
206    fn ir_traces(&self) -> &[TraceIR] {
207        &self.traces
208    }
209
210    fn ir_layout(&self) -> &LayoutIR {
211        &self.layout
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218    use crate::Plot;
219    use polars::prelude::*;
220
221    fn sample_df() -> DataFrame {
222        df![
223            "date" => ["2024-01-01", "2024-01-02", "2024-01-03"],
224            "open" => [100.0, 102.0, 101.0],
225            "high" => [105.0, 106.0, 104.0],
226            "low" => [98.0, 100.0, 99.0],
227            "close" => [103.0, 101.0, 102.0]
228        ]
229        .unwrap()
230    }
231
232    #[test]
233    fn test_basic_one_trace() {
234        let df = sample_df();
235        let plot = CandlestickPlot::builder()
236            .data(&df)
237            .dates("date")
238            .open("open")
239            .high("high")
240            .low("low")
241            .close("close")
242            .build();
243        assert_eq!(plot.ir_traces().len(), 1);
244    }
245
246    #[test]
247    fn test_trace_variant() {
248        let df = sample_df();
249        let plot = CandlestickPlot::builder()
250            .data(&df)
251            .dates("date")
252            .open("open")
253            .high("high")
254            .low("low")
255            .close("close")
256            .build();
257        assert!(matches!(plot.ir_traces()[0], TraceIR::CandlestickPlot(_)));
258    }
259
260    #[test]
261    fn test_layout_has_axes() {
262        let df = sample_df();
263        let plot = CandlestickPlot::builder()
264            .data(&df)
265            .dates("date")
266            .open("open")
267            .high("high")
268            .low("low")
269            .close("close")
270            .build();
271        assert!(plot.ir_layout().axes_2d.is_some());
272    }
273
274    #[test]
275    fn test_layout_title() {
276        let df = sample_df();
277        let plot = CandlestickPlot::builder()
278            .data(&df)
279            .dates("date")
280            .open("open")
281            .high("high")
282            .low("low")
283            .close("close")
284            .plot_title("Candlestick")
285            .build();
286        assert!(plot.ir_layout().title.is_some());
287    }
288}
plotlars_core/plots/candlestick.rs

plotlars_core/plots/
candlestick.rs
