ascot5_main.c

Go to the documentation of this file.

 
#include <getopt.h>
#include <math.h>
#include <omp.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "ascot5.h"
#include "consts.h"
#include "math.h"
#include "wall.h"
#include "diag.h"
#include "B_field.h"
#include "plasma.h"
#include "print.h"
#include "simulate.h"
#include "particle.h"
#include "endcond.h"
#include "hdf5_interface.h"
#include "offload.h"
#include "gitver.h"
#include "mpi_interface.h"
 
#include "ascot5_main.h"
 
#ifdef TRAP_FPE
#include <fenv.h>
#endif
 
int read_arguments(int argc, char** argv, sim_offload_data* sim);
 
int main(int argc, char** argv) {
 
#ifdef TRAP_FPE
    /* This will raise floating point exceptions */
    feenableexcept(FE_DIVBYZERO| FE_INVALID | FE_OVERFLOW);
    /*
     * If you are hunting a specific exception, you can disable the exceptions
     * in other parts of the code by surrounding it with: */
    /*
     * fedisableexcept(FE_DIVBYZERO | FE_INVALID | FE_OVERFLOW);
     *  --- your  code here ---
     * feenableexcept(FE_DIVBYZERO  | FE_INVALID | FE_OVERFLOW);
     *
     * */
#endif
 
    /* Read and parse command line arguments */
    sim_offload_data sim;
    if( read_arguments(argc, argv, &sim) ) {
        abort();
        return 1;
    }
 
    if(sim.mpi_size > 0) {
        /* This is a pseudo-mpi run, where rank and size were set on the command
         * line. Only set root equal to rank since there are no other processes
         */
        sim.mpi_root = sim.mpi_rank;
    }
    else {
        /* Init MPI if used, or run serial */
        int mpi_rank, mpi_size, mpi_root;
        mpi_interface_init(argc, argv, &mpi_rank, &mpi_size, &mpi_root);
        sim.mpi_rank = mpi_rank;
        sim.mpi_size = mpi_size;
        sim.mpi_root = mpi_root;
    }
 
    print_out0(VERBOSE_MINIMAL, sim.mpi_rank, sim.mpi_root, "ASCOT5_MAIN\n");
    print_out0(VERBOSE_MINIMAL, sim.mpi_rank, sim.mpi_root,
               "Tag %s\nBranch %s\n\n", GIT_VERSION, GIT_BRANCH);
    print_out(VERBOSE_NORMAL, "Initialized MPI, rank %d, size %d.\n",
              sim.mpi_rank, sim.mpi_size);
 
    /* Total number of markers to be simulated */
    int n_tot;
    /* Marker input struct */
    input_particle* p;
 
    /* Offload data arrays that are allocated when input is read */
    real* B_offload_array;
    real* E_offload_array;
    real* plasma_offload_array;
    real* neutral_offload_array;
    real* wall_offload_array;
    int* wall_int_offload_array;
    real* boozer_offload_array;
    real* mhd_offload_array;
    real* asigma_offload_array;
    real* diag_offload_array;
 
    /* Read input from the HDF5 file */
    if( hdf5_interface_read_input(&sim,
                                  hdf5_input_options | hdf5_input_bfield |
                                  hdf5_input_efield  | hdf5_input_plasma |
                                  hdf5_input_neutral | hdf5_input_wall   |
                                  hdf5_input_marker  | hdf5_input_boozer |
                                  hdf5_input_mhd     | hdf5_input_asigma,
                                  &B_offload_array, &E_offload_array,
                                  &plasma_offload_array, &neutral_offload_array,
                                  &wall_offload_array,  &wall_int_offload_array,
                                  &boozer_offload_array, &mhd_offload_array,
                                  &asigma_offload_array, NULL,
                                  &p, &n_tot) ) {
        print_out0(VERBOSE_MINIMAL, sim.mpi_rank, sim.mpi_root,
                   "\nInput reading or initializing failed.\n"
                   "See stderr for details.\n");
        mpi_interface_finalize();
        abort();
        return 1;
    };
 
    /* Initialize marker states array ps and free marker input p */
    int n_proc; /* Number of markers allocated for this MPI process */
    particle_state* ps;
    if( prepare_markers(&sim, n_tot, p, &ps, &n_proc, B_offload_array) ) {
        goto CLEANUP_FAILURE;
    }
 
    /* Combine input offload arrays to one */
    offload_package offload_data;
    real* offload_array;
    int* int_offload_array;
    if( pack_offload_array(
            &sim, &offload_data, &B_offload_array, &E_offload_array,
            &plasma_offload_array, &neutral_offload_array, &wall_offload_array,
            &wall_int_offload_array, &boozer_offload_array, &mhd_offload_array,
            &asigma_offload_array, &offload_array, &int_offload_array) ) {
        goto CLEANUP_FAILURE;
    }
 
    /* Initialize diagnostics offload data */
    diag_init_offload(&sim.diag_offload_data, &diag_offload_array, n_tot);
    print_out0(VERBOSE_NORMAL, sim.mpi_rank, sim.mpi_root,
               "Initialized diagnostics, %.1f MB.\n",
               sim.diag_offload_data.offload_array_length * sizeof(real)
               / (1024.0*1024.0));
 
    /* Write run group and inistate */
    char qid[11];
    hdf5_generate_qid(qid);
    if( write_rungroup(&sim, ps, n_tot, qid) ) {
        goto CLEANUP_FAILURE;
    }
 
    /* Actual simulation is done here; the results are stored in
     * diag_offload_array and pout contains end states from all processes.
     * n_gathered is the number of markers in the output array, which is n_tot
     * for most cases except when the simulation is run in condor-like manner,
     * in which case it is equal to n_proc. */
    int n_gathered;
    particle_state* pout;
    offload_and_simulate(
        &sim, n_tot, n_proc, ps, &offload_data, offload_array,
        int_offload_array, &n_gathered, &pout, diag_offload_array);
 
    /* Free input data */
    offload_free_offload(&offload_data, &offload_array, &int_offload_array);
 
    /* Write output and clean */
    if( write_output(&sim, pout, n_gathered, diag_offload_array) ) {
        goto CLEANUP_FAILURE;
    }
    diag_free_offload(&sim.diag_offload_data, &diag_offload_array);
 
    /* Display marker summary and free marker arrays */
    if(sim.mpi_rank == sim.mpi_root) {
        print_marker_summary(pout, n_tot);
    }
    free(pout);
 
    print_out0(VERBOSE_MINIMAL, sim.mpi_rank, sim.mpi_root, "\nDone.\n");
    mpi_interface_finalize();
    return 0;
 
/* GOTO this block to free resources in case simulation crashes */
CLEANUP_FAILURE:
    mpi_interface_finalize();
    free(p);
    free(ps);
    free(pout);
    free(offload_array);
    free(int_offload_array);
    free(diag_offload_array);
    abort();
    return 1;
}
 
