DREAM(KZS): Kalman-inspired proposal (ADR-0067 Stage 3, #358)#489
Merged
Conversation
The Kalman-inspired proposal (DREAM(KZS); Zhang, Vrugt et al. 2020) builds its gain from the archive's parameter<->output cross-covariance, so it needs each archive entry to carry the model *output vector* f(Z) -- not just the PSet and a scalar score. This is the ADR-0067 "implied axis 2b": the output-augmented archive turns on because a proposal declares it needs outputs, not via a user flag. Stage 3a lands that data plumbing, inert by default; Stage 3b adds the kalman proposal + burn-in switch on top. What lands here: - objective: `LikelihoodObjective.aligned_prediction_data(sim, exp, pset)` returns the aligned (prediction f(theta), observation d, variance sigma**2) vectors, walking the same scored points in the same deterministic order as evaluate_pointwise (reusing the row-match / measurement / scale-factor seams). Returns None unless every scored point is an ordinary additive-error (linear scale) Gaussian -- the Kalman R = diag(variance) is the Gaussian measurement covariance, so a log-scale (lognormal) or non-Gaussian family has no R to form. Base ObjectiveFunction returns None (the direct_pass / non-likelihood no-op gate). - dream: a parallel `archive_outputs` list index-aligned with `archive`, seeded with None for the initial random draws (unevaluated), plus `current_output_vec` (the accepted state's f(x) cached per chain, mirroring current_pointwise_loglik) set at all three accept sites (single-try got_result, _mt_bootstrap, _mt_accept) and appended at archive growth. All of it is gated on `_archive_stores_outputs` (False here; Stage 3b sets it True for proposal = kalman), so `de`/`whitened` runs store only PSets exactly as before -- carrying output vectors for a proposal that never reads them would spend memory/bandwidth for nothing. Byte-identical at defaults: the gate is off for de/whitened, so nothing new runs. Verified no regressions against the DREAM/P-DREAM/MT oracle suites (identical pre-existing BNG-env failures on a clean tree vs. this change). New BNG-free extractor unit tests in tests/test_kalman_dream.py cover the Gaussian aligned vectors, the pointwise-order alignment, and the direct_pass / lognormal None gates. Remaining (Stage 3b): the `kalman` proposal + `kalman_burnin_frac` key, the burn-in switch, the closed-form linear-Gaussian posterior-recovery oracle, docs, and the ADR-0067 Stage-3 status flip.
…sions (ADR-0067) Pins the Stage 3 (`kalman`, #358) proposal math against the primary source (Zhang, Vrugt et al. 2020, arXiv:1707.05431) and Vrugt's reference implementation (DREAM-Suite `dream_kzs`), so Stage 3b can be implemented from the ADR alone: the exact Kalman update, the current-state innovation, the paper-vs-code sign caveat, burn-in-only / no-Hastings, the internal M=20 ensemble (no new key), R = diag(σ²) from a Gaussian likelihood, kalman_burnin_frac semantics, the snooker-fixed renormalization, and the (kalman, n_try=1) scope. Also records the two confirmed build decisions (two commits; a test-only linear-Gaussian closed-form oracle) and marks Stage 3a (output-augmented archive plumbing) done on this branch.
wshlavacek
force-pushed
the
feat/358-dream-kzs-kalman
branch
from
July 17, 2026 23:47
d9e4308 to
77bbd95
Compare
Completes DREAM(KZS) (Zhang, Vrugt et al. 2020) on top of Stage 3a's output-augmented archive: the `kalman` proposal steers each burn-in move toward the data with a Kalman gain built from the archive's parameter<->output cross-covariance, then reverts to `de` for a reversible sampling phase. The last piece of the ADR-0067 two-axis unification; #358. The proposal (dream.py): - `_calculate_kalman_pset`: draws an internal M=20 ensemble {Z_K, f(Z_K)} from the archive entries that carry outputs, forms C_ZY / C_YY, and builds the gain K = C_ZY (C_YY + R)^-1 via `_kalman_gain` (a solve with PD jitter, never an explicit inverse). The jump is x_p = x_i + (1+lambda) K (d - f(x_i) + eps) + zeta with the *current* state's residual (paper sign, so the deterministic part reduces ||d - f||), a fresh perturbed-observations draw eps ~ N(0, R), and DREAM's standard e-randomization / small perturbation drawn exactly as the DE proposal draws them. Parameter vectors live in sampling space u (as for snooker / whitened), so the gain maps native output residuals to u-space jumps. - The classic DE body is extracted to `_calculate_de_pset`; `calculate_new_pset` now dispatches whitened -> kalman (`_kalman_active`, the burn-in window) -> de, byte-identical for de/whitened (the DREAM/P-DREAM/MT oracle suites pass unchanged). Kalman falls back to `de` before enough archive outputs accrue, mirroring how whitened falls back before its preconditioner warms up. - No Hastings correction: the Kalman jump breaks detailed balance in the window by design (plain Metropolis), so its samples are burn-in and discarded; after the window the non-snooker branch reverts to `de` and the binary snooker split renormalizes automatically. The window is `kalman_burnin_frac` * burn_in (a new float key, default 0.3; registered in parse.numkeys_float). - R = diag(sigma**2) and d are the current state's cached measurement variance and observation (current_output_obs / current_output_var, filled alongside current_output_vec at every accept). `kalman` therefore requires a linear-scale Gaussian likelihood and n_try == 1, both checked *before the run starts*: the new `ObjectiveFunction.is_linear_gaussian` gate (chi_sq / chi_sq_dynamic True; lognormal / laplace / direct_pass / sos False) and an n_try guard error clearly rather than silently degenerating to `de`. Oracle + unit tests (tests/): - A test-only linear-Gaussian forward model f(x) = A x in integration_harness.py emits observable columns scored by the *real* chi_sq against a generated multi-row .exp, so its posterior x | d ~ N(mu_post, Sigma_post) is closed form (linear_gaussian_posterior) -- the honest end-to-end oracle the scalar direct_pass menu cannot give. test_kalman_dream.py recovers N([2,-1], (1/3) I) (fast: mode; slow: full moments), plus pinned gain-math, deterministic-jump sign / residual reduction, de fallback, and burn-in-switch unit tests, and the is_linear_gaussian / n_try / kalman_burnin_frac config-gate errors. Docs: config_keys.rst (kalman + kalman_burnin_frac), CHANGELOG, ADR-0067 status flipped to Stages 1-3 implemented (Stage 3b done). Effective-config golden regenerated (a clean single-key insertion of kalman_burnin_frac: 0.3 into every dream/p_dream config).
…DR-0067 Stage 3b)
Two test follow-ups to the Kalman proposal commit:
- test_benchmark_harness: add kalman_burnin_frac to the pre-migration oracle's
_EXCLUDE set. The oracle is a frozen, independent witness that the ADR-0012
config migration reproduces the old per-method .conf files byte-for-byte; a key
added after the freeze is excluded rather than regenerated (exactly as proposal
and n_try were for Stages 1-2), so it stays an independent witness for the keys
the original confs actually carried. kalman_burnin_frac (default 0.3) is
meaningful only for the kalman proposal these de/whitened confs never selected,
so it is a no-op here.
- test_kalman_dream: fill the two gaps the unified two-axis engine opened but
nothing asserted as a whole:
* kalman x heavy snooker -- parametrize the mode-recovery oracle over
snooker_prob (0.1, 0.5). The Kalman jump is the non-snooker branch of DREAM's
binary split, so the snooker-heavy case is the composition check that
proposal = kalman co-exists with the snooker mix-in and still recovers the
mode (the analogue of multi-try's snooker-heavy-k3).
* a consolidated config-rejection matrix -- one parametrized test over every
unsupported config point (invalid proposal enum, kalman + non-Gaussian
objective, kalman + n_try > 1, kalman_burnin_frac out of [0,1] high and low),
each erroring at construction. Replaces the three separate kalman-rejection
tests and adds the invalid-proposal-string and negative-frac cases.
Full gate green (3067 passed, 9 skipped, 150 deselected slow/recovery).
wshlavacek
marked this pull request as ready for review
July 18, 2026 17:23
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.
Implements ADR-0067 Stage 3 — the
kalmanproposal (DREAM(KZS); Zhang, Vrugt et al. 2020), the last piece of unifying the DREAM family into one engine with two orthogonal axes. Closes #358 when complete.Draft: landing in two commits. 3a (the output-augmented archive plumbing) is here; 3b (the Kalman proposal itself) is still TODO.
Done — Commit 3a: output-augmented archive plumbing (the "implied axis 2b")
The Kalman gain is built from the archive's parameter↔output cross-covariance, so each archive entry must carry the model output vector
f(Z), not just the PSet + scalar score. This lands that data plumbing, inert by default.LikelihoodObjective.aligned_prediction_data(sim, exp, pset)returns the aligned(prediction f(θ), observation d, variance σ²)vectors, walking the same scored points in the same order asevaluate_pointwise. ReturnsNoneunless every point is an ordinary additive-error (linear-scale) Gaussian — the KalmanR = diag(variance)is the Gaussian measurement covariance, so a log-scale/non-Gaussian family has noR. BaseObjectiveFunctionreturnsNone(thedirect_pass/non-likelihood gate).archive_outputslist index-aligned witharchive(initial random draws seededNone), the accepted state'sf(x)cached per chain (current_output_vec, mirroringcurrent_pointwise_loglik) at all three accept sites, appended at archive growth. All gated on_archive_stores_outputs(Falseforde/whitened,Trueforkalmanin 3b), so a plain run stores only PSets exactly as before.tests/test_kalman_dream.py(Gaussian aligned vectors, pointwise-order alignment,direct_pass/lognormal→None).Byte-identical at defaults (gate off); no regressions vs. a clean tree.
TODO — Commit 3b: the
kalmanproposal + burn-in switch + oracleThe confirmed DREAM(KZS) algorithm and all build decisions are recorded in ADR-0067 → "Stage 3 — confirmed algorithm and build decisions" (pinned against arXiv:1707.05431 + Vrugt's DREAM-Suite
dream_kzs):DreamConfig.kalman_burnin_frac(default 0.3, a fraction ofburn_in) +proposal = 'kalman'validation._calculate_kalman_pset:K = C_ZY (C_YY + R)^-1(solve, jitter for PD), current-state innovationd − f(x_i) + ε, paper sign,defallback on a small archive, internalM = 20ensemble (no new user key).n_try == 1.f(x) = A xintests/integration_harness.pyas the closed-form posterior oracle, plus pinned-RNG gain-math unit tests.config_keys.rst), CHANGELOG, and the ADR Stage-3 status flip.