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 "grid_field_pack.hpp"
#include "timer_macro.hpp"
#include "plane_field_gatherer.hpp"

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

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

    // Field variables
    bool turb_efield; 
    bool E00_efield;   

    // Host Views
    Kokkos::View<Field<VarType::Vector2D,PhiInterpType::Planes>*,Kokkos::LayoutRight,HostType> E00_ff_h; 
    GridField<VarType::Scalar,PhiInterpType::Planes> ddpotdt; 

    //For EM
    Kokkos::View<double**,Kokkos::LayoutRight,HostType> Ah_h;
    GridField<VarType::Vector,PhiInterpType::Planes> dAh;
    GridField<VarType::Scalar,PhiInterpType::Planes> Ah;
    Kokkos::View<double**,Kokkos::LayoutRight,HostType> As_h;
    GridField<VarType::Vector,PhiInterpType::Planes> dAs;
    GridField<VarType::Scalar,PhiInterpType::Planes> As;
    GridField<VarType::Scalar,PhiInterpType::Planes> Ah_cv;
    GridField<VarType::Scalar,PhiInterpType::Planes> Epar_em;

    //Kokkos::View<double*,Kokkos::LayoutRight,HostType> dpot_n0_h;

#ifdef XGC1
    GridField<VarType::Scalar,PhiInterpType::Planes> pot;
    GridField<VarType::Vector,PhiInterpType::Planes> E;
    Kokkos::View<double**,Kokkos::LayoutRight,HostType> dpot_h;
    Kokkos::View<Field<VarType::Scalar,PhiInterpType::Planes>*,Kokkos::LayoutRight,HostType> dpot_ff_h;
#else
    GridField<VarType::Scalar,PhiInterpType::None> pot;
    GridField<VarType::Vector2D,PhiInterpType::None> E;
    Kokkos::View<double*,Kokkos::LayoutRight,HostType> dpot_h;
#  ifdef SONIC_GK
    GridField<VarType::Vector2D,PhiInterpType::None> dEr_B2; 
    GridField<VarType::Vector2D,PhiInterpType::None> dEz_B2; 
    GridField<VarType::Vector2D,PhiInterpType::None> du2_E ; 
#  endif
#endif
    Kokkos::View<double*,Kokkos::LayoutRight,HostType> pot0_h;

    // 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
#ifdef XGC1
      pot(nphi,nrho,nnode),
      E(nphi,nrho,nnode),
#  ifdef DELTAF_CONV 
      E00_ff_h("E00_ff_h",nnode),
      ddpotdt(nphi,nrho,nnode,GridFieldOpts::WithoutRhoFF),
#  endif
#  ifdef EXPLICIT_EM
      dAh(nphi,nrho,nnode),
      Ah(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),
#  ifdef SONIC_GK
      dEr_B2(nrho, nnode),
      dEz_B2(nrho, nnode),
      du2_E(nrho, nnode),
