grid_files.hpp

Go to the documentation of this file.

#ifndef GRID_FILES_HPP
#define GRID_FILES_HPP

// Sets the maximum number of planes that is read for stellarator simulations, e.g. 4 means up to 9999 planes.
// (This will probably be removed later with a different mesh format.)
#define STELLARATOR_PLANE_LIMIT 4

#ifdef USE_MPI
#include "my_mpi.hpp"
#endif
#include "file_reader.hpp"
#include "magnetic_field.hpp"
#include "grid_structs.hpp"
#include "space_settings.hpp"
#include "NamelistReader.hpp"
#include "check_aif_flx_format.hpp"
#include "send_recv_toroidal.hpp"

extern "C" void set_fortran_flx_aif_format(int is_flx_aif_format_int);

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

    // Contents of .node
    int nnodes;
    View<RZPair*, HostType> x;
    View<int*, HostType> rgn;

    // Contents of .ele
    int ntriangles;
    View<Vertex*, HostType> nodes;

    // Contents of .flx.aif
    View<int*, HostType> i_x;

    int nsurfaces;
    int nsurfaces1, nsurfaces2, nsurfaces3a, nsurfaces3b;
    View<int*, HostType> i_surf_separatrices;
    View<int*, HostType> nnodes_on_surface;
    int sum_nnodes_on_surfaces;
    int num_non_aligned;
    View<int*, HostType> non_aligned_vert;
    View<int*, HostType> non_aligned_nsurf;
    View<int**, CLayout, HostType> non_aligned_surf_idx;
    View<int**, CLayout, HostType> surf_idx;

    View<double*, HostType> psiN; 

    void convert_to_x_rgn_format(const View<double**, CLayout, HostType>& tmp_nodes, View<RZPair*, HostType>& x, View<int*, HostType>& rgn, bool is_stellarator){
        Kokkos::parallel_for("transpose_x_and_rgn", Kokkos::RangePolicy<HostExSpace>(0,tmp_nodes.extent(0)), KOKKOS_LAMBDA( const int i ){
            // tmp_nodes[i,0]; <-- first column is ignored
            x(i).r = tmp_nodes(i,1);
            x(i).z = tmp_nodes(i,2);
            rgn(i) = is_stellarator ? tmp_nodes(i,5) : tmp_nodes(i,3); // For stellarator .node files rgn is column 6, for tokamaks column 4
        });
    }

    void convert_to_node_format(const View<int**, CLayout, HostType>& tmp_ele, View<Vertex*, HostType>& nodes){
        Kokkos::parallel_for("transpose_x_and_rgn", Kokkos::RangePolicy<HostExSpace>(0,tmp_ele.extent(0)), KOKKOS_LAMBDA( const int i ){
            // tmp_ele[i,0]; <-- first column is ignored
            nodes(i).vertex[0] = tmp_ele(i,1);
            nodes(i).vertex[1] = tmp_ele(i,2);
            nodes(i).vertex[2] = tmp_ele(i,3);
        });
    }

    View<int*, HostType> pack_scalars(int max_nodes_on_surface, int n_xpts, int n_separatrices) const{
        int n_scalars = 11;
        View<int*, HostType> scalars_s("scalars", n_scalars);
        scalars_s(0) = nnodes;
        scalars_s(1) = ntriangles;
        scalars_s(2) = nsurfaces;
        scalars_s(3) = max_nodes_on_surface;
        scalars_s(4) = num_non_aligned;
        scalars_s(5) = n_xpts;
        scalars_s(6) = n_separatrices;
        scalars_s(7) = nsurfaces1;
        scalars_s(8) = nsurfaces2;
        scalars_s(9) = nsurfaces3a;
        scalars_s(10) = nsurfaces3b;
        return scalars_s;
    }

    void unpack_scalars(const View<int*, HostType>& scalars, int& max_nodes_on_surface, int& n_xpts, int& n_separatrices){
        nnodes = scalars(0);
        ntriangles = scalars(1);
        nsurfaces = scalars(2);
        max_nodes_on_surface = scalars(3);
        num_non_aligned = scalars(4);
        n_xpts = scalars(5);
        n_separatrices = scalars(6);
        nsurfaces1 = scalars(7);
        nsurfaces2 = scalars(8);
        nsurfaces3a = scalars(9);
        nsurfaces3b = scalars(10);
    }

    // Default constructor, to use to deallocate the file arrays after use
    PlaneFiles(){}

