electric_field.hpp

Go to the documentation of this file.

#ifndef ELECTRIC_FIELD_HPP
#define ELECTRIC_FIELD_HPP
#include "space_settings.hpp"
#include "NamelistReader.hpp"
#include "grid.hpp"
#include "linear_weights.hpp"
#include "local_fields.hpp"

//Struct used for E_phi_ff, dAs_phi_ff, dAh_phi_ff, As_phi_ff, Ah_phi_ff. 3 vector components, 2 contributing planes 
struct VectorFieldPlanes {
    double V[3][2];
};

//Struct used for E00_ff. 2 vector components, 2 contributing planes
struct Vector2DFieldPlanes {
    double V[2][2];
};

//Struct used for ddpotdt_phi, As_phi_ff, Ah_phi_ff, Ah_cv for EM.  2 contributing planes
struct ScalarFieldPlanes {
    double S[2];
}; 

// Struct used for E_rho. 2 vector components (r,z) --> can be combined to ScalarFieldPlanes
struct FieldXGCa {
    double E[2];
};

// Fortran calls
#ifdef XGC1
extern "C" void set_rho_ff_pointers(int nrho, int n, double* dpot, ScalarFieldPlanes* dpot_ff, ScalarFieldPlanes* pot_rho_ff, VectorFieldPlanes* E_rho_ff
#  ifdef EXPLICIT_EM
      , double* Ah, ScalarFieldPlanes* Ah_rho_ff, VectorFieldPlanes* dAh_rho_ff, double* As, ScalarFieldPlanes* As_rho_ff, VectorFieldPlanes* dAs_rho_ff, ScalarFieldPlanes* E_para_em_rho_ff
#  endif
      );
#else
extern "C" void set_rho_pointers(int nrho, int n, double* dpot, FieldXGCa* E_rho);
#endif

enum class GridFieldOpts{
    WithBothFF = 0,
    WithoutRhoFF,
    WithoutPhiFF
};

// 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<typename T>
struct GridField{
#ifdef XGC1
    Kokkos::View<T**,Kokkos::LayoutRight,HostType> ff_h; 
    Kokkos::View<T**,Kokkos::LayoutRight,HostType> phi_ff_h; 
    Kokkos::View<T**,Kokkos::LayoutRight,HostType> rho_ff_h; 

    GridField(int nphi, int nrho, int nnode, GridFieldOpts opt=GridFieldOpts::WithBothFF) {
        if(opt!=GridFieldOpts::WithoutRhoFF){
            rho_ff_h = Kokkos::View<T**,Kokkos::LayoutRight,HostType>("rho_ff_h",nnode,nrho+1);
        }
        if(opt!=GridFieldOpts::WithoutPhiFF){
            phi_ff_h = Kokkos::View<T**,Kokkos::LayoutRight,HostType>("phi_ff_h",nphi,nnode);
        }
    }
#else
    Kokkos::View<T**,Kokkos::LayoutRight,HostType> rho_h;

    GridField(int nrho, int nnode)
      : rho_h("rho_h",nnode, nrho+1) {}
#endif
};

// Electric field class
template<class Device>
struct ElectricField {

    // Dimensions
    int nnode;
    int nphi;
    int nrho;

    // Field variables
    bool turb_efield;      

    // Host Views
#ifdef XGC1
    GridField<ScalarFieldPlanes> pot;
    GridField<VectorFieldPlanes> E;
    Kokkos::View<double**,Kokkos::LayoutRight,HostType> dpot_h;
    Kokkos::View<ScalarFieldPlanes*,Kokkos::LayoutRight,HostType> dpot_ff_h;
#ifdef DELTAF_CONV
    Kokkos::View<Vector2DFieldPlanes*,Kokkos::LayoutRight,HostType> E00_ff_h; 
    GridField<ScalarFieldPlanes> ddpotdt; // Not field-following?
#endif
#ifdef EXPLICIT_EM
    //For EM
    Kokkos::View<double**,Kokkos::LayoutRight,HostType> Ah_h;
    GridField<VectorFieldPlanes> dAh;
    GridField<ScalarFieldPlanes> Ah;
    Kokkos::View<double**,Kokkos::LayoutRight,HostType> As_h;
    GridField<VectorFieldPlanes> dAs;
    GridField<ScalarFieldPlanes> As;
    GridField<ScalarFieldPlanes> Ah_cv;
    GridField<ScalarFieldPlanes> Epar_em;

