diff --git a/dmff/admp/qeq.py b/dmff/admp/qeq.py index 4f1ebea78..a2e1e00d2 100644 --- a/dmff/admp/qeq.py +++ b/dmff/admp/qeq.py @@ -97,13 +97,20 @@ def E_sr(pos, box, pairs, q, eta, buffer_scales, pbc_flag): @jit_condition(static_argnums=[6]) def E_sr2(pos, box, pairs, q, eta, buffer_scales, pbc_flag): + """ + eta: Gaussian width in angstrom + """ + # pairwise interaction ds = ds_pairs(pos, box, pairs, pbc_flag) etasqrt = jnp.sqrt(2 * (eta[pairs[:, 0]] ** 2 + eta[pairs[:, 1]] ** 2)) pre_pair = -eta_piecewise(etasqrt, ds) * DIELECTRIC - pre_self = etainv_piecewise(eta) / (jnp.sqrt(2 * jnp.pi)) * DIELECTRIC e_sr_pair = pre_pair * q[pairs[:, 0]] * q[pairs[:, 1]] / ds * buffer_scales e_sr_pair = mask_to_zero(e_sr_pair, buffer_scales) + + # self interaction + pre_self = etainv_piecewise(eta) / (2 * jnp.sqrt(jnp.pi)) * DIELECTRIC e_sr_self = pre_self * q * q + e_sr = jnp.sum(e_sr_pair) + jnp.sum(e_sr_self) return e_sr diff --git a/dmff/common/constants.py b/dmff/common/constants.py index 91bdeb4a3..bab320ddc 100644 --- a/dmff/common/constants.py +++ b/dmff/common/constants.py @@ -1,4 +1,17 @@ import numpy as np +from scipy import constants -DIELECTRIC = 1389.35455846 -SQRT_PI = np.sqrt(np.pi) \ No newline at end of file +j2ev = constants.physical_constants["joule-electron volt relationship"][0] +# kJ/mol to eV/particle +energy_coeff = j2ev * constants.kilo / constants.Avogadro +# vacuum electric permittivity in eV^-1 * angstrom^-1 +epsilon = constants.epsilon_0 / constants.elementary_charge * constants.angstrom +# qqrd2e = 1 / (4 * np.pi * epsilon) +# # eV +# dielectric = 1 / (4 * np.pi * epsilon) +# kJ/mol +DIELECTRIC = 1 / (4 * np.pi * epsilon) / energy_coeff +# DIELECTRIC = 1389.35455846 +SQRT_PI = np.sqrt(np.pi) + +__all__ = ["DIELECTRIC", "SQRT_PI"] diff --git a/dmff/settings.py b/dmff/settings.py index a880e8c52..e7a7fac7a 100644 --- a/dmff/settings.py +++ b/dmff/settings.py @@ -1,4 +1,9 @@ -from jax.config import config +try: + # jax < 0.4.20 + from jax.config import config +except: + # jax >= 0.4.20 + from jax import config PRECISION = 'double' # 'double' diff --git a/requirements.txt b/requirements.txt index 1980df549..691c4d4a5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,7 +6,7 @@ mkdocs-literate-nav>=0.4.1 mkdocstrings>=0.19.0 mkdocstrings-python>=0.7.0 pygments>=2.12 -jax>=0.4.1, <=0.4.20 +jax>=0.4.1 jaxlib>=0.4.1 pymbar>=4.0.0 tqdm