def test_batch_shape(self): batch_shape = [3, 2] partial_batch_shape = [2] seed = test_util.test_seed(sampler_type='stateless') num_seasons = 24 initial_state_prior = tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder( np.exp(np.random.randn(*(partial_batch_shape + [num_seasons]))))) drift_scale = self._build_placeholder( np.exp(np.random.randn(*batch_shape))) observation_noise_scale = self._build_placeholder( np.exp(np.random.randn(*partial_batch_shape))) ssm = SeasonalStateSpaceModel( num_timesteps=9, num_seasons=24, num_steps_per_season=2, drift_scale=drift_scale, observation_noise_scale=observation_noise_scale, initial_state_prior=initial_state_prior) # First check that the model's batch shape is the broadcast batch shape # of parameters, as expected. self.assertAllEqual(self.evaluate(ssm.batch_shape_tensor()), batch_shape) y_ = self.evaluate(ssm.sample(seed=seed)) self.assertAllEqual(y_.shape[:-2], batch_shape) # Next check that the broadcasting works as expected, and the batch log_prob # actually matches the log probs of independent models. individual_ssms = [ SeasonalStateSpaceModel( num_timesteps=9, num_seasons=num_seasons, num_steps_per_season=2, drift_scale=drift_scale[i, j, ...], observation_noise_scale=observation_noise_scale[j, ...], initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=initial_state_prior.scale.diag[j, ...])) for i in range(batch_shape[0]) for j in range(batch_shape[1]) ] batch_lps_ = self.evaluate(ssm.log_prob(y_)).flatten() individual_ys = [ y_[i, j, ...] for i in range(batch_shape[0]) for j in range(batch_shape[1]) ] individual_lps_ = self.evaluate([ individual_ssm.log_prob(individual_y) for (individual_ssm, individual_y) in zip(individual_ssms, individual_ys) ]) self.assertAllClose(individual_lps_, batch_lps_)
def test_day_of_week_example(self): # Test that the Seasonal SSM is equivalent to individually modeling # a random walk on each season's slice of timesteps. seed = test_util.test_seed(sampler_type='stateless') drift_scale = 0.6 observation_noise_scale = 0.1 day_of_week = SeasonalStateSpaceModel( num_timesteps=28, num_seasons=7, drift_scale=self._build_placeholder(drift_scale), observation_noise_scale=observation_noise_scale, initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder(np.ones([7]))), num_steps_per_season=1) random_walk_model = tfd.LinearGaussianStateSpaceModel( num_timesteps=4, transition_matrix=self._build_placeholder([[1.]]), transition_noise=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder([drift_scale])), observation_matrix=self._build_placeholder([[1.]]), observation_noise=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder([observation_noise_scale])), initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder([1.]))) sampled_time_series = day_of_week.sample(seed=seed) (sampled_time_series_, total_lp_, prior_mean_, prior_variance_) = self.evaluate([ sampled_time_series, day_of_week.log_prob(sampled_time_series), day_of_week.mean(), day_of_week.variance() ]) expected_daily_means_, expected_daily_variances_ = self.evaluate( [random_walk_model.mean(), random_walk_model.variance()]) # For the (noncontiguous) indices corresponding to each season, assert # that the model's mean, variance, and log_prob match a random-walk model. daily_lps = [] for day_idx in range(7): self.assertAllClose(prior_mean_[day_idx::7], expected_daily_means_) self.assertAllClose(prior_variance_[day_idx::7], expected_daily_variances_) daily_lps.append( self.evaluate( random_walk_model.log_prob( sampled_time_series_[day_idx::7]))) self.assertAlmostEqual(total_lp_, sum(daily_lps), places=3)
def test_day_of_week_example(self): # Test that the Seasonal SSM is equivalent to individually modeling # a random walk on each season's slice of timesteps. drift_scale = 0.6 observation_noise_scale = 0.1 day_of_week = SeasonalStateSpaceModel( num_timesteps=28, num_seasons=7, drift_scale=self._build_placeholder(drift_scale), observation_noise_scale=observation_noise_scale, initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder(np.ones([7]))), num_steps_per_season=1) random_walk_model = tfd.LinearGaussianStateSpaceModel( num_timesteps=4, transition_matrix=self._build_placeholder([[1.]]), transition_noise=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder([drift_scale])), observation_matrix=self._build_placeholder([[1.]]), observation_noise=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder([observation_noise_scale])), initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder([1.]))) sampled_time_series = day_of_week.sample() (sampled_time_series_, total_lp_, prior_mean_, prior_variance_) = self.evaluate([ sampled_time_series, day_of_week.log_prob(sampled_time_series), day_of_week.mean(), day_of_week.variance()]) expected_daily_means_, expected_daily_variances_ = self.evaluate([ random_walk_model.mean(), random_walk_model.variance()]) # For the (noncontiguous) indices corresponding to each season, assert # that the model's mean, variance, and log_prob match a random-walk model. daily_lps = [] for day_idx in range(7): self.assertAllClose(prior_mean_[day_idx::7], expected_daily_means_) self.assertAllClose(prior_variance_[day_idx::7], expected_daily_variances_) daily_lps.append(self.evaluate(random_walk_model.log_prob( sampled_time_series_[day_idx::7]))) self.assertAlmostEqual(total_lp_, sum(daily_lps), places=3)
def test_batch_shape(self): batch_shape = [3, 2] partial_batch_shape = [2] num_seasons = 24 initial_state_prior = tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder( np.exp(np.random.randn(*(partial_batch_shape + [num_seasons]))))) drift_scale = self._build_placeholder( np.exp(np.random.randn(*batch_shape))) observation_noise_scale = self._build_placeholder( np.exp(np.random.randn(*partial_batch_shape))) ssm = SeasonalStateSpaceModel( num_timesteps=9, num_seasons=24, num_steps_per_season=2, drift_scale=drift_scale, observation_noise_scale=observation_noise_scale, initial_state_prior=initial_state_prior) # First check that the model's batch shape is the broadcast batch shape # of parameters, as expected. self.assertAllEqual(self.evaluate(ssm.batch_shape_tensor()), batch_shape) y_ = self.evaluate(ssm.sample()) self.assertAllEqual(y_.shape[:-2], batch_shape) # Next check that the broadcasting works as expected, and the batch log_prob # actually matches the log probs of independent models. individual_ssms = [SeasonalStateSpaceModel( num_timesteps=9, num_seasons=num_seasons, num_steps_per_season=2, drift_scale=drift_scale[i, j, ...], observation_noise_scale=observation_noise_scale[j, ...], initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=initial_state_prior.scale.diag[j, ...])) for i in range(batch_shape[0]) for j in range(batch_shape[1])] batch_lps_ = self.evaluate(ssm.log_prob(y_)).flatten() individual_ys = [y_[i, j, ...] for i in range(batch_shape[0]) for j in range(batch_shape[1])] individual_lps_ = self.evaluate([ individual_ssm.log_prob(individual_y) for (individual_ssm, individual_y) in zip(individual_ssms, individual_ys)]) self.assertAllClose(individual_lps_, batch_lps_)
def test_month_of_year_example(self): num_days_per_month = np.array( [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31]) # put wildly different near-deterministic priors on the effect for each # month, so we can easily distinguish the months and diagnose off-by-one # errors. monthly_effect_prior_means = np.linspace(-1000, 1000, 12) monthly_effect_prior_scales = 0.1 * np.ones(12) drift_scale = 0.3 observation_noise_scale = 0.1 num_timesteps = 365 * 2 # 2 years. initial_step = 22 month_of_year = SeasonalStateSpaceModel( num_timesteps=num_timesteps, num_seasons=12, num_steps_per_season=num_days_per_month, drift_scale=self._build_placeholder(drift_scale), observation_noise_scale=observation_noise_scale, initial_state_prior=tfd.MultivariateNormalDiag( loc=self._build_placeholder(monthly_effect_prior_means), scale_diag=self._build_placeholder( monthly_effect_prior_scales)), initial_step=initial_step) sampled_series_, prior_mean_, prior_variance_ = self.evaluate( (month_of_year.sample()[..., 0], month_of_year.mean()[..., 0], month_of_year.variance()[..., 0])) # For each month, ensure the mean (and samples) of each day matches the # expected effect for that month, and that the variances all match # including expected drift from year to year current_idx = -initial_step # start at beginning of year current_drift_variance = 0. for _ in range( 3): # loop over three years to include entire 2-year period for (num_days_in_month, prior_effect_mean, prior_effect_scale) in zip(num_days_per_month, monthly_effect_prior_means, monthly_effect_prior_scales): month_indices = range(current_idx, current_idx + num_days_in_month) current_idx += num_days_in_month # Trim any indices outside of the observed window. month_indices = [ idx for idx in month_indices if 0 <= idx < num_timesteps ] if month_indices: ones = np.ones(len(month_indices)) self.assertAllClose(sampled_series_[month_indices], prior_effect_mean * ones, atol=20.) self.assertAllClose(prior_mean_[month_indices], prior_effect_mean * ones, atol=1e-4) self.assertAllClose( prior_variance_[month_indices], (prior_effect_scale**2 + current_drift_variance + observation_noise_scale**2) * ones, atol=1e-4) # At the end of each year, allow the effects to drift. current_drift_variance += drift_scale**2
def test_month_of_year_example(self): num_days_per_month = np.array( [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31]) # put wildly different near-deterministic priors on the effect for each # month, so we can easily distinguish the months and diagnose off-by-one # errors. monthly_effect_prior_means = np.linspace(-1000, 1000, 12) monthly_effect_prior_scales = 0.1 * np.ones(12) drift_scale = 0.3 observation_noise_scale = 0.1 num_timesteps = 365 * 2 # 2 years. initial_step = 22 month_of_year = SeasonalStateSpaceModel( num_timesteps=num_timesteps, num_seasons=12, num_steps_per_season=num_days_per_month, drift_scale=self._build_placeholder(drift_scale), observation_noise_scale=observation_noise_scale, initial_state_prior=tfd.MultivariateNormalDiag( loc=self._build_placeholder(monthly_effect_prior_means), scale_diag=self._build_placeholder(monthly_effect_prior_scales)), initial_step=initial_step) sampled_series_, prior_mean_, prior_variance_ = self.evaluate( (month_of_year.sample()[..., 0], month_of_year.mean()[..., 0], month_of_year.variance()[..., 0])) # For each month, ensure the mean (and samples) of each day matches the # expected effect for that month, and that the variances all match # including expected drift from year to year current_idx = -initial_step # start at beginning of year current_drift_variance = 0. for _ in range(3): # loop over three years to include entire 2-year period for (num_days_in_month, prior_effect_mean, prior_effect_scale) in zip( num_days_per_month, monthly_effect_prior_means, monthly_effect_prior_scales): month_indices = range(current_idx, current_idx + num_days_in_month) current_idx += num_days_in_month # Trim any indices outside of the observed window. month_indices = [idx for idx in month_indices if 0 <= idx < num_timesteps] if month_indices: ones = np.ones(len(month_indices)) self.assertAllClose(sampled_series_[month_indices], prior_effect_mean * ones, atol=20.) self.assertAllClose(prior_mean_[month_indices], prior_effect_mean * ones, atol=1e-4) self.assertAllClose(prior_variance_[month_indices], (prior_effect_scale**2 + current_drift_variance + observation_noise_scale**2) * ones, atol=1e-4) # At the end of each year, allow the effects to drift. current_drift_variance += drift_scale**2