grid_field.hpp

Go to the documentation of this file.

#ifndef GRID_FIELD_HPP
#define GRID_FIELD_HPP
#include "space_settings.hpp"
#include "NamelistReader.hpp"

#include "globals.hpp"
#include "grid.hpp"
#include "linear_weights.hpp"
#include "local_fields.hpp"
#include "push_controls.hpp"
#include "field_weights.hpp"
#include "gyro_radius.hpp"
#include "field.hpp"

/* TorType specifies whether the field is a single plane or needs a plane dimension
 * */
enum class TorType{
    OnePlane,
    MultiplePlanes
};

// General declaration
// 2x2 = 4 types of GridField available: With and without gyroaverage dimension, with and without plane dimension
template<class Device, VarType VT, PhiInterpType PIT, TorType TT, KinType KT, ScatterType ST = ScatterType::Atomic>
struct GridField;

// "ScalarGridField"
template<class Device, VarType VT, PhiInterpType PIT>
struct GridField<Device, VT, PIT, TorType::OnePlane, KinType::DriftKin,ScatterType::Atomic>{
    using field_type = Field<VT,PIT>;
    using device_type = Device;
    static constexpr KinType KT = KinType::DriftKin;

    Kokkos::View<field_type*,Kokkos::LayoutRight,Device> f;

    GridField(){}
    GridField(std::string name, int nphi, int nrho, int nnode) : f(name,nnode) {}
    GridField(std::string name, int nnode) : f(name,nnode) {}
    field_type* data() const {return f.data();}
    int size() const {return f.size();}
    int nnode() const {return f.extent(0);}
    int nrhop1() const {return 1;}
    int nphi() const {return 1;}

    KOKKOS_INLINE_FUNCTION bool is_allocated() const {return f.is_allocated();}
    KOKKOS_INLINE_FUNCTION field_type& operator ()(int inode) const {return f(inode);}

    // Scatter to a node
    KOKKOS_INLINE_FUNCTION void scatter(int node, const FieldWeights<DriftKin, PIT>& wts, double particle_weight) const{
        f(node).scatter(wts, particle_weight);
    }

    // Scatter to a triangle
    KOKKOS_INLINE_FUNCTION void scatter(const Grid<DeviceType>& grid, int ithread, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, int i_simd, const FieldWeights<DriftKin, PIT>& wts, double particle_weight) const{
        for (int j = 0; j<3; j++){
            // Find triangle's nodes
            int node=grid.get_node_index(grid_wts.itr[i_simd], j);
            double wp=grid_wts.p[j][i_simd];

            // Scatter to node
            scatter(node, wts, wp*particle_weight);
        }
    }

