tensorflow_probability/python/sts/decomposition.py
def _decompose_from_posterior_marginals(
    model, posterior_means, posterior_covs, parameter_samples):
  """Utility method to decompose a joint posterior into components.

  Args:
    model: `tfp.sts.Sum` instance defining an additive STS model.
    posterior_means: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps, latent_size])`
      representing the posterior mean over latents in an
      `AdditiveStateSpaceModel`.
    posterior_covs: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps,
      latent_size, latent_size])`
      representing the posterior marginal covariances over latents in an
      `AdditiveStateSpaceModel`.
    parameter_samples: Python `list` of `Tensors` representing posterior
      samples of model parameters, with shapes `[concat([
      [num_posterior_draws], param.prior.batch_shape,
      param.prior.event_shape]) for param in model.parameters]`. This may
      optionally also be a map (Python `dict`) of parameter names to
      `Tensor` values.

  Returns:
    component_dists: A `collections.OrderedDict` instance mapping
      component StructuralTimeSeries instances (elements of `model.components`)
      to `tfd.Distribution` instances representing the posterior marginal
      distributions on the process modeled by each component. Each distribution
      has batch shape matching that of `posterior_means`/`posterior_covs`, and
      event shape of `[num_timesteps]`.
  """

  try:
    model.components
  except AttributeError:
    raise ValueError('Model decomposed into components must be an instance of'
                     '`tfp.sts.Sum` (passed model {})'.format(model))

  with tf.compat.v1.name_scope('decompose_from_posterior_marginals'):

    # Extract the component means/covs from the joint latent posterior.
    latent_sizes = [component.latent_size for component in model.components]
    component_means = tf.split(posterior_means, latent_sizes, axis=-1)
    component_covs = _split_covariance_into_marginals(
        posterior_covs, latent_sizes)

    # Instantiate per-component state space models, and use them to push the
    # posterior means/covs through the observation model for each component.
    num_timesteps = dist_util.prefer_static_value(
        tf.shape(input=posterior_means))[-2]
    component_ssms = model.make_component_state_space_models(
        num_timesteps=num_timesteps,
        param_vals=parameter_samples)
    component_predictive_dists = collections.OrderedDict()
    for (component, component_ssm,
         component_mean, component_cov) in zip(model.components, component_ssms,
                                               component_means, component_covs):
      component_obs_mean, component_obs_cov = (
          component_ssm.latents_to_observations(
              latent_means=component_mean,
              latent_covs=component_cov))

      # Using the observation means and covs, build a mixture distribution
      # that integrates over the posterior draws.
      component_predictive_dists[component] = sts_util.mix_over_posterior_draws(
          means=component_obs_mean[..., 0],
          variances=component_obs_cov[..., 0, 0])
  return component_predictive_dists
def _decompose_from_posterior_marginals(
    model, posterior_means, posterior_covs, parameter_samples):
  """Utility method to decompose a joint posterior into components.

  Args:
    model: `tfp.sts.Sum` instance defining an additive STS model.
    posterior_means: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps, latent_size])`
      representing the posterior mean over latents in an
      `AdditiveStateSpaceModel`.
    posterior_covs: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps,
      latent_size, latent_size])`
      representing the posterior marginal covariances over latents in an
      `AdditiveStateSpaceModel`.
    parameter_samples: Python `list` of `Tensors` representing posterior
      samples of model parameters, with shapes `[concat([
      [num_posterior_draws], param.prior.batch_shape,
      param.prior.event_shape]) for param in model.parameters]`. This may
      optionally also be a map (Python `dict`) of parameter names to
      `Tensor` values.

  Returns:
    component_dists: A `collections.OrderedDict` instance mapping
      component StructuralTimeSeries instances (elements of `model.components`)
      to `tfd.Distribution` instances representing the posterior marginal
      distributions on the process modeled by each component. Each distribution
      has batch shape matching that of `posterior_means`/`posterior_covs`, and
      event shape of `[num_timesteps]`.
  """

  try:
    model.components
  except AttributeError:
    raise ValueError('Model decomposed into components must be an instance of'
                     '`tfp.sts.Sum` (passed model {})'.format(model))

  with tf.compat.v1.name_scope('decompose_from_posterior_marginals'):

    # Extract the component means/covs from the joint latent posterior.
    latent_sizes = [component.latent_size for component in model.components]
    component_means = tf.split(posterior_means, latent_sizes, axis=-1)
    component_covs = _split_covariance_into_marginals(
        posterior_covs, latent_sizes)

    # Instantiate per-component state space models, and use them to push the
    # posterior means/covs through the observation model for each component.
    num_timesteps = dist_util.prefer_static_value(
        tf.shape(input=posterior_means))[-2]
    component_ssms = model.make_component_state_space_models(
        num_timesteps=num_timesteps,
        param_vals=parameter_samples)
    component_predictive_dists = collections.OrderedDict()
    for (component, component_ssm,
         component_mean, component_cov) in zip(model.components, component_ssms,
                                               component_means, component_covs):
      component_obs_mean, component_obs_cov = (
          component_ssm.latents_to_observations(
              latent_means=component_mean,
              latent_covs=component_cov))

      # Using the observation means and covs, build a mixture distribution
      # that integrates over the posterior draws.
      component_predictive_dists[component] = sts_util.mix_over_posterior_draws(
          means=component_obs_mean[..., 0],
          variances=component_obs_cov[..., 0, 0])
  return component_predictive_dists
Utility method to decompose a joint posterior into components.

  Args:
    model: `tfp.sts.Sum` instance defining an additive STS model.
    posterior_means: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps, latent_size])`
      representing the posterior mean over latents in an
      `AdditiveStateSpaceModel`.
    posterior_covs: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps,
      latent_size, latent_size])`
      representing the posterior marginal covariances over latents in an
      `AdditiveStateSpaceModel`.
    parameter_samples: Python `list` of `Tensors` representing posterior
      samples of model parameters, with shapes `[concat([
      [num_posterior_draws], param.prior.batch_shape,
      param.prior.event_shape]) for param in model.parameters]`. This may
      optionally also be a map (Python `dict`) of parameter names to
      `Tensor` values.

  Returns:
    component_dists: A `collections.OrderedDict` instance mapping
      component StructuralTimeSeries instances (elements of `model.components`)
      to `tfd.Distribution` instances representing the posterior marginal
      distributions on the process modeled by each component. Each distribution
      has batch shape matching that of `posterior_means`/`posterior_covs`, and
      event shape of `[num_timesteps]`.