Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,988 @@
|
||||
// Copyright (c) 2024 RustyTorch++ Team
|
||||
// Licensed under the Apache License, Version 2.0
|
||||
|
||||
//! Sparse matrix implementations for finite element assembly.
|
||||
|
||||
use crate::error::{AssemblyError, FeaResult};
|
||||
use nalgebra::DVector;
|
||||
// use rtx_tensor::Tensor; // Optional dependency
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Compressed Sparse Row (CSR) matrix format.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SparseMatrix {
|
||||
/// Number of rows
|
||||
nrows: usize,
|
||||
/// Number of columns
|
||||
ncols: usize,
|
||||
/// Row pointers
|
||||
row_ptr: Vec<usize>,
|
||||
/// Column indices
|
||||
col_idx: Vec<usize>,
|
||||
/// Non-zero values
|
||||
values: Vec<f64>,
|
||||
/// Assembly map for efficient insertion during assembly
|
||||
assembly_map: HashMap<(usize, usize), usize>,
|
||||
/// Whether the matrix has been finalized
|
||||
finalized: bool,
|
||||
}
|
||||
|
||||
impl SparseMatrix {
|
||||
/// Create a new sparse matrix.
|
||||
pub fn new(nrows: usize, ncols: usize) -> Self {
|
||||
Self {
|
||||
nrows,
|
||||
ncols,
|
||||
row_ptr: vec![0; nrows + 1],
|
||||
col_idx: Vec::new(),
|
||||
values: Vec::new(),
|
||||
assembly_map: HashMap::new(),
|
||||
finalized: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create from triplet format (row, col, value).
|
||||
pub fn from_triplets(
|
||||
nrows: usize,
|
||||
ncols: usize,
|
||||
triplets: &[(usize, usize, f64)],
|
||||
) -> FeaResult<Self> {
|
||||
let mut matrix = Self::new(nrows, ncols);
|
||||
|
||||
for &(row, col, value) in triplets {
|
||||
matrix.add_entry(row, col, value)?;
|
||||
}
|
||||
|
||||
matrix.finalize()?;
|
||||
Ok(matrix)
|
||||
}
|
||||
|
||||
/// Get number of rows.
|
||||
pub fn nrows(&self) -> usize {
|
||||
self.nrows
|
||||
}
|
||||
|
||||
/// Get number of columns.
|
||||
pub fn ncols(&self) -> usize {
|
||||
self.ncols
|
||||
}
|
||||
|
||||
/// Get number of non-zeros.
|
||||
pub fn nnz(&self) -> usize {
|
||||
self.values.len()
|
||||
}
|
||||
|
||||
/// Get matrix density.
|
||||
pub fn density(&self) -> f64 {
|
||||
if self.nrows == 0 || self.ncols == 0 {
|
||||
0.0
|
||||
} else {
|
||||
self.nnz() as f64 / (self.nrows * self.ncols) as f64
|
||||
}
|
||||
}
|
||||
|
||||
/// Get matrix bandwidth.
|
||||
pub fn bandwidth(&self) -> usize {
|
||||
if !self.finalized {
|
||||
// For unfinalized matrix, compute from assembly_map
|
||||
let mut max_bandwidth = 0;
|
||||
for (row, col) in self.assembly_map.keys() {
|
||||
let distance = row.abs_diff(*col);
|
||||
max_bandwidth = max_bandwidth.max(distance);
|
||||
}
|
||||
max_bandwidth
|
||||
} else {
|
||||
// For finalized matrix, compute from CSR structure
|
||||
let mut max_bandwidth = 0;
|
||||
for row in 0..self.nrows {
|
||||
let start = self.row_ptr[row];
|
||||
let end = self.row_ptr[row + 1];
|
||||
for idx in start..end {
|
||||
let col = self.col_idx[idx];
|
||||
let distance = row.abs_diff(col);
|
||||
max_bandwidth = max_bandwidth.max(distance);
|
||||
}
|
||||
}
|
||||
max_bandwidth
|
||||
}
|
||||
}
|
||||
|
||||
/// Add entry to the matrix (accumulates if entry exists).
|
||||
pub fn add_entry(&mut self, row: usize, col: usize, value: f64) -> FeaResult<()> {
|
||||
if self.finalized {
|
||||
return Err(AssemblyError::SparseAssemblyFailed {
|
||||
reason: "Cannot add entries to finalized matrix".to_string(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
if row >= self.nrows || col >= self.ncols {
|
||||
return Err(AssemblyError::MatrixDimensionMismatch {
|
||||
expected_rows: self.nrows,
|
||||
expected_cols: self.ncols,
|
||||
actual_rows: row + 1,
|
||||
actual_cols: col + 1,
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
if value.abs() < 1e-15 {
|
||||
return Ok(()); // Skip near-zero entries
|
||||
}
|
||||
|
||||
let key = (row, col);
|
||||
if let Some(&idx) = self.assembly_map.get(&key) {
|
||||
// Entry exists, accumulate
|
||||
self.values[idx] += value;
|
||||
} else {
|
||||
// New entry
|
||||
let idx = self.values.len();
|
||||
self.values.push(value);
|
||||
self.col_idx.push(col);
|
||||
self.assembly_map.insert(key, idx);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Set entry in the matrix (overwrites if entry exists).
|
||||
pub fn set_entry(&mut self, row: usize, col: usize, value: f64) -> FeaResult<()> {
|
||||
if self.finalized {
|
||||
return Err(AssemblyError::SparseAssemblyFailed {
|
||||
reason: "Cannot set entries in finalized matrix".to_string(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
let key = (row, col);
|
||||
if let Some(&idx) = self.assembly_map.get(&key) {
|
||||
// Entry exists, overwrite
|
||||
self.values[idx] = value;
|
||||
} else if value.abs() >= 1e-15 {
|
||||
// New non-zero entry
|
||||
let idx = self.values.len();
|
||||
self.values.push(value);
|
||||
self.col_idx.push(col);
|
||||
self.assembly_map.insert(key, idx);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Apply Dirichlet boundary condition.
|
||||
pub fn set_dirichlet(&mut self, dof: usize, _value: f64) -> FeaResult<()> {
|
||||
if dof >= self.nrows {
|
||||
return Err(AssemblyError::GlobalDofOutOfBounds {
|
||||
index: dof,
|
||||
max_index: self.nrows.saturating_sub(1),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
// Zero out row and column
|
||||
for col in 0..self.ncols {
|
||||
if col != dof {
|
||||
self.set_entry(dof, col, 0.0)?;
|
||||
}
|
||||
}
|
||||
|
||||
for row in 0..self.nrows {
|
||||
if row != dof {
|
||||
self.set_entry(row, dof, 0.0)?;
|
||||
}
|
||||
}
|
||||
|
||||
// Set diagonal entry to 1
|
||||
self.set_entry(dof, dof, 1.0)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Finalize the matrix (convert to CSR format).
|
||||
pub fn finalize(&mut self) -> FeaResult<()> {
|
||||
if self.finalized {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Group entries by row
|
||||
let mut row_entries: Vec<Vec<(usize, f64)>> = vec![Vec::new(); self.nrows];
|
||||
|
||||
for (&(row, col), &idx) in &self.assembly_map {
|
||||
let value = self.values[idx];
|
||||
if value.abs() >= 1e-15 {
|
||||
row_entries[row].push((col, value));
|
||||
}
|
||||
}
|
||||
|
||||
// Sort columns within each row
|
||||
for row_data in &mut row_entries {
|
||||
row_data.sort_by_key(|&(col, _)| col);
|
||||
}
|
||||
|
||||
// Build CSR format
|
||||
self.values.clear();
|
||||
self.col_idx.clear();
|
||||
self.row_ptr = vec![0; self.nrows + 1];
|
||||
|
||||
let mut nnz = 0;
|
||||
for (row, row_data) in row_entries.iter().enumerate() {
|
||||
self.row_ptr[row] = nnz;
|
||||
for &(col, value) in row_data {
|
||||
self.col_idx.push(col);
|
||||
self.values.push(value);
|
||||
nnz += 1;
|
||||
}
|
||||
}
|
||||
self.row_ptr[self.nrows] = nnz;
|
||||
|
||||
// Clear assembly map to save memory
|
||||
self.assembly_map.clear();
|
||||
self.finalized = true;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get entry value.
|
||||
pub fn get_entry(&self, row: usize, col: usize) -> f64 {
|
||||
if !self.finalized {
|
||||
// Search in assembly map
|
||||
self.assembly_map
|
||||
.get(&(row, col))
|
||||
.map_or(0.0, |&idx| self.values[idx])
|
||||
} else {
|
||||
// Search in CSR format
|
||||
let start = self.row_ptr[row];
|
||||
let end = self.row_ptr[row + 1];
|
||||
|
||||
for i in start..end {
|
||||
if self.col_idx[i] == col {
|
||||
return self.values[i];
|
||||
}
|
||||
}
|
||||
0.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Matrix-vector multiplication: y = A * x.
|
||||
pub fn multiply_vector(&self, x: &DVector<f64>) -> FeaResult<DVector<f64>> {
|
||||
if !self.finalized {
|
||||
return Err(AssemblyError::SparseAssemblyFailed {
|
||||
reason: "Matrix must be finalized before multiplication".to_string(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
if x.len() != self.ncols {
|
||||
return Err(AssemblyError::MatrixDimensionMismatch {
|
||||
expected_rows: self.ncols,
|
||||
expected_cols: 1,
|
||||
actual_rows: x.len(),
|
||||
actual_cols: 1,
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
let mut y = DVector::zeros(self.nrows);
|
||||
|
||||
for row in 0..self.nrows {
|
||||
let start = self.row_ptr[row];
|
||||
let end = self.row_ptr[row + 1];
|
||||
|
||||
let mut sum = 0.0;
|
||||
for i in start..end {
|
||||
sum += self.values[i] * x[self.col_idx[i]];
|
||||
}
|
||||
y[row] = sum;
|
||||
}
|
||||
|
||||
Ok(y)
|
||||
}
|
||||
|
||||
/// Extract submatrix for given row and column indices.
|
||||
pub fn extract_submatrix(
|
||||
&self,
|
||||
row_indices: &[usize],
|
||||
col_indices: &[usize],
|
||||
) -> FeaResult<Self> {
|
||||
if !self.finalized {
|
||||
return Err(AssemblyError::SparseAssemblyFailed {
|
||||
reason: "Matrix must be finalized before extraction".to_string(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
let sub_nrows = row_indices.len();
|
||||
let sub_ncols = col_indices.len();
|
||||
let mut sub_matrix = Self::new(sub_nrows, sub_ncols);
|
||||
|
||||
// Create column index mapping
|
||||
let mut col_map = HashMap::new();
|
||||
for (new_col, &old_col) in col_indices.iter().enumerate() {
|
||||
col_map.insert(old_col, new_col);
|
||||
}
|
||||
|
||||
for (new_row, &old_row) in row_indices.iter().enumerate() {
|
||||
let start = self.row_ptr[old_row];
|
||||
let end = self.row_ptr[old_row + 1];
|
||||
|
||||
for i in start..end {
|
||||
let old_col = self.col_idx[i];
|
||||
if let Some(&new_col) = col_map.get(&old_col) {
|
||||
let value = self.values[i];
|
||||
sub_matrix.add_entry(new_row, new_col, value)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sub_matrix.finalize()?;
|
||||
Ok(sub_matrix)
|
||||
}
|
||||
|
||||
/// Convert to dense matrix (for debugging/small matrices).
|
||||
pub fn to_dense(&self) -> nalgebra::DMatrix<f64> {
|
||||
let mut dense = nalgebra::DMatrix::zeros(self.nrows, self.ncols);
|
||||
|
||||
if self.finalized {
|
||||
for row in 0..self.nrows {
|
||||
let start = self.row_ptr[row];
|
||||
let end = self.row_ptr[row + 1];
|
||||
|
||||
for i in start..end {
|
||||
dense[(row, self.col_idx[i])] = self.values[i];
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (&(row, col), &idx) in &self.assembly_map {
|
||||
dense[(row, col)] = self.values[idx];
|
||||
}
|
||||
}
|
||||
|
||||
dense
|
||||
}
|
||||
|
||||
/// Get matrix structure (row pointers and column indices).
|
||||
pub fn structure(&self) -> (&[usize], &[usize]) {
|
||||
(&self.row_ptr, &self.col_idx)
|
||||
}
|
||||
|
||||
/// Get matrix values.
|
||||
pub fn values(&self) -> &[f64] {
|
||||
&self.values
|
||||
}
|
||||
|
||||
/// Check if matrix is symmetric (structure-wise).
|
||||
pub fn is_symmetric_structure(&self) -> bool {
|
||||
if !self.finalized || self.nrows != self.ncols {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check if (i,j) exists whenever (j,i) exists
|
||||
let mut has_entry = vec![vec![false; self.ncols]; self.nrows];
|
||||
|
||||
for row in 0..self.nrows {
|
||||
let start = self.row_ptr[row];
|
||||
let end = self.row_ptr[row + 1];
|
||||
|
||||
for i in start..end {
|
||||
has_entry[row][self.col_idx[i]] = true;
|
||||
}
|
||||
}
|
||||
|
||||
for row in 0..self.nrows {
|
||||
for col in 0..self.ncols {
|
||||
if has_entry[row][col] != has_entry[col][row] {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// Check if matrix is symmetric (both structure and values).
|
||||
pub fn is_symmetric(&self) -> bool {
|
||||
if !self.finalized || self.nrows != self.ncols {
|
||||
return false;
|
||||
}
|
||||
|
||||
const EPSILON: f64 = 1e-10;
|
||||
|
||||
for row in 0..self.nrows {
|
||||
for col in row + 1..self.ncols {
|
||||
let val_ij = self.get_entry(row, col);
|
||||
let val_ji = self.get_entry(col, row);
|
||||
|
||||
if (val_ij - val_ji).abs() > EPSILON {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// Compute matrix norm (Frobenius norm).
|
||||
pub fn frobenius_norm(&self) -> f64 {
|
||||
self.values.iter().map(|&x| x * x).sum::<f64>().sqrt()
|
||||
}
|
||||
|
||||
/// Multiply transpose of matrix with vector: A^T * x
|
||||
pub fn transpose_multiply_vector(&self, x: &DVector<f64>) -> FeaResult<DVector<f64>> {
|
||||
if x.len() != self.nrows {
|
||||
return Err(AssemblyError::DimensionMismatch {
|
||||
expected: self.nrows,
|
||||
actual: x.len(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
let mut result = DVector::zeros(self.ncols);
|
||||
|
||||
for row in 0..self.nrows {
|
||||
for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
|
||||
let col = self.col_idx[idx];
|
||||
let val = self.values[idx];
|
||||
result[col] += val * x[row];
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Extract diagonal elements
|
||||
pub fn diagonal(&self) -> FeaResult<DVector<f64>> {
|
||||
let n = self.nrows.min(self.ncols);
|
||||
let mut diag = DVector::zeros(n);
|
||||
|
||||
for row in 0..n {
|
||||
for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
|
||||
if self.col_idx[idx] == row {
|
||||
diag[row] = self.values[idx];
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(diag)
|
||||
}
|
||||
|
||||
/// Forward substitution for lower triangular system L * x = b
|
||||
pub fn forward_solve(&self, b: &DVector<f64>) -> FeaResult<DVector<f64>> {
|
||||
if b.len() != self.nrows {
|
||||
return Err(AssemblyError::DimensionMismatch {
|
||||
expected: self.nrows,
|
||||
actual: b.len(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
let mut x = DVector::zeros(self.nrows);
|
||||
|
||||
for row in 0..self.nrows {
|
||||
let mut sum = b[row];
|
||||
let mut diag_val = 0.0;
|
||||
|
||||
for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
|
||||
let col = self.col_idx[idx];
|
||||
let val = self.values[idx];
|
||||
|
||||
if col < row {
|
||||
sum -= val * x[col];
|
||||
} else if col == row {
|
||||
diag_val = val;
|
||||
}
|
||||
}
|
||||
|
||||
if diag_val.abs() < 1e-12 {
|
||||
return Err(AssemblyError::SingularMatrix.into());
|
||||
}
|
||||
|
||||
x[row] = sum / diag_val;
|
||||
}
|
||||
|
||||
Ok(x)
|
||||
}
|
||||
|
||||
/// Backward substitution for upper triangular system U * x = b
|
||||
pub fn backward_solve(&self, b: &DVector<f64>) -> FeaResult<DVector<f64>> {
|
||||
if b.len() != self.nrows {
|
||||
return Err(AssemblyError::DimensionMismatch {
|
||||
expected: self.nrows,
|
||||
actual: b.len(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
let mut x = DVector::zeros(self.nrows);
|
||||
|
||||
for row in (0..self.nrows).rev() {
|
||||
let mut sum = b[row];
|
||||
let mut diag_val = 0.0;
|
||||
|
||||
for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
|
||||
let col = self.col_idx[idx];
|
||||
let val = self.values[idx];
|
||||
|
||||
if col > row {
|
||||
sum -= val * x[col];
|
||||
} else if col == row {
|
||||
diag_val = val;
|
||||
}
|
||||
}
|
||||
|
||||
if diag_val.abs() < 1e-12 {
|
||||
return Err(AssemblyError::SingularMatrix.into());
|
||||
}
|
||||
|
||||
x[row] = sum / diag_val;
|
||||
}
|
||||
|
||||
Ok(x)
|
||||
}
|
||||
|
||||
/// Solve linear system A * x = b using LU decomposition
|
||||
pub fn solve_vector(&self, b: &DVector<f64>) -> FeaResult<DVector<f64>> {
|
||||
// For now, use a simple iterative method (Conjugate Gradient)
|
||||
// In production, this should use a proper sparse solver
|
||||
if b.len() != self.nrows || self.nrows != self.ncols {
|
||||
return Err(AssemblyError::DimensionMismatch {
|
||||
expected: self.nrows,
|
||||
actual: b.len(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
// Simple Conjugate Gradient implementation for symmetric positive definite matrices
|
||||
let mut x = DVector::zeros(self.nrows);
|
||||
let mut r = b.clone();
|
||||
let mut p = r.clone();
|
||||
let tolerance = 1e-9;
|
||||
let max_iter = self.nrows * 2;
|
||||
|
||||
for _iter in 0..max_iter {
|
||||
let ap = self.multiply_vector(&p)?;
|
||||
let r_dot_r = r.dot(&r);
|
||||
|
||||
if r_dot_r < tolerance * tolerance {
|
||||
break;
|
||||
}
|
||||
|
||||
let alpha = r_dot_r / p.dot(&ap);
|
||||
x += alpha * &p;
|
||||
r -= alpha * ≈
|
||||
|
||||
let r_dot_r_new = r.dot(&r);
|
||||
let beta = r_dot_r_new / r_dot_r;
|
||||
p = &r + beta * p;
|
||||
}
|
||||
|
||||
Ok(x)
|
||||
}
|
||||
}
|
||||
|
||||
// Also implement Add for &SparseMatrix + &SparseMatrix for convenience
|
||||
impl std::ops::Add<&SparseMatrix> for &SparseMatrix {
|
||||
type Output = SparseMatrix;
|
||||
|
||||
fn add(self, rhs: &SparseMatrix) -> Self::Output {
|
||||
assert!(!(self.nrows != rhs.nrows || self.ncols != rhs.ncols), "Matrix dimensions must match for addition");
|
||||
|
||||
let mut result = SparseMatrix::new(self.nrows, self.ncols);
|
||||
|
||||
// Add entries from left matrix
|
||||
if self.finalized {
|
||||
// Use CSR format for finalized matrix
|
||||
for row in 0..self.nrows {
|
||||
for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
|
||||
let col = self.col_idx[idx];
|
||||
let value = self.values[idx];
|
||||
result.add_entry(row, col, value).unwrap();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Use assembly_map for unfinalized matrix
|
||||
for ((row, col), &idx) in &self.assembly_map {
|
||||
let value = self.values[idx];
|
||||
result.add_entry(*row, *col, value).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// Add entries from right matrix
|
||||
if rhs.finalized {
|
||||
// Use CSR format for finalized matrix
|
||||
for row in 0..rhs.nrows {
|
||||
for idx in rhs.row_ptr[row]..rhs.row_ptr[row + 1] {
|
||||
let col = rhs.col_idx[idx];
|
||||
let value = rhs.values[idx];
|
||||
result.add_entry(row, col, value).unwrap();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Use assembly_map for unfinalized matrix
|
||||
for ((row, col), &idx) in &rhs.assembly_map {
|
||||
let value = rhs.values[idx];
|
||||
result.add_entry(*row, *col, value).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// If both originals were finalized, finalize the result
|
||||
if self.finalized && rhs.finalized {
|
||||
result.finalize().unwrap();
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
/// GPU-accelerated sparse matrix operations.
|
||||
#[cfg(feature = "cuda")]
|
||||
use cudarc::driver::safe::{CudaContext as CudaDevice, CudaStream};
|
||||
|
||||
pub struct GpuSparseMatrix {
|
||||
/// CPU sparse matrix (for initial data only)
|
||||
cpu_matrix: SparseMatrix,
|
||||
/// CUDA device - REQUIRED
|
||||
#[cfg(feature = "cuda")]
|
||||
device: std::sync::Arc<CudaDevice>,
|
||||
/// CUDA stream for async operations
|
||||
#[cfg(feature = "cuda")]
|
||||
stream: std::sync::Arc<CudaStream>,
|
||||
}
|
||||
|
||||
impl GpuSparseMatrix {
|
||||
/// Create from CPU sparse matrix.
|
||||
/// PANICS if GPU is not available (GPU-only requirement)
|
||||
pub fn from_cpu_matrix(cpu_matrix: SparseMatrix) -> FeaResult<Self> {
|
||||
#[cfg(feature = "cuda")]
|
||||
{
|
||||
let device =
|
||||
CudaDevice::new(0).expect("GPU is REQUIRED for GpuSparseMatrix. No GPU found.");
|
||||
let stream = device.default_stream();
|
||||
|
||||
Ok(Self {
|
||||
cpu_matrix,
|
||||
device,
|
||||
stream,
|
||||
})
|
||||
}
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
{
|
||||
panic!("GPU is REQUIRED. Compile with 'cuda' feature enabled.")
|
||||
}
|
||||
}
|
||||
|
||||
/// Transfer matrix to GPU.
|
||||
pub fn to_gpu(&mut self) -> FeaResult<()> {
|
||||
if !self.cpu_matrix.finalized {
|
||||
return Err(AssemblyError::SparseAssemblyFailed {
|
||||
reason: "Matrix must be finalized before GPU transfer".to_string(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
{
|
||||
// GPU is always available here - we panic in from_cpu_matrix if not
|
||||
// Allocate device memory for CSR format
|
||||
let nnz = self.cpu_matrix.values.len();
|
||||
let nrows = self.cpu_matrix.nrows;
|
||||
|
||||
// TODO: Implement actual GPU transfer using cuSPARSE
|
||||
// This will use the device and stream fields
|
||||
tracing::info!(
|
||||
"Matrix transfer to GPU pending implementation (nnz: {}, rows: {})",
|
||||
nnz,
|
||||
nrows
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// GPU matrix-vector multiplication.
|
||||
pub fn gpu_multiply_vector(&self, _x: &DVector<f64>) -> FeaResult<DVector<f64>> {
|
||||
#[cfg(feature = "cuda")]
|
||||
{
|
||||
// Perform GPU SpMV using cuSPARSE
|
||||
// This would use cusparseSpMV() with optimal buffer allocation
|
||||
// and handle device-to-host memory transfers efficiently
|
||||
|
||||
// Steps for full implementation:
|
||||
// 1. Allocate device vectors for input/output
|
||||
// 2. Copy input vector to device
|
||||
// 3. Execute cusparseSpMV operation
|
||||
// 4. Copy result back to host
|
||||
|
||||
// TODO: Implement actual GPU SpMV
|
||||
// For now, return error to maintain GPU-only requirement
|
||||
Err(AssemblyError::SparseAssemblyFailed {
|
||||
reason: "GPU SpMV implementation pending".to_string(),
|
||||
}
|
||||
.into())
|
||||
}
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
{
|
||||
panic!("GPU is REQUIRED for matrix operations.")
|
||||
}
|
||||
}
|
||||
|
||||
/// Get CPU matrix reference.
|
||||
pub fn cpu_matrix(&self) -> &SparseMatrix {
|
||||
&self.cpu_matrix
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_sparse_matrix_creation() {
|
||||
let matrix = SparseMatrix::new(3, 3);
|
||||
assert_eq!(matrix.nrows(), 3);
|
||||
assert_eq!(matrix.ncols(), 3);
|
||||
assert_eq!(matrix.nnz(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sparse_matrix_assembly() {
|
||||
let mut matrix = SparseMatrix::new(3, 3);
|
||||
|
||||
matrix.add_entry(0, 0, 1.0).unwrap();
|
||||
matrix.add_entry(1, 1, 2.0).unwrap();
|
||||
matrix.add_entry(2, 2, 3.0).unwrap();
|
||||
matrix.add_entry(0, 1, 0.5).unwrap();
|
||||
|
||||
assert_eq!(matrix.nnz(), 4);
|
||||
|
||||
// Add to existing entry
|
||||
matrix.add_entry(0, 0, 1.0).unwrap();
|
||||
assert_eq!(matrix.get_entry(0, 0), 2.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sparse_matrix_finalization() {
|
||||
let mut matrix = SparseMatrix::new(3, 3);
|
||||
|
||||
matrix.add_entry(0, 0, 1.0).unwrap();
|
||||
matrix.add_entry(1, 1, 2.0).unwrap();
|
||||
matrix.add_entry(2, 2, 3.0).unwrap();
|
||||
|
||||
matrix.finalize().unwrap();
|
||||
|
||||
assert_eq!(matrix.get_entry(0, 0), 1.0);
|
||||
assert_eq!(matrix.get_entry(1, 1), 2.0);
|
||||
assert_eq!(matrix.get_entry(2, 2), 3.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_matrix_vector_multiplication() {
|
||||
let mut matrix = SparseMatrix::new(3, 3);
|
||||
|
||||
matrix.add_entry(0, 0, 2.0).unwrap();
|
||||
matrix.add_entry(1, 1, 3.0).unwrap();
|
||||
matrix.add_entry(2, 2, 4.0).unwrap();
|
||||
matrix.finalize().unwrap();
|
||||
|
||||
let x = DVector::from_vec(vec![1.0, 2.0, 3.0]);
|
||||
let y = matrix.multiply_vector(&x).unwrap();
|
||||
|
||||
assert_eq!(y[0], 2.0);
|
||||
assert_eq!(y[1], 6.0);
|
||||
assert_eq!(y[2], 12.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_submatrix_extraction() {
|
||||
let mut matrix = SparseMatrix::new(4, 4);
|
||||
|
||||
// Create a 4x4 diagonal matrix
|
||||
for i in 0..4 {
|
||||
matrix.add_entry(i, i, (i + 1) as f64).unwrap();
|
||||
}
|
||||
matrix.finalize().unwrap();
|
||||
|
||||
// Extract 2x2 submatrix
|
||||
let row_indices = vec![0, 2];
|
||||
let col_indices = vec![0, 2];
|
||||
let sub_matrix = matrix
|
||||
.extract_submatrix(&row_indices, &col_indices)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(sub_matrix.nrows(), 2);
|
||||
assert_eq!(sub_matrix.ncols(), 2);
|
||||
assert_eq!(sub_matrix.get_entry(0, 0), 1.0);
|
||||
assert_eq!(sub_matrix.get_entry(1, 1), 3.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dirichlet_boundary_condition() {
|
||||
let mut matrix = SparseMatrix::new(3, 3);
|
||||
|
||||
// Fill matrix
|
||||
for i in 0..3 {
|
||||
for j in 0..3 {
|
||||
matrix.add_entry(i, j, ((i + 1) * (j + 1)) as f64).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// Apply Dirichlet BC at DOF 1
|
||||
matrix.set_dirichlet(1, 5.0).unwrap();
|
||||
|
||||
// Check that row 1 and column 1 are zeroed except diagonal
|
||||
assert_eq!(matrix.get_entry(1, 0), 0.0);
|
||||
assert_eq!(matrix.get_entry(1, 1), 1.0);
|
||||
assert_eq!(matrix.get_entry(1, 2), 0.0);
|
||||
assert_eq!(matrix.get_entry(0, 1), 0.0);
|
||||
assert_eq!(matrix.get_entry(2, 1), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dense_conversion() {
|
||||
let mut matrix = SparseMatrix::new(2, 2);
|
||||
matrix.add_entry(0, 0, 1.0).unwrap();
|
||||
matrix.add_entry(0, 1, 2.0).unwrap();
|
||||
matrix.add_entry(1, 0, 3.0).unwrap();
|
||||
matrix.add_entry(1, 1, 4.0).unwrap();
|
||||
|
||||
let dense = matrix.to_dense();
|
||||
assert_eq!(dense[(0, 0)], 1.0);
|
||||
assert_eq!(dense[(0, 1)], 2.0);
|
||||
assert_eq!(dense[(1, 0)], 3.0);
|
||||
assert_eq!(dense[(1, 1)], 4.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_triplet_construction() {
|
||||
let triplets = vec![(0, 0, 1.0), (1, 1, 2.0), (0, 1, 0.5)];
|
||||
|
||||
let matrix = SparseMatrix::from_triplets(2, 2, &triplets).unwrap();
|
||||
assert_eq!(matrix.get_entry(0, 0), 1.0);
|
||||
assert_eq!(matrix.get_entry(1, 1), 2.0);
|
||||
assert_eq!(matrix.get_entry(0, 1), 0.5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_matrix_properties() {
|
||||
let mut matrix = SparseMatrix::new(3, 3);
|
||||
matrix.add_entry(0, 0, 1.0).unwrap();
|
||||
matrix.add_entry(1, 1, 4.0).unwrap();
|
||||
matrix.add_entry(2, 2, 9.0).unwrap();
|
||||
matrix.finalize().unwrap();
|
||||
|
||||
assert_eq!(matrix.nnz(), 3);
|
||||
assert!((matrix.density() - 3.0 / 9.0).abs() < 1e-12);
|
||||
assert!((matrix.frobenius_norm() - (1.0 + 16.0 + 81.0_f64).sqrt()).abs() < 1e-12);
|
||||
}
|
||||
}
|
||||
|
||||
// Implement scalar multiplication for SparseMatrix
|
||||
impl std::ops::Mul<f64> for SparseMatrix {
|
||||
type Output = Self;
|
||||
|
||||
fn mul(self, scalar: f64) -> Self::Output {
|
||||
let mut result = Self::new(self.nrows, self.ncols);
|
||||
|
||||
if self.finalized {
|
||||
// For finalized matrices, use CSR format data
|
||||
for row in 0..self.nrows {
|
||||
for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
|
||||
let col = self.col_idx[idx];
|
||||
let value = self.values[idx] * scalar;
|
||||
result.add_entry(row, col, value).unwrap();
|
||||
}
|
||||
}
|
||||
result.finalize().unwrap();
|
||||
} else {
|
||||
// For unfinalized matrices, use assembly_map
|
||||
for ((row, col), &idx) in &self.assembly_map {
|
||||
let value = self.values[idx] * scalar;
|
||||
result.add_entry(*row, *col, value).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
// Also implement for references to avoid unnecessary cloning
|
||||
impl std::ops::Mul<f64> for &SparseMatrix {
|
||||
type Output = SparseMatrix;
|
||||
|
||||
fn mul(self, scalar: f64) -> Self::Output {
|
||||
let mut result = SparseMatrix::new(self.nrows, self.ncols);
|
||||
|
||||
if self.finalized {
|
||||
// For finalized matrices, use CSR format data
|
||||
for row in 0..self.nrows {
|
||||
for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
|
||||
let col = self.col_idx[idx];
|
||||
let value = self.values[idx] * scalar;
|
||||
result.add_entry(row, col, value).unwrap();
|
||||
}
|
||||
}
|
||||
result.finalize().unwrap();
|
||||
} else {
|
||||
// For unfinalized matrices, use assembly_map
|
||||
for ((row, col), &idx) in &self.assembly_map {
|
||||
let value = self.values[idx] * scalar;
|
||||
result.add_entry(*row, *col, value).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
// Implement matrix addition for SparseMatrix
|
||||
impl std::ops::Add<&Self> for SparseMatrix {
|
||||
type Output = Self;
|
||||
|
||||
fn add(self, rhs: &Self) -> Self::Output {
|
||||
assert!(!(self.nrows != rhs.nrows || self.ncols != rhs.ncols), "Matrix dimensions must match for addition");
|
||||
|
||||
let mut result = Self::new(self.nrows, self.ncols);
|
||||
|
||||
// Add entries from left matrix
|
||||
if self.finalized {
|
||||
// Use CSR format for finalized matrix
|
||||
for row in 0..self.nrows {
|
||||
for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
|
||||
let col = self.col_idx[idx];
|
||||
let value = self.values[idx];
|
||||
result.add_entry(row, col, value).unwrap();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Use assembly_map for unfinalized matrix
|
||||
for ((row, col), &idx) in &self.assembly_map {
|
||||
let value = self.values[idx];
|
||||
result.add_entry(*row, *col, value).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// Add entries from right matrix
|
||||
if rhs.finalized {
|
||||
// Use CSR format for finalized matrix
|
||||
for row in 0..rhs.nrows {
|
||||
for idx in rhs.row_ptr[row]..rhs.row_ptr[row + 1] {
|
||||
let col = rhs.col_idx[idx];
|
||||
let value = rhs.values[idx];
|
||||
result.add_entry(row, col, value).unwrap();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Use assembly_map for unfinalized matrix
|
||||
for ((row, col), &idx) in &rhs.assembly_map {
|
||||
let value = rhs.values[idx];
|
||||
result.add_entry(*row, *col, value).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// If both originals were finalized, finalize the result
|
||||
if self.finalized && rhs.finalized {
|
||||
result.finalize().unwrap();
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user