Fix test issue.

.github/workflows/test_openmm_dmff_plugin.yml

@@ -19,7 +19,7 @@ jobs:
       - name: Install Dependencies
         run: |
           source $CONDA/bin/activate
-          conda create -n dmff_omm -y python=${{ matrix.python-version }} numpy openmm=7.7 -c conda-forge
+          conda create -n dmff_omm -y python=${{ matrix.python-version }} numpy openmm=7.7 rdkit openbabel mdtraj ambertools -c conda-forge
           conda activate dmff_omm
           conda install -y libtensorflow_cc=2.9.1 -c conda-forge
           pip install setuptools==59.5.0
@@ -41,10 +41,11 @@ jobs:
           export OPENMM_INSTALLED_DIR=$CONDA_PREFIX
           export CPPFLOW_INSTALLED_DIR=$CONDA_PREFIX
           export LIBTENSORFLOW_INSTALLED_DIR=$CONDA_PREFIX
-          cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR} -DUSE_HIGH_PRECISION=ON
+          cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR} -DUSE_HIGH_PRECISION=OFF
           make && make install 
-          make PythonInstall
+          make PythonInstall          
       - name: Run Tests
         run: |
           source $CONDA/bin/activate dmff_omm
-          python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nve -n 100
+          cd ${GITHUB_WORKSPACE}/backend/
+          python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nve -n 100 --pdb ../examples/water_fullpol/water_dimer.pdb --model ./openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/admp_water_dimer_aux --has_aux True

backend/openmm_dmff_plugin/platforms/reference/include/ReferenceDMFFKernels.h

