Skip to content

Commit

Permalink
Update omm plugin and save_dmff2tf.py for aux (U_ind especially) comp…
Browse files Browse the repository at this point in the history 
…atible. (#135)

* Add issue templates for feature request and bug-report

* Add script for dmff model saving.

* Remove issue template from devel branch.

* debug workflow

* remove debug

* Update ut.yml. Install mdtraj by conda.

* Add openmm dmff plugin to backend

* Update .gitignore

* Add test workflow for openmm dmff plugin

* Update workflow for openmm dmff plugin test.

* Update workflow for openmm dmff plugin test.

* Update workflow for openmm dmff plugin test.

* Update workflow for openmm dmff plugin test.

* Update workflow for openmm dmff plugin test.

* Specific setuptools version in workflow yaml

* Update tensorflow version

* Enable double precision in test omm dmff plugin workflow

* update nve assert checking in omm dmff plugin workflow

* Fix no axis type bug (#120)

* add rules of local axis for NoAxisType

* add test for NoAxisType

* Update omm plugin and save_dmff2tf.py for aux compatible.

* Fix test issue.

* Update test_openmm_dmff_plugin.yml

---------

Co-authored-by: plumbum082 <[email protected]>
  • Loading branch information
@dingye18 @plumbum082
dingye18 and plumbum082 authored Nov 8, 2023
1 parent 8ffb4fb commit ad1bf64
Show file tree
Hide file tree
Showing 19 changed files with 257 additions and 266 deletions.
7 changes: 4 additions & 3 deletions .github/workflows/test_openmm_dmff_plugin.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
16 changes: 15 additions & 1 deletion backend/openmm_dmff_plugin/openmmapi/include/DMFFForce.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ typedef double ENERGYTYPE;
#else
typedef float FORCETYPE;
typedef float COORDTYPE;
typedef float ENERGYTYPE;
typedef double ENERGYTYPE;
#endif

namespace DMFFPlugin {
Expand All @@ -71,6 +71,12 @@ class OPENMM_EXPORT_DMFF DMFFForce : public OpenMM::Force {
* @param energyCoefficient : the energy transform coefficient.
*/
void setUnitTransformCoefficients(const double coordCoefficient, const double forceCoefficient, const double energyCoefficient);
/**
* @brief Set the has_aux flag when model was saved with auxilary input.
*
* @param hasAux : true if model was saved with auxilary input.
*/
void setHasAux(const bool hasAux);

const std::string& getDMFFGraphFile() const;
/**
Expand All @@ -97,6 +103,13 @@ class OPENMM_EXPORT_DMFF DMFFForce : public OpenMM::Force {
* @return double
*/
double getCutoff() const;
/**
* @brief Get the Has Aux object
*
* @return true
* @return false
*/
bool getHasAux() const;
void updateParametersInContext(OpenMM::Context& context);
bool usesPeriodicBoundaryConditions() const {
return use_pbc;
Expand All @@ -106,6 +119,7 @@ class OPENMM_EXPORT_DMFF DMFFForce : public OpenMM::Force {
private:
string graph_file;
bool use_pbc = true;
bool has_aux = false;
double cutoff = 1.2;
double coordCoeff, forceCoeff, energyCoeff;

Expand Down
6 changes: 6 additions & 0 deletions backend/openmm_dmff_plugin/openmmapi/src/DMFFForce.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,12 +54,18 @@ void DMFFForce::setUnitTransformCoefficients(const double coordCoefficient, cons
energyCoeff = energyCoefficient;
}

void DMFFForce::setHasAux(const bool hasAux){
this->has_aux = hasAux;
}

double DMFFForce::getCoordUnitCoefficient() const {return coordCoeff;}
double DMFFForce::getForceUnitCoefficient() const {return forceCoeff;}
double DMFFForce::getEnergyUnitCoefficient() const {return energyCoeff;}

double DMFFForce::getCutoff() const {return cutoff;}

bool DMFFForce::getHasAux() const {return has_aux;}

const string& DMFFForce::getDMFFGraphFile() const{return graph_file;}


Expand Down
6 changes: 5 additions & 1 deletion backend/openmm_dmff_plugin/platforms/cuda/include/CudaDMFFKernels.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,19 +58,23 @@ class CudaCalcDMFFForceKernel : public CalcDMFFForceKernel{
std::string graph_file;
cppflow::model jax_model;
vector<int64_t> coord_shape = vector<int64_t>(2);
vector<int64_t> U_ind_shape = vector<int64_t>(2);
vector<int64_t> box_shape{3, 3};
vector<int64_t> pair_shape = vector<int64_t>(2);
vector<int32_t> pairs_v;
cppflow::tensor coord_tensor, box_tensor, pair_tensor;
cppflow::tensor coord_tensor, box_tensor, pair_tensor, U_ind_tensor;
vector<cppflow::tensor> output_tensors;
vector<double> last_U_ind;
vector<std::string> operations;
vector<std::string> input_node_names = vector<std::string>(3);
vector<std::string> output_node_names = vector<std::string>(2);

OpenMM::NeighborList neighborList;
vector<std::set<int>> exclusions;

int natoms;
double cutoff;
bool has_aux;
ENERGYTYPE dener;
vector<FORCETYPE> dforce;
vector<COORDTYPE> dcoord;
Expand Down
57 changes: 51 additions & 6 deletions backend/openmm_dmff_plugin/platforms/cuda/src/CudaDMFFKernels.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,25 +50,50 @@ void CudaCalcDMFFForceKernel::initialize(const System& system, const DMFFForce&
energyUnitCoeff = force.getEnergyUnitCoefficient();
coordUnitCoeff = force.getCoordUnitCoefficient();
cutoff = force.getCutoff();
this->has_aux = force.getHasAux();

natoms = system.getNumParticles();
coord_shape[0] = natoms;
coord_shape[1] = 3;
exclusions.resize(natoms);

if (this->has_aux){
U_ind_shape[0] = natoms;
U_ind_shape[1] = 3;
// Initialize the last_U_ind.
for(int ii = 0; ii < natoms * 3; ii ++){
last_U_ind.push_back(0.0);
}
}

// Load the ordinary graph firstly.
jax_model.init(graph_file);

operations = jax_model.get_operations();
for (int ii = 0; ii < operations.size(); ii++){
if (operations[ii].find("serving")!= std::string::npos){
if (operations[ii].find("0")!= std::string::npos){
if (operations[ii].find("serving") != std::string::npos){
if (operations[ii].find("0") != std::string::npos){
input_node_names[0] = operations[ii] + ":0";
} else if (operations[ii].find("1") != std::string::npos){
input_node_names[1] = operations[ii] + ":0";
} else if (operations[ii].find("2") != std::string::npos){
input_node_names[2] = operations[ii] + ":0";
}
// Set up the auxilary input node name. For U_ind
if(this->has_aux){
if (operations[ii].find("3") != std::string::npos){
input_node_names.push_back(operations[ii] + ":0");
}
}
}
// Set up the output names.
if (operations[ii].find("PartitionedCall") != std::string::npos){
output_node_names[0] = operations[ii] + ":0";
output_node_names[1] = operations[ii] + ":1";
if(this->has_aux){
output_node_names.push_back(operations[ii] + ":2");
}
break;
}
}

Expand Down Expand Up @@ -123,6 +148,9 @@ double CudaCalcDMFFForceKernel::execute(ContextImpl& context, bool includeForces
}
coord_tensor = cppflow::tensor(dcoord, coord_shape);

// Set input U_ind
U_ind_tensor = cppflow::tensor(last_U_ind, U_ind_shape);

// Fetch the neighbor list for input pairs tensor.
computeNeighborListVoxelHash(
neighborList,
Expand All @@ -145,10 +173,27 @@ double CudaCalcDMFFForceKernel::execute(ContextImpl& context, bool includeForces
pair_tensor = cppflow::tensor(pairs_v, pair_shape);

// Calculate the energy and forces.
output_tensors = 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_tensors[0].get_data<ENERGYTYPE>()[0];
dforce = output_tensors[1].get_data<FORCETYPE>();
if (!this->has_aux){
output_tensors = jax_model({
{input_node_names[0], coord_tensor},
{input_node_names[1], box_tensor},
{input_node_names[2], pair_tensor}},
{output_node_names[0], output_node_names[1]});
dener = output_tensors[0].get_data<ENERGYTYPE>()[0];
dforce = output_tensors[1].get_data<FORCETYPE>();
} else {
output_tensors = jax_model({
{input_node_names[0], coord_tensor},
{input_node_names[1], box_tensor},
{input_node_names[2], U_ind_tensor},
{input_node_names[3], pair_tensor}},
{output_node_names[0], output_node_names[1], output_node_names[2]});

dener = output_tensors[0].get_data<ENERGYTYPE>()[0];
dforce = output_tensors[1].get_data<FORCETYPE>();
// Save last U_ind for next step usage.
last_U_ind = output_tensors[2].get_data<double>();
}


// Transform the unit from eV/A to KJ/(mol*nm)
Expand Down
10 changes: 7 additions & 3 deletions backend/openmm_dmff_plugin/platforms/reference/include/ReferenceDMFFKernels.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,21 +75,25 @@ class ReferenceCalcDMFFForceKernel : public CalcDMFFForceKernel {
cppflow::model jax_model;

std::vector<int64_t> coord_shape = vector<int64_t>(2);
std::vector<int64_t> coord_shape_1 = vector<int64_t>(2);
std::vector<int64_t> coord_shape_2 = vector<int64_t>(2);
vector<int64_t> U_ind_shape = vector<int64_t>(2);

std::vector<int64_t> box_shape{3, 3};
std::vector<int64_t> pair_shape = vector<int64_t>(2);

std::vector<cppflow::tensor> output;
cppflow::tensor coord_tensor, box_tensor, pair_tensor;
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);
vector<std::string> output_node_names = vector<std::string>(2);

OpenMM::NeighborList neighborList;
std::vector<std::set<int>> exclusions;

int natoms;
double cutoff;
bool has_aux;
ENERGYTYPE dener;
vector<FORCETYPE> dforce;
vector<COORDTYPE> dcoord;
Expand Down
58 changes: 51 additions & 7 deletions backend/openmm_dmff_plugin/platforms/reference/src/ReferenceDMFFKernels.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,12 +70,22 @@ void ReferenceCalcDMFFForceKernel::initialize(const System& system, const DMFFFo
energyUnitCoeff = force.getEnergyUnitCoefficient();
coordUnitCoeff = force.getCoordUnitCoefficient();
cutoff = force.getCutoff();
this->has_aux = force.getHasAux();

natoms = system.getNumParticles();
coord_shape[0] = natoms;
coord_shape[1] = 3;
exclusions.resize(natoms);

if (this->has_aux){
U_ind_shape[0] = natoms;
U_ind_shape[1] = 3;
// Initialize the last_U_ind.
for(int ii = 0; ii < natoms * 3; ii ++){
last_U_ind.push_back(0.0);
}
}

// Load the ordinary graph firstly.
jax_model.init(graph_file);

Expand All @@ -88,9 +98,22 @@ void ReferenceCalcDMFFForceKernel::initialize(const System& system, const DMFFFo
input_node_names[1] = operations[ii]+":0";
} else if (operations[ii].find("2") != std::string::npos){
input_node_names[2] = operations[ii]+":0";
} else {
std::cout << "Warning: Unknown input node name: " << operations[ii] << std::endl;
}
// Set up the auxilary input node name. For U_ind
if(this->has_aux){
if (operations[ii].find("3") != std::string::npos){
input_node_names.push_back(operations[ii] + ":0");
}
}
}
// Set up the output names.
if (operations[ii].find("PartitionedCall") != std::string::npos){
output_node_names[0] = operations[ii] + ":0";
output_node_names[1] = operations[ii] + ":1";
if(this->has_aux){
output_node_names.push_back(operations[ii] + ":2");
}
break;
}
}

Expand Down Expand Up @@ -134,6 +157,9 @@ double ReferenceCalcDMFFForceKernel::execute(ContextImpl& context, bool includeF
}
coord_tensor = cppflow::tensor(dcoord, coord_shape);

// Set input U_ind.
U_ind_tensor = cppflow::tensor(last_U_ind, U_ind_shape);

// Set input pairs.
computeNeighborListVoxelHash(
neighborList,
Expand All @@ -156,12 +182,30 @@ double ReferenceCalcDMFFForceKernel::execute(ContextImpl& context, bool includeF
}
pair_shape[0] = totpairs;
pair_shape[1] = 2;
cppflow::tensor pair_tensor = cppflow::tensor(dpairs, pair_shape);

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"});
pair_tensor = cppflow::tensor(dpairs, pair_shape);

if (!this->has_aux){
output_tensors = jax_model({
{input_node_names[0], coord_tensor},
{input_node_names[1], box_tensor},
{input_node_names[2], pair_tensor}},
{output_node_names[0], output_node_names[1]});
dener = output_tensors[0].get_data<ENERGYTYPE>()[0];
dforce = output_tensors[1].get_data<FORCETYPE>();
} else {
output_tensors = jax_model({
{input_node_names[0], coord_tensor},
{input_node_names[1], box_tensor},
{input_node_names[2], U_ind_tensor},
{input_node_names[3], pair_tensor}},
{output_node_names[0], output_node_names[1], output_node_names[2]});

dener = output_tensors[0].get_data<ENERGYTYPE>()[0];
dforce = output_tensors[1].get_data<FORCETYPE>();
// Save last U_ind for next step usage.
last_U_ind = output_tensors[2].get_data<double>();
}

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 ++){
Expand Down
1 change: 1 addition & 0 deletions backend/openmm_dmff_plugin/python/OpenMMDMFFPlugin.i
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ public:
DMFFForce(const string& GraphFile);

void setUnitTransformCoefficients(const double coordCoefficient, const double forceCoefficient, const double energyCoefficient);
void setHasAux(const bool hasAux);

/*
* Add methods for casting a Force to a DMFFForce.
Expand Down
1 change: 1 addition & 0 deletions backend/openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/admp_water_dimer_aux/fingerprint.pb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
��ж��������ؓ������Ԭְ��� ��������-(�����ۄ��2
Binary file added backend/openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/admp_water_dimer_aux/saved_model.pb
View file
Binary file not shown.
0 ...u/variables/variables.data-00000-of-00001 → ...x/variables/variables.data-00000-of-00001
View file
File renamed without changes.
0 ...ta/lj_fluid_gpu/variables/variables.index → ...water_dimer_aux/variables/variables.index
View file
File renamed without changes.
Loading

0 comments on commit ad1bf64

Please sign in to comment.