Remove template of ESolver_KS_pw (#7024)

* small format changes

* refactor(esolver): extract charge density symmetrization to Symmetry_rho::symmetrize_rho

- Add static method symmetrize_rho() in Symmetry_rho class
- Replace 7 duplicate code blocks with single function call
- Simplify code from 35 lines to 7 lines (80% reduction)
- Improve code readability and maintainability

Modified files:
- source_estate/module_charge/symmetry_rho.h: add static method declaration
- source_estate/module_charge/symmetry_rho.cpp: implement static method
- source_esolver/esolver_ks_lcao.cpp: 2 calls updated
- source_esolver/esolver_ks_pw.cpp: 1 call updated
- source_esolver/esolver_ks_lcao_tddft.cpp: 1 call updated
- source_esolver/esolver_ks_lcaopw.cpp: 1 call updated
- source_esolver/esolver_of.cpp: 1 call updated
- source_esolver/esolver_sdft_pw.cpp: 1 call updated

This refactoring follows the ESolver cleanup principle:
keep ESolver focused on high-level workflow control.

* refactor(esolver): extract DeltaSpin lambda loop to deltaspin_lcao module

- Create new files deltaspin_lcao.h/cpp in module_deltaspin
- Extract DeltaSpin lambda loop logic from ESolver_KS_LCAO
- Simplify code from 18 lines to 1 line in hamilt2rho_single
- Separate LCAO and PW implementations for DeltaSpin

Modified files:
- source_esolver/esolver_ks_lcao.cpp: replace inline code with function call
- source_lcao/module_deltaspin/CMakeLists.txt: add new source file

New files:
- source_lcao/module_deltaspin/deltaspin_lcao.h: function declaration
- source_lcao/module_deltaspin/deltaspin_lcao.cpp: function implementation

This refactoring follows the ESolver cleanup principle:
keep ESolver focused on high-level workflow control.

* refactor(esolver): complete DeltaSpin refactoring in LCAO

- Add init_deltaspin_lcao() function for DeltaSpin initialization
- Add cal_mi_lcao_wrapper() function for magnetic moment calculation
- Refactor all DeltaSpin-related code in esolver_ks_lcao.cpp
- Simplify code from 29 lines to 3 lines (90% reduction)

Modified files:
- source_esolver/esolver_ks_lcao.cpp: replace 3 code blocks with function calls
- source_lcao/module_deltaspin/deltaspin_lcao.h: add 2 new function declarations
- source_lcao/module_deltaspin/deltaspin_lcao.cpp: implement 2 new functions

This completes the DeltaSpin refactoring for LCAO method:
1. init_deltaspin_lcao() - initialize DeltaSpin calculation
2. cal_mi_lcao_wrapper() - calculate magnetic moments
3. run_deltaspin_lambda_loop_lcao() - run lambda loop optimization

All functions follow the ESolver cleanup principle:
keep ESolver focused on high-level workflow control.

* refactor(esolver): extract DFT+U code to dftu_lcao module

- Create new files dftu_lcao.h/cpp in source_lcao directory
- Add init_dftu_lcao() function for DFT+U initialization
- Add finish_dftu_lcao() function for DFT+U finalization
- Simplify code from 32 lines to 2 lines in esolver_ks_lcao.cpp
- Remove conditional checks from ESolver, move them to functions

Modified files:
- source_esolver/esolver_ks_lcao.cpp: replace 2 code blocks with function calls
- source_lcao/CMakeLists.txt: add new source file

New files:
- source_lcao/dftu_lcao.h: function declarations
- source_lcao/dftu_lcao.cpp: function implementations

This refactoring prepares for unifying old and new DFT+U implementations:
- Old DFT+U: source_lcao/module_dftu/
- New DFT+U: source_lcao/module_operator_lcao/op_dftu_lcao.cpp

All functions follow ESolver cleanup principle:
keep ESolver focused on high-level workflow control.

* refactor(esolver): extract diagonalization parameters setup to hsolver module

- Create new files diago_params.h/cpp in source_hsolver directory
- Add setup_diago_params_pw() function for PW diagonalization parameters
- Simplify code from 11 lines to 1 line in esolver_ks_pw.cpp
- Encapsulate diagonalization parameter setup logic

Modified files:
- source_esolver/esolver_ks_pw.cpp: replace inline code with function call
- source_hsolver/CMakeLists.txt: add new source file

New files:
- source_hsolver/diago_params.h: function declaration
- source_hsolver/diago_params.cpp: function implementation

This refactoring follows ESolver cleanup principle:
keep ESolver focused on high-level workflow control.

* fix(deltaspin): add sc_mag_switch check in cal_mi_lcao_wrapper

- Add Input_para parameter to cal_mi_lcao_wrapper function
- Add sc_mag_switch check to avoid calling cal_mi_lcao when DeltaSpin is disabled
- Fix 'atomCounts is not set' error in non-DeltaSpin calculations
- Update function call in esolver_ks_lcao.cpp

This fix resolves the CI/CD failure caused by commit 2a520e3.
The root cause was that cal_mi_lcao_wrapper was called without checking
sc_mag_switch, leading to uninitialized atomCounts error.

Modified files:
- source_esolver/esolver_ks_lcao.cpp: update function call
- source_lcao/module_deltaspin/deltaspin_lcao.h: add parameter
- source_lcao/module_deltaspin/deltaspin_lcao.cpp: add check

This follows the refactoring principle: preserve original condition checks
when extracting code to wrapper functions.

* fix(deltaspin): add #ifdef __LCAO for conditional compilation

- Add #ifdef __LCAO conditional compilation in init_deltaspin_lcao and cal_mi_lcao_wrapper
- Fix parameter order in init_sc call for LCAO and non-LCAO builds
- Fix undefined reference to cal_mi_lcao in non-LCAO build

This fix resolves CI/CD compilation errors in both build_5pt (with __LCAO) and build_1p (without __LCAO) environments.

The The root cause was 1. init_sc has different parameter order in LCAO vs non-LCAO builds
   - LCAO: psi, dm, pelec
   - non-LCAO: psi, pelec
  2. cal_mi_lcao is only defined in LCAO build

Modified files:
- source_hsolver/diago_params.h: add setup_diago_params_sdft declaration
- source_lcao/module_deltaspin/deltaspin_lcao.cpp: add conditional compilation

This follows the refactoring principle: handle conditional compilation properly
when code has different implementations for different build configurations.

* refactor(esolver): extract SDFT diagonalization parameters setup

- Add setup_diago_params_sdft() function for SDFT diagonalization parameters
- Simplify code from 11 lines to 1 line in esolver_sdft_pw.cpp
- Encapsulate diagonalization parameter setup logic for SDFT

Modified files:
- source_esolver/esolver_sdft_pw.cpp: replace inline code with function call
- source_hsolver/diago_params.cpp: add setup_diago_params_sdft implementation

This refactoring follows ESolver cleanup principle:
keep ESolver focused on high-level workflow control.

Note: SDFT has different parameter setup logic compared to PW:
- Different need_subspace condition
- No SCF_ITER setting
- Always set PW_DIAG_NMAX (no nscf check)

* refactor(hamilt): introduce HamiltBase non-template base class

- Create HamiltBase as a non-template base class for Hamilt<T, Device>
- Modify Hamilt<T, Device> to inherit from HamiltBase
- Change ESolver_KS::p_hamilt type from Hamilt<T, Device>* to HamiltBase*
- Add static_cast where needed when passing p_hamilt to functions expecting Hamilt<T, Device>*

This is the first step towards removing template parameters from ESolver.

Modified files:
- source/source_esolver/esolver_ks.h
- source/source_esolver/esolver_ks_lcaopw.cpp
- source/source_esolver/esolver_ks_pw.cpp
- source/source_esolver/esolver_sdft_pw.cpp
- source/source_hamilt/hamilt.h

New files:
- source/source_hamilt/hamilt_base.h

* refactor(esolver): add static_cast for p_hamilt in esolver files

- Add static_cast<hamilt::Hamilt<T>*> when passing p_hamilt to functions
  expecting Hamilt<T, Device>* type
- Split long cast statements into multiple lines for better readability
- Files modified:
  - esolver_ks_pw.cpp: setup_pot, stp.init calls
  - esolver_ks_lcao.cpp: init_chg_hr, hsolver_lcao_obj.solve calls
  - esolver_ks_lcao_tddft.cpp: solve_psi, cal_edm_tddft, matrix calls
  - esolver_gets.cpp: ops access, output_SR call

This follows the HamiltBase refactoring strategy where p_hamilt is stored
as HamiltBase* and cast to Hamilt<T, Device>* when needed.

* refactor(esolver): remove psi member from ESolver_KS base class

Move psi::Psi<T>* psi from ESolver_KS base class to derived classes
to eliminate template parameter dependency and improve code organization.

Changes:
1. ESolver_KS base class:
   - Remove psi::Psi<T>* psi member variable
   - Remove Setup_Psi<T>::deallocate_psi() call in destructor
   - Remove unnecessary includes: psi.h and setup_psi.h

2. ESolver_KS_LCAO:
   - Add psi::Psi<TK>* psi member variable
   - Add Setup_Psi<TK>::deallocate_psi() in destructor
   - Add include: setup_psi.h

3. ESolver_KS_LCAO_TDDFT:
   - Improve psi_laststep deallocation with nullptr check
   - psi member inherited from ESolver_KS_LCAO

4. ESolver_KS_PW:
   - Use stp.psi_cpu directly instead of base class psi
   - Remove unnecessary memory allocation in after_scf()

5. pw_others.cpp (BUG FIX):
   - Fix gen_bessel: use *(this->stp.psi_cpu) instead of this->psi[0]
   - Previous code accessed uninitialized base class psi (nullptr)
   - This was a latent bug that could cause crashes

Benefits:
- Eliminates template parameter T dependency in ESolver_KS base class
- Clearer memory management: each derived class manages its own psi
- Reduces compilation dependencies
- Fixes potential memory access bug in pw_others.cpp

Tested: Compiled successfully in build_5pt and build_1p

* refactor(esolver): remove template parameters from ESolver_KS base class

This is a major milestone in ESolver refactoring!

ESolver_KS no longer needs template parameters because:
- All member variables are non-template types
- All member functions do not use T or Device parameters
- Template parameters were only needed for derived classes

Changes:
1. ESolver_KS base class:
   - Remove template <typename T, typename Device> declaration
   - Remove all template declarations from member functions
   - Remove template instantiation code at end of file
   - Fix Tab indentation to spaces for better readability

2. Derived classes:
   - ESolver_KS_PW: public ESolver_KS (was ESolver_KS<T, Device>)
   - ESolver_KS_LCAO: public ESolver_KS (was ESolver_KS<TK>)
   - ESolver_GetS: public ESolver_KS (was ESolver_KS<std::complex<double>>)
   - Update base class calls: ESolver_KS:: (was ESolver_KS<T, Device>::)

Code reduction:
- esolver_ks.h: 78 -> 77 lines (-1 line)
- esolver_ks.cpp: 346 -> 317 lines (-29 lines)
- Total ESolver code: 424 -> 394 lines (-30 lines)
- Overall: 8 files changed, 50 insertions(+), 80 deletions(-), net -30 lines

Benefits:
- Simpler base class without template complexity
- Faster compilation (no template instantiation needed)
- Clearer inheritance hierarchy
- Easier to extract common code in future refactoring
- Sets foundation for further ESolver template removal

Tested: Compiled successfully in build_5pt

* refactor(device): remove explicit template parameter from get_device_type calls

- Move get_device_type implementation to header file using std::is_same
- Add DEVICE_DSP support
- Remove template specialization declarations and definitions
- Update all call sites to use automatic template parameter deduction
- The compiler now deduces Device type from the ctx parameter

* refactor(esolver): remove device member variable from ESolver_KS_PW

- Modify copy_d2h to accept ctx parameter and call get_device_type internally
- Remove device parameter from ctrl_scf_pw function
- Remove device member variable from ESolver_KS_PW class
- Simplify function interfaces by using automatic template deduction

* style(esolver): explicitly initialize ctx to nullptr in constructor

* feat(device): add runtime device type support to DeviceContext

- Add device_type_ member variable to DeviceContext class
- Add set_device_type() and get_device_type() methods
- Add is_cpu(), is_gpu(), is_dsp() convenience methods
- Add get_device_type(const DeviceContext*) overload for runtime device type query
- Maintain backward compatibility with existing template-based get_device_type

* feat(device): add runtime device context overloads for gradual migration

- Add copy_d2h(const DeviceContext*) overload to Setup_Psi_pw
- Add ctrl_scf_pw(..., const DeviceContext*, ...) overload
- Add ctrl_runner_pw(..., const DeviceContext*, ...) overload
- Keep original functions for backward compatibility
- Replace tabs with spaces in modified files

* refactor(esolver): remove ctx member variable from ESolver_KS_PW

- Remove Device* ctx member variable from ESolver_KS_PW class
- Remove ctx parameter from ctrl_scf_pw and ctrl_runner_pw functions
- Add local ctx variable inside ctrl_scf_pw and ctrl_runner_pw functions
- Update all template instantiations to match new function signatures

This refactoring simplifies the code by moving the ctx variable from
a class member to a local variable within the functions that need it.
The ctx variable is only used for template parameter deduction in
copy_d2h and get_pchg_pw/get_wf_pw functions, so it doesn't need to
be stored as a member variable.

* refactor(psi): add runtime type information to Setup_Psi_pw

- Add runtime type information (device_type_ and precision_type_) to Setup_Psi_pw
- Add accessor functions for basic information (get_nbands, get_nk, get_nbasis, size)
- Add accessor functions for runtime type information
- Add get_psi_t() function for backward compatibility

This is the first step of a gradual refactoring to remove template parameters
from Setup_Psi_pw in the future. The current changes are backward compatible
and do not affect existing functionality.

* refactor(esolver): use get_psi_t() accessor instead of direct psi_t access

- Replace all direct access to stp.psi_t with stp.get_psi_t()
- Replace stp.psi_t->get_nbands() with stp.get_nbands()
- This is the second step of gradual refactoring to prepare for removing template parameters

Modified files:
- source/source_esolver/esolver_ks_pw.cpp
- source/source_esolver/esolver_sdft_pw.cpp
- source/source_esolver/esolver_ks_lcaopw.cpp
- source/source_io/module_ctrl/ctrl_output_pw.cpp

* refactor(psi): change psi_t from template pointer to void*

- Change psi_t from psi::Psi<T, Device>* to void*
- Add static_cast in get_psi_t() function for type conversion
- Update all functions that use psi_t to use get_psi_t() or static_cast
- This is the third step of gradual refactoring to remove template parameters

Modified functions:
- before_runner: use if-else instead of ternary operator for void* assignment
- update_psi_d: use get_psi_t() to access psi_t
- init: use get_psi_t() to access psi_t
- copy_d2h: use get_psi_t() to access psi_t
- clean: use get_psi_t() to delete psi_t

* style: replace Chinese comments with English in setup_psi_pw.h

- Replace '原来的模板版本' with 'Original template version'
- Replace '使用 void* 存储指针，运行时类型信息记录实际类型' with 'Use void* to store pointer, runtime type information records actual type'
- Follow ABACUS code style guidelines for English-only comments

* refactor(psi): change psi_d from template pointer to void*

- Change psi_d from psi::Psi<std::complex<double>, Device>* to void*
- Add get_psi_d() accessor function for type conversion
- Update all functions that use psi_d to use get_psi_d()
- This is part of step 1 in phase 4 of gradual refactoring

Modified files:
- source/source_psi/setup_psi_pw.h
- source/source_psi/setup_psi_pw.cpp
- source/source_esolver/esolver_ks_pw.cpp
- source/source_io/module_ctrl/ctrl_output_pw.cpp

* refactor(psi): introduce PSIPrepareBase base class for template removal

This is the first step towards removing template parameters from Setup_Psi_pw.

Changes:
1. Create PSIPrepareBase base class
   - Non-template base class for PSIPrepare<T, Device>
   - Similar approach to HamiltBase for Hamilt<T, Device>

2. Modify PSIPrepare to inherit from PSIPrepareBase
   - Add #include "source_psi/psi_prepare_base.h"
   - Change class declaration to inherit from PSIPrepareBase

3. Update Setup_Psi_pw to use PSIPrepareBase*
   - Change p_psi_init from PSIPrepare<T, Device>* to PSIPrepareBase*
   - Add static_cast when calling PSIPrepare methods

4. Update all PSIPrepare usage in ESolver files
   - esolver_ks_pw.cpp: add static_cast before prepare_init call
   - esolver_ks_lcaopw.cpp: add static_cast before method calls

Modified files:
- source/source_psi/psi_prepare_base.h (new)
- source/source_psi/psi_prepare.h
- source/source_psi/setup_psi_pw.h
- source/source_psi/setup_psi_pw.cpp
- source/source_esolver/esolver_ks_pw.cpp
- source/source_esolver/esolver_ks_lcaopw.cpp

Benefits:
- Eliminates p_psi_init template dependency from Setup_Psi_pw
- Paves the way for removing template parameters from Setup_Psi_pw
- Maintains type safety through static_cast
- Follows the same pattern as HamiltBase refactoring

Tested: Compiled successfully in build_5pt and build_1p

* refactor(psi): change init() parameter from template to HamiltBase*

This is the second step towards removing template parameters from Setup_Psi_pw.

Changes:
1. Modify init() function signature
   - Change parameter from hamilt::Hamilt<T, Device>* to hamilt::HamiltBase*
   - Eliminates template dependency in function signature

2. Update init() implementation
   - Add static_cast<hamilt::Hamilt<T, Device>*> inside function
   - Maintain type safety through explicit cast

3. Update call site in esolver_ks_pw.cpp
   - Remove static_cast from call site
   - Directly pass p_hamilt (which is already HamiltBase*)

Modified files:
- source/source_psi/setup_psi_pw.h
- source/source_psi/setup_psi_pw.cpp
- source/source_esolver/esolver_ks_pw.cpp

Benefits:
- init() function no longer depends on template parameters in signature
- Simplifies call sites (no cast needed)
- Follows the same pattern as p_hamilt storage in ESolver_KS
- One step closer to removing template parameters from Setup_Psi_pw

Tested: Compiled successfully in build_5pt and build_1p

* refactor(psi): remove template version of copy_d2h function

This is the third step towards removing template parameters from Setup_Psi_pw.

Changes:
1. Remove template version copy_d2h(const Device* ctx)
   - Delete the template version from setup_psi_pw.h
   - Delete the implementation from setup_psi_pw.cpp
   - Keep only the runtime version copy_d2h(const base_device::DeviceContext* ctx)

2. Update call site in ctrl_output_pw.cpp
   - Create Device* ctx = nullptr for template parameter deduction
   - Use DeviceContext::instance() for runtime device context
   - Call copy_d2h with DeviceContext* pointer

Modified files:
- source/source_psi/setup_psi_pw.h
- source/source_psi/setup_psi_pw.cpp
- source/source_io/module_ctrl/ctrl_output_pw.cpp

Benefits:
- Eliminates copy_d2h function's template dependency
- All member functions now use runtime device context
- One step closer to removing template parameters from Setup_Psi_pw
- Maintains backward compatibility with existing code

Technical details:
- get_pchg_pw and get_wf_pw still need Device* ctx for template deduction
- DeviceContext is used for actual device type information
- This follows the gradual migration pattern used in ESolver refactoring

Tested: Compiled successfully in build_5pt and build_1p

* refactor(psi): remove castmem_2d_d2h_op template type alias dependency

This is the fourth step towards removing template parameters from Setup_Psi_pw.

Changes:
1. Remove castmem_2d_d2h_op type alias from setup_psi_pw.h
   - The type alias depended on template parameters T and Device
   - Replaced with overloaded member functions

2. Add castmem_d2h_impl() overloaded functions
   - One overload for std::complex<double> source
   - One overload for std::complex<float> source
   - Each uses the appropriate cast_memory_op internally

3. Update copy_d2h() to use the new overloaded functions
   - Calls castmem_d2h_impl() instead of castmem_2d_d2h_op()
   - Compiler selects the correct overload based on T type

Modified files:
- source/source_psi/setup_psi_pw.h
- source/source_psi/setup_psi_pw.cpp

Benefits:
- All member variables now independent of template parameters
- castmem_d2h_impl encapsulates the type-dependent logic
- One step closer to removing template parameters from Setup_Psi_pw

Tested: Compiled successfully in build_5pt and build_1p

* delete useless files

---------

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

Author: mohanchen

source/source_esolver/esolver_ks_lcaopw.cpp

@@ -81,16 +81,18 @@ namespace ModuleESolver
     void ESolver_KS_LIP<T>::before_scf(UnitCell& ucell, const int istep)
     {
         ESolver_KS_PW<T>::before_scf(ucell, istep);
-        this->stp.p_psi_init->initialize_lcao_in_pw(this->psi_local, GlobalV::ofs_running);
+        auto* p_psi_init = static_cast<psi::PSIPrepare<T>*>(this->stp.p_psi_init);
+        p_psi_init->initialize_lcao_in_pw(this->psi_local, GlobalV::ofs_running);
     }
 
     template <typename T>
     void ESolver_KS_LIP<T>::before_all_runners(UnitCell& ucell, const Input_para& inp)
     {
         ESolver_KS_PW<T>::before_all_runners(ucell, inp);
+        auto* p_psi_init = static_cast<psi::PSIPrepare<T>*>(this->stp.p_psi_init);
         delete this->psi_local;
         this->psi_local = new psi::Psi<T>(this->stp.psi_cpu->get_nk(),
-                                          this->stp.p_psi_init->psi_initer->nbands_start(),
+                                          p_psi_init->psi_initer->nbands_start(),
                                           this->stp.psi_cpu->get_nbasis(),
                                           this->kv.ngk,
                                           true);
@@ -105,7 +107,7 @@ namespace ModuleESolver
                                                                            ucell.symm,
                                                                            &this->kv,
                                                                            this->psi_local,
-                                                                           this->stp.psi_t,
+                                                                           this->stp.get_psi_t(),
                                                                            this->pw_wfc,
                                                                            this->pw_rho,
                                                                            this->sf,
@@ -146,7 +148,7 @@ namespace ModuleESolver
         bool skip_charge = PARAM.inp.calculation == "nscf" ? true : false;
 
         hsolver::HSolverLIP<T> hsolver_lip_obj(this->pw_wfc);
-        hsolver_lip_obj.solve(static_cast<hamilt::Hamilt<T>*>(this->p_hamilt), this->stp.psi_t[0], this->pelec, 
+        hsolver_lip_obj.solve(static_cast<hamilt::Hamilt<T>*>(this->p_hamilt), *this->stp.get_psi_t(), this->pelec, 
           *this->psi_local, skip_charge,ucell.tpiba,ucell.nat);
 
         // add exx
@@ -240,7 +242,7 @@ namespace ModuleESolver
             ModuleIO::write_Vxc(PARAM.inp.nspin,
                                 PARAM.globalv.nlocal,
                                 GlobalV::DRANK,
-                                *this->stp.psi_t,
+                                *this->stp.get_psi_t(),
                                 ucell,
                                 this->sf,
                                 this->solvent,

source/source_esolver/esolver_ks_pw.cpp

@@ -43,7 +43,6 @@ ESolver_KS_PW<T, Device>::ESolver_KS_PW()
 {
     this->classname = "ESolver_KS_PW";
     this->basisname = "PW";
-    this->ctx = nullptr;
 }
 
 template <typename T, typename Device>
@@ -112,7 +111,8 @@ void ESolver_KS_PW<T, Device>::before_scf(UnitCell& ucell, const int istep)
 
     if (ucell.cell_parameter_updated)
     {
-        this->stp.p_psi_init->prepare_init(PARAM.inp.pw_seed);
+        auto* p_psi_init = static_cast<psi::PSIPrepare<T, Device>*>(this->stp.p_psi_init);
+        p_psi_init->prepare_init(PARAM.inp.pw_seed);
     }
 
     //! Init Hamiltonian (cell changed)
@@ -128,13 +128,13 @@ void ESolver_KS_PW<T, Device>::before_scf(UnitCell& ucell, const int istep)
     // init DFT+U is done in "before_all_runners" in LCAO basis. This should be refactored, mohan note 2025-11-06
     pw::setup_pot(istep, ucell, this->kv, this->sf, this->pelec, this->Pgrid,
               this->chr, this->locpp, this->ppcell, this->dftu, this->vsep_cell,
-              this->stp.psi_t, static_cast<hamilt::Hamilt<T, Device>*>(this->p_hamilt), this->pw_wfc, this->pw_rhod, PARAM.inp);
+              this->stp.get_psi_t(), static_cast<hamilt::Hamilt<T, Device>*>(this->p_hamilt), this->pw_wfc, this->pw_rhod, PARAM.inp);
 
     // setup psi (electronic wave functions)
-    this->stp.init(static_cast<hamilt::Hamilt<T, Device>*>(this->p_hamilt));
+    this->stp.init(this->p_hamilt);
 
     //! Setup EXX helper for Hamiltonian and psi
-    exx_helper.before_scf(this->p_hamilt, this->stp.psi_t, PARAM.inp);
+    exx_helper.before_scf(this->p_hamilt, this->stp.get_psi_t(), PARAM.inp);
 
     ModuleBase::timer::tick("ESolver_KS_PW", "before_scf");
 }
@@ -152,7 +152,7 @@ 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
-    pw::iter_init_dftu_pw(iter, istep, this->dftu, this->stp.psi_t, this->pelec->wg, ucell, PARAM.inp);
+    pw::iter_init_dftu_pw(iter, istep, this->dftu, this->stp.get_psi_t(), this->pelec->wg, ucell, PARAM.inp);
 }
 
 // Temporary, it should be replaced by hsolver later.
@@ -188,7 +188,7 @@ void ESolver_KS_PW<T, Device>::hamilt2rho_single(UnitCell& ucell, const int iste
                                                      hsolver::DiagoIterAssist<T, Device>::need_subspace,
                                                      PARAM.inp.use_k_continuity);
 
-        hsolver_pw_obj.solve(static_cast<hamilt::Hamilt<T, Device>*>(this->p_hamilt), this->stp.psi_t[0], this->pelec, this->pelec->ekb.c,
+        hsolver_pw_obj.solve(static_cast<hamilt::Hamilt<T, Device>*>(this->p_hamilt), *this->stp.get_psi_t(), this->pelec, this->pelec->ekb.c,
           GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, skip_charge, ucell.tpiba, ucell.nat);
     }
 
@@ -205,7 +205,7 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
     // Related to EXX
     if (GlobalC::exx_info.info_global.cal_exx && !exx_helper.op_exx->first_iter)
     {
-        this->pelec->set_exx(exx_helper.cal_exx_energy(this->stp.psi_t));
+        this->pelec->set_exx(exx_helper.cal_exx_energy(this->stp.get_psi_t()));
     }
 
     // deband is calculated from "output" charge density
@@ -224,7 +224,7 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
     }
 
     // Handle EXX-related operations after SCF iteration
