get_potential_grad.hpp

Go to the documentation of this file.

#ifndef GET_POTENTIAL_GRAD
#define GET_POTENTIAL_GRAD

#include "sml.hpp"
#include "grid.hpp"
#include "domain_decomposition.hpp"
#include "magnetic_field.hpp"
#include "my_mirror_view.hpp"
#include "gradparx2.hpp"
#include "field_following_coordinates.hpp"
#include "gradient_matrices.hpp"
#include "gyro_avg_mat.hpp"
#include "task_group.hpp"
#include "grid_field.hpp"
#include "grid_deriv.hpp"

// Fortran pointer retrievals

extern "C" rtype* get_psn_pbd0_2_iseg_loc();

extern "C" int get_psn_pbd0_2_nseg();

// Fortran routines
extern "C" void get_field_Ah_cv_ff(Field<VarType::Scalar,PIT_GLOBAL>* psn_Ah_cv_ff);
extern "C" void get_pot_epar_em_filter(double* tmp, double* E_para_em);


template<class Device>
class AlternatingStorage{
    View<double**,CLayout,Device> field;

    public:

    AlternatingStorage(const std::string& name, int nnode)
      : field(NoInit(name), 2, nnode)
    {}

    // Functions to access the "alternating" Views
    View<double*,CLayout, Device> left(int i_plane){
        int i_left_alt = i_plane%2;
        return my_subview(field,i_left_alt);
    }

    View<double*,CLayout, Device> right(int i_plane){
        int i_right_alt = (i_plane+1)%2;
        return my_subview(field,i_right_alt);
    }
};

template<class Device, class DeviceOut>
struct GetPotentialGradTemp{

    int nnode;
    int nrho;
    int nphi;
    int ndim;

    View<double**,CLayout,Device> input_potential;
    AlternatingStorage<Device> potential_alt;
    AlternatingStorage<Device> gradient_r_alt;
    AlternatingStorage<Device> gradient_z_alt;
    View<double**,CLayout,Device> potential;
    View<double***,CLayout, Device> gradient;
    View<double****,CLayout, Device> gradient_rho;
    View<double***,CLayout, Device> potential_rho;

    View<double*,CLayout,Device> scratch;

    // Additional objects that should go elsewhere eventually, but are here
    // because they are needed for get_pot_grad and can be reused for the different fields
    GyroAverageMatrices<DeviceType> gyro_avg_matrices;
    GradientMatrices<DeviceType> grad_matrices;
    FieldFollowingCoordinates ff;
    GradParXTmp gptx;

    GetPotentialGradTemp(const Simulation<DeviceType>& sml, const Grid<DeviceType>& grid, const DomainDecomposition<DeviceType>& pol_decomp, const MagneticField<DeviceType>& magnetic_field, int n_input_potential_planes, bool gyroavg_requested, const GyroAverageMatrices<HostType>& gyro_avg_matrices_in = GyroAverageMatrices<HostType>())
      : nnode(grid.nnode),
        nrho(gyroavg_requested ? gyro_avg_matrices_in.nrho : 0),
        nphi((PIT_GLOBAL==PhiInterpType::Planes) ? 2 : 1),
        ndim((vec2d_if_axisym<PIT_GLOBAL>()==VarType::Vector) ? 3 : 2),
        potential_alt("potential_alt", nnode),
        gradient_r_alt("gradient_r_alt", nnode),
        gradient_z_alt("gradient_r_alt", nnode),
        potential(NoInit("potential"), nphi, nnode),
        gradient(NoInit("gradient"), ndim, nphi, nnode),
        potential_rho(NoInit("potential_rho"), nrho+1, nphi, nnode),
        gradient_rho(NoInit("gradient_rho"), nrho+1, nphi, nnode, ndim),
        gyro_avg_matrices(gyro_avg_matrices_in.template mirror<DeviceType>()) // Send gyromatrices to GPU
    {
        // Use scratch for transpose on device if the fields can't be written directly
        if(!std::is_same<DeviceOut,Device>()){
            scratch = View<double*,CLayout,Device>(NoInit("scratch"), gradient_rho.size());
            input_potential = View<double**,CLayout,Device>(NoInit("input_potential"), n_input_potential_planes, nnode);
        }

        GPTLstart("GET_POT_GRAD_MAT_SETUP");
#ifdef NO_FORTRAN_MODULES
        grad_matrices = grid.gradient_matrices_h.template mirror<DeviceType>();
#else
        grad_matrices = GradientMatrices<DeviceType>(true); // Copy in fortran data
#endif
        GPTLstop("GET_POT_GRAD_MAT_SETUP");

        // This could be set up only once (if it doesnt take up too much device memory)
        GPTLstart("GET_POT_GRAD_FF_SETUP");
        if (!sml.is_XGCa) ff = FieldFollowingCoordinates(grid.half_plane_ff);
        GPTLstop("GET_POT_GRAD_FF_SETUP");

        // Set up for parallel gradient
        GPTLstart("GET_POT_GRAD_GPTX_SETUP");
#ifdef NO_FORTRAN_MODULES
        if(!sml.is_XGCa) gptx = GradParXTmp(grid, magnetic_field.bt_sign);
#else
        if(!sml.is_XGCa) gptx = GradParXTmp(grid, get_psn_pbd0_2_nseg(), magnetic_field.bt_sign, get_psn_pbd0_2_iseg_loc());
#endif
        GPTLstop("GET_POT_GRAD_GPTX_SETUP");
    }
};

