Add variance-budget and scale-invariant fit training diagnostics

ptgp/objectives.py

@@ -284,14 +284,88 @@ 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"],
+    [
+        "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`.
@@ -303,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,
@@ -324,4 +407,73 @@ 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,
+    )
+
+
+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,
     )

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():

tests/test_objectives.py

@@ -10,8 +10,15 @@
 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,
+    unapproximated_diagnostics,
+    variance_budget,
+    vfe_diagnostics,
+)
 
 
 def _eval(*tensors):
@@ -188,3 +195,134 @@ 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)
+
+
+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)
+
+
+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)