col_vgrids.hpp

Go to the documentation of this file.

#ifndef COL_VGRIDS_HPP
#define COL_VGRIDS_HPP

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

#include "col_species.hpp"


template<class Device>
struct CollisionVelocityGrids{
    Kokkos::DualView<int*,Kokkos::LayoutRight,Device> nspecies_in_grid;
    const int n; 
    const int mb_n_nodes; 
    const int n_species; 
    const int nvr; 
    const int nvz; 
    Kokkos::View<double***,Kokkos::LayoutRight,HostType> mesh_r;
    Kokkos::View<double***,Kokkos::LayoutRight,HostType> mesh_z;
    Kokkos::View<double***,Kokkos::LayoutRight,Device> vol;
    Kokkos::View<double***,Kokkos::LayoutRight,HostType> vol_h;
    Kokkos::View<double****,Device> delta_r;
    Kokkos::View<double****,Device> delta_z;

    // These three are sent to GPU, so use a dualview
    Kokkos::DualView<double***,Device> mesh_r_half;
    Kokkos::DualView<double***,Device> mesh_z_half;
    Kokkos::DualView<double***,Device> local_center_volume;
    Kokkos::DualView<double**,Device> vpar_beg;
    Kokkos::DualView<double**,Device> mesh_dz;
    Kokkos::DualView<double**,Device> mesh_dr;

    Kokkos::DualView<int**,Kokkos::LayoutRight,Device> map_grid_to_species;

    // Count how many species are using each grid
    inline Kokkos::DualView<int*,Kokkos::LayoutRight,Device> count_species_in_grid(const CollisionSpecies<Device>& col_spall) const{
        Kokkos::DualView<int*,Kokkos::LayoutRight,Device> nspecies_in_grid_tmp("nspecies_in_grid", col_spall.s.h_view.extent(0));
        // Loop over species and set up mapping
        for (int isp = 0; isp < col_spall.s.h_view.extent(0); isp++){
            int gri = col_spall.s.h_view(isp,0).grid_index; // Which grid this species is on
            nspecies_in_grid_tmp.h_view(gri)++;
        }
        return nspecies_in_grid_tmp;
    }

    // Count how many grids are used
    inline int count_grids() const{
        int ngrids=0;
        for(int i=0; i<nspecies_in_grid.h_view.size(); i++){
            if(nspecies_in_grid.h_view(i)>0) ngrids++;
        }
        return ngrids;
    }

    CollisionVelocityGrids<Device>(){}

    CollisionVelocityGrids<Device>(const CollisionSpecies<Device>& col_spall)
        : nspecies_in_grid(count_species_in_grid(col_spall)),
          n(count_grids()),
          mb_n_nodes(col_spall.pdf_n.h_view.extent(0)), // Use pdf_n dimensions: mb_n_nodes, n_species, nvz, nvr
          n_species(col_spall.pdf_n.h_view.extent(1)),
          nvz(col_spall.pdf_n.h_view.extent(2)),
          nvr(col_spall.pdf_n.h_view.extent(3)),
          map_grid_to_species("map_grid_to_species",n,n_species), // n_species is the max possible
          mesh_dr("mesh_dr", mb_n_nodes, n),
          mesh_dz("mesh_dz", mb_n_nodes, n),
          vpar_beg("vpar_beg", mb_n_nodes, n),
          mesh_r("mesh_r", mb_n_nodes, n, nvr),
          mesh_z("mesh_z", mb_n_nodes, n, nvz),
          mesh_r_half("mesh_r_half",mb_n_nodes,nvr-1,n),
          mesh_z_half("mesh_z_half",mb_n_nodes,nvz-1,n),
          local_center_volume("local_center_volume",mb_n_nodes,nvr-1,n),
          vol("vol",nvr,n, mb_n_nodes),
          vol_h("vol_h",mb_n_nodes,n, nvr), // Must have nvr fastest index for passing into fortran negative_f_correction routine
          delta_r("delta_r",n,nvr-1,n_species, mb_n_nodes),
          delta_z("delta_z",n,nvz-1,n_species, mb_n_nodes)
    {
        // Loop over species and set up mapping
        int isp_in_grid=0;
        int gri=0;
        for (int isp = 0; isp < n_species; isp++){
            // Move to next species or to next grid
            if(gri!=col_spall.s.h_view(isp,0).grid_index) isp_in_grid=0;

            gri = col_spall.s.h_view(isp,0).grid_index; // Which grid this species is on
            map_grid_to_species.h_view(gri,isp_in_grid) = isp;
            isp_in_grid++;
        }

#ifdef USE_GPU
        // Copy mapping views to device
        Kokkos::deep_copy(map_grid_to_species.d_view, map_grid_to_species.h_view);
        Kokkos::deep_copy(nspecies_in_grid.d_view, nspecies_in_grid.h_view);
#endif
    }

