Properly handle float64

Author: plainerman

main.py

@@ -90,7 +90,7 @@ 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]))
+    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)
@@ -104,7 +104,7 @@ def main():
     show_or_save_fig(args.save_dir, 'loss_plot.pdf')
 
     # TODO: how to plot this nicely?
-    t = args.T * jnp.linspace(0, 1, args.BS).reshape((-1, 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, setup.A.shape[-1]])
     mu_t, sigma_t, w_logits = state_q.apply_fn(state_q.params, t)
@@ -118,7 +118,7 @@ def main():
     # plt.show()
 
     key, init_key = jax.random.split(key)
-    x_0 = jnp.ones((args.num_paths, setup.A.shape[0])) * setup.A
+    x_0 = jnp.ones((args.num_paths, setup.A.shape[0]), dtype=jnp.float32) * setup.A
     eps = jax.random.normal(key, shape=x_0.shape)
     x_0 += args.base_sigma * eps
 

systems.py

@@ -10,7 +10,7 @@
 from typing import Self
 from utils.pdb import assert_same_molecule
 from utils.rmsd import kabsch_align
-
+from dmff import Hamiltonian, NeighborList  # This sets jax to use 64-bit precision
 
 class System:
     def __init__(self, U: Callable[[ArrayLike], ArrayLike], A: ArrayLike, B: ArrayLike, mass: ArrayLike, plot):
@@ -49,17 +49,14 @@ def from_name(cls, name: str) -> Self:
 
     @classmethod
     def from_pdb(cls, A: str, B: str, forcefield: [str], cv: Optional[str]) -> Self:
-        print("WARNING!!!! This changes jax to double precision")
-        from dmff import Hamiltonian, NeighborList
-
         A_pdb, B_pdb = app.PDBFile(A), app.PDBFile(B)
         assert_same_molecule(A_pdb, B_pdb)
 
         mass = [a.element.mass.value_in_unit(unit.dalton) for a in A_pdb.topology.atoms()]
-        mass = jnp.broadcast_to(jnp.array(mass).reshape(-1, 1), (len(mass), 3)).reshape(-1)
+        mass = jnp.broadcast_to(jnp.array(mass, dtype=jnp.float32).reshape(-1, 1), (len(mass), 3)).reshape(-1)
 
-        A = jnp.array(A_pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
-        B = jnp.array(B_pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
+        A = jnp.array(A_pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer), dtype=jnp.float32)
+        B = jnp.array(B_pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer), dtype=jnp.float32)
         A, B = kabsch_align(A, B)
         A, B = A.reshape(-1), B.reshape(-1)
 

training/diagonal.py

@@ -3,7 +3,7 @@
 from jax.typing import ArrayLike
 from flax import linen as nn
 import jax.numpy as jnp
-from typing import Union, Dict, Any, Callable, Tuple
+from typing import Union, Dict, Any, Callable, Tuple, Optional
 from flax.training.train_state import TrainState
 import jax
 from flax.typing import FrozenVariableDict
@@ -60,8 +60,8 @@ def loss_fn(params_q: Union[FrozenVariableDict, Dict[str, Any]], key: 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, 1, ndim])
+            t = self.T * jax.random.uniform(key[0], [BS, 1], dtype=jnp.float32)
+            eps = jax.random.normal(key[1], [BS, 1, ndim], dtype=jnp.float32)
 
             def v_t(_eps, _t):
                 """This function is equal to v_t * xi ** 2."""

training/qsetup.py

@@ -35,10 +35,10 @@ def sample_paths(self, state_q: TrainState, x_0: ArrayLike, dt: float, T: float,
 
         num_paths = x_0.shape[0]
         ndim = x_0.shape[1]
-        x_t = jnp.ones((num_paths, N, ndim))
+        x_t = jnp.zeros((num_paths, N, ndim), dtype=jnp.float32)
         x_t = x_t.at[:, 0, :].set(x_0)
 
-        t = jnp.zeros((BS, 1))
+        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: