consistent energy model

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -13,6 +13,7 @@
     DescrptSeA,
 )
 from deepmd.dpmodel.fitting import (  # noqa # TODO: should import all fittings!
+    EnergyFittingNet,
     InvarFitting,
 )
 from deepmd.dpmodel.output_def import (

deepmd/dpmodel/fitting/__init__.py

@@ -2,6 +2,9 @@
 from .dipole_fitting import (
     DipoleFitting,
 )
+from .ener_fitting import (
+    EnergyFittingNet,
+)
 from .invar_fitting import (
     InvarFitting,
 )
@@ -16,5 +19,6 @@
     "InvarFitting",
     "make_base_fitting",
     "DipoleFitting",
+    "EnergyFittingNet",
     "PolarFitting",
 ]

deepmd/dpmodel/model/dp_model.py

@@ -7,4 +7,7 @@
     make_model,
 )
 
-DPModel = make_model(DPAtomicModel)
+
+# use "class" to resolve "Variable not allowed in type expression"
+class DPModel(make_model(DPAtomicModel)):
+    pass

deepmd/dpmodel/model/make_model.py

@@ -6,6 +6,9 @@
 
 import numpy as np
 
+from deepmd.dpmodel.common import (
+    NativeOP,
+)
 from deepmd.dpmodel.output_def import (
     ModelOutputDef,
 )
@@ -45,7 +48,7 @@ def make_model(T_AtomicModel):
 
     """
 
-    class CM(T_AtomicModel):
+    class CM(T_AtomicModel, NativeOP):
         def __init__(
             self,
             *args,

deepmd/dpmodel/model/model.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.descriptor.se_e2_a import (
+    DescrptSeA,
+)
+from deepmd.dpmodel.fitting.ener_fitting import (
+    EnergyFittingNet,
+)
+from deepmd.dpmodel.model.dp_model import (
+    DPModel,
+)
+
+
+def get_model(data: dict) -> DPModel:
+    """Get a DPModel from a dictionary.
+
+    Parameters
+    ----------
+    data : dict
+        The data to construct the model.
+    """
+    descriptor_type = data["descriptor"].pop("type")
+    fitting_type = data["fitting_net"].pop("type")
+    if descriptor_type == "se_e2_a":
+        descriptor = DescrptSeA(
+            **data["descriptor"],
+        )
+    else:
+        raise ValueError(f"Unknown descriptor type {descriptor_type}")
+    if fitting_type == "ener":
+        fitting = EnergyFittingNet(
+            ntypes=descriptor.get_ntypes(),
+            dim_descrpt=descriptor.get_dim_out(),
+            **data["fitting_net"],
+        )
+    else:
+        raise ValueError(f"Unknown fitting type {fitting_type}")
+    return DPModel(
+        descriptor=descriptor,
+        fitting=fitting,
+        type_map=data["type_map"],
+    )

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -17,6 +17,7 @@
     DescrptSeA,
 )
 from deepmd.pt.model.task.ener import (  # noqa # TODO: should import all fittings!
+    EnergyFittingNet,
     InvarFitting,
 )
 from deepmd.pt.utils.utils import (

deepmd/pt/utils/utils.py

@@ -2,6 +2,7 @@
 from typing import (
     Callable,
     Optional,
+    overload,
 )
 
 import numpy as np
@@ -51,9 +52,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             raise RuntimeError(f"activation function {self.activation} not supported")
 
 
+@overload
+def to_numpy_array(xx: torch.Tensor) -> np.ndarray:
+    ...
+
+
+@overload
+def to_numpy_array(xx: None) -> None:
+    ...
+
+
 def to_numpy_array(
-    xx: torch.Tensor,
-) -> np.ndarray:
+    xx,
+):
     if xx is None:
         return None
     assert xx is not None
@@ -67,9 +78,19 @@ def to_numpy_array(
     return xx.detach().cpu().numpy().astype(prec)
 
 
+@overload
+def to_torch_tensor(xx: np.ndarray) -> torch.Tensor:
+    ...
+
+
+@overload
+def to_torch_tensor(xx: None) -> None:
+    ...
+
+
 def to_torch_tensor(
-    xx: np.ndarray,
-) -> torch.Tensor:
+    xx,
+):
     if xx is None:
         return None
     assert xx is not None

deepmd/tf/descriptor/descriptor.py

@@ -530,7 +530,7 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor":
             The deserialized descriptor
         """
         if cls is Descriptor:
