MOM_driver.F90

program mom_main

! This file is part of MOM6. See LICENSE.md for the license.

!********+*********+*********+*********+*********+*********+*********+**
!*                                                                     *
!*                  The Modular Ocean Model, version 6                 *
!*                               MOM6                                  *
!*                                                                     *
!*  By Alistair Adcroft, Stephen Griffies and Robert Hallberg          *
!*                                                                     *
!*    This file is the ocean-only driver for Version 6 of the Modular  *
!*  Ocean Model (MOM).  A separate ocean interface for use with        *
!*  coupled models is provided in ocean_model_MOM.F90.   These two     *
!*  drivers are kept in separate directories for convenience of code   *
!*  selection during compiling.  This file orchestrates the calls to   *
!*  the MOM initialization routines, to the subroutine that steps      *
!*  the model, and coordinates the output and saving restarts.  A      *
!*  description of all of the files that constitute MOM is found in    *
!*  the comments at the beginning of MOM.F90.  The arguments of each   *
!*  subroutine are described where the subroutine is defined.          *
!*                                                                     *
!*  Macros written all in capital letters are defined in MOM_memory.h. *
!*                                                                     *
!********+*********+*********+*********+*********+*********+*********+**

  use mom_cpu_clock,       only : cpu_clock_id, cpu_clock_begin, cpu_clock_end
  use mom_cpu_clock,       only : clock_component
  use mom_diag_mediator,   only : enable_averaging, disable_averaging, diag_mediator_end
  use mom_diag_mediator,   only : diag_ctrl, diag_mediator_close_registration
  use mom,                 only : initialize_mom, step_mom, mom_control_struct, mom_end
  use mom,                 only : extract_surface_state, finish_mom_initialization
  use mom,                 only : get_mom_state_elements, mom_state_is_synchronized
  use mom,                 only : step_offline
  use mom_domains,         only : mom_infra_init, mom_infra_end
  use mom_error_handler,   only : mom_error, mom_mesg, warning, fatal, is_root_pe
  use mom_error_handler,   only : calltree_enter, calltree_leave, calltree_waypoint
  use mom_file_parser,     only : read_param, get_param, log_param, log_version, param_file_type
  use mom_file_parser,     only : close_param_file
  use mom_forcing_type,    only : forcing, mech_forcing, forcing_diagnostics
  use mom_forcing_type,    only : mech_forcing_diags, mom_forcing_chksum, mom_mech_forcing_chksum
  use mom_get_input,       only : directories
  use mom_grid,            only : ocean_grid_type
  use mom_io,              only : file_exists, open_file, close_file
  use mom_io,              only : check_nml_error, io_infra_init, io_infra_end
  use mom_io,              only : append_file, ascii_file, readonly_file, single_file
  use mom_restart,         only : mom_restart_cs, save_restart
  use mom_string_functions,only : uppercase
  use mom_surface_forcing, only : set_forcing, forcing_save_restart
  use mom_surface_forcing, only : surface_forcing_init, surface_forcing_cs
  use mom_time_manager,    only : time_type, set_date, get_date
  use mom_time_manager,    only : real_to_time, time_type_to_real
  use mom_time_manager,    only : operator(+), operator(-), operator(*), operator(/)
  use mom_time_manager,    only : operator(>), operator(<), operator(>=)
  use mom_time_manager,    only : increment_date, set_calendar_type, month_name
  use mom_time_manager,    only : julian, gregorian, noleap, thirty_day_months
  use mom_time_manager,    only : no_calendar
  use mom_tracer_flow_control, only : tracer_flow_control_cs
  use mom_unit_scaling,    only : unit_scale_type
  use mom_variables,       only : surface
  use mom_verticalgrid,    only : verticalgrid_type
  use mom_write_cputime,   only : write_cputime, mom_write_cputime_init
  use mom_write_cputime,   only : write_cputime_start_clock, write_cputime_cs

  use ensemble_manager_mod, only : ensemble_manager_init, get_ensemble_size
  use ensemble_manager_mod, only : ensemble_pelist_setup
  use mpp_mod, only : set_current_pelist => mpp_set_current_pelist
  use time_interp_external_mod, only : time_interp_external_init
  use fms_affinity_mod,     only : fms_affinity_init, fms_affinity_set,fms_affinity_get

  use mom_ice_shelf, only : initialize_ice_shelf, ice_shelf_end, ice_shelf_cs
  use mom_ice_shelf, only : shelf_calc_flux, add_shelf_forces, ice_shelf_save_restart
