Refactor/unified coalescent

[rneher]

No description provided.

[ivan-aksamentov]

⚠️ AI-generated content below. Verify all claims.

First review of #869. Refactor/unified coalescent

• Head: 6783e0115d51b94d515d48053be89589fd910a0c on refactor/unified-coalescent
• Base: rust
• Scope: unify constant-Tc and skyline coalescent estimation (constant is the one-segment skyline), reparametrize the skyline objective to log-Tc, and make coalescent estimation fail loudly instead of falling back.

Overview

Click to expand

The branch collapses optimize_tc into a thin facade over optimize_skyline(n_points = 1), reparametrizes the skyline penalty from the inverse time scale x = 1/Tc to the log scale z = ln Tc (so smoothing charges squared log-fold-changes and is scale-independent), and replaces the previous per-round fallback with error propagation so a coalescent estimation failure stops the run. Changes touch coalescent/skyline.rs, coalescent/optimize_tc.rs, timetree/pipeline.rs, commands/timetree/args.rs, the coalescent tests, and several KB decision and issue files.

Observed

The math is correct: objective, gradient, tridiagonal Hessian, warm start $z_i = \ln(I_i/M_i)$, and the $T_c = I/M$ closed form all check out [src]. Unification is clean. Suite is green.

The solver is described as exact and numerically infallible, but solve_log_tc returns Ok unconditionally, so max_iter = 0, an exhausted budget, or a stalled line search all report the un-optimized warm start as success (F2, red repro branch below).

The log-$T_c$ rewrite also dropped the stiffness <= 0 short-circuit, leaving the documented stiffness > 0 precondition unenforced (F1). I built a counterexample expecting non-positive stiffness to emit a NaN; it did not. The solver degrades to the finite decoupled warm start, so F1 is a validation and non-convexity concern, not the singular-solve correctness bug the automated reviewers reported. It moves to non-blocking accordingly.

Several open correctness questions are already in the author's KB. They're credited under Author-tracked below, with reviewer extensions where the analysis goes beyond the existing note.

Background

Kingman coalescent constant-Tc estimator [click to expand]

Under the Kingman coalescent [1] the constant-Tc negative log-likelihood is $M \ln T_c + I / T_c$ where $I = \int k(k-1)/2,dt$ is the pairwise-rate integral over the tree and $M$ is the merger count. It is strictly convex in $\ln T_c$ with the unique minimizer $T_c = I/M$. The skyline generalizes this to a piecewise-constant $T_c(t)$ with a smoothness penalty across adjacent segments, analogous in spirit to the log effective population size formulation used by Bayesian GMRF skygrid methods (Gill et al. 2013 [2]; Hill and Baele 2019 [3]). Writing $z_i = \ln T_{c,i}$, the branch minimizes the regularized negative log-likelihood

$$C(z) = \sum_i \left( I_i e^{-z_i} + M_i z_i \right) + \frac{\gamma}{2} \sum_i (z_{i+1} - z_i)^2$$

where $T_c$ is the coalescent time scale (inverse merger rate), $z_i = \ln T_{c,i}$ is the optimizer's variable, $I_i = \int_{\text{seg } i} k(k-1)/2,dt$ is the pairwise-rate integral over segment $i$ (and $I$ the same integral over the whole tree), $M_i$ is the merger count in segment $i$ (and $M$ the count over the whole tree), $k$ is the number of extant lineages, and $\gamma$ is the smoothing stiffness on adjacent log-$T_c$ differences.

Blocking issues

Correctness concerns worth resolving before merge.

🔴 F2. Newton non-convergence and line-search stall are returned as success [click to expand]

solve_log_tc returns Ok(z) unconditionally after the iteration loop [src]. max_iter == 0 returns the warm start without optimizing; exhausting max_iter does not check the final gradient norm [src]; and an Armijo line search reaching alpha < 1e-12 breaks the inner loop leaving z unchanged [src], after which the outer loop recomputes an identical step and spins to max_iter, then returns the pre-stall iterate.

