Refactor SOGKSpace and SOG spline definitions for improved clarity and organization

- Updated copyright and author information in sog.h
- Reorganized member variables in SOGKSpace for better readability
- Introduced new parameters for SOG kernel and spline/grid method selection
- Enhanced method definitions and added new methods for kernel parameter finalization
- Refactored spline weight functions and node definitions in sog_spline.h
- Improved naming conventions for clarity and consistency

Author: Johnsonjz

deepmd/pt/model/atomic_model/sog_atomic_model.py

@@ -51,12 +51,16 @@ def __init__(
 
         self.descriptor = descriptor
         self.fitting_net = sog_energy_fitting
-        # self.sog_energy_fitting = self.fitting_net
         self.type_map = type_map
         self.ntypes = len(type_map)
         self.rcut = self.descriptor.get_rcut()
         self.sel = self.descriptor.get_sel()
 
+        # Apply SOG library defaults: if sigma was not explicitly set by user,
+        # compute it from the descriptor's r_cut (sigma = r_cut * nlayers / RCUT_TO_SIGMA).
+        nlayers_sog = getattr(self.descriptor, "nlayers", 1)
+        self.fitting_net.recompute_from_rcut(self.rcut, nlayers_sog)
+
         super().init_out_stat()
 
         self.enable_eval_descriptor_hook = False

deepmd/pt/model/model/sog_model.py

@@ -114,7 +114,14 @@ def _build_sog_lib_direct_kernel(
         )
 
         nlayers = getattr(self.atomic_model.descriptor, "nlayers", 1) if hasattr(self.atomic_model, "descriptor") else 1
-        # Build sog_arguments dict — prefer cubes2_phi_max, fall back to n_dl
+        # Auto-detect FFT vs direct:
+        # Explicit n_dl without cubes2_phi_max → direct k-sum (training/validation).
+        _has_n_dl = getattr(fitting, "n_dl", None) is not None
+        _has_phi = getattr(fitting, "cubes2_phi_max", None) is not None
+        _use_fft = bool(getattr(fitting, "use_cubes2_fft", False))
+        if _has_n_dl and not _has_phi:
+            _use_fft = False
+        # Build sog_arguments dict
         sog_args: dict = {
             "use_atomwise": False,
             "amp": amp_internal_runtime,
@@ -124,17 +131,21 @@ def _build_sog_lib_direct_kernel(
                 "remove_self_interaction": bool(fitting.remove_self_interaction),
                 "nufft": False,
                 "use_nufft": False,
-                "use_cubes2_fft": True,
+                "use_cubes2_fft": _use_fft,
                 "nlayers": nlayers,
                 "norm_factor": E2_PER_ANGSTROM_TO_EV,
                 "trainable_kernel": False,
+                "b": float(fitting.b),
             }
         # Prefer cubes2_phi_max (new API), fall back to n_dl (legacy)
         if getattr(fitting, "cubes2_phi_max", None) is not None:
             sog_args["cubes2_phi_max"] = float(fitting.cubes2_phi_max)
         elif getattr(fitting, "n_dl", None) is not None:
             sog_args["n_dl"] = float(fitting.n_dl)
         # else: auto-default from SOG lib's Table III
+        # Optional charge neutrality penalty (None = disabled, use physical k=0 instead)
+        if getattr(fitting, "charge_neutral_lambda", None) is not None:
+            sog_args["charge_neutral_lambda"] = float(fitting.charge_neutral_lambda)
 
         kernel = sog_lib.Sog(
             sog_arguments=sog_args,

deepmd/pt/model/task/sog_energy_fitting.py

@@ -31,9 +31,12 @@
     LRFittingNet,
 )
 
-SOG_DEFAULT_B = to_numpy_array(np.array(1.62976708826776469))
+SOG_DEFAULT_B = to_numpy_array(np.array(2.0))
 SOG_DEFAULT_SIGMA = to_numpy_array(np.array(2.180230445405648))
 SOG_DEFAULT_M = int(12)
+# r_cut / sigma = RCUT_TO_SIGMA for C¹ continuity at the cutoff.
+# Equivalent to the parameter in sog.module.gaussian.
+RCUT_TO_SIGMA = 1.9892536839080267
 
 
 @LRFittingNet.register("sog_energy")
@@ -111,6 +114,10 @@ class SOGEnergyFittingNet(LRFittingNet):
     external_kspace : bool
         If True, long-range correction is handled externally (e.g. kspace),
         and the model only provides latent charges.
+    use_cubes2_fft : bool
+        If True, use CubeS₂ + FFT for long-range computation (fast, default).
+        If False, use direct k-space summation (exact but slower for large systems).
+        The grid resolution is controlled by n_dl (k-space cutoff) when FFT is off.
     """
 
     def __init__(
@@ -147,7 +154,9 @@ def __init__(
         n_dl: float | int | None = None,
         cubes2_phi_max: float | None = None,
         remove_self_interaction: bool = False,
+        charge_neutral_lambda: float | None = None,
         external_kspace: bool = False,
+        use_cubes2_fft: bool = False,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -185,8 +194,11 @@ def __init__(
         if b_value <= 0.0:
             raise ValueError("`b` should be positive.")
 
+        self._sigma_user_set = sigma is not None
+        self.use_cubes2_fft = bool(use_cubes2_fft)
+        self.charge_neutral_lambda = charge_neutral_lambda
         if sigma is None:
-            sigma_value = SOG_DEFAULT_SIGMA  # will be overridden by sog lib via rcut
+            sigma_value = float(SOG_DEFAULT_SIGMA)  # placeholder, may be recomputed via rcut
         else:
             sigma_tensor = torch.as_tensor(sigma, dtype=dtype, device=device)
             sigma_value = float(sigma_tensor.reshape(-1)[0].item())
@@ -297,6 +309,9 @@ def serialize(self) -> dict:
             data["n_dl"] = self.n_dl  # legacy
         data["remove_self_interaction"] = bool(self.remove_self_interaction)
         data["external_kspace"] = bool(self.external_kspace)
+        data["use_cubes2_fft"] = bool(self.use_cubes2_fft)
+        if self.charge_neutral_lambda is not None:
+            data["charge_neutral_lambda"] = self.charge_neutral_lambda
         return data
 
     @classmethod
@@ -311,6 +326,8 @@ def deserialize(cls, data: dict) -> "SOGEnergyFittingNet":
 
         obj = super().deserialize(data)
 
+        obj.charge_neutral_lambda = data.get("charge_neutral_lambda", None)
+
         with torch.no_grad():
             if bandwidth_tensor is not None:
                 bw = bandwidth_tensor.to(
@@ -364,6 +381,39 @@ def deserialize(cls, data: dict) -> "SOGEnergyFittingNet":
     def _kernel_params(self) -> tuple[torch.Tensor, torch.Tensor]:
         return self.amp, self.bandwidth
 
+    def recompute_from_rcut(self, rcut: float, nlayers: int = 1) -> None:
+        """Recompute sigma, amp, bandwidth from the descriptor's r_cut.
+
+        This implements the SOG library's default sigma formula:
+            sigma = r_cut * nlayers / RCUT_TO_SIGMA
+
+        Only recomputes when sigma was NOT explicitly set by the user,
+        so explicit sigma in the config is always respected.
+        """
+        if self._sigma_user_set:
+            return  # user explicitly set sigma — don't override
+
+        new_sigma = rcut * nlayers / RCUT_TO_SIGMA
+        b_base = torch.tensor(self.b, dtype=self.amp.dtype, device=self.amp.device)
+        bw_tensor = new_sigma * torch.pow(
+            b_base,
+            torch.arange(self.M, dtype=self.amp.dtype, device=self.amp.device),
+        )
+        new_bandwidth = bw_tensor.square()
+        coef1 = float(4.0 * np.pi * np.log(self.b))
+        new_amp = torch.full_like(new_bandwidth, coef1)
+        new_amp *= new_bandwidth  # convert to sog-lib internal amplitude
+
+        self.sigma = new_sigma
+        self.amp = torch.nn.Parameter(
+            new_amp.to(device=self.amp.device, dtype=self.amp.dtype),
+            requires_grad=bool(self.trainable),
+        )
+        self.bandwidth = torch.nn.Parameter(
+            new_bandwidth.to(device=self.bandwidth.device, dtype=self.bandwidth.dtype),
+            requires_grad=bool(self.trainable),
+        )
+
     def forward(
         self,
         descriptor: torch.Tensor,

deepmd/utils/argcheck.py

@@ -2284,6 +2284,25 @@ def fitting_sog_energy() -> list[Argument]:
             + "φ = Δ/r_c grid control for CubeS₂ FFT solver. "
             + "Auto-defaults from Predescu 2020 Table III when not set.",
         ),
+        Argument(
+            "n_dl",
+            [float, int, type(None)],
+            optional=True,
+            default=None,
+            doc=doc_only_pt_supported
+            + "Legacy k-space grid density control. When set without cubes2_phi_max, "
+            + "triggers direct k-space summation (no FFT). "
+            + "Compute from ε via n_dl = 2π / sqrt(2·ln(1/ε) / bw[0]).",
+        ),
+        Argument(
+            "use_cubes2_fft",
+            bool,
+            optional=True,
+            default=False,
+            doc=doc_only_pt_supported
+            + "If True, use CubeS₂ + FFT solver. "
+            + "If False (default), use direct k-space summation (fully autograd-compatible).",
+        ),
         Argument(
             "remove_self_interaction",
             bool,

log.lammps

@@ -1,16 +1,83 @@
-LAMMPS (22 Jul 2025 - Update 4)
-units metal
-atom_style atomic
-lattice fcc 3.5
-Lattice spacing in x,y,z = 3.5 3.5 3.5
-region box block 0 1 0 1 0 1
-create_box 1 box
-Created orthogonal box = (0 0 0) to (3.5 3.5 3.5)
+LAMMPS (30 Mar 2026 - Development - 7f680de296)
+OMP_NUM_THREADS environment is not set. Defaulting to 1 thread.
+  using 1 OpenMP thread(s) per MPI task
+units           real
+atom_style      charge
+plugin load     /data/home/public/jiangzhen/dp/deepmd-kit/build/sog_lmp_develop_plugin/lmp/plugin/libdeepmd_lmp.so
+Loading plugin: deepmd pair style  by Han Wang
+Loading plugin: deepspin pair style  by Duo Zhang
+Loading plugin: compute deeptensor/atom  by Han Wang
+Loading plugin: fix dplr  by Han Wang
+Loading plugin: kspace pppm/dplr  by Han Wang
+Loading plugin: kspace les  by DeepMD contributors
+Loading plugin: kspace sog  by DeepMD contributors
+pair_style      coul/long 10.0
+kspace_style    sog 1e-5 2.0 5.027010924194599 12 spline cubes2_4 remove_self_interaction no
+read_data       /tmp/water_clean.data
+Reading data file ...
+  orthogonal box = (7.9947316 0 0) to (52.005268 44.121 54.219)
   1 by 1 by 1 MPI processor grid
-create_atoms 1 box
-Created 4 atoms
-  using lattice units in orthogonal box = (0 0 0) to (3.5 3.5 3.5)
-  create_atoms CPU = 0.000 seconds
-pair_style deepmd /tmp/nonexistent.pth
-ERROR: Unrecognized pair style 'deepmd' (src/force.cpp:275)
-Last input line: pair_style deepmd /tmp/nonexistent.pth
+  reading atoms ...
+  10368 atoms
+  read_data CPU = 0.023 seconds
+pair_coeff      * *
+thermo          1
+thermo_style    custom step pe ke etotal pxx pyy pzz pxy pxz pyz vol
+run             0
+WARNING: No fixes with time integration, atoms won't move
+For more information see https://docs.lammps.org/err0028 (src/verlet.cpp:60)
+SOG initialization ...
+  b = 2, sigma = 5.02701, M = 12
+  n_dl = 6.58236, accuracy = 1e-05
+  w0 = 0.99444649, rcut = 10
+  SOG-bandwidth grid: phi_max=0.2300 delta=2.3 -> 20x20x24
+  CubeS2 influence: grid 20x20x24 ngrid=9600
+  CubeS2 monomials built
+  CubeS2 1D integrals done, starting 3D loop
+  CubeS2 influence done
+Generated 0 of 1 mixed pair_coeff terms from geometric mixing rule
+Neighbor list info ...
+  update: every = 1 steps, delay = 0 steps, check = yes
+  max neighbors/atom: 2000, page size: 100000
+  master list distance cutoff = 12
+  ghost atom cutoff = 12
+  binsize = 6, bins = 8 8 10
+  1 neighbor lists, perpetual/occasional/extra = 1 0 0
+  (1) pair coul/long, perpetual
+      attributes: half, newton on
+      pair build: half/bin/atomonly/newton
+      stencil: half/bin/3d
+      bin: standard
+  SOG virial: max |force·r - Fourier|/max = 1.887629
+    force·r:       0.1158       0.0227      -0.0662      -0.0541       0.6101       1.1526
+    Fourier:      -0.2882      -0.6333      -0.7060       0.0480       0.1175       0.0562
+Per MPI rank memory allocation (min/avg/max) = 24.09 | 24.09 | 24.09 Mbytes
+   Step         PotEng         KinEng         TotEng          Pxx            Pyy            Pzz            Pxy            Pxz            Pyz           Volume    
+         0  -1484147        0             -1484147       -325151.18     -322689.17     -323022.85      1245.8592      623.70609     -62.88239       105281.85    
+Loop time of 5.4e-07 on 1 procs for 0 steps with 10368 atoms
+
+185.2% CPU use with 1 MPI tasks x 1 OpenMP threads
+
+MPI task timing breakdown:
+Section |  min time  |  avg time  |  max time  |%varavg| %total
+---------------------------------------------------------------
+Pair    | 0          | 0          | 0          |   0.0 |  0.00
+Kspace  | 0          | 0          | 0          |   0.0 |  0.00
+Neigh   | 0          | 0          | 0          |   0.0 |  0.00
+Comm    | 0          | 0          | 0          |   0.0 |  0.00
+Output  | 0          | 0          | 0          |   0.0 |  0.00
+Modify  | 0          | 0          | 0          |   0.0 |  0.00
+Other   |            | 5.4e-07    |            |       |100.00
+
+Nlocal:          10368 ave       10368 max       10368 min
+Histogram: 1 0 0 0 0 0 0 0 0 0
+Nghost:          25333 ave       25333 max       25333 min
+Histogram: 1 0 0 0 0 0 0 0 0 0
+Neighs:    3.69241e+06 ave 3.69241e+06 max 3.69241e+06 min
+Histogram: 1 0 0 0 0 0 0 0 0 0
+
+Total # of neighbors = 3692412
+Ave neighs/atom = 356.13542
+Neighbor list builds = 0
+Dangerous builds = 0
+Total wall time: 0:00:00

source/api_cc/src/DeepPotPT.cc

@@ -479,6 +479,7 @@ void DeepPotPT::compute_with_charge(
                         nall_real);
 
   atom_charge.clear();
+  int nchannels = 1;
   if (outputs.contains("latent_charge")) {
     torch::Tensor latent_tensor = outputs.at("latent_charge").toTensor();
     if (latent_tensor.dim() != 3 || latent_tensor.size(0) != nframes) {
@@ -489,25 +490,25 @@ void DeepPotPT::compute_with_charge(
       throw deepmd::deepmd_exception(
           "latent_charge must have last dimension >= 1 to map to atom->q.");
     }
-    // Keep LAMMPS q mapping consistent with Python-side validation path:
-    // use the first latent_charge channel as scalar q.
-    latent_tensor = latent_tensor.slice(/*dim=*/2, /*start=*/0, /*end=*/1);
+    // Preserve all latent_charge channels (dim_out_lr).
+    // Flatten as [atom0_ch0, atom0_ch1, ..., atom0_chN, atom1_ch0, ...]
+    nchannels = static_cast<int>(latent_tensor.size(2));
     torch::Tensor flat_latent_ =
-        latent_tensor.squeeze(-1).contiguous().view({-1}).to(floatType);
+        latent_tensor.contiguous().view({-1}).to(floatType);
     torch::Tensor cpu_latent_ = flat_latent_.to(torch::kCPU);
     std::vector<VALUETYPE> dcharge_local;
     dcharge_local.assign(cpu_latent_.data_ptr<VALUETYPE>(),
                          cpu_latent_.data_ptr<VALUETYPE>() +
                              cpu_latent_.numel());
-    if (dcharge_local.size() != static_cast<size_t>(nframes * nloc)) {
+    if (dcharge_local.size() != static_cast<size_t>(nframes * nloc * nchannels)) {
       throw deepmd::deepmd_exception(
           "latent_charge size is inconsistent with local atom count.");
     }
-    std::vector<VALUETYPE> dcharge_nall_real(nall_real, 0);
+    std::vector<VALUETYPE> dcharge_nall_real(nall_real * nchannels, 0);
     std::copy(dcharge_local.begin(), dcharge_local.end(),
               dcharge_nall_real.begin());
-    atom_charge.resize(static_cast<size_t>(nframes) * fwd_map.size(), 0);
-    select_map<VALUETYPE>(atom_charge, dcharge_nall_real, bkw_map, 1, nframes,
+    atom_charge.resize(static_cast<size_t>(nframes) * fwd_map.size() * nchannels, 0);
+    select_map<VALUETYPE>(atom_charge, dcharge_nall_real, bkw_map, nchannels, nframes,
                           fwd_map.size(), nall_real);
   }
 

source/lmp/pair_deepmd.cpp

@@ -120,6 +120,7 @@ static const char cite_user_deepmd_package[] =
 PairDeepMD::PairDeepMD(LAMMPS* lmp)
   : PairDeepBaseModel(lmp, cite_user_deepmd_package) {
   latent_charge_to_q = false;
+  ncharge_channels = 0;
   print_summary("  ");
 }
 
@@ -544,13 +545,25 @@ void PairDeepMD::compute(int eflag, int vflag) {
                  "latent_charge_to_q is enabled but model output has no "
                  "latent_charge");
     }
-    if (dcharge.size() < static_cast<size_t>(nlocal)) {
+    // Detect number of charge channels from dcharge size.
+    // dcharge is flat: [atom0_ch0, atom0_ch1, ..., atom1_ch0, ...]
+    // dcharge includes ghost atoms: size = nall * nchannels
+    int nall = nlocal + atom->nghost;
+    ncharge_channels = static_cast<int>(dcharge.size()) / nall;
+    if (dcharge.size() % nall != 0) {
       error->all(FLERR,
-                 "latent_charge size is smaller than local atom count");
+                 "latent_charge size is not a multiple of total atom count");
     }
+    if (ncharge_channels < 1) {
+      error->all(FLERR,
+                 "latent_charge size is smaller than atom count");
+    }
+    // Store full multi-channel charge for kspace access
+    dcharge_multi = dcharge;
+    // Set atom->q to channel 0 for backward compatibility (existing sog kspace)
     double* q = atom->q;
     for (int ii = 0; ii < nlocal; ++ii) {
-      q[ii] = dcharge[ii];
+      q[ii] = dcharge[ii * ncharge_channels];
     }
   }
 

source/lmp/pair_deepmd.h

@@ -57,6 +57,11 @@ class PairDeepMD : public PairDeepBaseModel {
  private:
   CommBrickDeepMD* commdata_;
   bool latent_charge_to_q;
+
+ public:
+  // Multi-channel charge access for kspace (fast_dp_sog)
+  int ncharge_channels;
+  std::vector<double> dcharge_multi;
 };
 
 }  // namespace LAMMPS_NS