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 "field_weights.hpp"
#include "gyro_radius.hpp"
#include "grid_field.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{
    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; 

    int nrho;         
    double inv_drho;  

    // Device Views
#ifdef XGC1
    GridField<Device, vec2d_if_axisym<PIT>(),PIT, TorType::MultiplePlanes, KinType::DriftKin> E_phi_ff;
    GridField<Device, vec2d_if_axisym<PIT>(),PIT, TorType::OnePlane, KinType::GyroKin> E_rho_ff;
    GridField<Device, VarType::Scalar,PIT, TorType::OnePlane, KinType::DriftKin> dpot_ff;
#  ifdef DELTAF_CONV
    GridField<Device, VarType::Vector2D,PIT, TorType::OnePlane, KinType::DriftKin> E00_ff;
    GridField<Device, VarType::Scalar,PIT, TorType::MultiplePlanes, KinType::DriftKin> ddpotdt_phi_ff;
#  endif
#  ifdef EXPLICIT_EM
    //For EM
    GridField<Device, vec2d_if_axisym<PIT>(),PIT, TorType::MultiplePlanes, KinType::DriftKin> dAh_phi_ff;
    GridField<Device, vec2d_if_axisym<PIT>(),PIT, TorType::OnePlane, KinType::GyroKin> dAh_rho_ff;
    GridField<Device, VarType::Scalar,PIT, TorType::MultiplePlanes, KinType::DriftKin> Ah_phi_ff;
    GridField<Device, VarType::Scalar,PIT, TorType::OnePlane, KinType::GyroKin> Ah_rho_ff;
    GridField<Device, vec2d_if_axisym<PIT>(),PIT, TorType::MultiplePlanes, KinType::DriftKin> dAs_phi_ff;
    GridField<Device, vec2d_if_axisym<PIT>(),PIT, TorType::OnePlane, KinType::GyroKin> dAs_rho_ff;
    GridField<Device, VarType::Scalar,PIT, TorType::MultiplePlanes, KinType::DriftKin> As_phi_ff;
    GridField<Device, VarType::Scalar,PIT, TorType::OnePlane, KinType::GyroKin> As_rho_ff;
    GridField<Device, VarType::Scalar,PIT, TorType::MultiplePlanes, KinType::DriftKin> Epar_em_phi_ff;
    GridField<Device, VarType::Scalar,PIT, TorType::OnePlane, KinType::GyroKin> Epar_em_rho_ff;
