bixverse-rs 0.6.2

//! faer matrix transformation utils. Concantenations, row removals, etc.

use faer::{Mat, MatRef, concat};
use rayon::prelude::*;
use std::borrow::Cow;

use crate::prelude::*;
use crate::utils::vec_utils::flatten_vector;

/// Rows per block when transposing between row-major buffers and `Mat`
const TRANSPOSE_ROW_BLOCK: usize = 64;

/// Columns written per task when filling a `Mat` from a row-major buffer
const TRANSPOSE_COL_BLOCK: usize = 8;

/// Transform a nested vector into a faer matrix
///
/// ### Params
///
/// * `nested_vec` - The nested vector
/// * `col_wise` - If set to `True` it will column bind (outer vector represents)
///   the columns. If set to `False` it will row bind (outer vector represents
///   the rows).
///
/// ### Returns
///
/// The row or column bound matrix.
pub fn nested_vector_to_faer_mat<T>(nested_vec: Vec<Vec<T>>, col_wise: bool) -> Mat<T>
where
    T: BixverseFloat,
{
    let (nrow, ncol) = if col_wise {
        (nested_vec[0].len(), nested_vec.len())
    } else {
        (nested_vec.len(), nested_vec[0].len())
    };

    let data = flatten_vector(nested_vec);

    if col_wise {
        Mat::from_fn(nrow, ncol, |i, j| data[i + j * nrow])
    } else {
        Mat::from_fn(nrow, ncol, |i, j| data[j + i * ncol])
    }
}

/// Create a diagonal matrix from vector values
///
/// ### Params
///
/// * `vec` - The vector of values to put in the diagonal of the matrix
///
/// ### Returns
///
/// The diagonal matrix as a faer Matrix.
pub fn faer_diagonal_from_vec<T>(vec: Vec<T>) -> Mat<T>
where
    T: BixverseFloat,
{
    let len = vec.len();
    Mat::from_fn(
        len,
        len,
        |row, col| if row == col { vec[row] } else { T::zero() },
    )
}

/// Get the index positions of the upper triangle of a symmetric matrix
///
/// Function will panic if offset > 1.
///
/// ### Params
///
/// * `n_dim` - The dimensions of the symmetric matrix
/// * `offset` - Do you want to include the diagonal values (offset = 0) or exclude
///   them (offset = 1).
///
/// ### Returns
///
/// A tuple of the row and column index positions of the upper triangle of the
/// matirx
pub fn upper_triangle_indices(n_dim: usize, offset: usize) -> (Vec<usize>, Vec<usize>) {
    if offset >= n_dim {
        return (Vec::new(), Vec::new());
    }
    assert!(offset <= 1, "The offset should be 0 or 1");

    // Precise calculation of total elements
    let total_elements: usize = (0..n_dim)
        .map(|row| n_dim.saturating_sub(row + offset))
        .sum();

    let mut row_indices = Vec::with_capacity(total_elements);
    let mut col_indices = Vec::with_capacity(total_elements);

    for row in 0..n_dim {
        let start_col = row + offset;
        if start_col < n_dim {
            let end_col = n_dim;
            let elements_in_row = end_col - start_col;
            // Use repeat_n for better performance and clarity
            row_indices.extend(std::iter::repeat_n(row, elements_in_row));
            col_indices.extend(start_col..end_col);
        }
    }

    (row_indices, col_indices)
}

/// Create from the upper triangle values a symmetric matrix
///
/// Generates the full dense matrix of values representing the upper triangle
/// of a symmetric matrix.
///
/// ### Params
///
/// * `data` - Slice of the values
/// * `shift` - Was the diagonal included (= 0) or not (= 1). If not included,
///   the diagonal is set to 1.
/// * `n` - Original dimension of the symmetric matrix.
///
/// ### Return
///
/// The symmetric, dense matrix.
pub fn upper_triangle_to_sym_faer<T>(data: &[T], shift: usize, n: usize) -> faer::Mat<T>
where
    T: BixverseFloat,
{
    let mut mat = Mat::<T>::zeros(n, n);
    let mut idx = 0;
    for i in 0..n {
        for j in i..n {
            if shift == 1 && i == j {
                mat[(i, j)] = T::one();
            } else {
                mat[(i, j)] = data[idx];
                mat[(j, i)] = data[idx];
                idx += 1;
            }
        }
    }

    mat
}

