Skip to main content

bixverse_core/utils/
vec_utils.rs

1//! Various utility functions for vectors: flattening them, min-max values,
2//! cumulative sums and scaling them.
3
4use num_traits::Float;
5use rand::SeedableRng;
6use rand::prelude::SliceRandom;
7use rand::rngs::StdRng;
8use rustc_hash::FxHashSet;
9use std::fmt::Debug;
10use std::hash::Hash;
11use std::ops::AddAssign;
12
13use crate::core::math::vector_helpers::*;
14use crate::prelude::*;
15
16////////////////////
17// Vector helpers //
18////////////////////
19
20pub use bixverse_sparse::thread_chunks;
21
22/// Flatten a nested vector
23///
24/// ### Params
25///
26/// * `vec` - The vector to flatten
27///
28/// ### Returns
29///
30/// The flattened vector
31pub fn flatten_vector<I, T>(vec: I) -> Vec<T>
32where
33    I: IntoIterator,
34    I::Item: IntoIterator<Item = T>,
35{
36    vec.into_iter().flatten().collect()
37}
38
39/// Get the maximum value of an array
40///
41/// ### Params
42///
43/// * `arr` - The array of values
44///
45/// ### Returns
46///
47/// The maximum value found in the array
48pub fn array_max<T>(arr: &[T]) -> T
49where
50    T: PartialOrd + Copy,
51{
52    let mut max_val = arr[0];
53    for number in arr {
54        if *number > max_val {
55            max_val = *number
56        }
57    }
58    max_val
59}
60
61/// Get the minimum value of an array
62///
63/// ### Params
64///
65/// * `arr` - The array of values
66///
67/// ### Returns
68///
69/// The minimum value found in the array
70pub fn array_min<T>(arr: &[T]) -> T
71where
72    T: PartialOrd + Copy,
73{
74    let mut min_val = arr[0];
75    for number in arr {
76        if *number < min_val {
77            min_val = *number
78        }
79    }
80    min_val
81}
82
83/// Get the maximum and minimum value of an array
84///
85/// ### Params
86///
87/// * `arr` - The array of values
88///
89/// ### Returns
90///
91/// Tuple of `(max, min)` in the array
92pub fn array_max_min<T>(arr: &[T]) -> (T, T)
93where
94    T: PartialOrd + Copy,
95{
96    let mut min_val = arr[0];
97    let mut max_val = arr[0];
98    for number in arr {
99        if *number < min_val {
100            min_val = *number
101        }
102        if *number > max_val {
103            max_val = *number
104        }
105    }
106
107    (max_val, min_val)
108}
109
110/// Get unique elements from a slice of any hashable, equatable numeric type.
111///
112/// ### Params
113///
114/// * `vec` - The slice of numerical values.
115///
116/// ### Returns
117///
118/// The unique elements of `vec` as a Vec.
119pub fn unique<T>(vec: &[T]) -> Vec<T>
120where
121    T: Copy + Eq + Hash + Debug,
122{
123    let mut set = FxHashSet::default();
124    vec.iter()
125        .filter(|&&item| set.insert(item))
126        .cloned()
127        .collect()
128}
129
130/// Calculate the cumulative sum over a vector
131///
132/// ### Params
133///
134/// * `x` - The slice of numerical values
135///
136/// ### Returns
137///
138/// The cumulative sum over the vector.
139pub fn cumsum<T>(x: &[T]) -> Vec<T>
140where
141    T: Float + AddAssign<T>,
142{
143    let mut sum = T::zero();
144    x.iter()
145        .map(|&val| {
146            sum += val;
147            sum
148        })
149        .collect()
150}
151
152/// Split a vector randomly into two chunks
153///
154/// Splits a vector randomly into two of [..x] and the other [x..]
155///
156/// ### Params
157///
158/// * `vec` - Slice of the vector you want to split
159/// * `x` - Length of the first vector; the rest will be put into the second vector
160/// * `seed` - Seed for reproducibility
161///
162/// ### Returns
163///
164/// A tuple of the pieces of the vector
165pub fn split_vector_randomly<T>(vec: &[T], x: usize, seed: u64) -> (Vec<T>, Vec<T>)
166where
167    T: Clone,
168{
169    let mut rng = StdRng::seed_from_u64(seed);
170    let mut shuffled = vec.to_vec();
171    shuffled.shuffle(&mut rng);
172
173    let (first_set, second_set) = shuffled.split_at(x);
174
175    (first_set.to_vec(), second_set.to_vec())
176}
177
178/////////////
179// Scalers //
180/////////////
181
182/// Min-max scale a vector to the range [0, 1].
183///
184/// ### Params
185///
186/// * `vec` - The slice of numerical values.
187///
188/// ### Returns
189///
190/// The scaled vector. Returns a zero vector if all values are identical.
191pub fn min_max_scale<T>(vec: &[T]) -> Vec<T>
192where
193    T: Float,
194{
195    if vec.is_empty() {
196        return Vec::new();
197    }
198    let (max_val, min_val) = array_max_min(vec);
199    let range = max_val - min_val;
200    if range == T::zero() {
201        return vec![T::zero(); vec.len()];
202    }
203    vec.iter().map(|&x| (x - min_val) / range).collect()
204}
205
206/// Rank-normalise a vector by ranking values and dividing by n.
207///
208/// Produces values in (0, 1] with tie correction via averaged ranks.
209///
210/// ### Params
211///
212/// * `vec` - The slice of numerical values.
213///
214/// ### Returns
215///
216/// The rank-normalised vector.
217pub fn rank_normalise<T>(vec: &[T]) -> Vec<T>
218where
219    T: Float,
220{
221    if vec.is_empty() {
222        return Vec::new();
223    }
224    let n = T::from(vec.len()).unwrap();
225    rank_vector(vec).into_iter().map(|r| r / n).collect()
226}
227
228/// Z-score normalise a vector (zero mean, unit variance).
229///
230/// ### Params
231///
232/// * `vec` - The slice of numerical values.
233///
234/// ### Returns
235///
236/// The normalised vector. Returns a zero vector if standard deviation is zero.
237pub fn zscore_normalise<T>(vec: &[T]) -> Vec<T>
238where
239    T: Float,
240{
241    if vec.is_empty() {
242        return Vec::new();
243    }
244    let n = T::from(vec.len()).unwrap();
245    let mean = vec.iter().copied().fold(T::zero(), |acc, x| acc + x) / n;
246    let variance = vec.iter().copied().fold(T::zero(), |acc, x| {
247        let d = x - mean;
248        acc + d * d
249    }) / n;
250    let std_dev = variance.sqrt();
251    if std_dev == T::zero() {
252        return vec![T::zero(); vec.len()];
253    }
254    vec.iter().map(|&x| (x - mean) / std_dev).collect()
255}
256
257/// Robust-scale a vector using median and MAD.
258///
259/// Computes `(x - median) / MAD` for each element. More resistant to outliers
260/// than z-score normalisation.
261///
262/// ### Params
263///
264/// * `vec` - The slice of numerical values.
265///
266/// ### Returns
267///
268/// The scaled vector, or `None` if the input is empty. Returns a zero vector
269/// if MAD is zero.
270pub fn robust_scale<T>(vec: &[T]) -> Option<Vec<T>>
271where
272    T: BixverseFloat,
273{
274    if vec.is_empty() {
275        return None;
276    }
277    let median_val = median(vec)?;
278    let mad_val = mad(vec, Some(T::from_f64(MAD_SCALE).unwrap()))?;
279    if mad_val == T::zero() {
280        return Some(vec![T::zero(); vec.len()]);
281    }
282    Some(vec.iter().map(|&x| (x - median_val) / mad_val).collect())
283}
284
285///////////
286// Tests //
287///////////
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292
293    const EPS: f64 = 1e-10;
294
295    fn approx_eq(a: f64, b: f64) -> bool {
296        (a - b).abs() < EPS
297    }
298
299    fn vec_approx_eq(a: &[f64], b: &[f64]) -> bool {
300        a.len() == b.len() && a.iter().zip(b).all(|(&x, &y)| approx_eq(x, y))
301    }
302
303    /// Order survives the nesting and empty inner vectors contribute nothing.
304    #[test]
305    fn test_flatten_vector() {
306        let nested = vec![vec![1, 2], vec![3], vec![4, 5]];
307        assert_eq!(flatten_vector(nested), vec![1, 2, 3, 4, 5]);
308
309        let with_empties: Vec<Vec<i32>> = vec![vec![], vec![1], vec![]];
310        assert_eq!(flatten_vector(with_empties), vec![1]);
311    }
312
313    /// The maximum, including the single-element slice the loop seeds from.
314    #[test]
315    fn test_array_max() {
316        assert_eq!(array_max(&[3, 1, 4, 1, 5, 9, 2, 6]), 9);
317        assert_eq!(array_max(&[42]), 42);
318    }
319
320    /// The minimum, including the single-element slice the loop seeds from.
321    #[test]
322    fn test_array_min() {
323        assert_eq!(array_min(&[3, 1, 4, 1, 5, 9, 2, 6]), 1);
324        assert_eq!(array_min(&[42]), 42);
325    }
326
327    /// The fused scan returns `(max, min)` in that order, which is easy to swap.
328    #[test]
329    fn test_array_max_min() {
330        let (max, min) = array_max_min(&[3.0_f64, 1.0, 4.0, 1.0, 5.0, 9.0]);
331        assert!(approx_eq(max, 9.0));
332        assert!(approx_eq(min, 1.0));
333
334        let (max, min) = array_max_min(&[7]);
335        assert_eq!(max, 7);
336        assert_eq!(min, 7);
337    }
338
339    /// Repeats collapse to one entry and every distinct value survives.
340    #[test]
341    fn test_unique() {
342        let mut result = unique(&[1, 2, 2, 3, 3, 3]);
343        result.sort();
344        assert_eq!(result, vec![1, 2, 3]);
345
346        let mut result = unique(&[4, 2, 1]);
347        result.sort();
348        assert_eq!(result, vec![1, 2, 4]);
349    }
350
351    /// The running sum is inclusive, so the last entry is the total.
352    #[test]
353    fn test_cumsum() {
354        let result = cumsum(&[1.0_f64, 2.0, 3.0, 4.0]);
355        assert!(vec_approx_eq(&result, &[1.0, 3.0, 6.0, 10.0]));
356
357        let result = cumsum(&[5.0_f64]);
358        assert!(vec_approx_eq(&result, &[5.0]));
359    }
360
361    /// The split honours the requested length, loses nothing, and repeats per seed.
362    #[test]
363    fn test_split_vector_randomly() {
364        let vec: Vec<i32> = (0..10).collect();
365
366        let (a, b) = split_vector_randomly(&vec, 4, 42);
367        assert_eq!(a.len(), 4);
368        assert_eq!(b.len(), 6);
369
370        // The split is a partition: nothing lost, nothing duplicated.
371        let (mut a, mut b) = split_vector_randomly(&vec, 4, 42);
372        a.append(&mut b);
373        a.sort();
374        assert_eq!(a, vec);
375
376        let (a1, _) = split_vector_randomly(&vec, 4, 99);
377        let (a2, _) = split_vector_randomly(&vec, 4, 99);
378        assert_eq!(a1, a2);
379    }
380
381    /// The extremes land on exactly 0 and 1 with everything between spaced linearly.
382    #[test]
383    fn test_min_max_scale() {
384        let result = min_max_scale(&[0.0_f64, 1.0, 2.0, 3.0, 4.0]);
385        assert!(vec_approx_eq(&result, &[0.0, 0.25, 0.5, 0.75, 1.0]));
386
387        // Zero range collapses to zeroes rather than dividing by zero.
388        let result = min_max_scale(&[3.0_f64, 3.0, 3.0]);
389        assert!(vec_approx_eq(&result, &[0.0, 0.0, 0.0]));
390    }
391
392    /// Ranks divided by n, with tied values taking the averaged rank.
393    #[test]
394    fn test_rank_normalise() {
395        // ranks: 1, 3, 2 -> divided by 3
396        let result = rank_normalise(&[10.0_f64, 30.0, 20.0]);
397        assert!(vec_approx_eq(&result, &[1.0 / 3.0, 1.0, 2.0 / 3.0]));
398
399        // ranks: avg(1,2)=1.5, 1.5, 3 -> divided by 3
400        let result = rank_normalise(&[1.0_f64, 1.0, 3.0]);
401        assert!(vec_approx_eq(&result, &[0.5, 0.5, 1.0]));
402    }
403
404    /// Checked as zero mean and unit variance rather than against fixed values.
405    #[test]
406    fn test_zscore_normalise() {
407        let result = zscore_normalise(&[2.0_f64, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]);
408        let mean: f64 = result.iter().sum::<f64>() / result.len() as f64;
409        let variance: f64 =
410            result.iter().map(|&x| (x - mean).powi(2)).sum::<f64>() / result.len() as f64;
411        assert!(approx_eq(mean, 0.0));
412        assert!(approx_eq(variance, 1.0));
413
414        // Zero standard deviation collapses to zeroes.
415        let result = zscore_normalise(&[5.0_f64, 5.0, 5.0]);
416        assert!(vec_approx_eq(&result, &[0.0, 0.0, 0.0]));
417    }
418
419    /// Median-and-MAD scaling, including the zero-MAD and empty-input branches.
420    #[test]
421    fn test_robust_scale() {
422        // values [1, 3, 2, 5, 4] -> median 3, raw MAD median(|2,0,1,2,1|) = 1,
423        // scaled MAD = 1 * MAD_SCALE
424        let result = robust_scale(&[1.0_f64, 3.0, 2.0, 5.0, 4.0]).unwrap();
425        assert!(vec_approx_eq(
426            &result,
427            &[
428                -1.3489795003921634,
429                0.0,
430                -0.6744897501960817,
431                1.3489795003921634,
432                0.6744897501960817
433            ]
434        ));
435
436        let result = robust_scale(&[4.0_f64, 4.0, 4.0]).unwrap();
437        assert!(vec_approx_eq(&result, &[0.0, 0.0, 0.0]));
438
439        assert!(robust_scale::<f64>(&[]).is_none());
440    }
441}