Started the reflected max example

Author: AdrienCorenflos

coupled_rejection_sampling/reflected_coupled_mh.py

@@ -0,0 +1,54 @@
+""" This algorithm implements Algorithm 4 of Maximal Couplings of the Metropolis-Hastings
+Algorithm to compare with our method in the manifold MALA case. It is however coded in a fairly generic manner, so I
+felt like putting it in the main code.
+"""
+import jax
+import jax.numpy as jnp
+
+from coupled_rejection_sampling.utils import logsubexp
+
+
+def reflection_coupled_mh(kernel, kernel_log_density):
+    def r_xy_fn(x, y, x_prime):
+        log_p_x_zp = kernel_log_density(x_prime, x)
+        log_p_y_zp = kernel_log_density(x_prime, y)
+        return logsubexp(log_p_x_zp, log_p_y_zp)
+
+    def tr_xy_fn(x, y, x_prime):
+        y_refl = _transport_fn(x, y, x_prime)
+        r_xy_xp = r_xy_fn(x, y, x_prime)
+        r_yx_yp = r_xy_fn(y, y, y_refl)
+        return logsubexp(r_xy_xp, r_yx_yp)
+
+
+    def _residual(x, y):
+        def cond(carry):
+            return ~carry[-1]
+
+        def body(carry):
+            op_key, y_prop, _ = carry
+            next_key, subkey1, subkey2 = jax.random.split(op_key)
+            y_prop, accepted_y = kernel(subkey2, y)
+
+            def if_accepted(_):
+                log_u = jnp.log(jax.random.uniform(subkey1))
+                log_p_y_yp = kernel_log_density(y_prop, y)
+                log_p_res_y_yp = tr_xy_fn(y, x, y_prop)
+                accept = log_u < log_p_res_y_yp - log_p_y_yp
+                return y_prop, accept
+
+            return jax.lax.cond(accepted_y, lambda _: (y_prop, False), if_accepted, None)
+
+    def step(key, x, y):
+        subkey1, subkey2, subkey3 = key
+        x_prime, accepted_x = kernel(subkey1, x)
+
+
+
+def _transport_fn(x, y, x_prime):
+    r_curr = jnp.linalg.norm(y - x)
+    e = jax.lax.select(r_curr < 1e-8, jnp.zeros_like(x), (y - x) / r_curr)
+    x_diff = x_prime - x
+    eta = x_diff - 2 * e * e.dot(x_diff)
+    y = x + eta
+    return y

examples/manifold_mala.py

@@ -10,16 +10,13 @@
 import jax.scipy.linalg as jlinalg
 import numpy as np
 import pandas as pd
-import tikzplotlib
 import tqdm.auto as tqdm
 from jax.scipy.stats import norm
 from matplotlib import pyplot as plt
-from matplotlib.ticker import FormatStrFormatter
 
 from coupled_rejection_sampling.mvn import coupled_mvns, mvn_logpdf
 from coupled_rejection_sampling.thorisson import modified_thorisson
 
-
 JAX_KEY = jax.random.PRNGKey(0)
 K = 10_000  # number of experiments
 CS = np.linspace(0.8, 0.99, 7)
@@ -121,6 +118,7 @@ def body(carry):
         op_key, sample_key = jax.random.split(op_key, 2)
         x, y, coupled = simplified_manifold_mala_step(sample_key, x, y, eps, sampler, log_pi)
         return op_key, x, y, iteration + 1, coupled
+
     *_, meeting_time, _ = jax.lax.while_loop(cond, body, (key, x0, y0, 0, False))
     return meeting_time
 
@@ -208,68 +206,15 @@ def log_target(theta, return_fisher=False):
     rejection_df.loc[("meeting time", "mean")] = [f"${v.mean(-1):.1f}$" for v in data["rejection_meeting_times"]]
     rejection_df.loc[("run time (s)", "mean")] = [f"${v.mean(-1):.1e}$" for v in data["rejection_runtime"][:, 1:]]
 
+    thorisson_df.loc[("meeting time", "standard deviation")] = [f"${v.std(-1):.1f}$" for v in
+                                                                data["thorisson_meeting_times"]]
+    thorisson_df.loc[("run time (s)", "standard deviation")] = [f"${v.std(-1):.1e}$" for v in
+                                                                data["thorisson_runtime"][:, 1:]]
 
