From a4224f9f24253222f4c24bc6e20a74f1589b2c01 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Fri, 20 Jan 2023 20:55:22 -0500
Subject: [PATCH 1/9] support assigning type_map in LAMMPS pair deepmd

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/lmp/pair_deepmd.cpp         | 38 +++++++++++++++++++++++++++++-
 source/lmp/pair_deepmd.h.in        |  1 +
 source/lmp/tests/data_type_map.lmp | 16 +++++++++++++
 source/lmp/tests/test_lammps.py    | 16 ++++++++++++-
 4 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 source/lmp/tests/data_type_map.lmp

diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index 0f1e7aa202..6bf903f7c9 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -328,7 +328,7 @@ void PairDeepMD::compute(int eflag, int vflag)
 
   vector<int > dtype (nall);
   for (int ii = 0; ii < nall; ++ii){
-    dtype[ii] = type[ii] - 1;
+    dtype[ii] = type_idx_map[type[ii] - 1];
   }  
 
   double dener (0);
@@ -810,6 +810,7 @@ is_key (const string& input)
   keys.push_back("atomic");
   keys.push_back("relative");
   keys.push_back("relative_v");
+  keys.push_back("type_map");
 
   for (int ii = 0; ii < keys.size(); ++ii){
     if (input == keys[ii]) {
@@ -863,6 +864,10 @@ void PairDeepMD::settings(int narg, char **arg)
     assert(dim_fparam == deep_pot.dim_fparam());
     assert(dim_aparam == deep_pot.dim_aparam());
   }
+  type_idx_map.resize(numb_types);
+  for (int ii = 0; ii < numb_types; ++ii){
+    type_idx_map[ii] = ii;
+  }
 
   out_freq = 100;
   out_file = "model_devi.out";
@@ -943,6 +948,37 @@ void PairDeepMD::settings(int narg, char **arg)
 #endif
       iarg += 2;
     }
+    else if (string(arg[iarg]) == string("type_map")) {
+      // type_map is a list of strings with undetermined length
+      // note: although we have numb_types from the model, we do not require
+      // the number of types in the system matches that in the model
+      std::vector<std::string> type_map;
+      std::string type_map_str;
+      deep_pot.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);
+      }
+
+      type_idx_map.clear();
+      while (iarg+1 < narg && !is_key(arg[iarg+1])) {
+        iarg += 1;
+        std::string type_name = arg[iarg];
+        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) {
+          error->all(FLERR, "type_map: element " + type_name + " not found in the model");
+        }
+      }
+    }
   }
   if (out_freq < 0) error->all(FLERR,"Illegal out_freq, should be >= 0");
   if (do_ttm && aparam.size() > 0) {
diff --git a/source/lmp/pair_deepmd.h.in b/source/lmp/pair_deepmd.h.in
index 45bcc03888..74fafcc0f7 100644
--- a/source/lmp/pair_deepmd.h.in
+++ b/source/lmp/pair_deepmd.h.in
@@ -109,6 +109,7 @@ private:
   int *counts,*displacements;
   tagint *tagsend, *tagrecv;
   double *stdfsend, *stdfrecv;
+  std::vector<int> type_idx_map;
 };
 
 }
diff --git a/source/lmp/tests/data_type_map.lmp b/source/lmp/tests/data_type_map.lmp
new file mode 100644
index 0000000000..c1f10c7c36
--- /dev/null
+++ b/source/lmp/tests/data_type_map.lmp
@@ -0,0 +1,16 @@
+# the first line must be comment
+6 atoms
+2 atom types
+0.0 13.0 xlo xhi
+0.0 13.0 ylo yhi
+0.0 13.0 zlo zhi
+0.0 0.0 0.0 xy xz yz
+
+Atoms
+
+1	2	12.83	2.56	2.18
+2	1	12.09	2.87	2.74
+3	1	0.25	3.32	1.68
+4	2	3.36	3.00	1.81
+5	1	3.51	2.51	2.60
+6	1	4.27	3.22	1.56 
diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py
index 014ef70d2a..ede6ce9e9e 100644
--- a/source/lmp/tests/test_lammps.py
+++ b/source/lmp/tests/test_lammps.py
@@ -13,6 +13,7 @@
 pb_file2 = Path(__file__).parent / "graph2.pb"
 system_file = Path(__file__).parent.parent.parent / "tests"
 data_file = Path(__file__).parent / "data.lmp"
