Skip to content

Commit

Permalink
Merge statespace module from http://github.com/jessegrabowski/pymc_st…
Browse files Browse the repository at this point in the history
…atespace (pymc-devs#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 <[email protected]>

* 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 <[email protected]>
  • Loading branch information
jessegrabowski and ricardoV94 authored Aug 18, 2023
1 parent 15c88e8 commit 67a9695
Show file tree
Hide file tree
Showing 48 changed files with 17,967 additions and 0 deletions.
1 change: 1 addition & 0 deletions conda-envs/environment-test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ dependencies:
- pytest>=3.0
- dask
- xhistogram
- statsmodels
- pip:
- pymc>=5.6.0 # CI was failing to resolve
- blackjax
Expand Down
1 change: 1 addition & 0 deletions conda-envs/windows-environment-test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ dependencies:
- pytest>=3.0
- dask
- xhistogram
- statsmodels
- pip:
- pymc>=5.6.0 # CI was failing to resolve
- scikit-learn
10 changes: 10 additions & 0 deletions docs/api_reference.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
10 changes: 10 additions & 0 deletions docs/statespace/core.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
********************
Statespace Core
********************

.. automodule:: pymc_experimental.statespace.core
.. autosummary::
:toctree: generated

PytensorRepresentation
PyMCStateSpace
15 changes: 15 additions & 0 deletions docs/statespace/filters.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
********************
Kalman Filters
********************

.. automodule:: pymc_experimental.statespace.filters
.. autosummary::
:toctree: generated

StandardFilter
UnivariateFilter
SteadyStateFilter
KalmanSmoother
SingleTimeseriesFilter
CholeskyFilter
LinearGaussianStateSpace
20 changes: 20 additions & 0 deletions docs/statespace/models.rst
Original file line number Diff line number Diff line change
@@ -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
16 changes: 16 additions & 0 deletions docs/statespace/models/structural.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
*********************
Structural Components
*********************

.. automodule:: pymc_experimental.statespace.models.structural
.. autosummary::
:toctree: generated

LevelTrendComponent
AutoregressiveComponent
TimeSeasonality
FrequencySeasonality
MeasurementError

StructuralTimeSeries
Component
Loading

0 comments on commit 67a9695

Please sign in to comment.