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)