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