Add convergence helper and error tracking to nested sampling #41
Conversation
- Add logZ_error, H (information), i_eff, and n_eff fields to NSState - Update runtime info calculation to compute these new metrics - Add converged() utility function for checking termination criteria - Update example to use new convergence function 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
yallup
left a comment
yallup
left a comment
There was a problem hiding this comment.
Very useful addition, have some queries about bloat of info, but if it is needed to be in the state then I’m happy with taking this on
| logX: Array # The current log-volume estimate | ||
| logZ: Array # The accumulated evidence estimate | ||
| logZ_live: Array # The current evidence estimate | ||
| logZ_error: Array # The current error estimate on logZ |
There was a problem hiding this comment.
Can any of this be calculated “offline”, i.e. in the utils rather than in the online function?
There was a problem hiding this comment.
It could all be calculated offline, but there is a virtue in accumulation here that we appreciated in the first place for logZ (this is not true for logZ_live, which we do not accumulate, and could be calculated as part of converged). My main reason for doing this is to be able to construct more advanced convergence criteria that e.g. incorporated an evidence error, or a KL-based stopping criterion (which John Skilling prefers).
| delta_logZ = logsumexp(log_delta_Z) | ||
| logZ = jnp.logaddexp(logZ, delta_logZ) | ||
| logZ = jnp.logaddexp(state.logZ, delta_logZ) | ||
| A = state.i_eff / state.n_eff + jnp.sum(1 / num_live) |
There was a problem hiding this comment.
If this has to be all calculated online, would it be useful to put any of the calculation logic into the utils, in case they are useful for various post hoc analysis bits
Summary
This PR introduces two key improvements to the nested sampling implementation:
Convergence helper function - Adds a
converged()utility function (discussed with @zwei-beiner) that properly handles the "full plateau" case where all live points have identical likelihoods. This addresses an edge case where our zero-weight check in the delete function doesn't prevent the algorithm from continuing - it just spreads the selection more fairly due tojax.random.choice's surprising behavior with zero/NaN weights.Error tracking and information metrics - Extends
NSStateto track:logZ_error: Error estimate on the log evidenceH: Information (negative entropy) in natsi_eff: Effective number of iterationsn_eff: Effective sample sizeChanges
blackjax/ns/base.pyNSState:logZ_error,H,i_eff,n_effinit()to initialize these fieldsupdate_ns_runtime_info()to compute error estimates and information metricsblackjax/ns/utils.pyconverged()function that checks:docs/examples/nested_sampling.pyconverged()utility functionTesting
The changes maintain backward compatibility and all existing tests pass. The convergence helper provides a cleaner termination condition that handles edge cases more robustly.
🤖 Generated with Claude Code
Co-Authored-By: Claude [email protected]