def _maximum_likelihood_parameters(cls, value):
return {
'loc':
tf.reduce_mean(value, axis=0),
'covariance_matrix':
tfp_stats.covariance(value, sample_axis=0, event_axis=-1)
}
def testSampleStatsMatchDistributionStats(self):
mvn = self._construct_mvn(
loc_shape=(2, 3),
diag_shape=(3, 2, 3),
update_shape=(3, 1),
)
n = 1000
samples = mvn.sample(n, seed=test_util.test_seed())
s_mean, mean, s_cov, cov = self.evaluate([
tf.reduce_mean(samples, axis=0),
mvn.mean(),
tfps.covariance(samples, sample_axis=0),
mvn.covariance(),
])
maxstddev = np.sqrt(np.max(cov))
self.assertAllClose(s_mean, mean, atol=5 * maxstddev / np.sqrt(n))
self.assertAllClose(s_cov, cov, atol=5 * maxstddev**2 / np.sqrt(n))
def testVersusMVNTriL(
self,
loc_shape,
diag_shape,
update_shape,
include_diag_update,
dtype,
):
mvn_pair = self._mvn_pair(loc=self._construct_loc(loc_shape,
dtype=dtype),
cov_operator=self._construct_cov_operator(
diag_shape,
update_shape,
include_diag_update=False,
dtype=dtype))
tril = mvn_pair.tril
low_rank_update = mvn_pair.low_rank_update
with self.subTest('Shapes are equal'):
self.assertAllEqual(tril.batch_shape, low_rank_update.batch_shape)
self.assertAllEqual(tril.event_shape, low_rank_update.event_shape)
self.assertAllEqual(*self.evaluate([
tril.batch_shape_tensor(),
low_rank_update.batch_shape_tensor()
]))
self.assertAllEqual(*self.evaluate([
tril.event_shape_tensor(),
low_rank_update.event_shape_tensor()
]))
with self.subTest('Statistics are almost equal'):
self.assertAllClose(*self.evaluate(
[tril.mode(), low_rank_update.mode()]))
self.assertAllClose(*self.evaluate(
[tril.mean(), low_rank_update.mean()]))
self.assertAllClose(*self.evaluate(
[tril.stddev(), low_rank_update.stddev()]))
self.assertAllClose(*self.evaluate(
[tril.covariance(),
low_rank_update.covariance()]))
self.assertAllClose(*self.evaluate(
[tril.variance(), low_rank_update.variance()]))
self.assertAllClose(*self.evaluate(
[tril.entropy(), low_rank_update.entropy()]))
with self.subTest('Samples are correct'):
n = 10000
samples = low_rank_update.sample(n, seed=test_util.test_seed())
sample_mean, sample_var, sample_cov = self.evaluate([
tf.reduce_mean(samples, axis=0),
tfps.variance(samples, sample_axis=0),
tfps.covariance(samples, sample_axis=0),
])
ref_samples = tril.sample(n, seed=test_util.test_seed())
self.assertAllEqual(ref_samples.shape, samples.shape)
maxstddev = np.max(self.evaluate(low_rank_update.stddev()))
self.assertAllClose(sample_mean,
self.evaluate(low_rank_update.mean()),
atol=5 * maxstddev / np.sqrt(n))
self.assertAllClose(sample_var,
self.evaluate(low_rank_update.variance()),
rtol=5 * maxstddev / np.sqrt(n))
self.assertAllClose(sample_cov,
self.evaluate(low_rank_update.covariance()),
atol=10 * maxstddev**2 / np.sqrt(n))
with self.subTest('prob(mean) is almost equal'):
x = sample_mean
self.assertAllClose(*self.evaluate(
[tril.prob(x), low_rank_update.prob(x)]),
rtol=1e-5)
self.assertAllClose(*self.evaluate(
[tril.log_prob(x),
low_rank_update.log_prob(x)]),
rtol=1e-5)
with self.subTest('prob(widely dispersed) is almost equal'):
x = self.evaluate(ref_samples * 5) # Widely dispersed samples.
self.assertAllClose(*self.evaluate(
[tril.log_prob(x),
low_rank_update.log_prob(x)]),
rtol=1e-5)
def _maximum_likelihood_parameters(cls, value):
return {'loc': tf.reduce_mean(value, axis=0),
'scale_tril': tf.linalg.cholesky(
tfp_stats.covariance(value, sample_axis=0, event_axis=-1))}