Make basic INLA interface and simple marginalisation routine

pymc_extras/inference/INLA/__init__.py

@@ -0,0 +1,3 @@
+from pymc_extras.inference.INLA.inla import fit_INLA
+
+__all__ = ["fit_INLA"]

pymc_extras/inference/INLA/inla.py

@@ -0,0 +1,73 @@
+import arviz as az
+import pymc as pm
+
+from pytensor.tensor import TensorLike, TensorVariable
+
+from pymc_extras.model.marginal.marginal_model import marginalize
+
+
+def fit_INLA(
+    x: TensorVariable,
+    Q: TensorLike,
+    minimizer_seed: int = 42,
+    model: pm.Model | None = None,
+    minimizer_kwargs: dict = {"method": "L-BFGS-B", "optimizer_kwargs": {"tol": 1e-8}},
+    return_latent_posteriors: bool = False,
+    **sampler_kwargs,
+) -> az.InferenceData:
+    r"""
+    Performs inference over a linear mixed model using Integrated Nested Laplace Approximations (INLA). Assumes a model of the form:
+
+    \begin{equation}
+    \theta \rightarrow x \rightarrow y
+    \end{equation}
+
+    Where the prior on the hyperparameters $\pi(\theta)$ is arbitrary, the prior on the latent field is Gaussian (and in precision form): $\pi(x) = N(\mu, Q^{-1})$ and the latent field is linked to the observables $y$ through some linear map.
+
+    As it stands, INLA in PyMC Extras has three main limitations:
+
+    - Does not support inference over the latent field, only the hyperparameters.
+    - Optimisation for $\mu^*$ is bottlenecked by calling `minimize`, and to a lesser extent, computing the hessian $f^"(x)$.
+    - Does not offer sparse support which can provide significant speedups.
+
+    Parameters
+    ----------
+    x: TensorVariable
+        The latent gaussian to marginalize out.
+    Q: TensorLike
+        Precision matrix of the latent field.
+    minimizer_seed: int
+        Seed for random initialisation of the minimum point x*.
+    model: pm.Model
+        PyMC model.
+    minimizer_kwargs:
+        Kwargs to pass to pytensor.optimize.minimize during the optimization step maximizing logp(x | y, params).
+    returned_latent_posteriors:
+        If True, also return posteriors for the latent Gaussian field (currently unsupported).
+    sampler_kwargs:
+        Kwargs to pass to pm.sample.
+    """
+    model = pm.modelcontext(model)
+
+    # TODO is there a better way to check if it's a RV?
+    # print(vars(Q.owner))
+    # if isinstance(Q, TensorVariable) and "module" in vars(Q.owner):
+    Q = model.rvs_to_values[Q] if isinstance(Q, TensorVariable) else Q
+
+    # Marginalize out the latent field
+    marginalize_kwargs = {
+        "Q": Q,
+        "minimizer_seed": minimizer_seed,
+        "minimizer_kwargs": minimizer_kwargs,
+    }
+    marginal_model = marginalize(model, x, use_laplace=True, **marginalize_kwargs)
+
+    # Sample over the hyperparameters
+    if not return_latent_posteriors:
+        idata = pm.sample(model=marginal_model, **sampler_kwargs)
+        return idata
+
+    # Unmarginalize stuff
+    raise NotImplementedError(
+        "Inference over the latent field with INLA is currently unsupported. Set return_latent_posteriors to False"
+    )

pymc_extras/inference/__init__.py

@@ -13,8 +13,9 @@
 #   limitations under the License.
 
 from pymc_extras.inference.fit import fit
+from pymc_extras.inference.INLA.inla import fit_INLA
 from pymc_extras.inference.laplace_approx.find_map import find_MAP
 from pymc_extras.inference.laplace_approx.laplace import fit_laplace
 from pymc_extras.inference.pathfinder.pathfinder import fit_pathfinder
 
-__all__ = ["find_MAP", "fit", "fit_laplace", "fit_pathfinder"]
+__all__ = ["fit", "fit_pathfinder", "fit_laplace", "find_MAP", "fit_INLA"]

pymc_extras/inference/fit.py

@@ -36,7 +36,17 @@ def fit(method: str, **kwargs) -> az.InferenceData:
 
         return fit_pathfinder(**kwargs)
 
-    if method == "laplace":
-        from pymc_extras.inference import fit_laplace
+    elif method == "laplace":
+        from pymc_extras.inference.laplace_approx import fit_laplace
 
         return fit_laplace(**kwargs)
+
+    elif method == "INLA":
+        from pymc_extras.inference.INLA import fit_INLA
+
+        return fit_INLA(**kwargs)
+
+    else:
+        raise ValueError(
+            f"method '{method}' not supported. Use one of 'pathfinder', 'laplace' or 'INLA'."
+        )

pymc_extras/inference/laplace_approx/__init__.py

@@ -0,0 +1,3 @@
+from pymc_extras.inference.laplace_approx.laplace import fit_laplace
+
+__all__ = ["fit_laplace"]

pymc_extras/inference/laplace_approx/laplace.py

@@ -15,27 +15,22 @@
 
 import logging
 
-from collections.abc import Callable
 from functools import partial
 from typing import Literal
 from typing import cast as type_cast
 
 import arviz as az
 import numpy as np
 import pymc as pm
-import pytensor
 import pytensor.tensor as pt
 import xarray as xr
 
 from better_optimize.constants import minimize_method
-from numpy.typing import ArrayLike
 from pymc.blocking import DictToArrayBijection
 from pymc.model.transform.optimization import freeze_dims_and_data
 from pymc.pytensorf import join_nonshared_inputs
 from pymc.util import get_default_varnames
 from pytensor.graph import vectorize_graph
-from pytensor.tensor import TensorVariable
-from pytensor.tensor.optimize import minimize
 from pytensor.tensor.type import Variable
 
 from pymc_extras.inference.laplace_approx.find_map import (
@@ -51,102 +46,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)
-
-    # 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})
-    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
-    )
-
-    # 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])
-
-
 def _unconstrained_vector_to_constrained_rvs(model):
     outputs = get_default_varnames(model.unobserved_value_vars, include_transformed=True)
     constrained_names = [