esolver_of.cpp

#include "esolver_of.h"

#include "source_io/module_parameter/parameter.h"
#include "source_io/cube_io.h"
#include "source_io/output_log.h"
#include "source_io/write_elecstat_pot.h"
//-----------temporary-------------------------
#include "source_base/global_function.h"
#include "source_estate/module_charge/symmetry_rho.h"
#include "source_hamilt/module_ewald/H_Ewald_pw.h"
#include "source_pw/module_pwdft/global.h"
#include "source_io/print_info.h"
#include "source_estate/cal_ux.h"
#include "source_pw/module_pwdft/forces.h"
#include "source_pw/module_ofdft/of_stress_pw.h"
// mohan add
#include "source_pw/module_ofdft/of_print_info.h"

namespace ModuleESolver
{

ESolver_OF::ESolver_OF()
{
    this->classname = "ESolver_OF";
    this->task_ = new char[60];
}

ESolver_OF::~ESolver_OF()
{
	//****************************************************
	// do not add any codes in this deconstructor funcion
	//****************************************************
	delete psi_;
    delete[] this->pphi_;

    for (int i = 0; i < PARAM.inp.nspin; ++i)
    {
        delete[] this->pdirect_[i];
        delete[] this->pdLdphi_[i];
        delete[] this->pdEdphi_[i];
        delete[] this->precip_dir_[i];
    }
    delete[] this->pdirect_;
    delete[] this->pdLdphi_;
    delete[] this->pdEdphi_;
    delete[] this->precip_dir_;

    delete[] this->nelec_;
    delete[] this->theta_;
    delete[] this->task_;
    delete this->ptemp_rho_;

    delete this->kedf_manager_;

    delete this->opt_cg_;
    delete this->opt_tn_;
    delete this->opt_dcsrch_;
    delete this->opt_cg_mag_;
}

void ESolver_OF::before_all_runners(UnitCell& ucell, const Input_para& inp)
{
    ESolver_FP::before_all_runners(ucell, inp);

    // save necessary parameters
    this->of_kinetic_ = inp.of_kinetic;
    this->of_method_ = inp.of_method;
    this->of_conv_ = inp.of_conv;
    this->of_tole_ = inp.of_tole;
    this->of_tolp_ = inp.of_tolp;
    this->max_iter_ = inp.scf_nmax;
    this->dV_ = ucell.omega / this->pw_rho->nxyz;
    this->bound_cal_potential_
        = std::bind(&ESolver_OF::cal_potential, this, std::placeholders::_1, std::placeholders::_2, std::ref(ucell));

    ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "SETUP UNITCELL");

//  XC_Functional::set_xc_type(ucell.atoms[0].ncpp.xc_func);
    int func_type = XC_Functional::get_func_type();
    if (func_type > 2)
    {
        ModuleBase::WARNING_QUIT("esolver_of", "meta-GGA and Hybrid functionals are not supported by OFDFT.");
    }

    this->chr.init_rho(ucell, this->Pgrid, this->sf.strucFac, ucell.symm, &this->kv);
    this->chr.check_rho(); // check the rho

    // initialize local pseudopotential
    this->locpp.init_vloc(ucell,pw_rho);
    ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "LOCAL POTENTIAL");


    // initialize elecstate, including potential
    this->init_elecstate(ucell);

    // calculate the total local pseudopotential in real space
    this->pelec->init_scf(ucell, Pgrid, sf.strucFac, locpp.numeric, ucell.symm); // atomic_rho, v_of_rho, set_vrs

    // liuyu move here 2023-10-09
    // D in uspp need vloc, thus behind init_scf()
    // calculate the effective coefficient matrix for non-local pseudopotential projectors
    ModuleBase::matrix veff = this->pelec->pot->get_eff_v();

    ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT POTENTIAL");

    // Initialize KEDF
    // Calculate electron numbers, which will be used to initialize WT KEDF
    this->nelec_ = new double[inp.nspin];
    if (inp.nspin == 1)
    {
        this->nelec_[0] = inp.nelec;
    }
    else if (inp.nspin == 2)
    {
        // in fact, nelec_spin will not be used anymore
        this->pelec->init_nelec_spin();
        this->nelec_[0] = this->pelec->nelec_spin[0];
        this->nelec_[1] = this->pelec->nelec_spin[1];
    }
    this->kedf_manager_ = new KEDF_Manager();
    this->kedf_manager_->init(inp, this->pw_rho, this->dV_, this->nelec_[0]);
    ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT KEDF");

    // Initialize optimization methods
    this->init_opt();
    ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT OPTIMIZATION");

    this->allocate_array();
}