    // Gather from a triangle
    KOKKOS_INLINE_FUNCTION void gather(const Grid<Device> &grid, const SimdGridWeights<Order::One, PhiInterpType::None>& grid_wts, const Simd<double>& phi, Simd<double>& fld) const{
        // Initialize local field to 0
        fld = 0.0;

        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            if(grid_wts.is_invalid(i_simd)) continue;

            // Loop over the 3 vertices of the grid triangle and add up the contributions from each
            for (int ip = 0; ip<3; ip++){
                int node = grid.get_node_index(grid_wts.itr[i_simd], ip);
                double wp = grid_wts.p[ip][i_simd];
                f(node).gather(fld,i_simd,wp);
            }
        }
    }

    // Gather from a triangle
    KOKKOS_INLINE_FUNCTION void gather(const Grid<Device> &grid, const SimdGridWeights<Order::One, PhiInterpType::Planes>& grid_wts, const Simd<double>& phi, Simd<double>& fld) const{
        // Initialize local field to 0
        fld = 0.0;

        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            if(grid_wts.is_invalid(i_simd)) continue;
            LinearWeights phi_wts(phi[i_simd],grid.inv_delta_phi,grid.nplanes);

            // Loop over the 3 vertices of the grid triangle and add up the contributions from each
            for (int ip = 0; ip<3; ip++){
                int node = grid.get_node_index(grid_wts.itr[i_simd], ip);
                double wp = grid_wts.p[ip][i_simd];
                f(node).gather(fld,i_simd,wp,phi_wts.w);
            }
        }
    }

    void reset_to_zero(){
        int ndoubles_in_grid_field = f.size()*field_type::SIZE();
        Kokkos::View<double*,Kokkos::LayoutRight,Device,Kokkos::MemoryTraits<Kokkos::Unmanaged>> view_1d_double((double*)(data()), ndoubles_in_grid_field);
        Kokkos::deep_copy(view_1d_double, 0.0);
    }

    // TODO: Generalize this copy and these transposes for ndim etc
    // Fill View with transposed data where inode is the contiguous index
    // Ultimately, there should probably be an implementation GridField that has the transposed indices
    void copy_to_double_view(const View<double*,CLayout,Device>& dest_view) const{
        // Make an unmanaged view around the input field
        View<double*,CLayout,Device,Kokkos::MemoryTraits<Kokkos::Unmanaged>> f_unmanaged((double*)data(),nnode());

        Kokkos::deep_copy(dest_view, f_unmanaged);
    }

    // Fill View with transposed data where inode is the contiguous index
    // Ultimately, there should probably be an implementation GridField that has the transposed indices
    void transpose_copy_to_double_view(const View<double**,CLayout,Device>& dest_view) const{
        // Retrieve nphi from the Field of this GridField
        constexpr int n_phi = field_type::NPHI();

        // Make an unmanaged view around the input field
        View<double**,CLayout,Device,Kokkos::MemoryTraits<Kokkos::Unmanaged>> f_unmanaged((double*)data(),nnode(), n_phi);

        // Copy f_unmanaged into iphi = 0 and iphi = 1 of dest_view
        Kokkos::parallel_for("var_transpose_copy", Kokkos::RangePolicy<ExSpace>( 0, nnode()), KOKKOS_LAMBDA(const int inode){
            for(int iphi=0; iphi<n_phi; iphi++){
                dest_view(iphi, inode) = f_unmanaged(inode, iphi);
            }
        });
        Kokkos::fence();
    }

    // Transpose and copy data from a view where inode is the contiguous index
    void transpose_copy_from_double_view(const View<double**,CLayout,Device>& src_view) const{
        // Retrieve nphi from the Field of this GridField
        constexpr int n_phi = field_type::NPHI();

        // Make an unmanaged view of the data
        View<double**,CLayout,Device,Kokkos::MemoryTraits<Kokkos::Unmanaged>> f_unmanaged((double*)data(),nnode(), n_phi);

        // Transpose from src_view into this gridfield
        Kokkos::parallel_for("var_transpose_copy", Kokkos::RangePolicy<ExSpace>( 0, nnode()), KOKKOS_LAMBDA(const int inode){
            for(int iphi=0; iphi<n_phi; iphi++){
                f_unmanaged(inode, iphi) = src_view(iphi, inode);
            }
        });
        Kokkos::fence();
    }

    static KOKKOS_INLINE_FUNCTION void view_gather3(const Grid<DeviceType>& grid, const View<double*,CLayout,Device>& view1, const View<double*,CLayout,Device>& view2, const View<double*,CLayout,Device>& view3, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, double& val1, double& val2, double& val3, int i_simd){
        val1=0.0;
        val2=0.0;
        val3=0.0;
        for (int ip=0; ip<3; ip++){
            int nd = grid.get_node_index(grid_wts.itr[i_simd], ip);
            val1 += view1(nd)*grid_wts.p[ip][i_simd];
            val2 += view2(nd)*grid_wts.p[ip][i_simd];
            val3 += view3(nd)*grid_wts.p[ip][i_simd];
        }
    }

    static KOKKOS_INLINE_FUNCTION void view_gather(const Grid<DeviceType>& grid, const View<double*,CLayout,Device>& view, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, Simd<double>& val){
        val = 0.0;
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            if(grid_wts.itr[i_simd]<=0) continue;
            for (int ip = 0; ip < 3; ip++){
                int node = grid.get_node_index(grid_wts.itr[i_simd], ip);
                val[i_simd] += grid_wts.p[ip][i_simd]*view(node);
            }   
        }
    }

    static KOKKOS_INLINE_FUNCTION void view_vec_odim_gather(const Grid<DeviceType>& grid, const View<Field<VarType::Vector,PhiInterpType::None>**,CLayout,Device>& view, int ind, const SimdGridWeights<Order::One, PhiInterpType::None>& grid_wts, SimdVector& val, int i_simd){
        val.r[i_simd] = 0.0;
        val.z[i_simd] = 0.0;
        val.phi[i_simd] = 0.0;
        for (int ip = 0; ip<3; ip++){
            int node = grid.get_node_index(grid_wts.itr[i_simd], ip);
            double wp = grid_wts.p[ip][i_simd];
            view(ind,node).gather(val, i_simd, wp);
        }
    }

    static KOKKOS_INLINE_FUNCTION void view_gyro_scatter3(const Grid<DeviceType>& grid, const View<double**,CLayout,Device>& view1, const View<double**,CLayout,Device>& view2, const View<double**,CLayout,Device>& view3, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, double val1, double val2, double val3, int i_simd, const LinearWeights& rho_wts, bool is_gyro){
        // Scatter moments into grid array
        for (int ip = 0; ip < 3; ip++){
            int node = grid.get_node_index(grid_wts.itr[i_simd], ip);
            Kokkos::atomic_add(&(view1(rho_wts.i,   node)), grid_wts.p[ip][i_simd]*rho_wts.w[0]*val1);
            Kokkos::atomic_add(&(view2(rho_wts.i,   node)), grid_wts.p[ip][i_simd]*rho_wts.w[0]*val2);
            Kokkos::atomic_add(&(view3(rho_wts.i,   node)), grid_wts.p[ip][i_simd]*rho_wts.w[0]*val3);
            if(is_gyro){
                Kokkos::atomic_add(&(view1(rho_wts.i+1, node)), grid_wts.p[ip][i_simd]*rho_wts.w[1]*val1);
                Kokkos::atomic_add(&(view2(rho_wts.i+1, node)), grid_wts.p[ip][i_simd]*rho_wts.w[1]*val2);
                Kokkos::atomic_add(&(view3(rho_wts.i+1, node)), grid_wts.p[ip][i_simd]*rho_wts.w[1]*val3);
            }   
        }
    }

    // view with outer dimension
    static KOKKOS_INLINE_FUNCTION void view_odim_scatter(const Grid<DeviceType>& grid, const View<double**,CLayout,Device>& view, int ind, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, double val, int i_simd){
        for (int j = 0; j<3; j++){
            int node=grid.get_node_index(grid_wts.itr[i_simd], j);
            double wp=grid_wts.p[j][i_simd];

            Kokkos::atomic_add(&(view(ind,node)), wp*val);
        }
    }

    static KOKKOS_INLINE_FUNCTION void view_scatter(const Grid<DeviceType>& grid, const View<double*,CLayout,Device>& view, const SimdGridWeights<Order::One, PhiInterpType::None>& grid_wts, double val, int i_simd){
        for (int j = 0; j<3; j++){
            int node=grid.get_node_index(grid_wts.itr[i_simd], j);
            double wp=grid_wts.p[j][i_simd];

            Kokkos::atomic_add(&(view(node)), wp*val);
        }
    }

    static KOKKOS_INLINE_FUNCTION void view_ff_scatter(const Grid<DeviceType>& grid, const View<double**,CLayout,Device>& view, const SimdGridWeights<Order::One, PhiInterpType::Planes>& grid_wts, double wphi, double val, int i_simd){
        for (int j = 0; j<3; j++){
            int node=grid.get_node_index(grid_wts.itr[i_simd], j);
            double wp=grid_wts.p[j][i_simd];

            Kokkos::atomic_add(&(view(0,node)), wp*wphi*val);
            Kokkos::atomic_add(&(view(1,node)), wp*(1.0-wphi)*val);
        }
    }

    static KOKKOS_INLINE_FUNCTION void view_odim_variance(const Grid<DeviceType>& grid, const View<double**,CLayout,Device>& view, const View<double**,CLayout,Device>& variance, int ind, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, double val, int i_simd){
        for (int j = 0; j<3; j++){
            int node=grid.get_node_index(grid_wts.itr[i_simd], j);
            double wp=grid_wts.p[j][i_simd];

            double tmp = view(ind,node) - wp*val;
            Kokkos::atomic_add(&(variance(ind,node)), tmp*tmp);
        }
    }

    static KOKKOS_INLINE_FUNCTION void view_rep_scatter3(const Grid<DeviceType>& grid, const View<double***,CLayout,Device>& view, int ind1, int ind2, int ind3, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, double val1, double val2, double val3, int ithread, int i_simd){
        // Scatter moments into grid array
        for (int ip = 0; ip < 3; ip++){
            int node = grid.get_node_index(grid_wts.itr[i_simd], ip);
            double wp = grid_wts.p[ip][i_simd];

            access_add(&(view(ithread,node,ind1)),wp*val1);
            access_add(&(view(ithread,node,ind2)),wp*val2);
            access_add(&(view(ithread,node,ind3)),wp*val3);
        }
    }
};

