Update after_scf in ESolver

source/module_esolver/esolver_fp.cpp

@@ -134,18 +134,18 @@ void ESolver_FP::after_scf(UnitCell& ucell, const int istep, const bool conv_eso
 {
     ModuleBase::TITLE("ESolver_FP", "after_scf");
 
-    // 0) output convergence information
+    // 1) output convergence information
     ModuleIO::output_convergence_after_scf(conv_esolver, this->pelec->f_en.etot);
 
-    // 1) write fermi energy
+    // 2) write fermi energy
     ModuleIO::output_efermi(conv_esolver, this->pelec->eferm.ef);
 
-    // 2) update delta rho for charge extrapolation
+    // 3) update delta rho for charge extrapolation
     CE.update_delta_rho(ucell, &(this->chr), &(this->sf));
 
     if (istep % PARAM.inp.out_interval == 0)
     {
-        // 3) write charge density
+        // 4) write charge density
         if (PARAM.inp.out_chg[0] > 0)
         {
             for (int is = 0; is < PARAM.inp.nspin; is++)
@@ -187,7 +187,7 @@ void ESolver_FP::after_scf(UnitCell& ucell, const int istep, const bool conv_eso
             }
         }
 
-        // 4) write potential
+        // 5) write potential
         if (PARAM.inp.out_pot == 1 || PARAM.inp.out_pot == 3)
         {
             for (int is = 0; is < PARAM.inp.nspin; is++)
@@ -223,7 +223,7 @@ void ESolver_FP::after_scf(UnitCell& ucell, const int istep, const bool conv_eso
                 this->solvent);
         }
 
-        // 5) write ELF
+        // 6) write ELF
         if (PARAM.inp.out_elf[0] > 0)
         {
             this->pelec->charge->cal_elf = true;

source/module_esolver/esolver_ks_lcao_tddft.cpp

@@ -225,13 +225,20 @@ void ESolver_KS_LCAO_TDDFT<Device>::update_pot(UnitCell& ucell,
     {
         if (this->psi_laststep == nullptr)
         {
+            int ncol_tmp = 0;
+            int nrow_tmp = 0;
 #ifdef __MPI
-            this->psi_laststep = new psi::Psi<std::complex<double>>(kv.get_nks(), ncol_nbands, nrow, kv.ngk, true);
+            ncol_tmp = ncol_nbands;
+            nrow_tmp = nrow;
 #else
-            this->psi_laststep = new psi::Psi<std::complex<double>>(kv.get_nks(), nbands, nlocal, kv.ngk, true);
+            ncol_tmp = nbands;
+            nrow_tmp = nlocal;
 #endif
+            this->psi_laststep = new psi::Psi<std::complex<double>>(kv.get_nks(), ncol_tmp, nrow_tmp, kv.ngk, true);
+
         }
 
+        // allocate memory for Hk_laststep and Sk_laststep
         if (td_htype == 1)
         {
             // Length of Hk_laststep and Sk_laststep, nlocal * nlocal for global, nloc for local
@@ -259,11 +266,14 @@ void ESolver_KS_LCAO_TDDFT<Device>::update_pot(UnitCell& ucell,
             }
         }
 
+        // put information to Hk_laststep and Sk_laststep
         for (int ik = 0; ik < kv.get_nks(); ++ik)
         {
             this->psi->fix_k(ik);
             this->psi_laststep->fix_k(ik);
-            int size0 = psi->get_nbands() * psi->get_nbasis();
+
+            // copy the data from psi to psi_laststep
+            const int size0 = psi->get_nbands() * psi->get_nbasis();
             for (int index = 0; index < size0; ++index)
             {
                 psi_laststep[0].get_pointer()[index] = psi[0].get_pointer()[index];
@@ -273,7 +283,8 @@ void ESolver_KS_LCAO_TDDFT<Device>::update_pot(UnitCell& ucell,
             if (td_htype == 1)
             {
                 this->p_hamilt->updateHk(ik);
-                hamilt::MatrixBlock<complex<double>> h_mat, s_mat;
+                hamilt::MatrixBlock<complex<double>> h_mat;
+                hamilt::MatrixBlock<complex<double>> s_mat;
                 this->p_hamilt->matrix(h_mat, s_mat);
 
                 if (use_tensor && use_lapack)
@@ -285,7 +296,8 @@ void ESolver_KS_LCAO_TDDFT<Device>::update_pot(UnitCell& ucell,
                     MPI_Comm_rank(MPI_COMM_WORLD, &myid);
                     MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
 
-                    Matrix_g<std::complex<double>> h_mat_g, s_mat_g; // Global matrix structure
+                    Matrix_g<std::complex<double>> h_mat_g; // Global matrix structure
+                    Matrix_g<std::complex<double>> s_mat_g; // Global matrix structure
 
                     // Collect H matrix
                     gatherMatrix(myid, 0, h_mat, h_mat_g);
@@ -343,6 +355,9 @@ void ESolver_KS_LCAO_TDDFT<Device>::after_scf(UnitCell& ucell, const int istep,
     ModuleBase::TITLE("ESolver_LCAO_TDDFT", "after_scf");
     ModuleBase::timer::tick("ESolver_LCAO_TDDFT", "after_scf");
 
+    ESolver_KS_LCAO<std::complex<double>, double>::after_scf(ucell, istep, conv_esolver);
+
+    // (1) write dipole information
     for (int is = 0; is < PARAM.inp.nspin; is++)
     {
         if (PARAM.inp.out_dipole == 1)
@@ -357,6 +372,8 @@ void ESolver_KS_LCAO_TDDFT<Device>::after_scf(UnitCell& ucell, const int istep,
                                    ss_dipole.str());
         }
     }
+
+     // (2) write current information
     if (TD_Velocity::out_current == true)
     {
         elecstate::DensityMatrix<std::complex<double>, double>* tmp_DM
@@ -373,7 +390,7 @@ void ESolver_KS_LCAO_TDDFT<Device>::after_scf(UnitCell& ucell, const int istep,
                                 orb_,
                                 this->RA);
     }
-    ESolver_KS_LCAO<std::complex<double>, double>::after_scf(ucell, istep, conv_esolver);
+
 
     ModuleBase::timer::tick("ESolver_LCAO_TDDFT", "after_scf");
 }
@@ -385,14 +402,18 @@ void ESolver_KS_LCAO_TDDFT<Device>::weight_dm_rho()
     {
         this->pelec->fixed_weights(PARAM.inp.ocp_kb, PARAM.inp.nbands, PARAM.inp.nelec);
     }
+
+    // calculate Eband energy
     this->pelec->calEBand();
 
+    // calculate the density matrix
     ModuleBase::GlobalFunc::NOTE("Calculate the density matrix.");
 
     auto _pes = dynamic_cast<elecstate::ElecStateLCAO<std::complex<double>>*>(this->pelec);
     elecstate::cal_dm_psi(_pes->DM->get_paraV_pointer(), _pes->wg, this->psi[0], *(_pes->DM));
     _pes->DM->cal_DMR();
 
+    // get the real-space charge density
     this->pelec->psiToRho(this->psi[0]);
 }
 

source/module_esolver/esolver_ks_pw.cpp

@@ -551,33 +551,45 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
     ModuleBase::TITLE("ESolver_KS_PW", "after_scf");
     ModuleBase::timer::tick("ESolver_KS_PW", "after_scf");
 
-    // 1) calculate the kinetic energy density tau, sunliang 2024-09-18
+    //------------------------------------------------------------------
+    // 1) calculate the kinetic energy density tau in pw basis
+    // sunliang 2024-09-18
+    //------------------------------------------------------------------
     if (PARAM.inp.out_elf[0] > 0)
     {
         this->pelec->cal_tau(*(this->psi));
     }
 
+    //------------------------------------------------------------------
     // 2) call after_scf() of ESolver_KS
+    //------------------------------------------------------------------
     ESolver_KS<T, Device>::after_scf(ucell, istep, conv_esolver);
 
+
+    //------------------------------------------------------------------
     // 3) output wavefunctions in pw basis
+    //------------------------------------------------------------------
     if (PARAM.inp.out_wfc_pw == 1 || PARAM.inp.out_wfc_pw == 2)
     {
         std::stringstream ssw;
         ssw << PARAM.globalv.global_out_dir << "WAVEFUNC";
         ModuleIO::write_wfc_pw(ssw.str(), this->psi[0], this->kv, this->pw_wfc);
     }
 
+    //------------------------------------------------------------------
     // 4) transfer data from GPU to CPU in pw basis
     // a question: the wavefunctions have been output, then the data transfer occurs? mohan 20250302
+    //------------------------------------------------------------------
     if (this->device == base_device::GpuDevice)
     {
         castmem_2d_d2h_op()(this->psi[0].get_pointer() - this->psi[0].get_psi_bias(),
                             this->kspw_psi[0].get_pointer() - this->kspw_psi[0].get_psi_bias(),
                             this->psi[0].size());
     }
 
+    //------------------------------------------------------------------
     // 5) calculate band-decomposed (partial) charge density in pw basis
+    //------------------------------------------------------------------
     const std::vector<int> bands_to_print = PARAM.inp.bands_to_print;
     if (bands_to_print.size() > 0)
     {
@@ -604,7 +616,9 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
                               PARAM.inp.if_separate_k);
     }
 
-    //! 6) calculate Wannier functions in PW basis
+    //------------------------------------------------------------------
+    //! 6) calculate Wannier functions in pw basis
+    //------------------------------------------------------------------
     if (PARAM.inp.calculation == "nscf" && PARAM.inp.towannier90)
     {
         std::cout << FmtCore::format("\n * * * * * *\n << Start %s.\n", "Wannier functions calculation");
@@ -620,7 +634,9 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
         std::cout << FmtCore::format(" >> Finish %s.\n * * * * * *\n", "Wannier functions calculation");
     }
 
-    //! 7) calculate Berry phase polarization
+    //------------------------------------------------------------------
+    //! 7) calculate Berry phase polarization in pw basis
+    //------------------------------------------------------------------
     if (PARAM.inp.calculation == "nscf" && berryphase::berry_phase_flag && ModuleSymmetry::Symmetry::symm_flag != 1)
     {
         std::cout << FmtCore::format("\n * * * * * *\n << Start %s.\n", "Berry phase polarization");
@@ -629,8 +645,10 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
         std::cout << FmtCore::format(" >> Finish %s.\n * * * * * *\n", "Berry phase polarization");
     }
 
-    // 8) write spin constrian results
+    //------------------------------------------------------------------
+    // 8) write spin constrian results in pw basis
     // spin constrain calculations, write atomic magnetization and magnetic force.
+    //------------------------------------------------------------------
     if (PARAM.inp.sc_mag_switch)
     {
         spinconstrain::SpinConstrain<std::complex<double>>& sc
@@ -639,7 +657,9 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
         sc.print_Mag_Force(GlobalV::ofs_running);
     }
 
+    //------------------------------------------------------------------
     // 9) write onsite occupations for charge and magnetizations
+    //------------------------------------------------------------------
     if (PARAM.inp.onsite_radius > 0)
     { // float type has not been implemented
         auto* onsite_p = projectors::OnsiteProjector<double, Device>::get_instance();

source/module_esolver/esolver_of.cpp

@@ -489,30 +489,47 @@ void ESolver_OF::after_opt(const int istep, UnitCell& ucell, const bool conv_eso
     ModuleBase::TITLE("ESolver_OF", "after_opt");
     ModuleBase::timer::tick("ESolver_OF", "after_opt");
 
-    // 1) calculate the kinetic energy density
+    //------------------------------------------------------------------
+    // 1) calculate kinetic energy density and ELF
+    //------------------------------------------------------------------
     if (PARAM.inp.out_elf[0] > 0)
     {
         this->kinetic_energy_density(this->pelec->charge->rho, this->pphi_, this->pelec->charge->kin_r);
     }
 
+    //------------------------------------------------------------------
+    // 2) call after_scf() of ESolver_FP
+    //------------------------------------------------------------------
+    ESolver_FP::after_scf(ucell, istep, conv_esolver);
+
+
+    // should not be here? mohan note 2025-03-03
     for (int ir = 0; ir < this->pw_rho->nrxx; ++ir)
     {
         this->pelec->charge->rho_save[0][ir] = this->pelec->charge->rho[0][ir];
     }
 
 #ifdef __MLKEDF
+    //------------------------------------------------------------------
     // Check the positivity of Pauli energy
+    //------------------------------------------------------------------
     if (this->of_kinetic_ == "ml")
     {
         this->tf_->get_energy(this->pelec->charge->rho);
-        std::cout << "ML Term = " << this->ml_->ml_energy << " Ry, TF Term = " << this->tf_->tf_energy << " Ry." << std::endl;
+
+        std::cout << "ML Term = " << this->ml_->ml_energy 
+                  << " Ry, TF Term = " << this->tf_->tf_energy 
+                  << " Ry." << std::endl;
+
         if (this->ml_->ml_energy >= this->tf_->tf_energy)
         {
             std::cout << "WARNING: ML >= TF" << std::endl;
         }
     }
 
+    //------------------------------------------------------------------
     // Generate data if needed
+    //------------------------------------------------------------------
     if (PARAM.inp.of_ml_gene_data)
     {
         this->pelec->pot->update_from_charge(pelec->charge, &ucell); // Hartree + XC + external
@@ -533,8 +550,6 @@ void ESolver_OF::after_opt(const int istep, UnitCell& ucell, const bool conv_eso
         this->ml_->generateTrainData(pelec->charge->rho, *(this->wt_), *(this->tf_), this->pw_rho, vr_eff);
     }
 #endif
-    // 2) call after_scf() of ESolver_FP
-    ESolver_FP::after_scf(ucell, istep, conv_esolver);
 
     ModuleBase::timer::tick("ESolver_OF", "after_opt");
 }