store X with ct::Tensor instead of Psi

source/module_lr/esolver_lrtd_lcao.cpp

@@ -136,6 +136,7 @@ void LR::ESolver_LR<T, TR>::reset_dim_spin2()
     if (nupdown == 0) { std::cout << "** Assuming the spin-up and spin-down states are degenerate. **" << std::endl; }
     else
     {
+        this->openshell = true;
         nupdown > 0 ? ((nocc[1] -= nupdown) && (nvirt[1] += nupdown)) : ((nocc[0] += nupdown) && (nvirt[0] -= nupdown));
         // npairs = { nocc[0] * nvirt[0], nocc[1] * nvirt[1] };
         std::cout << "** Solve the spin-up and spin-down states separately for open-shell system. **" << std::endl;
@@ -433,19 +434,19 @@ void LR::ESolver_LR<T, TR>::runner(int istep, UnitCell& cell)
     if (this->input.lr_solver != "spectrum")
     {
         // allocate and initialize A matrix and density matrix
-        std::vector<std::string> spin_type = { "singlet", "triplet" };
-        for (int is = 0;is < nspin;++is)
+        auto spin_types = (nspin == 2 && !openshell) ? std::vector<std::string>({ "singlet", "triplet" }) : std::vector<std::string>({ "updown" });
+        for (int is = 0;is < X.size();++is)
         {
-            if (nspin == 2) { std::cout << "Calculating " << spin_type[is] << " excitations" << std::endl; }
+            if (nspin == 2) { std::cout << "Calculating " << spin_types[is] << " excitations" << std::endl; }
             HamiltLR<T> hlr(xc_kernel, nspin, this->nbasis, this->nocc, this->nvirt, this->ucell, orb_cutoff_, GlobalC::GridD, this->psi_ks, this->eig_ks,
 #ifdef __EXX
                 this->exx_lri, this->exx_info.info_global.hybrid_alpha,
 #endif
                 this->gint_, this->pot[is], this->kv, this->paraX_, & this->paraC_, & this->paraMat_,
-                spin_type[is], input.ri_hartree_benchmark, (input.ri_hartree_benchmark == "aims" ? input.aims_nbasis : std::vector<int>({})));
+                spin_types[is], input.ri_hartree_benchmark, (input.ri_hartree_benchmark == "aims" ? input.aims_nbasis : std::vector<int>({})));
             // solve the Casida equation
             HSolverLR<T> hsol(nk, this->npairs[is], is, std::max(1e-13, this->input.lr_thr), this->input.out_wfc_lr);
-            hsol.solve(hlr, *this->X[is], this->pelec->ekb, this->input.lr_solver/*,
+            hsol.solve(hlr, this->X[is].template data<T>(), nloc_per_band, nstates, this->pelec->ekb, this->input.lr_solver/*,
                 !std::set<std::string>({ "hf", "hse" }).count(this->xc_kernel)*/);  //whether the kernel is Hermitian
         }
     }
@@ -477,14 +478,16 @@ void LR::ESolver_LR<T, TR>::after_all_runners()
     double lambda_diff = std::abs(abs_wavelen_range[1] - abs_wavelen_range[0]);
     double lambda_min = std::min(abs_wavelen_range[1], abs_wavelen_range[0]);
     for (int i = 0;i < freq.size();++i) { freq[i] = 91.126664 / (lambda_min + 0.01 * static_cast<double>(i + 1) * lambda_diff); }
