Implement Laplace (quadratic) approximation

conda-envs/environment-test.yml

@@ -10,6 +10,6 @@ dependencies:
 - xhistogram
 - statsmodels
 - pip:
-  - pymc>=5.13.0  # CI was failing to resolve
+  - pymc>=5.16.1  # CI was failing to resolve
   - blackjax
   - scikit-learn

conda-envs/windows-environment-test.yml

@@ -10,6 +10,6 @@ dependencies:
 - xhistogram
 - statsmodels
 - pip:
-  - pymc>=5.13.0  # CI was failing to resolve
+  - pymc>=5.16.1  # CI was failing to resolve
   - blackjax
   - scikit-learn

pymc_experimental/distributions/continuous.py

@@ -19,7 +19,7 @@
 The imports from pymc are not fully replicated here: add imports as necessary.
 """
 
-from typing import List, Tuple, Union
+from typing import Tuple, Union
 
 import numpy as np
 import pytensor.tensor as pt
@@ -37,8 +37,7 @@
 
 class GenExtremeRV(RandomVariable):
     name: str = "Generalized Extreme Value"
-    ndim_supp: int = 0
-    ndims_params: List[int] = [0, 0, 0]
+    signature = "(),(),()->()"
     dtype: str = "floatX"
     _print_name: Tuple[str, str] = ("Generalized Extreme Value", "\\operatorname{GEV}")
 
@@ -275,7 +274,7 @@ def chi_dist(nu: TensorVariable, size: TensorVariable) -> TensorVariable:
 
     def __new__(cls, name, nu, **kwargs):
         if "observed" not in kwargs:
-            kwargs.setdefault("transform", transforms.log)
+            kwargs.setdefault("default_transform", transforms.log)
         return CustomDist(name, nu, dist=cls.chi_dist, class_name="Chi", **kwargs)
 
     @classmethod

pymc_experimental/distributions/discrete.py

@@ -31,8 +31,7 @@ def log1mexp(x):
 
 class GeneralizedPoissonRV(RandomVariable):
     name = "generalized_poisson"
-    ndim_supp = 0
-    ndims_params = [0, 0]
+    signature = "(),()->()"
     dtype = "int64"
     _print_name = ("GeneralizedPoisson", "\\operatorname{GeneralizedPoisson}")
 

pymc_experimental/inference/fit.py

@@ -21,7 +21,7 @@ def fit(method, **kwargs):
     ----------
     method : str
         Which inference method to run.
-        Supported: pathfinder
+        Supported: pathfinder or laplace
 
     kwargs are passed on.
 
@@ -38,3 +38,9 @@ def fit(method, **kwargs):
         from pymc_experimental.inference.pathfinder import fit_pathfinder
 
         return fit_pathfinder(**kwargs)
+
+    if method == "laplace":
+
+        from pymc_experimental.inference.laplace import laplace
+
+        return laplace(**kwargs)

pymc_experimental/inference/laplace.py

@@ -0,0 +1,190 @@
+#   Copyright 2024 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+import warnings
+from collections.abc import Sequence
+from typing import Optional
+
+import arviz as az
+import numpy as np
+import pymc as pm
+import xarray as xr
+from arviz import dict_to_dataset
+from pymc.backends.arviz import (
+    coords_and_dims_for_inferencedata,
+    find_constants,
+    find_observations,
+)
+from pymc.model.transform.conditioning import remove_value_transforms
+from pymc.util import RandomSeed
+from pytensor import Variable
+
+
+def laplace(
+    vars: Sequence[Variable],
+    draws: Optional[int] = 1000,
+    model=None,
+    random_seed: Optional[RandomSeed] = None,
+    progressbar=True,
+):
+    """
+    Create a Laplace (quadratic) approximation for a posterior distribution.
+
+    This function generates a Laplace approximation for a given posterior distribution using a specified
+    number of draws. This is useful for obtaining a parametric approximation to the posterior distribution
+    that can be used for further analysis.
+
+    Parameters
+    ----------
+    vars : Sequence[Variable]
+        A sequence of variables for which the Laplace approximation of the posterior distribution
+        is to be created.
+    draws : Optional[int] with default=1_000
+        The number of draws to sample from the posterior distribution for creating the approximation.
+        For draws=None only the fit of the Laplace approximation is returned
+    model : object, optional, default=None
+        The model object that defines the posterior distribution. If None, the default model will be used.
+    random_seed : Optional[RandomSeed], optional, default=None
+        An optional random seed to ensure reproducibility of the draws. If None, the draws will be
+        generated using the current random state.
+    progressbar: bool, optional defaults to True
+        Whether to display a progress bar in the command line.
+
+    Returns
+    -------
+    arviz.InferenceData
+        An `InferenceData` object from the `arviz` library containing the Laplace
+        approximation of the posterior distribution. The inferenceData object also
+        contains constant and observed data as well as deterministic variables.
+        InferenceData also contains a group 'fit' with the mean and covariance
+        for the Laplace approximation.
+
+    Examples
+    --------
+
+    >>> import numpy as np
+    >>> import pymc as pm
+    >>> import arviz as az
+    >>> from pymc_experimental.inference.laplace import laplace
+    >>> y = np.array([2642, 3503, 4358]*10)
+    >>> with pm.Model() as m:
+    >>>     logsigma = pm.Uniform("logsigma", 1, 100)
+    >>>     mu = pm.Uniform("mu", -10000, 10000)
+    >>>     yobs = pm.Normal("y", mu=mu, sigma=pm.math.exp(logsigma), observed=y)
+    >>> idata = laplace([mu, logsigma], model=m)
+
+    Notes
+    -----
+    This method of approximation may not be suitable for all types of posterior distributions,
+    especially those with significant skewness or multimodality.
+
+    See Also
+    --------
+    fit : Calling the inference function 'fit' like pmx.fit(method="laplace", vars=[mu, logsigma], model=m)
+          will forward the call to 'laplace'.
+
+    """
+
+    rng = np.random.default_rng(seed=random_seed)
+
+    transformed_m = pm.modelcontext(model)
+
+    if len(vars) != len(transformed_m.free_RVs):
+        warnings.warn(
+            "Number of variables in vars does not eqaul the number of variables in the model.",
+            UserWarning,
+        )
+
+    map = pm.find_MAP(vars=vars, progressbar=progressbar, model=transformed_m)
+
+    # See https://www.pymc.io/projects/docs/en/stable/api/model/generated/pymc.model.transform.conditioning.remove_value_transforms.html
+    untransformed_m = remove_value_transforms(transformed_m)
+    untransformed_vars = [untransformed_m[v.name] for v in vars]
+    hessian = pm.find_hessian(point=map, vars=untransformed_vars, model=untransformed_m)
+
+    if np.linalg.det(hessian) == 0:
+        raise np.linalg.LinAlgError("Hessian is singular.")
+
+    cov = np.linalg.inv(hessian)
+    mean = np.concatenate([np.atleast_1d(map[v.name]) for v in vars])
+
+    chains = 1
+
+    if draws is not None:
+        samples = rng.multivariate_normal(mean, cov, size=(chains, draws))
+
+        data_vars = {}
+        for i, var in enumerate(vars):
+            data_vars[str(var)] = xr.DataArray(samples[:, :, i], dims=("chain", "draw"))
+
+        coords = {"chain": np.arange(chains), "draw": np.arange(draws)}
+        ds = xr.Dataset(data_vars, coords=coords)
+
+        idata = az.convert_to_inference_data(ds)
+        idata = addDataToInferenceData(model, idata, progressbar)
+    else:
+        idata = az.InferenceData()
+
+    idata = addFitToInferenceData(vars, idata, mean, cov)
+
+    return idata
+
+
+def addFitToInferenceData(vars, idata, mean, covariance):
+    coord_names = [v.name for v in vars]
+    # Convert to xarray DataArray
+    mean_dataarray = xr.DataArray(mean, dims=["rows"], coords={"rows": coord_names})
+    cov_dataarray = xr.DataArray(
+        covariance, dims=["rows", "columns"], coords={"rows": coord_names, "columns": coord_names}
+    )
+
+    # Create xarray dataset
+    dataset = xr.Dataset({"mean_vector": mean_dataarray, "covariance_matrix": cov_dataarray})
+
+    idata.add_groups(fit=dataset)
+
+    return idata
+
+
+def addDataToInferenceData(model, trace, progressbar):
+    # Add deterministic variables to inference data
+    trace.posterior = pm.compute_deterministics(
+        trace.posterior, model=model, merge_dataset=True, progressbar=progressbar
+    )
+
+    coords, dims = coords_and_dims_for_inferencedata(model)
+
+    observed_data = dict_to_dataset(
+        find_observations(model),
+        library=pm,
+        coords=coords,
+        dims=dims,
+        default_dims=[],
+    )
+
+    constant_data = dict_to_dataset(
+        find_constants(model),
+        library=pm,
+        coords=coords,
+        dims=dims,
+        default_dims=[],
+    )
+
+    trace.add_groups(
+        {"observed_data": observed_data, "constant_data": constant_data},
+        coords=coords,
+        dims=dims,
+    )
+
+    return trace

pymc_experimental/inference/smc/sampling.py

@@ -427,7 +427,7 @@ def build_smc_with_kernel(
     kernel_parameters,
     mcmc_kernel,
 ):
-    return blackjax.smc.adaptive_tempered.adaptive_tempered_smc(
+    return blackjax.adaptive_tempered_smc(
         prior_log_prob,
         loglikelihood,
         mcmc_kernel.build_kernel(),

pymc_experimental/model/marginal_model.py

@@ -21,7 +21,7 @@
 from pytensor.graph.replace import graph_replace, vectorize_graph
 from pytensor.scan import map as scan_map
 from pytensor.tensor import TensorType, TensorVariable
-from pytensor.tensor.elemwise import Elemwise
+from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.shape import Shape
 from pytensor.tensor.special import log_softmax
 
@@ -598,7 +598,18 @@ def is_elemwise_subgraph(rv_to_marginalize, other_input_rvs, output_rvs):
     fg = FunctionGraph(outputs=output_rvs, clone=False)
 
     non_elemwise_blockers = [
-        o for node in fg.apply_nodes if not isinstance(node.op, Elemwise) for o in node.outputs
+        o
+        for node in fg.apply_nodes
+        if not (
+            isinstance(node.op, Elemwise)
+            # Allow expand_dims on the left
+            or (
+                isinstance(node.op, DimShuffle)
+                and not node.op.drop
+                and node.op.shuffle == sorted(node.op.shuffle)
+            )
+        )
+        for o in node.outputs
     ]
     blocker_candidates = [rv_to_marginalize] + other_input_rvs + non_elemwise_blockers
     blockers = [var for var in blocker_candidates if var not in output_rvs]
@@ -698,16 +709,17 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
 
 def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> tuple[int, ...]:
     op = rv.owner.op
+    dist_params = rv.owner.op.dist_params(rv.owner)
     if isinstance(op, Bernoulli):
         return (0, 1)
     elif isinstance(op, Categorical):
-        p_param = rv.owner.inputs[3]
+        [p_param] = dist_params
         return tuple(range(pt.get_vector_length(p_param)))
     elif isinstance(op, DiscreteUniform):
-        lower, upper = constant_fold(rv.owner.inputs[3:])
+        lower, upper = constant_fold(dist_params)
         return tuple(np.arange(lower, upper + 1))
     elif isinstance(op, DiscreteMarkovChain):
-        P = rv.owner.inputs[0]
+        P, *_ = dist_params
         return tuple(range(pt.get_vector_length(P[-1])))
 
     raise NotImplementedError(f"Cannot compute domain for op {op}")
@@ -827,11 +839,15 @@ def marginal_hmm_logp(op, values, *inputs, **kwargs):
     # This is the "forward algorithm", alpha_t = p(y | s_t) * sum_{s_{t-1}}(p(s_t | s_{t-1}) * alpha_{t-1})
     # We do it entirely in logs, though.
 
-    # To compute the prior probabilities of each state, we evaluate the logp of the domain (all possible states) under
-    # the initial distribution. This is robust to everything the user can throw at it.
-    batch_logp_init_dist = pt.vectorize(lambda x: logp(init_dist_, x), "()->()")(
-        batch_chain_value[..., 0]
-    )
+    # To compute the prior probabilities of each state, we evaluate the logp of the domain (all possible states)
+    # under the initial distribution. This is robust to everything the user can throw at it.
+    init_dist_value = init_dist_.type()
+    logp_init_dist = logp(init_dist_, init_dist_value)
+    # There is a degerate batch dim for lags=1 (the only supported case),
+    # that we have to work around, by expanding the batch value and then squeezing it out of the logp
+    batch_logp_init_dist = vectorize_graph(
+        logp_init_dist, {init_dist_value: batch_chain_value[:, None, ..., 0]}
+    ).squeeze(1)
     log_alpha_init = batch_logp_init_dist + batch_logp_emissions[..., 0]
 
     def step_alpha(logp_emission, log_alpha, log_P):

pymc_experimental/model/transforms/autoreparam.py

@@ -7,6 +7,8 @@
 import pytensor
 import pytensor.tensor as pt
 import scipy.special
+from pymc.distributions import SymbolicRandomVariable
+from pymc.exceptions import NotConstantValueError
 from pymc.logprob.transforms import Transform
 from pymc.model.fgraph import (
     ModelDeterministic,
@@ -17,7 +19,7 @@
     model_from_fgraph,
     model_named,
 )
-from pymc.pytensorf import toposort_replace
+from pymc.pytensorf import constant_fold, toposort_replace
 from pytensor.graph.basic import Apply, Variable
 from pytensor.tensor.random.op import RandomVariable
 
@@ -170,14 +172,16 @@ def vip_reparam_node(
     dims: List[Variable],
     transform: Optional[Transform],
 ) -> Tuple[ModelDeterministic, ModelNamed]:
-    if not isinstance(node.op, RandomVariable):
+    if not isinstance(node.op, RandomVariable | SymbolicRandomVariable):
         raise TypeError("Op should be RandomVariable type")
-    size = node.inputs[1]
-    if not isinstance(size, pt.TensorConstant):
+    rv = node.default_output()
+    try:
+        [rv_shape] = constant_fold([rv.shape])
+    except NotConstantValueError:
         raise ValueError("Size should be static for autoreparametrization.")
     logit_lam_ = pytensor.shared(
-        np.zeros(size.data),
-        shape=size.data,
+        np.zeros(rv_shape),
+        shape=rv_shape,
         name=f"{name}::lam_logit__",
     )
     logit_lam = model_named(logit_lam_, *dims)
@@ -216,7 +220,7 @@ def _(
     transform: Optional[Transform],
     lam: pt.TensorVariable,
 ) -> ModelDeterministic:
-    rng, size, _, loc, scale = node.inputs
+    rng, size, loc, scale = node.inputs
     if transform is not None:
         raise NotImplementedError("Reparametrization of Normal with Transform is not implemented")
     vip_rv_ = pm.Normal.dist(