#include "esolver.h"

#include "esolver_ks_pw.h"

#include "esolver_sdft_pw.h"

#include "source_base/module_device/device.h"

#include "source_io/module_parameter/parameter.h"

#ifdef __LCAO

#include "esolver_dm2rho.h"

#include "esolver_gets.h"

#include "esolver_ks_lcao.h"

#include "esolver_ks_lcao_tddft.h"

#include "esolver_ks_lcaopw.h"

#include "source_lcao/module_lr/esolver_lrtd_lcao.h"

extern "C"

{

#include "source_base/module_external/blacs_connector.h"

}

#endif

#include "esolver_dp.h"

#include "esolver_lj.h"

#include "esolver_of.h"

#include "source_io/module_parameter/md_parameter.h"

#include <stdexcept>

namespace ModuleESolver

{

std::string determine_type()

{

std::string esolver_type = "none";

if (PARAM.inp.basis_type == "pw")

{

if (PARAM.inp.esolver_type == "sdft")

{

esolver_type = "sdft_pw";

}

else if (PARAM.inp.esolver_type == "ofdft")

{

esolver_type = "ofdft";

}

else if (PARAM.inp.esolver_type == "ksdft")

{

esolver_type = "ksdft_pw";

}

}

else if (PARAM.inp.basis_type == "lcao_in_pw")

{

#ifdef __LCAO

if (PARAM.inp.esolver_type == "sdft")

{

esolver_type = "sdft_pw";

}

else if (PARAM.inp.esolver_type == "ksdft")

{

esolver_type = "ksdft_lip";

}

#else

ModuleBase::WARNING_QUIT("ESolver", "Calculation involving numerical orbitals must be compiled with __LCAO");

#endif

}

else if (PARAM.inp.basis_type == "lcao")

{

#ifdef __LCAO

if (PARAM.inp.esolver_type == "tddft")

{

esolver_type = "ksdft_lcao_tddft";

}

else if (PARAM.inp.esolver_type == "ksdft")

{

esolver_type = "ksdft_lcao";

}

else if (PARAM.inp.esolver_type == "ks-lr")

{

esolver_type = "ksdft_lr_lcao";

}

else if (PARAM.inp.esolver_type == "lr")

{

esolver_type = "lr_lcao";

}

#else

ModuleBase::WARNING_QUIT("ESolver", "Calculation involving numerical orbitals must be compiled with __LCAO");

#endif

}

if (PARAM.inp.esolver_type == "lj")

{

esolver_type = "lj_pot";

}

else if (PARAM.inp.esolver_type == "dp")

{

esolver_type = "dp_pot";

}

else if (esolver_type == "none")

{

ModuleBase::WARNING_QUIT("ESolver", "No such esolver_type combined with basis_type");

}

GlobalV::ofs_running << "\n #ENERGY SOLVER# " << esolver_type << std::endl;

auto device_info = PARAM.inp.device;

for (char& c: device_info)

{

if (std::islower(c))

{

c = std::toupper(c);

}

}

if (GlobalV::MY_RANK == 0)

{

std::cout << " RUNNING WITH DEVICE : " << device_info << " / "

<< base_device::information::get_device_info(PARAM.inp.device) << std::endl;

}

GlobalV::ofs_running << "\n RUNNING WITH DEVICE : " << device_info << " / "

<< base_device::information::get_device_info(PARAM.inp.device) << std::endl;

/***auto end_time = std::chrono::high_resolution_clock::now();

auto duration = std::chrono::duration_cast<std::chrono::duration<double>>(end_time - start_time);

std::cout << "hipGetDeviceInfo took " << duration.count() << " seconds" << std::endl;***/

return esolver_type;

}

// Some API to operate E_Solver

ESolver* init_esolver(const Input_para& inp, UnitCell& ucell)

{

// determine type of esolver based on INPUT information

const std::string esolver_type = determine_type();

// initialize the corresponding Esolver child class

if (esolver_type == "ksdft_pw")

{

#if ((defined __CUDA) || (defined __ROCM))

if (PARAM.inp.device == "gpu")

{

if (PARAM.inp.precision == "single")

{

return new ESolver_KS_PW<std::complex<float>, base_device::DEVICE_GPU>();

}

else

{

return new ESolver_KS_PW<std::complex<double>, base_device::DEVICE_GPU>();

}

}

#endif

if (PARAM.inp.precision == "single")

{

return new ESolver_KS_PW<std::complex<float>, base_device::DEVICE_CPU>();

}

else

{

return new ESolver_KS_PW<std::complex<double>, base_device::DEVICE_CPU>();

}

}

else if (esolver_type == "sdft_pw")

{

#if ((defined __CUDA) || (defined __ROCM))

if (PARAM.inp.device == "gpu")

{

// if (PARAM.inp.precision == "single")

// {

// return new ESolver_SDFT_PW<std::complex<float>, base_device::DEVICE_GPU>();

// }

// else

// {

return new ESolver_SDFT_PW<std::complex<double>, base_device::DEVICE_GPU>();

// }

}

#endif

// if (PARAM.inp.precision == "single")

// {

// return new ESolver_SDFT_PW<std::complex<float>, base_device::DEVICE_CPU>();

// }

// else

// {

return new ESolver_SDFT_PW<std::complex<double>, base_device::DEVICE_CPU>();

// }

}

#ifdef __LCAO

else if (esolver_type == "ksdft_lip")

{

if (PARAM.inp.precision == "single")

{

return new ESolver_KS_LIP<std::complex<float>>();

}

else

{

return new ESolver_KS_LIP<std::complex<double>>();

}

}

else if (esolver_type == "ksdft_lcao")

{

if (PARAM.inp.calculation == "get_s")

{

if (PARAM.globalv.gamma_only_local)

{

ModuleBase::WARNING_QUIT("ESolver", "get_s is not implemented for gamma_only");

}

else

{

return new ESolver_GetS();

}

}

if (PARAM.globalv.gamma_only_local)

{

return new ESolver_KS_LCAO<double, double>();

}

else if (PARAM.inp.nspin < 4)

{

if (PARAM.inp.dm_to_rho)

{

return new ESolver_DM2rho<std::complex<double>, double>();

}

else

{

return new ESolver_KS_LCAO<std::complex<double>, double>();

}

}

else

{

if (PARAM.inp.dm_to_rho)

{

return new ESolver_DM2rho<std::complex<double>, std::complex<double>>();

}

else

{

return new ESolver_KS_LCAO<std::complex<double>, std::complex<double>>();

}

}

}

else if (esolver_type == "ksdft_lcao_tddft")

{

if (PARAM.inp.nspin < 4)

{

#if ((defined __CUDA) /* || (defined __ROCM) */)

if (PARAM.inp.device == "gpu")

{

return new ESolver_KS_LCAO_TDDFT<double, base_device::DEVICE_GPU>();

}

#endif

return new ESolver_KS_LCAO_TDDFT<double, base_device::DEVICE_CPU>();

}

else

{

#if ((defined __CUDA) /* || (defined __ROCM) */)

if (PARAM.inp.device == "gpu")

{

return new ESolver_KS_LCAO_TDDFT<std::complex<double>, base_device::DEVICE_GPU>();

}

#endif

return new ESolver_KS_LCAO_TDDFT<std::complex<double>, base_device::DEVICE_CPU>();

}

}

else if (esolver_type == "lr_lcao")

{

// use constructor rather than Init function to initialize reference (instead of pointers) to ucell

if (PARAM.globalv.gamma_only_local)

{

return new LR::ESolver_LR<double, double>(inp, ucell);

}

else

{

return new LR::ESolver_LR<std::complex<double>, double>(inp, ucell);

}

}

else if (esolver_type == "ksdft_lr_lcao")

{

// initialize the 1st ESolver_KS

ModuleESolver::ESolver* p_esolver = nullptr;

if (PARAM.globalv.gamma_only_local)

{

p_esolver = new ESolver_KS_LCAO<double, double>();

}

else if (PARAM.inp.nspin < 4)

{

p_esolver = new ESolver_KS_LCAO<std::complex<double>, double>();

}

else

{

p_esolver = new ESolver_KS_LCAO<std::complex<double>, std::complex<double>>();

}

p_esolver->before_all_runners(ucell, inp);

p_esolver->runner(ucell, 0); // scf-only

// force and stress is not needed currently,

// they will be supported after the analytical gradient

// of LR-TDDFT is implemented.

std::cout << " PREPARING FOR EXCITED STATES." << std::endl;

// initialize the 2nd ESolver_LR at the temporary pointer

ModuleESolver::ESolver* p_esolver_lr = nullptr;

if (PARAM.globalv.gamma_only_local)

{

p_esolver_lr = new LR::ESolver_LR<double, double>(

std::move(*dynamic_cast<ModuleESolver::ESolver_KS_LCAO<double, double>*>(p_esolver)),

inp,

ucell);

}

else

{

p_esolver_lr = new LR::ESolver_LR<std::complex<double>, double>(

std::move(*dynamic_cast<ModuleESolver::ESolver_KS_LCAO<std::complex<double>, double>*>(p_esolver)),

inp,

ucell);

}

// clean the 1st ESolver_KS and swap the pointer

ModuleESolver::clean_esolver(p_esolver, false); // do not call Cblacs_exit, remain it for the 2nd ESolver

return p_esolver_lr;

}

#endif

else if (esolver_type == "ofdft")

{

return new ESolver_OF();

}

else if (esolver_type == "lj_pot")

{

return new ESolver_LJ();

}

else if (esolver_type == "dp_pot")

{

return new ESolver_DP(PARAM.mdp.pot_file);

}

throw std::invalid_argument("esolver_type = " + std::string(esolver_type) + ". Wrong in " + std::string(__FILE__)

+ " line " + std::to_string(__LINE__));

}

void clean_esolver(ESolver*& pesolver, const bool lcao_cblacs_exit)

{

// Zhang Xiaoyang modified in 2024/7/6:

// Note: because of the init method of serial lcao hsolver

// it needs no release step for it, or this [delete] will cause Segmentation Fault

// Probably it will be modified later.

#ifdef __MPI

delete pesolver;

#ifdef __LCAO

if (lcao_cblacs_exit)

{

Cblacs_exit(1);

}

#endif

#endif

}

} // namespace ModuleESolver