Implement basic INLA interface (#533)

Author: Michal-Novomestsky

conda-envs/environment-test.yml

@@ -19,3 +19,5 @@ dependencies:
   - pip:
       - jax
       - blackjax
+      - pytensor>=3.0.4
+      - preliz>=0.25

notebooks/INLA Example.ipynb

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

pymc_extras/inference/INLA/__init__.py

@@ -0,0 +1,110 @@
+import warnings
+
+import pymc as pm
+
+from pytensor.tensor import TensorLike, TensorVariable, as_tensor
+from xarray import DataTree
+
+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,
+) -> DataTree:
+    r"""
+    Performs inference over a linear mixed model using Integrated Nested Laplace Approximations (INLA). Assumes a model of the form:
+
+    .. math::
+
+        \theta \rightarrow x \rightarrow y
+
+    Where the prior on the hyperparameters :math:`\pi(\theta)` is arbitrary, the prior on the latent field is Gaussian (and in precision form): :math:`\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 is currently experimental.
+
+    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.
+
+    Returns
+    -------
+    DataTree
+        The inference data containing the results of the INLA algorithm.
+
+    Examples
+    --------
+    .. code:: ipython
+
+        In [1]: rng = np.random.default_rng(123)
+           ...: n = 10000
+           ...: d = 3
+           ...: mu_mu = 10 * rng.random(d)
+           ...: mu_true = rng.random(d)
+           ...: tau = np.identity(d)
+           ...: cov = np.linalg.inv(tau)
+           ...: y_obs = rng.multivariate_normal(mean=mu_true, cov=cov, size=n)
+
+        In [2]: with pm.Model() as model:
+           ...:     mu = pm.MvNormal("mu", mu=mu_mu, tau=tau)
+           ...:     x = pm.MvNormal("x", mu=mu, tau=tau)
+           ...:     y = pm.MvNormal("y", mu=x, tau=tau, observed=y_obs)
+
+           ...:     idata = pmx.fit(
+           ...:     method="INLA",
+           ...:     x=x,
+           ...:     Q=tau,
+           ...:     return_latent_posteriors=False,
+           ...:     )
+
+        In[3]: posterior_mean_true = (mu_mu + mu_true) / 2
+           ...: posterior_mean_inla = idata.posterior.mu.mean(axis=(0, 1)).values
+           ...: print(posterior_mean_true)
+           ...: print(posterior_mean_inla)
+
+        Out[3]:
+            [3.50394522 0.35705804 1.50784662]
+            [3.48732847 0.35738072 1.46851421]
+
+    """
+    warnings.warn(
+        "INLA is currently experimental. Please see the INLA Roadmap for more info: https://github.com/pymc-devs/pymc-extras/issues/340.",
+        UserWarning,
+    )
+    model = pm.modelcontext(model)
+
+    # Marginalize out the latent field
+    marginalize_kwargs = {
+        "Q": as_tensor(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/INLA/inla.py

@@ -14,6 +14,7 @@
 
 from pymc_extras.inference.dadvi.dadvi import fit_dadvi
 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_blackjax_pathfinder, fit_pathfinder
@@ -25,4 +26,5 @@
     "fit_laplace",
     "fit_pathfinder",
     "fit_dadvi",
+    "fit_INLA",
 ]

pymc_extras/inference/__init__.py

@@ -36,12 +36,22 @@ def fit(method: str, **kwargs) -> DataTree:
 
         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)
 
-    if method == "dadvi":
+    elif method == "INLA":
+        from pymc_extras.inference.INLA import fit_INLA
+
+        return fit_INLA(**kwargs)
+
+    elif method == "dadvi":
         from pymc_extras.inference import fit_dadvi
 
         return fit_dadvi(**kwargs)
+
+    else:
+        raise ValueError(
+            f"method '{method}' not supported. Use one of 'pathfinder', 'laplace' or 'INLA'."
+        )

pymc_extras/inference/fit.py

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

pymc_extras/inference/laplace_approx/__init__.py

@@ -3,10 +3,12 @@
 from collections.abc import Sequence
 
 import numpy as np
+import pytensor
 import pytensor.tensor as pt
 
 from pymc.distributions import Bernoulli, Categorical, DiscreteUniform
 from pymc.distributions.distribution import _support_point, support_point
+from pymc.distributions.multivariate import _logdet_from_cholesky
 from pymc.logprob.abstract import MeasurableOp, _logprob
 from pymc.logprob.basic import conditional_logp, logp
 from pymc.pytensorf import constant_fold
@@ -17,7 +19,8 @@
 from pytensor.graph.replace import clone_replace, graph_replace
 from pytensor.scan import map as scan_map
 from pytensor.scan import scan
-from pytensor.tensor import TensorVariable
+from pytensor.tensor import TensorLike, TensorVariable
+from pytensor.tensor.optimize import minimize
 from pytensor.tensor.random.type import RandomType
 
 from pymc_extras.distributions import DiscreteMarkovChain
@@ -134,6 +137,24 @@ class MarginalDiscreteMarkovChainRV(MarginalRV):
     """Base class for Marginalized Discrete Markov Chain RVs"""
 
 
+class MarginalLaplaceRV(MarginalRV):
+    """Base class for Marginalized Laplace-Approximated RVs.
+
+    Estimates log likelihood using Laplace approximations.
+    """
+
+    def __init__(
+        self,
+        *args,
+        minimizer_seed: int,
+        minimizer_kwargs: dict = {"method": "L-BFGS-B", "optimizer_kwargs": {"tol": 1e-8}},
+        **kwargs,
+    ) -> None:
+        self.minimizer_seed = minimizer_seed
+        self.minimizer_kwargs = minimizer_kwargs
+        super().__init__(*args, **kwargs)
+
+
 def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> tuple[int, ...]:
     op = rv.owner.op
     dist_params = rv.owner.op.dist_params(rv.owner)
@@ -377,3 +398,137 @@ def step_alpha(logp_emission, log_alpha, log_P):
     warn_non_separable_logp(values)
     dummy_logps = (DUMMY_ZERO,) * (len(values) - 1)
     return joint_logp, *dummy_logps
