fix:

1.update n_dl parameter type to float and add validation for positive values in energy fitting models
2.nufft_shift and ishift in energy and force calculation

Author: Johnsonjz

deepmd/pt/model/model/les_model.py

@@ -3,6 +3,7 @@
     Any,
 )
 
+import math
 import pytorch_finufft
 import torch
 
@@ -29,6 +30,8 @@
     make_model,
 )
 
+E2_PER_ANGSTROM_TO_EV = 14.3996454784255
+
 LESEnergyModel_ = make_model(LESEnergyAtomicModel)
 
 
@@ -172,9 +175,18 @@ def _compute_les_frame_correction_bundle(
         ).reshape(-1)[0]
         sigma = torch.clamp(sigma, min=torch.finfo(real_dtype).eps)
         remove_self_interaction = bool(fitting.remove_self_interaction)
-        n_dl = int(fitting.n_dl)
+        n_dl = float(fitting.n_dl)
+        if (not math.isfinite(n_dl)) or n_dl <= 0.0:
+            raise ValueError("`n_dl` should be a positive finite number.")
         pi_tensor = torch.tensor(torch.pi, dtype=real_dtype, device=runtime_device)
         two_pi = torch.tensor(2.0 * torch.pi, dtype=real_dtype, device=runtime_device)
+        n_dl_tensor = torch.as_tensor(n_dl, dtype=real_dtype, device=runtime_device)
+        k_sq_max = (two_pi / n_dl_tensor) ** 2
+        coulomb_to_ev = torch.as_tensor(
+            E2_PER_ANGSTROM_TO_EV,
+            dtype=real_dtype,
+            device=runtime_device,
+        )
 
         nf, nloc, _ = coord.shape
         corr = torch.zeros((nf, 1), dtype=real_dtype, device=runtime_device)
