vgrid_distribution.hpp

Go to the documentation of this file.

#ifndef VGRID_DISTRIBUTION_HPP
#define VGRID_DISTRIBUTION_HPP

#include "space_settings.hpp"
#include "NamelistReader.hpp"
#include "velocity_grid.hpp"
#include "view_arithmetic.hpp"
#include "domain_decomposition.hpp"
#include "maxwellian.hpp"
#include "uniform_range.hpp"

enum class VGridDistributionOption{
    NoViewInit=0,
    ViewInit
};

template<class Device>
class VGridDistribution{
    public:

    using exec_space = typename Device::execution_space;

    View<double****,CLayout, Device> f;

    // Create a 1D unmanaged view looking at the same address
    View<double*,CLayout, Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f_1D;

    // Make f private when possible - ALS
    //public:

    double vp_max;  
    double dvp;     
    double smu_max; 
    double dsmu;    
    double inv_mu0_factor;

    VGridDistribution(){}

    // Analytical distribution for testing
    VGridDistribution(const VelocityGrid& vgrid, const DomainDecomposition<DeviceType>& pol_decomp, const std::vector<Maxwellian>& maxwellians)
      : f("f", maxwellians.size(), vgrid.nvr, pol_decomp.nnodes, vgrid.nvz),
        f_1D(f.data(), f.size()),
        vp_max(vgrid.vp_max),
        dvp(vgrid.dvp),
        smu_max(vgrid.smu_max),
        dsmu(vgrid.dsmu),
        inv_mu0_factor(vgrid.inv_mu0_factor)
    {
        for (int i=0; i<f.extent(0); i++){
            for (int imu=0; imu<f.extent(1); imu++){
                double vperp = imu*dsmu;
                for (int inode=0; inode<f.extent(2); inode++){
                    for (int ivp=0; ivp<f.extent(3); ivp++){
                        double vpara = ivp*dvp - vp_max;
                        f(i,imu,inode,ivp) = maxwellians[i].get_f(vpara, vperp);
                    }   
                }   
            }   
        }
    }

    VGridDistribution(int nsp, const VelocityGrid& vgrid, const DomainDecomposition<DeviceType>& pol_decomp, VGridDistributionOption option = VGridDistributionOption::ViewInit)
      : f(NoInit("f"), nsp, vgrid.nvr, pol_decomp.nnodes, vgrid.nvz),
        f_1D(f.data(), f.size()),
        vp_max(vgrid.vp_max),
        dvp(vgrid.dvp),
        smu_max(vgrid.smu_max),
        dsmu(vgrid.dsmu),
        inv_mu0_factor(vgrid.inv_mu0_factor)
    {
        if(option==VGridDistributionOption::ViewInit) Kokkos::deep_copy(f, 0.0);
    }

    template<class Device2>
    VGridDistribution(int nsp, const VGridDistribution<Device2>& dist_in, VGridDistributionOption option = VGridDistributionOption::ViewInit)
      : f(NoInit("f"), nsp, dist_in.n_vr(), dist_in.n_nodes(), dist_in.n_vz()),
        f_1D(f.data(), f.size()),
        vp_max(dist_in.vp_max),
        dvp(dist_in.dvp),
        smu_max(dist_in.smu_max),
        dsmu(dist_in.dsmu),
        inv_mu0_factor(dist_in.inv_mu0_factor)
    {
        if(option==VGridDistributionOption::ViewInit) Kokkos::deep_copy(f, 0.0);
    }

    // () operator: behave as if accessing the View directly
    KOKKOS_INLINE_FUNCTION double& operator()(int isp, int ivr, int inode, int ivz) const { return f(isp, ivr, inode, ivz);}

    // [] operator: access the data as if it has been squashed to 1 dimension
    KOKKOS_INLINE_FUNCTION double& operator[](int i) const { return f_1D(i);}
    
    double* data() const {return f.data();}

    template<typename F>
    inline void for_all_elements(const std::string label, F lambda_func) const {
        Kokkos::parallel_for(label, Kokkos::RangePolicy<exec_space>(0, f.size()), lambda_func);
    }

    template<typename F>
    inline void for_each_element(const std::string label, F lambda_func) const {
        Kokkos::parallel_for(label, Kokkos::MDRangePolicy<Kokkos::Rank<4>, exec_space>({0, 0, 0, 0}, {f.extent(0), f.extent(1), f.extent(2), f.extent(3)}), lambda_func);
    }


    // Scatter value onto vgrid distribution
    KOKKOS_INLINE_FUNCTION void scatter(int i_node, const VGridWeights& wt, double value) const{
        constexpr int ISP_ZERO = 0;
        Kokkos::atomic_add(&(f(ISP_ZERO, wt.i_vr+0,i_node,wt.i_vz+0)),value*wt.w_00);
        Kokkos::atomic_add(&(f(ISP_ZERO, wt.i_vr+0,i_node,wt.i_vz+1)),value*wt.w_01);
        Kokkos::atomic_add(&(f(ISP_ZERO, wt.i_vr+1,i_node,wt.i_vz+0)),value*wt.w_10);
        Kokkos::atomic_add(&(f(ISP_ZERO, wt.i_vr+1,i_node,wt.i_vz+1)),value*wt.w_11);
    }