void ESolver_OF::runner(UnitCell& ucell, const int istep)
{
    ModuleBase::timer::tick("ESolver_OF", "runner");
    // get Ewald energy, initial rho and phi if necessary
    this->before_opt(istep, ucell);
    this->iter_ = 0;

    bool conv_esolver = false; // this conv_esolver is added by mohan 20250302 
#ifdef __MPI
    this->iter_time = MPI_Wtime();
#else
    this->iter_time = std::chrono::system_clock::now();
#endif

    while (true)
    {
        // once we get a new rho and phi, update potential
        this->update_potential(ucell);

        // calculate the energy of new rho and phi
        this->energy_llast_ = this->energy_last_;
        this->energy_last_ = this->energy_current_;
        this->energy_current_ = this->cal_energy();


        // check if the job is done
        if (this->check_exit(conv_esolver))
        {
            break;
        }

        // find the optimization direction and step lenghth theta according to the potential
        this->optimize(ucell);

        // update the rho and phi based on the direction and theta
        this->update_rho();

        this->iter_++;

        ESolver_FP::iter_finish(ucell, istep, this->iter_, conv_esolver);
    }

    this->after_opt(istep, ucell, conv_esolver);

    ModuleBase::timer::tick("ESolver_OF", "runner");
}

/**
 * @brief Prepare to optimize the charge density,
 * update elecstate, kedf, and opts if needed
 * calculate ewald energy, initialize the rho, phi, theta
 *
 * @param istep
 * @param ucell
 */
void ESolver_OF::before_opt(const int istep, UnitCell& ucell)
{
    ModuleBase::TITLE("ESolver_OF", "before_opt");
    ModuleBase::timer::tick("ESolver_OF", "before_opt");

    //! 1) call before_scf() of ESolver_FP
    ESolver_FP::before_scf(ucell, istep);



    if (ucell.cell_parameter_updated)
    {
        this->dV_ = ucell.omega / this->pw_rho->nxyz;

        // initialize elecstate, including potential
        this->init_elecstate(ucell);

        // Initialize KEDF
        this->kedf_manager_->init(PARAM.inp, this->pw_rho, this->dV_, this->nelec_[0]);

        // Initialize optimization methods
        this->init_opt();

        // Refresh the arrays
        delete this->psi_;
        this->psi_ = new psi::Psi<double>(1, PARAM.inp.nspin, 
                                          this->pw_rho->nrxx, this->pw_rho->nrxx, true);

        for (int is = 0; is < PARAM.inp.nspin; ++is)
        {
            this->pphi_[is] = this->psi_->get_pointer(is);
        }

        delete this->ptemp_rho_;
        this->ptemp_rho_ = new Charge<double>();
		this->ptemp_rho_->set_rhopw(this->pw_rho);
		const bool kin_den = this->ptemp_rho_->kin_density(); // mohan add 20251202
		this->ptemp_rho_->allocate(PARAM.inp.nspin, kin_den);

        for (int is = 0; is < PARAM.inp.nspin; ++is)
        {
            delete[] this->pdLdphi_[is];
            delete[] this->pdEdphi_[is];
            delete[] this->pdirect_[is];
            delete[] this->precip_dir_[is];
            this->pdLdphi_[is] = new double[this->pw_rho->nrxx];
            this->pdEdphi_[is] = new double[this->pw_rho->nrxx];
            this->pdirect_[is] = new double[this->pw_rho->nrxx];
            this->precip_dir_[is] = new std::complex<double>[pw_rho->npw];
        }
    }

    this->pelec->init_scf(ucell, Pgrid, sf.strucFac, locpp.numeric, ucell.symm);

    Symmetry_rho srho;
    for (int is = 0; is < PARAM.inp.nspin; is++)
    {
        srho.begin(is, this->chr, this->pw_rho, ucell.symm);
    }

    for (int is = 0; is < PARAM.inp.nspin; ++is)
    {
        if (PARAM.inp.init_chg != "file")
        {
            for (int ibs = 0; ibs < this->pw_rho->nrxx; ++ibs)
            {
                // Here we initialize rho to be uniform,
                // because the rho got by pot.init_pot -> Charge::atomic_rho may contain minus elements.
                this->chr.rho[is][ibs] = this->nelec_[is] / ucell.omega;
                this->pphi_[is][ibs] = sqrt(this->chr.rho[is][ibs]);
            }
        }
        else
        {
            for (int ibs = 0; ibs < this->pw_rho->nrxx; ++ibs)
            {
                this->pphi_[is][ibs] = sqrt(this->chr.rho[is][ibs]);
            }
        }
    }

    for (int is = 0; is < PARAM.inp.nspin; ++is)
    {
        this->pelec->eferm.set_efval(is, 0);
        this->theta_[is] = 0.;
        ModuleBase::GlobalFunc::ZEROS(this->pdLdphi_[is], this->pw_rho->nrxx);
        ModuleBase::GlobalFunc::ZEROS(this->pdEdphi_[is], this->pw_rho->nrxx);
        ModuleBase::GlobalFunc::ZEROS(this->pdirect_[is], this->pw_rho->nrxx);
    }
    if (PARAM.inp.nspin == 1)
    {
        this->theta_[0] = 0.2;
    }

    ModuleBase::timer::tick("ESolver_OF", "before_opt");
}

