Add restart subroutine to MOM.F90

As described in issue #372, I would like to be able to create restart files that contain information about the particle location. These files will be written at the same time as other restart files. I cannot add these calls directly to the driver, because the driver does not have information about the particle location. We have added save_MOM6_internal_state as a subroutine in MOM.F90, and we added calls to this subroutine from each of the drivers. We hope this will allow for more new packages to write restart files in the future. Co-authored by Spencer Jones <[email protected]>
NOAA-GFDL · Jul 12, 2023 · 7970347 · 7970347
1 parent 77b5881
commit 7970347
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 9 deletions.
diff --git a/config_src/drivers/FMS_cap/ocean_model_MOM.F90 b/config_src/drivers/FMS_cap/ocean_model_MOM.F90
@@ -14,7 +14,7 @@ module ocean_model_mod
 use MOM, only : initialize_MOM, step_MOM, MOM_control_struct, MOM_end
 use MOM, only : extract_surface_state, allocate_surface_state, finish_MOM_initialization
 use MOM, only : get_MOM_state_elements, MOM_state_is_synchronized
-use MOM, only : get_ocean_stocks, step_offline
+use MOM, only : get_ocean_stocks, step_offline, save_MOM6_internal_state
 use MOM_coms,      only : field_chksum
 use MOM_constants, only : CELSIUS_KELVIN_OFFSET, hlf
 use MOM_coupler_types, only : coupler_1d_bc_type, coupler_2d_bc_type
@@ -700,6 +700,7 @@ subroutine ocean_model_restart(OS, timestamp)
     if (OS%use_ice_shelf) then
       call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, OS%dirs%restart_output_dir, .true.)
     endif
+    call save_MOM6_internal_state(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time, .true.)
   endif
   if (BTEST(OS%Restart_control,0)) then
     call save_restart(OS%dirs%restart_output_dir, OS%Time, OS%grid, &
@@ -709,6 +710,7 @@ subroutine ocean_model_restart(OS, timestamp)
     if (OS%use_ice_shelf) then
       call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, OS%dirs%restart_output_dir)
     endif
+    call save_MOM6_internal_state(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time)
   endif
 
 end subroutine ocean_model_restart
@@ -764,6 +766,8 @@ subroutine ocean_model_save_restart(OS, Time, directory, filename_suffix)
     call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, OS%dirs%restart_output_dir)
   endif
 
+  call save_MOM6_internal_state(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time)
+
 end subroutine ocean_model_save_restart
 
 !> Initialize the public ocean type

diff --git a/config_src/drivers/mct_cap/mom_ocean_model_mct.F90 b/config_src/drivers/mct_cap/mom_ocean_model_mct.F90
@@ -14,7 +14,7 @@ module MOM_ocean_model_mct
 use MOM,                      only : initialize_MOM, step_MOM, MOM_control_struct, MOM_end
 use MOM,                      only : extract_surface_state, allocate_surface_state, finish_MOM_initialization
 use MOM,                      only : get_MOM_state_elements, MOM_state_is_synchronized
-use MOM,                      only : get_ocean_stocks, step_offline
+use MOM,                      only : get_ocean_stocks, step_offline, save_MOM6_internal_state
 use MOM_coms,                 only : field_chksum
 use MOM_constants,            only : CELSIUS_KELVIN_OFFSET, hlf
 use MOM_diag_mediator,        only : diag_ctrl, enable_averages, disable_averaging
@@ -697,6 +697,7 @@ subroutine ocean_model_restart(OS, timestamp, restartname)
       call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, &
              OS%dirs%restart_output_dir)
     endif
+    call save_MOM6_internal_state(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time)
   else
     if (BTEST(OS%Restart_control,1)) then
       call save_restart(OS%dirs%restart_output_dir, OS%Time, OS%grid, &
@@ -706,6 +707,7 @@ subroutine ocean_model_restart(OS, timestamp, restartname)
       if (OS%use_ice_shelf) then
         call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, OS%dirs%restart_output_dir, .true.)
       endif
+      call save_MOM6_internal_state(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time, .true.)
     endif
     if (BTEST(OS%Restart_control,0)) then
       call save_restart(OS%dirs%restart_output_dir, OS%Time, OS%grid, &
@@ -715,6 +717,7 @@ subroutine ocean_model_restart(OS, timestamp, restartname)
       if (OS%use_ice_shelf) then
         call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, OS%dirs%restart_output_dir)
       endif
+      call save_MOM6_internal_state(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time)
     endif
   endif
 

diff --git a/config_src/drivers/nuopc_cap/mom_ocean_model_nuopc.F90 b/config_src/drivers/nuopc_cap/mom_ocean_model_nuopc.F90
@@ -14,7 +14,7 @@ module MOM_ocean_model_nuopc
 use MOM,                     only : initialize_MOM, step_MOM, MOM_control_struct, MOM_end
 use MOM,                     only : extract_surface_state, allocate_surface_state, finish_MOM_initialization
 use MOM,                     only : get_MOM_state_elements, MOM_state_is_synchronized
-use MOM,                     only : get_ocean_stocks, step_offline
+use MOM,                     only : get_ocean_stocks, step_offline, save_MOM6_internal_state
 use MOM_coms,                only : field_chksum
 use MOM_constants,           only : CELSIUS_KELVIN_OFFSET, hlf
 use MOM_diag_mediator,       only : diag_ctrl, enable_averages, disable_averaging