    // Temporary scatter for views not yet converted to VGridDistribution
    KOKKOS_INLINE_FUNCTION static void scatter(const View<double***,CLayout,Device>& view, int i_node, const VGridWeights& wt, double value){
        Kokkos::atomic_add(&(view(wt.i_vr+0,i_node,wt.i_vz+0)),value*wt.w_00);
        Kokkos::atomic_add(&(view(wt.i_vr+0,i_node,wt.i_vz+1)),value*wt.w_01);
        Kokkos::atomic_add(&(view(wt.i_vr+1,i_node,wt.i_vz+0)),value*wt.w_10);
        Kokkos::atomic_add(&(view(wt.i_vr+1,i_node,wt.i_vz+1)),value*wt.w_11);
    }

    // Temporary scatter for views not yet converted to VGridDistribution
    KOKKOS_INLINE_FUNCTION static void scatter(const View<double****,CLayout,Device>& view, int i_node, const VGridWeights& wt, double value){
        constexpr int ISP_ZERO = 0;
        Kokkos::atomic_add(&(view(ISP_ZERO,wt.i_vr+0,i_node,wt.i_vz+0)),value*wt.w_00);
        Kokkos::atomic_add(&(view(ISP_ZERO,wt.i_vr+0,i_node,wt.i_vz+1)),value*wt.w_01);
        Kokkos::atomic_add(&(view(ISP_ZERO,wt.i_vr+1,i_node,wt.i_vz+0)),value*wt.w_10);
        Kokkos::atomic_add(&(view(ISP_ZERO,wt.i_vr+1,i_node,wt.i_vz+1)),value*wt.w_11);
    }

    // Gather value from vgrid distribution
    KOKKOS_INLINE_FUNCTION double gather(int i_node, const VGridWeights& wt) const{
        constexpr int ISP_ZERO = 0;
        return ( f(ISP_ZERO,wt.i_vr+0,i_node,wt.i_vz+0)*wt.w_00 +
                 f(ISP_ZERO,wt.i_vr+0,i_node,wt.i_vz+1)*wt.w_01 +
                 f(ISP_ZERO,wt.i_vr+1,i_node,wt.i_vz+0)*wt.w_10 +
                 f(ISP_ZERO,wt.i_vr+1,i_node,wt.i_vz+1)*wt.w_11  );
    }

    // Temporary gather for views not yet converted to VGridDistribution
    template<class T>
    KOKKOS_INLINE_FUNCTION static double gather(const T& view, int i_node, const VGridWeights& wt){
        return ( view(wt.i_vr+0,i_node,wt.i_vz+0)*wt.w_00 +
                 view(wt.i_vr+0,i_node,wt.i_vz+1)*wt.w_01 +
                 view(wt.i_vr+1,i_node,wt.i_vz+0)*wt.w_10 +
                 view(wt.i_vr+1,i_node,wt.i_vz+1)*wt.w_11  );
    }

    // Temporary gather for views not yet converted to VGridDistribution
    // Divides by normalization view
    template<class T>
    KOKKOS_INLINE_FUNCTION static double normed_gather(const T& view, int i_node, const VGridWeights& wt, const T& norm_view){
        return (view(wt.i_vr+0,i_node,wt.i_vz+0)*wt.w_00/norm_view(wt.i_vr+0,i_node,wt.i_vz+0) +
                view(wt.i_vr+0,i_node,wt.i_vz+1)*wt.w_01/norm_view(wt.i_vr+0,i_node,wt.i_vz+1) +
                view(wt.i_vr+1,i_node,wt.i_vz+0)*wt.w_10/norm_view(wt.i_vr+1,i_node,wt.i_vz+0) +
                view(wt.i_vr+1,i_node,wt.i_vz+1)*wt.w_11/norm_view(wt.i_vr+1,i_node,wt.i_vz+1));
    }

    // Access dimensions
    KOKKOS_INLINE_FUNCTION int n_species() const{
        return f.extent(0);
    }

    KOKKOS_INLINE_FUNCTION int n_vr() const {
        return f.extent(1);
    }

    KOKKOS_INLINE_FUNCTION int n_nodes() const{
        return f.extent(2);
    }

    KOKKOS_INLINE_FUNCTION int n_vz() const{
        return f.extent(3);
    }

    KOKKOS_INLINE_FUNCTION int size() const{
        return f.size();
    }

    // Miscellaneous
    inline double get_smu_n(int imu) const{
        return (imu>0 ? dsmu*imu : dsmu*inv_mu0_factor);
    }

    // Returns factor to handle velocity space volume on edge of the v_perp velocity grid
    KOKKOS_INLINE_FUNCTION double mu_vol_fac(int ivr) const{
        return (ivr==0 || ivr==n_vr()-1) ? 0.5 : 1.0;
    }

    // Returns factor to handle velocity space volume on the v_parallel edge of velocity grid
    KOKKOS_INLINE_FUNCTION double vp_vol_fac(int ivz) const{
        return (ivz==0 || ivz==n_vz()-1) ? 0.5 : 1.0;
    }

    UniformRange vr_range() const{
        return UniformRange(n_vr(), 0.0, smu_max);
    }

    UniformRange vz_range() const{
        return UniformRange(n_vz(), -vp_max, vp_max);
    }

    // inode is the vertex to access
    // ip is the index of all other dimensions collapsed
    KOKKOS_INLINE_FUNCTION double& pull_node_index(int inode, int ip) const{
        int ivp = ip%n_vz(); // Fastest index
        int itmp = ip/n_vz();
        int imu = itmp%n_vr();
        int isp = itmp/n_vr(); // Slowest index
        return f(isp,imu,inode,ivp);
    }

    // Resize to new number of vertices
    void resize_n_vertices(int new_n_nodes){
        Kokkos::resize(f, n_species(), n_vr(), new_n_nodes, n_vz());
        f_1D = View<double*,CLayout, Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>>(f.data(), f.size());
    }
};

#endif