feat: implement per-k-mer sibling counts and central neighbor generation
Introduce the siblingannex module in obicompactvec to store per-slot minorant flags and sibling counts in a memory-mapped annex file. Add a scatter-gather pipeline in obikindex to compute these values across index layers and write them to .psib files. Implement central_canonical_neighbors in obikseq for generating strand-aware k-mer variants around the middle base. Expose rolling statistics in obiskbuilder and update dependency graphs accordingly.
This commit is contained in:
Eric Coissac
@@ -10,6 +10,7 @@ obiskio = { path = "../obiskio" }
|
||||
obisys = { path = "../obisys" }
|
||||
obicompactvec = { path = "../obicompactvec" }
|
||||
obilayeredmap = { path = "../obilayeredmap" }
|
||||
obiskbuilder = { path = "../obiskbuilder" }
|
||||
ndarray = "0.16"
|
||||
rayon = "1"
|
||||
crossbeam-channel = "0.5"
|
||||
@@ -19,6 +20,10 @@ indicatif = "0.17"
|
||||
tracing = "0.1.44"
|
||||
hwlocality = { version = "1.0.0-alpha.11", features = ["vendored"], optional = true }
|
||||
|
||||
[dev-dependencies]
|
||||
obiread = { path = "../obiread" }
|
||||
tempfile = "3"
|
||||
|
||||
[features]
|
||||
default = ["numa"]
|
||||
numa = ["hwlocality"]
|
||||
|
||||
@@ -9,6 +9,7 @@ mod numa;
|
||||
mod rebuild;
|
||||
mod reindex;
|
||||
mod select;
|
||||
mod siblings;
|
||||
mod stats;
|
||||
|
||||
pub use error::{OKIError, OKIResult};
|
||||
|
||||
@@ -0,0 +1,335 @@
|
||||
//! Sibling-count / minorant annex construction.
|
||||
//!
|
||||
//! See `docmd/theory/evolutionary_distances.md`, "Step 2b — sibling-count /
|
||||
//! minorant annex", for the full design discussion this implements.
|
||||
//!
|
||||
//! For each distinct k-mer of each layer of the (already built/merged)
|
||||
//! index, computes two facts about its "family" (the up to 4 k-mers sharing
|
||||
//! its flanks, differing only at the central base — well-defined for odd
|
||||
//! k), both properties of the whole current multi-genome index rather than
|
||||
//! of any one genome:
|
||||
//! - how many *other* family members (0-3) are observed anywhere in the
|
||||
//! index;
|
||||
//! - whether this k-mer is the "minorant" of its family — the smallest
|
||||
//! canonical encoding among the members actually observed.
|
||||
//!
|
||||
//! Implementation note (deviation from the fully staged `obipipeline` design
|
||||
//! discussed at length in the doc): this first implementation processes each
|
||||
//! layer with a straightforward sequential scatter (batch the layer's
|
||||
//! outgoing variant queries by destination partition) then gather (one
|
||||
//! `query_partition_with` call per destination partition) — not the
|
||||
//! multi-stage elementary `obipipeline` pipeline the design settled on. The
|
||||
//! external semantics (one annex file per layer, sequential outer loop over
|
||||
//! layers, order-independent reconciliation) match the design exactly; only
|
||||
//! the internal execution mechanism is simplified, as a scope trade-off.
|
||||
//! Revisiting this to use `obipipeline` with elementary stages, as designed,
|
||||
//! is a follow-up, not a behavioural change.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::{SiblingAnnexBuilder, SiblingInfo};
|
||||
use obikpartitionner::{KmerDesc, QueryHit};
|
||||
use obikseq::{CanonicalKmer, Minimizer};
|
||||
use obilayeredmap::{MphfLayer, OLMError};
|
||||
use obilayeredmap::meta::PartitionMeta;
|
||||
use obiskbuilder::rolling_stat::RollingStat;
|
||||
use obiskio::UnitigFileReader;
|
||||
|
||||
use crate::error::{OKIError, OKIResult};
|
||||
use crate::index::KmerIndex;
|
||||
|
||||
const INDEX_SUBDIR: &str = "index";
|
||||
const ANNEX_FILE_NAME: &str = "siblings.psib";
|
||||
|
||||
fn olm_to_ok(e: OLMError) -> OKIError {
|
||||
match e {
|
||||
OLMError::Io(e) => OKIError::Io(e),
|
||||
other => OKIError::InvalidInput(format!("layered-map error: {other}")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimiser of a single, isolated canonical k-mer (not part of a streamed
|
||||
/// sequence). `RollingStat` computes minimisers incrementally along a
|
||||
/// sequence; this feeds one k-mer's bases through a fresh instance to get
|
||||
/// the same selection for a single, disconnected k-mer. Not the leanest
|
||||
/// possible primitive (an O(1)-amortised dedicated scan, as originally
|
||||
/// sketched in the design doc's Step 0, would avoid the ASCII round-trip and
|
||||
/// `RollingStat` allocation) but correct and reuses already-tested logic;
|
||||
/// left as a follow-up optimisation.
|
||||
fn lone_kmer_minimizer(kmer: CanonicalKmer) -> Minimizer {
|
||||
let ascii = kmer.to_ascii();
|
||||
let mut rs = RollingStat::new(0);
|
||||
for b in ascii {
|
||||
rs.push(b);
|
||||
}
|
||||
rs.canonical_minimizer()
|
||||
.expect("RollingStat must be ready after k bases of a valid k-mer")
|
||||
}
|
||||
|
||||
/// Destination partition for a (possibly synthetic) canonical k-mer, using
|
||||
/// the same routing rule as the rest of the index (`minimiser.seq_hash() &
|
||||
/// mask`, `n_partitions` is a power of two).
|
||||
fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize {
|
||||
let mask = (n_partitions as u64) - 1;
|
||||
(lone_kmer_minimizer(kmer).seq_hash() & mask) as usize
|
||||
}
|
||||
|
||||
/// Running reconciliation state for one source k-mer, initialised to the
|
||||
/// trivial "no siblings observed yet" state and folded incrementally (in any
|
||||
/// order — commutative) as query answers come back.
|
||||
#[derive(Clone, Copy)]
|
||||
struct RunningState {
|
||||
minorant: bool,
|
||||
siblings: u8,
|
||||
}
|
||||
|
||||
impl Default for RunningState {
|
||||
fn default() -> Self {
|
||||
RunningState { minorant: true, siblings: 0 }
|
||||
}
|
||||
}
|
||||
|
||||
impl KmerIndex {
|
||||
/// Build the sibling-count/minorant annex for every layer of every
|
||||
/// partition of this (already built) index, writing one annex file per
|
||||
/// layer alongside its existing index files. Safe to call again later
|
||||
/// (e.g. after a fresh `merge`) — each run simply overwrites the annex
|
||||
/// files of the index it is called on.
|
||||
///
|
||||
/// Cross-partition/cross-layer lookups are required (a k-mer's siblings
|
||||
/// can live in any partition), but the layer loop itself — and thus the
|
||||
/// annex file this produces — stays local to one layer at a time.
|
||||
pub fn build_sibling_annex(&self) -> OKIResult<()> {
|
||||
let n_parts = self.n_partitions();
|
||||
|
||||
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}"));
|
||||
self.build_layer_sibling_annex(&layer_dir, &meta, n_parts)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn build_layer_sibling_annex(
|
||||
&self,
|
||||
layer_dir: &Path,
|
||||
meta: &PartitionMeta,
|
||||
n_parts: usize,
|
||||
) -> OKIResult<()> {
|
||||
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
|
||||
let n_slots = mphf.n();
|
||||
|
||||
// ── Enumerate this layer's distinct k-mers, one per slot ────────────
|
||||
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; n_slots];
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Scatter: bucket outgoing variant queries by destination partition ──
|
||||
// `KmerDesc.seq_idx` carries the origin (this layer's local slot);
|
||||
// `KmerDesc.pos` is repurposed as a 0/1 flag: 1 iff this specific
|
||||
// variant's own encoding is smaller than the source's — decided here,
|
||||
// at scatter time, since both encodings are already in hand; the
|
||||
// query response only needs to confirm existence (hit/miss).
|
||||
let mut outgoing: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> =
|
||||
(0..n_parts).map(|_| HashMap::new()).collect();
|
||||
|
||||
for (slot, maybe_kmer) in slot_kmer.iter().enumerate() {
|
||||
let Some(kmer) = maybe_kmer else { continue };
|
||||
for variant in kmer.central_canonical_neighbors() {
|
||||
if variant == *kmer {
|
||||
continue; // identity substitution — not a real variant
|
||||
}
|
||||
let smaller = variant.raw() < kmer.raw();
|
||||
let dest = partition_of(variant, n_parts);
|
||||
outgoing[dest].entry(variant).or_default().push(KmerDesc {
|
||||
seq_idx: slot as u32,
|
||||
pos: smaller as u32,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// ── Gather + reconcile ───────────────────────────────────────────────
|
||||
let mut state = vec![RunningState::default(); n_slots];
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
let with_counts = self.meta.config.with_counts;
|
||||
|
||||
for (dest, kmers) in outgoing.iter().enumerate() {
|
||||
if kmers.is_empty() {
|
||||
continue;
|
||||
}
|
||||
self.partition()
|
||||
.query_partition_with(dest, kmers, n_genomes, with_counts, |hit| {
|
||||
if let QueryHit::Found(descs) = hit {
|
||||
for d in descs {
|
||||
let slot = d.seq_idx as usize;
|
||||
let st = &mut state[slot];
|
||||
st.siblings = (st.siblings + 1).min(3);
|
||||
if d.pos == 1 {
|
||||
st.minorant = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.map_err(OKIError::Partition)?;
|
||||
}
|
||||
|
||||
// ── Write the layer's annex file ─────────────────────────────────────
|
||||
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
|
||||
let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?;
|
||||
for (slot, st) in state.iter().enumerate() {
|
||||
if slot_kmer[slot].is_none() {
|
||||
continue; // unused MPHF slot, if any — leave at the sentinel
|
||||
}
|
||||
builder.set(slot, SiblingInfo { minorant: st.minorant, siblings: st.siblings });
|
||||
}
|
||||
builder.close()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::io::Write;
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::SiblingAnnex;
|
||||
use obikseq::{Kmer, Sequence};
|
||||
use obisys::Reporter;
|
||||
use tempfile::tempdir;
|
||||
|
||||
use crate::meta::{GenomeInfo, IndexConfig};
|
||||
use crate::merge::MergeMode;
|
||||
|
||||
use super::*;
|
||||
|
||||
// k must be >= 11 (project constraint, "k ∈ [11,31]"); k=11, level_max=1,
|
||||
// theta=0.0 mirror `obiskbuilder`'s own tests (smaller k/level_max
|
||||
// combinations trip an unrelated pre-existing bug in `obikentropy`'s
|
||||
// sliding-window ring buffer — not this feature's concern).
|
||||
const K: usize = 11;
|
||||
const M: usize = 5;
|
||||
|
||||
/// Build a single-genome index from one in-memory FASTA sequence, driving
|
||||
/// the same primitives `obikmer`'s `scatter` step uses (minus the
|
||||
/// multi-file `obipipeline` wrapper — a single sequence needs none of
|
||||
/// that): normalise -> build superkmers -> route -> write.
|
||||
fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
|
||||
let fasta_path = dir.join(format!("{label}.fasta"));
|
||||
let mut f = std::fs::File::create(&fasta_path).unwrap();
|
||||
writeln!(f, ">{label}").unwrap();
|
||||
f.write_all(seq).unwrap();
|
||||
writeln!(f).unwrap();
|
||||
drop(f);
|
||||
|
||||
let index_path = dir.join(format!("{label}.idx"));
|
||||
let config = IndexConfig {
|
||||
kmer_size: K,
|
||||
minimizer_size: M,
|
||||
n_bits: 0, // 1 partition — keeps the test deterministic and simple
|
||||
with_counts: false,
|
||||
evidence: obilayeredmap::IndexMode::Exact,
|
||||
block_bits: 0,
|
||||
};
|
||||
let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)), false)
|
||||
.expect("create");
|
||||
|
||||
let mut rep = Reporter::new();
|
||||
let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
|
||||
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");
|
||||
}
|
||||
idx.partition_mut().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");
|
||||
idx
|
||||
}
|
||||
|
||||
fn canonical(ascii: &[u8]) -> CanonicalKmer {
|
||||
Kmer::from_ascii(ascii).unwrap().canonical()
|
||||
}
|
||||
|
||||
/// 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) -> SiblingInfo {
|
||||
let index_dir = idx.partition().part_dir(0).join(INDEX_SUBDIR);
|
||||
let meta = PartitionMeta::load(&index_dir).unwrap();
|
||||
for l in 0..meta.n_layers {
|
||||
let layer_dir = index_dir.join(format!("layer_{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();
|
||||
return annex.get(slot).expect("slot must have a computed annex entry");
|
||||
}
|
||||
}
|
||||
panic!("kmer not found in any layer of partition 0");
|
||||
}
|
||||
|
||||
fn merge_two(dir: &Path, g1: &KmerIndex, g2: &KmerIndex) -> KmerIndex {
|
||||
let mut rep = Reporter::new();
|
||||
KmerIndex::merge(
|
||||
&dir.join("merged.idx"),
|
||||
&[g1, g2],
|
||||
MergeMode::Presence,
|
||||
false,
|
||||
false,
|
||||
1.0,
|
||||
&mut rep,
|
||||
)
|
||||
.expect("merge")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sibling_annex_one_sibling_each() {
|
||||
// k=11, centre = index 5 (0-based). Two genomes, each exactly one
|
||||
// k-mer, sharing every base except the centre:
|
||||
// g1 = "AACCGCTTAAG" (centre 'C')
|
||||
// g2 = "AACCGGTTAAG" (centre 'G')
|
||||
// Hand-verified: both stay forward-oriented under canonicalisation
|
||||
// (each is lexicographically smaller than its own reverse
|
||||
// complement, since both start with "AA"), and raw(g1) < raw(g2)
|
||||
// (only differing base: C=0b01 < G=0b10 at the centre) — so g1 is
|
||||
// the minorant, g2 is not, and each is the other's one sibling.
|
||||
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");
|
||||
|
||||
let a = annex_info_for(&merged, canonical(b"AACCGCTTAAG"));
|
||||
assert_eq!(a, SiblingInfo { minorant: true, siblings: 1 }, "AACCGCTTAAG");
|
||||
|
||||
let b = annex_info_for(&merged, canonical(b"AACCGGTTAAG"));
|
||||
assert_eq!(b, SiblingInfo { minorant: false, siblings: 1 }, "AACCGGTTAAG");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sibling_annex_zero_siblings_when_identical_and_no_variant_exists() {
|
||||
// Same k-mer in both genomes, no other genome around to carry a
|
||||
// variant -> 0 siblings, trivially its own minorant.
|
||||
let dir = tempdir().unwrap();
|
||||
let g1 = build_single_genome_index(dir.path(), "g1", b"GATTACAGATC");
|
||||
let g2 = build_single_genome_index(dir.path(), "g2", b"GATTACAGATC");
|
||||
let merged = merge_two(dir.path(), &g1, &g2);
|
||||
merged.build_sibling_annex().expect("build_sibling_annex");
|
||||
|
||||
let info = annex_info_for(&merged, canonical(b"GATTACAGATC"));
|
||||
assert_eq!(info, SiblingInfo { minorant: true, siblings: 0 }, "GATTACAGATC");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user