def build_model_for_gibbs_fitting(observed_time_series, design_matrix,
                                  weights_prior, level_variance_prior,
                                  observation_noise_variance_prior):
    """Builds a StructuralTimeSeries model instance that supports Gibbs sampling.

  To support Gibbs sampling, a model must have have conjugate priors on all
  scale and weight parameters, and must be constructed so that
  `model.parameters` matches the parameters and ordering specified by the
  the `GibbsSamplerState` namedtuple. Currently, this includes (only) models
  consisting of the sum of a LocalLevel and a LinearRegression component.

  Args:
    observed_time_series: optional `float` `Tensor` of shape [..., T, 1]`
      (omitting the trailing unit dimension is also supported when `T > 1`),
      specifying an observed time series. May optionally be an instance of
      `tfp.sts.MaskedTimeSeries`, which includes a mask `Tensor` to specify
      timesteps with missing observations.
    design_matrix: float `Tensor` of shape `concat([batch_shape,
      [num_timesteps, num_features]])`. This may also optionally be
      an instance of `tf.linalg.LinearOperator`.
    weights_prior: An instance of `tfd.Normal` representing a scalar prior on
      each regression weight. May have batch shape broadcastable to the batch
      shape of `observed_time_series`.
    level_variance_prior: An instance of `tfd.InverseGamma` representing a prior
      on the level variance (`level_scale**2`) of a local level model. May have
      batch shape broadcastable to the batch shape of `observed_time_series`.
    observation_noise_variance_prior: An instance of `tfd.InverseGamma`
      representing a prior on the observation noise variance (
      `observation_noise_scale**2`). May have batch shape broadcastable to the
      batch shape of `observed_time_series`.
  Returns:
    model: A `tfp.sts.StructuralTimeSeries` model instance.
  """
    if not isinstance(weights_prior, tfd.Normal):
        raise ValueError(
            'Weights prior must be a univariate normal distribution.')
    if not isinstance(level_variance_prior, tfd.InverseGamma):
        raise ValueError(
            'Level variance prior must be an inverse gamma distribution.')
    if not isinstance(observation_noise_variance_prior, tfd.InverseGamma):
        raise ValueError('Observation noise variance prior must be an inverse '
                         'gamma distribution.')

    sqrt = tfb.Invert(
        tfb.Square())  # Converts variance priors to scale priors.
    local_level = sts.LocalLevel(observed_time_series=observed_time_series,
                                 level_scale_prior=sqrt(level_variance_prior),
                                 name='local_level')
    regression = sts.LinearRegression(design_matrix=design_matrix,
                                      weights_prior=weights_prior,
                                      name='regression')
    model = sts.Sum(
        [local_level, regression],
        observed_time_series=observed_time_series,
        observation_noise_scale_prior=sqrt(observation_noise_variance_prior),
        # The Gibbs sampling steps in this file do not account for an
        # offset to the observed series. Instead, we assume the
        # observed series has already been centered and
        # scale-normalized.
        constant_offset=0.)
    model.supports_gibbs_sampling = True
    return model
