Update statespace imports

Author: fonnesbeck

pymc_extras/model/marginal/marginal_model.py

@@ -25,16 +25,16 @@
 
 from pymc_extras.distributions import DiscreteMarkovChain
 from pymc_extras.model.marginal.distributions import (
-    FiniteDiscreteMarginalRV,
-    MarginalRV,
-    MarginalizedRV,
+    MarginalDiscreteMarkovChainRV,
+    MarginalFiniteDiscreteRV,
+    get_domain_of_finite_discrete_rv,
+    reduce_batch_dependent_logps,
 )
 from pymc_extras.model.marginal.graph_analysis import (
-    get_support_axes,
-    get_support_shape,
-    get_support_size,
-    get_support_values,
-    get_variable_support,
+    find_conditional_dependent_rvs,
+    find_conditional_input_rvs,
+    is_conditional_dependent,
+    subgraph_batch_dim_connection,
 )
 
 ModelRVs = TensorVariable | Sequence[TensorVariable] | str | Sequence[str]

pymc_extras/statespace/core/compile.py

@@ -3,9 +3,9 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pymc_experimental.statespace.core import PyMCStateSpace
-from pymc_experimental.statespace.filters.distributions import LinearGaussianStateSpace
-from pymc_experimental.statespace.utils.constants import SHORT_NAME_TO_LONG
+from pymc_extras.statespace.core import PyMCStateSpace
+from pymc_extras.statespace.filters.distributions import LinearGaussianStateSpace
+from pymc_extras.statespace.utils.constants import SHORT_NAME_TO_LONG
 
 
 def compile_statespace(

pymc_extras/statespace/core/representation.py

@@ -4,9 +4,11 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pymc_experimental.statespace.utils.constants import (
-    NEVER_TIME_VARYING,
-    VECTOR_VALUED,
+from pymc_extras.statespace.utils.constants import (
+    JITTER_DEFAULT,
+    LONG_MATRIX_NAMES,
+    MISSING_FILL,
+    SHORT_NAME_TO_LONG,
 )
 
 floatX = pytensor.config.floatX

pymc_extras/statespace/core/statespace.py

@@ -16,49 +16,42 @@
 from pytensor import Variable, graph_replace
 from pytensor.compile import get_mode
 
-from pymc_experimental.statespace.core.representation import PytensorRepresentation
-from pymc_experimental.statespace.filters import (
-    KalmanSmoother,
-    SquareRootFilter,
+from pymc_extras.statespace.core.representation import PytensorRepresentation
+from pymc_extras.statespace.filters import (
+    CholeskyFilter,
+    SingleTimeseriesFilter,
     StandardFilter,
+    SteadyStateFilter,
     UnivariateFilter,
 )
-from pymc_experimental.statespace.filters.distributions import (
+from pymc_extras.statespace.filters.distributions import (
     LinearGaussianStateSpace,
-    MvNormalSVD,
-    SequenceMvNormal,
+    LinearGaussianStateSpaceRV,
 )
-from pymc_experimental.statespace.filters.utilities import stabilize
-from pymc_experimental.statespace.utils.constants import (
-    ALL_STATE_AUX_DIM,
-    ALL_STATE_DIM,
-    FILTER_OUTPUT_DIMS,
-    FILTER_OUTPUT_TYPES,
+from pymc_extras.statespace.filters.utilities import stabilize
+from pymc_extras.statespace.utils.constants import (
     JITTER_DEFAULT,
-    MATRIX_DIMS,
-    MATRIX_NAMES,
-    OBS_STATE_DIM,
-    SHOCK_DIM,
+    LONG_MATRIX_NAMES,
+    MISSING_FILL,
     SHORT_NAME_TO_LONG,
-    TIME_DIM,
-    VECTOR_VALUED,
 )
-from pymc_experimental.statespace.utils.data_tools import register_data_with_pymc
+from pymc_extras.statespace.utils.data_tools import register_data_with_pymc
 
 _log = logging.getLogger("pymc.experimental.statespace")
 
 floatX = pytensor.config.floatX
 FILTER_FACTORY = {
     "standard": StandardFilter,
     "univariate": UnivariateFilter,
-    "cholesky": SquareRootFilter,
+    "cholesky": CholeskyFilter,
+    "steady_state": SteadyStateFilter,
 }
 
 
 def _validate_filter_arg(filter_arg):
-    if filter_arg.lower() not in FILTER_OUTPUT_TYPES:
+    if filter_arg.lower() not in FILTER_FACTORY.keys():
         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_FACTORY.keys())}, received {filter_arg}'
         )
 
 
@@ -752,7 +745,7 @@ def _register_matrices_with_pymc_model(self) -> list[pt.TensorVariable]:
         matrices = self.unpack_statespace()
 
         registered_matrices = []
-        for i, (matrix, name) in enumerate(zip(matrices, MATRIX_NAMES)):
+        for i, (matrix, name) in enumerate(zip(matrices, LONG_MATRIX_NAMES)):
             time_varying_ndim = 2 if name in VECTOR_VALUED else 3
             if not getattr(pm_mod, name, None):
                 shape, dims = self._get_matrix_shape_and_dims(name)
@@ -1473,7 +1466,7 @@ def sample_statespace_matrices(
         _verify_group(group)
 
         if matrix_names is None:
-            matrix_names = MATRIX_NAMES
+            matrix_names = LONG_MATRIX_NAMES
         elif isinstance(matrix_names, str):
             matrix_names = [matrix_names]
 
@@ -1486,7 +1479,7 @@ def sample_statespace_matrices(
 
             self._insert_data_variables()
             matrices = self.unpack_statespace()
-            for short_name, matrix in zip(MATRIX_NAMES, matrices):
+            for short_name, matrix in zip(LONG_MATRIX_NAMES, matrices):
                 long_name = SHORT_NAME_TO_LONG[short_name]
                 if (long_name in matrix_names) or (short_name in matrix_names):
                     name = long_name if long_name in matrix_names else short_name
@@ -2040,7 +2033,7 @@ def forecast(
                 }
 
                 matrices = graph_replace(matrices, replace=sub_dict, strict=True)
-                [setattr(matrix, "name", name) for name, matrix in zip(MATRIX_NAMES[2:], matrices)]
+                [setattr(matrix, "name", name) for name, matrix in zip(LONG_MATRIX_NAMES[2:], matrices)]
 
             _ = LinearGaussianStateSpace(
                 "forecast",

pymc_extras/statespace/models/ETS.py

@@ -7,9 +7,13 @@
 from pytensor import graph_replace
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
-from pymc_experimental.statespace.core.statespace import PyMCStateSpace, floatX
-from pymc_experimental.statespace.models.utilities import make_default_coords
-from pymc_experimental.statespace.utils.constants import (
+from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
+from pymc_extras.statespace.models.utilities import make_default_coords
+from pymc_extras.statespace.utils.constants import (
+    JITTER_DEFAULT,
+    LONG_MATRIX_NAMES,
+    MISSING_FILL,
+    SHORT_NAME_TO_LONG,
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
     ETS_SEASONAL_DIM,

pymc_extras/statespace/models/SARIMAX.py

@@ -6,22 +6,18 @@
 
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
-from pymc_experimental.statespace.core.statespace import PyMCStateSpace, floatX
-from pymc_experimental.statespace.models.utilities import (
+from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
+from pymc_extras.statespace.models.utilities import (
     make_default_coords,
-    make_harvey_state_names,
-    make_SARIMA_transition_matrix,
+    make_seasonal_harmonics,
 )
-from pymc_experimental.statespace.utils.constants import (
-    ALL_STATE_AUX_DIM,
-    ALL_STATE_DIM,
-    AR_PARAM_DIM,
-    MA_PARAM_DIM,
-    OBS_STATE_DIM,
-    SARIMAX_STATE_STRUCTURES,
-    SEASONAL_AR_PARAM_DIM,
-    SEASONAL_MA_PARAM_DIM,
+from pymc_extras.statespace.utils.constants import (
+    JITTER_DEFAULT,
+    LONG_MATRIX_NAMES,
+    MISSING_FILL,
+    SHORT_NAME_TO_LONG,
 )
+import pymc_extras.statespace as pmss
 
 
 def _verify_order(p, d, q, P, D, Q, S):
@@ -147,7 +143,7 @@ class BayesianSARIMA(PyMCStateSpace):
 
     .. code:: python
 
-        import pymc_experimental.statespace as pmss
+        import pymc_extras.statespace as pmss
         import pymc as pm
 
         ss_mod = pmss.BayesianSARIMA(order=(1, 0, 1), verbose=True)

pymc_extras/statespace/models/VARMAX.py

@@ -7,9 +7,9 @@
 
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
-from pymc_experimental.statespace.core.statespace import PyMCStateSpace
-from pymc_experimental.statespace.models.utilities import make_default_coords
-from pymc_experimental.statespace.utils.constants import (
+from pymc_extras.statespace.core.statespace import PyMCStateSpace
+from pymc_extras.statespace.models.utilities import make_default_coords
+from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
     AR_PARAM_DIM,
@@ -18,7 +18,12 @@
     OBS_STATE_DIM,
     SHOCK_AUX_DIM,
     SHOCK_DIM,
+    JITTER_DEFAULT,
+    LONG_MATRIX_NAMES,
+    MISSING_FILL,
+    SHORT_NAME_TO_LONG,
 )
+import pymc_extras.statespace as pmss
 
 floatX = pytensor.config.floatX
 
@@ -110,7 +115,7 @@ class BayesianVARMAX(PyMCStateSpace):
 
     .. code:: python
 
-        import pymc_experimental.statespace as pmss
+        import pymc_extras.statespace as pmss
         import pymc as pm
 
         # Create VAR Statespace Model

pymc_extras/statespace/models/structural.py

@@ -13,20 +13,19 @@
 
 from pytensor import Variable
 
-from pymc_experimental.statespace.core import PytensorRepresentation
-from pymc_experimental.statespace.core.statespace import PyMCStateSpace
-from pymc_experimental.statespace.models.utilities import (
-    conform_time_varying_and_time_invariant_matrices,
+from pymc_extras.statespace.core import PytensorRepresentation
+from pymc_extras.statespace.core.statespace import PyMCStateSpace
+from pymc_extras.statespace.models.utilities import (
     make_default_coords,
+    make_seasonal_harmonics,
 )
-from pymc_experimental.statespace.utils.constants import (
-    ALL_STATE_AUX_DIM,
-    ALL_STATE_DIM,
-    AR_PARAM_DIM,
+from pymc_extras.statespace.utils.constants import (
+    JITTER_DEFAULT,
     LONG_MATRIX_NAMES,
-    POSITION_DERIVATIVE_NAMES,
-    TIME_DIM,
+    MISSING_FILL,
+    SHORT_NAME_TO_LONG,
 )
+from pymc_extras.statespace import structural as st
 
 _log = logging.getLogger("pymc.experimental.statespace")
 
@@ -1481,9 +1480,11 @@ def __init__(
             k_endog=k_endog,
             k_states=k_states,
             k_posdef=k_posdef,
+            state_names=self.state_names,
             measurement_error=False,
             combine_hidden_states=True,
-            obs_state_idxs=obs_state_idx,
+            exog_names=[f"data_{name}"],
+            obs_state_idxs=np.ones(k_states),
         )
 
     def make_symbolic_graph(self) -> None:

pymc_extras/statespace/models/utilities.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytensor.tensor as pt
 
-from pymc_experimental.statespace.utils.constants import (
+from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
     LONG_MATRIX_NAMES,
@@ -11,6 +11,9 @@
     SHOCK_AUX_DIM,
     SHOCK_DIM,
     VECTOR_VALUED,
+    JITTER_DEFAULT,
+    MISSING_FILL,
+    SHORT_NAME_TO_LONG,
 )
 
 
@@ -233,8 +236,8 @@ def make_SARIMA_transition_matrix(
                         0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \end{bmatrix}
 
     When ARIMA differences and seasonal differences are mixed, the seasonal differences will be written in terms of the
-    highest ARIMA difference order, and recovery of the level state will require the use of all the ARIMA differences,
-    as well as the seasonal differences. In addition, the seasonal differences are needed to back out the ARIMA
+    highest ARIMA difference order, and recovery of the level state will require the use of all the ARIMA
+    differences, as well as the seasonal differences. In addition, the seasonal differences are needed to back out the ARIMA
     differences from :math:`x_t^\star`. Here is the differencing block for a SARIMA(0,2,0)x(0,2,0,4) -- the identites
     of the states is left an exercise for the motivated reader: