sto_iter.h

Go to the documentation of this file.

#ifndef STO_ITER_H
#define STO_ITER_H
#include "source_base/math_chebyshev.h"
#include "source_estate/elecstate_pw.h"
#include "source_hamilt/hamilt.h"
#include "source_pw/module_stodft/hamilt_sdft_pw.h"
#include "source_psi/psi.h"
#include "sto_che.h"
#include "sto_func.h"
#include "sto_wf.h"
 
//----------------------------------------------
// Solve for the new electron density and iterate
// until SCF loop converges (mu, epsilon, rho)
// mu: chemical potential
// epsilon: eigenvalues
// rho: charge density
//----------------------------------------------
 
template <typename T, typename Device = base_device::DEVICE_CPU>
class Stochastic_Iter
{
  private:
    using Real = typename GetTypeReal<T>::type;
  public:
    // constructor and deconstructor
    Stochastic_Iter();
    ~Stochastic_Iter();
 
    void init(K_Vectors* pkv_in,
              ModulePW::PW_Basis_K* wfc_basis,
              Stochastic_WF<T, Device>& stowf,
              StoChe<Real, Device>& stoche,
              hamilt::HamiltSdftPW<T, Device>* p_hamilt_sto);
 
    void sum_stoeband(Stochastic_WF<T, Device>& stowf,
                     elecstate::ElecStatePW<T, Device>* pes,
                     hamilt::Hamilt<T, Device>* pHamilt,
                     ModulePW::PW_Basis_K* wfc_basis);
 
    void cal_storho(const UnitCell& ucell,
                    Stochastic_WF<T, Device>& stowf,
                    elecstate::ElecStatePW<T, Device>* pes,
                    ModulePW::PW_Basis_K* wfc_basis);
 
    double calne(elecstate::ElecState* pes);
 
    void itermu(const int iter, elecstate::ElecState* pes);
 
    void orthog(const int& ik, psi::Psi<T, Device>& psi, Stochastic_WF<T, Device>& stowf);
 
    void checkemm(const int& ik, const int istep, const int iter, Stochastic_WF<T, Device>& stowf);
 
    void check_precision(const double ref, const double thr, const std::string info);
 
    ModuleBase::Chebyshev<double, Device>* p_che = nullptr;
 
    Sto_Func<double> stofunc;
    hamilt::HamiltSdftPW<T, Device>* p_hamilt_sto = nullptr;
 
    double mu0; // chemical potential; unit in Ry
    bool change;
    double targetne=0.0;
    Real* spolyv = nullptr;     //[Device] coefficients of Chebyshev expansion
    Real* spolyv_cpu = nullptr; //[CPU] coefficients of Chebyshev expansion
 
  public:
    int* nchip = nullptr;
    bool check = false;
    double th_ne=0.0;
    double KS_ne=0.0;
 
  public:
    int method; // different methods 1: slow, less memory  2: fast, more memory
    // cal shchi = \sqrt{f(\hat{H})}|\chi>
    void calHsqrtchi(Stochastic_WF<T, Device>& stowf);
    // cal Pn = \sum_\chi <\chi|Tn(\hat{h})|\chi>
    void calPn(const int& ik, Stochastic_WF<T, Device>& stowf);
    // cal Tnchi = \sum_n C_n*T_n(\hat{h})|\chi>
    void calTnchi_ik(const int& ik, Stochastic_WF<T, Device>& stowf);
 
  private:
    K_Vectors* pkv=nullptr;
    void dot(const int& n, const Real* x, const int& incx, const Real* y, const int& incy,  Real& result);
  private:
    const Device* ctx = {};
    const base_device::DEVICE_CPU* cpu_ctx = {};
    using ct_Device = typename container::PsiToContainer<Device>::type;
    using setmem_var_op = base_device::memory::set_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 cpymem_complex_op = base_device::memory::synchronize_memory_op<T, Device, Device>;
    using resmem_var_op = base_device::memory::resize_memory_op<Real, Device>;
    using delmem_var_op = base_device::memory::delete_memory_op<Real, Device>;
    using resmem_complex_op = base_device::memory::resize_memory_op<T, Device>;
    using delmem_complex_op = base_device::memory::delete_memory_op<T, Device>;
    using castmem_d2z_op = base_device::memory::cast_memory_op<T, Real, Device, Device>;
    using castmem_var_d2h_op = base_device::memory::cast_memory_op<double, Real, base_device::DEVICE_CPU, Device>;
    using gemv_op = ModuleBase::gemv_op<T, Device>;
};
 
#endif // Eelectrons_Iter