Chore: refactor InvarFitting

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -320,14 +320,16 @@ def _compute_weight(
             ),
             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

@@ -13,7 +13,7 @@
 log = logging.getLogger(__name__)
 
 
-class DipoleFittingNetType(Fitting):
+class DipoleFittingNet(Fitting):
     def __init__(
         self, ntypes, embedding_width, neuron, out_dim, resnet_dt=True, **kwargs
     ):