add npart_sto to avoid using too much memory

Author: Qianruipku

docs/input-main.md

@@ -642,6 +642,12 @@ This part of variables are used to control the parameters of stochastic DFT (SDF
 - **Description**: Frequency (once each initsto_freq steps) to generate new stochastic orbitals when running md.
 - **Default**:1000
 
+#### npart_sto
+
+- **Type**: Integer
+- **Description**: Make memory cost to 1/npart_sto times of previous one when running post process of SDFT like DOS with method_sto = 2.
+- **Default**:1
+
 ### Geometry relaxation
 
 This part of variables are used to control the geometry relaxation.

source/input.cpp

@@ -145,6 +145,7 @@ void Input::Default(void)
     kpar = 1;
     initsto_freq = 1000;
     method_sto = 2;
+    npart_sto = 1;
     cal_cond = false;
     dos_nche = 100;
     cond_nche = 20;
@@ -595,6 +596,10 @@ bool Input::Read(const std::string &fn)
         {
             read_value(ifs, method_sto);
         }
+        else if (strcmp("npart_sto", word) == 0)
+        {
+            read_value(ifs, npart_sto);
+        }
         else if (strcmp("cal_cond", word) == 0)
         {
             read_value(ifs, cal_cond);
@@ -1978,6 +1983,7 @@ void Input::Bcast()
     Parallel_Common::bcast_double(emin_sto);
     Parallel_Common::bcast_int(initsto_freq);
     Parallel_Common::bcast_int(method_sto);
+    Parallel_Common::bcast_int(npart_sto);
     Parallel_Common::bcast_bool(cal_cond);
     Parallel_Common::bcast_int(cond_nche);
     Parallel_Common::bcast_double(cond_dw);

source/input.h

@@ -72,6 +72,7 @@ class Input
     int bndpar; //parallel for stochastic/deterministic bands
     int initsto_freq; //frequency to init stochastic orbitals when running md
     int method_sto; //different methods for sdft, 1: slow, less memory  2: fast, more memory
+    int npart_sto; //for method_sto = 2, reduce memory
     bool cal_cond; //calculate electronic conductivities
     int cond_nche; //orders of Chebyshev expansions for conductivities
     double cond_dw; //d\omega for conductivities

source/module_esolver/esolver_sdft_pw.cpp

@@ -166,17 +166,21 @@ void ESolver_SDFT_PW::postprocess()
     if(INPUT.out_dos)
 	{
         double emax, emin;
-        if(INPUT.dos_setemax)	emax = INPUT.dos_emax_ev;
-		if(INPUT.dos_setemin)	emin = INPUT.dos_emin_ev;
+        if(INPUT.dos_setemax)	
+            emax = INPUT.dos_emax_ev;
+        else
+            emax =  ((hsolver::HSolverPW_SDFT*)phsol)->stoiter.stohchi.Emax*ModuleBase::Ry_to_eV;
+		if(INPUT.dos_setemin)
+        	emin = INPUT.dos_emin_ev;
+        else
+            emin =  ((hsolver::HSolverPW_SDFT*)phsol)->stoiter.stohchi.Emin*ModuleBase::Ry_to_eV;
 		if(!INPUT.dos_setemax && !INPUT.dos_setemin)
 		{
-            emax =  ((hsolver::HSolverPW_SDFT*)phsol)->stoiter.stohchi.Emax;
-            emin =  ((hsolver::HSolverPW_SDFT*)phsol)->stoiter.stohchi.Emin;
 			double delta=(emax-emin)*INPUT.dos_scale;
 			emax=emax+delta/2.0;
 			emin=emin-delta/2.0;
 		}
-        this->caldos(INPUT.dos_nche, INPUT.b_coef, emin, emax, INPUT.dos_edelta_ev );
+        this->caldos(INPUT.dos_nche, INPUT.b_coef, emin, emax, INPUT.dos_edelta_ev, INPUT.npart_sto );
     }
 }
 

source/module_esolver/esolver_sdft_pw.h

@@ -33,7 +33,7 @@ class ESolver_SDFT_PW: public ESolver_KS_PW
              const double dw_in, const int times);
     //calculate DOS
     void caldos(const int nche_dos, const double sigmain, 
-            const double emin, const double emax, const double de);
+            const double emin, const double emax, const double de, const int npart);
 
 private:
     int nche_sto; //norder of Chebyshev

source/module_esolver/esolver_sdft_pw_tool.cpp

@@ -454,7 +454,7 @@ void ESolver_SDFT_PW::sKG(const int nche_KG, const double fwhmin, const double w
     ModuleBase::timer::tick(this->classname,"sKG");
 }
 
