bixverse-rs 0.6.2

//! Various utility functions for vectors: flattening them, min-max values,
//! cumulative sums and scaling them.

use num_traits::Float;
use rand::SeedableRng;
use rand::prelude::SliceRandom;
use rand::rngs::StdRng;
use rustc_hash::FxHashSet;
use std::fmt::Debug;
use std::hash::Hash;
use std::ops::AddAssign;

use crate::core::math::vector_helpers::*;
use crate::prelude::*;

////////////////////
// Vector helpers //
////////////////////

/// Split `0..n` into one contiguous range per rayon worker.
///
/// For fan-outs where each worker carries scratch proportional to the whole
/// problem, one chunk per worker is what keeps the number of allocations and
/// merges down to the thread count. Rayon's own splitting is the better choice
/// whenever the per-item work is self-contained.
///
/// ### Params
///
/// * `n` - Length of the range to split
///
/// ### Returns
///
/// Half-open `(start, end)` ranges covering `0..n`. Empty when `n` is zero.
pub fn thread_chunks(n: usize) -> Vec<(usize, usize)> {
    let n_threads = rayon::current_num_threads().max(1);
    let chunk = n.div_ceil(n_threads).max(1);

    (0..n)
        .step_by(chunk)
        .map(|start| (start, (start + chunk).min(n)))
        .collect()
}

/// Flatten a nested vector
///
/// ### Params
///
/// * `vec` - The vector to flatten
///
/// ### Returns
///
/// The flattened vector
pub fn flatten_vector<I, T>(vec: I) -> Vec<T>
where
    I: IntoIterator,
    I::Item: IntoIterator<Item = T>,
{
    vec.into_iter().flatten().collect()
}

/// Get the maximum value of an array
///
/// ### Params
///
/// * `arr` - The array of values
///
/// ### Returns
///
/// The maximum value found in the array
pub fn array_max<T>(arr: &[T]) -> T
where
    T: PartialOrd + Copy,
{
    let mut max_val = arr[0];
    for number in arr {
        if *number > max_val {
            max_val = *number
        }
    }
    max_val
}

/// Get the minimum value of an array
///
/// ### Params
///
/// * `arr` - The array of values
///
/// ### Returns
///
/// The minimum value found in the array
pub fn array_min<T>(arr: &[T]) -> T
where
    T: PartialOrd + Copy,
{
    let mut min_val = arr[0];
    for number in arr {
        if *number < min_val {
            min_val = *number
        }
    }
    min_val
}

/// Get the maximum and minimum value of an array
///
/// ### Params
///
/// * `arr` - The array of values
///
/// ### Returns
///
/// Tuple of `(max, min)` in the array
pub fn array_max_min<T>(arr: &[T]) -> (T, T)
where
    T: PartialOrd + Copy,
{
    let mut min_val = arr[0];
    let mut max_val = arr[0];
    for number in arr {
        if *number < min_val {
            min_val = *number
        }
        if *number > max_val {
            max_val = *number
        }
    }

    (max_val, min_val)
}

/// Get unique elements from a slice of any hashable, equatable numeric type.
///
/// ### Params
///
/// * `vec` - The slice of numerical values.
///
/// ### Returns
///
/// The unique elements of `vec` as a Vec.
pub fn unique<T>(vec: &[T]) -> Vec<T>
where
    T: Copy + Eq + Hash + Debug,
{
    let mut set = FxHashSet::default();
    vec.iter()
        .filter(|&&item| set.insert(item))
        .cloned()
        .collect()
}

/// Calculate the cumulative sum over a vector
///
/// ### Params
///
/// * `x` - The slice of numerical values
///
/// ### Returns
///
/// The cumulative sum over the vector.
pub fn cumsum<T>(x: &[T]) -> Vec<T>
where
    T: Float + AddAssign<T>,
{
    let mut sum = T::zero();
    x.iter()
        .map(|&val| {
            sum += val;
            sum
        })
        .collect()
}

/// Split a vector randomly into two chunks
///
/// Splits a vector randomly into two of [..x] and the other [x..]
///
/// ### Params
///
/// * `vec` - Slice of the vector you want to split
/// * `x` - Length of the first vector; the rest will be put into the second vector
/// * `seed` - Seed for reproducibility
///
/// ### Returns
///
/// A tuple of the pieces of the vector
pub fn split_vector_randomly<T>(vec: &[T], x: usize, seed: u64) -> (Vec<T>, Vec<T>)
where
    T: Clone,
{
    let mut rng = StdRng::seed_from_u64(seed);
    let mut shuffled = vec.to_vec();
    shuffled.shuffle(&mut rng);

    let (first_set, second_set) = shuffled.split_at(x);

    (first_set.to_vec(), second_set.to_vec())
}

