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

Push zpwxxpnpktps #67

Merged
coissac merged 46 commits from push-zpwxxpnpktps into main 2026-08-17 09:41:42 +00:00
Conversation 0 Commits 46 Files Changed 405
+273 -361
7 changed files with 273 additions and 361 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit cd57cf0cbd - Show all commits
src/obikindex/src/siblings/alignment.rs
+9 -87
View File
@@ -1,19 +1,13 @@
use rayon::prelude::*; use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikpartitionner::KmerPartition; use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use obisys::progress_bar; use obisys::progress_bar;
use crate::error::{OKIError, OKIResult}; use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex; use crate::index::KmerIndex;
use super::cache::{Mat, PartitionCache}; use super::cache::PartitionCache;
use super::helpers::{central_base, is_minorant, partition_of}; use super::family_scan::scan_layer_families;
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set /// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set
/// iff the genome carries the member whose own central base is `b`): /// iff the genome carries the member whose own central base is `b`):
@@ -77,26 +71,7 @@ impl KmerIndex {
) )
.map_err(OKIError::Partition)?; .map_err(OKIError::Partition)?;
let cache = PartitionCache::build(&partition, n_parts, with_counts)?; let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
let layer_dirs = self.sibling_layer_dirs()?;
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
}
let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers"); let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers");
// `Vec<Vec<u8>>` per layer, one entry (column) per variable family; // `Vec<Vec<u8>>` per layer, one entry (column) per variable family;
@@ -106,65 +81,12 @@ impl KmerIndex {
let partials: Vec<Vec<Vec<u8>>> = layer_dirs let partials: Vec<Vec<Vec<u8>>> = layer_dirs
.par_iter() .par_iter()
.map(|layer_dir| -> OKIResult<Vec<Vec<u8>>> { .map(|layer_dir| -> OKIResult<Vec<Vec<u8>>> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?;
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; let columns = families
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; .into_iter()
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; .filter(|f| f.mask.family_size() >= 2) // monomorphic family — no signal, skip
.map(|f| f.genome_mask.iter().map(|&m| iupac_code(m)).collect())
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()]; .collect();
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
.map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
};
let n_cols = mat.n_cols().min(n_genomes);
let mut columns: Vec<Vec<u8>> = Vec::new();
let mut genome_mask: Vec<u8> = Vec::with_capacity(n_genomes);
for slot in 0..annex.len() {
let Some(mask) = annex.get(slot) else { continue };
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // family tallied once, at its minorant
}
if mask.family_size() < 2 {
continue; // monomorphic family — no signal, skip
}
genome_mask.clear();
genome_mask.resize(n_genomes, 0);
for other in kmer.central_canonical_neighbors() {
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let presence: Option<Vec<bool>> = if other == kmer {
Some((0..n_cols).map(|g| mat.carries(g, slot)).collect())
} else {
let dest = partition_of(other, n_parts);
cache.find_presence(dest, other, n_genomes)
};
let Some(presence) = presence else { continue };
for (g, &present) in presence.iter().enumerate() {
if present {
genome_mask[g] |= 1 << base;
}
}
}
columns.push(genome_mask.iter().map(|&m| iupac_code(m)).collect());
}
pb.inc(1); pb.inc(1);
Ok(columns) Ok(columns)
src/obikindex/src/siblings/cardinality.rs
+7 -80
View File
@@ -1,21 +1,15 @@
use ndarray::Array2; use ndarray::Array2;
use rayon::prelude::*; use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikpartitionner::KmerPartition; use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use obisys::progress_bar; use obisys::progress_bar;
use crate::error::{OKIError, OKIResult}; use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex; use crate::index::KmerIndex;
use super::cache::{Mat, PartitionCache}; use super::cache::PartitionCache;
use super::distance::RawSnpDistanceOutput; use super::distance::RawSnpDistanceOutput;
use super::helpers::{central_base, is_minorant, partition_of}; use super::family_scan::scan_layer_families;
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// See [`KmerIndex::cardinality_tally`]. /// See [`KmerIndex::cardinality_tally`].
pub struct CardinalityTally { pub struct CardinalityTally {
@@ -80,63 +74,17 @@ impl KmerIndex {
) )
.map_err(OKIError::Partition)?; .map_err(OKIError::Partition)?;
let cache = PartitionCache::build(&partition, n_parts, with_counts)?; let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
let layer_dirs = self.sibling_layer_dirs()?;
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
}
let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers"); let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers");
let partials: Vec<[[u64; 5]; 5]> = layer_dirs let partials: Vec<[[u64; 5]; 5]> = layer_dirs
.par_iter() .par_iter()
.map(|layer_dir| -> OKIResult<[[u64; 5]; 5]> { .map(|layer_dir| -> OKIResult<[[u64; 5]; 5]> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?;
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
.map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
};
let n_cols = mat.n_cols().min(n_genomes);
let mut counts = [[0u64; 5]; 5]; let mut counts = [[0u64; 5]; 5];
let mut genome_mask: Vec<u8> = Vec::with_capacity(n_genomes);
for slot in 0..annex.len() { for family in &families {
let Some(mask) = annex.get(slot) else { continue }; if family.mask.family_size() < 2 {
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // family tallied once, at its minorant
}
if mask.family_size() < 2 {
// Fully invariant family (never varies anywhere in // Fully invariant family (never varies anywhere in
// the index) — genome-wide background, not // the index) — genome-wide background, not
// SNP-adjacent signal; would otherwise swamp the // SNP-adjacent signal; would otherwise swamp the
@@ -148,28 +96,7 @@ impl KmerIndex {
continue; continue;
} }
genome_mask.clear(); let genome_mask = &family.genome_mask;
genome_mask.resize(n_genomes, 0);
for other in kmer.central_canonical_neighbors() {
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let presence: Option<Vec<bool>> = if other == kmer {
Some((0..n_cols).map(|g| mat.carries(g, slot)).collect())
} else {
let dest = partition_of(other, n_parts);
cache.find_presence(dest, other, n_genomes)
};
let Some(presence) = presence else { continue };
for (g, &present) in presence.iter().enumerate() {
if present {
genome_mask[g] |= 1 << base;
}
}
}
for i in 0..n_genomes { for i in 0..n_genomes {
let card_i = genome_mask[i].count_ones() as usize; let card_i = genome_mask[i].count_ones() as usize;
for j in (i + 1)..n_genomes { for j in (i + 1)..n_genomes {
src/obikindex/src/siblings/distance.rs
+18 -100
View File
@@ -1,20 +1,14 @@
use ndarray::Array2; use ndarray::Array2;
use rayon::prelude::*; use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikpartitionner::KmerPartition; use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use obisys::progress_bar; use obisys::progress_bar;
use crate::error::{OKIError, OKIResult}; use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex; use crate::index::KmerIndex;
use super::cache::{Mat, PartitionCache}; use super::cache::PartitionCache;
use super::helpers::{central_base, is_minorant, partition_of}; use super::family_scan::scan_layer_families;
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// Raw p-distance restricted to loci that are single-copy in **both** /// Raw p-distance restricted to loci that are single-copy in **both**
/// genomes of a pair — the "stringent / paralogy-aware" locus eligibility /// genomes of a pair — the "stringent / paralogy-aware" locus eligibility
@@ -84,108 +78,32 @@ impl KmerIndex {
) )
.map_err(OKIError::Partition)?; .map_err(OKIError::Partition)?;
let cache = PartitionCache::build(&partition, n_parts, with_counts)?; let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
let layer_dirs = self.sibling_layer_dirs()?;
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
}
let pb = progress_bar(label, layer_dirs.len() as u64, "layers"); let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
let partials: Vec<Acc> = layer_dirs let partials: Vec<Acc> = layer_dirs
.par_iter() .par_iter()
.map(|layer_dir| -> OKIResult<Acc> { .map(|layer_dir| -> OKIResult<Acc> {
let mut acc = zero(); let mut acc = zero();
let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?;
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); for family in &families {
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; let variable = family.mask.family_size() >= 2;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; let genome_mask = &family.genome_mask;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()]; // Per genome: which single form (if exactly one) it
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) // carries — `None` (a `popcount != 1` mask) once a
.map_err(OKIError::Partition)?; // second form is seen, ambiguous/not single-copy,
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { // ineligible for either side of a pair.
if let Some(slot) = mphf.find(kmer) { let single_form = |g: usize| -> Option<u8> {
slot_kmer[slot] = Some(kmer); let m = genome_mask[g];
} (m.count_ones() == 1).then(|| m.trailing_zeros() as u8)
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
}; };
let n_cols = mat.n_cols().min(n_genomes);
// Per family: which single form (if exactly one) each genome for i in 0..n_genomes {
// carries — `None` once a second form is seen (ambiguous, let Some(bi) = single_form(i) else { continue };
// not single-copy, ineligible for either side of a pair). for j in (i + 1)..n_genomes {
let mut single_form: Vec<Option<u8>> = Vec::with_capacity(n_cols); let Some(bj) = single_form(j) else { continue };
let mut ambiguous: Vec<bool> = Vec::with_capacity(n_cols);
for slot in 0..annex.len() {
let Some(mask) = annex.get(slot) else { continue };
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // family tallied once, at its minorant
}
let variable = mask.family_size() >= 2;
single_form.clear();
single_form.resize(n_cols, None);
ambiguous.clear();
ambiguous.resize(n_cols, false);
for other in kmer.central_canonical_neighbors() {
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let presence: Option<Vec<bool>> = if other == kmer {
Some((0..n_cols).map(|g| mat.carries(g, slot)).collect())
} else {
let dest = partition_of(other, n_parts);
cache.find_presence(dest, other, n_genomes)
};
let Some(presence) = presence else { continue };
for (g, &present) in presence.iter().enumerate() {
if !present {
continue;
}
if single_form[g].is_some() {
ambiguous[g] = true;
} else {
single_form[g] = Some(base);
}
}
}
for i in 0..n_cols {
if ambiguous[i] {
continue;
}
let Some(bi) = single_form[i] else { continue };
for j in (i + 1)..n_cols {
if ambiguous[j] {
continue;
}
let Some(bj) = single_form[j] else { continue };
on_pair(&mut acc, i, j, bi, bj, variable); on_pair(&mut acc, i, j, bi, bj, variable);
} }
} }
src/obikindex/src/siblings/family_scan.rs
+178
View File
@@ -0,0 +1,178 @@
//! Shared per-layer family traversal, used by every sibling-annex consumer
//! (`snp_pseudo_alignment`, `cardinality_tally`, `scan_family_pairs`,
//! `sibling_annex_stats`) — resolves each minorant family's per-genome
//! base-presence, with the same locality discipline as
//! `build_sibling_annex` (see `cache.rs`/`build.rs`'s docs): cross-partition
//! lookups are grouped by destination partition and resolved in one
//! contiguous sweep per partition, instead of one lookup at a time jumping
//! between partitions in family order. The naive, per-family inline lookup
//! was the shape all four consumers used to have independently; sampling a
//! real ~170 GB index against 137 GB of RAM showed it thrashing the page
//! cache (near-continuous ~2 GB/s pagein), the same failure mode
//! `build_sibling_annex` had already been fixed for.
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU8, Ordering};
use rayon::prelude::*;
use obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use super::cache::{Mat, PartitionCache};
use super::helpers::{central_base, is_minorant, partition_of};
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// One layer's minorant family, with its own annex mask and, per genome,
/// the 4-bit base-presence mask resolved across the whole index (bit `b`
/// set iff that genome carries the family member whose own canonical
/// central base is `b`).
pub(super) struct FamilyRow {
pub mask: FamilyMask,
pub genome_mask: Vec<u8>,
}
impl KmerIndex {
/// Every (partition, layer) directory carrying a sibling annex, checked
/// up front so a missing one is reported before any real work starts.
pub(super) fn sibling_layer_dirs(&self) -> OKIResult<Vec<PathBuf>> {
let n_parts = self.n_partitions();
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
}
Ok(layer_dirs)
}
}
/// Every minorant family of one layer, each with its per-genome
/// base-presence mask resolved against the whole (already-open) partition
/// cache.
///
/// Two passes, not one inline pass: the first is purely local (already-open
/// annex/mphf/matrix, no lookups) and both enumerates the layer's families
/// and collects every cross-partition query they need, grouped by
/// destination partition; the second resolves each partition's whole batch
/// in one contiguous sweep, keeping that partition's mmap'd pages hot for
/// its entire batch instead of faulting them in and out as lookups jump
/// between partitions in family order.
pub(super) fn scan_layer_families(
layer_dir: &Path,
n_parts: usize,
n_genomes: usize,
with_counts: bool,
k: usize,
cache: &PartitionCache,
) -> OKIResult<Vec<FamilyRow>> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")).map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
};
let n_cols = mat.n_cols().min(n_genomes);
// ── Pass 1: local only — enumerate minorant families and collect every
// cross-partition query they need, grouped by destination partition.
struct Family {
slot: usize,
mask: FamilyMask,
}
let mut families: Vec<Family> = Vec::new();
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
for slot in 0..annex.len() {
let Some(mask) = annex.get(slot) else { continue };
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // family tallied once, at its minorant
}
let family_idx = families.len();
families.push(Family { slot, mask });
for other in kmer.central_canonical_neighbors() {
if other == kmer {
continue; // local — resolved below straight from `mat`, no lookup
}
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let dest = partition_of(other, n_parts);
outgoing[dest].push((other, family_idx, base));
}
}
let genome_mask: Vec<AtomicU8> = (0..families.len() * n_genomes).map(|_| AtomicU8::new(0)).collect();
// Local presence (the family's own minorant slot) — already open, no
// lookup, so no locality concern either way.
for (family_idx, family) in families.iter().enumerate() {
let kmer = slot_kmer[family.slot].expect("family slot has a kmer");
let base = central_base(kmer, k);
for g in 0..n_cols {
if mat.carries(g, family.slot) {
genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
}
}
}
// ── Pass 2: resolve each destination partition's whole batch in one
// contiguous sweep — same rationale as `build_sibling_annex`'s
// `outgoing` grouping.
outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
for &(variant, family_idx, base) in queries {
let Some(presence) = cache.find_presence(dest, variant, n_genomes) else { continue };
for (g, &present) in presence.iter().enumerate() {
if present {
genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
}
}
}
});
Ok(families
.into_iter()
.enumerate()
.map(|(family_idx, family)| FamilyRow {
mask: family.mask,
genome_mask: (0..n_genomes)
.map(|g| genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed))
.collect(),
})
.collect())
}
src/obikindex/src/siblings/mod.rs
+1
View File
@@ -47,6 +47,7 @@ mod build;
mod cache; mod cache;
mod cardinality; mod cardinality;
mod distance; mod distance;
mod family_scan;
mod helpers; mod helpers;
mod stats; mod stats;
src/obikindex/src/siblings/stats.rs
+10 -93
View File
@@ -1,19 +1,13 @@
use rayon::prelude::*; use rayon::prelude::*;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
use obikpartitionner::KmerPartition; use obikpartitionner::KmerPartition;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obiskio::UnitigFileReader;
use obisys::progress_bar; use obisys::progress_bar;
use crate::error::{OKIError, OKIResult}; use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex; use crate::index::KmerIndex;
use super::cache::{Mat, PartitionCache}; use super::cache::PartitionCache;
use super::helpers::{central_base, is_minorant, partition_of}; use super::family_scan::scan_layer_families;
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
/// Distribution of family sizes (1-4), read back from an already-built /// Distribution of family sizes (1-4), read back from an already-built
/// annex (see [`KmerIndex::build_sibling_annex`]) plus the index's /// annex (see [`KmerIndex::build_sibling_annex`]) plus the index's
@@ -62,29 +56,7 @@ impl KmerIndex {
) )
.map_err(OKIError::Partition)?; .map_err(OKIError::Partition)?;
let cache = PartitionCache::build(&partition, n_parts, with_counts)?; let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
let layer_dirs = self.sibling_layer_dirs()?;
// Gather the (partition, layer) pairs to process — cheap metadata
// reads only, checking every annex file exists up front so a
// missing one is reported before any real work starts.
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() {
continue;
}
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
"no sibling annex at {} — run build_sibling_annex first",
annex_path.display()
)));
}
layer_dirs.push(layer_dir);
}
}
// One layer's worth of work, parallelised across layers with Rayon // One layer's worth of work, parallelised across layers with Rayon
// — independent, read-only, each producing its own partial tally // — independent, read-only, each producing its own partial tally
@@ -98,70 +70,15 @@ impl KmerIndex {
..Default::default() ..Default::default()
}; };
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?;
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; for family in &families {
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
// Need each slot's own k-mer to derive minorant — same
// enumeration as construction.
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
.map_err(OKIError::Partition)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
slot_kmer[slot] = Some(kmer);
}
}
let use_counts = with_counts && layer_dir.join("counts").exists();
let mat = if use_counts {
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
} else {
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
};
let n_cols = mat.n_cols().min(n_genomes);
for slot in 0..annex.len() {
let Some(mask) = annex.get(slot) else { continue };
let Some(kmer) = slot_kmer[slot] else { continue };
if !is_minorant(kmer, mask, k) {
continue; // this family is tallied at its minorant's slot only
}
let s = mask.siblings() as usize;
stats.counts[s] += 1;
// "Genome g represents this family" means g carries // "Genome g represents this family" means g carries
// *any* of its members, not just the minorant's own — // *any* of its members, not just the minorant's own —
// start from the minorant's own presence (already // `genome_mask[g] != 0` is exactly that.
// open, no lookup) and OR in every other present let s = family.mask.siblings() as usize;
// member's presence vector, resolved against the stats.counts[s] += 1;
// whole-run cache (no I/O) — exactly `mask.siblings()` for (g, &m) in family.genome_mask.iter().enumerate() {
// of them, the mask tells us precisely which to fetch. if m != 0 {
let mut carries = vec![false; n_cols];
for g in 0..n_cols {
carries[g] = mat.carries(g, slot);
}
for other in kmer.central_canonical_neighbors() {
if other == kmer {
continue;
}
let base = central_base(other, k);
if !mask.has(base) {
continue;
}
let dest = partition_of(other, n_parts);
if let Some(other_presence) = cache.find_presence(dest, other, n_genomes) {
for (g, &present) in other_presence.iter().enumerate() {
if present {
carries[g] = true;
}
}
}
}
for (g, &carried) in carries.iter().enumerate() {
if carried {
stats.per_genome[g][s] += 1; stats.per_genome[g][s] += 1;
} }
} }
src/obikindex/src/siblings/tests.rs
+49
View File
@@ -184,3 +184,52 @@ fn sibling_annex_stats_counts_each_family_once_and_per_genome() {
); );
} }
} }
/// Exercises the four sibling-annex consumers that were rewritten to share
/// `family_scan::scan_layer_families` (partition-grouped lookups instead of
/// one lookup per family) — same one-sibling-each fixture as the tests
/// above (g1 carries the family's centre-C member, g2 the centre-G one),
/// hand-verified expected output for each.
#[test]
fn family_scan_consumers_agree_on_one_sibling_each() {
let dir = tempdir().unwrap();
let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
let merged = merge_two(dir.path(), &g1, &g2);
merged.build_sibling_annex().expect("build_sibling_annex");
// Merge doesn't promise to preserve source order, so resolve each
// genome's index by label rather than assuming g1 -> 0, g2 -> 1.
let idx_of = |label: &str| merged.meta.genomes.iter().position(|g| g.label == label).unwrap();
let (i1, i2) = (idx_of("g1"), idx_of("g2"));
// snp_pseudo_alignment: one variable family, one column — g1's row
// reads 'C' (its own member), g2's reads 'G'.
let alignment = merged.snp_pseudo_alignment().expect("snp_pseudo_alignment");
assert_eq!(alignment.sequences[i1], vec![b'C']);
assert_eq!(alignment.sequences[i2], vec![b'G']);
// raw_snp_distance: g1's single form (C) != g2's (G) at the family's
// one eligible locus -> a SNP, not a shared site.
let raw = merged.raw_snp_distance().expect("raw_snp_distance");
assert_eq!(raw.snp[[i1, i2]], 1);
assert_eq!(raw.snp[[i2, i1]], 1);
assert_eq!(raw.shared[[i1, i2]], 0);
assert_eq!(raw.shared[[i2, i1]], 0);
// cardinality_tally: both genomes carry exactly one member of the
// family (cardinality 1 each) -> one co-occurrence at [1][1].
let cardinality = merged.cardinality_tally(&raw, 1.0).expect("cardinality_tally");
assert_eq!(cardinality.counts[1][1], 1);
let total: u64 = cardinality.counts.iter().flatten().sum();
assert_eq!(total, 1, "no other cardinality pair should be tallied");
// base_pair_tally: the one eligible, differing locus is C (base 1) vs
// G (base 2).
let base_pairs = merged.base_pair_tally(&raw, 1.0).expect("base_pair_tally");
assert_eq!(base_pairs.counts[1][2], 1);
assert_eq!(base_pairs.counts[2][1], 1);
let total: u64 = base_pairs.counts.iter().flatten().sum();
assert_eq!(total, 2, "no other base pair should be tallied");
assert_eq!(base_pairs.same, [0, 0, 0, 0], "the two genomes never agree at this locus");
}