ggplot_rs/stat/
heikin_ashi.rs1use crate::aes::Aesthetic;
11use crate::data::{DataFrame, Value};
12use crate::scale::ScaleSet;
13
14use super::Stat;
15
16#[derive(Clone, Copy, Debug, Default)]
18pub struct StatHeikinAshi;
19
20pub 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 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 assert_eq!(ha[1], (10.5, 13.0, 10.0, 11.5));
112 }
113}