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

refactor: merge KmerPartitions into KmerIndex and rename obikpartition

Browse Source 
Consolidates partition logic, metadata storage, and layer management directly into KmerIndex. Renames obikpartitionner to obikpartition, retaining only PartitionRouter for superkmer routing. Removes intermediate .partition() accessors in favor of direct methods on the index and updates PartitionCache::build to accept &KmerIndex directly. Derives n_partitions from config.n_bits and consolidates k-mer/minimizer sizes into IndexMeta.config. Fixes a regression where PartitionRouter::open incorrectly defaulted to closed.
This commit is contained in:
Eric Coissac
2026-08-20 20:23:58 +02:00
parent b5ec0122d0
commit 6c860f120f
49 changed files with 357 additions and 427 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obikphylo/src/siblings/alignment.rs
+2 -12
View File
@@ -1,10 +1,9 @@
use std::sync::Arc;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::cache::PartitionCache;
use super::family_scan::{Selection, scan_layer_families};
@@ -87,16 +86,7 @@ impl SnpAlignmentExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
let selections =
super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?;
src/obikphylo/src/siblings/build.rs
+5 -15
View File
@@ -4,7 +4,6 @@ use std::sync::atomic::Ordering;
use rayon::prelude::*;
use obikpartitionner::KmerPartitions;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::IndexMode;
@@ -12,7 +11,7 @@ use obipipeline::ThrottleGuard;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::cache::PartitionCache;
use super::helpers::{central_base, is_minorant};
@@ -76,19 +75,10 @@ pub trait SiblingAnnexBuildExt {
impl SiblingAnnexBuildExt for KmerIndex {
fn build_sibling_annex(&self) -> OKIResult<()> {
let n_parts = self.n_partitions();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
let cache = Arc::new(PartitionCache::build(
&partition,
self,
n_parts,
self.meta().config.with_counts,
)?);
@@ -96,18 +86,18 @@ impl SiblingAnnexBuildExt for KmerIndex {
let pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
let mut total_slots: u64 = 0;
for part in 0..n_parts {
let index_dir = self.partition().index_dir(part);
let index_dir = self.index_dir(part);
if !index_dir.exists() {
pb.inc(1);
continue;
}
let meta = self.partition().partition_meta(part)?;
let meta = self.partition_meta(part)?;
let mut part_slots: u64 = 0;
for l in 0..meta.n_layers {
part_slots += build_layer_sibling_annex(
self,
&self.partition().layer_dir(part, l),
&self.layer_dir(part, l),
&meta.mode,
n_parts,
l,
src/obikphylo/src/siblings/cache.rs
+5 -6
View File
@@ -3,13 +3,12 @@ use rayon::prelude::*;
use std::path::Path;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikpartitionner::KmerPartitions;
use obikseq::CanonicalKmer;
use obilayeredmap::meta::IndexMode;
use obilayeredmap::{Layer, OLMResult};
use obisys::progress_bar;
use obikindex::OKIResult;
use obikindex::{KmerIndex, OKIResult};
use super::SiblingAnnex;
use super::iter::SiblingLayerExt;
@@ -159,7 +158,7 @@ pub(super) struct PartitionCache {
impl PartitionCache {
pub(super) fn build(
partition: &KmerPartitions,
index: &KmerIndex,
n_parts: usize,
with_counts: bool,
) -> OKIResult<Self> {
@@ -167,15 +166,15 @@ impl PartitionCache {
let built: Vec<(Vec<Mat>, usize)> = (0..n_parts)
.into_par_iter()
.map(|part| -> OKIResult<(Vec<Mat>, usize)> {
let index_dir = partition.index_dir(part);
let index_dir = index.index_dir(part);
if !index_dir.exists() {
pb.inc(1);
return Ok((Vec::new(), 0));
}
let meta = partition.partition_meta(part)?;
let meta = index.partition_meta(part)?;
let mut mats = Vec::with_capacity(meta.n_layers);
for l in 0..meta.n_layers {
let Ok(mat) = Mat::open(&partition.layer_dir(part, l), &meta.mode, with_counts)
let Ok(mat) = Mat::open(&index.layer_dir(part, l), &meta.mode, with_counts)
else {
continue;
};
src/obikphylo/src/siblings/cardinality.rs
+2 -12
View File
@@ -2,11 +2,10 @@ use std::sync::Arc;
use ndarray::Array2;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::cache::PartitionCache;
use super::distance::RawSnpDistanceOutput;
@@ -69,8 +68,6 @@ impl CardinalityExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
if i == j {
return false;
@@ -80,14 +77,7 @@ impl CardinalityExt for KmerIndex {
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
});
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers");
src/obikphylo/src/siblings/distance.rs
+2 -12
View File
@@ -2,11 +2,10 @@ use std::sync::Arc;
use ndarray::Array2;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::cache::PartitionCache;
use super::family_scan::{Selection, scan_layer_families};
@@ -70,16 +69,7 @@ where
let n_genomes = index.meta().genomes.len();
let with_counts = index.meta().config.with_counts;
let k = index.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
index.root_path(),
index.kmer_size(),
index.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(index, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(index)?;
let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
src/obikphylo/src/siblings/entropy.rs
+2 -21
View File
@@ -13,7 +13,6 @@ use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
@@ -139,16 +138,7 @@ impl ShannonEntropyExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
let selections =
super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?;
@@ -240,16 +230,7 @@ pub(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf])
let n_genomes = index.meta().genomes.len();
let with_counts = index.meta().config.with_counts;
let k = index.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
index.root_path(),
index.kmer_size(),
index.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(index, n_parts, with_counts)?);
let minorant_counts = super::subsample::minorant_counts(layer_dirs)?;
let pb = progress_bar("entropy_annex_build", missing.len() as u64, "layers");
src/obikphylo/src/siblings/family_scan.rs
+3 -3
View File
@@ -105,13 +105,13 @@ pub(crate) fn sibling_layer_dirs(index: &KmerIndex) -> OKIResult<Vec<PathBuf>> {
let n_parts = index.n_partitions();
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = index.partition().index_dir(part);
let index_dir = index.index_dir(part);
if !index_dir.exists() {
continue;
}
let n_layers = index.partition().n_layers(part)?;
let n_layers = index.n_layers(part)?;
for l in 0..n_layers {
let this_layer_dir = index.partition().layer_dir(part, l);
let this_layer_dir = index.layer_dir(part, l);
let annex_path = this_layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
src/obikphylo/src/siblings/sankoff_bundle.rs
+2 -12
View File
@@ -29,11 +29,10 @@ use std::sync::Arc;
use ndarray::Array2;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::alignment::{SnpAlignment, iupac_code};
use super::cache::PartitionCache;
@@ -83,16 +82,7 @@ impl SankoffBundleExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = sibling_layer_dirs(self)?;
// Computed once — every pass below iterates the exact same
// families, in the exact same layers, per the shared-selection
src/obikphylo/src/siblings/stats.rs
+2 -12
View File
@@ -2,11 +2,10 @@ use std::sync::Arc;
use rayon::prelude::*;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::ANNEX_FILE_NAME;
use super::SiblingAnnex;
@@ -109,19 +108,10 @@ impl SiblingStatsExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
// Same whole-run cache as `build_sibling_annex` — see its docs for
// why re-opening per lookup (or per call to a batching helper) is
// not good enough on a real index.
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
// One layer at a time, not parallelised across layers — see
src/obikphylo/src/siblings/tests.rs
+9 -8
View File
@@ -69,11 +69,12 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
let mut rep = Reporter::new();
let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
let mut router = idx.partition_router().expect("partition_router");
for page in stream {
let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
idx.partition_mut().write_batch(batch).expect("write_batch");
router.write_batch(batch).expect("write_batch");
}
idx.partition_mut().close().expect("close partition writers");
router.close().expect("close partition writers");
idx.mark_scattered().expect("mark_scattered");
idx.dereplicate_and_count(false, &mut rep).expect("dereplicate_and_count");
idx.build_layers(1, None, false, &mut rep).expect("build_layers");
@@ -87,9 +88,9 @@ fn canonical(ascii: &[u8]) -> CanonicalKmer {
/// Read back the annex entry for a given canonical k-mer from the merged
/// index's (single) partition/layer, asserting it was found at all.
fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> FamilyMask {
let meta = idx.partition().partition_meta(0).unwrap();
let meta = idx.partition_meta(0).unwrap();
for l in 0..meta.n_layers {
let layer_dir = idx.partition().layer_dir(0, l);
let layer_dir = idx.layer_dir(0, l);
let mphf = MphfLayer::open(&layer_dir, &meta.mode).unwrap();
if let Some(slot) = mphf.find(kmer) {
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).unwrap();
@@ -165,7 +166,7 @@ fn sibling_annex_works_after_pack_sparse() {
let merged = merge_two(dir.path(), &g1, &g2);
merged.pack_matrices(true).expect("pack_matrices(sparse)");
let index_dir = merged.partition().index_dir(0);
let index_dir = merged.index_dir(0);
assert!(
index_dir.join("layer_0").join("presence").join("is_multi.prsb").exists(),
"pack_matrices(true) must leave the sparse marker file behind"
@@ -308,9 +309,9 @@ fn sibling_annex_no_empty_masks_after_build() {
let g1 = build_single_genome_index(dir.path(), "g1", &seq);
g1.build_sibling_annex().expect("build_sibling_annex");
let n_layers = g1.partition().n_layers(0).expect("partition meta");
let n_layers = g1.n_layers(0).expect("partition meta");
for l in 0..n_layers {
let layer_dir = g1.partition().layer_dir(0, l);
let layer_dir = g1.layer_dir(0, l);
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).expect("annex open");
for slot in 0..annex.len() {
let mask = annex.get(slot).expect("slot must have an entry");
@@ -358,7 +359,7 @@ fn sibling_annex_records_the_real_layer_of_each_family_member() {
let merged = merge_two(dir.path(), &g1, &g2);
merged.build_sibling_annex().expect("build_sibling_annex");
let n_layers = merged.partition().n_layers(0).unwrap();
let n_layers = merged.n_layers(0).unwrap();
assert_eq!(n_layers, 2, "fixture assumption: one merge, one new layer");
let g1_kmer = canonical(b"AACCGCTTAAG"); // own base C (1), sibling base G (2) — lives in layer 1