-    exx_helper.iter_finish(this->pelec, &this->chr, this->stp.psi_t, ucell, PARAM.inp, conv_esolver, iter);
+    exx_helper.iter_finish(this->pelec, &this->chr, this->stp.get_psi_t(), ucell, PARAM.inp, conv_esolver, iter);
 
     // check if oscillate for delta_spin method
     pw::check_deltaspin_oscillation(iter, this->drho, this->p_chgmix, PARAM.inp);
@@ -251,7 +251,7 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
     // Output quantities
     ModuleIO::ctrl_scf_pw<T, Device>(istep, ucell, this->pelec, this->chr, this->kv, this->pw_wfc,
               this->pw_rho, this->pw_rhod, this->pw_big, this->stp,
-              this->ctx, this->Pgrid, PARAM.inp);
+              this->Pgrid, PARAM.inp);
 
     ModuleBase::timer::tick("ESolver_KS_PW", "after_scf");
 }
@@ -273,7 +273,7 @@ void ESolver_KS_PW<T, Device>::cal_force(UnitCell& ucell, ModuleBase::matrix& fo
     // Calculate forces
     ff.cal_force(ucell, force, *this->pelec, this->pw_rhod, &ucell.symm,
                  &this->sf, this->solvent, &this->dftu, &this->locpp, &this->ppcell, 
-                 &this->kv, this->pw_wfc, this->stp.psi_d);
+                 &this->kv, this->pw_wfc, this->stp.get_psi_d());
 }
 
 template <typename T, typename Device>