    //> Convert the distribution function to the units used
    // internally by the collision operator
    // The distribution function set up with f0_module is in normalized velocity space
    // whereas the collision operator uses SI units.
    void setup_one(int isp, int mesh_ind, double vth, double dsmu, double dvp, double vp_max) {

        // vol
        double pi4bb0vth3_dd = TWOPI*vth*vth*vth*dsmu*dvp;

        // --- 0th component
        double smu_n = dsmu/3.0;
        vol_h(mesh_ind,isp,0) = 0.5*pi4bb0vth3_dd*smu_n;
        // --- 1st to (last-1)
        for (int imu = 1; imu < nvr-1; imu++){
            smu_n = dsmu*imu;
            vol_h(mesh_ind,isp,imu) = pi4bb0vth3_dd*smu_n;
        }
        // --- last component
        smu_n = dsmu*((nvr-1)-1.0/3.0);//rh imu or f0_nmu???
        vol_h(mesh_ind,isp,nvr-1) = 0.5*pi4bb0vth3_dd*smu_n;

        //vth_par   =  sqrt(sp_t_par(i)  * UNIT_CHARGE / mass)
        //vth_perp  =  sqrt(sp_t_perp(i) * UNIT_CHARGE / mass)

        // < Mesh-related quantities>
        double vpar_beg_l = -vp_max*vth;// ( = dlx)
        double mesh_dz_l = dvp*vth;// ( = mesh_dz)
        double mesh_dr_l = dsmu*vth;// ( = mesh_dr)

        vpar_beg.h_view(mesh_ind, isp) = vpar_beg_l;
        mesh_dz.h_view(mesh_ind, isp) = mesh_dz_l;
        mesh_dr.h_view(mesh_ind, isp) = mesh_dr_l;

        // mesh_r, mesh_r_half, mesh_z, mesh_z_half
        for (int j = 0; j < nvr-1; j++){
            mesh_r(mesh_ind,isp,j) = mesh_dr_l*j;
            mesh_r_half.h_view(mesh_ind,j,isp) = mesh_dr_l*(j+0.5);
        }
        mesh_r(mesh_ind,isp,nvr-1) = mesh_dr_l*(nvr-1);
        for (int j = 0; j < nvz-1; j++){
            mesh_z(mesh_ind,isp,j) = vpar_beg_l + mesh_dz_l*j;
            mesh_z_half.h_view(mesh_ind,j,isp) = vpar_beg_l + mesh_dz_l*(j+0.5);
        }
        mesh_z(mesh_ind,isp,nvz-1) = vpar_beg_l + mesh_dz_l*(nvz-1);
        // local_center_volume (mesh_r_half, mesh_dr, mesh_dz)
        // : volume centered at a cell
        for (int j = 0; j < (nvr-1); j++){
            local_center_volume.h_view(mesh_ind,j,isp) = mesh_r_half.h_view(mesh_ind,j,isp)*mesh_dr_l*mesh_dz_l;
        }

    }

    void setup_all(const std::vector<Species<Device>>& all_species, const CollisionSpecies<Device>& col_species,
                   const Kokkos::View<int**,HostType>& mesh_nodes, int mesh_batch_ind){

        for (int mesh_ind = 0; mesh_ind<col_species.s.h_view.extent(1); mesh_ind++){
            for (int gri = 0; gri < n; gri++){
                const int FIRST_SPECIES_OF_GRID = 0; // Initialize grid using the first species assigned to the grid
                int isp = map_grid_to_species.h_view(gri,FIRST_SPECIES_OF_GRID);

                // Use equilibrium thermal velocity to construct velocity grid
                int inode = mesh_nodes(mesh_batch_ind,mesh_ind);
                double vth = all_species[col_species.s.h_view(isp,mesh_ind).gid].get_f0_eq_thermal_velocity_lnode_h(inode);
                setup_one(isp, mesh_ind, vth, col_species.f.dsmu, col_species.f.dvp, col_species.f.vp_max);
            }
        }

#ifdef USE_GPU
        // Copy dual views to device
        Kokkos::deep_copy(mesh_dr.d_view, mesh_dr.h_view);
        Kokkos::deep_copy(mesh_dz.d_view, mesh_dz.h_view);
        Kokkos::deep_copy(vpar_beg.d_view, vpar_beg.h_view);
        Kokkos::deep_copy(mesh_r_half.d_view, mesh_r_half.h_view);
        Kokkos::deep_copy(mesh_z_half.d_view, mesh_z_half.h_view);
        Kokkos::deep_copy(local_center_volume.d_view, local_center_volume.h_view);
#endif
        // Transpose vol
        // First create a mirror on device of vol_h
        Kokkos::View<double***,Kokkos::LayoutRight,Device> vol_tmp("vol_tmp",vol_h.extent(0),vol_h.extent(1),vol_h.extent(2));
        // Deep copy that
        Kokkos::deep_copy(vol_tmp, vol_h);

        // Make a local copy of the View itself for the lambda function to copy to device
        Kokkos::View<double***,Kokkos::LayoutRight,Device> lvol = vol;

        // Naive transpose copy on device - unoptimized, should be small though
        int n2 = vol_h.extent(1)*nvr;
        int nvr_loc = nvr;
        Kokkos::parallel_for("transpose_vol", Kokkos::RangePolicy<ExSpace>( 0, n2*vol_h.extent(0)), KOKKOS_LAMBDA (const int idx){
            int i0 = idx/n2; // slowest index
            int rem = idx%n2;
            int i1 = rem/nvr_loc;
            int i2 = rem%nvr_loc; // fastest index
            lvol(i2,i1,i0) = vol_tmp(i0,i1,i2);
        });
    }
};

#endif