bixverse-rs 0.6.0

//! Dense backend for NMF. Borrows V as a faer `MatRef` and routes both data
//! products through `utils::gemm::gemm`.

use crate::utils::gemm::gemm;
use faer::{Accum, Mat, MatRef};

use super::NmfInput;
use crate::core::math::pca_svd::{SvdResults, randomised_svd};
use crate::prelude::*;
use crate::utils::faer_parallelism;

/// Dense, non-negative V wrapped behind the NMF trait. Validates non-negativity
/// and finiteness once at construction.
#[derive(Debug)]
pub struct DenseInput<'a, F: BixverseFloat> {
    /// The input matrix v
    v: MatRef<'a, F>,
    /// The squared Frobenius norm
    sq_frob: F,
}

impl<'a, F: BixverseFloat> DenseInput<'a, F> {
    /// Construct a [`DenseInput`] from a borrowed matrix.
    ///
    /// Validates that all entries are finite and non-negative, then computes
    /// `||V||_F^2`.
    ///
    /// ### Params
    ///
    /// * `v` - The input matrix to factorise.
    ///
    /// ### Returns
    ///
    /// A [`DenseInput`] wrapping `v`.
    pub fn new(v: MatRef<'a, F>) -> Result<Self, BixverseErrors> {
        let (m, n) = (v.nrows(), v.ncols());
        // Accumulated in f64 and cast back once. In `F` this drifts badly:
        // a naive f32 sum over positive terms stops registering the increment
        // once the running total dwarfs it, which measured 2.5% low over 1.5e7
        // entries and put that error straight into every relative loss the
        // solver reports. See [`sum_sq_f64`] for the same problem on the sparse
        // side. Validation stays fused into the same pass.
        let mut sq_frob = 0f64;
        for j in 0..n {
            for i in 0..m {
                let x = v[(i, j)];
                if !x.is_finite() {
                    return Err(BixverseErrors::NmfNonFinite);
                }
                if x < F::zero() {
                    return Err(BixverseErrors::NmfNonNegativeViolated);
                }
                let d = x.to_f64().unwrap();
                sq_frob += d * d;
            }
        }
        Ok(Self {
            v,
            sq_frob: F::from_f64(sq_frob).unwrap(),
        })
    }
}

