feat(pt): add charge density prediction support

deepmd/infer/deep_density.py

@@ -0,0 +1,121 @@
+# 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 density is predicted on grid points rather than on atoms, but it
+        is declared with the same per-site output definition as the fitting
+        net of the model.
+        """
+        return ModelOutputDef(
+            FittingOutputDef(
+                [
+                    OutputVariableDef(
+                        "density",
+                        [1],
+                        reducible=True,
+                        r_differentiable=True,
+                        c_differentiable=True,
+                    ),
+                ]
+            )
+        )
+
+    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,
+        **kwargs: dict[str, 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.
+        **kwargs : dict[str, Any]
+            Keyword arguments.
+
+        Returns
+        -------
+        density
+            The density on the grid points, in shape (nframes, ngrid).
+        """
+        (
+            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,
+            False,
+            fparam=fparam,
+            aparam=aparam,
+            grid=np.array(grid),
+            **kwargs,
+        )
+        return results["density"].reshape(nframes, -1)
+
+
+__all__ = ["DeepDensity"]

deepmd/infer/deep_pot.py

@@ -212,6 +212,11 @@ def eval(
             aparam=aparam,
             **kwargs,
         )
+        # TODO: if the grid is requested, we can directly return it without reshaping to energy, force and virial. We can also consider to return the grid in a separate key in the results dict, instead of reshaping it to energy, force and virial.
+        if "grid" in kwargs:
+            result = results["density"].reshape(nframes, -1)
+            return result
+
         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,93 @@
+# 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 "atomic" so the loader keeps the
+        # frame-major layout without reshaping to natoms; see the grid/density
+        # early return in DeepmdData._load_data.
+        data.add("grid", 3, atomic=True, must=True, high_prec=True)
+        data.add("density", 1, atomic=True, must=True, high_prec=True)
+        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
+            )
+
+    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
+        grid = test_data["grid"]
+
+        prediction = dp.eval(
+            coord,
+            box,
+            atype,
+            grid,
+            fparam=fparam,
+            aparam=aparam,
+            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/pt/infer/deep_eval.py

@@ -20,6 +20,9 @@
     OutputVariableCategory,
     OutputVariableDef,
 )
+from deepmd.infer.deep_density import (
+    DeepDensity,
+)
 from deepmd.infer.deep_dipole import (
     DeepDipole,
 )
@@ -440,6 +443,8 @@ def model_type(self) -> type["DeepEvalWrapper"]:
             return DeepWFC
         elif "population" in model_output_type:
             return DeepPopulation
+        elif "density" in model_output_type:
+            return DeepDensity
         elif self.get_var_name() in model_output_type:
             return DeepProperty
         else:
@@ -552,6 +557,17 @@ def eval(
             coords, atom_types, len(atom_types.shape) > 1
         )
         request_defs = self._get_request_defs(atomic)
+        if "grid" in kwargs and kwargs["grid"] is not None:
+            out = self._eval_func(self._eval_model_density, numb_test, natoms)(
+                coords,
+                cells,
+                atom_types,
+                np.array(kwargs["grid"]),
+                fparam,
+                aparam,
+                request_defs,
+            )
+            return {"density": out}
         if "spin" not in kwargs or kwargs["spin"] is None:
             out = self._eval_func(self._eval_model, numb_test, natoms)(
                 coords, cells, atom_types, fparam, aparam, request_defs, charge_spin
@@ -916,6 +932,80 @@ def _eval_model_spin(
                 )  # this is kinda hacky
         return tuple(results)
 
+    def _eval_model_density(
+        self,
+        coords: np.ndarray,
+        cells: np.ndarray | None,
+        atom_types: np.ndarray,
+        grid: np.ndarray,
+        fparam: np.ndarray | None,
+        aparam: np.ndarray | None,
+        request_defs: list[OutputVariableDef],
+    ) -> tuple[np.ndarray, ...]:
+        model = self.dp.to(DEVICE)
+
+        nframes = coords.shape[0]
+        if len(atom_types.shape) == 1:
+            natoms = len(atom_types)
+            atom_types = np.tile(atom_types, nframes).reshape(nframes, -1)
+        else:
+            natoms = len(atom_types[0])
+
+        coord_input = torch.tensor(
+            coords.reshape([nframes, natoms, 3]),
+            dtype=GLOBAL_PT_FLOAT_PRECISION,
+            device=DEVICE,
+        )
+        type_input = torch.tensor(atom_types, dtype=torch.long, device=DEVICE)
+        grid_input = torch.tensor(
+            grid.reshape([nframes, -1, 3]),
+            dtype=GLOBAL_PT_FLOAT_PRECISION,
+            device=DEVICE,
+        )
+        ngrid = grid_input.shape[1]
+        if cells is not None:
+            box_input = torch.tensor(
+                cells.reshape([nframes, 3, 3]),
+                dtype=GLOBAL_PT_FLOAT_PRECISION,
+                device=DEVICE,
+            )
+        else:
+            box_input = None
+        if fparam is not None:
+            fparam_input = to_torch_tensor(
+                fparam.reshape(nframes, self.get_dim_fparam())
+            )
+        else:
+            fparam_input = None
+        if aparam is not None:
+            aparam_input = to_torch_tensor(
+                aparam.reshape(nframes, natoms, self.get_dim_aparam())
+            )
+        else:
+            aparam_input = None
+
+        do_atomic_virial = any(
+            x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs
+        )
+        batch_output = model(
+            coord_input,
+            type_input,
+            grid=grid_input,
+            box=box_input,
+            do_atomic_virial=do_atomic_virial,
+            fparam=fparam_input,
+            aparam=aparam_input,
+        )
+        if isinstance(batch_output, tuple):
+            batch_output = batch_output[0]
+
+        results = []
+        pt_name = "density"
+        density_shape = [nframes, ngrid]
+        out = batch_output[pt_name].reshape(density_shape).detach().cpu().numpy()
+        results.append(out)
+        return tuple(results)
+
     def _get_output_shape(
         self, odef: OutputVariableDef, nframes: int, natoms: int
     ) -> list[int]:

deepmd/pt/loss/__init__.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from .charge import (
+    GridDensityLoss,
+)
 from .denoise import (
     DenoiseLoss,
 )
@@ -35,6 +38,7 @@
     "EnergyHessianStdLoss",
     "EnergySpinLoss",
     "EnergyStdLoss",
+    "GridDensityLoss",
     "PopulationLoss",
     "PropertyLoss",
     "TaskLoss",