def test_nesting_additive_ssms(self): ssm1 = self._dummy_model(batch_shape=[1, 2]) ssm2 = self._dummy_model(batch_shape=[3, 2]) observation_noise_scale = 0.1 additive_ssm = AdditiveStateSpaceModel( [ssm1, ssm2], observation_noise_scale=observation_noise_scale) nested_additive_ssm = AdditiveStateSpaceModel( [AdditiveStateSpaceModel([ssm1]), AdditiveStateSpaceModel([ssm2])], observation_noise_scale=observation_noise_scale) # Test that both models behave equivalently. y = self.evaluate(nested_additive_ssm.sample()) additive_lp = additive_ssm.log_prob(y) nested_additive_lp = nested_additive_ssm.log_prob(y) self.assertAllClose(self.evaluate(additive_lp), self.evaluate(nested_additive_lp)) additive_mean = additive_ssm.mean() nested_additive_mean = nested_additive_ssm.mean() self.assertAllClose(self.evaluate(additive_mean), self.evaluate(nested_additive_mean)) additive_variance = additive_ssm.variance() nested_additive_variance = nested_additive_ssm.variance() self.assertAllClose(self.evaluate(additive_variance), self.evaluate(nested_additive_variance))
def test_broadcasting_correctness(self): # This test verifies that broadcasting of component parameters works as # expected. We construct a SSM with no batch shape, and test that when we # add it to another SSM of batch shape [3], we get the same model # as if we had explicitly broadcast the parameters of the first SSM before # adding. num_timesteps = 5 transition_matrix = np.random.randn(2, 2) transition_noise_diag = np.exp(np.random.randn(2)) observation_matrix = np.random.randn(1, 2) observation_noise_diag = np.exp(np.random.randn(1)) initial_state_prior_diag = np.exp(np.random.randn(2)) # First build the model in which we let AdditiveSSM do the broadcasting. batchless_ssm = tfd.LinearGaussianStateSpaceModel( num_timesteps=num_timesteps, transition_matrix=self._build_placeholder(transition_matrix), transition_noise=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder(transition_noise_diag)), observation_matrix=self._build_placeholder(observation_matrix), observation_noise=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder(observation_noise_diag)), initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder(initial_state_prior_diag)) ) another_ssm = self._dummy_model(num_timesteps=num_timesteps, latent_size=4, batch_shape=[3]) broadcast_additive_ssm = AdditiveStateSpaceModel( [batchless_ssm, another_ssm]) # Next try doing our own broadcasting explicitly. broadcast_vector = np.ones([3, 1]) broadcast_matrix = np.ones([3, 1, 1]) batch_ssm = tfd.LinearGaussianStateSpaceModel( num_timesteps=num_timesteps, transition_matrix=self._build_placeholder( transition_matrix * broadcast_matrix), transition_noise=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder( transition_noise_diag * broadcast_vector)), observation_matrix=self._build_placeholder( observation_matrix * broadcast_matrix), observation_noise=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder( observation_noise_diag * broadcast_vector)), initial_state_prior=tfd.MultivariateNormalDiag( scale_diag=self._build_placeholder( initial_state_prior_diag * broadcast_vector))) manual_additive_ssm = AdditiveStateSpaceModel([batch_ssm, another_ssm]) # Both additive SSMs define the same model, so they should give the same # log_probs. y = self.evaluate(broadcast_additive_ssm.sample(seed=42)) self.assertAllEqual(self.evaluate(broadcast_additive_ssm.log_prob(y)), self.evaluate(manual_additive_ssm.log_prob(y)))
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_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))