[Feature Request] Support eval_descriptor for PT DeepEval

[njzjz]

Summary

Support eval_descriptor for PT DeepEval.

Detailed Description

The abstract method is here:

deepmd-kit/deepmd/infer/deep_eval.py

Lines 172 to 222 in d2b6243

	def eval_descriptor(
	self,
	coords: np.ndarray,
	cells: Optional[np.ndarray],
	atom_types: np.ndarray,
	fparam: Optional[np.ndarray] = None,
	aparam: Optional[np.ndarray] = None,
	efield: Optional[np.ndarray] = None,
	mixed_type: bool = False,
	**kwargs: Any,
	) -> np.ndarray:
	"""Evaluate descriptors by using this DP.
	Parameters
	----------
	coords
	The coordinates of atoms.
	The array should be of size nframes x natoms x 3
	cells
	The cell of the region.
	If None then non-PBC is assumed, otherwise using PBC.
	The array should be of size nframes x 9
	atom_types
	The atom types
	The list should contain natoms ints
	fparam
	The frame parameter.
	The array can be of size :
	- nframes x dim_fparam.
	- dim_fparam. Then all frames are assumed to be provided with the same fparam.
	aparam
	The atomic parameter
	The array can be of size :
	- nframes x natoms x dim_aparam.
	- natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
	- dim_aparam. Then all frames and atoms are provided with the same aparam.
	efield
	The external field on atoms.
	The array should be of size nframes x natoms x 3
	mixed_type
	Whether to perform the mixed_type mode.
	If True, the input data has the mixed_type format (see doc/model/train_se_atten.md),
	in which frames in a system may have different natoms_vec(s), with the same nloc.
	Returns
	-------
	descriptor
	Descriptors.
	"""
	raise NotImplementedError

However, how to do it for PT needs to be clarified. PT descriptor has a different forward method.

Further Information, Files, and Links

No response

[Chengqian-Zhang]

I implemented this function some time ago by copying all the forward related code before descriptor to eval_descriptor, but I think this implementation is rather redundant.

[njzjz]

    def fail(self, *args, **kwargs):
>       raise RuntimeError(name + " is not supported on ScriptModules")
E       RuntimeError: register_forward_hook is not supported on ScriptModules

../../anaconda3/lib/python3.10/site-packages/torch/jit/_script.py:981: RuntimeError