Merge pull request #930 from Qianruipku/SDFT

refactor: change Stochastic_Chebyshev to a math lib Chebyshev

Author: mohanchen

source/Makefile.Objects

@@ -59,8 +59,8 @@ unk_overlap_pw.o \
 berryphase.o \
 sto_iter.o\
 sto_wf.o\
+sto_func.o\
 sto_hchi.o\
-sto_che.o\
 sto_forces.o\
 sto_stress_pw.o
 
@@ -81,6 +81,7 @@ math_integral.o\
 math_ylmreal.o\
 mathzone_add1.o\
 math_bspline.o\
+math_chebyshev.o\
 integral.o \
 polint.o \
 sph_bessel.o \

source/module_base/CMakeLists.txt

@@ -16,7 +16,8 @@ add_library(
     math_polyint.cpp
     math_sphbes.cpp
     math_ylmreal.cpp
-	  math_bspline.cpp
+	math_bspline.cpp
+    math_chebyshev.cpp
     mathzone.cpp
     mathzone_add1.cpp
     matrix.cpp

source/module_base/math_chebyshev.cpp

@@ -0,0 +1,109 @@
+#include "./math_chebyshev.h" 
+#include "./constants.h"
+#include "./blas_connector.h"
+#include "./global_function.h"
+namespace ModuleBase
+{
+//we only have two examples: double and float.
+template class Chebyshev<double>;
+#ifdef __MIX_PRECISION
+template class Chebyshev<float>;
+#endif
+
+FFTW<double>::FFTW(const int norder2_in)
+{
+    ccoef = (fftw_complex *) fftw_malloc(sizeof(fftw_complex) * norder2_in);
+    dcoef = (double *) fftw_malloc(sizeof(double) * norder2_in);
+    coef_plan = fftw_plan_dft_r2c_1d(norder2_in, dcoef, ccoef, FFTW_ESTIMATE);
+}
+FFTW<double>::~FFTW()
+{
+    fftw_destroy_plan(coef_plan);
+	fftw_free(ccoef);
+	fftw_free(dcoef);
+}
+void FFTW<double>::execute_fftw()
+{
+    fftw_execute(this->coef_plan);
+}
+
+#ifdef __MIX_PRECISION
+FFTW<float>::FFTW(const int norder2_in)
+{
+    ccoef = (fftwf_complex *) fftw_malloc(sizeof(fftwf_complex) * norder2_in);
+    dcoef = (float *) fftw_malloc(sizeof(float) * norder2_in);
+    coef_plan = fftwf_plan_dft_r2c_1d(norder2_in, dcoef, ccoef, FFTW_ESTIMATE);
+}
+FFTW<float>::~FFTW()
+{
+    fftwf_destroy_plan(coef_plan);
+	fftw_free(ccoef);
+	fftw_free(dcoef);
+}
+void FFTW<float>::execute_fftw()
+{
+    fftwf_execute(this->coef_plan);
+}
+#endif
+
+//A number to control the number of grids in C_n integration
+#define EXTEND 16
+
+template<typename REAL>
+Chebyshev<REAL>::Chebyshev(const int norder_in) : fftw(2 * EXTEND * norder_in)
+{
+    this->norder = norder_in;
+    norder2 = 2 * norder * EXTEND;
+    if(this->norder < 1)
+    {
+        ModuleBase::WARNING_QUIT("Stochastic_Chebychev", "The Chebyshev expansion order should be at least 1!");
+    }
+    polytrace = new REAL [norder];
+    coef_real = new REAL [norder];
+    coef_complex = new std::complex<REAL> [norder];
+
+    // ndmin = ndmax = ndmax_in;
+
+    getcoef_complex = false;
+    getcoef_real = false;
+}
+
+template<typename REAL>
+Chebyshev<REAL>::~Chebyshev()
+{
+	delete [] polytrace;
+	delete [] coef_real;
+    delete [] coef_complex;
+}
+
+template<typename REAL>
+REAL Chebyshev<REAL>:: ddot_real(
+    const std::complex<REAL>* psi_L,
+    const std::complex<REAL>* psi_R,
+    const int N, const int LDA, const int m)
+{
+    REAL result = 0;
+    if(N == LDA || m==1)
+    {
+        int dim2=2 * N * m;
+        REAL *pL,*pR;
+        pL=(REAL *)psi_L;
+        pR=(REAL *)psi_R;
+        result=BlasConnector::dot(dim2,pL,1,pR,1);
+    }
+    else
+    {
+        REAL *pL,*pR;
+        pL=(REAL *)psi_L;
+        pR=(REAL *)psi_R;
+        for(int i = 0 ; i < m ; ++i)
+        {
+            int dim2=2 * N;
+            result +=  BlasConnector::dot(dim2,pL,1,pR,1);
+            pL += 2 * LDA;
+            pR += 2 * LDA;
+        }
+    }
+    return result;
+}
+}

source/module_base/math_chebyshev.h

@@ -0,0 +1,212 @@
+#ifndef STO_CHEBYCHEV_H
+#define STO_CHEBYCHEV_H
+#include <complex>
+#include "fftw3.h"
+
+namespace ModuleBase
+{
+//template class for fftw
+template<typename T>
+class FFTW;
+
+/**
+ * @brief A class to treat the Chebyshev expansion.
+ *
+ * @author qianrui on 2022-05-18
+ * @details
+ * Math:
+ * I. 
+ * Chebyshev polynomial:
+ * T_0(x) = 1;
+ * T_1(x) = x;
+ * T_2(x) = 2x^2 -1;
+ * T_3(x) = 4x^3-3x;
+ * T_{n+2}(x) = 2xT_{n+1}(x) - T_n(x)
+ * II. 
+ * Any analytical function f(x) can be expanded by Chebyshev polynomial:
+ * f(x) = \sum_{n=0}^{norder} C_n[f]*T_n(x) (|x| < 1),
+ * where C_n[f] = \frac{2-\delta_{n0}}{\pi} \int_0^\pi f(cos(\theta))cos(n\theta) d\theta
+ * Here C_n can be calculate with FFT.
+ * III. 
+ * Any functions of linear Operator or matrix f(\hat{A}) or f(A) can also be expanded as well:
+ * f(A) = \sum_{n=0}^{norder-1} C_n[f]*T_n(A) (|all eigenvalues of A| < 1).
+ * f(A)v = \sum_{n=0}^{norder-1} C_n[f]*T_n(A)v, where v is column vector
+ *       = \sum_{n=0}^{norder-1} C_n[f]*v_n, where v_n = T_n(A)v, v_0 = v
+ * v_{n+2} = 2Av_{n+1} - v_n
+ * IV.
+ * v^+f(A)v = \sum_{n=0}^{norder-1} C_n[f]*v^+v_n = \sum_{n=0}^{norder-1} C_n[f] * w_n, 
+ * where w_n = v^+ * v_n = v^+ * T_n(A) * v
+ * 
+ * USAGE：
+ * Chebyshev che(10); // constructe a chebyshev expansion of 10 orders (n=0,1,...,9)
+ * 1. che.calcoef_real(&a, &A::cos) 						// calculate C_n[f], where f is a.cos
+ *    for(inti=0;i<10;++i) cout<<che.coef_real[i]<<endl; 	//Then we print C_n[f]
+ * 
+ *    che.calcoef_complex(&b, &B::expi) 					// calculate C_n[g], where g is b.expi
+ *    for(inti=0;i<10;++i) cout<<che.coef_complex[i]<<endl; //Then we print C_n[g]
+ * 
+ *    che.calcoef_pair(&c, &C::cos, &C::sin) 				// calculate C_n[g], where g is (c.cos, c.sin)
+ *    for(inti=0;i<10;++i) cout<<che.coef_complex[i]<<endl; //Then we print C_n[g]
+ * 
+ * 2. che.calcoef_real(&occ, &Occupy::fd)
+ * 	  che.calfinalvec_real(&hamilt, &Hamilt::hpsi, psi_in, psi_out, npw);
+ *    //calculate f(H)|psi>, where f is occ.fd and H is hamilt.hpsi
+ * 
+ *    che.calcoef_complex(&b, &B::expi)
+ * 	  che.calfinalvec_complex(&hamilt, &Hamilt::hpsi, psi_in, psi_out, npw, npwx, nbands);
+ *    //calculate exp(iH)|psi_i>
+ * 
+ * 3. che.tracepolyA(&hamilt, &Hamilt::hpsi, psi_in, npw, npwx, nbands)
+ * 	  //calculate \sum_i^{nbands} <psi_i|T_n(H)|psi_i>
+ * 
+ * 4. che.recurs_complex(&hamilt, &Hamilt::hpsi, vp1, v, vm1, npw)
+ *    //calculate vp1: |vp1> = 2 H|v> - |vm1>;
+ * 
+ */
+template<typename REAL>
+class Chebyshev
+{
+
+public:
+
+    // constructor and deconstructor
+    Chebyshev(const int norder);
+    ~Chebyshev();
+
+public:
+	// I.
+	// Calculate coefficients C_n[f], where f is a function of real number
+	template<class T>
+    void calcoef_real(T *ptr, REAL (T::*fun)(REAL));
+	// Calculate coefficients C_n[g], where g is a function of complex number
+	template<class T>
+    void calcoef_complex(T *ptr, std::complex<REAL> (T::*fun)(std::complex<REAL>));
+	// Calculate coefficients C_n[g], where g is a general complex function g(x)=(g1(x), g2(x)) e.g. exp(ix)=(cos(x), sin(x))
+	template<class T>
+	void calcoef_pair(T *ptr, REAL (T::*fun1)(REAL), REAL (T::*fun2)(REAL));
+    
+	// II.
+	// Calculate the final vector f(A)v = \sum_{n=0}^{norder-1} C_n[f]*v_n
+	// Here funA(in, out) means the map v -> Av : funA(v, Av)
+	// Here m represents we treat m vectors at the same time: f(A)[v1,...,vm] and funA(in,out,m) means [v1,...,vm] -> A[v1,...,vm]
+	// N is dimension of vector, and LDA is the distance between the first number of v_n and v_{n+1}.
+	// LDA >= max(1, N). It is the same as the BLAS lib.
+	// calfinalvec_real uses C_n[f], where f is a function of real number and A is a real Operator.
+	template<class T>
+    void calfinalvec(T *ptr, 
+		void (T::*funA)(REAL *in, REAL *out, const int), 
+		REAL *wavein, REAL *waveout, 
+		const int N, const int LDA = 1,  const int m = 1);
+
+	// calfinalvec_real uses C_n[f], where f is a function of real number and A is a complex Operator.
+	template<class T>
+    void calfinalvec_real(T *ptr, 
+		void (T::*funA)(std::complex<REAL> *in, std::complex<REAL> *out, const int), 
+		std::complex<REAL> *wavein, std::complex<REAL> *waveout, 
+		const int N, const int LDA = 1,  const int m = 1);
+
+	// calfinalvec_complex uses C_n[g], where g is a function of complex number and A is a complex Operator.
+	template<class T>
+	void calfinalvec_complex(T *ptr,
+		void (T::*funA)(std::complex<REAL> *in, std::complex<REAL> *out, const int), 
+		std::complex<REAL> *wavein, std::complex<REAL> *waveout, 
+		const int N, const int LDA = 1,  const int m = 1);
+
+	// III.
+	// \sum_i v_i^+f(A)v_i = \sum_{i,n=0}^{norder-1} C_n[f]*v_i^+v_{i,n} = \sum_{n=0}^{norder-1} C_n[f] * w_n
+	// calculate the sum of diagonal elements (Trace) of T_n(A) in v-represent: w_n = \sum_i v_i^+ * T_n(A) * v_i
+	// i = 1,2,...m
+	template<class T>
+	void tracepolyA(
+		T *ptr, void (T::*funA)(std::complex<REAL> *in, std::complex<REAL> *out, const int), 
+		std::complex<REAL> *wavein, 
+		const int N, const int LDA = 1,  const int m = 1);
+	
+	// IV.
+	// recurs fomula: v_{n+1} = 2Av_n - v_{n-1}
+	// get v_{n+1} from v_n and v_{n-1}
+	// recurs_complex: A is a real operator
+	template<class T>
+    void recurs_real(
+		T *ptr, void (T::*funA)(REAL *in, REAL *out, const int),
+		REAL* arraynp1,  //v_{n+1}
+		REAL* arrayn,    //v_n
+		REAL* arrayn_1,  //v_{n-1}
+		const int N, const int LDA = 1,  const int m = 1);
+	// recurs_complex: A is a complex operator
+	template<class T>
+    void recurs_complex(
+		T *ptr, void (T::*funA)(std::complex<REAL> *in, std::complex<REAL> *out, const int),
+		std::complex<REAL>* arraynp1,  //v_{n+1}
+		std::complex<REAL>* arrayn,    //v_n
+		std::complex<REAL>* arrayn_1,  //v_{n-1}
+		const int N, const int LDA = 1,  const int m = 1);
+	
+	// V.
+	// auxiliary function
+	// Abs of all eigenvalues of A should be less than 1.
+	// Thus \hat(a) = \frac{(A - (tmax+tmin)/2)}{(tmax-tmin)/2}
+	// tmax >= all eigenvalues; tmin <= all eigenvalues
+	// Here we check if the trial number tmax(tmin) is the upper(lower) bound of eigenvalues and return it.
+    template<class T>
+	bool checkconverge(
+		T *ptr, void (T::*funA)(std::complex<REAL> *in, std::complex<REAL> *out, const int), 
+ 		std::complex<REAL> *wavein, const int N,
+		REAL& tmax, //trial number for upper bound
+		REAL& tmin, //trial number for lower bound
+		REAL stept); //tmax = max() + stept, tmin = min() - stept
+
+public:
+	//Members:
+    int norder;   // order of Chebyshev expansion
+    int norder2;  // 2 * norder * EXTEND
+
+    REAL* coef_real; // expansion coefficient of each order
+	std::complex<REAL>* coef_complex; // expansion coefficient of each order
+	FFTW<REAL> fftw; //use for fftw
+    REAL *polytrace; //w_n = \sum_i v^+ * T_n(A) * v
+
+	bool getcoef_real;    //coef_real has been calculated
+	bool getcoef_complex; //coef_complex has been calculated
+
+private:
+	//SI.
+	//calculate dot product <psi_L|psi_R>
+    REAL ddot_real(
+    const std::complex<REAL>* psi_L,
+    const std::complex<REAL>* psi_R,
+	const int N, const int LDA = 1,  const int m = 1);
+
+    
+};
+
+template<>
+class FFTW<double>
+{
+public:
+	FFTW(const int norder2_in);
+	~FFTW();
+	void execute_fftw();
+    double* dcoef; //[norder2]
+	fftw_complex *ccoef;
+	fftw_plan coef_plan;
+};
+
+#ifdef __MIX_PRECISION
+template<>
+class FFTW<float>
+{
+public:
+	FFTW(const int norder2_in);
+	~FFTW();
+	void execute_fftw();
+    double* dcoef; //[norder2]
+	fftwf_complex *ccoef;
+	fftwf_plan coef_plan;
+};
+#endif
+}
+
+#include "math_chebyshev_def.h"
+
+#endif