Rename index methods to hash and expose kmer iteration APIs
The layer API has been updated to rename `index` and `index_batch` to `hash` and `hash_batch`, aligning with raw MPHF slot hashing behavior. New public methods have been added to directly expose kmer hashing and iteration, enforcing strict state validation on empty layers. All other modifications are consistent code formatting adjustments without functional changes.
This commit is contained in:
@@ -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")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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 ──────────────────────────────────────────────────────────────
|
||||||
@@ -58,7 +69,8 @@ impl EvidenceKind {
|
|||||||
(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()
|
||||||
))),
|
))),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -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,32 +205,10 @@ 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.
|
||||||
@@ -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)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -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,8 +476,7 @@ 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()))?;
|
||||||
@@ -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());
|
||||||
|
|||||||
@@ -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};
|
||||||
@@ -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)
|
||||||
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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())
|
||||||
|
|||||||
@@ -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,7 +217,10 @@ 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 =
|
||||||
|
ctx.mat
|
||||||
|
.iter_minorants_batch(annex, FAMILY_BATCH)
|
||||||
|
.scan(0usize, |offset, batch| {
|
||||||
let start = *offset;
|
let start = *offset;
|
||||||
*offset += batch.len();
|
*offset += batch.len();
|
||||||
Some((start, batch))
|
Some((start, batch))
|
||||||
@@ -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,9 +332,15 @@ 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
|
||||||
|
.outgoing
|
||||||
|
.par_iter()
|
||||||
|
.enumerate()
|
||||||
|
.filter(|(_, q)| !q.is_empty())
|
||||||
|
.for_each(|(dest, queries)| {
|
||||||
let on_hit = |i: usize, base: u8, g: usize| {
|
let on_hit = |i: usize, base: u8, g: usize| {
|
||||||
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
|
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
|
||||||
};
|
};
|
||||||
@@ -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(())
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user