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;

template<class Device, VarType VT, PhiInterpType PIT>
struct GridField<Device, VT, PIT, TorType::OnePlane, KinType::DriftKin,ScatterType::Atomic>{
    using field_type = Field<VT,PIT>;
    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 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 Simd<int>& itr, const SimdGridVec& p, 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(itr[i_simd], j);
            double wp=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*)(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();
    }
};

template<class Device, VarType VT, PhiInterpType PIT>
struct GridField<Device, VT, PIT, TorType::OnePlane, KinType::GyroKin,ScatterType::Atomic>{
    using field_type = Field<VT,PIT>;
    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 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 Simd<int>& itr, const SimdGridVec& p, 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(itr[i_simd], j);
            double wp=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>;
    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 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 Simd<int>& itr, const SimdGridVec& p, 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(itr[i_simd], j);
            double wp=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>;
    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 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 Simd<int>& itr, const SimdGridVec& p, 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(itr[i_simd], j);
            double wp=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>;
    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 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(){
        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>;
    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 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;

#endif