bixverse_methods/methods/nmf_hals/
refit.rs

1//! Fixed-factor NNLS refits. Solves for one factor while the other is frozen,
2//! which is what a consensus NMF needs once the consensus factor has been
3//! assembled from the clustered restarts.
4//!
5//! Freezing a factor makes both of its data products constant, so `W^T V` (or
6//! `V H^T`) is computed once up front rather than once per iteration. That is
7//! the expensive part of a HALS sweep, so a refit is far cheaper than a full
8//! solve of the same iteration count.
9
10use faer::{Mat, MatRef};
11use rayon::prelude::*;
12
13use super::{
14    HalsOpts, NmfInput, compute_objective, gram_h_ht, gram_wt_w, hals_sweep_cols, hals_sweep_rows,
15};
16use crate::prelude::*;
17
18/////////////
19// Helpers //
20/////////////
21
22/// Denominator for the relative-decrease convergence test.
23///
24/// Matches [`super::nmf_hals`]: `||V||_F^2` floored at one, so `tol` reads as a
25/// fraction of the total energy and tiny-norm inputs do not divide by a value
26/// below one.
27///
28/// ### Params
29///
30/// * `sq_frob` - The squared Frobenius norm of V.
31///
32/// ### Returns
33///
34/// The floored denominator.
35#[inline]
36fn relative_denominator<F: BixverseFloat>(sq_frob: F) -> F {
37    if sq_frob > F::one() {
38        sq_frob
39    } else {
40        F::one()
41    }
42}
43
44/// Reconstruction error with H frozen.
45///
46/// Evaluates `||V - WH||_F^2` as
47/// `||V||_F^2 - 2<W, V H^T> + <W^T W, H H^T>`. The middle term uses the
48/// precomputed `V H^T` rather than `W^T V`, so a changing W costs no extra pass
49/// over V. Accumulated in `f64` regardless of `F`, matching
50/// [`super::compute_objective`], because the two large terms nearly cancel.
51///
52/// ### Params
53///
54/// * `sq_frob_v` - Precomputed `||V||_F^2`.
55/// * `w` - Left factor `m x k`.
56/// * `vht` - Product `V H^T`, shape `m x k`.
57/// * `wtw` - Gram matrix `W^T W`, shape `k x k`.
58/// * `hht` - Gram matrix `H H^T`, shape `k x k`.
59///
60/// ### Returns
61///
62/// The scalar reconstruction error, clamped at zero.
63fn objective_fixed_h<F: BixverseFloat + Send + Sync>(
64    sq_frob_v: F,
65    w: MatRef<F>,
66    vht: MatRef<F>,
67    wtw: MatRef<F>,
68    hht: MatRef<F>,
69) -> F {
70    let m = w.nrows();
71    let k = w.ncols();
72
73    let inner_w_vht: f64 = (0..m)
74        .into_par_iter()
75        .with_min_len(64)
76        .map(|i| {
77            let mut acc = 0f64;
78            for r in 0..k {
79                acc += w[(i, r)].to_f64().unwrap() * vht[(i, r)].to_f64().unwrap();
80            }
81            acc
82        })
83        .reduce(|| 0f64, |x, y| x + y);
84
85    let mut inner_grams = 0f64;
86    for j in 0..k {
87        for i in 0..k {
88            inner_grams += wtw[(i, j)].to_f64().unwrap() * hht[(i, j)].to_f64().unwrap();
89        }
90    }
91
92    let result = sq_frob_v.to_f64().unwrap() - 2.0 * inner_w_vht + inner_grams;
93    F::from_f64(result.max(0.0)).unwrap()
94}
95
96////////////
97// Refits //
98////////////
99
100/// Refit H against a frozen W.
101///
102/// Initialises H at `opts.eps` and runs HALS row sweeps until the relative
103/// objective decrease drops below `opts.tol` or `opts.max_iter` is reached.
104/// `W^T W` and `W^T V` are loop invariants here, so only `H H^T` is recomputed
105/// when the objective is evaluated. No column normalisation is applied: W is
106/// frozen, so its scale convention is the caller's.
107///
108/// ### Params
109///
110/// * `v` - Input matrix backend.
111/// * `w` - Frozen left factor, shape `m x k`. Must match `v`'s row count.
112/// * `opts` - Solver options. The `init` field is ignored.
113///
114/// ### Returns
115///
116/// Tuple of `(H of shape k x n, final ||V - WH||_F^2)`, or
117/// [`MethodsErrors::NmfDimensionMismatch`] if `w` does not match `v`.
118pub fn nmf_refit_h<F, In>(
119    v: &In,
120    w: MatRef<F>,
121    opts: &HalsOpts<F>,
122) -> Result<(Mat<F>, F), MethodsErrors>
123where
124    F: BixverseFloat + Send + Sync,
125    In: NmfInput<F> + Sync,
126{
127    let (m, n) = v.shape();
128    if w.nrows() != m {
129        return Err(MethodsErrors::NmfDimensionMismatch {
130            expected: m,
131            found: w.nrows(),
132        });
133    }
134    let k = w.ncols();
135    let sq_frob = v.sq_frob();
136    let denom = relative_denominator(sq_frob);
137
138    // Both loop invariants, because W never changes.
139    let mut wtw = Mat::<F>::zeros(k, k);
140    gram_wt_w(w, &mut wtw);
141    let mut wtv = Mat::<F>::zeros(k, n);
142    v.wt_v(w, &mut wtv);
143
144    let mut h = Mat::<F>::full(k, n, opts.eps);
145    let mut hht = Mat::<F>::zeros(k, k);
146
147    let mut final_loss = F::infinity();
148    let mut last_loss = F::infinity();
149    let mut n_iter = 0usize;
150
151    for iter in 0..opts.max_iter {
152        n_iter = iter + 1;
153        hals_sweep_rows(&mut h, wtw.as_ref(), wtv.as_ref(), opts.eps);
154
155        if n_iter.is_multiple_of(opts.check_every) {
156            gram_h_ht(h.as_ref(), &mut hht);
157            let loss = compute_objective(
158                sq_frob,
159                h.as_ref(),
160                wtv.as_ref(),
161                wtw.as_ref(),
162                hht.as_ref(),
163            );
164            final_loss = loss;
165
166            if n_iter > opts.check_every && (last_loss - loss).abs() / denom < opts.tol {
167                break;
168            }
169            last_loss = loss;
170        }
171    }
172
173    // `n_iter == 0` when `max_iter` is 0, and zero is a multiple of everything,
174    // so without the first clause the loss would be returned as infinity.
175    if n_iter == 0 || !n_iter.is_multiple_of(opts.check_every) {
176        gram_h_ht(h.as_ref(), &mut hht);
177        final_loss = compute_objective(
178            sq_frob,
179            h.as_ref(),
180            wtv.as_ref(),
181            wtw.as_ref(),
182            hht.as_ref(),
183        );
184    }
185
186    Ok((h, final_loss))
187}
188
189/// Refit W against a frozen H.
190///
191/// Mirror of [`nmf_refit_h`]. `H H^T` and `V H^T` are the loop invariants, and
192/// the objective goes through `objective_fixed_h` so that a changing W never
193/// forces another pass over V. No row normalisation is applied: H is frozen, so
194/// its scale convention is the caller's.
195///
196/// ### Params
197///
198/// * `v` - Input matrix backend.
199/// * `h` - Frozen right factor, shape `k x n`. Must match `v`'s column count.
200/// * `opts` - Solver options. The `init` field is ignored.
201///
202/// ### Returns
203///
204/// Tuple of `(W of shape m x k, final ||V - WH||_F^2)`, or
205/// [`MethodsErrors::NmfDimensionMismatch`] if `h` does not match `v`.
206pub fn nmf_refit_w<F, In>(
207    v: &In,
208    h: MatRef<F>,
209    opts: &HalsOpts<F>,
210) -> Result<(Mat<F>, F), MethodsErrors>
211where
212    F: BixverseFloat + Send + Sync,
213    In: NmfInput<F> + Sync,
214{
215    let (m, n) = v.shape();
216    if h.ncols() != n {
217        return Err(MethodsErrors::NmfDimensionMismatch {
218            expected: n,
219            found: h.ncols(),
220        });
221    }
222    let k = h.nrows();
223    let sq_frob = v.sq_frob();
224    let denom = relative_denominator(sq_frob);
225
226    // Both loop invariants, because H never changes.
227    let mut hht = Mat::<F>::zeros(k, k);
228    gram_h_ht(h, &mut hht);
229    let mut vht = Mat::<F>::zeros(m, k);
230    v.v_ht(h, &mut vht);
231
232    let mut w = Mat::<F>::full(m, k, opts.eps);
233    let mut wtw = Mat::<F>::zeros(k, k);
234
235    let mut final_loss = F::infinity();
236    let mut last_loss = F::infinity();
237    let mut n_iter = 0usize;
238
239    for iter in 0..opts.max_iter {
240        n_iter = iter + 1;
241        hals_sweep_cols(&mut w, hht.as_ref(), vht.as_ref(), opts.eps);
242
243        if n_iter.is_multiple_of(opts.check_every) {
244            gram_wt_w(w.as_ref(), &mut wtw);
245            let loss = objective_fixed_h(
246                sq_frob,
247                w.as_ref(),
248                vht.as_ref(),
249                wtw.as_ref(),
250                hht.as_ref(),
251            );
252            final_loss = loss;
253
254            if n_iter > opts.check_every && (last_loss - loss).abs() / denom < opts.tol {
255                break;
256            }
257            last_loss = loss;
258        }
259    }
260
261    // `n_iter == 0` when `max_iter` is 0, and zero is a multiple of everything,
262    // so without the first clause the loss would be returned as infinity.
263    if n_iter == 0 || !n_iter.is_multiple_of(opts.check_every) {
264        gram_wt_w(w.as_ref(), &mut wtw);
265        final_loss = objective_fixed_h(
266            sq_frob,
267            w.as_ref(),
268            vht.as_ref(),
269            wtw.as_ref(),
270            hht.as_ref(),
271        );
272    }
273
274    Ok((w, final_loss))
275}
276
277///////////
278// Tests //
279///////////
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284    use crate::methods::nmf_hals::NmfInit;
285    use crate::methods::nmf_hals::dense::DenseInput;
286    use approx::assert_relative_eq;
287
288    /// Exact rank-k non-negative product plus its squared Frobenius norm.
289    fn rank_k_product(m: usize, n: usize, k: usize) -> (Mat<f64>, Mat<f64>, Mat<f64>, f64) {
290        let w_true: Mat<f64> = Mat::from_fn(m, k, |i, j| ((i * 7 + j * 3) % 5 + 1) as f64);
291        let h_true: Mat<f64> = Mat::from_fn(k, n, |i, j| ((i * 5 + j * 2) % 4 + 1) as f64);
292        let v: Mat<f64> = Mat::from_fn(m, n, |i, j| {
293            (0..k).map(|r| w_true[(i, r)] * h_true[(r, j)]).sum()
294        });
295        let sq_frob = (0..m)
296            .map(|i| (0..n).map(|j| v[(i, j)].powi(2)).sum::<f64>())
297            .sum();
298        (w_true, h_true, v, sq_frob)
299    }
300
301    /// Squared Frobenius reconstruction error of an explicit `W H` product.
302    fn reconstruction_error(w: MatRef<f64>, h: MatRef<f64>, v: MatRef<f64>) -> f64 {
303        let k = w.ncols();
304        (0..v.nrows())
305            .map(|i| {
306                (0..v.ncols())
307                    .map(|j| {
308                        let wh: f64 = (0..k).map(|r| w[(i, r)] * h[(r, j)]).sum();
309                        (v[(i, j)] - wh).powi(2)
310                    })
311                    .sum::<f64>()
312            })
313            .sum()
314    }
315
316    /// With the true W frozen the system is exactly solvable, so the refit must
317    /// drive the reconstruction to zero. Note the contract is the reconstruction,
318    /// not recovery of `h_true`: the frozen factor need not have full rank, in
319    /// which case many exact non-negative H exist.
320    #[test]
321    fn refit_h_reconstructs_exactly() {
322        let (m, n, k) = (30, 20, 3);
323        let (w_true, _, v, sq_frob) = rank_k_product(m, n, k);
324
325        let input = DenseInput::new(v.as_ref()).unwrap();
326        let opts = HalsOpts::<f64>::new(600, 1e-12, 1e-12, 5, NmfInit::Nndsvd);
327        let (h, loss) = nmf_refit_h(&input, w_true.as_ref(), &opts).unwrap();
328
329        assert!(
330            loss / sq_frob < 1e-8,
331            "reported rel loss {}",
332            loss / sq_frob
333        );
334        // The reported loss must agree with an explicit W H, not just be small.
335        let direct = reconstruction_error(w_true.as_ref(), h.as_ref(), v.as_ref());
336        assert!(
337            direct / sq_frob < 1e-8,
338            "direct rel loss {}",
339            direct / sq_frob
340        );
341        assert!((0..k).all(|r| (0..n).all(|j| h[(r, j)] >= 0.0)));
342    }
343
344    /// Mirror of the above with the roles swapped.
345    #[test]
346    fn refit_w_reconstructs_exactly() {
347        let (m, n, k) = (30, 20, 3);
348        let (_, h_true, v, sq_frob) = rank_k_product(m, n, k);
349
350        let input = DenseInput::new(v.as_ref()).unwrap();
351        let opts = HalsOpts::<f64>::new(600, 1e-12, 1e-12, 5, NmfInit::Nndsvd);
352        let (w, loss) = nmf_refit_w(&input, h_true.as_ref(), &opts).unwrap();
353
354        assert!(
355            loss / sq_frob < 1e-8,
356            "reported rel loss {}",
357            loss / sq_frob
358        );
359        let direct = reconstruction_error(w.as_ref(), h_true.as_ref(), v.as_ref());
360        assert!(
361            direct / sq_frob < 1e-8,
362            "direct rel loss {}",
363            direct / sq_frob
364        );
365        assert!((0..m).all(|i| (0..k).all(|r| w[(i, r)] >= 0.0)));
366    }
367
368    /// When the frozen factor has full rank the NNLS solution is unique, so the
369    /// true partner is recovered outright.
370    #[test]
371    fn refit_recovers_unique_partner() {
372        // Full-rank 2 x 3 frozen H, so W is pinned by the normal equations.
373        let h_true: Mat<f64> = Mat::from_fn(2, 3, |i, j| if i == j { 2.0 } else { 0.5 });
374        let w_true: Mat<f64> = Mat::from_fn(5, 2, |i, j| (i + 2 * j + 1) as f64);
375        let v: Mat<f64> = Mat::from_fn(5, 3, |i, j| {
376            (0..2).map(|r| w_true[(i, r)] * h_true[(r, j)]).sum()
377        });
378
379        let input = DenseInput::new(v.as_ref()).unwrap();
380        let opts = HalsOpts::<f64>::new(2000, 1e-14, 1e-14, 5, NmfInit::Nndsvd);
381        let (w, _) = nmf_refit_w(&input, h_true.as_ref(), &opts).unwrap();
382
383        for i in 0..5 {
384            for r in 0..2 {
385                assert_relative_eq!(w[(i, r)], w_true[(i, r)], epsilon = 1e-6);
386            }
387        }
388    }
389
390    /// The objective must never increase across the refit, which is the property
391    /// that makes the convergence test meaningful.
392    #[test]
393    fn refit_loss_is_monotone() {
394        let (m, n, k) = (20, 15, 3);
395        let (w_true, _, v, _) = rank_k_product(m, n, k);
396        let input = DenseInput::new(v.as_ref()).unwrap();
397
398        let mut previous = f64::INFINITY;
399        for max_iter in [5, 10, 20, 40, 80] {
400            let opts = HalsOpts::<f64>::new(max_iter, 0.0, 1e-12, 5, NmfInit::Nndsvd);
401            let (_, loss) = nmf_refit_h(&input, w_true.as_ref(), &opts).unwrap();
402            assert!(
403                loss <= previous + 1e-9,
404                "loss rose from {previous} to {loss} at max_iter {max_iter}"
405            );
406            previous = loss;
407        }
408    }
409
410    /// The fixed-H objective shortcut must agree with the shared one, which goes
411    /// through `W^T V` instead of `V H^T`.
412    #[test]
413    fn fixed_h_objective_matches_shared() {
414        let (m, n, k) = (12, 9, 3);
415        let (w_true, h_true, v, sq_frob) = rank_k_product(m, n, k);
416        // Perturb W so the two terms do not cancel to exactly zero.
417        let w: Mat<f64> = Mat::from_fn(m, k, |i, j| w_true[(i, j)] + 0.3 * (i + j) as f64);
418
419        let input = DenseInput::new(v.as_ref()).unwrap();
420
421        let mut wtw = Mat::<f64>::zeros(k, k);
422        gram_wt_w(w.as_ref(), &mut wtw);
423        let mut hht = Mat::<f64>::zeros(k, k);
424        gram_h_ht(h_true.as_ref(), &mut hht);
425        let mut wtv = Mat::<f64>::zeros(k, n);
426        input.wt_v(w.as_ref(), &mut wtv);
427        let mut vht = Mat::<f64>::zeros(m, k);
428        input.v_ht(h_true.as_ref(), &mut vht);
429
430        let shared = compute_objective(
431            sq_frob,
432            h_true.as_ref(),
433            wtv.as_ref(),
434            wtw.as_ref(),
435            hht.as_ref(),
436        );
437        let fixed = objective_fixed_h(
438            sq_frob,
439            w.as_ref(),
440            vht.as_ref(),
441            wtw.as_ref(),
442            hht.as_ref(),
443        );
444
445        assert_relative_eq!(shared, fixed, max_relative = 1e-10);
446    }
447
448    /// `max_iter = 0` runs no sweeps, but the loss must still be evaluated on the
449    /// starting factor. Returning infinity would propagate into
450    /// `ConsensusNmfResult::error`, and `HalsOpts` is reachable from R.
451    #[test]
452    fn refit_with_zero_iterations_reports_a_finite_loss() {
453        let (m, n, k) = (12, 9, 2);
454        let (w_true, h_true, v, _) = rank_k_product(m, n, k);
455        let input = DenseInput::new(v.as_ref()).unwrap();
456        let opts = HalsOpts::<f64>::new(0, 1e-9, 1e-12, 10, NmfInit::Nndsvd);
457
458        let (h, loss_h) = nmf_refit_h(&input, w_true.as_ref(), &opts).unwrap();
459        assert!(loss_h.is_finite(), "H refit loss {loss_h}");
460        assert_eq!((h.nrows(), h.ncols()), (k, n));
461
462        let (w, loss_w) = nmf_refit_w(&input, h_true.as_ref(), &opts).unwrap();
463        assert!(loss_w.is_finite(), "W refit loss {loss_w}");
464        assert_eq!((w.nrows(), w.ncols()), (m, k));
465    }
466
467    /// A factor whose shared dimension does not match V is a caller error and
468    /// must not be silently truncated by the GEMMs.
469    #[test]
470    fn refit_rejects_dimension_mismatch() {
471        let (_, _, v, _) = rank_k_product(10, 8, 2);
472        let input = DenseInput::new(v.as_ref()).unwrap();
473        let opts = HalsOpts::<f64>::default();
474
475        let bad_w: Mat<f64> = Mat::full(7, 2, 1.0);
476        assert!(nmf_refit_h(&input, bad_w.as_ref(), &opts).is_err());
477
478        let bad_h: Mat<f64> = Mat::full(2, 5, 1.0);
479        assert!(nmf_refit_w(&input, bad_h.as_ref(), &opts).is_err());
480    }
481}
bixverse_methods/methods/nmf_hals/refit.rs

bixverse_methods/methods/nmf_hals/
refit.rs
