geographdb-core 0.5.4

//! Low-rank matrix approximation helpers.
//!
//! These are intentionally dependency-light: they use a deterministic
//! randomized subspace-iteration SVD with a tiny Jacobi eigen-decomposition on
//! the small projected covariance matrix.  The main intended use is to prototype
//! low-rank FFN weight approximations in the attention experiment.

/// Tiny deterministic xorshift RNG used for initialization matrices.
struct XorShift32 {
    state: u32,
}

impl XorShift32 {
    fn new(seed: u32) -> Self {
        Self {
            state: seed.wrapping_add(0x9e37_79b9),
        }
    }

    fn next(&mut self) -> u32 {
        let mut x = self.state;
        x ^= x << 13;
        x ^= x >> 17;
        x ^= x << 5;
        self.state = x;
        x
    }

    fn uniform_f32(&mut self) -> f32 {
        let u = self.next();
        (u as f32 / u32::MAX as f32) * 2.0 - 1.0
    }
}

fn matmul(a: &[f32], a_rows: usize, a_cols: usize, b: &[f32], b_cols: usize) -> Vec<f32> {
    assert_eq!(a.len(), a_rows * a_cols);
    assert_eq!(b.len(), a_cols * b_cols);
    let mut c = vec![0.0f32; a_rows * b_cols];
    for i in 0..a_rows {
        for k in 0..a_cols {
            let av = a[i * a_cols + k];
            for j in 0..b_cols {
                c[i * b_cols + j] += av * b[k * b_cols + j];
            }
        }
    }
    c
}

/// Multiply `a^T` (`[a_cols, a_rows]`) by `b` (`[a_rows, b_cols]`).
fn matmul_t_a(a: &[f32], a_rows: usize, a_cols: usize, b: &[f32], b_cols: usize) -> Vec<f32> {
    assert_eq!(a.len(), a_rows * a_cols);
    assert_eq!(b.len(), a_rows * b_cols);
    let mut c = vec![0.0f32; a_cols * b_cols];
    for j in 0..a_cols {
        for k in 0..a_rows {
            let av = a[k * a_cols + j];
            for l in 0..b_cols {
                c[j * b_cols + l] += av * b[k * b_cols + l];
            }
        }
    }
    c
}

fn transpose(mat: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    let mut out = vec![0.0f32; rows * cols];
    for i in 0..rows {
        for j in 0..cols {
            out[j * rows + i] = mat[i * cols + j];
        }
    }
    out
}

/// In-place modified Gram-Schmidt orthonormalization of a row-major matrix
/// `[rows, cols]` column by column.
fn orthonormalize_columns(q: &mut [f32], rows: usize, cols: usize) {
    for r in 0..cols {
        for prev in 0..r {
            let mut proj = 0.0f32;
            for i in 0..rows {
                proj += q[i * cols + prev] * q[i * cols + r];
            }
            for i in 0..rows {
                q[i * cols + r] -= proj * q[i * cols + prev];
            }
        }
        let mut norm = 0.0f32;
        for i in 0..rows {
            let v = q[i * cols + r];
            norm += v * v;
        }
        norm = norm.sqrt();
        if norm > 1e-12 {
            for i in 0..rows {
                q[i * cols + r] /= norm;
            }
        }
    }
}

