Add low rank training

plainerman · plainerman · commit 99d077c5a9fe · 2024-07-09T14:49:49.000+02:00
diff --git a/configs/toy/mueller_single_gaussian.yaml b/configs/toy/mueller_single_gaussian.yaml
@@ -1,7 +1,7 @@
 save_dir: ./out/toy/mueller_single_gaussian
 
 test_system: mueller_brown
-ode: second_order
+ode: first_order
 parameterization: diagonal
 T: 275e-4
 xi: 5.0
diff --git a/configs/toy/mueller_single_gaussian_low_rank.yaml b/configs/toy/mueller_single_gaussian_low_rank.yaml
@@ -1,8 +1,8 @@
-save_dir: ./out/toy/mueller_single_gaussian
+save_dir: ./out/toy/mueller_single_gaussian_low_rank
 
 test_system: mueller_brown
 ode: first_order
-parameterization: diagonal
+parameterization: low_rank
 T: 275e-4
 xi: 5.0
 gamma: 1.0
diff --git a/environment.yml b/environment.yml
@@ -15,6 +15,7 @@ dependencies:
   - openpathsampling=1.6.1
   - jax=0.4.23
   - flax=0.8.3
+  - notebook=7.0.8
   - pip:
     - dmff @ git+https://github.com/deepmodeling/DMFF@v1.0.0
     - rdkit==2023.3.3
diff --git a/main.py b/main.py
@@ -125,7 +125,9 @@ def main():
 
     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))
+    print(params_q)
 
     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)
@@ -146,6 +148,8 @@ def main():
         if checkpoint_manager.latest_step() is None:
             print("Warning: No checkpoint found.")
         else:
+            # TODO: fix this. At least for low rank it does not work
+
             print('Loading checkpoint:', checkpoint_manager.latest_step())
 
             state_restored = checkpoint_manager.restore(checkpoint_manager.latest_step())
@@ -165,12 +169,12 @@ def main():
     show_or_save_fig(args.save_dir, 'loss_plot.pdf')
 
     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, 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)
+    # 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, 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)
 
     # plot_energy_surface()
     # plt.scatter(samples[:, 0], samples[:, 1])
@@ -180,7 +184,7 @@ def main():
 
     key, init_key = jax.random.split(key)
     x_0 = jnp.ones((args.num_paths, A.shape[0]), dtype=jnp.float32) * A
-    eps = jax.random.normal(key, shape=x_0.shape)
+    eps = jax.random.normal(key, shape=x_0.shape, dtype=jnp.float32)
     x_0 += args.base_sigma * eps
 
     x_t_det = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, None)
diff --git a/prepare_molecule.py b/prepare_molecule.py
@@ -34,7 +34,7 @@ def minimize(pdb, out, steps):
         for mass_ in mass:
             for _ in range(3):
                 new_mass.append(mass_)
-        mass = jnp.array(new_mass)
+        mass = jnp.array(new_mass, dtype=jnp.float32)
 
         # Initialize the potential energy with amber forcefields
         ff = Hamiltonian('amber14/protein.ff14SB.xml', 'amber14/tip3p.xml')
diff --git a/systems.py b/systems.py
@@ -50,7 +50,7 @@ def from_name(cls, name: str, force_clip: float) -> Self:
         plot = partial(toy_plot_energy_surface,
                        U=U, states=list(zip(['A', 'B'], [A, B])), xlim=xlim, ylim=ylim, alpha=1.0
                        )
-        mass = jnp.array([1.0, 1.0])
+        mass = jnp.array([1.0, 1.0], dtype=jnp.float32)
         return cls(U, A, B, mass, plot, force_clip)
 
     @classmethod
diff --git a/training/qsetup.py b/training/qsetup.py
@@ -75,7 +75,7 @@ def u_t(self, state_q: TrainState, t: ArrayLike, x_t: ArrayLike, deterministic:
 
 def construct(system: System, model: nn.module, xi: float, A: ArrayLike, B: ArrayLike,
               args: argparse.Namespace) -> QSetup:
-    from training.setups import diagonal
+    from training.setups import diagonal, lowrank
 
     transform = None
     if args.internal_coordinates:
@@ -87,5 +87,10 @@ def construct(system: System, model: nn.module, xi: float, A: ArrayLike, B: Arra
             model, args.T, transform, A, B, args.num_gaussians, args.trainable_weights, args.base_sigma
         )
         return diagonal.DiagonalSetup(system, wrapped_module, xi, args.ode, args.T)
+    elif args.parameterization == 'low_rank':
+        wrapped_module = lowrank.LowRankWrapper(
+            model, args.T, transform, A, B, args.num_gaussians, args.trainable_weights, args.base_sigma
+        )
+        return lowrank.LowRankSetup(system, wrapped_module, xi, args.ode, args.T)
     else:
         raise ValueError(f"Unknown parameterization: {args.parameterization}")
diff --git a/training/setups/diagonal.py b/training/setups/diagonal.py
@@ -34,9 +34,9 @@ def _post_process(self, h: ArrayLike, t: ArrayLike):
         )
 
         if self.trainable_weights:
