def do_sampling(observed_time_series, is_missing):
return gibbs_sampler.fit_with_gibbs_sampling(
model,
tfp.sts.MaskedTimeSeries(observed_time_series, is_missing),
num_results=4,
num_warmup_steps=1,
seed=seed)
def do_sampling():
return gibbs_sampler.fit_with_gibbs_sampling(
model,
observed_time_series,
num_results=100,
num_warmup_steps=100,
seed=test_util.test_seed(sampler_type='stateless'))
def do_sampling_again(observed_time_series, is_missing):
return gibbs_sampler.fit_with_gibbs_sampling(
dummy_model,
tfp.sts.MaskedTimeSeries(observed_time_series, is_missing),
num_results=30,
num_warmup_steps=10,
seed=sample_seed)
def test_forecasts_are_sane(self):
seed = test_util.test_seed()
num_observed_steps = 5
num_forecast_steps = 3
model, observed_time_series, is_missing = self._build_test_model(
num_timesteps=num_observed_steps + num_forecast_steps,
batch_shape=[3])
samples = gibbs_sampler.fit_with_gibbs_sampling(
model,
tfp.sts.MaskedTimeSeries(
observed_time_series[..., :num_observed_steps, tf.newaxis],
is_missing[..., :num_observed_steps]),
num_results=5,
num_warmup_steps=10,
seed=seed,
compile_steps_with_xla=False)
predictive_dist = gibbs_sampler.one_step_predictive(
model,
samples,
num_forecast_steps=num_forecast_steps,
thin_every=1)
predictive_mean, predictive_stddev = self.evaluate(
(predictive_dist.mean(), predictive_dist.stddev()))
self.assertAllEqual(predictive_mean.shape,
[3, num_observed_steps + num_forecast_steps])
self.assertAllEqual(predictive_stddev.shape,
[3, num_observed_steps + num_forecast_steps])
# Uncertainty should increase over the forecast period.
self.assertTrue(
np.all(predictive_stddev[..., num_observed_steps + 1:] >
predictive_stddev[..., num_observed_steps:-1]))
def do_sampling():
return gibbs_sampler.fit_with_gibbs_sampling(
model,
observed_time_series,
num_results=100,
num_warmup_steps=100,
seed=test_util.test_seed(sampler_type='stateless'),
experimental_use_dynamic_cholesky=use_dyanamic_cholesky)
def do_sampling():
return gibbs_sampler.fit_with_gibbs_sampling(
model,
tfp.sts.MaskedTimeSeries(
observed_time_series[..., :num_observed_steps, tf.newaxis],
is_missing[..., :num_observed_steps]),
num_chains=num_chains,
num_results=num_results,
num_warmup_steps=100,
seed=seed)
def _run():
model, observed_time_series, is_missing = self._build_test_model(
num_timesteps=336, batch_shape=[])
return gibbs_sampler.fit_with_gibbs_sampling(
model,
tfp.sts.MaskedTimeSeries(observed_time_series[..., tf.newaxis],
is_missing),
num_results=500,
num_warmup_steps=100,
seed=seed)
def test_invalid_model_raises_error(self):
observed_time_series = tf.convert_to_tensor([1., 0., -1., 2.])
bad_model = tfp.sts.Sum(
[tfp.sts.LinearRegression(design_matrix=tf.ones([4, 2])),
tfp.sts.LocalLevel(observed_time_series=observed_time_series),],
observed_time_series=observed_time_series)
with self.assertRaisesRegexp(ValueError, 'does not support Gibbs sampling'):
gibbs_sampler.fit_with_gibbs_sampling(
bad_model, observed_time_series, seed=test_util.test_seed())
bad_model.supports_gibbs_sampling = True
with self.assertRaisesRegexp(
ValueError, 'Expected the first model component to be an instance of'):
gibbs_sampler.fit_with_gibbs_sampling(
bad_model, observed_time_series, seed=test_util.test_seed())
bad_model_with_correct_params = tfp.sts.Sum([
# A seasonal model with no drift has no parameters, so adding it
# won't break the check for correct params.
tfp.sts.Seasonal(num_seasons=2,
allow_drift=False,
observed_time_series=observed_time_series),
tfp.sts.LocalLevel(observed_time_series=observed_time_series),
tfp.sts.LinearRegression(design_matrix=tf.ones([5, 2]))])
bad_model_with_correct_params.supports_gibbs_sampling = True
with self.assertRaisesRegexp(ValueError,
'Expected the first model component to be an '
'instance of `tfp.sts.LocalLevel`'):
gibbs_sampler.fit_with_gibbs_sampling(bad_model_with_correct_params,
observed_time_series,
seed=test_util.test_seed())
def test_end_to_end_prediction_works_and_is_deterministic(
self, dtype, use_xla):
if not tf.executing_eagerly():
return
seed = test_util.test_seed()
model, observed_time_series, is_missing = self._build_test_model(
num_timesteps=5, batch_shape=[3])
samples = gibbs_sampler.fit_with_gibbs_sampling(
model,
tfp.sts.MaskedTimeSeries(observed_time_series[..., tf.newaxis],
is_missing),
num_results=4,
num_warmup_steps=1,
seed=seed,
compile_steps_with_xla=use_xla)
predictive_dist = gibbs_sampler.one_step_predictive(model,
samples,
thin_every=1)
# Test that the seeded calculation gives the same result on multiple runs.
samples2 = gibbs_sampler.fit_with_gibbs_sampling(
model,
tfp.sts.MaskedTimeSeries(observed_time_series, is_missing),
num_results=4,
num_warmup_steps=1,
seed=seed,
compile_steps_with_xla=use_xla)
predictive_dist2 = gibbs_sampler.one_step_predictive(model,
samples2,
thin_every=1)
(predictive_mean_, predictive_stddev_, predictive_mean2_,
predictive_stddev2_) = self.evaluate(
(predictive_dist.mean(), predictive_dist.stddev(),
predictive_dist2.mean(), predictive_dist2.stddev()))
self.assertAllEqual(predictive_mean_, predictive_mean2_)
self.assertAllEqual(predictive_stddev_, predictive_stddev2_)
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_benchmark_sampling_with_xla(self):
if not tf.executing_eagerly():
return
seed = tfp_test_util.test_seed()
model, observed_time_series, is_missing = self._build_test_model(
num_timesteps=336, batch_shape=[])
t0 = time.time()
samples = gibbs_sampler.fit_with_gibbs_sampling(
model,
tfp.sts.MaskedTimeSeries(observed_time_series[..., tf.newaxis],
is_missing),
num_results=500,
num_warmup_steps=100,
seed=seed,
compile_steps_with_xla=True)
t1 = time.time()
print('Drew (100+500) samples in time', t1 - t0)
print('Results:', samples)
def test_forecasts_match_reference(self, use_slope):
seed = test_util.test_seed()
num_observed_steps = 5
num_forecast_steps = 4
model, observed_time_series, is_missing = self._build_test_model(
num_timesteps=num_observed_steps + num_forecast_steps,
true_slope_scale=0.5 if use_slope else None,
batch_shape=[3])
samples = tf.function(lambda: gibbs_sampler.fit_with_gibbs_sampling( # pylint: disable=g-long-lambda
model,
tfp.sts.MaskedTimeSeries(
observed_time_series[..., :num_observed_steps, tf.newaxis],
is_missing[..., :num_observed_steps]),
num_results=10000,
num_warmup_steps=100,
seed=seed))()
predictive_dist = gibbs_sampler.one_step_predictive(
model,
samples,
num_forecast_steps=num_forecast_steps,
thin_every=1)
predictive_mean, predictive_stddev = self.evaluate(
(predictive_dist.mean(), predictive_dist.stddev()))
self.assertAllEqual(predictive_mean.shape,
[3, num_observed_steps + num_forecast_steps])
self.assertAllEqual(predictive_stddev.shape,
[3, num_observed_steps + num_forecast_steps])
if use_slope:
parameter_samples = (samples.observation_noise_scale,
samples.level_scale, samples.slope_scale,
samples.weights)
else:
parameter_samples = (samples.observation_noise_scale,
samples.level_scale, samples.weights)
# Note that although we expect the Gibbs-sampled forecasts to match a
# reference implementation, we *don't* expect the one-step predictions to
# match `tfp.sts.one_step_predictive`, because that makes predictions using
# a filtered posterior (i.e., given only previous observations) whereas the
# Gibbs-sampled latent `level`s will incorporate some information from
# future observations.
reference_forecast_dist = tfp.sts.forecast(
model,
observed_time_series=observed_time_series[
..., :num_observed_steps],
parameter_samples=parameter_samples,
num_steps_forecast=num_forecast_steps)
reference_forecast_mean = self.evaluate(
reference_forecast_dist.mean()[..., 0])
reference_forecast_stddev = self.evaluate(
reference_forecast_dist.stddev()[..., 0])
self.assertAllClose(predictive_mean[..., -num_forecast_steps:],
reference_forecast_mean,
atol=1.0 if use_slope else 0.3)
self.assertAllClose(predictive_stddev[..., -num_forecast_steps:],
reference_forecast_stddev,
atol=2.0 if use_slope else 1.0)
