math_chebyshev.h

Go to the documentation of this file.

#ifndef STO_CHEBYCHEV_H
#define STO_CHEBYCHEV_H
#include "fftw3.h"
#include "source_base/module_device/device.h"
#include "source_base/module_device/memory_op.h"
#include "source_base/module_container/ATen/core/tensor_types.h"
 
#include <complex>
#include <functional>
 
namespace ModuleBase
{
// template class for fftw
template <typename T>
class FFTW;
 
template <typename REAL, typename Device = base_device::DEVICE_CPU>
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
    void calcoef_real(std::function<REAL(REAL)> fun);
 
    // Calculate coefficients C_n[g], where g is a function of complex number
    void calcoef_complex(std::function<std::complex<REAL>(std::complex<REAL>)> fun);
 
    // 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))
    void calcoef_pair(std::function<REAL(REAL)> fun1, std::function<REAL(REAL)> fun2);
 
    // 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.
    void calfinalvec_real(std::function<void(REAL*, REAL*, const int)> funA,
                          REAL* wavein,
                          REAL* waveout,
                          const int N,
                          const int LDA = 1,
                          const int m = 1); // do not define yet
 
    // calfinalvec_real uses C_n[f], where f is a function of real number and A is a complex Operator.
    void calfinalvec_real(std::function<void(std::complex<REAL>*, std::complex<REAL>*, const int)> funA,
                          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.
    void calfinalvec_complex(std::function<void(std::complex<REAL>*, std::complex<REAL>*, const int)> funA,
                             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
    void tracepolyA(std::function<void(std::complex<REAL>* in, std::complex<REAL>* out, const int)> funA,
                    std::complex<REAL>* wavein,
                    const int N,
                    const int LDA = 1,
                    const int m = 1);
 
    // get T_n(x)
    void getpolyval(REAL x, REAL* polyval, const int N);
 
    // get each order of vector: {T_0(A)v, T_1(A)v, ..., T_n(A)v}
    // Note: use it carefully, it will cost a lot of memory!
    // calpolyvec_real: f(x) = \sum_n C_n*T_n(x), f is a real function
    void calpolyvec_real(std::function<void(REAL* in, REAL* out, const int)> funA,
                         REAL* wavein,
                         REAL* waveout,
                         const int N,
                         const int LDA = 1,
                         const int m = 1); // do not define yet
 
    // calpolyvec_complex: f(x) = \sum_n C_n*T_n(x), f is a complex function
    void calpolyvec_complex(std::function<void(std::complex<REAL>* in, std::complex<REAL>* out, const int)> funA,
                            std::complex<REAL>* wavein,
                            std::complex<REAL>* waveout,
                            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
    void recurs_real(std::function<void(REAL* in, REAL* out, const int)> funA,
                     REAL* arraynp1,
                     REAL* arrayn,
                     REAL* arrayn_1,
                     const int N,
                     const int LDA = 1,
                     const int m = 1);
 
    // recurs_complex: A is a complex operator
    void recurs_complex(std::function<void(std::complex<REAL>* in, std::complex<REAL>* out, const int)> funA,
                        std::complex<REAL>* arraynp1,
                        std::complex<REAL>* arrayn,
                        std::complex<REAL>* arrayn_1,
                        const int N,
                        const int LDA = 1,
                        const int m = 1);
 
    // return 2xTn-Tn_1
    REAL recurs(const REAL x, const REAL Tn, const REAL Tn_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.
    bool checkconverge(std::function<void(std::complex<REAL>* in, std::complex<REAL>* out, const int)> funA,
                       std::complex<REAL>* wavein,
                       const int N,
                       const int LDA,
                       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 = nullptr;                  //[Device] expansion coefficient of each order
    std::complex<REAL>* coef_complex = nullptr; //[Device] expansion coefficient of each order
    REAL* coefr_cpu = nullptr;                  //[CPU] expansion coefficient of each order
    std::complex<REAL>* coefc_cpu = nullptr;    //[CPU] expansion coefficient of each order
 
    FFTW<REAL> fftw;          // use for fftw
    REAL* polytrace;                  //[CPU] w_n = \sum_i v^+ * T_n(A) * v, only
 
    bool getcoef_real;    // coef_real has been calculated
    bool getcoef_complex; // coef_complex has been calculated
 
  public:
    // 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);
  
  private:
    Device* ctx = {};
    base_device::DEVICE_CPU* cpu_ctx = {};
    using ct_Device = typename container::PsiToContainer<Device>::type;
    using resmem_complex_op = base_device::memory::resize_memory_op<std::complex<REAL>, Device>; 
    using resmem_var_op = base_device::memory::resize_memory_op<REAL, Device>;
    using delmem_complex_op = base_device::memory::delete_memory_op<std::complex<REAL>, Device>;
    using delmem_var_op = base_device::memory::delete_memory_op<REAL, Device>;
    using syncmem_var_h2d_op = base_device::memory::synchronize_memory_op<REAL, Device, base_device::DEVICE_CPU>;
    using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<REAL, base_device::DEVICE_CPU, Device>;
    using syncmem_complex_h2d_op = base_device::memory::synchronize_memory_op<std::complex<REAL>, Device, base_device::DEVICE_CPU>;
    using syncmem_complex_d2h_op = base_device::memory::synchronize_memory_op<std::complex<REAL>, base_device::DEVICE_CPU, Device>;
    using memcpy_var_op = base_device::memory::synchronize_memory_op<REAL, Device, Device>;
    using memcpy_complex_op = base_device::memory::synchronize_memory_op<std::complex<REAL>, Device, Device>;
    using setmem_complex_op = base_device::memory::set_memory_op<std::complex<REAL>, Device>;
};
 
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 __ENABLE_FLOAT_FFTW
template <>
class FFTW<float>
{
  public:
    FFTW(const int norder2_in);
    ~FFTW();
    void execute_fftw();
    float* dcoef; //[norder2]
    fftwf_complex* ccoef;
    fftwf_plan coef_plan;
};
#endif
 
} // namespace ModuleBase
 
#endif