Add baselines for first order mueller

For one and two way shooting, and variational doobs

Author: plainerman

eval/path_metrics.py

@@ -3,12 +3,12 @@
 from tqdm import tqdm
 
 
-def plot_path_energy(paths, U, reduce=jnp.max, already_ln=False):
-    reduced = jnp.array([reduce(U(path)) for path in tqdm(paths, 'Computing path metric')])
+def plot_path_energy(paths, U, reduce=jnp.max, add=0, already_ln=False, **kwargs):
+    reduced = jnp.array([reduce(U(path)) for path in paths]) + add
 
     if already_ln:
         # Convert reduced to log10
         reduced = reduced / jnp.log(10)
-        plt.plot(jnp.arange(0, len(reduced), 1), reduced)
+        plt.plot(jnp.arange(0, len(reduced), 1), reduced, **kwargs)
     else:
-        plt.semilogy(jnp.arange(0, len(reduced), 1), reduced)
+        plt.semilogy(jnp.arange(0, len(reduced), 1), reduced, **kwargs)

evaluate_mueller.py

@@ -0,0 +1,67 @@
+import numpy as np
+import jax.numpy as jnp
+import jax
+from eval.path_metrics import plot_path_energy
+from tps_baseline_mueller import U, dUdx_fn, minima_points
+from scipy.optimize import minimize
+import matplotlib.pyplot as plt
+import os
+
+def load(path):
+    return jnp.array(np.load(path, allow_pickle=True).astype(np.float32)).squeeze()
+
+
+@jax.jit
+def log_prob_path(path):
+    rand = path[1:] - path[:-1] + dt * dUdx_fn(path[:-1])
+    return U(path[0]) + jax.scipy.stats.norm.logpdf(rand, scale=jnp.sqrt(dt) * xi).sum()
+
+
+if __name__ == '__main__':
+    savedir = './out/evaluation/mueller/'
+    os.makedirs(savedir, exist_ok=True)
+
+    all_paths = [
+        ('one-way-shooting', './out/baselines/mueller/paths-one-way-shooting.npy'),
+        ('two-way-shooting', './out/baselines/mueller/paths-two-way-shooting.npy'),
+        ('var-doobs', './out/var_doobs/mueller/paths.npy'),
+    ]
+
+    num_paths = 1000
+    xi = 5
+    dt = 1e-4
+    T = 275e-4
+    N = int(T / dt)
+
+    global_minimum_energy = U(minima_points[0])
+    for point in minima_points:
+        global_minimum_energy = min(global_minimum_energy, minimize(U, point).fun)
+    print("Global minimum energy", global_minimum_energy)
+
+    all_paths = [(name, load(path)) for name, path in all_paths]
+    [print(name, path.shape) for name, path in all_paths]
+
+    for name, paths in all_paths:
+        plot_path_energy(paths, U, add=-global_minimum_energy, label=name)
+
+    plt.legend()
+    plt.ylabel('Maximum energy')
+    plt.savefig(f'{savedir}/mueller-max-energy.pdf', bbox_inches='tight')
+    plt.show()
+
+    for name, paths in all_paths:
+        plot_path_energy(paths, U, add=-global_minimum_energy, reduce=jnp.median, label=name)
+
+    plt.legend()
+    plt.ylabel('Median energy')
+    plt.savefig(f'{savedir}/mueller-median-energy.pdf', bbox_inches='tight')
+    plt.show()
+
+    for name, paths in all_paths:
+        plot_path_energy(paths, log_prob_path, reduce=lambda x: x, label=name)
+        print('Median energy of:', name, jnp.median(jnp.array([log_prob_path(path) for path in paths])))
+
+    plt.legend()
+    plt.ylabel('log path likelihood')
+    plt.savefig(f'{savedir}/mueller-log-path-likelihood.pdf', bbox_inches='tight')
+    plt.show()

mueller.py

