fix errors

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/infer/backend.py

@@ -21,6 +21,7 @@ def detect_backend(filename: str) -> DPBackend:
     filename : str
         The model file name
     """
+    filename = str(filename).lower()
     if filename.endswith(".pb"):
         return DPBackend.TensorFlow
     elif filename.endswith(".pth") or filename.endswith(".pt"):

deepmd/infer/deep_eval.py

@@ -239,7 +239,7 @@ def get_has_efield(self):
         return False
 
     @abstractmethod
-    def get_ntypes_spin(self):
+    def get_ntypes_spin(self) -> int:
         """Get the number of spin atom types of this model."""
 
 
@@ -458,10 +458,10 @@ def _get_sel_natoms(self, atype) -> int:
         return np.sum(np.isin(atype, self.get_sel_type()).astype(int))
 
     @property
-    def has_efield(self):
+    def has_efield(self) -> bool:
         """Check if the model has efield."""
         return self.deep_eval.get_has_efield()
 
-    def get_ntypes_spin(self):
+    def get_ntypes_spin(self) -> int:
         """Get the number of spin atom types of this model."""
         return self.deep_eval.get_ntypes_spin()

deepmd/infer/deep_pot.py

@@ -148,7 +148,17 @@ def eval(
         virial = results["energy_derv_c_redu"].reshape(nframes, 9)
 
         if atomic:
-            atomic_energy = results["energy"].reshape(nframes, natoms, 1)
+            if self.get_ntypes_spin() > 0:
+                ntypes_real = self.get_ntypes() - self.get_ntypes_spin()
+                natoms_real = sum(
+                    [
+                        np.count_nonzero(np.array(atom_types[0]) == ii)
+                        for ii in range(ntypes_real)
+                    ]
+                )
+            else:
+                natoms_real = natoms
+            atomic_energy = results["energy"].reshape(nframes, natoms_real, 1)
             atomic_virial = results["energy_derv_c"].reshape(nframes, natoms, 9)
             return (
                 energy,

deepmd/pt/infer/deep_eval.py

@@ -112,7 +112,7 @@ def get_sel_type(self) -> List[int]:
 
     def get_numb_dos(self) -> int:
         """Get the number of DOS."""
-        raise 0
+        return 0
 
     def get_has_efield(self):
         """Check if the model has efield."""

deepmd/tf/infer/deep_eval.py

@@ -208,7 +208,7 @@ def _init_tensors(self):
             "efield": "t_efield:0",
             "fparam": "t_fparam:0",
             "aparam": "t_aparam:0",
-            "ntypes_spin": "descrpt_attr/ntypes_spin:0",
+            "ntypes_spin": "spin_attr/ntypes_spin:0",
             # descriptor
             "descriptor": "o_descriptor:0",
         }
@@ -435,7 +435,7 @@ def sort_input(
         """
         natoms = atom_type.shape[1]
         if sel_atoms is not None:
-            selection = [False] * natoms
+            selection = np.array([False] * natoms, dtype=bool)
             for ii in sel_atoms:
                 selection += atom_type[0] == ii
             sel_atom_type = atom_type[:, selection]
@@ -628,7 +628,7 @@ def get_ntypes(self) -> int:
         """Get the number of atom types of this model."""
         return self.ntypes
 
-    def get_ntypes_spin(self):
+    def get_ntypes_spin(self) -> int:
         """Get the number of spin atom types of this model."""
         return self.ntypes_spin
 
@@ -770,7 +770,6 @@ def eval(
             output_dict["energy_redu"] += me.reshape(e.shape)
             output_dict["energy_deri_r"] += mf.reshape(f.shape)
             output_dict["energy_deri_c_redu"] += mv.reshape(v.shape)
-            output = tuple(output)
         return output_dict
 
     def _prepare_feed_dict(
@@ -787,7 +786,7 @@ def _prepare_feed_dict(
             coords,
             atom_types,
         )
-        atom_types = np.array(atom_types, dtype=int).reshape([-1, natoms])
+        atom_types = np.array(atom_types, dtype=int).reshape([nframes, natoms])
         coords = np.reshape(np.array(coords), [nframes, natoms * 3])
         if cells is None:
             pbc = False
@@ -942,7 +941,7 @@ def _eval_inner(
             ntypes_real = self.ntypes - self.ntypes_spin
             natoms_real = sum(
                 [
-                    np.count_nonzero(np.array(atom_types) == ii)
+                    np.count_nonzero(np.array(atom_types[0]) == ii)
                     for ii in range(ntypes_real)
                 ]
             )
@@ -977,7 +976,6 @@ def _eval_inner(
                     odef_shape = self._get_output_shape(
                         odef.name, nframes, natoms_real, odef.shape
                     )
-                    # tmp_shape = [np.prod(odef_shape[:-2]), *odef_shape[-2:]]
                     v_out[ii] = self.reverse_map(
                         np.reshape(v_out[ii], odef_shape), sel_imap[:natoms_real]
                     )

deepmd/tf/model/frozen.py

@@ -7,6 +7,9 @@
     Union,
 )
 
+from deepmd.infer.deep_pot import (
+    DeepPot,
+)
 from deepmd.tf.env import (
     GLOBAL_TF_FLOAT_PRECISION,
     MODEL_VERSION,
@@ -40,7 +43,12 @@ def __init__(self, model_file: str, **kwargs):
         super().__init__(**kwargs)
         self.model_file = model_file
         self.model = DeepPotential(model_file)
-        self.model_type = self.model.model_type
+        if isinstance(self.model, DeepPot):
+            self.model_type = "ener"
+        else:
+            raise NotImplementedError(
+                "This model type has not been implemented. " "Contribution is welcome!"
+            )
 
     def build(
         self,