From 9abee87af3abcfbfe5c650171ea226538acf3dbe Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 20 Aug 2026 06:00:03 +0200 Subject: [PATCH] refactor: delegate distance computations and simplify matrix logic Consolidate bit matrix operations by delegating pairwise metric calculations to IntSliceView and col_view implementations. Introduce chunked_presence_count in colgroup to handle threshold-based accumulation efficiently. Add a direct point lookup method to the sparse matrix representation to eliminate buffer allocations, and simplify match arms in persistent accessors accordingly. All public APIs and behavioral contracts remain unchanged. --- src/obicompactvec/src/bitmatrix/group_ops.rs | 22 +----- src/obicompactvec/src/bitmatrix/persistent.rs | 6 +- src/obicompactvec/src/bitmatrix/sparse.rs | 12 +++ src/obicompactvec/src/colgroup.rs | 29 ++++++- src/obicompactvec/src/intmatrix.rs | 74 +++++------------ src/obicompactvec/src/reader.rs | 79 +++++-------------- 6 files changed, 81 insertions(+), 141 deletions(-) diff --git a/src/obicompactvec/src/bitmatrix/group_ops.rs b/src/obicompactvec/src/bitmatrix/group_ops.rs index 66f12edd..ac449cc2 100644 --- a/src/obicompactvec/src/bitmatrix/group_ops.rs +++ b/src/obicompactvec/src/bitmatrix/group_ops.rs @@ -1,6 +1,6 @@ use std::io; -use crate::colgroup::{ColGroup, MatrixGroupOps}; +use crate::colgroup::{chunked_presence_count, ColGroup, MatrixGroupOps}; use crate::tempbitvec::{TempBitVec, TempBitVecBuilder}; use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder}; @@ -11,25 +11,7 @@ use super::persistent::PersistentBitMatrix; impl MatrixGroupOps for PersistentBitMatrix { fn partial_group_presence_count(&self, g: &ColGroup, _threshold: u32) -> io::Result { // Bit matrices store 0/1 — threshold is structurally always 1. - let n = self.n(); - if g.indices.len() < 255 { - let mut builder = TempCompactIntVecBuilder::new(n)?; - for &c in &g.indices { - builder.inc_present_fast(self.col_view(c)); - } - builder.freeze() - } else { - let mut result = TempCompactIntVecBuilder::new(n)?; - for chunk in g.indices.chunks(254) { - let mut chunk_b = TempCompactIntVecBuilder::new(n)?; - for &c in chunk { - chunk_b.inc_present_fast(self.col_view(c)); - } - let frozen = chunk_b.freeze()?; - result.add(frozen.view()); - } - result.freeze() - } + chunked_presence_count(self.n(), &g.indices, |b, c| b.inc_present_fast(self.col_view(c))) } fn partial_group_sum(&self, g: &ColGroup) -> io::Result { diff --git a/src/obicompactvec/src/bitmatrix/persistent.rs b/src/obicompactvec/src/bitmatrix/persistent.rs index b6cc09c5..71fa4113 100644 --- a/src/obicompactvec/src/bitmatrix/persistent.rs +++ b/src/obicompactvec/src/bitmatrix/persistent.rs @@ -111,11 +111,7 @@ 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::Sparse(m) => m.get(c, slot), Self::Implicit { .. } => 1, } } diff --git a/src/obicompactvec/src/bitmatrix/sparse.rs b/src/obicompactvec/src/bitmatrix/sparse.rs index 64835954..409ef6fb 100644 --- a/src/obicompactvec/src/bitmatrix/sparse.rs +++ b/src/obicompactvec/src/bitmatrix/sparse.rs @@ -151,6 +151,18 @@ impl PersistentSparseBitMatrix { (start, end) } + /// Column-major point lookup: value at column `c`, slot `slot`, as 0/1. + /// + /// Avoids `fill_row`'s `n_cols`-wide buffer: walks the row's own decode + /// — 1 entry for a singleton row, k̄ entries for a multi-genome row — + /// via [`for_each_genome_in_row`](Self::for_each_genome_in_row), instead + /// of allocating and scanning a full-width buffer for a single cell. + pub fn get(&self, c: usize, slot: usize) -> u32 { + let mut found = false; + self.for_each_genome_in_row(slot, |g| found |= g == c); + found as u32 + } + pub fn row(&self, slot: usize) -> Box<[bool]> { let mut out = vec![false; self.n_cols]; self.fill_row_bool(slot, &mut out); diff --git a/src/obicompactvec/src/colgroup.rs b/src/obicompactvec/src/colgroup.rs index 49ca4770..11c0f73a 100644 --- a/src/obicompactvec/src/colgroup.rs +++ b/src/obicompactvec/src/colgroup.rs @@ -1,7 +1,7 @@ use std::io; use crate::tempbitvec::{TempBitVec, TempBitVecBuilder}; -use crate::tempintvec::TempCompactIntVec; +use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder}; // ── ColGroup ────────────────────────────────────────────────────────────────── @@ -21,6 +21,33 @@ impl ColGroup { } } +/// Shared chunked-accumulation body behind every `partial_group_presence_count` +/// impl: below 255 columns, accumulate directly (`u8` counters never +/// overflow); at or above, split into chunks of 254, accumulate each chunk +/// into its own builder, and merge chunks with `add`. `inc` supplies the +/// one thing that differs per matrix kind — how to bump the counter for +/// column `c` — so the chunking/merging logic itself is written once. +pub(crate) fn chunked_presence_count( + n: usize, + indices: &[usize], + mut inc: impl FnMut(&mut TempCompactIntVecBuilder, usize), +) -> io::Result { + if indices.len() < 255 { + let mut builder = TempCompactIntVecBuilder::new(n)?; + for &c in indices { inc(&mut builder, c); } + builder.freeze() + } else { + let mut result = TempCompactIntVecBuilder::new(n)?; + for chunk in indices.chunks(254) { + let mut chunk_b = TempCompactIntVecBuilder::new(n)?; + for &c in chunk { inc(&mut chunk_b, c); } + let frozen = chunk_b.freeze()?; + result.add(frozen.view()); + } + result.freeze() + } +} + // ── MatrixGroupOps ──────────────────────────────────────────────────────────── /// Per-matrix group aggregations. diff --git a/src/obicompactvec/src/intmatrix.rs b/src/obicompactvec/src/intmatrix.rs index f26f3228..deb2c31f 100644 --- a/src/obicompactvec/src/intmatrix.rs +++ b/src/obicompactvec/src/intmatrix.rs @@ -8,7 +8,7 @@ use rayon::prelude::*; use crate::bitmatrix::{pairwise_matrix, pairwise2_matrix}; use crate::builder::PersistentCompactIntVecBuilder; -use crate::colgroup::{ColGroup, MatrixGroupOps}; +use crate::colgroup::{chunked_presence_count, ColGroup, MatrixGroupOps}; use crate::format::{HEADER_SIZE, OVERFLOW_ENTRY_SIZE}; use crate::meta::MatrixMeta; use crate::reader::PersistentCompactIntVec; @@ -186,53 +186,33 @@ impl PackedCompactIntMatrix { ) } - fn pair_partial_bray(&self, i: usize, j: usize) -> u64 { - self.col_view(i).iter().zip(self.col_view(j).iter()).map(|(a, b)| a.min(b) as u64).sum() - } - fn pair_partial_euclidean(&self, i: usize, j: usize) -> f64 { - self.col_view(i).iter().zip(self.col_view(j).iter()) - .map(|(a, b)| { let d = a as f64 - b as f64; d * d }).sum() - } - fn pair_partial_threshold_jaccard(&self, i: usize, j: usize, t: u32) -> (u64, u64) { - self.col_view(i).iter().zip(self.col_view(j).iter()) - .fold((0u64, 0u64), |(inter, uni), (a, b)| { - let ap = a >= t; let bp = b >= t; - (inter + (ap & bp) as u64, uni + (ap | bp) as u64) - }) - } - fn pair_partial_relfreq_bray(&self, i: usize, j: usize, si: f64, sj: f64) -> f64 { - if si == 0.0 || sj == 0.0 { return 0.0; } - self.col_view(i).iter().zip(self.col_view(j).iter()) - .map(|(a, b)| (a as f64 / si).min(b as f64 / sj)).sum() - } - fn pair_partial_relfreq_euclidean(&self, i: usize, j: usize, si: f64, sj: f64) -> f64 { - if si == 0.0 || sj == 0.0 { return 0.0; } - self.col_view(i).iter().zip(self.col_view(j).iter()) - .map(|(a, b)| { let d = a as f64 / si - b as f64 / sj; d * d }).sum() - } - fn pair_partial_hellinger(&self, i: usize, j: usize, si: f64, sj: f64) -> f64 { - if si == 0.0 || sj == 0.0 { return 0.0; } - self.col_view(i).iter().zip(self.col_view(j).iter()) - .map(|(a, b)| { let d = (a as f64 / si).sqrt() - (b as f64 / sj).sqrt(); d * d }).sum() - } + // Every pair formula below delegates to `IntSliceView` (`views.rs`) — + // the canonical implementation, also used by `ColumnarCompactIntMatrix` + // via `PersistentCompactIntVec`. No formula is re-derived here. pub(crate) fn partial_bray_dist_matrix(&self) -> Array2 { - pairwise_matrix(self.n_cols, |i, j| self.pair_partial_bray(i, j)) + pairwise_matrix(self.n_cols, |i, j| self.col_view(i).partial_bray_dist(self.col_view(j))) } pub(crate) fn partial_euclidean_dist_matrix(&self) -> Array2 { - pairwise_matrix(self.n_cols, |i, j| self.pair_partial_euclidean(i, j)) + pairwise_matrix(self.n_cols, |i, j| self.col_view(i).partial_euclidean_dist(self.col_view(j))) } pub(crate) fn partial_threshold_jaccard_dist_matrix(&self, t: u32) -> (Array2, Array2) { - pairwise2_matrix(self.n_cols, |i, j| self.pair_partial_threshold_jaccard(i, j, t)) + pairwise2_matrix(self.n_cols, |i, j| self.col_view(i).partial_threshold_jaccard_dist(self.col_view(j), t)) } pub(crate) fn partial_relfreq_bray_dist_matrix(&self, col_sums: &Array1) -> Array2 { - pairwise_matrix(self.n_cols, |i, j| self.pair_partial_relfreq_bray(i, j, col_sums[i] as f64, col_sums[j] as f64)) + pairwise_matrix(self.n_cols, |i, j| { + self.col_view(i).partial_relfreq_bray_dist(self.col_view(j), col_sums[i] as f64, col_sums[j] as f64) + }) } pub(crate) fn partial_relfreq_euclidean_dist_matrix(&self, col_sums: &Array1) -> Array2 { - pairwise_matrix(self.n_cols, |i, j| self.pair_partial_relfreq_euclidean(i, j, col_sums[i] as f64, col_sums[j] as f64)) + pairwise_matrix(self.n_cols, |i, j| { + self.col_view(i).partial_relfreq_euclidean_dist(self.col_view(j), col_sums[i] as f64, col_sums[j] as f64) + }) } pub(crate) fn partial_hellinger_euclidean_dist_matrix(&self, col_sums: &Array1) -> Array2 { - pairwise_matrix(self.n_cols, |i, j| self.pair_partial_hellinger(i, j, col_sums[i] as f64, col_sums[j] as f64)) + pairwise_matrix(self.n_cols, |i, j| { + self.col_view(i).partial_hellinger_euclidean_dist(self.col_view(j), col_sums[i] as f64, col_sums[j] as f64) + }) } } @@ -531,25 +511,9 @@ impl PersistentCompactIntMatrixBuilder { impl MatrixGroupOps for PersistentCompactIntMatrix { fn partial_group_presence_count(&self, g: &ColGroup, threshold: u32) -> io::Result { - let n = self.n(); - if g.indices.len() < 255 { - let mut builder = TempCompactIntVecBuilder::new(n)?; - for &c in &g.indices { - builder.inc_predicate_fast(self.col_view(c), |v| v >= threshold); - } - builder.freeze() - } else { - let mut result = TempCompactIntVecBuilder::new(n)?; - for chunk in g.indices.chunks(254) { - let mut chunk_b = TempCompactIntVecBuilder::new(n)?; - for &c in chunk { - chunk_b.inc_predicate_fast(self.col_view(c), |v| v >= threshold); - } - let frozen = chunk_b.freeze()?; - result.add(frozen.view()); - } - result.freeze() - } + chunked_presence_count(self.n(), &g.indices, |b, c| { + b.inc_predicate_fast(self.col_view(c), |v| v >= threshold) + }) } fn partial_group_sum(&self, g: &ColGroup) -> io::Result { diff --git a/src/obicompactvec/src/reader.rs b/src/obicompactvec/src/reader.rs index 78f51e4a..e69ae0ec 100644 --- a/src/obicompactvec/src/reader.rs +++ b/src/obicompactvec/src/reader.rs @@ -138,107 +138,66 @@ impl PersistentCompactIntVec { } // ── Distance methods ────────────────────────────────────────────────────── + // + // Thin delegates to `IntSliceView` (see `views.rs`) — the canonical, + // zero-copy implementation of every formula below. Kept here only for + // API convenience (`&PersistentCompactIntVec` in, no `.view()` needed at + // call sites); no formula is re-derived. pub fn bray_dist(&self, other: &PersistentCompactIntVec) -> f64 { - let sum_min = self.partial_bray_dist(other); - let denom = self.sum() + other.sum(); - if denom == 0 { 0.0 } else { 1.0 - 2.0 * sum_min as f64 / denom as f64 } + self.view().bray_dist(other.view()) } pub fn partial_bray_dist(&self, other: &PersistentCompactIntVec) -> u64 { - assert_eq!(self.n, other.len(), "length mismatch"); - self.iter().zip(other.iter()).map(|(a, b)| a.min(b) as u64).sum() + self.view().partial_bray_dist(other.view()) } pub fn relfreq_bray_dist(&self, other: &PersistentCompactIntVec) -> f64 { - assert_eq!(self.n, other.len(), "length mismatch"); - let sa = self.sum() as f64; - let sb = other.sum() as f64; - if sa == 0.0 && sb == 0.0 { return 0.0; } - 1.0 - self.partial_relfreq_bray_dist(other, sa, sb) + self.view().relfreq_bray_dist(other.view()) } pub fn partial_relfreq_bray_dist(&self, other: &PersistentCompactIntVec, sum_a: f64, sum_b: f64) -> f64 { - assert_eq!(self.n, other.len(), "length mismatch"); - self.iter().zip(other.iter()) - .map(|(a, b)| { - let pa = if sum_a > 0.0 { a as f64 / sum_a } else { 0.0 }; - let pb = if sum_b > 0.0 { b as f64 / sum_b } else { 0.0 }; - pa.min(pb) - }) - .sum() + self.view().partial_relfreq_bray_dist(other.view(), sum_a, sum_b) } pub fn euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 { - self.partial_euclidean_dist(other).sqrt() + self.view().euclidean_dist(other.view()) } pub fn partial_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 { - assert_eq!(self.n, other.len(), "length mismatch"); - self.iter().zip(other.iter()) - .map(|(a, b)| { let d = a as f64 - b as f64; d * d }) - .sum() + self.view().partial_euclidean_dist(other.view()) } pub fn relfreq_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 { - let sa = self.sum() as f64; - let sb = other.sum() as f64; - if sa == 0.0 && sb == 0.0 { return 0.0; } - self.partial_relfreq_euclidean_dist(other, sa, sb).sqrt() + self.view().relfreq_euclidean_dist(other.view()) } pub fn partial_relfreq_euclidean_dist(&self, other: &PersistentCompactIntVec, sum_a: f64, sum_b: f64) -> f64 { - assert_eq!(self.n, other.len(), "length mismatch"); - self.iter().zip(other.iter()) - .map(|(a, b)| { - let pa = if sum_a > 0.0 { a as f64 / sum_a } else { 0.0 }; - let pb = if sum_b > 0.0 { b as f64 / sum_b } else { 0.0 }; - let d = pa - pb; - d * d - }) - .sum() + self.view().partial_relfreq_euclidean_dist(other.view(), sum_a, sum_b) } pub fn hellinger_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 { - let sa = self.sum() as f64; - let sb = other.sum() as f64; - if sa == 0.0 && sb == 0.0 { return 0.0; } - self.partial_hellinger_euclidean_dist(other, sa, sb).sqrt() + self.view().hellinger_euclidean_dist(other.view()) } pub fn partial_hellinger_euclidean_dist(&self, other: &PersistentCompactIntVec, sum_a: f64, sum_b: f64) -> f64 { - assert_eq!(self.n, other.len(), "length mismatch"); - self.iter().zip(other.iter()) - .map(|(a, b)| { - let pa = if sum_a > 0.0 { (a as f64 / sum_a).sqrt() } else { 0.0 }; - let pb = if sum_b > 0.0 { (b as f64 / sum_b).sqrt() } else { 0.0 }; - let d = pa - pb; - d * d - }) - .sum() + self.view().partial_hellinger_euclidean_dist(other.view(), sum_a, sum_b) } pub fn hellinger_dist(&self, other: &PersistentCompactIntVec) -> f64 { - self.hellinger_euclidean_dist(other) / std::f64::consts::SQRT_2 + self.view().hellinger_dist(other.view()) } pub fn threshold_jaccard_dist(&self, other: &PersistentCompactIntVec, threshold: u32) -> f64 { - let (intersection, union) = self.partial_threshold_jaccard_dist(other, threshold); - if union == 0 { 0.0 } else { 1.0 - intersection as f64 / union as f64 } + self.view().threshold_jaccard_dist(other.view(), threshold) } pub fn partial_threshold_jaccard_dist(&self, other: &PersistentCompactIntVec, threshold: u32) -> (u64, u64) { - assert_eq!(self.n, other.len(), "length mismatch"); - self.iter().zip(other.iter()) - .fold((0u64, 0u64), |(inter, uni), (a, b)| { - let ap = a >= threshold; - let bp = b >= threshold; - (inter + (ap & bp) as u64, uni + (ap | bp) as u64) - }) + self.view().partial_threshold_jaccard_dist(other.view(), threshold) } pub fn jaccard_dist(&self, other: &PersistentCompactIntVec) -> f64 { - self.threshold_jaccard_dist(other, 1) + self.view().jaccard_dist(other.view()) } }