From c05c94ec03f5d17808c9459fd1f877b74bba14a7 Mon Sep 17 00:00:00 2001 From: Michal-Novomestsky Date: Wed, 2 Jul 2025 17:11:28 +1000 Subject: [PATCH 1/5] changed laplace approx to return MvNormal --- pymc_extras/inference/laplace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc_extras/inference/laplace.py b/pymc_extras/inference/laplace.py index d64d2ada..1b9f3048 100644 --- a/pymc_extras/inference/laplace.py +++ b/pymc_extras/inference/laplace.py @@ -149,7 +149,7 @@ def get_conditional_gaussian_approximation( # Currently x is passed both as the query point for f(x, args) = logp(x | y, params) AND as an initial guess for x0. This may cause issues if the query point is # far from the mode x0 or in a neighbourhood which results in poor convergence. - return pytensor.function(args, [x0, conditional_gaussian_approx]) + return pytensor.function(args, pm.MvNormal(mu=x0, tau=Q-hess)) def laplace_draws_to_inferencedata( From 37fa1edc96f0d7cada50ec1773279209b4614738 Mon Sep 17 00:00:00 2001 From: Michal-Novomestsky Date: Wed, 2 Jul 2025 17:28:48 +1000 Subject: [PATCH 2/5] added seperate line for evaluating Q-hess --- pymc_extras/inference/laplace.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/pymc_extras/inference/laplace.py b/pymc_extras/inference/laplace.py index 1b9f3048..996f5ac0 100644 --- a/pymc_extras/inference/laplace.py +++ b/pymc_extras/inference/laplace.py @@ -143,13 +143,16 @@ def get_conditional_gaussian_approximation( # Full log(p(x | y, params)) using the Laplace approximation (up to a constant) _, logdetQ = pt.nlinalg.slogdet(Q) - conditional_gaussian_approx = ( - -0.5 * x.T @ (-hess + Q) @ x + x.T @ (Q @ mu + jac - hess @ x0) + 0.5 * logdetQ - ) + # conditional_gaussian_approx = ( + # -0.5 * x.T @ (-hess + Q) @ x + x.T @ (Q @ mu + jac - hess @ x0) + 0.5 * logdetQ + # ) + + # In the future, this could be made more efficient with only adding the diagonal of -hess + tau = Q - hess # Currently x is passed both as the query point for f(x, args) = logp(x | y, params) AND as an initial guess for x0. This may cause issues if the query point is # far from the mode x0 or in a neighbourhood which results in poor convergence. - return pytensor.function(args, pm.MvNormal(mu=x0, tau=Q-hess)) + return pytensor.function(args, [x0, pm.MvNormal(mu=x0, tau=tau)]) def laplace_draws_to_inferencedata( From 83bef75972e67249b10c3e7b66afbf4160414774 Mon Sep 17 00:00:00 2001 From: Michal-Novomestsky Date: Fri, 4 Jul 2025 16:15:36 +1000 Subject: [PATCH 3/5] WIP: minor refactor --- pymc_extras/inference/laplace.py | 20 ++++++++++++++------ tests/test_laplace.py | 2 +- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/pymc_extras/inference/laplace.py b/pymc_extras/inference/laplace.py index 996f5ac0..11858c27 100644 --- a/pymc_extras/inference/laplace.py +++ b/pymc_extras/inference/laplace.py @@ -121,11 +121,13 @@ def get_conditional_gaussian_approximation( # f = log(p(y | x, params)) f_x = model.logp() - jac = pytensor.gradient.grad(f_x, x) - hess = pytensor.gradient.jacobian(jac.flatten(), x) + # jac = pytensor.gradient.grad(f_x, x) + # hess = pytensor.gradient.jacobian(jac.flatten(), x) # log(p(x | y, params)) only including terms that depend on x for the minimization step (logdet(Q) ignored as it is a constant wrt x) - log_x_posterior = f_x - 0.5 * (x - mu).T @ Q @ (x - mu) + log_x_posterior = f_x - 0.5 * (x - mu).T @ Q @ ( + x - mu + ) # TODO could be f + x.logp - IS X.LOGP DUPLICATED IN F? # Maximize log(p(x | y, params)) wrt x to find mode x0 x0, _ = minimize( @@ -138,11 +140,13 @@ def get_conditional_gaussian_approximation( ) # require f'(x0) and f''(x0) for Laplace approx - jac = pytensor.graph.replace.graph_replace(jac, {x: x0}) + # jac = pytensor.graph.replace.graph_replace(jac, {x: x0}) + jac = pytensor.gradient.grad(f_x, x) + hess = pytensor.gradient.jacobian(jac.flatten(), x) hess = pytensor.graph.replace.graph_replace(hess, {x: x0}) # Full log(p(x | y, params)) using the Laplace approximation (up to a constant) - _, logdetQ = pt.nlinalg.slogdet(Q) + # _, logdetQ = pt.nlinalg.slogdet(Q) # conditional_gaussian_approx = ( # -0.5 * x.T @ (-hess + Q) @ x + x.T @ (Q @ mu + jac - hess @ x0) + 0.5 * logdetQ # ) @@ -152,7 +156,11 @@ def get_conditional_gaussian_approximation( # Currently x is passed both as the query point for f(x, args) = logp(x | y, params) AND as an initial guess for x0. This may cause issues if the query point is # far from the mode x0 or in a neighbourhood which results in poor convergence. - return pytensor.function(args, [x0, pm.MvNormal(mu=x0, tau=tau)]) + return ( + x0, + pm.MvNormal(f"{x.name}_laplace_approx", mu=x0, tau=tau), + tau, + ) # pytensor.function(args, [x0, pm.MvNormal(mu=x0, tau=tau)]) def laplace_draws_to_inferencedata( diff --git a/tests/test_laplace.py b/tests/test_laplace.py index 72ff3e93..9aafa6e8 100644 --- a/tests/test_laplace.py +++ b/tests/test_laplace.py @@ -323,7 +323,7 @@ def test_get_conditional_gaussian_approximation(): Q = pm.MvNormal("Q", mu=Q_mu, cov=Q_cov) # Pytensor currently doesn't support autograd for pt inverses, so we use a numeric Q instead - x = pm.MvNormal("x", mu=mu_param, cov=np.linalg.inv(Q_val)) + x = pm.MvNormal("x", mu=mu_param, tau=Q) # cov=np.linalg.inv(Q_val)) y = pm.MvNormal( "y", From 13961fab907209b3d1b7d00907f3ab5c62288b9a Mon Sep 17 00:00:00 2001 From: Michal Novomestsky Date: Sun, 6 Jul 2025 21:11:37 +1000 Subject: [PATCH 4/5] started writing fit_INLA routine --- pymc_extras/inference/fit.py | 12 ++- pymc_extras/inference/inla.py | 164 +++++++++++++++++++++++++++++++ pymc_extras/inference/laplace.py | 111 --------------------- tests/test_inla.py | 105 ++++++++++++++++++++ tests/test_laplace.py | 83 ---------------- 5 files changed, 280 insertions(+), 195 deletions(-) create mode 100644 pymc_extras/inference/inla.py create mode 100644 tests/test_inla.py diff --git a/pymc_extras/inference/fit.py b/pymc_extras/inference/fit.py index 5b83ff1f..8814ba3b 100644 --- a/pymc_extras/inference/fit.py +++ b/pymc_extras/inference/fit.py @@ -36,7 +36,17 @@ def fit(method: str, **kwargs) -> az.InferenceData: return fit_pathfinder(**kwargs) - if method == "laplace": + elif method == "laplace": from pymc_extras.inference.laplace import fit_laplace return fit_laplace(**kwargs) + + elif method == "INLA": + from pymc_extras.inference.laplace import fit_INLA + + return fit_INLA(**kwargs) + + else: + raise ValueError( + f"method '{method}' not supported. Use one of 'pathfinder', 'laplace' or 'INLA'." + ) diff --git a/pymc_extras/inference/inla.py b/pymc_extras/inference/inla.py new file mode 100644 index 00000000..0eaa649c --- /dev/null +++ b/pymc_extras/inference/inla.py @@ -0,0 +1,164 @@ +import arviz as az +import numpy as np +import pymc as pm +import pytensor +import pytensor.tensor as pt + +from better_optimize.constants import minimize_method +from numpy.typing import ArrayLike +from pytensor.tensor import TensorVariable +from pytensor.tensor.optimize import minimize + + +def get_conditional_gaussian_approximation( + x: TensorVariable, + Q: TensorVariable | ArrayLike, + mu: TensorVariable | ArrayLike, + model: pm.Model | None = None, + method: minimize_method = "BFGS", + use_jac: bool = True, + use_hess: bool = False, + optimizer_kwargs: dict | None = None, +) -> list[TensorVariable]: + """ + Returns an estimate the a posteriori probability of a latent Gaussian field x and its mode x0 using the Laplace approximation. + + That is: + y | x, sigma ~ N(Ax, sigma^2 W) + x | params ~ N(mu, Q(params)^-1) + + We seek to estimate p(x | y, params) with a Gaussian: + + log(p(x | y, params)) = log(p(y | x, params)) + log(p(x | params)) + const + + Let f(x) = log(p(y | x, params)). From the definition of our model above, we have log(p(x | params)) = -0.5*(x - mu).T Q (x - mu) + 0.5*logdet(Q). + + This gives log(p(x | y, params)) = f(x) - 0.5*(x - mu).T Q (x - mu) + 0.5*logdet(Q). We will estimate this using the Laplace approximation by Taylor expanding f(x) about the mode. + + Thus: + + 1. Maximize log(p(x | y, params)) = f(x) - 0.5*(x - mu).T Q (x - mu) wrt x (note that logdet(Q) does not depend on x) to find the mode x0. + + 2. Use the Laplace approximation expanded about the mode: p(x | y, params) ~= N(mu=x0, tau=Q - f''(x0)). + + Parameters + ---------- + x: TensorVariable + The parameter with which to maximize wrt (that is, find the mode in x). In INLA, this is the latent Gaussian field x~N(mu,Q^-1). + Q: TensorVariable | ArrayLike + The precision matrix of the latent field x. + mu: TensorVariable | ArrayLike + The mean of the latent field x. + model: Model + PyMC model to use. + method: minimize_method + Which minimization algorithm to use. + use_jac: bool + If true, the minimizer will compute the gradient of log(p(x | y, params)). + use_hess: bool + If true, the minimizer will compute the Hessian log(p(x | y, params)). + optimizer_kwargs: dict + Kwargs to pass to scipy.optimize.minimize. + + Returns + ------- + x0, p(x | y, params): list[TensorVariable] + Mode and Laplace approximation for posterior. + """ + model = pm.modelcontext(model) + + # f = log(p(y | x, params)) + f_x = model.logp() + + # log(p(x | y, params)) only including terms that depend on x for the minimization step (logdet(Q) ignored as it is a constant wrt x) + log_x_posterior = f_x - 0.5 * (x - mu).T @ Q @ (x - mu) + + # Maximize log(p(x | y, params)) wrt x to find mode x0 + x0, _ = minimize( + objective=-log_x_posterior, + x=x, + method=method, + jac=use_jac, + hess=use_hess, + optimizer_kwargs=optimizer_kwargs, + ) + + # require f''(x0) for Laplace approx + hess = pytensor.gradient.hessian(f_x, x) + hess = pytensor.graph.replace.graph_replace(hess, {x: x0}) + + # Could be made more efficient with adding diagonals only + tau = Q - hess + + # Currently x is passed both as the query point for f(x, args) = logp(x | y, params) AND as an initial guess for x0. This may cause issues if the query point is + # far from the mode x0 or in a neighbourhood which results in poor convergence. + return x0, pm.MvNormal(f"{x.name}_laplace_approx", mu=x0, tau=tau) + + +def get_log_marginal_likelihood( + x: TensorVariable, + Q: TensorVariable | ArrayLike, + mu: TensorVariable | ArrayLike, + model: pm.Model | None = None, + method: minimize_method = "BFGS", + use_jac: bool = True, + use_hess: bool = False, + optimizer_kwargs: dict | None = None, +) -> TensorVariable: + model = pm.modelcontext(model) + + x0, laplace_approx = get_conditional_gaussian_approximation( + x, Q, mu, model, method, use_jac, use_hess, optimizer_kwargs + ) + log_laplace_approx = pm.logp(laplace_approx, model.rvs_to_values[x]) + + _, logdetQ = pt.nlinalg.slogdet(Q) + log_x_likelihood = ( + -0.5 * (x - mu).T @ Q @ (x - mu) + 0.5 * logdetQ - 0.5 * x.shape[0] * np.log(2 * np.pi) + ) + + log_likelihood = ( # logp(y | params) = + model.logp() # logp(y | x, params) + + log_x_likelihood # * logp(x | params) + - log_laplace_approx # / logp(x | y, params) + ) + + return log_likelihood + + +def fit_INLA( + x: TensorVariable, + Q: TensorVariable | ArrayLike, + mu: TensorVariable | ArrayLike, + model: pm.Model | None = None, + method: minimize_method = "BFGS", + use_jac: bool = True, + use_hess: bool = False, + optimizer_kwargs: dict | None = None, +) -> az.InferenceData: + model = pm.modelcontext(model) + + # logp(y | params) + log_likelihood = get_log_marginal_likelihood( + x, Q, mu, model, method, use_jac, use_hess, optimizer_kwargs + ) + + # TODO How to obtain prior? It can parametrise Q, mu, y, etc. Not sure if we could extract from model.logp somehow. Otherwise simply specify as a user input + prior = None + params = None + log_prior = pm.logp(prior, model.rvs_to_values[params]) + + # logp(params | y) = logp(y | params) + logp(params) + const + log_posterior = log_likelihood + log_prior + + # TODO log_marginal_x_likelihood is almost the same as log_likelihood, but need to do some sampling? + log_marginal_x_likelihood = None + log_marginal_x_posterior = log_marginal_x_likelihood + log_prior + + # TODO can we sample over log likelihoods? + # Marginalize params + idata_params = log_posterior.sample() # TODO something like NUTS, QMC, etc.? + idata_x = log_marginal_x_posterior.sample() + + # Bundle up idatas somehow + return idata_params, idata_x diff --git a/pymc_extras/inference/laplace.py b/pymc_extras/inference/laplace.py index 11858c27..9c0a0d27 100644 --- a/pymc_extras/inference/laplace.py +++ b/pymc_extras/inference/laplace.py @@ -15,7 +15,6 @@ import logging -from collections.abc import Callable from functools import reduce from importlib.util import find_spec from itertools import product @@ -30,7 +29,6 @@ from arviz import dict_to_dataset from better_optimize.constants import minimize_method -from numpy.typing import ArrayLike from pymc import DictToArrayBijection from pymc.backends.arviz import ( coords_and_dims_for_inferencedata, @@ -41,8 +39,6 @@ from pymc.model.transform.conditioning import remove_value_transforms from pymc.model.transform.optimization import freeze_dims_and_data from pymc.util import get_default_varnames -from pytensor.tensor import TensorVariable -from pytensor.tensor.optimize import minimize from scipy import stats from pymc_extras.inference.find_map import ( @@ -56,113 +52,6 @@ _log = logging.getLogger(__name__) -def get_conditional_gaussian_approximation( - x: TensorVariable, - Q: TensorVariable | ArrayLike, - mu: TensorVariable | ArrayLike, - args: list[TensorVariable] | None = None, - model: pm.Model | None = None, - method: minimize_method = "BFGS", - use_jac: bool = True, - use_hess: bool = False, - optimizer_kwargs: dict | None = None, -) -> Callable: - """ - Returns a function to estimate the a posteriori log probability of a latent Gaussian field x and its mode x0 using the Laplace approximation. - - That is: - y | x, sigma ~ N(Ax, sigma^2 W) - x | params ~ N(mu, Q(params)^-1) - - We seek to estimate log(p(x | y, params)): - - log(p(x | y, params)) = log(p(y | x, params)) + log(p(x | params)) + const - - Let f(x) = log(p(y | x, params)). From the definition of our model above, we have log(p(x | params)) = -0.5*(x - mu).T Q (x - mu) + 0.5*logdet(Q). - - This gives log(p(x | y, params)) = f(x) - 0.5*(x - mu).T Q (x - mu) + 0.5*logdet(Q). We will estimate this using the Laplace approximation by Taylor expanding f(x) about the mode. - - Thus: - - 1. Maximize log(p(x | y, params)) = f(x) - 0.5*(x - mu).T Q (x - mu) wrt x (note that logdet(Q) does not depend on x) to find the mode x0. - - 2. Substitute x0 into the Laplace approximation expanded about the mode: log(p(x | y, params)) ~= -0.5*x.T (-f''(x0) + Q) x + x.T (Q.mu + f'(x0) - f''(x0).x0) + 0.5*logdet(Q). - - Parameters - ---------- - x: TensorVariable - The parameter with which to maximize wrt (that is, find the mode in x). In INLA, this is the latent field x~N(mu,Q^-1). - Q: TensorVariable | ArrayLike - The precision matrix of the latent field x. - mu: TensorVariable | ArrayLike - The mean of the latent field x. - args: list[TensorVariable] - Args to supply to the compiled function. That is, (x0, logp) = f(x, *args). If set to None, assumes the model RVs are args. - model: Model - PyMC model to use. - method: minimize_method - Which minimization algorithm to use. - use_jac: bool - If true, the minimizer will compute the gradient of log(p(x | y, params)). - use_hess: bool - If true, the minimizer will compute the Hessian log(p(x | y, params)). - optimizer_kwargs: dict - Kwargs to pass to scipy.optimize.minimize. - - Returns - ------- - f: Callable - A function which accepts a value of x and args and returns [x0, log(p(x | y, params))], where x0 is the mode. x is currently both the point at which to evaluate logp and the initial guess for the minimizer. - """ - model = pm.modelcontext(model) - - if args is None: - args = model.continuous_value_vars + model.discrete_value_vars - - # f = log(p(y | x, params)) - f_x = model.logp() - # jac = pytensor.gradient.grad(f_x, x) - # hess = pytensor.gradient.jacobian(jac.flatten(), x) - - # log(p(x | y, params)) only including terms that depend on x for the minimization step (logdet(Q) ignored as it is a constant wrt x) - log_x_posterior = f_x - 0.5 * (x - mu).T @ Q @ ( - x - mu - ) # TODO could be f + x.logp - IS X.LOGP DUPLICATED IN F? - - # Maximize log(p(x | y, params)) wrt x to find mode x0 - x0, _ = minimize( - objective=-log_x_posterior, - x=x, - method=method, - jac=use_jac, - hess=use_hess, - optimizer_kwargs=optimizer_kwargs, - ) - - # require f'(x0) and f''(x0) for Laplace approx - # jac = pytensor.graph.replace.graph_replace(jac, {x: x0}) - jac = pytensor.gradient.grad(f_x, x) - hess = pytensor.gradient.jacobian(jac.flatten(), x) - hess = pytensor.graph.replace.graph_replace(hess, {x: x0}) - - # Full log(p(x | y, params)) using the Laplace approximation (up to a constant) - # _, logdetQ = pt.nlinalg.slogdet(Q) - # conditional_gaussian_approx = ( - # -0.5 * x.T @ (-hess + Q) @ x + x.T @ (Q @ mu + jac - hess @ x0) + 0.5 * logdetQ - # ) - - # In the future, this could be made more efficient with only adding the diagonal of -hess - tau = Q - hess - - # Currently x is passed both as the query point for f(x, args) = logp(x | y, params) AND as an initial guess for x0. This may cause issues if the query point is - # far from the mode x0 or in a neighbourhood which results in poor convergence. - return ( - x0, - pm.MvNormal(f"{x.name}_laplace_approx", mu=x0, tau=tau), - tau, - ) # pytensor.function(args, [x0, pm.MvNormal(mu=x0, tau=tau)]) - - def laplace_draws_to_inferencedata( posterior_draws: list[np.ndarray[float | int]], model: pm.Model | None = None ) -> az.InferenceData: diff --git a/tests/test_inla.py b/tests/test_inla.py new file mode 100644 index 00000000..cda24bb5 --- /dev/null +++ b/tests/test_inla.py @@ -0,0 +1,105 @@ +# Copyright 2024 The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import numpy as np +import pymc as pm +import pytensor + +from pymc_extras.inference.inla import get_conditional_gaussian_approximation + + +def test_get_conditional_gaussian_approximation(): + """ + Consider the trivial case of: + + y | x ~ N(x, cov_param) + x | param ~ N(mu_param, Q^-1) + + cov_param ~ N(cov_mu, cov_cov) + mu_param ~ N(mu_mu, mu_cov) + Q ~ N(Q_mu, Q_cov) + + This has an analytic solution at the mode which we can compare against. + """ + rng = np.random.default_rng(12345) + n = 10000 + d = 10 + + # Initialise arrays + mu_true = rng.random(d) + cov_true = np.diag(rng.random(d)) + Q_val = np.diag(rng.random(d)) + cov_param_val = np.diag(rng.random(d)) + + x_val = rng.random(d) + mu_val = rng.random(d) + + mu_mu = rng.random(d) + mu_cov = np.diag(np.ones(d)) + cov_mu = rng.random(d**2) + cov_cov = np.diag(np.ones(d**2)) + Q_mu = rng.random(d**2) + Q_cov = np.diag(np.ones(d**2)) + + with pm.Model() as model: + y_obs = rng.multivariate_normal(mean=mu_true, cov=cov_true, size=n) + + mu_param = pm.MvNormal("mu_param", mu=mu_mu, cov=mu_cov) + cov_param = pm.MvNormal("cov_param", mu=cov_mu, cov=cov_cov) + Q = pm.MvNormal("Q", mu=Q_mu, cov=Q_cov) + + x = pm.MvNormal("x", mu=mu_param, tau=Q_val) + + y = pm.MvNormal( + "y", + mu=x, + cov=cov_param.reshape((d, d)), + observed=y_obs, + ) + + args = model.continuous_value_vars + model.discrete_value_vars + + # logp(x | y, params) + x0, x_g = get_conditional_gaussian_approximation( + x=model.rvs_to_values[x], + Q=Q.reshape((d, d)), + mu=mu_param, + optimizer_kwargs={"tol": 1e-25}, + ) + + cga = pytensor.function(args, [x0, pm.logp(x_g, model.rvs_to_values[x])]) + + x0, log_x_posterior = cga( + x=x_val, mu_param=mu_val, cov_param=cov_param_val.flatten(), Q=Q_val.flatten() + ) + + # Get analytic values of the mode and Laplace-approximated log posterior + cov_param_inv = np.linalg.inv(cov_param_val) + + x0_true = np.linalg.inv(n * cov_param_inv + 2 * Q_val) @ ( + cov_param_inv @ y_obs.sum(axis=0) + 2 * Q_val @ mu_val + ) + + hess_true = -n * cov_param_inv - Q_val + tau_true = Q_val - hess_true + + log_x_taylor = ( + -0.5 * (x_val - x0_true).T @ tau_true @ (x_val - x0_true) + + 0.5 * np.log(np.linalg.det(tau_true)) + - 0.5 * d * np.log(2 * np.pi) + ) + + np.testing.assert_allclose(x0, x0_true, atol=0.1, rtol=0.1) + np.testing.assert_allclose(log_x_posterior, log_x_taylor, atol=0.1, rtol=0.1) diff --git a/tests/test_laplace.py b/tests/test_laplace.py index 9aafa6e8..8f7a4c01 100644 --- a/tests/test_laplace.py +++ b/tests/test_laplace.py @@ -23,7 +23,6 @@ from pymc_extras.inference.laplace import ( fit_laplace, fit_mvn_at_MAP, - get_conditional_gaussian_approximation, sample_laplace_posterior, ) @@ -280,85 +279,3 @@ def test_laplace_scalar(): assert idata_laplace.fit.covariance_matrix.shape == (1, 1) np.testing.assert_allclose(idata_laplace.fit.mean_vector.values.item(), data.mean(), atol=0.1) - - -def test_get_conditional_gaussian_approximation(): - """ - Consider the trivial case of: - - y | x ~ N(x, cov_param) - x | param ~ N(mu_param, Q^-1) - - cov_param ~ N(cov_mu, cov_cov) - mu_param ~ N(mu_mu, mu_cov) - Q ~ N(Q_mu, Q_cov) - - This has an analytic solution at the mode which we can compare against. - """ - rng = np.random.default_rng(12345) - n = 10000 - d = 10 - - # Initialise arrays - mu_true = rng.random(d) - cov_true = np.diag(rng.random(d)) - Q_val = np.diag(rng.random(d)) - cov_param_val = np.diag(rng.random(d)) - - x_val = rng.random(d) - mu_val = rng.random(d) - - mu_mu = rng.random(d) - mu_cov = np.diag(np.ones(d)) - cov_mu = rng.random(d**2) - cov_cov = np.diag(np.ones(d**2)) - Q_mu = rng.random(d**2) - Q_cov = np.diag(np.ones(d**2)) - - with pm.Model() as model: - y_obs = rng.multivariate_normal(mean=mu_true, cov=cov_true, size=n) - - mu_param = pm.MvNormal("mu_param", mu=mu_mu, cov=mu_cov) - cov_param = pm.MvNormal("cov_param", mu=cov_mu, cov=cov_cov) - Q = pm.MvNormal("Q", mu=Q_mu, cov=Q_cov) - - # Pytensor currently doesn't support autograd for pt inverses, so we use a numeric Q instead - x = pm.MvNormal("x", mu=mu_param, tau=Q) # cov=np.linalg.inv(Q_val)) - - y = pm.MvNormal( - "y", - mu=x, - cov=cov_param.reshape((d, d)), - observed=y_obs, - ) - - # logp(x | y, params) - cga = get_conditional_gaussian_approximation( - x=model.rvs_to_values[x], - Q=Q.reshape((d, d)), - mu=mu_param, - optimizer_kwargs={"tol": 1e-25}, - ) - - x0, log_x_posterior = cga( - x=x_val, mu_param=mu_val, cov_param=cov_param_val.flatten(), Q=Q_val.flatten() - ) - - # Get analytic values of the mode and Laplace-approximated log posterior - cov_param_inv = np.linalg.inv(cov_param_val) - - x0_true = np.linalg.inv(n * cov_param_inv + 2 * Q_val) @ ( - cov_param_inv @ y_obs.sum(axis=0) + 2 * Q_val @ mu_val - ) - - jac_true = cov_param_inv @ (y_obs - x0_true).sum(axis=0) - Q_val @ (x0_true - mu_val) - hess_true = -n * cov_param_inv - Q_val - - log_x_posterior_laplace_true = ( - -0.5 * x_val.T @ (-hess_true + Q_val) @ x_val - + x_val.T @ (Q_val @ mu_val + jac_true - hess_true @ x0_true) - + 0.5 * np.log(np.linalg.det(Q_val)) - ) - - np.testing.assert_allclose(x0, x0_true, atol=0.1, rtol=0.1) - np.testing.assert_allclose(log_x_posterior, log_x_posterior_laplace_true, atol=0.1, rtol=0.1) From 43fb626ed223a3d704b25b4da6dbd58fd33392c9 Mon Sep 17 00:00:00 2001 From: Michal Novomestsky Date: Sun, 6 Jul 2025 21:43:54 +1000 Subject: [PATCH 5/5] changed minimizer tol to 1e-8 --- tests/test_inla.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_inla.py b/tests/test_inla.py index cda24bb5..58ff9ae1 100644 --- a/tests/test_inla.py +++ b/tests/test_inla.py @@ -76,7 +76,7 @@ def test_get_conditional_gaussian_approximation(): x=model.rvs_to_values[x], Q=Q.reshape((d, d)), mu=mu_param, - optimizer_kwargs={"tol": 1e-25}, + optimizer_kwargs={"tol": 1e-8}, ) cga = pytensor.function(args, [x0, pm.logp(x_g, model.rvs_to_values[x])])