@@ -285,7 +285,7 @@ void ESolver_KS_PW<T, Device>::cal_stress(UnitCell& ucell, ModuleBase::matrix& s
     this->stp.update_psi_d();
 
     ss.cal_stress(stress, ucell, this->dftu, this->locpp, this->ppcell, this->pw_rhod,
-                  &ucell.symm, &this->sf, &this->kv, this->pw_wfc, this->stp.psi_d);
+                  &ucell.symm, &this->sf, &this->kv, this->pw_wfc, this->stp.get_psi_d());
 
     // external stress
     double unit_transform = 0.0;
@@ -304,7 +304,7 @@ void ESolver_KS_PW<T, Device>::after_all_runners(UnitCell& ucell)
 
     ModuleIO::ctrl_runner_pw<T, Device>(ucell, this->pelec, this->pw_wfc, 
             this->pw_rho, this->pw_rhod, this->chr, this->kv, this->stp, 
-            this->sf, this->ppcell, this->solvent, this->ctx, this->Pgrid, PARAM.inp); 
+            this->sf, this->ppcell, this->solvent, this->Pgrid, PARAM.inp); 
 
     elecstate::teardown_estate_pw<T, Device>(this->pelec, this->vsep_cell);
 

source/source_esolver/esolver_ks_pw.h

