use faer::{Mat, MatMut, MatRef};
use num_traits::Float;
use rand::prelude::*;
use rand_distr::Normal;
use rayon::prelude::*;
use super::*;
use crate::prelude::*;
#[derive(Clone, Debug)]
pub struct RandomSvdResults<T> {
pub u: Mat<T>,
pub v: Mat<T>,
pub s: Vec<T>,
}
#[derive(Clone, Debug)]
pub struct SvdResults<T> {
pub u: Mat<T>,
pub v: Mat<T>,
pub s: Vec<T>,
}
pub trait SvdResult<T> {
fn u(&self) -> &faer::Mat<T>;
fn v(&self) -> &faer::Mat<T>;
fn s(&self) -> &[T];
}
impl<T> SvdResult<T> for RandomSvdResults<T> {
fn u(&self) -> &faer::Mat<T> {
&self.u
}
fn v(&self) -> &faer::Mat<T> {
&self.v
}
fn s(&self) -> &[T] {
&self.s
}
}
impl<T> SvdResult<T> for SvdResults<T> {
fn u(&self) -> &faer::Mat<T> {
&self.u
}
fn v(&self) -> &faer::Mat<T> {
&self.v
}
fn s(&self) -> &[T] {
&self.s
}
}
pub fn compute_pc_scores<T, S>(svd_results: &S) -> Mat<T>
where
T: Float,
S: SvdResult<T>,
{
let n_cells = svd_results.u().nrows();
let n_pcs = svd_results.s().len();
Mat::from_fn(n_cells, n_pcs, |i, j| {
svd_results.u()[(i, j)] * svd_results.s()[j]
})
}
pub fn get_top_eigenvalues<T>(
matrix: &Mat<T>,
top_n: usize,
) -> Result<Vec<(T, Vec<T>)>, BixverseErrors>
where
T: BixverseFloat,
{
assert_symmetric_mat!(matrix);
let eigendecomp = matrix.eigen().map_err(|_| BixverseErrors::FaerEigenError)?;
let s = eigendecomp.S();
let u = eigendecomp.U();
let mut eigenpairs = s
.column_vector()
.iter()
.zip(u.col_iter())
.map(|(l, v)| {
let l_real = l.re;
let v_real = v.iter().map(|v_i| v_i.re).collect::<Vec<T>>();
(l_real, v_real)
})
.collect::<Vec<(T, Vec<T>)>>();
eigenpairs.sort_by(|a, b| b.0.total_cmp(&a.0));
let res: Vec<(T, Vec<T>)> = eigenpairs.into_iter().take(top_n).collect();
Ok(res)
}
pub fn randomised_svd<T>(
x: MatRef<T>,
rank: usize,
seed: usize,
oversampling: Option<usize>,
n_power_iter: Option<usize>,
) -> Result<RandomSvdResults<T>, BixverseErrors>
where
T: BixverseFloat,
{
let ncol = x.ncols();
let nrow = x.nrows();
let os = oversampling.unwrap_or(DEFAULT_OVERSAMPLING_RAND_SVD);
let sample_size = (rank + os).min(ncol.min(nrow));
let n_iter = n_power_iter.unwrap_or(DEFAULT_N_POWER_ITERS_RAND_SVD);
let mut rng = StdRng::seed_from_u64(seed as u64);
let normal = Normal::new(0.0, 1.0).unwrap();
let omega = Mat::from_fn(ncol, sample_size, |_, _| {
T::from_f64(normal.sample(&mut rng)).unwrap()
});
let y = x * omega;
let mut q = y.qr().compute_thin_Q();
for _ in 0..n_iter {
let z = x.transpose() * q;
q = (x * z).qr().compute_thin_Q();
}
let b = q.transpose() * x;
let svd = b
.thin_svd()
.map_err(|e| BixverseErrors::FaerSvdError(format!("{e:?}")))?;
Ok(RandomSvdResults {
u: q * svd.U(),
v: svd.V().cloned(),
s: svd.S().column_vector().iter().copied().collect(),
})
}
#[allow(clippy::too_many_arguments)]
pub fn randomised_svd_matfree<T, FA, FT>(
n_rows: usize,
n_cols: usize,
rank: usize,
seed: usize,
oversampling: Option<usize>,
n_power_iter: Option<usize>,
apply: FA,
apply_t: FT,
) -> Result<RandomSvdResults<T>, BixverseErrors>
where
T: BixverseFloat,
FA: Fn(MatRef<T>) -> Mat<T>,
FT: Fn(MatRef<T>) -> Mat<T>,
{
let os = oversampling.unwrap_or(DEFAULT_OVERSAMPLING_RAND_SVD);
let sample_size = (rank + os).min(n_cols.min(n_rows));
let n_iter = n_power_iter.unwrap_or(DEFAULT_N_POWER_ITERS_RAND_SVD);
let mut rng = StdRng::seed_from_u64(seed as u64);
let normal = Normal::new(0.0, 1.0).unwrap();
let omega = Mat::from_fn(n_cols, sample_size, |_, _| {
T::from_f64(normal.sample(&mut rng)).unwrap()
});
let y = apply(omega.as_ref());
let mut q = y.qr().compute_thin_Q();
for _ in 0..n_iter {
let z = apply_t(q.as_ref());
q = apply(z.as_ref()).qr().compute_thin_Q();
}
let bt = apply_t(q.as_ref());
let b = bt.transpose().to_owned();
let svd = b
.thin_svd()
.map_err(|e| BixverseErrors::FaerSvdError(format!("{e:?}")))?;
Ok(RandomSvdResults {
u: q * svd.U(),
v: svd.V().cloned(),
s: svd.S().column_vector().iter().copied().collect(),
})
}
#[allow(clippy::too_many_arguments)]
pub fn randomised_sparse_svd<T, F>(
matrix: CompressedSparseData2<T>,
rank: usize,
seed: u64,
use_second_layer: bool,
oversampling: Option<usize>,
n_power_iter: Option<usize>,
col_means: Option<&[F]>,
col_stds: Option<&[F]>,
row_offsets: Option<&[F]>,
) -> Result<RandomSvdResults<F>, BixverseErrors>
where
T: BixverseNumeric + Into<F>,
F: BixverseFloat,
{
let (n, m) = matrix.shape;
let os = oversampling.unwrap_or(DEFAULT_OVERSAMPLING_RAND_SVD);
let sample_size = (rank + os).min(m).min(n);
let n_iter = n_power_iter.unwrap_or(DEFAULT_N_POWER_ITERS_RAND_SVD);
let csr = match matrix.cs_type {
CompressedSparseFormat::Csr => matrix,
CompressedSparseFormat::Csc => matrix.transform_single_layer(use_second_layer)?,
};
let active_data: &[T] = if use_second_layer {
csr.data_2
.as_ref()
.ok_or(BixverseErrors::Data2NotAvailable)?
.as_slice()
} else {
csr.data.as_slice()
};
let val_f64 = |idx: usize| -> f64 {
let v: F = active_data[idx].into();
v.to_f64().unwrap()
};
let prescale = |x: MatRef<F>| -> Option<Mat<F>> {
col_stds.map(|sd| Mat::from_fn(x.nrows(), x.ncols(), |i, col| x[(i, col)] / sd[i]))
};
let sparse_matvec_a = |x_scaled: MatRef<F>, y: MatMut<F>| {
let ncols = x_scaled.ncols();
let mean_dots: Vec<f64> = if let Some(mu) = col_means {
(0..ncols)
.map(|col| {
(0..m)
.map(|j| mu[j].to_f64().unwrap() * x_scaled[(j, col)].to_f64().unwrap())
.sum::<f64>()
})
.collect()
} else {
Vec::new()
};
let x_sums: Vec<f64> = if row_offsets.is_some() {
(0..ncols)
.map(|col| {
(0..m)
.map(|j| x_scaled[(j, col)].to_f64().unwrap())
.sum::<f64>()
})
.collect()
} else {
Vec::new()
};
let y_ptr = y.as_ptr_mut() as usize;
let y_row_stride = y.row_stride();
let y_col_stride = y.col_stride();
(0..n).into_par_iter().for_each(|i| {
let base = y_ptr as *mut F;
let mut row_acc = vec![0f64; ncols];
for idx in csr.indptr[i]..csr.indptr[i + 1] {
let j = csr.indices[idx as usize] as usize;
let idx_usize = idx as usize;
let a_val = val_f64(idx_usize);
for col in 0..ncols {
row_acc[col] += a_val * x_scaled[(j, col)].to_f64().unwrap();
}
}
if col_means.is_some() {
for col in 0..ncols {
row_acc[col] -= mean_dots[col];
}
}
if let Some(off) = row_offsets {
let m_i = off[i].to_f64().unwrap();
for col in 0..ncols {
row_acc[col] -= m_i * x_sums[col];
}
}
for col in 0..ncols {
unsafe {
let ptr = base.offset(i as isize * y_row_stride + col as isize * y_col_stride);
*ptr = F::from_f64(row_acc[col]).unwrap();
}
}
});
};
let sparse_matvec_at = |x: MatRef<F>, mut y: MatMut<F>| {
let ncols = x.ncols();
let m_dot_x: Vec<f64> = if let Some(off) = row_offsets {
(0..ncols)
.map(|col| {
(0..n)
.map(|i| off[i].to_f64().unwrap() * x[(i, col)].to_f64().unwrap())
.sum::<f64>()
})
.collect()
} else {
Vec::new()
};
let (result, col_sums) = (0..n)
.into_par_iter()
.fold(
|| (vec![0f64; m * ncols], vec![0f64; ncols]),
|(mut acc, mut cs), i| {
for col in 0..ncols {
cs[col] += x[(i, col)].to_f64().unwrap();
}
for idx in csr.indptr[i]..csr.indptr[i + 1] {
let idx_usize = idx as usize;
let j = csr.indices[idx_usize] as usize;
let a_val = val_f64(idx_usize);
for col in 0..ncols {
acc[j * ncols + col] += a_val * x[(i, col)].to_f64().unwrap();
}
}
(acc, cs)
},
)
.reduce(
|| (vec![0f64; m * ncols], vec![0f64; ncols]),
|(mut a, mut cs_a), (b, cs_b)| {
for i in 0..a.len() {
a[i] += b[i];
}
for i in 0..ncols {
cs_a[i] += cs_b[i];
}
(a, cs_a)
},
);
for j in 0..m {
for col in 0..ncols {
let mut val = result[j * ncols + col];
if let Some(mu) = col_means {
val -= mu[j].to_f64().unwrap() * col_sums[col];
}
if row_offsets.is_some() {
val -= m_dot_x[col];
}
if let Some(sd) = col_stds {
val /= sd[j].to_f64().unwrap();
}
y[(j, col)] = F::from_f64(val).unwrap();
}
}
};
let mut rng = StdRng::seed_from_u64(seed);
let normal = Normal::new(0.0, 1.0).unwrap();
let omega = Mat::from_fn(m, sample_size, |_, _| {
F::from_f64(normal.sample(&mut rng)).unwrap()
});
let omega_scaled = prescale(omega.as_ref());
let mut y = Mat::<F>::zeros(n, sample_size);
sparse_matvec_a(
omega_scaled
.as_ref()
.map(|m| m.as_ref())
.unwrap_or(omega.as_ref()),
y.as_mut(),
);
drop(omega);
drop(omega_scaled);
let mut q = y.qr().compute_thin_Q();
let mut z = Mat::<F>::zeros(m, sample_size);
let mut y_new = Mat::<F>::zeros(n, sample_size);
for _ in 0..n_iter {
sparse_matvec_at(q.as_ref(), z.as_mut());
let z_scaled = prescale(z.as_ref());
sparse_matvec_a(
z_scaled.as_ref().map(|m| m.as_ref()).unwrap_or(z.as_ref()),
y_new.as_mut(),
);
q = y_new.qr().compute_thin_Q();
}
drop(z);
drop(y_new);
drop(y);
let mut b_t = Mat::<F>::zeros(m, sample_size);
sparse_matvec_at(q.as_ref(), b_t.as_mut());
let b = b_t.transpose().to_owned();
let svd = b
.thin_svd()
.map_err(|e| BixverseErrors::FaerSvdError(format!("{e:?}")))?;
let u = &q * svd.U();
let s: Vec<F> = svd.S().column_vector().iter().copied().collect();
let v = svd.V().to_owned();
Ok(RandomSvdResults { u, s, v })
}
#[cfg(test)]
mod tests {
use super::*;
use faer::Mat;
#[test]
fn test_compute_pc_scores() {
let u: Mat<f64> = Mat::from_fn(2, 2, |i, j| if i == j { 1.0 } else { 0.0 });
let v: Mat<f64> = Mat::from_fn(2, 2, |i, j| if i == j { 1.0 } else { 0.0 });
let s = vec![3.0, 1.5];
let svd_res = SvdResults { u, v, s };
let scores = compute_pc_scores(&svd_res);
assert!((scores[(0, 0)] - 3.0).abs() < 1e-6);
assert!((scores[(0, 1)] - 0.0).abs() < 1e-6);
assert!((scores[(1, 0)] - 0.0).abs() < 1e-6);
assert!((scores[(1, 1)] - 1.5).abs() < 1e-6);
}
#[test]
fn test_get_top_eigenvalues() {
let mat: Mat<f64> = Mat::from_fn(
2,
2,
|i, j| if i == j { 3.0 - (i as f64 * 2.0) } else { 0.0 },
);
let top_eigen = get_top_eigenvalues(&mat, 2).unwrap();
assert_eq!(top_eigen.len(), 2);
assert!((top_eigen[0].0 - 3.0).abs() < 1e-6);
assert!((top_eigen[1].0 - 1.0).abs() < 1e-6);
}
#[test]
fn test_randomised_svd_logic() {
let x = [1.0, 2.0, 3.0, 4.0];
let y = [1.0, 0.5, 0.25];
let mut mat: Mat<f64> = Mat::zeros(4, 3);
for i in 0..4 {
for j in 0..3 {
mat[(i, j)] = x[i] * y[j];
}
}
let svd = randomised_svd(mat.as_ref(), 1, 42, Some(5), Some(4)).unwrap();
let u_col = svd.u.col(0);
let x_norm: f64 = x.iter().map(|v| v * v).sum::<f64>().sqrt();
let mut dot_u = 0.0;
for i in 0..4 {
dot_u += u_col[i] * (x[i] / x_norm);
}
let v_col = svd.v.col(0);
let y_norm: f64 = y.iter().map(|v| v * v).sum::<f64>().sqrt();
let mut dot_v = 0.0;
for j in 0..3 {
dot_v += v_col[j] * (y[j] / y_norm);
}
assert!(dot_u.abs() > 0.999);
assert!(dot_v.abs() > 0.999);
}
#[test]
fn test_randomised_svd_matfree_matches_dense() {
let n_hvg = 5;
let n1 = 6;
let n2 = 4;
let rank = 2;
let x1: Mat<f64> = Mat::from_fn(n_hvg, n1, |i, j| ((i + 1) as f64 * 0.3 + j as f64).sin());
let x2: Mat<f64> = Mat::from_fn(n_hvg, n2, |i, j| ((j + 2) as f64 * 0.7 + i as f64).cos());
let a = x1.transpose() * &x2;
let dense = randomised_svd(a.as_ref(), rank, 42, Some(5), Some(4)).unwrap();
let apply = |v: MatRef<f64>| -> Mat<f64> { x1.transpose() * (&x2 * v) };
let apply_t = |u: MatRef<f64>| -> Mat<f64> { x2.transpose() * (&x1 * u) };
let matfree =
randomised_svd_matfree(n1, n2, rank, 42, Some(5), Some(4), apply, apply_t).unwrap();
for k in 0..rank {
assert!(
(dense.s[k] - matfree.s[k]).abs() < 1e-6,
"singular value {k}: dense {} vs matfree {}",
dense.s[k],
matfree.s[k]
);
}
for k in 0..rank {
let mut dot_u = 0.0_f64;
for i in 0..n1 {
dot_u += dense.u[(i, k)] * matfree.u[(i, k)];
}
let mut dot_v = 0.0_f64;
for j in 0..n2 {
dot_v += dense.v[(j, k)] * matfree.v[(j, k)];
}
assert!(dot_u.abs() > 0.999, "u[{k}] |dot| = {}", dot_u.abs());
assert!(dot_v.abs() > 0.999, "v[{k}] |dot| = {}", dot_v.abs());
}
}
#[test]
fn test_randomised_sparse_svd_logic() {
let data = vec![2.0, 1.0, 4.0, 2.0];
let indices = vec![0, 2, 0, 2];
let indptr = vec![0, 0, 2, 2, 4];
let shape = (4, 3);
let csr = CompressedSparseData2::<f64, f64>::new_csr(&data, &indices, &indptr, None, shape);
let no_params: Option<&[f64]> = None;
let svd = randomised_sparse_svd(
csr,
1,
42,
false,
Some(5),
Some(4),
no_params,
no_params,
None,
)
.unwrap();
let u_col = svd.u.col(0);
let x_norm = (2.0_f64.powi(2) + 4.0_f64.powi(2)).sqrt(); let dot_u = (u_col[1] * 2.0 + u_col[3] * 4.0) / x_norm;
assert!(dot_u.abs() > 0.999);
let v_col = svd.v.col(0);
let y_norm = (1.0_f64.powi(2) + 0.5_f64.powi(2)).sqrt(); let dot_v = (v_col[0] * 1.0 + v_col[2] * 0.5) / y_norm;
assert!(dot_v.abs() > 0.999);
}
}