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 "get_potential_grad.hpp"
#include "plane_field_gatherer.hpp"
#ifdef SONIC_GK
#include "get_sonic_fields.hpp"
#endif
#ifdef EXPLICIT_EM
#include "globals.hpp"
#endif

// Fortran calls
#ifdef XGC1
extern "C" void set_rho_ff_pointers(int nrho, int n, Field<VarType::Scalar,PhiInterpType::None>* pot0, Field<VarType::Scalar,PhiInterpType::None>* 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
      , Field<VarType::Scalar,PhiInterpType::None>* Ah, Field<VarType::Scalar,PhiInterpType::Planes>* Ah_rho_ff, Field<VarType::Vector,PhiInterpType::Planes>* dAh_rho_ff, Field<VarType::Scalar,PhiInterpType::None>* 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, Field<VarType::Scalar,PhiInterpType::None>* pot0, Field<VarType::Scalar,PhiInterpType::None>* dpot, Field<VarType::Vector2D,PhiInterpType::None>* E_rho
#  ifdef SONIC_GK
      , Field<VarType::Vector2D,PhiInterpType::None>* dEr_B2_rho
      , Field<VarType::Vector2D,PhiInterpType::None>* dEz_B2_rho
      , Field<VarType::Vector2D,PhiInterpType::None>* du2_E_rho
#  endif
      );
#endif

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

    // Dimensions
    int nnode;
    int nphi;

    // Gyromatrix info (to be moved soon)
    int nrho;
    double inv_drho;

    // Optimization options
    bool calculate_phi_ff_on_device;

    // Field variables
    bool turb_efield; 
    bool E00_efield;   

    // Host Views
    GridField<HostType,VarType::Vector2D,PIT_GLOBAL,TorType::OnePlane,KinType::DriftKin> E00_ff_h; 
    GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::DriftKin> ddpotdt; 

    //For EM
    GridField<HostType,VarType::Scalar,PhiInterpType::None,TorType::MultiplePlanes,KinType::DriftKin> Ah_h;
    GridField<HostType,vec2d_if_axisym<PIT_GLOBAL>(),PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> dAh;
    GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> Ah;
    GridField<HostType,VarType::Scalar,PhiInterpType::None,TorType::MultiplePlanes,KinType::DriftKin> As_h;
    GridField<HostType,vec2d_if_axisym<PIT_GLOBAL>(),PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> dAs;
    GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> As;
    GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::DriftKin> Ah_cv;
    GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> Epar_em;

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

    GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> pot;
    GridField<HostType,vec2d_if_axisym<PIT_GLOBAL>(),PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> E;
#ifdef XGC1
    GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::DriftKin> dpot_ff_h;
#else
#  ifdef SONIC_GK
    GridField<HostType,VarType::Vector2D,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> dEr_B2; 
    GridField<HostType,VarType::Vector2D,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> dEz_B2; 
    GridField<HostType,VarType::Vector2D,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> du2_E ; 
#  endif
#endif
    GridField<HostType,VarType::Scalar,PhiInterpType::None,TorType::MultiplePlanes,KinType::DriftKin> dpot_h;
    GridField<HostType,VarType::Scalar,PhiInterpType::None,TorType::OnePlane,KinType::DriftKin> 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, double inv_drho_in) : //Grid<HostTypeType>& grid) :
      /*nnode(grid.nnode),
      nphi(grid.nplanes),
      nrho(grid.nrho),*/
      nnode(nnode_in),
      nphi(nplanes_in),
      nrho(nrho_in),
      inv_drho(inv_drho_in),
      // Host Views
#ifdef XGC1
      pot("pot",nphi,nrho,nnode),
      E("E",nphi,nrho,nnode),
#  ifdef DELTAF_CONV 
      E00_ff_h("E00_ff_h",nnode),
      ddpotdt("ddpotdt",nphi,nrho,nnode),
#  endif
#  ifdef EXPLICIT_EM
      dAh("dAh",nphi,nrho,nnode),
      Ah("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("E",nrho, nnode),
      dpot_h("dpot_h",1,nnode),
#  ifdef SONIC_GK
      dEr_B2("dEr_B2",nrho, nnode),
      dEz_B2("dEz_B2",nrho, nnode),
      du2_E("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",true);
        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_mixed_variable)
            exit_XGC("\nError: sml_em_mixed_variable is explicitly set to .false. which is incompatible with EXPLICIT_EM\n");
#endif

        if(nlr.namelist_present("performance_param")){
            nlr.use_namelist("performance_param"); 
            calculate_phi_ff_on_device = nlr.get<bool>("calculate_phi_ff_on_device", false);
        }else{
            calculate_phi_ff_on_device = false;
        }

