use super::{FixedVector, FixedMatrix};
use super::tensor::Tensor;
use super::decompose::svd_decompose;
use super::compute_matrix::{compute_lu_decompose, ComputeMatrix};
use super::kernels;
use super::kernels::Orient;
use crate::fixed_point::fixed_num::{Compute, Fixed, FixedNum};
use crate::fixed_point::width::{Storage, WideInt};
use crate::fixed_point::core_types::errors::OverflowDetected;
pub struct TruncatedSVD<T> {
pub u: FixedMatrix<T>,
pub sigma: FixedVector<T>,
pub vt: FixedMatrix<T>,
}
impl<S: Storage, const F: u32> TruncatedSVD<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
pub fn reconstruct(&self) -> FixedMatrix<Fixed<S, F>> {
const OVERFLOW: &str = "TruncatedSVD::reconstruct: entry exceeds storage";
let m = self.u.rows();
let n = self.vt.cols();
let k = self.sigma.len();
let mut u_sigma = Vec::with_capacity(m * k);
for i in 0..m {
for r in 0..k {
u_sigma.push(kernels::exact_dot::<Fixed<S, F>>(&[self.u.get(i, r).raw()], &[self.sigma[r].raw()]).expect(OVERFLOW));
}
}
let mut result = FixedMatrix::<Fixed<S, F>>::new(m, n);
for i in 0..m {
for j in 0..n {
let mut sum = <Orient<Fixed<S, F>> as WideInt>::zero();
for r in 0..k {
sum = sum.add_exact(S::widen_product(u_sigma[i * k + r], S::widen(self.vt.get(r, j).raw()))).expect(OVERFLOW);
}
result.set(i, j, Fixed::from_raw(kernels::narrow_triple_nearest::<Fixed<S, F>>(sum).expect(OVERFLOW)));
}
}
result
}
pub fn compression_ratio(&self, m: usize, n: usize) -> f64 {
let k = self.sigma.len();
(m * n) as f64 / (m * k + k + k * n) as f64
}
}
pub fn truncated_svd<S: Storage, const F: u32>(a: &FixedMatrix<Fixed<S, F>>, k: usize) -> Result<TruncatedSVD<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let svd = svd_decompose(a)?;
let full_k = svd.sigma.len();
let k = k.min(full_k);
let m = svd.u.rows();
let mut u_k = FixedMatrix::<Fixed<S, F>>::new(m, k);
for i in 0..m {
for j in 0..k {
u_k.set(i, j, svd.u.get(i, j));
}
}
let mut sigma_k = FixedVector::<Fixed<S, F>>::new(k);
for i in 0..k {
sigma_k[i] = svd.sigma[i];
}
let n = svd.vt.cols();
let mut vt_k = FixedMatrix::<Fixed<S, F>>::new(k, n);
for i in 0..k {
for j in 0..n {
vt_k.set(i, j, svd.vt.get(i, j));
}
}
Ok(TruncatedSVD::<Fixed<S, F>> { u: u_k, sigma: sigma_k, vt: vt_k })
}
pub fn truncated_svd_auto<S: Storage, const F: u32>(a: &FixedMatrix<Fixed<S, F>>, threshold: Option<Fixed<S, F>>) -> Result<TruncatedSVD<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let svd = svd_decompose(a)?;
let thresh = threshold.unwrap_or_else(|| {
if svd.sigma.len() == 0 { return Fixed::one(); }
let sigma_max = svd.sigma[0];
let eps = kernels::convergence_threshold(sigma_max);
eps.mul_count(a.rows().max(a.cols()))
});
let mut k = 0;
for i in 0..svd.sigma.len() {
if svd.sigma[i] > thresh { k += 1; } else { break; }
}
if k == 0 { k = 1; }
truncated_svd(a, k)
}
pub struct TuckerDecomposition<T> {
pub core: Tensor<T>,
pub factors: Vec<FixedMatrix<T>>,
}
impl<S: Storage, const F: u32> TuckerDecomposition<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
pub fn reconstruct(&self) -> Tensor<Fixed<S, F>> {
let chain = || -> Result<Tensor<Fixed<S, F>>, OverflowDetected> {
let mut data = to_compute(&self.core);
let mut shape = self.core.shape().to_vec();
for (n, u) in self.factors.iter().enumerate() {
(data, shape) = mode_n_product_compute(&data, &shape, u, n)?;
}
to_storage_tensor(&data, &shape)
};
chain().expect("TuckerDecomposition::reconstruct: entry exceeds the compute tier or storage")
}
pub fn compression_ratio(&self, original_shape: &[usize]) -> f64 {
let orig: usize = original_shape.iter().product();
let core_size: usize = self.core.shape().iter().product();
let factor_size: usize = self.factors.iter().enumerate()
.map(|(n, f)| original_shape[n] * f.cols())
.sum();
orig as f64 / (core_size + factor_size) as f64
}
}
pub fn tucker_decompose<S: Storage, const F: u32>(t: &Tensor<Fixed<S, F>>, ranks: &[usize]) -> Result<TuckerDecomposition<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let ndim = t.rank();
assert_eq!(ranks.len(), ndim, "ranks must have one entry per tensor mode");
let mut factors: Vec<FixedMatrix<Fixed<S, F>>> = Vec::with_capacity(ndim);
for n in 0..ndim {
let unfolded = mode_unfold(t, n);
let k = ranks[n].min(unfolded.rows()).min(unfolded.cols());
let tsvd = truncated_svd(&unfolded, k)?;
factors.push(tsvd.u); }
let mut data = to_compute(t);
let mut shape = t.shape().to_vec();
for (n, u) in factors.iter().enumerate() {
(data, shape) = mode_n_product_compute(&data, &shape, &u.transpose(), n)?;
}
let core = to_storage_tensor(&data, &shape)?;
Ok(TuckerDecomposition::<Fixed<S, F>> { core, factors })
}
pub struct CPDecomposition<T> {
pub weights: FixedVector<T>,
pub factors: Vec<FixedMatrix<T>>,
}
impl<S: Storage, const F: u32> CPDecomposition<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
pub fn reconstruct(&self, shape: &[usize]) -> Tensor<Fixed<S, F>> {
let rank = self.weights.len();
let total: usize = shape.iter().product();
let mut data = vec![Fixed::<S, F>::make_compute_int(0); total];
for r in 0..rank {
let w = self.weights[r];
add_rank1_to_flat(&mut data, shape, &self.factors, r, w);
}
to_storage_tensor(&data, shape).expect("CPDecomposition::reconstruct: entry exceeds storage")
}
}
pub fn cp_decompose<S: Storage, const F: u32>(
t: &Tensor<Fixed<S, F>>,
rank: usize,
max_iter: usize,
_tol: Fixed<S, F>,
) -> Result<CPDecomposition<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let ndim = t.rank();
let shape = t.shape().to_vec();
let mut factors: Vec<ComputeMatrix<Fixed<S, F>>> = Vec::with_capacity(ndim);
for n in 0..ndim {
let unfolded = mode_unfold(t, n);
let k = rank.min(unfolded.rows()).min(unfolded.cols());
let svd = svd_decompose(&unfolded)?;
let mut f = ComputeMatrix::<Fixed<S, F>>::new(shape[n], rank);
for i in 0..shape[n] {
for r in 0..rank {
if r < k {
f.set(i, r, Fixed::<S, F>::upscale_to_compute(svd.u.get(i, r).raw()));
}
}
}
factors.push(f);
}
for _iter in 0..max_iter {
for n in 0..ndim {
let v = khatri_rao_except(&factors, n, &shape)?;
let unfolded = mode_unfold(t, n);
let cols: Vec<Vec<Compute<Fixed<S, F>>>> = (0..rank).map(|r| v.col_vec(r)).collect();
let mut vtv = ComputeMatrix::<Fixed<S, F>>::new(rank, rank);
for a in 0..rank {
for b in 0..rank {
vtv.set(a, b, kernels::narrow_product_to_compute::<Fixed<S, F>>(kernels::exact_dot_compute::<Fixed<S, F>>(&cols[a], &cols[b])?)?);
}
}
let lu = match compute_lu_decompose(&vtv) {
Ok(lu) => lu,
Err(_) => continue,
};
let mut next = ComputeMatrix::<Fixed<S, F>>::new(shape[n], rank);
let mut solved = true;
for i in 0..shape[n] {
let row: Vec<Compute<Fixed<S, F>>> = (0..unfolded.cols()).map(|c| Fixed::<S, F>::upscale_to_compute(unfolded.get(i, c).raw())).collect();
let mut rhs = Vec::with_capacity(rank);
for col in &cols {
rhs.push(kernels::narrow_product_to_compute::<Fixed<S, F>>(kernels::exact_dot_compute::<Fixed<S, F>>(&row, col)?)?);
}
match lu.solve(&rhs) {
Ok(x) => for r in 0..rank { next.set(i, r, x[r]); },
Err(_) => { solved = false; break; }
}
}
if solved {
factors[n] = next;
}
}
}
let mut weights = FixedVector::<Fixed<S, F>>::new(rank);
let mut out: Vec<FixedMatrix<Fixed<S, F>>> = (0..ndim).map(|n| FixedMatrix::<Fixed<S, F>>::new(shape[n], rank)).collect();
for r in 0..rank {
let mut norm_product = Fixed::<S, F>::make_compute_int(1);
for n in 0..ndim {
let col = factors[n].col_vec(r);
let col_norm = Fixed::<S, F>::sqrt_at_compute_tier(kernels::narrow_product_to_compute::<Fixed<S, F>>(kernels::exact_dot_compute::<Fixed<S, F>>(&col, &col)?)?);
if !Fixed::<S, F>::compute_is_zero(&col_norm) {
for i in 0..shape[n] {
out[n].set(i, r, Fixed::from_raw(kernels::divide_to_storage_nearest::<Fixed<S, F>>(lift_to_triple(col[i]), col_norm)?));
}
norm_product = kernels::compute_product::<Fixed<S, F>>(norm_product, col_norm)?;
} else {
for i in 0..shape[n] {
out[n].set(i, r, Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(col[i])?));
}
}
}
weights[r] = Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(norm_product)?);
}
Ok(CPDecomposition::<Fixed<S, F>> { weights, factors: out })
}
fn mode_unfold<S: Storage, const F: u32>(t: &Tensor<Fixed<S, F>>, mode: usize) -> FixedMatrix<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let shape = t.shape();
let ndim = shape.len();
let rows = shape[mode];
let cols: usize = shape.iter().enumerate()
.filter(|&(i, _)| i != mode)
.map(|(_, &d)| d)
.product();
let mut result = FixedMatrix::<Fixed<S, F>>::new(rows, cols);
let mut perm: Vec<usize> = vec![mode];
for i in 0..ndim {
if i != mode { perm.push(i); }
}
let mut indices = vec![0usize; ndim];
let total: usize = shape.iter().product();
for flat in 0..total {
let mut rem = flat;
for d in (0..ndim).rev() {
indices[d] = rem % shape[d];
rem /= shape[d];
}
let row = indices[mode];
let mut col = 0;
let mut stride = 1;
for &p in perm[1..].iter().rev() {
col += indices[p] * stride;
stride *= shape[p];
}
result.set(row, col, t.get(&indices));
}
result
}
#[inline]
fn lift_to_triple<S: Storage, const F: u32>(c: Compute<Fixed<S, F>>) -> Orient<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
S::widen_product(c, S::widen(Fixed::one().raw()))
}
fn to_compute<S: Storage, const F: u32>(t: &Tensor<Fixed<S, F>>) -> Vec<Compute<Fixed<S, F>>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
t.data().iter().map(|x| Fixed::<S, F>::upscale_to_compute(x.raw())).collect()
}
fn to_storage_tensor<S: Storage, const F: u32>(data: &[Compute<Fixed<S, F>>], shape: &[usize]) -> Result<Tensor<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let entries = data.iter().map(|c| Fixed::<S, F>::downscale_to_storage(*c).map(Fixed::from_raw)).collect::<Result<Vec<_>, _>>()?;
Ok(Tensor::<Fixed<S, F>>::from_data(shape, &entries))
}
fn mode_n_product_compute<S: Storage, const F: u32>(
data: &[Compute<Fixed<S, F>>],
shape: &[usize],
m: &FixedMatrix<Fixed<S, F>>,
mode: usize,
) -> Result<(Vec<Compute<Fixed<S, F>>>, Vec<usize>), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let ndim = shape.len();
let d_n = shape[mode];
let r = m.rows();
assert_eq!(m.cols(), d_n, "Matrix cols must match tensor mode dimension");
let mut strides = vec![1usize; ndim];
for d in (0..ndim.saturating_sub(1)).rev() {
strides[d] = strides[d + 1] * shape[d + 1];
}
let mut new_shape = shape.to_vec();
new_shape[mode] = r;
let m_rows: Vec<Vec<Compute<Fixed<S, F>>>> = (0..r)
.map(|i| (0..d_n).map(|k| Fixed::<S, F>::upscale_to_compute(m.get(i, k).raw())).collect())
.collect();
let total: usize = new_shape.iter().product();
let mut result = Vec::with_capacity(total);
let mut out_indices = vec![0usize; ndim];
for flat in 0..total {
let mut rem = flat;
for d in (0..ndim).rev() {
out_indices[d] = rem % new_shape[d];
rem /= new_shape[d];
}
let i = out_indices[mode];
let base: usize = (0..ndim).filter(|&d| d != mode).map(|d| out_indices[d] * strides[d]).sum();
let fiber: Vec<Compute<Fixed<S, F>>> = (0..d_n).map(|k| data[base + k * strides[mode]]).collect();
result.push(kernels::narrow_product_to_compute::<Fixed<S, F>>(kernels::exact_dot_compute::<Fixed<S, F>>(&m_rows[i], &fiber)?)?);
}
Ok((result, new_shape))
}
fn khatri_rao_except<S: Storage, const F: u32>(factors: &[ComputeMatrix<Fixed<S, F>>], skip: usize, shape: &[usize]) -> Result<ComputeMatrix<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let rank = factors[0].cols();
let ndim = factors.len();
let rows: usize = shape.iter().enumerate()
.filter(|&(i, _)| i != skip)
.map(|(_, &d)| d)
.product();
let mut result = ComputeMatrix::<Fixed<S, F>>::new(rows, rank);
for r in 0..rank {
let active_modes: Vec<usize> = (0..ndim).filter(|&i| i != skip).collect();
for row in 0..rows {
let mut rem = row;
let mut val = Fixed::<S, F>::make_compute_int(1);
for &m in active_modes.iter().rev() {
let idx = rem % shape[m];
rem /= shape[m];
val = kernels::compute_product::<Fixed<S, F>>(val, factors[m].get(idx, r))?;
}
result.set(row, r, val);
}
}
Ok(result)
}
fn add_rank1_to_flat<S: Storage, const F: u32>(
data: &mut [Compute<Fixed<S, F>>],
shape: &[usize],
factors: &[FixedMatrix<Fixed<S, F>>],
r: usize,
weight: Fixed<S, F>,
)
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let ndim = shape.len();
let total = data.len();
let mut indices = vec![0usize; ndim];
for flat in 0..total {
let mut rem = flat;
for d in (0..ndim).rev() {
indices[d] = rem % shape[d];
rem /= shape[d];
}
let mut val = Fixed::<S, F>::upscale_to_compute(weight.raw());
for n in 0..ndim {
val = kernels::compute_product::<Fixed<S, F>>(val, Fixed::<S, F>::upscale_to_compute(factors[n].get(indices[n], r).raw()))
.expect("CPDecomposition::reconstruct: term exceeds the compute tier");
}
data[flat] = Fixed::<S, F>::compute_add(data[flat], val);
}
}