OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 34 Wiki Activity

refactor: delegate distance computations and simplify matrix logic

Browse Source 
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.
This commit is contained in:
Eric Coissac
2026-08-20 14:27:32 +02:00
parent 19f9954050
commit 9abee87af3
6 changed files with 81 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obicompactvec/src/bitmatrix/group_ops.rs
+2 -20
View File
@@ -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<TempCompactIntVec> {
// 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<TempCompactIntVec> {
src/obicompactvec/src/bitmatrix/persistent.rs
+1 -5
View File
@@ -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,
}
}
src/obicompactvec/src/bitmatrix/sparse.rs
+12
View File
@@ -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);
src/obicompactvec/src/colgroup.rs
+28 -1
View File
@@ -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<TempCompactIntVec> {
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.
src/obicompactvec/src/intmatrix.rs
+19 -55
View File
@@ -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<u64> {
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<f64> {
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<u64>, Array2<u64>) {
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<u64>) -> Array2<f64> {
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<u64>) -> Array2<f64> {
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<u64>) -> Array2<f64> {
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<TempCompactIntVec> {
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<TempCompactIntVec> {
src/obicompactvec/src/reader.rs
+19 -60
View File
@@ -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())
}
}