pymc-extras 0.5.0__tar.gz → 0.7.0__tar.gz

pymc_extras-0.7.0/CONTRIBUTING.md

@@ -0,0 +1,24 @@
+# Contributing guide
+
+Page in construction, for now go to https://github.com/pymc-devs/pymc-extras#questions.
+
+## Building the documentation
+
+To build the documentation locally, you need to install the necessary
+dependencies and then use `make` to build the HTML files.
+
+First, install the package with the optional documentation dependencies:
+
+```bash
+pip install ".[docs]"
+```
+
+Then, navigate to the `docs` directory and run `make html`:
+
+```bash
+cd docs
+make html
+```
+
+The generated HTML files will be in the `docs/_build/html` directory. You can
+open the `index.html` file in that directory to view the documentation.

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pymc-extras
-Version: 0.5.0
+Version: 0.7.0
 Summary: A home for new additions to PyMC, which may include unusual probability distribitions, advanced model fitting algorithms, or any code that may be inappropriate to include in the pymc repository, but may want to be made available to users.
 Project-URL: Documentation, https://pymc-extras.readthedocs.io/
 Project-URL: Repository, https://github.com/pymc-devs/pymc-extras.git
@@ -235,8 +235,8 @@ Requires-Python: >=3.11
 Requires-Dist: better-optimize>=0.1.5
 Requires-Dist: preliz>=0.20.0
 Requires-Dist: pydantic>=2.0.0
-Requires-Dist: pymc>=5.24.1
-Requires-Dist: pytensor>=2.31.4
+Requires-Dist: pymc>=5.27.0
+Requires-Dist: pytensor>=2.36.3
 Requires-Dist: scikit-learn
 Provides-Extra: complete
 Requires-Dist: dask[complete]<2025.1.1; extra == 'complete'
@@ -245,7 +245,7 @@ Provides-Extra: dask-histogram
 Requires-Dist: dask[complete]<2025.1.1; extra == 'dask-histogram'
 Requires-Dist: xhistogram; extra == 'dask-histogram'
 Provides-Extra: dev
-Requires-Dist: blackjax; extra == 'dev'
+Requires-Dist: blackjax>=0.12; extra == 'dev'
 Requires-Dist: dask[all]<2025.1.1; extra == 'dev'
 Requires-Dist: pytest-mock; extra == 'dev'
 Requires-Dist: pytest>=6.0; extra == 'dev'

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.5.0'
-__version_tuple__ = version_tuple = (0, 5, 0)
+__version__ = version = '0.7.0'
+__version_tuple__ = version_tuple = (0, 7, 0)
 
 __commit_id__ = commit_id = None

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/conda-envs/environment-test.yml

@@ -1,10 +1,10 @@
-name: pymc-extras-test
+name: pymc-extras
 channels:
 - conda-forge
 - nodefaults
 dependencies:
-  - pymc>=5.24.1
-  - pytensor>=2.31.4
+  - pymc>=5.27.0
+  - pytensor>=2.36.3
   - scikit-learn
   - better-optimize>=0.1.5
   - dask<2025.1.1

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/docs/api_reference.rst

@@ -56,6 +56,7 @@ Prior
    create_dim_handler
    handle_dims
    Prior
+   register_tensor_transform
    VariableFactory
    sample_prior
    Censored

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/docs/statespace/models.rst

@@ -6,7 +6,8 @@ Statespace Models
 .. autosummary::
    :toctree: generated
 
-    BayesianSARIMA
+    BayesianETS
+    BayesianSARIMAX
     BayesianVARMAX
 
 *********************

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/pymc_extras/deserialize.py

@@ -13,10 +13,7 @@ Make use of the already registered deserializers:
 
     from pymc_extras.deserialize import deserialize
 
-    prior_class_data = {
-        "dist": "Normal",
-        "kwargs": {"mu": 0, "sigma": 1}
-    }
+    prior_class_data = {"dist": "Normal", "kwargs": {"mu": 0, "sigma": 1}}
     prior = deserialize(prior_class_data)
     # Prior("Normal", mu=0, sigma=1)
 
