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))