template<class Device>
struct Field00{
    bool is_provided;
    bool discard_when_basis_is_one;

    View<Field<VarType::Scalar,PhiInterpType::None>*,CLayout, Device> pot_managed; // Use if device != host
    View<double*,CLayout, Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>> pot;
    View<double*,CLayout, Device> r;
    View<double*,CLayout, Device> z;

    Field00()
     : is_provided(false)
    {}

    template<class DeviceIn>
    Field00(const GridField<DeviceIn,VarType::Scalar,PhiInterpType::None,TorType::OnePlane,KinType::DriftKin>& field_in, bool discard_when_basis_is_one_in)
     : r(NoInit("r"), field_in.f.extent(0)),
       z(NoInit("z"), field_in.f.extent(0)),
       is_provided(true),
       discard_when_basis_is_one(discard_when_basis_is_one_in)
    {
        // Create a mirror if Device is different from the input DeviceIn
        pot_managed = my_mirror_view(field_in.f, Device());
        mirror_copy(pot_managed, field_in.f);
        pot = View<double*,CLayout,Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>>((double*)(pot_managed.data()), pot_managed.layout());
    }

    void calculate_gradient(const Grid<DeviceType>& grid, const GradientMatrices<DeviceType>& grad_matrices){
        grid_deriv(grad_matrices, grid,pot,r,z, discard_when_basis_is_one);

        auto r_tmp = r; // Use these to avoid using member function in kernel
        auto z_tmp = z; // Use these to avoid using member function in kernel

        // rh This is general now --->
        // We compute the electric field --> flip sign
        // E=-E
        Kokkos::parallel_for("get_pot_grad_flip_00", Kokkos::RangePolicy<ExSpace>(0,grid.nnode), KOKKOS_LAMBDA( const int i ){
            r_tmp(i)=-r_tmp(i);
            z_tmp(i)=-z_tmp(i);
        });
    }

    template<class DeviceOut, PhiInterpType PIT>
    void set_field(const Grid<DeviceType>& grid, const FieldFollowingCoordinates& ff, View<Field<VarType::Vector2D,PIT>*,CLayout,DeviceOut>& field00_ff_h){
        if (is_provided){
            // Copy to two planes if field following is needed
            int nphi = (PIT==PhiInterpType::Planes) ? 2 : 1;
            View<double**,CLayout, Device> r_tmp(NoInit("r_tmp"), nphi, r.size());
            View<double**,CLayout, Device> z_tmp(NoInit("z_tmp"), nphi, z.size());

            for(int iphi=0; iphi<nphi; iphi++){
                Kokkos::deep_copy(my_subview(r_tmp, iphi), r);
                Kokkos::deep_copy(my_subview(z_tmp, iphi), z);
            }

            if(PIT==PhiInterpType::Planes){
                ff.cnvt_grid_real2ff(grid,r_tmp);
                ff.cnvt_grid_real2ff(grid,z_tmp);
            }

            auto field00_ff = my_mirror_view(field00_ff_h, Device());
            for(int iphi=0; iphi<nphi; iphi++){
                Kokkos::parallel_for("set_field", Kokkos::RangePolicy<ExSpace>(0,field00_ff_h.extent(0)), KOKKOS_LAMBDA( const int i ){
#ifdef XGC1
                    field00_ff(i).V[iphi][0]=r_tmp(iphi,i);
                    field00_ff(i).V[iphi][1]=z_tmp(iphi,i);
#else
                    field00_ff(i).E[0]=r_tmp(iphi,i);
                    field00_ff(i).E[1]=z_tmp(iphi,i);
#endif
                });
            }
            Kokkos::fence();
            // Copy if not a mirror
            mirror_copy(field00_ff_h, field00_ff);
        }else{
            Kokkos::parallel_for("set_field", Kokkos::RangePolicy<typename DeviceOut::execution_space>(0,field00_ff_h.extent(0)), KOKKOS_LAMBDA( const int i ){
#ifdef XGC1
                field00_ff_h(i).V[0][0]=0.0;
                field00_ff_h(i).V[0][1]=0.0;
                field00_ff_h(i).V[1][0]=0.0;
                field00_ff_h(i).V[1][1]=0.0;
#else
                field00_ff_h(i).E[0]=0.0;
                field00_ff_h(i).E[1]=0.0;
#endif
            });
        }
    }
};

