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:
Eric Coissac
@@ -0,0 +1,11 @@
|
||||
[package]
|
||||
name = "obikfilter"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
obikindex = { path = "../obikindex" }
|
||||
obicompactvec = { path = "../obicompactvec" }
|
||||
obikseq = { path = "../obikseq" }
|
||||
obiskio = { path = "../obiskio" }
|
||||
obitaxonomy = { path = "../obitaxonomy" }
|
||||
@@ -0,0 +1,192 @@
|
||||
use obikindex::layer::MphfLayer;
|
||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
||||
use obikseq::CanonicalKmer;
|
||||
use obiskio::UnitigFileReader;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use crate::filter::{KmerFilter, passes_all};
|
||||
use obikindex::KmerIndex;
|
||||
|
||||
impl KmerIndex {
|
||||
/// Iterate all indexed kmers in partition `part`, calling `cb(kmer, row)` for each
|
||||
/// kmer that passes every filter in `filters`.
|
||||
///
|
||||
/// `use_counts = true` → reads count columns (u32 values per genome).
|
||||
/// `use_counts = false` → reads presence columns, converted to 0/1 u32.
|
||||
///
|
||||
/// If no data matrix exists for a layer (pure set-membership, single genome),
|
||||
/// a row of `n_genomes` ones is emitted for every kmer in that layer — unless
|
||||
/// the filter rejects it, in which case the whole layer is skipped.
|
||||
/// Like [`iter_partition_kmers`] but the callback returns `false` to stop early.
|
||||
/// Returns `Ok(true)` if all kmers were visited, `Ok(false)` if the callback halted.
|
||||
pub fn iter_partition_kmers(
|
||||
&self,
|
||||
part: usize,
|
||||
use_counts: bool,
|
||||
n_genomes: usize,
|
||||
filters: &[Box<dyn KmerFilter>],
|
||||
mut cb: impl FnMut(CanonicalKmer, Box<[u32]>) -> bool,
|
||||
) -> OKIResult<bool> {
|
||||
let index_dir = self.index_dir(part);
|
||||
if !index_dir.exists() {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let mut l = 0;
|
||||
loop {
|
||||
let layer_dir = self.layer_dir(part, l)?;
|
||||
if !layer_dir.exists() {
|
||||
break;
|
||||
}
|
||||
l += 1;
|
||||
let mphf = MphfLayer::open(&layer_dir)?;
|
||||
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?;
|
||||
|
||||
let counts_dir = layer_dir.join("counts");
|
||||
let presence_dir = layer_dir.join("presence");
|
||||
|
||||
let cont = if use_counts && counts_dir.exists() {
|
||||
let mat = PersistentCompactIntMatrix::open(&layer_dir).map_err(OKIError::Io)?;
|
||||
let mut cont = true;
|
||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
let row = mat.row(slot);
|
||||
if passes_all(filters, kmer, &row, n_genomes) {
|
||||
cont = cb(kmer, row);
|
||||
if !cont {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
cont
|
||||
} else if !use_counts && presence_dir.exists() {
|
||||
let mat = PersistentBitMatrix::open(&layer_dir).map_err(OKIError::Io)?;
|
||||
let mut cont = true;
|
||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
|
||||
if passes_all(filters, kmer, &row, n_genomes) {
|
||||
cont = cb(kmer, row);
|
||||
if !cont {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
cont
|
||||
} else {
|
||||
// No data matrix: implicit presence — all values are 1. `row`
|
||||
// is identical for every kmer, but a filter can still depend
|
||||
// on the kmer's own sequence (e.g. MinComplexity), so this
|
||||
// cannot be evaluated once for the whole layer — filters must
|
||||
// still be tested per kmer.
|
||||
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
|
||||
let mut cont = true;
|
||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||
if mphf.find(kmer).is_some()
|
||||
&& passes_all(filters, kmer, &all_present, n_genomes)
|
||||
{
|
||||
cont = cb(kmer, all_present.clone());
|
||||
if !cont {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
cont
|
||||
};
|
||||
|
||||
if !cont {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// Like [`iter_partition_kmers`] but the callback also receives `(partition, layer)`
|
||||
/// indices, enabling debug output that identifies where each kmer was stored.
|
||||
/// Returns `Ok(true)` if all kmers were visited, `Ok(false)` if the callback halted.
|
||||
pub fn iter_partition_kmers_located(
|
||||
&self,
|
||||
part: usize,
|
||||
use_counts: bool,
|
||||
n_genomes: usize,
|
||||
filters: &[Box<dyn KmerFilter>],
|
||||
mut cb: impl FnMut(usize, usize, CanonicalKmer, Box<[u32]>) -> bool,
|
||||
) -> OKIResult<bool> {
|
||||
let index_dir = self.index_dir(part);
|
||||
if !index_dir.exists() {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let mut layer = 0;
|
||||
loop {
|
||||
let layer_dir = self.layer_dir(part, layer);
|
||||
if !layer_dir.exists() {
|
||||
break;
|
||||
}
|
||||
let mphf = MphfLayer::open(&layer_dir)?;
|
||||
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?;
|
||||
|
||||
let counts_dir = layer_dir.join("counts");
|
||||
let presence_dir = layer_dir.join("presence");
|
||||
|
||||
let cont = if use_counts && counts_dir.exists() {
|
||||
let mat = PersistentCompactIntMatrix::open(&layer_dir).map_err(OKIError::Io)?;
|
||||
let mut cont = true;
|
||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
let row = mat.row(slot);
|
||||
if passes_all(filters, kmer, &row, n_genomes) {
|
||||
cont = cb(part, layer, kmer, row);
|
||||
if !cont {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
cont
|
||||
} else if !use_counts && presence_dir.exists() {
|
||||
let mat = PersistentBitMatrix::open(&layer_dir).map_err(OKIError::Io)?;
|
||||
let mut cont = true;
|
||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
|
||||
if passes_all(filters, kmer, &row, n_genomes) {
|
||||
cont = cb(part, layer, kmer, row);
|
||||
if !cont {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
cont
|
||||
} else {
|
||||
// Same as iter_partition_kmers: row is constant but a filter
|
||||
// may still depend on the kmer's own sequence, so this must
|
||||
// be tested per kmer, not once for the whole layer.
|
||||
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
|
||||
let mut cont = true;
|
||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||
if mphf.find(kmer).is_some()
|
||||
&& passes_all(filters, kmer, &all_present, n_genomes)
|
||||
{
|
||||
cont = cb(part, layer, kmer, all_present.clone());
|
||||
if !cont {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
cont
|
||||
};
|
||||
|
||||
if !cont {
|
||||
return Ok(false);
|
||||
}
|
||||
layer += 1;
|
||||
}
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,298 @@
|
||||
use obicompactvec::FilterMask;
|
||||
use obikseq::CanonicalKmer;
|
||||
|
||||
/// Trait for kmer filters.
|
||||
///
|
||||
/// `kmer` is the k-mer's own canonical sequence, reconstructed from the
|
||||
/// source index's `unitigs.bin` (always present — see `rebuild_layer.rs`);
|
||||
/// `row` contains raw per-genome counts (or 0/1 for presence/absence data).
|
||||
/// `n_genomes` equals `row.len()`. Most filters only need `row` — `kmer` is
|
||||
/// there for filters that reason about the k-mer's sequence itself (e.g.
|
||||
/// [`MinComplexity`]).
|
||||
pub trait KmerFilter: Send + Sync {
|
||||
fn passes(&self, kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool;
|
||||
|
||||
/// Express this filter as a [`FilterMask`] column-operation expression.
|
||||
///
|
||||
/// Returns `Some(expr)` if the filter can be evaluated solely from matrix
|
||||
/// column aggregates (no per-kmer row scan needed). Returns `None` if the
|
||||
/// filter requires row-level inspection — always the case for a filter
|
||||
/// that needs the k-mer's sequence, since a `FilterMask` only expresses
|
||||
/// per-genome column aggregates, never per-slot sequence data.
|
||||
///
|
||||
/// `threshold` semantics in the returned mask use `>= threshold`, matching
|
||||
/// [`obicompactvec::MatrixGroupOps`]. Implementations must add 1 to any
|
||||
/// row-level threshold that uses strict `>` comparison.
|
||||
fn column_mask_expr(&self, _n_genomes: usize) -> Option<FilterMask> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// True when `row` passes every filter in `filters`.
|
||||
/// Returns `true` if `filters` is empty.
|
||||
pub fn passes_all(
|
||||
filters: &[Box<dyn KmerFilter>],
|
||||
kmer: CanonicalKmer,
|
||||
row: &[u32],
|
||||
n_genomes: usize,
|
||||
) -> bool {
|
||||
filters.iter().all(|f| f.passes(kmer, row, n_genomes))
|
||||
}
|
||||
|
||||
// ── Quorum filters ─────────────────────────────────────────────────────────────
|
||||
|
||||
fn present_count(row: &[u32], threshold: u32) -> usize {
|
||||
row.iter().filter(|&&v| v > threshold).count()
|
||||
}
|
||||
|
||||
/// At least `frac` fraction of genomes contain this kmer (count > `threshold`).
|
||||
pub struct MinGenomeFraction {
|
||||
pub frac: f64,
|
||||
pub threshold: u32,
|
||||
}
|
||||
|
||||
impl KmerFilter for MinGenomeFraction {
|
||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool {
|
||||
let p = present_count(row, self.threshold);
|
||||
p as f64 / n_genomes as f64 >= self.frac
|
||||
}
|
||||
|
||||
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
|
||||
let t = self.threshold.checked_add(1)?;
|
||||
let min_count = (self.frac * n_genomes as f64).ceil() as usize;
|
||||
Some(FilterMask::PresenceGeq {
|
||||
indices: (0..n_genomes).collect(),
|
||||
threshold: t,
|
||||
min_count,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// At most `frac` fraction of genomes contain this kmer (count > `threshold`).
|
||||
pub struct MaxGenomeFraction {
|
||||
pub frac: f64,
|
||||
pub threshold: u32,
|
||||
}
|
||||
|
||||
impl KmerFilter for MaxGenomeFraction {
|
||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool {
|
||||
let p = present_count(row, self.threshold);
|
||||
p as f64 / n_genomes as f64 <= self.frac
|
||||
}
|
||||
|
||||
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
|
||||
let t = self.threshold.checked_add(1)?;
|
||||
let max_count = (self.frac * n_genomes as f64).floor() as usize;
|
||||
Some(FilterMask::PresenceLeq {
|
||||
indices: (0..n_genomes).collect(),
|
||||
threshold: t,
|
||||
max_count,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// At least `count` genomes contain this kmer (count > `threshold`).
|
||||
pub struct MinGenomeCount {
|
||||
pub count: usize,
|
||||
pub threshold: u32,
|
||||
}
|
||||
|
||||
impl KmerFilter for MinGenomeCount {
|
||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
||||
present_count(row, self.threshold) >= self.count
|
||||
}
|
||||
|
||||
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
|
||||
let t = self.threshold.checked_add(1)?;
|
||||
Some(FilterMask::PresenceGeq {
|
||||
indices: (0..n_genomes).collect(),
|
||||
threshold: t,
|
||||
min_count: self.count,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// At most `count` genomes contain this kmer (count > `threshold`).
|
||||
pub struct MaxGenomeCount {
|
||||
pub count: usize,
|
||||
pub threshold: u32,
|
||||
}
|
||||
|
||||
impl KmerFilter for MaxGenomeCount {
|
||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
||||
present_count(row, self.threshold) <= self.count
|
||||
}
|
||||
|
||||
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
|
||||
let t = self.threshold.checked_add(1)?;
|
||||
Some(FilterMask::PresenceLeq {
|
||||
indices: (0..n_genomes).collect(),
|
||||
threshold: t,
|
||||
max_count: self.count,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ── Total-count filters (count indexes only) ───────────────────────────────────
|
||||
|
||||
/// Sum of counts across all genomes >= `total`.
|
||||
pub struct MinTotalCount {
|
||||
pub total: u32,
|
||||
}
|
||||
|
||||
impl KmerFilter for MinTotalCount {
|
||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
||||
row.iter().sum::<u32>() >= self.total
|
||||
}
|
||||
|
||||
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
|
||||
Some(FilterMask::SumGeq {
|
||||
indices: (0..n_genomes).collect(),
|
||||
min_sum: self.total,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Sum of counts across all genomes <= `total`.
|
||||
pub struct MaxTotalCount {
|
||||
pub total: u32,
|
||||
}
|
||||
|
||||
impl KmerFilter for MaxTotalCount {
|
||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
||||
row.iter().sum::<u32>() <= self.total
|
||||
}
|
||||
|
||||
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
|
||||
Some(FilterMask::SumLeq {
|
||||
indices: (0..n_genomes).collect(),
|
||||
max_sum: self.total,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ── Group-based quorum filter ─────────────────────────────────────────────────
|
||||
|
||||
/// Quorum filter operating on pre-classified genome groups.
|
||||
///
|
||||
/// `ingroup_idx` / `outgroup_idx` are column indices into the per-genome row.
|
||||
/// When `ingroup_idx` is empty, no ingroup quorum is checked.
|
||||
/// When `outgroup_idx` is empty, no outgroup quorum is checked.
|
||||
pub struct GroupQuorumFilter {
|
||||
pub ingroup_idx: Vec<usize>,
|
||||
pub outgroup_idx: Vec<usize>,
|
||||
pub threshold: u32,
|
||||
pub min_count: usize,
|
||||
pub max_count: usize,
|
||||
pub min_frac: f64,
|
||||
pub max_frac: f64,
|
||||
pub min_outgroup_count: usize,
|
||||
pub max_outgroup_count: usize,
|
||||
pub min_outgroup_frac: f64,
|
||||
pub max_outgroup_frac: f64,
|
||||
}
|
||||
|
||||
impl GroupQuorumFilter {
|
||||
// Build PresenceGeq/PresenceLeq constraints for one group (ingroup or outgroup).
|
||||
fn group_mask_parts(
|
||||
indices: &[usize],
|
||||
threshold: u32,
|
||||
min_count: usize,
|
||||
max_count: usize,
|
||||
min_frac: f64,
|
||||
max_frac: f64,
|
||||
parts: &mut Vec<FilterMask>,
|
||||
) {
|
||||
let n = indices.len();
|
||||
let geq = min_count.max((min_frac * n as f64).ceil() as usize);
|
||||
if geq > 0 {
|
||||
parts.push(FilterMask::PresenceGeq {
|
||||
indices: indices.to_vec(),
|
||||
threshold,
|
||||
min_count: geq,
|
||||
});
|
||||
}
|
||||
let leq = max_count.min((max_frac * n as f64).floor() as usize);
|
||||
if leq < n {
|
||||
parts.push(FilterMask::PresenceLeq {
|
||||
indices: indices.to_vec(),
|
||||
threshold,
|
||||
max_count: leq,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl KmerFilter for GroupQuorumFilter {
|
||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
||||
if !self.ingroup_idx.is_empty() {
|
||||
let n = self.ingroup_idx.iter()
|
||||
.filter(|&&i| row.get(i).copied().unwrap_or(0) > self.threshold)
|
||||
.count();
|
||||
let denom = self.ingroup_idx.len();
|
||||
if n < self.min_count { return false; }
|
||||
if n > self.max_count { return false; }
|
||||
let frac = n as f64 / denom as f64;
|
||||
if frac < self.min_frac { return false; }
|
||||
if frac > self.max_frac { return false; }
|
||||
}
|
||||
if !self.outgroup_idx.is_empty() {
|
||||
let n = self.outgroup_idx.iter()
|
||||
.filter(|&&i| row.get(i).copied().unwrap_or(0) > self.threshold)
|
||||
.count();
|
||||
let denom = self.outgroup_idx.len();
|
||||
if n < self.min_outgroup_count { return false; }
|
||||
if n > self.max_outgroup_count { return false; }
|
||||
let frac = n as f64 / denom as f64;
|
||||
if frac < self.min_outgroup_frac { return false; }
|
||||
if frac > self.max_outgroup_frac { return false; }
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn column_mask_expr(&self, _n_genomes: usize) -> Option<FilterMask> {
|
||||
let t = self.threshold.checked_add(1)?;
|
||||
let mut parts: Vec<FilterMask> = Vec::new();
|
||||
if !self.ingroup_idx.is_empty() {
|
||||
Self::group_mask_parts(
|
||||
&self.ingroup_idx, t,
|
||||
self.min_count, self.max_count,
|
||||
self.min_frac, self.max_frac,
|
||||
&mut parts,
|
||||
);
|
||||
}
|
||||
if !self.outgroup_idx.is_empty() {
|
||||
Self::group_mask_parts(
|
||||
&self.outgroup_idx, t,
|
||||
self.min_outgroup_count, self.max_outgroup_count,
|
||||
self.min_outgroup_frac, self.max_outgroup_frac,
|
||||
&mut parts,
|
||||
);
|
||||
}
|
||||
Some(FilterMask::And(parts))
|
||||
}
|
||||
}
|
||||
|
||||
// ── Complexity filter (post-hoc, sequence-based) ──────────────────────────────
|
||||
|
||||
/// Reject k-mers with normalized entropy below `theta` — the same complexity
|
||||
/// metric `obikmer index`'s `--theta`/`--level-max` apply *during* superkmer
|
||||
/// construction (see [`obikentropy::KmerEntropy`]), applied here after the
|
||||
/// fact, to k-mers already committed to a built index.
|
||||
///
|
||||
/// Unlike every other filter in this module, this one needs the k-mer's own
|
||||
/// sequence, not its per-genome row — `column_mask_expr` is never overridden
|
||||
/// (stays `None`), so this filter always forces the row-level scan path in
|
||||
/// `rebuild_layer.rs` (which reconstructs the sequence from `unitigs.bin`
|
||||
/// regardless, so no extra I/O beyond what filtering already requires).
|
||||
pub struct MinComplexity {
|
||||
pub level_max: usize,
|
||||
pub theta: f64,
|
||||
}
|
||||
|
||||
impl KmerFilter for MinComplexity {
|
||||
fn passes(&self, kmer: CanonicalKmer, _row: &[u32], _n_genomes: usize) -> bool {
|
||||
use obikentropy::KmerEntropy;
|
||||
kmer.entropy(self.level_max) >= self.theta
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
//! Per-kmer filtering: `KmerFilter` judges one k-mer (its canonical
|
||||
//! sequence + its per-genome row) at a time, deciding whether it passes.
|
||||
//! Orthogonal to genome-column selection/aggregation (`obikselect`), which
|
||||
//! operates on already-retained k-mers.
|
||||
//!
|
||||
//! [`filter`] (the `KmerFilter` trait + its implementations) depends only
|
||||
//! on `obicompactvec`/`obikseq`, not on `obikindex`. [`dump_layer`] is the
|
||||
//! extension over `obikindex::KmerIndex` that actually iterates a
|
||||
//! partition's k-mers through those filters (`iter_partition_kmers`,
|
||||
//! `iter_partition_kmers_located`) — every k-mer, filtered or not, goes
|
||||
//! through this same path (`passes_all` on an empty filter list is always
|
||||
//! `true`), so this crate depends one-way on `obikindex`, not the reverse.
|
||||
|
||||
mod filter;
|
||||
mod dump_layer;
|
||||
mod predicate;
|
||||
|
||||
pub use filter::{
|
||||
GroupQuorumFilter, KmerFilter, MaxGenomeCount, MaxGenomeFraction, MaxTotalCount,
|
||||
MinComplexity, MinGenomeCount, MinGenomeFraction, MinTotalCount, passes_all,
|
||||
};
|
||||
pub use predicate::{GroupFilterParams, MetaPred};
|
||||
@@ -0,0 +1,254 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::filter::GroupQuorumFilter;
|
||||
use obitaxonomy::{TaxPath, TaxPattern};
|
||||
|
||||
use obikindex::{GenomeInfo, IndexMeta};
|
||||
|
||||
// ── Operator ──────────────────────────────────────────────────────────────────
|
||||
|
||||
enum PredOp { Wildcard, Eq, Ne, Matches, NotMatches }
|
||||
|
||||
// ── MetaPred ──────────────────────────────────────────────────────────────────
|
||||
|
||||
/// A single predicate on genome metadata: `key OP val1|val2|…`
|
||||
///
|
||||
/// Operators: `=` (exact), `!=` (not equal), `~` (path ancestor), `!~` (not ancestor).
|
||||
/// Multiple values separated by `|` are OR'd.
|
||||
pub struct MetaPred {
|
||||
key: String,
|
||||
op: PredOp,
|
||||
values: Vec<String>,
|
||||
}
|
||||
|
||||
impl MetaPred {
|
||||
/// Parse a predicate string of the form `key=v1|v2`, `key!=v`, `key~path`, `key!~path`.
|
||||
/// The special values `*` and `all` (case-insensitive) match every genome.
|
||||
pub fn parse(s: &str) -> Result<Self, String> {
|
||||
let t = s.trim();
|
||||
if t == "*" || t.eq_ignore_ascii_case("all") {
|
||||
return Ok(Self { key: String::new(), op: PredOp::Wildcard, values: vec![] });
|
||||
}
|
||||
|
||||
let (op, key, rhs) =
|
||||
if let Some(pos) = s.find("!=") {
|
||||
(PredOp::Ne, &s[..pos], &s[pos+2..])
|
||||
} else if let Some(pos) = s.find("!~") {
|
||||
(PredOp::NotMatches, &s[..pos], &s[pos+2..])
|
||||
} else if let Some(pos) = s.find('=') {
|
||||
(PredOp::Eq, &s[..pos], &s[pos+1..])
|
||||
} else if let Some(pos) = s.find('~') {
|
||||
(PredOp::Matches, &s[..pos], &s[pos+1..])
|
||||
} else {
|
||||
return Err(format!("no operator found in predicate: {s}"));
|
||||
};
|
||||
|
||||
let key = key.trim().to_string();
|
||||
if key.is_empty() { return Err(format!("empty key in predicate: {s}")); }
|
||||
|
||||
let values: Vec<String> = rhs.split('|').map(|v| v.trim().to_string()).collect();
|
||||
if values.iter().any(|v| v.is_empty()) {
|
||||
return Err(format!("empty value in predicate: {s}"));
|
||||
}
|
||||
Ok(Self { key, op, values })
|
||||
}
|
||||
|
||||
/// Evaluate against one genome's metadata.
|
||||
/// Returns `None` when the key is absent (NA propagation).
|
||||
pub(crate) fn eval(&self, meta: &HashMap<String, String>) -> Option<bool> {
|
||||
if matches!(self.op, PredOp::Wildcard) { return Some(true); }
|
||||
let value = meta.get(&self.key)?;
|
||||
Some(match self.op {
|
||||
PredOp::Wildcard => unreachable!(),
|
||||
PredOp::Eq => self.values.iter().any(|v| v == value),
|
||||
PredOp::Ne => self.values.iter().all(|v| v != value),
|
||||
PredOp::Matches => self.values.iter().any(|v| path_matches(value, v)),
|
||||
PredOp::NotMatches => self.values.iter().all(|v| !path_matches(value, v)),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl GenomeInfo {
|
||||
/// Evaluate a single metadata predicate against this genome.
|
||||
/// Returns `None` when the predicate's key is absent (NA propagation).
|
||||
pub fn matches(&self, pred: &MetaPred) -> Option<bool> {
|
||||
pred.eval(&self.meta)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Path matching ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// True if the stored taxonomy `value` matches `pattern`.
|
||||
///
|
||||
/// `value` must be a valid `TaxPath` (starts with `taxonomy:/`).
|
||||
/// `pattern` is a `TaxPattern` query (see `obitaxonomy::TaxPattern` for syntax).
|
||||
/// Returns `false` if either fails to parse.
|
||||
fn path_matches(value: &str, pattern: &str) -> bool {
|
||||
let Ok(path) = TaxPath::parse(value) else { return false };
|
||||
let Ok(pat) = TaxPattern::parse(pattern) else { return false };
|
||||
pat.matches(&path)
|
||||
}
|
||||
|
||||
// ── Three-value group evaluation ──────────────────────────────────────────────
|
||||
|
||||
/// AND of all predicates (ingroup semantics).
|
||||
/// Short-circuits on `Some(false)`; propagates `None` if no predicate returns `false`.
|
||||
fn eval_and(preds: &[MetaPred], meta: &HashMap<String, String>) -> Option<bool> {
|
||||
let mut has_na = false;
|
||||
for pred in preds {
|
||||
match pred.eval(meta) {
|
||||
Some(false) => return Some(false),
|
||||
Some(true) => {}
|
||||
None => has_na = true,
|
||||
}
|
||||
}
|
||||
if has_na { None } else { Some(true) }
|
||||
}
|
||||
|
||||
/// OR of all predicates (outgroup semantics).
|
||||
/// Short-circuits on `Some(true)`; propagates `None` if no predicate returns `true`.
|
||||
fn eval_or(preds: &[MetaPred], meta: &HashMap<String, String>) -> Option<bool> {
|
||||
let mut has_na = false;
|
||||
for pred in preds {
|
||||
match pred.eval(meta) {
|
||||
Some(true) => return Some(true),
|
||||
Some(false) => {}
|
||||
None => has_na = true,
|
||||
}
|
||||
}
|
||||
if has_na { None } else { Some(false) }
|
||||
}
|
||||
|
||||
// ── Genome classification ─────────────────────────────────────────────────────
|
||||
|
||||
enum Membership { Ingroup, Outgroup, Uncategorized }
|
||||
|
||||
fn classify(
|
||||
genomes: &[GenomeInfo],
|
||||
ingroup: &[MetaPred],
|
||||
outgroup: &[MetaPred],
|
||||
) -> Vec<Membership> {
|
||||
genomes.iter().map(|g| {
|
||||
let in_r = if ingroup.is_empty() { None } else { eval_and(ingroup, &g.meta) };
|
||||
let out_r = if outgroup.is_empty() { None } else { eval_or(outgroup, &g.meta) };
|
||||
|
||||
// Ingroup wins over outgroup.
|
||||
if in_r == Some(true) { return Membership::Ingroup; }
|
||||
if out_r == Some(true) { return Membership::Outgroup; }
|
||||
Membership::Uncategorized
|
||||
}).collect()
|
||||
}
|
||||
|
||||
// ── Group quorum filter construction ──────────────────────────────────────────
|
||||
|
||||
pub struct GroupFilterParams {
|
||||
pub threshold: u32,
|
||||
pub min_count: Option<isize>,
|
||||
pub max_count: Option<isize>,
|
||||
pub min_frac: Option<f64>,
|
||||
pub max_frac: Option<f64>,
|
||||
pub min_outgroup_count: Option<isize>,
|
||||
pub max_outgroup_count: Option<isize>,
|
||||
pub min_outgroup_frac: Option<f64>,
|
||||
pub max_outgroup_frac: Option<f64>,
|
||||
}
|
||||
|
||||
impl IndexMeta {
|
||||
/// Returns indices of genomes matching `pred_str` (single predicate).
|
||||
pub fn matching_genome_indices(&self, pred_str: &str) -> Result<Vec<usize>, String> {
|
||||
let pred = MetaPred::parse(pred_str)?;
|
||||
let genomes = self.genomes().map_err(|e| e.to_string())?;
|
||||
Ok(genomes.iter().enumerate()
|
||||
.filter_map(|(i, g)| {
|
||||
if g.matches(&pred) == Some(true) { Some(i) } else { std::option::Option::None }
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Build a `GroupQuorumFilter` from parsed predicates, evaluated against `self.genomes`.
|
||||
///
|
||||
/// - No groups defined: `ingroup_idx` = all genomes (implicit ingroup).
|
||||
/// - `ingroup` predicates only: outgroup indices are empty.
|
||||
/// - `outgroup` predicates only: ingroup indices are empty.
|
||||
/// - Both defined: ingroup wins on overlap; uncategorized genomes are ignored.
|
||||
pub fn build_group_filter(
|
||||
&self,
|
||||
ingroup_preds: &[MetaPred],
|
||||
outgroup_preds: &[MetaPred],
|
||||
p: GroupFilterParams,
|
||||
) -> Result<GroupQuorumFilter, String> {
|
||||
let genomes = self.genomes().map_err(|e| e.to_string())?;
|
||||
let (ingroup_idx, outgroup_idx) = if ingroup_preds.is_empty() && outgroup_preds.is_empty() {
|
||||
((0..genomes.len()).collect(), vec![])
|
||||
} else {
|
||||
let members = classify(&genomes, ingroup_preds, outgroup_preds);
|
||||
let in_idx: Vec<usize> = members.iter().enumerate()
|
||||
.filter(|(_, m)| matches!(m, Membership::Ingroup))
|
||||
.map(|(i, _)| i).collect();
|
||||
let out_idx: Vec<usize> = members.iter().enumerate()
|
||||
.filter(|(_, m)| matches!(m, Membership::Outgroup))
|
||||
.map(|(i, _)| i).collect();
|
||||
(in_idx, out_idx)
|
||||
};
|
||||
|
||||
let in_size = ingroup_idx.len();
|
||||
let out_size = outgroup_idx.len();
|
||||
|
||||
let ingroup_quorum_explicit = p.min_count.is_some() || p.max_count.is_some()
|
||||
|| p.min_frac.is_some() || p.max_frac.is_some();
|
||||
let outgroup_quorum_explicit = p.min_outgroup_count.is_some() || p.max_outgroup_count.is_some()
|
||||
|| p.min_outgroup_frac.is_some() || p.max_outgroup_frac.is_some();
|
||||
|
||||
let default_min_frac = if !ingroup_preds.is_empty() && !ingroup_quorum_explicit { 1.0 } else { 0.0 };
|
||||
let default_max_outgroup_count = if !outgroup_preds.is_empty() && !outgroup_quorum_explicit { 0 } else { out_size };
|
||||
|
||||
// Resolve a signed count: negative means an offset from the group size
|
||||
// (e.g. -1 = all but one), floored at 1 so the negative form always keeps
|
||||
// constraining the group — even a singleton group, where n-1 would be 0
|
||||
// and would otherwise drop the constraint entirely.
|
||||
let resolve = |v: isize, size: usize| -> usize {
|
||||
if v < 0 { (size as isize + v).max(1) as usize } else { v as usize }
|
||||
};
|
||||
|
||||
let min_count = p.min_count.map(|v| resolve(v, in_size)).unwrap_or(0);
|
||||
let max_count = p.max_count.map(|v| resolve(v, in_size)).unwrap_or(in_size);
|
||||
let min_frac = p.min_frac.unwrap_or(default_min_frac);
|
||||
let max_frac = p.max_frac.unwrap_or(1.0);
|
||||
let min_outgroup_count = p.min_outgroup_count.map(|v| resolve(v, out_size)).unwrap_or(0);
|
||||
let max_outgroup_count = p.max_outgroup_count.map(|v| resolve(v, out_size)).unwrap_or(default_max_outgroup_count);
|
||||
let min_outgroup_frac = p.min_outgroup_frac.unwrap_or(0.0);
|
||||
let max_outgroup_frac = p.max_outgroup_frac.unwrap_or(1.0);
|
||||
|
||||
for (v, lo, hi) in [
|
||||
("--min-frac/--max-frac", min_frac, max_frac),
|
||||
("--min-outgroup-frac/--max-outgroup-frac", min_outgroup_frac, max_outgroup_frac),
|
||||
] {
|
||||
if !(0.0..=1.0).contains(&lo) || !(0.0..=1.0).contains(&hi) {
|
||||
return Err(format!("{v}: fraction values must be in [0.0, 1.0]"));
|
||||
}
|
||||
if lo > hi {
|
||||
return Err(format!("{v}: min ({lo}) is greater than max ({hi})"));
|
||||
}
|
||||
}
|
||||
if min_count > max_count {
|
||||
return Err(format!("--min-count/--max-count: min ({min_count}) is greater than max ({max_count})"));
|
||||
}
|
||||
if min_outgroup_count > max_outgroup_count {
|
||||
return Err(format!("--min-outgroup-count/--max-outgroup-count: min ({min_outgroup_count}) is greater than max ({max_outgroup_count})"));
|
||||
}
|
||||
|
||||
Ok(GroupQuorumFilter {
|
||||
ingroup_idx,
|
||||
outgroup_idx,
|
||||
threshold: p.threshold,
|
||||
min_count,
|
||||
max_count,
|
||||
min_frac,
|
||||
max_frac,
|
||||
min_outgroup_count,
|
||||
max_outgroup_count,
|
||||
min_outgroup_frac,
|
||||
max_outgroup_frac,
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user