Remove find_reasonable_step_size (#553)

Author: fehiepsi

numpyro/infer/hmc_util.py

@@ -328,7 +328,7 @@ def _identity_step_size(inverse_mass_matrix, z, rng_key, step_size):
     return step_size
 
 
-def warmup_adapter(num_adapt_steps, find_reasonable_step_size=_identity_step_size,
+def warmup_adapter(num_adapt_steps, find_reasonable_step_size=None,
                    adapt_step_size=True, adapt_mass_matrix=True,
                    dense_mass=False, target_accept_prob=0.8):
     """
@@ -349,6 +349,8 @@ def warmup_adapter(num_adapt_steps, find_reasonable_step_size=_identity_step_siz
         step size, hence the sampling will be slower but more robust. Default to 0.8.
     :return: a pair of (`init_fn`, `update_fn`).
     """
+    if find_reasonable_step_size is None:
+        find_reasonable_step_size = _identity_step_size
     ss_init, ss_update = dual_averaging()
     mm_init, mm_update, mm_final = welford_covariance(diagonal=not dense_mass)
     adaptation_schedule = np.array(build_adaptation_schedule(num_adapt_steps))

numpyro/infer/mcmc.py

@@ -192,6 +192,7 @@ def init_kernel(init_params,
                     target_accept_prob=0.8,
                     trajectory_length=2*math.pi,
                     max_tree_depth=10,
+                    find_heuristic_step_size=False,
                     model_args=(),
                     model_kwargs=None,
                     rng_key=PRNGKey(0)):
@@ -221,6 +222,8 @@ def init_kernel(init_params,
             value is :math:`2\\pi`.
         :param int max_tree_depth: Max depth of the binary tree created during the doubling
             scheme of NUTS sampler. Defaults to 10.
+        :param bool find_heuristic_step_size: whether to a heuristic function to adjust the
+            step size at the beginning of each adaptation window. Defaults to False.
         :param tuple model_args: Model arguments if `potential_fn_gen` is specified.
         :param dict model_kwargs: Model keyword arguments if `potential_fn_gen` is specified.
         :param jax.random.PRNGKey rng_key: random key to be used as the source of
@@ -243,8 +246,12 @@ def init_kernel(init_params,
                 kwargs = {} if model_kwargs is None else model_kwargs
                 pe_fn = potential_fn_gen(*model_args, **kwargs)
 
-        find_reasonable_ss = partial(find_reasonable_step_size,
-                                     pe_fn, kinetic_fn, momentum_generator)
+        find_reasonable_ss = None
+        if find_heuristic_step_size:
+            find_reasonable_ss = partial(find_reasonable_step_size,
+                                         pe_fn,
+                                         kinetic_fn,
+                                         momentum_generator)
 
         wa_init, wa_update = warmup_adapter(num_warmup,
                                             adapt_step_size=adapt_step_size,
@@ -437,6 +444,8 @@ class HMC(MCMCKernel):
         value is :math:`2\\pi`.
     :param callable init_strategy: a per-site initialization function.
         See :ref:`init_strategy` section for available functions.
+    :param bool find_heuristic_step_size: whether to a heuristic function to adjust the
+        step size at the beginning of each adaptation window. Defaults to False.
     """
     def __init__(self,
                  model=None,
@@ -448,7 +457,8 @@ def __init__(self,
                  dense_mass=False,
                  target_accept_prob=0.8,
                  trajectory_length=2 * math.pi,
-                 init_strategy=init_to_uniform()):
+                 init_strategy=init_to_uniform(),
+                 find_heuristic_step_size=False):
         if not (model is None) ^ (potential_fn is None):
             raise ValueError('Only one of `model` or `potential_fn` must be specified.')
         self._model = model
@@ -463,6 +473,7 @@ def __init__(self,
         self._algo = 'HMC'
         self._max_tree_depth = 10
         self._init_strategy = init_strategy
+        self._find_heuristic_step_size = find_heuristic_step_size
         # Set on first call to init
         self._init_fn = None
         self._postprocess_fn = None
@@ -525,9 +536,10 @@ def init(self, rng_key, num_warmup, init_params=None, model_args=(), model_kwarg
             target_accept_prob=self._target_accept_prob,
             trajectory_length=self._trajectory_length,
             max_tree_depth=self._max_tree_depth,
-            rng_key=rng_key,
+            find_heuristic_step_size=self._find_heuristic_step_size,
             model_args=model_args,
             model_kwargs=model_kwargs,
+            rng_key=rng_key,
         )
         if rng_key.ndim == 1:
             init_state = hmc_init_fn(init_params, rng_key)
@@ -600,6 +612,8 @@ class NUTS(HMC):
         scheme of NUTS sampler. Defaults to 10.
     :param callable init_strategy: a per-site initialization function.
         See :ref:`init_strategy` section for available functions.
+    :param bool find_heuristic_step_size: whether to a heuristic function to adjust the
+        step size at the beginning of each adaptation window. Defaults to False.
     """
     def __init__(self,
                  model=None,
@@ -612,12 +626,15 @@ def __init__(self,
                  target_accept_prob=0.8,
                  trajectory_length=None,
                  max_tree_depth=10,
-                 init_strategy=init_to_uniform()):
+                 init_strategy=init_to_uniform(),
+                 find_heuristic_step_size=False):
         super(NUTS, self).__init__(potential_fn=potential_fn, model=model, kinetic_fn=kinetic_fn,
                                    step_size=step_size, adapt_step_size=adapt_step_size,
                                    adapt_mass_matrix=adapt_mass_matrix, dense_mass=dense_mass,
                                    target_accept_prob=target_accept_prob,
-                                   trajectory_length=trajectory_length, init_strategy=init_strategy)
+                                   trajectory_length=trajectory_length,
+                                   init_strategy=init_strategy,
+                                   find_heuristic_step_size=find_heuristic_step_size)
         self._max_tree_depth = max_tree_depth
         self._algo = 'NUTS'
 

test/test_mcmc.py

@@ -87,7 +87,7 @@ def model(labels):
     if kernel_cls is SA:
         kernel = SA(model=model, adapt_state_size=9)
     else:
-        kernel = kernel_cls(model=model, trajectory_length=8)
+        kernel = kernel_cls(model=model, trajectory_length=8, find_heuristic_step_size=True)
     mcmc = MCMC(kernel, warmup_steps, num_samples, progress_bar=False)
     mcmc.run(random.PRNGKey(2), labels)
     mcmc.print_summary()