grid_field_pack.hpp

Go to the documentation of this file.

#ifndef GRID_FIELD_PACK_HPP
#define GRID_FIELD_PACK_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 "gyro_radius.hpp"
#include "grid_field.hpp"
#include "label.hpp"
#include "pack.hpp"

// Base struct for GridFieldPack so that non-GPU code doesn't need to be templated just because there is a GridFieldPack argument
struct GridFieldPackGeneric{
};

template<class Device, PhiInterpType PIT>
struct GridFieldPack : GridFieldPackGeneric{
    // Temporary implementation of this to make the PRs more legible
#ifdef EXPLICIT_EM
    using pack_type = Pack<Labeled<EfieldType, Label::E>,
                           Labeled<EfieldType, Label::dAh>,
                           Labeled<EfieldType, Label::dAs>,
                           Labeled<PotType, Label::Ah>,
                           Labeled<PotType, Label::As>,
                           Labeled<PotType, Label::Epar_em>,

                           Labeled<EfieldGyroType, Label::E_gyro>,
                           Labeled<EfieldGyroType, Label::dAh_gyro>,
                           Labeled<EfieldGyroType, Label::dAs_gyro>,
                           Labeled<PotGyroType, Label::Ah_gyro>,
                           Labeled<PotGyroType, Label::As_gyro>,
                           Labeled<PotGyroType, Label::Epar_em_gyro>>;
#else
# ifdef DELTAF_CONV
    using pack_type = Pack<Labeled<EfieldType, Label::E>,
                           Labeled<EfieldGyroType, Label::E_gyro>,
                           Labeled<E00Type, Label::E00>,
                           Labeled<PotType, Label::ddpotdt>>;
# else
    using pack_type = Pack<Labeled<EfieldType, Label::E>,
                           Labeled<EfieldGyroType, Label::E_gyro>>;
# endif
#endif

    private:
    KinType kintype;
    bool is_initialized;

    public:

    bool turb_efield; // Ideally this would not be here but it is tedious to pass into the gather otherwise

    int phi_offset; 
    int node_offset; 

#ifdef XGC1
    // Views stored in Pack
    pack_type pack;

    // Const access to pack contents for inside kernels (can modify data but not structure)
    template<Label FN>
    KOKKOS_INLINE_FUNCTION const auto& get() const{
        return pack.template get<FN>();
    }

    // Non-const access to pack contents (can modify structure)
    template<Label FN>
    inline auto& get(){
        return pack.template get<FN>();
    }
#endif


#ifdef XGC1
    GridField<Device, VarType::Scalar,PIT, TorType::OnePlane, KinType::DriftKin> dpot_ff;
#else
    GridField<Device, vec2d_if_axisym<PIT>(),PIT, TorType::OnePlane, KinType::GyroKin> E_rho;
    GridField<Device, VarType::Scalar,PIT, TorType::OnePlane, KinType::DriftKin> dpot;
#  ifdef SONIC_GK
    GridField<Device, VarType::Vector2D,PIT, TorType::OnePlane, KinType::GyroKin> dEr_B2_rho;
    GridField<Device, VarType::Vector2D,PIT, TorType::OnePlane, KinType::GyroKin> dEz_B2_rho;
    GridField<Device, VarType::Vector2D,PIT, TorType::OnePlane, KinType::GyroKin> du2_E_rho;
#  endif
#endif
    // This will be a ScalarGridField in case of XGCa and a grid field with field-following
    // interpolation in case of XGC1
    GridField<Device, VarType::Scalar,PIT, TorType::OnePlane, KinType::DriftKin> loop_voltage;

    GridFieldPack()
        : is_initialized(false) {}

    GridFieldPack(KinType kintype_in, bool turb_efield_in)
        : is_initialized(true),
          kintype(kintype_in),
          turb_efield(turb_efield_in),
          phi_offset(0),
          node_offset(0) {}

    bool can_reuse(KinType kintype_in){
        if(is_initialized){
#ifdef XGC1
            // For XGC1, can reuse if the existing GridFieldPack is the same KinType
            return kintype_in==kintype;
#else
            // For XGCa, the drift kinetic GridFieldPack is a subset of the gyrokinetic one,
            // so is always reused
            return true;
#endif
        }else{
            // Cant reuse if uninitialized
            return false;
        }
    }

    // Device functions:

    KOKKOS_INLINE_FUNCTION void phi_from_para(SimdVector& vec, int i_simd, const SimdVector& B, double Bmag) const{
        vec.phi[i_simd]=(vec.phi[i_simd]*Bmag - vec.r[i_simd]*B.r[i_simd] - vec.z[i_simd]*B.z[i_simd]) / B.phi[i_simd];
    }