template<class Device, VarType VT, PhiInterpType PIT>
struct GridField<Device, VT, PIT, TorType::OnePlane, KinType::GyroKin,ScatterType::Atomic>{
    using field_type = Field<VT,PIT>;
    using device_type = Device;
    static constexpr KinType KT = KinType::GyroKin;

    Kokkos::View<field_type**,Kokkos::LayoutRight,Device> f;

    GridField(){}
    GridField(std::string name, int nphi, int nrho, int nnode) : f(name,nnode, nrho+1) {}
    GridField(std::string name, int nrho, int nnode) : f(name,nnode, nrho+1) {}
    field_type* data() const {return f.data();}
    int size() const {return f.size();}
    int nnode() const {return f.extent(0);}
    int nrhop1() const {return f.extent(1);}
    int nphi() const {return 1;}
    KOKKOS_INLINE_FUNCTION bool is_allocated() const {return f.is_allocated();}
    KOKKOS_INLINE_FUNCTION field_type& operator ()(int inode, int irho) const {return f(inode, irho);}

    // Scatter to a node
    KOKKOS_INLINE_FUNCTION void scatter(int node, const FieldWeights<GyroKin, PIT>& wts, double particle_weight) const{
#ifdef NEWGYROMATRIX
        // If NEWGYROMATRIX, there are 2 rho wts so must specify (even though the i indices are identical)
        int irho = wts.rho[0].i;
#else
        int irho = wts.rho.i;
#endif
        f(node,irho).scatter(wts, particle_weight, f(node,irho+1));
    }

