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.
This commit is contained in:
Eric Coissac
@@ -26,6 +26,7 @@ benchmark/specific_index_count
|
||||
benchmark/specific_index_presence
|
||||
TNT
|
||||
phyg
|
||||
biblio
|
||||
*.tnt
|
||||
*.tre
|
||||
*.phy
|
||||
|
||||
+182
-44
@@ -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<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.
|
||||
|
||||
Generated
+1
@@ -1858,6 +1858,7 @@ dependencies = [
|
||||
"obiskbuilder",
|
||||
"obiskio",
|
||||
"obisys",
|
||||
"rand 0.8.6",
|
||||
"rayon",
|
||||
"tempfile",
|
||||
"tracing",
|
||||
|
||||
@@ -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<usize>,
|
||||
|
||||
/// Write <prefix>_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 (`<prefix>_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: <prefix>_dist.csv, <prefix>_shared.csv,
|
||||
/// <prefix>_siblings.csv, <prefix>_sibling_hist.csv, <prefix>_rawsnp.csv,
|
||||
/// <prefix>_rawsnp_counts.csv,
|
||||
/// <prefix>_snp.fasta, <prefix>_family_overlap.csv,
|
||||
/// <prefix>_snp.fasta, <prefix>_family_overlap.csv, <prefix>_shannon.csv,
|
||||
/// <prefix>_sankoff_matrix.csv, <prefix>_sankoff_params.yaml,
|
||||
/// <prefix>_sankoff.fasta, <prefix>_sankoff.tnt, <prefix>_sankoff.tcm,
|
||||
/// <prefix>_sankoff.pg, <prefix>_iqtree.model, <prefix>_iqtree.fasta,
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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<SnpAlignment>;
|
||||
///
|
||||
/// `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<usize>) -> OKIResult<SnpAlignment>;
|
||||
}
|
||||
|
||||
impl SnpAlignmentExt for KmerIndex {
|
||||
fn snp_pseudo_alignment(&self) -> OKIResult<SnpAlignment> {
|
||||
fn snp_pseudo_alignment(&self, subsample: Option<usize>) -> OKIResult<SnpAlignment> {
|
||||
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<u8>> = 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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<usize>) -> OKIResult<()>;
|
||||
}
|
||||
|
||||
impl ShannonEntropyExt for KmerIndex {
|
||||
fn shannon_entropy_csv(&self, path: &Path, subsample: Option<usize>) -> 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(())
|
||||
}
|
||||
}
|
||||
@@ -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<usize>),
|
||||
}
|
||||
|
||||
/// Shared by the pipeline worker (via the `Arc`-cloned `Option<HashSet>`,
|
||||
/// 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<std::collections::HashSet<usize>>, 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<PartitionCache>,
|
||||
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<Option<std::collections::HashSet<usize>>> = 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;
|
||||
}
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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<Vec<u64>> {
|
||||
layer_dirs
|
||||
.par_iter()
|
||||
.map(|layer_dir| -> OKIResult<u64> {
|
||||
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<HashSet<usize>> {
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
|
||||
let mut reservoir: Vec<usize> = 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<HashSet<usize>>;
|
||||
|
||||
/// 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<usize>) -> OKIResult<Vec<LayerSelection>> {
|
||||
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<Vec<LayerSelection>> {
|
||||
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<LayerSelection> {
|
||||
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()
|
||||
}
|
||||
@@ -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<Vec<u32>, 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<u32> = (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<u64> = 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,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user