src/obikmer/src/cmd/phylo/iqtree.rs

@@ -5,7 +5,7 @@ use obifastwrite::{JsonVal, write_record};
 use obikphylo::siblings::SnpAlignment;
 use tracing::info;
 
-use super::sankoff::{state_index_table, STATE_SYMBOL};
+use super::sankoff::{STATE_SYMBOL, state_index_table};
 
 // ── Sankoff-calibrated data → IQ-TREE custom ML model + recoded alignment ──
 //
@@ -85,7 +85,8 @@ impl CompactAlphabet {
 /// two exports.
 fn drop_ascertainment_noninformative(alignment: &SnpAlignment) -> SnpAlignment {
     let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
-    let keep: Vec<bool> = (0..n_sites).map(|site| {
+    let keep: Vec<bool> = (0..n_sites)
+        .map(|site| {
             let mut first: Option<u8> = None;
             for seq in &alignment.sequences {
                 let b = seq[site];
@@ -99,10 +100,19 @@ fn drop_ascertainment_noninformative(alignment: &SnpAlignment) -> SnpAlignment {
                 }
             }
             false // all calls missing, or all calls agree — non-informative
-    }).collect();
+        })
+        .collect();
 
-    let sequences = alignment.sequences.iter()
+    let sequences = alignment
-        .map(|seq| seq.iter().zip(keep.iter()).filter(|&(_, &k)| k).map(|(&b, _)| b).collect())
+        .sequences
+        .iter()
+        .map(|seq| {
+            seq.iter()
+                .zip(keep.iter())
+                .filter(|&(_, &k)| k)
+                .map(|(&b, _)| b)
+                .collect()
+        })
         .collect();
     SnpAlignment { sequences }
 }
@@ -114,8 +124,14 @@ fn drop_ascertainment_noninformative(alignment: &SnpAlignment) -> SnpAlignment {
 /// states into the missing-data treatment before a second
 /// `compact_alphabet` pass, without duplicating that treatment's logic.
 fn recode_symbols_as_absent(alignment: &SnpAlignment, symbols: &[u8]) -> SnpAlignment {
-    let sequences = alignment.sequences.iter()
+    let sequences = alignment
-        .map(|seq| seq.iter().map(|&b| if symbols.contains(&b) { b'-' } else { b }).collect())
+        .sequences
+        .iter()
+        .map(|seq| {
+            seq.iter()
+                .map(|&b| if symbols.contains(&b) { b'-' } else { b })
+                .collect()
+        })
         .collect();
     SnpAlignment { sequences }
 }
@@ -151,11 +167,16 @@ fn compact_alphabet(alignment: &SnpAlignment, free_loss: bool) -> CompactAlphabe
     }
 
     let total: u64 = compact_to_old.iter().map(|&old| counts[old as usize]).sum();
-    let freq: Vec<f64> = compact_to_old.iter()
+    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 }
+    CompactAlphabet {
+        old_to_compact,
+        compact_to_old,
+        freq,
+    }
 }
 
 /// Write `<prefix>_iqtree_states.csv`: the mapping from IQ-TREE's own
@@ -168,7 +189,8 @@ fn compact_alphabet(alignment: &SnpAlignment, free_loss: bool) -> CompactAlphabe
 /// (e.g. "state 0 has zero exchangeability with everything else") can't be
 /// traced back to which real state that is.
 fn write_iqtree_states_csv(alphabet: &CompactAlphabet, output: &Option<PathBuf>) -> String {
-    let path = output.as_ref()
+    let path = output
+        .as_ref()
         .map(|p| format!("{}_iqtree_states.csv", p.display()))
         .unwrap_or_else(|| "iqtree_states.csv".into());
     let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
@@ -177,7 +199,12 @@ fn write_iqtree_states_csv(alphabet: &CompactAlphabet, output: &Option<PathBuf>)
     }));
     writeln!(f, "iqtree_symbol,canonical_symbol,frequency").unwrap();
     for (compact, &old) in alphabet.compact_to_old.iter().enumerate() {
-        writeln!(f, "{},{},{}", IQTREE_STATE_SYMBOL[compact], STATE_SYMBOL[old as usize], alphabet.freq[compact]).unwrap();
+        writeln!(
+            f,
+            "{},{},{}",
+            IQTREE_STATE_SYMBOL[compact], STATE_SYMBOL[old as usize], alphabet.freq[compact]
+        )
+        .unwrap();
     }
     path
 }
@@ -186,10 +213,15 @@ fn write_iqtree_states_csv(alphabet: &CompactAlphabet, output: &Option<PathBuf>)
 /// `-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 {
