Skip to content

DFT+U refactor, step 6 - #7879

Open
mohanchen wants to merge 22 commits into
deepmodeling:developfrom
mohanchen:2026-08-29
Open

DFT+U refactor, step 6#7879
mohanchen wants to merge 22 commits into
deepmodeling:developfrom
mohanchen:2026-08-29

Conversation

@mohanchen

@mohanchen mohanchen commented Aug 29, 2026

Copy link
Copy Markdown
Collaborator

DFT+U Slimming: PW occ-path free-function refactor + LCAO occ signature collapse

Overview

This PR continues the DFT+U refactoring by moving the remaining
occupation-matrix member functions of Plus_U_Base/Plus_U out into
namespace free functions, and by collapsing the over-long
cal_occ_mat_k/gamma parameter lists. It further shrinks the class
interface and reduces header coupling.

It consists of four independently compilable commits, each verified by a
full build before committing.

Per-commit breakdown

1. c90489cf2 — slim dftu_lcao.h and turn pot_uterm_HR into free functions

Header slimming (forward declarations)

  • dftu_lcao.h: drop 7 transitive
    includes (klist/unitcell/parallel_orbitals/orb_read/hamilt/hcontainer/ density_matrix), replaced with forward declarations for UnitCell,
    Parallel_Orbitals, LCAO_Orbitals, hamilt::HContainer<T> and
    elecstate::DensityMatrix<TK,TR>.
  • Add direct includes where the removed transitives were actually used:
    dftu_lcao.cpp, dftu_lcao_pots.cpp, dftu_lcao_energy.cpp,
    force_stress_lcao.h, write_hs_r.h, spar_u.cpp.

Free-function refactor

  • Move Plus_U::cal_eff_pot_mat_R_double / cal_eff_pot_mat_R_complex_double
    out of the class into namespace DFTU_LCAO, renamed to pot_uterm_HR_real /
    pot_uterm_HR_complex to match the existing pot_uterm_real/complex family.
  • spar_u.cpp now calls DFTU_LCAO::pot_uterm_HR_real/complex(dftu, ...).

Note: this commit merges the originally planned step 1 (free functions) and
step 2 (header slimming). The two are tightly coupled in dftu_lcao.h
(removing the member declarations + forward-declaration change in the same
file, and spar_u.cpp relies on the header's Parallel_Orbitals forward
declaration), so splitting them would have produced a non-compiling
intermediate state. Merging keeps the history clean and buildable.

2. 8f2d19a83 — sink occupation mixing to source_estate and drop forwarding methods

  • Add free function elecstate::mix_occ_with_save() in
    occ_matrix.h /
    occ_matrix.cpp:
    occ = beta*occ + (1-beta)*occ_save applied to each atom's correlated
    orbital, nspin-aware (nspin=4 mixes the single Pauli block, nspin=1/2 mixes
    both spin channels).
  • LCAO cal_occ_mat_k/gamma: replace the two duplicated hand-written mixing
    loops with a single mix_occ_with_save() call (~60 lines removed).
  • Plus_U_Base: delete the pure-forwarding methods copy_occ_mat /
    zero_occ_mat / set_occ_mat and inline their call sites
    (init_base, cal_occ_pw) to call occmat_.copy_to_save / write_save_to_flat / zero / read_from_flat directly.

3. 816f3020d — turn Plus_U_Base PW occ path into DFTU_BASE free functions

Move the remaining PW occupation-matrix members out of Plus_U_Base into
namespace DFTU_BASE (declared in
dftu_base_tools.h, defined in
dftu_base_occ.cpp):

  • reduce_occ_mat(cell, nspin, kpar, orbital_corr, occmat): the k-pool reduce
    still goes through Parallel_Reduce::reduce_double_allpool (MPI wrapper
    unchanged).
  • compute_eff_pot_and_energycompute_pot_uterm_and_energy(...), renamed to
    match the pot_uterm_* family; writes energy_u by reference.
  • accumulate_occ_one_k<Device>: templated free function, with explicit
    instantiations for DEVICE_CPU (and DEVICE_GPU) moved to the .cpp.
  • sync_occ_to_uom removed; cal_occ_pw now calls occmat_.write_to_flat
    directly.