+
+
+def _precision_mv_normal_logp(value: TensorLike, mean: TensorLike, tau: TensorLike):
+    """
+    Compute the log likelihood of a multivariate normal distribution in precision form. May be phased out - see https://github.com/pymc-devs/pymc/pull/7895
+
+    Parameters
+    ----------
+    value: TensorLike
+        Query point to compute the log prob at.
+    mean: TensorLike
+        Mean vector of the Gaussian,
+    tau: TensorLike
+        Precision matrix of the Gaussian (i.e. cov = inv(tau))
+
+    Returns
+    -------
+    logp: TensorLike
+        Log likelihood at value.
+    posdef: TensorLike
+        Boolean indicating whether the precision matrix is positive definite.
+    """
+    k = value.shape[-1].astype("floatX")
+
+    delta = value - mean
+    quadratic_form = delta.T @ tau @ delta
+    logdet, posdef = _logdet_from_cholesky(pt.linalg.cholesky(tau, lower=True))
+    logp = -0.5 * (k * pt.log(2 * np.pi) + quadratic_form) + logdet
+
+    return logp, posdef
+
+
+def get_laplace_approx(
+    log_likelihood: TensorVariable,
+    logp_objective: TensorVariable,
+    x: TensorVariable,
+    x0_init: TensorLike,
+    Q: TensorLike,
+    minimizer_kwargs: dict = {"method": "L-BFGS-B", "optimizer_kwargs": {"tol": 1e-8}},
+):
+    """
+    Compute the laplace approximation logp_G(x | y, params) of some variable x.
+
+    Parameters
+    ----------
+    log_likelihood: TensorVariable
+        Model likelihood logp(y | x, params).
+    logp_objective: TensorVariable
+        Obective log likelihood to maximize, logp(x | y, params) (up to some constant in x).
+    x: TensorVariable
+        Variable to be laplace approximated.
+    x0_init: TensorLike
+        Initial guess for minimization.
+    Q: TensorLike
+        Precision matrix of x.
+    minimizer_kwargs:
+        Kwargs to pass to pytensor.optimize.minimize.
+
+    Returns
+    -------
+    x0: TensorVariable
+        x*, the maximizer of logp(x | y, params) in x.
+    log_laplace_approx: TensorVariable
+        Laplace approximation of logp(x | y, params) evaluated at x.
+    """
+    # Maximize log(p(x | y, params)) wrt x to find mode x0
+    # This step is currently bottlenecking the logp calculation.
+    x0, _ = minimize(
+        objective=-logp_objective,  # logp(x | y, params) = logp(y | x, params) + logp(x | params) + const (const omitted during minimization)
+        x=x,
+        use_vectorized_jac=True,
+        **minimizer_kwargs,
+    )
+
+    # Set minimizer initialisation to be random
+    x0 = pytensor.graph.replace.graph_replace(x0, {x: x0_init})
+
+    # This step is also expensive (but not as much as minimize). Could be made more efficient by recycling hessian from the minimizer step, however that requires a bespoke algorithm described in Rasmussen & Williams
+    # since the general optimisation scheme maximises logp(x | y, params) rather than logp(y | x, params), and thus the hessian that comes out of methods
+    # like L-BFGS-B is in fact not the hessian of logp(y | x, params)
+    # TODO: Use vectorized hessian?
+    hess = pytensor.gradient.hessian(log_likelihood, x)
+
+    # Evaluate logp of Laplace approx of logp(x | y, params) at some point x
+    tau = Q - hess
+    mu = x0
+    log_laplace_approx, _ = _precision_mv_normal_logp(x, mu, tau)
+
+    return x0, log_laplace_approx
+
+
+@_logprob.register(MarginalLaplaceRV)
+def laplace_marginal_rv_logp(op: MarginalLaplaceRV, values, *inputs_and_Q, **kwargs):
+    # Get Q and remove it from the graph (stored as a dummy input)
+    *inputs, Q = inputs_and_Q
+
+    # Clone the inner RV graph of the Marginalized RV
+    x, *inner_rvs = inline_ofg_outputs(op, inputs)
+
+    # Obtain the joint_logp graph of the inner RV graph
+    inner_rv_values = dict(zip(inner_rvs, values))
+
+    marginalized_vv = x.clone()
+    rv_values = inner_rv_values | {x: marginalized_vv}
+    logps_dict = conditional_logp(rv_values=rv_values, **kwargs)
+
+    # logp(x | params)
+    logp_x = logps_dict.pop(marginalized_vv).sum()
+
+    # logp(y | x, params)
+    logp_y = pt.sum([logp_term.sum() for value, logp_term in logps_dict.items()])
+
+    # logp_total = logp(y | x, params) + logp(x | params) (i.e. logp(x | y, params) up to a constant in x)
+    logp_total = logp_x + logp_y
+
+    # Set minimizer initialisation to be random (TODO: Let pymc accept this one, maybe when rng is constant)
+    # TODO: Use newer pytensor helper
+    d = pt.prod(constant_fold(tuple(x.shape), raise_not_constant=True))
+    x0_init = pt.ones(d)
+
+    # Obtain laplace approx for logp(x | y, params)
+    x0, log_laplace_approx = get_laplace_approx(
+        logp_y,
+        logp_total,
+        x=marginalized_vv,
+        x0_init=x0_init,
+        Q=Q,
+        minimizer_kwargs=op.minimizer_kwargs,
+    )
+
+    # logp(y | params) = logp(y | x, params) + logp(x | params) - logp(x | y, params)
+    # TODO: Can we recover the elementwise logp?
+    marginal_likelihood = logp_total - log_laplace_approx
+    return graph_replace(marginal_likelihood, {marginalized_vv: x0})