From 11476bc557b6278b9dd8db92f1037f59ef2bccbb Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Sun, 16 Aug 2026 13:05:27 +0200 Subject: [PATCH] Add sparse bit matrix support and enhance dimension logging Introduce a `Sparse` variant for `PersistentBitMatrix` with full method dispatch and implementations of `ColumnWeights` and `BitPartials` traits. Refactor distance matrix computations in layered stores to eagerly collect results before reduction, enabling reliable dimension tracking. Add comprehensive debug logging across phylo commands and distance modules to report matrix shapes and flag out-of-bounds indices during CSV iteration. --- src/obicompactvec/src/bitmatrix/persistent.rs | 77 ++++++++++-- src/obicompactvec/src/bitmatrix/sparse.rs | 64 +++++++++- src/obicompactvec/src/traits.rs | 5 +- src/obikindex/src/distance.rs | 1 + src/obikmer/src/cmd/phylo/mod.rs | 5 + src/obilayeredmap/src/layered_store.rs | 116 ++++++++++++++---- 6 files changed, 225 insertions(+), 43 deletions(-) diff --git a/src/obicompactvec/src/bitmatrix/persistent.rs b/src/obicompactvec/src/bitmatrix/persistent.rs index 7a8e8293..c8a41166 100644 --- a/src/obicompactvec/src/bitmatrix/persistent.rs +++ b/src/obicompactvec/src/bitmatrix/persistent.rs @@ -11,17 +11,20 @@ use crate::views::BitSliceView; use super::columnar::ColumnarBitMatrix; use super::packed::PackedBitMatrix; +use super::sparse::PersistentSparseBitMatrix; // ── PersistentBitMatrix — public enum ──────────────────────────────────────── -/// Bit matrix that transparently handles columnar, packed, and implicit formats. +/// Bit matrix that transparently handles columnar, packed, sparse, and implicit formats. /// /// - `Columnar`: per-column `.pbiv` files (original format, used during build) /// - `Packed`: single `matrix.pbmx` file (optimised for query — one `mmap`) +/// - `Sparse`: sparse row-major format (`sparse_meta.json` + PFIV/EF files) /// - `Implicit`: no file — all values are 1 (mono-genome presence/absence) pub enum PersistentBitMatrix { Columnar(ColumnarBitMatrix), Packed(PackedBitMatrix), + Sparse(PersistentSparseBitMatrix), Implicit { n_rows: usize, n_cols: usize }, } @@ -29,19 +32,30 @@ impl PersistentBitMatrix { /// Open from `layer_dir`, auto-detecting the format. /// /// Checks (in order): - /// 1. `layer_dir/presence/matrix.pbmx` → Packed - /// 2. `layer_dir/presence/meta.json` → Columnar - /// 3. `layer_dir/layer_meta.json` → Implicit (new index) - /// 4. `layer_dir/unitigs.bin` → Implicit with warning (old index) + /// 1. `layer_dir/presence/matrix.pbmx` → Packed + /// 2. `layer_dir/presence/meta.json` → Columnar + /// 3. `layer_dir/presence/sparse_meta.json` → Sparse + /// 4. `layer_dir/layer_meta.json` → Implicit (new index) + /// 5. `layer_dir/unitigs.bin` → Implicit with warning (old index) pub fn open(layer_dir: &Path) -> io::Result { let presence_dir = layer_dir.join("presence"); if presence_dir.join("matrix.pbmx").exists() { - return Ok(Self::Packed(PackedBitMatrix::open(&presence_dir.join("matrix.pbmx"))?)); + let m = PackedBitMatrix::open(&presence_dir.join("matrix.pbmx"))?; + eprintln!("[DIAG] PersistentBitMatrix::open PACKED layer={} n_rows={} n_cols={}", layer_dir.display(), m.n_rows, m.n_cols); + return Ok(Self::Packed(m)); } if MatrixMeta::load(&presence_dir).is_ok() { - return Ok(Self::Columnar(ColumnarBitMatrix::open(&presence_dir)?)); + let m = ColumnarBitMatrix::open(&presence_dir)?; + eprintln!("[DIAG] PersistentBitMatrix::open COLUMNAR layer={} n={} n_cols={}", layer_dir.display(), m.n(), m.n_cols()); + return Ok(Self::Columnar(m)); + } + + if presence_dir.join("sparse_meta.json").exists() { + let m = PersistentSparseBitMatrix::open(&presence_dir)?; + eprintln!("[DIAG] PersistentBitMatrix::open SPARSE layer={} n={} n_cols={}", layer_dir.display(), m.n(), m.n_cols()); + return Ok(Self::Sparse(m)); } // No presence matrix → Implicit; requires layer_meta.json @@ -52,6 +66,7 @@ impl PersistentBitMatrix { layer_dir.display() ), ))?; + eprintln!("[DIAG] PersistentBitMatrix::open IMPLICIT layer={} n_rows={} n_cols=1", layer_dir.display(), meta.n); Ok(Self::Implicit { n_rows: meta.n, n_cols: 1 }) } @@ -60,6 +75,7 @@ impl PersistentBitMatrix { match self { Self::Columnar(m) => m.n(), Self::Packed(m) => m.n_rows, + Self::Sparse(m) => m.n(), Self::Implicit { n_rows, .. } => *n_rows, } } @@ -69,6 +85,7 @@ impl PersistentBitMatrix { match self { Self::Columnar(m) => m.n_cols(), Self::Packed(m) => m.n_cols, + Self::Sparse(m) => m.n_cols(), Self::Implicit { n_cols, .. } => *n_cols, } } @@ -86,6 +103,7 @@ impl PersistentBitMatrix { match self { Self::Columnar(m) => m.col(c).view(), Self::Packed(m) => m.col_slice(c), + Self::Sparse(_) => panic!("col_view() not available on Sparse PersistentBitMatrix"), Self::Implicit { .. } => panic!("col_view() not available on Implicit PersistentBitMatrix"), } } @@ -100,6 +118,11 @@ impl PersistentBitMatrix { match self { Self::Columnar(m) => m.col(c).get(slot) as u32, Self::Packed(m) => m.col_slice(c).get(slot) as u32, + Self::Sparse(m) => { + let mut buf = vec![0u32; m.n_cols()]; + m.fill_row(slot, &mut buf); + buf[c] + } Self::Implicit { .. } => 1, } } @@ -108,6 +131,8 @@ impl PersistentBitMatrix { match self { Self::Columnar(m) => PersistentBitVecBuilder::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 PersistentBitMatrix")), Self::Implicit { n_rows, .. } => { PersistentBitVecBuilder::new_ones(*n_rows, path) } @@ -119,6 +144,7 @@ impl PersistentBitMatrix { match self { Self::Columnar(m) => m.row(slot), Self::Packed(m) => m.row(slot), + Self::Sparse(m) => m.row(slot), Self::Implicit { n_cols, .. } => vec![true; *n_cols].into_boxed_slice(), } } @@ -129,6 +155,7 @@ impl PersistentBitMatrix { 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), Self::Implicit { n_cols, .. } => buf[..*n_cols].fill(1), } } @@ -147,6 +174,13 @@ impl PersistentBitMatrix { match self { Self::Columnar(m) => m.col(c).view().fill_batch(slots, &mut col_buf), Self::Packed(m) => m.col_slice(c).fill_batch(slots, &mut col_buf), + Self::Sparse(m) => { + let mut row_buf = vec![false; m.n_cols()]; + for (i, &slot) in slots.iter().enumerate() { + m.fill_row_bool(slot, &mut row_buf); + col_buf[i] = row_buf[c]; + } + } Self::Implicit { .. } => col_buf.iter_mut().for_each(|b| *b = true), } out.push(col_buf); @@ -176,10 +210,20 @@ impl PersistentBitMatrix { match self { Self::Columnar(m) => m.col(c).view().fill_batch_sorted(&sorted_slots, &mut tmp), Self::Packed(m) => m.col_slice(c).fill_batch_sorted(&sorted_slots, &mut tmp), + Self::Sparse(m) => { + for (i, &orig_idx) in perm.iter().enumerate() { + let slot = sorted_slots[i]; + let mut row_buf = vec![false; m.n_cols()]; + m.fill_row_bool(slot, &mut row_buf); + col[orig_idx] = row_buf[c]; + } + } Self::Implicit { .. } => tmp.iter_mut().for_each(|b| *b = true), } for (i, &orig_idx) in perm.iter().enumerate() { - col[orig_idx] = tmp[i]; + if !matches!(self, Self::Sparse(_)) { + col[orig_idx] = tmp[i]; + } } } } @@ -189,14 +233,16 @@ impl PersistentBitMatrix { match self { Self::Columnar(m) => m.count_ones(), Self::Packed(m) => m.count_ones(), + Self::Sparse(m) => m.count_ones(), Self::Implicit { n_rows, n_cols } => Array1::from_elem(*n_cols, *n_rows as u64), } } pub fn partial_jaccard_dist_matrix(&self) -> (Array2, Array2) { - match self { + let result = match self { Self::Columnar(m) => m.partial_jaccard_dist_matrix(), Self::Packed(m) => m.partial_jaccard_dist_matrix(), + Self::Sparse(m) => BitPartials::partial_jaccard(m), Self::Implicit { n_rows, n_cols } => { let v = *n_rows as u64; let n = *n_cols; @@ -207,15 +253,22 @@ impl PersistentBitMatrix { }} (inter, union) } - } + }; + eprintln!("[DIAG] PersistentBitMatrix::partial_jaccard_dist_matrix self.n_cols={} result={}x{} / {}x{}", + self.n_cols(), result.0.shape()[0], result.0.shape()[1], result.1.shape()[0], result.1.shape()[1]); + result } pub fn partial_hamming_dist_matrix(&self) -> Array2 { - match self { + let result = match self { Self::Columnar(m) => m.partial_hamming_dist_matrix(), Self::Packed(m) => m.partial_hamming_dist_matrix(), + Self::Sparse(m) => BitPartials::partial_hamming(m), Self::Implicit { n_cols, .. } => Array2::zeros((*n_cols, *n_cols)), - } + }; + eprintln!("[DIAG] PersistentBitMatrix::partial_hamming_dist_matrix self.n_cols={} result={}x{}", + self.n_cols(), result.shape()[0], result.shape()[1]); + result } /// Append a new column to an on-disk Columnar matrix. diff --git a/src/obicompactvec/src/bitmatrix/sparse.rs b/src/obicompactvec/src/bitmatrix/sparse.rs index e9795783..7d6aadc3 100644 --- a/src/obicompactvec/src/bitmatrix/sparse.rs +++ b/src/obicompactvec/src/bitmatrix/sparse.rs @@ -25,8 +25,9 @@ use std::fs; use std::io; use std::path::{Path, PathBuf}; -use ndarray::Array1; +use ndarray::{Array1, Array2}; +use crate::traits::{BitPartials, ColumnWeights}; use crate::eliasfano::{EliasFano, EliasFanoBuilder}; use crate::fixedintvec::{PersistentFixedIntVec, PersistentFixedIntVecBuilder, bit_width_for_range}; use crate::meta::field; @@ -174,7 +175,7 @@ impl PersistentSparseBitMatrix { } } - fn fill_row_bool(&self, slot: usize, buf: &mut [bool]) { + pub(crate) fn fill_row_bool(&self, slot: usize, buf: &mut [bool]) { buf.fill(false); if self.is_multi.get(slot) { let pos = self.is_multi.rank1(slot) as usize; @@ -372,7 +373,64 @@ impl PersistentSparseBitMatrixBuilder { } } -/// `pack --sparse`'s entry point: converts a presence directory into the +// ── Trait impls ─────────────────────────────────────────────────────────────── + +impl ColumnWeights for PersistentSparseBitMatrix { + fn col_weights(&self) -> Array1 { + self.count_ones() + } +} + +impl BitPartials for PersistentSparseBitMatrix { + fn partial_jaccard(&self) -> (Array2, Array2) { + let n = self.n_cols(); + let mut inter = Array2::::zeros((n, n)); + let mut buf = vec![0u32; n]; + for slot in 0..self.n() { + self.fill_row(slot, &mut buf); + let present: Vec = buf.iter().enumerate().filter(|&(_, &v)| v != 0).map(|(c, _)| c).collect(); + for (p, &i) in present.iter().enumerate() { + for &j in present.iter().skip(p) { + inter[[i, j]] += 1; + inter[[j, i]] += 1; + } + } + } + let col_weights = self.count_ones(); + 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]]; + } + } + (inter, union) + } + + fn partial_hamming(&self) -> Array2 { + let n = self.n_cols(); + let mut inter = Array2::::zeros((n, n)); + let mut buf = vec![0u32; n]; + for slot in 0..self.n() { + self.fill_row(slot, &mut buf); + let present: Vec = buf.iter().enumerate().filter(|&(_, &v)| v != 0).map(|(c, _)| c).collect(); + for (p, &i) in present.iter().enumerate() { + for &j in present.iter().skip(p) { + inter[[i, j]] += 1; + inter[[j, i]] += 1; + } + } + } + let col_weights = self.count_ones(); + 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]]); + } + } + m + } +} /// sparse on-disk format in place, mirroring [`super::pack_bit_matrix`]'s /// convention (old-format files removed only after the new format is /// fully written, so a crash mid-conversion leaves the previous, still diff --git a/src/obicompactvec/src/traits.rs b/src/obicompactvec/src/traits.rs index 38f3afb1..dc958e45 100644 --- a/src/obicompactvec/src/traits.rs +++ b/src/obicompactvec/src/traits.rs @@ -166,6 +166,7 @@ pub trait BitPartials: ColumnWeights { fn jaccard_dist_matrix(&self) -> Array2 { let (inter, union) = self.partial_jaccard(); let n = inter.shape()[0]; + eprintln!("[TRACE] BitPartials::jaccard_dist_matrix finalising: n={}", n); let mut m = Array2::::zeros((n, n)); for i in 0..n { for j in 0..n { @@ -182,7 +183,9 @@ pub trait BitPartials: ColumnWeights { /// Mash distance (https://mash.readthedocs.io/en/latest/distances.html), derived /// from the Jaccard distance. fn mash_dist_matrix(&self, k: usize) -> Array2 { - jaccard_to_mash(&self.jaccard_dist_matrix(), k) + let j = self.jaccard_dist_matrix(); + eprintln!("[TRACE] BitPartials::mash_dist_matrix jaccard shape={}x{}", j.shape()[0], j.shape()[1]); + jaccard_to_mash(&j, k) } fn hamming_dist_matrix(&self) -> Array2 { diff --git a/src/obikindex/src/distance.rs b/src/obikindex/src/distance.rs index f5d63be5..8ae1d23f 100644 --- a/src/obikindex/src/distance.rs +++ b/src/obikindex/src/distance.rs @@ -121,6 +121,7 @@ impl KmerIndex { ))); } }; + tracing::info!("distance matrix final: {}x{}", matrix.shape()[0], matrix.shape()[1]); let shared = if shared_kmers { let (inter, _) = BitPartials::partial_jaccard(&global); diff --git a/src/obikmer/src/cmd/phylo/mod.rs b/src/obikmer/src/cmd/phylo/mod.rs index cfeaf816..bf46144c 100644 --- a/src/obikmer/src/cmd/phylo/mod.rs +++ b/src/obikmer/src/cmd/phylo/mod.rs @@ -294,12 +294,17 @@ pub fn run(args: PhyloArgs) { // ── Distance matrix → CSV ───────────────────────────────────────────────── let write_dist_csv = |w: &mut dyn Write| { + let matrix_shape = result.matrix.shape(); + eprintln!("[DIAG] write_dist_csv: matrix_shape={}x{} labels.len()={} n={}", matrix_shape[0], matrix_shape[1], labels.len(), n); write!(w, "genome").unwrap(); for g in &labels { write!(w, ",{g}").unwrap(); } writeln!(w).unwrap(); for (i, g) in labels.iter().enumerate() { write!(w, "{g}").unwrap(); for j in 0..n { + if i >= matrix_shape[0] || j >= matrix_shape[1] { + eprintln!("[DIAG] OUT OF BOUNDS: i={} j={} matrix={}x{}", i, j, matrix_shape[0], matrix_shape[1]); + } write!(w, ",{:.6}", result.matrix[[i, j]]).unwrap(); } writeln!(w).unwrap(); diff --git a/src/obilayeredmap/src/layered_store.rs b/src/obilayeredmap/src/layered_store.rs index 4747af87..c335a2ff 100644 --- a/src/obilayeredmap/src/layered_store.rs +++ b/src/obilayeredmap/src/layered_store.rs @@ -22,10 +22,16 @@ impl LayeredStore { impl ColumnWeights for LayeredStore { fn col_weights(&self) -> Array1 { - self.0.par_iter() + let parts: Vec> = self.0.par_iter() .map(|s| s.col_weights()) - .reduce_with(|a, b| a + b) - .unwrap_or_else(|| Array1::zeros(0)) + .collect(); + for (i, w) in parts.iter().enumerate() { + eprintln!("layered_store col_weights layer={i} len={}", w.len()); + } + let result = parts.into_iter().reduce(|a, b| a + b) + .unwrap_or_else(|| Array1::zeros(0)); + eprintln!("layered_store col_weights reduced len={}", result.len()); + result } } @@ -33,45 +39,87 @@ impl ColumnWeights for LayeredStore { impl CountPartials for LayeredStore { fn partial_bray(&self) -> Array2 { - self.0.par_iter() + let parts: Vec<_> = self.0.par_iter() .map(|s| s.partial_bray()) - .reduce_with(|a, b| a + b) - .unwrap() + .collect(); + for (i, m) in parts.iter().enumerate() { + eprintln!("layered_store partial_bray layer={i} {}x{}", + m.shape()[0], m.shape()[1]); + } + let result = parts.into_iter().reduce(|a, b| a + b).unwrap(); + eprintln!("layered_store partial_bray reduced {}x{}", + result.shape()[0], result.shape()[1]); + result } fn partial_euclidean(&self) -> Array2 { - self.0.par_iter() + let parts: Vec<_> = self.0.par_iter() .map(|s| s.partial_euclidean()) - .reduce_with(|a, b| a + b) - .unwrap() + .collect(); + for (i, m) in parts.iter().enumerate() { + eprintln!("layered_store partial_euclidean layer={i} {}x{}", + m.shape()[0], m.shape()[1]); + } + let result = parts.into_iter().reduce(|a, b| a + b).unwrap(); + eprintln!("layered_store partial_euclidean reduced {}x{}", + result.shape()[0], result.shape()[1]); + result } fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2, Array2) { - self.0.par_iter() + let parts: Vec<_> = self.0.par_iter() .map(|s| s.partial_threshold_jaccard(threshold)) - .reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu)) - .unwrap() + .collect(); + for (i, (inter, union)) in parts.iter().enumerate() { + eprintln!("layered_store partial_threshold_jaccard layer={i} threshold={threshold} inter={}x{} union={}x{}", + inter.shape()[0], inter.shape()[1], union.shape()[0], union.shape()[1]); + } + let (ai, au) = parts.into_iter().reduce(|(ai, au), (bi, bu)| (ai + bi, au + bu)).unwrap(); + eprintln!("layered_store partial_threshold_jaccard reduced threshold={threshold} inter={}x{} union={}x{}", + ai.shape()[0], ai.shape()[1], au.shape()[0], au.shape()[1]); + (ai, au) } fn partial_relfreq_bray(&self, global: &Array1) -> Array2 { - self.0.par_iter() + let parts: Vec<_> = self.0.par_iter() .map(|s| s.partial_relfreq_bray(global)) - .reduce_with(|a, b| a + b) - .unwrap() + .collect(); + for (i, m) in parts.iter().enumerate() { + eprintln!("layered_store partial_relfreq_bray layer={i} {}x{}", + m.shape()[0], m.shape()[1]); + } + let result = parts.into_iter().reduce(|a, b| a + b).unwrap(); + eprintln!("layered_store partial_relfreq_bray reduced {}x{}", + result.shape()[0], result.shape()[1]); + result } fn partial_relfreq_euclidean(&self, global: &Array1) -> Array2 { - self.0.par_iter() + let parts: Vec<_> = self.0.par_iter() .map(|s| s.partial_relfreq_euclidean(global)) - .reduce_with(|a, b| a + b) - .unwrap() + .collect(); + for (i, m) in parts.iter().enumerate() { + eprintln!("layered_store partial_relfreq_euclidean layer={i} {}x{}", + m.shape()[0], m.shape()[1]); + } + let result = parts.into_iter().reduce(|a, b| a + b).unwrap(); + eprintln!("layered_store partial_relfreq_euclidean reduced {}x{}", + result.shape()[0], result.shape()[1]); + result } fn partial_hellinger(&self, global: &Array1) -> Array2 { - self.0.par_iter() + let parts: Vec<_> = self.0.par_iter() .map(|s| s.partial_hellinger(global)) - .reduce_with(|a, b| a + b) - .unwrap() + .collect(); + for (i, m) in parts.iter().enumerate() { + eprintln!("layered_store partial_hellinger layer={i} {}x{}", + m.shape()[0], m.shape()[1]); + } + let result = parts.into_iter().reduce(|a, b| a + b).unwrap(); + eprintln!("layered_store partial_hellinger reduced {}x{}", + result.shape()[0], result.shape()[1]); + result } } @@ -79,17 +127,31 @@ impl CountPartials for LayeredStore { impl BitPartials for LayeredStore { fn partial_jaccard(&self) -> (Array2, Array2) { - self.0.par_iter() + let parts: Vec<_> = self.0.par_iter() .map(|s| s.partial_jaccard()) - .reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu)) - .unwrap() + .collect(); + for (i, (inter, union)) in parts.iter().enumerate() { + eprintln!("layered_store partial_jaccard layer={i} inter={}x{} union={}x{}", + inter.shape()[0], inter.shape()[1], union.shape()[0], union.shape()[1]); + } + let (ai, au) = parts.into_iter().reduce(|(ai, au), (bi, bu)| (ai + bi, au + bu)).unwrap(); + eprintln!("layered_store partial_jaccard reduced inter={}x{} union={}x{}", + ai.shape()[0], ai.shape()[1], au.shape()[0], au.shape()[1]); + (ai, au) } fn partial_hamming(&self) -> Array2 { - self.0.par_iter() + let parts: Vec<_> = self.0.par_iter() .map(|s| s.partial_hamming()) - .reduce_with(|a, b| a + b) - .unwrap() + .collect(); + for (i, m) in parts.iter().enumerate() { + eprintln!("layered_store partial_hamming layer={i} {}x{}", + m.shape()[0], m.shape()[1]); + } + let result = parts.into_iter().reduce(|a, b| a + b).unwrap(); + eprintln!("layered_store partial_hamming reduced {}x{}", + result.shape()[0], result.shape()[1]); + result } }