Skip to content

Fix store restore for mutable tips in beast 2.7 - #8

Merged
kche309 merged 11 commits into
CompEvol:v2.7from
alexeid:fix-store-restore-for-mutable-tips-beast27
May 13, 2026
Merged

kche309 merged 11 commits into
CompEvol:v2.7from
alexeid:fix-store-restore-for-mutable-tips-beast27

Conversation

@kche309

@kche309 kche309 commented May 13, 2026

Copy link
Copy Markdown
Collaborator

Fixes store and restore in BeagleMATreeLikelihood.java to work properly with Gibbs operator in phylonco.

alexeid added 11 commits April 13, 2026 19:36
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.
@kche309
kche309 merged commit dec8279 into CompEvol:v2.7 May 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants