fix(tf): fix DeepEval degradation for virtual types

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

deepmd/infer/deep_eval.py

@@ -242,6 +242,10 @@ def _check_mixed_types(self, atom_types: np.ndarray) -> bool:
         atom_types : np.ndarray
             The atom types of all frames, in shape nframes * natoms.
         """
+        if np.count_nonzero(atom_types[0] == -1) > 0:
+            # assume mixed_types if there are virtual types, even when
+            # the atom types of all frames are the same
+            return False
         return np.all(np.equal(atom_types, atom_types[0]))
 
     @property

deepmd/tf/infer/deep_eval.py

@@ -489,6 +489,11 @@ def make_natoms_vec(
         natoms_vec[1] = natoms
         for ii in range(self.ntypes):
             natoms_vec[ii + 2] = np.count_nonzero(atom_types[0] == ii)
+        if np.count_nonzero(atom_types[0] == -1) > 0:
+            # contains virtual atoms
+            # energy fitting sums over natoms_vec[2:] instead of reading from natoms_vec[0]
+            # causing errors for shape mismatch
+            natoms_vec[2] += np.count_nonzero(atom_types[0] == -1)
         return natoms_vec
 
     def eval_typeebd(self) -> np.ndarray: