pymc-extras 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

pymc_extras/__init__.py

@@ -13,18 +13,17 @@
 #   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.find_map import find_MAP
-from pymc_extras.inference.fit import fit
-from pymc_extras.inference.laplace import fit_laplace
+from pymc_extras.inference import find_MAP, fit, fit_laplace, fit_pathfinder
 from pymc_extras.model.marginal.marginal_model import (
     MarginalModel,
     marginalize,
     recover_marginals,
 )
 from pymc_extras.model.model_api import as_model
-from pymc_extras.version import __version__
 
 _log = logging.getLogger("pmx")
 
@@ -33,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/distributions/__init__.py

@@ -26,6 +26,7 @@ from pymc_extras.distributions.discrete import (
 from pymc_extras.distributions.histogram_utils import histogram_approximation
 from pymc_extras.distributions.multivariate import R2D2M2CP
 from pymc_extras.distributions.timeseries import DiscreteMarkovChain
+from pymc_extras.distributions.transforms import PartialOrder
 
 __all__ = [
     "Chi",
@@ -37,4 +38,5 @@ __all__ = [
     "R2D2M2CP",
     "Skellam",
     "histogram_approximation",
+    "PartialOrder",
 ]

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/distributions/transforms/__init__.py

@@ -0,0 +1,3 @@
+from pymc_extras.distributions.transforms.partial_order import PartialOrder
+
+__all__ = ["PartialOrder"]

pymc_extras/distributions/transforms/partial_order.py

@@ -0,0 +1,227 @@
+#   Copyright 2025 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.
+import numpy as np
+import pytensor.tensor as pt
+
+from pymc.logprob.transforms import Transform
+
+__all__ = ["PartialOrder"]
+
+
+def dtype_minval(dtype):
+    """Find the minimum value for a given dtype"""
+    return np.iinfo(dtype).min if np.issubdtype(dtype, np.integer) else np.finfo(dtype).min
+
+
+def padded_where(x, to_len, padval=-1):
+    """A padded version of np.where"""
+    w = np.where(x)
+    return np.concatenate([w[0], np.full(to_len - len(w[0]), padval)])
+
+
+class PartialOrder(Transform):
+    """Create a PartialOrder transform
+
+    A more flexible version of the pymc ordered transform that
+    allows specifying a (strict) partial order on the elements.
+
+    Examples
+    --------
+    .. code:: python
+
+        import numpy as np
+        import pymc as pm
+        import pymc_extras as pmx
+
+        # Define two partial orders on 4 elements
+        # am[i,j] = 1 means i < j
+        adj_mats = np.array([
+            # 0 < {1, 2} < 3
+            [[0, 1, 1, 0],
+            [0, 0, 0, 1],
+            [0, 0, 0, 1],
+            [0, 0, 0, 0]],
+
+            # 1 < 0 < 3 < 2
+            [[0, 0, 0, 1],
+            [1, 0, 0, 0],
+            [0, 0, 0, 0],
+            [0, 0, 1, 0]],
+        ])
+
+        # Create the partial order from the adjacency matrices
+        po = pmx.PartialOrder(adj_mats)
+
+        with pm.Model() as model:
+            # Generate 3 samples from both partial orders
+            pm.Normal("po_vals", shape=(3,2,4), transform=po,
+                        initval=po.initvals((3,2,4)))
+
+            idata = pm.sample()
+
+        # Verify that for first po, the zeroth element is always the smallest
+        assert (idata.posterior['po_vals'][:,:,:,0,0] <
+                idata.posterior['po_vals'][:,:,:,0,1:]).all()
+
+        # Verify that for second po, the second element is always the largest
+        assert (idata.posterior['po_vals'][:,:,:,1,2] >=
+                idata.posterior['po_vals'][:,:,:,1,:]).all()
+
+    Technical notes
+    ----------------
+    Partial order needs to be strict, i.e. without equalities.
+    A DAG defining the partial order is sufficient, as transitive closure is automatically computed.
+    Code works in O(N*D) in runtime, but takes O(N^3) in initialization,
+    where N is the number of nodes in the dag and D is the maximum
+    in-degree of a node in the transitive reduction.
+    """
+
+    name = "partial_order"
+
+    def __init__(self, adj_mat):
+        """
+        Initialize the PartialOrder transform
+
+        Parameters
+        ----------
+        adj_mat: ndarray
+            adjacency matrix for the DAG that generates the partial order,
+            where ``adj_mat[i][j] = 1`` denotes ``i < j``.
+            Note this also accepts multiple DAGs if RV is multidimensional
+        """
+
+        # Basic input checks
+        if adj_mat.ndim < 2:
+            raise ValueError("Adjacency matrix must have at least 2 dimensions")
+        if adj_mat.shape[-2] != adj_mat.shape[-1]:
+            raise ValueError("Adjacency matrix is not square")
+        if adj_mat.min() != 0 or adj_mat.max() != 1:
+            raise ValueError("Adjacency matrix must contain only 0s and 1s")
+
+        # Create index over the first ellipsis dimensions
+        idx = np.ix_(*[np.arange(s) for s in adj_mat.shape[:-2]])
+
+        # Transitive closure using Floyd-Warshall
+        tc = adj_mat.astype(bool)
+        for k in range(tc.shape[-1]):
+            tc |= np.logical_and(tc[..., :, k, None], tc[..., None, k, :])
+
+        # Check if the dag is acyclic
+        if np.any(tc.diagonal(axis1=-2, axis2=-1)):
+            raise ValueError("Partial order contains equalities")
+
+        # Transitive reduction using the closure
+        # This gives the minimum description of the partial order
+        # This is to minmax the input degree
+        adj_mat = tc * (1 - np.matmul(tc, tc))
+
+        # Find the maximum in-degree of the reduced dag
+        dag_idim = adj_mat.sum(axis=-2).max()
+
+        # Topological sort
+        ts_inds = np.zeros(adj_mat.shape[:-1], dtype=int)
+        dm = adj_mat.copy()
+        for i in range(adj_mat.shape[1]):
+            assert dm.sum(axis=-2).min() == 0  # DAG is acyclic
+            nind = np.argmin(dm.sum(axis=-2), axis=-1)
+            dm[(*idx, slice(None), nind)] = 1  # Make nind not show up again
+            dm[(*idx, nind, slice(None))] = 0  # Allow it's children to show
+            ts_inds[(*idx, i)] = nind
+        self.ts_inds = ts_inds
+
+        # Change the dag to adjacency lists (with -1 for NA)
+        dag_T = np.apply_along_axis(padded_where, axis=-2, arr=adj_mat, padval=-1, to_len=dag_idim)
+        self.dag = np.swapaxes(dag_T, -2, -1)
+        self.is_start = np.all(self.dag[..., :, :] == -1, axis=-1)
+
+    def initvals(self, shape=None, lower=-1, upper=1):
+        """
+        Create a set of appropriate initial values for the variable.
+        NB! It is important that proper initial values are used,
+        as only properly ordered values are in the range of the transform.
+
+        Parameters
+        ----------
+        shape: tuple, default None
+            shape of the initial values. If None, adj_mat[:-1] is used
+        lower: float, default -1
+            lower bound for the initial values
+        upper: float, default 1
+            upper bound for the initial values
+
+        Returns
+        -------
+        vals: ndarray
+            initial values for the transformed variable
+        """
+
+        if shape is None:
+            shape = self.dag.shape[:-1]
+
+        if shape[-len(self.dag.shape[:-1]) :] != self.dag.shape[:-1]:
+            raise ValueError("Shape must match the shape of the adjacency matrix")
+
+        # Create the initial values
+        vals = np.linspace(lower, upper, self.dag.shape[-2])
+        inds = np.argsort(self.ts_inds, axis=-1)
+        ivals = vals[inds]
+
+        # Expand the initial values to the extra dimensions
+        extra_dims = shape[: -len(self.dag.shape[:-1])]
+        ivals = np.tile(ivals, extra_dims + tuple([1] * len(self.dag.shape[:-1])))
+
+        return ivals
+
+    def backward(self, value, *inputs):
+        minv = dtype_minval(value.dtype)
+        x = pt.concatenate(
+            [pt.zeros_like(value), pt.full(value.shape[:-1], minv)[..., None]], axis=-1
+        )
+
+        # Indices to allow broadcasting the max over the last dimension
+        idx = np.ix_(*[np.arange(s) for s in self.dag.shape[:-2]])
+        idx2 = tuple(np.tile(i[:, None], self.dag.shape[-1]) for i in idx)
+
+        # Has to be done stepwise as next steps depend on previous values
+        # Also has to be done in topological order, hence the ts_inds
+        for i in range(self.dag.shape[-2]):
+            tsi = self.ts_inds[..., i]
+            if len(tsi.shape) == 0:
+                tsi = int(tsi)  # if shape 0, it's a scalar
+            ni = (*idx, tsi)  # i-th node in topological order
+            eni = (Ellipsis, *ni)
+            ist = self.is_start[ni]
+
+            mval = pt.max(x[(Ellipsis, *idx2, self.dag[ni])], axis=-1)
+            x = pt.set_subtensor(x[eni], ist * value[eni] + (1 - ist) * (mval + pt.exp(value[eni])))
+        return x[..., :-1]
+
+    def forward(self, value, *inputs):
+        y = pt.zeros_like(value)
+
+        minv = dtype_minval(value.dtype)
+        vx = pt.concatenate([value, pt.full(value.shape[:-1], minv)[..., None]], axis=-1)
+
+        # Indices to allow broadcasting the max over the last dimension
+        idx = np.ix_(*[np.arange(s) for s in self.dag.shape[:-2]])
+        idx = tuple(np.tile(i[:, None, None], self.dag.shape[-2:]) for i in idx)
+
+        y = self.is_start * value + (1 - self.is_start) * (
+            pt.log(value - pt.max(vx[(Ellipsis, *idx, self.dag[..., :])], axis=-1))
+        )
+
+        return y
+
+    def log_jac_det(self, value, *inputs):
+        return pt.sum(value * (1 - self.is_start), axis=-1)

pymc_extras/inference/__init__.py

@@ -12,7 +12,9 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-
+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.inference.pathfinder.pathfinder import fit_pathfinder
 
-__all__ = ["fit"]
+__all__ = ["fit", "fit_pathfinder", "fit_laplace", "find_MAP"]

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/fit.py

@@ -11,11 +11,13 @@
 #   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.
+import arviz as az
 
 
-def fit(method, **kwargs):
+def fit(method: str, **kwargs) -> az.InferenceData:
     """
-    Fit a model with an inference algorithm
+    Fit a model with an inference algorithm.
+    See :func:`fit_pathfinder` and :func:`fit_laplace` for more details.
 
     Parameters
     ----------
@@ -23,11 +25,11 @@ def fit(method, **kwargs):
         Which inference method to run.
         Supported: pathfinder or laplace
 
-    kwargs are passed on.
+    kwargs: keyword arguments are passed on to the inference method.
 
     Returns
     -------
-    arviz.InferenceData
+    :class:`~arviz.InferenceData`
     """
     if method == "pathfinder":
         from pymc_extras.inference.pathfinder import fit_pathfinder

pymc_extras/inference/laplace.py

@@ -377,7 +377,10 @@ def sample_laplace_posterior(
     posterior_dist = stats.multivariate_normal(
         mean=mu.data, cov=H_inv, allow_singular=True, seed=rng
     )
+
     posterior_draws = posterior_dist.rvs(size=(chains, draws))
+    if mu.data.shape == (1,):
+        posterior_draws = np.expand_dims(posterior_draws, -1)
 
     if transform_samples:
         constrained_rvs, unconstrained_vector = _unconstrained_vector_to_constrained_rvs(model)
@@ -413,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,
@@ -446,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``.
@@ -497,16 +503,16 @@ 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.
 
     Returns
     -------
-    idata: az.InferenceData
+    :class:`~arviz.InferenceData`
         An InferenceData object containing the approximated posterior samples.
 
     Examples