Example #1
0
    def test_empirical_statistics_accepts_masked_values(self):

        # Ensure that masks broadcast over batch shape by creating a batch of
        # time series.
        time_series = np.random.randn(3, 2, 5)
        mask = np.array([True, False, False, True, False])

        masked_series = missing_values_util.MaskedTimeSeries(
            time_series=time_series, is_missing=mask)
        mean, stddev, initial = self.evaluate(
            sts_util.empirical_statistics(masked_series))

        broadcast_mask = np.broadcast_to(mask, time_series.shape)
        unmasked_series = time_series[~broadcast_mask].reshape([3, 2, 3])
        unmasked_mean, unmasked_stddev, unmasked_initial = self.evaluate(
            sts_util.empirical_statistics(unmasked_series))
        self.assertAllClose(mean, unmasked_mean)
        self.assertAllClose(stddev, unmasked_stddev)
        self.assertAllClose(initial, unmasked_initial)

        # Run the same tests without batch shape.
        unbatched_time_series = time_series[0, 0, :]
        masked_series = missing_values_util.MaskedTimeSeries(
            time_series=unbatched_time_series, is_missing=mask)
        mean, stddev, initial = self.evaluate(
            sts_util.empirical_statistics(masked_series))
        unmasked_mean, unmasked_stddev, unmasked_initial = self.evaluate(
            sts_util.empirical_statistics(unbatched_time_series[~mask]))
        self.assertAllClose(mean, unmasked_mean)
        self.assertAllClose(stddev, unmasked_stddev)
        self.assertAllClose(initial, unmasked_initial)
def _detect_anomalies_inner(observed_time_series,
                            seasonal_structure,
                            anomaly_threshold=0.01,
                            num_warmup_steps=50,
                            num_samples=100,
                            use_gibbs_predictive_dist=True,
                            seed=None):
    """Helper function for `detect_anomalies` to cache `tf.function` traces."""
    with tf.name_scope('build_default_model_for_gibbs_sampling'):
        observed_mean, observed_stddev, _ = sts_util.empirical_statistics(
            observed_time_series)
        # Center the series to have mean 0 and stddev 1. Alternately, we could
        # rescale the priors to match the series, but this is simpler.
        observed_mean = observed_mean[...,
                                      tf.newaxis]  # Broadcast with num_steps.
        observed_stddev = observed_stddev[..., tf.newaxis]
        observed_time_series = observed_time_series._replace(
            time_series=(observed_time_series.time_series -
                         observed_mean[..., tf.newaxis]) /
            observed_stddev[..., tf.newaxis])

        model, posterior_samples = _fit_seasonal_model_with_gibbs_sampling(
            observed_time_series,
            seasonal_structure=seasonal_structure,
            num_results=num_samples,
            num_warmup_steps=num_warmup_steps,
            seed=seed)
        parameter_samples = _parameter_samples_from_gibbs_posterior(
            model, posterior_samples)
        if use_gibbs_predictive_dist:
            predictive_dist = gibbs_sampler.one_step_predictive(
                model, posterior_samples)
        else:
            # Rebuild the model using appropriate Seasonal components, throwing away
            # the seasonal effects fit by the Gibbs sampler. Instead, the predictive
            # model incorporates the seasonal effects in its state space, so their
            # posterior is tracked exactly (rather than by sampling) and is updated
            # with each step of the series. This avoids the learning-from-the-future
            # behavior of the Gibbs-estimated effects.
            seasonal_model = (
                model.components[0] +  # LocalLinearTrend component.
                default_model.model_from_seasonal_structure(
                    seasonal_structure,
                    observed_time_series,
                    # Gibbs sampling didn't fit a drift scale(s).
                    allow_drift=False))
            predictive_dist = one_step_predictive(
                seasonal_model,
                observed_time_series,
                timesteps_are_event_shape=False,
                parameter_samples=parameter_samples)

        prob_lower = predictive_dist.cdf(observed_time_series.time_series[...,
                                                                          0])
        tail_probabilities = 2 * tf.minimum(prob_lower, 1 - prob_lower)
        lower_limit, upper_limit, predictive_mean = compute_predictive_bounds(
            predictive_dist, anomaly_threshold=anomaly_threshold)
        restore_scale = lambda x: x * observed_stddev + observed_mean
        return (restore_scale(lower_limit), restore_scale(upper_limit),
                restore_scale(predictive_mean), tail_probabilities)
Example #3
0
    def test_empirical_statistics_accepts_masked_values(self):

        # Ensure that masks broadcast over batch shape by creating a batch of
        # time series.
        time_series = np.random.randn(3, 2, 5)
        mask = np.array([[True, False, False, True, False],
                         [True, True, True, True, True]])

        masked_series = missing_values_util.MaskedTimeSeries(
            time_series=time_series, is_missing=mask)
        mean, stddev, initial = self.evaluate(
            sts_util.empirical_statistics(masked_series))

        # Should return default values when the series is completely masked.
        self.assertAllClose(mean[:, 1], tf.zeros_like(mean[:, 1]))
        self.assertAllClose(stddev[:, 1], tf.ones_like(stddev[:, 1]))
        self.assertAllClose(initial[:, 1], tf.zeros_like(initial[:, 1]))

        # Otherwise, should return the actual mean/stddev/initial values.
        time_series = time_series[:, 0, :]
        mask = mask[0, :]
        broadcast_mask = np.broadcast_to(mask, time_series.shape)
        unmasked_series = time_series[~broadcast_mask].reshape([3, 3])
        unmasked_mean, unmasked_stddev, unmasked_initial = self.evaluate(
            sts_util.empirical_statistics(unmasked_series))
        self.assertAllClose(mean[:, 0], unmasked_mean)
        self.assertAllClose(stddev[:, 0], unmasked_stddev)
        self.assertAllClose(initial[:, 0], unmasked_initial)

        # Run the same tests without batch shape.
        unbatched_time_series = time_series[0, :]
        masked_series = missing_values_util.MaskedTimeSeries(
            time_series=unbatched_time_series, is_missing=mask)
        mean, stddev, initial = self.evaluate(
            sts_util.empirical_statistics(masked_series))
        unmasked_mean, unmasked_stddev, unmasked_initial = self.evaluate(
            sts_util.empirical_statistics(unbatched_time_series[~mask]))
        self.assertAllClose(mean, unmasked_mean)
        self.assertAllClose(stddev, unmasked_stddev)
        self.assertAllClose(initial, unmasked_initial)