#  endif
#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

    GridFieldPack()
        : is_initialized(false) {}

    GridFieldPack(KinType kintype_in, bool turb_efield_in, int nrho_in, double inv_drho_in)
        : is_initialized(true),
          kintype(kintype_in),
          turb_efield(turb_efield_in),
          nrho(nrho_in),
          inv_drho(inv_drho_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 FieldWeights<DriftKin, PIT>& 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
#else
        const int i_phi = (wts.phi.i - phi_offset)%E_phi_ff.f.extent(0); //should be grid.nplanes (?);
        E_phi_ff(i_phi,node).gather(fld.E,i_simd, wp, wts.phi.w, corr);
#ifdef EXPLICIT_EM
        Ah_phi_ff(i_phi,node).gather(fld.Ah,i_simd, wp, wts.phi.w);
        dAh_phi_ff(i_phi,node).gather(fld.dAh,i_simd, wp, wts.phi.w, corr);

        if(push_controls.em_mixed_variable) {
            dAs_phi_ff(i_phi,node).gather(fld.dAs,i_simd, wp, wts.phi.w, corr);
            As_phi_ff(i_phi,node).gather(fld.As,i_simd, wp, wts.phi.w);
        }
        if(push_controls.em_mixed_variable && (push_controls.em_pullback_mode==4)) {
            Epar_em_phi_ff(i_phi,node).gather(fld.Epar_em,i_simd,wp,wts.phi.w);
        }
#endif
#endif

#if defined(XGC1) && defined(DELTAF_CONV)
        E00_ff(node).gather(fld.E00,i_simd, wp, wts.phi.w, corr);
        ddpotdt_phi_ff(i_phi,node).gather(fld.ddpotdt,i_simd, wp, wts.phi.w);
#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 FieldWeights<GyroKin, PIT>& wts, LocalFields& fld) 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
#ifdef XGCA
        E_rho(node,irho).gather(fld.E,i_simd, wp, 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
#else
        E_rho_ff(node,irho).gather(fld.E,i_simd, wp, wts, corr, E_rho_ff(node,irho+1));
#ifdef EXPLICIT_EM
        Ah_rho_ff(node,irho).gather(fld.Ah,i_simd, wp, wts, Ah_rho_ff(node,irho+1));
        dAh_rho_ff(node,irho).gather(fld.dAh,i_simd, wp, wts, corr, dAh_rho_ff(node,irho+1));

        if(push_controls.em_mixed_variable) {
            dAs_rho_ff(node,irho).gather(fld.dAs,i_simd, wp, wts, corr, dAs_rho_ff(node,irho+1));
            As_rho_ff(node,irho).gather(fld.As,i_simd, wp, wts, As_rho_ff(node,irho+1));
        }
        if(push_controls.em_mixed_variable && (push_controls.em_pullback_mode==4)) {
            Epar_em_rho_ff(node,irho).gather(fld.Epar_em,i_simd,wp,wts, Epar_em_rho_ff(node,irho+1));
        }
#endif
#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 Simd<double>& fld_phi, const Simd<int>& itr, const SimdGridVec &p, SimdGyroRadius<PT>& rho, LocalFields& fld) const
    {
        // Initialize local fields to 0
        fld.E.zero();
#ifdef DELTAF_CONV
        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 = itr[i_simd]>0 ? itr[i_simd] : 1;

            double phi= itr[i_simd]>0 ? fld_phi[i_simd] : 0.0;

            // Get rho and phi weights
            FieldWeights<PT, PIT> wts(grid, phi, i_simd, rho, inv_drho, nrho);

            // 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 = itr[i_simd]>0 ? grid.get_node_index(itr_work, ip) : 0;
                double wp = itr[i_simd]>0 ? 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, 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 Simd<double>& fld_phi, const Simd<int>& itr, const SimdGridVec &p, SimdGyroRadius<GyroKin>& rho, 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 = itr[i_simd]>0 ? itr[i_simd] : 1;
            double phi= itr[i_simd]>0 ? fld_phi[i_simd] : 0.0;

            // Get rho and phi weights
            FieldWeights<GyroKin, PIT> wts(grid, phi, i_simd, rho, inv_drho, nrho);
#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

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

                Ah_rho_ff(node,irho).gather(Ah,i_simd, wp, wts, Ah_rho_ff(node,irho+1));
            }
        }
#endif
    }

    KOKKOS_INLINE_FUNCTION void get_Ah(const Grid<Device> &grid, const Simd<double>& fld_phi, const Simd<int>& itr, const SimdGridVec &p, SimdGyroRadius<DriftKin>& rho, 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 = itr[i_simd]>0 ? itr[i_simd] : 1;
            double phi= itr[i_simd]>0 ? fld_phi[i_simd] : 0.0;

            // Get rho and phi weights
            FieldWeights<DriftKin, PIT> wts(grid, phi, i_simd, rho, inv_drho, nrho);

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

                const int i_phi = (wts.phi.i - phi_offset)%grid.nplanes;
                const int node = i_node - node_offset;

                Ah_phi_ff(i_phi,node).gather(Ah,i_simd, wp, wts.phi.w);
            }
        }
#endif
    }

    KOKKOS_INLINE_FUNCTION void get_dpot(const Grid<Device> &grid, const Simd<double>& fld_phi, const Simd<int>& itr, const SimdGridVec &p, 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 = itr[i_simd]>0 ? itr[i_simd] : 1;

            double phi = itr[i_simd]>0 ? fld_phi[i_simd] : 0.0;
            SimdGyroRadius<DriftKin> rho;
            FieldWeights<DriftKin, PIT> wts(grid, phi, i_simd, rho, inv_drho, nrho);

            for(int ip=0; ip<3; ip++){
                int node = itr[i_simd]>0 ? grid.get_node_index(itr_work, ip) : 0;
                double wp = itr[i_simd]>0 ? 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, wts.phi.w);
    //#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