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"

/* VarType specifies whether the field is a vector, a 2D vector, or a scalar
 * */
enum class VarType{
    Vector,
    Vector2D,
    Scalar
};

template<VarType VT, PhiInterpType PIT>
struct Field{};

//Struct used for E_phi_ff, dAs_phi_ff, dAh_phi_ff, As_phi_ff, Ah_phi_ff. 3 vector components, 2 contributing planes 
template<>
struct Field<VarType::Vector,PhiInterpType::Planes>{
    double V[3][2];

    KOKKOS_INLINE_FUNCTION void gather(SimdVector& vec, int i_simd, double wp, const double (&wphi)[2], const double (&rvec)[2], const double (&zvec)[2]) const{
        vec.r[i_simd]   += wp*(wphi[0]*( rvec[0]*V[PIR][0] + zvec[0]*V[PIZ][0] )
                             + wphi[1]*( rvec[0]*V[PIR][1] + zvec[0]*V[PIZ][1] ) );
        vec.z[i_simd]   += wp*(wphi[0]*( rvec[1]*V[PIR][0] + zvec[1]*V[PIZ][0] )
                             + wphi[1]*( rvec[1]*V[PIR][1] + zvec[1]*V[PIZ][1] ) );
        vec.phi[i_simd] += wp*(wphi[0]*  V[PIP][0]
                             + wphi[1]*  V[PIP][1] );
    }

    KOKKOS_INLINE_FUNCTION void gather(SimdVector& vec, int i_simd, double wp, const FieldWeights<GyroKin, PhiInterpType::Planes>& wts, const double (&rvec)[2], const double (&zvec)[2], const Field<VarType::Vector,PhiInterpType::Planes>& field2) const{
# ifdef NEWGYROMATRIX
        vec.r[i_simd]   += wp*(wts.phi.w[0]*( rvec[0]*(wts.rho[0].w[0]*V[PIR][0] + wts.rho[0].w[1]*field2.V[PIR][0])
                                            + zvec[0]*(wts.rho[0].w[0]*V[PIZ][0] + wts.rho[0].w[1]*field2.V[PIZ][0]) )
                             + wts.phi.w[1]*( rvec[0]*(wts.rho[1].w[0]*V[PIR][1] + wts.rho[1].w[1]*field2.V[PIR][1])
                                            + zvec[0]*(wts.rho[1].w[0]*V[PIZ][1] + wts.rho[1].w[1]*field2.V[PIZ][1]) ) );
        vec.z[i_simd]   += wp*(wts.phi.w[0]*( rvec[1]*(wts.rho[0].w[0]*V[PIR][0] + wts.rho[0].w[1]*field2.V[PIR][0])
                                            + zvec[1]*(wts.rho[0].w[0]*V[PIZ][0] + wts.rho[0].w[1]*field2.V[PIZ][0]) )
                             + wts.phi.w[1]*( rvec[1]*(wts.rho[1].w[0]*V[PIR][1] + wts.rho[1].w[1]*field2.V[PIR][1])
                                            + zvec[1]*(wts.rho[1].w[0]*V[PIZ][1] + wts.rho[1].w[1]*field2.V[PIZ][1]) ) );
        vec.phi[i_simd] += wp*(wts.phi.w[0]*  (wts.rho[0].w[0]*V[PIP][0] + wts.rho[0].w[1]*field2.V[PIP][0])
                             + wts.phi.w[1]*  (wts.rho[1].w[0]*V[PIP][1] + wts.rho[1].w[1]*field2.V[PIP][1]) );
# else
        vec.r[i_simd]   += wp*(wts.phi.w[0]*( rvec[0]*(wts.rho.w[0]*V[PIR][0] + wts.rho.w[1]*field2.V[PIR][0])
                                            + zvec[0]*(wts.rho.w[0]*V[PIZ][0] + wts.rho.w[1]*field2.V[PIZ][0]) )
                             + wts.phi.w[1]*( rvec[0]*(wts.rho.w[0]*V[PIR][1] + wts.rho.w[1]*field2.V[PIR][1])
                                            + zvec[0]*(wts.rho.w[0]*V[PIZ][1] + wts.rho.w[1]*field2.V[PIZ][1]) ) );
        vec.z[i_simd]   += wp*(wts.phi.w[0]*( rvec[1]*(wts.rho.w[0]*V[PIR][0] + wts.rho.w[1]*field2.V[PIR][0])
                                            + zvec[1]*(wts.rho.w[0]*V[PIZ][0] + wts.rho.w[1]*field2.V[PIZ][0]) )
                             + wts.phi.w[1]*( rvec[1]*(wts.rho.w[0]*V[PIR][1] + wts.rho.w[1]*field2.V[PIR][1])
                                            + zvec[1]*(wts.rho.w[0]*V[PIZ][1] + wts.rho.w[1]*field2.V[PIZ][1]) ) );
        vec.phi[i_simd] += wp*(wts.phi.w[0]*  (wts.rho.w[0]*V[PIP][0] + wts.rho.w[1]*field2.V[PIP][0])
                             + wts.phi.w[1]*  (wts.rho.w[0]*V[PIP][1] + wts.rho.w[1]*field2.V[PIP][1]) );
# endif
    }

};