/// Eigendecomposition of a symmetric matrix using cyclic Jacobi rotations.
///
/// `a` is a row-major `[n, n]` symmetric matrix.  Returns `(eigenvalues, eigenvectors)`
/// where eigenvectors are stored row-major `[n, n]` with eigenvectors as columns.
fn jacobi_eigh(a: &[f32], n: usize) -> (Vec<f32>, Vec<f32>) {
    let mut a_mat = a.to_vec();
    let mut v_mat = vec![0.0f32; n * n];
    for i in 0..n {
        v_mat[i * n + i] = 1.0;
    }

    let max_iter = 50.max(n * n);
    let eps = 1e-10f32;

    for _ in 0..max_iter {
        let mut max_val = 0.0f32;
        let mut p = 0usize;
        let mut q = 1usize;
        for i in 0..n {
            for j in (i + 1)..n {
                let val = a_mat[i * n + j].abs();
                if val > max_val {
                    max_val = val;
                    p = i;
                    q = j;
                }
            }
        }
        if max_val < eps {
            break;
        }

        let app = a_mat[p * n + p];
        let aqq = a_mat[q * n + q];
        let apq = a_mat[p * n + q];
        let phi = 0.5 * (2.0 * apq).atan2(aqq - app);
        let c = phi.cos();
        let s = phi.sin();

        // Rotate columns/rows of A.
        for i in 0..n {
            if i == p || i == q {
                continue;
            }
            let aip = a_mat[i * n + p];
            let aiq = a_mat[i * n + q];
            a_mat[i * n + p] = c * aip - s * aiq;
            a_mat[p * n + i] = a_mat[i * n + p];
            a_mat[i * n + q] = s * aip + c * aiq;
            a_mat[q * n + i] = a_mat[i * n + q];
        }
        let app_new = c * c * app - 2.0 * s * c * apq + s * s * aqq;
        let aqq_new = s * s * app + 2.0 * s * c * apq + c * c * aqq;
        a_mat[p * n + p] = app_new;
        a_mat[q * n + q] = aqq_new;
        a_mat[p * n + q] = 0.0;
        a_mat[q * n + p] = 0.0;

        // Accumulate eigenvectors: V <- V * J.
        for i in 0..n {
            let vip = v_mat[i * n + p];
            let viq = v_mat[i * n + q];
            v_mat[i * n + p] = c * vip - s * viq;
            v_mat[i * n + q] = s * vip + c * viq;
        }
    }

    let mut evals = Vec::with_capacity(n);
    for i in 0..n {
        evals.push(a_mat[i * n + i]);
    }
    (evals, v_mat)
}

/// Truncated randomized SVD.
///
/// Returns `(u, s, v)` such that `a ≈ u * diag(s) * v^T`.
/// - `u` is row-major `[rows, rank]`
/// - `s` has length `rank`
/// - `v` is row-major `[cols, rank]` (columns are right singular vectors)
pub fn truncated_svd(
    a: &[f32],
    rows: usize,
    cols: usize,
    rank: usize,
    n_iter: usize,
    seed: u32,
) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
    assert_eq!(a.len(), rows * cols, "truncated_svd: size mismatch");
    assert!(rank > 0, "truncated_svd: rank must be positive");
    assert!(
        rank <= rows.min(cols),
        "truncated_svd: rank exceeds matrix dimensions"
    );

    // Random Gaussian test matrix and initial subspace estimate.
    let mut rng = XorShift32::new(seed);
    let mut omega: Vec<f32> = (0..cols * rank).map(|_| rng.uniform_f32()).collect();
    orthonormalize_columns(&mut omega, cols, rank);

    let mut q = matmul(a, rows, cols, &omega, rank);
    orthonormalize_columns(&mut q, rows, rank);

    for _ in 0..n_iter {
        let atq = matmul_t_a(a, rows, cols, &q, rank);
        let aatq = matmul(a, rows, cols, &atq, rank);
        q = aatq;
        orthonormalize_columns(&mut q, rows, rank);
    }

    // Small projected matrix B = Q^T A.
    let qt = transpose(&q, rows, rank);
    let b = matmul(&qt, rank, rows, a, cols);

    // Eigendecomposition of B B^T gives singular values/vectors of B.
    let bt = transpose(&b, rank, cols);
    let c = matmul(&b, rank, cols, &bt, rank);
    let (evals, evecs) = jacobi_eigh(&c, rank);

    // Sort eigenvalues (and corresponding eigenvectors) descending.
    let mut order: Vec<usize> = (0..rank).collect();
    order.sort_by(|&i, &j| evals[j].partial_cmp(&evals[i]).unwrap());

    let mut u = vec![0.0f32; rows * rank];
    let mut s = vec![0.0f32; rank];
    let mut eb = vec![0.0f32; rank * rank]; // eigenvectors of C sorted as columns.
    for (new_idx, &old_idx) in order.iter().enumerate() {
        s[new_idx] = evals[old_idx].max(0.0).sqrt();
        for p in 0..rank {
            eb[p * rank + new_idx] = evecs[p * rank + old_idx];
        }
    }

    // U = Q * E_B (rows x rank).
    u = matmul(&q, rows, rank, &eb, rank);

    // V = (B^T * E_B) * diag(1/sigma).
    let mut v = vec![0.0f32; cols * rank];
    let bt_eb = matmul(&bt, cols, rank, &eb, rank);
    for r in 0..rank {
        let inv_sigma = if s[r] > 1e-12 { 1.0 / s[r] } else { 0.0 };
        for j in 0..cols {
            v[j * rank + r] = bt_eb[j * rank + r] * inv_sigma;
        }
    }

    (u, s, v)
}