Effect: a non-converged or stalled solve is reported as an optimized skyline. This contradicts the pipeline comment that the solves "are exact" and "fail only on a tree that is degenerate for the coalescent" [src] and the decision doc's "cannot fail numerically" claim, and it undercuts the fail-loud intent of this branch because the pipeline can only propagate an error the solver never raises. The problem is convex so this should be rare, but "should converge" is not verification.

Repro (red counterexample, base is this PR head): #871 drives a three-segment skyline with max_iter = 0 and asserts an error; on the current head it returns Ok with the un-optimized warm start.

Suggestions:

• after the loop, recompute the gradient infinity-norm and return an error (or set a converged flag on SkylineResult) when it exceeds tolerance.
• treat a stalled line search (alpha < 1e-12 with no accepted step) as a failure rather than continuing; validate max_iter >= 1 and finite tolerance > 0.

Non-blocking issues

Test deficiencies, reinvention, and convention drift. Worth addressing but not merge-blocking.

🟡 F1. stiffness > 0 precondition is unenforced; negative stiffness is non-convex [click to expand]

The rewrite changed the solver short-circuit from if n == 1 || stiffness <= 0.0 to if n == 1 [src], so for n_points > 1 any stiffness <= 0 enters the Newton loop. SkylineParams::stiffness documents "Must be positive for more than one segment" [src], but nothing enforces it: optimize_skyline validates only n_points, and --skyline-stiffness is a bare f64 with no value_parser [src]. A negative stiffness makes the penalty $-\tfrac{|s|}{2}\sum(\Delta \ln T_c)^2$, unbounded below, so the objective is non-convex and the Thomas solve loses the positive-definite Hessian it assumes.

Effect: an unenforced precondition, and a latent ill-posed solve for stiffness < 0. The automated reviewers claimed this produces a NaN/garbage Tc; the repro below shows it does not on a representative tree -- the Armijo line search rejects the indefinite Newton step and optimize_skyline returns the finite decoupled warm start (stiffness = 0 is in fact behavior-preserving, since the warm start is already the minimizer there). So the concern is input validation and a non-convex objective reaching the solver, not an observed wrong number. That downgrades it from the blocking severity the automated tools assigned.

Repro (red counterexample, base is this PR head): #870 asserts a five-segment solve with stiffness = -1.0 is rejected; on the current head it returns Ok.

Suggestions:

• reject non-finite or <= 0 stiffness for n_points > 1 at the optimize_skyline boundary, and add a clap value_parser enforcing > 0.
• optionally restore the stiffness <= 0 fast path as defense in depth for programmatic callers.

🟡 F7. NaN-panicking float sort; hand-written Default [click to expand]

merger_quantile_boundaries sorts merger times with a.partial_cmp(b).expect("merger times must be comparable") [src], which panics on a NaN time and violates the project float-sort rule (OrderedFloat [doc] / total_cmp). Separately, impl Default for SkylineParams is a plain field-default block [src] that SmartDefault [doc] (already a dependency, used in args.rs) generates.

Suggestions:

• times.sort_by_key(|&t| OrderedFloat(t)).
• derive SmartDefault with field-level #[default = ...].

🟡 F8. Newton, line search, and tridiagonal solve are hand-rolled [click to expand]

solve_log_tc inlines a damped Newton method with Armijo backtracking (Nocedal and Wright 2006 [4]; Armijo 1966 [5]), and solve_symmetric_tridiagonal applies the Thomas algorithm [wiki], the general tridiagonal solve, to the (symmetric) tridiagonal Hessian without exploiting the symmetry [src]. Maintained equivalents are already linked in the crate: argmin's Newton + BacktrackingLineSearch + ArmijoCondition [doc] and ndarray-linalg's SolveTridiagonal [doc], a LAPACK ?gtsv wrapper. The Armijo constants 1e-4, 0.5, 1e-12 are unnamed literals.

Effect: maintenance surface and no error reporting from the linear solve. The local SPD tridiagonal solve is O(n) and specialized, so replacing it with a library is a real design decision rather than an obvious win.

Suggestions:

• decide library-versus-local explicitly; if kept local, name the Armijo constants and add a zero-pivot guard (ties into F1).

🟡 F9. Numeric core computes on Vec, with ndarray used only to wrap outputs [click to expand]

