col_grid.hpp

Go to the documentation of this file.

#ifndef COL_GRID_HPP
#define COL_GRID_HPP

#include <vector>
#include <stdio.h>
#include <Kokkos_Core.hpp>
#include <Kokkos_DualView.hpp>

#include "xgc_io.hpp"
#include "domain_decomposition.hpp"
#include "velocity_grid.hpp"
#include "vgrid_distribution.hpp"
#include "magnetic_field.hpp"

#include "col_grid_matrix.hpp"
#include "col_species.hpp"
#include "col_vgrids.hpp"

//> XGC's multi-species Fokker-Planck-Landau collision operator
// see R. Hager et al., Journal of Computational Physics 315 (2016), pp. 644-660
//     E. S. Yoon et al., Physics of Plasmas 21 (2014), 032503
//
// The headers of the subroutines in this module refer to the corresponding parts of Hager et al.
// where applicable in order to make the source code easier to understand.


// Enum for derivatives of ED array
namespace ED{
enum{
    RR=0,
    RZ,
    ZZ,
    ER,
    EZ,
    N
};
}

// <SOLVER-related>
const int vpic_inner_iter_max=20;

namespace Convergence{

// Ideally this would be an enum class,
// but the are MPI and I/O operations that expect ints
enum Status{
    NaNOrInfFound=-5,
    MaxedIter=-4,
    TooManyNegValues=-3,
    PosCorrectionFail=-2,
    FirstStepFail=-1,
    NotAttempted=0,
    Success=1,
    InProgress=2
};

inline bool is_okay(Status convergence_status){
    return (convergence_status!=MaxedIter &&
            convergence_status!=NaNOrInfFound &&
            convergence_status!=FirstStepFail);
}

class Moments{
    // Values
    double w_sum; //> sum of energy over species
    double p_sum; //> sum of momentum over species
    double min_p_thres; //> momentum error criterion

    // Deltas
    double dw_sum; //> sum of energy change over species (i.e. error in energy)
    double dp_sum; //> sum of momentum change over species (i.e. error in momentum)
    double dn_n_max; //> maximum fraction of density changes among species (i.e. maximum relative error in density in all species )
    bool dens_exit_ok;
    bool en_exit_ok;
    bool mom_exit_ok;

    //> Checks the convergence criterion (i.e. conservation of mass, energy and momentum)
    template<class Device>
    inline void eval(int isp, int mesh_ind, const CollisionVelocityGrids<Device>& cvg, const CollisionSpecies<Device>& cs,double& dw,double& dp,double& dn_n) const{
        double vpic_dn = 0.0;
        double vpic_dw = 0.0;
        double vpic_dp = 0.0;

        int gri = cs.s.h_view(isp, mesh_ind).grid_index;

        for (int index_i = 0; index_i<cvg.mesh_z.extent(2); index_i++){
            double mesh_z2 = cvg.mesh_z(mesh_ind,gri,index_i)*cvg.mesh_z(mesh_ind,gri,index_i);
            for (int index_j = 0; index_j<cvg.mesh_r.extent(2); index_j++){
                double mesh_r2 = cvg.mesh_r(mesh_ind,gri,index_j)*cvg.mesh_r(mesh_ind,gri,index_j);
                double vpic_dfc = (cs.pdf_np1.h_view(mesh_ind,isp,index_i,index_j) - cs.pdf_n.h_view(mesh_ind,isp,index_i,index_j)) * cvg.vol_h(mesh_ind,gri,index_j);

                vpic_dn += vpic_dfc;
                vpic_dw += vpic_dfc*(mesh_z2 + mesh_r2);
                vpic_dp += vpic_dfc*cvg.mesh_z(mesh_ind,gri,index_i);
            }
        }

        dw = vpic_dw*cs.s.h_view(isp,mesh_ind).mass;
        dp = vpic_dp*cs.s.h_view(isp,mesh_ind).mass;
        dn_n = std::abs(vpic_dn/cs.s.h_view(isp,mesh_ind).dens);
    }

    public:

    template<class Device>
    inline void set(int mesh_ind, const CollisionSpecies<Device>& col_spall){
        w_sum = 0;
        p_sum = 0;
        min_p_thres = 1e99; //just large value to be reinitialized with a small value
        for (int spi = 0; spi<col_spall.n(); spi++){
            w_sum += col_spall.s.h_view(spi,mesh_ind).ens;
            p_sum += col_spall.s.h_view(spi,mesh_ind).mom;
            min_p_thres = std::min( col_spall.s.h_view(spi,mesh_ind).mass*col_spall.s.h_view(spi,mesh_ind).dens*sqrt(col_spall.s.h_view(spi,mesh_ind).numeric_vth2),
                                    min_p_thres );
        }
    }

