Skip to main content

polars_compute/ewm/
options.rs

1use std::hash::{Hash, Hasher};
2
3#[cfg(feature = "serde")]
4use serde::{Deserialize, Serialize};
5
6#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
7#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
8#[derive(Debug, Copy, Clone, PartialEq)]
9#[must_use]
10pub struct EWMOptions {
11    pub alpha: f64,
12    pub adjust: bool,
13    pub bias: bool,
14    pub min_periods: usize,
15    pub ignore_nulls: bool,
16}
17
18impl Default for EWMOptions {
19    fn default() -> Self {
20        Self {
21            alpha: 0.5,
22            adjust: true,
23            bias: false,
24            min_periods: 1,
25            ignore_nulls: true,
26        }
27    }
28}
29
30impl Hash for EWMOptions {
31    fn hash<H: Hasher>(&self, state: &mut H) {
32        self.alpha.to_bits().hash(state);
33        self.adjust.hash(state);
34        self.bias.hash(state);
35        self.min_periods.hash(state);
36        self.ignore_nulls.hash(state);
37    }
38}
39
40impl EWMOptions {
41    pub fn and_min_periods(mut self, min_periods: usize) -> Self {
42        self.min_periods = min_periods;
43        self
44    }
45    pub fn and_adjust(mut self, adjust: bool) -> Self {
46        self.adjust = adjust;
47        self
48    }
49    pub fn and_span(mut self, span: usize) -> Self {
50        assert!(span >= 1);
51        self.alpha = 2.0 / (span as f64 + 1.0);
52        self
53    }
54    pub fn and_half_life(mut self, half_life: f64) -> Self {
55        assert!(half_life > 0.0);
56        self.alpha = 1.0 - (-(2.0f64.ln()) / half_life).exp();
57        self
58    }
59    pub fn and_com(mut self, com: f64) -> Self {
60        assert!(com > 0.0);
61        self.alpha = 1.0 / (1.0 + com);
62        self
63    }
64    pub fn and_ignore_nulls(mut self, ignore_nulls: bool) -> Self {
65        self.ignore_nulls = ignore_nulls;
66        self
67    }
68}