manifold mala example done, need to regenerate the data on the same GPU

Author: AdrienCorenflos

examples/manifold_mala.py

@@ -19,23 +19,24 @@
 from coupled_rejection_sampling.mvn import coupled_mvns, mvn_logpdf
 from coupled_rejection_sampling.thorisson import modified_thorisson
 
+
 JAX_KEY = jax.random.PRNGKey(0)
-B = 1_000  # number of parallel coupled chains
-K = 50  # number of experiments
-CS = np.linspace(0.8, 0.99, 5)
-NS = [16, 32, 64]
+K = 10_000  # number of experiments
+CS = np.linspace(0.8, 0.99, 7)
+NS = [4, 8, 16, 32, 64, 128, 256]
 
 ALPHA = 100.  # Same wide prior as in mMALA paper
-EPS = 0.05
+EPS = 1.
 
-RUN = True
+RUN = False
 PLOT = True
 
 data_path = "data/heart.csv"
 save_path = "out/mmala.npz"
 
 
 def thorisson_sample(C):
+    @jax.jit
     def sampler(key, x_mean, x_chol, y_mean, y_chol):
         def p_xy(k, mean, scale):
             return mean + scale @ jax.random.normal(k, (mean.shape[0],))
@@ -52,25 +53,26 @@ def p_xy(k, mean, scale):
 
 
 def rejection_sample(N):
+    @jax.jit
     def sampler(key, x_mean, x_chol, y_mean, y_chol):
         X, Y, coupled, _ = coupled_mvns(key, x_mean, x_chol, y_mean, y_chol, N)
         return X, Y, coupled
 
     return sampler
 
 
+@partial(jax.jit, static_argnums=(2,))
 def _get_manifold_langevin_discretisation(theta, eps, log_pi):
     """Gets the mean and cholesky covariance of the proposal. This only considers flat manifolds"""
     d = theta.shape[0]
 
     # Compute gradient and Fisher information
-    value, grad = jax.value_and_grad(log_pi)(theta)
+    grad = jax.grad(log_pi)(theta)
     fisher = -jax.hessian(log_pi)(theta)
-
+    # id_print(fisher, what="fisher autograd")
     # Get sqrt of inverse Fisher
     chol_fisher = jlinalg.cholesky(fisher, lower=True)
     inv_chol_fisher = jlinalg.solve_triangular(chol_fisher, jnp.eye(d), lower=True)
-
     # Get mean
     mean = 0.5 * eps ** 2 * jlinalg.cho_solve((chol_fisher, True), grad)
 
@@ -101,16 +103,15 @@ def simplified_manifold_mala_step(key, x, y, eps, sampler, log_pi):
 
     x = jax.lax.select(accept_x, x_star, x)
     y = jax.lax.select(accept_y, y_star, y)
-
     return x, y, accept_x & accept_y & coupled
 
 
 @partial(jax.jit, static_argnums=(2, 3, 4))
 def sample_coupled_chain(key, eps, sampler, log_pi, D):
     key, init_x_key, init_y_key = jax.random.split(key, 3)
 
-    x0 = jax.random.normal(init_x_key, (D,))
-    y0 = jax.random.normal(init_y_key, (D,))
+    x0 = 0.1 * jax.random.normal(init_x_key, (D,))
+    y0 = 0.1 * jax.random.normal(init_y_key, (D,))
 
     def cond(carry):
         return ~carry[-1]
@@ -120,7 +121,6 @@ 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
 
@@ -135,59 +135,52 @@ def experiment():
 
     X = np.pad(X, [(0, 0), (0, 1)], constant_values=1)
 
-    y = df.values[:, -1].astype(bool)
+    y = df.values[:, -1].astype(float)
 
-    def log_target(theta):
+    @partial(jax.jit, static_argnames=("return_fisher",))
+    def log_target(theta, return_fisher=False):
         prior = norm.logpdf(theta, 0., ALPHA ** 0.5).sum()
+        temp = X @ theta
+        log_lik = jnp.dot(y, temp) - jnp.sum(jax.nn.softplus(temp))
 
