pymc-extras 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

pymc_extras/__init__.py

@@ -15,7 +15,9 @@ import logging
 
 from pymc_extras import gp, statespace, utils
 from pymc_extras.distributions import *
+from pymc_extras.inference.find_map import find_MAP
 from pymc_extras.inference.fit import fit
+from pymc_extras.inference.laplace import fit_laplace
 from pymc_extras.model.marginal.marginal_model import (
     MarginalModel,
     marginalize,

pymc_extras/inference/find_map.py

@@ -1,9 +1,9 @@
 import logging
 
 from collections.abc import Callable
+from importlib.util import find_spec
 from typing import Literal, cast, get_args
 
-import jax
 import numpy as np
 import pymc as pm
 import pytensor
@@ -30,13 +30,29 @@ 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:
+        _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
 
 
@@ -59,7 +75,7 @@ def get_nearest_psd(A: np.ndarray) -> np.ndarray:
         The nearest positive semi-definite matrix to the input matrix.
     """
     C = (A + A.T) / 2
-    eigval, eigvec = np.linalg.eig(C)
+    eigval, eigvec = np.linalg.eigh(C)
     eigval[eigval < 0] = 0
 
     return eigvec @ np.diag(eigval) @ eigvec.T
@@ -97,7 +113,7 @@ def _create_transformed_draws(H_inv, slices, out_shapes, posterior_draws, model,
     return f_untransform(posterior_draws)
 
 
-def _compile_jax_gradients(
+def _compile_grad_and_hess_to_jax(
     f_loss: Function, use_hess: bool, use_hessp: bool
 ) -> tuple[Callable | None, Callable | None]:
     """
@@ -122,6 +138,8 @@ def _compile_jax_gradients(
     f_hessp: Callable | None
         The compiled hessian-vector product function, or None if use_hessp is False.
     """
+    import jax
+
     f_hess = None
     f_hessp = None
 
@@ -152,7 +170,7 @@ def _compile_jax_gradients(
     return f_loss_and_grad, f_hess, f_hessp
 
 
-def _compile_functions(
+def _compile_functions_for_scipy_optimize(
     loss: TensorVariable,
     inputs: list[TensorVariable],
     compute_grad: bool,
@@ -177,7 +195,7 @@ def _compile_functions(
     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.
+        Additional keyword arguments to pass to the ``pm.compile`` function.
 
     Returns
     -------
@@ -193,19 +211,19 @@ def _compile_functions(
     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)
+        f_loss_and_grad = pm.compile(inputs, [loss, grad], **compile_kwargs)
     else:
-        f_loss = pm.compile_pymc(inputs, loss, **compile_kwargs)
+        f_loss = pm.compile(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)
+        f_hess = pm.compile(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)
+        f_hessp = pm.compile([*inputs, p], hessp[0], **compile_kwargs)
 
     return [f_loss_and_grad, f_hess, f_hessp]
 
@@ -240,7 +258,7 @@ def scipy_optimize_funcs_from_loss(
     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.
+        Additional keyword arguments to pass to the ``pm.compile`` function.
 
     Returns
     -------
@@ -265,6 +283,8 @@ def scipy_optimize_funcs_from_loss(
         )
 
     use_jax_gradients = (gradient_backend == "jax") and use_grad
+    if use_jax_gradients and not find_spec("jax"):
+        raise ImportError("JAX must be installed to use JAX gradients")
 
     mode = compile_kwargs.get("mode", None)
     if mode is None and use_jax_gradients:
@@ -285,7 +305,7 @@ def scipy_optimize_funcs_from_loss(
     compute_hess = use_hess and not use_jax_gradients
     compute_hessp = use_hessp and not use_jax_gradients
 
-    funcs = _compile_functions(
+    funcs = _compile_functions_for_scipy_optimize(
         loss=loss,
         inputs=[flat_input],
         compute_grad=compute_grad,
@@ -301,7 +321,7 @@ def scipy_optimize_funcs_from_loss(
 
     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)
+        f_loss, f_hess, f_hessp = _compile_grad_and_hess_to_jax(f_loss, use_hess, use_hessp)
 
     return f_loss, f_hess, f_hessp
 

pymc_extras/inference/fit.py

@@ -11,7 +11,6 @@
 #   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):
@@ -31,9 +30,6 @@ def fit(method, **kwargs):
     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)

pymc_extras/inference/laplace.py

@@ -16,6 +16,7 @@
 import logging
 
 from functools import reduce
+from importlib.util import find_spec
 from itertools import product
 from typing import Literal
 
@@ -231,7 +232,7 @@ def add_data_to_inferencedata(
     return idata
 
 
-def fit_mvn_to_MAP(
+def fit_mvn_at_MAP(
     optimized_point: dict[str, np.ndarray],
     model: pm.Model | None = None,
     on_bad_cov: Literal["warn", "error", "ignore"] = "ignore",
@@ -276,6 +277,9 @@ def fit_mvn_to_MAP(
     inverse_hessian: np.ndarray
         The inverse Hessian matrix of the log-posterior evaluated at the MAP estimate.
     """
+    if gradient_backend == "jax" and not find_spec("jax"):
+        raise ImportError("JAX must be installed to use JAX gradients")
+
     model = pm.modelcontext(model)
     compile_kwargs = {} if compile_kwargs is None else compile_kwargs
     frozen_model = freeze_dims_and_data(model)
@@ -344,8 +348,10 @@ def sample_laplace_posterior(
 
     Parameters
     ----------
-    mu
-    H_inv
+    mu: RaveledVars
+        The MAP estimate of the model parameters.
+    H_inv: np.ndarray
+        The inverse Hessian matrix of the log-posterior evaluated at the MAP estimate.
     model : Model
         A PyMC model
     chains : int
@@ -384,9 +390,7 @@ def sample_laplace_posterior(
             constrained_rvs, replace={unconstrained_vector: batched_values}
         )
 
-        f_constrain = pm.compile_pymc(
-            inputs=[batched_values], outputs=batched_rvs, **compile_kwargs
-        )
+        f_constrain = pm.compile(inputs=[batched_values], outputs=batched_rvs, **compile_kwargs)
         posterior_draws = f_constrain(posterior_draws)
 
     else:
@@ -472,15 +476,17 @@ def fit_laplace(
         and 1).
 
         .. warning::
-            This argumnet should be considered highly experimental. It has not been verified if this method produces
+            This argument should be considered highly experimental. It has not been verified if this method produces
             valid draws from the posterior. **Use at your own risk**.
 
     gradient_backend: str, default "pytensor"
         The backend to use for gradient computations. Must be one of "pytensor" or "jax".
     chains: int, default: 2
-        The number of sampling chains running in parallel.
+        The number of chain dimensions to sample. Note that this is *not* the number of chains to run in parallel,
+        because the Laplace approximation is not an MCMC method. This argument exists to ensure that outputs are
+        compatible with the ArviZ library.
     draws: int, default: 500
-        The number of samples to draw from the approximated posterior.
+        The number of samples to draw from the approximated posterior. Totals samples will be chains * draws.
     on_bad_cov : str, one of 'ignore', 'warn', or 'error', default: 'ignore'
         What to do when ``H_inv`` (inverse Hessian) is not positive semi-definite.
         If 'ignore' or 'warn', the closest positive-semi-definite matrix to ``H_inv`` (in L1 norm) will be returned.
@@ -547,11 +553,12 @@ def fit_laplace(
         **optimizer_kwargs,
     )
 
-    mu, H_inv = fit_mvn_to_MAP(
+    mu, H_inv = fit_mvn_at_MAP(
         optimized_point=optimized_point,
         model=model,
         on_bad_cov=on_bad_cov,
         transform_samples=fit_in_unconstrained_space,
+        gradient_backend=gradient_backend,
         zero_tol=zero_tol,
         diag_jitter=diag_jitter,
         compile_kwargs=compile_kwargs,

pymc_extras/inference/pathfinder/__init__.py

@@ -0,0 +1,3 @@
+from pymc_extras.inference.pathfinder.pathfinder import fit_pathfinder
+
+__all__ = ["fit_pathfinder"]

pymc_extras/inference/pathfinder/importance_sampling.py

@@ -0,0 +1,139 @@
+import logging
+import warnings as _warnings
+
+from dataclasses import dataclass, field
+from typing import Literal
+
+import arviz as az
+import numpy as np
+
+from numpy.typing import NDArray
+from scipy.special import logsumexp
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ImportanceSamplingResult:
+    """container for importance sampling results"""
+
+    samples: NDArray
+    pareto_k: float | None = None
+    warnings: list[str] = field(default_factory=list)
+    method: str = "none"
+
+
+def importance_sampling(
+    samples: NDArray,
+    logP: NDArray,
+    logQ: NDArray,
+    num_draws: int,
+    method: Literal["psis", "psir", "identity", "none"] | None,
+    random_seed: int | None = None,
+) -> ImportanceSamplingResult:
+    """Pareto Smoothed Importance Resampling (PSIR)
+    This implements the Pareto Smooth Importance Resampling (PSIR) method, as described in Algorithm 5 of Zhang et al. (2022). The PSIR follows a similar approach to Algorithm 1 PSIS diagnostic from Yao et al., (2018). However, before computing the the importance ratio r_s, the logP and logQ are adjusted to account for the number multiple estimators (or paths). The process involves resampling from the original sample with replacement, with probabilities proportional to the computed importance weights from PSIS.
+
+    Parameters
+    ----------
+    samples : NDArray
+        samples from proposal distribution, shape (L, M, N)
+    logP : NDArray
+        log probability values of target distribution, shape (L, M)
+    logQ : NDArray
+        log probability values of proposal distribution, shape (L, M)
+    num_draws : int
+        number of draws to return where num_draws <= samples.shape[0]
+    method : str, optional
+        importance sampling method to use. Options are "psis" (default), "psir", "identity", "none. Pareto Smoothed Importance Sampling (psis) is recommended in many cases for more stable results than Pareto Smoothed Importance Resampling (psir). identity applies the log importance weights directly without resampling. none applies no importance sampling weights and returns the samples as is of size num_draws_per_path * num_paths.
+    random_seed : int | None
+
+    Returns
+    -------
+    ImportanceSamplingResult
+        importance sampled draws and other info based on the specified method
+
+    Future work!
+    ----------
+    - Implement the 3 sampling approaches and 5 weighting functions from Elvira et al. (2019)
+    - Implement Algorithm 2 VSBC marginal diagnostics from Yao et al. (2018)
+    - Incorporate these various diagnostics, sampling approaches and weighting functions into VI algorithms.
+
+    References
+    ----------
+    Elvira, V., Martino, L., Luengo, D., & Bugallo, M. F. (2019). Generalized Multiple Importance Sampling. Statistical Science, 34(1), 129-155. https://doi.org/10.1214/18-STS668
+
+    Yao, Y., Vehtari, A., Simpson, D., & Gelman, A. (2018). Yes, but Did It Work?: Evaluating Variational Inference. arXiv:1802.02538 [Stat]. http://arxiv.org/abs/1802.02538
+
+    Zhang, L., Carpenter, B., Gelman, A., & Vehtari, A. (2022). Pathfinder: Parallel quasi-Newton variational inference. Journal of Machine Learning Research, 23(306), 1-49.
+    """
+
+    warnings = []
+    num_paths, _, N = samples.shape
+
+    if method == "none":
+        warnings.append(
+            "Importance sampling is disabled. The samples are returned as is which may include samples from failed paths with non-finite logP or logQ values. It is recommended to use importance_sampling='psis' for better stability."
+        )
+        return ImportanceSamplingResult(samples=samples, warnings=warnings)
+    else:
+        samples = samples.reshape(-1, N)
+        logP = logP.ravel()
+        logQ = logQ.ravel()
+
+        # adjust log densities
+        log_I = np.log(num_paths)
+        logP -= log_I
+        logQ -= log_I
+        logiw = logP - logQ
+
+        with _warnings.catch_warnings():
+            _warnings.filterwarnings(
+                "ignore", category=RuntimeWarning, message="overflow encountered in exp"
+            )
+            if method == "psis":
+                replace = False
+                logiw, pareto_k = az.psislw(logiw)
+            elif method == "psir":
+                replace = True
+                logiw, pareto_k = az.psislw(logiw)
+            elif method == "identity":
+                replace = False
+                pareto_k = None
+            else:
+                raise ValueError(f"Invalid importance sampling method: {method}")
+
+    # NOTE: Pareto k is normally bad for Pathfinder even when the posterior is close to the NUTS posterior or closer to NUTS than ADVI.
+    # Pareto k may not be a good diagnostic for Pathfinder.
+    # TODO: Find replacement diagnostics for Pathfinder.
+
+    p = np.exp(logiw - logsumexp(logiw))
+    rng = np.random.default_rng(random_seed)
+
+    try:
+        resampled = rng.choice(samples, size=num_draws, replace=replace, p=p, shuffle=False, axis=0)
+        return ImportanceSamplingResult(
+            samples=resampled, pareto_k=pareto_k, warnings=warnings, method=method
+        )
+    except ValueError as e1:
+        if "Fewer non-zero entries in p than size" in str(e1):
+            num_nonzero = np.where(np.nonzero(p)[0], 1, 0).sum()
+            warnings.append(
+                f"Not enough valid samples: {num_nonzero} available out of {num_draws} requested. Switching to psir importance sampling."
+            )
+            try:
+                resampled = rng.choice(
+                    samples, size=num_draws, replace=True, p=p, shuffle=False, axis=0
+                )
+                return ImportanceSamplingResult(
+                    samples=resampled, pareto_k=pareto_k, warnings=warnings, method=method
+                )
+            except ValueError as e2:
+                logger.error(
+                    "Importance sampling failed even with psir importance sampling. "
+                    "This might indicate invalid probability weights or insufficient valid samples."
+                )
+                raise ValueError(
+                    "Importance sampling failed for both with and without replacement"
+                ) from e2
+        raise

pymc_extras/inference/pathfinder/lbfgs.py

@@ -0,0 +1,190 @@
+import logging
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from enum import Enum, auto
+
+import numpy as np
+
+from numpy.typing import NDArray
+from scipy.optimize import minimize
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class LBFGSHistory:
+    """History of LBFGS iterations."""
+
+    x: NDArray[np.float64]
+    g: NDArray[np.float64]
+    count: int
+
+    def __post_init__(self):
+        self.x = np.ascontiguousarray(self.x, dtype=np.float64)
+        self.g = np.ascontiguousarray(self.g, dtype=np.float64)
+
+
+@dataclass(slots=True)
+class LBFGSHistoryManager:
+    """manages and stores the history of lbfgs optimisation iterations.
+
+    Parameters
+    ----------
+    value_grad_fn : Callable
+        function that returns tuple of (value, gradient) given input x
+    x0 : NDArray
+        initial position
+    maxiter : int
+        maximum number of iterations to store
+    """
+
+    value_grad_fn: Callable[[NDArray[np.float64]], tuple[np.float64, NDArray[np.float64]]]
+    x0: NDArray[np.float64]
+    maxiter: int
+    x_history: NDArray[np.float64] = field(init=False)
+    g_history: NDArray[np.float64] = field(init=False)
+    count: int = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.x_history = np.empty((self.maxiter + 1, self.x0.shape[0]), dtype=np.float64)
+        self.g_history = np.empty((self.maxiter + 1, self.x0.shape[0]), dtype=np.float64)
+        self.count = 0
+
+        value, grad = self.value_grad_fn(self.x0)
+        if np.all(np.isfinite(grad)) and np.isfinite(value):
+            self.add_entry(self.x0, grad)
+
+    def add_entry(self, x: NDArray[np.float64], g: NDArray[np.float64]) -> None:
+        """adds new position and gradient to history.
+
+        Parameters
+        ----------
+        x : NDArray
+            position vector
+        g : NDArray
+            gradient vector
+        """
+        self.x_history[self.count] = x
+        self.g_history[self.count] = g
+        self.count += 1
+
+    def get_history(self) -> LBFGSHistory:
+        """returns history of optimisation iterations."""
+        return LBFGSHistory(
+            x=self.x_history[: self.count], g=self.g_history[: self.count], count=self.count
+        )
+
+    def __call__(self, x: NDArray[np.float64]) -> None:
+        value, grad = self.value_grad_fn(x)
+        if np.all(np.isfinite(grad)) and np.isfinite(value) and self.count < self.maxiter + 1:
+            self.add_entry(x, grad)
+
+
+class LBFGSStatus(Enum):
+    CONVERGED = auto()
+    MAX_ITER_REACHED = auto()
+    DIVERGED = auto()
+    # Statuses that lead to Exceptions:
+    INIT_FAILED = auto()
+    LBFGS_FAILED = auto()
+
+
+class LBFGSException(Exception):
+    DEFAULT_MESSAGE = "LBFGS failed."
+
+    def __init__(self, message=None, status: LBFGSStatus = LBFGSStatus.LBFGS_FAILED):
+        super().__init__(message or self.DEFAULT_MESSAGE)
+        self.status = status
+
+
+class LBFGSInitFailed(LBFGSException):
+    DEFAULT_MESSAGE = "LBFGS failed to initialise."
+
+    def __init__(self, message=None):
+        super().__init__(message or self.DEFAULT_MESSAGE, LBFGSStatus.INIT_FAILED)
+
+
+class LBFGS:
+    """L-BFGS optimizer wrapper around scipy's implementation.
+
+    Parameters
+    ----------
+    value_grad_fn : Callable
+        function that returns tuple of (value, gradient) given input x
+    maxcor : int
+        maximum number of variable metric corrections
+    maxiter : int, optional
+        maximum number of iterations, defaults to 1000
+    ftol : float, optional
+        function tolerance for convergence, defaults to 1e-5
+    gtol : float, optional
+        gradient tolerance for convergence, defaults to 1e-8
+    maxls : int, optional
+        maximum number of line search steps, defaults to 1000
+    """
+
+    def __init__(
+        self, value_grad_fn, maxcor, maxiter=1000, ftol=1e-5, gtol=1e-8, maxls=1000
+    ) -> None:
+        self.value_grad_fn = value_grad_fn
+        self.maxcor = maxcor
+        self.maxiter = maxiter
+        self.ftol = ftol
+        self.gtol = gtol
+        self.maxls = maxls
+
+    def minimize(self, x0) -> tuple[NDArray, NDArray, int, LBFGSStatus]:
+        """minimizes objective function starting from initial position.
+
+        Parameters
+        ----------
+        x0 : array_like
+            initial position
+
+        Returns
+        -------
+        x : NDArray
+            history of positions
+        g : NDArray
+            history of gradients
+        count : int
+            number of iterations
+        status : LBFGSStatus
+            final status of optimisation
+        """
+
+        x0 = np.array(x0, dtype=np.float64)
+
+        history_manager = LBFGSHistoryManager(
+            value_grad_fn=self.value_grad_fn, x0=x0, maxiter=self.maxiter
+        )
+
+        result = minimize(
+            self.value_grad_fn,
+            x0,
+            method="L-BFGS-B",
+            jac=True,
+            callback=history_manager,
+            options={
+                "maxcor": self.maxcor,
+                "maxiter": self.maxiter,
+                "ftol": self.ftol,
+                "gtol": self.gtol,
+                "maxls": self.maxls,
+            },
+        )
+        history = history_manager.get_history()
+
+        # warnings and suggestions for LBFGSStatus are displayed at the end
+        if result.status == 1:
+            lbfgs_status = LBFGSStatus.MAX_ITER_REACHED
+        elif (result.status == 2) or (history.count <= 1):
+            if result.nit <= 1:
+                lbfgs_status = LBFGSStatus.INIT_FAILED
+            elif result.fun == np.inf:
+                lbfgs_status = LBFGSStatus.DIVERGED
+        else:
+            lbfgs_status = LBFGSStatus.CONVERGED
+
+        return history.x, history.g, history.count, lbfgs_status