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
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