lmp: fix_dplr use the same type_map from pair_deepmd (deepmodel…

…ing#2776)

When `pair_coeff` has set LAMMPS `type_map` (e.g. `pair_coeff * * O H`),
`fix_dplr` uses the same LAMMPS `type_map` (i.e. `O H` in this case) to
map from LAMMPS types to the DP model types, considering LAMMPS types
should always be the same among different commands.

---------

Signed-off-by: Jinzhe Zeng <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

doc/model/dplr.md

@@ -171,8 +171,8 @@ fix_modify	0 virial yes
 ```
 
 The fix command `dplr` calculates the position of WCs by the DW model and back-propagates the long-range interaction on virtual atoms to real toms.
-At this time, the training parameter {ref}`type_map <model/type_map>` will be mapped to LAMMPS atom types.
-
+The atom names specified in [pair_style `deepmd`](../third-party/lammps-command.md#pair_style-deepmd) will be used to determine elements.
+If it is not set, the training parameter {ref}`type_map <model/type_map>` will be mapped to LAMMPS atom types.
 
 To use a time-dependent electric field, LAMMPS's `variable` feature can be utilized:
 ```lammps

source/api_c/include/c_api.h

@@ -1033,6 +1033,13 @@ int* DP_DeepTensorGetSelTypes(DP_DeepTensor* dt);
  */
 int DP_DeepTensorGetNumbSelTypes(DP_DeepTensor* dt);
 
+/**
+ * @brief Get the type map of a Deep Tensor.
+ * @param[in] dt The Deep Tensor to use.
+ * @return The type map of the Deep Tensor.
+ */
+const char* DP_DeepTensorGetTypeMap(DP_DeepTensor* dt);
+
 /**
  * @brief Check if there is any exceptions throw.
  *

source/api_c/include/deepmd.hpp

@@ -1837,6 +1837,15 @@ class DeepTensor {
   void print_summary(const std::string &pre) const {
     DP_PrintSummary(pre.c_str());
   }
+  /**
+   * @brief Get the type map (element name of the atom types) of this model.
+   * @param[out] type_map The type map of this model.
+   **/
+  void get_type_map(std::string &type_map) {
+    const char *type_map_c = DP_DeepTensorGetTypeMap(dt);
+    type_map.assign(type_map_c);
+    delete[] type_map_c;
+  };
 
  private:
   DP_DeepTensor *dt;

source/api_c/src/c_api.cc

@@ -1267,6 +1267,12 @@ int DP_DeepTensorGetNumbSelTypes(DP_DeepTensor* dt) {
   return dt->dt.sel_types().size();
 }
 
+const char* DP_DeepTensorGetTypeMap(DP_DeepTensor* dt) {
+  std::string type_map;
+  dt->dt.get_type_map(type_map);
+  return string_to_char(type_map);
+}
+
 const char* DP_DeepTensorCheckOK(DP_DeepTensor* dt) {
   return string_to_char(dt->exception);
 }

source/api_cc/include/DeepTensor.h

@@ -39,7 +39,6 @@ class DeepTensor {
    **/
   void print_summary(const std::string& pre) const;
 
- public:
   /**
    * @brief Evaluate the value by using this model.
    * @param[out] value The value to evalute, usually would be the atomic tensor.
@@ -198,6 +197,11 @@ class DeepTensor {
     assert(inited);
     return sel_type;
   };
+  /**
+   * @brief Get the type map (element name of the atom types) of this model.
+   * @param[out] type_map The type map of this model.
+   **/
+  void get_type_map(std::string& type_map);
 
  private:
   tensorflow::Session* session;

source/api_cc/src/DeepTensor.cc

@@ -792,3 +792,7 @@ template void DeepTensor::compute_inner<float>(
     const std::vector<float> &dbox,
     const int nghost,
     const InputNlist &nlist_);
+
+void DeepTensor::get_type_map(std::string &type_map) {
+  type_map = get_scalar<STRINGTYPE>("model_attr/tmap");
+}

source/lmp/fix_dplr.cpp

@@ -4,6 +4,7 @@
 #include <iomanip>
 #include <iostream>
 #include <limits>
+#include <sstream>
 
 #include "atom.h"
 #include "comm.h"
@@ -63,16 +64,15 @@ FixDPLR::FixDPLR(LAMMPS *lmp, int narg, char **arg)
   if (strcmp(update->unit_style, "metal") != 0) {
     error->all(
         FLERR,
-        "Pair deepmd requires metal unit, please set it by \"units metal\"");
+        "Fix dplr requires metal unit, please set it by \"units metal\"");
   }
 
   int iarg = 3;
   vector<int> map_vec;
   bond_type.clear();
   while (iarg < narg) {
     if (!is_key(arg[iarg])) {
-      error->all(FLERR,
-                 "Illegal pair_style command\nwrong number of parameters\n");
+      error->all(FLERR, "Illegal fix command\nwrong number of parameters\n");
     }
     if (string(arg[iarg]) == string("model")) {
       if (iarg + 1 > narg) {
@@ -128,10 +128,6 @@ FixDPLR::FixDPLR(LAMMPS *lmp, int narg, char **arg)
   }
   assert(map_vec.size() % 2 == 0),
       "number of ints provided by type_associate should be even";
-  for (int ii = 0; ii < map_vec.size() / 2; ++ii) {
-    type_asso[map_vec[ii * 2 + 0]] = map_vec[ii * 2 + 1];
-    bk_type_asso[map_vec[ii * 2 + 1]] = map_vec[ii * 2 + 0];
-  }
 
   // dpt.init(model);
   // dtm.init("frozen_model.pb");
@@ -142,18 +138,70 @@ FixDPLR::FixDPLR(LAMMPS *lmp, int narg, char **arg)
     error->one(FLERR, e.what());
   }
 
+  pair_deepmd = (PairDeepMD *)force->pair_match("deepmd", 1);
+  if (!pair_deepmd) {
+    error->all(FLERR, "pair_style deepmd should be set before this fix\n");
+  }
+
+  int n = atom->ntypes;
+  std::vector<std::string> type_names = pair_deepmd->type_names;
+  std::vector<std::string> type_map;
+  std::string type_map_str;
+  dpt.get_type_map(type_map_str);
+  // convert the string to a vector of strings
+  std::istringstream iss(type_map_str);
+  std::string type_name;
+  while (iss >> type_name) {
+    type_map.push_back(type_name);
+  }
+  if (type_names.size() == 0 || type_map.size() == 0) {
+    type_idx_map.resize(n);
+    for (int ii = 0; ii < n; ++ii) {
+      type_idx_map[ii] = ii;
+    }
+  } else {
+    type_idx_map.clear();
+    for (std::string type_name : type_names) {
+      bool found_element = false;
+      for (int ii = 0; ii < type_map.size(); ++ii) {
+        if (type_map[ii] == type_name) {
+          type_idx_map.push_back(ii);
+          found_element = true;
+          break;
+        }
+      }
+      if (!found_element && "NULL" == type_name) {
+        type_idx_map.push_back(type_map.size());  // ghost type
+        found_element = true;
+      }
+      if (!found_element) {
+        error->all(FLERR, "pair_coeff: element " + type_name +
+                              " not found in the DPLR model");
+      }
+    }
+    int numb_types = type_idx_map.size();
+    if (numb_types < n) {
+      type_idx_map.resize(n);
+      for (int ii = numb_types; ii < n; ++ii) {
+        type_idx_map[ii] = ii;
+      }
+    }
+  }
+
+  for (int ii = 0; ii < map_vec.size() / 2; ++ii) {
+    type_asso[type_idx_map[map_vec[ii * 2 + 0]]] =
+        type_idx_map[map_vec[ii * 2 + 1]];
+    bk_type_asso[type_idx_map[map_vec[ii * 2 + 1]]] =
+        type_idx_map[map_vec[ii * 2 + 0]];
+  }
+
   sel_type = dpt.sel_types();
   sort(sel_type.begin(), sel_type.end());
   dpl_type.clear();
   for (int ii = 0; ii < sel_type.size(); ++ii) {
     dpl_type.push_back(type_asso[sel_type[ii]]);
   }
 
-  pair_deepmd = (PairDeepMD *)force->pair_match("deepmd", 1);
-  if (!pair_deepmd) {
-    error->all(FLERR, "pair_style deepmd should be set before this fix\n");
-  }
-
   // set comm size needed by this fix
   comm_reverse = 3;
 }
@@ -284,7 +332,7 @@ void FixDPLR::get_valid_pairs(vector<pair<int, int> > &pairs) {
   // get type
   int *type = atom->type;
   for (int ii = 0; ii < nall; ++ii) {
-    dtype[ii] = type[ii] - 1;
+    dtype[ii] = type_idx_map[type[ii] - 1];
   }
 
   int **bondlist = neighbor->bondlist;
@@ -394,7 +442,7 @@ void FixDPLR::pre_force(int vflag) {
   vector<FLOAT_PREC> dcoord(nall * 3, 0.);
   // get type
   for (int ii = 0; ii < nall; ++ii) {
-    dtype[ii] = type[ii] - 1;
+    dtype[ii] = type_idx_map[type[ii] - 1];
   }
   // get box
   dbox[0] = domain->h[0];  // xx
@@ -518,7 +566,7 @@ void FixDPLR::post_force(int vflag) {
   {
     int *type = atom->type;
     for (int ii = 0; ii < nall; ++ii) {
-      dtype[ii] = type[ii] - 1;
+      dtype[ii] = type_idx_map[type[ii] - 1];
     }
     dbox[0] = domain->h[0];  // xx
     dbox[4] = domain->h[1];  // yy

source/lmp/fix_dplr.h

@@ -78,6 +78,7 @@ class FixDPLR : public Fix {
   double qe2f;
   void update_efield_variables();
   enum { NONE, CONSTANT, EQUAL };
+  std::vector<int> type_idx_map;
 };
 }  // namespace LAMMPS_NS
 

source/lmp/pair_deepmd.cpp

@@ -1163,8 +1163,10 @@ void PairDeepMD::coeff(int narg, char **arg) {
     }
 
     type_idx_map.clear();
+    type_names.clear();
     while (iarg < narg) {
       std::string type_name = arg[iarg];
+      type_names.push_back(type_name);
       bool found_element = false;
       for (int ii = 0; ii < type_map.size(); ++ii) {
         if (type_map[ii] == type_name) {

source/lmp/pair_deepmd.h

@@ -75,6 +75,7 @@ class PairDeepMD : public Pair {
   std::string get_file_content(const std::string &model);
   std::vector<std::string> get_file_content(
       const std::vector<std::string> &models);
+  std::vector<std::string> type_names;
 
  protected:
   virtual void allocate();

source/lmp/tests/test_dplr.py

@@ -20,6 +20,7 @@
 dipole_pbtxt_file = Path(__file__).parent / "lrdipole.pbtxt"
 dipole_pb_file = Path(__file__).parent / "lrdipole.pb"
 data_file = Path(__file__).parent / "data.lmp"
+data_type_map_file = Path(__file__).parent / "data_type_map.lmp"
 
 # this is as the same as python and c++ tests, test_deeppot_a.py
 expected_e_sr = -40.56538550
@@ -252,6 +253,7 @@
 )
 mol_list = np.array([1, 2, 1, 1, 2, 2, 1, 2])
 type_OH = np.array([1, 1, 2, 2, 2, 2, 3, 3])
+type_HO = np.array([2, 2, 1, 1, 1, 1, 3, 3])
 charge = np.array([6, 6, 1, 1, 1, 1, -8, -8])
 bond_list = (((1, 7), (2, 8)),)
 mass_list = np.array([15.99940, 1.00794, 15.99940])
@@ -275,19 +277,22 @@ def setup_module():
     write_lmp_data_full(
         box, coord, mol_list, type_OH, charge, data_file, bond_list, mass_list
     )
+    write_lmp_data_full(
+        box, coord, mol_list, type_HO, charge, data_type_map_file, bond_list, mass_list
+    )
 
 
 def teardown_module():
     os.remove(data_file)
 
 
-def _lammps(data_file) -> PyLammps:
+def _lammps(data_file, exclude_type="1 3") -> PyLammps:
     lammps = PyLammps()
     lammps.units("metal")
     lammps.boundary("p p p")
     lammps.atom_style("full")
     lammps.neighbor("0.2 bin")
-    lammps.neigh_modify("every 1 delay 0 check no exclude type 1 3")
+    lammps.neigh_modify("every 1 delay 0 check no exclude type " + exclude_type)
     lammps.read_data(data_file.resolve())
     lammps.timestep(0.0005)
     lammps.fix("1 all nve")
@@ -301,6 +306,13 @@ def lammps():
     lmp.close()
 
 
+@pytest.fixture
+def lammps_type_map():
+    lmp = _lammps(data_file=data_type_map_file, exclude_type="2 3")
+    yield lmp
+    lmp.close()
+
+
 def test_pair_deepmd_sr(lammps):
     lammps.pair_style(f"deepmd {pb_file.resolve()}")
     lammps.pair_coeff("* *")
@@ -460,3 +472,33 @@ def test_min_dplr(lammps):
         assert lammps.atoms[ii].force == pytest.approx(
             expected_f_min_step1[lammps.atoms[ii].id - 1]
         )
+
+
+def test_pair_deepmd_lr_type_map(lammps_type_map):
+    lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}")
+    lammps_type_map.pair_coeff("* * H O")
+    lammps_type_map.bond_style("zero")
+    lammps_type_map.bond_coeff("*")
+    lammps_type_map.special_bonds("lj/coul 1 1 1 angle no")
+    lammps_type_map.kspace_style("pppm/dplr 1e-5")
+    lammps_type_map.kspace_modify(
+        f"gewald {beta:.2f} diff ik mesh {mesh:d} {mesh:d} {mesh:d}"
+    )
+    lammps_type_map.fix(
+        f"0 all dplr model {pb_file.resolve()} type_associate 2 3 bond_type 1"
+    )
+    lammps_type_map.fix_modify("0 virial yes")
+    lammps_type_map.run(0)
+    for ii in range(8):
+        if lammps_type_map.atoms[ii].id > 6:
+            assert lammps_type_map.atoms[ii].position == pytest.approx(
+                expected_WC[lammps_type_map.atoms[ii].id - 7]
+            )
+    assert lammps_type_map.eval("elong") == pytest.approx(expected_e_kspace)
+    assert lammps_type_map.eval("pe") == pytest.approx(expected_e_lr)
+    for ii in range(8):
+        if lammps_type_map.atoms[ii].id <= 6:
+            assert lammps_type_map.atoms[ii].force == pytest.approx(
+                expected_f_lr[lammps_type_map.atoms[ii].id - 1]
+            )
+    lammps_type_map.run(1)