CoCalc -- rolling.rs

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-python/src/expr/rolling.rs
⁷⁸⁸⁹ views
1
use polars::prelude::*;
2
use polars_utils::python_function::PythonObject;
3
use pyo3::prelude::*;
4

5
use crate::PyExpr;
6
use crate::conversion::Wrap;
7
use crate::error::PyPolarsErr;
8

9
#[pymethods]
10
impl PyExpr {
11
    #[pyo3(signature = (window_size, weights, min_periods, center))]
12
    fn rolling_sum(
13
        &self,
14
        window_size: usize,
15
        weights: Option<Vec<f64>>,
16
        min_periods: Option<usize>,
17
        center: bool,
18
    ) -> Self {
19
        let min_periods = min_periods.unwrap_or(window_size);
20
        let options = RollingOptionsFixedWindow {
21
            window_size,
22
            weights,
23
            min_periods,
24
            center,
25
            ..Default::default()
26
        };
27
        self.inner.clone().rolling_sum(options).into()
28
    }
29

30
    #[pyo3(signature = (by, window_size, min_periods, closed))]
31
    fn rolling_sum_by(
32
        &self,
33
        by: PyExpr,
34
        window_size: &str,
35
        min_periods: usize,
36
        closed: Wrap<ClosedWindow>,
37
    ) -> PyResult<Self> {
38
        let options = RollingOptionsDynamicWindow {
39
            window_size: Duration::try_parse(window_size).map_err(PyPolarsErr::from)?,
40
            min_periods,
41
            closed_window: closed.0,
42
            fn_params: None,
43
        };
44
        Ok(self.inner.clone().rolling_sum_by(by.inner, options).into())
45
    }
46

47
    #[pyo3(signature = (window_size, weights, min_periods, center))]
48
    fn rolling_min(
49
        &self,
50
        window_size: usize,
51
        weights: Option<Vec<f64>>,
52
        min_periods: Option<usize>,
53
        center: bool,
54
    ) -> Self {
55
        let min_periods = min_periods.unwrap_or(window_size);
56
        let options = RollingOptionsFixedWindow {
57
            window_size,
58
            weights,
59
            min_periods,
60
            center,
61
            ..Default::default()
62
        };
63
        self.inner.clone().rolling_min(options).into()
64
    }
65

66
    #[pyo3(signature = (by, window_size, min_periods, closed))]
67
    fn rolling_min_by(
68
        &self,
69
        by: PyExpr,
70
        window_size: &str,
71
        min_periods: usize,
72
        closed: Wrap<ClosedWindow>,
73
    ) -> PyResult<Self> {
74
        let options = RollingOptionsDynamicWindow {
75
            window_size: Duration::try_parse(window_size).map_err(PyPolarsErr::from)?,
76
            min_periods,
77
            closed_window: closed.0,
78
            fn_params: None,
79
        };
80
        Ok(self.inner.clone().rolling_min_by(by.inner, options).into())
81
    }
82

83
    #[pyo3(signature = (window_size, weights, min_periods, center))]
84
    fn rolling_max(
85
        &self,
86
        window_size: usize,
87
        weights: Option<Vec<f64>>,
88
        min_periods: Option<usize>,
89
        center: bool,
90
    ) -> Self {
91
        let min_periods = min_periods.unwrap_or(window_size);
92
        let options = RollingOptionsFixedWindow {
93
            window_size,
94
            weights,
95
            min_periods,
96
            center,
97
            ..Default::default()
98
        };
99
        self.inner.clone().rolling_max(options).into()
100
    }
101
    #[pyo3(signature = (by, window_size, min_periods, closed))]
102
    fn rolling_max_by(
103
        &self,
104
        by: PyExpr,
105
        window_size: &str,
106
        min_periods: usize,
107
        closed: Wrap<ClosedWindow>,
108
    ) -> PyResult<Self> {
109
        let options = RollingOptionsDynamicWindow {
110
            window_size: Duration::try_parse(window_size).map_err(PyPolarsErr::from)?,
111
            min_periods,
112
            closed_window: closed.0,
113
            fn_params: None,
114
        };
115
        Ok(self.inner.clone().rolling_max_by(by.inner, options).into())
116
    }
117

118
    #[pyo3(signature = (window_size, weights, min_periods, center))]
119
    fn rolling_mean(
120
        &self,
121
        window_size: usize,
122
        weights: Option<Vec<f64>>,
123
        min_periods: Option<usize>,
124
        center: bool,
125
    ) -> Self {
126
        let min_periods = min_periods.unwrap_or(window_size);
127
        let options = RollingOptionsFixedWindow {
128
            window_size,
129
            weights,
130
            min_periods,
131
            center,
132
            ..Default::default()
133
        };
134

135
        self.inner.clone().rolling_mean(options).into()
136
    }
137

138
    #[pyo3(signature = (by, window_size, min_periods, closed))]
139
    fn rolling_mean_by(
140
        &self,
141
        by: PyExpr,
142
        window_size: &str,
143
        min_periods: usize,
144
        closed: Wrap<ClosedWindow>,
145
    ) -> PyResult<Self> {
146
        let options = RollingOptionsDynamicWindow {
147
            window_size: Duration::try_parse(window_size).map_err(PyPolarsErr::from)?,
148
            min_periods,
149
            closed_window: closed.0,
150
            fn_params: None,
151
        };
152

153
        Ok(self.inner.clone().rolling_mean_by(by.inner, options).into())
154
    }
155

156
    #[pyo3(signature = (window_size, weights, min_periods, center, ddof))]
157
    fn rolling_std(
158
        &self,
159
        window_size: usize,
160
        weights: Option<Vec<f64>>,
161
        min_periods: Option<usize>,
162
        center: bool,
163
        ddof: u8,
164
    ) -> Self {
165
        let min_periods = min_periods.unwrap_or(window_size);
166
        let options = RollingOptionsFixedWindow {
167
            window_size,
168
            weights,
169
            min_periods,
170
            center,
171
            fn_params: Some(RollingFnParams::Var(RollingVarParams { ddof })),
172
        };
173

174
        self.inner.clone().rolling_std(options).into()
175
    }
176

177
    #[pyo3(signature = (by, window_size, min_periods, closed, ddof))]
178
    fn rolling_std_by(
179
        &self,
180
        by: PyExpr,
181
        window_size: &str,
182
        min_periods: usize,
183
        closed: Wrap<ClosedWindow>,
184
        ddof: u8,
185
    ) -> PyResult<Self> {
186
        let options = RollingOptionsDynamicWindow {
187
            window_size: Duration::try_parse(window_size).map_err(PyPolarsErr::from)?,
188
            min_periods,
189
            closed_window: closed.0,
190
            fn_params: Some(RollingFnParams::Var(RollingVarParams { ddof })),
191
        };
192

193
        Ok(self.inner.clone().rolling_std_by(by.inner, options).into())
194
    }
195

196
    #[pyo3(signature = (window_size, weights, min_periods, center, ddof))]
197
    fn rolling_var(
198
        &self,
199
        window_size: usize,
200
        weights: Option<Vec<f64>>,
201
        min_periods: Option<usize>,
202
        center: bool,
203
        ddof: u8,
204
    ) -> Self {
205
        let min_periods = min_periods.unwrap_or(window_size);
206
        let options = RollingOptionsFixedWindow {
207
            window_size,
208
            weights,
209
            min_periods,
210
            center,
211
            fn_params: Some(RollingFnParams::Var(RollingVarParams { ddof })),
212
        };
213

214
        self.inner.clone().rolling_var(options).into()
215
    }
216

217
    #[pyo3(signature = (by, window_size, min_periods, closed, ddof))]
218
    fn rolling_var_by(
219
        &self,
220
        by: PyExpr,
221
        window_size: &str,
222
        min_periods: usize,
223
        closed: Wrap<ClosedWindow>,
224
        ddof: u8,
225
    ) -> PyResult<Self> {
226
        let options = RollingOptionsDynamicWindow {
227
            window_size: Duration::try_parse(window_size).map_err(PyPolarsErr::from)?,
228
            min_periods,
229
            closed_window: closed.0,
230
            fn_params: Some(RollingFnParams::Var(RollingVarParams { ddof })),
231
        };
232

233
        Ok(self.inner.clone().rolling_var_by(by.inner, options).into())
234
    }
235

236
    #[pyo3(signature = (window_size, weights, min_periods, center))]
237
    fn rolling_median(
238
        &self,
239
        window_size: usize,
240
        weights: Option<Vec<f64>>,
241
        min_periods: Option<usize>,
242
        center: bool,
243
    ) -> Self {
244
        let min_periods = min_periods.unwrap_or(window_size);
245
        let options = RollingOptionsFixedWindow {
246
            window_size,
247
            min_periods,
248
            weights,
249
            center,
250
            fn_params: None,
251
        };
252
        self.inner.clone().rolling_median(options).into()
253
    }
254

255
    #[pyo3(signature = (by, window_size, min_periods, closed))]
256
    fn rolling_median_by(
257
        &self,
258
        by: PyExpr,
259
        window_size: &str,
260
        min_periods: usize,
261
        closed: Wrap<ClosedWindow>,
262
    ) -> PyResult<Self> {
263
        let options = RollingOptionsDynamicWindow {
264
            window_size: Duration::try_parse(window_size).map_err(PyPolarsErr::from)?,
265
            min_periods,
266
            closed_window: closed.0,
267
            fn_params: None,
268
        };
269
        Ok(self
270
            .inner
271
            .clone()
272
            .rolling_median_by(by.inner, options)
273
            .into())
274
    }
275

276
    #[pyo3(signature = (quantile, interpolation, window_size, weights, min_periods, center))]
277
    fn rolling_quantile(
278
        &self,
279
        quantile: f64,
280
        interpolation: Wrap<QuantileMethod>,
281
        window_size: usize,
282
        weights: Option<Vec<f64>>,
283
        min_periods: Option<usize>,
284
        center: bool,
285
    ) -> Self {
286
        let min_periods = min_periods.unwrap_or(window_size);
287
        let options = RollingOptionsFixedWindow {
288
            window_size,
289
            weights,
290
            min_periods,
291
            center,
292
            fn_params: None,
293
        };
294

295
        self.inner
296
            .clone()
297
            .rolling_quantile(interpolation.0, quantile, options)
298
            .into()
299
    }
300

301
    #[pyo3(signature = (by, quantile, interpolation, window_size, min_periods, closed))]
302
    fn rolling_quantile_by(
303
        &self,
304
        by: PyExpr,
305
        quantile: f64,
306
        interpolation: Wrap<QuantileMethod>,
307
        window_size: &str,
308
        min_periods: usize,
309
        closed: Wrap<ClosedWindow>,
310
    ) -> PyResult<Self> {
311
        let options = RollingOptionsDynamicWindow {
312
            window_size: Duration::try_parse(window_size).map_err(PyPolarsErr::from)?,
313
            min_periods,
314
            closed_window: closed.0,
315
            fn_params: None,
316
        };
317

318
        Ok(self
319
            .inner
320
            .clone()
321
            .rolling_quantile_by(by.inner, interpolation.0, quantile, options)
322
            .into())
323
    }
324

325
    #[pyo3(signature = (window_size, method, seed, min_samples, center))]
326
    fn rolling_rank(
327
        &self,
328
        window_size: usize,
329
        method: Wrap<RollingRankMethod>,
330
        seed: Option<u64>,
331
        min_samples: Option<usize>,
332
        center: bool,
333
    ) -> Self {
334
        let min_samples = min_samples.unwrap_or(window_size);
335
        let options = RollingOptionsFixedWindow {
336
            window_size,
337
            min_periods: min_samples,
338
            weights: None,
339
            center,
340
            fn_params: Some(RollingFnParams::Rank {
341
                method: method.0,
342
                seed,
343
            }),
344
        };
345

346
        self.inner.clone().rolling_rank(options).into()
347
    }
348

349
    #[pyo3(signature = (by, window_size, method, seed, min_samples, closed))]
350
    fn rolling_rank_by(
351
        &self,
352
        by: PyExpr,
353
        window_size: &str,
354
        method: Wrap<RollingRankMethod>,
355
        seed: Option<u64>,
356
        min_samples: usize,
357
        closed: Wrap<ClosedWindow>,
358
    ) -> PyResult<Self> {
359
        let options = RollingOptionsDynamicWindow {
360
            window_size: Duration::try_parse(window_size).map_err(PyPolarsErr::from)?,
361
            min_periods: min_samples,
362
            closed_window: closed.0,
363
            fn_params: Some(RollingFnParams::Rank {
364
                method: method.0,
365
                seed,
366
            }),
367
        };
368

369
        Ok(self.inner.clone().rolling_rank_by(by.inner, options).into())
370
    }
371

372
    #[pyo3(signature = (window_size, bias, min_periods, center))]
373
    fn rolling_skew(
374
        &self,
375
        window_size: usize,
376
        bias: bool,
377
        min_periods: Option<usize>,
378
        center: bool,
379
    ) -> Self {
380
        let min_periods = min_periods.unwrap_or(window_size);
381
        let options = RollingOptionsFixedWindow {
382
            window_size,
383
            weights: None,
384
            min_periods,
385
            center,
386
            fn_params: Some(RollingFnParams::Skew { bias }),
387
        };
388

389
        self.inner.clone().rolling_skew(options).into()
390
    }
391

392
    #[pyo3(signature = (window_size, fisher, bias, min_periods, center))]
393
    fn rolling_kurtosis(
394
        &self,
395
        window_size: usize,
396
        fisher: bool,
397
        bias: bool,
398
        min_periods: Option<usize>,
399
        center: bool,
400
    ) -> Self {
401
        let min_periods = min_periods.unwrap_or(window_size);
402
        let options = RollingOptionsFixedWindow {
403
            window_size,
404
            weights: None,
405
            min_periods,
406
            center,
407
            fn_params: Some(RollingFnParams::Kurtosis { fisher, bias }),
408
        };
409

410
        self.inner.clone().rolling_kurtosis(options).into()
411
    }
412

413
    #[pyo3(signature = (lambda, window_size, weights, min_periods, center))]
414
    fn rolling_map(
415
        &self,
416
        lambda: Py<PyAny>,
417
        window_size: usize,
418
        weights: Option<Vec<f64>>,
419
        min_periods: Option<usize>,
420
        center: bool,
421
    ) -> Self {
422
        let min_periods = min_periods.unwrap_or(window_size);
423
        let options = RollingOptionsFixedWindow {
424
            window_size,
425
            weights,
426
            min_periods,
427
            center,
428
            ..Default::default()
429
        };
430
        let function = PlanCallback::new_python(PythonObject(lambda));
431

432
        self.inner.clone().rolling_map(function, options).into()
433
    }
434
}
435

436
Product

Resources

Company
