How to Decompose and Visualise Slope Component in Tensorflow Probability

Question

I'm running tensorflow 2.1 and tensorflow_probability 0.9. I have fit a Structural Time Series Model with a seasonal component. I am using code from the Tensorflow Probability Structural Time Series Probability example: Tensorflow Github.

In the example there is a great plot where the decomposition is visualised:

# Get the distributions over component outputs from the posterior marginals on
# training data, and from the forecast model.
component_dists = sts.decompose_by_component(
    demand_model,
    observed_time_series=demand_training_data,
    parameter_samples=q_samples_demand_)

forecast_component_dists = sts.decompose_forecast_by_component(
    demand_model,
    forecast_dist=demand_forecast_dist,
    parameter_samples=q_samples_demand_)




demand_component_means_, demand_component_stddevs_ = (
    {k.name: c.mean() for k, c in component_dists.items()},
    {k.name: c.stddev() for k, c in component_dists.items()})

(
    demand_forecast_component_means_,
    demand_forecast_component_stddevs_
) = (
    {k.name: c.mean() for k, c in forecast_component_dists.items()},
    {k.name: c.stddev() for k, c in forecast_component_dists.items()}
    )

When using a trend component, is it possible to decompose and visualise both:

trend/_level_scale & trend/_slope_scale

I have tried many permutations to extract the nested element of the trend component with no luck.

Thanks for your time in advance.

Answer 1

We didn't write a separate STS interface for this, but you can access the posterior on latent states (in this case, both the level and slope) by directly querying the underlying state-space model for its marginal means and covariances:

ssm = model.make_state_space_model(
        num_timesteps=num_timesteps,
        param_vals=parameter_samples)
posterior_means, posterior_covs = (
  ssm.posterior_marginals(observed_time_series))

You should also be able to draw samples from the joint posterior by running ssm.posterior_sample(observed_time_series, num_samples).

It looks like there's currently a glitch when drawing posterior samples from a model with no batch shape (Could not find valid device for node. Node:{{node Reshape}}): while we fix that, it should work to add an artificial batch dimension as a workaround: ssm.posterior_sample(observed_time_series[tf.newaxis, ...], num_samples).