diagnose_poincare.f90 Source File


Source Code

program diagnose_poincare
    !! This program generates poincare data, based on a parameter file
    !! specified via the command line, and writes that data to NetCDF.

    use MPI
    !$ use omp_lib
    use netcdf
    use precision_m, only : FP, DP, NF90_FP, FP_NAN
    use constants_m, only : TWO_PI
    use screen_io_m, only : get_stdout
    use comm_handling_m, only: get_communicator
    use status_codes_m, only : PARALLAX_SUCCESS, PARALLAX_ERR_POINCARE
    use error_handling_m, only : handle_error, handle_error_netcdf, error_info_t
    use array_generation_m, only : linspace
    use coords_polar_m, only : polar_to_cart
    use fieldline_tracer_m, only : trace
    use equilibrium_m, only : equilibrium_t
    use equilibrium_factory_m, only : create_equilibrium, &
        get_equilibrium_identifier

    implicit none

    ! Default parameters
    character(len=64) :: geometry = 'CIRCULAR'
    !! Specifies equilibrium type
    character(len=256) :: equi_filename = ""
    !! Equilibrium parameter filename (if blank, the equilibrium will be
    !! initialized with its default parameters)
    integer :: n_planes = 1
    !! Number of toroidal planes
    integer :: n_turns = 50
    !! Number of toroidal rotations 2pi
    integer :: n_surfaces = 6
    !! Number of flux surfaces
    logical :: trace_start_loc_polar = .false.
    !! Switch whether the starting locations for tracing trace_start_loc are
    !! specified in (x, y) or (rho, theta) coordinates
    real(FP) :: phi_min = 0.0_FP
    !! Lower bound of default phi array
    real(FP) :: phi_max = TWO_PI
    !! Upper bound of default phi_array
    character(len=256) :: poincare_filename = "poincare_data.nc"
    !! File name under which to save the Poincare data

    real(FP), dimension(:), allocatable :: phi_array
    !! Toroidal angles phi of Poincare data. By default, this array is
    !! size(n_planes) and spans [phi_min, phi_max).
    ! NOTE: The phi values in phi_array must be given in either ascending or
    !       descending order, so that the trace direction always remains the
    !       same. Additionally, the phi values must lie within one toroidal
    !       turn.
    real(FP), dimension(:,:), allocatable :: trace_start_loc
    !! Trace starting locations for each surface, for the first plane. By
    !! default, the starting locations will span horizontally from the magnetic
    !! axis to either xmax (polar=false) or rhomax (polar=true). The starting
    !! locations for the subsequent planes will be calculated using tracing, so
    !! that the same flux surfaces are shown on all planes.

    namelist / params_poincare / geometry, equi_filename, n_planes, &
                                 n_surfaces, n_turns, trace_start_loc_polar, &
                                 phi_min, phi_max, poincare_filename

    namelist / params_poincare_phi_array / phi_array

    namelist / params_poincare_trace_start_loc / trace_start_loc

    class(equilibrium_t), allocatable :: equi

    integer :: ierror, nprocs
    integer :: n_threads
    integer :: argc
    integer :: eq_id
    character(len=256), allocatable, dimension(:) :: argv
    character(len=256) :: parameter_file, io_errmsg
    integer :: funit, io_error
    integer :: i
    real(FP) :: x, y

    ! Initialize MPI. This program is currently only designed for one process,
    ! but we need to init MPI for PARALLAX internal functionality to work.
    call MPI_init(ierror)
    call MPI_comm_size(get_communicator(), nprocs, ierror)
    if (nprocs > 1) then
        call handle_error("Poincare not designed for more than one MPI &
                          &process!", PARALLAX_ERR_POINCARE, &
                          __LINE__, __FILE__)
    end if

    !$omp parallel private(n_threads)
    !$ n_threads = omp_get_num_threads()
    !$omp master
    !$ write(get_stdout(), "(A, I3)") 'OpenMP on, n_threads = ', n_threads
    !$omp end master
    !$omp end parallel

    ! Determine parameter file given by command line argument
    argc = command_argument_count()
    if(argc /= 1) then
        ! The only command line argument currently is the parameter file
        call print_usage()
        stop
    endif

    ! Retrieve all arguments
    allocate(argv(argc))
    do i = 1, argc
        call get_command_argument(i, argv(i))
    enddo

    ! Get the parameter file from the last command line argument
    if(argv(argc)(1:1) == "-") then
        ! The parameter file may not be an optional variable
        call print_usage()
        stop
    endif
    parameter_file = argv(argc)

    ! Read parameters from file
    open(newunit=funit, file=parameter_file, status="old", &
         action="read", iostat=io_error, iomsg=io_errmsg)
    if (io_error /= 0) then
        call handle_error(io_errmsg, PARALLAX_ERR_POINCARE, __LINE__, &
                          __FILE__, additional_info=&
                                  error_info_t("IO error status: ", [io_error]))
    end if

    read(funit, nml=params_poincare, iostat=io_error, iomsg=io_errmsg)
    if (io_error /= 0) then
        call handle_error(io_errmsg, PARALLAX_ERR_POINCARE, __LINE__, &
                          __FILE__, additional_info=&
                                  error_info_t("IO error status: ", [io_error]))
    end if

    ! Allocate equilibrium
    eq_id = get_equilibrium_identifier(geometry)
    if (equi_filename /= "") then
        call create_equilibrium(equi, eq_id, equi_filename)
    else
        call create_equilibrium(equi, eq_id)
    end if

    ! Default phi array. In case n_planes = 1, [phi_min] is returned
    phi_array = linspace(phi_min, phi_max, n_planes)

    ! Read phi_array
    read(funit, nml=params_poincare_phi_array, iostat=io_error, iomsg=io_errmsg)
    if (io_error /= 0) then
        call handle_error(io_errmsg, PARALLAX_ERR_POINCARE, __LINE__, &
                          __FILE__, additional_info=&
                                  error_info_t("IO error status: ", [io_error]))
    end if

    ! Default trace starting locations span horizontally from magnetic axis to
    ! xmax or rhomax
    allocate(trace_start_loc(1:2, 1:n_surfaces))
    if (trace_start_loc_polar) then
        trace_start_loc(1, :) = linspace(0.0_FP, equi%rhomax, n_surfaces, &
                                         endpoint=.true.)
        trace_start_loc(2, :) = 0.0_FP
    else
        trace_start_loc(1, :) = linspace(equi%x0, equi%xmax, n_surfaces, &
                                         endpoint=.true.)
        trace_start_loc(2, :) = equi%y0
    end if

    ! Read trace_start_loc
    read(funit, nml=params_poincare_trace_start_loc, iostat=io_error, &
         iomsg=io_errmsg)
    if (io_error /= 0) then
        call handle_error(io_errmsg, PARALLAX_ERR_POINCARE, __LINE__, &
                          __FILE__, additional_info=&
                                  error_info_t("IO error status: ", [io_error]))
    end if

    ! If trace_start_loc are specified in (rho, theta), convert to (x, y) for
    ! generate routine
    if (trace_start_loc_polar) then
        do i = 1, n_surfaces
            call polar_to_cart(equi, trace_start_loc(1, i), &
                               trace_start_loc(2, i), phi_array(1), x, y)
            trace_start_loc(1, i) = x
            trace_start_loc(2, i) = y
        end do
    end if

    call generate_poincare_data(equi, n_surfaces, n_turns, n_planes, &
                                phi_array, trace_start_loc, &
                                trim(poincare_filename))

