diff --git a/docs/docs/reference.md b/docs/docs/reference.md index d255b8d98..f4348e8e3 100644 --- a/docs/docs/reference.md +++ b/docs/docs/reference.md @@ -31,6 +31,7 @@ If you have TreeTime CLI installed, you can type `treetime --help` to read the l * [`treetime clock`↴](#treetime-clock) * [`treetime homoplasy`↴](#treetime-homoplasy) * [`treetime mugration`↴](#treetime-mugration) +* [`treetime pipeline`↴](#treetime-pipeline) * [`treetime arg`↴](#treetime-arg) * [`treetime schema`↴](#treetime-schema) * [`treetime help-markdown`↴](#treetime-help-markdown) @@ -54,6 +55,7 @@ Publication: https://academic.oup.com/ve/article/4/1/vex042/4794731 * `clock` — Calculates the root-to-tip regression and quantifies the 'clock-i-ness' of the tree. It will reroot the tree to maximize the clock-like signal and recalculate branch length unless run with --keep_root * `homoplasy` — Reconstructs ancestral sequences and maps mutations to the tree. The tree is then scanned for homoplasies. An excess number of homoplasies might suggest contamination, recombination, culture adaptation or similar * `mugration` — Reconstructs discrete ancestral states, for example geographic location, host, or similar. In addition to ancestral states, a GTR model of state transitions is inferred +* `pipeline` — Runs an ordered list of analysis commands from one config file, on one dataset, in one process * `arg` — Estimates ancestral reassortment graph (ARG) * `schema` — Write JSON Schema definitions for TreeTime data types * `help-markdown` — Print CLI reference documentation in Markdown format @@ -105,6 +107,11 @@ Estimates time trees from an initial tree topology, a set of date constraints (e ###### **Options:** +* `--config ` — Load command configuration from a file (optionally compressed: `.gz`, `.bz2`, `.xz`, `.zst`). + + The file holds this command's configuration object, in the same shape the command serializes to. It is parsed as YAML, which also accepts JSON. Use `-` to read from stdin. + + Explicit command-line flags override values from the file; the file overrides defaults. A boolean enabled in the config cannot be disabled from the command line (a flag has no `false` spelling); edit the config instead. * `-a`, `--alignment ` [alias: `aln`] — Path to one or multiple FASTA files with aligned input sequences Accepts plain or compressed FASTA files. If a compressed fasta file is provided, it will be transparently decompressed. Supported compression formats: `gz`, `bz2`, `xz`, `zstd`. Decompressor is chosen based on file extension. If there's multiple input files, then different files can have different compression formats. @@ -247,7 +254,9 @@ Estimates time trees from an initial tree topology, a set of date constraints (e Possible values: `only-terminal`, `all`, `none` * `--zero-based` — Zero-based mutation indexing -* `--reconstruct-tip-states` — Overwrite ambiguous states on tips with the most likely inferred state +* `--include-leaves` — Emit reconstructed leaf (tip) sequences in addition to internal nodes +* `--impute-missing-data` — Resolve ambiguous and unknown tip states (`N` and IUPAC codes such as `R`) to the most likely inferred state. Gaps are left as deletions +* `--reconstruct-tip-states` — v0-compatible alias for `--include-leaves --impute-missing-data` * `--report-ambiguous` — Include transitions involving ambiguous states * `--no-indels` — Disable indel (insertion/deletion) contributions to branch-length optimization and branch-length distributions. @@ -267,6 +276,11 @@ Estimates time trees from an initial tree topology, a set of date constraints (e Takes precedence over paths configured with `--output-all` and `--output-selection`. * `--output-gtr ` — Path to output GTR model JSON. + Takes precedence over paths configured with `--output-all` and `--output-selection`. +* `--output-reconstructed-nuc-fasta ` — Path to output reconstructed ancestral-sequence nucleotide FASTA. + + The v1 equivalent of TreeTime v0's `ancestral_sequences.fasta`: internal-node sequences reconstructed by the marginal pass, plus reconstructed tip sequences when `--include-leaves` (or `--reconstruct-tip-states`) is set. `--impute-missing-data` resolves ambiguous tip states. + Takes precedence over paths configured with `--output-all` and `--output-selection`. * `--output-clock-model ` — Path to output clock model JSON. @@ -364,7 +378,7 @@ Estimates time trees from an initial tree topology, a set of date constraints (e Restricts which outputs `--output-all` writes. Special value `all` expands to every output available for this command. Requires `--output-all`. Per-file flags are always honored regardless of this selection. - Possible values: `all`, `nwk`, `nexus`, `auspice`, `phyloxml`, `phyloxml-json`, `mat-pb`, `mat-json`, `graph-json`, `dot`, `augur-node-data`, `gtr`, `clock-model`, `confidence-tsv`, `tracelog`, `coalescent-tsv`, `coalescent-csv`, `coalescent-json` + Possible values: `all`, `nwk`, `nexus`, `auspice`, `phyloxml`, `phyloxml-json`, `mat-pb`, `mat-json`, `graph-json`, `dot`, `augur-node-data`, `gtr`, `reconstructed-nuc-fasta`, `clock-model`, `confidence-tsv`, `tracelog`, `coalescent-tsv`, `coalescent-csv`, `coalescent-json` * `--ladderize ` — Order tree topology before writing output files @@ -398,10 +412,15 @@ Estimates time trees from an initial tree topology, a set of date constraints (e Optimizes the branch lengths and likelihood of a phylogenetic tree given aligned sequences -**Usage:** `treetime optimize [OPTIONS] --tree ` +**Usage:** `treetime optimize [OPTIONS]` ###### **Options:** +* `--config ` — Load command configuration from a file (optionally compressed: `.gz`, `.bz2`, `.xz`, `.zst`). + + The file holds this command's configuration object, in the same shape the command serializes to. It is parsed as YAML, which also accepts JSON. Use `-` to read from stdin. + + Explicit command-line flags override values from the file; the file overrides defaults. A boolean enabled in the config cannot be disabled from the command line (a flag has no `false` spelling); edit the config instead. * `-a`, `--alignment ` [alias: `aln`] — Path to one or multiple FASTA files with aligned input sequences Accepts plain or compressed FASTA files. If a compressed fasta file is provided, it will be transparently decompressed. Supported compression formats: `gz`, `bz2`, `xz`, `zstd`. Decompressor is chosen based on file extension. If there's multiple input files, then different files can have different compression formats. @@ -621,6 +640,15 @@ Optimizes the branch lengths and likelihood of a phylogenetic tree given aligned * `--keep-root` — Keep the input tree root instead of rerooting. Optimize keeps the input root by default; this flag is the explicit form and is mutually exclusive with the reroot options. +* `--no-collapse-short-branches` — 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. +* `--no-merge-siblings` — 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. +* `--no-flip-parent-child` — 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. * `--gap-fill ` — How to handle gap characters in input sequences 'only-terminal': replace leading and trailing gap characters with the ambiguous character (default, matches v0). 'all': replace all gap characters with the ambiguous character. 'none': leave all gap characters unchanged. @@ -636,10 +664,15 @@ Optimizes the branch lengths and likelihood of a phylogenetic tree given aligned Prunes short branches and/or branches without mutations from a phylogenetic tree -**Usage:** `treetime prune [OPTIONS] --tree ` +**Usage:** `treetime prune [OPTIONS]` ###### **Options:** +* `--config ` — Load command configuration from a file (optionally compressed: `.gz`, `.bz2`, `.xz`, `.zst`). + + The file holds this command's configuration object, in the same shape the command serializes to. It is parsed as YAML, which also accepts JSON. Use `-` to read from stdin. + + Explicit command-line flags override values from the file; the file overrides defaults. A boolean enabled in the config cannot be disabled from the command line (a flag has no `false` spelling); edit the config instead. * `-a`, `--alignment ` [alias: `aln`] — Path to one or multiple FASTA files with aligned input sequences Accepts plain or compressed FASTA files. If a compressed fasta file is provided, it will be transparently decompressed. Supported compression formats: `gz`, `bz2`, `xz`, `zstd`. Decompressor is chosen based on file extension. If there's multiple input files, then different files can have different compression formats. @@ -799,10 +832,15 @@ Prunes short branches and/or branches without mutations from a phylogenetic tree Reconstructs ancestral sequences and maps mutations to the tree. The output consists of a file 'ancestral.fasta' with ancestral sequences and a tree 'ancestral.nexus' with mutations added as comments like A45G,G136T,..., number in SNPs used 1-based index by default. The inferred GTR model is written to stdout -**Usage:** `treetime ancestral [OPTIONS] --tree ` +**Usage:** `treetime ancestral [OPTIONS]` ###### **Options:** +* `--config ` — Load command configuration from a file (optionally compressed: `.gz`, `.bz2`, `.xz`, `.zst`). + + The file holds this command's configuration object, in the same shape the command serializes to. It is parsed as YAML, which also accepts JSON. Use `-` to read from stdin. + + Explicit command-line flags override values from the file; the file overrides defaults. A boolean enabled in the config cannot be disabled from the command line (a flag has no `false` spelling); edit the config instead. * `-a`, `--alignment ` [alias: `aln`] — Path to one or multiple FASTA files with aligned input sequences Accepts plain or compressed FASTA files. If a compressed fasta file is provided, it will be transparently decompressed. Supported compression formats: `gz`, `bz2`, `xz`, `zstd`. Decompressor is chosen based on file extension. If there's multiple input files, then different files can have different compression formats. @@ -863,7 +901,13 @@ Reconstructs ancestral sequences and maps mutations to the tree. The output cons Possible values: `only-terminal`, `all`, `none` * `--zero-based` — Zero-based mutation indexing -* `--reconstruct-tip-states` — Overwrite ambiguous states on tips with the most likely inferred state +* `--include-leaves` — Emit reconstructed leaf (tip) sequences in addition to internal nodes +* `--impute-missing-data` — Resolve ambiguous and unknown tip states (`N` and IUPAC codes such as `R`) to the most likely inferred state. + + Gaps are left as deletions (inferred structure, not missing data). Only defined for marginal reconstruction; a no-op with a warning under `--method-anc=parsimony`. +* `--reconstruct-tip-states` — v0-compatible alias for `--include-leaves --impute-missing-data`. + + Emits tip sequences and resolves ambiguous/unknown tip states to the most likely inferred state. * `--report-ambiguous` — Include transitions involving ambiguous states * `--ignore-missing-alns` — Treat tree tips that have no sequence in the alignment as fully ambiguous (missing data) instead of aborting. @@ -1026,10 +1070,15 @@ Reconstructs ancestral sequences and maps mutations to the tree. The output cons Calculates the root-to-tip regression and quantifies the 'clock-i-ness' of the tree. It will reroot the tree to maximize the clock-like signal and recalculate branch length unless run with --keep_root -**Usage:** `treetime clock [OPTIONS] --metadata ` +**Usage:** `treetime clock [OPTIONS]` ###### **Options:** +* `--config ` — Load command configuration from a file (optionally compressed: `.gz`, `.bz2`, `.xz`, `.zst`). + + The file holds this command's configuration object, in the same shape the command serializes to. It is parsed as YAML, which also accepts JSON. Use `-` to read from stdin. + + Explicit command-line flags override values from the file; the file overrides defaults. A boolean enabled in the config cannot be disabled from the command line (a flag has no `false` spelling); edit the config instead. * `-a`, `--alignment ` [alias: `aln`] — Path to one or multiple FASTA files with aligned input sequences Accepts plain or compressed FASTA files. If a compressed fasta file is provided, it will be transparently decompressed. Supported compression formats: `gz`, `bz2`, `xz`, `zstd`. Decompressor is chosen based on file extension. If there's multiple input files, then different files can have different compression formats. @@ -1255,10 +1304,15 @@ Calculates the root-to-tip regression and quantifies the 'clock-i-ness' of the t Reconstructs ancestral sequences and maps mutations to the tree. The tree is then scanned for homoplasies. An excess number of homoplasies might suggest contamination, recombination, culture adaptation or similar -**Usage:** `treetime homoplasy [OPTIONS] --tree ` +**Usage:** `treetime homoplasy [OPTIONS]` ###### **Options:** +* `--config ` — Load command configuration from a file (optionally compressed: `.gz`, `.bz2`, `.xz`, `.zst`). + + The file holds this command's configuration object, in the same shape the command serializes to. It is parsed as YAML, which also accepts JSON. Use `-` to read from stdin. + + Explicit command-line flags override values from the file; the file overrides defaults. A boolean enabled in the config cannot be disabled from the command line (a flag has no `false` spelling); edit the config instead. * `-a`, `--alignment ` [alias: `aln`] — Path to one or multiple FASTA files with aligned input sequences Accepts plain or compressed FASTA files. If a compressed fasta file is provided, it will be transparently decompressed. Supported compression formats: `gz`, `bz2`, `xz`, `zstd`. Decompressor is chosen based on file extension. If there's multiple input files, then different files can have different compression formats. @@ -1319,7 +1373,13 @@ Reconstructs ancestral sequences and maps mutations to the tree. The tree is the Possible values: `only-terminal`, `all`, `none` * `--zero-based` — Zero-based mutation indexing -* `--reconstruct-tip-states` — Overwrite ambiguous states on tips with the most likely inferred state +* `--include-leaves` — Emit reconstructed leaf (tip) sequences in addition to internal nodes +* `--impute-missing-data` — Resolve ambiguous and unknown tip states (`N` and IUPAC codes such as `R`) to the most likely inferred state. + + Gaps are left as deletions (inferred structure, not missing data). Only defined for marginal reconstruction; a no-op with a warning under `--method-anc=parsimony`. +* `--reconstruct-tip-states` — v0-compatible alias for `--include-leaves --impute-missing-data`. + + Emits tip sequences and resolves ambiguous/unknown tip states to the most likely inferred state. * `--report-ambiguous` — Include transitions involving ambiguous states * `--ignore-missing-alns` — Treat tree tips that have no sequence in the alignment as fully ambiguous (missing data) instead of aborting. @@ -1489,10 +1549,15 @@ Reconstructs ancestral sequences and maps mutations to the tree. The tree is the Reconstructs discrete ancestral states, for example geographic location, host, or similar. In addition to ancestral states, a GTR model of state transitions is inferred -**Usage:** `treetime mugration [OPTIONS] --attribute --metadata ` +**Usage:** `treetime mugration [OPTIONS]` ###### **Options:** +* `--config ` — Load command configuration from a file (optionally compressed: `.gz`, `.bz2`, `.xz`, `.zst`). + + The file holds this command's configuration object, in the same shape the command serializes to. It is parsed as YAML, which also accepts JSON. Use `-` to read from stdin. + + Explicit command-line flags override values from the file; the file overrides defaults. A boolean enabled in the config cannot be disabled from the command line (a flag has no `false` spelling); edit the config instead. * `-t`, `--tree ` — Name of file containing the tree in newick, nexus, or phylip format. If none is provided, treetime will attempt to build a tree from the alignment using fasttree, iqtree, or raxml (assuming they are installed) @@ -1643,6 +1708,24 @@ Reconstructs discrete ancestral states, for example geographic location, host, o +## `treetime pipeline` + +Runs an ordered list of analysis commands from one config file, on one dataset, in one process. + +The config file (JSON or YAML) lists named steps, each an analysis command with its arguments. Steps run sequentially; a step may reference an earlier step's outputs with `{{ steps..outputs. }}`. Use `--check` to print the resolved plan without running anything, and `--steps` to run only a subset. + +**Usage:** `treetime pipeline [OPTIONS] --config ` + +###### **Options:** + +* `--config ` — Pipeline configuration file (JSON or YAML) describing an ordered list of steps +* `--steps ` — Run only these named steps (comma-separated), in the pipeline's list order. + + A referenced upstream step that is not selected must already have its outputs on disk. +* `--check` — Resolve and print the plan (steps, inputs, outputs) without running any step + + + ## `treetime arg` Estimates ancestral reassortment graph (ARG) @@ -1659,11 +1742,11 @@ Write JSON Schema definitions for TreeTime data types ###### **Options:** -* `--for ` — Which schema to generate +* `--for ` — Which schema to generate Default value: `all` - Possible values: `all`, `version-info`, `progress-event`, `error-response` + Possible values: `all`, `version-info`, `progress-event`, `error-response`, `pipeline`, `timetree`, `optimize`, `prune`, `ancestral`, `clock`, `mugration` * `-o`, `--output ` — Output file or directory (use "-" for stdout). Directory required when --for=all diff --git a/kb/features/optimize.md b/kb/features/optimize.md index 593bb8dc5..913725a75 100644 --- a/kb/features/optimize.md +++ b/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)). diff --git a/packages/app-server/src/args.rs b/packages/app-server/src/args.rs index 5e86ee28a..f50866e07 100644 --- a/packages/app-server/src/args.rs +++ b/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 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, diff --git a/packages/schemas/input-config-optimize.schema.json b/packages/schemas/input-config-optimize.schema.json index 0999f05da..96ce86d71 100644 --- a/packages/schemas/input-config-optimize.schema.json +++ b/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", diff --git a/packages/schemas/input-config-pipeline.schema.json b/packages/schemas/input-config-pipeline.schema.json index b4fa71a35..753157eef 100644 --- a/packages/schemas/input-config-pipeline.schema.json +++ b/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", diff --git a/packages/treetime/src/commands/optimize/__tests__/test_augur_node_data.rs b/packages/treetime/src/commands/optimize/__tests__/test_augur_node_data.rs index a316dfc23..25c256266 100644 --- a/packages/treetime/src/commands/optimize/__tests__/test_augur_node_data.rs +++ b/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 { diff --git a/packages/treetime/src/commands/optimize/args.rs b/packages/treetime/src/commands/optimize/args.rs index d460ca750..06d799f46 100644 --- a/packages/treetime/src/commands/optimize/args.rs +++ b/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, pub reroot_tips: Vec, pub keep_root: bool, + pub topology_ops: TopologyOps, pub gap_fill_args: GapFillArgs, } @@ -285,6 +312,11 @@ impl TryFrom 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, }) } diff --git a/packages/treetime/src/commands/optimize/run.rs b/packages/treetime/src/commands/optimize/run.rs index 3bb85b65c..fc75b8cb5 100644 --- a/packages/treetime/src/commands/optimize/run.rs +++ b/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 { diff --git a/packages/treetime/src/optimize/__tests__/test_convergence_conditions.rs b/packages/treetime/src/optimize/__tests__/test_convergence_conditions.rs index 12cc1d7cf..22a3f4f9f 100644 --- a/packages/treetime/src/optimize/__tests__/test_convergence_conditions.rs +++ b/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!( diff --git a/packages/treetime/src/optimize/__tests__/test_convergence_sc2.rs b/packages/treetime/src/optimize/__tests__/test_convergence_sc2.rs index 64498e95f..71dae86fa 100644 --- a/packages/treetime/src/optimize/__tests__/test_convergence_sc2.rs +++ b/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!( diff --git a/packages/treetime/src/optimize/__tests__/test_damping.rs b/packages/treetime/src/optimize/__tests__/test_damping.rs index f7e92d237..7f08d4b34 100644 --- a/packages/treetime/src/optimize/__tests__/test_damping.rs +++ b/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. diff --git a/packages/treetime/src/optimize/__tests__/test_gm_optimize.rs b/packages/treetime/src/optimize/__tests__/test_gm_optimize.rs index 6b209caf9..185a355eb 100644 --- a/packages/treetime/src/optimize/__tests__/test_gm_optimize.rs +++ b/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 diff --git a/packages/treetime/src/optimize/__tests__/test_no_indels.rs b/packages/treetime/src/optimize/__tests__/test_no_indels.rs index 7c23d845b..2e040937f 100644 --- a/packages/treetime/src/optimize/__tests__/test_no_indels.rs +++ b/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!( diff --git a/packages/treetime/src/optimize/__tests__/test_pipeline_gtr_normalized.rs b/packages/treetime/src/optimize/__tests__/test_pipeline_gtr_normalized.rs index c56523636..f7fa1743d 100644 --- a/packages/treetime/src/optimize/__tests__/test_pipeline_gtr_normalized.rs +++ b/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 { diff --git a/packages/treetime/src/optimize/__tests__/test_pipeline_reroot.rs b/packages/treetime/src/optimize/__tests__/test_pipeline_reroot.rs index 7633722e1..db336b600 100644 --- a/packages/treetime/src/optimize/__tests__/test_pipeline_reroot.rs +++ b/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(), } } diff --git a/packages/treetime/src/optimize/__tests__/test_run_optimize_loop.rs b/packages/treetime/src/optimize/__tests__/test_run_optimize_loop.rs index c97159a1a..16154b3b0 100644 --- a/packages/treetime/src/optimize/__tests__/test_run_optimize_loop.rs +++ b/packages/treetime/src/optimize/__tests__/test_run_optimize_loop.rs @@ -4,7 +4,7 @@ mod tests { use crate::optimize::__tests__::test_convergence::test_convergence_support::tests::{ TREE_NEWICK, setup_partitions, simple_alignment, }; - use crate::optimize::params::BranchOptMethod; + use crate::optimize::params::{BranchOptMethod, TopologyOps}; use crate::optimize::run_loop::{ConvergenceReason, run_optimize_loop}; use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; use crate::payload::ancestral::GraphAncestral; @@ -103,6 +103,7 @@ mod tests { 0.75, BranchOptMethod::BrentSqrt, false, + TopologyOps::default(), )?; // One entry per executed iteration, regardless of how the loop stopped. @@ -145,6 +146,7 @@ mod tests { 0.75, BranchOptMethod::BrentSqrt, false, + TopologyOps::default(), )?; assert_eq!(result.lh_history.len(), 1); @@ -178,6 +180,7 @@ mod tests { 0.0, BranchOptMethod::BrentSqrt, false, + TopologyOps::default(), )?; assert_eq!(result.stopped_at, Some((1, ConvergenceReason::Converged))); @@ -202,6 +205,7 @@ mod tests { 0.75, BranchOptMethod::BrentSqrt, false, + TopologyOps::default(), )?; assert!(result.lh_history.is_empty()); @@ -229,6 +233,7 @@ mod tests { 0.75, BranchOptMethod::BrentSqrt, false, + TopologyOps::default(), )?; for (i, lh) in result.lh_history.iter().map(|log_lh| log_lh.value()).enumerate() { @@ -272,6 +277,7 @@ mod tests { 0.75, BranchOptMethod::BrentSqrt, false, + TopologyOps::default(), )?; let final_sparse_lh = update_marginal(&graph, &sparse_partitions)?.value(); diff --git a/packages/treetime/src/optimize/__tests__/test_topology_cleanup.rs b/packages/treetime/src/optimize/__tests__/test_topology_cleanup.rs index a6c5c3d49..3ded31973 100644 --- a/packages/treetime/src/optimize/__tests__/test_topology_cleanup.rs +++ b/packages/treetime/src/optimize/__tests__/test_topology_cleanup.rs @@ -6,9 +6,9 @@ mod tests { use crate::gtr::get_gtr::{JC69Params, jc69}; use crate::optimize::dispatch::{initial_guess_mixed, run_optimize_mixed}; 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::collect_optimize_partitions; - use crate::optimize::run_loop::{find_zero_optimal_internal_edges, prune_and_merge_in_loop}; + use crate::optimize::run_loop::{find_zero_optimal_internal_edges, prune_and_merge_in_loop, run_optimize_loop}; use crate::optimize::topology::merge_shared_mutations::merge_shared_mutation_branches; use crate::partition::marginal::dense::partition::PartitionMarginalDense; use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; @@ -114,7 +114,7 @@ mod tests { let sparse: Vec>> = vec![]; let dense: Vec>> = vec![]; - let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[])?; + let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[], TopologyOps::default())?; assert!(!changed); assert_eq!(graph.get_nodes().len(), 4); Ok(()) @@ -172,7 +172,7 @@ mod tests { edge.write_arc().payload().write_arc().set_branch_length(Some(0.0)); } - let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[ri_key])?; + let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[ri_key], TopologyOps::default())?; assert!(changed); // I should be gone @@ -248,7 +248,7 @@ mod tests { apply_damping(&graph, &old_branch_lengths, 0.75, i); - prune_and_merge_in_loop(&mut graph, &sparse_partitions, &dense_partitions, &zero_optimal_edges)?; + prune_and_merge_in_loop(&mut graph, &sparse_partitions, &dense_partitions, &zero_optimal_edges, TopologyOps::default())?; lh_prev = total_lh; } @@ -321,7 +321,7 @@ mod tests { let zero_optimal_edges = find_zero_optimal_internal_edges(&graph, &sparse_partitions); apply_damping(&graph, &old_branch_lengths, 0.75, i); - prune_and_merge_in_loop(&mut graph, &sparse_partitions, &dense_partitions, &zero_optimal_edges)?; + prune_and_merge_in_loop(&mut graph, &sparse_partitions, &dense_partitions, &zero_optimal_edges, TopologyOps::default())?; lh_prev = total_lh; } @@ -429,7 +429,7 @@ mod tests { let dense: Vec>> = vec![]; // Empty zero-optimal list: the old loop was a no-op here. The hoist must still fire. - let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[])?; + let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[], TopologyOps::default())?; assert!(changed, "reversion polytomy must be resolved even without a collapse"); let p = sparse[0].read_arc(); @@ -467,7 +467,13 @@ mod tests { let sparse: Vec>> = vec![]; let dense: Vec>> = vec![]; - let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[ri1_key, i1i2_key])?; + let changed = prune_and_merge_in_loop( + &mut graph, + &sparse, + &dense, + &[ri1_key, i1i2_key], + TopologyOps::default(), + )?; assert!(changed); // Both I1 and I2 should be gone. A, B become children of root. @@ -535,7 +541,7 @@ mod tests { let zero_optimal_edges = find_zero_optimal_internal_edges(&graph, &sparse_partitions); apply_damping(&graph, &old_branch_lengths, 0.75, i); - prune_and_merge_in_loop(&mut graph, &sparse_partitions, &dense_partitions, &zero_optimal_edges)?; + prune_and_merge_in_loop(&mut graph, &sparse_partitions, &dense_partitions, &zero_optimal_edges, TopologyOps::default())?; lh_prev = dense_lh; } @@ -603,7 +609,7 @@ mod tests { edge.write_arc().payload().write_arc().set_branch_length(Some(0.0)); } - let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[ri_key])?; + let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[ri_key], TopologyOps::default())?; assert!(changed); let mut names: Vec = graph @@ -618,4 +624,245 @@ mod tests { Ok(()) } + + #[test] + fn test_optimize_prune_and_merge_merge_disabled_keeps_polytomy() -> Result<(), Report> { + // Same setup as the collapse+merge test, but with merge-siblings disabled. Collapsing the + // zero-length internal edge still forms the polytomy; without merge, the shared-mutation + // siblings A, B, C stay as direct children of root rather than being grouped under a node. + let mut graph: GraphAncestral = nwk_read_str("((A:0.1,B:0.1)I:0.0,C:0.1,D:0.1)root;")?; + + let ri_key = find_edge_key(&graph, "root", "I").unwrap(); + + let mut partition = PartitionMarginalSparse { + index: 0, + gtr: jc69(JC69Params::default())?, + alphabet: Alphabet::new(AlphabetName::Nuc)?, + length: 100, + nodes: btreemap! {}, + edges: btreemap! {}, + root_sequence: seq![], + }; + populate_test_nodes(&mut partition, &graph); + partition.edges.insert(ri_key, SparseEdgePartition::default()); + + let ia_key = find_edge_key(&graph, "I", "A").unwrap(); + let ib_key = find_edge_key(&graph, "I", "B").unwrap(); + let rc_key = find_edge_key(&graph, "root", "C").unwrap(); + let rd_key = find_edge_key(&graph, "root", "D").unwrap(); + partition + .edges + .insert(ia_key, SparseEdgePartition::with_fitch_subs(vec![sub(b'A', 0, b'T')])); + partition + .edges + .insert(ib_key, SparseEdgePartition::with_fitch_subs(vec![sub(b'A', 0, b'T')])); + partition + .edges + .insert(rc_key, SparseEdgePartition::with_fitch_subs(vec![sub(b'A', 0, b'T')])); + partition + .edges + .insert(rd_key, SparseEdgePartition::with_fitch_subs(vec![sub(b'G', 5, b'C')])); + + let sparse = vec![Arc::new(RwLock::new(partition))]; + let dense: Vec>> = vec![]; + + if let Some(edge) = graph.get_edge(ri_key) { + edge.write_arc().payload().write_arc().set_branch_length(Some(0.0)); + } + + let ops = TopologyOps { + merge_siblings: false, + ..TopologyOps::default() + }; + let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[ri_key], ops)?; + assert!(changed, "collapse still fires even with merge disabled"); + + // I collapsed away. + assert!(find_node_key_by_name(&graph, "I").is_none()); + + // Without merge, root keeps all four children A, B, C, D (no grouping under a new node). + let root_key = find_node_key_by_name(&graph, "root").unwrap(); + let root_node = graph.get_node(root_key).unwrap(); + assert_eq!(root_node.read_arc().degree_out(), 4); + + Ok(()) + } + + #[test] + fn test_optimize_prune_and_merge_flip_disabled_keeps_reversion() -> Result<(), Report> { + // Reversion polytomy. With flip-parent-child disabled, merge still groups the two reverting + // children, but the reverting mutation is not hoisted away, so it remains in the tree. + let mut graph: GraphAncestral = nwk_read_str("(((C1:0.1,C2:0.1,C3:0.1)V:0.2)U:0.1)root:0.0;")?; + + let mut partition = PartitionMarginalSparse { + index: 0, + gtr: jc69(JC69Params::default())?, + alphabet: Alphabet::new(AlphabetName::Nuc)?, + length: 100, + nodes: btreemap! {}, + edges: btreemap! {}, + root_sequence: seq![], + }; + populate_test_nodes(&mut partition, &graph); + + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let vc1 = find_edge_key(&graph, "V", "C1").unwrap(); + let vc2 = find_edge_key(&graph, "V", "C2").unwrap(); + let vc3 = find_edge_key(&graph, "V", "C3").unwrap(); + partition.edges.insert( + uv, + SparseEdgePartition::with_fitch_subs(vec![sub(b'A', 0, b'T'), sub(b'C', 5, b'G')]), + ); + partition + .edges + .insert(vc1, SparseEdgePartition::with_fitch_subs(vec![sub(b'T', 0, b'A')])); + partition + .edges + .insert(vc2, SparseEdgePartition::with_fitch_subs(vec![sub(b'T', 0, b'A')])); + partition.edges.insert(vc3, SparseEdgePartition::default()); + + let sparse = vec![Arc::new(RwLock::new(partition))]; + let dense: Vec>> = vec![]; + + let ops = TopologyOps { + flip_parent_child: false, + ..TopologyOps::default() + }; + let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[], ops)?; + assert!(changed, "merge still groups the reverting siblings"); + + let p = sparse[0].read_arc(); + let reversion_remains = graph + .get_edges() + .iter() + .filter_map(|e| p.edges.get(&e.read_arc().key())) + .any(|e| e.fitch_subs().contains(&sub(b'T', 0, b'A'))); + assert!( + reversion_remains, + "reversion is kept when flip-parent-child is disabled" + ); + + Ok(()) + } + + #[test] + fn test_optimize_prune_and_merge_all_ops_disabled_is_noop() -> Result<(), Report> { + // With every topology step disabled, the reversion polytomy is left untouched: no collapse, + // no merge, no hoist. The tree shape and its mutation content are unchanged. + let mut graph: GraphAncestral = nwk_read_str("(((C1:0.1,C2:0.1,C3:0.1)V:0.2)U:0.1)root:0.0;")?; + + let mut partition = PartitionMarginalSparse { + index: 0, + gtr: jc69(JC69Params::default())?, + alphabet: Alphabet::new(AlphabetName::Nuc)?, + length: 100, + nodes: btreemap! {}, + edges: btreemap! {}, + root_sequence: seq![], + }; + populate_test_nodes(&mut partition, &graph); + + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let vc1 = find_edge_key(&graph, "V", "C1").unwrap(); + let vc2 = find_edge_key(&graph, "V", "C2").unwrap(); + let vc3 = find_edge_key(&graph, "V", "C3").unwrap(); + partition.edges.insert( + uv, + SparseEdgePartition::with_fitch_subs(vec![sub(b'A', 0, b'T'), sub(b'C', 5, b'G')]), + ); + partition + .edges + .insert(vc1, SparseEdgePartition::with_fitch_subs(vec![sub(b'T', 0, b'A')])); + partition + .edges + .insert(vc2, SparseEdgePartition::with_fitch_subs(vec![sub(b'T', 0, b'A')])); + partition.edges.insert(vc3, SparseEdgePartition::default()); + + let sparse = vec![Arc::new(RwLock::new(partition))]; + let dense: Vec>> = vec![]; + + let node_count_before = graph.get_nodes().len(); + let ops = TopologyOps { + collapse_short_branches: false, + merge_siblings: false, + flip_parent_child: false, + }; + let changed = prune_and_merge_in_loop(&mut graph, &sparse, &dense, &[], ops)?; + assert!(!changed, "no topology step runs when all are disabled"); + assert_eq!(graph.get_nodes().len(), node_count_before); + + let p = sparse[0].read_arc(); + let total_subs: usize = graph + .get_edges() + .iter() + .filter_map(|e| p.edges.get(&e.read_arc().key())) + .map(|e| e.fitch_subs().len()) + .sum(); + assert_eq!(total_subs, 4, "mutation content is unchanged"); + + Ok(()) + } + + #[rustfmt::skip] + #[rstest] + #[case::brent_sqrt(BranchOptMethod::BrentSqrt)] + #[case::newton( BranchOptMethod::Newton)] + #[trace] + fn test_run_optimize_loop_collapse_disabled_keeps_zero_edge(#[case] method: BranchOptMethod) -> Result<(), Report> { + // A and B are identical, so the AB internal edge optimizes toward zero. With collapse + // disabled, run_optimize_loop must leave that edge in place: the node count is unchanged. + let nuc = Alphabet::new(AlphabetName::Nuc)?; + let aln = read_many_fasta_str( + indoc! {r#" + >A + ACGTACGTACGT + >B + ACGTACGTACGT + >C + ACGTACGTACGG + >D + TCGTACGTACGT + "#}, + &nuc, + )?; + + let mut graph: GraphAncestral = nwk_read_str("((A:0.01,B:0.01)AB:0.01,(C:0.01,D:0.01)CD:0.01)root:0.0;")?; + + let fitch = create_fitch_partition(&graph, 0, nuc, &aln)?; + let sparse_partitions = vec![Arc::new(RwLock::new( + fitch.into_marginal_sparse(jc69(JC69Params::default())?, &graph)?, + ))]; + update_marginal(&graph, &sparse_partitions)?.value(); + + let dense_partitions: Vec>> = vec![]; + let mixed_partitions = collect_optimize_partitions(&dense_partitions, &sparse_partitions); + initial_guess_mixed(&graph, &mixed_partitions, true, false)?; + + let initial_node_count = graph.get_nodes().len(); + + let ops = TopologyOps { + collapse_short_branches: false, + ..TopologyOps::default() + }; + run_optimize_loop( + &mut graph, + &sparse_partitions, + &dense_partitions, + &mixed_partitions, + 10, + 1e-2, + 0.75, + method, + false, + ops, + )?; + + assert_eq!( + graph.get_nodes().len(), + initial_node_count, + "no collapse when collapse_short_branches is disabled" + ); + + Ok(()) + } } diff --git a/packages/treetime/src/optimize/params.rs b/packages/treetime/src/optimize/params.rs index a38b216bd..fcebbf0e1 100644 --- a/packages/treetime/src/optimize/params.rs +++ b/packages/treetime/src/optimize/params.rs @@ -2,6 +2,29 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use smart_default::SmartDefault; +/// Which topology-cleanup steps run during the optimize loop. +/// +/// Each step is applied once per iteration unless disabled here. All default to +/// on, so an optimize run with no topology flags behaves as it did before the +/// switches existed. The steps map to the loop's cleanup passes: +/// +/// - `collapse_short_branches`: contract internal edges the per-edge optimizer +/// drove to exactly zero that carry no mutations (v0's `prune_short_branches`). +/// - `merge_siblings`: group polytomy siblings that share substitutions under a +/// new internal node. Requires discrete per-edge mutations, so sparse only. +/// - `flip_parent_child`: hoist a child that carries the exact reversion of a +/// parent-edge substitution under a new node, removing one mutation per +/// reverted position. Requires discrete per-edge mutations, so sparse only. +#[derive(Copy, Clone, Debug, PartialEq, Eq, SmartDefault)] +pub struct TopologyOps { + #[default = true] + pub collapse_short_branches: bool, + #[default = true] + pub merge_siblings: bool, + #[default = true] + pub flip_parent_child: bool, +} + #[derive(Copy, Debug, Clone, PartialEq, Eq, PartialOrd, Ord, SmartDefault, Serialize, Deserialize, JsonSchema)] #[cfg_attr(feature = "clap", derive(clap::ValueEnum))] #[serde(rename_all = "kebab-case")] diff --git a/packages/treetime/src/optimize/pipeline.rs b/packages/treetime/src/optimize/pipeline.rs index 11029e08c..0a70bd54c 100644 --- a/packages/treetime/src/optimize/pipeline.rs +++ b/packages/treetime/src/optimize/pipeline.rs @@ -5,7 +5,7 @@ use crate::gtr::get_gtr::GtrModelName; use crate::gtr::gtr::GTR; use crate::optimize::dispatch::{run_optimize_mixed, run_optimize_mixed_inner}; use crate::optimize::iteration::{apply_damping, save_branch_lengths}; -use crate::optimize::params::{BranchOptMethod, InitialGuessMode}; +use crate::optimize::params::{BranchOptMethod, InitialGuessMode, TopologyOps}; use crate::optimize::run_loop::{ apply_initial_guess_mode, collect_optimize_partitions, normalize_partition_rates, run_optimize_loop, }; @@ -21,7 +21,7 @@ use crate::reroot::orchestrate::{RerootTopologyParams, reroot_at_node, reroot_in use crate::reroot::params::BrentParams; use crate::reroot::variance::VarianceModel; use eyre::Report; -use log::info; +use log::{info, warn}; use parking_lot::RwLock; use serde::Serialize; use std::sync::Arc; @@ -47,6 +47,8 @@ pub struct OptimizeParams { pub no_indels: bool, /// Reroot policy before optimization. `None` keeps the input root. pub reroot_spec: Option, + /// Which per-iteration topology-cleanup steps run. All on by default. + pub topology_ops: TopologyOps, } pub struct OptimizeInput { @@ -101,6 +103,22 @@ pub fn run( }, } + // `merge_siblings` and `flip_parent_child` operate on discrete per-edge mutation lists, which + // only sparse partitions carry. Under a dense build these steps never run, so disabling them + // has no effect; warn the user their flag did nothing rather than failing silently. + if sparse_partitions.is_empty() { + if !params.topology_ops.merge_siblings { + warn!( + "--no-merge-siblings has no effect in dense mode: sibling merging requires the sparse sequence representation" + ); + } + if !params.topology_ops.flip_parent_child { + warn!( + "--no-flip-parent-child has no effect in dense mode: the reversion hoist requires the sparse sequence representation" + ); + } + } + update_marginal(&input.graph, &sparse_partitions)?; if !dense_partitions.is_empty() { initialize_marginal(&input.graph, &dense_partitions, &input.sequences)?; @@ -145,6 +163,7 @@ pub fn run( params.damping, params.opt_method, params.no_indels, + params.topology_ops, )?; info!("Re-running marginal to populate subs_ml after optimization loop"); diff --git a/packages/treetime/src/optimize/run_loop.rs b/packages/treetime/src/optimize/run_loop.rs index d4d6c77c5..a183759aa 100644 --- a/packages/treetime/src/optimize/run_loop.rs +++ b/packages/treetime/src/optimize/run_loop.rs @@ -4,7 +4,7 @@ use crate::optimize::dispatch::initial_guess_mixed; use crate::optimize::dispatch::run_optimize_mixed_inner; use crate::optimize::indel::{estimate_indel_rate, total_indel_log_lh}; use crate::optimize::iteration::{apply_damping, restore_branch_lengths, save_branch_lengths}; -use crate::optimize::params::{BranchOptMethod, InitialGuessMode}; +use crate::optimize::params::{BranchOptMethod, InitialGuessMode, TopologyOps}; use crate::optimize::topology::collapse::collapse_edge; use crate::optimize::topology::resolve_polytomy::resolve_polytomies; use crate::partition::marginal::dense::partition::PartitionMarginalDense; @@ -69,6 +69,7 @@ pub fn run_optimize_loop( damping: f64, opt_method: BranchOptMethod, no_indels: bool, + topology_ops: TopologyOps, ) -> Result { let indel_rate = if no_indels { 0.0 @@ -141,11 +142,21 @@ pub fn run_optimize_loop( let old_branch_lengths = save_branch_lengths(graph); run_optimize_mixed_inner(graph, mixed_partitions, opt_method, indel_rate, no_indels)?; - let zero_optimal_edges = find_zero_optimal_internal_edges(graph, sparse_partitions); + let zero_optimal_edges = if topology_ops.collapse_short_branches { + find_zero_optimal_internal_edges(graph, sparse_partitions) + } else { + vec![] + }; apply_damping(graph, &old_branch_lengths, damping, i); - let topology_changed = prune_and_merge_in_loop(graph, sparse_partitions, dense_partitions, &zero_optimal_edges)?; + let topology_changed = prune_and_merge_in_loop( + graph, + sparse_partitions, + dense_partitions, + &zero_optimal_edges, + topology_ops, + )?; if topology_changed { best_lh = LogLh::IMPOSSIBLE; } @@ -321,6 +332,7 @@ pub fn prune_and_merge_in_loop( sparse_partitions: &[Arc>], dense_partitions: &[Arc>], zero_optimal_edges: &[GraphEdgeKey], + topology_ops: TopologyOps, ) -> Result { let mut topology_changed = false; @@ -351,7 +363,7 @@ pub fn prune_and_merge_in_loop( } } - if resolve_polytomies(graph, sparse_partitions, dense_partitions)? > 0 { + if resolve_polytomies(graph, sparse_partitions, dense_partitions, topology_ops)? > 0 { topology_changed = true; } diff --git a/packages/treetime/src/optimize/topology/__tests__/test_prop_resolve_polytomy.rs b/packages/treetime/src/optimize/topology/__tests__/test_prop_resolve_polytomy.rs index 5d6a9dd74..30ac14836 100644 --- a/packages/treetime/src/optimize/topology/__tests__/test_prop_resolve_polytomy.rs +++ b/packages/treetime/src/optimize/topology/__tests__/test_prop_resolve_polytomy.rs @@ -1,5 +1,6 @@ #[cfg(test)] mod tests { + use crate::optimize::params::TopologyOps; use crate::optimize::topology::resolve_polytomy::resolve_polytomies; use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; use crate::payload::ancestral::GraphAncestral; @@ -25,7 +26,7 @@ mod tests { let (mut graph, partition, before) = helpers::build_case(n_children, k, &revert_masks, &own_counts); let sparse = vec![partition]; - let changed = resolve_polytomies(&mut graph, &sparse, &helpers::no_dense()).unwrap(); + let changed = resolve_polytomies(&mut graph, &sparse, &helpers::no_dense(), TopologyOps::default()).unwrap(); let after = helpers::total_subs(&graph, &sparse[0].read_arc()); prop_assert!(after <= before, "mutation count increased: before={before} after={after}"); @@ -48,7 +49,7 @@ mod tests { let leaves_before = helpers::leaf_names(&graph); let sparse = vec![partition]; - resolve_polytomies(&mut graph, &sparse, &helpers::no_dense()).unwrap(); + resolve_polytomies(&mut graph, &sparse, &helpers::no_dense(), TopologyOps::default()).unwrap(); prop_assert_eq!(helpers::leaf_names(&graph), leaves_before); diff --git a/packages/treetime/src/optimize/topology/__tests__/test_resolve_polytomy.rs b/packages/treetime/src/optimize/topology/__tests__/test_resolve_polytomy.rs index 041478235..1f1595942 100644 --- a/packages/treetime/src/optimize/topology/__tests__/test_resolve_polytomy.rs +++ b/packages/treetime/src/optimize/topology/__tests__/test_resolve_polytomy.rs @@ -1,5 +1,6 @@ #[cfg(test)] mod tests { + use crate::optimize::params::TopologyOps; use crate::optimize::topology::resolve_polytomy::resolve_polytomies; use crate::partition::marginal::dense::partition::PartitionMarginalDense; use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; @@ -36,7 +37,7 @@ mod tests { ); let sparse = vec![partition]; - let changed = resolve_polytomies(&mut graph, &sparse, &no_dense())?; + let changed = resolve_polytomies(&mut graph, &sparse, &no_dense(), TopologyOps::default())?; assert!(changed > 0); let p = sparse[0].read_arc(); @@ -76,7 +77,7 @@ mod tests { let sparse = vec![partition]; let before = total_subs(&graph, &sparse[0].read_arc()); - resolve_polytomies(&mut graph, &sparse, &no_dense())?; + resolve_polytomies(&mut graph, &sparse, &no_dense(), TopologyOps::default())?; let after = total_subs(&graph, &sparse[0].read_arc()); assert_eq!(before, 5); @@ -102,7 +103,7 @@ mod tests { ); let sparse = vec![partition]; - resolve_polytomies(&mut graph, &sparse, &no_dense())?; + resolve_polytomies(&mut graph, &sparse, &no_dense(), TopologyOps::default())?; assert!( find_node_key_by_name(&graph, "W").is_some(), @@ -138,7 +139,7 @@ mod tests { let sparse = vec![partition]; let nodes_before = graph.get_nodes().len(); - let changed = resolve_polytomies(&mut graph, &sparse, &no_dense())?; + let changed = resolve_polytomies(&mut graph, &sparse, &no_dense(), TopologyOps::default())?; assert_eq!(changed, 0); assert_eq!(graph.get_nodes().len(), nodes_before); @@ -164,7 +165,7 @@ mod tests { let sparse = vec![partition]; let nodes_before = graph.get_nodes().len(); - let changed = resolve_polytomies(&mut graph, &sparse, &no_dense())?; + let changed = resolve_polytomies(&mut graph, &sparse, &no_dense(), TopologyOps::default())?; assert_eq!(changed, 0); assert_eq!(graph.get_nodes().len(), nodes_before); diff --git a/packages/treetime/src/optimize/topology/resolve_polytomy.rs b/packages/treetime/src/optimize/topology/resolve_polytomy.rs index 5c96071b8..d954a6225 100644 --- a/packages/treetime/src/optimize/topology/resolve_polytomy.rs +++ b/packages/treetime/src/optimize/topology/resolve_polytomy.rs @@ -1,3 +1,4 @@ +use crate::optimize::params::TopologyOps; use crate::optimize::topology::collapse::collapse_edge; use crate::optimize::topology::hoist_reversions::{count_child_reversions, hoist_reverting_child}; use crate::optimize::topology::merge_shared_mutations::merge_single_polytomy; @@ -32,8 +33,10 @@ pub fn resolve_polytomies( graph: &mut GraphAncestral, sparse: &[Arc>], dense: &[Arc>], + topology_ops: TopologyOps, ) -> Result { - if sparse.is_empty() { + // Sparse-only, and skipped entirely when both moves this routine performs are disabled. + if sparse.is_empty() || !(topology_ops.merge_siblings || topology_ops.flip_parent_child) { return Ok(0); } @@ -42,7 +45,7 @@ pub fn resolve_polytomies( let polytomy_keys = find_polytomy_nodes(graph); let mut round_changed = 0; for node_key in polytomy_keys { - if resolve_one(graph, sparse, dense, node_key)? { + if resolve_one(graph, sparse, dense, node_key, topology_ops)? { round_changed += 1; } } @@ -80,13 +83,16 @@ fn resolve_one( sparse: &[Arc>], dense: &[Arc>], node_key: GraphNodeKey, + topology_ops: TopologyOps, ) -> Result { let preexisting: BTreeSet = graph.get_nodes().iter().map(|node| node.read_arc().key()).collect(); let mut any_changed = false; loop { - let merged = merge_single_polytomy(graph, sparse, node_key)? > 0; - let hoisted = try_hoist_reverting_child(graph, sparse, dense, node_key)?; + // Each move is gated independently; retirement stays unconditional because it only ever + // collapses helper edges the enabled moves just created, so it is a no-op when neither ran. + let merged = topology_ops.merge_siblings && merge_single_polytomy(graph, sparse, node_key)? > 0; + let hoisted = topology_ops.flip_parent_child && try_hoist_reverting_child(graph, sparse, dense, node_key)?; let retired = retire_created_helpers(graph, sparse, dense, &preexisting)?; if !(merged || hoisted || retired) {