refactor DeepEval

deepmd/__init__.py

@@ -14,6 +14,16 @@
         __version__,
     )
 
+
+def DeepPotential(*args, **kwargs):
+    from deepmd.infer.deep_eval import (
+        DeepEval,
+    )
+
+    return DeepEval(*args, **kwargs)
+
+
 __all__ = [
     "__version__",
+    "DeepPotential",
 ]

deepmd/infer/deep_dipole.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.infer.deep_tensor import (
+    DeepTensor,
+)
+
+
+class DeepDipole(DeepTensor):
+    @property
+    def output_tensor_name(self) -> str:
+        return "dipole"

deepmd/infer/deep_dos.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+    Dict,
+    Optional,
+    Tuple,
+)
+
+import numpy as np
+
+from deepmd.model_format.output_def import (
+    FittingOutputDef,
+    ModelOutputDef,
+    OutputVariableDef,
+)
+
+from .deep_eval import (
+    DeepEval,
+)
+
+
+class DeepDOS(DeepEval):
+    """Potential energy model.
+
+    Parameters
+    ----------
+    model_file : Path
+        The name of the frozen model file.
+    auto_batch_size : bool or int or AutoBatchSize, default: True
+        If True, automatic batch size will be used. If int, it will be used
+        as the initial batch size.
+    neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional
+        The ASE neighbor list class to produce the neighbor list. If None, the
+        neighbor list will be built natively in the model.
+    """
+
+    @property
+    def output_def(self) -> ModelOutputDef:
+        """Get the output definition of this model."""
+        return ModelOutputDef(
+            FittingOutputDef(
+                [
+                    OutputVariableDef(
+                        "dos",
+                        shape=[-1],
+                        reduciable=True,
+                        atomic=True,
+                    ),
+                ]
+            )
+        )
+
+    def eval(
+        self,
+        coords: np.ndarray,
+        cells: Optional[np.ndarray],
+        atom_types: np.ndarray,
+        atomic: bool = False,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
+        mixed_type: bool = False,
+        **kwargs: Dict[str, Any],
+    ) -> Tuple[np.ndarray, ...]:
+        """Evaluate energy, force, and virial. If atomic is True,
+        also return atomic energy and atomic virial.
+
+        Parameters
+        ----------
+        coords : np.ndarray
+            The coordinates of the atoms, in shape (nframes, natoms, 3).
+        cells : np.ndarray
+            The cell vectors of the system, in shape (nframes, 9). If the system
+            is not periodic, set it to None.
+        atom_types : List[int]
+            The types of the atoms. If mixed_type is False, the shape is (natoms,);
+            otherwise, the shape is (nframes, natoms).
+        atomic : bool, optional
+            Whether to return atomic energy and atomic virial, by default False.
+        fparam : np.ndarray, optional
+            The frame parameters, by default None.
+        aparam : np.ndarray, optional
+            The atomic parameters, by default None.
+        mixed_type : bool, optional
+            Whether the atom_types is mixed type, by default False.
+        **kwargs : Dict[str, Any]
+            Keyword arguments.
+
+        Returns
+        -------
+        energy
+            The energy of the system, in shape (nframes,).
+        force
+            The force of the system, in shape (nframes, natoms, 3).
+        virial
+            The virial of the system, in shape (nframes, 9).
+        atomic_energy
+            The atomic energy of the system, in shape (nframes, natoms). Only returned
+            when atomic is True.
+        atomic_virial
+            The atomic virial of the system, in shape (nframes, natoms, 9). Only returned
+            when atomic is True.
+        """
+        (
+            coords,
+            cells,
+            atom_types,
+            fparam,
+            aparam,
+            nframes,
+            natoms,
+        ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type)
+        results = self.deep_eval.eval(
+            coords,
+            cells,
+            atom_types,
+            atomic,
+            fparam=fparam,
+            aparam=aparam,
+            **kwargs,
+        )
+        # energy = results["dos_redu"].reshape(nframes, self.get_numb_dos())
+        atomic_energy = results["dos"].reshape(nframes, natoms, self.get_numb_dos())
+        # not same as dos_redu... why?
+        energy = np.sum(atomic_energy, axis=1)
+
+        if atomic:
+            return (
+                energy,
+                atomic_energy,
+            )
+        else:
+            return (energy,)
+
+    def get_numb_dos(self) -> int:
+        return self.deep_eval.get_numb_dos()
+
+
+__all__ = ["DeepDOS"]