updates to the indicators and fix minor bugs

1 files changed, 39 insertions, 0 deletions

diff --git a/src/indicators.rs b/src/indicators.rs
index 4378757..ec52f2a 100644
--- a/src/indicators.rs
+++ b/src/indicators.rs
@@ -27,6 +27,39 @@ pub fn rsi(candles: &[Candle], period: usize) -> Vec<Option<f64>> {
     out
 }
 
+/// Average true range (Wilder). Before `period` bars it is the mean true range so far,
+/// so every bar has a value.
+pub fn atr(candles: &[Candle], period: usize) -> Vec<f64> {
+    let n = period as f64;
+    let mut out = Vec::with_capacity(candles.len());
+    let mut avg = 0.0;
+    for (i, c) in candles.iter().enumerate() {
+        let tr = if i == 0 {
+            c.high - c.low
+        } else {
+            let pc = candles[i - 1].close;
+            (c.high - c.low).max((c.high - pc).abs()).max((c.low - pc).abs())
+        };
+        avg = if i < period { (avg * i as f64 + tr) / (i as f64 + 1.0) } else { (avg * (n - 1.0) + tr) / n };
+        out.push(avg);
+    }
+    out
+}
+
+/// Causal exponential moving average of closes, returned as flat candles
+/// (used only to locate candidate swings).
+pub fn ema_candles(candles: &[Candle], n: usize) -> Vec<Candle> {
+    let alpha = 2.0 / (n as f64 + 1.0);
+    let mut e = candles[0].close;
+    candles
+        .iter()
+        .map(|c| {
+            e = alpha * c.close + (1.0 - alpha) * e;
+            Candle { open: e, high: e, low: e, close: e, ..c.clone() }
+        })
+        .collect()
+}
+
 /// KD: RSV from `ta`'s FastStochastic over high/low/close, then
 /// K = 2/3 K_prev + 1/3 RSV, D = 2/3 D_prev + 1/3 K, both seeded at 50.
 pub fn kd(candles: &[Candle], period: usize) -> (Vec<Option<f64>>, Vec<Option<f64>>) {
@@ -88,6 +121,12 @@ mod tests {
     }
 
     #[test]
+    fn atr_constant_range() {
+        let cs = [candle(12.0, 10.0, 11.0), candle(12.0, 10.0, 11.0), candle(12.0, 10.0, 11.0)];
+        assert!(atr(&cs, 14).iter().all(|&a| (a - 2.0).abs() < 1e-9));
+    }
+
+    #[test]
     fn rsi_all_gains_is_100() {
         let cs: Vec<_> = (1..=20).map(|i| candle(i as f64 + 1.0, i as f64, i as f64)).collect();
         let r = rsi(&cs, 14);