def test_logprob(self): y = self._build_placeholder([1.0, 2.5, 4.3, 6.1, 7.8]) ssm = LocalLinearTrendStateSpaceModel( num_timesteps=5, level_scale=0.5, slope_scale=0.5, initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder([1., 1.]))) lp = ssm.log_prob(y[..., np.newaxis]) expected_lp = -5.801624298095703 self.assertAllClose(self.evaluate(lp), expected_lp)
def test_joint_sample(self): strm = test_util.test_seed_stream() batch_shape = [4, 3] level_scale = self._build_placeholder(2 * np.ones(batch_shape)) slope_scale = self._build_placeholder(0.2 * np.ones(batch_shape)) observation_noise_scale = self._build_placeholder(1.) initial_state_prior = tfd.MultivariateNormalDiag( loc=self._build_placeholder([-3, 0.5]), scale_diag=self._build_placeholder([0.1, 0.2])) ssm = LocalLinearTrendStateSpaceModel( num_timesteps=10, level_scale=level_scale, slope_scale=slope_scale, observation_noise_scale=observation_noise_scale, initial_state_prior=initial_state_prior) num_samples = 10000 sampled_latents, sampled_obs = ssm._joint_sample_n(n=num_samples, seed=strm()) latent_mean, obs_mean = ssm._joint_mean() latent_cov, obs_cov = ssm._joint_covariances() (sampled_latents_, sampled_obs_, latent_mean_, obs_mean_, latent_level_std_, level_slope_std_, obs_std_) = self.evaluate( (sampled_latents, sampled_obs, latent_mean, obs_mean, tf.sqrt(latent_cov[..., 0, 0]), tf.sqrt(latent_cov[..., 1, 1]), tf.sqrt(obs_cov[..., 0]))) latent_std_ = np.stack([latent_level_std_, level_slope_std_], axis=-1) # Instead of directly comparing means and stddevs, we normalize by stddev # to make the stderr constant. self.assertAllClose(np.mean(sampled_latents_, axis=0) / latent_std_, latent_mean_ / latent_std_, atol=4. / np.sqrt(num_samples)) self.assertAllClose(np.mean(sampled_obs_, axis=0) / obs_std_, obs_mean_ / obs_std_, atol=4. / np.sqrt(num_samples)) self.assertAllClose(np.std(sampled_latents_, axis=0) / latent_std_, np.ones(latent_std_.shape, dtype=latent_std_.dtype), atol=4. / np.sqrt(num_samples)) self.assertAllClose(np.std(sampled_obs_, axis=0) / obs_std_, np.ones(obs_std_.shape, dtype=obs_std_.dtype), atol=4. / np.sqrt(num_samples))
def test_batch_shape(self): batch_shape = [4, 2] partial_batch_shape = [2] level_scale = self._build_placeholder( np.exp(np.random.randn(*partial_batch_shape))) slope_scale = self._build_placeholder(np.exp(np.random.randn(*batch_shape))) initial_state_prior = tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder([1., 1.])) ssm = LocalLinearTrendStateSpaceModel( num_timesteps=10, level_scale=level_scale, slope_scale=slope_scale, initial_state_prior=initial_state_prior) self.assertAllEqual(self.evaluate(ssm.batch_shape_tensor()), batch_shape) y = ssm.sample() self.assertAllEqual(self.evaluate(tf.shape(input=y))[:-2], batch_shape)
def test_matches_locallineartrend(self): """SemiLocalLinearTrend with trivial AR process is a LocalLinearTrend.""" level_scale = self._build_placeholder(0.5) slope_scale = self._build_placeholder(0.5) initial_level = self._build_placeholder(3.) initial_slope = self._build_placeholder(-2.) num_timesteps = 5 y = self._build_placeholder([1.0, 2.5, 4.3, 6.1, 7.8]) semilocal_ssm = SemiLocalLinearTrendStateSpaceModel( num_timesteps=num_timesteps, level_scale=level_scale, slope_scale=slope_scale, slope_mean=self._build_placeholder(0.), autoregressive_coef=self._build_placeholder(1.), initial_state_prior=tfd.MultivariateNormalDiag( loc=[initial_level, initial_slope], scale_diag=self._build_placeholder([1., 1.]))) local_ssm = LocalLinearTrendStateSpaceModel( num_timesteps=num_timesteps, level_scale=level_scale, slope_scale=slope_scale, initial_state_prior=tfd.MultivariateNormalDiag( loc=[initial_level, initial_slope], scale_diag=self._build_placeholder([1., 1.]))) semilocal_lp = semilocal_ssm.log_prob(y[:, tf.newaxis]) local_lp = local_ssm.log_prob(y[:, tf.newaxis]) self.assertAllClose(self.evaluate(semilocal_lp), self.evaluate(local_lp)) semilocal_mean = semilocal_ssm.mean() local_mean = local_ssm.mean() self.assertAllClose(self.evaluate(semilocal_mean), self.evaluate(local_mean)) semilocal_variance = semilocal_ssm.variance() local_variance = local_ssm.variance() self.assertAllClose(self.evaluate(semilocal_variance), self.evaluate(local_variance))
def test_identity(self): # Test that an additive SSM with a single component defines the same # distribution as the component model. y = self._build_placeholder([1.0, 2.5, 4.3, 6.1, 7.8]) local_ssm = LocalLinearTrendStateSpaceModel( num_timesteps=5, level_scale=0.3, slope_scale=0.6, observation_noise_scale=0.1, initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder([1., 1.]))) additive_ssm = AdditiveStateSpaceModel([local_ssm]) local_lp = local_ssm.log_prob(y[:, np.newaxis]) additive_lp = additive_ssm.log_prob(y[:, np.newaxis]) self.assertAllClose(self.evaluate(local_lp), self.evaluate(additive_lp))
def test_stats(self): # Build a model with expected initial loc 0 and slope 1. level_scale = self._build_placeholder(1.0) slope_scale = self._build_placeholder(1.0) initial_state_prior = tfd.MultivariateNormalDiag( loc=self._build_placeholder([0, 1.]), scale_diag=self._build_placeholder([1., 1.])) ssm = LocalLinearTrendStateSpaceModel( num_timesteps=10, level_scale=level_scale, slope_scale=slope_scale, initial_state_prior=initial_state_prior) # In expectation, the process grows linearly. mean = self.evaluate(ssm.mean()) self.assertAllClose(mean, np.arange(0, 10)[:, np.newaxis]) # slope variance at time T is linear: T * slope_scale expected_variance = [1, 3, 8, 18, 35, 61, 98, 148, 213, 295] variance = self.evaluate(ssm.variance()) self.assertAllClose(variance, np.array(expected_variance)[:, np.newaxis])
def test_sum_of_local_linear_trends(self): # We know analytically that the sum of two local linear trends is # another local linear trend, with means and variances scaled # accordingly, so the additive model should match this behavior. level_scale = 0.5 slope_scale = 1.1 initial_level = 3. initial_slope = -2. observation_noise_scale = 0. num_timesteps = 5 y = self._build_placeholder([1.0, 2.5, 4.3, 6.1, 7.8]) # Combine two local linear trend models, one a full model, the other # with just a moving mean (zero slope). local_ssm = LocalLinearTrendStateSpaceModel( num_timesteps=num_timesteps, level_scale=level_scale, slope_scale=slope_scale, initial_state_prior=tfd.MultivariateNormalDiag( loc=self._build_placeholder([initial_level, initial_slope]), scale_diag=self._build_placeholder([1., 1.]))) second_level_scale = 0.3 second_initial_level = 1.1 moving_level_ssm = LocalLinearTrendStateSpaceModel( num_timesteps=num_timesteps, level_scale=second_level_scale, slope_scale=0., initial_state_prior=tfd.MultivariateNormalDiag( loc=self._build_placeholder([second_initial_level, 0.]), scale_diag=self._build_placeholder([1., 0.]))) additive_ssm = AdditiveStateSpaceModel( [local_ssm, moving_level_ssm], observation_noise_scale=observation_noise_scale) # Build the analytical sum of the two processes. target_ssm = LocalLinearTrendStateSpaceModel( num_timesteps=num_timesteps, level_scale=np.float32( np.sqrt(level_scale**2 + second_level_scale**2)), slope_scale=np.float32(slope_scale), observation_noise_scale=observation_noise_scale, initial_state_prior=tfd.MultivariateNormalDiag( loc=self._build_placeholder( [initial_level + second_initial_level, initial_slope + 0.]), scale_diag=self._build_placeholder(np.sqrt([2., 1.])))) # Test that both models behave equivalently. additive_mean = additive_ssm.mean() target_mean = target_ssm.mean() self.assertAllClose(self.evaluate(additive_mean), self.evaluate(target_mean)) additive_variance = additive_ssm.variance() target_variance = target_ssm.variance() self.assertAllClose(self.evaluate(additive_variance), self.evaluate(target_variance)) additive_lp = additive_ssm.log_prob(y[:, np.newaxis]) target_lp = target_ssm.log_prob(y[:, np.newaxis]) self.assertAllClose(self.evaluate(additive_lp), self.evaluate(target_lp))