Feature: Add INPUT parameter bands_to_print for LCAO basis set (#3934)

* Add INPUT param out_band_index

* Add INPUT param out_band_index

* Add INPUT param out_band_index

* Fix extra spacing within  in doc

* Fix several bugs of band-decomposed charge densities using PW basis set

* Change INPUT param from out_band_index to bands_to_print and modified the corresponding docs

---------

Co-authored-by: Mohan Chen <mohan.chen.chen.mohan@gmail.com>

docs/advanced/input_files/input-main.md

@@ -64,7 +64,6 @@
     - [basis\_type](#basis_type)
     - [ks\_solver](#ks_solver)
     - [nbands](#nbands)
-    - [nbands\_istate](#nbands_istate)
     - [nspin](#nspin)
     - [smearing\_method](#smearing_method)
     - [smearing\_sigma](#smearing_sigma)
@@ -148,12 +147,12 @@
     - [out\_app\_flag](#out_app_flag)
     - [out\_ndigits](#out_ndigits)
     - [out\_interval](#out_interval)
-    - [band\_print\_num](#band_print_num)
-    - [bands\_to\_print](#bands_to_print)
     - [out\_element\_info](#out_element_info)
     - [restart\_save](#restart_save)
     - [restart\_load](#restart_load)
     - [rpa](#rpa)
+    - [nbands\_istate](#nbands_istate)
+    - [bands\_to\_print](#bands_to_print)
   - [Density of states](#density-of-states)
     - [dos\_edelta\_ev](#dos_edelta_ev)
     - [dos\_sigma](#dos_sigma)
@@ -413,7 +412,7 @@ These variables are used to control general system parameters.
   - relax: do structure relaxation calculation, one can use `relax_nmax` to decide how many ionic relaxations you want
   - cell-relax: do variable-cell relaxation calculation
   - nscf: do the non self-consistent electronic structure calculations. For this option, you need a charge density file. For nscf calculations with planewave basis set, pw_diag_thr should be <= 1e-3
-  - get_pchg: For LCAO basis. Please see the explanation for variable `nbands_istate`
+  - get_pchg: For LCAO basis. Please see the explanation for variable `nbands_istate` and `bands_to_print`
   - get_wf: Envelope function for LCAO basis. Please see the explanation for variable `nbands_istate`
   - md: molecular dynamics
   - test_memory : checks memory required for the calculation. The number is not quite reliable, please use it with care
@@ -937,13 +936,6 @@ calculations.
   - nspin=2: max(1.2\*nelec_spin, nelec_spin + 10), in which nelec_spin = max(nelec_spin_up, nelec_spin_down)
   - nspin=4: max(1.2\*nelec, nelec + 20)
 
-### nbands_istate
-
-- **Type**: Integer
-- **Availability**: Only used when `calculation = get_wf` or `calculation = get_pchg`.
-- **Description**: The number of bands around the Fermi level you would like to calculate. `get_wf` means to calculate the envelope functions of wave functions $\Psi_{i}=\Sigma_{\mu}C_{i\mu}\Phi_{\mu}$, where $\Psi_{i}$ is the ith wave function with the band index $i$ and $\Phi_{\mu}$ is the localized atomic orbital set. `get_pchg` means to calculate the density of each wave function $|\Psi_{i}|^{2}$. Specifically, suppose we have highest occupied bands at 100th wave functions. And if you set this variable to 5, it will print five wave functions from 96th to 105th. But before all this can be carried out, the wave functions coefficients  should be first calculated and written into a file by setting the flag `out_wfc_lcao = 1`.
-- **Default**: 5
-
 ### nspin
 
 - **Type**: Integer
@@ -1635,20 +1627,6 @@ These variables are used to control the output of properties.
 - **Description**: Control the interval for printing Mulliken population analysis, $r(R)$, $H(R)$, $S(R)$, $T(R)$, $dH(R)$, $H(k)$, $S(k)$ and $wfc(k)$ matrices during molecular dynamics calculations. Check input parameters [out_mul](#out_mul), [out_mat_r](#out_mat_r), [out_mat_hs2](#out_mat_hs2), [out_mat_t](#out_mat_t), [out_mat_dh](#out_mat_dh), [out_mat_hs](#out_mat_hs) and [out_wfc_lcao](#out_wfc_lcao) for more information, respectively.
 - **Default**: 1
 
-### band_print_num
-
-- **Type**: Integer
-- **Availability**: PW basis
-- **Description**: If you want to plot a partial charge density contributed from some chosen bands. `band_print_num` define the number of band list. The result can be found in "band*.cube".
-- **Default**: 0
-
-### bands_to_print
-
-- **Type**: vector
-- **Availability**: band_print_num > 0
-- **Description**: define which band you want to choose for partial charge density.
-- **Default**: []
-
 ### out_element_info
 
 - **Type**: Boolean
@@ -1681,6 +1659,20 @@ These variables are used to control the output of properties.
 - **Description**: Generate output files used in rpa calculations.
 - **Default**: False
 
+### nbands_istate
+
+- **Type**: Integer
+- **Availability**: Only for LCAO, used when `calculation = get_wf` or `calculation = get_pchg`.
+- **Description**: The number of bands around the Fermi level you would like to calculate. `get_wf` means to calculate the envelope functions of wave functions $\Psi_{i}=\Sigma_{\mu}C_{i\mu}\Phi_{\mu}$, where $\Psi_{i}$ is the ith wave function with the band index $i$ and $\Phi_{\mu}$ is the localized atomic orbital set. `get_pchg` means to calculate the density of each wave function $|\Psi_{i}|^{2}$. Specifically, suppose we have highest occupied bands at 100th wave functions. And if you set this variable to 5, it will print five wave functions from 96th to 105th. But before all this can be carried out, the wave functions coefficients should be first calculated and written into a file by setting the flag `out_wfc_lcao = 1`.
+- **Default**: 5
+
+### bands_to_print
+
+- **Type**: String
+- **Availability**: For both PW and LCAO. When `basis_type = lcao`, only used when `calculation = get_pchg`.
+- **Description**: Specifies the bands to calculate the charge density for, using a space-separated string of 0s and 1s, providing a more flexible selection compared to `nbands_istate`. Each digit in the string corresponds to a band, starting from the first band. A `1` indicates that the charge density should be calculated for that band, while a `0` means the band will be ignored. The parameter allows a compact and flexible notation (similar to [`ocp_set`](#ocp_set)), for example the syntax `1 4*0 5*1 0` is used to denote the selection of bands: `1` means calculate for the first band, `4*0` skips the next four bands, `5*1` means calculate for the following five bands, and the final `0` skips the next band. It's essential that the total count of bands does not exceed the total number of bands (`nbands`); otherwise, it results in an error, and the process exits. The input string must contain only numbers and the asterisk (`*`) for repetition, ensuring correct format and intention of band selection.
+- **Default**: none
+
 [back to top](#full-list-of-input-keywords)
 
 ## Density of states

source/module_esolver/esolver_ks_lcao_elec.cpp

@@ -408,6 +408,7 @@ void ESolver_KS_LCAO<TK, TR>::others(const int istep)
                   GlobalV::NBANDS,
                   GlobalV::nelec,
                   GlobalV::NSPIN,
+                  GlobalV::NLOCAL,
                   GlobalV::global_out_dir,
                   GlobalV::MY_RANK,
                   GlobalV::ofs_warning);

source/module_esolver/esolver_ks_pw.cpp

@@ -7,6 +7,7 @@
 #include "module_io/write_istate_info.h"
 #include "module_io/write_wfc_pw.h"
 #include "module_io/output_log.h"
+#include "module_io/input_conv.h"
 
 //--------------temporary----------------------------
 #include "module_elecstate/module_charge/symmetry_rho.h"
@@ -1023,36 +1024,81 @@ void ESolver_KS_PW<T, Device>::after_scf(const int istep)
                             this->psi[0].size());
     }
 
-    if(INPUT.band_print_num > 0)
+    // Get bands_to_print through public function of INPUT (returns a const pointer to string)
+    std::string bands_to_print = *INPUT.get_bands_to_print();
+    if(!bands_to_print.empty())
     {
-        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++)
+        std::vector<double> out_band_kb;
+        Input_Conv::parse_expression(bands_to_print, out_band_kb);
+
+        // bands_picked is a vector of 0s and 1s, where 1 means the band is picked to output
+        std::vector<int> bands_picked;
+        bands_picked.resize(this->kspw_psi->get_nbands());
+        ModuleBase::GlobalFunc::ZEROS(bands_picked.data(), this->kspw_psi->get_nbands());
+
+        // Check if length of out_band_kb is valid
+        if (static_cast<int>(out_band_kb.size()) > this->kspw_psi->get_nbands())
+        {
+            ModuleBase::WARNING_QUIT(
+                "ESolver_KS_PW::after_scf",
+                "The number of bands specified by `bands_to_print` in the INPUT file exceeds `nbands`!");
+        }
+
+        // Check if all elements in bands_picked are 0 or 1
+        for (int value: out_band_kb)
+        {
+            if (value != 0 && value != 1)
+            {
+                ModuleBase::WARNING_QUIT(
+                    "ESolver_KS_PW::after_scf",
+                    "The elements of `bands_to_print` must be either 0 or 1. Invalid values found!");
+            }
+        }
+
+        // Fill bands_picked with values from out_band_kb, converting to int
+        // Remaining bands are already set to 0
+        int length = std::min(static_cast<int>(out_band_kb.size()), this->kspw_psi->get_nbands());
+        for (int i = 0; i < length; ++i)
         {
-            rho_band[i] = 0.0;
+            // out_band_kb rely on function parse_expression from input_conv.cpp
+            // Initially designed for ocp_set, which can be double
+            bands_picked[i] = static_cast<int>(out_band_kb[i]);
         }
 
-        for(int i = 0; i < INPUT.band_print_num; i++)
+        std::complex<double>* wfcr = new std::complex<double>[this->pw_rho->nxyz];
+        double* rho_band = new double[this->pw_rho->nxyz];
+
+        for (int ib = 0; ib < this->kspw_psi->get_nbands(); ++ib)
         {
-            int ib = INPUT.bands_to_print[i];
-            for(int ik = 0; ik < this->kv.nks; ik++)
+            // Skip the loop iteration if bands_picked[ib] is 0
+            if (!bands_picked[ib])
+            {
+                continue;
+            } 
+
+            for (int i = 0; i < this->pw_rho->nxyz; i++)
+            {
+                // Initialize rho_band to zero for each band
+                rho_band[i] = 0.0;
+            }
+
+            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);
+                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++)
+                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";     
+            ssc << GlobalV::global_out_dir << "band" << ib + 1 << ".cube"; // band index starts from 1
 
-            ModuleIO::write_rho
-            (
+            ModuleIO::write_rho(
 #ifdef __MPI
                 this->pw_big->bz,
                 this->pw_big->nbz,

source/module_io/input.cpp

@@ -162,6 +162,7 @@ void Input::Default(void)
     nbands_sto = 256;
     nbndsto_str = "256";
     nbands_istate = 5;
+    bands_to_print_ = "";
     pw_seed = 1;
     emin_sto = 0.0;
     emax_sto = 0.0;
@@ -332,8 +333,6 @@ 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;
@@ -784,6 +783,10 @@ bool Input::Read(const std::string& fn)
             // if (nbands_istate < 0)
             // 	ModuleBase::WARNING_QUIT("Input", "NBANDS_ISTATE must > 0");
         }
+        else if (strcmp("bands_to_print", word) == 0)
+        {
+            getline(ifs, bands_to_print_);
+        }
         else if (strcmp("nche_sto", word) == 0) // Chebyshev expansion order
         {
             read_value(ifs, nche_sto);
@@ -1351,14 +1354,6 @@ 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);
@@ -2416,29 +2411,6 @@ 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
     //----------------------------------------------------------
@@ -3256,6 +3228,7 @@ void Input::Bcast()
     Parallel_Common::bcast_int(nbands);
     Parallel_Common::bcast_int(nbands_sto);
     Parallel_Common::bcast_int(nbands_istate);
+    Parallel_Common::bcast_string(bands_to_print_);
     for (int i = 0; i < 3; i++)
     {
         Parallel_Common::bcast_double(kspacing[i]);
@@ -3625,17 +3598,6 @@ 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)
     //-----------------------------------------------------------------------------------

source/module_io/input.h

@@ -48,6 +48,7 @@ class Input
     int ntype; // number of atom types
     int nbands; // number of bands
     int nbands_istate; // number of bands around fermi level for get_pchg calculation.
+
     int pw_seed; // random seed for initializing wave functions qianrui 2021-8-12
 
     bool init_vel;             // read velocity from STRU or not  liuyu 2021-07-14
@@ -262,8 +263,6 @@ 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
@@ -640,6 +639,8 @@ class Input
 
     int count_ntype(const std::string &fn); // sunliang add 2022-12-06
 
+    std::string bands_to_print_; // specify the bands to be calculated in the get_pchg calculation, formalism similar to ocp_set.
+
   public:
     template <class T> static void read_value(std::ifstream &ifs, T &var)
     {
@@ -697,6 +698,12 @@ class Input
     typename std::enable_if<std::is_same<T, std::string>::value, T>::type cast_string(const std::string& str) { return str; }
     void strtolower(char *sa, char *sb);
     void read_bool(std::ifstream &ifs, bool &var);
+
+    // Return the const string pointer of private member bands_to_print_
+    const std::string* get_bands_to_print() const
+    {
+        return &bands_to_print_;
+    }
 };
 
 extern Input INPUT;

source/module_io/input_conv.cpp

@@ -305,6 +305,7 @@ void Input_Conv::Convert(void)
     GlobalV::MIN_DIST_COEF = INPUT.min_dist_coef;
     GlobalV::NBANDS = INPUT.nbands;
     GlobalV::NBANDS_ISTATE = INPUT.nbands_istate;
+
     GlobalV::device_flag = psi::device::get_device_flag(INPUT.device, INPUT.ks_solver, INPUT.basis_type);
 
     if (GlobalV::device_flag == "gpu")