Example #4
0
    def __init__(self,
                 components,
                 constant_offset=None,
                 observation_noise_scale_prior=None,
                 observed_time_series=None,
                 name=None):
        """Specify a structural time series model representing a sum of components.

    Args:
      components: Python `list` of one or more StructuralTimeSeries instances.
        These must have unique names.
      constant_offset: optional `float` `Tensor` of shape broadcasting to
        `concat([batch_shape, [num_timesteps]]`) specifying a constant value
        added to the sum of outputs from the component models.
        This allows the components to model the shifted series
        `observed_time_series - constant_offset`. If `None`, this is set to the
        mean of the provided `observed_time_series`.
        Default value: `None`.
      observation_noise_scale_prior: optional `tfd.Distribution` instance
        specifying a prior on `observation_noise_scale`. If `None`, a heuristic
        default prior is constructed based on the provided
        `observed_time_series`.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series. This is
        used to set the constant offset, if not provided, and to construct a
        default heuristic `observation_noise_scale_prior` if not provided. May
        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
        a mask `Tensor` to specify timesteps with missing observations.
        Default value: `None`.
      name: Python `str` name of this model component; used as `name_scope`
        for ops created by this class.
        Default value: 'Sum'.

    Raises:
      ValueError: if components do not have unique names.
    """

        with tf.name_scope(name or 'Sum') as name:
            if observed_time_series is not None:
                observed_mean, observed_stddev, _ = (
                    sts_util.empirical_statistics(observed_time_series))
            else:
                observed_mean, observed_stddev = 0., 1.

            if observation_noise_scale_prior is None:
                observation_noise_scale_prior = tfd.LogNormal(loc=tf.math.log(
                    .01 * observed_stddev),
                                                              scale=2.)

            dtype = dtype_util.common_dtype([
                constant_offset, observation_noise_scale_prior, observed_mean,
                observed_stddev
            ])

            # Ensure that offsets have canonical shape `[..., num_timesteps]`.
            if constant_offset is None:
                constant_offset = tf.convert_to_tensor(observed_mean,
                                                       dtype=dtype)[...,
                                                                    tf.newaxis]
            constant_offset *= tf.ones([1], dtype=dtype)

            # Check that components have unique names, to ensure that inherited
            # parameters will be assigned unique names.
            component_names = [c.name for c in components]
            if len(component_names) != len(set(component_names)):
                raise ValueError(
                    'Components must have unique names: {}'.format(
                        component_names))
            components_by_name = collections.OrderedDict([(c.name, c)
                                                          for c in components])

            # Build parameters list for the combined model, by inheriting parameters
            # from the component models in canonical order.
            parameters = [
                Parameter(
                    'observation_noise_scale', observation_noise_scale_prior,
                    tfb.Chain(
                        [tfb.Scale(scale=observed_stddev),
                         tfb.Softplus()]))
            ]
            for component in components:
                for parameter in component.parameters:
                    parameters.append(
                        Parameter(name='{}_{}'.format(component.name,
                                                      parameter.name),
                                  prior=parameter.prior,
                                  bijector=parameter.bijector))

            self._components = components
            self._components_by_name = components_by_name
            self._constant_offset = constant_offset

            super(Sum, self).__init__(parameters=parameters,
                                      latent_size=sum([
                                          component.latent_size
                                          for component in components
                                      ]),
                                      name=name)
