powerio_matrix/matrix/
triplet.rs

1//! `CooBuilder` — a small HashMap backed coordinate format accumulator.
2//! Deduplicates `(i, j)` entries on insert (each `add` is O(1) amortized,
3//! independent of `nnz`). Replaces the previous Vec linear scan
4//! accumulator, which was O(nnz²) per case.
5//!
6//! Square by default (`new`), rectangular via `new_rect` for the incidence,
7//! flow map, and generator→bus matrices.
8
9use rustc_hash::{FxBuildHasher, FxHashMap};
10use sprs::{CsMat, TriMat};
11
12type CoordinateMap = FxHashMap<usize, f64>;
13
14#[derive(Debug, Clone)]
15pub struct CooBuilder {
16    rows: usize,
17    cols: usize,
18    entries: CoordinateMap,
19}
20
21impl CooBuilder {
22    /// Square `n × n` accumulator.
23    pub fn new(n: usize) -> Self {
24        Self::new_rect(n, n)
25    }
26
27    /// Square `n × n` accumulator with a pre-sized entry table.
28    pub fn with_capacity(n: usize, capacity: usize) -> Self {
29        Self::with_capacity_rect(n, n, capacity)
30    }
31
32    /// Rectangular `rows × cols` accumulator.
33    pub fn new_rect(rows: usize, cols: usize) -> Self {
34        Self {
35            rows,
36            cols,
37            entries: CoordinateMap::default(),
38        }
39    }
40
41    /// Rectangular `rows × cols` accumulator with a pre-sized entry table.
42    pub fn with_capacity_rect(rows: usize, cols: usize, capacity: usize) -> Self {
43        Self {
44            rows,
45            cols,
46            entries: CoordinateMap::with_capacity_and_hasher(capacity, FxBuildHasher),
47        }
48    }
49
50    /// Side length for a square builder (row count in general).
51    #[inline]
52    pub fn n(&self) -> usize {
53        self.rows
54    }
55
56    /// `(rows, cols)`.
57    #[inline]
58    pub fn shape(&self) -> (usize, usize) {
59        (self.rows, self.cols)
60    }
61
62    /// Accumulate `v` into entry `(i, j)`. Skips the insert if `v == 0.0`.
63    ///
64    /// # Panics
65    /// Panics if `(i, j)` is outside the matrix shape or the packed coordinate
66    /// key overflows `usize`.
67    #[inline]
68    pub fn add(&mut self, i: usize, j: usize, v: f64) {
69        if v == 0.0 {
70            return;
71        }
72        assert!(
73            i < self.rows && j < self.cols,
74            "COO coordinate ({i}, {j}) out of bounds for shape {}x{}",
75            self.rows,
76            self.cols
77        );
78        let key = i
79            .checked_mul(self.cols)
80            .and_then(|base| base.checked_add(j))
81            .expect("COO matrix dimensions overflow usize");
82        *self.entries.entry(key).or_insert(0.0) += v;
83    }
84
85    /// Symmetrically accumulate `v` into both `(i, j)` and `(j, i)`. Square
86    /// builders only.
87    #[inline]
88    pub fn add_sym(&mut self, i: usize, j: usize, v: f64) {
89        if i == j {
90            self.add(i, j, v);
91        } else {
92            self.add(i, j, v);
93            self.add(j, i, v);
94        }
95    }
96
97    /// Materialize as a `CsMat<f64>` (CSR) with explicit zeros pruned.
98    pub fn finish_csr(self) -> CsMat<f64> {
99        let mut tri = TriMat::with_capacity((self.rows, self.cols), self.entries.len());
100        for (key, v) in self.entries {
101            if v != 0.0 {
102                let i = key / self.cols;
103                let j = key % self.cols;
104                tri.add_triplet(i, j, v);
105            }
106        }
107        tri.to_csr()
108    }
109
110    /// Materialize as a CSC matrix.
111    pub fn finish_csc(self) -> CsMat<f64> {
112        self.finish_csr().to_csc()
113    }
114}
powerio_matrix/matrix/triplet.rs

powerio_matrix/matrix/
triplet.rs
