feat: toggle optimize topology-cleanup steps

kb/features/optimize.md

@@ -66,6 +66,7 @@ Once per iteration, `prune_and_merge_in_loop` simplifies the tree. Sparse partit
 - [x] Shared-mutation merge: siblings in a polytomy that carry identical substitutions are grouped under a new internal node (`merge_shared_mutation_branches`).
 - [x] Reversion hoist: when a child edge reverts a substitution on the node's parent edge, a new node is inserted grouping that child with its sibling subtree, lifting the non-reverted substitutions above it. Removes one mutation per reverted position and adds none; branch lengths split proportionally to preserve root-to-node distances ([kb/decisions/optimize-polytomy-reversion-resolution.md](../decisions/optimize-polytomy-reversion-resolution.md)).
 - [x] Helper-node retirement: mutation-free edges to nodes created during the routine are collapsed, dissolving the transient helpers the merge and hoist leave behind.
+- [x] Per-step toggles: `--no-collapse-short-branches`, `--no-merge-siblings`, `--no-flip-parent-child` disable individual steps (all default on, so output is unchanged without them). Carried in-memory as `TopologyOps` on `OptimizeParams`. Merge and flip-parent-child are sparse-only, so both warn when disabled under `--dense` where they cannot run. When collapse is disabled the optimizer may leave zero-length internal edges in the output; they are kept as-is.
 
 The merge, hoist, and retire steps form one per-polytomy routine (`resolve_polytomies`) run every iteration and driven to a fixpoint by a monotone (mutation count, node count) potential. A polytomy whose children all revert the same position collapses to a single reverting child that the two-children guard leaves in place ([kb/issues/M-optimize-reversion-hoist-single-child-residual.md](../issues/M-optimize-reversion-hoist-single-child-residual.md)).
 

packages/app-server/src/args.rs

@@ -10,7 +10,7 @@ use treetime::commands::optimize::args::OptimizeRerootMethod;
 use treetime::commands::shared::reroot::RerootArgs;
 use treetime::commands::timetree::args::TimeMarginalMode;
 use treetime::gtr::get_gtr::GtrModelName;
-use treetime::optimize::params::{BranchLengthMode, BranchOptMethod, InitialGuessMode};
+use treetime::optimize::params::{BranchLengthMode, BranchOptMethod, InitialGuessMode, TopologyOps};
 use treetime::seq::gap_fill::GapFill;
 
 use app_api::{
@@ -483,6 +483,7 @@ impl From<ServerOptimizeArgs> for TreetimeOptimizeArgs {
       reroot: s.reroot,
       reroot_tips: s.reroot_tips,
       keep_root: s.keep_root,
+      topology_ops: TopologyOps::default(),
       divergence_units: DivergenceUnits::default(),
       gap_fill_args: GapFillArgs {
         gap_fill: s.gap_fill,

packages/schemas/input-config-optimize.schema.json

@@ -278,6 +278,21 @@
       "type": "boolean",
       "default": false
     },
+    "no_collapse_short_branches": {
+      "description": "Disable collapsing of internal branches whose optimized length is zero.\n\nBy default the optimize loop contracts internal edges the per-edge optimizer drove to\nexactly zero that carry no substitutions or indels, turning the resulting binary nodes into\npolytomies. When set, such edges are kept in the output tree with length zero.",
+      "type": "boolean",
+      "default": false
+    },
+    "no_merge_siblings": {
+      "description": "Disable merging of polytomy siblings that share substitutions.\n\nBy default, sibling branches in a polytomy that carry identical substitutions are grouped\nunder a new internal node. Requires the sparse sequence representation, so this step has no\neffect under `--dense` and the flag is then a no-op.",
+      "type": "boolean",
+      "default": false
+    },
+    "no_flip_parent_child": {
+      "description": "Disable the flip-parent-child step (reversion hoist) in polytomies.\n\nBy default, when a child branch carries the exact reversion of a substitution on the node's\nparent branch, a new node is inserted that groups the node with that child, removing one\nmutation per reverted position. When set, those reversions are left in place. Requires the\nsparse sequence representation, so this step has no effect under `--dense` and the flag is\nthen a no-op.",
+      "type": "boolean",
+      "default": false
+    },
     "gap_fill": {
       "description": "How to handle gap characters in input sequences\n\n'only-terminal': replace leading and trailing gap characters with the ambiguous character (default, matches v0).\n'all': replace all gap characters with the ambiguous character.\n'none': leave all gap characters unchanged.",
       "$ref": "#/$defs/GapFill",

packages/schemas/input-config-pipeline.schema.json

@@ -1621,6 +1621,45 @@
             }
           ]
         },
