grid_files.hpp

Go to the documentation of this file.

#ifndef GRID_FILES_HPP
#define GRID_FILES_HPP

#include <vector>
#ifdef USE_MPI
#include "my_mpi.hpp"
#endif
#include "file_reader.hpp"
#include "magnetic_field.hpp"
#include "grid_structs.hpp"
#include "grid_setup.hpp"


struct GridFiles{
    // Contents of .node
    int nnodes;
    std::vector<RZPair> x;
    std::vector<int> rgn;

    // Contents of .ele
    int ntriangles;
    std::vector<Vertex> nodes;

    // Contents of .flx.aif
    int nsurfaces;
    std::vector<int> nnodes_on_surface;
    int sum_nnodes_on_surfaces;
    std::vector<int> nodes_of_surfaces;

    // Read the files
    GridFiles(const std::string& node_file, const std::string& element_file, const std::string& flx_aif_file){
        // Only one MPI rank reads the files, then broadcasts to the others
        if(is_rank_zero()){

            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
            std::vector<double> tmp_nodes(nnodes*4); // 4 entries per node: index, r, z, region
            file_reader.read(tmp_nodes.data(), tmp_nodes.size());

            // Reformat to x and rgn
            x = std::vector<RZPair>(nnodes);
            rgn = std::vector<int>(nnodes);
            for(int i=0; i<nnodes; i++){
                // tmp_nodes[i*4 + 0]; <-- first column is ignored
                x[i].r = tmp_nodes[i*4 + 1];
                x[i].z = tmp_nodes[i*4 + 2];
                rgn[i] = tmp_nodes[i*4 + 3];
            }

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

            // Use dummy array since reformatting is needed
            std::vector<int> tmp_ele(ntriangles*4); // 4 entries per node: index, r, z, region

            file_reader.read(tmp_ele.data(), tmp_ele.size());

            // Reformat since first column is ignored
            nodes = std::vector<Vertex>(ntriangles);
            for(int i=0; i<ntriangles; i++){
                // tmp_ele[i*4 + 0]; <-- first column is ignored
                nodes[i].vertex[0] = tmp_ele[i*4 + 1];
                nodes[i].vertex[1] = tmp_ele[i*4 + 2];
                nodes[i].vertex[2] = tmp_ele[i*4 + 3];
            }

#ifndef NEW_FLX_AIF
            exit_XGC("\nError: C++ grid file reader can't read old flux files (i.e. NEW_FLX_AIF must be On)");
#endif
            file_reader.open(flx_aif_file);
            file_reader.skip(3); // Skip 3 unused values in file
            int nsurfaces1, nsurfaces2, nsurfaces3a, nsurfaces3b;
            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;

            file_reader.skip(2); // Skip 2 unused values in file

            // Read in nnodes on each surface
            nnodes_on_surface = std::vector<int>(nsurfaces);
            file_reader.read(nnodes_on_surface.data(), nnodes_on_surface.size());

            // Loop through surfaces to calculate sum total of nodes on surfaces
            sum_nnodes_on_surfaces = 0;
            for(int i=0; i<nsurfaces; i++) sum_nnodes_on_surfaces += nnodes_on_surface[i];

            // Allocate nodes_of_surfaces
            nodes_of_surfaces = std::vector<int>(sum_nnodes_on_surfaces);

            // Read in values
            file_reader.read(nodes_of_surfaces.data(), nodes_of_surfaces.size());
        }

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

        // First broadcast scalars so we know the sizes of the arrays
        std::vector<int> scalars(4);
        scalars[0] = nnodes;
        scalars[1] = ntriangles;
        scalars[2] = nsurfaces;
        scalars[3] = sum_nnodes_on_surfaces;

        // Broadcast file size to other processes
#ifdef USE_MPI
        int root_rank = 0;
        MPI_Bcast(scalars.data(), 4, MPI_INT, root_rank, SML_COMM_WORLD);
#endif

        if(!is_rank_zero()){
            // Set scalars on non-root ranks
            nnodes = scalars[0];
            ntriangles = scalars[1];
            nsurfaces = scalars[2];
            sum_nnodes_on_surfaces = scalars[3];

            // Allocate vectors on non-root ranks
            x.resize(nnodes);
            rgn.resize(nnodes);
            nodes.resize(ntriangles);
            nnodes_on_surface.resize(nsurfaces);
            nodes_of_surfaces.resize(sum_nnodes_on_surfaces);
        }

        // Broadcast grid arrays
#ifdef USE_MPI
        MPI_Bcast( x.data(), x.size()*2, MPI_DOUBLE, root_rank, SML_COMM_WORLD);
        MPI_Bcast( rgn.data(), rgn.size(), MPI_INT, root_rank, SML_COMM_WORLD);
        MPI_Bcast( nodes.data(), nodes.size()*3, MPI_INT, root_rank, SML_COMM_WORLD);
        MPI_Bcast( nnodes_on_surface.data(), nnodes_on_surface.size(), MPI_INT, root_rank, SML_COMM_WORLD);
        MPI_Bcast( nodes_of_surfaces.data(), nodes_of_surfaces.size(), MPI_INT, root_rank, SML_COMM_WORLD);
#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 };

        // Work from center outward
        x[0].r = raxis;
        x[0].z = zaxis;
        rgn[0] = 0; // not wall
        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

                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++;
        }

        // 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]++;
        }

        // nsurfaces is 0 because the hex grid is not surface-aligned
        nsurfaces=0;
    }

    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

        nnodes_on_surface[0] = 1;
        nodes_of_surfaces[0] = 0;
        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
                }

                // 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){
                    nodes_of_surfaces[inode] = inode;
                    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]++;
        }
    }


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

    GridFiles(MagneticField<HostType>& magnetic_field, int nshells, GridCreateOpt grid_opt=Hexagonal)
      : nnodes(1 + 3*(nshells+1)*nshells),
        x(nnodes),
        rgn(nnodes),
        ntriangles(6*nshells*nshells),
        nodes(ntriangles),
        nsurfaces(nshells+1),
        nnodes_on_surface(nsurfaces),
        sum_nnodes_on_surfaces(nnodes),
        nodes_of_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.bd_max_r - magnetic_field.equil.axis_r,
                                   magnetic_field.equil.axis_r - magnetic_field.bd_min_r);

        double height = 2*std::min(magnetic_field.bd_max_z - magnetic_field.equil.axis_z,
                                   magnetic_field.equil.axis_z - magnetic_field.bd_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);
        }
    }

    // Read from files
    GridFiles(const std::string node_file, const std::string ele_file){
    }
};

#endif