Add two features to stop short range failure:

1. hard core potential associated with Slater function
2. bounded dipole (use soft truncation)

Author: KuangYu

dmff/admp/pairwise.py

@@ -165,5 +165,11 @@ def slater_sr_kernel(dr, m, ai, aj, bi, bj):
     br = b * dr
     br2 = br * br
     P = 1/3 * br2 + br + 1 
-    return a * P * jnp.exp(-br) * m
+    # hard core potential
+    x = 3.9/br
+    x2 = x * x
+    x4 = x2 * x2
+    x8 = x4 * x4
+    x14 = x8 * x2 * x4
+    return a * (P * jnp.exp(-br) + x14) * m
 

dmff/admp/pme.py

@@ -32,6 +32,15 @@
 
 DEFAULT_THOLE_WIDTH = 5.0
 
+# variables used in soft dipole truncation
+MAX_DIP = 1.0
+TRUNCATION_HARDNESS = 25 # the smaller, the softer
+def SOFT_TRUNCATION(x):
+    x2 = x * x
+    x_abs = jnp.sqrt(x2 + 1e-6)
+    val = -1/TRUNCATION_HARDNESS * jnp.log(1 + jnp.exp(-TRUNCATION_HARDNESS*(x_abs-MAX_DIP))) + MAX_DIP
+    return val * x/x_abs
+
 
 class ADMPPmeForce:
     """
@@ -416,6 +425,8 @@ def update_U(i, U):
                     dScales,
                 )
                 U = U - field * pol[:, jnp.newaxis] / DIELECTRIC
+                # soft truncation: stop polarization catastrophe
+                U = SOFT_TRUNCATION(U)
                 return U
 
             U = jax.lax.fori_loop(0, steps_pol, update_U, U)