Refactor: RT-TDDFT ESolver ESolver_KS_LCAO_TDDFT (#6668)

* Remove redundant includes in RT-TDDFT

* Refactor store_h_s_psi

* Refactor Hk and Sk with Tensor

* Refactor MPI utility functions

* Refactor gather and distribute Psi function

* Modify the output suffix of some text files from .dat to .txt

* Change dipole file name from SPIN*_DIPOLE to dipole_s*.txt

* Move the output functions in after_scf to ctrl_output_td

* Only calculate EDM in RT-TDDFT when the electronic step ends

* Fix MPI bug

* Fix LCAO macro bug

Author: AsTonyshment

docs/advanced/input_files/input-main.md

@@ -3185,14 +3185,14 @@ These variables are used to control molecular dynamics calculations. For more in
 
 - **Type**: Boolean
 - **Description**: Control whether to restart molecular dynamics calculations and time-dependent density functional theory calculations.
-  - True: ABACUS will read in `${read_file_dir}/Restart_md.dat` to determine the current step `${md_step}`, then read in the corresponding `STRU_MD_${md_step}` in the folder `OUT.$suffix/STRU/` automatically. For tddft, ABACUS will also read in `WFC_NAO_K${kpoint}` of the last step (You need to set out_wfc_lcao=1 and out_app_flag=0 to obtain this file).
+  - True: ABACUS will read in `${read_file_dir}/Restart_md.txt` to determine the current step `${md_step}`, then read in the corresponding `STRU_MD_${md_step}` in the folder `OUT.$suffix/STRU/` automatically. For tddft, ABACUS will also read in `WFC_NAO_K${kpoint}` of the last step (You need to set out_wfc_lcao=1 and out_app_flag=0 to obtain this file).
   - False: ABACUS will start molecular dynamics calculations normally from the first step.
 - **Default**: False
 
 ### md_restartfreq
 
 - **Type**: Integer
-- **Description**: The output frequency of `OUT.${suffix}/Restart_md.dat` and structural files in the directory `OUT.${suffix}/STRIU/`, which are used to restart molecular dynamics calculations, see [md_restart](#md_restart) in detail.
+- **Description**: The output frequency of `OUT.${suffix}/Restart_md.txt` and structural files in the directory `OUT.${suffix}/STRIU/`, which are used to restart molecular dynamics calculations, see [md_restart](#md_restart) in detail.
 - **Default**: 5
 
 ### md_dumpfreq

source/Makefile.Objects

@@ -560,6 +560,7 @@ OBJS_IO=input_conv.o\
     ctrl_iter_lcao.o\
     ctrl_output_fp.o\
     ctrl_output_pw.o\
+    ctrl_output_td.o\
     para_json.o\
     abacusjson.o\
     general_info.o\

source/source_esolver/esolver_ks_lcao_tddft.cpp

@@ -2,52 +2,13 @@
 #define ESOLVER_KS_LCAO_TDDFT_H
 #include "esolver_ks.h"
 #include "esolver_ks_lcao.h"
-#include "source_base/module_external/scalapack_connector.h" // Cpxgemr2d
-#include "source_lcao/record_adj.h"
-#include "source_psi/psi.h"
-#include "source_lcao/module_rt/velocity_op.h"
+#include "source_base/module_container/ATen/core/tensor.h" // ct::Tensor
+#include "source_lcao/module_rt/gather_mat.h"              // MPI gathering and distributing functions
 #include "source_lcao/module_rt/td_info.h"
+#include "source_lcao/module_rt/velocity_op.h"
 
 namespace ModuleESolver
 {
-//------------------------ MPI gathering and distributing functions ------------------------//
-// This struct is used for collecting matrices from all processes to root process
-template <typename T>
-struct Matrix_g
-{
-    std::shared_ptr<T> p;
-    size_t row;
-    size_t col;
-    std::shared_ptr<int> desc;
-};
-
-// Collect matrices from all processes to root process
-template <typename T>
-void gatherMatrix(const int myid, const int root_proc, const hamilt::MatrixBlock<T>& mat_l, Matrix_g<T>& mat_g)
-{
-    const int* desca = mat_l.desc; // Obtain the descriptor of the local matrix
-    int ctxt = desca[1];           // BLACS context
-    int nrows = desca[2];          // Global matrix row number
-    int ncols = desca[3];          // Global matrix column number
-
-    if (myid == root_proc)
-    {
-        mat_g.p.reset(new T[nrows * ncols]); // No need to delete[] since it is a shared_ptr
-    }
-    else
-    {
-        mat_g.p.reset(new T[nrows * ncols]); // Placeholder for non-root processes
-    }
-
-    // Set the descriptor of the global matrix
-    mat_g.desc.reset(new int[9]{1, ctxt, nrows, ncols, nrows, ncols, 0, 0, nrows});
-    mat_g.row = nrows;
-    mat_g.col = ncols;
-
-    // Call the Cpxgemr2d function in ScaLAPACK to collect the matrix data
-    Cpxgemr2d(nrows, ncols, mat_l.p, 1, 1, const_cast<int*>(desca), mat_g.p.get(), 1, 1, mat_g.desc.get(), ctxt);
-}
-//------------------------ MPI gathering and distributing functions ------------------------//
 
 template <typename TR, typename Device = base_device::DEVICE_CPU>
 class ESolver_KS_LCAO_TDDFT : public ESolver_KS_LCAO<std::complex<double>, TR>
@@ -64,29 +25,36 @@ class ESolver_KS_LCAO_TDDFT : public ESolver_KS_LCAO<std::complex<double>, TR>
 
     virtual void hamilt2rho_single(UnitCell& ucell, const int istep, const int iter, const double ethr) override;
 
-    // mohan change update_pot to save2, 2025-10-17
-    void save2(UnitCell& ucell, const int istep, const int iter, const bool conv_esolver);
+    void store_h_s_psi(UnitCell& ucell, const int istep, const int iter, const bool conv_esolver);
 
-    virtual void iter_finish(UnitCell& ucell, const int istep, int& iter, bool& conv_esolver) override;
+    void iter_finish(UnitCell& ucell,
+                     const int istep,
+                     const int estep,
+                     const int estep_max,
+                     int& iter,
+                     bool& conv_esolver);
 
     virtual void after_scf(UnitCell& ucell, const int istep, const bool conv_esolver) override;
 
     void print_step();
-    //! wave functions of last time step
-    psi::Psi<std::complex<double>>* psi_laststep = nullptr;
 
-    //! Hamiltonian of last time step
-    std::complex<double>** Hk_laststep = nullptr;
+    //! Wave function for all k-points of last time step
+    psi::Psi<std::complex<double>>* psi_laststep = nullptr;
 
-    //! Overlap matrix of last time step
-    std::complex<double>** Sk_laststep = nullptr;
+    //! Hamiltonian for all k-points of last time step
+    ct::Tensor Hk_laststep = ct::Tensor(ct::DataType::DT_COMPLEX_DOUBLE);
 
-    const int td_htype = 1;
+    //! Overlap matrix for all k-points of last time step
+    ct::Tensor Sk_laststep = ct::Tensor(ct::DataType::DT_COMPLEX_DOUBLE);
 
     //! Control heterogeneous computing of the TDDFT solver
     bool use_tensor = false;
     bool use_lapack = false;
 
+    // Control the device type for Hk_laststep and Sk_laststep
+    // Set to CPU temporarily, should wait for further GPU development
+    static constexpr ct::DeviceType ct_device_type_hs = ct::DeviceType::CpuDevice;
+
     //! Total steps for evolving the wave function
     int totstep = -1;
 
@@ -95,13 +63,12 @@ class ESolver_KS_LCAO_TDDFT : public ESolver_KS_LCAO<std::complex<double>, TR>
 
     TD_info* td_p = nullptr;
 
-    //! doubt
+    //! Restart flag
     bool restart_done = false;
 
   private:
     void weight_dm_rho(const UnitCell& ucell);
 };
 
 } // namespace ModuleESolver
-#endif
-
+#endif // ESOLVER_KS_LCAO_TDDFT_H

source/source_esolver/esolver_ks_lcao_tddft.h

@@ -2,6 +2,7 @@ list(APPEND objects
     input_conv.cpp
     ctrl_output_fp.cpp
     ctrl_output_pw.cpp
+    ctrl_output_td.cpp
     bessel_basis.cpp
     cal_test.cpp
     cal_dos.cpp

source/source_io/CMakeLists.txt

@@ -0,0 +1,163 @@
+#include "ctrl_output_td.h"
+
+#include "source_base/parallel_global.h"
+#include "source_io/dipole_io.h"
+#include "source_io/module_parameter/parameter.h"
+#include "source_io/td_current_io.h"
+
+namespace ModuleIO
+{
+
+template <typename TR>
+void ctrl_output_td(const UnitCell& ucell,
+                    double** rho_save,
+                    const ModulePW::PW_Basis* rhopw,
+                    const int istep,
+                    const psi::Psi<std::complex<double>>* psi,
+                    const elecstate::ElecState* pelec,
+                    const K_Vectors& kv,
+                    const TwoCenterIntegrator* intor,
+                    const Parallel_Orbitals* pv,
+                    const LCAO_Orbitals& orb,
+                    const Velocity_op<TR>* velocity_mat,
+                    Record_adj& RA,
+                    TD_info* td_p)
+{
+    ModuleBase::TITLE("ModuleIO", "ctrl_output_td");
+
+    // Original code commented out, might need reference later
+
+    // // (1) Write dipole information
+    // for (int is = 0; is < PARAM.inp.nspin; is++)
+    // {
+    //     if (PARAM.inp.out_dipole == 1)
+    //     {
+    //         std::stringstream ss_dipole;
+    //         ss_dipole << PARAM.globalv.global_out_dir << "dipole_s" << is + 1 << ".txt";
+    //         ModuleIO::write_dipole(ucell, this->chr.rho_save[is], this->chr.rhopw, is, istep, ss_dipole.str());
+    //     }
+    // }
+
+    // // (2) Write current information
+    // elecstate::DensityMatrix<std::complex<double>, double>* tmp_DM
+    //     = dynamic_cast<elecstate::ElecStateLCAO<std::complex<double>>*>(this->pelec)->get_DM();
+    // if (TD_info::out_current)
+    // {
+    //     if (TD_info::out_current_k)
+    //     {
+    //         ModuleIO::write_current_eachk(ucell,
+    //                                       istep,
+    //                                       this->psi,
+    //                                       this->pelec,
+    //                                       this->kv,
+    //                                       this->two_center_bundle_.overlap_orb.get(),
+    //                                       tmp_DM->get_paraV_pointer(),
+    //                                       this->orb_,
+    //                                       this->velocity_mat,
+    //                                       this->RA);
+    //     }
+    //     else
+    //     {
+    //         ModuleIO::write_current(ucell,
+    //                                 istep,
+    //                                 this->psi,
+    //                                 this->pelec,
+    //                                 this->kv,
+    //                                 this->two_center_bundle_.overlap_orb.get(),
+    //                                 tmp_DM->get_paraV_pointer(),
+    //                                 this->orb_,
+    //                                 this->velocity_mat,
+    //                                 this->RA);
+    //     }
+    // }
+
+    // // (3) Output file for restart
+    // if (PARAM.inp.out_freq_ion > 0) // default value of out_freq_ion is 0
+    // {
+    //     if (istep % PARAM.inp.out_freq_ion == 0)
+    //     {
+    //         td_p->out_restart_info(istep, elecstate::H_TDDFT_pw::At, elecstate::H_TDDFT_pw::At_laststep);
+    //     }
+    // }
+
+#ifdef __LCAO
+    // (1) Write dipole information
+    for (int is = 0; is < PARAM.inp.nspin; ++is)
+    {
+        if (PARAM.inp.out_dipole == 1)
+        {
+            std::stringstream ss_dipole;
+            ss_dipole << PARAM.globalv.global_out_dir << "dipole_s" << is + 1 << ".txt";
+            ModuleIO::write_dipole(ucell, rho_save[is], rhopw, is, istep, ss_dipole.str());
+        }
+    }
+
+    // (2) Write current information
+    const elecstate::ElecStateLCAO<std::complex<double>>* pelec_lcao
+        = dynamic_cast<const elecstate::ElecStateLCAO<std::complex<double>>*>(pelec);
+
+    if (!pelec_lcao)
+    {
+        ModuleBase::WARNING_QUIT("ModuleIO::ctrl_output_td", "Failed to cast ElecState to ElecStateLCAO");
+    }
+
+    if (TD_info::out_current)
+    {
+        if (TD_info::out_current_k)
+        {
+            ModuleIO::write_current_eachk<TR>(ucell, istep, psi, pelec, kv, intor, pv, orb, velocity_mat, RA);
+        }
+        else
+        {
+            ModuleIO::write_current<TR>(ucell, istep, psi, pelec, kv, intor, pv, orb, velocity_mat, RA);
+        }
+    }
+
+    // (3) Output file for restart
+    if (PARAM.inp.out_freq_ion > 0) // default value of out_freq_ion is 0
+    {
+        if (istep % PARAM.inp.out_freq_ion == 0)
+        {
+            if (td_p != nullptr)
+            {
+                td_p->out_restart_info(istep, elecstate::H_TDDFT_pw::At, elecstate::H_TDDFT_pw::At_laststep);
+            }
+            else
+            {
+                ModuleBase::WARNING_QUIT("ModuleIO::ctrl_output_td",
+                                         "TD_info pointer is null, cannot output restart info.");
+            }
+        }
+    }
+#endif // __LCAO
+}
+
+template void ctrl_output_td<double>(const UnitCell&,
+                                     double**,
+                                     const ModulePW::PW_Basis*,
+                                     const int,
+                                     const psi::Psi<std::complex<double>>*,
+                                     const elecstate::ElecState*,
+                                     const K_Vectors&,
+                                     const TwoCenterIntegrator*,
+                                     const Parallel_Orbitals*,
+                                     const LCAO_Orbitals&,
+                                     const Velocity_op<double>*,
+                                     Record_adj&,
+                                     TD_info*);
+
+template void ctrl_output_td<std::complex<double>>(const UnitCell&,
+                                                   double**,
+                                                   const ModulePW::PW_Basis*,
+                                                   const int,
+                                                   const psi::Psi<std::complex<double>>*,
+                                                   const elecstate::ElecState*,
+                                                   const K_Vectors&,
+                                                   const TwoCenterIntegrator*,
+                                                   const Parallel_Orbitals*,
+                                                   const LCAO_Orbitals&,
+                                                   const Velocity_op<std::complex<double>>*,
+                                                   Record_adj&,
+                                                   TD_info*);
+
+} // namespace ModuleIO

source/source_io/ctrl_output_td.cpp

@@ -0,0 +1,35 @@
+#ifndef CTRL_OUTPUT_TD_H
+#define CTRL_OUTPUT_TD_H
+
+#include "source_basis/module_ao/ORB_read.h"
+#include "source_basis/module_ao/parallel_orbitals.h"
+#include "source_basis/module_nao/two_center_bundle.h"
+#include "source_cell/unitcell.h"
+#include "source_estate/elecstate_lcao.h"
+#include "source_estate/module_pot/H_TDDFT_pw.h"
+#include "source_lcao/module_rt/td_info.h"
+#include "source_lcao/module_rt/velocity_op.h"
+#include "source_lcao/record_adj.h"
+#include "source_psi/psi.h"
+
+namespace ModuleIO
+{
+
+template <typename TR>
+void ctrl_output_td(const UnitCell& ucell,
+                    double** rho_save,
+                    const ModulePW::PW_Basis* rhopw,
+                    const int istep,
+                    const psi::Psi<std::complex<double>>* psi,
+                    const elecstate::ElecState* pelec,
+                    const K_Vectors& kv,
+                    const TwoCenterIntegrator* intor,
+                    const Parallel_Orbitals* pv,
+                    const LCAO_Orbitals& orb,
+                    const Velocity_op<TR>* velocity_mat,
+                    Record_adj& RA,
+                    TD_info* td_p);
+
+} // namespace ModuleIO
+
+#endif // CTRL_OUTPUT_TD_H

source/source_io/ctrl_output_td.h

@@ -488,12 +488,12 @@ void ReadInput::check_ntype(const std::string& fn, int& param_ntype)
 int ReadInput::current_md_step(const std::string& file_dir)
 {
     std::stringstream ssc;
-    ssc << file_dir << "Restart_md.dat";
+    ssc << file_dir << "Restart_md.txt";
     std::ifstream file(ssc.str().c_str());
 
     if (!file)
     {
-        ModuleBase::WARNING_QUIT("current_md_step", "no Restart_md.dat");
+        ModuleBase::WARNING_QUIT("current_md_step", "no Restart_md.txt");
     }
 
     int md_step;

source/source_io/read_input.cpp

@@ -68,7 +68,7 @@ class ReadInput
     /**
      * @brief determine the md step in restart case
      *
-     * @param file_dir directory of Restart_md.dat
+     * @param file_dir directory of Restart_md.txt
      * @return md step
      */
     int current_md_step(const std::string& file_dir);