def testShape(self):
x_shape = tf.TensorShape([5, 4, 6])
y_shape = tf.TensorShape([5, 4, 4, 4])
b = tfb.CorrelationCholesky(validate_args=True)
x = tf.ones(shape=x_shape, dtype=tf.float32)
y_ = b.forward(x)
self.assertAllEqual(
tensorshape_util.as_list(y_.shape), tensorshape_util.as_list(y_shape))
x_ = b.inverse(y_)
self.assertAllEqual(
tensorshape_util.as_list(x_.shape), tensorshape_util.as_list(x_shape))
y_shape_ = b.forward_event_shape(x_shape)
self.assertAllEqual(
tensorshape_util.as_list(y_shape_), tensorshape_util.as_list(y_shape))
x_shape_ = b.inverse_event_shape(y_shape)
self.assertAllEqual(
tensorshape_util.as_list(x_shape_), tensorshape_util.as_list(x_shape))
y_shape_tensor = self.evaluate(
b.forward_event_shape_tensor(tensorshape_util.as_list(x_shape)))
self.assertAllEqual(y_shape_tensor, tensorshape_util.as_list(y_shape))
x_shape_tensor = self.evaluate(
b.inverse_event_shape_tensor(tensorshape_util.as_list(y_shape)))
self.assertAllEqual(x_shape_tensor, tensorshape_util.as_list(x_shape))
def testBijectorBatch(self):
x = np.float32([[7., -5., 5., 1., 2., -2.], [1., 3., -5., 1., -4., 8.]])
y = np.float32([
[[1., 0., 0., 0.], [0.707107, 0.707107, 0., 0.],
[-0.666667, 0.666667, 0.333333, 0.], [0.5, -0.5, 0.7, 0.1]],
[[1., 0., 0., 0.], [0.707107, 0.707107, 0., 0.],
[0.888889, -0.444444, 0.111111, 0.],
[-0.833333, 0.5, 0.166667, 0.166667]],
])
b = tfb.CorrelationCholesky()
y_ = self.evaluate(b.forward(x))
self.assertAllClose(y, y_, atol=1e-5, rtol=1e-5)
x_ = self.evaluate(b.inverse(y))
self.assertAllClose(x, x_, atol=1e-5, rtol=1e-5)
expected_fldj = -0.5 * np.sum(
[3, 4, 5] * np.log([[2, 9, 100], [2, 81, 36]]), axis=-1)
fldj = self.evaluate(b.forward_log_det_jacobian(x, event_ndims=1))
self.assertAllClose(expected_fldj, fldj)
ildj = self.evaluate(b.inverse_log_det_jacobian(y, event_ndims=2))
self.assertAllClose(-expected_fldj, ildj)
def testBijectorWithVariables(self):
x_ = np.array([1.], dtype=np.float32)
y_ = np.array([[1., 0.], [0.707107, 0.707107]], dtype=np.float32)
x = tf.Variable(x_, dtype=tf.float32)
y = tf.Variable(y_, dtype=tf.float32)
forward_event_ndims = tf.Variable(1, dtype=tf.int32)
inverse_event_ndims = tf.Variable(2, dtype=tf.int32)
self.evaluate([
v.initializer
for v in (x, y, forward_event_ndims, inverse_event_ndims)
])
bijector = tfb.CorrelationCholesky()
self.assertAllClose(y_,
self.evaluate(bijector.forward(x)),
atol=1e-5,
rtol=1e-5)
self.assertAllClose(x_,
self.evaluate(bijector.inverse(y)),
atol=1e-5,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(
x, event_ndims=forward_event_ndims)
self.assertAllClose(-3 * 0.5 * np.log(2), self.evaluate(fldj))
ildj = bijector.inverse_log_det_jacobian(
y, event_ndims=inverse_event_ndims)
self.assertAllClose(3 * 0.5 * np.log(2), ildj)
def testWithLKJSamples(self, dimension, concentration):
bijector = tfb.CorrelationCholesky()
lkj_dist = lkj.LKJ(dimension=dimension,
concentration=np.float64(concentration),
input_output_cholesky=True)
batch_size = 10
y = self.evaluate(lkj_dist.sample([batch_size]))
x = self.evaluate(bijector.inverse(y))
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=1,
inverse_event_ndims=2,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector,
x,
event_ndims=1,
inverse_event_ndims=2,
output_to_unconstrained=tfb.Invert(tfb.FillTriangular()))
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testBijectiveWithLKJSamples(self, dimension, concentration):
bijector = tfb.CorrelationCholesky()
lkj_dist = lkj.LKJ(dimension=dimension,
concentration=np.float64(concentration),
input_output_cholesky=True)
batch_size = 10
y = self.evaluate(lkj_dist.sample([batch_size]))
x = self.evaluate(bijector.inverse(y))
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=1,
inverse_event_ndims=2,
rtol=1e-5)
def testBijector(self):
x = np.float32(np.array([7., -5., 5., 1., 2., -2.]))
y = np.float32(
np.array([[1., 0., 0., 0.], [0.707107, 0.707107, 0., 0.],
[-0.666667, 0.666667, 0.333333, 0.], [0.5, -0.5, 0.7, 0.1]]))
b = tfb.CorrelationCholesky()
y_ = self.evaluate(b.forward(x))
self.assertAllClose(y, y_, atol=1e-5, rtol=1e-5)
x_ = self.evaluate(b.inverse(y))
self.assertAllClose(x, x_, atol=1e-5, rtol=1e-5)
expected_fldj = -0.5 * np.sum([3, 4, 5] * np.log([2, 9, 100]))
fldj = self.evaluate(b.forward_log_det_jacobian(x, event_ndims=1))
self.assertAllClose(expected_fldj, fldj)
ildj = self.evaluate(b.inverse_log_det_jacobian(y, event_ndims=2))
self.assertAllClose(-expected_fldj, ildj)
def testJacobianWithLKJSamples(self, dimension, concentration):
bijector = tfb.CorrelationCholesky()
lkj_dist = lkj.LKJ(dimension=dimension,
concentration=np.float64(concentration),
input_output_cholesky=True)
batch_size = 10
y = self.evaluate(
lkj_dist.sample([batch_size], seed=test_util.test_seed()))
x = self.evaluate(bijector.inverse(y))
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector,
x,
event_ndims=1,
inverse_event_ndims=2,
output_to_unconstrained=OutputToUnconstrained())
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testTheoreticalFldj(self):
bijector = tfb.CorrelationCholesky()
x = np.linspace(-50, 50, num=30).reshape(5, 6).astype(np.float64)
y = self.evaluate(bijector.forward(x))
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=1,
inverse_event_ndims=2,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector,
x,
event_ndims=1,
inverse_event_ndims=2,
output_to_unconstrained=OutputToUnconstrained())
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testSampleMarginals(self):
# Verify that the marginals of the LKJ distribution are distributed
# according to a (scaled) Beta distribution. The LKJ distributed samples are
# obtained by sampling a CholeskyLKJ distribution using HMC and the
# CorrelationCholesky bijector.
dim = 4
concentration = np.array(2.5, dtype=np.float64)
beta_concentration = np.array(.5 * dim + concentration - 1, np.float64)
beta_dist = beta.Beta(
concentration0=beta_concentration, concentration1=beta_concentration)
inner_kernel = hmc.HamiltonianMonteCarlo(
target_log_prob_fn=cholesky_lkj.CholeskyLKJ(
dimension=dim, concentration=concentration).log_prob,
num_leapfrog_steps=3,
step_size=0.3)
kernel = transformed_kernel.TransformedTransitionKernel(
inner_kernel=inner_kernel, bijector=tfb.CorrelationCholesky())
num_chains = 10
num_total_samples = 30000
# Make sure that we have enough samples to catch a wrong sampler to within
# a small enough discrepancy.
self.assertLess(
self.evaluate(
st.min_num_samples_for_dkwm_cdf_test(
discrepancy=0.04, false_fail_rate=1e-9, false_pass_rate=1e-9)),
num_total_samples)
@tf.function # Ensure that MCMC sampling is done efficiently.
def sample_mcmc_chain():
return sample.sample_chain(
num_results=num_total_samples // num_chains,
num_burnin_steps=1000,
current_state=tf.eye(dim, batch_shape=[num_chains], dtype=tf.float64),
trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
kernel=kernel,
seed=test_util.test_seed())
# Draw samples from the HMC chains.
chol_lkj_samples, is_accepted = self.evaluate(sample_mcmc_chain())
# Ensure that the per-chain acceptance rate is high enough.
self.assertAllGreater(np.mean(is_accepted, axis=0), 0.8)
# Transform from Cholesky LKJ samples to LKJ samples.
lkj_samples = tf.matmul(chol_lkj_samples, chol_lkj_samples, adjoint_b=True)
lkj_samples = tf.reshape(lkj_samples, shape=[num_total_samples, dim, dim])
# Only look at the entries strictly below the diagonal which is achieved by
# the OutputToUnconstrained bijector. Also scale the marginals from the
# range [-1,1] to [0,1].
scaled_lkj_samples = .5 * (OutputToUnconstrained().forward(lkj_samples) + 1)
# Each of the off-diagonal marginals should be distributed according to a
# Beta distribution.
for i in range(dim * (dim - 1) // 2):
self.evaluate(
st.assert_true_cdf_equal_by_dkwm(
scaled_lkj_samples[..., i],
cdf=beta_dist.cdf,
false_fail_rate=1e-9))