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.

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<u8>; 16],
+    /// Compact index -> canonical state index, order-preserving.
+    compact_to_old: Vec<u8>,
+    /// Empirical frequency of each compact-indexed state (sums to 1).
+    freq: Vec<f64>,
+}
+
+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<u8>; 16] = [None; 16];
+    let mut compact_to_old: Vec<u8> = 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<f64> = 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 <file>+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<PathBuf>) -> 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<String> = (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::<Vec<_>>().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<PathBuf>,
+) -> (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<u8> = 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<PathBuf>,
+) {
+    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()
+    );
+}

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 <prefix>_iqtree.model and <prefix>_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: <prefix>_dist.csv, <prefix>_shared.csv,
-    /// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_snp.fasta,
+    /// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_rawsnp_counts.csv,
+    /// <prefix>_snp.fasta,
     /// <prefix>_sankoff_matrix.csv, <prefix>_sankoff_params.yaml,
     /// <prefix>_sankoff.fasta, <prefix>_sankoff.tnt, <prefix>_sankoff.tcm,
-    /// <prefix>_sankoff.pg, <prefix>_nj.nwk, <prefix>_upgma.nwk.
+    /// <prefix>_sankoff.pg, <prefix>_iqtree.model, <prefix>_iqtree.fasta,
+    /// <prefix>_nj.nwk, <prefix>_upgma.nwk.
     /// If omitted, the distance matrix is written to stdout.
     #[arg(short, long)]
     pub output: Option<PathBuf>,
@@ -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<PathBuf>) {
+    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

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,

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();