/// Store the upper triangle values as a flat vector from a faer matrix
///
/// ### Params
///
/// * `x` The faer matrix
/// * `shift` Shall the diagonal be included (shift = 0) or not (shift = 1).
///
/// ### Returns
///
/// A vector representing the upper triangle values (row major ordered)
pub fn faer_mat_to_upper_triangle<T>(x: MatRef<T>, shift: usize) -> Vec<T>
where
    T: BixverseFloat,
{
    assert!(shift <= 1, "The shift should be 0 or 1");

    let n = x.ncols();
    let total_elements = if shift == 0 {
        n * (n + 1) / 2
    } else {
        n * (n - 1) / 2
    };
    let mut vals = Vec::with_capacity(total_elements);
    for i in 0..n {
        let start_j = i + shift;
        for j in start_j..n {
            vals.push(*x.get(i, j));
        }
    }

    vals
}

/// Slice out a single row and return the remaining matrix
///
/// ### Params
///
/// * `x` - The matrix from which to remove a single row
/// * `idx_to_remove` - The index of the row to remove.
///
/// ### Returns
///
/// The matrix minus the specified row.
pub fn mat_rm_row<T>(x: MatRef<T>, idx_to_remove: usize) -> Mat<T>
where
    T: BixverseFloat,
{
    assert!(
        idx_to_remove <= x.nrows(),
        "The specified index is larger than the matrix"
    );

    let total_rows = x.nrows();

    if idx_to_remove == 0 {
        x.subrows(1, total_rows - 1).to_owned()
    } else if idx_to_remove == total_rows - 1 {
        x.subrows(0, total_rows - 1).to_owned()
    } else {
        let upper = x.subrows(0, idx_to_remove);
        let lower = x.subrows(idx_to_remove + 1, total_rows - idx_to_remove - 1);
        concat![[upper], [lower]]
    }
}

/// Rowbind a vector of faer Matrices
///
/// The function will panic if the number of columns of the matrices differ in
/// the vector
///
/// ### Params
///
/// * `matrices` - Vector of faer matrix to row bind
///
/// ### Returns
///
/// One row bound matrix from the initial matrices
#[allow(dead_code)]
pub fn rowbind_matrices<T>(matrices: Vec<Mat<T>>) -> Mat<T>
where
    T: BixverseFloat,
{
    let ncols = matrices[0].ncols();
    let total_row = matrices.iter().map(|m| m.nrows()).sum();
    let mut result: Mat<T> = Mat::zeros(total_row, ncols);
    let mut row_offset = 0;
    for matrix in matrices {
        assert_eq!(
            matrix.ncols(),
            ncols,
            "All matrices must have the same number of columns"
        );
        let nrows = matrix.nrows();
        result
            .as_mut()
            .submatrix_mut(row_offset, 0, nrows, ncols)
            .copy_from(matrix.as_ref());
        row_offset += nrows;
    }

    result
}

/// Colbind a vector of faer Matrices
///
/// The function will panic if the number of rows of the matrices differ in
/// the vector
///
/// ### Params
///
/// * `matrices` - Vector of faer matrix to column bind
///
/// ### Returns
///
/// One column bound matrix from the initial matrices
pub fn colbind_matrices<T>(matrices: &[Mat<T>]) -> Mat<T>
where
    T: BixverseFloat,
{
    let nrows = matrices[0].nrows();
    let total_col = matrices.iter().map(|m| m.ncols()).sum();
    let mut result: Mat<T> = Mat::zeros(nrows, total_col);
    let mut col_offset = 0;
    for matrix in matrices {
        assert_eq!(
            matrix.nrows(),
            nrows,
            "All matrices must have the same number of columns"
        );
        let ncols = matrix.ncols();
        result
            .as_mut()
            .submatrix_mut(0, col_offset, nrows, ncols)
            .copy_from(matrix.as_ref());
        col_offset += ncols;
    }

    result
}

