1use std::fmt::Display;
17
18use arraydeque::{ArrayDeque, Wrapping};
19use nautilus_core::correctness::{FAILED, check_predicate_true};
20use nautilus_model::{
21 data::{Bar, QuoteTick, TradeTick},
22 enums::PriceType,
23};
24
25use crate::indicator::{Indicator, MovingAverage};
26
27const MAX_PERIOD: usize = 8_192;
28
29#[repr(C)]
31#[derive(Debug)]
32#[cfg_attr(
33 feature = "python",
34 pyo3::pyclass(module = "nautilus_trader.indicators")
35)]
36#[cfg_attr(
37 feature = "python",
38 pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.indicators")
39)]
40pub struct WeightedMovingAverage {
41 pub period: usize,
43 pub weights: Vec<f64>,
45 pub price_type: PriceType,
47 pub value: f64,
49 pub initialized: bool,
51 pub inputs: ArrayDeque<f64, MAX_PERIOD, Wrapping>,
53}
54
55impl Display for WeightedMovingAverage {
56 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57 write!(f, "{}({},{:?})", self.name(), self.period, self.weights)
58 }
59}
60
61impl WeightedMovingAverage {
62 #[must_use]
71 pub fn new(period: usize, weights: Vec<f64>, price_type: Option<PriceType>) -> Self {
72 Self::new_checked(period, weights, price_type).expect(FAILED)
73 }
74
75 pub fn new_checked(
84 period: usize,
85 weights: Vec<f64>,
86 price_type: Option<PriceType>,
87 ) -> anyhow::Result<Self> {
88 const EPS: f64 = f64::EPSILON;
89
90 check_predicate_true(period > 0, "`period` must be positive")?;
91
92 check_predicate_true(
93 period == weights.len(),
94 "`period` must equal `weights.len()`",
95 )?;
96
97 let weight_sum: f64 = weights.iter().copied().sum();
98 check_predicate_true(
99 weight_sum > EPS,
100 "`weights` sum must be positive and > f64::EPSILON",
101 )?;
102
103 Ok(Self {
104 period,
105 weights,
106 price_type: price_type.unwrap_or(PriceType::Last),
107 value: 0.0,
108 inputs: ArrayDeque::new(),
109 initialized: false,
110 })
111 }
112
113 fn weighted_average(&self) -> f64 {
114 let n = self.inputs.len();
115 let weights_slice = &self.weights[self.period - n..];
116
117 let mut sum = 0.0;
118 let mut weight_sum = 0.0;
119
120 for (input, weight) in self.inputs.iter().rev().zip(weights_slice.iter().rev()) {
121 sum += input * weight;
122 weight_sum += weight;
123 }
124 sum / weight_sum
125 }
126}
127
128impl Indicator for WeightedMovingAverage {
129 fn name(&self) -> String {
130 stringify!(WeightedMovingAverage).to_string()
131 }
132
133 fn has_inputs(&self) -> bool {
134 !self.inputs.is_empty()
135 }
136
137 fn initialized(&self) -> bool {
138 self.initialized
139 }
140
141 fn handle_quote(&mut self, quote: &QuoteTick) -> anyhow::Result<()> {
142 self.update_raw(quote.extract_price(self.price_type)?.into());
143 Ok(())
144 }
145
146 fn handle_trade(&mut self, trade: &TradeTick) {
147 self.update_raw((&trade.price).into());
148 }
149
150 fn handle_bar(&mut self, bar: &Bar) {
151 self.update_raw((&bar.close).into());
152 }
153
154 fn reset(&mut self) {
155 self.value = 0.0;
156 self.initialized = false;
157 self.inputs.clear();
158 }
159}
160
161impl MovingAverage for WeightedMovingAverage {
162 fn value(&self) -> f64 {
163 self.value
164 }
165
166 fn count(&self) -> usize {
167 self.inputs.len()
168 }
169
170 fn update_raw(&mut self, value: f64) {
171 if self.inputs.len() == self.period.min(MAX_PERIOD) {
172 self.inputs.pop_front();
173 }
174 let _ = self.inputs.push_back(value);
175
176 self.value = self.weighted_average();
177 self.initialized = self.count() >= self.period;
178 }
179}
180
181#[cfg(test)]
182mod tests {
183
184 use arraydeque::{ArrayDeque, Wrapping};
185 use rstest::rstest;
186
187 use crate::{
188 average::wma::WeightedMovingAverage,
189 indicator::{Indicator, MovingAverage},
190 stubs::*,
191 testing::assert_approx_equal,
192 };
193
194 #[rstest]
195 fn test_wma_initialized(indicator_wma_10: WeightedMovingAverage) {
196 let display_str = format!("{indicator_wma_10}");
197 assert_eq!(
198 display_str,
199 "WeightedMovingAverage(10,[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])"
200 );
201 assert_eq!(indicator_wma_10.name(), "WeightedMovingAverage");
202 assert!(!indicator_wma_10.has_inputs());
203 assert!(!indicator_wma_10.initialized());
204 }
205
206 #[rstest]
207 #[should_panic]
208 fn test_different_weights_len_and_period_error() {
209 let _ = WeightedMovingAverage::new(10, vec![0.5, 0.5, 0.5], None);
210 }
211
212 #[rstest]
213 fn test_value_with_one_input(mut indicator_wma_10: WeightedMovingAverage) {
214 indicator_wma_10.update_raw(1.0);
215 assert_eq!(indicator_wma_10.value, 1.0);
216 }
217
218 #[rstest]
219 fn test_value_with_two_inputs_equal_weights() {
220 let mut wma = WeightedMovingAverage::new(2, vec![0.5, 0.5], None);
221 wma.update_raw(1.0);
222 wma.update_raw(2.0);
223 assert_eq!(wma.value, 1.5);
224 }
225
226 #[rstest]
227 fn test_value_with_four_inputs_equal_weights() {
228 let mut wma = WeightedMovingAverage::new(4, vec![0.25, 0.25, 0.25, 0.25], None);
229 wma.update_raw(1.0);
230 wma.update_raw(2.0);
231 wma.update_raw(3.0);
232 wma.update_raw(4.0);
233 assert_eq!(wma.value, 2.5);
234 }
235
236 #[rstest]
237 fn test_value_with_two_inputs(mut indicator_wma_10: WeightedMovingAverage) {
238 indicator_wma_10.update_raw(1.0);
239 indicator_wma_10.update_raw(2.0);
240 let result = 2.0f64.mul_add(1.0, 1.0 * 0.9) / 1.9;
241 assert_eq!(indicator_wma_10.value, result);
242 }
243
244 #[rstest]
245 fn test_value_with_three_inputs(mut indicator_wma_10: WeightedMovingAverage) {
246 indicator_wma_10.update_raw(1.0);
247 indicator_wma_10.update_raw(2.0);
248 indicator_wma_10.update_raw(3.0);
249 let result = 1.0f64.mul_add(0.8, 3.0f64.mul_add(1.0, 2.0 * 0.9)) / (1.0 + 0.9 + 0.8);
250 assert_eq!(indicator_wma_10.value, result);
251 }
252
253 #[rstest]
254 fn test_value_expected_with_exact_period(mut indicator_wma_10: WeightedMovingAverage) {
255 for i in 1..11 {
256 indicator_wma_10.update_raw(f64::from(i));
257 }
258 assert_eq!(indicator_wma_10.value, 7.0);
259 }
260
261 #[rstest]
262 fn test_value_expected_with_more_inputs(mut indicator_wma_10: WeightedMovingAverage) {
263 for i in 1..=11 {
264 indicator_wma_10.update_raw(f64::from(i));
265 }
266 assert_approx_equal(indicator_wma_10.value(), 8.0);
267 }
268
269 #[rstest]
270 fn test_reset(mut indicator_wma_10: WeightedMovingAverage) {
271 indicator_wma_10.update_raw(1.0);
272 indicator_wma_10.update_raw(2.0);
273 indicator_wma_10.reset();
274 assert_eq!(indicator_wma_10.value, 0.0);
275 assert_eq!(indicator_wma_10.count(), 0);
276 assert!(!indicator_wma_10.initialized);
277 }
278
279 #[rstest]
280 #[should_panic]
281 fn new_panics_on_zero_period() {
282 let _ = WeightedMovingAverage::new(0, vec![1.0], None);
283 }
284
285 #[rstest]
286 fn new_checked_err_on_zero_period() {
287 let res = WeightedMovingAverage::new_checked(0, vec![1.0], None);
288 assert!(res.is_err());
289 }
290
291 #[rstest]
292 #[should_panic]
293 fn new_panics_on_zero_weight_sum() {
294 let _ = WeightedMovingAverage::new(3, vec![0.0, 0.0, 0.0], None);
295 }
296
297 #[rstest]
298 fn new_checked_err_on_zero_weight_sum() {
299 let res = WeightedMovingAverage::new_checked(3, vec![0.0, 0.0, 0.0], None);
300 assert!(res.is_err());
301 }
302
303 #[rstest]
304 #[should_panic]
305 fn new_panics_when_weight_sum_below_epsilon() {
306 let tiny = f64::EPSILON / 10.0;
307 let _ = WeightedMovingAverage::new(3, vec![tiny; 3], None);
308 }
309
310 #[rstest]
311 fn initialized_flag_transitions() {
312 let period = 3;
313 let weights = vec![1.0, 2.0, 3.0];
314 let mut wma = WeightedMovingAverage::new(period, weights, None);
315
316 assert!(!wma.initialized());
317
318 for i in 0..period {
319 wma.update_raw(i as f64);
320 let expected = (i + 1) >= period;
321 assert_eq!(wma.initialized(), expected);
322 }
323 assert!(wma.initialized());
324 }
325
326 #[rstest]
327 fn count_matches_inputs_and_has_inputs() {
328 let mut wma = WeightedMovingAverage::new(4, vec![0.25; 4], None);
329
330 assert_eq!(wma.count(), 0);
331 assert!(!wma.has_inputs());
332
333 wma.update_raw(1.0);
334 wma.update_raw(2.0);
335 assert_eq!(wma.count(), 2);
336 assert!(wma.has_inputs());
337 }
338
339 #[rstest]
340 fn reset_restores_pristine_state() {
341 let mut wma = WeightedMovingAverage::new(2, vec![0.5, 0.5], None);
342 wma.update_raw(1.0);
343 wma.update_raw(2.0);
344 assert!(wma.initialized());
345
346 wma.reset();
347
348 assert_eq!(wma.count(), 0);
349 assert_eq!(wma.value(), 0.0);
350 assert!(!wma.initialized());
351 assert!(!wma.has_inputs());
352 }
353
354 #[rstest]
355 fn weighted_average_with_non_uniform_weights() {
356 let mut wma = WeightedMovingAverage::new(3, vec![1.0, 2.0, 3.0], None);
357 wma.update_raw(10.0);
358 wma.update_raw(20.0);
359 wma.update_raw(30.0);
360 let expected = 23.333_333_333_333_332;
361 let tol = f64::EPSILON.sqrt();
362 assert!(
363 (wma.value() - expected).abs() < tol,
364 "value = {}, expected ≈ {}",
365 wma.value(),
366 expected
367 );
368 }
369
370 #[rstest]
371 fn test_window_never_exceeds_period(mut indicator_wma_10: WeightedMovingAverage) {
372 for i in 0..100 {
373 indicator_wma_10.update_raw(f64::from(i));
374 assert!(indicator_wma_10.count() <= indicator_wma_10.period);
375 }
376 }
377
378 #[rstest]
379 fn test_negative_weights_positive_sum() {
380 let period = 3;
381 let weights = vec![-1.0, 2.0, 2.0];
382 let mut wma = WeightedMovingAverage::new(period, weights, None);
383 wma.update_raw(1.0);
384 wma.update_raw(2.0);
385 wma.update_raw(3.0);
386
387 let expected = 2.0f64.mul_add(3.0, 2.0f64.mul_add(2.0, -1.0)) / 3.0;
388 let tol = f64::EPSILON.sqrt();
389 assert!((wma.value() - expected).abs() < tol);
390 }
391
392 #[rstest]
393 fn test_nan_input_propagates() {
394 let mut wma = WeightedMovingAverage::new(2, vec![0.5, 0.5], None);
395 wma.update_raw(1.0);
396 wma.update_raw(f64::NAN);
397
398 assert!(wma.value().is_nan());
399 }
400
401 #[rstest]
402 #[should_panic]
403 fn new_panics_when_weight_sum_equals_epsilon() {
404 let eps_third = f64::EPSILON / 3.0;
405 let _ = WeightedMovingAverage::new(3, vec![eps_third; 3], None);
406 }
407
408 #[rstest]
409 fn new_checked_err_when_weight_sum_equals_epsilon() {
410 let eps_third = f64::EPSILON / 3.0;
411 let res = WeightedMovingAverage::new_checked(3, vec![eps_third; 3], None);
412 assert!(res.is_err());
413 }
414
415 #[rstest]
416 fn new_checked_err_when_weight_sum_below_epsilon() {
417 let w = f64::EPSILON * 0.9;
418 let res = WeightedMovingAverage::new_checked(1, vec![w], None);
419 assert!(res.is_err());
420 }
421
422 #[rstest]
423 fn new_ok_when_weight_sum_above_epsilon() {
424 let w = f64::EPSILON * 1.1;
425 let res = WeightedMovingAverage::new_checked(1, vec![w], None);
426 assert!(res.is_ok());
427 }
428
429 #[rstest]
430 #[should_panic]
431 fn new_panics_on_cancelled_weights_sum() {
432 let _ = WeightedMovingAverage::new(3, vec![1.0, -1.0, 0.0], None);
433 }
434
435 #[rstest]
436 fn new_checked_err_on_cancelled_weights_sum() {
437 let res = WeightedMovingAverage::new_checked(3, vec![1.0, -1.0, 0.0], None);
438 assert!(res.is_err());
439 }
440
441 #[rstest]
442 fn single_period_returns_latest_input() {
443 let mut wma = WeightedMovingAverage::new(1, vec![1.0], None);
444
445 for i in 0..5 {
446 let v = f64::from(i);
447 wma.update_raw(v);
448 assert_eq!(wma.value(), v);
449 }
450 }
451
452 #[rstest]
453 fn value_with_sparse_weights() {
454 let mut wma = WeightedMovingAverage::new(3, vec![0.0, 1.0, 0.0], None);
455 wma.update_raw(10.0);
456 wma.update_raw(20.0);
457 wma.update_raw(30.0);
458 assert_eq!(wma.value(), 20.0);
459 }
460
461 #[rstest]
462 fn warm_up_len1() {
463 let mut wma = WeightedMovingAverage::new(4, vec![1.0, 2.0, 3.0, 4.0], None);
464 wma.update_raw(42.0);
465 assert_eq!(wma.value(), 42.0);
466 }
467
468 #[rstest]
469 fn warm_up_len2() {
470 let mut wma = WeightedMovingAverage::new(4, vec![1.0, 2.0, 3.0, 4.0], None);
471 wma.update_raw(10.0);
472 wma.update_raw(20.0);
473 let expected = 20.0f64.mul_add(4.0, 10.0 * 3.0) / (4.0 + 3.0);
474 assert_eq!(wma.value(), expected);
475 }
476
477 #[rstest]
478 fn warm_up_len3() {
479 let mut wma = WeightedMovingAverage::new(4, vec![1.0, 2.0, 3.0, 4.0], None);
480 wma.update_raw(1.0);
481 wma.update_raw(2.0);
482 wma.update_raw(3.0);
483 let expected = 1.0f64.mul_add(2.0, 3.0f64.mul_add(4.0, 2.0 * 3.0)) / (4.0 + 3.0 + 2.0);
484 assert_eq!(wma.value(), expected);
485 }
486
487 #[rstest]
488 fn input_window_contains_latest_period() {
489 let period = 3;
490 let mut wma = WeightedMovingAverage::new(period, vec![1.0; period], None);
491 let vals = [1.0, 2.0, 3.0, 4.0];
492 for v in vals {
493 wma.update_raw(v);
494 }
495 let expected: Vec<f64> = vals[vals.len() - period..].to_vec();
496 assert_eq!(wma.inputs.iter().copied().collect::<Vec<_>>(), expected);
497 }
498
499 #[rstest]
500 fn window_slides_correctly() {
501 let mut wma = WeightedMovingAverage::new(2, vec![1.0; 2], None);
502 wma.update_raw(1.0);
503 assert_eq!(wma.inputs.iter().copied().collect::<Vec<_>>(), vec![1.0]);
504 wma.update_raw(2.0);
505 assert_eq!(
506 wma.inputs.iter().copied().collect::<Vec<_>>(),
507 vec![1.0, 2.0]
508 );
509 wma.update_raw(3.0);
510 assert_eq!(
511 wma.inputs.iter().copied().collect::<Vec<_>>(),
512 vec![2.0, 3.0]
513 );
514 }
515
516 #[rstest]
517 fn window_len_constant_after_many_updates() {
518 let period = 5;
519 let mut wma = WeightedMovingAverage::new(period, vec![1.0; period], None);
520 for i in 0..100 {
521 wma.update_raw(i as f64);
522 assert_eq!(wma.inputs.len(), period.min(i + 1));
523 }
524 }
525
526 #[rstest]
527 fn arraydeque_wraps_when_full() {
528 const CAP: usize = 3;
529 let mut buf: ArrayDeque<usize, CAP, Wrapping> = ArrayDeque::new();
530 for i in 0..=CAP {
531 let _ = buf.push_back(i);
532 }
533 assert_eq!(buf.len(), CAP);
534 assert_eq!(buf.front().copied(), Some(1));
535 assert_eq!(buf.back().copied(), Some(3));
536 }
537
538 #[rstest]
539 fn arraydeque_sliding_window_with_pop() {
540 const CAP: usize = 3;
541 let mut buf: ArrayDeque<usize, CAP, Wrapping> = ArrayDeque::new();
542 for i in 0..10 {
543 if buf.len() == CAP {
544 buf.pop_front();
545 }
546 let _ = buf.push_back(i);
547 assert!(buf.len() <= CAP);
548 }
549 assert_eq!(buf.len(), CAP);
550 }
551
552 #[rstest]
553 fn new_ok_with_infinite_weight() {
554 let res = WeightedMovingAverage::new_checked(2, vec![f64::INFINITY, 1.0], None);
555 assert!(res.is_ok());
556 }
557
558 #[rstest]
559 #[should_panic]
560 fn new_panics_on_nan_weight() {
561 let _ = WeightedMovingAverage::new(2, vec![f64::NAN, 1.0], None);
562 }
563
564 #[rstest]
565 #[should_panic]
566 fn new_panics_on_empty_weights() {
567 let _ = WeightedMovingAverage::new(1, Vec::new(), None);
568 }
569
570 #[rstest]
571 fn inf_input_propagates() {
572 let mut wma = WeightedMovingAverage::new(2, vec![0.5, 0.5], None);
573 wma.update_raw(1.0);
574 wma.update_raw(f64::INFINITY);
575 assert!(wma.value().is_infinite());
576 }
577
578 #[rstest]
579 fn warm_up_with_front_zero_weights() {
580 let mut wma = WeightedMovingAverage::new(4, vec![0.0, 0.0, 1.0, 1.0], None);
581 wma.update_raw(10.0);
582 wma.update_raw(20.0);
583 let expected = 20.0f64.mul_add(1.0, 10.0 * 1.0) / 2.0;
584 assert_eq!(wma.value(), expected);
585 }
586}