DeePKS: support collinear spin (nspin=2) (#7433)

* deepks(nspin2): add magnetization-channel data members

Add charge/magnetization data members for collinear nspin=2 DeePKS without
changing the nspin=1 path:
  * dm_r_mag : (rho_up - rho_dn) real-space density matrix
  * pdm_mag  : magnetization projected density matrix
  * gedm_mag : dE/d(pdm_mag)
Allocated and freed only when nspin==2; for nspin==1 they stay null/empty and
the existing dm_r/pdm/gedm behave exactly as before. Scaffolding only, no
physics yet. tests/09_DeePKS: 284/284.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* deepks(nspin2): compute charge/magnetization channels, feed 2-channel model

Traditional (non-equivariant) nspin=2 collinear DeePKS now computes both the
charge and magnetization channels and evaluates a 2-channel model:
  * update_dmr gains nspin/mag options; setup_deepks fills dm_r_mag = rho_up - rho_dn
  * the operator builds pdm_mag and descriptor_mag, stacks charge+magnetization
    descriptors into model input (1, nat, 2, des), and autograds in one pass to
    fill gedm = dE/d(pdm) and gedm_mag = dE/d(pdm_mag)
The equivariant version stays a separate top-level branch; all magnetization
allocation/use is gated on nspin==2 && !deepks_equiv. The correction is not yet
applied per spin to the Hamiltonian (next step), so nspin=1 is unchanged.
tests/09_DeePKS: 284/284.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* deepks(nspin2): apply per-spin V_delta and forces/stress

SCF correction: calculate_HR takes the gedm to assemble from; for nspin=2
(traditional) the operator rebuilds V_delta_R per spin as
|alpha>(gedm +/- gedm_mag)<alpha| and toggles current_spin like Veff. The
label/output interface now feeds the 2-channel model for nspin=2 (previously
crashed a real 2-channel model post-SCF).

Forces/stress: the single live path is getForceStress -> cal_f_delta
(integral_part/ftable is dead code). cal_f_delta gains an optional
magnetization channel (dmr_mag, gedm_mag) and adds its contribution in the same
pass, so F_delta = dm_r * nlm(gedm) + dm_r_mag * nlm(gedm_mag) with correct
single symmetrize+weight on the stress.

Validation (synthetic 2-channel model):
  * non-magnetic CH4: nspin=2 == nspin=1 total energy to 2.5e-9 eV, forces to ~1e-5
  * polarized CH4 (nupdown=2): finite-difference force matches analytic to 0.05%
nspin=1 tests/09_DeePKS: 284/284 throughout. nspin=1 path unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* deepks(nspin2): write the dm_eig descriptor label with the spin channel

For nspin=2 (traditional), the interface now builds the magnetization-channel
descriptor before saving labels and reuses it for the 2-channel model. save_npy_d
gains an optional descriptor_mag: when present it writes dm_eig as (nat, 2, des)
- channel 0 = charge, channel 1 = magnetization - matching the 2-channel model
input, so deepks-kit can train on nspin=2 descriptors. nspin=1 writes (nat, des)
as before.

Validated: nspin=2 deepks_out_labels=1 on CH4 produces deepks_dm_eig.npy of
shape (5, 2, 18). nspin=1 tests/09_DeePKS: 284/284.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* deepks(nspin2): spin channel for precalc labels and bandgap

For nspin=2 (traditional), every precalc/output label now carries the
charge/magnetization channel (channel 0 = charge, 1 = magnetization), inserted
before the descriptor dimension to match dm_eig (nat, 2, des):
  * gevdm computed for both pdm and pdm_mag
  * gradvx (force) and gvepsl (stress): gdmx/gdmepsl built from dm_r_mag, then
    gvx/gvepsl per channel, stacked
  * orbpre (bandgap): orbital_precalc per channel, stacked (cal_o_delta already
    sums spin)
  * vdpre / vdrpre (v_delta precalc) and gevdm export: per-channel + stacked
The precalc kernels are unchanged (channel-agnostic); only the interface calls
them per channel and stacks. deepks_spin2 hoisted to function scope.

Validated (synthetic 2-channel model, nupdown=2): label shapes gain the size-2
channel dim (dm_eig (5,2,18), gradvx (5,3,5,2,18), gvepsl (6,5,2,18),
orbpre (2,5,2,18), vdpre (2,33,33,5,2,18)); nspin=1 shapes unchanged.
tests/09_DeePKS (nspin=1): 284/284.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(deepks): add nspin=2 integration tests (29-31)

Add spin-polarized DeePKS integration tests under tests/09_DeePKS, using a
small synthetic 2-channel (charge + magnetization) demo model committed as
Model_ProjOrb/model_nspin2_demo.ptg:
  * 29_NO_GO_deepks_scf_nspin2  : gamma SCF + force + stress (net moment)
  * 30_NO_KP_deepks_scf_nspin2  : multi-k SCF (net moment)
  * 31_NO_GO_deepks_bandgap_nspin2 : output labels (all precalc) + bandgap
All use nupdown to give a real moment so the magnetization channel is active.
tests/09_DeePKS now passes 313/313 (28 original + 3 new) under the standard
np=4 harness.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* deepks(nspin2): guard precalc/bandgap mag-channel stacks on rank 0

The cal_gvx/cal_gvepsl/cal_orbital_precalc/cal_v_delta_precalc helpers
assemble their output tensor only on the root rank; on other ranks the
returned tensor is undefined. Stacking the charge and magnetization
channels unconditionally therefore aborts on non-root ranks with
"tensor does not have a device". Restrict each two-channel torch::stack
to rank 0, matching where the labels are actually written.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* deepks(nspin2): guard undefined autograd grads and fix indentation

torch::autograd::grad(..., allow_unused=true) returns an undefined
gradient for any PDM block the model output does not depend on, which is
now reachable because the magnetization-channel PDM blocks are appended
to grad_inputs. Calling .accessor() on such a tensor aborts; check
.defined() and write zeros for missing gradients before accessing.

Also re-indent the deepks_v_delta == -2 magnetization block to match the
surrounding scope.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* deepks(nspin2): make _mag parameters adjacent and type-consistent

Group each magnetization-channel argument directly after its charge
counterpart and pass it the same way as the base, so the paired
quantities stay together and similarly-named parameters no longer use
different calling conventions (which was easy to misuse):

  cal_edelta_gedm: descriptor/descriptor_mag, pdm/pdm_mag, gedm/gedm_mag
  cal_f_delta:     dmr/dmr_mag, gedm/gedm_mag
  save_npy_d:      descriptor/descriptor_mag

descriptor_mag and pdm_mag now take const& like their bases; the
two-channel path is selected by descriptor_mag being non-empty. As the
paired arguments are no longer trailing they carry no defaults, so the
nspin=1 / equivariant call sites pass empty containers and nullptr
explicitly.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* deepks(nspin2): make _mag arguments a trailing optional group

Move the magnetization-channel arguments into a single optional group at
the end of each signature, with their charge counterparts grouped
immediately before. The _mag group carries defaults (empty container /
nullptr) so nspin=1 and equivariant call sites simply omit them instead
of passing explicit empty/null placeholders.

  cal_edelta_gedm(..., descriptor, pdm, gedm, descriptor_mag, pdm_mag, gedm_mag)
  cal_f_delta(..., dmr, gedm, dmr_mag, gedm_mag)
  save_npy_d(..., descriptor, descriptor_mag)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* deepks(nspin2): apply clang-format to changed lines

No functional change; bring the nspin=2 additions into line with the
project .clang-format (Microsoft style, one argument per line, expanded
single-line braces).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* deepks(nspin2): drop per-call gedm buffer, clarify contributeHR

Two improvements to the nspin=2 path in contributeHR, no functional
change:

- calculate_HR now folds the per-spin combination gedm + sigma*gedm_mag
  into the small per-atom gedm copy it already builds, taking the
  magnetization channel and sign as optional arguments. This removes the
  inlmax x pdm_size temporary buffer (and its inlmax per-call heap
  allocations) that was rebuilt on every contributeHR call during SCF.

- The magnetization-channel cal_pdm/cal_descriptor is hoisted into its
  own block parallel to the charge-channel preparation, so the flow reads
  prepare-charge -> prepare-mag (nspin=2) -> evaluate-model, and the
  traditional model evaluation is a single cal_edelta_gedm call (an empty
  descriptor_mag selects the single-channel path for nspin=1).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* deepks: bracket whole contributeHR with its timer

The DeePKS contributeHR timer was started and stopped inside the
get_hr_cal() guard. After the nspin=2 per-spin V_delta_R build was moved
out of that guard, the timer ended before that work ran (and was not
started at all on the second-spin call). Start the timer at function
entry and stop it at the end so it brackets all work in every call,
matching Veff::contributeHR.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* test(deepks): loosen tolerance for nspin=2 bandgap precalc case

The deepks_fpre/spre labels in case 31 are large-magnitude sums (~113,
~47) whose last digits vary at the ~1e-7 level run to run, due to the
nondeterministic OpenMP critical accumulation in calculate_HR. That
occasionally exceeds the default 1e-7 tolerance. Add a per-case
threshold of 1e-5; the small-magnitude labels (edelta, deltas) are
bit-stable and remain tightly checked, and a real regression shifts
them far more than 1e-5.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: ErjieWu

source/source_lcao/FORCE_STRESS.cpp

@@ -276,17 +276,41 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
         const int nks = (PARAM.inp.nspin == 1 || PARAM.inp.nspin == 2) ? 1 : kv.get_nks();
         if (PARAM.globalv.gamma_only_local)
         {
-            DeePKS_domain::cal_f_delta<double>(deepks.ld.dm_r, ucell, orb, gd,
-                                               *flk.ParaV, nks, deepks.ld.deepks_param,
-                                               kv.kvec_d, deepks.ld.phialpha, deepks.ld.gedm,
-                                               fvnl_dalpha, isstress, svnl_dalpha);
+            DeePKS_domain::cal_f_delta<double>(
+                ucell,
+                orb,
+                gd,
+                *flk.ParaV,
+                nks,
+                deepks.ld.deepks_param,
+                kv.kvec_d,
+                deepks.ld.phialpha,
+                fvnl_dalpha,
+                isstress,
+                svnl_dalpha,
+                deepks.ld.dm_r,
+                deepks.ld.gedm,
+                (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) ? deepks.ld.dm_r_mag : nullptr,
+                (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) ? deepks.ld.gedm_mag : nullptr);
         }
         else
         {
-            DeePKS_domain::cal_f_delta<std::complex<double>>(deepks.ld.dm_r, ucell, orb, gd,
-                                                              *flk.ParaV, nks, deepks.ld.deepks_param,
-                                                              kv.kvec_d, deepks.ld.phialpha, deepks.ld.gedm,
-                                                              fvnl_dalpha, isstress, svnl_dalpha);
+            DeePKS_domain::cal_f_delta<std::complex<double>>(
+                ucell,
+                orb,
+                gd,
+                *flk.ParaV,
+                nks,
+                deepks.ld.deepks_param,
+                kv.kvec_d,
+                deepks.ld.phialpha,
+                fvnl_dalpha,
+                isstress,
+                svnl_dalpha,
+                deepks.ld.dm_r,
+                deepks.ld.gedm,
+                (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) ? deepks.ld.dm_r_mag : nullptr,
+                (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) ? deepks.ld.gedm_mag : nullptr);
         }
 
         if (isforce)

source/source_lcao/FORCE_gamma.cpp

@@ -217,19 +217,19 @@ void Force_LCAO<double>::ftable(const bool isforce,
     {
         // No need to update E_delta here since it have been done in LCAO_Deepks_Interface in after_scf
         const int nks = 1;
-        DeePKS_domain::cal_f_delta<double>(deepks.ld.dm_r,
-                                           ucell,
+        DeePKS_domain::cal_f_delta<double>(ucell,
                                            orb,
                                            gd,
                                            *this->ParaV,
                                            nks,
                                            deepks.ld.deepks_param,
                                            kv->kvec_d,
                                            deepks.ld.phialpha,
-                                           deepks.ld.gedm,
                                            fvnl_dalpha,
                                            isstress,
-                                           svnl_dalpha);
+                                           svnl_dalpha,
+                                           deepks.ld.dm_r,
+                                           deepks.ld.gedm);
     }
 #endif
 

source/source_lcao/FORCE_k.cpp

@@ -224,19 +224,19 @@ void Force_LCAO<std::complex<double>>::ftable(const bool isforce,
     if (PARAM.inp.deepks_scf)
     {
         // No need to update E_delta since it have been done in LCAO_Deepks_Interface in after_scf
-        DeePKS_domain::cal_f_delta<std::complex<double>>(deepks.ld.dm_r,
-                                                         ucell,
+        DeePKS_domain::cal_f_delta<std::complex<double>>(ucell,
                                                          orb,
                                                          gd,
                                                          pv,
                                                          kv->get_nks(),
                                                          deepks.ld.deepks_param,
                                                          kv->kvec_d,
                                                          deepks.ld.phialpha,
-                                                         deepks.ld.gedm,
                                                          fvnl_dalpha,
                                                          isstress,
-                                                         svnl_dalpha);
+                                                         svnl_dalpha,
+                                                         deepks.ld.dm_r,
+                                                         deepks.ld.gedm);
     }
 #endif
 

source/source_lcao/module_deepks/LCAO_deepks.cpp

@@ -11,6 +11,7 @@ LCAO_Deepks<T>::LCAO_Deepks()
 {
     deepks_param.inl_index = new ModuleBase::IntArray[1];
     gedm = nullptr;
+    gedm_mag = nullptr;
     this->phialpha.resize(1);
 }
 
@@ -33,6 +34,16 @@ LCAO_Deepks<T>::~LCAO_Deepks()
         }
         delete[] gedm;
     }
+    if (gedm_mag)
+    {
+        for (int inl = 0; inl < this->deepks_param.inlmax; inl++)
+        {
+            delete[] gedm_mag[inl];
+        }
+        delete[] gedm_mag;
+    }
+    pdm_mag.clear();
+    delete dm_r_mag;
 }
 
 template <typename T>