Example #5
0
  def __init__(self,
               order,
               coefficients_prior=None,
               level_scale_prior=None,
               initial_state_prior=None,
               coefficient_constraining_bijector=None,
               observed_time_series=None,
               name=None):
    """Specify an autoregressive model.

    Args:
      order: scalar Python positive `int` specifying the number of past
        timesteps to regress on.
      coefficients_prior: optional `tfd.Distribution` instance specifying a
        prior on the `coefficients` parameter. If `None`, a default standard
        normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([order]))`) prior
        is used.
        Default value: `None`.
      level_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `level_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      initial_state_prior: optional `tfd.Distribution` instance specifying a
        prior on the initial state, corresponding to the values of the process
        at a set of size `order` of imagined timesteps before the initial step.
        If `None`, a heuristic default prior is constructed based on the
        provided `observed_time_series`.
        Default value: `None`.
      coefficient_constraining_bijector: optional `tfb.Bijector` instance
        representing a constraining mapping for the autoregressive coefficients.
        For example, `tfb.Tanh()` constrains the coefficients to lie in
        `(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
        `tfb.Identity()` implies no constraint. If `None`, the default behavior
        constrains the coefficients to lie in `(-1, 1)` using a `Tanh` bijector.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series). May
        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
        a mask `Tensor` to specify timesteps with missing observations.
        Default value: `None`.
      name: the name of this model component.
        Default value: 'Autoregressive'.
    """
    with tf.name_scope(name or 'Autoregressive') as name:
      masked_time_series = None
      if observed_time_series is not None:
        masked_time_series = (
            sts_util.canonicalize_observed_time_series_with_mask(
                observed_time_series))

      dtype = dtype_util.common_dtype(
          [(masked_time_series.time_series
            if masked_time_series is not None else None),
           coefficients_prior,
           level_scale_prior,
           initial_state_prior], dtype_hint=tf.float32)

      if observed_time_series is not None:
        _, observed_stddev, observed_initial = sts_util.empirical_statistics(
            masked_time_series)
      else:
        observed_stddev, observed_initial = (
            tf.convert_to_tensor(value=1., dtype=dtype),
            tf.convert_to_tensor(value=0., dtype=dtype))
      batch_ones = tf.ones(tf.concat([
          tf.shape(observed_initial),  # Batch shape
          [order]], axis=0), dtype=dtype)

      # Heuristic default priors. Overriding these may dramatically
      # change inference performance and results.
      if coefficients_prior is None:
        coefficients_prior = tfd.MultivariateNormalDiag(
            scale_diag=batch_ones)
      if level_scale_prior is None:
        level_scale_prior = tfd.LogNormal(
            loc=tf.math.log(0.05 *  observed_stddev), scale=3.)

      if (coefficients_prior.event_shape.is_fully_defined() and
          order != coefficients_prior.event_shape[0]):
        raise ValueError("Prior dimension {} doesn't match order {}.".format(
            coefficients_prior.event_shape[0], order))

      if initial_state_prior is None:
        initial_state_prior = tfd.MultivariateNormalDiag(
            loc=observed_initial[..., tf.newaxis] * batch_ones,
            scale_diag=(tf.abs(observed_initial) +
                        observed_stddev)[..., tf.newaxis] * batch_ones)

      self._order = order
      self._coefficients_prior = coefficients_prior
      self._level_scale_prior = level_scale_prior
      self._initial_state_prior = initial_state_prior

      if coefficient_constraining_bijector is None:
        coefficient_constraining_bijector = tfb.Tanh()
      super(Autoregressive, self).__init__(
          parameters=[
              Parameter('coefficients',
                        coefficients_prior,
                        coefficient_constraining_bijector),
              Parameter('level_scale', level_scale_prior,
                        tfb.Chain([tfb.AffineScalar(scale=observed_stddev),
                                   tfb.Softplus()]))
          ],
          latent_size=order,
          name=name)
    def __init__(self,
                 level_scale_prior=None,
                 slope_mean_prior=None,
                 slope_scale_prior=None,
                 autoregressive_coef_prior=None,
                 initial_level_prior=None,
                 initial_slope_prior=None,
                 observed_time_series=None,
                 constrain_ar_coef_stationary=True,
                 constrain_ar_coef_positive=False,
                 name=None):
        """Specify a semi-local linear trend model.

    Args:
      level_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `level_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      slope_mean_prior: optional `tfd.Distribution` instance specifying a prior
        on the `slope_mean` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      slope_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `slope_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      autoregressive_coef_prior: optional `tfd.Distribution` instance specifying
        a prior on the `autoregressive_coef` parameter. If `None`, the default
        prior is a standard `Normal(0., 1.)`. Note that the prior may be
        implicitly truncated by `constrain_ar_coef_stationary` and/or
        `constrain_ar_coef_positive`.
        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`.
      initial_slope_prior: optional `tfd.Distribution` instance specifying a
        prior on the initial slope. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series). May
        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
        a mask `Tensor` to specify timesteps with missing observations.
        Default value: `None`.
      constrain_ar_coef_stationary: if `True`, perform inference using a
        parameterization that restricts `autoregressive_coef` to the interval
        `(-1, 1)`, or `(0, 1)` if `force_positive_ar_coef` is also `True`,
        corresponding to stationary processes. This will implicitly truncates
        the support of `autoregressive_coef_prior`.
        Default value: `True`.
      constrain_ar_coef_positive: if `True`, perform inference using a
        parameterization that restricts `autoregressive_coef` to be positive,
        or in `(0, 1)` if `constrain_ar_coef_stationary` is also `True`. This
        will implicitly truncate the support of `autoregressive_coef_prior`.
        Default value: `False`.
      name: the name of this model component.
        Default value: 'SemiLocalLinearTrend'.
    """

        with tf.name_scope(name or 'SemiLocalLinearTrend') as name:
            if observed_time_series is not None:
                _, observed_stddev, observed_initial = sts_util.empirical_statistics(
                    observed_time_series)
            else:
                observed_stddev, observed_initial = 1., 0.

            # Heuristic default priors. Overriding these may dramatically
            # change inference performance and results.
            if level_scale_prior is None:
                level_scale_prior = tfd.LogNormal(loc=tf.math.log(
                    .01 * observed_stddev),
                                                  scale=2.)
            if slope_mean_prior is None:
                slope_mean_prior = tfd.Normal(loc=0., scale=observed_stddev)
            if slope_scale_prior is None:
                slope_scale_prior = tfd.LogNormal(loc=tf.math.log(
                    .01 * observed_stddev),
                                                  scale=2.)
            if autoregressive_coef_prior is None:
                autoregressive_coef_prior = tfd.Normal(
                    loc=0., scale=tf.ones_like(observed_initial))
            if initial_level_prior is None:
                initial_level_prior = tfd.Normal(
                    loc=observed_initial,
                    scale=tf.abs(observed_initial) + observed_stddev)
            if initial_slope_prior is None:
                initial_slope_prior = tfd.Normal(loc=0., scale=observed_stddev)

            self._initial_state_prior = tfd.MultivariateNormalDiag(
                loc=tf.stack(
                    [initial_level_prior.mean(),
                     initial_slope_prior.mean()],
                    axis=-1),
                scale_diag=tf.stack([
                    initial_level_prior.stddev(),
                    initial_slope_prior.stddev()
                ],
                                    axis=-1))

            # Constrain the support of the autoregressive coefficient.
            if constrain_ar_coef_stationary and constrain_ar_coef_positive:
                autoregressive_coef_bijector = tfb.Sigmoid(
                )  # support in (0, 1)
            elif constrain_ar_coef_positive:
                autoregressive_coef_bijector = tfb.Softplus(
                )  # support in (0, infty)
            elif constrain_ar_coef_stationary:
                autoregressive_coef_bijector = tfb.Tanh()  # support in (-1, 1)
            else:
                autoregressive_coef_bijector = tfb.Identity()  # unconstrained

            stddev_preconditioner = tfb.Scale(scale=observed_stddev)
            scaled_softplus = tfb.Chain(
                [stddev_preconditioner, tfb.Softplus()])
            super(SemiLocalLinearTrend, self).__init__(parameters=[
                Parameter('level_scale', level_scale_prior, scaled_softplus),
                Parameter('slope_mean', slope_mean_prior,
                          stddev_preconditioner),
                Parameter('slope_scale', slope_scale_prior, scaled_softplus),
                Parameter('autoregressive_coef', autoregressive_coef_prior,
                          autoregressive_coef_bijector),
            ],
                                                       latent_size=2,
                                                       name=name)
Example #7
0
  def __init__(self,
               components,
               observation_noise_scale_prior=None,
               observed_time_series=None,
               name=None):
    """Specify a structural time series model representing a sum of components.

    Args:
      components: Python `list` of one or more StructuralTimeSeries instances.
        These must have unique names.
      observation_noise_scale_prior: optional `tfd.Distribution` instance
        specifying a prior on `observation_noise_scale`. If `None`, a heuristic
        default prior is constructed based on the provided
        `observed_time_series`.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series. This is
        used only if `observation_noise_scale_prior` is not provided, to
        construct a default heuristic prior.
        Default value: `None`.
      name: Python `str` name of this model component; used as `name_scope`
        for ops created by this class.
        Default value: 'Sum'.

    Raises:
      ValueError: if components do not have unique names.
    """

    with tf.compat.v1.name_scope(
        name, 'Sum', values=[observed_time_series]) as name:
      if observation_noise_scale_prior is None:
        observed_stddev, _ = (
            sts_util.empirical_statistics(observed_time_series)
            if observed_time_series is not None else (1., 0.))
        observation_noise_scale_prior = tfd.LogNormal(
            loc=tf.math.log(.01 * observed_stddev), scale=2.)

      # Check that components have unique names, to ensure that inherited
      # parameters will be assigned unique names.
      component_names = [c.name for c in components]
      if len(component_names) != len(set(component_names)):
        raise ValueError(
            'Components must have unique names: {}'.format(component_names))
      components_by_name = collections.OrderedDict(
          [(c.name, c) for c in components])

      # Build parameters list for the combined model, by inheriting parameters
      # from the component models in canonical order.
      parameters = [
          Parameter('observation_noise_scale', observation_noise_scale_prior,
                    tfb.Softplus()),
      ] + [Parameter(name='{}_{}'.format(component.name, parameter.name),
                     prior=parameter.prior,
                     bijector=parameter.bijector)
           for component in components for parameter in component.parameters]

      self._components = components
      self._components_by_name = components_by_name

      super(Sum, self).__init__(
          parameters=parameters,
          latent_size=sum(
              [component.latent_size for component in components]),
          name=name)
