qust_ds/
roll.rs

1use crate::types::*;
2use itertools::Itertools;
3use num_traits::Pow;
4use serde::{Deserialize, Serialize};
5use std::iter::Sum;
6use std::ops::{Deref, Sub};
7use num_traits::cast::AsPrimitive;
8
9/* #region Agg Func  */
10
11pub trait AggFunc {
12    type T;
13    fn min(&self) -> Self::T;
14    fn max(&self) -> Self::T;
15    fn sum(&self) -> Self::T;
16    fn mean(&self) -> Self::T;
17    fn var(&self) -> Self::T
18    where
19        Self::T: Sub<Self::T, Output = Self::T> + num_traits::Pow<Self::T, Output = Self::T>;
20    fn std(&self) -> Self::T
21    where
22        Self::T:
23            Sub<Self::T, Output = Self::T> + num_traits::Pow<Self::T, Output = Self::T> + From<f32>;
24}
25
26impl<E> AggFunc for [E]
27where
28    E: PartialOrd + Copy + std::ops::Div<Output = E> + 'static,
29    for<'l> E: Sum<&'l E>,
30    usize: AsPrimitive<E>,
31{
32    type T = E;
33    fn min(&self) -> Self::T {
34        *self
35            .iter()
36            .min_by(|a, b| a.partial_cmp(b).unwrap())
37            .unwrap()
38    }
39    fn max(&self) -> Self::T {
40        *self
41            .iter()
42            .max_by(|a, b| a.partial_cmp(b).unwrap())
43            .unwrap()
44    }
45    fn sum(&self) -> Self::T {
46        self.iter().sum()
47    }
48    fn mean(&self) -> Self::T {
49        let len: E = self.len().as_();
50        self.sum() / len
51    }
52    fn var(&self) -> Self::T
53    where
54        Self::T: Sub<Self::T, Output = Self::T> + num_traits::Pow<Self::T, Output = Self::T>,
55    {
56        let m = self.mean();
57        let var_sum = self
58            .iter()
59            .map(|x| E::pow(*x - m, 2usize.as_()))
60            .collect::<Vec<E>>()
61            .sum();
62        var_sum / self.len().as_()
63    }
64    fn std(&self) -> Self::T
65    where
66        Self::T:
67            Sub<Self::T, Output = Self::T> + num_traits::Pow<Self::T, Output = Self::T> + From<f32>,
68    {
69        E::pow(self.var(), E::from(0.5f32))
70    }
71}
72
73pub trait RollApply<T, N> {
74    fn roll_func(&self, f: fn(x: &[T]) -> T, n: usize) -> Vec<N>;
75}
76
77impl RollApply<f32, f32> for [f32] {
78    fn roll_func(&self, f: fn(x: &[f32]) -> f32, n: usize) -> Vec<f32> {
79        let mut res = Vec::with_capacity(self.len());
80        for i in 0..self.len() {
81            let start_i = if i < n - 1 { 0 } else { i + 1 - n };
82            let res_ = f(&self[start_i..i + 1]);
83            res.push(res_);
84        }
85        res
86    }
87}
88
89impl<T, N> RollApply<T, Vec<T>> for Vec<N>
90where
91    [T]: RollApply<T, T>,
92    N: AsRef<[T]>,
93{
94    fn roll_func(&self, f: fn(x: &[T]) -> T, n: usize) -> Vec<Vec<T>> {
95        self.iter().map(|x| x.as_ref().roll_func(f, n)).collect()
96    }
97}
98
99fn _max(data: &[f32]) -> f32 {
100    *data
101        .iter()
102        .max_by(|a, b| a.partial_cmp(b).unwrap())
103        .unwrap()
104}
105
106pub trait Roll<T, N> {
107    fn roll_max(&self, n: usize) -> Vec<N>
108    where
109        Self: RollApply<T, N>,
110        [T]: AggFunc<T = T>,
111    {
112        self.roll_func(<[T] as AggFunc>::max, n)
113    }
114    fn roll_min(&self, n: usize) -> Vec<N>
115    where
116        Self: RollApply<T, N>,
117        [T]: AggFunc<T = T>,
118    {
119        self.roll_func(<[T] as AggFunc>::min, n)
120    }
121    fn roll_sum(&self, n: usize) -> Vec<N>
122    where
123        Self: RollApply<T, N>,
124        [T]: AggFunc<T = T>,
125    {
126        self.roll_func(<[T] as AggFunc>::sum, n)
127    }
128    fn roll_mean(&self, n: usize) -> Vec<N>
129    where
130        Self: RollApply<T, N>,
131        [T]: AggFunc<T = T>,
132    {
133        self.roll_func(<[T] as AggFunc>::mean, n)
134    }
135    fn roll_std(&self, n: usize) -> Vec<N>
136    where
137        Self: RollApply<T, N>,
138        [T]: AggFunc<T = T>,
139        T: Sub<T, Output = T> + num_traits::Pow<T, Output = T> + From<f32>,
140    {
141        self.roll_func(<[T] as AggFunc>::std, n)
142    }
143}
144
145impl Roll<f32, f32> for [f32] {}
146impl Roll<f32, v32> for vv32 {}
147impl Roll<f32, v32> for Vec<&[f32]> {}
148impl Roll<f32, v32> for Vec<&v32> {}
149/* #endregion */
150
151/* #region RollFunc */
152#[derive(Debug, Clone, Copy, PartialEq, Hash, Eq, Serialize, Deserialize)]
153pub enum RollFunc {
154    Sum,
155    Mean,
156    Min,
157    Max,
158    Var,
159    Std,
160    Momentum,
161    Skewness,
162}
163use RollFunc::*;
164
165pub trait AggFunc2 {
166    type T;
167    fn agg(&self, t: RollFunc) -> Self::T;
168}
169
170impl AggFunc2 for [f32] {
171    type T = f32;
172    fn agg(&self, t: RollFunc) -> Self::T {
173        match t {
174            Sum => self.iter().sum(),
175            Mean => self.agg(Sum) / (self.len() as f32),
176            Min => *self
177                .iter()
178                .min_by(|a, b| a.partial_cmp(b).unwrap())
179                .unwrap(),
180            Max => *self
181                .iter()
182                .max_by(|a, b| a.partial_cmp(b).unwrap())
183                .unwrap(),
184            Var => {
185                let m = self.agg(Mean);
186                let var_sum = self.iter().map(|x| f32::pow(*x - m, 2f32)).sum::<f32>();
187                var_sum / (self.len() as f32 - 1.)
188            }
189            Std => f32::pow(self.agg(Var), 0.5),
190            Momentum => {
191                if self.is_empty() {
192                    f32::NAN
193                } else {
194                    self[self.len() - 1] / self[0] - 1.
195                }
196            }
197            Skewness => {
198                if self.len() < 2 {
199                    f32::NAN
200                } else {
201                    let mut cm2 = 0f32;
202                    let mut cm3 = 0f32;
203                    let m = self.agg(Mean);
204                    for i in self {
205                        let z = i - m;
206                        let z2 = z * z;
207                        cm2 += z2;
208                        cm3 += z2 * z;
209                    }
210                    cm3 /= self.len() as f32;
211                    cm2 /= self.len() as f32;
212                    cm3 / f32::pow(cm2, 1.5f32)
213                }
214            }
215        }
216    }
217}
218
219/* #region RollCalc */
220#[derive(Debug, Clone, Serialize, Deserialize)]
221pub enum RollOps {
222    N(usize),
223    InitMiss(usize),
224    Vary(Box<Vec<usize>>)
225}
226impl AsRef<RollOps> for RollOps {
227    fn as_ref(&self) -> &RollOps {
228        self
229    }
230}
231impl From<usize> for RollOps {
232    fn from(value: usize) -> Self {
233        RollOps::N(value)
234    }
235}
236
237impl RollOps {
238    pub fn roll(&self, f: RollFunc, data: &[f32]) -> v32 {
239        match self {
240            RollOps::N(n) => {
241                let mut res = vec![f32::NAN; data.len()];
242                for i in 0..data.len() {
243                    let start_i = if i < *n { 0 } else { i + 1 - n };
244                    let data_part = &data[start_i..i + 1];
245                    res[i] = data_part.agg(f);
246                }
247                res
248            },
249            RollOps::InitMiss(n) => {
250                let mut res = RollOps::N(*n).roll(f, data);
251                res.iter_mut()
252                    .take(n - 1)
253                    .for_each(|x| *x = f32::NAN);
254                res
255            }
256            RollOps::Vary(v) => {
257                data.rolling(&**v).map(|x| x.agg(f)).collect_vec()
258            }
259        }
260    }
261}
262
263
264pub trait RollCalc<T> {
265    fn roll<N: AsRef<RollOps> + Clone>(&self, f: RollFunc, n: N) -> Vec<T>;
266}
267
268impl RollCalc<f32> for [f32] {
269    fn roll<N: AsRef<RollOps> + Clone>(&self, f: RollFunc, n: N) -> Vec<f32> {
270        n.as_ref().roll(f, self)
271    }
272}
273
274impl<T> RollCalc<v32> for Vec<&T>
275where
276    T: AsRef<[f32]> + RollCalc<f32> + ?Sized,
277{
278    fn roll<N: AsRef<RollOps> + Clone>(&self, f: RollFunc, n: N) -> vv32
279    {
280        self.deref()
281            .iter()
282            .map(|&x| x.as_ref().roll(f, n.clone()))
283            .collect_vec()
284    }
285}
286/* #endregion */
287
288/* #region RollSetp */
289#[derive(Debug, Clone, PartialEq, Hash, Eq, Serialize, Deserialize)]
290pub struct RollStep(pub usize, pub usize);
291
292impl RollStep {
293    pub fn roll(&self, data: &[f32], f: impl Fn(&[f32], usize) -> v32) -> v32 {
294        let mut res = Vec::with_capacity(data.len());
295        let slice_vec = self.get_slice(data.len());
296        for (start_loc, end_loc, start_i) in slice_vec {
297            let mut res_step = f(&data[start_loc..end_loc], start_i);
298            res.append(&mut res_step);
299        }
300        res
301    }
302
303    pub fn get_slice(&self, n: usize) -> Vec<(usize, usize, usize)> {
304        let window = self.0;
305        let step = self.1;
306        if n <= window {
307            vec![(0, n, 0)]
308        } else {
309            let partition_size = (n - window) / step + 1 + 1;
310            let mut res: Vec<(usize, usize, usize)> = Vec::with_capacity(partition_size);
311            for i in 0..partition_size {
312                let res_i = if i == 0 {
313                    (0, window, 0usize)
314                } else {
315                    let start_loc = i * step;
316                    let end_loc = if i == partition_size - 1 {
317                        n
318                    } else {
319                        start_loc + window
320                    };
321                    let start_i = res[i - 1].1 - start_loc;
322                    (start_loc, end_loc, start_i)
323                };
324                res.push(res_i);
325            }
326            res
327        }
328    }
329}
330
331/* #endregion */
332
333/* #region Rolling */
334use std::iter::Iterator;
335
336#[derive(Clone)]
337pub struct MovingWindow<'a, T, N> {
338    slice: &'a [T],
339    n: N,
340    count: usize,
341    init_size: usize
342}
343
344impl<'a, T> Iterator for MovingWindow<'a, T, usize> {
345    type Item = &'a [T];
346
347    #[inline]
348    fn next(&mut self) -> Option<Self::Item> {
349        if self.n > self.slice.len() {
350            if self.count < self.n.min(self.init_size + 1) {
351                let ret = Some(&self.slice[..self.count]);
352                self.count += 1;
353                ret
354            } else {
355                None
356            }
357        } else if self.count < self.n {
358            let ret = Some(&self.slice[..self.count]);
359            self.count += 1;
360            ret
361        } else {
362            let ret = Some(&self.slice[..self.n]);
363            self.slice = &self.slice[1..];
364            ret
365        }
366    }
367}
368
369impl<'a, 'b, T> Iterator for MovingWindow<'a, T, &'b Vec<usize>> {
370    type Item = &'a [T];
371
372    #[inline]
373    fn next(&mut self) -> Option<Self::Item> {
374        if self.count <= self.init_size {
375            let i = self.count;
376            let win = self.n[i - 1];
377            let start_i = if i < win { 0usize } else { i - win };
378            // println!("i: {i}, win: {win}, start_i: {start_i}");
379            let ret = Some(&self.slice[start_i..i]);
380            self.count += 1;
381            ret
382        } else {
383            None
384        }
385    }
386}
387
388pub trait Rolling {
389    type T;
390    fn rolling<N>(&self, n: N) -> MovingWindow<Self::T, N>;
391}
392
393impl<N> Rolling for [N] {
394    type T = N;
395    fn rolling<K>(&self, n: K) -> MovingWindow<Self::T, K>
396    {
397        MovingWindow { slice: self, n, count: 1, init_size: self.len()}
398    }
399}
400/* #endregion */
qust_ds/roll.rs

qust_ds/
roll.rs
