UPDATE

This repo is no longer maintained! Statespace models are now a part of pymc-experimental, and are maintained by the PyMC development team. Please look over there for the most up-to-date Bayesian State Space models!

PyMC StateSpace

A system for Bayesian estimation of state space models using PyMC 5.0. This package is designed to mirror the functionality of the Statsmodels.api tsa.statespace module, except within a Bayesian estimation framework. To accomplish this, PyMC Statespace has a Kalman filter written in Pytensor, allowing the gradients of the iterative Kalman filter likelihood to be computed and provided to the PyMC NUTS sampler.

State Space Models

This package follows Statsmodels in using the Durbin and Koopman (2012) nomenclature for a linear state space model. Under this nomenclature, the model is written as:

The objects in the above equation have the following shapes and meanings:

• , observation vector

• , state vector

• , observation noise vector

• , state innovation vector

• , the design matrix

• , observation noise covariance matrix

• , the transition matrix

• , the selection matrix

• , the state innovation covariance matrix

• , the state intercept vector

• , the observation intercept vector

The linear state space model is a workhorse in many disciplines, and is flexible enough to represent a wide range of models, including Box-Jenkins SARIMAX class models, time series decompositions, and model of multiple time series (VARMAX) models. Use of a Kalman filter allows for estimation of unobserved and missing variables. Taken together, these are a powerful class of models that can be further augmented by Bayesian estimation. This allows the researcher to integrate uncertainty about the true model when estimating model parameteres.

Example Usage

Currently, local level and ARMA models are implemented. To initialize a model, simply create a model object, provide data, and any additional arguments unique to that model. In addition, you can choose from several Kalman Filter implementations. Currently these are standard, steady_state, univariate, single, and cholesky. For more information, see the docs (they're on the to-do list)

import pymc_statespace as pmss
#make sure data is a 2d numpy array!
arma_model = pmss.BayesianARMA(data = data[:, None], order=(1, 1), filter='standard', stationary_initalization=True)

You will see a message letting you know the model was created, as well as telling you the expected parameter names you will need to create inside a PyMC model block.

Model successfully initialized! The following parameters should be assigned priors inside a PyMC model block: x0, sigma_state, rho, theta

Next, a PyMC model is declared as usual, and the parameters can be passed into the state space model. This is done as follows:

with pm.Model() as arma_model:
    x0 = pm.Normal('x0', mu=0, sigma=1.0, shape=mod.k_states)
    sigma_state = pm.HalfNormal('sigma_state', sigma=1)

    rho = pm.TruncatedNormal('rho', mu=0.0, sigma=0.5, lower=-1.0, upper=1.0, shape=p)
    theta = pm.Normal('theta', mu=0.0, sigma=0.5, shape=q)

    arma_model.build_statespace_graph()
    trace_1 = pm.sample(init='jitter+adapt_diag_grad', target_accept=0.9)

After you place priors over the requested parameters, call arma_model.build_statespace_graph() to automatically put everything together into a state space system. If you are missing any parameters, you will get a error letting you know what else needs to be defined.

And that's it! After this, you can sample the PyMC model as normal.

Creating your own state space model

Creating a custom state space model isn't complicated. Once again, the API follows the Statsmodels implementation. All models need to subclass the PyMCStateSpace class, and pass three values into the class super construction: data (from which p is inferred), k_states (this is "m" in the shapes above), and k_posdef (this is "r" above). The user also needs to declare any required state space matrices. Here is an example of a simple local linear trend model:

def __init__(self, data):
    # Model shapes
    k_states = k_posdef = 2

    super().__init__(data, k_states, k_posdef)

    # Initialize the matrices
    self.ssm['design'] = np.array([[1.0, 0.0]])
    self.ssm['transition'] = np.array([[1.0, 1.0],
                                       [0.0, 1.0]])
    self.ssm['selection'] = np.eye(k_states)

    self.ssm['initial_state'] = np.array([[0.0],
                                          [0.0]])
    self.ssm['initial_state_cov'] = np.array([[1.0, 0.0],
                                              [0.0, 1.0]])

You will also need to set up the param_names class method, which returns a list of parameter names. This lets users know what priors they need to define, and lets the model gather the required parameters from the PyMC model context.

@property
def param_names(self):
    return ['x0', 'P0', 'sigma_obs', 'sigma_state']

self.ssm is an PytensorRepresentation class object that is created by the super constructor. Every model has a self.ssm and a self.kalman_filter created after the super constructor is called. All the matrices stored in self.ssm are Pytensor tensor variables, but numpy arrays can be passed to them for convenience. Behind the scenes, they will be converted to Pytensor tensors.

Note that the names of the matrices correspond to the names listed above. They are (in the same order):

• Z = design
• H = obs_cov
• T = transition
• R = selection
• Q = state_cov
• c = state_intercept
• d = obs_intercept
• a1 = initial_state
• P1 = initial_state_cov

Indexing by name only will expose the entire matrix. A name can also be followed by the usual numpy slice notation to get a specific element, row, or column.

The user also needs to implement an update method, which takes in a single Pytensor tensor as an argument. This method routes the parameters estimated by PyMC into the right spots in the state space matrices. The local level has at least two parameters to estimate: the variance of the level state innovations, and the variance of the trend state innovations. Here is the corresponding update method:

def update(self, theta: at.TensorVariable) -> None:
    # Observation covariance
    self.ssm['obs_cov', 0, 0] = theta[0]

    # State covariance
    self.ssm['state_cov', np.diag_indices(2)] = theta[1:]

This function is why the order matters when flattening and concatenating the random variables inside the PyMC model. In this case, we must first pass sigma_obs, followed by sigma_level, then sigma_trend.

But that's it! Obviously this API isn't great, and will be subject to change as the package evolves, but it should be enough to get a motivated research going. Happy estimation, and let me know all the bugs you find by opening an issue.