Feat: add DOS net

deepmd/dpmodel/fitting/__init__.py

@@ -2,6 +2,9 @@
 from .dipole_fitting import (
     DipoleFitting,
 )
+from .dos_fitting import (
+    DOSFittingNet,
+)
 from .ener_fitting import (
     EnergyFittingNet,
 )
@@ -21,4 +24,5 @@
     "DipoleFitting",
     "EnergyFittingNet",
     "PolarFitting",
+    "DOSFittingNet",
 ]

deepmd/dpmodel/fitting/dos_fitting.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
+from typing import (
+    TYPE_CHECKING,
+    List,
+    Optional,
+    Union,
+)
+
+import numpy as np
+
+from deepmd.dpmodel.common import (
+    DEFAULT_PRECISION,
+)
+from deepmd.dpmodel.fitting.invar_fitting import (
+    InvarFitting,
+)
+
+if TYPE_CHECKING:
+    from deepmd.dpmodel.fitting.general_fitting import (
+        GeneralFitting,
+    )
+
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
+
+
+@InvarFitting.register("dos")
+class DOSFittingNet(InvarFitting):
+    def __init__(
+        self,
+        ntypes: int,
+        dim_descrpt: int,
+        numb_dos: int = 300,
+        neuron: List[int] = [120, 120, 120],
+        resnet_dt: bool = True,
+        numb_fparam: int = 0,
+        numb_aparam: int = 0,
+        bias_dos: Optional[np.ndarray] = None,
+        rcond: Optional[float] = None,
+        trainable: Union[bool, List[bool]] = True,
+        activation_function: str = "tanh",
+        precision: str = DEFAULT_PRECISION,
+        mixed_types: bool = False,
+        exclude_types: List[int] = [],
+        # not used
+        seed: Optional[int] = None,
+    ):
+        if bias_dos is not None:
+            self.bias_dos = bias_dos
+        else:
+            self.bias_dos = np.zeros((ntypes, numb_dos), dtype=float)
+        super().__init__(
+            var_name="dos",
+            ntypes=ntypes,
+            dim_descrpt=dim_descrpt,
+            dim_out=numb_dos,
+            neuron=neuron,
+            resnet_dt=resnet_dt,
+            bias_atom=bias_dos,
+            numb_fparam=numb_fparam,
+            numb_aparam=numb_aparam,
+            rcond=rcond,
+            trainable=trainable,
+            activation_function=activation_function,
+            precision=precision,
+            mixed_types=mixed_types,
+            exclude_types=exclude_types,
+        )
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "GeneralFitting":
+        data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("var_name")
+        data["numb_dos"] = data.pop("dim_out")
+        data.pop("tot_ener_zero")
+        data.pop("layer_name")
+        data.pop("use_aparam_as_mask")
+        data.pop("spin")
+        data.pop("atom_ener")
+        return super().deserialize(data)
+
+    def serialize(self) -> dict:
+        """Serialize the fitting to dict."""
+        dd = {
+            **super().serialize(),
+            "type": "dos",
+        }
+        dd["@variables"]["bias_atom_e"] = self.bias_atom_e
+
+        return dd

deepmd/dpmodel/fitting/general_fitting.py

@@ -40,6 +40,8 @@
             The dimension of the input descriptor.
     neuron
             Number of neurons :math:`N` in each hidden layer of the fitting net
+    bias_atom_e
+            Average enery per atom for each element.
     resnet_dt
             Time-step `dt` in the resnet construction:
             :math:`y = x + dt * \phi (Wx + b)`
@@ -85,6 +87,7 @@
         resnet_dt: bool = True,
         numb_fparam: int = 0,
         numb_aparam: int = 0,
+        bias_atom_e: Optional[np.ndarray] = None,
         rcond: Optional[float] = None,
         tot_ener_zero: bool = False,
         trainable: Optional[List[bool]] = None,
@@ -125,7 +128,11 @@
 
         net_dim_out = self._net_out_dim()
         # init constants
-        self.bias_atom_e = np.zeros([self.ntypes, net_dim_out])
+        if bias_atom_e is None:
+            self.bias_atom_e = np.zeros([self.ntypes, net_dim_out])
+        else:
+            assert bias_atom_e.shape == (self.ntypes, net_dim_out)
+            self.bias_atom_e = bias_atom_e
         if self.numb_fparam > 0:
             self.fparam_avg = np.zeros(self.numb_fparam)
             self.fparam_inv_std = np.ones(self.numb_fparam)

deepmd/dpmodel/fitting/invar_fitting.py

@@ -82,6 +82,8 @@ class InvarFitting(GeneralFitting):
             Number of atomic parameter
     rcond
             The condition number for the regression of atomic energy.
+    bias_atom
+            Bias for each element.
     tot_ener_zero
             Force the total energy to zero. Useful for the charge fitting.
     trainable
