bixverse-rs 0.6.3

//! Information theory-based measures

use faer::{ColRef, Mat, MatRef};
use rayon::prelude::*;

use crate::prelude::BixverseFloat;

///////////////
// MiScratch //
///////////////

/// Reusable count tables for the pairwise mutual information kernels
pub(crate) struct MiScratch {
    joint: Vec<usize>,
    marginal_i: Vec<usize>,
    marginal_j: Vec<usize>,
    n_bins: usize,
}

impl MiScratch {
    /// Allocate tables for `n_bins` bins
    ///
    /// ### Params
    ///
    /// * `n_bins` - Number of bins per variable
    ///
    /// ### Returns
    ///
    /// Zeroed scratch tables
    pub(crate) fn new(n_bins: usize) -> Self {
        Self {
            joint: vec![0; n_bins * n_bins],
            marginal_i: vec![0; n_bins],
            marginal_j: vec![0; n_bins],
            n_bins,
        }
    }

    /// Tabulate the joint and marginal counts of two binned columns
    fn count(&mut self, col_i: ColRef<usize>, col_j: ColRef<usize>) {
        self.joint.fill(0);
        self.marginal_i.fill(0);
        self.marginal_j.fill(0);
        for (&bin_i, &bin_j) in col_i.iter().zip(col_j.iter()) {
            self.joint[bin_i * self.n_bins + bin_j] += 1;
            self.marginal_i[bin_i] += 1;
            self.marginal_j[bin_j] += 1;
        }
    }

    /// Mutual information from the current tables
    fn mutual_information<T: BixverseFloat>(&self, n_rows: usize) -> T {
        let n = T::from_usize(n_rows).unwrap();
        let prob_i: Vec<T> = self
            .marginal_i
            .iter()
            .map(|&c| T::from_usize(c).unwrap() / n)
            .collect();
        let prob_j: Vec<T> = self
            .marginal_j
            .iter()
            .map(|&c| T::from_usize(c).unwrap() / n)
            .collect();
        let mut mi = T::zero();
        for i in 0..self.n_bins {
            for j in 0..self.n_bins {
                let joint_prob = T::from_usize(self.joint[i * self.n_bins + j]).unwrap() / n;
                if joint_prob > T::zero() {
                    mi += joint_prob * (joint_prob / (prob_i[i] * prob_j[j])).ln();
                }
            }
        }
        mi
    }

    /// Joint entropy from the current tables
    fn joint_entropy<T: BixverseFloat>(&self, n_rows: usize) -> T {
        let n = T::from_usize(n_rows).unwrap();
        let mut joint_entropy = T::zero();
        for &count in &self.joint {
            let joint_prob = T::from_usize(count).unwrap() / n;
            if joint_prob > T::zero() {
                joint_entropy -= joint_prob * joint_prob.ln();
            }
        }
        joint_entropy
    }

    /// Mutual information, and optionally the joint entropy, of a column pair
    ///
    /// ### Params
    ///
    /// * `col_i` - First binned column
    /// * `col_j` - Second binned column
    /// * `with_joint_entropy` - Also compute the joint entropy
    ///
    /// ### Returns
    ///
    /// `(mutual information, joint entropy)`; the entropy is zero when not
    /// requested.
    pub(crate) fn pair<T: BixverseFloat>(
        &mut self,
        col_i: ColRef<usize>,
        col_j: ColRef<usize>,
        with_joint_entropy: bool,
    ) -> (T, T) {
        self.count(col_i, col_j);
        let n_rows = col_i.nrows();
        let entropy = if with_joint_entropy {
            self.joint_entropy(n_rows)
        } else {
            T::zero()
        };
        (self.mutual_information(n_rows), entropy)
    }
}

/////////////
// Helpers //
/////////////

/// Resolve the number of bins, defaulting to `sqrt(nrows)`
///
/// ### Params
///
/// * `n_bins` - Optional number of bins
/// * `n_rows` - Number of rows of the binned data
///
/// ### Returns
///
/// The number of bins
pub(crate) fn resolve_n_bins<T: BixverseFloat>(n_bins: Option<usize>, n_rows: usize) -> usize {
    n_bins.unwrap_or_else(|| T::from_usize(n_rows).unwrap().sqrt().to_usize().unwrap())
}

/// Calculate the mutual information between two column references
///
/// The columns need to be binned
///
/// ### Params
///
/// * `col_i` - Column reference to the first column to compare
/// * `col_j` - Column reference to the second column to compare
/// * `n_bins` - Optional number of bins. If not provided, will default to
///   `sqrt(nrows)`.
///
/// ### Returns
///
/// The mutual information between the two columns
pub fn calculate_mi<T>(col_i: ColRef<usize>, col_j: ColRef<usize>, n_bins: Option<usize>) -> T
where
    T: BixverseFloat,
{
    let n_bins = resolve_n_bins::<T>(n_bins, col_i.nrows());
    MiScratch::new(n_bins).pair(col_i, col_j, false).0
}

/// Calculates the joint entropy between two column references
///
/// ### Params
///
/// * `col_i` - Column reference to the first column to compare
/// * `col_j` - Column reference to the second column to compare
/// * `n_bins` - Optional number of bins. If not provided, will default to
///   `sqrt(nrows)`.
///
/// ### Returns
///
/// The joint entropy
pub fn calculate_joint_entropy<T>(
    col_i: ColRef<usize>,
    col_j: ColRef<usize>,
    n_bins: Option<usize>,
) -> T
where
    T: BixverseFloat,
{
    let n_bins = resolve_n_bins::<T>(n_bins, col_i.nrows());
    MiScratch::new(n_bins).pair(col_i, col_j, true).1
}

/// Calculates the entropy of a column reference
///
/// The column needs to be binned
///
/// ### Params
///
/// * `col` - Column reference for which to calculate entropy
/// * `n_bins` - Optional number of bins. If not provided, will default to
///   `sqrt(nrows)`.
///
/// ### Returns
///
/// The entropy of the column
pub fn calculate_entropy<T>(col: ColRef<usize>, n_bins: Option<usize>) -> T
where
    T: BixverseFloat,
{
    let n_rows = col.nrows();
    let n_bins =
        n_bins.unwrap_or_else(|| T::from_usize(n_rows).unwrap().sqrt().to_usize().unwrap());
    let mut counts = vec![0usize; n_bins];
    for i in 0..n_rows {
        counts[col[i]] += 1;
    }
    let n = T::from_usize(n_rows).unwrap();
    let mut entropy = T::zero();
    for &count in &counts {
        if count > 0 {
            let prob = T::from_usize(count).unwrap() / n;
            entropy -= prob * prob.ln();
        }
    }
    entropy
}

/////////////
// Binning //
/////////////

/// Binning strategy enum
#[derive(Debug, Clone, Default, Copy)]
pub enum BinningStrategy {
    /// Equal width bins (equal distance between bin edges)
    #[default]
    EqualWidth,
    /// Equal frequency bins (approximately equal number of values per bin)
    EqualFrequency,
}

/// Parsing the binning strategy
///
/// ### Params
///
/// * `s` - string defining the binning strategy
///
/// ### Returns
///
/// The `BinningStrategy`.
pub fn parse_bin_strategy_type(s: &str) -> Option<BinningStrategy> {
    match s.to_lowercase().as_str() {
        "equal_width" | "equalwidth" => Some(BinningStrategy::EqualWidth),
        "equal_frequency" | "equal_freq" | "equalfreq" | "equalfrequency" => {
            Some(BinningStrategy::EqualFrequency)
        }
        _ => None,
    }
}

/// Equal width binning for a single column
///
/// ### Params
///
/// * `col` - Column refence to bin with equal width strategy
/// * `n_bins` - Number of bins to use
///
/// ### Returns
///
/// Binned vector
fn bin_equal_width<T>(col: &ColRef<T>, n_bins: usize) -> Vec<usize>
where
    T: BixverseFloat,
{
    let (min_val, max_val) = col.iter().fold(
        (
            T::from_f64(f64::INFINITY).unwrap(),
            T::from_f64(f64::NEG_INFINITY).unwrap(),
        ),
        |(min, max), x| (min.min(*x), max.max(*x)),
    );
    let range = max_val - min_val;
    if range == T::zero() {
        return vec![0; col.nrows()];
    }
    let step = range / T::from_usize(n_bins).unwrap();
    col.iter()
        .map(|x| {
            let bin = ((*x - min_val) / step).floor().to_usize().unwrap();
            bin.min(n_bins - 1)
        })
        .collect()
}

/// Equal frequency binning for a single column
///
/// ### Params
///
/// * `col` - Column refence to bin with equal frequency strategy
/// * `n_bins` - Number of bins to use
///
/// ### Returns
///
/// Binned vector
fn bin_equal_frequency<T>(col: &ColRef<T>, n_bins: usize) -> Vec<usize>
where
    T: BixverseFloat,
{
    let n_rows = col.nrows();
    let mut sorted_values: Vec<T> = col.iter().copied().collect();
    sorted_values.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    let mut split_points = Vec::with_capacity(n_bins);
    for k in 1..n_bins {
        let idx = (k * n_rows) / n_bins;
        split_points.push(sorted_values[idx.min(n_rows - 1)]);
    }

    let mut bin_assignments = vec![0; n_rows];
    for (i, &value) in col.iter().enumerate() {
        let mut bin = 0;
        for &split in &split_points {
            if value >= split {
                bin += 1;
            } else {
                break;
            }
        }
        bin_assignments[i] = bin.min(n_bins - 1);
    }
    bin_assignments
}

/// Column wise binning
///
/// ### Params
///
/// * `mat` - The matrix on which to apply column-wise binning
/// * `n_bins` - Optional number of bins. If not provided, will default to
///   `sqrt(nrow)`
///
/// ### Returns
///
/// The matrix with the columns being binned into equal distances.
pub fn bin_matrix_cols<T>(
    mat: &MatRef<T>,
    n_bins: Option<usize>,
    strategy: BinningStrategy,
) -> Mat<usize>
where
    T: BixverseFloat,
{
    let (n_rows, n_cols) = mat.shape();
    let n_bins =
        n_bins.unwrap_or_else(|| T::from_usize(n_rows).unwrap().sqrt().to_usize().unwrap());
    let binned_vals: Vec<Vec<usize>> = mat
        .par_col_iter()
        .map(|col| match strategy {
            BinningStrategy::EqualWidth => bin_equal_width(&col, n_bins),
            BinningStrategy::EqualFrequency => bin_equal_frequency(&col, n_bins),
        })
        .collect();
    Mat::from_fn(n_rows, n_cols, |i, j| binned_vals[j][i])
}

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

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

    fn assert_approx_eq(a: f64, b: f64) {
        assert!((a - b).abs() < 1e-10, "{} != {}", a, b);
    }

    /// Two equally frequent symbols hit the maximum entropy for two bins, in nats.
    #[test]
    fn test_calculate_entropy_uniform() {
        let data = mat![[0], [1], [0], [1]]; // usize matrix
        let col = data.col(0);

        // Explicitly set 2 bins
        let ent: f64 = calculate_entropy(col, Some(2));
        assert_approx_eq(ent, 2.0_f64.ln());
    }

    /// A single occupied bin has zero entropy, and the empty bins must not contribute NaN.
    #[test]
    fn test_calculate_entropy_zero() {
        let data = mat![[0], [0], [0], [0]];
        let col = data.col(0);

        let ent: f64 = calculate_entropy(col, Some(1));
        assert_approx_eq(ent, 0.0);
    }

    /// Joint entropy equals H(X) for identical columns and adds up for independent ones.
    #[test]
    fn test_calculate_joint_entropy() {
        let data = mat![[0, 0], [1, 1], [0, 0], [1, 1]]; // usize matrix

        let je: f64 = calculate_joint_entropy(data.col(0), data.col(1), Some(2));
        assert_approx_eq(je, 2.0_f64.ln());

        // Independent variables
        let data_indep = mat![[0, 0], [0, 1], [1, 0], [1, 1]];
        let je_indep: f64 = calculate_joint_entropy(data_indep.col(0), data_indep.col(1), Some(2));
        assert_approx_eq(je_indep, 4.0_f64.ln());
    }

    /// MI saturates at H(X) for identical columns and drops to zero for independent ones.
    #[test]
    fn test_calculate_mi() {
        let data = mat![[0, 0], [1, 1]];
        let mi: f64 = calculate_mi(data.col(0), data.col(1), Some(2));
        assert_approx_eq(mi, 2.0_f64.ln());

        // Independent -> MI should be 0
        let data_indep = mat![[0, 0], [0, 1], [1, 0], [1, 1]];
        let mi_indep: f64 = calculate_mi(data_indep.col(0), data_indep.col(1), Some(2));
        assert_approx_eq(mi_indep, 0.0);
    }

    /// Equal-width binning clamps the top edge back into the last bin rather than overflowing it.
    #[test]
    fn test_bin_equal_width() {
        let mat = mat![[0.0], [5.0], [10.0]];

        let binned = bin_matrix_cols(&mat.as_ref(), Some(3), BinningStrategy::EqualWidth);

        assert_eq!(binned.nrows(), 3);
        assert_eq!(binned.get(0, 0), &0);
        assert_eq!(binned.get(1, 0), &1);
        assert_eq!(binned.get(2, 0), &2);
    }

    /// A constant column has zero range, so the bin width must not become a division by zero.
    #[test]
    fn test_bin_equal_width_constant_column() {
        // Edge case: Range is 0
        let mat = mat![[1.0], [1.0], [1.0]];
        let binned = bin_matrix_cols(&mat.as_ref(), Some(3), BinningStrategy::EqualWidth);

        // Should default to bin 0 for all
        assert_eq!(binned.get(0, 0), &0);
        assert_eq!(binned.get(1, 0), &0);
        assert_eq!(binned.get(2, 0), &0);
    }

    /// Equal-frequency binning splits on rank yet writes the labels back in input order.
    #[test]
    fn test_bin_equal_frequency() {
        let mat = mat![[1.0], [9.0], [2.0], [8.0]];

        let binned = bin_matrix_cols(&mat.as_ref(), Some(2), BinningStrategy::EqualFrequency);

        // 1.0 and 2.0 should be in bin 0
        // 8.0 and 9.0 should be in bin 1
        // Indices in input: 0->1.0, 1->9.0, 2->2.0, 3->8.0
        assert_eq!(binned.get(0, 0), &0); // 1.0
        assert_eq!(binned.get(1, 0), &1); // 9.0
        assert_eq!(binned.get(2, 0), &0); // 2.0
        assert_eq!(binned.get(3, 0), &1); // 8.0
    }

    /// The R-facing strategy strings still map onto the enum variants they name.
    #[test]
    fn test_parse_bin_strategy() {
        assert!(matches!(
            parse_bin_strategy_type("equal_width"),
            Some(BinningStrategy::EqualWidth)
        ));
        assert!(matches!(
            parse_bin_strategy_type("equal_freq"),
            Some(BinningStrategy::EqualFrequency)
        ));
        assert!(parse_bin_strategy_type("random").is_none());
    }
}