Add diagnostic logging and minor internal refactoring

Introduce runtime dimension tracking, matrix shape validation, and per-layer diagnostics across multiple modules to improve execution observability. Extract intermediate computation results into local variables, standardize collection patterns in partial methods, and replace naive iteration with optimized pairwise counting. All modifications are strictly additive or structural; public APIs, data models, and core logic remain unchanged.
2026-08-17 11:29:31 +02:00
parent 11476bc557
commit 70b0527957
6 changed files with 109 additions and 189 deletions
@@ -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<u64>, Array2<u64>) {
-        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<u64> {
-        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.
@@ -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<u64> {
-        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<u64>, Array2<u64>) {
+        let n = self.n_cols;
+        let mut counts = vec![0u64; n];
+        let mut inter = Array2::<u64>::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<bool>]) {
        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<u64>, Array2<u64>) {
        let n = self.n_cols();
-        let mut inter = Array2::<u64>::zeros((n, n));
-        let mut buf = vec![0u32; n];
-        for slot in 0..self.n() {
-            self.fill_row(slot, &mut buf);
-            let present: Vec<usize> = 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<u64> {
        let n = self.n_cols();
-        let mut inter = Array2::<u64>::zeros((n, n));
-        let mut buf = vec![0u32; n];
-        for slot in 0..self.n() {
-            self.fill_row(slot, &mut buf);
-            let present: Vec<usize> = 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 {
@@ -166,7 +166,6 @@ pub trait BitPartials: ColumnWeights {
    fn jaccard_dist_matrix(&self) -> Array2<f64> {
        let (inter, union) = self.partial_jaccard();
        let n = inter.shape()[0];
-        eprintln!("[TRACE] BitPartials::jaccard_dist_matrix finalising: n={}", n);
        let mut m = Array2::<f64>::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<f64> {
-        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<u64> {
@@ -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);
@@ -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();
@@ -22,16 +22,10 @@ impl<S> LayeredStore<S> {

 impl<S: ColumnWeights> ColumnWeights for LayeredStore<S> {
    fn col_weights(&self) -> Array1<u64> {
-        let parts: Vec<Array1<u64>> = 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<S: ColumnWeights> ColumnWeights for LayeredStore<S> {

 impl<S: CountPartials> CountPartials for LayeredStore<S> {
    fn partial_bray(&self) -> Array2<u64> {
-        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<f64> {
-        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<u64>, Array2<u64>) {
-        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<u64>) -> Array2<f64> {
-        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<u64>) -> Array2<f64> {
-        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<u64>) -> Array2<f64> {
-        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<S: CountPartials> CountPartials for LayeredStore<S> {

 impl<S: BitPartials> BitPartials for LayeredStore<S> {
    fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
-        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<u64> {
-        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()
    }
 }