1use std::collections::HashMap;
21
22use crate::error::IndicatorError;
23use crate::indicator::{Indicator, IndicatorOutput};
24use crate::registry::param_usize;
25use crate::types::Candle;
26
27#[derive(Debug, Clone)]
30pub struct StochParams {
31 pub k_period: usize,
33 pub smooth_k: usize,
35 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#[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 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 let smooth_k = if sk <= 1 {
108 raw_k.clone()
109 } else {
110 sma_of(&raw_k, sk)
111 };
112
113 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
123fn sma_of(src: &[f64], period: usize) -> Vec<f64> {
126 let n = src.len();
127 let mut out = vec![f64::NAN; n];
128 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
144pub 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#[cfg(test)]
157mod tests {
158 use super::*;
159
160 fn make_candles(data: &[(f64, f64, f64)]) -> Vec<Candle> {
161 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 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 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 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 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 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 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}