-            return Descriptor.get_class_by_input(data).deserialize(data)
+            return Descriptor.get_class_by_input(data).deserialize(data, suffix=suffix)
         raise NotImplementedError("Not implemented in class %s" % cls.__name__)
 
     def serialize(self, suffix: str = "") -> dict:

deepmd/tf/fit/ener.py

@@ -929,7 +929,7 @@ def get_loss(self, loss: dict, lr) -> Loss:
             raise RuntimeError("unknown loss type")
 
     @classmethod
-    def deserialize(cls, data: dict, suffix: str):
+    def deserialize(cls, data: dict, suffix: str = ""):
         """Deserialize the model.
 
         Parameters
@@ -956,7 +956,7 @@ def deserialize(cls, data: dict, suffix: str):
             fitting.aparam_inv_std = data["@variables"]["aparam_inv_std"]
         return fitting
 
-    def serialize(self, suffix: str) -> dict:
+    def serialize(self, suffix: str = "") -> dict:
         """Serialize the model.
 
         Returns

deepmd/tf/fit/fitting.py

@@ -6,6 +6,7 @@
 from typing import (
     Callable,
     List,
+    Type,
 )
 
 from deepmd.dpmodel.utils.network import (
@@ -50,16 +51,33 @@ class SomeFitting(Fitting):
         """
         return Fitting.__plugins.register(key)
 
+    @classmethod
+    def get_class_by_input(cls, data: dict) -> Type["Fitting"]:
+        """Get the fitting class by the input data.
+
+        Parameters
+        ----------
+        data : dict
+            The input data
+
+        Returns
+        -------
+        Fitting
+            The fitting class
+        """
+        try:
+            fitting_type = data["type"]
+        except KeyError:
+            raise KeyError("the type of fitting should be set by `type`")
+        if fitting_type in Fitting.__plugins.plugins:
+            cls = Fitting.__plugins.plugins[fitting_type]
+        else:
+            raise RuntimeError("Unknown descriptor type: " + fitting_type)
+        return cls
+
     def __new__(cls, *args, **kwargs):
         if cls is Fitting:
-            try:
-                fitting_type = kwargs["type"]
-            except KeyError:
-                raise KeyError("the type of fitting should be set by `type`")
-            if fitting_type in Fitting.__plugins.plugins:
-                cls = Fitting.__plugins.plugins[fitting_type]
-            else:
-                raise RuntimeError("Unknown descriptor type: " + fitting_type)
+            cls = cls.get_class_by_input(kwargs)
         return super().__new__(cls)
 
     @property
@@ -110,6 +128,44 @@ def get_loss(self, loss: dict, lr) -> Loss:
             the loss function
         """
 
+    @classmethod
+    def deserialize(cls, data: dict, suffix: str = "") -> "Fitting":
+        """Deserialize the fitting.
+
+        There is no suffix in a native DP model, but it is important
+        for the TF backend.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+        suffix : str, optional
+            Name suffix to identify this fitting
+
+        Returns
+        -------
+        Fitting
+            The deserialized fitting
+        """
+        if cls is Fitting:
+            return Fitting.get_class_by_input(data).deserialize(data, suffix=suffix)
+        raise NotImplementedError("Not implemented in class %s" % cls.__name__)
+
+    def serialize(self, suffix: str = "") -> dict:
+        """Serialize the fitting.
+
+        There is no suffix in a native DP model, but it is important
+        for the TF backend.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        suffix : str, optional
+            Name suffix to identify this fitting
+        """
+        raise NotImplementedError("Not implemented in class %s" % self.__name__)
+
     def serialize_network(
         self,
         ntypes: int,