Files
obikmer/docmd/architecture/siblings.md
T
Eric Coissac 45b19503a1 Add proportional subsampling and Shannon entropy calculation
Introduces `--subsample N` and `--shannon` CLI flags to cap retained variable families via proportional reservoir sampling and compute per-family Shannon entropy. Updates the family scanning API to support explicit selection filtering with early-exit optimization, resolving an indexing drift issue in monomorphic layers. Streams entropy metrics for 15-state and 4-nucleotide spaces directly to CSV while maintaining parallel processing across sibling layers.
2026-08-16 21:31:06 +02:00

24 KiB
Raw Blame History

Sibling annex — architecture (discussion)

Status: architecture decided (2026-08-14). Implementation not yet mandated.

Two index spaces, uncorrelated

Every kmer stored in a Layer lives in two independent index spaces:

  • Iteration order: its position when enumerating unitigs.bin (the superkmer file), deterministic but arbitrary with respect to slot.
  • MPHF slot: MphfLayer::index(kmer), the number the MPHF assigns.

The two are not correlated by any formula. Converting from one to the other requires either recomputing the MPHF (kmer → slot) or scanning the iteration stream (kmer → order). There is no slot → kmer operation: the MPHF is a one-way function, not an invertible bijection with a stored inverse. Any method that reconstructs a kmer from a bare slot number is wrong by construction, regardless of the mechanism used (MPHF re-hash, or evidence decode + direct unitig read). See MphfLayer::kmer_at (obilayeredmap/src/mphf_layer.rs) — flagged for removal, currently called from obikphylo/siblings/build.rs and family_scan.rs (since removed — see "Pending work" status below).

Two pipelines, never mixed

origin of the kmer membership known? correct mapping
query pipeline external (caller-supplied) no query/find/find_strict — MPHF + evidence check
iteration pipeline enumerated from this layer's own unitigs.bin yes, by construction index/index_batch — MPHF only, no evidence

Evidence exists solely to answer "is this external kmer a member of the layer" for the query pipeline. Using it (or the MPHF) to go the other way — recover a kmer from a slot, or re-verify a kmer that was just produced by iterating the layer — is a conceptual error: evidence can be probabilistic (Approx mode), so any slot→kmer attempt is unsound in general, and pointless even in Exact/Hybrid mode since the kmer was already known.

Sibling annex: an iteration-pipeline artifact only

The sibling annex (FamilyMask/SiblingAnnex, .psib, obicompactvec/src/siblingannex.rs) records, per kmer, whether it is a family minorant and which family members are present in the index. Its only consumers (obikphylo/siblings/stats.rs, family_scan.rs) enumerate it exhaustively (0..annex.len()); no query-pipeline code path touches it.

Decision: the annex must be persisted in iteration order, not slot order. This lets readers zip-iterate Layer::iter_kmers() and the annex file directly — one linear, cache-friendly pass, no MPHF/slot indirection, no kmer_at. It also enables specialized iterators building on this zip: minorants-only iteration, batch-of-kmers → batch-of-family-members, etc.

Today the annex is built and stored in slot order (build_layer_sibling_annex, siblings/build.rs): slot_kmer is populated via (0..n_slots).map(|slot| mphf.kmer_at(slot)), and the origin slot is threaded through the whole cross-partition reconciliation pipeline (variant generation, query_partition_with, final mask[slot].fetch_or(...)). This must change to iterating iter_kmers()/enumerate_kmers() and threading the iteration index instead of the slot end to end — eliminating kmer_at from the build path entirely, not just the read path. No slot-indexed intermediate is needed even during construction; the iteration-order id is sufficient throughout.

