Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update the doc for openmm dmff plugin usage. #146

Merged
merged 35 commits into from
Nov 9, 2023
Merged
Changes from 1 commit
Commits
Show all changes
35 commits
Select commit Hold shift + click to select a range
f533581
Add issue templates for feature request and bug-report
dingye18 Feb 13, 2023
3f4d29b
Merge branch 'devel' of https://github.com/dingye18/DMFF into devel
dingye18 Jul 3, 2023
f1a7919
Merge branch 'devel' of https://github.com/deepmodeling/DMFF into devel
dingye18 Jul 3, 2023
3cdf6a1
Merge branch 'devel' of https://github.com/deepmodeling/DMFF into devel
dingye18 Jul 30, 2023
bec7eb4
Merge branch 'devel' of https://github.com/deepmodeling/DMFF into devel
dingye18 Aug 9, 2023
d11057a
Merge branch 'devel' of https://github.com/dingye18/DMFF into devel
dingye18 Aug 18, 2023
5fbd560
Add script for dmff model saving.
dingye18 Aug 18, 2023
63ffd61
Remove issue template from devel branch.
dingye18 Aug 18, 2023
38b2b2f
debug workflow
dingye18 Aug 18, 2023
cbb35b8
remove debug
dingye18 Aug 18, 2023
14a45cd
Update ut.yml. Install mdtraj by conda.
dingye18 Aug 20, 2023
e816bd1
Merge branch 'deepmodeling:devel' into devel
dingye18 Oct 8, 2023
8e9bede
Add openmm dmff plugin to backend
dingye18 Oct 9, 2023
7eedaad
Update .gitignore
dingye18 Oct 9, 2023
0d2aaad
Add test workflow for openmm dmff plugin
dingye18 Oct 10, 2023
1a93eed
Update workflow for openmm dmff plugin test.
dingye18 Oct 10, 2023
c3b060b
Update workflow for openmm dmff plugin test.
dingye18 Oct 10, 2023
925a09a
Update workflow for openmm dmff plugin test.
dingye18 Oct 10, 2023
fc5aa00
Update workflow for openmm dmff plugin test.
dingye18 Oct 10, 2023
c60de88
Update workflow for openmm dmff plugin test.
dingye18 Oct 10, 2023
22bf802
Specific setuptools version in workflow yaml
dingye18 Oct 10, 2023
2ba5d30
Update tensorflow version
dingye18 Oct 10, 2023
3f04d60
Enable double precision in test omm dmff plugin workflow
dingye18 Oct 10, 2023
86d47e2
update nve assert checking in omm dmff plugin workflow
dingye18 Oct 10, 2023
e0877ee
Fix no axis type bug (#120)
plumbum082 Oct 20, 2023
ac94a96
Merge branch 'deepmodeling:devel' into devel
dingye18 Nov 5, 2023
fa695a6
Merge branch 'wangxy/v1.0.0-devel' of https://github.com/deepmodeling…
dingye18 Nov 5, 2023
447a73a
Update omm plugin and save_dmff2tf.py for aux compatible.
dingye18 Nov 5, 2023
33d622b
Fix test issue.
dingye18 Nov 6, 2023
26b44f6
Update test_openmm_dmff_plugin.yml
dingye18 Nov 6, 2023
f8fe2fe
Update the doc for openmm dmff plugin usage
dingye18 Nov 8, 2023
236cb1c
Merge branch 'wangxy/v1.0.0-devel' of https://github.com/deepmodeling…
dingye18 Nov 8, 2023
6416580
Revision on the doc of DMFF plugin
dingye18 Nov 8, 2023
1f7522d
Upload tf c directory.
dingye18 Nov 8, 2023
d740579
Remove tf2.9.1_c.tar
dingye18 Nov 8, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Update omm plugin and save_dmff2tf.py for aux compatible.
dingye18 committed Nov 5, 2023
commit 447a73a46513164b51123e4f97380857f67273b6
16 changes: 15 additions & 1 deletion backend/openmm_dmff_plugin/openmmapi/include/DMFFForce.h
Original file line number Diff line number Diff line change
@@ -50,7 +50,7 @@ typedef double ENERGYTYPE;
#else
typedef float FORCETYPE;
typedef float COORDTYPE;
typedef float ENERGYTYPE;
typedef double ENERGYTYPE;
#endif

namespace DMFFPlugin {
@@ -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;
/**
@@ -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;
@@ -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;

6 changes: 6 additions & 0 deletions backend/openmm_dmff_plugin/openmmapi/src/DMFFForce.cpp
Original file line number Diff line number Diff line change
@@ -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;}


Original file line number Diff line number Diff line change
@@ -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;
57 changes: 51 additions & 6 deletions backend/openmm_dmff_plugin/platforms/cuda/src/CudaDMFFKernels.cpp
Original file line number Diff line number Diff line change
@@ -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;
}
}

@@ -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,
@@ -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)
Original file line number Diff line number Diff line change
@@ -75,21 +75,24 @@ 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<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;
Original file line number Diff line number Diff line change
@@ -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);

@@ -88,10 +98,23 @@ 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;
}
}

// Initialize the ordinary input and output array.
@@ -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,
@@ -156,7 +182,29 @@ double ReferenceCalcDMFFForceKernel::execute(ContextImpl& context, bool includeF
}
pair_shape[0] = totpairs;
pair_shape[1] = 2;
cppflow::tensor pair_tensor = cppflow::tensor(dpairs, pair_shape);
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>();
}

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"});

1 change: 1 addition & 0 deletions backend/openmm_dmff_plugin/python/OpenMMDMFFPlugin.i
Original file line number Diff line number Diff line change
@@ -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.
9 changes: 9 additions & 0 deletions backend/openmm_dmff_plugin/python/OpenMMDMFFPlugin/tools.py
Original file line number Diff line number Diff line change
@@ -70,6 +70,15 @@ def setUnitTransformCoefficients(self, coordinatesCoefficient, forceCoefficient,
self.dmff_force.setUnitTransformCoefficients(coordinatesCoefficient, forceCoefficient, energyCoefficient)
return

def setHasAux(self, has_aux = False):
"""Set whether the DMFF model has auxilary output.
Used when model was saved with has_aux = True.

Args:
has_aux (bool, optional): Defaults to False.
"""
self.dmff_force.setHasAux(has_aux)

def createSystem(self, topology):
"""Create the OpenMM System object for the DMFF model.

Loading