@@ -57,9 +57,6 @@ class ESolver_KS_PW : public ESolver_KS
     // DFT-1/2 method
     VSep* vsep_cell = nullptr;
 
-    // for get_pchg and get_wf, use ctx as input of fft
-    Device* ctx = {};
-
 };
 } // namespace ModuleESolver
 #endif

source/source_esolver/esolver_sdft_pw.cpp

@@ -168,7 +168,7 @@ void ESolver_SDFT_PW<T, Device>::hamilt2rho_single(UnitCell& ucell, int istep, i
 
     hsolver_pw_sdft_obj.solve(ucell,
                               static_cast<hamilt::Hamilt<T, Device>*>(this->p_hamilt),
-                              this->stp.psi_t[0],
+                              *this->stp.get_psi_t(),
                               this->stp.psi_cpu[0],
                               this->pelec,
                               this->pw_wfc,
@@ -221,7 +221,7 @@ void ESolver_SDFT_PW<T, Device>::cal_force(UnitCell& ucell, ModuleBase::matrix&
                     this->locpp,
                     this->ppcell,
                     ucell,
-                    *this->stp.psi_t,
+                    *this->stp.get_psi_t(),
                     this->stowf);
 }
 
@@ -236,7 +236,7 @@ void ESolver_SDFT_PW<T, Device>::cal_stress(UnitCell& ucell, ModuleBase::matrix&
                   &this->sf,
                   &this->kv,
                   this->pw_wfc,
-                  *this->stp.psi_t,
+                  *this->stp.get_psi_t(),
                   this->stowf,
                   &this->chr,
                   &this->locpp,
@@ -289,7 +289,7 @@ void ESolver_SDFT_PW<T, Device>::after_all_runners(UnitCell& ucell)
                                               &this->kv,
                                               this->pelec,
                                               this->pw_wfc,
-                                              this->stp.psi_t,
+                                              this->stp.get_psi_t(),
                                               &this->ppcell,
                                               static_cast<hamilt::Hamilt<std::complex<double>, Device>*>(this->p_hamilt),
                                               this->stoche,

source/source_io/module_ctrl/ctrl_output_pw.cpp

@@ -91,15 +91,19 @@ void ModuleIO::ctrl_scf_pw(const int istep,
 		const ModulePW::PW_Basis *pw_rhod,
 		const ModulePW::PW_Basis_Big *pw_big,
         Setup_Psi_pw<T, Device> &stp,
-        const Device* ctx,
         const Parallel_Grid &para_grid,
         const Input_para& inp)
 {
     ModuleBase::TITLE("ModuleIO", "ctrl_scf_pw");
     ModuleBase::timer::tick("ModuleIO", "ctrl_scf_pw");
 
+    // Create local ctx for device type deduction
+    Device* ctx = nullptr;
+
     // Transfer data from device (GPU) to host (CPU) in pw basis
-    stp.copy_d2h(ctx);
+    base_device::DeviceContext* device_ctx = &base_device::DeviceContext::instance();
+    device_ctx->set_device_type(stp.get_device_type());
+    stp.copy_d2h(device_ctx);
 
     //----------------------------------------------------------
     //! 4) Compute density of states (DOS)
@@ -164,7 +168,7 @@ void ModuleIO::ctrl_scf_pw(const int istep,
         // update psi_d
         stp.update_psi_d();
 
-        const int nbands = stp.psi_t->get_nbands();
+        const int nbands = stp.get_nbands();
         const int ngmc = chr.ngmc;
 
         ModuleIO::get_pchg_pw(inp.out_pchg,
@@ -173,7 +177,7 @@ void ModuleIO::ctrl_scf_pw(const int istep,
                               pw_rhod->nxyz,
                               ngmc,
                               &ucell,
-                              stp.psi_d,
+                              stp.get_psi_d(),
                               pw_rhod,
                               pw_wfc,
                               ctx,
@@ -235,7 +239,7 @@ void ModuleIO::ctrl_scf_pw(const int istep,
     if (inp.onsite_radius > 0)
     { // float type has not been implemented
         auto* onsite_p = projectors::OnsiteProjector<double, Device>::get_instance();
-        onsite_p->cal_occupations(reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(stp.psi_t),
+        onsite_p->cal_occupations(reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(stp.get_psi_t()),
                                   pelec->wg);
     }
 
@@ -255,13 +259,15 @@ void ModuleIO::ctrl_runner_pw(UnitCell& ucell,
         Structure_Factor &sf,
         pseudopot_cell_vnl &ppcell,
 		surchem &solvent,
-        const Device* ctx,
         Parallel_Grid &para_grid,
         const Input_para& inp)
 {
     ModuleBase::TITLE("ModuleIO", "ctrl_runner_pw");
     ModuleBase::timer::tick("ModuleIO", "ctrl_runner_pw");
 
+    // Create local ctx for device type deduction
+    Device* ctx = nullptr;
+
 	//----------------------------------------------------------
 	//! 1) Compute LDOS
 	//----------------------------------------------------------
@@ -303,11 +309,11 @@ void ModuleIO::ctrl_runner_pw(UnitCell& ucell,
 
         ModuleIO::get_wf_pw(inp.out_wfc_norm,
                             inp.out_wfc_re_im,
-                            stp.psi_t->get_nbands(),
+                            stp.get_nbands(),
                             inp.nspin,
                             pw_rhod->nxyz,
                             &ucell,
-                            stp.psi_d,
+                            stp.get_psi_d(),
                             pw_wfc,
                             ctx,
                             para_grid,
@@ -323,7 +329,7 @@ void ModuleIO::ctrl_runner_pw(UnitCell& ucell,
     if (inp.cal_cond)
     {
         using Real = typename GetTypeReal<T>::type;
-        EleCond<Real, Device> elec_cond(&ucell, &kv, pelec, pw_wfc, stp.psi_t, &ppcell);
+        EleCond<Real, Device> elec_cond(&ucell, &kv, pelec, pw_wfc, stp.get_psi_t(), &ppcell);
         elec_cond.KG(inp.cond_smear,
                      inp.cond_fwhm,
                      inp.cond_wcut,
@@ -360,7 +366,7 @@ void ModuleIO::ctrl_runner_pw(UnitCell& ucell,
                                              pw_rho);
 
         write_mlkedf_desc.generateTrainData_KS(PARAM.globalv.global_mlkedf_descriptor_dir,
-                                               stp.psi_t,
+                                               stp.get_psi_t(),
                                                pelec,
                                                pw_wfc,
                                                pw_rho,
@@ -384,7 +390,6 @@ template void ModuleIO::ctrl_scf_pw<std::complex<float>, base_device::DEVICE_CPU
     const ModulePW::PW_Basis *pw_rhod,
     const ModulePW::PW_Basis_Big *pw_big,
     Setup_Psi_pw<std::complex<float>, base_device::DEVICE_CPU> &stp,
-    const base_device::DEVICE_CPU* ctx,
     const Parallel_Grid &para_grid,
     const Input_para& inp);
 
@@ -400,7 +405,6 @@ template void ModuleIO::ctrl_scf_pw<std::complex<double>, base_device::DEVICE_CP
     const ModulePW::PW_Basis *pw_rhod,
     const ModulePW::PW_Basis_Big *pw_big,
     Setup_Psi_pw<std::complex<double>, base_device::DEVICE_CPU> &stp,
-    const base_device::DEVICE_CPU* ctx,
     const Parallel_Grid &para_grid,
     const Input_para& inp);
 
@@ -417,13 +421,12 @@ template void ModuleIO::ctrl_scf_pw<std::complex<float>, base_device::DEVICE_GPU
     const ModulePW::PW_Basis *pw_rhod,
     const ModulePW::PW_Basis_Big *pw_big,
     Setup_Psi_pw<std::complex<float>, base_device::DEVICE_GPU> &stp,
-    const base_device::DEVICE_GPU* ctx,
     const Parallel_Grid &para_grid,
     const Input_para& inp);
 
 // complex<double> + GPU
 template void ModuleIO::ctrl_scf_pw<std::complex<double>, base_device::DEVICE_GPU>(
-	const int nstep,
+    const int nstep,
 	UnitCell& ucell,
     elecstate::ElecState* pelec,
     const Charge &chr,
@@ -433,7 +436,6 @@ template void ModuleIO::ctrl_scf_pw<std::complex<double>, base_device::DEVICE_GP
     const ModulePW::PW_Basis *pw_rhod,
     const ModulePW::PW_Basis_Big *pw_big,
     Setup_Psi_pw<std::complex<double>, base_device::DEVICE_GPU> &stp,
-    const base_device::DEVICE_GPU* ctx,
     const Parallel_Grid &para_grid,
     const Input_para& inp);
 #endif
@@ -444,14 +446,13 @@ template void ModuleIO::ctrl_runner_pw<std::complex<float>, base_device::DEVICE_
 	elecstate::ElecState* pelec,	
     ModulePW::PW_Basis_K* pw_wfc,
     ModulePW::PW_Basis* pw_rho,
-	ModulePW::PW_Basis* pw_rhod,
+    ModulePW::PW_Basis* pw_rhod,
 	Charge &chr,
-	K_Vectors &kv,
+    K_Vectors &kv,
     Setup_Psi_pw<std::complex<float>, base_device::DEVICE_CPU> &stp,
     Structure_Factor &sf,
     pseudopot_cell_vnl &ppcell,
 	surchem &solvent,
-    const base_device::DEVICE_CPU* ctx,
     Parallel_Grid &para_grid,
     const Input_para& inp);
 
@@ -461,14 +462,13 @@ template void ModuleIO::ctrl_runner_pw<std::complex<double>, base_device::DEVICE
 	elecstate::ElecState* pelec,	
     ModulePW::PW_Basis_K* pw_wfc,
     ModulePW::PW_Basis* pw_rho,
-	ModulePW::PW_Basis* pw_rhod,
+    ModulePW::PW_Basis* pw_rhod,
 	Charge &chr,
-	K_Vectors &kv,
+    K_Vectors &kv,
     Setup_Psi_pw<std::complex<double>, base_device::DEVICE_CPU> &stp,
     Structure_Factor &sf,
     pseudopot_cell_vnl &ppcell,
 	surchem &solvent,
-    const base_device::DEVICE_CPU* ctx,
     Parallel_Grid &para_grid,
     const Input_para& inp);
 
@@ -481,12 +481,11 @@ template void ModuleIO::ctrl_runner_pw<std::complex<float>, base_device::DEVICE_
     ModulePW::PW_Basis* pw_rho,
     ModulePW::PW_Basis* pw_rhod,
 	Charge &chr,
-	K_Vectors &kv,
+    K_Vectors &kv,
     Setup_Psi_pw<std::complex<float>, base_device::DEVICE_GPU> &stp,
     Structure_Factor &sf,
     pseudopot_cell_vnl &ppcell,
 	surchem &solvent,
-    const base_device::DEVICE_GPU* ctx,
     Parallel_Grid &para_grid,
     const Input_para& inp);
 
@@ -498,12 +497,11 @@ template void ModuleIO::ctrl_runner_pw<std::complex<double>, base_device::DEVICE
     ModulePW::PW_Basis* pw_rho,
     ModulePW::PW_Basis* pw_rhod,
 	Charge &chr,
-	K_Vectors &kv,
+    K_Vectors &kv,
     Setup_Psi_pw<std::complex<double>, base_device::DEVICE_GPU> &stp,
     Structure_Factor &sf,
     pseudopot_cell_vnl &ppcell,
 	surchem &solvent,
-    const base_device::DEVICE_GPU* ctx,
     Parallel_Grid &para_grid,
     const Input_para& inp);
 #endif