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

pymc_extras/statespace/filters/distributions.py

@@ -69,9 +69,9 @@ class _LinearGaussianStateSpace(Continuous):
         H,
         Q,
         steps=None,
-        mode=None,
         sequence_names=None,
         append_x0=True,
+        method="svd",
         **kwargs,
     ):
         # Ignore dims in support shape because they are just passed along to the "observed" and "latent" distributions
@@ -97,9 +97,9 @@ class _LinearGaussianStateSpace(Continuous):
             H,
             Q,
             steps=steps,
-            mode=mode,
             sequence_names=sequence_names,
             append_x0=append_x0,
+            method=method,
             **kwargs,
         )
 
@@ -116,9 +116,9 @@ class _LinearGaussianStateSpace(Continuous):
         H,
         Q,
         steps=None,
-        mode=None,
         sequence_names=None,
         append_x0=True,
+        method="svd",
         **kwargs,
     ):
         steps = get_support_shape_1d(
@@ -132,9 +132,9 @@ class _LinearGaussianStateSpace(Continuous):
 
         return super().dist(
             [a0, P0, c, d, T, Z, R, H, Q, steps],
-            mode=mode,
             sequence_names=sequence_names,
             append_x0=append_x0,
+            method=method,
             **kwargs,
         )
 
@@ -152,9 +152,9 @@ class _LinearGaussianStateSpace(Continuous):
         Q,
         steps,
         size=None,
-        mode=None,
         sequence_names=None,
         append_x0=True,
+        method="svd",
     ):
         if sequence_names is None:
             sequence_names = []
@@ -205,10 +205,10 @@ class _LinearGaussianStateSpace(Continuous):
             a = state[:k]
 
             middle_rng, a_innovation = pm.MvNormal.dist(
-                mu=0, cov=Q, rng=rng, method="svd"
+                mu=0, cov=Q, rng=rng, method=method
             ).owner.outputs
             next_rng, y_innovation = pm.MvNormal.dist(
-                mu=0, cov=H, rng=middle_rng, method="svd"
+                mu=0, cov=H, rng=middle_rng, method=method
             ).owner.outputs
 
             a_mu = c + T @ a
@@ -224,8 +224,8 @@ class _LinearGaussianStateSpace(Continuous):
         Z_init = Z_ if Z_ in non_sequences else Z_[0]
         H_init = H_ if H_ in non_sequences else H_[0]
 
-        init_x_ = pm.MvNormal.dist(a0_, P0_, rng=rng, method="svd")
-        init_y_ = pm.MvNormal.dist(Z_init @ init_x_, H_init, rng=rng, method="svd")
+        init_x_ = pm.MvNormal.dist(a0_, P0_, rng=rng, method=method)
+        init_y_ = pm.MvNormal.dist(Z_init @ init_x_, H_init, rng=rng, method=method)
 
         init_dist_ = pt.concatenate([init_x_, init_y_], axis=0)
 
@@ -235,7 +235,6 @@ class _LinearGaussianStateSpace(Continuous):
             sequences=None if len(sequences) == 0 else sequences,
             non_sequences=[*non_sequences, rng],
             n_steps=steps,
-            mode=mode,
             strict=True,
         )
 
@@ -279,8 +278,8 @@ class LinearGaussianStateSpace(Continuous):
         steps,
         k_endog=None,
         sequence_names=None,
-        mode=None,
         append_x0=True,
+        method="svd",
         **kwargs,
     ):
         dims = kwargs.pop("dims", None)
@@ -307,9 +306,9 @@ class LinearGaussianStateSpace(Continuous):
             H,
             Q,
             steps=steps,
-            mode=mode,
             sequence_names=sequence_names,
             append_x0=append_x0,
+            method=method,
             **kwargs,
         )
         latent_obs_combined = pt.specify_shape(latent_obs_combined, (steps + int(append_x0), None))
@@ -368,11 +367,11 @@ class SequenceMvNormal(Continuous):
         return super().__new__(cls, *args, **kwargs)
 
     @classmethod
-    def dist(cls, mus, covs, logp, **kwargs):
-        return super().dist([mus, covs, logp], **kwargs)
+    def dist(cls, mus, covs, logp, method="svd", **kwargs):
+        return super().dist([mus, covs, logp], method=method, **kwargs)
 
     @classmethod
-    def rv_op(cls, mus, covs, logp, size=None):
+    def rv_op(cls, mus, covs, logp, method="svd", size=None):
         # Batch dimensions (if any) will be on the far left, but scan requires time to be there instead
         if mus.ndim > 2:
             mus = pt.moveaxis(mus, -2, 0)
@@ -385,7 +384,7 @@ class SequenceMvNormal(Continuous):
         rng = pytensor.shared(np.random.default_rng())
 
         def step(mu, cov, rng):
