Chore: refactor InvarFitting

deepmd/dpmodel/model/linear_atomic_model.py

@@ -318,14 +318,16 @@
             ),
             axis=-1,
         )  # handle masked nnei.
-        sigma = numerator / denominator
+        with np.errstate(divide="ignore", invalid="ignore"):
+            sigma = numerator / denominator
         u = (sigma - self.sw_rmin) / (self.sw_rmax - self.sw_rmin)
         coef = np.zeros_like(u)
         left_mask = sigma < self.sw_rmin
         mid_mask = (self.sw_rmin <= sigma) & (sigma < self.sw_rmax)
         right_mask = sigma >= self.sw_rmax
         coef[left_mask] = 1
-        smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1
+        with np.errstate(invalid="ignore"):
+            smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1
         coef[mid_mask] = smooth[mid_mask]
         coef[right_mask] = 0
         self.zbl_weight = coef

deepmd/pt/model/task/__init__.py

@@ -9,7 +9,7 @@
     DenoiseNet,
 )
 from .dipole import (
-    DipoleFittingNetType,
+    DipoleFittingNet,
 )
 from .ener import (
     EnergyFittingNet,
@@ -25,7 +25,7 @@
 __all__ = [
     "FittingNetAttenLcc",
     "DenoiseNet",
-    "DipoleFittingNetType",
+    "DipoleFittingNet",
     "EnergyFittingNet",
     "EnergyFittingNetDirect",
     "Fitting",

deepmd/pt/model/task/dipole.py

@@ -1,21 +1,47 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
+from typing import (
+    List,
+    Optional,
+)
 
 import torch
 
-from deepmd.pt.model.network.network import (
-    ResidualDeep,
+from deepmd.pt.model.network.mlp import (
+    FittingNet,
+    NetworkCollection,
 )
 from deepmd.pt.model.task.fitting import (
-    Fitting,
+    GeneralFitting,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.env import (
+    DEFAULT_PRECISION,
 )
 
+dtype = env.GLOBAL_PT_FLOAT_PRECISION
+device = env.DEVICE
+
 log = logging.getLogger(__name__)
 
 
-class DipoleFittingNetType(Fitting):
+class DipoleFittingNet(GeneralFitting):
     def __init__(
-        self, ntypes, embedding_width, neuron, out_dim, resnet_dt=True, **kwargs
+        self,
+        var_name: str,
+        ntypes: int,
+        dim_descrpt: int,
+        dim_out: int,
+        neuron: List[int] = [128, 128, 128],
+        resnet_dt: bool = True,
+        numb_fparam: int = 0,
+        numb_aparam: int = 0,
+        activation_function: str = "tanh",
+        precision: str = DEFAULT_PRECISION,
+        distinguish_types: bool = False,
+        **kwargs,
     ):
         """Construct a fitting net for dipole.
 
@@ -26,23 +52,32 @@
         - bias_atom_e: Average enery per atom for each element.
         - resnet_dt: Using time-step in the ResNet construction.
         """
-        super().__init__()
-        self.ntypes = ntypes
-        self.embedding_width = embedding_width
-        self.out_dim = out_dim
-
-        filter_layers = []
-        one = ResidualDeep(
-            0, embedding_width, neuron, 0.0, out_dim=self.out_dim, resnet_dt=resnet_dt
+        super().__init__(
+            var_name=var_name,
+            ntypes=ntypes,
+            dim_descrpt=dim_descrpt,
+            dim_out=dim_out,
+            neuron=neuron,
+            resnet_dt=resnet_dt,
+            numb_fparam=numb_fparam,
+            numb_aparam=numb_aparam,
+            activation_function=activation_function,
+            precision=precision,
+            distinguish_types=distinguish_types,
+            **kwargs,
         )
-        filter_layers.append(one)
-        self.filter_layers = torch.nn.ModuleList(filter_layers)
+        self.old_impl = False  # this only supports the new implementation.
 
-        if "seed" in kwargs:
-            log.info("Set seed to %d in fitting net.", kwargs["seed"])
-            torch.manual_seed(kwargs["seed"])
-
-    def forward(self, inputs, atype, atype_tebd, rot_mat):
+    def forward(
+        self,
+        descriptor: torch.Tensor,
+        atype: torch.Tensor,
+        gr: Optional[torch.Tensor] = None,
+        g2: Optional[torch.Tensor] = None,
+        h2: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+    ):
         """Based on embedding net output, alculate total energy.
 
         Args:
@@ -55,13 +90,34 @@
         -------
         - vec_out: output vector. Its shape is [nframes, nloc, 3].
         """
-        nframes, nloc, _ = inputs.size()
-        if atype_tebd is not None:
-            inputs = torch.concat([inputs, atype_tebd], dim=-1)
-        vec_out = self.filter_layers[0](inputs)  # Shape is [nframes, nloc, m1]
-        assert list(vec_out.size()) == [nframes, nloc, self.out_dim]
-        vec_out = vec_out.view(-1, 1, self.out_dim)
-        vec_out = (
-            torch.bmm(vec_out, rot_mat).squeeze(-2).view(nframes, nloc, 3)
-        )  # Shape is [nframes, nloc, 3]
-        return vec_out
+        in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam
+
+        nframes, nloc, _ = descriptor.shape
+        gr = gr.view(nframes, nloc, -1, 3)
+        out_dim = gr.shape[2]  # m1
+        self.filter_layers = NetworkCollection(
+            1 if self.distinguish_types else 0,
+            self.ntypes,
+            network_type="fitting_network",
+            networks=[
+                FittingNet(
+                    in_dim,
+                    out_dim,
+                    self.neuron,
+                    self.activation_function,
+                    self.resnet_dt,
+                    self.precision,
+                    bias_out=True,
+                )
+                for ii in range(self.ntypes if self.distinguish_types else 1)
+            ],
+        )
+        # (nframes, nloc, m1)
+        out = self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam)[
+            self.var_name
+        ]
+        # (nframes * nloc, 1, m1)
+        out = out.view(-1, 1, out_dim)
+        # (nframes, nloc, 3)
+        out = torch.bmm(out, gr).squeeze(-2).view(nframes, nloc, 3)
+        return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)}