From 45b19503a1114ca72a58bbbeccf2144861a0accf Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Sat, 15 Aug 2026 09:50:02 +0200 Subject: [PATCH] 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. --- .DS_Store | Bin 14340 -> 14340 bytes .gitignore | 1 + docmd/architecture/siblings.md | 226 ++++++++++++++++---- src/Cargo.lock | 1 + src/obikmer/src/cmd/phylo/args.rs | 24 ++- src/obikmer/src/cmd/phylo/family_overlap.rs | 2 +- src/obikmer/src/cmd/phylo/mod.rs | 23 +- src/obikphylo/Cargo.toml | 1 + src/obikphylo/src/siblings/alignment.rs | 22 +- src/obikphylo/src/siblings/cardinality.rs | 4 +- src/obikphylo/src/siblings/distance.rs | 4 +- src/obikphylo/src/siblings/entropy.rs | 170 +++++++++++++++ src/obikphylo/src/siblings/family_scan.rs | 62 +++++- src/obikphylo/src/siblings/mod.rs | 3 + src/obikphylo/src/siblings/stats.rs | 4 +- src/obikphylo/src/siblings/subsample.rs | 158 ++++++++++++++ src/obikphylo/src/siblings/tests.rs | 206 +++++++++++++++++- 17 files changed, 841 insertions(+), 70 deletions(-) create mode 100644 src/obikphylo/src/siblings/entropy.rs create mode 100644 src/obikphylo/src/siblings/subsample.rs diff --git a/.DS_Store b/.DS_Store index 34be0f9da62fb45f2230c2abb804bbf385535890..4373600fa9ad52b079125e98c415ac6ae5447bf4 100644 GIT binary patch delta 385 zcmZoEXepQ=#l+AyQN~zQn1O+TnSnuofuV#UkwKTC7)TdQEL5IsV8X)62$W=*?5iR* z*+HF$kByz3lbxGAHaH`{Jh)`Dqtq$ZjSZEIoFG*|z`+d^5@*@0&hvzsk#n+-`lZbR zEK$sx`2{X>3vvP#urVYtWCC#xkj$SfE8HMrVq&4AU}S1stD{hDX=I?IU}9`Gd84p# zJ&Mwtbi?4}{M-VtEf4}lJ~!XRB`GIA3FL+i4q?i!5+@wdt;@|%A*iMx17?pB+$t1{ z%7TmXa`N-ifwGLCP+6>|z&?2ai}U6kB5LfD*$sr*7}^*Z6nHiZa=d4r%x|I#4w?fX T!zSDCNKcM2nY+1FnVlH`3?NgR delta 192 zcmZoEXepQ=#l*rjQN~!5mw|zSnSnuofgzb8mmwvmxXd{zXXC}i%##gFSXdc>vMihL z2}ChZexx#s2gn17009U0#*5-Co7H)qFf+1mHWj(Xz43u23nS~~L?x-k^Ay=Z^u$#> zlX(O#P5vt)wK-8{74yaisw|V)4TRZPxEL4|cs5>q&pesmL{|o69>~}R2o1IYNH77p S43lF_7H%$7(`7;!xB~#j^)4C! diff --git a/.gitignore b/.gitignore index addd10f0..bfe28adc 100644 --- a/.gitignore +++ b/.gitignore @@ -26,6 +26,7 @@ benchmark/specific_index_count benchmark/specific_index_presence TNT phyg +biblio *.tnt *.tre *.phy diff --git a/docmd/architecture/siblings.md b/docmd/architecture/siblings.md index 785f18a7..3677cd9f 100644 --- a/docmd/architecture/siblings.md +++ b/docmd/architecture/siblings.md @@ -148,6 +148,22 @@ Baseline: mostly one active core, with short multi-core bursts — average 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`, @@ -201,7 +217,15 @@ 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. -## Pseudo-alignment at scale — pruning is unavoidable (discussion, 2026-08-14) +**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 @@ -219,49 +243,163 @@ 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. -**Criterion under discussion**: fix a global cell budget -(`n_genomes × n_columns_retained ≤ threshold`) and retain the -highest-entropy families first until the budget is spent — column count -self-adjusts to genome count automatically. Entropy of a family: Shannon -entropy over the observed base distribution across genomes carrying that -family, **decided 2026-08-14: genomes where the family is absent are -excluded from the calculation** (denominator = genomes where present, not -all genomes) — measures signal purity where it exists, independent of -coverage rate. +**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. -Key consequence for storage: entropy needs per-variant genome counts, which -cannot be derived from `FamilyMask`'s bits alone — it requires the same -cross-partition resolution work the alignment/stats pipeline already pays -for. This favors computing it once at annex-build time and persisting it in -an **auxiliary vector alongside the annex** (one value per retained -minorant) over encoding a fixed keep/discard decision as one of -`FamilyMask`'s 3 remaining free bits (bits 13-15, `siblingannex.rs:77`): the -bit approach bakes a single threshold in permanently (any different budget -needs a full annex rebuild), the vector approach pays the expensive -resolution once and lets budget/threshold be chosen freely per analysis -afterward. +## `--subsample` / `--shannon` — sampling strategy (decided 2026-08-15) -**Open, explicitly deferred**: scope of the top-K selection — global across -the whole index (needs a first pass computing/storing every family's -entropy, then a global threshold from the full distribution, e.g. a -quantile, before building the final alignment: two full passes, but the -size budget is honored precisely) vs. local per layer/partition (stays -within the current single-pass streaming model, but the global budget is -no longer exactly guaranteed — depends on how unevenly entropy is -distributed across layers). +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. -**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. +**`--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` 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. diff --git a/src/Cargo.lock b/src/Cargo.lock index 1cafb680..113db77b 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1858,6 +1858,7 @@ dependencies = [ "obiskbuilder", "obiskio", "obisys", + "rand 0.8.6", "rayon", "tempfile", "tracing", diff --git a/src/obikmer/src/cmd/phylo/args.rs b/src/obikmer/src/cmd/phylo/args.rs index aa294e6e..b166ab25 100644 --- a/src/obikmer/src/cmd/phylo/args.rs +++ b/src/obikmer/src/cmd/phylo/args.rs @@ -146,6 +146,28 @@ pub struct PhyloArgs { #[arg(long)] pub snp: bool, + /// Cap the number of variable families (non-monomorphic minorants, + /// `family_size() >= 2`) retained by `--snp`/`--sankoff` (and everything + /// `--sankoff` implies) and `--shannon`, to (approximately) this many — + /// sampled proportionally per layer, so the pseudo-alignment/entropy + /// report stays usable on an index far larger than the sample itself + /// (mandatory, not optional, once the index is large enough that a full + /// pseudo-alignment can't be materialized at all). See + /// `docmd/architecture/siblings.md`, "`--subsample`/`--shannon`". If the + /// index has fewer non-monomorphic minorants than this, every one of + /// them is kept — no error, no under/over-shoot handling needed. + #[arg(long, value_name = "N")] + pub subsample: Option, + + /// Write _shannon.csv: per-family Shannon entropy (bits, over + /// the 15 non-empty subsets of `{A,C,G,T}`, `∅`/absent genomes excluded + /// from the denominator — see `docmd/architecture/siblings.md`, + /// "Entropy definition") of every non-monomorphic minorant, one row per + /// family. Combine with `--subsample N` for a bounded diagnostic sample + /// instead of a full-index pass. Same annex requirement as `--snp`. + #[arg(long)] + pub shannon: bool, + /// Write an NxN CSV (`_family_overlap.csv`) of, for each genome /// pair, how many variable families (same set `--snp`'s pseudo-alignment /// uses — `family_size() >= 2`) both genomes actually carry a call for @@ -244,7 +266,7 @@ pub struct PhyloArgs { /// Output prefix: _dist.csv, _shared.csv, /// _siblings.csv, _sibling_hist.csv, _rawsnp.csv, /// _rawsnp_counts.csv, - /// _snp.fasta, _family_overlap.csv, + /// _snp.fasta, _family_overlap.csv, _shannon.csv, /// _sankoff_matrix.csv, _sankoff_params.yaml, /// _sankoff.fasta, _sankoff.tnt, _sankoff.tcm, /// _sankoff.pg, _iqtree.model, _iqtree.fasta, diff --git a/src/obikmer/src/cmd/phylo/family_overlap.rs b/src/obikmer/src/cmd/phylo/family_overlap.rs index 7a8f9d7b..00732aed 100644 --- a/src/obikmer/src/cmd/phylo/family_overlap.rs +++ b/src/obikmer/src/cmd/phylo/family_overlap.rs @@ -80,7 +80,7 @@ pub(super) fn apply_min_shared_family_exclusion( threshold: f64, mask: &mut [bool], ) { - let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { + let alignment = idx.snp_pseudo_alignment(None).unwrap_or_else(|e| { eprintln!("error computing SNP pseudo-alignment for --min-shared-family: {e}"); std::process::exit(1); }); diff --git a/src/obikmer/src/cmd/phylo/mod.rs b/src/obikmer/src/cmd/phylo/mod.rs index e7dd41f0..8f7715d1 100644 --- a/src/obikmer/src/cmd/phylo/mod.rs +++ b/src/obikmer/src/cmd/phylo/mod.rs @@ -13,8 +13,8 @@ use obikindex::KmerIndex; use obikphylo::{ cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix, siblings::{ - CardinalityExt, DistanceExt, RawSnpDistanceOutput, SiblingAnnexBuildExt, SiblingStatsExt, - SnpAlignment, SnpAlignmentExt, + CardinalityExt, DistanceExt, RawSnpDistanceOutput, ShannonEntropyExt, SiblingAnnexBuildExt, + SiblingStatsExt, SnpAlignment, SnpAlignmentExt, }, }; use obisys::{Reporter, Stage}; @@ -177,7 +177,7 @@ pub fn run(args: PhyloArgs) { } if args.snp { let t = Stage::start("snp_pseudo_alignment"); - let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { + let alignment = idx.snp_pseudo_alignment(args.subsample).unwrap_or_else(|e| { eprintln!("error computing SNP pseudo-alignment: {e}"); std::process::exit(1); }); @@ -187,7 +187,7 @@ pub fn run(args: PhyloArgs) { } if args.family_overlap { let t = Stage::start("snp_pseudo_alignment"); - let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { + let alignment = idx.snp_pseudo_alignment(args.subsample).unwrap_or_else(|e| { eprintln!("error computing SNP pseudo-alignment: {e}"); std::process::exit(1); }); @@ -195,6 +195,18 @@ pub fn run(args: PhyloArgs) { let (alignment, kept_labels) = drop_excluded(alignment); write_family_overlap_csv(&alignment, &kept_labels, &args.output); } + if args.shannon { + let t = Stage::start("shannon_entropy"); + let path = args.output.as_ref() + .map(|p| format!("{}_shannon.csv", p.display())) + .unwrap_or_else(|| "shannon.csv".into()); + idx.shannon_entropy_csv(std::path::Path::new(&path), args.subsample).unwrap_or_else(|e| { + eprintln!("error computing Shannon entropy: {e}"); + std::process::exit(1); + }); + rep.push(t.stop()); + info!("per-family Shannon entropy → {path}"); + } if args.sankoff || args.tnt || args.phyg || args.iqtree { let t = Stage::start("raw_snp_distance"); let mut raw = idx.raw_snp_distance().unwrap_or_else(|e| { @@ -223,7 +235,7 @@ pub fn run(args: PhyloArgs) { write_sankoff_params(&card_tally, &p_card, &base_tally, &p_comp, args.sankoff_ratio_ceiling, &args.output); let t = Stage::start("snp_pseudo_alignment"); - let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { + let alignment = idx.snp_pseudo_alignment(args.subsample).unwrap_or_else(|e| { eprintln!("error computing SNP pseudo-alignment: {e}"); std::process::exit(1); }); @@ -256,6 +268,7 @@ pub fn run(args: PhyloArgs) { || args.raw_snp_counts || args.snp || args.family_overlap + || args.shannon || args.sankoff || args.tnt || args.phyg diff --git a/src/obikphylo/Cargo.toml b/src/obikphylo/Cargo.toml index 5dc257ba..282a852f 100644 --- a/src/obikphylo/Cargo.toml +++ b/src/obikphylo/Cargo.toml @@ -15,6 +15,7 @@ obiskbuilder = { path = "../obiskbuilder" } obipipeline = { path = "../obipipeline" } memmap2 = "0.9" ndarray = "0.16" +rand = "0.8" rayon = "1" tracing = "0.1.44" diff --git a/src/obikphylo/src/siblings/alignment.rs b/src/obikphylo/src/siblings/alignment.rs index 123b14d3..bd503f0c 100644 --- a/src/obikphylo/src/siblings/alignment.rs +++ b/src/obikphylo/src/siblings/alignment.rs @@ -7,7 +7,7 @@ use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; use super::cache::PartitionCache; -use super::family_scan::scan_layer_families; +use super::family_scan::{scan_layer_families, Selection}; /// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set /// iff the genome carries the member whose own central base is `b`): @@ -59,11 +59,18 @@ pub struct SnpAlignment { pub trait SnpAlignmentExt { /// Build the SNP-only pseudo-alignment from an already-built sibling /// annex (run [`build_sibling_annex`](super::build::SiblingAnnexBuildExt::build_sibling_annex) first). - fn snp_pseudo_alignment(&self) -> OKIResult; + /// + /// `subsample`: `--subsample N` — cap the number of variable families + /// retained to (approximately) `N`, sampled proportionally per layer + /// among non-monomorphic minorants (see + /// `docmd/architecture/siblings.md`, "`--subsample`/`--shannon`"). + /// `None` keeps today's behaviour: every non-monomorphic minorant of + /// every layer. + fn snp_pseudo_alignment(&self, subsample: Option) -> OKIResult; } impl SnpAlignmentExt for KmerIndex { - fn snp_pseudo_alignment(&self) -> OKIResult { + fn snp_pseudo_alignment(&self, subsample: Option) -> OKIResult { let n_parts = self.n_partitions(); let n_genomes = self.meta().genomes.len(); let with_counts = self.meta().config.with_counts; @@ -79,6 +86,7 @@ impl SnpAlignmentExt for KmerIndex { .map_err(OKIError::Partition)?; let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); let layer_dirs = super::family_scan::sibling_layer_dirs(self)?; + let selections = super::subsample::compute_selections(&layer_dirs, subsample)?; let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers"); // One layer at a time, not `par_iter()` over layers — same @@ -92,8 +100,12 @@ impl SnpAlignmentExt for KmerIndex { // whole-layer) — no intermediate per-layer buffer, layer order // giving a single deterministic column order across the whole index. let mut sequences: Vec> = vec![Vec::new(); n_genomes]; - for layer_dir in &layer_dirs { - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| { + for (layer_dir, layer_selection) in layer_dirs.iter().zip(selections.iter()) { + let selection = match layer_selection { + None => Selection::All, + Some(set) => Selection::Some(set), + }; + scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |_family_idx, mask, genome_mask| { if mask.family_size() < 2 { return; // monomorphic family — no signal, skip } diff --git a/src/obikphylo/src/siblings/cardinality.rs b/src/obikphylo/src/siblings/cardinality.rs index c76888fd..0e63910c 100644 --- a/src/obikphylo/src/siblings/cardinality.rs +++ b/src/obikphylo/src/siblings/cardinality.rs @@ -10,7 +10,7 @@ use obikindex::KmerIndex; use super::cache::PartitionCache; use super::distance::RawSnpDistanceOutput; -use super::family_scan::scan_layer_families; +use super::family_scan::{scan_layer_families, Selection}; /// See [`KmerIndex::cardinality_tally`]. pub struct CardinalityTally { @@ -89,7 +89,7 @@ impl CardinalityExt for KmerIndex { // family comes back — no per-layer buffer. let mut total = [[0u64; 5]; 5]; for layer_dir in &layer_dirs { - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| { + scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |_family_idx, mask, genome_mask| { if mask.family_size() < 2 { // Fully invariant family (never varies anywhere in // the index) — genome-wide background, not diff --git a/src/obikphylo/src/siblings/distance.rs b/src/obikphylo/src/siblings/distance.rs index 99347600..8ce640a9 100644 --- a/src/obikphylo/src/siblings/distance.rs +++ b/src/obikphylo/src/siblings/distance.rs @@ -9,7 +9,7 @@ use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; use super::cache::PartitionCache; -use super::family_scan::scan_layer_families; +use super::family_scan::{scan_layer_families, Selection}; /// Raw p-distance restricted to loci that are single-copy in **both** /// genomes of a pair — the "stringent / paralogy-aware" locus eligibility @@ -90,7 +90,7 @@ where for layer_dir in &layer_dirs { let mut acc = zero(); - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| { + scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |_family_idx, mask, genome_mask| { let variable = mask.family_size() >= 2; // Per genome: which single form (if exactly one) it diff --git a/src/obikphylo/src/siblings/entropy.rs b/src/obikphylo/src/siblings/entropy.rs new file mode 100644 index 00000000..0b73b5c5 --- /dev/null +++ b/src/obikphylo/src/siblings/entropy.rs @@ -0,0 +1,170 @@ +//! `--shannon` — per-family Shannon entropy over the 15 non-empty subsets +//! of `{A,C,G,T}`. See `docmd/architecture/siblings.md`, "Entropy +//! definition — settled 2026-08-15" and "`--subsample`/`--shannon`", for +//! the design discussion this implements: the project's calibrated 16-state +//! Sankoff cost matrix (`cardcomp::pairwise_cost_matrix`) treats every +//! observed combination of bases as its own first-class state, so entropy +//! over that same 15-symbol space (state `0`/`∅` excluded — a family is +//! only informative among the genomes where it is actually observed) is the +//! consistent informativeness proxy for this project, not a 4-symbol +//! reduction. + +use std::io::{BufWriter, Write}; +use std::path::Path; +use std::sync::Arc; + +use obikpartitionner::KmerPartition; +use obisys::progress_bar; + +use obikindex::{OKIError, OKIResult}; +use obikindex::KmerIndex; + +use super::cache::PartitionCache; +use super::family_scan::{Selection, scan_layer_families}; + +/// Shannon entropy (bits) of one family's per-genome states, over the 15 +/// non-empty subsets of `{A,C,G,T}` actually observed among the genomes +/// carrying it (`genome_mask[g] != 0`) — genomes where the family is absent +/// are excluded from the denominator, not scored as a 16th state. `None` +/// only if the family is absent from every genome (shouldn't happen for a +/// real minorant, guarded against rather than dividing by zero). Second +/// value: the denominator (`n_genomes_present`) used, for traceability. +pub(super) fn family_entropy(genome_mask: &[u8]) -> Option<(f64, usize)> { + let mut freq = [0u32; 16]; // index 0 (∅) never incremented, kept for direct `m as usize` indexing + let mut present = 0u32; + for &m in genome_mask { + if m != 0 { + freq[m as usize] += 1; + present += 1; + } + } + if present == 0 { + return None; + } + let mut h = 0.0f64; + for &c in &freq[1..] { + if c == 0 { + continue; + } + let p = f64::from(c) / f64::from(present); + h -= p * p.log2(); + } + Some((h, present as usize)) +} + +/// Shannon entropy (bits) of one family's per-genome states, reduced to the +/// 4 plain nucleotide symbols instead of the 15-state space above — kept +/// alongside it (not instead of it) purely to measure how much the two +/// diverge on real data, per the discussion in +/// `docmd/architecture/siblings.md`, "`--subsample`/`--shannon`". A genome +/// carrying several bases at once (`genome_mask[g]` with more than one bit +/// set) contributes to *each* of those bases' counts — counted once per +/// base present, not split fractionally, and not folded into a combined +/// state the way [`family_entropy`] does: a genome polymorphic for the +/// family is, by construction, present at more than one base, so it simply +/// counts more than once here. The denominator is therefore the total +/// number of base-occurrences summed over every present genome, not the +/// genome count — the two coincide only when no genome carries more than +/// one base. `None` under the same condition as [`family_entropy`]. +pub(super) fn family_entropy_4(genome_mask: &[u8]) -> Option<(f64, usize)> { + let mut freq = [0u32; 4]; + let mut total = 0u32; + for &m in genome_mask { + for (base, count) in freq.iter_mut().enumerate() { + if m & (1 << base) != 0 { + *count += 1; + total += 1; + } + } + } + if total == 0 { + return None; + } + let mut h = 0.0f64; + for &c in &freq { + if c == 0 { + continue; + } + let p = f64::from(c) / f64::from(total); + h -= p * p.log2(); + } + Some((h, total as usize)) +} + +/// Adds [`shannon_entropy_csv`](Self::shannon_entropy_csv) to `KmerIndex`. +pub trait ShannonEntropyExt { + /// Writes a CSV + /// (`layer,family_idx,entropy15,entropy4,family_size,n_genomes_present`) + /// of per-family Shannon entropy — both [`family_entropy`] (15 non-empty + /// states, the project's settled definition) and [`family_entropy_4`] + /// (plain 4-symbol reduction, kept alongside for comparison) — one row + /// per non-monomorphic minorant visited, from an already-built sibling + /// annex (run + /// [`build_sibling_annex`](super::build::SiblingAnnexBuildExt::build_sibling_annex) + /// first). `layer` is this layer's ordinal among every layer carrying an + /// annex (not a partition/layer pair — an opaque but stable id, unique + /// together with `family_idx`); `family_idx` is the family's + /// iteration-order index within that layer, the same numbering + /// `--subsample`'s selection is drawn from. + /// + /// `subsample`: same meaning as + /// [`SnpAlignmentExt::snp_pseudo_alignment`](super::SnpAlignmentExt::snp_pseudo_alignment)'s + /// — `None` streams every non-monomorphic minorant of the whole index + /// (a time cost, not a memory one: entropy is computed and written + /// per-family as soon as it resolves), `Some(n)` bounds the index's + /// cross-partition resolution work to ~`n` sampled families. + fn shannon_entropy_csv(&self, path: &Path, subsample: Option) -> OKIResult<()>; +} + +impl ShannonEntropyExt for KmerIndex { + fn shannon_entropy_csv(&self, path: &Path, subsample: Option) -> OKIResult<()> { + let n_parts = self.n_partitions(); + let n_genomes = self.meta().genomes.len(); + let with_counts = self.meta().config.with_counts; + let k = self.kmer_size(); + let n_bits = n_parts.trailing_zeros() as usize; + + let partition = KmerPartition::open_with_config( + self.root_path(), + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); + let layer_dirs = super::family_scan::sibling_layer_dirs(self)?; + let selections = super::subsample::compute_selections(&layer_dirs, subsample)?; + + let mut f = BufWriter::new(std::fs::File::create(path).map_err(OKIError::Io)?); + writeln!(f, "layer,family_idx,entropy15,entropy4,family_size,n_genomes_present").map_err(OKIError::Io)?; + + let pb = progress_bar("shannon_entropy", layer_dirs.len() as u64, "layers"); + for (layer_ord, (layer_dir, layer_selection)) in layer_dirs.iter().zip(selections.iter()).enumerate() { + let selection = match layer_selection { + None => Selection::All, + Some(set) => Selection::Some(set), + }; + let mut write_err = None; + scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |family_idx, mask, genome_mask| { + if write_err.is_some() { + return; + } + if mask.family_size() < 2 { + return; // monomorphic family — entropy trivially 0, no signal, skip + } + let Some((h15, n_present)) = family_entropy(genome_mask) else { return }; + let h4 = family_entropy_4(genome_mask).map_or(0.0, |(h, _)| h); + if let Err(e) = writeln!(f, "{layer_ord},{family_idx},{h15:.6},{h4:.6},{},{n_present}", mask.family_size()) { + write_err = Some(e); + } + })?; + if let Some(e) = write_err { + return Err(OKIError::Io(e)); + } + pb.inc(1); + } + pb.finish_and_clear(); + + Ok(()) + } +} diff --git a/src/obikphylo/src/siblings/family_scan.rs b/src/obikphylo/src/siblings/family_scan.rs index 46e5f76d..338a1b8c 100644 --- a/src/obikphylo/src/siblings/family_scan.rs +++ b/src/obikphylo/src/siblings/family_scan.rs @@ -74,6 +74,29 @@ use super::{olm_to_ok, FamilyMask, SiblingAnnex, ANNEX_FILE_NAME, INDEX_SUBDIR}; /// queries per batch to amortise against. const FAMILY_BATCH: usize = 65536; +/// Restricts [`scan_layer_families`] to a subset of a layer's families, +/// identified by their iteration-order index — see +/// `docmd/architecture/siblings.md`, "`--subsample`/`--shannon`", for why: +/// bounding which families actually pay the expensive cross-partition +/// resolution cost is what makes a sampled run cheap on an index far larger +/// than the sample itself. `All` (the default for every pre-existing +/// caller) keeps today's behaviour exactly. +pub(super) enum Selection<'a> { + All, + Some(&'a std::collections::HashSet), +} + +/// Shared by the pipeline worker (via the `Arc`-cloned `Option`, +/// see below — a `Selection` itself doesn't cross the worker thread +/// boundary) and the final replay loop, so both sides agree on membership +/// without duplicating the match. +fn is_selected(selected: &Option>, family_idx: usize) -> bool { + match selected { + None => true, + Some(set) => set.contains(&family_idx), + } +} + /// Every (partition, layer) directory carrying a sibling annex, checked /// up front so a missing one is reported before any real work starts. /// Shared by every sibling-annex consumer's extension-trait impl (`alignment`, @@ -152,11 +175,12 @@ enum FamData { /// Visits every minorant family of one layer, in iteration order, batched /// [`FAMILY_BATCH`] at a time — see the module docs for why generation and /// resolution use different concurrency. `on_family` is called once per -/// family, in iteration order, with its own annex mask and its per-genome -/// base-presence (`genome_mask[g]`: bit `b` set iff genome `g` carries the -/// member whose own canonical central base is `b`), backed by a scratch -/// buffer reused across every call — callers that need to keep data past -/// the call must copy it themselves. +/// family, in iteration order, with that family's iteration-order index +/// (the same numbering [`Selection`] indices are drawn from), its own annex +/// mask, and its per-genome base-presence (`genome_mask[g]`: bit `b` set +/// iff genome `g` carries the member whose own canonical central base is +/// `b`), backed by a scratch buffer reused across every call — callers that +/// need to keep data past the call must copy it themselves. pub(super) fn scan_layer_families( layer_dir: &Path, n_parts: usize, @@ -164,7 +188,8 @@ pub(super) fn scan_layer_families( with_counts: bool, k: usize, cache: &Arc, - mut on_family: impl FnMut(FamilyMask, &[u8]), + selection: &Selection, + mut on_family: impl FnMut(usize, FamilyMask, &[u8]), ) -> OKIResult<()> { 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)?; @@ -203,7 +228,18 @@ pub(super) fn scan_layer_families( _permit: t.guard, }); + // `Selection::Some` borrows a `HashSet` whose lifetime doesn't span the + // pipeline's worker threads — cloned once into an `Arc` so every worker + // can share it cheaply instead of requiring `selection` itself to be + // `'static`. `None` means [`Selection::All`], checked with a plain + // `map_or` at each use instead of allocating an always-true set. + let selected: Arc>> = Arc::new(match selection { + Selection::All => None, + Selection::Some(set) => Some((*set).clone()), + }); + let worker_ctx = Arc::clone(&ctx); + let worker_selected = Arc::clone(&selected); let pipe = obipipeline::make_pipe! { FamData : SourceBatch => GeneratedBatch, | { @@ -222,11 +258,19 @@ pub(super) fn scan_layer_families( // (not a hand-rolled `central_canonical_neighbors` + // `mask.has` loop) already filters to present members and // hands back each one's annex-recorded layer alongside. + // Families outside `selected` (sampling only) still get a + // mask/base entry — pass 2 indexes uniformly by `i` — but + // never an outgoing query: that's the expensive part a + // sampled run exists to avoid paying for every family. for (i, entry) in batch.entries.iter().enumerate() { let (kmer, mask) = (entry.kmer, entry.mask); masks.push(mask); let base = central_base(kmer, ctx.k); bases.push(base); + let family_idx = batch.start_family_idx + i; + if !is_selected(&worker_selected, family_idx) { + continue; + } for (member, layer) in mask.family_members(kmer, ctx.k) { if member == kmer { continue; // local — resolved below straight from `mat`, no lookup @@ -296,10 +340,14 @@ pub(super) fn scan_layer_families( }); for i in 0..n { + let family_idx = batch.start_family_idx + i; + if !is_selected(&selected, family_idx) { + continue; // not in the sample — never resolved above, nothing to report + } for (g, dst) in scratch.iter_mut().enumerate() { *dst = genome_mask[i * n_genomes + g].load(Ordering::Relaxed); } - on_family(batch.masks[i], &scratch); + on_family(family_idx, batch.masks[i], &scratch); } next_expected += n; } diff --git a/src/obikphylo/src/siblings/mod.rs b/src/obikphylo/src/siblings/mod.rs index 9fa04f5e..a5c368bb 100644 --- a/src/obikphylo/src/siblings/mod.rs +++ b/src/obikphylo/src/siblings/mod.rs @@ -47,11 +47,13 @@ mod build; mod cache; mod cardinality; mod distance; +mod entropy; mod family_scan; mod helpers; mod iter; mod siblingannex; mod stats; +mod subsample; #[cfg(test)] mod tests; @@ -60,6 +62,7 @@ pub use alignment::{SnpAlignment, SnpAlignmentExt}; pub use build::SiblingAnnexBuildExt; pub use cardinality::{CardinalityExt, CardinalityTally}; pub use distance::{BasePairTally, DistanceExt, RawSnpDistanceOutput}; +pub use entropy::ShannonEntropyExt; pub use iter::{MinorantBatchIter, MinorantIter, SiblingBatchIter, SiblingEntry, SiblingIter, SiblingLayerExt}; pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder}; pub use stats::{SiblingAnnexStats, SiblingStatsExt}; diff --git a/src/obikphylo/src/siblings/stats.rs b/src/obikphylo/src/siblings/stats.rs index e612257f..06cf4698 100644 --- a/src/obikphylo/src/siblings/stats.rs +++ b/src/obikphylo/src/siblings/stats.rs @@ -11,7 +11,7 @@ use obikindex::KmerIndex; use super::ANNEX_FILE_NAME; use super::SiblingAnnex; use super::cache::PartitionCache; -use super::family_scan::scan_layer_families; +use super::family_scan::{scan_layer_families, Selection}; /// Distribution of family sizes (1-4), read back from an already-built /// annex (see [`super::build::SiblingAnnexBuildExt::build_sibling_annex`]) @@ -131,7 +131,7 @@ impl SiblingStatsExt for KmerIndex { ..Default::default() }; for layer_dir in &layer_dirs { - scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| { + scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |_family_idx, mask, genome_mask| { // "Genome g represents this family" means g carries // *any* of its members, not just the minorant's own — // `genome_mask[g] != 0` is exactly that. diff --git a/src/obikphylo/src/siblings/subsample.rs b/src/obikphylo/src/siblings/subsample.rs new file mode 100644 index 00000000..80b593af --- /dev/null +++ b/src/obikphylo/src/siblings/subsample.rs @@ -0,0 +1,158 @@ +//! Proportional per-layer sampling of non-monomorphic minorants +//! (`family_size() >= 2`) — see `docmd/architecture/siblings.md`, +//! "`--subsample`/`--shannon`", for the full design discussion. +//! +//! Three steps, in order, each cheaper than the one before because it +//! narrows what the next one has to touch: +//! +//! 1. [`non_monomorphic_counts`] — one cheap structural pass per layer +//! (annex bits only, no cross-partition resolution), giving `count_layer` +//! for every layer. +//! 2. [`sample_layers`] — given a target `n` and those counts, computes +//! `n_layer = round(n * count_layer / total_count)` per layer (never +//! exceeds `count_layer` as long as `n <= total_count`; if +//! `total_count <= n`, sampling is skipped entirely — every +//! non-monomorphic minorant of every layer is kept) and reservoir-samples +//! (Algorithm R, one more cheap structural pass, `O(n_layer)` memory) the +//! selected iteration-order indices. +//! 3. The result feeds [`super::family_scan::scan_layer_families`]'s +//! `Selection` parameter, which bounds the expensive cross-partition +//! resolution to just the selected families. + +use std::collections::HashSet; +use std::path::{Path, PathBuf}; + +use rand::Rng; +use rayon::prelude::*; + +use obikindex::OKIResult; + +use super::ANNEX_FILE_NAME; +use super::SiblingAnnex; + +/// A family is non-monomorphic (informative-eligible) iff at least one +/// sibling is registered alongside it — `family_size() >= 2`, i.e. +/// `siblings() >= 1`. Mirrors the filter already applied post-hoc in +/// `snp_pseudo_alignment`, but checked here *before* any resolution work. +fn is_non_monomorphic_minorant(mask: super::FamilyMask) -> bool { + mask.is_minorant() && mask.siblings() >= 1 +} + +/// Step 1: per-layer count of non-monomorphic minorants — cheap, parallel +/// across layers, annex bits only (same shape as +/// `SiblingStatsExt::sibling_family_size_histogram`, but keeping the +/// per-layer breakdown instead of collapsing it into one index-wide sum, +/// since proportional sampling needs each layer's own share). +pub(super) fn non_monomorphic_counts(layer_dirs: &[PathBuf]) -> OKIResult> { + layer_dirs + .par_iter() + .map(|layer_dir| -> OKIResult { + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + let mut count = 0u64; + for slot in 0..annex.len() { + if annex.get(slot).is_some_and(is_non_monomorphic_minorant) { + count += 1; + } + } + Ok(count) + }) + .collect() +} + +/// Step 2, per layer: Algorithm-R reservoir sampling of `n_layer` +/// non-monomorphic minorant iteration-indices — single pass, no +/// cross-partition resolution, `O(n_layer)` memory regardless of the +/// layer's true size (never materialises the full non-monomorphic index +/// list, only the reservoir itself). +/// +/// The stored values are **`family_idx`, `scan_layer_families`'s own +/// numbering** — the index among *every* minorant of the layer (monomorphic +/// ones included, since `iter_minorants_batch` only filters on +/// `is_minorant()`), not the raw annex slot (`for slot in 0..annex.len()` +/// spans every k-mer, minorant or not) and not a counter over +/// non-monomorphic minorants alone. Conflating any of these with +/// `family_idx` silently drifts apart as soon as the layer contains a +/// monomorphic minorant — i.e. almost immediately, since ~98% of minorants +/// are monomorphic — so the two counters below are deliberately kept +/// separate: `family_idx` tracks every minorant (matching +/// `scan_layer_families`), `seen` tracks only the non-monomorphic ones +/// (what Algorithm R actually samples over). +fn reservoir_sample_layer(layer_dir: &Path, n_layer: usize) -> OKIResult> { + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + let mut reservoir: Vec = Vec::with_capacity(n_layer); + let mut seen: u64 = 0; + let mut family_idx: usize = 0; + let mut rng = rand::thread_rng(); + for slot in 0..annex.len() { + let Some(mask) = annex.get(slot) else { continue }; + if !mask.is_minorant() { + continue; // not a minorant at all — doesn't advance `family_idx` either + } + let this_family_idx = family_idx; + family_idx += 1; + if mask.siblings() < 1 { + continue; // monomorphic minorant — consumed a family_idx above, but not sample-eligible + } + if reservoir.len() < n_layer { + reservoir.push(this_family_idx); + } else { + let j = rng.gen_range(0..=seen); + if j < n_layer as u64 { + reservoir[j as usize] = this_family_idx; + } + } + seen += 1; + } + Ok(reservoir.into_iter().collect()) +} + +/// Per-layer sampling outcome — `None` at a given layer means "keep every +/// non-monomorphic minorant of this layer" (either the `total_count <= n` +/// case, or a layer whose entire non-monomorphic set was already smaller +/// than its computed share and got fully absorbed by rounding). +pub(super) type LayerSelection = Option>; + +/// Full entry point for every caller (`snp_pseudo_alignment`, the Shannon +/// entropy report): `n = None` means "no `--subsample`", i.e. keep every +/// non-monomorphic minorant of every layer — skips steps 1 and 2 entirely, +/// since there is nothing to proportion. +pub(super) fn compute_selections(layer_dirs: &[PathBuf], n: Option) -> OKIResult> { + match n { + None => Ok(vec![None; layer_dirs.len()]), + Some(n) => { + let counts = non_monomorphic_counts(layer_dirs)?; + sample_layers(layer_dirs, &counts, n) + } + } +} + +/// Step 2 entry point: `n` is the total number of families requested across +/// the whole index (`--subsample N`); `layer_dirs` and `counts` must be the +/// same length and in the same order as returned by +/// [`super::family_scan::sibling_layer_dirs`] / +/// [`non_monomorphic_counts`]. +fn sample_layers(layer_dirs: &[PathBuf], counts: &[u64], n: usize) -> OKIResult> { + let total_count: u64 = counts.iter().sum(); + + if total_count <= n as u64 { + // Fewer non-monomorphic minorants in the whole index than requested + // — no sampling needed, keep everything everywhere. + return Ok(vec![None; layer_dirs.len()]); + } + + layer_dirs + .par_iter() + .zip(counts.par_iter()) + .map(|(layer_dir, &count_layer)| -> OKIResult { + if count_layer == 0 { + return Ok(Some(HashSet::new())); + } + let n_layer = ((n as u128 * count_layer as u128) / total_count as u128) as usize; + if n_layer == 0 { + return Ok(Some(HashSet::new())); + } + let selected = reservoir_sample_layer(layer_dir, n_layer)?; + Ok(Some(selected)) + }) + .collect() +} diff --git a/src/obikphylo/src/siblings/tests.rs b/src/obikphylo/src/siblings/tests.rs index 43139e38..3b5359d3 100644 --- a/src/obikphylo/src/siblings/tests.rs +++ b/src/obikphylo/src/siblings/tests.rs @@ -13,6 +13,7 @@ use super::alignment::SnpAlignmentExt; use super::build::SiblingAnnexBuildExt; use super::cardinality::CardinalityExt; use super::distance::DistanceExt; +use super::entropy::ShannonEntropyExt; use super::helpers::is_minorant; use super::stats::SiblingStatsExt; use super::{FamilyMask, SiblingAnnex, ANNEX_FILE_NAME, INDEX_SUBDIR}; @@ -229,7 +230,7 @@ fn family_scan_consumers_agree_on_one_sibling_each() { // snp_pseudo_alignment: one variable family, one column — g1's row // reads 'C' (its own member), g2's reads 'G'. - let alignment = merged.snp_pseudo_alignment().expect("snp_pseudo_alignment"); + let alignment = merged.snp_pseudo_alignment(None).expect("snp_pseudo_alignment"); assert_eq!(alignment.sequences[i1], vec![b'C']); assert_eq!(alignment.sequences[i2], vec![b'G']); @@ -331,3 +332,206 @@ fn sibling_annex_records_the_real_layer_of_each_family_member() { assert_eq!(b.layer_value(2), Some(0), "g2's own base: seeded with its own layer (0)"); assert_eq!(b.layer_value(1), Some(1), "g2's sibling (g1's form): resolved to its real layer (1)"); } + +#[test] +fn subsample_and_shannon_on_one_variable_family() { + let dir = tempdir().unwrap(); + let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG"); + let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG"); + let merged = merge_two(dir.path(), &g1, &g2); + merged.build_sibling_annex().expect("build_sibling_annex"); + + let idx_of = |label: &str| merged.meta().genomes.iter().position(|g| g.label == label).unwrap(); + let (i1, i2) = (idx_of("g1"), idx_of("g2")); + + // This fixture has exactly one non-monomorphic family (see + // `family_scan_consumers_agree_on_one_sibling_each`) — total_count = 1. + // Asking for far more than that must fall back to "keep everything", + // identical to the unsampled (`None`) result. + let sampled = merged.snp_pseudo_alignment(Some(1000)).expect("snp_pseudo_alignment (subsample >> total)"); + assert_eq!(sampled.sequences[i1], vec![b'C']); + assert_eq!(sampled.sequences[i2], vec![b'G']); + + // Asking for 0 must keep nothing — an empty (but still one-row-per-genome) + // alignment, not an error. + let empty = merged.snp_pseudo_alignment(Some(0)).expect("snp_pseudo_alignment (subsample 0)"); + assert!(empty.sequences[i1].is_empty()); + assert!(empty.sequences[i2].is_empty()); + + // shannon_entropy_csv: one row, both entropy15 and entropy4 = 1 bit + // exactly (2 genomes present, each carrying exactly one distinct base, + // uniform split -> -2*(0.5*log2(0.5)) = 1 — the two definitions + // coincide here since no genome carries more than one base). + let csv_path = dir.path().join("shannon.csv"); + merged.shannon_entropy_csv(&csv_path, Some(1000)).expect("shannon_entropy_csv"); + let csv = std::fs::read_to_string(&csv_path).unwrap(); + let mut lines = csv.lines(); + assert_eq!(lines.next(), Some("layer,family_idx,entropy15,entropy4,family_size,n_genomes_present")); + let row = lines.next().expect("exactly one data row"); + assert!(lines.next().is_none(), "exactly one non-monomorphic family in this fixture"); + let fields: Vec<&str> = row.split(',').collect(); + let entropy15: f64 = fields[2].parse().unwrap(); + let entropy4: f64 = fields[3].parse().unwrap(); + assert!((entropy15 - 1.0).abs() < 1e-6, "entropy15 should be exactly 1 bit, got {entropy15}"); + assert!((entropy4 - 1.0).abs() < 1e-6, "entropy4 should be exactly 1 bit, got {entropy4}"); + assert_eq!(fields[4], "2", "family_size"); + assert_eq!(fields[5], "2", "n_genomes_present"); +} + +#[test] +#[ignore] +fn diag_real_index_layer_distribution() { + let idx = KmerIndex::open("/Users/coissac/Sync/travail/__MOI__/obikmer/benchmark/global_index_presence") + .expect("open real index"); + let layer_dirs = super::family_scan::sibling_layer_dirs(&idx).expect("layer dirs"); + let counts = super::subsample::non_monomorphic_counts(&layer_dirs).expect("counts"); + let total: u64 = counts.iter().sum(); + let n_zero = counts.iter().filter(|&&c| c == 0).count(); + let max = counts.iter().max().unwrap(); + let min_nonzero = counts.iter().filter(|&&c| c > 0).min().unwrap_or(&0); + println!("n_layers={} total={} n_zero_layers={} max={} min_nonzero={}", + counts.len(), total, n_zero, max, min_nonzero); + let n = 100_000usize; + let selections = super::subsample::compute_selections(&layer_dirs, Some(n)).expect("selections"); + let selected_total: usize = selections.iter().map(|s| match s { + None => 0, // shouldn't happen when total_count > n + Some(set) => set.len(), + }).sum(); + println!("requested={n} selected_total={selected_total}"); + let mut sorted_counts = counts.clone(); + sorted_counts.sort_unstable_by(|a, b| b.cmp(a)); + println!("top10 counts: {:?}", &sorted_counts[..10.min(sorted_counts.len())]); +} + +#[test] +#[ignore] +fn diag_plant_index_presence_matrix_sparsity() { + let idx = KmerIndex::open("/Users/coissac/travail/obiskim/data/phyloalps/phyloskims_sal_vac") + .expect("open real plant index"); + let layer_dirs = super::family_scan::sibling_layer_dirs(&idx).expect("layer dirs"); + println!("n_layers={}", layer_dirs.len()); + + // Sample a spread of layers rather than just the first few (partition + // order, not necessarily representative) — every 37th layer, capped at 20. + let sample: Vec<&std::path::PathBuf> = layer_dirs.iter().step_by(37).take(20).collect(); + + let mut total_ones: u64 = 0; + let mut total_cells: u64 = 0; + for layer_dir in &sample { + let mat = obicompactvec::PersistentBitMatrix::open(layer_dir).expect("open presence matrix"); + let n = mat.n() as u64; + let n_cols = mat.n_cols() as u64; + let ones: u64 = mat.count_ones().iter().sum(); + let cells = n * n_cols; + let density = ones as f64 / cells as f64; + println!( + "{}: n_slots={n} n_cols={n_cols} ones={ones} cells={cells} density={:.6} sparsity={:.6}", + layer_dir.display(), density, 1.0 - density, + ); + total_ones += ones; + total_cells += cells; + } + let overall_density = total_ones as f64 / total_cells as f64; + println!( + "OVERALL sampled_layers={} total_ones={total_ones} total_cells={total_cells} density={:.6} sparsity={:.6}", + sample.len(), overall_density, 1.0 - overall_density, + ); +} + +#[test] +#[ignore] +fn diag_plant_index_multigenome_set_duplication() { + use std::collections::HashMap; + + let idx = KmerIndex::open("/Users/coissac/travail/obiskim/data/phyloalps/phyloskims_sal_vac") + .expect("open real plant index"); + let layer_dirs = super::family_scan::sibling_layer_dirs(&idx).expect("layer dirs"); + + // Same spread-sampling strategy as the sparsity diagnostic — every + // 37th layer, capped at 8 (this one is heavier per layer: builds a + // hash map of every distinct multi-genome set). + let sample: Vec<&std::path::PathBuf> = layer_dirs.iter().step_by(37).take(8).collect(); + + let mut total_rows: u64 = 0; + let mut total_multi_rows: u64 = 0; + let mut total_distinct_multi_sets: u64 = 0; + let mut total_singleton_rows: u64 = 0; + + for layer_dir in &sample { + let mat = obicompactvec::PersistentBitMatrix::open(layer_dir).expect("open presence matrix"); + let n = mat.n(); + let n_cols = mat.n_cols(); + let mut seen: HashMap, u32> = HashMap::new(); + let mut singleton = 0u64; + let mut multi = 0u64; + let mut buf = vec![0u32; n_cols]; + for slot in 0..n { + mat.fill_row(slot, &mut buf); + let set: Vec = (0..n_cols).filter(|&c| buf[c] != 0).map(|c| c as u32).collect(); + match set.len() { + 0 => {} + 1 => singleton += 1, + _ => { + multi += 1; + *seen.entry(set).or_insert(0) += 1; + } + } + } + let distinct = seen.len() as u64; + println!( + "{}: n_slots={n} n_cols={n_cols} singleton_rows={singleton} multi_rows={multi} distinct_multi_sets={distinct} dedup_ratio={:.4}", + layer_dir.display(), + if multi > 0 { distinct as f64 / multi as f64 } else { 0.0 }, + ); + total_rows += n as u64; + total_singleton_rows += singleton; + total_multi_rows += multi; + total_distinct_multi_sets += distinct; + } + + println!( + "OVERALL sampled_layers={} total_rows={total_rows} singleton_rows={total_singleton_rows} multi_rows={total_multi_rows} distinct_multi_sets={total_distinct_multi_sets} dedup_ratio={:.4}", + sample.len(), + if total_multi_rows > 0 { total_distinct_multi_sets as f64 / total_multi_rows as f64 } else { 0.0 }, + ); +} + +#[test] +#[ignore] +fn diag_plant_index_cardinality_distribution() { + let idx = KmerIndex::open("/Users/coissac/travail/obiskim/data/phyloalps/phyloskims_sal_vac") + .expect("open real plant index"); + let layer_dirs = super::family_scan::sibling_layer_dirs(&idx).expect("layer dirs"); + let sample: Vec<&std::path::PathBuf> = layer_dirs.iter().step_by(37).take(8).collect(); + + let mut hist: Vec = Vec::new(); // hist[k] = number of rows with cardinality k + let mut total_rows: u64 = 0; + + for layer_dir in &sample { + let mat = obicompactvec::PersistentBitMatrix::open(layer_dir).expect("open presence matrix"); + let n = mat.n(); + let n_cols = mat.n_cols(); + if hist.len() < n_cols + 1 { + hist.resize(n_cols + 1, 0); + } + let mut buf = vec![0u32; n_cols]; + for slot in 0..n { + mat.fill_row(slot, &mut buf); + let card = buf.iter().filter(|&&v| v != 0).count(); + hist[card] += 1; + } + total_rows += n as u64; + } + + println!("total_rows={total_rows}"); + let mut cum: u64 = 0; + for (card, &count) in hist.iter().enumerate() { + if count == 0 { continue; } + cum += count; + println!( + "cardinality={card:3} count={count:10} frac={:.6} cumulative_frac={:.6}", + count as f64 / total_rows as f64, + cum as f64 / total_rows as f64, + ); + } +}