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

bixverse_core/utils/
vec_utils.rs