The cross-partition side of the same pipeline is unaffected: checking whether a generated family-variant kmer exists in another partition is a genuine query-pipeline operation (the variant's membership in the target partition is unknown) and must keep going through KmerPartition::query_partition_with (MPHF + evidence), never a raw index().

Pending work — done

The plan above shipped: obikphylo (a new crate — phylo-domain extension traits over obikindex::KmerIndex/obilayeredmap::Layer<D>, replacing the old obikindex::siblings module) builds and reads the annex purely in iteration order (SiblingLayerExt::iter_siblings/iter_minorants, both with batch variants, mirroring Layer<D>'s own KmerIter/KmerBatchIter shape). MphfLayer::kmer_at has no remaining callers.

A separate, unrelated bug surfaced during this work and was fixed (2026-08-14): MphfLayer::enumerate_kmers_batch computed its batch_start_index via the stdlib .enumerate() adapter, which counts batches (0, 1, 2…), not the cumulative k-mer offset the annex is actually keyed on — every batch past the first wrote its mask/annex entries at the wrong iteration-order position. Fixed by tracking a running offset instead; regression tests added (sibling_annex_no_empty_masks_after_build, sibling_histogram_does_not_panic_on_partial_last_batch).

Performance: build_sibling_annex parallelism (2026-08-14)

Investigated on a real multi-genome run (phyloskims_sal_vac, k=31/m=11). Baseline: mostly one active core, with short multi-core bursts — average ~3 cores.

Fixes that helped, kept:

  • CanonicalKmerOf::minimizer() (obikseq/src/kmer.rs) — a direct O(k) bit-arithmetic minimiser for a single isolated k-mer, replacing a RollingStat instance fed byte-by-byte through an ASCII round-trip (used by helpers::partition_of, called for every generated family variant). ~3x wall-clock improvement on its own, confirmed by sampling (obiskbuilder::rolling_stat/obikentropy frames disappeared from the hot path). CanonicalKmerOf::partition() added alongside it (wraps minimizer().seq_hash() & mask, the same routing rule KmerPartition/RoutableSuperKmer use).
  • Cross-partition resolution (outgoing.par_iter() in build_layer_sibling_annex) parallelised at the partition level — one Rayon task per non-empty outgoing[dest] bucket. For k=31/m=11, a central-base substitution changes the winning minimiser (and thus the destination partition) only when that window overlaps the central base: ~11 of the 21 possible windows do, so ~10/21 (≈48%) of generated variants route right back to the partition already being built. That self bucket ends up far larger than any other, so the per-partition split pinned one thread to it alone while the rest of the pool finished instantly — confirmed by sampling: one thread solid in MphfLayer::find, everyone else idle. Fixed by splitting each non-empty bucket into total_queries / n_workers (capped 4096) chunks before par_iter(), preserving per-partition mmap locality (each chunk stays contiguous within one partition) while letting Rayon spread an oversized bucket across several threads. Net effect of both fixes together: ~3 cores average → ~10-13 cores average on the same run, and a projected total build time of ~1h15 down to ~30min on the real phyloskims_sal_vac run this was measured against.
  • TracedBar's ETA (obisys/src/progress.rs) was silently starved: the custom progress message and the self-computed ETA text used to share one pb.set_message() slot, with the ETA holding off for 2s after any custom message — fine when custom messages are rare, broken once build_sibling_annex's per-partition callback fires more often than that. Fixed by keeping the two texts in separate fields, composed together on every render instead of one overwriting the other.

Tried and reverted — do not repeat blindly:

  • Parallelising the outer partition loop in build_sibling_annex with obikindex::PartitionRunner (already used by merge/build_layers), splitting a fixed core budget between outer (partition) and inner (pipeline + resolution) concurrency so their product wouldn't exceed the budget. Measured worse: throughput dropped over time (26 partitions/5min → 38/11-12min) and peak resolution concurrency fell from ~11-12 cores to ~7-8. Cause: this capped the resolution burst — which scales very well on its own — to make room for outer concurrency, and running several partitions' resolution at once scatters access across multiple partitions' mmap regions at once, working against the locality outgoing's per-partition grouping exists for. PartitionRunner stayed exported from obikindex (new_capped too) since it's general-purpose, but nothing in obikphylo calls it.
  • Splitting resolution chunks even finer (/(n_workers*8), cap 1024, instead of /n_workers, cap 4096) to smooth the residual sawtooth. Measured ~10% slower, wider dips, not narrower. Reverted to the original chunk sizing.

Known remaining limitation, not yet worth fixing: within one layer, the four stages (sequential unitigs.bin read → parallel generation → parallel resolution → sequential annex write) never overlap — confirmed by 1s-interval sampling: generation alone occupies ~17 threads evenly, but the next layer's read/generation never starts until the current layer's resolution and write are both done. This produces a real, periodic (~layer duration) alternation between "many cores" and "few cores" that neither of the fixes above touches, since both operate within one layer's resolution step. The only remaining lever is overlapping consecutive layers (e.g. a depth-2 pipeline: start layer N+1's read/generation while layer N's resolution/write is still running) — a real restructuring, not a parameter tweak, and explicitly not to be combined with the reverted budget-capping idea above (let each phase use however many cores it naturally wants; only the scheduling needs to overlap). Deferred, not started.

Cross-partition batch resolution — current state vs. the batched-accumulator design (discussion, 2026-08-14)

family_scan.rs::scan_layer_families (shared by snp_pseudo_alignment, sibling_annex_stats, cardinality_tally, scan_family_pairs) already implements most of a dispatch/accumulate/resolve pipeline: generation (cheap, CPU-only — builds outgoing[dest_partition] from FamilyMask and buckets cross-partition queries) runs on an obipipeline::throttle + make_pipe! stage, decoupled from resolution (I/O-bound, rayon::par_iter across partitions, one generated batch resolved at a time, never several concurrently — this ordering is deliberate, see the module's own docs on a reverted concurrent-batch-resolution attempt that scattered mmap access). The fast/slow mode gate (PartitionCache::fast_mode, cache.rs:162-163) already exists: n_layers <= 7 (checked once from the first non-empty partition's PartitionMeta::n_layers, documented as identical across every partition of an index — a structural, build-time property, never a per-partition state) decides whether FamilyMask's recorded layer_value can be trusted to skip straight to the right layer (find_presence_batch_fast) or must fall back to scanning every layer of the destination partition (find_presence_batch).

Real gap, confirmed not implemented: resolution is triggered by the source batch finishing (FAMILY_BATCH = 65536 minorants read from the scanned layer), not by an output accumulator filling up. Since most central-base variants of a family route back to the same partition being scanned (~48% per the k=31/m=11 measurement above), a FAMILY_BATCH's outgoing[dest] is large for the local/self partition and thin for the other ~255 (or however many) destination partitions — each of those gets resolved at low query density every batch instead of being accumulated across several source batches until resolving it is worthwhile. This is distinct from, and not fixed by, the fast/slow layer gate above.

Redesign sketched (not built): per-destination accumulators decoupled from FAMILY_BATCH, flushed on reaching a size threshold instead of on source-batch completion — a "hot" accumulator for the partition being scanned (sharded one-per-generation-worker, no lock, since all n_workers pipeline workers write to it concurrently — this differs from an earlier, simpler mental model of "one thread owns one layer's local collector," which doesn't hold here since n_workers threads cooperate on scanning one layer at a time, not one thread per layer) and "cold" mutex-per-partition accumulators for the rest, low contention expected since traffic to any single cold destination is a small fraction of total.

This breaks the current strict-iteration-order delivery of on_family (today: a reorder buffer keyed by batch, since a whole FAMILY_BATCH resolves atomically). With cross-batch accumulation, a family only becomes complete once every accumulator holding one of its outgoing queries has flushed, at unpredictable, independent times — no longer streamable strictly in order without a large, unbounded pending buffer. Resolution sketched: replace order-dependent consumers with coordinate-addressed writes instead of order-dependent appends (see PseudoAlignment idea below) wherever possible, since sibling_annex_stats's reduction (plain counts) is already order-independent and needs nothing here.

Superseded 2026-08-15 by the --subsample/--shannon design below, which sidesteps the accumulator redesign for now: bounding the number of families actually resolved per layer (via sampling) keeps per-layer resolution volume small enough that the batch-density problem above stops mattering in practice for these two consumers. The accumulator redesign remains relevant for a future unsampled, full-index run, but is not required to ship --subsample/--shannon.

Pseudo-alignment at scale — pruning is unavoidable (discussion, 2026-08-14/15)

The reference run (phyloskims_sal_vac-scale bacterial test set, iqtree.fasta) produced a dense alignment for 13 genomes × 383,965 sites (4.8 MB) — trivially small. The in-progress plant index is expected to carry on the order of 9 billion minorant families; a dense byte-per-cell alignment at that column count is unbuildable regardless of genome count (hundreds of GB even at a handful of genomes). Long-term ambition is 6,000 8,000 genomes on a large machine, which makes the per-cell cost dominant in the other dimension too. Pruning the retained family set before materializing anything is mandatory, not an optimization.

Already free: family_size() < 2 (no sibling variant registered at all) is a zero-cost structural filter, read directly off FamilyMask bits, already applied in snp_pseudo_alignment. Insufficient alone — per the ~80-85% mono-family estimate from earlier discussion, this only brings 9 billion down to roughly 1.3-1.8 billion, still unusable.

Entropy definition — settled 2026-08-15, correcting an earlier wrong turn. The project does not encode families as IUPAC ambiguity codes interpreted the classical way (Fitch-parsimony subset-compatibility, or ML's "one true state, uncertain which"); see docmd/theory/evolutionary_distances.md ("Why the IUPAC/DNA encoding used for the first --snp test was wrong") and the Sankoff resolution that followed it. The real model is a genuine 16-state alphabet (the powerset of {A,C,G,T}, included as a real state) scored with a calibrated pairwise cost matrix (obikphylo::cardcomp::pairwise_cost_matrix, cmd/phylo/sankoff.rs), not a compatibility/subset relation between states. Under that model, each of the 16 states — including multi-bit ones like AC — is a first-class, independently-costed state, not an uncertainty encoding of a single true base. So: entropy over the 15 non-empty states ( excluded, matching the earlier decision to exclude genomes where the family is absent) is the correct informativeness measure for this project — not a 4-symbol reduction, which would discard exactly the cardinality/composition information the calibrated cost matrix is built to exploit.

--subsample / --shannon — sampling strategy (decided 2026-08-15)

Goal: make both the pseudo-alignment (--snp) and a Shannon-entropy diagnostic usable at any index scale, from the 13-genome bacterial reference run up to the 9-billion-family plant index, without requiring the batched-accumulator redesign above.

--subsample N (integer, families to retain): bounds the pseudo-alignment to N minorant families, sampled proportionally per layer among non-monomorphic minorants (family_size >= 2) — this sidesteps the need for a true global reservoir merge across layers while still approximating a uniform sample over the whole index, and directly answers the earlier open question of global-vs-per-layer selection scope.

Three passes, in order:

  1. Global count (cheap, structural, parallel across layers — same shape as the existing sibling_family_size_histogram, extended to report a per-layer breakdown rather than one index-wide aggregate): for each layer, count_layer = number of non-monomorphic minorants. Gives total_count = Σ count_layer.
  2. Per-layer proportional reservoir sampling (cheap, structural, one pass per layer, no cross-partition resolution): N_layer = round(N × count_layer / total_count). Since N_layer is a proportion of count_layer, it can never exceed it as long as N <= total_count — the one edge case is total_count <= N, in which case sampling is skipped entirely and every non-monomorphic minorant of every layer is kept (no reservoir needed, N was never a real constraint). Otherwise: Algorithm-R reservoir sampling over the layer's non-monomorphic minorant indices, producing N_layer iteration-order indices directly, no intermediate full list ever materialized.
  3. Filtered resolution (the expensive step, the existing scan_layer_families engine, unchanged): re-scan the layer, generating and resolving cross-partition queries only for the indices selected in step 2 (cheap membership test against a small per-layer index set) — this is what keeps --subsample cheap even on an unsampled-scale index, since the cross-partition resolution volume is bounded by N, not by the layer's true size.

Steps 2 and 3 cannot be merged into one pass: true single-pass reservoir sampling would waste step-3's expensive resolution work on candidates later evicted by the reservoir. Step 1 must fully complete (every layer) before step 2 can start for any layer, since total_count is a global quantity.

--shannon (no argument): emits a CSV of per-family Shannon entropy (15 non-empty states, /absent genomes excluded from the denominator, per the settled definition above). Independent of --subsample — entropy is computed and written per family as soon as its genome_mask resolves, O(1) memory per family, so it streams fine even unsampled at full index scale (a time cost, not a memory one). Combined with --subsample N, it delivers the original exploratory diagnostic (e.g. --subsample 1000000 --shannon) directly from this general machinery, rather than a purpose-built one-off script.

Validated end-to-end (2026-08-15) against real data: --sibling-hist on phyloskims_sal_vac (91 real genomes, k=31/m=11, 256 partitions × 2 layers) confirms the ~9-billion-family estimate almost exactly (8,925,068,238 total, 97.9% monomorphic — a sharper mono fraction than the ~80-85% earlier guess, corrected here). --subsample/--shannon on the smaller 20-genome bacterial reference (benchmark/global_index_presence) produced a sample size within rounding of the request (99,742/100,000) and a [0.8,1.2)-bucket share (40.7%) matching the full unsampled population (41.6%) — the two histograms only diverged wildly (5‰ vs 41.6%) under a real bug in reservoir_sample_layer (see next section), now fixed.

Bug found and fixed (2026-08-15): family_idx numbering mismatch. scan_layer_families's family_idx counts every minorant of a layer (monomorphic ones included, since iter_minorants_batch filters only on is_minorant()), not the raw annex slot (SiblingAnnex::get(slot) spans every k-mer, minorant or not) and not a counter over non-monomorphic minorants alone. subsample.rs's reservoir_sample_layer originally stored raw slot numbers in its HashSet<usize> selection, which drifts away from family_idx as soon as any monomorphic minorant is seen — i.e. almost immediately, since ~98% of minorants are monomorphic. Fixed by tracking two separate counters: family_idx (every minorant, matching scan_layer_families) and seen (non-monomorphic minorants only, what Algorithm R actually samples over) — only family_idx values are ever stored in the selection set. The existing unit test never exercised this (its fixture has exactly one non-monomorphic family, always hitting the "keep everything" shortcut) — a stronger fixture with several interleaved monomorphic/non-monomorphic families would be needed to catch a regression here automatically; not yet written.

Cheap entropy pre-filtering — row-marginal sums (idea, not implemented, 2026-08-15)

Motivation: on real data (bacterial reference, full unsampled run), only ~5‰ of non-monomorphic minorants fall in the [0.5, 1.5] bit band judged phylogenetically interesting (entropy too low = uninformative near-invariant site; too high = saturated/noisy, see family_entropy's 15-state discussion) — roughly 1 in 10,000 minorants overall. Computing exact entropy for every candidate just to discard 99.99% of them is wasteful at the full 9-billion scale.

The idea: PersistentBitMatrix::col_view(c) gives a genome's whole presence column as a BitSliceView (sequential, no MPHF, no cross-partition routing — purely local to one layer's own matrix). Accumulating TempCompactIntVecBuilder::inc_present(col) (already exists in obicompactvec/src/builder.rs:121, along with add/min/max/diff on IntSliceView — no new low-level API needed) over every column of a layer produces coverage[slot]: how many genomes carry each exact k-mer, in one sequential per-layer pass, entirely decoupled from family/sibling structure. Persisted once per layer, a family's members' coverage could then be looked up via a plain local MPHF index() (cheap) instead of a full cross-partition presence resolution (find_presence_batch) — the expensive part today is specifically the cross-partition/cross-layer routing to a sibling's own matrix, not the bit-reading itself, and coverage[slot] sidesteps that routing entirely by moving the cost into a one-time, purely local, embarrassingly-parallel build step.

Why not implemented: coverage[slot] is a per-member marginal — summing members' coverages to approximate a family's entropy silently assumes no genome carries more than one member at once. It cannot represent or detect joint co-occurrence (a genome carrying both A and C at once, i.e. a combined 15-symbol state) at all, which is exactly the phenomenon family_entropy's 15-state definition exists to capture (see CardinalityTally/cardinality_transition_probs, the project's own existing machinery for this same co-occurrence structure, built for the Sankoff matrix calibration). A family that is in truth uniformly AC across every carrying genome would look like a well-balanced 2-state split under the marginal approximation (entropy ≈ 1) while its true 15-state entropy is 0 — i.e. the marginal proxy's failure mode lands on exactly the "saturated, uninformative" tail this pre-filter would need to catch, undermining the point. It stays plausible as a coarse filter for the low tail only (a dominant single member's marginal share reliably predicts low true entropy too), but not as a stand-in for the high tail — not pursued further for now.

Two entropy definitions kept side by side, for comparison (2026-08-15)

--shannon's CSV carries both entropy15 (family_entropy — the settled 15-non-empty-state definition, see above) and entropy4 (family_entropy_4 — plain nucleotide reduction), computed from the same already-resolved genome_mask, not from the marginal approximation above. A genome carrying several bases at once contributes to each base's count (counted once per base present, not fractionally split, not folded into one combined state) — a genome polymorphic for the family is present at more than one base by construction, so it is expected to count more than once; the denominator is the total base-occurrence count, not the genome count (the two coincide only when no genome carries more than one base). Kept side by side specifically to measure, on real data, how much the two diverge — not yet analyzed.