    // Scatter to a triangle (same as for DriftKin - consolidate)
    KOKKOS_INLINE_FUNCTION void scatter(const Grid<DeviceType>& grid, int ithread, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, int i_simd, const FieldWeights<GyroKin, PIT>& wts, double particle_weight) const{
        for (int j = 0; j<3; j++){
            // Find triangle's nodes
            int node=grid.get_node_index(grid_wts.itr[i_simd], j);
            double wp=grid_wts.p[j][i_simd];
            
            // Scatter to node
            scatter(node, wts, wp*particle_weight);
        }   
    }

    void reset_to_zero(){
        int ndoubles_in_grid_field = f.size()*field_type::SIZE();
        Kokkos::View<double*,Kokkos::LayoutRight,Device,Kokkos::MemoryTraits<Kokkos::Unmanaged>> view_1d_double((double*)(f.data()), ndoubles_in_grid_field);
        Kokkos::deep_copy(view_1d_double, 0.0);
    }
};

/******** GridField with array replication ********/
template<class Device, VarType VT, PhiInterpType PIT>
struct GridField<Device, VT, PIT, TorType::OnePlane, KinType::DriftKin,ScatterType::Replication>{
    using field_type = Field<VT,PIT>;
    using device_type = Device;
    static constexpr KinType KT = KinType::DriftKin;

