Refactor first and second order

Author: plainerman

model/utils.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from typing import Tuple, Any
 from flax import linen as nn
 from jax.typing import ArrayLike
 
@@ -16,9 +17,15 @@ class WrappedModule(ABC, nn.Module):
     def __call__(self, t: ArrayLike):
         t = t / self.T
 
-        h = self.other(t)
-        return self._post_process(t, h)
+        h, args = self._pre_process(t)
+        h = self.other(h)
+        return self._post_process(h, *args)
+
+    def _pre_process(self, t: ArrayLike) -> Tuple[ArrayLike, Tuple[Any, ...]]:
+        """This function returns a tuple. The first element will be used as an input to the other module,
+        and the second value will be passed to the post process function."""
+        return t, (t,)
 
     @abstractmethod
-    def _post_process(self, t: ArrayLike, h: ArrayLike):
+    def _post_process(self, h: ArrayLike, *args):
         raise NotImplementedError

training/qsetup.py

@@ -83,21 +83,30 @@ def B(self):
 
 def construct(system: System, model: nn.module, ode: str, parameterization: str, xi: ArrayLike,
               args: argparse.Namespace) -> QSetup:
-    from training import diagonal
+    from training.setups import diagonal
 
     if ode == 'first_order':
-        if parameterization == 'diagonal':
-            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, xi, args.T, args.base_sigma, args.num_gaussians,
-                                             args.trainable_weights)
-        else:
-            raise NotImplementedError("Second-order ODE not implemented")
+        order = 'first'
+        A = system.A
+        B = system.B
+    elif 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)])
+
+        xi_velocity = jnp.ones_like(system.A) * xi
+        xi_pos = jnp.zeros_like(xi_velocity) + 1e-4
+
+        xi = jnp.concatenate((xi_pos, xi_velocity), axis=-1)
+    else:
+        raise ValueError(f"Unknown ODE: {ode}")
+
+    if 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)
     else:
-        raise ValueError(f"Unknown ODE: {args.ode}")
+        raise ValueError(f"Unknown parameterization: {parameterization}")

training/diagonal.py renamed to training/setups/diagonal.py

@@ -1,15 +1,14 @@
-from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from jax.typing import ArrayLike
 from flax import linen as nn
 import jax.numpy as jnp
-from typing import Union, Dict, Any, Callable, Tuple, Optional
+from typing import Union, Dict, Any, Callable
 from flax.training.train_state import TrainState
 import jax
 from flax.typing import FrozenVariableDict
 from model.utils import WrappedModule
-from training.qsetup import QSetup
 from systems import System
+from training.setups.drift import DriftedSetup
 from training.utils import forward_and_derivatives
 
 
@@ -21,7 +20,7 @@ class DiagonalWrapper(WrappedModule):
     base_sigma: float
 
     @nn.compact
-    def _post_process(self, t: ArrayLike, h: ArrayLike):
+    def _post_process(self, h: ArrayLike, t: ArrayLike):
         ndim = self.A.shape[0]
         num_mixtures = self.num_mixtures
         h = nn.Dense(2 * ndim * num_mixtures)(h)
@@ -43,15 +42,13 @@ def _post_process(self, t: ArrayLike, h: ArrayLike):
 
 
 @dataclass
-class DiagonalSetup(QSetup, ABC):
+class DiagonalSetup(DriftedSetup):
     model_q: DiagonalWrapper
     T: float
-    base_sigma: float
-    num_mixtures: int
 
-    @abstractmethod
-    def _drift(self, _x: ArrayLike, gamma: float) -> ArrayLike:
-        raise NotImplementedError
+    def __init__(self, system: System, model_q: DiagonalWrapper, xi: ArrayLike, order: str, T: float):
+        super().__init__(system, model_q, xi, order)
+        self.T = T
 
     def construct_loss(self, state_q: TrainState, gamma: float, BS: int) -> Callable[
         [Union[FrozenVariableDict, Dict[str, Any]], ArrayLike], ArrayLike]:
@@ -70,7 +67,7 @@ def v_t(_eps, _t):
 
                 _x = _mu_t[jnp.arange(BS), _i, None] + _sigma_t[jnp.arange(BS), _i, None] * eps
 
-                if self.num_mixtures == 1:
+                if _mu_t.shape[1] == 1:
                     # This completely ignores the weights and saves some time
                     relative_mixture_weights = 1
                 else:
@@ -102,43 +99,3 @@ def u_t(self, state_q: TrainState, t: ArrayLike, x_t: ArrayLike, deterministic:
         log_q_t = -(relative_mixture_weights / (_sigma_t ** 2) * (_x - _mu_t)).sum(axis=1)
 
         return _u_t + 0.5 * (self.xi ** 2) * log_q_t
-
-
-class FirstOrderSetup(DiagonalSetup):
-    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, 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, 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, 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])
-
-    @property
-    def A(self):
-        return self._A
-
-    @property
-    def B(self):
-        return self._B

training/setups/drift.py

@@ -0,0 +1,37 @@
+from abc import ABC
+from flax import linen as nn
+import jax.numpy as jnp
+from systems import System
+from training.qsetup import QSetup
+from jax.typing import ArrayLike
+
+
+class DriftedSetup(QSetup, ABC):
+    """A QSetup that has a drift term. This drift term can be either first or second order."""
+
+    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)
+
+    def _drift(self, _x: ArrayLike, gamma: float) -> ArrayLike:
+        if self.order == 'first':
+            return -self.system.dUdx(_x / (gamma * self.system.mass))
+        else:
+            # 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])
+
+    @property
+    def A(self):
+        return self._A
+
+    @property
+    def B(self):
+        return self._B

training/setups/lowrank.py

@@ -29,9 +29,9 @@ def train_step(_state_q: TrainState, _key: ArrayLike) -> (TrainState, float):
                 log_loss = True
 
             if log_loss:
-                pbar.set_postfix(log_loss=jnp.log(loss))
+                pbar.set_postfix(log_loss=f"{jnp.log(loss):.4f}")
             else:
-                pbar.set_postfix(loss=loss)
+                pbar.set_postfix(loss=f"{loss:.4f}")
             ckpt['losses'].append(loss.item())
 
             if checkpoint_manager.should_save(i + 1):