pymc-extras 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

pymc_extras/inference/pathfinder/pathfinder.py

@@ -16,12 +16,13 @@
 import collections
 import logging
 import time
+import warnings
 
 from collections import Counter
 from collections.abc import Callable, Iterator
 from dataclasses import asdict, dataclass, field, replace
 from enum import Enum, auto
-from typing import Literal, TypeAlias
+from typing import Literal, Self, TypeAlias
 
 import arviz as az
 import filelock
@@ -59,9 +60,6 @@ from rich.progress import BarColumn, TextColumn, TimeElapsedColumn, TimeRemainin
 from rich.table import Table
 from rich.text import Text
 
-# TODO: change to typing.Self after Python versions greater than 3.10
-from typing_extensions import Self
-
 from pymc_extras.inference.laplace_approx.idata import add_data_to_inference_data
 from pymc_extras.inference.pathfinder.importance_sampling import (
     importance_sampling as _importance_sampling,
@@ -280,12 +278,13 @@ def alpha_recover(
     z = pt.diff(g, axis=0)
     alpha_l_init = pt.ones(N)
 
-    alpha, _ = pytensor.scan(
+    alpha = pytensor.scan(
         fn=compute_alpha_l,
         outputs_info=alpha_l_init,
         sequences=[s, z],
         n_steps=Lp1 - 1,
         allow_gc=False,
+        return_updates=False,
     )
 
     # assert np.all(alpha.eval() > 0), "alpha cannot be negative"
@@ -336,11 +335,12 @@ def inverse_hessian_factors(
             return pt.set_subtensor(chi_l[j_last], diff_l)
 
         chi_init = pt.zeros((J, N))
-        chi_mat, _ = pytensor.scan(
+        chi_mat = pytensor.scan(
             fn=chi_update,
             outputs_info=chi_init,
             sequences=[diff],
             allow_gc=False,
+            return_updates=False,
         )
 
         chi_mat = pt.matrix_transpose(chi_mat)
@@ -379,14 +379,14 @@ def inverse_hessian_factors(
     eta = pt.diagonal(E, axis1=-2, axis2=-1)
 
     # beta: (L, N, 2J)
-    alpha_diag, _ = pytensor.scan(lambda a: pt.diag(a), sequences=[alpha])
+    alpha_diag = pytensor.scan(lambda a: pt.diag(a), sequences=[alpha], return_updates=False)
     beta = pt.concatenate([alpha_diag @ Z, S], axis=-1)
 
     # more performant and numerically precise to use solve than inverse: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.inv.html
 
     # E_inv: (L, J, J)
     E_inv = pt.slinalg.solve_triangular(E, Ij, check_finite=False)
-    eta_diag, _ = pytensor.scan(pt.diag, sequences=[eta])
+    eta_diag = pytensor.scan(pt.diag, sequences=[eta], return_updates=False)
 
     # block_dd: (L, J, J)
     block_dd = (
@@ -532,7 +532,9 @@ def bfgs_sample_sparse(
 
     # qr_input: (L, N, 2J)
     qr_input = inv_sqrt_alpha_diag @ beta
-    (Q, R), _ = pytensor.scan(fn=pt.nlinalg.qr, sequences=[qr_input], allow_gc=False)
+    Q, R = pytensor.scan(
+        fn=pt.linalg.qr, sequences=[qr_input], allow_gc=False, return_updates=False
+    )
 
     IdN = pt.eye(R.shape[1])[None, ...]
     IdN += IdN * REGULARISATION_TERM
@@ -625,10 +627,11 @@ def bfgs_sample(
 
     L, N, JJ = beta.shape
 
-    (alpha_diag, inv_sqrt_alpha_diag, sqrt_alpha_diag), _ = pytensor.scan(
+    alpha_diag, inv_sqrt_alpha_diag, sqrt_alpha_diag = pytensor.scan(
         lambda a: [pt.diag(a), pt.diag(pt.sqrt(1.0 / a)), pt.diag(pt.sqrt(a))],
         sequences=[alpha],
         allow_gc=False,
+        return_updates=False,
     )
 
     u = pt.random.normal(size=(L, num_samples, N))
@@ -1398,6 +1401,7 @@ def multipath_pathfinder(
     random_seed: RandomSeed,
     pathfinder_kwargs: dict = {},
     compile_kwargs: dict = {},
+    display_summary: bool = True,
 ) -> MultiPathfinderResult:
     """
     Fit the Pathfinder Variational Inference algorithm using multiple paths with PyMC/PyTensor backend.
@@ -1556,8 +1560,9 @@ def multipath_pathfinder(
                     compute_time=compute_end - compute_start,
                 )
             )
-            # TODO: option to disable summary, save to file, etc.
-            mpr.display_summary()
+            # Display summary conditionally
+            if display_summary:
+                mpr.display_summary()
             if mpr.all_paths_failed:
                 raise ValueError(
                     "All paths failed. Consider decreasing the jitter or reparameterizing the model."
@@ -1600,6 +1605,14 @@ def fit_pathfinder(
     pathfinder_kwargs: dict = {},
     compile_kwargs: dict = {},
     initvals: dict | None = None,
+    # New pathfinder result integration options
+    add_pathfinder_groups: bool = True,
+    display_summary: bool | Literal["auto"] = "auto",
+    store_diagnostics: bool = False,
+    pathfinder_group: str = "pathfinder",
+    paths_group: str = "pathfinder_paths",
+    diagnostics_group: str = "pathfinder_diagnostics",
+    config_group: str = "pathfinder_config",
 ) -> az.InferenceData:
     """
     Fit the Pathfinder Variational Inference algorithm.
@@ -1658,6 +1671,22 @@ def fit_pathfinder(
     initvals: dict | None = None
         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.
+    add_pathfinder_groups : bool, optional
+        Whether to add pathfinder results as additional groups to the InferenceData (default is True).
+        When True, adds pathfinder and pathfinder_paths groups with optimization diagnostics.
+    display_summary : bool or "auto", optional
+        Whether to display the pathfinder results summary (default is "auto").
+        "auto" preserves current behavior, False suppresses console output.
+    store_diagnostics : bool, optional
+        Whether to include potentially large diagnostic arrays in the pathfinder groups (default is False).
+    pathfinder_group : str, optional
+        Name for the main pathfinder results group (default is "pathfinder").
+    paths_group : str, optional
+        Name for the per-path results group (default is "pathfinder_paths").
+    diagnostics_group : str, optional
+        Name for the diagnostics group (default is "pathfinder_diagnostics").
+    config_group : str, optional
+        Name for the configuration group (default is "pathfinder_config").
 
     Returns
     -------
@@ -1694,6 +1723,9 @@ def fit_pathfinder(
         maxcor = np.ceil(3 * np.log(N)).astype(np.int32)
         maxcor = max(maxcor, 5)
 
+    # Handle display_summary logic
+    should_display_summary = display_summary == "auto" or display_summary is True
+
     if inference_backend == "pymc":
         mp_result = multipath_pathfinder(
             model,
@@ -1714,6 +1746,7 @@ def fit_pathfinder(
             random_seed=random_seed,
             pathfinder_kwargs=pathfinder_kwargs,
             compile_kwargs=compile_kwargs,
+            display_summary=should_display_summary,
         )
         pathfinder_samples = mp_result.samples
     elif inference_backend == "blackjax":
@@ -1760,4 +1793,30 @@ def fit_pathfinder(
 
     idata = add_data_to_inference_data(idata, progressbar, model, compile_kwargs)
 
+    # Add pathfinder results to InferenceData if requested
+    if add_pathfinder_groups:
+        if inference_backend == "pymc":
+            from pymc_extras.inference.pathfinder.idata import add_pathfinder_to_inference_data
+
+            idata = add_pathfinder_to_inference_data(
+                idata=idata,
+                result=mp_result,
+                model=model,
+                group=pathfinder_group,
+                paths_group=paths_group,
+                diagnostics_group=diagnostics_group,
+                config_group=config_group,
+                store_diagnostics=store_diagnostics,
+            )
+        else:
+            warnings.warn(
+                f"Pathfinder diagnostic groups are only supported with the PyMC backend. "
+                f"Current backend is '{inference_backend}', which does not support adding "
+                "pathfinder diagnostics to InferenceData. The InferenceData will only contain "
+                "posterior samples. To add diagnostic groups, use inference_backend='pymc', "
+                "or set add_pathfinder_groups=False to suppress this warning.",
+                UserWarning,
+                stacklevel=2,
+            )
+
     return idata

pymc_extras/inference/smc/sampling.py

@@ -238,7 +238,7 @@ class SMCDiagnostics(NamedTuple):
     def update_diagnosis(i, history, info, state):
         le, lli, ancestors, weights_evolution = history
         return SMCDiagnostics(
-            le.at[i].set(state.lmbda),
+            le.at[i].set(state.tempering_param),
             lli.at[i].set(info.log_likelihood_increment),
             ancestors.at[i].set(info.ancestors),
             weights_evolution.at[i].set(state.weights),
@@ -265,7 +265,7 @@ def inference_loop(rng_key, initial_state, kernel, iterations_to_diagnose, n_par
 
     def cond(carry):
         i, state, _, _ = carry
-        return state.lmbda < 1
+        return state.tempering_param < 1
 
     def one_step(carry):
         i, state, k, previous_info = carry

pymc_extras/model/marginal/distributions.py

@@ -282,11 +282,12 @@ def finite_discrete_marginal_rv_logp(op: MarginalFiniteDiscreteRV, values, *inpu
         def logp_fn(marginalized_rv_const, *non_sequences):
             return graph_replace(joint_logp, replace={marginalized_vv: marginalized_rv_const})
 
-        joint_logps, _ = scan_map(
+        joint_logps = scan_map(
             fn=logp_fn,
             sequences=marginalized_rv_domain_tensor,
             non_sequences=[*values, *inputs],
             mode=Mode().including("local_remove_check_parameter"),
+            return_updates=False,
         )
 
     joint_logp = pt.logsumexp(joint_logps, axis=0)
@@ -350,12 +351,13 @@ def marginal_hmm_logp(op, values, *inputs, **kwargs):
 
     P_bcast_dims = (len(chain_shape) - 1) - (P.type.ndim - 2)
     log_P = pt.shape_padright(pt.log(P), P_bcast_dims)
-    log_alpha_seq, _ = scan(
+    log_alpha_seq = scan(
         step_alpha,
         non_sequences=[log_P],
         outputs_info=[log_alpha_init],
         # Scan needs the time dimension first, and we already consumed the 1st logp computing the initial value
         sequences=pt.moveaxis(batch_logp_emissions[..., 1:], -1, 0),
+        return_updates=False,
     )
     # Final logp is just the sum of the last scan state
     joint_logp = pt.logsumexp(log_alpha_seq[-1], axis=0)

pymc_extras/model/marginal/graph_analysis.py

@@ -6,8 +6,8 @@ from itertools import zip_longest
 from pymc import SymbolicRandomVariable
 from pymc.model.fgraph import ModelVar
 from pymc.variational.minibatch_rv import MinibatchRandomVariable
-from pytensor.graph import Variable, ancestors
-from pytensor.graph.basic import io_toposort
+from pytensor.graph.basic import Variable
+from pytensor.graph.traversal import ancestors, io_toposort
 from pytensor.tensor import TensorType, TensorVariable
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise

pymc_extras/model/marginal/marginal_model.py

@@ -11,7 +11,7 @@ from pymc.backends.arviz import coords_and_dims_for_inferencedata, dataset_to_po
 from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
 from pymc.distributions.transforms import Chain
 from pymc.logprob.transforms import IntervalTransform
-from pymc.model import Model
+from pymc.model import Model, modelcontext
 from pymc.model.fgraph import (
     ModelFreeRV,
     ModelValuedVar,
@@ -337,8 +337,9 @@ def transform_posterior_pts(model, posterior_pts):
 
 
 def recover_marginals(
-    model: Model,
     idata: InferenceData,
+    *,
+    model: Model | None = None,
     var_names: Sequence[str] | None = None,
     return_samples: bool = True,
     extend_inferencedata: bool = True,
@@ -389,6 +390,15 @@ def recover_marginals(
 
 
     """
+    # Temporary error message for helping with migration
+    # Will be removed in a future release
+    if isinstance(idata, Model):
+        raise TypeError(
+            "The order of arguments of `recover_marginals` changed. The first input must be an idata"
+        )
+
+    model = modelcontext(model)
+
     unmarginal_model = unmarginalize(model)
 
     # Find the names of the marginalized variables

pymc_extras/model_builder.py

@@ -334,7 +334,9 @@ class ModelBuilder:
         >>> model = MyModel(ModelBuilder)
         >>> idata = az.InferenceData(your_dataset)
         >>> model.set_idata_attrs(idata=idata)
-        >>> assert "id" in idata.attrs #this and the following lines are part of doctest, not user manual
+        >>> assert (
+        ...     "id" in idata.attrs
+        ... )  # this and the following lines are part of doctest, not user manual
         >>> assert "model_type" in idata.attrs
         >>> assert "version" in idata.attrs
         >>> assert "sampler_config" in idata.attrs
@@ -382,7 +384,7 @@ class ModelBuilder:
         >>>         super().__init__()
         >>> model = MyModel()
         >>> model.fit(data)
-        >>> model.save('model_results.nc')  # This will call the overridden method in MyModel
+        >>> model.save("model_results.nc")  # This will call the overridden method in MyModel
         """
         if self.idata is not None and "posterior" in self.idata:
             file = Path(str(fname))
@@ -432,7 +434,7 @@ class ModelBuilder:
         --------
         >>> class MyModel(ModelBuilder):
         >>>     ...
-        >>> name = './mymodel.nc'
+        >>> name = "./mymodel.nc"
         >>> imported_model = MyModel.load(name)
         """
         filepath = Path(str(fname))
@@ -444,6 +446,7 @@ class ModelBuilder:
             sampler_config=json.loads(idata.attrs["sampler_config"]),
         )
         model.idata = idata
+        model.is_fitted_ = True
         dataset = idata.fit_data.to_dataframe()
         X = dataset.drop(columns=[model.output_var])
         y = dataset[model.output_var]
@@ -524,6 +527,8 @@ class ModelBuilder:
             )
             self.idata.add_groups(fit_data=combined_data.to_xarray())  # type: ignore
 
+        self.is_fitted_ = True
+
         return self.idata  # type: ignore
 
     def predict(
@@ -554,7 +559,7 @@ class ModelBuilder:
         >>> model = MyModel()
         >>> idata = model.fit(data)
         >>> x_pred = []
-        >>> prediction_data = pd.DataFrame({'input':x_pred})
+        >>> prediction_data = pd.DataFrame({"input": x_pred})
         >>> pred_mean = model.predict(prediction_data)
         """
 

pymc_extras/prior.py

@@ -70,8 +70,10 @@ Create a prior with a custom transform function by registering it with
 
     from pymc_extras.prior import register_tensor_transform
 
+
     def custom_transform(x):
-        return x ** 2
+        return x**2
+
 
     register_tensor_transform("square", custom_transform)
 
@@ -138,9 +140,23 @@ def handle_dims(x: pt.TensorLike, dims: Dims, desired_dims: Dims) -> pt.TensorVa
 
     Doesn't check for validity of the dims
 
+    Parameters
+    ----------
+    x : pt.TensorLike
+        The tensor to align.
+    dims : Dims
+        The current dimensions of the tensor.
+    desired_dims : Dims
+        The desired dimensions of the tensor.
+
+    Returns
+    -------
+    pt.TensorVariable
+        The aligned tensor.
+
     Examples
     --------
-    1D to 2D with new dim
+    Handle transpose 1D to 2D with new dimension.
 
     .. code-block:: python
 
@@ -177,10 +193,40 @@ def handle_dims(x: pt.TensorLike, dims: Dims, desired_dims: Dims) -> pt.TensorVa
 
 
 DimHandler = Callable[[pt.TensorLike, Dims], pt.TensorLike]
+"""A function that takes a tensor and its current dims and makes it broadcastable to the desired dims."""
 
 
 def create_dim_handler(desired_dims: Dims) -> DimHandler:
-    """Wrap the `handle_dims` function to act like the previous `create_dim_handler` function."""
+    """Wrap the :func:`handle_dims` function to always use the same desired_dims.
+
+    Parameters
+    ----------
+    desired_dims : Dims
+        The desired dimensions to align to.
+
+    Returns
+    -------
+    DimHandler
+        A function that takes a tensor and its current dims and aligns it to
+        the desired dims.
+
+
+    Examples
+    --------
+    Create a dim handler to align to ("channel", "group").
+
+    .. code-block:: python
+
+        import numpy as np
+
+        from pymc_extras.prior import create_dim_handler
+
+        dim_handler = create_dim_handler(("channel", "group"))
+
+        result = dim_handler(np.array([1, 2, 3]), dims="channel")
+
+
+    """
 
     def func(x: pt.TensorLike, dims: Dims) -> pt.TensorVariable:
         return handle_dims(x, dims, desired_dims)
@@ -228,8 +274,10 @@ def register_tensor_transform(name: str, transform: Transform) -> None:
             register_tensor_transform,
         )
 
+
         def custom_transform(x):
-            return x ** 2
+            return x**2
+
 
         register_tensor_transform("square", custom_transform)
 
@@ -268,9 +316,50 @@ def _get_pymc_parameters(distribution: pm.Distribution) -> set[str]:
 
 @runtime_checkable
 class VariableFactory(Protocol):
-    """Protocol for something that works like a Prior class."""
+    '''Protocol for something that works like a Prior class.
+
+    Sample with :func:`sample_prior`.
+
+    Examples
+    --------
+    Create a custom variable factory.
+
+    .. code-block:: python
+
+        import pymc as pm
+
+        import pytensor.tensor as pt
+
+        from pymc_extras.prior import sample_prior, VariableFactory
+
+
+        class PowerSumDistribution:
+            """Create a distribution that is the sum of powers of a base distribution."""
+
+            def __init__(self, distribution: VariableFactory, n: int):
+                self.distribution = distribution
+                self.n = n
+
+            @property
+            def dims(self):
+                return self.distribution.dims
+
+            def create_variable(self, name: str) -> "TensorVariable":
+                raw = self.distribution.create_variable(f"{name}_raw")
+                return pm.Deterministic(
+                    name,
+                    pt.sum([raw**n for n in range(1, self.n + 1)], axis=0),
+                    dims=self.dims,
+                )
+
+
+        cubic = PowerSumDistribution(Prior("Normal"), n=3)
+        samples = sample_prior(cubic)
+
+    '''
 
     dims: tuple[str, ...]
+    """The dimensions of the variable to create."""
 
     def create_variable(self, name: str) -> pt.TensorVariable:
         """Create a TensorVariable."""
@@ -316,6 +405,7 @@ def sample_prior(
 
         from pymc_extras.prior import sample_prior
 
+
         class CustomVariableDefinition:
             def __init__(self, dims, n: int):
                 self.dims = dims
@@ -323,7 +413,8 @@ def sample_prior(
 
             def create_variable(self, name: str) -> "TensorVariable":
                 x = pm.Normal(f"{name}_x", mu=0, sigma=1, dims=self.dims)
-                return pt.sum([x ** n for n in range(1, self.n + 1)], axis=0)
+                return pt.sum([x**n for n in range(1, self.n + 1)], axis=0)
+
 
         cubic = CustomVariableDefinition(dims=("channel",), n=3)
         coords = {"channel": ["C1", "C2", "C3"]}
@@ -381,6 +472,82 @@ class Prior:
         be registered with `register_tensor_transform` function or
         be available in either `pytensor.tensor` or `pymc.math`.
 
+    Examples
+    --------
+    Create a normal prior.
+
+    .. code-block:: python
+
+        from pymc_extras.prior import Prior
+
+        normal = Prior("Normal")
+
+    Create a hierarchical normal prior by using distributions for the parameters
+    and specifying the dims.
+
+    .. code-block:: python
+
+        hierarchical_normal = Prior(
+            "Normal",
+            mu=Prior("Normal"),
+            sigma=Prior("HalfNormal"),
+            dims="channel",
+        )
+
+    Create a non-centered hierarchical normal prior with the `centered` parameter.
+
+    .. code-block:: python
+
+        non_centered_hierarchical_normal = Prior(
+            "Normal",
+            mu=Prior("Normal"),
+            sigma=Prior("HalfNormal"),
+            dims="channel",
+            # Only change needed to make it non-centered
+            centered=False,
+        )
+
+    Create a hierarchical beta prior by using Beta distribution, distributions for
+    the parameters, and specifying the dims.
+
+    .. code-block:: python
+
+        hierarchical_beta = Prior(
+            "Beta",
+            alpha=Prior("HalfNormal"),
+            beta=Prior("HalfNormal"),
+            dims="channel",
+        )
+
+    Create a transformed hierarchical normal prior by using the `transform`
+    parameter. Here the "sigmoid" transformation comes from `pm.math`.
+
+    .. code-block:: python
+
+        transformed_hierarchical_normal = Prior(
+            "Normal",
+            mu=Prior("Normal"),
+            sigma=Prior("HalfNormal"),
+            transform="sigmoid",
+            dims="channel",
+        )
+
+    Create a prior with a custom transform function by registering it with
+    :func:`register_tensor_transform`.
+
+    .. code-block:: python
+
+        from pymc_extras.prior import register_tensor_transform
+
+
+        def custom_transform(x):
+            return x**2
+
+
+        register_tensor_transform("square", custom_transform)
+
+        custom_distribution = Prior("Normal", transform="square")
+
     """
 
     # Taken from https://en.wikipedia.org/wiki/Location%E2%80%93scale_family
@@ -389,9 +556,13 @@ class Prior:
         "StudentT": {"mu": 0, "sigma": 1},
         "ZeroSumNormal": {"sigma": 1},
     }
+    """Available non-centered distributions and their default parameters."""
 
     pymc_distribution: type[pm.Distribution]
+    """The PyMC distribution class."""
+
     pytensor_transform: Callable[[pt.TensorLike], pt.TensorLike] | None
+    """The PyTensor transform function."""
 
     @validate_call
     def __init__(
@@ -1317,9 +1488,33 @@ class Censored:
 
 
 class Scaled:
-    """Scaled distribution for numerical stability."""
+    """Scaled distribution for numerical stability.
+
+    This is the same as multiplying the variable by a constant factor.
+
+    Parameters
+    ----------
+    dist : Prior
+        The prior distribution to scale.
+    factor : pt.TensorLike
+        The scaling factor. This will have to be broadcastable to the
+        dimensions of the distribution.
+
+    Examples
+    --------
+    Create a scaled normal distribution.
+
+    .. code-block:: python
+
+        from pymc_extras.prior import Prior, Scaled
+
+        normal = Prior("Normal", mu=0, sigma=1)
+        # Same as Normal(mu=0, sigma=10)
+        scaled_normal = Scaled(normal, factor=10)
+
+    """
 
-    def __init__(self, dist: Prior, factor: float | pt.TensorVariable) -> None:
+    def __init__(self, dist: Prior, factor: pt.TensorLike) -> None:
         self.dist = dist
         self.factor = factor
 

pymc_extras/statespace/core/compile.py

@@ -28,7 +28,7 @@ def compile_statespace(
         x0, P0, c, d, T, Z, R, H, Q, steps=steps, sequence_names=sequence_names
     )
 
-    inputs = list(pytensor.graph.basic.explicit_graph_inputs(outputs))
+    inputs = list(pytensor.graph.traversal.explicit_graph_inputs(outputs))
 
     _f = pm.compile(inputs, outputs, on_unused_input="ignore", **compile_kwargs)
 

pymc_extras/statespace/core/statespace.py

@@ -2500,13 +2500,14 @@ class PyMCStateSpace:
                 next_x = c + T @ x + R @ shock
                 return next_x
 
-            irf, updates = pytensor.scan(
+            irf = pytensor.scan(
                 irf_step,
                 sequences=[shock_trajectory],
                 outputs_info=[x0],
                 non_sequences=[c, T, R],
                 n_steps=n_steps,
                 strict=True,
+                return_updates=False,
             )
 
             pm.Deterministic("irf", irf, dims=[TIME_DIM, ALL_STATE_DIM])