1use crate::error::{SparseError, SparseResult};
7use scirs2_core::numeric::{Float, SparseElement, Zero};
8use scirs2_core::GpuDataType;
9use std::cmp::PartialEq;
10
11#[derive(Clone)]
16pub struct CsrMatrix<T> {
17 rows: usize,
19 cols: usize,
21 pub indptr: Vec<usize>,
23 pub indices: Vec<usize>,
25 pub data: Vec<T>,
27}
28
29impl<T> CsrMatrix<T>
30where
31 T: Clone + Copy + Zero + PartialEq + SparseElement,
32{
33 pub fn get(&self, row: usize, col: usize) -> T {
35 if row >= self.rows || col >= self.cols {
37 return T::sparse_zero();
38 }
39
40 for j in self.indptr[row]..self.indptr[row + 1] {
42 if self.indices[j] == col {
43 return self.data[j];
44 }
45 }
46
47 T::sparse_zero()
49 }
50
51 pub fn get_triplets(&self) -> (Vec<usize>, Vec<usize>, Vec<T>) {
53 let mut rows = Vec::new();
54 let mut cols = Vec::new();
55 let mut values = Vec::new();
56
57 for i in 0..self.rows {
58 for j in self.indptr[i]..self.indptr[i + 1] {
59 rows.push(i);
60 cols.push(self.indices[j]);
61 values.push(self.data[j]);
62 }
63 }
64
65 (rows, cols, values)
66 }
67 pub fn new(
94 data: Vec<T>,
95 rowindices: Vec<usize>,
96 colindices: Vec<usize>,
97 shape: (usize, usize),
98 ) -> SparseResult<Self> {
99 if data.len() != rowindices.len() || data.len() != colindices.len() {
101 return Err(SparseError::DimensionMismatch {
102 expected: data.len(),
103 found: std::cmp::min(rowindices.len(), colindices.len()),
104 });
105 }
106
107 let (rows, cols) = shape;
108
109 if rowindices.iter().any(|&i| i >= rows) {
111 return Err(SparseError::ValueError(
112 "Row index out of bounds".to_string(),
113 ));
114 }
115
116 if colindices.iter().any(|&i| i >= cols) {
117 return Err(SparseError::ValueError(
118 "Column index out of bounds".to_string(),
119 ));
120 }
121
122 let mut triplets: Vec<(usize, usize, T)> = rowindices
125 .into_iter()
126 .zip(colindices)
127 .zip(data)
128 .map(|((r, c), v)| (r, c, v))
129 .collect();
130 triplets.sort_by_key(|&(r, c_, _)| (r, c_));
131
132 let nnz = triplets.len();
134 let mut indptr = vec![0; rows + 1];
135 let mut indices = Vec::with_capacity(nnz);
136 let mut data_out = Vec::with_capacity(nnz);
137
138 for &(r_, _, _) in &triplets {
140 indptr[r_ + 1] += 1;
141 }
142
143 for i in 1..=rows {
145 indptr[i] += indptr[i - 1];
146 }
147
148 for (_r, c, v) in triplets {
150 indices.push(c);
151 data_out.push(v);
152 }
153
154 Ok(CsrMatrix {
155 rows,
156 cols,
157 indptr,
158 indices,
159 data: data_out,
160 })
161 }
162
163 pub fn from_triplets(
194 nrows: usize,
195 ncols: usize,
196 row_indices: Vec<usize>,
197 col_indices: Vec<usize>,
198 values: Vec<T>,
199 ) -> SparseResult<Self> {
200 Self::new(values, row_indices, col_indices, (nrows, ncols))
201 }
202
203 pub fn try_from_triplets(
238 nrows: usize,
239 ncols: usize,
240 triplets: &[(usize, usize, T)],
241 ) -> SparseResult<Self> {
242 let mut row_indices = Vec::with_capacity(triplets.len());
243 let mut col_indices = Vec::with_capacity(triplets.len());
244 let mut values = Vec::with_capacity(triplets.len());
245
246 for &(r, c, v) in triplets {
247 row_indices.push(r);
248 col_indices.push(c);
249 values.push(v);
250 }
251
252 Self::from_triplets(nrows, ncols, row_indices, col_indices, values)
253 }
254
255 pub fn from_raw_csr(
268 data: Vec<T>,
269 indptr: Vec<usize>,
270 indices: Vec<usize>,
271 shape: (usize, usize),
272 ) -> SparseResult<Self> {
273 let (rows, cols) = shape;
274
275 if indptr.len() != rows + 1 {
277 return Err(SparseError::DimensionMismatch {
278 expected: rows + 1,
279 found: indptr.len(),
280 });
281 }
282
283 if data.len() != indices.len() {
284 return Err(SparseError::DimensionMismatch {
285 expected: data.len(),
286 found: indices.len(),
287 });
288 }
289
290 for i in 1..indptr.len() {
292 if indptr[i] < indptr[i - 1] {
293 return Err(SparseError::ValueError(
294 "Row pointer array must be monotonically increasing".to_string(),
295 ));
296 }
297 }
298
299 if indptr[rows] != data.len() {
301 return Err(SparseError::ValueError(
302 "Last row pointer entry must match data length".to_string(),
303 ));
304 }
305
306 if indices.iter().any(|&i| i >= cols) {
308 return Err(SparseError::ValueError(
309 "Column index out of bounds".to_string(),
310 ));
311 }
312
313 Ok(CsrMatrix {
314 rows,
315 cols,
316 indptr,
317 indices,
318 data,
319 })
320 }
321
322 pub fn empty(shape: (usize, usize)) -> Self {
332 let (rows, cols) = shape;
333 let indptr = vec![0; rows + 1];
334
335 CsrMatrix {
336 rows,
337 cols,
338 indptr,
339 indices: Vec::new(),
340 data: Vec::new(),
341 }
342 }
343
344 pub fn rows(&self) -> usize {
346 self.rows
347 }
348
349 pub fn cols(&self) -> usize {
351 self.cols
352 }
353
354 pub fn shape(&self) -> (usize, usize) {
356 (self.rows, self.cols)
357 }
358
359 pub fn nnz(&self) -> usize {
361 self.data.len()
362 }
363
364 pub fn to_dense(&self) -> Vec<Vec<T>>
366 where
367 T: Zero + Copy + SparseElement,
368 {
369 let mut result = vec![vec![T::sparse_zero(); self.cols]; self.rows];
370
371 for (row_idx, row) in result.iter_mut().enumerate() {
372 for j in self.indptr[row_idx]..self.indptr[row_idx + 1] {
373 let col_idx = self.indices[j];
374 row[col_idx] = self.data[j];
375 }
376 }
377
378 result
379 }
380
381 pub fn transpose(&self) -> Self {
383 let mut col_counts = vec![0; self.cols];
385 for &col in &self.indices {
386 col_counts[col] += 1;
387 }
388
389 let mut col_ptrs = vec![0; self.cols + 1];
391 for i in 0..self.cols {
392 col_ptrs[i + 1] = col_ptrs[i] + col_counts[i];
393 }
394
395 let nnz = self.nnz();
397 let mut indices_t = vec![0; nnz];
398 let mut data_t = vec![T::sparse_zero(); nnz];
399 let mut col_counts = vec![0; self.cols];
400
401 for row in 0..self.rows {
402 for j in self.indptr[row]..self.indptr[row + 1] {
403 let col = self.indices[j];
404 let dest = col_ptrs[col] + col_counts[col];
405
406 indices_t[dest] = row;
407 data_t[dest] = self.data[j];
408 col_counts[col] += 1;
409 }
410 }
411
412 CsrMatrix {
413 rows: self.cols,
414 cols: self.rows,
415 indptr: col_ptrs,
416 indices: indices_t,
417 data: data_t,
418 }
419 }
420}
421
422impl<
423 T: Clone
424 + Copy
425 + std::ops::AddAssign
426 + std::ops::MulAssign
427 + std::cmp::PartialEq
428 + std::fmt::Debug
429 + scirs2_core::numeric::Zero
430 + std::ops::Add<Output = T>
431 + std::ops::Mul<Output = T>
432 + SparseElement,
433 > CsrMatrix<T>
434{
435 pub fn is_symmetric(&self) -> bool {
441 if self.rows != self.cols {
442 return false;
443 }
444
445 let transposed = self.transpose();
447
448 if self.nnz() != transposed.nnz() {
450 return false;
451 }
452
453 for row in 0..self.rows {
455 let self_start = self.indptr[row];
456 let self_end = self.indptr[row + 1];
457 let trans_start = transposed.indptr[row];
458 let trans_end = transposed.indptr[row + 1];
459
460 if self_end - self_start != trans_end - trans_start {
461 return false;
462 }
463
464 let mut self_entries: Vec<(usize, &T)> = (self_start..self_end)
466 .map(|j| (self.indices[j], &self.data[j]))
467 .collect();
468 self_entries.sort_by_key(|(col_, _)| *col_);
469
470 let mut trans_entries: Vec<(usize, &T)> = (trans_start..trans_end)
471 .map(|j| (transposed.indices[j], &transposed.data[j]))
472 .collect();
473 trans_entries.sort_by_key(|(col_, _)| *col_);
474
475 for i in 0..self_entries.len() {
477 if self_entries[i].0 != trans_entries[i].0
478 || self_entries[i].1 != trans_entries[i].1
479 {
480 return false;
481 }
482 }
483 }
484
485 true
486 }
487
488 pub fn matmul(&self, other: &CsrMatrix<T>) -> SparseResult<CsrMatrix<T>> {
498 if self.cols != other.rows {
499 return Err(SparseError::DimensionMismatch {
500 expected: self.cols,
501 found: other.rows,
502 });
503 }
504
505 let a_dense = self.to_dense();
508 let b_dense = other.to_dense();
509
510 let m = self.rows;
511 let n = other.cols;
512 let k = self.cols;
513
514 let mut c_dense = vec![vec![T::sparse_zero(); n]; m];
515
516 for (i, c_row) in c_dense.iter_mut().enumerate().take(m) {
517 for (j, val) in c_row.iter_mut().enumerate().take(n) {
518 for (l, &a_val) in a_dense[i].iter().enumerate().take(k) {
519 let prod = a_val * b_dense[l][j];
520 *val += prod;
521 }
522 }
523 }
524
525 let mut rowindices = Vec::new();
527 let mut colindices = Vec::new();
528 let mut values = Vec::new();
529
530 for (i, row) in c_dense.iter().enumerate() {
531 for (j, val) in row.iter().enumerate() {
532 if *val != T::sparse_zero() {
533 rowindices.push(i);
534 colindices.push(j);
535 values.push(*val);
536 }
537 }
538 }
539
540 CsrMatrix::new(values, rowindices, colindices, (m, n))
541 }
542
543 pub fn row_range(&self, row: usize) -> std::ops::Range<usize> {
553 assert!(row < self.rows, "Row index out of bounds");
554 self.indptr[row]..self.indptr[row + 1]
555 }
556
557 pub fn colindices(&self) -> &[usize] {
559 &self.indices
560 }
561}
562
563impl CsrMatrix<f64> {
564 pub fn dot(&self, vec: &[f64]) -> SparseResult<Vec<f64>> {
574 if vec.len() != self.cols {
575 return Err(SparseError::DimensionMismatch {
576 expected: self.cols,
577 found: vec.len(),
578 });
579 }
580
581 let mut result = vec![0.0; self.rows];
582
583 for (row_idx, result_val) in result.iter_mut().enumerate() {
584 for j in self.indptr[row_idx]..self.indptr[row_idx + 1] {
585 let col_idx = self.indices[j];
586 *result_val += self.data[j] * vec[col_idx];
587 }
588 }
589
590 Ok(result)
591 }
592
593 #[allow(dead_code)]
621 pub fn gpu_dot(&self, vec: &[f64]) -> SparseResult<Vec<f64>> {
622 let gpu_spmv = crate::gpu_spmv_implementation::GpuSpMV::new()?;
624 gpu_spmv.spmv(
625 self.rows,
626 self.cols,
627 &self.indptr,
628 &self.indices,
629 &self.data,
630 vec,
631 )
632 }
633
634 #[allow(dead_code)]
645 pub fn gpu_dot_with_backend(
646 &self,
647 vec: &[f64],
648 backend: scirs2_core::gpu::GpuBackend,
649 ) -> SparseResult<Vec<f64>> {
650 let gpu_spmv = crate::gpu_spmv_implementation::GpuSpMV::with_backend(backend)?;
652 gpu_spmv.spmv(
653 self.rows,
654 self.cols,
655 &self.indptr,
656 &self.indices,
657 &self.data,
658 vec,
659 )
660 }
661}
662
663impl<T> CsrMatrix<T>
664where
665 T: scirs2_core::numeric::Float
666 + std::fmt::Debug
667 + Copy
668 + Default
669 + GpuDataType
670 + Send
671 + Sync
672 + SparseElement
673 + std::ops::AddAssign
674 + std::ops::Mul<Output = T>
675 + 'static,
676{
677 #[allow(dead_code)]
687 pub fn gpu_dot_generic(&self, vec: &[T]) -> SparseResult<Vec<T>>
688where {
689 if vec.len() != self.cols {
691 return Err(SparseError::DimensionMismatch {
692 expected: self.cols,
693 found: vec.len(),
694 });
695 }
696
697 let mut result = vec![T::sparse_zero(); self.rows];
698
699 for (row_idx, result_val) in result.iter_mut().enumerate() {
700 let start = self.indptr[row_idx];
701 let end = self.indptr[row_idx + 1];
702
703 for idx in start..end {
704 let col = self.indices[idx];
705 *result_val += self.data[idx] * vec[col];
706 }
707 }
708
709 Ok(result)
710 }
711
712 pub fn should_use_gpu(&self) -> bool {
718 let nnz_threshold = 10000;
721 let density = self.nnz() as f64 / (self.rows * self.cols) as f64;
722
723 self.nnz() > nnz_threshold && density < 0.5
724 }
725
726 #[allow(dead_code)]
732 pub fn gpu_backend_info() -> SparseResult<(crate::gpu_ops::GpuBackend, String)> {
733 Ok((crate::gpu_ops::GpuBackend::Cpu, "CPU Fallback".to_string()))
735 }
736}
737
738#[cfg(test)]
739mod tests {
740 use super::*;
741 use approx::assert_relative_eq;
742
743 #[test]
744 fn test_csr_create() {
745 let rows = vec![0, 0, 1, 2, 2];
747 let cols = vec![0, 2, 2, 0, 1];
748 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
749 let shape = (3, 3);
750
751 let matrix = CsrMatrix::new(data, rows, cols, shape).unwrap();
752
753 assert_eq!(matrix.shape(), (3, 3));
754 assert_eq!(matrix.nnz(), 5);
755 }
756
757 #[test]
758 fn test_csr_to_dense() {
759 let rows = vec![0, 0, 1, 2, 2];
761 let cols = vec![0, 2, 2, 0, 1];
762 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
763 let shape = (3, 3);
764
765 let matrix = CsrMatrix::new(data, rows, cols, shape).unwrap();
766 let dense = matrix.to_dense();
767
768 let expected = vec![
769 vec![1.0, 0.0, 2.0],
770 vec![0.0, 0.0, 3.0],
771 vec![4.0, 5.0, 0.0],
772 ];
773
774 assert_eq!(dense, expected);
775 }
776
777 #[test]
778 fn test_csr_dot() {
779 let rows = vec![0, 0, 1, 2, 2];
781 let cols = vec![0, 2, 2, 0, 1];
782 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
783 let shape = (3, 3);
784
785 let matrix = CsrMatrix::new(data, rows, cols, shape).unwrap();
786
787 let vec = vec![1.0, 2.0, 3.0];
793 let result = matrix.dot(&vec).unwrap();
794
795 let expected = [7.0, 9.0, 14.0];
800
801 assert_eq!(result.len(), expected.len());
802 for (a, b) in result.iter().zip(expected.iter()) {
803 assert_relative_eq!(a, b, epsilon = 1e-10);
804 }
805 }
806
807 #[test]
808 fn test_csr_transpose() {
809 let rows = vec![0, 0, 1, 2, 2];
811 let cols = vec![0, 2, 2, 0, 1];
812 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
813 let shape = (3, 3);
814
815 let matrix = CsrMatrix::new(data, rows, cols, shape).unwrap();
816 let transposed = matrix.transpose();
817
818 assert_eq!(transposed.shape(), (3, 3));
819 assert_eq!(transposed.nnz(), 5);
820
821 let dense = transposed.to_dense();
822 let expected = vec![
823 vec![1.0, 0.0, 4.0],
824 vec![0.0, 0.0, 5.0],
825 vec![2.0, 3.0, 0.0],
826 ];
827
828 assert_eq!(dense, expected);
829 }
830
831 #[test]
832 fn test_gpu_dot() {
833 let rows = vec![0, 0, 1, 2, 2];
835 let cols = vec![0, 2, 2, 0, 1];
836 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
837 let shape = (3, 3);
838
839 let matrix = CsrMatrix::new(data, rows, cols, shape).unwrap();
840 let vec = vec![1.0, 2.0, 3.0];
841
842 match matrix.gpu_dot(&vec) {
844 Ok(result) => {
845 let expected = [7.0, 9.0, 14.0];
846 assert_eq!(result.len(), expected.len());
847 for (a, b) in result.iter().zip(expected.iter()) {
848 assert_relative_eq!(a, b, epsilon = 1e-10);
849 }
850 }
851 Err(crate::error::SparseError::ComputationError(_))
852 | Err(crate::error::SparseError::OperationNotSupported(_)) => {
853 }
855 Err(e) => panic!("Unexpected error in GPU SpMV: {:?}", e),
856 }
857 }
858
859 #[test]
860 fn test_should_use_gpu() {
861 let small_matrix = CsrMatrix::new(vec![1.0, 2.0], vec![0, 1], vec![0, 1], (2, 2)).unwrap();
863 assert!(
864 !small_matrix.should_use_gpu(),
865 "Small matrix should not use GPU"
866 );
867
868 let large_data = vec![1.0; 15000];
870 let large_rows: Vec<usize> = (0..15000).collect();
871 let large_cols: Vec<usize> = (0..15000).collect();
872 let large_matrix =
873 CsrMatrix::new(large_data, large_rows, large_cols, (15000, 15000)).unwrap();
874 assert!(
875 large_matrix.should_use_gpu(),
876 "Large sparse matrix should use GPU"
877 );
878 }
879
880 #[test]
881 fn test_gpu_backend_info() {
882 let backend_info = CsrMatrix::<f64>::gpu_backend_info();
883 assert!(
884 backend_info.is_ok(),
885 "Should be able to get GPU backend info"
886 );
887
888 if let Ok((backend, name)) = backend_info {
889 assert!(!name.is_empty(), "Backend name should not be empty");
890 match backend {
892 crate::gpu_ops::GpuBackend::Cuda
893 | crate::gpu_ops::GpuBackend::OpenCL
894 | crate::gpu_ops::GpuBackend::Metal
895 | crate::gpu_ops::GpuBackend::Cpu
896 | crate::gpu_ops::GpuBackend::Rocm
897 | crate::gpu_ops::GpuBackend::Wgpu => {}
898 }
899 }
900 }
901
902 #[test]
903 fn test_gpu_dot_generic_f32() {
904 let rows = vec![0, 0, 1, 2, 2];
906 let cols = vec![0, 2, 2, 0, 1];
907 let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
908 let shape = (3, 3);
909
910 let matrix = CsrMatrix::new(data, rows, cols, shape).unwrap();
911 let vec = vec![1.0f32, 2.0, 3.0];
912
913 match matrix.gpu_dot_generic(&vec) {
914 Ok(result) => {
915 let expected = [7.0f32, 9.0, 14.0];
916 assert_eq!(result.len(), expected.len());
917 for (a, b) in result.iter().zip(expected.iter()) {
918 assert_relative_eq!(a, b, epsilon = 1e-6);
919 }
920 }
921 Err(crate::error::SparseError::ComputationError(_))
922 | Err(crate::error::SparseError::OperationNotSupported(_)) => {}
923 Err(e) => panic!("Unexpected error in generic GPU SpMV: {:?}", e),
924 }
925 }
926}