-        temp = -X @ theta
-        log_probs_true = -jnp.logaddexp(0, temp)
-        log_probs_false = temp + log_probs_true
-        log_lik = jnp.where(y, log_probs_true, log_probs_false)
-        return prior + jnp.nansum(log_lik)
+        if not return_fisher:
+            return log_lik + prior
 
-    rejection_meeting_times_res = np.empty((len(NS), K, B))
-    thorisson_meeting_times_res = np.empty((len(CS), K, B))
+    rejection_meeting_times_res = np.empty((len(NS), K))
+    thorisson_meeting_times_res = np.empty((len(CS), K))
 
     rejection_runtime_res = np.empty((len(NS), K))
     thorisson_runtime_res = np.empty((len(CS), K))
 
-    # Rejection:
     rej_key = JAX_KEY
-    for i, N in enumerate(tqdm.tqdm(NS, leave=False)):
+    for i, N in enumerate(tqdm.tqdm(NS, leave=True)):
         rej_sampler = rejection_sample(N)
-        rej_experiment_fun = jax.jit(
-            jax.vmap(lambda op_key: sample_coupled_chain(op_key, EPS, rej_sampler, log_target, dim)))
+        rej_experiment_fun = jax.jit(lambda op_key: sample_coupled_chain(op_key, EPS, rej_sampler, log_target, dim))
 
         # run it once to compile
-        batched_keys = jax.random.split(JAX_KEY, B)
-        block = rej_experiment_fun(batched_keys)
+        block = rej_experiment_fun(rej_key)
         block.block_until_ready()
-        for k in enumerate(tqdm.trange(K, leave=False)):
+        for k in range(K):
             tic = time.time()
             rej_key, subkey = jax.random.split(rej_key)
-            batched_keys = jax.random.split(subkey, B)
 
-            rej_out = rej_experiment_fun(batched_keys)
+            rej_out = rej_experiment_fun(rej_key)
             rej_out.block_until_ready()
             rejection_meeting_times_res[i, k] = rej_out
             rejection_runtime_res[i, k] = time.time() - tic
 
     thor_key = JAX_KEY
-    for j, C in enumerate(tqdm.tqdm(CS, leave=False)):
+    for j, C in enumerate(tqdm.tqdm(CS, leave=True)):
         thor_sampler = thorisson_sample(C)
-        thor_experiment_fun = jax.jit(
-            jax.vmap(lambda op_key: sample_coupled_chain(op_key, EPS, thor_sampler, log_target, dim)))
+        thor_experiment_fun = jax.jit(lambda op_key: sample_coupled_chain(op_key, EPS, thor_sampler, log_target, dim))
 
         # run it once to compile
-        batched_keys = jax.random.split(JAX_KEY, B)
-        block = thor_experiment_fun(batched_keys)
+        block = thor_experiment_fun(thor_key)
         block.block_until_ready()
-        for k in tqdm.trange(K):
+        for k in range(K):
             tic = time.time()
             thor_key, subkey = jax.random.split(thor_key)
-            batched_keys = jax.random.split(subkey, B)
-            thor_out = thor_experiment_fun(batched_keys)
+            thor_out = thor_experiment_fun(thor_key)
             thor_out.block_until_ready()
             thorisson_meeting_times_res[j, k] = thor_out
             thorisson_runtime_res[j, k] = time.time() - tic
@@ -205,60 +198,78 @@ def log_target(theta):
     cmap = plt.get_cmap("tab10")
     data = np.load(save_path)
 
-    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_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")
+    index = pd.MultiIndex.from_product([["meeting time", "run time (s)"], ["mean", "standard deviation"]])
+    thorisson_df = pd.DataFrame(columns=data["CS"], index=index)
+    rejection_df = pd.DataFrame(columns=data["NS"], index=index)
+
+    thorisson_df.loc[("meeting time", "mean")] = [f"${v.mean(-1):.1f}$" for v in data["thorisson_meeting_times"]]
+    thorisson_df.loc[("run time (s)", "mean")] = [f"${v.mean(-1):.1e}$" for v in data["thorisson_runtime"][:, 1:]]
+
+    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:]]
+
+    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")