/**
 * @brief Get dL/dphi = dL/drho * drho/dphi = (dE/drho - mu) * 2 * phi,
 * as well as normdLdphi = sqrt{<dL/dphi|dL/dphi>}
 *
 * @param ucell
 */
void ESolver_OF::update_potential(UnitCell& ucell)
{
    // (1) get dL/dphi
    elecstate::cal_ux(ucell);

    this->pelec->pot->update_from_charge(&this->chr, &ucell); // Hartree + XC + external
    this->kedf_manager_->get_potential(this->chr.rho,
                                       this->pphi_,
                                       this->pw_rho,
                                       this->pelec->pot->get_eff_v()); // KEDF potential
    for (int is = 0; is < PARAM.inp.nspin; ++is)
    {
        const double* vr_eff = this->pelec->pot->get_eff_v(is);
        for (int ir = 0; ir < this->pw_rho->nrxx; ++ir)
        {
            this->pdEdphi_[is][ir] = vr_eff[ir];
        }
        this->pelec->eferm.set_efval(is, this->cal_mu(this->pphi_[is], this->pdEdphi_[is], this->nelec_[is]));
        for (int ir = 0; ir < this->pw_rho->nrxx; ++ir)
        {
            this->pdLdphi_[is][ir]
                = this->pdEdphi_[is][ir] - 2. * this->pelec->eferm.get_efval(is) * this->pphi_[is][ir];
        }
    }

    // (2) get the norm of dLdphi
    // ===== temporary solution of potential convergence when of_full_pw = 0 =====
    this->normdLdphi_llast_ = this->normdLdphi_last_;
    this->normdLdphi_last_ = this->normdLdphi_;
    // ===========================================================================
    this->normdLdphi_ = 0.;

    for (int is = 0; is < PARAM.inp.nspin; ++is)
    {
        this->normdLdphi_ += this->inner_product(this->pdLdphi_[is], this->pdLdphi_[is], this->pw_rho->nrxx, 1.0);
    }
    Parallel_Reduce::reduce_all(this->normdLdphi_);
    this->normdLdphi_ = sqrt(this->normdLdphi_ / this->pw_rho->nxyz / PARAM.inp.nspin);
}

/**
 * @brief Get the optimization direction (this->pdirection_) and the step length (this->theta)
 *
 * @param ucell
 */
void ESolver_OF::optimize(UnitCell& ucell)
{
    // (1) get |d0> with optimization algorithm
    this->get_direction(ucell);
    // initialize temp_phi and temp_rho used in line search
    double** ptemp_phi = new double*[PARAM.inp.nspin];
    for (int is = 0; is < PARAM.inp.nspin; ++is)
    {
        ptemp_phi[is] = new double[this->pw_rho->nrxx];
        for (int ir = 0; ir < this->pw_rho->nrxx; ++ir)
        {
            ptemp_phi[is][ir] = this->pphi_[is][ir];
            this->ptemp_rho_->rho[is][ir] = ptemp_phi[is][ir] * ptemp_phi[is][ir];
        }
    }

    // (2) rotate and renormalize the direction
    this->adjust_direction();

    // (3) make sure that dEdtheta<0 at theta = 0
    double* dEdtheta = new double[PARAM.inp.nspin]; // dE/dtheta of tempPhi
    ModuleBase::GlobalFunc::ZEROS(dEdtheta, PARAM.inp.nspin);

    this->check_direction(dEdtheta, ptemp_phi, ucell);
    // this->test_direction(dEdtheta, ptemp_phi, ucell);

    // (4) call line search to find the best theta (step length)
    this->get_step_length(dEdtheta, ptemp_phi, ucell);

    for (int is = 0; is < PARAM.inp.nspin; ++is)
    {
        delete[] ptemp_phi[is];
    }
    delete[] ptemp_phi;
    delete[] dEdtheta;
}

