def kalman_filter(): (_, _, filtered_covs, _, _, _, _) = (parallel_kalman_filter_lib.kalman_filter( transition_matrix=transition_matrix, transition_mean=transition_mean, transition_cov=transition_cov, observation_matrix=observation_matrix, observation_mean=observation_mean, observation_cov=observation_cov, initial_mean=initial_mean, initial_cov=initial_cov, y=y, mask=None)) return filtered_covs
def lp(initial_mean, initial_cov, transition_matrix, transition_mean, transition_cov, observation_matrix, observation_mean, observation_cov, y): (log_likelihoods, _, _, _, _, _, _) = (parallel_kalman_filter_lib.kalman_filter( transition_matrix=transition_matrix, transition_mean=transition_mean, transition_cov=transition_cov, observation_matrix=observation_matrix, observation_mean=observation_mean, observation_cov=observation_cov, initial_mean=initial_mean, initial_cov=initial_cov, y=y, mask=None)) return tf.reduce_sum(log_likelihoods)
def test_basic_example_time_dependent(self): nsteps = 7 initial_mean = np.array([0.], dtype=self.dtype) initial_cov = np.eye(1, dtype=self.dtype) transition_matrix = np.array(nsteps * [[[1.]]], dtype=self.dtype) transition_mean = np.array(nsteps * [[0.]], dtype=self.dtype) transition_cov = np.array(nsteps * [[[1.]]], dtype=self.dtype) observation_matrix = np.array(nsteps * [[[1.]]], dtype=self.dtype) observation_mean = np.array(nsteps * [[0.]], dtype=self.dtype) observation_cov = np.array(nsteps * [[[1.]]], dtype=self.dtype) _, y = parallel_kalman_filter_lib.sample_walk( transition_matrix=transition_matrix, transition_mean=transition_mean, transition_scale_tril=tf.linalg.cholesky(transition_cov), observation_matrix=observation_matrix, observation_mean=observation_mean, observation_scale_tril=tf.linalg.cholesky(observation_cov), initial_mean=initial_mean, initial_scale_tril=tf.linalg.cholesky(initial_cov), seed=test_util.test_seed()) (_, _, filtered_covs, _, _, _, _) = (parallel_kalman_filter_lib.kalman_filter( transition_matrix=transition_matrix, transition_mean=transition_mean, transition_cov=transition_cov, observation_matrix=observation_matrix, observation_mean=observation_mean, observation_cov=observation_cov, initial_mean=initial_mean, initial_cov=initial_cov, y=y, mask=None)) fibonacci = np.array( [1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610], dtype=self.dtype) approximants = fibonacci[::2] / fibonacci[1::2] self.assertAllClose(filtered_covs[..., 0, 0], approximants)
def test_basic_example_time_dependent_batched(self): batch_shape = (2, 3) ndim = 7 # Dimension of latent space mdim = 5 # Dimension of observation space nsteps = 9 Batches = collections.namedtuple('Batches', [ 'initial_mean', 'initial_cov', 'transition_matrix', 'transition_mean', 'transition_cov', 'observation_matrix', 'observation_mean', 'observation_cov', 'mask' ]) def batch_generator(): # Skipping 'mask' case because it isn't used in sample generation. for skip in range(8): batch_list = skip * [()] + [batch_shape ] + (9 - skip - 1) * [()] yield Batches(*batch_list) # Test the broadcasting by ensuring each parameter individually # can be broadcast up to the full batch size. seed = test_util.test_seed(sampler_type='stateless') for batches in batch_generator(): iter_seed, seed = samplers.split_seed(seed, n=2, salt='') s = samplers.split_seed(iter_seed, n=10, salt='') initial_mean = _random_vector(ndim, batches.initial_mean, dtype=self.dtype, seed=s[0]) initial_cov = _random_variance(ndim, batches.initial_cov, dtype=self.dtype, seed=s[1]) transition_matrix = 0.2 * _random_matrix( # Avoid blowup (eigvals > 1). ndim, ndim, (nsteps, ) + batches.transition_matrix, dtype=self.dtype, seed=s[2]) transition_mean = _random_vector(ndim, (nsteps, ) + batches.transition_mean, dtype=self.dtype, seed=s[3]) transition_cov = _random_variance(ndim, (nsteps, ) + batches.transition_cov, dtype=self.dtype, seed=s[4]) observation_matrix = _random_matrix(mdim, ndim, (nsteps, ) + batches.observation_matrix, dtype=self.dtype, seed=s[5]) observation_mean = _random_vector(mdim, (nsteps, ) + batches.observation_mean, dtype=self.dtype, seed=s[6]) observation_cov = _random_variance(mdim, (nsteps, ) + batches.observation_cov, dtype=self.dtype, seed=s[7]) mask = _random_mask((nsteps, ) + batches.mask, dtype=tf.bool, seed=s[8]) _, y = parallel_kalman_filter_lib.sample_walk( transition_matrix=transition_matrix, transition_mean=transition_mean, transition_scale_tril=tf.linalg.cholesky(transition_cov), observation_matrix=observation_matrix, observation_mean=observation_mean, observation_scale_tril=tf.linalg.cholesky(observation_cov), initial_mean=initial_mean, initial_scale_tril=tf.linalg.cholesky(initial_cov), seed=s[9]) my_filter_results = parallel_kalman_filter_lib.kalman_filter( transition_matrix=transition_matrix, transition_mean=transition_mean, transition_cov=transition_cov, observation_matrix=observation_matrix, observation_mean=observation_mean, observation_cov=observation_cov, initial_mean=initial_mean, initial_cov=initial_cov, y=y, mask=mask) ((my_log_likelihoods, my_filtered_means, my_filtered_covs, my_predicted_means, my_predicted_covs, my_observation_means, my_observation_covs), y, mask) = tf.nest.map_structure( lambda x, r: distribution_util.move_dimension(x, 0, -r), (my_filter_results, y, mask), (type(my_filter_results)(1, 2, 3, 2, 3, 2, 3), 2, 1)) # pylint: disable=g-long-lambda,cell-var-from-loop mvn = tfd.MultivariateNormalFullCovariance dist = tfd.LinearGaussianStateSpaceModel( num_timesteps=nsteps, transition_matrix=lambda t: tf.linalg.LinearOperatorFullMatrix( tf.gather(transition_matrix, t, axis=0)), transition_noise=lambda t: mvn( loc=tf.gather(transition_mean, t, axis=0), covariance_matrix=tf.gather(transition_cov, t, axis=0)), observation_matrix=lambda t: tf.linalg. LinearOperatorFullMatrix( tf.gather(observation_matrix, t, axis=0)), observation_noise=lambda t: mvn( loc=tf.gather(observation_mean, t, axis=0), covariance_matrix=tf.gather(observation_cov, t, axis=0)), initial_state_prior=mvn(loc=initial_mean, covariance_matrix=initial_cov), experimental_parallelize=False ) # Compare against sequential filter. # pylint: enable=g-long-lambda,cell-var-from-loop (log_likelihoods, filtered_means, filtered_covs, predicted_means, predicted_covs, observation_means, observation_covs) = dist.forward_filter(y, mask) rtol = (1e-6 if self.dtype == np.float64 else 1e-1) atol = (1e-6 if self.dtype == np.float64 else 1e-3) self.assertAllClose(log_likelihoods, my_log_likelihoods, rtol=rtol, atol=atol) rtol = (1e-6 if self.dtype == np.float64 else 1e-3) atol = (1e-6 if self.dtype == np.float64 else 1e-3) self.assertAllClose(filtered_means, my_filtered_means, rtol=rtol, atol=atol) self.assertAllClose(filtered_covs, my_filtered_covs, rtol=rtol, atol=atol) self.assertAllClose(predicted_means, my_predicted_means, rtol=rtol, atol=atol) self.assertAllClose(predicted_covs, my_predicted_covs, rtol=rtol, atol=atol) self.assertAllClose(observation_means, my_observation_means, rtol=rtol, atol=atol) self.assertAllClose(observation_covs, my_observation_covs, rtol=rtol, atol=atol)