! , add_shelf_flux_forcing, add_shelf_flux_IOB

  use mom_wave_interface, only: wave_parameters_cs, mom_wave_interface_init
  use mom_wave_interface, only: mom_wave_interface_init_lite, update_surface_waves

  implicit none

#include <MOM_memory.h>

  ! A structure with the driving mechanical surface forces
  type(mech_forcing) :: forces
  ! A structure containing pointers to the thermodynamic forcing fields
  ! at the ocean surface.
  type(forcing) :: fluxes

  ! A structure containing pointers to the ocean surface state fields.
  type(surface) :: sfc_state

  ! A pointer to a structure containing metrics and related information.
  type(ocean_grid_type), pointer :: grid
  type(verticalGrid_type), pointer :: GV
  ! A pointer to a structure containing dimensional unit scaling factors.
  type(unit_scale_type), pointer :: US

  ! If .true., use the ice shelf model for part of the domain.
  logical :: use_ice_shelf

  ! If .true., use surface wave coupling
  logical :: use_waves = .false.

  ! This is .true. if incremental restart files may be saved.
  logical :: permit_incr_restart = .true.

  integer :: ns

  ! nmax is the number of iterations after which to stop so that the
  ! simulation does not exceed its CPU time limit.  nmax is determined by
  ! evaluating the CPU time used between successive calls to write_cputime.
  ! Initially it is set to be very large.
  integer :: nmax=2000000000

  ! A structure containing several relevant directory paths.
  type(directories) :: dirs

  ! A suite of time types for use by MOM
  type(time_type), target :: Time       ! A copy of the ocean model's time.
                                        ! Other modules can set pointers to this and
                                        ! change it to manage diagnostics.
  type(time_type) :: Master_Time        ! The ocean model's master clock. No other
                                        ! modules are ever given access to this.
  type(time_type) :: Time1              ! The value of the ocean model's time at the
                                        ! start of a call to step_MOM.
  type(time_type) :: Start_time         ! The start time of the simulation.
  type(time_type) :: segment_start_time ! The start time of this run segment.
  type(time_type) :: Time_end           ! End time for the segment or experiment.
  type(time_type) :: restart_time       ! The next time to write restart files.
  type(time_type) :: Time_step_ocean    ! A time_type version of dt_forcing.

  real    :: elapsed_time = 0.0   ! Elapsed time in this run  [s].
  logical :: elapsed_time_master  ! If true, elapsed time is used to set the
                                  ! model's master clock (Time).  This is needed
                                  ! if Time_step_ocean is not an exact
                                  ! representation of dt_forcing.
  real :: dt_forcing              ! The coupling time step [s].
  real :: dt                      ! The nominal baroclinic dynamics time step [s].
  real :: dt_off                  ! Offline time step [s].
  integer :: ntstep               ! The number of baroclinic dynamics time steps
                                  ! within dt_forcing.
  real :: dt_therm                ! The thermodynamic timestep [s]
  real :: dt_dyn                  ! The actual dynamic timestep used [s].  The value of dt_dyn is
                                  ! chosen so that dt_forcing is an integer multiple of dt_dyn.
  real :: dtdia                   ! The diabatic timestep [s]
  real :: t_elapsed_seg           ! The elapsed time in this run segment [s]
  integer :: n, n_max, nts, n_last_thermo
  logical :: diabatic_first, single_step_call
  type(time_type) :: Time2, time_chg

  integer :: Restart_control    ! An integer that is bit-tested to determine whether
                                ! incremental restart files are saved and whether they
                                ! have a time stamped name.  +1 (bit 0) for generic
                                ! files and +2 (bit 1) for time-stamped files.  A
                                ! restart file is saved at the end of a run segment
                                ! unless Restart_control is negative.

  real            :: Time_unit       ! The time unit for the following input fields [s].
  type(time_type) :: restint         ! The time between saves of the restart file.
  type(time_type) :: daymax          ! The final day of the simulation.

  integer :: CPU_steps          ! The number of steps between writing CPU time.
  integer :: date_init(6)=0                ! The start date of the whole simulation.
  integer :: date(6)=-1                    ! Possibly the start date of this run segment.
  integer :: years=0, months=0, days=0     ! These may determine the segment run
  integer :: hours=0, minutes=0, seconds=0 ! length, if read from a namelist.
  integer :: yr, mon, day, hr, mins, sec   ! Temp variables for writing the date.
  type(param_file_type) :: param_file      ! The structure indicating the file(s)
                                           ! containing all run-time parameters.
  character(len=9)  :: month
  character(len=16) :: calendar = 'julian'
  integer :: calendar_type=-1

  integer :: unit, io_status, ierr
  integer :: ensemble_size, nPEs_per, ensemble_info(6)

  integer, dimension(0) :: atm_PElist, land_PElist, ice_PElist
  integer, dimension(:), allocatable :: ocean_PElist
  logical :: unit_in_use
  integer :: initClock, mainClock, termClock

  logical :: debug               ! If true, write verbose checksums for debugging purposes.
  logical :: offline_tracer_mode ! If false, use the model in prognostic mode where
                                 ! the barotropic and baroclinic dynamics, thermodynamics,
                                 ! etc. are stepped forward integrated in time.
                                 ! If true, then all of the above are bypassed with all
                                 ! fields necessary to integrate only the tracer advection
                                 ! and diffusion equation are read in from files stored from
                                 ! a previous integration of the prognostic model

  type(MOM_control_struct),  pointer :: MOM_CSp => null()
  !> A pointer to the tracer flow control structure.
  type(tracer_flow_control_CS), pointer :: &
    tracer_flow_CSp => null()  !< A pointer to the tracer flow control structure
  type(surface_forcing_CS),  pointer :: surface_forcing_CSp => null()
  type(write_cputime_CS),    pointer :: write_CPU_CSp => null()
  type(ice_shelf_CS),        pointer :: ice_shelf_CSp => null()
  type(wave_parameters_cs),  pointer :: waves_CSp => null()
  type(MOM_restart_CS),      pointer :: &
    restart_CSp => null()     !< A pointer to the restart control structure
                              !! that will be used for MOM restart files.
  type(diag_ctrl), pointer :: &
    diag => null()            !< A pointer to the diagnostic regulatory structure
  !-----------------------------------------------------------------------

  character(len=4), parameter :: vers_num = 'v2.0'