/**
 * @brief Update the charge density and "wavefunction" (phi) after one step of optimization
 * phi = cos(theta) * phi + sin(theta) * direction,
 * rho = phi^2
 */
void ESolver_OF::update_rho()
{
    for (int is = 0; is < PARAM.inp.nspin; ++is)
    {
        for (int ir = 0; ir < this->pw_rho->nrxx; ++ir)
        {
            this->pphi_[is][ir]
                = this->pphi_[is][ir] * cos(this->theta_[is]) + this->pdirect_[is][ir] * sin(this->theta_[is]);
            this->chr.rho[is][ir] = this->pphi_[is][ir] * this->pphi_[is][ir];
        }
    }
    // // ------------ turn on symmetry may cause instability in optimization ------------
    // if (ModuleSymmetry::Symmetry::symm_flag == 1)
    // {
    //     Symmetry_rho srho;
    //     for (int is = 0; is < PARAM.inp.nspin; is++)
    //     {
    //         srho.begin(is, *(this->chr), this->pw_rho, Pgrid, ucell.symm);
    //         for (int ibs = 0; ibs < this->pw_rho->nrxx; ++ibs)
    //         {
    //             this->pphi_[is][ibs] = sqrt(this->chr.rho[is][ibs]);
    //         }
    //     }
    // }
    // // --------------------------------------------------------------------------------
}

/**
 * @brief Check convergence, return ture if converge or iter >= max_iter_,
 * and print the necessary information
 *
 * @return exit or not
 */
bool ESolver_OF::check_exit(bool& conv_esolver)
{
    conv_esolver = false;
    bool potConv = false;
    bool potHold = false; // if normdLdphi nearly remains unchanged
    bool energyConv = false;

    if (this->normdLdphi_ < this->of_tolp_)
    {
        potConv = true;
    }
    if (this->iter_ >= 3 && std::abs(this->normdLdphi_ - this->normdLdphi_last_) < 1e-10
        && std::abs(this->normdLdphi_ - this->normdLdphi_llast_) < 1e-10)
    {
        potHold = true;
    }

    if (this->iter_ >= 3 && std::abs(this->energy_current_ - this->energy_last_) < this->of_tole_
        && std::abs(this->energy_current_ - this->energy_llast_) < this->of_tole_)
    {
        energyConv = true;
    }

    conv_esolver = (this->of_conv_ == "energy" && energyConv) || (this->of_conv_ == "potential" && potConv)
                         || (this->of_conv_ == "both" && potConv && energyConv);

    OFDFT::print_info(this->iter_, this->iter_time, this->energy_current_, this->energy_last_, 
                      this->normdLdphi_, this->pelec, this->kedf_manager_, conv_esolver);

    if (conv_esolver || this->iter_ >= this->max_iter_)
    {
        return true;
    }
    // ============ temporary solution of potential convergence ===========
    else if (this->of_conv_ == "potential" && potHold)
    {
        GlobalV::ofs_warning << "ESolver_OF WARNING: "
                             << "The convergence of potential has not been reached, but the norm of potential nearly "
                                "remains unchanged, set of_full_pw = 1 may work."
                             << std::endl;
        return true;
    }
    // ====================================================================
    else
    {
        return false;
    }
}

/**
 * @brief After optimization, output the charge density, effective potential, ..., if needed.
 *
 * @param istep
 * @param ucell
 */
