def test_pair_neighbor_list_force_scalar_diverging_potential(
self, spatial_dimension, dtype):
key = random.PRNGKey(0)
def potential(dr, sigma):
return np.where(dr < sigma, dr**-6, f32(0.))
N = NEIGHBOR_LIST_PARTICLE_COUNT
box_size = 4. * N**(1. / spatial_dimension)
key, split = random.split(key)
disp, _ = space.periodic(box_size)
d = space.metric(disp)
neighbor_square = jit(
quantity.force(smap.pair_neighbor_list(potential, d)))
mapped_square = jit(quantity.force(smap.pair(potential, d)))
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = box_size * random.uniform(split, (N, spatial_dimension),
dtype=dtype)
sigma = random.uniform(key, (), minval=0.5, maxval=4.5)
neighbor_fn = partition.neighbor_list(disp, box_size, sigma, 0.0)
nbrs = neighbor_fn(R)
self.assertAllClose(mapped_square(R, sigma=sigma),
neighbor_square(R, nbrs, sigma=sigma))
def test_lennard_jones_neighbor_list_force(self, spatial_dimension, dtype,
format):
key = random.PRNGKey(1)
box_size = f32(15.0)
displacement, _ = space.periodic(box_size)
metric = space.metric(displacement)
exact_force_fn = quantity.force(
energy.lennard_jones_pair(displacement))
r = box_size * random.uniform(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
neighbor_fn, energy_fn = energy.lennard_jones_neighbor_list(
displacement, box_size, format=format)
force_fn = quantity.force(energy_fn)
nbrs = neighbor_fn.allocate(r)
if dtype == f32 and format is partition.OrderedSparse:
self.assertAllClose(np.array(exact_force_fn(r), dtype=dtype),
force_fn(r, nbrs),
atol=5e-5,
rtol=5e-5)
else:
self.assertAllClose(np.array(exact_force_fn(r), dtype=dtype),
force_fn(r, nbrs))
def test_lennard_jones_cell_list_force(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f32(15.0)
displacement, _ = space.periodic(box_size)
metric = space.metric(displacement)
exact_force_fn = quantity.force(
energy.lennard_jones_pair(displacement))
R = box_size * random.uniform(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
force_fn = quantity.force(
energy.lennard_jones_cell_list(displacement, box_size, R))
self.assertAllClose(np.array(exact_force_fn(R), dtype=dtype),
force_fn(R), True)
def test_morse_neighbor_list_force(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f32(15.0)
displacement, _ = space.periodic(box_size)
metric = space.metric(displacement)
exact_force_fn = quantity.force(energy.morse_pair(displacement))
r = box_size * random.uniform(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
neighbor_fn, energy_fn = energy.morse_neighbor_list(
displacement, box_size)
force_fn = quantity.force(energy_fn)
nbrs = neighbor_fn(r)
self.assertAllClose(np.array(exact_force_fn(r), dtype=dtype),
force_fn(r, nbrs))
def test_cell_list_direct_force_jit(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f32(9.0)
cell_size = f32(1.0)
displacement, _ = space.periodic(box_size)
energy_fn = energy.soft_sphere_pair(displacement)
force_fn = quantity.force(energy_fn)
R = box_size * random.uniform(
key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
grid_energy_fn = smap.cartesian_product(
energy.soft_sphere, space.metric(displacement))
grid_force_fn = quantity.force(grid_energy_fn)
grid_force_fn = jit(smap.cell_list(grid_force_fn, box_size, cell_size, R))
self.assertAllClose(
np.array(force_fn(R), dtype=dtype), grid_force_fn(R), True)
def test_cell_list_force_nonuniform(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
if spatial_dimension == 2:
box_size = f32(np.array([[8.0, 10.0]]))
else:
box_size = f32(np.array([[8.0, 10.0, 12.0]]))
cell_size = f32(2.0)
displacement, _ = space.periodic(box_size[0])
energy_fn = energy.soft_sphere_pair(displacement)
force_fn = quantity.force(energy_fn)
R = box_size * random.uniform(
key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
cell_energy_fn = smap.cartesian_product(
energy.soft_sphere, space.metric(displacement))
cell_force_fn = quantity.force(cell_energy_fn)
cell_force_fn = smap.cell_list(cell_force_fn, box_size, cell_size, R)
df = np.sum((force_fn(R) - cell_force_fn(R)) ** 2, axis=1)
self.assertAllClose(
np.array(force_fn(R), dtype=dtype), cell_force_fn(R), True)
def potential(R: space.Array, neighbor: NeighborList, *args,
**kwargs) -> PotentialProperties:
# this simply wraps the passed energy_fn to achieve a single consistent signature for all following lambdas.
atomwise_energy_fn = lambda R, neighbor, *args, **kwargs: energy_fn(
R, *args, **kwargs, neighbor=neighbor)
total_energy_fn = lambda R, neighbor, *args, **kwargs: jnp.sum(
energy_fn(R, *args, **kwargs, neighbor=neighbor))
force_fn = quantity.force(total_energy_fn)
atomwise_energies = atomwise_energy_fn(R, neighbor, *args, **kwargs)
total_energy = total_energy_fn(R, neighbor, *args, **kwargs)
forces = force_fn(R, neighbor, *args, **kwargs)
return total_energy, atomwise_energies, forces, None
def test_pair_grid_force_incommensurate(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f32(12.1)
cell_size = f32(3.0)
displacement, _ = space.periodic(box_size)
energy_fn = energy.soft_sphere_pair(displacement, quantity.Dynamic)
force_fn = quantity.force(energy_fn)
R = box_size * random.uniform(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
grid_force_fn = jit(smap.grid(force_fn, box_size, cell_size, R))
species = np.zeros((PARTICLE_COUNT, ), dtype=np.int64)
self.assertAllClose(np.array(force_fn(R, species, 1), dtype=dtype),
grid_force_fn(R), True)
def main(unused_argv):
key = random.PRNGKey(0)
# Setup some variables describing the system.
N = 500
dimension = 2
box_size = f32(25.0)
# Create helper functions to define a periodic box of some size.
displacement, shift = space.periodic(box_size)
metric = space.metric(displacement)
# Use JAX's random number generator to generate random initial positions.
key, split = random.split(key)
R = random.uniform(split, (N, dimension),
minval=0.0,
maxval=box_size,
dtype=f32)
# The system ought to be a 50:50 mixture of two types of particles, one
# large and one small.
sigma = np.array([[1.0, 1.2], [1.2, 1.4]], dtype=f32)
N_2 = int(N / 2)
species = np.array([0] * N_2 + [1] * N_2, dtype=i32)
# Create an energy function.
energy_fn = energy.soft_sphere_pair(displacement, species, sigma)
force_fn = quantity.force(energy_fn)
# Create a minimizer.
init_fn, apply_fn = minimize.fire_descent(energy_fn, shift)
opt_state = init_fn(R)
# Minimize the system.
minimize_steps = 50
print_every = 10
print('Minimizing.')
print('Step\tEnergy\tMax Force')
print('-----------------------------------')
for step in range(minimize_steps):
opt_state = apply_fn(opt_state)
if step % print_every == 0:
R = opt_state.position
print('{:.2f}\t{:.2f}\t{:.2f}'.format(step, energy_fn(R),
np.max(force_fn(R))))
def test_pair_grid_force_nonuniform(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
if spatial_dimension == 2:
box_size = f32(np.array([[8.0, 10.0]]))
else:
box_size = f32(np.array([[8.0, 10.0, 12.0]]))
cell_size = f32(2.0)
displacement, _ = space.periodic(box_size[0])
energy_fn = energy.soft_sphere_pair(displacement, quantity.Dynamic)
force_fn = quantity.force(energy_fn)
R = box_size * random.uniform(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
grid_force_fn = smap.grid(force_fn, box_size, cell_size, R)
species = np.zeros((PARTICLE_COUNT, ), dtype=np.int64)
self.assertAllClose(np.array(force_fn(R, species, 1), dtype=dtype),
grid_force_fn(R), True)
args = parser.parse_args()
edge_length = pow(args.parts / args.dense, 1.0 / 3.0)
# edge_length*=2
spatial_dimension = 3
box_size = onp.asarray([edge_length] * spatial_dimension)
displacement_fn, shift_fn = space.periodic(box_size)
key = random.PRNGKey(0)
R = random.uniform(key, (args.parts, spatial_dimension),
minval=0.0,
maxval=box_size[0],
dtype=np.float64)
# print(R)
energy_fn = energy.lennard_jones_pair(displacement_fn)
print('E = {}'.format(energy_fn(R)))
force_fn = quantity.force(energy_fn)
print('Total Squared Force = {}'.format(np.sum(force_fn(R)**2)))
init, apply = simulate.nve(energy_fn, shift_fn, args.time / args.steps)
apply = jit(apply)
state = init(key, R, velocity_scale=0.0)
PE = []
KE = []
print_every = args.log
old_time = time.perf_counter()
print('Step\tKE\tPE\tTotal Energy\ttime/step')
print('----------------------------------------')
for i in range(args.steps // print_every):
state = lax.fori_loop(0, print_every, lambda _, state: apply(state), state)
def npt_nose_hoover(energy_fn: Callable[..., Array],
shift_fn: ShiftFn,
dt: float,
pressure: float,
kT: float,
barostat_kwargs: Optional[Dict] = None,
thermostat_kwargs: Optional[Dict] = None) -> Simulator:
"""Simulation in the NPT ensemble using a pair of Nose Hoover Chains.
Samples from the canonical ensemble in which the number of particles (N),
the system pressure (P), and the temperature (T) are held constant. We use a
pair of Nose Hoover Chains (NHC) described in [1, 2, 3] coupled to the
barostat and the thermostat respectively. We follow the direct translation
method outlined in [3] and the interested reader might want to look at that
paper as a reference.
Args:
energy_fn: A function that produces either an energy from a set of particle
positions specified as an ndarray of shape [n, spatial_dimension].
shift_fn: A function that displaces positions, R, by an amount dR. Both R
and dR should be ndarrays of shape [n, spatial_dimension].
dt: Floating point number specifying the timescale (step size) of the
simulation.
pressure: Floating point number specifying the target pressure. To update
the pressure dynamically during a simulation one should pass `pressure`
as a keyword argument to the step function.
kT: Floating point number specifying the temperature in units of Boltzmann
constant. To update the temperature dynamically during a simulation one
should pass `kT` as a keyword argument to the step function.
barostat_kwargs: A dictionary of keyword arguments passed to the barostat
NHC. Any parameters not set are drawn from a relatively robust default
set.
thermostat_kwargs: A dictionary of keyword arguments passed to the
thermostat NHC. Any parameters not set are drawn from a relatively robust
default set.
Returns:
See above.
[1] Martyna, Glenn J., Michael L. Klein, and Mark Tuckerman.
"Nose-Hoover chains: The canonical ensemble via continuous dynamics."
The Journal of chemical physics 97, no. 4 (1992): 2635-2643.
[2] Martyna, Glenn, Mark Tuckerman, Douglas J. Tobias, and Michael L. Klein.
"Explicit reversible integrators for extended systems dynamics."
Molecular Physics 87. (1998) 1117-1157.
[3] Tuckerman, Mark E., Jose Alejandre, Roberto Lopez-Rendon,
Andrea L. Jochim, and Glenn J. Martyna.
"A Liouville-operator derived measure-preserving integrator for molecular
dynamics simulations in the isothermal-isobaric ensemble."
Journal of Physics A: Mathematical and General 39, no. 19 (2006): 5629.
"""
t = f32(dt)
dt_2 = f32(dt / 2)
force_fn = quantity.force(energy_fn)
barostat_kwargs = default_nhc_kwargs(1000 * dt, barostat_kwargs)
barostat = nose_hoover_chain(dt, **barostat_kwargs)
thermostat_kwargs = default_nhc_kwargs(100 * dt, thermostat_kwargs)
thermostat = nose_hoover_chain(dt, **thermostat_kwargs)
def init_fn(key, R, box, mass=f32(1.0), **kwargs):
N, dim = R.shape
_kT = kT if 'kT' not in kwargs else kwargs['kT']
mass = quantity.canonicalize_mass(mass)
V = jnp.sqrt(_kT / mass) * random.normal(key, R.shape, dtype=R.dtype)
V = V - jnp.mean(V * mass, axis=0, keepdims=True) / mass
KE = quantity.kinetic_energy(V, mass)
# The box position is defined via pos = (1 / d) log V / V_0.
zero = jnp.zeros((), dtype=R.dtype)
one = jnp.ones((), dtype=R.dtype)
box_position = zero
box_velocity = zero
box_mass = dim * (N + 1) * kT * barostat_kwargs['tau']**2 * one
KE_box = quantity.kinetic_energy(box_velocity, box_mass)
if jnp.isscalar(box) or box.ndim == 0:
# TODO(schsam): This is necessary because of JAX issue #5849.
box = jnp.eye(R.shape[-1]) * box
return NPTNoseHooverState(R, V, force_fn(R, box=box, **kwargs), mass,
box, box_position, box_velocity, box_mass,
barostat.initialize(1, KE_box, _kT),
thermostat.initialize(R.size, KE, _kT)) # pytype: disable=wrong-arg-count
def update_box_mass(state, kT):
N, dim = state.position.shape
dtype = state.position.dtype
box_mass = jnp.array(dim * (N + 1) * kT * state.barostat.tau**2, dtype)
return dataclasses.replace(state, box_mass=box_mass)
def box_force(alpha, vol, box_fn, position, velocity, mass, force,
pressure, **kwargs):
N, dim = position.shape
def U(vol):
return energy_fn(position, box=box_fn(vol), **kwargs)
dUdV = grad(U)
KE2 = util.high_precision_sum(velocity**2 * mass)
R = space.transform(box_fn(vol), position)
RdotF = util.high_precision_sum(R * force)
return alpha * KE2 + RdotF - dim * vol * dUdV(
vol) - pressure * vol * dim
def sinhx_x(x):
"""Taylor series for sinh(x) / x as x -> 0."""
return 1 + x**2 / 6 + x**4 / 120
def exp_iL1(box, R, V, V_b, **kwargs):
x = V_b * dt
x_2 = x / 2
sinhV = sinhx_x(x_2) # jnp.sinh(x_2) / x_2
return shift_fn(R * jnp.exp(x),
dt * V * jnp.exp(x_2) * sinhV,
box=box,
**kwargs) # pytype: disable=wrong-keyword-args
def exp_iL2(alpha, V, A, V_b):
x = alpha * V_b * dt_2
x_2 = x / 2
sinhV = sinhx_x(x_2) # jnp.sinh(x_2) / x_2
return V * jnp.exp(-x) + dt_2 * A * sinhV * jnp.exp(-x_2)
def inner_step(state, **kwargs):
_pressure = kwargs.pop('pressure', pressure)
R, V, M, F = state.position, state.velocity, state.mass, state.force
R_b, V_b, M_b = state.box_position, state.box_velocity, state.box_mass
N, dim = R.shape
vol, box_fn = _npt_box_info(state)
alpha = 1 + 1 / N
G_e = box_force(alpha, vol, box_fn, R, V, M, F, _pressure, **kwargs)
V_b = V_b + dt_2 * G_e / M_b
V = exp_iL2(alpha, V, F / M, V_b)
R_b = R_b + V_b * dt
state = dataclasses.replace(state, box_position=R_b)
vol, box_fn = _npt_box_info(state)
box = box_fn(vol)
R = exp_iL1(box, R, V, V_b)
F = force_fn(R, box=box, **kwargs)
V = exp_iL2(alpha, V, F / M, V_b)
G_e = box_force(alpha, vol, box_fn, R, V, M, F, _pressure, **kwargs)
V_b = V_b + dt_2 * G_e / M_b
return dataclasses.replace(state,
position=R,
velocity=V,
mass=M,
force=F,
box_position=R_b,
box_velocity=V_b,
box_mass=M_b)
def apply_fn(state, **kwargs):
S = state
_kT = kT if 'kT' not in kwargs else kwargs['kT']
bc = barostat.update_mass(S.barostat, _kT)
tc = thermostat.update_mass(S.thermostat, _kT)
S = update_box_mass(S, _kT)
V_b, bc = barostat.half_step(S.box_velocity, bc, _kT)
V, tc = thermostat.half_step(S.velocity, tc, _kT)
S = dataclasses.replace(S, velocity=V, box_velocity=V_b)
S = inner_step(S, **kwargs)
KE = quantity.kinetic_energy(S.velocity, S.mass)
tc = dataclasses.replace(tc, kinetic_energy=KE)
KE_box = quantity.kinetic_energy(S.box_velocity, S.box_mass)
bc = dataclasses.replace(bc, kinetic_energy=KE_box)
V, tc = thermostat.half_step(S.velocity, tc, _kT)
V_b, bc = barostat.half_step(S.box_velocity, bc, _kT)
S = dataclasses.replace(S,
thermostat=tc,
barostat=bc,
velocity=V,
box_velocity=V_b)
return S
return init_fn, apply_fn
