add aparam_from_compute to pair deepmd (#2929)

Signed-off-by: Jia-Xin Zhu <53895049+ChiahsinChu@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

doc/third-party/lammps-command.md

@@ -38,7 +38,7 @@ pair_style deepmd models ... keyword value ...
 - models = frozen model(s) to compute the interaction.
 If multiple models are provided, then only the first model serves to provide energy and force prediction for each timestep of molecular dynamics,
 and the model deviation will be computed among all models every `out_freq` timesteps.
-- keyword = *out_file* or *out_freq* or *fparam* or *fparam_from_compute* or *atomic* or *relative* or *relative_v* or *aparam* or *ttm*
+- keyword = *out_file* or *out_freq* or *fparam* or *fparam_from_compute* or *aparam_from_compute* or *atomic* or *relative* or *relative_v* or *aparam* or *ttm*
 <pre>
     <i>out_file</i> value = filename
         filename = The file name for the model deviation output. Default is model_devi.out
@@ -48,6 +48,8 @@ and the model deviation will be computed among all models every `out_freq` times
         parameters = one or more frame parameters required for model evaluation.
     <i>fparam_from_compute</i> value = id
         id = compute id used to update the frame parameter.
+    <i>aparam_from_compute</i> value = id
+        id = compute id used to update the atom parameter.
     <i>atomic</i> = no value is required.
         If this keyword is set, the force model deviation of each atom will be output.
     <i>relative</i> value = level
@@ -69,6 +71,9 @@ pair_coeff * * O H
 
 pair_style deepmd cp.pb fparam_from_compute TEMP
 compute    TEMP all temp
+
+pair_style deepmd ener.pb aparam_from_compute 1
+compute    1 all ke/atom
 ```
 
 ### Description
@@ -89,6 +94,7 @@ $$E_{v_i}=\frac{\left|D_{v_i}\right|}{\left|v_i\right|+l}$$
 
 If the keyword `fparam` is set, the given frame parameter(s) will be fed to the model.
 If the keyword `fparam_from_compute` is set, the global parameter(s) from compute command (e.g., temperature from [compute temp command](https://docs.lammps.org/compute_temp.html)) will be fed to the model as the frame parameter(s).
+If the keyword `aparam_from_compute` is set, the atomic parameter(s) from compute command (e.g., per-atom translational kinetic energy from [compute ke/atom command](https://docs.lammps.org/compute_ke_atom.html)) will be fed to the model as the atom parameter(s).
 If the keyword `aparam` is set, the given atomic parameter(s) will be fed to the model, where each atom is assumed to have the same atomic parameter(s).
 If the keyword `ttm` is set, electronic temperatures from [fix ttm command](https://docs.lammps.org/fix_ttm.html) will be fed to the model as the atomic parameters.
 

source/lmp/pair_deepmd.cpp

@@ -213,26 +213,59 @@ static void make_uniform_aparam(vector<double> &daparam,
 }
 
 void PairDeepMD::make_fparam_from_compute(vector<double> &fparam) {
-  assert(do_compute);
+  assert(do_compute_fparam);
 
-  int icompute = modify->find_compute(compute_id);
+  int icompute = modify->find_compute(compute_fparam_id);
   Compute *compute = modify->compute[icompute];
 
   assert(compute);
   fparam.resize(dim_fparam);
 
   if (dim_fparam == 1) {
-    compute->compute_scalar();
+    if (!(compute->invoked_flag & Compute::INVOKED_SCALAR)) {
+      compute->compute_scalar();
+      compute->invoked_flag |= Compute::INVOKED_SCALAR;
+    }
     fparam[0] = compute->scalar;
   } else if (dim_fparam > 1) {
-    compute->compute_vector();
+    if (!(compute->invoked_flag & Compute::INVOKED_VECTOR)) {
+      compute->compute_vector();
+      compute->invoked_flag |= Compute::INVOKED_VECTOR;
+    }
     double *cvector = compute->vector;
     for (int jj = 0; jj < dim_fparam; ++jj) {
       fparam[jj] = cvector[jj];
     }
   }
 }
 
+void PairDeepMD::make_aparam_from_compute(vector<double> &aparam) {
+  assert(do_compute_aparam);
+
+  int icompute = modify->find_compute(compute_aparam_id);
+  Compute *compute = modify->compute[icompute];
+
+  assert(compute);
+  int nlocal = atom->nlocal;
+  aparam.resize(dim_aparam * nlocal);
+
+  if (!(compute->invoked_flag & Compute::INVOKED_PERATOM)) {
+    compute->compute_peratom();
+    compute->invoked_flag |= Compute::INVOKED_PERATOM;
+  }
+  if (dim_aparam == 1) {
+    double *cvector = compute->vector_atom;
+    aparam.assign(cvector, cvector + nlocal);
+  } else if (dim_aparam > 1) {
+    double **carray = compute->array_atom;
+    for (int ii = 0; ii < nlocal; ++ii) {
+      for (int jj = 0; jj < dim_aparam; ++jj) {
+        aparam[ii * dim_aparam + jj] = carray[ii][jj];
+      }
+    }
+  }
+}
+
 #ifdef USE_TTM
 void PairDeepMD::make_ttm_fparam(vector<double> &fparam) {
   assert(do_ttm);
@@ -379,7 +412,8 @@ PairDeepMD::PairDeepMD(LAMMPS *lmp)
   eps_v = 0.;
   scale = NULL;
   do_ttm = false;
-  do_compute = false;
+  do_compute_fparam = false;
+  do_compute_aparam = false;
   single_model = false;
   multi_models_mod_devi = false;
   multi_models_no_mod_devi = false;
@@ -492,8 +526,10 @@ void PairDeepMD::compute(int eflag, int vflag) {
     }
   }
 
-  // uniform aparam
-  if (aparam.size() > 0) {
+  if (do_compute_aparam) {
+    make_aparam_from_compute(daparam);
+  } else if (aparam.size() > 0) {
+    // uniform aparam
     make_uniform_aparam(daparam, aparam, nlocal);
   } else if (do_ttm) {
 #ifdef USE_TTM
@@ -505,7 +541,7 @@ void PairDeepMD::compute(int eflag, int vflag) {
 #endif
   }
 
-  if (do_compute) {
+  if (do_compute_fparam) {
     make_fparam_from_compute(fparam);
   }
 
@@ -884,6 +920,7 @@ static bool is_key(const string &input) {
   keys.push_back("fparam");
   keys.push_back("aparam");
   keys.push_back("fparam_from_compute");
+  keys.push_back("aparam_from_compute");
   keys.push_back("ttm");
   keys.push_back("atomic");
   keys.push_back("relative");
@@ -1019,15 +1056,24 @@ void PairDeepMD::settings(int narg, char **arg) {
         if (iarg + 1 + ii >= narg || is_key(arg[iarg + 1 + ii])) {
           error->all(FLERR,
                      "invalid fparam_from_compute key: should be "
-                     "fparam_from_compute compute_id(str)");
+                     "fparam_from_compute compute_fparam_id(str)");
         }
       }
-      do_compute = true;
-      compute_id = arg[iarg + 1];
+      do_compute_fparam = true;
+      compute_fparam_id = arg[iarg + 1];
       iarg += 1 + 1;
-    }
-
-    else if (string(arg[iarg]) == string("atomic")) {
+    } else if (string(arg[iarg]) == string("aparam_from_compute")) {
+      for (int ii = 0; ii < 1; ++ii) {
+        if (iarg + 1 + ii >= narg || is_key(arg[iarg + 1 + ii])) {
+          error->all(FLERR,
+                     "invalid aparam_from_compute key: should be "
+                     "aparam_from_compute compute_aparam_id(str)");
+        }
+      }
+      do_compute_aparam = true;
+      compute_aparam_id = arg[iarg + 1];
+      iarg += 1 + 1;
+    } else if (string(arg[iarg]) == string("atomic")) {
       out_each = 1;
       iarg += 1;
     } else if (string(arg[iarg]) == string("relative")) {
@@ -1056,10 +1102,12 @@ void PairDeepMD::settings(int narg, char **arg) {
   if (out_freq < 0) {
     error->all(FLERR, "Illegal out_freq, should be >= 0");
   }
-  if (do_ttm && aparam.size() > 0) {
-    error->all(FLERR, "aparam and ttm should NOT be set simultaneously");
+  if ((int)do_ttm + (int)do_compute_aparam + (int)(aparam.size() > 0) > 1) {
+    error->all(FLERR,
+               "aparam, aparam_from_compute, and ttm should NOT be set "
+               "simultaneously");
   }
-  if (do_compute && fparam.size() > 0) {
+  if (do_compute_fparam && fparam.size() > 0) {
     error->all(
         FLERR,
         "fparam and fparam_from_compute should NOT be set simultaneously");
@@ -1104,9 +1152,13 @@ void PairDeepMD::settings(int narg, char **arg) {
       }
       cout << endl;
     }
-    if (do_compute) {
-      cout << pre << "using compute id:      ";
-      cout << compute_id << "  " << endl;
+    if (do_compute_fparam) {
+      cout << pre << "using compute id (fparam):      ";
+      cout << compute_fparam_id << "  " << endl;
+    }
+    if (do_compute_aparam) {
+      cout << pre << "using compute id (aparam):      ";
+      cout << compute_aparam_id << "  " << endl;
     }
     if (aparam.size() > 0) {
       cout << pre << "using aparam(s):    ";

source/lmp/pair_deepmd.h

@@ -121,8 +121,11 @@ class PairDeepMD : public Pair {
   double eps_v;
 
   void make_fparam_from_compute(std::vector<double> &fparam);
-  bool do_compute;
-  std::string compute_id;
+  bool do_compute_fparam;
+  std::string compute_fparam_id;
+  void make_aparam_from_compute(std::vector<double> &aparam);
+  bool do_compute_aparam;
+  std::string compute_aparam_id;
 
   void make_ttm_fparam(std::vector<double> &fparam);