+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()
+    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| {
@@ -198,12 +230,30 @@ fn write_iqtree_model(matrix: &[[f64; 16]; 16], alphabet: &CompactAlphabet, outp
     }));
     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])))
+            .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();
+    writeln!(
-    info!("IQ-TREE model file → {model_path} ({} of 16 states present in the alignment)", alphabet.k());
+        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
 }
 
@@ -220,7 +270,8 @@ fn write_iqtree_alignment(
 ) -> (String, usize) {
     let iupac_to_state = state_index_table();
 
-    let fasta_path = output.as_ref()
+    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| {
@@ -229,7 +280,9 @@ fn write_iqtree_alignment(
     }));
     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 recoded: Vec<u8> = seq
+            .iter()
+            .map(|&b| {
                 if free_loss && b == b'-' {
                     return b'?';
                 }
@@ -238,8 +291,15 @@ fn write_iqtree_alignment(
                 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| {
+            .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);
         });
@@ -283,7 +343,10 @@ pub(super) fn write_iqtree(
     // instability warnings).
     let refiltered;
     let alignment = if free_loss {
-        let low_freq_symbols: Vec<u8> = alphabet.compact_to_old.iter().zip(alphabet.freq.iter())
+        let low_freq_symbols: Vec<u8> = alphabet
+            .compact_to_old
+            .iter()
+            .zip(alphabet.freq.iter())
             .filter(|&(_, &f)| f < min_freq)
             .map(|(&old, _)| STATE_SYMBOL[old as usize] as u8)
             .collect();
@@ -297,7 +360,10 @@ pub(super) fn write_iqtree(
             info!(
                 "--iqtree-min-freq {min_freq}: {} rare state(s) ({}) recoded as missing, {before} → {after} sites",
                 low_freq_symbols.len(),
-                low_freq_symbols.iter().map(|&b| b as char).collect::<String>(),
+                low_freq_symbols
+                    .iter()
+                    .map(|&b| b as char)
+                    .collect::<String>(),
             );
             alphabet = compact_alphabet(&refiltered, free_loss);
             &refiltered
@@ -308,9 +374,11 @@ pub(super) fn write_iqtree(
 
     let states_path = write_iqtree_states_csv(&alphabet, output);
     let model_path = write_iqtree_model(matrix, &alphabet, output);
-    let (fasta_path, n_sites) = write_iqtree_alignment(alignment, labels, &alphabet, output, free_loss);
+    let (fasta_path, n_sites) =
+        write_iqtree_alignment(alignment, labels, &alphabet, output, free_loss);
 
-    let prefix_name = output.as_ref()
+    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());
@@ -318,7 +386,9 @@ pub(super) fn write_iqtree(
         "IQ-TREE alignment → {fasta_path} ({n_sites} sites, {} states)\n\
          IQ-TREE state mapping → {states_path}\n\
          Run with:\n  \
-         iqtree3 -s {fasta_path} --seqtype MORPH -m {model_path}+ASC --prefix {prefix_name} -T AUTO",
+         iqtree3 -s {fasta_path} --seqtype MORPH -m {model_path}+ASC --prefix {prefix_name} -T AUTO\n\
+         \n\
+         options -alrt 1000 -B 1000 can be added to evaluate robustness of the tree",
         alphabet.k()
     );
 }
@@ -336,11 +406,7 @@ mod tests {
         // ("absent") was still being counted despite being recoded to `?`
         // (IQ-TREE's own missing symbol) in the alignment actually written.
         let alignment = SnpAlignment {
-            sequences: vec![
+            sequences: vec![vec![b'A', b'-'], vec![b'C', b'-'], vec![b'-', b'G']],
-                vec![b'A', b'-'],
-                vec![b'C', b'-'],
-                vec![b'-', b'G'],
-            ],
         };
 
         let alphabet = compact_alphabet(&alignment, true);
@@ -351,18 +417,18 @@ mod tests {
             alphabet.compact_to_old
         );
         let sum: f64 = alphabet.freq.iter().sum();
-        assert!((sum - 1.0).abs() < 1e-9, "frequencies must sum to 1, got {sum} ({:?})", alphabet.freq);
+        assert!(
+            (sum - 1.0).abs() < 1e-9,
+            "frequencies must sum to 1, got {sum} ({:?})",
+            alphabet.freq
+        );
         assert_eq!(alphabet.k(), 3, "A, C, G — 3 real states, `-` excluded");
     }
 
     #[test]
     fn without_free_loss_absent_state_is_counted_normally() {
         let alignment = SnpAlignment {
-            sequences: vec![
+            sequences: vec![vec![b'A', b'-'], vec![b'C', b'-'], vec![b'-', b'G']],
-                vec![b'A', b'-'],
-                vec![b'C', b'-'],
-                vec![b'-', b'G'],
-            ],
         };
 
         let alphabet = compact_alphabet(&alignment, false);
@@ -372,7 +438,11 @@ mod tests {
             "state 0 (absent, recoded from '-') must be counted when --free-loss is off"
         );
         let sum: f64 = alphabet.freq.iter().sum();
-        assert!((sum - 1.0).abs() < 1e-9, "frequencies must sum to 1, got {sum} ({:?})", alphabet.freq);
+        assert!(
+            (sum - 1.0).abs() < 1e-9,
+            "frequencies must sum to 1, got {sum} ({:?})",
+            alphabet.freq
+        );
     }
 
     #[test]
@@ -380,19 +450,21 @@ mod tests {
         // 'A' (state 1) and 'G' (state 4) occur, '-' (state 0) excluded by
         // --free-loss — compact index 0 -> 'A', compact index 1 -> 'G'.
         let alignment = SnpAlignment {
-            sequences: vec![
+            sequences: vec![vec![b'A', b'-'], vec![b'-', b'G']],
-                vec![b'A', b'-'],
-                vec![b'-', b'G'],
-            ],
         };
         let alphabet = compact_alphabet(&alignment, true);
-        let output = Some(std::env::temp_dir().join(format!("obikmer_test_iqtree_states_{}", std::process::id())));
+        let output = Some(
+            std::env::temp_dir().join(format!("obikmer_test_iqtree_states_{}", std::process::id())),
+        );
 
         let path = write_iqtree_states_csv(&alphabet, &output);
         let csv = std::fs::read_to_string(&path).unwrap();
         std::fs::remove_file(&path).ok();
         let mut lines = csv.lines();
-        assert_eq!(lines.next(), Some("iqtree_symbol,canonical_symbol,frequency"));
+        assert_eq!(
+            lines.next(),
+            Some("iqtree_symbol,canonical_symbol,frequency")
+        );
         assert_eq!(lines.next(), Some("0,A,0.5"));
         assert_eq!(lines.next(), Some("1,G,0.5"));
         assert!(lines.next().is_none());
@@ -407,7 +479,11 @@ mod tests {
         // at 1/62, well below the 0.05 threshold used here.
         let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); 3];
         for i in 0..20 {
-            let (a, b, c) = if i % 2 == 0 { (b'A', b'C', b'A') } else { (b'C', b'A', b'C') };
+            let (a, b, c) = if i % 2 == 0 {
+                (b'A', b'C', b'A')
+            } else {
+                (b'C', b'A', b'C')
+            };
             sequences[0].push(a);
             sequences[1].push(b);
             sequences[2].push(c);
@@ -418,15 +494,24 @@ mod tests {
         let alignment = SnpAlignment { sequences };
         let labels = vec!["g1".to_string(), "g2".to_string(), "g3".to_string()];
         let matrix = [[0.0f64; 16]; 16];
-        let prefix = std::env::temp_dir().join(format!("obikmer_test_iqtree_minfreq_{}", std::process::id()));
+        let prefix = std::env::temp_dir().join(format!(
+            "obikmer_test_iqtree_minfreq_{}",
+            std::process::id()
+        ));
         let output = Some(prefix.clone());
 
         write_iqtree(&matrix, &alignment, &labels, &output, true, 0.05);
 
         let states_path = format!("{}_iqtree_states.csv", prefix.display());
         let csv = std::fs::read_to_string(&states_path).unwrap();
-        assert!(!csv.contains(",M,"), "M (freq ~1/62) must be folded into missing under --iqtree-min-freq 0.05, got:\n{csv}");
+        assert!(
-        assert!(csv.contains(",A,") && csv.contains(",C,"), "A/C must survive (well above threshold), got:\n{csv}");
+            !csv.contains(",M,"),
+            "M (freq ~1/62) must be folded into missing under --iqtree-min-freq 0.05, got:\n{csv}"
+        );
+        assert!(
+            csv.contains(",A,") && csv.contains(",C,"),
+            "A/C must survive (well above threshold), got:\n{csv}"
+        );
 
         for suffix in ["_iqtree_states.csv", "_iqtree.model", "_iqtree.fasta"] {
             std::fs::remove_file(format!("{}{suffix}", prefix.display())).ok();