//Struct used for E00_ff. 2 vector components, 2 contributing planes
template<>
struct Field<VarType::Vector2D,PhiInterpType::Planes>{
    double V[2][2];

    KOKKOS_INLINE_FUNCTION void gather(SimdVector& vec, int i_simd, double wp, const double (&wphi)[2], const double (&rvec)[2], const double (&zvec)[2]) const{
        vec.r[i_simd]   += wp*(wphi[0]*( rvec[0]*V[PIR][0] + zvec[0]*V[PIZ][0] )
                             + wphi[1]*( rvec[0]*V[PIR][1] + zvec[0]*V[PIZ][1] ) );
        vec.z[i_simd]   += wp*(wphi[0]*( rvec[1]*V[PIR][0] + zvec[1]*V[PIZ][0] )
                             + wphi[1]*( rvec[1]*V[PIR][1] + zvec[1]*V[PIZ][1] ) );
    }
};

//Struct used for ddpotdt_phi_ff, As_phi_ff, Ah_phi_ff, Ah_cv for EM.  2 contributing planes
template<>
struct Field<VarType::Scalar,PhiInterpType::Planes>{
    double S[2];

    KOKKOS_INLINE_FUNCTION void gather(Simd<double>& sca, int i_simd, double wp, const double (&wphi)[2]) const{
        sca[i_simd] += wp* ( wphi[0]*S[0] + wphi[1]*S[1] );
    }

    KOKKOS_INLINE_FUNCTION void gather(Simd<double>& sca, int i_simd, double wp, const FieldWeights<GyroKin, PhiInterpType::Planes>& wts, const Field<VarType::Scalar,PhiInterpType::Planes>& field2) const{
# ifdef NEWGYROMATRIX
        sca[i_simd] += wp* (  wts.phi.w[0]*wts.rho[0].w[0]*S[0]
                            + wts.phi.w[1]*wts.rho[1].w[0]*S[1]
                            + wts.phi.w[0]*wts.rho[0].w[1]*field2.S[0]
                            + wts.phi.w[1]*wts.rho[1].w[1]*field2.S[1]);
# else
        sca[i_simd] += wp* (  wts.phi.w[0]*wts.rho.w[0]*S[0]
                            + wts.phi.w[1]*wts.rho.w[0]*S[1]
                            + wts.phi.w[0]*wts.rho.w[1]*field2.S[0]
                            + wts.phi.w[1]*wts.rho.w[1]*field2.S[1]);
# endif
    }
};

// Struct used for E_rho and du2_E_rho. 2 vector components (r,z)
template<>
struct Field<VarType::Vector2D,PhiInterpType::None>{
    double E[2];

    KOKKOS_INLINE_FUNCTION void gather(SimdVector& vec, int i_simd, double wp, const double (&rvec)[2], const double (&zvec)[2]) const{
        vec.r[i_simd] += wp*( rvec[0]*E[PIR] + zvec[0]*E[PIZ] );
        vec.z[i_simd] += wp*( rvec[1]*E[PIR] + zvec[1]*E[PIZ] );
    }

    KOKKOS_INLINE_FUNCTION void gather(SimdVector& vec, int i_simd, double wp, const FieldWeights<GyroKin, PhiInterpType::None>& wts, const double (&rvec)[2], const double (&zvec)[2], const Field<VarType::Vector2D,PhiInterpType::None>& field2) const{
        vec.r[i_simd] += wp*( rvec[0]*(wts.rho.w[0]*E[PIR] + wts.rho.w[1]*field2.E[PIR])
                             +zvec[0]*(wts.rho.w[0]*E[PIZ] + wts.rho.w[1]*field2.E[PIZ]) );
        vec.z[i_simd] += wp*( rvec[1]*(wts.rho.w[0]*E[PIR] + wts.rho.w[1]*field2.E[PIR])
                             +zvec[1]*(wts.rho.w[0]*E[PIZ] + wts.rho.w[1]*field2.E[PIZ]) );
    }

};

// Struct for scalar without planar interpolation (i.e. just a double)
template<>
struct Field<VarType::Scalar,PhiInterpType::None>{
    double S;

    KOKKOS_INLINE_FUNCTION void gather(Simd<double>& sca, int i_simd, double wp){
        sca[i_simd] += wp*S;
    }
};

