pymc-extras 0.3.1__tar.gz → 0.4.1__tar.gz

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pymc-extras
-Version: 0.3.1
+Version: 0.4.1
 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
@@ -232,9 +232,11 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Requires-Python: >=3.11
-Requires-Dist: better-optimize>=0.1.4
+Requires-Dist: better-optimize>=0.1.5
+Requires-Dist: preliz>=0.20.0
 Requires-Dist: pydantic>=2.0.0
-Requires-Dist: pymc>=5.21.1
+Requires-Dist: pymc>=5.24.1
+Requires-Dist: pytensor>=2.31.4
 Requires-Dist: scikit-learn
 Provides-Extra: complete
 Requires-Dist: dask[complete]<2025.1.1; extra == 'complete'
@@ -245,7 +247,6 @@ Requires-Dist: xhistogram; extra == 'dask-histogram'
 Provides-Extra: dev
 Requires-Dist: blackjax; extra == 'dev'
 Requires-Dist: dask[all]<2025.1.1; extra == 'dev'
-Requires-Dist: preliz>=0.5.0; extra == 'dev'
 Requires-Dist: pytest-mock; extra == 'dev'
 Requires-Dist: pytest>=6.0; extra == 'dev'
 Requires-Dist: statsmodels; extra == 'dev'

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/_version.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.1'
-__version_tuple__ = version_tuple = (0, 3, 1)
+__version__ = version = '0.4.1'
+__version_tuple__ = version_tuple = (0, 4, 1)
+
+__commit_id__ = commit_id = None

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/conda-envs/environment-test.yml

@@ -3,9 +3,10 @@ channels:
 - conda-forge
 - nodefaults
 dependencies:
-  - pymc>=5.21.1
+  - pymc>=5.24.1
+  - pytensor>=2.31.4
   - scikit-learn
-  - better-optimize>=0.1.2
+  - better-optimize>=0.1.5
   - dask<2025.1.1
   - xhistogram
   - statsmodels
@@ -13,7 +14,7 @@ dependencies:
   - pytest
   - pytest-cov
   - pydantic>=2.0.0
-  - preliz>=0.5.0
+  - preliz>=0.20.0
   - pip
   - pip:
       - jax

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/distributions/histogram_utils.py

@@ -18,7 +18,7 @@ import pymc as pm
 
 from numpy.typing import ArrayLike
 
-__all__ = ["quantile_histogram", "discrete_histogram", "histogram_approximation"]
+__all__ = ["discrete_histogram", "histogram_approximation", "quantile_histogram"]
 
 
 def quantile_histogram(

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/inference/__init__.py

@@ -17,4 +17,4 @@ from pymc_extras.inference.laplace_approx.find_map import find_MAP
 from pymc_extras.inference.laplace_approx.laplace import fit_laplace
 from pymc_extras.inference.pathfinder.pathfinder import fit_pathfinder
 
-__all__ = ["fit", "fit_pathfinder", "fit_laplace", "find_MAP"]
+__all__ = ["find_MAP", "fit", "fit_laplace", "fit_pathfinder"]

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/inference/laplace_approx/find_map.py

@@ -326,7 +326,7 @@ def find_MAP(
     )
 
     raveled_optimized = RaveledVars(optimizer_result.x, initial_params.point_map_info)
-    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed)
+    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed=True)
     unobserved_vars_values = model.compile_fn(unobserved_vars, mode="FAST_COMPILE")(
         DictToArrayBijection.rmap(raveled_optimized)
     )
@@ -335,13 +335,20 @@ def find_MAP(
         var.name: value for var, value in zip(unobserved_vars, unobserved_vars_values)
     }
 
-    idata = map_results_to_inference_data(optimized_point, frozen_model)
-    idata = add_fit_to_inference_data(idata, raveled_optimized, H_inv)
+    idata = map_results_to_inference_data(
+        map_point=optimized_point, model=frozen_model, include_transformed=include_transformed
+    )
+
+    idata = add_fit_to_inference_data(
+        idata=idata, mu=raveled_optimized, H_inv=H_inv, model=frozen_model
+    )
+
     idata = add_optimizer_result_to_inference_data(
-        idata, optimizer_result, method, raveled_optimized, model
+        idata=idata, result=optimizer_result, method=method, mu=raveled_optimized, model=model
     )