The gradient, Hessian assembly, objective, and solve run on Vec<f64> with manual index loops; results are lifted into Array1 only for SkylineResult, then build_tc_distribution drops back to Vec via tc_values.to_vec() for a closure that only indexes [src]. The project rule is to keep computation in ndarray. The clear wins are cheap: capture the Array1 in the closure directly, use .dot() for the Armijo slope, and an ndarray reduction for the gradient norm (the manual fold with f64::max also silently absorbs NaN). The full solver migration to Array1/Zip is larger and the scalar Thomas/Newton loops are defensibly scalar.

Suggestions:

• apply the closure-capture, .dot(), and norm fixes now; treat the full ndarray migration as a separate change.

🟡 F10. The added facade test is tautological; no independent value oracle [click to expand]

test_optimize_tc_matches_single_segment_skyline compares optimize_tc(&graph) against optimize_skyline(&graph, n_points = 1) [src], but optimize_tc is exactly that call [src], so the assertion is f(x) == f(x) and cannot fail on a wrong Tc. No test pins Tc = I/M against a hand-derived oracle, exercises the degenerate-tree error, or checks convergence. test_skyline_beats_or_matches_constant_tc also uses a 1e-6 tolerance, looser than the project 1e-7 [src].

Suggestions:

• replace the facade test with an analytical Tc = I/M oracle on the fixed fixture; add a degenerate-tree assert_err! test and a stationarity or scale-invariance property test; tighten the 1e-6 tolerance to the tightest passing 1e-N.

🔵 F11. Per-iteration allocation churn in the solver [click to expand]

Each Newton iteration allocates g, diag, off; each tridiagonal solve allocates c, d, x; each Armijo step allocates a fresh z_new [src]. Bounded at the default n_points = 10 and dominated by the per-round O(E) graph rebuild, so no hot-path regression, but a single reusable workspace removes it. Per-segment info! loops could also drop to debug!.

🔵 F12. Stale docs and retained resolved KB issues; unnamed constants [click to expand]

CLI help still calls the skyline "piecewise linear" and n_skyline "grid points" [src] though it is now piecewise-constant; TimetreeParams::skyline_stiffness and parts of kb/algo/timetree.md still describe the inverse-Tc, Brent, and Nelder-Mead pipeline. Three issue files marked resolved/obsolete (M-timetree-skyline-nelder-mead-optimizer, N-coalescent-initial-tc-hardcoded-fallback, N-coalescent-skyline-simplex-initialization-undecided) are retained with historical bodies rather than deleted per the KB lifecycle. The max(0.5, k - 1) floor in the rate integrand [src] is asserted inert by prose but is an unnamed clamp; worth a named constant plus an assertion that covered intervals have k >= 2.

🔵 F13. KB math notation and citations [click to expand]

The coalescent decision and feature docs use several undeclared symbols (P, N, k, gamma, I, M, z), incomplete where-clauses, and ASCII where KaTeX is used elsewhere. The Kingman likelihood derivations lack a citation to Kingman 1982, and the GMRF skygrid analogy is uncited (Gill et al. 2013; Hill and Baele 2019). Two v0 code-reference links show a line range in the text that the href omits.

Author-tracked

Items the author already documented in KB, with reviewer assessment where the analysis extends or disputes the existing note.

🔴 F3. Merger-quantile boundaries can be non-monotone / merger-free; polytomy weighting mismatch (tracked) [click to expand]

Author documented: kb/issues/N-coalescent-skyline-grid-validation-incomplete.md.

Author's assessment: the note states "the remaining validation gap is narrower -- internal boundaries indexing still lacks a typed nonempty/ordering guard" and proposes parsing the grid into a validated type (O1).

