end-to-end implementation

Author: Michal-Novomestsky

notebooks/INLA_testing.ipynb

@@ -6,7 +6,9 @@
 
 from better_optimize.constants import minimize_method
 from numpy.typing import ArrayLike
+from pymc.distributions.multivariate import MvNormal
 from pytensor.tensor import TensorVariable
+from pytensor.tensor.linalg import inv as matrix_inverse
 from pytensor.tensor.optimize import minimize
 
 from pymc_extras.model.marginal.marginal_model import marginalize
@@ -67,6 +69,8 @@ def get_conditional_gaussian_approximation(
     x0, p(x | y, params): list[TensorVariable]
         Mode and Laplace approximation for posterior.
     """
+    raise DeprecationWarning("Legacy code. Please use fit_INLA instead.")
+
     model = pm.modelcontext(model)
 
     # f = log(p(y | x, params))
@@ -108,6 +112,8 @@ def get_log_marginal_likelihood(
     use_hess: bool = False,
     optimizer_kwargs: dict | None = None,
 ) -> TensorVariable:
+    raise DeprecationWarning("Legacy code. Please use fit_INLA instead.")
+
     model = pm.modelcontext(model)
 
     x0, log_laplace_approx = get_conditional_gaussian_approximation(
@@ -134,43 +140,40 @@ def get_log_marginal_likelihood(
 
 def fit_INLA(
     x: TensorVariable,
-    Q: TensorVariable | ArrayLike,
-    # mu: TensorVariable | ArrayLike,
+    temp_kwargs=None,  # TODO REMOVE. DEBUGGING TOOL
     model: pm.Model | None = None,
     minimizer_kwargs: dict | None = None,
+    return_latent_posteriors: bool = True,
     **sampler_kwargs,
 ) -> az.InferenceData:
     model = pm.modelcontext(model)
 
-    # Marginalize out the latent field
-    marginalize(model, [x], Q, minimizer_kwargs, method="INLA")
+    # Check if latent field is Gaussian
+    if not isinstance(x.owner.op, MvNormal):
+        raise ValueError(
+            f"Latent field {x} is not instance of MvNormal. Has distribution {x.owner.op}."
+        )
 
-    # Sample over the hyperparameters
-    pm.sample(model=model, **sampler_kwargs)
+    _, _, _, tau = x.owner.inputs
 
-    # # logp(y | params)
-    # x0, log_likelihood = get_log_marginal_likelihood(
-    #     x, Q, mu, model, method, use_jac, use_hess, optimizer_kwargs
-    # )
+    # Latent field should use precison rather than covariance
+    if not tau.owner or tau.owner.op != matrix_inverse:
+        raise ValueError(
+            f"Latent field {x} is not in precision matrix form. Use MvNormal(tau=Q) instead."
+        )
 
-    # # 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
-    # # Perhaps obtain as RVs which y depends on which aren't x?
-    # prior = None
-    # params = None
-    # log_prior = pm.logp(prior, model.rvs_to_values[params])
+    Q = tau.owner.inputs[0]
 
-    # # logp(params | y) = logp(y | params) + logp(params) + const
-    # log_posterior = log_likelihood + log_prior
-    # log_posterior = pytensor.graph.replace.graph_replace(log_posterior, {x: x0})
+    # Marginalize out the latent field
+    minimizer_kwargs = {"method": "L-BFGS-B", "optimizer_kwargs": {"tol": 1e-8}}
+    marginalize_kwargs = {"Q": Q, "temp_kwargs": temp_kwargs, "minimizer_kwargs": minimizer_kwargs}
+    marginal_model = marginalize(model, x, use_laplace=True, **marginalize_kwargs)
 
-    # # 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
+    # Sample over the hyperparameters
+    idata = pm.sample(model=marginal_model, **sampler_kwargs)
 
-    # # 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()
+    if not return_latent_posteriors:
+        return idata
 
-    # Bundle up idatas somehow
-    # return idata_params, idata_x
+    # TODO Unmarginalize stuff
+    raise NotImplementedError("Latent posteriors not supported yet, WIP.")

pymc_extras/inference/inla.py

@@ -8,6 +8,7 @@
 
 from pymc.distributions import Bernoulli, Categorical, DiscreteUniform
 from pymc.distributions.distribution import _support_point, support_point
+from pymc.distributions.multivariate import _precision_mv_normal_logp
 from pymc.logprob.abstract import MeasurableOp, _logprob
 from pymc.logprob.basic import conditional_logp, logp
 from pymc.pytensorf import constant_fold
@@ -414,7 +415,7 @@ def laplace_marginal_rv_logp(op: MarginalLaplaceRV, values, *inputs, **kwargs):
     minimizer_kwargs = (
         op.minimizer_kwargs
         if op.minimizer_kwargs is not None
-        else {"method": "BFGS", "optimizer_kwargs": {"tol": 1e-8}}
+        else {"method": "L-BFGS-B", "optimizer_kwargs": {"tol": 1e-8}}
     )
 
     x0, _ = minimize(
@@ -423,23 +424,23 @@ def laplace_marginal_rv_logp(op: MarginalLaplaceRV, values, *inputs, **kwargs):
         **minimizer_kwargs,
     )
 
-    # # Set minimizer initialisation to be random
+    # Set minimizer initialisation to be random
     d = 3  # 10000 # TODO pull this from x.shape (or similar) somehow
     rng = np.random.default_rng(12345)
     x0 = pytensor.graph.replace.graph_replace(x0, {marginalized_vv: rng.random(d)})
 
     # TODO USE CLOSED FORM SOLUTION FOR NOW
-    # n, y_obs = op.temp_kwargs
-    # mu_param = pytensor.graph.basic.get_var_by_name(x, "mu_param")[0]
-    # x0 = (y_obs.sum(axis=0) - mu_param) / (n - 1)
+    n, y_obs = op.temp_kwargs
+    mu_param = pytensor.graph.basic.get_var_by_name(x, "mu")[0]
+    x0 = (y_obs.sum(axis=0) - mu_param) / (n - 1)
 
     # logp(x | y, params) using laplace approx evaluated at x0
-    hess = pytensor.gradient.hessian(log_likelihood, marginalized_vv)
+    hess = pytensor.gradient.hessian(
+        log_likelihood, marginalized_vv
+    )  # TODO check how stan makes this quicker
     tau = op.Q - hess
-    _, logdetTau = pt.nlinalg.slogdet(tau)
-    log_laplace_approx = 0.5 * logdetTau - 0.5 * marginalized_vv.shape[0] * np.log(
-        2 * np.pi
-    )  # At x = x0, the quadratic term becomes 0
+    mu = x0  # TODO double check with Theo
+    log_laplace_approx, _ = _precision_mv_normal_logp(x0, mu, tau)
 
     # logp(y | params) = logp(y | x, params) + logp(x | params) - logp(x | y, params)
     marginal_likelihood = logp - log_laplace_approx

pymc_extras/model/marginal/distributions.py

@@ -9,6 +9,7 @@
 from arviz import InferenceData, dict_to_dataset
 from pymc.backends.arviz import coords_and_dims_for_inferencedata, dataset_to_point_list
 from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
+from pymc.distributions.multivariate import MvNormal
 from pymc.distributions.transforms import Chain
 from pymc.logprob.transforms import IntervalTransform
 from pymc.model import Model
@@ -45,6 +46,7 @@
 from pymc_extras.model.marginal.distributions import (
     MarginalDiscreteMarkovChainRV,
     MarginalFiniteDiscreteRV,
+    MarginalLaplaceRV,
     MarginalRV,
     NonSeparableLogpWarning,
     get_domain_of_finite_discrete_rv,
@@ -144,7 +146,9 @@ def _unique(seq: Sequence) -> list:
     return [x for x in seq if not (x in seen or seen_add(x))]
 
 
-def marginalize(model: Model, rvs_to_marginalize: ModelRVs) -> MarginalModel:
+def marginalize(
+    model: Model, rvs_to_marginalize: ModelRVs, use_laplace: bool = False, **marginalize_kwargs
+) -> MarginalModel:
     """Marginalize a subset of variables in a PyMC model.
 
     This creates a class of `MarginalModel` from an existing `Model`, with the specified
@@ -158,6 +162,8 @@ def marginalize(model: Model, rvs_to_marginalize: ModelRVs) -> MarginalModel:
         PyMC model to marginalize. Original variables well be cloned.
     rvs_to_marginalize : Sequence[TensorVariable]
         Variables to marginalize in the returned model.
+    use_laplace : bool
+        Whether to use Laplace appoximations to marginalize out rvs_to_marginalize.
 
     Returns
     -------
@@ -186,7 +192,12 @@ def marginalize(model: Model, rvs_to_marginalize: ModelRVs) -> MarginalModel:
                 raise NotImplementedError(
                     "Marginalization for DiscreteMarkovChain with non-matrix transition probability is not supported"
                 )
-        elif not isinstance(rv_op, Bernoulli | Categorical | DiscreteUniform):
+        elif use_laplace and not isinstance(rv_op, MvNormal):
+            raise ValueError(
+                f"Marginalisation method set to Laplace but RV {rv_to_marginalize} is not instance of MvNormal. Has distribution {rv_to_marginalize.owner.op}"
+            )
+
+        elif not use_laplace and not isinstance(rv_op, Bernoulli | Categorical | DiscreteUniform):
             raise NotImplementedError(
                 f"Marginalization of RV with distribution {rv_to_marginalize.owner.op} is not supported"
             )
@@ -241,7 +252,9 @@ def marginalize(model: Model, rvs_to_marginalize: ModelRVs) -> MarginalModel:
         ]
         input_rvs = _unique((*marginalized_rv_input_rvs, *other_direct_rv_ancestors))
 
-        replace_finite_discrete_marginal_subgraph(fg, rv_to_marginalize, dependent_rvs, input_rvs)
+        replace_marginal_subgraph(
+            fg, rv_to_marginalize, dependent_rvs, input_rvs, use_laplace, **marginalize_kwargs
+        )
 
     return model_from_fgraph(fg, mutate_fgraph=True)
 
@@ -551,22 +564,32 @@ def remove_model_vars(vars):
     return fgraph.outputs
 
 
-def replace_finite_discrete_marginal_subgraph(
-    fgraph, rv_to_marginalize, dependent_rvs, input_rvs
+def replace_marginal_subgraph(
+    fgraph,
+    rv_to_marginalize,
+    dependent_rvs,
+    input_rvs,
+    use_laplace=False,
+    **marginalize_kwargs,
 ) -> None:
     # If the marginalized RV has multiple dimensions, check that graph between
     # marginalized RV and dependent RVs does not mix information from batch dimensions
     # (otherwise logp would require enumerating over all combinations of batch dimension values)
-    try:
-        dependent_rvs_dim_connections = subgraph_batch_dim_connection(
-            rv_to_marginalize, dependent_rvs
-        )
-    except (ValueError, NotImplementedError) as e:
-        # For the perspective of the user this is a NotImplementedError
-        raise NotImplementedError(
-            "The graph between the marginalized and dependent RVs cannot be marginalized efficiently. "
-            "You can try splitting the marginalized RV into separate components and marginalizing them separately."
-        ) from e
+    if not use_laplace:
+        try:
+            dependent_rvs_dim_connections = subgraph_batch_dim_connection(
+                rv_to_marginalize, dependent_rvs
+            )
+        except (ValueError, NotImplementedError) as e:
+            # For the perspective of the user this is a NotImplementedError
+            raise NotImplementedError(
+                "The graph between the marginalized and dependent RVs cannot be marginalized efficiently. "
+                "You can try splitting the marginalized RV into separate components and marginalizing them separately."
+            ) from e
+    else:
+        dependent_rvs_dim_connections = [
+            (None,),
+        ]
 
     output_rvs = [rv_to_marginalize, *dependent_rvs]
     rng_updates = collect_default_updates(output_rvs, inputs=input_rvs, must_be_shared=False)
@@ -581,6 +604,8 @@ def replace_finite_discrete_marginal_subgraph(
 
     if isinstance(inner_outputs[0].owner.op, DiscreteMarkovChain):
         marginalize_constructor = MarginalDiscreteMarkovChainRV
+    elif use_laplace:
+        marginalize_constructor = MarginalLaplaceRV
     else:
         marginalize_constructor = MarginalFiniteDiscreteRV
 
@@ -590,6 +615,7 @@ def replace_finite_discrete_marginal_subgraph(
         outputs=inner_outputs,
         dims_connections=dependent_rvs_dim_connections,
         dims=dims,
+        **marginalize_kwargs,
     )
 
     new_outputs = marginalization_op(*inputs)