Repository navigation
Fix store restore for mutable tips in beast 2.7 - #8
Merged
kche309 merged 11 commits intoMay 13, 2026
Merged
Conversation
getLogProbsForStateSequence / getLogProbsForPartialsSequence mutate the likelihoodCore's tip states and internal partials in place without respecting BEAST's store/restore contract: - calcPatternLogLikelihoods writes partials to whichever slot is "current" without first flipping to the scratch slot, so the stored partials (captured by store() before the proposal runs) are clobbered and restore() cannot recover them. - setNodeStates / setTipStates overwrite tip states in place, and there is no mechanism to restore the original states on a rejected proposal if the mutated tip wasn't in alignment.getDirtySequenceIndices(). This corrupts the "old" state after the first rejected proposal and is the likely cause of 100% acceptance rates in operators that use these probe methods (e.g. phylonco's ExchangeGibbsOperator). Fix: - Flip the partials buffer to the scratch slot before each probe write, tracking flipped nodes so each is flipped at most once per proposal (toggling twice would rotate back to the stored slot and clobber it). - Track tip nodes transiently mutated via getLogProbs*Sequence, and resync them from the (post-restore) alignment on restore() -- a superset of alignment.getDirtySequenceIndices() since probes may touch tips that weren't alignment-dirty. - Clear both tracking sets in store(), accept(), and restore().
Previous fix prevented probe-to-probe poisoning of the stored partials slot, but not probe-to-framework poisoning. With hasDirt set to Tree.IS_DIRTY by requiresRecalculation(), TreeLikelihood.traverse() calls setNodePartialsForUpdate() for every internal ancestor. For ancestors previously flipped by getLogProbs*Sequence, this is a second toggle in the same proposal: current rotates back to the stored slot, calculatePartials writes on top of it, and restore() (an index swap) cannot recover. Fix: in calculateLogP(), undo the flippedNodes toggles before super.calculateLogP() runs, so the framework's own traverse-time toggles land on the scratch slot instead of the stored slot. Same change for BeagleMATreeLikelihood using partialBufferHelper.flipOffset(). This is the remaining cause of 100% acceptance in operators that use the probe methods (e.g. phylonco's ExchangeGibbsOperator. EOF )
The flip-undo machinery added in 0a33bb7/c99c8f8 didn't actually run: MCMC calls state.storeCalculationNodes() (and hence MATreeLikelihood.store) between operator.proposal() and posterior.calculateLogP() whenever the operator requires state initialisation (the default). Probes happen during proposal, so clearing tempTipNodes / flippedNodes in store() wiped the tracking before calculateLogP() could consume it. Result: probe-time buffer flips were never undone, traverse's flip rotated current back to the stored slot, and restore() rolled the polluted slot into "current", producing the ~30 logP drift observed in phylonco's ExchangeGibbsOperator (detected as a likelihood mismatch at sample 0). Fix: leave the sets populated through the proposal cycle. They are still cleared by calculateLogP() (after the undo), accept(), and restore(). Verified: 50K-step short reproducer in beast-phylonco previously failed with cached -541.026 vs fresh -570.892 (delta 29.87) at sample 0; now passes. 3M-step long test passes cleanly in ~17 minutes.
Sequence diagrams of the buggy flow (cleared flippedNodes in store) and the fixed flow, plus rationale for why a simpler localised design isn't adopted (single-buffered tip states + same constraint on BEAGLE path). Notes the remaining edge case for operators that probe then return Double.NEGATIVE_INFINITY.
The term was used throughout without definition. Added a short opening paragraph explaining it means a speculative getLogProbsForStateSequence call that side-effects the buffer.
GitHub's mermaid renderer rejected the first sequence diagram, parse error around the trailing restore line. Removed the trailing-parens "(no-op)" suffix and the c↔s arrow character; substituted plain "and" and dropped the inline emoji from the BUG markers. Diagram content is the same.
Replaces the class-level flippedNodes Set + cross-method undo machinery with a localised "flip up, compute, flip back" pattern inside calcPatternLogLikelihoods. After the probe call returns, the partials indices are exactly what they were on entry, so the framework's store/restore can never mishandle them. Also removes the don't-clear-in-store trap: since there is no longer any cross-method probe state for partials, store() can clear normally. Only tempTipNodes (for single-buffered tip states) still has to survive across the proposal-then-store-then-calculateLogP boundary. BeagleMATreeLikelihood threads a per-call `flipped` set through the rescale recursion so flips are still deduped across rescale attempts and flipped back exactly once on success. Verified: ExchangeGibbsOperatorShortTest still passes.
Tree leaf order and alignment column order are independent: they are linked only by taxon name, not by integer index. The previous code in updateTipData() and the getLogProbs*Sequence probe APIs conflated the two, treating tree node number and alignment column index as interchangeable. This is correct for synthetic tests where both happen to be in the same order (e.g. taxa "0".."4"), but corrupts state when the orderings differ (typical for real data, where the alignment is written in input/data order while the TaxonSet is alphabetised). When a probe-based operator (ExchangeGibbsOperator) committed Gibbs- sampled sequence to taxon-by-tree-node-number, the alignment wrote it to the wrong column. updateTipData() then re-synced tips by taxon name, which read the (now stale) right column and ignored the new data sitting in the wrong column. Cumulative state corruption across pre-burnin tripped the MCMC integrity check. Build the alignmentIdx <-> treeNodeNr maps once in initAndValidate (by name lookup) and use them in updateTipData and the probe APIs. dirtySequences and tempTipNodes both now hold alignment column indices consistently.
The previous restore() wiped dirtySequences and re-derived it from alignment.getDirtySequenceIndices(), but by the time MATreeLikelihood. restore() runs the framework has already called alignment.restore(), which clears the alignment's edit list. So the re-derivation returned nothing for non-probe edits and tip data was never re-synced from the rolled-back alignment. The pre-hermetic-refactor code reused the dirtySequences captured during the previous calculateLogP (when the edit list was still populated). Restore that pattern. Probe-touched tips (tempTipNodes) are still added in case they were not alignment-dirty. Was masked by ExchangeGibbsOperator passing all edited tips through tempTipNodes too, so MCMC scenarios with only that operator looked fine. Surfaces the moment any non-probe alignment-edit operator is rejected, or any direct setSiteValue test path. Verified: ExchangeGibbsOperatorShortTest still passes; CRC09 50k chain runs to completion without integrity-check failures.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes store and restore in BeagleMATreeLikelihood.java to work properly with Gibbs operator in phylonco.