From f7eb395e94e848664bd434bf7e026211f8cf02d3 Mon Sep 17 00:00:00 2001
From: Jia-Xin Zhu <jiaxinzhu@stu.xmu.edu.cn>
Date: Tue, 17 Oct 2023 09:33:52 +0800
Subject: [PATCH 1/6] add `aparam_from_compute` to `pair deepmd`

---
 doc/third-party/lammps-command.md |  8 +++-
 source/lmp/pair_deepmd.cpp        | 73 ++++++++++++++++++++++++-------
 source/lmp/pair_deepmd.h          |  7 ++-
 3 files changed, 69 insertions(+), 19 deletions(-)

diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md
index a9d849bc7c..cdfa4b87d6 100644
--- a/doc/third-party/lammps-command.md
+++ b/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.
 
diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index 6dba95d533..06b1b0b87a 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -213,9 +213,9 @@ 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);
@@ -233,6 +233,30 @@ void PairDeepMD::make_fparam_from_compute(vector<double> &fparam) {
   }
 }
 
+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);
+  
+  compute->compute_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 +403,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 +517,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 +532,7 @@ void PairDeepMD::compute(int eflag, int vflag) {
 #endif
   }
 
-  if (do_compute) {
+  if (do_compute_fparam) {
     make_fparam_from_compute(fparam);
   }
 
@@ -884,6 +911,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 +1047,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")) {
@@ -1059,7 +1096,7 @@ void PairDeepMD::settings(int narg, char **arg) {
   if (do_ttm && aparam.size() > 0) {
     error->all(FLERR, "aparam 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 +1141,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):    ";
diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h
index bde7745d36..a75a53cded 100644
--- a/source/lmp/pair_deepmd.h
+++ b/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> &fparam);
+  bool do_compute_aparam;
+  std::string compute_aparam_id;
 
   void make_ttm_fparam(std::vector<double> &fparam);
 

From 1e1e60a4d701ff9b7d490a3250c23d8dc1eabf0b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 17 Oct 2023 01:35:26 +0000
Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/lmp/pair_deepmd.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index 06b1b0b87a..d48abce5a9 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -242,7 +242,7 @@ void PairDeepMD::make_aparam_from_compute(vector<double> &aparam) {
   assert(compute);
   int nlocal = atom->nlocal;
   aparam.resize(dim_aparam * nlocal);
-  
+
   compute->compute_peratom();
   if (dim_aparam == 1) {
     double *cvector = compute->vector_atom;

From bb5511f17d8b43c77de1135495d252cf46ff67ba Mon Sep 17 00:00:00 2001
From: Jia-Xin Zhu <jiaxinzhu@stu.xmu.edu.cn>
Date: Tue, 17 Oct 2023 10:09:12 +0800
Subject: [PATCH 3/6] add invoked_flag into make_fparam_from_compute and
 make_aparam_from_compute

---
 source/lmp/pair_deepmd.cpp | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index d48abce5a9..e9f226baa3 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -222,10 +222,16 @@ void PairDeepMD::make_fparam_from_compute(vector<double> &fparam) {
   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];
@@ -242,8 +248,11 @@ void PairDeepMD::make_aparam_from_compute(vector<double> &aparam) {
   assert(compute);
   int nlocal = atom->nlocal;
   aparam.resize(dim_aparam * nlocal);
-
-  compute->compute_peratom();
+  
+  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);

From 43e74bc778b54d2e406af92ebbd592579b7e8947 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 17 Oct 2023 02:12:20 +0000
Subject: [PATCH 4/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/lmp/pair_deepmd.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index e9f226baa3..3f88bfa667 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -248,7 +248,7 @@ void PairDeepMD::make_aparam_from_compute(vector<double> &aparam) {
   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;

From 8654d546c2d8e25a8244703aaeb1ba2b622f2e37 Mon Sep 17 00:00:00 2001
From: Jia-Xin Zhu <53895049+ChiahsinChu@users.noreply.github.com>
Date: Tue, 17 Oct 2023 15:09:20 +0800
Subject: [PATCH 5/6] Apply suggestions from code review

Co-authored-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Signed-off-by: Jia-Xin Zhu <53895049+ChiahsinChu@users.noreply.github.com>
---
 source/lmp/pair_deepmd.cpp | 4 ++--
 source/lmp/pair_deepmd.h   | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index 3f88bfa667..cf6468c3d3 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -1102,8 +1102,8 @@ 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 && fparam.size() > 0) {
     error->all(
diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h
index a75a53cded..0f704ab45c 100644
--- a/source/lmp/pair_deepmd.h
+++ b/source/lmp/pair_deepmd.h
@@ -123,7 +123,7 @@ class PairDeepMD : public Pair {
   void make_fparam_from_compute(std::vector<double> &fparam);
   bool do_compute_fparam;
   std::string compute_fparam_id;
-  void make_aparam_from_compute(std::vector<double> &fparam);
+  void make_aparam_from_compute(std::vector<double> &aparam);
   bool do_compute_aparam;
   std::string compute_aparam_id;
 

From cd31a0029d1377ae7e3c5ec08359610abbb82fa6 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 17 Oct 2023 07:09:46 +0000
Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/lmp/pair_deepmd.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index cf6468c3d3..432077de5b 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -1103,7 +1103,9 @@ void PairDeepMD::settings(int narg, char **arg) {
     error->all(FLERR, "Illegal out_freq, should be >= 0");
   }
   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");
+    error->all(FLERR,
+               "aparam, aparam_from_compute, and ttm should NOT be set "
+               "simultaneously");
   }
   if (do_compute_fparam && fparam.size() > 0) {
     error->all(