contains

    subroutine print_usage()
        !! Prints the intended usage of the executable
        write(get_stdout(), "(A)") &
            "usage: diagnose_poincare [-h] PARAMETER_FILE"
        write(get_stdout(), *) ""
        write(get_stdout(), "(A)") "positional arguments:"
        write(get_stdout(), "(2X, A, T24, A)") &
            "PARAMETER_FILE", "path to Poincare parameter file"
        write(get_stdout(), *) ""
        write(get_stdout(), "(A)") "optional arguments:"
        write(get_stdout(), "(2X, A, T24, A)") &
            "-h, --help", "show this help message and exit"
    end subroutine

    subroutine generate_poincare_data(equi, n_surfaces, n_turns, n_planes, &
                                      phi_array, trace_start_loc, filename)
        !! For a given equilibrium `equi`, this routine generates Poincare data
        !! on the toroidal planes with angles given in `phi_array`.
        !!
        !! This is done by tracing `n_surfaces` number of magnetic field lines
        !! `n_turns` number of full toroidal turns 2pi around the device. The
        !! starting locations of each field line, for the first plane, are
        !! given in `trace_start_loc`.
        !!
        !! Each time the trace intersects one of the specified planes, the x
        !! and y locations of the field line are saved. This data is then saved
        !! to netcdf file with the name `filename`.
        class(equilibrium_t), intent(inout) :: equi
        !! Equilibrium instance
        integer, intent(in) :: n_surfaces
        !! Number of flux surfaces
        integer, intent(in) :: n_turns
        !! Number of toroidal rotations 2pi
        integer, intent(in) :: n_planes
        !! Number of toroidal planes
        real(FP), intent(in), dimension(n_planes) :: phi_array
        !! Toroidal angles at which to save data
        real(FP), intent(in), dimension(2, n_surfaces) :: trace_start_loc
        !! Trace starting locations (x, y) for each surface, for the first plane
        character(len=*), intent(in) :: filename
        !! File name under which to save the Poincare data

        real(FP), dimension(2, n_planes, n_turns + 1, n_surfaces) :: &
                                                                  poincare_data

        real(FP), dimension(n_planes) :: dphi_array
        real(FP) :: x_init, y_init, x_end, y_end
        integer :: i_surf, i_turn, i_plane
        real(DP) :: start_wtime, dt
        integer :: surface_trace_status, i_plane_error, i_turn_error
        logical :: error_printed
        character(len=:), allocatable :: fmt

        ! Calculate toroidal distance between given planes
        dphi_array(1:n_planes - 1) = phi_array(2:n_planes) &
                                   - phi_array(1:n_planes - 1)

        ! Ensure trace direction is always either positive or negative
        if (.not. all(dphi_array(1:n_planes - 1) > 0.0_FP) .and. &
            .not. all(dphi_array(1:n_planes - 1) < 0.0_FP)) then
            call handle_error("The phi_array must be given in either &
                              &ascending or descending order!", &
                              PARALLAX_ERR_POINCARE, __LINE__, __FILE__)
        end if

        ! Ensure all phi planes occur within one toroidal turn
        if (abs(sum(dphi_array(1:n_planes - 1))) >= TWO_PI) then
            call handle_error("The phi_array must span less than one &
                              &toroidal turn!", PARALLAX_ERR_POINCARE, &
                              __LINE__, __FILE__)
        endif

        ! Calculate toroidal distance required to complete a full toroidal turn
        ! 2pi, to wrap from the last plane back to the first. The sign of this
        ! final segment should match the sign of the rest of the dphi_array.
        dphi_array(n_planes) = sign(TWO_PI, dphi_array(1)) &
                             - sum(dphi_array(1:n_planes - 1))

        start_wtime = MPI_wtime()
        !$omp parallel do default(none) &
        !$omp private(i_surf, i_turn, i_plane, x_init, y_init, x_end, y_end, &
        !$omp         surface_trace_status, error_printed, fmt, &
        !$omp         i_plane_error, i_turn_error) &
        !$omp firstprivate(equi, n_surfaces, n_turns, n_planes, phi_array, &
        !$omp              dphi_array, trace_start_loc) &
        !$omp shared(poincare_data)
        do i_surf = 1, n_surfaces
            ! The first data point for each surface is the initial x, y position
            poincare_data(1, 1, 1, i_surf) = trace_start_loc(1, i_surf)
            poincare_data(2, 1, 1, i_surf) = trace_start_loc(2, i_surf)

            write(get_stdout(), "(A,I3.3,A,I3.3,A,F6.3,A,F6.3)") "Surface ", &
                i_surf, "/", n_surfaces, &
                ": x_init = ", poincare_data(1, 1, 1, i_surf), &
                ", y_init = ", poincare_data(2, 1, 1, i_surf)

            surface_trace_status = PARALLAX_SUCCESS
            error_printed = .false.

            do i_turn = 1, n_turns + 1
            do i_plane = 1, n_planes
                if (surface_trace_status /= PARALLAX_SUCCESS) then

                    ! Trace error occured in the previous call. Print message
                    ! once, then skip all remaining trace calls for this surface
                    if (.not. error_printed) then
                        if (i_plane == 1) then
                            i_plane_error = n_planes
                            i_turn_error = i_turn - 1
                        else
                            i_plane_error = n_planes - 1
                            i_turn_error = i_turn
                        endif

                        fmt = "(A,I3,A,I3,A,I3,A,I3,A,I3,A,I3,A)"
                        write(get_stdout(), fmt) "Trace error encountered on &
                            &surface = ", i_surf, " / ", n_surfaces, &
                            ", turn = ", i_turn_error, " / ", n_turns + 1, &
                            ", plane = ", i_plane_error, " / ", n_planes, "! &
                            &Remaining R and Z values for this surface will be &
                            &set to NaN."

                        error_printed = .true.
                    endif

                    x_end = FP_NAN
                    y_end = FP_NAN

                else
                    x_init = poincare_data(1, i_plane, i_turn, i_surf)
                    y_init = poincare_data(2, i_plane, i_turn, i_surf)

                    call trace(x_init=x_init, &
                               y_init=y_init, &
                               phi_init=phi_array(i_plane), &
                               dphi=dphi_array(i_plane), &
                               equi=equi, &
                               x_end=x_end, &
                               y_end=y_end, &
                               istat=surface_trace_status)
                endif

                ! Store final trace location for each turn. If we have wrapped
                ! around from last plane to first (i_plane == n_planes), then
                ! we have already gone one full turn and are back at the first
                ! plane.
                if (i_plane < n_planes) then
                    poincare_data(1, i_plane + 1, i_turn, i_surf) = x_end
                    poincare_data(2, i_plane + 1, i_turn, i_surf) = y_end
                elseif (i_turn /= n_turns + 1) then
                    poincare_data(1, 1, i_turn + 1, i_surf) = x_end
                    poincare_data(2, 1, i_turn + 1, i_surf) = y_end
                endif
            end do
            end do
        end do
        !$omp end parallel do

        dt = MPI_wtime() - start_wtime
        write(get_stdout(), "(A, F10.3, A)") "Poincare data generation &
                                             &completed in ", dt, " seconds"

        call write_poincare(filename, n_surfaces, n_turns, n_planes, &
                            phi_array, poincare_data)

    end subroutine generate_poincare_data

    subroutine write_poincare(filename, n_surfaces, n_turns, n_planes, &
                              phi_array, poincare_data)
        !! Writes generated poincare data to a NetCDF file
        character(len=*), intent(in) :: filename
        !! File name to save to
        integer, intent(in) :: n_surfaces
        !! Number of flux surfaces
        integer, intent(in) :: n_turns
        !! Number of toroidal rotations
        integer, intent(in) :: n_planes
        !! Number of toroidal planes
        real(FP), intent(in), dimension(n_planes) :: phi_array
        !! Toroidal angles of each plane
        real(FP), intent(in), &
            dimension(2, n_planes, n_turns + 1, n_surfaces) :: poincare_data
        !! Generated Poincare data: x and y locations of the field lines for
        !! every surface, turn, and plane

        integer :: ierr, file_id
        integer :: id_dim_xy, id_dim_surf, id_dim_turn, id_dim_plane, &
                   id_dim_poincare(4), id_phi, id_poincare

        ! Open file
        ierr = nf90_create(filename, NF90_NETCDF4, file_id)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ! Write global attributes
        ierr = nf90_put_att(file_id, NF90_GLOBAL, 'equi_id', eq_id)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ierr = nf90_put_att(file_id, NF90_GLOBAL, 'n_surfaces', n_surfaces)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ierr = nf90_put_att(file_id, NF90_GLOBAL, 'n_turns', n_turns)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ierr = nf90_put_att(file_id, NF90_GLOBAL, 'n_planes', n_planes)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ! Definitions
        ierr = nf90_def_dim(file_id, 'dim_xy', 2, id_dim_xy)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ierr = nf90_def_dim(file_id, 'dim_surf', n_surfaces, id_dim_surf)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ierr = nf90_def_dim(file_id, 'dim_turn', n_turns + 1, id_dim_turn)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ierr = nf90_def_dim(file_id, 'dim_plane', n_planes, id_dim_plane)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        id_dim_poincare = [id_dim_xy, id_dim_plane, id_dim_turn, id_dim_surf]

        ierr = nf90_def_var(file_id, 'phi_array', NF90_FP, id_dim_plane, id_phi)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ierr = nf90_def_var(file_id, 'poincare_data', NF90_FP, &
                            id_dim_poincare, id_poincare)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ! Put
        ierr = nf90_put_var(file_id, id_phi, phi_array)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ierr = nf90_put_var(file_id, id_poincare, poincare_data)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

        ! Close file
        ierr = nf90_close(file_id)
        call handle_error_netcdf(ierr, __LINE__, __FILE__)

    end subroutine write_poincare

end program