+        "no_collapse_short_branches": {
+          "anyOf": [
+            {
+              "description": "Disable collapsing of internal branches whose optimized length is zero.\n\nBy default the optimize loop contracts internal edges the per-edge optimizer drove to\nexactly zero that carry no substitutions or indels, turning the resulting binary nodes into\npolytomies. When set, such edges are kept in the output tree with length zero.",
+              "type": "boolean",
+              "default": false
+            },
+            {
+              "type": "string",
+              "pattern": "\\{\\{.*\\}\\}"
+            }
+          ]
+        },
+        "no_merge_siblings": {
+          "anyOf": [
+            {
+              "description": "Disable merging of polytomy siblings that share substitutions.\n\nBy default, sibling branches in a polytomy that carry identical substitutions are grouped\nunder a new internal node. Requires the sparse sequence representation, so this step has no\neffect under `--dense` and the flag is then a no-op.",
+              "type": "boolean",
+              "default": false
+            },
+            {
+              "type": "string",
+              "pattern": "\\{\\{.*\\}\\}"
+            }
+          ]
+        },
+        "no_flip_parent_child": {
+          "anyOf": [
+            {
+              "description": "Disable the flip-parent-child step (reversion hoist) in polytomies.\n\nBy default, when a child branch carries the exact reversion of a substitution on the node's\nparent branch, a new node is inserted that groups the node with that child, removing one\nmutation per reverted position. When set, those reversions are left in place. Requires the\nsparse sequence representation, so this step has no effect under `--dense` and the flag is\nthen a no-op.",
+              "type": "boolean",
+              "default": false
+            },
+            {
+              "type": "string",
+              "pattern": "\\{\\{.*\\}\\}"
+            }
+          ]
+        },
         "gap_fill": {
           "description": "How to handle gap characters in input sequences\n\n'only-terminal': replace leading and trailing gap characters with the ambiguous character (default, matches v0).\n'all': replace all gap characters with the ambiguous character.\n'none': leave all gap characters unchanged.",
           "$ref": "#/$defs/GapFill",

packages/treetime/src/commands/optimize/__tests__/test_augur_node_data.rs

@@ -127,7 +127,7 @@ mod tests {
   fn test_augur_node_data_optimize_end_to_end() {
     use crate::alphabet::alphabet::Alphabet;
     use crate::gtr::get_gtr::GtrModelName;
-    use crate::optimize::params::{BranchOptMethod, InitialGuessMode};
+    use crate::optimize::params::{BranchOptMethod, InitialGuessMode, TopologyOps};
     use crate::optimize::pipeline::{self, OptimizeInput, OptimizeParams};
     use crate::progress::NoopProgress;
     use treetime_io::fasta::read_many_fasta;
@@ -148,6 +148,7 @@ mod tests {
       initial_guess: InitialGuessMode::default(),
       no_indels: false,
       reroot_spec: None,
+      topology_ops: TopologyOps::default(),
     };
 
     let input = OptimizeInput {

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

@@ -6,7 +6,7 @@ use crate::commands::shared::gap_fill::GapFillArgs;
 use crate::commands::shared::model::ModelArgs;
 use crate::commands::shared::output::{DivergenceUnits, OptimizeOutputSelection, OutputCoreArgs, TopologyOrderArgs};
 use crate::commands::shared::required::missing_required_args;
-use crate::optimize::params::{BranchOptMethod, InitialGuessMode};
+use crate::optimize::params::{BranchOptMethod, InitialGuessMode, TopologyOps};
 #[cfg(feature = "clap")]
 use clap::ValueHint;
 use eyre::Report;
@@ -196,6 +196,32 @@ pub struct TreetimeOptimizeArgsRaw {
   #[cfg_attr(feature = "clap", clap(long, conflicts_with_all = ["reroot", "reroot_tips"]))]
   pub keep_root: bool,
 
+  /// Disable collapsing of internal branches whose optimized length is zero.
+  ///
+  /// By default the optimize loop contracts internal edges the per-edge optimizer drove to
+  /// exactly zero that carry no substitutions or indels, turning the resulting binary nodes into
+  /// polytomies. When set, such edges are kept in the output tree with length zero.
+  #[cfg_attr(feature = "clap", clap(long))]
+  pub no_collapse_short_branches: bool,
+
+  /// Disable merging of polytomy siblings that share substitutions.
+  ///
+  /// By default, sibling branches in a polytomy that carry identical substitutions are grouped
+  /// under a new internal node. Requires the sparse sequence representation, so this step has no
+  /// effect under `--dense` and the flag is then a no-op.
+  #[cfg_attr(feature = "clap", clap(long))]
+  pub no_merge_siblings: bool,
+
+  /// Disable the flip-parent-child step (reversion hoist) in polytomies.
+  ///
+  /// By default, when a child branch carries the exact reversion of a substitution on the node's
+  /// parent branch, a new node is inserted that groups the node with that child, removing one
+  /// mutation per reverted position. When set, those reversions are left in place. Requires the
+  /// sparse sequence representation, so this step has no effect under `--dense` and the flag is
+  /// then a no-op.
+  #[cfg_attr(feature = "clap", clap(long))]
+  pub no_flip_parent_child: bool,
+
   #[cfg_attr(feature = "clap", clap(flatten))]
   #[serde(flatten)]
   pub gap_fill_args: GapFillArgs,
@@ -227,6 +253,7 @@ pub struct TreetimeOptimizeArgs {
   pub reroot: Option<OptimizeRerootMethod>,
   pub reroot_tips: Vec<String>,
   pub keep_root: bool,
+  pub topology_ops: TopologyOps,
   pub gap_fill_args: GapFillArgs,
 }
 
@@ -285,6 +312,11 @@ impl TryFrom<TreetimeOptimizeArgsRaw> for TreetimeOptimizeArgs {
       reroot: raw.reroot,
       reroot_tips: raw.reroot_tips,
       keep_root: raw.keep_root,
+      topology_ops: TopologyOps {
+        collapse_short_branches: !raw.no_collapse_short_branches,
+        merge_siblings: !raw.no_merge_siblings,
+        flip_parent_child: !raw.no_flip_parent_child,
+      },
       gap_fill_args: raw.gap_fill_args,
     })
   }

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

@@ -45,6 +45,7 @@ pub fn run_optimize(
     initial_guess: args.branch_length_initial_guess,
     no_indels: args.no_indels,
     reroot_spec: args.reroot_spec(),
+    topology_ops: args.topology_ops,
   };
 
   let input = OptimizeInput {

packages/treetime/src/optimize/__tests__/test_convergence_conditions.rs

@@ -4,7 +4,7 @@ mod tests {
     TREE_NEWICK, setup_partitions, simple_alignment,
   };
   use crate::optimize::iteration::{DAMPING_FLOOR, apply_damping, restore_branch_lengths, save_branch_lengths};
-  use crate::optimize::params::BranchOptMethod;
+  use crate::optimize::params::{BranchOptMethod, TopologyOps};
   use crate::optimize::run_loop::{ConvergenceReason, collect_optimize_partitions, run_optimize_loop};
   use crate::payload::ancestral::GraphAncestral;
   use approx::assert_abs_diff_eq;
@@ -117,6 +117,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       false,
+      TopologyOps::default(),
     )?;
 
     let (iter, reason) = result.stopped_at.expect("loop should have stopped");
@@ -147,6 +148,7 @@ mod tests {
       0.0,
       BranchOptMethod::BrentSqrt,
       false,
+      TopologyOps::default(),
     )?;
 
     match result.stopped_at {
@@ -191,6 +193,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       false,
+      TopologyOps::default(),
     )?;
 
     let (iter, reason) = result.stopped_at.expect("loop should have stopped");
@@ -223,6 +226,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       false,
+      TopologyOps::default(),
     )?;
 
     assert_eq!(result.lh_history.len(), 2);
@@ -255,6 +259,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       false,
+      TopologyOps::default(),
     )?;
 
     assert!(

packages/treetime/src/optimize/__tests__/test_convergence_sc2.rs

@@ -6,7 +6,7 @@ mod tests {
   use crate::ancestral::marginal::update_marginal;
   use crate::gtr::get_gtr::{JC69Params, jc69};
   use crate::optimize::dispatch::initial_guess_mixed;
-  use crate::optimize::params::BranchOptMethod;
+  use crate::optimize::params::{BranchOptMethod, TopologyOps};
   use crate::optimize::run_loop::{collect_optimize_partitions, run_optimize_loop};
 
   use eyre::Report;
@@ -58,6 +58,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       false,
+      TopologyOps::default(),
     )?;
 
     assert!(
@@ -106,6 +107,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       false,
+      TopologyOps::default(),
     )?;
 
     assert!(

packages/treetime/src/optimize/__tests__/test_damping.rs

@@ -4,7 +4,7 @@ mod tests {
     TREE_NEWICK, compute_total_lh, setup_partitions, simple_alignment,
   };
   use crate::optimize::iteration::{apply_damping, save_branch_lengths};
-  use crate::optimize::params::BranchOptMethod;
+  use crate::optimize::params::{BranchOptMethod, TopologyOps};
   use crate::optimize::run_loop::run_optimize_loop;
   use crate::payload::ancestral::GraphAncestral;
   use approx::assert_abs_diff_eq;
@@ -171,6 +171,7 @@ mod tests {
       damping,
       method,
       false,
+      TopologyOps::default(),
     )?;
 
     assert!(
@@ -225,6 +226,7 @@ mod tests {
       0.75,
       method,
       false,
+      TopologyOps::default(),
     )?;
 
     // Strict non-regression: damped optimization must not degrade likelihood.

packages/treetime/src/optimize/__tests__/test_gm_optimize.rs

@@ -176,7 +176,7 @@ mod tests {
     use crate::ancestral::marginal::{initialize_marginal, update_marginal};
     use crate::gtr::get_gtr::{JC69Params, jc69};
     use crate::optimize::dispatch::initial_guess_mixed;
-    use crate::optimize::params::BranchOptMethod;
+    use crate::optimize::params::{BranchOptMethod, TopologyOps};
     use crate::optimize::run_loop::{collect_optimize_partitions, run_optimize_loop};
     use crate::partition::marginal::dense::partition::PartitionMarginalDense;
     use crate::seq::alignment::get_common_length;
@@ -273,6 +273,7 @@ mod tests {
         case.damping,
         method,
         false,
+        TopologyOps::default(),
       )?;
 
       // Append a trailing likelihood measurement so `lh_history.last()` reflects the state

packages/treetime/src/optimize/__tests__/test_no_indels.rs

@@ -12,7 +12,7 @@ mod tests {
     inject_indels_on_first_edge, setup_identical_partitions,
   };
   use crate::optimize::dispatch::run_optimize_mixed_inner;
-  use crate::optimize::params::{BranchOptMethod, InitialGuessMode};
+  use crate::optimize::params::{BranchOptMethod, InitialGuessMode, TopologyOps};
   use crate::optimize::run_loop::apply_initial_guess_mode;
   use crate::optimize::run_loop::run_optimize_loop;
   use crate::payload::ancestral::GraphAncestral;
@@ -69,6 +69,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       false,
+      TopologyOps::default(),
     )?;
 
     let result_without = run_optimize_loop(
@@ -81,6 +82,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       true,
+      TopologyOps::default(),
     )?;
 
     assert!(
@@ -150,6 +152,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       false,
+      TopologyOps::default(),
     )?;
 
     let result_flag = run_optimize_loop(
@@ -162,6 +165,7 @@ mod tests {
       0.75,
       BranchOptMethod::BrentSqrt,
       true,
+      TopologyOps::default(),
     )?;
 
     assert_eq!(

packages/treetime/src/optimize/__tests__/test_pipeline_gtr_normalized.rs

@@ -2,7 +2,7 @@
 mod tests {
   use crate::alphabet::alphabet::Alphabet;
   use crate::gtr::get_gtr::GtrModelName;
-  use crate::optimize::params::{BranchOptMethod, InitialGuessMode};
+  use crate::optimize::params::{BranchOptMethod, InitialGuessMode, TopologyOps};
   use crate::optimize::pipeline::{OptimizeInput, OptimizeParams, run};
   use crate::partition::traits::HasGtr;
   use crate::payload::ancestral::GraphAncestral;
@@ -44,6 +44,7 @@ mod tests {
       initial_guess: InitialGuessMode::default(),
       no_indels: false,
       reroot_spec: None,
+      topology_ops: TopologyOps::default(),
     };
 
     let input = OptimizeInput {

packages/treetime/src/optimize/__tests__/test_pipeline_reroot.rs

@@ -3,7 +3,7 @@ mod tests {
   use crate::alphabet::alphabet::Alphabet;
   use crate::clock::find_best_root::params::{RerootMethod, RerootSpec};
   use crate::gtr::get_gtr::GtrModelName;
-  use crate::optimize::params::{BranchOptMethod, InitialGuessMode};
+  use crate::optimize::params::{BranchOptMethod, InitialGuessMode, TopologyOps};
   use crate::optimize::pipeline::{OptimizeInput, OptimizeParams, run};
   use crate::payload::ancestral::GraphAncestral;
   use crate::progress::NoopProgress;
@@ -37,6 +37,7 @@ mod tests {
       initial_guess: InitialGuessMode::default(),
       no_indels: false,
       reroot_spec,
+      topology_ops: TopologyOps::default(),
     }
   }