def nonbonded_v2(conf, params, box, lamb, exclusion_idxs, scales, beta, cutoff,
lambda_offset_idxs):
# assert box is None
conf_4d = convert_to_4d(conf, lamb, lambda_offset_idxs)
# print(conf_4d)
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
charge_params = params[:, 0]
lj_params = params[:, 1:]
charge_scales = scales[:, 0]
lj_scales = scales[:, 1]
lj = lennard_jones(conf_4d, lj_params, box_4d, cutoff)
lj_exc = lennard_jones_exclusion(conf_4d, lj_params, box_4d,
exclusion_idxs, lj_scales, cutoff)
es = simple_energy(conf_4d, box_4d, charge_params, exclusion_idxs,
charge_scales, beta, cutoff)
return lj - lj_exc + es
def lennard_jones_v2(conf, lj_params, box, lamb, exclusion_idxs, lj_scales,
cutoff, lambda_plane_idxs, lambda_offset_idxs):
conf_4d = convert_to_4d(conf, lamb, lambda_plane_idxs, lambda_offset_idxs,
cutoff)
box_4d = np.eye(4) * 1000
box_4d = index_update(box_4d, index[:3, :3], box)
lj = lennard_jones(conf_4d, lj_params, box_4d, cutoff)
lj_exc = lennard_jones_exclusion(conf_4d, lj_params, box_4d,
exclusion_idxs, lj_scales, cutoff)
return lj - lj_exc
def _nonbonded_v3_clone(
conf,
params,
box,
lamb,
charge_rescale_mask,
lj_rescale_mask,
beta,
cutoff,
lambda_plane_idxs,
lambda_offset_idxs,
):
"""See docstring of nonbonded_v3 for more details
This is here just for testing purposes, to mimic the signature of nonbonded_v3 but to use
nonbonded_v3_on_specific_pairs under the hood.
"""
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
# TODO: len(inds_i) == n_interactions -- may want to break this
# up into more manageable blocks if n_interactions is large
inds_i, inds_j = get_all_pairs_indices(N)
lj, coulomb = nonbonded_v3_on_specific_pairs(conf, params, box, inds_i,
inds_j, beta, cutoff)
# keep only eps > 0
eps = params[:, 2]
lj = np.where(eps[inds_i] > 0, lj, 0)
lj = np.where(eps[inds_j] > 0, lj, 0)
eij_total = lj * lj_rescale_mask[
inds_i, inds_j] + coulomb * charge_rescale_mask[inds_i, inds_j]
return np.sum(eij_total)
def nongroup_electrostatics(
conf,
lamb,
charge_params,
exclusion_idxs,
charge_scales,
cutoff,
lambda_plane_idxs,
lambda_offset_idxs):
# assert box is None
conf_4d = convert_to_4d(conf, lamb, lambda_plane_idxs, lambda_offset_idxs, cutoff)
return simple_energy(conf_4d, charge_params, exclusion_idxs, charge_scales, cutoff)
def electrostatics_v2(conf, charge_params, box, lamb, exclusion_idxs,
charge_scales, beta, cutoff, lambda_offset_idxs):
# assert box is None
conf_4d = convert_to_4d(conf, lamb, lambda_offset_idxs)
# print(conf_4d)
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
return simple_energy(conf_4d, box_4d, charge_params, exclusion_idxs,
charge_scales, beta, cutoff)
def nonbonded(conf, lamb, charge_params, lj_params, exclusion_idxs,
charge_scales, lj_scales, cutoff, lambda_plane_idxs,
lambda_offset_idxs):
# assert box is None
conf_4d = convert_to_4d(conf, lamb, lambda_plane_idxs, lambda_offset_idxs,
cutoff)
lj = lennard_jones(conf_4d, lj_params, cutoff)
lj_exc = lennard_jones_exclusion(conf_4d, lj_params, exclusion_idxs,
lj_scales, cutoff)
es = simple_energy(conf_4d, charge_params, exclusion_idxs, charge_scales,
cutoff)
return lj - lj_exc + es
def group_lennard_jones(
conf,
lamb,
lj_params,
exclusion_idxs,
lj_scales,
cutoff,
lambda_plane_idxs,
lambda_offset_idxs,
lambda_group_idxs):
conf_4d = convert_to_4d(conf, lamb, lambda_plane_idxs, lambda_offset_idxs, cutoff)
lj = lennard_jones(conf_4d, lj_params, cutoff, lambda_group_idxs)
lj_exc = lennard_jones_exclusion(conf_4d, lj_params, exclusion_idxs, lj_scales, cutoff, lambda_group_idxs)
return lj - lj_exc
def gbsa_obc(
coords,
# params,
lamb,
# box,
charge_params,
gb_params,
# charge_idxs,
# radii_idxs,
# scale_idxs,
alpha,
beta,
gamma,
cutoff_radii,
cutoff_force,
lambda_plane_idxs,
lambda_offset_idxs,
dielectric_offset=0.009,
surface_tension=28.3919551,
solute_dielectric=1.0,
solvent_dielectric=78.5,
probe_radius=0.14):
box = None
assert cutoff_radii == cutoff_force
coords_4d = convert_to_4d(coords, lamb, lambda_plane_idxs,
lambda_offset_idxs, cutoff_radii)
N = len(charge_params)
radii = gb_params[:, 0]
scales = gb_params[:, 1]
ri = np.expand_dims(coords_4d, 0)
rj = np.expand_dims(coords_4d, 1)
dij = distance(ri, rj, box)
eye = np.eye(N, dtype=dij.dtype)
r = dij + eye # so I don't have divide-by-zero nonsense
or1 = radii.reshape((N, 1)) - dielectric_offset
or2 = radii.reshape((1, N)) - dielectric_offset
sr2 = scales.reshape((1, N)) * or2
L = np.maximum(or1, abs(r - sr2))
U = r + sr2
I = 1 / L - 1 / U + 0.25 * (r - sr2**2 / r) * (1 / (U**2) - 1 /
(L**2)) + 0.5 * np.log(
L / U) / r
# handle the interior case
I = np.where(or1 < (sr2 - r), I + 2 * (1 / or1 - 1 / L), I)
I = step(r + sr2 - or1) * 0.5 * I # note the extra 0.5 here
I -= np.diag(np.diag(I))
# switch I only for now
# inner = (np.pi*np.power(dij,8))/(2*cutoff_radii)
# sw = np.power(np.cos(inner), 2)
# I = I*sw
I = np.where(dij > cutoff_radii, 0, I)
I = np.sum(I, axis=1)
# okay, next compute born radii
offset_radius = radii - dielectric_offset
psi = I * offset_radius
psi_coefficient = alpha
psi2_coefficient = beta
psi3_coefficient = gamma
psi_term = (psi_coefficient * psi) - (psi2_coefficient *
psi**2) + (psi3_coefficient * psi**3)
B = 1 / (1 / offset_radius - np.tanh(psi_term) / radii)
E = 0.0
# single particle
# ACE
E += np.sum(surface_tension * (radii + probe_radius)**2 * (radii / B)**6)
# on-diagonal
charges = charge_params
E += np.sum(-0.5 * (1 / solute_dielectric - 1 / solvent_dielectric) *
charges**2 / B)
# particle pair
f = np.sqrt(r**2 + np.outer(B, B) * np.exp(-r**2 / (4 * np.outer(B, B))))
charge_products = np.outer(charges, charges)
ixns = -(1 / solute_dielectric -
1 / solvent_dielectric) * charge_products / f
# sw = np.power(np.cos((np.pi*dij)/(2*cutoff_radii)), 2)
# ixns = ixns*sw
ixns = np.where(dij > cutoff_force, 0, ixns)
E += np.sum(np.triu(ixns, k=1))
return E
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)