Plus_U_Base::cal_occ_pw remains as the thin orchestration wrapper.

4. 8045f3691 — collapse cal_occ_mat_k/gamma signatures onto Plus_U&

  • cal_occ_mat_k: 13 params → 9 (add Plus_U& dftu, drop
    nspin/npol/nlocal/ks_solver/iatlnmipol2iwt/orbital_corr/occ_mat/ occ_mat_save/occ_mat_initialized).
  • cal_occ_mat_gamma: 13 params → 7 (add Plus_U& dftu).
  • Both now read all occupation-matrix state from dftu.occmat()
    (mat/mat_save/data/data_save/iatlnmipol2iwt/nspin/npol) and the
    Plus_U_Base accessors (get_orbital_corr_vec / is_occ_mat_initialized / mark_occ_mat_initialized), and reuse occmat().copy_to_save()/zero()
    instead of the duplicated hand-written copy/zero loops.
  • The intermediate cal_occ_mat template forwarders shrink to a single direct
    call.

Note: the MPI_Allreduce(MPI_COMM_WORLD) inside the accumulation loop is
intentionally left unchanged — it has no exact Parallel_Reduce counterpart
and was kept per the "only replace what maps cleanly" rule.

Verification

Each commit was verified before committing with:

This PR reorganizes the DFT+U force and stress implementation in the LCAO module, with two main goals: splitting the monolithic force/stress file into focused, well-documented units, and renaming all LCAO-specific files to explicitly indicate the NAO (numerical atomic orbital) basis set.

First, the original dftu_fs.cpp, which contained the entire force and stress calculation in a single file, is split into three new files. dftu_nao_fs_r.h/cpp provides the unified entry point cal_fs_nao_r(), handling OpenMP parallelization over atom pairs and the final MPI reduction. dftu_nao_for_r.h/cpp contains the single atom-pair force kernel cal_for_IJR_nao_r(). dftu_nao_str_r.h/cpp contains the single atom-pair stress kernel cal_str_IJR_nao_r(). Each file header includes the full mathematical formulas for the Hubbard correction energy, forces, and stress tensor, and explicitly documents that the implementation operates in real space using the real-space density matrix DMR. Because DMR already contains the Brillouin-zone integration, this implementation is unified for both gamma-only and multi-k-point calculations, eliminating the need for separate gamma and k-space branches.

abacus_fixer added 2 commits August 29, 2026 21:40
- Remove dead members with no callers: get_u_target,
  get_orbital_corr_data, deprecated get_pot_uterm_pw(iat),
  element-wise set_occ_mat, and mix_occ_mat (occupation mixing now
  goes through Charge_Mixing::mix_uom + sync_occ_to_uom)
- Merge duplicate occ_mat_initialized accessors: keep the is_*/mark_*
  pair and update dftu_lcao_occ.cpp call sites accordingly
- Replace unitcell.h/charge_mixing.h includes with forward
  declarations in dftu_base.h; add explicit includes to units that
  relied on them transitively (dftu_base.cpp, dftu_cal_occ_pw.cpp,
  dftu_output.cpp) and parallel_comm.h to stodft/hsolver units that
  used BP_WORLD/POOL_WORLD/INT_BGROUP indirectly

Verified: full incremental build of all targets passes with no
errors; repo-wide grep confirms no remaining references to the
removed symbols.
test_hsolver_sdft.cpp and diago_bpcg_test.cpp use BP_WORLD without
including source_base/parallel_comm.h, which broke the MPI build.
@mohanchen mohanchen added DFT+U Issues related to DFT plus U function Refactor Refactor ABACUS codes labels Aug 29, 2026
abacus_fixer and others added 20 commits August 29, 2026 22:39
- Rename dftu_output.{h,cpp} -> dftu_base_io.{h,cpp},
  dftu_tools_pw.{h,cpp} -> dftu_base_tools.{h,cpp},
  dftu_cal_occ_pw.cpp -> dftu_base_occ.cpp (git mv)
