add static adjusted mclmc

blackjax/adaptation/adjusted_mclmc_adaptation.py

@@ -74,7 +74,7 @@ def adjusted_mclmc_find_L_and_step_size(
     dim = pytree_size(state.position)
     if params is None:
         params = MCLMCAdaptationState(
-            jnp.sqrt(dim), jnp.sqrt(dim) * 0.2, sqrt_diag_cov=jnp.ones((dim,))
+            jnp.sqrt(dim), jnp.sqrt(dim) * 0.2, inverse_mass_matrix=jnp.ones((dim,))
         )
 
     part1_key, part2_key = jax.random.split(rng_key, 2)
@@ -152,7 +152,7 @@ def step(iteration_state, weight_and_key):
                 state=previous_state,
                 avg_num_integration_steps=avg_num_integration_steps,
                 step_size=params.step_size,
-                sqrt_diag_cov=params.sqrt_diag_cov,
+                inverse_mass_matrix=params.inverse_mass_matrix,
             )
 
             # step updating
@@ -283,9 +283,7 @@ def L_step_size_adaptation(state, params, num_steps, rng_key):
                 L=params.L * change, step_size=params.step_size * change
             )
             if diagonal_preconditioning:
-                params = params._replace(
-                    sqrt_diag_cov=jnp.sqrt(variances), L=jnp.sqrt(dim)
-                )
+                params = params._replace(inverse_mass_matrix=variances, L=jnp.sqrt(dim))
 
             initial_da, update_da, final_da = dual_averaging_adaptation(target=target)
             (
@@ -323,7 +321,7 @@ def step(state, key):
                 state=state,
                 step_size=params.step_size,
                 avg_num_integration_steps=params.L / params.step_size,
-                sqrt_diag_cov=params.sqrt_diag_cov,
+                inverse_mass_matrix=params.inverse_mass_matrix,
             )
             return next_state, next_state.position
 

blackjax/adaptation/mclmc_adaptation.py

@@ -30,13 +30,13 @@ class MCLMCAdaptationState(NamedTuple):
         The momentum decoherent rate for the MCLMC algorithm.
     step_size
         The step size used for the MCLMC algorithm.
-    sqrt_diag_cov
+    inverse_mass_matrix
         A matrix used for preconditioning.
     """
 
     L: float
     step_size: float
-    sqrt_diag_cov: float
+    inverse_mass_matrix: float
 
 
 def mclmc_find_L_and_step_size(
@@ -87,10 +87,10 @@ def mclmc_find_L_and_step_size(
     Example
     -------
     .. code::
-        kernel = lambda sqrt_diag_cov : blackjax.mcmc.mclmc.build_kernel(
+        kernel = lambda inverse_mass_matrix : blackjax.mcmc.mclmc.build_kernel(
         logdensity_fn=logdensity_fn,
         integrator=integrator,
-        sqrt_diag_cov=sqrt_diag_cov,
+        inverse_mass_matrix=inverse_mass_matrix,
         )
 
         (
@@ -106,7 +106,7 @@ def mclmc_find_L_and_step_size(
     """
     dim = pytree_size(state.position)
     params = MCLMCAdaptationState(
-        jnp.sqrt(dim), jnp.sqrt(dim) * 0.25, sqrt_diag_cov=jnp.ones((dim,))
+        jnp.sqrt(dim), jnp.sqrt(dim) * 0.25, inverse_mass_matrix=jnp.ones((dim,))
     )
     part1_key, part2_key = jax.random.split(rng_key, 2)
 