@@ -124,6 +135,15 @@ void LCAO_Deepks<T>::init(const LCAO_Orbitals& orb,
         }
     }
 
+    if (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) // magnetization-channel PDM
+    {
+        this->pdm_mag.resize(this->deepks_param.inlmax);
+        for (int inl = 0; inl < this->deepks_param.inlmax; ++inl)
+        {
+            this->pdm_mag[inl] = torch::zeros_like(this->pdm[inl]);
+        }
+    }
+
     this->pv = &pv_in;
 
     ModuleBase::timer::end("LCAO_Deepks", "init");
@@ -202,6 +222,16 @@ void LCAO_Deepks<T>::allocate_V_delta(const int nat, const int nks)
         ModuleBase::GlobalFunc::ZEROS(this->gedm[inl], pdm_size);
     }
 
+    if (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) // magnetization-channel gedm
+    {
+        this->gedm_mag = new double*[this->deepks_param.inlmax];
+        for (int inl = 0; inl < this->deepks_param.inlmax; inl++)
+        {
+            this->gedm_mag[inl] = new double[pdm_size];
+            ModuleBase::GlobalFunc::ZEROS(this->gedm_mag[inl], pdm_size);
+        }
+    }
+
     ModuleBase::timer::end("LCAO_Deepks", "allocate_V_delta");
     return;
 }
@@ -249,6 +279,11 @@ void LCAO_Deepks<T>::init_DMR(const UnitCell& ucell,
                                    this->dm_r->insert_pair(dm_pair);
                                });
     this->dm_r->allocate(nullptr, true);
+
+    if (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) // magnetization-channel real-space DM
+    {
+        this->dm_r_mag = new hamilt::HContainer<double>(*this->dm_r);
+    }
 }
 
 template <typename T>

source/source_lcao/module_deepks/LCAO_deepks.h

@@ -79,6 +79,11 @@ class LCAO_Deepks
     /// dE/dD, autograd from loaded model(E: Ry)
     double** gedm = nullptr; //[tot_Inl][(2l+1)*(2l+1)]
 
+    // magnetization-channel (rho_up - rho_dn) counterparts, used only for nspin=2
+    hamilt::HContainer<double>* dm_r_mag = nullptr;
+    std::vector<torch::Tensor> pdm_mag;
+    double** gedm_mag = nullptr;
+
     // functions for hr status: 1. get value; 2. set value;
     int get_hr_cal()
     {