writecheckpoint.F90

Go to the documentation of this file.

module checkpointer_write_checkpoint_mod
#ifndef TEST_MODE
  use netcdf, only : nf90_double, nf90_real, nf90_int, nf90_char, nf90_global, nf90_clobber, nf90_netcdf4, nf90_mpiio, &
       nf90_collective, nf90_def_var, nf90_var_par_access, nf90_def_var_fill, nf90_put_att, nf90_create, nf90_put_var, &
       nf90_def_dim, nf90_enddef, nf90_close, nf90_inq_dimid, nf90_inq_varid
#else
  use dummy_netcdf_mod, only : nf90_double, nf90_real, nf90_int, nf90_char, nf90_global, nf90_clobber, nf90_netcdf4, nf90_mpiio, &
       nf90_def_var, nf90_put_att, nf90_create, nf90_put_var, nf90_def_dim, nf90_enddef, nf90_close
#endif
  use state_mod, only : model_state_type
  use grids_mod, only : local_grid_type, global_grid_type, vertical_grid_configuration_type, x_index, y_index, z_index
  use prognostics_mod, only : prognostic_field_type
  use conversions_mod, only : conv_to_string
  use optionsdatabase_mod, only : options_size, options_value_at, options_key_at
  use checkpointer_common_mod, only : empty_dim_key, key_value_pair_key, options_dim_key, options_key, string_dim_key, &
       x_dim_key, y_dim_key, z_dim_key, zn_dim_key, q_dim_key, q_key, zq_key, th_key, zth_key, p_key, u_key, v_key, w_key, &
       zu_key, zv_key, zw_key, x_key, y_key, z_key, zn_key, nqfields, ugal, vgal, time_key, timestep, max_string_length, &
       created_attribute_key, title_attribute_key, absolute_new_dtm_key, dtm_key, dtm_new_key, q_indices_key, &
       q_indices_dim_key, x_resolution, y_resolution,x_top, y_top, x_bottom, y_bottom, thref, olubar, olzubar, olvbar, &
       olzvbar, olthbar, olzthbar, olqbar, olzqbar, check_status
  use datadefn_mod, only : default_precision, single_precision, double_precision, string_length
  use q_indices_mod, only : q_metadata_type, get_max_number_q_indices, get_indices_descriptor, get_number_active_q_indices
  use mpi, only : mpi_info_null
  implicit none

#ifndef TEST_MODE
  private
#endif

  character(len=*), parameter :: checkpoint_title = "MONC checkpoint file"

  public write_checkpoint_file

contains

  subroutine write_checkpoint_file(current_state, filename)
    type(model_state_type), intent(inout) :: current_state
    character(len=*), intent(in) :: filename

    integer :: ncid, z_dim_id, y_dim_id, x_dim_id, q_dim_id, x_id, y_id, z_id, th_id, p_id, time_id,&
         u_id, v_id, w_id, q_id, zu_id, zv_id, zw_id, zth_id, zq_id, timestep_id, ugal_id, &
         vgal_id, number_q_fields_id, string_dim_id, key_value_dim_id, options_id, q_indices_id, &
         dtm_id, dtm_new_id, absolute_new_dtm_id
    logical :: q_indices_declared

#ifdef SINGLE_MONC_DO_SEQUENTIAL_NETCDF
    if (current_state%parallel%processes .gt. 1) then
      call check_status(nf90_create(filename, ior(nf90_netcdf4, nf90_mpiio), ncid, &
           comm = current_state%parallel%monc_communicator, info = mpi_info_null))
    else
      call check_status(nf90_create(filename, nf90_clobber, ncid))
    end if
#else
    call check_status(nf90_create(filename, ior(nf90_netcdf4, nf90_mpiio), ncid, &
         comm = current_state%parallel%monc_communicator, info = mpi_info_null))