@@ -26,6 +23,7 @@ Register custom class deserialization:
 
     from pymc_extras.deserialize import register_deserialization
 
+
     class MyClass:
         def __init__(self, value: int):
             self.value = value
@@ -34,6 +32,7 @@ Register custom class deserialization:
             # 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"]),
@@ -80,18 +79,23 @@ class Deserializer:
 
         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)
 
     """
@@ -196,6 +200,7 @@ def register_deserialization(is_type: IsType, deserialize: Deserialize) -> None:
 
         from pymc_extras.deserialize import register_deserialization
 
+
         class MyClass:
             def __init__(self, value: int):
                 self.value = value
@@ -204,6 +209,7 @@ def register_deserialization(is_type: IsType, deserialize: Deserialize) -> None:
                 # 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"]),

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/pymc_extras/distributions/continuous.py

@@ -265,7 +265,7 @@ class Chi:
         from pymc_extras.distributions import Chi
 
         with pm.Model():
-            x = Chi('x', nu=1)
+            x = Chi("x", nu=1)
     """
 
     @staticmethod

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/pymc_extras/distributions/histogram_utils.py

@@ -130,8 +130,7 @@ def histogram_approximation(name, dist, *, observed, **h_kwargs):
     ...     m = pm.Normal("m", dims="tests")
     ...     s = pm.LogNormal("s", dims="tests")
     ...     pot = pmx.distributions.histogram_approximation(
-    ...         "pot", pm.Normal.dist(m, s),
-    ...         observed=measurements, n_quantiles=50
+    ...         "pot", pm.Normal.dist(m, s), observed=measurements, n_quantiles=50
     ...     )
 
     For special cases like Zero Inflation in Continuous variables there is a flag.
@@ -143,8 +142,11 @@ def histogram_approximation(name, dist, *, observed, **h_kwargs):
     ...     m = pm.Normal("m", dims="tests")
     ...     s = pm.LogNormal("s", dims="tests")
     ...     pot = pmx.distributions.histogram_approximation(
-    ...         "pot", pm.Normal.dist(m, s),
-    ...         observed=measurements, n_quantiles=50, zero_inflation=True
+    ...         "pot",
+    ...         pm.Normal.dist(m, s),
+    ...         observed=measurements,
+    ...         n_quantiles=50,
+    ...         zero_inflation=True,
     ...     )
     """
     try:

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/pymc_extras/distributions/multivariate/r2d2m2cp.py

@@ -305,6 +305,7 @@ def R2D2M2CP(
         import pymc_extras as pmx
         import pymc as pm
         import numpy as np
+
         X = np.random.randn(10, 3)
         b = np.random.randn(3)
         y = X @ b + np.random.randn(10) * 0.04 + 5
@@ -339,7 +340,7 @@ def R2D2M2CP(
                 # "c" - a must have in the relation
                 variables_importance=[10, 1, 34],
                 # NOTE: try both
-                centered=True
+                centered=True,
             )
             # intercept prior centering should be around prior predictive mean
             intercept = y.mean()
@@ -365,7 +366,7 @@ def R2D2M2CP(
                 r2_std=0.2,
                 # NOTE: if you know where a variable should go
                 # if you do not know, leave as 0.5
-                centered=False
+                centered=False,
             )
             # intercept prior centering should be around prior predictive mean
             intercept = y.mean()
@@ -394,7 +395,7 @@ def R2D2M2CP(
                 # if you do not know, leave as 0.5
                 positive_probs=[0.8, 0.5, 0.1],
                 # NOTE: try both
-                centered=True
+                centered=True,
             )
             intercept = y.mean()
             obs = pm.Normal("obs", intercept + X @ beta, eps, observed=y)

{pymc_extras-0.5.0 → pymc_extras-0.7.0}/pymc_extras/distributions/timeseries.py

@@ -113,8 +113,10 @@ class DiscreteMarkovChain(Distribution):
 
         with pm.Model() as markov_chain:
             P = pm.Dirichlet("P", a=[1, 1, 1], size=(3,))
-            init_dist = pm.Categorical.dist(p = np.full(3, 1 / 3))
-            markov_chain = pmx.DiscreteMarkovChain("markov_chain", P=P, init_dist=init_dist, shape=(100,))
+            init_dist = pm.Categorical.dist(p=np.full(3, 1 / 3))
+            markov_chain = pmx.DiscreteMarkovChain(
+                "markov_chain", P=P, init_dist=init_dist, shape=(100,)
+            )
 
     """
 