    //Kokkos::View<double*,Kokkos::LayoutRight,HostType> dpot_n0_h;
#endif
#else
    GridField<FieldXGCa> E;
    Kokkos::View<double*,Kokkos::LayoutRight,HostType> dpot_h;
#endif

    // Device Views
#ifdef XGC1
    Kokkos::View<VectorFieldPlanes**,Kokkos::LayoutRight,Device> E_phi_ff;
    Kokkos::View<VectorFieldPlanes**,Kokkos::LayoutRight,Device> E_rho_ff;
    Kokkos::View<ScalarFieldPlanes*,Kokkos::LayoutRight,Device> dpot_ff;
#ifdef DELTAF_CONV
    Kokkos::View<Vector2DFieldPlanes*,Kokkos::LayoutRight,Device> E00_ff;
    Kokkos::View<ScalarFieldPlanes**,Kokkos::LayoutRight,Device> ddpotdt_phi;
#endif
#ifdef EXPLICIT_EM
    //For EM
    Kokkos::View<VectorFieldPlanes**,Kokkos::LayoutRight,Device> dAh_phi_ff;
    Kokkos::View<VectorFieldPlanes**,Kokkos::LayoutRight,Device> dAh_rho_ff;
    Kokkos::View<ScalarFieldPlanes**,Kokkos::LayoutRight,Device> Ah_phi_ff;
    Kokkos::View<ScalarFieldPlanes**,Kokkos::LayoutRight,Device> Ah_rho_ff;
    Kokkos::View<VectorFieldPlanes**,Kokkos::LayoutRight,Device> dAs_phi_ff;
    Kokkos::View<VectorFieldPlanes**,Kokkos::LayoutRight,Device> dAs_rho_ff;
    Kokkos::View<ScalarFieldPlanes**,Kokkos::LayoutRight,Device> As_phi_ff;
    Kokkos::View<ScalarFieldPlanes**,Kokkos::LayoutRight,Device> As_rho_ff;
    Kokkos::View<ScalarFieldPlanes**,Kokkos::LayoutRight,Device> Ah_cv_phi_ff;
    Kokkos::View<ScalarFieldPlanes**,Kokkos::LayoutRight,Device> Epar_em_phi_ff;
    Kokkos::View<ScalarFieldPlanes**,Kokkos::LayoutRight,Device> Epar_em_rho_ff;

    //Kokkos::View<double*,Kokkos::LayoutRight,Device> dpot_n0;
#endif
#else
    Kokkos::View<FieldXGCa**,Kokkos::LayoutRight,Device> E_rho;
    Kokkos::View<double*,Kokkos::LayoutRight,Device> dpot;
#endif

    // Default constructor
    ElectricField(){}

    // Makes sense to use the Grid in the constructor here, but the Grid still copies from fortran so until that changes,
    // it's simpler to not use it here so that ElectricField can be defined before the restart, which uses psn%As
    ElectricField(NLReader::NamelistReader& nlr, int nnode_in, int nplanes_in, int nrho_in) : //Grid<DeviceType>& grid) :
      /*nnode(grid.nnode),
      nphi(grid.nplanes),
      nrho(grid.nrho),*/
      nnode(nnode_in),
      nphi(nplanes_in),
      nrho(nrho_in),
      // Host Views (only phi_ff during this intermediate PR)
#ifdef XGC1
      pot(nphi,nrho,nnode,GridFieldOpts::WithoutPhiFF),
      E(nphi,nrho,nnode),
#  ifdef DELTAF_CONV 
      E00_ff_h("E00_ff_h",nnode),
      ddpotdt(nphi,nrho,nnode,GridFieldOpts::WithoutRhoFF),
#  endif
#  ifdef EXPLICIT_EM
      dAs(nphi,nrho,nnode),
      dAh(nphi,nrho,nnode),
      As(nphi,nrho,nnode),
      Ah(nphi,nrho,nnode),
      Ah_cv(nphi,nrho,nnode,GridFieldOpts::WithoutRhoFF),
      Epar_em(nphi,nrho,nnode),
      Ah_h("Ah_h",4,nnode), // -1:2 in fortran
      As_h("As_h",4,nnode), // -1:2 in fortran
#  endif
      dpot_h("dpot_h",4,nnode), // -1:2 in fortran
      dpot_ff_h("dpot_ff_h",nnode),
#else
      E(nrho, nnode),
      dpot_h("dpot_h",nnode),
#endif