// Should remove this USE_MPI at some point
#ifdef USE_MPI

    // Copy data from a PlaneFiles object on a neighboring rank
    PlaneFiles(const MyMPI& my_mpi, const PlaneFiles& plane_files){
        /*** Broadcast grid contents to other MPI ranks ***/

        // First broadcast scalars so we know the sizes of the arrays
        int max_nodes_on_surface = plane_files.surf_idx.extent(1);
        int n_xpts = plane_files.i_x.size();
        int n_separatrices = plane_files.i_surf_separatrices.size();
        auto scalars_s = plane_files.pack_scalars(max_nodes_on_surface, n_xpts, n_separatrices);
        View<int*, HostType> scalars("scalars", scalars_s.layout());
        send_right(my_mpi, scalars_s, scalars);
        unpack_scalars(scalars, max_nodes_on_surface, n_xpts, n_separatrices);

        // node and ele views
        x = View<RZPair*, HostType>(NoInit("x"),nnodes);
        rgn = View<int*, HostType>(NoInit("rgn"),nnodes);
        nodes = View<Vertex*, HostType>("nodes",ntriangles);

        // flx.aif views
        i_x = View<int*, HostType>("i_x", n_xpts);
        i_surf_separatrices = View<int*, HostType>(NoInit("i_surf_separatrices"), n_separatrices);
        non_aligned_vert = View<int*, HostType>(NoInit("non_aligned_vert"), num_non_aligned);
        non_aligned_nsurf = View<int*, HostType>(NoInit("non_aligned_nsurf"), num_non_aligned);
        non_aligned_surf_idx = View<int**, CLayout, HostType>(NoInit("non_aligned_surf_idx"), num_non_aligned, 4);
        nnodes_on_surface = View<int*, HostType>(NoInit("nnodes_on_surface"), nsurfaces);
        surf_idx = View<int**, CLayout, HostType>(NoInit("surf_idx"),nsurfaces, max_nodes_on_surface);

        // Send grid arrays
        send_right(my_mpi, plane_files.rgn, rgn);
        // Cast these as double and int respectively for MPI simplicity (or could create MPI_Types)
        View<double*,CLayout, HostType, Kokkos::MemoryTraits<Kokkos::Unmanaged>> x_dbl((double*)(x.data()), x.size()*2);
        View<double*,CLayout, HostType, Kokkos::MemoryTraits<Kokkos::Unmanaged>> src_x_dbl((double*)(plane_files.x.data()), plane_files.x.size()*2);
        send_right(my_mpi, src_x_dbl, x_dbl);
        View<int*,CLayout, HostType, Kokkos::MemoryTraits<Kokkos::Unmanaged>> nodes_dbl((int*)(nodes.data()), nodes.size()*3);
        View<int*,CLayout, HostType, Kokkos::MemoryTraits<Kokkos::Unmanaged>> src_nodes_dbl((int*)(plane_files.nodes.data()), plane_files.nodes.size()*3);
        send_right(my_mpi, src_nodes_dbl, nodes_dbl);

        // These might be size 0 on one plane but not the other
        send_right(my_mpi, plane_files.i_x, i_x);
        send_right(my_mpi, plane_files.i_surf_separatrices, i_surf_separatrices);
        send_right(my_mpi, plane_files.non_aligned_vert, non_aligned_vert);
        send_right(my_mpi, plane_files.non_aligned_nsurf, non_aligned_nsurf);
        send_right(my_mpi, plane_files.non_aligned_surf_idx, non_aligned_surf_idx);
        send_right(my_mpi, plane_files.nnodes_on_surface, nnodes_on_surface);
        send_right(my_mpi, plane_files.surf_idx, surf_idx);
    }

    // Read the files
    PlaneFiles(const std::string& node_file, const std::string& element_file, const std::string& flx_aif_file, bool is_stellarator, const MPI_Comm& comm){
        // MPI setup
        int my_rank;
#ifdef USE_MPI
        MPI_Comm_rank(comm, &my_rank);
#else
        my_rank = 0;
#endif
        int ROOT_RANK = 0;

        int max_nodes_on_surface;
        int n_xpts;
        int n_separatrices;
        // Only one MPI rank reads the files, then broadcasts to the others
        if(my_rank==ROOT_RANK){

            FileReader file_reader;

            file_reader.open(node_file);
            file_reader.read(&nnodes, 1); // Get nnodes from file
            file_reader.skip(3); // Skip 3 unused values in file

            // Use dummy array since reformatting is needed
            View<double**, CLayout, HostType> tmp_nodes(NoInit("tmp_nodes"),nnodes,4); // 4 entries per node: index, r, z, region
            if(is_stellarator){
                // AM: I think (psi/psi_a)^(1/2) and poloidal angle are only used for debugging in Fortran XGC-S so do not store them.
                Kokkos::resize(tmp_nodes,nnodes,6); // 6 entries per node: index, r, z, (psi/psi_a)^(1/2), poloidal angle, region
            }
            file_reader.read(tmp_nodes.data(), tmp_nodes.size());

            // Reformat to x and rgn
            x = View<RZPair*, HostType>(NoInit("x"),nnodes);
            rgn = View<int*, HostType>(NoInit("rgn"),nnodes);
            convert_to_x_rgn_format(tmp_nodes, x, rgn, is_stellarator);

            file_reader.open(element_file);
            file_reader.read(&ntriangles, 1); // Get ntriangles from file
            file_reader.skip(2); // Skip 2 unused values in file

            // Use dummy array since reformatting is needed
            View<int**, CLayout, HostType> tmp_ele(NoInit("tmp_ele"),ntriangles,4); // 4 entries per node: index, 3 vertices of the triangle
            file_reader.read(tmp_ele.data(), tmp_ele.size());

            // Reformat since first column is ignored
            nodes = View<Vertex*, HostType>("nodes",ntriangles);
            convert_to_node_format(tmp_ele, nodes);

            if(nnodes==1) exit_XGC("\nGrid must have more than 1 nnode!\n"); // This is assumed in the flx format check
            bool old_flx_aif_format = is_stellarator ? false : is_old_flx_format(flx_aif_file);
            set_fortran_flx_aif_format(old_flx_aif_format ? 1 : 0);

            file_reader.open(flx_aif_file);

            if(!is_stellarator){
                if(old_flx_aif_format){
                    file_reader.read(&nsurfaces, 1); // Get number of surfaces
                    n_xpts = 0;
                    n_separatrices = 0;
                    i_x = View<int*, HostType>("i_x",0); // No xpoints specified
                    i_surf_separatrices = View<int*, HostType>(NoInit("i_surf_separatrices"), 0); // No separatrices specified

                    // These don't get read or used by the old format
                    nsurfaces1 = 0;
                    nsurfaces2 = 0;
                    nsurfaces3a = 0;
                    nsurfaces3b = 0;
                }else{
                    file_reader.read(&n_xpts, 1); // Get number of x-points
                    i_x = View<int*, HostType>("i_x", n_xpts);

                    // Current flx.aif files list 2 x-point indices, regardless of n_xpts
                    int i_xpt1, i_xpt2;
                    file_reader.read(&i_xpt1, 1); // Get index of 1st x-point
                    file_reader.read(&i_xpt2, 1); // Get index of 2nd x-point

                    // Check for consistency
                    if(n_xpts==1 && i_xpt2 != -1) exit_XGC("Error in flx.aif file: n_xpts=1, but i_xpt2 is specified.\n");
                    if(n_xpts==2 && i_xpt2 <= 0) exit_XGC("Error in flx.aif file: n_xpts=2, but i_xpt2 is not specified.\n");

                    if(n_xpts>=1) i_x(0) = i_xpt1 - 1; // switch to zero-index
                    if(n_xpts>=2) i_x(1) = i_xpt2 - 1;

                    // Next are the nsurface variables
                    file_reader.read(&nsurfaces1, 1); // Get number of surfaces
                    file_reader.read(&nsurfaces2, 1); // Get number of surfaces
                    file_reader.read(&nsurfaces3a, 1); // Get number of surfaces
                    file_reader.read(&nsurfaces3b, 1); // Get number of surfaces

                    // nsurfaces is actually the sum of these 4 inputs
                    nsurfaces = nsurfaces1 + nsurfaces2 + nsurfaces3a + nsurfaces3b;

                    // Get surface index of separatrices
                    // Hard-coded to 2, but generalize the read just in case
                    n_separatrices = 2;
                    i_surf_separatrices = View<int*, HostType>(NoInit("i_surf_separatrices"), n_separatrices);
                    file_reader.read(&i_surf_separatrices(0), 1);
                    file_reader.read(&i_surf_separatrices(1), 1);
                }
            }else{ // Stellarator .grid file
                n_xpts = 0;
                n_separatrices = 0;
                i_x = View<int*, HostType>("i_x",0); // No xpoints specified
                i_surf_separatrices = View<int*, HostType>(NoInit("i_surf_separatrices"), 0); // No separatrices specified

                file_reader.read(&nsurfaces1, 1); // Get number of surfaces in the core
                nsurfaces2 = 0; // No surfaces in the SOL
                nsurfaces3a = 0; // No surfaces in the lower private region
                nsurfaces3b = 0; // No surfaces in the upper private region

                nsurfaces = nsurfaces1 + nsurfaces2 + nsurfaces3a + nsurfaces3b;   
            }
            // Read in nnodes on each surface
            nnodes_on_surface = View<int*, HostType>(NoInit("nnodes_on_surface"), nsurfaces);
            if(!is_stellarator){
                file_reader.read(nnodes_on_surface.data(), nnodes_on_surface.size());
            }else{ // Stellarator .grid file
                psiN = View<double*, HostType>(NoInit("normalized_psi_on_surface"), nsurfaces);
                int dummy_idx;
                for(int i=0; i<nnodes_on_surface.size(); i++){
                    file_reader.read(&nnodes_on_surface(i), 1);
                    file_reader.read(&psiN(i), 1);
                    if(i<nnodes_on_surface.size() - 1) file_reader.read(&dummy_idx, 1);
                }
            }

            // Loop through surfaces to find max number of nodes in a surface
            max_nodes_on_surface = 0;
            for(int i=0; i<nnodes_on_surface.size(); i++) max_nodes_on_surface = std::max(nnodes_on_surface(i), max_nodes_on_surface);

            // Allocate surf_idx and initialize to 0
            surf_idx = View<int**, CLayout, HostType>("surf_idx",nsurfaces, max_nodes_on_surface);

            // Read in values
            if(!is_stellarator){
                for(int i=0; i<nnodes_on_surface.size(); i++){
                    file_reader.read(&surf_idx(i,0), nnodes_on_surface(i));
                }
            }else{ // The stellarator .grid file has no enumeration of surface indices so we assume they are in order
                int surf_idx_counter = 0;
                for(int i=0; i<nnodes_on_surface.size(); i++){
                    for(int j=0; j<nnodes_on_surface(i); j++){
                        surf_idx_counter++;
                        surf_idx(i,j) = surf_idx_counter;
                    }
                }
            }

            int should_be_negative_one;
            if(!is_stellarator){
                // There is a "-1" in the file here to mark the end of the section
                file_reader.read(&should_be_negative_one, 1);
                if(should_be_negative_one != -1){
                    printf("\nExpected -1 at end of section in surface file, encountered %d instead\n", should_be_negative_one);
                    exit_XGC("Error reading surface file.\n");
                }

                file_reader.read(&num_non_aligned, 1);
                printf(" -  reading num_non_aligned: %d",num_non_aligned);
            }else{ // The stellarator .grid file has no non-aligned vertex data
                num_non_aligned = 0;
            }

            // Read the non-aligned vertex data
            non_aligned_vert = View<int*, HostType>(NoInit("non_aligned_vert"), num_non_aligned);
            non_aligned_nsurf = View<int*, HostType>(NoInit("non_aligned_nsurf"), num_non_aligned);
            // Initialize to 0
            non_aligned_surf_idx = View<int**, CLayout, HostType>("non_aligned_surf_idx", num_non_aligned, 4);
            for(int i=0; i<num_non_aligned; i++){
                file_reader.read(&non_aligned_vert(i), 1);
                if(old_flx_aif_format){
                    file_reader.read(&non_aligned_nsurf(i), 1);
                }else{
                    // New format does not contain non_aligned_nsurf
                    non_aligned_nsurf(i)=2;
                }
                file_reader.read(&non_aligned_surf_idx(i,0), non_aligned_nsurf(i));
            }

            if(!is_stellarator){
                // Another "-1" at the end of the file
                file_reader.read(&should_be_negative_one, 1);
                if(should_be_negative_one != -1){
                    printf("\nExpected -1 at end of the surface file, encountered %d instead\n", should_be_negative_one);
                    exit_XGC("Error reading surface file.\n");
                }
            }
        }

        /*** Broadcast grid contents to other MPI ranks ***/

        // First broadcast scalars so we know the sizes of the arrays
        auto scalars = pack_scalars(max_nodes_on_surface, n_xpts, n_separatrices);

        // Broadcast file size to other processes
#ifdef USE_MPI
        MPI_Bcast(scalars.data(), scalars.size(), MPI_INT, ROOT_RANK, comm);
#endif

        if(my_rank!=ROOT_RANK){
            // Set scalars on non-root ranks
            unpack_scalars(scalars, max_nodes_on_surface, n_xpts, n_separatrices);

            // node and ele views
            x = View<RZPair*, HostType>(NoInit("x"),nnodes);
            rgn = View<int*, HostType>(NoInit("rgn"),nnodes);
            nodes = View<Vertex*, HostType>("nodes",ntriangles);

            // flx.aif views
            i_x = View<int*, HostType>("i_x", n_xpts);
            i_surf_separatrices = View<int*, HostType>(NoInit("i_surf_separatrices"), n_separatrices);
            non_aligned_vert = View<int*, HostType>(NoInit("non_aligned_vert"), num_non_aligned);
            non_aligned_nsurf = View<int*, HostType>(NoInit("non_aligned_nsurf"), num_non_aligned);
            non_aligned_surf_idx = View<int**, CLayout, HostType>(NoInit("non_aligned_surf_idx"), num_non_aligned, 4);
            nnodes_on_surface = View<int*, HostType>(NoInit("nnodes_on_surface"), nsurfaces);
            surf_idx = View<int**, CLayout, HostType>(NoInit("surf_idx"),nsurfaces, max_nodes_on_surface);
        }

        // Broadcast grid arrays
#ifdef USE_MPI
        MPI_Bcast( x.data(), x.size()*2, MPI_DOUBLE, ROOT_RANK, comm);
        MPI_Bcast( rgn.data(), rgn.size(), MPI_INT, ROOT_RANK, comm);
        MPI_Bcast( nodes.data(), nodes.size()*3, MPI_INT, ROOT_RANK, comm);

        if(i_x.size()>0) MPI_Bcast( i_x.data(), i_x.size(), MPI_INT, ROOT_RANK, comm);
        if(i_surf_separatrices.size()>0) MPI_Bcast( i_surf_separatrices.data(), i_surf_separatrices.size(), MPI_INT, ROOT_RANK, comm);
        if(non_aligned_vert.size()>0) MPI_Bcast( non_aligned_vert.data(), non_aligned_vert.size(), MPI_INT, ROOT_RANK, comm);
        if(non_aligned_nsurf.size()>0) MPI_Bcast( non_aligned_nsurf.data(), non_aligned_nsurf.size(), MPI_INT, ROOT_RANK, comm);
        if(non_aligned_surf_idx.size()>0) MPI_Bcast( non_aligned_surf_idx.data(), non_aligned_surf_idx.size(), MPI_INT, ROOT_RANK, comm);
        MPI_Bcast( nnodes_on_surface.data(), nnodes_on_surface.size(), MPI_INT, ROOT_RANK, comm);
        MPI_Bcast( surf_idx.data(), surf_idx.size(), MPI_INT, ROOT_RANK, comm);
#endif
    }