// Fortran calls
#ifdef XGC1
extern "C" void set_rho_ff_pointers(int nrho, int n, double* pot0, double* dpot, Field<VarType::Scalar,PhiInterpType::Planes>* dpot_ff, Field<VarType::Scalar,PhiInterpType::Planes>* pot_rho_ff, Field<VarType::Vector,PhiInterpType::Planes>* E_rho_ff
#  ifdef EXPLICIT_EM
      , double* Ah, Field<VarType::Scalar,PhiInterpType::Planes>* Ah_rho_ff, Field<VarType::Vector,PhiInterpType::Planes>* dAh_rho_ff, double* As, Field<VarType::Scalar,PhiInterpType::Planes>* As_rho_ff, Field<VarType::Vector,PhiInterpType::Planes>* dAs_rho_ff, Field<VarType::Scalar,PhiInterpType::Planes>* E_para_em_rho_ff
#  endif
      );
#else
extern "C" void set_rho_pointers(int nrho, int n, double* pot0, double* dpot, Field<VarType::Vector2D,PhiInterpType::None>* E_rho
#  ifdef SONIC_GK
      , Field<VarType::Vector2D,PhiInterpType::None>* du2_E_rho
#  endif
      );
#endif

enum class GridFieldOpts{
    WithRhoFF = 0,
    WithoutRhoFF
};

// Struct for each field that contains different variations of the data: global vs one plane, field-following vs not, maybe host vs device, maybe XGCA vs XGC1
template<VarType VT, PhiInterpType PIT>
struct GridField{
    GridField(){}

    // XGC1
    Kokkos::View<Field<VT,PIT>*,Kokkos::LayoutRight,HostType> ff_h; 
    Kokkos::View<Field<VT,PIT>**,Kokkos::LayoutRight,HostType> rho_ff_h; 

    GridField(int nphi, int nrho, int nnode, GridFieldOpts opt=GridFieldOpts::WithRhoFF) {
        if(opt==GridFieldOpts::WithoutRhoFF){
            ff_h = Kokkos::View<Field<VT,PIT>*,Kokkos::LayoutRight,HostType>("ff_h",nnode);
        }else{
            rho_ff_h = Kokkos::View<Field<VT,PIT>**,Kokkos::LayoutRight,HostType>("rho_ff_h",nnode,nrho+1);
        }
    }

    // XGCa
    Kokkos::View<Field<VT,PIT>**,Kokkos::LayoutRight,HostType> rho_h;

    // XGCa constructor is distinct since it has 2 arguments vs 3-4 in above XGC1 constructor
    GridField(int nrho, int nnode)
      : rho_h("rho_h",nnode, nrho+1) {}
};


template<class Device, PhiInterpType PIT>
struct GridFieldPack {
    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; 

    // Device Views
#ifdef XGC1
    Kokkos::View<Field<VarType::Vector,PIT>**,Kokkos::LayoutRight,Device> E_phi_ff;
    Kokkos::View<Field<VarType::Vector,PIT>**,Kokkos::LayoutRight,Device> E_rho_ff;
    Kokkos::View<Field<VarType::Scalar,PIT>*,Kokkos::LayoutRight,Device> dpot_ff;
#  ifdef DELTAF_CONV
    Kokkos::View<Field<VarType::Vector2D,PIT>*,Kokkos::LayoutRight,Device> E00_ff;
    Kokkos::View<Field<VarType::Scalar,PIT>**,Kokkos::LayoutRight,Device> ddpotdt_phi_ff;
#  endif
#  ifdef EXPLICIT_EM
    //For EM
    Kokkos::View<Field<VarType::Vector,PIT>**,Kokkos::LayoutRight,Device> dAh_phi_ff;
    Kokkos::View<Field<VarType::Vector,PIT>**,Kokkos::LayoutRight,Device> dAh_rho_ff;
    Kokkos::View<Field<VarType::Scalar,PIT>**,Kokkos::LayoutRight,Device> Ah_phi_ff;
    Kokkos::View<Field<VarType::Scalar,PIT>**,Kokkos::LayoutRight,Device> Ah_rho_ff;
    Kokkos::View<Field<VarType::Vector,PIT>**,Kokkos::LayoutRight,Device> dAs_phi_ff;
    Kokkos::View<Field<VarType::Vector,PIT>**,Kokkos::LayoutRight,Device> dAs_rho_ff;
    Kokkos::View<Field<VarType::Scalar,PIT>**,Kokkos::LayoutRight,Device> As_phi_ff;
    Kokkos::View<Field<VarType::Scalar,PIT>**,Kokkos::LayoutRight,Device> As_rho_ff;
    Kokkos::View<Field<VarType::Scalar,PIT>**,Kokkos::LayoutRight,Device> Ah_cv_phi_ff;
    Kokkos::View<Field<VarType::Scalar,PIT>**,Kokkos::LayoutRight,Device> Epar_em_phi_ff;
    Kokkos::View<Field<VarType::Scalar,PIT>**,Kokkos::LayoutRight,Device> Epar_em_rho_ff;
#  endif
#else
    Kokkos::View<Field<VarType::Vector2D,PIT>**,Kokkos::LayoutRight,Device> E_rho;
    Kokkos::View<double*,Kokkos::LayoutRight,Device> dpot;
#  ifdef SONIC_GK
    View<Field<VarType::Vector2D,PIT>**,CLayout,Device> du2_E_rho;
#  endif
#endif

    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:

