From 33c0bd853ddc520014ea3c067286328b59eda661 Mon Sep 17 00:00:00 2001 From: ivan-aksamentov Date: Wed, 2 Sep 2026 19:04:01 +0200 Subject: [PATCH 1/4] feat(optimize): resolve reversion polytomies in the optimize loop - Add a reversion-driven hoist: when a child edge reverts a substitution on the parent edge, insert a node grouping the child with its sibling subtree and lift the non-reverted substitutions above it, removing one mutation per reversion and adding none. Branch lengths split proportionally to preserve root-to-node distances; indels use an all-or-nothing distance-preserving rule. - Fold the shared-mutation merge, the hoist, and helper-node retirement into a single per-polytomy routine driven to a fixpoint by a monotone (mutation count, node count) potential. Run it every loop iteration, independent of the zero-optimal collapse, so reversions in the input and in polytomies formed by earlier iterations are resolved. - Relocate the two moved edges with reparent_edge to keep their keys and partition state; register the fresh node and edge in dense partitions. --- packages/treetime/src/optimize/run_loop.rs | 81 ++-- .../src/optimize/topology/hoist_reversions.rs | 364 ++++++++++++++++++ .../topology/merge_shared_mutations.rs | 2 +- .../treetime/src/optimize/topology/mod.rs | 2 + .../src/optimize/topology/resolve_polytomy.rs | 212 ++++++++++ 5 files changed, 625 insertions(+), 36 deletions(-) create mode 100644 packages/treetime/src/optimize/topology/hoist_reversions.rs create mode 100644 packages/treetime/src/optimize/topology/resolve_polytomy.rs diff --git a/packages/treetime/src/optimize/run_loop.rs b/packages/treetime/src/optimize/run_loop.rs index a529d43c8..d4d6c77c5 100644 --- a/packages/treetime/src/optimize/run_loop.rs +++ b/packages/treetime/src/optimize/run_loop.rs @@ -6,7 +6,7 @@ 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::topology::collapse::collapse_edge; -use crate::optimize::topology::merge_shared_mutations::merge_shared_mutation_branches; +use crate::optimize::topology::resolve_polytomy::resolve_polytomies; use crate::partition::marginal::dense::partition::PartitionMarginalDense; use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; use crate::partition::traits::{HasGtr, PartitionOptimizeOps, PartitionOptimizeVec}; @@ -297,13 +297,23 @@ pub fn find_zero_optimal_internal_edges( .collect_vec() } -/// Collapse zero-optimal internal edges and merge shared mutations in resulting polytomies. +/// Collapse zero-optimal internal edges, then resolve polytomies in the resulting tree. /// -/// Analogous to v0's `prune_short_branches()` inside the optimization loop: -/// edges whose optimal branch length is zero are collapsed, simplifying the tree -/// progressively across iterations. After collapsing, sibling branches in newly -/// formed polytomies that share identical substitutions are merged under new -/// internal nodes (sparse partitions only, requires discrete substitution data). +/// Two topology-cleanup steps run per iteration: +/// +/// 1. Collapse zero-optimal internal edges. Analogous to v0's `prune_short_branches()`: +/// edges whose optimal branch length is zero are collapsed, simplifying the tree +/// progressively across iterations and producing polytomies. +/// 2. Resolve polytomies ([`resolve_polytomies`], sparse partitions only): merge siblings +/// that share substitutions, hoist a reverting child under a new node to remove one +/// mutation per reversion, and retire the helper nodes left behind. +/// +/// Step 2 runs every iteration, not only after a collapse fired, because reverting-child and +/// shared-mutation polytomies exist in the input and in polytomies formed by earlier +/// iterations, independent of any zero-optimal collapse in the current one. The combined +/// (mutation count, node count) potential is monotone non-increasing, so the two steps cannot +/// oscillate across iterations even though the hoist deliberately relocates mutation-carrying +/// edges that [`find_zero_optimal_internal_edges`] refuses to collapse. /// /// Returns true if any topology change occurred. pub fn prune_and_merge_in_loop( @@ -312,44 +322,45 @@ pub fn prune_and_merge_in_loop( dense_partitions: &[Arc>], zero_optimal_edges: &[GraphEdgeKey], ) -> Result { - if zero_optimal_edges.is_empty() { - return Ok(false); - } + let mut topology_changed = false; + + if !zero_optimal_edges.is_empty() { + // Override damped branch lengths back to zero for edges the optimizer identified + // as zero-optimal. Damping is a convergence aid for continuous values; it should + // not prevent collapsing degenerate edges. + for &edge_key in zero_optimal_edges { + if let Some(edge) = graph.get_edge(edge_key) { + edge.write_arc().payload().write_arc().set_branch_length(Some(0.0)); + } + } - // Override damped branch lengths back to zero for edges the optimizer identified - // as zero-optimal. Damping is a convergence aid for continuous values; it should - // not prevent collapsing degenerate edges. - for &edge_key in zero_optimal_edges { - if let Some(edge) = graph.get_edge(edge_key) { - edge.write_arc().payload().write_arc().set_branch_length(Some(0.0)); + let mut collapsed = 0_usize; + for &edge_key in zero_optimal_edges { + // Edge may already be gone if a prior collapse in this batch removed it + // (e.g., collapsing a parent also removed a child edge that was zero-optimal) + if graph.get_edge(edge_key).is_none() { + continue; + } + collapse_edge(graph, sparse_partitions, dense_partitions, edge_key)?; + collapsed += 1; } - } - let mut collapsed = 0_usize; - for &edge_key in zero_optimal_edges { - // Edge may already be gone if a prior collapse in this batch removed it - // (e.g., collapsing a parent also removed a child edge that was zero-optimal) - if graph.get_edge(edge_key).is_none() { - continue; + if collapsed > 0 { + debug!("Collapsed {collapsed} zero-optimal internal edges"); + topology_changed = true; } - collapse_edge(graph, sparse_partitions, dense_partitions, edge_key)?; - collapsed += 1; } - if collapsed == 0 { - return Ok(false); + if resolve_polytomies(graph, sparse_partitions, dense_partitions)? > 0 { + topology_changed = true; } - debug!("Collapsed {collapsed} zero-optimal internal edges"); - - // Merge shared mutations in newly formed polytomies (sparse only) - if !sparse_partitions.is_empty() { - merge_shared_mutation_branches(graph, sparse_partitions)?; + if topology_changed { + graph.build()?; + assign_node_names(graph)?; } - graph.build()?; - assign_node_names(graph)?; - Ok(true) + Ok(topology_changed) } /// Whether any edge that carries indels has a zero branch length. diff --git a/packages/treetime/src/optimize/topology/hoist_reversions.rs b/packages/treetime/src/optimize/topology/hoist_reversions.rs new file mode 100644 index 000000000..82ca65aab --- /dev/null +++ b/packages/treetime/src/optimize/topology/hoist_reversions.rs @@ -0,0 +1,364 @@ +use crate::partition::marginal::dense::partition::PartitionMarginalDense; +use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; +use crate::partition::storage::dense::{DenseNodePartition, DenseSeqDistribution, DenseSeqInfo}; +use crate::partition::storage::sparse::SparseNodePartition; +use crate::payload::ancestral::{EdgeAncestral, GraphAncestral, NodeAncestral}; +use crate::seq::indel::{InDel, compose_indels, sort_indels}; +use crate::seq::mutation::Sub; +use eyre::Report; +use parking_lot::RwLock; +use std::cmp::Ordering; +use std::sync::Arc; +use treetime_graph::edge::{GraphEdgeKey, HasBranchLength}; +use treetime_graph::node::GraphNodeKey; + +/// Count the substitution reversions a child edge applies to its parent edge, summed +/// across sparse partitions. +/// +/// A position is a reversion when the parent edge carries $a \to b$ and the child edge +/// carries the exact inverse $b \to a$, so composing the two edges cancels the change. +/// This is the gain $\lvert R\rvert$ of hoisting that child (see [`hoist_reverting_child`]); +/// each reversion is one mutation the move removes. +pub(crate) fn count_child_reversions( + sparse: &[Arc>], + parent_edge_key: GraphEdgeKey, + child_edge_key: GraphEdgeKey, +) -> usize { + sparse + .iter() + .map(|partition| { + let partition = partition.read_arc(); + let parent_subs: &[Sub] = partition.edges.get(&parent_edge_key).map_or(&[], |e| e.fitch_subs()); + let child_subs: &[Sub] = partition.edges.get(&child_edge_key).map_or(&[], |e| e.fitch_subs()); + count_reversions(parent_subs, child_subs) + }) + .sum() +} + +/// Insert a new node $N$ between $u$ and $v$ that groups $v$ and one reverting child $c$, +/// hoisting the non-reverted substitutions above $N$. +/// +/// Let $e_p = u \to v$ carry the parent substitutions $M_v$ and $e_c = v \to c$ carry $M_c$. +/// Each position of $M_v$ falls into one of three disjoint sets relative to $M_c$: +/// +/// - $T$: positions untouched by the child. Hoisted onto $u \to N$. +/// - $H$: chained positions ($a \to b$ then $b \to d$, $d \neq a$). Kept on $N \to v$; the +/// composed $a \to d$ moves to $N \to c$. +/// - $R$: reverted positions ($a \to b$ then $b \to a$). Kept on $N \to v$; nothing on +/// $N \to c$, because the composition cancels. This is where the move removes one mutation +/// per reverted position. +/// +/// The resulting edges carry: +/// +/// | edge | substitutions | +/// | --- | --- | +/// | $u \to N$ | $T$ | +/// | $N \to v$ | $H \cup R$ (the original $M_v$ entries at child-shared positions) | +/// | $N \to c$ | $\mathrm{compose}(M_v, M_c)$ at $M_c$ positions ($H' \cup D$) | +/// +/// where $D$ are the child's own positions absent from $M_v$. The net substitution change is +/// $\Delta = -\lvert R\rvert$: no mutation is ever added. +/// +/// Branch lengths are split in proportion to substitution count so that root-to-$v$ and +/// root-to-$c$ distances are preserved exactly; the next optimizer iteration re-fits them. +/// +/// Indels use an all-or-nothing rule (see [`split_indels`]) that is always distance +/// preserving. The substitution gain is unaffected by the indel handling. +/// +/// The relocated edges $e_p$ and $e_c$ keep their edge keys via +/// [`Graph::reparent_edge`](treetime_graph::graph::Graph::reparent_edge), so their partition +/// entries stay valid and only their substitution and indel content is rewritten. Only the +/// fresh $u \to N$ edge and the node $N$ are new keys to register. +/// +/// Returns the key of the new node $N$. +pub(crate) fn hoist_reverting_child( + graph: &mut GraphAncestral, + sparse: &[Arc>], + dense: &[Arc>], + parent_edge_key: GraphEdgeKey, + child_edge_key: GraphEdgeKey, +) -> Result { + let u_key = graph.get_source_node_key(parent_edge_key)?; + + // Compute the per-partition split before touching the graph, so the reads see the + // pre-move edge state. + let mut splits = Vec::with_capacity(sparse.len()); + let mut total_parent_subs = 0_usize; + let mut total_hoisted_subs = 0_usize; + for partition in sparse { + let partition = partition.read_arc(); + let parent_subs = partition + .edges + .get(&parent_edge_key) + .map_or(Vec::new(), |e| e.fitch_subs().to_vec()); + let child_subs = partition + .edges + .get(&child_edge_key) + .map_or(Vec::new(), |e| e.fitch_subs().to_vec()); + let parent_indels = partition + .edges + .get(&parent_edge_key) + .map_or(Vec::new(), |e| e.indels.clone()); + let child_indels = partition + .edges + .get(&child_edge_key) + .map_or(Vec::new(), |e| e.indels.clone()); + + let sub_split = split_subs(&parent_subs, &child_subs)?; + let indel_split = split_indels(&parent_indels, &child_indels); + + total_parent_subs += parent_subs.len(); + total_hoisted_subs += sub_split.hoisted.len(); + + splits.push(EdgeSplit { + hoisted: sub_split.hoisted, + kept: sub_split.kept, + composed: sub_split.composed, + indels: indel_split, + }); + } + + // Distance-preserving branch-length split, proportional to substitution count. The move + // only fires when R is non-empty, so `total_parent_subs >= 1` and the ratio is well defined. + let bl_uv = edge_branch_length(graph, parent_edge_key); + let bl_vc = edge_branch_length(graph, child_edge_key); + let bl_un = if total_parent_subs > 0 { + bl_uv * (total_hoisted_subs as f64) / (total_parent_subs as f64) + } else { + 0.0 + }; + let bl_nv = bl_uv - bl_un; + let bl_nc = bl_nv + bl_vc; + + // Graph surgery: add N, connect u -> N, then relocate the two existing edges under N. + let n_key = graph.add_node(NodeAncestral::default()); + let un_edge_key = graph.add_edge( + u_key, + n_key, + EdgeAncestral { + branch_length: Some(bl_un), + }, + )?; + graph.reparent_edge(parent_edge_key, n_key)?; // e_p becomes N -> v + graph.reparent_edge(child_edge_key, n_key)?; // e_c becomes N -> c + set_edge_branch_length(graph, parent_edge_key, bl_nv); + set_edge_branch_length(graph, child_edge_key, bl_nc); + + // Sparse partition bookkeeping. The relocated edges keep their keys, so only their content + // is rewritten; the u -> N edge and node N are inserted fresh. + for (partition, split) in sparse.iter().zip(splits) { + let mut partition = partition.write_arc(); + + let mut node_n = SparseNodePartition::empty(&partition.alphabet); + node_n.seq.composition = partition.nodes[&u_key].seq.composition.clone(); + partition.nodes.entry(n_key).or_insert(node_n); + + let un_edge = partition.edges.entry(un_edge_key).or_default(); + un_edge.set_fitch_subs(split.hoisted); + un_edge.indels = split.indels.hoisted; + + let nv_edge = partition.edges.entry(parent_edge_key).or_default(); + nv_edge.set_fitch_subs(split.kept); + nv_edge.indels = split.indels.kept; + + let nc_edge = partition.edges.entry(child_edge_key).or_default(); + nc_edge.set_fitch_subs(split.composed); + nc_edge.indels = split.indels.composed; + } + + // Dense partitions carry no mutation lists, but they key node and edge state by graph key + // and must learn about the new node and edge. `apply_reroot` registers reroot-created keys + // the same way; mirroring it keeps dense state consistent when both families coexist. + for partition in dense { + let mut partition = partition.write_arc(); + partition.data.nodes.entry(n_key).or_insert_with(|| DenseNodePartition { + seq: DenseSeqInfo::default(), + profile: DenseSeqDistribution::default(), + }); + partition.data.edges.entry(un_edge_key).or_default(); + } + + Ok(n_key) +} + +/// One partition's substitution split of $M_v$ against $M_c$ (see [`hoist_reverting_child`]). +struct SubSplit { + /// $T$: parent positions untouched by the child. Goes to $u \to N$. + hoisted: Vec, + /// $H \cup R$: original parent entries at positions the child also touches. Goes to $N \to v$. + kept: Vec, + /// $H' \cup D$: $\mathrm{compose}(M_v, M_c)$ restricted to child positions. Goes to $N \to c$. + composed: Vec, +} + +/// Split parent-edge substitutions against one child edge into the hoist's three edges. +/// +/// Both inputs are position-sorted with at most one entry per position (the fitch-subs +/// invariant). The single merge-walk keeps every output position-sorted. +fn split_subs(parent_subs: &[Sub], child_subs: &[Sub]) -> Result { + debug_assert!( + parent_subs.is_sorted_by(|a, b| a.pos() < b.pos()), + "parent_subs not sorted by unique position" + ); + debug_assert!( + child_subs.is_sorted_by(|a, b| a.pos() < b.pos()), + "child_subs not sorted by unique position" + ); + + let mut hoisted = Vec::new(); + let mut kept = Vec::new(); + let mut composed = Vec::new(); + let mut pi = 0; + let mut ci = 0; + + while pi < parent_subs.len() && ci < child_subs.len() { + let ps = &parent_subs[pi]; + let cs = &child_subs[ci]; + match ps.pos().cmp(&cs.pos()) { + Ordering::Less => { + hoisted.push(ps.clone()); // T: parent-only position + pi += 1; + }, + Ordering::Greater => { + composed.push(cs.clone()); // D: child-only position + ci += 1; + }, + Ordering::Equal => { + debug_assert_eq!( + ps.qry(), + cs.reff(), + "Substitution chain broken at position {}: parent produces {} but child expects {}", + ps.pos(), + ps.qry(), + cs.reff() + ); + kept.push(ps.clone()); // H or R: keep original parent entry on N -> v + if ps.reff() != cs.qry() { + composed.push(Sub::new(ps.reff(), ps.pos(), cs.qry())?); // H': net a -> d on N -> c + } + // ps.reff() == cs.qry(): reversion, cancels, nothing on N -> c + pi += 1; + ci += 1; + }, + } + } + + hoisted.extend_from_slice(&parent_subs[pi..]); // remaining parent-only -> T + composed.extend_from_slice(&child_subs[ci..]); // remaining child-only -> D + + Ok(SubSplit { + hoisted, + kept, + composed, + }) +} + +/// Count reversions between one parent edge and one child edge (single partition). +fn count_reversions(parent_subs: &[Sub], child_subs: &[Sub]) -> usize { + debug_assert!( + parent_subs.is_sorted_by(|a, b| a.pos() < b.pos()), + "parent_subs not sorted by unique position" + ); + debug_assert!( + child_subs.is_sorted_by(|a, b| a.pos() < b.pos()), + "child_subs not sorted by unique position" + ); + + let mut count = 0; + let mut pi = 0; + let mut ci = 0; + while pi < parent_subs.len() && ci < child_subs.len() { + let ps = &parent_subs[pi]; + let cs = &child_subs[ci]; + match ps.pos().cmp(&cs.pos()) { + Ordering::Less => pi += 1, + Ordering::Greater => ci += 1, + Ordering::Equal => { + if ps.reff() == cs.qry() { + count += 1; + } + pi += 1; + ci += 1; + }, + } + } + count +} + +/// One partition's indel split for the hoist. +struct IndelSplit { + hoisted: Vec, + kept: Vec, + composed: Vec, +} + +/// Split parent-edge indels against a child edge, all-or-nothing. +/// +/// `compose_indels` merges overlapping and adjacent ranges rather than being position-keyed, +/// so an exact three-way split is not always definable. When no child indel overlaps or is +/// adjacent to any parent indel, the parent indels move cleanly above $N$ and $N \to c$ carries +/// only the child's own indels. Otherwise the parent indels stay on $N \to v$ and $N \to c$ +/// carries the full composition. Both branches preserve the root-to-$v$ and root-to-$c$ indel +/// content exactly. +fn split_indels(parent_indels: &[InDel], child_indels: &[InDel]) -> IndelSplit { + if indels_interact(parent_indels, child_indels) { + let mut parent = parent_indels.to_vec(); + let mut child = child_indels.to_vec(); + sort_indels(&mut parent); + sort_indels(&mut child); + let composed = compose_indels(&parent, &child); + IndelSplit { + hoisted: Vec::new(), + kept: parent, + composed, + } + } else { + IndelSplit { + hoisted: parent_indels.to_vec(), + kept: Vec::new(), + composed: child_indels.to_vec(), + } + } +} + +/// Whether any parent indel overlaps or is adjacent to any child indel. +/// +/// Adjacency matters because `compose_indels` merges touching deletions, so a clean hoist is +/// only possible when the ranges are strictly separated. +fn indels_interact(parent_indels: &[InDel], child_indels: &[InDel]) -> bool { + parent_indels.iter().any(|p| { + child_indels + .iter() + .any(|c| ranges_overlap_or_adjacent(p.range, c.range)) + }) +} + +/// Whether two half-open ranges overlap or touch at an endpoint. +fn ranges_overlap_or_adjacent((a_lo, a_hi): (usize, usize), (b_lo, b_hi): (usize, usize)) -> bool { + a_lo <= b_hi && b_lo <= a_hi +} + +/// Combined substitution and indel split for one partition. +struct EdgeSplit { + hoisted: Vec, + kept: Vec, + composed: Vec, + indels: IndelSplit, +} + +fn edge_branch_length(graph: &GraphAncestral, edge_key: GraphEdgeKey) -> f64 { + graph + .get_edge(edge_key) + .and_then(|edge| edge.read_arc().payload().read_arc().branch_length()) + .unwrap_or(0.0) +} + +fn set_edge_branch_length(graph: &GraphAncestral, edge_key: GraphEdgeKey, branch_length: f64) { + if let Some(edge) = graph.get_edge(edge_key) { + edge + .write_arc() + .payload() + .write_arc() + .set_branch_length(Some(branch_length)); + } +} diff --git a/packages/treetime/src/optimize/topology/merge_shared_mutations.rs b/packages/treetime/src/optimize/topology/merge_shared_mutations.rs index ac3eac74c..3794f55ba 100644 --- a/packages/treetime/src/optimize/topology/merge_shared_mutations.rs +++ b/packages/treetime/src/optimize/topology/merge_shared_mutations.rs @@ -61,7 +61,7 @@ pub fn merge_shared_mutation_branches( /// internal nodes that would result from repeated pairwise merging. /// /// Returns number of new internal nodes created. -fn merge_single_polytomy( +pub(crate) fn merge_single_polytomy( graph: &mut GraphAncestral, partitions: &[Arc>], node_key: GraphNodeKey, diff --git a/packages/treetime/src/optimize/topology/mod.rs b/packages/treetime/src/optimize/topology/mod.rs index 9f83792bc..9fff951b8 100644 --- a/packages/treetime/src/optimize/topology/mod.rs +++ b/packages/treetime/src/optimize/topology/mod.rs @@ -2,5 +2,7 @@ mod __tests__; pub mod collapse; +pub mod hoist_reversions; pub mod merge_shared_mutations; pub mod polytomy_nodes; +pub mod resolve_polytomy; diff --git a/packages/treetime/src/optimize/topology/resolve_polytomy.rs b/packages/treetime/src/optimize/topology/resolve_polytomy.rs new file mode 100644 index 000000000..99ce2e9f0 --- /dev/null +++ b/packages/treetime/src/optimize/topology/resolve_polytomy.rs @@ -0,0 +1,212 @@ +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; +use crate::optimize::topology::polytomy_nodes::find_polytomy_nodes; +use crate::partition::marginal::dense::partition::PartitionMarginalDense; +use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; +use crate::payload::ancestral::GraphAncestral; +use eyre::Report; +use log::debug; +use parking_lot::RwLock; +use std::collections::BTreeSet; +use std::sync::Arc; +use treetime_graph::edge::GraphEdgeKey; +use treetime_graph::node::GraphNodeKey; + +/// Resolve reversion-driven and shared-mutation polytomies across the whole tree. +/// +/// Runs [`resolve_one`] over every polytomy and repeats until a full pass changes nothing. +/// Each applied move (a shared-mutation merge, a reverting-child hoist, or a helper-node +/// retirement) strictly decreases the lexicographic potential (total fitch mutation count, +/// then node count), so the fixpoint is reached in a bounded number of rounds. Retirement can +/// turn a former helper's parent into a new polytomy, which the outer loop then picks up. +/// +/// Sparse-only: dense partitions carry no per-edge mutation lists, so the routine is inert +/// when no sparse partition is present, matching [`merge_shared_mutation_branches`]. +/// +/// Returns the number of polytomies whose local structure changed, summed across rounds; a +/// non-zero result means the caller must rebuild the graph and reassign node names. +/// +/// [`merge_shared_mutation_branches`]: crate::optimize::topology::merge_shared_mutations::merge_shared_mutation_branches +pub fn resolve_polytomies( + graph: &mut GraphAncestral, + sparse: &[Arc>], + dense: &[Arc>], +) -> Result { + if sparse.is_empty() { + return Ok(0); + } + + let mut total_changed = 0; + loop { + 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)? { + round_changed += 1; + } + } + if round_changed == 0 { + break; + } + total_changed += round_changed; + } + + if total_changed > 0 { + debug!("Resolved {total_changed} polytomies via merge/hoist/retire"); + } + + Ok(total_changed) +} + +/// Apply the merge -> hoist -> retire routine at one polytomy until it stops changing. +/// +/// 1. Merge siblings that share substitutions under helper nodes +/// ([`merge_single_polytomy`]). This canonicalizes several children reverting the same +/// position into a single reverting child. +/// 2. Hoist the child with the largest reversion count against the node's parent edge +/// ([`hoist_reverting_child`]), removing one mutation per reverted position. One child per +/// round keeps the move greedy and deterministic (ties broken by edge key). +/// 3. Retire helper nodes: collapse mutation-free edges whose target was created during this +/// invocation ([`retire_created_helpers`]). This dissolves the empty helper the hoist +/// leaves behind, turning its children into genuine siblings. +/// +/// The routine skips the root (no parent edge to revert). Nodes created before this call are +/// recorded so retirement never dissolves a pre-existing subtree root on a mutation-free edge. +/// +/// Returns whether anything changed. +fn resolve_one( + graph: &mut GraphAncestral, + sparse: &[Arc>], + dense: &[Arc>], + node_key: GraphNodeKey, +) -> 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)?; + let retired = retire_created_helpers(graph, sparse, dense, &preexisting)?; + + if !(merged || hoisted || retired) { + break; + } + any_changed = true; + } + + Ok(any_changed) +} + +/// Hoist the best reverting child of a node, if the node has a parent edge and any child +/// reverts one of its substitutions. Returns whether a hoist was applied. +fn try_hoist_reverting_child( + graph: &mut GraphAncestral, + sparse: &[Arc>], + dense: &[Arc>], + node_key: GraphNodeKey, +) -> Result { + let Some(parent_edge_key) = single_inbound_edge(graph, node_key) else { + return Ok(false); + }; + let Some(child_edge_key) = best_reverting_child(graph, sparse, node_key, parent_edge_key) else { + return Ok(false); + }; + hoist_reverting_child(graph, sparse, dense, parent_edge_key, child_edge_key)?; + Ok(true) +} + +/// Pick the child edge with the most reversions against the parent edge. +/// +/// Chooses the largest reversion count, breaking ties by smallest edge key for determinism. +/// Returns `None` when no child reverts any parent substitution. +fn best_reverting_child( + graph: &GraphAncestral, + sparse: &[Arc>], + node_key: GraphNodeKey, + parent_edge_key: GraphEdgeKey, +) -> Option { + let child_edges = graph.get_node(node_key)?.read_arc().outbound().to_vec(); + + let mut best: Option<(usize, GraphEdgeKey)> = None; + for child_edge_key in child_edges { + let reversions = count_child_reversions(sparse, parent_edge_key, child_edge_key); + if reversions == 0 { + continue; + } + let better = match best { + Some((best_reversions, best_key)) => { + best_reversions > reversions || (best_reversions == reversions && best_key <= child_edge_key) + }, + None => false, + }; + if !better { + best = Some((reversions, child_edge_key)); + } + } + + best.map(|(_, key)| key) +} + +/// Collapse mutation-free edges whose target was created during the current [`resolve_one`]. +/// +/// The hoist leaves an empty edge to the reverting child exactly when that child consisted of +/// nothing but reversions, which is the normal case after a merge. Collapsing it retires the +/// helper node and makes its children genuine siblings. +/// +/// Restricting the target to nodes created in this invocation is essential: a merge's helper +/// edges to pre-existing subtree roots are frequently mutation-free too, and collapsing those +/// would flatten topology the input asserted and discard its branch lengths. Within this scope +/// the branch-length-zero requirement of the loop's zero-optimal collapse is relaxed, because +/// these edges were synthesized moments earlier and carry no optimizer decision to override. +/// +/// Returns whether any edge was retired. +fn retire_created_helpers( + graph: &mut GraphAncestral, + sparse: &[Arc>], + dense: &[Arc>], + preexisting: &BTreeSet, +) -> Result { + let mut retired = false; + loop { + let candidate = graph.get_edges().iter().find_map(|edge| { + let edge = edge.read_arc(); + let target_key = edge.target(); + if preexisting.contains(&target_key) { + return None; + } + let target_is_leaf = graph.get_node(target_key).is_some_and(|node| node.read_arc().is_leaf()); + if target_is_leaf { + return None; + } + let edge_key = edge.key(); + let mutation_free = sparse.iter().all(|partition| { + let partition = partition.read_arc(); + match partition.edges.get(&edge_key) { + Some(edge_data) => edge_data.fitch_subs().is_empty() && edge_data.indels.is_empty(), + None => true, + } + }); + mutation_free.then_some(edge_key) + }); + + match candidate { + Some(edge_key) => { + collapse_edge(graph, sparse, dense, edge_key)?; + retired = true; + }, + None => break, + } + } + Ok(retired) +} + +/// The single parent edge of a node, or `None` for the root. +fn single_inbound_edge(graph: &GraphAncestral, node_key: GraphNodeKey) -> Option { + let node = graph.get_node(node_key)?; + let node = node.read_arc(); + match node.inbound() { + [edge_key] => Some(*edge_key), + _ => None, + } +} From 90221e3ebb0d5af004683835d24c50f5f30d3ff4 Mon Sep 17 00:00:00 2001 From: ivan-aksamentov Date: Wed, 2 Sep 2026 19:32:07 +0200 Subject: [PATCH 2/4] fix(optimize): guard reversion hoist against childless nodes - Require the polytomy node to keep at least two children before hoisting: moving the reverting child under the new node must leave a sibling behind, or the node becomes a childless stub (a spurious leaf). A node whose children all revert the same position merges to a single reverting child and is now left with its residual reversion rather than dissolved. --- .../src/optimize/topology/resolve_polytomy.rs | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/packages/treetime/src/optimize/topology/resolve_polytomy.rs b/packages/treetime/src/optimize/topology/resolve_polytomy.rs index 99ce2e9f0..5c96071b8 100644 --- a/packages/treetime/src/optimize/topology/resolve_polytomy.rs +++ b/packages/treetime/src/optimize/topology/resolve_polytomy.rs @@ -98,8 +98,15 @@ fn resolve_one( Ok(any_changed) } -/// Hoist the best reverting child of a node, if the node has a parent edge and any child -/// reverts one of its substitutions. Returns whether a hoist was applied. +/// Hoist the best reverting child of a node, if the node has a parent edge, at least two +/// children, and any child reverts one of its substitutions. Returns whether a hoist applied. +/// +/// The two-children requirement keeps the node a valid internal node: the hoist moves the +/// reverting child under the new node `N`, so the node must retain at least one other child +/// or it would become a childless stub. A node whose children all revert the same position is +/// merged into a single reverting child first, dropping it below this threshold; the residual +/// reversion on that single child is left in place (a greedy limitation, tracked in the +/// knowledge base), rather than removing the pre-existing node the input asserted. fn try_hoist_reverting_child( graph: &mut GraphAncestral, sparse: &[Arc>], @@ -109,6 +116,10 @@ fn try_hoist_reverting_child( let Some(parent_edge_key) = single_inbound_edge(graph, node_key) else { return Ok(false); }; + let degree_out = graph.get_node(node_key).map_or(0, |node| node.read_arc().degree_out()); + if degree_out < 2 { + return Ok(false); + } let Some(child_edge_key) = best_reverting_child(graph, sparse, node_key, parent_edge_key) else { return Ok(false); }; From 4e613151c422e5f3474aa095b28ee940e183550d Mon Sep 17 00:00:00 2001 From: ivan-aksamentov Date: Wed, 2 Sep 2026 19:23:10 +0200 Subject: [PATCH 3/4] test(optimize): cover the reversion hoist and polytomy resolution - Unit-test the hoist move on constructed trees: the untouched substitutions land once above the new node without duplication, chained positions compose onto both output edges, pure reversions drop one mutation, distances are preserved, and per-partition splits are independent. - Unit-test the indel rule for the cancelling, overlapping-fallback, and clean-hoist cases. - Integration-test the merge -> hoist -> retire routine end to end: the worked example reaches the parsimony optimum, incompatible splits stop at the greedy bound, helper retirement spares pre-existing internal nodes, and the root and reversion-free polytomies are left untouched. - Add a loop test proving the hoist fires through prune_and_merge_in_loop with no zero-optimal collapse. - Property-test that the potential never rises and strictly falls on any change, and that leaves, non-negative branch lengths, and the single-root tree shape are preserved. The committed regression seed pins the all-children-revert degenerate case. --- .../__tests__/test_prop_resolve_polytomy.txt | 7 + .../__tests__/test_topology_cleanup.rs | 61 +++ .../src/optimize/topology/__tests__/mod.rs | 3 + .../__tests__/test_hoist_reversions.rs | 431 ++++++++++++++++++ .../__tests__/test_prop_resolve_polytomy.rs | 194 ++++++++ .../__tests__/test_resolve_polytomy.rs | 257 +++++++++++ 6 files changed, 953 insertions(+) create mode 100644 packages/treetime/proptest-regressions/optimize/topology/__tests__/test_prop_resolve_polytomy.txt create mode 100644 packages/treetime/src/optimize/topology/__tests__/test_hoist_reversions.rs create mode 100644 packages/treetime/src/optimize/topology/__tests__/test_prop_resolve_polytomy.rs create mode 100644 packages/treetime/src/optimize/topology/__tests__/test_resolve_polytomy.rs diff --git a/packages/treetime/proptest-regressions/optimize/topology/__tests__/test_prop_resolve_polytomy.txt b/packages/treetime/proptest-regressions/optimize/topology/__tests__/test_prop_resolve_polytomy.txt new file mode 100644 index 000000000..ad76cbb91 --- /dev/null +++ b/packages/treetime/proptest-regressions/optimize/topology/__tests__/test_prop_resolve_polytomy.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc b8997137ea3a9ecf433f623ed27d82479f141625b2a9909d727b181938469386 # shrinks to n_children = 3, k = 1, revert_masks = [11, 29, 19], own_counts = [0, 0, 0] diff --git a/packages/treetime/src/optimize/__tests__/test_topology_cleanup.rs b/packages/treetime/src/optimize/__tests__/test_topology_cleanup.rs index 8b714224d..a6c5c3d49 100644 --- a/packages/treetime/src/optimize/__tests__/test_topology_cleanup.rs +++ b/packages/treetime/src/optimize/__tests__/test_topology_cleanup.rs @@ -390,6 +390,67 @@ mod tests { Ok(()) } + #[test] + fn test_optimize_prune_and_merge_hoists_reversion_without_collapse() -> Result<(), Report> { + // Reversion polytomy with no zero-optimal edge to collapse. The loop must still resolve + // it: merge C1+C2 (shared reversion), hoist the reverting group, retire the helper. + // Tree: root -> U -> V -> {C1, C2, C3}. U->V carries {A0T, C5G}; C1 and C2 revert A0T, + // C3 keeps it. Parsimony optimum is 2 mutations. + 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![]; + + // 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, &[])?; + assert!(changed, "reversion polytomy must be resolved even without a collapse"); + + 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, 2, "reaches the parsimony optimum"); + + 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 must be removed"); + + Ok(()) + } + #[test] fn test_optimize_cascading_collapse_parent_child_both_zero() -> Result<(), Report> { // Parent and child internal edges are both zero-optimal. diff --git a/packages/treetime/src/optimize/topology/__tests__/mod.rs b/packages/treetime/src/optimize/topology/__tests__/mod.rs index d0cab9fcb..c67fffdc5 100644 --- a/packages/treetime/src/optimize/topology/__tests__/mod.rs +++ b/packages/treetime/src/optimize/topology/__tests__/mod.rs @@ -1,4 +1,7 @@ mod test_collapse_edge; +mod test_hoist_reversions; mod test_merge_shared_mutations; mod test_prop_merge_shared_mutations; +mod test_prop_resolve_polytomy; mod test_reroot; +mod test_resolve_polytomy; diff --git a/packages/treetime/src/optimize/topology/__tests__/test_hoist_reversions.rs b/packages/treetime/src/optimize/topology/__tests__/test_hoist_reversions.rs new file mode 100644 index 000000000..ccfad5a67 --- /dev/null +++ b/packages/treetime/src/optimize/topology/__tests__/test_hoist_reversions.rs @@ -0,0 +1,431 @@ +#[cfg(test)] +mod tests { + use crate::optimize::topology::hoist_reversions::hoist_reverting_child; + use crate::partition::marginal::dense::partition::PartitionMarginalDense; + use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; + use crate::payload::ancestral::GraphAncestral; + use crate::seq::indel::InDel; + use crate::seq::indel::InDelKind; + use crate::seq::mutation::Sub; + use crate::test_utils::{find_edge_key, find_node_key_by_name}; + use approx::assert_abs_diff_eq; + use eyre::Report; + use parking_lot::RwLock; + use pretty_assertions::assert_eq; + use std::sync::Arc; + use treetime_graph::edge::{GraphEdgeKey, HasBranchLength}; + use treetime_graph::node::GraphNodeKey; + use treetime_io::nwk::nwk_read_str; + + use helpers::{Hoisted, c, edge_indels, edge_subs, make_partition, no_dense, sub}; + + // Tree: root -> U -> V -> {A, B, Z}. The hoist inserts N between U and V, grouping V with A. + const NWK: &str = "(((A:0.1,B:0.1,Z:0.1)V:0.2)U:0.1)root:0.0;"; + + #[test] + fn test_hoist_reversions_large_t_not_duplicated() -> Result<(), Report> { + // M_v has three substitutions; the child reverts only one. The two untouched + // substitutions (T) must land on u->N once and NOT be duplicated onto N->c + // (which distinguishes the move from re-attaching the child to the parent). + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let va = find_edge_key(&graph, "V", "A").unwrap(); + + let partition = make_partition( + &graph, + 0, + 100, + &[ + ( + "U", + "V", + vec![sub(b'A', 0, b'T'), sub(b'C', 5, b'G'), sub(b'G', 10, b'A')], + ), + ("V", "A", vec![sub(b'T', 0, b'A')]), + ], + ); + let sparse = vec![partition]; + + hoist_reverting_child(&mut graph, &sparse, &no_dense(), uv, va)?; + + let h = Hoisted::locate(&graph, "V", "A"); + let p = sparse[0].read_arc(); + assert_eq!(edge_subs(&p, h.un), vec![sub(b'C', 5, b'G'), sub(b'G', 10, b'A')]); + assert_eq!(edge_subs(&p, h.nv), vec![sub(b'A', 0, b'T')]); + assert_eq!(edge_subs(&p, h.nc), Vec::::new()); + + // V keeps its other children B and Z. + assert_eq!(graph.get_node(h.v).unwrap().read_arc().degree_out(), 2); + Ok(()) + } + + #[test] + fn test_hoist_reversions_chain_composed() -> Result<(), Report> { + // Chain: parent A0T at pos 0, child T0G at pos 0 -> net A0G. The original A0T stays + // on N->v; the composed A0G moves to N->c. + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let va = find_edge_key(&graph, "V", "A").unwrap(); + + let partition = make_partition( + &graph, + 0, + 100, + &[ + ("U", "V", vec![sub(b'A', 0, b'T')]), + ("V", "A", vec![sub(b'T', 0, b'G')]), + ], + ); + let sparse = vec![partition]; + + hoist_reverting_child(&mut graph, &sparse, &no_dense(), uv, va)?; + + let h = Hoisted::locate(&graph, "V", "A"); + let p = sparse[0].read_arc(); + assert_eq!(edge_subs(&p, h.un), Vec::::new()); + assert_eq!(edge_subs(&p, h.nv), vec![sub(b'A', 0, b'T')]); + assert_eq!(edge_subs(&p, h.nc), vec![sub(b'A', 0, b'G')]); + Ok(()) + } + + #[test] + fn test_hoist_reversions_reversion_removed_reduces_count() -> Result<(), Report> { + // Pure reversion: A0T then T0A. Two mutations before, one after (delta = -1). + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let va = find_edge_key(&graph, "V", "A").unwrap(); + + let partition = make_partition( + &graph, + 0, + 100, + &[ + ("U", "V", vec![sub(b'A', 0, b'T')]), + ("V", "A", vec![sub(b'T', 0, b'A')]), + ], + ); + let sparse = vec![partition]; + + let before = helpers::total_subs(&graph, &sparse[0].read_arc()); + hoist_reverting_child(&mut graph, &sparse, &no_dense(), uv, va)?; + let after = helpers::total_subs(&graph, &sparse[0].read_arc()); + + assert_eq!(before, 2); + assert_eq!(after, 1); + + let h = Hoisted::locate(&graph, "V", "A"); + let p = sparse[0].read_arc(); + assert_eq!(edge_subs(&p, h.nv), vec![sub(b'A', 0, b'T')]); + assert_eq!(edge_subs(&p, h.nc), Vec::::new()); + Ok(()) + } + + #[test] + fn test_hoist_reversions_branch_length_distance_preserved() -> Result<(), Report> { + // Distances root->V and root->A are unchanged by the move; the parent edge is split + // proportionally to substitution count (|T|/|M_v| = 2/3). + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let va = find_edge_key(&graph, "V", "A").unwrap(); + let ru = find_edge_key(&graph, "root", "U").unwrap(); + + let partition = make_partition( + &graph, + 0, + 100, + &[ + ( + "U", + "V", + vec![sub(b'A', 0, b'T'), sub(b'C', 5, b'G'), sub(b'G', 10, b'A')], + ), + ("V", "A", vec![sub(b'T', 0, b'A')]), + ], + ); + let sparse = vec![partition]; + + hoist_reverting_child(&mut graph, &sparse, &no_dense(), uv, va)?; + + let h = Hoisted::locate(&graph, "V", "A"); + let bl = |ek: GraphEdgeKey| helpers::branch_length(&graph, ek); + let root_to_v = bl(ru) + bl(h.un) + bl(h.nv); + let root_to_a = bl(ru) + bl(h.un) + bl(h.nc); + + assert_abs_diff_eq!(root_to_v, 0.1 + 0.2, epsilon = 1e-9); + assert_abs_diff_eq!(root_to_a, 0.1 + 0.2 + 0.1, epsilon = 1e-9); + Ok(()) + } + + #[test] + fn test_hoist_reversions_multi_partition() -> Result<(), Report> { + // Two partitions revert independent positions. Each partition's edges are split on + // its own positions; T is per-partition (present in p0, empty in p1). + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let va = find_edge_key(&graph, "V", "A").unwrap(); + + let p0 = make_partition( + &graph, + 0, + 100, + &[ + ("U", "V", vec![sub(b'A', 0, b'T'), sub(b'G', 10, b'C')]), + ("V", "A", vec![sub(b'T', 0, b'A')]), + ], + ); + let p1 = make_partition( + &graph, + 1, + 100, + &[ + ("U", "V", vec![sub(b'C', 5, b'G')]), + ("V", "A", vec![sub(b'G', 5, b'C')]), + ], + ); + let sparse = vec![p0, p1]; + + hoist_reverting_child(&mut graph, &sparse, &no_dense(), uv, va)?; + + let h = Hoisted::locate(&graph, "V", "A"); + let g0 = sparse[0].read_arc(); + assert_eq!(edge_subs(&g0, h.un), vec![sub(b'G', 10, b'C')]); + assert_eq!(edge_subs(&g0, h.nv), vec![sub(b'A', 0, b'T')]); + assert_eq!(edge_subs(&g0, h.nc), Vec::::new()); + + let g1 = sparse[1].read_arc(); + assert_eq!(edge_subs(&g1, h.un), Vec::::new()); + assert_eq!(edge_subs(&g1, h.nv), vec![sub(b'C', 5, b'G')]); + assert_eq!(edge_subs(&g1, h.nc), Vec::::new()); + Ok(()) + } + + #[test] + fn test_hoist_reversions_indel_cancellation() -> Result<(), Report> { + // A deletion on the parent edge and its inverse insertion on the child edge interact, + // so the parent indel stays on N->v and the composition cancels on N->c. + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let va = find_edge_key(&graph, "V", "A").unwrap(); + + let partition = make_partition( + &graph, + 0, + 100, + &[ + ("U", "V", vec![sub(b'A', 0, b'T')]), + ("V", "A", vec![sub(b'T', 0, b'A')]), + ], + ); + let del = InDel::del((20, 23), [c(b'A'), c(b'A'), c(b'A')].as_slice())?; + let ins = InDel::ins((20, 23), [c(b'A'), c(b'A'), c(b'A')].as_slice())?; + { + let mut p = partition.write_arc(); + p.edges.get_mut(&uv).unwrap().indels = vec![del.clone()]; + p.edges.get_mut(&va).unwrap().indels = vec![ins]; + } + let sparse = vec![partition]; + + hoist_reverting_child(&mut graph, &sparse, &no_dense(), uv, va)?; + + let h = Hoisted::locate(&graph, "V", "A"); + let p = sparse[0].read_arc(); + assert_eq!(edge_indels(&p, h.un), Vec::::new()); + assert_eq!(edge_indels(&p, h.nv), vec![del]); + assert_eq!(edge_indels(&p, h.nc), Vec::::new()); + Ok(()) + } + + #[test] + fn test_hoist_reversions_indel_overlap_fallback() -> Result<(), Report> { + // Overlapping deletions cannot be cleanly hoisted: the parent deletion stays on N->v + // and the merged deletion lands on N->c. + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let va = find_edge_key(&graph, "V", "A").unwrap(); + + let partition = make_partition( + &graph, + 0, + 100, + &[ + ("U", "V", vec![sub(b'A', 0, b'T')]), + ("V", "A", vec![sub(b'T', 0, b'A')]), + ], + ); + let parent_del = InDel::del((20, 25), [c(b'A'); 5].as_slice())?; + let child_del = InDel::del((22, 28), [c(b'A'); 6].as_slice())?; + { + let mut p = partition.write_arc(); + p.edges.get_mut(&uv).unwrap().indels = vec![parent_del.clone()]; + p.edges.get_mut(&va).unwrap().indels = vec![child_del]; + } + let sparse = vec![partition]; + + hoist_reverting_child(&mut graph, &sparse, &no_dense(), uv, va)?; + + let h = Hoisted::locate(&graph, "V", "A"); + let p = sparse[0].read_arc(); + assert_eq!(edge_indels(&p, h.un), Vec::::new()); + assert_eq!(edge_indels(&p, h.nv), vec![parent_del]); + + let nc = edge_indels(&p, h.nc); + assert_eq!(nc.len(), 1); + assert_eq!(nc[0].range, (20, 28)); + assert_eq!(nc[0].kind, InDelKind::Deletion); + Ok(()) + } + + #[test] + fn test_hoist_reversions_indel_no_interaction_hoisted() -> Result<(), Report> { + // A parent indel disjoint from the child's indels is hoisted cleanly to u->N, leaving + // N->v free of indels and N->c carrying only the child's own indel. + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let uv = find_edge_key(&graph, "U", "V").unwrap(); + let va = find_edge_key(&graph, "V", "A").unwrap(); + + let partition = make_partition( + &graph, + 0, + 100, + &[ + ("U", "V", vec![sub(b'A', 0, b'T')]), + ("V", "A", vec![sub(b'T', 0, b'A')]), + ], + ); + let parent_del = InDel::del((20, 23), [c(b'A'); 3].as_slice())?; + let child_del = InDel::del((50, 53), [c(b'A'); 3].as_slice())?; + { + let mut p = partition.write_arc(); + p.edges.get_mut(&uv).unwrap().indels = vec![parent_del.clone()]; + p.edges.get_mut(&va).unwrap().indels = vec![child_del.clone()]; + } + let sparse = vec![partition]; + + hoist_reverting_child(&mut graph, &sparse, &no_dense(), uv, va)?; + + let h = Hoisted::locate(&graph, "V", "A"); + let p = sparse[0].read_arc(); + assert_eq!(edge_indels(&p, h.un), vec![parent_del]); + assert_eq!(edge_indels(&p, h.nv), Vec::::new()); + assert_eq!(edge_indels(&p, h.nc), vec![child_del]); + Ok(()) + } + + mod helpers { + use super::*; + use crate::alphabet::alphabet::{Alphabet, AlphabetName}; + use crate::gtr::get_gtr::{JC69Params, jc69}; + use crate::partition::storage::sparse::{SparseEdgePartition, SparseNodePartition}; + use maplit::btreemap; + use treetime_primitives::{AsciiChar, Seq, seq}; + + pub fn c(b: u8) -> AsciiChar { + AsciiChar::from_byte_unchecked(b) + } + + pub fn sub(reff: u8, pos: usize, qry: u8) -> Sub { + Sub::new(c(reff), pos, c(qry)).unwrap() + } + + pub fn no_dense() -> Vec>> { + vec![] + } + + /// Edge keys of the three edges the hoist produces, located by the child/node names. + pub struct Hoisted { + pub v: GraphNodeKey, + pub un: GraphEdgeKey, + pub nv: GraphEdgeKey, + pub nc: GraphEdgeKey, + } + + impl Hoisted { + pub fn locate(graph: &GraphAncestral, v_name: &str, c_name: &str) -> Self { + let v = find_node_key_by_name(graph, v_name).unwrap(); + let c = find_node_key_by_name(graph, c_name).unwrap(); + let nv = single_inbound(graph, v); + let n = graph.get_source_node_key(nv).unwrap(); + let un = single_inbound(graph, n); + let nc = single_inbound(graph, c); + Self { v, un, nv, nc } + } + } + + fn single_inbound(graph: &GraphAncestral, node_key: GraphNodeKey) -> GraphEdgeKey { + let node = graph.get_node(node_key).unwrap(); + let node = node.read_arc(); + match node.inbound() { + [edge_key] => *edge_key, + other => panic!("expected exactly one inbound edge, found {}", other.len()), + } + } + + pub fn edge_subs(partition: &PartitionMarginalSparse, edge_key: GraphEdgeKey) -> Vec { + partition.edges[&edge_key].fitch_subs().to_vec() + } + + pub fn edge_indels(partition: &PartitionMarginalSparse, edge_key: GraphEdgeKey) -> Vec { + partition.edges[&edge_key].indels.clone() + } + + pub fn total_subs(graph: &GraphAncestral, partition: &PartitionMarginalSparse) -> usize { + graph + .get_edges() + .iter() + .filter_map(|e| partition.edges.get(&e.read_arc().key())) + .map(|e| e.fitch_subs().len()) + .sum() + } + + pub fn branch_length(graph: &GraphAncestral, edge_key: GraphEdgeKey) -> f64 { + graph + .get_edge(edge_key) + .and_then(|e| e.read_arc().payload().read_arc().branch_length()) + .unwrap() + } + + pub fn make_partition( + graph: &GraphAncestral, + index: usize, + length: usize, + edge_mutations: &[(&str, &str, Vec)], + ) -> Arc> { + let mut partition = PartitionMarginalSparse { + index, + gtr: jc69(JC69Params::default()).unwrap(), + alphabet: Alphabet::new(AlphabetName::Nuc).unwrap(), + length, + nodes: btreemap! {}, + edges: btreemap! {}, + root_sequence: seq![], + }; + + let mut ref_seq: Seq = std::iter::repeat_with(|| c(b'A')).take(length).collect(); + for (_, _, subs) in edge_mutations { + for s in subs { + if s.pos() < length { + ref_seq[s.pos()] = s.reff(); + } + } + } + partition.root_sequence = ref_seq.clone(); + + for node in graph.get_nodes() { + let key = node.read_arc().key(); + let mut node_part = SparseNodePartition::empty(&partition.alphabet); + node_part.seq.sequence = ref_seq.clone(); + partition.nodes.insert(key, node_part); + } + + for (source, target, subs) in edge_mutations { + let edge_key = + find_edge_key(graph, source, target).unwrap_or_else(|| panic!("edge {source}->{target} missing")); + partition + .edges + .insert(edge_key, SparseEdgePartition::with_fitch_subs(subs.clone())); + } + + Arc::new(RwLock::new(partition)) + } + } +} 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 new file mode 100644 index 000000000..5d6a9dd74 --- /dev/null +++ b/packages/treetime/src/optimize/topology/__tests__/test_prop_resolve_polytomy.rs @@ -0,0 +1,194 @@ +#[cfg(test)] +mod tests { + use crate::optimize::topology::resolve_polytomy::resolve_polytomies; + use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; + use crate::payload::ancestral::GraphAncestral; + use crate::seq::mutation::Sub; + use parking_lot::RwLock; + use proptest::prelude::*; + use std::collections::BTreeSet; + use std::sync::Arc; + use treetime_graph::edge::HasBranchLength; + + proptest! { + /// The routine never increases the total mutation count, and when it changes anything + /// the count strictly decreases: the first component of the (mutation count, node count) + /// potential falls on every applied merge or hoist. This also exercises termination - + /// a non-decreasing potential would hang the test. + #[test] + fn test_prop_resolve_polytomy_potential_decreases( + n_children in 3_usize..7, + k in 1_usize..5, + revert_masks in prop::collection::vec(0_u32..32, 3..7), + own_counts in prop::collection::vec(0_usize..3, 3..7), + ) { + 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 after = helpers::total_subs(&graph, &sparse[0].read_arc()); + + prop_assert!(after <= before, "mutation count increased: before={before} after={after}"); + if changed > 0 { + prop_assert!(after < before, "changed but mutation count did not fall: before={before} after={after}"); + } + } + + /// Structural invariants hold after resolution: leaves are preserved, branch lengths stay + /// non-negative, and the result is still a single-rooted tree (each non-root node keeps + /// exactly one parent edge). + #[test] + fn test_prop_resolve_polytomy_preserves_tree( + n_children in 3_usize..7, + k in 1_usize..5, + revert_masks in prop::collection::vec(0_u32..32, 3..7), + own_counts in prop::collection::vec(0_usize..3, 3..7), + ) { + let (mut graph, partition, _before) = helpers::build_case(n_children, k, &revert_masks, &own_counts); + let leaves_before = helpers::leaf_names(&graph); + let sparse = vec![partition]; + + resolve_polytomies(&mut graph, &sparse, &helpers::no_dense()).unwrap(); + + prop_assert_eq!(helpers::leaf_names(&graph), leaves_before); + + for edge in graph.get_edges() { + if let Some(bl) = edge.read_arc().payload().read_arc().branch_length() { + prop_assert!(bl >= 0.0, "negative branch length {bl}"); + } + } + + let mut roots = 0; + for node in graph.get_nodes() { + let inbound = node.read_arc().inbound().len(); + if inbound == 0 { + roots += 1; + } else { + prop_assert_eq!(inbound, 1, "non-root node has {} parents", inbound); + } + } + prop_assert_eq!(roots, 1, "expected exactly one root"); + } + } + + mod helpers { + use super::*; + use crate::alphabet::alphabet::{Alphabet, AlphabetName}; + use crate::gtr::get_gtr::{JC69Params, jc69}; + use crate::partition::marginal::dense::partition::PartitionMarginalDense; + use crate::partition::storage::sparse::{SparseEdgePartition, SparseNodePartition}; + use crate::test_utils::find_edge_key; + use itertools::Itertools; + use maplit::btreemap; + use treetime_io::nwk::nwk_read_str; + use treetime_primitives::{AsciiChar, Seq, seq}; + + fn c(b: u8) -> AsciiChar { + AsciiChar::from_byte_unchecked(b) + } + + pub fn no_dense() -> Vec>> { + vec![] + } + + pub fn leaf_names(graph: &GraphAncestral) -> BTreeSet { + graph + .get_nodes() + .iter() + .filter(|n| n.read_arc().is_leaf()) + .filter_map(|n| n.read_arc().payload().read_arc().name.clone()) + .collect() + } + + pub fn total_subs(graph: &GraphAncestral, partition: &PartitionMarginalSparse) -> usize { + graph + .get_edges() + .iter() + .filter_map(|e| partition.edges.get(&e.read_arc().key())) + .map(|e| e.fitch_subs().len()) + .sum() + } + + /// Build a `root -> U -> V -> {C0..}` case. The parent edge U->V carries `k` + /// substitutions `A->C` at positions `0..k`. Each child reverts the M_v positions set in + /// its mask and adds its own distinct substitutions at high positions. Returns the graph, + /// the sparse partition, and the initial total mutation count. + pub fn build_case( + n_children: usize, + k: usize, + revert_masks: &[u32], + own_counts: &[usize], + ) -> (GraphAncestral, Arc>, usize) { + let n_children = n_children.min(revert_masks.len()).min(own_counts.len()).max(3); + let length = 200_usize; + + let child_names: Vec = (0..n_children).map(|i| format!("C{i}")).collect(); + let inner = child_names.iter().map(|name| format!("{name}:0.1")).join(","); + let newick = format!("((({inner})V:0.2)U:0.1)root:0.0;"); + let graph: GraphAncestral = nwk_read_str(&newick).unwrap(); + + // M_v: k substitutions A->C at positions 0..k. + let parent_subs: Vec = (0..k).map(|pos| Sub::new(c(b'A'), pos, c(b'C')).unwrap()).collect(); + + let mut edge_mutations: Vec<(String, String, Vec)> = vec![("U".to_owned(), "V".to_owned(), parent_subs)]; + + let mut own_pos = k + 1; + for (i, name) in child_names.iter().enumerate() { + let mask = revert_masks[i]; + let mut subs: Vec = Vec::new(); + for pos in 0..k { + if mask & (1 << pos) != 0 { + // Revert the parent substitution: C->A. + subs.push(Sub::new(c(b'C'), pos, c(b'A')).unwrap()); + } + } + for _ in 0..own_counts[i] { + subs.push(Sub::new(c(b'G'), own_pos, c(b'T')).unwrap()); + own_pos += 2; + } + subs.sort_by_key(Sub::pos); + edge_mutations.push(("V".to_owned(), name.clone(), subs)); + } + + let total: usize = edge_mutations.iter().map(|(_, _, subs)| subs.len()).sum(); + let partition = make_partition(&graph, length, &edge_mutations); + (graph, partition, total) + } + + fn make_partition( + graph: &GraphAncestral, + length: usize, + edge_mutations: &[(String, String, Vec)], + ) -> Arc> { + let mut partition = PartitionMarginalSparse { + index: 0, + gtr: jc69(JC69Params::default()).unwrap(), + alphabet: Alphabet::new(AlphabetName::Nuc).unwrap(), + length, + nodes: btreemap! {}, + edges: btreemap! {}, + root_sequence: seq![], + }; + + let ref_seq: Seq = std::iter::repeat_with(|| c(b'A')).take(length).collect(); + partition.root_sequence = ref_seq.clone(); + + for node in graph.get_nodes() { + let key = node.read_arc().key(); + let mut node_part = SparseNodePartition::empty(&partition.alphabet); + node_part.seq.sequence = ref_seq.clone(); + partition.nodes.insert(key, node_part); + } + + for (source, target, subs) in edge_mutations { + let edge_key = + find_edge_key(graph, source, target).unwrap_or_else(|| panic!("edge {source}->{target} missing")); + partition + .edges + .insert(edge_key, SparseEdgePartition::with_fitch_subs(subs.clone())); + } + + Arc::new(RwLock::new(partition)) + } + } +} diff --git a/packages/treetime/src/optimize/topology/__tests__/test_resolve_polytomy.rs b/packages/treetime/src/optimize/topology/__tests__/test_resolve_polytomy.rs new file mode 100644 index 000000000..041478235 --- /dev/null +++ b/packages/treetime/src/optimize/topology/__tests__/test_resolve_polytomy.rs @@ -0,0 +1,257 @@ +#[cfg(test)] +mod tests { + use crate::optimize::topology::resolve_polytomy::resolve_polytomies; + use crate::partition::marginal::dense::partition::PartitionMarginalDense; + use crate::partition::marginal::sparse::partition::PartitionMarginalSparse; + use crate::payload::ancestral::GraphAncestral; + use crate::seq::mutation::Sub; + use crate::test_utils::find_node_key_by_name; + use eyre::Report; + use parking_lot::RwLock; + use pretty_assertions::assert_eq; + use std::sync::Arc; + + use helpers::{no_dense, reversion_present, sub, total_subs}; + use treetime_io::nwk::nwk_read_str; + + // root -> U -> V -> {C1, C2, C3}. V is the polytomy under test. + const NWK: &str = "(((C1:0.1,C2:0.1,C3:0.1)V:0.2)U:0.1)root:0.0;"; + + #[test] + fn test_resolve_polytomy_merge_hoist_retire_worked_example() -> Result<(), Report> { + // M_v = {A0T (p), C5G (q)}; C1 and C2 both revert p, C3 keeps it. The routine merges + // C1+C2, hoists the reverting group, and retires the helper, reaching the parsimony + // optimum of two mutations (q above, p only on the C3 lineage). + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let partition = helpers::make_partition( + &graph, + 0, + 100, + &[ + ("U", "V", vec![sub(b'A', 0, b'T'), sub(b'C', 5, b'G')]), + ("V", "C1", vec![sub(b'T', 0, b'A')]), + ("V", "C2", vec![sub(b'T', 0, b'A')]), + ("V", "C3", vec![]), + ], + ); + let sparse = vec![partition]; + + let changed = resolve_polytomies(&mut graph, &sparse, &no_dense())?; + assert!(changed > 0); + + let p = sparse[0].read_arc(); + assert_eq!(total_subs(&graph, &p), 2); + assert!(!reversion_present(&graph, &p, &sub(b'T', 0, b'A'))); + + for leaf in ["C1", "C2", "C3"] { + assert!( + find_node_key_by_name(&graph, leaf).is_some(), + "leaf {leaf} must survive" + ); + } + Ok(()) + } + + #[test] + fn test_resolve_polytomy_incompatible_splits_five_to_four() -> Result<(), Report> { + // C1 reverts p1, C2 reverts p2 (different positions): the two required splits are + // incompatible. One hoist takes the total from 5 to 4; the residual reversion is + // irreducible homoplasy, and the routine stops there. + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let partition = helpers::make_partition( + &graph, + 0, + 100, + &[ + ( + "U", + "V", + vec![sub(b'A', 0, b'T'), sub(b'C', 5, b'G'), sub(b'G', 10, b'A')], + ), + ("V", "C1", vec![sub(b'T', 0, b'A')]), + ("V", "C2", vec![sub(b'G', 5, b'C')]), + ("V", "C3", vec![]), + ], + ); + let sparse = vec![partition]; + + let before = total_subs(&graph, &sparse[0].read_arc()); + resolve_polytomies(&mut graph, &sparse, &no_dense())?; + let after = total_subs(&graph, &sparse[0].read_arc()); + + assert_eq!(before, 5); + assert_eq!(after, 4); + Ok(()) + } + + #[test] + fn test_resolve_polytomy_retirement_preserves_preexisting_internal_node() -> Result<(), Report> { + // W is a pre-existing internal node reached by a mutation-free edge from V. Helper + // retirement must dissolve only nodes it created, never W, even though V->W is empty. + let mut graph: GraphAncestral = nwk_read_str("((((X1:0.1,X2:0.1)W:0.0,C1:0.1,C2:0.1)V:0.2)U:0.1)root:0.0;")?; + let partition = helpers::make_partition( + &graph, + 0, + 100, + &[ + ("U", "V", vec![sub(b'A', 0, b'T')]), + ("V", "C1", vec![sub(b'T', 0, b'A')]), + ("V", "C2", vec![sub(b'T', 0, b'A')]), + ("V", "W", vec![]), + ], + ); + let sparse = vec![partition]; + + resolve_polytomies(&mut graph, &sparse, &no_dense())?; + + assert!( + find_node_key_by_name(&graph, "W").is_some(), + "pre-existing internal node W must survive helper retirement" + ); + for leaf in ["X1", "X2", "C1", "C2"] { + assert!( + find_node_key_by_name(&graph, leaf).is_some(), + "leaf {leaf} must survive" + ); + } + let p = sparse[0].read_arc(); + assert_eq!(total_subs(&graph, &p), 1); + assert!(!reversion_present(&graph, &p, &sub(b'T', 0, b'A'))); + Ok(()) + } + + #[test] + fn test_resolve_polytomy_root_polytomy_skipped() -> Result<(), Report> { + // A polytomy at the root has no parent edge to revert, so no hoist fires. With no + // shared substitutions there is nothing to do; the routine leaves the tree untouched. + let mut graph: GraphAncestral = nwk_read_str("(A:0.1,B:0.1,C:0.1)root;")?; + let partition = helpers::make_partition( + &graph, + 0, + 100, + &[ + ("root", "A", vec![sub(b'A', 0, b'T')]), + ("root", "B", vec![sub(b'C', 5, b'G')]), + ("root", "C", vec![sub(b'G', 10, b'A')]), + ], + ); + let sparse = vec![partition]; + let nodes_before = graph.get_nodes().len(); + + let changed = resolve_polytomies(&mut graph, &sparse, &no_dense())?; + + assert_eq!(changed, 0); + assert_eq!(graph.get_nodes().len(), nodes_before); + assert_eq!(total_subs(&graph, &sparse[0].read_arc()), 3); + Ok(()) + } + + #[test] + fn test_resolve_polytomy_no_change_without_reversions() -> Result<(), Report> { + // Distinct, non-shared, non-reverting child substitutions: nothing to merge or hoist. + let mut graph: GraphAncestral = nwk_read_str(NWK)?; + let partition = helpers::make_partition( + &graph, + 0, + 100, + &[ + ("U", "V", vec![sub(b'A', 0, b'T')]), + ("V", "C1", vec![sub(b'C', 5, b'G')]), + ("V", "C2", vec![sub(b'G', 10, b'A')]), + ("V", "C3", vec![sub(b'T', 15, b'A')]), + ], + ); + let sparse = vec![partition]; + let nodes_before = graph.get_nodes().len(); + + let changed = resolve_polytomies(&mut graph, &sparse, &no_dense())?; + + assert_eq!(changed, 0); + assert_eq!(graph.get_nodes().len(), nodes_before); + assert_eq!(total_subs(&graph, &sparse[0].read_arc()), 4); + Ok(()) + } + + mod helpers { + use super::*; + use crate::alphabet::alphabet::{Alphabet, AlphabetName}; + use crate::gtr::get_gtr::{JC69Params, jc69}; + use crate::partition::storage::sparse::{SparseEdgePartition, SparseNodePartition}; + use crate::test_utils::find_edge_key; + use maplit::btreemap; + use treetime_primitives::{AsciiChar, Seq, seq}; + + pub fn c(b: u8) -> AsciiChar { + AsciiChar::from_byte_unchecked(b) + } + + pub fn sub(reff: u8, pos: usize, qry: u8) -> Sub { + Sub::new(c(reff), pos, c(qry)).unwrap() + } + + pub fn no_dense() -> Vec>> { + vec![] + } + + pub fn total_subs(graph: &GraphAncestral, partition: &PartitionMarginalSparse) -> usize { + graph + .get_edges() + .iter() + .filter_map(|e| partition.edges.get(&e.read_arc().key())) + .map(|e| e.fitch_subs().len()) + .sum() + } + + pub fn reversion_present(graph: &GraphAncestral, partition: &PartitionMarginalSparse, needle: &Sub) -> bool { + graph + .get_edges() + .iter() + .filter_map(|e| partition.edges.get(&e.read_arc().key())) + .any(|e| e.fitch_subs().contains(needle)) + } + + pub fn make_partition( + graph: &GraphAncestral, + index: usize, + length: usize, + edge_mutations: &[(&str, &str, Vec)], + ) -> Arc> { + let mut partition = PartitionMarginalSparse { + index, + gtr: jc69(JC69Params::default()).unwrap(), + alphabet: Alphabet::new(AlphabetName::Nuc).unwrap(), + length, + nodes: btreemap! {}, + edges: btreemap! {}, + root_sequence: seq![], + }; + + let mut ref_seq: Seq = std::iter::repeat_with(|| c(b'A')).take(length).collect(); + for (_, _, subs) in edge_mutations { + for s in subs { + if s.pos() < length { + ref_seq[s.pos()] = s.reff(); + } + } + } + partition.root_sequence = ref_seq.clone(); + + for node in graph.get_nodes() { + let key = node.read_arc().key(); + let mut node_part = SparseNodePartition::empty(&partition.alphabet); + node_part.seq.sequence = ref_seq.clone(); + partition.nodes.insert(key, node_part); + } + + for (source, target, subs) in edge_mutations { + let edge_key = + find_edge_key(graph, source, target).unwrap_or_else(|| panic!("edge {source}->{target} missing")); + partition + .edges + .insert(edge_key, SparseEdgePartition::with_fitch_subs(subs.clone())); + } + + Arc::new(RwLock::new(partition)) + } + } +} From 020c1fc9fb395e50bfa5e46956f832492774e52f Mon Sep 17 00:00:00 2001 From: ivan-aksamentov Date: Wed, 2 Sep 2026 19:31:28 +0200 Subject: [PATCH 4/4] docs(kb): update kb --- .../optimize-polytomy-reversion-resolution.md | 50 +++ kb/features/optimize.md | 11 + ...e-reversion-hoist-single-child-residual.md | 28 ++ .../optimize-polytomy-reversion-resolution.md | 310 ------------------ ...timetree-stochastic-polytomy-resolution.md | 24 +- 5 files changed, 101 insertions(+), 322 deletions(-) create mode 100644 kb/decisions/optimize-polytomy-reversion-resolution.md create mode 100644 kb/issues/M-optimize-reversion-hoist-single-child-residual.md delete mode 100644 kb/proposals/optimize-polytomy-reversion-resolution.md diff --git a/kb/decisions/optimize-polytomy-reversion-resolution.md b/kb/decisions/optimize-polytomy-reversion-resolution.md new file mode 100644 index 000000000..7a2af134d --- /dev/null +++ b/kb/decisions/optimize-polytomy-reversion-resolution.md @@ -0,0 +1,50 @@ +# Optimize loop resolves reversion-driven polytomies + +The `optimize` loop adds a topology move with no v0 counterpart: when a polytomy forces a substitution onto an internal edge and then reverts it on a child, the move relocates the child so the reversion cancels. This removes homoplasy that an arbitrary binary resolution of a polytomy introduces, and that v0's loop never targets. + +## What v0 does + +v0's branch-length optimization loop (`TreeAnc.optimize_tree_marginal`) collapses short internal branches (`prune_short_branches`) but has no step that inspects or removes reversions. A reversion is only ever dropped as an incidental side effect when an edge is collapsed for an unrelated reason and its substitutions compose with the child's. v0 has no shared-mutation merge and no reversion-aware polytomy handling. + +## What v1 does + +Once per iteration, [`prune_and_merge_in_loop()`](../../packages/treetime/src/optimize/run_loop.rs#L319) runs a single per-polytomy routine, [`resolve_polytomies()`](../../packages/treetime/src/optimize/topology/resolve_polytomy.rs#L31), after the zero-optimal edge collapse. Sparse partitions only; dense partitions carry no per-edge mutation lists. The routine applies three steps at each polytomy and iterates to a fixpoint: + +- Merge siblings that share substitutions under a helper node ([`merge_single_polytomy()`](../../packages/treetime/src/optimize/topology/merge_shared_mutations.rs#L64)). This canonicalizes several children reverting the same position into one reverting child. +- Hoist the reverting child with the largest reversion count ([`hoist_reverting_child()`](../../packages/treetime/src/optimize/topology/hoist_reversions.rs#L74)), ties broken by edge key. +- Retire helper nodes: collapse mutation-free edges to nodes the routine created in this pass. + +The hoist inserts a new node $N$ between parent $u$ and node $v$, grouping $v$ with one reverting child $c$. The parent-edge substitutions $M_v$ partition by position against the child substitutions $M_c$ into three disjoint sets: $T$ (untouched by the child), $H$ (chained, $a \to b$ then $b \to d$), and $R$ (reverted, $a \to b$ then $b \to a$). The resulting edges carry: + +| edge | substitutions | +| --------- | ------------------------------------------------------------------------------------------------------ | +| $u \to N$ | $T$ | +| $N \to v$ | $H \cup R$ | +| $N \to c$ | $\mathrm{compose}(M_v, M_c)$ at $M_c$ positions ($H' \cup D$, where $D$ are the child's own positions) | + +The net substitution change is $\Delta = -\lvert R\rvert$: one mutation removed per reverted position, none added. The two relocated edges keep their keys via [`Graph::reparent_edge()`](../../packages/treetime-graph/src/graph_ops.rs#L123), so their partition entries stay valid and only their content is rewritten. + +Branch lengths split in proportion to substitution count so that root-to-$v$ and root-to-$c$ distances are preserved exactly and re-fit on the next iteration: $b(u \to N) = b(u \to v)\cdot\lvert T\rvert / \lvert M_v\rvert$, $b(N \to v) = b(u \to v) - b(u \to N)$, $b(N \to c) = b(N \to v) + b(v \to c)$. This deliberately differs from the Jukes-Cantor recomputation used by the `prune` merge ([prune-merge-jukes-cantor-branch-length.md](prune-merge-jukes-cantor-branch-length.md)), which would discard the converged branch lengths inside the loop. + +Indels use an all-or-nothing rule, because `compose_indels` merges overlapping and adjacent ranges rather than being position-keyed. When no child indel overlaps or is adjacent to a parent indel, the parent indels move cleanly above $N$ and $N \to c$ carries the child's own indels; otherwise the parent indels stay on $N \to v$ and $N \to c$ carries the full composition. Both branches preserve the root-to-$v$ and root-to-$c$ indel content, and the substitution gain is unaffected. + +## Why v1 differs + +An arbitrary binary resolution of a polytomy can spend two mutations (a substitution plus its reversion) where one suffices, and tree builders produce such resolutions routinely. Fitch's forward pass never emits a reversion on a binary node's child edge, so the pattern is inherently a polytomy phenomenon; that is why the move belongs in the per-polytomy routine rather than as a general edge-level rewrite. Inserting a node beats re-attaching the child directly to the parent, which would duplicate $T$ onto the re-attached edge and give $\Delta = \lvert R\rvert - \lvert T\rvert$, almost never a gain because $\lvert T\rvert$ is typically much larger than $\lvert R\rvert$. + +The routine runs every iteration, independent of whether a collapse fired, so reversions present in the input and in polytomies formed by earlier iterations are resolved. It cannot oscillate with the zero-optimal collapse: the lexicographic (total fitch mutation count, node count) potential strictly decreases on every applied merge, hoist, or retirement, even though the hoist deliberately relocates mutation-carrying edges that the collapse step refuses to touch. + +## Greedy limitation + +When every child of a polytomy reverts the same position, the merge reduces them to a single reverting child and the hoist declines (moving that child would leave the node childless). One residual reversion remains rather than dissolving a pre-existing node ([kb/issues/M-optimize-reversion-hoist-single-child-residual.md](../issues/M-optimize-reversion-hoist-single-child-residual.md)). Disagreeing reverting children likewise leave irreducible homoplasy; exact resolution is subtree parsimony re-optimization, which is out of scope. + +## Affected commands + +- [`optimize`](../../packages/treetime/src/optimize/run_loop.rs#L319) - runs the routine in its per-iteration topology-cleanup step +- `prune` is unchanged: it still calls the shared-mutation merge directly and does not adopt the reversion hoist + +## Tests + +- Unit (hoist mechanics): [test_hoist_reversions.rs](../../packages/treetime/src/optimize/topology/__tests__/test_hoist_reversions.rs) - no duplication of $T$, chained composition, reversion removal, distance preservation, per-partition splits, indel cases. +- Integration (merge to hoist to retire): [test_resolve_polytomy.rs](../../packages/treetime/src/optimize/topology/__tests__/test_resolve_polytomy.rs) - the worked example reaching the parsimony optimum, incompatible splits stopping at the greedy bound, helper retirement sparing pre-existing internal nodes, root and reversion-free polytomies left untouched. +- Property: [test_prop_resolve_polytomy.rs](../../packages/treetime/src/optimize/topology/__tests__/test_prop_resolve_polytomy.rs) - the potential never rises and strictly falls on any change, and leaves, non-negative branch lengths, and the single-root tree shape are preserved. diff --git a/kb/features/optimize.md b/kb/features/optimize.md index ec431cb19..593bb8dc5 100644 --- a/kb/features/optimize.md +++ b/kb/features/optimize.md @@ -58,6 +58,17 @@ v0 uses Brent's method (`scipy.optimize.minimize_scalar`) in sqrt(t) space with - [ ] Bifurcating root special handling (v0 optimizes combined root-children length, preserves ratio) - [ ] Convergence by sequence change count (v0 joint mode: stops when zero nucleotides change) +## Topology Cleanup (v1-only) + +Once per iteration, `prune_and_merge_in_loop` simplifies the tree. Sparse partitions only; dense partitions carry no per-edge mutation lists, so the reversion and merge steps are inert under a dense-only run. + +- [x] Zero-optimal edge collapse: internal edges the per-edge optimizer drove to exactly zero, carrying no substitutions or indels, are contracted into polytomies. Mirrors v0's `prune_short_branches` inside the loop. +- [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. + +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)). + ## GTR Integration - [x] `--model` flag wired through `get_gtr_sparse()`/`get_gtr_dense()` for all named models and inference diff --git a/kb/issues/M-optimize-reversion-hoist-single-child-residual.md b/kb/issues/M-optimize-reversion-hoist-single-child-residual.md new file mode 100644 index 000000000..ba15e3c15 --- /dev/null +++ b/kb/issues/M-optimize-reversion-hoist-single-child-residual.md @@ -0,0 +1,28 @@ +# Reversion hoist leaves a residual reversion when every child reverts the same position + +## Symptom and reproduction + +In the optimize loop, `resolve_polytomies` resolves a reversion polytomy by merging siblings that share substitutions, then hoisting the reverting group under a new node. When _every_ child of a polytomy node reverts the same parent-edge substitution, the shared-mutation merge groups all of them under one helper node, leaving the polytomy node with a single child. The hoist then declines to fire, because moving that single reverting child under a new node would leave the polytomy node childless (a spurious leaf). + +The result keeps one residual reversion on the merged child edge. The total mutation count still drops relative to the input (the merge removes the shared reversion from every duplicate), but it does not reach the parsimony optimum, which would remove the reversion entirely. + +Reproduction: a node `V` under parent edge `U -> V = {A0C}` whose children all carry `C0A`. After the merge, `V` has a single child (the helper) carrying `C0A`; the optimum is zero mutations (drop `A0C` and the reversion), but the routine stops at two (`A0C` on `U -> V`, `C0A` on the helper edge). + +## Impact and scope + +Narrow. It affects only polytomies where all children revert the same position, and only leaves a reversion that the merge already reduced from many to one. The output is never worse than the input and never invalid. Correctness (no added mutations, distance preservation, single-rooted tree) is unaffected. + +## Root cause + +`fn try_hoist_reverting_child()` requires the node to keep at least two children so the hoisted child leaves a sibling behind. A single-child node drops below that threshold. The optimal resolution here is to collapse the now-degree-two node into its parent (composing the two edges cancels the reversion), but collapsing a pre-existing node contradicts the routine's helper-retirement scope, which deliberately never dissolves nodes the input asserted. + +## Fix approach + +Recognize a degree-two node whose single child reverts its parent edge and collapse that node's parent edge (compose `M_v` with `M_c`, cancelling the reversion). This removes a topologically vacuous node (a single-child node carries no phylogenetic split), so it preserves all bipartitions. Requires distinguishing this case from the helper-retirement scope, and a decision on whether removing pre-existing degree-two nodes during `optimize` is acceptable. + +This is a greedy limitation in the same family as the incompatible-splits case the routine already accepts (disagreeing reverting children leave irreducible homoplasy); exact resolution is subtree parsimony re-optimization, which is out of scope. + +## Locations + +- `fn try_hoist_reverting_child()` two-children guard: `packages/treetime/src/optimize/topology/resolve_polytomy.rs` +- Hoist move: `packages/treetime/src/optimize/topology/hoist_reversions.rs` diff --git a/kb/proposals/optimize-polytomy-reversion-resolution.md b/kb/proposals/optimize-polytomy-reversion-resolution.md deleted file mode 100644 index 8f5ae4728..000000000 --- a/kb/proposals/optimize-polytomy-reversion-resolution.md +++ /dev/null @@ -1,310 +0,0 @@ -# Proposal: reversion-driven polytomy resolution in the optimize loop - -## Summary - -Extend the optimize loop's topology cleanup with a third local move: when a child edge of -node $v$ carries the reversion of a substitution mapped to $v$'s parent edge, insert a new -node above $v$ that groups $v$ and that child, moving the non-reverted mutations above the -new node. The move removes exactly one mutation per reverted position and never adds any. - -Combined with the existing shared-mutation merge and a helper-node cleanup step, this turns -per-polytomy topology cleanup into a single three-step routine: **merge shared mutations → -hoist reverting child → retire helper nodes**. - -## Current state - -`fn run_optimize_loop()` performs topology cleanup once per iteration through -`fn prune_and_merge_in_loop()` -[`packages/treetime/src/optimize/run_loop.rs#L309-L353`](../../packages/treetime/src/optimize/run_loop.rs#L309-L353): - -1. `fn find_zero_optimal_internal_edges()` collects internal edges the per-edge optimizer - drove to exactly zero, excluding edges that carry substitutions or indels. -2. `fn collapse_edge()` contracts each such edge, producing polytomies. -3. `fn merge_shared_mutation_branches()` groups siblings in those polytomies that share - identical substitutions under a new internal node. - -Step 3 is gated on step 2 having fired: `prune_and_merge_in_loop` returns early when the -zero-optimal edge list is empty. - -Nothing in the loop acts on reversions. A reversion is currently only ever removed as a side -effect of `fn compose_substitutions()` -[`packages/treetime/src/seq/mutation.rs#L176-L227`](../../packages/treetime/src/seq/mutation.rs#L176-L227) -when an edge happens to be collapsed for an unrelated reason. - -## Motivation - -An arbitrary binary resolution of a polytomy can force a substitution onto an internal edge -and then require its reversion on one child. The pattern costs two mutations where a -different resolution costs one. Tree builders produce such resolutions routinely, and the -optimize loop already creates fresh polytomies by collapsing zero-length edges, so the -pattern arises both in the input and during the run. - -Fitch's forward pass resolves an internal node toward the parent state when its children -conflict, so a *binary* node never emerges from reconstruction with a reversion on one child -edge — the reconstruction would place a single mutation on the other child instead. The -pattern is therefore inherently a **polytomy** phenomenon, which is why the move belongs in -the per-polytomy routine rather than as a general edge-level rewrite. - -## Move definition - -Let $v$ be an internal non-root node, $u$ its parent, $e_p = u \to v$ carrying substitutions -$M_v$, and $e_c = v \to c$ a child edge carrying $M_c$. Because every edge holds at most one -substitution per position (asserted in `fn compose_substitutions()`), $M_v$ partitions by -position against $M_c$ into three disjoint sets: - -- $R$ — positions where $M_c$ holds the exact inverse (**reversions**) -- $H$ — positions where $M_c$ holds a different substitution (**chains**: $a \to b$ then $b \to d$) -- $T$ — the remainder of $M_v$, untouched by the child - -Let $D$ be the positions of $M_c$ absent from $M_v$, and let $H'$ denote the composed -chains $\{a \to d\}$ for $p \in H$. - -The move inserts a new node $N$ between $u$ and $v$: - -| edge | substitutions | count | -| --- | --- | --- | -| $u \to N$ | $T$ | $\lvert T\rvert$ | -| $N \to v$ | $H \cup R$ | $\lvert H\rvert + \lvert R\rvert$ | -| $N \to c$ | $H' \cup D$ | $\lvert H\rvert + \lvert D\rvert$ | - -Before the move the two edges carry $(\lvert T\rvert + \lvert H\rvert + \lvert R\rvert) + -(\lvert R\rvert + \lvert H\rvert + \lvert D\rvert)$. After, they carry -$\lvert T\rvert + 2\lvert H\rvert + \lvert R\rvert + \lvert D\rvert$. - -$$\Delta = -\lvert R\rvert$$ - -**Trigger: $R \neq \emptyset$. Gain: $\lvert R\rvert$, independent of $\lvert T\rvert$.** - -$N \to c$ is exactly $\mathrm{compose}(M_v, M_c) \setminus T$, so it can be produced with the -existing `fn chain_fitch_subs()` rather than new composition logic. - -### Why not re-attach the child directly to the parent - -The simpler move — detach $c$ and make it a child of $u$ with -$\mathrm{compose}(M_v, M_c) = T \cup H' \cup D$ — gives -$\Delta = \lvert R\rvert - \lvert T\rvert$, because it duplicates $T$ onto the re-attached -edge. On real data $\lvert T\rvert$ is typically much larger than $\lvert R\rvert$, so that -form would almost never fire. It is rejected. - -`fn merge_shared_mutation_branches()` would in principle recover the node-insertion state -from the re-attached state on a later iteration: the shared set between $u \to v = M_v$ and -$u \to c = T \cup H' \cup D$ is precisely $T$. But it cannot when $u$ is left with two -children — `fn find_polytomy_nodes()` requires `degree_out > 2` and -`fn merge_single_polytomy()` breaks at `child_edges.len() <= 2` -[`packages/treetime/src/optimize/topology/merge_shared_mutations.rs#L64-L100`](../../packages/treetime/src/optimize/topology/merge_shared_mutations.rs#L64-L100). -That is a wrong fixed point, not a delayed one. Inserting $N$ directly avoids it and also -avoids optimizing branch lengths against the inflated intermediate tree for one iteration. - -### Branch lengths - -Distance-preserving split, proportional to substitution count: - -$$b(u \to N) = b(u \to v)\cdot\frac{\lvert T\rvert}{\lvert M_v\rvert},\quad -b(N \to v) = b(u \to v) - b(u \to N),\quad -b(N \to c) = b(N \to v) + b(v \to c)$$ - -Root-to-$v$ and root-to-$c$ distances are both preserved exactly; the next iteration re-fits. - -This deliberately differs from `fn merge_sibling_group()`, which recomputes branch lengths -from Jukes-Cantor distance -[`packages/treetime/src/optimize/topology/merge_shared_mutations.rs#L312-L331`](../../packages/treetime/src/optimize/topology/merge_shared_mutations.rs#L312-L331). -That is appropriate for `prune`, which has no optimizer state, but inside the optimize loop it -would discard converged branch lengths. - -### Indels - -`fn compose_indels()` merges overlapping and adjacent ranges rather than being position-keyed, -so an exact three-way split is not always definable. Conservative rule: hoist $e_p$'s indels -to $u \to N$ only when the child's indels neither cancel nor overlap them; otherwise leave -them on $N \to v$ and let $N \to c$ carry $\mathrm{compose}(I_v, I_c)$. Always correct, and -leaves at most $\lvert I_v\rvert$ unclaimed. The substitution gain $\lvert R\rvert$ is -unaffected either way. - -## Interaction with the shared-mutation merge - -Running the merge **first** canonicalizes the input to the hoist. When several siblings revert -the same position, their edges literally share the mutation $b \to a$ at $p$, so the merge -groups them under a helper node. The hoist then sees a single reverting child. - -Worked example: $M_v = \{p, q\}$, children $c_1$ and $c_2$ both reverting $p$, $c_3$ keeping it. - -| step | edges | total | -| --- | --- | --- | -| start | $u\to v=\{p,q\}$, $v\to c_1=\{\lnot p\}$, $v\to c_2=\{\lnot p\}$, $v\to c_3=\emptyset$ | 4 | -| 1. merge | $u\to v=\{p,q\}$, $v\to N'=\{\lnot p\}$, $N'\to c_1=\emptyset$, $N'\to c_2=\emptyset$ | 3 | -| 2. hoist $N'$ | $u\to N=\{q\}$, $N\to v=\{p\}$, $N\to N'=\emptyset$, $v\to c_3=\emptyset$ | 2 | -| 3. retire $N'$ | $u\to N=\{q\}$, $N\to v=\{p\}$, $N\to c_1=\emptyset$, $N\to c_2=\emptyset$ | 2 | - -The parsimony optimum is 2 ($q$ on every lineage, $p$ only on $c_3$'s). The routine reaches it. - -Step 3 is not cosmetic. The hoist produces an **empty** $N \to c$ edge exactly when the child -edge consisted of nothing but reversions, which is the normal case after a merge. Collapsing -it is what dissolves the helper node and makes $c_1$ and $c_2$ genuine siblings of $v$. The -merge creates the helper, the hoist relocates it, the collapse retires it. - -The merge cannot undo a hoist: after a merge places shared substitution $s$ on the group's -parent edge, $s$ is removed from every group member's edge, and since an edge holds at most -one substitution per position, no member can then carry $s$'s reversion. - -## Greedy limits - -When reverting children disagree on *which* positions they revert, one hoist is often already -optimal. With $M_v = \{p_1, p_2, q\}$, $c_1$ reverting $p_1$, $c_2$ reverting $p_2$, $c_3$ -keeping both, hoisting $c_1$ takes the total from 5 to 4. No tree does better: $p_1$ is needed -by $\{c_2, c_3\}$ and $p_2$ by $\{c_1, c_3\}$, which are incompatible splits, so the residual -reversion on $v \to c_2$ is irreducible homoplasy. - -Genuine greedy failures exist. With $c_1$ reverting $\{p_1, p_2\}$ and $c_2$ reverting -$\{p_2, p_3\}$, greedy reaches 5 against an optimum of 4. Exact resolution is subtree -parsimony re-optimization and is out of scope; `fn merge_shared_mutation_branches()` carries -the same greedy caveat via `fn greedy_disjoint_group_matching()`. - -Rule adopted: **one reverting child per node per round**, choosing the child with the largest -$\lvert R\rvert$, ties broken by edge key for determinism. Iterate the node to a fixpoint. - -## Termination - -Potential function: (total fitch mutation count, node count), lexicographic. - -- hoist: strictly decreases the first component by $\lvert R\rvert \geq 1$ -- merge: strictly decreases the first component by $\text{total\_shared}\cdot(k-1)$ -- helper retirement and zero-optimal collapse: first component unchanged, second decreases - -No cycling. This has to be stated explicitly because the new move deliberately relocates -mutation-carrying edges, which `fn find_zero_optimal_internal_edges()` documents as previously -forbidden precisely to avoid a merge/collapse oscillation -[`packages/treetime/src/optimize/run_loop.rs#L270-L272`](../../packages/treetime/src/optimize/run_loop.rs#L270-L272). - -## Design - -### New module - -`packages/treetime/src/optimize/topology/resolve_polytomy.rs` - -```rust -/// Merge → hoist → retire, at one polytomy. Returns whether anything changed. -fn resolve_one( - graph: &mut GraphAncestral, - sparse_partitions: &[Arc>], - dense_partitions: &[Arc>], - node_key: GraphNodeKey, -) -> Result; - -/// Drive `resolve_one` over all polytomies to a fixpoint. Returns moves applied. -pub fn resolve_polytomies( - graph: &mut GraphAncestral, - sparse_partitions: &[Arc>], - dense_partitions: &[Arc>], -) -> Result; -``` - -`packages/treetime/src/optimize/topology/hoist_reversions.rs` holds the $(T, H, R)$ split and -the node insertion. The split is a single merge-walk over two position-sorted `Vec`, -the same shape as `fn compose_substitutions()`. - -Sparse-only, matching `fn merge_shared_mutation_branches()`: dense partitions carry no -mutation lists. `optimize` builds either a sparse or a dense partition, never both -[`packages/treetime/src/optimize/pipeline.rs#L93-L102`](../../packages/treetime/src/optimize/pipeline.rs#L93-L102), -so the routine is inert under `--dense`. - -### Helper retirement scope - -Step 3 collapses a mutation-free internal edge **only when its target is a node created -during this invocation of `resolve_one`**. Track created keys in a local set. - -The restriction matters: a merge's $N' \to c_i$ edges are frequently mutation-free too, but -their targets are pre-existing subtree roots. Collapsing those would flatten topology the -input asserted and discard its branch lengths. - -Within that scope the rule relaxes `fn find_zero_optimal_internal_edges()`, which additionally -requires `bl == 0`. Justified because these edges were synthesized moments earlier by the -routine itself, so there is no optimizer decision to override. - -### Partition bookkeeping - -Node insertion follows `fn merge_sibling_group()`: `add_node`, `add_edge`, `remove_edge`, then -rewrite the sparse `partition.edges` entries. No `treetime-graph` change is required. - -Reparenting via a new `Graph::reparent_edge()` primitive (mutating `set_source`, as -`Graph::collapse_edge()` already does internally) would avoid edge-key churn and preserve -dense edge state. Deferred: `fn merge_sibling_group()` already establishes remove-and-add as -the accepted pattern, and matching it keeps the two routines uniform. - -Register the new node and edge keys in dense partitions, mirroring -`PartitionMarginalDense::apply_reroot()` -[`packages/treetime/src/partition/marginal_dense.rs#L96-L120`](../../packages/treetime/src/partition/marginal_dense.rs#L96-L120). -`fn merge_sibling_group()` omits this; it is harmless today only because dense and sparse -partitions never coexist in `optimize`. Not worth adding a second instance of the same latent -trap. - -Stale `msg_to_parent` / `msg_to_child` / `msg_from_child` on rewritten edges need no handling: -`fn marginal_backward()` overwrites them wholesale -[`packages/treetime/src/partition/marginal_passes.rs#L386`](../../packages/treetime/src/partition/marginal_passes.rs#L386), -which is what the collapse path already relies on. - -### Loop integration - -`fn prune_and_merge_in_loop()` becomes: - -1. zero-optimal collapse (existing, optimizer-driven — it is what creates polytomies) -2. `resolve_polytomies()` — **every iteration**, no longer gated on step 1 having fired -3. `graph.build()` + `assign_node_names()` - -`fn merge_shared_mutation_branches()` is no longer called directly from the loop; -`fn merge_single_polytomy()` is promoted to `pub(crate)` for reuse by `resolve_one`. The -public wrapper stays for `prune` -[`packages/treetime/src/prune/pipeline.rs#L79`](../../packages/treetime/src/prune/pipeline.rs#L79). - -Each applied move sets `topology_changed`, which resets `best_lh`. The monotone potential -bounds how many iterations can keep firing. - -## Open axes - -- **`prune` adoption.** `prune` already calls the merge and would benefit from the same - routine, but adopting it changes `prune`'s output. Separate decision. -- **Accepted wart.** In the common case the merge creates a helper node that step 3 destroys - moments later. Fusing the two passes to skip the transient would entangle them for no - parsimony gain. -- **Per-iteration cost.** The pass is $O(\text{edges} \times \text{mutations per edge})$ per - round. `Graph::remove_edge()` scans all nodes, so node insertion is $O(V)$ — inherited from - `fn merge_sibling_group()`. Measure before optimizing. - -## Validation plan - -Unit tests in `packages/treetime/src/optimize/topology/__tests__/`, reusing the fixture -helpers from `test_collapse_edge.rs` (`fn make_sparse_partition()`, -`fn populate_test_nodes()`, `fn nwk_read_str()`): - -- polytomy with large $\lvert T\rvert$ — asserts $T$ is *not* duplicated, distinguishing the - move from re-attachment to the parent -- chained ($H$) positions composed correctly on both output edges -- the worked merge → hoist → retire example above, end to end, asserting the final mutation - count is 2 -- incompatible-splits example: total goes 5 → 4 and stops -- helper retirement does not dissolve pre-existing internal nodes on mutation-free edges -- indel cancellation, and the overlapping-indel fallback -- root skipped; distance preservation of root-to-node paths -- multi-partition - -Loop-level: - -- planted reversion converges with the reversion removed, no oscillation, iteration count not - regressed -- property test in the style of `test_prop_merge_shared_mutations.rs`: the potential - (mutation count, node count) decreases lexicographically on every applied move - -Datasets: - -- ebola/20 and sc2/4500: count moves applied, total mutation count before and after, final - log-likelihood, iterations to convergence -- confirm no change on trees with no reversions - -## Related - -- [kb/features/optimize.md](../features/optimize.md) — needs a topology section; the existing - shared-mutation merge is undocumented there -- [kb/proposals/optimize-short-branch-pruning.md](optimize-short-branch-pruning.md) — - complementary post-loop cleanup on branch length rather than mutation content -- [kb/decisions/prune-merge-jukes-cantor-branch-length.md](../decisions/prune-merge-jukes-cantor-branch-length.md) — - the JC branch-length convention this proposal declines to reuse inside the loop -- [kb/decisions/timetree-no-zero-branch-collapse-in-loop.md](../decisions/timetree-no-zero-branch-collapse-in-loop.md) diff --git a/kb/proposals/timetree-stochastic-polytomy-resolution.md b/kb/proposals/timetree-stochastic-polytomy-resolution.md index bfe07cc12..78e8f6cdd 100644 --- a/kb/proposals/timetree-stochastic-polytomy-resolution.md +++ b/kb/proposals/timetree-stochastic-polytomy-resolution.md @@ -64,16 +64,16 @@ parent's time is reached; survivors attach to the parent as a residual polytomy. This is the main translation hazard. v0 works in `time_before_present`, increasing into the past; v1 uses calendar time with `parent.time < child.time`. Every comparison inverts: -| v0 | v1 | -| --- | --- | -| `tmax = parent.time_before_present` | `t_stop = parent.time` (lower bound) | -| sort children ascending by `time_before_present` | sort **descending** by `time` | -| `t` starts at most recent child, increases | `t` starts at `max(child.time)`, **decreases** | -| `while t < tmax` | `while t > t_stop` | -| `t += dt` | `t -= dt` | -| pop when `t > to_come[0].tbp` | pop when `t < to_come[0].time` | -| `b.branch_length = tmax - b.time` | `time_length = b.time - t_stop` | -| early return `if t >= tmax` | early return if `max(child.time) <= t_stop` | +| v0 | v1 | +| ------------------------------------------------ | ---------------------------------------------- | +| `tmax = parent.time_before_present` | `t_stop = parent.time` (lower bound) | +| sort children ascending by `time_before_present` | sort **descending** by `time` | +| `t` starts at most recent child, increases | `t` starts at `max(child.time)`, **decreases** | +| `while t < tmax` | `while t > t_stop` | +| `t += dt` | `t -= dt` | +| pop when `t > to_come[0].tbp` | pop when `t < to_come[0].time` | +| `b.branch_length = tmax - b.time` | `time_length = b.time - t_stop` | +| early return `if t >= tmax` | early return if `max(child.time) <= t_stop` | ### Rates @@ -297,5 +297,5 @@ characterise the change in resolved topology. - [kb/issues/N-timetree-unused-cli-flags.md](../issues/N-timetree-unused-cli-flags.md) -- `--keep-polytomies` - [kb/features/timetree.md](../features/timetree.md) -- polytomy resolution checklist -- [kb/proposals/optimize-polytomy-reversion-resolution.md](optimize-polytomy-reversion-resolution.md) -- - the `optimize` loop's polytomy handling, which also wants `Graph::reparent_edge` +- [kb/decisions/optimize-polytomy-reversion-resolution.md](../decisions/optimize-polytomy-reversion-resolution.md) -- + the `optimize` loop's polytomy handling, which also uses `Graph::reparent_edge`