def soft_sphere_neighbor_list(displacement_or_metric,
box_size,
species=None,
sigma=1.0,
epsilon=1.0,
alpha=2.0,
dr_threshold=0.2):
"""Convenience wrapper to compute soft spheres using a neighbor list."""
sigma = np.array(sigma, dtype=f32)
epsilon = np.array(epsilon, dtype=f32)
alpha = np.array(alpha, dtype=f32)
list_cutoff = f32(np.max(sigma))
dr_threshold = f32(list_cutoff * dr_threshold)
neighbor_fn = partition.neighbor_list(displacement_or_metric, box_size,
list_cutoff, dr_threshold)
energy_fn = smap.pair_neighbor_list(
soft_sphere,
space.canonicalize_displacement_or_metric(displacement_or_metric),
species=species,
sigma=sigma,
epsilon=epsilon,
alpha=alpha)
return neighbor_fn, energy_fn
def periodic_displacement(side, dR):
"""Wraps displacement vectors into a hypercube.
Args:
side: Specification of hypercube size. Either,
(a) float if all sides have equal length.
(b) ndarray(spatial_dim) if sides have different lengths.
dR: Matrix of displacements; ndarray(shape=[..., spatial_dim]).
Returns:
Matrix of wrapped displacements; ndarray(shape=[..., spatial_dim]).
"""
return np.mod(dR + side * f32(0.5), side) - f32(0.5) * side
def load_lammps_eam_parameters(f):
"""Reads EAM parameters from a LAMMPS file and returns relevant spline fits.
This function reads single-element EAM potential fit parameters from a file
in DYNAMO funcl format. In summary, the file contains:
Line 1-3: comments,
Line 4: Number of elements and the element type,
Line 5: The number of charge values that the embedding energy is evaluated
on (num_drho), interval between the charge values (drho), the number of
distances the pairwise energy and the charge density is evaluated on (num_dr),
the interval between these distances (dr), and the cutoff distance (cutoff).
The lines that come after are the embedding function evaluated on num_drho
charge values, charge function evaluated at num_dr distance values, and
pairwise energy evaluated at num_dr distance values. Note that the pairwise
energy is multiplied by distance (in units of eV x Angstroms). For more
details of the DYNAMO file format, see:
https://sites.google.com/a/ncsu.edu/cjobrien/tutorials-and-guides/eam
Args:
f: File handle for the EAM parameters text file.
Returns:
charge_fn: A function that takes an ndarray of shape [n, m] of distances
between particles and returns a matrix of charge contributions.
embedding_fn: Function that takes an ndarray of shape [n] of charges and
returns an ndarray of shape [n] of the energy cost of embedding an atom
into the charge.
pairwise_fn: A function that takes an ndarray of shape [n, m] of distances
and returns an ndarray of shape [n, m] of pairwise energies.
cutoff: Cutoff distance for the embedding_fn and pairwise_fn.
"""
raw_text = f.read().split('\n')
if 'setfl' not in raw_text[0]:
raise ValueError(
'File format is incorrect, expected LAMMPS setfl format.')
temp_params = raw_text[4].split()
num_drho, num_dr = int(temp_params[0]), int(temp_params[2])
drho, dr, cutoff = float(temp_params[1]), float(temp_params[3]), float(
temp_params[4])
data = np.array(map(float, raw_text[6:-1]))
embedding_fn = spline(data[:num_drho], drho)
charge_fn = spline(data[num_drho:num_drho + num_dr], dr)
# LAMMPS EAM parameters file lists pairwise energies after multiplying by
# distance, in units of eV*Angstrom. We divide the energy by distance below,
distances = np.arange(num_dr) * dr
# Prevent dividing by zero at zero distance, which will not
# affect the calculation
distances = np.where(distances == 0, f32(0.001), distances)
pairwise_fn = spline(
data[num_dr + num_drho:num_drho + f32(2) * num_dr] / distances, dr)
return charge_fn, embedding_fn, pairwise_fn, cutoff
def bks_neighbor_list(displacement_or_metric,
box_size,
species,
Q_sq,
exp_coeff,
exp_decay,
attractive_coeff,
repulsive_coeff,
coulomb_alpha,
cutoff,
dr_threshold=0.8):
Q_sq = np.array(Q_sq, f32)
exp_coeff = np.array(exp_coeff, f32)
exp_decay = np.array(exp_decay, f32)
attractive_coeff = np.array(attractive_coeff, f32)
repulsive_coeff = np.array(repulsive_coeff, f32)
dr_threshold = f32(dr_threshold)
neighbor_fn = partition.neighbor_list(displacement_or_metric, box_size,
cutoff, dr_threshold)
energy_fn = smap.pair_neighbor_list(
bks,
space.canonicalize_displacement_or_metric(displacement_or_metric),
species=species,
Q_sq=Q_sq,
exp_coeff=exp_coeff,
exp_decay=exp_decay,
attractive_coeff=attractive_coeff,
repulsive_coeff=repulsive_coeff,
coulomb_alpha=coulomb_alpha,
cutoff=cutoff)
return neighbor_fn, energy_fn
def multiplicative_isotropic_cutoff(fn, r_onset, r_cutoff):
"""Takes an isotropic function and constructs a truncated function.
Given a function f:R -> R, we construct a new function f':R -> R such that
f'(r) = f(r) for r < r_onset, f'(r) = 0 for r > r_cutoff, and f(r) is C^1
everywhere. To do this, we follow the approach outlined in HOOMD Blue [1]
(thanks to Carl Goodrich for the pointer). We construct a function S(r) such
that S(r) = 1 for r < r_onset, S(r) = 0 for r > r_cutoff, and S(r) is C^1.
Then f'(r) = S(r)f(r).
Args:
fn: A function that takes an ndarray of distances of shape [n, m] as well
as varargs.
r_onset: A float specifying the onset radius of deformation.
r_cutoff: A float specifying the cutoff radius.
Returns:
A new function with the same signature as fn, with the properties outlined
above.
[1] HOOMD Blue documentation. Accessed on 05/31/2019.
https://hoomd-blue.readthedocs.io/en/stable/module-md-pair.html#hoomd.md.pair.pair
"""
r_c = r_cutoff ** f32(2)
r_o = r_onset ** f32(2)
def smooth_fn(dr):
r = dr ** f32(2)
return np.where(
dr < r_onset,
f32(1),
np.where(
dr < r_cutoff,
(r_c - r) ** f32(2) * (r_c + f32(2) * r - f32(3) * r_o) / (
r_c - r_o) ** f32(3),
f32(0)
)
)
@wraps(fn)
def cutoff_fn(dr, *args, **kwargs):
return smooth_fn(dr) * fn(dr, *args, **kwargs)
return cutoff_fn
def energy(R, **kwargs):
dr = metric(R, R, **kwargs)
total_charge = smap._high_precision_sum(charge_fn(dr), axis=1)
embedding_energy = embedding_fn(total_charge)
pairwise_energy = smap._high_precision_sum(
smap._diagonal_mask(pairwise_fn(dr)), axis=1) / f32(2.0)
return smap._high_precision_sum(embedding_energy + pairwise_energy,
axis=axis)
def morse(dr, sigma=1.0, epsilon=5.0, alpha=5.0, **unused_kwargs):
"""Morse interaction between particles with a minimum at r0.
Args:
dr: An ndarray of shape [n, m] of pairwise distances between particles.
sigma: Distance between particles where the energy has a minimum. Should
either be a floating point scalar or an ndarray whose shape is [n, m].
epsilon: Interaction energy scale. Should either be a floating point scalar
or an ndarray whose shape is [n, m].
alpha: Range parameter. Should either be a floating point scalar or an
ndarray whose shape is [n, m].
unused_kwargs: Allows extra data (e.g. time) to be passed to the energy.
Returns:
Matrix of energies of shape [n, m].
"""
check_kwargs_time_dependence(unused_kwargs)
U = epsilon * (f32(1) - np.exp(-alpha * (dr - sigma)))**f32(2) - epsilon
return np.nan_to_num(np.array(U, dtype=dr.dtype))
def smooth_fn(dr):
r = dr**f32(2)
return np.where(
dr < r_onset, f32(1),
np.where(dr < r_cutoff, (r_c - r)**f32(2) *
(r_c + f32(2) * r - f32(3) * r_o) / (r_c - r_o)**f32(3),
f32(0)))
def lennard_jones(dr, sigma=1, epsilon=1, **unused_kwargs):
"""Lennard-Jones interaction between particles with a minimum at sigma.
Args:
dr: An ndarray of shape [n, m] of pairwise distances between particles.
sigma: Distance between particles where the energy has a minimum. Should
either be a floating point scalar or an ndarray whose shape is [n, m].
epsilon: Interaction energy scale. Should either be a floating point scalar
or an ndarray whose shape is [n, m].
unused_kwargs: Allows extra data (e.g. time) to be passed to the energy.
Returns:
Matrix of energies of shape [n, m].
"""
check_kwargs_time_dependence(unused_kwargs)
dr = (sigma / dr)**f32(2)
idr6 = dr**f32(3)
idr12 = idr6**f32(2)
# TODO(schsam): This seems potentially dangerous. We should do ErrorChecking
# here.
return np.nan_to_num(f32(4) * epsilon * (idr12 - idr6))
def soft_sphere(dr, sigma=1, epsilon=1, alpha=2, **unused_kwargs):
"""Finite ranged repulsive interaction between soft spheres.
Args:
dr: An ndarray of shape [n, m] of pairwise distances between particles.
sigma: Particle diameter. Should either be a floating point scalar or an
ndarray whose shape is [n, m].
epsilon: Interaction energy scale. Should either be a floating point scalar
or an ndarray whose shape is [n, m].
alpha: Exponent specifying interaction stiffness. Should either be a float
point scalar or an ndarray whose shape is [n, m].
unused_kwargs: Allows extra data (e.g. time) to be passed to the energy.
Returns:
Matrix of energies whose shape is [n, m].
"""
check_kwargs_time_dependence(unused_kwargs)
dr = dr / sigma
U = epsilon * np.where(dr < 1.0,
f32(1.0) / alpha * (f32(1.0) - dr)**alpha, f32(0.0))
return U
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 displacement_fn(Ra, Rb, **kwargs):
_box, _inv_box = box, inv_box
if 'box' in kwargs:
_box = kwargs['box']
if not fractional_coordinates:
_inv_box = inverse(_box)
if 'new_box' in kwargs:
_box = kwargs['new_box']
if not fractional_coordinates:
Ra = transform(_inv_box, Ra)
Rb = transform(_inv_box, Rb)
dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb))
return transform(_box, dR)
def displacement_fn(Ra, Rb, perturbation=None, **kwargs):
_box, _inv_box = box, inv_box
if 'box' in kwargs:
_box = kwargs['box']
if not fractional_coordinates:
_inv_box = inverse(_box)
if 'new_box' in kwargs:
_box = kwargs['new_box']
if not fractional_coordinates:
Ra = transform(_inv_box, Ra)
Rb = transform(_inv_box, Rb)
dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb))
dR = transform(_box, dR)
if perturbation is not None:
dR = raw_transform(perturbation, dR)
return dR
def shift(R, dR, **unused_kwargs): check_kwargs_time_dependence(unused_kwargs) return periodic_shift(f32(1.0), R, transform(T_inv, dR))
def displacement(Ra, Rb, **unused_kwargs): check_kwargs_time_dependence(unused_kwargs) dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb)) return transform(T, dR)
def shift(R, dR, t=None, **unused_kwargs): _check_time_dependence(t) check_kwargs_empty(unused_kwargs) return periodic_shift(f32(1.0), R, transform(_small_inverse(T(t)), dR))
def displacement(Ra, Rb, t=None, **unused_kwargs): _check_time_dependence(t) check_kwargs_empty(unused_kwargs) dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb)) return transform(T(t), dR)
def u(R, dR):
if wrapped:
return periodic_shift(f32(1.0), R, dR)
return R + dR
def shift(R: Array, dR: Array, **kwargs) -> Array:
return periodic_shift(f32(1.0),
R,
transform(_small_inverse(T(**kwargs)), dR))
def displacement(Ra: Array, Rb: Array, **unused_kwargs) -> Array: dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb)) return transform(T, dR)
def shift(R: Array, dR: Array, **unused_kwargs) -> Array: return periodic_shift(f32(1.0), R, transform(T_inv, dR))
