def test_atypical_nesting(self): results = FakeAtypicalNestingResults( unique_atypical_nesting=0, atypical_inner_results=FakeResults(1)) self.assertTrue(unnest.has_nested(results, 'unique_atypical_nesting')) self.assertTrue(unnest.has_nested(results, 'unique_core')) self.assertFalse(hasattr(results, 'unique_core')) self.assertFalse(unnest.has_nested(results, 'foo')) self.assertEqual( results.unique_atypical_nesting, unnest.get_innermost(results, 'unique_atypical_nesting')) self.assertEqual( results.unique_atypical_nesting, unnest.get_outermost(results, 'unique_atypical_nesting')) self.assertEqual(results.atypical_inner_results.unique_core, unnest.get_innermost(results, 'unique_core')) self.assertEqual(results.atypical_inner_results.unique_core, unnest.get_outermost(results, 'unique_core')) self.assertRaises(AttributeError, lambda: unnest.get_innermost(results, 'foo')) self.assertRaises(AttributeError, lambda: unnest.get_outermost(results, 'foo')) self.assertIs(unnest.get_innermost(results, 'foo', SINGLETON), SINGLETON) self.assertIs(unnest.get_outermost(results, 'foo', SINGLETON), SINGLETON)
def _process_results(self): """Process outputs to extract useful data.""" if unnest.has_nested(self.kernel, 'reducer'): reducers = unnest.get_outermost(self.kernel, 'reducer') # Finalize streaming calculations. self.reductions = nest.map_structure_up_to( reducers, lambda r, s: r.finalize(s), reducers, unnest.get_outermost(self.results, 'reduction_results'), check_types=False) # Grab useful reductions. def process_reductions(reducer, reduction): if isinstance(reducer, tracing_reducer.TracingReducer): self.all_states, self.trace = reduction nest.map_structure_up_to(reducers, process_reductions, reducers, self.reductions, check_types=False) if unnest.has_nested(self.results, 'new_step_size'): self.new_step_size = unnest.get_outermost(self.results, 'new_step_size')
def _slow_window(*, kind, proposal_kernel_kwargs, dual_averaging_kwargs, num_draws, initial_position, initial_running_variance, bijector, trace_fn, seed): """Sample using both step size and mass matrix adaptation.""" dual_averaging_kwargs = dict(dual_averaging_kwargs) dual_averaging_kwargs.setdefault('num_adaptation_steps', num_draws) kernel = make_slow_adapt_kernel( kind=kind, proposal_kernel_kwargs=proposal_kernel_kwargs, dual_averaging_kwargs=dual_averaging_kwargs, initial_running_variance=initial_running_variance) with warnings.catch_warnings(): warnings.simplefilter('ignore') draws, trace, fkr = sample.sample_chain( num_draws, initial_position, kernel=kernel, return_final_kernel_results=True, # pylint: disable=g-long-lambda trace_fn=lambda state, pkr: trace_fn(state, bijector, tf.constant(True), pkr.inner_results.inner_results), seed=seed) # pylint: enable=g-long-lambda draw_and_chain_axes = [0, 1] weighted_running_variance = [] for d in list(draws): prev_mean, prev_var = tf.nn.moments(d[-num_draws // 2:], axes=draw_and_chain_axes) num_samples = tf.cast( num_draws / 2, dtype=dtype_util.common_dtype([prev_mean, prev_var], tf.float32)) weighted_running_variance.append(sample_stats.RunningVariance.from_stats( num_samples=num_samples, mean=prev_mean, variance=prev_var)) step_size = unnest.get_outermost(fkr, 'step_size') momentum_distribution = unnest.get_outermost(fkr, 'momentum_distribution') return draws, trace, step_size, weighted_running_variance, momentum_distribution
def test_flat(self): results = FakeResults(0) self.assertTrue(unnest.has_nested(results, 'unique_core')) self.assertFalse(unnest.has_nested(results, 'foo')) self.assertEqual(results.unique_core, unnest.get_innermost(results, 'unique_core')) self.assertEqual(results.unique_core, unnest.get_outermost(results, 'unique_core')) self.assertRaises( AttributeError, lambda: unnest.get_innermost(results, 'foo')) self.assertRaises( AttributeError, lambda: unnest.get_outermost(results, 'foo')) self.assertIs(unnest.get_innermost(results, 'foo', SINGLETON), SINGLETON) self.assertIs(unnest.get_outermost(results, 'foo', SINGLETON), SINGLETON)
def realized_acceptance_rate(self): """Return realized acceptance rate of the samples.""" try: is_accepted = unnest.get_outermost(self.trace, 'is_accepted') except AttributeError: return return tf.math.reduce_mean(tf.cast(is_accepted, tf.float32), axis=0)
def _fast_window(*, kind, proposal_kernel_kwargs, dual_averaging_kwargs, num_draws, initial_position, bijector, trace_fn, seed): """Sample using just step size adaptation.""" dual_averaging_kwargs.update({'num_adaptation_steps': num_draws}) kernel = make_fast_adapt_kernel( kind=kind, proposal_kernel_kwargs=proposal_kernel_kwargs, dual_averaging_kwargs=dual_averaging_kwargs) with warnings.catch_warnings(): warnings.simplefilter('ignore') draws, trace, fkr = sample.sample_chain( num_draws, initial_position, kernel=kernel, return_final_kernel_results=True, # pylint: disable=g-long-lambda trace_fn=lambda state, pkr: trace_fn( state, bijector, tf.constant(True), pkr.inner_results), seed=seed) # pylint: enable=g-long-lambda draw_and_chain_axes = [0, 1] prev_mean, prev_var = tf.nn.moments(draws[-num_draws // 2:], axes=draw_and_chain_axes) num_samples = tf.cast(num_draws / 2, dtype=dtype_util.common_dtype([prev_mean, prev_var], tf.float32)) weighted_running_variance = sample_stats.RunningVariance.from_stats( num_samples=num_samples, mean=prev_mean, variance=prev_var) step_size = unnest.get_outermost(fkr, 'step_size') return draws, trace, step_size, weighted_running_variance
def _fast_adapt_window( num_draws, joint_log_prob_fn, initial_position, hmc_kernel_kwargs, dual_averaging_kwargs, event_kernel_kwargs, trace_fn=None, seed=None, ): """ In the fast adaptation window, we use the `DualAveragingStepSizeAdaptation` kernel to wrap an HMC kernel. :param num_draws: Number of MCMC draws in window :param joint_log_prob_fn: joint log posterior function :param initial_position: initial state of the Markov chain :param hmc_kernel_kwargs: `HamiltonianMonteCarlo` kernel keywords args :param dual_averaging_kwargs: `DualAveragingStepSizeAdaptation` keyword args :param event_kernel_kwargs: EventTimesMH and Occult kernel args :param trace_fn: function to trace kernel results :param seed: optional random seed. :returns: draws, kernel results, the adapted HMC step size, and variance accumulator """ kernel_list = [ ( 0, make_hmc_fast_adapt_kernel( hmc_kernel_kwargs=hmc_kernel_kwargs, dual_averaging_kwargs=dual_averaging_kwargs, ), ), (1, make_event_multiscan_gibbs_step(**event_kernel_kwargs)), ] kernel = GibbsKernel( target_log_prob_fn=joint_log_prob_fn, kernel_list=kernel_list, name="fast_adapt", ) draws, trace, fkr = tfp.mcmc.sample_chain( num_draws, initial_position, kernel=kernel, return_final_kernel_results=True, trace_fn=trace_fn, seed=seed, ) weighted_running_variance = get_weighted_running_variance(draws[0]) step_size = unnest.get_outermost(fkr.inner_results[0], "step_size") return draws, trace, step_size, weighted_running_variance
def test_deeply_nested(self): results = _build_deeply_nested(0, 1, 2, 3, 4) self.assertTrue(unnest.has_nested(results, 'unique_nesting')) self.assertTrue(unnest.has_nested(results, 'unique_atypical_nesting')) self.assertTrue(unnest.has_nested(results, 'unique_core')) self.assertFalse(hasattr(self, 'unique_core')) self.assertFalse(unnest.has_nested(results, 'foo')) self.assertEqual(unnest.get_innermost(results, 'unique_nesting'), 3) self.assertEqual(unnest.get_outermost(results, 'unique_nesting'), 0) self.assertEqual( unnest.get_innermost(results, 'unique_atypical_nesting'), 2) self.assertEqual( unnest.get_outermost(results, 'unique_atypical_nesting'), 1) self.assertEqual(unnest.get_innermost(results, 'unique_core'), 4) self.assertEqual(unnest.get_outermost(results, 'unique_core'), 4) self.assertRaises( AttributeError, lambda: unnest.get_innermost(results, 'foo')) self.assertRaises( AttributeError, lambda: unnest.get_outermost(results, 'foo')) self.assertIs(unnest.get_innermost(results, 'foo', SINGLETON), SINGLETON) self.assertIs(unnest.get_outermost(results, 'foo', SINGLETON), SINGLETON)
def _slow_window(*, target_log_prob_fn, num_leapfrog_steps, num_draws, initial_position, initial_running_variance, initial_step_size, target_accept_prob, bijector, trace_fn, seed): """Sample using both step size and mass matrix adaptation.""" kernel = make_slow_adapt_kernel( target_log_prob_fn=target_log_prob_fn, initial_running_variance=initial_running_variance, initial_step_size=initial_step_size, num_leapfrog_steps=num_leapfrog_steps, num_adaptation_steps=num_draws, target_accept_prob=target_accept_prob) with warnings.catch_warnings(): warnings.simplefilter('ignore') draws, trace, fkr = sample.sample_chain( num_draws, initial_position, kernel=kernel, return_final_kernel_results=True, # pylint: disable=g-long-lambda trace_fn=lambda state, pkr: trace_fn( state, bijector, tf.constant(True), pkr.inner_results. inner_results), seed=seed) # pylint: enable=g-long-lambda draw_and_chain_axes = [0, 1] prev_mean, prev_var = tf.nn.moments(draws[-num_draws // 2:], axes=draw_and_chain_axes) num_samples = tf.cast(num_draws / 2, dtype=dtype_util.common_dtype([prev_mean, prev_var], tf.float32)) weighted_running_variance = sample_stats.RunningVariance.from_stats( num_samples=num_samples, mean=prev_mean, variance=prev_var) step_size = unnest.get_outermost(fkr, 'step_size') momentum_distribution = unnest.get_outermost(fkr, 'momentum_distribution') return draws, trace, step_size, weighted_running_variance, momentum_distribution
def trace_results_fn(_, results): """Packs results into a dictionary""" results_dict = {} root_results = results.inner_results step_size = tf.convert_to_tensor( unnest.get_outermost(root_results[0], "step_size") ) results_dict["hmc"] = { "is_accepted": unnest.get_innermost(root_results[0], "is_accepted"), "target_log_prob": unnest.get_innermost( root_results[0], "target_log_prob" ), "step_size": step_size, } def get_move_results(results): return { "is_accepted": results.is_accepted, "target_log_prob": results.accepted_results.target_log_prob, "proposed_delta": tf.stack( [ results.accepted_results.m, results.accepted_results.t, results.accepted_results.delta_t, results.accepted_results.x_star, ] ), } res1 = root_results[1].inner_results results_dict["move/S->E"] = get_move_results(res1[0]) results_dict["move/E->I"] = get_move_results(res1[1]) results_dict["occult/S->E"] = get_move_results(res1[2]) results_dict["occult/E->I"] = get_move_results(res1[3]) return results_dict
def _slow_adapt_window( num_draws, joint_log_prob_fn, initial_position, initial_running_variance, hmc_kernel_kwargs, dual_averaging_kwargs, event_kernel_kwargs, trace_fn=None, seed=None, ): """In the slow adaptation phase, we adapt the HMC step size and mass matrix together. :param num_draws: number of MCMC iterations :param joint_log_prob_fn: the joint posterior density function :param initial_position: initial Markov chain state :param initial_running_variance: initial variance accumulator :param hmc_kernel_kwargs: `HamiltonianMonteCarlo` kernel kwargs :param dual_averaging_kwargs: `DualAveragingStepSizeAdaptation` kwargs :param event_kernel_kwargs: EventTimesMH and Occults kwargs :param trace_fn: result trace function :param seed: optional random seed :returns: draws, kernel results, adapted step size, the variance accumulator, and "learned" momentum distribution for the HMC. """ kernel_list = [ ( 0, make_hmc_slow_adapt_kernel( initial_running_variance, hmc_kernel_kwargs, dual_averaging_kwargs, ), ), (1, make_event_multiscan_gibbs_step(**event_kernel_kwargs)), ] kernel = GibbsKernel( target_log_prob_fn=joint_log_prob_fn, kernel_list=kernel_list, name="slow_adapt", ) draws, trace, fkr = tfp.mcmc.sample_chain( num_draws, current_state=initial_position, kernel=kernel, return_final_kernel_results=True, trace_fn=trace_fn, ) step_size = unnest.get_outermost(fkr.inner_results[0], "step_size") momentum_distribution = unnest.get_outermost( fkr.inner_results[0], "momentum_distribution" ) weighted_running_variance = get_weighted_running_variance(draws[0]) return ( draws, trace, step_size, weighted_running_variance, momentum_distribution, )