Make lambda value to be gloabl parameter, so that we can change dynamically in OpenMM context.

.conda/build.sh

@@ -8,6 +8,11 @@ cd build
 wget https://github.com/deepmodeling/deepmd-kit/releases/latest/download/libdeepmd_c.tar.gz
 tar -xf libdeepmd_c.tar.gz -C ${PREFIX}
 
+mkdir -p ${PREFIX}/lib/libdeepmd_c
+mkdir -p ${PREFIX}/include/libdeepmd_c
+cp -r ${PREFIX}/libdeepmd_c/include/* ${PREFIX}/include/
+cp -r ${PREFIX}/libdeepmd_c/lib/* ${PREFIX}/lib/
+
 cmake -DOPENMM_DIR=${PREFIX} -DDEEPMD_DIR=${PREFIX}/libdeepmd_c ..
 
 make #-j${NUM_CPUS}

.conda/meta.yaml

@@ -1,5 +1,5 @@
 {% set name = "openmm_deepmd_plugin" %}
-{% set version = "0.2.0" %}
+{% set version = "0.3.0" %}
 {% set cuda_compiler_version = "12.0" %}
 {% set py_version = "3.11" %}
 
@@ -9,12 +9,13 @@ package:
   version: {{ version }}
 
 source:
-  git_url: https://github.com/JingHuangLab/openmm_deepmd_plugin.git
-  git_rev: 43e319f75e8597eac7d2060ee475e1bd6d30e674
-  patches:
-    - cuda_platform.patch
+  #git_url: https://github.com/JingHuangLab/openmm_deepmd_plugin.git
+  #git_rev: 43e319f75e8597eac7d2060ee475e1bd6d30e674
+  #patches:
+  #  - cuda_platform.patch
+  path: ..
   build:
-    number: 2
+    number: 1
     string: "py{{ py_version }}_cuda{{ cuda_compiler_version }}_{{ build_number|default(1) }}"
     skip: True  # [not linux]
     missing_dso_whitelist:

openmmapi/include/DeepmdForce.h

@@ -65,9 +65,9 @@ class OPENMM_EXPORT_DEEPMD DeepmdForce : public OpenMM::Force {
      * @brief Construct a new Deepmd Force object. Used when running with specific lambda.
      * 
      * @param GraphFile 
-     * @param lambda
+     * @param lambda_name : the name of the lambda parameter in the OpenMM context.
      */
-    DeepmdForce(const string& GraphFile, const double& lambda);
+    DeepmdForce(const string& GraphFile, const string& lambda_name);
     /**
      * @brief Destroy the Deepmd Force object.
      * 
@@ -199,18 +199,6 @@ class OPENMM_EXPORT_DEEPMD DeepmdForce : public OpenMM::Force {
     void addAtom(int resIndex, string AtomName, string AtomElement, int atomIndex, int atomId);
 
     Topology* getTopology() const;
-    /**
-     * @brief Set the lambda value for this alchemical simulation.
-     * 
-     * @param lambda 
-     */
-    void setLambda(const double lambda);
-    /**
-     * @brief Get the lambda value for DP force scale weights in simulation.
-     * 
-     * @return double 
-     */
-    double getLambda() const;
 
     /**
      * @brief Set the GPU rank index for DP model evaluation.
@@ -225,6 +213,11 @@ class OPENMM_EXPORT_DEEPMD DeepmdForce : public OpenMM::Force {
      */
     int getGPURank() const;
 
