Fix state plotting

Author: plainerman

main.py

@@ -1,8 +1,13 @@
 from argparse import ArgumentParser
-
 from utils.args import parse_args
 from systems import System
 import matplotlib.pyplot as plt
+import jax.numpy as jnp
+import jax
+from flax.training import train_state
+import optax
+import model.diagonal as diagonal
+from model.train import train
 
 parser = ArgumentParser()
 parser.add_argument('--out', type=str, default=None, help="Specify a path where the data will be stored.")
@@ -59,18 +64,9 @@
         raise NotImplementedError
         # system = System.from_forcefield(args.start, args.target)
 
-    import jax.numpy as jnp
-    import jax
-    from tqdm import trange
-    from flax.training import train_state
-    import optax
-    import model.diagonal as diagonal
-    from model.train import train
-    from model import MLPq
-
-    N = int(args.T / args.dt)
-
     # You can play around with any model here
+    # The chosen setup will append a final layer so that the output is mu, sigma, and weights
+    from model import MLPq
     model = MLPq([128, 128, 128])
 
     # TODO: parameterize mixtures, weights, and base_sigma
@@ -113,3 +109,12 @@
 
     key, path_key = jax.random.split(key)
     x_t_stoch = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, args.xi, path_key)
+
+    if system.plot:
+        system.plot(title='Deterministic Paths', trajectories=x_t_det)
+        plt.show()
+        plt.clf()
+
+        system.plot(title='Stochastic Paths', trajectories=x_t_stoch)
+        plt.show()
+        plt.clf()

model/qsetup.py

@@ -44,6 +44,7 @@ def sample_paths(self, state_q: TrainState, x_0: ArrayLike, dt: float, T: float,
 
         for i in trange(N):
             for j in range(0, num_paths, BS):
+                # If the BS does not divide the number of paths, we need to pad the last batch
                 if j + BS > num_paths:
                     j_end = num_paths
                     cur_x_t = jnp.pad(x_t[j:, i], pad_width=((0, BS - (num_paths - j)), (0, 0)))
@@ -56,6 +57,7 @@ def sample_paths(self, state_q: TrainState, x_0: ArrayLike, dt: float, T: float,
                 if key is None:
                     noise = 0
                 else:
+                    # For stochastic sampling we compute the noise
                     key, iter_key = jax.random.split(key)
                     noise = xi * jax.random.normal(iter_key, shape=(BS, ndim))
 

systems.py

@@ -37,7 +37,7 @@ def from_name(cls, name: str) -> Self:
         else:
             raise ValueError(f"Unknown system: {name}")
 
-        plot = partial(toy.plot_energy_surface, U=U, states=zip(['A', 'B'], [A, B]), xlim=xlim, ylim=ylim, alpha=1.0)
+        plot = partial(toy.plot_energy_surface, U=U, states=list(zip(['A', 'B'], [A, B])), xlim=xlim, ylim=ylim, alpha=1.0)
         mass = jnp.array([1.0, 1.0])
         return cls(U, A, B, mass, plot)
 

tps_baseline_mueller.py

@@ -48,7 +48,7 @@ def interpolate_two_points(start, stop, steps):
     return interpolation
 
 
-plot_energy_surface = partial(toy.plot_energy_surface, U=U, states=zip(['A', 'B'], minima_points),
+plot_energy_surface = partial(toy.plot_energy_surface, U=U, states=list(zip(['A', 'B'], minima_points)),
                               xlim=jnp.array((-1.5, 0.9)), ylim=jnp.array((-0.5, 1.7)), alpha=1.0)
 
 if __name__ == '__main__':

utils/toy_plot_helpers.py

@@ -2,7 +2,11 @@
 import matplotlib.pyplot as plt
 
 
-def plot_energy_surface(U, states, xlim, ylim, points=[], trajectories=[], bins=150, levels=30, alpha=0.7, radius=0.1):
+def plot_energy_surface(U, states, xlim, ylim, points=[], trajectories=[],
+                        title=None, bins=150, levels=30, alpha=0.7, radius=0.1):
+    if title:
+        plt.title(title)
+
     x, y = jnp.linspace(xlim[0], xlim[1], bins), jnp.linspace(ylim[0], ylim[1], bins)
     x, y = jnp.meshgrid(x, y, indexing='ij')
     z = U(jnp.stack([x, y], -1).reshape(-1, 2)).reshape([bins, bins])
@@ -42,7 +46,7 @@ def plot_energy_surface(U, states, xlim, ylim, points=[], trajectories=[], bins=
         plt.scatter(p[0], p[1], marker='*')
 
     for name, pos in states:
-        pos = pos.reshape(2,)
+        pos = pos.reshape(2, )
         c = plt.Circle(pos, radius=radius, edgecolor='gray', alpha=alpha, facecolor='white', ls='--', lw=0.7, zorder=10)
         plt.gca().add_patch(c)
         plt.gca().annotate(name, xy=pos, ha="center", va="center", fontsize=14, zorder=11)