#ifdef EXPLICIT_EM
        if (em_control_variate){
            Ah_cv = GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::DriftKin>("Ah_cv",nphi,nrho,nnode);
        }

        if (em_mixed_variable){
            As = GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin>("As",nphi,nrho,nnode);
            dAs = GridField<HostType,vec2d_if_axisym<PIT_GLOBAL>(),PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin>("dAs",nphi,nrho,nnode);
            if (em_pullback_mode==4){
                Epar_em = GridField<HostType,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin>("Epar_em",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.data(), E.data()
#  ifdef EXPLICIT_EM
          , Ah_h.data(), Ah.data(), dAh.data(), As_h.data(), As.data(), dAs.data(), Epar_em.data()
#  endif
        );
#else
        set_rho_pointers(nrho, nnode, pot0_h.data(), dpot_h.data(), E.data()
#  ifdef SONIC_GK
          , dEr_B2.data()
          , dEz_B2.data()
          , du2_E.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_grid_field, const T& host_grid_field) const{
        device_grid_field = host_grid_field;
    }

    // Allocate device grid_field and copy data from host
    template<typename T1, typename T2>
    void copy_to_pack(T1& device_grid_field, const T2& host_grid_field) const{
        // Allocate device array
        device_grid_field = T1("grid_field",host_grid_field.nphi(), host_grid_field.nrhop1()-1, host_grid_field.nnode());
        // Deep copy from host to device
        Kokkos::deep_copy(device_grid_field.f, host_grid_field.f);
    }

    // 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 Simulation<DeviceType>& sml, const DomainDecomposition<DeviceType>& pol_decomp, const MagneticField<DeviceType>& magnetic_field, const Grid<DeviceType>& grid, bool near_field=false) const{
        // Re-initialize new gridpack - deallocates whatever Views were already allocated
        GridFieldPack<Device, PIT> gfpack(kintype_in, turb_efield, nrho, inv_drho);
#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){
                if(near_field){
                    // Gathering the "near field", defined as the field decomposition region that contains this rank's
                    // global domain decomposition
                    int local_pid = pol_decomp.field_decomp.find_domain_owner(pol_decomp.plane_index, grid.nplanes, pol_decomp.node_offset, grid.nnode);
                    gfpack.phi_offset = pol_decomp.field_decomp.all_first_plane(local_pid);
                    gfpack.node_offset = pol_decomp.field_decomp.all_first_node(local_pid);
                }else{
                    gfpack.phi_offset = pol_decomp.field_decomp.first_plane;
                    gfpack.node_offset = pol_decomp.field_decomp.first_node;
                }
            }
            PlaneFieldGatherer pfg(pol_decomp, grid, near_field);

#  ifdef DELTAF_CONV 
            if(ddpotdt.f.size()>0) pfg.gather_phi_ff_on_device(ddpotdt.f, gfpack.ddpotdt_phi_ff.f);
            copy_to_pack(gfpack.E00_ff, E00_ff_h); // No gather_phi needed due to axisymmetry
#  endif

            if(calculate_phi_ff_on_device){
                // Whether to ignore dpot
                bool ignore_poloidal_dpot = !(turb_efield);
                bool use_field00 = E00_efield;
                bool potential_is_requested = false; // Not used in push
                GridField<Device, VarType::Scalar,PIT, TorType::MultiplePlanes, KinType::DriftKin> unused_pot("unused_pot",0,0);
                pfg.calculate_phi_ff_on_device(sml,grid, magnetic_field, my_subview(dpot_h.f, 2), pot0_h, gfpack.E_phi_ff, unused_pot, potential_is_requested, use_field00, ignore_poloidal_dpot);
#  ifdef EXPLICIT_EM
                GridField<HostType,VarType::Scalar,PhiInterpType::None,TorType::OnePlane,KinType::DriftKin> pot0_em_h("pot0_em_h",0); // unused
                pfg.calculate_phi_ff_on_device(sml,grid, magnetic_field, my_subview(Ah_h.f, 2), pot0_em_h, gfpack.dAh_phi_ff, gfpack.Ah_phi_ff);
                if(As.f.size()>0) pfg.calculate_phi_ff_on_device(sml,grid, magnetic_field, my_subview(As_h.f, 2), pot0_em_h, gfpack.dAs_phi_ff, gfpack.As_phi_ff);
#  endif
            }else{
                pfg.gather_phi_ff_on_device(E.f, gfpack.E_phi_ff.f);
#  ifdef EXPLICIT_EM
                pfg.gather_phi_ff_on_device(dAh.f, gfpack.dAh_phi_ff.f);
                pfg.gather_phi_ff_on_device(Ah.f, gfpack.Ah_phi_ff.f);
                if(dAs.f.size()>0) pfg.gather_phi_ff_on_device(dAs.f, gfpack.dAs_phi_ff.f);
                if(As.f.size()>0) pfg.gather_phi_ff_on_device(As.f, gfpack.As_phi_ff.f);
#  endif
            }
#  ifdef EXPLICIT_EM
            if(Epar_em.f.size()>0) pfg.gather_phi_ff_on_device(Epar_em.f, gfpack.Epar_em_phi_ff.f);
#  endif
            GPTLstop("GAT_FIELD_INFO");
        }else{
            GPTLstart("Copy_rho_ff_fields_to_device");
            copy_to_pack(gfpack.E_rho_ff, E);
#  ifdef EXPLICIT_EM
            copy_to_pack(gfpack.dAs_rho_ff, dAs);
            copy_to_pack(gfpack.dAh_rho_ff, dAh);
            copy_to_pack(gfpack.As_rho_ff, As);
            copy_to_pack(gfpack.Ah_rho_ff, Ah);
            copy_to_pack(gfpack.Epar_em_rho_ff, Epar_em);
#  endif
            GPTLstop("Copy_rho_ff_fields_to_device");
        }
