Merge pull request #4 from Feiyang472/master

Completed reciprocal part force calculation.

demo.py

@@ -1,12 +1,15 @@
 from email import generator
+from python.utils import convert_cart2harm
 import numpy as np
-from python.ADMPForce import ADMPGenerator
+from simtk.openmm.openmm import Platform_getPluginLoadFailures
+from python.ADMPForce import ADMPGenerator, get_mtpls_global
 import scipy
 from scipy.stats import special_ortho_group
 from simtk.openmm import *
 from simtk.openmm.app import *
 from simtk.unit import *
 import mpidplugin
+import python.pme as pme
 
 mScales = np.array([0.0, 0.0, 0.0, 1.0])
 pScales = np.array([0.0, 0.0, 0.0, 1.0])
@@ -33,26 +36,37 @@
 force.update()
 force.kappa = 0.328532611
 
-# test real_space_energy
-real_e = force.calc_real_space_energy()
+multipoles_lc = np.concatenate((np.expand_dims(force.mpid_params['charges'], axis=1), force.mpid_params['dipoles'], force.mpid_params['quadrupoles']), axis=1)
+Q_lh = convert_cart2harm(multipoles_lc, lmax=2)
+axis_types = force.mpid_params['axis_types']
+axis_indices = force.mpid_params['axis_indices']
 
 # test reci_space_energy
+print('============reciprocal energy===========')
 reci_e = force.calc_reci_space_energy()
+print(reci_e)
+print('============reciprocal force============')
+reci_f = force.calc_reci_space_force()
+print(reci_f)
+
+
+
+# test real_space_energy
+# real_e = force.calc_real_space_energy()
 
-# test self energy
-self_e = force.calc_self_energy()
+# # test self energy
+# self_e = force.calc_self_energy()
 
 
-print('real space', real_e)
-print('reciprocal', reci_e)
-print('self', self_e)
+# print('real space', real_e)
+# print('self', self_e)
 
-pdb = PDBFile(pdb)
-forcefield = ForceField(xml)
-system = forcefield.createSystem(pdb.topology, nonbondedMethod=LJPME, nonbondedCutoff=8*angstrom, constraints=HBonds, defaultTholeWidth=8)
-integrator = VerletIntegrator(1e-10*femtoseconds)
-simulation = Simulation(pdb.topology, system, integrator)
-context = simulation.context
-context.setPositions(positions*angstrom)
-state = context.getState(getEnergy=True, getPositions=True)
-Etot = state.getPotentialEnergy().value_in_unit(kilojoules_per_mole)
+# pdb = PDBFile(pdb)
+# forcefield = ForceField(xml)
+# system = forcefield.createSystem(pdb.topology, nonbondedMethod=LJPME, nonbondedCutoff=8*angstrom, constraints=HBonds, defaultTholeWidth=8)
+# integrator = VerletIntegrator(1e-10*femtoseconds)
+# simulation = Simulation(pdb.topology, system, integrator)
+# context = simulation.context
+# context.setPositions(positions*angstrom)
+# state = context.getState(getEnergy=True, getPositions=True)
+# Etot = state.getPotentialEnergy().value_in_unit(kilojoules_per_mole)

docs/README.md

@@ -4,4 +4,15 @@ home: true
 heroText: Automatic Differentiation Multipole Moment Molecular Forcefield
 ---
 
-　 ADMP因何而发生. 既然如何，了解清楚ADMP到底是一种怎么样的存在，是解决一切问题的关键。 我们都知道，只要有意义，那么就必须慎重考虑。那么，ADMP因何而发生?德国在不经意间这样说过，只有在人群中间，才能认识自己。这启发了我，既然如何，问题的关键究竟为何? 洛克在不经意间这样说过，学到很多东西的诀窍，就是一下子不要学很多。带着这句话，我们还要更加慎重的审视这个问题： 一般来讲，我们都必须务必慎重的考虑考虑。孔子曾经说过，知之者不如好之者，好之者不如乐之者。这不禁令我深思。而这些并不是完全重要，更加重要的问题是，经过上述讨论一般来讲，我们都必须务必慎重的考虑考虑。就我个人来说，ADMP对我的意义，不能不说非常重大。我们一般认为，抓住了问题的关键，其他一切则会迎刃而解。 ADMP，发生了会如何，不发生又会如何。了解清楚ADMP到底是一种怎么样的存在，是解决一切问题的关键。
+# ADMP Generator:
+Factory pattern for `ADMPforce` objects
+
+- `force`
+  - created by `ADMPGenerator`
+  - `self.update()` creates global frame multipoles and neighbour lists.
+  - 
+
+## Demo:
+Constructs test case 
+
+## Derivative quantities

python/ADMPForce.py

@@ -10,9 +10,11 @@
 from simtk.unit import *
 from .neighborList import construct_nblist
 from .utils import *
-from .pme import pme_real, pme_reciprocal, pme_self
+from .pme import pme_real, pme_self, gen_pme_reciprocal
 import mpidplugin
 
+from jax import value_and_grad, jit
+
 # see the python/mpidplugin.i code 
 ZThenX = 0
 Bisector = 1
@@ -83,6 +85,7 @@ def get_mtpls_global(positions, box, params, lmax=2):
     z_atoms = params['axis_indices'][:, 0]
     x_atoms = params['axis_indices'][:, 1]
     y_atoms = params['axis_indices'][:, 2]
+
     vec_z = positions[z_atoms] - positions
     vec_z = pbc_shift(vec_z, box, box_inv)
     vec_z = normalize(vec_z, axis=1)
@@ -172,7 +175,7 @@ def read_mpid_inputs(pdb, xml):
     system = forcefield.createSystem(pdb.topology, nonbondedMethod=LJPME, defaultTholeWidth=5)
 
     generators = {}
-    print([i.__repr__() for i in forcefield.getGenerators()])
+    # print([i.__repr__() for i in forcefield.getGenerators()])
     for force_gen in forcefield.getGenerators():
         generator_name = re.search('.([A-Za-z]+) +object', force_gen.__repr__()).group(1)
         generators[generator_name] = force_gen
@@ -306,7 +309,20 @@ def calc_real_space_energy(self):
     def calc_self_energy(self):
         return pme_self(self.Q, self.lmax, self.kappa)
 
-    def calc_reci_space_energy(self):
-        N = np.array([self.K1, self.K2, self.K3])
-        Q = self.Q[:, :(self.lmax+1)**2].reshape(self.Q.shape[0], (self.lmax+1)**2)
-        return pme_reciprocal(self.positions, self.box, Q, self.lmax, self.kappa, N)
+    def compile_reci_space_energy_and_force(self):
+        # Load the static parameters from self
+        axis_types = self.mpid_params['axis_types']
+        axis_indices = self.mpid_params['axis_indices']
+
+        # Do the calculations
+        pme_reciprocal_energy = gen_pme_reciprocal(axis_types, axis_indices)
+        self.__reci_space_energy_and_force = jit(value_and_grad(pme_reciprocal_energy), static_argnums=(4,5,6,7))
+
+    def calc_reci_space_energy_and_force(self):
+        Q_lc = np.concatenate((np.expand_dims(self.mpid_params['charges'], axis=1), self.mpid_params['dipoles'], self.mpid_params['quadrupoles']), axis=1)
+        Q_lh = convert_cart2harm(Q_lc, lmax=2)
+
+        # Do the calculations
+
+        ene, f = self.__reci_space_energy_and_force(self.positions, self.box, Q_lh, self.kappa, self.lmax, self.K1, self.K2, self.K3)
+        return ene, f