! This include declares and sets the variable "version".
#include "version_variable.h"
  character(len=40)  :: mod_name = "MOM_main (MOM_driver)" ! This module's name.

  integer :: ocean_nthreads = 1
  logical :: use_hyper_thread = .false.
  integer :: omp_get_num_threads,omp_get_thread_num
  namelist /ocean_solo_nml/ date_init, calendar, months, days, hours, minutes, seconds,&
                            ocean_nthreads, use_hyper_thread

  !=====================================================================

  call write_cputime_start_clock(write_cpu_csp)

  call mom_infra_init() ; call io_infra_init()

  ! Initialize the ensemble manager.  If there are no settings for ensemble_size
  ! in input.nml(ensemble.nml), these should not do anything.  In coupled
  ! configurations, this all occurs in the external driver.
  call ensemble_manager_init() ; ensemble_info(:) =  get_ensemble_size()
  ensemble_size=ensemble_info(1) ; npes_per=ensemble_info(2)
  if (ensemble_size > 1) then ! There are multiple ensemble members.
    allocate(ocean_pelist(npes_per))
    call ensemble_pelist_setup(.true., 0, npes_per, 0, 0, atm_pelist, ocean_pelist, &
                               land_pelist, ice_pelist)
    call set_current_pelist(ocean_pelist)
    deallocate(ocean_pelist)
  endif

  ! These clocks are on the global pelist.
  initclock = cpu_clock_id( 'Initialization' )
  mainclock = cpu_clock_id( 'Main loop' )
  termclock = cpu_clock_id( 'Termination' )
  call cpu_clock_begin(initclock)

  call mom_mesg('======== Model being driven by MOM_driver ========', 2)
  call calltree_waypoint("Program MOM_main, MOM_driver.F90")

  if (file_exists('input.nml')) then
    ! Provide for namelist specification of the run length and calendar data.
    call open_file(unit, 'input.nml', form=ascii_file, action=readonly_file)
    read(unit, ocean_solo_nml, iostat=io_status)
    call close_file(unit)
    ierr = check_nml_error(io_status,'ocean_solo_nml')
    if (years+months+days+hours+minutes+seconds > 0) then
      if (is_root_pe()) write(*,ocean_solo_nml)
    endif
  endif

