def testLinearGaussianObservation(self):
prior_mean = np.asarray([[-2], [.4]], dtype=np.float32)
prior_cov = np.asarray([[.5, -.2], [-.2, .9]], dtype=np.float32)
x_observed = np.asarray([[4.], [-1.]], dtype=np.float32)
prior_mean_tensor = self.build_tensor(prior_mean)
prior_cov_tensor = self.build_tensor(prior_cov)
x_observed_tensor = self.build_tensor(x_observed)
observation_matrix = self.get_observation_matrix_for_timestep(0)
observation_noise = self.get_observation_noise_for_timestep(0)
(posterior_mean,
posterior_cov,
predictive_dist) = linear_gaussian_update(
prior_mean_tensor, prior_cov_tensor,
observation_matrix, observation_noise,
x_observed_tensor)
expected_posterior_mean = [[-0.373276], [1.65086186]]
expected_posterior_cov = [[0.24379307, 0.02689655],
[0.02689655, 0.13344827]]
expected_predicted_mean = [-2.5, -0.77999997]
expected_predicted_cov = [[1.99999988, -0.39999998],
[-0.39999998, 0.65999997]]
self.assertAllClose(self.evaluate(posterior_mean),
expected_posterior_mean)
self.assertAllClose(self.evaluate(posterior_cov),
expected_posterior_cov)
self.assertAllClose(self.evaluate(predictive_dist.mean()),
expected_predicted_mean)
self.assertAllClose(self.evaluate(predictive_dist.covariance()),
expected_predicted_cov)
def testLinearGaussianObservation(self):
prev_mean = np.asarray([[-2], [.4]], dtype=np.float32)
prev_cov = np.asarray([[.5, -.2], [-.2, .9]], dtype=np.float32)
x_observed = np.asarray([[4.]], dtype=np.float32)
prev_mean_tensor = self.build_tensor(prev_mean)
prev_cov_tensor = self.build_tensor(prev_cov)
x_observed_tensor = self.build_tensor(x_observed)
observation_matrix = self.get_observation_matrix_for_timestep(0)
observation_noise = self.get_observation_noise_for_timestep(0)
(posterior_mean,
posterior_cov,
predictive_dist) = linear_gaussian_update(
prev_mean_tensor, prev_cov_tensor,
observation_matrix, observation_noise,
x_observed_tensor)
expected_posterior_mean = [[-1.025], [2.675]]
expected_posterior_cov = [[0.455, -0.305], [-0.305, 0.655]]
expected_predicted_mean = [-2.5]
expected_predicted_cov = [[2.]]
self.assertAllClose(self.evaluate(posterior_mean),
expected_posterior_mean)
self.assertAllClose(self.evaluate(posterior_cov),
expected_posterior_cov)
self.assertAllClose(self.evaluate(predictive_dist.mean()),
expected_predicted_mean)
self.assertAllClose(self.evaluate(predictive_dist.covariance()),
expected_predicted_cov)
def testLinearGaussianObservation(self):
prev_mean = np.asarray([[-2], [.4]], dtype=np.float32)
prev_cov = np.asarray([[.5, -.2], [-.2, .9]], dtype=np.float32)
x_observed = np.asarray([[4.]], dtype=np.float32)
prev_mean_tensor = self.build_tensor(prev_mean)
prev_cov_tensor = self.build_tensor(prev_cov)
x_observed_tensor = self.build_tensor(x_observed)
observation_matrix = self.get_observation_matrix_for_timestep(0)
observation_noise = self.get_observation_noise_for_timestep(0)
(posterior_mean,
posterior_cov,
predictive_dist) = linear_gaussian_update(
prev_mean_tensor, prev_cov_tensor,
observation_matrix, observation_noise,
x_observed_tensor)
expected_posterior_mean = [[-1.025], [2.675]]
expected_posterior_cov = [[0.455, -0.305], [-0.305, 0.655]]
expected_predicted_mean = [-2.5]
expected_predicted_cov = [[2.]]
self.assertAllClose(self.evaluate(posterior_mean),
expected_posterior_mean)
self.assertAllClose(self.evaluate(posterior_cov),
expected_posterior_cov)
self.assertAllClose(self.evaluate(predictive_dist.mean()),
expected_predicted_mean)
self.assertAllClose(self.evaluate(predictive_dist.covariance()),
expected_predicted_cov)
def testLinearGaussianObservation(self):
prior_mean = np.asarray([[-2], [.4]], dtype=np.float32)
prior_cov = np.asarray([[.5, -.2], [-.2, .9]], dtype=np.float32)
x_observed = np.asarray([[4.], [-1.]], dtype=np.float32)
prior_mean_tensor = self.build_tensor(prior_mean)
prior_cov_tensor = self.build_tensor(prior_cov)
x_observed_tensor = self.build_tensor(x_observed)
observation_matrix = self.get_observation_matrix_for_timestep(0)
observation_noise = self.get_observation_noise_for_timestep(0)
(posterior_mean,
posterior_cov,
predictive_dist) = linear_gaussian_update(
prior_mean_tensor, prior_cov_tensor,
observation_matrix, observation_noise,
x_observed_tensor)
expected_posterior_mean = [[-0.373276], [1.65086186]]
expected_posterior_cov = [[0.24379307, 0.02689655],
[0.02689655, 0.13344827]]
expected_predicted_mean = [-2.5, -0.77999997]
expected_predicted_cov = [[1.99999988, -0.39999998],
[-0.39999998, 0.65999997]]
self.assertAllClose(self.evaluate(posterior_mean),
expected_posterior_mean)
self.assertAllClose(self.evaluate(posterior_cov),
expected_posterior_cov)
self.assertAllClose(self.evaluate(predictive_dist.mean()),
expected_predicted_mean)
self.assertAllClose(self.evaluate(predictive_dist.covariance()),
expected_predicted_cov)
def test_sampled_weights_follow_correct_distribution(self):
seed = test_util.test_seed(sampler_type='stateless')
design_seed, true_weights_seed, sampled_weights_seed = samplers.split_seed(
seed, 3, 'test_sampled_weights_follow_correct_distribution')
num_timesteps = 10
num_features = 2
batch_shape = [3, 1]
design_matrix = self.evaluate(samplers.normal(
batch_shape + [num_timesteps, num_features], seed=design_seed))
true_weights = self.evaluate(samplers.normal(
batch_shape + [num_features, 1], seed=true_weights_seed) * 10.0)
targets = np.matmul(design_matrix, true_weights)
is_missing = np.array([False, False, False, True, True,
False, False, True, False, False])
prior_scale = tf.convert_to_tensor(5.)
likelihood_scale = tf.convert_to_tensor(0.1)
# Analytically compute the true posterior distribution on weights.
valid_design_matrix = design_matrix[..., ~is_missing, :]
valid_targets = targets[..., ~is_missing, :]
num_valid_observations = tf.shape(valid_design_matrix)[-2]
weights_posterior_mean, weights_posterior_cov, _ = linear_gaussian_update(
prior_mean=tf.zeros([num_features, 1]),
prior_cov=tf.eye(num_features) * prior_scale**2,
observation_matrix=tfl.LinearOperatorFullMatrix(valid_design_matrix),
observation_noise=tfd.MultivariateNormalDiag(
loc=tf.zeros([num_valid_observations]),
scale_diag=likelihood_scale * tf.ones([num_valid_observations])),
x_observed=valid_targets)
# Check that the empirical moments of sampled weights match the true values.
sampled_weights = tf.vectorized_map(
lambda seed: gibbs_sampler._resample_weights( # pylint: disable=g-long-lambda
design_matrix=tf.where(is_missing[..., tf.newaxis],
tf.zeros_like(design_matrix),
design_matrix),
target_residuals=targets[..., 0],
observation_noise_scale=likelihood_scale,
weights_prior_scale=tf.linalg.LinearOperatorScaledIdentity(
num_features, prior_scale),
seed=seed),
tfp.random.split_seed(sampled_weights_seed, tf.constant(10000)))
sampled_weights_mean = tf.reduce_mean(sampled_weights, axis=0)
centered_weights = sampled_weights - weights_posterior_mean[..., 0]
sampled_weights_cov = tf.reduce_mean(centered_weights[..., :, tf.newaxis] *
centered_weights[..., tf.newaxis, :],
axis=0)
(sampled_weights_mean_, weights_posterior_mean_,
sampled_weights_cov_, weights_posterior_cov_) = self.evaluate((
sampled_weights_mean, weights_posterior_mean[..., 0],
sampled_weights_cov, weights_posterior_cov))
self.assertAllClose(sampled_weights_mean_, weights_posterior_mean_,
atol=0.01, rtol=0.05)
self.assertAllClose(sampled_weights_cov_, weights_posterior_cov_,
atol=0.01, rtol=0.05)
def testLinearGaussianObservationScalarPath(self):
# Construct observed data with a scalar observation.
prior_mean_tensor = self.build_tensor(
np.asarray([[-2], [.4]], dtype=np.float32))
prior_cov_tensor = self.build_tensor(
np.asarray([[.5, -.2], [-.2, .9]], dtype=np.float32))
x_observed_tensor = self.build_tensor(
np.asarray([[4.]], dtype=np.float32))
observation_matrix = tfl.LinearOperatorFullMatrix(
self.build_tensor([[1., 1.]]))
observation_noise = tfd.MultivariateNormalDiag(
loc=self.build_tensor([-.9]), scale_diag=self.build_tensor([1.]))
(posterior_mean,
posterior_cov,
predictive_dist) = linear_gaussian_update(
prior_mean_tensor, prior_cov_tensor,
observation_matrix, observation_noise,
x_observed_tensor)
# Ensure we take the scalar-optimized path when shape is static.
self.assertIsInstance(predictive_dist,
(tfd.Independent if self.use_static_shape
else tfd.MultivariateNormalTriL))
self.assertAllEqual(
self.evaluate(predictive_dist.event_shape_tensor()), [1])
self.assertAllEqual(
self.evaluate(predictive_dist.batch_shape_tensor()), [])
# Extend `x_observed` with an extra dimension to force the vector path.
# The added dimension is non-informative, so we can check that the scalar
# and vector paths yield the same posterior.
observation_matrix_extended = tfl.LinearOperatorFullMatrix(
self.build_tensor([[1., 1.], [0., 0.]]))
observation_noise_extended = tfd.MultivariateNormalDiag(
loc=self.build_tensor([-.9, 0.]),
scale_diag=self.build_tensor([1., 1e15]))
x_observed_extended_tensor = self.build_tensor(
np.asarray([[4.], [0.]], dtype=np.float32))
(posterior_mean_extended,
posterior_cov_extended,
predictive_dist_extended) = linear_gaussian_update(
prior_mean_tensor, prior_cov_tensor,
observation_matrix_extended, observation_noise_extended,
x_observed_extended_tensor)
# Ensure we took the vector path.
self.assertIsInstance(predictive_dist_extended, tfd.MultivariateNormalTriL)
self.assertAllEqual(
self.evaluate(predictive_dist_extended.event_shape_tensor()), [2])
self.assertAllEqual(
self.evaluate(predictive_dist_extended.batch_shape_tensor()), [])
# Test that the results are the same.
self.assertAllClose(*self.evaluate((posterior_mean,
posterior_mean_extended)))
self.assertAllClose(*self.evaluate((posterior_cov,
posterior_cov_extended)))
self.assertAllClose(*self.evaluate((predictive_dist.mean(),
predictive_dist_extended.mean()[:1])))
self.assertAllClose(
*self.evaluate((predictive_dist.covariance(),
predictive_dist_extended.covariance()[:1, :1])))
def testLinearGaussianObservationScalarPath(self):
# Construct observed data with a scalar observation.
prior_mean_tensor = self.build_tensor(
np.asarray([[-2], [.4]], dtype=np.float32))
prior_cov_tensor = self.build_tensor(
np.asarray([[.5, -.2], [-.2, .9]], dtype=np.float32))
x_observed_tensor = self.build_tensor(
np.asarray([[4.]], dtype=np.float32))
observation_matrix = tfl.LinearOperatorFullMatrix(
self.build_tensor([[1., 1.]]))
observation_noise = tfd.MultivariateNormalDiag(
loc=self.build_tensor([-.9]), scale_diag=self.build_tensor([1.]))
(posterior_mean,
posterior_cov,
predictive_dist) = linear_gaussian_update(
prior_mean_tensor, prior_cov_tensor,
observation_matrix, observation_noise,
x_observed_tensor)
# Ensure we take the scalar-optimized path when shape is static.
self.assertIsInstance(predictive_dist,
(tfd.Independent if self.use_static_shape
else tfd.MultivariateNormalTriL))
self.assertAllEqual(
self.evaluate(predictive_dist.event_shape_tensor()), [1])
self.assertAllEqual(
self.evaluate(predictive_dist.batch_shape_tensor()), [])
# Extend `x_observed` with an extra dimension to force the vector path.
# The added dimension is non-informative, so we can check that the scalar
# and vector paths yield the same posterior.
observation_matrix_extended = tfl.LinearOperatorFullMatrix(
self.build_tensor([[1., 1.], [0., 0.]]))
observation_noise_extended = tfd.MultivariateNormalDiag(
loc=self.build_tensor([-.9, 0.]),
scale_diag=self.build_tensor([1., 1e15]))
x_observed_extended_tensor = self.build_tensor(
np.asarray([[4.], [0.]], dtype=np.float32))
(posterior_mean_extended,
posterior_cov_extended,
predictive_dist_extended) = linear_gaussian_update(
prior_mean_tensor, prior_cov_tensor,
observation_matrix_extended, observation_noise_extended,
x_observed_extended_tensor)
# Ensure we took the vector path.
self.assertIsInstance(predictive_dist_extended, tfd.MultivariateNormalTriL)
self.assertAllEqual(
self.evaluate(predictive_dist_extended.event_shape_tensor()), [2])
self.assertAllEqual(
self.evaluate(predictive_dist_extended.batch_shape_tensor()), [])
# Test that the results are the same.
self.assertAllClose(*self.evaluate((posterior_mean,
posterior_mean_extended)))
self.assertAllClose(*self.evaluate((posterior_cov,
posterior_cov_extended)))
self.assertAllClose(*self.evaluate((predictive_dist.mean(),
predictive_dist_extended.mean()[:1])))
self.assertAllClose(
*self.evaluate((predictive_dist.covariance(),
predictive_dist_extended.covariance()[:1, :1])))
