code: def one_step_predictive(
      model,
      observed_time_series,
      parameter_samples,
      timesteps_are_event_shape=True,
  ):
      """Compute one-step-ahead predictive distributions for all timesteps.
  
      Given samples from the posterior over parameters, return the predictive
      distribution over observations at each time `T`, given observations up
      through time `T-1`.
  
      Args:
        model: An instance of `StructuralTimeSeries` representing a
          time-series model. This represents a joint distribution over
          time-series and their parameters with batch shape `[b1, ..., bN]`.
        observed_time_series: `float` `Tensor` of shape
          `concat([sample_shape, model.batch_shape, [num_timesteps, 1]])` where
          `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
          dimension may (optionally) be omitted if `num_timesteps > 1`. Any `NaN`s
            are interpreted as missing observations; missingness may be also be
            explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
        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.
        timesteps_are_event_shape: Deprecated, for backwards compatibility only.
          If `False`, the predictive distribution will return per-timestep
          probabilities
          Default value: `True`.
  
      Returns:
        predictive_dist: a `tfd.MixtureSameFamily` instance with event shape
          `[num_timesteps] if timesteps_are_event_shape else []` and
          batch shape `concat([sample_shape, model.batch_shape,
          [] if timesteps_are_event_shape else [num_timesteps])`, with
          `num_posterior_draws` mixture components. The `t`th step represents the
          forecast distribution `p(observed_time_series[t] |
          observed_time_series[0:t-1], parameter_samples)`.
  
      #### Examples
  
      Suppose we've built a model and fit it to data using HMC:
  
      ```python
        day_of_week = tfp.sts.Seasonal(
            num_seasons=7,
            observed_time_series=observed_time_series,
            name='day_of_week')
        local_linear_trend = tfp.sts.LocalLinearTrend(
            observed_time_series=observed_time_series,
            name='local_linear_trend')
        model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                            observed_time_series=observed_time_series)
  
        samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
      ```
  
      Passing the posterior samples into `one_step_predictive`, we construct a
      one-step-ahead predictive distribution:
  
      ```python
        one_step_predictive_dist = tfp.sts.one_step_predictive(
          model, observed_time_series, parameter_samples=samples)
  
        predictive_means = one_step_predictive_dist.mean()
        predictive_scales = one_step_predictive_dist.stddev()
      ```
  
      If using variational inference instead of HMC, we'd construct a forecast using
      samples from the variational posterior:
  
      ```python
        surrogate_posterior = tfp.sts.build_factored_surrogate_posterior(
          model=model)
        loss_curve = tfp.vi.fit_surrogate_posterior(
          target_log_prob_fn=model.joint_distribution(observed_time_series).log_prob,
          surrogate_posterior=surrogate_posterior,
          optimizer=tf_keras.optimizers.Adam(learning_rate=0.1),
          num_steps=200)
        samples = surrogate_posterior.sample(30)
  
        one_step_predictive_dist = tfp.sts.one_step_predictive(
          model, observed_time_series, parameter_samples=samples)
      ```
  
      We can visualize the forecast by plotting:
  
      ```python
        from matplotlib import pylab as plt
        def plot_one_step_predictive(observed_time_series,
                                     forecast_mean,
                                     forecast_scale):
          plt.figure(figsize=(12, 6))
          num_timesteps = forecast_mean.shape[-1]
          c1, c2 = (0.12, 0.47, 0.71), (1.0, 0.5, 0.05)
          plt.plot(observed_time_series, label="observed time series", color=c1)
          plt.plot(forecast_mean, label="one-step prediction", color=c2)
          plt.fill_between(np.arange(num_timesteps),
                           forecast_mean - 2 * forecast_scale,
                           forecast_mean + 2 * forecast_scale,
                           alpha=0.1, color=c2)
          plt.legend()
  
        plot_one_step_predictive(observed_time_series,
                                 forecast_mean=predictive_means,
                                 forecast_scale=predictive_scales)
      ```
  
      To detect anomalous timesteps, we check whether the observed value at each
      step is within a 95% predictive interval, i.e., two standard deviations from
      the mean:
  
      ```python
        z_scores = ((observed_time_series[..., 1:] - predictive_means[..., :-1])
                     / predictive_scales[..., :-1])
        anomalous_timesteps = tf.boolean_mask(
            tf.range(1, num_timesteps),
            tf.abs(z_scores) > 2.0)
      ```
  
      """
  
      with tf.name_scope("one_step_predictive"):
          [
              observed_time_series,
              is_missing,
          ] = sts_util.canonicalize_observed_time_series_with_mask(
              observed_time_series
          )
  
          # Run filtering over the training timesteps to extract the
          # predictive means and variances.
          num_timesteps = ps.dimension_size(observed_time_series, -2)
          lgssm = tfe_util.JitPublicMethods(
              model.make_state_space_model(
                  num_timesteps=num_timesteps,
                  param_vals=parameter_samples,
              ),
              trace_only=True,
          )  # Avoid eager overhead w/o introducing XLA dependence.
          (
              _,
              _,
              _,
              _,
              _,
              observation_means,
              observation_covs,
          ) = lgssm.forward_filter(
              observed_time_series, mask=is_missing
          )
  
          # Squeeze dims to convert from LGSSM's event shape `[num_timesteps, 1]`
          # to a scalar time series.
          predictive_dist = sts_util.mix_over_posterior_draws(
              means=observation_means[..., 0],
              variances=observation_covs[..., 0, 0],
          )
          if timesteps_are_event_shape:
              predictive_dist = independent.Independent(
                  predictive_dist, reinterpreted_batch_ndims=1
              )
          return predictive_dist
  
file_path: probability/tensorflow_probability/python/sts/forecast.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.experimental import util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.sts.internal import util
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 2 for calling a global function with example usage: one_step_predictive(model, observed_time_series, parameter_samples, timesteps_are_event_shape) and returns: predictive_dist