bixverse_methods/methods/nmf_hals/
dense.rs

1//! Dense backend for NMF. Borrows V as a faer `MatRef` and routes both data
2//! products through `utils::gemm::gemm`.
3
4use bixverse_core::utils::gemm::gemm;
5use faer::{Accum, Mat, MatRef};
6
7use super::NmfInput;
8use crate::prelude::*;
9use bixverse_core::core::math::pca_svd::{SvdResults, randomised_svd};
10use bixverse_core::utils::faer_parallelism;
11
12/// Dense, non-negative V wrapped behind the NMF trait. Validates non-negativity
13/// and finiteness once at construction.
14#[derive(Debug)]
15pub struct DenseInput<'a, F: BixverseFloat> {
16    /// The input matrix v
17    v: MatRef<'a, F>,
18    /// The squared Frobenius norm
19    sq_frob: F,
20}
21
22impl<'a, F: BixverseFloat> DenseInput<'a, F> {
23    /// Construct a [`DenseInput`] from a borrowed matrix.
24    ///
25    /// Validates that all entries are finite and non-negative, then computes
26    /// `||V||_F^2`.
27    ///
28    /// ### Params
29    ///
30    /// * `v` - The input matrix to factorise.
31    ///
32    /// ### Returns
33    ///
34    /// A [`DenseInput`] wrapping `v`.
35    pub fn new(v: MatRef<'a, F>) -> Result<Self, MethodsErrors> {
36        let (m, n) = (v.nrows(), v.ncols());
37        // Accumulated in f64 and cast back once. In `F` this drifts badly:
38        // a naive f32 sum over positive terms stops registering the increment
39        // once the running total dwarfs it, which measured 2.5% low over 1.5e7
40        // entries and put that error straight into every relative loss the
41        // solver reports. See [`sum_sq_f64`] for the same problem on the sparse
42        // side. Validation stays fused into the same pass.
43        let mut sq_frob = 0f64;
44        for j in 0..n {
45            for i in 0..m {
46                let x = v[(i, j)];
47                if !x.is_finite() {
48                    return Err(MethodsErrors::NmfNonFinite);
49                }
50                if x < F::zero() {
51                    return Err(MethodsErrors::NmfNonNegativeViolated);
52                }
53                let d = x.to_f64().unwrap();
54                sq_frob += d * d;
55            }
56        }
57        Ok(Self {
58            v,
59            sq_frob: F::from_f64(sq_frob).unwrap(),
60        })
61    }
62}
63
64/// NmfInput trait implementation for [DenseInput]
65impl<F: BixverseFloat> NmfInput<F> for DenseInput<'_, F> {
66    fn shape(&self) -> (usize, usize) {
67        (self.v.nrows(), self.v.ncols())
68    }
69
70    fn sq_frob(&self) -> F {
71        self.sq_frob
72    }
73
74    fn wt_v(&self, w: MatRef<F>, out: &mut Mat<F>) {
75        gemm(
76            out.as_mut(),
77            Accum::Replace,
78            w.transpose(),
79            self.v,
80            F::one(),
81            faer_parallelism(),
82        );
83    }
84
85    fn v_ht(&self, h: MatRef<F>, out: &mut Mat<F>) {
86        gemm(
87            out.as_mut(),
88            Accum::Replace,
89            self.v,
90            h.transpose(),
91            F::one(),
92            faer_parallelism(),
93        );
94    }
95
96    fn top_k_svd(&self, k: usize) -> Result<SvdResults<F>, MethodsErrors> {
97        let (n, m) = (self.v.nrows(), self.v.ncols());
98        let min_dim = m.min(n);
99
100        if min_dim < k {
101            return Err(MethodsErrors::NmfRankTooLarge {
102                requested: k,
103                available: min_dim,
104            });
105        }
106
107        // randomised when k is small relative to min(m, n) and the matrix is
108        // large enough that the projection pays off; otherwise thin_svd.
109        let use_randomised = k * 4 < min_dim && min_dim > 64;
110
111        if use_randomised {
112            let rsvd = randomised_svd(self.v, k, 42, None, None)?;
113            // randomised_svd returns rank + oversampling triplets; truncate to k.
114            let u = Mat::<F>::from_fn(rsvd.u.nrows(), k, |i, j| rsvd.u[(i, j)]);
115            let v = Mat::<F>::from_fn(rsvd.v.nrows(), k, |i, j| rsvd.v[(i, j)]);
116            let s: Vec<F> = rsvd.s.into_iter().take(k).collect();
117            Ok(SvdResults { u, v, s })
118        } else {
119            let svd = self
120                .v
121                .thin_svd()
122                .map_err(|e| CoreErrors::FaerSvdError(format!("{e:?}")))?;
123
124            let u_full = svd.U();
125            let v_full = svd.V();
126            let s_full = svd.S().column_vector();
127
128            let u = Mat::<F>::from_fn(u_full.nrows(), k, |i, j| u_full[(i, j)]);
129            let v = Mat::<F>::from_fn(v_full.nrows(), k, |i, j| v_full[(i, j)]);
130            let s: Vec<F> = (0..k).map(|i| s_full[i]).collect();
131
132            Ok(SvdResults { u, v, s })
133        }
134    }
135}
136
137///////////
138// Tests //
139///////////
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use faer::Mat;
145
146    /// `||V||_F^2` has to be accurate on a matrix where accumulating in `F`
147    /// visibly is not, because it is the denominator of every relative loss the
148    /// solver and the consensus path report.
149    ///
150    /// The leading entry squares to `1e8`, where the f32 spacing is 8, so a
151    /// running f32 total swallows every one of the `1.0` entries that follow. The
152    /// second assertion pins that the reference really is wrong, so the first one
153    /// is testing something.
154    #[test]
155    fn sq_frob_survives_a_swamped_accumulator() {
156        let (m, n) = (1_000usize, 1_000usize);
157        let v = Mat::<f32>::from_fn(m, n, |i, j| if i == 0 && j == 0 { 1e4 } else { 1.0 });
158        let input = DenseInput::new(v.as_ref()).unwrap();
159
160        let exact = 1e8 + (m * n - 1) as f64;
161        let got = input.sq_frob() as f64;
162        assert!(
163            (got - exact).abs() <= 1e-5 * exact,
164            "sq_frob is off: got {got}, exact {exact}"
165        );
166
167        // The f32 accumulation this replaced, in the same column-major order.
168        let mut naive = 0f32;
169        for j in 0..n {
170            for i in 0..m {
171                let x = v[(i, j)];
172                naive += x * x;
173            }
174        }
175        assert!(
176            (naive as f64 - exact).abs() > 0.005 * exact,
177            "the f32 reference was supposed to be visibly wrong, so this test proves nothing"
178        );
179    }
180
181    /// Shape is reported as (rows, cols) and the cached norm is the sum of squared entries.
182    #[test]
183    fn shape_and_sq_frob() {
184        let v: Mat<f64> = Mat::from_fn(3, 2, |i, j| (i + j) as f64);
185        let input = DenseInput::new(v.as_ref()).unwrap();
186        assert_eq!(input.shape(), (3, 2));
187        assert!((input.sq_frob() - 19.0).abs() < 1e-12);
188    }
189
190    /// A negative entry is refused at construction rather than silently factorised.
191    #[test]
192    fn rejects_negative() {
193        let v: Mat<f64> = Mat::from_fn(2, 2, |i, _| if i == 0 { 1.0 } else { -1.0 });
194        let err = DenseInput::new(v.as_ref()).unwrap_err();
195        assert!(matches!(err, MethodsErrors::NmfNonNegativeViolated));
196    }
197
198    /// Both NaN and infinity are refused at construction.
199    #[test]
200    fn rejects_non_finite() {
201        let v: Mat<f64> = Mat::from_fn(2, 2, |i, _| if i == 0 { 1.0 } else { f64::NAN });
202        let err = DenseInput::new(v.as_ref()).unwrap_err();
203        assert!(matches!(err, MethodsErrors::NmfNonFinite));
204
205        let v: Mat<f64> = Mat::from_fn(2, 2, |i, _| if i == 0 { 1.0 } else { f64::INFINITY });
206        let err = DenseInput::new(v.as_ref()).unwrap_err();
207        assert!(matches!(err, MethodsErrors::NmfNonFinite));
208    }
209
210    /// The faer-routed `W^T V` matches an explicit dot product reference, transpose included.
211    #[test]
212    fn wt_v_matches_reference() {
213        let (m, n, k) = (4, 3, 2);
214        let v_mat: Mat<f64> = Mat::from_fn(m, n, |i, j| (i * n + j + 1) as f64);
215        let w: Mat<f64> = Mat::from_fn(m, k, |i, j| (i + j + 1) as f64);
216        let input = DenseInput::new(v_mat.as_ref()).unwrap();
217        let mut out = Mat::<f64>::zeros(k, n);
218        input.wt_v(w.as_ref(), &mut out);
219        for r in 0..k {
220            for j in 0..n {
221                let expected: f64 = (0..m).map(|i| w[(i, r)] * v_mat[(i, j)]).sum();
222                assert!((out[(r, j)] - expected).abs() < 1e-9);
223            }
224        }
225    }
226
227    /// The faer-routed `V H^T` matches an explicit dot product reference, transpose included.
228    #[test]
229    fn v_ht_matches_reference() {
230        let (m, n, k) = (4, 3, 2);
231        let v_mat: Mat<f64> = Mat::from_fn(m, n, |i, j| (i * n + j + 1) as f64);
232        let h: Mat<f64> = Mat::from_fn(k, n, |i, j| (i + j + 1) as f64);
233        let input = DenseInput::new(v_mat.as_ref()).unwrap();
234        let mut out = Mat::<f64>::zeros(m, k);
235        input.v_ht(h.as_ref(), &mut out);
236        for i in 0..m {
237            for r in 0..k {
238                let expected: f64 = (0..n).map(|j| v_mat[(i, j)] * h[(r, j)]).sum();
239                assert!((out[(i, r)] - expected).abs() < 1e-9);
240            }
241        }
242    }
243
244    /// The truncated SVD returns k triplets of the right shape, descending and non-negative.
245    #[test]
246    fn top_k_svd_shapes_and_ordering() {
247        let v_mat: Mat<f64> = Mat::from_fn(5, 4, |i, j| (i * 4 + j + 1) as f64);
248        let input = DenseInput::new(v_mat.as_ref()).unwrap();
249        let svd = input.top_k_svd(2).unwrap();
250        assert_eq!(svd.u.shape(), (5, 2));
251        assert_eq!(svd.v.shape(), (4, 2));
252        assert_eq!(svd.s.len(), 2);
253        assert!(svd.s[0] >= svd.s[1]);
254        assert!(svd.s[1] >= 0.0);
255    }
256
257    /// A k above min(m, n) errors instead of indexing past the available triplets.
258    #[test]
259    fn top_k_svd_rejects_too_large_k() {
260        let v_mat: Mat<f64> = Mat::from_fn(3, 2, |i, j| (i + j + 1) as f64);
261        let input = DenseInput::new(v_mat.as_ref()).unwrap();
262        let err = input.top_k_svd(3).unwrap_err();
263        assert!(matches!(err, MethodsErrors::NmfRankTooLarge { .. }));
264    }
265}
bixverse_methods/methods/nmf_hals/dense.rs

bixverse_methods/methods/nmf_hals/
dense.rs
