facett_plotters/

candlestick.rs

//! **Candlestick** (OHLC) chart. Each bar is an open/high/low/close tuple drawn
//! as a wick (high→low) + a body (open↔close), coloured up/down from theme
//! colours. The price→pixel mapping ([`OhlcLayout`]) is a pure function with
//! tests so the geometry is verifiable headlessly.

use egui::{Align2, FontId, Pos2, Rect, Sense, Stroke, Ui, vec2};
use facett_core::{Facet, FacetCaps, theme};

use crate::boxplot::paint_tooltip;

/// One OHLC bar. `label` is the x-axis tick (e.g. a date).
#[derive(Clone, Debug, PartialEq)]
pub struct Ohlc {
    pub label: String,
    pub open: f64,
    pub high: f64,
    pub low: f64,
    pub close: f64,
}

impl Ohlc {
    pub fn new(label: impl Into<String>, open: f64, high: f64, low: f64, close: f64) -> Self {
        Self { label: label.into(), open, high, low, close }
    }
    /// A candle is "up" when it closed at or above its open.
    pub fn bullish(&self) -> bool {
        self.close >= self.open
    }
}

/// The price→pixel mapping for a candlestick chart: where each bar's x-centre
/// sits, the body half-width, and a `y_of(price)` closure source. Pure data so
/// it can be asserted without a painter.
#[derive(Clone, Copy, Debug)]
pub struct OhlcLayout {
    plot: Rect,
    pub price_min: f64,
    pub price_max: f64,
    pub count: usize,
    pub slot: f32,
    pub body_half: f32,
}

impl OhlcLayout {
    /// Build the layout for `count` bars in `plot`, given the price range.
    pub fn new(plot: Rect, price_min: f64, price_max: f64, count: usize) -> Self {
        let slot = if count == 0 { plot.width() } else { plot.width() / count as f32 };
        let body_half = (slot * 0.3).clamp(1.0, 30.0);
        Self { plot, price_min, price_max, count, slot, body_half }
    }

    /// X-centre of bar `i`.
    pub fn x_center(&self, i: usize) -> f32 {
        self.plot.left() + self.slot * (i as f32 + 0.5)
    }

    /// Map a price to a y pixel (higher price → higher on screen).
    pub fn y_of(&self, price: f64) -> f32 {
        let span = (self.price_max - self.price_min).max(1e-9);
        self.plot.bottom() - ((price - self.price_min) / span) as f32 * self.plot.height()
    }

    /// Which bar index (if any) the x coordinate falls in.
    pub fn bar_at_x(&self, x: f32) -> Option<usize> {
        if x < self.plot.left() || x > self.plot.right() || self.count == 0 {
            return None;
        }
        let i = ((x - self.plot.left()) / self.slot).floor() as usize;
        if i < self.count { Some(i) } else { None }
    }
}

/// A candlestick chart over a series of OHLC bars.
pub struct CandlestickView {
    pub title: String,
    pub bars: Vec<Ohlc>,
}

impl CandlestickView {
    pub fn new(title: impl Into<String>, bars: Vec<Ohlc>) -> Self {
        Self { title: title.into(), bars }
    }

    /// `(min low, max high)` across all bars.
    pub fn price_range(&self) -> Option<(f64, f64)> {
        let mut r: Option<(f64, f64)> = None;
        for b in &self.bars {
            r = Some(match r {
                None => (b.low, b.high),
                Some((lo, hi)) => (lo.min(b.low), hi.max(b.high)),
            });
        }
        r
    }

    fn to_csv(&self) -> String {
        let mut s = String::from("label,open,high,low,close\n");
        for b in &self.bars {
            s.push_str(&format!("{},{},{},{},{}\n", b.label, b.open, b.high, b.low, b.close));
        }
        s
    }
}

