Add force clipping

Author: plainerman

main.py

@@ -46,6 +46,7 @@
 parser.add_argument('--epochs', type=int, default=10_000, help="Number of epochs the system is training for.")
 parser.add_argument('--BS', type=int, default=512, help="Batch size used for training.")
 parser.add_argument('--lr', type=float, default=1e-4, help="Learning rate")
+parser.add_argument('--force_clip', type=float, default=1e8, help="Clipping value for the force")
 
 parser.add_argument('--seed', type=int, default=1, help="The seed that will be used for initialization")
 
@@ -55,8 +56,9 @@
 
 
 def main():
-    # TODO: force clipping
     print("!!!!Next todos: plot ALDP")
+    # TODO: internal coordinates
+    # TODO: neural network parameterization
 
     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"
@@ -70,9 +72,9 @@ def main():
     os.makedirs(args.save_dir, exist_ok=True)
 
     if args.test_system:
-        system = System.from_name(args.test_system)
+        system = System.from_name(args.test_system, args.force_clip)
     else:
-        system = System.from_pdb(args.start, args.target, args.forcefield, args.cv)
+        system = System.from_pdb(args.start, args.target, args.forcefield, args.cv, args.force_clip)
 
     if args.xi:
         xi = args.xi

potentials.py

@@ -1,14 +1,11 @@
-import jax
 import jax.numpy as jnp
 
 
-@jax.jit
 def U_double_well(xs, a=1.0, b=-4.0, c=0, d=1.0, beta=1.0):
     x, y = xs[:, 0], xs[:, 1]
     return beta * (a * (x ** 4) + b * (x ** 2) + c * x + 0.5 * d * (y ** 2))
 
 
-@jax.jit
 def U_double_well_hard(xs, beta=1.0):
     A = jnp.array([[-3, 0]])
     B = jnp.array([[3, 0]])
@@ -18,7 +15,6 @@ def U_double_well_hard(xs, beta=1.0):
     return beta * out
 
 
-@jax.jit
 def U_double_well_dual_channel(xs, beta=1.0):
     x, y = xs[:, 0], xs[:, 1]
     borders = x ** 6 + y ** 6
@@ -28,7 +24,6 @@ def U_double_well_dual_channel(xs, beta=1.0):
     return beta * (borders + e1 + e2 + e3)
 
 
-@jax.jit
 def U_mueller_brown(xs, beta=1.0):
     x, y = xs[:, 0], xs[:, 1]
     e1 = -200 * jnp.exp(-(x - 1) ** 2 - 10 * y ** 2)

systems.py

@@ -11,21 +11,26 @@
 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
+import mdtraj as md
+
 
 class System:
-    def __init__(self, U: Callable[[ArrayLike], ArrayLike], A: ArrayLike, B: ArrayLike, mass: ArrayLike, plot):
+    def __init__(self, U: Callable[[ArrayLike], ArrayLike], A: ArrayLike, B: ArrayLike, mass: ArrayLike, plot,
+                 force_clip: float):
         assert A.shape == B.shape == mass.shape
 
         self.U = U
-        self.dUdx = jax.jit(jax.grad(lambda _x: U(_x).sum()))
+
+        dUdx = jax.grad(lambda _x: U(_x).sum())
+        self.dUdx = jax.jit(jax.jit(lambda _x: jnp.clip(dUdx(_x), -force_clip, force_clip)))
 
         self.A, self.B = A, B
         self.mass = mass
 
         self.plot = plot
 
     @classmethod
-    def from_name(cls, name: str) -> Self:
+    def from_name(cls, name: str, force_clip: float) -> Self:
         if name == 'double_well':
             U, A, B = potentials.double_well
         elif name == 'double_well_hard':
@@ -45,10 +50,10 @@ def from_name(cls, name: str) -> Self:
                        U=U, states=list(zip(['A', 'B'], [A, B])), xlim=xlim, ylim=ylim, alpha=1.0
                        )
         mass = jnp.array([1.0, 1.0])
-        return cls(U, A, B, mass, plot)
+        return cls(U, A, B, mass, plot, force_clip)
 
     @classmethod
-    def from_pdb(cls, A: str, B: str, forcefield: [str], cv: Optional[str]) -> Self:
+    def from_pdb(cls, A: str, B: str, forcefield: [str], cv: Optional[str], force_clip: float) -> Self:
         A_pdb, B_pdb = app.PDBFile(A), app.PDBFile(B)
         assert_same_molecule(A_pdb, B_pdb)
 
@@ -87,4 +92,4 @@ def U(_x):
         else:
             raise ValueError(f"Unknown cv: {cv}")
 
-        return cls(U, A, B, mass, None)
+        return cls(U, A, B, mass, plot, force_clip)