OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 34 Wiki Activity

Add sparse storage variant to PersistentCompactIntMatrix

Browse Source 
Introduce a new `Sparse` format alongside existing `Columnar` and `Packed` variants, enabling optimized row-major pairwise counting for distance and similarity metrics via the `CountPartials` trait. Update storage detection priorities, extend matrix dispatch logic to sparse backends, and correct diagonal/off-diagonal formulas in bit matrix partial computations. Expand layer APIs with format-agnostic `nonzero_iter`, update usage documentation for the `--sparse` flag, and add comprehensive tests verifying roundtrip integrity and metric equivalence against dense implementations.
This commit is contained in:
Eric Coissac
2026-08-28 19:06:39 +02:00
parent 904d85f33b
commit b0890d1781
13 changed files with 751 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obicompactvec/src/bitmatrix/sparse.rs
+31 -4
View File
@@ -525,26 +525,53 @@ impl ColumnWeights for PersistentSparseBitMatrix {
}
impl BitPartials for PersistentSparseBitMatrix {
/// `col_weights_and_pair_counts`'s own `inter` never pairs a column with
/// itself (its dictionary decode loop skips same-index pairs by
/// construction), so its diagonal is `0`, not a genuine
/// self-intersection. Dense `partial_jaccard_dist_matrix` (`columnar.rs`/
/// `packed.rs`, via `pairwise2_matrix`) diagonal is `col(i).
/// partial_jaccard_dist(col(i))` — i.e. `f(i,i) = (col_weights[i],
/// col_weights[i])`, a full self-intersection/self-union — so that's the
/// diagonal patched in here, off the raw `col_weights_and_pair_counts`
/// result rather than the wrong `2·col_weights[i]` the generic
/// off-diagonal formula would otherwise produce.
fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
let n = self.n_cols();
let (col_weights, inter) = self.col_weights_and_pair_counts();
let (col_weights, mut inter) = self.col_weights_and_pair_counts();
let mut union = Array2::zeros((n, n));
for i in 0..n {
for j in 0..n {
union[[i, j]] = col_weights[i] + col_weights[j] - inter[[i, j]];
if i != j {
union[[i, j]] = col_weights[i] + col_weights[j] - inter[[i, j]];
}
}
}
for i in 0..n {
inter[[i, i]] = col_weights[i];
union[[i, i]] = col_weights[i];
}
(inter, union)
}
/// Hamming distance is the symmetric-difference size: `|A|+|B|-2|A∩B|`
/// (`total - union`, the previous formula here, is the count of rows
/// where *neither* column is present — a different, unrelated quantity
/// that only coincides with Hamming distance when `col_weights[i] +
/// col_weights[j] == total`; verified against dense
/// `partial_hamming_dist_matrix` on real asymmetric-presence data, see
/// `tests::sparse::partial_jaccard_and_hamming_match_dense_including_diagonal`).
/// Diagonal is `col(i).hamming_dist(col(i))`, i.e. `0` — a column XORed
/// with itself, matching dense — left at `Array2::zeros`'s default via
/// the `i != j` guard rather than computed through the general formula.
fn partial_hamming(&self) -> Array2<u64> {
let n = self.n_cols();
let (col_weights, inter) = self.col_weights_and_pair_counts();
let total = self.n() as u64;
let mut m = Array2::zeros((n, n));
for i in 0..n {
for j in 0..n {
m[[i, j]] = total - (col_weights[i] + col_weights[j] - inter[[i, j]]);
if i != j {
m[[i, j]] = col_weights[i] + col_weights[j] - 2 * inter[[i, j]];
}
}
}
m
src/obicompactvec/src/intmatrix.rs
+86 -35
View File
@@ -11,6 +11,7 @@ use crate::colgroup::{chunked_presence_count, ColGroup, MatrixGroupOps};
use crate::format::{HEADER_SIZE, OVERFLOW_ENTRY_SIZE};
use crate::meta::MatrixMeta;
use crate::reader::PersistentCompactIntVec;
use crate::sparse_intmatrix::{is_present as sparse_is_present, PersistentSparseCompactIntMatrix};
use crate::storage_kind::StorageKind;
use crate::tempbitvec::{TempBitVec, TempBitVecBuilder};
use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
@@ -271,9 +272,14 @@ pub fn pack_compact_int_matrix(dir: &Path) -> io::Result<()> {
pub enum PersistentCompactIntMatrix {
Columnar(ColumnarCompactIntMatrix),
Packed(PackedCompactIntMatrix),
Sparse(PersistentSparseCompactIntMatrix),
}
impl PersistentCompactIntMatrix {
/// Checks (in order): `counts/matrix.pcmx` → Packed; `counts/meta.json`
/// → Columnar; `counts/singleton_values.pciv` → Sparse. Mirrors
/// `PersistentBitMatrix::open`'s own Packed → Columnar → Sparse
/// priority order.
pub fn open(layer_dir: &Path) -> io::Result<Self> {
let counts_dir = layer_dir.join("counts");
if counts_dir.join("matrix.pcmx").exists() {
@@ -282,6 +288,9 @@ impl PersistentCompactIntMatrix {
if MatrixMeta::load(&counts_dir).is_ok() {
return Ok(Self::Columnar(ColumnarCompactIntMatrix::open(&counts_dir)?));
}
if sparse_is_present(&counts_dir) {
return Ok(Self::Sparse(PersistentSparseCompactIntMatrix::open(&counts_dir)?));
}
Err(io::Error::new(
io::ErrorKind::NotFound,
format!("no count matrix found in {} — run 'obikmer upgrade'", layer_dir.display()),
@@ -295,6 +304,7 @@ impl PersistentCompactIntMatrix {
match self {
Self::Columnar(_) => StorageKind::Columnar,
Self::Packed(_) => StorageKind::Packed,
Self::Sparse(_) => StorageKind::Sparse,
}
}
@@ -309,6 +319,9 @@ impl PersistentCompactIntMatrix {
if MatrixMeta::load(&counts_dir).is_ok() {
return Ok(StorageKind::Columnar);
}
if sparse_is_present(&counts_dir) {
return Ok(StorageKind::Sparse);
}
Err(io::Error::new(
io::ErrorKind::NotFound,
format!("no count matrix found in {}", layer_dir.display()),
@@ -317,11 +330,11 @@ impl PersistentCompactIntMatrix {
#[inline]
pub fn n(&self) -> usize {
match self { Self::Columnar(m) => m.n(), Self::Packed(m) => m.n_rows }
match self { Self::Columnar(m) => m.n(), Self::Packed(m) => m.n_rows, Self::Sparse(m) => m.n() }
}
#[inline]
pub fn n_cols(&self) -> usize {
match self { Self::Columnar(m) => m.n_cols(), Self::Packed(m) => m.n_cols }
match self { Self::Columnar(m) => m.n_cols(), Self::Packed(m) => m.n_cols, Self::Sparse(m) => m.n_cols() }
}
#[inline]
@@ -337,6 +350,7 @@ impl PersistentCompactIntMatrix {
match self {
Self::Columnar(m) => m.col(c).view(),
Self::Packed(m) => m.col_view(c),
Self::Sparse(_) => panic!("col_view() not available on Sparse PersistentCompactIntMatrix"),
}
}
@@ -344,16 +358,22 @@ impl PersistentCompactIntMatrix {
match self {
Self::Columnar(m) => PersistentCompactIntVecBuilder::build_from(m.col(c), path),
Self::Packed(m) => m.col_persist(c, path),
Self::Sparse(_) => Err(io::Error::new(io::ErrorKind::Unsupported,
"col_persist not available on Sparse PersistentCompactIntMatrix")),
}
}
#[inline]
pub fn row(&self, slot: usize) -> Box<[u32]> {
match self { Self::Columnar(m) => m.row(slot), Self::Packed(m) => m.row(slot) }
match self { Self::Columnar(m) => m.row(slot), Self::Packed(m) => m.row(slot), Self::Sparse(m) => m.row(slot) }
}
#[inline]
pub fn fill_row(&self, slot: usize, buf: &mut [u32]) {
match self { Self::Columnar(m) => m.fill_row(slot, buf), Self::Packed(m) => m.fill_row(slot, buf) }
match self {
Self::Columnar(m) => m.fill_row(slot, buf),
Self::Packed(m) => m.fill_row(slot, buf),
Self::Sparse(m) => m.fill_row(slot, buf),
}
}
/// Extract a sub-matrix containing only the rows at `slots`.
@@ -374,67 +394,98 @@ impl PersistentCompactIntMatrix {
/// `out` must have length `self.n_cols()`. Each `out[c]` is cleared,
/// resized to `slots.len()`, and filled with the values for column `c`
/// in the same order as `slots`.
///
/// Derived from [`nonzero_iter`](Self::nonzero_iter) — one traversal per
/// format, no separate `Sparse`-specific scatter loop duplicated here.
/// Mirrors `PersistentBitMatrix::fill_sub_matrix`'s own reasoning (see
/// that method's doc comment).
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<u32>]) {
assert_eq!(out.len(), self.n_cols());
let n = slots.len();
if n == 0 {
for col in out.iter_mut() { col.clear(); }
return;
for col in out.iter_mut() {
col.clear();
col.resize(slots.len(), 0);
}
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted_slots: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
for (c, col) in out.iter_mut().enumerate() {
col.resize(n, 0);
for (i, v) in self.col_view(c).enumerate_slots_values(&sorted_slots) {
col[perm[i]] = v;
for (i, c, v) in self.nonzero_iter(slots) {
out[c][i] = v;
}
}
/// Yields every nonzero `(idx into slots, col, value)` triple.
///
/// One primitive per format, each choosing its own natural traversal —
/// same shape as `PersistentBitMatrix::nonzero_iter`: `Sparse` delegates
/// to its own row-major decode (no `n_cols`-wide buffer, ever);
/// `Columnar`/`Packed` reuse the shared sorted-slot batching in
/// `views::nonzero_triples`. Boxed, not `impl Iterator`, because the
/// match arms are genuinely different concrete types.
pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> Box<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
match self {
Self::Sparse(m) => Box::new(m.nonzero_iter(slots)),
Self::Columnar(_) | Self::Packed(_) => {
Box::new(nonzero_triples(slots, self.n_cols(), |c| self.col_view(c), false))
}
}
}
/// Yields every nonzero `(idx into slots, col, value)` triple. Same
/// primitive as `PersistentBitMatrix::nonzero_iter` (see
/// `DevDocMD/architecture/siblings.md`, "`query` never benefits from
/// sparse row-major access" — counts explicitly not excluded from that
/// design even though no sparse count format exists yet). No native
/// low-effort case here the way `PersistentSparseBitMatrix` has one —
/// both variants reuse the shared sorted-slot batching in
/// `views::nonzero_triples`.
pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> impl Iterator<Item = (usize, usize, u32)> + 'a {
nonzero_triples(slots, self.n_cols(), |c| self.col_view(c), false)
}
#[inline]
pub fn sum(&self) -> Array1<u64> {
match self { Self::Columnar(m) => m.sum(), Self::Packed(m) => m.sum() }
match self { Self::Columnar(m) => m.sum(), Self::Packed(m) => m.sum(), Self::Sparse(m) => m.sum() }
}
#[inline]
pub fn count_nonzero(&self) -> Array1<u64> {
match self { Self::Columnar(m) => m.count_nonzero(), Self::Packed(m) => m.count_nonzero() }
match self {
Self::Columnar(m) => m.count_nonzero(),
Self::Packed(m) => m.count_nonzero(),
Self::Sparse(m) => m.count_nonzero(),
}
}
#[inline]
pub fn partial_bray_dist_matrix(&self) -> Array2<u64> {
match self { Self::Columnar(m) => m.partial_bray_dist_matrix(), Self::Packed(m) => m.partial_bray_dist_matrix() }
match self {
Self::Columnar(m) => m.partial_bray_dist_matrix(),
Self::Packed(m) => m.partial_bray_dist_matrix(),
Self::Sparse(m) => CountPartials::partial_bray(m),
}
}
#[inline]
pub fn partial_euclidean_dist_matrix(&self) -> Array2<f64> {
match self { Self::Columnar(m) => m.partial_euclidean_dist_matrix(), Self::Packed(m) => m.partial_euclidean_dist_matrix() }
match self {
Self::Columnar(m) => m.partial_euclidean_dist_matrix(),
Self::Packed(m) => m.partial_euclidean_dist_matrix(),
Self::Sparse(m) => CountPartials::partial_euclidean(m),
}
}
#[inline]
pub fn partial_threshold_jaccard_dist_matrix(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
match self { Self::Columnar(m) => m.partial_threshold_jaccard_dist_matrix(threshold), Self::Packed(m) => m.partial_threshold_jaccard_dist_matrix(threshold) }
match self {
Self::Columnar(m) => m.partial_threshold_jaccard_dist_matrix(threshold),
Self::Packed(m) => m.partial_threshold_jaccard_dist_matrix(threshold),
Self::Sparse(m) => CountPartials::partial_threshold_jaccard(m, threshold),
}
}
#[inline]
pub fn partial_relfreq_bray_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
match self { Self::Columnar(m) => m.partial_relfreq_bray_dist_matrix(col_sums), Self::Packed(m) => m.partial_relfreq_bray_dist_matrix(col_sums) }
match self {
Self::Columnar(m) => m.partial_relfreq_bray_dist_matrix(col_sums),
Self::Packed(m) => m.partial_relfreq_bray_dist_matrix(col_sums),
Self::Sparse(m) => CountPartials::partial_relfreq_bray(m, col_sums),
}
}
#[inline]
pub fn partial_relfreq_euclidean_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
match self { Self::Columnar(m) => m.partial_relfreq_euclidean_dist_matrix(col_sums), Self::Packed(m) => m.partial_relfreq_euclidean_dist_matrix(col_sums) }
match self {
Self::Columnar(m) => m.partial_relfreq_euclidean_dist_matrix(col_sums),
Self::Packed(m) => m.partial_relfreq_euclidean_dist_matrix(col_sums),
Self::Sparse(m) => CountPartials::partial_relfreq_euclidean(m, col_sums),
}
}
#[inline]
pub fn partial_hellinger_euclidean_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
match self { Self::Columnar(m) => m.partial_hellinger_euclidean_dist_matrix(col_sums), Self::Packed(m) => m.partial_hellinger_euclidean_dist_matrix(col_sums) }
match self {
Self::Columnar(m) => m.partial_hellinger_euclidean_dist_matrix(col_sums),
Self::Packed(m) => m.partial_hellinger_euclidean_dist_matrix(col_sums),
Self::Sparse(m) => CountPartials::partial_hellinger(m, col_sums),
}
}
#[inline]
pub fn append_column(dir: &Path, value_of: impl Fn(usize) -> u32) -> io::Result<()> {
src/obicompactvec/src/sparse_intmatrix.rs
+244 -11
View File
@@ -28,19 +28,28 @@ use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use ndarray::Array1;
use ndarray::{Array1, Array2};
use crate::bitmatrix::RowRank;
use crate::builder::PersistentCompactIntVecBuilder;
use crate::eliasfano::{EliasFano, EliasFanoBuilder};
use crate::reader::PersistentCompactIntVec;
use crate::traits::ColumnWeights;
use crate::traits::{BitPartials, ColumnWeights, CountPartials};
use crate::{PersistentCompactIntMatrix, PersistentSparseBitMatrix, PersistentSparseBitMatrixBuilder};
fn singleton_values_path(dir: &Path) -> PathBuf { dir.join("singleton_values.pciv") }
fn multi_values_path(dir: &Path) -> PathBuf { dir.join("multi_values.pciv") }
fn multi_offsets_base(dir: &Path) -> PathBuf { dir.join("multi_offsets") }
/// Lightweight disk probe: true iff a sparse count matrix has already been
/// written at `dir` — the marker [`PersistentCompactIntMatrix::open`]/
/// `detect_storage` check to route into this format, mirroring the
/// `sparse_meta.json` check the bit side does for `PersistentSparseBitMatrix`.
/// Also backs [`pack_sparse_compact_int_matrix`]'s own idempotency check.
pub(crate) fn is_present(dir: &Path) -> bool {
singleton_values_path(dir).exists()
}
/// Row batch size for [`PersistentSparseCompactIntMatrixBuilder::build_from_dense`].
const BUILD_FROM_DENSE_BATCH: usize = 4096;
@@ -153,14 +162,7 @@ impl PersistentSparseCompactIntMatrix {
})
}
/// Per-column sum, via a naive row-first accumulation. Column-pairwise
/// distance matrices (`CountPartials`) are explicitly deferred here,
/// same as `BitPartials` was initially for
/// `PersistentSparseBitMatrix`: unlike the support, values aren't
/// deduplicated across rows, so the dict-driven co-occurrence shortcut
/// (`col_weights_and_pair_counts`) doesn't carry over — a row-major
/// pairwise rewrite for counts is future work, not part of this type
/// yet.
/// Per-column sum, via a naive row-first accumulation.
pub fn sum(&self) -> Array1<u64> {
let mut sums = vec![0u64; self.n_cols()];
for slot in 0..self.n() {
@@ -169,6 +171,20 @@ impl PersistentSparseCompactIntMatrix {
Array1::from(sums)
}
/// Per-column sum of squares — the marginal `partial_euclidean`/
/// `partial_relfreq_euclidean`/`partial_hellinger` reconstruct their
/// squared-difference formulas from, via `Σ(a-b)² = Σa²+Σb²-2Σab` (see
/// [`row_major_pairwise`](Self::row_major_pairwise)'s doc comment for
/// why that identity is what makes the row-major approach correct for
/// those three).
fn sum_sq(&self) -> Array1<u64> {
let mut sums = vec![0u64; self.n_cols()];
for slot in 0..self.n() {
self.for_each_cell_in_row(slot, |g, v| sums[g] += (v as u64) * (v as u64));
}
Array1::from(sums)
}
pub fn count_nonzero(&self) -> Array1<u64> {
let mut counts = vec![0u64; self.n_cols()];
for slot in 0..self.n() {
@@ -176,6 +192,71 @@ impl PersistentSparseCompactIntMatrix {
}
Array1::from(counts)
}
/// Per-column count of values `>= threshold` — the marginal
/// `partial_threshold_jaccard` needs both for its `union` formula
/// (`count_geq[i] + count_geq[j] - inter[i,j]`) and for its diagonal.
/// `threshold == 1` is `count_nonzero()` exactly (`v >= 1 ⟺ v != 0`),
/// but this is only ever called with `threshold >= 1` by
/// [`partial_threshold_jaccard`](CountPartials::partial_threshold_jaccard)
/// (the `threshold == 0` case is closed-form there, no marginal needed).
fn count_geq(&self, threshold: u32) -> Array1<u64> {
let mut counts = vec![0u64; self.n_cols()];
for slot in 0..self.n() {
self.for_each_cell_in_row(slot, |g, v| if v >= threshold { counts[g] += 1; });
}
Array1::from(counts)
}
/// Shared row-major traversal behind every `CountPartials` formula
/// below: decodes each row once via
/// [`for_each_cell_in_row`](Self::for_each_cell_in_row) — the same
/// decode `sum`/`count_nonzero` already use — and for every pair of
/// columns actually co-present in that row, folds `kernel(i, a, j, b)`
/// into both `[i,j]` and `[j,i]` of the returned `n_cols × n_cols`
/// matrix. `kernel` must be symmetric in its (index, value) pair
/// arguments (every kernel used below is).
///
/// Unlike [`PersistentSparseBitMatrix::col_weights_and_pair_counts`],
/// this can't weight by how many rows share a `dict_id`: two rows can
/// share the same column *set* while carrying different *values* (see
/// this module's own doc comment), so every row still costs its own
/// `O(k̄²)` pass. What it still buys over the naive `O(n_cols² × n)`
/// column-pair scan: total cost is `O(Σ k̄²)` over actually-populated
/// rows, not `O(n_cols²)` full-column rescans — the same complexity
/// class the bit side reaches, just without the dict-multiplicity
/// shortcut.
///
/// Kernels here are all zero whenever either side is absent (`min`,
/// `a·b`, `√(a·b)`, a `threshold`-AND indicator) — cells absent from a
/// row's decode correctly contribute nothing, no separate handling
/// needed. Squared-difference formulas (`euclidean`/`relfreq_euclidean`/
/// `hellinger`) do *not* have that property — `(a-0)² = a² ≠ 0` — so
/// they never use this loop directly for their asymmetric-presence
/// contributions; they use it only for the `Σab`/`Σ√(ab)` cross term and
/// reconstruct the rest from per-column marginals (`sum`/`sum_sq`) at
/// the call site.
fn row_major_pairwise<T: Copy + Default + std::ops::AddAssign>(
&self,
mut kernel: impl FnMut(usize, u32, usize, u32) -> T,
) -> Array2<T> {
let n_cols = self.n_cols();
let mut acc = Array2::<T>::default((n_cols, n_cols));
let mut present: Vec<(usize, u32)> = Vec::new();
for slot in 0..self.n() {
present.clear();
self.for_each_cell_in_row(slot, |c, v| present.push((c, v)));
for p in 0..present.len() {
let (i, a) = present[p];
for &(j, b) in &present[p + 1..] {
let t = kernel(i, a, j, b);
acc[[i, j]] += t;
acc[[j, i]] += t;
}
}
}
acc
}
}
impl ColumnWeights for PersistentSparseCompactIntMatrix {
@@ -185,6 +266,158 @@ impl ColumnWeights for PersistentSparseCompactIntMatrix {
fn partial_kmer_counts(&self) -> Array1<u64> { self.count_nonzero() }
}
impl CountPartials for PersistentSparseCompactIntMatrix {
/// `Σ min(a,b)` — zero whenever either side is absent, so
/// [`row_major_pairwise`](Self::row_major_pairwise) needs no correction
/// off-diagonal. Diagonal is `min(a,a) = a`, i.e. the column's own sum —
/// not `T::default()` — matching the dense `pairwise_matrix` convention
/// (`f(i,i)` is a genuine self-comparison, see that helper's doc
/// comment); `row_major_pairwise` never calls its kernel with `i == j`
/// (a row's decode never lists the same column twice), so it's patched
/// in here explicitly.
fn partial_bray(&self) -> Array2<u64> {
let mut m = self.row_major_pairwise(|_, a, _, b| (a as u64).min(b as u64));
let sum = self.sum();
for i in 0..self.n_cols() { m[[i, i]] = sum[i]; }
m
}
/// `Σ(a-b)²`, reconstructed via `Σa²+Σb²-2Σab` — the `Σab` cross term is
/// the one part of this formula that's zero when either side is absent
/// (a row can't multiply an absent cell into anything), so it's the
/// only part [`row_major_pairwise`](Self::row_major_pairwise) computes;
/// the `Σa²`/`Σb²` marginals come from [`sum_sq`](Self::sum_sq).
fn partial_euclidean(&self) -> Array2<f64> {
let n = self.n_cols();
let dot = self.row_major_pairwise(|_, a, _, b| (a as f64) * (b as f64));
let sq = self.sum_sq();
let mut m = Array2::<f64>::zeros((n, n));
for i in 0..n {
for j in 0..n {
if i != j {
m[[i, j]] = sq[i] as f64 + sq[j] as f64 - 2.0 * dot[[i, j]];
}
}
}
m
}
/// `threshold == 1` is exactly presence (`a >= 1 ⟺ a != 0 ⟺` present in
/// the support) — shortcuts straight to `support`'s own
/// [`BitPartials::partial_jaccard`], which already has the dict-driven
/// shortcut this type's own values can't reuse (see
/// [`row_major_pairwise`](Self::row_major_pairwise)'s doc comment).
/// `threshold == 0` is the other degenerate case — every value is `>= 0`
/// (`u32`), so every column pair trivially co-occurs on every row,
/// closed-form, no traversal needed. General `threshold > 1` walks the
/// row-major decode, since only values actually `>= threshold` can ever
/// be absent from it (any lower value is already excluded from the
/// intersection/union, same reasoning as `partial_bray`).
fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
let n = self.n_cols();
if threshold == 0 {
// Diagonal included: f(i,i) = n rows too, trivially.
let full = Array2::from_elem((n, n), self.n() as u64);
return (full.clone(), full);
}
let count_geq = self.count_geq(threshold);
if threshold == 1 {
// `count_geq[i]` at threshold 1 is exactly `support`'s own
// `col_weights[i]` (both count columns present, i.e. nonzero),
// so `support.partial_jaccard`'s own diagonal already matches
// this method's `f(i,i) = count_geq[i]` convention — nothing
// left to patch.
return BitPartials::partial_jaccard(&self.support);
}
let mut inter = self.row_major_pairwise(|_, a, _, b| (a >= threshold && b >= threshold) as u64);
let mut union = Array2::<u64>::zeros((n, n));
for i in 0..n {
for j in 0..n {
if i != j {
union[[i, j]] = count_geq[i] + count_geq[j] - inter[[i, j]];
}
}
}
for i in 0..n {
inter[[i, i]] = count_geq[i];
union[[i, i]] = count_geq[i];
}
(inter, union)
}
/// `Σ min(a/sa, b/sb)` — `min` isn't bilinear, so unlike
/// `partial_relfreq_euclidean` there's no marginal-based shortcut; it's
/// computed directly in the row-major pass, same shape as `partial_bray`
/// but with each side pre-scaled by its column's global sum (`sa`/`sb`,
/// looked up from `global` by the kernel's own column indices). Diagonal
/// is `Σ a/sa = sum[i]/global[i]` — same self-comparison convention as
/// `partial_bray`'s diagonal, patched in afterwards for the same reason
/// (`row_major_pairwise` never pairs a column with itself).
fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64> {
let mut m = self.row_major_pairwise(|i, a, j, b| {
let sa = global[i] as f64;
let sb = global[j] as f64;
let pa = if sa > 0.0 { a as f64 / sa } else { 0.0 };
let pb = if sb > 0.0 { b as f64 / sb } else { 0.0 };
pa.min(pb)
});
let sum = self.sum();
for i in 0..self.n_cols() {
m[[i, i]] = if global[i] > 0 { sum[i] as f64 / global[i] as f64 } else { 0.0 };
}
m
}
/// `Σ(a/sa - b/sb)²`, reconstructed the same way as `partial_euclidean`
/// (`Σpa²+Σpb²-2Σpa·pb`), reusing the *same* raw `Σab`/`Σa²` passes —
/// just rescaled by the column's global sum at the finalisation step
/// below, not inside the row-major loop.
fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64> {
let n = self.n_cols();
let dot = self.row_major_pairwise(|_, a, _, b| (a as f64) * (b as f64));
let sq = self.sum_sq();
let mut m = Array2::<f64>::zeros((n, n));
for i in 0..n {
for j in 0..n {
if i != j {
let sa = global[i] as f64;
let sb = global[j] as f64;
let sq_a = if sa > 0.0 { sq[i] as f64 / (sa * sa) } else { 0.0 };
let sq_b = if sb > 0.0 { sq[j] as f64 / (sb * sb) } else { 0.0 };
let cross = if sa > 0.0 && sb > 0.0 { dot[[i, j]] / (sa * sb) } else { 0.0 };
m[[i, j]] = sq_a + sq_b - 2.0 * cross;
}
}
}
m
}
/// `Σ(√pa-√pb)² = Σpa+Σpb-2Σ√(pa·pb)` — `Σpa` is a plain marginal
/// (`sum()[i]/sa`, no co-presence restriction needed: it doesn't depend
/// on column `j` at all), only the cross term `Σ√(pa·pb) = Σ√(ab)/√(sa·sb)`
/// goes through the row-major pass (its `√(a·b)` kernel is zero exactly
/// when either side is absent, like every other cross term above).
fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64> {
let n = self.n_cols();
let sqrt_dot = self.row_major_pairwise(|_, a, _, b| ((a as f64) * (b as f64)).sqrt());
let sum = self.sum();
let mut m = Array2::<f64>::zeros((n, n));
for i in 0..n {
for j in 0..n {
if i != j {
let sa = global[i] as f64;
let sb = global[j] as f64;
let pa_sum = if sa > 0.0 { sum[i] as f64 / sa } else { 0.0 };
let pb_sum = if sb > 0.0 { sum[j] as f64 / sb } else { 0.0 };
let cross = if sa > 0.0 && sb > 0.0 { sqrt_dot[[i, j]] / (sa * sb).sqrt() } else { 0.0 };
m[[i, j]] = pa_sum + pb_sum - 2.0 * cross;
}
}
}
m
}
}
// ── PersistentSparseCompactIntMatrixBuilder ─────────────────────────────────
pub struct PersistentSparseCompactIntMatrixBuilder {
@@ -313,7 +546,7 @@ impl PersistentSparseCompactIntMatrixBuilder {
pub fn pack_sparse_compact_int_matrix(dir: &Path) -> io::Result<()> {
use crate::intmatrix::{col_path, ColumnarCompactIntMatrix, PackedCompactIntMatrix};
if singleton_values_path(dir).exists() {
if is_present(dir) {
return Ok(());
}
src/obicompactvec/src/storage_kind.rs
+7 -4
View File
@@ -1,8 +1,10 @@
//! How a persistent matrix's data is physically laid out on disk —
//! orthogonal to *what* it stores (count vs. presence, `obikindex::layer`'s
//! `LayerContent` concern, one crate up). `PersistentCompactIntMatrix` only
//! ever reports `Columnar`/`Packed`; `PersistentBitMatrix` is the only type
//! that can also report `Sparse`/`Implicit`.
//! `LayerContent` concern, one crate up). Both `PersistentCompactIntMatrix`
//! and `PersistentBitMatrix` can report `Columnar`/`Packed`/`Sparse`;
//! `Implicit` (no file at all, mono-genome) is `PersistentBitMatrix`-only —
//! counts always have at least one on-disk column, so there's no implicit
//! count matrix.
/// On-disk storage format of a persistent matrix.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -11,7 +13,8 @@ pub enum StorageKind {
Columnar,
/// Single mmap'd file (`matrix.pbmx`/`matrix.pcmx`), query-optimised.
Packed,
/// Row-major deduplicated sparse format (`PersistentBitMatrix` only).
/// Row-major deduplicated sparse format (`PersistentBitMatrix`/
/// `PersistentCompactIntMatrix` both).
Sparse,
/// No file at all — mono-genome, all values present
/// (`PersistentBitMatrix` only).
src/obicompactvec/src/tests/intmatrix.rs
+146 -1
View File
@@ -1,6 +1,6 @@
use tempfile::tempdir;
use crate::{pack_compact_int_matrix, PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, PersistentCompactIntVec, PersistentCompactIntVecBuilder};
use crate::{pack_compact_int_matrix, pack_sparse_compact_int_matrix, PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, PersistentCompactIntVec, PersistentCompactIntVecBuilder, StorageKind};
use crate::traits::CountPartials;
fn make_matrix(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentCompactIntMatrix) {
@@ -377,6 +377,151 @@ fn nonzero_iter_matches_row() {
drop(dir_col);
}
/// `PersistentCompactIntMatrix::open` must transparently pick up the
/// `Sparse` on-disk format written by `pack_sparse_compact_int_matrix` —
/// the regression this étape 1 work fixes: before the `Sparse` enum variant
/// existed, `open` had no code path back to a layer packed with
/// `pack_matrices(sparse=true)` (see `obikindex::index::kmer_index::
/// KmerIndex::pack_matrices`), which called `pack_sparse_compact_int_matrix`
/// on `counts/` without anything downstream able to read the result back.
#[test]
fn sparse_roundtrip_matches_columnar() {
let data: &[&[u32]] = &[&[0, 5, 0, 3, 7], &[2, 0, 0, 4, 0], &[0, 0, 9, 0, 1]];
let (dir_col, m_col) = make_matrix(data);
let (dir_sparse, _) = make_matrix(data);
pack_sparse_compact_int_matrix(&dir_sparse.path().join("counts")).unwrap();
let m_sparse = PersistentCompactIntMatrix::open(dir_sparse.path()).unwrap();
assert_eq!(m_sparse.storage_kind(), StorageKind::Sparse);
assert_eq!(m_sparse.n(), m_col.n());
assert_eq!(m_sparse.n_cols(), m_col.n_cols());
for slot in 0..m_col.n() {
assert_eq!(&*m_sparse.row(slot), &*m_col.row(slot), "slot={slot}");
}
let slots = [4usize, 0, 3, 1];
let mut expected: Vec<(usize, usize, u32)> = Vec::new();
for (i, &slot) in slots.iter().enumerate() {
for c in 0..data.len() {
let v = data[c][slot];
if v != 0 {
expected.push((i, c, v));
}
}
}
expected.sort();
let mut got: Vec<(usize, usize, u32)> = m_sparse.nonzero_iter(&slots).collect();
got.sort();
assert_eq!(got, expected);
assert_eq!(m_sparse.sub_matrix(&slots), m_col.sub_matrix(&slots));
assert_eq!(m_sparse.sum(), m_col.sum());
assert_eq!(m_sparse.count_nonzero(), m_col.count_nonzero());
drop(dir_col);
}
/// The `CountPartials` formulas on `Sparse` must agree with the
/// already-established `Columnar` reference implementation — this is what
/// exercises the row-major reconstruction (`Σ(a-b)² = Σa²+Σb²-2Σab`) against
/// real asymmetric-presence data, not just a hand re-derivation of the same
/// formula. Data deliberately mixes: cells present in only one of a pair of
/// columns (exercises the squared-difference identity's correction term),
/// cells present in both (exercises the row-major cross-term itself), and
/// singleton vs. multi-genome rows (exercises both branches of
/// `for_each_cell_in_row`).
#[test]
fn sparse_count_partials_match_dense() {
let data: &[&[u32]] = &[
&[0, 5, 0, 3, 7, 2],
&[2, 0, 0, 4, 0, 6],
&[0, 0, 9, 0, 1, 3],
&[4, 4, 0, 0, 2, 0],
];
let (dir_col, m_col) = make_matrix(data);
let (dir_sparse, _) = make_matrix(data);
pack_sparse_compact_int_matrix(&dir_sparse.path().join("counts")).unwrap();
let m_sparse = PersistentCompactIntMatrix::open(dir_sparse.path()).unwrap();
assert_eq!(m_sparse.storage_kind(), StorageKind::Sparse);
let n = m_col.n_cols();
let global = m_col.sum(); // same data on both sides, so the global sums agree too
let close = |a: f64, b: f64, ctx: &str| {
assert!((a - b).abs() < 1e-9, "{ctx}: {a} vs {b}");
};
// partial_bray
let bray_col = m_col.partial_bray_dist_matrix();
let bray_sparse = m_sparse.partial_bray_dist_matrix();
assert_eq!(bray_col, bray_sparse, "partial_bray");
// partial_euclidean
let eucl_col = m_col.partial_euclidean_dist_matrix();
let eucl_sparse = m_sparse.partial_euclidean_dist_matrix();
for i in 0..n {
for j in 0..n {
close(eucl_col[[i, j]], eucl_sparse[[i, j]], &format!("partial_euclidean[{i},{j}]"));
}
}
// partial_threshold_jaccard, across thresholds 0 (degenerate), 1
// (support shortcut), and >1 (general row-major path).
for threshold in [0u32, 1, 2, 3] {
let (inter_col, union_col) = m_col.partial_threshold_jaccard_dist_matrix(threshold);
let (inter_sparse, union_sparse) = m_sparse.partial_threshold_jaccard_dist_matrix(threshold);
assert_eq!(inter_col, inter_sparse, "partial_threshold_jaccard inter, threshold={threshold}");
assert_eq!(union_col, union_sparse, "partial_threshold_jaccard union, threshold={threshold}");
}
// partial_relfreq_bray
let rfb_col = m_col.partial_relfreq_bray_dist_matrix(&global);
let rfb_sparse = m_sparse.partial_relfreq_bray_dist_matrix(&global);
for i in 0..n {
for j in 0..n {
close(rfb_col[[i, j]], rfb_sparse[[i, j]], &format!("partial_relfreq_bray[{i},{j}]"));
}
}
// partial_relfreq_euclidean
let rfe_col = m_col.partial_relfreq_euclidean_dist_matrix(&global);
let rfe_sparse = m_sparse.partial_relfreq_euclidean_dist_matrix(&global);
for i in 0..n {
for j in 0..n {
close(rfe_col[[i, j]], rfe_sparse[[i, j]], &format!("partial_relfreq_euclidean[{i},{j}]"));
}
}
// partial_hellinger
let hel_col = m_col.partial_hellinger_euclidean_dist_matrix(&global);
let hel_sparse = m_sparse.partial_hellinger_euclidean_dist_matrix(&global);
for i in 0..n {
for j in 0..n {
close(hel_col[[i, j]], hel_sparse[[i, j]], &format!("partial_hellinger[{i},{j}]"));
}
}
drop(dir_col);
}
/// `pack_sparse_compact_int_matrix` must also handle an already-`Packed`
/// source (not just `Columnar`) — mirrors
/// `pack_sparse_bit_matrix_from_already_packed_matrix` on the bit side.
#[test]
fn sparse_roundtrip_from_packed() {
let data: &[&[u32]] = &[&[10, 300, 500], &[200, 50, 1000]];
let (dir_col, m_col) = make_matrix(data);
let (dir_sparse, _) = make_matrix(data);
pack_compact_int_matrix(&dir_sparse.path().join("counts")).unwrap();
pack_sparse_compact_int_matrix(&dir_sparse.path().join("counts")).unwrap();
let m_sparse = PersistentCompactIntMatrix::open(dir_sparse.path()).unwrap();
assert_eq!(m_sparse.storage_kind(), StorageKind::Sparse);
for slot in 0..m_col.n() {
assert_eq!(&*m_sparse.row(slot), &*m_col.row(slot), "slot={slot}");
}
drop(dir_col);
}
#[test]
fn partial_relfreq_bray_additive_across_split() {
// Split rows [1,2,3,4,5] between two matrices; partial sums should add up.
src/obicompactvec/src/tests/sparse.rs
+43
View File
@@ -372,3 +372,46 @@ fn pack_sparse_bit_matrix_from_already_packed_matrix() {
assert_eq!(&*sparse.row(slot), &**expected, "slot {slot}");
}
}
/// Regression for the `partial_jaccard`/`partial_hamming` diagonal bug:
/// `col_weights_and_pair_counts`'s own `inter` never pairs a column with
/// itself, so before this fix `partial_jaccard`'s diagonal was `(0,
/// 2×col_weights[i])` instead of the genuine self-comparison
/// `(col_weights[i], col_weights[i])`, and `partial_hamming`'s diagonal
/// wasn't `0` (a column XORed with itself) either — invisible through
/// `jaccard_dist_matrix`/`hamming_dist_matrix` (both explicitly zero their
/// own diagonal at finalisation), but wrong for any caller of the raw
/// `partial_*` methods directly (documented as additive components, see
/// `traits.rs`).
#[test]
fn partial_jaccard_and_hamming_match_dense_including_diagonal() {
use crate::traits::BitPartials;
let cols: &[&[bool]] = &[
&[true, false, true, true, false],
&[false, true, true, false, true],
&[true, true, false, false, true],
];
let (dir_dense, dense) = make_dense(cols);
let n = cols.first().unwrap().len();
let n_cols = cols.len();
let rows: Vec<Vec<bool>> = (0..n).map(|s| cols.iter().map(|c| c[s]).collect()).collect();
let row_refs: Vec<&[bool]> = rows.iter().map(|r| r.as_slice()).collect();
let (dir_sparse, sparse) = make_sparse(&row_refs, n_cols);
let (dense_inter, dense_union) = dense.partial_jaccard_dist_matrix();
let (sparse_inter, sparse_union) = sparse.partial_jaccard();
assert_eq!(dense_inter, sparse_inter, "partial_jaccard inter");
assert_eq!(dense_union, sparse_union, "partial_jaccard union");
let dense_hamming = dense.partial_hamming_dist_matrix();
let sparse_hamming = sparse.partial_hamming();
assert_eq!(dense_hamming, sparse_hamming, "partial_hamming");
for i in 0..n_cols {
assert_eq!(sparse_hamming[[i, i]], 0, "hamming self-distance must be 0, col {i}");
}
drop(dir_dense);
drop(dir_sparse);
}
src/obicompactvec/src/traits.rs
+9 -7
View File
@@ -3,13 +3,15 @@ use ndarray::{Array1, Array2};
/// Minimal shared surface between `PersistentBitMatrix` (dense) and
/// `PersistentSparseBitMatrix` (row-major, deduplicated) — exactly what
/// real consumers use today (`obikphylo::siblings::cache::Mat`), not the
/// two types' full individual APIs. Column-oriented operations
/// (`col`/`col_view`, the `BitPartials`/`ColumnWeights` distance-matrix
/// traits above) are *not* part of this trait — `PersistentSparseBitMatrix`
/// only offers a naive, row-scanning `count_ones` for now (see
/// `DevDocMD/architecture/siblings.md` and the sparse-matrix design plan,
/// "Explicitly deferred": a row-major co-occurrence rewrite of the
/// pairwise distance matrices is future work, not part of this trait).
/// two types' full individual APIs. Column-oriented operations (`col`/
/// `col_view`) are *not* part of this trait — `PersistentSparseBitMatrix`
/// has no on-disk column representation to hand one out from at all. Note
/// this is no longer true of `BitPartials`/`ColumnWeights`:
/// `PersistentSparseBitMatrix` implements both today, via the same
/// dict-driven row-major pass `count_ones` itself goes through
/// (`col_weights_and_pair_counts`) — they're just not folded into this
/// trait's own surface, since `BinaryMatrix` only ever needed to cover what
/// `Mat` actually calls.
pub trait BinaryMatrix {
/// Number of rows (k-mer slots).
fn n(&self) -> usize;
src/obikindex/src/layer/content_layer.rs
+12
View File
@@ -217,6 +217,18 @@ impl KmerLayer {
}
}
/// Every nonzero `(idx into slots, col, value)` triple among `slots` —
/// format-agnostic column-major fetch, delegating to whichever matrix
/// this layer actually holds (see `obicompactvec::PersistentBitMatrix`/
/// `PersistentCompactIntMatrix::nonzero_iter`).
pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> Box<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
match self {
KmerLayer::Presence { layer, .. } => layer.nonzero_iter(slots),
KmerLayer::Count { layer, .. } => Box::new(layer.nonzero_iter(slots)),
KmerLayer::Empty { .. } => panic!("Layer::nonzero_iter() called on an Empty layer"),
}
}
/// Per-genome column weights — count sum or presence k-mer count,
/// depending on content (see `obicompactvec::ColumnWeights`).
pub fn col_weights(&self) -> ndarray::Array1<u64> {
src/obikindex/src/layer/tests/typed_layer.rs
+62
View File
@@ -98,6 +98,68 @@ fn presence_layer_generic_over_sparse_matches_dense() {
}
}
// ── `PersistentCompactIntMatrix` transparently absorbing `Sparse` ────────────
//
// Unlike presence (`PersistentSparseBitMatrix` is a standalone type, plugged
// into `TypedLayer` via its own `LayerData`/`HasLayerContent` impls, see
// `presence_layer_generic_over_sparse_matches_dense` above), counts don't
// get a second `TypedLayer<PersistentSparseCompactIntMatrix>` instantiation
// — sparsity is absorbed inside the `PersistentCompactIntMatrix` enum
// itself, so `TypedLayer<PersistentCompactIntMatrix>` (the only impl block
// for counts) must keep working unchanged once the on-disk format is
// repacked to `Sparse`. This is the regression test for the bug that
// motivated adding the `Sparse` variant: before it existed, `open()` had no
// path back to a layer packed with `pack_matrices(sparse=true)` — see
// `obikindex::index::kmer_index::KmerIndex::pack_matrices`, which already
// called `pack_sparse_compact_int_matrix` on `counts/` without anything
// downstream able to read the result back.
#[test]
fn count_layer_transparently_reads_sparse_after_pack() {
set_k(4);
let dir = tempdir().unwrap();
write_unitigs(dir.path(), &[b"AAATCTA", b"CTTCGCC", b"TGATACG"]);
let n_genomes = 3;
let mode = IndexMode::Exact;
// Deterministic, arbitrary count function — same shape as the presence
// test's arbitrary predicate, just producing a small count instead of a
// bool.
TypedLayer::<()>::build_with_matrix(dir.path(), DEFAULT_BLOCK_BITS, &mode, false, n_genomes, |kmer| {
(0..n_genomes)
.map(|g| ((kmer.raw().wrapping_add(g as u64)) % 5) as u32)
.collect()
})
.unwrap();
let dense_layer = TypedLayer::<PersistentCompactIntMatrix>::open(dir.path()).unwrap();
assert_eq!(dense_layer.storage_kind(), obicompactvec::StorageKind::Columnar);
let n = dense_layer.n();
let slots: Vec<usize> = (0..n).collect();
let dense_sub = dense_layer.sub_matrix(&slots);
let dense_weights = dense_layer.col_weights();
let dense_slots: Vec<Option<usize>> = all_canonical_kmers(dir.path())
.iter()
.map(|&kmer| dense_layer.find_slot(kmer))
.collect();
// Repack in place — same directory, matching how `pack --sparse` works
// in production.
obicompactvec::pack_sparse_compact_int_matrix(&dir.path().join(COUNTS_DIR)).unwrap();
let sparse_layer = TypedLayer::<PersistentCompactIntMatrix>::open(dir.path()).unwrap();
assert_eq!(sparse_layer.storage_kind(), obicompactvec::StorageKind::Sparse);
assert_eq!(sparse_layer.n(), n);
assert_eq!(sparse_layer.n_cols(), dense_layer.n_cols());
assert_eq!(sparse_layer.sub_matrix(&slots), dense_sub);
assert_eq!(sparse_layer.col_weights(), dense_weights);
let sparse_slots: Vec<Option<usize>> = all_canonical_kmers(dir.path())
.iter()
.map(|&kmer| sparse_layer.find_slot(kmer))
.collect();
assert_eq!(sparse_slots, dense_slots);
}
// ── content / storage_kind / evidence_kind (runtime introspection) ─────────
#[test]
src/obikindex/src/layer/typed_layer.rs
+12
View File
@@ -415,6 +415,12 @@ impl TypedLayer<PersistentCompactIntMatrix> {
pub fn col_weights(&self) -> ndarray::Array1<u64> {
obicompactvec::ColumnWeights::col_weights(&self.data)
}
/// Every nonzero `(idx into slots, col, value)` triple among `slots` —
/// delegates to `PersistentCompactIntMatrix::nonzero_iter`.
pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> impl Iterator<Item = (usize, usize, u32)> + 'a {
self.data.nonzero_iter(slots)
}
}
// ── Mode 3 — presence/absence matrix ─────────────────────────────────────────
@@ -469,6 +475,12 @@ impl<D: LayerData<Item = Box<[bool]>> + BinaryMatrix + obicompactvec::ColumnWeig
}
impl TypedLayer<PersistentBitMatrix> {
/// Every nonzero `(idx into slots, col, value)` triple among `slots` —
/// delegates to `PersistentBitMatrix::nonzero_iter`.
pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> Box<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
self.data.nonzero_iter(slots)
}
pub fn append_genome_column(
layer_dir: &Path,
value_of: impl Fn(usize) -> bool,