Skip to main content

ggplot_rs/stat/
heikin_ashi.rs

1//! `StatHeikinAshi`: Heikin-Ashi transform of OHLC rows, for
2//! `geom_heikin_ashi` (candles over the smoothed values).
3//!
4//! Rows are taken in `x` order; for each period
5//! `HA close = (o + h + l + c) / 4`, `HA open = (prev HA open + prev HA
6//! close) / 2` (the first period: `(o + c) / 2`), `HA high = max(h, HA open,
7//! HA close)` and `HA low = min(l, HA open, HA close)`. The `open`/`high`/
8//! `low`/`close` columns are replaced; every other column is carried.
9
10use crate::aes::Aesthetic;
11use crate::data::{DataFrame, Value};
12use crate::scale::ScaleSet;
13
14use super::Stat;
15
16/// Heikin-Ashi transform (see the module docs).
17#[derive(Clone, Copy, Debug, Default)]
18pub struct StatHeikinAshi;
19
20/// Heikin-Ashi `(open, high, low, close)` for consecutive OHLC periods.
21pub fn heikin_ashi(ohlc: &[(f64, f64, f64, f64)]) -> Vec<(f64, f64, f64, f64)> {
22    let mut out = Vec::with_capacity(ohlc.len());
23    let mut prev: Option<(f64, f64)> = None;
24    for &(o, h, l, c) in ohlc {
25        let close = (o + h + l + c) / 4.0;
26        let open = match prev {
27            Some((po, pc)) => (po + pc) / 2.0,
28            None => (o + c) / 2.0,
29        };
30        out.push((open, h.max(open).max(close), l.min(open).min(close), close));
31        prev = Some((open, close));
32    }
33    out
34}
35
36impl Stat for StatHeikinAshi {
37    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
38        let cols: Option<[&[Value]; 5]> = (|| {
39            Some([
40                data.column("x")?,
41                data.column("open")?,
42                data.column("high")?,
43                data.column("low")?,
44                data.column("close")?,
45            ])
46        })();
47        let Some([x, o, h, l, c]) = cols else {
48            return DataFrame::new();
49        };
50        let key = |v: &Value| v.as_f64().unwrap_or(f64::NAN);
51        let mut order: Vec<usize> = (0..data.nrows())
52            .filter(|&i| {
53                [&o[i], &h[i], &l[i], &c[i]]
54                    .iter()
55                    .all(|v| v.as_f64().is_some())
56            })
57            .collect();
58        order.sort_by(|&a, &b| key(&x[a]).total_cmp(&key(&x[b])));
59        let ha = heikin_ashi(
60            &order
61                .iter()
62                .map(|&i| (key(&o[i]), key(&h[i]), key(&l[i]), key(&c[i])))
63                .collect::<Vec<_>>(),
64        );
65        let mut out = DataFrame::new();
66        for name in data.column_names() {
67            let vals: Vec<Value> = match name {
68                "open" => ha.iter().map(|v| Value::Float(v.0)).collect(),
69                "high" => ha.iter().map(|v| Value::Float(v.1)).collect(),
70                "low" => ha.iter().map(|v| Value::Float(v.2)).collect(),
71                "close" => ha.iter().map(|v| Value::Float(v.3)).collect(),
72                _ => {
73                    let col = data.column(name).unwrap_or_default();
74                    order.iter().map(|&i| col[i].clone()).collect()
75                }
76            };
77            out.add_column(name.to_string(), vals);
78        }
79        out
80    }
81
82    fn required_aes(&self) -> Vec<Aesthetic> {
83        vec![
84            Aesthetic::X,
85            Aesthetic::Open,
86            Aesthetic::High,
87            Aesthetic::Low,
88            Aesthetic::Close,
89        ]
90    }
91
92    /// Each period depends on the previous one: the whole series at once.
93    fn panelwise(&self) -> bool {
94        true
95    }
96
97    fn name(&self) -> &str {
98        "heikin_ashi"
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[test]
107    fn matches_the_reference_formulas() {
108        let ha = heikin_ashi(&[(10.0, 12.0, 9.0, 11.0), (11.0, 13.0, 10.0, 12.0)]);
109        assert_eq!(ha[0], (10.5, 12.0, 9.0, 10.5));
110        // open = (10.5 + 10.5) / 2, close = (11 + 13 + 10 + 12) / 4 = 11.5.
111        assert_eq!(ha[1], (10.5, 13.0, 10.0, 11.5));
112    }
113}