edge_rs/utils/
sparse.rs

1//! Compressed sparse storage for count matrices.
2//!
3//! Modelled on `CompressedSparseData` in `manifolds-rs` and
4//! `CompressedSparseData2` in `bixverse-rs`, with `u32` indices as in the
5//! latter, which halves the bandwidth of every scan.
6//!
7//! * [`CompressedSparse::transpose`] flips the format label and swaps the
8//!   shape. No data moves.
9//! * [`CompressedSparse::convert`] stores the same matrix in the other format.
10//!   It moves data.
11//!
12//! This keeps the `bixverse-rs` boundary cheap: its chunks arrive as CSC over
13//! `(n_cells, n_genes)`, which is bit-for-bit a CSR over `(n_genes, n_cells)`,
14//! the layout every algorithm here wants. `transpose` is the whole adapter.
15
16use crate::prelude::*;
17
18//////////////////
19// SparseFormat //
20//////////////////
21
22/// Which axis a compressed sparse matrix is compressed along.
23#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24pub enum SparseFormat {
25    /// Compressed sparse row: `indptr` walks rows, `indices` holds column
26    /// indices.
27    Csr,
28    /// Compressed sparse column: `indptr` walks columns, `indices` holds row
29    /// indices.
30    Csc,
31}
32
33impl SparseFormat {
34    /// The other format.
35    ///
36    /// ### Returns
37    ///
38    /// `Csc` for `Csr` and vice versa.
39    #[inline]
40    pub fn flipped(&self) -> Self {
41        match self {
42            SparseFormat::Csr => SparseFormat::Csc,
43            SparseFormat::Csc => SparseFormat::Csr,
44        }
45    }
46}
47
48//////////////////////
49// CompressedSparse //
50//////////////////////
51
52/// A sparse matrix in compressed row or column form.
53///
54/// Fields are public so buffers can go straight to a kernel, but only
55/// [`CompressedSparse::validate`] guarantees consistency. Prefer the
56/// constructors.
57#[derive(Clone, Debug)]
58pub struct CompressedSparse<T> {
59    /// Non-zero values, ordered by the compressed axis then by `indices`.
60    pub data: Vec<T>,
61    /// Index along the non-compressed axis for each value in `data`.
62    pub indices: Vec<u32>,
63    /// Offsets into `data`/`indices`, one per compressed slice plus a trailing total.
64    pub indptr: Vec<u32>,
65    /// Whether `indptr` walks rows or columns.
66    pub format: SparseFormat,
67    /// Logical shape as `(n_rows, n_cols)`, independent of `format`.
68    pub shape: (usize, usize),
69}
70
71impl<T: Copy> CompressedSparse<T> {
72    /// Assembles a matrix from its parts without copying.
73    ///
74    /// ### Params
75    ///
76    /// * `data` - Non-zero values
77    /// * `indices` - Index along the non-compressed axis per value
78    /// * `indptr` - Offsets per compressed slice, length `outer_len + 1`
79    /// * `format` - Which axis `indptr` walks
80    /// * `shape` - Logical `(n_rows, n_cols)`
81    ///
82    /// ### Returns
83    ///
84    /// The matrix, or [`EdgeErrors::MalformedSparse`] if the parts disagree.
85    pub fn from_parts(
86        data: Vec<T>,
87        indices: Vec<u32>,
88        indptr: Vec<u32>,
89        format: SparseFormat,
90        shape: (usize, usize),
91    ) -> Result<Self, EdgeErrors> {
92        let out = Self {
93            data,
94            indices,
95            indptr,
96            format,
97            shape,
98        };
99        out.validate()?;
100        Ok(out)
101    }
102
103    /// Checks the structural invariants.
104    ///
105    /// `data` and `indices` must match in length, `indptr` must have the right
106    /// length, be non-decreasing and end at the non-zero count, and every index
107    /// must be inside the non-compressed axis.
108    ///
109    /// ### Returns
110    ///
111    /// `Ok(())` when consistent, otherwise [`EdgeErrors::MalformedSparse`].
112    pub fn validate(&self) -> Result<(), EdgeErrors> {
113        if self.data.len() != self.indices.len() {
114            return Err(EdgeErrors::MalformedSparse(format!(
115                "data has {} values but indices has {}",
116                self.data.len(),
117                self.indices.len()
118            )));
119        }
120
121        let expected_indptr = self.outer_len() + 1;
122        if self.indptr.len() != expected_indptr {
123            return Err(EdgeErrors::MalformedSparse(format!(
124                "indptr has length {} but the {:?} layout of a {:?} matrix needs {}",
125                self.indptr.len(),
126                self.format,
127                self.shape,
128                expected_indptr
129            )));
130        }
131
132        if self.indptr[0] != 0 {
133            return Err(EdgeErrors::MalformedSparse(format!(
134                "indptr must start at 0, starts at {}",
135                self.indptr[0]
136            )));
137        }
138
139        for w in self.indptr.windows(2) {
140            if w[1] < w[0] {
141                return Err(EdgeErrors::MalformedSparse(format!(
142                    "indptr is not non-decreasing: {} follows {}",
143                    w[1], w[0]
144                )));
145            }
146        }
147
148        let nnz = *self.indptr.last().unwrap_or(&0) as usize;
149        if nnz != self.data.len() {
150            return Err(EdgeErrors::MalformedSparse(format!(
151                "indptr ends at {} but data has {} values",
152                nnz,
153                self.data.len()
154            )));
155        }
156
157        let inner = self.inner_len();
158        if let Some(&bad) = self.indices.iter().find(|&&i| i as usize >= inner) {
159            return Err(EdgeErrors::MalformedSparse(format!(
160                "index {bad} is out of range for a non-compressed axis of length {inner}"
161            )));
162        }
163
164        Ok(())
165    }
166
167    /// Number of rows.
168    ///
169    /// ### Returns
170    ///
171    /// The first component of the shape.
172    #[inline]
173    pub fn nrows(&self) -> usize {
174        self.shape.0
175    }
176
177    /// Number of columns.
178    ///
179    /// ### Returns
180    ///
181    /// The second component of the shape.
182    #[inline]
183    pub fn ncols(&self) -> usize {
184        self.shape.1
185    }
186
187    /// Number of stored non-zeros.
188    ///
189    /// ### Returns
190    ///
191    /// The length of `data`.
192    #[inline]
193    pub fn nnz(&self) -> usize {
194        self.data.len()
195    }
196
197    /// Length of the compressed axis, that is, how many slices `indptr` walks.
198    ///
199    /// ### Returns
200    ///
201    /// Row count for CSR, column count for CSC.
202    #[inline]
203    pub fn outer_len(&self) -> usize {
204        match self.format {
205            SparseFormat::Csr => self.shape.0,
206            SparseFormat::Csc => self.shape.1,
207        }
208    }
209
210    /// Length of the non-compressed axis, that is, the range `indices` spans.
211    ///
212    /// ### Returns
213    ///
214    /// Column count for CSR, row count for CSC.
215    #[inline]
216    pub fn inner_len(&self) -> usize {
217        match self.format {
218            SparseFormat::Csr => self.shape.1,
219            SparseFormat::Csc => self.shape.0,
220        }
221    }
222
223    /// Borrows one slice along the compressed axis.
224    ///
225    /// The per-gene accessor: for a gene-major matrix, that gene's non-zero
226    /// cells and counts, no copy.
227    ///
228    /// ### Params
229    ///
230    /// * `outer` - Index along the compressed axis
231    ///
232    /// ### Returns
233    ///
234    /// The indices and values of that slice, in `indices` order.
235    #[inline]
236    pub fn outer(&self, outer: usize) -> (&[u32], &[T]) {
237        let start = self.indptr[outer] as usize;
238        let end = self.indptr[outer + 1] as usize;
239        (&self.indices[start..end], &self.data[start..end])
240    }
241
242    /// The transposed matrix, for free.
243    ///
244    /// A CSR over `(m, n)` describes the same bytes as a CSC over `(n, m)`.
245    ///
246    /// ### Returns
247    ///
248    /// A matrix representing the transpose, sharing this one's layout.
249    pub fn transpose(self) -> Self {
250        Self {
251            data: self.data,
252            indices: self.indices,
253            indptr: self.indptr,
254            format: self.format.flipped(),
255            shape: (self.shape.1, self.shape.0),
256        }
257    }
258
259    /// The same matrix stored in the other format.
260    ///
261    /// A pass over the non-zeros plus a counting sort. Sequential; meant to be
262    /// called once at ingest.
263    ///
264    /// ### Returns
265    ///
266    /// An equivalent matrix with the opposite `format`.
267    pub fn convert(&self) -> Self
268    where
269        T: Default,
270    {
271        let outer_new = self.inner_len();
272        let nnz = self.nnz();
273
274        // Counting sort: how many entries land in each new compressed slice.
275        let mut counts = vec![0u32; outer_new + 1];
276        for &idx in &self.indices {
277            counts[idx as usize + 1] += 1;
278        }
279        for i in 0..outer_new {
280            counts[i + 1] += counts[i];
281        }
282        let indptr_new = counts.clone();
283
284        let mut data_new = vec![T::default(); nnz];
285        let mut indices_new = vec![0u32; nnz];
286        let mut cursor = counts;
287
288        for outer in 0..self.outer_len() {
289            let start = self.indptr[outer] as usize;
290            let end = self.indptr[outer + 1] as usize;
291            for k in start..end {
292                let inner = self.indices[k] as usize;
293                let dst = cursor[inner] as usize;
294                data_new[dst] = self.data[k];
295                indices_new[dst] = outer as u32;
296                cursor[inner] += 1;
297            }
298        }
299
300        Self {
301            data: data_new,
302            indices: indices_new,
303            indptr: indptr_new,
304            format: self.format.flipped(),
305            shape: self.shape,
306        }
307    }
308
309    /// Builds a compressed matrix from a dense row-major buffer.
310    ///
311    /// For tests and small inputs.
312    ///
313    /// ### Params
314    ///
315    /// * `dense` - `n_rows * n_cols` values, row-major
316    /// * `n_rows` - Number of rows
317    /// * `n_cols` - Number of columns
318    /// * `format` - Storage format to build
319    /// * `is_zero` - Predicate deciding which values to drop
320    ///
321    /// ### Returns
322    ///
323    /// The compressed matrix, or [`EdgeErrors::LengthMismatch`] if `dense` does
324    /// not match the stated shape.
325    pub fn from_dense<F>(
326        dense: &[T],
327        n_rows: usize,
328        n_cols: usize,
329        format: SparseFormat,
330        is_zero: F,
331    ) -> Result<Self, EdgeErrors>
332    where
333        F: Fn(&T) -> bool,
334    {
335        if dense.len() != n_rows * n_cols {
336            return Err(EdgeErrors::LengthMismatch {
337                name: "dense",
338                expected: n_rows * n_cols,
339                got: dense.len(),
340            });
341        }
342
343        let (outer_len, inner_len) = match format {
344            SparseFormat::Csr => (n_rows, n_cols),
345            SparseFormat::Csc => (n_cols, n_rows),
346        };
347
348        let mut data = Vec::new();
349        let mut indices = Vec::new();
350        let mut indptr = Vec::with_capacity(outer_len + 1);
351        indptr.push(0u32);
352
353        for o in 0..outer_len {
354            for i in 0..inner_len {
355                let flat = match format {
356                    SparseFormat::Csr => o * n_cols + i,
357                    SparseFormat::Csc => i * n_cols + o,
358                };
359                let v = dense[flat];
360                if !is_zero(&v) {
361                    data.push(v);
362                    indices.push(i as u32);
363                }
364            }
365            indptr.push(data.len() as u32);
366        }
367
368        Self::from_parts(data, indices, indptr, format, (n_rows, n_cols))
369    }
370
371    /// Materialises a dense row-major buffer.
372    ///
373    /// Tests and interop only.
374    ///
375    /// ### Params
376    ///
377    /// * `zero` - Value to fill the structural zeros with
378    ///
379    /// ### Returns
380    ///
381    /// An `n_rows * n_cols` row-major vector.
382    pub fn to_dense(&self, zero: T) -> Vec<T> {
383        let mut out = vec![zero; self.nrows() * self.ncols()];
384        for o in 0..self.outer_len() {
385            let (idx, vals) = self.outer(o);
386            for (&i, &v) in idx.iter().zip(vals.iter()) {
387                let flat = match self.format {
388                    SparseFormat::Csr => o * self.ncols() + i as usize,
389                    SparseFormat::Csc => i as usize * self.ncols() + o,
390                };
391                out[flat] = v;
392            }
393        }
394        out
395    }
396}
397
398///////////
399// Tests //
400///////////
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405
406    /// 3 by 4, five non-zeros, deliberately including an all-zero row.
407    fn dense_fixture() -> (Vec<f64>, usize, usize) {
408        let dense = vec![
409            1.0, 0.0, 2.0, 0.0, //
410            0.0, 0.0, 0.0, 0.0, //
411            0.0, 3.0, 4.0, 5.0, //
412        ];
413        (dense, 3, 4)
414    }
415
416    #[test]
417    fn test_from_dense_csr_round_trips() {
418        let (dense, r, c) = dense_fixture();
419        let m =
420            CompressedSparse::from_dense(&dense, r, c, SparseFormat::Csr, |v| *v == 0.0).unwrap();
421        assert_eq!(m.nnz(), 5);
422        assert_eq!(m.to_dense(0.0), dense);
423    }
424
425    #[test]
426    fn test_from_dense_csc_round_trips() {
427        let (dense, r, c) = dense_fixture();
428        let m =
429            CompressedSparse::from_dense(&dense, r, c, SparseFormat::Csc, |v| *v == 0.0).unwrap();
430        assert_eq!(m.nnz(), 5);
431        assert_eq!(m.to_dense(0.0), dense);
432    }
433
434    #[test]
435    fn test_convert_preserves_the_matrix() {
436        let (dense, r, c) = dense_fixture();
437        let csr =
438            CompressedSparse::from_dense(&dense, r, c, SparseFormat::Csr, |v| *v == 0.0).unwrap();
439        let csc = csr.convert();
440        assert_eq!(csc.format, SparseFormat::Csc);
441        assert_eq!(csc.shape, (r, c));
442        assert_eq!(csc.to_dense(0.0), dense);
443        // Round trip.
444        assert_eq!(csc.convert().to_dense(0.0), dense);
445    }
446
447    /// A CSR over (m, n) is a CSC over (n, m) with the same bytes (the
448    /// bixverse-rs boundary).
449    #[test]
450    fn test_transpose_is_a_relabel_not_a_copy() {
451        let (dense, r, c) = dense_fixture();
452        let csr =
453            CompressedSparse::from_dense(&dense, r, c, SparseFormat::Csr, |v| *v == 0.0).unwrap();
454        let data_before = csr.data.clone();
455        let t = csr.transpose();
456
457        assert_eq!(t.format, SparseFormat::Csc);
458        assert_eq!(t.shape, (c, r));
459        assert_eq!(t.data, data_before);
460
461        // Transposing a 3x4 gives a 4x3 whose (i, j) is the original (j, i).
462        let dense_t = t.to_dense(0.0);
463        for i in 0..c {
464            for j in 0..r {
465                assert_eq!(dense_t[i * r + j], dense[j * c + i]);
466            }
467        }
468    }
469
470    #[test]
471    fn test_outer_borrows_a_gene_slice() {
472        let (dense, r, c) = dense_fixture();
473        let csr =
474            CompressedSparse::from_dense(&dense, r, c, SparseFormat::Csr, |v| *v == 0.0).unwrap();
475
476        let (idx, vals) = csr.outer(0);
477        assert_eq!(idx, &[0, 2]);
478        assert_eq!(vals, &[1.0, 2.0]);
479
480        // The empty row is a valid empty slice.
481        let (idx, vals) = csr.outer(1);
482        assert!(idx.is_empty());
483        assert!(vals.is_empty());
484
485        let (idx, vals) = csr.outer(2);
486        assert_eq!(idx, &[1, 2, 3]);
487        assert_eq!(vals, &[3.0, 4.0, 5.0]);
488    }
489
490    #[test]
491    fn test_validate_rejects_out_of_range_index() {
492        let err = CompressedSparse::from_parts(
493            vec![1.0_f64],
494            vec![99],
495            vec![0, 1, 1, 1],
496            SparseFormat::Csr,
497            (3, 4),
498        )
499        .unwrap_err();
500        assert!(matches!(err, EdgeErrors::MalformedSparse(_)));
501    }
502
503    #[test]
504    fn test_validate_rejects_bad_indptr_length() {
505        let err = CompressedSparse::from_parts(
506            vec![1.0_f64],
507            vec![0],
508            vec![0, 1],
509            SparseFormat::Csr,
510            (3, 4),
511        )
512        .unwrap_err();
513        assert!(matches!(err, EdgeErrors::MalformedSparse(_)));
514    }
515
516    #[test]
517    fn test_convert_handles_empty_matrix() {
518        let m: CompressedSparse<f64> =
519            CompressedSparse::from_parts(vec![], vec![], vec![0; 4], SparseFormat::Csr, (3, 4))
520                .unwrap();
521        let c = m.convert();
522        assert_eq!(c.nnz(), 0);
523        assert_eq!(c.indptr.len(), 5);
524        assert_eq!(c.to_dense(0.0), vec![0.0; 12]);
525    }
526}
edge_rs/utils/sparse.rs

edge_rs/utils/
sparse.rs
