<commit_before>#include "potential.h"

namespace sirius {

#ifdef _GPU_
extern "C" void generate_phase_factors_gpu(int num_gvec_loc__,
                                           int num_atoms__,
                                           int const* gvec__,
                                           double const* atom_pos__,
                                           cuDoubleComplex* phase_factors__);
#endif

/** The following operation is performed:
 *  \f[
 *    q_{\ell m}^{\alpha} = \sum_{\bf G} 4\pi \rho({\bf G}) e^{i{\bf G}{\bf r}_{\alpha}}i^{\ell}f_{\ell}^{\alpha}(G) Y_{\ell m}^{*}(\hat{\bf G})
 *  \f]
 */
void Potential::poisson_sum_G(int lmmax__,
                              double_complex* fpw__,
                              mdarray<double, 3>& fl__,
                              matrix<double_complex>& flm__)
{
    Timer t("sirius::Potential::poisson_sum_G");
    
    auto rl = parameters_.reciprocal_lattice();
    auto uc = parameters_.unit_cell();
    int ngv_loc = (int)rl->spl_num_gvec().local_size();

    int na_max = 0;
    for (int iat = 0; iat < uc->num_atom_types(); iat++) na_max = std::max(na_max, uc->atom_type(iat)->num_atoms());
    
    matrix<double_complex> phase_factors(ngv_loc, na_max);
    matrix<double_complex> zm(lmmax__, ngv_loc);
    matrix<double_complex> tmp(lmmax__, na_max);

    if (parameters_.processing_unit() == CPU)
    {
        for (int iat = 0; iat < uc->num_atom_types(); iat++)
        {
            int na = uc->atom_type(iat)->num_atoms();
            #pragma omp parallel for
            for (int igloc = 0; igloc < ngv_loc; igloc++)
            {
                int ig = parameters_.reciprocal_lattice()->spl_num_gvec(igloc);
                for (int i = 0; i < na; i++)
                {
                    int ia = uc->atom_type(iat)->atom_id(i);
                    phase_factors(igloc, i) = rl->gvec_phase_factor<local>(igloc, ia);
                }
                for (int lm = 0; lm < lmmax__; lm++)
                {
                    int l = l_by_lm_[lm];
                    zm(lm, igloc) = fourpi * fpw__[ig] * zilm_[lm] *
                                    fl__(l, iat, rl->gvec_shell(ig)) * std::conj(rl->gvec_ylm(lm, igloc));
                }
            }
            linalg<CPU>::gemm(0, 0, lmmax__, na, ngv_loc, zm.at<CPU>(), zm.ld(), phase_factors.at<CPU>(), phase_factors.ld(),
                              tmp.at<CPU>(), tmp.ld());
            for (int i = 0; i < na; i++)
            {
                int ia = uc->atom_type(iat)->atom_id(i);
                for (int lm = 0; lm < lmmax__; lm++) flm__(lm, ia) = tmp(lm, i);
            }
        }
    }

    if (parameters_.processing_unit() == GPU)
    {
        #ifdef _GPU_
        auto gvec = mdarray<int, 2>(3, ngv_loc);
        for (int igloc = 0; igloc < ngv_loc; igloc++)
        {
            for (int x = 0; x < 3; x++) gvec(x, igloc) = rl->gvec(rl->spl_num_gvec(igloc))[x];
        }
        gvec.allocate_on_device();
        gvec.copy_to_device();

        phase_factors.allocate_on_device();
        zm.allocate_on_device();
        tmp.allocate_on_device();

        double_complex alpha(1, 0);
        double_complex beta(0, 0);

        for (int iat = 0; iat < uc->num_atom_types(); iat++)
        {
            int na = uc->atom_type(iat)->num_atoms();
            
            mdarray<double, 2> atom_pos(3, na);
            for (int i = 0; i < na; i++)
            {
                int ia = uc->atom_type(iat)->atom_id(i);
                for (int x = 0; x < 3; x++) atom_pos(x, i) = uc->atom(ia)->position(x);
            }
            atom_pos.allocate_on_device();
            atom_pos.copy_to_device();

            generate_phase_factors_gpu(ngv_loc, na, gvec.at<GPU>(), atom_pos.at<GPU>(), phase_factors.at<GPU>());

            #pragma omp parallel for
            for (int igloc = 0; igloc < ngv_loc; igloc++)
            {
                int ig = rl->spl_num_gvec(igloc);
                for (int lm = 0; lm < lmmax__; lm++)
                {
                    int l = l_by_lm_[lm];
                    zm(lm, igloc) = fourpi * fpw__[ig] * zilm_[lm] *
                                    fl__(l, iat, rl->gvec_shell(ig)) * std::conj(rl->gvec_ylm(lm, igloc));
                }
            }
            zm.copy_to_device();
            linalg<GPU>::gemm(0, 0, lmmax__, na, ngv_loc, &alpha, zm.at<GPU>(), zm.ld(), phase_factors.at<GPU>(), phase_factors.ld(),
                              &beta, tmp.at<GPU>(), tmp.ld());
            tmp.copy_to_host();
            for (int i = 0; i < na; i++)
            {
                int ia = uc->atom_type(iat)->atom_id(i);
                for (int lm = 0; lm < lmmax__; lm++) flm__(lm, ia) = tmp(lm, i);
            }
        }
        #else
        TERMINATE_NO_GPU
        #endif
    }
    
    parameters_.comm().allreduce(&flm__(0, 0), (int)flm__.size());
}

void Potential::poisson_add_pseudo_pw(mdarray<double_complex, 2>& qmt, mdarray<double_complex, 2>& qit, double_complex* rho_pw)
{
    Timer t("sirius::Potential::poisson_add_pseudo_pw");
    std::vector<double_complex> pseudo_pw(parameters_.reciprocal_lattice()->num_gvec());
    memset(&pseudo_pw[0], 0, parameters_.reciprocal_lattice()->num_gvec() * sizeof(double_complex));
    
    /* The following term is added to the plane-wave coefficients of the charge density:
     * Integrate[SphericalBesselJ[l,a*x]*p[x,R]*x^2,{x,0,R},Assumptions->{l>=0,n>=0,R>0,a>0}] / 
     *  Integrate[p[x,R]*x^(2+l),{x,0,R},Assumptions->{h>=0,n>=0,R>0}]
     * i.e. contributon from pseudodensity to l-th channel of plane wave expansion multiplied by 
     * the difference bethween true and interstitial-in-the-mt multipole moments and divided by the 
     * moment of the pseudodensity.
     */
    #pragma omp parallel default(shared)
    {
        std::vector<double_complex> pseudo_pw_pt(parameters_.reciprocal_lattice()->spl_num_gvec().local_size(), double_complex(0, 0));

        #pragma omp for
        for (int ia = 0; ia < parameters_.unit_cell()->num_atoms(); ia++)
        {
            int iat = parameters_.unit_cell()->atom(ia)->type_id();

            double R = parameters_.unit_cell()->atom(ia)->mt_radius();

            /* compute G-vector independent prefactor */
            std::vector<double_complex> zp(parameters_.lmmax_rho());
            for (int l = 0, lm = 0; l <= parameters_.lmax_rho(); l++)
            {
                for (int m = -l; m <= l; m++, lm++)
                    zp[lm] = (qmt(lm, ia) - qit(lm, ia)) * conj(zil_[l]) * gamma_factors_R_(l, iat);
            }

            for (int igloc = 0; igloc < (int)parameters_.reciprocal_lattice()->spl_num_gvec().local_size(); igloc++)
            {
                int ig = parameters_.reciprocal_lattice()->spl_num_gvec(igloc);
                
                double gR = parameters_.reciprocal_lattice()->gvec_len(ig) * R;
                
                double_complex zt = fourpi * conj(parameters_.reciprocal_lattice()->gvec_phase_factor<local>(igloc, ia)) / parameters_.unit_cell()->omega();

                if (ig)
                {
                    double_complex zt2(0, 0);
                    for (int l = 0, lm = 0; l <= parameters_.lmax_rho(); l++)
                    {
                        double_complex zt1(0, 0);
                        for (int m = -l; m <= l; m++, lm++) zt1 += parameters_.reciprocal_lattice()->gvec_ylm(lm, igloc) * zp[lm];

                        zt2 += zt1 * sbessel_mt_(l + pseudo_density_order + 1, iat, parameters_.reciprocal_lattice()->gvec_shell(ig));
                    }

                    pseudo_pw_pt[igloc] += zt * zt2 * pow(2.0 / gR, pseudo_density_order + 1);
                }
                else // for |G|=0
                {
                    pseudo_pw_pt[igloc] += zt * y00 * (qmt(0, ia) - qit(0, ia));
                }
            }
        }
        #pragma omp critical
        for (int igloc = 0; igloc < (int)parameters_.reciprocal_lattice()->spl_num_gvec().local_size(); igloc++) 
            pseudo_pw[parameters_.reciprocal_lattice()->spl_num_gvec(igloc)] += pseudo_pw_pt[igloc];
    }

    parameters_.comm().allgather(&pseudo_pw[0], (int)parameters_.reciprocal_lattice()->spl_num_gvec().global_offset(), 
                                 (int)parameters_.reciprocal_lattice()->spl_num_gvec().local_size());
        
    // add pseudo_density to interstitial charge density; now rho(G) has the correct multipole moments in the muffin-tins
    for (int ig = 0; ig < parameters_.reciprocal_lattice()->num_gvec(); ig++) rho_pw[ig] += pseudo_pw[ig];
}

void Potential::poisson_vmt(Periodic_function<double>* rho__, 
                            Periodic_function<double>* vh__,
                            mdarray<double_complex, 2>& qmt__)
{
    Timer t("sirius::Potential::poisson_vmt");

    qmt__.zero();
    
    for (int ialoc = 0; ialoc < (int)parameters_.unit_cell()->spl_num_atoms().local_size(); ialoc++)
    {
        int ia = parameters_.unit_cell()->spl_num_atoms(ialoc);

        double R = parameters_.unit_cell()->atom(ia)->mt_radius();
        int nmtp = parameters_.unit_cell()->atom(ia)->num_mt_points();

        std::vector<double> qmt(parameters_.lmmax_rho(), 0);
       
        #pragma omp parallel default(shared)<commit_msg>simplify notation<commit_after>#include "potential.h"

namespace sirius {

#ifdef _GPU_
extern "C" void generate_phase_factors_gpu(int num_gvec_loc__,
                                           int num_atoms__,
                                           int const* gvec__,
                                           double const* atom_pos__,
                                           cuDoubleComplex* phase_factors__);
#endif

/** The following operation is performed:
 *  \f[
 *    q_{\ell m}^{\alpha} = \sum_{\bf G} 4\pi \rho({\bf G}) e^{i{\bf G}{\bf r}_{\alpha}}i^{\ell}f_{\ell}^{\alpha}(G) Y_{\ell m}^{*}(\hat{\bf G})
 *  \f]
 */
void Potential::poisson_sum_G(int lmmax__,
                              double_complex* fpw__,
                              mdarray<double, 3>& fl__,
                              matrix<double_complex>& flm__)
{
    Timer t("sirius::Potential::poisson_sum_G");
    
    auto rl = parameters_.reciprocal_lattice();
    auto uc = parameters_.unit_cell();
    int ngv_loc = (int)rl->spl_num_gvec().local_size();

    int na_max = 0;
    for (int iat = 0; iat < uc->num_atom_types(); iat++) na_max = std::max(na_max, uc->atom_type(iat)->num_atoms());
    
    matrix<double_complex> phase_factors(ngv_loc, na_max);
    matrix<double_complex> zm(lmmax__, ngv_loc);
    matrix<double_complex> tmp(lmmax__, na_max);

    if (parameters_.processing_unit() == CPU)
    {
        for (int iat = 0; iat < uc->num_atom_types(); iat++)
        {
            int na = uc->atom_type(iat)->num_atoms();
            #pragma omp parallel for
            for (int igloc = 0; igloc < ngv_loc; igloc++)
            {
                int ig = parameters_.reciprocal_lattice()->spl_num_gvec(igloc);
                for (int i = 0; i < na; i++)
                {
                    int ia = uc->atom_type(iat)->atom_id(i);
                    phase_factors(igloc, i) = rl->gvec_phase_factor<local>(igloc, ia);
                }
                for (int lm = 0; lm < lmmax__; lm++)
                {
                    int l = l_by_lm_[lm];
                    zm(lm, igloc) = fourpi * fpw__[ig] * zilm_[lm] *
                                    fl__(l, iat, rl->gvec_shell(ig)) * std::conj(rl->gvec_ylm(lm, igloc));
                }
            }
            linalg<CPU>::gemm(0, 0, lmmax__, na, ngv_loc, zm.at<CPU>(), zm.ld(), phase_factors.at<CPU>(), phase_factors.ld(),
                              tmp.at<CPU>(), tmp.ld());
            for (int i = 0; i < na; i++)
            {
                int ia = uc->atom_type(iat)->atom_id(i);
                for (int lm = 0; lm < lmmax__; lm++) flm__(lm, ia) = tmp(lm, i);
            }
        }
    }

    if (parameters_.processing_unit() == GPU)
    {
        #ifdef _GPU_
        auto gvec = mdarray<int, 2>(3, ngv_loc);
        for (int igloc = 0; igloc < ngv_loc; igloc++)
        {
            for (int x = 0; x < 3; x++) gvec(x, igloc) = rl->gvec(rl->spl_num_gvec(igloc))[x];
        }
        gvec.allocate_on_device();
        gvec.copy_to_device();

        phase_factors.allocate_on_device();
        zm.allocate_on_device();
        tmp.allocate_on_device();

        double_complex alpha(1, 0);
        double_complex beta(0, 0);

        for (int iat = 0; iat < uc->num_atom_types(); iat++)
        {
            int na = uc->atom_type(iat)->num_atoms();
            
            mdarray<double, 2> atom_pos(3, na);
            for (int i = 0; i < na; i++)
            {
                int ia = uc->atom_type(iat)->atom_id(i);
                for (int x = 0; x < 3; x++) atom_pos(x, i) = uc->atom(ia)->position(x);
            }
            atom_pos.allocate_on_device();
            atom_pos.copy_to_device();

            generate_phase_factors_gpu(ngv_loc, na, gvec.at<GPU>(), atom_pos.at<GPU>(), phase_factors.at<GPU>());

            #pragma omp parallel for
            for (int igloc = 0; igloc < ngv_loc; igloc++)
            {
                int ig = rl->spl_num_gvec(igloc);
                for (int lm = 0; lm < lmmax__; lm++)
                {
                    int l = l_by_lm_[lm];
                    zm(lm, igloc) = fourpi * fpw__[ig] * zilm_[lm] *
                                    fl__(l, iat, rl->gvec_shell(ig)) * std::conj(rl->gvec_ylm(lm, igloc));
                }
            }
            zm.copy_to_device();
            linalg<GPU>::gemm(0, 0, lmmax__, na, ngv_loc, &alpha, zm.at<GPU>(), zm.ld(), phase_factors.at<GPU>(), phase_factors.ld(),
                              &beta, tmp.at<GPU>(), tmp.ld());
            tmp.copy_to_host();
            for (int i = 0; i < na; i++)
            {
                int ia = uc->atom_type(iat)->atom_id(i);
                for (int lm = 0; lm < lmmax__; lm++) flm__(lm, ia) = tmp(lm, i);
            }
        }
        #else
        TERMINATE_NO_GPU
        #endif
    }
    
    parameters_.comm().allreduce(&flm__(0, 0), (int)flm__.size());
}

void Potential::poisson_add_pseudo_pw(mdarray<double_complex, 2>& qmt, mdarray<double_complex, 2>& qit, double_complex* rho_pw)
{
    Timer t("sirius::Potential::poisson_add_pseudo_pw");
    auto rl = parameters_.reciprocal_lattice();
    auto uc = parameters_.unit_cell();
    std::vector<double_complex> pseudo_pw(rl->num_gvec());
    memset(&pseudo_pw[0], 0, rl->num_gvec() * sizeof(double_complex));
    
    /* The following term is added to the plane-wave coefficients of the charge density:
     * Integrate[SphericalBesselJ[l,a*x]*p[x,R]*x^2,{x,0,R},Assumptions->{l>=0,n>=0,R>0,a>0}] / 
     *  Integrate[p[x,R]*x^(2+l),{x,0,R},Assumptions->{h>=0,n>=0,R>0}]
     * i.e. contributon from pseudodensity to l-th channel of plane wave expansion multiplied by 
     * the difference bethween true and interstitial-in-the-mt multipole moments and divided by the 
     * moment of the pseudodensity.
     */
    #pragma omp parallel default(shared)
    {
        std::vector<double_complex> pseudo_pw_pt(rl->spl_num_gvec().local_size(), double_complex(0, 0));

        #pragma omp for
        for (int ia = 0; ia < uc->num_atoms(); ia++)
        {
            int iat = uc->atom(ia)->type_id();

            double R = uc->atom(ia)->mt_radius();

            /* compute G-vector independent prefactor */
            std::vector<double_complex> zp(parameters_.lmmax_rho());
            for (int l = 0, lm = 0; l <= parameters_.lmax_rho(); l++)
            {
                for (int m = -l; m <= l; m++, lm++)
                    zp[lm] = (qmt(lm, ia) - qit(lm, ia)) * conj(zil_[l]) * gamma_factors_R_(l, iat);
            }

            for (int igloc = 0; igloc < (int)rl->spl_num_gvec().local_size(); igloc++)
            {
                int ig = rl->spl_num_gvec(igloc);
                
                double gR = rl->gvec_len(ig) * R;
                
                double_complex zt = fourpi * conj(rl->gvec_phase_factor<local>(igloc, ia)) / uc->omega();

                if (ig)
                {
                    double_complex zt2(0, 0);
                    for (int l = 0, lm = 0; l <= parameters_.lmax_rho(); l++)
                    {
                        double_complex zt1(0, 0);
                        for (int m = -l; m <= l; m++, lm++) zt1 += rl->gvec_ylm(lm, igloc) * zp[lm];

                        zt2 += zt1 * sbessel_mt_(l + pseudo_density_order + 1, iat, rl->gvec_shell(ig));
                    }

                    pseudo_pw_pt[igloc] += zt * zt2 * pow(2.0 / gR, pseudo_density_order + 1);
                }
                else // for |G|=0
                {
                    pseudo_pw_pt[igloc] += zt * y00 * (qmt(0, ia) - qit(0, ia));
                }
            }
        }
        #pragma omp critical
        for (int igloc = 0; igloc < (int)rl->spl_num_gvec().local_size(); igloc++) 
            pseudo_pw[rl->spl_num_gvec(igloc)] += pseudo_pw_pt[igloc];
    }

    parameters_.comm().allgather(&pseudo_pw[0], (int)rl->spl_num_gvec().global_offset(), (int)rl->spl_num_gvec().local_size());
        
    /* add pseudo_density to interstitial charge density;
     * now rho(G) has the correct multipole moments in the muffin-tins
     */
    for (int ig = 0; ig < rl->num_gvec(); ig++) rho_pw[ig] += pseudo_pw[ig];
}

void Potential::poisson_vmt(Periodic_function<double>* rho__, 
                            Periodic_function<double>* vh__,
                            mdarray<double_complex, 2>& qmt__)
{
    Timer t("sirius::Potential::poisson_vmt");

    qmt__.zero();
    
    for (int ialoc = 0; ialoc < (int)parameters_.unit_cell()->spl_num_atoms().local_size(); ialoc++)
    {
        int ia = parameters_.unit_cell()->spl_num_atoms(ialoc);

        double R = parameters_.unit_cell()->atom(ia)->mt_radius();
        int nmtp = parameters_.unit_cell()->atom(ia)->num_mt_points();

        std::vector<double> qmt(parameters_.lmmax_rho(), 0);
       
        #pragma omp parallel default(shared)
{'fraction_non_alphanumeric': 0.1502810536757412, 'fraction_numerical': 0.009148021602557038, 'mean_word_length': 1.8372420262664164, 'pattern_counts': {'":': 0, '<': 111, '<?xml version=': 0, '>': 189, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 21, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6783085', 'n_tokens_mistral': 6394, 'n_tokens_neox': 6013, 'n_words': 1198}