-void ESolver_SDFT_PW:: caldos( const int nche_dos, const double sigmain, const double emin, const double emax, const double de)
+void ESolver_SDFT_PW:: caldos( const int nche_dos, const double sigmain, const double emin, const double emax, const double de, const int npart)
 {
     cout<<"========================="<<endl;
     cout<<"###Calculating Dos....###"<<endl;
@@ -476,7 +476,8 @@ void ESolver_SDFT_PW:: caldos( const int nche_dos, const double sigmain, const d
     {
         spolyv = new double [nche_dos*nche_dos];
         ModuleBase::GlobalFunc::ZEROS(spolyv, nche_dos*nche_dos);
-        allorderchi = new std::complex<double> [this->stowf.nchip_max * npwx * nche_dos];
+        int nchip_new = ceil((double)this->stowf.nchip_max / npart);
+        allorderchi = new std::complex<double> [nchip_new * npwx * nche_dos];
     }
     cout<<"1. TracepolyA:"<<endl;
     for (int ik = 0;ik < nk;ik++)
@@ -490,7 +491,7 @@ void ESolver_SDFT_PW:: caldos( const int nche_dos, const double sigmain, const d
         const int npw = GlobalC::kv.ngk[ik];
         const int nchipk = this->stowf.nchip[ik];
 
-        complex<double> * pchi;
+        std::complex<double> * pchi;
         if(GlobalV::NBANDS > 0)
             pchi = stowf.chiortho[ik].c;
         else
@@ -505,17 +506,28 @@ void ESolver_SDFT_PW:: caldos( const int nche_dos, const double sigmain, const d
         }
         else
         {
-            ModuleBase::GlobalFunc::ZEROS(allorderchi, this->stowf.nchip_max * npwx * nche_dos);
-            che.calpolyvec_complex(&stohchi, &Stochastic_hchi::hchi_norm, pchi, allorderchi, npw, npwx, nchipk);
-            double* vec_all= (double *) allorderchi;
+            int N = nche_dos;
+            double kweight = GlobalC::kv.wk[ik] / 2;
             char trans = 'T';
             char normal = 'N';
             double one = 1;
-            int LDA = npwx * nchipk * 2;
-            int M = npwx * nchipk * 2;
-            int N = nche_dos;
-            double kweight = GlobalC::kv.wk[ik] / 2;
-            dgemm_(&trans,&normal, &N,&N,&M,&kweight,vec_all,&LDA,vec_all,&LDA,&one,spolyv,&N);
+            for(int ipart = 0 ; ipart < npart ; ++ipart)
+            {
+                int nchipk_new = nchipk / npart;
+                int start_nchipk = ipart * nchipk_new + nchipk % npart;
+                if(ipart < nchipk % npart)
+                {
+                    nchipk_new++;
+                    start_nchipk = ipart * nchipk_new;
+                }
+                ModuleBase::GlobalFunc::ZEROS(allorderchi, nchipk_new * npwx * nche_dos);
+                std::complex<double> *tmpchi = pchi + start_nchipk * npwx;
+                che.calpolyvec_complex(&stohchi, &Stochastic_hchi::hchi_norm, tmpchi, allorderchi, npw, npwx, nchipk_new);
+                double* vec_all= (double *) allorderchi;
+                int LDA = npwx * nchipk_new * 2;
+                int M = npwx * nchipk_new * 2;
+                dgemm_(&trans,&normal, &N,&N,&M,&kweight,vec_all,&LDA,vec_all,&LDA,&one,spolyv,&N);
+            }
         }
     }
     if(stoiter.method == 2) delete[] allorderchi;

source/src_io/write_input.cpp

@@ -111,6 +111,7 @@ void Input::Print(const std::string &fn) const
 
     ofs << "\n#Parameters (3.Stochastic DFT)" << std::endl;
     ModuleBase::GlobalFunc::OUTP(ofs, "method_sto", method_sto, "1: slow and save memory, 2: fast and waste memory");
+    ModuleBase::GlobalFunc::OUTP(ofs, "npart_sto", npart_sto, "Reduce memory when calculating Stochastic DOS");
     ModuleBase::GlobalFunc::OUTP(ofs, "nbands_sto", nbands_sto, "number of stochstic orbitals");
     ModuleBase::GlobalFunc::OUTP(ofs, "nche_sto", nche_sto, "Chebyshev expansion orders");
     ModuleBase::GlobalFunc::OUTP(ofs, "emin_sto", emin_sto, "trial energy to guess the lower bound of eigen energies of the Hamitonian operator");

tests/integrate/186_PW_SDOS_10D10S/INPUT

@@ -39,4 +39,4 @@ dos_emax_ev		 100
 dos_edelta_ev    0.1
 dos_sigma        4
 dos_nche         240
-
+npart_sto        2