Description
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 andxarray.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.