feat(pt): add charge density prediction support

deepmd/infer/deep_density.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
+import numpy as np
+
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+    ModelOutputDef,
+    OutputVariableDef,
+)
+
+from .deep_eval import (
+    DeepEval,
+)
+
+
+class DeepDensity(DeepEval):
+    """Charge density evaluated on grid points.
+
+    Parameters
+    ----------
+    model_file : Path
+        The name of the frozen model file.
+    *args : list
+        Positional arguments.
+    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.
+    **kwargs : dict
+        Keyword arguments.
+    """
+
+    @property
+    def output_def(self) -> ModelOutputDef:
+        """Get the output definition of this model.
+
+        The output definition is identical to that of the density
+        fitting net of the model: the density is predicted on grid
+        points, and is neither reducible nor differentiable.
+        """
+        return ModelOutputDef(
+            FittingOutputDef(
+                [
+                    OutputVariableDef(
+                        "density",
+                        [1],
+                        reducible=False,
+                        r_differentiable=False,
+                        c_differentiable=False,
+                    ),
+                ]
+            )
+        )
+
+    def eval(
+        self,
+        coords: np.ndarray,
+        cells: np.ndarray | None,
+        atom_types: list[int] | np.ndarray,
+        grid: np.ndarray,
+        fparam: np.ndarray | None = None,
+        aparam: np.ndarray | None = None,
+        mixed_type: bool = False,
+        grid_batch_size: int | None = 4096,
+        **kwargs: Any,
+    ) -> np.ndarray:
+        """Evaluate the density on grid points.
+
+        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] or np.ndarray
+            The types of the atoms. If mixed_type is False, the shape is (natoms,);
+            otherwise, the shape is (nframes, natoms).
+        grid : np.ndarray
+            The coordinates of the grid points, in shape (nframes, ngrid, 3).
+        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.
+        grid_batch_size : int or None, optional
+            Maximum number of grid points evaluated in one model call. Grid
+            chunks are concatenated in their original order. The default is
+            4096; set to None to evaluate the full grid in one call.
+        **kwargs : dict[str, Any]
+            Keyword arguments. ``charge_spin`` (nframes, 2) is forwarded to
+            models with charge spin embedding; when omitted, such models use
+            their ``default_chg_spin`` if set.
+
+        Returns
+        -------
+        density
+            The density on the grid points, in shape (nframes, ngrid).
+        """
+        (
+            coords,
+            cells,
+            atom_types,
+            fparam,
+            aparam,
+            nframes,
+            _,
+        ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type)
+        grid = self._standard_grid(grid, nframes)
+        if grid_batch_size is not None and grid_batch_size <= 0:
+            raise ValueError("grid_batch_size must be positive or None")
+        if grid_batch_size is None or grid.shape[1] <= grid_batch_size:
+            grid_slices = (grid,)
+        else:
+            grid_slices = tuple(
+                grid[:, begin : begin + grid_batch_size]
+                for begin in range(0, grid.shape[1], grid_batch_size)
+            )
+        density_chunks = []
+        for grid_chunk in grid_slices:
+            results = self.deep_eval.eval(
+                coords,
+                cells,
+                atom_types,
+                False,
+                fparam=fparam,
+                aparam=aparam,
+                grid=grid_chunk,
+                **kwargs,
+            )
+            density_chunks.append(results["density"].reshape(nframes, -1))
+        return np.concatenate(density_chunks, axis=1)
+
+    @staticmethod
+    def _standard_grid(grid: np.ndarray, nframes: int) -> np.ndarray:
+        """Normalise the grid coordinates to ``(nframes, ngrid, 3)``.
+
+        The auto batcher slices every argument with ``ndim > 1`` along axis
+        0, so a natural single-frame ``(ngrid, 3)`` input would silently be
+        truncated to one grid point; carry the frame dimension explicitly.
+        """
+        arr = np.asarray(grid)
+        if arr.ndim == 2:
+            if nframes != 1:
+                raise ValueError(
+                    f"grid of shape {arr.shape} is ambiguous for {nframes} "
+                    "frames; pass grid with shape (nframes, ngrid, 3)"
+                )
+            arr = arr[None, ...]
+        if arr.ndim != 3 or arr.shape[-1] != 3:
+            raise ValueError(
+                f"grid must have shape (nframes, ngrid, 3), got {arr.shape}"
+            )
+        if arr.shape[0] != nframes:
+            raise ValueError(f"grid has {arr.shape[0]} frames but coord has {nframes}")
+        if arr.shape[1] == 0:
+            raise ValueError("grid must contain at least one point per frame")
+        return arr
+
+
+__all__ = ["DeepDensity"]

deepmd/infer/deep_pot.py

@@ -212,6 +212,7 @@ def eval(
             aparam=aparam,
             **kwargs,
         )
+
         energy = results["energy_redu"].reshape(nframes, 1)
         force = results["energy_derv_r"].reshape(nframes, natoms, 3)
         virial = results["energy_derv_c_redu"].reshape(nframes, 9)

deepmd/infer/model_test/__init__.py

@@ -14,6 +14,9 @@
     Any,
 )
 
