Skip to main content

legume_numeric/matrix/
dmatrix_rsvd.rs

1use crate::matrix::traits::*;
2use nalgebra::{DMatrix, DVector};
3use nalgebra_sparse::{csc::CscMatrix, csr::CsrMatrix};
4
5/// Fixed start-vector seed for the randomized-SVD subspace iteration. The
6/// iteration converges onto the dominant subspace, so pinning the start makes
7/// `rsvd` reproducible without altering the subspace it recovers.
8const RSVD_SUBSPACE_SEED: u64 = 0x5253_5644_5342_5350; // "RSVDSBSP"
9
10/// Compute the Nystrom basis: `U * diag(1 / (s + eps))`.
11///
12/// Given the left singular vectors `u` and singular values `s` from an SVD,
13/// returns the pseudo-inverted projection matrix used for out-of-sample
14/// Nystrom extension.
15pub fn nystrom_basis(u: &DMatrix<f32>, s: &DVector<f32>) -> DMatrix<f32> {
16    let eps = 1e-8;
17    let sinv = DVector::from_iterator(s.len(), s.iter().map(|&si| 1.0 / (si + eps)));
18    u * DMatrix::from_diagonal(&sinv)
19}
20
21trait IntoDense<OutMat> {
22    fn matmul(&self, other: &OutMat) -> OutMat;
23    fn transpose_matmul(&self, other: &OutMat) -> OutMat;
24    fn num_rows(&self) -> usize;
25    fn num_columns(&self) -> usize;
26}
27
28impl<T> IntoDense<DMatrix<T>> for DMatrix<T>
29where
30    T: nalgebra::RealField + num_traits::Float + Copy,
31{
32    fn matmul(&self, other: &DMatrix<T>) -> DMatrix<T> {
33        self * other
34    }
35
36    fn transpose_matmul(&self, other: &DMatrix<T>) -> DMatrix<T> {
37        self.transpose() * other
38    }
39
40    fn num_rows(&self) -> usize {
41        self.nrows()
42    }
43    fn num_columns(&self) -> usize {
44        self.ncols()
45    }
46}
47
48impl<T> IntoDense<DMatrix<T>> for CscMatrix<T>
49where
50    T: nalgebra::RealField + num_traits::Float + Copy,
51{
52    fn matmul(&self, other: &DMatrix<T>) -> DMatrix<T> {
53        self * other
54    }
55    fn transpose_matmul(&self, other: &DMatrix<T>) -> DMatrix<T> {
56        self.transpose() * other
57    }
58    fn num_rows(&self) -> usize {
59        self.nrows()
60    }
61    fn num_columns(&self) -> usize {
62        self.ncols()
63    }
64}
65
66impl<T> IntoDense<DMatrix<T>> for CsrMatrix<T>
67where
68    T: nalgebra::RealField + num_traits::Float + Copy,
69{
70    fn matmul(&self, other: &DMatrix<T>) -> DMatrix<T> {
71        self * other
72    }
73    fn transpose_matmul(&self, other: &DMatrix<T>) -> DMatrix<T> {
74        self.transpose() * other
75    }
76    fn num_rows(&self) -> usize {
77        self.nrows()
78    }
79    fn num_columns(&self) -> usize {
80        self.ncols()
81    }
82}
83
84fn _subspace_iteration<T, D>(
85    xx: &D,
86    rank_and_oversample: usize,
87    power_iters: usize,
88) -> anyhow::Result<DMatrix<T>>
89where
90    T: nalgebra::RealField + num_traits::Float + Copy,
91    D: IntoDense<DMatrix<T>>,
92{
93    let nc = xx.num_columns();
94    // Fixed seed: with enough power iterations the subspace converges onto
95    // the dominant one regardless of the start, so a pinned (rather than
96    // entropy) draw makes the whole randomized SVD reproducible run-to-run —
97    // which in turn pins every downstream consumer (binary-sketch collapse,
98    // layout, SVD fits) — without changing what subspace it recovers. With
99    // very few iterations the result does depend on this start.
100    let mut qq = DMatrix::<T>::runif_seeded(nc, rank_and_oversample, RSVD_SUBSPACE_SEED);
101    let half = T::from(0.5).expect("no half found");
102    qq.iter_mut().for_each(|x| *x -= half);
103
104    // Each half-step re-orthonormalises the iterate with a thin QR. The
105    // basis must span exactly the range of the product it came from: a
106    // pivoted LU factor does not (its permutation is lost), and iterating on
107    // a row-permuted range does not converge onto the dominant subspace.
108    for _i in 0..power_iters {
109        let ll = xx.matmul(&qq).qr().q();
110        qq = xx.transpose_matmul(&ll).qr().q();
111    }
112
113    // let qq = DMatrix::<T>::runif(nc, rank_and_oversample);
114
115    let qr_q = xx.matmul(&qq).qr().q();
116    let kk = rank_and_oversample.min(qr_q.ncols());
117    let ret = qr_q.columns(0, kk).into_owned();
118
119    Ok(ret)
120}
121
122fn _randomized_svd<T, D>(
123    xx: &D,
124    max_rank: usize,
125    args: &RsvdArgs,
126) -> anyhow::Result<(DMatrix<T>, DVector<T>, DMatrix<T>)>
127where
128    T: nalgebra::RealField + num_traits::Float + Copy,
129    D: IntoDense<DMatrix<T>>,
130{
131    let nr = xx.num_rows();
132    let nc = xx.num_columns();
133
134    let mut rank = nr.min(nc);
135    let mut oversample = 0;
136
137    if max_rank > 0 && rank > max_rank {
138        rank = max_rank;
139        oversample = args.oversample;
140    }
141
142    anyhow::ensure!(rank > 0, "randomized SVD of an empty {nr} x {nc} matrix");
143
144    // Keep the oversampled basis through the projection: its columns are
145    // not ordered by singular value, so truncating here would discard part
146    // of the dominant subspace. The rank is applied to the small SVD below.
147    // Columns beyond the matrix's own rank add nothing; the cap sits above
148    // anything the default oversampling reaches, so it changes no default
149    // result and only stops an oversized `oversample` from overflowing or
150    // allocating for nothing.
151    let width = rank
152        .saturating_add(oversample)
153        .min(nr.min(nc) + RsvdArgs::default().oversample);
154    let qq = _subspace_iteration(xx, width, args.power_iters)?;
155    let rank = rank.min(qq.ncols());
156
157    // let bb = qq.transpose() * xx
158    let bb = xx.transpose_matmul(&qq).transpose();
159
160    let svd = bb.svd(true, true);
161
162    if let (Some(svd_u), Some(svd_vt)) = (svd.u, svd.v_t) {
163        return Ok((
164            qq.clone() * svd_u.columns(0, rank).into_owned(),
165            svd.singular_values.rows(0, rank).into_owned(),
166            svd_vt.transpose().columns(0, rank).into_owned(),
167        ));
168    }
169    Err(anyhow::anyhow!("randomized SVD failed"))
170}
171
172impl<T> RandomizedAlgs for DMatrix<T>
173where
174    T: nalgebra::RealField + num_traits::Float + Copy,
175{
176    type InMat = DMatrix<T>;
177    type OutMat = DMatrix<T>;
178    type DVec = DVector<T>;
179    type Scalar = T;
180
181    fn rsvd_with(
182        &self,
183        max_rank: usize,
184        args: &RsvdArgs,
185    ) -> anyhow::Result<(Self::OutMat, Self::DVec, Self::OutMat)> {
186        _randomized_svd(self, max_rank, args)
187    }
188}
189
190impl<T> RandomizedAlgs for CscMatrix<T>
191where
192    T: nalgebra::RealField + num_traits::Float + Copy,
193{
194    type InMat = CscMatrix<T>;
195    type OutMat = DMatrix<T>;
196    type DVec = DVector<T>;
197    type Scalar = T;
198
199    fn rsvd_with(
200        &self,
201        max_rank: usize,
202        args: &RsvdArgs,
203    ) -> anyhow::Result<(Self::OutMat, Self::DVec, Self::OutMat)> {
204        _randomized_svd(self, max_rank, args)
205    }
206}
207
208impl<T> RandomizedAlgs for CsrMatrix<T>
209where
210    T: nalgebra::RealField + num_traits::Float + Copy,
211{
212    type InMat = CsrMatrix<T>;
213    type OutMat = DMatrix<T>;
214    type DVec = DVector<T>;
215    type Scalar = T;
216
217    fn rsvd_with(
218        &self,
219        max_rank: usize,
220        args: &RsvdArgs,
221    ) -> anyhow::Result<(Self::OutMat, Self::DVec, Self::OutMat)> {
222        _randomized_svd(self, max_rank, args)
223    }
224}
225
226#[cfg(test)]
227#[path = "dmatrix_rsvd_tests.rs"]
228mod tests;