    template<class Device>
    inline bool evaluate_conservation(int mesh_ind, const CollisionVelocityGrids<Device>& col_vgrids, const CollisionSpecies<Device>& col_spall, bool& values_are_all_finite){
        dw_sum = 0.0;
        dp_sum = 0.0;
        dn_n_max = 0.0;
        dens_exit_ok = true;

        for (int spi = 0; spi<col_spall.n(); spi++){
            // Re-compute conservation errors after correction
            double dw,dp,dn_n;
            eval(spi, mesh_ind, col_vgrids, col_spall,dw,dp,dn_n);// moments evalulation from difference betw. pdf_n and pdf_np1
            dw_sum += dw;
            dp_sum += dp;
            if(dn_n_max<dn_n) dn_n_max = dn_n;
            dens_exit_ok = dens_exit_ok && (dn_n<1e-10);
        }

        // Check for NaNs and Infs
        values_are_all_finite = (std::isfinite(dw_sum) && std::isfinite(dp_sum) && std::isfinite(dn_n_max));

        en_exit_ok = std::abs(dw_sum/w_sum) <= 1e-7;
        //        bool mom_exit_ok = std::abs(col_dp_sum)/(std::max(std::abs(col_p_sum),min_p_thres)) <= (col_f_vp_max/col_f_nvz);// v_par  =  max_v_par of simulation domain / vth
        // ALS : momentum tolerance changed to the tolerance in XGC's former collisions module:
        // For momentum tolerance use main ion mass and electron density as normalization (when there are more than 2 species).
        double mass_use = (col_spall.s.h_view(0,mesh_ind).is_electron ? col_spall.s.h_view(1,mesh_ind).mass : col_spall.s.h_view(0,mesh_ind).mass);
        mom_exit_ok = std::abs(dp_sum)/(std::max(std::abs(p_sum),1e-3*mass_use*col_spall.s.h_view(0,mesh_ind).dens)) <= 1.0e-7;

        return (dens_exit_ok && en_exit_ok && mom_exit_ok);
    }
};

}


template<class Device>
struct TmpColData{

    // These views are both on CPU and GPU - dual views
    Kokkos::View<double***,Device> gammac_spall;
    Kokkos::DualView<int[4][4],Device> mat_pos_rel_indx; // Shouldn't need to be a dual view

    Kokkos::View<double***,Device> fhalf;
    Kokkos::View<double***,Device> dfdr;
    Kokkos::View<double***,Device> dfdz;
    Kokkos::View<double****,Device> EDs;
    Kokkos::View<double******,Device> M_ab;
    Kokkos::View<double*****,Device> M_s;

    TmpColData<Device>(int nvr, int nvz, int sp_num, int n_vgrids, int mb_n_nodes)
        : // Note: layout left
          gammac_spall("gammac_spall",mb_n_nodes,sp_num,sp_num),
          fhalf("fhalf",mb_n_nodes,nvz-1,nvr-1),
          dfdr("dfdr",mb_n_nodes,nvz-1,nvr-1),
          dfdz("dfdz",mb_n_nodes,nvz-1,nvr-1),
          EDs("EDs",mb_n_nodes,nvz-1,nvr-1,ED::N),
          M_s("M_s",mb_n_nodes,n_vgrids, (nvr-1)*(nvz-1), ED::N, nvr-1),
          M_ab("M_ab",mb_n_nodes,n_vgrids, n_vgrids-1, (nvr-1)*(nvz-1), 3, (nvr-1)*(nvz-1)),
          mat_pos_rel_indx("mat_pos_rel_indx")
        {
            // Set up mat_pos_rel_indx and send to GPU
            int mat_loc[16] = {4,5,7,8,
                               1,2,4,5,
                               3,4,6,7,
                               0,1,3,4};

            for (int i=0;i<16;i++)
                mat_pos_rel_indx.h_view(i/4,i%4)=mat_loc[i];

#ifdef USE_GPU
            Kokkos::deep_copy(mat_pos_rel_indx.d_view, mat_pos_rel_indx.h_view);
#endif
        }
};


