Merge statespace module from http://github.com/jessegrabowski/pymc_statespace (#174)

* move http://github.com/jessegrabowski/pymc_statespace to pymc-experimental

* Make tests compatible with float32

* Fix remaining tests that failed when dtype is float32

* Fix remaining tests that failed when dtype is float32

* Replace print with `logger.info`

* Use `getattr` to access model variables in `gather_required_random_variables`

* Reduce code duplication, eliminate use of `pathlib` in favor of explict long paths

* Refactor `PyMCStateSpace` and `PytensorRepresentation` to no longer require data in `__init__`

* Refactor kalman filters to remove singleton dimensions everywhere possible

* Add tests for sarimax, add first pass at state space distribution

* Add more helpers to statespace class for building PyMC model

* Changes to SARIMAX to make states interpretable, plus bad draft pymc distributions for kf outputs

* Distribution attempt 2

* Distribution attempt 3

* Fix broken tests

* Fix remaining tests

* Preserve name/shape information after updates to matrices

* Add test for `SequenceMvNormal`

* Add logp as parameter to `SequenceMvNormal`

* Add logp argument to `SequenceMvNormal` test

* Refactor shape checking to use `type.shape`
Refactor `statespace.build_statespace_graph()` to use new distribution, preparing for simulation refactor

* Refactor shape checking to use `type.shape`
Refactor `statespace.build_statespace_graph()` to use new distribution, preparing for simulation refactor

* Add float32 test skip back to `test_filters_match_statsmodel_output`

* Refactor `LinearGaussianStateSpace` for use in prior/posterior predictive sampling

* Infer coordinates for statespace matrices from parameters

* Refactor posterior predictive sampling, coords, and data handling

* Remove typecheck with `Union`

* Fix float32 error in tests

* Fix float32 error in tests

* Fix bugs in SARMIAX, add measurement error option to SARIMAX, update notebooks.

* always add `TIME_DIM` to model coords, even when data aren't given dims

* always add `TIME_DIM` to model coords, even when data aren't given dims

* Add forecast method to `Statespace`, add coord/dim info to `VARMAX`, update ARIMA notebook to show forecasting

* Add `impulse_response_function` method to `Statespace`, update `ARIMA` example notebook

