feat: emit reconstructed ancestral sequences from timetree

dev/run-smoke-tests

@@ -214,6 +214,10 @@ TIMETREE_BATCH_RUNS = [
   ("flu/h3n2/20",                 "flu/h3n2/20/branch-input",          "--branch-length-mode=input"),
   ("ebola/20",                    "ebola/20/branch-input",             "--branch-length-mode=input"),
 
+  # Tip-state reconstruction (reconstructed ancestral-sequence FASTA)
+  ("flu/h3n2/20",                 "flu/h3n2/20/reconstruct-tips",      "--reconstruct-tip-states"),
+  ("ebola/20",                    "ebola/20/reconstruct-tips",         "--include-leaves --impute-missing-data"),
+
   # Time marginal modes
   ("flu/h3n2/20",                 "flu/h3n2/20/marginal-always",       "--time-marginal=always"),
   ("flu/h3n2/200",                "flu/h3n2/200/marginal-always",      "--time-marginal=always"),

kb/features/timetree.md

@@ -134,7 +134,7 @@
 - [x] Clock model JSON with `timetree.*` basename
 - [x] Confidence TSV
 - [x] Coalescent time scale TSV/CSV/JSON (`--output-coalescent-{tsv,csv,json}`; per-segment $T_c$, $N_e$, and confidence band for skyline, optimized-constant, and fixed $T_c$; TSV in the default `--output-all` set, CSV and JSON opt-in; the `skyline.pdf` plot and the grid-points default remain, see [kb/issues/N-timetree-missing-skyline-output.md](../issues/N-timetree-missing-skyline-output.md))
-- [ ] Ancestral sequences FASTA ([kb/issues/N-timetree-ancestral-sequences-output-unimplemented.md](../issues/N-timetree-ancestral-sequences-output-unimplemented.md))
+- [x] Ancestral sequences FASTA (`--output-reconstructed-nuc-fasta`, default under `--output-all`; reuses the `ancestral` marginal reconstruction, internal nodes always, tips with `--include-leaves`)
 - [ ] Branch mutations table ([kb/issues/N-timetree-branch-mutations-output-unimplemented.md](../issues/N-timetree-branch-mutations-output-unimplemented.md))
 - [ ] Molecular clock text output or approved replacement ([kb/issues/N-timetree-molecular-clock-text-output-undecided.md](../issues/N-timetree-molecular-clock-text-output-undecided.md))
 - [ ] Sequence-evolution model text output or approved replacement ([kb/issues/N-timetree-sequence-evolution-model-text-output-undecided.md](../issues/N-timetree-sequence-evolution-model-text-output-undecided.md))
@@ -159,7 +159,7 @@
 - [ ] `--method-anc`
 - [ ] `--keep-overhangs` (gap handling not implemented)
 - [ ] `--zero-based` ([kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md](../issues/M-core-mutation-representation-and-format-projection-inconsistent.md))
-- [ ] `--include-leaves` / `--impute-missing-data` / `--reconstruct-tip-states` (flags accepted, alias enabling both, threaded to `TimetreeParams`; tip-state output is not wired into `timetree`: [kb/issues/M-timetree-tip-state-flags-not-wired.md](../issues/M-timetree-tip-state-flags-not-wired.md))
+- [x] `--include-leaves` / `--impute-missing-data` / `--reconstruct-tip-states` (drive the reconstructed FASTA: emit tips, impute ambiguous tip states, alias enabling both; see [kb/decisions/ancestral-marginal-tip-reconstruction-and-imputation.md](../decisions/ancestral-marginal-tip-reconstruction-and-imputation.md))
 - [ ] `--report-ambiguous`
 - [ ] `--seed`
 - [x] `--gen-per-year` (generations per year, default 50.0; reports effective population size `N_e = Tc * gen_per_year` to the log for the constant, opt, and skyline modes, and to the coalescent output files above)

kb/issues/N-clock-unnamed-root-after-reroot.md

@@ -0,0 +1,23 @@
+# Rerooted root and polytomy-resolution nodes stay unnamed in the clock command
+
+After rerooting, the new root node (and any nodes created during polytomy resolution) have no name. v0 assigns `NODE_XXXXXXX` names to unnamed internal nodes after every reroot via `_prepare_nodes()` (`treeanc.py:471-478`), called from `prepare_tree()` at the end of the `reroot()` method. v1's `assign_node_names` (`packages/treetime-graph/src/assign_node_names.rs:7`) runs during Newick parsing (`nwk.rs:100`) but is not called by the reroot path itself, so a command that reroots without a later naming pass emits the unnamed root.
+
+The `timetree` command no longer has this defect: it names every unnamed node once after the pipeline, before serialization (`packages/treetime/src/commands/timetree/run.rs`, immediately after `pipeline::run`). The `clock` command reroots (`packages/treetime/src/commands/clock/run.rs`) and still has no post-reroot naming pass.
+
+## Current state
+
+`create_new_root_node` (`packages/treetime/src/clock/reroot.rs:163`) creates a new node via `N::default()` which has `name=None`. In the `clock` command this node stays unnamed through output. The auspice writer assigns a fallback name `node_<key>` for display, but other consumers (TSV output, Newick annotations) see the unnamed state.
+
+## Impact
+
+- Clock `confidence_intervals.tsv` includes unnamed nodes with an empty name column
+- Clock Newick/Nexus output uses default formatting for the unnamed root
+- Differs from v0, where all internal nodes have `NODE_` names after any topology change
+
+## Fix
+
+Give the `clock` command a post-reroot naming pass, or move `assign_node_names` into the shared reroot path so every rerooting command names its new nodes, matching v0's `prepare_tree()` pattern.
+
+## Related tickets
+
+- [kb/tickets/clock-unnamed-root-and-polytomy-nodes-after-reroot.md](../tickets/clock-unnamed-root-and-polytomy-nodes-after-reroot.md)

kb/tickets/clock-unnamed-root-and-polytomy-nodes-after-reroot.md

@@ -0,0 +1,23 @@
+# Name the rerooted root and polytomy nodes in the clock command
+
+After rerooting, the new root node (and any nodes created during polytomy resolution) have no name. v0 assigns `NODE_XXXXXXX` names to unnamed internal nodes after every reroot via `_prepare_nodes()` (`treeanc.py:471-478`), called from `prepare_tree()` at the end of the `reroot()` method. v1's `assign_node_names` (`packages/treetime-graph/src/assign_node_names.rs:7`) runs during Newick parsing (`nwk.rs:100`) but is not called by the reroot path itself.
+
+The `timetree` command already names every unnamed node once after its pipeline (`packages/treetime/src/commands/timetree/run.rs`). Apply the equivalent to `clock`.
+
+## Current state
+
+`create_new_root_node` (`packages/treetime/src/clock/reroot.rs:163`) creates a new node via `N::default()` which has `name=None`. In the `clock` command this node stays unnamed through output. The auspice writer assigns a fallback name `node_<key>` for display, but other consumers (TSV output, Newick annotations) see the unnamed state.
+
+## Task
+
+Call `assign_node_names` in the `clock` command after rerooting completes, or move the call into the shared reroot path so every rerooting command names its new nodes uniformly, matching v0's `prepare_tree()` pattern. Prefer the shared path if it does not regress the `timetree` command's existing post-pipeline call.
+
+## Impact
+
+- Clock `confidence_intervals.tsv` includes unnamed nodes with an empty name column
+- Clock Newick/Nexus output uses default formatting for the unnamed root
+- Differs from v0, where all internal nodes have `NODE_` names after any topology change
+
+## Related issues
+
+- Source: [kb/issues/N-clock-unnamed-root-after-reroot.md](../issues/N-clock-unnamed-root-after-reroot.md) -- delete after full resolution

packages/app-cli/src/cli/pipeline/resolve.rs

@@ -479,7 +479,7 @@ mod tests {
     let result = resolve(config);
     assert_error!(
       result,
-      "step `anc` references `traits-csv` from step `tt`, which does not produce it; it produces: `augur-node-data`, `auspice`, `clock-model`, `coalescent-tsv`, `gtr`, `nexus`, `nwk`"
+      "step `anc` references `traits-csv` from step `tt`, which does not produce it; it produces: `augur-node-data`, `auspice`, `clock-model`, `coalescent-tsv`, `gtr`, `nexus`, `nwk`, `reconstructed-nuc-fasta`"
     );
   }
 

packages/app-server/src/args.rs

@@ -332,6 +332,7 @@ impl From<ServerTimetreeArgs> for TreetimeTimetreeArgs {
       divergence_units: DivergenceUnits::default(),
       output_augur_node_data: None,
       output_gtr: None,
+      output_reconstructed_nuc_fasta: None,
       output_clock_model: None,
       output_confidence_tsv: None,
       output_coalescent_tsv: None,

packages/treetime/src/commands/shared/output.rs

@@ -261,6 +261,7 @@ per_command_output_selection!(AncestralOutputSelection {
 per_command_output_selection!(TimetreeOutputSelection {
   AugurNodeData,
   Gtr,
+  ReconstructedNucFasta,
   ClockModel,
   ConfidenceTsv,
   Tracelog,
@@ -341,6 +342,7 @@ impl CommandKind {
       Self::Timetree => btreeset![
         AugurNodeData,
         Gtr,
+        ReconstructedNucFasta,
         ClockModel,
         ConfidenceTsv,
         Tracelog,

packages/treetime/src/commands/shared/resolve_outputs.rs

@@ -63,6 +63,10 @@ impl ResolveOutputs for TreetimeTimetreeArgs {
       &[
         (OutputSelection::AugurNodeData, self.output_augur_node_data.as_deref()),
         (OutputSelection::Gtr, self.output_gtr.as_deref()),
+        (
+          OutputSelection::ReconstructedNucFasta,
+          self.output_reconstructed_nuc_fasta.as_deref(),
+        ),
         (OutputSelection::ClockModel, self.output_clock_model.as_deref()),
         (OutputSelection::ConfidenceTsv, self.output_confidence_tsv.as_deref()),
         (OutputSelection::Tracelog, self.output_tracelog.as_deref()),

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

@@ -1 +1,2 @@
 mod test_pipeline;
+mod test_reconstructed_fasta;

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

@@ -0,0 +1,160 @@
+#[cfg(test)]
+mod tests {
+  use crate::alphabet::alphabet::Alphabet;
+  use eyre::Report;
+  use pretty_assertions::assert_eq;
+  use std::collections::BTreeSet;
+  use treetime_primitives::AsciiChar;
+
+  // Without a tip-state flag the reconstructed FASTA holds internal nodes only, and every record
+  // carries a real name. Rerooting introduces a fresh root node, so this also guards that the root
+  // reaches output named (`NODE_<n>`) rather than as an empty FASTA header.
+  #[test]
+  fn test_timetree_reconstructed_fasta_internal_only_and_named() -> Result<(), Report> {
+    let records = helpers::run_reconstructed_fasta("test-tt-recon-internal", |_| {})?;
+    let leaves = helpers::input_leaves()?;
+
+    assert!(!records.is_empty(), "reconstruction must emit internal-node sequences");
+    for record in &records {
+      assert!(
+        !record.seq_name.is_empty(),
+        "every reconstructed record must have a non-empty name (the rerooted root included)"
+      );
+      assert!(
+        !leaves.contains_key(&record.seq_name),
+        "without --include-leaves the FASTA must not contain leaf {}",
+        record.seq_name
+      );
+    }
+    assert!(
+      records.iter().any(|record| record.seq_name.starts_with("NODE_")),
+      "internal nodes carry NODE_<n> names, matching the ancestral command and v0"
+    );
+    Ok(())
+  }
+
+  // `--include-leaves` adds every leaf, and without imputation each reconstructed leaf echoes its
+  // observed input sequence exactly.
+  #[test]
+  fn test_timetree_reconstructed_fasta_include_leaves_echoes_observed() -> Result<(), Report> {
+    let records = helpers::run_reconstructed_fasta("test-tt-recon-leaves", |args| args.include_leaves = true)?;
+    let leaves = helpers::input_leaves()?;
+
+    let emitted_leaf_names: BTreeSet<String> = records
+      .iter()
+      .filter(|record| leaves.contains_key(&record.seq_name))
+      .map(|record| record.seq_name.clone())
+      .collect();
+    let expected_leaf_names: BTreeSet<String> = leaves.keys().cloned().collect();
+    assert_eq!(
+      expected_leaf_names, emitted_leaf_names,
+      "--include-leaves must emit exactly the observed leaves"
+    );
+
+    for record in &records {
+      if let Some(observed) = leaves.get(&record.seq_name) {
+        assert_eq!(
+          *observed, record.seq,
+          "leaf {} must echo its observed sequence when imputation is off",
+          record.seq_name
+        );
+      }
+    }
+    Ok(())
+  }
+
+  // `--reconstruct-tip-states` enables imputation: every ambiguous observed tip position resolves to
+  // an inferred canonical state, while gaps stay gaps.
+  #[test]
+  fn test_timetree_reconstructed_fasta_imputes_ambiguous_tips() -> Result<(), Report> {
+    let records = helpers::run_reconstructed_fasta("test-tt-recon-impute", |args| args.reconstruct_tip_states = true)?;
+    let leaves = helpers::input_leaves()?;
+    let alphabet = Alphabet::default();
+
+    let mut resolved = 0_usize;
+    for record in &records {
+      let Some(observed) = leaves.get(&record.seq_name) else {
+        continue;
+      };
+      let observed = observed.as_str().as_bytes();
+      let imputed = record.seq.as_str().as_bytes();
+      assert_eq!(observed.len(), imputed.len(), "imputation preserves sequence length");
+      for (obs, imp) in observed.iter().zip(imputed) {
+        let obs = AsciiChar::from_byte_unchecked(*obs);
+        let imp = AsciiChar::from_byte_unchecked(*imp);
+        if alphabet.is_gap(obs) {
+          assert!(alphabet.is_gap(imp), "gaps are inferred structure and stay gaps");
+        } else if alphabet.is_ambiguous(obs) || alphabet.is_unknown(obs) {
+          assert!(
+            !alphabet.is_ambiguous(imp) && !alphabet.is_unknown(imp),
+            "imputed tip position must resolve to a canonical state, got {}",
+            char::from(imp.inner())
+          );
+          resolved += 1;
+        }
+      }
+    }
+    assert!(
+      resolved > 0,
+      "the test dataset must contain ambiguous tip positions for imputation to resolve"
+    );
+    Ok(())
+  }
+
+  mod helpers {
+    use crate::alphabet::alphabet::Alphabet;
+    use crate::commands::shared::alignment::AlignmentArgs;
+    use crate::commands::timetree::args::TreetimeTimetreeArgs;
+    use crate::commands::timetree::run::run_timetree_estimation;
+    use crate::progress::NoopProgress;
+    use eyre::Report;
+    use std::collections::BTreeMap;
+    use std::path::PathBuf;
+    use treetime_io::fasta::{FastaRecord, read_many_fasta};
+    use treetime_primitives::Seq;
+
+    pub fn project_root() -> PathBuf {
+      PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .parent()
+        .and_then(|p| p.parent())
+        .map(PathBuf::from)
+        .expect("project has workspace root")
+    }
+
+    pub fn input_leaves() -> Result<BTreeMap<String, Seq>, Report> {
+      let alignment = project_root().join("data/flu/h3n2/20/aln.fasta.xz");
+      let records = read_many_fasta(&[alignment], &Alphabet::default())?;
+      Ok(
+        records
+          .into_iter()
+          .map(|record| (record.seq_name, record.seq))
+          .collect(),
+      )
+    }
+
+    pub fn run_reconstructed_fasta(
+      subdir: &str,
+      configure: impl FnOnce(&mut TreetimeTimetreeArgs),
+    ) -> Result<Vec<FastaRecord>, Report> {
+      let root = project_root();
+      let outdir = root.join("tmp").join(subdir);
+      std::fs::create_dir_all(&outdir)?;
+      let fasta = outdir.join("ancestral_sequences.fasta");
+
+      let mut 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,
+        output_reconstructed_nuc_fasta: Some(fasta.clone()),
+        ..TreetimeTimetreeArgs::default()
+      };
+      configure(&mut args);
+
+      run_timetree_estimation(&args, &NoopProgress)?;
+      read_many_fasta(&[fasta], &Alphabet::default())
+    }
+  }
+}