template<class Device, class DeviceOut>
struct EMParField{
    bool requested;

    View<double**,CLayout, Device> field;
    View<double***,CLayout, Device> field_rho;

    GridField<DeviceOut,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin> field_out;

    EMParField()
     : requested(false)
    {}

    void request(const GridField<DeviceOut,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin>& output_field){
        requested = true;
        field_out = output_field;
        field_rho = View<double***,CLayout, Device>(NoInit("E_para_em_rho"), field_out.f.extent(1), 2, field_out.f.extent(0));
        field = View<double**,CLayout, Device>(NoInit("E_para_em"), 2, field_out.f.extent(0));
    }

    void calculate(const Grid<DeviceType>& grid, const View<double**,CLayout, Device>& field_para){
        // Obtain E_para with the same filters that are applied in push_As
        // This filtered E_para must be used in the equation of motion
        // for dA_s/dt in case of pullback-mode 4.
        GPTLstart("GET_POT_EPAR_EM");
        // tmp_copy gets copied in and modified internally
        View<double**,CLayout, HostType> tmp_copy(NoInit("tmp"),2,grid.nnode);
        Kokkos::deep_copy(tmp_copy, field_para);

        // Output:
        auto field_h = my_mirror_view(field, HostType());

        get_pot_epar_em_filter(tmp_copy.data(), field_h.data());

        // Copy result of get_pot_epar_em_filter back to device
        mirror_copy(field, field_h);
        GPTLstop("GET_POT_EPAR_EM");
    }
};

template<class DeviceIn, class DeviceOut, VarType VT, PhiInterpType PIT, TorType TT, KinType KT>
struct GetPotGradFieldArgs{
    View<double**,CLayout,DeviceIn, Kokkos::MemoryTraits<Kokkos::Unmanaged>> input_potential;
    GridField<DeviceOut,VT,PIT,TT,KT> gradient;
    GridField<DeviceOut,VarType::Scalar,PIT,TT,KT> potential;
    bool ignore_poloidal_dpot;
    bool potential_is_requested;
    bool gradient_is_requested;
    Field00<DeviceType> field00;
    EMParField<DeviceType, DeviceOut> E_para_em;

    // Constructor if passed an unmanaged view
    GetPotGradFieldArgs(const View<double**,CLayout,DeviceIn, Kokkos::MemoryTraits<Kokkos::Unmanaged>>& input_potential_in, bool ignore_poloidal_dpot_in=false)
     : input_potential(input_potential_in),
       ignore_poloidal_dpot(ignore_poloidal_dpot_in),
       potential_is_requested(false),
       gradient_is_requested(false)
    {}

    // Constructor if passed a 1D managed view
    GetPotGradFieldArgs(const View<double*,CLayout,DeviceIn>& input_potential_in, bool ignore_poloidal_dpot_in=false)
     : input_potential(View<double**,CLayout,DeviceIn,Kokkos::MemoryTraits<Kokkos::Unmanaged>>((double*)(input_potential_in.data()), 1, input_potential_in.extent(0))),
      ignore_poloidal_dpot(ignore_poloidal_dpot_in),
      potential_is_requested(false),
      gradient_is_requested(false)
    {}

    void request_potential(const GridField<DeviceOut,VarType::Scalar,PIT,TT,KT>& potential_in){
        potential = potential_in;
        potential_is_requested = true;
    }

    void request_gradient(const GridField<DeviceOut,VT,PIT,TT,KT>& gradient_in){
        gradient = gradient_in;
        gradient_is_requested = true;
    }

    void request_para_em(const GridField<DeviceOut,VarType::Scalar,PIT_GLOBAL,TorType::OnePlane,KinType::GyroKin>& output_field){
        E_para_em.request(output_field);
    }
};

template<class DeviceIn, class DeviceOut, VarType VT, PhiInterpType PIT, TorType TT, KinType KT>
void get_field_grad(const Grid<DeviceType>& grid,
                    const DomainDecomposition<DeviceType>& pol_decomp,
                    const MagneticField<DeviceType>& magnetic_field,
                    GetPotGradFieldArgs<DeviceIn, DeviceOut, VT, PIT, TT, KT>& field_args,
                    GetPotentialGradTemp<DeviceType, DeviceOut>& tmp);

#endif