Push zunrplorkwkt #70

Merged
coissac merged 93 commits from push-zunrplorkwkt into main 2026-08-28 23:15:38 +00:00
5 changed files with 294 additions and 145 deletions
Showing only changes of commit 346095b9eb - Show all commits
+89 -16
View File
@@ -28,9 +28,11 @@ use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer; use obikseq::CanonicalKmer;
use crate::layer::error::OLMResult; use crate::layer::error::OLMResult;
use crate::layer::typed_layer::{LayerContent, TypedLayer, COUNTS_DIR, PRESENCE_DIR};
use crate::layer::meta::IndexMode; use crate::layer::meta::IndexMode;
use crate::layer::mphf_layer::{EvidenceKind, EVIDENCE_FILE, FINGERPRINT_FILE, MPHF_FILE, UNITIGS_FILE}; use crate::layer::mphf_layer::{
EVIDENCE_FILE, EvidenceKind, FINGERPRINT_FILE, MPHF_FILE, UNITIGS_FILE,
};
use crate::layer::typed_layer::{COUNTS_DIR, LayerContent, PRESENCE_DIR, TypedLayer};
/// One layer, at any point in its life — see the module docs. Only /// One layer, at any point in its life — see the module docs. Only
/// [`Empty`](Layer::Empty) and the two ready-to-read states /// [`Empty`](Layer::Empty) and the two ready-to-read states
@@ -42,7 +44,9 @@ pub enum Layer {
/// the path accessors panics on this variant: calling them means the /// the path accessors panics on this variant: calling them means the
/// caller assumed a layer was ready when it wasn't, an implementation /// caller assumed a layer was ready when it wasn't, an implementation
/// error to surface loudly, not paper over with a default value. /// error to surface loudly, not paper over with a default value.
Empty { dir: PathBuf }, Empty {
dir: PathBuf,
},
Count(TypedLayer<PersistentCompactIntMatrix>), Count(TypedLayer<PersistentCompactIntMatrix>),
Presence(TypedLayer<PersistentBitMatrix>), Presence(TypedLayer<PersistentBitMatrix>),
} }
@@ -53,7 +57,9 @@ impl Layer {
/// building a new layer. /// building a new layer.
pub fn create(dir: &Path) -> std::io::Result<Self> { pub fn create(dir: &Path) -> std::io::Result<Self> {
std::fs::create_dir_all(dir)?; std::fs::create_dir_all(dir)?;
Ok(Layer::Empty { dir: dir.to_owned() }) Ok(Layer::Empty {
dir: dir.to_owned(),
})
} }
/// Open one layer, auto-detecting count vs. presence from what is /// Open one layer, auto-detecting count vs. presence from what is
@@ -66,7 +72,8 @@ impl Layer {
/// here. /// here.
pub fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult<Self> { pub fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult<Self> {
if with_counts && layer_dir.join(COUNTS_DIR).exists() { if with_counts && layer_dir.join(COUNTS_DIR).exists() {
return TypedLayer::<PersistentCompactIntMatrix>::open(layer_dir, mode).map(Layer::Count); return TypedLayer::<PersistentCompactIntMatrix>::open(layer_dir, mode)
.map(Layer::Count);
} }
TypedLayer::<PersistentBitMatrix>::open(layer_dir, mode).map(Layer::Presence) TypedLayer::<PersistentBitMatrix>::open(layer_dir, mode).map(Layer::Presence)
} }
@@ -91,12 +98,24 @@ impl Layer {
} }
} }
pub fn mphf_path(&self) -> PathBuf { self.dir().join(MPHF_FILE) } pub fn mphf_path(&self) -> PathBuf {
pub fn unitigs_path(&self) -> PathBuf { self.dir().join(UNITIGS_FILE) } self.dir().join(MPHF_FILE)
pub fn evidence_path(&self) -> PathBuf { self.dir().join(EVIDENCE_FILE) } }
pub fn fingerprint_path(&self) -> PathBuf { self.dir().join(FINGERPRINT_FILE) } pub fn unitigs_path(&self) -> PathBuf {
pub fn counts_dir(&self) -> PathBuf { self.dir().join(COUNTS_DIR) } self.dir().join(UNITIGS_FILE)
pub fn presence_dir(&self) -> PathBuf { self.dir().join(PRESENCE_DIR) } }
pub fn evidence_path(&self) -> PathBuf {
self.dir().join(EVIDENCE_FILE)
}
pub fn fingerprint_path(&self) -> PathBuf {
self.dir().join(FINGERPRINT_FILE)
}
pub fn counts_dir(&self) -> PathBuf {
self.dir().join(COUNTS_DIR)
}
pub fn presence_dir(&self) -> PathBuf {
self.dir().join(PRESENCE_DIR)
}
// ── Ready-only surface ─────────────────────────────────────────────── // ── Ready-only surface ───────────────────────────────────────────────
@@ -135,11 +154,11 @@ impl Layer {
/// Raw MPHF batch lookup: kmer → slot, no membership check — for /// Raw MPHF batch lookup: kmer → slot, no membership check — for
/// callers that already know every kmer is a member of *this* layer, so /// callers that already know every kmer is a member of *this* layer, so
/// the evidence check `find_slot`/`find` would perform is redundant. /// the evidence check `find_slot`/`find` would perform is redundant.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> { pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
match self { match self {
Layer::Count(l) => l.index_batch(kmers), Layer::Count(l) => l.hash_batch(kmers),
Layer::Presence(l) => l.index_batch(kmers), Layer::Presence(l) => l.hash_batch(kmers),
Layer::Empty { .. } => panic!("Layer::index_batch() called on an Empty layer"), Layer::Empty { .. } => panic!("Layer::hash_batch() called on an Empty layer"),
} }
} }
@@ -168,7 +187,61 @@ impl Layer {
o.extend(c.iter().map(|&v| v != 0)); o.extend(c.iter().map(|&v| v != 0));
} }
} }
Layer::Empty { .. } => panic!("Layer::fill_sub_matrix_carries() called on an Empty layer"), Layer::Empty { .. } => {
panic!("Layer::fill_sub_matrix_carries() called on an Empty layer")
}
}
}
/// Raw MPHF lookup: kmer → slot, no membership check.
pub fn hash(&self, kmer: CanonicalKmer) -> usize {
match self {
Layer::Count(l) => l.hash(kmer),
Layer::Presence(l) => l.hash(kmer),
Layer::Empty { .. } => panic!("Layer::hash() called on an Empty layer"),
}
}
/// Iterate over all canonical kmers in the layer, in deterministic order.
pub fn iter_kmers(&self) -> crate::layer::mphf_layer::KmerIter {
match self {
Layer::Count(l) => l.iter_kmers(),
Layer::Presence(l) => l.iter_kmers(),
Layer::Empty { .. } => panic!("Layer::iter_kmers() called on an Empty layer"),
}
}
/// Iterate over all canonical kmers, each paired with its zero-based
/// sequence index in `unitigs.bin`.
pub fn enumerate_kmers(&self) -> std::iter::Enumerate<crate::layer::mphf_layer::KmerIter> {
match self {
Layer::Count(l) => l.enumerate_kmers(),
Layer::Presence(l) => l.enumerate_kmers(),
Layer::Empty { .. } => panic!("Layer::enumerate_kmers() called on an Empty layer"),
}
}
/// Iterate over the layer's canonical kmers in batches of `n`.
pub fn iter_kmers_batch(&self, n: usize) -> crate::layer::mphf_layer::KmerBatchIter {
match self {
Layer::Count(l) => l.iter_kmers_batch(n),
Layer::Presence(l) => l.iter_kmers_batch(n),
Layer::Empty { .. } => panic!("Layer::iter_kmers_batch() called on an Empty layer"),
}
}
/// Iterate over batches, each paired with the zero-based index of the
/// first kmer in the batch.
pub fn enumerate_kmers_batch(
&self,
n: usize,
) -> Box<dyn Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static> {
match self {
Layer::Count(l) => Box::new(l.enumerate_kmers_batch(n)),
Layer::Presence(l) => Box::new(l.enumerate_kmers_batch(n)),
Layer::Empty { .. } => {
panic!("Layer::enumerate_kmers_batch() called on an Empty layer")
}
} }
} }
} }
+118 -78
View File
@@ -15,9 +15,9 @@ use crate::layer::evidence::{Evidence, EvidenceWriter};
use crate::layer::fingerprint::{FingerprintVec, FingerprintVecWriter}; use crate::layer::fingerprint::{FingerprintVec, FingerprintVecWriter};
use crate::layer::meta::IndexMode; use crate::layer::meta::IndexMode;
pub(crate) const MPHF_FILE: &str = "mphf.bin"; pub(crate) const MPHF_FILE: &str = "mphf.bin";
pub(crate) const UNITIGS_FILE: &str = "unitigs.bin"; pub(crate) const UNITIGS_FILE: &str = "unitigs.bin";
pub(crate) const EVIDENCE_FILE: &str = "evidence.bin"; pub(crate) const EVIDENCE_FILE: &str = "evidence.bin";
pub(crate) const FINGERPRINT_FILE: &str = "fingerprint.bin"; pub(crate) const FINGERPRINT_FILE: &str = "fingerprint.bin";
/// Owned MPHF — used only at build time (construction + store). /// Owned MPHF — used only at build time (construction + store).
@@ -30,9 +30,20 @@ type MphfEps = PtrHash<u64, CubicEps, CachelineEfVec<&'static [CachelineEf]>, Xx
// ── LayerEvidence ───────────────────────────────────────────────────────────── // ── LayerEvidence ─────────────────────────────────────────────────────────────
enum LayerEvidence { enum LayerEvidence {
Exact { evidence: Evidence, unitigs: Arc<UnitigFileReader> }, Exact {
Approx { fingerprint: FingerprintVec, unitigs: Arc<UnitigFileReader>, unitigs_path: PathBuf }, evidence: Evidence,
Hybrid { evidence: Evidence, unitigs: Arc<UnitigFileReader>, fingerprint: FingerprintVec }, unitigs: Arc<UnitigFileReader>,
},
Approx {
fingerprint: FingerprintVec,
unitigs: Arc<UnitigFileReader>,
unitigs_path: PathBuf,
},
Hybrid {
evidence: Evidence,
unitigs: Arc<UnitigFileReader>,
fingerprint: FingerprintVec,
},
} }
// ── EvidenceKind ────────────────────────────────────────────────────────────── // ── EvidenceKind ──────────────────────────────────────────────────────────────
@@ -51,14 +62,15 @@ impl EvidenceKind {
/// `evidence.bin`/`fingerprint.bin` alone determines the mode — same /// `evidence.bin`/`fingerprint.bin` alone determines the mode — same
/// signal `MphfLayer::open` uses, just without opening either file. /// signal `MphfLayer::open` uses, just without opening either file.
pub fn detect(layer_dir: &Path) -> OLMResult<EvidenceKind> { pub fn detect(layer_dir: &Path) -> OLMResult<EvidenceKind> {
let has_evidence = layer_dir.join(EVIDENCE_FILE).exists(); let has_evidence = layer_dir.join(EVIDENCE_FILE).exists();
let has_fingerprint = layer_dir.join(FINGERPRINT_FILE).exists(); let has_fingerprint = layer_dir.join(FINGERPRINT_FILE).exists();
match (has_evidence, has_fingerprint) { match (has_evidence, has_fingerprint) {
(true, false) => Ok(EvidenceKind::Exact), (true, false) => Ok(EvidenceKind::Exact),
(false, true) => Ok(EvidenceKind::Approx), (false, true) => Ok(EvidenceKind::Approx),
(true, true) => Ok(EvidenceKind::Hybrid), (true, true) => Ok(EvidenceKind::Hybrid),
(false, false) => Err(OLMError::InvalidLayer(format!( (false, false) => Err(OLMError::InvalidLayer(format!(
"no evidence.bin or fingerprint.bin in {}", layer_dir.display() "no evidence.bin or fingerprint.bin in {}",
layer_dir.display()
))), ))),
} }
} }
@@ -74,8 +86,8 @@ impl EvidenceKind {
/// O(n) sequential scan on Approx layers. /// O(n) sequential scan on Approx layers.
pub struct MphfLayer { pub struct MphfLayer {
mphf: MemCase<MphfEps>, mphf: MemCase<MphfEps>,
ev: LayerEvidence, ev: LayerEvidence,
n: usize, n: usize,
} }
impl MphfLayer { impl MphfLayer {
@@ -96,14 +108,28 @@ impl MphfLayer {
let n = fingerprint.n(); let n = fingerprint.n();
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?); let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
let unitigs_path = dir.join(UNITIGS_FILE); let unitigs_path = dir.join(UNITIGS_FILE);
(LayerEvidence::Approx { fingerprint, unitigs, unitigs_path }, n) (
LayerEvidence::Approx {
fingerprint,
unitigs,
unitigs_path,
},
n,
)
} }
IndexMode::Hybrid { .. } => { IndexMode::Hybrid { .. } => {
let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?; let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?;
let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?; let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?;
let n = evidence.len(); let n = evidence.len();
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?); let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
(LayerEvidence::Hybrid { evidence, unitigs, fingerprint }, n) (
LayerEvidence::Hybrid {
evidence,
unitigs,
fingerprint,
},
n,
)
} }
}; };
Ok(Self { mphf, ev, n }) Ok(Self { mphf, ev, n })
@@ -118,9 +144,13 @@ impl MphfLayer {
#[inline] #[inline]
pub fn find(&self, kmer: CanonicalKmer) -> Option<usize> { pub fn find(&self, kmer: CanonicalKmer) -> Option<usize> {
let slot = self.mphf.index(&kmer.raw()); let slot = self.mphf.index(&kmer.raw());
if slot >= self.n { return None; } if slot >= self.n {
return None;
}
match &self.ev { match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } => { LayerEvidence::Exact {
evidence, unitigs, ..
} => {
let (chunk_id, rank) = evidence.decode(slot); let (chunk_id, rank) = evidence.decode(slot);
if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) { if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
Some(slot) Some(slot)
@@ -128,9 +158,13 @@ impl MphfLayer {
None None
} }
} }
LayerEvidence::Approx { fingerprint, .. } | LayerEvidence::Approx { fingerprint, .. }
LayerEvidence::Hybrid { fingerprint, .. } => { | LayerEvidence::Hybrid { fingerprint, .. } => {
if fingerprint.matches(slot, kmer.seq_hash()) { Some(slot) } else { None } if fingerprint.matches(slot, kmer.seq_hash()) {
Some(slot)
} else {
None
}
} }
} }
} }
@@ -142,10 +176,16 @@ impl MphfLayer {
/// that owns the slot, then exact comparison. /// that owns the slot, then exact comparison.
pub fn find_strict(&self, kmer: CanonicalKmer) -> Option<usize> { pub fn find_strict(&self, kmer: CanonicalKmer) -> Option<usize> {
let slot = self.mphf.index(&kmer.raw()); let slot = self.mphf.index(&kmer.raw());
if slot >= self.n { return None; } if slot >= self.n {
return None;
}
match &self.ev { match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } | LayerEvidence::Exact {
LayerEvidence::Hybrid { evidence, unitigs, .. } => { evidence, unitigs, ..
}
| LayerEvidence::Hybrid {
evidence, unitigs, ..
} => {
let (chunk_id, rank) = evidence.decode(slot); let (chunk_id, rank) = evidence.decode(slot);
if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) { if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
Some(slot) Some(slot)
@@ -165,38 +205,16 @@ impl MphfLayer {
} }
} }
/// Reconstruct the canonical k-mer stored at `slot` — the inverse of /// Number of slots in the MPHF layer — equal to the number of stored k-mers.
/// [`find`](Self::find)/[`find_strict`](Self::find_strict) (k-mer → slot). pub fn n(&self) -> usize {
/// O(1) on `Exact`/`Hybrid` layers: `evidence.decode(slot)` gives self.n
/// `(chunk_id, rank)` directly (no MPHF hashing, no file scan), then
/// `unitigs.canonical_raw_kmer` is a direct-access read. `None` on
/// `Approx` layers — fingerprints alone can't recover a k-mer, and
/// falling back to a full sequential scan here would silently make one
/// call cost O(n); callers needing this on an `Approx` layer should
/// scan `unitigs.bin` themselves and decide how to handle that cost
/// explicitly.
pub fn kmer_at(&self, slot: usize) -> Option<CanonicalKmer> {
if slot >= self.n {
return None;
}
match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } |
LayerEvidence::Hybrid { evidence, unitigs, .. } => {
let (chunk_id, rank) = evidence.decode(slot);
let raw = unitigs.canonical_raw_kmer(chunk_id as usize, rank as usize);
Some(CanonicalKmer::from_raw_unchecked(raw))
}
LayerEvidence::Approx { .. } => None,
}
} }
pub fn n(&self) -> usize { self.n }
/// This already-open layer's evidence mode — reads the discriminant /// This already-open layer's evidence mode — reads the discriminant
/// already in memory, no disk access. /// already in memory, no disk access.
pub fn evidence_kind(&self) -> EvidenceKind { pub fn evidence_kind(&self) -> EvidenceKind {
match &self.ev { match &self.ev {
LayerEvidence::Exact { .. } => EvidenceKind::Exact, LayerEvidence::Exact { .. } => EvidenceKind::Exact,
LayerEvidence::Approx { .. } => EvidenceKind::Approx, LayerEvidence::Approx { .. } => EvidenceKind::Approx,
LayerEvidence::Hybrid { .. } => EvidenceKind::Hybrid, LayerEvidence::Hybrid { .. } => EvidenceKind::Hybrid,
} }
@@ -208,7 +226,7 @@ impl MphfLayer {
/// check against the layer's evidence. The returned slot is only meaningful /// check against the layer's evidence. The returned slot is only meaningful
/// when the kmer is actually present in the layer; callers that need an /// when the kmer is actually present in the layer; callers that need an
/// existence test should use [`find`](Self::find) instead. /// existence test should use [`find`](Self::find) instead.
pub fn index(&self, kmer: CanonicalKmer) -> usize { pub fn hash(&self, kmer: CanonicalKmer) -> usize {
self.mphf.index(&kmer.raw()) self.mphf.index(&kmer.raw())
} }
@@ -217,7 +235,7 @@ impl MphfLayer {
/// Returns a [`Vec`] of slots, one per input kmer, in the same order as the /// Returns a [`Vec`] of slots, one per input kmer, in the same order as the
/// input slice. Like [`index`](Self::index), no membership check is /// input slice. Like [`index`](Self::index), no membership check is
/// performed. /// performed.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> { pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
kmers.iter().map(|k| self.mphf.index(&k.raw())).collect() kmers.iter().map(|k| self.mphf.index(&k.raw())).collect()
} }
@@ -237,7 +255,9 @@ impl MphfLayer {
LayerEvidence::Approx { unitigs, .. } => unitigs, LayerEvidence::Approx { unitigs, .. } => unitigs,
LayerEvidence::Hybrid { unitigs, .. } => unitigs, LayerEvidence::Hybrid { unitigs, .. } => unitigs,
}; };
KmerIter { inner: Box::new(reader.iter_indexed_canonical_kmers_owned()) } KmerIter {
inner: Box::new(reader.iter_indexed_canonical_kmers_owned()),
}
} }
/// Iterate over all canonical kmers, each paired with its zero-based /// Iterate over all canonical kmers, each paired with its zero-based
@@ -255,7 +275,10 @@ impl MphfLayer {
/// Each call to [`next`](Iterator::next) returns a [`Vec`] of up to `n` /// Each call to [`next`](Iterator::next) returns a [`Vec`] of up to `n`
/// kmers. The final batch may be shorter when the layer is exhausted. /// kmers. The final batch may be shorter when the layer is exhausted.
pub fn iter_kmers_batch(&self, n: usize) -> KmerBatchIter { pub fn iter_kmers_batch(&self, n: usize) -> KmerBatchIter {
KmerBatchIter { inner: self.iter_kmers(), batch_size: n } KmerBatchIter {
inner: self.iter_kmers(),
batch_size: n,
}
} }
/// Iterate over batches, each paired with the zero-based index of the /// Iterate over batches, each paired with the zero-based index of the
@@ -265,7 +288,10 @@ impl MphfLayer {
/// `batch_start_index` is the iteration-order offset of the first kmer /// `batch_start_index` is the iteration-order offset of the first kmer
/// in that batch within the full layer sequence — i.e. a multiple of `n` /// in that batch within the full layer sequence — i.e. a multiple of `n`
/// except for the final (possibly shorter) batch. /// except for the final (possibly shorter) batch.
pub fn enumerate_kmers_batch(&self, n: usize) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static { pub fn enumerate_kmers_batch(
&self,
n: usize,
) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
let mut offset = 0usize; let mut offset = 0usize;
self.iter_kmers_batch(n).map(move |batch| { self.iter_kmers_batch(n).map(move |batch| {
let base = offset; let base = offset;
@@ -320,11 +346,7 @@ impl Iterator for KmerBatchIter {
break; break;
} }
} }
if batch.is_empty() { if batch.is_empty() { None } else { Some(batch) }
None
} else {
Some(batch)
}
} }
} }
@@ -373,7 +395,7 @@ impl MphfLayer {
let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE)) let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?; .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
let mut ev = EvidenceWriter::new(n); let mut ev = EvidenceWriter::new(n);
let mut seen = vec![0u8; (n + 7) / 8]; let mut seen = vec![0u8; (n + 7) / 8];
for (kmer, chunk_id, rank) in unitigs.iter_indexed_canonical_kmers() { for (kmer, chunk_id, rank) in unitigs.iter_indexed_canonical_kmers() {
@@ -382,7 +404,7 @@ impl MphfLayer {
return Err(OLMError::Mphf("slot out of bounds".into())); return Err(OLMError::Mphf("slot out of bounds".into()));
} }
let byte = slot / 8; let byte = slot / 8;
let bit = 1u8 << (slot % 8); let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 { if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into())); return Err(OLMError::Mphf("duplicate slot".into()));
} }
@@ -398,7 +420,9 @@ impl MphfLayer {
/// Build `fingerprint.bin` from `unitigs.bin` + `mphf.bin`. /// Build `fingerprint.bin` from `unitigs.bin` + `mphf.bin`.
pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OLMResult<usize> { pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
if b == 0 || b > 64 { if b == 0 || b > 64 {
return Err(OLMError::InvalidLayer("fingerprint width must be 1..=64".into())); return Err(OLMError::InvalidLayer(
"fingerprint width must be 1..=64".into(),
));
} }
if z == 0 { if z == 0 {
return Err(OLMError::InvalidLayer("z must be ≥ 1".into())); return Err(OLMError::InvalidLayer("z must be ≥ 1".into()));
@@ -452,9 +476,8 @@ impl MphfLayer {
// ── Empty layer ─────────────────────────────────────────────────────── // ── Empty layer ───────────────────────────────────────────────────────
if n == 0 { if n == 0 {
let mphf: Mphf = let mphf: Mphf = Mphf::try_new(&[] as &[u64], PtrHashParams::<CubicEps>::default())
Mphf::try_new(&[] as &[u64], PtrHashParams::<CubicEps>::default()) .ok_or_else(|| OLMError::Mphf("construction failed".into()))?;
.ok_or_else(|| OLMError::Mphf("construction failed".into()))?;
mphf.store(&dir.join(MPHF_FILE)) mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?; .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
match mode { match mode {
@@ -474,9 +497,11 @@ impl MphfLayer {
// ── Pass 1: MPHF via clonable mmap iterator ─────────────────────────── // ── Pass 1: MPHF via clonable mmap iterator ───────────────────────────
let keys = CanonicalKmerIter::new(&unitig_path).map_err(sk_to_olm)?; let keys = CanonicalKmerIter::new(&unitig_path).map_err(sk_to_olm)?;
let mphf: Mphf = let mphf: Mphf = Mphf::new_from_par_iter(
Mphf::new_from_par_iter(n, keys.map(|k| k.raw()).par_bridge(), n,
PtrHashParams::<CubicEps>::default()); keys.map(|k| k.raw()).par_bridge(),
PtrHashParams::<CubicEps>::default(),
);
mphf.store(&dir.join(MPHF_FILE)) mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?; .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
@@ -489,9 +514,14 @@ impl MphfLayer {
let mut ev = EvidenceWriter::new(n); let mut ev = EvidenceWriter::new(n);
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() { for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw()); let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); } if slot >= n {
let byte = slot / 8; let bit = 1u8 << (slot % 8); return Err(OLMError::Mphf("slot out of bounds".into()));
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); } }
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
}
seen[byte] |= bit; seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8); ev.set(slot, chunk_id as u32, rank as u8);
fill_slot(slot, kmer)?; fill_slot(slot, kmer)?;
@@ -504,9 +534,14 @@ impl MphfLayer {
let mut fw = FingerprintVecWriter::new(n, *b); let mut fw = FingerprintVecWriter::new(n, *b);
for (kmer, _, _) in unitigs2.iter_indexed_canonical_kmers() { for (kmer, _, _) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw()); let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); } if slot >= n {
let byte = slot / 8; let bit = 1u8 << (slot % 8); return Err(OLMError::Mphf("slot out of bounds".into()));
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); } }
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
}
seen[byte] |= bit; seen[byte] |= bit;
fw.set(slot, kmer.seq_hash()); fw.set(slot, kmer.seq_hash());
fill_slot(slot, kmer)?; fill_slot(slot, kmer)?;
@@ -519,9 +554,14 @@ impl MphfLayer {
let mut fw = FingerprintVecWriter::new(n, *b); let mut fw = FingerprintVecWriter::new(n, *b);
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() { for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw()); let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); } if slot >= n {
let byte = slot / 8; let bit = 1u8 << (slot % 8); return Err(OLMError::Mphf("slot out of bounds".into()));
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); } }
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
}
seen[byte] |= bit; seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8); ev.set(slot, chunk_id as u32, rank as u8);
fw.set(slot, kmer.seq_hash()); fw.set(slot, kmer.seq_hash());
+41 -21
View File
@@ -3,10 +3,8 @@ use std::fs;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use obicompactvec::{ use obicompactvec::{
BinaryMatrix, BinaryMatrix, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrix,
PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrixBuilder, PersistentSparseBitMatrix,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
PersistentSparseBitMatrix,
}; };
use obikseq::CanonicalKmer; use obikseq::CanonicalKmer;
use obiskio::{UnitigFileReader, UnitigFileWriter}; use obiskio::{UnitigFileReader, UnitigFileWriter};
@@ -16,7 +14,7 @@ use crate::layer::meta::IndexMode;
use crate::layer::mphf_layer::MphfLayer; use crate::layer::mphf_layer::MphfLayer;
pub(crate) use crate::layer::mphf_layer::UNITIGS_FILE; pub(crate) use crate::layer::mphf_layer::UNITIGS_FILE;
pub(crate) const COUNTS_DIR: &str = "counts"; pub(crate) const COUNTS_DIR: &str = "counts";
pub(crate) const PRESENCE_DIR: &str = "presence"; pub(crate) const PRESENCE_DIR: &str = "presence";
// ── Trait ───────────────────────────────────────────────────────────────────── // ── Trait ─────────────────────────────────────────────────────────────────────
@@ -82,7 +80,9 @@ pub fn open_data<D: LayerData>(root: &Path, i: usize) -> OLMResult<D> {
impl LayerData for () { impl LayerData for () {
type Item = (); type Item = ();
fn open(_layer_dir: &Path) -> OLMResult<Self> { Ok(()) } fn open(_layer_dir: &Path) -> OLMResult<Self> {
Ok(())
}
fn read(&self, _slot: usize) {} fn read(&self, _slot: usize) {}
} }
@@ -91,7 +91,9 @@ impl LayerData for PersistentCompactIntMatrix {
fn open(layer_dir: &Path) -> OLMResult<Self> { fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentCompactIntMatrix::open(layer_dir).map_err(OLMError::Io) PersistentCompactIntMatrix::open(layer_dir).map_err(OLMError::Io)
} }
fn read(&self, slot: usize) -> Box<[u32]> { self.row(slot) } fn read(&self, slot: usize) -> Box<[u32]> {
self.row(slot)
}
} }
impl LayerData for PersistentBitMatrix { impl LayerData for PersistentBitMatrix {
@@ -99,7 +101,9 @@ impl LayerData for PersistentBitMatrix {
fn open(layer_dir: &Path) -> OLMResult<Self> { fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentBitMatrix::open(layer_dir).map_err(OLMError::Io) PersistentBitMatrix::open(layer_dir).map_err(OLMError::Io)
} }
fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) } fn read(&self, slot: usize) -> Box<[bool]> {
self.row(slot)
}
} }
impl LayerData for PersistentSparseBitMatrix { impl LayerData for PersistentSparseBitMatrix {
@@ -107,7 +111,9 @@ impl LayerData for PersistentSparseBitMatrix {
fn open(layer_dir: &Path) -> OLMResult<Self> { fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentSparseBitMatrix::open(&layer_dir.join(PRESENCE_DIR)).map_err(OLMError::Io) PersistentSparseBitMatrix::open(&layer_dir.join(PRESENCE_DIR)).map_err(OLMError::Io)
} }
fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) } fn read(&self, slot: usize) -> Box<[bool]> {
self.row(slot)
}
} }
// ── LayerContent ───────────────────────────────────────────────────────────── // ── LayerContent ─────────────────────────────────────────────────────────────
@@ -202,7 +208,10 @@ impl<D: LayerData> TypedLayer<D> {
} }
pub fn query(&self, kmer: CanonicalKmer) -> Option<Hit<D::Item>> { pub fn query(&self, kmer: CanonicalKmer) -> Option<Hit<D::Item>> {
self.mphf.find(kmer).map(|slot| Hit { slot, data: self.data.read(slot) }) self.mphf.find(kmer).map(|slot| Hit {
slot,
data: self.data.read(slot),
})
} }
/// MPHF + evidence membership check only — no data read. For callers /// MPHF + evidence membership check only — no data read. For callers
@@ -214,16 +223,18 @@ impl<D: LayerData> TypedLayer<D> {
self.mphf.find(kmer) self.mphf.find(kmer)
} }
pub fn n(&self) -> usize { self.mphf.n() } pub fn n(&self) -> usize {
self.mphf.n()
}
/// Raw MPHF lookup: kmer → slot, no membership check. /// Raw MPHF lookup: kmer → slot, no membership check.
pub fn index(&self, kmer: CanonicalKmer) -> usize { pub fn hash(&self, kmer: CanonicalKmer) -> usize {
self.mphf.index(kmer) self.mphf.hash(kmer)
} }
/// Batch raw MPHF lookup: kmers → slots, no membership check. /// Batch raw MPHF lookup: kmers → slots, no membership check.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> { pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
self.mphf.index_batch(kmers) self.mphf.hash_batch(kmers)
} }
/// Iterate over all canonical kmers in the layer, in deterministic order. /// Iterate over all canonical kmers in the layer, in deterministic order.
@@ -244,7 +255,10 @@ impl<D: LayerData> TypedLayer<D> {
/// Iterate over batches, each paired with the zero-based index of the /// Iterate over batches, each paired with the zero-based index of the
/// first kmer in the batch. /// first kmer in the batch.
pub fn enumerate_kmers_batch(&self, n: usize) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static { pub fn enumerate_kmers_batch(
&self,
n: usize,
) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
self.mphf.enumerate_kmers_batch(n) self.mphf.enumerate_kmers_batch(n)
} }
@@ -301,8 +315,12 @@ impl TypedLayer<()> {
pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> { pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> {
let presence_dir = layer_dir.join(PRESENCE_DIR); let presence_dir = layer_dir.join(PRESENCE_DIR);
fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?; fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?;
let mut mb = PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?; let mut mb =
mb.add_col_ones().map_err(OLMError::Io)?.close().map_err(OLMError::Io)?; PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?;
mb.add_col_ones()
.map_err(OLMError::Io)?
.close()
.map_err(OLMError::Io)?;
mb.close().map_err(OLMError::Io) mb.close().map_err(OLMError::Io)
} }
} }
@@ -318,8 +336,8 @@ impl TypedLayer<PersistentCompactIntMatrix> {
) -> OLMResult<usize> { ) -> OLMResult<usize> {
let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers(); let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
let counts_dir = out_dir.join(COUNTS_DIR); let counts_dir = out_dir.join(COUNTS_DIR);
let mut mb = PersistentCompactIntMatrixBuilder::new(n, &counts_dir) let mut mb =
.map_err(OLMError::Io)?; PersistentCompactIntMatrixBuilder::new(n, &counts_dir).map_err(OLMError::Io)?;
let mut col = mb.add_col().map_err(OLMError::Io)?; let mut col = mb.add_col().map_err(OLMError::Io)?;
let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| { let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
col.set(slot, count_of(kmer)); col.set(slot, count_of(kmer));
@@ -336,7 +354,9 @@ impl TypedLayer<PersistentCompactIntMatrix> {
mode: &IndexMode, mode: &IndexMode,
counts: &HashMap<CanonicalKmer, u32>, counts: &HashMap<CanonicalKmer, u32>,
) -> OLMResult<usize> { ) -> OLMResult<usize> {
Self::build(out_dir, block_bits, mode, |kmer| counts.get(&kmer).copied().unwrap_or(0)) Self::build(out_dir, block_bits, mode, |kmer| {
counts.get(&kmer).copied().unwrap_or(0)
})
} }
} }
+3 -7
View File
@@ -1,7 +1,7 @@
use rayon::prelude::*; use rayon::prelude::*;
use obikseq::CanonicalKmer;
use obikindex::layer::Layer; use obikindex::layer::Layer;
use obikseq::CanonicalKmer;
use obisys::progress_bar; use obisys::progress_bar;
use obikindex::{KmerIndex, OKIResult}; use obikindex::{KmerIndex, OKIResult};
@@ -45,11 +45,7 @@ pub(super) struct PartitionCache {
} }
impl PartitionCache { impl PartitionCache {
pub(super) fn build( pub(super) fn build(index: &KmerIndex, n_parts: usize, with_counts: bool) -> OKIResult<Self> {
index: &KmerIndex,
n_parts: usize,
with_counts: bool,
) -> OKIResult<Self> {
let pb = progress_bar("open_partitions", n_parts as u64, "partitions"); let pb = progress_bar("open_partitions", n_parts as u64, "partitions");
let built: Vec<(Vec<Layer>, usize)> = (0..n_parts) let built: Vec<(Vec<Layer>, usize)> = (0..n_parts)
.into_par_iter() .into_par_iter()
@@ -200,7 +196,7 @@ impl PartitionCache {
let mat = &mats[li]; let mat = &mats[li];
let variants: Vec<CanonicalKmer> = let variants: Vec<CanonicalKmer> =
entries.iter().map(|&(variant, _, _)| variant).collect(); entries.iter().map(|&(variant, _, _)| variant).collect();
let slots = mat.index_batch(&variants); let slots = mat.hash_batch(&variants);
let hits: Vec<(usize, usize, u8)> = slots let hits: Vec<(usize, usize, u8)> = slots
.into_iter() .into_iter()
.zip(entries.iter()) .zip(entries.iter())
+43 -23
View File
@@ -53,19 +53,19 @@ use std::sync::atomic::{AtomicU8, Ordering};
use rayon::prelude::*; use rayon::prelude::*;
use obikseq::CanonicalKmer;
use obikindex::layer::meta::PartitionMeta; use obikindex::layer::meta::PartitionMeta;
use obikseq::CanonicalKmer;
use obipipeline::{ThrottleGuard, throttle}; use obipipeline::{ThrottleGuard, throttle};
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex; use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::layer::Layer; use obikindex::layer::Layer;
use super::cache::PartitionCache; use super::cache::PartitionCache;
use super::helpers::central_base; use super::helpers::central_base;
use super::iter::{SiblingEntry, SiblingLayerExt}; use super::iter::{SiblingEntry, SiblingLayerExt};
use super::{olm_to_ok, FamilyMask, SiblingAnnex, ANNEX_FILE_NAME}; use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex, olm_to_ok};
/// Families per batch — see the module docs for the memory-vs-per-partition- /// 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 /// density trade-off this picks a point on. At ~90 genomes and a few
@@ -191,14 +191,22 @@ pub(super) fn scan_layer_families(
selection: &Selection, selection: &Selection,
mut on_family: impl FnMut(usize, FamilyMask, &[u8]), mut on_family: impl FnMut(usize, FamilyMask, &[u8]),
) -> OKIResult<()> { ) -> OKIResult<()> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); 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).map_err(olm_to_ok)?;
let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?); let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?);
let mat = Layer::open(layer_dir, &meta.mode, with_counts).map_err(olm_to_ok)?; let mat = Layer::open(layer_dir, &meta.mode, with_counts).map_err(olm_to_ok)?;
let n_cols = mat.n_cols().min(n_genomes); let n_cols = mat.n_cols().min(n_genomes);
let ctx = Arc::new(LayerCtx { mat, n_parts, n_genomes, n_cols, k }); let ctx = Arc::new(LayerCtx {
mat,
n_parts,
n_genomes,
n_cols,
k,
});
// Streamed straight from `iter_minorants_batch` (zips this layer's own // Streamed straight from `iter_minorants_batch` (zips this layer's own
// `iter_kmers()` with the annex, both in iteration order — never an // `iter_kmers()` with the annex, both in iteration order — never an
@@ -209,11 +217,14 @@ pub(super) fn scan_layer_families(
// (see the module docs). `.scan()` computes each batch's starting // (see the module docs). `.scan()` computes each batch's starting
// family index lazily, mirroring what the eager `chunks()`+running // family index lazily, mirroring what the eager `chunks()`+running
// `offset` used to do. // `offset` used to do.
let batches = ctx.mat.iter_minorants_batch(annex, FAMILY_BATCH).scan(0usize, |offset, batch| { let batches =
let start = *offset; ctx.mat
*offset += batch.len(); .iter_minorants_batch(annex, FAMILY_BATCH)
Some((start, batch)) .scan(0usize, |offset, batch| {
}); let start = *offset;
*offset += batch.len();
Some((start, batch))
});
let n_workers = obisys::effective_parallelism(); let n_workers = obisys::effective_parallelism();
let capacity = 4; let capacity = 4;
@@ -288,7 +299,7 @@ pub(super) fn scan_layer_families(
// — a pure MPHF lookup, no evidence check, since these are // — a pure MPHF lookup, no evidence check, since these are
// this layer's own kmers, known members by construction. // this layer's own kmers, known members by construction.
let kmers: Vec<CanonicalKmer> = batch.entries.iter().map(|e| e.kmer).collect(); let kmers: Vec<CanonicalKmer> = batch.entries.iter().map(|e| e.kmer).collect();
let slots = ctx.mat.index_batch(&kmers); let slots = ctx.mat.hash_batch(&kmers);
let mut carries: Vec<Vec<bool>> = (0..ctx.n_cols).map(|_| Vec::new()).collect(); let mut carries: Vec<Vec<bool>> = (0..ctx.n_cols).map(|_| Vec::new()).collect();
ctx.mat.fill_sub_matrix_carries(&slots, &mut carries); ctx.mat.fill_sub_matrix_carries(&slots, &mut carries);
for (g, col) in carries.iter().enumerate() { for (g, col) in carries.iter().enumerate() {
@@ -321,18 +332,24 @@ pub(super) fn scan_layer_families(
// (one batch at a time, never several concurrently — see the // (one batch at a time, never several concurrently — see the
// module docs), each thread owning one partition's queries // module docs), each thread owning one partition's queries
// contiguously until this batch is done. // contiguously until this batch is done.
let genome_mask: Vec<AtomicU8> = batch.genome_mask.into_iter().map(AtomicU8::new).collect(); let genome_mask: Vec<AtomicU8> =
batch.genome_mask.into_iter().map(AtomicU8::new).collect();
let fast_mode = cache.fast_mode(); let fast_mode = cache.fast_mode();
batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { batch
let on_hit = |i: usize, base: u8, g: usize| { .outgoing
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); .par_iter()
}; .enumerate()
if fast_mode { .filter(|(_, q)| !q.is_empty())
cache.find_presence_batch_fast(dest, queries, n_genomes, on_hit); .for_each(|(dest, queries)| {
} else { let on_hit = |i: usize, base: u8, g: usize| {
cache.find_presence_batch(dest, queries, n_genomes, on_hit); genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
} };
}); if fast_mode {
cache.find_presence_batch_fast(dest, queries, n_genomes, on_hit);
} else {
cache.find_presence_batch(dest, queries, n_genomes, on_hit);
}
});
for i in 0..n { for i in 0..n {
let family_idx = batch.start_family_idx + i; let family_idx = batch.start_family_idx + i;
@@ -347,7 +364,10 @@ pub(super) fn scan_layer_families(
next_expected += n; next_expected += n;
} }
} }
debug_assert!(pending.is_empty(), "every generated batch must have been replayed"); debug_assert!(
pending.is_empty(),
"every generated batch must have been replayed"
);
Ok(()) Ok(())
} }