Updated nonbonded code so that it caches fits to "jj-nar-dna-pmf-100Mg.dat"

mrdna/config.py

@@ -10,6 +10,10 @@ _USER_CONF_DIR = Path(appdirs.user_data_dir())
 _USER_CONF_DIR.mkdir(parents=True, exist_ok=True)
 _USER_CONF = _USER_CONF_DIR / "mrdna.conf"
 
+CACHE_DIR = _USER_CONF_DIR / "mrdna.cache"
+CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+
 if _USER_CONF.exists():
     with open(_USER_CONF) as ch:
         config = json.load(ch)

mrdna/model/nbPot.py

@@ -6,9 +6,11 @@ from scipy.signal import savgol_filter as savgol
 from mrdna.model.nonbonded import NonbondedScheme
 from mrdna import get_resource_path
 
+from mrdna.config import CACHE_DIR
+
 dxParam = -1
 
-def maxForce(x,u0,maxForce=40):
+def _maxForce(x,u0,maxForce=40):
     maxForce = maxForce * 0.014393265 # convert pN to kcal_mol/AA
     dx = np.diff(x)
     f = -np.diff(u0)/dx
@@ -27,7 +29,6 @@ def maxForce(x,u0,maxForce=40):
     assert( np.isnan(u).sum() == 0 )
     return u
 
-# d = np.loadtxt("resources/jj-nar-dna-pmf-20Mg-200Na.dat")
 d = np.loadtxt(get_resource_path("jj-nar-dna-pmf-100Mg.dat"))
 
 x0 = d[:,1] + dxParam
@@ -51,26 +52,6 @@ def integrate_potential(potfn,x,y,*parms):
         u = u + potfn(r,*parms) # * x/r
     return u
 
-
-# def integrated_potential(potfn,x,numBp,*parms):
-#     u = np.zeros(np.shape(x))
-#     for h in range(-numBp,numBp+1):
-#         r = np.sqrt(x**2 + (3.4*h)**2)
-#         r[r<1] = 1
-#         u = u + potfn(r,*parms) # * x/r
-#     return u
-
-def fitFun(x,*parms):
-    ## 1-turn iteracting with 4 turns (more is not needed)
-    y = np.arange(-20,21)*3.4
-    return 10 * integrate_potential(parametric_potential,x,y,*parms)
-
-def fitFun5(x,*parms):
-    ## 1-turn iteracting with 4 turns (more is not needed)
-    y = np.arange(-20,21)*(5*3.4)
-    return 5 * 5 * integrate_potential(parametric_potential,x,y,*parms)
-
-
 xinterp = np.linspace(np.min(x0),np.max(x0),1000) # interpolate potentialb
 yinterp = interp1d(x0,y0)(xinterp)
 
@@ -83,49 +64,52 @@ def fit_potential(bp1,bp2):
         return (10.0/bp1) * integrate_potential(parametric_potential,x,y,*parms)
 
     popt, pcov = opt.curve_fit(fitFun, xinterp, yinterp, p0=y0/10**2 ) # find fit to interpolated potential
-    return lambda x: parametric_potential(x,*popt)
-
+    return popt
+    # return lambda x: parametric_potential(x,*popt)
 
-def nbPot(x,bps1,bps2):
+def _round_bp(bp):
+    return 10**(np.around(np.log10(bp)*100)/100.0)
+ 
+_pot_dict = {}
+_pot_dict_full = {}
+def nbPot_full(x,bps1,bps2):
     larger  = np.max([bps1,bps2])
     smaller = np.min([bps1,bps2])
-    if larger >= 2:
-        largerInt = int(0.5*larger)
-        smaller *= larger/(2*largerInt)
-        y = np.arange(-largerInt,largerInt+1)*3.4
-        u = smaller * integrate_potential(parametric_potential,x,y,*popt)
-    else:
-        u = smaller * larger * parametric_potential(x,*popt)
-    u1 = maxForce(x,u, 100)
+    key = (smaller,larger)
+    if key not in _pot_dict:
+        _pot_dict_full[key] = fit_potential(*key)
+    u = parametric_potential(x,*_pot_dict_full[key])
+
+    u1 = _maxForce(x,u, 100)
     # u2 = savgol(u1, int(5.0/(x0[1]-x0[0])), 3)    
     u2 = u
     assert(np.sum(np.isnan(u2)) == 0)
     return u2
 
-_pot_dict = {}
-
 def nbPot(x,bps1,bps2):
     larger  = np.max([bps1,bps2])
     smaller = np.min([bps1,bps2])