+    auto spin_types = (nspin == 2 && !openshell) ? std::vector<std::string>({ "singlet", "triplet" }) : std::vector<std::string>({ "updown" });
     for (int is = 0;is < this->nspin;++is)
     {
-        LR_Spectrum<T> spectrum(&this->pelec->ekb.c[is * this->nstates], *this->X[is],
-            this->nspin, this->nbasis, this->nocc[is], this->nvirt[is], this->gint_, *this->pw_rho, *this->psi_ks,
-            this->ucell, this->kv, this->paraX_[is], this->paraC_, this->paraMat_);
-        spectrum.oscillator_strength();
-        spectrum.transition_analysis(is);
-        spectrum.optical_absorption(freq, input.abs_broadening, is);
+        const int is_in_x = openshell ? 0 : is;
+        LR_Spectrum<T> spectrum(this->nspin, this->nbasis, this->nocc[is], this->nvirt[is], this->gint_, *this->pw_rho, *this->psi_ks,
+            this->ucell, this->kv, this->paraX_[is], this->paraC_, this->paraMat_,
+            &this->pelec->ekb.c[is_in_x * nstates], this->X[is_in_x].template data<T>(),
+            nstates, nloc_per_band, /*offset=*/openshell ? nk * paraX_[0].get_local_size() : 0);
+        spectrum.transition_analysis(spin_types[is]);
+        spectrum.optical_absorption(freq, input.abs_broadening, spin_types[is]);
     }
 }
 
@@ -503,29 +506,24 @@ void LR::ESolver_LR<T, TR>::setup_eigenvectors_X()
         );//nvirt - row, nocc - col 
         this->paraX_.emplace_back(std::move(px));
     }
-    this->X.resize(this->nspin);
-    const std::vector<std::string> spin_types = { "singlet", "triplet" };
+    this->nloc_per_band = nk * (openshell ? paraX_[0].get_local_size() + paraX_[1].get_local_size() : paraX_[0].get_local_size());
+
+    this->X.resize(openshell ? 1 : nspin, LR_Util::newTensor<T>({ nstates, nloc_per_band }));
+    for (auto& x : X) { x.zero(); }
+
+    auto spin_types = (nspin == 2 && !openshell) ? std::vector<std::string>({ "singlet", "triplet" }) : std::vector<std::string>({ "updown" });
     // if spectrum-only, read the LR-eigenstates from file and return
     if (this->input.lr_solver == "spectrum")
     {
         std::cout << "reading the excitation amplitudes from file: \n";
-        for (int is = 0; is < this->nspin; ++is)
+        for (int is = 0; is < this->X.size(); ++is)
         {
-            this->X[is] = std::make_shared<psi::Psi<T>>(LR_Util::read_psi_bandfirst<T>(
-                PARAM.globalv.global_readin_dir + "Excitation_Amplitude_" + spin_types[is], GlobalV::MY_RANK));
+            std::ifstream ifs(PARAM.globalv.global_readin_dir + "Excitation_Amplitude_" + spin_types[is] + "_" + std::to_string(GlobalV::MY_RANK) + ".dat");
+            assert(nstates * nloc_per_band == LR_Util::read_value(ifs, X[is].data<T>(), nstates, nloc_per_band));
         }
     }
     else
     {
-        for (int is = 0; is < this->nspin; ++is)
-        {
-            this->X[is] = std::make_shared<psi::Psi<T>>(this->nk,
-                this->nstates,
-                this->paraX_[is].get_local_size(),
-                nullptr,
-                false); // band(state)-first
-            this->X[is]->zero_out();
-        }
         set_X_initial_guess();
     }
 }
@@ -544,7 +542,7 @@ void LR::ESolver_LR<T, TR>::set_X_initial_guess()
         // if (E_{lumo}-E_{homo-1} < E_{lumo+1}-E{homo}), mode = 0, else 1(smaller first)
         bool ix_mode = false;   //default
         if (this->eig_ks.nc > no + 1 && no >= 2 && eig_ks(is, no) - eig_ks(is, no - 2) - 1e-5 > eig_ks(is, no + 1) - eig_ks(is, no - 1)) { ix_mode = true; }
-        GlobalV::ofs_running << "setting the initial guess of X: " << std::endl;
+        GlobalV::ofs_running << "setting the initial guess of X of spin" << is << std::endl;
         if (no >= 2 && eig_ks.nc > no) { GlobalV::ofs_running << "E_{lumo}-E_{homo-1}=" << eig_ks(is, no) - eig_ks(is, no - 2) << std::endl; }
         if (no >= 1 && eig_ks.nc > no + 1) { GlobalV::ofs_running << "E_{lumo+1}-E{homo}=" << eig_ks(is, no + 1) - eig_ks(is, no - 1) << std::endl; }
         GlobalV::ofs_running << "mode of X-index: " << ix_mode << std::endl;