- Unify all DFTU-related free functions in module_pwdft under a single
  namespace DFTU_BASE (formerly dftu_io, dftu_pw, and pw::iter_init_dftu_pw)
- Extract read_occup_m and local_occup_bcast from Plus_U_Base as
  DFTU_BASE:: free functions declared in dftu_base_io.h; dependencies
  (occ_mat, orbital_corr, occ_mat_ctrl) are now passed explicitly
- Replace element-wise MPI_Bcast in local_occup_bcast with whole-matrix
  Parallel_Common::bcast_double calls (rule: no direct MPI usage)
- dftu_lcao_test links one extra light source (dftu_base_io.cpp) so that
  the heavy PW implementation dftu_base_occ.cpp stays out of the test
  closure
- Move JacobiRotate/CalculateEigenvalues helpers into an anonymous
  namespace and drop an unused local variable during migration

Verification: full incremental build passes in build/ (abacus_basic_para,
100%); in build_max_para_test/ the targets dftu_lcao_test and
dftu_pw_test build and link successfully (100%); grep confirms no
references remain to dftu_io, dftu_pw::, dftu_output, dftu_tools_pw,
dftu_cal_occ_pw, or the removed Plus_U_Base members.
Introduce YukawaScreening to own the Yukawa screening length, Slater
integrals Fk and derived U_Yukawa/J_Yukawa values. The implementation is
moved verbatim from DFTU_LCAO free functions; no call sites are changed
yet, so dftu_yukawa.* stays in place and behavior is unchanged.

Verification: make -j 30 dftu (build_max_para_test) passed.
Move yukawa_screening.* from module_dftu to module_pwdft so the base
layer can own it without depending on LCAO. Plus_U_Base now constructs a
YukawaScreening when use_yukawa_ is set and exposes it via yukawa(); the
legacy lambda/Fk/U_Yukawa/J_Yukawa members are kept in place until call
sites migrate. The cal_slater_Fk orb dependency stays guarded by __LCAO.

Verification: make -j 30 module_pwdft (build_max_para_test) passed.
Replace the DFTU_LCAO free Yukawa functions with the YukawaScreening
member owned by Plus_U_Base. setup_dftu_lcao now drives
yukawa().cal_slater_UJ and writes U-J back to u_current; the pots,
energy and IO paths read U/J through yukawa(). The yukawa_lambda config
is threaded through init_base into YukawaScreening::init, and
dftu_yukawa.h/.cpp are removed.

Verification: make -j 30 (build_max_para_test) passed;
ctest -R dftu: 4/4 passed (dftu_pw_test, dftu_core_test,
dftu_operator_test, dftu_lcao_test).
All Yukawa state (lambda, Slater Fk integrals, derived U/J) now lives
exclusively in YukawaScreening, owned by Plus_U_Base via yukawa_.
use_yukawa() is derived from the pointer instead of a separate flag,
and the redundant yukawa_lambda member/accessor on Plus_U is removed.

Verified: make -j 30 in build_max_para_test passes;
OMP_NUM_THREADS=1 ctest -R dftu passes (4/4).
New class owning the nested occ[iat][l][n][spin] matrices, their saved
copy for mixing, and the iat->(l,n,m,ipol)->iwt lookup table. Provides
element/matrix access, flat (de)serialization and bulk zero/copy/flat
operations covering everything the legacy Plus_U_Base code paths need.
Not yet referenced by existing code; call sites switch in later steps.

Verified: make -j 30 in build_max_para_test passes.
occmat_ is allocated alongside the legacy occ_mat/occ_mat_save/
iatlnmipol2iwt arrays in init_base, and exposed via occmat(). Existing
read/write paths are unchanged; writers switch to occmat_ next.

Verified: make -j 30 in build_max_para_test passes;
OMP_NUM_THREADS=1 ctest -R dftu passes (4/4).
Switch all internal writers and public read accessors of Plus_U_Base from
the legacy occ_mat/occ_mat_save/iatlnmipol2iwt members to the new
occmat_ (OccupationMatrix) object:

