Fix energy jumps occur in rt-TDDFT calculations

source/Makefile.Objects

@@ -622,6 +622,7 @@ OBJS_LCAO=evolve_elec.o\
       velocity_op.o\
       snap_psibeta_half_tddft.o\
       solve_propagation.o\
+      boundary_fix.o\
       upsi.o\
       FORCE_STRESS.o\
 	  FORCE_gamma.o\

source/source_cell/atom_spec.cpp

@@ -101,6 +101,7 @@ void Atom::bcast_atom()
             this->tau.resize(na, ModuleBase::Vector3<double>(0, 0, 0));
             this->dis.resize(na, ModuleBase::Vector3<double>(0, 0, 0));
             this->taud.resize(na, ModuleBase::Vector3<double>(0, 0, 0));
+            this->boundary_shift.resize(na, ModuleBase::Vector3<int>(0, 0, 0));
             this->vel.resize(na, ModuleBase::Vector3<double>(0, 0, 0));
             this->mag.resize(na, 0);
             this->angle1.resize(na, 0);

source/source_cell/atom_spec.h

@@ -36,6 +36,7 @@ class Atom
     std::vector<ModuleBase::Vector3<double>> tau; // Cartesian coordinates of each atom in this type.
     std::vector<ModuleBase::Vector3<double>> dis; // direct displacements of each atom in this type in current step  liuyu modift 2023-03-22
     std::vector<ModuleBase::Vector3<double>> taud;  // Direct coordinates of each atom in this type.
+    std::vector<ModuleBase::Vector3<int>> boundary_shift;  // record for periodic boundary adjustment.
     std::vector<ModuleBase::Vector3<double>> vel;   // velocities of each atom in this type.
     std::vector<ModuleBase::Vector3<double>> force; // force acting on each atom in this type.
     std::vector<ModuleBase::Vector3<double>> lambda; // Lagrange multiplier for each atom in this type. used in deltaspin

source/source_cell/read_atoms.cpp