@@ -194,21 +196,20 @@ class DiscreteMarkovChain(Distribution):
         state_rng = pytensor.shared(np.random.default_rng())
 
         def transition(*args):
-            *states, transition_probs, old_rng = args
+            old_rng, *states, transition_probs = args
             p = transition_probs[tuple(states)]
             next_rng, next_state = pm.Categorical.dist(p=p, rng=old_rng).owner.outputs
-            return next_state, {old_rng: next_rng}
+            return next_rng, next_state
 
-        markov_chain, state_updates = pytensor.scan(
+        state_next_rng, markov_chain = pytensor.scan(
             transition,
-            non_sequences=[P_, state_rng],
-            outputs_info=_make_outputs_info(n_lags, init_dist_),
+            outputs_info=[state_rng, *_make_outputs_info(n_lags, init_dist_)],
+            non_sequences=[P_],
             n_steps=steps_,
             strict=True,
+            return_updates=False,
         )
 
-        (state_next_rng,) = tuple(state_updates.values())
-
         discrete_mc_ = pt.moveaxis(pt.concatenate([init_dist_, markov_chain], axis=0), 0, -1)
 
         discrete_mc_op = DiscreteMarkovChainRV(
@@ -237,16 +238,17 @@ def discrete_mc_moment(op, rv, P, steps, init_dist, state_rng):
     n_lags = op.n_lags
 
     def greedy_transition(*args):
-        *states, transition_probs, old_rng = args
+        *states, transition_probs = args
         p = transition_probs[tuple(states)]
         return pt.argmax(p)
 
-    chain_moment, moment_updates = pytensor.scan(
+    chain_moment = pytensor.scan(
         greedy_transition,
-        non_sequences=[P, state_rng],
+        non_sequences=[P],
         outputs_info=_make_outputs_info(n_lags, init_dist),
         n_steps=steps,
         strict=True,
+        return_updates=False,
     )
     chain_moment = pt.concatenate([init_dist_moment, chain_moment])
     return chain_moment

pymc_extras-0.7.0/pymc_extras/inference/dadvi/dadvi.py

@@ -0,0 +1,282 @@
+import arviz as az
+import numpy as np
+import pymc
+import pytensor
+import pytensor.tensor as pt
+
+from arviz import InferenceData
+from better_optimize import basinhopping, minimize
+from better_optimize.constants import minimize_method
+from pymc import DictToArrayBijection, Model, join_nonshared_inputs
+from pymc.blocking import RaveledVars
+from pymc.util import RandomSeed
+from pytensor.tensor.variable import TensorVariable
+
+from pymc_extras.inference.laplace_approx.idata import (
+    add_data_to_inference_data,
+    add_optimizer_result_to_inference_data,
+)
+from pymc_extras.inference.laplace_approx.laplace import draws_from_laplace_approx
+from pymc_extras.inference.laplace_approx.scipy_interface import (
+    scipy_optimize_funcs_from_loss,
+    set_optimizer_function_defaults,
+)
+
+
+def fit_dadvi(
+    model: Model | None = None,
+    n_fixed_draws: int = 30,
+    n_draws: int = 1000,
+    include_transformed: bool = False,
+    optimizer_method: minimize_method = "trust-ncg",
+    use_grad: bool | None = None,
+    use_hessp: bool | None = None,
+    use_hess: bool | None = None,
+    gradient_backend: str = "pytensor",
+    compile_kwargs: dict | None = None,
+    random_seed: RandomSeed = None,
+    progressbar: bool = True,
+    **optimizer_kwargs,
+) -> az.InferenceData:
+    """
+    Does inference using Deterministic ADVI (Automatic Differentiation Variational Inference), DADVI for short.
+
+    For full details see the paper cited in the references: https://www.jmlr.org/papers/v25/23-1015.html
+
+    Parameters
+    ----------
+    model : pm.Model
+        The PyMC model to be fit. If None, the current model context is used.
+
+    n_fixed_draws : int
+        The number of fixed draws to use for the optimisation. More draws will result in more accurate estimates, but
+        also increase inference time. Usually, the default of 30 is a good tradeoff between speed and accuracy.
+
+    random_seed: int
+        The random seed to use for the fixed draws. Running the optimisation twice with the same seed should arrive at
+        the same result.
+
+    n_draws: int
+        The number of draws to return from the variational approximation.
+
+    include_transformed: bool
+        Whether or not to keep the unconstrained variables (such as logs of positive-constrained parameters) in the
+        output.
+
+    optimizer_method: str
+        Which optimization method to use. The function calls ``scipy.optimize.minimize``, so any of the methods there
+        can be used. The default is trust-ncg, which uses second-order information and is generally very reliable.
+        Other methods such as L-BFGS-B might be faster but potentially more brittle and may not converge exactly to
+        the optimum.
+
+    gradient_backend: str
+        Which backend to use to compute gradients. Must be one of "jax" or "pytensor". Default is "pytensor".
+
+    compile_kwargs: dict, optional
+        Additional keyword arguments to pass to `pytensor.function`
+
+    use_grad: bool, optional
+        If True, pass the gradient function to `scipy.optimize.minimize` (where it is referred to as `jac`).
+
+    use_hessp: bool, optional
+        If True, pass the hessian vector product to `scipy.optimize.minimize`.
+
+    use_hess: bool, optional
+        If True, pass the hessian to `scipy.optimize.minimize`. Note that this is generally not recommended since its
+        computation can be slow and memory-intensive if there are many parameters.
+
+    progressbar: bool
+        Whether or not to show a progress bar during optimization. Default is True.
+
+    optimizer_kwargs:
+        Additional keyword arguments to pass to the ``scipy.optimize.minimize`` function. See the documentation of
+        that function for details.
+
+    Returns
+    -------
+    :class:`~arviz.InferenceData`
+        The inference data containing the results of the DADVI algorithm.
+
+    References
+    ----------
+    Giordano, R., Ingram, M., & Broderick, T. (2024). Black Box Variational Inference with a Deterministic Objective:
+    Faster, More Accurate, and Even More Black Box. Journal of Machine Learning Research, 25(18), 1–39.
+    """
+
+    model = pymc.modelcontext(model) if model is None else model
+    do_basinhopping = optimizer_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.
+
+        optimizer_method = minimizer_kwargs.pop("method", "L-BFGS-B")
+        minimizer_kwargs["method"] = optimizer_method
+
+    initial_point_dict = model.initial_point()
+    initial_point = DictToArrayBijection.map(initial_point_dict)
+    n_params = initial_point.data.shape[0]
+
+    var_params, objective = create_dadvi_graph(
+        model,
+        n_fixed_draws=n_fixed_draws,
+        random_seed=random_seed,
+        n_params=n_params,
+    )
+
+    use_grad, use_hess, use_hessp = set_optimizer_function_defaults(
+        optimizer_method, use_grad, use_hess, use_hessp
+    )
+
+    f_fused, f_hessp = scipy_optimize_funcs_from_loss(
+        loss=objective,
+        inputs=[var_params],
+        initial_point_dict=None,
+        use_grad=use_grad,
+        use_hessp=use_hessp,
+        use_hess=use_hess,
+        gradient_backend=gradient_backend,
+        compile_kwargs=compile_kwargs,
+        inputs_are_flat=True,
+    )
+
+    dadvi_initial_point = {
+        f"{var_name}_mu": np.zeros_like(value).ravel()
+        for var_name, value in initial_point_dict.items()
+    }
+    dadvi_initial_point.update(
+        {
+            f"{var_name}_sigma__log": np.zeros_like(value).ravel()
+            for var_name, value in initial_point_dict.items()
+        }
+    )
+
+    dadvi_initial_point = DictToArrayBijection.map(dadvi_initial_point)
+    args = optimizer_kwargs.pop("args", ())
+
+    if do_basinhopping:
+        if "args" not in minimizer_kwargs:
+            minimizer_kwargs["args"] = args
+        if "hessp" not in minimizer_kwargs:
+            minimizer_kwargs["hessp"] = f_hessp
+        if "method" not in minimizer_kwargs:
+            minimizer_kwargs["method"] = optimizer_method
+
+        result = basinhopping(
+            func=f_fused,
+            x0=dadvi_initial_point.data,
+            progressbar=progressbar,
+            minimizer_kwargs=minimizer_kwargs,
+            **optimizer_kwargs,
+        )
+
+    else:
+        result = minimize(
+            f=f_fused,
+            x0=dadvi_initial_point.data,
+            args=args,
+            method=optimizer_method,
+            hessp=f_hessp,
+            progressbar=progressbar,
+            **optimizer_kwargs,
+        )
+
+    raveled_optimized = RaveledVars(result.x, dadvi_initial_point.point_map_info)
+
+    opt_var_params = result.x
+    opt_means, opt_log_sds = np.split(opt_var_params, 2)
+
+    posterior, unconstrained_posterior = draws_from_laplace_approx(
+        mean=opt_means,
+        standard_deviation=np.exp(opt_log_sds),
+        draws=n_draws,
+        model=model,
+        vectorize_draws=False,
+        return_unconstrained=include_transformed,
+        random_seed=random_seed,
+    )
+    idata = InferenceData(posterior=posterior)
+    if include_transformed:
+        idata.add_groups(unconstrained_posterior=unconstrained_posterior)
+
+    var_name_to_model_var = {f"{var_name}_mu": var_name for var_name in initial_point_dict.keys()}
+    var_name_to_model_var.update(
+        {f"{var_name}_sigma__log": var_name for var_name in initial_point_dict.keys()}
+    )
+
+    idata = add_optimizer_result_to_inference_data(
+        idata=idata,
+        result=result,
+        method=optimizer_method,
+        mu=raveled_optimized,
+        model=model,
+        var_name_to_model_var=var_name_to_model_var,
+    )
+
+    idata = add_data_to_inference_data(
+        idata=idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
+    )
+
+    return idata
+
+
+def create_dadvi_graph(
+    model: Model,
+    n_params: int,
+    n_fixed_draws: int = 30,
+    random_seed: RandomSeed = None,
+) -> tuple[TensorVariable, TensorVariable]:
+    """
+    Sets up the DADVI graph in pytensor and returns it.
+
+    Parameters
+    ----------
+    model : pm.Model
+        The PyMC model to be fit.
+
+    n_params: int
+        The total number of parameters in the model.
+
+    n_fixed_draws : int
+        The number of fixed draws to use.
+
+    random_seed: int
+        The random seed to use for the fixed draws.
+
+    Returns
+    -------
+    Tuple[TensorVariable, TensorVariable]
+        A tuple whose first element contains the variational parameters,
+        and whose second contains the DADVI objective.
+    """
+
+    # Make the fixed draws
+    generator = np.random.default_rng(seed=random_seed)
+    draws = generator.standard_normal(size=(n_fixed_draws, n_params))
+
+    inputs = model.continuous_value_vars + model.discrete_value_vars
+    initial_point_dict = model.initial_point()
+    logp = model.logp()
+
+    # Graph in terms of a flat input
+    [logp], flat_input = join_nonshared_inputs(
+        point=initial_point_dict, outputs=[logp], inputs=inputs
+    )
+
+    var_params = pt.vector(name="eta", shape=(2 * n_params,))
+
+    means, log_sds = pt.split(var_params, axis=0, splits_size=[n_params, n_params], n_splits=2)
+
+    draw_matrix = pt.constant(draws)
+    samples = means + pt.exp(log_sds) * draw_matrix
+
+    logp_vectorized_draws = pytensor.graph.vectorize_graph(logp, replace={flat_input: samples})
+
+    mean_log_density = pt.mean(logp_vectorized_draws)
+    entropy = pt.sum(log_sds)
+
+    objective = -mean_log_density - entropy
+
+    return var_params, objective