- init_base: read_occup_m/local_occup_bcast now fill occmat_.data()
- copy_occ_mat -> occmat_.copy_to_save + write_save_to_flat
- zero_occ_mat  -> occmat_.zero
- set_occ_mat   -> occmat_.read_from_flat
- get/set_occ_mat_flat -> occmat_.get_flat/set_flat
- dftu_base_occ.cpp: reduce_occ_mat, sync_occ_to_uom,
  compute_eff_pot_and_energy, accumulate_occ_one_k use occmat_.mat()
- header accessors get_occ_mat/get_occ_mat_save/get_occ_mat_data/
  get_occ_mat_save_data/get_iatlnmipol2iwt forward to occmat_
- dftu_lcao_test builds and reads the matrices through occmat()

The legacy members are still allocated but now dead; they are removed in
the follow-up commit. Behavior is unchanged: make -j 30 and
OMP_NUM_THREADS=1 ctest -R dftu (4/4) pass.
All readers and writers now go through occmat_ (OccupationMatrix), so the
legacy Plus_U_Base members occ_mat, occ_mat_save and iatlnmipol2iwt are
deleted together with their allocation block in init_base (the
pot_uterm_pw_index / num_locale bookkeeping is kept).

The IO free functions read_occup_m / local_occup_bcast now take an
OccupationMatrix& instead of the raw nested vector, and write into it
through the set()/mat() interface. write_occup_m / output already used
the public get_occ_mat() accessors and are unchanged.

Behavior is unchanged: make -j 30 and OMP_NUM_THREADS=1 ctest -R dftu
(4/4) pass.
Lesson from the DFT+U OccupationMatrix refactor (B2-B4): migrate a
base-class nested-vector member in three steps so no commit mixes
old-storage writes with new-storage reads.
OccupationMatrix depends only on UnitCell and ModuleBase::matrix, not on
any PW- or LCAO-specific layer. It is an electronic-state data container,
so source_estate is a more natural home than source_pw/module_pwdft.

Verified with make -j 30 in build_max_para_test and
OMP_NUM_THREADS=1 ctest -R dftu (4/4 passed).
Remove 8 accessors on Plus_U_Base that only forwarded to the
OccupationMatrix member; call sites now use dftu.occmat() directly:
  get_iatlnmipol2iwt, get_occ_mat, get_occ_mat_save, get_occ_mat_data,
  get_occ_mat_save_data, get_occ_mat_flat, set_occ_mat_flat
(get_occ_mat_flat/set_occ_mat_flat definitions in dftu_base.cpp removed too).

Kept occmat(), get_occ_mat_ctrl(), and the copy_occ_mat/zero_occ_mat/
set_occ_mat/sync_occ_to_uom wrappers, which still touch the uom_array /
uom_save / pot_uterm_pw_index base-class members.

Verified with make -j 30 in build_max_para_test and
OMP_NUM_THREADS=1 ctest -R dftu (4/4 passed).
cal_slater_Fk dereferences const LCAO_Orbitals*, but the complete type
(orb_read.h) is only included under #ifdef __LCAO while the function body
was not guarded. PW-only CI builds (without __LCAO) failed with
"invalid use of incomplete type 'const class LCAO_Orbitals'".

Guard the body with #ifdef __LCAO and WARNING_QUIT in the non-LCAO stub,
matching the existing codebase pattern (vnl_pw_alpha.cpp, esolver_factory).

Verified: make -j 30 (LCAO build) passed;
c++ -fsyntax-only -std=gnu++11 (no __LCAO) passed;
OMP_NUM_THREADS=1 ctest -R dftu 4/4 passed.
…tions

Header slimming (forward declarations):
- dftu_lcao.h: drop 7 transitive includes (klist/unitcell/parallel_orbitals/
  orb_read/hamilt/hcontainer/density_matrix), replace with forward
  declarations for UnitCell, Parallel_Orbitals, LCAO_Orbitals,
  hamilt::HContainer<T> and elecstate::DensityMatrix<TK,TR>.