#  endif
#endif
      pot0_h("pot0_h",nnode)
    {
        nlr.use_namelist("sml_param");
        turb_efield = nlr.get<bool>("sml_turb_efield",true); 
        E00_efield = nlr.get<bool>("sml_00_efield",true);

#ifdef EXPLICIT_EM
        // Pass these (or sml) as arguments
        bool em_mixed_variable    = nlr.get<bool>("sml_em_mixed_variable",false);
        bool em_control_variate   = nlr.get<bool>("sml_em_control_variate",false);
        int em_pullback_mode     = nlr.get<int>("sml_em_pullback_mode",3);
        if (em_control_variate){
            Ah_cv = GridField<VarType::Scalar,PhiInterpType::Planes>(nphi,nrho,nnode,GridFieldOpts::WithoutRhoFF);
        }

        if (em_mixed_variable){
            As = GridField<VarType::Scalar,PhiInterpType::Planes>(nphi,nrho,nnode);
            dAs = GridField<VarType::Vector,PhiInterpType::Planes>(nphi,nrho,nnode);
            if (em_pullback_mode==4){
                Epar_em = GridField<VarType::Scalar,PhiInterpType::Planes>(nphi,nrho,nnode);
            }
        }
#endif

        // Set pointers in Fortran to rho arrays
#ifdef XGC1
        set_rho_ff_pointers(nrho, nnode, pot0_h.data(), 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, pot0_h.data(), dpot_h.data(), E.rho_h.data()
#  ifdef SONIC_GK
          , dEr_B2.rho_h.data()
          , dEz_B2.rho_h.data()
          , du2_E.rho_h.data()
#  endif
        );
#endif
    }


    /**************************************************/
    /************  Copy to GridFieldPack  *************/
    /**************************************************/

    // Should these functions be in GridFieldPack instead? - ALS


    // Shallow copy if DeviceType and HostType are the same
    template<typename T>
    void copy_to_pack(T& device_view, const T& host_view) const{
        device_view = host_view;
    }

    // Allocate device view and copy data from host
    template<typename T1, typename T2>
    void copy_to_pack(T1& device_view, const T2& host_view) const{
        // Allocate device array
        device_view = T1("gfpack_view",host_view.layout());
        // Deep copy from host to device
        Kokkos::deep_copy(device_view, host_view);
    }

    // Resizes allocations for large arrays that have a phi dimension (i.e used for electron push)
    template<PhiInterpType PIT>
    GridFieldPack<Device, PIT> copy_push_fields_to_device(KinType kintype_in, const DomainDecomposition<DeviceType>& pol_decomp, const Grid<DeviceType>& grid) const{
        // Re-initialize new gridpack - deallocates whatever Views were already allocated
        GridFieldPack<Device, PIT> gfpack(kintype_in, turb_efield);
#ifdef XGC1
        // Copy field data to device (or shallow copy if no device)
        if(kintype_in==DriftKin){
            GPTLstart("GAT_FIELD_INFO");
            if(pol_decomp.decompose_fields){
                gfpack.phi_offset = pol_decomp.field_decomp.first_plane;
                gfpack.node_offset = pol_decomp.field_decomp.first_node;
//printf("\ngfpack.phi_offset: %d, gfpack.node_offset: %d\n", gfpack.phi_offset, gfpack.node_offset);
            }
            PlaneFieldGatherer pfg(pol_decomp, grid);
//checkpoint("\n\n");
//checkpoint("Starting GATHER");
//checkpoint("\n\n");
            pfg.gather_phi_ff_on_device(E.rho_ff_h, gfpack.E_phi_ff);
//Kokkos::fence();
//checkpoint("\n\n");
//checkpoint("Finished GATHER");
//checkpoint("\n\n");
#  ifdef DELTAF_CONV 
            if(ddpotdt.ff_h.size()>0) pfg.gather_phi_ff_on_device(ddpotdt.ff_h, gfpack.ddpotdt_phi_ff);
            copy_to_pack(gfpack.E00_ff, E00_ff_h); // No gather_phi needed due to axisymmetry
#  endif
#  ifdef EXPLICIT_EM
            pfg.gather_phi_ff_on_device(dAh.rho_ff_h, gfpack.dAh_phi_ff);
            pfg.gather_phi_ff_on_device(Ah.rho_ff_h, gfpack.Ah_phi_ff);
            if(dAs.rho_ff_h.size()>0) pfg.gather_phi_ff_on_device(dAs.rho_ff_h, gfpack.dAs_phi_ff);
            if(As.rho_ff_h.size()>0) pfg.gather_phi_ff_on_device(As.rho_ff_h, gfpack.As_phi_ff);
            if(Ah_cv.rho_ff_h.size()>0) pfg.gather_phi_ff_on_device(Ah_cv.rho_ff_h, gfpack.Ah_cv_phi_ff);
            if(Epar_em.rho_ff_h.size()>0) pfg.gather_phi_ff_on_device(Epar_em.rho_ff_h, gfpack.Epar_em_phi_ff);
#  endif
            GPTLstop("GAT_FIELD_INFO");
        }else{
            GPTLstart("Copy_rho_ff_fields_to_device");
            copy_to_pack(gfpack.E_rho_ff, E.rho_ff_h);
#  ifdef EXPLICIT_EM
            copy_to_pack(gfpack.dAs_rho_ff, dAs.rho_ff_h);
            copy_to_pack(gfpack.dAh_rho_ff, dAh.rho_ff_h);
            copy_to_pack(gfpack.As_rho_ff, As.rho_ff_h);
            copy_to_pack(gfpack.Ah_rho_ff, Ah.rho_ff_h);
            copy_to_pack(gfpack.Epar_em_rho_ff, Epar_em.rho_ff_h);
#  endif
            GPTLstop("Copy_rho_ff_fields_to_device");
        }
#else
        copy_to_pack(gfpack.E_rho, E.rho_h);
#  ifdef SONIC_GK
        copy_to_pack(gfpack.dEr_B2_rho, dEr_B2.rho_h);
        copy_to_pack(gfpack.dEz_B2_rho, dEz_B2.rho_h);
        copy_to_pack(gfpack.du2_E_rho, du2_E.rho_h);
#  endif
#endif
        return gfpack;
    }

    /* Set of functions used to allocate/deallocate gpu memory for Ah during particle scatters */
    // Resizes allocation of Ah_cv
    template<PhiInterpType PIT>
    GridFieldPack<Device, PIT> copy_Ah_cv_to_device(KinType kintype_in, const DomainDecomposition<DeviceType>& pol_decomp, const Grid<DeviceType>& grid) const{
        GridFieldPack<Device, PIT> gfpack(kintype_in, turb_efield);
#ifdef EXPLICIT_EM
        PlaneFieldGatherer pfg(pol_decomp, grid);
        pfg.gather_phi_ff_on_device(Ah_cv.rho_ff_h, gfpack.Ah_cv_phi_ff);
#endif
        return gfpack;
    }

    template<PhiInterpType PIT>
    GridFieldPack<Device, PIT> copy_Ah_ff_to_device(KinType kintype_in, const DomainDecomposition<DeviceType>& pol_decomp, const Grid<DeviceType>& grid) const{
        GridFieldPack<Device, PIT> gfpack(kintype_in, turb_efield);
#ifdef EXPLICIT_EM
        if(kintype_in==DriftKin){
            PlaneFieldGatherer pfg(pol_decomp, grid);
            pfg.gather_phi_ff_on_device(Ah.rho_ff_h, gfpack.Ah_phi_ff);
        }else{
            copy_to_pack(gfpack.Ah_rho_ff, Ah.rho_ff_h);
        }
#endif
        return gfpack;
    }


    // Resizes allocation of dpot_ff and copies to device
    template<PhiInterpType PIT>
    GridFieldPack<Device, PIT> copy_dpot_to_device(KinType kintype_in) const{
        GridFieldPack<Device, PIT> gfpack(kintype_in, turb_efield);
#ifdef XGC1
        copy_to_pack(gfpack.dpot_ff, dpot_ff_h);
#else
        copy_to_pack(gfpack.dpot, dpot_h);
#endif
        return gfpack;
    }

};

#endif