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
};
}
 
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);
}
 
}
 
 
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; 
    View<float*,CLayout,HostType> timing_all; // Vertex timing
    View<bool[1],CLayout,HostType> timing_available; // Whether timing has begun yet
    VertexList vertices;
 
    bool async_reassign;
 
    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),
        timing_all("timing_all", grid.nnode), // Init to 0
        timing_available("timing_available")
    {
        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. @parent: col_param:col_mode=4
        double outer_psi_bound = nlr.get<double>("col_pout",magnetic_field.outpsi); // Maximum of psi range where collisions are performed. @parent: col_param:col_mode=4
 
        // 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. @parent: col_param:col_mode=4
 
        start_step = nlr.get<int>("col_f_start", 1); // The time step at which collisions begin @parent: col_param:col_mode=4
 
        // Initially, no timing data is available
        timing_available(0) = false;
 
        // Start off with every vertex having 1 subcycle
        Kokkos::deep_copy(n_subcycles, 1);
 
        int default_batch_size = get_default_batch_size();
 
        // Set default solver option
#if defined(USE_GINKGO) && defined(USE_GPU)
        labackend = Collisions::LinAlgBackend::ginkgo;
#else
        labackend = Collisions::LinAlgBackend::lapack;
#endif
 
        nlr.use_namelist("performance_param", NLReader::NamelistReader::Optional);
        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. @parent: col_param:col_mode=4
        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
                                                                                                  // @parent: col_param:col_mode=4
            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 @parent: performance_param:collisions_solver=ginkgo
        ginkgo_max_iterations = nlr.get<int>("ginkgo_max_iterations", 300); // For the Ginkgo solver, the maximum number of iterations @parent: performance_param:collisions_solver=ginkgo
        async_reassign = nlr.get<bool>("col_async_reassign", false); // Asynchronously transfer collision workload between MPI ranks @parent: col_param:col_mode=4
 
#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
        }
    }
 
    bool timing_is_available() const{
        return timing_available(0);
    }
 
    static int get_default_batch_size(){
        // These defaults have not been systematically optimized
#ifdef USE_GPU
        return 64;
#else
#ifdef USE_OMP 
        int n_threads = omp_get_max_threads();
        int dbs=1; while ( dbs < n_threads ) dbs*=2; // get next power of 2 above n_threads so that all threads are busy(?)
        return dbs;
#else
        return 1;
#endif
#endif
    }
 
    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);
        int default_batch_size = get_default_batch_size();
 
        int col_grid_batch_size;
        nlr.use_namelist("performance_param", NLReader::NamelistReader::Optional);
        col_grid_batch_size = nlr.get<int>("mesh_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;
    }
 
    // 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, const CollisionSpecies<Device>& col_spall, TmpColData<Device>& tcd, const std::unique_ptr<Collisions::GridMatrix<Device>>& matrix, int mb_n_nodes, Kokkos::View<Convergence::Status*,HostType>& convergence_status) 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, const CollisionSpecies<Device>& col_spall, TmpColData<Device>& tcd, const std::unique_ptr<Collisions::GridMatrix<Device>>& solve_matrix, int mb_n_nodes) const;
 
 
    // Carries out collisions
    void collision(const CollisionSpecies<Device>& col_spall, const VertexList& assigned, const View<int*,CLayout,HostType>& n_subcycles_local, double dt, const VGridDistribution<HostType>& df_out, const View<int*,CLayout,HostType>& converged_local, const View<double*,CLayout,HostType>& node_cost) const;
};
 
#endif