diff --git a/deepmd/infer/backend.py b/deepmd/infer/backend.py index 809e19466b..26eef22eb4 100644 --- a/deepmd/infer/backend.py +++ b/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"): diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 92b449eb36..266b3b2d4d 100644 --- a/deepmd/infer/deep_eval.py +++ b/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() diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index 3f207d03f8..d77fd1172a 100644 --- a/deepmd/infer/deep_pot.py +++ b/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, diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 31f61f68ca..8eb7e06868 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/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.""" diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 93417a10b3..8cb9229905 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/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] ) diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 8732fec8f4..4151b1b0e4 100644 --- a/deepmd/tf/model/frozen.py +++ b/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,