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