#endif

    // Options: hex vs circular
    enum GridCreateOpt{
        Hexagonal=0,
        Circular=1
    };

    void construct_hex_grid(int nshells, double raxis, double zaxis, double rscale, double zscale) {
        // First, build construction of what .node and .ele files would look like

        // Build hexagonal sample grid of concentric shells
        // Minimum 1 shell, which produces 7 nodes and 6 triangles

        // Grid construction uses this enum for directions
        enum { NE = 0, N, NW, SW, S, SE };

    // Node indices of each surface
        int max_nodes_on_surface = 6*nshells;
        surf_idx = View<int**, CLayout, HostType>(NoInit("surf_idx"),nsurfaces, max_nodes_on_surface);

        // Work from center outward
        x(0).r = raxis;
        x(0).z = zaxis;
        rgn(0) = 0; // not wall
    nnodes_on_surface(0) = 1;
        surf_idx(0, 0) = 1; // Since surf_idx is 1-indexed
        int inode = 1;
        int itri = 0;
        for(int ish = 1; ish<=nshells; ish++){
            int nodes_in_shell = 6*ish;
            int first_inner_shell_node = ish==1 ? 0 : inode - 6*(ish-1);
            int first_node = inode;
            int inner_shell_node = first_inner_shell_node;
            // Starting position:
            double r_track = raxis + rscale*ish;
            double z_track = zaxis;
            for (int inodesh = 0; inodesh < nodes_in_shell; inodesh++){
                x(inode).r = r_track;
                x(inode).z = z_track;
                rgn(inode) = (ish==nshells? OnWall : Inside); // last shell is wall

        surf_idx(ish, inodesh) = inode+1; // Since surf_idx is 1-indexed

                if (inodesh>0){
                    nodes(itri).vertex[0] = inner_shell_node;
                    nodes(itri).vertex[1] = inode-1;
                    nodes(itri).vertex[2] = inode;
                    itri++;
                    if (inodesh%ish>0){ // corners only add one triangle
                        nodes(itri).vertex[0] = inner_shell_node;
                        nodes(itri).vertex[1] = (inodesh == nodes_in_shell - 1) ? first_inner_shell_node : (inner_shell_node+1);
                        nodes(itri).vertex[2] = inode;
                        inner_shell_node++;
                        itri++;
                    }
                }

                // Determine position of next node
                int hex_region = inodesh/ish;
                double r_move, z_move;
                if(hex_region == NE)
                    {r_move = -0.5; z_move = 1.0;}
                else if (hex_region == N)
                    {r_move = -1.0; z_move = 0.0;}
                else if (hex_region == NW)
                    {r_move = -0.5; z_move = -1.0;}
                else if (hex_region == SW)
                    {r_move = 0.5; z_move = -1.0;}
                else if (hex_region == S)
                    {r_move = 1.0; z_move = 0.0;}
                else if (hex_region == SE)
                    {r_move = 0.5; z_move = 1.0;}
                r_track += rscale*r_move;
                z_track += zscale*z_move;
                inode++;
            }

            // Add last element to close ring
            nodes(itri).vertex[0] = first_inner_shell_node;
            nodes(itri).vertex[1] = inode-1;
            nodes(itri).vertex[2] = first_node;
            itri++;

        nnodes_on_surface(ish) = nodes_in_shell;
        }

        // Shift node list to 1-indexing (input files are 1-indexed)
        for(int i = 0; i<ntriangles; i++){
            nodes(i).vertex[0]++;
            nodes(i).vertex[1]++;
            nodes(i).vertex[2]++;
        }

    // Field geometry
        int n_xpts = 0; // Assume no X-points for now
        int n_separatrices = 0; // Assume no separatrices for now
        i_x = View<int*, HostType>("i_x", n_xpts);
        i_surf_separatrices = View<int*, HostType>(NoInit("i_surf_separatrices"), n_separatrices);

        // Assume no non-aligned surfaces for now
        num_non_aligned = 0;
        non_aligned_nsurf = View<int*, HostType>(NoInit("non_aligned_nsurf"), num_non_aligned);
        non_aligned_surf_idx = View<int**, CLayout, HostType>(NoInit("non_aligned_surf_idx"), num_non_aligned, 4);
    }

    double node_to_node_dist(RZPair a, RZPair b){
        return sqrt((a.r-b.r)*(a.r-b.r) + (a.z-b.z)*(a.z-b.z));
    }

    void construct_circular_grid(int nshells, double raxis, double zaxis, double rscale, double zscale) {
        // First, build construction of what .node and .ele files would look like

        // Build circular sample grid of concentric shells
        // Minimum 1 shell, which produces 7 nodes and 6 triangles

        // Work from center outward
        x(0).r = raxis;
        x(0).z = zaxis;
        rgn(0) = 0; // not wall

        // Node indices of each surface
        int max_nodes_on_surface = 6*nshells;
        surf_idx = View<int**, CLayout, HostType>(NoInit("surf_idx"),nsurfaces, max_nodes_on_surface);

        nnodes_on_surface(0) = 1;
        surf_idx(0, 0) = 1; // Since surf_idx is 1-indexed
        int inode = 1;
        int itri = 0;
        for(int ish = 1; ish<=nshells; ish++){
            int nodes_in_shell = 6*ish;
            int nodes_in_inner_shell = ish==1 ? 1 : 6*(ish-1);
            int first_inner_shell_node = ish==1 ? 0 : inode - 6*(ish-1);
            int first_node = inode;
            int inner_shell_node = first_inner_shell_node;
            // Starting position:
            for (int inodesh = 0; inodesh <= nodes_in_shell; inodesh++){
                // Create node
                if(inodesh<nodes_in_shell){
                    double angle = TWOPI*double(inodesh)/double(nodes_in_shell);
                    x(inode).r = raxis + ish*rscale*cos(angle);
                    x(inode).z = zaxis + ish*zscale*sin(angle);
                    rgn(inode) = (ish==nshells? OnWall : Inside); // last shell is wall

                    surf_idx(ish, inodesh) = inode+1; // Since surf_idx is 1-indexed
                }

                // Create element(s)
                if (inodesh>0){
                    int outer_node = (inodesh < nodes_in_shell) ? inode : first_node;
                    int prev_outer_node = inode-1;

                    // First, check cross-distances
                    int next_inner_shell_node = inner_shell_node+1;
                    if(next_inner_shell_node-first_inner_shell_node==nodes_in_inner_shell) next_inner_shell_node = first_inner_shell_node;
                    double distance1 = node_to_node_dist(x(outer_node), x(inner_shell_node));
                    double distance2 = node_to_node_dist(x(prev_outer_node), x(next_inner_shell_node));

                    // The next node is far away enough that we need to add two elements
                    if(distance1>distance2 && next_inner_shell_node!=inner_shell_node){
                        if(itri>=nodes.size()) {printf("\nMore elements than expected in circular grid generation. Dynamically resize?"); exit(1); }
                        nodes(itri).vertex[0] = inner_shell_node;
                        nodes(itri).vertex[1] = next_inner_shell_node;
                        nodes(itri).vertex[2] = prev_outer_node;
                        inner_shell_node = next_inner_shell_node;
                        itri++;
                    }

                    if(itri>=nodes.size()) {printf("\nMore elements than expected in circular grid generation. Dynamically resize?"); exit(1); }
                    nodes(itri).vertex[0] = inner_shell_node;
                    nodes(itri).vertex[1] = prev_outer_node;
                    nodes(itri).vertex[2] = outer_node;
                    itri++;
                }

                if(inodesh<nodes_in_shell){
                    inode++;
                }
            }

            nnodes_on_surface(ish) = nodes_in_shell;
        }

        if(itri!=nodes.size()) {printf("\nFewer elements than expected in circular grid generation. Dynamically resize?"); exit(1); }

        // Shift node list to 1-indexing (input files are 1-indexed)
        for(int i = 0; i<ntriangles; i++){
            nodes(i).vertex[0]++;
            nodes(i).vertex[1]++;
            nodes(i).vertex[2]++;
        }

        // Field geometry
        int n_xpts = 0; // Assume no X-points for now
        int n_separatrices = 0; // Assume no separatrices for now
        i_x = View<int*, HostType>("i_x", n_xpts);
        i_surf_separatrices = View<int*, HostType>(NoInit("i_surf_separatrices"), n_separatrices);

        nsurfaces1 = nsurfaces;
        nsurfaces2 = 0;
        nsurfaces3a = 0;
        nsurfaces3b = 0;

        // Assume no non-aligned surfaces for now
        num_non_aligned = 0;
        non_aligned_nsurf = View<int*, HostType>(NoInit("non_aligned_nsurf"), num_non_aligned);
        non_aligned_surf_idx = View<int**, CLayout, HostType>(NoInit("non_aligned_surf_idx"), num_non_aligned, 4);
        non_aligned_vert = View<int*, HostType>(NoInit("non_aligned_vert"), num_non_aligned);
    }

    // Returns the number of shells needed to create a grid with at least nnodes nodes.
    static int nshells_from_nnodes(int nnodes){
        // nnodes = 1 + 3*(nshells+1)*nshells
        //     0  = 3s^2 + 3s + 1 - n
        // s = (-3 +/- sqrt(9 + 4*3*(n-1)))/6
        // s = sqrt(9 + 12*(n-1))/6 - 0.5
        if(nnodes<1) return 0;
        else return std::ceil(sqrt(12*nnodes - 3)/6 - 0.5);
    }

    PlaneFiles(int nshells, double raxis, double zaxis, double rscale, double zscale, GridCreateOpt grid_opt=Hexagonal)
      : nnodes(1 + 3*(nshells+1)*nshells),
        x(NoInit("x"),nnodes),
        rgn(NoInit("rgn"),nnodes),
        ntriangles(6*nshells*nshells),
        nodes(NoInit("nodes"),ntriangles),
        nsurfaces(nshells+1),
        nnodes_on_surface(NoInit("nnodes_on_surface"),nsurfaces),
        sum_nnodes_on_surfaces(nnodes)
    {
        if(grid_opt==Hexagonal){
            construct_hex_grid(nshells, raxis, zaxis, rscale, zscale);
        }else{
            construct_circular_grid(nshells, raxis, zaxis, rscale, zscale);
        }
    }

    PlaneFiles(MagneticField<HostType>& magnetic_field, int nshells, GridCreateOpt grid_opt=Hexagonal)
      : nnodes(1 + 3*(nshells+1)*nshells),
        x(NoInit("x"),nnodes),
        rgn(NoInit("rgn"),nnodes),
        ntriangles(6*nshells*nshells),
        nodes(NoInit("nodes"),ntriangles),
        nsurfaces(nshells+1),
        nnodes_on_surface(NoInit("nnodes_on_surface"),nsurfaces),
        sum_nnodes_on_surfaces(nnodes)
    {
        // Magnetic axis will be the center of the grid
        double raxis = magnetic_field.equil.axis.r;
        double zaxis = magnetic_field.equil.axis.z;

        // Since the axis is not necessarily centered in the field map, take the shorter
        // distance to the edge to use as the dimensions of the grid
        double width  = 2*std::min(magnetic_field.bounds.max_r - magnetic_field.equil.axis.r,
                                   magnetic_field.equil.axis.r - magnetic_field.bounds.min_r);

        double height = 2*std::min(magnetic_field.bounds.max_z - magnetic_field.equil.axis.z,
                                   magnetic_field.equil.axis.z - magnetic_field.bounds.min_z);

        // Make sure the grid is within the psi bounds
        width *= 0.99;
        height *= 0.99;

        if(grid_opt==Hexagonal){
            double triangle_width = width/(2*nshells);
            double triangle_height = height/(2*nshells);
            construct_hex_grid(nshells, raxis, zaxis, triangle_width, triangle_height);
        }else{
            // Circular grid; take the smaller of the width or height to make sure it fits
            double scale = std::min(width,height)/(2*nshells);
            construct_circular_grid(nshells, raxis, zaxis, scale, scale);
        }
    }

    int n_specified_xpts() const{
        return i_x.size();
    }

    View<Equil::XPoint*, HostType> get_xpts() const{
        View<Equil::XPoint*, HostType> xpts(NoInit("xpts"), i_x.size());
        for(int i = 0; i<i_x.size(); i++){
            xpts(i).r = x(i_x(i)).r;
            xpts(i).z = x(i_x(i)).z;
        }
        return xpts;
    }

    // Return first vertex coordinates. Assumes this is the axis
    RZPair get_axis() const{
        return x(0);
    }

    // Returns last vertex coordinates. Assumes this is on the LCFS for stellarators
    RZPair get_reference_point() const{
        return x(x.size()-1);
    }

    // Returns normalized psi coordinates of surfaces temporarily needed for the stellarator version
    View<double*, HostType> get_psiN() const{
        return psiN;
    }
};

struct GridFiles{
    using GridCreateOpt = PlaneFiles::GridCreateOpt;
    static constexpr GridCreateOpt Hexagonal = PlaneFiles::Hexagonal;
    static constexpr GridCreateOpt Circular = PlaneFiles::Circular;

    std::vector<PlaneFiles> plane_files_vec;

    // Allows array-like access pattern
    const PlaneFiles& operator [](int i) const {return plane_files_vec[i];}

    // For test grids
    GridFiles(int nshells, double raxis, double zaxis, double rscale, double zscale, GridCreateOpt grid_opt=Hexagonal){
        plane_files_vec.push_back(PlaneFiles(nshells, raxis, zaxis, rscale, zscale, grid_opt));
    }

    GridFiles(MagneticField<HostType>& magnetic_field, int nshells, GridCreateOpt grid_opt=Hexagonal){
        plane_files_vec.push_back(PlaneFiles(magnetic_field, nshells, grid_opt));
    }

#ifdef USE_MPI
    /* Returns PlaneFiles object using the file-reader constructor. This may be updated so the input file
     * can specify an analytic grid.
     * */
    GridFiles(NLReader::NamelistReader& nlr, bool is_stellarator, const MyMPI& my_mpi){
        // Read file names and path from input file
        nlr.use_namelist("sml_param");

        if(nlr.get<bool>("sml_create_analytic_grid", false)==true){ // Option to use a generated grid instead of reading a grid from file. The grid is composed of concentric circles of vertices. analytic_grid_param can be used to control basic grid properties.
            // Keep analytic grid options simple: number of surfaces, and axis_r
            int nshells = 30;
            double axis_r = 1.5;
            if(nlr.namelist_present("analytic_grid_param")){
                nlr.use_namelist("analytic_grid_param");

                nshells = nlr.get<int>("nsurfaces", 30); // Number of surfaces in the analytic circular grid.
                axis_r = nlr.get<double>("axis_r", 1.5);
            }

            if(is_rank_zero()) printf("\nCreating a circular analytic grid with %d surfaces and %d nodes.\n", nshells, 1 + 3*(nshells+1)*nshells);

            double axis_z = 0.0;

            // Leave a margin so that all vertices are inside equilibrium
            double scale = 0.999/nshells; // Height of each triangle
            double rscale = scale; // rscale = zscale for circular
            double zscale = scale;
            plane_files_vec.push_back(PlaneFiles(nshells, axis_r, axis_z, rscale, zscale, GridCreateOpt::Circular));
            
            if(is_stellarator){
                // Push back a copy of the grid since stellarator needs two
                plane_files_vec.push_back(PlaneFiles(nshells, axis_r, axis_z, rscale, zscale, GridCreateOpt::Circular));
            }
        }else{
            nlr.use_namelist("sml_param");
            std::string node_file = nlr.get<std::string>("sml_node_file","example_file.node"); // The grid file
            std::string ele_file = nlr.get<std::string>("sml_ele_file","example_file.ele"); // The element file
            std::string surf_file = nlr.get<std::string>("sml_surf_file","example_file.flx.aif"); // The flux surface file
            std::string input_file_dir = nlr.get<std::string>("sml_input_file_dir","./"); // Directory containing the grid files

            if(is_stellarator){
                surf_file = nlr.get<std::string>("sml_surf_file","example_file.grid"); // The surface file for stellarators
                // Add path to file names
                // Remove the extensions ".node" and ".ele" since we have the add plane numbers
                node_file = input_file_dir + "/equilibrium/" + node_file.substr(0, node_file.size() - 5);
                ele_file = input_file_dir + "/equilibrium/" + ele_file.substr(0, ele_file.size() - 4);
                // At the moment the surf_file (.grid) is identical for all stellarator planes
                surf_file = input_file_dir + "/equilibrium/" + surf_file;

                int iplane = my_mpi.my_intpl_rank;
                int ifile = iplane*2 + 1; // Midplane is .1, Edge is .2, next midplane is .3, etc.

                size_t n_char = STELLARATOR_PLANE_LIMIT;
                if(std::to_string(ifile + 1).length() > n_char)
                    exit_XGC("Error, plane number " + std::to_string(ifile + 1) + " is larger than the maximum number of allowed planes.\n");

                std::string ifile_str = std::to_string(ifile);
                ifile_str = std::string(n_char - std::min(n_char, ifile_str.length()), '0') + ifile_str;

                std::string midplane_node_file = node_file + "_" + ifile_str + ".node";
                std::string midplane_ele_file  = ele_file  + "_" + ifile_str + ".ele";
                std::string midplane_surf_file = surf_file;

                // Read in midplane grid files (one rank per plane reads)
                plane_files_vec.push_back(PlaneFiles(midplane_node_file, midplane_ele_file, midplane_surf_file, is_stellarator, my_mpi.plane_comm));

                ifile += 1; // shift to edge plane
                ifile_str = std::to_string(ifile);
                ifile_str = std::string(n_char - std::min(n_char, ifile_str.length()), '0') + ifile_str;
                node_file = node_file + "_" + ifile_str + ".node";
                ele_file = ele_file  + "_" + ifile_str + ".ele";

                // Read in edgeplane grid files (one rank per plane reads)
                plane_files_vec.push_back(PlaneFiles(node_file, ele_file, surf_file, is_stellarator, my_mpi.plane_comm));

                // Send right edgeplane to next MPI plane
                plane_files_vec.push_back(PlaneFiles(my_mpi, plane_files_vec[1]));
            }else{
                // Add path to file names
                node_file = input_file_dir + "/" + node_file;
                ele_file = input_file_dir + "/" + ele_file;
                surf_file = input_file_dir + "/" + surf_file;

                // Read in grid files of single plane
                plane_files_vec.push_back(PlaneFiles(node_file, ele_file, surf_file, is_stellarator, SML_COMM_WORLD));
            }
        }
    }
#endif

    int n_planes() const{
        return plane_files_vec.size();
    }

    void clear(){
        plane_files_vec.clear();
    }

    // Returns the number of shells needed to create a grid with at least nnodes nodes.
    static int nshells_from_nnodes(int nnodes){
        return PlaneFiles::nshells_from_nnodes(nnodes);
    }

};

#endif