+data_type_map_file = Path(__file__).parent / "data_type_map.lmp"
 md_file = Path(__file__).parent / "md.out"
 
 # this is as the same as python and c++ tests, test_deeppot_a.py
@@ -74,7 +75,7 @@
 
 
 @pytest.fixture
-def lammps() -> PyLammps:
+def lammps(data_file=data_file) -> PyLammps:
     lammps = PyLammps()
     lammps.units("metal")
     lammps.boundary("p p p")
@@ -89,6 +90,10 @@ def lammps() -> PyLammps:
     yield lammps
 
 
+def lammps_type_map():
+    lammps(data_file=data_type_map_file)
+
+
 def test_pair_deepmd(lammps):
     lammps.pair_style("deepmd {}".format(pb_file.resolve()))
     lammps.pair_coeff("* *")
@@ -205,3 +210,12 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps):
     assert md[1] == pytest.approx(np.max(expected_md_v))
     assert md[2] == pytest.approx(np.min(expected_md_v))
     assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v))))
+
+def test_pair_deepmd_type_map(lammps_type_map):
+    lammps_type_map.pair_style("deepmd {} type_map H O".format(pb_file.resolve()))
+    lammps_type_map.pair_coeff("* *")
+    lammps_type_map.run(0)
+    assert lammps_type_map.eval("pe") == pytest.approx(expected_e)
+    for ii in range(6):
+        assert lammps_type_map.atoms[ii].force == pytest.approx(expected_f[ii])
+    lammps_type_map.run(1)

From 3873c80fd81fb9f43d00afd6b18e2e266c8c98eb Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Fri, 20 Jan 2023 21:44:10 -0500
Subject: [PATCH 2/9] bugfix

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/lmp/pair_deepmd.cpp      | 1 +
 source/lmp/tests/test_lammps.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index 6bf903f7c9..428bb130b6 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -978,6 +978,7 @@ void PairDeepMD::settings(int narg, char **arg)
           error->all(FLERR, "type_map: element " + type_name + " not found in the model");
         }
       }
+      iarg += 1;
     }
   }
   if (out_freq < 0) error->all(FLERR,"Illegal out_freq, should be >= 0");
diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py
index ede6ce9e9e..80941776ee 100644
--- a/source/lmp/tests/test_lammps.py
+++ b/source/lmp/tests/test_lammps.py
@@ -90,6 +90,7 @@ def lammps(data_file=data_file) -> PyLammps:
     yield lammps
 
 
+@pytest.fixture
 def lammps_type_map():
     lammps(data_file=data_type_map_file)
 

From 84daa6fb65e797a0c96287c8d47fe501208f8341 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sat, 21 Jan 2023 00:08:28 -0500
Subject: [PATCH 3/9] fix tests

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/lmp/tests/test_lammps.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py
index 80941776ee..e91703f661 100644
--- a/source/lmp/tests/test_lammps.py
+++ b/source/lmp/tests/test_lammps.py
@@ -73,9 +73,7 @@
     ).split())
 
 
-
-@pytest.fixture
-def lammps(data_file=data_file) -> PyLammps:
+def _lammps(data_file) -> PyLammps:
     lammps = PyLammps()
     lammps.units("metal")
     lammps.boundary("p p p")
@@ -90,9 +88,13 @@ def lammps(data_file=data_file) -> PyLammps:
     yield lammps
 
 
+@pytest.fixture
+def lammps():
+    yield _lammps(data_file=data_file)
+
 @pytest.fixture
 def lammps_type_map():
-    lammps(data_file=data_type_map_file)
+    yield lammps(data_file=data_type_map_file)
 
 
 def test_pair_deepmd(lammps):

From c687c5aa280be2d4c351524a361dd091e5794fbc Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sat, 21 Jan 2023 00:20:22 -0500
Subject: [PATCH 4/9] fix tests

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/lmp/tests/test_lammps.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py
index e91703f661..14d09aaff6 100644
--- a/source/lmp/tests/test_lammps.py
+++ b/source/lmp/tests/test_lammps.py
@@ -85,7 +85,7 @@ def _lammps(data_file) -> PyLammps:
     lammps.mass("2 2")
     lammps.timestep(0.0005)
     lammps.fix("1 all nve")
-    yield lammps
+    return lammps
 
 
 @pytest.fixture

From 68570db0195566fee3b5cde2bda8ccdbb8a65db6 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sat, 21 Jan 2023 00:40:36 -0500
Subject: [PATCH 5/9] fix typo

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/lmp/tests/test_lammps.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py
index 14d09aaff6..2b75ab7415 100644
--- a/source/lmp/tests/test_lammps.py
+++ b/source/lmp/tests/test_lammps.py
@@ -94,7 +94,7 @@ def lammps():
 
 @pytest.fixture
 def lammps_type_map():
-    yield lammps(data_file=data_type_map_file)
+    yield _lammps(data_file=data_type_map_file)
 
 
 def test_pair_deepmd(lammps):

From 683003a5c36f320f6e95215de63696886834bedf Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sat, 21 Jan 2023 01:47:21 -0500
Subject: [PATCH 6/9] add docs

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 deepmd/utils/argcheck.py            |  2 +-
 doc/data/system.md                  |  4 ++--
 doc/model/dplr.md                   |  1 +
 doc/third-party/lammps-command.md   | 10 +++++++++-
 doc/third-party/lammps.md           |  4 ++--
 examples/water/lmp/in.lammps        |  2 +-
 examples/water/lmp/in.plugin.lammps |  2 +-
 7 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 1a81a2fefe..13c351731f 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -457,7 +457,7 @@ def model_compression_type_args():
 
 
 def model_args ():    
-    doc_type_map = 'A list of strings. Give the name to each type of atoms. It is noted that the number of atom type of training system must be less than 128 in a GPU environment.'
+    doc_type_map = 'A list of strings. Give the name to each type of atoms. It is noted that the number of atom type of training system must be less than 128 in a GPU environment. If not given, type.raw in each system should use the same type indexes, and type_map.raw will take no effect.'
     doc_data_stat_nbatch = 'The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics.'
     doc_data_stat_protect = 'Protect parameter for atomic energy regression.'
     doc_data_bias_nsample = 'The number of training samples in a system to compute and change the energy bias.'
diff --git a/doc/data/system.md b/doc/data/system.md
index d09761b675..af925a8d85 100644
--- a/doc/data/system.md
+++ b/doc/data/system.md
@@ -6,8 +6,8 @@ A system should contain system properties, input frame properties, and labeled f
 
 ID       | Property                | Raw file     | Required/Optional    | Shape                    | Description
 -------- | ----------------------  | ------------ | -------------------- | -----------------------  | -----------
-type     | Atom type indexes       | type.raw     | Required             | Natoms                   | Integers that start with 0
-type_map | Atom type names         | type_map.raw | Optional             | Ntypes                   | Atom names that map to atom type, which is unnecessart to be contained in the periodic table
+type     | Atom type indexes       | type.raw     | Required             | Natoms                   | Integers that start with 0. If the training parameter {ref}`type_map <model/type_map>` is set and `type_map.raw` is provided, type indexes map to `type_map.raw`; otherwise, the system type indexes map to model type indexes (whether {ref}`type_map <model/type_map>` is set or not)
+type_map | Atom type names         | type_map.raw | Optional             | Ntypes                   | Atom names that map to atom type, which is unnecessary to be contained in the periodic table. Only works when the training parameter {ref}`type_map <model/type_map>` is set
 nopbc    | Non-periodic system     | nopbc        | Optional             | 1                        | If True, this system is non-periodic; otherwise it's periodic
 
 The input frame properties contain the following property, the first axis of which is the number of frames:
diff --git a/doc/model/dplr.md b/doc/model/dplr.md
index bb6ea66084..c5d3031f86 100644
--- a/doc/model/dplr.md
+++ b/doc/model/dplr.md
@@ -143,6 +143,7 @@ fix		0 all dplr model ener.pb type_associate 1 3 bond_type 1
 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 used to map with LAMMPS atom types.
 
 ```lammps
 # compute the temperature of real atoms, excluding virtual atom contribution
diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md
index 5ec51c6830..7106c69945 100644
--- a/doc/third-party/lammps-command.md
+++ b/doc/third-party/lammps-command.md
@@ -29,8 +29,10 @@ 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 *atomic* or *relative* or *relative_v* or *aparam* or *ttm*
+- keyword = *type_map* or *out_file* or *out_freq* or *fparam* or *atomic* or *relative* or *relative_v* or *aparam* or *ttm*
 <pre>
+    <i>type_map</i> value = types
+        types = Atom names that map to LAMMPS atom type. Default follows the type map of the model
     <i>out_file</i> value = filename
         filename = The file name for the model deviation output. Default is model_devi.out
     <i>out_freq</i> value = freq
@@ -61,6 +63,10 @@ Evaluate the interaction of the system by using [Deep Potential][DP] or [Deep Po
 
 This pair style takes the deep potential defined in a model file that usually has the .pb extension. The model can be trained and frozen by package [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit).
 
+`type_map` maps atom names with LAMMPS atom types (integers from 1 to Ntypes).
+If the `pair_style` argument `type_map` is not set, the training parameter {ref}`type_map <model/type_map>` will be used by default.
+If the training parameter {ref}`type_map <model/type_map>` is not set, the `pair_style` argument `type_map` cannot be used. In this case, atom type indexes in [`type.raw`](../data/system.md) (integers from 0 to Ntypes-1) will map with LAMMPS atom types.
+
 The model deviation evalulates the consistency of the force predictions from multiple models. By default, only the maximal, minimal and average model deviations are output. If the key `atomic` is set, then the model deviation of force prediction of each atom will be output.
 
 By default, the model deviation is output in absolute value. If the keyword `relative` is set, then the relative model deviation of the force will be output, including values output by the keyword `atomic`. The relative model deviation of the force on atom $i$ is defined by
@@ -92,6 +98,8 @@ compute ID group-ID deeptensor/atom model_file
 - deeptensor/atom: the style of this compute
 - model_file: the name of the binary model file.
 
+At this time, the training parameter {ref}`type_map <model/type_map>` will be used to map with LAMMPS atom types.
+
 ### Examples
 ```lammps
 compute         dipole all deeptensor/atom dipole.pb
diff --git a/doc/third-party/lammps.md b/doc/third-party/lammps.md
index 5b9d3fd728..0154152b9d 100644
--- a/doc/third-party/lammps.md
+++ b/doc/third-party/lammps.md
@@ -3,7 +3,7 @@
 Running an MD simulation with LAMMPS is simpler. In the LAMMPS input file, one needs to specify the pair style as follows
 
 ```lammps
-pair_style     deepmd graph.pb
+pair_style     deepmd graph.pb type_map O H
 pair_coeff     * *
 ```
-where `graph.pb` is the file name of the frozen model. It should be noted that LAMMPS counts atom types starting from 1, therefore, all LAMMPS atom types will be firstly subtracted by 1, and then passed into the DeePMD-kit engine to compute the interactions. 
+where `graph.pb` is the file name of the frozen model. `type_map` maps atom names with LAMMPS atom types (integers from 1 to Ntypes).
diff --git a/examples/water/lmp/in.lammps b/examples/water/lmp/in.lammps
index 81864aad0e..e45641b0f4 100644
--- a/examples/water/lmp/in.lammps
+++ b/examples/water/lmp/in.lammps
@@ -11,7 +11,7 @@ read_data	water.lmp
 mass 		1 16
 mass		2 2
 
-pair_style	deepmd frozen_model.pb
+pair_style	deepmd frozen_model.pb type_map O H
 pair_coeff  * *	
 
 velocity        all create 330.0 23456789
diff --git a/examples/water/lmp/in.plugin.lammps b/examples/water/lmp/in.plugin.lammps
index 7161735c23..9b086b41f4 100644
--- a/examples/water/lmp/in.plugin.lammps
+++ b/examples/water/lmp/in.plugin.lammps
@@ -14,7 +14,7 @@ mass		2 2
 # load the deepmd plugin
 plugin load libdeepmd_lmp.so
 
-pair_style	deepmd frozen_model.pb
+pair_style	deepmd frozen_model.pb type_map O H
 pair_coeff  * *	
 
 velocity        all create 330.0 23456789

From 334b802f780c8085bbfca21e05791afba7cfe26e Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sat, 21 Jan 2023 15:31:24 -0500
Subject: [PATCH 7/9] move to pair_coeff, like others such as reaxff, eam

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 doc/model/dplr.md                   |  2 +-
 doc/third-party/lammps-command.md   | 15 +++---
 doc/third-party/lammps.md           |  6 +--
 examples/water/lmp/in.lammps        |  4 +-
 examples/water/lmp/in.plugin.lammps |  4 +-
 source/lmp/pair_deepmd.cpp          | 83 +++++++++++++++--------------
 source/lmp/tests/test_lammps.py     |  4 +-
 7 files changed, 61 insertions(+), 57 deletions(-)

diff --git a/doc/model/dplr.md b/doc/model/dplr.md
index c5d3031f86..9893d30713 100644
--- a/doc/model/dplr.md
+++ b/doc/model/dplr.md
@@ -143,7 +143,7 @@ fix		0 all dplr model ener.pb type_associate 1 3 bond_type 1
 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 used to map with LAMMPS atom types.
+At this time, the training parameter {ref}`type_map <model/type_map>` will be mapped to LAMMPS atom types.
 
 ```lammps
 # compute the temperature of real atoms, excluding virtual atom contribution
diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md
index 7106c69945..23c23136f8 100644
--- a/doc/third-party/lammps-command.md
+++ b/doc/third-party/lammps-command.md
@@ -29,10 +29,8 @@ 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 = *type_map* or *out_file* or *out_freq* or *fparam* or *atomic* or *relative* or *relative_v* or *aparam* or *ttm*
+- keyword = *out_file* or *out_freq* or *fparam* or *atomic* or *relative* or *relative_v* or *aparam* or *ttm*
 <pre>
-    <i>type_map</i> value = types
-        types = Atom names that map to LAMMPS atom type. Default follows the type map of the model
     <i>out_file</i> value = filename
         filename = The file name for the model deviation output. Default is model_devi.out
     <i>out_freq</i> value = freq
@@ -56,6 +54,7 @@ and the model deviation will be computed among all models every `out_freq` times
 pair_style deepmd graph.pb
 pair_style deepmd graph.pb fparam 1.2
 pair_style deepmd graph_0.pb graph_1.pb graph_2.pb out_file md.out out_freq 10 atomic relative 1.0
+pair_coeff * * O H
 ```
 
 ### Description
@@ -63,10 +62,6 @@ Evaluate the interaction of the system by using [Deep Potential][DP] or [Deep Po
 
 This pair style takes the deep potential defined in a model file that usually has the .pb extension. The model can be trained and frozen by package [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit).
 
-`type_map` maps atom names with LAMMPS atom types (integers from 1 to Ntypes).
-If the `pair_style` argument `type_map` is not set, the training parameter {ref}`type_map <model/type_map>` will be used by default.
-If the training parameter {ref}`type_map <model/type_map>` is not set, the `pair_style` argument `type_map` cannot be used. In this case, atom type indexes in [`type.raw`](../data/system.md) (integers from 0 to Ntypes-1) will map with LAMMPS atom types.
-
 The model deviation evalulates the consistency of the force predictions from multiple models. By default, only the maximal, minimal and average model deviations are output. If the key `atomic` is set, then the model deviation of force prediction of each atom will be output.
 
 By default, the model deviation is output in absolute value. If the keyword `relative` is set, then the relative model deviation of the force will be output, including values output by the keyword `atomic`. The relative model deviation of the force on atom $i$ is defined by
@@ -82,6 +77,10 @@ If the keyword `fparam` is set, the given frame parameter(s) will be fed to the
 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.
 
+Only a single `pair_coeff` command is used with the deepmd style which specifies atom names. These are mapped to LAMMPS atom types (integers from 1 to Ntypes) by specifying Ntypes additional arguments after `* *` in the `pair_coeff` command.
+If atom names are not set in the `pair_coeff` command, the training parameter {ref}`type_map <model/type_map>` will be used by default.
+If the training parameter {ref}`type_map <model/type_map>` is not set, atom names in the `pair_coeff` command cannot be set. In this case, atom type indexes in [`type.raw`](../data/system.md) (integers from 0 to Ntypes-1) will map to LAMMPS atom types.
+
 ### Restrictions
 - The `deepmd` pair style is provided in the USER-DEEPMD package, which is compiled from the DeePMD-kit, visit the [DeePMD-kit website](https://github.com/deepmodeling/deepmd-kit) for more information.
 
@@ -98,7 +97,7 @@ compute ID group-ID deeptensor/atom model_file
 - deeptensor/atom: the style of this compute
 - model_file: the name of the binary model file.
 
-At this time, the training parameter {ref}`type_map <model/type_map>` will be used to map with LAMMPS atom types.
+At this time, the training parameter {ref}`type_map <model/type_map>` will be mapped to LAMMPS atom types.
 
 ### Examples
 ```lammps
diff --git a/doc/third-party/lammps.md b/doc/third-party/lammps.md
index 0154152b9d..0020db01c5 100644
--- a/doc/third-party/lammps.md
+++ b/doc/third-party/lammps.md
@@ -3,7 +3,7 @@
 Running an MD simulation with LAMMPS is simpler. In the LAMMPS input file, one needs to specify the pair style as follows
 
 ```lammps
-pair_style     deepmd graph.pb type_map O H
-pair_coeff     * *
+pair_style     deepmd graph.pb
+pair_coeff     * * O H
 ```
-where `graph.pb` is the file name of the frozen model. `type_map` maps atom names with LAMMPS atom types (integers from 1 to Ntypes).
+where `graph.pb` is the file name of the frozen model. `pair_coeff` maps atom names (`O H`) with LAMMPS atom types (integers from 1 to Ntypes, i.e. `1 2`).
diff --git a/examples/water/lmp/in.lammps b/examples/water/lmp/in.lammps
index e45641b0f4..828beb4c68 100644
--- a/examples/water/lmp/in.lammps
+++ b/examples/water/lmp/in.lammps
@@ -11,8 +11,8 @@ read_data	water.lmp
 mass 		1 16
 mass		2 2
 
-pair_style	deepmd frozen_model.pb type_map O H
-pair_coeff  * *	
+pair_style	deepmd frozen_model.pb
+pair_coeff  * *	O H
 
 velocity        all create 330.0 23456789
 
diff --git a/examples/water/lmp/in.plugin.lammps b/examples/water/lmp/in.plugin.lammps
index 9b086b41f4..7791c7ae77 100644
--- a/examples/water/lmp/in.plugin.lammps
+++ b/examples/water/lmp/in.plugin.lammps
@@ -14,8 +14,8 @@ mass		2 2
 # load the deepmd plugin
 plugin load libdeepmd_lmp.so
 
-pair_style	deepmd frozen_model.pb type_map O H
-pair_coeff  * *	
+pair_style	deepmd frozen_model.pb
+pair_coeff  * *	O H
 
 velocity        all create 330.0 23456789
 
diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index 428bb130b6..e236c0d862 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -810,7 +810,6 @@ is_key (const string& input)
   keys.push_back("atomic");
   keys.push_back("relative");
   keys.push_back("relative_v");
-  keys.push_back("type_map");
 
   for (int ii = 0; ii < keys.size(); ++ii){
     if (input == keys[ii]) {
@@ -864,10 +863,6 @@ void PairDeepMD::settings(int narg, char **arg)
     assert(dim_fparam == deep_pot.dim_fparam());
     assert(dim_aparam == deep_pot.dim_aparam());
   }
-  type_idx_map.resize(numb_types);
-  for (int ii = 0; ii < numb_types; ++ii){
-    type_idx_map[ii] = ii;
-  }
 
   out_freq = 100;
   out_file = "model_devi.out";
@@ -948,38 +943,6 @@ void PairDeepMD::settings(int narg, char **arg)
 #endif
       iarg += 2;
     }
-    else if (string(arg[iarg]) == string("type_map")) {
-      // type_map is a list of strings with undetermined length
-      // note: although we have numb_types from the model, we do not require
-      // the number of types in the system matches that in the model
-      std::vector<std::string> type_map;
-      std::string type_map_str;
-      deep_pot.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);
-      }
-
-      type_idx_map.clear();
-      while (iarg+1 < narg && !is_key(arg[iarg+1])) {
-        iarg += 1;
-        std::string type_name = arg[iarg];
-        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) {
-          error->all(FLERR, "type_map: element " + type_name + " not found in the model");
-        }
-      }
-      iarg += 1;
-    }
   }
   if (out_freq < 0) error->all(FLERR,"Illegal out_freq, should be >= 0");
   if (do_ttm && aparam.size() > 0) {
@@ -1065,13 +1028,55 @@ void PairDeepMD::coeff(int narg, char **arg)
   jlo = 0;
   ihi = n;
   jhi = n;
-  if (narg == 2) {
+  if (narg >= 2) {
     utils::bounds(FLERR,arg[0],1,atom->ntypes,ilo,ihi,error);
     utils::bounds(FLERR,arg[1],1,atom->ntypes,jlo,jhi,error);
     if (ilo != 1 || jlo != 1 || ihi != n || jhi != n) {
       error->all(FLERR,"deepmd requires that the scale should be set to all atom types, i.e. pair_coeff * *.");
     }
-  }  
+  }
+  if (narg <= 2) {
+    type_idx_map.resize(numb_types);
+    for (int ii = 0; ii < numb_types; ++ii){
+      type_idx_map[ii] = ii;
+    }
+  } else {
+    int iarg = 2;
+
+    // type_map is a list of strings with undetermined length
+    // note: although we have numb_types from the model, we do not require
+    // the number of types in the system matches that in the model
+    std::vector<std::string> type_map;
+    std::string type_map_str;
+    deep_pot.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);
+    }
+
+    type_idx_map.clear();
+    while (iarg+1 < narg && !is_key(arg[iarg+1])) {
+      iarg += 1;
+      std::string type_name = arg[iarg];
+      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) {
+        error->all(FLERR, "pair_coeff: element " + type_name + " not found in the model");
+      }
+    }
+    numb_types = type_idx_map.size();
+  }
+  if (numb_types < n) {
+    error->all(FLERR, "number of types assigned by pair_coeff or in the model is less than the number of types in the system");
+  }
   for (int i = ilo; i <= ihi; i++) {
     for (int j = MAX(jlo,i); j <= jhi; j++) {
       setflag[i][j] = 1;
diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py
index 2b75ab7415..9815a962a5 100644
--- a/source/lmp/tests/test_lammps.py
+++ b/source/lmp/tests/test_lammps.py
@@ -215,8 +215,8 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps):
     assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v))))
 
 def test_pair_deepmd_type_map(lammps_type_map):
-    lammps_type_map.pair_style("deepmd {} type_map H O".format(pb_file.resolve()))
-    lammps_type_map.pair_coeff("* *")
+    lammps_type_map.pair_style("deepmd {}".format(pb_file.resolve()))
+    lammps_type_map.pair_coeff("* * H O")
     lammps_type_map.run(0)
     assert lammps_type_map.eval("pe") == pytest.approx(expected_e)
     for ii in range(6):

From 93879db799282046d899fd315eaec4226ff71f20 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sat, 21 Jan 2023 16:11:34 -0500
Subject: [PATCH 8/9] fix iarg

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/lmp/pair_deepmd.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index e236c0d862..e44bf03b7b 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -1057,8 +1057,7 @@ void PairDeepMD::coeff(int narg, char **arg)
     }
 
     type_idx_map.clear();
-    while (iarg+1 < narg && !is_key(arg[iarg+1])) {
-      iarg += 1;
+    while (iarg < narg) {
       std::string type_name = arg[iarg];
       bool found_element = false;
       for (int ii = 0; ii < type_map.size(); ++ii) {
@@ -1071,6 +1070,7 @@ void PairDeepMD::coeff(int narg, char **arg)
       if (!found_element) {
         error->all(FLERR, "pair_coeff: element " + type_name + " not found in the model");
       }
+      iarg += 1;
     }
     numb_types = type_idx_map.size();
   }

From 20851eb7f530201afe9f08dfba25ec4a5d6c65fd Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Tue, 24 Jan 2023 04:07:34 -0500
Subject: [PATCH 9/9] update docs

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 doc/data/system.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/data/system.md b/doc/data/system.md
index af925a8d85..6a6fb4b58f 100644
--- a/doc/data/system.md
+++ b/doc/data/system.md
@@ -6,7 +6,7 @@ A system should contain system properties, input frame properties, and labeled f
 
 ID       | Property                | Raw file     | Required/Optional    | Shape                    | Description
 -------- | ----------------------  | ------------ | -------------------- | -----------------------  | -----------
-type     | Atom type indexes       | type.raw     | Required             | Natoms                   | Integers that start with 0. If the training parameter {ref}`type_map <model/type_map>` is set and `type_map.raw` is provided, type indexes map to `type_map.raw`; otherwise, the system type indexes map to model type indexes (whether {ref}`type_map <model/type_map>` is set or not)
+type     | Atom type indexes       | type.raw     | Required             | Natoms                   | Integers that start with 0. If both the training parameter {ref}`type_map <model/type_map>` is set and `type_map.raw` is provided, the system atom type should be mapped to `type_map.raw` in `type.raw` and will be mapped to the model atom type when training; otherwise, the system atom type will be always mapped to the model atom type (whether {ref}`type_map <model/type_map>` is set or not)
 type_map | Atom type names         | type_map.raw | Optional             | Ntypes                   | Atom names that map to atom type, which is unnecessary to be contained in the periodic table. Only works when the training parameter {ref}`type_map <model/type_map>` is set
 nopbc    | Non-periodic system     | nopbc        | Optional             | 1                        | If True, this system is non-periodic; otherwise it's periodic