Example #8
0
    def __init__(self,
                 level_scale_prior=None,
                 initial_level_prior=None,
                 observed_time_series=None,
                 name=None):
        """Specify a local level model.

    Args:
      level_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `level_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        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`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series). May
        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
        a mask `Tensor` to specify timesteps with missing observations.
        Default value: `None`.
      name: the name of this model component.
        Default value: 'LocalLevel'.
    """

        with tf.name_scope(name or 'LocalLevel') as name:

            dtype = dtype_util.common_dtype(
                [level_scale_prior, initial_level_prior])

            if observed_time_series is not None:
                _, observed_stddev, observed_initial = (
                    sts_util.empirical_statistics(observed_time_series))
            else:
                observed_stddev, observed_initial = (tf.convert_to_tensor(
                    value=1.,
                    dtype=dtype), tf.convert_to_tensor(value=0., dtype=dtype))

            # Heuristic default priors. Overriding these may dramatically
            # change inference performance and results.
            if level_scale_prior is None:
                level_scale_prior = tfd.LogNormal(loc=tf.math.log(
                    .05 * observed_stddev),
                                                  scale=3.,
                                                  name='level_scale_prior')
            if initial_level_prior is None:
                self._initial_state_prior = tfd.MultivariateNormalDiag(
                    loc=observed_initial[..., tf.newaxis],
                    scale_diag=(tf.abs(observed_initial) +
                                observed_stddev)[..., tf.newaxis],
                    name='initial_level_prior')
            else:
                self._initial_state_prior = tfd.MultivariateNormalDiag(
                    loc=initial_level_prior.mean()[..., tf.newaxis],
                    scale_diag=initial_level_prior.stddev()[..., tf.newaxis])

            super(LocalLevel, self).__init__(parameters=[
                Parameter(
                    'level_scale', level_scale_prior,
                    tfb.Chain([
                        tfb.AffineScalar(scale=observed_stddev),
                        tfb.Softplus()
                    ])),
            ],
                                             latent_size=1,
                                             name=name)
Example #9
0
  def __init__(self,
               components,
               observation_noise_scale_prior=None,
               observed_time_series=None,
               name=None):
    """Specify a structural time series model representing a sum of components.

    Args:
      components: Python `list` of one or more StructuralTimeSeries instances.
        These must have unique names.
      observation_noise_scale_prior: optional `tfd.Distribution` instance
        specifying a prior on `observation_noise_scale`. If `None`, a heuristic
        default prior is constructed based on the provided
        `observed_time_series`.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series. This is
        used only if `observation_noise_scale_prior` is not provided, to
        construct a default heuristic prior.
        Default value: `None`.
      name: Python `str` name of this model component; used as `name_scope`
        for ops created by this class.
        Default value: 'Sum'.

    Raises:
      ValueError: if components do not have unique names.
    """

    with tf.name_scope(
        name, 'Sum', values=[observed_time_series]) as name:
      if observation_noise_scale_prior is None:
        observed_stddev, _ = (
            sts_util.empirical_statistics(observed_time_series)
            if observed_time_series is not None else (1., 0.))
        observation_noise_scale_prior = tfd.LogNormal(
            loc=tf.log(.01 * observed_stddev), scale=2.)

      # Check that components have unique names, to ensure that inherited
      # parameters will be assigned unique names.
      component_names = [c.name for c in components]
      if len(component_names) != len(set(component_names)):
        raise ValueError(
            'Components must have unique names: {}'.format(component_names))
      components_by_name = collections.OrderedDict(
          [(c.name, c) for c in components])

      # Build parameters list for the combined model, by inheriting parameters
      # from the component models in canonical order.
      parameters = [
          Parameter('observation_noise_scale', observation_noise_scale_prior,
                    tfb.Softplus()),
      ] + [Parameter(name='{}_{}'.format(component.name, parameter.name),
                     prior=parameter.prior,
                     bijector=parameter.bijector)
           for component in components for parameter in component.parameters]

      self._components = components
      self._components_by_name = components_by_name

      super(Sum, self).__init__(
          parameters=parameters,
          latent_size=sum(
              [component.latent_size for component in components]),
          name=name)