int prepare_markers(
    sim_offload_data* sim, int n_tot, input_particle* pin,
    particle_state** pout, int* n_proc, real* B_offload_array) {
 
    /* This sets up GC transformation order etc. so it must be called before
     * initializing markers. */
    simulate_init_offload(sim);
 
    /* Choose which markers are used in this MPI process. Simply put, markers
     * are divided into mpi_size sequential blocks and the mpi_rank:th block
     * is chosen for this simulation. */
    int start_index;
    mpi_my_particles(&start_index, n_proc, n_tot, sim->mpi_rank, sim->mpi_size);
    pin += start_index;
 
    /* Set up particlestates on host, needs magnetic field evaluation */
    print_out0(VERBOSE_NORMAL, sim->mpi_rank, sim->mpi_root,
               "\nInitializing marker states.\n");
    B_field_data Bdata;
    B_field_init(&Bdata, &sim->B_offload_data, B_offload_array);
 
    *pout = (particle_state*) malloc(*n_proc * sizeof(particle_state));
    for(int i = 0; i < *n_proc; i++) {
        particle_input_to_state(&(pin[i]), &((*pout)[i]), &Bdata);
    }
 
    print_out0(VERBOSE_NORMAL, sim->mpi_rank, sim->mpi_root,
               "Estimated memory usage %.1f MB.\n",
               (sizeof(real) * (*n_proc)) / (1024.0*1024.0));
    print_out0(VERBOSE_NORMAL, sim->mpi_rank, sim->mpi_root,
               "Marker states initialized.\n");
 
    /* We can now free the input markers */
    free(pin-start_index);
 
    return 0;
}
 
 
int pack_offload_array(
    sim_offload_data* sim, offload_package* offload_data,
    real** B_offload_array, real** E_offload_array, real** plasma_offload_array,
    real** neutral_offload_array, real** wall_offload_array,
    int** wall_int_offload_array, real** boozer_offload_array,
    real** mhd_offload_array, real** asigma_offload_array, real** offload_array,
    int** int_offload_array) {
 
    /* Pack offload data into single array and free individual offload arrays */
    offload_init_offload(offload_data, offload_array, int_offload_array);
 
    offload_pack(offload_data, offload_array, *B_offload_array,
                 sim->B_offload_data.offload_array_length,
                 int_offload_array, NULL, 0);
    B_field_free_offload(&sim->B_offload_data, B_offload_array);
 
    offload_pack(offload_data, offload_array, *E_offload_array,
                 sim->E_offload_data.offload_array_length,
                 int_offload_array, NULL, 0);
    E_field_free_offload(&sim->E_offload_data, E_offload_array);
 
    offload_pack(offload_data, offload_array, *plasma_offload_array,
                 sim->plasma_offload_data.offload_array_length,
                 int_offload_array, NULL, 0);
    plasma_free_offload(&sim->plasma_offload_data, plasma_offload_array);
 
    offload_pack(offload_data, offload_array, *neutral_offload_array,
                 sim->neutral_offload_data.offload_array_length,
                 int_offload_array, NULL, 0);
    neutral_free_offload(&sim->neutral_offload_data, neutral_offload_array);
 
    offload_pack(offload_data, offload_array, *wall_offload_array,
                 sim->wall_offload_data.offload_array_length,
                 int_offload_array, *wall_int_offload_array,
                 sim->wall_offload_data.int_offload_array_length);
    wall_free_offload(&sim->wall_offload_data, wall_offload_array,
                      wall_int_offload_array);
 
    offload_pack(offload_data, offload_array, *boozer_offload_array,
                 sim->boozer_offload_data.offload_array_length,
                 int_offload_array, NULL, 0);
    boozer_free_offload(&sim->boozer_offload_data, boozer_offload_array);
 
    offload_pack(offload_data, offload_array, *mhd_offload_array,
                 sim->mhd_offload_data.offload_array_length,
                 int_offload_array, NULL, 0);
    mhd_free_offload(&sim->mhd_offload_data, mhd_offload_array);
 
    offload_pack(offload_data, offload_array, *asigma_offload_array,
                 sim->asigma_offload_data.offload_array_length,
                 int_offload_array, NULL, 0);
    asigma_free_offload(&sim->asigma_offload_data, asigma_offload_array);
 
    return 0;
}
 
 
int write_rungroup(sim_offload_data* sim, particle_state* ps, int n_tot,
                   char* qid) {
 
    if(sim->mpi_rank == sim->mpi_root) {
        /* Initialize results group in the output file */
        print_out0(VERBOSE_IO, sim->mpi_rank, sim->mpi_root,
                   "\nPreparing output.\n");
        if( hdf5_interface_init_results(sim, qid, "run") ) {
            print_out0(VERBOSE_MINIMAL, sim->mpi_rank, sim->mpi_root,
                       "\nInitializing output failed.\n"
                       "See stderr for details.\n");
            /* Free offload data and terminate */
            return 1;
        }
        strcpy(sim->qid, qid);
    }
 
    /* Gather particle states so that we can write inistate */
    int n_gather;
    particle_state* ps_gather;
    mpi_gather_particlestate(ps, &ps_gather, &n_gather, n_tot,
                             sim->mpi_rank, sim->mpi_size, sim->mpi_root);
 
    if(sim->mpi_rank == sim->mpi_root) {
        /* Write inistate */
        if(hdf5_interface_write_state(
            sim->hdf5_out, "inistate", n_gather, ps_gather)) {
            print_out0(VERBOSE_MINIMAL, sim->mpi_rank, sim->mpi_root,
                       "\n"
                       "Writing inistate failed.\n"
                       "See stderr for details.\n"
                       "\n");
            /* Free offload data and terminate */
            return 1;
        }
        print_out0(VERBOSE_NORMAL, sim->mpi_rank, sim->mpi_root,
                   "\nInistate written.\n");
    }
    free(ps_gather);
 
    return 0;
}
 
 
int offload_and_simulate(
    sim_offload_data* sim, int n_tot, int n_proc, particle_state* pin,
    offload_package* offload_data, real* offload_array, int* int_offload_array,
    int* n_gather, particle_state** pout, real* diag_offload_array) {
 
    /* Empty message buffer before proceeding to actual simulation */
    fflush(stdout);
 
    /* Actual marker simulation happens here. */
    real t_sim_start = omp_get_wtime();
    simulate(0, n_proc, pin, sim, offload_data,
        offload_array, int_offload_array, diag_offload_array);
 
    mpi_interface_barrier();
    real t_sim_end = omp_get_wtime();
    print_out0(VERBOSE_NORMAL, sim->mpi_rank, sim->mpi_root,
        "Simulation finished in %lf s\n", t_sim_end-t_sim_start);
 
    /* Gather output data */
    mpi_gather_particlestate(pin, pout, n_gather, n_tot, sim->mpi_rank,
                             sim->mpi_size, sim->mpi_root);
    free(pin);
 
    mpi_gather_diag(&sim->diag_offload_data, diag_offload_array, n_tot,
                    sim->mpi_rank, sim->mpi_size, sim->mpi_root);
    return 0;
}
 
 
int write_output(sim_offload_data* sim, particle_state* ps, int n_tot,
                 real* diag_offload_array){
 
    if(sim->mpi_rank == sim->mpi_root) {
        /* Write endstate */
        if( hdf5_interface_write_state(
                sim->hdf5_out, "endstate", n_tot, ps)) {
            print_out0(VERBOSE_MINIMAL, sim->mpi_rank, sim->mpi_root,
                   "\nWriting endstate failed.\n"
                   "See stderr for details.\n");
            return 1;
        }
        print_out0(VERBOSE_NORMAL, sim->mpi_rank, sim->mpi_root,
                   "Endstate written.\n");
    }
 
    /* Combine diagnostic data and write it to HDF5 file */
    print_out0(VERBOSE_MINIMAL, sim->mpi_rank, sim->mpi_root,
                   "\nCombining and writing diagnostics.\n");
    int err_writediag = 0;
 
    if(sim->mpi_rank == sim->mpi_root) {
        err_writediag = hdf5_interface_write_diagnostics(sim,
            diag_offload_array, sim->hdf5_out);
    }
    if(err_writediag) {
        print_out0(VERBOSE_MINIMAL, sim->mpi_rank, sim->mpi_root,
                   "\nWriting diagnostics failed.\n"
                   "See stderr for details.\n");
        return 1;
    }
    else {
        print_out0(VERBOSE_MINIMAL, sim->mpi_rank, sim->mpi_root,
                   "Diagnostics written.\n");
    }
 
    return 0;
}
 
 
int read_arguments(int argc, char** argv, sim_offload_data* sim) {
    struct option longopts[] = {
        {"in", required_argument, 0, 1},
        {"out", required_argument, 0, 2},
        {"mpi_size", required_argument, 0, 3},
        {"mpi_rank", required_argument, 0, 4},
        {"d", required_argument, 0, 5},
        {"options", required_argument, 0, 6},
        {"bfield",  required_argument, 0, 7},
        {"efield",  required_argument, 0, 8},
        {"marker",  required_argument, 0, 9},
        {"wall",    required_argument, 0, 10},
        {"plasma",  required_argument, 0, 11},
        {"neutral", required_argument, 0, 12},
        {"boozer",  required_argument, 0, 13},
        {"mhd",     required_argument, 0, 14},
        {"asigma",  required_argument, 0, 15},
        {0, 0, 0, 0}
    };
 
    /* Initialize default values */
    sim->hdf5_in[0]     = '\0';
    sim->hdf5_out[0]    = '\0';
    sim->mpi_rank       = 0;
    sim->mpi_size       = 0;
    strcpy(sim->description, "No description.");
    sim->qid_options[0] = '\0';
    sim->qid_bfield[0]  = '\0';
    sim->qid_efield[0]  = '\0';
    sim->qid_marker[0]  = '\0';
    sim->qid_wall[0]    = '\0';
    sim->qid_plasma[0]  = '\0';
    sim->qid_neutral[0] = '\0';
    sim->qid_boozer[0]  = '\0';
    sim->qid_mhd[0]     = '\0';
    sim->qid_asigma[0]  = '\0';
    sim->qid_nbi[0]     = '\0';
 
    /* Read user input */
    int c;
    int slen;  /* String length */
    while((c = getopt_long(argc, argv, "", longopts, NULL)) != -1) {
        switch(c) {
            case 1:
                /* The .hdf5 filename can be specified with or without the
                 * trailing .h5 */
                slen = strlen(optarg);
                if ( slen > 3 && !strcmp(optarg+slen-3, ".h5") ) {
                    strcpy(sim->hdf5_in, optarg);
                    (sim->hdf5_in)[slen-3]='\0';
                }
                else {
                    strcpy(sim->hdf5_in, optarg);
                }
                break;
            case 2:
                /* The .hdf5 filename can be specified with or without the
                 * trailing .h5 */
                slen = strlen(optarg);
                if ( slen > 3 && !strcmp(optarg+slen-3, ".h5") ) {
                    strcpy(sim->hdf5_out, optarg);
                    (sim->hdf5_out)[slen-3]='\0';
                }
                else {
                    strcpy(sim->hdf5_out, optarg);
                }
                break;
            case 3:
                sim->mpi_size = atoi(optarg);
                break;
            case 4:
                sim->mpi_rank = atoi(optarg);
                break;
            case 5:
                strcpy(sim->description, optarg);
                break;
            case 6:
                strcpy(sim->qid_options, optarg);
                break;
            case 7:
                strcpy(sim->qid_bfield, optarg);
                break;
            case 8:
                strcpy(sim->qid_efield, optarg);
                break;
            case 9:
                strcpy(sim->qid_marker, optarg);
                break;
            case 10:
                strcpy(sim->qid_wall, optarg);
                break;
            case 11:
                strcpy(sim->qid_plasma, optarg);
                break;
            case 12:
                strcpy(sim->qid_neutral, optarg);
                break;
            case 13:
                strcpy(sim->qid_boozer, optarg);
                break;
            case 14:
                strcpy(sim->qid_mhd, optarg);
                break;
            case 15:
                strcpy(sim->qid_asigma, optarg);
                break;
            default:
                // Unregonizable argument(s). Tell user how to run ascot5_main
                print_out(VERBOSE_MINIMAL,
                          "\nUnrecognized argument. The valid arguments are:\n");
                print_out(VERBOSE_MINIMAL,
                          "--in input file (default: ascot.h5)\n");
                print_out(VERBOSE_MINIMAL,
                          "--out output file (default: same as input)\n");
                print_out(VERBOSE_MINIMAL,
                          "--mpi_size number of independent processes\n");
                print_out(VERBOSE_MINIMAL,
                          "--mpi_rank rank of independent process\n");
                print_out(VERBOSE_MINIMAL,
                          "--d run description maximum of 250 characters\n");
                return 1;
        }
    }
 
    /* Default value for input file is ascot.h5, and for output same as input
     * file. Adujust hdf5_in and hdf5_out accordingly. For output file, we don't
     * add the .h5 extension here. */
    if(sim->hdf5_in[0] == '\0' && sim->hdf5_out[0] == '\0') {
        /* No input, use default values for both */
        strcpy(sim->hdf5_in,  "ascot.h5");
        strcpy(sim->hdf5_out, "ascot.h5");
    }
    else if(sim->hdf5_in[0] == '\0' && sim->hdf5_out[0] != '\0') {
        /* Output file is given but the input file is not */
        strcpy(sim->hdf5_in,  "ascot.h5");
        strcat(sim->hdf5_out, ".h5");
    }
    else if(sim->hdf5_in[0] != '\0' && sim->hdf5_out[0] == '\0') {
        /* Input file is given but the output is not */
        strcat(sim->hdf5_in, ".h5");
        strcpy(sim->hdf5_out, sim->hdf5_in);
    }
    else {
        /* Both input and output files are given */
        strcat(sim->hdf5_in,  ".h5");
        strcat(sim->hdf5_out, ".h5");
    }
    return 0;
}
 
