field_following_coordinates.hpp

Go to the documentation of this file.

#ifndef FIELD_FOLLOWING_COORDINATES_HPP
#define FIELD_FOLLOWING_COORDINATES_HPP

#include "grid.hpp"

extern "C" int* get_psn_ff_hdp_tr_loc();
extern "C" double* get_psn_ff_hdp_p_loc();
extern "C" double* get_grid_node_vol_ff_loc();

class FieldFollowingCoordinates{

    Kokkos::View<int**,Kokkos::LayoutRight,DeviceType> tr;
    Kokkos::View<double***,Kokkos::LayoutRight,DeviceType> p;

    // View for temporarily storing result data for in-place conversions
    Kokkos::View<double**,Kokkos::LayoutRight,DeviceType> result;

    public:

    FieldFollowingCoordinates(){}

    FieldFollowingCoordinates(const Grid<DeviceType>& grid)
        : tr(Kokkos::ViewAllocateWithoutInitializing("tr"),2,grid.nnode),
          p(Kokkos::ViewAllocateWithoutInitializing("p"),2,grid.nnode,3),
          result(Kokkos::ViewAllocateWithoutInitializing("result"),2,grid.nnode)
    {
        // Copy grid conversion info to device
        array_deep_copy(tr,get_psn_ff_hdp_tr_loc());
        array_deep_copy(p,get_psn_ff_hdp_p_loc());
    }

    /* Conversion from field-following coordinates to toroidal coordinates
     * @param[in] grid is the Grid object needed here to map elements to nodes
     * @param[in] input is the array one wants to convert
     * @param[out] output is the resulting converted array
     * */
    void cnvt_grid_ff2real(const Grid<DeviceType>& grid, const Kokkos::View<double**,Kokkos::LayoutRight,DeviceType>& input, const Kokkos::View<double**,Kokkos::LayoutRight,DeviceType>& output){
        // Check the input and output views are the correct size
        if(input.extent(0) != 2 || input.extent(1) != grid.nnode ||
           output.extent(0)!= 2 || output.extent(1)!= grid.nnode  ) exit_XGC("\nError in cnvt_grid_ff2real: Unexpected View sizes. \n");

        // Initialize output to zero
        Kokkos::deep_copy(output, 0.0);

        // Initialize vol to zero
        Kokkos::View<double**,Kokkos::LayoutRight,DeviceType> vol("vol",input.layout());

        // Copy grid%node_vol_ff from fortran
        Kokkos::View<double**,Kokkos::LayoutRight,DeviceType> node_vol_ff(Kokkos::ViewAllocateWithoutInitializing("node_vol_ff"),input.layout());
        array_deep_copy(node_vol_ff, get_grid_node_vol_ff_loc());

        // Make a (SHALLOW) copy so the lambda doesn't have to implicitly capture *this
        FieldFollowingCoordinates ff = *this;

        // for all node point of ff
        Kokkos::parallel_for("cnvt_grid_ff2real", Kokkos::RangePolicy<ExSpace>(0,grid.nnode), KOKKOS_LAMBDA( const int i ){
            for (int dir=0;dir<2;dir++){
                for (int j=0;j<3;j++){
                    // node and weight
                    int nd=grid.nodes(ff.tr(dir,i)-1).vertex[j] - 1; // 0-index
                    double wt=ff.p(dir,i,j)*node_vol_ff(dir,nd);

                    // accumulate values
                    Kokkos::atomic_add(&(output(dir,nd)), wt*input(dir,i));
                    Kokkos::atomic_add(&(vol(dir,nd)), wt);
                }
            }
        });

        // Normalize to vol
        Kokkos::parallel_for("cnvt_grid_ff2real_norm", Kokkos::RangePolicy<ExSpace>(0,grid.nnode), KOKKOS_LAMBDA( const int i ){
            for (int dir=0;dir<2;dir++){
                if(vol(dir, i) == 0.0){
                    output(dir, i) = 0.0;
                }else{
                    output(dir, i) /= vol(dir, i);
                }
            }   
        });

        // fill empty area --- replace nan values
        Kokkos::parallel_for("cnvt_grid_ff2real_fillempty", Kokkos::RangePolicy<ExSpace>(0,grid.nnode), KOKKOS_LAMBDA( const int i ){
            for (int dir=0; dir<2; dir++){
                if(vol(dir,i)==0.0){
                    int cdir = 1-dir;
                    output(dir,i)=0.0;
                    for (int j=0;j<3;j++){
                        int nd=grid.nodes(ff.tr(cdir,i)-1).vertex[j] - 1; // 0-index
                        double wt=ff.p(cdir,i,j);

                        // obtain interpolated values
                        output(dir,i) += wt*input(dir,nd);
                    }
                }
            }
        });
        Kokkos::fence();
    }

    /* Conversion from real coordinates to field-following coordinates
     * @param[in] grid is the Grid object needed here to map elements to nodes
     * @param[in] input is the array one wants to convert
     * @param[out] output is the resulting converted array
     * */    
    void cnvt_grid_real2ff(const Grid<DeviceType>& grid, const Kokkos::View<double**,Kokkos::LayoutRight,DeviceType>& input, const Kokkos::View<double**,Kokkos::LayoutRight,DeviceType>& output){
        // Check the input and output views are the correct size
        if(input.extent(0) != 2 || input.extent(1) != grid.nnode ||
           output.extent(0)!= 2 || output.extent(1)!= grid.nnode  ) exit_XGC("\nError in cnvt_grid_real2ff: Unexpected View sizes. \n");

        // Initialize output to zero
        Kokkos::deep_copy(output, 0.0);


        // Make a (SHALLOW) copy so the lambda doesn't have to implicitly capture *this
        FieldFollowingCoordinates ff = *this;

        Kokkos::parallel_for("cnvt_grid_real2ff", Kokkos::RangePolicy<ExSpace>(0,grid.nnode), KOKKOS_LAMBDA( const int i ){
            for (int dir=0;dir<2;dir++){
                for (int j=0;j<3;j++){
                    int nd=grid.nodes(ff.tr(dir,i)-1).vertex[j] - 1; // 0-index

                    double wt=ff.p(dir,i,j);

                    // obtain interpolated values
                    output(dir,i) += wt*input(dir,nd);
                }
            }
        });
        Kokkos::fence();
    }

    /* In-place conversion from real coordinates to field-following coordinates, using a buffer array to store the result temporarily
     * before copying back to the input array
     * @param[in] grid is the Grid object needed here to map elements to nodes
     * @param[inout] input is the array one wants to convert, and the resulting converted array
     * */
    void cnvt_grid_real2ff(const Grid<DeviceType>& grid, const Kokkos::View<double**,Kokkos::LayoutRight,DeviceType>& input){
        cnvt_grid_real2ff(grid, input, result);

        // Copy back to original array
        Kokkos::deep_copy(input, result);
    }
};

#endif