polars-ops 0.46.0

More operations on Polars data structures
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

pub use arrow::legacy::kernels::ewm::EWMOptions;
use arrow::legacy::kernels::ewm::{
    ewm_mean as kernel_ewm_mean, ewm_std as kernel_ewm_std, ewm_var as kernel_ewm_var,
};
use polars_core::prelude::*;

fn check_alpha(alpha: f64) -> PolarsResult<()> {
    polars_ensure!((0.0..=1.0).contains(&alpha), ComputeError: "alpha must be in [0; 1]");
    Ok(())
}

pub fn ewm_mean(s: &Series, options: EWMOptions) -> PolarsResult<Series> {
    check_alpha(options.alpha)?;
    match s.dtype() {
        DataType::Float32 => {
            let xs = s.f32().unwrap();
            let result = kernel_ewm_mean(
                xs,
                options.alpha as f32,
                options.adjust,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        DataType::Float64 => {
            let xs = s.f64().unwrap();
            let result = kernel_ewm_mean(
                xs,
                options.alpha,
                options.adjust,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        _ => ewm_mean(&s.cast(&DataType::Float64)?, options),
    }
}

pub fn ewm_std(s: &Series, options: EWMOptions) -> PolarsResult<Series> {
    check_alpha(options.alpha)?;
    match s.dtype() {
        DataType::Float32 => {
            let xs = s.f32().unwrap();
            let result = kernel_ewm_std(
                xs,
                options.alpha as f32,
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        DataType::Float64 => {
            let xs = s.f64().unwrap();
            let result = kernel_ewm_std(
                xs,
                options.alpha,
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        _ => ewm_std(&s.cast(&DataType::Float64)?, options),
    }
}

pub fn ewm_var(s: &Series, options: EWMOptions) -> PolarsResult<Series> {
    check_alpha(options.alpha)?;
    match s.dtype() {
        DataType::Float32 => {
            let xs = s.f32().unwrap();
            let result = kernel_ewm_var(
                xs,
                options.alpha as f32,
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        DataType::Float64 => {
            let xs = s.f64().unwrap();
            let result = kernel_ewm_var(
                xs,
                options.alpha,
                options.adjust,
                options.bias,
                options.min_periods,
                options.ignore_nulls,
            );
            Series::try_from((s.name().clone(), Box::new(result) as ArrayRef))
        },
        _ => ewm_var(&s.cast(&DataType::Float64)?, options),
    }
}