/// NmfInput trait implementation for [DenseInput]
impl<F: BixverseFloat> NmfInput<F> for DenseInput<'_, F> {
    fn shape(&self) -> (usize, usize) {
        (self.v.nrows(), self.v.ncols())
    }

    fn sq_frob(&self) -> F {
        self.sq_frob
    }

    fn wt_v(&self, w: MatRef<F>, out: &mut Mat<F>) {
        gemm(
            out.as_mut(),
            Accum::Replace,
            w.transpose(),
            self.v,
            F::one(),
            faer_parallelism(),
        );
    }

    fn v_ht(&self, h: MatRef<F>, out: &mut Mat<F>) {
        gemm(
            out.as_mut(),
            Accum::Replace,
            self.v,
            h.transpose(),
            F::one(),
            faer_parallelism(),
        );
    }

    fn top_k_svd(&self, k: usize) -> Result<SvdResults<F>, BixverseErrors> {
        let (n, m) = (self.v.nrows(), self.v.ncols());
        let min_dim = m.min(n);

        if min_dim < k {
            return Err(BixverseErrors::NmfRankTooLarge {
                requested: k,
                available: min_dim,
            });
        }

        // randomised when k is small relative to min(m, n) and the matrix is
        // large enough that the projection pays off; otherwise thin_svd.
        let use_randomised = k * 4 < min_dim && min_dim > 64;

        if use_randomised {
            let rsvd = randomised_svd(self.v, k, 42, None, None)?;
            // randomised_svd returns rank + oversampling triplets; truncate to k.
            let u = Mat::<F>::from_fn(rsvd.u.nrows(), k, |i, j| rsvd.u[(i, j)]);
            let v = Mat::<F>::from_fn(rsvd.v.nrows(), k, |i, j| rsvd.v[(i, j)]);
            let s: Vec<F> = rsvd.s.into_iter().take(k).collect();
            Ok(SvdResults { u, v, s })
        } else {
            let svd = self
                .v
                .thin_svd()
                .map_err(|e| BixverseErrors::FaerSvdError(format!("{e:?}")))?;

            let u_full = svd.U();
            let v_full = svd.V();
            let s_full = svd.S().column_vector();

            let u = Mat::<F>::from_fn(u_full.nrows(), k, |i, j| u_full[(i, j)]);
            let v = Mat::<F>::from_fn(v_full.nrows(), k, |i, j| v_full[(i, j)]);
            let s: Vec<F> = (0..k).map(|i| s_full[i]).collect();

            Ok(SvdResults { u, v, s })
        }
    }
}

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

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

    /// `||V||_F^2` has to be accurate on a matrix where accumulating in `F`
    /// visibly is not, because it is the denominator of every relative loss the
    /// solver and the consensus path report.
    ///
    /// The leading entry squares to `1e8`, where the f32 spacing is 8, so a
    /// running f32 total swallows every one of the `1.0` entries that follow. The
    /// second assertion pins that the reference really is wrong, so the first one
    /// is testing something.
    #[test]
    fn sq_frob_survives_a_swamped_accumulator() {
        let (m, n) = (1_000usize, 1_000usize);
        let v = Mat::<f32>::from_fn(m, n, |i, j| if i == 0 && j == 0 { 1e4 } else { 1.0 });
        let input = DenseInput::new(v.as_ref()).unwrap();

        let exact = 1e8 + (m * n - 1) as f64;
        let got = input.sq_frob() as f64;
        assert!(
            (got - exact).abs() <= 1e-5 * exact,
            "sq_frob is off: got {got}, exact {exact}"
        );

        // The f32 accumulation this replaced, in the same column-major order.
        let mut naive = 0f32;
        for j in 0..n {
            for i in 0..m {
                let x = v[(i, j)];
                naive += x * x;
            }
        }
        assert!(
            (naive as f64 - exact).abs() > 0.005 * exact,
            "the f32 reference was supposed to be visibly wrong, so this test proves nothing"
        );
    }

    /// Shape is reported as (rows, cols) and the cached norm is the sum of squared entries.
    #[test]
    fn shape_and_sq_frob() {
        let v: Mat<f64> = Mat::from_fn(3, 2, |i, j| (i + j) as f64);
        let input = DenseInput::new(v.as_ref()).unwrap();
        assert_eq!(input.shape(), (3, 2));
        assert!((input.sq_frob() - 19.0).abs() < 1e-12);
    }

    /// A negative entry is refused at construction rather than silently factorised.
    #[test]
    fn rejects_negative() {
        let v: Mat<f64> = Mat::from_fn(2, 2, |i, _| if i == 0 { 1.0 } else { -1.0 });
        let err = DenseInput::new(v.as_ref()).unwrap_err();
        assert!(matches!(err, BixverseErrors::NmfNonNegativeViolated));
    }

    /// Both NaN and infinity are refused at construction.
    #[test]
    fn rejects_non_finite() {
        let v: Mat<f64> = Mat::from_fn(2, 2, |i, _| if i == 0 { 1.0 } else { f64::NAN });
        let err = DenseInput::new(v.as_ref()).unwrap_err();
        assert!(matches!(err, BixverseErrors::NmfNonFinite));

        let v: Mat<f64> = Mat::from_fn(2, 2, |i, _| if i == 0 { 1.0 } else { f64::INFINITY });
        let err = DenseInput::new(v.as_ref()).unwrap_err();
        assert!(matches!(err, BixverseErrors::NmfNonFinite));
    }

    /// The faer-routed `W^T V` matches an explicit dot product reference, transpose included.
    #[test]
    fn wt_v_matches_reference() {
        let (m, n, k) = (4, 3, 2);
        let v_mat: Mat<f64> = Mat::from_fn(m, n, |i, j| (i * n + j + 1) as f64);
        let w: Mat<f64> = Mat::from_fn(m, k, |i, j| (i + j + 1) as f64);
        let input = DenseInput::new(v_mat.as_ref()).unwrap();
        let mut out = Mat::<f64>::zeros(k, n);
        input.wt_v(w.as_ref(), &mut out);
        for r in 0..k {
            for j in 0..n {
                let expected: f64 = (0..m).map(|i| w[(i, r)] * v_mat[(i, j)]).sum();
                assert!((out[(r, j)] - expected).abs() < 1e-9);
            }
        }
    }

    /// The faer-routed `V H^T` matches an explicit dot product reference, transpose included.
    #[test]
    fn v_ht_matches_reference() {
        let (m, n, k) = (4, 3, 2);
        let v_mat: Mat<f64> = Mat::from_fn(m, n, |i, j| (i * n + j + 1) as f64);
        let h: Mat<f64> = Mat::from_fn(k, n, |i, j| (i + j + 1) as f64);
        let input = DenseInput::new(v_mat.as_ref()).unwrap();
        let mut out = Mat::<f64>::zeros(m, k);
        input.v_ht(h.as_ref(), &mut out);
        for i in 0..m {
            for r in 0..k {
                let expected: f64 = (0..n).map(|j| v_mat[(i, j)] * h[(r, j)]).sum();
                assert!((out[(i, r)] - expected).abs() < 1e-9);
            }
        }
    }

    /// The truncated SVD returns k triplets of the right shape, descending and non-negative.
    #[test]
    fn top_k_svd_shapes_and_ordering() {
        let v_mat: Mat<f64> = Mat::from_fn(5, 4, |i, j| (i * 4 + j + 1) as f64);
        let input = DenseInput::new(v_mat.as_ref()).unwrap();
        let svd = input.top_k_svd(2).unwrap();
        assert_eq!(svd.u.shape(), (5, 2));
        assert_eq!(svd.v.shape(), (4, 2));
        assert_eq!(svd.s.len(), 2);
        assert!(svd.s[0] >= svd.s[1]);
        assert!(svd.s[1] >= 0.0);
    }

    /// A k above min(m, n) errors instead of indexing past the available triplets.
    #[test]
    fn top_k_svd_rejects_too_large_k() {
        let v_mat: Mat<f64> = Mat::from_fn(3, 2, |i, j| (i + j + 1) as f64);
        let input = DenseInput::new(v_mat.as_ref()).unwrap();
        let err = input.top_k_svd(3).unwrap_err();
        assert!(matches!(err, BixverseErrors::NmfRankTooLarge { .. }));
    }
}