/// Convert faer matrix to flat row-major layout
///
/// Helper function that is used to flatten structures into a row-major layout
/// for certain other functions that expect flat row-major layout.
///
/// ### Params
///
/// * `mat` - Matrix in N×d layout (cells × features)
///
/// ### Returns
///
/// Flat vector containing all data in row-major order
pub fn mat_to_flat_row_major<T>(mat: MatRef<T>) -> Vec<T>
where
    T: BixverseFloat,
{
    let n = mat.nrows();
    let d = mat.ncols();
    let mut flat = vec![T::zero(); n * d];

    // Row blocks, column outer: reads stay on contiguous column slices and the
    // strided writes stay inside one cache-resident block.
    flat.par_chunks_mut((TRANSPOSE_ROW_BLOCK * d).max(1))
        .enumerate()
        .for_each(|(b, block)| {
            let r0 = b * TRANSPOSE_ROW_BLOCK;
            for c in 0..d {
                for r in 0..block.len() / d {
                    block[r * d + c] = mat[(r0 + r, c)];
                }
            }
        });

    flat
}

/// Convert flat row-major layout back to faer matrix
///
/// ### Params
///
/// * `flat` - Flat data in row-major order
/// * `nrows` - Number of rows
/// * `ncols` - Number of columns
///
/// ### Returns
///
/// Faer Mat in nrows × ncols shape
pub fn flat_row_major_to_mat<T>(flat: &[T], nrows: usize, ncols: usize) -> Mat<T>
where
    T: BixverseFloat,
{
    assert_eq!(
        flat.len(),
        nrows * ncols,
        "Flat data size doesn't match dimensions"
    );

    let mut out = Mat::<T>::zeros(nrows, ncols);
    out.par_col_chunks_mut(TRANSPOSE_COL_BLOCK)
        .enumerate()
        .for_each(|(b, mut block)| {
            let c0 = b * TRANSPOSE_COL_BLOCK;
            for i in 0..nrows {
                for c in 0..block.ncols() {
                    block[(i, c)] = flat[i * ncols + c0 + c];
                }
            }
        });

    out
}

/// Borrow every column of a matrix as a contiguous slice
///
/// Columns of an owned `Mat` or a column-major `MatRef` are borrowed; a
/// strided column is copied.
///
/// ### Params
///
/// * `mat` - The matrix
///
/// ### Returns
///
/// One slice per column.
pub fn column_slices<T>(mat: MatRef<'_, T>) -> Vec<Cow<'_, [T]>>
where
    T: BixverseFloat,
{
    (0..mat.ncols())
        .map(|j| match mat.col(j).try_as_col_major() {
            Some(col) => Cow::Borrowed(col.as_slice()),
            None => Cow::Owned(mat.col(j).iter().cloned().collect()),
        })
        .collect()
}

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

#[cfg(test)]
mod tests_matrix_utils {
    use super::*;
    use approx::assert_relative_eq;

    /// The flat layout must be row-major, not faer's native column-major order.
    #[test]
    fn test_mat_to_flat_row_major() {
        let mat = Mat::from_fn(2, 3, |i, j| (i * 3 + j) as f32);
        let flat = mat_to_flat_row_major(mat.as_ref());
        assert_eq!(flat, vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]);
    }

    /// Flattening and rebuilding must return the matrix it started from.
    #[test]
    fn test_round_trip() {
        let original = Mat::from_fn(3, 4, |i, j| (i as f32) * 0.5 + (j as f32) * 0.1);
        let flat = mat_to_flat_row_major(original.as_ref());
        let restored = flat_row_major_to_mat(&flat, 3, 4);

        for i in 0..3 {
            for j in 0..4 {
                assert_relative_eq!(original[(i, j)], restored[(i, j)], epsilon = 1e-7);
            }
        }
    }
}