-            new_rng, mvn = pm.MvNormal.dist(mu=mu, cov=cov, rng=rng, method="svd").owner.outputs
+            new_rng, mvn = pm.MvNormal.dist(mu=mu, cov=cov, rng=rng, method=method).owner.outputs
             return mvn, {rng: new_rng}
 
         mvn_seq, updates = pytensor.scan(

pymc_extras/statespace/filters/kalman_filter.py

@@ -5,7 +5,6 @@ import pytensor
 import pytensor.tensor as pt
 
 from pymc.pytensorf import constant_fold
-from pytensor.compile.mode import get_mode
 from pytensor.graph.basic import Variable
 from pytensor.raise_op import Assert
 from pytensor.tensor import TensorVariable
@@ -28,15 +27,10 @@ assert_time_varying_dim_correct = Assert(
 
 
 class BaseFilter(ABC):
-    def __init__(self, mode=None):
+    def __init__(self):
         """
         Kalman Filter.
 
-        Parameters
-        ----------
-        mode : str, optional
-            The mode used for Pytensor compilation. Defaults to None.
-
         Notes
         -----
         The BaseFilter class is an abstract base class (ABC) for implementing kalman filters.
@@ -44,9 +38,6 @@ class BaseFilter(ABC):
 
         Attributes
         ----------
-        mode : str or None
-            The mode used for Pytensor compilation.
-
         seq_names : list[str]
             A list of name representing time-varying statespace matrices. That is, inputs that will need to be
             provided to the `sequences` argument of `pytensor.scan`
@@ -56,7 +47,6 @@ class BaseFilter(ABC):
             to the `non_sequences` argument of `pytensor.scan`
         """
 
-        self.mode: str = mode
         self.seq_names: list[str] = []
         self.non_seq_names: list[str] = []
 
@@ -153,7 +143,6 @@ class BaseFilter(ABC):
         R,
         H,
         Q,
-        mode=None,
         return_updates=False,
         missing_fill_value=None,
         cov_jitter=None,
@@ -166,9 +155,6 @@ class BaseFilter(ABC):
         data : TensorVariable
             Data to be filtered
 
-        mode : optional, str
-            Pytensor compile mode, passed to pytensor.scan
-
         return_updates: bool, default False
             Whether to return updates associated with the pytensor scan. Should only be requried to debug pruposes.
 
@@ -199,7 +185,6 @@ class BaseFilter(ABC):
         if cov_jitter is None:
             cov_jitter = JITTER_DEFAULT
 
-        self.mode = mode
         self.missing_fill_value = missing_fill_value
         self.cov_jitter = cov_jitter
 
@@ -227,7 +212,6 @@ class BaseFilter(ABC):
             outputs_info=[None, a0, None, None, P0, None, None],
             non_sequences=non_sequences,
             name="forward_kalman_pass",
-            mode=get_mode(self.mode),
             strict=False,
         )
 
@@ -800,7 +784,6 @@ class UnivariateFilter(BaseFilter):
             self._univariate_inner_filter_step,
             sequences=[y_masked, Z_masked, d, pt.diag(H_masked), nan_mask],
             outputs_info=[a, P, None, None, None],
-            mode=get_mode(self.mode),
             name="univariate_inner_scan",
         )
 

pymc_extras/statespace/filters/kalman_smoother.py

@@ -1,7 +1,6 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pytensor.compile import get_mode
 from pytensor.tensor.nlinalg import matrix_dot
 
 from pymc_extras.statespace.filters.utilities import (
@@ -18,8 +17,7 @@ class KalmanSmoother:
 
     """
 
-    def __init__(self, mode: str | None = None):
-        self.mode = mode
+    def __init__(self):
         self.cov_jitter = JITTER_DEFAULT
         self.seq_names = []
         self.non_seq_names = []
@@ -64,9 +62,8 @@ class KalmanSmoother:
         return a, P, a_smooth, P_smooth, T, R, Q
 
     def build_graph(
-        self, T, R, Q, filtered_states, filtered_covariances, mode=None, cov_jitter=JITTER_DEFAULT
+        self, T, R, Q, filtered_states, filtered_covariances, cov_jitter=JITTER_DEFAULT
     ):
-        self.mode = mode
         self.cov_jitter = cov_jitter
 
         n, k = filtered_states.type.shape
@@ -88,7 +85,6 @@ class KalmanSmoother:
             non_sequences=non_sequences,
             go_backwards=True,
             name="kalman_smoother",
-            mode=get_mode(self.mode),
         )
 
         smoothed_states, smoothed_covariances = smoother_result

pymc_extras/statespace/models/ETS.py

@@ -5,6 +5,7 @@ import numpy as np
 import pytensor.tensor as pt
 
 from pytensor import graph_replace
+from pytensor.compile.mode import Mode
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
@@ -35,6 +36,7 @@ class BayesianETS(PyMCStateSpace):
         initialization_dampening: float = 0.8,
         filter_type: str = "standard",
         verbose: bool = True,
+        mode: str | Mode | None = None,
     ):
         r"""
         Exponential Smoothing State Space Model
@@ -212,6 +214,13 @@ class BayesianETS(PyMCStateSpace):
             and "cholesky". See the docs for kalman filters for more details.
         verbose: bool, default True
             If true, a message will be logged to the terminal explaining the variable names, dimensions, and supports.
+        mode: str or Mode, optional
+            Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+            ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+            Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+            to all sampling methods.
+
 
         References
         ----------
@@ -284,6 +293,7 @@ class BayesianETS(PyMCStateSpace):
             filter_type,
             verbose=verbose,
             measurement_error=measurement_error,
+            mode=mode,
         )
 
     @property

