From 708608996bb94bfe7ce4a58a8577438786334d2e Mon Sep 17 00:00:00 2001
From: wenfei-li <liwenfei@gmail.com>
Date: Tue, 16 Jan 2024 16:28:54 +0800
Subject: [PATCH] Feature : printing band density

---
 source/module_esolver/esolver_ks_pw.cpp | 52 +++++++++++++++++++++++++
 source/module_io/input.cpp              | 44 +++++++++++++++++++++
 source/module_io/input.h                |  2 +
 3 files changed, 98 insertions(+)

diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
index 294f1636f7..58ff6e76c3 100644
--- a/source/module_esolver/esolver_ks_pw.cpp
+++ b/source/module_esolver/esolver_ks_pw.cpp
@@ -894,6 +894,58 @@ void ESolver_KS_PW<T, Device>::afterscf(const int istep)
                             this->kspw_psi[0].get_pointer() - this->kspw_psi[0].get_psi_bias(),
                             this->psi[0].size());
     }
+
+    if(INPUT.band_print_num > 0)
+    {
+        std::complex<double> * wfcr = new std::complex<double>[this->pw_rho->nxyz];
+        double * rho_band = new double [this->pw_rho->nxyz];
+        for(int i = 0; i < this->pw_rho->nxyz; i++)
+        {
+            rho_band[i] = 0.0;
+        }
+
+        for(int i = 0; i < INPUT.band_print_num; i++)
+        {
+            int ib = INPUT.bands_to_print[i];
+            for(int ik = 0; ik < this->kv.nks; ik++)
+            {
+                this->psi->fix_k(ik);
+                this->pw_wfc->recip_to_real(this->ctx,&psi[0](ib,0),wfcr,ik);
+
+                double w1 = static_cast<double>(this->kv.wk[ik] / GlobalC::ucell.omega);
+
+                for(int i = 0; i < this->pw_rho->nxyz; i++)
+                {
+                    rho_band[i] += std::norm(wfcr[i]) * w1;
+                }
+            }
+
+            std::stringstream ssc;
+            ssc << GlobalV::global_out_dir << "band" << ib << ".cube";     
+
+            ModuleIO::write_rho
+            (
+#ifdef __MPI
+                this->pw_big->bz,
+                this->pw_big->nbz,
+                this->pw_big->nplane,
+                this->pw_big->startz_current,
+#endif
+                rho_band,
+                0,
+                GlobalV::NSPIN,
+                0,
+                ssc.str(),
+                this->pw_rho->nx,
+                this->pw_rho->ny,
+                this->pw_rho->nz,
+                0.0,
+                &(GlobalC::ucell),
+                11);
+        }
+        delete[] wfcr;
+        delete[] rho_band;
+    }
 }
 
 template <typename T, typename Device>
diff --git a/source/module_io/input.cpp b/source/module_io/input.cpp
index 2685ba540e..2f7db3457f 100644
--- a/source/module_io/input.cpp
+++ b/source/module_io/input.cpp
@@ -327,6 +327,8 @@ void Input::Default(void)
 
     out_bandgap = 0; // QO added for bandgap printing
 
+    band_print_num = 0;
+
     deepks_out_labels = 0; // caoyu added 2020-11-24, mohan added 2021-01-03
     deepks_scf = 0;
     deepks_bandgap = 0;
@@ -1327,6 +1329,14 @@ bool Input::Read(const std::string& fn)
         {
             read_bool(ifs, out_chg);
         }
+        else if (strcmp("band_print_num", word) == 0)
+        {
+            read_value(ifs, band_print_num);
+        }
+        else if (strcmp("bands_to_print", word) == 0)
+        {
+            ifs.ignore(150, '\n');
+        }
         else if (strcmp("out_dm", word) == 0)
         {
             read_bool(ifs, out_dm);
@@ -2369,6 +2379,29 @@ bool Input::Read(const std::string& fn)
         ModuleBase::WARNING_QUIT("Input", "The ntype in INPUT is not equal to the ntype counted in STRU, check it.");
     }
 
+    if(band_print_num > 0)
+    {
+        bands_to_print.resize(band_print_num);
+        ifs.clear();
+        ifs.seekg(0); // move to the beginning of the file
+        ifs.rdstate();
+        while (ifs.good())
+        {
+            ifs >> word1;
+            if (ifs.eof() != 0)
+                break;
+            strtolower(word1, word); // convert uppercase std::string to lower case; word1 --> word
+
+            if (strcmp("bands_to_print", word) == 0)
+            {
+                for(int i = 0; i < band_print_num; i ++)
+                {
+                    ifs >> bands_to_print[i];
+                }
+            }
+        }
+    }
+
     //----------------------------------------------------------
     //       DFT+U    Xin Qu  added on 2020-10-29
     //----------------------------------------------------------
@@ -3523,6 +3556,17 @@ void Input::Bcast()
     Parallel_Common::bcast_bool(restart_save);  // Peize Lin add 2020.04.04
     Parallel_Common::bcast_bool(restart_load);  // Peize Lin add 2020.04.04
 
+    Parallel_Common::bcast_int(band_print_num);
+    if(GlobalV::MY_RANK != 0)
+    {
+        bands_to_print.resize(band_print_num);
+    }
+
+    for(int i = 0; i < band_print_num; i++)
+    {
+        Parallel_Common::bcast_int(bands_to_print[i]);
+    }
+
     //-----------------------------------------------------------------------------------
     // DFT+U (added by Quxin 2020-10-29)
     //-----------------------------------------------------------------------------------
diff --git a/source/module_io/input.h b/source/module_io/input.h
index b4e983abad..f5aaacc541 100644
--- a/source/module_io/input.h
+++ b/source/module_io/input.h
@@ -259,6 +259,8 @@ class Input
     bool out_chg; // output charge density. 0: no; 1: yes
     bool out_dm; // output density matrix.
     bool out_dm1;
+    int band_print_num;
+    std::vector<int> bands_to_print;
     int out_pot; // yes or no
     int out_wfc_pw; // 0: no; 1: txt; 2: dat
     bool out_wfc_r; // 0: no; 1: yes