@@ -123,7 +123,7 @@ def mclmc_find_L_and_step_size(
 
     if frac_tune3 != 0:
         state, params = make_adaptation_L(
-            mclmc_kernel(params.sqrt_diag_cov), frac=frac_tune3, Lfactor=0.4
+            mclmc_kernel(params.inverse_mass_matrix), frac=frac_tune3, Lfactor=0.4
         )(state, params, num_steps, part2_key)
 
     return state, params
@@ -152,7 +152,7 @@ def predictor(previous_state, params, adaptive_state, rng_key):
         rng_key, nan_key = jax.random.split(rng_key)
 
         # dynamics
-        next_state, info = kernel(params.sqrt_diag_cov)(
+        next_state, info = kernel(params.inverse_mass_matrix)(
             rng_key=rng_key,
             state=previous_state,
             L=params.L,
@@ -247,15 +247,15 @@ def L_step_size_adaptation(state, params, num_steps, rng_key):
 
         L = params.L
         # determine L
-        sqrt_diag_cov = params.sqrt_diag_cov
+        inverse_mass_matrix = params.inverse_mass_matrix
         if num_steps2 > 1:
             x_average, x_squared_average = average[0], average[1]
             variances = x_squared_average - jnp.square(x_average)
             L = jnp.sqrt(jnp.sum(variances))
 
             if diagonal_preconditioning:
-                sqrt_diag_cov = jnp.sqrt(variances)
-                params = params._replace(sqrt_diag_cov=sqrt_diag_cov)
+                inverse_mass_matrix = variances
+                params = params._replace(inverse_mass_matrix=inverse_mass_matrix)
                 L = jnp.sqrt(dim)
 
                 # readjust the stepsize
@@ -265,7 +265,7 @@ def L_step_size_adaptation(state, params, num_steps, rng_key):
                     xs=(jnp.ones(steps), keys), state=state, params=params
                 )
 
-        return state, MCLMCAdaptationState(L, params.step_size, sqrt_diag_cov)
+        return state, MCLMCAdaptationState(L, params.step_size, inverse_mass_matrix)
 
     return L_step_size_adaptation
 

blackjax/mcmc/adjusted_mclmc.py

@@ -40,7 +40,7 @@ def build_kernel(
     num_integration_steps: int,
     integrator: Callable = integrators.isokinetic_mclachlan,
     divergence_threshold: float = 1000,
-    sqrt_diag_cov=1.0,
+    inverse_mass_matrix=1.0,
 ):
     """Build an MHMCHMC kernel where the number of integration steps is chosen randomly.
 
@@ -76,7 +76,9 @@ def kernel(
         momentum = generate_unit_vector(key_momentum, state.position)
         proposal, info, _ = adjusted_mclmc_proposal(
             integrator=integrators.with_isokinetic_maruyama(
-                integrator(logdensity_fn=logdensity_fn, sqrt_diag_cov=sqrt_diag_cov)
+                integrator(
+                    logdensity_fn=logdensity_fn, inverse_mass_matrix=inverse_mass_matrix
+                )
             ),
             step_size=step_size,
             L_proposal_factor=L_proposal_factor * (num_integration_steps * step_size),
@@ -105,7 +107,7 @@ def as_top_level_api(
     logdensity_fn: Callable,
     step_size: float,
     L_proposal_factor: float = jnp.inf,
-    sqrt_diag_cov=1.0,
+    inverse_mass_matrix=1.0,
     *,
     divergence_threshold: int = 1000,
     integrator: Callable = integrators.isokinetic_mclachlan,
@@ -140,7 +142,7 @@ def as_top_level_api(
     kernel = build_kernel(
         num_integration_steps,
         integrator=integrator,
-        sqrt_diag_cov=sqrt_diag_cov,
+        inverse_mass_matrix=inverse_mass_matrix,
         divergence_threshold=divergence_threshold,
     )
 

blackjax/mcmc/adjusted_mclmc_dynamic.py

@@ -38,7 +38,7 @@ def build_kernel(
     integrator: Callable = integrators.isokinetic_mclachlan,
     divergence_threshold: float = 1000,
     next_random_arg_fn: Callable = lambda key: jax.random.split(key)[1],
-    sqrt_diag_cov=1.0,
+    inverse_mass_matrix=1.0,
 ):
     """Build a Dynamic MHMCHMC kernel where the number of integration steps is chosen randomly.
 
@@ -76,7 +76,9 @@ def kernel(
         momentum = generate_unit_vector(key_momentum, state.position)
         proposal, info, _ = adjusted_mclmc_proposal(
             integrator=integrators.with_isokinetic_maruyama(
-                integrator(logdensity_fn=logdensity_fn, sqrt_diag_cov=sqrt_diag_cov)
+                integrator(
+                    logdensity_fn=logdensity_fn, inverse_mass_matrix=inverse_mass_matrix
+                )
             ),
             step_size=step_size,
             L_proposal_factor=L_proposal_factor * (num_integration_steps * step_size),
@@ -106,7 +108,7 @@ def as_top_level_api(
     logdensity_fn: Callable,
     step_size: float,
     L_proposal_factor: float = jnp.inf,
-    sqrt_diag_cov=1.0,
+    inverse_mass_matrix=1.0,
     *,
     divergence_threshold: int = 1000,
     integrator: Callable = integrators.isokinetic_mclachlan,
@@ -143,7 +145,7 @@ def as_top_level_api(
         integration_steps_fn=integration_steps_fn,
         integrator=integrator,
         next_random_arg_fn=next_random_arg_fn,
-        sqrt_diag_cov=sqrt_diag_cov,
+        inverse_mass_matrix=inverse_mass_matrix,
         divergence_threshold=divergence_threshold,
     )
 

blackjax/mcmc/integrators.py

@@ -311,7 +311,9 @@ def _normalized_flatten_array(x, tol=1e-13):
     return jnp.where(norm > tol, x / norm, x), norm
 
 
-def esh_dynamics_momentum_update_one_step(sqrt_diag_cov=1.0):
+def esh_dynamics_momentum_update_one_step(inverse_mass_matrix=1.0):
+    sqrt_inverse_mass_matrix = jax.tree_util.tree_map(jnp.sqrt, inverse_mass_matrix)
+
     def update(
         momentum: ArrayTree,
         logdensity_grad: ArrayTree,
@@ -330,7 +332,7 @@ def update(
 
         logdensity_grad = logdensity_grad
         flatten_grads, unravel_fn = ravel_pytree(logdensity_grad)
-        flatten_grads = flatten_grads * sqrt_diag_cov
+        flatten_grads = flatten_grads * sqrt_inverse_mass_matrix
         flatten_momentum, _ = ravel_pytree(momentum)
         dims = flatten_momentum.shape[0]
         normalized_gradient, gradient_norm = _normalized_flatten_array(flatten_grads)
@@ -342,7 +344,7 @@ def update(
             + 2 * zeta * flatten_momentum
         )
         new_momentum_normalized, _ = _normalized_flatten_array(new_momentum_raw)
-        gr = unravel_fn(new_momentum_normalized * sqrt_diag_cov)
+        gr = unravel_fn(new_momentum_normalized * sqrt_inverse_mass_matrix)
         next_momentum = unravel_fn(new_momentum_normalized)
         kinetic_energy_change = (
             delta
@@ -374,11 +376,11 @@ def format_isokinetic_state_output(
 
 def generate_isokinetic_integrator(coefficients):
     def isokinetic_integrator(
-        logdensity_fn: Callable, sqrt_diag_cov: ArrayTree = 1.0
+        logdensity_fn: Callable, inverse_mass_matrix: ArrayTree = 1.0
     ) -> GeneralIntegrator:
         position_update_fn = euclidean_position_update_fn(logdensity_fn)
         one_step = generalized_two_stage_integrator(
-            esh_dynamics_momentum_update_one_step(sqrt_diag_cov),
+            esh_dynamics_momentum_update_one_step(inverse_mass_matrix),
             position_update_fn,
             coefficients,
             format_output_fn=format_isokinetic_state_output,

blackjax/mcmc/mclmc.py

@@ -60,7 +60,7 @@ def init(position: ArrayLike, logdensity_fn, rng_key):
     )
 
 
-def build_kernel(logdensity_fn, sqrt_diag_cov, integrator):
+def build_kernel(logdensity_fn, inverse_mass_matrix, integrator):
     """Build a HMC kernel.
 
     Parameters
@@ -81,7 +81,7 @@ def build_kernel(logdensity_fn, sqrt_diag_cov, integrator):
     """
 
     step = with_isokinetic_maruyama(
-        integrator(logdensity_fn=logdensity_fn, sqrt_diag_cov=sqrt_diag_cov)
+        integrator(logdensity_fn=logdensity_fn, inverse_mass_matrix=inverse_mass_matrix)
     )
 
     def kernel(
@@ -107,7 +107,7 @@ def as_top_level_api(
     L,
     step_size,
     integrator=isokinetic_mclachlan,
-    sqrt_diag_cov=1.0,
+    inverse_mass_matrix=1.0,
 ) -> SamplingAlgorithm:
     """The general mclmc kernel builder (:meth:`blackjax.mcmc.mclmc.build_kernel`, alias `blackjax.mclmc.build_kernel`) can be
     cumbersome to manipulate. Since most users only need to specify the kernel
@@ -155,7 +155,7 @@ def as_top_level_api(
     A ``SamplingAlgorithm``.
     """
 
-    kernel = build_kernel(logdensity_fn, sqrt_diag_cov, integrator)
+    kernel = build_kernel(logdensity_fn, inverse_mass_matrix, integrator)
 
     def init_fn(position: ArrayLike, rng_key: PRNGKey):
         return init(position, logdensity_fn, rng_key)

tests/mcmc/test_integrators.py

@@ -238,7 +238,7 @@ def test_esh_momentum_update(self, dims):
 
         # Efficient implementation
         update_stable = self.variant(
-            esh_dynamics_momentum_update_one_step(sqrt_diag_cov=1.0)
+            esh_dynamics_momentum_update_one_step(inverse_mass_matrix=1.0)
         )
         next_momentum1, *_ = update_stable(momentum, gradient, step_size, 1.0)
         np.testing.assert_array_almost_equal(next_momentum, next_momentum1)
@@ -263,7 +263,7 @@ def test_isokinetic_velocity_verlet(self):
         next_state, kinetic_energy_change = step(initial_state, step_size)
 
         # explicit integration
-        op1 = esh_dynamics_momentum_update_one_step(sqrt_diag_cov=1.0)
+        op1 = esh_dynamics_momentum_update_one_step(inverse_mass_matrix=1.0)
         op2 = integrators.euclidean_position_update_fn(logdensity_fn)
         position, momentum, _, logdensity_grad = initial_state
         momentum, kinetic_grad, kinetic_energy_change0 = op1(