/// Build balanced low-rank factors from a truncated SVD.
///
/// Returns `(left, right)` such that `left * right^T ≈ a`.  `left` has shape
/// `[rows, rank]` and `right` has shape `[cols, rank]`, both stored row-major.
pub fn balanced_low_rank_factors(
    a: &[f32],
    rows: usize,
    cols: usize,
    rank: usize,
    n_iter: usize,
    seed: u32,
) -> (Vec<f32>, Vec<f32>) {
    let (u, s, v) = truncated_svd(a, rows, cols, rank, n_iter, seed);
    let mut left = Vec::with_capacity(rows * rank);
    let mut right = Vec::with_capacity(cols * rank);
    for i in 0..rows {
        for r in 0..rank {
            left.push(u[i * rank + r] * s[r].sqrt());
        }
    }
    for j in 0..cols {
        for r in 0..rank {
            right.push(v[j * rank + r] * s[r].sqrt());
        }
    }
    (left, right)
}

/// Multiply a row vector `[cols]` by a row-major matrix `[cols, out]` to get
/// `[out]`.
pub fn vec_matmul(v: &[f32], cols: usize, m: &[f32], out: usize) -> Vec<f32> {
    assert_eq!(v.len(), cols);
    assert_eq!(m.len(), cols * out);
    let mut res = vec![0.0f32; out];
    for c in 0..cols {
        let vc = v[c];
        for o in 0..out {
            res[o] += vc * m[c * out + o];
        }
    }
    res
}

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

    #[test]
    fn rank_one_reconstruction() {
        let rows = 8;
        let cols = 10;
        let u: Vec<f32> = (0..rows).map(|i| (i + 1) as f32).collect();
        let v: Vec<f32> = (0..cols).map(|j| (j + 2) as f32).collect();
        let a: Vec<f32> = u
            .iter()
            .flat_map(|&ui| v.iter().map(move |&vj| ui * vj))
            .collect();
        let (u_out, s, v_out) = truncated_svd(&a, rows, cols, 1, 10, 1);
        let mut recon = vec![0.0f32; rows * cols];
        for i in 0..rows {
            for j in 0..cols {
                recon[i * cols + j] = u_out[i] * s[0] * v_out[j];
            }
        }
        let mse: f32 = a
            .iter()
            .zip(recon.iter())
            .map(|(a, b)| (a - b) * (a - b))
            .sum::<f32>()
            / a.len() as f32;
        assert!(mse < 1e-3, "rank-1 reconstruction MSE = {}", mse);
    }

    #[test]
    fn low_rank_factors_reconstruct() {
        let rows = 12;
        let cols = 16;
        let u1: Vec<f32> = (0..rows)
            .map(|i| if i < rows / 2 { 1.0 } else { -1.0 })
            .collect();
        let u2: Vec<f32> = (0..rows)
            .map(|i| if i % 2 == 0 { 1.0 } else { -1.0 })
            .collect();
        let v1: Vec<f32> = (0..cols)
            .map(|j| if j < cols / 2 { 1.0 } else { -1.0 })
            .collect();
        let v2: Vec<f32> = (0..cols)
            .map(|j| if j % 2 == 0 { 1.0 } else { -1.0 })
            .collect();
        let mut a = vec![0.0f32; rows * cols];
        for i in 0..rows {
            for j in 0..cols {
                a[i * cols + j] = u1[i] * v1[j] + 2.0 * u2[i] * v2[j];
            }
        }
        let (left, right) = balanced_low_rank_factors(&a, rows, cols, 2, 10, 7);
        let mut recon = vec![0.0f32; rows * cols];
        for i in 0..rows {
            for j in 0..cols {
                let mut acc = 0.0f32;
                for r in 0..2 {
                    acc += left[i * 2 + r] * right[j * 2 + r];
                }
                recon[i * cols + j] = acc;
            }
        }
        let mse = a
            .iter()
            .zip(recon.iter())
            .map(|(x, y)| (x - y) * (x - y))
            .sum::<f32>()
            / a.len() as f32;
        assert!(mse < 1e-2, "low-rank reconstruction MSE = {}", mse);
    }
}