Fix xi for second order

Author: plainerman

main.py

@@ -56,7 +56,8 @@
 
 def main():
     # TODO: force clipping
-    # TODO: temperature
+    print("!!!!Next todos: plot ALDP")
+
     args = parse_args(parser)
     assert args.test_system or args.start and args.target, "Either specify a test system or provide start and target structures"
     assert not (
@@ -86,24 +87,26 @@ def main():
     from model import MLP
 
     model = MLP([128, 128, 128])
-    setup = qsetup.construct(system, model, args.ode, args.parameterization, args)
+    setup = qsetup.construct(system, model, args.ode, args.parameterization, xi, args)
 
     key = jax.random.PRNGKey(args.seed)
     key, init_key = jax.random.split(key)
     params_q = setup.model_q.init(init_key, jnp.zeros([args.BS, 1], dtype=jnp.float32))
 
     optimizer_q = optax.adam(learning_rate=args.lr)
     state_q = train_state.TrainState.create(apply_fn=setup.model_q.apply, params=params_q, tx=optimizer_q)
-    loss_fn = setup.construct_loss(state_q, xi, args.gamma, args.BS)
+    loss_fn = setup.construct_loss(state_q, args.gamma, args.BS)
 
     key, train_key = jax.random.split(key)
     state_q, loss_plot = train(state_q, loss_fn, args.epochs, train_key)
     print("Number of potential evaluations", args.BS * args.epochs)
 
+    if jnp.isnan(jnp.array(loss_plot)).any():
+        print("Warning: Loss contains NaNs")
     plt.plot(loss_plot)
     show_or_save_fig(args.save_dir, 'loss_plot.pdf')
 
-    # TODO: how to plot this nicely?
+    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]])
@@ -122,15 +125,15 @@ def main():
     eps = jax.random.normal(key, shape=x_0.shape)
     x_0 += args.base_sigma * eps
 
-    x_t_det = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, None, None)
+    x_t_det = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, None)
 
     if system.plot:
         # In case we have a second order integration scheme, we remove the velocity for plotting
         system.plot(title='Deterministic Paths', trajectories=x_t_det[:, :, :system.A.shape[0]])
         show_or_save_fig(args.save_dir, 'paths_deterministic.pdf')
 
     key, path_key = jax.random.split(key)
-    x_t_stoch = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, xi, path_key)
+    x_t_stoch = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, path_key)
 
     if system.plot:
         system.plot(title='Stochastic Paths', trajectories=x_t_stoch[:, :, :system.A.shape[0]])

training/diagonal.py

@@ -53,7 +53,7 @@ class DiagonalSetup(QSetup, ABC):
     def _drift(self, _x: ArrayLike, gamma: float) -> ArrayLike:
         raise NotImplementedError
 
-    def construct_loss(self, state_q: TrainState, xi: ArrayLike, gamma: float, BS: int) -> Callable[
+    def construct_loss(self, state_q: TrainState, gamma: float, BS: int) -> Callable[
         [Union[FrozenVariableDict, Dict[str, Any]], ArrayLike], ArrayLike]:
 
         def loss_fn(params_q: Union[FrozenVariableDict, Dict[str, Any]], key: ArrayLike) -> ArrayLike:
@@ -80,14 +80,14 @@ def v_t(_eps, _t):
                 log_q_t = -(relative_mixture_weights / (_sigma_t ** 2) * (_x - _mu_t)).sum(axis=1)
                 u_t = (relative_mixture_weights * (1 / _sigma_t * _dsigmadt * (_x - _mu_t) + _dmudt)).sum(axis=1)
 
-                return u_t - self._drift(_x.reshape(BS, ndim), gamma) + 0.5 * (xi ** 2) * log_q_t
+                return u_t - self._drift(_x.reshape(BS, ndim), gamma) + 0.5 * (self.xi ** 2) * log_q_t
 
-            loss = 0.5 * ((v_t(eps, t) / xi) ** 2).sum(-1, keepdims=True)
+            loss = 0.5 * ((v_t(eps, t) / self.xi) ** 2).sum(-1, keepdims=True)
             return loss.mean()
 
         return loss_fn
 
-    def u_t(self, state_q: TrainState, t: ArrayLike, x_t: ArrayLike, xi: ArrayLike, *args, **kwargs) -> ArrayLike:
+    def u_t(self, state_q: TrainState, t: ArrayLike, x_t: ArrayLike, deterministic: bool, *args, **kwargs) -> ArrayLike:
         _mu_t, _sigma_t, _w_logits, _dmudt, _dsigmadt = forward_and_derivatives(state_q, t)
         _x = x_t[:, None, :]
 
@@ -96,56 +96,45 @@ def u_t(self, state_q: TrainState, t: ArrayLike, x_t: ArrayLike, xi: ArrayLike,
 
         _u_t = (relative_mixture_weights * (1 / _sigma_t * _dsigmadt * (_x - _mu_t) + _dmudt)).sum(axis=1)
 
-        if xi == 0:
+        if deterministic:
             return _u_t
 
         log_q_t = -(relative_mixture_weights / (_sigma_t ** 2) * (_x - _mu_t)).sum(axis=1)
 
-        return _u_t + 0.5 * (xi ** 2) * log_q_t
+        return _u_t + 0.5 * (self.xi ** 2) * log_q_t
 
 
 class FirstOrderSetup(DiagonalSetup):
-    def __init__(self, system: System, model: nn.module, T: float, base_sigma: float, num_mixtures: int,
+    def __init__(self, system: System, model: nn.module, xi: ArrayLike, T: float, base_sigma: float, num_mixtures: int,
                  trainable_weights: bool):
         model_q = DiagonalWrapper(model, T, system.A, system.B, num_mixtures, trainable_weights, base_sigma)
-        super().__init__(system, model_q, T, base_sigma, num_mixtures)
+        super().__init__(system, model_q, xi, T, base_sigma, num_mixtures)
 
     def _drift(self, _x: ArrayLike, gamma: float) -> ArrayLike:
         return -self.system.dUdx(_x / (gamma * self.system.mass))
 
 
 class SecondOrderSetup(DiagonalSetup):
-    def __init__(self, system: System, model: nn.module, T: float, base_sigma: float, num_mixtures: int,
+    def __init__(self, system: System, model: nn.module, xi: ArrayLike, T: float, base_sigma: float, num_mixtures: int,
                  trainable_weights: bool):
         # We pad the A and B matrices with zeros to account for the velocity
         self._A = jnp.hstack([system.A, jnp.zeros_like(system.A)])
         self._B = jnp.hstack([system.B, jnp.zeros_like(system.B)])
 
+        xi_velocity = jnp.ones_like(system.A) * xi
+        xi_pos = jnp.zeros_like(xi_velocity) + 1e-4
+
+        xi_second_order = jnp.concatenate((xi_pos, xi_velocity), axis=-1)
+
         model_q = DiagonalWrapper(model, T, self._A, self._B, num_mixtures, trainable_weights, base_sigma)
-        super().__init__(system, model_q, T, base_sigma, num_mixtures)
+        super().__init__(system, model_q, xi_second_order, T, base_sigma, num_mixtures)
 
     def _drift(self, _x: ArrayLike, gamma: float) -> ArrayLike:
         # number of dimensions without velocity
         ndim = self.system.A.shape[0]
 
         return jnp.hstack([_x[:, ndim:] / self.system.mass, -self.system.dUdx(_x[:, :ndim]) - _x[:, ndim:] * gamma])
 
-    def _xi_to_second_order(self, xi: ArrayLike) -> ArrayLike:
-        if xi.shape == self.model_q.A.shape:
-            return xi
-
-        xi_velocity = jnp.ones_like(self.system.A) * xi
-        xi_pos = jnp.zeros_like(xi_velocity) + 1e-4
-
-        return jnp.concatenate((xi_pos, xi_velocity), axis=-1)
-
-    def construct_loss(self, state_q: TrainState, xi: ArrayLike, gamma: float, BS: int) -> Callable[
-        [Union[FrozenVariableDict, Dict[str, Any]], ArrayLike], ArrayLike]:
-        return super().construct_loss(state_q, self._xi_to_second_order(xi), gamma, BS)
-
-    def u_t(self, state_q: TrainState, t: ArrayLike, x_t: ArrayLike, xi: ArrayLike, *args, **kwargs) -> ArrayLike:
-        return super().u_t(state_q, t, x_t, self._xi_to_second_order(xi), *args, **kwargs)
-
     @property
     def A(self):
         return self._A

training/qsetup.py

@@ -21,16 +21,17 @@ class QSetup(ABC):
     """
     system: System
     model_q: nn.Module
+    xi: ArrayLike
 
     @abstractmethod
     def construct_loss(self, *args, **kwargs) -> Callable:
         raise NotImplementedError
 
     def sample_paths(self, state_q: TrainState, x_0: ArrayLike, dt: float, T: float, BS: int,
-                     xi: Optional[float], key: Optional[ArrayLike], *args, **kwargs) -> ArrayLike:
+                     key: Optional[ArrayLike], *args, **kwargs) -> ArrayLike:
+        """Sample paths. If key is None, the sampling is deterministic. Otherwise, it is stochastic."""
         assert x_0.ndim == 2
         assert T / dt == int(T / dt), "dt must divide T evenly"
-        assert (xi is None) == (key is None), "xi and key must be both None or both specified"
         N = int(T / dt)
 
         num_paths = x_0.shape[0]
@@ -39,10 +40,7 @@ def sample_paths(self, state_q: TrainState, x_0: ArrayLike, dt: float, T: float,
         x_t = x_t.at[:, 0, :].set(x_0)
 
         t = jnp.zeros((BS, 1), dtype=jnp.float32)
-        if key is None:
-            u = jax.jit(lambda _t, _x: self.u_t(state_q, _t, _x, 0, *args, **kwargs))
-        else:
-            u = jax.jit(lambda _t, _x: self.u_t(state_q, _t, _x, xi, *args, **kwargs))
+        u = jax.jit(lambda _t, _x: self.u_t(state_q, _t, _x, key is None, *args, **kwargs))
 
         for i in trange(N):
             for j in range(0, num_paths, BS):
@@ -61,17 +59,17 @@ def sample_paths(self, state_q: TrainState, x_0: ArrayLike, dt: float, T: float,
                 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))
+                    noise = self.xi * jax.random.normal(iter_key, shape=(BS, ndim))
 
-                new_x = cur_x_t + dt * u(t, cur_x_t, *args, **kwargs) + jnp.sqrt(dt) * noise
+                new_x = cur_x_t + dt * u(t, cur_x_t) + jnp.sqrt(dt) * noise
                 x_t = x_t.at[j:j_end, i + 1, :].set(new_x[:j_end - j])
 
             t += dt
 
         return x_t
 
     @abstractmethod
-    def u_t(self, state_q: TrainState, t: ArrayLike, x_t: ArrayLike, xi: ArrayLike, *args, **kwargs) -> ArrayLike:
+    def u_t(self, state_q: TrainState, t: ArrayLike, x_t: ArrayLike, deterministic: bool, *args, **kwargs) -> ArrayLike:
         raise NotImplementedError
 
     @property
@@ -83,20 +81,21 @@ def B(self):
         return self.system.B
 
 
-def construct(system: System, model: nn.module, ode: str, parameterization: str, args: argparse.Namespace) -> QSetup:
+def construct(system: System, model: nn.module, ode: str, parameterization: str, xi: ArrayLike,
+              args: argparse.Namespace) -> QSetup:
     from training import diagonal
 
     if ode == 'first_order':
         if parameterization == 'diagonal':
-            return diagonal.FirstOrderSetup(system, model, args.T, args.base_sigma, args.num_gaussians,
+            return diagonal.FirstOrderSetup(system, model, xi, args.T, args.base_sigma, args.num_gaussians,
                                             args.trainable_weights)
         elif args.parameterization == 'low_rank':
             raise NotImplementedError("Low-rank parameterization not implemented")
         else:
             raise ValueError(f"Unknown parameterization: {args.parameterization}")
     elif args.ode == 'second_order':
         if parameterization == 'diagonal':
-            return diagonal.SecondOrderSetup(system, model, args.T, args.base_sigma, args.num_gaussians,
+            return diagonal.SecondOrderSetup(system, model, xi, args.T, args.base_sigma, args.num_gaussians,
                                              args.trainable_weights)
         else:
             raise NotImplementedError("Second-order ODE not implemented")