      // Device Views
#ifdef XGC1
      E_phi_ff("E_phi_ff",0,0),
      E_rho_ff("E_rho_ff",0,0),
#  ifdef DELTAF_CONV 
      E00_ff("E00_ff",0),
      ddpotdt_phi("ddpotdt_phi",0,0),
#  endif
#  ifdef EXPLICIT_EM
      dAs_phi_ff("dAs_phi_ff",0,0),
      dAs_rho_ff("dAs_rho_ff",0,0),
      dAh_phi_ff("dAh_phi_ff",0,0),
      dAh_rho_ff("dAh_rho_ff",0,0),
      As_phi_ff("As_phi_ff",0,0),
      As_rho_ff("As_rho_ff",0,0),
      Ah_phi_ff("Ah_phi_ff",0,0),
      Ah_rho_ff("Ah_rho_ff",0,0),
      Ah_cv_phi_ff("Ah_cv_phi_ff",0,0),
      Epar_em_phi_ff("Epar_em_phi_ff",0,0),
      Epar_em_rho_ff("Epar_em_rho_ff",0,0),
      //dpot_n0("dpot_n0",nnode),
#  endif
      dpot_ff("dpot_ff",0)
#else
      E_rho("E_rho",0,0),
      dpot("dpot",0)
#endif
    {
        nlr.use_namelist("sml_param");
        turb_efield = nlr.get<bool>("sml_turb_efield",true); 
#ifdef XGC1
        // Set pointers in Fortran to rho_ff arrays
        set_rho_ff_pointers(nrho, nnode, dpot_h.data(), dpot_ff_h.data(), pot.rho_ff_h.data(), E.rho_ff_h.data()
#  ifdef EXPLICIT_EM
          , Ah_h.data(), Ah.rho_ff_h.data(), dAh.rho_ff_h.data(), As_h.data(), As.rho_ff_h.data(), dAs.rho_ff_h.data(), Epar_em.rho_ff_h.data()
#  endif
        );
#else
        set_rho_pointers(nrho, nnode, dpot_h.data(), E.rho_h.data());
#endif
    }

#ifdef XGC1
    // Resizes allocations for large arrays that have a phi dimension (i.e used for electron push)
    void copy_phi_fields_to_device(){
        Kokkos::resize(E_phi_ff,nphi, nnode);
        Kokkos::deep_copy(E_phi_ff, E.phi_ff_h);
#  ifdef DELTAF_CONV 
        Kokkos::resize(ddpotdt_phi,nphi,nnode);
        Kokkos::deep_copy(ddpotdt_phi, ddpotdt.phi_ff_h);
        Kokkos::resize(E00_ff, nnode);
        Kokkos::deep_copy(E00_ff, E00_ff_h);
#  endif
#  ifdef EXPLICIT_EM
        Kokkos::resize(dAs_phi_ff,nphi,nnode);
        Kokkos::resize(dAh_phi_ff,nphi,nnode);
        Kokkos::resize(As_phi_ff,nphi,nnode);
        Kokkos::resize(Ah_phi_ff,nphi,nnode);
        Kokkos::resize(Ah_cv_phi_ff,nphi,nnode);
        Kokkos::resize(Epar_em_phi_ff,nphi,nnode);
        Kokkos::deep_copy(dAs_phi_ff, dAs.phi_ff_h);
        Kokkos::deep_copy(dAh_phi_ff, dAh.phi_ff_h);
        Kokkos::deep_copy(As_phi_ff, As.phi_ff_h);
        Kokkos::deep_copy(Ah_phi_ff, Ah.phi_ff_h);
        Kokkos::deep_copy(Ah_cv_phi_ff, Ah_cv.phi_ff_h);
        Kokkos::deep_copy(Epar_em_phi_ff, Epar_em.phi_ff_h);
#  endif
    }

    // Resizes allocations that have a phi dimension to 0, to make room for other data on the device
    void deallocate_phi_fields_on_device(){
        Kokkos::resize(E_phi_ff,0,0);
#  ifdef DELTAF_CONV
        Kokkos::resize(ddpotdt_phi,0,0);
        Kokkos::resize(E00_ff, 0);
#  endif
#  ifdef EXPLICIT_EM
        Kokkos::resize(dAs_phi_ff,0,0);
        Kokkos::resize(dAh_phi_ff,0,0);
        Kokkos::resize(As_phi_ff,0,0);
        Kokkos::resize(Ah_phi_ff,0,0);
        Kokkos::resize(Ah_cv_phi_ff,0,0);
        Kokkos::resize(Epar_em_phi_ff,0,0);
#  endif
    }

