pymc-extras 0.2.0__py3-none-any.whl

pymc_extras/inference/find_map.py

@@ -0,0 +1,431 @@
+import logging
+
+from collections.abc import Callable
+from typing import Literal, cast, get_args
+
+import jax
+import numpy as np
+import pymc as pm
+import pytensor
+import pytensor.tensor as pt
+
+from better_optimize import minimize
+from better_optimize.constants import MINIMIZE_MODE_KWARGS, minimize_method
+from pymc.blocking import DictToArrayBijection, RaveledVars
+from pymc.initial_point import make_initial_point_fn
+from pymc.model.transform.optimization import freeze_dims_and_data
+from pymc.pytensorf import join_nonshared_inputs
+from pymc.util import get_default_varnames
+from pytensor.compile import Function
+from pytensor.compile.mode import Mode
+from pytensor.tensor import TensorVariable
+from scipy.optimize import OptimizeResult
+
+_log = logging.getLogger(__name__)
+
+GradientBackend = Literal["pytensor", "jax"]
+VALID_BACKENDS = get_args(GradientBackend)
+
+
+def set_optimizer_function_defaults(method, use_grad, use_hess, use_hessp):
+    method_info = MINIMIZE_MODE_KWARGS[method].copy()
+
+    use_grad = use_grad if use_grad is not None else method_info["uses_grad"]
+    use_hess = use_hess if use_hess is not None else method_info["uses_hess"]
+    use_hessp = use_hessp if use_hessp is not None else method_info["uses_hessp"]
+
+    if use_hess and use_hessp:
+        use_hess = False
+
+    return use_grad, use_hess, use_hessp
+
+
+def get_nearest_psd(A: np.ndarray) -> np.ndarray:
+    """
+    Compute the nearest positive semi-definite matrix to a given matrix.
+
+    This function takes a square matrix and returns the nearest positive semi-definite matrix using
+    eigenvalue decomposition. It ensures all eigenvalues are non-negative. The "nearest" matrix is defined in terms
+    of the Frobenius norm.
+
+    Parameters
+    ----------
+    A : np.ndarray
+        Input square matrix.
+
+    Returns
+    -------
+    np.ndarray
+        The nearest positive semi-definite matrix to the input matrix.
+    """
+    C = (A + A.T) / 2
+    eigval, eigvec = np.linalg.eig(C)
+    eigval[eigval < 0] = 0
+
+    return eigvec @ np.diag(eigval) @ eigvec.T
+
+
+def _unconstrained_vector_to_constrained_rvs(model):
+    constrained_rvs, unconstrained_vector = join_nonshared_inputs(
+        model.initial_point(),
+        inputs=model.value_vars,
+        outputs=get_default_varnames(model.unobserved_value_vars, include_transformed=False),
+    )
+
+    unconstrained_vector.name = "unconstrained_vector"
+    return constrained_rvs, unconstrained_vector
+
+
+def _create_transformed_draws(H_inv, slices, out_shapes, posterior_draws, model, chains, draws):
+    X = pt.tensor("transformed_draws", shape=(chains, draws, H_inv.shape[0]))
+    out = []
+    for rv, idx in slices.items():
+        f = model.rvs_to_transforms[rv]
+        untransformed_X = f.backward(X[..., idx]) if f is not None else X[..., idx]
+
+        if rv in out_shapes:
+            new_shape = (chains, draws) + out_shapes[rv]
+            untransformed_X = untransformed_X.reshape(new_shape)
+
+        out.append(untransformed_X)
+
+    f_untransform = pytensor.function(
+        inputs=[pytensor.In(X, borrow=True)],
+        outputs=pytensor.Out(out, borrow=True),
+        mode=Mode(linker="py", optimizer="FAST_COMPILE"),
+    )
+    return f_untransform(posterior_draws)
+
+
+def _compile_jax_gradients(
+    f_loss: Function, use_hess: bool, use_hessp: bool
+) -> tuple[Callable | None, Callable | None]:
+    """
+    Compile loss function gradients using JAX.
+
+    Parameters
+    ----------
+    f_loss: Function
+        The loss function to compile gradients for. Expected to be a pytensor function that returns a scalar loss,
+        compiled with mode="JAX".
+    use_hess: bool
+        Whether to compile a function to compute the hessian of the loss function.
+    use_hessp: bool
+        Whether to compile a function to compute the hessian-vector product of the loss function.
+
+    Returns
+    -------
+    f_loss_and_grad: Callable
+        The compiled loss function and gradient function.
+    f_hess: Callable | None
+        The compiled hessian function, or None if use_hess is False.
+    f_hessp: Callable | None
+        The compiled hessian-vector product function, or None if use_hessp is False.
+    """
+    f_hess = None
+    f_hessp = None
+
+    orig_loss_fn = f_loss.vm.jit_fn
+
+    @jax.jit
+    def loss_fn_jax_grad(x, *shared):
+        return jax.value_and_grad(lambda x: orig_loss_fn(x)[0])(x)
+
+    f_loss_and_grad = loss_fn_jax_grad
+
+    if use_hessp:
+
+        def f_hessp_jax(x, p):
+            y, u = jax.jvp(lambda x: f_loss_and_grad(x)[1], (x,), (p,))
+            return jax.numpy.stack(u)
+
+        f_hessp = jax.jit(f_hessp_jax)
+
+    if use_hess:
+        _f_hess_jax = jax.jacfwd(lambda x: f_loss_and_grad(x)[1])
+
+        def f_hess_jax(x):
+            return jax.numpy.stack(_f_hess_jax(x))
+
+        f_hess = jax.jit(f_hess_jax)
+
+    return f_loss_and_grad, f_hess, f_hessp
+
+
+def _compile_functions(
+    loss: TensorVariable,
+    inputs: list[TensorVariable],
+    compute_grad: bool,
+    compute_hess: bool,
+    compute_hessp: bool,
+    compile_kwargs: dict | None = None,
+) -> list[Function] | list[Function, Function | None, Function | None]:
+    """
+    Compile loss functions for use with scipy.optimize.minimize.
+
+    Parameters
+    ----------
+    loss: TensorVariable
+        The loss function to compile.
+    inputs: list[TensorVariable]
+        A single flat vector input variable, collecting all inputs to the loss function. Scipy optimize routines
+        expect the function signature to be f(x, *args), where x is a 1D array of parameters.
+    compute_grad: bool
+        Whether to compile a function that computes the gradients of the loss function.
+    compute_hess: bool
+        Whether to compile a function that computes the Hessian of the loss function.
+    compute_hessp: bool
+        Whether to compile a function that computes the Hessian-vector product of the loss function.
+    compile_kwargs: dict, optional
+        Additional keyword arguments to pass to the ``pm.compile_pymc`` function.
+
+    Returns
+    -------
+    f_loss: Function
+
+    f_hess: Function | None
+    f_hessp: Function | None
+    """
+    loss = pm.pytensorf.rewrite_pregrad(loss)
+    f_hess = None
+    f_hessp = None
+
+    if compute_grad:
+        grads = pytensor.gradient.grad(loss, inputs)
+        grad = pt.concatenate([grad.ravel() for grad in grads])
+        f_loss_and_grad = pm.compile_pymc(inputs, [loss, grad], **compile_kwargs)
+    else:
+        f_loss = pm.compile_pymc(inputs, loss, **compile_kwargs)
+        return [f_loss]
+
+    if compute_hess:
+        hess = pytensor.gradient.jacobian(grad, inputs)[0]
+        f_hess = pm.compile_pymc(inputs, hess, **compile_kwargs)
+
+    if compute_hessp:
+        p = pt.tensor("p", shape=inputs[0].type.shape)
+        hessp = pytensor.gradient.hessian_vector_product(loss, inputs, p)
+        f_hessp = pm.compile_pymc([*inputs, p], hessp[0], **compile_kwargs)
+
+    return [f_loss_and_grad, f_hess, f_hessp]
+
+
+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,
+    gradient_backend: GradientBackend = "pytensor",
+    compile_kwargs: dict | None = None,
+) -> tuple[Callable, ...]:
+    """
+    Compile loss functions for use with scipy.optimize.minimize.
+
+    Parameters
+    ----------
+    loss: TensorVariable
+        The loss function to compile.
+    inputs: list[TensorVariable]
+        The input variables to the loss function.
+    initial_point_dict: dict[str, np.ndarray | float | int]
+        Dictionary mapping variable names to initial values. Used to determine the shapes of the input variables.
+    use_grad: bool
+        Whether to compile a function that computes the gradients of the loss function.
+    use_hess: bool
+        Whether to compile a function that computes the Hessian of the loss function.
+    use_hessp: bool
+        Whether to compile a function that computes the Hessian-vector product of the loss function.
+    gradient_backend: str, default "pytensor"
+        Which backend to use to compute gradients. Must be one of "jax" or "pytensor"
+    compile_kwargs:
+        Additional keyword arguments to pass to the ``pm.compile_pymc`` function.
+
+    Returns
+    -------
+    f_loss: Callable
+        The compiled loss function.
+    f_hess: Callable | None
+        The compiled hessian function, or None if use_hess is False.
+    f_hessp: Callable | None
+        The compiled hessian-vector product function, or None if use_hessp is False.
+    """
+
+    compile_kwargs = {} if compile_kwargs is None else compile_kwargs
+
+    if (use_hess or use_hessp) and not use_grad:
+        raise ValueError(
+            "Cannot compute hessian or hessian-vector product without also computing the gradient"
+        )
+
+    if gradient_backend not in VALID_BACKENDS:
+        raise ValueError(
+            f"Invalid gradient backend: {gradient_backend}. Must be one of {VALID_BACKENDS}"
+        )
+
+    use_jax_gradients = (gradient_backend == "jax") and use_grad
+
+    mode = compile_kwargs.get("mode", None)
+    if mode is None and use_jax_gradients:
+        compile_kwargs["mode"] = "JAX"
+    elif mode != "JAX" and use_jax_gradients:
+        raise ValueError(
+            'jax gradients can only be used when ``compile_kwargs["mode"]`` is set to "JAX"'
+        )
+
+    if not isinstance(inputs, list):
+        inputs = [inputs]
+
+    [loss], flat_input = join_nonshared_inputs(
+        point=initial_point_dict, outputs=[loss], inputs=inputs
+    )
+
+    compute_grad = use_grad and not use_jax_gradients
+    compute_hess = use_hess and not use_jax_gradients
+    compute_hessp = use_hessp and not use_jax_gradients
+
+    funcs = _compile_functions(
+        loss=loss,
+        inputs=[flat_input],
+        compute_grad=compute_grad,
+        compute_hess=compute_hess,
+        compute_hessp=compute_hessp,
+        compile_kwargs=compile_kwargs,
+    )
+
+    # f_loss here is f_loss_and_grad if compute_grad = True. The name is unchanged to simplify the return values
+    f_loss = funcs.pop(0)
+    f_hess = funcs.pop(0) if compute_grad else None
+    f_hessp = funcs.pop(0) if compute_grad else None
+
+    if use_jax_gradients:
+        # f_loss here is f_loss_and_grad; the name is unchanged to simplify the return values
+        f_loss, f_hess, f_hessp = _compile_jax_gradients(f_loss, use_hess, use_hessp)
+
+    return f_loss, f_hess, f_hessp
+
+
+def find_MAP(
+    method: minimize_method,
+    *,
+    model: pm.Model | None = None,
+    use_grad: bool | None = None,
+    use_hessp: bool | None = None,
+    use_hess: bool | None = None,
+    initvals: dict | None = None,
+    random_seed: int | np.random.Generator | None = None,
+    return_raw: bool = False,
+    jitter_rvs: list[TensorVariable] | None = None,
+    progressbar: bool = True,
+    include_transformed: bool = True,
+    gradient_backend: GradientBackend = "pytensor",
+    compile_kwargs: dict | None = None,
+    **optimizer_kwargs,
+) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], OptimizeResult]:
+    """
+    Fit a PyMC model via maximum a posteriori (MAP) estimation using JAX and scipy.minimize.
+
+    Parameters
+    ----------
+    model : pm.Model
+        The PyMC model to be fit. If None, the current model context is used.
+    method : str
+        The optimization method to use. See scipy.optimize.minimize documentation for details.
+    use_grad : bool | None, optional
+        Whether to use gradients in the optimization. Defaults to None, which determines this automatically based on
+        the ``method``.
+    use_hessp : bool | None, optional
+        Whether to use Hessian-vector products in the optimization. Defaults to None, which determines this automatically based on
+        the ``method``.
+    use_hess : bool | None, optional
+        Whether to use the Hessian matrix in the optimization. Defaults to None, which determines this automatically based on
+        the ``method``.
+    initvals : None | dict, optional
+        Initial values for the model parameters, as str:ndarray key-value pairs. Paritial initialization is permitted.
+         If None, the model's default initial values are used.
+    random_seed : None | int | np.random.Generator, optional
+        Seed for the random number generator or a numpy Generator for reproducibility
+    return_raw: bool | False, optinal
+        Whether to also return the full output of `scipy.optimize.minimize`
+    jitter_rvs : list of TensorVariables, optional
+        Variables whose initial values should be jittered. If None, all variables are jittered.
+    progressbar : bool, optional
+        Whether to display a progress bar during optimization. Defaults to True.
+    include_transformed: bool, optional
+        Whether to include transformed variable values in the returned dictionary. Defaults to True.
+    gradient_backend: str, default "pytensor"
+        Which backend to use to compute gradients. Must be one of "pytensor" or "jax".
+    compile_kwargs: dict, optional
+        Additional options to pass to the ``pytensor.function`` function when compiling loss functions.
+    **optimizer_kwargs
+        Additional keyword arguments to pass to the ``scipy.optimize.minimize`` function.
+
+    Returns
+    -------
+    optimizer_result: dict[str, np.ndarray] or tuple[dict[str, np.ndarray], OptimizerResult]
+        Dictionary with names of random variables as keys, and optimization results as values. If return_raw is True,
+        also returns the object returned by ``scipy.optimize.minimize``.
+    """
+    model = pm.modelcontext(model)
+    frozen_model = freeze_dims_and_data(model)
+
+    jitter_rvs = [] if jitter_rvs is None else jitter_rvs
+    compile_kwargs = {} if compile_kwargs is None else compile_kwargs
+
+    ipfn = make_initial_point_fn(
+        model=frozen_model,
+        jitter_rvs=set(jitter_rvs),
+        return_transformed=True,
+        overrides=initvals,
+    )
+
+    start_dict = ipfn(random_seed)
+    vars_dict = {var.name: var for var in frozen_model.continuous_value_vars}
+    initial_params = DictToArrayBijection.map(
+        {var_name: value for var_name, value in start_dict.items() if var_name in vars_dict}
+    )
+    use_grad, use_hess, use_hessp = set_optimizer_function_defaults(
+        method, use_grad, use_hess, use_hessp
+    )
+
+    f_logp, f_hess, f_hessp = scipy_optimize_funcs_from_loss(
+        loss=-frozen_model.logp(jacobian=False),
+        inputs=frozen_model.continuous_value_vars + frozen_model.discrete_value_vars,
+        initial_point_dict=start_dict,
+        use_grad=use_grad,
+        use_hess=use_hess,
+        use_hessp=use_hessp,
+        gradient_backend=gradient_backend,
+        compile_kwargs=compile_kwargs,
+    )
+
+    args = optimizer_kwargs.pop("args", None)
+
+    # better_optimize.minimize will check if f_logp is a fused loss+grad Op, and automatically assign the jac argument
+    # if so. That is why it is not set here, regardless of user settings.
+    optimizer_result = minimize(
+        f=f_logp,
+        x0=cast(np.ndarray[float], initial_params.data),
+        args=args,
+        hess=f_hess,
+        hessp=f_hessp,
+        progressbar=progressbar,
+        method=method,
+        **optimizer_kwargs,
+    )
+
+    raveled_optimized = RaveledVars(optimizer_result.x, initial_params.point_map_info)
+    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed)
+    unobserved_vars_values = model.compile_fn(unobserved_vars, mode="FAST_COMPILE")(
+        DictToArrayBijection.rmap(raveled_optimized)
+    )
+
+    optimized_point = {
+        var.name: value for var, value in zip(unobserved_vars, unobserved_vars_values)
+    }
+
+    if return_raw:
+        return optimized_point, optimizer_result
+
+    return optimized_point

pymc_extras/inference/fit.py

@@ -0,0 +1,44 @@
+#   Copyright 2022 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.
+from importlib.util import find_spec
+
+
+def fit(method, **kwargs):
+    """
+    Fit a model with an inference algorithm
+
+    Parameters
+    ----------
+    method : str
+        Which inference method to run.
+        Supported: pathfinder or laplace
+
+    kwargs are passed on.
+
+    Returns
+    -------
+    arviz.InferenceData
+    """
+    if method == "pathfinder":
+        if find_spec("blackjax") is None:
+            raise RuntimeError("Need BlackJAX to use `pathfinder`")
+
+        from pymc_extras.inference.pathfinder import fit_pathfinder
+
+        return fit_pathfinder(**kwargs)
+
+    if method == "laplace":
+        from pymc_extras.inference.laplace import fit_laplace
+
+        return fit_laplace(**kwargs)