    KOKKOS_INLINE_FUNCTION void gather_all_fields(const PushControls &push_controls, int i_simd, int i_node, double wp, const FieldCorrection& corr, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, const SimdGyroWeights<DriftKin>& rho_wts, LocalFields& fld) const {
        const int node = i_node - node_offset;
#ifdef XGCA
        E_rho(node,0).gather(fld.E,i_simd, wp, corr); // irho = 0
        // Loop voltage
        //loop_voltage(node).gather(fld.loop_voltage,i_simd,wp);
#else
        const int i_phi = (grid_wts.phi_wts.phi.i[i_simd] - phi_offset)%get<Label::E>().f.extent(0); //should be grid.nplanes (?);
        get<Label::E>()(i_phi,node).gather(fld.E,i_simd, wp, grid_wts.phi_wts, corr);
#ifdef EXPLICIT_EM
        get<Label::Ah>()(i_phi,node).gather(fld.Ah,i_simd, wp, grid_wts.phi_wts);
        get<Label::dAh>()(i_phi,node).gather(fld.dAh,i_simd, wp, grid_wts.phi_wts, corr);

        if(push_controls.em_mixed_variable) {
            get<Label::dAs>()(i_phi,node).gather(fld.dAs,i_simd, wp, grid_wts.phi_wts, corr);
            get<Label::As>()(i_phi,node).gather(fld.As,i_simd, wp, grid_wts.phi_wts);
        }
        if(push_controls.em_mixed_variable && (push_controls.em_pullback_mode==4)) {
            get<Label::Epar_em>()(i_phi,node).gather(fld.Epar_em,i_simd,wp,grid_wts.phi_wts);
        }
#endif
        // Loop voltage
        //loop_voltage(node).gather(fld.loop_voltage,i_simd,wp,grid_wts.phi_wts);
#endif

#if defined(DELTAF_CONV) && !defined(EXPLICIT_EM)
        get<Label::E00>()(node).gather(fld.E00,i_simd, wp, grid_wts.phi_wts, corr);
        get<Label::ddpotdt>()(i_phi,node).gather(fld.ddpotdt,i_simd, wp, grid_wts.phi_wts);
#endif
    }

    // Gather all fields (ion)
    KOKKOS_INLINE_FUNCTION void gather_all_fields(const PushControls &push_controls, int i_simd, int node, double wp, const FieldCorrection& corr, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, const SimdGyroWeights<GyroKin>& rho_wts, LocalFields& fld) const {
        int irho = rho_wts.irho(i_simd);
#ifdef XGCA
        E_rho(node,irho).gather(fld.E,i_simd, wp, rho_wts, corr, E_rho(node,irho+1));
#  ifdef SONIC_GK
        dEr_B2_rho(node,0).gather(fld.dEr_B2,i_simd, wp, corr); // irho = 0
        dEz_B2_rho(node,0).gather(fld.dEz_B2,i_simd, wp, corr); // irho = 0
        du2_E_rho(node,0).gather(fld.du2_E,i_simd, wp, corr); // irho = 0
#  endif
        // Loop voltage
        //loop_voltage(node).gather(fld.loop_voltage,i_simd,wp);
#else
        get<Label::E_gyro>()(node,irho).gather(fld.E,i_simd, wp, rho_wts, grid_wts.phi_wts, corr, get<Label::E_gyro>()(node,irho+1));
#ifdef EXPLICIT_EM
        get<Label::Ah_gyro>()(node,irho).gather(fld.Ah,i_simd, wp, rho_wts, grid_wts.phi_wts, get<Label::Ah_gyro>()(node,irho+1));
        get<Label::dAh_gyro>()(node,irho).gather(fld.dAh,i_simd, wp, rho_wts, grid_wts.phi_wts, corr, get<Label::dAh_gyro>()(node,irho+1));

        if(push_controls.em_mixed_variable) {
            get<Label::dAs_gyro>()(node,irho).gather(fld.dAs,i_simd, wp, rho_wts, grid_wts.phi_wts, corr, get<Label::dAs_gyro>()(node,irho+1));
            get<Label::As_gyro>()(node,irho).gather(fld.As,i_simd, wp, rho_wts, grid_wts.phi_wts, get<Label::As_gyro>()(node,irho+1));
        }
        if(push_controls.em_mixed_variable && (push_controls.em_pullback_mode==4)) {
            get<Label::Epar_em_gyro>()(node,irho).gather(fld.Epar_em,i_simd,wp,rho_wts, grid_wts.phi_wts, get<Label::Epar_em_gyro>()(node,irho+1));
        }
#endif
        // Loop voltage
        //loop_voltage(node).gather(fld.loop_voltage,i_simd,wp,grid_wts.phi_wts);
#endif
    }