void ESolver_OF::after_opt(const int istep, UnitCell& ucell, const bool conv_esolver)
{
    ModuleBase::TITLE("ESolver_OF", "after_opt");
    ModuleBase::timer::tick("ESolver_OF", "after_opt");

    //------------------------------------------------------------------
    // 1) calculate kinetic energy density and ELF
    //------------------------------------------------------------------
    if (PARAM.inp.out_elf[0] > 0)
    {
        this->kedf_manager_->get_energy_density(this->chr.rho, this->pphi_, this->pw_rho, this->chr.kin_r);
    }

    // should not be here? mohan note 2025-03-03
    for (int ir = 0; ir < this->pw_rho->nrxx; ++ir)
    {
        this->chr.rho_save[0][ir] = this->chr.rho[0][ir];
    }

    //------------------------------------------------------------------
    // 2) call after_scf() of ESolver_FP
    //------------------------------------------------------------------
    ESolver_FP::after_scf(ucell, istep, conv_esolver);

#ifdef __MLALGO
    //------------------------------------------------------------------
    // Generate data if needed
    //------------------------------------------------------------------
    if (PARAM.inp.of_ml_gene_data)
    {
        this->pelec->pot->update_from_charge(&this->chr, &ucell); // Hartree + XC + external
    this->kedf_manager_->get_potential(this->chr.rho,
                                       this->pphi_,
                                       this->pw_rho,
                                       this->pelec->pot->get_eff_v()); // KEDF potential
        
        const double* vr_eff = this->pelec->pot->get_eff_v(0);
        for (int ir = 0; ir < this->pw_rho->nrxx; ++ir)
        {
            this->pdEdphi_[0][ir] = vr_eff[ir];
        }
        this->pelec->eferm.set_efval(0, this->cal_mu(this->pphi_[0], this->pdEdphi_[0], this->nelec_[0]));

        std::cout << "Generating Training data..." << std::endl;
        std::cout << "mu = " << this->pelec->eferm.get_efval(0) << std::endl;
        this->kedf_manager_->generate_ml_target(this->chr.rho, this->pw_rho, vr_eff);
    }
#endif

    ModuleBase::timer::tick("ESolver_OF", "after_opt");
}

/**
 * @brief Output the FINAL_ETOT
 */
void ESolver_OF::after_all_runners(UnitCell& ucell)
{
    ESolver_FP::after_all_runners(ucell);
}

/**
 * @brief Calculate the total energy.
 * NOTE THIS FUNCTION SHOULD BE CALLEDD AFTER POTENTIAL HAS BEEN UPDATED
 *
 * @return total energy
 */
double ESolver_OF::cal_energy()
{
    this->pelec->cal_energies(2);
    double kinetic_energy = this->kedf_manager_->get_energy(); // kinetic energy
    double pseudopot_energy = 0.;                   // electron-ion interaction energy
    for (int is = 0; is < PARAM.inp.nspin; ++is)
    {
        pseudopot_energy += this->inner_product(this->pelec->pot->get_fixed_v(),
                                                this->chr.rho[is],
                                                this->pw_rho->nrxx,
                                                this->dV_);
    }
    Parallel_Reduce::reduce_pool(pseudopot_energy);
    this->pelec->f_en.ekinetic = kinetic_energy;
    this->pelec->f_en.e_local_pp = pseudopot_energy;
    this->pelec->f_en.etot += kinetic_energy + pseudopot_energy;
    return this->pelec->f_en.etot;
}

/**
 * @brief Calculate the force
 *
 * @param [out] force
 */
void ESolver_OF::cal_force(UnitCell& ucell, ModuleBase::matrix& force)
{
    Forces<double> ff(ucell.nat);
 
    // here nullptr is for DFT+U, which may cause bugs, mohan note 2025-11-07
    // solvent can be used? mohan ask 2025-11-07
    ff.cal_force(ucell, force, *pelec, this->pw_rho, &ucell.symm, &sf, this->solvent, nullptr, &this->locpp);
}

/**
 * @brief Calculate the stress
 *
 * @param [out] stress
 */
void ESolver_OF::cal_stress(UnitCell& ucell, ModuleBase::matrix& stress)
{
    ModuleBase::matrix kinetic_stress_;
    kinetic_stress_.create(3, 3);
    this->kedf_manager_->get_stress(ucell.omega, this->chr.rho,
                         this->pphi_, this->pw_rho, kinetic_stress_); // kinetic stress

    OF_Stress_PW ss(this->pelec, this->pw_rho);
    ss.cal_stress(stress, kinetic_stress_, ucell, &ucell.symm, this->locpp, &sf, &kv);
}
} // namespace ModuleESolver