template<class Device>
class CollisionGrid{

    public:

    int batch_size;

    Collisions::LinAlgBackend labackend;

    double ginkgo_residual_reduction;
    int ginkgo_max_iterations;

    double inner_psi_bound; 
    double outer_psi_bound; 
    
    int start_step; 

    bool diag_on; // Switches file-output of convergence status of the collision operator on/off
    std::shared_ptr<XGC_IO_Stream> io_stream;

    CollisionGrid(){}

    //> Set up global variables of the collision operator
    CollisionGrid(NLReader::NamelistReader& nlr, const MagneticField<DeviceType>& magnetic_field, bool overwrite_existing_files){
        nlr.use_namelist("col_param");
        inner_psi_bound = nlr.get<double>("col_pin",magnetic_field.inpsi);
        outer_psi_bound = nlr.get<double>("col_pout",magnetic_field.outpsi);

        // Normalize by xpt
        // Default values are already normalized, so only normalize if not default
        if(nlr.present("col_pin")) inner_psi_bound *= magnetic_field.equil.xpt_psi;
        if(nlr.present("col_pout")) outer_psi_bound *= magnetic_field.equil.xpt_psi;

        start_step = nlr.get<int>("col_f_start", 1);

        // These defaults have not been systematically optimized
#ifdef USE_GPU
        int default_batch_size = 64;
#else
        int default_batch_size = 1;//omp_get_max_threads();
#endif
        std::string linalg_backend_str;
        if(nlr.namelist_present("performance_param")){
            nlr.use_namelist("performance_param");
            batch_size = nlr.get<int>("mesh_batch_size", default_batch_size);
            linalg_backend_str = nlr.get<std::string>("collisions_solver", "lapack");
            ginkgo_residual_reduction = nlr.get<double>("ginkgo_residual_reduction", 1.0e-16);
            ginkgo_max_iterations = nlr.get<int>("ginkgo_max_iterations", 300);
        }else{
            batch_size = default_batch_size;
            linalg_backend_str = "lapack";
            ginkgo_residual_reduction = 1.0e-16;
            ginkgo_max_iterations = 300;
        }

        labackend = (linalg_backend_str=="lapack" ? Collisions::LinAlgBackend::lapack : Collisions::LinAlgBackend::ginkgo);

        // Diagnostic
        nlr.use_namelist("diag_param");
        diag_on = nlr.get<bool>("diag_col_convergence_stat_on",false); 

        if(diag_on){
            io_stream = std::make_shared<XGC_IO_Stream>();
            io_stream->Init("collision_stats");
            XGC_IO_Mode mode = (overwrite_existing_files ? XGC_IO_Mode::Write : XGC_IO_Mode::Append);
            io_stream->Open("xgc.col_conv_status.bp", mode);
        }
    }

    static MemoryPrediction estimate_memory_usage(NLReader::NamelistReader& nlr){
        nlr.use_namelist("ptl_param");
        int plasma_nspecies = nlr.get<int>("ptl_nsp",1);

        nlr.use_namelist("sml_param");
        int sml_special = nlr.get<int>("sml_special", 0);
        bool sml_electron_on = nlr.get<bool>("sml_electron_on", false);
        if(sml_special==4) sml_electron_on = false;
        if(sml_electron_on) plasma_nspecies += 1; // Add electrons

        VelocityGrid vgrid(nlr);
#ifdef USE_GPU
        int default_batch_size = 64;
#else
        int default_batch_size = 1;//omp_get_max_threads();
#endif
        int col_grid_batch_size;
        if(nlr.namelist_present("performance_param")){
            nlr.use_namelist("performance_param");
            col_grid_batch_size = nlr.get<int>("mesh_batch_size", default_batch_size);
        }else{
            col_grid_batch_size = default_batch_size;
        }

        double device_col_GB_per_vertex = plasma_nspecies*std::max(plasma_nspecies-1,1)*(vgrid.nvp*vgrid.nmu)*(vgrid.nvp*vgrid.nmu)*8*BYTES_TO_GB;
        double gpu_memory_usage = (col_grid_batch_size*device_col_GB_per_vertex);

        MemoryPrediction memory_prediction{"Fields", 0.0, gpu_memory_usage};

        return memory_prediction;
    }

    // Core init
    void core_init(int isp, int mesh_ind, const CollisionVelocityGrids<Device>& col_vgrids, const CollisionSpecies<Device>& col_spall ) const;

    // Get maxw_fac
    static KOKKOS_INLINE_FUNCTION double get_maxw_fac(double mesh_dr, double mesh_r, double numeric_vth2) ;

    // Get numeric v_thermal equilibrium for a given species against a given grid
    View<double*,Device> get_numeric_v_thermal_equil(int mb_n_nodes, int spi, int grj, const CollisionVelocityGrids<Device>& col_vgrids, const CollisionSpecies<Device>& col_spall) const;

    // Delta init
    void core_delta_init(int mb_n_nodes, int gri, int grj, int spi, CollisionVelocityGrids<Device>& col_vgrids, const CollisionSpecies<Device>& col_spall) const;

    // LU_matrix_ftn
    static KOKKOS_INLINE_FUNCTION void LU_matrix_ftn(int mesh_ind, int gri, int grj, int spi, int cell_i,int cell_j,const CollisionVelocityGrids<Device>& col_vgrids,int mprl_col, int mat_pos,double coeff1,double coeff2, const TmpColData<Device>& tcd, const Kokkos::View<int**,Device>& index_map_LU_d,typename Collisions::GridMatrix<Device>::values_array_t LU_values);

    // LU_matrix
    void LU_matrix(int mb_n_nodes, int gri, int grj, int spi, const CollisionVelocityGrids<Device>& col_vgrids,const TmpColData<Device>& tcd, Collisions::GridMatrix<Device>* const mtx) const;

    // the core loop including the picard loop is now inside this function 
    void picard_loop(int vpic_inner_iter_max, const CollisionVelocityGrids<Device>& col_vgrids, CollisionSpecies<Device>& col_spall, TmpColData<Device>& tcd, int mb_n_nodes, Kokkos::View<Convergence::Status*,HostType>& convergence_status, Kokkos::View<Convergence::Moments*,HostType>& moments) const;

    // dispatch for E_and_D calc
    void E_and_D(int mb_n_nodes, int gri, int grj, const CollisionVelocityGrids<Device>& col_vgrids, TmpColData<Device>& tcd) const;

    // E_and_D_s
    static KOKKOS_INLINE_FUNCTION void E_and_D_s(int idx, int mb_n_nodes, int nvrm1, int nvzm1, const TmpColData<Device>& tcd, int gri) ;

    // dispatch for angle_avg calc
    void angle_avg(int mb_n_nodes, int gri, int grj, CollisionVelocityGrids<Device>& col_vgrids, TmpColData<Device>& tcd) const;

    // angle_avg_s
    static KOKKOS_INLINE_FUNCTION void angle_avg_s(int idx, int mb_n_nodes, int nvrm1, int nvzm1, const CollisionVelocityGrids<Device>& col_vgrids, const TmpColData<Device>& tcd, int gri);

    // E_and_D_ab
    static KOKKOS_INLINE_FUNCTION void E_and_D_ab(int idx, int mb_n_nodes, int nvrm1, int nvzm1, const CollisionVelocityGrids<Device>& col_vgrids, const TmpColData<Device>& tcd, int gri, int grj);

    // angle_avg_ab
    static KOKKOS_INLINE_FUNCTION void angle_avg_ab(int idx, int mb_n_nodes, int nvrm1, int nvzm1, const CollisionVelocityGrids<Device>& col_vgrids, const TmpColData<Device>& tcd, int gri, int grj);

    void f_df(int mb_n_nodes, const CollisionVelocityGrids<Device>& col_vgrids, const CollisionSpecies<Device>& col_spall,int spi,int grj,TmpColData<Device>& tcd) const;

    // Main collision algorithm
    Kokkos::View<Convergence::Status*,HostType> core(CollisionVelocityGrids<Device>& col_vgrids, CollisionSpecies<Device>& col_spall, TmpColData<Device>& tcd, int mb_n_nodes) const;


    // Carries out collisions
    void collision(std::vector<Species<DeviceType>> &all_species, const Moments& moments, const DomainDecomposition<DeviceType>& pol_decomp, const View<bool*, HostType>& in_range, bool symmetric_f, const VGridDistribution<HostType>& f0_f, VGridDistribution<HostType>& df_out, double dt, View<int*,CLayout,HostType>& converged_all, View<double*,CLayout,HostType>& node_cost) const;
};

#include "col_grid.tpp"

#endif