pymc-extras 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

pymc_extras/__init__.py

@@ -13,6 +13,8 @@
 #   limitations under the License.
 import logging
 
+from importlib.metadata import version
+
 from pymc_extras import gp, statespace, utils
 from pymc_extras.distributions import *
 from pymc_extras.inference import find_MAP, fit, fit_laplace, fit_pathfinder
@@ -22,7 +24,6 @@ from pymc_extras.model.marginal.marginal_model import (
     recover_marginals,
 )
 from pymc_extras.model.model_api import as_model
-from pymc_extras.version import __version__
 
 _log = logging.getLogger("pmx")
 
@@ -31,3 +32,6 @@ if not logging.root.handlers:
     if len(_log.handlers) == 0:
         handler = logging.StreamHandler()
         _log.addHandler(handler)
+
+
+__version__ = version("pymc-extras")

pymc_extras/deserialize.py

@@ -0,0 +1,224 @@
+"""Deserialize dictionaries into Python objects.
+
+This is a two step process:
+
+1. Determine if the data is of the correct type.
+2. Deserialize the data into a python object.
+
+Examples
+--------
+Make use of the already registered deserializers:
+
+.. code-block:: python
+
+    from pymc_extras.deserialize import deserialize
+
+    prior_class_data = {
+        "dist": "Normal",
+        "kwargs": {"mu": 0, "sigma": 1}
+    }
+    prior = deserialize(prior_class_data)
+    # Prior("Normal", mu=0, sigma=1)
+
+Register custom class deserialization:
+
+.. code-block:: python
+
+    from pymc_extras.deserialize import register_deserialization
+
+    class MyClass:
+        def __init__(self, value: int):
+            self.value = value
+
+        def to_dict(self) -> dict:
+            # Example of what the to_dict method might look like.
+            return {"value": self.value}
+
+    register_deserialization(
+        is_type=lambda data: data.keys() == {"value"} and isinstance(data["value"], int),
+        deserialize=lambda data: MyClass(value=data["value"]),
+    )
+
+Deserialize data into that custom class:
+
+.. code-block:: python
+
+    from pymc_extras.deserialize import deserialize
+
+    data = {"value": 42}
+    obj = deserialize(data)
+    assert isinstance(obj, MyClass)
+
+
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+IsType = Callable[[Any], bool]
+Deserialize = Callable[[Any], Any]
+
+
+@dataclass
+class Deserializer:
+    """Object to store information required for deserialization.
+
+    All deserializers should be stored via the :func:`register_deserialization` function
+    instead of creating this object directly.
+
+    Attributes
+    ----------
+    is_type : IsType
+        Function to determine if the data is of the correct type.
+    deserialize : Deserialize
+        Function to deserialize the data.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        from typing import Any
+
+        class MyClass:
+            def __init__(self, value: int):
+                self.value = value
+
+        from pymc_extras.deserialize import Deserializer
+
+        def is_type(data: Any) -> bool:
+            return data.keys() == {"value"} and isinstance(data["value"], int)
+
+        def deserialize(data: dict) -> MyClass:
+            return MyClass(value=data["value"])
+
+        deserialize_logic = Deserializer(is_type=is_type, deserialize=deserialize)
+
+    """
+
+    is_type: IsType
+    deserialize: Deserialize
+
+
+DESERIALIZERS: list[Deserializer] = []
+
+
+class DeserializableError(Exception):
+    """Error raised when data cannot be deserialized."""
+
+    def __init__(self, data: Any):
+        self.data = data
+        super().__init__(
+            f"Couldn't deserialize {data}. Use register_deserialization to add a deserialization mapping."
+        )
+
+
+def deserialize(data: Any) -> Any:
+    """Deserialize a dictionary into a Python object.
+
+    Use the :func:`register_deserialization` function to add custom deserializations.
+
+    Deserialization is a two step process due to the dynamic nature of the data:
+
+    1. Determine if the data is of the correct type.
+    2. Deserialize the data into a Python object.
+
+    Each registered deserialization is checked in order until one is found that can
+    deserialize the data. If no deserialization is found, a :class:`DeserializableError` is raised.
+
+    A :class:`DeserializableError` is raised when the data fails to be deserialized
+    by any of the registered deserializers.
+
+    Parameters
+    ----------
+    data : Any
+        The data to deserialize.
+
+    Returns
+    -------
+    Any
+        The deserialized object.
+
+    Raises
+    ------
+    DeserializableError
+        Raised when the data doesn't match any registered deserializations
+        or fails to be deserialized.
+
+    Examples
+    --------
+    Deserialize a :class:`pymc_extras.prior.Prior` object:
+
+    .. code-block:: python
+
+        from pymc_extras.deserialize import deserialize
+
+        data = {"dist": "Normal", "kwargs": {"mu": 0, "sigma": 1}}
+        prior = deserialize(data)
+        # Prior("Normal", mu=0, sigma=1)
+
+    """
+    for mapping in DESERIALIZERS:
+        try:
+            is_type = mapping.is_type(data)
+        except Exception:
+            is_type = False
+
+        if not is_type:
+            continue
+
+        try:
+            return mapping.deserialize(data)
+        except Exception as e:
+            raise DeserializableError(data) from e
+    else:
+        raise DeserializableError(data)
+
+
+def register_deserialization(is_type: IsType, deserialize: Deserialize) -> None:
+    """Register an arbitrary deserialization.
+
+    Use the :func:`deserialize` function to then deserialize data using all registered
+    deserialize functions.
+
+    Parameters
+    ----------
+    is_type : Callable[[Any], bool]
+        Function to determine if the data is of the correct type.
+    deserialize : Callable[[dict], Any]
+        Function to deserialize the data of that type.
+
+    Examples
+    --------
+    Register a custom class deserialization:
+
+    .. code-block:: python
+
+        from pymc_extras.deserialize import register_deserialization
+
+        class MyClass:
+            def __init__(self, value: int):
+                self.value = value
+
+            def to_dict(self) -> dict:
+                # Example of what the to_dict method might look like.
+                return {"value": self.value}
+
+        register_deserialization(
+            is_type=lambda data: data.keys() == {"value"} and isinstance(data["value"], int),
+            deserialize=lambda data: MyClass(value=data["value"]),
+        )
+
+    Use that custom class deserialization:
+
+    .. code-block:: python
+
+        from pymc_extras.deserialize import deserialize
+
+        data = {"value": 42}
+        obj = deserialize(data)
+        assert isinstance(obj, MyClass)
+
+    """
+    mapping = Deserializer(is_type=is_type, deserialize=deserialize)
+    DESERIALIZERS.append(mapping)

pymc_extras/distributions/continuous.py

@@ -81,7 +81,7 @@ class GenExtreme(Continuous):
 
         \left\{x: 1 + \xi\left(\frac{x-\mu}{\sigma}\right) > 0 \right\}.
 
-    Note that this parametrization is per Coles (2001), and differs from that of
+    Note that this parametrization is per Coles (2001) [1]_, and differs from that of
     Scipy in the sign of the shape parameter, :math:`\xi`.
 
     .. plot::
@@ -132,7 +132,7 @@ class GenExtreme(Continuous):
 
     References
     ----------
-    .. [Coles2001] Coles, S.G. (2001).
+    .. [1] Coles, S.G. (2001).
         An Introduction to the Statistical Modeling of Extreme Values
         Springer-Verlag, London
 
@@ -260,6 +260,7 @@ class Chi:
     Examples
     --------
     .. code-block:: python
+
         import pymc as pm
         from pymc_extras.distributions import Chi
 

pymc_extras/distributions/discrete.py

@@ -116,6 +116,7 @@ class GeneralizedPoisson(pm.distributions.Discrete):
 
     .. math:: f(x \mid \mu, \lambda) =
                   \frac{\mu (\mu + \lambda x)^{x-1} e^{-\mu - \lambda x}}{x!}
+
     ========  ======================================
     Support   :math:`x \in \mathbb{N}_0`
     Mean      :math:`\frac{\mu}{1 - \lambda}`
@@ -135,9 +136,10 @@ class GeneralizedPoisson(pm.distributions.Discrete):
     When lam < 0, the mean is greater than the variance (underdispersion).
     When lam > 0, the mean is less than the variance (overdispersion).
 
+    The PMF is taken from [1]_ and the random generator function is adapted from [2]_.
+
     References
     ----------
-    The PMF is taken from [1] and the random generator function is adapted from [2].
     .. [1] Consul, PoC, and Felix Famoye. "Generalized Poisson regression model."
        Communications in Statistics-Theory and Methods 21.1 (1992): 89-109.
     .. [2] Famoye, Felix. "Generalized Poisson random variate generation." American

pymc_extras/inference/find_map.py

@@ -9,7 +9,7 @@ import pymc as pm
 import pytensor
 import pytensor.tensor as pt
 
-from better_optimize import minimize
+from better_optimize import basinhopping, 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
@@ -146,7 +146,7 @@ def _compile_grad_and_hess_to_jax(
     orig_loss_fn = f_loss.vm.jit_fn
 
     @jax.jit
-    def loss_fn_jax_grad(x, *shared):
+    def loss_fn_jax_grad(x):
         return jax.value_and_grad(lambda x: orig_loss_fn(x)[0])(x)
 
     f_loss_and_grad = loss_fn_jax_grad
@@ -301,6 +301,14 @@ def scipy_optimize_funcs_from_loss(
         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
+    # away.
+    if use_jax_gradients:
+        from pymc.sampling.jax import _replace_shared_variables
+
+        [loss] = _replace_shared_variables([loss])
+
     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
@@ -327,7 +335,7 @@ def scipy_optimize_funcs_from_loss(
 
 
 def find_MAP(
-    method: minimize_method,
+    method: minimize_method | Literal["basinhopping"],
     *,
     model: pm.Model | None = None,
     use_grad: bool | None = None,
@@ -344,14 +352,17 @@ def find_MAP(
     **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.
+    Fit a PyMC model via maximum a posteriori (MAP) estimation using JAX and scipy.optimize.
 
     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.
+        The optimization method to use. Valid choices are: Nelder-Mead, Powell, CG, BFGS, L-BFGS-B, TNC, SLSQP,
+        trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov, and basinhopping.
+
+        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``.
@@ -379,7 +390,9 @@ def find_MAP(
     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.
+        Additional keyword arguments to pass to the ``scipy.optimize`` function being used. Unless
+        ``method = "basinhopping"``, ``scipy.optimize.minimize`` will be used. For ``basinhopping``,
+        ``scipy.optimize.basinhopping`` will be used. See the documentation of these functions for details.
 
     Returns
     -------
@@ -405,6 +418,18 @@ def find_MAP(
     initial_params = DictToArrayBijection.map(
         {var_name: value for var_name, value in start_dict.items() if var_name in vars_dict}
     )
+
+    do_basinhopping = method == "basinhopping"
+    minimizer_kwargs = optimizer_kwargs.pop("minimizer_kwargs", {})
+
+    if do_basinhopping:
+        # For a nice API, we let the user set method="basinhopping", but if we're doing basinhopping we still need
+        # another method for the inner optimizer. This will be set in the minimizer_kwargs, but also needs a default
+        # if one isn't provided.
+
+        method = minimizer_kwargs.pop("method", "L-BFGS-B")
+        minimizer_kwargs["method"] = method
+
     use_grad, use_hess, use_hessp = set_optimizer_function_defaults(
         method, use_grad, use_hess, use_hessp
     )
@@ -423,17 +448,37 @@ def find_MAP(
     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,
-    )
+    # if so. That is why the jac argument is not passed here in either branch.
+
+    if do_basinhopping:
+        if "args" not in minimizer_kwargs:
+            minimizer_kwargs["args"] = args
+        if "hess" not in minimizer_kwargs:
+            minimizer_kwargs["hess"] = f_hess
+        if "hessp" not in minimizer_kwargs:
+            minimizer_kwargs["hessp"] = f_hessp
+        if "method" not in minimizer_kwargs:
+            minimizer_kwargs["method"] = method
+
+        optimizer_result = basinhopping(
+            func=f_logp,
+            x0=cast(np.ndarray[float], initial_params.data),
+            progressbar=progressbar,
+            minimizer_kwargs=minimizer_kwargs,
+            **optimizer_kwargs,
+        )
+
+    else:
+        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)

pymc_extras/inference/laplace.py

@@ -416,7 +416,7 @@ def sample_laplace_posterior(
 
 
 def fit_laplace(
-    optimize_method: minimize_method = "BFGS",
+    optimize_method: minimize_method | Literal["basinhopping"] = "BFGS",
     *,
     model: pm.Model | None = None,
     use_grad: bool | None = None,
@@ -449,8 +449,11 @@ def fit_laplace(
     ----------
     model : pm.Model
         The PyMC model to be fit. If None, the current model context is used.
-    optimize_method : str
-        The optimization method to use. See scipy.optimize.minimize documentation for details.
+    method : str
+        The optimization method to use. Valid choices are: Nelder-Mead, Powell, CG, BFGS, L-BFGS-B, TNC, SLSQP,
+        trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov, and basinhopping.
+
+        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``.
@@ -500,10 +503,10 @@ def fit_laplace(
     diag_jitter: float | None
         A small value added to the diagonal of the inverse Hessian matrix to ensure it is positive semi-definite.
         If None, no jitter is added. Default is 1e-8.
-    optimizer_kwargs: dict, optional
-        Additional keyword arguments to pass to scipy.minimize. See the documentation for scipy.optimize.minimize for
-        details. Arguments that are typically passed via ``options`` will be automatically extracted without the need
-        to use a nested dictionary.
+    optimizer_kwargs
+        Additional keyword arguments to pass to the ``scipy.optimize`` function being used. Unless
+        ``method = "basinhopping"``, ``scipy.optimize.minimize`` will be used. For ``basinhopping``,
+        ``scipy.optimize.basinhopping`` will be used. See the documentation of these functions for details.
     compile_kwargs: dict, optional
         Additional keyword arguments to pass to pytensor.function.