vgrid_distribution.hpp

Go to the documentation of this file.

#ifndef VGRID_DISTRIBUTION_HPP
#define VGRID_DISTRIBUTION_HPP

#include "space_settings.hpp"
#include "velocity_grid.hpp"
#include "view_arithmetic.hpp"
#include "velocity_grid.hpp"
#include "domain_decomposition.hpp"
#include "linear_weights.hpp"

// Struct containing indices and weights on velocity grid
struct VGridWeights{
    int i_vr;
    int i_vz;
    double w_00;
    double w_01;
    double w_10;
    double w_11;

    KOKKOS_INLINE_FUNCTION VGridWeights(const LinearWeights& vr_wt, const LinearWeights& vz_wt)
     : i_vr(vr_wt.i),
       i_vz(vz_wt.i),
       w_00(vr_wt.w[0]*vz_wt.w[0]),
       w_01(vr_wt.w[0]*vz_wt.w[1]),
       w_10(vr_wt.w[1]*vz_wt.w[0]),
       w_11(vr_wt.w[1]*vz_wt.w[1])
    {}
};

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(){}

    VGridDistribution(int nsp, const VelocityGrid& vgrid, const DomainDecomposition<DeviceType>& pol_decomp)
      : f("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)
    {}

    // () 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(){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);
    }


    // Determine vgrid weights from input smu, vp and vgrid_vol
    KOKKOS_INLINE_FUNCTION VGridWeights get_weights(double smu, double vp, double inv_vgrid_vol) const{
        // Transformation from mu to vr:
        const double& vr = smu; // square root mu for sqrt(mu) grid, v_perp/v_th else

        // Transformation from vp to vz:
        double vz = vp + vp_max; // Shift so that value is (probably) between 0 and 2*vp_max

        // Calculate linear weights for velocity grid
        LinearWeights vr_wt(vr, 1.0/dsmu);
        LinearWeights vz_wt(vz, 1.0/dvp);

        //rh For f0_f: conversion from mu-vpar to sqrt(mu)-vpar or vperp-vpar
        vr_wt.w[0] *= inv_vgrid_vol * (vr_wt.i==0 ? 2.0 : 1.0);
        vr_wt.w[1] *= inv_vgrid_vol * (vr_wt.i==n_vr()-2 ? 2.0 : 1.0);

        return VGridWeights(vr_wt, vz_wt);
    }

    // Scatter value onto vgrid distribution
    KOKKOS_INLINE_FUNCTION void scatter(int i_node, const VGridWeights& wt, double value) const{
        Kokkos::atomic_add(&(f(wt.i_vr+0,i_node,wt.i_vz+0)),value*wt.w_00);
        Kokkos::atomic_add(&(f(wt.i_vr+0,i_node,wt.i_vz+1)),value*wt.w_01);
        Kokkos::atomic_add(&(f(wt.i_vr+1,i_node,wt.i_vz+0)),value*wt.w_10);
        Kokkos::atomic_add(&(f(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{
        return ( f(wt.i_vr+0,i_node,wt.i_vz+0)*wt.w_00 +
                 f(wt.i_vr+0,i_node,wt.i_vz+1)*wt.w_01 +
                 f(wt.i_vr+1,i_node,wt.i_vz+0)*wt.w_10 +
                 f(wt.i_vr+1,i_node,wt.i_vz+1)*wt.w_11  );
    }

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

#endif