#else
        copy_to_pack(gfpack.E_rho, E);
#  ifdef SONIC_GK
        copy_to_pack(gfpack.dEr_B2_rho, dEr_B2);
        copy_to_pack(gfpack.dEz_B2_rho, dEz_B2);
        copy_to_pack(gfpack.du2_E_rho, du2_E);
#  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_as_Ah(KinType kintype_in, const DomainDecomposition<DeviceType>& pol_decomp, const Grid<DeviceType>& grid) const{
        GridFieldPack<Device, PIT> gfpack(kintype_in, turb_efield, nrho, inv_drho);
#ifdef EXPLICIT_EM
        PlaneFieldGatherer pfg(pol_decomp, grid);
        pfg.gather_phi_ff_on_device(Ah_cv.f, gfpack.Ah_phi_ff.f);
#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, nrho, inv_drho);
#ifdef EXPLICIT_EM
        if(kintype_in==DriftKin){
            PlaneFieldGatherer pfg(pol_decomp, grid);
            pfg.gather_phi_ff_on_device(Ah.f, gfpack.Ah_phi_ff.f);
        }else{
            copy_to_pack(gfpack.Ah_rho_ff, Ah);
        }
#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, nrho, inv_drho);
#ifdef XGC1
        copy_to_pack(gfpack.dpot_ff, dpot_ff_h);
#else
        const int PLANE_ZERO=0;
        copy_to_pack(gfpack.dpot, dpot_h.subfield(PLANE_ZERO));