#endif
    call write_out_global_attributes(ncid)
    call define_grid_dimensions(current_state, ncid, z_dim_id, y_dim_id, x_dim_id)    
    if (current_state%number_q_fields .gt. 0) call define_q_field_dimension(current_state, ncid, q_dim_id)

    call define_options_variable(current_state, ncid, string_dim_id, key_value_dim_id, options_id)
    q_indices_declared=define_q_indices_variable(ncid, string_dim_id, key_value_dim_id, q_indices_id)
    call define_grid_variables(current_state, ncid)
    call define_mean_fields(current_state, ncid)
    if (current_state%number_q_fields .gt. 0) call define_q_variable(ncid, current_state%parallel%processes .gt. 1, &
         q_dim_id, z_dim_id, y_dim_id, x_dim_id, q_id, zq_id)
    call define_prognostic_variables(current_state, current_state%parallel%processes .gt. 1, ncid, z_dim_id, y_dim_id, &
         x_dim_id, u_id, v_id, w_id, th_id, p_id, zu_id, zv_id, zw_id, zth_id)
    call define_misc_variables(ncid, timestep_id, time_id, ugal_id, vgal_id, number_q_fields_id, &
         dtm_id, dtm_new_id, absolute_new_dtm_id)

    call check_status(nf90_enddef(ncid))

    if (current_state%parallel%my_rank==0) call write_out_grid(ncid, current_state%global_grid)
    if (current_state%parallel%my_rank==0) call write_out_mean_fields(ncid, current_state%global_grid)
    call write_out_all_fields(current_state, ncid, u_id, v_id, w_id, zu_id, zv_id, zw_id, th_id, zth_id, q_id, zq_id, p_id)
    if (current_state%parallel%my_rank==0) then
      call write_out_options(current_state, ncid, options_id)
      if (q_indices_declared) call write_out_q_indices(ncid, q_indices_id)
      call write_out_misc_variables(current_state, ncid, timestep_id, time_id, &
           ugal_id, vgal_id, number_q_fields_id, dtm_id, dtm_new_id, absolute_new_dtm_id)
    end if

    call check_status(nf90_close(ncid))
  end subroutine write_checkpoint_file

  subroutine write_out_global_attributes(ncid)
    integer, intent(in) :: ncid

    integer :: date_values(8)

    call date_and_time(values=date_values)

    call check_status(nf90_put_att(ncid, nf90_global, title_attribute_key, checkpoint_title))
    call check_status(nf90_put_att(ncid, nf90_global, created_attribute_key, trim(conv_to_string(date_values(3)))//"/"//&
         trim(conv_to_string(date_values(2)))//"/"//trim(conv_to_string(date_values(1)))//" "//trim(conv_to_string(&
         date_values(5)))// ":"//trim(conv_to_string(date_values(6)))//":"//trim(conv_to_string(date_values(7)))))
  end subroutine write_out_global_attributes

  subroutine write_out_q_indices(ncid, q_indices_id)
    integer, intent(in) :: ncid, q_indices_id

    integer :: i, current_index
    type(q_metadata_type) :: specific_q_data

    current_index=1
    do i=1, get_max_number_q_indices()
      specific_q_data=get_indices_descriptor(i)
      if (specific_q_data%l_used) then
        call check_status(nf90_put_var(ncid, q_indices_id, trim(specific_q_data%name), (/ 1, 1, current_index /)))
        call check_status(nf90_put_var(ncid, q_indices_id, trim(conv_to_string(i)), (/ 1, 2, current_index /)))
        current_index=current_index+1
      end if
    end do    
  end subroutine write_out_q_indices

  subroutine write_out_options(current_state, ncid, options_id)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: ncid, options_id

    integer :: i
    character(len=STRING_LENGTH), pointer :: sized_raw_character
    class(*), pointer :: raw_data, raw_to_string

    do i=1,options_size(current_state%options_database)
      raw_data=> options_value_at(current_state%options_database, i)
      raw_to_string=>raw_data
      call check_status(nf90_put_var(ncid, options_id, trim(options_key_at(current_state%options_database, i)), (/ 1, 1, i /)))
      select type (raw_data)
      type is(integer)
        call check_status(nf90_put_var(ncid, options_id, trim(conv_to_string(raw_data)), (/ 1, 2, i /)))
      type is(real(kind=single_precision))
        call check_status(nf90_put_var(ncid, options_id, trim(conv_to_string(raw_data)), (/ 1, 2, i /)))
      type is(real(kind=double_precision))
        call check_status(nf90_put_var(ncid, options_id, trim(conv_to_string(raw_data)), (/ 1, 2, i /)))
      type is(logical)
        call check_status(nf90_put_var(ncid, options_id, trim(conv_to_string(raw_data)), (/ 1, 2, i /)))
      type is(character(len=*))
        ! Done this way to give the character size information and keep the (unsafe) cast in the conversion module
        sized_raw_character=>conv_to_string(raw_to_string, .false., string_length)
        call check_status(nf90_put_var(ncid, options_id, trim(sized_raw_character), (/ 1, 2, i /)))
      end select
    end do
  end subroutine write_out_options

  subroutine write_out_all_fields(current_state, ncid, u_id, v_id, w_id, zu_id, zv_id, zw_id, th_id, zth_id, q_id, zq_id, p_id)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: ncid, u_id, v_id, w_id, zu_id, zv_id, zw_id, th_id, zth_id, q_id, zq_id, p_id

    integer :: i
    logical :: multi_process

    multi_process = current_state%parallel%processes .gt. 1
