Update simulate.py

jax_md/quantity.py

@@ -470,3 +470,15 @@ def update_fn(state: PHopState, position: Array) -> PHopState:
     return PHopState(buff, phop)  # pytype: disable=wrong-arg-count
 
   return init_fn, update_fn
+
+def gamma_from_stokes_law_3d(eta, radii, mass):
+  """
+  Compute gamma value for a sphere in fluid using stokes' law in three dimensions.
+
+  Args:
+    eta: A float specifying the dynamic viscosity of the medium
+    radii: A float or an Array specifying the radii of the spheres
+    mass: A float or an Array specifying the mass of the spheres
+  """
+
+  return 6 * jnp.pi * eta * radii / mass

jax_md/simulate.py

@@ -930,15 +930,14 @@ def apply_fn(state, **kwargs):
     return NVTLangevinState(R, V, F_new, mass, key)  # pytype: disable=wrong-arg-count
   return init_fn, apply_fn
 
-
 @dataclasses.dataclass
 class BrownianState:
   """A tuple containing state information for Brownian dynamics.
 
   Attributes:
     position: The current position of the particles. An ndarray of floats with
       shape [n, spatial_dimension].
-    mass: The mmass of particles. Will either be a float or an ndarray of floats
+    mass: The mass of particles. Will either be a float or an ndarray of floats
       with shape [n].
     rng: The current state of the random number generator.
   """
@@ -972,7 +971,9 @@ def brownian(energy_or_force: Callable[..., Array],
       constant. To update the temperature dynamically during a simulation one
       should pass `kT` as a keyword argument to the step function.
     gamma: A float specifying the friction coefficient between the particles
-      and the solvent.
+      and the solvent. The gamma here is the friction coeifficient divided by
+      the mass. For example, when the particles are 3 diemsional spheres
+      gamma = 6 pi eta R/mass. See quantity.gamma_from_stokes_law for detail.
 
   Returns:
     See above.
@@ -993,6 +994,7 @@ def init_fn(key, R, mass=f32(1)):
 
   def apply_fn(state, **kwargs):
     _kT = kT if 'kT' not in kwargs else kwargs['kT']
+    _gamma = gamma if 'gamma' not in kwargs else kwargs['gamma']
 
     R, mass, key = dataclasses.astuple(state)
 
@@ -1001,7 +1003,7 @@ def apply_fn(state, **kwargs):
     F = force_fn(R, **kwargs)
     xi = random.normal(split, R.shape, R.dtype)
 
-    nu = f32(1) / (mass * gamma)
+    nu = f32(1) / (mass * _gamma)
 
     dR = F * dt * nu + jnp.sqrt(f32(2) * _kT * dt * nu) * xi
     R = shift(R, dR, **kwargs)
@@ -1019,7 +1021,6 @@ def apply_fn(state, **kwargs):
 as the more polished environments above.
 """
 
-
 @dataclasses.dataclass
 class SwapMCState:
   """A struct containing state information about a Hybrid Swap MC simulation.