Example #10
0
  def __init__(self,
               level_scale_prior=None,
               slope_scale_prior=None,
               initial_level_prior=None,
               initial_slope_prior=None,
               observed_time_series=None,
               name=None):
    """Specify a local linear trend model.

    Args:
      level_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `level_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      slope_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `slope_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        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`.
      initial_slope_prior: optional `tfd.Distribution` instance specifying a
        prior on the initial slope. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series).
        Default value: `None`.
      name: the name of this model component.
        Default value: 'LocalLinearTrend'.
    """

    with tf.name_scope(
        name, 'LocalLinearTrend', values=[observed_time_series]) as name:

      observed_stddev, observed_initial = (
          sts_util.empirical_statistics(observed_time_series)
          if observed_time_series is not None else (1., 0.))

      # Heuristic default priors. Overriding these may dramatically
      # change inference performance and results.
      if level_scale_prior is None:
        level_scale_prior = tfd.LogNormal(
            loc=tf.log(.05 * observed_stddev),
            scale=3.,
            name='level_scale_prior')
      if slope_scale_prior is None:
        slope_scale_prior = tfd.LogNormal(
            loc=tf.log(.05 * observed_stddev),
            scale=3.,
            name='slope_scale_prior')
      if initial_level_prior is None:
        initial_level_prior = tfd.Normal(
            loc=observed_initial,
            scale=tf.abs(observed_initial) + observed_stddev,
            name='initial_level_prior')
      if initial_slope_prior is None:
        initial_slope_prior = tfd.Normal(
            loc=0., scale=observed_stddev, name='initial_slope_prior')

      tf.assert_same_float_dtype([
          level_scale_prior, slope_scale_prior, initial_level_prior,
          initial_slope_prior
      ])

      self._initial_state_prior = tfd.MultivariateNormalDiag(
          loc=tf.stack(
              [initial_level_prior.mean(),
               initial_slope_prior.mean()
              ], axis=-1),
          scale_diag=tf.stack([
              initial_level_prior.stddev(),
              initial_slope_prior.stddev()
          ], axis=-1))

      super(LocalLinearTrend, self).__init__(
          parameters=[
              Parameter('level_scale', level_scale_prior, tfb.Softplus()),
              Parameter('slope_scale', slope_scale_prior, tfb.Softplus())
          ],
          latent_size=2,
          name=name)
    def __init__(self,
                 design_matrix,
                 drift_scale_prior=None,
                 initial_weights_prior=None,
                 observed_time_series=None,
                 name=None):
        """Specify a dynamic linear regression.

    Args:
      design_matrix: float `Tensor` of shape `concat([batch_shape,
        [num_timesteps, num_features]])`.
      drift_scale_prior: instance of `tfd.Distribution` specifying a prior on
        the `drift_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      initial_weights_prior: instance of `tfd.MultivariateNormal` representing
        the prior distribution on the latent states (the regression weights).
        Must have event shape `[num_features]`. If `None`, a weakly-informative
        Normal(0., 10.) prior is used.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series. Any `NaN`s
        are interpreted as missing observations; missingness may be also be
        explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series).
        Default value: `None`.
      name: Python `str` for the name of this component.
        Default value: 'DynamicLinearRegression'.

    """
        init_parameters = dict(locals())
        with tf.name_scope(name or 'DynamicLinearRegression') as name:
            dtype = dtype_util.common_dtype(
                [design_matrix, drift_scale_prior, initial_weights_prior])

            num_features = prefer_static.shape(design_matrix)[-1]

            # Default to a weakly-informative Normal(0., 10.) for the initital state
            if initial_weights_prior is None:
                initial_weights_prior = tfd.MultivariateNormalDiag(
                    scale_diag=10. * tf.ones([num_features], dtype=dtype))

            # Heuristic default priors. Overriding these may dramatically
            # change inference performance and results.
            if drift_scale_prior is None:
                if observed_time_series is None:
                    observed_stddev = tf.constant(1.0, dtype=dtype)
                else:
                    _, observed_stddev, _ = sts_util.empirical_statistics(
                        observed_time_series)

                drift_scale_prior = tfd.LogNormal(loc=tf.math.log(
                    .05 * observed_stddev),
                                                  scale=3.,
                                                  name='drift_scale_prior')

            self._initial_state_prior = initial_weights_prior
            self._design_matrix = design_matrix

            super(DynamicLinearRegression,
                  self).__init__(parameters=[
                      Parameter(
                          'drift_scale', drift_scale_prior,
                          tfb.Chain([
                              tfb.Scale(scale=observed_stddev),
                              tfb.Softplus()
                          ]))
                  ],
                                 latent_size=num_features,
                                 init_parameters=init_parameters,
                                 name=name)
  def __init__(self,
               level_scale_prior=None,
               slope_scale_prior=None,
               initial_level_prior=None,
               initial_slope_prior=None,
               observed_time_series=None,
               name=None):
    """Specify a local linear trend model.

    Args:
      level_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `level_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      slope_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `slope_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        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`.
      initial_slope_prior: optional `tfd.Distribution` instance specifying a
        prior on the initial slope. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series).
        Default value: `None`.
      name: the name of this model component.
        Default value: 'LocalLinearTrend'.
    """

    with tf.name_scope(
        name, 'LocalLinearTrend', values=[observed_time_series]) as name:

      observed_stddev, observed_initial = (
          sts_util.empirical_statistics(observed_time_series)
          if observed_time_series is not None else (1., 0.))

      # Heuristic default priors. Overriding these may dramatically
      # change inference performance and results.
      if level_scale_prior is None:
        level_scale_prior = tfd.LogNormal(
            loc=tf.log(.05 * observed_stddev),
            scale=3.,
            name='level_scale_prior')
      if slope_scale_prior is None:
        slope_scale_prior = tfd.LogNormal(
            loc=tf.log(.05 * observed_stddev),
            scale=3.,
            name='slope_scale_prior')
      if initial_level_prior is None:
        initial_level_prior = tfd.Normal(
            loc=observed_initial,
            scale=observed_stddev,
            name='initial_level_prior')
      if initial_slope_prior is None:
        initial_slope_prior = tfd.Normal(
            loc=0., scale=observed_stddev, name='initial_slope_prior')

      tf.assert_same_float_dtype([
          level_scale_prior, slope_scale_prior, initial_level_prior,
          initial_slope_prior
      ])

      self._initial_state_prior = tfd.MultivariateNormalDiag(
          loc=tf.stack(
              [initial_level_prior.mean(),
               initial_slope_prior.mean()
              ], axis=-1),
          scale_diag=tf.stack([
              initial_level_prior.stddev(),
              initial_slope_prior.stddev()
          ], axis=-1))

      super(LocalLinearTrend, self).__init__(
          parameters=[
              Parameter('level_scale', level_scale_prior, tfb.Softplus()),
              Parameter('slope_scale', slope_scale_prior, tfb.Softplus())
          ],
          latent_size=2,
          name=name)
Example #13
0
  def _build_test_model(self,
                        num_timesteps=5,
                        num_features=2,
                        batch_shape=(),
                        missing_prob=0,
                        true_noise_scale=0.1,
                        true_level_scale=0.04,
                        true_slope_scale=0.02,
                        prior_class=tfd.InverseGamma,
                        weights=None,
                        weights_prior_scale=10.,
                        sparse_weights_nonzero_prob=None,
                        time_series_shift=0.,
                        dtype=tf.float32,
                        design_matrix=False,
                        seed=None):
    if seed is None:
      seed = test_util.test_seed(sampler_type='stateless')
    (design_seed,
     weights_seed,
     noise_seed,
     level_seed,
     slope_seed,
     is_missing_seed) = samplers.split_seed(seed, 6, salt='_build_test_model')

    if weights is None:
      weights = samplers.normal(
          list(batch_shape) + [num_features], dtype=dtype, seed=weights_seed)
    if design_matrix is None:
      regression = tf.zeros(num_timesteps, dtype)
    else:
      if isinstance(design_matrix, bool) and not design_matrix:
        design_matrix = samplers.normal([num_timesteps, num_features],
                                        dtype=dtype,
                                        seed=design_seed)
      regression = tf.linalg.matvec(design_matrix, weights)
    noise = samplers.normal(
        list(batch_shape) + [num_timesteps],
        dtype=dtype, seed=noise_seed) * true_noise_scale

    level_residuals = samplers.normal(
        list(batch_shape) + [num_timesteps],
        dtype=dtype, seed=level_seed) * true_level_scale
    if true_slope_scale is not None:
      slope = tf.cumsum(samplers.normal(
          list(batch_shape) + [num_timesteps],
          dtype=dtype, seed=slope_seed) * true_slope_scale, axis=-1)
      level_residuals += slope
    level = tf.cumsum(level_residuals, axis=-1)
    time_series = (regression + noise + level + time_series_shift)
    is_missing = samplers.uniform(
        list(batch_shape) + [num_timesteps],
        dtype=dtype, seed=is_missing_seed) < missing_prob

    observation_noise_variance_prior = prior_class(
        concentration=tf.cast(0.01, dtype),
        scale=tf.cast(0.01 * 0.01, dtype))
    observation_noise_variance_prior.upper_bound = 100.0

    observed_time_series = tfp.sts.MaskedTimeSeries(
        time_series[..., tf.newaxis], is_missing)
    if time_series_shift != 0.:
      observed_mean, observed_stddev, observed_initial = (
          sts_util.empirical_statistics(observed_time_series))
      initial_level_prior = tfd.Normal(
          loc=observed_mean + observed_initial,
          scale=tf.abs(observed_initial) + observed_stddev)
    else:
      initial_level_prior = None

    model = gibbs_sampler.build_model_for_gibbs_fitting(
        observed_time_series=observed_time_series,
        design_matrix=design_matrix,
        weights_prior=(
            None if weights_prior_scale is None
            else tfd.Normal(loc=tf.cast(0., dtype),
                            scale=tf.cast(weights_prior_scale, dtype))),
        level_variance_prior=prior_class(
            concentration=tf.cast(0.01, dtype),
            scale=tf.cast(0.01 * 0.01, dtype)),
        slope_variance_prior=None if true_slope_scale is None else prior_class(
            concentration=tf.cast(0.01, dtype),
            scale=tf.cast(0.01 * 0.01, dtype)),
        initial_level_prior=initial_level_prior,
        observation_noise_variance_prior=observation_noise_variance_prior,
        sparse_weights_nonzero_prob=sparse_weights_nonzero_prob)
    return model, time_series, is_missing
