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