-    key = (smaller,larger)
-    if key not in _pot_dict:
-        _pot_dict[key] = fit_potential(*key)
-    u = _pot_dict[key](x)
+    smaller_shift, larger_shift = [_round_bp(bp) for bp in (smaller,larger)]
 
-    u1 = maxForce(x,u, 100)
-    # u2 = savgol(u1, int(5.0/(x0[1]-x0[0])), 3)    
+    key = (smaller_shift,larger_shift)
+    if key not in _pot_dict:
+        cache_file = CACHE_DIR / "nbPot-{}-{}.npz".format(*key)
+        try:
+            _pot_dict[key] = np.load(cache_file)['arr_0'].tolist()
+        except:
+            print("Could not find cache file for {}".format("nbPot-{}-{}.npz".format(*key)))
+            _pot_dict[key] = fit_potential(*key)
+            np.savez(cache_file, _pot_dict[key])
+
+    u = parametric_potential(x,*_pot_dict[key])
+    u = u * ( (smaller/smaller_shift)*(larger/larger_shift) )
+
+    u1 = _maxForce(x,u, 100)
+    u2 = savgol(u1, int(5.0/(x0[1]-x0[0])), 3)    
     u2 = u
     assert(np.sum(np.isnan(u2)) == 0)
     return u2
 
-x = np.linspace(0,100,1001)
-
-# np.savetxt("potentials/nb.orig.dat",np.array([x0,y0 ]).T)
-# np.savetxt("potentials/nb.orig-interp.dat",np.array([xinterp,yinterp ]).T)
-# np.savetxt("potentials/nb.fit.dat",np.array([x,fitFun(x,*popt)]).T)
-# np.savetxt("potentials/nb.fit2.dat",np.array([x,nbPot(x,10,40)]).T)
-
-
 def _writeNonbondedParameterFiles(self, prefix):
     x = np.arange(0, 50, 0.1)
     for i,j,t1,t2 in self._particleTypePairIter():
@@ -149,35 +133,21 @@ class nbDnaScheme(NonbondedScheme):
         return u
 nbDnaScheme = nbDnaScheme()
 
-if __name__ == "__main__":
-    ## run some tests
+def _run():
     np.savetxt("jj-filtered.dat", np.array((x0,y0)).T)
+    x = np.arange(0, 50, 0.1)
+    ## run some tests
+    for i in (0.5,1,1.05,2,3.333,5):
+        for j in (1,1.05,1.1,2,3.333, 4,5):
+            if i > j: continue
+            
+            # y = (10.0/i)*integrate_potential(nbPot_full, x, np.arange(-20,20)*(3.4*j) , i, j)
+            # np.savetxt("test1-{}-{}.dat".format(i,j), np.array((x,y)).T),
 
-    # y = single_bp_pot = fitFun(x,*popt)
-    # np.savetxt("1-bp-fit.dat", np.array((x,y)).T),
-
-    # y = fitFun5(x,*popt5)
-    # np.savetxt("5-bp-fit.dat", np.array((x,y)).T),
-
-
-    # tmp = nbPot(np.linspace(20,45,100), 1,1)
-    # import pdb
-    # pdb.set_trace()
-    y = 10*integrate_potential(nbPot, x, np.arange(-20,20)*3.4 , 1, 1)
-    np.savetxt("nb-1-1.dat", np.array((x,y)).T),
-
-    y = 10*integrate_potential(nbPot, x, np.arange(-20,20)*(3.4*2) , 1, 2)
-    np.savetxt("nb-1-2.dat", np.array((x,y)).T),
-
-    y = 10*integrate_potential(nbPot, x, np.arange(-20,20)*(3.4*5) , 1, 5)
-    np.savetxt("nb-1-5.dat", np.array((x,y)).T),
-
-    y = 5*integrate_potential(nbPot, x, np.arange(-20,20)*(3.4*2) , 2, 2)
-    np.savetxt("nb-2-2.dat", np.array((x,y)).T),
-
-    y = 2*integrate_potential(nbPot, x, np.arange(-20,20)*(3.4*5) , 5, 5)
-    np.savetxt("nb-5-5.dat", np.array((x,y)).T),
+            y = (10.0/i)*integrate_potential(nbPot, x, np.arange(-20,20)*(3.4*j) , i, j)
+            np.savetxt("test2-{}-{}.dat".format(i,j), np.array((x,y)).T),
 
-    # y = 10*5*integrate_potential(parametric_potential, x, np.arange(-20,20)*(3.4*5) , *popt)
-    # np.savetxt("scale-5-5.dat", np.array((x,y)).T),
 
+if __name__ == "__main__":
+    import cProfile
+    cProfile.run("_run()")