diff --git a/dpdata/system.py b/dpdata/system.py
index 3c9ccec8c..945996f2a 100644
--- a/dpdata/system.py
+++ b/dpdata/system.py
@@ -331,7 +331,7 @@ def dump(self,filename,indent=4):
         dumpfn(self.as_dict(),filename,indent=indent)
 
 
-    def map_atom_types(self,type_map=None):
+    def map_atom_types(self, type_map=None) -> np.ndarray:
         """
         Map the atom types of the system
 
@@ -346,7 +346,7 @@ def map_atom_types(self,type_map=None):
 
         Returns
         -------
-        new_atom_types : list
+        new_atom_types : np.ndarray
             The mapped atom types
         """
         if isinstance(type_map,dict) or type_map is None:
@@ -366,7 +366,7 @@ def map_atom_types(self,type_map=None):
         atom_types_list=[]
         for name, numb  in  zip(self.get_atom_names(), self.get_atom_numbs()):
             atom_types_list.extend([name]*numb)
-        new_atom_types=np.array([type_map[ii] for ii in atom_types_list],dtype=np.int)
+        new_atom_types = np.array([type_map[ii] for ii in atom_types_list], dtype=int)
 
         return new_atom_types