powerio_matrix/matrix/
triplet.rs1use std::collections::HashMap;
10
11use sprs::{CsMat, TriMat};
12
13#[derive(Debug, Clone)]
14pub struct CooBuilder {
15 rows: usize,
16 cols: usize,
17 entries: HashMap<(usize, usize), f64>,
18}
19
20impl CooBuilder {
21 pub fn new(n: usize) -> Self {
23 Self::new_rect(n, n)
24 }
25
26 pub fn with_capacity(n: usize, capacity: usize) -> Self {
28 Self::with_capacity_rect(n, n, capacity)
29 }
30
31 pub fn new_rect(rows: usize, cols: usize) -> Self {
33 Self {
34 rows,
35 cols,
36 entries: HashMap::new(),
37 }
38 }
39
40 pub fn with_capacity_rect(rows: usize, cols: usize, capacity: usize) -> Self {
42 Self {
43 rows,
44 cols,
45 entries: HashMap::with_capacity(capacity),
46 }
47 }
48
49 #[inline]
51 pub fn n(&self) -> usize {
52 self.rows
53 }
54
55 #[inline]
57 pub fn shape(&self) -> (usize, usize) {
58 (self.rows, self.cols)
59 }
60
61 #[inline]
63 pub fn add(&mut self, i: usize, j: usize, v: f64) {
64 if v == 0.0 {
65 return;
66 }
67 debug_assert!(i < self.rows && j < self.cols);
68 *self.entries.entry((i, j)).or_insert(0.0) += v;
69 }
70
71 #[inline]
74 pub fn add_sym(&mut self, i: usize, j: usize, v: f64) {
75 if i == j {
76 self.add(i, j, v);
77 } else {
78 self.add(i, j, v);
79 self.add(j, i, v);
80 }
81 }
82
83 pub fn finish_csr(self) -> CsMat<f64> {
85 let mut tri = TriMat::with_capacity((self.rows, self.cols), self.entries.len());
86 for ((i, j), v) in self.entries {
87 if v != 0.0 {
88 tri.add_triplet(i, j, v);
89 }
90 }
91 tri.to_csr()
92 }
93
94 pub fn finish_csc(self) -> CsMat<f64> {
96 self.finish_csr().to_csc()
97 }
98}