Skip to main content

rhai_sci/
moving.rs

1use rhai::plugin::*;
2
3#[export_module]
4pub mod moving_functions {
5    use crate::if_list_do;
6    use rhai::{Array, Dynamic, EvalAltResult, INT};
7
8    fn mov<G>(arr: &mut Array, k: INT, mut f: G) -> Result<Array, Box<EvalAltResult>>
9    where
10        G: FnMut(&mut Array) -> Dynamic,
11    {
12        if_list_do(arr, |arr| {
13            // First, validate the inputs
14            let mut new_arr = vec![];
15            let n = arr.len() as INT;
16            for i in 0..n {
17                new_arr.push(f(&mut arr
18                    .get(if k % 2 != 0 {
19                        (std::cmp::max(i - (k - 1) / 2, 0) as usize)
20                            ..=(std::cmp::min(i + (k - 1) / 2, n - 1) as usize)
21                    } else {
22                        (std::cmp::max(i - k / 2, 0) as usize)
23                            ..=(std::cmp::min(i + k / 2 - 1, n - 1) as usize)
24                    })
25                    .unwrap()
26                    .to_vec()))
27            }
28            Ok(new_arr)
29        })
30    }
31
32    /// Returns an array of the moving minimum (with a given width) across the input array.
33    /// ```typescript
34    /// let data = [1, 2, 4, -1, -2, -3, -1, 3, 2, 1];
35    /// let m = movmin(data, 3);
36    /// assert_eq(m, [1, 1, -1, -2, -3, -3, -3, -1, 1, 1]);
37    /// ```
38    #[rhai_fn(name = "movmin", return_raw, pure)]
39    pub fn movmin(arr: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
40        mov(arr, k, |x| crate::stats::array_min(x).unwrap())
41    }
42
43    /// Returns an array of the moving maximum (with a given width) across the input array.
44    /// ```typescript
45    /// let data = [1, 2, 4, -1, -2, -3, -1, 3, 2, 1];
46    /// let m = movmax(data, 3);
47    /// assert_eq(m, [2, 4, 4, 4, -1, -1, 3, 3, 3, 2]);
48    /// ```
49    #[rhai_fn(name = "movmax", return_raw, pure)]
50    pub fn movmax(arr: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
51        mov(arr, k, |x| crate::stats::array_max(x).unwrap())
52    }
53
54    /// Returns an array of the moving maximum absolute deviation (with a given width) across the input array.
55    /// ```typescript
56    /// let data = [1, 2, 4, -1, -2, -3, -1, 3, 2, 1];
57    /// let m = movmad(data, 3);
58    /// assert_eq(m, [0.5, 1.0, 2.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 0.5]);
59    /// ```
60    #[rhai_fn(name = "movmad", return_raw, pure)]
61    pub fn movmad(arr: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
62        mov(arr, k, |x| crate::stats::mad(x).unwrap())
63    }
64
65    /// Returns an array of the moving average (with a given width) across the input array.
66    /// ```typescript
67    /// let data = [1, 2, 3, 4, 5, 6];
68    /// let m = movmean(data, 3);
69    /// assert_eq(m, [1.5, 2.0, 3.0, 4.0, 5.0, 5.5]);
70    /// ```
71    #[rhai_fn(name = "movmean", return_raw, pure)]
72    pub fn movmean(arr: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
73        mov(arr, k, |x| crate::stats::mean(x).unwrap())
74    }
75
76    /// Returns an array of the moving median (with a given width) across the input array.
77    /// ```typescript
78    /// let data = [1, 2, 3, 4, 5, 6];
79    /// let m = movmedian(data, 3);
80    /// assert_eq(m, [1.5, 2.0, 3.0, 4.0, 5.0, 5.5]);
81    /// ```
82    #[rhai_fn(name = "movmedian", return_raw, pure)]
83    pub fn movmedian(arr: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
84        mov(arr, k, |x| crate::stats::median(x).unwrap())
85    }
86
87    /// Returns an array of the moving product (with a given width) across the input array.
88    /// ```typescript
89    /// let data = [1, 2, 3, 4, 5, 6];
90    /// let m = movprod(data, 3);
91    /// assert_eq(m, [2, 6, 24, 60, 120, 30]);
92    /// ```
93    #[rhai_fn(name = "movprod", return_raw, pure)]
94    pub fn movprod(arr: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
95        mov(arr, k, |x| crate::stats::prod(x).unwrap())
96    }
97
98    /// Returns an array of the moving standard deviation (with a given width) across the input array.
99    /// ```typescript
100    /// let data = [1, 2, 3, 4, 5, 6];
101    /// let m = movstd(data, 3);
102    /// assert_eq(m, [0.7071067811865476, 1.0, 1.0, 1.0, 1.0, 0.7071067811865476]);
103    /// ```
104    #[rhai_fn(name = "movstd", return_raw, pure)]
105    pub fn movstd(arr: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
106        mov(arr, k, |x| crate::stats::std(x).unwrap())
107    }
108
109    /// Returns an array of the moving variance (with a given width) across the input array.
110    /// ```typescript
111    /// let data = [1, 2, 3, 4, 5, 6];
112    /// let m = movvar(data, 3);
113    /// assert_eq(m, [0.5, 1.0, 1.0, 1.0, 1.0, 0.5]);
114    /// ```
115    #[rhai_fn(name = "movvar", return_raw, pure)]
116    pub fn movvar(arr: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
117        mov(arr, k, |x| crate::stats::variance(x).unwrap())
118    }
119
120    /// Returns an array of the moving sum (with a given width) across the input array.
121    /// ```typescript
122    /// let data = [1, 2, 3, 4, 5, 6];
123    /// let m = movsum(data, 3);
124    /// assert_eq(m, [3, 6, 9, 12, 15, 11]);
125    /// ```
126    #[rhai_fn(name = "movsum", return_raw, pure)]
127    pub fn movsum(arr: &mut Array, k: INT) -> Result<Array, Box<EvalAltResult>> {
128        mov(arr, k, |x| crate::stats::sum(x).unwrap())
129    }
130}