@@ -0,0 +1,189 @@
+from tps_baseline_mueller import U, A, B, plot_energy_surface
+from flax import linen as nn
+from flax.training import train_state
+import optax
+import jax
+import jax.numpy as jnp
+from tqdm import trange
+import matplotlib.pyplot as plt
+import os
+import numpy as np
+
+
+class MLPq(nn.Module):
+    @nn.compact
+    def __call__(self, t):
+        t = t / T
+        h = nn.Dense(128)(t - 0.5)
+        h = nn.swish(h)
+        h = nn.Dense(128)(h)
+        h = nn.swish(h)
+        h = nn.Dense(128)(h)
+        h = nn.swish(h)
+        h = nn.Dense(4)(h)
+        mu = (1 - t) * A + t * B + (1 - t) * t * h[:, :2]
+        sigma = (1 - t) * 2.5 * 1e-2 + t * 2.5 * 1e-2 + (1 - t) * t * jnp.exp(h[:, 2:])
+        return mu, sigma
+
+
+if __name__ == '__main__':
+    savedir = f"out/var_doobs/mueller"
+    os.makedirs(savedir, exist_ok=True)
+
+    num_paths = 1000
+    xi = 5
+    dt = 1e-4
+    T = 275e-4
+    N = int(T / dt)
+    epochs = 2_500
+
+    q = MLPq()
+
+    BS = 512
+    key = jax.random.PRNGKey(1)
+    key, *init_key = jax.random.split(key, 3)
+    params_q = q.init(init_key[0], jnp.ones([BS, 1]))
+
+    optimizer_q = optax.adam(learning_rate=1e-4)
+    state_q = train_state.TrainState.create(apply_fn=q.apply,
+                                            params=params_q,
+                                            tx=optimizer_q)
+
+
+    def loss_fn(params_q, key):
+        key = jax.random.split(key)
+        t = T * jax.random.uniform(key[0], [BS, 1])
+        eps = jax.random.normal(key[1], [BS, 2])
+
+        mu_t = lambda _t: state_q.apply_fn(params_q, _t)[0]
+        sigma_t = lambda _t: state_q.apply_fn(params_q, _t)[1]
+
+        def dmudt(_t):
+            _dmudt = jax.jacrev(lambda _t: mu_t(_t).sum(0))
+            return _dmudt(_t).squeeze().T
+
+        def dsigmadt(_t):
+            _dsigmadt = jax.jacrev(lambda _t: sigma_t(_t).sum(0))
+            return _dsigmadt(_t).squeeze().T
+
+        dUdx_fn = jax.grad(lambda _x: U(_x).sum())
+
+        def v_t(_eps, _t):
+            u_t = dmudt(_t) + dsigmadt(_t) * _eps
+            _x = mu_t(_t) + sigma_t(_t) * _eps
+            out = (u_t + dUdx_fn(_x)) - 0.5 * (xi ** 2) * _eps / sigma_t(t)
+            return out
+
+        loss = 0.5 * ((v_t(eps, t) / xi) ** 2).sum(1, keepdims=True)
+        print(loss.shape, 'loss.shape', flush=True)
+        return loss.mean()
+
+
+    @jax.jit
+    def train_step(state_q, key):
+        grad_fn = jax.value_and_grad(loss_fn, argnums=0)
+        loss, grads = grad_fn(state_q.params, key)
+        state_q = state_q.apply_gradients(grads=grads)
+        return state_q, loss
+
+
+    key, loc_key = jax.random.split(key)
+    state_q, loss = train_step(state_q, loc_key)
+
+    loss_plot = []
+    for i in trange(epochs):
+        key, loc_key = jax.random.split(key)
+        state_q, loss = train_step(state_q, loc_key)
+        loss_plot.append(loss)
+
+    plt.plot(loss_plot)
+    plt.show()
+
+    t = T * jnp.linspace(0, 1, BS).reshape((-1, 1))
+    key, path_key = jax.random.split(key)
+    eps = jax.random.normal(path_key, [BS, 2])
+    mu_t, sigma_t = state_q.apply_fn(state_q.params, t)
+    samples = mu_t + sigma_t * eps
+    plot_energy_surface()
+    plt.scatter(samples[:, 0], samples[:, 1])
+    plt.scatter(A[0, 0], A[0, 1], color='red')
+    plt.scatter(B[0, 0], B[0, 1], color='orange')
+    plt.show()
+
+    print("Number of potential evaluations", BS * epochs)
+
+    mu_t = lambda _t: state_q.apply_fn(state_q.params, _t)[0]
+    sigma_t = lambda _t: state_q.apply_fn(state_q.params, _t)[1]
+
+
+    def dmudt(_t):
+        _dmudt = jax.jacrev(lambda _t: mu_t(_t).sum(0), argnums=0)
+        return _dmudt(_t).squeeze().T
+
+
+    def dsigmadt(_t):
+        _dsigmadt = jax.jacrev(lambda _t: sigma_t(_t).sum(0))
+        return _dsigmadt(_t).squeeze().T
+
+
+    u_t = jax.jit(lambda _t, _x: dmudt(_t) + dsigmadt(_t) / sigma_t(_t) * (_x - mu_t(_t)))
+
+    key, loc_key = jax.random.split(key)
+    x_t = jnp.ones((BS, N + 1, 2)) * A
+    eps = jax.random.normal(key, shape=(BS, 2))
+    x_t = x_t.at[:, 0, :].set(x_t[:, 0, :] + sigma_t(jnp.zeros((BS, 1))) * eps)
+    t = jnp.zeros((BS, 1))
+    for i in trange(N):
+        dx = dt * u_t(t, x_t[:, i, :])
+        x_t = x_t.at[:, i + 1, :].set(x_t[:, i, :] + dx)
+        t += dt
+
+    x_t_det = x_t.copy()
+
+    u_t = jax.jit(
+        lambda _t, _x: dmudt(_t) + (dsigmadt(_t) / sigma_t(_t) - 0.5 * (xi / sigma_t(_t)) ** 2) * (_x - mu_t(_t)))
+
+    BS = num_paths
+    key, loc_key = jax.random.split(key)
+    x_t = jnp.ones((BS, N + 1, 2)) * A
+    eps = jax.random.normal(key, shape=(BS, 2))
+    x_t = x_t.at[:, 0, :].set(x_t[:, 0, :] + sigma_t(jnp.zeros((BS, 1))) * eps)
+    t = jnp.zeros((BS, 1))
+    for i in trange(N):
+        key, loc_key = jax.random.split(key)
+        eps = jax.random.normal(key, shape=(BS, 2))
+        dx = dt * u_t(t, x_t[:, i, :]) + jnp.sqrt(dt) * xi * eps
+        x_t = x_t.at[:, i + 1, :].set(x_t[:, i, :] + dx)
+        t += dt
+
+    x_t_stoch = x_t.copy()
+
+    np.save(f'{savedir}/paths.npy', np.array([jnp.array(p) for p in x_t_stoch], dtype=object), allow_pickle=True)
+
+    plt.figure(figsize=(16, 5))
+    plt.subplot(121)
+    plot_energy_surface()
+    plt.plot(x_t_det[:10, :, 0].T, x_t_det[:10, :, 1].T)
+    plt.scatter(A[0, 0], A[0, 1], color='red')
+    plt.scatter(B[0, 0], B[0, 1], color='orange')
+
+    plt.subplot(122)
+    plot_energy_surface()
+    plt.plot(x_t_stoch[:10, :, 0].T, x_t_stoch[:10, :, 1].T)
+    plt.scatter(A[0, 0], A[0, 1], color='red')
+    plt.scatter(B[0, 0], B[0, 1], color='orange')
+    plt.savefig(f'{savedir}/selected_paths_det_vs_stoch.png', bbox_inches='tight')
+    plt.show()
+
+    plt.figure(figsize=(16, 5))
+    plt.subplot(121)
+    plot_energy_surface(trajectories=x_t_det)
+
+    plt.subplot(122)
+    plot_energy_surface(trajectories=x_t_stoch)
+    plt.savefig(f'{savedir}/paths_det_vs_stoch.png', bbox_inches='tight')
+    plt.show()
+
+    plot_energy_surface(trajectories=x_t_stoch)
+    plt.savefig(f'{savedir}/mueller-variational-doobs.pdf', bbox_inches='tight')
+    plt.show()