@@ -117,10 +119,11 @@ def __init__(
         resnet_dt: bool = True,
         numb_fparam: int = 0,
         numb_aparam: int = 0,
+        bias_atom: Optional[np.ndarray] = None,
         rcond: Optional[float] = None,
         tot_ener_zero: bool = False,
         trainable: Optional[List[bool]] = None,
-        atom_ener: Optional[List[float]] = [],
+        atom_ener: Optional[List[float]] = None,
         activation_function: str = "tanh",
         precision: str = DEFAULT_PRECISION,
         layer_name: Optional[List[Optional[str]]] = None,
@@ -152,6 +155,7 @@ def __init__(
             numb_fparam=numb_fparam,
             numb_aparam=numb_aparam,
             rcond=rcond,
+            bias_atom_e=bias_atom,
             tot_ener_zero=tot_ener_zero,
             trainable=trainable,
             activation_function=activation_function,

deepmd/pt/model/model/dos_model.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+    Optional,
+)
+
+import torch
+
+from .dp_model import (
+    DPModel,
+)
+
+
+class DOSModel(DPModel):
+    model_type = "dos"
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self,
+        coord,
+        atype,
+        box: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        model_ret = self.forward_common(
+            coord,
+            atype,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        if self.get_fitting_net() is not None:
+            model_predict = {}
+            model_predict["atom_dos"] = model_ret["dos"]
+            model_predict["dos"] = model_ret["dos_redu"]
+
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
+        else:
+            model_predict = model_ret
+            model_predict["updated_coord"] += coord
+        return model_predict
+
+    @torch.jit.export
+    def forward_lower(
+        self,
+        extended_coord,
+        extended_atype,
+        nlist,
+        mapping: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ):
+        model_ret = self.forward_common_lower(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        if self.get_fitting_net() is not None:
+            model_predict = {}
+            model_predict["atom_dos"] = model_ret["dos"]
+            model_predict["dos"] = model_ret["energy_redu"]
+
+        else:
+            model_predict = model_ret
+        return model_predict

deepmd/pt/model/model/dp_model.py

@@ -18,6 +18,9 @@
 from deepmd.pt.model.task.dipole import (
     DipoleFittingNet,
 )
+from deepmd.pt.model.task.dos import (
+    DOSFittingNet,
+)
 from deepmd.pt.model.task.ener import (
     EnergyFittingNet,
     EnergyFittingNetDirect,
@@ -45,6 +48,9 @@
         from deepmd.pt.model.model.dipole_model import (
             DipoleModel,
         )
+        from deepmd.pt.model.model.dos_model import (
+            DOSModel,
+        )
         from deepmd.pt.model.model.ener_model import (
             EnergyModel,
         )
@@ -68,6 +74,8 @@
                 cls = DipoleModel
             elif isinstance(fitting, PolarFittingNet):
                 cls = PolarModel
+            elif isinstance(fitting, DOSFittingNet):
+                cls = DOSModel
             # else: unknown fitting type, fall back to DPModel
         return super().__new__(cls)
 

deepmd/pt/model/task/dos.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
+import logging
+from typing import (
+    List,
+    Optional,
+    Union,
+)
+
+import torch
+
+from deepmd.pt.model.task.ener import (
+    InvarFitting,
+)
+from deepmd.pt.model.task.fitting import (
+    Fitting,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.env import (
+    DEFAULT_PRECISION,
+)
+from deepmd.pt.utils.utils import (
+    to_numpy_array,
+)
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
+
+dtype = env.GLOBAL_PT_FLOAT_PRECISION
+device = env.DEVICE
+
+log = logging.getLogger(__name__)
+
+
+@Fitting.register("dos")
+class DOSFittingNet(InvarFitting):
+    def __init__(
+        self,
+        ntypes: int,
+        dim_descrpt: int,
+        numb_dos: int = 300,
+        neuron: List[int] = [128, 128, 128],
+        resnet_dt: bool = True,
+        numb_fparam: int = 0,
+        numb_aparam: int = 0,
+        rcond: Optional[float] = None,
+        bias_dos: Optional[torch.Tensor] = None,
+        trainable: Union[bool, List[bool]] = True,
+        seed: Optional[int] = None,
+        activation_function: str = "tanh",
+        precision: str = DEFAULT_PRECISION,
+        exclude_types: List[int] = [],
+        mixed_types: bool = True,
+    ):
+        if bias_dos is not None:
+            self.bias_dos = bias_dos
+        else:
+            self.bias_dos = torch.zeros((ntypes, numb_dos), dtype=float)
+        super().__init__(
+            var_name="dos",
+            ntypes=ntypes,
+            dim_descrpt=dim_descrpt,
+            dim_out=numb_dos,
+            neuron=neuron,
+            bias_atom_e=bias_dos,
+            resnet_dt=resnet_dt,
+            numb_fparam=numb_fparam,
+            numb_aparam=numb_aparam,
+            activation_function=activation_function,
+            precision=precision,
+            mixed_types=mixed_types,
+            rcond=rcond,
+            seed=seed,
+            exclude_types=exclude_types,
+            trainable=trainable,
+        )
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "DOSFittingNet":
+        data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("@class")
+        data.pop("var_name")
+        data.pop("tot_ener_zero")
+        data.pop("layer_name")
+        data.pop("use_aparam_as_mask")
+        data.pop("spin")
+        data.pop("atom_ener")
+        data["numb_dos"] = data.pop("dim_out")
+        obj = super().deserialize(data)
+
+        return obj
+
+    def serialize(self) -> dict:
+        """Serialize the fitting to dict."""
+        # dd = super(InvarFitting, self).serialize()
+        dd = {
+            **InvarFitting.serialize(self),
+            "type": "dos",
+            "dim_out": self.dim_out,
+        }
+        dd["@variables"]["bias_atom_e"] = to_numpy_array(self.bias_atom_e)
+
+        return dd
+
+    # make jit happy with torch 2.0.0
+    exclude_types: List[int]

deepmd/pt/model/task/fitting.py

@@ -403,7 +403,7 @@
 
     def __setitem__(self, key, value):
         if key in ["bias_atom_e"]:
-            value = value.view([self.ntypes, self._net_out_dim()])
+            value = value.view([self.ntypes, -1])
             self.bias_atom_e = value
         elif key in ["fparam_avg"]:
             self.fparam_avg = value