+    const std::string& getLambdaName() const;
+    void setLambdaName(const std::string& name);
+    void addLambdaParameter(const std::string& name, double defaultValue);
+    map<string, double> getGlobalParameters() const;
+
     void updateParametersInContext(OpenMM::Context& context);
     bool usesPeriodicBoundaryConditions() const {
         return use_pbc;
@@ -233,14 +226,17 @@ class OPENMM_EXPORT_DEEPMD DeepmdForce : public OpenMM::Force {
     OpenMM::ForceImpl* createImpl() const;
 private:
     string graph_file = "";
-    double lambda = 1.0;
     bool use_pbc = true;
     int gpu_rank = 0;
 
     int numb_types = 0;
     string type_map = "";
     double cutoff = 0.;
-
+
+    string lambda_name = "dp_alchem_lambda";
+    double lambda = 1.0; // Default value for lambda, used for alchemical simulations.
+    map<string, double> globalParameters;
+
     map<int, string> type4EachParticle;
     map<string, vector<int>> particleGroup4EachType;
     map<string, int> typesIndexMap;

openmmapi/include/internal/DeepmdForceImpl.h

@@ -59,9 +59,7 @@ class OPENMM_EXPORT_DEEPMD DeepmdForceImpl : public OpenMM::ForceImpl {
         // This force field doesn't update the state directly.
     }
     double calcForcesAndEnergy(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy, int groups);
-    std::map<std::string, double> getDefaultParameters() {
-        return std::map<std::string, double>(); // This force field doesn't define any parameters.
-    }
+    std::map<std::string, double> getDefaultParameters();
     std::vector<std::string> getKernelNames();
     vector<pair<int, int>> getBondedParticles() const; 
     //void updateParametersInContext(OpenMM::ContextImpl& context);

openmmapi/src/DeepmdForce.cpp

@@ -51,20 +51,24 @@ DeepmdForce::DeepmdForce(const string& GraphFile){
     if (!exists(graph_file)){
         throw OpenMMException("Graph file not found: "+graph_file);
     }
+    globalParameters[lambda_name] = lambda; // Initialize the lambda parameter with default value 1.0
+
     // Initialize dp model
     DeepPot tmp_dp = DeepPot(graph_file);
     this->numb_types = tmp_dp.numb_types();
     this->cutoff = tmp_dp.cutoff();
     tmp_dp.get_type_map(this->type_map);
 }
 
-DeepmdForce::DeepmdForce(const string& GraphFile, const double& lambda = 0){
+DeepmdForce::DeepmdForce(const string& GraphFile, const string& lambda_name = "dp_alchem_lambda"){
     this->graph_file  = GraphFile;
     this->topology = new Topology();
-    this->lambda = lambda;
+    this->lambda_name = lambda_name;
     if (!exists(graph_file)){
         throw OpenMMException("Graph file not found: "+graph_file);
     }
+    globalParameters[lambda_name] = lambda; // Initialize the lambda parameter with default value 1.0
+
     // Initialize dp model
     DeepPot* tmp_dp = new DeepPot(graph_file);
     this->numb_types = tmp_dp->numb_types();
@@ -138,12 +142,6 @@ const vector<pair<int, int>> DeepmdForce::getBondsList() const{
     return bondsList;
 }
 
-void DeepmdForce::setLambda(const double lambda){
-    this->lambda = lambda;
-}
-
-double DeepmdForce::getLambda() const {return lambda;}
-
 void DeepmdForce::setGPURank(const int gpu_rank) {
     this->gpu_rank = gpu_rank;
 }
@@ -197,3 +195,26 @@ void DeepmdForce::updateParametersInContext(Context& context) {
     return;
 }
 
+map<string, double> DeepmdForce::getGlobalParameters() const {
+    map<string, double> parameters;
+    for (const auto& param : globalParameters) {
+        parameters[param.first] = param.second;
+    }
+    return parameters;
+}
+
+void DeepmdForce::addLambdaParameter(const string& name, double defaultValue) {
+    lambda_name = name; // Set the name for the lambda parameter
+    globalParameters[name] = defaultValue; // Add or update the parameter with the default value
+    //globalParameters.push_back(make_pair(name, defaultValue));
+}
+
+void DeepmdForce::setLambdaName(const string& name) {
+    lambda_name = name;
+    globalParameters[name] = lambda; // Initialize the lambda parameter with default value 0.0
+    //globalParameters.push_back(make_pair(name, lambda)); // Add lambda parameter with default value 0.0
+}
+
+const std::string& DeepmdForce::getLambdaName() const {
+    return lambda_name;
+}

openmmapi/src/DeepmdForceImpl.cpp

@@ -66,6 +66,10 @@ vector<pair<int, int>> DeepmdForceImpl::getBondedParticles() const{
     return owner.getBondsList();
 }
 
+std::map<std::string, double> DeepmdForceImpl::getDefaultParameters() {
+    return owner.getGlobalParameters();
+}
+
 std::vector<std::string> DeepmdForceImpl::getKernelNames() {
     std::vector<std::string> names;
     names.push_back(CalcDeepmdForceKernel::Name());

platforms/cuda/include/CudaDeepmdKernels.h

@@ -59,7 +59,9 @@ class CudaCalcDeepmdForceKernel : public CalcDeepmdForceKernel{
     DeepPot dp;    
 
     int natoms, tot_atoms;
-    double lambda = 0.0; // interpolate the DP generated forces.
+    double lambda = 1.0; // interpolate the DP generated forces.
+    string lambda_name = "";
+
     ENERGYTYPE dener;
     vector<VALUETYPE> dforce;
     vector<VALUETYPE> dvirial;

platforms/cuda/src/CudaDeepmdKernels.cpp

@@ -51,7 +51,8 @@ void CudaCalcDeepmdForceKernel::initialize(const System& system, const DeepmdFor
     forceUnitCoeff = force.getForceUnitCoefficient();
     energyUnitCoeff = force.getEnergyUnitCoefficient();
     coordUnitCoeff = force.getCoordUnitCoefficient();
-    lambda = force.getLambda();
+    lambda_name = force.getLambdaName();
+
     natoms = type4EachParticle.size();
     tot_atoms = system.getNumParticles();
 
@@ -126,6 +127,8 @@ void CudaCalcDeepmdForceKernel::initialize(const System& system, const DeepmdFor
 double CudaCalcDeepmdForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
     vector<Vec3> pos;
     context.getPositions(pos);
+    lambda = context.getParameter(lambda_name);
+
     Vec3 box[3];
 
     if (isFixedRegion){

platforms/reference/include/ReferenceDeepmdKernels.h

@@ -75,6 +75,7 @@ class ReferenceCalcDeepmdForceKernel : public CalcDeepmdForceKernel {
 
     int natoms, tot_atoms;
     double lambda = 0.0;
+    string lambda_name = "dp_alchem_lambda";
     ENERGYTYPE dener;
     vector<VALUETYPE> dforce;
     vector<VALUETYPE> dvirial;

platforms/reference/src/ReferenceDeepmdKernels.cpp

@@ -72,7 +72,7 @@ void ReferenceCalcDeepmdForceKernel::initialize(const System& system, const Deep
     forceUnitCoeff = force.getForceUnitCoefficient();
     energyUnitCoeff = force.getEnergyUnitCoefficient();
     coordUnitCoeff = force.getCoordUnitCoefficient();
-    lambda = force.getLambda();
+    lambda_name = force.getLambdaName();
     natoms = type4EachParticle.size();
     tot_atoms = system.getNumParticles();
 
@@ -142,6 +142,8 @@ double ReferenceCalcDeepmdForceKernel::execute(ContextImpl& context, bool includ
     vector<RealVec>& pos = extractPositions(context);
     vector<RealVec>& force = extractForces(context);
 
+    lambda = context.getParameter(lambda_name);
+
     if (isFixedRegion){
         // Set box size.
         if (context.getSystem().usesPeriodicBoundaryConditions()){

python/OpenMMDeepmdPlugin.i

@@ -56,14 +56,16 @@ class DeepmdForce : public OpenMM::Force {
 public:
     //DeepmdForce::DeepmdForce(const string& GraphFile, const string& GraphFile_1, const string& GraphFile_2, const bool used4Alchemical);
     DeepmdForce(const string& GraphFile);
-    DeepmdForce(const string& GraphFile, const double& lambda);
+
+    const std::string& getLambdaName() const;
+    void setLambdaName(const std::string& name);
+    void addLambdaParameter(const std::string& name, double defaultValue);
 
     void addParticle(const int particleIndex, const string particleType);
     void addType(const int typeIndex, const string Type);
     void addBond(const int particle1, const int particle2);
     void setPBC(const bool use_pbc);
     void setUnitTransformCoefficients(const double coordCoefficient, const double forceCoefficient, const double energyCoefficient);
-    void setLambda(const double lambda);
     void setGPURank(const int gpu_rank);
 
     // Extract the model info from dp model.    

python/OpenMMDeepmdPlugin/tools.py

@@ -87,20 +87,23 @@ def DrawScatter(x, y, name, xlabel="Time", ylabel="Force, unit is KJ/(mol*nm)",
 
 
 class DeepPotentialModel():
-    def __init__(self, model_file, Lambda = 1.0) -> None:
+    def __init__(self, model_file, LambdaName = "dp_alchem_lambda", Lambda = 1.0) -> None:
         """Initialize the Deep Potential model.
 
         Args:
             model_file (string): Path to the Deep Potential model file.
             Lambda (float, optional): Weight of the Deep Potential models in use. Output forces and energy from the Deep Potential model will be multiplied by Lambda and return to the OpenMM context. Defaults to 1.0.
         """
         self.model_file = model_file
-        self.dp_force = DeepmdForce(model_file, Lambda)
+        self.dp_force = DeepmdForce(model_file)
         self.cutoff = self.dp_force.getCutoff()
         self.numb_types = self.dp_force.getNumberTypes()
         self.type_map_raw = self.dp_force.getTypesMap()
         self.type_map_dict, self.dp_model_types = self.__decode_type_map(self.type_map_raw)
-
+
+        self.dp_force.addLambdaParameter(LambdaName, Lambda)
+        self.lambda_name = LambdaName
+
         # Set up the atom type
         for atom_type in self.type_map_dict.keys():
             self.dp_force.addType(self.type_map_dict[atom_type], atom_type)
@@ -258,4 +261,14 @@ def addCenterParticlesToAdaptiveDPRegion(
             self.dp_force.setSelNum4EachType(self.dp_model_types, num4type)        
 
         return self.dp_force
-
+
+    def setLambdaName(lambda_name):
+        """Set the name of the lambda parameter for the DP model.
+
+        Args:
+            lambda_name (str): Name of the lambda parameter.
+        """
+        self.lambda_name = lambda_name
+        self.dp_force.setLambdaName(lambda_name)
+        assert self.dp_force.getLambdaName() == lambda_name, "Failed to set the lambda name."
+        return

python/examples/water_solvation_energy/water_solvation_free_energy.py

@@ -25,14 +25,14 @@ def create_alchemical_builder(pdb_file, dp_model_file, dp_model_file1, dp_model_
 
     def get_alchemical_system(lambda_value):
         # Set up the dp_system with the dp_model.    
-        dp_model = DeepPotentialModel(dp_model_file, Lambda = lambda_value)
+        dp_model = DeepPotentialModel(dp_model_file, LambdaName = "dp_alchem_lambda_0", Lambda = lambda_value)
         dp_model.setUnitTransformCoefficients(10.0, 964.8792534459, 96.48792534459)
         # By default, createSystem from dp_model will put all atoms in topology into the DP particles for dp_model.
         dp_system = dp_model.createSystem(topology)
 
         # Initial the other two dp_models for alchemical simulation.
-        dp_model_1 = DeepPotentialModel(dp_model_file1, Lambda = 1 - lambda_value)
-        dp_model_2 = DeepPotentialModel(dp_model_file2, Lambda = 1 - lambda_value)
+        dp_model_1 = DeepPotentialModel(dp_model_file1, LambdaName = "dp_alchem_lambda_1", Lambda = 1 - lambda_value)
+        dp_model_2 = DeepPotentialModel(dp_model_file2, LambdaName = "dp_alchem_lambda_2", Lambda = 1 - lambda_value)
         dp_model_1.setUnitTransformCoefficients(10.0, 964.8792534459, 96.48792534459)
         dp_model_2.setUnitTransformCoefficients(10.0, 964.8792534459, 96.48792534459)
 
@@ -127,11 +127,11 @@ def reevaluate_energy(simulation, dcd_files, pdb_file):
 
 if __name__ == "__main__":
 
-    pdb_file = os.path.join(os.path.dirname(__file__), "openmm_deepmd_plugin/python/OpenMMDeepmdPlugin/data", "lw_256_test.pdb")
+    pdb_file = os.path.join(os.path.dirname(__file__), "../../OpenMMDeepmdPlugin/data", "lw_256_test.pdb")
 
-    dp_model_file = os.path.join(os.path.dirname(__file__), "openmm_deepmd_plugin/python/OpenMMDeepmdPlugin/data", "water.pb")
-    dp_model_file1 = os.path.join(os.path.dirname(__file__), "openmm_deepmd_plugin/python/OpenMMDeepmdPlugin/data", "water.pb")
-    dp_model_file2 = os.path.join(os.path.dirname(__file__), "openmm_deepmd_plugin/python/OpenMMDeepmdPlugin/data", "water.pb")
+    dp_model_file = os.path.join(os.path.dirname(__file__), "../../OpenMMDeepmdPlugin/data", "water.pb")
+    dp_model_file1 = os.path.join(os.path.dirname(__file__), "../../OpenMMDeepmdPlugin/data", "water.pb")
+    dp_model_file2 = os.path.join(os.path.dirname(__file__), "../../OpenMMDeepmdPlugin/data", "water.pb")
 
     output_temp_dir = "/tmp/omm_dp_water_solvation_free_energy"
     if not os.path.exists(output_temp_dir):
@@ -173,13 +173,15 @@ def reevaluate_energy(simulation, dcd_files, pdb_file):
     num_lambdas = len(lambda_list)
     energy_matrix = {}
 
-
+    alchemical_system, topology, positions = alchemical_builder(lambda_list[0])
+    alchem_simulation = build_alchemical_simulation(alchemical_system, topology, box, positions, nsteps, time_step, temperature, report_frequency, None, None)    
     for ii, lambda_value in enumerate(lambda_list):
         energy_matrix[ii] = {"00": None, "01": None, "10": None, "11": None}
 
         print("Re-evaluating energy for lambda: ", lambda_value)
-        alchemical_system, topology, positions = alchemical_builder(lambda_value)
-        alchem_simulation = build_alchemical_simulation(alchemical_system, topology, box, positions, nsteps, time_step, temperature, report_frequency, None, None)
+        alchem_simulation.context.setParameter("dp_alchem_lambda_0", lambda_value)
+        alchem_simulation.context.setParameter("dp_alchem_lambda_1", 1 - lambda_value)
+        alchem_simulation.context.setParameter("dp_alchem_lambda_2", 1 - lambda_value)
 
         if ii == 0:
             input_dcd_files = [dcd_files[ii], dcd_files[ii+1]]

python/setup.py

@@ -32,7 +32,7 @@
 
 
 setup(name='OpenMMDeepmdPlugin',
-      version="0.2.0",
+      version="0.3.0",
       ext_modules=[extension],
       packages=['OpenMMDeepmdPlugin', "OpenMMDeepmdPlugin.tests"],
       package_data={"OpenMMDeepmdPlugin":['data/*.pb', 'data/*.pdb']},