pymc-extras 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

pymc_extras/distributions/__init__.py

@@ -29,14 +29,14 @@ from pymc_extras.distributions.timeseries import DiscreteMarkovChain
 from pymc_extras.distributions.transforms import PartialOrder
 
 __all__ = [
+    "R2D2M2CP",
+    "BetaNegativeBinomial",
     "Chi",
-    "Maxwell",
     "DiscreteMarkovChain",
-    "GeneralizedPoisson",
-    "BetaNegativeBinomial",
     "GenExtreme",
-    "R2D2M2CP",
+    "GeneralizedPoisson",
+    "Maxwell",
+    "PartialOrder",
     "Skellam",
     "histogram_approximation",
-    "PartialOrder",
 ]

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/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/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/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/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/statespace/core/representation.py

@@ -153,19 +153,19 @@ class PytensorRepresentation:
     """
 
     __slots__ = (
+        "design",
+        "initial_state",
+        "initial_state_cov",
         "k_endog",
-        "k_states",
         "k_posdef",
-        "shapes",
-        "design",
-        "obs_intercept",
+        "k_states",
         "obs_cov",
-        "transition",
-        "state_intercept",
+        "obs_intercept",
         "selection",
+        "shapes",
         "state_cov",
-        "initial_state",
-        "initial_state_cov",
+        "state_intercept",
+        "transition",
     )
 
     def __init__(

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/statespace/filters/__init__.py

@@ -7,9 +7,9 @@ from pymc_extras.statespace.filters.kalman_filter import (
 from pymc_extras.statespace.filters.kalman_smoother import KalmanSmoother
 
 __all__ = [
-    "StandardFilter",
-    "UnivariateFilter",
     "KalmanSmoother",
-    "SquareRootFilter",
     "LinearGaussianStateSpace",
+    "SquareRootFilter",
+    "StandardFilter",
+    "UnivariateFilter",
 ]

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"