* Update VARMAX notebook, remove `pm.DiracDelta` from `Statespace.forecast` (JAX doesn't support it)

* Update VARMAX example notebook

* Add docstrings to statespace methods and utilities

* More docstrings

* documentation

* More docs, fix failing tests

* Implement conditional and unconditonal sampling from prior idata

* Trying to get the API docs to render

* Trying to get the API docs to render

* Revert "Trying to get the API docs to render"

This reverts commit 3f766d3.

* Trying to get the API docs to render, revert local change to `pyproject.toml`

* Trying to get the API docs to render

* Remove `specify_broadcast` in `SingleTimeseriesFilter.update`

* Fixing docstrings

* Fixing docstrings

* create toctree hierarchy

* create toctree hierarchy

* create toctree hierarchy

* docstrings

* docstrings

* docstrings

* Making a table in a docstring

* Making a table in a docstring

* More docs

* More docstrings

* More docstrings

* More docstrings

* Remove `BayesianLocalLevel`, add a new `StructuralTimeSeries` module

* Remove `BayesianLocalLevel`, add a new `StructuralTimeSeries` module

* Bug fixes in `structural`, new example notebook

* Bug fixes in `structural`, update structural example

* Kalman filter no longer returns a forecast on the predicted states
Predicted state dimensions are no longer incorrectly shifted backwards 1 step
Remove `extended_time` dim from everywhere
Seed all tests

* Fix bugs in IRF API
Update VARMAX example notebook

* Fix bugs in IRF API
Update VARMAX example notebook

* Add `airpass.csv` test data
Update example notebooks
More docstrings

* More docstrings

* More docstrings

* More docstrings

* More docstrings

* Fix VARMAX tests

* More docs

* More docs, allow user to name states in seasonal components

* more docs, rename `test_structural_model` to `test_structural`

* Allow non-integer seasonal length in `FrequencySeasonality` (breaks tests)

* Seed `sample_posterior` and `sample_prior` tests, unbreak `test_structural.py`

* Set test seed as a global constant

* Remove numba dependency

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

* Add helper function to `StructuralTimeSeries` to reconstruct components from hidden states

* Add `fast_eval` as test utility

* Changes from review feedback

* Try to fix failing test

* Catch expected test warnings with `pytest.mark.filterwarnings`

* Scope fixtures in sampling tests to speed them up a bit

* More fixes from review

* Trying to fix the failing test

* Remove numba from `windows-environment-test.yml`

* Test for NaNs in sampled statespace objects

* Redo the `Custom Statespace` notebooks

* Add `ImputationWarning` to `mask_missing_values_in_data`

* Add `ImputationWarning` to `mask_missing_values_in_data`

Raise if data contains the missing mask value

Add `missing_variable_fill` argument to `build_statespace_graph`

* Add off-diagonal averaging to covariance `stabilize` function.

* Covariance matrices output by kalman filter not always PSD

* Add some stability tricks to Kalman filters

Add test that all filter outputs are PSD

* Add `measurement_error` kwarg to `PyMCStateSpace`

* Allow `eig == 0.0` to pass PSD test

* Add stability tests to VARMAX and SARIMAX

* LGSS distribution doesn't need to know anything about measurement error

Add jitter to `H` in unconditional_sampling when there is no measurement error

Stabilize `P_filtered` and `P_smoothed` as every kalman step

Loosen `atol` and `rtol` for kalman filter tests when comparing with statsmodels to account for all the diagonal jittering

* Re-run Structural Timeseries Modeling.ipynb with some new features

* Delete out-of-date example notebook

* Updates to Making a Custom Statespace Model.ipynb

* Link to the bVAR example notebook in the VARMAX Example.ipynb

* Make the IRF plots more readable in ARIMA Example.ipynb

* Increase `JITTER_DEFAULT` when pytensor is in `float32` mode

* Increase `JITTER_DEFAULT` when pytensor is in `float32` mode

* Add stabilization to univariate filter covariance

* Trying to get the last tests to pass

* Different stability strategy for UnivariateFilter

* Skip overly sensitive float32 tests

* Remove `update` function

`make_symbolic_graph` is now the core method for `PyMCStateSpace`

Add shape and name info to `KalmanFilter` outputs

* All tests pass

Update Structural Timeseries Modeling.ipynb to reflect refactor

* Update Making a Custom Statespace Model.ipynb to reflect refactor

* Relax float32 test tolerance in `test_structural.py`

* Set dtype on numpy arrays used in test_structural.py

* Remove dictionary merge with pipe

Tweak tolerances for test in `test_structural.py`

* Disable all PSD tests for univariate filter when floatX=float32

* Remove `variable_by_shape` helper, use `pt.tensor` directly.

* Remove unused helper functions

Add tests for impulse_response_function and forecast

Add tests for SARIMAX in "interpretable" mode

* Adjust new tests for float32

* Add test for equivalence between SARIMA representations

Add reference to Harvey (1989) in SARIMA docstring

* Add test for equivalence between SARIMA representations

Add reference to Harvey (1989) in SARIMA docstring

* Use `self.mode` in all kalman filter scans

* Begin adding seasonal components to SARIMAX.py

* Implement seasonal lags and differences in SARIMAX

* Add an Exogenous Regression component to `structural.py`

* Expand support for exogenous variables in statespace models

Re-run notebooks

* Tweak tests

* Tweak tests

---------

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

Author: jessegrabowski

conda-envs/environment-test.yml

@@ -9,6 +9,7 @@ dependencies:
 - pytest>=3.0
 - dask
 - xhistogram
+- statsmodels
 - pip:
   - pymc>=5.6.0  # CI was failing to resolve
   - blackjax

conda-envs/windows-environment-test.yml

@@ -8,6 +8,7 @@ dependencies:
 - pytest>=3.0
 - dask
 - xhistogram
+- statsmodels
 - pip:
   - pymc>=5.6.0  # CI was failing to resolve
   - scikit-learn

docs/api_reference.rst

@@ -60,3 +60,13 @@ Utils
    prior.prior_from_idata
    model_fgraph.fgraph_from_model
    model_fgraph.model_from_fgraph
+
+Statespace Models
+=================
+.. automodule:: pymc_experimental.statespace
+.. toctree::
+   :maxdepth: 1
+
+   statespace/core
+   statespace/filters
+   statespace/models

docs/statespace/core.rst

@@ -0,0 +1,10 @@
+********************
+Statespace Core
+********************
+
+.. automodule:: pymc_experimental.statespace.core
+.. autosummary::
+   :toctree: generated
+
+    PytensorRepresentation
+    PyMCStateSpace

docs/statespace/filters.rst

@@ -0,0 +1,15 @@
+********************
+Kalman Filters
+********************
+
+.. automodule:: pymc_experimental.statespace.filters
+.. autosummary::
+   :toctree: generated
+
+    StandardFilter
+    UnivariateFilter
+    SteadyStateFilter
+    KalmanSmoother
+    SingleTimeseriesFilter
+    CholeskyFilter
+    LinearGaussianStateSpace

docs/statespace/models.rst

@@ -0,0 +1,20 @@
+*****************
+Statespace Models
+*****************
+
+.. automodule:: pymc_experimental.statespace.models
+.. autosummary::
+   :toctree: generated
+
+    BayesianARIMA
+    BayesianVARMAX
+
+*********************
+Structural Components
+*********************
+
+.. currentmodule:: pymc_experimental.statespace.models
+.. toctree::
+   :maxdepth: 2
+
+   models/structural

docs/statespace/models/structural.rst

@@ -0,0 +1,16 @@
+*********************
+Structural Components
+*********************
+
+.. automodule:: pymc_experimental.statespace.models.structural
+.. autosummary::
+   :toctree: generated
+
+   LevelTrendComponent
+   AutoregressiveComponent
+   TimeSeasonality
+   FrequencySeasonality
+   MeasurementError
+
+   StructuralTimeSeries
+   Component