@@ -558,19 +556,21 @@ void LR::ESolver_LR<T, TR>::set_X_initial_guess()
         ioiv2ix = std::move(std::get<0>(indexmap));
         ix2ioiv = std::move(std::get<1>(indexmap));
 
-        for (int s = 0; s < nstates; ++s)
+        for (int ib = 0; ib < nstates; ++ib)
         {
-            this->X[is]->fix_b(s);
-            int ipair = s % np;
-            int occ_global = std::get<0>(ix2ioiv[ipair]);   // occ
-            int virt_global = std::get<1>(ix2ioiv[ipair]);   // virt
-            int ik = s / np;
+            const int ipair = ib % np;
+            const int occ_global = std::get<0>(ix2ioiv[ipair]);   // occ
+            const int virt_global = std::get<1>(ix2ioiv[ipair]);   // virt
+            const int ik = ib / np;
+            const int xstart_b = ib * nloc_per_band;    //start index of band ib
+            const int xstart_bs = (openshell && is == 1) ? xstart_b + nk * paraX_[0].get_local_size() : xstart_b;  // start index of band ib, spin is
+            const int is_in_x = openshell ? 1 : is;     // if openshell, spin-up and spin-down are put together
             if (px.in_this_processor(virt_global, occ_global))
-                (*X[is])(ik, px.global2local_col(occ_global) * px.get_row_size()
-                    + px.global2local_row(virt_global))
-                = (static_cast<T>(1.0) / static_cast<T>(nk));
+            {
+                const int ipair_loc = px.global2local_col(occ_global) * px.get_row_size() + px.global2local_row(virt_global);
+                X[is_in_x].data<T>()[xstart_bs + ipair_loc] = (static_cast<T>(1.0) / static_cast<T>(nk));
+            }
         }
-        this->X[is]->fix_b(0);  //recover the pointer
     }
 }
 
@@ -585,8 +585,8 @@ void LR::ESolver_LR<T, TR>::init_pot(const Charge& chg_gs)
         this->pot[0] = std::make_shared<PotHxcLR>(xc_kernel, this->pw_rho, &ucell, &chg_gs, PotHxcLR::SpinType::S1);
         break;
     case 2:
-        this->pot[0] = std::make_shared<PotHxcLR>(xc_kernel, this->pw_rho, &ucell, &chg_gs, nupdown ? PotHxcLR::SpinType::S2_up : PotHxcLR::SpinType::S2_singlet);
-        this->pot[1] = std::make_shared<PotHxcLR>(xc_kernel, this->pw_rho, &ucell, &chg_gs, nupdown ? PotHxcLR::SpinType::S2_down : PotHxcLR::SpinType::S2_triplet);
+        this->pot[0] = std::make_shared<PotHxcLR>(xc_kernel, this->pw_rho, &ucell, &chg_gs, openshell ? PotHxcLR::SpinType::S2_up : PotHxcLR::SpinType::S2_singlet);
+        this->pot[1] = std::make_shared<PotHxcLR>(xc_kernel, this->pw_rho, &ucell, &chg_gs, openshell ? PotHxcLR::SpinType::S2_down : PotHxcLR::SpinType::S2_triplet);
         break;
     default:
         throw std::invalid_argument("ESolver_LR: nspin must be 1 or 2");

source/module_lr/esolver_lrtd_lcao.h

@@ -65,8 +65,12 @@ namespace LR
         /// @brief ground state bands, read from the file, or moved from ESolver_FP::pelec.ekb
         ModuleBase::matrix eig_ks;///< energy of ground state
 
-        /// @brief Excited state wavefunction, size: [nspin][nstates * nk * (nocc(local) * nvirt (local))]
-        std::vector<std::shared_ptr<psi::Psi<T>>> X;
+        /// @brief Excited state wavefunction (locc, lvirt are local size of nocc and nvirt in each process)
+        /// size of X: [neq][{nstate, nloc_per_band}], namely:
+        /// - [nspin][{nstates, nk* (locc* lvirt}] for close- shell,
+        /// -  [1][{nstates, nk * (locc[0] * lvirt[0]) + nk * (locc[1] * lvirt[1])}] for open-shell
+        std::vector<ct::Tensor> X;
+        int nloc_per_band = 1;
 
         std::vector<int> nocc = { 1, 1 };   ///< number of occupied orbitals for each spin used in the calculation
         int nocc_in = 1;    ///< nocc read from input (adjusted by nelec): max(spin-up, spindown)
@@ -82,6 +86,7 @@ namespace LR
         int nspin = 1;
         int nk = 1;
         int nupdown = 0;
+        bool openshell = false;
         std::string xc_kernel;
 
         Grid_Technique gt_;

source/module_lr/hsolver_lrtd.cpp

@@ -19,20 +19,21 @@ namespace LR
     }
     template<typename T>
     void HSolverLR<T>::solve(const HamiltLR<T>& hm,
-        psi::Psi<T>& psi,
+        T* psi,
+        const int& dim,
+        const int& nband,
         ModuleBase::matrix& ekb,
         const std::string method,
         const bool hermitian)
     {
         ModuleBase::TITLE("HSolverLR", "solve");
-        assert(psi.get_nk() == nk);
         const std::vector<std::string> spin_types = { "singlet", "triplet" };
         // note: if not TDA, the eigenvalues will be complex
         // then we will need a new constructor of DiagoDavid
 
         // 1. allocate precondition and eigenvalue
-        std::vector<Real> precondition(psi.get_nk() * psi.get_nbasis());
-        std::vector<Real> eigenvalue(psi.get_nbands());   //nstates
+        std::vector<Real> precondition(dim);
+        std::vector<Real> eigenvalue(nband);   //nstates
         // 2. select the method
 #ifdef __MPI
         const hsolver::diag_comm_info comm_info = { POOL_WORLD, GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL };
@@ -52,25 +53,19 @@ namespace LR
                 print_eigs(eig_complex, "Right eigenvalues: of the non-Hermitian matrix: (Ry)");
                 for (int i = 0; i < nk * npairs; i++) { eigenvalue[i] = eig_complex[i].real(); }
             }
-            psi.fix_kb(0, 0);
             // copy eigenvectors
