set input DeepmdData.type_map to input type_map (deepmodeling#924)

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: amcadmus

deepmd/utils/data.py

@@ -57,8 +57,7 @@ def __init__ (self,
         if type_map is not None and self.type_map is not None:
             atom_type_ = [type_map.index(self.type_map[ii]) for ii in self.atom_type]
             self.atom_type = np.array(atom_type_, dtype = np.int32)
-            ntypes = len(self.type_map)
-            self.type_map = type_map[:ntypes]
+            self.type_map = type_map
         # make idx map
         self.idx_map = self._make_idx_map(self.atom_type)
         # train dirs

source/tests/test_deepmd_data.py

@@ -147,7 +147,7 @@ def test_init_type_map (self) :
         self.assertEqual(dd.idx_map[1], 1)
         self.assertEqual(dd.atom_type[0], 0)
         self.assertEqual(dd.atom_type[1], 1)
-        self.assertEqual(dd.type_map, ['bar', 'foo'])
+        self.assertEqual(dd.type_map, ['bar', 'foo', 'tar'])
 
     def test_load_set(self) :
         dd = DeepmdData(self.data_name)\