@@ -82,6 +82,7 @@ class ReferenceCalcDMFFForceKernel : public CalcDMFFForceKernel {
 
     std::vector<cppflow::tensor> output;
     cppflow::tensor coord_tensor, box_tensor, pair_tensor, U_ind_tensor;
+    vector<cppflow::tensor> output_tensors;
     vector<std::string> operations;
     vector<double> last_U_ind;
     vector<std::string> input_node_names = vector<std::string>(3);

backend/openmm_dmff_plugin/platforms/reference/src/ReferenceDMFFKernels.cpp

@@ -206,10 +206,6 @@ double ReferenceCalcDMFFForceKernel::execute(ContextImpl& context, bool includeF
         last_U_ind = output_tensors[2].get_data<double>();
     }
 
-    output = jax_model({{input_node_names[0], coord_tensor}, {input_node_names[1], box_tensor}, {input_node_names[2], pair_tensor}}, {"PartitionedCall:0", "PartitionedCall:1"});
-
-    dener = output[0].get_data<ENERGYTYPE>()[0];
-    dforce = output[1].get_data<FORCETYPE>();
 
     // Transform the unit from output units to KJ/(mol*nm)
     for(int ii = 0; ii < natoms; ii ++){

backend/openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/admp_water_dimer_aux/fingerprint.pb

@@ -0,0 +1 @@
+������
���ؓ����
��Ԭְ���
 ��������-(�����ۄ��
2

backend/openmm_dmff_plugin/python/OpenMMDMFFPlugin/tests/test_dmff_plugin_nve.py

@@ -15,24 +15,20 @@
 from OpenMMDMFFPlugin import DMFFModel
 
 
-def test_dmff_nve(nsteps = 1000, time_step = 0.2, platform_name = "Reference", output_temp_dir = "/tmp/openmm_dmff_plugin_test_nve_output", energy_std_tol = 0.0005 ):
+def test_dmff_nve(nsteps = 1000, time_step = 0.2, pdb_file = None, model_dir = None, platform_name = "Reference", output_temp_dir = "/tmp/openmm_dmff_plugin_test_nve_output", energy_std_tol = 0.005, has_aux = False ):
     if not os.path.exists(output_temp_dir):
         os.mkdir(output_temp_dir)
 
-    pdb_file = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid.pdb")
-    if platform_name == "Reference":
-        dmff_model_file = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid_gpu")
-    elif platform_name == "CUDA":
-        dmff_model_file = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid_gpu")
-
-    output_dcd = os.path.join(output_temp_dir, "lj_fluid_test.nve.dcd")
-    output_log = os.path.join(output_temp_dir, "lj_fluid_test.nve.log")
+    dmff_model_file = model_dir
+
+    output_dcd = os.path.join(output_temp_dir, "test.nve.dcd")
+    output_log = os.path.join(output_temp_dir, "test.nve.log")
 
     # Set up the simulation parameters.
     nsteps = nsteps
     time_step = time_step # unit is femtosecond.
-    report_frequency = 10
-    box = [24.413, 0, 0, 0, 24.413, 0, 0, 0, 24.413]
+    report_frequency = 1
+    box = [31.289, 0, 0, 0, 31.289, 0, 0, 0, 31.289]
     box = [mm.Vec3(box[0], box[1], box[2]), mm.Vec3(box[3], box[4], box[5]), mm.Vec3(box[6], box[7], box[8])] * u.angstroms
 
     liquid_water = PDBFile(pdb_file)
@@ -43,6 +39,8 @@ def test_dmff_nve(nsteps = 1000, time_step = 0.2, platform_name = "Reference", o
     # Set up the dmff_system with the dmff_model.    
     dmff_model = DMFFModel(dmff_model_file)
     dmff_model.setUnitTransformCoefficients(1, 1, 1)
+    if has_aux:
+        dmff_model.setHasAux(True)
     dmff_system = dmff_model.createSystem(topology)
 
     integrator = mm.VerletIntegrator(time_step*u.femtoseconds)
@@ -83,19 +81,30 @@ def test_dmff_nve(nsteps = 1000, time_step = 0.2, platform_name = "Reference", o
     # Check the total energy fluctuations over # of atoms is smaller than energy_std_tol, unit in kJ/mol.
     print("Total energy std: %.4f kJ/mol"%(np.std(total_energy)))
     print("Mean total energy: %.4f kJ/mol"%(np.mean(total_energy)))
-    #assert(np.std(total_energy) / num_atoms < energy_std_tol)    
+    assert(np.std(total_energy) / num_atoms < energy_std_tol)    
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('-n', '--nsteps', type = int, dest='nsteps', help='Number of steps', default=100)
     parser.add_argument('--dt', type = float, dest='timestep', help='Time step for simulation, unit is femtosecond', default=0.2)
+    parser.add_argument('--pdb', type = str, dest='pdb', help='PDB file for simulation.', default=None)
+    parser.add_argument('--model', type = str, dest='model', help='DMFF model dir for simulation. Saved by backend/save_dmff2tf.py.', default=None)
     parser.add_argument('--platform', type = str, dest='platform', help='Platform for simulation.', default="Reference")
+    parser.add_argument('--has_aux', type = bool, dest='has_aux', help='Whether the model has aux output.', default=False)
 
     args = parser.parse_args()
 
     nsteps = args.nsteps
     time_step = args.timestep
     platform_name = args.platform
 
-    test_dmff_nve(nsteps=nsteps, time_step=time_step, platform_name=platform_name)
+    pdb = args.pdb
+    model_dir = args.model
+
+    if pdb is None:
+        pdb = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid.pdb")
+    if model_dir is None:
+        model_dir = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid_gpu")
+
+    test_dmff_nve(nsteps=nsteps, time_step=time_step, pdb_file=pdb, model_dir=model_dir, platform_name=platform_name, has_aux=args.has_aux)
 

backend/openmm_dmff_plugin/python/tests/test_dmff_plugin_nve.py

@@ -15,18 +15,14 @@
 from OpenMMDMFFPlugin import DMFFModel
 
 
-def test_dmff_nve(nsteps = 1000, time_step = 0.2, platform_name = "Reference", output_temp_dir = "/tmp/openmm_dmff_plugin_test_nve_output", energy_std_tol = 0.005 ):
+def test_dmff_nve(nsteps = 1000, time_step = 0.2, pdb_file = None, model_dir = None, platform_name = "Reference", output_temp_dir = "/tmp/openmm_dmff_plugin_test_nve_output", energy_std_tol = 0.005, has_aux = False ):
     if not os.path.exists(output_temp_dir):
         os.mkdir(output_temp_dir)
 
-    pdb_file = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid.pdb")
-    if platform_name == "Reference":
-        dmff_model_file = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid_gpu")
-    elif platform_name == "CUDA":
-        dmff_model_file = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid_gpu")
-
-    output_dcd = os.path.join(output_temp_dir, "lj_fluid_test.nve.dcd")
-    output_log = os.path.join(output_temp_dir, "lj_fluid_test.nve.log")
+    dmff_model_file = model_dir
+
+    output_dcd = os.path.join(output_temp_dir, "test.nve.dcd")
+    output_log = os.path.join(output_temp_dir, "test.nve.log")
 
     # Set up the simulation parameters.
     nsteps = nsteps
@@ -43,6 +39,8 @@ def test_dmff_nve(nsteps = 1000, time_step = 0.2, platform_name = "Reference", o
     # Set up the dmff_system with the dmff_model.    
     dmff_model = DMFFModel(dmff_model_file)
     dmff_model.setUnitTransformCoefficients(1, 1, 1)
+    if has_aux:
+        dmff_model.setHasAux()
     dmff_system = dmff_model.createSystem(topology)
 
     integrator = mm.VerletIntegrator(time_step*u.femtoseconds)
@@ -89,13 +87,24 @@ def test_dmff_nve(nsteps = 1000, time_step = 0.2, platform_name = "Reference", o
     parser = argparse.ArgumentParser()
     parser.add_argument('-n', '--nsteps', type = int, dest='nsteps', help='Number of steps', default=100)
     parser.add_argument('--dt', type = float, dest='timestep', help='Time step for simulation, unit is femtosecond', default=0.2)
+    parser.add_argument('--pdb', type = str, dest='pdb', help='PDB file for simulation.', default=None)
+    parser.add_argument('--model', type = str, dest='model', help='DMFF model dir for simulation. Saved by backend/save_dmff2tf.py.', default=None)
     parser.add_argument('--platform', type = str, dest='platform', help='Platform for simulation.', default="Reference")
+    parser.add_argument('--has_aux', type = bool, dest='has_aux', help='Whether the model has aux output.', default=False)
 
     args = parser.parse_args()
 
     nsteps = args.nsteps
     time_step = args.timestep
     platform_name = args.platform
 
-    test_dmff_nve(nsteps=nsteps, time_step=time_step, platform_name=platform_name)
+    pdb = args.pdb
+    model_dir = args.model
+
+    if pdb is None:
+        pdb = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid.pdb")
+    if model_dir is None:
+        model_dir = os.path.join(os.path.dirname(__file__), "../data", "lj_fluid_gpu")
+
+    test_dmff_nve(nsteps=nsteps, time_step=time_step, pdb_file=pdb, model_dir=model_dir, platform_name=platform_name)
 

backend/save_dmff2tf.py

@@ -17,9 +17,12 @@
 for gpu in gpus:
   tf.config.experimental.set_memory_growth(gpu, True)
 
-def create_dmff_potential(input_pdb_file, ff_xml_files, has_aux=False):
+def create_dmff_potential(input_pdb_file, ff_xml_files, bond_definitions_xml = None, has_aux = False):
     pdb = app.PDBFile(input_pdb_file)
     h = dmff.Hamiltonian(*ff_xml_files)
+    if bond_definitions_xml is not None:
+        app.Topology.loadBondDefinitions(bond_definitions_xml)
+
     if has_aux:# Used when using ADMP with DMFF.
         pot = h.createPotential(pdb.topology,
                             nonbondedMethod=app.PME,
@@ -32,6 +35,7 @@ def create_dmff_potential(input_pdb_file, ff_xml_files, has_aux=False):
                             nonbondedMethod=app.PME,
                             nonbondedCutoff=1.2 *
                             unit.nanometer)
+
     pot_func = pot.getPotentialFunc()
     a, b, c = pdb.topology.getPeriodicBoxVectors()
     a = a.value_in_unit(unit.nanometer)
@@ -71,20 +75,22 @@ def potential_engrad(positions, box, pairs):
     parser.add_argument("--input_pdb", dest="input_pdb", help="input pdb file. Box information is required in the pdb file.")
     parser.add_argument("--xml_files", dest="xml_files", nargs="+", help=".xml files with parameters are derived from DMFF.")
     parser.add_argument("--output", dest="output", help="output directory")
+    parser.add_argument("--bond_definitions_xml", dest="bond_definitions_xml", help=".xml file that contains bond definitions. Optional", default=None)
     parser.add_argument("--has_aux", dest="has_aux", default=False, help="Enable aux output in the model. Used when using ADMP with DMFF, and the output would be U_ind.")
     args = parser.parse_args()
 
     input_pdb = args.input_pdb
     ff_xml_files = args.xml_files
     output_dir = args.output
     has_aux = args.has_aux
+    bond_definitions_xml = args.bond_definitions_xml
 
     if output_dir[-1] == "/":
         output_dir = output_dir[:-1]
     if not os.path.exists(output_dir):
         os.mkdir(output_dir)
 
-    pdb, pot_grad, covalent_map, pot, h = create_dmff_potential(input_pdb, ff_xml_files, has_aux=has_aux)
+    pdb, pot_grad, covalent_map, pot, h = create_dmff_potential(input_pdb, ff_xml_files, bond_definitions_xml=bond_definitions_xml, has_aux=has_aux)
 
     natoms = pdb.getTopology().getNumAtoms()