CoCalc -- rolling.rs

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-ops/src/series/ops/rolling.rs
⁸⁴³³ views
1
use polars_core::prelude::*;
2
#[cfg(feature = "moment")]
3
use {
4
    num_traits::{Float, pow::Pow},
5
    std::ops::SubAssign,
6
};
7

8
#[cfg(feature = "moment")]
9
fn rolling_skew_ca<T>(
10
    ca: &ChunkedArray<T>,
11
    window_size: usize,
12
    min_periods: usize,
13
    center: bool,
14
    params: Option<RollingFnParams>,
15
) -> PolarsResult<ChunkedArray<T>>
16
where
17
    T: PolarsFloatType,
18
    T::Native: Float + SubAssign + Pow<T::Native, Output = T::Native>,
19
{
20
    use arrow::array::Array;
21

22
    let ca = ca.rechunk();
23
    let arr = ca.downcast_get(0).unwrap();
24
    let arr = if arr.has_nulls() {
25
        polars_compute::rolling::nulls::rolling_skew(arr, window_size, min_periods, center, params)
26
    } else {
27
        let values = arr.values();
28
        polars_compute::rolling::no_nulls::rolling_skew(
29
            values,
30
            window_size,
31
            min_periods,
32
            center,
33
            params,
34
        )?
35
    };
36
    Ok(unsafe { ca.with_chunks(vec![arr]) })
37
}
38

39
#[cfg(feature = "moment")]
40
pub fn rolling_skew(s: &Series, options: RollingOptionsFixedWindow) -> PolarsResult<Series> {
41
    let window_size = options.window_size;
42
    let min_periods = options.min_periods;
43
    let center = options.center;
44
    let params = options.fn_params;
45

46
    match s.dtype() {
47
        #[cfg(feature = "dtype-f16")]
48
        DataType::Float16 => {
49
            let ca = s.f16().unwrap();
50
            rolling_skew_ca(ca, window_size, min_periods, center, params).map(|ca| ca.into_series())
51
        },
52
        DataType::Float64 => {
53
            let ca = s.f64().unwrap();
54
            rolling_skew_ca(ca, window_size, min_periods, center, params).map(|ca| ca.into_series())
55
        },
56
        DataType::Float32 => {
57
            let ca = s.f32().unwrap();
58
            rolling_skew_ca(ca, window_size, min_periods, center, params).map(|ca| ca.into_series())
59
        },
60
        dt if dt.is_primitive_numeric() => {
61
            let s = s.cast(&DataType::Float64).unwrap();
62
            rolling_skew(&s, options)
63
        },
64
        dt => polars_bail!(opq = rolling_skew, dt),
65
    }
66
}
67

68
#[cfg(feature = "moment")]
69
fn rolling_kurtosis_ca<T>(
70
    ca: &ChunkedArray<T>,
71
    window_size: usize,
72
    params: Option<RollingFnParams>,
73
    min_periods: usize,
74
    center: bool,
75
) -> PolarsResult<ChunkedArray<T>>
76
where
77
    T: PolarsFloatType,
78
    T::Native: Float + SubAssign + Pow<T::Native, Output = T::Native>,
79
{
80
    use arrow::array::Array;
81

82
    let ca = ca.rechunk();
83
    let arr = ca.downcast_get(0).unwrap();
84
    let arr = if arr.has_nulls() {
85
        polars_compute::rolling::nulls::rolling_kurtosis(
86
            arr,
87
            window_size,
88
            min_periods,
89
            center,
90
            params,
91
        )
92
    } else {
93
        let values = arr.values();
94
        polars_compute::rolling::no_nulls::rolling_kurtosis(
95
            values,
96
            window_size,
97
            min_periods,
98
            center,
99
            params,
100
        )?
101
    };
102
    Ok(unsafe { ca.with_chunks(vec![arr]) })
103
}
104

105
#[cfg(feature = "moment")]
106
pub fn rolling_kurtosis(s: &Series, options: RollingOptionsFixedWindow) -> PolarsResult<Series> {
107
    let window_size = options.window_size;
108
    let min_periods = options.min_periods;
109
    let center = options.center;
110
    let params = options.fn_params;
111

112
    match s.dtype() {
113
        DataType::Float64 => {
114
            let ca = s.f64().unwrap();
115
            rolling_kurtosis_ca(ca, window_size, params, min_periods, center)
116
                .map(|ca| ca.into_series())
117
        },
118
        DataType::Float32 => {
119
            let ca = s.f32().unwrap();
120
            rolling_kurtosis_ca(ca, window_size, params, min_periods, center)
121
                .map(|ca| ca.into_series())
122
        },
123
        #[cfg(feature = "dtype-f16")]
124
        DataType::Float16 => {
125
            let ca = s.f16().unwrap();
126
            rolling_kurtosis_ca(ca, window_size, params, min_periods, center)
127
                .map(|ca| ca.into_series())
128
        },
129
        dt if dt.is_primitive_numeric() => {
130
            let s = s.cast(&DataType::Float64).unwrap();
131
            rolling_kurtosis(&s, options)
132
        },
133
        dt => polars_bail!(opq = rolling_kurtosis, dt),
134
    }
135
}
136

137
Product

Resources

Company
