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