Fix model-devi with mixed_type format (deepmodeling#2433)

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

---------

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

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)