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);
            double vp_vol_fac = cvg.vp_vol_fac(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)*vp_vol_fac;

                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
        }
};


inline VertexList vertices_in_range(const MagneticField<DeviceType>& magnetic_field, const Grid<DeviceType>& grid, bool exclude_private_region, double inner_psi_bound, double outer_psi_bound){
    return VertexList(grid.nnode, KOKKOS_LAMBDA(const int inode){
        if(inner_psi_bound <= grid.psi(inode) &&
           outer_psi_bound >= grid.psi(inode) &&
           grid.node_is_in_included_region(inode, exclude_private_region))
        {
            return true;
        }else{
            return false;
        }
    });
}

template<class Device>
class CollisionGrid{

    public:

    int batch_size;

    Collisions::LinAlgBackend labackend;

    double ginkgo_residual_reduction;
    int ginkgo_max_iterations;

    int start_step; 

    int max_n_subcycles; 
    View<int*,CLayout,HostType> n_subcycles; 
    VertexList vertices;

    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, const Grid<DeviceType>& grid, bool exclude_private_region, bool overwrite_existing_files)
      : n_subcycles(NoInit("n_subcycles"), grid.nnode)
    {
        nlr.use_namelist("col_param");
        double inner_psi_bound = nlr.get<double>("col_pin",magnetic_field.inpsi); // Minimum of psi range where collisions are performed.
        double outer_psi_bound = nlr.get<double>("col_pout",magnetic_field.outpsi); // Maximum of psi range where collisions are performed.

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

        max_n_subcycles = nlr.get<int>("col_max_n_subcycles", 1); // Enables collision subcycling. Collision time step of an unconverged vertex will be reduced, increasing the number of cycles of that vertex up to a maximum of col_max_n_subcycles. If it still does not converge, the vertex will not be attempted again.

        start_step = nlr.get<int>("col_f_start", 1); // The time step at which collisions begin

        // Start off with every vertex having 1 subcycle
        Kokkos::deep_copy(n_subcycles, 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

        // Set default solver option
#if defined(USE_GINKGO) && defined(USE_GPU)
        labackend = Collisions::LinAlgBackend::ginkgo;
#else
        labackend = Collisions::LinAlgBackend::lapack;
#endif

        if(nlr.namelist_present("performance_param")){
            nlr.use_namelist("performance_param");
            batch_size = nlr.get<int>("mesh_batch_size", default_batch_size); // Collisions are batched (multiple vertices run in parallel) for better performance. The best value may vary by platform. Typically higher is better (for GPU simulations) but memory usage is linearly proportional to this value.
            if(nlr.present("collisions_solver")){
                std::string linalg_backend_str = nlr.get<std::string>("collisions_solver", "lapack"); // Which collisions solver to use.
                                                                                                      // "lapack" is always used for CPU-only simulations
                                                                                                      // "ginkgo" is available for GPU and improves performance
                if(linalg_backend_str=="lapack") labackend = Collisions::LinAlgBackend::lapack;
                else if(linalg_backend_str=="ginkgo") labackend = Collisions::LinAlgBackend::ginkgo;
                else exit_XGC("\nError: Invalid input option. collisions_solver must be 'lapack' or 'ginkgo'.\n");
            }
            ginkgo_residual_reduction = nlr.get<double>("ginkgo_residual_reduction", 1.0e-16); // For the Ginkgo solver, the residual reduction
            ginkgo_max_iterations = nlr.get<int>("ginkgo_max_iterations", 300); // For the Ginkgo solver, the maximum number of iterations
        }else{
            batch_size = default_batch_size;
            ginkgo_residual_reduction = 1.0e-16;
            ginkgo_max_iterations = 300;
        }

#ifndef USE_GINKGO
        if(labackend == Collisions::LinAlgBackend::ginkgo) exit_XGC("\nError: collisions_solver=='ginkgo' was specified, but XGC was configured without Ginkgo.\n");
#endif
#ifndef USE_GPU
        if(labackend == Collisions::LinAlgBackend::ginkgo) exit_XGC("\nError: The CPU-only version of XGC cannot use collisions_solver=='ginkgo'.\n");
#endif

        if(is_rank_zero()){
            printf("\nUsing %s for collision solver.\n", (labackend == Collisions::LinAlgBackend::ginkgo ? "Ginkgo" : "LAPACK"));
        }


        // Save list of vertices on which to do collisions
        vertices = vertices_in_range(magnetic_field, grid, exclude_private_region, inner_psi_bound, outer_psi_bound);

        // 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");
#ifdef USE_MPI
            io_stream->Open("xgc.col_conv_status.bp", XGC_IO_Mode::Write, SML_COMM_WORLD);
#else
            io_stream->Open("xgc.col_conv_status.bp", XGC_IO_Mode::Write);
#endif
        }
    }

    int largest_n_subcycles() const{
        return (max_n_subcycles>1 ? max_view(n_subcycles) : 1);
    }

    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");
        bool sml_electron_on = nlr.get<bool>("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 VertexList& assigned, 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