    // A correction to the field
    KOKKOS_INLINE_FUNCTION void field_correction(int i_simd, double basis, double gamma_psi, const SimdVector2D &gradpsi, double* rvec, double* zvec) const {
        rvec[0]=basis + (1.0-basis) * gamma_psi *  gradpsi.r[i_simd];
        rvec[1]=        (1.0-basis) * gamma_psi *  gradpsi.z[i_simd];
        zvec[0]=        (1.0-basis) * gamma_psi *(-gradpsi.z[i_simd]);
        zvec[1]=basis + (1.0-basis) * gamma_psi *  gradpsi.r[i_simd];
    }

    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 double (&rvec)[2], const double (&zvec)[2], 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, rvec, zvec); // irho = 0
#else
        const int i_phi = (wts.phi.i - phi_offset)%E_phi_ff.extent(0); //grid.nplanes;
        E_phi_ff(i_phi,node).gather(fld.E,i_simd, wp, wts.phi.w, rvec, zvec);
#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, rvec, zvec);

        if(push_controls.em_control_variate) {
            Ah_cv_phi_ff(i_phi,node).gather(fld.Ah_cv,i_simd, wp, wts.phi.w);
        }

        if(push_controls.em_mixed_variable) {
            dAs_phi_ff(i_phi,node).gather(fld.dAs,i_simd, wp, wts.phi.w, rvec, zvec);
            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, rvec, zvec);
        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 double (&rvec)[2], const double (&zvec)[2], 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, rvec, zvec, E_rho(node,irho+1));
#  ifdef SONIC_GK
        du2_E_rho(node,0).gather(fld.du2_E,i_simd, wp, rvec, zvec); // irho = 0
#  endif
#else
        E_rho_ff(node,irho).gather(fld.E,i_simd, wp, wts, rvec, zvec, 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, rvec, zvec, dAh_rho_ff(node,irho+1));

        if(push_controls.em_mixed_variable) {
            dAs_rho_ff(node,irho).gather(fld.dAs,i_simd, wp, wts, rvec, zvec, 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.Ah_cv.zero();
        fld.Epar_em.zero();
#endif
#ifdef SONIC_GK
        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);

            double gamma_psi=1.0/sqrt(gradpsi.r[i_simd]*gradpsi.r[i_simd] + gradpsi.z[i_simd]*gradpsi.z[i_simd]);

            double rvec[2];
            double zvec[2];

            // 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.nodes(itr_work - 1).vertex[ip] - 1) : 0; // - 1 since its 0-indexed now
                double wp = itr[i_simd]>0 ? p[ip][i_simd] : 0.0; // all these trinary ops are probably unnecessary

                field_correction(i_simd, grid.basis(node), gamma_psi, gradpsi, rvec, zvec);

                gather_all_fields(push_controls, i_simd, node, wp, rvec, zvec, 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=sqrt(B.r[i_simd]*B.r[i_simd] + B.z[i_simd]*B.z[i_simd] + B.phi[i_simd]*B.phi[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);
#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.nodes(itr_work - 1).vertex[ip] - 1) : 0; // - 1 since its 0-indexed now
                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);

            // 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.nodes(itr_work - 1).vertex[ip] - 1) : 0; // - 1 since its 0-indexed now
                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_Ah_cv(const Grid<Device> &grid, const Simd<double>& fld_phi, const Simd<int>& itr, const SimdGridVec &p, SimdGyroRadius<DriftKin>& rho, Simd<double>& Ah_cv) const
    {
        // Initialize local fields to 0
#ifdef EXPLICIT_EM
        Ah_cv = 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);

            // 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.nodes(itr_work - 1).vertex[ip] - 1) : 0; // - 1 since its 0-indexed now
                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_cv_phi_ff(i_phi,node).gather(Ah_cv,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);

            for(int ip=0; ip<3; ip++){
                int node = itr[i_simd]>0 ? (grid.nodes(itr_work - 1).vertex[ip] - 1) : 0; // - 1 since its 0-indexed now
                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_out[i_simd]+=wp*dpot(node);
#endif
            }
        }
    }
};
#endif