-            for (int i = 0;i < psi.size();++i) { psi.get_pointer()[i] = Amat_full[i]; }
+            std::memcpy(psi, Amat_full.data(), sizeof(T) * dim * nband);
         }
         else
         {
             // 3. set precondition and diagethr
-            for (int i = 0; i < psi.get_nk() * psi.get_nbasis(); ++i) { precondition[i] = static_cast<Real>(1.0); }
-
-            // wrap band-first psi as k1-first psi_k1_dav
-            psi::Psi<T> psi_k1_dav = LR_Util::bfirst_to_k1_wrapper(psi);
-            assert(psi_k1_dav.get_nbands() == psi.get_nbands());
-            assert(psi_k1_dav.get_nbasis() == psi.get_nbasis() * psi.get_nk());
+            for (int i = 0; i < dim; ++i) { precondition[i] = static_cast<Real>(1.0); }
 
             const int david_maxiter = hsolver::DiagoIterAssist<T>::PW_DIAG_NMAX;
 
             auto hpsi_func = [&hm](T* psi_in, T* hpsi, const int ld_psi, const int nvec) {hm.hPsi(psi_in, hpsi, ld_psi, nvec);};
-            auto spsi_func = [&hm](const T* psi_in, T* spsi, const int ld_psi, const int nbands)
-                { std::memcpy(spsi, psi_in, sizeof(T) * ld_psi * nbands); };
+            auto spsi_func = [&hm](const T* psi_in, T* spsi, const int ld_psi, const int nband)
+                { std::memcpy(spsi, psi_in, sizeof(T) * ld_psi * nband); };
 
             if (method == "dav")
             {
@@ -80,35 +75,33 @@ namespace LR
                 // converged.
                 const int notconv_max = ("nscf" == PARAM.inp.calculation) ? 0 : 5;
                 // do diag and add davidson iteration counts up to avg_iter
-                const int& dim = psi_k1_dav.get_nbasis();   //equals to leading dimension here
-                const int& nband = psi_k1_dav.get_nbands();
                 hsolver::DiagoDavid<T> david(precondition.data(), nband, dim, PARAM.inp.pw_diag_ndim, PARAM.inp.use_paw, comm_info);
                 hsolver::DiagoIterAssist<T>::avg_iter += static_cast<double>(david.diag(hpsi_func, spsi_func,
-                    dim, psi_k1_dav.get_pointer(), eigenvalue.data(), this->diag_ethr, david_maxiter, ntry_max, 0));
+                    dim, psi, eigenvalue.data(), this->diag_ethr, david_maxiter, ntry_max, 0));
             }
             else if (method == "dav_subspace") //need refactor
             {
                 hsolver::Diago_DavSubspace<T> dav_subspace(precondition,
-                    psi_k1_dav.get_nbands(),
-                    psi_k1_dav.get_nbasis(),
+                    nband,
+                    dim,
                     PARAM.inp.pw_diag_ndim,
                     this->diag_ethr,
                     david_maxiter,
                     false, //always do the subspace diag (check the implementation)
                     comm_info);
                 hsolver::DiagoIterAssist<T>::avg_iter
                     += static_cast<double>(dav_subspace.diag(
-                        hpsi_func, psi_k1_dav.get_pointer(),
-                        psi_k1_dav.get_nbasis(),
+                        hpsi_func, psi,
+                        dim,
                         eigenvalue.data(),
-                        std::vector<bool>(psi_k1_dav.get_nbands(), true),
+                        std::vector<bool>(nband, true),
                         false /*scf*/));
             }
             else {throw std::runtime_error("HSolverLR::solve: method not implemented");}
         }
 
         // 5. copy eigenvalue to pes
-        for (int ist = 0;ist < psi.get_nbands();++ist) { ekb(ispin_solve, ist) = eigenvalue[ist]; }
+        for (int ist = 0;ist < nband;++ist) { ekb(ispin_solve, ist) = eigenvalue[ist]; }
 
 
         // 6. output eigenvalues and eigenvectors
@@ -123,12 +116,12 @@ namespace LR
                 for (auto& e : eigenvalue) {ofs << e << " ";}
                 ofs.close();
             }
-            LR_Util::write_psi_bandfirst(psi, PARAM.globalv.global_out_dir + "Excitation_Amplitude_" + spin_types[ispin_solve], GlobalV::MY_RANK);
+            LR_Util::write_psi_bandfirst(psi, nband, nk, npairs, PARAM.globalv.global_out_dir + "Excitation_Amplitude_" + spin_types[ispin_solve], GlobalV::MY_RANK);
         }
 
         // normalization is already satisfied
         // std::cout << "check normalization of eigenvectors:" << std::endl;
-        // for (int ist = 0;ist < psi.get_nbands();++ist)
+        // for (int ist = 0;ist < nband;++ist)
         // {
         //     double norm2 = 0;
         //     for (int ik = 0;ik < psi.get_nk();++ik)

source/module_lr/hsolver_lrtd.h

@@ -1,7 +1,6 @@
 #pragma once
 #include "module_lr/hamilt_casida.h"
 #include "module_hsolver/diago_iter_assist.h"
-#include "module_psi/psi.h"
 namespace LR
 {
     template<typename T>
@@ -18,7 +17,9 @@ namespace LR
 
         /// eigensolver for common Hamilt
         void solve(const HamiltLR<T>& hm,
-            psi::Psi<T>& psi,
+            T* psi,
+            const int& dim, ///< local leading dimension (or nbasis)
+            const int& nband,   ///< nstates in LR-TDDFT, not (nocc+nvirt)
             ModuleBase::matrix& ekb,
             const std::string method_in,
             const bool hermitian = true);