def simple_spring(dr, length=1, epsilon=1, alpha=2, **unused_kwargs):
"""Isotropic spring potential with a given rest length.
We define `simple_spring` to be a generalized Hookian spring with
agreement when alpha = 2.
"""
check_kwargs_time_dependence(unused_kwargs)
return epsilon / alpha * (dr - length)**alpha
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 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)
idr = (sigma / dr)
idr = idr * idr
idr6 = idr * idr * idr
idr12 = idr6 * idr6
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 shift(R, dR, **unused_kwargs): check_kwargs_time_dependence(unused_kwargs) return R + transform(T_inv, 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_fn(R, dR, **unused_kwargs): check_kwargs_time_dependence(unused_kwargs) return R + dR
def shift_fn(R, dR, **unused_kwargs): check_kwargs_time_dependence(unused_kwargs) return periodic_shift(side, R, dR)
def displacement_fn(Ra, Rb, **unused_kwargs): check_kwargs_time_dependence(unused_kwargs) return periodic_displacement(side, pairwise_displacement(Ra, Rb))
