Fix second order path sampling shapes

Author: plainerman

configs/toy/mueller_second_single_gaussian.yaml

@@ -0,0 +1,21 @@
+save_dir: ./out/toy/mueller_second_order_single_gaussian
+
+test_system: mueller_brown
+ode: second_order
+parameterization: diagonal
+T: 20.0
+xi: 1.4142135
+xi_pos_noise: 1e-3
+gamma: 1.0
+
+num_gaussians: 1
+trainable_weights: False
+base_sigma: 5e-3
+
+epochs: 10000
+BS: 512
+
+num_paths: 1000
+dt: 0.02
+
+log_plots: True

main.py

@@ -30,6 +30,8 @@
 parser.add_argument('--T', type=float, required=True,
                     help="Transition time in the base unit of the system. For molecular simulations, this is in picoseconds.")
 parser.add_argument('--xi', type=float)
+parser.add_argument('--xi_pos_noise', type=float, default=1e-4,
+                    help="For second order SDEs we have to add a small noise to the positional xi. This is the value of this noise.")
 parser.add_argument('--temperature', type=float,
                     help="The temperature of the system in Kelvin. Either specify this or xi.")
 parser.add_argument('--gamma', type=float, required=True)
@@ -98,7 +100,7 @@ def main():
     from model import MLP
 
     model = MLP([128, 128, 128])
-    setup = qsetup.construct(system, model, args.ode, args.parameterization, xi, args)
+    setup, A, B = qsetup.construct(system, model, xi, args)
 
     key = jax.random.PRNGKey(args.seed)
     key, init_key = jax.random.split(key)
@@ -144,7 +146,7 @@ def main():
     print("!!!TODO: how to plot this nicely?")
     t = args.T * jnp.linspace(0, 1, args.BS, dtype=jnp.float32).reshape((-1, 1))
     key, path_key = jax.random.split(key)
-    eps = jax.random.normal(path_key, [args.BS, args.num_gaussians, setup.A.shape[-1]])
+    eps = jax.random.normal(path_key, [args.BS, args.num_gaussians, A.shape[-1]])
     mu_t, sigma_t, w_logits = state_q.apply_fn(state_q.params, t)
     w = jax.nn.softmax(w_logits)[None, :, None]
     samples = (w * (mu_t + sigma_t * eps)).sum(axis=1)
@@ -156,7 +158,7 @@ def main():
     # plt.show()
 
     key, init_key = jax.random.split(key)
-    x_0 = jnp.ones((args.num_paths, setup.A.shape[0]), dtype=jnp.float32) * setup.A
+    x_0 = jnp.ones((args.num_paths, A.shape[0]), dtype=jnp.float32) * A
     eps = jax.random.normal(key, shape=x_0.shape)
     x_0 += args.base_sigma * eps
 

training/qsetup.py

@@ -1,7 +1,7 @@
 import argparse
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Callable, Optional
+from typing import Callable, Optional, Tuple
 from flax import linen as nn
 from flax.training.train_state import TrainState
 
@@ -81,32 +81,37 @@ def B(self):
         return self.system.B
 
 
-def construct(system: System, model: nn.module, ode: str, parameterization: str, xi: ArrayLike,
-              args: argparse.Namespace) -> QSetup:
+def construct(system: System, model: nn.module, xi: float, args: argparse.Namespace) -> Tuple[
+    QSetup, ArrayLike, ArrayLike]:
+    """
+    Construct a QSetup object based on the given arguments.
+    return: QSetup, A, B
+    """
     from training.setups import diagonal
 
-    if ode == 'first_order':
+    if args.ode == 'first_order':
         order = 'first'
         A = system.A
         B = system.B
-    elif ode == 'second_order':
+    elif args.ode == 'second_order':
         order = 'second'
 
         # We pad the A and B matrices with zeros to account for the velocity
-        A = jnp.hstack([system.A, jnp.zeros_like(system.A)])
-        B = jnp.hstack([system.B, jnp.zeros_like(system.B)])
+        A = jnp.hstack([system.A, jnp.zeros_like(system.A)], dtype=jnp.float32)
+        B = jnp.hstack([system.B, jnp.zeros_like(system.B)], dtype=jnp.float32)
 
         xi_velocity = jnp.ones_like(system.A) * xi
-        xi_pos = jnp.zeros_like(xi_velocity) + 1e-4
+        xi_pos = jnp.zeros_like(xi_velocity) + args.xi_pos_noise
 
-        xi = jnp.concatenate((xi_pos, xi_velocity), axis=-1)
+        xi = jnp.concatenate((xi_pos, xi_velocity), axis=-1, dtype=jnp.float32)
+        print("Setting xi to", xi)
     else:
-        raise ValueError(f"Unknown ODE: {ode}")
+        raise ValueError(f"Unknown ODE: {args.ode}")
 
-    if parameterization == 'diagonal':
+    if args.parameterization == 'diagonal':
         wrapped_module = diagonal.DiagonalWrapper(
             model, args.T, A, B, args.num_gaussians, args.trainable_weights, args.base_sigma
         )
-        return diagonal.DiagonalSetup(system, wrapped_module, xi, order, args.T)
+        return diagonal.DiagonalSetup(system, wrapped_module, xi, order, args.T), A, B
     else:
-        raise ValueError(f"Unknown parameterization: {parameterization}")
+        raise ValueError(f"Unknown parameterization: {args.parameterization}")

training/setups/drift.py

@@ -12,10 +12,7 @@ class DriftedSetup(QSetup, ABC):
     def __init__(self, system: System, model_q: nn.Module, xi: ArrayLike, order: str):
         """Either instantiate with first or second order drift."""
         assert order == 'first' or order == 'second', "Order must be either 'first' or 'second'."
-
         self.order = order
-        self._A = system.A
-        self._B = system.B
 
         super().__init__(system, model_q, xi)
 
@@ -27,11 +24,3 @@ def _drift(self, _x: ArrayLike, gamma: float) -> ArrayLike:
             ndim = self.system.A.shape[0]
 
             return jnp.hstack([_x[:, ndim:] / self.system.mass, -self.system.dUdx(_x[:, :ndim]) - _x[:, ndim:] * gamma])
-
-    @property
-    def A(self):
-        return self._A
-
-    @property
-    def B(self):
-        return self._B