Files
rustytorch/archive/legacy_files_backup/mod_original_backup.rs
T
2026-03-04 00:08:42 +00:00

2527 lines
88 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! 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(&regularization)?;
}
// 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(())
}