@@ -738,6 +738,8 @@ subroutine ocean_model_restart(OS, timestamp, restartname, stoch_restartname, nu
       call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, &
            OS%dirs%restart_output_dir)
     endif
+
+    call save_MOM6_internal_state(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time)
   else
     if (BTEST(OS%Restart_control,1)) then
       call save_restart(OS%dirs%restart_output_dir, OS%Time, OS%grid, &
@@ -747,6 +749,7 @@ subroutine ocean_model_restart(OS, timestamp, restartname, stoch_restartname, nu
       if (OS%use_ice_shelf) then
         call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, OS%dirs%restart_output_dir, .true.)
       endif
+      call save_MOM6_internal_state(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time, .true.)
     endif
     if (BTEST(OS%Restart_control,0)) then
       call save_restart(OS%dirs%restart_output_dir, OS%Time, OS%grid, &
@@ -756,6 +759,7 @@ subroutine ocean_model_restart(OS, timestamp, restartname, stoch_restartname, nu
       if (OS%use_ice_shelf) then
         call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, OS%dirs%restart_output_dir)
       endif
+      call save_MOM6_internal_state(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time)
     endif
   endif
   if (present(stoch_restartname)) then
@@ -818,6 +822,8 @@ subroutine ocean_model_save_restart(OS, Time, directory, filename_suffix)
     call ice_shelf_save_restart(OS%Ice_shelf_CSp, OS%Time, OS%dirs%restart_output_dir)
   endif
 
+  call save_MOM6_internal_start(OS%MOM_CSp, OS%dirs%restart_output_dir, OS%Time)
+
 end subroutine ocean_model_save_restart
 
 !> Initialize the public ocean type

diff --git a/config_src/drivers/solo_driver/MOM_driver.F90 b/config_src/drivers/solo_driver/MOM_driver.F90
@@ -31,7 +31,7 @@ program MOM6
   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,                 only : step_offline, save_MOM6_internal_state
   use MOM_coms,            only : Set_PElist
   use MOM_domains,         only : MOM_infra_init, MOM_infra_end, set_MOM_thread_affinity
   use MOM_ensemble_manager, only : ensemble_manager_init, get_ensemble_size
@@ -570,6 +570,7 @@ program MOM6
                             dirs%restart_output_dir, .true.)
         if (use_ice_shelf) call ice_shelf_save_restart(ice_shelf_CSp, Time, &
                                     dirs%restart_output_dir, .true.)
+        call save_MOM6_internal_state(MOM_CSp, dirs%restart_output_dir, Time, .true.)
       endif
       if (BTEST(Restart_control,0)) then
         call save_restart(dirs%restart_output_dir, Time, grid, &
@@ -578,6 +579,7 @@ program MOM6
                             dirs%restart_output_dir)
         if (use_ice_shelf) call ice_shelf_save_restart(ice_shelf_CSp, Time, &
                                     dirs%restart_output_dir)
+        call save_MOM6_internal_state(MOM_CSp, dirs%restart_output_dir, Time)
       endif
       restart_time = restart_time + restint
     endif
@@ -601,6 +603,8 @@ program MOM6
     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)
+    call save_MOM6_internal_state(MOM_CSp, dirs%restart_output_dir, Time)
+
     ! Write the ocean solo restart file.
     call write_ocean_solo_res(Time, Start_time, calendar_type, &
                               trim(dirs%restart_output_dir)//'ocean_solo.res')

diff --git a/src/core/MOM.F90 b/src/core/MOM.F90
@@ -438,7 +438,7 @@ module MOM
 end type MOM_control_struct
 
 public initialize_MOM, finish_MOM_initialization, MOM_end
-public step_MOM, step_offline
+public step_MOM, step_offline, save_MOM6_internal_state
 public extract_surface_state, get_ocean_stocks
 public get_MOM_state_elements, MOM_state_is_synchronized
 public allocate_surface_state, deallocate_surface_state
@@ -3904,14 +3904,34 @@ subroutine get_ocean_stocks(CS, mass, heat, salt, on_PE_only)
 
 end subroutine get_ocean_stocks
 
+
+!> Trigger a writing of restarts for the MOM6 internal state
+!!
+!! Currently this applies to the state that does not take the form
+!! of simple arrays for which the generic save_restart() function
+!! can be used.
+!!
+!! Todo:
+!! [ ] update particles to use Time and directories
+!! [ ] move the call to generic save_restart() in here.
+subroutine save_MOM6_internal_state(CS, dirs, time, stamp_time)
+  type(MOM_control_struct), intent(inout) :: CS         !< MOM control structure
+  character(len=*),         intent(in)    :: dirs  !< The directory where the restart
+                                                        !! files are to be written
+  type(time_type),          intent(in)    :: time       !< The current model time
+  logical,       optional,  intent(in)    :: stamp_time !< If present and true, add time-stamp
+
+    ! Could call save_restart(CS%restart_CSp) here
+
+    if (CS%use_particles) call particles_save_restart(CS%particles)
+
+end subroutine save_MOM6_internal_state
+
+
 !> End of ocean model, including memory deallocation
 subroutine MOM_end(CS)
   type(MOM_control_struct), intent(inout) :: CS   !< MOM control structure
 
-  if (CS%use_particles) then
-    call particles_save_restart(CS%particles)
-  endif
-
   call MOM_sum_output_end(CS%sum_output_CSp)
 
   if (CS%use_ALE_algorithm) call ALE_end(CS%ALE_CSp)