    Kokkos::View<field_type**,Kokkos::LayoutRight,Device> f;

    GridField(){}
    GridField(std::string name, int nphi, int nrho, int nnode) : f(name,get_num_cpu_threads(), nnode) {}
    GridField(std::string name, int nnode) : f(name,get_num_cpu_threads(), nnode) {}
    field_type* data() const {return f.data();}
    int size() const {return f.size();}
    int nnode() const {return f.extent(1);}
    int nrhop1() const {return 1;}
    int nphi() const {return 1;}
    KOKKOS_INLINE_FUNCTION bool is_allocated() const {return f.is_allocated();}
    KOKKOS_INLINE_FUNCTION field_type& operator ()(int ithread, int inode) const {return f(ithread, inode);}

    // Scatter to a node
    KOKKOS_INLINE_FUNCTION void scatter(int ithread, int node, const FieldWeights<DriftKin, PIT>& wts, double particle_weight) const{
        f(ithread,node).scatter(wts, particle_weight);
    }

    // Scatter to a triangle
    KOKKOS_INLINE_FUNCTION void scatter(const Grid<DeviceType>& grid, int ithread, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, int i_simd, const FieldWeights<DriftKin, PIT>& wts, double particle_weight) const{
        for (int j = 0; j<3; j++){
            // Find triangle's nodes
            int node=grid.get_node_index(grid_wts.itr[i_simd], j);
            double wp=grid_wts.p[j][i_simd];

            // Scatter to node
            scatter(ithread, node, wts, wp*particle_weight);
        }
    }

    void reset_to_zero(){
        int ndoubles_in_grid_field = f.size()*field_type::SIZE();
        Kokkos::View<double*,Kokkos::LayoutRight,Device,Kokkos::MemoryTraits<Kokkos::Unmanaged>> view_1d_double((double*)(f.data()), ndoubles_in_grid_field);
        Kokkos::deep_copy(view_1d_double, 0.0);
    }
};

template<class Device, VarType VT, PhiInterpType PIT>
struct GridField<Device, VT, PIT, TorType::OnePlane, KinType::GyroKin,ScatterType::Replication>{
    using field_type = Field<VT,PIT>;
    using device_type = Device;
    static constexpr KinType KT = KinType::GyroKin;

    Kokkos::View<field_type***,Kokkos::LayoutRight,Device> f;

    GridField(){}
    GridField(std::string name, int nphi, int nrho, int nnode) : f(name,get_num_cpu_threads(),nnode, nrho+1) {}
    GridField(std::string name, int nrho, int nnode) : f(name,get_num_cpu_threads(),nnode, nrho+1) {}
    field_type* data() const {return f.data();}
    int size() const {return f.size();}
    int nnode() const {return f.extent(1);}
    int nrhop1() const {return f.extent(2);}
    int nphi() const {return 1;}
    KOKKOS_INLINE_FUNCTION bool is_allocated() const {return f.is_allocated();}
    KOKKOS_INLINE_FUNCTION field_type& operator ()(int ithread, int inode, int irho) const {return f(ithread,inode, irho);}

    // Scatter to a node
    KOKKOS_INLINE_FUNCTION void scatter(int ithread, int node, const FieldWeights<GyroKin, PIT>& wts, double particle_weight) const{
#ifdef NEWGYROMATRIX
        // If NEWGYROMATRIX, there are 2 rho wts so must specify (even though the i indices are identical)
        int irho = wts.rho[0].i;
#else
        int irho = wts.rho.i;
#endif
        f(ithread,node,irho).scatter(wts, particle_weight, f(ithread,node,irho+1));
    }

    // Scatter to a triangle (same as for DriftKin - consolidate)
    KOKKOS_INLINE_FUNCTION void scatter(const Grid<DeviceType>& grid, int ithread, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, int i_simd, const FieldWeights<GyroKin, PIT>& wts, double particle_weight) const{
        for (int j = 0; j<3; j++){
            // Find triangle's nodes
            int node=grid.get_node_index(grid_wts.itr[i_simd], j);
            double wp=grid_wts.p[j][i_simd];

            // Scatter to node
            scatter(ithread, node, wts, wp*particle_weight);
        }
    }

