Skip to content

Add sample_filter_outputs helper method to StateSpaceModel #521

Open
@jessegrabowski

Description

@jessegrabowski

Copied from @Dekermanjian 's comment here:

Getting the variance covariance matrix for the predicted|filtered|smoothed states using the conditional posterior is not trivial because there is no idata.cov() method and xarray.cov() only computes the covariances. This means you will need to compute the covariances then the variances and then build the variance covariance matrix manually. It would be nice to have a utility function in state space that would handle this for the user.

Here is a thumbnail sketch of what we're looking for:

with pm.Model(coords=ss_mod.coords) as m:
    ss_mod._build_dummy_graph()
    ss_mod._insert_random_variables()
    
    x0, P0, c, d, T, Z, R, H, Q = ss_matrices = ss_mod.unpack_statespace()
    filter_outputs = ss_mod.kalman_filter.build_graph(pt.as_tensor(ss_mod._fit_data), 
                                                      *ss_matrices)
    smoother_outputs = ss_mod.kalman_smoother.build_graph(T, R, Q, filter_outputs[0], filter_outputs[3])
    for output in filter_outputs[:-1] + list(smoother_outputs):
        pm.Deterministic(output.name, output)
        
    idata_filter = pm.sample_posterior_predictive(idata, 
                                                  var_names=[x.name for x in m.deterministics],
                                                  compile_kwargs={'mode':'JAX'})

Obviously it needs a bit more to support exogenous data and dims and stuff. But this is enough to get someone started on a PR :)

After we have this utility, we should remove the return_kalman_filter_outputs_in_idata flag in build_statespace_graph. People should just use this utility to get the filtered outputs from the idata if they want it.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions