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.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
      if(is_rank_zero()) 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;
        }
    }
}

void order_wall(const View<RZPair*,CLayout,HostType>& gx, const View<int*, CLayout, HostType>& rgn, const View<Vertex*, CLayout, HostType>& nodes, const View<int*,CLayout,DeviceType>& num_t_node_d, const View<int**,CLayout,DeviceType>& tr_node_d,  const View<int**,CLayout,HostType>& adj, double z_axis, const View<int*,CLayout,HostType>& wall_nodes, const View<int*,CLayout,HostType>& node_to_wall){
    // Copy num_t_node and tr_node to host
    View<int*,CLayout,HostType> num_t_node("num_t_node", num_t_node_d.layout());
    View<int**,CLayout,HostType> tr_node("tr_node", tr_node_d.layout());
    Kokkos::deep_copy(num_t_node, num_t_node_d);
    Kokkos::deep_copy(tr_node, tr_node_d);

    int nwall = wall_nodes.size();

    // Allocate temporary array for the re-ordered wall nodes
    View<int*, CLayout, HostType> ordered_wall_nodes("ordered_wall_nodes", wall_nodes.layout());
    int lastnode = -1;
    int thisnode = -1;
    int nextnode = wall_nodes(0);
    for(int w=0; w<nwall; w++){
        if (nextnode==lastnode && is_rank_zero()) printf("WARNING: nextnode and lastnode the same!\n");
        ordered_wall_nodes(w) = nextnode;
        lastnode = thisnode;
        thisnode = nextnode;
        nextnode = -1;

        // loop through triangles touching this node
        for(int j=0; j<num_t_node(thisnode-1); j++){
            int itr = tr_node(j,thisnode-1);

            // Find index of wall_node in triangle
            int iv = nodes(itr).vertex[0]==thisnode ? 0 :
                     nodes(itr).vertex[1]==thisnode ? 1 : 2;

            // Loop through nodes of this triangle
            for(int k=0; k<3; k++){
                int candidate = nodes(itr).vertex[k];
                if(rgn(candidate-1)==Wall &&
                   candidate!=thisnode && // Skip triangle node of the wall_node in question
                   // adj matrix setup so index 0 -> edge 12,  index 1 -> edge 02, index 2 -> edge 01.
                   adj(itr, 3 - (iv + k))==-1 && // Not a boundary if there is an adjacent triangle
                   candidate!=lastnode){ // Skip the lastnode entered (move forward) if after first node

                    // Candidate is the correct next node
                    nextnode = candidate;
                    break;
                }
            }
            if(nextnode!=-1) break; // next node has been found
        }
    }

    if (nextnode != ordered_wall_nodes(0) && is_rank_zero()) printf("WARNING: wall_nodes ordered doesnt form closed surface\n");

    // First node should be right above the inboard midplane, and last node right below
    // Find inner-most intersection of z-axis
    int crossing_segment = -1;
    double min_r_crossing;
    bool clockwise;
    for(int w=0; w<nwall; w++){
        // Use next point to define segment
        int w2 = (w+1)%nwall;
        double z1 = gx(ordered_wall_nodes(w)-1).z;
        double z2 = gx(ordered_wall_nodes(w2)-1).z;
        // Check if z-axis is crossed
        bool z_axis_crossed = (z1>=z_axis && z2<z_axis) || (z1<z_axis && z2>=z_axis);
        if(z_axis_crossed){
            double r1 = gx(ordered_wall_nodes(w)-1).r;
            double r2 = gx(ordered_wall_nodes(w2)-1).r;
            double dr = r2-r1;
            double dz = z2-z1; // z_axis_crossed condition guarantees this is not exactly zero
            double r_crossing = r1 + dr*(z_axis-z1)/dz;
            // Bound r_crossing to reside on the segment: this could be an issue if dz!= but is so small
            // there is a numerical issue. (It may also be impossible.)
            r_crossing = std::max(std::min(r1, r2), r_crossing);
            r_crossing = std::min(std::max(r1, r2), r_crossing);
            if(crossing_segment==-1 || r_crossing<min_r_crossing){
                crossing_segment = w;
                min_r_crossing = r_crossing;
                clockwise = (z2>z1);
            }
        }
    }

    // Edge case: wall does not intersect z-axis. In this case, keep the existing first node.
    // Because tr_node order is not deterministic, some MPI ranks may choose clockwise/counterclockwise differently
    // So choose an arbitrary but consistent order now to make sure wall_nodes is the same for every rank
    if(crossing_segment==-1){
        if(is_rank_zero()) printf("WARNING: Wall does not appear to intersect z-axis.\n");

        crossing_segment = 0;
        int p0 = crossing_segment;
        int p1 = (crossing_segment+1)%nwall; // Second point in current ordering
        int pm1 = (nwall+crossing_segment-1)%nwall; // Final point in current ordering
        double r0 = gx(ordered_wall_nodes(p0)-1).r;
        double r1 = gx(ordered_wall_nodes(p1)-1).r;
        double rm1 = gx(ordered_wall_nodes(pm1)-1).r;
        double z0 = gx(ordered_wall_nodes(p0)-1).z;
        double z1 = gx(ordered_wall_nodes(p1)-1).z;
        double zm1 = gx(ordered_wall_nodes(pm1)-1).z;

        // Determine direction based on looking at the 2 neighboring points. This may put us in the wrong global
        // order, but the important thing is to be consistent.
        double det = (r1-r0)*(zm1-z0)-(rm1-r0)*(z1-z0);
        if(det!=0.0){
            clockwise = (det<0.0);
        }else{
            if(r1!=rm1){
                // If the three points are colinear, choose direction based on which point has lower r
                clockwise = (r1<rm1);
            }else{
                if(z1!=zm1){
                    clockwise = (z1<zm1);
                }else{
                    if(nwall<=2){
                        clockwise = true; // This doesn't matter if there are only 1 or 2 wall nodes
                    }else{
                        exit_XGC("\nError: Two wall nodes are the same.\n");
                    }
                }
            }
        }
    }

    if(clockwise){
        // If already clockwise, we just need to cyclically shift the wall
        int first_point = (crossing_segment+1)%nwall;
        for(int w=0; w<nwall; w++){
            wall_nodes(w) = ordered_wall_nodes((first_point+w)%nwall);
        }
    }else{
        // If counterclockwise, we need to reverse the order
        int first_point = crossing_segment;
        for(int w=0; w<nwall; w++){
            wall_nodes(w) = ordered_wall_nodes((nwall+first_point-w)%nwall);
        }
    }

    // Reset node_to_wall
    Kokkos::deep_copy(node_to_wall, 0);
    for(int i = 0; i<nwall; i++){
        node_to_wall(wall_nodes(i)-1) = i+1; // 1-indexed
    }
}


