Add EPAnalysisFactor for cavity-message injection#1248
Merged
Conversation
EPAnalysisFactor is a thin AnalysisFactor subclass that exposes the EP
cavity distribution to its Analysis on each fit, via a guarded
``set_cavity_dist`` hook called from ``factor_step``. This lets a
hierarchical / population-level Analysis read per-variable cavity
messages directly inside ``log_likelihood_function`` — useful for
"global" likelihoods that compare predictions to the per-dataset
posterior summaries produced by upstream local fits, e.g. the
EP-message-comparison form documented in the cancer-IC50 use case:
log L = -0.5 * sum_i || (pred_i - cavity_mean_i) / cavity_sigma_i ||^2
The hook is duck-typed so plain AnalysisFactor behaviour is unchanged.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
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
1 participant
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.
Summary
Adds
EPAnalysisFactor, an additiveAnalysisFactorsubclass that exposes the EP cavity distribution to itsAnalysison each fit. This unblocks hierarchical / population-level likelihoods that need to read per-variable cavity messages directly insidelog_likelihood_function— for example, a "global" Analysis that compares predictions to per-dataset posterior summaries produced by upstream local fits, of the formlog L = -0.5 * sum_i || (pred_i - cavity_mean_i) / cavity_sigma_i ||^2.The change is purely additive: the cavity hook in
factor_stepis duck-typed viahasattr(factor, "set_cavity_dist"), so plainAnalysisFactorand every other existing factor type are unaffected.API Changes
Added
af.EPAnalysisFactor— anAnalysisFactorsubclass with aset_cavity_dist(cavity_dist)hook that the EP optimiser invokes before every fit, attaching the cavityMeanFieldtoanalysis._cavity_mean_fieldso the user'slog_likelihood_functioncan read per-variable cavity messages (.mean,.sigma) at evaluation time. No symbols removed, renamed, or behaviour-changed for existing factors. See full details below.Test Plan
pytest test_autofit/graphical/test_ep_analysis_factor.py— 4 new unit tests passpytest test_autofit/graphical -q— full graphical suite passes (184 tests)EPAnalysisFactor.set_cavity_distfires before every global step, populatesanalysis._cavity_mean_fieldwith non-trivial messages from the local fits, and the per-variable means/sigmas tighten across EP iterations as expected.Full API Changes (for automation & release notes)
Added
autofit.EPAnalysisFactor(prior_model, analysis, optimiser=None, name=None)—AnalysisFactorsubclass exposing the EP cavity distribution to itsAnalysisviaset_cavity_dist(cavity_dist). Stores the cavityMeanFieldonanalysis._cavity_mean_fieldsolog_likelihood_functioncan read per-variable cavity messages.autofit.graphical.EPAnalysisFactor— same class re-exported under the graphical namespace.Changed Behaviour
autofit.graphical.expectation_propagation.optimiser.factor_step— now callsfactor.set_cavity_dist(factor_approx.cavity_dist)if the factor implements it. Duck-typed: no-op for plainAnalysisFactorand every other factor type that does not opt in.Migration
None required — purely additive. Existing
AnalysisFactorusers see no change.🤖 Generated with Claude Code