    template<KinType PT>
    KOKKOS_INLINE_FUNCTION void fields_at_point(const PushControls &push_controls, const Grid<Device> &grid, const SimdVector &B, const SimdVector2D &gradpsi, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, const SimdGyroWeights<PT>& rho_wts, LocalFields& fld) const
    {
        // Initialize local fields to 0
        fld.E.zero();
#if defined(DELTAF_CONV) && !defined(EXPLICIT_EM)
        fld.E00.zero(); 
        fld.ddpotdt.zero();
#endif
#ifdef EXPLICIT_EM
        fld.dAh.zero(); 
        fld.dAs.zero();
        fld.Ah.zero();
        fld.As.zero();
        fld.Epar_em.zero();
#endif
#ifdef SONIC_GK
        fld.dEr_B2.zero();
        fld.dEz_B2.zero();
        fld.du2_E.zero();
#endif

        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            // If particle is inactive, set itr_work=1, phi=0D0, iphi=0
            int itr_work = grid_wts.is_valid(i_simd) ? grid_wts.itr[i_simd] : 1;

            // Field basis correction
            FieldCorrection field_correction(gradpsi, i_simd);

            // 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_wts.is_valid(i_simd) ? grid.get_node_index(itr_work, ip) : 0;
                double wp = grid_wts.is_valid(i_simd) ? grid_wts.p[ip][i_simd] : 0.0; // all these trinary ops are probably unnecessary

                field_correction.set(i_simd, grid.uses_rz_basis(node), gradpsi);

                gather_all_fields(push_controls, i_simd, node, wp, field_correction, grid_wts, rho_wts, fld);
            }

#ifdef NEOCLASSICAL_TEST
            fld.E.phi[i_simd]=0.0;
#else
#ifdef XGC1
            if(turb_efield){
                // Correct phi component of E-field
                double Bmag = B.magnitude(i_simd);

                // Get phi from para
                phi_from_para(fld.E, i_simd, B, Bmag);
#ifdef EXPLICIT_EM
                phi_from_para(fld.dAh, i_simd, B, Bmag);
                phi_from_para(fld.dAs, i_simd, B, Bmag);
#endif
            }else{
                fld.E.phi[i_simd]=0.0;
#ifdef EXPLICIT_EM
                fld.dAh.phi[i_simd]=0.0;
                fld.dAs.phi[i_simd]=0.0;
#endif
            }
#endif
#endif
        }
    }


    KOKKOS_INLINE_FUNCTION void get_Ah(const Grid<Device> &grid, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, const SimdGyroWeights<GyroKin>& rho_wts, Simd<double>& Ah) const
    {
        // Initialize local fields to 0
#ifdef EXPLICIT_EM
        Ah = 0.0;

        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            // If particle is inactive, set itr_work=1, phi=0D0, iphi=0
            int itr_work = grid_wts.is_valid(i_simd) ? grid_wts.itr[i_simd] : 1;

            int irho = rho_wts.irho(i_simd);

            // 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_wts.is_valid(i_simd) ? grid.get_node_index(itr_work, ip) : 0;
                double wp = grid_wts.is_valid(i_simd) ? grid_wts.p[ip][i_simd] : 0.0; // all these trinary ops are probably unnecessary

                get<Label::Ah_gyro>()(node,irho).gather(Ah,i_simd, wp, rho_wts, grid_wts.phi_wts, get<Label::Ah_gyro>()(node,irho+1));
            }
        }
#endif
    }

    KOKKOS_INLINE_FUNCTION void get_Ah(const Grid<Device> &grid, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, const SimdGyroWeights<DriftKin>& rho_wts, Simd<double>& Ah) const
    {
        // Initialize local fields to 0
#ifdef EXPLICIT_EM
        Ah = 0.0;

        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            // If particle is inactive, set itr_work=1, phi=0D0, iphi=0
            int itr_work = grid_wts.is_valid(i_simd) ? grid_wts.itr[i_simd] : 1;

            // Loop over the 3 vertices of the grid triangle and add up the contributions from each
            for (int ip = 0; ip<3; ip++){
                int i_node = grid_wts.is_valid(i_simd) ? grid.get_node_index(itr_work, ip) : 0;
                double wp = grid_wts.is_valid(i_simd) ? grid_wts.p[ip][i_simd] : 0.0; // all these trinary ops are probably unnecessary

                const int node = i_node - node_offset;
                constexpr int PHI_ZERO=0;
                get<Label::Ah>()(PHI_ZERO,node).gather(Ah,i_simd, wp, grid_wts.phi_wts);
            }
        }
#endif
    }

    KOKKOS_INLINE_FUNCTION void get_dpot(const Grid<Device> &grid, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, Simd<double>& dpot_out) const {

        dpot_out.zero();

        
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            // If particle is inactive, set itr_work=1, phi=0D0, iphi=0
            int itr_work = grid_wts.is_valid(i_simd) ? grid_wts.itr[i_simd] : 1;

            for(int ip=0; ip<3; ip++){
                int node = grid_wts.is_valid(i_simd) ? grid.get_node_index(itr_work, ip) : 0;
                double wp = grid_wts.is_valid(i_simd) ? grid_wts.p[ip][i_simd] : 0.0; // all these trinary ops are probably unnecessary

#ifdef XGC1
    //#ifndef EXPLICIT_EM
                // for electron -- rho=0 case, optimized.
                dpot_ff(node).gather(dpot_out,i_simd, wp, grid_wts.phi_wts);
    //#else
    //            dpot_out[i_simd]+=wp*dpot_n0(node);
    //#endif
#elif defined(XGCA)
                // Simple gather
                dpot(node).gather(dpot_out,i_simd,wp);
#endif
            }
        }
    }
};
#endif