From 28d841c7beda8b7fe740680c925724c037c1f30e Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 12 Aug 2026 18:23:51 +0200 Subject: [PATCH] feat(distance): add IQ-TREE output and optimize state index mapping Introduce `--iqtree` and `--raw-snp-counts` flags to generate IQ-TREE model files, recoded FASTA alignments, and per-pair diagnostic counts. Centralize alphabet conversion by extracting a precomputed state index lookup table into the Sankoff module, eliminating redundant iterations across downstream adapters. --- .gitignore | 1 + src/obikmer/src/cmd/distance/iqtree.rs | 183 ++++++++++++++++++++++++ src/obikmer/src/cmd/distance/mod.rs | 91 +++++++++++- src/obikmer/src/cmd/distance/sankoff.rs | 11 ++ src/obikmer/src/cmd/distance/tnt.rs | 7 +- 5 files changed, 285 insertions(+), 8 deletions(-) create mode 100644 src/obikmer/src/cmd/distance/iqtree.rs diff --git a/.gitignore b/.gitignore index 3cd5920d..32189eb2 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,4 @@ benchmark/specific_index_count benchmark/specific_index_presence TNT phyg +*.tnt diff --git a/src/obikmer/src/cmd/distance/iqtree.rs b/src/obikmer/src/cmd/distance/iqtree.rs new file mode 100644 index 00000000..0a1ac655 --- /dev/null +++ b/src/obikmer/src/cmd/distance/iqtree.rs @@ -0,0 +1,183 @@ +use std::io::{BufWriter, Write}; +use std::path::PathBuf; + +use obifastwrite::{JsonVal, write_record}; +use obikindex::SnpAlignment; +use tracing::info; + +use super::sankoff::state_index_table; + +// ── Sankoff-calibrated data → IQ-TREE custom ML model + recoded alignment ── +// +// Not itself a Sankoff computation — IQ-TREE does maximum likelihood, not +// parsimony. Only the *source data* is shared with `--tnt`/`--phyg` (the +// calibrated 16-state cost matrix, the pseudo-alignment); the operation +// performed on it here is different, hence no "sankoff" in these names, +// unlike `tnt::write_sankoff_tnt`/`phyg::write_sankoff_phyg`. +// +// Verified against the locally installed `iqtree3` binary/source (not just +// its docs — see `docmd/theory/evolutionary_distances.md`, "Next +// direction: genuine ML branch lengths"), because the web documentation's +// `-mdef` NEXUS route turned out not to apply to a plain (non-mixture) +// custom morphology model. The real mechanism: pass a **file path** +// directly as `-m`, containing (as whitespace/newline-separated numbers) +// the lower-triangular exchangeability matrix `R` (`k(k-1)/2` values, PAML +// row-major order) immediately followed by the `k` state frequencies `π` +// on the same stream — `ModelMarkov::readRates`/`readStateFreq` read them +// in that exact order, no header, no separator required. +// +// `R` is recovered from the calibrated Sankoff cost matrix via +// `R(a,b) = exp(-cost(a,b))` (the cost is `-ln(rate)`, see "A concrete +// Sankoff cost matrix"), symmetric by construction (the underlying tally +// never captured direction). `π` is the real, empirical, non-uniform +// marginal frequency of each state across the whole alignment — precise +// enough at this sample size (908k+ sites) without spending IQ-TREE's own +// `+FO` ML degrees of freedom re-estimating it (checked: IQ-TREE's `+F` +// doesn't work as a shortcut here either, a custom-file model always +// requires the frequency line in the file itself). IQ-TREE reconstructs +// the (generally asymmetric) rate matrix internally as +// `Q(i,j) = R(i,j)·π_j` — reversible for *any* `π`, not just uniform, +// because `R` is symmetric. +// +// IQ-TREE infers its state count from the highest-ordinal symbol actually +// present in the alignment, not from a declared count (`--seqtype +// MORPH{N}` was tested and does not override this for real ML analysis, +// only for the `--alisim` simulator). So states that never occur anywhere +// in this particular alignment are dropped, and the survivors are +// renumbered compactly (`0..k-1`, order preserved) rather than leaving +// gaps that would silently misalign every value IQ-TREE reads. Both the +// model and the alignment must agree on this same renumbering, so it's +// computed once (`CompactAlphabet`) and shared between them. + +const IQTREE_STATE_SYMBOL: [char; 16] = [ + '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', +]; + +struct CompactAlphabet { + /// Canonical (0..16) state index -> compact index, for states that occur. + old_to_compact: [Option; 16], + /// Compact index -> canonical state index, order-preserving. + compact_to_old: Vec, + /// Empirical frequency of each compact-indexed state (sums to 1). + freq: Vec, +} + +impl CompactAlphabet { + fn k(&self) -> usize { + self.compact_to_old.len() + } +} + +fn compact_alphabet(alignment: &SnpAlignment) -> CompactAlphabet { + let iupac_to_state = state_index_table(); + + let mut occurs = [false; 16]; + let mut counts = [0u64; 16]; + for seq in &alignment.sequences { + for &b in seq { + let b = if b == b'-' { b'0' } else { b }; + let state = iupac_to_state[b as usize] as usize; + occurs[state] = true; + counts[state] += 1; + } + } + + let mut old_to_compact: [Option; 16] = [None; 16]; + let mut compact_to_old: Vec = Vec::new(); + for old in 0..16 { + if occurs[old] { + old_to_compact[old] = Some(compact_to_old.len() as u8); + compact_to_old.push(old as u8); + } + } + + let total: u64 = compact_to_old.iter().map(|&old| counts[old as usize]).sum(); + let freq: Vec = compact_to_old.iter() + .map(|&old| counts[old as usize] as f64 / total as f64) + .collect(); + + CompactAlphabet { old_to_compact, compact_to_old, freq } +} + +/// Write the `R` (exchangeability) + `π` (frequencies) model file IQ-TREE's +/// `-m +ASC` reads. Returns the path, so the caller can print a +/// single combined "how to run this" message once the alignment is also +/// written. +fn write_iqtree_model(matrix: &[[f64; 16]; 16], alphabet: &CompactAlphabet, output: &Option) -> String { + let rate = |old_i: u8, old_j: u8| (-matrix[old_i as usize][old_j as usize]).exp(); + + let model_path = output.as_ref() + .map(|p| format!("{}_iqtree.model", p.display())) + .unwrap_or_else(|| "iqtree.model".into()); + let mut f = BufWriter::new(std::fs::File::create(&model_path).unwrap_or_else(|e| { + eprintln!("error creating {model_path}: {e}"); + std::process::exit(1); + })); + for i in 1..alphabet.k() { + let row: Vec = (0..i) + .map(|j| format!("{:.6}", rate(alphabet.compact_to_old[i], alphabet.compact_to_old[j]))) + .collect(); + writeln!(f, "{}", row.join(" ")).unwrap(); + } + writeln!(f, "{}", alphabet.freq.iter().map(|p| format!("{p:.6}")).collect::>().join(" ")).unwrap(); + info!("IQ-TREE model file → {model_path} ({} of 16 states present in the alignment)", alphabet.k()); + model_path +} + +/// Write the pseudo-alignment recoded to the same compact `0..k-1` alphabet +/// as `write_iqtree_model`'s matrix — not `--sankoff`'s own IUPAC alphabet, +/// since IQ-TREE needs the symbol ordinal itself to match the surviving +/// state count (see this module's own doc comment on `MORPH{N}`). +fn write_iqtree_alignment( + alignment: &SnpAlignment, + labels: &[String], + alphabet: &CompactAlphabet, + output: &Option, +) -> (String, usize) { + let iupac_to_state = state_index_table(); + + let fasta_path = output.as_ref() + .map(|p| format!("{}_iqtree.fasta", p.display())) + .unwrap_or_else(|| "iqtree.fasta".into()); + let mut f = BufWriter::new(std::fs::File::create(&fasta_path).unwrap_or_else(|e| { + eprintln!("error creating {fasta_path}: {e}"); + std::process::exit(1); + })); + let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); + for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { + let recoded: Vec = seq.iter().map(|&b| { + let b = if b == b'-' { b'0' } else { b }; + let old = iupac_to_state[b as usize] as usize; + let compact = alphabet.old_to_compact[old] + .expect("state occurs in the alignment, so it must have a compact index"); + IQTREE_STATE_SYMBOL[compact as usize] as u8 + }).collect(); + write_record(&recoded, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| { + eprintln!("error writing {fasta_path}: {e}"); + std::process::exit(1); + }); + } + (fasta_path, n_sites) +} + +pub(super) fn write_iqtree( + matrix: &[[f64; 16]; 16], + alignment: &SnpAlignment, + labels: &[String], + output: &Option, +) { + let alphabet = compact_alphabet(alignment); + let model_path = write_iqtree_model(matrix, &alphabet, output); + let (fasta_path, n_sites) = write_iqtree_alignment(alignment, labels, &alphabet, output); + + let prefix_name = output.as_ref() + .and_then(|p| p.file_name()) + .map(|n| format!("{}_iqtree", n.to_string_lossy())) + .unwrap_or_else(|| "iqtree".into()); + info!( + "IQ-TREE alignment → {fasta_path} ({n_sites} sites, {} states)\n\ + Run with:\n \ + iqtree3 -s {fasta_path} --seqtype MORPH -m {model_path}+ASC --prefix {prefix_name} -T AUTO", + alphabet.k() + ); +} diff --git a/src/obikmer/src/cmd/distance/mod.rs b/src/obikmer/src/cmd/distance/mod.rs index bd698caa..3dfe08d3 100644 --- a/src/obikmer/src/cmd/distance/mod.rs +++ b/src/obikmer/src/cmd/distance/mod.rs @@ -1,3 +1,4 @@ +mod iqtree; mod phyg; mod sankoff; mod tnt; @@ -16,6 +17,7 @@ use obikindex::{ use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; +use iqtree::write_iqtree; use phyg::write_sankoff_phyg; use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params}; use tnt::write_sankoff_tnt; @@ -98,6 +100,19 @@ pub struct DistanceArgs { #[arg(long)] pub raw_snp_distance: bool, + /// Write the raw per-pair counts (`n_snp`, `n_shared`, `n_eligible`) + /// behind `--raw-snp-distance`'s ratio, one row per genome pair — a + /// diagnostic table, not a matrix. The ratio alone can't distinguish + /// "identical at every eligible locus" from "almost no eligible loci + /// at all" (e.g. `0.0` from 0/2 looks the same as `0.0` from 0/2000), + /// and that distinction matters a lot for genome pairs near the edge + /// of what central-position families can resolve (see + /// `docmd/theory/evolutionary_distances.md`, "Run 3" and the + /// IQ-TREE/Mash comparison). Same annex requirement as + /// `--raw-snp-distance`. + #[arg(long)] + pub raw_snp_counts: bool, + /// Write a SNP-only pseudo-alignment (FASTA, IUPAC-coded) from an /// already-built sibling annex — one row per genome, one column per /// variable family (monomorphic families skipped), no flanking @@ -141,6 +156,24 @@ pub struct DistanceArgs { #[arg(long)] pub phyg: bool, + /// Also write _iqtree.model and _iqtree.fasta, a + /// custom-model file and a matching + /// recoded alignment for genuine maximum-likelihood inference with + /// IQ-TREE (`iqtree3 -s ... --seqtype MORPH -m ...+ASC`) — real branch + /// lengths, unlike `--tnt`/`--phyg`'s parsimony step counts. The model + /// is the reversible `Q(i,j) = R(i,j)·π_j` construction: `R` + /// (exchangeability, symmetric) recovered from the same calibrated + /// cost matrix `--sankoff` computes, `π` the real empirical state + /// frequencies counted from the alignment (not IQ-TREE's `+FO`/`+F` — + /// neither applies to a custom-file model, see + /// `docmd/theory/evolutionary_distances.md`). Only the states that + /// actually occur in this alignment are kept, compactly renumbered + /// (IQ-TREE infers its state count from the alignment itself, and a + /// gap in the numbering would silently misalign the model file). + /// Implies `--sankoff`. + #[arg(long)] + pub iqtree: bool, + /// Scale factor applied before rounding real-valued costs to the /// integers both `--tnt`'s smatrix/cost commands and `--phyg`'s `tcm:` /// matrix require. Keep this small: the total tree score is this scale @@ -157,10 +190,12 @@ pub struct DistanceArgs { pub sankoff_cost_scale: f64, /// Output prefix: _dist.csv, _shared.csv, - /// _siblings.csv, _rawsnp.csv, _snp.fasta, + /// _siblings.csv, _rawsnp.csv, _rawsnp_counts.csv, + /// _snp.fasta, /// _sankoff_matrix.csv, _sankoff_params.yaml, /// _sankoff.fasta, _sankoff.tnt, _sankoff.tcm, - /// _sankoff.pg, _nj.nwk, _upgma.nwk. + /// _sankoff.pg, _iqtree.model, _iqtree.fasta, + /// _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] pub output: Option, @@ -208,6 +243,13 @@ pub fn run(args: DistanceArgs) { }); write_raw_snp_distance_csv(&result, &labels, &args.output); } + if args.raw_snp_counts { + let result = idx.raw_snp_distance().unwrap_or_else(|e| { + eprintln!("error computing raw SNP distance: {e}"); + std::process::exit(1); + }); + write_raw_snp_counts_csv(&result, &labels, &args.output); + } if args.snp { let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { eprintln!("error computing SNP pseudo-alignment: {e}"); @@ -215,7 +257,7 @@ pub fn run(args: DistanceArgs) { }); write_snp_fasta(&alignment, &labels, &args.output); } - if args.sankoff || args.tnt || args.phyg { + if args.sankoff || args.tnt || args.phyg || args.iqtree { let raw = idx.raw_snp_distance().unwrap_or_else(|e| { eprintln!("error computing raw SNP distance: {e}"); std::process::exit(1); @@ -248,6 +290,9 @@ pub fn run(args: DistanceArgs) { if args.phyg { write_sankoff_phyg(&matrix, &args.output, args.sankoff_cost_scale); } + if args.iqtree { + write_iqtree(&matrix, &alignment, &labels, &args.output); + } } // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp`/ @@ -260,10 +305,12 @@ pub fn run(args: DistanceArgs) { if args.sibling_annex || args.sibling_stats || args.raw_snp_distance + || args.raw_snp_counts || args.snp || args.sankoff || args.tnt || args.phyg + || args.iqtree { return; } @@ -450,6 +497,44 @@ fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: &[String], info!("raw single-copy SNP distance matrix → {path}"); } +// ── Raw single-copy SNP distance → per-pair diagnostic counts ────────────── +// +// A pair table (one row per unordered genome pair), not a matrix: the ratio +// alone can't distinguish "identical across every eligible locus" from +// "almost no eligible locus at all" — both can read `0.0`/`NA` in +// `--raw-snp-distance`'s output. Distinguishing them matters most exactly +// where it's easy to miss: genome pairs near the edge of what +// central-position families can resolve at all (deep cross-lineage splits, +// see `docmd/theory/evolutionary_distances.md`, "Run 3" and the later +// IQ-TREE/Mash comparison — a `ratio=0.0` backed by 2 eligible loci is not +// the same claim as one backed by 2000). + +fn write_raw_snp_counts_csv(result: &RawSnpDistanceOutput, labels: &[String], output: &Option) { + let path = output.as_ref() + .map(|p| format!("{}_rawsnp_counts.csv", p.display())) + .unwrap_or_else(|| "rawsnp_counts.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); + })); + let n = labels.len(); + writeln!(f, "genome_a,genome_b,n_snp,n_shared,n_eligible,ratio").unwrap(); + for i in 0..n { + for j in (i + 1)..n { + let snp = result.snp[[i, j]]; + let shared = result.shared[[i, j]]; + let eligible = snp + shared; + write!(f, "{},{},{snp},{shared},{eligible}", labels[i], labels[j]).unwrap(); + if eligible == 0 { + writeln!(f, ",NA").unwrap(); + } else { + writeln!(f, ",{:.6}", snp as f64 / eligible as f64).unwrap(); + } + } + } + info!("raw single-copy SNP distance counts (diagnostic) → {path}"); +} + // ── SNP-only pseudo-alignment → FASTA ─────────────────────────────────────── // // One record per genome, IUPAC-coded, no flanking sequence — see diff --git a/src/obikmer/src/cmd/distance/sankoff.rs b/src/obikmer/src/cmd/distance/sankoff.rs index 76676237..5e02fda6 100644 --- a/src/obikmer/src/cmd/distance/sankoff.rs +++ b/src/obikmer/src/cmd/distance/sankoff.rs @@ -52,6 +52,17 @@ pub(super) const STATE_SYMBOL: [char; 16] = [ '0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N', ]; +/// `STATE_SYMBOL` byte -> state index (0..16), for adapters that need to +/// translate an alignment written in this alphabet into their own. Shared +/// rather than rebuilt per adapter (`tnt`, `iqtree`). +pub(super) fn state_index_table() -> [u8; 128] { + let mut table = [0u8; 128]; + for (state, &sym) in STATE_SYMBOL.iter().enumerate() { + table[sym as usize] = state as u8; + } + table +} + pub(super) fn write_sankoff_matrix_csv( matrix: &[[f64; 16]; 16], estimate: &PHatEstimate, diff --git a/src/obikmer/src/cmd/distance/tnt.rs b/src/obikmer/src/cmd/distance/tnt.rs index 977994c4..48f5ad53 100644 --- a/src/obikmer/src/cmd/distance/tnt.rs +++ b/src/obikmer/src/cmd/distance/tnt.rs @@ -4,7 +4,7 @@ use std::path::PathBuf; use obikindex::SnpAlignment; use tracing::info; -use super::sankoff::{STATE_SYMBOL, scaled_metric_matrix}; +use super::sankoff::{scaled_metric_matrix, state_index_table}; // ── Sankoff cost matrix + alignment → ready-to-run TNT script ────────────── // @@ -40,10 +40,7 @@ pub(super) fn write_sankoff_tnt( // IUPAC-ish symbol -> bitmask, to translate the alignment (which uses // `STATE_SYMBOL`, `-` already normalised to `0` by `snp_pseudo_alignment` // callers) into TNT's alphabet without re-deriving state indices. - let mut iupac_to_state = [0u8; 128]; - for (state, &sym) in STATE_SYMBOL.iter().enumerate() { - iupac_to_state[sym as usize] = state as u8; - } + let iupac_to_state = state_index_table(); let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); writeln!(f, "xread").unwrap();