optimizer.f90

Go to the documentation of this file.

 
 
module optimizer
  use json_module, only: json_file
  use simulation_m, only: simulation_t
  use problem, only: problem_t
  use design, only: design_t
  use num_types, only: rp
  use logger, only: neko_log
  use profiler, only: profiler_start_region, profiler_end_region
  use mpi_f08, only: mpi_wtime, mpi_allreduce, mpi_max
  use utils, only: neko_error, filename_suffix
  use csv_file, only: csv_file_t
  use vector, only: vector_t
  use json_utils, only: json_get_or_default
  use comm, only: pe_rank, mpi_real_precision, neko_comm
  use continuation_scheduler, only: nekotop_continuation
 
  implicit none
  private
 
  type, abstract, public :: optimizer_t
 
     character(len=64), private :: optimizer_type = ''
     integer, private :: max_iterations = 0
     integer, private :: current_iteration = 0
 
     ! ----------------------------------------------------------------------- !
     ! Restart related members
 
     character(len=256), private :: checkpoint_file = ''
 
     ! Checkpoint related information
     character(len=256), private :: checkpoint_path = './checkpoints/'
     character(len=256), private :: checkpoint_base = 'optimizer_checkpoint'
     character(len=256), private :: checkpoint_format = 'hdf5'
     integer, private :: checkpoint_interval = -1
 
     ! Variables for the runtime-based stopping criteria
     real(kind=rp), private :: max_runtime = -1.0_rp
     real(kind=rp), private :: start_time = 0.0_rp
     real(kind=rp), private :: average_time = 0.0_rp
     real(kind=rp), private :: step_count = 0.0_rp
     ! Logging state
     logical, private :: log_initialized = .false.
     logical, private :: log_include_constraints = .true.
     integer, private :: log_extra_size = 0
     type(csv_file_t), private :: log_file
     type(vector_t), private :: log_data
 
   contains
 
     !  ---------------------------------------------------------------------- !
     ! Deferred procedures for specific optimizers
 
     procedure(optimizer_init_from_json), pass(this), public, deferred :: &
          init_from_json
     procedure(optimizer_free), pass(this), public, deferred :: free
 
     procedure(optimizer_initialize), pass(this), public, deferred :: initialize
     procedure(optimizer_step), pass(this), public, deferred :: step
     procedure(optimizer_validate), pass(this), public, deferred :: validate
 
     procedure(optimizer_write), pass(this), public, deferred :: write
     procedure(optimizer_save_checkpoint_components), pass(this), deferred :: &
          save_checkpoint_components
     procedure(optimizer_load_checkpoint_components), pass(this), deferred :: &
          load_checkpoint_components
 
     procedure, pass(this) :: save_checkpoint => optimizer_save_checkpoint
     procedure, pass(this) :: load_checkpoint => optimizer_load_checkpoint
 
     ! ----------------------------------------------------------------------- !
     ! Public procedures
 
     procedure, pass(this), public :: run => optimizer_run
 
     ! ----------------------------------------------------------------------- !
     ! Private procedures
 
     procedure, pass(this) :: init_base => optimizer_init_base
     procedure, pass(this) :: free_base => optimizer_free_base
     procedure, pass(this) :: read_base_settings => optimizer_read_base_settings
     procedure, pass(this) :: print_status => optimizer_print_status
     procedure, pass(this) :: out_of_time => optimizer_out_of_time
     procedure, pass(this) :: init_log => optimizer_init_log
     procedure, pass(this) :: write_log => optimizer_write_log
  end type optimizer_t
 
  ! -------------------------------------------------------------------------- !
  ! Interface for the optimizer module.
 
  abstract interface
 
     subroutine optimizer_init_from_json(this, parameters, problem, design, &
          simulation)
       import optimizer_t, json_file, simulation_t, problem_t, design_t, rp
       class(optimizer_t), intent(inout) :: this
       type(json_file), intent(inout) :: parameters
       class(problem_t), intent(inout) :: problem
       class(design_t), intent(in) :: design
       type(simulation_t), optional, intent(in) :: simulation
     end subroutine optimizer_init_from_json
 
     subroutine optimizer_initialize(this, problem, design, simulation)
       import optimizer_t, simulation_t, problem_t, design_t
       class(optimizer_t), intent(inout) :: this
       class(problem_t), intent(inout) :: problem
       class(design_t), intent(inout) :: design
       type(simulation_t), optional, intent(inout) :: simulation
     end subroutine optimizer_initialize
 
     subroutine optimizer_free(this)
       import optimizer_t
       class(optimizer_t), intent(inout) :: this
     end subroutine optimizer_free
 
     logical function optimizer_step(this, iter, problem, design, simulation)
       import optimizer_t, simulation_t, problem_t, design_t
       class(optimizer_t), intent(inout) :: this
       integer, intent(in) :: iter
       class(problem_t), intent(inout) :: problem
       class(design_t), intent(inout) :: design
       type(simulation_t), optional, intent(inout) :: simulation
     end function optimizer_step
 
     subroutine optimizer_validate(this, problem, design)
       import optimizer_t, problem_t, design_t
       class(optimizer_t), intent(inout) :: this
       class(problem_t), intent(in) :: problem
       class(design_t), intent(in) :: design
     end subroutine optimizer_validate
 
     subroutine optimizer_write(this, iter, problem)
       import optimizer_t, simulation_t, problem_t, design_t
       class(optimizer_t), intent(inout) :: this
       integer, intent(in) :: iter
       class(problem_t), intent(inout) :: problem
     end subroutine optimizer_write
 
     subroutine optimizer_save_checkpoint_components(this, filename, overwrite)
       import optimizer_t
       class(optimizer_t), intent(inout) :: this
       character(len=*), intent(in) :: filename
       logical, intent(in), optional :: overwrite
     end subroutine optimizer_save_checkpoint_components
 
     subroutine optimizer_load_checkpoint_components(this, filename)
       import optimizer_t
       class(optimizer_t), intent(inout) :: this
       character(len=*), intent(in) :: filename
     end subroutine optimizer_load_checkpoint_components
 
  end interface
 
  ! -------------------------------------------------------------------------- !
  ! Interfaces for the factory functions
 
  interface optimizer_factory
     module subroutine optimizer_factory(object, parameters, problem, design, &
          simulation)
       class(optimizer_t), allocatable, intent(inout) :: object
       type(json_file), intent(inout) :: parameters
       class(problem_t), intent(inout) :: problem
       class(design_t), intent(in) :: design
       type(simulation_t), optional, intent(in) :: simulation
     end subroutine optimizer_factory
  end interface optimizer_factory
 
  ! -------------------------------------------------------------------------- !
  ! IO routines for HDF5 checkpoints
 
  interface
 
     module subroutine optimizer_save_checkpoint_hdf5(object, filename, iter, &
          overwrite)
       class(optimizer_t), intent(inout) :: object
       character(len=*), intent(in) :: filename
       integer, intent(in) :: iter
       logical, intent(in), optional :: overwrite
     end subroutine optimizer_save_checkpoint_hdf5
 
     module subroutine optimizer_load_checkpoint_hdf5(object, filename, iter)
       class(optimizer_t), intent(inout) :: object
       character(len=*), intent(in) :: filename
       integer, intent(out) :: iter
     end subroutine optimizer_load_checkpoint_hdf5
  end interface
 
  public :: optimizer_factory
 