    void reset_to_zero(){
        int ndoubles_in_grid_field = f.size()*field_type::SIZE();
        Kokkos::View<double*,Kokkos::LayoutRight,Device,Kokkos::MemoryTraits<Kokkos::Unmanaged>> view_1d_double((double*)(f.data()), ndoubles_in_grid_field);
        Kokkos::deep_copy(view_1d_double, 0.0);
    }
};
/************/

template<class Device, VarType VT, PhiInterpType PIT>
struct GridField<Device, VT, PIT, TorType::MultiplePlanes, KinType::DriftKin, ScatterType::Atomic>{
    using field_type = Field<VT,PIT>;
    using device_type = Device;
    static constexpr KinType KT = KinType::DriftKin;

    Kokkos::View<field_type**,Kokkos::LayoutRight,Device> f;

    GridField(){}
    GridField(std::string name, int nphi, int nrho, int nnode) : f(name,nphi, nnode) {}
    GridField(std::string name, int nphi, int nnode) : f(name,nphi, nnode) {}
    field_type* data() const {return f.data();}
    int size() const {return f.size();}
    int nnode() const {return f.extent(1);}
    int nrhop1() const {return 1;}
    int nphi() const {return f.extent(0);}
    KOKKOS_INLINE_FUNCTION bool is_allocated() const {return f.is_allocated();}
    KOKKOS_INLINE_FUNCTION field_type& operator ()(int iphi, int inode) const {return f(iphi, inode);}

    GridField<Device, VT, PIT, TorType::OnePlane, KinType::DriftKin> subfield(int subfield_idx) const{
        GridField<Device, VT, PIT, TorType::OnePlane, KinType::DriftKin> new_field("subfield", nnode());
        new_field.f = my_subview(f, subfield_idx);
        return new_field;
    }

    Kokkos::View<double**,Kokkos::LayoutRight,Device,Kokkos::MemoryTraits<Kokkos::Unmanaged>> unmanaged() const{
        return Kokkos::View<double**,Kokkos::LayoutRight,Device,Kokkos::MemoryTraits<Kokkos::Unmanaged>>((double*)(f.data()), f.layout());
    }
};

template<class Device, VarType VT, PhiInterpType PIT>
struct GridField<Device, VT, PIT, TorType::MultiplePlanes, KinType::GyroKin, ScatterType::Atomic>{
    using field_type = Field<VT,PIT>;
    using device_type = Device;
    static constexpr KinType KT = KinType::GyroKin;

    Kokkos::View<field_type***,Kokkos::LayoutRight,Device> f;

    GridField(){}
    GridField(std::string name, int nphi, int nrho, int nnode) : f(name,nphi, nnode, nrho+1) {}
    field_type* data() const {return f.data();}
    int size() const {return f.size();}
    int nnode() const {return f.extent(1);}
    int nrhop1() const {return f.extent(2);}
    int nphi() const {return f.extent(0);}
    KOKKOS_INLINE_FUNCTION bool is_allocated() const {return f.is_allocated();}
    KOKKOS_INLINE_FUNCTION field_type& operator ()(int iphi, int inode, int irho) const {return f(iphi, inode, irho);}
};

// For convenience
typedef GridField<DeviceType, VarType::Scalar,PhiInterpType::None, TorType::OnePlane, KinType::DriftKin> ScalarGridField;


// Shallow copy if types are the same
template<typename T>
inline void grid_field_copy(T& dest, const T& src){
    dest = src;
}

// Allocate dest grid_field and copy data from src
template<typename T1, typename T2>
inline void grid_field_copy(T1& dest, const T2& src){
    // Allocate device array
    dest = T1("grid_field",src.nphi(), src.nrhop1()-1, src.nnode());
    // Deep copy from host to device
    Kokkos::deep_copy(dest.f, src.f);
}

#endif