@@ -163,6 +163,7 @@ bool unitcell::read_atom_positions(UnitCell& ucell,
                 ucell.atoms[it].tau.resize(na, ModuleBase::Vector3<double>(0,0,0));
                 ucell.atoms[it].dis.resize(na, ModuleBase::Vector3<double>(0,0,0));
                 ucell.atoms[it].taud.resize(na, ModuleBase::Vector3<double>(0,0,0));
+                ucell.atoms[it].boundary_shift.resize(na, ModuleBase::Vector3<int>(0,0,0));
                 ucell.atoms[it].vel.resize(na, ModuleBase::Vector3<double>(0,0,0));
                 ucell.atoms[it].mbl.resize(na, ModuleBase::Vector3<int>(0,0,0));
                 ucell.atoms[it].mag.resize(na, 0);

source/source_cell/update_cell.cpp

@@ -482,17 +482,20 @@ void periodic_boundary_adjustment(Atom* atoms,
 	for (int it = 0; it < ntype; it++) 
 	{
 		Atom* atom = &atoms[it];
+        atom->boundary_shift.assign(atom->na, {0,0,0});
 		for (int ia = 0; ia < atom->na; ia++) 
 		{
             // mohan update 2011-03-21
             for (int ik = 0; ik < 3; ik++) 
             {
                 if (atom->taud[ia][ik] < 0) 
                 {
+                    atom->boundary_shift[ia][ik] += 1;
                     atom->taud[ia][ik] += 1.0;
                 }
                 if (atom->taud[ia][ik] >= 1.0) 
                 {
+                    atom->boundary_shift[ia][ik] -= 1;
                     atom->taud[ia][ik] -= 1.0;
                 }
             }

source/source_esolver/esolver_ks_lcao_tddft.cpp

@@ -107,7 +107,15 @@ void ESolver_KS_LCAO_TDDFT<TR, Device>::runner(UnitCell& ucell, const int istep)
     {
         estep_max = PARAM.inp.estep_per_md + 1;
     }
-    // int estep_max = PARAM.inp.estep_per_md;
+    // reset laststep matrix and wfc, if any atom cross the boundary
+    const size_t len_hs_ik = use_tensor && use_lapack ? PARAM.globalv.nlocal * PARAM.globalv.nlocal : this->pv.nloc;
+    module_rt::reset_matrix_boundary(ucell,
+                                     this->kv,
+                                     &(this->pv),
+                                     this->Hk_laststep,
+                                     this->Sk_laststep,
+                                     this->psi_laststep,
+                                     len_hs_ik);
     for (int estep = 0; estep < estep_max; estep++)
     {
         // calculate total time step

source/source_esolver/esolver_ks_lcao_tddft.h

@@ -6,6 +6,7 @@
 #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"
+#include "source_lcao/module_rt/boundary_fix.h"
 
 namespace ModuleESolver
 {

source/source_lcao/module_rt/CMakeLists.txt

@@ -15,6 +15,7 @@ if(ENABLE_LCAO)
         snap_psibeta_half_tddft.cpp
         td_folding.cpp
         solve_propagation.cpp
+        boundary_fix.cpp
     )
 
     add_library(

source/source_lcao/module_rt/boundary_fix.cpp

@@ -0,0 +1,99 @@
+#include "boundary_fix.h"
+#include "source_base/libm/libm.h"
+#include "source_base/constants.h"
+#include "source_base/vector3.h"
+
+namespace module_rt{
+
+void reset_matrix_boundary(const UnitCell& ucell,
+                           const K_Vectors& kv,
+                           const Parallel_Orbitals* pv,
+                           ct::Tensor& hk_last,
+                           ct::Tensor& sk_last,
+                           psi::Psi<std::complex<double>>* psi_last,
+                           const size_t len_hs)
+{
+    ModuleBase::TITLE("module_rt", "reset_matrix_boundary");
+    ModuleBase::timer::tick("module_rt", "reset_matrix_boundary");
+    const ModuleBase::Vector3<int> zero = {0, 0, 0};
+    for(size_t iat = 0; iat < ucell.nat; iat++)
+    {
+        const size_t it = ucell.iat2it[iat];
+        const size_t ia = ucell.iat2ia[iat];
+        if(ucell.atoms[it].boundary_shift[ia]!=zero)
+        {
+            const auto& rshift = ucell.atoms[it].boundary_shift[ia];
+#ifdef _OPENMP
+#pragma omp parallel for schedule(dynamic)
+#endif
+            for(int ik = 0; ik < kv.get_nks(); ik++)
+            {
+                const ModuleBase::Vector3<double> tmp_rshift(rshift.x, rshift.y, rshift.z);
+                const double arg = -kv.kvec_d[ik] * tmp_rshift * ModuleBase::TWO_PI;
+                //skip unrelevent ik
+                if(arg==0)continue;
+                //calculate correction phase
+                double sinp, cosp;
+                ModuleBase::libm::sincos(arg, &sinp, &cosp);
+                const std::complex<double> phase = std::complex<double>(cosp, sinp);
+                //phase correction for Hamiltionian, overlap matrix and c vec.
+                module_rt::boundary_shift_mat(phase, hk_last.template data<std::complex<double>>() + ik * len_hs, pv, iat);
+                module_rt::boundary_shift_mat(phase, sk_last.template data<std::complex<double>>() + ik * len_hs, pv, iat);
+                psi_last->fix_k(ik);
+                module_rt::boundary_shift_c(phase, psi_last[0].get_pointer(), pv, iat);
+            }
+        }
+    }
+    ModuleBase::timer::tick("module_rt", "reset_matrix_boundary");
+    return;
+}
+
+void boundary_shift_mat(const std::complex<double>& phase,
+                        std::complex<double>* matk,
+                        const Parallel_Orbitals* pv,
+                        const size_t iat)
+{
+    const std::complex<double> phase_conj = std::conj(phase);
+    size_t row0 = pv->atom_begin_row[iat];
+    size_t col0 = pv->atom_begin_col[iat];
+    std::complex<double>* p_matkc = matk + col0 * pv->get_row_size();
+    for(size_t nu = 0; nu < pv->get_col_size(iat); ++nu)
+    {
+
+        BlasConnector::scal(pv->get_row_size(),
+                            phase,
+                            p_matkc,
+                            1);
+        p_matkc += pv->get_row_size();
+    }
+    std::complex<double>* p_matkr = matk + row0;
+    for(size_t mu = 0; mu < pv->get_row_size(iat); ++mu)
+    {
+        BlasConnector::scal(pv->get_col_size(),
+                            phase_conj,
+                            p_matkr,
+                            pv->get_row_size());
+        p_matkr += 1;
+    }
+    return;
+}
+
+void boundary_shift_c(const std::complex<double>& phase,
+                      std::complex<double>* psi_k_last,
+                      const Parallel_Orbitals* pv,
+                      const size_t iat)
+{
+    const std::complex<double> phase_conj = std::conj(phase);
+    size_t row0 = pv->atom_begin_row[iat];
+    std::complex<double>* p_ck = psi_k_last + row0;
+    for(size_t nu = 0; nu < pv->get_row_size(iat); ++nu)
+    {
+        BlasConnector::scal(pv->ncol_bands,
+                            phase_conj,
+                            p_ck,
+                            pv->get_row_size());
+        p_ck+=1;
+    }
+    return;
+}
+} //namespace module_rt

source/source_lcao/module_rt/boundary_fix.h

@@ -0,0 +1,64 @@
+/**
+ * @file boundary_fix.h
+ * @brief Correct the discontinuity that occurs when crossing periodic boundary conditions
+ */
+#ifndef BOUNDARY_FIX_H
+#define BOUNDARY_FIX_H
+
+#include "source_cell/unitcell.h"
+#include "source_cell/klist.h"
+#include "source_basis/module_ao/parallel_orbitals.h"
+#include "source_psi/psi.h"
+#include "source_base/module_container/ATen/core/tensor.h"
+namespace module_rt{
+
+/**
+*  @brief Add phases to the matrix and coefficient from the previous step to correct the boundary discontinuity.
+*
+* @param[in] ucell Unitcell information
+* @param[in] kv K-point vectors
+* @param[in] pv information of parallel
+* @param[in] hk_last Hamiltonian matrix from last step 
+* @param[in] sk_last Overlap matrix from last step 
+* @param[in] psi_last Wavefunctions from last step
+* @param[in] len_hs size of matrix element in this processor
+* @param[out] hk_last the fixed hk matrix
+* @param[out] sk_last the fixed sk matrix
+* @param[out] sk_last the fixed wavefunctions
+*/
+void reset_matrix_boundary(const UnitCell& ucell,
+                           const K_Vectors& kv,
+                           const Parallel_Orbitals* pv,
+                           ct::Tensor& hk_last,
+                           ct::Tensor& sk_last,
+                           psi::Psi<std::complex<double>>* psi_last,
+                           const size_t len_hs);
+
+/**
+*  @brief Add extra phase to the matrix element belong to iat
+*
+* @param[in] phase extra phase
+* @param[in] matk the matrix need to be fixed
+* @param[in] pv information of parallel
+* @param[in] iat atom index
+* @param[out] matk the fixed matrix
+*/
+void boundary_shift_mat(const std::complex<double>& phase,
+                        std::complex<double>* matk,
+                        const Parallel_Orbitals* pv,
+                        const size_t iat);
+/**
+*  @brief Add extra phase to the wfc coefficient belong to iat
+*
+* @param[in] phase extra phase
+* @param[in] psi_k_last psi of last step
+* @param[in] pv information of parallel
+* @param[in] iat atom index
+* @param[out] psi_k_last fixed psi of last step
+*/
+void boundary_shift_c(const std::complex<double>& phase,
+                      std::complex<double>* psi_k_last,
+                      const Parallel_Orbitals* pv,
+                      const size_t iat);
+}// namespace module_rt
+#endif // BOUNDARY_FIX_H