removed temp_kwargs, made Q amenable to RVs, removed dependency on MvNormal

Author: Michal-Novomestsky

pymc_extras/inference/INLA/inla.py

@@ -1,49 +1,54 @@
-import warnings
-
 import arviz as az
 import pymc as pm
 
-from pymc.distributions.multivariate import MvNormal
 from pytensor.tensor import TensorVariable
-from pytensor.tensor.linalg import inv as matrix_inverse
 
 from pymc_extras.model.marginal.marginal_model import marginalize
 
 
 def fit_INLA(
     x: TensorVariable,
-    temp_kwargs=None,  # TODO REMOVE. DEBUGGING TOOL
+    Q: TensorVariable,
+    minimizer_seed: int = 42,
     model: pm.Model | None = None,
     minimizer_kwargs: dict | None = None,
     return_latent_posteriors: bool = True,
     **sampler_kwargs,
 ) -> az.InferenceData:
-    warnings.warn("Currently only valid for a nested normal model. WIP.", UserWarning)
-
     model = pm.modelcontext(model)
 
     # 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}."
-        )
+    # if not isinstance(x.owner.op, MvNormal):
+    #     raise ValueError(
+    #         f"Latent field {x} is not instance of MvNormal. Has distribution {x.owner.op}."
+    #     )
+
+    # _, _, _, tau = x.owner.inputs
 
-    _, _, _, tau = x.owner.inputs
+    # # Latent field should use precison rather than covariance
+    # if not (tau.owner and tau.owner.op == matrix_inverse):
+    #     raise ValueError(
+    #         f"Latent field {x} is not in precision matrix form. Use MvNormal(tau=Q) instead."
+    #     )
 
-    # Latent field should use precison rather than covariance
-    if not (tau.owner and tau.owner.op == matrix_inverse):
-        raise ValueError(
-            f"Latent field {x} is not in precision matrix form. Use MvNormal(tau=Q) instead."
-        )
+    # Q = tau.owner.inputs[0]
 
-    Q = tau.owner.inputs[0]
+    # 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]
 
     # 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}
+    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
+    # marginal_model.logp().dprint()
     idata = pm.sample(model=marginal_model, **sampler_kwargs)
 
     if not return_latent_posteriors:

pymc_extras/model/marginal/distributions.py

@@ -9,6 +9,7 @@
 from pymc.distributions import Bernoulli, Categorical, DiscreteUniform
 from pymc.distributions.distribution import _support_point, support_point
 from pymc.distributions.multivariate import _logdet_from_cholesky, nan_lower_cholesky
+from pymc.logprob import ValuedRV
 from pymc.logprob.abstract import MeasurableOp, _logprob
 from pymc.logprob.basic import conditional_logp, logp
 from pymc.pytensorf import constant_fold
@@ -142,12 +143,12 @@ def __init__(
         self,
         *args,
         Q: TensorVariable,
-        temp_kwargs: list,
+        minimizer_seed: int,
         minimizer_kwargs: dict | None = None,
         **kwargs,
     ) -> None:
-        self.temp_kwargs = temp_kwargs  # TODO REMOVE
         self.Q = Q
+        self.minimizer_seed = minimizer_seed
         self.minimizer_kwargs = minimizer_kwargs
         super().__init__(*args, **kwargs)
 
@@ -440,21 +441,42 @@ def laplace_marginal_rv_logp(op: MarginalLaplaceRV, values, *inputs, **kwargs):
     # Set minimizer initialisation to be random
     # TODO Assumes that the observed variable y is the first/only element of values, and that d is shape[-1]
     d = values[0].data.shape[-1]
-    rng = np.random.default_rng(12345)
+    rng = np.random.default_rng(op.minimizer_seed)
     x0_init = rng.random(d)
     x0 = pytensor.graph.replace.graph_replace(x0, {marginalized_vv: x0_init})
 
-    # TODO USE CLOSED FORM SOLUTION FOR NOW
-    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
     )  # TODO check how stan makes this quicker
-    tau = op.Q - hess
-    mu = x0  # TODO double check with Theo
+
+    # Get Q from the list of inputs
+    Q = None
+    if isinstance(op.Q, TensorVariable):
+        for var in inputs:
+            if var.owner is not None and isinstance(var.owner.op, ValuedRV):
+                for inp in var.owner.inputs:
+                    if (
+                        inp.name is not None
+                        and inp.name == op.Q.name
+                        or inp.name == op.Q.name + "_log"
+                    ):
+                        Q = var
+                        break
+
+            if var.name is not None and var.name == op.Q.name or var.name == op.Q.name + "_log":
+                Q = var
+                break
+
+        if Q is None:
+            raise ValueError(f"No inputs could be matched to precision matrix {op.Q}: {inputs}.")
+
+    # Q is an array
+    else:
+        Q = op.Q
+
+    tau = Q - hess
+    mu = x0
     log_laplace_approx, _ = _precision_mv_normal_logp(x0, mu, tau)
 
     # logp(y | params) = logp(y | x, params) + logp(x | params) - logp(x | y, params)

pymc_extras/model/marginal/marginal_model.py

@@ -9,7 +9,6 @@
 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
@@ -192,10 +191,10 @@ def marginalize(
                 raise NotImplementedError(
                     "Marginalization for DiscreteMarkovChain with non-matrix transition probability is not supported"
                 )
-        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 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(
@@ -587,9 +586,7 @@ def replace_marginal_subgraph(
                 "You can try splitting the marginalized RV into separate components and marginalizing them separately."
             ) from e
     else:
-        dependent_rvs_dim_connections = [
-            (None,),
-        ]
+        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)

tests/inference/INLA/test_inla.py

@@ -97,7 +97,6 @@ def test_3_layer_normal(rng):
         idata = pmx.fit(
             method="INLA",
             x=x,
-            temp_kwargs=[n, y_obs],  # TODO REMOVE LATER - DEBUGGING TOOL
             return_latent_posteriors=False,
         )