Merge pull request #1179 from deepmodeling/SDFT

Feature: Optimize SDFT to use less Chebyshev expansions

Author: mohanchen

docs/input-main.md

@@ -597,9 +597,9 @@ This part of variables are used to control the parameters of stochastic DFT (SDF
 - **Description**:
   - Different method to do SDFT.
   - 1: SDFT calculates $T_n(\hat{h})\ket{\chi}$ twice, where $T_n(x)$ is the n-th order Chebyshev polynomial and $\hat{h}=\frac{\hat{H}-\bar{E}}{\Delta E}$ owning eigen-value $\in(-1,1)$. This method cost less memory but slow.
-  - 2: SDFT calculates $T_n(\hat{h})\ket{\chi}$ once but need much more memory. This method is fast but when memory is not enough, only method 1 can be used.
-  - other: use 1
-- **Default**: 1
+  - 2: SDFT calculates $T_n(\hat{h})\ket{\chi}$ once but need much more memory. This method is much faster. Besides, it calculate $N_e$ with $\bra{\chi}\sqrt{\hat f}\sqrt{\hat f}\ket{\chi}$, which needs smaller [nche_sto](#nche_sto). However, when memory is not enough, only method 1 can be used.
+  - other: use 2
+- **Default**: 2
 
 #### nbands_sto
 
@@ -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

@@ -144,7 +144,8 @@ void Input::Default(void)
     bndpar = 1;
     kpar = 1;
     initsto_freq = 1000;
-    method_sto = 1;
+    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);
@@ -1940,6 +1945,10 @@ void Input::Default_2(void) // jiyy add 2019-08-04
     }
     if(calculation.substr(0,3) != "sto")    bndpar = 1;
     if(bndpar > GlobalV::NPROC) bndpar = GlobalV::NPROC;
+    if(method_sto != 1 && method_sto != 2) 
+    {
+        method_sto = 2;
+    }
 }
 #ifdef __MPI
 void Input::Bcast()
@@ -1974,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_base/math_chebyshev_def.h

