polars_arrow/legacy/kernels/ewm/
mod.rs

1mod variance;
2
3use std::hash::{Hash, Hasher};
4
5#[cfg(feature = "serde")]
6use serde::{Deserialize, Serialize};
7pub use variance::*;
8
9#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
10#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
11#[derive(Debug, Copy, Clone, PartialEq)]
12#[must_use]
13pub struct EWMOptions {
14    pub alpha: f64,
15    pub adjust: bool,
16    pub bias: bool,
17    pub min_periods: usize,
18    pub ignore_nulls: bool,
19}
20
21impl Default for EWMOptions {
22    fn default() -> Self {
23        Self {
24            alpha: 0.5,
25            adjust: true,
26            bias: false,
27            min_periods: 1,
28            ignore_nulls: true,
29        }
30    }
31}
32
33impl Hash for EWMOptions {
34    fn hash<H: Hasher>(&self, state: &mut H) {
35        self.alpha.to_bits().hash(state);
36        self.adjust.hash(state);
37        self.bias.hash(state);
38        self.min_periods.hash(state);
39        self.ignore_nulls.hash(state);
40    }
41}
42
43impl EWMOptions {
44    pub fn and_min_periods(mut self, min_periods: usize) -> Self {
45        self.min_periods = min_periods;
46        self
47    }
48    pub fn and_adjust(mut self, adjust: bool) -> Self {
49        self.adjust = adjust;
50        self
51    }
52    pub fn and_span(mut self, span: usize) -> Self {
53        assert!(span >= 1);
54        self.alpha = 2.0 / (span as f64 + 1.0);
55        self
56    }
57    pub fn and_half_life(mut self, half_life: f64) -> Self {
58        assert!(half_life > 0.0);
59        self.alpha = 1.0 - (-(2.0f64.ln()) / half_life).exp();
60        self
61    }
62    pub fn and_com(mut self, com: f64) -> Self {
63        assert!(com > 0.0);
64        self.alpha = 1.0 / (1.0 + com);
65        self
66    }
67    pub fn and_ignore_nulls(mut self, ignore_nulls: bool) -> Self {
68        self.ignore_nulls = ignore_nulls;
69        self
70    }
71}
72
73#[cfg(test)]
74macro_rules! assert_allclose {
75    ($xs:expr, $ys:expr, $tol:expr) => {
76        assert!($xs.iter().zip($ys.iter()).all(|(x, z)| {
77            match (x, z) {
78                (Some(a), Some(b)) => (a - b).abs() < $tol,
79                (None, None) => true,
80                _ => false,
81            }
82        }));
83    };
84}
85#[cfg(test)]
86pub(crate) use assert_allclose;
87
88use crate::array::Array;
89
90pub trait EwmStateUpdate {
91    fn ewm_state_update(&mut self, values: &dyn Array) -> Box<dyn Array>;
92}