-    thorisson_df.loc[("meeting time", "standard deviation")] = [f"${v.std(-1):.1f}$" for v in data["thorisson_meeting_times"]]
-    thorisson_df.loc[("run time (s)", "standard deviation")] = [f"${v.std(-1):.1e}$" for v in data["thorisson_runtime"][:, 1:]]
-
-    rejection_df.loc[("meeting time", "standard deviation")] = [f"${v.std(-1):.1f}$" for v in data["rejection_meeting_times"]]
-    rejection_df.loc[("run time (s)", "standard deviation")] = [f"${v.std(-1):.1e}$" for v in data["rejection_runtime"][:, 1:]]
-
+    rejection_df.loc[("meeting time", "standard deviation")] = [f"${v.std(-1):.1f}$" for v in
+                                                                data["rejection_meeting_times"]]
+    rejection_df.loc[("run time (s)", "standard deviation")] = [f"${v.std(-1):.1e}$" for v in
+                                                                data["rejection_runtime"][:, 1:]]
 
     print(rejection_df.to_latex("out/rejection_mmala.tex"))
     print(thorisson_df.to_latex("out/thorisson_mmala.tex"))
-    # for i, D in enumerate(DS):
-    #     axes[0].set_title("Rejection")
-    #     axes[0].set_xscale("log")
-    #     axes[0].set_yscale("log")
-    #     axes[0].plot(data["NS"], data["rejection_meeting_times"][i].mean(-1).mean(-1),
-    #                  color=cmap(i), label=f"$D={D}$")
-    #     axes[0].fill_between(data["NS"],
-    #                          data["rejection_meeting_times"][i].mean(-1).mean(-1) - 1.96 *
-    #                          data["rejection_meeting_times"][i].mean(-1).std(-1),
-    #                          data["rejection_meeting_times"][i].mean(-1).mean(-1) + 1.96 *
-    #                          data["rejection_meeting_times"][i].mean(-1).std(-1),
-    #                          color=cmap(i), alpha=0.66)
-    #
-    #     axes[1].set_title("Thorisson")
-    #     axes[1].plot(data["CS"], data["thorisson_meeting_times"][i].mean(-1).mean(-1),
-    #                  color=cmap(i), label=f"$D={D}$")
-    #     axes[1].fill_between(data["CS"],
-    #                          data["thorisson_meeting_times"][i].mean(-1).mean(-1) - 1.96 *
-    #                          data["thorisson_meeting_times"][i].mean(-1).std(-1),
-    #                          data["thorisson_meeting_times"][i].mean(-1).mean(-1) + 1.96 *
-    #                          data["thorisson_meeting_times"][i].mean(-1).std(-1),
-    #                          color=cmap(i), alpha=0.66)
-    #     axes[1].set_yscale("log")
-    #     axes[1].xaxis.set_major_formatter(FormatStrFormatter('%.1f'))
-    # axes[1].legend()
-    # tikzplotlib.save("out/gibbs_meeting_time.tikz")
-    #
-    # fig, axes = plt.subplots(ncols=2, figsize=(15, 6), sharey=True)
-    #
-    # for i, D in enumerate(DS):
-    #     axes[0].set_title("Rejection")
-    #     axes[0].set_xscale("log")
-    #     axes[0].set_yscale("log")
-    #     axes[0].plot(data["NS"], data["rejection_runtime"][i].mean(-1),
-    #                  color=cmap(i), label=f"$D={D}$")
-    #     axes[0].fill_between(data["NS"],
-    #                          data["rejection_runtime"][i].mean(-1) - 1.96 *
-    #                          data["rejection_runtime"][i].std(-1),
-    #                          data["rejection_runtime"][i].mean(-1) + 1.96 *
-    #                          data["rejection_runtime"][i].std(-1),
-    #                          color=cmap(i), alpha=0.66)
-    #     axes[1].set_title("Thorisson")
-    #     axes[1].plot(data["CS"], data["thorisson_runtime"][i].mean(-1),
-    #                  color=cmap(i), label=f"$D={D}$")
-    #     axes[1].fill_between(data["CS"],
-    #                          data["thorisson_runtime"][i].mean(-1) - 1.96 *
-    #                          data["thorisson_runtime"][i].std(-1),
-    #                          data["thorisson_runtime"][i].mean(-1) + 1.96 *
-    #                          data["thorisson_runtime"][i].std(-1),
-    #                          color=cmap(i), alpha=0.66)
-    #     axes[1].set_yscale("log")
-    #     axes[1].xaxis.set_major_formatter(FormatStrFormatter('%.1f'))
-    # axes[1].legend()
-    #
-    # tikzplotlib.save("out/gibbs_run_time.tikz")