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))