Example #14
0
  def __init__(self,
               num_seasons,
               num_steps_per_season=1,
               drift_scale_prior=None,
               initial_effect_prior=None,
               observed_time_series=None,
               name=None):
    """Specify a seasonal effects model.

    Args:
      num_seasons: Scalar Python `int` number of seasons.
      num_steps_per_season: Python `int` number of steps in each
        season. This may be either a scalar (shape `[]`), in which case all
        seasons have the same length, or a NumPy array of shape `[num_seasons]`.
        Default value: 1.
      drift_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `drift_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      initial_effect_prior: optional `tfd.Distribution` instance specifying a
        normal prior on the initial effect of each season. This may be either
        a scalar `tfd.Normal` prior, in which case it applies independently to
        every season, or it may be multivariate normal (e.g.,
        `tfd.MultivariateNormalDiag`) with event shape `[num_seasons]`, in
        which case it specifies a joint prior across all seasons. If `None`, a
        heuristic default prior is constructed based on the provided
        `observed_time_series`.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series).
        Default value: `None`.
      name: the name of this model component.
        Default value: 'Seasonal'.
    """

    with tf.name_scope(name, 'Seasonal', values=[observed_time_series]) as name:

      observed_stddev, observed_initial = (
          sts_util.empirical_statistics(observed_time_series)
          if observed_time_series is not None else (1., 0.))

      # Heuristic default priors. Overriding these may dramatically
      # change inference performance and results.
      if drift_scale_prior is None:
        drift_scale_prior = tfd.LogNormal(loc=tf.log(.01 * observed_stddev),
                                          scale=3.)
      if initial_effect_prior is None:
        initial_effect_prior = tfd.Normal(loc=observed_initial,
                                          scale=observed_stddev)

      self._num_seasons = num_seasons
      self._num_steps_per_season = num_steps_per_season

      tf.assert_same_float_dtype([drift_scale_prior, initial_effect_prior])

      if isinstance(initial_effect_prior, tfd.Normal):
        self._initial_state_prior = tfd.MultivariateNormalDiag(
            loc=tf.stack([initial_effect_prior.mean()] * num_seasons, axis=-1),
            scale_diag=tf.stack([initial_effect_prior.stddev()] * num_seasons,
                                axis=-1))
      else:
        self._initial_state_prior = initial_effect_prior

      super(Seasonal, self).__init__(
          parameters=[
              Parameter('drift_scale', drift_scale_prior, tfb.Softplus()),
          ],
          latent_size=num_seasons,
          name=name)
Example #15
0
    def __init__(self,
                 num_seasons,
                 num_steps_per_season=1,
                 drift_scale_prior=None,
                 initial_effect_prior=None,
                 observed_time_series=None,
                 name=None):
        """Specify a seasonal effects model.

    Args:
      num_seasons: Scalar Python `int` number of seasons.
      num_steps_per_season: Python `int` number of steps in each
        season. This may be either a scalar (shape `[]`), in which case all
        seasons have the same length, or a NumPy array of shape `[num_seasons]`.
        Default value: 1.
      drift_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `drift_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      initial_effect_prior: optional `tfd.Distribution` instance specifying a
        normal prior on the initial effect of each season. This may be either
        a scalar `tfd.Normal` prior, in which case it applies independently to
        every season, or it may be multivariate normal (e.g.,
        `tfd.MultivariateNormalDiag`) with event shape `[num_seasons]`, in
        which case it specifies a joint prior across all seasons. If `None`, a
        heuristic default prior is constructed based on the provided
        `observed_time_series`.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series).
        Default value: `None`.
      name: the name of this model component.
        Default value: 'Seasonal'.
    """

        with tf.compat.v1.name_scope(name,
                                     'Seasonal',
                                     values=[observed_time_series]) as name:

            observed_stddev, observed_initial = (
                sts_util.empirical_statistics(observed_time_series)
                if observed_time_series is not None else (1., 0.))

            # Heuristic default priors. Overriding these may dramatically
            # change inference performance and results.
            if drift_scale_prior is None:
                drift_scale_prior = tfd.LogNormal(loc=tf.math.log(
                    .01 * observed_stddev),
                                                  scale=3.)
            if initial_effect_prior is None:
                initial_effect_prior = tfd.Normal(
                    loc=observed_initial,
                    scale=tf.abs(observed_initial) + observed_stddev)

            self._num_seasons = num_seasons
            self._num_steps_per_season = num_steps_per_season

            tf.debugging.assert_same_float_dtype(
                [drift_scale_prior, initial_effect_prior])

            if isinstance(initial_effect_prior, tfd.Normal):
                self._initial_state_prior = tfd.MultivariateNormalDiag(
                    loc=tf.stack([initial_effect_prior.mean()] * num_seasons,
                                 axis=-1),
                    scale_diag=tf.stack([initial_effect_prior.stddev()] *
                                        num_seasons,
                                        axis=-1))
            else:
                self._initial_state_prior = initial_effect_prior

            super(Seasonal, self).__init__(parameters=[
                Parameter('drift_scale', drift_scale_prior, tfb.Softplus()),
            ],
                                           latent_size=num_seasons,
                                           name=name)
