Unify DFT+U occupation-matrix output across PW and LCAO

source/Makefile.Objects

@@ -1012,8 +1012,8 @@ OBJS_SRCPW=h_ewald_pw.o\
     update_cell_pw.o\
     dftu_base.o\
     dftu_base_io.o\
-    dftu_base_io_read.o\
-    dftu_base_io_write.o\
+    dftu_base_read.o\
+    dftu_base_write.o\
     dftu_pw.o\
     dftu_pw_tools.o\
     yukawa_screening.o\

source/source_esolver/esolver_ks.cpp

@@ -19,6 +19,7 @@
 #include "source_io/module_output/output_log.h" // use write_head
 #include "source_estate/elecstate_print.h" // print_etot
 #include "source_pw/module_pwdft/dftu_base.h" // Plus_U_Base::u_converged in iter_finish
+#include "source_pw/module_pwdft/dftu_base_io.h"
 #include "source_hamilt/module_xc/general_exx_info.h" // for init_general_exx_info
 
 namespace ModuleESolver
@@ -271,6 +272,16 @@ void ESolver_KS::before_scf(UnitCell& ucell, const int istep)
 
 void ESolver_KS::iter_init(UnitCell& ucell, const int istep, const int iter)
 {
+    if (this->inp_->dft_plus_u > 0)
+    {
+        const DFTU_BASE::OccmatOutputCfg occmat_cfg{this->inp_->out_freq_ion,
+                                                    this->inp_->out_freq_elec,
+                                                    this->scf_iter_limit(istep),
+                                                    this->inp_->out_occ_mat,
+                                                    this->inp_->dft_plus_u};
+        DFTU_BASE::prepare_ion_step_output(PARAM.globalv.global_out_dir, istep, iter, occmat_cfg);
+    }
+
     if(this->inp_->esolver_type != "tddft")
     {
         ModuleIO::write_head(GlobalV::ofs_running, istep, iter, this->basisname);

source/source_esolver/esolver_ks_lcao.cpp

@@ -690,12 +690,8 @@ void ESolver_KS_LCAO<TK, TR>::iter_finish(UnitCell& ucell, const int istep, int&
                          this->kv,
                          this->p_chgmix->get_mixing_beta(),
                          hamilt_lcao,
-                         global_out_dir,
                          this->inp_->nspin,
-                         gamma_only_local,
-                         istep,
-                         iter,
-                         occmat_cfg);
+                         gamma_only_local);
 
     // mohan add 2025-11: push DFT+U energy from Plus_U_Base to ElecState.
     // Covers both dft_plus_u==1 (new method, energy accumulated by DFTU_onsite::contributeHR

source/source_esolver/esolver_ks_pw.cpp

@@ -260,20 +260,14 @@ void ESolver_KS_PW<T, Device>::iter_init(UnitCell& ucell, const int istep, const
 
     // update local occupations for DFT+U
     // should before lambda loop in DeltaSpin
-    const DFTU_BASE::OccmatOutputCfg init_occmat_cfg{this->inp_->out_freq_ion,
-                                                     this->inp_->out_freq_elec,
-                                                     iter_limit,
-                                                     this->inp_->out_occ_mat,
-                                                     this->inp_->dft_plus_u};
     DFTU_BASE::iter_init_dftu_pw(iter,
                                  istep,
                                  *this->dftu_,
                                  this->stp.template get_psi_t<T, Device>(),
                                  this->pelec->wg,
                                  ucell,
                                  this->p_chgmix,
-                                 PARAM.globalv.global_out_dir,
-                                 init_occmat_cfg,
+                                 this->inp_->dft_plus_u,
                                  this->kv.isk.data(),
                                  this->inp_->nspin,
                                  this->inp_->kpar);
@@ -398,9 +392,8 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
 
     // overwrite occ_mat.txt with the latest occupation matrix and the
     // actual charge-density residual of this electronic step. At the very
-    // first step (istep 0, iter 1) the PW occupation matrix does not exist
-    // yet (iter_init_dftu_pw returns before calculating it), so skip it and
-    // keep the previous behavior where occ_mat.txt first appears at iter 2.
+    // first step (istep 0, iter 1) the PW occupation matrix exists only if
+    // it was loaded from file; otherwise it is first available at iter 2.
     const int iter_limit = this->scf_iter_limit(istep);
     const DFTU_BASE::OccmatOutputCfg occmat_cfg{this->inp_->out_freq_ion,
                                                 this->inp_->out_freq_elec,
@@ -414,8 +407,8 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
     // Read globalv values once into locals so the repeated DFT+U call sites
     // below do not each re-enter the global dependency surface.
     const std::string& global_out_dir = PARAM.globalv.global_out_dir;
-    const bool latest_ready = (iter > 1 || istep > 0);
-    if (latest_ready)
+    const bool occmat_ready = (iter > 1 || istep > 0 || this->dftu_->is_occmat_ready());
+    if (occmat_ready)
     {
         DFTU_BASE::write_latest_occmat(*this->dftu_,
                                       ucell,
@@ -429,7 +422,6 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
     // At istep 0 / iter 1 the PW occupation matrix does not exist yet,
     // unless it was loaded from dm_onsite_ini.txt or a previous
     // occ_mat.txt; snapshot an "N/A" placeholder in that case.
-    const bool occmat_ready = latest_ready || this->dftu_->is_occmat_ready();
     DFTU_BASE::append_ion_step_snapshot(*this->dftu_,
                                         ucell,
                                         global_out_dir,

source/source_lcao/module_dftu/unittests/CMakeLists.txt

@@ -94,8 +94,8 @@ if(ENABLE_MPI)
         ../../../source_estate/module_dm/dmr_full.cpp
         ../../../source_pw/module_pwdft/dftu_base.cpp
         ../../../source_pw/module_pwdft/dftu_base_io.cpp
-        ../../../source_pw/module_pwdft/dftu_base_io_read.cpp
-        ../../../source_pw/module_pwdft/dftu_base_io_write.cpp
+        ../../../source_pw/module_pwdft/dftu_base_read.cpp
+        ../../../source_pw/module_pwdft/dftu_base_write.cpp
         ../../../source_pw/module_pwdft/yukawa_screening.cpp
         ../../../source_estate/occ_matrix.cpp
         ../../../source_estate/occ_mixer.cpp
@@ -148,8 +148,8 @@ if(ENABLE_MPI)
         ../../../source_estate/module_dm/dmr_full.cpp
         ../../../source_pw/module_pwdft/dftu_base.cpp
         ../../../source_pw/module_pwdft/dftu_base_io.cpp
-        ../../../source_pw/module_pwdft/dftu_base_io_read.cpp
-        ../../../source_pw/module_pwdft/dftu_base_io_write.cpp
+        ../../../source_pw/module_pwdft/dftu_base_read.cpp
+        ../../../source_pw/module_pwdft/dftu_base_write.cpp
         ../../../source_pw/module_pwdft/yukawa_screening.cpp
         ../../../source_estate/occ_matrix.cpp
         ../../../source_estate/occ_mixer.cpp

source/source_lcao/setup_dftu_lcao.cpp

@@ -2,7 +2,6 @@
 #include "source_pw/module_pwdft/dftu_base.h"
 #include "source_lcao/module_dftu/dftu_nao_occ.h"
 #include "source_lcao/module_dftu/dftu_nao_energy.h"
-#include "source_pw/module_pwdft/dftu_base_io.h" // mohan add 2025-11-08
 #include "source_io/module_parameter/parameter.h"
 #include "source_lcao/hamilt_lcao.h"
 
@@ -47,12 +46,8 @@ void finish_dftu_lcao(const bool conv_esolver,
                        const K_Vectors& kv,
                        const double mixing_beta,
                        void* hamilt_lcao,
-                       const std::string& global_out_dir,
                        int nspin,
-                       const bool gamma_only_local,
-                       int istep,
-                       int iter,
-                       const DFTU_BASE::OccmatOutputCfg& occmat_cfg)
+                       const bool gamma_only_local)
 {
     if (!dft_plus_u)
     {
@@ -81,8 +76,6 @@ void finish_dftu_lcao(const bool conv_esolver,
             DFTU_LCAO::cal_energy_correction(*dftu_ptr, ucell, PARAM.inp.nspin);
         }
     }
-    DFTU_BASE::prepare_ion_step_output(global_out_dir, istep, iter, occmat_cfg);
-
     /// use the converged occupation matrix for next MD/Relax SCF calculation
     if (conv_esolver)
     {
@@ -99,12 +92,8 @@ template void finish_dftu_lcao<double>(const bool conv_esolver,
                                         const K_Vectors& kv,
                                         const double mixing_beta,
                                         void* hamilt_lcao,
-                                        const std::string& global_out_dir,
                                         int nspin,
-                                        const bool gamma_only_local,
-                                        int istep,
-                                        int iter,
-                                        const DFTU_BASE::OccmatOutputCfg& occmat_cfg);
+                                        const bool gamma_only_local);
 
 template void finish_dftu_lcao<std::complex<double>>(const bool conv_esolver,
                                                       int dft_plus_u,
@@ -114,11 +103,7 @@ template void finish_dftu_lcao<std::complex<double>>(const bool conv_esolver,
                                                       const K_Vectors& kv,
                                                       const double mixing_beta,
                                                       void* hamilt_lcao,
-                                                      const std::string& global_out_dir,
                                                       int nspin,
-                                                      const bool gamma_only_local,
-                                                      int istep,
-                                                      int iter,
-                                                      const DFTU_BASE::OccmatOutputCfg& occmat_cfg);
+                                                      const bool gamma_only_local);
 
 } // namespace ModuleESolver

source/source_lcao/setup_dftu_lcao.h

@@ -1,18 +1,12 @@
 #ifndef SETUP_DFTU_LCAO_H
 #define SETUP_DFTU_LCAO_H
 
-#include <string>
 #include <vector>
 #include "source_cell/unitcell.h"
 #include "source_cell/klist.h"
 
 class LCAO_Orbitals;
 
-namespace DFTU_BASE
-{
-struct OccmatOutputCfg; // defined in source_pw/module_pwdft/dftu_base_io.h
-}
-
 namespace ModuleESolver
 {
 
@@ -54,13 +48,8 @@ void init_dftu_lcao(int dft_plus_u,
  * @param kv K-vectors
  * @param mixing_beta Mixing beta parameter
  * @param hamilt_lcao Hamiltonian LCAO object
- * @param global_out_dir Output directory for occ_mat.txt and occ_matg{#}.txt
  * @param nspin Number of spin channels (1, 2, or 4)
  * @param gamma_only_local Whether only the Gamma point is used for LCAO
- * @param istep Ionic-step index, starting from 0
- * @param iter Electronic-iteration index, starting from 1
- * @param occmat_cfg Configuration for occupation-matrix output; its out_occ_mat
- *        member is the master switch of the occ_mat*.txt files
  */
 template <typename TK>
 void finish_dftu_lcao(const bool conv_esolver,
@@ -71,12 +60,8 @@ void finish_dftu_lcao(const bool conv_esolver,
                        const K_Vectors& kv,
                        const double mixing_beta,
                        void* hamilt_lcao,
-                       const std::string& global_out_dir,
                        int nspin,
-                       const bool gamma_only_local,
-                       int istep,
-                       int iter,
-                       const DFTU_BASE::OccmatOutputCfg& occmat_cfg);
+                       const bool gamma_only_local);
 
 } // namespace ModuleESolver
 

source/source_pw/module_pwdft/CMakeLists.txt

@@ -30,8 +30,8 @@ list(
   objects
   dftu_base.cpp
   dftu_base_io.cpp
-  dftu_base_io_read.cpp
-  dftu_base_io_write.cpp
+  dftu_base_read.cpp
+  dftu_base_write.cpp
   dftu_pw.cpp
   dftu_pw_tools.cpp
   yukawa_screening.cpp

source/source_pw/module_pwdft/setup_dftu_pw.cpp

@@ -1,6 +1,5 @@
 #include "source_pw/module_pwdft/setup_dftu_pw.h"
 #include "source_pw/module_pwdft/dftu_base.h" // mohan add 2025-11-06
-#include "source_pw/module_pwdft/dftu_base_io.h" // mohan add 2025-11-08
 #include "source_pw/module_pwdft/dftu_pw.h"
 #include "source_pw/module_pwdft/dftu_pw_tools.h"
 
@@ -14,20 +13,16 @@ void iter_init_dftu_pw(const int iter,
                         const ModuleBase::matrix& wg,
                         const UnitCell& ucell,
                         Charge_Mixing* p_chgmix,
-                        const std::string& global_out_dir,
-                        const OccmatOutputCfg& occmat_cfg,
+                        const int dft_plus_u,
                         const int* isk,
                         const int nspin,
                         const int kpar)
 {
-    if (!p_chgmix || occmat_cfg.dft_plus_u <= 0)
+    if (!p_chgmix || dft_plus_u <= 0)
     {
         return;
     }
 
-    // Prepare the per-ionic-step file before the first snapshot is appended.
-    DFTU_BASE::prepare_ion_step_output(global_out_dir, istep, iter, occmat_cfg);
-
     const bool first_scf_step = (iter == 1 && istep == 0);
 
     // Build the +U potential from the file-loaded occupation matrix when:

source/source_pw/module_pwdft/setup_dftu_pw.h

@@ -5,25 +5,20 @@
 #include "source_base/matrix.h"
 #include "source_estate/module_charge/chg_mix.h"
 
-#include <string>
-
 struct Input_para;
 class Plus_U_Base; // mohan add 2025-11-06
 
 namespace DFTU_BASE
 {
 
-struct OccmatOutputCfg;
-
 void iter_init_dftu_pw(const int iter,
                        const int istep,
                        Plus_U_Base& dftu, // mohan add 2025-11-06
                        const void* psi,
                        const ModuleBase::matrix& wg,
                        const UnitCell& ucell,
                        Charge_Mixing* p_chgmix,
-                       const std::string& global_out_dir,
-                       const OccmatOutputCfg& occmat_cfg,
+                       const int dft_plus_u,
                        const int* isk,
                        const int nspin,
                        const int kpar);

source/source_pw/module_pwdft/unittests/CMakeLists.txt

@@ -30,8 +30,8 @@ addtest(
   SOURCES test_dftu_base.cpp
           ../dftu_base.cpp
           ../dftu_base_io.cpp
-          ../dftu_base_io_read.cpp
-          ../dftu_base_io_write.cpp
+          ../dftu_base_read.cpp
+          ../dftu_base_write.cpp
           ../dftu_pw_tools.cpp
           ../yukawa_screening.cpp
           ../../../source_estate/occ_matrix.cpp

tests/01_PW/CMakeLists.txt

@@ -1,6 +1,14 @@
 enable_testing()
 
 find_program(BASH bash)
+if(ENABLE_MPI)
+  find_package(Python3 COMPONENTS Interpreter REQUIRED)
+  add_test(
+    NAME 01_PW_dftu_first_step_restart
+    COMMAND ${Python3_EXECUTABLE}
+            ${ABACUS_TEST_DIR}/integrate/validation_tools/test_dftu_first_step_restart.py
+            ${ABACUS_BIN_PATH} ${CMAKE_CURRENT_SOURCE_DIR}/relax_dftu_spin2)
+endif()
 if(ENABLE_ASAN)
   add_test(
     NAME 01_PW_test_with_asan