- Add direct includes where the removed transitives were actually used:
  dftu_lcao.cpp, dftu_lcao_pots.cpp, dftu_lcao_energy.cpp,
  force_stress_lcao.h, write_hs_r.h, spar_u.cpp.

Free-function refactor:
- Move Plus_U::cal_eff_pot_mat_R_double / cal_eff_pot_mat_R_complex_double
  out of the class into namespace DFTU_LCAO, renamed to pot_uterm_HR_real /
  pot_uterm_HR_complex to match the existing pot_uterm_real/complex family.
- spar_u.cpp now calls DFTU_LCAO::pot_uterm_HR_real/complex(dftu, ...).

Verified: cmake --build build -j4 (exit 0). This build has no test targets.
…arding methods

- Add free function elecstate::mix_occ_with_save() in occ_matrix.{h,cpp}:
  occ = beta*occ + (1-beta)*occ_save over each atom's correlated orbital,
  nspin-aware (nspin=4 single Pauli block, nspin=1/2 both spin channels).
- LCAO cal_occ_mat_k/gamma: replace the two duplicated hand-written mixing
  loops with a single mix_occ_with_save() call (~60 lines removed).
- Plus_U_Base: delete the pure-forwarding methods copy_occ_mat /
  zero_occ_mat / set_occ_mat and inline their call sites
  (init_base in dftu_base.cpp, cal_occ_pw in dftu_base_occ.cpp) to call
  occmat_.copy_to_save / write_save_to_flat / zero / read_from_flat directly.

Verified: cmake --build build -j4 (exit 0). This build has no test targets.
…tions

Move the remaining PW occupation-matrix members out of Plus_U_Base into
namespace DFTU_BASE (declared in dftu_base_tools.h, defined in
dftu_base_occ.cpp):
- reduce_occ_mat(cell, nspin, kpar, orbital_corr, occmat): k-pool reduce
  still goes through Parallel_Reduce::reduce_double_allpool (unchanged MPI
  wrapper).
- compute_eff_pot_and_energy -> compute_pot_uterm_and_energy(...), renamed to
  match the pot_uterm_* family; takes u_current/orbital_corr/
  pot_uterm_pw_index/occmat/pot_uterm_pw and writes energy_u by reference.
- accumulate_occ_one_k<Device>: templated free function, explicit
  instantiations for DEVICE_CPU (and DEVICE_GPU) moved to the .cpp.
- sync_occ_to_uom removed; cal_occ_pw now calls occmat_.write_to_flat
  directly.

Plus_U_Base::cal_occ_pw remains as the thin orchestration wrapper.

Verified: cmake --build build -j4 (exit 0). This build has no test targets.
cal_occ_mat_k: 13 params -> 9 (add Plus_U& dftu, drop nspin/npol/nlocal/
ks_solver/iatlnmipol2iwt/orbital_corr/occ_mat/occ_mat_save/occ_mat_initialized).
cal_occ_mat_gamma: 13 params -> 7 (add Plus_U& dftu).

Both now read all occupation-matrix state from dftu.occmat()
(mat/mat_save/data/data_save/iatlnmipol2iwt/nspin/npol) and the Plus_U_Base
accessors (get_orbital_corr_vec / is_occ_mat_initialized /
mark_occ_mat_initialized), and use occmat().copy_to_save()/zero() instead of
the duplicated hand-written copy/zero loops. The intermediate cal_occ_mat
template forwarders shrink to a single direct call.

MPI_Allreduce(MPI_COMM_WORLD) inside the accumulation is intentionally left
unchanged: it has no exact Parallel_Reduce counterpart and was kept per the
"only replace what maps cleanly" rule.

Verified: cmake --build build -j4 (exit 0). This build has no test targets.
…ename for clarity

Rename and reorganize DFT+U force/stress implementation:

