From e015362ce62fe5b34751c420ef4118e58f2be353 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 27 Aug 2026 16:28:57 +0200 Subject: [PATCH] Refactor sampling and reduction to use batched SurvivingFamily data Replace per-family callbacks with a per-layer batch callback that accumulates surviving families into an Arc>. Update the alignment reduction routine to iterate over this batched structure, enabling zero-copy sharing for parallel operations and structuring the code for future extensibility. --- .../src/siblings/algorithms/alignment.rs | 39 ++++++++--- .../src/siblings/algorithms/subsample.rs | 65 ++++++++++++++----- 2 files changed, 80 insertions(+), 24 deletions(-) diff --git a/src/obikphylo/src/siblings/algorithms/alignment.rs b/src/obikphylo/src/siblings/algorithms/alignment.rs index 71094742..6624cc75 100644 --- a/src/obikphylo/src/siblings/algorithms/alignment.rs +++ b/src/obikphylo/src/siblings/algorithms/alignment.rs @@ -7,7 +7,7 @@ use obikidxcache::index_cache::IndexCache; use obikindex::OKIResult; use super::masking::{iupac_code, masked_state}; -use super::subsample::{EntropyBias, sample_index}; +use super::subsample::{EntropyBias, SurvivingFamily, sample_index}; /// `sequences[i]` is `genome_indices[i]`'s IUPAC-coded row — every row the /// same length (one byte per sampled family, in the order [`sample_index`] @@ -27,7 +27,12 @@ pub struct SnpAlignment { } /// See [`crate::siblings::extensions::SiblingExt::snp_pseudo_alignment`] -/// for the public-facing docs — this is its implementation. +/// for the public-facing docs — this is its implementation. A single +/// "reduce" ([`reduce_alignment`]) over [`sample_index`]'s per-layer +/// batches — the only consumer today, but structured so a future sibling +/// reduce (e.g. the Sankoff cost-matrix calibration's cardinality/ +/// composition tallies) can run over the exact same batches, each cloning +/// its own `Arc`, without this function changing at all. pub(crate) fn snp_pseudo_alignment( cache: &IndexCache, n: usize, @@ -49,15 +54,31 @@ pub(crate) fn snp_pseudo_alignment( no_ambiguity, excluded, entropy_bias, - |_partition, _layer, _family_idx, _mask, genome_mask| { - for (row, &g) in genome_indices.iter().enumerate() { - let ch = masked_state(genome_mask[g], free_loss, no_ambiguity) - .map(iupac_code) - .unwrap_or(b'?'); - sequences[row].push(ch); - } + |_partition, _layer, survivors| { + reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences); }, )?; Ok(SnpAlignment { sequences, genome_indices }) } + +/// Appends one layer's worth of columns onto `sequences` — a "reduce" over +/// [`sample_index`]'s per-layer batch, self-contained (no sampling/masking +/// logic of its own, just character encoding) so it can be called +/// independently of, and alongside, any other reduce over the same batch. +fn reduce_alignment( + survivors: &[SurvivingFamily], + genome_indices: &[usize], + free_loss: bool, + no_ambiguity: bool, + sequences: &mut [Vec], +) { + for family in survivors { + for (row, &g) in genome_indices.iter().enumerate() { + let ch = masked_state(family.genome_mask[g], free_loss, no_ambiguity) + .map(iupac_code) + .unwrap_or(b'?'); + sequences[row].push(ch); + } + } +} diff --git a/src/obikphylo/src/siblings/algorithms/subsample.rs b/src/obikphylo/src/siblings/algorithms/subsample.rs index 1767707c..c017fa42 100644 --- a/src/obikphylo/src/siblings/algorithms/subsample.rs +++ b/src/obikphylo/src/siblings/algorithms/subsample.rs @@ -44,6 +44,7 @@ use std::collections::HashSet; use std::path::Path; +use std::sync::Arc; use obicompactvec::TempBitVecBuilder; use obikidxcache::index_cache::IndexCache; @@ -63,6 +64,19 @@ use crate::siblings::extensions::SiblingBuilder; /// accepted families rather than resolving them one at a time. const MAX_TOPUP_ROUNDS: usize = 5; +/// One surviving family — already drawn (Bernoulli/entropy-bias), already +/// resolved (`scan_layer_families`), already passed post-hoc masking +/// (`super::masking::survives_masking`). `genome_mask` is an owned copy: the +/// borrowed slice `scan_layer_families` hands its own callback doesn't +/// outlive that callback, but a `SurvivingFamily` is meant to be handed to +/// [`sample_index`]'s caller as part of a whole layer's batch, well after +/// that callback returns. +pub(crate) struct SurvivingFamily { + pub family_idx: usize, + pub mask: FamilyMask, + pub genome_mask: Vec, +} + /// Gaussian-kernel entropy bias parameters — `mu`/`sigma` are the /// `--entropy`/`--entropy-sd` CLI values (defaulted by the caller, not /// here). `pub`, not `pub(crate)`: part of the public signature of @@ -89,17 +103,29 @@ fn circular_entropy_values( } /// Samples up to `n` non-monomorphic families index-wide, proportionally -/// per layer, calling `on_family(partition, layer, family_idx, mask, -/// genome_mask)` once for every family that was both drawn *and* still -/// carries real polymorphism after `free_loss`/`no_ambiguity`/`excluded` -/// masking (`algorithms::masking::survives_masking`) — i.e. every family -/// this produces is already a usable alignment site, no second resolution -/// pass needed by the caller. `excluded[g]` (see `survives_masking`'s own -/// docs) never counts genome `g` toward that check — `genome_mask` itself -/// still carries every genome's own resolved state, excluded or not, so a -/// caller building per-genome output (e.g. an alignment row) must apply the -/// same exclusion itself when consuming it. Returns the actual number of -/// sites kept, which may be less than `n` (see the module docs). +/// per layer, calling `on_layer(partition, layer, survivors)` once per +/// cached layer that contributed anything, with every family from that +/// layer that was both drawn *and* still carries real polymorphism after +/// `free_loss`/`no_ambiguity`/`excluded` masking +/// (`algorithms::masking::survives_masking`) — i.e. every +/// [`SurvivingFamily`] this produces is already a usable alignment site, no +/// second resolution pass needed by the caller. `excluded[g]` (see +/// `survives_masking`'s own docs) never counts genome `g` toward that check +/// — `genome_mask` itself still carries every genome's own resolved state, +/// excluded or not, so a caller building per-genome output (e.g. an +/// alignment row) must apply the same exclusion itself when consuming it. +/// +/// `survivors` is a single `Arc>` — one allocation per +/// layer (bounded by that layer's own quota, never the whole index), handed +/// out already shared: a caller that wants several independent reduces over +/// the same batch (an alignment builder, a tally builder, ...) just clones +/// the `Arc` once per reduce (an `O(1)` refcount bump, no data copy) instead +/// of re-deriving or re-resolving anything — including handing a clone to +/// its own `rayon::join`/`scope` worker, since `Arc` (unlike a plain +/// borrowed slice) satisfies the `Send + 'static` those need. +/// +/// Returns the actual number of sites kept, which may be less than `n` (see +/// the module docs). pub(crate) fn sample_index( cache: &IndexCache, n: usize, @@ -107,7 +133,7 @@ pub(crate) fn sample_index( no_ambiguity: bool, excluded: &[bool], entropy_bias: Option, - mut on_family: impl FnMut(usize, usize, usize, FamilyMask, &[u8]), + mut on_layer: impl FnMut(usize, usize, Arc>), ) -> OKIResult { let n_genomes = cache.meta().genomes().len(); let fast_mode = is_fast_mode(cache); @@ -164,7 +190,7 @@ pub(crate) fn sample_index( no_ambiguity, excluded, entropy_bias, - |family_idx, mask, genome_mask| on_family(part, l, family_idx, mask, genome_mask), + &mut on_layer, )?; } @@ -187,7 +213,7 @@ fn sample_layer( no_ambiguity: bool, excluded: &[bool], entropy_bias: Option, - mut on_family: impl FnMut(usize, FamilyMask, &[u8]), + on_layer: &mut impl FnMut(usize, usize, Arc>), ) -> OKIResult { if n_minorants == 0 { return Ok(0); @@ -208,6 +234,7 @@ fn sample_layer( let mut visited = 0usize; let mut kept = 0usize; let mut resolved = 0usize; // total accepted (and resolved) so far, across every round + let mut survivors: Vec = Vec::new(); for _round in 0..MAX_TOPUP_ROUNDS { // Stop on success (quota reached) or true exhaustion (every @@ -286,12 +313,20 @@ fn sample_layer( "eligible bitset must only accept non-monomorphic minorants" ); if survives_masking(genome_mask, excluded, free_loss, no_ambiguity) { - on_family(family_idx, mask, genome_mask); + survivors.push(SurvivingFamily { + family_idx, + mask, + genome_mask: genome_mask.to_vec(), + }); kept += 1; } }, )?; } + if !survivors.is_empty() { + on_layer(partition, layer_idx, Arc::new(survivors)); + } + Ok(kept) }