Fix(dftu-pw): fix DFT+U for PW basis (nspin=1/2/4) and occupation matrix mixing

- Relax DFT+U nspin validation to support nspin=1/2/4 for PW basis
- Encapsulate Plus_U with typed accessors (get/set_locale, get_orbital_corr, get_hubbard_u, is_locale_initialized, mark_locale_dirty, enable_mixing)
- Rewrite cal_occ_pw for nspin=1/2/4 with correct becp indexing
- Add DFT+U occupation matrix mixing via Broyden method in Charge_Mixing
- Fix DFT+U locale double-counting when kpar>1
- Fix spin-channel selection: use isk[ik] instead of ik>=nk/2 in cal_occ_pw
- Extract setup_pw_dftu_indices() from cal_ps_dftu
- Propagate ld_psi for correct GEMM strides when ngk[ik] < npwx
- Restructure eff_pot_pw layout: nspin=2 uses split spin_up|spin_down
- Add get_eff_pot_pw_spin(isk) for nspin-aware access
- Add becp_ready caching in OnsiteProjector
- Update dftu_lcao.cpp for new Plus_U accessor interface
- Add unit tests for PW DFT+U (dftu_pw_test.cpp)
- Update dftu_io.cpp to use new accessors and onsite.dm format

Author: dyzheng

source/source_estate/module_charge/charge_mixing.cpp

@@ -263,11 +263,13 @@ void Charge_Mixing::allocate_mixing_uom(int uom_size)
 {
     ModuleBase::TITLE("Charge_Mixing", "allocate_mixing_uom");
     ModuleBase::timer::start("Charge_Mixing", "allocate_mixing_uom");
+    ModuleBase::timer::end("Charge_Mixing", "allocate_mixing_uom");
     // For nspin=2, uom_size already includes both spin channels
     // (eff_pot_pw.size() = pot_index * 2 for nspin=2)
     // So uom_fold should always be 1
     this->mixing->init_mixing_data(this->uom_mdata, uom_size, sizeof(double));
     this->uom_mdata.reset();
+    ModuleBase::timer::start("Charge_Mixing", "allocate_mixing_uom");
     ModuleBase::timer::end("Charge_Mixing", "allocate_mixing_uom");
     return;
 }
