DFT+U refactor, step 7 (#7909)

* feat(estate): add OccMatMixer skeleton for DFT+U occupation mixing

Introduce OccMatMixer in source_estate to own the flattened
occupation-matrix buffers (uom/uom_save) currently embedded in
Plus_U_Base. Step 1 only adds the class skeleton (init/seed_save/
begin_iter/collect/mix) and wires occ_mixer.cpp into CMake; the
migration of state and call sites follows in subsequent commits.

* refactor(pwdft): attach OccMatMixer to Plus_U_Base (transitional)

Compose Plus_U_Base with a std::unique_ptr<OccMatMixer>, constructed in
init_base only when mixing_dftu != 0 so that "mixer exists" doubles as
the mixing-on flag. The legacy uom_array/uom_save members and the
is_mixing_enabled/enable_mixing accessors are kept for now; they are
removed in a later commit once all call sites migrate.

Invert the Charge_Mixing dependency: OccMatMixer no longer calls
Charge_Mixing::mix_uom itself. It only owns the flat buffers, exposes
them via uom()/uom_save(), and provides write_back() to deserialize the
mixed result. The PW driver (which already links charge_mixing) feeds
the buffers to mix_uom. This keeps occ_mixer.cpp dependent solely on
occ_matrix.h, so the PW and LCAO dftu unit tests link with just
occ_matrix.cpp + occ_mixer.cpp instead of dragging in the planewave/xc
dependency chain (PW_Basis::real2recip, XC_Functional::ked_flag).

Wire occ_mixer.cpp into the dftu_base_test and dftu_lcao_test sources.

* refactor(pwdft): route cal_occ_pw mixing through OccMatMixer

cal_occ_pw now drives the PW mixing orchestration via the composed
OccMatMixer instead of the legacy uom_array/uom_save members:
begin_iter flattens the saved occ into uom_save, collect flattens the
fresh occ into uom, the buffers are fed to Charge_Mixing::mix_uom, and
write_back deserializes the mixed result. The enable condition switches
from is_mixing_enabled() to has_occ_mixer().

The flat save/new writes now run only when a mixer exists; they were
previously executed unconditionally but their results were only read
inside the mixing branch, so behavior is bit-identical. uom_array and
uom_save are now dead members and are removed in a later commit.

OccMatMixer::begin_iter no longer performs copy_to_save itself; that
snapshot stays in cal_occ_pw because it must precede occmat_.zero().

* refactor(estate): switch PW uom allocation to OccMatMixer flat_size

In chgmixing_ks_pw, drop the now-dead dftu.enable_mixing() call (the
mixer is constructed in init_base and cal_occ_pw reads has_occ_mixer()),
gate the uom_mdata allocation on has_occ_mixer() instead of
inp.mixing_dftu, and size it with occ_mixer().flat_size() which equals
the previous get_size_pot_uterm_pw(). The LCAO branch's dead
enable_mixing() call is left for the PARAM-dependency cleanup step.

* refactor(pwdft): drop dead uom_array/uom_save/is_mixing_enabled

Remove the flat occupation buffers uom_array/uom_save from Plus_U_Base
(now owned by OccMatMixer) together with their resize in init_base, and
the is_mixing_enabled() accessor that has no callers left. enable_mixing()
and the mixing_dftu member stay for now because the LCAO path still calls
enable_mixing(); they are removed alongside the PARAM.inp.mixing_dftu
dependency in the cleanup step.

* refactor(dftu): route LCAO plain mixing through OccMatMixer::mix_plain

Add OccMatMixer::mix_plain(occmat, beta), which delegates to
elecstate::mix_occ_with_save on the nested occupation blocks (the flat
buffers are PW/Broyden-only). Replace the two duplicated LCAO k/gamma
call sites with dftu.occ_mixer().mix_plain(dftu.occmat(), mixing_beta)
and gate them on has_occ_mixer() instead of PARAM.inp.mixing_dftu,
removing that cross-layer global read from dftu_nao_occ.cpp.

* refactor(pwdft): remove mixing_dftu member and dead enable_mixing calls

The mutable workflow switch is fully retired: drop the LCAO branch's dead
dftu.enable_mixing() block in chgmixing_ks_lcao (the flag it set was never
read), the enable_mixing() accessor, and the mixing_dftu member of
Plus_U_Base. The init_base mixing_dftu parameter is kept as the sole
decision input for constructing the OccMatMixer; presence of the mixer is
now the only "mixing on" signal.

* add ut for occ_mixer

* fix(dftu): serialize occ_save in write_save_to_flat; fix test double-free

OccupationMatrix::write_save_to_flat is documented to flatten occ_save_
into the uom_save mixing buffer, but it was reading occ_ instead, so the
first mixing step after seed_save/begin_iter used the wrong history
matrix.

In test_occ_mixer, Statistics::~Statistics() owns and delete[]s
iat2it/iat2ia, so backing them with std::vector storage caused a
double-free. Allocate them with new[] instead and drop the extra
backing vectors.

Verified: cmake --build build --target MODULE_ESTATE_occ_mixer &&
OMP_NUM_THREADS=1 ./build/source/source_estate/test/MODULE_ESTATE_occ_mixer
(4/4 tests pass).

* build(Makefile): add occ_mixer.o to OBJS_ELECSTAT

The Makefile build links dftu_base_occ.o/dftu_nao_occ.o against
OccMatMixer, but occ_mixer.cpp was missing from Makefile.Objects,
causing undefined-reference link errors for write_back/mix_plain.
CMake already wires it via source_estate/CMakeLists.txt.

* fix(base): update test call sites to world_comm_domain()

PR #7888 tests were added before PR #7900 renamed
ModuleBase::world_communication_domain() to
ModuleBase::world_comm_domain(); update the three affected
test files so the parallel test targets build again.

Verified: cmake --build build --target
MODULE_BASE_parallel_domain_grid MODULE_BASE_para_gemm
MODULE_BASE_parallel_device all succeed.

---------

Co-authored-by: abacus_fixer <mohanchen@pku.eud.cn>

Author: mohanchen

source/Makefile.Objects

@@ -283,7 +283,8 @@ OBJS_ELECSTAT=elecstate.o\
     read_orb.o\
     setup_estate_pw.o\
     update_pot.o\
-    occ_matrix.o
+    occ_matrix.o\
+    occ_mixer.o
 
 OBJS_ELECSTAT_LCAO=elecstate_lcao.o\
       init_dm.o\

source/source_base/test_parallel/parallel_device_test.cpp

@@ -124,7 +124,7 @@ TEST(ParallelDevice, CoversGpuStagingWithoutAccelerator)
 
 TEST(ParallelDevice, CoversMpiTypeOverloads)
 {
-    const ModuleBase::CommunicationDomain domain = ModuleBase::world_communication_domain();
+    const ModuleBase::CommunicationDomain domain = ModuleBase::world_comm_domain();
     exercise_mpi_wrappers<float>(domain);
     exercise_mpi_wrappers<double>(domain);
     exercise_mpi_wrappers<std::complex<float>>(domain);

source/source_base/test_parallel/parallel_domain_grid_test.cpp

@@ -14,7 +14,7 @@ TEST(CommunicationDomainTest, ReportsDefaultAndWorldDomains)
     EXPECT_EQ(local_domain.rank(), 0);
     EXPECT_EQ(local_domain.communicator(), MPI_COMM_NULL);
 
-    const ModuleBase::CommunicationDomain world_domain = ModuleBase::world_communication_domain();
+    const ModuleBase::CommunicationDomain world_domain = ModuleBase::world_comm_domain();
     MPICommGroup world_group(world_domain.communicator());
     EXPECT_EQ(world_domain.communicator(), MPI_COMM_WORLD);
     EXPECT_GE(world_domain.rank(), 0);
@@ -28,7 +28,7 @@ TEST(CommunicationDomainTest, ReportsDefaultAndWorldDomains)
 
 TEST(MPICommGroupTest, DividesWorldIntoEvenGroups)
 {
-    const ModuleBase::CommunicationDomain world_domain = ModuleBase::world_communication_domain();
+    const ModuleBase::CommunicationDomain world_domain = ModuleBase::world_comm_domain();
     MPICommGroup group(MPI_COMM_WORLD);
     EXPECT_EQ(group.grank, world_domain.rank());
 
@@ -46,7 +46,7 @@ TEST(MPICommGroupTest, DividesWorldIntoEvenGroups)
 
 TEST(ParallelGridTest, BroadcastsAndReducesDistributedGrid)
 {
-    const ModuleBase::CommunicationDomain world_domain = ModuleBase::world_communication_domain();
+    const ModuleBase::CommunicationDomain world_domain = ModuleBase::world_comm_domain();
     MPICommGroup world_group(world_domain.communicator());
     const int nx = 2;
     const int ny = 1;

source/source_base/test_parallel/test_para_gemm.cpp

@@ -74,7 +74,7 @@ void expect_near_value(const T& actual, const T& expected)
 template <typename T>
 void test_additional_type_paths()
 {
-    const ModuleBase::CommunicationDomain domain = ModuleBase::world_communication_domain();
+    const ModuleBase::CommunicationDomain domain = ModuleBase::world_comm_domain();
     MPI_Comm world = domain.communicator();
     const int rank = domain.rank();
     MPICommGroup world_group(world);

source/source_estate/CMakeLists.txt

@@ -46,6 +46,7 @@ list(APPEND objects
     fp_energy.cpp
     occupy.cpp
     occ_matrix.cpp
+    occ_mixer.cpp
     param_update.cpp
     setup_estate_pw.cpp
     update_pot.cpp

source/source_estate/module_charge/chgmixing.cpp

@@ -130,12 +130,10 @@ void module_charge::chgmixing_ks_pw(const int iter, // scf iteration number
     {
         p_chgmix->init_mixing();
         p_chgmix->mixing_restart_step = inp.scf_nmax + 1;
-        if (inp.dft_plus_u && inp.mixing_dftu)
+        if (inp.dft_plus_u && dftu.has_occ_mixer())
         {
-            // enable mixing_dftu for DFT+U occupation mixing
-            dftu.enable_mixing();
-            // allocate memory for uom_mdata
-            p_chgmix->allocate_mixing_uom(dftu.get_size_pot_uterm_pw());
+            // allocate memory for uom_mdata sized to the flat occupation buffer
+            p_chgmix->allocate_mixing_uom(dftu.occ_mixer().flat_size());
         }
     }
 
@@ -193,11 +191,6 @@ void module_charge::chgmixing_ks_lcao(const int iter, // scf iteration number
         p_chgmix->mix_reset(); // init mixing
         p_chgmix->mixing_restart_step = inp.scf_nmax + 1;
         p_chgmix->mixing_restart_count = 0;
-        // enable mixing_dftu for DFT+U occupation mixing
-        if (inp.dft_plus_u && inp.mixing_dftu)
-        {
-            dftu.enable_mixing();
-        }
         // this output will be removed once the feeature is stable
         if (dftu.get_uramping() > 0.01)
         {

source/source_estate/occ_matrix.cpp

@@ -280,27 +280,27 @@ void OccupationMatrix::write_save_to_flat(const UnitCell& cell,
         for (int I = 0; I < cell.atoms[T].na; I++)
         {
             const int iat = cell.itia2iat(T, I);
-            const int size = occ_[iat][target_l][0][0].nr * occ_[iat][target_l][0][0].nc;
+            const int size = occ_save_[iat][target_l][0][0].nr * occ_save_[iat][target_l][0][0].nc;
 
             if (nspin_ == 4)
             {
                 for (int mm = 0; mm < size; mm++)
                 {
-                    uom_save[index[iat] + mm] = occ_[iat][target_l][0][0].c[mm];
+                    uom_save[index[iat] + mm] = occ_save_[iat][target_l][0][0].c[mm];
                 }
             }
             else if (nspin_ == 1 || nspin_ == 2)
             {
                 for (int mm = 0; mm < size; mm++)
                 {
-                    uom_save[index[iat] + mm] = occ_[iat][target_l][0][0].c[mm];
+                    uom_save[index[iat] + mm] = occ_save_[iat][target_l][0][0].c[mm];
                 }
                 if (nspin_ == 2)
                 {
                     const int half_size = uom_save.size() / 2;
                     for (int mm = 0; mm < size; mm++)
                     {
-                        uom_save[half_size + index[iat] + mm] = occ_[iat][target_l][0][1].c[mm];
+                        uom_save[half_size + index[iat] + mm] = occ_save_[iat][target_l][0][1].c[mm];
                     }
                 }
             }

source/source_estate/occ_mixer.cpp

@@ -0,0 +1,49 @@
+#include "source_estate/occ_mixer.h"
+
+void OccMatMixer::init(const UnitCell* cell,
+                       const std::vector<int>* orbital_corr,
+                       const std::vector<int>* flat_index,
+                       const int nspin,
+                       const int total_size)
+{
+    this->cell_ = cell;
+    this->orbital_corr_ = orbital_corr;
+    this->index_ = flat_index;
+    this->nspin_ = nspin;
+    this->uom_.resize(total_size, 0.0);
+    this->uom_save_.resize(total_size, 0.0);
+}
+
+void OccMatMixer::seed_save(const OccupationMatrix& occmat)
+{
+    occmat.write_save_to_flat(*this->cell_, *this->orbital_corr_,
+                              *this->index_, this->uom_save_);
+}
+
+void OccMatMixer::begin_iter(OccupationMatrix& occmat)
+{
+    // the caller has already snapshotted occ into occ_save via
+    // OccupationMatrix::copy_to_save; here we only flatten that snapshot
+    // into uom_save_ for the mixing history.
+    occmat.write_save_to_flat(*this->cell_, *this->orbital_corr_,
+                              *this->index_, this->uom_save_);
+}
+
+void OccMatMixer::collect(const OccupationMatrix& occmat)
+{
+    occmat.write_to_flat(*this->cell_, *this->orbital_corr_,
+                         *this->index_, this->uom_);
+}
+
+void OccMatMixer::write_back(OccupationMatrix& occmat)
+{
+    occmat.read_from_flat(*this->cell_, *this->orbital_corr_,
+                          *this->index_, this->uom_);
+}
+
+void OccMatMixer::mix_plain(OccupationMatrix& occmat, const double beta)
+{
+    elecstate::mix_occ_with_save(occmat.data(), occmat.data_save(),
+                                 *this->cell_, *this->orbital_corr_,
+                                 this->nspin_, beta);
+}

source/source_estate/occ_mixer.h

@@ -0,0 +1,106 @@
+#ifndef OCC_MIXER_H
+#define OCC_MIXER_H
+
+#include "source_estate/occ_matrix.h"
+
+#include <vector>
+
+class UnitCell;
+
+/**
+ * @brief Mixing of the DFT+U on-site occupation matrix.
+ *
+ * Owns the flattened occupation-matrix buffers used by the charge-mixing
+ * machinery (PW path) and the plain linear mixing kernel (LCAO path).
+ *
+ * The flat layout reuses the pot_uterm_pw_index offset table: for nspin=2
+ * the buffer is split into [all_up | all_dn] halves; for nspin=1/4 a single
+ * block per atom is used. Serialization to/from the nested OccupationMatrix
+ * is delegated to OccupationMatrix::{write_to_flat, read_from_flat,
+ * write_save_to_flat}.
+ *
+ * An OccMatMixer instance exists only when mixing is enabled, so its
+ * presence doubles as the "mixing on" flag (no mutable workflow switch).
+ *
+ * This class deliberately does NOT call Charge_Mixing itself; it only owns
+ * the flat buffers and exposes them via uom()/uom_save(). The caller (the
+ * PW driver, which already links charge_mixing) feeds these buffers to
+ * Charge_Mixing::mix_uom and then calls write_back(). Keeping Charge_Mixing
+ * out of this translation unit avoids dragging the planewave/xc dependency
+ * chain into lightweight unit tests.
+ */
+class OccMatMixer
+{
+  public:
+    OccMatMixer() = default;
+    ~OccMatMixer() = default;
+
+    /**
+     * @brief Allocate the flat buffers and bind the layout table.
+     * @param cell          unit cell (borrowed, must outlive this object)
+     * @param orbital_corr  per-type correlated-l table (borrowed)
+     * @param flat_index    per-atom offset table, i.e. pot_uterm_pw_index (borrowed)
+     * @param nspin         spin channels (1, 2 or 4)
+     * @param total_size    total flat-buffer size (== pot_uterm_pw.size())
+     */
+    void init(const UnitCell* cell,
+              const std::vector<int>* orbital_corr,
+              const std::vector<int>* flat_index,
+              int nspin,
+              int total_size);
+
+    /**
+     * @brief Seed uom_save from an occupation matrix loaded from file.
+     *
+     * Used when occ_mat_ctrl != 0 (restart from dm_onsite_ini.txt) so that
+     * the first mixing step has a meaningful "previous" matrix.
+     */
+    void seed_save(const OccupationMatrix& occmat);
+
+    /**
+     * @brief Begin an SCF iteration: flatten the saved occ into uom_save_.
+     *
+     * The caller must already have snapshotted occ into occ_save via
+     * OccupationMatrix::copy_to_save; this only flattens that snapshot into
+     * uom_save_ for the mixing history.
+     */
+    void begin_iter(OccupationMatrix& occmat);
+
+    /**
+     * @brief Flatten the freshly-computed occupation matrix into uom_.
+     */
+    void collect(const OccupationMatrix& occmat);
+
+    /**
+     * @brief Write the (already mixed) uom_ buffer back into the occupation
+     *        matrix. Called after the caller has run Charge_Mixing::mix_uom.
+     */
+    void write_back(OccupationMatrix& occmat);
+
+    /**
+     * @brief Plain linear mixing for the nested-matrix (LCAO) path:
+     *        occ = beta * occ + (1 - beta) * occ_save.
+     *
+     * Operates directly on the nested OccupationMatrix blocks; the flat
+     * buffers are not used. Replaces the duplicated LCAO k/gamma call sites.
+     */
+    void mix_plain(OccupationMatrix& occmat, double beta);
+
+    /// Total flat-buffer size (== Charge_Mixing::allocate_mixing_uom argument).
+    int flat_size() const { return static_cast<int>(uom_.size()); }
+
+    /// Mutable access to the new / mixed flat buffer (fed to mix_uom).
+    std::vector<double>& uom() { return uom_; }
+    /// Mutable access to the previous flat buffer (fed to mix_uom).
+    std::vector<double>& uom_save() { return uom_save_; }
+
+  private:
+    std::vector<double> uom_;      ///< new / mixed flat occupation matrix
+    std::vector<double> uom_save_; ///< previous flat occupation matrix
+    const std::vector<int>* index_ = nullptr;        ///< borrowed pot_uterm_pw_index
+    const UnitCell* cell_ = nullptr;                 ///< borrowed unit cell
+    const std::vector<int>* orbital_corr_ = nullptr; ///< borrowed correlated-l table
+    int nspin_ = 0;
+};
+
+#endif

source/source_estate/test/CMakeLists.txt

@@ -85,6 +85,12 @@ AddTest(
   SOURCES potentials_base_test.cpp
 )
 
+AddTest(
+  TARGET MODULE_ESTATE_occ_mixer
+  LIBS parameter base device cell_info
+  SOURCES test_occ_mixer.cpp ../occ_mixer.cpp ../occ_matrix.cpp
+)
+
 AddTest(
   TARGET MODULE_ESTATE_potentials_new
   LIBS parameter  base device planewave_serial