Variational inference with PyMC (#1306)

* variational inference fit

* remove variational from sample

* make pymc object accessible

* save as McmcPtResult

* tests added

* add warning in write_result()

pypesto/store/save_to_hdf5.py

@@ -316,3 +316,9 @@ def write_result(
  if sample:
  pypesto_sample_writer = SamplingResultHDF5Writer(filename)
  pypesto_sample_writer.write(result, overwrite=overwrite)
+
+ if hasattr(result, "variational_result"):
+ logger.warning(
+ "Results from variational inference are not saved in the hdf5 file. "
+ "You have to save them manually."
+ )

pypesto/variational/__init__.py

@@ -0,0 +1,9 @@
+"""
+Variational inference
+======
+
+Find the best variational approximation in a given family to a distribution from which we can sample.
+"""
+
+from .pymc import PymcVariational
+from .variational_inference import variational_fit

pypesto/variational/pymc.py

@@ -0,0 +1,196 @@
+"""Pymc v4 Sampler for Variational Inference."""
+
+import logging
+from typing import Optional
+
+import numpy as np
+import pytensor.tensor as pt
+from scipy import stats
+
+from ..objective import FD
+from ..result import McmcPtResult
+from ..sample.pymc import PymcObjectiveOp, PymcSampler
+from ..sample.sampler import SamplerImportError
+
+logger = logging.getLogger(__name__)
+
+
+# implementation based on the pymc sampler code in pypesto and:
+# https://www.pymc.io/projects/examples/en/latest/variational_inference/variational_api_quickstart.html
+
+
+class PymcVariational(PymcSampler):
+ """Wrapper around Pymc v4 variational inference.
+
+ Parameters
+ ----------
+ step_function:
+ A pymc step function, e.g. NUTS, Slice. If not specified, pymc
+ determines one automatically (preferable).
+ **kwargs:
+ Options are directly passed on to `pymc.fit`.
+ """
+
+ def fit(
+ self,
+ n_iterations: int,
+ method: str = "advi",
+ random_seed: Optional[int] = None,
+ start_sigma: Optional = None,
+ inf_kwargs: Optional = None,
+ beta: float = 1.0,
+ **kwargs,
+ ):
+ """
+ Sample the problem.
+
+ Parameters
+ ----------
+ n_iterations:
+ Number of iterations.
+ method: str or :class:`Inference` of pymc
+ string name is case-insensitive in:
+ - 'advi' for ADVI
+ - 'fullrank_advi' for FullRankADVI
+ - 'svgd' for Stein Variational Gradient Descent
+ - 'asvgd' for Amortized Stein Variational Gradient Descent
+ random_seed: int
+ random seed for reproducibility
+ start_sigma: `dict[str, np.ndarray]`
+ starting standard deviation for inference, only available for method 'advi'
+ inf_kwargs: dict
+ additional kwargs passed to pymc.Inference
+ beta:
+ Inverse temperature (e.g. in parallel tempering).
+ """
+ try:
+ import pymc
+ except ImportError:
+ raise SamplerImportError("pymc") from None
+
+ problem = self.problem
+ if not problem.objective.has_grad:
+ logger.info(
+ "The objective function does not provide gradients. "
+ "Finite differences will be used."
+ )
+ problem.objective = FD(obj=problem.objective)
+ log_post = PymcObjectiveOp.create_instance(problem.objective, beta)
+
+ x0 = None
+ x_names_free = problem.get_reduced_vector(problem.x_names)
+ if self.x0 is not None:
+ x0 = {
+ x_name: val
+ for x_name, val in zip(problem.x_names, self.x0)
+ if x_name in x_names_free
+ }
+
+ # create model context
+ with pymc.Model():
+ # parameter bounds as uniform prior
+ _k = [
+ pymc.Uniform(x_name, lower=lb, upper=ub)
+ for x_name, lb, ub in zip(
+ x_names_free,
+ problem.lb,
+ problem.ub,
+ )
+ ]
+
+ # convert parameters to PyTensor tensor variable
+ theta = pt.as_tensor_variable(_k)
+
+ # define distribution with log-posterior as density
+ pymc.Potential("potential", log_post(theta))
+
+ # record function values
+ pymc.Deterministic("loggyposty", log_post(theta))
+
+ # perform the actual sampling
+ data = pymc.fit(
+ n=int(n_iterations),
+ method=method,
+ random_seed=random_seed,
+ start=x0,
+ start_sigma=start_sigma,
+ inf_kwargs=inf_kwargs,
+ **kwargs,
+ )
+
+ self.data = data
+
+ def sample(self, n_samples: int, beta: float = 1.0) -> McmcPtResult:
+ """
+ Sample from the variational approximation and return McmcPtResult object.
+
+ Parameters
+ ----------
+ n_samples:
+ Number of samples to be computed.
+ """
+ # get InferenceData object
+ pymc_data = self.data.sample(n_samples)
+ x_names_free = self.problem.get_reduced_vector(self.problem.x_names)
+ post_samples = np.concatenate(
+ [pymc_data.posterior[name].values for name in x_names_free]
+ ).T
+ return McmcPtResult(
+ trace_x=post_samples[np.newaxis, :],
+ trace_neglogpost=pymc_data.posterior.loggyposty.values,
+ trace_neglogprior=np.full(
+ pymc_data.posterior.loggyposty.values.shape, np.nan
+ ),
+ betas=np.array([1.0] * post_samples.shape[0]),
+ burn_in=0,
+ auto_correlation=0,
+ effective_sample_size=n_samples,
+ message="variational inference results",
+ )
+
+ def get_variational_parameters(self) -> (list, list):
+ """Get the internal pymc variational parameters."""
+ return (
+ [param.name for param in self.data.params],
+ [param.eval() for param in self.data.params],
+ )
+
+ def set_variational_parameters(self, param_list: list):
+ """
+ Set the internal pymc variational parameters.
+
+ Parameters
+ ----------
+ param_list:
+ List of tuples of the form (param_name, param_value).
+ """
+ if len(param_list) != len(self.data.params):
+ raise ValueError(
+ "The number of parameters does not match the number of variational parameters."
+ )
+ for i, param in enumerate(param_list):
+ self.data.params[i].set_value(param)
+
+ def eval_variational_log_density(self, x: np.ndarray) -> np.ndarray:
+ """
+ Evaluate the log density of the variational approximation at x_points.
+
+ Parameters
+ ----------
+ x:
+ The points at which to evaluate the log density.
+ """
+ # TODO: add support for other methods
+ logger.warning(
+ "currently only supports the methods `advi` and `fullrank_advi`"
+ )
+
+ if x.ndim == 1:
+ x = x.reshape(1, -1)
+ log_density_at_points = np.zeros_like(x)
+ for i, point in enumerate(x):
+ log_density_at_points[i] = stats.multivariate_normal.logpdf(
+ point, mean=self.data.mean.eval(), cov=self.data.cov.eval()
+ )
+ vi_log_density = np.sum(log_density_at_points, axis=-1)
+ return vi_log_density

pypesto/variational/variational_inference.py

@@ -0,0 +1,136 @@
+"""Functions for variational inference accessible to the user. Currently only pymc is supported."""
+
+import logging
+from time import process_time
+from typing import Callable, List, Optional, Union
+
+import numpy as np
+
+from ..problem import Problem
+from ..result import Result
+from ..sample.util import bound_n_samples_from_env
+from ..store import autosave
+from .pymc import PymcVariational
+
+logger = logging.getLogger(__name__)
+
+
+def variational_fit(
+ problem: Problem,
+ n_iterations: int,
+ method: str = "advi",
+ n_samples: Optional[int] = None,
+ random_seed: Optional[int] = None,
+ start_sigma: Optional[dict[str, np.ndarray]] = None,
+ x0: Union[np.ndarray, List[np.ndarray]] = None,
+ result: Result = None,
+ filename: Union[str, Callable, None] = None,
+ overwrite: bool = False,
+ **kwargs,
+) -> Result:
+ """
+ Call to do parameter sampling.
+
+ Parameters
+ ----------
+ problem:
+ The problem to be solved. If None is provided, a
+ :class:`pypesto.AdaptiveMetropolisSampler` is used.
+ n_iterations:
+ Number of iterations for the optimization.
+ method: str or :class:`Inference` of pymc (only interface currently supported)
+ string name is case-insensitive in:
+ - 'advi' for ADVI
+ - 'fullrank_advi' for FullRankADVI
+ - 'svgd' for Stein Variational Gradient Descent
+ - 'asvgd' for Amortized Stein Variational Gradient Descent
+ n_samples:
+ Number of samples to generate after optimization.
+ random_seed: int
+ random seed for reproducibility
+ start_sigma: `dict[str, np.ndarray]`
+ starting standard deviation for inference, only available for method 'advi'
+ x0:
+ Initial parameter for the variational optimization. If None, the best parameter
+ found in optimization is used.
+ result:
+ A result to write to. If None provided, one is created from the
+ problem.
+ filename:
+ Name of the hdf5 file, where the result will be saved. Default is
+ None, which deactivates automatic saving. If set to
+ "Auto" it will automatically generate a file named
+ `year_month_day_profiling_result.hdf5`.
+ Optionally a method, see docs for `pypesto.store.auto.autosave`.
+ overwrite:
+ Whether to overwrite `result/sampling` in the autosave file
+ if it already exists.
+
+ Returns
+ -------
+ result:
+ A result with filled in sample_options part.
+ """
+ # prepare result object
+ if result is None:
+ result = Result(problem)
+
+ # number of samples
+ if n_iterations is not None:
+ n_iterations = bound_n_samples_from_env(n_iterations)
+
+ # try to find initial parameters
+ if x0 is None:
+ result.optimize_result.sort()
+ if len(result.optimize_result.list) > 0:
+ x0 = problem.get_reduced_vector(
+ result.optimize_result.list[0]["x"]
+ )
+
+ # set variational inference
+ # currently we only support pymc
+ variational = PymcVariational()
+
+ # initialize sampler to problem
+ variational.initialize(problem=problem, x0=x0)
+
+ # perform the sampling and track time
+ t_start = process_time()
+ variational.fit(
+ n_iterations=n_iterations,
+ method=method,
+ random_seed=random_seed,
+ start_sigma=start_sigma,
+ **kwargs,
+ )
+ t_elapsed = process_time() - t_start
+ logger.info("Elapsed time: " + str(t_elapsed))
+
+ # extract results and save samples to pypesto result
+ if n_samples is None or n_samples == 0:
+ # constructing a McmcPtResult object with nearly empty trace_x
+ n_samples = 1
+
+ result.sample_result = variational.sample(n_samples)
+ result.sample_result.time = t_elapsed
+
+ autosave(
+ filename=filename,
+ result=result,
+ store_type="sample",
+ overwrite=overwrite,
+ )
+
+ # make pymc object available in result
+ # TODO: if needed, we can add a result object for variational inference methods
+ result.variational_result = variational
+ (
+ result.sample_result.variational_parameters_names,
+ result.sample_result.variational_parameters,
+ ) = variational.get_variational_parameters()
+ if filename is not None:
+ logger.warning(
+ "Variational parameters are not saved in the hdf5 file. You have to save them manually."
+ )
+
+ return result

test/variational/__init__.py

@@ -0,0 +1 @@
+"""Variational inference tests."""