update damp function in qeq module

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
 

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)
+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"]

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'
 

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