Reviewer assessment (extends): the deficiency is not only a missing typed guard. merger_quantile_boundaries uses candidate.clamp(prev, t_max), which permits candidate == prev [src], so tied merger times or n_points exceeding the number of distinct merger times produce zero-width segments with M_i = 0. The existing small-tree test requests 5 segments from parent times [2000, 2000, 2005, 2005] [src], which yields [2000, 2000, 2000, 2002.5, 2005, 2012] -- two zero-width, merger-free segments -- while the docstring claims every segment owns mergers [src]. The zero-width segments are unreachable via segment_index but their $z_i$ still enter the smoothing penalty, so the fitted trajectory depends on phantom transitions. Separately, boundaries weight a node by its child-edge count while the objective's merger mass per node is $(m-1)/m$ summed [src], so for polytomies the boundaries are not quantiles of the same statistic the objective uses. Suggest building boundaries from distinct merger-event groups weighted consistently with $M_i$, enforcing strict monotonicity, and capping or rejecting unsupported n_points.

Repro (red counterexample, base is this PR head): #872 asserts segment_boundaries is strictly increasing for the over-segmented request; on the current head it fails on the duplicate grid above.

🔴 F4. Midpoint interval attribution vs exact integration across segment boundaries (undecided) [click to expand]

Author documented: kb/issues/N-coalescent-skyline-quadrature-contract-undecided.md.

Author's assessment: the note argues the midpoint rule is now "exact for the piecewise-constant skyline (no quadrature error, and the analytic optimum matches the model-evaluated likelihood)" because the optimizer and the model share the same midpoint convention, and narrows the open quadrature contract to genuinely continuous Tc(t).

Reviewer assessment (disputes the "exact" scope): the shared convention establishes self-consistency (the analytic optimum equals the model-evaluated maximum), which is what the tests confirm, but it does not make $I_i$ the exact integral of $k(k-1)/2$ over segment $i$'s time domain. accumulate_segment_terms assigns each whole lineage-count interval to the segment of its midpoint [src], and merger-quantile boundaries are midpoints between merger times [src], which do not coincide with lineage-count breakpoints. When an interval straddles a segment boundary its entire pairwise-rate mass is attributed to one segment, so $I_i \neq \int_{\text{seg } i} k(k-1)/2,dt$ and the docstring identity [src] holds only for n_points = 1, where boundaries coincide with the tree span. Whether this discretization is acceptable is the open contract in the note; the extension here is that the "exact for piecewise-constant" phrasing overstates it for the multi-segment case, and a test with a boundary strictly inside a constant-k interval would make the difference observable.

🔴 F5. Fail-loud unification aborts the run on node-time inversion; tracking note now stale for the constant path (tracked) [click to expand]

Author documented: kb/issues/M-timetree-skyline-aborts-on-node-time-inversion.md.

Author's assessment: the note says skyline aborts on child_time < parent_time inversions (e.g. flu/h3n2/200) while the constant-Tc path "catches the error and falls back to the initial Tc with a warning", and proposes making skyline degrade like the constant path.

Reviewer assessment (extends / note now inaccurate): this branch removes that asymmetry in the opposite direction. estimate_coalescent_tc now routes both constant and skyline modes through optimize_skyline and propagates every error with ? [src], so the constant-Tc fallback the note relies on as the graceful comparison no longer exists -- an inversion at i >= 2 now aborts the whole run in both modes. collect_coalescent_edges still hard-errors on inversion [src], and the new pipeline comment stating failure occurs only for no-span/no-merger trees does not mention inversion. This is consistent with the deliberate fail-loud decision, but it is a behavior change versus the state the KB note describes, the inversion policy is still undecided, and the note's "constant path degrades gracefully" text is now inaccurate. Worth confirming the intended policy and refreshing the note.

🟡 F6. Retained stiffness = 2.0 default under a reinterpreted, dimensionless penalty (undecided) [click to expand]

Author documented: kb/decisions/coalescent-skyline-convex-log-tc.md.

Author's assessment: the decision records the move to a scale-independent log-Tc penalty and acknowledges reinterpreting the numeric default.

Reviewer assessment (extends): the penalty changed from v0's $\text{stiffness}\sum(\Delta\cdot)^2$ to $\tfrac{\text{stiffness}}{2}\sum(\Delta \ln T_c)^2$ [src], so for identical log-Tc coordinates the retained default 2.0 [src] applies half of v0's coefficient, on top of the piecewise-linear to piecewise-constant change. The decision doc gives no calibration, golden-master, or sensitivity basis for keeping 2.0. Worth an explicit approval of the new regularization contract, or a derivation mapping 2.0 to the intended smoothing strength.

