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 routines
 
extern "C" void ptb_3db_replace_b_grad_phi(double* tmp);
 
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, const Boundary& boundary, int n_input_potential_planes, bool gradient_requested, bool gyroavg_requested, const GyroAverageMatrices<DeviceType>& gyro_avg_matrices_in = GyroAverageMatrices<DeviceType>())
      : 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", gradient_requested ? nnode : 0),
        gradient_z_alt("gradient_r_alt", gradient_requested ? nnode : 0),
        potential(NoInit("potential"), nphi, nnode),
        gradient(NoInit("gradient"), ndim, nphi, gradient_requested ? nnode : 0),
        potential_rho(NoInit("potential_rho"), nrho+1, nphi, nnode),
        gradient_rho(NoInit("gradient_rho"), nrho+1, nphi, gradient_requested ? nnode : 0, ndim),
        gyro_avg_matrices(gyro_avg_matrices_in) // 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_requested ? gradient_rho.size() : potential_rho.size() );
            input_potential = View<double**,CLayout,Device>(NoInit("input_potential"), n_input_potential_planes, nnode);
        }
 
        GPTLstart("GET_POT_GRAD_MAT_SETUP");
        if(gradient_requested){
#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");
        if(gradient_requested && !sml.is_XGCa){
            gptx = GradParXTmp(grid, boundary, magnetic_field.bt_sign);
        }
        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 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;
    View<double**,CLayout, DeviceType> para_grad;
    bool ignore_poloidal_dpot;
    bool potential_is_requested;
    bool gradient_is_requested;
    bool para_grad_is_requested;
    Field00<DeviceType> field00;
 
    // 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),
       para_grad_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),
      para_grad_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_grad(const View<double**,CLayout, DeviceType>& para_grad_in){
        para_grad = para_grad_in;
        para_grad_is_requested = true;
    }
};
 
void get_field_Ah_cv_ff(const Grid<DeviceType>& grid, const DomainDecomposition<DeviceType>& pol_decomp, const GridField<HostType,VarType::Scalar,PhiInterpType::None,TorType::MultiplePlanes,KinType::DriftKin>& Ah_cv_h, const GridField<HostType,VarType::Scalar,PhiInterpType::Planes,TorType::OnePlane,KinType::DriftKin>& Ah_cv_ff);
 
template<class DeviceIn, class DeviceOut, VarType VT, PhiInterpType PIT, TorType TT, KinType KT>
void get_field_grad(const Grid<DeviceType>& grid, GetPotGradFieldArgs<DeviceIn, DeviceOut, VT, PIT, TT, KT>& args, GetPotentialGradTemp<DeviceType, DeviceOut>& tmp);
 
#endif