diff --git a/src/obicompactvec/src/bitmatrix/persistent.rs b/src/obicompactvec/src/bitmatrix/persistent.rs index c8a41166..5e9a45ea 100644 --- a/src/obicompactvec/src/bitmatrix/persistent.rs +++ b/src/obicompactvec/src/bitmatrix/persistent.rs @@ -41,21 +41,15 @@ impl PersistentBitMatrix { let presence_dir = layer_dir.join("presence"); if presence_dir.join("matrix.pbmx").exists() { - 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)); + return Ok(Self::Packed(PackedBitMatrix::open(&presence_dir.join("matrix.pbmx"))?)); } if MatrixMeta::load(&presence_dir).is_ok() { - 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)); + return Ok(Self::Columnar(ColumnarBitMatrix::open(&presence_dir)?)); } 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)); + return Ok(Self::Sparse(PersistentSparseBitMatrix::open(&presence_dir)?)); } // No presence matrix → Implicit; requires layer_meta.json @@ -66,7 +60,6 @@ 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 }) } @@ -239,7 +232,7 @@ impl PersistentBitMatrix { } pub fn partial_jaccard_dist_matrix(&self) -> (Array2, Array2) { - let result = match self { + 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), @@ -253,22 +246,16 @@ 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 { - let result = match self { + 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 7d6aadc3..1e98352e 100644 --- a/src/obicompactvec/src/bitmatrix/sparse.rs +++ b/src/obicompactvec/src/bitmatrix/sparse.rs @@ -157,75 +157,103 @@ impl PersistentSparseBitMatrix { out.into_boxed_slice() } + /// Calls `f` once per genome index present at `slot` — the shared + /// decode branch (singleton vs. varint-encoded multi set) behind + /// `fill_row`, `fill_row_bool`, and `fill_sub_matrix`. + #[inline] + fn for_each_genome_in_row(&self, slot: usize, mut f: impl FnMut(usize)) { + if self.is_multi.get(slot) { + let pos = self.is_multi.rank1(slot) as usize; + let dict_id = self.multi.get(pos) as usize; + let (start, end) = self.dict_entry_range(dict_id); + let mut p = start; + while p < end { + f(read_varint(&self.dict_values, &mut p) as usize); + } + } else { + let pos = self.is_multi.rank0(slot) as usize; + f(self.singleton.get(pos) as usize); + } + } + /// Fill `buf[i]` with `1` iff genome `i` is present at `slot`, else `0` /// — mirrors [`super::PersistentBitMatrix::fill_row`]'s signature. pub fn fill_row(&self, slot: usize, buf: &mut [u32]) { buf[..self.n_cols].fill(0); - if self.is_multi.get(slot) { - let pos = self.is_multi.rank1(slot) as usize; - let dict_id = self.multi.get(pos) as usize; - let (start, end) = self.dict_entry_range(dict_id); - let mut p = start; - while p < end { - buf[read_varint(&self.dict_values, &mut p) as usize] = 1; - } - } else { - let pos = self.is_multi.rank0(slot) as usize; - buf[self.singleton.get(pos) as usize] = 1; - } + self.for_each_genome_in_row(slot, |g| buf[g] = 1); } 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; - let dict_id = self.multi.get(pos) as usize; - let (start, end) = self.dict_entry_range(dict_id); - let mut p = start; - while p < end { - buf[read_varint(&self.dict_values, &mut p) as usize] = true; - } - } else { - let pos = self.is_multi.rank0(slot) as usize; - buf[self.singleton.get(pos) as usize] = true; - } + self.for_each_genome_in_row(slot, |g| buf[g] = true); } - /// Column-oriented per-genome k-mer totals — a naive row-by-row scan, - /// not the O(1)-per-column reduction the dense matrix's `count_ones` - /// is. Deliberately not optimised: see `DevDocMD/architecture/siblings.md` - /// and the sparse-matrix plan's "Explicitly deferred" — column-side - /// access stays correct but slow on this type for now. pub fn count_ones(&self) -> Array1 { - let mut counts = vec![0u64; self.n_cols]; - let mut buf = vec![false; self.n_cols]; - for slot in 0..self.n { - self.fill_row_bool(slot, &mut buf); - for (c, &present) in buf.iter().enumerate() { - if present { - counts[c] += 1; + self.col_weights_and_pair_counts().0 + } + + /// Row-first accumulation that goes straight through the sparse + /// encoding instead of `fill_row`+full-width scan: singleton rows + /// (`is_multi[slot] == false`) can never contribute a pair, so they + /// only bump `col_weights`; each *distinct* multi-genome set is + /// decoded and paired exactly once — weighted by how many k-mer slots + /// share it — instead of replaying the pairing work once per + /// occurrence. Complexity `O(n_singleton + n_distinct_multi * k̄²)` + /// instead of `O(n * n_cols)`. + fn col_weights_and_pair_counts(&self) -> (Array1, Array2) { + let n = self.n_cols; + let mut counts = vec![0u64; n]; + let mut inter = Array2::::zeros((n, n)); + + for pos in 0..self.singleton.len() { + counts[self.singleton.get(pos) as usize] += 1; + } + + let mut weight = vec![0u64; self.n_distinct_multi]; + for pos in 0..self.multi.len() { + weight[self.multi.get(pos) as usize] += 1; + } + + let mut present = Vec::new(); + for (dict_id, &w) in weight.iter().enumerate() { + if w == 0 { + continue; + } + let (start, end) = self.dict_entry_range(dict_id); + present.clear(); + let mut p = start; + while p < end { + present.push(read_varint(&self.dict_values, &mut p) as usize); + } + for &g in &present { + counts[g] += w; + } + for p in 0..present.len() { + let i = present[p]; + for &j in &present[p + 1..] { + inter[[i, j]] += w; + inter[[j, i]] += w; } } } - Array1::from(counts) + + (Array1::from(counts), inter) } /// Like [`super::PersistentBitMatrix::fill_sub_matrix`]: `out` has one /// entry per genome column, each filled with that column's values at - /// `slots`, in `slots` order. Naive (row-major decode, scattered into - /// column buffers) — see [`count_ones`](Self::count_ones)'s doc comment. + /// `slots`, in `slots` order. Row-major decode via + /// [`for_each_genome_in_row`](Self::for_each_genome_in_row) — touches + /// only the columns actually present per row, not an `n_cols`-wide + /// buffer. pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec]) { assert_eq!(out.len(), self.n_cols); for col in out.iter_mut() { col.clear(); col.resize(slots.len(), false); } - let mut buf = vec![false; self.n_cols]; for (i, &slot) in slots.iter().enumerate() { - self.fill_row_bool(slot, &mut buf); - for (c, &present) in buf.iter().enumerate() { - out[c][i] = present; - } + self.for_each_genome_in_row(slot, |g| out[g][i] = true); } } } @@ -384,19 +412,7 @@ impl ColumnWeights for PersistentSparseBitMatrix { 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 (col_weights, inter) = self.col_weights_and_pair_counts(); let mut union = Array2::zeros((n, n)); for i in 0..n { for j in 0..n { @@ -408,19 +424,7 @@ impl BitPartials for PersistentSparseBitMatrix { 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 (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 { diff --git a/src/obicompactvec/src/traits.rs b/src/obicompactvec/src/traits.rs index dc958e45..38f3afb1 100644 --- a/src/obicompactvec/src/traits.rs +++ b/src/obicompactvec/src/traits.rs @@ -166,7 +166,6 @@ 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 { @@ -183,9 +182,7 @@ 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 { - 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) + jaccard_to_mash(&self.jaccard_dist_matrix(), k) } fn hamming_dist_matrix(&self) -> Array2 { diff --git a/src/obikindex/src/distance.rs b/src/obikindex/src/distance.rs index 8ae1d23f..f5d63be5 100644 --- a/src/obikindex/src/distance.rs +++ b/src/obikindex/src/distance.rs @@ -121,7 +121,6 @@ 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 bf46144c..cfeaf816 100644 --- a/src/obikmer/src/cmd/phylo/mod.rs +++ b/src/obikmer/src/cmd/phylo/mod.rs @@ -294,17 +294,12 @@ 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 c335a2ff..4747af87 100644 --- a/src/obilayeredmap/src/layered_store.rs +++ b/src/obilayeredmap/src/layered_store.rs @@ -22,16 +22,10 @@ impl LayeredStore { impl ColumnWeights for LayeredStore { fn col_weights(&self) -> Array1 { - let parts: Vec> = self.0.par_iter() + self.0.par_iter() .map(|s| s.col_weights()) - .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 + .reduce_with(|a, b| a + b) + .unwrap_or_else(|| Array1::zeros(0)) } } @@ -39,87 +33,45 @@ impl ColumnWeights for LayeredStore { impl CountPartials for LayeredStore { fn partial_bray(&self) -> Array2 { - let parts: Vec<_> = self.0.par_iter() + self.0.par_iter() .map(|s| s.partial_bray()) - .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 + .reduce_with(|a, b| a + b) + .unwrap() } fn partial_euclidean(&self) -> Array2 { - let parts: Vec<_> = self.0.par_iter() + self.0.par_iter() .map(|s| s.partial_euclidean()) - .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 + .reduce_with(|a, b| a + b) + .unwrap() } fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2, Array2) { - let parts: Vec<_> = self.0.par_iter() + self.0.par_iter() .map(|s| s.partial_threshold_jaccard(threshold)) - .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) + .reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu)) + .unwrap() } fn partial_relfreq_bray(&self, global: &Array1) -> Array2 { - let parts: Vec<_> = self.0.par_iter() + self.0.par_iter() .map(|s| s.partial_relfreq_bray(global)) - .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 + .reduce_with(|a, b| a + b) + .unwrap() } fn partial_relfreq_euclidean(&self, global: &Array1) -> Array2 { - let parts: Vec<_> = self.0.par_iter() + self.0.par_iter() .map(|s| s.partial_relfreq_euclidean(global)) - .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 + .reduce_with(|a, b| a + b) + .unwrap() } fn partial_hellinger(&self, global: &Array1) -> Array2 { - let parts: Vec<_> = self.0.par_iter() + self.0.par_iter() .map(|s| s.partial_hellinger(global)) - .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 + .reduce_with(|a, b| a + b) + .unwrap() } } @@ -127,31 +79,17 @@ impl CountPartials for LayeredStore { impl BitPartials for LayeredStore { fn partial_jaccard(&self) -> (Array2, Array2) { - let parts: Vec<_> = self.0.par_iter() + self.0.par_iter() .map(|s| s.partial_jaccard()) - .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) + .reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu)) + .unwrap() } fn partial_hamming(&self) -> Array2 { - let parts: Vec<_> = self.0.par_iter() + self.0.par_iter() .map(|s| s.partial_hamming()) - .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 + .reduce_with(|a, b| a + b) + .unwrap() } }