Skip to main content

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 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    /// Square `n × n` accumulator.
22    pub fn new(n: usize) -> Self {
23        Self::new_rect(n, n)
24    }
25
26    /// Square `n × n` accumulator with a pre-sized entry table.
27    pub fn with_capacity(n: usize, capacity: usize) -> Self {
28        Self::with_capacity_rect(n, n, capacity)
29    }
30
31    /// Rectangular `rows × cols` accumulator.
32    pub fn new_rect(rows: usize, cols: usize) -> Self {
33        Self {
34            rows,
35            cols,
36            entries: HashMap::new(),
37        }
38    }
39
40    /// Rectangular `rows × cols` accumulator with a pre-sized entry table.
41    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    /// Side length for a square builder (row count in general).
50    #[inline]
51    pub fn n(&self) -> usize {
52        self.rows
53    }
54
55    /// `(rows, cols)`.
56    #[inline]
57    pub fn shape(&self) -> (usize, usize) {
58        (self.rows, self.cols)
59    }
60
61    /// Accumulate `v` into entry `(i, j)`. Skips the insert if `v == 0.0`.
62    #[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    /// Symmetrically accumulate `v` into both `(i, j)` and `(j, i)`. Square
72    /// builders only.
73    #[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    /// Materialize as a `CsMat<f64>` (CSR) with explicit zeros pruned.
84    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    /// Materialize as a CSC matrix.
95    pub fn finish_csc(self) -> CsMat<f64> {
96        self.finish_csr().to_csc()
97    }
98}