use std::fmt::Debug;
use std::ops::{Add, Mul, Sub};
use super::extract_block::CscBlock;
use super::sparsity::MatrixSparsityRef;
use super::utils::*;
use super::{Matrix, MatrixCommon, MatrixSparsity};
use crate::error::{DiffsolError, MatrixError};
use crate::{DefaultSolver, FaerScalar, FaerSparseLU, IndexType, Scalar, Scale};
use crate::{FaerContext, FaerVec, FaerVecIndex, Vector, VectorIndex};
use faer::reborrow::{Reborrow, ReborrowMut};
use faer::sparse::ops::{ternary_op_assign_into, union_symbolic};
use faer::sparse::{Pair, SparseColMat, SymbolicSparseColMat, SymbolicSparseColMatRef, Triplet};
#[derive(Clone, Debug)]
pub struct FaerSparseMat<T: FaerScalar> {
pub(crate) data: SparseColMat<IndexType, T>,
pub(crate) context: FaerContext,
}
impl<T: FaerScalar> DefaultSolver for FaerSparseMat<T> {
type LS = FaerSparseLU<T>;
}
impl_matrix_common!(
FaerSparseMat<T>,
FaerVec<T>,
FaerContext,
SparseColMat<IndexType, T>,
FaerScalar
);
macro_rules! impl_mul_scalar {
($mat_type:ty, $out:ty) => {
impl<'a, T: FaerScalar> Mul<Scale<T>> for $mat_type {
type Output = $out;
fn mul(self, rhs: Scale<T>) -> Self::Output {
let scale: faer::Scale<T> = rhs.into();
Self::Output {
data: &self.data * scale,
context: self.context,
}
}
}
};
}
impl_mul_scalar!(FaerSparseMat<T>, FaerSparseMat<T>);
impl_mul_scalar!(&FaerSparseMat<T>, FaerSparseMat<T>);
impl_add!(
FaerSparseMat<T>,
&FaerSparseMat<T>,
FaerSparseMat<T>,
FaerScalar
);
impl_sub!(
FaerSparseMat<T>,
&FaerSparseMat<T>,
FaerSparseMat<T>,
FaerScalar
);
impl<T: FaerScalar> MatrixSparsity<FaerSparseMat<T>> for SymbolicSparseColMat<IndexType> {
fn union(
self,
other: SymbolicSparseColMatRef<IndexType>,
) -> Result<SymbolicSparseColMat<IndexType>, DiffsolError> {
union_symbolic(self.rb(), other).map_err(|e| DiffsolError::Other(e.to_string()))
}
fn as_ref(&self) -> SymbolicSparseColMatRef<'_, IndexType> {
self.rb()
}
fn nrows(&self) -> IndexType {
self.nrows()
}
fn ncols(&self) -> IndexType {
self.ncols()
}
fn is_sparse() -> bool {
true
}
fn indices(&self) -> Vec<(IndexType, IndexType)> {
let mut indices = Vec::with_capacity(self.compute_nnz());
for col_i in 0..self.ncols() {
for row_j in self.col_range(col_i) {
indices.push((row_j, col_i));
}
}
indices
}
fn new_diagonal(n: IndexType) -> Self {
let indices = (0..n).map(|i| Pair::new(i, i)).collect::<Vec<_>>();
SymbolicSparseColMat::try_new_from_indices(n, n, indices.as_slice())
.unwrap()
.0
}
fn try_from_indices(
nrows: IndexType,
ncols: IndexType,
indices: Vec<(IndexType, IndexType)>,
) -> Result<Self, DiffsolError> {
let indices = indices
.iter()
.map(|(i, j)| Pair::new(*i, *j))
.collect::<Vec<_>>();
match Self::try_new_from_indices(nrows, ncols, indices.as_slice()) {
Ok((sparsity, _)) => Ok(sparsity),
Err(e) => Err(DiffsolError::Other(e.to_string())),
}
}
fn get_index(
&self,
indices: &[(IndexType, IndexType)],
ctx: FaerContext,
) -> <<FaerSparseMat<T> as MatrixCommon>::V as Vector>::Index {
let col_ptrs = self.col_ptr();
let row_indices = self.row_idx();
let mut ret = Vec::with_capacity(indices.len());
for &(i, j) in indices.iter() {
let col_ptr = col_ptrs[j];
let next_col_ptr = col_ptrs[j + 1];
for (ii, &ri) in row_indices
.iter()
.enumerate()
.take(next_col_ptr)
.skip(col_ptr)
{
if ri == i {
ret.push(ii);
break;
}
}
}
FaerVecIndex {
data: ret,
context: ctx,
}
}
}
impl<'a, T: FaerScalar> MatrixSparsityRef<'a, FaerSparseMat<T>>
for SymbolicSparseColMatRef<'a, IndexType>
{
fn to_owned(&self) -> SymbolicSparseColMat<IndexType> {
self.to_owned().unwrap()
}
fn nrows(&self) -> IndexType {
self.nrows()
}
fn ncols(&self) -> IndexType {
self.ncols()
}
fn is_sparse() -> bool {
true
}
fn split(
&self,
indices: &<<FaerSparseMat<T> as MatrixCommon>::V as Vector>::Index,
) -> [(
SymbolicSparseColMat<IndexType>,
<<FaerSparseMat<T> as MatrixCommon>::V as Vector>::Index,
); 4] {
let (_ni, _nj, col_ptrs, _col_nnz, row_idx) = self.parts();
let ctx = indices.context();
let (ul_blk, ur_blk, ll_blk, lr_blk) = CscBlock::split(row_idx, col_ptrs, indices);
let ul_sym = SymbolicSparseColMat::new_checked(
ul_blk.nrows,
ul_blk.ncols,
ul_blk.col_pointers,
None,
ul_blk.row_indices,
);
let ur_sym = SymbolicSparseColMat::new_checked(
ur_blk.nrows,
ur_blk.ncols,
ur_blk.col_pointers,
None,
ur_blk.row_indices,
);
let ll_sym = SymbolicSparseColMat::new_checked(
ll_blk.nrows,
ll_blk.ncols,
ll_blk.col_pointers,
None,
ll_blk.row_indices,
);
let lr_sym = SymbolicSparseColMat::new_checked(
lr_blk.nrows,
lr_blk.ncols,
lr_blk.col_pointers,
None,
lr_blk.row_indices,
);
[
(
ul_sym,
FaerVecIndex {
data: ul_blk.src_indices,
context: *ctx,
},
),
(
ur_sym,
FaerVecIndex {
data: ur_blk.src_indices,
context: *ctx,
},
),
(
ll_sym,
FaerVecIndex {
data: ll_blk.src_indices,
context: *ctx,
},
),
(
lr_sym,
FaerVecIndex {
data: lr_blk.src_indices,
context: *ctx,
},
),
]
}
fn indices(&self) -> Vec<(IndexType, IndexType)> {
let mut indices = Vec::with_capacity(self.compute_nnz());
for col_i in 0..self.ncols() {
for row_j in self.col_range(col_i) {
indices.push((row_j, col_i));
}
}
indices
}
}
impl<T: FaerScalar> Matrix for FaerSparseMat<T> {
type Sparsity = SymbolicSparseColMat<IndexType>;
type SparsityRef<'a> = SymbolicSparseColMatRef<'a, IndexType>;
fn sparsity(&self) -> Option<Self::SparsityRef<'_>> {
Some(self.data.symbolic())
}
fn context(&self) -> &FaerContext {
&self.context
}
fn gather(&mut self, other: &Self, indices: &<Self::V as Vector>::Index) {
let dst_data = self.data.val_mut();
let src_data = other.data.val();
for (dst_i, idx) in dst_data.iter_mut().zip(indices.data.iter()) {
*dst_i = src_data[*idx];
}
}
fn set_data_with_indices(
&mut self,
dst_indices: &<Self::V as Vector>::Index,
src_indices: &<Self::V as Vector>::Index,
data: &Self::V,
) {
let values = self.data.val_mut();
for (dst_i, src_i) in dst_indices.data.iter().zip(src_indices.data.iter()) {
values[*dst_i] = data[*src_i];
}
}
fn add_column_to_vector(&self, j: IndexType, v: &mut Self::V) {
for i in self.data.col_range(j) {
let row = self.data.row_idx()[i];
v[row] += self.data.val()[i];
}
}
fn triplet_iter(&self) -> impl Iterator<Item = (IndexType, IndexType, Self::T)> {
(0..self.ncols()).flat_map(move |j| {
self.data.col_range(j).map(move |i| {
let row = self.data.row_idx()[i];
(row, j, self.data.val()[i])
})
})
}
fn try_from_triplets(
nrows: IndexType,
ncols: IndexType,
triplets: Vec<(IndexType, IndexType, T)>,
ctx: Self::C,
) -> Result<Self, DiffsolError> {
let triplets = triplets
.iter()
.map(|(i, j, v)| Triplet::new(*i, *j, *v))
.collect::<Vec<_>>();
match faer::sparse::SparseColMat::try_new_from_triplets(nrows, ncols, triplets.as_slice()) {
Ok(mat) => Ok(Self {
data: mat,
context: ctx,
}),
Err(e) => Err(DiffsolError::from(
MatrixError::FailedToCreateMatrixFromTriplets(e),
)),
}
}
fn gemv(&self, alpha: Self::T, x: &Self::V, beta: Self::T, y: &mut Self::V) {
let tmp = Self::V {
data: &self.data * &x.data,
context: self.context,
};
y.axpy(alpha, &tmp, beta);
}
fn zeros(nrows: IndexType, ncols: IndexType, ctx: Self::C) -> Self {
Self {
data: SparseColMat::try_new_from_triplets(nrows, ncols, &[]).unwrap(),
context: ctx,
}
}
fn copy_from(&mut self, other: &Self) {
self.data = faer::sparse::SparseColMat::new(
other.data.symbolic().to_owned().unwrap(),
other.data.val().to_vec(),
)
}
fn from_diagonal(v: &FaerVec<T>) -> Self {
let dim = v.len();
let triplets = (0..dim)
.map(|i| Triplet::new(i, i, v[i]))
.collect::<Vec<_>>();
Self {
data: SparseColMat::try_new_from_triplets(dim, dim, &triplets).unwrap(),
context: *v.context(),
}
}
fn partition_indices_by_zero_diagonal(
&self,
) -> (<Self::V as Vector>::Index, <Self::V as Vector>::Index) {
let mut indices_zero_diag = vec![];
let mut indices_non_zero_diag = vec![];
'outer: for j in 0..self.ncols() {
for (i, v) in self.data.row_idx_of_col(j).zip(self.data.val_of_col(j)) {
if i == j && *v != T::zero() {
indices_non_zero_diag.push(j);
continue 'outer;
} else if i > j {
break;
}
}
indices_zero_diag.push(j);
}
(
<Self::V as Vector>::Index::from_vec(indices_zero_diag, self.context),
<Self::V as Vector>::Index::from_vec(indices_non_zero_diag, self.context),
)
}
fn set_column(&mut self, j: IndexType, v: &Self::V) {
assert_eq!(v.len(), self.nrows());
for i in self.data.col_range(j) {
let row_i = self.data.row_idx()[i];
self.data.val_mut()[i] = v[row_i];
}
}
fn scale_add_and_assign(&mut self, x: &Self, beta: Self::T, y: &Self) {
ternary_op_assign_into(self.data.rb_mut(), x.data.rb(), y.data.rb(), |s, x, y| {
*s = *x.unwrap_or(&T::zero()) + beta * *y.unwrap_or(&T::zero())
});
}
fn new_from_sparsity(
nrows: IndexType,
ncols: IndexType,
sparsity: Option<Self::Sparsity>,
ctx: Self::C,
) -> Self {
let sparsity = sparsity.expect("Sparsity pattern required for sparse matrix");
assert_eq!(sparsity.nrows(), nrows);
assert_eq!(sparsity.ncols(), ncols);
let nnz = sparsity.row_idx().len();
Self {
data: SparseColMat::new(sparsity, vec![T::zero(); nnz]),
context: ctx,
}
}
}
#[cfg(test)]
mod tests {
use crate::{FaerSparseMat, Matrix};
#[test]
fn test_triplet_iter() {
let triplets = vec![(0, 0, 1.0), (1, 0, 2.0), (2, 2, 3.0), (3, 2, 4.0)];
let mat =
FaerSparseMat::<f64>::try_from_triplets(4, 3, triplets.clone(), Default::default())
.unwrap();
let mut iter = mat.triplet_iter();
for triplet in triplets {
let (i, j, val) = iter.next().unwrap();
assert_eq!(i, triplet.0);
assert_eq!(j, triplet.1);
assert_eq!(val, triplet.2);
}
}
#[test]
fn test_partition_indices_by_zero_diagonal() {
super::super::tests::test_partition_indices_by_zero_diagonal::<FaerSparseMat<f64>>();
}
}