@@ -349,15 +349,19 @@ void Chebyshev<REAL>::calpolyvec_complex(T *ptr,
 {
 
     assert(N>=0 && LDA >= N);
-    int ndmxt;
-    if(m == 1) ndmxt = N * m;
-    else       ndmxt = LDA * m; 
+    const int ndmxt = LDA * m; 
 
     std::complex<REAL> *arraynp1 = polywaveout + 2 * ndmxt;
 	std::complex<REAL> *arrayn = polywaveout + ndmxt;
 	std::complex<REAL> *arrayn_1 = polywaveout;
-  
-    ModuleBase::GlobalFunc::DCOPY(wavein, arrayn_1, ndmxt);
+    
+    std::complex<REAL> *tmpin = wavein, *tmpout = arrayn_1;
+    for(int i = 0 ; i < m ; ++i)
+    {
+        ModuleBase::GlobalFunc::DCOPY(tmpin, tmpout, N);
+        tmpin += LDA;
+        tmpout += LDA;
+    }
 
     //1-st order
     (ptr->*funA)(arrayn_1, arrayn, m);

source/module_cell/read_atoms.cpp

@@ -567,6 +567,7 @@ bool UnitCell_pseudo::read_atom_positions(std::ifstream &ifpos, std::ofstream &o
 			if (na > 0)
 			{
        			delete[] atoms[it].tau;
+				delete[] atoms[it].tau_original;
 				delete[] atoms[it].taud;
 				delete[] atoms[it].vel;
        			delete[] atoms[it].mbl;

source/module_esolver/esolver_ks_pw_tool.cpp

@@ -158,7 +158,7 @@ void ESolver_KS_PW::calcondw(const int nt,const double dt,const double fwhmin,co
     double * cw11 = new double [nw];
     double * cw12 = new double [nw];
     double * cw22 = new double [nw];
-    double * kappa = new double [(int)ceil(wcut/dw_in)];
+    double * kappa = new double [nw];
     ModuleBase::GlobalFunc::ZEROS(cw11,nw);
     ModuleBase::GlobalFunc::ZEROS(cw12,nw);
     ModuleBase::GlobalFunc::ZEROS(cw22,nw);

source/module_esolver/esolver_sdft_pw.cpp

@@ -144,7 +144,6 @@ void ESolver_SDFT_PW::postprocess()
 {
     GlobalC::en.print_occ();
 
-    ((hsolver::HSolverPW_SDFT*)phsol)->stoiter.cleanchiallorder();//release lots of memories
     if(this->maxniter == 0)
     {
         int iter = 1;
@@ -153,6 +152,7 @@ void ESolver_SDFT_PW::postprocess()
         hsolver::DiagoIterAssist::PW_DIAG_THR = std::max(std::min(1e-5, 0.1 * GlobalV::SCF_THR / std::max(1.0, GlobalC::CHR.nelec)),1e-12);
         hsolver::DiagoIterAssist::need_subspace = false;
         this->phsol->solve(this->phami, this->psi[0], this->pelec,this->stowf,istep, iter, GlobalV::KS_SOLVER, true);
+        ((hsolver::HSolverPW_SDFT*)phsol)->stoiter.cleanchiallorder();//release lots of memories
     }
     int nche_test = 0;
     if(INPUT.cal_cond)  nche_test = std::max(nche_test, INPUT.cond_nche);
@@ -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

@@ -56,7 +56,7 @@ void ESolver_SDFT_PW::check_che(const int nche_in)
             while(1)
             {
                 bool converge;
-                converge= chetest.checkconverge(&stohchi, &Stochastic_hchi::hchi_reciprocal, 
+                converge= chetest.checkconverge(&stohchi, &Stochastic_hchi::hchi_norm, 
 	        	    	pchi, npw, stohchi.Emax, stohchi.Emin, 5.0);
 
                 if(!converge)
@@ -263,10 +263,10 @@ void ESolver_SDFT_PW::sKG(const int nche_KG, const double fwhmin, const double w
         ModuleBase::GlobalFunc::COPYARRAY(hj1sfpsi_out, j2sfpsi.get_pointer(), ndim*totbands_per*npwx);
 
         /*
-        // stohchi.hchi_reciprocal(psi0.get_pointer(), hpsi0.get_pointer(), totbands_per);
-        // stohchi.hchi_reciprocal(sfpsi0.get_pointer(), hsfpsi0.get_pointer(), totbands_per);
-        // stohchi.hchi_reciprocal(j1psi.get_pointer(), j2psi.get_pointer(), ndim*totbands_per);
-        // stohchi.hchi_reciprocal(j1sfpsi.get_pointer(), j2sfpsi.get_pointer(), ndim*totbands_per);
+        // stohchi.hchi_norm(psi0.get_pointer(), hpsi0.get_pointer(), totbands_per);
+        // stohchi.hchi_norm(sfpsi0.get_pointer(), hsfpsi0.get_pointer(), totbands_per);
+        // stohchi.hchi_norm(j1psi.get_pointer(), j2psi.get_pointer(), ndim*totbands_per);
+        // stohchi.hchi_norm(j1sfpsi.get_pointer(), j2sfpsi.get_pointer(), ndim*totbands_per);
         // double Ebar = (stohchi.Emin + stohchi.Emax)/2;
 	    // double DeltaE = (stohchi.Emax - stohchi.Emin)/2;
 	    // for(int ib = 0 ; ib < totbands_per ; ++ib)
@@ -307,8 +307,8 @@ void ESolver_SDFT_PW::sKG(const int nche_KG, const double fwhmin, const double w
 
         //(1-f)
         che.calcoef_real(&stoiter.stofunc,&Sto_Func<double>::n_fd);
-        che.calfinalvec_real(&stohchi, &Stochastic_hchi::hchi_reciprocal, j1sfpsi.get_pointer(), j1sfpsi.get_pointer(), npw, npwx, totbands_per*ndim);
-        che.calfinalvec_real(&stohchi, &Stochastic_hchi::hchi_reciprocal, j2sfpsi.get_pointer(), j2sfpsi.get_pointer(), npw, npwx, totbands_per*ndim);
+        che.calfinalvec_real(&stohchi, &Stochastic_hchi::hchi_norm, j1sfpsi.get_pointer(), j1sfpsi.get_pointer(), npw, npwx, totbands_per*ndim);
+        che.calfinalvec_real(&stohchi, &Stochastic_hchi::hchi_norm, j2sfpsi.get_pointer(), j2sfpsi.get_pointer(), npw, npwx, totbands_per*ndim);
 
         psi::Psi<std::complex<double>> *p_j1psi = &j1psi; 
         psi::Psi<std::complex<double>> *p_j2psi = &j2psi; 
@@ -361,9 +361,9 @@ void ESolver_SDFT_PW::sKG(const int nche_KG, const double fwhmin, const double w
             }
 
             //exp(iHdt)|chi>
-            chet.calfinalvec_complex(&stohchi, &Stochastic_hchi::hchi_reciprocal, &exppsi(ksbandper,0), &exppsi(ksbandper,0), npw, npwx, nchip);
+            chet.calfinalvec_complex(&stohchi, &Stochastic_hchi::hchi_norm, &exppsi(ksbandper,0), &exppsi(ksbandper,0), npw, npwx, nchip);
             //exp(-iHdt)|shchi>
-            chet2.calfinalvec_complex(&stohchi, &Stochastic_hchi::hchi_reciprocal, &expsfpsi(ksbandper,0), &expsfpsi(ksbandper,0), npw, npwx, nchip);
+            chet2.calfinalvec_complex(&stohchi, &Stochastic_hchi::hchi_norm, &expsfpsi(ksbandper,0), &expsfpsi(ksbandper,0), npw, npwx, nchip);
             psi::Psi<std::complex<double>> *p_exppsi = &exppsi;
 #ifdef __MPI
             psi::Psi<std::complex<double>> exppsi_tot;
@@ -454,53 +454,111 @@ 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<<"Calculating Dos...."<<endl;
+    cout<<"========================="<<endl;
+    cout<<"###Calculating Dos....###"<<endl;
+    cout<<"========================="<<endl;
     ModuleBase::Chebyshev<double> che(nche_dos);
     const int nk = GlobalC::kv.nks;
     Stochastic_Iter& stoiter = ((hsolver::HSolverPW_SDFT*)phsol)->stoiter;
     Stochastic_hchi& stohchi = stoiter.stohchi;
     const int npwx = GlobalC::wf.npwx;
 
-    double * spolyv = new double [nche_dos];
-    ModuleBase::GlobalFunc::ZEROS(spolyv, nche_dos);
+    double * spolyv = nullptr;
+    std::complex<double> *allorderchi = nullptr;
+    if(stoiter.method == 1)
+    {
+        spolyv = new double [nche_dos];
+        ModuleBase::GlobalFunc::ZEROS(spolyv, nche_dos);
+    }
+    else
+    {
+        spolyv = new double [nche_dos*nche_dos];
+        ModuleBase::GlobalFunc::ZEROS(spolyv, nche_dos*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++)
 	{
+        cout<<"ik: "<<ik+1<<endl;
 		if(nk > 1) 
         {
             this->phami->updateHk(ik);
         }
         stohchi.current_ik = ik;
         const int npw = GlobalC::kv.ngk[ik];
-        const int nchip = this->stowf.nchip[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
             pchi = stowf.chi0[ik].c;
-        che.tracepolyA(&stohchi, &Stochastic_hchi::hchi_reciprocal, pchi, npw, npwx, nchip);
-        for(int i = 0 ; i < nche_dos ; ++i)
+        if(stoiter.method == 1)
         {
-            spolyv[i] += che.polytrace[i] * GlobalC::kv.wk[ik] / 2 ;
+            che.tracepolyA(&stohchi, &Stochastic_hchi::hchi_norm, pchi, npw, npwx, nchipk);
+            for(int i = 0 ; i < nche_dos ; ++i)
+            {
+                spolyv[i] += che.polytrace[i] * GlobalC::kv.wk[ik] / 2 ;
+            }
+        }
+        else
+        {
+            int N = nche_dos;
+            double kweight = GlobalC::kv.wk[ik] / 2;
+            char trans = 'T';
+            char normal = 'N';
+            double one = 1;
+            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;
+
     string dosfile = GlobalV::global_out_dir+"DOS1_smearing.dat";
     ofstream ofsdos(dosfile.c_str());
     int ndos = int((emax-emin) / de)+1;
     double *dos = new double [ndos];
     ModuleBase::GlobalFunc::ZEROS(dos,ndos);
     stoiter.stofunc.sigma = sigmain / ModuleBase::Ry_to_eV;
     double sum = 0; 
-    double error = 0;
+    double maxerror = 0;
     ofsdos<<setw(8)<<"## E(eV) "<<setw(20)<<"dos(eV^-1)"<<setw(20)<<"sum"<<setw(20)<<"Error(eV^-1)"<<endl;
+    cout<<"2. Dos:"<<endl;
+    int n10 = ndos/10;
+    int percent = 10;
 	for(int ie = 0; ie < ndos; ++ie)
 	{
-		stoiter.stofunc.targ_e = (emin + ie * de) / ModuleBase::Ry_to_eV;
-        che.calcoef_real(&stoiter.stofunc, &Sto_Func<double>::ngauss);
         double KS_dos = 0;
-		double sto_dos = BlasConnector::dot(nche_dos,che.coef_real,1,spolyv,1);
+        double sto_dos = 0;
+        stoiter.stofunc.targ_e = (emin + ie * de) / ModuleBase::Ry_to_eV;
+        if(stoiter.method == 1)
+        {
+            che.calcoef_real(&stoiter.stofunc, &Sto_Func<double>::ngauss);
+		    sto_dos = BlasConnector::dot(nche_dos,che.coef_real,1,spolyv,1);
+        }
+        else
+        {
+            che.calcoef_real(&stoiter.stofunc, &Sto_Func<double>::nroot_gauss);
+            sto_dos = stoiter.vTMv(che.coef_real,spolyv,nche_dos);
+        }
         if(GlobalV::NBANDS > 0)
         {
             for(int ik = 0; ik < nk; ++ik)
@@ -517,16 +575,35 @@ void ESolver_SDFT_PW:: caldos( const int nche_dos, const double sigmain, const d
 	    MPI_Allreduce(MPI_IN_PLACE, &KS_dos, 1, MPI_DOUBLE, MPI_SUM , STO_WORLD);
         MPI_Allreduce(MPI_IN_PLACE, &sto_dos, 1, MPI_DOUBLE, MPI_SUM , MPI_COMM_WORLD);
 #endif
-        double tmpre = che.coef_real[nche_dos-1] * spolyv[nche_dos-1];
+        double tmpre = 0;
+        if(stoiter.method == 1)
+        {
+            tmpre = che.coef_real[nche_dos-1] * spolyv[nche_dos-1];
+        }
+        else
+        {
+            const int norder = nche_dos;
+            double last_coef = che.coef_real[norder-1];
+            double last_spolyv = spolyv[norder*norder - 1];
+            tmpre = last_coef *(BlasConnector::dot(norder,che.coef_real,1,spolyv+norder*(norder-1),1)
+                        + BlasConnector::dot(norder,che.coef_real,1,spolyv+norder-1,norder)-last_coef*last_spolyv);
+        } 
 #ifdef __MPI
         MPI_Allreduce(MPI_IN_PLACE, &tmpre, 1, MPI_DOUBLE, MPI_SUM , MPI_COMM_WORLD);
 #endif
-        if(error < tmpre) error = tmpre;
+        if(maxerror < tmpre) maxerror = tmpre;
         dos[ie] = (KS_dos + sto_dos) / ModuleBase::Ry_to_eV;
         sum += dos[ie];
-		ofsdos <<setw(8)<< emin + ie * de <<setw(20)<<dos[ie]<<setw(20)<<sum * de <<setw(20) <<error <<endl;
+		ofsdos <<setw(8)<< emin + ie * de <<setw(20)<<dos[ie]<<setw(20)<<sum * de <<setw(20) <<tmpre <<endl;
+        if(ie%n10 == n10 -1) 
+        {
+            cout<<percent<<"%"<<" ";
+            percent+=10;
+        }
 	}
-    cout<<scientific<<"DOS max Chebyshev Error: "<<error<<endl;
+    cout<<endl;
+    cout<<"Finish DOS"<<endl;
+    cout<<scientific<<"DOS max Chebyshev Error: "<<maxerror<<endl;
     delete[] dos;
     delete[] spolyv;
     return;

source/module_hsolver/hsolver_pw_sdft.cpp

@@ -100,7 +100,7 @@ namespace hsolver
 			}
 		}
 		// calculate stochastic rho
-		stoiter.sum_stoband(stowf,pes);
+		stoiter.sum_stoband(stowf,pes,pHamilt);
 
 
 		//(6) calculate the delta_harris energy