Extract index modules into specialized workspace subcrates
This commit partitions the obikindex crate into multiple focused subcrates (obikfilter, obikmerge, obikquery, obikrebuild, obikselect, obikstats, obikdump, and obikidxcache) to reduce coupling and clarify module boundaries. It standardizes error handling across the workspace using OKIError and OKIResult, updates index APIs to support lazy, disk-backed partition access, and migrates NUMA system utilities to a new obisys crate. All modifications are structural, focusing on dependency graph expansion, import path updates, and API surface reorganization without altering core runtime behavior.
This commit is contained in:
@@ -9,6 +9,7 @@ obikseq = { path = "../obikseq" }
|
||||
obiskio = { path = "../obiskio" }
|
||||
obisys = { path = "../obisys" }
|
||||
obicompactvec = { path = "../obicompactvec" }
|
||||
obikidxcache = { path = "../obikidxcache" }
|
||||
obiskbuilder = { path = "../obiskbuilder" }
|
||||
obipipeline = { path = "../obipipeline" }
|
||||
memmap2 = "0.9"
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
use ndarray::Array2;
|
||||
use obicompactvec::traits::{BitPartials, CountPartials};
|
||||
use obikidxcache::LayeredStore;
|
||||
use rayon::prelude::*;
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
|
||||
// ── Public API ────────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum DistanceMetric {
|
||||
/// Jaccard distance on presence/absence data.
|
||||
Jaccard,
|
||||
/// Hamming distance (number of differing kmer positions) on presence/absence data.
|
||||
Hamming,
|
||||
/// Mash distance on presence/absence data (Jaccard-derived mutation-rate estimate).
|
||||
Mash,
|
||||
/// Bray-Curtis dissimilarity on raw counts.
|
||||
BrayCurtis,
|
||||
/// Bray-Curtis dissimilarity normalised by per-genome total counts.
|
||||
RelfreqBrayCurtis,
|
||||
/// Euclidean distance on raw counts.
|
||||
Euclidean,
|
||||
/// Euclidean distance on relative frequencies.
|
||||
RelfreqEuclidean,
|
||||
/// Hellinger distance on counts.
|
||||
Hellinger,
|
||||
/// Euclidean distance in the Hellinger (√relative-frequency) space (unnormalised variant).
|
||||
HellingerEuclidean,
|
||||
}
|
||||
|
||||
pub struct DistanceOutput {
|
||||
/// n×n pairwise distance matrix (genomes in index order).
|
||||
pub matrix: Array2<f64>,
|
||||
/// n×n shared-kmer count matrix (intersection), if requested.
|
||||
pub shared_kmers: Option<Array2<u64>>,
|
||||
}
|
||||
|
||||
impl DistanceMetric {
|
||||
pub fn requires_counts(self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
DistanceMetric::BrayCurtis
|
||||
| DistanceMetric::RelfreqBrayCurtis
|
||||
| DistanceMetric::Euclidean
|
||||
| DistanceMetric::RelfreqEuclidean
|
||||
| DistanceMetric::Hellinger
|
||||
| DistanceMetric::HellingerEuclidean
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// ── KmerIndex::distance ───────────────────────────────────────────────────────
|
||||
|
||||
impl KmerIndex {
|
||||
pub fn distance(&self, metric: DistanceMetric, shared_kmers: bool, presence_threshold: u32) -> OKIResult<DistanceOutput> {
|
||||
let n_genomes = self.meta.genomes().map_err(OKIError::Io)?.len();
|
||||
if n_genomes < 2 {
|
||||
return Err(OKIError::InvalidInput(
|
||||
"distance requires at least 2 genomes in the index".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let use_counts = self.meta.config.with_counts;
|
||||
if metric.requires_counts() && !use_counts {
|
||||
return Err(OKIError::InvalidInput(format!(
|
||||
"{metric:?} requires a count index (with_counts = true)"
|
||||
)));
|
||||
}
|
||||
|
||||
let n_parts = self.n_partitions();
|
||||
|
||||
if use_counts {
|
||||
let stores: Vec<_> = (0..n_parts)
|
||||
.into_par_iter()
|
||||
.map(|i| self.count_store(i).map_err(OKIError::Partition))
|
||||
.collect::<OKIResult<_>>()?;
|
||||
let global = LayeredStore::new(stores);
|
||||
|
||||
let matrix = match metric {
|
||||
DistanceMetric::BrayCurtis => CountPartials::bray_dist_matrix(&global),
|
||||
DistanceMetric::RelfreqBrayCurtis => CountPartials::relfreq_bray_dist_matrix(&global),
|
||||
DistanceMetric::Euclidean => CountPartials::euclidean_dist_matrix(&global),
|
||||
DistanceMetric::RelfreqEuclidean => CountPartials::relfreq_euclidean_dist_matrix(&global),
|
||||
DistanceMetric::Hellinger => CountPartials::hellinger_dist_matrix(&global),
|
||||
DistanceMetric::HellingerEuclidean => CountPartials::hellinger_euclidean_dist_matrix(&global),
|
||||
DistanceMetric::Jaccard => CountPartials::threshold_jaccard_dist_matrix(&global, presence_threshold),
|
||||
DistanceMetric::Mash => CountPartials::threshold_mash_dist_matrix(&global, self.kmer_size(), presence_threshold),
|
||||
DistanceMetric::Hamming => {
|
||||
return Err(OKIError::InvalidInput(
|
||||
"Hamming is only available for presence/absence indexes".into(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let shared = if shared_kmers {
|
||||
let (inter, _) = CountPartials::partial_threshold_jaccard(&global, presence_threshold);
|
||||
Some(inter)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(DistanceOutput { matrix, shared_kmers: shared })
|
||||
} else {
|
||||
let stores: Vec<_> = (0..n_parts)
|
||||
.into_par_iter()
|
||||
.map(|i| self.presence_store(i).map_err(OKIError::Partition))
|
||||
.collect::<OKIResult<_>>()?;
|
||||
let global = LayeredStore::new(stores);
|
||||
|
||||
let matrix = match metric {
|
||||
DistanceMetric::Jaccard => BitPartials::jaccard_dist_matrix(&global),
|
||||
DistanceMetric::Mash => BitPartials::mash_dist_matrix(&global, self.kmer_size()),
|
||||
DistanceMetric::Hamming => {
|
||||
BitPartials::hamming_dist_matrix(&global).mapv(|v| v as f64)
|
||||
}
|
||||
other => {
|
||||
return Err(OKIError::InvalidInput(format!(
|
||||
"{other:?} requires a count index; use --metric jaccard or --metric hamming"
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
let shared = if shared_kmers {
|
||||
let (inter, _) = BitPartials::partial_jaccard(&global);
|
||||
Some(inter)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(DistanceOutput { matrix, shared_kmers: shared })
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,9 @@
|
||||
//! incrementally.
|
||||
|
||||
mod cardcomp;
|
||||
mod distance;
|
||||
mod matrix_store;
|
||||
pub mod siblings;
|
||||
|
||||
pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix};
|
||||
pub use distance::{DistanceMetric, DistanceOutput};
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
use obikindex::OKIResult;
|
||||
use obikindex::layer::open_data;
|
||||
use obikidxcache::LayeredStore;
|
||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
||||
|
||||
use obikindex::load_meta;
|
||||
use obikindex::KmerIndex;
|
||||
|
||||
impl KmerIndex {
|
||||
/// Open all count matrices for partition `part`, one per layer.
|
||||
/// Layers without a `counts/` directory are skipped.
|
||||
pub fn count_store(&self, part: usize) -> OKIResult<LayeredStore<PersistentCompactIntMatrix>> {
|
||||
let index_dir = self.index_dir(part);
|
||||
if !index_dir.exists() {
|
||||
return Ok(LayeredStore::new(vec![]));
|
||||
}
|
||||
let n_layers = self.n_layers();
|
||||
let matrices = (0..n_layers)
|
||||
.filter_map(|l| {
|
||||
self.layer_dir(part, l)
|
||||
.join("counts")
|
||||
.exists()
|
||||
.then(|| open_data(&index_dir, l))
|
||||
})
|
||||
.collect::<OKIResult<Vec<_>>>()?;
|
||||
Ok(LayeredStore::new(matrices))
|
||||
}
|
||||
|
||||
/// Open all presence matrices for partition `part`, one per layer.
|
||||
/// Layers without a `presence/` directory are skipped.
|
||||
pub fn presence_store(&self, part: usize) -> OKIResult<LayeredStore<PersistentBitMatrix>> {
|
||||
let index_dir = self.index_dir(part);
|
||||
if !index_dir.exists() {
|
||||
return Ok(LayeredStore::new(vec![]));
|
||||
}
|
||||
let n_layers = load_meta(&index_dir)?.n_layers;
|
||||
let matrices = (0..n_layers)
|
||||
.filter_map(|l| {
|
||||
self.layer_dir(part, l)
|
||||
.join("presence")
|
||||
.exists()
|
||||
.then(|| open_data(&index_dir, l))
|
||||
})
|
||||
.collect::<OKIResult<Vec<_>>>()?;
|
||||
Ok(LayeredStore::new(matrices))
|
||||
}
|
||||
}
|
||||
@@ -14,7 +14,7 @@ use obikindex::OKIResult;
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::helpers::{central_base, is_minorant};
|
||||
use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnexBuilder, olm_to_ok};
|
||||
use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnexBuilder};
|
||||
|
||||
// ── obipipeline data types ─────────────────────────────────────────────────
|
||||
|
||||
@@ -125,7 +125,7 @@ fn build_layer_sibling_annex(
|
||||
l: usize,
|
||||
cache: &Arc<PartitionCache>,
|
||||
) -> OKIResult<u64> {
|
||||
let mphf = MphfLayer::open(layer_dir).map_err(olm_to_ok)?;
|
||||
let mphf = MphfLayer::open(layer_dir)?;
|
||||
let k = index.kmer_size();
|
||||
let n = mphf.n();
|
||||
|
||||
|
||||
@@ -65,7 +65,7 @@ use obikindex::layer::KmerLayer;
|
||||
use super::cache::PartitionCache;
|
||||
use super::helpers::central_base;
|
||||
use super::iter::{SiblingEntry, SiblingLayerExt};
|
||||
use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex, olm_to_ok};
|
||||
use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex};
|
||||
|
||||
/// Families per batch — see the module docs for the memory-vs-per-partition-
|
||||
/// density trade-off this picks a point on. At ~90 genomes and a few
|
||||
@@ -194,10 +194,10 @@ pub(super) fn scan_layer_families(
|
||||
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 meta = PartitionMeta::load(index_dir)?;
|
||||
let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?);
|
||||
|
||||
let mat = KmerLayer::open(layer_dir, &meta.mode, with_counts).map_err(olm_to_ok)?;
|
||||
let mat = KmerLayer::open(layer_dir, &meta.mode, with_counts)?;
|
||||
let n_cols = mat.n_cols().min(n_genomes);
|
||||
|
||||
let ctx = Arc::new(LayerCtx {
|
||||
|
||||
@@ -71,15 +71,4 @@ pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
|
||||
pub use stats::{SiblingAnnexStats, SiblingStatsExt};
|
||||
pub use subsample::EntropyBias;
|
||||
|
||||
use obikindex::layer::OLMError;
|
||||
|
||||
use obikindex::OKIError;
|
||||
|
||||
pub(super) const ANNEX_FILE_NAME: &str = "siblings.psib";
|
||||
|
||||
pub(super) fn olm_to_ok(e: OLMError) -> OKIError {
|
||||
match e {
|
||||
OLMError::Io(e) => OKIError::Io(e),
|
||||
other => OKIError::InvalidInput(format!("layered-map error: {other}")),
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user