fix: fit coalescent prior on final node times

kb/issues/M-timetree-skyline-timing-v0-divergence.md

@@ -1,10 +1,10 @@
 # Timetree skyline coalescent timing may diverge from v0
 
-v1 runs skyline coalescent optimization after the iteration loop exits, then re-runs timetree with the fitted skyline prior. v0 activates skyline on the last iteration within the loop, so the timetree inference runs once with the skyline prior during the loop itself.
+v1 re-fits the coalescent prior (constant or skyline) once per optimization round on that round's node times, so the output node times are conditioned on a prior fit on refined times. After the loop and any final marginal pass, it re-fits the prior once more on the final node times and reports that value, so the reported prior is the maximum-likelihood fit of the tree actually written out. v0 activates the skyline on the last iteration within the loop (constant Tc on earlier iterations), so its final node times are conditioned on the skyline prior during the loop itself.
 
-The v1 code at `packages/treetime/src/timetree/pipeline.rs` (skyline block after the optimizer loop) has a comment claiming "matching v0 behavior where skyline optimization happens only on the final iteration after times have converged." The v0 code (`packages/legacy/treetime/treetime/treetime.py` line 312) uses `if Tc == 'skyline' and niter < max_iter - 1: tmpTc = 'const'`, which means the last iteration (when `niter == max_iter - 1`) runs with `Tc = 'skyline'` inside the loop, not after it. The comment may be inaccurate.
+The v0 code (`packages/legacy/treetime/treetime/treetime.py` line 312) uses `if Tc == 'skyline' and niter < max_iter - 1: tmpTc = 'const'`, so the last iteration (`niter == max_iter - 1`) runs with `Tc = 'skyline'` inside the loop. v1 instead re-fits every round rather than switching from constant to skyline only on the last iteration.
 
-Both approaches produce self-consistent results. The difference is whether the final node times are conditioned on the skyline prior during the optimization iteration or via a separate post-loop pass. For datasets where the skyline shape deviates from constant Tc, final node time estimates may differ.
+Both approaches produce self-consistent results. The remaining difference is the constant-Tc-until-last-iteration schedule: v1 uses the skyline (or constant optimum) from round 1, v0 uses a constant Tc until the final iteration. For datasets where the skyline shape deviates from constant Tc, final node time estimates may differ.
 
 ## Investigation needed
 

packages/treetime/src/commands/timetree/__tests__/test_pipeline.rs