!$  call fms_affinity_init
!$  call fms_affinity_set('OCEAN', use_hyper_thread, ocean_nthreads)
!$  call omp_set_num_threads(ocean_nthreads)
!$OMP PARALLEL
!$  write(6,*) "ocean_solo OMPthreading ", fms_affinity_get(), omp_get_thread_num(), omp_get_num_threads()
!$  call flush(6)
!$OMP END PARALLEL

  ! Read ocean_solo restart, which can override settings from the namelist.
  if (file_exists(trim(dirs%restart_input_dir)//'ocean_solo.res')) then
    call open_file(unit,trim(dirs%restart_input_dir)//'ocean_solo.res', &
                   form=ascii_file,action=readonly_file)
    read(unit,*) calendar_type
    read(unit,*) date_init
    read(unit,*) date
    call close_file(unit)
  else
    calendar = uppercase(calendar)
    if (calendar(1:6) == 'JULIAN') then ;         calendar_type = julian
    elseif (calendar(1:9) == 'GREGORIAN') then ; calendar_type = gregorian
    elseif (calendar(1:6) == 'NOLEAP') then ;    calendar_type = noleap
    elseif (calendar(1:10)=='THIRTY_DAY') then ; calendar_type = thirty_day_months
    elseif (calendar(1:11)=='NO_CALENDAR') then; calendar_type = no_calendar
    elseif (calendar(1:1) /= ' ') then
      call mom_error(fatal,'MOM_driver: Invalid namelist value '//trim(calendar)//' for calendar')
    else
      call mom_error(fatal,'MOM_driver: No namelist value for calendar')
    endif
  endif
  call set_calendar_type(calendar_type)


  if (sum(date_init) > 0) then
    start_time = set_date(date_init(1),date_init(2), date_init(3), &
         date_init(4),date_init(5),date_init(6))
  else
    start_time = real_to_time(0.0)
  endif

  call time_interp_external_init

  if (sum(date) >= 0) then
    ! In this case, the segment starts at a time fixed by ocean_solo.res
    segment_start_time = set_date(date(1),date(2),date(3),date(4),date(5),date(6))
    time = segment_start_time
    call initialize_mom(time, start_time, param_file, dirs, mom_csp, restart_csp, &
                        segment_start_time, offline_tracer_mode=offline_tracer_mode, &
                        diag_ptr=diag, tracer_flow_csp=tracer_flow_csp)
  else
    ! In this case, the segment starts at a time read from the MOM restart file
    ! or left as Start_time by MOM_initialize.
    time = start_time
    call initialize_mom(time, start_time, param_file, dirs, mom_csp, restart_csp, &
                        offline_tracer_mode=offline_tracer_mode, diag_ptr=diag, &
                        tracer_flow_csp=tracer_flow_csp)
  endif

  call get_mom_state_elements(mom_csp, g=grid, gv=gv, us=us, c_p_scaled=fluxes%C_p)
  master_time = time

  call calltree_waypoint("done initialize_MOM")

  call extract_surface_state(mom_csp, sfc_state)

  call surface_forcing_init(time, grid, us, param_file, diag, &
                            surface_forcing_csp, tracer_flow_csp)
  call calltree_waypoint("done surface_forcing_init")

  call get_param(param_file, mod_name, "ICE_SHELF", use_ice_shelf, &
                 "If true, enables the ice shelf model.", default=.false.)
  if (use_ice_shelf) then
    ! These arrays are not initialized in most solo cases, but are needed
    ! when using an ice shelf
    call initialize_ice_shelf(param_file, grid, time, ice_shelf_csp, &
                              diag, forces, fluxes)
  endif

  call get_param(param_file,mod_name,"USE_WAVES",use_waves,&
       "If true, enables surface wave modules.",default=.false.)
  if (use_waves) then
    call mom_wave_interface_init(time, grid, gv, us, param_file, waves_csp, diag)
  else
    call mom_wave_interface_init_lite(param_file)
  endif

  segment_start_time = time
  elapsed_time = 0.0

  ! Read all relevant parameters and write them to the model log.
  call log_version(param_file, mod_name, version, "")
  call get_param(param_file, mod_name, "DT", dt, fail_if_missing=.true.)
  call get_param(param_file, mod_name, "DT_FORCING", dt_forcing, &
                 "The time step for changing forcing, coupling with other "//&
                 "components, or potentially writing certain diagnostics. "//&
                 "The default value is given by DT.", units="s", default=dt)
  if (offline_tracer_mode) then
    call get_param(param_file, mod_name, "DT_OFFLINE", dt_forcing, &
                   "Time step for the offline time step")
    dt = dt_forcing
  endif
  ntstep = max(1,ceiling(dt_forcing/dt - 0.001))

  time_step_ocean = real_to_time(dt_forcing)
  elapsed_time_master = (abs(dt_forcing - time_type_to_real(time_step_ocean)) > 1.0e-12*dt_forcing)
  if (elapsed_time_master) &
    call mom_mesg("Using real elapsed time for the master clock.", 2)

  ! Determine the segment end time, either from the namelist file or parsed input file.
  call get_param(param_file, mod_name, "TIMEUNIT", time_unit, &
                 "The time unit for DAYMAX, ENERGYSAVEDAYS, and RESTINT.", &
                 units="s", default=86400.0)
  if (years+months+days+hours+minutes+seconds > 0) then
    time_end = increment_date(time, years, months, days, hours, minutes, seconds)
    call mom_mesg('Segment run length determined from ocean_solo_nml.', 2)
    call get_param(param_file, mod_name, "DAYMAX", daymax, timeunit=time_unit, &
                   default=time_end, do_not_log=.true.)
    call log_param(param_file, mod_name, "DAYMAX", daymax, &
                 "The final time of the whole simulation, in units of "//&
                 "TIMEUNIT seconds.  This also sets the potential end "//&
                 "time of the present run segment if the end time is "//&
                 "not set via ocean_solo_nml in input.nml.", &
                 timeunit=time_unit)
  else
    call get_param(param_file, mod_name, "DAYMAX", daymax, &
                 "The final time of the whole simulation, in units of "//&
                 "TIMEUNIT seconds.  This also sets the potential end "//&
                 "time of the present run segment if the end time is "//&
                 "not set via ocean_solo_nml in input.nml.", &
                 timeunit=time_unit, fail_if_missing=.true.)
    time_end = daymax
  endif

  call get_param(param_file, mod_name, "SINGLE_STEPPING_CALL", single_step_call, &
                 "If true, advance the state of MOM with a single step "//&
                 "including both dynamics and thermodynamics.  If false "//&
                 "the two phases are advanced with separate calls.", default=.true.)
  call get_param(param_file, mod_name, "DT_THERM", dt_therm, &
                 "The thermodynamic and tracer advection time step. "//&
                 "Ideally DT_THERM should be an integer multiple of DT "//&
                 "and less than the forcing or coupling time-step, unless "//&
                 "THERMO_SPANS_COUPLING is true, in which case DT_THERM "//&
                 "can be an integer multiple of the coupling timestep.  By "//&
                 "default DT_THERM is set to DT.", units="s", default=dt)
  call get_param(param_file, mod_name, "DIABATIC_FIRST", diabatic_first, &
                 "If true, apply diabatic and thermodynamic processes, "//&
                 "including buoyancy forcing and mass gain or loss, "//&
                 "before stepping the dynamics forward.", default=.false.)


  if (time >= time_end) call mom_error(fatal, &
    "MOM_driver: The run has been started at or after the end time of the run.")

  call get_param(param_file, mod_name, "RESTART_CONTROL", restart_control, &
                 "An integer whose bits encode which restart files are "//&
                 "written. Add 2 (bit 1) for a time-stamped file, and odd "//&
                 "(bit 0) for a non-time-stamped file. A non-time-stamped "//&
                 "restart file is saved at the end of the run segment "//&
                 "for any non-negative value.", default=1)
  call get_param(param_file, mod_name, "RESTINT", restint, &
                 "The interval between saves of the restart file in units "//&
                 "of TIMEUNIT.  Use 0 (the default) to not save "//&
                 "incremental restart files at all.", default=real_to_time(0.0), &
                 timeunit=time_unit)
  call get_param(param_file, mod_name, "WRITE_CPU_STEPS", cpu_steps, &
                 "The number of coupled timesteps between writing the cpu "//&
                 "time. If this is not positive, do not check cpu time, and "//&
                 "the segment run-length can not be set via an elapsed CPU time.", &
                 default=1000)
  call get_param(param_file, "MOM", "DEBUG", debug, &
                 "If true, write out verbose debugging data.", &
                 default=.false., debuggingparam=.true.)

  call log_param(param_file, mod_name, "ELAPSED TIME AS MASTER", elapsed_time_master)

  if (cpu_steps > 0) &
    call mom_write_cputime_init(param_file, dirs%output_directory, start_time, &
                                write_cpu_csp)

  ! Close the param_file.  No further parsing of input is possible after this.
  call close_param_file(param_file)
  call diag_mediator_close_registration(diag)

  ! Write out a time stamp file.
  if (calendar_type /= no_calendar) then
    call open_file(unit, 'time_stamp.out', form=ascii_file, action=append_file, &
                   threading=single_file)
    call get_date(time, date(1), date(2), date(3), date(4), date(5), date(6))
    month = month_name(date(2))
    if (is_root_pe()) write(unit,'(6i4,2x,a3)') date, month(1:3)
    call get_date(time_end, date(1), date(2), date(3), date(4), date(5), date(6))
    month = month_name(date(2))
    if (is_root_pe()) write(unit,'(6i4,2x,a3)') date, month(1:3)
    call close_file(unit)
  endif

  if (cpu_steps > 0) call write_cputime(time, 0, write_cpu_csp)

  if (((.not.btest(restart_control,1)) .and. (.not.btest(restart_control,0))) &
      .or. (restart_control < 0)) permit_incr_restart = .false.

  if (restint > real_to_time(0.0)) then
    ! restart_time is the next integral multiple of restint.
    restart_time = start_time + restint * &
        (1 + ((time + time_step_ocean) - start_time) / restint)
  else
    ! Set the time so late that there is no intermediate restart.
    restart_time = time_end + time_step_ocean
    permit_incr_restart = .false.
  endif

  call cpu_clock_end(initclock) !end initialization

  call cpu_clock_begin(mainclock) !begin main loop

  ns = 1
  do while ((ns < nmax) .and. (time < time_end))
    call calltree_enter("Main loop, MOM_driver.F90",ns)

    ! Set the forcing for the next steps.
    if (.not. offline_tracer_mode) then
        call set_forcing(sfc_state, forces, fluxes, time, time_step_ocean, grid, us, &
                     surface_forcing_csp)
    endif
    if (debug) then
      call mom_mech_forcing_chksum("After set forcing", forces, grid, us, haloshift=0)
      call mom_forcing_chksum("After set forcing", fluxes, grid, us, haloshift=0)
    endif

    if (use_ice_shelf) then
      call shelf_calc_flux(sfc_state, fluxes, time, dt_forcing, ice_shelf_csp)
      call add_shelf_forces(grid, us, ice_shelf_csp, forces)
    endif
    fluxes%fluxes_used = .false.
    fluxes%dt_buoy_accum = us%s_to_T*dt_forcing

    if (use_waves) then
      call update_surface_waves(grid, gv, us, time, time_step_ocean, waves_csp)
    endif

    if (ns==1) then
      call finish_mom_initialization(time, dirs, mom_csp, restart_csp)
    endif

    ! This call steps the model over a time dt_forcing.
    time1 = master_time ; time = master_time
    if (offline_tracer_mode) then
      call step_offline(forces, fluxes, sfc_state, time1, dt_forcing, mom_csp)
    elseif (single_step_call) then
      call step_mom(forces, fluxes, sfc_state, time1, dt_forcing, mom_csp, waves=waves_csp)
    else
      n_max = 1 ; if (dt_forcing > dt) n_max = ceiling(dt_forcing/dt - 0.001)
      dt_dyn = dt_forcing / real(n_max)

      nts = max(1,min(n_max,floor(dt_therm/dt_dyn + 0.001)))
      n_last_thermo = 0

      time2 = time1 ; t_elapsed_seg = 0.0
      do n=1,n_max
        if (diabatic_first) then
          if (modulo(n-1,nts)==0) then
            dtdia = dt_dyn*min(ntstep,n_max-(n-1))
            call step_mom(forces, fluxes, sfc_state, time2, dtdia, mom_csp, &
                          do_dynamics=.false., do_thermodynamics=.true., &
                          start_cycle=(n==1), end_cycle=.false., cycle_length=dt_forcing)
          endif

          call step_mom(forces, fluxes, sfc_state, time2, dt_dyn, mom_csp, &
                        do_dynamics=.true., do_thermodynamics=.false., &
                        start_cycle=.false., end_cycle=(n==n_max), cycle_length=dt_forcing)
        else
          call step_mom(forces, fluxes, sfc_state, time2, dt_dyn, mom_csp, &
                        do_dynamics=.true., do_thermodynamics=.false., &
                        start_cycle=(n==1), end_cycle=.false., cycle_length=dt_forcing)

          if ((modulo(n,nts)==0) .or. (n==n_max)) then
            dtdia = dt_dyn*(n - n_last_thermo)
            ! Back up Time2 to the start of the thermodynamic segment.
            if (n > n_last_thermo+1) &
              time2 = time2 - real_to_time(dtdia - dt_dyn)
            call step_mom(forces, fluxes, sfc_state, time2, dtdia, mom_csp, &
                          do_dynamics=.false., do_thermodynamics=.true., &
                          start_cycle=.false., end_cycle=(n==n_max), cycle_length=dt_forcing)
            n_last_thermo = n
          endif
        endif

        t_elapsed_seg = t_elapsed_seg + dt_dyn
        time2 = time1 + real_to_time(t_elapsed_seg)
      enddo
    endif

!   Time = Time + Time_step_ocean
!   This is here to enable fractional-second time steps.
    elapsed_time = elapsed_time + dt_forcing
    if (elapsed_time > 2e9) then
      ! This is here to ensure that the conversion from a real to an integer can be accurately
      ! represented in long runs (longer than ~63 years). It will also ensure that elapsed time
      ! does not lose resolution of order the timetype's resolution, provided that the timestep and
      ! tick are larger than 10-5 seconds.  If a clock with a finer resolution is used, a smaller
      ! value would be required.
      time_chg = real_to_time(elapsed_time)
      segment_start_time = segment_start_time + time_chg
      elapsed_time = elapsed_time - time_type_to_real(time_chg)
    endif
    if (elapsed_time_master) then
      master_time = segment_start_time + real_to_time(elapsed_time)
    else
      master_time = master_time + time_step_ocean
    endif
    time = master_time

    if (cpu_steps > 0) then ; if (mod(ns, cpu_steps) == 0) then
      call write_cputime(time, ns+ntstep-1, write_cpu_csp, nmax)
    endif ; endif

    call mech_forcing_diags(forces, dt_forcing, grid, time, diag, surface_forcing_csp%handles)

    if (.not. offline_tracer_mode) then
      if (fluxes%fluxes_used) then
        call forcing_diagnostics(fluxes, sfc_state, grid, us, time, &
                                 diag, surface_forcing_csp%handles)
      else
        call mom_error(fatal, "The solo MOM_driver is not yet set up to handle "//&
               "thermodynamic time steps that are longer than the coupling timestep.")
      endif
    endif

!  See if it is time to write out a restart file - timestamped or not.
    if ((permit_incr_restart) .and. (fluxes%fluxes_used) .and. &
        (time + (time_step_ocean/2) > restart_time)) then
      if (btest(restart_control,1)) then
        call save_restart(dirs%restart_output_dir, time, grid, &
                          restart_csp, .true., gv=gv)
        call forcing_save_restart(surface_forcing_csp, grid, time, &
                            dirs%restart_output_dir, .true.)
        if (use_ice_shelf) call ice_shelf_save_restart(ice_shelf_csp, time, &
                                    dirs%restart_output_dir, .true.)
      endif
      if (btest(restart_control,0)) then
        call save_restart(dirs%restart_output_dir, time, grid, &
                          restart_csp, gv=gv)
        call forcing_save_restart(surface_forcing_csp, grid, time, &
                            dirs%restart_output_dir)
        if (use_ice_shelf) call ice_shelf_save_restart(ice_shelf_csp, time, &
                                    dirs%restart_output_dir)
      endif
      restart_time = restart_time + restint
    endif

    ns = ns + ntstep
    call calltree_leave("Main loop")
  enddo

  call cpu_clock_end(mainclock)
  call cpu_clock_begin(termclock)
  if (restart_control>=0) then
    if (.not.mom_state_is_synchronized(mom_csp)) &
      call mom_error(warning, "End of MOM_main reached with inconsistent "//&
         "dynamics and advective times.  Additional restart fields "//&
         "that have not been coded yet would be required for reproducibility.")
    if (.not.fluxes%fluxes_used .and. .not.offline_tracer_mode) call mom_error(fatal, &
         "End of MOM_main reached with unused buoyancy fluxes. "//&
         "For conservation, the ocean restart files can only be "//&
         "created after the buoyancy forcing is applied.")

    call save_restart(dirs%restart_output_dir, time, grid, restart_csp, gv=gv)
    if (use_ice_shelf) call ice_shelf_save_restart(ice_shelf_csp, time, &
                                dirs%restart_output_dir)
    ! Write ocean solo restart file.
    call open_file(unit, trim(dirs%restart_output_dir)//'ocean_solo.res', nohdrs=.true.)
    if (is_root_pe())then
        write(unit, '(i6,8x,a)') calendar_type, &
             '(Calendar: no_calendar=0, thirty_day_months=1, julian=2, gregorian=3, noleap=4)'

        call get_date(start_time, yr, mon, day, hr, mins, sec)
        write(unit, '(6i6,8x,a)') yr, mon, day, hr, mins, sec, &
             'Model start time:   year, month, day, hour, minute, second'
        call get_date(time, yr, mon, day, hr, mins, sec)
        write(unit, '(6i6,8x,a)') yr, mon, day, hr, mins, sec, &
             'Current model time: year, month, day, hour, minute, second'
    endif
    call close_file(unit)
  endif

  if (is_root_pe()) then
    do unit=10,1967
      INQUIRE(unit,opened=unit_in_use)
      if (.not.unit_in_use) exit
    enddo
    open(unit,file="exitcode",form="FORMATTED",status="REPLACE",action="WRITE")
    if (time < daymax) then
      write(unit,*) 9
    else
      write(unit,*) 0
    endif
    close(unit)
  endif

  call calltree_waypoint("End MOM_main")
  call diag_mediator_end(time, diag, end_diag_manager=.true.)
  if (cpu_steps > 0) call write_cputime(time, ns-1, write_cpu_csp, call_end=.true.)
  call cpu_clock_end(termclock)

  call io_infra_end ; call mom_infra_end

  call mom_end(mom_csp)
  if (use_ice_shelf) call ice_shelf_end(ice_shelf_csp)

end program mom_main