def testRunningCovarianceMaxPoints(self): window_size = 100 rng = np.random.RandomState(_test_seed()) data = tf.convert_to_tensor( np.concatenate( [ rng.randn(window_size, 2), np.array([1., 2.]) + np.array([2., 3.]) * rng.randn(window_size * 10, 2) ], axis=0, )) def kernel(rvs, idx): rvs, _ = fun_mcmc.running_covariance_step( rvs, data[idx], window_size=window_size) return (rvs, idx + 1), (rvs.mean, rvs.covariance) _, (mean, cov) = fun_mcmc.trace( state=(fun_mcmc.running_covariance_init([2], data.dtype), 0), fn=kernel, num_steps=len(data), ) # Up to window_size, we compute the running mean/variance exactly. self.assertAllClose( np.mean(data[:window_size], axis=0), mean[window_size - 1]) self.assertAllClose( _gen_cov(data[:window_size], axis=0), cov[window_size - 1]) # After window_size, we're doing exponential moving average, and pick up the # mean/variance after the change in the distribution. Since the moving # average is computed only over ~window_size points, this test is rather # noisy. self.assertAllClose(np.array([1., 2.]), mean[-1], atol=0.2) self.assertAllClose(np.array([[4., 0.], [0., 9.]]), cov[-1], atol=1.)
def testRunningCovariance(self, shape, aggregation): data = tf.convert_to_tensor(np.random.randn(*shape)) true_aggregation = (0, ) + (() if aggregation is None else tuple( [a + 1 for a in util.flatten_tree(aggregation)])) true_mean = np.mean(data, true_aggregation) true_cov = _gen_cov(data, true_aggregation) def kernel(rcs, idx): rcs, _ = fun_mcmc.running_covariance_step(rcs, data[idx], axis=aggregation) return (rcs, idx + 1), () (rcs, _), _ = fun_mcmc.trace(state=(fun_mcmc.running_covariance_init( true_mean.shape, data[0].dtype), 0), fn=kernel, num_steps=len(data), trace_fn=lambda *args: ()) self.assertAllClose(true_mean, rcs.mean) self.assertAllClose(true_cov, rcs.covariance)