pymc_extras/statespace/models/SARIMAX.py

@@ -4,6 +4,7 @@ from typing import Any
 import numpy as np
 import pytensor.tensor as pt
 
+from pytensor.compile.mode import Mode
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
@@ -91,6 +92,13 @@ class BayesianSARIMA(PyMCStateSpace):
     verbose: bool, default True
         If true, a message will be logged to the terminal explaining the variable names, dimensions, and supports.
 
+    mode: str or Mode, optional
+        Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+        ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+        Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+        to all sampling methods.
+
     Notes
     -----
         The ARIMAX model is a univariate time series model that posits the future evolution of a stationary time series will
@@ -158,7 +166,7 @@ class BayesianSARIMA(PyMCStateSpace):
             rho = pm.Beta("ar_params", alpha=5, beta=1, dims=ss_mod.param_dims["ar_params"])
             theta = pm.Normal("ma_params", mu=0.0, sigma=0.5, dims=ss_mod.param_dims["ma_params"])
 
-            ss_mod.build_statespace_graph(df, mode="JAX")
+            ss_mod.build_statespace_graph(df)
             idata = pm.sample(nuts_sampler='numpyro')
 
     References
@@ -180,7 +188,21 @@ class BayesianSARIMA(PyMCStateSpace):
         state_structure: str = "fast",
         measurement_error: bool = False,
         verbose=True,
+        mode: str | Mode | None = None,
     ):
+        """
+
+        Parameters
+        ----------
+        order
+        seasonal_order
+        stationary_initialization
+        filter_type
+        state_structure
+        measurement_error
+        verbose
+        mode
+        """
         # Model order
         self.p, self.d, self.q = order
         if seasonal_order is None:
@@ -228,6 +250,7 @@ class BayesianSARIMA(PyMCStateSpace):
             filter_type,
             verbose=verbose,
             measurement_error=measurement_error,
+            mode=mode,
         )
 
     @property
@@ -366,7 +389,7 @@ class BayesianSARIMA(PyMCStateSpace):
 
         return coords
 
-    def _stationary_initialization(self, mode=None):
+    def _stationary_initialization(self):
         # Solve for matrix quadratic for P0
         T = self.ssm["transition"]
         R = self.ssm["selection"]
@@ -374,9 +397,7 @@ class BayesianSARIMA(PyMCStateSpace):
         c = self.ssm["state_intercept"]
 
         x0 = pt.linalg.solve(pt.identity_like(T) - T, c, assume_a="gen", check_finite=True)
-
-        method = "direct" if (self.k_states < 5) or (mode == "JAX") else "bilinear"
-        P0 = solve_discrete_lyapunov(T, pt.linalg.matrix_dot(R, Q, R.T), method=method)
+        P0 = solve_discrete_lyapunov(T, pt.linalg.matrix_dot(R, Q, R.T), method="bilinear")
 
         return x0, P0
 

pymc_extras/statespace/models/VARMAX.py

@@ -5,6 +5,7 @@ import numpy as np
 import pytensor
 import pytensor.tensor as pt
 
+from pytensor.compile.mode import Mode
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
@@ -72,6 +73,13 @@ class BayesianVARMAX(PyMCStateSpace):
     verbose: bool, default True
         If true, a message will be logged to the terminal explaining the variable names, dimensions, and supports.
 
+    mode: str or Mode, optional
+        Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+        ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+        Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+        to all sampling methods.
+
     Notes
     -----
     The VARMA model is a multivariate extension of the SARIMAX model. Given a set of timeseries :math:`\{x_t\}_{t=0}^T`,
@@ -135,7 +143,7 @@ class BayesianVARMAX(PyMCStateSpace):
             ar_params = pm.Normal("ar_params", mu=0, sigma=1, dims=ar_dims)
             state_cov = pm.Deterministic("state_cov", state_chol @ state_chol.T, dims=state_cov_dims)
 