@@ -276,17 +278,14 @@ void Charge_Mixing::mix_uom(std::vector<double>& uom_in, std::vector<double>& uo
 {
     ModuleBase::TITLE("Charge_Mixing", "mix_uom");
     ModuleBase::timer::start("Charge_Mixing", "mix_uom");
+    ModuleBase::timer::end("Charge_Mixing", "mix_uom");
     double* uom_value_out = uom_in.data();
     double* uom_value_in = uom_save_in.data();
     // For all nspin cases, uom_array layout is already fully sized
     // and mixing operates on the entire array
     this->mixing->push_data(this->uom_mdata, uom_value_in, uom_value_out, nullptr, false);
     this->mixing->mix_data(this->uom_mdata, uom_value_out);
+    ModuleBase::timer::start("Charge_Mixing", "mix_uom");
     ModuleBase::timer::end("Charge_Mixing", "mix_uom");
-#ifdef __MPI
-    // Synchronize mixed uom across all ranks to prevent divergence
-    // after multiple Pulay steps (same pattern as mix_dmr)
-    Parallel_Common::bcast_double(uom_in.data(), uom_in.size());
-#endif
     return;
 }

source/source_estate/module_charge/charge_mixing.h

@@ -51,6 +51,7 @@ class Charge_Mixing
 
     void close_kerker_gg0() { mixing_gg0 = 0.0; mixing_gg0_mag = 0.0; }
     void conserve_setting() { mixing_beta = 0.01; mixing_beta_mag = 0.04; }
+    void set_mixing_beta(double beta) { mixing_beta = beta; }
     /**
      * @brief initialize mixing, including constructing mixing and allocating memory for mixing data
      * @brief this function should be called at eachiterinit()

source/source_lcao/module_dftu/dftu.cpp

@@ -85,36 +85,8 @@ void Plus_U::init(UnitCell& cell,
 #endif
 
     Plus_U::nspin = nspin;
-    Plus_U::orbital_corr = orbital_corr;
-    Plus_U::Yukawa = yukawa_potential;
-    this->yukawa_lambda = yukawa_lambda;
-
-    this->global_readin_dir = global_readin_dir;
-    this->global_out_dir = global_out_dir;
-    this->init_chg = init_chg;
     this->npol = npol;
-
-    if (pv != nullptr)
-    {
-        const int global_rows = pv->get_global_row_size();
-        const int global_cols = pv->get_global_col_size();
-        if (global_rows != global_cols)
-        {
-            ModuleBase::WARNING_QUIT("Plus_U::init", "Global row and column dimensions do not match");
-        }
-        if (nlocal != global_rows)
-        {
-            ModuleBase::WARNING_QUIT("Plus_U::init", "nlocal does not match global matrix dimension");
-        }
-    }
-    this->nlocal = nlocal;
-
-    this->gamma_only_local = gamma_only_local;
-    this->ks_solver = ks_solver;
-    this->cal_force = cal_force;
-    this->cal_stress = cal_stress;
     this->device = device;
-    this->kpar = kpar;
 
     // mohan update 2025-11-06
     Plus_U::energy_u = 0.0;
@@ -131,6 +103,10 @@ void Plus_U::init(UnitCell& cell,
     // it:index of type of atom
     for (int it = 0; it < cell.ntype; ++it)
     {
+        if(!has_correlated_orbital(it))
+        {
+            continue;
+        }
         for (int ia = 0; ia < cell.atoms[it].na; ia++)
         {
             // ia:index of atoms of this type
@@ -140,14 +116,6 @@ void Plus_U::init(UnitCell& cell,
             locale[iat].resize(cell.atoms[it].nwl + 1);
             locale_save[iat].resize(cell.atoms[it].nwl + 1);
 
-            // initialize the arrry iatlnm2iwt[iat][l][n][m]
-            this->iatlnmipol2iwt[iat].resize(cell.atoms[it].nwl + 1);
-
-            if(!has_correlated_orbital(it))
-            {
-                continue;
-            }
-
             const int tlp1_npol = (get_orbital_corr(it)*2+1)*npol;
             const int tlp1 = 2 * get_orbital_corr(it) + 1;
             const int elem_size = tlp1 * tlp1;
@@ -160,16 +128,16 @@ void Plus_U::init(UnitCell& cell,
     //          spin-down at eff_pot_pw[size/2 + eff_pot_pw_index[iat] + mm]
     // nspin=4: offset = sum(tlp1_npol^2) where tlp1_npol = (2l+1)*npol = 2*(2l+1)
     //          each atom occupies (2*tlp1)^2 = 4*tlp1^2 entries for 4 Pauli blocks
-            if(nspin == 4)
-            {
-                this->eff_pot_pw_index[iat] = pot_index;
-                pot_index += tlp1_npol * tlp1_npol;
-            }
-            else // nspin=1 or nspin=2: one tlp1^2 block per atom per spin channel
-            {
-                this->eff_pot_pw_index[iat] = pot_index;
-                pot_index += elem_size;
-            }
+    if(nspin == 4)
+    {
+        this->eff_pot_pw_index[iat] = pot_index;
+        pot_index += tlp1_npol * tlp1_npol;
+    }
+    else // nspin=1 or nspin=2: one tlp1^2 block per atom per spin channel
+    {
+        this->eff_pot_pw_index[iat] = pot_index;
+        pot_index += elem_size;
+    }
 
             for (int l = 0; l <= cell.atoms[it].nwl; l++)
             {
@@ -291,7 +259,7 @@ void Plus_U::init(UnitCell& cell,
         if (this->init_chg == "file")
         {
             std::stringstream sst;
-            sst << this->global_readin_dir << "dm_onsite.txt";
+            sst << PARAM.globalv.global_readin_dir << "onsite.dm";
             this->read_occup_m(cell, sst.str(), this->init_chg, nspin, npol);
 #ifdef __MPI
             this->local_occup_bcast(cell, nspin, npol);

source/source_lcao/module_dftu/dftu.h

@@ -8,7 +8,7 @@
 #ifdef __LCAO
 #include "source_basis/module_ao/ORB_read.h"
 #include "source_hamilt/hamilt.h"
-#include "source_hamilt/module_hcontainer/hcontainer.h"
+#include "source_lcao/module_hcontainer/hcontainer.h"
 #include "source_estate/module_dm/density_matrix.h"
 #include "source_lcao/force_stress_arrays.h" // mohan add 2024-06-15
 #endif

source/source_lcao/module_dftu/dftu_folding.cpp

@@ -4,8 +4,8 @@
 #include "source_io/module_parameter/parameter.h"
 #include "source_cell/module_neighbor/sltk_grid_driver.h"
 #include "source_lcao/hamilt_lcao.h"
-#include "source_hamilt/module_hcontainer/hcontainer.h"
-#include "source_hamilt/module_hcontainer/hcontainer_funcs.h"
+#include "source_lcao/module_hcontainer/hcontainer.h"
+#include "source_lcao/module_hcontainer/hcontainer_funcs.h"
 
 void Plus_U::fold_dSR_gamma(const UnitCell& ucell,
                           const Parallel_Orbitals& pv,

source/source_lcao/module_dftu/dftu_io.cpp

@@ -27,10 +27,10 @@ void Plus_U::output(const UnitCell& ucell,
 
             if (L >= get_orbital_corr(T) && has_correlated_orbital(T))
             {
-                if (L != get_orbital_corr(T))
-                {
-                    continue;
-                }
+				if (L != get_orbital_corr(T)) 
+				{
+					continue;
+				}
 
                 if (!Yukawa)
                 {
@@ -97,11 +97,11 @@ void Plus_U::write_occup_m(const UnitCell& ucell,
 
     for (int T = 0; T < ucell.ntype; T++)
     {
-        if (!has_correlated_orbital(T))
-        {
-            continue;
-        }
-        const int NL = ucell.atoms[T].nwl + 1;
+		if (!has_correlated_orbital(T)) 
+		{
+			continue;
+		}
+		const int NL = ucell.atoms[T].nwl + 1;
         const int LC = get_orbital_corr(T);
 
         for (int I = 0; I < ucell.atoms[T].na; I++)
@@ -110,10 +110,10 @@ void Plus_U::write_occup_m(const UnitCell& ucell,
 
             for (int l = 0; l < NL; l++)
             {
-                if (l != get_orbital_corr(T))
-                {
-                    continue;
-                }
+				if (l != get_orbital_corr(T)) 
+				{
+					continue;
+				}
 
                 const int N = ucell.atoms[T].l_nchi[l];
 
@@ -316,7 +316,14 @@ void Plus_U::read_occup_m(const UnitCell& ucell,
 
             for (int l = 0; l < NL; l++)
             {
-                if (l != get_orbital_corr(T))
+				if (l != get_orbital_corr(T)) 
+				{
+					continue;
+				}
+
+                ifdftu >> word;
+
+                if (strcmp("L", word) == 0)
                 {
                     continue;
                 }
@@ -397,10 +404,10 @@ void Plus_U::local_occup_bcast(const UnitCell& ucell,
 
     for (int T = 0; T < ucell.ntype; T++)
     {
-        if (!has_correlated_orbital(T))
-        {
-            continue;
-        }
+		if (!has_correlated_orbital(T)) 
+		{
+			continue;
+		}
 
         for (int I = 0; I < ucell.atoms[T].na; I++)
         {
@@ -409,10 +416,10 @@ void Plus_U::local_occup_bcast(const UnitCell& ucell,
 
             for (int l = 0; l <= ucell.atoms[T].nwl; l++)
             {
-                if (l != get_orbital_corr(T))
-                {
-                    continue;
-                }
+				if (l != get_orbital_corr(T)) 
+				{
+					continue;
+				}
 
                 for (int n = 0; n < ucell.atoms[T].l_nchi[l]; n++)
                 {

source/source_lcao/module_dftu/dftu_occup.cpp

@@ -27,9 +27,10 @@ void Plus_U::copy_locale(const UnitCell& ucell)
         {
             const int iat = ucell.itia2iat(T, I);
 
-            if (Plus_U::nspin == 4)
+            if (PARAM.inp.nspin == 4)
             {
                 locale_save[iat][target_l][0][0] = locale[iat][target_l][0][0];
+                // nspin=4 locale matrix already contains all spin components interleaved
                 if(this->uom_save.size() != 0)
                 {
                     const int size = locale[iat][target_l][0][0].nr * locale[iat][target_l][0][0].nc;
@@ -39,10 +40,11 @@ void Plus_U::copy_locale(const UnitCell& ucell)
                     }
                 }
             }
-            else if (Plus_U::nspin == 1 || Plus_U::nspin == 2)
+            else if (PARAM.inp.nspin == 1 || PARAM.inp.nspin == 2)
             {
                 locale_save[iat][target_l][0][0] = locale[iat][target_l][0][0];
                 locale_save[iat][target_l][0][1] = locale[iat][target_l][0][1];
+                // save locale matrix for spin=0,1 to uom_save
                 if(this->uom_save.size() != 0)
                 {
                     const int size = locale[iat][target_l][0][0].nr * locale[iat][target_l][0][0].nc;
@@ -107,13 +109,15 @@ void Plus_U::mix_locale(const UnitCell& ucell,
 
     for (int T = 0; T < ucell.ntype; T++)
     {
-        int target_l = get_orbital_corr(T);
-        if (target_l == -1)
-            continue;
+		if (!has_correlated_orbital(T))
+		{
+			continue;
+		}
 
         for (int I = 0; I < ucell.atoms[T].na; I++)
         {
             const int iat = ucell.itia2iat(T, I);
+            int target_l = get_orbital_corr(T);
 
             if (Plus_U::nspin == 4)
             {
@@ -171,18 +175,18 @@ void Plus_U::set_locale(const UnitCell& ucell)
         for (int I = 0; I < ucell.atoms[T].na; I++)
         {
             const int iat = ucell.itia2iat(T, I);
-            if (Plus_U::nspin == 4)
+            if (PARAM.inp.nspin == 4)
             {
                 for(int mm = 0; mm < locale[iat][l][0][0].nr * locale[iat][l][0][0].nc; mm++)
                     locale[iat][l][0][0].c[mm] = this->uom_array[eff_pot_pw_index[iat] + mm];
             }
-            else if (Plus_U::nspin == 1 || Plus_U::nspin == 2)
+            else if (PARAM.inp.nspin == 1 || PARAM.inp.nspin == 2)
             {
                 const int half_size = this->uom_array.size() / 2;
                 for(int mm = 0; mm < locale[iat][l][0][0].nr * locale[iat][l][0][0].nc; mm++)
                 {
                     locale[iat][l][0][0].c[mm] = this->uom_array[eff_pot_pw_index[iat] + mm];
-                    if (Plus_U::nspin == 2)
+                    if (PARAM.inp.nspin == 2)
                     {
                         locale[iat][l][0][1].c[mm] = this->uom_array[half_size + eff_pot_pw_index[iat] + mm];
                     }