Fix the two examples:

* fluctuated_leading_term_waterff
* peg_slater_isa
Make map_atomtype and map_poltype available in generator

dmff/generators/admp.py

@@ -139,6 +139,7 @@ def createPotential(self, topdata: DMFFTopology, nonbondedMethod, nonbondedCutof
         for i in range(n_atoms):
             atype = atoms[i].meta[self.key_type]
             map_atomtype[i] = self._find_atype_key_index(atype)
+        self.map_atomtype = map_atomtype
         # here box is only used to setup ewald parameters, no need to be differentiable
         if lpme:
             box = topdata.getPeriodicBoxVectors() * 10
@@ -301,6 +302,8 @@ def createPotential(self, topdata: DMFFTopology, nonbondedMethod, nonbondedCutof
             atype = atoms[i].meta[self.key_type]
             map_atomtype[i] = self._find_atype_key_index(atype)
 
+        self.map_atomtype = map_atomtype
+
         # here box is only used to setup ewald parameters, no need to be differentiable
         if lpme:
             box = topdata.getPeriodicBoxVectors() * 10
@@ -435,6 +438,7 @@ def createPotential(self, topdata: DMFFTopology, nonbondedMethod, nonbondedCutof
             atype = atoms[i].meta[self.key_type]
             map_atomtype[i] = np.where(self.atom_keys == atype)[0][0]
 
+        self.map_atomtype = map_atomtype
         pot_fn_sr = generate_pairwise_interaction(TT_damping_qq_kernel,
                                                   static_args={})
 
@@ -554,6 +558,8 @@ def createPotential(self, topdata: DMFFTopology, nonbondedMethod, nonbondedCutof
             atype = atoms[i].meta[self.key_type]
             map_atomtype[i] = np.where(self.atom_keys == atype)[0][0]
 
+        self.map_atomtype = map_atomtype
+
         # WORKING
         pot_fn_sr = generate_pairwise_interaction(slater_disp_damping_kernel,
                                                   static_args={})
@@ -663,6 +669,8 @@ def createPotential(self, topdata: DMFFTopology, nonbondedMethod, nonbondedCutof
             atype = atoms[i].meta[self.key_type]
             map_atomtype[i] = np.where(self.atom_keys == atype)[0][0]
 
+        self.map_atomtype = map_atomtype
+
         pot_fn_sr = generate_pairwise_interaction(slater_sr_kernel,
                                                   static_args={})
 
@@ -709,6 +717,8 @@ def createPotential(self, topdata: DMFFTopology, nonbondedMethod, nonbondedCutof
             atype = atoms[i].meta[self.key_type]
             map_atomtype[i] = np.where(self.atom_keys == atype)[0][0]
 
+        self.map_atomtype = map_atomtype
+
         pot_fn_sr = generate_pairwise_interaction(slater_sr_kernel,
                                                   static_args={})
 
@@ -982,6 +992,9 @@ def createPotential(self, topdata: DMFFTopology, nonbondedMethod, nonbondedCutof
             if self.lpol:
                 map_poltype[i] = self._find_polarize_key_index(atype)
 
+        self.map_poltype = map_poltype
+        self.map_atomtype = map_atomtype
+
         # here box is only used to setup ewald parameters, no need to be differentiable
         if lpme:
             box = topdata.getPeriodicBoxVectors() * 10

examples/fluctuated_leading_term_waterff/ref_out

@@ -0,0 +1,9 @@
+# Tot Interaction Energy:
+# -41.437696011272465 kJ/mol
+# Tot force :
+# [[  6.99817736 -10.73540559   7.44495642]
+ [-74.774497    57.81233046  24.98315842]
+ [  2.00037973   4.60480058 -13.86528301]
+ [ 63.0415243  -40.32457937 -15.57667743]
+ [  2.60863633   6.53827371  -8.38934949]
+ [  0.13944051 -17.89996202   5.35314178]]

examples/fluctuated_leading_term_waterff/run.py

@@ -54,28 +54,29 @@ def compute_leading_terms(positions,box):
     c6_list = c6_list.at[2::3].set(jnp.sqrt(C6_H2))
     return c0, c6_list
 
-def generate_calculator(pot_disp, pot_pme, pot_sr, disp_generator, pme_generator):
+def generate_calculator(pot_disp, pot_pme, pot_sr, disp_generator, pme_generator, params, covalent_map):
     def admp_calculator(positions, box, pairs):
-        params_pme = pme_generator.paramtree['ADMPPmeForce']
-        params_disp = disp_generator.paramtree['ADMPDispForce']
+        # params_pme = pme_generator.paramtree['ADMPPmeForce']
+        # params_disp = disp_generator.paramtree['ADMPDispForce']
+        params_pme = params['ADMPPmeForce']
+        params_disp = params['ADMPDispForce']
         c0, c6_list = compute_leading_terms(positions,box) # compute fluctuated leading terms
         Q_local = params_pme["Q_local"][pme_generator.map_atomtype]
         Q_local = Q_local.at[:,0].set(c0)  # change fixed charge into fluctuated one
         pol = params_pme["pol"][pme_generator.map_atomtype]
-        tholes = params_pme["tholes"][pme_generator.map_atomtype]
+        tholes = params_pme["thole"][pme_generator.map_atomtype]
         c8_list = jnp.sqrt(params_disp["C8"][disp_generator.map_atomtype]*1e8)
         c10_list = jnp.sqrt(params_disp["C10"][disp_generator.map_atomtype]*1e10)
         c_list = jnp.vstack((c6_list, c8_list, c10_list))
-        covalent_map = disp_generator.covalent_map    
         a_list = (params_disp["A"][disp_generator.map_atomtype] / 2625.5)
         b_list = params_disp["B"][disp_generator.map_atomtype] * 0.0529177249 
         q_list = params_disp["Q"][disp_generator.map_atomtype]
 
         E_pme = pme_generator.pme_force.get_energy(
-                positions, box, pairs, Q_local, pol, tholes, params_pme["mScales"], params_pme["pScales"], params_pme["dScales"]
+                positions, box, pairs, Q_local, pol, tholes, pme_generator.mScales, pme_generator.pScales, pme_generator.dScales
                 )    
-        E_disp = disp_generator.disp_pme_force.get_energy(positions, box, pairs, c_list.T, params_disp["mScales"])
-        E_sr = pot_sr(positions, box, pairs, params_disp["mScales"], a_list, b_list, q_list, c_list[0])
+        E_disp = disp_generator.disp_pme_force.get_energy(positions, box, pairs, c_list.T, disp_generator.mScales)
+        E_sr = pot_sr(positions, box, pairs, disp_generator.mScales, a_list, b_list, q_list, c_list[0])
         return E_pme 
     return jit(value_and_grad(admp_calculator,argnums=(0)))
 
@@ -84,27 +85,27 @@ def admp_calculator(positions, box, pairs):
     H = Hamiltonian('forcefield.xml')
     app.Topology.loadBondDefinitions("residues.xml")
     pdb = app.PDBFile("water_dimer.pdb")
-    rc = 4
+    rc = 0.4
 
     # generator stores all force field parameters     
     disp_generator, pme_generator = H.getGenerators()
-    pots = H.createPotential(pdb.topology, nonbondedCutoff=rc*unit.angstrom, step_pol=5)
+    pots = H.createPotential(pdb.topology, nonbondedCutoff=rc*unit.nanometer, step_pol=5)
     pot_disp = pots.dmff_potentials['ADMPDispForce']
     pot_pme = pots.dmff_potentials['ADMPPmeForce']
     pot_sr = generate_pairwise_interaction(TT_damping_qq_c6_kernel, static_args={})
 
     # construct inputs
-    positions = jnp.array(pdb.positions._value) * 10
+    positions = jnp.array(pdb.positions._value)
     a, b, c = pdb.topology.getPeriodicBoxVectors()
-    box = jnp.array([a._value, b._value, c._value]) * 10
+    box = jnp.array([a._value, b._value, c._value])
     # neighbor list
-    nbl = nblist.NeighborList(box, rc, H.getGenerators()[0].covalent_map)
+    nbl = nblist.NeighborList(box, rc, pots.meta["cov_map"])
     nbl.allocate(positions)
     pairs = nbl.pairs
 
 
-    admp_calc = generate_calculator(pot_disp, pot_pme, pot_sr, disp_generator, pme_generator)
-    tot_ene, tot_force = admp_calc(positions, box, pairs)
+    admp_calc = generate_calculator(pot_disp, pot_pme, pot_sr, disp_generator, pme_generator, H.getParameters(), pots.meta["cov_map"])
+    tot_ene, tot_force = admp_calc(positions*10, box*10, pairs)
     print('# Tot Interaction Energy:')
     print('#', tot_ene, 'kJ/mol')
     print('# Tot force :')

examples/peg_slater_isa/calc_energy_comps.py

@@ -23,10 +23,10 @@
     H_A = Hamiltonian(ff)
     H_B = Hamiltonian(ff)
 
-    rc = 15
+    rc = 1.45
 
     # get potential functions
-    pots_AB = H_AB.createPotential(pdb_AB.topology, nonbondedCutoff=rc*angstrom, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
+    pots_AB = H_AB.createPotential(pdb_AB.topology, nonbondedCutoff=rc*nanometer, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
     pot_pme_AB = pots_AB.dmff_potentials['ADMPPmeForce']
     pot_disp_AB = pots_AB.dmff_potentials['ADMPDispPmeForce']
     pot_ex_AB = pots_AB.dmff_potentials['SlaterExForce']
@@ -36,7 +36,7 @@
     pot_dhf_AB = pots_AB.dmff_potentials['SlaterDhfForce']
     pot_dmp_es_AB = pots_AB.dmff_potentials['QqTtDampingForce']
     pot_dmp_disp_AB = pots_AB.dmff_potentials['SlaterDampingForce']
-    pots_A = H_A.createPotential(pdb_A.topology, nonbondedCutoff=rc*angstrom, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
+    pots_A = H_A.createPotential(pdb_A.topology, nonbondedCutoff=rc*nanometer, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
     pot_pme_A = pots_A.dmff_potentials['ADMPPmeForce']
     pot_disp_A = pots_A.dmff_potentials['ADMPDispPmeForce']
     pot_ex_A = pots_A.dmff_potentials['SlaterExForce']
@@ -46,7 +46,7 @@
     pot_dhf_A = pots_A.dmff_potentials['SlaterDhfForce']
     pot_dmp_es_A = pots_A.dmff_potentials['QqTtDampingForce']
     pot_dmp_disp_A = pots_A.dmff_potentials['SlaterDampingForce']
-    pots_B = H_B.createPotential(pdb_B.topology, nonbondedCutoff=rc*angstrom, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
+    pots_B = H_B.createPotential(pdb_B.topology, nonbondedCutoff=rc*nanometer, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
     pot_pme_B = pots_B.dmff_potentials['ADMPPmeForce']
     pot_disp_B = pots_B.dmff_potentials['ADMPDispPmeForce']
     pot_ex_B = pots_B.dmff_potentials['SlaterExForce']
@@ -60,21 +60,21 @@
     pme_generator_A = H_A.getGenerators()[0]
     pme_generator_B = H_B.getGenerators()[0]
 
-    pos_AB0 = jnp.array(pdb_AB.positions._value) * 10
+    pos_AB0 = jnp.array(pdb_AB.positions._value)
     n_atoms = len(pos_AB0)
     n_atoms_A = n_atoms // 2
     n_atoms_B = n_atoms // 2
-    pos_A0 = jnp.array(pdb_AB.positions._value[:n_atoms_A]) * 10
-    pos_B0 = jnp.array(pdb_AB.positions._value[n_atoms_A:n_atoms]) * 10
-    box = jnp.array(pdb_AB.topology.getPeriodicBoxVectors()._value) * 10
+    pos_A0 = jnp.array(pdb_AB.positions._value[:n_atoms_A])
+    pos_B0 = jnp.array(pdb_AB.positions._value[n_atoms_A:n_atoms])
+    box = jnp.array(pdb_AB.topology.getPeriodicBoxVectors()._value)
     # nn list initial allocation
-    nbl_AB = nblist.NeighborList(box, rc, H_AB.getGenerators()[0].covalent_map)
+    nbl_AB = nblist.NeighborList(box, rc, pots_AB.meta["cov_map"])
     nbl_AB.allocate(pos_AB0)
     pairs_AB = nbl_AB.pairs
-    nbl_A = nblist.NeighborList(box, rc, H_A.getGenerators()[0].covalent_map)
+    nbl_A = nblist.NeighborList(box, rc, pots_A.meta["cov_map"])
     nbl_A.allocate(pos_A0)
     pairs_A = nbl_A.pairs
-    nbl_B = nblist.NeighborList(box, rc, H_B.getGenerators()[0].covalent_map)
+    nbl_B = nblist.NeighborList(box, rc, pots_B.meta["cov_map"])
     nbl_B.allocate(pos_B0)
     pairs_B = nbl_B.pairs
 
@@ -107,8 +107,8 @@
             E_tot_ref = scan_res['tot'][ipt]
 
             # get position array
-            pos_A = jnp.array(scan_res['posA'][ipt])
-            pos_B = jnp.array(scan_res['posB'][ipt])
+            pos_A = jnp.array(scan_res['posA'][ipt]) / 10
+            pos_B = jnp.array(scan_res['posB'][ipt]) / 10
             pos_AB = jnp.concatenate([pos_A, pos_B], axis=0)
 
 
@@ -135,10 +135,15 @@
             map_poltypes = pme_generator_AB.map_poltype
             Q_local = params_pme['Q_local'][map_atypes]
             pol = params_pme['pol'][map_poltypes]
-            tholes = params_pme['tholes'][map_poltypes]
+            tholes = params_pme['thole'][map_poltypes]
             pme_force = pme_generator_AB.pme_force
-            E_AB_nonpol = pme_force.energy_fn(pos_AB, box, pairs_AB, Q_local, U_ind_AB, pol, tholes, params_pme['mScales'], params_pme['pScales'], params_pme['dScales'])
+
+
+            E_AB_nonpol = pme_force.energy_fn(pos_AB*10, box*10, pairs_AB, Q_local, U_ind_AB, pol, tholes, \
+                    pme_generator_AB.mScales, pme_generator_AB.pScales, pme_generator_AB.dScales)
             E_es = E_AB_nonpol - E_A - E_B
+
+
             E_dmp_es = pot_dmp_es_AB(pos_AB, box, pairs_AB, params) \
                      - pot_dmp_es_A(pos_A, box, pairs_A, params) \
                      - pot_dmp_es_B(pos_B, box, pairs_B, params)

examples/peg_slater_isa/check.py

@@ -27,10 +27,10 @@
     H_A = Hamiltonian(ff)
     H_B = Hamiltonian(ff)
 
-    rc = 15
+    rc = 1.45
 
     # get potential functions
-    pots_AB = H_AB.createPotential(pdb_AB.topology, nonbondedCutoff=rc*angstrom, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
+    pots_AB = H_AB.createPotential(pdb_AB.topology, nonbondedCutoff=rc*nanometer, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
     pot_pme_AB = pots_AB.dmff_potentials['ADMPPmeForce']
     pot_disp_AB = pots_AB.dmff_potentials['ADMPDispPmeForce']
     pot_ex_AB = pots_AB.dmff_potentials['SlaterExForce']
@@ -40,7 +40,7 @@
     pot_dhf_AB = pots_AB.dmff_potentials['SlaterDhfForce']
     pot_dmp_es_AB = pots_AB.dmff_potentials['QqTtDampingForce']
     pot_dmp_disp_AB = pots_AB.dmff_potentials['SlaterDampingForce']
-    pots_A = H_A.createPotential(pdb_A.topology, nonbondedCutoff=rc*angstrom, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
+    pots_A = H_A.createPotential(pdb_A.topology, nonbondedCutoff=rc*nanometer, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
     pot_pme_A = pots_A.dmff_potentials['ADMPPmeForce']
     pot_disp_A = pots_A.dmff_potentials['ADMPDispPmeForce']
     pot_ex_A = pots_A.dmff_potentials['SlaterExForce']
@@ -50,7 +50,7 @@
     pot_dhf_A = pots_A.dmff_potentials['SlaterDhfForce']
     pot_dmp_es_A = pots_A.dmff_potentials['QqTtDampingForce']
     pot_dmp_disp_A = pots_A.dmff_potentials['SlaterDampingForce']
-    pots_B = H_B.createPotential(pdb_B.topology, nonbondedCutoff=rc*angstrom, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
+    pots_B = H_B.createPotential(pdb_B.topology, nonbondedCutoff=rc*nanometer, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
     pot_pme_B = pots_B.dmff_potentials['ADMPPmeForce']
     pot_disp_B = pots_B.dmff_potentials['ADMPDispPmeForce']
     pot_ex_B = pots_B.dmff_potentials['SlaterExForce']
@@ -66,22 +66,22 @@
     pme_generator_B = H_B.getGenerators()[0]
 
     # init positions used to set up neighbor list
-    pos_AB0 = jnp.array(pdb_AB.positions._value) * 10
+    pos_AB0 = jnp.array(pdb_AB.positions._value)
     n_atoms = len(pos_AB0)
     n_atoms_A = n_atoms // 2
     n_atoms_B = n_atoms // 2
-    pos_A0 = jnp.array(pdb_AB.positions._value[:n_atoms_A]) * 10
-    pos_B0 = jnp.array(pdb_AB.positions._value[n_atoms_A:n_atoms]) * 10
-    box = jnp.array(pdb_AB.topology.getPeriodicBoxVectors()._value) * 10
+    pos_A0 = jnp.array(pdb_AB.positions._value[:n_atoms_A])
+    pos_B0 = jnp.array(pdb_AB.positions._value[n_atoms_A:n_atoms])
+    box = jnp.array(pdb_AB.topology.getPeriodicBoxVectors()._value)
 
     # nn list initial allocation
-    nbl_AB = nblist.NeighborList(box, rc, H_AB.getGenerators()[0].covalent_map)
+    nbl_AB = nblist.NeighborList(box, rc, pots_AB.meta["cov_map"])
     nbl_AB.allocate(pos_AB0)
     pairs_AB = nbl_AB.pairs
-    nbl_A = nblist.NeighborList(box, rc, H_A.getGenerators()[0].covalent_map)
+    nbl_A = nblist.NeighborList(box, rc, pots_A.meta["cov_map"])
     nbl_A.allocate(pos_A0)
     pairs_A = nbl_A.pairs
-    nbl_B = nblist.NeighborList(box, rc, H_B.getGenerators()[0].covalent_map)
+    nbl_B = nblist.NeighborList(box, rc, pots_B.meta["cov_map"])
     nbl_B.allocate(pos_B0)
     pairs_B = nbl_B.pairs
 
@@ -123,7 +123,7 @@
     params_dmp_disp['C10'] = params['C10']
     # long range parameters
     params_espol = {}
-    for k in ['mScales', 'pScales', 'dScales', 'Q_local', 'pol', 'tholes']:
+    for k in ['mScales', 'pScales', 'dScales', 'Q_local', 'pol', 'thole']:
         params_espol[k] = params[k]
     params_disp = {}
     for k in ['B', 'C6', 'C8', 'C10', 'mScales']:
@@ -173,8 +173,8 @@
             E_dhf_ref = scan_res['dhf'][ipt]
             E_tot_ref = scan_res['tot'][ipt]
 
-            pos_A = jnp.array(scan_res['posA'][ipt])
-            pos_B = jnp.array(scan_res['posB'][ipt])
+            pos_A = jnp.array(scan_res['posA'][ipt]) / 10
+            pos_B = jnp.array(scan_res['posB'][ipt]) / 10
             pos_AB = jnp.concatenate([pos_A, pos_B], axis=0)
 
             #####################
@@ -200,7 +200,7 @@
             map_poltype = pme_generator_AB.map_poltype
             Q_local = params['Q_local'][map_atypes]
             pol = params['pol'][map_poltype]
-            tholes = params['tholes'][map_poltype]
+            tholes = params['thole'][map_poltype]
             pme_force = pme_generator_AB.pme_force
             E_AB_nonpol = pme_force.energy_fn(pos_AB, box, pairs_AB, Q_local, U_ind_AB, pol, tholes, params['mScales'], params['pScales'], params['dScales'])
             E_es = E_AB_nonpol - E_A - E_B

examples/peg_slater_isa/check_calc.py

@@ -19,9 +19,9 @@
     ff = 'peg.xml'
     pdb_AB = PDBFile('peg2.pdb')
     H_AB = Hamiltonian(ff)
-    rc = 15
+    rc = 1.45
     # get potential functions
-    pots_AB = H_AB.createPotential(pdb_AB.topology, nonbondedCutoff=rc*angstrom, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
+    pots_AB = H_AB.createPotential(pdb_AB.topology, nonbondedCutoff=rc*nanometer, nonbondedMethod=CutoffPeriodic, ethresh=1e-4)
     pot_pme_AB = pots_AB.dmff_potentials['ADMPPmeForce']
     pot_disp_AB = pots_AB.dmff_potentials['ADMPDispPmeForce']
     pot_ex_AB = pots_AB.dmff_potentials['SlaterExForce']
@@ -35,12 +35,12 @@
     paramtree = H_AB.getParameters()
 
     # init positions used to set up neighbor list
-    pos_AB0 = jnp.array(pdb_AB.positions._value) * 10
+    pos_AB0 = jnp.array(pdb_AB.positions._value)
     n_atoms = len(pos_AB0)
-    box = jnp.array(pdb_AB.topology.getPeriodicBoxVectors()._value) * 10
+    box = jnp.array(pdb_AB.topology.getPeriodicBoxVectors()._value)
 
     # nn list initial allocation
-    nbl_AB = nblist.NeighborList(box, rc, H_AB.getGenerators()[0].covalent_map)
+    nbl_AB = nblist.NeighborList(box, rc, pots_AB.meta['cov_map'])
     nbl_AB.allocate(pos_AB0)
     pairs_AB = nbl_AB.pairs
     pairs_AB =  pairs_AB[pairs_AB[:, 0] < pairs_AB[:, 1]]