    // Resizes allocations for large arrays that have a rho dimension (i.e used for ions)
    void copy_rho_fields_to_device(){
        Kokkos::resize(E_rho_ff, nnode,nrho+1);
        Kokkos::deep_copy(E_rho_ff, E.rho_ff_h);
#  ifdef EXPLICIT_EM
        Kokkos::resize(dAs_rho_ff,nnode,nrho+1);
        Kokkos::resize(dAh_rho_ff,nnode,nrho+1);
        Kokkos::resize(As_rho_ff,nnode,nrho+1);
        Kokkos::resize(Ah_rho_ff,nnode,nrho+1);
        Kokkos::resize(Epar_em_rho_ff,nnode,nrho+1);
        Kokkos::deep_copy(dAs_rho_ff, dAs.rho_ff_h);
        Kokkos::deep_copy(dAh_rho_ff, dAh.rho_ff_h);
        Kokkos::deep_copy(As_rho_ff, As.rho_ff_h);
        Kokkos::deep_copy(Ah_rho_ff, Ah.rho_ff_h);
        Kokkos::deep_copy(Epar_em_rho_ff, Epar_em.rho_ff_h);
#  endif
    }

    // Resizes allocations that have a rho dimension to 0, to make room for other data on the device
    void deallocate_rho_fields_on_device(){
        Kokkos::resize(E_rho_ff,0,0);
#  ifdef EXPLICIT_EM
        Kokkos::resize(dAs_rho_ff,0,0);
        Kokkos::resize(dAh_rho_ff,0,0);
        Kokkos::resize(As_rho_ff,0,0);
        Kokkos::resize(Ah_rho_ff,0,0);
        Kokkos::resize(Epar_em_rho_ff,0,0);
#  endif
    }

#  ifdef EXPLICIT_EM
    /* Set of functions used to allocate/deallocate gpu memory for Ah during particle scatters */
    // Resizes allocation of Ah_cv
    void copy_Ah_cv_to_device(){
        Kokkos::resize(Ah_cv_phi_ff,nphi,nnode);
        Kokkos::deep_copy(Ah_cv_phi_ff, Ah_cv.phi_ff_h);
    }

    // Resizes Ah_cv to 0
    void deallocate_Ah_cv_on_device(){
        Kokkos::resize(Ah_cv_phi_ff,0,0);
    }

    void copy_Ah_phi_ff_to_device(){
        Kokkos::resize(Ah_phi_ff,nphi,nnode);
        Kokkos::deep_copy(Ah_phi_ff, Ah.phi_ff_h);
    }

    // Resizes Ah_phi_ff to 0
    void deallocate_Ah_phi_ff_on_device(){
        Kokkos::resize(Ah_phi_ff,0,0);
    }

    // Resizes allocation of Ah_rho_ff
    void copy_Ah_rho_ff_to_device(){
        Kokkos::resize(Ah_rho_ff,nnode,nrho+1);
        Kokkos::deep_copy(Ah_rho_ff, Ah.rho_ff_h);
    }

    // Resizes Ah_rho_ff to 0
    void deallocate_Ah_rho_ff_on_device(){
        Kokkos::resize(Ah_rho_ff,0,0);
    }
#  endif

    // Resizes allocation of dpot_ff and copies to device
    void copy_dpot_to_device(){
        Kokkos::resize(dpot_ff,nnode);
        Kokkos::deep_copy(dpot_ff, dpot_ff_h);
    }

    // Resizes dpot_ff to 0
    void deallocate_dpot_on_device(){
        Kokkos::resize(dpot_ff,0);
    }
#else
    // Resizes allocations for large arrays that have a rho dimension (i.e used for ions)
    void copy_rho_fields_to_device(){
        Kokkos::resize(E_rho, nnode,nrho+1);
        Kokkos::deep_copy(E_rho, E.rho_h);
    }

    // Resizes allocations that have a rho dimension to 0, to make room for other data on the device
    void deallocate_rho_fields_on_device(){
        Kokkos::resize(E_rho,0,0);
    }

    // Resizes allocation of dpot and copies to device
    void copy_dpot_to_device(){
        Kokkos::resize(dpot,nnode);
        Kokkos::deep_copy(dpot, dpot_h);
    }

    // Resizes dpot to 0
    void deallocate_dpot_on_device(){
        Kokkos::resize(dpot,0);
    }
#endif

};

#endif