template<class Device, class GridDevice>
View<double*,CLayout,GridDevice> setup_grid_psi(const MagneticField<Device>& magnetic_field, const View<RZPair*,CLayout,GridDevice>& gx_h, double phi){
    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.get_psi(gx(idx).r,gx(idx).z, phi);
        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;
}

/*
 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,
                     View<int*, CLayout, HostType>& psi_map_surf_to_surf2,
                     View<int*, CLayout, HostType>& psi_map_surf2_to_surf){
    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);
            }
        }

#ifdef STELLARATOR
        // AM: This is a temporary solution for the 1D diagnostics in the stellarator case which should be ok for the core.
        // If we don't use it different planes could end up having a different number of flux surfaces which doesn't work.
        if (true){
#else
        if (zmin <= use_surf_thresh){
#endif
            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);
    psi_map_surf_to_surf2 = View<int*, CLayout, HostType>(NoInit("psi_map_surf_to_surf2"),nsurfaces);
    psi_map_surf2_to_surf = View<int*, CLayout, HostType>(NoInit("psi_map_surf2_to_surf"),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);

            // Save mapping
            psi_map_surf_to_surf2(i) = j;
            psi_map_surf2_to_surf(j) = i;
            j++;
        }else{
            // If this surface does not occur in psi_surf2
            psi_map_surf_to_surf2(i) = -1;
        }
    }

    // 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, double phi){
    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,phi,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;
}

// Calculate magnetic field vector and magnitude at nodes
template<class Device, class GridDevice>
View<double**,CLayout,GridDevice> setup_bfield_vecmag(const MagneticField<Device>& magnetic_field, const View<RZPair*,CLayout,GridDevice>& gx_h, double phi){
    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", 4, 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,phi,br,bz,bphi);
        bfield(0, idx) = br;
        bfield(1, idx) = bz;
        bfield(2, idx) = bphi;
        bfield(3, idx) = sqrt(br*br + bz*bz + bphi*bphi);
    });
    Kokkos::fence();

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

    return bfield_h;
}

// Calculate magnetic field vector and magnitude at nodes
template<class Device, class GridDevice>
View<double**,CLayout,GridDevice> setup_gradpsi(const MagneticField<Device>& magnetic_field, const View<RZPair*,CLayout,GridDevice>& gx_h, double phi){
    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> gradpsi("gradpsi", 2, nnode);
    Kokkos::parallel_for("node_gradpsi", Kokkos::RangePolicy<typename Device::execution_space>(0,nnode), KOKKOS_LAMBDA( const int idx ){
        double psi,dpsi_dr,dpsi_dz, dpsi_dphi;
        magnetic_field.get_psi_and_derivs(gx(idx).r, gx(idx).z, phi, psi,dpsi_dr,dpsi_dz, dpsi_dphi );
        gradpsi(0,idx) = dpsi_dr;
        gradpsi(1,idx) = dpsi_dz;
    });
    Kokkos::fence();

    View<double**,CLayout,GridDevice> gradpsi_h("gradpsi", 2, nnode);
    Kokkos::deep_copy(gradpsi_h, gradpsi);

    return gradpsi_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