void print_marker_summary(particle_state* ps, int n_tot) {
 
    print_out(VERBOSE_MINIMAL, "\nSummary of results:\n");
 
    /* Temporary arrays that are needed to store unique end conditions and
     * errors. We can have at most n_tot different values. */
    int* temp = (int*)malloc(n_tot*sizeof(int));
    int* unique = (int*)malloc(n_tot*sizeof(int));
    int* count = (int*)malloc((n_tot+1)*sizeof(int));
 
    /* First we find out what the end conditions are */
    for(int i=0; i<n_tot; i++) {
        temp[i] = ps[i].endcond;
    }
    math_uniquecount(temp, unique, count, n_tot);
    count[n_tot] = 0;/* This ensures the following while loops are terminated.*/
 
    /* Print the end conditions, if marker has multiple end conditions, separate
     * those with "and". */
    int i = 0;
    while(count[i] > 0) {
        /* Initialize */
        int endconds[32];
        for(int j=0; j<32;j++) {
            endconds[j] = 0;
        }
        char endcondstr[256];
        endcondstr[0] = '\0';
 
        /* Represent all end conditions with a single string and print it */
        int j = 0;
        endcond_parse(unique[i], endconds);
        while(endconds[j]) {
            if(j>0) {
                strcat(endcondstr, " and ");
            }
            char temp[256];
            endcond_parse2str(endconds[j], temp);
            strcat(endcondstr, temp);
            j++;
        }
        if(j == 0) {
            sprintf(endcondstr, "Aborted");
        }
        print_out(VERBOSE_MINIMAL, "%9d markers had end condition %s\n",
                  count[i], endcondstr);
        i++;
    }
 
    /* Empty line between end conditions and errors */
    print_out(VERBOSE_MINIMAL, "\n");
 
    /* Find all errors */
    for(int i=0; i<n_tot; i++) {
        temp[i] = (int)(ps[i].err);
    }
    math_uniquecount(temp, unique, count, n_tot);
 
    /* Go through each unique error and represent is a string */
    i = 0;
    while(count[i] > 0) {
        if(unique[i] == 0) {
            /* Value 0 indicates no error occurred so skip that */
            i++;
            continue;
        }
        char msg[256];
        char line[256];
        char file[256];
        error_parse2str(unique[i], msg, line, file);
        print_out(VERBOSE_MINIMAL,
                  "%9d markers were aborted with an error message:\n"
                  "          %s\n"
                  "          at line %s in %s\n",
                  count[i], msg, line, file);
        i++;
    }
 
    /* If count[0] equals to number of markers and their error field is zero,
     * we have no markers that were aborted. */
    if(count[0] == n_tot && unique[0] == 0) {
        print_out(VERBOSE_MINIMAL,
                  "          No markers were aborted.\n");
    }
 
    /* Free temporary arrays */
    free(temp);
    free(unique);
    free(count);
}