+from deepmd.infer.deep_density import (
+    DeepDensity,
+)
 from deepmd.infer.deep_dipole import (
     DeepDipole,
 )
@@ -36,6 +39,9 @@
     save_txt_file,
     test_chunk_atoms,
 )
+from deepmd.infer.model_test.density import (
+    DensityTester,
+)
 from deepmd.infer.model_test.dos import (
     DosTester,
 )
@@ -56,6 +62,7 @@
 
 __all__ = [
     "ChunkContext",
+    "DensityTester",
     "DipoleTester",
     "DosTester",
     "EnerTester",
@@ -96,6 +103,8 @@ def build_tester(dp: Any, *, atomic: bool) -> ModelTester:
         return tester(dp, atomic=atomic)
     if isinstance(dp, DeepDOS):
         return DosTester(dp, atomic=atomic)
+    if isinstance(dp, DeepDensity):
+        return DensityTester(dp, atomic=atomic)
     if isinstance(dp, DeepProperty):
         return PropertyTester(dp, atomic=atomic)
     if isinstance(dp, DeepGlobalPolar):

deepmd/infer/model_test/density.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Testing of models predicting charge density on grid points."""
+
+from deepmd.infer.model_test.base import (
+    ChunkContext,
+    ModelTester,
+    _write_per_frame_details,
+)
+from deepmd.utils.data import (
+    DeepmdData,
+)
+from deepmd.utils.eval_metrics import (
+    mae,
+    rmse,
+)
+
+__all__ = ["DensityTester"]
+
+
+class DensityTester(ModelTester):
+    """Test a model of charge density on grid points."""
+
+    report = (
+        ("mae_density", "DENSITY MAE             : {} units"),
+        ("rmse_density", "DENSITY RMSE            : {} units"),
+    )
+
+    def add_data_requirements(self, data: DeepmdData) -> None:
+        """Declare the labels a density test consumes."""
+        dp = self.dp
+        # The grid and the density are defined on grid points rather than on
+        # atoms, and their extent (ngrid) is not known until the data is
+        # loaded. They are declared frame-major so the loader keeps the
+        # (nframes, ngrid, ndof) layout without reshaping to natoms.
+        data.add(
+            "grid",
+            3,
+            atomic=False,
+            must=True,
+            high_prec=True,
+            special_shape="frame_major",
+        )
+        data.add(
+            "density",
+            1,
+            atomic=False,
+            must=True,
+            high_prec=True,
+            special_shape="frame_major",
+        )
+        if dp.get_dim_fparam() > 0:
+            data.add(
+                "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False
+            )
+        if dp.get_dim_aparam() > 0:
+            data.add(
+                "aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False
+            )
+        if dp.has_chg_spin_ebd():
+            data.add(
+                "charge_spin",
+                2,
+                atomic=False,
+                must=not dp.has_default_chg_spin(),
+                high_prec=False,
+            )
+
+    def evaluate_chunk(
+        self,
+        data: DeepmdData,
+        test_data: dict,
+        context: ChunkContext,
+    ) -> dict[str, tuple[float, float]]:
+        """Evaluate one chunk of a density test."""
+        dp = self.dp
+        mixed_type = data.mixed_type
+        nframes = test_data["box"].shape[0]
+
+        coord = test_data["coord"].reshape([nframes, -1])
+        box = test_data["box"] if data.pbc else None
+        if mixed_type:
+            atype = test_data["type"].reshape([nframes, -1])
+        else:
+            atype = test_data["type"][0]
+        fparam = test_data["fparam"] if dp.get_dim_fparam() > 0 else None
+        aparam = test_data["aparam"] if dp.get_dim_aparam() > 0 else None
+        charge_spin = (
+            test_data["charge_spin"]
+            if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 0.0) != 0.0
+            else None
+        )
+        # the loader accepts both (nframes, ngrid, 3) and the flattened
+        # (nframes, ngrid*3) layout; DeepDensity.eval only accepts the
+        # former, so normalize here and keep one contract at the eval call
+        grid = test_data["grid"].reshape(nframes, -1, 3)
+
+        prediction = dp.eval(
+            coord,
+            box,
+            atype,
+            grid,
+            fparam=fparam,
+            aparam=aparam,
+            charge_spin=charge_spin,
+            mixed_type=mixed_type,
+        ).reshape(nframes, -1)
+        label = test_data["density"].reshape(nframes, -1)
+
+        diff = prediction - label
+        errors: dict[str, tuple[float, float]] = {
+            "mae_density": (mae(diff), diff.size),
+            "rmse_density": (rmse(diff), diff.size),
+        }
+
+        if context.detail_path is not None:
+            _write_per_frame_details(
+                context,
+                suffix="density",
+                reference=label,
+                prediction=prediction,
+            )
+
+        return errors

deepmd/pd/utils/dataset.py

@@ -57,4 +57,5 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N
                 default=data_item["default"],
                 dtype=data_item["dtype"],
                 output_natoms_for_type_sel=data_item["output_natoms_for_type_sel"],
+                special_shape=getattr(data_item, "special_shape", None),
             )