@@ -224,14 +236,17 @@ def _compute_les_frame_correction_bundle(
             cell_inv_group = cell_inv_all[frame_ids]
             g_cart_group = two_pi * torch.einsum("bik,k...->bi...", cell_inv_group, k_grid_int)
             k_sq_group = torch.sum(g_cart_group**2, dim=1)
+            k_in_cutoff = k_sq_group <= k_sq_max
 
             k_sq_safe_group = torch.where(
-                zero_mask_expand,
+                zero_mask_expand | (~k_in_cutoff),
                 torch.ones_like(k_sq_group),
                 k_sq_group,
             )
             kfac_group = torch.exp(-0.5 * (sigma**2) * k_sq_safe_group) / k_sq_safe_group
-            kfac_group = kfac_group.to(dtype=real_dtype).masked_fill(zero_mask_expand, 0.0)
+            kfac_group = kfac_group.to(dtype=real_dtype).masked_fill(
+                zero_mask_expand | (~k_in_cutoff), 0.0
+            )
 
             for local_idx, ff in enumerate(frame_ids):
                 r_raw = coord[ff]
@@ -253,6 +268,9 @@ def _compute_les_frame_correction_bundle(
                     eps=1e-4,
                     isign=-1,
                 )
+                # FINUFFT coefficients are returned in FFT order; align to centered
+                # mode ordering (-nk..nk) used by k_grid_int/kfac/g_cart.
+                recon = torch.fft.fftshift(recon, dim=(1, 2, 3))
 
                 rho_sq = recon.real.square() + recon.imag.square()
                 corr[ff, 0] = (kfac.unsqueeze(0) * rho_sq).sum() * two_pi / volume
@@ -266,6 +284,8 @@ def _compute_les_frame_correction_bundle(
                     grad_conv = (
                         1j * g_cart.unsqueeze(1).to(dtype=complex_dtype)
                     ) * conv.unsqueeze(0)
+                    # Convert back to FINUFFT FFT order before type-2 evaluation.
+                    grad_conv = torch.fft.ifftshift(grad_conv, dim=(2, 3, 4))
                     grad_field = pytorch_finufft.functional.finufft_type2(
                         nufft_points,
                         grad_conv,
@@ -288,8 +308,17 @@ def _compute_les_frame_correction_bundle(
                         ).reshape(nloc, 1, 9)
 
                 if remove_self_interaction:
-                    diag_sum = kfac.sum() * two_pi / volume
-                    corr[ff, 0] -= torch.sum(latent_charge[ff] ** 2) * diag_sum
+                    self_corr = torch.sum(latent_charge[ff] ** 2) / (
+                        sigma * torch.sqrt(two_pi)
+                    )
+                    corr[ff, 0] -= self_corr
+
+        # Convert electrostatic unit from e^2/A to eV.
+        corr = corr * coulomb_to_ev
+        if force_local is not None:
+            force_local = force_local * coulomb_to_ev
+        if virial_local is not None:
+            virial_local = virial_local * coulomb_to_ev
 
         out: dict[str, torch.Tensor] = {"corr_redu": corr}
         if force_local is not None:

deepmd/pt/model/model/sog_model.py

@@ -3,6 +3,7 @@
     Any,
 )
 
+import math
 import pytorch_finufft
 import torch
 
@@ -29,6 +30,8 @@
     make_model,
 )
 
+E2_PER_ANGSTROM_TO_EV = 14.3996454784255
+
 SOGEnergyModel_ = make_model(SOGEnergyAtomicModel)
 
 
@@ -189,9 +192,18 @@ def _compute_sog_frame_correction_bundle(
             raise ValueError("Invalid SOG `amp` value in fitting net.")
         if bandwidth.ndim != 1 or bandwidth.numel() == 0:
             raise ValueError("Invalid SOG `bandwidth` in fitting net.")
-        n_dl = int(fitting.n_dl)
+        n_dl = float(fitting.n_dl)
+        if (not math.isfinite(n_dl)) or n_dl <= 0.0:
+            raise ValueError("`n_dl` should be a positive finite number.")
         pi_tensor = torch.tensor(torch.pi, dtype=real_dtype, device=runtime_device)
         two_pi = 2.0 * pi_tensor
+        n_dl_tensor = torch.as_tensor(n_dl, dtype=real_dtype, device=runtime_device)
+        k_sq_max = (two_pi / n_dl_tensor) ** 2
+        coulomb_to_ev = torch.as_tensor(
+            E2_PER_ANGSTROM_TO_EV,
+            dtype=real_dtype,
+            device=runtime_device,
+        )
 
         nf, nloc, _ = coord.shape
         corr = torch.zeros((nf, 1), dtype=real_dtype, device=runtime_device)
@@ -242,9 +254,12 @@ def _compute_sog_frame_correction_bundle(
             cell_inv_group = cell_inv_all[frame_ids]
             g_cart_group = two_pi * torch.einsum("bik,k...->bi...", cell_inv_group, k_grid_int)
             k_sq_group = torch.sum(g_cart_group**2, dim=1)
+            k_in_cutoff = k_sq_group <= k_sq_max
 
             kfac_group = amp * bw2 * torch.exp(-0.5 * bw2 * k_sq_group.unsqueeze(-1))
-            kfac_group = kfac_group.sum(dim=-1).masked_fill(zero_mask_expand, 0.0)
+            kfac_group = kfac_group.sum(dim=-1).masked_fill(
+                zero_mask_expand | (~k_in_cutoff), 0.0
+            )
 
             for local_idx, ff in enumerate(frame_ids):
                 r_raw = coord[ff]
@@ -266,6 +281,9 @@ def _compute_sog_frame_correction_bundle(
                     eps=1e-4,
                     isign=-1,
                 )
+                # FINUFFT coefficients are returned in FFT order; align to centered
+                # mode ordering (-nk..nk) used by k_grid_int/kfac/g_cart.
+                recon = torch.fft.fftshift(recon, dim=(1, 2, 3))
 
                 rho_sq = recon.real.square() + recon.imag.square()
                 corr[ff, 0] = (kfac.unsqueeze(0) * rho_sq).sum() / (2.0 * volume)
@@ -279,6 +297,8 @@ def _compute_sog_frame_correction_bundle(
                     grad_conv = (
                         1j * g_cart.unsqueeze(1).to(dtype=complex_dtype)
                     ) * conv.unsqueeze(0)
+                    # Convert back to FINUFFT FFT order before type-2 evaluation.
+                    grad_conv = torch.fft.ifftshift(grad_conv, dim=(2, 3, 4))
                     grad_field = pytorch_finufft.functional.finufft_type2(
                         nufft_points,
                         grad_conv,
@@ -304,6 +324,13 @@ def _compute_sog_frame_correction_bundle(
                     diag_sum = kfac.sum() / (2.0 * volume)
                     corr[ff, 0] -= torch.sum(latent_charge[ff] ** 2) * diag_sum
 
+        # Convert electrostatic unit from e^2/A to eV.
+        corr = corr * coulomb_to_ev
+        if force_local is not None:
+            force_local = force_local * coulomb_to_ev
+        if virial_local is not None:
+            virial_local = virial_local * coulomb_to_ev
+
         out: dict[str, torch.Tensor] = {"corr_redu": corr}
         if force_local is not None:
             out["force_local"] = force_local

deepmd/pt/model/task/les_energy_fitting.py

@@ -93,7 +93,7 @@ class LESEnergyFittingNet(LRFittingNet):
     default_fparam: list[float], optional
         The default frame parameter. If set, when `fparam.npy` files are not included in the data system,
         this value will be used as the default value for the frame parameter in the fitting net.
-    n_dl : int
+    n_dl : float
         NUFFT long-range grid density control factor.
     remove_self_interaction : bool
         If True, remove self interaction term in long-range correction.
@@ -126,7 +126,7 @@ def __init__(
         use_aparam_as_mask: bool = False,
         default_fparam: list[float] | None = None,
         sigma: float | list[float] | torch.Tensor | None = None,
-        n_dl: int = 1,
+        n_dl: float | int = 1.0,
         remove_self_interaction: bool = False,
         **kwargs: Any,
     ) -> None:
@@ -169,7 +169,11 @@ def __init__(
             min=torch.finfo(sigma_tensor.dtype).eps,
         )
 
-        self.n_dl = max(1, int(n_dl))
+        n_dl_value = float(n_dl)
+        if (not np.isfinite(n_dl_value)) or n_dl_value <= 0.0:
+            raise ValueError("`n_dl` should be a positive finite number.")
+
+        self.n_dl = n_dl_value
         self.sigma = torch.nn.Parameter(
             sigma_tensor,
             requires_grad=bool(self.trainable),

deepmd/pt/model/task/sog_energy_fitting.py

@@ -101,7 +101,7 @@ class SOGEnergyFittingNet(LRFittingNet):
         Base bandwidth used by SOG parameterization.
     M : int
         Number of geometric bandwidth levels.
-    n_dl : int
+    n_dl : float
         NUFFT long-range grid density control factor.
     remove_self_interaction : bool
         If True, remove self interaction term in long-range correction.
@@ -138,7 +138,7 @@ def __init__(
         b: float | torch.Tensor | None = None,
         sigma: float | torch.Tensor | None = None,
         M: int | None = None,
-        n_dl: int = 1,
+        n_dl: float | int = 1.0,
         remove_self_interaction: bool = False,
         **kwargs: Any,
     ) -> None:
@@ -217,7 +217,11 @@ def __init__(
         if torch.any(bandwidth_tensor <= 0.0):
             raise ValueError("`bandwidth` values should be positive.")
 
-        self.n_dl = max(1, int(n_dl))
+        n_dl_value = float(n_dl)
+        if (not np.isfinite(n_dl_value)) or n_dl_value <= 0.0:
+            raise ValueError("`n_dl` should be a positive finite number.")
+
+        self.n_dl = n_dl_value
         self.amp = torch.nn.Parameter(
             torch.tensor([amp_value], dtype=dtype, device=device),
             requires_grad=bool(self.trainable),

deepmd/utils/argcheck.py

@@ -2263,9 +2263,9 @@ def fitting_sog_energy() -> list[Argument]:
         ),
         Argument(
             "n_dl",
-            int,
+            [float, int],
             optional=True,
-            default=1,
+            default=1.0,
             doc=doc_only_pt_supported + doc_n_dl,
         ),
         Argument(
@@ -2426,9 +2426,9 @@ def fitting_les_energy() -> list[Argument]:
         ),
         Argument(
             "n_dl",
-            int,
+            [float, int],
             optional=True,
-            default=1,
+            default=1.0,
             doc=doc_only_pt_supported + doc_n_dl,
         ),
         Argument(