-            w_logits = self.param('w_logits', nn.initializers.zeros_init(), (num_mixtures,))
+            w_logits = self.param('w_logits', nn.initializers.zeros_init(), (num_mixtures,), dtype=jnp.float32)
         else:
-            w_logits = jnp.zeros(num_mixtures)
+            w_logits = jnp.zeros(num_mixtures, dtype=jnp.float32)
 
         return mu, sigma, w_logits
 
diff --git a/training/setups/lowrank.py b/training/setups/lowrank.py
@@ -19,15 +19,16 @@ class LowRankWrapper(WrappedModule):
     trainable_weights: bool
     base_sigma: float
 
-    def _pre_process(self, t: ArrayLike) -> Tuple[ArrayLike, Tuple[ArrayLike, ArrayLike, ArrayLike]]:
+    def _pre_process(self, t: ArrayLike) -> Tuple[Tuple[ArrayLike], Tuple[ArrayLike, ArrayLike, ArrayLike]]:
         ndim = self.A.shape[0]
 
         h_mu = (1 - t) * self.A + t * self.B
-        S_0 = jnp.eye(ndim)
-        S_0 = S_0 * jnp.vstack([self.base_sigma * jnp.ones((ndim // 2, 1)), self.base_sigma * jnp.ones((ndim // 2, 1))])
+        S_0 = jnp.eye(ndim, dtype=jnp.float32)
+        S_0 = S_0 * jnp.vstack([self.base_sigma * jnp.ones((ndim // 2, 1), dtype=jnp.float32),
+                                self.base_sigma * jnp.ones((ndim // 2, 1), dtype=jnp.float32)])
         S_0 = S_0[None, ...]
         h_S = (1 - 2 * t * (1 - t))[..., None] * S_0
-        return jnp.hstack([h_mu, h_S.reshape(-1, ndim * ndim), t]), (h_mu, h_S, t)
+        return (jnp.hstack([h_mu, h_S.reshape(-1, ndim * ndim), t]),), (h_mu, h_S, t)
 
     @nn.compact
     def _post_process(self, h: ArrayLike, h_mu: ArrayLike, h_S: ArrayLike, t: ArrayLike):
@@ -43,18 +44,18 @@ def _post_process(self, h: ArrayLike, h_mu: ArrayLike, h_S: ArrayLike, t: ArrayL
 
         @jax.vmap
         def get_tril(v):
-            a = jnp.zeros((ndim, ndim))
+            a = jnp.zeros((ndim, ndim), dtype=jnp.float32)
             a = a.at[jnp.tril_indices(ndim)].set(v)
             return a
 
         S = get_tril(h[:, ndim:])
-        S = jnp.tril(2 * jax.nn.sigmoid(S) - 1.0, k=-1) + jnp.eye(ndim)[None, ...] * jnp.exp(S)
+        S = jnp.tril(2 * jax.nn.sigmoid(S) - 1.0, k=-1) + jnp.eye(ndim, dtype=jnp.float32)[None, ...] * jnp.exp(S)
         S = h_S + 2 * ((1 - t) * t)[..., None] * S
 
         if self.trainable_weights:
-            w_logits = self.param('w_logits', nn.initializers.zeros_init(), (num_mixtures,))
+            w_logits = self.param('w_logits', nn.initializers.zeros_init(), (num_mixtures,), dtype=jnp.float32)
         else:
-            w_logits = jnp.zeros(num_mixtures)
+            w_logits = jnp.zeros(num_mixtures, dtype=jnp.float32)
 
         print('mu.shape', mu.shape)
         print('S.shape', S.shape)
@@ -65,93 +66,47 @@ def get_tril(v):
 @dataclass
 class LowRankSetup(DriftedSetup):
     model_q: LowRankWrapper
-    T: float
 
     def __init__(self, system: System, model_q: LowRankWrapper, xi: ArrayLike, order: str, T: float):
-        super().__init__(system, model_q, xi, order)
-        self.T = T
+        super().__init__(system, model_q, xi, order, T)
 
     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:
             ndim = self.model_q.A.shape[-1]
 
             key = jax.random.split(key)
-            t = self.T * jax.random.uniform(key[0], [BS, 1])
-            eps = jax.random.normal(key[1], [BS, ndim, 1])
-
-            mu_t = lambda _t: state_q.apply_fn(params_q, _t)[0]
-            S_t = lambda _t: state_q.apply_fn(params_q, _t)[1]
-
-            def dmudt(_t):
-                _dmudt = jax.jacrev(lambda _t: mu_t(_t).sum(0))
-                return _dmudt(_t).squeeze().T
-
-            def dSdt(_t):
-                _dSdt = jax.jacrev(lambda _t: S_t(_t).sum(0))
-                return _dSdt(_t).squeeze().T
+            t = self.T * jax.random.uniform(key[0], [BS, 1], dtype=jnp.float32)
+            eps = jax.random.normal(key[1], [BS, ndim, 1], dtype=jnp.float32)
 
             def v_t(_eps, _t):
-                S_t_val, dSdt_val = S_t(_t), dSdt(_t)
-                _x = mu_t(_t) + jax.lax.batch_matmul(S_t_val, _eps).squeeze()
-                dlogdx = -jax.scipy.linalg.solve_triangular(jnp.transpose(S_t_val, (0, 2, 1)), _eps)
+                _mu_t, _S_t_val, _w_logits, _dmudt, _dSdt_val = forward_and_derivatives(state_q, _t, params_q)
+
+                _x = _mu_t + jax.lax.batch_matmul(_S_t_val, _eps).squeeze()
+                dlogdx = -jax.scipy.linalg.solve_triangular(jnp.transpose(_S_t_val, (0, 2, 1)), _eps)
                 # S_t_val_inv = jnp.transpose(jnp.linalg.inv(S_t_val), (0,2,1))
                 # dlogdx = -jax.lax.batch_matmul(S_t_val_inv, _eps)
-                dSigmadt = jax.lax.batch_matmul(dSdt_val, jnp.transpose(S_t_val, (0, 2, 1)))
-                dSigmadt += jax.lax.batch_matmul(S_t_val, jnp.transpose(dSdt_val, (0, 2, 1)))
-                u_t = dmudt(_t) - 0.5 * jax.lax.batch_matmul(dSigmadt, dlogdx).squeeze()
+                dSigmadt = jax.lax.batch_matmul(_dSdt_val, jnp.transpose(_S_t_val, (0, 2, 1)))
+                dSigmadt += jax.lax.batch_matmul(_S_t_val, jnp.transpose(_dSdt_val, (0, 2, 1)))
+                u_t = _dmudt - 0.5 * jax.lax.batch_matmul(dSigmadt, dlogdx).squeeze()
                 out = (u_t - self._drift(_x.reshape(BS, ndim), gamma)) + 0.5 * (self.xi ** 2) * dlogdx.squeeze()
                 return out
 
             loss = 0.5 * ((v_t(eps, t) / self.xi) ** 2).sum(1, keepdims=True)
-            print(loss.shape, 'loss.shape', flush=True)
+            print(loss.shape, 'loss.shape', 'loss.dtype', loss.dtype, flush=True)
             return loss.mean()
 
-            # ndim = self.model_q.A.shape[-1]
-            # key = jax.random.split(key)
-            #
-            # t = self.T * jax.random.uniform(key[0], [BS, 1], dtype=jnp.float32)
-            # #TODO: the following needs to be changed for num gaussians. It should be BS, num_mitures, ndim
-            # eps = jax.random.normal(key[1], [BS, ndim, 1], dtype=jnp.float32)
-            #
-            # def v_t(_eps, _t):
-            #     """This function is equal to v_t * xi ** 2."""
-            #     _mu_t, _sigma_t, _w_logits, _dmudt, _dsigmadt = forward_and_derivatives(state_q, _t, params_q)
-            #     _i = jax.random.categorical(key[2], _w_logits, shape=[BS, ])
-            #
-            #     _x = _mu_t[jnp.arange(BS), _i, None] + _sigma_t[jnp.arange(BS), _i, None] * eps
-            #
-            #     if _mu_t.shape[1] == 1:
-            #         # This completely ignores the weights and saves some time
-            #         relative_mixture_weights = 1
-            #     else:
-            #         log_q_i = jax.scipy.stats.norm.logpdf(_x, _mu_t, _sigma_t).sum(-1)
-            #         relative_mixture_weights = jax.nn.softmax(_w_logits + log_q_i)[:, :, None]
-            #
-            #     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 * (self.xi ** 2) * log_q_t
-            #
-            # 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, deterministic: bool, *args, **kwargs) -> ArrayLike:
-        raise NotImplementedError
-
-        # _mu_t, _sigma_t, _w_logits, _dmudt, _dsigmadt = forward_and_derivatives(state_q, t)
-        # _x = x_t[:, None, :]
-        #
-        # log_q_i = jax.scipy.stats.norm.logpdf(_x, _mu_t, _sigma_t).sum(-1)
-        # relative_mixture_weights = jax.nn.softmax(_w_logits + log_q_i)[:, :, None]
-        #
-        # _u_t = (relative_mixture_weights * (1 / _sigma_t * _dsigmadt * (_x - _mu_t) + _dmudt)).sum(axis=1)
-        #
-        # 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 * (self.xi ** 2) * log_q_t
+        _mu_t, _S_t_val, _w_logits, _dmudt, _dSdt_val = forward_and_derivatives(state_q, t)
+
+        dSigmadt = jax.lax.batch_matmul(_dSdt_val, jnp.transpose(_S_t_val, (0, 2, 1)))
+        dSigmadt += jax.lax.batch_matmul(_S_t_val, jnp.transpose(_dSdt_val, (0, 2, 1)))
+        STdlogdx = jax.scipy.linalg.solve_triangular(_S_t_val, (x_t - _mu_t)[..., None])
+        dlogdx = -jax.scipy.linalg.solve_triangular(jnp.transpose(_S_t_val, (0, 2, 1)), STdlogdx)
+
+        if deterministic:
+            return _dmudt + (-0.5 * jax.lax.batch_matmul(dSigmadt, dlogdx)).squeeze()
+
+        return _dmudt + (-0.5 * jax.lax.batch_matmul(dSigmadt, dlogdx) + 0.5 * self.xi ** 2 * dlogdx).squeeze()
diff --git a/training/train.py b/training/train.py
@@ -20,16 +20,12 @@ def train_step(_state_q: TrainState, _key: ArrayLike) -> (TrainState, float):
         _state_q = _state_q.apply_gradients(grads=grads)
         return _state_q, loss
 
-    log_loss = False
     with trange(ckpt['model'].step, epochs) as pbar:
         for i in pbar:
             key, loc_key = jax.random.split(key)
             ckpt['model'], loss = train_step(ckpt['model'], loc_key)
             if loss > 1e4:
-                log_loss = True
-
-            if log_loss:
-                pbar.set_postfix(log_loss=f"{jnp.log(loss):.4f}")
+                pbar.set_postfix(log_loss=f"{jnp.log10(loss):.4f}")
             else:
                 pbar.set_postfix(loss=f"{loss:.4f}")
             ckpt['losses'].append(loss.item())

Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,7 @@ def from_name(cls, name: str, force_clip: float) -> Self:`
`50`	`50`	`plot = partial(toy_plot_energy_surface,`
`51`	`51`	`U=U, states=list(zip(['A', 'B'], [A, B])), xlim=xlim, ylim=ylim, alpha=1.0`
`52`	`52`	`)`
`53`		`- mass = jnp.array([1.0, 1.0])`
	`53`	`+ mass = jnp.array([1.0, 1.0], dtype=jnp.float32)`
`54`	`54`	`return cls(U, A, B, mass, plot, force_clip)`
`55`	`55`
`56`	`56`	`@classmethod`
Original file line number	Diff line number	Diff line change
`@@ -34,9 +34,9 @@ def _post_process(self, h: ArrayLike, t: ArrayLike):`
`34`	`34`	`)`
`35`	`35`
`36`	`36`	`if self.trainable_weights:`
`37`		`- w_logits = self.param('w_logits', nn.initializers.zeros_init(), (num_mixtures,))`
	`37`	`+ w_logits = self.param('w_logits', nn.initializers.zeros_init(), (num_mixtures,), dtype=jnp.float32)`
`38`	`38`	`else:`
`39`		`- w_logits = jnp.zeros(num_mixtures)`
	`39`	`+ w_logits = jnp.zeros(num_mixtures, dtype=jnp.float32)`
`40`	`40`
`41`	`41`	`return mu, sigma, w_logits`
`42`	`42`