Skip to main content

nautilus_indicators/average/
wma.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use 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
27/// Maximum supported rolling window period (bounded by the fixed-capacity input buffer).
28pub(crate) const MAX_PERIOD: usize = 8_192;
29
30/// An indicator which calculates a weighted moving average across a rolling window.
31#[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    /// The rolling window period for the indicator (> 0).
43    pub period: usize,
44    /// The weights for the moving average calculation
45    pub weights: Vec<f64>,
46    /// Price type
47    pub price_type: PriceType,
48    /// The last indicator value.
49    pub value: f64,
50    /// Whether the indicator is initialized.
51    pub initialized: bool,
52    /// Inputs
53    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    /// Creates a new [`WeightedMovingAverage`] instance.
64    ///
65    /// # Panics
66    ///
67    /// This function panics if:
68    /// - `period` is zero.
69    /// - `period` exceeds `MAX_PERIOD`.
70    /// - `weights.len()` does not equal `period`.
71    /// - `weights` sum is effectively zero.
72    #[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    /// Creates a new [`WeightedMovingAverage`] instance with the given period and weights.
78    ///
79    /// # Errors
80    ///
81    /// Returns an error if **any** of the validation rules fails:
82    /// - `period` must be **positive**.
83    /// - `period` must not exceed `MAX_PERIOD`.
84    /// - `weights` must be **exactly** `period` elements long.
85    /// - `weights` must contain at least one non-zero value (∑wᵢ > ε).
86    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        // `MAX_PERIOD` is not reachable from Python, so the message has to carry
608        // both the offending period and the bound it exceeded.
609        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        // The boundary itself stays valid: the buffer holds exactly `period`
617        // inputs, so the indicator can still reach `initialized`.
618        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}