2527 lines
88 KiB
Rust
2527 lines
88 KiB
Rust
//! Advanced linear algebra operations for RTX Tensor
|
||
//!
|
||
//! This module provides numerically stable implementations of linear algebra decompositions:
|
||
//!
|
||
//! # Decompositions
|
||
//! - **SVD** (Singular Value Decomposition): A = U*S*V^T
|
||
//! - **QR** decomposition using Householder reflections: A = Q*R
|
||
//! - **Cholesky** decomposition for positive definite matrices: A = L*L^T
|
||
//! - **Eigendecomposition** for symmetric matrices: A = V*Λ*V^T
|
||
//! - **Pseudo-inverse** using SVD with numerical tolerance
|
||
//! - **Truncated SVD** for dimensionality reduction
|
||
//!
|
||
//! # Features
|
||
//! - Numerical stability with proper conditioning
|
||
//! - Support for f32 and f64 precision
|
||
//! - Error handling for ill-conditioned matrices
|
||
//! - Integration with autograd system
|
||
//! - Memory-efficient implementations
|
||
//!
|
||
//! # Example
|
||
//! ```rust
|
||
//! use rtx_tensor::{Tensor, Device};
|
||
//!
|
||
//! let device = Device::cpu();
|
||
//! let a = Tensor::from_data(vec![3.0, 1.0, 1.0, 3.0], vec![2, 2], &device).unwrap();
|
||
//!
|
||
//! // SVD decomposition
|
||
//! let svd = a.svd(true, true).unwrap();
|
||
//! println!("Singular values: {:?}", svd.s.to_cpu().unwrap());
|
||
//!
|
||
//! // QR decomposition
|
||
//! let qr = a.qr().unwrap();
|
||
//!
|
||
//! // Eigendecomposition (symmetric matrices only)
|
||
//! let eigen = a.symeig(true).unwrap();
|
||
//! ```
|
||
|
||
use crate::{Tensor, TensorError, Result};
|
||
use std::fmt;
|
||
|
||
#[cfg(test)]
|
||
mod tests;
|
||
|
||
#[cfg(test)]
|
||
mod simple_test;
|
||
|
||
#[cfg(test)]
|
||
mod matrix_power_tests;
|
||
|
||
// Re-export test module types for external testing (when testing)
|
||
#[cfg(test)]
|
||
pub use tests::*;
|
||
|
||
// Public types are defined below and available for import
|
||
|
||
/// Result of Singular Value Decomposition
|
||
#[derive(Debug, Clone)]
|
||
pub struct SVDResult {
|
||
/// Left singular vectors (U matrix)
|
||
pub u: Tensor,
|
||
/// Singular values (diagonal of S matrix)
|
||
pub s: Tensor,
|
||
/// Right singular vectors transposed (V^T matrix)
|
||
pub vt: Tensor,
|
||
}
|
||
|
||
impl fmt::Display for SVDResult {
|
||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
write!(
|
||
f,
|
||
"SVDResult {{ U: {}, S: {}, V^T: {} }}",
|
||
self.u.shape(),
|
||
self.s.shape(),
|
||
self.vt.shape()
|
||
)
|
||
}
|
||
}
|
||
|
||
/// Result of QR Decomposition
|
||
#[derive(Debug, Clone)]
|
||
pub struct QRResult {
|
||
/// Orthogonal matrix Q
|
||
pub q: Tensor,
|
||
/// Upper triangular matrix R
|
||
pub r: Tensor,
|
||
}
|
||
|
||
impl fmt::Display for QRResult {
|
||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
write!(
|
||
f,
|
||
"QRResult {{ Q: {}, R: {} }}",
|
||
self.q.shape(),
|
||
self.r.shape()
|
||
)
|
||
}
|
||
}
|
||
|
||
/// Result of Eigenvalue Decomposition
|
||
#[derive(Debug, Clone)]
|
||
pub struct EigenResult {
|
||
/// Eigenvalues (Λ)
|
||
pub eigenvalues: Tensor,
|
||
/// Eigenvectors (V)
|
||
pub eigenvectors: Tensor,
|
||
}
|
||
|
||
/// Result of LU Decomposition
|
||
#[derive(Debug, Clone)]
|
||
pub struct LUResult {
|
||
/// Lower triangular matrix L (with 1s on diagonal in Doolittle algorithm)
|
||
pub l: Tensor,
|
||
/// Upper triangular matrix U
|
||
pub u: Tensor,
|
||
/// Permutation vector P (indices of row exchanges)
|
||
pub p: Tensor,
|
||
/// Permutation matrix P (optional, returned when requested)
|
||
pub p_matrix: Option<Tensor>,
|
||
/// Determinant of the original matrix (computed from LU)
|
||
pub determinant: f32,
|
||
/// Whether the matrix is invertible (non-singular)
|
||
pub is_invertible: bool,
|
||
/// Condition number estimate
|
||
pub condition_number: f32,
|
||
/// Number of row exchanges performed during pivoting
|
||
pub num_pivots: usize,
|
||
}
|
||
|
||
/// Options for LU decomposition
|
||
#[derive(Debug, Clone)]
|
||
pub struct LUOptions {
|
||
/// Whether to use partial pivoting for numerical stability
|
||
pub pivoting: bool,
|
||
/// Whether to perform decomposition in-place (memory efficient)
|
||
pub in_place: bool,
|
||
/// Whether to return permutation matrix P in addition to permutation vector
|
||
pub return_permutation_matrix: bool,
|
||
}
|
||
|
||
impl fmt::Display for EigenResult {
|
||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
write!(
|
||
f,
|
||
"EigenResult {{ eigenvalues: {}, eigenvectors: {} }}",
|
||
self.eigenvalues.shape(),
|
||
self.eigenvectors.shape()
|
||
)
|
||
}
|
||
}
|
||
|
||
impl fmt::Display for LUResult {
|
||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
write!(
|
||
f,
|
||
"LUResult {{ L: {}, U: {}, det: {:.6}, invertible: {}, pivots: {} }}",
|
||
self.l.shape(),
|
||
self.u.shape(),
|
||
self.determinant,
|
||
self.is_invertible,
|
||
self.num_pivots
|
||
)
|
||
}
|
||
}
|
||
|
||
impl Default for LUOptions {
|
||
fn default() -> Self {
|
||
Self {
|
||
pivoting: true,
|
||
in_place: false,
|
||
return_permutation_matrix: false,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Matrix structure type for triangular solving
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum TriangularType {
|
||
/// Lower triangular matrix (non-zero entries below or on diagonal)
|
||
Lower,
|
||
/// Upper triangular matrix (non-zero entries above or on diagonal)
|
||
Upper,
|
||
}
|
||
|
||
/// Options for triangular system solving
|
||
#[derive(Debug, Clone)]
|
||
pub struct TriangularSolveOptions {
|
||
/// Whether the matrix is lower or upper triangular
|
||
pub triangular_type: TriangularType,
|
||
/// Whether to solve A^T x = b instead of A x = b
|
||
pub transpose: bool,
|
||
/// Whether the diagonal elements are assumed to be 1 (unit diagonal)
|
||
pub unit_diagonal: bool,
|
||
/// Whether to validate that the matrix is actually triangular
|
||
pub validate_triangular: bool,
|
||
/// Whether to use strict triangular validation (exact zeros) or relaxed (within tolerance)
|
||
pub strict_validation: bool,
|
||
}
|
||
|
||
/// Method to use for least squares solving
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum LeastSquaresMethod {
|
||
/// Automatic method selection based on matrix properties
|
||
Auto,
|
||
/// QR decomposition method (best for overdetermined systems)
|
||
QR,
|
||
/// SVD method (most robust, handles rank-deficient matrices)
|
||
SVD,
|
||
/// Normal equations (A^T A)x = A^T b (fastest but least stable)
|
||
NormalEquations,
|
||
}
|
||
|
||
impl Default for LeastSquaresMethod {
|
||
fn default() -> Self {
|
||
Self::Auto
|
||
}
|
||
}
|
||
|
||
/// Options for least squares solving
|
||
#[derive(Debug, Clone)]
|
||
pub struct LeastSquaresOptions {
|
||
/// Method to use for solving
|
||
pub method: LeastSquaresMethod,
|
||
/// Tikhonov regularization parameter (λ in (A^T A + λI)x = A^T b)
|
||
pub regularization_lambda: Option<f32>,
|
||
/// Diagonal weight matrix for weighted least squares
|
||
pub weights: Option<Tensor>,
|
||
/// Relative condition number threshold for rank determination
|
||
pub rcond: f32,
|
||
/// Maximum number of iterations for iterative methods
|
||
pub max_iterations: Option<usize>,
|
||
/// Convergence tolerance for iterative methods
|
||
pub tolerance: Option<f32>,
|
||
}
|
||
|
||
impl Default for LeastSquaresOptions {
|
||
fn default() -> Self {
|
||
Self {
|
||
method: LeastSquaresMethod::Auto,
|
||
regularization_lambda: None,
|
||
weights: None,
|
||
rcond: 1e-12,
|
||
max_iterations: None,
|
||
tolerance: None,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Result of least squares solving
|
||
#[derive(Debug, Clone)]
|
||
pub struct LeastSquaresResult {
|
||
/// Solution vector(s) x
|
||
pub solution: Tensor,
|
||
/// Residual norm ||Ax - b||
|
||
pub residual_norm: f32,
|
||
/// Condition number of the coefficient matrix
|
||
pub condition_number: f32,
|
||
/// Effective rank of the coefficient matrix
|
||
pub rank: usize,
|
||
/// Method actually used for solving
|
||
pub method_used: LeastSquaresMethod,
|
||
/// R-squared goodness-of-fit statistic (if applicable)
|
||
pub r_squared: Option<f32>,
|
||
/// Sum of squared residuals
|
||
pub sum_squared_residuals: f32,
|
||
}
|
||
|
||
impl fmt::Display for LeastSquaresResult {
|
||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
write!(
|
||
f,
|
||
"LeastSquaresResult {{ solution: {}, residual_norm: {:.6}, condition_number: {:.6}, rank: {}, method: {:?} }}",
|
||
self.solution.shape(),
|
||
self.residual_norm,
|
||
self.condition_number,
|
||
self.rank,
|
||
self.method_used
|
||
)
|
||
}
|
||
}
|
||
|
||
/// Linear algebra operations implementation
|
||
impl Tensor {
|
||
// Note: Helper functions unravel_index, ravel_index, and broadcast_index are defined in tensor::utilities
|
||
/// Singular Value Decomposition (SVD)
|
||
///
|
||
/// Decomposes a matrix A into A = U * S * V^T where:
|
||
/// - U contains the left singular vectors
|
||
/// - S contains the singular values (non-negative, decreasing)
|
||
/// - V^T contains the right singular vectors transposed
|
||
///
|
||
/// # Arguments
|
||
/// - `compute_u`: Whether to compute the U matrix (left singular vectors)
|
||
/// - `compute_vt`: Whether to compute the V^T matrix (right singular vectors)
|
||
///
|
||
/// # Returns
|
||
/// `SVDResult` containing U, S, and V^T tensors
|
||
///
|
||
/// # Numerical Stability
|
||
/// Uses iterative power method with Gram-Schmidt orthogonalization.
|
||
/// Handles rank-deficient and ill-conditioned matrices gracefully.
|
||
pub fn svd(&self, compute_u: bool, compute_vt: bool) -> Result<SVDResult> {
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("SVD requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
let m = dims[0]; // rows
|
||
let n = dims[1]; // cols
|
||
let min_dim = m.min(n);
|
||
|
||
if m == 0 || n == 0 {
|
||
return Err(TensorError::shape("SVD requires non-empty matrix"));
|
||
}
|
||
|
||
let data = self.to_cpu()?;
|
||
|
||
// Use power iteration method for SVD
|
||
// This is a simplified implementation - production would use more sophisticated algorithms
|
||
let mut a = data.clone();
|
||
let mut u_data = vec![0.0f32; if compute_u { m * min_dim } else { 0 }];
|
||
let mut s_data = vec![0.0f32; min_dim];
|
||
let mut vt_data = vec![0.0f32; if compute_vt { min_dim * n } else { 0 }];
|
||
|
||
// For each singular value/vector pair
|
||
for k in 0..min_dim {
|
||
// Power iteration to find dominant singular value and vectors
|
||
let mut v = vec![1.0f32; n];
|
||
normalize_vector(&mut v);
|
||
|
||
const MAX_ITER: usize = 100;
|
||
const TOLERANCE: f32 = 1e-6;
|
||
|
||
for _iter in 0..MAX_ITER {
|
||
// v = A^T * (A * v)
|
||
let mut av = vec![0.0f32; m];
|
||
matrix_vector_multiply(&a, &v, &mut av, m, n);
|
||
|
||
let mut atav = vec![0.0f32; n];
|
||
matrix_transpose_vector_multiply(&a, &av, &mut atav, m, n);
|
||
|
||
let old_v = v.clone();
|
||
v = atav;
|
||
normalize_vector(&mut v);
|
||
|
||
// Check convergence
|
||
let mut diff = 0.0f32;
|
||
for i in 0..n {
|
||
diff += (v[i] - old_v[i]).abs();
|
||
}
|
||
if diff < TOLERANCE {
|
||
break;
|
||
}
|
||
}
|
||
|
||
// Compute Av to get left singular vector
|
||
let mut av = vec![0.0f32; m];
|
||
matrix_vector_multiply(&a, &v, &mut av, m, n);
|
||
let sigma = vector_norm(&av);
|
||
s_data[k] = sigma;
|
||
|
||
if sigma > 1e-12 {
|
||
// Normalize to get left singular vector
|
||
for i in 0..m {
|
||
av[i] /= sigma;
|
||
}
|
||
|
||
// Store singular vectors
|
||
if compute_u {
|
||
for i in 0..m {
|
||
u_data[i * min_dim + k] = av[i];
|
||
}
|
||
}
|
||
|
||
if compute_vt {
|
||
for j in 0..n {
|
||
vt_data[k * n + j] = v[j];
|
||
}
|
||
}
|
||
|
||
// Deflate matrix: A = A - sigma * u * v^T
|
||
for i in 0..m {
|
||
for j in 0..n {
|
||
a[i * n + j] -= sigma * av[i] * v[j];
|
||
}
|
||
}
|
||
} else {
|
||
// Zero singular value - fill with zeros
|
||
if compute_u {
|
||
for i in 0..m {
|
||
u_data[i * min_dim + k] = 0.0;
|
||
}
|
||
}
|
||
|
||
if compute_vt {
|
||
for j in 0..n {
|
||
vt_data[k * n + j] = 0.0;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Create result tensors
|
||
let u = if compute_u {
|
||
if min_dim < m {
|
||
// Thin SVD - U is m x min_dim
|
||
Tensor::from_data(u_data, vec![m, min_dim], &self.device)?
|
||
} else {
|
||
// Complete orthogonal basis for U (pad with random orthogonal vectors)
|
||
let mut full_u_data = vec![0.0f32; m * m];
|
||
|
||
// Copy existing U data
|
||
for i in 0..m {
|
||
for j in 0..min_dim {
|
||
full_u_data[i * m + j] = u_data[i * min_dim + j];
|
||
}
|
||
}
|
||
|
||
// Fill remaining columns with orthogonal vectors using Gram-Schmidt
|
||
for col in min_dim..m {
|
||
let mut new_col = vec![0.0f32; m];
|
||
if col < m {
|
||
new_col[col] = 1.0; // Start with standard basis vector
|
||
}
|
||
|
||
// Gram-Schmidt orthogonalization against existing columns
|
||
for prev_col in 0..col {
|
||
let mut dot = 0.0f32;
|
||
for i in 0..m {
|
||
dot += new_col[i] * full_u_data[i * m + prev_col];
|
||
}
|
||
for i in 0..m {
|
||
new_col[i] -= dot * full_u_data[i * m + prev_col];
|
||
}
|
||
}
|
||
|
||
// Normalize
|
||
let norm = (new_col.iter().map(|x| x * x).sum::<f32>()).sqrt();
|
||
if norm > 1e-12 {
|
||
for i in 0..m {
|
||
new_col[i] /= norm;
|
||
full_u_data[i * m + col] = new_col[i];
|
||
}
|
||
}
|
||
}
|
||
|
||
Tensor::from_data(full_u_data, vec![m, m], &self.device)?
|
||
}
|
||
} else {
|
||
Tensor::zeros([m, min_dim], &self.device)?
|
||
};
|
||
|
||
let s = Tensor::from_data(s_data, vec![min_dim], &self.device)?;
|
||
|
||
let vt = if compute_vt {
|
||
if min_dim < n {
|
||
// Thin SVD - V^T is min_dim x n
|
||
Tensor::from_data(vt_data, vec![min_dim, n], &self.device)?
|
||
} else {
|
||
// Complete orthogonal basis for V^T
|
||
let mut full_vt_data = vec![0.0f32; n * n];
|
||
|
||
// Copy existing V^T data
|
||
for i in 0..min_dim {
|
||
for j in 0..n {
|
||
full_vt_data[i * n + j] = vt_data[i * n + j];
|
||
}
|
||
}
|
||
|
||
// Fill remaining rows with orthogonal vectors
|
||
for row in min_dim..n {
|
||
let mut new_row = vec![0.0f32; n];
|
||
if row < n {
|
||
new_row[row] = 1.0;
|
||
}
|
||
|
||
// Gram-Schmidt orthogonalization
|
||
for prev_row in 0..row {
|
||
let mut dot = 0.0f32;
|
||
for j in 0..n {
|
||
dot += new_row[j] * full_vt_data[prev_row * n + j];
|
||
}
|
||
for j in 0..n {
|
||
new_row[j] -= dot * full_vt_data[prev_row * n + j];
|
||
}
|
||
}
|
||
|
||
// Normalize
|
||
let norm = (new_row.iter().map(|x| x * x).sum::<f32>()).sqrt();
|
||
if norm > 1e-12 {
|
||
for j in 0..n {
|
||
new_row[j] /= norm;
|
||
full_vt_data[row * n + j] = new_row[j];
|
||
}
|
||
}
|
||
}
|
||
|
||
Tensor::from_data(full_vt_data, vec![n, n], &self.device)?
|
||
}
|
||
} else {
|
||
Tensor::zeros([min_dim, n], &self.device)?
|
||
};
|
||
|
||
Ok(SVDResult { u, s, vt })
|
||
}
|
||
|
||
/// Truncated SVD for dimensionality reduction
|
||
///
|
||
/// Computes only the k largest singular values and corresponding vectors.
|
||
/// More efficient than full SVD when only the dominant components are needed.
|
||
///
|
||
/// # Arguments
|
||
/// - `k`: Number of singular values/vectors to compute
|
||
///
|
||
/// # Returns
|
||
/// `SVDResult` with dimensions reduced to rank k
|
||
pub fn svd_truncated(&self, k: usize) -> Result<SVDResult> {
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("Truncated SVD requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
let m = dims[0];
|
||
let n = dims[1];
|
||
let min_dim = m.min(n);
|
||
|
||
if k > min_dim {
|
||
return Err(TensorError::shape(format!(
|
||
"k ({}) cannot exceed min(m, n) ({})", k, min_dim
|
||
)));
|
||
}
|
||
|
||
let data = self.to_cpu()?;
|
||
|
||
// Use power iteration to find top k singular values/vectors
|
||
let mut a = data.clone();
|
||
let mut u_data = vec![0.0f32; m * k];
|
||
let mut s_data = vec![0.0f32; k];
|
||
let mut vt_data = vec![0.0f32; k * n];
|
||
|
||
for i in 0..k {
|
||
// Power iteration for i-th singular value
|
||
let mut v = vec![1.0f32; n];
|
||
normalize_vector(&mut v);
|
||
|
||
for _iter in 0..50 {
|
||
let mut av = vec![0.0f32; m];
|
||
matrix_vector_multiply(&a, &v, &mut av, m, n);
|
||
|
||
let mut atav = vec![0.0f32; n];
|
||
matrix_transpose_vector_multiply(&a, &av, &mut atav, m, n);
|
||
|
||
v = atav;
|
||
normalize_vector(&mut v);
|
||
}
|
||
|
||
// Compute singular value and vectors
|
||
let mut av = vec![0.0f32; m];
|
||
matrix_vector_multiply(&a, &v, &mut av, m, n);
|
||
let sigma = vector_norm(&av);
|
||
s_data[i] = sigma;
|
||
|
||
if sigma > 1e-12 {
|
||
// Normalize left singular vector
|
||
for j in 0..m {
|
||
av[j] /= sigma;
|
||
u_data[j * k + i] = av[j];
|
||
}
|
||
|
||
// Store right singular vector
|
||
for j in 0..n {
|
||
vt_data[i * n + j] = v[j];
|
||
}
|
||
|
||
// Deflate matrix
|
||
for row in 0..m {
|
||
for col in 0..n {
|
||
a[row * n + col] -= sigma * av[row] * v[col];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
let u = Tensor::from_data(u_data, vec![m, k], &self.device)?;
|
||
let s = Tensor::from_data(s_data, vec![k], &self.device)?;
|
||
let vt = Tensor::from_data(vt_data, vec![k, n], &self.device)?;
|
||
|
||
Ok(SVDResult { u, s, vt })
|
||
}
|
||
|
||
/// QR Decomposition using Householder reflections
|
||
///
|
||
/// Decomposes matrix A into A = Q * R where:
|
||
/// - Q is an orthogonal matrix
|
||
/// - R is an upper triangular matrix
|
||
///
|
||
/// # Returns
|
||
/// `QRResult` containing Q and R tensors
|
||
///
|
||
/// # Numerical Stability
|
||
/// Uses Householder reflections for better numerical stability than Gram-Schmidt.
|
||
pub fn qr(&self) -> Result<QRResult> {
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("QR requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
let m = dims[0]; // rows
|
||
let n = dims[1]; // cols
|
||
|
||
if m == 0 || n == 0 {
|
||
return Err(TensorError::shape("QR requires non-empty matrix"));
|
||
}
|
||
|
||
let data = self.to_cpu()?;
|
||
let mut r = data.clone(); // Will become R matrix
|
||
let mut q = vec![0.0f32; m * m]; // Full Q matrix
|
||
|
||
// Initialize Q as identity matrix
|
||
for i in 0..m {
|
||
q[i * m + i] = 1.0;
|
||
}
|
||
|
||
let k = n.min(m); // Number of Householder reflections
|
||
|
||
// Apply Householder reflections
|
||
for col in 0..k {
|
||
// Extract column vector starting from diagonal
|
||
let mut x = vec![0.0f32; m - col];
|
||
for i in col..m {
|
||
x[i - col] = r[i * n + col];
|
||
}
|
||
|
||
let norm_x = vector_norm(&x);
|
||
if norm_x < 1e-12 {
|
||
continue; // Skip near-zero columns
|
||
}
|
||
|
||
// Create Householder vector
|
||
let mut v = x.clone();
|
||
let sign = if v[0] >= 0.0 { 1.0 } else { -1.0 };
|
||
v[0] += sign * norm_x;
|
||
|
||
let norm_v = vector_norm(&v);
|
||
if norm_v < 1e-12 {
|
||
continue;
|
||
}
|
||
|
||
// Normalize Householder vector
|
||
for i in 0..v.len() {
|
||
v[i] /= norm_v;
|
||
}
|
||
|
||
let tau = 2.0;
|
||
|
||
// Apply Householder reflection to R: R = (I - tau * v * v^T) * R
|
||
// Only update the submatrix starting from (col, col)
|
||
for j in col..n {
|
||
let mut dot = 0.0f32;
|
||
for i in 0..(m - col) {
|
||
dot += v[i] * r[(col + i) * n + j];
|
||
}
|
||
|
||
for i in 0..(m - col) {
|
||
r[(col + i) * n + j] -= tau * dot * v[i];
|
||
}
|
||
}
|
||
|
||
// Apply Householder reflection to Q: Q = Q * (I - tau * v * v^T)
|
||
// This builds up the Q matrix from the right
|
||
for i in 0..m {
|
||
let mut dot = 0.0f32;
|
||
for j in 0..(m - col) {
|
||
dot += q[i * m + (col + j)] * v[j];
|
||
}
|
||
|
||
for j in 0..(m - col) {
|
||
q[i * m + (col + j)] -= tau * dot * v[j];
|
||
}
|
||
}
|
||
}
|
||
|
||
// For rectangular matrices, we want reduced QR
|
||
let q_result = if m > n {
|
||
// Extract first n columns of Q
|
||
let mut q_reduced = vec![0.0f32; m * n];
|
||
for i in 0..m {
|
||
for j in 0..n {
|
||
q_reduced[i * n + j] = q[i * m + j];
|
||
}
|
||
}
|
||
Tensor::from_data(q_reduced, vec![m, n], &self.device)?
|
||
} else {
|
||
Tensor::from_data(q, vec![m, m], &self.device)?
|
||
};
|
||
|
||
let r_result = if m > n {
|
||
// Extract first n rows of R
|
||
let mut r_reduced = vec![0.0f32; n * n];
|
||
for i in 0..n {
|
||
for j in 0..n {
|
||
r_reduced[i * n + j] = r[i * n + j];
|
||
}
|
||
}
|
||
Tensor::from_data(r_reduced, vec![n, n], &self.device)?
|
||
} else {
|
||
Tensor::from_data(r, vec![m, n], &self.device)?
|
||
};
|
||
|
||
Ok(QRResult { q: q_result, r: r_result })
|
||
}
|
||
|
||
/// Cholesky Decomposition
|
||
///
|
||
/// Decomposes a positive definite matrix A into A = L * L^T where L is lower triangular.
|
||
///
|
||
/// # Arguments
|
||
/// - `upper`: If true, return upper triangular U such that A = U^T * U
|
||
///
|
||
/// # Returns
|
||
/// Lower triangular matrix L (or upper triangular U if upper=true)
|
||
///
|
||
/// # Errors
|
||
/// Returns error if matrix is not positive definite
|
||
pub fn cholesky(&self, upper: bool) -> Result<Tensor> {
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("Cholesky requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
if dims[0] != dims[1] {
|
||
return Err(TensorError::shape("Cholesky requires square matrix"));
|
||
}
|
||
|
||
let n = dims[0];
|
||
let data = self.to_cpu()?;
|
||
let mut l = vec![0.0f32; n * n];
|
||
|
||
// Cholesky-Banachiewicz algorithm
|
||
for i in 0..n {
|
||
for j in 0..=i {
|
||
if i == j {
|
||
// Diagonal elements: L[i,i] = sqrt(A[i,i] - sum(L[i,k]^2 for k < i))
|
||
let mut sum = 0.0f32;
|
||
for k in 0..i {
|
||
sum += l[i * n + k] * l[i * n + k];
|
||
}
|
||
|
||
let diagonal_val = data[i * n + i] - sum;
|
||
if diagonal_val <= 0.0 {
|
||
return Err(TensorError::numerical(
|
||
"Matrix is not positive definite - negative diagonal element in Cholesky"
|
||
));
|
||
}
|
||
|
||
l[i * n + i] = diagonal_val.sqrt();
|
||
} else {
|
||
// Off-diagonal elements: L[i,j] = (A[i,j] - sum(L[i,k]*L[j,k] for k < j)) / L[j,j]
|
||
let mut sum = 0.0f32;
|
||
for k in 0..j {
|
||
sum += l[i * n + k] * l[j * n + k];
|
||
}
|
||
|
||
if l[j * n + j].abs() < 1e-12 {
|
||
return Err(TensorError::numerical(
|
||
"Matrix is not positive definite - zero diagonal element in Cholesky"
|
||
));
|
||
}
|
||
|
||
l[i * n + j] = (data[i * n + j] - sum) / l[j * n + j];
|
||
}
|
||
}
|
||
}
|
||
|
||
if upper {
|
||
// Return upper triangular matrix U = L^T
|
||
let mut u = vec![0.0f32; n * n];
|
||
for i in 0..n {
|
||
for j in i..n {
|
||
u[i * n + j] = l[j * n + i];
|
||
}
|
||
}
|
||
Tensor::from_data(u, vec![n, n], &self.device)
|
||
} else {
|
||
Tensor::from_data(l, vec![n, n], &self.device)
|
||
}
|
||
}
|
||
|
||
/// Symmetric Eigenvalue Decomposition
|
||
///
|
||
/// Computes eigenvalues and eigenvectors for symmetric matrices.
|
||
/// For symmetric matrix A: A = V * Λ * V^T
|
||
///
|
||
/// # Arguments
|
||
/// - `eigenvectors`: Whether to compute eigenvectors (true) or just eigenvalues (false)
|
||
///
|
||
/// # Returns
|
||
/// `EigenResult` containing eigenvalues and optionally eigenvectors
|
||
///
|
||
/// # Note
|
||
/// Only works for symmetric matrices. Eigenvalues are returned in ascending order.
|
||
pub fn symeig(&self, eigenvectors: bool) -> Result<EigenResult> {
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("Eigendecomposition requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
if dims[0] != dims[1] {
|
||
return Err(TensorError::shape("Eigendecomposition requires square matrix"));
|
||
}
|
||
|
||
let n = dims[0];
|
||
let data = self.to_cpu()?;
|
||
|
||
// Check symmetry
|
||
for i in 0..n {
|
||
for j in 0..n {
|
||
if (data[i * n + j] - data[j * n + i]).abs() > 1e-8 {
|
||
return Err(TensorError::shape("symeig requires symmetric matrix"));
|
||
}
|
||
}
|
||
}
|
||
|
||
// Use power iteration method for eigenvalues/eigenvectors
|
||
// This is simplified - production code would use QR algorithm or Jacobi method
|
||
let mut eigenvalues = vec![0.0f32; n];
|
||
let mut eigenvectors_data = if eigenvectors {
|
||
vec![0.0f32; n * n]
|
||
} else {
|
||
vec![]
|
||
};
|
||
|
||
let mut a = data.clone();
|
||
|
||
// Find eigenvalues using deflation
|
||
for k in 0..n {
|
||
// Power iteration for dominant eigenvalue
|
||
let mut v = vec![1.0f32; n - k];
|
||
for i in 1..(n - k) {
|
||
v[i] = i as f32; // Different starting vector for each eigenvalue
|
||
}
|
||
normalize_vector(&mut v);
|
||
|
||
let mut lambda = 0.0f32;
|
||
|
||
for _iter in 0..100 {
|
||
let mut av = vec![0.0f32; n - k];
|
||
|
||
// Matrix-vector multiply with current submatrix
|
||
for i in 0..(n - k) {
|
||
for j in 0..(n - k) {
|
||
av[i] += a[(k + i) * n + (k + j)] * v[j];
|
||
}
|
||
}
|
||
|
||
// Rayleigh quotient
|
||
let mut vav = 0.0f32;
|
||
let mut vv = 0.0f32;
|
||
for i in 0..(n - k) {
|
||
vav += v[i] * av[i];
|
||
vv += v[i] * v[i];
|
||
}
|
||
|
||
let new_lambda = if vv > 1e-12 { vav / vv } else { 0.0 };
|
||
|
||
// Normalize av to get new v
|
||
normalize_vector(&mut av);
|
||
|
||
// Check convergence
|
||
if (new_lambda - lambda).abs() < 1e-8 {
|
||
lambda = new_lambda;
|
||
v = av;
|
||
break;
|
||
}
|
||
|
||
lambda = new_lambda;
|
||
v = av;
|
||
}
|
||
|
||
eigenvalues[k] = lambda;
|
||
|
||
if eigenvectors {
|
||
// Store eigenvector (pad with zeros for deflated dimensions)
|
||
for i in 0..n {
|
||
for j in 0..n {
|
||
if i >= k && j == k && i - k < v.len() {
|
||
eigenvectors_data[i * n + j] = v[i - k];
|
||
} else if i < k && j == k {
|
||
eigenvectors_data[i * n + j] = 0.0;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Deflate matrix: A = A - lambda * v * v^T
|
||
if k < n - 1 {
|
||
for i in 0..(n - k) {
|
||
for j in 0..(n - k) {
|
||
a[(k + i) * n + (k + j)] -= lambda * v[i] * v[j];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Sort eigenvalues in ascending order (and reorder eigenvectors accordingly)
|
||
let mut indices: Vec<usize> = (0..n).collect();
|
||
indices.sort_by(|&a, &b| eigenvalues[a].partial_cmp(&eigenvalues[b]).unwrap());
|
||
|
||
let mut sorted_eigenvalues = vec![0.0f32; n];
|
||
for (i, &idx) in indices.iter().enumerate() {
|
||
sorted_eigenvalues[i] = eigenvalues[idx];
|
||
}
|
||
|
||
let sorted_eigenvectors = if eigenvectors {
|
||
let mut sorted = vec![0.0f32; n * n];
|
||
for (new_col, &old_col) in indices.iter().enumerate() {
|
||
for row in 0..n {
|
||
sorted[row * n + new_col] = eigenvectors_data[row * n + old_col];
|
||
}
|
||
}
|
||
Tensor::from_data(sorted, vec![n, n], &self.device)?
|
||
} else {
|
||
Tensor::zeros([n, n], &self.device)?
|
||
};
|
||
|
||
Ok(EigenResult {
|
||
eigenvalues: Tensor::from_data(sorted_eigenvalues, vec![n], &self.device)?,
|
||
eigenvectors: sorted_eigenvectors,
|
||
})
|
||
}
|
||
|
||
/// Moore-Penrose Pseudo-inverse using SVD
|
||
///
|
||
/// Computes the pseudo-inverse A^+ such that:
|
||
/// - If A has full row rank: A^+ = A^T(AA^T)^(-1)
|
||
/// - If A has full column rank: A^+ = (A^TA)^(-1)A^T
|
||
/// - General case: A^+ = V * S^+ * U^T (via SVD)
|
||
///
|
||
/// # Arguments
|
||
/// - `rcond`: Relative condition number. Singular values below rcond*max(S) are set to zero
|
||
///
|
||
/// # Returns
|
||
/// Pseudo-inverse tensor with dimensions transposed relative to input
|
||
pub fn pinv(&self, rcond: f32) -> Result<Tensor> {
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("Pseudo-inverse requires 2D tensor"));
|
||
}
|
||
|
||
// Use SVD to compute pseudo-inverse
|
||
let svd_result = self.svd(true, true)?;
|
||
let s_data = svd_result.s.to_cpu()?;
|
||
|
||
// Find cutoff threshold
|
||
let s_max = s_data.iter().fold(0.0f32, |a, &b| a.max(b));
|
||
let cutoff = rcond * s_max;
|
||
|
||
// Compute reciprocal of singular values (with threshold)
|
||
let mut s_inv_data = vec![0.0f32; s_data.len()];
|
||
for (i, &s_val) in s_data.iter().enumerate() {
|
||
s_inv_data[i] = if s_val > cutoff { 1.0 / s_val } else { 0.0 };
|
||
}
|
||
|
||
let s_inv = Tensor::from_data(s_inv_data, svd_result.s.shape().dims().to_vec(), &self.device)?;
|
||
|
||
// A^+ = V * S^+ * U^T
|
||
// Note: SVD gives us V^T, so we need to transpose it
|
||
let v = svd_result.vt.transpose(0, 1)?; // V^T -> V
|
||
let u_t = svd_result.u.transpose(0, 1)?; // U -> U^T
|
||
let s_inv_diag = Tensor::diag(&s_inv)?;
|
||
|
||
let result = v.matmul(&s_inv_diag)?.matmul(&u_t)?;
|
||
Ok(result)
|
||
}
|
||
|
||
/// Compute matrix condition number (ratio of largest to smallest singular value)
|
||
pub fn condition_number(&self) -> Result<f32> {
|
||
let svd_result = self.svd(false, false)?;
|
||
let s_data = svd_result.s.to_cpu()?;
|
||
|
||
if s_data.is_empty() {
|
||
return Ok(1.0);
|
||
}
|
||
|
||
let s_max = s_data[0]; // Singular values are sorted in descending order
|
||
let s_min = s_data[s_data.len() - 1];
|
||
|
||
if s_min < 1e-12 {
|
||
Ok(f32::INFINITY)
|
||
} else {
|
||
Ok(s_max / s_min)
|
||
}
|
||
}
|
||
|
||
/// Compute matrix rank (number of non-zero singular values)
|
||
pub fn matrix_rank(&self, tol: f32) -> Result<usize> {
|
||
let svd_result = self.svd(false, false)?;
|
||
let s_data = svd_result.s.to_cpu()?;
|
||
|
||
let s_max = s_data.iter().fold(0.0f32, |a, &b| a.max(b));
|
||
let cutoff = tol * s_max;
|
||
|
||
let rank = s_data.iter().filter(|&&s| s > cutoff).count();
|
||
Ok(rank)
|
||
}
|
||
|
||
/// Compute matrix norm
|
||
pub fn matrix_norm(&self, ord: &str) -> Result<f32> {
|
||
match ord {
|
||
"fro" => {
|
||
// Frobenius norm: sqrt(sum of squares of all elements)
|
||
let data = self.to_cpu()?;
|
||
let sum_squares: f32 = data.iter().map(|&x| x * x).sum();
|
||
Ok(sum_squares.sqrt())
|
||
}
|
||
"2" => {
|
||
// Spectral norm: largest singular value
|
||
let svd_result = self.svd(false, false)?;
|
||
let s_data = svd_result.s.to_cpu()?;
|
||
Ok(s_data.get(0).copied().unwrap_or(0.0))
|
||
}
|
||
_ => Err(TensorError::value(format!("Unsupported matrix norm: {}", ord)))
|
||
}
|
||
}
|
||
|
||
/// Compute determinant using LU decomposition
|
||
pub fn det(&self) -> Result<f32> {
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("Determinant requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
if dims[0] != dims[1] {
|
||
return Err(TensorError::shape("Determinant requires square matrix"));
|
||
}
|
||
|
||
let n = dims[0];
|
||
if n == 0 {
|
||
return Ok(1.0);
|
||
}
|
||
if n == 1 {
|
||
let data = self.to_cpu()?;
|
||
return Ok(data[0]);
|
||
}
|
||
if n == 2 {
|
||
let data = self.to_cpu()?;
|
||
return Ok(data[0] * data[3] - data[1] * data[2]);
|
||
}
|
||
|
||
// For larger matrices, use LU decomposition
|
||
let data = self.to_cpu()?;
|
||
let mut a = data.clone();
|
||
let mut det = 1.0f32;
|
||
|
||
// Gaussian elimination with partial pivoting
|
||
for i in 0..n {
|
||
// Find pivot
|
||
let mut max_val = 0.0f32;
|
||
let mut pivot_row = i;
|
||
for k in i..n {
|
||
let val = a[k * n + i].abs();
|
||
if val > max_val {
|
||
max_val = val;
|
||
pivot_row = k;
|
||
}
|
||
}
|
||
|
||
if max_val < 1e-12 {
|
||
return Ok(0.0); // Singular matrix
|
||
}
|
||
|
||
// Swap rows if needed
|
||
if pivot_row != i {
|
||
for j in 0..n {
|
||
let temp = a[i * n + j];
|
||
a[i * n + j] = a[pivot_row * n + j];
|
||
a[pivot_row * n + j] = temp;
|
||
}
|
||
det = -det; // Row swap changes sign
|
||
}
|
||
|
||
det *= a[i * n + i];
|
||
|
||
// Eliminate below diagonal
|
||
for k in (i + 1)..n {
|
||
let factor = a[k * n + i] / a[i * n + i];
|
||
for j in i..n {
|
||
a[k * n + j] -= factor * a[i * n + j];
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(det)
|
||
}
|
||
|
||
/// Compute trace (sum of diagonal elements)
|
||
pub fn trace(&self) -> Result<f32> {
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("Trace requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
let n = dims[0].min(dims[1]);
|
||
let data = self.to_cpu()?;
|
||
let cols = dims[1];
|
||
|
||
let mut trace = 0.0f32;
|
||
for i in 0..n {
|
||
trace += data[i * cols + i];
|
||
}
|
||
|
||
Ok(trace)
|
||
}
|
||
|
||
/// Create diagonal matrix from vector
|
||
pub fn diag(diagonal: &Tensor) -> Result<Tensor> {
|
||
if diagonal.ndim() != 1 {
|
||
return Err(TensorError::shape("diag requires 1D tensor"));
|
||
}
|
||
|
||
let n = diagonal.shape().dims()[0];
|
||
let diag_data = diagonal.to_cpu()?;
|
||
let mut matrix_data = vec![0.0f32; n * n];
|
||
|
||
for i in 0..n {
|
||
matrix_data[i * n + i] = diag_data[i];
|
||
}
|
||
|
||
Tensor::from_data(matrix_data, vec![n, n], &diagonal.device)
|
||
}
|
||
|
||
// Note: eye function is already defined in tensor::core
|
||
|
||
/// Safe reciprocal with threshold for pseudo-inverse computation
|
||
pub fn reciprocal_safe(&self, threshold: f32) -> Result<Tensor> {
|
||
let data = self.to_cpu()?;
|
||
let result_data: Vec<f32> = data.iter()
|
||
.map(|&x| if x.abs() > threshold { 1.0 / x } else { 0.0 })
|
||
.collect();
|
||
|
||
Tensor::from_data(result_data, self.shape().dims().to_vec(), &self.device)
|
||
}
|
||
|
||
/// Solve triangular systems of linear equations: A x = b
|
||
///
|
||
/// Efficiently solves triangular systems using forward or back substitution.
|
||
/// This is much faster than general matrix solve for triangular matrices.
|
||
///
|
||
/// # Arguments
|
||
/// * `b` - Right-hand side tensor(s). Can be 1D (single RHS) or 2D (multiple RHS)
|
||
/// * `options` - Configuration for the solve operation
|
||
///
|
||
/// # Returns
|
||
/// Solution tensor x such that A x = b
|
||
///
|
||
/// # Algorithm Complexity
|
||
/// - Forward substitution (lower): O(n²) for single RHS, O(n²m) for m RHS
|
||
/// - Back substitution (upper): O(n²) for single RHS, O(n²m) for m RHS
|
||
///
|
||
/// # Numerical Stability
|
||
/// - Checks for singular matrices (zero diagonal elements)
|
||
/// - Optional triangular validation with configurable tolerance
|
||
/// - Handles unit diagonal matrices efficiently
|
||
pub fn solve_triangular(&self, b: &Tensor, options: &TriangularSolveOptions) -> Result<Tensor> {
|
||
// Validate inputs
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("solve_triangular requires 2D matrix"));
|
||
}
|
||
|
||
let matrix_dims = self.shape().dims();
|
||
if matrix_dims[0] != matrix_dims[1] {
|
||
return Err(TensorError::shape("solve_triangular requires square matrix"));
|
||
}
|
||
|
||
let n = matrix_dims[0];
|
||
if n == 0 {
|
||
return Err(TensorError::shape("solve_triangular requires non-empty matrix"));
|
||
}
|
||
|
||
// Validate RHS dimensions
|
||
let b_dims = b.shape().dims();
|
||
let is_multiple_rhs = b.ndim() == 2;
|
||
let num_rhs = if is_multiple_rhs { b_dims[1] } else { 1 };
|
||
|
||
if (is_multiple_rhs && b_dims[0] != n) || (!is_multiple_rhs && b.ndim() != 1) || (!is_multiple_rhs && b_dims[0] != n) {
|
||
return Err(TensorError::shape(format!(
|
||
"RHS dimensions {} do not match matrix dimensions [{}x{}]",
|
||
format!("{:?}", b_dims), n, n
|
||
)));
|
||
}
|
||
|
||
// Check device compatibility
|
||
if self.device != b.device {
|
||
return Err(TensorError::device(format!(
|
||
"Matrix and RHS must be on same device: {:?} vs {:?}",
|
||
self.device, b.device
|
||
)));
|
||
}
|
||
|
||
// Get data for computation
|
||
let matrix_data = self.to_cpu()?;
|
||
let b_data = b.to_cpu()?;
|
||
|
||
// Triangular validation if requested
|
||
if options.validate_triangular {
|
||
validate_triangular_matrix(&matrix_data, n, options.triangular_type, options.strict_validation)?;
|
||
}
|
||
|
||
// Prepare result tensor shape and data
|
||
let result_shape = b.shape().dims().to_vec();
|
||
let mut result_data = vec![0.0f32; b.numel()];
|
||
|
||
// Choose algorithm based on options
|
||
if options.transpose {
|
||
// Solve A^T x = b
|
||
match options.triangular_type {
|
||
TriangularType::Lower => {
|
||
// L^T is upper triangular, use back substitution
|
||
solve_upper_triangular_batch(&matrix_data, &b_data, &mut result_data, n, num_rhs, true, options.unit_diagonal)?;
|
||
}
|
||
TriangularType::Upper => {
|
||
// U^T is lower triangular, use forward substitution
|
||
solve_lower_triangular_batch(&matrix_data, &b_data, &mut result_data, n, num_rhs, true, options.unit_diagonal)?;
|
||
}
|
||
}
|
||
} else {
|
||
// Solve A x = b
|
||
match options.triangular_type {
|
||
TriangularType::Lower => {
|
||
solve_lower_triangular_batch(&matrix_data, &b_data, &mut result_data, n, num_rhs, false, options.unit_diagonal)?;
|
||
}
|
||
TriangularType::Upper => {
|
||
solve_upper_triangular_batch(&matrix_data, &b_data, &mut result_data, n, num_rhs, false, options.unit_diagonal)?;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Create result tensor and copy data back
|
||
Tensor::from_data(result_data, result_shape, &self.device)
|
||
}
|
||
|
||
/// Solve least squares problem min ||Ax - b||²
|
||
///
|
||
/// Solves linear least squares problems for overdetermined and underdetermined systems:
|
||
/// - Overdetermined (m > n): Finds x that minimizes ||Ax - b||²
|
||
/// - Underdetermined (m < n): Finds minimum norm solution x such that Ax = b
|
||
/// - Square well-conditioned: Exact solution if it exists
|
||
/// - Rank-deficient: Minimum norm least squares solution
|
||
///
|
||
/// # Arguments
|
||
/// * `b` - Right-hand side vector or matrix (multiple RHS)
|
||
/// * `options` - Optional solving configuration
|
||
///
|
||
/// # Returns
|
||
/// `LeastSquaresResult` containing solution and diagnostic information
|
||
///
|
||
/// # Methods
|
||
/// - **QR**: Best for overdetermined systems (m > n), numerically stable
|
||
/// - **SVD**: Most robust, handles rank-deficient matrices, slower
|
||
/// - **Normal Equations**: Fastest for well-conditioned problems, less stable
|
||
/// - **Auto**: Automatically selects best method based on matrix properties
|
||
///
|
||
/// # Features
|
||
/// - Tikhonov regularization: (A^T A + λI)x = A^T b
|
||
/// - Weighted least squares with diagonal weight matrix
|
||
/// - Automatic condition number estimation and rank detection
|
||
/// - Residual computation and goodness-of-fit metrics
|
||
///
|
||
/// # Example
|
||
/// ```rust
|
||
/// use rtx_tensor::{Tensor, Device};
|
||
/// use rtx_tensor::linalg::LeastSquaresOptions;
|
||
///
|
||
/// let device = Device::cpu();
|
||
/// let a = Tensor::from_data(vec![1.0, 1.0, 2.0, 1.0, 3.0, 1.0], vec![3, 2], &device).unwrap();
|
||
/// let b = Tensor::from_data(vec![6.0, 8.0, 10.0], vec![3], &device).unwrap();
|
||
///
|
||
/// let result = a.least_squares(&b, None).unwrap();
|
||
/// println!("Solution: {:?}", result.solution.to_cpu().unwrap());
|
||
/// println!("Residual norm: {}", result.residual_norm);
|
||
/// ```
|
||
pub fn least_squares(&self, b: &Tensor, options: Option<&LeastSquaresOptions>) -> Result<LeastSquaresResult> {
|
||
// Validate inputs
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("least_squares requires 2D coefficient matrix"));
|
||
}
|
||
|
||
let matrix_dims = self.shape().dims();
|
||
let m = matrix_dims[0]; // number of equations
|
||
let n = matrix_dims[1]; // number of unknowns
|
||
|
||
if m == 0 || n == 0 {
|
||
return Err(TensorError::shape("least_squares requires non-empty matrix"));
|
||
}
|
||
|
||
// Validate RHS dimensions
|
||
let b_dims = b.shape().dims();
|
||
let is_multiple_rhs = b.ndim() == 2;
|
||
let num_rhs = if is_multiple_rhs { b_dims[1] } else { 1 };
|
||
|
||
if (is_multiple_rhs && b_dims[0] != m) || (!is_multiple_rhs && b.ndim() != 1) || (!is_multiple_rhs && b_dims[0] != m) {
|
||
return Err(TensorError::shape(format!(
|
||
"RHS dimensions {:?} do not match matrix rows {}",
|
||
b_dims, m
|
||
)));
|
||
}
|
||
|
||
// Check device compatibility
|
||
if self.device != b.device {
|
||
return Err(TensorError::device(format!(
|
||
"Matrix and RHS must be on same device: {:?} vs {:?}",
|
||
self.device, b.device
|
||
)));
|
||
}
|
||
|
||
// Use default options if none provided
|
||
let default_opts = LeastSquaresOptions::default();
|
||
let opts = options.unwrap_or(&default_opts);
|
||
|
||
// Determine method to use
|
||
let method_to_use = match opts.method {
|
||
LeastSquaresMethod::Auto => {
|
||
// Auto-select based on matrix properties
|
||
if m >= n {
|
||
// Overdetermined or square - use QR for stability
|
||
LeastSquaresMethod::QR
|
||
} else {
|
||
// Underdetermined - use SVD for minimum norm solution
|
||
LeastSquaresMethod::SVD
|
||
}
|
||
}
|
||
specified_method => specified_method,
|
||
};
|
||
|
||
// Solve based on selected method
|
||
match method_to_use {
|
||
LeastSquaresMethod::QR => {
|
||
self.solve_least_squares_qr(b, opts, num_rhs, is_multiple_rhs)
|
||
}
|
||
LeastSquaresMethod::SVD => {
|
||
self.solve_least_squares_svd(b, opts, num_rhs, is_multiple_rhs)
|
||
}
|
||
LeastSquaresMethod::NormalEquations => {
|
||
self.solve_least_squares_normal_equations(b, opts, num_rhs, is_multiple_rhs)
|
||
}
|
||
LeastSquaresMethod::Auto => unreachable!(), // Already handled above
|
||
}
|
||
}
|
||
|
||
/// Solve least squares using QR decomposition
|
||
/// Best for overdetermined systems (m >= n)
|
||
fn solve_least_squares_qr(&self, b: &Tensor, opts: &LeastSquaresOptions, num_rhs: usize, is_multiple_rhs: bool) -> Result<LeastSquaresResult> {
|
||
// Apply weights if provided
|
||
let (weighted_a, weighted_b) = if let Some(ref weights) = opts.weights {
|
||
self.apply_weights(b, weights)?
|
||
} else {
|
||
(self.clone(), b.clone())
|
||
};
|
||
|
||
// Compute QR decomposition
|
||
let qr_result = weighted_a.qr()?;
|
||
|
||
// Solve Rx = Q^T b using back substitution
|
||
let qt = qr_result.q.transpose(0, 1)?;
|
||
let qtb = qt.matmul(&weighted_b)?;
|
||
|
||
// Solve upper triangular system R x = Q^T b
|
||
let options = TriangularSolveOptions {
|
||
triangular_type: TriangularType::Upper,
|
||
transpose: false,
|
||
unit_diagonal: false,
|
||
validate_triangular: false, // We trust QR decomposition
|
||
strict_validation: false,
|
||
};
|
||
|
||
let solution = qr_result.r.solve_triangular(&qtb, &options)?;
|
||
|
||
// Compute residual and diagnostic information
|
||
let residual_result = self.compute_residuals_and_metrics(&solution, b, &weighted_a)?;
|
||
|
||
Ok(LeastSquaresResult {
|
||
solution,
|
||
residual_norm: residual_result.residual_norm,
|
||
condition_number: residual_result.condition_number,
|
||
rank: residual_result.rank,
|
||
method_used: LeastSquaresMethod::QR,
|
||
r_squared: residual_result.r_squared,
|
||
sum_squared_residuals: residual_result.sum_squared_residuals,
|
||
})
|
||
}
|
||
|
||
/// Solve least squares using SVD decomposition
|
||
/// Most robust method, handles rank-deficient matrices
|
||
fn solve_least_squares_svd(&self, b: &Tensor, opts: &LeastSquaresOptions, num_rhs: usize, is_multiple_rhs: bool) -> Result<LeastSquaresResult> {
|
||
// Apply weights if provided
|
||
let (weighted_a, weighted_b) = if let Some(ref weights) = opts.weights {
|
||
self.apply_weights(b, weights)?
|
||
} else {
|
||
(self.clone(), b.clone())
|
||
};
|
||
|
||
// Compute SVD: A = U S V^T
|
||
let svd_result = weighted_a.svd(true, true)?;
|
||
|
||
// Compute pseudo-inverse using SVD with rcond threshold
|
||
let s_data = svd_result.s.to_cpu()?;
|
||
let s_max = s_data.iter().fold(0.0f32, |a, &b| a.max(b));
|
||
let cutoff = opts.rcond * s_max;
|
||
|
||
// Create S^+ (pseudo-inverse of singular values)
|
||
let mut s_pinv_data = vec![0.0f32; s_data.len()];
|
||
let mut rank = 0;
|
||
for (i, &s_val) in s_data.iter().enumerate() {
|
||
if s_val > cutoff {
|
||
s_pinv_data[i] = 1.0 / s_val;
|
||
rank += 1;
|
||
} else {
|
||
s_pinv_data[i] = 0.0;
|
||
}
|
||
}
|
||
|
||
let s_pinv = Tensor::from_data(s_pinv_data, svd_result.s.shape().dims().to_vec(), &self.device)?;
|
||
let s_pinv_diag = Tensor::diag(&s_pinv)?;
|
||
|
||
// Compute x = V S^+ U^T b
|
||
let ut = svd_result.u.transpose(0, 1)?;
|
||
let v = svd_result.vt.transpose(0, 1)?;
|
||
|
||
let utb = ut.matmul(&weighted_b)?;
|
||
let s_pinv_utb = s_pinv_diag.matmul(&utb)?;
|
||
let solution = v.matmul(&s_pinv_utb)?;
|
||
|
||
// Compute residual and diagnostic information
|
||
let residual_result = self.compute_residuals_and_metrics(&solution, b, &weighted_a)?;
|
||
|
||
Ok(LeastSquaresResult {
|
||
solution,
|
||
residual_norm: residual_result.residual_norm,
|
||
condition_number: if s_data[s_data.len() - 1] > 1e-15 { s_max / s_data[s_data.len() - 1] } else { f32::INFINITY },
|
||
rank,
|
||
method_used: LeastSquaresMethod::SVD,
|
||
r_squared: residual_result.r_squared,
|
||
sum_squared_residuals: residual_result.sum_squared_residuals,
|
||
})
|
||
}
|
||
|
||
/// Solve least squares using normal equations (A^T A)x = A^T b
|
||
/// Fastest but least numerically stable method
|
||
fn solve_least_squares_normal_equations(&self, b: &Tensor, opts: &LeastSquaresOptions, num_rhs: usize, is_multiple_rhs: bool) -> Result<LeastSquaresResult> {
|
||
// Apply weights if provided
|
||
let (weighted_a, weighted_b) = if let Some(ref weights) = opts.weights {
|
||
self.apply_weights(b, weights)?
|
||
} else {
|
||
(self.clone(), b.clone())
|
||
};
|
||
|
||
// Compute A^T A and A^T b
|
||
let at = weighted_a.transpose(0, 1)?;
|
||
let mut ata = at.matmul(&weighted_a)?;
|
||
let atb = at.matmul(&weighted_b)?;
|
||
|
||
// Apply Tikhonov regularization if specified
|
||
if let Some(lambda) = opts.regularization_lambda {
|
||
let n = ata.shape().dims()[0];
|
||
let identity = Tensor::eye(n, &self.device)?;
|
||
let regularization = identity.mul_scalar(lambda)?;
|
||
ata = ata.add(®ularization)?;
|
||
}
|
||
|
||
// Solve (A^T A + λI)x = A^T b using Cholesky decomposition
|
||
let l = ata.cholesky(false)?;
|
||
|
||
// Forward substitution: solve L y = A^T b
|
||
let forward_opts = TriangularSolveOptions {
|
||
triangular_type: TriangularType::Lower,
|
||
transpose: false,
|
||
unit_diagonal: false,
|
||
validate_triangular: false,
|
||
strict_validation: false,
|
||
};
|
||
let y = l.solve_triangular(&atb, &forward_opts)?;
|
||
|
||
// Back substitution: solve L^T x = y
|
||
let back_opts = TriangularSolveOptions {
|
||
triangular_type: TriangularType::Lower,
|
||
transpose: true,
|
||
unit_diagonal: false,
|
||
validate_triangular: false,
|
||
strict_validation: false,
|
||
};
|
||
let solution = l.solve_triangular(&y, &back_opts)?;
|
||
|
||
// Compute residual and diagnostic information
|
||
let residual_result = self.compute_residuals_and_metrics(&solution, b, &weighted_a)?;
|
||
|
||
Ok(LeastSquaresResult {
|
||
solution,
|
||
residual_norm: residual_result.residual_norm,
|
||
condition_number: residual_result.condition_number,
|
||
rank: residual_result.rank,
|
||
method_used: LeastSquaresMethod::NormalEquations,
|
||
r_squared: residual_result.r_squared,
|
||
sum_squared_residuals: residual_result.sum_squared_residuals,
|
||
})
|
||
}
|
||
|
||
/// Matrix power (exponentiation) A^n
|
||
///
|
||
/// Computes matrix exponentiation using efficient algorithms:
|
||
/// - For n = 0: Returns identity matrix
|
||
/// - For n > 0: Uses binary exponentiation (repeated squaring)
|
||
/// - For n < 0: Computes A^(-|n|) = (A^(-1))^|n| using matrix inverse
|
||
///
|
||
/// # Arguments
|
||
/// - `n`: Integer power (can be positive, negative, or zero)
|
||
///
|
||
/// # Returns
|
||
/// Matrix A raised to power n
|
||
///
|
||
/// # Algorithm
|
||
/// Uses binary exponentiation for efficiency:
|
||
/// - If n = 2k: A^n = (A^k)^2
|
||
/// - If n = 2k+1: A^n = A * (A^k)^2
|
||
/// - Complexity: O(log|n|) matrix multiplications
|
||
///
|
||
/// # Errors
|
||
/// - Returns error if matrix is not square
|
||
/// - Returns error if n < 0 and matrix is singular (not invertible)
|
||
pub fn matrix_power(&self, n: i32) -> Result<Tensor> {
|
||
// Input validation
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("matrix_power requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
if dims[0] != dims[1] {
|
||
return Err(TensorError::shape("matrix_power requires square matrix"));
|
||
}
|
||
|
||
let size = dims[0];
|
||
if size == 0 {
|
||
return Err(TensorError::shape("matrix_power requires non-empty matrix"));
|
||
}
|
||
|
||
// Special case: A^0 = I (identity matrix)
|
||
if n == 0 {
|
||
return Tensor::eye(size, &self.device);
|
||
}
|
||
|
||
// Special case: A^1 = A
|
||
if n == 1 {
|
||
return Ok(self.clone());
|
||
}
|
||
|
||
// Handle negative powers: A^n = (A^-1)^|n|
|
||
if n < 0 {
|
||
// Check if matrix is invertible by computing determinant
|
||
let det = self.det()?;
|
||
if det.abs() < 1e-12 {
|
||
return Err(TensorError::numerical(
|
||
"Cannot compute negative power of singular matrix (determinant ≈ 0)"
|
||
));
|
||
}
|
||
|
||
// Compute inverse and then raise to positive power
|
||
let inverse = self.pinv(1e-12)?;
|
||
return inverse.matrix_power((-n) as i32);
|
||
}
|
||
|
||
// For positive powers, use binary exponentiation
|
||
self.matrix_power_binary_exponentiation(n as u32)
|
||
}
|
||
|
||
/// Matrix power with floating-point exponent using eigendecomposition
|
||
///
|
||
/// Computes A^p for real power p using eigendecomposition.
|
||
/// For matrix A with eigendecomposition A = V*D*V^-1,
|
||
/// we have A^p = V*D^p*V^-1 where D^p applies power element-wise.
|
||
///
|
||
/// # Arguments
|
||
/// - `p`: Real-valued power
|
||
///
|
||
/// # Returns
|
||
/// Matrix A raised to power p
|
||
///
|
||
/// # Requirements
|
||
/// - Matrix must be square and diagonalizable
|
||
/// - For non-integer powers, matrix should have positive eigenvalues for real result
|
||
///
|
||
/// # Note
|
||
/// This is more expensive than integer powers but supports fractional exponents.
|
||
pub fn matrix_power_f32(&self, p: f32) -> Result<Tensor> {
|
||
// Input validation
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("matrix_power_f32 requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
if dims[0] != dims[1] {
|
||
return Err(TensorError::shape("matrix_power_f32 requires square matrix"));
|
||
}
|
||
|
||
// Special case: p = 0 gives identity
|
||
if (p - 0.0).abs() < 1e-12 {
|
||
return Tensor::eye(dims[0], &self.device);
|
||
}
|
||
|
||
// Special case: p = 1 gives original matrix
|
||
if (p - 1.0).abs() < 1e-12 {
|
||
return Ok(self.clone());
|
||
}
|
||
|
||
// For integer powers, use more efficient integer algorithm
|
||
if (p.fract()).abs() < 1e-12 {
|
||
return self.matrix_power(p as i32);
|
||
}
|
||
|
||
// Use eigendecomposition for fractional powers
|
||
// For now, implement simplified version for diagonal matrices
|
||
// Full eigendecomposition would be more complex
|
||
|
||
// Check if matrix is diagonal (simplified case)
|
||
if self.is_diagonal()? {
|
||
return self.matrix_power_diagonal_fractional(p);
|
||
}
|
||
|
||
// For general matrices, use eigendecomposition via symmetric case
|
||
// This is a simplified implementation - production would handle non-symmetric matrices
|
||
let eigen_result = self.symeig(true)?;
|
||
|
||
// Compute D^p (eigenvalues raised to power p)
|
||
let eigenvals_data = eigen_result.eigenvalues.to_cpu()?;
|
||
let mut powered_eigenvals = vec![0.0f32; eigenvals_data.len()];
|
||
|
||
for (i, &lambda) in eigenvals_data.iter().enumerate() {
|
||
if lambda < 0.0 && (p.fract()).abs() > 1e-12 {
|
||
return Err(TensorError::numerical(
|
||
"Fractional power of matrix with negative eigenvalues not supported"
|
||
));
|
||
}
|
||
powered_eigenvals[i] = lambda.powf(p);
|
||
}
|
||
|
||
let powered_eigenvals_tensor = Tensor::from_data(powered_eigenvals, vec![eigenvals_data.len()], &self.device)?;
|
||
let d_powered = Tensor::diag(&powered_eigenvals_tensor)?;
|
||
|
||
// Reconstruct: A^p = V * D^p * V^-1
|
||
let v_inv = eigen_result.eigenvectors.pinv(1e-12)?;
|
||
let result = eigen_result.eigenvectors.matmul(&d_powered)?.matmul(&v_inv)?;
|
||
|
||
Ok(result)
|
||
}
|
||
|
||
/// Binary exponentiation for positive integer powers
|
||
/// Implements efficient O(log n) algorithm using repeated squaring
|
||
fn matrix_power_binary_exponentiation(&self, mut n: u32) -> Result<Tensor> {
|
||
if n == 0 {
|
||
return Tensor::eye(self.shape().dims()[0], &self.device);
|
||
}
|
||
|
||
let mut base = self.clone();
|
||
let mut result = Tensor::eye(self.shape().dims()[0], &self.device)?;
|
||
|
||
while n > 0 {
|
||
// If n is odd, multiply result by current base
|
||
if n % 2 == 1 {
|
||
result = result.matmul(&base)?;
|
||
}
|
||
|
||
// Square the base and halve the exponent
|
||
if n > 1 {
|
||
base = base.matmul(&base)?;
|
||
}
|
||
n /= 2;
|
||
}
|
||
|
||
Ok(result)
|
||
}
|
||
|
||
/// Check if matrix is diagonal
|
||
fn is_diagonal(&self) -> Result<bool> {
|
||
let dims = self.shape().dims();
|
||
let n = dims[0];
|
||
let data = self.to_cpu()?;
|
||
|
||
for i in 0..n {
|
||
for j in 0..n {
|
||
if i != j && data[i * n + j].abs() > 1e-12 {
|
||
return Ok(false);
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(true)
|
||
}
|
||
|
||
/// Compute fractional power for diagonal matrix
|
||
fn matrix_power_diagonal_fractional(&self, p: f32) -> Result<Tensor> {
|
||
let dims = self.shape().dims();
|
||
let n = dims[0];
|
||
let data = self.to_cpu()?;
|
||
let mut result_data = vec![0.0f32; n * n];
|
||
|
||
// For diagonal matrix, just raise each diagonal element to power p
|
||
for i in 0..n {
|
||
let diagonal_val = data[i * n + i];
|
||
if diagonal_val < 0.0 && (p.fract()).abs() > 1e-12 {
|
||
return Err(TensorError::numerical(
|
||
"Fractional power of negative diagonal element not supported"
|
||
));
|
||
}
|
||
result_data[i * n + i] = diagonal_val.powf(p);
|
||
}
|
||
|
||
Tensor::from_data(result_data, dims.to_vec(), &self.device)
|
||
}
|
||
|
||
/// LU Decomposition with partial pivoting
|
||
///
|
||
/// Decomposes a matrix A into P*A = L*U where:
|
||
/// - P is a permutation matrix representing row exchanges
|
||
/// - L is a lower triangular matrix with 1s on the diagonal (Doolittle algorithm)
|
||
/// - U is an upper triangular matrix
|
||
///
|
||
/// # Arguments
|
||
/// - `options`: Configuration options for the decomposition
|
||
///
|
||
/// # Returns
|
||
/// `LUResult` containing L, U, permutation information, and computed properties
|
||
///
|
||
/// # Algorithms
|
||
/// Uses Gaussian elimination with partial pivoting for numerical stability.
|
||
/// The algorithm selects the largest element in each column as the pivot
|
||
/// to minimize roundoff errors.
|
||
///
|
||
/// # Features
|
||
/// - Partial pivoting for numerical stability
|
||
/// - Supports rectangular matrices (m×n)
|
||
/// - Efficient determinant computation
|
||
/// - Singularity detection
|
||
/// - In-place option for memory efficiency
|
||
///
|
||
/// # Example
|
||
/// ```rust
|
||
/// use rtx_tensor::{Tensor, Device};
|
||
/// use rtx_tensor::linalg::LUOptions;
|
||
///
|
||
/// let device = Device::cpu();
|
||
/// let a = Tensor::from_data(vec![2.0, 1.0, 1.0, 1.0], vec![2, 2], &device).unwrap();
|
||
///
|
||
/// let options = LUOptions::default();
|
||
/// let lu_result = a.lu_decomposition(&options).unwrap();
|
||
///
|
||
/// println!("L matrix: {:?}", lu_result.l.to_cpu().unwrap());
|
||
/// println!("U matrix: {:?}", lu_result.u.to_cpu().unwrap());
|
||
/// println!("Determinant: {}", lu_result.determinant);
|
||
/// ```
|
||
pub fn lu_decomposition(&self, options: &LUOptions) -> Result<LUResult> {
|
||
// Input validation
|
||
if self.ndim() != 2 {
|
||
return Err(TensorError::shape("LU decomposition requires 2D tensor"));
|
||
}
|
||
|
||
let dims = self.shape().dims();
|
||
let m = dims[0]; // rows
|
||
let n = dims[1]; // cols
|
||
|
||
if m == 0 || n == 0 {
|
||
return Err(TensorError::shape("LU decomposition requires non-empty matrix"));
|
||
}
|
||
|
||
// Get matrix data
|
||
let data = self.to_cpu()?;
|
||
|
||
// Choose algorithm based on options
|
||
if options.pivoting {
|
||
self.lu_decomposition_with_pivoting(&data, m, n, options)
|
||
} else {
|
||
self.lu_decomposition_without_pivoting(&data, m, n, options)
|
||
}
|
||
}
|
||
|
||
/// Batch LU decomposition for multiple matrices
|
||
///
|
||
/// Performs LU decomposition on a batch of matrices efficiently.
|
||
/// Input should be a tensor of shape (batch_size, m, n).
|
||
///
|
||
/// # Arguments
|
||
/// - `options`: Configuration options applied to all matrices in the batch
|
||
///
|
||
/// # Returns
|
||
/// `LUResult` with batched L, U, and permutation tensors
|
||
pub fn lu_decomposition_batch(&self, options: &LUOptions) -> Result<LUResult> {
|
||
// This will be implemented later
|
||
Err(TensorError::not_implemented("Batch LU decomposition not yet implemented"))
|
||
}
|
||
|
||
/// LU decomposition without pivoting using Doolittle algorithm
|
||
///
|
||
/// Performs A = L*U decomposition where L has 1s on diagonal
|
||
/// This method assumes the matrix doesn't need pivoting for numerical stability
|
||
fn lu_decomposition_without_pivoting(&self, data: &[f32], m: usize, n: usize, options: &LUOptions) -> Result<LUResult> {
|
||
let min_dim = m.min(n);
|
||
|
||
// Create working matrix (copy of input)
|
||
let mut a = data.to_vec();
|
||
|
||
// Initialize L and U matrices
|
||
let mut l_data = vec![0.0f32; m * min_dim];
|
||
let mut u_data = if m <= n {
|
||
a.clone() // U will be stored in-place in upper part
|
||
} else {
|
||
vec![0.0f32; min_dim * n] // U is smaller for tall matrices
|
||
};
|
||
|
||
// Perform Doolittle LU decomposition: A = L*U
|
||
// L has 1s on diagonal, U has the computed values
|
||
for k in 0..min_dim {
|
||
// Check for zero pivot (singularity)
|
||
if a[k * n + k].abs() < 1e-15 {
|
||
// Matrix is singular - continue but mark as non-invertible
|
||
}
|
||
|
||
// Set L diagonal to 1 (Doolittle)
|
||
l_data[k * min_dim + k] = 1.0;
|
||
|
||
// Compute L column k (below diagonal)
|
||
for i in (k + 1)..m {
|
||
if a[k * n + k].abs() > 1e-15 {
|
||
l_data[i * min_dim + k] = a[i * n + k] / a[k * n + k];
|
||
} else {
|
||
l_data[i * min_dim + k] = 0.0;
|
||
}
|
||
}
|
||
|
||
// Update remaining submatrix: A[i,j] -= L[i,k] * U[k,j]
|
||
for i in (k + 1)..m {
|
||
for j in (k + 1)..n {
|
||
a[i * n + j] -= l_data[i * min_dim + k] * a[k * n + j];
|
||
}
|
||
}
|
||
}
|
||
|
||
// Extract U matrix from upper triangular part of A
|
||
if m <= n {
|
||
// Square or wide matrix: U is min_dim x n
|
||
for i in 0..min_dim {
|
||
for j in i..n {
|
||
u_data[i * n + j] = a[i * n + j];
|
||
}
|
||
// Zero out below diagonal
|
||
for j in 0..i {
|
||
u_data[i * n + j] = 0.0;
|
||
}
|
||
}
|
||
} else {
|
||
// Tall matrix: U is min_dim x n
|
||
for i in 0..min_dim {
|
||
for j in 0..n {
|
||
if i <= j {
|
||
u_data[i * n + j] = a[i * n + j];
|
||
} else {
|
||
u_data[i * n + j] = 0.0;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Create result tensors
|
||
let l_shape = vec![m, min_dim];
|
||
let u_shape = if m <= n { vec![min_dim, n] } else { vec![min_dim, n] };
|
||
|
||
let l_tensor = Tensor::from_data(l_data, l_shape, &self.device)?;
|
||
let u_tensor = Tensor::from_data(u_data, u_shape, &self.device)?;
|
||
|
||
// Create identity permutation (no row swaps)
|
||
let p_data: Vec<f32> = (0..m).map(|i| i as f32).collect();
|
||
let p_tensor = Tensor::from_data(p_data, vec![m], &self.device)?;
|
||
|
||
// Compute determinant (product of U diagonal elements)
|
||
let mut determinant = 1.0f32;
|
||
for i in 0..min_dim.min(n) {
|
||
determinant *= a[i * n + i];
|
||
}
|
||
|
||
// Check if matrix is invertible (non-singular)
|
||
let is_invertible = determinant.abs() > 1e-12 && m == n;
|
||
|
||
// Estimate condition number (simplified)
|
||
let condition_number = self.condition_number().unwrap_or(1.0);
|
||
|
||
Ok(LUResult {
|
||
l: l_tensor,
|
||
u: u_tensor,
|
||
p: p_tensor,
|
||
p_matrix: if options.return_permutation_matrix {
|
||
Some(Tensor::eye(m, &self.device)?) // Identity matrix for no pivoting
|
||
} else {
|
||
None
|
||
},
|
||
determinant,
|
||
is_invertible,
|
||
condition_number,
|
||
num_pivots: 0, // No pivots in this algorithm
|
||
})
|
||
}
|
||
|
||
/// LU decomposition with partial pivoting using Doolittle algorithm
|
||
///
|
||
/// Performs P*A = L*U decomposition with row pivoting for numerical stability
|
||
fn lu_decomposition_with_pivoting(&self, data: &[f32], m: usize, n: usize, options: &LUOptions) -> Result<LUResult> {
|
||
let min_dim = m.min(n);
|
||
|
||
// Create working matrix (copy of input)
|
||
let mut a = data.to_vec();
|
||
|
||
// Track permutation
|
||
let mut permutation: Vec<usize> = (0..m).collect();
|
||
let mut num_pivots = 0;
|
||
|
||
// Initialize L matrix
|
||
let mut l_data = vec![0.0f32; m * min_dim];
|
||
|
||
// Perform LU decomposition with partial pivoting
|
||
for k in 0..min_dim {
|
||
// Find the best pivot in column k (largest absolute value)
|
||
let mut pivot_row = k;
|
||
let mut max_val = 0.0f32;
|
||
|
||
for i in k..m {
|
||
let val = a[i * n + k].abs();
|
||
if val > max_val {
|
||
max_val = val;
|
||
pivot_row = i;
|
||
}
|
||
}
|
||
|
||
// Check for singularity
|
||
if max_val < 1e-15 {
|
||
// Column is effectively zero - matrix is singular
|
||
// Continue with zero pivot but mark as non-invertible
|
||
}
|
||
|
||
// Swap rows if needed
|
||
if pivot_row != k {
|
||
// Swap rows in matrix A
|
||
for j in 0..n {
|
||
let temp = a[k * n + j];
|
||
a[k * n + j] = a[pivot_row * n + j];
|
||
a[pivot_row * n + j] = temp;
|
||
}
|
||
|
||
// Swap rows in L matrix (for already computed part)
|
||
for j in 0..k {
|
||
let temp = l_data[k * min_dim + j];
|
||
l_data[k * min_dim + j] = l_data[pivot_row * min_dim + j];
|
||
l_data[pivot_row * min_dim + j] = temp;
|
||
}
|
||
|
||
// Update permutation
|
||
permutation.swap(k, pivot_row);
|
||
num_pivots += 1;
|
||
}
|
||
|
||
// Set L diagonal to 1 (Doolittle)
|
||
l_data[k * min_dim + k] = 1.0;
|
||
|
||
// Compute L column k (below diagonal)
|
||
for i in (k + 1)..m {
|
||
if a[k * n + k].abs() > 1e-15 {
|
||
l_data[i * min_dim + k] = a[i * n + k] / a[k * n + k];
|
||
} else {
|
||
l_data[i * min_dim + k] = 0.0;
|
||
}
|
||
}
|
||
|
||
// Update remaining submatrix: A[i,j] -= L[i,k] * U[k,j]
|
||
for i in (k + 1)..m {
|
||
for j in (k + 1)..n {
|
||
a[i * n + j] -= l_data[i * min_dim + k] * a[k * n + j];
|
||
}
|
||
}
|
||
}
|
||
|
||
// Extract U matrix from upper triangular part of A
|
||
let u_shape = if m <= n { vec![min_dim, n] } else { vec![min_dim, n] };
|
||
let mut u_data = vec![0.0f32; u_shape[0] * u_shape[1]];
|
||
|
||
for i in 0..min_dim {
|
||
for j in i..n {
|
||
if j < n {
|
||
u_data[i * n + j] = a[i * n + j];
|
||
}
|
||
}
|
||
}
|
||
|
||
// Create result tensors
|
||
let l_tensor = Tensor::from_data(l_data, vec![m, min_dim], &self.device)?;
|
||
let u_tensor = Tensor::from_data(u_data, u_shape, &self.device)?;
|
||
|
||
// Create permutation vector
|
||
let p_data: Vec<f32> = permutation.iter().map(|&i| i as f32).collect();
|
||
let p_tensor = Tensor::from_data(p_data, vec![m], &self.device)?;
|
||
|
||
// Create permutation matrix if requested
|
||
let p_matrix = if options.return_permutation_matrix {
|
||
let mut p_matrix_data = vec![0.0f32; m * m];
|
||
for (i, &perm_i) in permutation.iter().enumerate() {
|
||
p_matrix_data[i * m + perm_i] = 1.0;
|
||
}
|
||
Some(Tensor::from_data(p_matrix_data, vec![m, m], &self.device)?)
|
||
} else {
|
||
None
|
||
};
|
||
|
||
// Compute determinant (product of U diagonal * sign of permutation)
|
||
let mut determinant = 1.0f32;
|
||
for i in 0..min_dim.min(n) {
|
||
determinant *= a[i * n + i];
|
||
}
|
||
// Adjust for row swaps (each swap changes sign)
|
||
if num_pivots % 2 == 1 {
|
||
determinant = -determinant;
|
||
}
|
||
|
||
// Check if matrix is invertible
|
||
let is_invertible = determinant.abs() > 1e-12 && m == n;
|
||
|
||
// Estimate condition number
|
||
let condition_number = self.condition_number().unwrap_or(1.0);
|
||
|
||
Ok(LUResult {
|
||
l: l_tensor,
|
||
u: u_tensor,
|
||
p: p_tensor,
|
||
p_matrix,
|
||
determinant,
|
||
is_invertible,
|
||
condition_number,
|
||
num_pivots,
|
||
})
|
||
}
|
||
|
||
/// Apply diagonal weights to the least squares problem
|
||
/// Transform Ax = b to sqrt(W)A x = sqrt(W)b where W is diagonal weight matrix
|
||
fn apply_weights(&self, b: &Tensor, weights: &Tensor) -> Result<(Tensor, Tensor)> {
|
||
// Weights should be a 1D tensor matching number of rows
|
||
if weights.ndim() != 1 {
|
||
return Err(TensorError::shape("Weights must be 1D tensor"));
|
||
}
|
||
|
||
let m = self.shape().dims()[0];
|
||
if weights.shape().dims()[0] != m {
|
||
return Err(TensorError::shape("Weights length must match number of equations"));
|
||
}
|
||
|
||
// Compute sqrt of weights for scaling
|
||
let weights_data = weights.to_cpu()?;
|
||
let sqrt_weights_data: Vec<f32> = weights_data.iter()
|
||
.map(|&w| {
|
||
if w < 0.0 {
|
||
panic!("Weights must be non-negative");
|
||
}
|
||
w.sqrt()
|
||
})
|
||
.collect();
|
||
|
||
let sqrt_weights = Tensor::from_data(sqrt_weights_data, vec![m], &self.device)?;
|
||
|
||
// Scale each row of A and b by sqrt(weight)
|
||
let a_data = self.to_cpu()?;
|
||
let b_data = b.to_cpu()?;
|
||
let n = self.shape().dims()[1];
|
||
|
||
let mut weighted_a_data = vec![0.0f32; m * n];
|
||
let mut weighted_b_data = vec![0.0f32; b.numel()];
|
||
|
||
let sqrt_weights_cpu = sqrt_weights.to_cpu()?;
|
||
|
||
for i in 0..m {
|
||
let weight = sqrt_weights_cpu[i];
|
||
// Scale row i of A
|
||
for j in 0..n {
|
||
weighted_a_data[i * n + j] = a_data[i * n + j] * weight;
|
||
}
|
||
// Scale row i of b
|
||
if b.ndim() == 1 {
|
||
weighted_b_data[i] = b_data[i] * weight;
|
||
} else {
|
||
// Multiple RHS case
|
||
let b_cols = b.shape().dims()[1];
|
||
for j in 0..b_cols {
|
||
weighted_b_data[i * b_cols + j] = b_data[i * b_cols + j] * weight;
|
||
}
|
||
}
|
||
}
|
||
|
||
let weighted_a = Tensor::from_data(weighted_a_data, self.shape().dims().to_vec(), &self.device)?;
|
||
let weighted_b = Tensor::from_data(weighted_b_data, b.shape().dims().to_vec(), &self.device)?;
|
||
|
||
Ok((weighted_a, weighted_b))
|
||
}
|
||
}
|
||
|
||
impl LUResult {
|
||
/// Apply permutation to a matrix or vector
|
||
///
|
||
/// Computes P*A where P is the permutation matrix from LU decomposition.
|
||
/// This is used to verify the decomposition: P*A = L*U
|
||
pub fn apply_permutation(&self, matrix: &Tensor) -> Result<Tensor> {
|
||
if matrix.ndim() == 0 || matrix.ndim() > 2 {
|
||
return Err(TensorError::shape("Permutation can only be applied to 1D or 2D tensors"));
|
||
}
|
||
|
||
let matrix_data = matrix.to_cpu()?;
|
||
let p_data = self.p.to_cpu()?;
|
||
let permutation: Vec<usize> = p_data.iter().map(|&x| x as usize).collect();
|
||
|
||
if matrix.ndim() == 1 {
|
||
// Apply permutation to vector
|
||
let n = matrix.shape().dims()[0];
|
||
if n != permutation.len() {
|
||
return Err(TensorError::shape("Vector length doesn't match permutation size"));
|
||
}
|
||
|
||
let mut result_data = vec![0.0f32; n];
|
||
for (i, &perm_i) in permutation.iter().enumerate() {
|
||
if perm_i < n {
|
||
result_data[i] = matrix_data[perm_i];
|
||
} else {
|
||
return Err(TensorError::shape("Invalid permutation index"));
|
||
}
|
||
}
|
||
|
||
Tensor::from_data(result_data, vec![n], &matrix.device)
|
||
} else {
|
||
// Apply permutation to matrix (permute rows)
|
||
let dims = matrix.shape().dims();
|
||
let m = dims[0];
|
||
let n = dims[1];
|
||
|
||
if m != permutation.len() {
|
||
return Err(TensorError::shape("Matrix rows don't match permutation size"));
|
||
}
|
||
|
||
let mut result_data = vec![0.0f32; m * n];
|
||
for (new_row, &old_row) in permutation.iter().enumerate() {
|
||
if old_row < m {
|
||
for j in 0..n {
|
||
result_data[new_row * n + j] = matrix_data[old_row * n + j];
|
||
}
|
||
} else {
|
||
return Err(TensorError::shape("Invalid permutation index"));
|
||
}
|
||
}
|
||
|
||
Tensor::from_data(result_data, dims.to_vec(), &matrix.device)
|
||
}
|
||
}
|
||
|
||
/// Apply permutation for batch operations
|
||
///
|
||
/// Applies permutation from batch index i to the corresponding matrix
|
||
pub fn apply_permutation_batch(&self, matrix: &Tensor, batch_idx: usize) -> Result<Tensor> {
|
||
// This will be implemented
|
||
Err(TensorError::not_implemented("Batch permutation application not yet implemented"))
|
||
}
|
||
|
||
/// Solve linear system using LU decomposition
|
||
///
|
||
/// Solves Ax = b using the precomputed LU decomposition.
|
||
/// This is much faster than recomputing the decomposition for each solve.
|
||
///
|
||
/// # Algorithm
|
||
/// 1. Solve Ly = Pb using forward substitution
|
||
/// 2. Solve Ux = y using back substitution
|
||
///
|
||
/// # Arguments
|
||
/// - `b`: Right-hand side vector or matrix (multiple RHS supported)
|
||
///
|
||
/// # Returns
|
||
/// Solution tensor x such that Ax = b
|
||
pub fn solve(&self, b: &Tensor) -> Result<Tensor> {
|
||
if b.ndim() == 0 || b.ndim() > 2 {
|
||
return Err(TensorError::shape("RHS must be 1D or 2D tensor"));
|
||
}
|
||
|
||
let b_dims = b.shape().dims();
|
||
let n = b_dims[0];
|
||
|
||
// Check dimensions compatibility
|
||
let l_dims = self.l.shape().dims();
|
||
let u_dims = self.u.shape().dims();
|
||
|
||
if l_dims[0] != n {
|
||
return Err(TensorError::shape(format!(
|
||
"L matrix rows {} don't match RHS rows {}",
|
||
l_dims[0], n
|
||
)));
|
||
}
|
||
|
||
if u_dims[0] != u_dims[1] {
|
||
return Err(TensorError::shape("Cannot solve with rectangular U matrix - need square system"));
|
||
}
|
||
|
||
// Step 1: Apply permutation to b to get Pb
|
||
let pb = self.apply_permutation(b)?;
|
||
|
||
// Step 2: Solve Ly = Pb using forward substitution
|
||
let forward_options = TriangularSolveOptions {
|
||
triangular_type: TriangularType::Lower,
|
||
transpose: false,
|
||
unit_diagonal: true, // L has 1s on diagonal (Doolittle)
|
||
validate_triangular: false,
|
||
strict_validation: false,
|
||
};
|
||
|
||
let y = self.l.solve_triangular(&pb, &forward_options)?;
|
||
|
||
// Step 3: Solve Ux = y using back substitution
|
||
let back_options = TriangularSolveOptions {
|
||
triangular_type: TriangularType::Upper,
|
||
transpose: false,
|
||
unit_diagonal: false, // U has computed values on diagonal
|
||
validate_triangular: false,
|
||
strict_validation: false,
|
||
};
|
||
|
||
let x = self.u.solve_triangular(&y, &back_options)?;
|
||
|
||
Ok(x)
|
||
}
|
||
|
||
/// Compute matrix inverse using LU decomposition
|
||
///
|
||
/// Computes A^(-1) using the LU decomposition by solving AX = I
|
||
/// where I is the identity matrix.
|
||
pub fn inverse(&self) -> Result<Tensor> {
|
||
// This will be implemented
|
||
Err(TensorError::not_implemented("LU inverse not yet implemented"))
|
||
}
|
||
}
|
||
|
||
// Helper struct for residual computation results
|
||
struct ResidualMetrics {
|
||
residual_norm: f32,
|
||
condition_number: f32,
|
||
rank: usize,
|
||
r_squared: Option<f32>,
|
||
sum_squared_residuals: f32,
|
||
}
|
||
|
||
impl Tensor {
|
||
/// Compute residuals and goodness-of-fit metrics
|
||
fn compute_residuals_and_metrics(&self, solution: &Tensor, b_original: &Tensor, a_used: &Tensor) -> Result<ResidualMetrics> {
|
||
// Compute residual r = b - A*x
|
||
let ax = self.matmul(solution)?;
|
||
let residual = b_original.sub(&ax)?;
|
||
|
||
// Compute residual norm
|
||
let residual_data = residual.to_cpu()?;
|
||
let residual_norm = (residual_data.iter().map(|&r| r * r).sum::<f32>()).sqrt();
|
||
let sum_squared_residuals = residual_data.iter().map(|&r| r * r).sum::<f32>();
|
||
|
||
// Estimate condition number using existing method
|
||
let condition_number = self.condition_number()?;
|
||
|
||
// Estimate rank using existing method
|
||
let rank = self.matrix_rank(1e-12)?;
|
||
|
||
// Compute R-squared if this is not a weighted problem
|
||
let r_squared = if b_original.ndim() == 1 {
|
||
let b_data = b_original.to_cpu()?;
|
||
let b_mean = b_data.iter().sum::<f32>() / b_data.len() as f32;
|
||
let total_sum_squares: f32 = b_data.iter().map(|&b_val| (b_val - b_mean) * (b_val - b_mean)).sum();
|
||
|
||
if total_sum_squares > 1e-15 {
|
||
Some(1.0 - sum_squared_residuals / total_sum_squares)
|
||
} else {
|
||
None
|
||
}
|
||
} else {
|
||
None // R-squared not computed for multiple RHS
|
||
};
|
||
|
||
Ok(ResidualMetrics {
|
||
residual_norm,
|
||
condition_number,
|
||
rank,
|
||
r_squared,
|
||
sum_squared_residuals,
|
||
})
|
||
}
|
||
}
|
||
|
||
// Helper functions for linear algebra computations
|
||
|
||
fn normalize_vector(v: &mut [f32]) {
|
||
let norm = vector_norm(v);
|
||
if norm > 1e-12 {
|
||
for x in v.iter_mut() {
|
||
*x /= norm;
|
||
}
|
||
}
|
||
}
|
||
|
||
fn vector_norm(v: &[f32]) -> f32 {
|
||
v.iter().map(|&x| x * x).sum::<f32>().sqrt()
|
||
}
|
||
|
||
fn matrix_vector_multiply(matrix: &[f32], vector: &[f32], result: &mut [f32], rows: usize, cols: usize) {
|
||
for i in 0..rows {
|
||
result[i] = 0.0;
|
||
for j in 0..cols {
|
||
result[i] += matrix[i * cols + j] * vector[j];
|
||
}
|
||
}
|
||
}
|
||
|
||
fn matrix_transpose_vector_multiply(matrix: &[f32], vector: &[f32], result: &mut [f32], rows: usize, cols: usize) {
|
||
for j in 0..cols {
|
||
result[j] = 0.0;
|
||
for i in 0..rows {
|
||
result[j] += matrix[i * cols + j] * vector[i];
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Validate that a matrix is triangular according to the specified type
|
||
fn validate_triangular_matrix(data: &[f32], n: usize, tri_type: TriangularType, strict: bool) -> Result<()> {
|
||
let tolerance = if strict { 0.0 } else { 1e-9 };
|
||
|
||
match tri_type {
|
||
TriangularType::Lower => {
|
||
// Check that upper triangle is zero (or within tolerance)
|
||
for i in 0..n {
|
||
for j in (i + 1)..n {
|
||
let idx = i * n + j;
|
||
if data[idx].abs() > tolerance {
|
||
return Err(TensorError::value(format!(
|
||
"Matrix is not lower triangular: element at [{},{}] = {} exceeds tolerance {}",
|
||
i, j, data[idx], tolerance
|
||
)));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
TriangularType::Upper => {
|
||
// Check that lower triangle is zero (or within tolerance)
|
||
for i in 0..n {
|
||
for j in 0..i {
|
||
let idx = i * n + j;
|
||
if data[idx].abs() > tolerance {
|
||
return Err(TensorError::value(format!(
|
||
"Matrix is not upper triangular: element at [{},{}] = {} exceeds tolerance {}",
|
||
i, j, data[idx], tolerance
|
||
)));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Forward substitution for lower triangular systems (with batch support)
|
||
/// Solves L x = b or L^T x = b (if transpose=true)
|
||
fn solve_lower_triangular_batch(
|
||
matrix: &[f32],
|
||
b: &[f32],
|
||
result: &mut [f32],
|
||
n: usize,
|
||
num_rhs: usize,
|
||
transpose: bool,
|
||
unit_diagonal: bool
|
||
) -> Result<()> {
|
||
|
||
for rhs_idx in 0..num_rhs {
|
||
if transpose {
|
||
// Solve L^T x = b (L^T is upper triangular)
|
||
// Back substitution from bottom to top
|
||
for i in (0..n).rev() {
|
||
let b_idx = if num_rhs == 1 { i } else { i * num_rhs + rhs_idx };
|
||
let result_idx = if num_rhs == 1 { i } else { i * num_rhs + rhs_idx };
|
||
|
||
let mut sum = 0.0f32;
|
||
|
||
// Sum contributions from already solved variables
|
||
for j in (i + 1)..n {
|
||
let matrix_idx = j * n + i; // L^T[i,j] = L[j,i]
|
||
let x_idx = if num_rhs == 1 { j } else { j * num_rhs + rhs_idx };
|
||
sum += matrix[matrix_idx] * result[x_idx];
|
||
}
|
||
|
||
// Solve for current variable
|
||
let diagonal_val = if unit_diagonal {
|
||
1.0
|
||
} else {
|
||
matrix[i * n + i]
|
||
};
|
||
|
||
if !unit_diagonal && diagonal_val.abs() < 1e-12 {
|
||
return Err(TensorError::numerical(format!(
|
||
"Matrix is singular: diagonal element at [{},{}] = {}",
|
||
i, i, diagonal_val
|
||
)));
|
||
}
|
||
|
||
result[result_idx] = (b[b_idx] - sum) / diagonal_val;
|
||
}
|
||
} else {
|
||
// Standard forward substitution for L x = b
|
||
for i in 0..n {
|
||
let b_idx = if num_rhs == 1 { i } else { i * num_rhs + rhs_idx };
|
||
let result_idx = if num_rhs == 1 { i } else { i * num_rhs + rhs_idx };
|
||
|
||
let mut sum = 0.0f32;
|
||
|
||
// Sum contributions from already solved variables
|
||
for j in 0..i {
|
||
let matrix_idx = i * n + j;
|
||
let x_idx = if num_rhs == 1 { j } else { j * num_rhs + rhs_idx };
|
||
sum += matrix[matrix_idx] * result[x_idx];
|
||
}
|
||
|
||
// Solve for current variable
|
||
let diagonal_val = if unit_diagonal {
|
||
1.0
|
||
} else {
|
||
matrix[i * n + i]
|
||
};
|
||
|
||
if !unit_diagonal && diagonal_val.abs() < 1e-12 {
|
||
return Err(TensorError::numerical(format!(
|
||
"Matrix is singular: diagonal element at [{},{}] = {}",
|
||
i, i, diagonal_val
|
||
)));
|
||
}
|
||
|
||
result[result_idx] = (b[b_idx] - sum) / diagonal_val;
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Back substitution for upper triangular systems (with batch support)
|
||
/// Solves U x = b or U^T x = b (if transpose=true)
|
||
fn solve_upper_triangular_batch(
|
||
matrix: &[f32],
|
||
b: &[f32],
|
||
result: &mut [f32],
|
||
n: usize,
|
||
num_rhs: usize,
|
||
transpose: bool,
|
||
unit_diagonal: bool
|
||
) -> Result<()> {
|
||
|
||
for rhs_idx in 0..num_rhs {
|
||
if transpose {
|
||
// Solve U^T x = b (U^T is lower triangular)
|
||
// Forward substitution from top to bottom
|
||
for i in 0..n {
|
||
let b_idx = if num_rhs == 1 { i } else { i * num_rhs + rhs_idx };
|
||
let result_idx = if num_rhs == 1 { i } else { i * num_rhs + rhs_idx };
|
||
|
||
let mut sum = 0.0f32;
|
||
|
||
// Sum contributions from already solved variables
|
||
for j in 0..i {
|
||
let matrix_idx = j * n + i; // U^T[i,j] = U[j,i]
|
||
let x_idx = if num_rhs == 1 { j } else { j * num_rhs + rhs_idx };
|
||
sum += matrix[matrix_idx] * result[x_idx];
|
||
}
|
||
|
||
// Solve for current variable
|
||
let diagonal_val = if unit_diagonal {
|
||
1.0
|
||
} else {
|
||
matrix[i * n + i]
|
||
};
|
||
|
||
if !unit_diagonal && diagonal_val.abs() < 1e-12 {
|
||
return Err(TensorError::numerical(format!(
|
||
"Matrix is singular: diagonal element at [{},{}] = {}",
|
||
i, i, diagonal_val
|
||
)));
|
||
}
|
||
|
||
result[result_idx] = (b[b_idx] - sum) / diagonal_val;
|
||
}
|
||
} else {
|
||
// Standard back substitution for U x = b
|
||
for i in (0..n).rev() {
|
||
let b_idx = if num_rhs == 1 { i } else { i * num_rhs + rhs_idx };
|
||
let result_idx = if num_rhs == 1 { i } else { i * num_rhs + rhs_idx };
|
||
|
||
let mut sum = 0.0f32;
|
||
|
||
// Sum contributions from already solved variables
|
||
for j in (i + 1)..n {
|
||
let matrix_idx = i * n + j;
|
||
let x_idx = if num_rhs == 1 { j } else { j * num_rhs + rhs_idx };
|
||
sum += matrix[matrix_idx] * result[x_idx];
|
||
}
|
||
|
||
// Solve for current variable
|
||
let diagonal_val = if unit_diagonal {
|
||
1.0
|
||
} else {
|
||
matrix[i * n + i]
|
||
};
|
||
|
||
if !unit_diagonal && diagonal_val.abs() < 1e-12 {
|
||
return Err(TensorError::numerical(format!(
|
||
"Matrix is singular: diagonal element at [{},{}] = {}",
|
||
i, i, diagonal_val
|
||
)));
|
||
}
|
||
|
||
result[result_idx] = (b[b_idx] - sum) / diagonal_val;
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
} |