@@ -1,13 +1,17 @@
 #[cfg(test)]
 mod tests {
+  use crate::coalescent::total_lh::compute_coalescent_total_lh;
   use crate::commands::shared::alignment::AlignmentArgs;
   use crate::commands::shared::output::{LadderizeArg, OutputCoreArgs, TimetreeOutputSelection, TopologyOrderArgs};
   use crate::commands::timetree::args::TreetimeTimetreeArgs;
-  use crate::commands::timetree::run::run_timetree_estimation;
+  use crate::commands::timetree::initialization::load_input_data;
+  use crate::commands::timetree::run::{run_timetree_estimation, timetree_params_from_args};
   use crate::progress::NoopProgress;
+  use crate::timetree::pipeline::{self, TimetreeInput};
   use eyre::Report;
   use std::fs::read_to_string;
   use std::path::PathBuf;
+  use treetime_distribution::Distribution;
   use treetime_io::auspice_types::{AuspiceTree, AuspiceTreeNode};
   use treetime_utils::io::json::json_read_file;
 
@@ -115,6 +119,58 @@ mod tests {
     Ok(())
   }
 
+  #[test]
+  fn test_pipeline_coalescent_opt_prior_maximizes_output_tree() -> Result<(), Report> {
+    // B1 regression: the reported coalescent prior must be fit on the final output
+    // node times, so it must maximize the output tree's coalescent likelihood. A
+    // prior fit on the first, least-refined tree (the pre-fix behavior) would not.
+    let root = project_root();
+    let args = TreetimeTimetreeArgs {
+      alignment: AlignmentArgs {
+        alignment: vec![root.join("data/flu/h3n2/20/aln.fasta.xz")],
+      },
+      tree: Some(root.join("data/flu/h3n2/20/tree.nwk")),
+      metadata: Some(root.join("data/flu/h3n2/20/metadata.tsv")),
+      max_iter: 2,
+      coalescent_opt: true,
+      ..TreetimeTimetreeArgs::default()
+    };
+
+    let input_data = load_input_data(&args)?;
+    let params = timetree_params_from_args(&args);
+    let input = TimetreeInput {
+      graph: input_data.graph,
+      alphabet: input_data.alphabet,
+      sequences: input_data.aln,
+      dates: input_data.dates,
+    };
+
+    let output = pipeline::run(&params, input, None, &NoopProgress)?;
+
+    let tc_dist = output.coalescent_tc.expect("coalescent-opt must report a fitted Tc");
+    // A constant prior evaluates to the same Tc across its domain.
+    let tc = tc_dist.eval(2005.0)?;
+    assert!(
+      tc.is_finite() && tc > 0.0,
+      "reported Tc must be finite and positive, got {tc}"
+    );
+
+    // Optimality oracle (independent of optimize_tc): the coalescent log-likelihood
+    // L(Tc) = -M ln Tc - I/Tc is strictly unimodal, so the MLE of the output tree
+    // beats any perturbed Tc. This fails if the prior was fit on a different tree.
+    let lh_at = |scale: f64| compute_coalescent_total_lh(&output.graph, &Distribution::constant(tc * scale));
+    let lh_opt = lh_at(1.0)?;
+    let lh_lo = lh_at(0.8)?;
+    let lh_hi = lh_at(1.25)?;
+    assert!(
+      lh_opt >= lh_lo && lh_opt >= lh_hi,
+      "reported Tc={tc} must maximize the output-tree coalescent likelihood: \
+       L(Tc)={lh_opt}, L(0.8·Tc)={lh_lo}, L(1.25·Tc)={lh_hi}"
+    );
+
+    Ok(())
+  }
+
   fn project_root() -> PathBuf {
     PathBuf::from(env!("CARGO_MANIFEST_DIR"))
       .parent()

packages/treetime/src/commands/timetree/args.rs

@@ -122,7 +122,10 @@ pub struct TreetimeTimetreeArgs {
   ///
   /// When set, TreeTime will find the optimal Tc using Brent's method. This is similar
   /// to Python v0's `--coalescent=opt`, but uses a closed-form solution.
-  #[cfg_attr(feature = "clap", clap(long, conflicts_with = "coalescent", conflicts_with = "coalescent_skyline"))]
+  #[cfg_attr(
+    feature = "clap",
+    clap(long, conflicts_with = "coalescent", conflicts_with = "coalescent_skyline")
+  )]
   pub coalescent_opt: bool,
 
   /// Use skyline coalescent model instead of constant Tc.

packages/treetime/src/commands/timetree/run.rs

@@ -17,6 +17,42 @@ use std::path::PathBuf;
 use treetime_io::nwk::CommentProviders;
 use treetime_utils::io::file::create_file_or_stdout;
 
+/// Maps parsed CLI arguments onto the pipeline's [`TimetreeParams`].
+pub(crate) fn timetree_params_from_args(args: &TreetimeTimetreeArgs) -> TimetreeParams {
+  TimetreeParams {
+    model: args.model_args.model,
+    alphabet_name: args.alphabet_args.alphabet.unwrap_or_default(),
+    dense: args.dense,
+    gap_fill: args.gap_fill_args.effective_gap_fill(),
+    branch_length_mode: args.branch_length_mode,
+    no_indels: args.no_indels,
+    sequence_length: args.sequence_length,
+    clock_rate: args.clock_rate,
+    clock_std_dev: args.clock_std_dev,
+    keep_root: args.keep_root,
+    reroot_spec: args.reroot.spec(),
+    allow_negative_rate: args.allow_negative_rate,
+    clock_filter: args.clock_filter,
+    covariation: args.covariation,
+    tip_slack: args.tip_slack,
+    max_iter: args.max_iter,
+    resolve_polytomies: args.resolve_polytomies,
+    keep_polytomies: args.keep_polytomies,
+    relax: args.relax.clone(),
+    coalescent: args.coalescent,
+    coalescent_opt: args.coalescent_opt,
+    coalescent_skyline: args.coalescent_skyline,
+    n_skyline: args.n_skyline,
+    skyline_stiffness: args.skyline_stiffness,
+    n_branches_posterior: args.n_branches_posterior,
+    time_marginal: args.time_marginal,
+    confidence: args.confidence,
+    reconstruct_tip_states: args.reconstruct_tip_states,
+    report_ambiguous: args.report_ambiguous,
+    zero_based: args.zero_based,
+  }
+}
+
 pub fn run_timetree_estimation(
   args: &TreetimeTimetreeArgs,
   progress: &dyn crate::progress::ProgressSink,
@@ -52,38 +88,7 @@ pub fn run_timetree_estimation(
     None => None,
   };
 
-  let params = TimetreeParams {
-    model: args.model_args.model,
-    alphabet_name: args.alphabet_args.alphabet.unwrap_or_default(),
-    dense: args.dense,
-    gap_fill: args.gap_fill_args.effective_gap_fill(),
-    branch_length_mode: args.branch_length_mode,
-    no_indels: args.no_indels,
-    sequence_length: args.sequence_length,
-    clock_rate: args.clock_rate,
-    clock_std_dev: args.clock_std_dev,
-    keep_root: args.keep_root,
-    reroot_spec: args.reroot.spec(),
-    allow_negative_rate: args.allow_negative_rate,
-    clock_filter: args.clock_filter,
-    covariation: args.covariation,
-    tip_slack: args.tip_slack,
-    max_iter: args.max_iter,
-    resolve_polytomies: args.resolve_polytomies,
-    keep_polytomies: args.keep_polytomies,
-    relax: args.relax.clone(),
-    coalescent: args.coalescent,
-    coalescent_opt: args.coalescent_opt,
-    coalescent_skyline: args.coalescent_skyline,
-    n_skyline: args.n_skyline,
-    skyline_stiffness: args.skyline_stiffness,
-    n_branches_posterior: args.n_branches_posterior,
-    time_marginal: args.time_marginal,
-    confidence: args.confidence,
-    reconstruct_tip_states: args.reconstruct_tip_states,
-    report_ambiguous: args.report_ambiguous,
-    zero_based: args.zero_based,
-  };
+  let params = timetree_params_from_args(args);
 
   let input = TimetreeInput {
     graph: input_data.graph,
@@ -116,6 +121,7 @@ pub fn run_timetree_estimation(
     dates,
     gtr,
     model_name,
+    coalescent_tc: _,
   } = output;
   let mut graph = graph.map_data(TimetreeGraphData::new(
     clock_model,

packages/treetime/src/timetree/pipeline.rs

@@ -98,6 +98,11 @@ pub struct TimetreeOutput {
   pub gtr: Option<GTR>,
   #[serde(skip)]
   pub model_name: Option<GtrModelName>,
+  /// Coalescent prior Tc(t) fit on the final output node times, or `None` when the
+  /// coalescent model is disabled or fixed. This is the maximum-likelihood prior of
+  /// the tree that is written out, so it stays consistent with the reported node times.
+  #[serde(skip)]
+  pub coalescent_tc: Option<Distribution>,
 }
 
 pub fn run(
@@ -293,7 +298,12 @@ pub fn run(
       &format!("iteration {}/{max_iter}", i + 1),
     );
 
-    if coalescent.is_optimized() && i >= 2 {
+    // Re-fit the coalescent prior every round on the current node times. The
+    // previous `i >= 2` guard, combined with early convergence, could leave the
+    // whole run using the pre-loop prior (fit on the first, least-refined tree),
+    // so the reported prior disagreed with the output node times. v0 likewise
+    // re-fits inside the loop (constant Tc each round, skyline on the last).
+    if coalescent.is_optimized() {
       coalescent_tc = estimate_coalescent_tc(&coalescent, &input.graph, &skyline_params, coalescent_tc.as_ref())?;
     }
 
@@ -363,6 +373,15 @@ pub fn run(
     || rate_std.is_some())
   .then(|| extract_confidence_intervals(&input.graph));
 
+  // Report the coalescent prior as the maximum-likelihood fit on the final output
+  // node times. The optimization loop already conditioned the node times on a
+  // per-round prior; this last fit, after any final marginal pass, makes the
+  // *reported* prior consistent with the tree that is written out instead of the
+  // first, least-refined tree the pre-loop estimate saw.
+  if coalescent.is_optimized() {
+    coalescent_tc = estimate_coalescent_tc(&coalescent, &input.graph, &skyline_params, coalescent_tc.as_ref())?;
+  }
+
   progress.report("Done", 1.0, "");
   Ok(TimetreeOutput {
     graph: input.graph,
@@ -372,6 +391,7 @@ pub fn run(
     dates: input.dates,
     gtr: partition_gtr,
     model_name: partition_model_name,
+    coalescent_tc,
   })
 }