Miscellaneous fixes (#335)

1. In the past our SpectralABF examples needed to be manually shifted, but this works fine now, so I'm updating the examples to reflect that.

2. In some cases the synchronization between OpenMM, jax and cupy CUDA contexts/streams might throw an error, so fixing this as well.

examples/hoomd-blue/spectral_abf/Butane-SpectralABF.ipynb

@@ -638,8 +638,7 @@
     "A = result[\"free_energy\"]\n",
     "# Alternatively:\n",
     "# fes_fn = result[\"fes_fn\"]\n",
-    "# A = fes_fn(mesh)\n",
-    "A = A.max() - A"
+    "# A = fes_fn(mesh)"
    ]
   },
   {

examples/hoomd-blue/spectral_abf/Butane-SpectralABF.md

@@ -386,7 +386,6 @@ A = result["free_energy"]
 # Alternatively:
 # fes_fn = result["fes_fn"]
 # A = fes_fn(mesh)
-A = A.max() - A
 ```
 
 ```python colab={"base_uri": "https://localhost:8080/", "height": 302} id="7_d_XfVLLkbI" outputId="e35db259-31f8-4a3b-b1fa-7e91a8a5c88a"

examples/hoomd-blue/spectral_abf/butane.py

@@ -257,7 +257,7 @@ def get_args(argv):
         ("time-steps", "t", int, 5e5, "Number of simulation steps"),
     ]
     parser = argparse.ArgumentParser(description="Example script to run SpectralABF")
-    for (name, short, T, val, doc) in available_args:
+    for name, short, T, val, doc in available_args:
         parser.add_argument("--" + name, "-" + short, type=T, default=T(val), help=doc)
 
     return parser.parse_args(argv)
@@ -283,7 +283,6 @@ def main(argv=[]):
     mesh = result["mesh"]
     fes_fn = result["fes_fn"]
     A = fes_fn(mesh)
-    A = A.max() - A
 
     # plot the free energy
     fig, ax = plt.subplots()

examples/hoomd3/spectral_abf/butane.py

@@ -298,7 +298,6 @@ def main(argv=[]):
     mesh = result["mesh"]
     fes_fn = result["fes_fn"]
     A = fes_fn(mesh)
-    A = A.max() - A
 
     # plot the free energy
     fig, ax = plt.subplots()

examples/openmm/spectral_abf/ADP_SpectralABF.ipynb

@@ -342,7 +342,6 @@
     "# mesh = result[\"mesh\"]\n",
     "# fes_fn = result[\"fes_fn\"]\n",
     "# A = fes_fn(mesh)\n",
-    "A = A.max() - A\n",
     "A = A.reshape(grid.shape)"
    ]
   },

examples/openmm/spectral_abf/ADP_SpectralABF.md

@@ -211,7 +211,6 @@ A = result["free_energy"]
 # mesh = result["mesh"]
 # fes_fn = result["fes_fn"]
 # A = fes_fn(mesh)
-A = A.max() - A
 A = A.reshape(grid.shape)
 ```
 

examples/openmm/spectral_abf/alanine-dipeptide.py

@@ -78,7 +78,7 @@ def get_args(argv):
         ("time-steps", "t", int, 5e5, "Number of simulation steps"),
     ]
     parser = argparse.ArgumentParser(description="Example script to run Spectral ABF")
-    for (name, short, T, val, doc) in available_args:
+    for name, short, T, val, doc in available_args:
         parser.add_argument("--" + name, "-" + short, type=T, default=T(val), help=doc)
     return parser.parse_args(argv)
 
@@ -108,7 +108,6 @@ def main(argv=[]):
 
     # Set min free energy to zero
     A = fes_fn(xi)
-    A = A.max() - A
     A = A.reshape(plot_grid.shape)
 
     # plot and save free energy to a PNG file

pysages/backends/openmm.py

@@ -157,13 +157,12 @@ def bias(snapshot, state, sync_backend):
         """Adds the computed bias to the forces."""
         if state.bias is None:
             return
-        biases = adapt(state.bias)
         # Forces may be computed asynchronously on the GPU, so we need to
         # synchronize them before applying the bias.
         sync_backend()
+        biases = adapt(state.bias)
         forces = view(snapshot.forces)
-        biases = view(biases.block_until_ready())
-        forces += biases
+        forces += view(biases.block_until_ready())
         sync_forces()
 
     def dimensionality():