field_following_coordinates.hpp

Go to the documentation of this file.

#ifndef FIELD_FOLLOWING_COORDINATES_HPP
#define FIELD_FOLLOWING_COORDINATES_HPP

#include "my_subview.hpp"
#include "grid.hpp"
#include "grid_field.hpp"


class FieldFollowingCoordinates{
    View<int**,CLayout,DeviceType> tr;
    View<double***,CLayout,DeviceType> p;

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

    public:

    FieldFollowingCoordinates(){}

    FieldFollowingCoordinates(const Projection<HostType>& projection)
        : tr(NoInit("tr"),projection.tr.layout()),
          p(NoInit("p"),projection.p.layout()),
          result(NoInit("result"),tr.layout())
    {
        // Copy grid conversion info to device
        Kokkos::deep_copy(tr, projection.tr);
        Kokkos::deep_copy(p, projection.p);
    }

    /* 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. It is at the midplane so must be size midplane.nnode
     * @param[out] output is the resulting converted array. It is at the edge planes so must be size max(rplane.nnode, lplane.nnode)
     * */
    void cnvt_grid_ff2real(const Grid<DeviceType>& grid, const View<double**,CLayout,DeviceType>& input, const View<double**,Kokkos::LayoutRight,DeviceType>& output) const{
        // Check the input and output views are the correct size
        if(input.extent(0) != 2 || input.extent(1) != grid.midplane.nnode ||
           output.extent(0)!= 2 || output.extent(1)!= std::max(grid.rplane.nnode, grid.lplane.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
        View<double**,CLayout,DeviceType> vol("vol",input.layout());

        // Copy grid%node_vol_ff from fortran
        View<double**,CLayout,DeviceType> node_vol_ff(NoInit("node_vol_ff"),input.layout());
        // For stellarators, Dir 0 uses left plane node volume, Dir 1 uses right plane node volume
        Kokkos::deep_copy(my_subview(node_vol_ff,0), my_subview(grid.volumes_and_areas.node_vol_ff_l_h,0));
        Kokkos::deep_copy(my_subview(node_vol_ff,1), my_subview(grid.volumes_and_areas.node_vol_ff_r_h,1));

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

        // for all node point of ff
        constexpr int LEFTWARD=0;
        Kokkos::parallel_for("cnvt_grid_ff2real", Kokkos::RangePolicy<ExSpace>(0,grid.midplane.nnode), KOKKOS_LAMBDA( const int i ){
            for (int j=0;j<3;j++){
                // node and weight
                int nd=grid.lplane.get_node_index(ff.tr(LEFTWARD,i), j);
                double wt=ff.p(LEFTWARD,i,j)*node_vol_ff(LEFTWARD,nd);

                // accumulate values
                Kokkos::atomic_add(&(output(LEFTWARD,nd)), wt*input(LEFTWARD,i));
                Kokkos::atomic_add(&(vol(LEFTWARD,nd)), wt);
            }
        });
        constexpr int RIGHTWARD=1;
        Kokkos::parallel_for("cnvt_grid_ff2real", Kokkos::RangePolicy<ExSpace>(0,grid.midplane.nnode), KOKKOS_LAMBDA( const int i ){
            for (int j=0;j<3;j++){
                // node and weight
                int nd=grid.rplane.get_node_index(ff.tr(RIGHTWARD,i), j);
                double wt=ff.p(RIGHTWARD,i,j)*node_vol_ff(RIGHTWARD,nd);

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

        // Normalize to vol
        Kokkos::parallel_for("cnvt_grid_ff2real_norm", Kokkos::RangePolicy<ExSpace>(0,output.extent(1)), 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);
                }
            }   
        });
#ifndef STELLARATOR
        // I believe this logic assumes axisymmetric grid:
        // - ff.tr is size midplane.nnode, but this loop is over edge plane vertices.
        // - ff.tr(RIGHTWARD) cannot be interpreted on lplane.
        // Don't use for stellarators yet - ALS

        // fill empty area --- replace nan values
        Kokkos::parallel_for("cnvt_grid_ff2real_fillempty", Kokkos::RangePolicy<ExSpace>(0,output.extent(1)), KOKKOS_LAMBDA( const int i ){
            if(vol(LEFTWARD,i)==0.0){
                output(LEFTWARD,i)=0.0;
                for (int j=0;j<3;j++){
                    int nd=grid.lplane.get_node_index(ff.tr(RIGHTWARD,i), j);
                    double wt=ff.p(RIGHTWARD,i,j);

                    // obtain interpolated values
                    output(LEFTWARD,i) += wt*input(LEFTWARD,nd);
                }
            }
        });
        Kokkos::parallel_for("cnvt_grid_ff2real_fillempty", Kokkos::RangePolicy<ExSpace>(0,output.extent(1)), KOKKOS_LAMBDA( const int i ){
            if(vol(RIGHTWARD,i)==0.0){
                output(RIGHTWARD,i)=0.0;
                for (int j=0;j<3;j++){
                    int nd=grid.rplane.get_node_index(ff.tr(LEFTWARD,i), j);
                    double wt=ff.p(LEFTWARD,i,j);
                    
                    // obtain interpolated values
                    output(RIGHTWARD,i) += wt*input(RIGHTWARD,nd);
                }
            }
        });
#endif
        Kokkos::fence();
    }

    /* Conversion from field-following coordinates to real coordinates when views are on host
     * User also provides a temporary device allocation
     * @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
     * @param[in] tmp is an array with the same size as input/output used for temporary storage
     * */
    void cnvt_grid_ff2real(const Grid<DeviceType>& grid, const View<double**,CLayout,HostType>& input, const View<double**,CLayout,HostType>& output, const View<double**,CLayout,DeviceType>& tmp) const{
        // Copy input to tmp
        Kokkos::deep_copy(tmp, input);

        // Convert
        cnvt_grid_ff2real(grid, tmp, result);

        // Copy to output
        Kokkos::deep_copy(output, result);
    }

#ifdef USE_GPU
    /* Conversion from field-following coordinates to real coordinates when input view is on host but output view is on device
     * User also provides a temporary device allocation
     * @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 View<double**,CLayout,HostType>& input, const View<double**,CLayout,DeviceType>& output) const{
        // Copy input to output
        Kokkos::deep_copy(output, input);

        // Convert
        cnvt_grid_ff2real(grid, output, result);

        // Copy to output
        Kokkos::deep_copy(output, result);
    }
#endif

    /* 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. The input values are at the two edge planes. If the two planes have a different
     *            number of vertices, it must be of size max(lplane.nnode, rplane.nnode)
     * @param[out] output is the resulting converted array. It corresponds to the midplane so it must be size midplane.nnode
     * */    
    void cnvt_grid_real2ff(const Grid<DeviceType>& grid, const View<double**,CLayout,DeviceType>& input, const View<double**,Kokkos::LayoutRight,DeviceType>& output) const{
        // Check the input and output views are the correct size
        if(input.extent(0) != 2 || input.extent(1) != std::max(grid.rplane.nnode, grid.lplane.nnode) ||
           output.extent(0)!= 2 || output.extent(1)!= grid.midplane.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;

        constexpr int LEFTWARD=0;
        Kokkos::parallel_for("cnvt_grid_real2ff", Kokkos::RangePolicy<ExSpace>(0,grid.midplane.nnode), KOKKOS_LAMBDA( const int i ){
            for (int j=0;j<3;j++){
                int nd=grid.lplane.get_node_index(ff.tr(LEFTWARD,i), j);

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

                // obtain interpolated values
                output(LEFTWARD,i) += wt*input(LEFTWARD,nd);
            }
        });
        constexpr int RIGHTWARD=1;
        Kokkos::parallel_for("cnvt_grid_real2ff", Kokkos::RangePolicy<ExSpace>(0,grid.midplane.nnode), KOKKOS_LAMBDA( const int i ){
            for (int j=0;j<3;j++){
                int nd=grid.rplane.get_node_index(ff.tr(RIGHTWARD,i), j); 
       
                double wt=ff.p(RIGHTWARD,i,j);
    
                // obtain interpolated values
                output(RIGHTWARD,i) += wt*input(RIGHTWARD,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 View<double**,CLayout,DeviceType>& input) const{
        cnvt_grid_real2ff(grid, input, result);

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

    /* Conversion from real coordinates to field-following coordinates for vectors
     * @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 View<double***,CLayout,DeviceType>& input, const View<double***,Kokkos::LayoutRight,DeviceType>& output) const{
        for(int i = 0; i<input.extent(0); i++){
            auto input_i = my_subview(input, i);
            auto output_i = my_subview(output, i);

            cnvt_grid_real2ff(grid, input_i, output_i);
        }
    }

    /* In-place conversion from real coordinates to field-following coordinates for vectors, 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 View<double***,CLayout,DeviceType>& input)const{
        for(int i = 0; i<input.extent(0); i++){
            auto input_i = my_subview(input, i);
            cnvt_grid_real2ff(grid, input_i);
        }
    }


    /* Conversion from real coordinates to field-following coordinates from a GridField with PhiInterpType::None to a GridField with PhiInterpType::Planes.
     * This implementation could be improved. It uses the existing algorithm which requires the input to have nphi=2, so there is a lot of superfluous copying etc.  */
    void cnvt_grid_real2ff(const Grid<DeviceType>& grid, const ScalarGridField& field_in,
                                                         const GridField<DeviceType, VarType::Scalar,PhiInterpType::Planes, TorType::OnePlane, KinType::DriftKin>& field_out) const{

        // Field-following algorithm requires argument to be nphi=2, and inode as the contiguous index
        int nphi = 2;
        Kokkos::View<double**,Kokkos::LayoutRight,DeviceType> dest_view(NoInit("dest_view"), nphi, field_in.nnode());

        // Since the input is a ScalarGridField (nphi=1), copy the field into both the iphi=0 and iphi=1 indices
        field_in.copy_to_double_view(my_subview(dest_view, 0)); // iphi=0;
        field_in.copy_to_double_view(my_subview(dest_view, 1)); // iphi=1;

        // Convert to field-following
        cnvt_grid_real2ff(grid, dest_view, result);

        // Now transpose and copy into the output argument:
        field_out.transpose_copy_from_double_view(result);
    }

};

#endif