Implement a minimizer for INLA

pymc_extras/inference/laplace.py

@@ -39,6 +39,8 @@
 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 TensorLike, TensorVariable
+from pytensor.tensor.optimize import minimize
 from scipy import stats
 
 from pymc_extras.inference.find_map import (
@@ -415,6 +417,128 @@ def sample_laplace_posterior(
     return idata
 
 
+def find_mode(
+    x: TensorVariable,
+    args: dict,
+    x0: TensorLike | None = None,
+    model: pm.Model | None = None,
+    method: minimize_method = "BFGS",
+    use_jac: bool = True,
+    use_hess: bool = False,  # TODO Tbh we can probably just remove this arg and pass True to the minimizer all the time, but if this is the case, it will throw a warning when the hessian doesn't need to be computed for a particular optimisation routine.
+    optimizer_kwargs: dict | None = None,
+) -> list[TensorLike]:
+    """
+    Estimates the mode and hessian of a model by minimizing negative log likelihood. Wrapper for (pytensor-native) scipy.optimize.minimize.
+
+    Parameters
+    ----------
+    x: TensorVariable
+        The parameter with which to minimize wrt (that is, find the mode in x).
+    args: dict
+        A dictionary of the form {tensorvariable_name: TensorLike}, where tensorvariable_name is the (exact) name of TensorVariable which is to be provided some numerical value. Same usage as args in scipy.optimize.minimize.
+    x0: TensorLike
+        Initial guess for the mode (in x). Initialised over a uniform distribution if unspecified.
+    model: Model
+        PyMC model to use.
+    method: minimize_method
+        Which minimization algorithm to use.
+    use_jac: bool
+        If true, the minimizer will compute and store the Jacobian.
+    use_hess: bool
+        If true, the minimizer will compute and store the Hessian (note that the Hessian will be computed explicitely even if this is False).
+    optimizer_kwargs: dict
+        Kwargs to pass to scipy.optimize.minimize.
+
+    Returns
+    -------
+    mu: TensorLike
+        The mode of the model.
+    hess:
+        Hessian evalulated at mu.
+    """
+    model = pm.modelcontext(model)
+
+    # # TODO I would like to generate a random initialisation for x0 if set to None. Ideally this would be done using something like np.random.rand(x.shape), however I don't believe x.shape is
+    # # immediately accessible in pytensor. model.initial_point() throws the following error:
+
+    # # MissingInputError: Input 0 (X) of the graph (indices start from 0), used to compute Transpose{axes=[1, 0]}(X), was not provided and not given a value. Use the PyTensor flag exception_verbosity='high', for more information on this error.
+
+    # # Instead I've tried to follow what find_MAP does (below), but this doesn't really get me anywhere unfortunately.
+    # # if x0 is None:
+    # # Yes ik this is here, just for debugging purposes
+    # from pymc.initial_point import make_initial_point_fn
+
+    # frozen_model = freeze_dims_and_data(model)
+    # ipfn = make_initial_point_fn(
+    #     model=frozen_model,
+    #     jitter_rvs=set(),  # (jitter_rvs),
+    #     return_transformed=True,
+    #     overrides={x.name: x0}, # x0 is here for debugging purposes
+    # )
+
+    # random_seed = None
+    # start_dict = ipfn(random_seed)
+    # vars_dict = {var.name: var for var in frozen_model.continuous_value_vars}
+    # initial_params = DictToArrayBijection.map(
+    #     {var_name: value for var_name, value in start_dict.items() if var_name in vars_dict}
+    # )
+    # # Printing this should return {'name of x TensorVariable': initialised values for x}
+    # print(initial_params)
+
+    # Minimise negative log likelihood
+    nll = -model.logp()
+    soln, _ = minimize(
+        objective=nll,
+        x=x,
+        method=method,
+        jac=use_jac,
+        hess=use_hess,
+        optimizer_kwargs=optimizer_kwargs,
+    )
+
+    # TODO To prevent needing to user to pass in the TensorVariables alongside their names (e.g. args = {'X': [X, [0,0]], 'beta': [beta_val, [1]], ...}) the codeblock below digs up the
+    # TensorVariables associated with the names listed in args from the graph. The sensible graph to use here would be nll, because that is what is going into the minimize function, however doing
+    # so results in the following error:
+    #
+    # TypeError: TensorType does not support iteration.
+    # Did you pass a PyTensor variable to a function that expects a list?
+    # Maybe you are using builtins.sum instead of pytensor.tensor.sum?
+    #
+    # Using model.basic_RVs[1] instead works, but note that this is a hardcoded fix because the model I'm testing on happens to have all of the relevant TensorVariables in the graph of model.basic_RVs[1],
+    # but this isn't true in general.
+
+    # Get arg TensorVariables
+    arg_tensorvars = [
+        pytensor.graph.basic.get_var_by_name(model.basic_RVs[1], target_var_id=var)[0]
+        # pytensor.graph.basic.get_var_by_name(nll, target_var_id=var)[0]
+        for var in args
+    ]
+    for i, var in enumerate(arg_tensorvars):
+        try:
+            arg_tensorvars[i] = model.rvs_to_values[var]
+        except KeyError:
+            pass
+    arg_tensorvars.insert(0, x)
+
+    # TODO: Jesse suggested I use this graph_replace function, but it seems that "mode" here is a different type to soln:
+    #
+    # TypeError: Cannot convert Type Vector(float64, shape=(10,)) (of Variable MinimizeOp(method=BFGS, jac=True, hess=True, hessp=False).0) into Type Scalar(float64, shape=()). You can try to manually convert MinimizeOp(method=BFGS, jac=True, hess=True, hessp=False).0 into a Scalar(float64, shape=()).
+    #
+    # My understanding here is that for some function which evaluates the hessian at x, we're replacing "x" in the hess graph with the subgraph that computes "x" (i.e. soln)?
+
+    # Obtain the Hessian (re-use graph if already computed in minimize)
+    if use_hess:
+        mode, _, hess = (
+            soln.owner.op.inner_outputs
+        )  # Note that this mode, _, hess will need to be slightly more elaborate for when use_jac is False (2 items to unpack instead of 3). Just a few if-blocks, but not implemented for now while we're debugging
+        hess = pytensor.graph.replace.graph_replace(hess, {mode: soln})
+    else:
+        hess = pytensor.gradient.hessian(nll, x)
+
+    get_mode_and_hessian = pytensor.function(arg_tensorvars, [soln, hess])
+    return get_mode_and_hessian(x0, **args)
+
+
 def fit_laplace(
     optimize_method: minimize_method | Literal["basinhopping"] = "BFGS",
     *,

tests/test_laplace.py

@@ -15,12 +15,14 @@
 
 import numpy as np
 import pymc as pm
+import pytensor as pt
 import pytest
 
 import pymc_extras as pmx
 
 from pymc_extras.inference.find_map import GradientBackend, find_MAP
 from pymc_extras.inference.laplace import (
+    find_mode,
     fit_laplace,
     fit_mvn_at_MAP,
     sample_laplace_posterior,
@@ -279,3 +281,160 @@ 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_find_mode():
+    k = 10
+    N = 10000
+    y = pt.vector("y", dtype="int64")
+    X = pt.matrix("X", shape=(N, k))
+
+    # TODO Pre-commit formatted it like this. Quite ugly. Should compute hess in code rather than storing a hardcoded array.
+    true_hess = np.array(
+        [
+            [
+                2.50100000e03,
+                -1.78838742e00,
+                1.59484217e01,
+                -9.78343803e00,
+                2.86125467e01,
+                -7.38071788e00,
+                -4.97729126e01,
+                3.53243810e01,
+                1.69071769e01,
+                -1.30755942e01,
+            ],
+            [
+                -1.78838742e00,
+                2.54687995e03,
+                8.99456512e-02,
+                -1.33603390e01,
+                -2.37641179e01,
+                4.57780742e01,
+                -1.22640681e01,
+                2.70879664e01,
+                4.04435512e01,
+                2.08826556e00,
+            ],
+            [
+                1.59484217e01,
+                8.99456512e-02,
+                2.46908384e03,
+                -1.80358232e01,
+                1.14131535e01,
+                2.21632317e01,
+                1.25443469e00,
+                1.50344618e01,
+                -3.59940488e01,
+                -1.05191328e01,
+            ],
+            [
+                -9.78343803e00,
+                -1.33603390e01,
+                -1.80358232e01,
+                2.50546496e03,
+                3.27545028e01,
+                -3.33517501e01,
+                -2.68735672e01,
+                -2.69114305e01,
+                -1.20464337e01,
+                9.02338622e00,
+            ],
+            [
+                2.86125467e01,
+                -2.37641179e01,
+                1.14131535e01,
+                3.27545028e01,
+                2.49959736e03,
+                -3.98220135e00,
+                -4.09495199e00,
+                -1.51115257e01,
+                -5.77436126e01,
+                -2.98600447e00,
+            ],
+            [
+                -7.38071788e00,
+                4.57780742e01,
+                2.21632317e01,
+                -3.33517501e01,
+                -3.98220135e00,
+                2.48169432e03,
+                -1.26885014e01,
+                -3.53524089e01,
+                5.89656794e00,
+                1.67164400e01,
+            ],
+            [
+                -4.97729126e01,
+                -1.22640681e01,
+                1.25443469e00,
+                -2.68735672e01,
+                -4.09495199e00,
+                -1.26885014e01,
+                2.47216241e03,
+                8.16935659e00,
+                -4.89399152e01,
+                -1.11646138e01,
+            ],
+            [
+                3.53243810e01,
+                2.70879664e01,
+                1.50344618e01,
+                -2.69114305e01,
+                -1.51115257e01,
+                -3.53524089e01,
+                8.16935659e00,
+                2.52940405e03,
+                3.07751540e00,
+                -8.60023392e00,
+            ],
+            [
+                1.69071769e01,
+                4.04435512e01,
+                -3.59940488e01,
+                -1.20464337e01,
+                -5.77436126e01,
+                5.89656794e00,
+                -4.89399152e01,
+                3.07751540e00,
+                2.49452594e03,
+                6.06984410e01,
+            ],
+            [
+                -1.30755942e01,
+                2.08826556e00,
+                -1.05191328e01,
+                9.02338622e00,
+                -2.98600447e00,
+                1.67164400e01,
+                -1.11646138e01,
+                -8.60023392e00,
+                6.06984410e01,
+                2.49290175e03,
+            ],
+        ]
+    )
+
+    with pm.Model() as model:
+        beta = pm.MvNormal("beta", mu=np.zeros(k), cov=np.identity(k), shape=(k,))
+        p = pm.math.invlogit(beta @ X.T)
+        y = pm.Bernoulli("y", p)
+
+        rng = np.random.default_rng(123)
+        Xval = rng.normal(size=(10000, 9))
+        Xval = np.c_[np.ones(10000), Xval]
+
+        true_beta = rng.normal(scale=0.1, size=(10,))
+        true_p = pm.math.invlogit(Xval @ true_beta).eval()
+        ynum = rng.binomial(1, true_p)
+
+        beta_val = model.rvs_to_values[beta]
+        x0 = np.zeros(k)
+        args = {"y": ynum, "X": Xval}
+
+        beta_mode, beta_hess = find_mode(
+            x=beta_val, x0=x0, args=args, method="BFGS", optimizer_kwargs={"tol": 1e-8}
+        )
+
+    np.testing.assert_allclose(beta_mode, true_beta, atol=0.1, rtol=0.1)
+    np.testing.assert_allclose(beta_hess, true_hess, atol=0.1, rtol=0.1)