Example #2
0
def build_model_for_gibbs_fitting(observed_time_series,
                                  design_matrix,
                                  weights_prior,
                                  level_variance_prior,
                                  observation_noise_variance_prior,
                                  slope_variance_prior=None,
                                  initial_level_prior=None,
                                  sparse_weights_nonzero_prob=None):
    """Builds a StructuralTimeSeries model instance that supports Gibbs sampling.

  To support Gibbs sampling, a model must have have conjugate priors on all
  scale and weight parameters, and must be constructed so that
  `model.parameters` matches the parameters and ordering specified by the
  `GibbsSamplerState` namedtuple. Currently, this includes (only) models
  consisting of the sum of a LocalLevel or LocalLinearTrend component with
  (optionally) a LinearRegression or SpikeAndSlabSparseLinearRegression
  component.

  Args:
    observed_time_series: optional `float` `Tensor` of shape [..., T, 1]`
      (omitting the trailing unit dimension is also supported when `T > 1`),
      specifying an observed time series. May optionally be an instance of
      `tfp.sts.MaskedTimeSeries`, which includes a mask `Tensor` to specify
      timesteps with missing observations.
    design_matrix: Optional float `Tensor` of shape `concat([batch_shape,
      [num_timesteps, num_features]])`. This may also optionally be an instance
      of `tf.linalg.LinearOperator`. If None, no regression is done.
    weights_prior: Optional distribution instance specifying a normal prior on
      weights. This may be a multivariate normal instance with event shape
      `[num_features]`, or a scalar normal distribution with event shape `[]`.
      In either case, the batch shape must broadcast to the batch shape of
      `observed_time_series`. If a `sparse_weights_nonzero_prob` is specified,
      requesting sparse regression, then the `weights_prior` mean is ignored
      (because nonzero means are not currently implemented by the spike-and-slab
      sampler). In this case, `weights_prior=None` is also valid, and will use
      the default prior of the spike-and-slab sampler.
    level_variance_prior: An instance of `tfd.InverseGamma` representing a prior
      on the level variance (`level_scale**2`) of a local level model. May have
      batch shape broadcastable to the batch shape of `observed_time_series`.
    observation_noise_variance_prior: An instance of `tfd.InverseGamma`
      representing a prior on the observation noise variance (
      `observation_noise_scale**2`). May have batch shape broadcastable to the
      batch shape of `observed_time_series`.
    slope_variance_prior: Optional instance of `tfd.InverseGamma` representing a
      prior on slope variance (`slope_scale**2`) of a local linear trend model.
      May have batch shape broadcastable to the batch shape of
      `observed_time_series`. If specified, a local linear trend model is used
      rather than a local level model.
      Default value: `None`.
    initial_level_prior: optional `tfd.Distribution` instance specifying a
      prior on the initial level. If `None`, a heuristic default prior is
      constructed based on the provided `observed_time_series`.
      Default value: `None`.
    sparse_weights_nonzero_prob: Optional scalar float `Tensor` prior
      probability that any given feature has nonzero weight. If specified, this
      triggers a sparse regression with a spike-and-slab prior, where
      `sparse_weights_nonzero_prob` is the prior probability of the 'slab'
      component.
      Default value: `None`.

  Returns:
    model: A `tfp.sts.StructuralTimeSeries` model instance.
  """
    if design_matrix is None:
        if sparse_weights_nonzero_prob is not None:
            raise ValueError(
                'Design matrix is None thus sparse_weights_nonzero_prob should '
                'not be defined, as it will not be used.')
        if weights_prior is not None:
            raise ValueError(
                'Design matrix is None thus weights_prior should not be defined, '
                'as it will not be used.')

    if isinstance(weights_prior, tfd.Normal):
        # Canonicalize scalar normal priors as diagonal MVNs.
        # design_matrix must be defined, otherwise we threw an exception earlier.
        if isinstance(design_matrix, tf.linalg.LinearOperator):
            num_features = design_matrix.shape_tensor()[-1]
        else:
            num_features = prefer_static.dimension_size(design_matrix, -1)
        weights_prior = _tile_normal_to_mvn_diag(weights_prior, num_features)
    elif weights_prior is not None and not _is_multivariate_normal(
            weights_prior):
        raise ValueError(
            'Weights prior must be a normal distribution or `None`.')
    if not isinstance(level_variance_prior, tfd.InverseGamma):
        raise ValueError(
            'Level variance prior must be an inverse gamma distribution.')
    if (slope_variance_prior is not None
            and not isinstance(slope_variance_prior, tfd.InverseGamma)):
        raise ValueError(
            'Slope variance prior must be an inverse gamma distribution; got: {}.'
            .format(slope_variance_prior))
    if not isinstance(observation_noise_variance_prior, tfd.InverseGamma):
        raise ValueError('Observation noise variance prior must be an inverse '
                         'gamma distribution.')

    sqrt = tfb.Invert(
        tfb.Square())  # Converts variance priors to scale priors.
    components = []

    # Level or trend component.
    if slope_variance_prior:
        components.append(
            sts.LocalLinearTrend(observed_time_series=observed_time_series,
                                 level_scale_prior=sqrt(level_variance_prior),
                                 slope_scale_prior=sqrt(slope_variance_prior),
                                 initial_level_prior=initial_level_prior,
                                 name='local_linear_trend'))
    else:
        components.append(
            sts.LocalLevel(observed_time_series=observed_time_series,
                           level_scale_prior=sqrt(level_variance_prior),
                           initial_level_prior=initial_level_prior,
                           name='local_level'))

    # Regression component.
    if design_matrix is None:
        pass
    elif sparse_weights_nonzero_prob is not None:
        components.append(
            SpikeAndSlabSparseLinearRegression(
                design_matrix=design_matrix,
                weights_prior=weights_prior,
                sparse_weights_nonzero_prob=sparse_weights_nonzero_prob,
                name='sparse_regression'))
    else:
        components.append(
            sts.LinearRegression(design_matrix=design_matrix,
                                 weights_prior=weights_prior,
                                 name='regression'))
    model = sts.Sum(
        components,
        observed_time_series=observed_time_series,
        observation_noise_scale_prior=sqrt(observation_noise_variance_prior),
        # The Gibbs sampling steps in this file do not account for an
        # offset to the observed series. Instead, we assume the
        # observed series has already been centered and
        # scale-normalized.
        constant_offset=0.)
    model.supports_gibbs_sampling = True
    return model