diago_cg.h

Go to the documentation of this file.

#ifndef MODULE_HSOLVER_DIAGO_CG_H_
#define MODULE_HSOLVER_DIAGO_CG_H_
 
#include <functional>
 
#include <source_base/macros.h>
#include <source_base/kernels/math_kernel_op.h>
 
#include <ATen/core/tensor.h>
#include <ATen/core/tensor_types.h>
 
namespace hsolver {
 
template <typename T, typename Device = base_device::DEVICE_CPU>
class DiagoCG final
{
    // private: accessibility within class is private by default
    // Note GetTypeReal<T>::type will
    // return T if T is real type(float, double),
    // otherwise return the real type of T(complex<float>, std::complex<double>)
    using Real = typename GetTypeReal<T>::type;
    using ct_Device = typename ct::PsiToContainer<Device>::type;
  public:
    using Func = std::function<void(const ct::Tensor&, ct::Tensor&)>;
    // Constructor need:
    // 1. temporary mock of Hamiltonian "Hamilt_PW"
    // 2. precondition pointer should point to place of precondition array.
    DiagoCG(const std::string& basis_type, const std::string& calculation);
    DiagoCG(
        const std::string& basis_type,
        const std::string& calculation,
        const bool& need_subspace,
        const Func& subspace_func,
        const Real& pw_diag_thr,
        const int& pw_diag_nmax,
        const int& nproc_in_pool);
    
    ~DiagoCG();
 
    // virtual void init(){};
    // refactor hpsi_info
    // this is the diag() function for CG method
    void diag(const Func& hpsi_func,
              const Func& spsi_func,
              ct::Tensor& psi,
              ct::Tensor& eigen,
              const std::vector<double>& ethr_band,
              const ct::Tensor& prec = {});
 
  private:
    Device * ctx_ = {};
    int notconv_ = 0;
    int n_band_ = 0;
    int n_basis_ = 0;
    int avg_iter_ = 0;
    Real pw_diag_thr_ = 1e-5;
    int pw_diag_nmax_ = 0;
    int nproc_in_pool_ = 0;
    std::string basis_type_ = {};
    std::string calculation_ = {};
 
    bool need_subspace_ = false;
    std::function<void(const ct::Tensor&, ct::Tensor&)> hpsi_func_ = nullptr;
    std::function<void(const ct::Tensor&, ct::Tensor&)> spsi_func_ = nullptr;
    std::function<void(const ct::Tensor&, ct::Tensor&)> subspace_func_ = nullptr;
 
    void calc_grad(
        const ct::Tensor& prec,
        ct::Tensor& grad,
        ct::Tensor& hphi,
        ct::Tensor& sphi,
        ct::Tensor& pphi);
 
    void orth_grad(
        const ct::Tensor& psi, 
        const int& m, 
        ct::Tensor& grad, 
        ct::Tensor& scg,
        ct::Tensor& lagrange);
 
    void calc_gamma_cg(
        const int& iter,
        const Real& cg_norm, 
        const Real& theta,
        const ct::Tensor& prec,
        const ct::Tensor& scg,
        const ct::Tensor& grad,
        const ct::Tensor& phi_m,
        Real& gg_last,
        ct::Tensor& g0,
        ct::Tensor& cg);
 
    bool update_psi(
        const ct::Tensor& pphi,
        const ct::Tensor& cg,
        const ct::Tensor& scg,
        const double& ethreshold,
        Real &cg_norm, 
        Real &theta, 
        Real &eigen,
        ct::Tensor& phi_m,
        ct::Tensor& sphi,
        ct::Tensor& hphi);
 
    void schmit_orth(const int& m, const ct::Tensor& psi, const ct::Tensor& sphi, ct::Tensor& phi_m);
 
    // used in diag() for template replace Hamilt with Hamilt_PW
    void diag_mock(const ct::Tensor& prec,
                   ct::Tensor& psi,
                   ct::Tensor& eigen,
                   const std::vector<double>& ethr_band);
 
    bool test_exit_cond(const int& ntry, const int& notconv) const;
 
    using dot_real_op = ModuleBase::dot_real_op<T, Device>;
    const T * one_ = nullptr, * zero_ = nullptr, * neg_one_ = nullptr;
};
 
} // namespace hsolver
 
#endif // MODULE_HSOLVER_DIAGO_CG_H_