Example #16
0
    def __init__(self,
                 period,
                 frequency_multipliers,
                 allow_drift=True,
                 drift_scale_prior=None,
                 initial_state_prior=None,
                 observed_time_series=None,
                 name=None):
        """Specify a smooth seasonal effects model.

    Args:
      period: positive scalar `float` `Tensor` giving the number of timesteps
        required for the longest cyclic effect to repeat.
      frequency_multipliers: One-dimensional `float` `Tensor` listing the
        frequencies (cyclic components) included in the model, as multipliers of
        the base/fundamental frequency `2. * pi / period`. Each component is
        specified by the number of times it repeats per period, and adds two
        latent dimensions to the model. A smooth seasonal model that can
        represent any periodic function is given by `frequency_multipliers = [1,
        2, ..., floor(period / 2)]`. However, it is often desirable to enforce a
        smoothness assumption (and reduce the computational burden) by dropping
        some of the higher frequencies.
      allow_drift: optional Python `bool` specifying whether the seasonal
        effects can drift over time.  Setting this to `False`
        removes the `drift_scale` parameter from the model. This is
        mathematically equivalent to
        `drift_scale_prior = tfd.Deterministic(0.)`, but removing drift
        directly is preferred because it avoids the use of a degenerate prior.
        Default value: `True`.
      drift_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `drift_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      initial_state_prior: instance of `tfd.MultivariateNormal` representing
        the prior distribution on the latent states. Must have event shape
        `[2 * len(frequency_multipliers)]`. If `None`, a heuristic default prior
        is constructed based on the provided `observed_time_series`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series). May
        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
        a mask `Tensor` to specify timesteps with missing observations.
        Default value: `None`.
      name: the name of this model component.
        Default value: 'SmoothSeasonal'.

    """

        with tf.name_scope(name or 'SmoothSeasonal') as name:

            _, observed_stddev, observed_initial = (
                sts_util.empirical_statistics(observed_time_series)
                if observed_time_series is not None else (0., 1., 0.))

            latent_size = 2 * static_num_frequencies(frequency_multipliers)

            # Heuristic default priors. Overriding these may dramatically
            # change inference performance and results.
            if drift_scale_prior is None:
                drift_scale_prior = tfd.LogNormal(loc=tf.math.log(
                    .01 * observed_stddev),
                                                  scale=3.)

            if initial_state_prior is None:
                initial_state_scale = (tf.abs(observed_initial) +
                                       observed_stddev)[..., tf.newaxis]
                ones = tf.ones([latent_size], dtype=drift_scale_prior.dtype)
                initial_state_prior = tfd.MultivariateNormalDiag(
                    scale_diag=initial_state_scale * ones)

            self._initial_state_prior = initial_state_prior
            self._period = period
            self._frequency_multipliers = frequency_multipliers

            parameters = []
            if allow_drift:
                parameters.append(
                    Parameter(
                        'drift_scale', drift_scale_prior,
                        tfb.Chain([
                            tfb.AffineScalar(scale=observed_stddev),
                            tfb.Softplus()
                        ])))
            self._allow_drift = allow_drift

            super(SmoothSeasonal, self).__init__(parameters=parameters,
                                                 latent_size=latent_size,
                                                 name=name)
