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 } }