impl Facet for CandlestickView {
    fn title(&self) -> &str {
        &self.title
    }
    fn caps(&self) -> FacetCaps {
        FacetCaps::NONE.themeable().copyable().resizable()
    }
    fn copy(&mut self) -> Option<String> {
        Some(self.to_csv())
    }
    fn ui(&mut self, ui: &mut Ui) {
        let th = theme(ui);
        let (rect, resp) = ui.allocate_exact_size(ui.available_size(), Sense::hover());
        let painter = ui.painter_at(rect);
        painter.rect_filled(rect, egui::CornerRadius::ZERO, th.bg);
        let Some((pmin, pmax)) = self.price_range() else {
            painter.text(rect.center(), Align2::CENTER_CENTER, "no data", FontId::proportional(13.0), th.text_dim);
            return;
        };
        let plot = rect.shrink2(vec2(36.0, 16.0));
        let lay = OhlcLayout::new(plot, pmin, pmax, self.bars.len());
        // Up = theme point/accent, down = a desaturated text_dim/red-ish; derive
        // both from the theme so no palette is hardcoded.
        let up = th.point;
        let down = crate::ramp::lerp_color(th.text_dim, th.accent, 0.0);

        for (i, b) in self.bars.iter().enumerate() {
            let cx = lay.x_center(i);
            let color = if b.bullish() { up } else { down };
            // Wick.
            painter.line_segment([Pos2::new(cx, lay.y_of(b.high)), Pos2::new(cx, lay.y_of(b.low))], Stroke::new(1.0, color));
            // Body.
            let (top, bot) = (b.open.max(b.close), b.open.min(b.close));
            let body = Rect::from_min_max(
                Pos2::new(cx - lay.body_half, lay.y_of(top)),
                Pos2::new(cx + lay.body_half, lay.y_of(bot)),
            );
            // A doji (open==close) collapses to a thin line — give it min height.
            let body = if body.height() < 1.0 {
                Rect::from_min_max(body.min, Pos2::new(body.right(), body.top() + 1.0))
            } else {
                body
            };
            if b.bullish() {
                painter.rect_filled(body, egui::CornerRadius::same(1), color.gamma_multiply(0.7));
            } else {
                painter.rect_filled(body, egui::CornerRadius::same(1), th.bg);
                painter.rect_stroke(body, egui::CornerRadius::same(1), Stroke::new(1.2, color), egui::StrokeKind::Middle);
            }
        }

        // Hover readout for the bar under the cursor.
        if let Some(hp) = resp.hover_pos() {
            if let Some(i) = lay.bar_at_x(hp.x) {
                let b = &self.bars[i];
                let cx = lay.x_center(i);
                painter.line_segment([Pos2::new(cx, plot.top()), Pos2::new(cx, plot.bottom())], Stroke::new(0.5, th.text_dim.gamma_multiply(0.5)));
                paint_tooltip(
                    &painter,
                    &th,
                    hp,
                    &format!("{}\nO {:.2}\nH {:.2}\nL {:.2}\nC {:.2}", b.label, b.open, b.high, b.low, b.close),
                );
            }
        }
    }
    fn state_json(&self) -> serde_json::Value {
        let (mn, mx) = self.price_range().unwrap_or((0.0, 0.0));
        let up = self.bars.iter().filter(|b| b.bullish()).count();
        serde_json::json!({
            "kind": "candlestick",
            "bars": self.bars.len(),
            "bullish": up,
            "bearish": self.bars.len() - up,
            "low": mn, "high": mx,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn plot() -> Rect {
        Rect::from_min_size(Pos2::new(0.0, 0.0), vec2(200.0, 100.0))
    }

    #[test]
    fn bullish_classification() {
        assert!(Ohlc::new("d", 10.0, 12.0, 9.0, 11.0).bullish());
        assert!(!Ohlc::new("d", 11.0, 12.0, 9.0, 10.0).bullish());
        // open == close is treated as up (non-decreasing)
        assert!(Ohlc::new("d", 10.0, 12.0, 9.0, 10.0).bullish());
    }

    #[test]
    fn y_mapping_is_inverted_and_spans_plot() {
        let lay = OhlcLayout::new(plot(), 0.0, 100.0, 4);
        // max price → top of plot, min price → bottom
        assert!((lay.y_of(100.0) - plot().top()).abs() < 1e-3);
        assert!((lay.y_of(0.0) - plot().bottom()).abs() < 1e-3);
        // midpoint maps to vertical centre
        assert!((lay.y_of(50.0) - plot().center().y).abs() < 1e-3);
    }

    #[test]
    fn x_centers_are_evenly_spaced() {
        let lay = OhlcLayout::new(plot(), 0.0, 1.0, 4);
        let c0 = lay.x_center(0);
        let c1 = lay.x_center(1);
        assert!((c1 - c0 - lay.slot).abs() < 1e-4);
        // first centre is half a slot in
        assert!((c0 - lay.slot * 0.5).abs() < 1e-4);
    }

    #[test]
    fn bar_at_x_round_trips_through_centers() {
        let lay = OhlcLayout::new(plot(), 0.0, 1.0, 5);
        for i in 0..5 {
            assert_eq!(lay.bar_at_x(lay.x_center(i)), Some(i));
        }
        assert_eq!(lay.bar_at_x(plot().left() - 5.0), None);
        assert_eq!(lay.bar_at_x(plot().right() + 5.0), None);
    }

    #[test]
    fn price_range_covers_low_and_high() {
        let v = CandlestickView::new(
            "t",
            vec![Ohlc::new("a", 5.0, 9.0, 3.0, 7.0), Ohlc::new("b", 7.0, 12.0, 6.0, 8.0)],
        );
        assert_eq!(v.price_range(), Some((3.0, 12.0)));
    }

    #[test]
    fn csv_header_and_rows() {
        let v = CandlestickView::new("t", vec![Ohlc::new("a", 1.0, 2.0, 0.5, 1.5)]);
        let csv = v.to_csv();
        assert!(csv.starts_with("label,open,high,low,close\n"));
        assert_eq!(csv.lines().count(), 2);
    }
}