pymc-extras 0.4.1__py3-none-any.whl → 0.6.0__py3-none-any.whl

pymc_extras/inference/laplace_approx/idata.py

@@ -22,10 +22,15 @@ def make_default_labels(name: str, shape: tuple[int, ...]) -> list:
     return [list(range(dim)) for dim in shape]
 
 
-def make_unpacked_variable_names(names: list[str], model: pm.Model) -> list[str]:
+def make_unpacked_variable_names(
+    names: list[str], model: pm.Model, var_name_to_model_var: dict[str, str] | None = None
+) -> list[str]:
     coords = model.coords
     initial_point = model.initial_point()
 
+    if var_name_to_model_var is None:
+        var_name_to_model_var = {}
+
     value_to_dim = {
         value.name: model.named_vars_to_dims.get(model.values_to_rvs[value].name, None)
         for value in model.value_vars
@@ -37,6 +42,7 @@ def make_unpacked_variable_names(names: list[str], model: pm.Model) -> list[str]
 
     unpacked_variable_names = []
     for name in names:
+        name = var_name_to_model_var.get(name, name)
         shape = initial_point[name].shape
         if shape:
             dims = dims_dict.get(name)
@@ -109,7 +115,7 @@ def map_results_to_inference_data(
         x.name for x in get_default_varnames(model.unobserved_value_vars, include_transformed=True)
     ]
 
-    unconstrained_names = set(all_varnames) - set(constrained_names)
+    unconstrained_names = sorted(set(all_varnames) - set(constrained_names))
 
     idata = az.from_dict(
         posterior={
@@ -136,7 +142,10 @@ def map_results_to_inference_data(
 
 
 def add_fit_to_inference_data(
-    idata: az.InferenceData, mu: RaveledVars, H_inv: np.ndarray, model: pm.Model | None = None
+    idata: az.InferenceData,
+    mu: RaveledVars,
+    H_inv: np.ndarray | None,
+    model: pm.Model | None = None,
 ) -> az.InferenceData:
     """
     Add the mean vector and covariance matrix of the Laplace approximation to an InferenceData object.
@@ -147,7 +156,7 @@ def add_fit_to_inference_data(
         An InferenceData object containing the approximated posterior samples.
     mu: RaveledVars
         The MAP estimate of the model parameters.
-    H_inv: np.ndarray
+    H_inv: np.ndarray, optional
         The inverse Hessian matrix of the log-posterior evaluated at the MAP estimate.
     model: Model, optional
         A PyMC model. If None, the model is taken from the current model context.
@@ -255,6 +264,7 @@ def optimizer_result_to_dataset(
     method: minimize_method | Literal["basinhopping"],
     mu: RaveledVars | None = None,
     model: pm.Model | None = None,
+    var_name_to_model_var: dict[str, str] | None = None,
 ) -> xr.Dataset:
     """
     Convert an OptimizeResult object to an xarray Dataset object.
@@ -265,6 +275,9 @@ def optimizer_result_to_dataset(
         The result of the optimization process.
     method: minimize_method or "basinhopping"
         The optimization method used.
+    var_name_to_model_var: dict, optional
+        Mapping between variables in the optimization result and the model variable names. Used when auxiliary
+        variables were introduced, e.g. in DADVI.
 
     Returns
     -------
@@ -276,7 +289,9 @@ def optimizer_result_to_dataset(
 
     model = pm.modelcontext(model) if model is None else model
     variable_names, *_ = zip(*mu.point_map_info)
-    unpacked_variable_names = make_unpacked_variable_names(variable_names, model)
+    unpacked_variable_names = make_unpacked_variable_names(
+        variable_names, model, var_name_to_model_var
+    )
 
     data_vars = {}
 
@@ -365,6 +380,7 @@ def add_optimizer_result_to_inference_data(
     method: minimize_method | Literal["basinhopping"],
     mu: RaveledVars | None = None,
     model: pm.Model | None = None,
+    var_name_to_model_var: dict[str, str] | None = None,
 ) -> az.InferenceData:
     """
     Add the optimization result to an InferenceData object.
@@ -381,13 +397,18 @@ def add_optimizer_result_to_inference_data(
         The MAP estimate of the model parameters.
     model: Model, optional
         A PyMC model. If None, the model is taken from the current model context.
+    var_name_to_model_var: dict, optional
+        Mapping between variables in the optimization result and the model variable names. Used when auxiliary
+        variables were introduced, e.g. in DADVI.
 
     Returns
     -------
     idata: az.InferenceData
         The provided InferenceData, with the optimization results added to the "optimizer" group.
     """
-    dataset = optimizer_result_to_dataset(result, method=method, mu=mu, model=model)
+    dataset = optimizer_result_to_dataset(
+        result, method=method, mu=mu, model=model, var_name_to_model_var=var_name_to_model_var
+    )
     idata.add_groups({"optimizer_result": dataset})
 
     return idata

pymc_extras/inference/laplace_approx/laplace.py

@@ -168,9 +168,13 @@ def _unconstrained_vector_to_constrained_rvs(model):
     unconstrained_vector.name = "unconstrained_vector"
 
     # Redo the names list to ensure it is sorted to match the return order
-    names = [*constrained_names, *unconstrained_names]
+    constrained_rvs_and_names = [(rv, name) for rv, name in zip(constrained_rvs, constrained_names)]
+    value_rvs_and_names = [
+        (rv, name) for rv, name in zip(value_rvs, names) for name in unconstrained_names
+    ]
+    # names = [*constrained_names, *unconstrained_names]
 
-    return names, constrained_rvs, value_rvs, unconstrained_vector
+    return constrained_rvs_and_names, value_rvs_and_names, unconstrained_vector
 
 
 def model_to_laplace_approx(
@@ -182,8 +186,11 @@ def model_to_laplace_approx(
 
     # temp_chain and temp_draw are a hack to allow sampling from the Laplace approximation. We only have one mu and cov,
     # so we add batch dims (which correspond to chains and draws). But the names "chain" and "draw" are reserved.
-    names, constrained_rvs, value_rvs, unconstrained_vector = (
-        _unconstrained_vector_to_constrained_rvs(model)
+
+    # The model was frozen during the find_MAP procedure. To ensure we're operating on the same model, freeze it again.
+    frozen_model = freeze_dims_and_data(model)
+    constrained_rvs_and_names, _, unconstrained_vector = _unconstrained_vector_to_constrained_rvs(
+        frozen_model
     )
 
     coords = model.coords | {
@@ -204,12 +211,13 @@ def model_to_laplace_approx(
         )
 
         cast_to_var = partial(type_cast, Variable)
+        constrained_rvs, constrained_names = zip(*constrained_rvs_and_names)
         batched_rvs = vectorize_graph(
             type_cast(list[Variable], constrained_rvs),
             replace={cast_to_var(unconstrained_vector): cast_to_var(laplace_approximation)},
         )
 
-        for name, batched_rv in zip(names, batched_rvs):
+        for name, batched_rv in zip(constrained_names, batched_rvs):
             batch_dims = ("temp_chain", "temp_draw")
             if batched_rv.ndim == 2:
                 dims = batch_dims
@@ -285,6 +293,7 @@ def fit_laplace(
     jitter_rvs: list[pt.TensorVariable] | None = None,
     progressbar: bool = True,
     include_transformed: bool = True,
+    freeze_model: bool = True,
     gradient_backend: GradientBackend = "pytensor",
     chains: int = 2,
     draws: int = 500,
@@ -328,6 +337,10 @@ def fit_laplace(
     include_transformed: bool, default True
         Whether to include transformed variables in the output. If True, transformed variables will be included in the
         output InferenceData object. If False, only the original variables will be included.
+    freeze_model: bool, optional
+        If True, freeze_dims_and_data will be called on the model before compiling the loss functions. This is
+        sometimes necessary for JAX, and can sometimes improve performance by allowing constant folding. Defaults to
+        True.
     gradient_backend: str, default "pytensor"
         The backend to use for gradient computations. Must be one of "pytensor" or "jax".
     chains: int, default: 2
@@ -354,7 +367,7 @@ def fit_laplace(
     >>> import numpy as np
     >>> import pymc as pm
     >>> import arviz as az
-    >>> y = np.array([2642, 3503, 4358]*10)
+    >>> y = np.array([2642, 3503, 4358] * 10)
     >>> with pm.Model() as m:
     >>>     logsigma = pm.Uniform("logsigma", 1, 100)
     >>>     mu = pm.Uniform("mu", -10000, 10000)
@@ -376,6 +389,9 @@ def fit_laplace(
     optimizer_kwargs = {} if optimizer_kwargs is None else optimizer_kwargs
     model = pm.modelcontext(model) if model is None else model
 
+    if freeze_model:
+        model = freeze_dims_and_data(model)
+
     idata = find_MAP(
         method=optimize_method,
         model=model,
@@ -387,8 +403,10 @@ def fit_laplace(
         jitter_rvs=jitter_rvs,
         progressbar=progressbar,
         include_transformed=include_transformed,
+        freeze_model=False,
         gradient_backend=gradient_backend,
         compile_kwargs=compile_kwargs,
+        compute_hessian=True,
         **optimizer_kwargs,
     )
 

pymc_extras/inference/laplace_approx/scipy_interface.py

@@ -1,3 +1,5 @@
+import logging
+
 from collections.abc import Callable
 from importlib.util import find_spec
 from typing import Literal, get_args
@@ -6,6 +8,7 @@ import numpy as np
 import pymc as pm
 import pytensor
 
+from better_optimize.constants import MINIMIZE_MODE_KWARGS
 from pymc import join_nonshared_inputs
 from pytensor import tensor as pt
 from pytensor.compile import Function
@@ -14,6 +17,39 @@ from pytensor.tensor import TensorVariable
 GradientBackend = Literal["pytensor", "jax"]
 VALID_BACKENDS = get_args(GradientBackend)
 
+_log = logging.getLogger(__name__)
+
+
+def set_optimizer_function_defaults(
+    method: str, use_grad: bool | None, use_hess: bool | None, use_hessp: bool | None
+):
+    method_info = MINIMIZE_MODE_KWARGS[method].copy()
+
+    if use_hess and use_hessp:
+        _log.warning(
+            'Both "use_hess" and "use_hessp" are set to True, but scipy.optimize.minimize never uses both at the '
+            'same time. When possible "use_hessp" is preferred because its is computationally more efficient. '
+            'Setting "use_hess" to False.'
+        )
+        use_hess = False
+
+    use_grad = use_grad if use_grad is not None else method_info["uses_grad"]
+
+    if use_hessp is not None and use_hess is None:
+        use_hess = not use_hessp
+
+    elif use_hess is not None and use_hessp is None:
+        use_hessp = not use_hess
+
+    elif use_hessp is None and use_hess is None:
+        use_hessp = method_info["uses_hessp"]
+        use_hess = method_info["uses_hess"]
+        if use_hessp and use_hess:
+            # If a method could use either hess or hessp, we default to using hessp
+            use_hess = False
+
+    return use_grad, use_hess, use_hessp
+
 
 def _compile_grad_and_hess_to_jax(
     f_fused: Function, use_hess: bool, use_hessp: bool
@@ -144,12 +180,13 @@ def _compile_functions_for_scipy_optimize(
 def scipy_optimize_funcs_from_loss(
     loss: TensorVariable,
     inputs: list[TensorVariable],
-    initial_point_dict: dict[str, np.ndarray | float | int],
-    use_grad: bool,
-    use_hess: bool,
-    use_hessp: bool,
+    initial_point_dict: dict[str, np.ndarray | float | int] | None = None,
+    use_grad: bool | None = None,
+    use_hess: bool | None = None,
+    use_hessp: bool | None = None,
     gradient_backend: GradientBackend = "pytensor",
     compile_kwargs: dict | None = None,
+    inputs_are_flat: bool = False,
 ) -> tuple[Callable, ...]:
     """
     Compile loss functions for use with scipy.optimize.minimize.
@@ -206,9 +243,12 @@ def scipy_optimize_funcs_from_loss(
     if not isinstance(inputs, list):
         inputs = [inputs]
 
-    [loss], flat_input = join_nonshared_inputs(
-        point=initial_point_dict, outputs=[loss], inputs=inputs
-    )
+    if inputs_are_flat:
+        [flat_input] = inputs
+    else:
+        [loss], flat_input = join_nonshared_inputs(
+            point=initial_point_dict, outputs=[loss], inputs=inputs
+        )
 
     # If we use pytensor gradients, we will use the pytensor function wrapper that handles shared variables. When
     # computing jax gradients, we discard the function wrapper, so we can't handle shared variables --> rewrite them