Fix model-devi with mixed_type format (#2433)

Fix model-devi with mixed_type format;
Add UTs;
Add detection of whether models support mixed_type inference.

---------

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

Author: 3 people

deepmd/infer/deep_pot.py

@@ -164,6 +164,16 @@ def __init__(
         except (ValueError, KeyError):
             self.modifier_type = None
 
+        try:
+            t_jdata = self._get_tensor("train_attr/training_script:0")
+            jdata = run_sess(self.sess, t_jdata).decode("UTF-8")
+            import json
+
+            jdata = json.loads(jdata)
+            self.descriptor_type = jdata["model"]["descriptor"]["type"]
+        except (ValueError, KeyError):
+            self.descriptor_type = None
+
         if self.modifier_type == "dipole_charge":
             t_mdl_name = self._get_tensor("modifier_attr/mdl_name:0")
             t_mdl_charge_map = self._get_tensor("modifier_attr/mdl_charge_map:0")
@@ -243,6 +253,10 @@ def get_sel_type(self) -> List[int]:
         """Unsupported in this model."""
         raise NotImplementedError("This model type does not support this attribute")
 
+    def get_descriptor_type(self) -> List[int]:
+        """Get the descriptor type of this model."""
+        return self.descriptor_type
+
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this DP."""
         return self.dfparam

deepmd/infer/model_devi.py

@@ -143,6 +143,7 @@ def calc_model_devi(
     models,
     fname=None,
     frequency=1,
+    mixed_type=False,
 ):
     """Python interface to calculate model deviation.
 
@@ -160,6 +161,8 @@ def calc_model_devi(
         File to dump results, default None
     frequency : int
         Steps between frames (if the system is given by molecular dynamics engine), default 1
+    mixed_type : bool
+        Whether the input atype is in mixed_type format or not
 
     Returns
     -------
@@ -187,6 +190,7 @@ def calc_model_devi(
             coord,
             box,
             atype,
+            mixed_type=mixed_type,
         )
         energies.append(ret[0] / natom)
         forces.append(ret[1])
@@ -246,17 +250,21 @@ def make_model_devi(
     for system in all_sys:
         # create data-system
         dp_data = DeepmdData(system, set_prefix, shuffle_test=False, type_map=tmap)
+        mixed_type = dp_data.mixed_type
 
         data_sets = [dp_data._load_set(set_name) for set_name in dp_data.dirs]
         nframes_tot = 0
         devis = []
         for data in data_sets:
             coord = data["coord"]
             box = data["box"]
-            atype = data["type"][0]
+            if mixed_type:
+                atype = data["type"]
+            else:
+                atype = data["type"][0]
             if not dp_data.pbc:
                 box = None
-            devi = calc_model_devi(coord, box, atype, dp_models)
+            devi = calc_model_devi(coord, box, atype, dp_models, mixed_type=mixed_type)
             nframes_tot += coord.shape[0]
             devis.append(devi)
         devis = np.vstack(devis)