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)
