def testRunningMeanMaxPoints(self): window_size = 100 rng = np.random.RandomState(_test_seed()) data = tf.convert_to_tensor( np.concatenate([ rng.randn(window_size), 1. + 2. * rng.randn(window_size * 10) ], axis=0)) def kernel(rms, idx): rms, _ = fun_mcmc.running_mean_step(rms, data[idx], window_size=window_size) return (rms, idx + 1), rms.mean _, mean = fun_mcmc.trace( state=(fun_mcmc.running_mean_init([], data.dtype), 0), fn=kernel, num_steps=len(data), ) # Up to window_size, we compute the running mean exactly. self.assertAllClose(np.mean(data[:window_size]), mean[window_size - 1]) # After window_size, we're doing exponential moving average, and pick up the # mean after the change in the distribution. Since the moving average is # computed only over ~window_size points, this test is rather noisy. self.assertAllClose(1., mean[-1], atol=0.2)
def SanitizedAutoCorrelationMean(x, axis, reduce_axis, max_lags=None, **kwargs): shape_arr = np.array(list(x.shape)) axes = list(sorted(set(range(len(shape_arr))) - set([reduce_axis]))) mean_shape = shape_arr[axes] if max_lags is not None: mean_shape[axis] = max_lags + 1 mean_state = fun_mcmc.running_mean_init(mean_shape, x.dtype) new_order = list(range(len(shape_arr))) new_order[0] = new_order[reduce_axis] new_order[reduce_axis] = 0 x = tf.transpose(x, new_order) x_arr = tf.TensorArray(x.dtype, x.shape[0]).unstack(x) mean_state, _ = fun_mcmc.trace( state=mean_state, fn=lambda state: fun_mcmc.running_mean_step( # pylint: disable=g-long-lambda state, SanitizedAutoCorrelation(x_arr.read(state.num_points), axis, max_lags=max_lags, **kwargs)), num_steps=x.shape[0], trace_fn=lambda *_: ()) return mean_state.mean
def testRunningMean(self, shape, aggregation): rng = np.random.RandomState(_test_seed()) data = tf.convert_to_tensor(rng.randn(*shape)) def kernel(rms, idx): rms, _ = fun_mcmc.running_mean_step(rms, data[idx], axis=aggregation) return (rms, idx + 1), () 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) (rms, _), _ = fun_mcmc.trace( state=(fun_mcmc.running_mean_init(true_mean.shape, data.dtype), 0), fn=kernel, num_steps=len(data), trace_fn=lambda *args: ()) self.assertAllClose(true_mean, rms.mean)
def testSimpleDualAverages(self): def loss_fn(x, y): return tf.square(x - 1.) + tf.square(y - 2.), [] def kernel(sda_state, rms_state): sda_state, _ = fun_mcmc.simple_dual_averages_step(sda_state, loss_fn, 1.) rms_state, _ = fun_mcmc.running_mean_step(rms_state, sda_state.state) return (sda_state, rms_state), rms_state.mean _, (x, y) = fun_mcmc.trace( ( fun_mcmc.simple_dual_averages_init( [self._constant(0.), self._constant(0.)]), fun_mcmc.running_mean_init([[], []], [self._dtype, self._dtype]), ), kernel, num_steps=1000, ) self.assertAllClose(1., x[-1], atol=1e-1) self.assertAllClose(2., y[-1], atol=1e-1)