def _fit_seasonal_model_with_gibbs_sampling(observed_time_series, seasonal_structure, num_warmup_steps=50, num_results=100, seed=None): """Builds a seasonality-as-regression model and fits it by Gibbs sampling.""" with tf.name_scope('fit_seasonal_model_with_gibbs_sampling'): observed_time_series = sts_util.canonicalize_observed_time_series_with_mask( observed_time_series) dtype = observed_time_series.time_series.dtype design_matrix = seasonality_util.build_fixed_effects( num_steps=ps.shape(observed_time_series.time_series)[-2], seasonal_structure=seasonal_structure, dtype=dtype) # Default priors. # pylint: disable=protected-access one = tf.ones([], dtype=dtype) level_variance_prior = tfd.InverseGamma(concentration=16, scale=16. * 0.001**2 * one) level_variance_prior._upper_bound = one slope_variance_prior = tfd.InverseGamma(concentration=16, scale=16. * 0.05**2 * one) slope_variance_prior._upper_bound = 0.01 * one observation_noise_variance_prior = tfd.InverseGamma(concentration=0.05, scale=0.05 * one) observation_noise_variance_prior._upper_bound = 1.2 * one # pylint: enable=protected-access model = gibbs_sampler.build_model_for_gibbs_fitting( observed_time_series=observed_time_series, design_matrix=design_matrix, weights_prior=tfd.Normal(loc=0., scale=one), level_variance_prior=level_variance_prior, slope_variance_prior=slope_variance_prior, observation_noise_variance_prior=observation_noise_variance_prior) return [ model, gibbs_sampler.fit_with_gibbs_sampling( model, observed_time_series, num_results=num_results, num_warmup_steps=num_warmup_steps, seed=seed) ]
def test_log_prob_matches_linear_gaussian_ssm(self): dim = 2 batch_shape = [3, 1] seed, *model_seeds = samplers.split_seed(test_util.test_seed(), n=6) # Sample a random linear Gaussian process. prior_loc = self.evaluate( tfd.Normal(0., 1.).sample(batch_shape + [dim], seed=model_seeds[0])) prior_scale = self.evaluate( tfd.InverseGamma(1., 1.).sample(batch_shape + [dim], seed=model_seeds[1])) transition_matrix = self.evaluate( tfd.Normal(0., 1.).sample([dim, dim], seed=model_seeds[2])) transition_bias = self.evaluate( tfd.Normal(0., 1.).sample(batch_shape + [dim], seed=model_seeds[3])) transition_scale_tril = self.evaluate( tf.linalg.cholesky( tfd.WishartTriL( df=dim, scale_tril=tf.eye(dim)).sample(seed=model_seeds[4]))) initial_state_prior = tfd.MultivariateNormalDiag( loc=prior_loc, scale_diag=prior_scale, name='initial_state_prior') lgssm = tfd.LinearGaussianStateSpaceModel( num_timesteps=7, transition_matrix=transition_matrix, transition_noise=tfd.MultivariateNormalTriL( loc=transition_bias, scale_tril=transition_scale_tril), # Trivial observation model to pass through the latent state. observation_matrix=tf.eye(dim), observation_noise=tfd.MultivariateNormalDiag( loc=tf.zeros(dim), scale_diag=tf.zeros(dim)), initial_state_prior=initial_state_prior) markov_chain = tfd.MarkovChain( initial_state_prior=initial_state_prior, transition_fn=lambda _, x: tfd.MultivariateNormalTriL( # pylint: disable=g-long-lambda loc=tf.linalg.matvec(transition_matrix, x) + transition_bias, scale_tril=transition_scale_tril), num_steps=7) x = markov_chain.sample(5, seed=seed) self.assertAllClose(lgssm.log_prob(x), markov_chain.log_prob(x), rtol=1e-5)
def _resample_scale(prior_concentration, prior_scale, observed_residuals, is_missing=None, seed=None): """Samples a scale parameter from its conditional posterior. We assume the conjugate InverseGamma->Normal model: ``` scale ~ Sqrt(InverseGamma(prior_concentration, prior_scale)) for i in [1, ..., num_observations]: x[i] ~ Normal(loc, scale) ``` in which `loc` is known, and return a sample from `p(scale | x)`. Args: prior_concentration: Float `Tensor` concentration parameter of the InverseGamma prior distribution. prior_scale: Float `Tensor` scale parameter of the InverseGamma prior distribution. observed_residuals: Float `Tensor` of shape `[..., num_observations]`, specifying the centered observations `(x - loc)`. is_missing: Optional `bool` `Tensor` of shape `[..., num_observations]`. A `True` value indicates that the corresponding observation is missing. seed: Optional `Python` `int` seed controlling the sampled value. Returns: sampled_scale: A `Tensor` sample from the posterior `p(scale | x)`. """ if is_missing is not None: num_missing = tf.reduce_sum(tf.cast(is_missing, observed_residuals.dtype), axis=-1) num_observations = prefer_static.shape(observed_residuals)[-1] if is_missing is not None: observed_residuals = tf.where(is_missing, tf.zeros_like(observed_residuals), observed_residuals) num_observations -= num_missing variance_posterior = tfd.InverseGamma( concentration=prior_concentration + num_observations / 2., scale=prior_scale + tf.reduce_sum( tf.square(observed_residuals), axis=-1) / 2.) return tf.sqrt(variance_posterior.sample(seed=seed))
def __init__(self, design_matrix, weights_prior_scale=0.1, weights_batch_shape=None, name=None): """Specify a sparse linear regression model. Args: 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_scale: float `Tensor` defining the scale of the Horseshoe prior on regression weights. Small values encourage the weights to be sparse. The shape must broadcast with `weights_batch_shape`. Default value: `0.1`. weights_batch_shape: if `None`, defaults to `design_matrix.batch_shape_tensor()`. Must broadcast with the batch shape of `design_matrix`. Default value: `None`. name: the name of this model component. Default value: 'SparseLinearRegression'. """ with tf.compat.v1.name_scope( name, 'SparseLinearRegression', values=[design_matrix, weights_prior_scale]) as name: if not isinstance(design_matrix, tfl.LinearOperator): design_matrix = tfl.LinearOperatorFullMatrix( tf.convert_to_tensor(value=design_matrix, name='design_matrix'), name='design_matrix_linop') if tf.compat.dimension_value(design_matrix.shape[-1]) is not None: num_features = design_matrix.shape[-1] else: num_features = design_matrix.shape_tensor()[-1] if weights_batch_shape is None: weights_batch_shape = design_matrix.batch_shape_tensor() else: weights_batch_shape = tf.convert_to_tensor(value=weights_batch_shape, dtype=tf.int32) weights_shape = tf.concat([weights_batch_shape, [num_features]], axis=0) dtype = design_matrix.dtype self._design_matrix = design_matrix self._weights_prior_scale = weights_prior_scale ones_like_weights_batch = tf.ones(weights_batch_shape, dtype=dtype) ones_like_weights = tf.ones(weights_shape, dtype=dtype) super(SparseLinearRegression, self).__init__( parameters=[ Parameter('global_scale_variance', prior=tfd.InverseGamma( 0.5 * ones_like_weights_batch, 0.5 * ones_like_weights_batch), bijector=tfb.Softplus()), Parameter('global_scale_noncentered', prior=tfd.HalfNormal( scale=ones_like_weights_batch), bijector=tfb.Softplus()), Parameter('local_scale_variances', prior=tfd.Independent(tfd.InverseGamma( 0.5 * ones_like_weights, 0.5 * ones_like_weights), reinterpreted_batch_ndims=1), bijector=tfb.Softplus()), Parameter('local_scales_noncentered', prior=tfd.Independent(tfd.HalfNormal( scale=ones_like_weights), reinterpreted_batch_ndims=1), bijector=tfb.Softplus()), Parameter('weights_noncentered', prior=tfd.Independent(tfd.Normal( loc=tf.zeros_like(ones_like_weights), scale=ones_like_weights), reinterpreted_batch_ndims=1), bijector=tfb.Identity()) ], latent_size=0, name=name)