tps/first_order.py

@@ -101,10 +101,26 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, key, fixed_le
     # pick an initial trajectory
     trajectories = [initial_trajectory]
 
+    statistics = {
+        'num_force_evaluations': 0,
+        'num_tries': 0,
+        'num_metropolis_rejected': 0,
+        'warmup': warmup,
+        'num_paths': num_paths,
+    }
+    if fixed_length > 0:
+        statistics['fixed_length'] = fixed_length
+    else:
+        statistics['max_steps'] = MAX_STEPS
+
     with tqdm(total=num_paths) as pbar:
         while len(trajectories) <= num_paths + warmup:
+            statistics['num_tries'] += 1
+
             key, iter_key, accept_key = jax.random.split(key, 3)
             found, new_trajectory = proposal(system, trajectories[-1], fixed_length, iter_key)
+            statistics['num_force_evaluations'] += len(new_trajectory) - 1
+
             if not found:
                 continue
 
@@ -115,19 +131,30 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, key, fixed_le
 
                 if len(trajectories) > warmup:
                     pbar.update(1)
+            else:
+                statistics['num_metropolis_rejected'] += 1
 
-    return trajectories[warmup + 1:]
+    return trajectories[warmup + 1:], statistics
 
 
 def unguided_md(system, initial_point, num_paths, key, fixed_length=0):
     trajectories = []
     current_frame = initial_point.clone()
     current_trajectory = []
 
+    statistics = {
+        'num_force_evaluations': 0,
+        'num_paths': num_paths,
+        'max_steps': MAX_STEPS,
+    }
+    if fixed_length > 0:
+        statistics['fixed_length'] = fixed_length
+
     with tqdm(total=num_paths) as pbar:
         while len(trajectories) < num_paths:
             key, iter_key = jax.random.split(key)
             next_frame = system.step(current_frame, iter_key)
+            statistics['num_force_evaluations'] += 1
 
             is_transition = not (system.start_state(next_frame) or system.target_state(next_frame))
             if is_transition:
@@ -153,4 +180,4 @@ def unguided_md(system, initial_point, num_paths, key, fixed_length=0):
 
             current_frame = next_frame
 
-    return trajectories
+    return trajectories, statistics