From d2d1b7ca4f4e6d2c617aa4afe0dfd9d60834ce26 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Tue, 23 Jun 2026 14:31:33 -0700 Subject: [PATCH 1/4] Add variance_budget: mean/GP/noise variance decomposition --- ptgp/objectives.py | 62 +++++++++++++++++++++++++++++++++++++ tests/test_objectives.py | 66 ++++++++++++++++++++++++++++++++++++++-- 2 files changed, 126 insertions(+), 2 deletions(-) diff --git a/ptgp/objectives.py b/ptgp/objectives.py index 32c0f4f..5a88fff 100644 --- a/ptgp/objectives.py +++ b/ptgp/objectives.py @@ -284,6 +284,68 @@ def objective(vfe, X, y): return logdet_Kuu +VarianceBudget = namedtuple( + "VarianceBudget", + [ + "mean_var", + "signal_var", + "noise_var", + "total_var", + "frac_mean", + "frac_signal", + "frac_noise", + "var_ratio", + ], +) + + +def variance_budget(gp, X, y): + """Decompose the model-implied response variance into mean / GP / noise parts. + + Under the model ``y = m(x) + f(x) + e`` with ``f ~ GP(0, K)`` and + ``e ~ N(0, sigma^2(x))``, the law of total variance gives:: + + Var(y) = Var_x(m(x)) + E_x[K(x, x)] + E_x[sigma^2(x)] + + Returns a ``VarianceBudget`` namedtuple of symbolic TensorVariables. The + fractions are invariant to the mean and scale of ``y`` and are well defined + for any mean function, composed kernel (the GP term is the prior signal + variance ``mean(diag(K))``, so no single amplitude is needed), and scalar or + ``x``-dependent (heteroskedastic) ``sigma``. + + Fields + ------ + mean_var, signal_var, noise_var + Variance contributed by the mean function, the prior GP signal, and the + observation noise. + total_var + Their sum — the model-implied marginal variance of ``y``. + frac_mean, frac_signal, frac_noise + Each contribution as a fraction of ``total_var`` (sum to one). + var_ratio + ``total_var / Var(y)`` — calibration against the empirical data variance + (~1 when calibrated, >1 over-dispersed, <1 under-dispersed). + """ + N = X.shape[0] + mean_var = pt.var(gp.mean(X)) + signal_var = pt.mean(gp.kernel.diag(X)) + # sigma may be a scalar or a length-N heteroskedastic vector; the broadcast + # against ones(N) handles both, and mean(sigma**2) is the noise contribution. + sigma_vec = gp.likelihood.sigma * pt.ones(N) + noise_var = pt.mean(sigma_vec**2) + total_var = mean_var + signal_var + noise_var + return VarianceBudget( + mean_var=mean_var, + signal_var=signal_var, + noise_var=noise_var, + total_var=total_var, + frac_mean=mean_var / total_var, + frac_signal=signal_var / total_var, + frac_noise=noise_var / total_var, + var_ratio=total_var / pt.var(y), + ) + + VFEDiagnostics = namedtuple( "VFEDiagnostics", ["elbo", "fit", "trace_penalty", "nystrom_residual", "sigma", "fit_per_n", "excess_fit_per_n"], diff --git a/tests/test_objectives.py b/tests/test_objectives.py index d5f4fa3..2a5527a 100644 --- a/tests/test_objectives.py +++ b/tests/test_objectives.py @@ -10,8 +10,8 @@ from ptgp.inducing import Points from ptgp.kernels import ExpQuad from ptgp.likelihoods import Gaussian -from ptgp.mean import Zero -from ptgp.objectives import collapsed_elbo, elbo, marginal_log_likelihood +from ptgp.mean import Constant, Linear, Zero +from ptgp.objectives import collapsed_elbo, elbo, marginal_log_likelihood, variance_budget def _eval(*tensors): @@ -188,3 +188,65 @@ def test_collapsed_elbo_less_than_mll(self, regression_data, inducing_points): ) assert celbo <= mll_val + 1e-6 # collapsed ELBO <= MLL + + +class TestVarianceBudget: + @pytest.fixture + def data(self): + rng = np.random.default_rng(0) + X = np.sort(rng.uniform(0, 5, 30))[:, None].astype(np.float64) + # nonzero offset (2.0) to exercise mean-invariance + y = np.sin(X.ravel()) + 0.1 * rng.standard_normal(30) + 2.0 + return X, y + + def _gp(self, eta=1.0, sigma=0.3, mean=None): + return Unapproximated( + kernel=eta**2 * ExpQuad(input_dim=1, ls=1.0), + mean=mean if mean is not None else Zero(), + sigma=sigma, + ) + + def _budget(self, gp, X, y): + b = variance_budget(gp, pt.as_tensor_variable(X), pt.as_tensor_variable(y)) + return b._make(_eval(*b)) + + def test_fractions_sum_to_one(self, data): + X, y = data + b = self._budget(self._gp(), X, y) + assert np.isclose(b.frac_mean + b.frac_signal + b.frac_noise, 1.0) + + def test_mean_invariance(self, data): + X, y = data + b0 = self._budget(self._gp(), X, y) + b1 = self._budget(self._gp(), X, y + 10.0) + for f in ("frac_mean", "frac_signal", "frac_noise", "var_ratio"): + assert np.isclose(getattr(b0, f), getattr(b1, f)), f + + def test_scale_invariance(self, data): + X, y = data + a = 7.0 + b0 = self._budget(self._gp(eta=1.0, sigma=0.3), X, y) + # scale y, and the model's amplitude and noise, by a + b1 = self._budget(self._gp(eta=a, sigma=a * 0.3), X, a * y) + for f in ("frac_mean", "frac_signal", "frac_noise", "var_ratio"): + assert np.isclose(getattr(b0, f), getattr(b1, f)), f + + def test_linear_mean_contributes_variance(self, data): + X, y = data + gp_lin = self._gp(mean=Linear(coeffs=pt.as_tensor_variable(np.array([1.0])))) + assert self._budget(gp_lin, X, y).frac_mean > 0.0 + # a constant mean sets the level, not the variance -> 0 contribution + gp_const = self._gp(mean=Constant(c=5.0)) + assert np.isclose(self._budget(gp_const, X, y).frac_mean, 0.0) + + def test_heteroskedastic_sigma(self, data): + X, y = data + Xt = pt.as_tensor_variable(X) + sigma = 0.1 + 0.05 * Xt[:, 0] ** 2 # length-N vector + gp = Unapproximated(kernel=ExpQuad(input_dim=1, ls=1.0), mean=Zero(), sigma=sigma) + b = variance_budget(gp, Xt, pt.as_tensor_variable(y))._make( + _eval(*variance_budget(gp, Xt, pt.as_tensor_variable(y))) + ) + expected_noise = float(np.mean((0.1 + 0.05 * X[:, 0] ** 2) ** 2)) + assert np.isclose(b.noise_var, expected_noise) + assert np.isclose(b.frac_mean + b.frac_signal + b.frac_noise, 1.0) From a33d0d599de5f4b80ff64c18329d56e3e18a1b58 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Tue, 23 Jun 2026 14:38:44 -0700 Subject: [PATCH 2/4] Make excess_fit_per_n scale-invariant; add budget to vfe_diagnostics --- ptgp/objectives.py | 41 ++++++++++++++++++++++++++++++++-------- tests/test_objectives.py | 40 ++++++++++++++++++++++++++++++++++++++- 2 files changed, 72 insertions(+), 9 deletions(-) diff --git a/ptgp/objectives.py b/ptgp/objectives.py index 5a88fff..e80ccdf 100644 --- a/ptgp/objectives.py +++ b/ptgp/objectives.py @@ -348,12 +348,24 @@ def variance_budget(gp, X, y): VFEDiagnostics = namedtuple( "VFEDiagnostics", - ["elbo", "fit", "trace_penalty", "nystrom_residual", "sigma", "fit_per_n", "excess_fit_per_n"], + [ + "elbo", + "fit", + "trace_penalty", + "nystrom_residual", + "sigma", + "fit_per_n", + "excess_fit_per_n", + "frac_mean", + "frac_signal", + "frac_noise", + "var_ratio", + ], ) def vfe_diagnostics(vfe, X, y): - """Collapsed ELBO terms plus sigma and two normalised fit metrics. + """Collapsed ELBO terms, fit metrics, and the mean/GP/noise variance budget. Returns a ``VFEDiagnostics`` namedtuple of symbolic TensorVariables, suitable for use with :func:`ptgp.optim.compile_scipy_diagnostics`. @@ -365,19 +377,28 @@ def vfe_diagnostics(vfe, X, y): nystrom_residual ``tr(Kff - Qff) / N`` — per-point Nyström approximation error. sigma - Likelihood noise (constrained space). + Likelihood noise (constrained space); the mean of ``sigma`` when it is + heteroskedastic. fit_per_n - ``fit / N`` — scale-invariant data fit. + ``fit / N`` — per-point data fit. excess_fit_per_n - ``fit_per_n + 0.5 * log(2π σ²)`` — how much better than noise floor. - Goes to zero when the model fits at the noise level only. + ``fit_per_n + 0.5 * log(2π * Var(y - m(X))) + 0.5`` — per-point fit + relative to a constant-mean Gaussian at the residual variance. Reads 0 + when the kernel does no better than a flat mean and grows as it explains + structure. Referencing the residual variance (not ``sigma**2``) makes it + invariant to the scale of ``y``; pair it with the variance budget for the + mean-invariant view. + frac_mean, frac_signal, frac_noise, var_ratio + The mean/GP/noise variance budget — see :func:`variance_budget`. """ terms = collapsed_elbo(vfe, X, y) + budget = variance_budget(vfe, X, y) N = X.shape[0] sigma_vec = vfe.likelihood.sigma * pt.ones(N) - sigma_mean = pt.mean(sigma_vec) # scalar; mean of a constant vector = that constant + sigma_mean = pt.mean(sigma_vec) fit_per_n = terms.fit / N - excess_fit_per_n = fit_per_n + 0.5 * pt.log(2.0 * np.pi * sigma_mean**2) + resid_var = pt.var(y - vfe.mean(X)) + excess_fit_per_n = fit_per_n + 0.5 * pt.log(2.0 * np.pi * resid_var) + 0.5 return VFEDiagnostics( elbo=terms.elbo, fit=terms.fit, @@ -386,4 +407,8 @@ def vfe_diagnostics(vfe, X, y): sigma=sigma_mean, fit_per_n=fit_per_n, excess_fit_per_n=excess_fit_per_n, + frac_mean=budget.frac_mean, + frac_signal=budget.frac_signal, + frac_noise=budget.frac_noise, + var_ratio=budget.var_ratio, ) diff --git a/tests/test_objectives.py b/tests/test_objectives.py index 2a5527a..56de43b 100644 --- a/tests/test_objectives.py +++ b/tests/test_objectives.py @@ -11,7 +11,13 @@ from ptgp.kernels import ExpQuad from ptgp.likelihoods import Gaussian from ptgp.mean import Constant, Linear, Zero -from ptgp.objectives import collapsed_elbo, elbo, marginal_log_likelihood, variance_budget +from ptgp.objectives import ( + collapsed_elbo, + elbo, + marginal_log_likelihood, + variance_budget, + vfe_diagnostics, +) def _eval(*tensors): @@ -250,3 +256,35 @@ def test_heteroskedastic_sigma(self, data): expected_noise = float(np.mean((0.1 + 0.05 * X[:, 0] ** 2) ** 2)) assert np.isclose(b.noise_var, expected_noise) assert np.isclose(b.frac_mean + b.frac_signal + b.frac_noise, 1.0) + + +class TestVFEDiagnostics: + @pytest.fixture + def setup(self): + rng = np.random.default_rng(1) + X = np.sort(rng.uniform(0, 5, 30))[:, None].astype(np.float64) + y = np.sin(X.ravel()) + 0.1 * rng.standard_normal(30) + 3.0 # offset + Z = np.linspace(0.5, 4.5, 6)[:, None].astype(np.float64) + return X, y, Z + + def _diag(self, X, y, Z, eta=1.0, sigma=0.3): + vfe = VFE( + kernel=eta**2 * ExpQuad(input_dim=1, ls=1.0), + mean=Zero(), + sigma=sigma, + inducing_variable=Points(pt.as_tensor_variable(Z)), + ) + d = vfe_diagnostics(vfe, pt.as_tensor_variable(X), pt.as_tensor_variable(y)) + return d._make(_eval(*d)) + + def test_budget_fields_present(self, setup): + X, y, Z = setup + d = self._diag(X, y, Z) + assert np.isclose(d.frac_mean + d.frac_signal + d.frac_noise, 1.0) + + def test_excess_fit_scale_invariant(self, setup): + X, y, Z = setup + a = 5.0 + d0 = self._diag(X, y, Z, eta=1.0, sigma=0.3) + d1 = self._diag(X, a * y, Z, eta=a, sigma=a * 0.3) + assert np.isclose(d0.excess_fit_per_n, d1.excess_fit_per_n) From 3714a9d76d1846eb8ee613c163a44518c9bf970c Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Tue, 23 Jun 2026 14:44:04 -0700 Subject: [PATCH 3/4] Add unapproximated_diagnostics for the exact GP --- ptgp/objectives.py | 65 ++++++++++++++++++++++++++++++++++++++++ tests/test_objectives.py | 38 +++++++++++++++++++++++ 2 files changed, 103 insertions(+) diff --git a/ptgp/objectives.py b/ptgp/objectives.py index e80ccdf..2b9e143 100644 --- a/ptgp/objectives.py +++ b/ptgp/objectives.py @@ -412,3 +412,68 @@ def vfe_diagnostics(vfe, X, y): frac_noise=budget.frac_noise, var_ratio=budget.var_ratio, ) + + +UnapproximatedDiagnostics = namedtuple( + "UnapproximatedDiagnostics", + [ + "mll", + "fit", + "logdet", + "sigma", + "fit_per_n", + "logdet_per_n", + "excess_fit_per_n", + "frac_mean", + "frac_signal", + "frac_noise", + "var_ratio", + ], +) + + +def unapproximated_diagnostics(gp, X, y): + """Exact-GP marginal-likelihood terms, fit metrics, and the variance budget. + + The exact-GP analogue of :func:`vfe_diagnostics`, for + :class:`ptgp.gp.Unapproximated`. Returns an ``UnapproximatedDiagnostics`` + namedtuple of symbolic TensorVariables, for use with + :func:`ptgp.optim.compile_scipy_diagnostics`. + + Fields + ------ + mll, fit, logdet + Direct from :func:`marginal_log_likelihood` (``mll = fit + logdet``; + ``fit`` is the data-fit quadratic, ``logdet`` the Occam complexity term). + sigma + Likelihood noise (the mean of ``sigma`` when it is heteroskedastic). + fit_per_n, logdet_per_n + ``fit / N`` and ``logdet / N`` — per-point data fit and complexity. + excess_fit_per_n + ``mll / N + 0.5 * log(2π * Var(y - m(X))) + 0.5`` — per-point evidence + relative to a constant-mean Gaussian at the residual variance. Reads 0 at + that baseline and is invariant to the scale of ``y`` (the residual-variance + reference cancels the log-determinant's scale dependence). + frac_mean, frac_signal, frac_noise, var_ratio + The mean/GP/noise variance budget — see :func:`variance_budget`. + """ + terms = marginal_log_likelihood(gp, X, y) + budget = variance_budget(gp, X, y) + N = X.shape[0] + sigma_vec = gp.likelihood.sigma * pt.ones(N) + sigma_mean = pt.mean(sigma_vec) + resid_var = pt.var(y - gp.mean(X)) + excess_fit_per_n = terms.mll / N + 0.5 * pt.log(2.0 * np.pi * resid_var) + 0.5 + return UnapproximatedDiagnostics( + mll=terms.mll, + fit=terms.fit, + logdet=terms.logdet, + sigma=sigma_mean, + fit_per_n=terms.fit / N, + logdet_per_n=terms.logdet / N, + excess_fit_per_n=excess_fit_per_n, + frac_mean=budget.frac_mean, + frac_signal=budget.frac_signal, + frac_noise=budget.frac_noise, + var_ratio=budget.var_ratio, + ) diff --git a/tests/test_objectives.py b/tests/test_objectives.py index 56de43b..8d8db90 100644 --- a/tests/test_objectives.py +++ b/tests/test_objectives.py @@ -15,6 +15,7 @@ collapsed_elbo, elbo, marginal_log_likelihood, + unapproximated_diagnostics, variance_budget, vfe_diagnostics, ) @@ -288,3 +289,40 @@ def test_excess_fit_scale_invariant(self, setup): d0 = self._diag(X, y, Z, eta=1.0, sigma=0.3) d1 = self._diag(X, a * y, Z, eta=a, sigma=a * 0.3) assert np.isclose(d0.excess_fit_per_n, d1.excess_fit_per_n) + + +class TestUnapproximatedDiagnostics: + @pytest.fixture + def data(self): + rng = np.random.default_rng(2) + X = np.sort(rng.uniform(0, 5, 25))[:, None].astype(np.float64) + y = np.sin(X.ravel()) + 0.1 * rng.standard_normal(25) + 4.0 # offset + return X, y + + def _diag(self, X, y, eta=1.0, sigma=0.3): + gp = Unapproximated(kernel=eta**2 * ExpQuad(input_dim=1, ls=1.0), mean=Zero(), sigma=sigma) + d = unapproximated_diagnostics(gp, pt.as_tensor_variable(X), pt.as_tensor_variable(y)) + return d._make(_eval(*d)) + + def test_terms_and_budget(self, data): + X, y = data + d = self._diag(X, y) + assert np.isclose(d.mll, d.fit + d.logdet) + assert np.isclose(d.frac_mean + d.frac_signal + d.frac_noise, 1.0) + + def test_excess_fit_scale_invariant(self, data): + X, y = data + a = 6.0 + d0 = self._diag(X, y, eta=1.0, sigma=0.3) + d1 = self._diag(X, a * y, eta=a, sigma=a * 0.3) + assert np.isclose(d0.excess_fit_per_n, d1.excess_fit_per_n) + + def test_heteroskedastic_sigma(self, data): + X, y = data + Xt = pt.as_tensor_variable(X) + sigma = 0.1 + 0.05 * Xt[:, 0] ** 2 # length-N vector + gp = Unapproximated(kernel=ExpQuad(input_dim=1, ls=1.0), mean=Zero(), sigma=sigma) + d = unapproximated_diagnostics(gp, Xt, pt.as_tensor_variable(y)) + d = d._make(_eval(*d)) + assert np.isfinite(d.mll) + assert np.isclose(d.frac_mean + d.frac_signal + d.frac_noise, 1.0) From c0f36d1b253e878fa2875c56dd8a7060ee26c215 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Tue, 23 Jun 2026 14:49:01 -0700 Subject: [PATCH 4/4] Test variance-budget fields reach idata optimizer_result --- tests/test_idata.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_idata.py b/tests/test_idata.py index d9b2b1c..7f9d70b 100644 --- a/tests/test_idata.py +++ b/tests/test_idata.py @@ -117,6 +117,10 @@ def test_optimizer_result_combines_scalars_and_trajectory(): assert opt.sizes["iteration"] == len(history) assert opt.elbo.dims == ("iteration",) assert opt.trace_penalty.dims == ("iteration",) + # The mean/GP/noise variance budget and the redefined excess fit flow through too. + assert opt.excess_fit_per_n.dims == ("iteration",) + assert opt.frac_signal.dims == ("iteration",) + assert opt.var_ratio.dims == ("iteration",) def test_no_result_no_history_omits_optimizer_result():