Improve plotting and mueller baselines

Author: plainerman

tps_baseline_mueller.py

@@ -9,7 +9,9 @@
 import numpy as np
 import utils.toy_plot_helpers as toy
 
-minima_points = jnp.array([[-0.55828035, 1.44169], [-0.05004308, 0.46666032], [0.62361133, 0.02804632]])
+minima_points = jnp.array([[-0.55828035, 1.44169],
+                           #[-0.05004308, 0.46666032],
+                           [0.62361133, 0.02804632]])
 A, B = minima_points[None, 0], minima_points[None, 2]
 
 
@@ -46,20 +48,23 @@ def interpolate_two_points(start, stop, steps):
     return interpolation
 
 
-plot_energy_surface = partial(toy.plot_energy_surface, U=U, states=zip(['A', 'B', 'C'], minima_points),
-                              xlim=jnp.array((-1.5, 0.9)), ylim=jnp.array((-0.5, 1.7)))
+plot_energy_surface = partial(toy.plot_energy_surface, U=U, states=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__':
+    variable = True
     savedir = f"out/baselines/mueller"
+    if variable:
+        savedir += "-variable"
+
     os.makedirs(savedir, exist_ok=True)
 
     num_paths = 1000
     xi = 5
     dt = 1e-4
     T = 275e-4
-    N = int(T / dt)
-    initial_trajectory = [t.reshape(1, 2) for t in interpolate(minima_points, 100 if N == 0 else N)]
-
+    N = 0 if variable else int(T / dt)
+    initial_trajectory = [t.reshape(1, 2) for t in interpolate(minima_points, 100 if variable else N)]
 
     @jax.jit
     def step(_x, _key):
@@ -97,3 +102,4 @@ def step(_x, _key):
         plot_energy_surface(trajectories=paths)
         plt.savefig(f'{savedir}/mueller-{name}.pdf', bbox_inches='tight')
         plt.show()
+        plt.clf()

utils/toy_plot_helpers.py

@@ -2,13 +2,13 @@
 import matplotlib.pyplot as plt
 
 
-def plot_energy_surface(U, states, xlim, ylim, points=[], trajectories=[], bins=150, levels=30, alpha=0.7):
+def plot_energy_surface(U, states, xlim, ylim, points=[], trajectories=[], bins=150, levels=30, alpha=0.7, radius=0.1):
     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])
 
     # black and white contour plot
-    plt.contour(x, y, z, levels=levels, cmap='gray')
+    plt.contour(x, y, z, levels=levels, colors='black')
 
     plt.xlim(xlim[0], xlim[1])
     plt.ylim(ylim[0], ylim[1])
@@ -35,12 +35,13 @@ def plot_energy_surface(U, states, xlim, ylim, points=[], trajectories=[], bins=
             rasterized=True
         )
 
-        plt.colorbar()
+    plt.xticks([])
+    plt.yticks([])
 
     for p in points:
         plt.scatter(p[0], p[1], marker='*')
 
     for name, pos in states:
-        c = plt.Circle(pos, radius=0.1, edgecolor='gray', alpha=alpha, facecolor='white', ls='--', lw=0.7)
+        c = plt.Circle(pos, radius=radius, edgecolor='gray', alpha=alpha, facecolor='white', ls='--', lw=0.7, zorder=100)
         plt.gca().add_patch(c)
-        plt.gca().annotate(name, xy=pos, ha="center", va="center")
+        plt.gca().annotate(name, xy=pos, ha="center", va="center", fontsize=14, zorder=101)