def testSampleGammaLogRateLogSpaceDerivatives(self): conc = tf.constant(np.linspace(.8, 1.2, 5), tf.float64) rate = np.linspace(.5, 2, 5) np.random.shuffle(rate) rate = tf.constant(rate, tf.float64) n = int(1e5) seed = test_util.test_seed() # pylint: disable=g-long-lambda lambdas = [ # Each should sample the same distribution. lambda c, r: gamma_lib.random_gamma( [n], c, r, seed=seed, log_space=True), lambda c, r: gamma_lib.random_gamma( [n], c, log_rate=tf.math.log(r), seed=seed, log_space=True), lambda c, r: tf.math.log(gamma_lib.random_gamma( [n], c, r, seed=seed)), lambda c, r: tf.math.log(gamma_lib.random_gamma( [n], c, log_rate=tf.math.log(r), seed=seed)), ] # pylint: enable=g-long-lambda samps = [] dconc = [] drate = [] for fn in lambdas: # Take samples without the nonlinearity. samps.append(fn(conc, rate)) # We compute gradient through a nonlinearity to catch a class of errors. _, (dc_i, dr_i) = tfp.math.value_and_gradient( lambda c, r: tf.reduce_mean(tf.square(fn(c, r))), (conc, rate)) # pylint: disable=cell-var-from-loop dconc.append(dc_i) drate.append(dr_i) # Assert d rate correctness. Note that the non-logspace derivative for rate # depends on the realized sample whereas the logspace one does not. Also, # comparing grads with differently-placed log/exp is numerically perilous. self.assertAllClose(drate[0], drate[1], rtol=0.06) self.assertAllClose(drate[0], drate[2], rtol=0.06) self.assertAllClose(drate[1], drate[3], rtol=0.06) # Assert sample correctness. If incorrect, dconc will be incorrect. 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)), n) equiv_dist = tfb.Log()(tfd.Gamma(conc, rate)) self.evaluate(st.assert_true_cdf_equal_by_dkwm( samps[0], equiv_dist.cdf, false_fail_rate=1e-9)) self.evaluate(st.assert_true_cdf_equal_by_dkwm( samps[1], equiv_dist.cdf, false_fail_rate=1e-9)) self.evaluate(st.assert_true_cdf_equal_by_dkwm( samps[2], equiv_dist.cdf, false_fail_rate=1e-9)) self.evaluate(st.assert_true_cdf_equal_by_dkwm( samps[3], equiv_dist.cdf, false_fail_rate=1e-9)) # Assert d concentration correctness. These are sensitive to sample values, # which are more strongly effected by the log/exp, thus looser tolerances. self.assertAllClose(dconc[0], dconc[1], rtol=0.06) self.assertAllClose(dconc[0], dconc[2], rtol=0.06) self.assertAllClose(dconc[1], dconc[3], rtol=0.06)
def test_dkwm_cdf_one_sample_assertion(self, dtype): rng = np.random.RandomState(seed=0) num_samples = 13000 d = st.min_discrepancy_of_true_cdfs_detectable_by_dkwm( num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6) d = self.evaluate(d) self.assertLess(d, 0.05) # Test that the test assertion agrees that the cdf of the standard # uniform distribution is the identity. samples = rng.uniform(size=num_samples).astype(dtype=dtype) self.evaluate(st.assert_true_cdf_equal_by_dkwm( samples, lambda x: x, false_fail_rate=1e-6)) # Test that the test assertion confirms that the cdf of a # scaled uniform distribution is not the identity. with self.assertRaisesOpError('Empirical CDF outside K-S envelope'): samples = rng.uniform( low=0., high=0.9, size=num_samples).astype(dtype=dtype) self.evaluate(st.assert_true_cdf_equal_by_dkwm( samples, lambda x: x, false_fail_rate=1e-6)) # Test that the test assertion confirms that the cdf of a # shifted uniform distribution is not the identity. with self.assertRaisesOpError('Empirical CDF outside K-S envelope'): samples = rng.uniform( low=0.1, high=1.1, size=num_samples).astype(dtype=dtype) self.evaluate(st.assert_true_cdf_equal_by_dkwm( samples, lambda x: x, false_fail_rate=1e-6))
def propSampleCorrectMarginals(self, dist, special_class, under_hypothesis=False): # Property: When projected on one class, multinomial should sample the # binomial distribution. seed = test_util.test_seed() num_samples = 120000 needed = self.evaluate( st.min_num_samples_for_dkwm_cdf_test(0.02, false_fail_rate=1e-9, false_pass_rate=1e-9)) self.assertGreater(num_samples, needed) samples = dist.sample(num_samples, seed=seed) successes = samples[..., special_class] prob_success = dist._probs_parameter_no_checks()[..., special_class] if under_hypothesis: hp.note('Expected probability of success {}'.format(prob_success)) hp.note('Successes obtained {}'.format(successes)) expected_dist = tfd.Binomial(dist.total_count, probs=prob_success) self.evaluate( st.assert_true_cdf_equal_by_dkwm( successes, expected_dist.cdf, st.left_continuous_cdf_discrete_distribution(expected_dist), false_fail_rate=1e-9))
def testSamplePoissonLowAndHighRates(self): rate = [1., 3., 5., 6., 7., 10., 13.0, 14., 15., 18.] log_rate = np.log(rate) num_samples = int(1e5) 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_samples) samples = self.evaluate( poisson_dist.random_poisson_rejection_sampler( [num_samples, 10], log_rate, seed=test_util.test_seed())) poisson = tfd.Poisson(log_rate=log_rate, validate_args=True) self.evaluate( st.assert_true_cdf_equal_by_dkwm( samples, poisson.cdf, st.left_continuous_cdf_discrete_distribution(poisson), false_fail_rate=1e-9)) self.assertAllClose(self.evaluate(tf.math.reduce_mean(samples, axis=0)), stats.poisson.mean(rate), rtol=0.01) self.assertAllClose(self.evaluate( tf.math.reduce_variance(samples, axis=0)), stats.poisson.var(rate), rtol=0.05)
def testSampleGammaLogSpace(self): concentration = np.linspace(.1, 2., 10) rate = np.linspace(.5, 2, 10) np.random.shuffle(rate) num_samples = int(1e5) 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_samples) samples = gamma_lib.random_gamma([num_samples], concentration, rate, seed=test_util.test_seed(), log_space=True) exp_gamma = tfb.Log()(tfd.Gamma(concentration=concentration, rate=rate, validate_args=True)) self.evaluate( st.assert_true_cdf_equal_by_dkwm(samples, exp_gamma.cdf, false_fail_rate=1e-9)) self.assertAllClose(self.evaluate(tf.math.reduce_mean(samples, axis=0)), tf.math.digamma(concentration) - tf.math.log(rate), rtol=0.02, atol=0.01) self.assertAllClose(self.evaluate( tf.math.reduce_variance(samples, axis=0)), tf.math.polygamma(1., concentration), rtol=0.05)
def test_dkwm_cdf_one_sample_batch_discrete_assertion(self, dtype): rng = np.random.RandomState(seed=0) num_samples = 13000 batch_shape = [3, 2] shape = [num_samples] + batch_shape probs = [0.1, 0.2, 0.3, 0.4] samples = rng.choice(4, size=shape, p=probs).astype(dtype=dtype) def cdf(x): ones = tf.ones_like(x) answer = tf.where(x < 3, 0.6 * ones, ones) answer = tf.where(x < 2, 0.3 * ones, answer) answer = tf.where(x < 1, 0.1 * ones, answer) return tf.where(x < 0, 0 * ones, answer) def left_continuous_cdf(x): ones = tf.ones_like(x) answer = tf.where(x <= 3, 0.6 * ones, ones) answer = tf.where(x <= 2, 0.3 * ones, answer) answer = tf.where(x <= 1, 0.1 * ones, answer) return tf.where(x <= 0, 0 * ones, answer) self.evaluate(st.assert_true_cdf_equal_by_dkwm( samples, cdf, left_continuous_cdf=left_continuous_cdf, false_fail_rate=1e-6)) d = st.min_discrepancy_of_true_cdfs_detectable_by_dkwm( tf.ones(batch_shape) * num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6) self.evaluate(d < 0.05)
def testSamplePoissonHighRates(self): # High rate (>= log(10.)) samples would use rejection sampling. rate = [10., 10.5, 11., 11.5, 12.0, 12.5, 13.0, 13.5, 14.0, 14.5] log_rate = np.log(rate) num_samples = int(1e5) 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_samples) samples = poisson_lib._random_poisson_noncpu( shape=[num_samples], log_rates=log_rate, output_dtype=tf.float64, seed=test_util.test_seed()) poisson = tfd.Poisson(log_rate=log_rate, validate_args=True) self.evaluate( st.assert_true_cdf_equal_by_dkwm( samples, poisson.cdf, st.left_continuous_cdf_discrete_distribution(poisson), false_fail_rate=1e-9)) self.assertAllClose(self.evaluate(tf.math.reduce_mean(samples, axis=0)), stats.poisson.mean(rate), rtol=0.01) self.assertAllClose(self.evaluate( tf.math.reduce_variance(samples, axis=0)), stats.poisson.var(rate), rtol=0.05)
def testSampleGammaHighConcentration(self): concentration = np.linspace(10., 20., 10) rate = np.float64(1.) num_samples = int(1e5) 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_samples) samples = gamma_lib._random_gamma_noncpu(shape=[num_samples, 10], concentration=concentration, rate=rate, seed=test_util.test_seed()) gamma = tfd.Gamma(concentration=concentration, rate=rate, validate_args=True) self.evaluate( st.assert_true_cdf_equal_by_dkwm(samples, gamma.cdf, false_fail_rate=1e-9)) self.assertAllClose(self.evaluate(tf.math.reduce_mean(samples, axis=0)), sp_stats.gamma.mean(concentration, scale=1 / rate), rtol=0.01) self.assertAllClose(self.evaluate( tf.math.reduce_variance(samples, axis=0)), sp_stats.gamma.var(concentration, scale=1 / rate), rtol=0.05)
def assert_univariate_target_conservation(test, mk_target, step_size, stackless): # Sample count limited partly by memory reliably available on Forge. The test # remains reasonable even if the nuts recursion limit is severely curtailed # (e.g., 3 or 4 levels), so use that to recover some memory footprint and bump # the sample count. num_samples = int(5e4) num_steps = 1 target_d = mk_target() strm = tfp.util.SeedStream(salt='univariate_nuts_test', seed=1) # We wrap the initial values in `tf.identity` to avoid broken gradients # resulting from a bijector cache hit, since bijectors of the same # type/parameterization now share a cache. # TODO(b/72831017): Fix broken gradients caused by bijector caching. initialization = tf.identity(target_d.sample([num_samples], seed=strm())) def target(*args): # TODO(axch): Just use target_d.log_prob directly, and accept target_d # itself as an argument instead of a maker function. Blocked by # b/128932888. It would then also be nice not to eta-expand # target_d.log_prob; that was blocked by b/122414321, but maybe tfp's port # of value_and_gradients_function fixed that bug. return mk_target().log_prob(*args) operator = tfp.experimental.mcmc.NoUTurnSampler(target, step_size=step_size, max_tree_depth=3, use_auto_batching=True, stackless=stackless, unrolled_leapfrog_steps=2, seed=strm()) result, extra = tfp.mcmc.sample_chain(num_results=num_steps, num_burnin_steps=0, current_state=initialization, kernel=operator) # Note: sample_chain puts the chain history on top, not the (independent) # chains. test.assertAllEqual([num_steps, num_samples], result.shape) answer = result[0] check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm(answer, target_d.cdf, false_fail_rate=1e-6) check_enough_power = tf1.assert_less( st.min_discrepancy_of_true_cdfs_detectable_by_dkwm( num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.025) test.assertAllEqual([num_samples], extra.leapfrogs_taken[0].shape) unique, _ = tf.unique(extra.leapfrogs_taken[0]) check_leapfrogs_vary = tf1.assert_greater_equal( tf.shape(input=unique)[0], 3) avg_leapfrogs = tf.math.reduce_mean(input_tensor=extra.leapfrogs_taken[0]) check_leapfrogs = tf1.assert_greater_equal( avg_leapfrogs, tf.constant(4, dtype=avg_leapfrogs.dtype)) movement = tf.abs(answer - initialization) test.assertAllEqual([num_samples], movement.shape) # This movement distance (1 * step_size) was selected by reducing until 100 # runs with independent seeds all passed. check_movement = tf1.assert_greater_equal( tf.reduce_mean(input_tensor=movement), 1 * step_size) return (check_cdf_agrees, check_enough_power, check_leapfrogs_vary, check_leapfrogs, check_movement)
def testXLAFriendlySampler(self): if tf.executing_eagerly(): msg = 'XLA requires tf.function, mode switching is meaningless.' self.skipTest(msg) dist = tfd.BetaBinomial(total_count=50, concentration0=1e-7, concentration1=1e-5) seed = test_util.test_seed(sampler_type='stateless') num_samples = 20000 sample = self.evaluate( tf.function(jit_compile=True)(dist.sample)(num_samples, seed=seed)) self.assertAllEqual(np.zeros_like(sample), np.isnan(sample)) # Beta(1e-7, 1e-5) should basically always be either 1 or 0, and 1 should # occur with probability 100/101. # Ergo, the beta binomial samples should basically always be either 50 or 0, # and 50 should occur with probability 100/101. high_samples_mask = sample == 50 low_samples_mask = sample == 0 self.assertAllEqual(np.ones_like(sample), high_samples_mask | low_samples_mask) expect = tfd.Bernoulli(probs=100.0 / 101.0) self.evaluate( st.assert_true_cdf_equal_by_dkwm( samples=tf.cast(high_samples_mask, tf.float32), cdf=expect.cdf, left_continuous_cdf=st. left_continuous_cdf_discrete_distribution(expect), false_fail_rate=1e-9)) self.assertGreater( num_samples, self.evaluate( st.min_num_samples_for_dkwm_cdf_test(0.05, false_fail_rate=1e-9, false_pass_rate=1e-9)))
def testSampleHighConcentration(self): concentration = np.linspace(10., 20., 10) rate = np.float64(1.) num_samples = int(1e5) 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_samples) d = tfd.ExpGamma(concentration=concentration, rate=rate, validate_args=True) samples = d.sample(num_samples, seed=test_util.test_seed()) self.evaluate( st.assert_true_cdf_equal_by_dkwm(samples, d.cdf, false_fail_rate=1e-9)) self.assertAllClose(self.evaluate(tf.math.reduce_mean(samples, axis=0)), d.mean(), rtol=0.01) self.assertAllClose(self.evaluate( tf.math.reduce_variance(samples, axis=0)), d.variance(), rtol=0.05)
def testSamplePoissonLowRates(self): # Low log rate (< log(10.)) samples would use Knuth's algorithm. rate = [1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5] log_rate = np.log(rate) num_samples = int(1e5) 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_samples) samples = self.evaluate( poisson_dist.random_poisson_rejection_sampler( [num_samples, 10], log_rate, seed=test_util.test_seed())) poisson = tfd.Poisson(log_rate=log_rate, validate_args=True) self.evaluate( st.assert_true_cdf_equal_by_dkwm( samples, poisson.cdf, st.left_continuous_cdf_discrete_distribution(poisson), false_fail_rate=1e-9)) self.assertAllClose(self.evaluate(tf.math.reduce_mean(samples, axis=0)), stats.poisson.mean(rate), rtol=0.01) self.assertAllClose(self.evaluate( tf.math.reduce_variance(samples, axis=0)), stats.poisson.var(rate), rtol=0.05)
def check_catches_mistake(wrong_probs): wrong_samples = rng.choice( len(wrong_probs), size=shape, p=wrong_probs).astype(dtype=dtype) with self.assertRaisesOpError('Empirical CDF outside K-S envelope'): self.evaluate(st.assert_true_cdf_equal_by_dkwm( wrong_samples, cdf, left_continuous_cdf=left_continuous_cdf, false_fail_rate=1e-6))
def testSampleEmpiricalCDF(self): num_samples = 300000 dist = tfd.HalfStudentT(df=5., loc=10., scale=2., validate_args=True) samples = dist.sample(num_samples, seed=test_util.test_seed()) check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm( samples, dist.cdf, false_fail_rate=1e-6) check_enough_power = assert_util.assert_less( st.min_discrepancy_of_true_cdfs_detectable_by_dkwm( num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.01) self.evaluate([check_cdf_agrees, check_enough_power])
def testSampleEmpiricalCDF(self): num_samples = 300000 temperature, low, peak, high = 2., 1., 7., 10. dist = tfd.PERT(low, peak, high, temperature, validate_args=True) samples = dist.sample(num_samples, seed=test_util.test_seed()) check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm( samples, dist.cdf, false_fail_rate=1e-6) check_enough_power = assert_util.assert_less( st.min_discrepancy_of_true_cdfs_detectable_by_dkwm( num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.01) self.evaluate([check_cdf_agrees, check_enough_power])
def test_dkwm_cdf_one_sample_batch_discrete_assertion(self, dtype): rng = np.random.RandomState(seed=0) num_samples = 13000 batch_shape = [3, 2] shape = [num_samples] + batch_shape probs = [0.1, 0.2, 0.3, 0.4] samples = rng.choice(4, size=shape, p=probs).astype(dtype=dtype) def cdf(x): ones = tf.ones_like(x) answer = tf1.where(x < 3, 0.6 * ones, ones) answer = tf1.where(x < 2, 0.3 * ones, answer) answer = tf1.where(x < 1, 0.1 * ones, answer) return tf1.where(x < 0, 0 * ones, answer) def left_continuous_cdf(x): ones = tf.ones_like(x) answer = tf1.where(x <= 3, 0.6 * ones, ones) answer = tf1.where(x <= 2, 0.3 * ones, answer) answer = tf1.where(x <= 1, 0.1 * ones, answer) return tf1.where(x <= 0, 0 * ones, answer) self.evaluate( st.assert_true_cdf_equal_by_dkwm( samples, cdf, left_continuous_cdf=left_continuous_cdf, false_fail_rate=1e-6)) d = st.min_discrepancy_of_true_cdfs_detectable_by_dkwm( tf.ones(batch_shape) * num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6) self.assertTrue(np.all(self.evaluate(d) < 0.05)) def check_catches_mistake(wrong_probs): wrong_samples = rng.choice(len(wrong_probs), size=shape, p=wrong_probs).astype(dtype=dtype) with self.assertRaisesOpError( 'Empirical CDF outside K-S envelope'): self.evaluate( st.assert_true_cdf_equal_by_dkwm( wrong_samples, cdf, left_continuous_cdf=left_continuous_cdf, false_fail_rate=1e-6)) check_catches_mistake([0.1, 0.2, 0.3, 0.3, 0.1]) check_catches_mistake([0.2, 0.2, 0.3, 0.3])
def assert_univariate_target_conservation(test, target_d, step_size): # Sample count limited partly by memory reliably available on Forge. The test # remains reasonable even if the nuts recursion limit is severely curtailed # (e.g., 3 or 4 levels), so use that to recover some memory footprint and bump # the sample count. num_samples = int(5e4) num_steps = 1 strm = test_util.test_seed_stream() # We wrap the initial values in `tf.identity` to avoid broken gradients # resulting from a bijector cache hit, since bijectors of the same # type/parameterization now share a cache. # TODO(b/72831017): Fix broken gradients caused by bijector caching. initialization = tf.identity(target_d.sample([num_samples], seed=strm())) @tf.function(autograph=False) def run_chain(): nuts = tfp.experimental.mcmc.PreconditionedNoUTurnSampler( target_d.log_prob, step_size=step_size, max_tree_depth=3, unrolled_leapfrog_steps=2) result = tfp.mcmc.sample_chain(num_results=num_steps, num_burnin_steps=0, current_state=initialization, trace_fn=None, kernel=nuts, seed=strm()) return result result = run_chain() test.assertAllEqual([num_steps, num_samples], result.shape) answer = result[0] check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm(answer, target_d.cdf, false_fail_rate=1e-6) check_enough_power = assert_util.assert_less( st.min_discrepancy_of_true_cdfs_detectable_by_dkwm( num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.025) movement = tf.abs(answer - initialization) test.assertAllEqual([num_samples], movement.shape) # This movement distance (1 * step_size) was selected by reducing until 100 # runs with independent seeds all passed. check_movement = assert_util.assert_greater_equal(tf.reduce_mean(movement), 1 * step_size) return (check_cdf_agrees, check_enough_power, check_movement)
def testSample(self): a = tf.constant(1.0) b = tf.constant(2.0) n = 500000 d = tfd.SigmoidBeta(concentration0=a, concentration1=b, validate_args=True) samples = d.sample(n, seed=test_util.test_seed()) sample_values = self.evaluate(samples) self.assertEqual(samples.shape, (n,)) self.assertEqual(sample_values.shape, (n,)) self.assertTrue(self._kstest(a, b, sample_values)) check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm( samples, d.cdf, false_fail_rate=1e-6) self.evaluate(check_cdf_agrees) check_enough_power = assert_util.assert_less( st.min_discrepancy_of_true_cdfs_detectable_by_dkwm( n, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.01) self.evaluate(check_enough_power)
def assert_univariate_target_conservation(test, target_d, step_size): # Sample count limited partly by memory reliably available on Forge. The test # remains reasonable even if the nuts recursion limit is severely curtailed # (e.g., 3 or 4 levels), so use that to recover some memory footprint and bump # the sample count. num_samples = int(5e4) num_steps = 1 strm = tfp.util.SeedStream(salt='univariate_nuts_test', seed=1) initialization = target_d.sample([num_samples], seed=strm()) @tf.function(autograph=False) def run_chain(): nuts = tfp.mcmc.NoUTurnSampler( target_d.log_prob, step_size=step_size, max_tree_depth=3, unrolled_leapfrog_steps=2, seed=strm()) result, _ = tfp.mcmc.sample_chain( num_results=num_steps, num_burnin_steps=0, current_state=initialization, kernel=nuts) return result result = run_chain() test.assertAllEqual([num_steps, num_samples], result.shape) answer = result[0] check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm( answer, target_d.cdf, false_fail_rate=1e-6) check_enough_power = assert_util.assert_less( st.min_discrepancy_of_true_cdfs_detectable_by_dkwm( num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.025) movement = tf.abs(answer - initialization) test.assertAllEqual([num_samples], movement.shape) # This movement distance (1 * step_size) was selected by reducing until 100 # runs with independent seeds all passed. check_movement = assert_util.assert_greater_equal( tf.reduce_mean(movement), 1 * step_size) return (check_cdf_agrees, check_enough_power, check_movement)
def VerifySampleAndPdfConsistency(self, uniform): """Verifies samples are consistent with the PDF using importance sampling. In particular, we verify an estimate the surface area of the n-dimensional hypersphere, and the surface areas of the spherical caps demarcated by a handful of survival rates. Args: uniform: A `SphericalUniform` distribution instance. """ dim = tf.compat.dimension_value(uniform.event_shape[-1]) nsamples = int(6e4) 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)), nsamples) samples = uniform.sample(sample_shape=[nsamples], seed=test_util.test_seed()) samples = tf.debugging.check_numerics(samples, 'samples') log_prob = uniform.log_prob(samples) log_prob = self.evaluate( tf.debugging.check_numerics(log_prob, 'log_prob')) true_sphere_surface_area = 2 * (np.pi)**(dim / 2) * self.evaluate( tf.exp(-tf.math.lgamma(dim / 2))) true_sphere_surface_area += np.zeros_like(log_prob) # Check the log prob is a constant and is the reciprocal of the surface # area. self.assertAllClose(np.exp(log_prob), 1. / true_sphere_surface_area) # For sampling, let's check the marginals. x_i**2 ~ Beta(0.5, d - 1 / 2) beta_dist = tfp.distributions.Beta(self.dtype(0.5), self.dtype((dim - 1.) / 2.)) for i in range(dim): self.evaluate( st.assert_true_cdf_equal_by_dkwm(samples[..., i]**2, cdf=beta_dist.cdf, false_fail_rate=1e-9))
def testSamplePoissonLowAndHighRates(self): rate = [1., 3., 5., 6., 7., 10., 13.0, 14., 15., 18.] log_rate = np.log(rate) num_samples = int(1e5) poisson = tfd.Poisson(log_rate=log_rate, validate_args=True) 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_samples) samples = poisson_lib._random_poisson_noncpu( shape=[num_samples], log_rates=log_rate, output_dtype=tf.float64, seed=test_util.test_seed()) self.evaluate( st.assert_true_cdf_equal_by_dkwm( samples, poisson.cdf, st.left_continuous_cdf_discrete_distribution(poisson), false_fail_rate=1e-9))
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))