Feat: add DOS net

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,
@@ -51,6 +54,7 @@
         from deepmd.pt.model.model.polar_model import (
             PolarModel,
         )
+        from deepmd.pt.model.model.dos_model
 
         if atomic_model_ is not None:
             fitting = atomic_model_.fitting_net
@@ -68,6 +72,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,86 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
+import logging
+from typing import (
+    List,
+    Optional,
+)
+
+import torch
+
+from deepmd.pt.model.task.fitting import (
+    InvarFitting,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.env import (
+    DEFAULT_PRECISION,
+)
+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,
+        neuron: List[int] = [128, 128, 128],
+        resnet_dt: bool = True,
+        numb_fparam: int = 0,
+        numb_aparam: int = 0,
+        numb_dos: int = 300,
+        rcond: Optional[float] = None,
+        bias_dos: Optional[torch.Tensor] = None,
+        trainable: Optional[List[bool]] = None,
+        seed: Optional[int] = None,
+        activation_function: str = "tanh",
+        precision: str = DEFAULT_PRECISION,
+        exclude_types: List[int] = [],
+        mixed_types: bool = True,
+        **kwargs,
+    ):
+        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,
+            **kwargs,
+        )
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "DOSFittingNet":
+        data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("var_name")
+        data.pop("dim_out")
+        return super().deserialize(data)
+
+    def serialize(self) -> dict:
+        """Serialize the fitting to dict."""
+        return {
+            **super().serialize(),
+            "type": "dos",
+        }
+
+    # make jit happy with torch 2.0.0
+    exclude_types: List[int]