Add double_well to toys

Author: plainerman

configs/toy/double_well.yaml

@@ -0,0 +1,18 @@
+save_dir: ./out/toy/double_well
+
+test_system: double_well
+ode: first_order
+parameterization: diagonal
+T: 1.0
+xi: 1.0
+gamma: 1.0
+
+num_gaussians: 1
+trainable_weights: False
+base_sigma: 0.1
+
+epochs: 2500
+BS: 512
+
+num_paths: 100
+dt: 5e-4

configs/toy/double_well_hard.yaml

@@ -0,0 +1,18 @@
+save_dir: ./out/toy/double_well_hard
+
+test_system: double_well_hard
+ode: first_order
+parameterization: diagonal
+T: 1.0
+xi: 1.0
+gamma: 1.0
+
+num_gaussians: 1
+trainable_weights: False
+base_sigma: 0.1
+
+epochs: 2500
+BS: 512
+
+num_paths: 100
+dt: 5e-4

plot_test.py

@@ -33,7 +33,7 @@ def U_mueller_brown(xs, beta=1.0):
     return beta * (e1 + e2 + e3 + e4)
 
 
-double_well = (U_double_well,)
-double_well_hard = (U_double_well_hard,)
+double_well = (U_double_well, jnp.array([-jnp.sqrt(2), 0]), jnp.array([jnp.sqrt(2), 0]))
+double_well_hard = (U_double_well_hard, jnp.array([-3, 0]), jnp.array([3, 0]))
 double_well_dual_channel = (U_double_well_dual_channel, jnp.array([-0.5, 0]), jnp.array([0.5, 0]))
 mueller_brown = (U_mueller_brown, jnp.array([-0.55828035, 1.44169]), jnp.array([0.62361133, 0.02804632]))

potentials.py

@@ -34,8 +34,12 @@ def __init__(self, U: Callable[[ArrayLike], ArrayLike], A: ArrayLike, B: ArrayLi
     def from_name(cls, name: str, force_clip: float) -> Self:
         if name == 'double_well':
             U, A, B = potentials.double_well
+            xlim = jnp.array((-2.0, 2.0))
+            ylim = jnp.array((-2.5, 2.5))
         elif name == 'double_well_hard':
             U, A, B = potentials.double_well_hard
+            xlim = jnp.array((-6.0, 6.0))
+            ylim = jnp.array((-4.0, 4.0))
         elif name == 'double_well_dual_channel':
             U, A, B = potentials.double_well_dual_channel
             xlim = jnp.array((-1.0, 1.0))

systems.py

@@ -83,6 +83,7 @@ def step(_x, _key):
             with open(f'{save_dir}/stats-{name}.json', 'r') as fp:
                 statistics = json.load(fp)
         else:
+            print('Generating paths for', name)
             paths, statistics = tps1.mcmc_shooting(tps_config, method, initial_trajectory, num_paths,
                                                    jax.random.PRNGKey(1), warmup=0, fixed_length=N)