Skip to main content

indicators/momentum/
stochastic.rs

1//! Stochastic Oscillator (%K and %D).
2//!
3//! Python source: `indicators/momentum/stochastic.py :: class Stochastic`
4//!
5//! # Algorithm
6//!
7//! 1. **Raw %K**:
8//!    `%K[i] = 100 * (close[i] - lowest_low) / (highest_high - lowest_low)`
9//!    where the window is `k_period` bars ending at `i`.
10//!    Yields `NaN` when `highest_high == lowest_low`.
11//!
12//! 2. **Smooth %K** (optional): SMA of raw %K over `smooth_k` bars.
13//!    `smooth_k = 1` means no smoothing (fast stochastic).
14//!    `smooth_k = 3` is the standard slow stochastic.
15//!
16//! 3. **%D**: SMA of smooth %K over `d_period` bars.
17//!
18//! Output columns: `"Stoch_K"`, `"Stoch_D"`.
19
20use std::collections::HashMap;
21
22use crate::error::IndicatorError;
23use crate::indicator::{Indicator, IndicatorOutput};
24use crate::registry::param_usize;
25use crate::types::Candle;
26
27// ── Params ────────────────────────────────────────────────────────────────────
28
29#[derive(Debug, Clone)]
30pub struct StochParams {
31    /// Look-back window for highest-high / lowest-low. Default: 14.
32    pub k_period: usize,
33    /// Smoothing of raw %K. 1 = no smoothing. Default: 3.
34    pub smooth_k: usize,
35    /// SMA period for %D. Default: 3.
36    pub d_period: usize,
37}
38
39impl Default for StochParams {
40    fn default() -> Self {
41        Self {
42            k_period: 14,
43            smooth_k: 3,
44            d_period: 3,
45        }
46    }
47}
48
49// ── Indicator struct ──────────────────────────────────────────────────────────
50
51#[derive(Debug, Clone)]
52pub struct Stochastic {
53    pub params: StochParams,
54}
55
56impl Stochastic {
57    pub fn new(params: StochParams) -> Self {
58        Self { params }
59    }
60}
61
62impl Default for Stochastic {
63    fn default() -> Self {
64        Self::new(StochParams::default())
65    }
66}
67
68impl Indicator for Stochastic {
69    fn name(&self) -> &'static str {
70        "Stochastic"
71    }
72
73    fn required_len(&self) -> usize {
74        self.params.k_period + self.params.smooth_k + self.params.d_period - 2
75    }
76
77    fn required_columns(&self) -> &[&'static str] {
78        &["high", "low", "close"]
79    }
80
81    fn calculate(&self, candles: &[Candle]) -> Result<IndicatorOutput, IndicatorError> {
82        self.check_len(candles)?;
83
84        let n = candles.len();
85        let kp = self.params.k_period;
86        let sk = self.params.smooth_k;
87        let dp = self.params.d_period;
88
89        // ── Step 1: raw %K ────────────────────────────────────────────────────
90        let mut raw_k = vec![f64::NAN; n];
91        for i in (kp - 1)..n {
92            let window = &candles[(i + 1 - kp)..=i];
93            let hh = window
94                .iter()
95                .map(|c| c.high)
96                .fold(f64::NEG_INFINITY, f64::max);
97            let ll = window.iter().map(|c| c.low).fold(f64::INFINITY, f64::min);
98            let range = hh - ll;
99            raw_k[i] = if range == 0.0 {
100                f64::NAN
101            } else {
102                100.0 * (candles[i].close - ll) / range
103            };
104        }
105
106        // ── Step 2: smooth %K (SMA) ───────────────────────────────────────────
107        let smooth_k = if sk <= 1 {
108            raw_k.clone()
109        } else {
110            sma_of(&raw_k, sk)
111        };
112
113        // ── Step 3: %D (SMA of smooth_k) ─────────────────────────────────────
114        let d = sma_of(&smooth_k, dp);
115
116        Ok(IndicatorOutput::from_pairs([
117            ("Stoch_K".to_string(), smooth_k),
118            ("Stoch_D".to_string(), d),
119        ]))
120    }
121}
122
123/// Rolling SMA over a `Vec<f64>` that may contain leading NaN values.
124/// The first valid window requires `period` consecutive non-NaN values.
125fn sma_of(src: &[f64], period: usize) -> Vec<f64> {
126    let n = src.len();
127    let mut out = vec![f64::NAN; n];
128    // Find the first index where `period` consecutive non-NaN values end.
129    let mut consecutive = 0usize;
130    for i in 0..n {
131        if src[i].is_nan() {
132            consecutive = 0;
133        } else {
134            consecutive += 1;
135            if consecutive >= period {
136                let sum: f64 = src[(i + 1 - period)..=i].iter().sum();
137                out[i] = sum / period as f64;
138            }
139        }
140    }
141    out
142}
143
144// ── Registry factory ──────────────────────────────────────────────────────────
145
146pub fn factory<S: ::std::hash::BuildHasher>(params: &HashMap<String, String, S>) -> Result<Box<dyn Indicator>, IndicatorError> {
147    Ok(Box::new(Stochastic::new(StochParams {
148        k_period: param_usize(params, "k_period", 14)?,
149        smooth_k: param_usize(params, "smooth_k", 3)?,
150        d_period: param_usize(params, "d_period", 3)?,
151    })))
152}
153
154// ── Tests ─────────────────────────────────────────────────────────────────────
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    fn make_candles(data: &[(f64, f64, f64)]) -> Vec<Candle> {
161        // (high, low, close)
162        data.iter()
163            .enumerate()
164            .map(|(i, &(h, l, c))| Candle {
165                time: i64::try_from(i).expect("time index fits i64"),
166                open: c,
167                high: h,
168                low: l,
169                close: c,
170                volume: 1.0,
171            })
172            .collect()
173    }
174
175    fn uniform_candles(n: usize, high: f64, low: f64, close: f64) -> Vec<Candle> {
176        make_candles(&vec![(high, low, close); n])
177    }
178
179    #[test]
180    fn stoch_insufficient_data() {
181        let err = Stochastic::default()
182            .calculate(&uniform_candles(5, 12.0, 8.0, 10.0))
183            .unwrap_err();
184        assert!(matches!(err, IndicatorError::InsufficientData { .. }));
185    }
186
187    #[test]
188    fn stoch_output_columns_exist() {
189        let out = Stochastic::default()
190            .calculate(&uniform_candles(30, 12.0, 8.0, 10.0))
191            .unwrap();
192        assert!(out.get("Stoch_K").is_some());
193        assert!(out.get("Stoch_D").is_some());
194    }
195
196    #[test]
197    fn stoch_known_value_midpoint() {
198        // high=12, low=8, close=10 for all bars.
199        // raw %K = 100*(10-8)/(12-8) = 50.0.
200        // smooth_k=3 SMA of [50,50,50,...] = 50. %D = 50.
201        let out = Stochastic::new(StochParams {
202            k_period: 5,
203            smooth_k: 3,
204            d_period: 3,
205        })
206        .calculate(&uniform_candles(20, 12.0, 8.0, 10.0))
207        .unwrap();
208        let k = out.get("Stoch_K").unwrap();
209        let d = out.get("Stoch_D").unwrap();
210        let last_k = k.iter().rev().find(|v| !v.is_nan()).copied().unwrap();
211        let last_d = d.iter().rev().find(|v| !v.is_nan()).copied().unwrap();
212        assert!(
213            (last_k - 50.0).abs() < 1e-9,
214            "K expected 50.0, got {last_k}"
215        );
216        assert!(
217            (last_d - 50.0).abs() < 1e-9,
218            "D expected 50.0, got {last_d}"
219        );
220    }
221
222    #[test]
223    fn stoch_close_at_high_is_100() {
224        // close == high → raw %K = 100.
225        let out = Stochastic::new(StochParams {
226            k_period: 5,
227            smooth_k: 1,
228            d_period: 1,
229        })
230        .calculate(&uniform_candles(10, 12.0, 8.0, 12.0))
231        .unwrap();
232        let k = out.get("Stoch_K").unwrap();
233        for &v in k.iter().filter(|v| !v.is_nan()) {
234            assert!((v - 100.0).abs() < 1e-9, "expected 100.0, got {v}");
235        }
236    }
237
238    #[test]
239    fn stoch_close_at_low_is_0() {
240        // close == low → raw %K = 0.
241        let out = Stochastic::new(StochParams {
242            k_period: 5,
243            smooth_k: 1,
244            d_period: 1,
245        })
246        .calculate(&uniform_candles(10, 12.0, 8.0, 8.0))
247        .unwrap();
248        let k = out.get("Stoch_K").unwrap();
249        for &v in k.iter().filter(|v| !v.is_nan()) {
250            assert!(v.abs() < 1e-9, "expected 0.0, got {v}");
251        }
252    }
253
254    #[test]
255    fn stoch_range_0_to_100() {
256        // Rising then falling sequence.
257        let mut data = vec![];
258        for i in 0..15 {
259            let f = i as f64;
260            data.push((f + 1.0, f - 1.0, f));
261        }
262        for i in (0..10).rev() {
263            let f = i as f64;
264            data.push((f + 1.0, f - 1.0, f));
265        }
266        let out = Stochastic::default()
267            .calculate(&make_candles(&data))
268            .unwrap();
269        for &v in out.get("Stoch_K").unwrap() {
270            if !v.is_nan() {
271                assert!((0.0..=100.0).contains(&v), "K out of range: {v}");
272            }
273        }
274        for &v in out.get("Stoch_D").unwrap() {
275            if !v.is_nan() {
276                assert!((0.0..=100.0).contains(&v), "D out of range: {v}");
277            }
278        }
279    }
280
281    #[test]
282    fn stoch_no_smoothing_fast_stochastic() {
283        // smooth_k=1 → raw %K passed through directly.
284        let out = Stochastic::new(StochParams {
285            k_period: 3,
286            smooth_k: 1,
287            d_period: 1,
288        })
289        .calculate(&uniform_candles(10, 10.0, 0.0, 6.0))
290        .unwrap();
291        // close=6, range=10 → 60.0.
292        let k = out.get("Stoch_K").unwrap();
293        for &v in k.iter().filter(|v| !v.is_nan()) {
294            assert!((v - 60.0).abs() < 1e-9, "expected 60.0, got {v}");
295        }
296    }
297
298    #[test]
299    fn factory_creates_stochastic() {
300        let ind = factory(&HashMap::new()).unwrap();
301        assert_eq!(ind.name(), "Stochastic");
302    }
303}