#ifdef U_ACTIVE
    call write_out_velocity_field(ncid, current_state%local_grid, current_state%u, u_id, multi_process)
    call write_out_velocity_field(ncid, current_state%local_grid, current_state%zu, zu_id, multi_process)
#endif
#ifdef V_ACTIVE
    call write_out_velocity_field(ncid, current_state%local_grid, current_state%v, v_id, multi_process)
    call write_out_velocity_field(ncid, current_state%local_grid, current_state%zv, zv_id, multi_process)
#endif
#ifdef W_ACTIVE
    call write_out_velocity_field(ncid, current_state%local_grid, current_state%w, w_id, multi_process)
    call write_out_velocity_field(ncid, current_state%local_grid, current_state%zw, zw_id, multi_process)
#endif
    if (current_state%th%active) then
      call write_out_velocity_field(ncid, current_state%local_grid, current_state%th, th_id, multi_process)
      call write_out_velocity_field(ncid, current_state%local_grid, current_state%zth, zth_id, multi_process)
    end if
    if (current_state%p%active) call write_out_velocity_field(ncid, current_state%local_grid, current_state%p, &
         p_id, multi_process)
    do i=1,current_state%number_q_fields
      if (current_state%q(i)%active) then
        call write_out_velocity_field(ncid, current_state%local_grid, current_state%q(i), q_id, multi_process, i)
        call write_out_velocity_field(ncid, current_state%local_grid, current_state%zq(i), zq_id, multi_process, i)
      end if
    end do
  end subroutine write_out_all_fields

  subroutine write_out_velocity_field(ncid, local_grid, field, variable_id, multi_process, fourth_dim_loc)
    integer, intent(in) :: ncid, variable_id
    type(prognostic_field_type), intent(in) :: field
    type(local_grid_type), intent(inout) :: local_grid
    logical, intent(in) :: multi_process
    integer, optional, intent(in) :: fourth_dim_loc

    integer :: start(4), count(4), i, map(4)

    if (multi_process .or. present(fourth_dim_loc)) then
      do i=1,3
        if (i==1) then
          map(i)=1
        else
          map(i)=map(i-1)*local_grid%size(i-1)
        end if
        start(i) = local_grid%start(i)
        count(i) = local_grid%size(i)
      end do
      if (present(fourth_dim_loc)) then
        start(4) = fourth_dim_loc
        count(4) = 1
        map(4)=map(3)*local_grid%size(3)
      end if

      call check_status(nf90_put_var(ncid, variable_id, field%data(local_grid%local_domain_start_index(z_index):&
           local_grid%local_domain_end_index(z_index),local_grid%local_domain_start_index(y_index):&
           local_grid%local_domain_end_index(y_index), local_grid%local_domain_start_index(x_index):&
           local_grid%local_domain_end_index(x_index)), start=start, count=count))
    else
      call check_status(nf90_put_var(ncid, variable_id, field%data(local_grid%local_domain_start_index(z_index):&
           local_grid%local_domain_end_index(z_index),local_grid%local_domain_start_index(y_index):&
           local_grid%local_domain_end_index(y_index), local_grid%local_domain_start_index(x_index):&
           local_grid%local_domain_end_index(x_index))))
    end if
  end subroutine write_out_velocity_field

  subroutine write_out_grid(ncid, grid)
    integer, intent(in) :: ncid
    type(global_grid_type), intent(in) :: grid

    integer :: var_id

    if (grid%active(z_index)) then
      call write_z_grid_gimension(ncid, grid%configuration%vertical)
    end if
    if (grid%active(y_index)) then
      call check_status(nf90_inq_varid(ncid, y_resolution, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%resolution(y_index)))
      call check_status(nf90_inq_varid(ncid, y_top, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%top(y_index)))
      call check_status(nf90_inq_varid(ncid, y_bottom, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%bottom(y_index)))
    end if
    if (grid%active(x_index)) then
      call check_status(nf90_inq_varid(ncid, x_resolution, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%resolution(x_index)))
      call check_status(nf90_inq_varid(ncid, x_top, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%top(x_index)))
      call check_status(nf90_inq_varid(ncid, x_bottom, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%bottom(x_index)))
    end if
  end subroutine write_out_grid

  subroutine write_out_mean_fields(ncid, grid)
    integer, intent(in) :: ncid
    type(global_grid_type), intent(in) :: grid

    integer :: var_id

    if (allocated(grid%configuration%vertical%olubar)) then
      call check_status(nf90_inq_varid(ncid, olubar, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%configuration%vertical%olubar))
    end if
    if (allocated(grid%configuration%vertical%olzubar)) then
      call check_status(nf90_inq_varid(ncid, olzubar, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%configuration%vertical%olzubar))
    end if
    if (allocated(grid%configuration%vertical%olvbar)) then
      call check_status(nf90_inq_varid(ncid, olvbar, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%configuration%vertical%olvbar))
    end if    
    if (allocated(grid%configuration%vertical%olzvbar)) then
      call check_status(nf90_inq_varid(ncid, olzvbar, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%configuration%vertical%olzvbar))
    end if
    if (allocated(grid%configuration%vertical%olthbar)) then
      call check_status(nf90_inq_varid(ncid, olthbar, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%configuration%vertical%olthbar))
    end if
    if (allocated(grid%configuration%vertical%olzthbar)) then
      call check_status(nf90_inq_varid(ncid, olzthbar, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%configuration%vertical%olzthbar))
    end if    
    if (allocated(grid%configuration%vertical%olqbar)) then
      call check_status(nf90_inq_varid(ncid, olqbar, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%configuration%vertical%olqbar))
    end if
    if (allocated(grid%configuration%vertical%olzqbar)) then
      call check_status(nf90_inq_varid(ncid, olzqbar, var_id))
      call check_status(nf90_put_var(ncid, var_id, grid%configuration%vertical%olzqbar))
    end if    
  end subroutine write_out_mean_fields  

  subroutine write_z_grid_gimension(ncid, vertical_grid)
    type(vertical_grid_configuration_type), intent(in) :: vertical_grid
    integer, intent(in) :: ncid

    integer :: z_var_id, zn_var_id, thref_var_id

    call check_status(nf90_inq_varid(ncid, z_key, z_var_id))
    call check_status(nf90_put_var(ncid, z_var_id, vertical_grid%z))
    call check_status(nf90_inq_varid(ncid, zn_key, zn_var_id))
    call check_status(nf90_put_var(ncid, zn_var_id, vertical_grid%zn))
    call check_status(nf90_inq_varid(ncid, thref, thref_var_id))
    call check_status(nf90_put_var(ncid, thref_var_id, vertical_grid%thref))
  end subroutine write_z_grid_gimension

  subroutine define_options_variable(current_state, ncid, string_dim_id, key_value_dim_id, options_id)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: ncid
    integer, intent(out) :: string_dim_id, key_value_dim_id, options_id

    integer :: options_dim_id, command_dimensions(3)

    call check_status(nf90_def_dim(ncid, string_dim_key, string_length, string_dim_id))
    call check_status(nf90_def_dim(ncid, key_value_pair_key, 2, key_value_dim_id))
    call check_status(nf90_def_dim(ncid, options_dim_key, options_size(current_state%options_database), options_dim_id))

    command_dimensions = (/ string_dim_id, key_value_dim_id, options_dim_id /)

    call check_status(nf90_def_var(ncid, options_key, nf90_char, command_dimensions, options_id))
  end subroutine define_options_variable

  logical function define_q_indices_variable(ncid, string_dim_id, key_value_dim_id, q_indices_id)
    integer, intent(in) :: ncid, string_dim_id, key_value_dim_id
    integer, intent(out) :: q_indices_id

    integer :: q_indices_dim_id, command_dimensions(3), number_active_q

    number_active_q=get_number_active_q_indices()

    if (number_active_q == 0) then
      define_q_indices_variable=.false.
    else
      call check_status(nf90_def_dim(ncid, q_indices_dim_key, number_active_q, q_indices_dim_id))

      command_dimensions = (/ string_dim_id, key_value_dim_id, q_indices_dim_id /)

      call check_status(nf90_def_var(ncid, q_indices_key, nf90_char, command_dimensions, q_indices_id))
      define_q_indices_variable=.true.
    end if
  end function define_q_indices_variable  

  subroutine define_q_field_dimension(current_state, ncid, q_dim_id)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: ncid
    integer, intent(out) :: q_dim_id

    call check_status(nf90_def_dim(ncid, q_dim_key, current_state%number_q_fields, q_dim_id))
  end subroutine define_q_field_dimension

  subroutine define_grid_dimensions(current_state, ncid, z_dim_id, y_dim_id, x_dim_id)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: ncid
    integer, intent(out) :: z_dim_id, y_dim_id, x_dim_id

    integer :: empty_dim_id

    call check_status(nf90_def_dim(ncid, empty_dim_key, 1, empty_dim_id))

    if (current_state%global_grid%active(z_index)) then
      call check_status(nf90_def_dim(ncid, z_dim_key, current_state%global_grid%size(z_index), z_dim_id))
      call check_status(nf90_def_dim(ncid, zn_dim_key, current_state%global_grid%size(z_index), z_dim_id))
    else
      z_dim_id = empty_dim_id
    end if
    if (current_state%global_grid%active(y_index)) then
      call check_status(nf90_def_dim(ncid, y_dim_key, current_state%global_grid%size(y_index), y_dim_id))
    else
      y_dim_id = empty_dim_id
    end if
    if (current_state%global_grid%active(x_index)) then
      call check_status(nf90_def_dim(ncid, x_dim_key, current_state%global_grid%size(x_index), x_dim_id))
    else
      x_dim_id = empty_dim_id
    end if
  end subroutine define_grid_dimensions

  subroutine define_grid_variables(current_state, ncid)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: ncid

    integer :: var_id, z_dim_id

    if (current_state%global_grid%active(x_index)) then
      call check_status(nf90_def_var(ncid, x_resolution, nf90_double, var_id))
      call check_status(nf90_put_att(ncid, var_id, "units", "m"))
      call check_status(nf90_def_var(ncid, x_top, nf90_double, var_id))
      call check_status(nf90_put_att(ncid, var_id, "units", "m"))
      call check_status(nf90_def_var(ncid, x_bottom, nf90_double, var_id))
      call check_status(nf90_put_att(ncid, var_id, "units", "m"))
    end if
    if (current_state%global_grid%active(y_index)) then
      call check_status(nf90_def_var(ncid, y_resolution, nf90_double, var_id))
      call check_status(nf90_put_att(ncid, var_id, "units", "m"))
      call check_status(nf90_def_var(ncid, y_top, nf90_double, var_id))
      call check_status(nf90_put_att(ncid, var_id, "units", "m"))
      call check_status(nf90_def_var(ncid, y_bottom, nf90_double, var_id))
      call check_status(nf90_put_att(ncid, var_id, "units", "m"))
    end if
    if (current_state%global_grid%active(z_index)) then
      call check_status(nf90_inq_dimid(ncid, z_dim_key, z_dim_id))
      call check_status(nf90_def_var(ncid, z_key, nf90_double, z_dim_id, var_id))
      call check_status(nf90_put_att(ncid, var_id, "units", "m"))
      call check_status(nf90_inq_dimid(ncid, zn_dim_key, z_dim_id))
      call check_status(nf90_def_var(ncid, zn_key, nf90_double, z_dim_id, var_id))
      call check_status(nf90_put_att(ncid, var_id, "units", "m"))
    end if
    call check_status(nf90_inq_dimid(ncid, z_dim_key, z_dim_id))
    call check_status(nf90_def_var(ncid, thref, nf90_double, z_dim_id, var_id))
  end subroutine define_grid_variables

  subroutine define_mean_fields(current_state, ncid)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: ncid

    integer :: var_id, z_dim_id, q_dim_id, qdimids(2)

    call check_status(nf90_inq_dimid(ncid, z_dim_key, z_dim_id))

    if (allocated(current_state%global_grid%configuration%vertical%olubar)) then
      call check_status(nf90_def_var(ncid, olubar, nf90_double, z_dim_id, var_id))
    end if    
    if (allocated(current_state%global_grid%configuration%vertical%olzubar)) then
      call check_status(nf90_def_var(ncid, olzubar, nf90_double, z_dim_id, var_id))
    end if
    if (allocated(current_state%global_grid%configuration%vertical%olvbar)) then
      call check_status(nf90_def_var(ncid, olvbar, nf90_double, z_dim_id, var_id))
    end if
    if (allocated(current_state%global_grid%configuration%vertical%olzvbar)) then
      call check_status(nf90_def_var(ncid, olzvbar, nf90_double, z_dim_id, var_id))
    end if
    if (allocated(current_state%global_grid%configuration%vertical%olthbar)) then
      call check_status(nf90_def_var(ncid, olthbar, nf90_double, z_dim_id, var_id))
    end if    
    if (allocated(current_state%global_grid%configuration%vertical%olzthbar)) then
      call check_status(nf90_def_var(ncid, olzthbar, nf90_double, z_dim_id, var_id))
    end if
    if (allocated(current_state%global_grid%configuration%vertical%olqbar) .or. &
         allocated(current_state%global_grid%configuration%vertical%olzqbar)) then
      call check_status(nf90_inq_dimid(ncid, q_dim_key, q_dim_id))
      qdimids=(/ z_dim_id, q_dim_id /)
      if (allocated(current_state%global_grid%configuration%vertical%olqbar)) then
        call check_status(nf90_def_var(ncid, olqbar, nf90_double, qdimids, var_id))
      end if
      if (allocated(current_state%global_grid%configuration%vertical%olzqbar)) then
        call check_status(nf90_def_var(ncid, olzqbar, nf90_double, qdimids, var_id))
      end if
    end if
  end subroutine define_mean_fields  

  subroutine define_q_variable(ncid, multi_process, q_dim_id, z_dim_id, y_dim_id, x_dim_id, q_id, zq_id)
    logical, intent(in) :: multi_process
    integer, intent(in) :: ncid, z_dim_id, y_dim_id, x_dim_id, q_dim_id
    integer, intent(out) :: q_id, zq_id

    integer, dimension(:), allocatable :: dimids

    allocate(dimids(4))
    dimids = (/ z_dim_id, y_dim_id, x_dim_id, q_dim_id /)

    call check_status(nf90_def_var(ncid, q_key, merge(nf90_double, nf90_real, &
         default_precision == double_precision), dimids, q_id))
    call check_status(nf90_def_var(ncid, zq_key, merge(nf90_double, nf90_real, &
         default_precision == double_precision), dimids, zq_id))
    
    if (multi_process) then
      call check_status(nf90_def_var_fill(ncid, q_id, 1, 1))
      call check_status(nf90_def_var_fill(ncid, zq_id, 1, 1))
      call check_status(nf90_var_par_access(ncid, q_id, nf90_collective))
      call check_status(nf90_var_par_access(ncid, zq_id, nf90_collective))
    end if
  end subroutine define_q_variable

  subroutine define_prognostic_variables(current_state, multi_process, ncid, z_dim_id, &
       y_dim_id, x_dim_id, u_id, v_id, w_id, th_id, p_id, zu_id, zv_id, zw_id, zth_id)
    type(model_state_type), intent(inout) :: current_state
    logical, intent(in) :: multi_process
    integer, intent(in) :: ncid, z_dim_id, y_dim_id, x_dim_id
    integer, intent(out) :: u_id, v_id, w_id, th_id, p_id, zu_id, zv_id, zw_id, zth_id

#ifdef U_ACTIVE
    call define_velocity_variable(ncid, multi_process, z_dim_id, y_dim_id, x_dim_id, field_name=u_key, field_id=u_id)
    call define_velocity_variable(ncid, multi_process, z_dim_id, y_dim_id, x_dim_id, field_name=zu_key, field_id=zu_id)
#endif
#ifdef V_ACTIVE
    call define_velocity_variable(ncid, multi_process, z_dim_id, y_dim_id, x_dim_id, field_name=v_key, field_id=v_id)
    call define_velocity_variable(ncid, multi_process, z_dim_id, y_dim_id, x_dim_id, field_name=zv_key, field_id=zv_id)
#endif
#ifdef W_ACTIVE
    call define_velocity_variable(ncid, multi_process, z_dim_id, y_dim_id, x_dim_id, field_name=w_key, field_id=w_id)
    call define_velocity_variable(ncid, multi_process, z_dim_id, y_dim_id, x_dim_id, field_name=zw_key, field_id=zw_id)
#endif
    if (current_state%th%active) then
      call define_velocity_variable(ncid, multi_process, z_dim_id, y_dim_id, x_dim_id, field_name=th_key, field_id=th_id)
      call define_velocity_variable(ncid, multi_process, z_dim_id, y_dim_id, x_dim_id, field_name=zth_key, field_id=zth_id)
    end if
    if (current_state%p%active) then 
      call define_velocity_variable(ncid, multi_process, z_dim_id, y_dim_id, x_dim_id, field_name=p_key, field_id=p_id)
    end if
  end subroutine define_prognostic_variables

  subroutine define_misc_variables(ncid, timestep_id, time_id, ugal_id, vgal_id, number_q_fields_id, &
       dtm_id, dtm_new_id, absolute_new_dtm_id)
    integer, intent(in) :: ncid
    integer, intent(out) :: timestep_id, time_id, ugal_id, vgal_id, number_q_fields_id, dtm_id, dtm_new_id, absolute_new_dtm_id

    call check_status(nf90_def_var(ncid, timestep, nf90_int, timestep_id))
    call check_status(nf90_def_var(ncid, time_key, nf90_double, time_id))
    call check_status(nf90_def_var(ncid, ugal, nf90_double, ugal_id))
    call check_status(nf90_def_var(ncid, vgal, nf90_double, vgal_id))
    call check_status(nf90_def_var(ncid, nqfields, nf90_int, number_q_fields_id))
    call check_status(nf90_def_var(ncid, dtm_key, nf90_double, dtm_id))
    call check_status(nf90_def_var(ncid, dtm_new_key, nf90_double, dtm_new_id))
    call check_status(nf90_def_var(ncid, absolute_new_dtm_key, nf90_double, absolute_new_dtm_id))
  end subroutine define_misc_variables

  subroutine write_out_misc_variables(current_state, ncid, timestep_id, time_id, ugal_id, &
       vgal_id, number_q_fields_id, dtm_id, dtm_new_id, absolute_new_dtm_id)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: ncid, timestep_id, time_id, ugal_id, vgal_id, number_q_fields_id, &
         dtm_id, dtm_new_id, absolute_new_dtm_id

    call check_status(nf90_put_var(ncid, timestep_id, current_state%timestep))
    ! The time is incremented with dtm as the model was about to increment for the next step and this is needed for diagnostics
    call check_status(nf90_put_var(ncid, time_id, current_state%time+current_state%dtm))
    call check_status(nf90_put_var(ncid, ugal_id, current_state%ugal))
    call check_status(nf90_put_var(ncid, vgal_id, current_state%vgal))
    call check_status(nf90_put_var(ncid, number_q_fields_id, current_state%number_q_fields))
    call check_status(nf90_put_var(ncid, dtm_id, current_state%dtm))
    call check_status(nf90_put_var(ncid, dtm_new_id, current_state%dtm_new))
    call check_status(nf90_put_var(ncid, absolute_new_dtm_id, current_state%absolute_new_dtm))
  end subroutine write_out_misc_variables

  subroutine define_velocity_variable(ncid, multi_process, dimone, dimtwo, dimthree, field_name, field_id)
    integer, intent(in) :: ncid, dimone
    integer, intent(in), optional :: dimtwo, dimthree
    integer, intent(out) :: field_id
    character(len=*), intent(in) :: field_name
    logical, intent(in) :: multi_process

    integer, dimension(:), allocatable :: dimids

    if (present(dimtwo) .and. present(dimthree)) then
      allocate(dimids(3))
      dimids = (/ dimone, dimtwo, dimthree /)
    else if (present(dimtwo) .or. present(dimthree)) then
      allocate(dimids(2))
      dimids = (/ dimone, merge(dimtwo, dimthree, present(dimtwo)) /)
    else
      allocate(dimids(1))
      dimids = (/ dimone /)
    end if

    call check_status(nf90_def_var(ncid, field_name, merge(nf90_double, nf90_real, default_precision == double_precision), &
         dimids, field_id))
    if (multi_process) then
      call check_status(nf90_def_var_fill(ncid, field_id, 1, 1))
      call check_status(nf90_var_par_access(ncid, field_id, nf90_collective))
    end if
    call check_status(nf90_put_att(ncid, field_id, "units", "m/s"))
  end subroutine define_velocity_variable
end module checkpointer_write_checkpoint_mod