Add execution timing and parallelize sibling stats
Instruments the phylo command pipeline with structured execution timing, wrapping major computational blocks with stage hooks and printing aggregated metrics upon completion. Additionally, parallelizes sibling counting logic using Rayon to process independent layer directories concurrently, preserving identical functionality and public API contracts.
This commit is contained in:
Eric Coissac
@@ -17,6 +17,7 @@ use obikphylo::{
|
||||
SnpAlignment, SnpAlignmentExt,
|
||||
},
|
||||
};
|
||||
use obisys::{Reporter, Stage};
|
||||
use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
|
||||
use tracing::info;
|
||||
|
||||
@@ -38,6 +39,8 @@ pub fn run(args: PhyloArgs) {
|
||||
let labels: Vec<String> = idx.meta().genomes.iter().map(|g| g.label.clone()).collect();
|
||||
let n = labels.len();
|
||||
|
||||
let mut rep = Reporter::new();
|
||||
|
||||
// ── Genome exclusion (`--exclude-genome`) ───────────────────────────────
|
||||
// Applied by zeroing a `RawSnpDistanceOutput`'s excluded rows/columns
|
||||
// (`zero_excluded_pairs`) — `base_pair_tally`/`cardinality_tally`
|
||||
@@ -127,82 +130,104 @@ pub fn run(args: PhyloArgs) {
|
||||
std::process::exit(1);
|
||||
});
|
||||
info!("building sibling-count/minorant annex");
|
||||
let t = Stage::start("sibling_annex");
|
||||
idx.build_sibling_annex().unwrap_or_else(|e| {
|
||||
eprintln!("error building sibling annex: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
}
|
||||
if args.sibling_stats {
|
||||
let t = Stage::start("sibling_stats");
|
||||
let stats = idx.sibling_annex_stats().unwrap_or_else(|e| {
|
||||
eprintln!("error computing sibling-annex stats: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
write_sibling_stats_csv(&stats, &labels, &args.output);
|
||||
}
|
||||
if args.sibling_hist {
|
||||
let t = Stage::start("sibling_hist");
|
||||
let counts = idx.sibling_family_size_histogram().unwrap_or_else(|e| {
|
||||
eprintln!("error computing sibling family-size histogram: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
write_sibling_hist_csv(&counts, &args.output);
|
||||
}
|
||||
if args.raw_snp_distance {
|
||||
let t = Stage::start("raw_snp_distance");
|
||||
let mut result = idx.raw_snp_distance().unwrap_or_else(|e| {
|
||||
eprintln!("error computing raw SNP distance: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
zero_excluded_pairs(&mut result);
|
||||
write_raw_snp_distance_csv(&result, &labels, &args.output);
|
||||
}
|
||||
if args.raw_snp_counts {
|
||||
let t = Stage::start("raw_snp_distance");
|
||||
let mut result = idx.raw_snp_distance().unwrap_or_else(|e| {
|
||||
eprintln!("error computing raw SNP distance: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
zero_excluded_pairs(&mut result);
|
||||
write_raw_snp_counts_csv(&result, &labels, &args.output);
|
||||
}
|
||||
if args.snp {
|
||||
let t = Stage::start("snp_pseudo_alignment");
|
||||
let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| {
|
||||
eprintln!("error computing SNP pseudo-alignment: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
let (alignment, kept_labels) = drop_excluded(alignment);
|
||||
write_snp_fasta(&alignment, &kept_labels, &args.output);
|
||||
}
|
||||
if args.family_overlap {
|
||||
let t = Stage::start("snp_pseudo_alignment");
|
||||
let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| {
|
||||
eprintln!("error computing SNP pseudo-alignment: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
let (alignment, kept_labels) = drop_excluded(alignment);
|
||||
write_family_overlap_csv(&alignment, &kept_labels, &args.output);
|
||||
}
|
||||
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| {
|
||||
eprintln!("error computing raw SNP distance: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
zero_excluded_pairs(&mut raw);
|
||||
|
||||
let t = Stage::start("base_pair_tally");
|
||||
let base_tally = idx.base_pair_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| {
|
||||
eprintln!("error computing base-pair tally: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
let t = Stage::start("cardinality_tally");
|
||||
let card_tally = idx.cardinality_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| {
|
||||
eprintln!("error computing cardinality tally: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
let p_card = cardinality_transition_probs(&card_tally);
|
||||
let p_comp = composition_transition_probs(&base_tally);
|
||||
let matrix = pairwise_cost_matrix(&p_card, &p_comp, args.free_loss);
|
||||
write_sankoff_matrix_csv(&matrix, &args.output);
|
||||
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| {
|
||||
eprintln!("error computing SNP pseudo-alignment: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
let (alignment, kept_labels) = drop_excluded(alignment);
|
||||
write_sankoff_alignment_fasta(&alignment, &kept_labels, &args.output, args.free_loss);
|
||||
|
||||
@@ -226,6 +251,7 @@ pub fn run(args: PhyloArgs) {
|
||||
// defaults to) pass and printing an unrequested matrix.
|
||||
if args.sibling_annex
|
||||
|| args.sibling_stats
|
||||
|| args.sibling_hist
|
||||
|| args.raw_snp_distance
|
||||
|| args.raw_snp_counts
|
||||
|| args.snp
|
||||
@@ -235,6 +261,7 @@ pub fn run(args: PhyloArgs) {
|
||||
|| args.phyg
|
||||
|| args.iqtree
|
||||
{
|
||||
rep.print();
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -244,12 +271,14 @@ pub fn run(args: PhyloArgs) {
|
||||
);
|
||||
|
||||
let need_shared = args.shared_kmers || args.nj || args.upgma;
|
||||
let t = Stage::start("distance");
|
||||
let result = idx
|
||||
.distance(args.metric.into(), need_shared, args.presence_threshold)
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!("error computing distances: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
|
||||
// ── Distance matrix → CSV ─────────────────────────────────────────────────
|
||||
let write_dist_csv = |w: &mut dyn Write| {
|
||||
@@ -347,4 +376,6 @@ pub fn run(args: PhyloArgs) {
|
||||
});
|
||||
info!("UPGMA tree → {path}");
|
||||
}
|
||||
|
||||
rep.print();
|
||||
}
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use rayon::prelude::*;
|
||||
|
||||
use obicompactvec::SiblingAnnex;
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obisys::progress_bar;
|
||||
@@ -68,8 +70,15 @@ impl SiblingStatsExt for KmerIndex {
|
||||
fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]> {
|
||||
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
|
||||
|
||||
let mut counts = [0u64; 4];
|
||||
for layer_dir in &layer_dirs {
|
||||
// Each layer's annex file is independent — no shared `PartitionCache`
|
||||
// and no `scan_layer_families` partition-grouped locality to protect
|
||||
// (unlike `sibling_annex_stats`/`distance`/`cardinality_tally`), so
|
||||
// layers can be scanned concurrently.
|
||||
layer_dirs
|
||||
.par_iter()
|
||||
.try_fold(
|
||||
|| [0u64; 4],
|
||||
|mut counts, layer_dir| -> OKIResult<[u64; 4]> {
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
|
||||
for slot in 0..annex.len() {
|
||||
let Some(mask) = annex.get(slot) else { continue };
|
||||
@@ -78,9 +87,19 @@ impl SiblingStatsExt for KmerIndex {
|
||||
}
|
||||
counts[mask.siblings() as usize] += 1;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(counts)
|
||||
},
|
||||
)
|
||||
.try_reduce(
|
||||
|| [0u64; 4],
|
||||
|a, b| {
|
||||
let mut sum = a;
|
||||
for i in 0..4 {
|
||||
sum[i] += b[i];
|
||||
}
|
||||
Ok(sum)
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats> {
|
||||
|
||||
Reference in New Issue
Block a user