1use std::fmt::Debug;
2use std::ops::{Add, Mul, Sub};
3
4use super::extract_block::CscBlock;
5use super::sparsity::MatrixSparsityRef;
6use super::utils::*;
7use super::{Matrix, MatrixCommon, MatrixSparsity};
8use crate::error::{LaError, MatrixError};
9use crate::{DefaultSolver, FaerScalar, FaerSparseLU, IndexType, Scalar, Scale};
10use crate::{FaerContext, FaerVec, FaerVecIndex, Vector, VectorIndex};
11
12use faer::reborrow::{Reborrow, ReborrowMut};
13use faer::sparse::ops::{ternary_op_assign_into, union_symbolic};
14use faer::sparse::{Pair, SparseColMat, SymbolicSparseColMat, SymbolicSparseColMatRef, Triplet};
15
16#[derive(Clone, Debug)]
17pub struct FaerSparseMat<T: FaerScalar> {
18 pub(crate) data: SparseColMat<IndexType, T>,
19 pub(crate) context: FaerContext,
20}
21
22impl<T: FaerScalar> DefaultSolver for FaerSparseMat<T> {
23 type LS = FaerSparseLU<T>;
24}
25
26impl_matrix_common!(
27 FaerSparseMat<T>,
28 FaerVec<T>,
29 FaerContext,
30 SparseColMat<IndexType, T>,
31 FaerScalar
32);
33
34macro_rules! impl_mul_scalar {
35 ($mat_type:ty, $out:ty) => {
36 impl<'a, T: FaerScalar> Mul<Scale<T>> for $mat_type {
37 type Output = $out;
38
39 fn mul(self, rhs: Scale<T>) -> Self::Output {
40 let scale: faer::Scale<T> = rhs.into();
41 Self::Output {
42 data: &self.data * scale,
43 context: self.context,
44 }
45 }
46 }
47 };
48}
49
50impl_mul_scalar!(FaerSparseMat<T>, FaerSparseMat<T>);
51impl_mul_scalar!(&FaerSparseMat<T>, FaerSparseMat<T>);
52
53impl_add!(
54 FaerSparseMat<T>,
55 &FaerSparseMat<T>,
56 FaerSparseMat<T>,
57 FaerScalar
58);
59
60impl_sub!(
61 FaerSparseMat<T>,
62 &FaerSparseMat<T>,
63 FaerSparseMat<T>,
64 FaerScalar
65);
66
67impl<T: FaerScalar> MatrixSparsity<FaerSparseMat<T>> for SymbolicSparseColMat<IndexType> {
68 fn union(
69 self,
70 other: SymbolicSparseColMatRef<IndexType>,
71 ) -> Result<SymbolicSparseColMat<IndexType>, LaError> {
72 union_symbolic(self.rb(), other).map_err(|e| LaError::Other(e.to_string()))
73 }
74
75 fn as_ref(&self) -> SymbolicSparseColMatRef<'_, IndexType> {
76 self.rb()
77 }
78
79 fn nrows(&self) -> IndexType {
80 self.nrows()
81 }
82
83 fn ncols(&self) -> IndexType {
84 self.ncols()
85 }
86
87 fn is_sparse() -> bool {
88 true
89 }
90
91 fn indices(&self) -> Vec<(IndexType, IndexType)> {
92 let mut indices = Vec::with_capacity(self.compute_nnz());
93 for col_i in 0..self.ncols() {
94 for row_j in self.col_range(col_i) {
95 indices.push((row_j, col_i));
96 }
97 }
98 indices
99 }
100
101 fn new_diagonal(n: IndexType) -> Self {
102 let indices = (0..n).map(|i| Pair::new(i, i)).collect::<Vec<_>>();
103 SymbolicSparseColMat::try_new_from_indices(n, n, indices.as_slice())
104 .unwrap()
105 .0
106 }
107
108 fn try_from_indices(
109 nrows: IndexType,
110 ncols: IndexType,
111 indices: Vec<(IndexType, IndexType)>,
112 ) -> Result<Self, LaError> {
113 let indices = indices
114 .iter()
115 .map(|(i, j)| Pair::new(*i, *j))
116 .collect::<Vec<_>>();
117 match Self::try_new_from_indices(nrows, ncols, indices.as_slice()) {
118 Ok((sparsity, _)) => Ok(sparsity),
119 Err(e) => Err(LaError::Other(e.to_string())),
120 }
121 }
122
123 fn get_index(
124 &self,
125 indices: &[(IndexType, IndexType)],
126 ctx: FaerContext,
127 ) -> <<FaerSparseMat<T> as MatrixCommon>::V as Vector>::Index {
128 let col_ptrs = self.col_ptr();
129 let row_indices = self.row_idx();
130 let mut ret = Vec::with_capacity(indices.len());
131 for &(i, j) in indices.iter() {
132 let col_ptr = col_ptrs[j];
133 let next_col_ptr = col_ptrs[j + 1];
134 for (ii, &ri) in row_indices
135 .iter()
136 .enumerate()
137 .take(next_col_ptr)
138 .skip(col_ptr)
139 {
140 if ri == i {
141 ret.push(ii);
142 break;
143 }
144 }
145 }
146 FaerVecIndex {
147 data: ret,
148 context: ctx,
149 }
150 }
151}
152
153impl<'a, T: FaerScalar> MatrixSparsityRef<'a, FaerSparseMat<T>>
154 for SymbolicSparseColMatRef<'a, IndexType>
155{
156 fn to_owned(&self) -> SymbolicSparseColMat<IndexType> {
157 self.to_owned().unwrap()
158 }
159 fn nrows(&self) -> IndexType {
160 self.nrows()
161 }
162
163 fn ncols(&self) -> IndexType {
164 self.ncols()
165 }
166
167 fn is_sparse() -> bool {
168 true
169 }
170
171 fn split(
172 &self,
173 indices: &<<FaerSparseMat<T> as MatrixCommon>::V as Vector>::Index,
174 ) -> [(
175 SymbolicSparseColMat<IndexType>,
176 <<FaerSparseMat<T> as MatrixCommon>::V as Vector>::Index,
177 ); 4] {
178 let (_ni, _nj, col_ptrs, _col_nnz, row_idx) = self.parts();
179 let ctx = indices.context();
180 let (ul_blk, ur_blk, ll_blk, lr_blk) = CscBlock::split(row_idx, col_ptrs, indices);
181 let ul_sym = SymbolicSparseColMat::new_checked(
182 ul_blk.nrows,
183 ul_blk.ncols,
184 ul_blk.col_pointers,
185 None,
186 ul_blk.row_indices,
187 );
188 let ur_sym = SymbolicSparseColMat::new_checked(
189 ur_blk.nrows,
190 ur_blk.ncols,
191 ur_blk.col_pointers,
192 None,
193 ur_blk.row_indices,
194 );
195 let ll_sym = SymbolicSparseColMat::new_checked(
196 ll_blk.nrows,
197 ll_blk.ncols,
198 ll_blk.col_pointers,
199 None,
200 ll_blk.row_indices,
201 );
202 let lr_sym = SymbolicSparseColMat::new_checked(
203 lr_blk.nrows,
204 lr_blk.ncols,
205 lr_blk.col_pointers,
206 None,
207 lr_blk.row_indices,
208 );
209 [
210 (
211 ul_sym,
212 FaerVecIndex {
213 data: ul_blk.src_indices,
214 context: *ctx,
215 },
216 ),
217 (
218 ur_sym,
219 FaerVecIndex {
220 data: ur_blk.src_indices,
221 context: *ctx,
222 },
223 ),
224 (
225 ll_sym,
226 FaerVecIndex {
227 data: ll_blk.src_indices,
228 context: *ctx,
229 },
230 ),
231 (
232 lr_sym,
233 FaerVecIndex {
234 data: lr_blk.src_indices,
235 context: *ctx,
236 },
237 ),
238 ]
239 }
240
241 fn indices(&self) -> Vec<(IndexType, IndexType)> {
242 let mut indices = Vec::with_capacity(self.compute_nnz());
243 for col_i in 0..self.ncols() {
244 for row_j in self.col_range(col_i) {
245 indices.push((row_j, col_i));
246 }
247 }
248 indices
249 }
250}
251
252impl<T: FaerScalar> Matrix for FaerSparseMat<T> {
253 type Sparsity = SymbolicSparseColMat<IndexType>;
254 type SparsityRef<'a> = SymbolicSparseColMatRef<'a, IndexType>;
255
256 fn sparsity(&self) -> Option<Self::SparsityRef<'_>> {
257 Some(self.data.symbolic())
258 }
259 fn context(&self) -> &FaerContext {
260 &self.context
261 }
262 fn inner_mut(&mut self) -> &mut Self::Inner {
263 &mut self.data
264 }
265
266 fn gather(&mut self, other: &Self, indices: &<Self::V as Vector>::Index) {
267 let dst_data = self.data.val_mut();
268 let src_data = other.data.val();
269 for (dst_i, idx) in dst_data.iter_mut().zip(indices.data.iter()) {
270 *dst_i = src_data[*idx];
271 }
272 }
273
274 fn set_data_with_indices(
275 &mut self,
276 dst_indices: &<Self::V as Vector>::Index,
277 src_indices: &<Self::V as Vector>::Index,
278 data: &Self::V,
279 ) {
280 let values = self.data.val_mut();
281 for (dst_i, src_i) in dst_indices.data.iter().zip(src_indices.data.iter()) {
282 values[*dst_i] = data[*src_i];
283 }
284 }
285
286 fn add_column_to_vector(&self, j: IndexType, v: &mut Self::V) {
287 for i in self.data.col_range(j) {
288 let row = self.data.row_idx()[i];
289 v[row] += self.data.val()[i];
290 }
291 }
292
293 fn triplet_iter(
294 &self,
295 ) -> (
296 impl Iterator<Item = (IndexType, IndexType)> + '_,
297 impl Iterator<Item = Self::T> + '_,
298 ) {
299 let indices: Vec<_> = (0..self.ncols())
300 .flat_map(move |j| {
301 self.data
302 .col_range(j)
303 .map(move |i| (self.data.row_idx()[i], j))
304 })
305 .collect();
306 let values: Vec<_> = indices
307 .iter()
308 .enumerate()
309 .map(|(k, _)| self.data.val()[k])
310 .collect();
311 (indices.into_iter(), values.into_iter())
312 }
313
314 fn try_from_triplets(
315 nrows: IndexType,
316 ncols: IndexType,
317 indices: Vec<(IndexType, IndexType)>,
318 values: Vec<Self::T>,
319 ctx: Self::C,
320 ) -> Result<Self, LaError> {
321 assert_eq!(
322 values.len(),
323 indices.len(),
324 "values.len() must equal indices.len() for non-batched backend"
325 );
326 let triplets = indices
327 .iter()
328 .zip(values)
329 .map(|((i, j), v)| Triplet::new(*i, *j, v))
330 .collect::<Vec<_>>();
331 match faer::sparse::SparseColMat::try_new_from_triplets(nrows, ncols, triplets.as_slice()) {
332 Ok(mat) => Ok(Self {
333 data: mat,
334 context: ctx,
335 }),
336 Err(e) => Err(LaError::from(
337 MatrixError::FailedToCreateMatrixFromTriplets(e),
338 )),
339 }
340 }
341 fn gemv(&self, alpha: Self::T, x: &Self::V, beta: Self::T, y: &mut Self::V) {
342 let tmp = Self::V {
343 data: &self.data * &x.data,
344 context: self.context,
345 };
346 y.axpy(alpha, &tmp, beta);
347 }
348 fn zeros(nrows: IndexType, ncols: IndexType, ctx: Self::C) -> Self {
349 Self {
350 data: SparseColMat::try_new_from_triplets(nrows, ncols, &[]).unwrap(),
351 context: ctx,
352 }
353 }
354 fn copy_from(&mut self, other: &Self) {
355 self.data = faer::sparse::SparseColMat::new(
356 other.data.symbolic().to_owned().unwrap(),
357 other.data.val().to_vec(),
358 )
359 }
360 fn from_diagonal(v: &FaerVec<T>) -> Self {
361 let dim = v.len();
362 let triplets = (0..dim)
363 .map(|i| Triplet::new(i, i, v[i]))
364 .collect::<Vec<_>>();
365 Self {
366 data: SparseColMat::try_new_from_triplets(dim, dim, &triplets).unwrap(),
367 context: *v.context(),
368 }
369 }
370
371 fn partition_indices_by_zero_diagonal(
372 &self,
373 ) -> (<Self::V as Vector>::Index, <Self::V as Vector>::Index) {
374 let mut indices_zero_diag = vec![];
375 let mut indices_non_zero_diag = vec![];
376 'outer: for j in 0..self.ncols() {
377 for (i, v) in self.data.row_idx_of_col(j).zip(self.data.val_of_col(j)) {
378 if i == j && *v != T::zero() {
379 indices_non_zero_diag.push(j);
380 continue 'outer;
381 } else if i > j {
382 break;
383 }
384 }
385 indices_zero_diag.push(j);
386 }
387 (
388 <Self::V as Vector>::Index::from_vec(indices_zero_diag, self.context),
389 <Self::V as Vector>::Index::from_vec(indices_non_zero_diag, self.context),
390 )
391 }
392
393 fn set_column(&mut self, j: IndexType, v: &Self::V) {
394 assert_eq!(v.len(), self.nrows());
395 for i in self.data.col_range(j) {
396 let row_i = self.data.row_idx()[i];
397 self.data.val_mut()[i] = v[row_i];
398 }
399 }
400
401 fn scale_add_and_assign(&mut self, x: &Self, beta: Self::T, y: &Self) {
402 ternary_op_assign_into(self.data.rb_mut(), x.data.rb(), y.data.rb(), |s, x, y| {
403 *s = *x.unwrap_or(&T::zero()) + beta * *y.unwrap_or(&T::zero())
404 });
405 }
406
407 fn new_from_sparsity(
408 nrows: IndexType,
409 ncols: IndexType,
410 sparsity: Option<Self::Sparsity>,
411 ctx: Self::C,
412 ) -> Self {
413 let sparsity = sparsity.expect("Sparsity pattern required for sparse matrix");
414 assert_eq!(sparsity.nrows(), nrows);
415 assert_eq!(sparsity.ncols(), ncols);
416 let nnz = sparsity.row_idx().len();
417 Self {
418 data: SparseColMat::new(sparsity, vec![T::zero(); nnz]),
419 context: ctx,
420 }
421 }
422}
423
424#[cfg(test)]
425mod tests {
426 use crate::{FaerSparseMat, Matrix};
427 #[test]
428 fn test_triplet_iter() {
429 let indices = vec![(0, 0), (1, 0), (2, 2), (3, 2)];
430 let values = vec![1.0, 2.0, 3.0, 4.0];
431 let mat = FaerSparseMat::<f64>::try_from_triplets(
432 4,
433 3,
434 indices.clone(),
435 values.clone(),
436 Default::default(),
437 )
438 .unwrap();
439 let (iter_indices, iter_values): (Vec<_>, Vec<_>) = {
440 let (i, v) = mat.triplet_iter();
441 (i.collect(), v.collect())
442 };
443 assert_eq!(iter_indices, indices);
444 assert_eq!(iter_values, values);
445 }
446
447 #[test]
448 fn test_set_data_with_indices() {
449 use crate::{FaerVec, FaerVecIndex, Vector, VectorIndex};
450
451 let triplets_full = vec![
453 (0, 0, 1.0),
454 (2, 0, 2.0),
455 (1, 1, 3.0),
456 (0, 2, 4.0),
457 (2, 2, 5.0),
458 ];
459 let (indices, init_values): (Vec<_>, Vec<_>) =
460 triplets_full.iter().map(|&(i, j, v)| ((i, j), v)).unzip();
461 let mut mat = FaerSparseMat::<f64>::try_from_triplets(
462 3,
463 3,
464 indices.clone(),
465 init_values,
466 Default::default(),
467 )
468 .unwrap();
469
470 assert_eq!(mat.inner_mut().val_mut().len(), indices.len());
471
472 let new_values = [10.0, 11.0, 12.0, 13.0, 14.0];
473 mat.inner_mut().val_mut().copy_from_slice(&new_values);
474 let (indices_iter, values_iter) = mat.triplet_iter();
475 let got: Vec<(usize, usize, f64)> = indices_iter
476 .zip(values_iter)
477 .map(|((i, j), v)| (i, j, v))
478 .collect();
479 let expected: Vec<(usize, usize, f64)> = triplets_full
480 .iter()
481 .zip(new_values.iter())
482 .map(|(&(i, j, _), &v)| (i, j, v))
483 .collect();
484 assert_eq!(got, expected);
485
486 let mut via_set_data = FaerSparseMat::<f64>::try_from_triplets(
487 3,
488 3,
489 indices,
490 new_values.iter().map(|_| 0.0).collect(),
491 Default::default(),
492 )
493 .unwrap();
494 let nnz = new_values.len();
495 let identity = FaerVecIndex::from_vec((0..nnz).collect(), Default::default());
496 let data = FaerVec::from_vec(new_values.to_vec(), Default::default());
497 via_set_data.set_data_with_indices(&identity, &identity, &data);
498 assert_eq!(
499 mat.inner_mut().val_mut(),
500 via_set_data.inner_mut().val_mut()
501 );
502 }
503
504 #[test]
505 fn test_partition_indices_by_zero_diagonal() {
506 super::super::tests::test_partition_indices_by_zero_diagonal::<FaerSparseMat<f64>>();
507 }
508
509 super::super::generate_matrix_tests_nonbatched!(faer_sparse, FaerSparseMat<f64>);
510}