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.

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<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;
         }