Example #17
0
    def __init__(self,
                 num_seasons,
                 num_steps_per_season=1,
                 drift_scale_prior=None,
                 initial_effect_prior=None,
                 constrain_mean_effect_to_zero=True,
                 observed_time_series=None,
                 name=None):
        """Specify a seasonal effects model.

    Args:
      num_seasons: Scalar Python `int` number of seasons.
      num_steps_per_season: Python `int` number of steps in each
        season. This may be either a scalar (shape `[]`), in which case all
        seasons have the same length, or a NumPy array of shape `[num_seasons]`,
        in which seasons have different length, but remain constant around
        different cycles, or a NumPy array of shape `[num_cycles, num_seasons]`,
        in which num_steps_per_season for each season also varies in different
        cycle (e.g., a 4 years cycle with leap day).
        Default value: 1.
      drift_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `drift_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      initial_effect_prior: optional `tfd.Distribution` instance specifying a
        normal prior on the initial effect of each season. This may be either
        a scalar `tfd.Normal` prior, in which case it applies independently to
        every season, or it may be multivariate normal (e.g.,
        `tfd.MultivariateNormalDiag`) with event shape `[num_seasons]`, in
        which case it specifies a joint prior across all seasons. If `None`, a
        heuristic default prior is constructed based on the provided
        `observed_time_series`.
        Default value: `None`.
      constrain_mean_effect_to_zero: if `True`, use a model parameterization
        that constrains the mean effect across all seasons to be zero. This
        constraint is generally helpful in identifying the contributions of
        different model components and can lead to more interpretable
        posterior decompositions. It may be undesirable if you plan to directly
        examine the latent space of the underlying state space model.
        Default value: `True`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series). May
        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
        a mask `Tensor` to specify timesteps with missing observations.
        Default value: `None`.
      name: the name of this model component.
        Default value: 'Seasonal'.
    """

        with tf.compat.v1.name_scope(name,
                                     'Seasonal',
                                     values=[observed_time_series]) as name:

            _, observed_stddev, observed_initial = (
                sts_util.empirical_statistics(observed_time_series)
                if observed_time_series is not None else (0., 1., 0.))

            # Heuristic default priors. Overriding these may dramatically
            # change inference performance and results.
            if drift_scale_prior is None:
                drift_scale_prior = tfd.LogNormal(loc=tf.math.log(
                    .01 * observed_stddev),
                                                  scale=3.)
            if initial_effect_prior is None:
                initial_effect_prior = tfd.Normal(
                    loc=observed_initial,
                    scale=tf.abs(observed_initial) + observed_stddev)

            dtype = tf.debugging.assert_same_float_dtype(
                [drift_scale_prior, initial_effect_prior])

            if isinstance(initial_effect_prior, tfd.Normal):
                initial_state_prior = tfd.MultivariateNormalDiag(
                    loc=tf.stack([initial_effect_prior.mean()] * num_seasons,
                                 axis=-1),
                    scale_diag=tf.stack([initial_effect_prior.stddev()] *
                                        num_seasons,
                                        axis=-1))
            else:
                initial_state_prior = initial_effect_prior

            if constrain_mean_effect_to_zero:
                # Transform the prior to the residual parameterization used by
                # `ConstrainedSeasonalStateSpaceModel`, imposing a zero-sum constraint.
                # This doesn't change the marginal prior on individual effects, but
                # does introduce dependence between the effects.
                (effects_to_residuals,
                 _) = build_effects_to_residuals_matrix(num_seasons,
                                                        dtype=dtype)
                effects_to_residuals_linop = tf.linalg.LinearOperatorFullMatrix(
                    effects_to_residuals
                )  # Use linop so that matmul broadcasts.
                initial_state_prior_loc = effects_to_residuals_linop.matvec(
                    initial_state_prior.mean())
                initial_state_prior_scale_linop = effects_to_residuals_linop.matmul(
                    initial_state_prior.scale)  # returns LinearOperator
                initial_state_prior = tfd.MultivariateNormalFullCovariance(
                    loc=initial_state_prior_loc,
                    covariance_matrix=initial_state_prior_scale_linop.matmul(
                        initial_state_prior_scale_linop.to_dense(),
                        adjoint_arg=True))

            self._constrain_mean_effect_to_zero = constrain_mean_effect_to_zero
            self._initial_state_prior = initial_state_prior
            self._num_seasons = num_seasons
            self._num_steps_per_season = num_steps_per_season

            super(Seasonal, self).__init__(
                parameters=[
                    Parameter('drift_scale', drift_scale_prior,
                              tfb.Softplus()),
                ],
                latent_size=(num_seasons -
                             1 if self.constrain_mean_effect_to_zero else
                             num_seasons),
                name=name)
  def __init__(self,
               ar_order,
               ma_order,
               integration_degree=0,
               ar_coefficients_prior=None,
               ma_coefficients_prior=None,
               level_drift_prior=None,
               level_scale_prior=None,
               initial_state_prior=None,
               ar_coefficient_constraining_bijector=None,
               ma_coefficient_constraining_bijector=None,
               observed_time_series=None,
               name=None):
    """Specifies an ARIMA(p=ar_order, d=integration_degree, q=ma_order) model.

    Args:
      ar_order: scalar Python positive `int` specifying the order of the
        autoregressive process (`p` in `ARIMA(p, d, q)`).
      ma_order: scalar Python positive `int` specifying the order of the
        moving-average process (`q` in `ARIMA(p, d, q)`).
      integration_degree: scalar Python positive `int` specifying the number
        of times to integrate an ARMA process. (`d` in `ARIMA(p, d, q)`).
        Default value: `0`.
      ar_coefficients_prior: optional `tfd.Distribution` instance specifying a
        prior on the `ar_coefficients` parameter. If `None`, a default standard
        normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([ar_order]))`)
        prior is used.
        Default value: `None`.
      ma_coefficients_prior: optional `tfd.Distribution` instance specifying a
        prior on the `ma_coefficients` parameter. If `None`, a default standard
        normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([ma_order]))`)
        prior is used.
        Default value: `None`.
      level_drift_prior: optional `tfd.Distribution` instance specifying a prior
        on the `level_drift` parameter. If `None`, the parameter is not inferred
        and is instead fixed to zero.
        Default value: `None`.
      level_scale_prior: optional `tfd.Distribution` instance specifying a prior
        on the `level_scale` parameter. If `None`, a heuristic default prior is
        constructed based on the provided `observed_time_series`.
        Default value: `None`.
      initial_state_prior: optional `tfd.Distribution` instance specifying a
        prior on the initial state, corresponding to the values of the process
        at a set of size `order` of imagined timesteps before the initial step.
        If `None`, a heuristic default prior is constructed based on the
        provided `observed_time_series`.
        Default value: `None`.
      ar_coefficient_constraining_bijector: optional `tfb.Bijector` instance
        representing a constraining mapping for the autoregressive coefficients.
        For example, `tfb.Tanh()` constrains the coefficients to lie in
        `(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
        `tfb.Identity()` implies no constraint. If `None`, the default behavior
        constrains the coefficients to lie in `(-1, 1)` using a `Tanh` bijector.
        Default value: `None`.
      ma_coefficient_constraining_bijector: optional `tfb.Bijector` instance
        representing a constraining mapping for the moving average coefficients.
        For example, `tfb.Tanh()` constrains the coefficients to lie in
        `(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
        `tfb.Identity()` implies no constraint. If `None`, the default behavior
        is to apply no constraint.
        Default value: `None`.
      observed_time_series: optional `float` `Tensor` of shape
        `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
        supported when `T > 1`), specifying an observed time series. Any `NaN`s
        are interpreted as missing observations; missingness may be also be
        explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
        Any priors not explicitly set will be given default values according to
        the scale of the observed time series (or batch of time series).
        Default value: `None`.
      name: the name of this model component.
        Default value: 'ARIMA'.
    """
    init_parameters = dict(locals())
    with tf.name_scope(name or 'ARIMA') as name:
      masked_time_series = None
      if observed_time_series is not None:
        masked_time_series = (
            sts_util.canonicalize_observed_time_series_with_mask(
                observed_time_series))
      dtype = dtype_util.common_dtype(
          [(masked_time_series.time_series
            if masked_time_series is not None else None),
           ar_coefficients_prior,
           ma_coefficients_prior,
           level_scale_prior,
           initial_state_prior], dtype_hint=tf.float32)

      if observed_time_series is not None:
        for _ in range(integration_degree):
          # Compute statistics using `integration_order`-order differences.
          masked_time_series = (
              missing_values_util.differentiate_masked_time_series(
                  masked_time_series))
        _, observed_stddev, observed_initial = sts_util.empirical_statistics(
            masked_time_series)
      else:
        observed_stddev, observed_initial = (
            tf.convert_to_tensor(value=1., dtype=dtype),
            tf.convert_to_tensor(value=0., dtype=dtype))
      batch_ones = ps.ones(ps.concat([
          ps.shape(observed_initial),  # Batch shape
          [1]], axis=0), dtype=dtype)

      # Heuristic default priors. Overriding these may dramatically
      # change inference performance and results.
      if ar_coefficients_prior is None:
        ar_coefficients_prior = tfd.MultivariateNormalDiag(
            scale_diag=batch_ones * ps.ones([ar_order]))
      if ma_coefficients_prior is None:
        ma_coefficients_prior = tfd.MultivariateNormalDiag(
            scale_diag=batch_ones * ps.ones([ma_order]))
      if level_scale_prior is None:
        level_scale_prior = tfd.LogNormal(
            loc=tf.math.log(0.05 *  observed_stddev), scale=3.)

      if (ar_coefficients_prior.event_shape.is_fully_defined() and
          ar_order != ar_coefficients_prior.event_shape[0]):
        raise ValueError(
            "Autoregressive prior dimension {} doesn't match order {}.".format(
                ar_coefficients_prior.event_shape[0], ar_order))
      if (ma_coefficients_prior.event_shape.is_fully_defined() and
          ma_order != ma_coefficients_prior.event_shape[0]):
        raise ValueError(
            "Moving average prior dimension {} doesn't match order {}.".format(
                ma_coefficients_prior.event_shape[0], ma_order))

      latent_size = ps.maximum(ar_order, ma_order + 1) + integration_degree
      if initial_state_prior is None:
        initial_state_prior = tfd.MultivariateNormalDiag(
            loc=sts_util.pad_tensor_with_trailing_zeros(
                observed_initial[..., tf.newaxis] * batch_ones,
                num_zeros=latent_size - 1),
            scale_diag=sts_util.pad_tensor_with_trailing_zeros(
                (tf.abs(observed_initial) +
                 observed_stddev)[..., tf.newaxis] * batch_ones,
                num_zeros=latent_size - 1))

      self._ar_order = ar_order
      self._ma_order = ma_order
      self._integration_degree = integration_degree
      self._ar_coefficients_prior = ar_coefficients_prior
      self._ma_coefficients_prior = ma_coefficients_prior
      self._level_scale_prior = level_scale_prior
      self._initial_state_prior = initial_state_prior

      parameters = []
      if ar_order > 0:
        parameters.append(
            Parameter('ar_coefficients',
                      ar_coefficients_prior,
                      (ar_coefficient_constraining_bijector
                       if ar_coefficient_constraining_bijector
                       else tfb.Tanh())))
      if ma_order > 0:
        parameters.append(
            Parameter('ma_coefficients',
                      ma_coefficients_prior,
                      (ma_coefficient_constraining_bijector
                       if ma_coefficient_constraining_bijector
                       else tfb.Identity())))
      if level_drift_prior is not None:
        parameters.append(
            Parameter(
                'level_drift',
                level_drift_prior,
                tfb.Chain([
                    tfb.Scale(scale=observed_stddev),
                    (level_drift_prior.
                     experimental_default_event_space_bijector())])))
      super(AutoregressiveIntegratedMovingAverage, self).__init__(
          parameters=parameters + [
              Parameter('level_scale', level_scale_prior,
                        tfb.Chain([tfb.Scale(scale=observed_stddev),
                                   tfb.Softplus(low=dtype_util.eps(dtype))]))
          ],
          latent_size=latent_size,
          init_parameters=init_parameters,
          name=name)