#endif
        return gfpack;
    }

    void get_potential_grad(const Simulation<DeviceType>& sml, const Grid<DeviceType>& grid, const DomainDecomposition<DeviceType>& pol_decomp, const MagneticField<DeviceType>& magnetic_field, const Species<DeviceType>& species){
        GPTLstart("GET_POT_GRAD_EXCL_DESTR"); // Timer excludes the destructors which are non-negligible

        // Allocate temporary views
        GPTLstart("GET_POT_GRAD_ALLOC_VIEWS");
        int n_input_field_planes = dpot_h.nphi();
        bool gyroaverage_requested = true;
        GetPotentialGradTemp<DeviceType, HostType> tmp(sml, grid, pol_decomp, magnetic_field, n_input_field_planes, gyroaverage_requested, species.gyro_avg_matrices);
        GPTLstop("GET_POT_GRAD_ALLOC_VIEWS");

        /****************************************/
        // Calculation of the electric potential E
        /****************************************/
        GPTLstart("GET_POT_GRAD_PHI");

        auto gpg_field_args = GetPotGradFieldArgs<HostType, HostType, vec2d_if_axisym<PIT_GLOBAL>(), PIT_GLOBAL, TorType::OnePlane, KinType::GyroKin>
                                            (dpot_h.unmanaged(), !(turb_efield));

        // Set up E00
        if(E00_efield){
            gpg_field_args.field00 = Field00<DeviceType>(pot0_h, sml.grad_psitheta);
            gpg_field_args.field00.calculate_gradient(grid, tmp.grad_matrices);
        }
        // Write results of E00
        if(sml.reduced_deltaf && sml.electron_on) gpg_field_args.field00.set_field(grid, tmp.ff, E00_ff_h.f);

        // Set up EMPar
        if(sml.explicit_electromagnetic && (sml.em_mixed_variable && sml.em_pullback_mode==4)) gpg_field_args.request_para_em(Epar_em);

        // Request potential and/or gradient outputs
        // XGC1 algorithm uses potential but XGCa doesnt
        if(!sml.is_XGCa) gpg_field_args.request_potential(pot);
        gpg_field_args.request_gradient(E);

        // Get potentials
        get_field_grad(grid, pol_decomp, magnetic_field, gpg_field_args, tmp);

        GPTLstop("GET_POT_GRAD_PHI");

        if(sml.explicit_electromagnetic){
            GPTLstart("GET_POT_GRAD_APAR");

            /****************************************/
            // Calculation of the Hamiltonian vector potential A_h
            /****************************************/
            gpg_field_args = GetPotGradFieldArgs<HostType, HostType, vec2d_if_axisym<PIT_GLOBAL>(), PIT_GLOBAL, TorType::OnePlane, KinType::GyroKin>
                                        (Ah_h.unmanaged());
            gpg_field_args.request_potential(Ah);
            gpg_field_args.request_gradient(dAh);
            get_field_grad(grid, pol_decomp, magnetic_field, gpg_field_args, tmp);

            /****************************************/
            // Calculation of the symplectic vector potential A_s
            /****************************************/
            if (sml.em_mixed_variable){
                gpg_field_args = GetPotGradFieldArgs<HostType, HostType, vec2d_if_axisym<PIT_GLOBAL>(), PIT_GLOBAL, TorType::OnePlane, KinType::GyroKin>
                                            (As_h.unmanaged());
                gpg_field_args.request_potential(As);
                gpg_field_args.request_gradient(dAs);
                get_field_grad(grid, pol_decomp, magnetic_field, gpg_field_args, tmp);
            }

            /****************************************/
            // Calculation of the Hamiltonian vector potential A_h for the control-variate method
            /****************************************/
            if (sml.em_control_variate){
                TIMER("GET_POT_APAR_CV",
                    get_field_Ah_cv_ff(Ah_cv.data()) );
            }
            GPTLstop("GET_POT_GRAD_APAR");
        }

#ifdef SONIC_GK
        GPTLstart("GET_POT_GRAD_SONIC");
        View<double*,CLayout,HostType> Er_B2_h("Er_B2_h",grid.nnode); 
        View<double*,CLayout,HostType> Ez_B2_h("Ez_B2_h",grid.nnode); 
        View<double*,CLayout,HostType> u2_E_h("u2_E_h",grid.nnode); 

        TIMER("GET_POT_SONIC",
        get_sonic_fields(grid, E, Er_B2_h, Ez_B2_h, u2_E_h) );

        gpg_field_args = GetPotGradFieldArgs<HostType, HostType, vec2d_if_axisym<PIT_GLOBAL>(), PIT_GLOBAL, TorType::OnePlane, KinType::GyroKin>(Er_B2_h);
        gpg_field_args.request_gradient(dEr_B2);
        get_field_grad(grid, pol_decomp, magnetic_field, gpg_field_args, tmp);

        gpg_field_args = GetPotGradFieldArgs<HostType, HostType, vec2d_if_axisym<PIT_GLOBAL>(), PIT_GLOBAL, TorType::OnePlane, KinType::GyroKin>(Ez_B2_h);
        gpg_field_args.request_gradient(dEz_B2);
        get_field_grad(grid, pol_decomp, magnetic_field, gpg_field_args, tmp);

        gpg_field_args = GetPotGradFieldArgs<HostType, HostType, vec2d_if_axisym<PIT_GLOBAL>(), PIT_GLOBAL, TorType::OnePlane, KinType::GyroKin>(u2_E_h);
        gpg_field_args.request_gradient(du2_E);
        get_field_grad(grid, pol_decomp, magnetic_field, gpg_field_args, tmp);

        GPTLstop("GET_POT_GRAD_SONIC");
#endif

        GPTLstop("GET_POT_GRAD_EXCL_DESTR");
    }

    static MemoryPrediction estimate_memory_usage(NLReader::NamelistReader& nlr, const Grid<DeviceType> &grid){
        nlr.use_namelist("sml_param");
        int field_n_rho = nlr.get<int>("sml_grid_nrho", 6);

#ifdef EXPLICIT_EM
        int n_fields = 3 + 1; // E, dAs, dAh, + scalars As,Ah
#else
        int n_fields = 1;
#endif
        double host_field_GB_per_vertex = (53 + 8*grid.nplanes + 11*(field_n_rho + 1))*8*BYTES_TO_GB;
        double cpu_memory_usage = grid.nnode*host_field_GB_per_vertex;

        double device_field_GB_per_vertex = std::max(n_fields*8*grid.nplanes, n_fields*8*(field_n_rho+1))*8*BYTES_TO_GB;
        double gpu_memory_usage = grid.nnode*device_field_GB_per_vertex;

        MemoryPrediction memory_prediction{"Collisions", cpu_memory_usage, gpu_memory_usage};

        return memory_prediction;
    }
};

#endif