+
     idata = add_data_to_inference_data(
-        idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
+        idata=idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
     )
 
     return idata

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/inference/laplace_approx/idata.py

@@ -59,6 +59,7 @@ def make_unpacked_variable_names(names: list[str], model: pm.Model) -> list[str]
 def map_results_to_inference_data(
     map_point: dict[str, float | int | np.ndarray],
     model: pm.Model | None = None,
+    include_transformed: bool = True,
 ):
     """
     Add the MAP point to an InferenceData object in the posterior group.
@@ -68,13 +69,13 @@ def map_results_to_inference_data(
 
     Parameters
     ----------
-    idata: az.InferenceData
-        An InferenceData object to which the MAP point will be added.
     map_point: dict
         A dictionary containing the MAP point estimates for each variable. The keys should be the variable names, and
         the values should be the corresponding MAP estimates.
     model: Model, optional
         A PyMC model. If None, the model is taken from the current model context.
+    include_transformed: bool
+        Whether to return transformed (unconstrained) variables in the constrained_posterior group. Default is True.
 
     Returns
     -------
@@ -118,7 +119,7 @@ def map_results_to_inference_data(
         dims=dims,
     )
 
-    if unconstrained_names:
+    if unconstrained_names and include_transformed:
         unconstrained_posterior = az.from_dict(
             posterior={
                 k: np.expand_dims(v, (0, 1))

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/inference/laplace_approx/laplace.py

@@ -302,7 +302,7 @@ def fit_laplace(
     ----------
     model : pm.Model
         The PyMC model to be fit. If None, the current model context is used.
-    method : str
+    optimize_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.
 
@@ -441,9 +441,11 @@ def fit_laplace(
             .rename({"temp_chain": "chain", "temp_draw": "draw"})
         )
 
-        idata.unconstrained_posterior = unstack_laplace_draws(
-            new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
-        )
+        if include_transformed:
+            idata.unconstrained_posterior = unstack_laplace_draws(
+                new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
+            )
+
         idata.posterior = new_posterior.drop_vars(
             ["laplace_approximation", "unpacked_variable_names"]
         )

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/inference/pathfinder/pathfinder.py

@@ -38,16 +38,15 @@ from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.initial_point import make_initial_point_fn
 from pymc.model import modelcontext
 from pymc.model.core import Point
+from pymc.progress_bar import CustomProgress, default_progress_theme
 from pymc.pytensorf import (
     compile,
     find_rng_nodes,
     reseed_rngs,
 )
 from pymc.util import (
-    CustomProgress,
     RandomSeed,
     _get_seeds_per_chain,
-    default_progress_theme,
     get_default_varnames,
 )
 from pytensor.compile.function.types import Function

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/printing.py

@@ -166,7 +166,7 @@ def model_table(
 
         for var in group:
             var_name = var.name
-            sep = f'[b]{" ~" if (var in model.basic_RVs) else " ="}[/b]'
+            sep = f"[b]{' ~' if (var in model.basic_RVs) else ' ='}[/b]"
             var_expr = variable_expression(model, var, truncate_deterministic)
             dims_expr = dims_expression(model, var)
             if dims_expr == "[]":

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/statespace/__init__.py

@@ -1,13 +1,13 @@
 from pymc_extras.statespace.core.compile import compile_statespace
 from pymc_extras.statespace.models import structural
 from pymc_extras.statespace.models.ETS import BayesianETS
-from pymc_extras.statespace.models.SARIMAX import BayesianSARIMA
+from pymc_extras.statespace.models.SARIMAX import BayesianSARIMAX
 from pymc_extras.statespace.models.VARMAX import BayesianVARMAX
 
 __all__ = [
-    "compile_statespace",
-    "structural",
     "BayesianETS",
-    "BayesianSARIMA",
+    "BayesianSARIMAX",
     "BayesianVARMAX",
+    "compile_statespace",
+    "structural",
 ]

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/statespace/core/__init__.py

@@ -4,4 +4,4 @@ from pymc_extras.statespace.core.representation import PytensorRepresentation
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
 from pymc_extras.statespace.core.compile import compile_statespace
 
-__all__ = ["PytensorRepresentation", "PyMCStateSpace", "compile_statespace"]
+__all__ = ["PyMCStateSpace", "PytensorRepresentation", "compile_statespace"]

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/statespace/core/statespace.py

@@ -60,7 +60,7 @@ FILTER_FACTORY = {
 def _validate_filter_arg(filter_arg):
     if filter_arg.lower() not in FILTER_OUTPUT_TYPES:
         raise ValueError(
-            f'filter_output should be one of {", ".join(FILTER_OUTPUT_TYPES)}, received {filter_arg}'
+            f"filter_output should be one of {', '.join(FILTER_OUTPUT_TYPES)}, received {filter_arg}"
         )
 
 
@@ -233,10 +233,9 @@ class PyMCStateSpace:
         self._fit_coords: dict[str, Sequence[str]] | None = None
         self._fit_dims: dict[str, Sequence[str]] | None = None
         self._fit_data: pt.TensorVariable | None = None
+        self._fit_exog_data: dict[str, dict] = {}
 
         self._needs_exog_data = None
-        self._exog_names = []
-        self._exog_data_info = {}
         self._name_to_variable = {}
         self._name_to_data = {}
 
@@ -671,7 +670,7 @@ class PyMCStateSpace:
         pymc_mod = modelcontext(None)
         for data_name in self.data_names:
             data = pymc_mod[data_name]
-            self._exog_data_info[data_name] = {
+            self._fit_exog_data[data_name] = {
                 "name": data_name,
                 "value": data.get_value(),
                 "dims": pymc_mod.named_vars_to_dims.get(data_name, None),
@@ -685,7 +684,7 @@ class PyMCStateSpace:
         --------
         .. code:: python
 
-            ss_mod = pmss.BayesianSARIMA(order=(2, 0, 2), verbose=False, stationary_initialization=True)
+            ss_mod = pmss.BayesianSARIMAX(order=(2, 0, 2), verbose=False, stationary_initialization=True)
             with pm.Model():
                  x0 = pm.Normal('x0', size=ss_mod.k_states)
                  ar_params = pm.Normal('ar_params', size=ss_mod.p)
@@ -805,16 +804,16 @@ class PyMCStateSpace:
         states, covs = outputs[:4], outputs[4:]
 
         state_names = [
-            "filtered_state",
-            "predicted_state",
-            "predicted_observed_state",
-            "smoothed_state",
+            "filtered_states",
+            "predicted_states",
+            "predicted_observed_states",
+            "smoothed_states",
         ]
         cov_names = [
-            "filtered_covariance",
-            "predicted_covariance",
-            "predicted_observed_covariance",
-            "smoothed_covariance",
+            "filtered_covariances",
+            "predicted_covariances",
+            "predicted_observed_covariances",
+            "smoothed_covariances",
         ]
 
         with mod:
@@ -939,7 +938,7 @@ class PyMCStateSpace:
             all_kf_outputs = [*states, smooth_states, *covs, smooth_covariances]
             self._register_kalman_filter_outputs_with_pymc_model(all_kf_outputs)
 
-        obs_dims = FILTER_OUTPUT_DIMS["predicted_observed_state"]
+        obs_dims = FILTER_OUTPUT_DIMS["predicted_observed_states"]
         obs_dims = obs_dims if all([dim in pm_mod.coords.keys() for dim in obs_dims]) else None
 
         SequenceMvNormal(
@@ -1082,7 +1081,7 @@ class PyMCStateSpace:
 
         for name in self.data_names:
             if name not in pm_mod:
-                pm.Data(**self._exog_data_info[name])
+                pm.Data(**self._fit_exog_data[name])
 
         self._insert_data_variables()
 
@@ -1229,7 +1228,7 @@ class PyMCStateSpace:
                     method=mvn_method,
                 )
 
-                obs_mu = (Z @ mu[..., None]).squeeze(-1)
+                obs_mu = d + (Z @ mu[..., None]).squeeze(-1)
                 obs_cov = Z @ cov @ pt.swapaxes(Z, -2, -1) + H
 
                 SequenceMvNormal(
@@ -1351,7 +1350,7 @@ class PyMCStateSpace:
             self._insert_random_variables()
 
             for name in self.data_names:
-                pm.Data(**self._exog_data_info[name])
+                pm.Data(**self._fit_exog_data[name])
 
             self._insert_data_variables()
 
@@ -1651,7 +1650,7 @@ class PyMCStateSpace:
             self._insert_random_variables()
 
             for name in self.data_names:
-                pm.Data(**self._exog_data_info[name])
+                pm.Data(**self.data_info[name])
 
             self._insert_data_variables()
             matrices = self.unpack_statespace()
@@ -1678,6 +1677,78 @@ class PyMCStateSpace:
 
         return matrix_idata
 
+    def sample_filter_outputs(
+        self, idata, filter_output_names: str | list[str] | None, group: str = "posterior", **kwargs
+    ):
+        if isinstance(filter_output_names, str):
+            filter_output_names = [filter_output_names]
+
+        if filter_output_names is None:
+            filter_output_names = list(FILTER_OUTPUT_DIMS.keys())
+        else:
+            unknown_filter_output_names = np.setdiff1d(
+                filter_output_names, list(FILTER_OUTPUT_DIMS.keys())
+            )
+            if unknown_filter_output_names.size > 0:
+                raise ValueError(f"{unknown_filter_output_names} not a valid filter output name!")
+            filter_output_names = [x for x in FILTER_OUTPUT_DIMS.keys() if x in filter_output_names]
+
+        compile_kwargs = kwargs.pop("compile_kwargs", {})
+        compile_kwargs.setdefault("mode", self.mode)
+
+        with pm.Model(coords=self.coords) as m:
+            self._build_dummy_graph()
+            self._insert_random_variables()
+
+            if self.data_names:
+                for name in self.data_names:
+                    pm.Data(**self._fit_exog_data[name])
+
+            self._insert_data_variables()
+
+            x0, P0, c, d, T, Z, R, H, Q = self.unpack_statespace()
+            data = self._fit_data
+
+            obs_coords = m.coords.get(OBS_STATE_DIM, None)
+
+            data, nan_mask = register_data_with_pymc(
+                data,
+                n_obs=self.ssm.k_endog,
+                obs_coords=obs_coords,
+                register_data=True,
+            )
+
+            filter_outputs = self.kalman_filter.build_graph(
+                data,
+                x0,
+                P0,
+                c,
+                d,
+                T,
+                Z,
+                R,
+                H,
+                Q,
+            )
+
+            smoother_outputs = self.kalman_smoother.build_graph(
+                T, R, Q, filter_outputs[0], filter_outputs[3]
+            )
+
+            filter_outputs = filter_outputs[:-1] + list(smoother_outputs)
+            for output in filter_outputs:
+                if output.name in filter_output_names:
+                    dims = FILTER_OUTPUT_DIMS[output.name]
+                    pm.Deterministic(output.name, output, dims=dims)
+
+        with freeze_dims_and_data(m):
+            return pm.sample_posterior_predictive(
+                idata if group == "posterior" else idata.prior,
+                var_names=filter_output_names,
+                compile_kwargs=compile_kwargs,
+                **kwargs,
+            )
+
     @staticmethod
     def _validate_forecast_args(
         time_index: pd.RangeIndex | pd.DatetimeIndex,
@@ -1774,7 +1845,7 @@ class PyMCStateSpace:
         }
 
         if self._needs_exog_data and scenario is None:
-            exog_str = ",".join(self._exog_names)
+            exog_str = ",".join(self.data_names)
             suffix = "s" if len(exog_str) > 1 else ""
             raise ValueError(
                 f"This model was fit using exogenous data. Forecasting cannot be performed without "
@@ -1783,7 +1854,7 @@ class PyMCStateSpace:
 
         if isinstance(scenario, dict):
             for name, data in scenario.items():
-                if name not in self._exog_names:
+                if name not in self.data_names:
                     raise ValueError(
                         f"Scenario data provided for variable '{name}', which is not an exogenous variable "
                         f"used to fit the model."
@@ -1824,12 +1895,12 @@ class PyMCStateSpace:
                 # name should only be None on the first non-recursive call. We only arrive to this branch in that case
                 # if a non-dictionary was passed, which in turn should only happen if only a single exogenous data
                 # needs to be set.
-                if len(self._exog_names) > 1:
+                if len(self.data_names) > 1:
                     raise ValueError(
                         "Multiple exogenous variables were used to fit the model. Provide a dictionary of "
                         "scenario data instead."
                     )
-                name = self._exog_names[0]
+                name = self.data_names[0]
 
             # Omit dataframe from this basic shape check so we can give more detailed information about missing columns
             # in the next check
@@ -2031,7 +2102,7 @@ class PyMCStateSpace:
             return scenario
 
         # This was already checked as valid
-        name = self._exog_names[0] if name is None else name
+        name = self.data_names[0] if name is None else name
 
         # Small tidying up in the case we just have a single scenario that's already a dataframe.
         if isinstance(scenario, pd.DataFrame | pd.Series):

{pymc_extras-0.3.1 → pymc_extras-0.4.1}/pymc_extras/statespace/filters/kalman_filter.py

@@ -15,10 +15,15 @@ from pymc_extras.statespace.filters.utilities import (
     split_vars_into_seq_and_nonseq,
     stabilize,
 )
-from pymc_extras.statespace.utils.constants import JITTER_DEFAULT, MISSING_FILL
+from pymc_extras.statespace.utils.constants import (
+    FILTER_OUTPUT_NAMES,
+    JITTER_DEFAULT,
+    MATRIX_NAMES,
+    MISSING_FILL,
+)
 
 MVN_CONST = pt.log(2 * pt.constant(np.pi, dtype="float64"))
-PARAM_NAMES = ["c", "d", "T", "Z", "R", "H", "Q"]
+PARAM_NAMES = MATRIX_NAMES[2:]
 
 assert_time_varying_dim_correct = Assert(
     "The first dimension of a time varying matrix (the time dimension) must be "
@@ -119,7 +124,7 @@ class BaseFilter(ABC):
         # There are always two outputs_info wedged between the seqs and non_seqs
         seqs, (a0, P0), non_seqs = args[:n_seq], args[n_seq : n_seq + 2], args[n_seq + 2 :]
         return_ordered = []
-        for name in ["c", "d", "T", "Z", "R", "H", "Q"]:
+        for name in PARAM_NAMES:
             if name in self.seq_names:
                 idx = self.seq_names.index(name)
                 return_ordered.append(seqs[idx])
@@ -253,28 +258,28 @@ class BaseFilter(ABC):
         )
 
         filtered_states = pt.specify_shape(filtered_states, (n, self.n_states))
-        filtered_states.name = "filtered_states"
+        filtered_states.name = FILTER_OUTPUT_NAMES[0]
 
         predicted_states = pt.specify_shape(predicted_states, (n, self.n_states))
-        predicted_states.name = "predicted_states"
-
-        observed_states = pt.specify_shape(observed_states, (n, self.n_endog))
-        observed_states.name = "observed_states"
+        predicted_states.name = FILTER_OUTPUT_NAMES[1]
 
         filtered_covariances = pt.specify_shape(
             filtered_covariances, (n, self.n_states, self.n_states)
         )
-        filtered_covariances.name = "filtered_covariances"
+        filtered_covariances.name = FILTER_OUTPUT_NAMES[2]
 
         predicted_covariances = pt.specify_shape(
             predicted_covariances, (n, self.n_states, self.n_states)
         )
-        predicted_covariances.name = "predicted_covariances"
+        predicted_covariances.name = FILTER_OUTPUT_NAMES[3]
+
+        observed_states = pt.specify_shape(observed_states, (n, self.n_endog))
+        observed_states.name = FILTER_OUTPUT_NAMES[4]
 
         observed_covariances = pt.specify_shape(
             observed_covariances, (n, self.n_endog, self.n_endog)
         )
-        observed_covariances.name = "observed_covariances"
+        observed_covariances.name = FILTER_OUTPUT_NAMES[5]
 
         loglike_obs = pt.specify_shape(loglike_obs.squeeze(), (n,))
         loglike_obs.name = "loglike_obs"