grid_setup.hpp

Go to the documentation of this file.

#ifndef GRID_SETUP_HPP
#define GRID_SETUP_HPP

#include "globals.hpp"
#include "grid.hpp"

extern "C" void my_qsort(int left,int right,int* idx,double* psi_surf2_tmp,int npsi_surf2);

extern "C" void init_guess_xtable_fort( int* nguess_2, double* grid_guess_max, double* grid_guess_min, double* grid_inv_guess_d, int* ntriangle, int* nnodes, const RZPair* grid_x,
                            Vertex* grid_nd, int* guess_xtable, int* guess_count, int* guess_list_len_in);


extern "C" void init_guess_list_fort( int* nguess_2, double* grid_guess_max, double* grid_guess_min, double* grid_inv_guess_d, int* ntriangle, int* nnodes, const RZPair* grid_x,
                                     Vertex* grid_nd, int* guess_list, int* guess_xtable, int* guess_count, int* guess_list_len_in);

// Provides setup and checks for an analytic example of a grid

namespace{

// Mapping matrix init
void setup_mapping(const View<Vertex*, CLayout, HostType> nd, const View<RZPair*, CLayout, HostType> x, View<VertexMap*, CLayout, HostType>& mapping){
    Kokkos::parallel_for("setup_mapping", Kokkos::RangePolicy<HostExSpace>(0,mapping.extent(0)), KOKKOS_LAMBDA( const int i ){
       double dx1r=x[nd[i].vertex[0]-1].r - x[nd[i].vertex[2]-1].r;
       double dx1z=x[nd[i].vertex[0]-1].z - x[nd[i].vertex[2]-1].z;
       double dx2r=x[nd[i].vertex[1]-1].r - x[nd[i].vertex[2]-1].r;
       double dx2z=x[nd[i].vertex[1]-1].z - x[nd[i].vertex[2]-1].z;

       double det = 1.0/( dx1r*dx2z - dx2r*dx1z);

       mapping[i].vertex[0].r = dx2z*det;
       mapping[i].vertex[1].r =-dx2r*det;
       mapping[i].vertex[0].z =-dx1z*det;
       mapping[i].vertex[1].z = dx1r*det;


       mapping[i].vertex[2].r = x[nd[i].vertex[2]-1].r;
       mapping[i].vertex[2].z = x[nd[i].vertex[2]-1].z;
    });
}

// These are the regions designated in the .node file
enum FileRegion{
    Inside=0,
    OnWall
};

// These are the regions used in XGC
enum Region{
    RegionOne=1,
    RegionTwo,
    PrivateRegion,
    Wall=100 // temporary, since 100 is the value used elsewhere
};

void set_psi(){
    //for(int i = 0; i<psi.size(); i++){
    //grid%psi(i)=psi_interpol(grid%x(1,i),grid%x(2,i),0,0);
}

template<class Device, class GridDevice>
void set_rgn(const MagneticField<Device>& bfield, const View<RZPair*,CLayout,GridDevice>& gx_h,
             const View<double*,CLayout,GridDevice>& psi_h, View<int*, CLayout, GridDevice>& rgn_h){

    int nnode = gx_h.extent(0);
    View<RZPair*,CLayout,Device> gx("gx",nnode);
    View<double*,CLayout,Device> psi("psi",nnode);
    View<int*,CLayout,Device> rgn("rgn",nnode);
    Kokkos::deep_copy(gx, gx_h);
    Kokkos::deep_copy(psi, psi_h);
    Kokkos::deep_copy(rgn, rgn_h);

    Kokkos::parallel_for("set_rgn", Kokkos::RangePolicy<typename Device::execution_space>(0,rgn.extent(0)), KOKKOS_LAMBDA( const int i ){
        if(rgn(i)==Inside){
            if(bfield.equil.is_in_region_1_or_2(gx(i).r,gx(i).z,psi(i))){
                if(psi(i) > bfield.equil.xpt_psi - bfield.equil.epsil_psi){
                    rgn(i)=RegionTwo;
                }else{
                    rgn(i)=RegionOne;
                }
            }else{
                rgn(i)=PrivateRegion;
            }
        } else {
            rgn(i)=Wall;
        }
    });
    Kokkos::fence();

    Kokkos::deep_copy(gx_h, gx);
    Kokkos::deep_copy(psi_h, psi);
    Kokkos::deep_copy(rgn_h, rgn);
}

// Counts number of wall nodes
int get_n_wall_nodes(const View<int*, HostType>& rgn){
    int count=0;
    for(int i = 0; i<rgn.size(); i++){
        if(rgn[i]==Wall) count++;
    }

    if (count == 0){
      // No wall nodes found --> print warning and return
      printf("\ninit_wall: No wall nodes found!\n");
    }

    return count;
}

void setup_wall(const View<int*, CLayout, HostType>& rgn, View<int*, CLayout, HostType>& wall_nodes, View<int*, CLayout, HostType>& node_to_wall){
    int count=0;
    for(int i = 0; i<rgn.size(); i++){
        if(rgn[i]==Wall){
            wall_nodes[count]=i + 1; // 1-indexed
            node_to_wall[i]=count + 1; // 1-indexed
            count++;
        }else{
            node_to_wall[i]=0;
        }
    }
}


template<class Device, class GridDevice>
View<double*,CLayout,GridDevice> setup_grid_psi(const MagneticField<Device>& magnetic_field, const View<RZPair*,CLayout,GridDevice>& gx_h){
    int nnode = gx_h.extent(0);
    View<RZPair*,CLayout,Device> gx("gx",nnode);
    Kokkos::deep_copy(gx, gx_h);

    // Calculate psi at grid points
    View<double*,CLayout,Device> psi("psi", nnode);
    Kokkos::parallel_for("psi_calc", Kokkos::RangePolicy<typename Device::execution_space>(0, nnode), KOKKOS_LAMBDA( const int idx ){
        double psi_local;
        magnetic_field.psi_bicub.der_zero(gx(idx).r,gx(idx).z, psi_local);
        psi(idx) = max(1e-70, psi_local);
    });

    // Check whether the first vertex is the magnetic axis
    // If so, set psi exactly to 0. This is needed for 1D
    // psi search on irregular grid
    double psi_zero_eps = 5.0e-3;
    Kokkos::parallel_for("psi_calc_check_first", Kokkos::RangePolicy<typename Device::execution_space>(0, 1), KOKKOS_LAMBDA( const int idx ){
        double dr = gx(idx).r - magnetic_field.equil.axis_r;
        double dz = gx(idx).z - magnetic_field.equil.axis_z;
        if(sqrt(dr*dr + dz*dz) < psi_zero_eps){
            psi(idx) = 0.0;
        }
    });

    Kokkos::fence();

    View<double*,CLayout,GridDevice> psi_h("psi", nnode);
    Kokkos::deep_copy(psi_h, psi);

    return psi_h;
}

// Returns volume-weighted average of node psi values along each flux surface
template<class Device>
View<double*,CLayout,HostType> get_psi_surf(const MagneticField<Device>& magnetic_field,
                     const View<double*,CLayout,HostType>& psi_h,
                     const View<int*, CLayout, HostType>& nnodes_on_surface_h,
                     const View<int**,CLayout,HostType>& surf_idx_h,
                     const View<double*,CLayout,HostType>& node_vol_h){

    // Copy host views to device
    View<double*,CLayout,Device> psi(NoInit("psi"), psi_h.layout());
    View<int*, CLayout,Device> nnodes_on_surface(NoInit("nnodes_on_surface"),nnodes_on_surface_h.layout());
    View<int**,CLayout,Device> surf_idx(NoInit("surf_idx"),surf_idx_h.layout());
    View<double*,CLayout,Device> node_vol(NoInit("node_vol"),node_vol_h.layout());
    Kokkos::deep_copy(psi, psi_h);
    Kokkos::deep_copy(nnodes_on_surface, nnodes_on_surface_h);
    Kokkos::deep_copy(node_vol, node_vol_h);
    Kokkos::deep_copy(surf_idx, surf_idx_h);

    // Allocate output view
    View<double*,CLayout,Device> psi_surf(NoInit("psi_surf"), nnodes_on_surface.size());

    // Loop over surfaces
    Kokkos::parallel_for("get_psi_surf", Kokkos::RangePolicy<typename Device::execution_space>(0,psi_surf.size()), KOKKOS_LAMBDA( const int i ){
        double surf_vol = 0.0;
        double psi_sum =  0.0;
        for(int j=0; j<nnodes_on_surface(i); j++){
            int k=surf_idx(i, j)-1; // 1-indexed

            // Calculate psi_surf by taking a volume-weighted average of psi values along
            // the nodes of the surface
            psi_sum += psi(k)*node_vol(k);
            surf_vol += node_vol(k);
        }
        psi_surf(i) = psi_sum/surf_vol;
    });
    Kokkos::fence();

    View<double*,CLayout,HostType> psi_surf_h(NoInit("psi_surf"), psi_surf.layout());
    Kokkos::deep_copy(psi_surf_h, psi_surf);
    return psi_surf_h;
}

/*
 int nsurfaces;
    std::vector<int> nnodes_on_surface;
    int sum_nnodes_on_surfaces;
    std::vector<int> nodes_of_surfaces;
    */
template<class Device>
View<double*,CLayout,HostType> setup_psi_surf2(const MagneticField<Device>& magnetic_field,
                     const View<int*, CLayout, HostType>& nnodes_on_surface_h,
                     const View<int**,CLayout, HostType>& surf_idx_h,
                     const View<double*,CLayout,HostType>& psi_surf_h,
                     const View<RZPair*,CLayout,HostType>& gx_h){ 
    const double use_surf_thresh=3.0e-2;
    int nnode = gx_h.extent(0);
    View<RZPair*,CLayout,Device> gx("gx",nnode);
    View<int*, CLayout, Device> nnodes_on_surface(NoInit("nnodes_on_surface"),nnodes_on_surface_h.layout());
    View<int**,CLayout, Device> surf_idx(NoInit("surf_idx"),surf_idx_h.layout());
    Kokkos::deep_copy(nnodes_on_surface, nnodes_on_surface_h);
    Kokkos::deep_copy(surf_idx, surf_idx_h);
    Kokkos::deep_copy(gx, gx_h);

    int nsurfaces = nnodes_on_surface.size();

    int npsi_surf2=0;
    View<bool*,CLayout,Device> use_surf("use_surf", nsurfaces);
    Kokkos::parallel_reduce("get_psi_surf2", Kokkos::RangePolicy<typename Device::execution_space>(0,psi_surf_h.size()), KOKKOS_LAMBDA( const int i, int& npsi_surf2_l){
        double zmin=1e10;
        for(int j=0; j<nnodes_on_surface(i); j++){
            int k=surf_idx(i, j)-1; // 1-indexed
            // Determine which flux surfaces to use for 1D diagnosis
            // Only use flux surfaces that cross outer midplane
            if ( (std::abs(gx(k).z-magnetic_field.equil.axis_z) < zmin) && (gx(k).r-(magnetic_field.equil.axis_r - 5.0e-5) >= 0.0) ){
                zmin=std::abs(gx(k).z -magnetic_field.equil.axis_z);
            }
        }

        if (zmin <= use_surf_thresh){
            use_surf(i) = true;
            npsi_surf2_l++;
        }else{
            use_surf(i) = false;
        }
    }, npsi_surf2);
    Kokkos::fence();

    // Bring results to host
    View<bool*,CLayout,HostType> use_surf_h("use_surf", nsurfaces);
    Kokkos::deep_copy(use_surf_h, use_surf);

    // Populate psi_surf2 with desired surfaces from psi_surf
    View<double*, CLayout, HostType> psi_surf2_h(NoInit("psi_surf2"),npsi_surf2);
    int j=0;
    for(int i=0; i<nsurfaces; i++){
        if (use_surf_h(i)){
            psi_surf2_h(j) = psi_surf_h(i);
            j++;
        }
    }

    // Make sure surfaces are in ascending order
    View<int*, HostType> idx(NoInit("idx"),npsi_surf2);
    for(int i=0; i<npsi_surf2; i++) idx(i) = i + 1; // idx is one-indexed because it is reordered in Fortran
    int left=1;
    int right=npsi_surf2;
    my_qsort(left,right,idx.data(),psi_surf2_h.data(),npsi_surf2);
    View<double*, HostType> tmp(NoInit("tmp"),npsi_surf2);
    for(int i=0; i<npsi_surf2; i++) tmp(i)=psi_surf2_h(idx(i) - 1); // -1 since idx is one-indexed
    for(int i=0; i<npsi_surf2; i++) psi_surf2_h(i)=tmp(i);

    return psi_surf2_h;
}

// Calculate and store magnetic field at nodes
template<class Device, class GridDevice>
View<double*,CLayout,GridDevice> setup_grid_bfield(const MagneticField<Device>& magnetic_field, const View<RZPair*,CLayout,GridDevice>& gx_h){
    int nnode = gx_h.extent(0);
    View<RZPair*,CLayout,Device> gx("gx",nnode);
    Kokkos::deep_copy(gx, gx_h);

    // Calculate psi at grid points
    View<double*,CLayout,Device> bfield("bfield", nnode);
    Kokkos::parallel_for("node_bfield", Kokkos::RangePolicy<typename Device::execution_space>(0,nnode), KOKKOS_LAMBDA( const int idx ){
        double br, bz, bphi;
        magnetic_field.bvec_interpol(gx(idx).r,gx(idx).z,0.0,br,bz,bphi);
        bfield(idx) = sqrt(br*br + bz*bz + bphi*bphi);
    });
    Kokkos::fence();

    View<double*,CLayout,GridDevice> bfield_h("bfield", nnode);
    Kokkos::deep_copy(bfield_h, bfield);

    return bfield_h;
}


/*void plot_grid(Grid<HostType>& grid){
    int nx = 80;
    int ny = 40;
    double dr = (equil.max_r - equil.min_r)/nx;
    double dz = (equil.max_z - equil.min_z)/ny;
    printf("\n");
    std::vector<int> prevrow(nx+1,-1);

    for (int j=ny;j>=0;j--){
        printf("\n");
        int prevnum = -1;
        for (int i=0;i<=nx;i++){
            double r = equil.min_r + i*dr;
            double z = equil.min_z + j*dz;
            int num = get_analytic_psi(equil,r,z)*5.0;
            char mychar = 48+num;
            if (prevnum!=num || prevrow[i]!=num) printf("%c",mychar);
            else printf(" ");
            prevnum = num;
            prevrow[i] = num;
        }
    }
    printf("\n");
}
*/

}
#endif