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(