/////////////
// Scalers //
/////////////

/// Min-max scale a vector to the range [0, 1].
///
/// ### Params
///
/// * `vec` - The slice of numerical values.
///
/// ### Returns
///
/// The scaled vector. Returns a zero vector if all values are identical.
pub fn min_max_scale<T>(vec: &[T]) -> Vec<T>
where
    T: Float,
{
    if vec.is_empty() {
        return Vec::new();
    }
    let (max_val, min_val) = array_max_min(vec);
    let range = max_val - min_val;
    if range == T::zero() {
        return vec![T::zero(); vec.len()];
    }
    vec.iter().map(|&x| (x - min_val) / range).collect()
}

/// Rank-normalise a vector by ranking values and dividing by n.
///
/// Produces values in (0, 1] with tie correction via averaged ranks.
///
/// ### Params
///
/// * `vec` - The slice of numerical values.
///
/// ### Returns
///
/// The rank-normalised vector.
pub fn rank_normalise<T>(vec: &[T]) -> Vec<T>
where
    T: Float,
{
    if vec.is_empty() {
        return Vec::new();
    }
    let n = T::from(vec.len()).unwrap();
    rank_vector(vec).into_iter().map(|r| r / n).collect()
}

/// Z-score normalise a vector (zero mean, unit variance).
///
/// ### Params
///
/// * `vec` - The slice of numerical values.
///
/// ### Returns
///
/// The normalised vector. Returns a zero vector if standard deviation is zero.
pub fn zscore_normalise<T>(vec: &[T]) -> Vec<T>
where
    T: Float,
{
    if vec.is_empty() {
        return Vec::new();
    }
    let n = T::from(vec.len()).unwrap();
    let mean = vec.iter().copied().fold(T::zero(), |acc, x| acc + x) / n;
    let variance = vec.iter().copied().fold(T::zero(), |acc, x| {
        let d = x - mean;
        acc + d * d
    }) / n;
    let std_dev = variance.sqrt();
    if std_dev == T::zero() {
        return vec![T::zero(); vec.len()];
    }
    vec.iter().map(|&x| (x - mean) / std_dev).collect()
}