New files (r-space, NAO basis):
- dftu_nao_fs_r.h/cpp    : unified force+stress entry with full formula docs
- dftu_nao_for_r.h/cpp   : single-pair force core (cal_for_IJR_nao_r)
- dftu_nao_str_r.h/cpp   : single-pair stress core (cal_str_IJR_nao_r)

Renamed files:
- dftu_lcao_op.h/cpp          -> dftu_nao_op.h/cpp
- dftu_lcao_op_legacy.cpp     -> dftu_nao_op_legacy.cpp
- dftu_force.h/cpp (legacy k) -> dftu_nao_fs_k.h/cpp
- dftu_folding.h/cpp          -> dftu_nao_folding.h/cpp

Deleted:
- dftu_fs.cpp (split into dftu_nao_fs_r.cpp + dftu_nao_for_r.cpp + dftu_nao_str_r.cpp)

Updated:
- dftu_nao_op.h: expose getters and cal_pot_onsite/transfer_pot_onsite for free functions
- force_stress_lcao.cpp, hamilt_lcao.cpp, dftu_lcao_occ.cpp: update includes
- test/CMakeLists.txt, test/dftu_lcao_test.cpp: update source list and include

Naming convention:
- nao : NAO/LCAO basis set
- _r  : real-space implementation (uses DMR, k-point independent)
- _k  : k-space legacy implementation (deprecated, broken)
- for : force
- str : stress
- fs  : force+stress combined

Co-Authored-By: TRAE <noreply@trae.ai>
…s clarity

Rename all LCAO-specific DFT+U files in module_dftu to use 'nao' instead
of 'lcao', making the basis-set dependency explicit in the filename:

- dftu_lcao.h/cpp          -> dftu_nao.h/cpp
- dftu_lcao_occ.h/cpp      -> dftu_nao_occ.h/cpp
- dftu_lcao_pots.h/cpp     -> dftu_nao_pots.h/cpp
- dftu_lcao_energy.h/cpp   -> dftu_nao_energy.h/cpp
- dftu_lcao_op_legacy.h    -> dftu_nao_op_legacy.h

Update all #include references across 29 files in:
  source_lcao, source_esolver, source_estate, source_io, module_dftu

Update CMakeLists.txt source lists accordingly.

Note: setup_dftu_lcao.cpp/h keep their names (they are not in module_dftu
and represent the DFT+U setup workflow, not a basis-specific implementation).

Co-Authored-By: TRAE <noreply@trae.ai>
@mohanchen
mohanchen requested a review from lanshuyue August 30, 2026 14:50
@mohanchen mohanchen added the The Absolute Zero Reduce the "entropy" of the code to 0 label Aug 30, 2026
@lanshuyue

Copy link
Copy Markdown
Collaborator

There seems to be an existing out-of-bounds write for nspin == 1 in write_save_to_flat().

For nspin == 1, init_base() allocates only one occupation-matrix block in uom_save: pot_index is doubled only when nspin == 2. However, write_save_to_flat() handles nspin == 1 and nspin == 2 together and unconditionally writes both spin channels:

const int half_size = uom_save.size() / 2;

uom_save[index[iat] + mm] = ...[0]...;
uom_save[half_size + index[iat] + mm] = ...[1]...;

For example, for a single correlated d orbital, uom_save.size() == 25 when nspin == 1. The first matrix is written to [0, 24], while the second write goes to [12, 36], so part of the first block is overwritten and [25, 36] is out of bounds. Since operator[] is used, this results in undefined behavior rather than a guaranteed crash.

This does not appear to be newly introduced by this PR: the legacy Plus_U_Base::copy_occ_mat() contains the same logic. However, now that the flat serialization has been extracted into OccupationMatrix, I think it would be good to fix the existing inconsistency here.

In particular, write_to_flat() and read_from_flat() already treat the layout differently: they access the second flat block only for nspin == 2. Should write_save_to_flat() follow the same logic, i.e. always write spin 0 for nspin == 1/2, but write spin 1 only when nspin == 2?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

DFT+U Issues related to DFT plus U function Refactor Refactor ABACUS codes The Absolute Zero Reduce the "entropy" of the code to 0

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants