feat: add Sankoff cost-matrix calibration CLI and serialization

Introduce the --sankoff flag to generate subsampled calibration bundles and pairwise cost matrices. Expose internal sibling algorithm types as public to support external consumption. Implement YAML, CSV, and FASTA output serialization using serde, gated by a configurable substitution saturation ceiling and requiring the --subsample flag.
This commit is contained in:
Eric Coissac
2026-08-28 21:19:20 +02:00
parent 9043868228
commit 1536217ac1
11 changed files with 320 additions and 35 deletions
+2
View File
@@ -1689,6 +1689,8 @@ dependencies = [
"obiread",
"obiskbuilder",
"obisys",
"serde",
"serde_yaml",
"tracing",
"tracing-subscriber",
]
+2
View File
@@ -29,6 +29,8 @@ obifastwrite = { path = "../obifastwrite" }
obiskbuilder = { path = "../obiskbuilder" }
clap = { version = "4", features = ["derive"] }
csv = "1"
serde = { version = "1", features = ["derive"] }
serde_yaml = "0.9"
tracing = "0.1.44"
tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
+23 -5
View File
@@ -38,12 +38,13 @@ impl From<MetricArg> for DistanceMetric {
/// Partial transfer of `obikmer`'s `phylo` command: the plain distance-metric
/// path (`--metric`/NJ/UPGMA), annex construction (`--sibling-annex`),
/// annex diagnostics (`--sibling-stats`, `--sibling-hist`), entropy
/// reporting (`--shannon`) and SNP pseudo-alignment sampling
/// reporting (`--shannon`), SNP pseudo-alignment sampling
/// (`--pseudo-alignment`, `--subsample`, `--free-loss`, `--no-ambiguity`,
/// `--entropy`/`--entropy-sd`) — everything else sibling-annex-based
/// (`--sankoff`, `--tnt`/`--phyg`/`--iqtree`, raw SNP distance, family
/// overlap, ...) stays in `obikmer` until the rest of
/// `obikphylo::siblings` is reconnected (see the project memory on this).
/// `--entropy`/`--entropy-sd`) and Sankoff cost-matrix calibration
/// (`--sankoff`, `--sankoff-ratio-ceiling`) — everything else sibling-annex-based
/// (`--tnt`/`--phyg`/`--iqtree`, raw SNP distance, family overlap, ...)
/// stays in `obikmer` until the rest of `obikphylo::siblings` is
/// reconnected (see the project memory on this).
#[derive(Args)]
pub struct PhyloArgs {
/// Index directory
@@ -124,6 +125,23 @@ pub struct PhyloArgs {
#[arg(long)]
pub entropy_sd: Option<f64>,
/// Calibrate a 16-state Sankoff cost matrix (and its matching
/// pseudo-alignment) from an already-built sibling annex — requires
/// `--subsample <N>`, and shares `--free-loss`/`--no-ambiguity`/
/// `--entropy`/`--entropy-sd` with `--pseudo-alignment` (one draw, same
/// selection feeds both the alignment and every calibration tally).
#[arg(long)]
pub sankoff: bool,
/// Exclude genome pairs whose raw SNP ratio exceeds this value from the
/// base-pair (composition) calibration `--sankoff` pools — a pair this
/// close to substitution saturation carries no information about the
/// true substitution spectrum. Does *not* gate the cardinality
/// calibration (see `obikphylo::siblings::CardinalityTally`'s own
/// docs for why).
#[arg(long, default_value = "0.5")]
pub sankoff_ratio_ceiling: f64,
/// Distance metric to compute
#[arg(long, value_enum, default_value = "jaccard")]
pub metric: MetricArg,
+59 -12
View File
@@ -1,15 +1,21 @@
mod args;
mod sankoff;
use std::io::{self, BufWriter, Write};
use std::sync::Arc;
use obikidxcache::index_cache::IndexCache;
use obikindex::KmerIndex;
use obikphylo::siblings::{EntropyBias, SiblingExt};
use obikphylo::siblings::{
EntropyBias, SiblingExt, cardinality_transition_probs, composition_transition_probs,
pairwise_cost_matrix,
};
use obikphylo::{Metrics, neighbor_joining, upgma};
use obisys::{Reporter, Stage};
use tracing::info;
use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params};
pub use args::PhyloArgs;
pub fn run(args: PhyloArgs) {
@@ -152,23 +158,24 @@ pub fn run(args: PhyloArgs) {
info!("entropy report → {path}");
}
// Shared by `--pseudo-alignment` and `--sankoff` — same activation rule:
// either flag given activates entropy-biased sampling, the other
// defaults to 1.0/0.5.
let entropy_bias = if args.entropy.is_some() || args.entropy_sd.is_some() {
Some(EntropyBias {
mu: args.entropy.unwrap_or(1.0),
sigma: args.entropy_sd.unwrap_or(0.5),
})
} else {
None
};
// ── SNP pseudo-alignment (`--pseudo-alignment`) ─────────────────────────────
if args.pseudo_alignment {
let Some(subsample_n) = args.subsample else {
eprintln!("error: --pseudo-alignment requires --subsample <N>");
std::process::exit(1);
};
// Same activation rule as `--shannon`'s own entropy-biased path
// would use: either flag given activates biasing, the other
// defaults to 1.0/0.5.
let entropy_bias = if args.entropy.is_some() || args.entropy_sd.is_some() {
Some(EntropyBias {
mu: args.entropy.unwrap_or(1.0),
sigma: args.entropy_sd.unwrap_or(0.5),
})
} else {
None
};
info!("sampling SNP pseudo-alignment (target {subsample_n} site(s))");
let t = Stage::start("pseudo_alignment");
@@ -194,6 +201,46 @@ pub fn run(args: PhyloArgs) {
info!("pseudo-alignment ({n_sites} site(s), {} genome(s)) → {path}", alignment.genome_indices.len());
}
// ── Sankoff cost-matrix calibration (`--sankoff`) ───────────────────────────
if args.sankoff {
let Some(subsample_n) = args.subsample else {
eprintln!("error: --sankoff requires --subsample <N>");
std::process::exit(1);
};
info!("sampling Sankoff calibration bundle (target {subsample_n} site(s))");
let t = Stage::start("sankoff_bundle");
let bundle = cache
.sankoff_bundle(
subsample_n,
args.free_loss,
args.no_ambiguity,
&exclude_mask,
entropy_bias,
args.sankoff_ratio_ceiling,
)
.unwrap_or_else(|e| {
eprintln!("error computing Sankoff calibration bundle: {e}");
std::process::exit(1);
});
rep.push(t.stop());
let p_card = cardinality_transition_probs(&bundle.cardinality_tally);
let p_comp = composition_transition_probs(&bundle.base_pair_tally);
let matrix = pairwise_cost_matrix(&p_card, &p_comp, args.free_loss);
write_sankoff_matrix_csv(&matrix, &args.output);
write_sankoff_params(
&bundle.cardinality_tally,
&p_card,
&bundle.base_pair_tally,
&p_comp,
args.sankoff_ratio_ceiling,
&args.output,
);
write_sankoff_alignment_fasta(&bundle.alignment, &labels, &args.output, args.free_loss);
}
info!("computing {:?} distances for {} genome(s)", args.metric, n);
let need_shared = args.shared_kmers || args.nj || args.upgma;
+168
View File
@@ -0,0 +1,168 @@
//! Output writers for `--sankoff` — no calibration logic here, just
//! formatting: `obikphylo::siblings::SiblingExt::sankoff_bundle` and the
//! `cardinality_transition_probs`/`composition_transition_probs`/
//! `pairwise_cost_matrix` calibration functions do all the actual work in
//! `mod.rs`, this module only serialises their results.
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use obifastwrite::{JsonVal, write_record};
use obikphylo::siblings::{BasePairTally, CardinalityTally, SnpAlignment};
use tracing::info;
// ── Sankoff pseudo-alignment → FASTA ────────────────────────────────────────
//
// Same data as `--pseudo-alignment`'s output (`SnpAlignment`/
// `snp_pseudo_alignment`), re-coded so its symbols match the accompanying
// `--sankoff` matrix output exactly: `0` for the empty/absent state instead
// of `-`, which TNT/PhyG would otherwise read as their own gap character
// rather than our "family absent" state. Unless `free_loss` (`--free-loss`)
// is set, in which case `∅` is recoded to `?` instead — TNT/PhyG's own
// missing-data symbol, deliberately *not* `-` (still gap/indel semantics in
// both tools) — so non-detection costs nothing rather than being scored as
// an ordinary, calibrated state transition.
pub(super) fn write_sankoff_alignment_fasta(
alignment: &SnpAlignment,
labels: &[String],
output: &Option<PathBuf>,
free_loss: bool,
) {
let path = output.as_ref()
.map(|p| format!("{}_sankoff.fasta", p.display()))
.unwrap_or_else(|| "sankoff.fasta".into());
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
eprintln!("error creating {path}: {e}");
std::process::exit(1);
}));
let absent_symbol = if free_loss { b'?' } else { b'0' };
let n_sites = alignment.sequences.first().map_or(0, Vec::len);
for (&g, seq) in alignment.genome_indices.iter().zip(alignment.sequences.iter()) {
let recoded: Vec<u8> = seq.iter().map(|&b| if b == b'-' { absent_symbol } else { b }).collect();
write_record(&recoded, &labels[g], &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f)
.unwrap_or_else(|e| {
eprintln!("error writing {path}: {e}");
std::process::exit(1);
});
}
info!("Sankoff pseudo-alignment ({n_sites} site(s)) → {path}");
}
// ── Sankoff cost matrix → CSV ────────────────────────────────────────────────
//
// 16 states indexed by bitmask (bit 0=A, 1=C, 2=G, 3=T; state 0 is `∅`),
// matching the convention used for `--pseudo-alignment`'s IUPAC-coded output
// and for the external TNT/PhyG scripts this feeds.
/// IUPAC ambiguity code per state (same mapping
/// `obikphylo::siblings::algorithms::masking::iupac_code` uses internally
/// for `--pseudo-alignment`), with `0` standing in for the empty state (`-`
/// would collide with TNT/PhyG's own gap/range syntax). Bit order: 0=A,
/// 1=C, 2=G, 3=T. This project's canonical alphabet for every Sankoff
/// export (`--tnt`/`--phyg`/`--iqtree` each recode it to their own alphabet
/// at their own adapter boundary, rather than using it directly).
pub(super) const STATE_SYMBOL: [char; 16] = [
'0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N',
];
pub(super) fn write_sankoff_matrix_csv(matrix: &[[f64; 16]; 16], output: &Option<PathBuf>) {
let path = output.as_ref()
.map(|p| format!("{}_sankoff_matrix.csv", p.display()))
.unwrap_or_else(|| "sankoff_matrix.csv".into());
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
eprintln!("error creating {path}: {e}");
std::process::exit(1);
}));
write!(f, "state").unwrap();
for sym in STATE_SYMBOL { write!(f, ",{sym}").unwrap(); }
writeln!(f).unwrap();
for (s, row) in matrix.iter().enumerate() {
write!(f, "{}", STATE_SYMBOL[s]).unwrap();
for cost in row { write!(f, ",{cost:.4}").unwrap(); }
writeln!(f).unwrap();
}
info!("Sankoff cost matrix → {path}");
}
// ── Sankoff calibration parameters → YAML report ────────────────────────────
//
// Everything `--sankoff` estimates from real data, in one durable,
// machine-readable file: the cardinality and base-pair transition tallies
// (raw counts, not just the derived probabilities) — costs are a modelling
// choice built *from* the counts, and reproducing/re-deriving them later
// needs the counts, not just their current derived value.
#[derive(serde::Serialize)]
struct CardinalityTransition {
from: usize,
to: usize,
count: u64,
probability: f64,
}
#[derive(serde::Serialize)]
struct CompositionTransition {
from: char,
to: char,
count: u64,
probability: f64,
}
#[derive(serde::Serialize)]
struct SankoffParamsReport {
ratio_ceiling: f64,
cardinality_transitions: Vec<CardinalityTransition>,
composition_transitions: Vec<CompositionTransition>,
}
pub(super) fn write_sankoff_params(
card_tally: &CardinalityTally,
p_card: &[[f64; 5]; 5],
base_tally: &BasePairTally,
p_comp: &[[f64; 4]; 4],
ratio_ceiling: f64,
output: &Option<PathBuf>,
) {
const BASE_LETTER: [char; 4] = ['A', 'C', 'G', 'T'];
let mut cardinality_transitions = Vec::with_capacity(25);
for a in 0..5 {
for b in 0..5 {
cardinality_transitions.push(CardinalityTransition {
from: a,
to: b,
count: card_tally.counts[a][b],
probability: p_card[a][b],
});
}
}
let mut composition_transitions = Vec::with_capacity(16);
for a in 0..4 {
for b in 0..4 {
let count = if a == b { base_tally.same[a] } else { base_tally.counts[a][b] };
composition_transitions.push(CompositionTransition {
from: BASE_LETTER[a],
to: BASE_LETTER[b],
count,
probability: p_comp[a][b],
});
}
}
let report = SankoffParamsReport { ratio_ceiling, cardinality_transitions, composition_transitions };
let path = output.as_ref()
.map(|p| format!("{}_sankoff_params.yaml", p.display()))
.unwrap_or_else(|| "sankoff_params.yaml".into());
let f = std::fs::File::create(&path).unwrap_or_else(|e| {
eprintln!("error creating {path}: {e}");
std::process::exit(1);
});
serde_yaml::to_writer(f, &report).unwrap_or_else(|e| {
eprintln!("error writing {path}: {e}");
std::process::exit(1);
});
info!("Sankoff calibration parameters → {path}");
}
@@ -22,7 +22,7 @@ use super::pairwise::{BasePairTally, CardinalityTally};
/// Row-stochastic 5×5 cardinality transition probabilities (`0..=4`),
/// diagonal included ("stay at the same cardinality"), from
/// [`CardinalityTally`]'s pooled co-occurrence counts.
pub(crate) fn cardinality_transition_probs(tally: &CardinalityTally) -> [[f64; 5]; 5] {
pub fn cardinality_transition_probs(tally: &CardinalityTally) -> [[f64; 5]; 5] {
let mut p = [[0.0f64; 5]; 5];
for a in 0..5 {
let row_sum: u64 = tally.counts[a].iter().sum();
@@ -40,7 +40,7 @@ pub(crate) fn cardinality_transition_probs(tally: &CardinalityTally) -> [[f64; 5
/// (`A,C,G,T`), diagonal included ("stay the same base"), from
/// [`BasePairTally`]'s pooled substitution (off-diagonal) and agreement
/// (`same`, diagonal) counts — unambiguous, cardinality-1 loci only.
pub(crate) fn composition_transition_probs(tally: &BasePairTally) -> [[f64; 4]; 4] {
pub fn composition_transition_probs(tally: &BasePairTally) -> [[f64; 4]; 4] {
let mut p = [[0.0f64; 4]; 4];
for a in 0..4 {
let row_sum = tally.same[a] + (0..4).map(|b| tally.counts[a][b]).sum::<u64>();
@@ -143,7 +143,7 @@ fn best_pairing_cost(lost: &[u8], gained: &[u8], p_comp: &[[f64; 4]; 4]) -> f64
/// factor dropped, cost is driven only by composition matching
/// (shared-base retention and paired substitutions), never by a state
/// pair's cardinality difference alone.
pub(crate) fn pairwise_cost_matrix(
pub fn pairwise_cost_matrix(
p_card: &[[f64; 5]; 5],
p_comp: &[[f64; 4]; 4],
free_loss: bool,
+4 -3
View File
@@ -27,17 +27,18 @@ mod subsample;
use obikidxcache::index_cache::IndexCache;
pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix};
pub use alignment::SnpAlignment;
pub use pairwise::{BasePairTally, CardinalityTally, RawSnpDistanceOutput};
pub use sankoff::SankoffBundle;
pub use stats::SiblingAnnexStats;
pub use subsample::EntropyBias;
pub(crate) use alignment::snp_pseudo_alignment;
pub(crate) use annex::build_layer_sibling_annex;
pub(crate) use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix};
pub(crate) use entropy::{ensure_layer_entropy_annex, family_entropy, family_entropy_4, iter_full_entropy};
pub(crate) use family_scan::{Selection, scan_layer_families};
pub(crate) use pairwise::{BasePairTally, CardinalityTally, RawSnpDistanceOutput};
pub(crate) use sankoff::{SankoffBundle, sankoff_bundle};
pub(crate) use sankoff::sankoff_bundle;
pub(crate) use stats::{sibling_annex_stats, sibling_family_size_histogram};
pub(crate) use subsample::sample_index;
@@ -178,8 +178,9 @@ impl PairwiseTally {
}
/// Raw p-distance restricted to loci that are single-copy in **both**
/// genomes of a pair — see `PairwiseTally`'s module docs.
pub(crate) struct RawSnpDistanceOutput {
/// genomes of a pair — see `PairwiseTally`'s module docs. `pub`: part of
/// [`super::SankoffBundle`]'s public signature.
pub struct RawSnpDistanceOutput {
/// n×n count of eligible loci where the two genomes' single forms differ.
pub snp: Array2<u64>,
/// n×n count of eligible loci where the two genomes' single forms agree.
@@ -187,8 +188,9 @@ pub(crate) struct RawSnpDistanceOutput {
}
/// Symmetric 6-category base-pair substitution tally (indexed
/// `0=A,1=C,2=G,3=T`), pooled over [`PairwiseTally::included`] genome pairs.
pub(crate) struct BasePairTally {
/// `0=A,1=C,2=G,3=T`), pooled over [`PairwiseTally::included`] genome
/// pairs. `pub`: part of [`super::SankoffBundle`]'s public signature.
pub struct BasePairTally {
/// `counts[a][b] == counts[b][a]` = number of eligible loci, pooled over
/// included genome pairs, where the two genomes' single forms are `a`
/// and `b`. Diagonal always `0` — an `a == b` locus is counted in
@@ -207,8 +209,9 @@ pub(crate) struct BasePairTally {
/// [`PairwiseTally::cardinality_tally`]'s own docs for why this is *not*
/// gated by `ratio_ceiling` the way [`BasePairTally`] is), restricted to
/// variable families (`family_size() >= 2`) — matching
/// `snp_pseudo_alignment`'s own scope.
pub(crate) struct CardinalityTally {
/// `snp_pseudo_alignment`'s own scope. `pub`: part of
/// [`super::SankoffBundle`]'s public signature.
pub struct CardinalityTally {
/// `counts[a][b] == counts[b][a]` = number of family sites, pooled over
/// included genome pairs, where one genome's family cardinality is `a`
/// and the other's is `b`. Diagonal is real data here (both genomes at
@@ -19,8 +19,9 @@ use super::subsample::{EntropyBias, sample_index};
/// Every output the `--sankoff`/`--tnt`/`--phyg`/`--iqtree` pipeline needs,
/// computed together from one shared, possibly-subsampled/entropy-biased
/// selection — see the module docs.
pub(crate) struct SankoffBundle {
/// selection — see the module docs. `pub`: part of the public signature of
/// [`crate::siblings::extensions::SiblingExt::sankoff_bundle`].
pub struct SankoffBundle {
pub alignment: SnpAlignment,
pub raw: RawSnpDistanceOutput,
pub base_pair_tally: BasePairTally,
@@ -14,9 +14,10 @@ use obikindex::{OKIError, OKIResult};
use obisys::progress_bar;
use crate::siblings::algorithms::{
EntropyBias, Selection, SiblingAnnexStats, SnpAlignment, build_layer_sibling_annex,
family_entropy, family_entropy_4, is_fast_mode, scan_layer_families,
sibling_annex_stats, sibling_family_size_histogram, snp_pseudo_alignment,
EntropyBias, SankoffBundle, Selection, SiblingAnnexStats, SnpAlignment,
build_layer_sibling_annex, family_entropy, family_entropy_4, is_fast_mode,
sankoff_bundle, scan_layer_families, sibling_annex_stats, sibling_family_size_histogram,
snp_pseudo_alignment,
};
use crate::siblings::extensions::SiblingBuilder;
use crate::siblings::ENTROPY_ANNEX_FILE_NAME;
@@ -94,6 +95,32 @@ pub trait SiblingExt {
/// use [`sibling_family_size_histogram`](Self::sibling_family_size_histogram)
/// instead when only the global histogram is needed.
fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats>;
/// Fused entry point for the `--sankoff`/`--tnt`/`--phyg`/`--iqtree`
/// pipeline: one shared, possibly-subsampled/entropy-biased selection
/// (`n`/`free_loss`/`no_ambiguity`/`excluded`/`entropy_bias` — same
/// meaning as [`snp_pseudo_alignment`](Self::snp_pseudo_alignment)'s
/// own) drives a single scan producing the pseudo-alignment *and* every
/// tally the Sankoff cost-matrix calibration
/// (`crate::siblings::cardinality_transition_probs`/
/// `composition_transition_probs`/`pairwise_cost_matrix`) needs —
/// never two independent draws of the same index. `ratio_ceiling`
/// (`--sankoff-ratio-ceiling`) excludes genome pairs too close to
/// substitution saturation from `SankoffBundle::base_pair_tally`'s
/// pool (a saturated pair's base composition is noise, not signal) —
/// `SankoffBundle::cardinality_tally` is *not* gated by it (see
/// `crate::siblings::CardinalityTally`'s own docs for why that
/// wouldn't make sense: cardinality reflects each genome's own
/// coverage/duplication structure, not the pair's mutual divergence).
fn sankoff_bundle(
&self,
n: usize,
free_loss: bool,
no_ambiguity: bool,
excluded: &[bool],
entropy_bias: Option<EntropyBias>,
ratio_ceiling: f64,
) -> OKIResult<SankoffBundle>;
}
impl SiblingExt for IndexCache {
@@ -226,4 +253,16 @@ impl SiblingExt for IndexCache {
fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats> {
sibling_annex_stats(self)
}
fn sankoff_bundle(
&self,
n: usize,
free_loss: bool,
no_ambiguity: bool,
excluded: &[bool],
entropy_bias: Option<EntropyBias>,
ratio_ceiling: f64,
) -> OKIResult<SankoffBundle> {
sankoff_bundle(self, n, free_loss, no_ambiguity, excluded, entropy_bias, ratio_ceiling)
}
}
+5 -1
View File
@@ -29,7 +29,11 @@ mod siblingannex;
pub(crate) use entropy_annex::{ENTROPY_ANNEX_FILE_NAME, EntropyAnnex, EntropyAnnexBuilder};
pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
pub use algorithms::{EntropyBias, SiblingAnnexStats, SnpAlignment};
pub use algorithms::{
BasePairTally, CardinalityTally, EntropyBias, RawSnpDistanceOutput, SankoffBundle,
SiblingAnnexStats, SnpAlignment, cardinality_transition_probs, composition_transition_probs,
pairwise_cost_matrix,
};
pub use extensions::SiblingExt;
pub(crate) const ANNEX_FILE_NAME: &str = "siblings.psib";