diff --git a/.conda/build.sh b/.conda/build.sh index 2a4d06c..7acbd6a 100644 --- a/.conda/build.sh +++ b/.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} diff --git a/.conda/meta.yaml b/.conda/meta.yaml index 54c1d62..cbc6794 100644 --- a/.conda/meta.yaml +++ b/.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: diff --git a/openmmapi/include/DeepmdForce.h b/openmmapi/include/DeepmdForce.h index 7565bf0..d0dfa1e 100644 --- a/openmmapi/include/DeepmdForce.h +++ b/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 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 globalParameters; + map type4EachParticle; map> particleGroup4EachType; map typesIndexMap; diff --git a/openmmapi/include/internal/DeepmdForceImpl.h b/openmmapi/include/internal/DeepmdForceImpl.h index 50e0b9d..5997491 100644 --- a/openmmapi/include/internal/DeepmdForceImpl.h +++ b/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 getDefaultParameters() { - return std::map(); // This force field doesn't define any parameters. - } + std::map getDefaultParameters(); std::vector getKernelNames(); vector> getBondedParticles() const; //void updateParametersInContext(OpenMM::ContextImpl& context); diff --git a/openmmapi/src/DeepmdForce.cpp b/openmmapi/src/DeepmdForce.cpp index 1ea11e9..4d11dc5 100644 --- a/openmmapi/src/DeepmdForce.cpp +++ b/openmmapi/src/DeepmdForce.cpp @@ -51,6 +51,8 @@ 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(); @@ -58,13 +60,15 @@ DeepmdForce::DeepmdForce(const string& GraphFile){ 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> 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 DeepmdForce::getGlobalParameters() const { + map 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; +} \ No newline at end of file diff --git a/openmmapi/src/DeepmdForceImpl.cpp b/openmmapi/src/DeepmdForceImpl.cpp index 5e1005e..870d285 100644 --- a/openmmapi/src/DeepmdForceImpl.cpp +++ b/openmmapi/src/DeepmdForceImpl.cpp @@ -66,6 +66,10 @@ vector> DeepmdForceImpl::getBondedParticles() const{ return owner.getBondsList(); } +std::map DeepmdForceImpl::getDefaultParameters() { + return owner.getGlobalParameters(); +} + std::vector DeepmdForceImpl::getKernelNames() { std::vector names; names.push_back(CalcDeepmdForceKernel::Name()); diff --git a/platforms/cuda/include/CudaDeepmdKernels.h b/platforms/cuda/include/CudaDeepmdKernels.h index 56efe9f..b4dea0a 100644 --- a/platforms/cuda/include/CudaDeepmdKernels.h +++ b/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 dforce; vector dvirial; diff --git a/platforms/cuda/src/CudaDeepmdKernels.cpp b/platforms/cuda/src/CudaDeepmdKernels.cpp index 534416e..0f9f7f8 100644 --- a/platforms/cuda/src/CudaDeepmdKernels.cpp +++ b/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 pos; context.getPositions(pos); + lambda = context.getParameter(lambda_name); + Vec3 box[3]; if (isFixedRegion){ diff --git a/platforms/reference/include/ReferenceDeepmdKernels.h b/platforms/reference/include/ReferenceDeepmdKernels.h index feb2a03..4594683 100644 --- a/platforms/reference/include/ReferenceDeepmdKernels.h +++ b/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 dforce; vector dvirial; diff --git a/platforms/reference/src/ReferenceDeepmdKernels.cpp b/platforms/reference/src/ReferenceDeepmdKernels.cpp index ee4da74..c680d5f 100644 --- a/platforms/reference/src/ReferenceDeepmdKernels.cpp +++ b/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& pos = extractPositions(context); vector& force = extractForces(context); + lambda = context.getParameter(lambda_name); + if (isFixedRegion){ // Set box size. if (context.getSystem().usesPeriodicBoundaryConditions()){ diff --git a/python/OpenMMDeepmdPlugin.i b/python/OpenMMDeepmdPlugin.i index ad539de..0d5427f 100644 --- a/python/OpenMMDeepmdPlugin.i +++ b/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. diff --git a/python/OpenMMDeepmdPlugin/tools.py b/python/OpenMMDeepmdPlugin/tools.py index 8ad246d..4c1ac6b 100644 --- a/python/OpenMMDeepmdPlugin/tools.py +++ b/python/OpenMMDeepmdPlugin/tools.py @@ -87,7 +87,7 @@ 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: @@ -95,12 +95,15 @@ def __init__(self, model_file, Lambda = 1.0) -> None: 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) @@ -164,7 +167,7 @@ def createSystem(self, topology, particleNameLabeler = "element"): # Add bond information into DeepmdForce for the PBC issue - # during the trajectory saving. + # during the trajectory saving. For visualization, the bond information is not necessary. for bond in topology.bonds(): self.dp_force.addBond(bond[0].index, bond[1].index) @@ -258,4 +261,14 @@ def addCenterParticlesToAdaptiveDPRegion( self.dp_force.setSelNum4EachType(self.dp_model_types, num4type) return self.dp_force - \ No newline at end of file + + 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 \ No newline at end of file diff --git a/python/examples/water_solvation_energy/water_solvation_free_energy.py b/python/examples/water_solvation_energy/water_solvation_free_energy.py index f94084f..e314128 100644 --- a/python/examples/water_solvation_energy/water_solvation_free_energy.py +++ b/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]] diff --git a/python/setup.py b/python/setup.py index 04451c1..3ff9203 100644 --- a/python/setup.py +++ b/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']},