From 9724842a419656891ce1b06219869252b20b6836 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 10 Jul 2026 10:08:02 +0100 Subject: [PATCH] =?UTF-8?q?test(graphical):=20EP=20integration=20coverage?= =?UTF-8?q?=20=E2=80=94=20parity,=20deterministic,=20exact=20(#81)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit EP framework review Phase 3 (test leg): ep_parity.py (EP vs joint Dynesty parity + truth recovery), ep_deterministic.py (low-level factor_out API through an EP fit), ep_exact.py (ExactFactorFit conjugate path vs analytic product). All verified PASS, re-validated against post-#1345/#1348 main. Development surfaced audit finding F10 (undamped sigma collapse) on PyAutoFit#1332. Co-Authored-By: Claude Opus 4.8 --- scripts/graphical/ep_deterministic.py | 164 +++++++++++++++++ scripts/graphical/ep_exact.py | 120 +++++++++++++ scripts/graphical/ep_parity.py | 243 ++++++++++++++++++++++++++ 3 files changed, 527 insertions(+) create mode 100644 scripts/graphical/ep_deterministic.py create mode 100644 scripts/graphical/ep_exact.py create mode 100644 scripts/graphical/ep_parity.py diff --git a/scripts/graphical/ep_deterministic.py b/scripts/graphical/ep_deterministic.py new file mode 100644 index 0000000..9ed0bc4 --- /dev/null +++ b/scripts/graphical/ep_deterministic.py @@ -0,0 +1,164 @@ +""" +Integration Test: EP Deterministic Variable (Low-Level `factor_out` API) +========================================================================= + +This script exercises the LOW-LEVEL graph API used to build a hand-crafted `FactorGraph` with a +deterministic variable, and runs expectation propagation (EP) on it directly (i.e. without the +declarative `AnalysisFactor` / `FactorGraphModel` API used in `ep.py` and `simultaneous.py`). + +A deterministic variable is a node on the factor graph whose value is a (deterministic) function of +other variables, rather than a free parameter with its own prior. It is declared via the `factor_out` +keyword argument of `graph.Factor`, following the canonical pattern used throughout +`PyAutoFit/test_autofit/graphical/` (e.g. `test_factor_graph.py::make_plus`, +`regression/test_exact.py::make_model`, `regression/conftest.py::make_linear_factor`). + +__Graph__ + +We build a small graph with two free scalar variables `x` and `y`, each with an independent Gaussian +prior, and a deterministic variable `z = x + y` declared via `factor_out`. A Gaussian likelihood +message-factor is placed directly on `z`, informing the graph of an observed value of `x + y`. + + prior_x: NormalMessage(mean=1.0, sigma=2.0).as_factor(x) + prior_y: NormalMessage(mean=1.0, sigma=2.0).as_factor(y) + linear_factor: Factor(lambda x, y: x + y, x, y, factor_out=z) + likelihood: NormalMessage(mean=3.0, sigma=0.1).as_factor(z) + + model = likelihood * linear_factor * prior_x * prior_y + +This mirrors the pattern used in `test_autofit/graphical/regression/test_linear_regression.py` +(`test_exact_updates`), which builds `likelihood_factor * linear_factor * prior_a * prior_b` and runs +`EPOptimiser.from_meanfield(model_approx, default_optimiser=laplace)`. The deterministic `linear_factor` +here is a simple Python function with no closed-form (exact) projection, so it is fit via +`af.LaplaceOptimiser()`; the Gaussian prior/likelihood factors are exact and are auto-assigned an +`ExactFactorFit` optimiser by `EPOptimiser.from_meanfield`. +""" + +import numpy as np + +import autofit as af +from autofit import graphical as graph +from autofit.messages.normal import NormalMessage + +""" +__Variables__ + +`x` and `y` are free scalar variables. `z` is the deterministic output variable representing `x + y`. +""" +x_, y_, z_ = graph.variables("x, y, z") + +""" +__Priors__ + +Independent Gaussian priors on `x` and `y`, each built via `NormalMessage(...).as_factor(variable)`, +following the pattern in `regression/test_linear_regression.py::make_normal_model_approx`. +""" +prior_mean = 1.0 +prior_sigma = 2.0 + +prior_x = NormalMessage(prior_mean, prior_sigma).as_factor(x_) +prior_y = NormalMessage(prior_mean, prior_sigma).as_factor(y_) + +""" +__Deterministic Factor__ + +`z = x + y`, declared via `factor_out=z_` following `test_factor_graph.py::make_plus` +(`graph.Factor(plus_two, x, factor_out=y)`) and `regression/conftest.py::make_linear_factor` +(`graph.Factor(linear, x_, a_, b_, factor_out=z_)`). +""" + + +def sum_of_x_and_y(x, y): + return x + y + + +linear_factor = graph.Factor(sum_of_x_and_y, x_, y_, factor_out=z_) + +""" +__Likelihood__ + +A Gaussian message-factor on `z`, centred on an observed value `z_obs = 3.0` with a tight `sigma = 0.1`, +informing the graph that `x + y` is observed to be close to `3.0`. +""" +z_obs = 3.0 +z_obs_sigma = 0.1 + +likelihood_factor = NormalMessage(z_obs, z_obs_sigma).as_factor(z_) + +""" +__Factor Graph__ + +Combine every factor via `*`, following every example in `test_autofit/graphical/` (e.g. +`likelihood_factor * linear_factor * prior_a * prior_b`). +""" +model = likelihood_factor * linear_factor * prior_x * prior_y + +""" +__Initial Mean Field__ + +An initial approximating distribution is required for every variable in the graph, including the +deterministic variable `z` (see `regression/test_exact.py::make_model`, which supplies an initial dist +for the deterministic variable `f_`). The initial `z` dist is the (rough) implied distribution of +`x + y` under the priors on `x` and `y` (mean = mean_x + mean_y, variance = var_x + var_y). +""" +approx0 = { + x_: NormalMessage(prior_mean, prior_sigma), + y_: NormalMessage(prior_mean, prior_sigma), + z_: NormalMessage(2.0 * prior_mean, (2.0 * prior_sigma**2) ** 0.5), +} + +z_prior_variance = float(approx0[z_].variance) + +model_approx = graph.EPMeanField.from_approx_dists(model, approx0) + +""" +__Expectation Propagation__ + +`EPOptimiser.from_meanfield` auto-assigns `ExactFactorFit` to the exact Gaussian prior/likelihood +factors and the supplied `default_optimiser` (`af.LaplaceOptimiser()`) to the deterministic +`linear_factor`, following `regression/test_linear_regression.py::test_exact_updates`. +""" +laplace = af.LaplaceOptimiser() + +ep_opt = graph.EPOptimiser.from_meanfield( + model_approx, default_optimiser=laplace, paths=False +) + +fit_approx = ep_opt.run(model_approx) + +mean_field = fit_approx.mean_field + +print(mean_field) +print() +print(f"x: mean = {mean_field[x_].mean}, variance = {mean_field[x_].variance}") +print(f"y: mean = {mean_field[y_].mean}, variance = {mean_field[y_].variance}") +print(f"z: mean = {mean_field[z_].mean}, variance = {mean_field[z_].variance}") + +x_mean = float(mean_field[x_].mean) +y_mean = float(mean_field[y_].mean) +x_variance = float(mean_field[x_].variance) +y_variance = float(mean_field[y_].variance) +z_variance = float(mean_field[z_].variance) + +sum_mean = x_mean + y_mean +combined_sigma = (x_variance + y_variance) ** 0.5 + +""" +__Assertions__ + +1) The posterior mean of `x + y` should be close to the observed value `z_obs = 3.0`, within 3 combined + (x, y) sigma. +2) The posterior uncertainty on `z` should have shrunk relative to its (rough) prior-implied variance, + since the tight likelihood on `z` (sigma = 0.1) is far more informative than the priors on `x`/`y`. +""" +assert np.isfinite(sum_mean) +assert np.isfinite(combined_sigma) and combined_sigma > 0.0 +assert abs(sum_mean - z_obs) < 3.0 * combined_sigma, ( + f"posterior mean of x + y ({sum_mean}) not within 3 sigma ({3.0 * combined_sigma}) " + f"of observed z_obs ({z_obs})" +) +assert z_variance < z_prior_variance, ( + f"posterior variance on z ({z_variance}) did not shrink relative to its prior-implied " + f"variance ({z_prior_variance})" +) + +print("ep_deterministic.py: PASS") diff --git a/scripts/graphical/ep_exact.py b/scripts/graphical/ep_exact.py new file mode 100644 index 0000000..8d937b5 --- /dev/null +++ b/scripts/graphical/ep_exact.py @@ -0,0 +1,120 @@ +""" +Integration Test: EP Exact Conjugate Updates (`ExactFactorFit`) +================================================================ + +This script exercises the exact conjugate-update path of the expectation propagation (EP) framework. + +When a factor on a hand-built `FactorGraph` is itself a message distribution (e.g. a `NormalMessage` +turned into a factor via `.as_factor(variable)`), the factor supports an EXACT projection: the EP +update is computed in closed form via conjugate Gaussian multiplication rather than by running a +numerical optimiser. `EPOptimiser.from_meanfield(...)` detects this via `factor.has_exact_projection` +and auto-assigns the `ExactFactorFit` optimiser to such factors (see +`PyAutoFit/autofit/graphical/expectation_propagation/optimiser.py::EPOptimiser.from_meanfield` and the +canonical usage in `PyAutoFit/test_autofit/graphical/regression/test_linear_regression.py::test_exact_updates` +and `regression/test_exact.py::test_probit_regression`). + +__Graph__ + +The simplest possible conjugate problem: a single scalar variable `x` with + + prior: NormalMessage(m1, s1).as_factor(x) + likelihood: NormalMessage(m2, s2).as_factor(x) + + model = likelihood * prior + +The exact posterior is the product of two Gaussians: + + posterior_precision = 1/s1^2 + 1/s2^2 + posterior_mean = (m1/s1^2 + m2/s2^2) / posterior_precision + +Because every factor on this graph is exact, `EPOptimiser.from_meanfield` runs the whole EP fit via +`ExactFactorFit` optimisers (no `LaplaceOptimiser` is involved) and the EP mean field should match the +analytic posterior to machine-level precision. +""" + +import numpy as np + +from autofit import graphical as graph +from autofit.messages.normal import NormalMessage + +""" +__Variable__ +""" +x_ = graph.Variable("x") + +""" +__Factors__ + +Both the prior and the likelihood are `NormalMessage`s converted to factors via `.as_factor`, following +`regression/test_linear_regression.py::make_normal_model_approx`. +""" +m1, s1 = 1.0, 2.0 # prior +m2, s2 = 3.0, 0.5 # likelihood (observation) + +prior_factor = NormalMessage(m1, s1).as_factor(x_, name="prior_x") +likelihood_factor = NormalMessage(m2, s2).as_factor(x_, name="likelihood_x") + +model = likelihood_factor * prior_factor + +""" +__Initial Mean Field__ + +A deliberately broad, offset starting approximation, so the assertions below genuinely test the exact +EP updates rather than the initialisation. +""" +approx0 = { + x_: NormalMessage(0.0, 10.0), +} + +model_approx = graph.EPMeanField.from_approx_dists(model, approx0) + +""" +Both factors must support the exact projection for `ExactFactorFit` to be auto-selected. +""" +factor_mean_field = model_approx.factor_mean_field + +assert prior_factor.has_exact_projection(factor_mean_field[prior_factor]) +assert likelihood_factor.has_exact_projection(factor_mean_field[likelihood_factor]) + +""" +__Expectation Propagation__ + +No `default_optimiser` is passed: every factor is exact, so `EPOptimiser.from_meanfield` assigns +`ExactFactorFit` to each of them. +""" +ep_opt = graph.EPOptimiser.from_meanfield(model_approx, paths=False) + +fit_approx = ep_opt.run(model_approx) + +mean_field = fit_approx.mean_field + +""" +__Analytic Conjugate Posterior__ + +The product of two Gaussians N(m1, s1) x N(m2, s2). +""" +posterior_precision = 1.0 / s1**2 + 1.0 / s2**2 +posterior_mean = (m1 / s1**2 + m2 / s2**2) / posterior_precision +posterior_sigma = posterior_precision**-0.5 + +ep_mean = float(mean_field[x_].mean) +ep_sigma = float(mean_field[x_].sigma) + +print(mean_field) +print() +print(f"EP posterior: mean = {ep_mean}, sigma = {ep_sigma}") +print(f"Analytic posterior: mean = {posterior_mean}, sigma = {posterior_sigma}") + +""" +__Assertions__ + +The EP mean field must match the analytic conjugate posterior to rtol 1e-6. +""" +assert np.isclose(ep_mean, posterior_mean, rtol=1e-6), ( + f"EP mean ({ep_mean}) does not match analytic posterior mean ({posterior_mean}) to rtol 1e-6" +) +assert np.isclose(ep_sigma, posterior_sigma, rtol=1e-6), ( + f"EP sigma ({ep_sigma}) does not match analytic posterior sigma ({posterior_sigma}) to rtol 1e-6" +) + +print("ep_exact.py: PASS") diff --git a/scripts/graphical/ep_parity.py b/scripts/graphical/ep_parity.py new file mode 100644 index 0000000..47ba362 --- /dev/null +++ b/scripts/graphical/ep_parity.py @@ -0,0 +1,243 @@ +""" +Integration Test: EP vs Joint-Fit Parity (Shared Centre) +========================================================= + +This script fits the same factor graph -- 3 noisy 1D Gaussian datasets whose `Gaussian` models share a +single global `centre` -- in two independent ways and checks the two posteriors on the shared centre +agree: + + (a) a JOINT fit of the `global_prior_model` of the factor graph, sampling all parameters + simultaneously with a small `DynestyStatic` search (the pattern of `graphical/simultaneous.py`); + + (b) an EXPECTATION PROPAGATION (EP) fit of the factor graph via `factor_graph.optimise` with + `af.LaplaceOptimiser()` (the pattern of `graphical/ep.py`). + +Both fits use identical models and priors, so their posteriors on the shared `centre` should agree +within their combined uncertainties, and both should recover the true simulated value of 50.0. + +__Dataset__ + +The first two datasets are the `gaussian_x1__low_snr` datasets on disk (true centre 50.0). Only two of +these exist in the workspace, so the third dataset is simulated in-memory with a fixed seed, using the +same generator settings as `scripts/simulators/util.py` (centre=50.0, normalization=25.0, sigma=10.0, +signal-to-noise ratio 25). It is not written to disk so the tracked `dataset` folder is unchanged. +""" + +import os + +cwd = os.getcwd() + +import numpy as np +from os import path + +import autofit as af + +""" +__Dataset__ +""" +total_datasets_on_disk = 2 + +data_list = [] +noise_map_list = [] + +""" +__Dataset Auto-Simulation__ + +If the on-disk datasets do not already exist on your system, they are created by running the +corresponding simulator script. +""" +if not path.exists( + path.join("dataset", "example_1d", "gaussian_x1__low_snr", "dataset_0") +): + import subprocess + import sys + + subprocess.run( + [sys.executable, "scripts/simulators/simulators.py"], + check=True, + ) + +for dataset_index in range(total_datasets_on_disk): + dataset_name = f"dataset_{dataset_index}" + + dataset_path = path.join( + "dataset", "example_1d", "gaussian_x1__low_snr", dataset_name + ) + + data = af.util.numpy_array_from_json(file_path=path.join(dataset_path, "data.json")) + noise_map = af.util.numpy_array_from_json( + file_path=path.join(dataset_path, "noise_map.json") + ) + + data_list.append(data) + noise_map_list.append(noise_map) + +""" +__Third Dataset (In-Memory)__ + +Simulated with the same settings as `scripts/simulators/util.py` but with a fixed seed, and not written +to disk. +""" +rng = np.random.default_rng(seed=1) + +pixels = 100 +signal_to_noise_ratio = 25.0 +xvalues = np.arange(pixels) + +gaussian_true = af.ex.Gaussian(centre=50.0, normalization=25.0, sigma=10.0) +model_data_1d = gaussian_true.model_data_from(xvalues=xvalues) + +data_2 = model_data_1d + rng.normal(0.0, 1.0 / signal_to_noise_ratio, pixels) +noise_map_2 = (1.0 / signal_to_noise_ratio) * np.ones(pixels) + +data_list.append(data_2) +noise_map_list.append(noise_map_2) + +total_datasets = len(data_list) + +""" +__Analysis__ +""" +analysis_list = [] + +for data, noise_map in zip(data_list, noise_map_list): + analysis = af.ex.Analysis(data=data, noise_map=noise_map) + + analysis_list.append(analysis) + +""" +__Model__ + +All Gaussians share the same `centre`, set up as a single shared `GaussianPrior` assigned to every +model's `centre` (the pattern of `graphical/ep.py`). + +The `normalization` and `sigma` priors are plain `GaussianPrior`s, following the canonical pure-Laplace +declarative EP pattern in `PyAutoFit/test_autofit/graphical/gaussian/test_declarative.py::test_gaussian`. +(The `TruncatedGaussianPrior`s used by `graphical/ep.py` are numerically unstable when the analysis +factors themselves are optimised with `LaplaceOptimiser`; `ep.py` avoids this by assigning each +`AnalysisFactor` its own `DynestyStatic` optimiser.) +""" +centre_shared_prior = af.GaussianPrior(mean=50.0, sigma=30.0) + +model_list = [] + +for model_index in range(total_datasets): + gaussian = af.Model(af.ex.Gaussian) + + gaussian.centre = centre_shared_prior # This prior is used by all 3 Gaussians! + + gaussian.normalization = af.GaussianPrior(mean=25.0, sigma=10.0) + gaussian.sigma = af.GaussianPrior(mean=10.0, sigma=10.0) + + model_list.append(gaussian) + +""" +__Analysis Factors + Factor Graph__ +""" +analysis_factor_list = [] + +for dataset_index, (model, analysis) in enumerate(zip(model_list, analysis_list)): + analysis_factor = af.AnalysisFactor( + prior_model=model, analysis=analysis, name=f"dataset_{dataset_index}" + ) + + analysis_factor_list.append(analysis_factor) + +factor_graph = af.FactorGraphModel(*analysis_factor_list) + +""" +__(a) Joint Fit__ + +Fit the `global_prior_model` of the factor graph with a small `DynestyStatic` budget, following +`graphical/simultaneous.py`. Parity tolerances below are loose, so the budget is kept small for speed. + +Any previous output of this search is removed first: resuming a completed run loads a samples summary +whose `errors_at_sigma` are zero, which would invalidate the sigma-based parity assertions below. +""" +import shutil + +shutil.rmtree(path.join("output", "graphical", "ep_parity_joint"), ignore_errors=True) + +search = af.DynestyStatic( + path_prefix=path.join("graphical"), + name="ep_parity_joint", + nlive=50, + sample="rwalk", + walks=10, + maxcall=3000, + maxiter=3000, +) + +joint_result = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph) + +joint_instance = joint_result.samples.median_pdf() +joint_median = joint_instance[0].centre + +joint_error_instance = joint_result.samples.errors_at_sigma(sigma=1.0) +joint_sigma = float(np.mean(joint_error_instance[0].centre)) + +""" +__(b) EP Fit__ + +Fit the same factor graph with the EP framework via `factor_graph.optimise`, following +`graphical/ep.py`. + +`max_steps` (forwarded by `optimise(**kwargs)` to `EPOptimiser.run`) is capped at 3. With the default +undamped updater (`SimplerUpdater(delta=1.0)`), pure-Laplace EP on this graph does not trigger the +`kl_tol` termination and repeated cycles progressively over-count shared-variable information, +collapsing every posterior sigma towards zero (~1e-15 after the default 100 steps). The mean field is +stable and sensible after 2-3 sweeps. +""" +laplace = af.LaplaceOptimiser() + +paths = af.DirectoryPaths(name=path.join("graphical", "ep_parity")) + +factor_graph_result = factor_graph.optimise( + optimiser=laplace, + paths=paths, + ep_history=af.EPHistory(kl_tol=0.05), + max_steps=3, +) + +mean_field = factor_graph_result.updated_ep_mean_field.mean_field + +ep_mean = float(mean_field.mean[centre_shared_prior]) +ep_sigma = float(mean_field.scale[centre_shared_prior]) + +""" +__Posteriors__ +""" +true_centre = 50.0 + +print() +print(f"True centre = {true_centre}") +print(f"Joint (Dynesty) centre: median = {joint_median}, sigma (1D, 1 sigma) = {joint_sigma}") +print(f"EP (Laplace) centre: mean = {ep_mean}, sigma = {ep_sigma}") +print() + +""" +__Assertions__ + +1) The EP posterior mean of the shared centre is within 3 (EP) sigma of the truth. +2) The joint posterior median of the shared centre is within 3 (joint) sigma of the truth. +3) The EP and joint posteriors agree within 3 combined sigma. +""" +assert np.isfinite(ep_mean) and np.isfinite(ep_sigma) and ep_sigma > 0.0 +assert np.isfinite(joint_median) and np.isfinite(joint_sigma) and joint_sigma > 0.0 + +assert abs(ep_mean - true_centre) < 3.0 * ep_sigma, ( + f"EP centre mean ({ep_mean}) not within 3 sigma ({3.0 * ep_sigma}) of truth ({true_centre})" +) +assert abs(joint_median - true_centre) < 3.0 * joint_sigma, ( + f"joint centre median ({joint_median}) not within 3 sigma ({3.0 * joint_sigma}) " + f"of truth ({true_centre})" +) + +combined_sigma = (ep_sigma**2 + joint_sigma**2) ** 0.5 + +assert abs(ep_mean - joint_median) < 3.0 * combined_sigma, ( + f"EP mean ({ep_mean}) and joint median ({joint_median}) disagree by more than " + f"3 combined sigma ({3.0 * combined_sigma})" +) + +print("ep_parity.py: PASS")