diff --git a/src/obikmer/src/cmd/phylo/iqtree.rs b/src/obikmer/src/cmd/phylo/iqtree.rs index 3ec1b712..bfe92863 100644 --- a/src/obikmer/src/cmd/phylo/iqtree.rs +++ b/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,24 +85,34 @@ 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 = (0..n_sites).map(|site| { - let mut first: Option = None; - for seq in &alignment.sequences { - let b = seq[site]; - if b == b'-' { - continue; + let keep: Vec = (0..n_sites) + .map(|site| { + let mut first: Option = None; + for seq in &alignment.sequences { + let b = seq[site]; + if b == b'-' { + continue; + } + match first { + None => first = Some(b), + Some(f) if f != b => return true, + _ => {} + } } - match first { - None => first = Some(b), - Some(f) if f != b => return true, - _ => {} - } - } - false // all calls missing, or all calls agree — non-informative - }).collect(); + false // all calls missing, or all calls agree — non-informative + }) + .collect(); - let sequences = alignment.sequences.iter() - .map(|seq| seq.iter().zip(keep.iter()).filter(|&(_, &k)| k).map(|(&b, _)| b).collect()) + let sequences = alignment + .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() - .map(|seq| seq.iter().map(|&b| if symbols.contains(&b) { b'-' } else { b }).collect()) + let sequences = alignment + .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 = compact_to_old.iter() + 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 } + CompactAlphabet { + old_to_compact, + compact_to_old, + freq, + } } /// Write `_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) -> 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) })); 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) /// `-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 { +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() + 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 = (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::>().join(" ")).unwrap(); - info!("IQ-TREE model file → {model_path} ({} of 16 states present in the alignment)", alphabet.k()); + 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 } @@ -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,17 +280,26 @@ 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 = seq.iter().map(|&b| { - if free_loss && b == b'-' { - return 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| { + let recoded: Vec = seq + .iter() + .map(|&b| { + if free_loss && b == b'-' { + return 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); }); @@ -283,7 +343,10 @@ pub(super) fn write_iqtree( // instability warnings). let refiltered; let alignment = if free_loss { - let low_freq_symbols: Vec = alphabet.compact_to_old.iter().zip(alphabet.freq.iter()) + let low_freq_symbols: Vec = 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::(), + low_freq_symbols + .iter() + .map(|&b| b as char) + .collect::(), ); 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![ - vec![b'A', b'-'], - vec![b'C', b'-'], - vec![b'-', b'G'], - ], + sequences: vec![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![ - vec![b'A', b'-'], - vec![b'C', b'-'], - vec![b'-', b'G'], - ], + sequences: vec![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![ - vec![b'A', b'-'], - vec![b'-', b'G'], - ], + sequences: vec![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![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!(csv.contains(",A,") && csv.contains(",C,"), "A/C must survive (well above threshold), got:\n{csv}"); + assert!( + !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();