def simple_energy(conf, box, charge_params, exclusion_idxs, charge_scales,
beta, cutoff):
"""
Numerically stable implementation of the pairwise term:
eij = qi*qj/dij
"""
charges = charge_params
qi = np.expand_dims(charges, 0) # (1, N)
qj = np.expand_dims(charges, 1) # (N, 1)
qij = np.multiply(qi, qj)
ri = np.expand_dims(conf, 0)
rj = np.expand_dims(conf, 1)
dij = distance(ri, rj, box)
# (ytz): trick used to avoid nans in the diagonal due to the 1/dij term.
keep_mask = 1 - np.eye(conf.shape[0])
qij = np.where(keep_mask, qij, np.zeros_like(qij))
dij = np.where(keep_mask, dij, np.zeros_like(dij))
# funny enough lim_{x->0} erfc(x)/x = 0
eij = np.where(keep_mask,
qij * erfc(beta * dij) / dij,
np.zeros_like(dij)) # zero out diagonals
# print(dij)
if cutoff is not None:
# sw = switch_fn(dij, cutoff)
# eij = eij*sw
eij = np.where(dij > cutoff, np.zeros_like(eij), eij)
src_idxs = exclusion_idxs[:, 0]
dst_idxs = exclusion_idxs[:, 1]
ri = conf[src_idxs]
rj = conf[dst_idxs]
dij = distance(ri, rj, box)
qi = charges[src_idxs]
qj = charges[dst_idxs]
qij = np.multiply(qi, qj)
scale_ij = charge_scales
eij_exc = scale_ij * qij * erfc(beta * dij) / dij
if cutoff is not None:
# sw = switch_fn(dij, cutoff)
# eij_exc = eij_exc*sw
eij_exc = np.where(dij > cutoff, np.zeros_like(eij_exc), eij_exc)
eij_exc = np.where(src_idxs == dst_idxs, np.zeros_like(eij_exc),
eij_exc)
return np.sum(eij / 2) - np.sum(eij_exc)
def dsf_coulomb(r, Q_sq, alpha=0.25, cutoff=8.0):
qqr2e = 332.06371 #coulmbic conversion factor:1/(4*pi*epo)
cutoffsq = cutoff * cutoff
erfcc = erfc(alpha * cutoff)
erfcd = np.exp(-alpha * alpha * cutoffsq)
f_shift = -(erfcc / cutoffsq +
2.0 / np.sqrt(np.pi) * alpha * erfcd / cutoff)
e_shift = erfcc / cutoff - f_shift * cutoff
coulomb_en = qqr2e * Q_sq / r * (erfc(alpha * r) - r * e_shift -
r**2 * f_shift)
return np.where(r < cutoff, coulomb_en, 0.0)
def dsf_coulomb(r: Array,
Q_sq: Array,
alpha: Array=0.25,
cutoff: float=8.0) -> Array:
"""Damped-shifted-force approximation of the coulombic interaction."""
qqr2e = 332.06371 # Coulmbic conversion factor: 1/(4*pi*epo).
cutoffsq = cutoff*cutoff
erfcc = erfc(alpha*cutoff)
erfcd = np.exp(-alpha*alpha*cutoffsq)
f_shift = -(erfcc/cutoffsq + 2.0/np.sqrt(np.pi)*alpha*erfcd/cutoff)
e_shift = erfcc/cutoff - f_shift*cutoff
coulomb_en = qqr2e*Q_sq/r * (erfc(alpha*r) - r*e_shift - r**2*f_shift)
return np.where(r < cutoff, coulomb_en, 0.0)
def direct_space_pme(dij, qij, beta):
"""Direct-space contribution from eq 2 of:
Darden, York, Pedersen, 1993, J. Chem. Phys.
"Particle mesh Ewald: An N log(N) method for Ewald sums in large systems"
https://aip.scitation.org/doi/abs/10.1063/1.470117
"""
return qij * erfc(beta * dij) / dij
def burgers(x, t, v, A):
R = A / (2 * v)
z = x / jnp.sqrt(4 * v * t)
u = (jnp.sqrt(v / (jnp.pi * t)) * ((jnp.exp(R) - 1) * jnp.exp(-(z**2))) /
(1 + (jnp.exp(R) - 1) / 2 * erfc(z)))
return u
def _bks_silica_self(Q_sq: Array, alpha: Array, cutoff: float) -> Array: cutoffsq = cutoff * cutoff erfcc = erfc(alpha * cutoff) erfcd = np.exp(-alpha * alpha * cutoffsq) f_shift = -(erfcc / cutoffsq + 2.0 / np.sqrt(np.pi) * alpha * erfcd / cutoff) e_shift = erfcc / cutoff - f_shift * cutoff qqr2e = 332.06371 #kcal/mol #coulmbic conversion factor:1/(4*pi*epo) return -(e_shift / 2.0 + alpha / np.sqrt(np.pi)) * Q_sq * qqr2e
def validate_coulomb_cutoff(cutoff=1.0, beta=2.0, threshold=1e-2):
"""check whether f(r) = erfc(beta * r) <= threshold at r = cutoff
following https://github.com/proteneer/timemachine/pull/424#discussion_r629678467"""
if erfc(beta * cutoff) > threshold:
print(
UserWarning(
f"erfc(beta * cutoff) = {erfc(beta * cutoff)} > threshold = {threshold}"
))
def e_self(Q_sq, alpha=0.25, cutoff=8.0):
cutoffsq = cutoff * cutoff
erfcc = erfc(alpha * cutoff)
erfcd = np.exp(-alpha * alpha * cutoffsq)
f_shift = -(erfcc / cutoffsq +
2.0 / np.sqrt(np.pi) * alpha * erfcd / cutoff)
e_shift = erfcc / cutoff - f_shift * cutoff
qqr2e = 332.06371 #kcal/mol #coulmbic conversion factor:1/(4*pi*epo)
return -(e_shift / 2.0 + alpha / np.sqrt(np.pi)) * Q_sq * qqr2e
def function(self, x):
mean, var = self.model.predict(x)
s = jnp.sqrt(var)
u = (self.fmin - mean + self.jitter) / s
phi = jnp.exp(-0.5 * u**2) / jnp.sqrt(2 * jnp.pi)
Phi = 0.5 * erfc(-u / jnp.sqrt(2))
return -jnp.sum(s * (u * Phi + phi))
def logphi(z):
"""
Calculate the log Gaussian CDF, used for closed form moment matching when the EP power is 1,
logΦ(z) = log[(1 + erf(z / √2)) / 2]
for erf(z) = (2/√π) ∫ exp(-x²) dx, where the integral is over [0, z]
and its derivative w.r.t. z
dlogΦ(z)/dz = 𝓝(z|0,1) / Φ(z)
:param z: input value, typically z = (my) / √(1 + v) [scalar]
:return:
lp: logΦ(z) [scalar]
dlp: dlogΦ(z)/dz [scalar]
"""
z = np.real(z)
# erfc(z) = 1 - erf(z) is the complementary error function
lp = np.log(erfc(-z / np.sqrt(2.0)) / 2.0) # log Φ(z)
dlp = np.exp(-z * z / 2.0 - lp) / np.sqrt(2.0 * pi) # derivative w.r.t. z
return lp, dlp
def ewald_energy(conf, box, charges, scale_matrix, cutoff, alpha, kmax):
eij = pairwise_energy(conf, box, charges, cutoff)
assert cutoff is not None
# 1. Assume scale matrix is not used at all (no exceptions, no exclusions)
# 1a. Direct Space
eij_direct = eij * erfc(alpha*eij)
eij_direct = ONE_4PI_EPS0*np.sum(eij_direct)/2
# 1b. Reciprocal Space
eij_recip = reciprocal_energy(conf, box, charges, alpha, kmax)
# 2. Remove over estimated scale matrix contribution scaled by erf
eij_offset = (1-scale_matrix) * eij
eij_offset *= erf(alpha*eij_offset)
eij_offset = ONE_4PI_EPS0*np.sum(eij_offset)/2
return eij_direct + eij_recip - eij_offset - self_energy(conf, charges, alpha)
def nonbonded_v3(
conf,
params,
box,
lamb,
charge_rescale_mask,
lj_rescale_mask,
beta,
cutoff,
lambda_plane_idxs,
lambda_offset_idxs,
runtime_validate=True,
):
"""Lennard-Jones + Coulomb, with a few important twists:
* distances are computed in 4D, controlled by lambda, lambda_plane_idxs, lambda_offset_idxs
* each pairwise LJ and Coulomb term can be multiplied by an adjustable rescale_mask parameter
* Coulomb terms are multiplied by erfc(beta * distance)
Parameters
----------
conf : (N, 3) or (N, 4) np.array
3D or 4D coordinates
if 3D, will be converted to 4D using (x,y,z) -> (x,y,z,w)
where w = cutoff * (lambda_plane_idxs + lambda_offset_idxs * lamb)
params : (N, 3) np.array
columns [charges, sigmas, epsilons], one row per particle
box : Optional 3x3 np.array
lamb : float
charge_rescale_mask : (N, N) np.array
the Coulomb contribution of pair (i,j) will be multiplied by charge_rescale_mask[i,j]
lj_rescale_mask : (N, N) np.array
the Lennard-Jones contribution of pair (i,j) will be multiplied by lj_rescale_mask[i,j]
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 4D coordinates exceeds cutoff
lambda_plane_idxs : Optional (N,) np.array
lambda_offset_idxs : Optional (N,) np.array
runtime_validate: bool
check whether beta is compatible with cutoff
(if True, this function will currently not play nice with Jax JIT)
TODO: is there a way to conditionally print a runtime warning inside
of a Jax JIT-compiled function, without triggering a Jax ConcretizationTypeError?
Returns
-------
energy : float
References
----------
* Rodinger, Howell, Pomès, 2005, J. Chem. Phys. "Absolute free energy calculations by thermodynamic integration in four spatial
dimensions" https://aip.scitation.org/doi/abs/10.1063/1.1946750
* Darden, York, Pedersen, 1993, J. Chem. Phys. "Particle mesh Ewald: An N log(N) method for Ewald sums in large
systems" https://aip.scitation.org/doi/abs/10.1063/1.470117
* Coulomb interactions are treated using the direct-space contribution from eq 2
"""
if runtime_validate:
assert (charge_rescale_mask == charge_rescale_mask.T).all()
assert (lj_rescale_mask == lj_rescale_mask.T).all()
N = conf.shape[0]
if conf.shape[-1] == 3:
conf = convert_to_4d(conf, lamb, lambda_plane_idxs, lambda_offset_idxs,
cutoff)
# make 4th dimension of box large enough so its roughly aperiodic
if box is not None:
if box.shape[-1] == 3:
box_4d = np.eye(4) * 1000
box_4d = index_update(box_4d, index[:3, :3], box)
else:
box_4d = box
else:
box_4d = None
box = box_4d
charges = params[:, 0]
sig = params[:, 1]
eps = params[:, 2]
sig_i = np.expand_dims(sig, 0)
sig_j = np.expand_dims(sig, 1)
sig_ij = sig_i + sig_j
eps_i = np.expand_dims(eps, 0)
eps_j = np.expand_dims(eps, 1)
eps_ij = eps_i * eps_j
dij = distance(conf, box)
keep_mask = np.ones((N, N)) - np.eye(N)
keep_mask = np.where(eps_ij != 0, keep_mask, 0)
if cutoff is not None:
if runtime_validate:
validate_coulomb_cutoff(cutoff, beta, threshold=1e-2)
eps_ij = np.where(dij < cutoff, eps_ij, 0)
# (ytz): this avoids a nan in the gradient in both jax and tensorflow
sig_ij = np.where(keep_mask, sig_ij, 0)
eps_ij = np.where(keep_mask, eps_ij, 0)
inv_dij = 1 / dij
inv_dij = np.where(np.eye(N), 0, inv_dij)
sig2 = sig_ij * inv_dij
sig2 *= sig2
sig6 = sig2 * sig2 * sig2
eij_lj = 4 * eps_ij * (sig6 - 1.0) * sig6
eij_lj = np.where(keep_mask, eij_lj, 0)
qi = np.expand_dims(charges, 0) # (1, N)
qj = np.expand_dims(charges, 1) # (N, 1)
qij = np.multiply(qi, qj)
# (ytz): trick used to avoid nans in the diagonal due to the 1/dij term.
keep_mask = 1 - np.eye(N)
qij = np.where(keep_mask, qij, 0)
dij = np.where(keep_mask, dij, 0)
# funny enough lim_{x->0} erfc(x)/x = 0
eij_charge = np.where(keep_mask,
qij * erfc(beta * dij) * inv_dij,
0) # zero out diagonals
if cutoff is not None:
eij_charge = np.where(dij > cutoff, 0, eij_charge)
eij_total = eij_lj * lj_rescale_mask + eij_charge * charge_rescale_mask
return np.sum(eij_total / 2)
def nonbonded_v3(conf, params, box, lamb, charge_rescale_mask, lj_rescale_mask,
scales, beta, cutoff, lambda_plane_idxs, lambda_offset_idxs):
N = conf.shape[0]
conf = convert_to_4d(conf, lamb, lambda_plane_idxs, lambda_offset_idxs,
cutoff)
# make 4th dimension of box large enough so its roughly aperiodic
if box is not None:
box_4d = np.eye(4) * 1000
box_4d = index_update(box_4d, index[:3, :3], box)
else:
box_4d = None
box = box_4d
charges = params[:, 0]
sig = params[:, 1]
eps = params[:, 2]
sig_i = np.expand_dims(sig, 0)
sig_j = np.expand_dims(sig, 1)
sig_ij = sig_i + sig_j
sig_ij_raw = sig_ij
eps_i = np.expand_dims(eps, 0)
eps_j = np.expand_dims(eps, 1)
eps_ij = eps_i * eps_j
ri = np.expand_dims(conf, 0)
rj = np.expand_dims(conf, 1)
dij = distance(ri, rj, box)
N = conf.shape[0]
keep_mask = np.ones((N, N)) - np.eye(N)
keep_mask = np.where(eps_ij != 0, keep_mask, 0)
if cutoff is not None:
eps_ij = np.where(dij < cutoff, eps_ij, np.zeros_like(eps_ij))
# (ytz): this avoids a nan in the gradient in both jax and tensorflow
sig_ij = np.where(keep_mask, sig_ij, np.zeros_like(sig_ij))
eps_ij = np.where(keep_mask, eps_ij, np.zeros_like(eps_ij))
sig2 = sig_ij / dij
sig2 *= sig2
sig6 = sig2 * sig2 * sig2
eij_lj = 4 * eps_ij * (sig6 - 1.0) * sig6
eij_lj = np.where(keep_mask, eij_lj, np.zeros_like(eij_lj))
qi = np.expand_dims(charges, 0) # (1, N)
qj = np.expand_dims(charges, 1) # (N, 1)
qij = np.multiply(qi, qj)
# (ytz): trick used to avoid nans in the diagonal due to the 1/dij term.
keep_mask = 1 - np.eye(conf.shape[0])
qij = np.where(keep_mask, qij, np.zeros_like(qij))
dij = np.where(keep_mask, dij, np.zeros_like(dij))
# funny enough lim_{x->0} erfc(x)/x = 0
eij_charge = np.where(keep_mask,
qij * erfc(beta * dij) / dij,
np.zeros_like(dij)) # zero out diagonals
if cutoff is not None:
eij_charge = np.where(dij > cutoff, np.zeros_like(eij_charge),
eij_charge)
eij_total = (eij_lj * lj_rescale_mask + eij_charge * charge_rescale_mask)
return np.sum(eij_total / 2)
