From 090bd1a215fb8839cda498e8ba2e47af45d7d4c2 Mon Sep 17 00:00:00 2001
From: maki49 <1579492865@qq.com>
Date: Sun, 14 Jul 2024 03:29:46 +0800
Subject: [PATCH] truncate nocc for KS-LR

---
 source/module_lr/esolver_lrtd_lcao.cpp | 32 ++++++++++++++++++--------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/source/module_lr/esolver_lrtd_lcao.cpp b/source/module_lr/esolver_lrtd_lcao.cpp
index 194efbe21b..06fedbceef 100644
--- a/source/module_lr/esolver_lrtd_lcao.cpp
+++ b/source/module_lr/esolver_lrtd_lcao.cpp
@@ -12,6 +12,7 @@
 #include "module_io/print_info.h"
 #include "module_cell/module_neighbor/sltk_atom_arrange.h"
 #include "module_lr/utils/lr_util_print.h"
+#include "module_base/scalapack_connector.h"
 
 #ifdef __EXX
 template<>
@@ -123,22 +124,35 @@ LR::ESolver_LR<T, TR>::ESolver_LR(ModuleESolver::ESolver_KS_LCAO<T, TR>&& ks_sol
 
     this->parameter_check();
 
-    // set up the 2d division for nbasis*nbasis matrices
-    this->nbasis = GlobalV::NLOCAL;
+    this->set_dimension();
 
     LR_Util::setup_2d_division(this->paraMat_, 1, this->nbasis, this->nbasis);
     this->paraMat_.atom_begin_row = std::move(ks_sol.ParaV.atom_begin_row);
     this->paraMat_.atom_begin_col = std::move(ks_sol.ParaV.atom_begin_col);
     this->paraMat_.iat2iwt_ = ucell.get_iat2iwt();
 
-    // move the ground state info 
-    this->psi_ks = ks_sol.psi;
-    ks_sol.psi = nullptr;
-    //only need the eigenvalues. the 'elecstates' of excited states is different from ground state.
-    this->eig_ks = std::move(ks_sol.pelec->ekb);
+    LR_Util::setup_2d_division(this->paraC_, 1, this->nbasis, this->nbands, this->paraMat_.blacs_ctxt);
 
-    this->set_dimension();
-    LR_Util::setup_2d_division(this->paraC_, 1, this->nbasis, this->nocc + this->nvirt, this->paraMat_.blacs_ctxt);
+    if (this->nbands == GlobalV::NBANDS)    // move the ground state info 
+    {
+        this->psi_ks = ks_sol.psi;
+        ks_sol.psi = nullptr;
+        //only need the eigenvalues. the 'elecstates' of excited states is different from ground state.
+        this->eig_ks = std::move(ks_sol.pelec->ekb);
+    }
+    else    // copy the part of ground state info according to paraC_
+    {
+        this->psi_ks = new psi::Psi<T>(this->kv.get_nks(), this->paraC_.get_col_size(), this->paraC_.get_row_size());
+        this->eig_ks.create(this->kv.get_nks(), this->nbands);
+        const int start_band = this->nocc_max - this->nocc;
+        for (int ik = 0;ik < this->kv.get_nks();++ik)
+        {
+            Cpxgemr2d(this->nbasis, this->nbands, &(*ks_sol.psi)(ik, 0, 0), 1, start_band + 1, ks_sol.ParaV.desc_wfc,
+                &(*this->psi_ks)(ik, 0, 0), 1, 1, this->paraC_.desc, this->paraC_.blacs_ctxt);
+            for (int ib = 0;ib < this->nbands;++ib)
+                this->eig_ks(ik, ib) = ks_sol.pelec->ekb(ik, start_band + ib);
+        }
+    }
 
     //grid integration
     this->gt_ = std::move(ks_sol.GridT);