contains
 
  ! -------------------------------------------------------------------------- !
  ! Base initializer and free routines
 
  subroutine optimizer_init_base(this, optimizer_type, max_iterations, &
       max_runtime, checkpoint_file, checkpoint_path, checkpoint_base, &
       checkpoint_format, checkpoint_interval)
    class(optimizer_t), intent(inout) :: this
    character(len=*), intent(in) :: optimizer_type
    integer, intent(in) :: max_iterations
    real(kind=rp), intent(in), optional :: max_runtime
    character(len=*), intent(in), optional :: checkpoint_file
    character(len=*), intent(in), optional :: checkpoint_path
    character(len=*), intent(in), optional :: checkpoint_base
    character(len=*), intent(in), optional :: checkpoint_format
    integer, intent(in), optional :: checkpoint_interval
 
    ! Mandatory settings
    this%optimizer_type = optimizer_type
    this%max_iterations = max_iterations
 
    ! Optional settings
    if (present(max_runtime)) this%max_runtime = max_runtime
    if (present(checkpoint_file)) this%checkpoint_file = checkpoint_file
    if (present(checkpoint_path)) this%checkpoint_path = checkpoint_path
    if (present(checkpoint_base)) this%checkpoint_base = checkpoint_base
    if (present(checkpoint_format)) this%checkpoint_format = checkpoint_format
    if (present(checkpoint_interval)) then
       this%checkpoint_interval = checkpoint_interval
    end if
 
    ! Initialize internals
    this%start_time = mpi_wtime()
 
  end subroutine optimizer_init_base
 
  subroutine optimizer_free_base(this)
    class(optimizer_t), intent(inout) :: this
 
    this%optimizer_type = ''
    this%max_iterations = 0
    this%max_runtime = -1.0_rp
    this%checkpoint_file = ''
    this%checkpoint_path = './checkpoints/'
    this%checkpoint_base = 'optimizer_checkpoint'
    this%checkpoint_format = 'hdf5'
    this%checkpoint_interval = -1
 
    this%start_time = 0.0_rp
    this%current_iteration = 0
    call this%log_data%free()
    this%log_initialized = .false.
    this%log_extra_size = 0
    this%log_include_constraints = .true.
 
  end subroutine optimizer_free_base
 
  subroutine optimizer_read_base_settings(this, solver_params)
    class(optimizer_t), intent(inout) :: this
    type(json_file), intent(inout) :: solver_params
    integer :: read_int
    real(kind=rp) :: read_real
    character(len=:), allocatable :: read_str
 
    call json_get_or_default(solver_params, 'max_runtime', read_real, &
         this%max_runtime)
    this%max_runtime = read_real
    call json_get_or_default(solver_params, 'restart_file', read_str, &
         this%checkpoint_file)
    this%checkpoint_file = read_str
 
    call json_get_or_default(solver_params, 'checkpoint.path', read_str, &
         this%checkpoint_path)
    this%checkpoint_path = read_str
    call json_get_or_default(solver_params, 'checkpoint.base', read_str, &
         this%checkpoint_base)
    this%checkpoint_base = read_str
    call json_get_or_default(solver_params, 'checkpoint.format', read_str, &
         this%checkpoint_format)
    this%checkpoint_format = read_str
    call json_get_or_default(solver_params, 'checkpoint.interval', read_int, &
         this%checkpoint_interval)
    this%checkpoint_interval = read_int
 
  end subroutine optimizer_read_base_settings
 
 
  ! -------------------------------------------------------------------------- !
  ! Optimization loop routine
 
  subroutine optimizer_run(this, problem, design, simulation)
    class(optimizer_t), intent(inout) :: this
    class(problem_t), intent(inout) :: problem
    class(design_t), intent(inout) :: design
    type(simulation_t), optional, intent(inout) :: simulation
    real(kind=rp) :: iteration_time
    character(len=1024) :: checkpoint_file
    logical :: converged, file_exists
    integer :: stop_flag
 
    ! Initialize variables
    stop_flag = 1
    converged = .false.
 
    ! Restart from checkpoint if available
    if (trim(this%checkpoint_file) .ne. '') then
       checkpoint_file = trim(this%checkpoint_file)
    else
       select case (trim(this%checkpoint_format))
       case ('h5', 'hdf5', 'hf5', 'hdf')
          checkpoint_file = trim(this%checkpoint_path) // &
               'optimizer_rt_checkpoint.h5'
       end select
    end if
 
    inquire(file = checkpoint_file, exist = file_exists)
    if (file_exists) then
       call this%load_checkpoint(checkpoint_file, this%current_iteration, &
            design)
    end if
 
    ! compute potential internals for for a potential restart
    if (this%current_iteration .ne. 0) then
       call design%set_output_counter(this%current_iteration - 1)
       if (present(simulation)) then
          call simulation%set_output_counter(this%current_iteration - 1)
       end if
    end if
 
    ! Prepare the problem state before starting the optimization
    call this%initialize(problem, design, simulation)
 
    call this%write(this%current_iteration, problem)
    call design%write(this%current_iteration)
 
    call neko_log%section('Optimization Loop')
 
    do while (this%current_iteration .lt. this%max_iterations)
       this%current_iteration = this%current_iteration + 1
       if (pe_rank .eq. 0) then
          write(*,*) 'Starting iteration ', this%current_iteration
       end if
 
       call profiler_start_region('Optimizer iteration')
       iteration_time = mpi_wtime()
 
       ! Update the parameters in continuation scheduler
       call nekotop_continuation%update(this%current_iteration)
 
       converged = this%step(this%current_iteration, problem, design, &
            simulation)
 
       iteration_time = mpi_wtime() - iteration_time
       call profiler_end_region('Optimizer iteration')
 
       ! Log the progress and outputs
       call this%write(this%current_iteration, problem)
       call design%write(this%current_iteration)
 
       ! Save checkpoint if enabled
       if (this%checkpoint_interval .gt. 0 .and. &
            mod(this%current_iteration, this%checkpoint_interval) == 0) then
          call this%save_checkpoint(this%current_iteration, design, .false.)
       end if
 
       ! --------------------------------------------------------------------- !
       ! Check stopping criteria
 
       if (converged) then
          stop_flag = 0
          exit
       else if (this%out_of_time(iteration_time)) then
          call this%save_checkpoint(this%current_iteration, design, .true., &
               basename = 'optimizer_rt_checkpoint')
          stop_flag = 2
          exit
       end if
    end do
 
    ! Check that the final design is valid
    call this%validate(problem, design)
    call this%print_status(stop_flag, this%current_iteration)
 
    call neko_log%end_section()
 
  end subroutine optimizer_run
 
  ! ========================================================================== !
  ! Helper routines
 
  subroutine optimizer_print_status(this, stop_flag, iter)
    class(optimizer_t), intent(in) :: this
    integer, intent(in) :: stop_flag
    integer, intent(in) :: iter
    character(len=256) :: msg
 
    select case (stop_flag)
    case (0)
       write(msg, '(A,I0,A)') 'Optimizer converged successfully after ', &
            iter, ' iterations.'
       call neko_log%message(msg)
    case (1)
       write(msg, '(A,I0,A)') 'Optimizer did not converge in ', &
            this%max_iterations, ' iterations.'
       call neko_log%warning(msg)
    case (2)
       write(msg, '(A,A,F8.2,A)') 'Optimizer stopped after reaching the ', &
            'maximum runtime of ', this%max_runtime, ' seconds.'
       call neko_error(trim(msg))
 
    case default
       write(msg, '(A)') 'Optimizer stopped for an unknown reason.'
       call neko_error(msg)
    end select
  end subroutine optimizer_print_status
 
  function optimizer_out_of_time(this, step_time) result(out_of_time)
    class(optimizer_t), intent(inout) :: this
    real(kind=rp), intent(in) :: step_time
    logical :: out_of_time
    real(kind=rp) :: elapsed_time, time, old_avg_weight
 
    out_of_time = .false.
 
    if (this%max_runtime .lt. 0.0_rp) then
       return
    end if
 
    call mpi_allreduce(step_time, time, 1, mpi_real_precision, mpi_max, &
         neko_comm)
 
    elapsed_time = mpi_wtime() - this%start_time
    this%step_count = this%step_count + 1.0_rp
    old_avg_weight = (this%step_count - 1) / this%step_count
 
    ! Estimate Cumulative Average iteration time
    this%average_time = time / this%step_count + &
         this%average_time * old_avg_weight
 
    ! Determine if next iteration would exceed max runtime
    out_of_time = (elapsed_time + this%average_time) .gt. this%max_runtime
 
  end function optimizer_out_of_time
 
  ! ========================================================================== !
  ! Logging helpers
 
  subroutine optimizer_init_log(this, problem, extra_headers, &
       include_constraints, filename)
    class(optimizer_t), intent(inout) :: this
    class(problem_t), intent(in) :: problem
    character(len=*), intent(in), optional :: extra_headers(:)
    logical, intent(in), optional :: include_constraints
    character(len=*), intent(in), optional :: filename
 
    character(len=4096) :: header
    integer :: total_size, base_size, i, n_cont
    character(len=256) :: log_name
 
    if (present(include_constraints)) then
       this%log_include_constraints = include_constraints
    end if
    n_cont = nekotop_continuation%get_n_params()
 
    base_size = problem%get_log_size(this%log_include_constraints)
 
    this%log_extra_size = 0
    if (present(extra_headers)) this%log_extra_size = size(extra_headers)
 
    total_size = 1 + base_size + this%log_extra_size + n_cont
    call this%log_data%init(total_size)
 
    if (present(filename)) then
       log_name = trim(filename)
    else
       log_name = 'optimization_data.csv'
    end if
 
    call this%log_file%init(trim(log_name))
 
    header = 'iter, ' // &
         trim(problem%get_log_header(this%log_include_constraints))
    if (present(extra_headers)) then
       do i = 1, size(extra_headers)
          if (trim(extra_headers(i)) .eq. '') then
             call neko_error('some headers are empty')
          end if
          header = trim(header) // ', ' // trim(extra_headers(i))
       end do
    end if
 
    ! continuation parameters
    do i = 1, n_cont
       header = trim(header) // ', ' // &
            trim(nekotop_continuation%get_param_name(i))
    end do
    call this%log_file%set_header(trim(header))
 
    this%log_initialized = .true.
  end subroutine optimizer_init_log
 
  subroutine optimizer_write_log(this, iter, problem, extra_values)
    class(optimizer_t), intent(inout) :: this
    integer, intent(in) :: iter
    class(problem_t), intent(in) :: problem
    real(kind=rp), intent(in), optional :: extra_values(:)
    integer :: base_size, offset, n_cont, i
 
    if (.not. this%log_initialized) return
 
    ! Number of continuation parameters
    n_cont = nekotop_continuation%get_n_params()
 
    base_size = problem%get_log_size(this%log_include_constraints)
    this%log_data%x(1) = real(iter, kind=rp)
 
    call problem%get_log_values( &
         this%log_data%x(2:1 + base_size), &
         this%log_include_constraints)
 
    offset = 2 + base_size
    if (present(extra_values)) then
       if (this%log_extra_size .eq. 0) then
          call neko_error('got extra values but no headers')
       end if
       this%log_data%x(offset:offset + size(extra_values) - 1) = extra_values
    end if
 
    ! Continuation parameter values
    do i = 1, n_cont
       this%log_data%x(offset + size(extra_values) - 1 + i) = &
            nekotop_continuation%params(i)%target
    end do
 
    call this%log_file%write(this%log_data)
 
  end subroutine optimizer_write_log
 
  ! ========================================================================== !
  ! IO Functions
 
  subroutine optimizer_save_checkpoint(this, iter, design, overwrite, &
       path, basename, format)
    class(optimizer_t), intent(inout) :: this
    integer, intent(in) :: iter
    class(design_t), intent(inout) :: design
    logical, intent(in) :: overwrite
    character(len=*), intent(in), optional :: path
    character(len=*), intent(in), optional :: basename
    character(len=*), intent(in), optional :: format
    character(len=:), allocatable :: checkpoint_format
    character(len=256) :: file_path, file_base, file_ext, file_full
    logical :: exist
 
    ! Set default behaviour, read from object if not provided
    if (.not. present(path)) file_path = trim(this%checkpoint_path)
    if (.not. present(basename)) file_base = trim(this%checkpoint_base)
    if (.not. present(format)) checkpoint_format = trim(this%checkpoint_format)
 
    if (present(path)) file_path = trim(path)
    if (present(basename)) file_base = trim(basename)
    if (present(format)) checkpoint_format = trim(format)
 
    ! Make sure path is valid and exists
    if (len_trim(file_path) .eq. 0) then
       file_path = './'
    else if (file_path(len_trim(file_path):len_trim(file_path)) .ne. '/') then
       file_path = trim(file_path) // '/'
    end if
 
    inquire(file=file_path, exist=exist)
    if (.not. exist) then
       call execute_command_line('mkdir -p "' // trim(file_path) // '"')
    end if
 
    select case (trim(checkpoint_format))
    case ('h5', 'hdf5', 'hf5', 'hdf')
       file_ext = 'h5'
    case default
       call neko_error('optimizer: Unsupported checkpoint format: "' // &
            trim(checkpoint_format) // '"')
    end select
 
    ! Construct the full filename based on overwrite flag
    if (overwrite) then
       write(file_full, '(4A)') &
            trim(file_path), trim(file_base), ".", trim(file_ext)
    else
       write(file_full, '(3A,I5.5,2A)') &
            trim(file_path), trim(file_base), "_", iter, ".", trim(file_ext)
    end if
 
    select case (trim(file_ext))
    case ('h5', 'hdf5', 'hf5')
       call optimizer_save_checkpoint_hdf5(this, file_full, iter, overwrite)
    case default
       call neko_error('optimizer: Unsupported checkpoint format: "' // &
            trim(file_ext) // '"')
    end select
 
    call this%save_checkpoint_components(file_full, overwrite)
    call design%save_checkpoint(file_full, overwrite)
 
  end subroutine optimizer_save_checkpoint
 
  subroutine optimizer_load_checkpoint(this, filename, iter, design)
    class(optimizer_t), intent(inout) :: this
    character(len=*), intent(in) :: filename
    integer, intent(out) :: iter
    class(design_t), intent(inout) :: design
    character(len=12) :: file_ext
 
    ! Get the file extension
    call filename_suffix(filename, file_ext)
 
    select case (trim(file_ext))
    case ('h5', 'hdf5', 'hf5')
       call optimizer_load_checkpoint_hdf5(this, filename, iter)
    case default
       call neko_error('optimizer: Unsupported checkpoint format: "' // &
            trim(file_ext) // '"')
    end select
 
    call this%load_checkpoint_components(filename)
    call design%load_checkpoint(filename)
 
    ! Set the current iteration to the loaded iteration
    this%current_iteration = iter
 
    if (pe_rank .eq. 0) then
       write(*,*) 'Restarted simulation from checkpoint.'
       write(*,*) '    Checkpoint file: "', trim(filename), '"'
       write(*,*) '    Iteration      : ', this%current_iteration
    end if
 
  end subroutine optimizer_load_checkpoint
 
  ! ========================================================================== !
  ! Dummy implementations for module procedures
 
#if !HAVE_HDF5
  module subroutine optimizer_save_checkpoint_hdf5(object, filename, iter, &
       overwrite)
    class(optimizer_t), intent(inout) :: object
    character(len=*), intent(in) :: filename
    integer, intent(in) :: iter
    logical, intent(in), optional :: overwrite
    call neko_error('optimizer: HDF5 support not enabled rebuild with ' // &
         'HAVE_HDF5')
  end subroutine optimizer_save_checkpoint_hdf5
 
  module subroutine optimizer_load_checkpoint_hdf5(object, filename, iter)
    class(optimizer_t), intent(inout) :: object
    character(len=*), intent(in) :: filename
    integer, intent(out) :: iter
    call neko_error('optimizer: HDF5 support not enabled rebuild with ' // &
         'HAVE_HDF5')
  end subroutine optimizer_load_checkpoint_hdf5
#endif
 
end module optimizer