feat: full energy model (but not exportable)

Author: Han Wang

deepmd/dpmodel/model/make_model.py

@@ -21,6 +21,7 @@
     PRECISION_DICT,
     RESERVED_PRECISION_DICT,
     NativeOP,
+    get_xp_precision,
 )
 from deepmd.dpmodel.model.base_model import (
     BaseModel,
@@ -103,7 +104,8 @@ def model_call_from_call_lower(
             bb.reshape(nframes, 3, 3),
         )
     else:
-        coord_normalized = cc.copy()
+        xp = array_api_compat.array_namespace(cc)
+        coord_normalized = xp.reshape(cc, (nframes, nloc, 3))
     extended_coord, extended_atype, mapping = extend_coord_with_ghosts(
         coord_normalized, atype, bb, rcut
     )
@@ -371,53 +373,73 @@ def input_type_cast(
             box: Array | None = None,
             fparam: Array | None = None,
             aparam: Array | None = None,
-        ) -> tuple[Array, Array, np.ndarray | None, np.ndarray | None, str]:
+        ) -> tuple[Array, Array | None, Array | None, Array | None, Any]:
             """Cast the input data to global float type."""
-            input_prec = RESERVED_PRECISION_DICT[self.precision_dict[coord.dtype.name]]
+            xp = array_api_compat.array_namespace(coord)
+            input_dtype = coord.dtype
+            global_dtype = get_xp_precision(
+                xp, RESERVED_PRECISION_DICT[self.global_np_float_precision]
+            )
             ###
             ### type checking would not pass jit, convert to coord prec anyway
             ###
-            _lst: list[np.ndarray | None] = [
-                vv.astype(coord.dtype) if vv is not None else None
+            _lst: list[Array | None] = [
+                xp.astype(vv, input_dtype) if vv is not None else None
                 for vv in [box, fparam, aparam]
             ]
             box, fparam, aparam = _lst
-            if input_prec == RESERVED_PRECISION_DICT[self.global_np_float_precision]:
-                return coord, box, fparam, aparam, input_prec
+            if input_dtype == global_dtype:
+                return coord, box, fparam, aparam, input_dtype
             else:
-                pp = self.global_np_float_precision
                 return (
-                    coord.astype(pp),
-                    box.astype(pp) if box is not None else None,
-                    fparam.astype(pp) if fparam is not None else None,
-                    aparam.astype(pp) if aparam is not None else None,
-                    input_prec,
+                    xp.astype(coord, global_dtype),
+                    xp.astype(box, global_dtype) if box is not None else None,
+                    xp.astype(fparam, global_dtype) if fparam is not None else None,
+                    xp.astype(aparam, global_dtype) if aparam is not None else None,
+                    input_dtype,
                 )
 
         def output_type_cast(
             self,
             model_ret: dict[str, Array],
-            input_prec: str,
+            input_prec: Any,
         ) -> dict[str, Array]:
-            """Convert the model output to the input prec."""
-            do_cast = (
-                input_prec != RESERVED_PRECISION_DICT[self.global_np_float_precision]
+            """Convert the model output to the input prec.
+
+            Parameters
+            ----------
+            model_ret
+                The model output.
+            input_prec
+                The input dtype returned by ``input_type_cast``.
+            """
+            model_ret_not_none = [vv for vv in model_ret.values() if vv is not None]
+            if not model_ret_not_none:
+                return model_ret
+            xp = array_api_compat.array_namespace(model_ret_not_none[0])
+            global_dtype = get_xp_precision(
+                xp, RESERVED_PRECISION_DICT[self.global_np_float_precision]
+            )
+            ener_dtype = get_xp_precision(
+                xp, RESERVED_PRECISION_DICT[self.global_ener_float_precision]
             )
-            pp = self.precision_dict[input_prec]
+            do_cast = input_prec != global_dtype
             odef = self.model_output_def()
             for kk in odef.keys():
                 if kk not in model_ret.keys():
                     # do not return energy_derv_c if not do_atomic_virial
                     continue
                 if check_operation_applied(odef[kk], OutputVariableOperation.REDU):
                     model_ret[kk] = (
-                        model_ret[kk].astype(self.global_ener_float_precision)
+                        xp.astype(model_ret[kk], ener_dtype)
                         if model_ret[kk] is not None
                         else None
                     )
                 elif do_cast:
                     model_ret[kk] = (
-                        model_ret[kk].astype(pp) if model_ret[kk] is not None else None
+                        xp.astype(model_ret[kk], input_prec)
+                        if model_ret[kk] is not None
+                        else None
                     )
             return model_ret
 

deepmd/dpmodel/model/transform_output.py

@@ -98,6 +98,7 @@ def communicate_extended_output(
 
     """
     xp = array_api_compat.get_namespace(mapping)
+    device = array_api_compat.device(mapping)
     mapping_ = mapping
     new_ret = {}
     for kk in model_output_def.keys_outp():
@@ -117,7 +118,9 @@ def communicate_extended_output(
                         mapping, tuple(mldims + [1] * len(derv_r_ext_dims))
                     )
                     mapping = xp.tile(mapping, [1] * len(mldims) + derv_r_ext_dims)
-                    force = xp.zeros(vldims + derv_r_ext_dims, dtype=vv.dtype)
+                    force = xp.zeros(
+                        vldims + derv_r_ext_dims, dtype=vv.dtype, device=device
+                    )
                     force = xp_scatter_sum(
                         force,
                         1,
@@ -149,7 +152,9 @@ def communicate_extended_output(
                         nall = hess_1.shape[1]
                         # (1) -> [nf, nloc1, nall2, *def, 3(1), 3(2)]
                         hessian1 = xp.zeros(
-                            [*vldims, nall, *vdef.shape, 3, 3], dtype=vv.dtype
+                            [*vldims, nall, *vdef.shape, 3, 3],
+                            dtype=vv.dtype,
+                            device=device,
                         )
                         mapping_hess = xp.reshape(
                             mapping_, (mldims + [1] * (len(vdef.shape) + 3))
@@ -172,7 +177,9 @@ def communicate_extended_output(
                         nloc = hessian1.shape[2]
                         # (2) -> [nf, nloc2, nloc1, *def, 3(1), 3(2)]
                         hessian = xp.zeros(
-                            [*vldims, nloc, *vdef.shape, 3, 3], dtype=vv.dtype
+                            [*vldims, nloc, *vdef.shape, 3, 3],
+                            dtype=vv.dtype,
+                            device=device,
                         )
                         mapping_hess = xp.reshape(
                             mapping_, (mldims + [1] * (len(vdef.shape) + 3))
@@ -218,21 +225,14 @@ def communicate_extended_output(
                     virial = xp.zeros(
                         vldims + derv_c_ext_dims,
                         dtype=vv.dtype,
+                        device=device,
+                    )
+                    virial = xp_scatter_sum(
+                        virial,
+                        1,
+                        mapping,
+                        model_ret[kk_derv_c],
                     )
-                    # jax only
-                    if array_api_compat.is_jax_array(virial):
-                        from deepmd.jax.common import (
-                            scatter_sum,
-                        )
-
-                        virial = scatter_sum(
-                            virial,
-                            1,
-                            mapping,
-                            model_ret[kk_derv_c],
-                        )
-                    else:
-                        raise NotImplementedError("Only JAX arrays are supported.")
                     new_ret[kk_derv_c] = virial
                     new_ret[kk_derv_c + "_redu"] = xp.sum(new_ret[kk_derv_c], axis=1)
                 else:

deepmd/dpmodel/utils/network.py

@@ -286,11 +286,11 @@ def call(self, x):  # noqa: ANN001, ANN201
             y = xp.astype(y, x.dtype)
         y = fn(y)
         if self.idt is not None:
-            y *= self.idt
+            y = y * self.idt
         if self.resnet and self.w.shape[1] == self.w.shape[0]:
-            y += x
+            y = y + x
         elif self.resnet and self.w.shape[1] == 2 * self.w.shape[0]:
-            y += xp.concat([x, x], axis=-1)
+            y = y + xp.concat([x, x], axis=-1)
         return y
 
 

deepmd/pt_expt/model/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .ener_model import (
+    EnergyModel,
+)
+
+__all__ = [
+    "EnergyModel",
+]

deepmd/pt_expt/model/ener_model.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
+import torch
+
+from deepmd.dpmodel.model.dp_model import (
+    DPModelCommon,
+)
+from deepmd.pt_expt.atomic_model import (
+    DPEnergyAtomicModel,
+)
+
+from .make_model import (
+    make_model,
+)
+
+DPEnergyModel_ = make_model(DPEnergyAtomicModel)
+
+
+class EnergyModel(DPModelCommon, DPEnergyModel_):
+    model_type = "ener"
+
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        DPModelCommon.__init__(self)
+        DPEnergyModel_.__init__(self, *args, **kwargs)
+
+    def forward(
+        self,
+        coord: torch.Tensor,
+        atype: torch.Tensor,
+        box: torch.Tensor | None = None,
+        fparam: torch.Tensor | None = None,
+        aparam: torch.Tensor | None = None,
+        do_atomic_virial: bool = False,
+    ) -> dict[str, torch.Tensor]:
+        model_ret = self.call(
+            coord,
+            atype,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        model_predict = {}
+        model_predict["atom_energy"] = model_ret["energy"]
+        model_predict["energy"] = model_ret["energy_redu"]
+        if self.do_grad_r("energy"):
+            model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
+        if self.do_grad_c("energy"):
+            model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
+            if do_atomic_virial:
+                model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
+        return model_predict

deepmd/pt_expt/model/make_model.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
+import torch
+
+from deepmd.dpmodel.atomic_model.base_atomic_model import (
+    BaseAtomicModel,
+)
+from deepmd.dpmodel.model.make_model import make_model as make_model_dp
+from deepmd.pt_expt.common import (
+    dpmodel_setattr,
+)
+
+from .transform_output import (
+    fit_output_to_model_output,
+)
+
+
+def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type:
+    """Make a model as a derived class of an atomic model.
+
+    Wraps dpmodel's make_model with torch.nn.Module and overrides
+    forward_common_atomic to use autograd-based derivatives.
+
+    Parameters
+    ----------
+    T_AtomicModel
+        The atomic model.
+
+    Returns
+    -------
+    CM
+        The model.
+
+    """
+    DPModel = make_model_dp(T_AtomicModel)
+
+    class CM(DPModel, torch.nn.Module):
+        def __init__(
+            self,
+            *args: Any,
+            **kwargs: Any,
+        ) -> None:
+            torch.nn.Module.__init__(self)
+            DPModel.__init__(self, *args, **kwargs)
+
+        def __call__(self, *args: Any, **kwargs: Any) -> Any:
+            # Ensure torch.nn.Module.__call__ drives forward() for export/tracing.
+            return torch.nn.Module.__call__(self, *args, **kwargs)
+
+        def __setattr__(self, name: str, value: Any) -> None:
+            handled, value = dpmodel_setattr(self, name, value)
+            if not handled:
+                super().__setattr__(name, value)
+
+        def forward(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
+            """Default forward delegates to call().
+
+            Subclasses (e.g. EnergyModel) override this with output translation.
+            """
+            return self.call(*args, **kwargs)
+
+        def forward_common_atomic(
+            self,
+            extended_coord: torch.Tensor,
+            extended_atype: torch.Tensor,
+            nlist: torch.Tensor,
+            mapping: torch.Tensor | None = None,
+            fparam: torch.Tensor | None = None,
+            aparam: torch.Tensor | None = None,
+            do_atomic_virial: bool = False,
+        ) -> dict[str, torch.Tensor]:
+            atomic_ret = self.atomic_model.forward_common_atomic(
+                extended_coord,
+                extended_atype,
+                nlist,
+                mapping=mapping,
+                fparam=fparam,
+                aparam=aparam,
+            )
+            return fit_output_to_model_output(
+                atomic_ret,
+                self.atomic_output_def(),
+                extended_coord,
+                do_atomic_virial=do_atomic_virial,
+                create_graph=self.training,
+                mask=atomic_ret.get("mask"),
+            )
+
+    return CM