/// Robust-scale a vector using median and MAD.
///
/// Computes `(x - median) / MAD` for each element. More resistant to outliers
/// than z-score normalisation.
///
/// ### Params
///
/// * `vec` - The slice of numerical values.
///
/// ### Returns
///
/// The scaled vector, or `None` if the input is empty. Returns a zero vector
/// if MAD is zero.
pub fn robust_scale<T>(vec: &[T]) -> Option<Vec<T>>
where
    T: BixverseFloat,
{
    if vec.is_empty() {
        return None;
    }
    let median_val = median(vec)?;
    let mad_val = mad(vec, Some(T::from_f64(MAD_SCALE).unwrap()))?;
    if mad_val == T::zero() {
        return Some(vec![T::zero(); vec.len()]);
    }
    Some(vec.iter().map(|&x| (x - median_val) / mad_val).collect())
}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    use super::*;

    const EPS: f64 = 1e-10;

    fn approx_eq(a: f64, b: f64) -> bool {
        (a - b).abs() < EPS
    }

    fn vec_approx_eq(a: &[f64], b: &[f64]) -> bool {
        a.len() == b.len() && a.iter().zip(b).all(|(&x, &y)| approx_eq(x, y))
    }

    /// Order survives the nesting and empty inner vectors contribute nothing.
    #[test]
    fn test_flatten_vector() {
        let nested = vec![vec![1, 2], vec![3], vec![4, 5]];
        assert_eq!(flatten_vector(nested), vec![1, 2, 3, 4, 5]);

        let with_empties: Vec<Vec<i32>> = vec![vec![], vec![1], vec![]];
        assert_eq!(flatten_vector(with_empties), vec![1]);
    }

    /// The maximum, including the single-element slice the loop seeds from.
    #[test]
    fn test_array_max() {
        assert_eq!(array_max(&[3, 1, 4, 1, 5, 9, 2, 6]), 9);
        assert_eq!(array_max(&[42]), 42);
    }

    /// The minimum, including the single-element slice the loop seeds from.
    #[test]
    fn test_array_min() {
        assert_eq!(array_min(&[3, 1, 4, 1, 5, 9, 2, 6]), 1);
        assert_eq!(array_min(&[42]), 42);
    }

    /// The fused scan returns `(max, min)` in that order, which is easy to swap.
    #[test]
    fn test_array_max_min() {
        let (max, min) = array_max_min(&[3.0_f64, 1.0, 4.0, 1.0, 5.0, 9.0]);
        assert!(approx_eq(max, 9.0));
        assert!(approx_eq(min, 1.0));

        let (max, min) = array_max_min(&[7]);
        assert_eq!(max, 7);
        assert_eq!(min, 7);
    }

    /// Repeats collapse to one entry and every distinct value survives.
    #[test]
    fn test_unique() {
        let mut result = unique(&[1, 2, 2, 3, 3, 3]);
        result.sort();
        assert_eq!(result, vec![1, 2, 3]);

        let mut result = unique(&[4, 2, 1]);
        result.sort();
        assert_eq!(result, vec![1, 2, 4]);
    }

    /// The running sum is inclusive, so the last entry is the total.
    #[test]
    fn test_cumsum() {
        let result = cumsum(&[1.0_f64, 2.0, 3.0, 4.0]);
        assert!(vec_approx_eq(&result, &[1.0, 3.0, 6.0, 10.0]));

        let result = cumsum(&[5.0_f64]);
        assert!(vec_approx_eq(&result, &[5.0]));
    }

    /// The split honours the requested length, loses nothing, and repeats per seed.
    #[test]
    fn test_split_vector_randomly() {
        let vec: Vec<i32> = (0..10).collect();

        let (a, b) = split_vector_randomly(&vec, 4, 42);
        assert_eq!(a.len(), 4);
        assert_eq!(b.len(), 6);

        // The split is a partition: nothing lost, nothing duplicated.
        let (mut a, mut b) = split_vector_randomly(&vec, 4, 42);
        a.append(&mut b);
        a.sort();
        assert_eq!(a, vec);

        let (a1, _) = split_vector_randomly(&vec, 4, 99);
        let (a2, _) = split_vector_randomly(&vec, 4, 99);
        assert_eq!(a1, a2);
    }

    /// The extremes land on exactly 0 and 1 with everything between spaced linearly.
    #[test]
    fn test_min_max_scale() {
        let result = min_max_scale(&[0.0_f64, 1.0, 2.0, 3.0, 4.0]);
        assert!(vec_approx_eq(&result, &[0.0, 0.25, 0.5, 0.75, 1.0]));

        // Zero range collapses to zeroes rather than dividing by zero.
        let result = min_max_scale(&[3.0_f64, 3.0, 3.0]);
        assert!(vec_approx_eq(&result, &[0.0, 0.0, 0.0]));
    }

    /// Ranks divided by n, with tied values taking the averaged rank.
    #[test]
    fn test_rank_normalise() {
        // ranks: 1, 3, 2 -> divided by 3
        let result = rank_normalise(&[10.0_f64, 30.0, 20.0]);
        assert!(vec_approx_eq(&result, &[1.0 / 3.0, 1.0, 2.0 / 3.0]));

        // ranks: avg(1,2)=1.5, 1.5, 3 -> divided by 3
        let result = rank_normalise(&[1.0_f64, 1.0, 3.0]);
        assert!(vec_approx_eq(&result, &[0.5, 0.5, 1.0]));
    }

    /// Checked as zero mean and unit variance rather than against fixed values.
    #[test]
    fn test_zscore_normalise() {
        let result = zscore_normalise(&[2.0_f64, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]);
        let mean: f64 = result.iter().sum::<f64>() / result.len() as f64;
        let variance: f64 =
            result.iter().map(|&x| (x - mean).powi(2)).sum::<f64>() / result.len() as f64;
        assert!(approx_eq(mean, 0.0));
        assert!(approx_eq(variance, 1.0));

        // Zero standard deviation collapses to zeroes.
        let result = zscore_normalise(&[5.0_f64, 5.0, 5.0]);
        assert!(vec_approx_eq(&result, &[0.0, 0.0, 0.0]));
    }

    /// Median-and-MAD scaling, including the zero-MAD and empty-input branches.
    #[test]
    fn test_robust_scale() {
        // values [1, 3, 2, 5, 4] -> median 3, raw MAD median(|2,0,1,2,1|) = 1,
        // scaled MAD = 1 * MAD_SCALE
        let result = robust_scale(&[1.0_f64, 3.0, 2.0, 5.0, 4.0]).unwrap();
        assert!(vec_approx_eq(
            &result,
            &[
                -1.3489795003921634,
                0.0,
                -0.6744897501960817,
                1.3489795003921634,
                0.6744897501960817
            ]
        ));

        let result = robust_scale(&[4.0_f64, 4.0, 4.0]).unwrap();
        assert!(vec_approx_eq(&result, &[0.0, 0.0, 0.0]));

        assert!(robust_scale::<f64>(&[]).is_none());
    }
}