-            bvar_mod.build_statespace_graph(data, mode="JAX")
+            bvar_mod.build_statespace_graph(data)
             idata = pm.sample(nuts_sampler="numpyro")
     """
 
@@ -147,7 +155,8 @@ class BayesianVARMAX(PyMCStateSpace):
         stationary_initialization: bool = False,
         filter_type: str = "standard",
         measurement_error: bool = False,
-        verbose=True,
+        verbose: bool = True,
+        mode: str | Mode | None = None,
     ):
         if (endog_names is None) and (k_endog is None):
             raise ValueError("Must specify either endog_names or k_endog")
@@ -174,6 +183,7 @@ class BayesianVARMAX(PyMCStateSpace):
             filter_type,
             verbose=verbose,
             measurement_error=measurement_error,
+            mode=mode,
         )
 
         # Save counts of the number of parameters in each category

pymc_extras/statespace/models/structural.py

@@ -12,6 +12,7 @@ import pytensor.tensor as pt
 import xarray as xr
 
 from pytensor import Variable
+from pytensor.compile.mode import Mode
 
 from pymc_extras.statespace.core import PytensorRepresentation
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
@@ -81,6 +82,7 @@ class StructuralTimeSeries(PyMCStateSpace):
         name: str | None = None,
         verbose: bool = True,
         filter_type: str = "standard",
+        mode: str | Mode | None = None,
     ):
         # Add the initial state covariance to the parameters
         if name is None:
@@ -112,6 +114,7 @@ class StructuralTimeSeries(PyMCStateSpace):
             filter_type=filter_type,
             verbose=verbose,
             measurement_error=measurement_error,
+            mode=mode,
         )
         self.ssm = ssm.copy()
 
@@ -644,7 +647,9 @@ class Component(ABC):
 
         return new_comp
 
-    def build(self, name=None, filter_type="standard", verbose=True):
+    def build(
+        self, name=None, filter_type="standard", verbose=True, mode: str | Mode | None = None
+    ):
         """
         Build a StructuralTimeSeries statespace model from the current component(s)
 
@@ -660,6 +665,13 @@ class Component(ABC):
         verbose : bool, optional
             If True, displays information about the initialized model. Defaults to True.
 
+        mode: str or Mode, optional
+            Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+            ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+            Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+            to all sampling methods.
+
         Returns
         -------
         PyMCStateSpace
@@ -685,6 +697,7 @@ class Component(ABC):
             name_to_data=self._name_to_data,
             filter_type=filter_type,
             verbose=verbose,
+            mode=mode,
         )
 
 
@@ -908,7 +921,7 @@ class MeasurementError(Component):
             intitial_trend = pm.Normal('initial_trend', sigma=10, dims=ss_mod.param_dims['initial_trend'])
             sigma_obs = pm.Exponential('sigma_obs', 1, dims=ss_mod.param_dims['sigma_obs'])
 
-            ss_mod.build_statespace_graph(data, mode='JAX')
+            ss_mod.build_statespace_graph(data)
             idata = pm.sample(nuts_sampler='numpyro')
     """
 
@@ -991,7 +1004,7 @@ class AutoregressiveComponent(Component):
             ar_params = pm.Normal('ar_params', dims=ss_mod.param_dims['ar_params'])
             sigma_ar = pm.Exponential('sigma_ar', 1, dims=ss_mod.param_dims['sigma_ar'])
 
-            ss_mod.build_statespace_graph(data, mode='JAX')
+            ss_mod.build_statespace_graph(data)
             idata = pm.sample(nuts_sampler='numpyro')
 
     """
@@ -1153,7 +1166,7 @@ class TimeSeasonality(Component):
             intitial_trend = pm.Deterministic('initial_trend', pt.zeros(1), dims=ss_mod.param_dims['initial_trend'])
             annual_coefs = pm.Normal('annual_coefs', sigma=1e-2, dims=ss_mod.param_dims['annual_coefs'])
             trend_sigmas = pm.HalfNormal('trend_sigmas', sigma=1e-6, dims=ss_mod.param_dims['trend_sigmas'])
-            ss_mod.build_statespace_graph(data, mode='JAX')
+            ss_mod.build_statespace_graph(data)
             idata = pm.sample(nuts_sampler='numpyro')
 
     References
@@ -1451,7 +1464,7 @@ class CycleComponent(Component):
             cycle_length = pm.Uniform('business_cycle_length', lower=6, upper=12)
 
             sigma_cycle = pm.HalfNormal('sigma_business_cycle', sigma=1)
-            ss_mod.build_statespace_graph(data, mode='JAX')
+            ss_mod.build_statespace_graph(data)
 
             idata = pm.sample(nuts_sampler='numpyro')