Dismissed

Reported by automated reviewers, investigated and found not to be issues.

⚪ X1. "ndarray-linalg exposes no tridiagonal solver" [click to expand]

Reported: an automated reviewer justified keeping the hand-rolled Thomas solver by claiming the pinned ndarray-linalg 0.17 offers only a dense .solve() with no banded or tridiagonal driver.

Dismissed: the premise is false. SolveTridiagonal<A, Ix1> for Tridiagonal<A> and solve_tridiagonal() exist at the pinned revision (ndarray-linalg/src/tridiagonal.rs) [doc], backed by LAPACK ?gtsv. The reinvention observation itself stands (F8); only this justification for it is wrong.

⚪ X2. "Use approx instead of assert_eq! for the determinism test" [click to expand]

Reported: an automated reviewer flagged the exact assert_eq! on f64 in test_optimize_tc_is_deterministic as violating the float-tolerance rule.

Dismissed: that test asserts bit-identical output from the same deterministic code path across two calls [src]; the invariant is reproducibility, for which exact equality is the correct oracle. An approx comparison would weaken it. (The facade test's real defect is tautology, F10.)

Notes

Click to expand

• N1. The core convex derivation is correct: gradient, tridiagonal Hessian, decoupled warm start, and single-segment closed form $T_c = I/M$ all check out [src].
• N2. "Constant Tc is the one-segment skyline" is an elegant unification that removes a parallel code path [src].
• N3. Propagating the estimator error instead of substituting an invented Tc is the right failure policy once the solver actually reports its failures (F2) [src].
• N4. The commit history pairs each functional change with its KB update, and the KB decision/issue set documents most of the open scientific questions candidly.

Glossary

Click to expand

• Skyline. A time-varying coalescent history, here a piecewise-constant $T_c(t)$ over segments of the tree's time span, generalizing a single constant $T_c$. The frequentist analogue of a Bayesian skygrid population-size history.
• GMRF skygrid. Bayesian nonparametric coalescent model that estimates log effective population size on a fixed time grid under a Gaussian Markov random field smoothing prior (Gill et al. 2013; Hill and Baele 2019).
• Merger. A coalescence event: an internal node where two or more lineages join, counted by $M$.
• Merger-quantile boundaries. Segment boundaries placed at quantiles of the internal-node (merger) times, so each segment spans a comparable share of merger events.
• Warm start. The initial iterate handed to the Newton solver, here the decoupled per-segment optimum $z_i = \ln(I_i / M_i)$.
• Thomas algorithm. Linear-time direct solver for a tridiagonal linear system via forward elimination and back-substitution (Tridiagonal matrix algorithm).
• Armijo backtracking. Line-search rule that shrinks the step until the objective decreases by at least a fixed fraction of the decrease predicted by the gradient (Armijo 1966).

References

Click to expand

1. Kingman, J. F. C. 1982. "The Coalescent." Stochastic Processes and Their Applications 13:235-248. https://doi.org/10.1016/0304-4149(82)90011-4 ↩
2. Gill, Mandev S., Philippe Lemey, Nuno R. Faria, Andrew Rambaut, Beth Shapiro, and Marc A. Suchard. 2013. "Improving Bayesian Population Dynamics Inference: A Coalescent-Based Model for Multiple Loci." Molecular Biology and Evolution 30:713-724. https://doi.org/10.1093/molbev/mss265 ↩
3. Hill, Verity, and Guy Baele. 2019. "Bayesian Estimation of Past Population Dynamics in BEAST 1.10 Using the Skygrid Coalescent Model." Molecular Biology and Evolution 36:2620-2628. https://doi.org/10.1093/molbev/msz172 ↩
4. Nocedal, Jorge, and Stephen J. Wright. 2006. Numerical Optimization. 2nd ed. Springer Series in Operations Research and Financial Engineering. Springer. https://doi.org/10.1007/978-0-387-40065-5 ↩
5. Armijo, Larry. 1966. "Minimization of Functions Having Lipschitz Continuous First Partial Derivatives." Pacific Journal of Mathematics 16:1-3. https://doi.org/10.2140/pjm.1966.16.1 ↩