def test_delta_r():
"""assert that
* delta_r(ri, rj, box) == - delta_r(rj, ri, box)
* delta_r agrees with jit(delta_r)
* jit(norm(delta_r)) symmetric
on a few random inputs of varying size
"""
@jit
def _distances(ri, rj, box):
return np.linalg.norm(delta_r(ri, rj, box), axis=1)
for _ in range(5):
n_atoms = onp.random.randint(50, 1000)
dim = onp.random.randint(3, 5)
ri, rj = onp.random.randn(2, n_atoms, dim)
box = np.eye(dim)
dr_ij_1 = delta_r(ri, rj, box)
dr_ji_1 = delta_r(rj, ri, box)
onp.testing.assert_allclose(dr_ij_1, -dr_ji_1)
dr_ij_2 = jit(delta_r)(ri, rj, box)
dr_ji_2 = jit(delta_r)(rj, ri, box)
onp.testing.assert_allclose(dr_ij_1, dr_ij_2)
onp.testing.assert_allclose(dr_ji_1, dr_ji_2)
dij = _distances(ri, rj, box)
dji = _distances(rj, ri, box)
onp.testing.assert_allclose(dij, dji)
def harmonic_angle(conf, params, box, angle_idxs, param_idxs, cos_angles=True):
"""
Compute the harmonic bond energy given a collection of molecules.
This implements a harmonic angle potential: V(t) = k*(t - t0)^2 or V(t) = k*(cos(t)-cos(t0))^2
Parameters:
-----------
conf: shape [num_atoms, 3] np.array
atomic coordinates
params: shape [num_params,] np.array
unique parameters
box: shape [3, 3] np.array
periodic boundary vectors, if not None
angle_idxs: shape [num_angles, 3] np.array
each element (a, b, c) is a unique angle in the conformation. atom b is defined
to be the middle atom.
param_idxs: shape [num_angles, 2] np.array
each element (k_idx, t_idx) maps into params for angle constants and ideal angles
cos_angles: True (default)
if True, then this instead implements V(t) = k*(cos(t)-cos(t0))^2. This is far more
numerically stable when the angle is pi.
"""
ci = conf[angle_idxs[:, 0]]
cj = conf[angle_idxs[:, 1]]
ck = conf[angle_idxs[:, 2]]
kas = params[param_idxs[:, 0]]
a0s = params[param_idxs[:, 1]]
vij = delta_r(ci, cj, box)
vjk = delta_r(ck, cj, box)
top = np.sum(np.multiply(vij, vjk), -1)
bot = np.linalg.norm(vij, axis=-1) * np.linalg.norm(vjk, axis=-1)
tb = top / bot
# (ytz): we used the squared version so that we make this energy being strictly positive
if cos_angles:
energies = kas / 2 * np.power(tb - np.cos(a0s), 2)
else:
angle = np.arccos(tb)
energies = kas / 2 * np.power(angle - a0s, 2)
return np.sum(energies, -1) # reduce over all angles
def signed_torsion_angle(ci, cj, ck, cl):
"""
Batch compute the signed angle of a torsion angle. The torsion angle
between two planes should be periodic but not necessarily symmetric.
Parameters
----------
ci: shape [num_torsions, 3] np.array
coordinates of the 1st atom in the 1-4 torsion angle
cj: shape [num_torsions, 3] np.array
coordinates of the 2nd atom in the 1-4 torsion angle
ck: shape [num_torsions, 3] np.array
coordinates of the 3rd atom in the 1-4 torsion angle
cl: shape [num_torsions, 3] np.array
coordinates of the 4th atom in the 1-4 torsion angle
Returns
-------
shape [num_torsions,] np.array
array of torsion angles.
"""
# Taken from the wikipedia arctan2 implementation:
# https://en.wikipedia.org/wiki/Dihedral_angle
# We use an identical but numerically stable arctan2
# implementation as opposed to the OpenMM energy function to
# avoid asingularity when the angle is zero.
rij = delta_r(cj, ci)
rkj = delta_r(cj, ck)
rkl = delta_r(cl, ck)
n1 = np.cross(rij, rkj)
n2 = np.cross(rkj, rkl)
lhs = np.linalg.norm(n1, axis=-1)
rhs = np.linalg.norm(n2, axis=-1)
bot = lhs * rhs
y = np.sum(np.multiply(np.cross(n1, n2),
rkj / np.linalg.norm(rkj, axis=-1, keepdims=True)),
axis=-1)
x = np.sum(np.multiply(n1, n2), -1)
return np.arctan2(y, x)
def nonbonded_block(xi, xj, box, params_i, params_j, beta, cutoff):
"""
This is a modified version of nonbonded_v3 that computes a block of
interactions between two sets of particles x_i and x_j. It is assumed that
there are no exclusions between the two particle sets. Typical use cases
include computing the interaction energy between the environment and a ligand.
This is mainly used for testing, as it does not support 4D decoupling or
alchemical semantics yet.
Parameters
----------
xi : (N,3) np.ndarray
Coordinates
xj : (N,3) np.ndarray
Coordinates
box : Optional 3x3 np.array
Periodic boundary conditions
beta : float
the charge product q_ij will be multiplied by erfc(beta*d_ij)
cutoff : Optional float
a pair of particles (i,j) will be considered non-interacting if the distance d_ij
between their 3D coordinates exceeds cutoff
Returns
-------
scalar
Interaction energy
"""
ri = np.expand_dims(xi, axis=1)
rj = np.expand_dims(xj, axis=0)
d2ij = np.sum(np.power(delta_r(ri, rj, box), 2), axis=-1)
dij = np.sqrt(d2ij)
sig_i = np.expand_dims(params_i[:, 1], axis=1)
sig_j = np.expand_dims(params_j[:, 1], axis=0)
eps_i = np.expand_dims(params_i[:, 2], axis=1)
eps_j = np.expand_dims(params_j[:, 2], axis=0)
sig_ij = sig_i + sig_j
eps_ij = eps_i * eps_j
qi = np.expand_dims(params_i[:, 0], axis=1)
qj = np.expand_dims(params_j[:, 0], axis=0)
qij = np.multiply(qi, qj)
es = direct_space_pme(dij, qij, beta)
lj = lennard_jones(dij, sig_ij, eps_ij)
nrg = np.where(dij > cutoff, 0, es + lj)
return np.sum(nrg)
def _distances(ri, rj, box):
return np.linalg.norm(delta_r(ri, rj, box), axis=1)