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{
    virtual ~GridFieldPackGeneric() = default; // Virtual destructor
};

using GridFieldPackPtr = std::unique_ptr<GridFieldPackGeneric>;

template<class Device, typename GFPT>
struct GridFieldPack : GridFieldPackGeneric{
    using pack_type = GFPT;

    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; 

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

    // 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_GLOBAL, TorType::OnePlane, KinType::DriftKin> loop_voltage;

    GridFieldPack() {}

    GridFieldPack(bool turb_efield_in)
        : turb_efield(turb_efield_in),
          phi_offset(0),
          node_offset(0) {}

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

#ifndef XGCA
    template<typename FieldT>
    KOKKOS_INLINE_FUNCTION void gather_deltaf_conv(int i_simd, int node, int i_phi, double wp, const FieldCorrection& corr, const SimdGridWeights<Order::One, PIT_GLOBAL>& grid_wts, LocalFields<GFPT>& fld, const FieldT& labeled_field) const {
        if constexpr(FieldT::name == Label::E00){
            get<Label::E00>()(node).gather(fld.template get<Label::E00>(),i_simd, wp, grid_wts.phi_wts, corr);
            get<Label::ddpotdt>()(i_phi,node).gather(fld.template get<Label::ddpotdt>(),i_simd, wp, grid_wts.phi_wts);
        }
    }
#endif

    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<GFPT>& fld) const {
        const int node = i_node - node_offset;
#ifdef XGCA
        get<Label::E>()(node,0).gather(fld.template get<Label::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.template get<Label::E>(),i_simd, wp, grid_wts.phi_wts, corr);
#ifdef EXPLICIT_EM
        get<Label::Ah>()(i_phi,node).gather(fld.template get<Label::Ah>(),i_simd, wp, grid_wts.phi_wts);
        get<Label::dAh>()(i_phi,node).gather(fld.template get<Label::dAh>(),i_simd, wp, grid_wts.phi_wts, corr);

        if(push_controls.em_mixed_variable) {
            get<Label::dAs>()(i_phi,node).gather(fld.template get<Label::dAs>(),i_simd, wp, grid_wts.phi_wts, corr);
            get<Label::As>()(i_phi,node).gather(fld.template get<Label::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.template get<Label::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);


        // Could replace with new function like: if constexpr(pack.is_present<Label::E00>())
        fld.pack.for_each_labeled([&](auto& labeled_field) {
            gather_deltaf_conv(i_simd, node, i_phi, wp, corr, grid_wts, fld, labeled_field);
        });
#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<GFPT>& fld) const {
        int irho = rho_wts.irho(i_simd);
#ifdef XGCA
        get<Label::E>()(node,irho).gather(fld.template get<Label::E>(),i_simd, wp, rho_wts, corr, get<Label::E>()(node,irho+1));
#  ifdef SONIC_GK
        get<Label::dEr_B2>()(node,0).gather(fld.template get<Label::dEr_B2>(),i_simd, wp, corr); // irho = 0
        get<Label::dEz_B2>()(node,0).gather(fld.template get<Label::dEz_B2>(),i_simd, wp, corr); // irho = 0
        get<Label::du2_E>()(node,0).gather(fld.template get<Label::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>()(node,irho).gather(fld.template get<Label::E>(),i_simd, wp, rho_wts, grid_wts.phi_wts, corr, get<Label::E>()(node,irho+1));
#ifdef EXPLICIT_EM
        get<Label::Ah>()(node,irho).gather(fld.template get<Label::Ah>(),i_simd, wp, rho_wts, grid_wts.phi_wts, get<Label::Ah>()(node,irho+1));
        get<Label::dAh>()(node,irho).gather(fld.template get<Label::dAh>(),i_simd, wp, rho_wts, grid_wts.phi_wts, corr, get<Label::dAh>()(node,irho+1));

        if(push_controls.em_mixed_variable) {
            get<Label::dAs>()(node,irho).gather(fld.template get<Label::dAs>(),i_simd, wp, rho_wts, grid_wts.phi_wts, corr, get<Label::dAs>()(node,irho+1));
            get<Label::As>()(node,irho).gather(fld.template get<Label::As>(),i_simd, wp, rho_wts, grid_wts.phi_wts, get<Label::As>()(node,irho+1));
        }
        if(push_controls.em_mixed_variable && (push_controls.em_pullback_mode==4)) {
            get<Label::Epar_em>()(node,irho).gather(fld.template get<Label::Epar_em>(),i_simd,wp,rho_wts, grid_wts.phi_wts, get<Label::Epar_em>()(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<GFPT>& fld) const
    {
        // Initialize local fields to 0
        fld.pack.for_each([&](auto& field) { field.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;

            // 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.template get<Label::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.template get<Label::E>(), i_simd, B, Bmag);
#ifdef EXPLICIT_EM
                phi_from_para(fld.template get<Label::dAh>(), i_simd, B, Bmag);
                phi_from_para(fld.template get<Label::dAs>(), i_simd, B, Bmag);
#endif
            }else{
                fld.template get<Label::E>().phi[i_simd]=0.0;
#ifdef EXPLICIT_EM
                fld.template get<Label::dAh>().phi[i_simd]=0.0;
                fld.template get<Label::dAs>().phi[i_simd]=0.0;
#endif
            }
#endif
#endif
        }
    }
};
#endif