Ensemble filename and stoch restart fixes

config_src/drivers/nuopc_cap/mom_cap.F90

@@ -4,6 +4,7 @@ module MOM_cap_mod
 
 use MOM_domains,              only: get_domain_extent
 use MOM_io,                   only: stdout, io_infra_end
+use MOM_io,                   only: insert_ensemble_appendix
 use mpp_domains_mod,          only: mpp_get_compute_domains
 use mpp_domains_mod,          only: mpp_get_ntile_count, mpp_get_pelist, mpp_get_global_domain
 use mpp_domains_mod,          only: mpp_get_domain_npes
@@ -24,6 +25,7 @@ module MOM_cap_mod
 use MOM_ocean_model_nuopc,    only: ocean_model_init_sfc, ocean_model_flux_init
 use MOM_ocean_model_nuopc,    only: ocean_model_init, update_ocean_model, ocean_model_end
 use MOM_ocean_model_nuopc,    only: get_ocean_grid, get_eps_omesh, query_ocean_state
+use MOM_ocean_model_nuopc,    only: stoch_restart_needed
 use MOM_cap_time,             only: AlarmInit
 use MOM_cap_methods,          only: mom_import, mom_export, mom_set_geomtype, mod2med_areacor
 use MOM_cap_methods,          only: med2mod_areacor, state_diagnose
@@ -1728,7 +1730,7 @@ subroutine ModelAdvance(gcomp, rc)
   character(240)                         :: msgString
   character(ESMF_MAXSTR)                 :: casename
   integer                                :: iostat
-  integer                                :: writeunit
+  integer                                :: rpointer_unit
   integer                                :: localPet
   type(ESMF_VM)                          :: vm
   integer                                :: n, i
@@ -1940,30 +1942,27 @@ subroutine ModelAdvance(gcomp, rc)
           rpointer_filename = trim(rpointer_filename//timestamp)
         endif
 
-        write(restartname,'(A,".mom6.r",A)') &
-             trim(casename), timestamp
-        write(stoch_restartname,'(A,".mom6.r_stoch",A,".nc")') &
-             trim(casename), timestamp
+        write(restartname,'(A,".mom6.r",A)') trim(casename), timestamp
+        write(stoch_restartname,'(A,".mom6.r_stoch",A,".nc")')  trim(casename), timestamp
+
+        call insert_ensemble_appendix(stoch_restartname, ".mom6")
+
         call ESMF_LogWrite("MOM_cap: Writing restart :  "//trim(restartname), ESMF_LOGMSG_INFO)
         ! write restart file(s)
         call ocean_model_restart(ocean_state, restartname=restartname, &
                 stoch_restartname=stoch_restartname, num_rest_files=num_rest_files)
         if (localPet == 0) then
            ! Write name of restart file in the rpointer file - this is currently hard-coded for the ocean
-          open(newunit=writeunit, file=rpointer_filename, form='formatted', status='unknown', iostat=iostat)
+          open(newunit=rpointer_unit, file=rpointer_filename, form='formatted', status='unknown', iostat=iostat)
           if (iostat /= 0) then
             call ESMF_LogSetError(ESMF_RC_FILE_OPEN, &
                  msg=subname//' ERROR opening '//rpointer_filename, line=__LINE__, file=u_FILE_u, rcToReturn=rc)
             return
           endif
 
-          ! update restart file name to include the instance suffix
-          if (len_trim(inst_suffix) > 0) then
-            write(restartname, '(A,".mom6",A,".r",A)') trim(casename), trim(inst_suffix), timestamp
-          endif
-
-          write(writeunit,'(a)') trim(restartname)//'.nc'
+          call insert_ensemble_appendix(restartname, ".mom6")
 
+          write(rpointer_unit,'(a)') trim(restartname)//'.nc'
           if (num_rest_files > 1) then
             ! append i.th restart file name to rpointer
             do i=1, num_rest_files-1
@@ -1972,10 +1971,15 @@ subroutine ModelAdvance(gcomp, rc)
               else
                 write(suffix,'("_",I2)') i
               endif
-              write(writeunit,'(a)') trim(restartname) // trim(suffix) // '.nc'
+              write(rpointer_unit,'(a)') trim(restartname) // trim(suffix) // '.nc'
             enddo
           endif
-          close(writeunit)
+
+          if (stoch_restart_needed(ocean_state)) then
+            write(rpointer_unit,'(a)') trim(stoch_restartname)
+          endif
+
+          close(rpointer_unit)
         endif
       else  ! not cesm_coupled
         write(restartname,'(i4.4,2(i2.2),A,3(i2.2),A)') year, month, day,".", hour, minute, seconds, &

config_src/drivers/nuopc_cap/mom_ocean_model_nuopc.F90

@@ -82,6 +82,7 @@ module MOM_ocean_model_nuopc
 public ocean_public_type_chksum
 public get_ocean_grid, query_ocean_state
 public get_eps_omesh
+public stoch_restart_needed
 
 !> This type is used for communication with other components via the FMS coupler.
 !! The element names and types can be changed only with great deliberation, hence
@@ -790,7 +791,7 @@ subroutine ocean_model_restart(OS, timestamp, restartname, stoch_restartname, nu
     endif
   endif
   if (present(stoch_restartname)) then
-    if (OS%do_sppt .OR. OS%pert_epbl .OR. OS%do_skeb) then
+    if (stoch_restart_needed(OS)) then
       call write_stoch_restart_ocn(trim(stoch_restartname))
     endif
   endif
@@ -1148,7 +1149,6 @@ end subroutine Ocean_stock_pe
 
 !> Write out checksums for fields from the ocean surface state
 subroutine ocean_public_type_chksum(id, timestep, ocn)
-
   character(len=*),        intent(in) :: id  !< An identifying string for this call
   integer,                 intent(in) :: timestep !< The number of elapsed timesteps
   type(ocean_public_type), intent(in) :: ocn !< A structure containing various publicly
@@ -1174,17 +1174,23 @@ end subroutine ocean_public_type_chksum
 
 subroutine get_ocean_grid(OS, Gridp)
   ! Obtain the ocean grid.
-  type(ocean_state_type) :: OS
-  type(ocean_grid_type) , pointer :: Gridp
+  type(ocean_state_type), intent(in) :: OS
+  type(ocean_grid_type) , pointer, intent(out) :: Gridp
 
   Gridp => OS%grid
   return
 end subroutine get_ocean_grid
 
 !> Returns eps_omesh read from param file
 real function get_eps_omesh(OS)
-  type(ocean_state_type) :: OS
+  type(ocean_state_type), intent(in) :: OS
   get_eps_omesh = OS%eps_omesh; return
 end function
 
+!> Returns true if a stochastic restart file is needed
+logical function stoch_restart_needed(OS)
+  type(ocean_state_type), intent(in) :: OS
+  stoch_restart_needed = OS%do_sppt .OR. OS%pert_epbl .OR. OS%do_skeb
+end function stoch_restart_needed
+
 end module MOM_ocean_model_nuopc

src/framework/MOM_io.F90

@@ -58,6 +58,7 @@ module MOM_io
 public :: file_exists, open_ASCII_file, close_file
 public :: MOM_file, MOM_infra_file, MOM_netcdf_file
 public :: field_exists, get_filename_appendix
+public :: insert_ensemble_appendix
 public :: fieldtype, field_size, get_field_atts
 public :: axistype, get_axis_data
 public :: MOM_read_data, MOM_read_vector, read_field_chksum
@@ -2997,6 +2998,47 @@ subroutine MOM_write_field_0d(IO_handle, field_md, field, tstamp, fill_value, sc
   call IO_handle%write_field(field_md, scaled_val, tstamp=tstamp)
 end subroutine MOM_write_field_0d
 
+!> Insert the ensemble appendix into a filename. If provided, the appendix is inserted after
+!! the last occurrence of the insert_after substring in the filename.
+subroutine insert_ensemble_appendix(filename, insert_after)
+  character(len=*),           intent(inout) :: filename     !< The filename to which the appendix is inserted
+  character(len=*), optional, intent(in)    :: insert_after !< The string after which the appendix is inserted.
+                                                            !! If not provided or found, the appendix is inserted
+                                                            !! at the end of the filename.
+  ! Local variables
+  character(len=32) :: filename_appendix            ! ensemble appendix to be inserted into the filename
+  character(len=:), allocatable :: filename_tr      ! trimmed filename
+  character(len=:), allocatable :: insert_after_tr  ! trimmed insert_after
+  integer :: pos ! The filename string index after which the appendix is to be inserted
+
+  call get_filename_appendix(filename_appendix)
+  if (len_trim(filename_appendix) == 0) return
+
+  filename_tr = trim(adjustl(filename))
+  pos = len(filename_tr)
+
+  ! If insert_after is provided, find the last occurrence of insert_after in the filename and set pos accordingly.
+  if (present(insert_after)) then
+    insert_after_tr = trim(adjustl(insert_after))
+    pos = index(filename_tr, insert_after_tr, back=.true.)
+    if (pos == 0) then
+      call MOM_error(FATAL, "insert_ensemble_appendix: The string " // insert_after_tr // &
+                     " was not found in the filename " // filename_tr)
+    endif
+    pos = pos + len(insert_after_tr) - 1
+  endif
+
+  ! Insert the ensemble appendix into the filename. If the appendix is to be added to
+  ! the end of the filename, do so before the .nc extension if it exists.
+  if (pos>3 .and. pos == len(filename_tr)) then
+    if (filename_tr(pos-2:pos) == ".nc") then
+      pos = pos - 3 ! Position before the .nc extension
+    endif
+  endif
+  filename = filename_tr(1:pos) // trim(filename_appendix) // filename_tr(pos+1:)
+
+end subroutine insert_ensemble_appendix
+
 !> Given filename and fieldname, this subroutine returns the size of the field in the file
 subroutine field_size(filename, fieldname, sizes, field_found, no_domain, ndims, ncid_in)
   character(len=*),      intent(in)    :: filename  !< The name of the file to read

src/framework/MOM_restart.F90

@@ -13,10 +13,11 @@ module MOM_restart
 use MOM_io, only : create_MOM_file, file_exists
 use MOM_io, only : MOM_infra_file, MOM_field
 use MOM_io, only : MOM_read_data, read_data, MOM_write_field, field_exists
-use MOM_io, only : vardesc, var_desc, query_vardesc, modify_vardesc, get_filename_appendix
+use MOM_io, only : vardesc, var_desc, query_vardesc, modify_vardesc
 use MOM_io, only : MULTIPLE, READONLY_FILE, SINGLE_FILE
 use MOM_io, only : CENTER, CORNER, NORTH_FACE, EAST_FACE
 use MOM_io, only : axis_info, get_axis_info
+use MOM_io, only : insert_ensemble_appendix
 use MOM_string_functions, only : lowercase
 use MOM_time_manager,  only : time_type, time_type_to_real, real_to_time
 use MOM_time_manager,  only : days_in_month, get_date, set_date
@@ -132,7 +133,8 @@ module MOM_restart
   type(p4d), pointer :: var_ptr4d(:) => NULL()
   !>@}
   integer :: max_fields !< The maximum number of restart fields
-  character(len=32) :: restartfile_appx_prefix !< The prefix for the restart file appendix (i.e., ensemble id)
+  character(len=32) :: ensemble_appendix_prefix !< The prefix after which the ensemble id appendix is added
+                                                !! in output file names.
 end type MOM_restart_CS
 
 !> Register fields for restarts
@@ -1614,8 +1616,6 @@ subroutine save_restart(directory, time, G, CS, time_stamped, filename, GV, num_
   integer :: turns                      ! Number of quarter turns from input to model domain
   integer, parameter :: nmax_extradims = 5
   type(axis_info), dimension(:), allocatable :: extra_axes
-  integer :: prefix_index    ! The index of the first occurrence of prefix in the restart filename.
-  integer :: prefix_length   ! The length of the prefix string.
 
   turns = CS%turns
 
@@ -1654,31 +1654,7 @@ subroutine save_restart(directory, time, G, CS, time_stamped, filename, GV, num_
   endif ; endif
 
   ! Determine if there is a filename_appendix (used for ensemble runs).
-  call get_filename_appendix(filename_appendix)
-  if (len_trim(filename_appendix) > 0) then
-    length = len_trim(restartname)
-
-    ! Determine if a valid prefix for appendix is provided.
-    prefix_index = 0
-    prefix_length = len_trim(CS%restartfile_appx_prefix)
-    if (prefix_length > 0) prefix_index = index(restartname, trim(CS%restartfile_appx_prefix))
-
-    if (prefix_index == 0) then ! No prefix is found
-      if (restartname(length-2:length) == '.nc') then
-        restartname = restartname(1:length-3)//'.'//trim(filename_appendix)//'.nc'
-      else
-        restartname = restartname(1:length)  //'.'//trim(filename_appendix)
-      endif
-    else ! Prefix is found
-      if (restartname(length-2:length) == '.nc') then
-        restartname = restartname(1:prefix_index-1+prefix_length) // &
-                  trim(filename_appendix) // restartname(prefix_index+prefix_length:length-3) // '.nc'
-      else
-        restartname = restartname(1:prefix_index-1+prefix_length) // &
-                  trim(filename_appendix) // restartname(prefix_index+prefix_length:)
-      endif
-    endif
-  endif
+  call insert_ensemble_appendix(restartname, CS%ensemble_appendix_prefix)
 
   next_var = 1
   do while (next_var <= CS%novars )
@@ -2150,17 +2126,7 @@ function open_restart_units(filename, directory, G, CS, IO_handles, file_paths,
       still_looking = (num_restart <= 0) ! Avoid going through the file list twice.
       do while (still_looking)
         restartname = trim(CS%restartfile)
-
-        ! Determine if there is a filename_appendix (used for ensemble runs).
-        call get_filename_appendix(filename_appendix)
-        if (len_trim(filename_appendix) > 0) then
-          length = len_trim(restartname)
-          if (restartname(length-2:length) == '.nc') then
-            restartname = restartname(1:length-3)//'.'//trim(filename_appendix)//'.nc'
-          else
-            restartname = restartname(1:length)  //'.'//trim(filename_appendix)
-          endif
-        endif
+        call insert_ensemble_appendix(restartname, CS%ensemble_appendix_prefix)
         filepath = trim(directory) // trim(restartname)
 
         write(suffix,'("_",I0)') num_restart
@@ -2314,11 +2280,11 @@ subroutine restart_init(param_file, CS, restart_root)
                  "made from a run with a different mask_table than the current run, "//&
                  "in which case the checksums will not match and cause crash.",&
                  default=.true.)
-  call get_param(param_file, mdl, "RESTARTFILE_APPENDIX_PREFIX", CS%restartfile_appx_prefix, &
-                 "The prefix for the restart file appendix (i.e., ensemble id for ensemble runs). "// &
-                 "If this prefix is found in the restart file name, the appendix is added right after the "// &
-                 "first occurrence of the prefix. If not found, the appendix is added to the end of the "// &
-                 "file name. This parameter is ignored for non-ensemble runs.", &
+  call get_param(param_file, mdl, "ENSEMBLE_APPENDIX_PREFIX", CS%ensemble_appendix_prefix, &
+                 "If set to a non-empty string, this value specifies the substring after which "//&
+                 "the ensemble appendix is inserted in restart, initial conditions, and ocean "//&
+                 "geometry file names. If the specified substring is not found in any of those "//&
+                 "output file names, the model terminates with an error.", &
                  default="")
   call get_param(param_file, mdl, "RESTART_SYMMETRIC_CHECKSUMS", CS%symmetric_checksums, &
                  "If true, do the restart checksums on all the edge points for a non-reentrant "//&

src/initialization/MOM_shared_initialization.F90

@@ -16,6 +16,7 @@ module MOM_shared_initialization
 use MOM_io, only : MOM_read_data, MOM_read_vector, read_variable, stdout
 use MOM_io, only : open_file_to_read, close_file_to_read, SINGLE_FILE, MULTIPLE
 use MOM_io, only : slasher, vardesc, MOM_write_field, var_desc
+use MOM_io, only : insert_ensemble_appendix
 use MOM_string_functions, only : uppercase
 use MOM_unit_scaling, only : unit_scale_type
 
@@ -1342,6 +1343,7 @@ subroutine write_ocean_geometry_file(G, param_file, directory, US, geom_file)
   character(len=240) :: filepath  ! The full path to the file to write
   character(len=40)  :: mdl = "write_ocean_geometry_file"
   character(len=32)  :: filename_appendix = '' ! Appendix to geom filename for ensemble runs
+  character(len=32)  :: ensemble_appendix_prefix ! The prefix after which the ensemble id appendix is added
   type(vardesc),   dimension(:), allocatable :: &
     vars     ! Types with metadata about the variables and their staggering
   type(MOM_field), dimension(:), allocatable :: &
@@ -1407,16 +1409,13 @@ subroutine write_ocean_geometry_file(G, param_file, directory, US, geom_file)
     filepath = trim(directory) // "ocean_geometry"
   endif
 
-  ! Append ensemble run number to filename if it is an ensemble run
-  call get_filename_appendix(filename_appendix)
-  if (len_trim(filename_appendix) > 0) then
-    geom_file_len = len_trim(filepath)
-    if (filepath(geom_file_len-2:geom_file_len) == ".nc") then
-      filepath = filepath(1:geom_file_len-3) // '.' // trim(filename_appendix) // ".nc"
-    else
-      filepath = filepath // '.' // trim(filename_appendix)
-    endif
-  endif
+  call get_param(param_file, mdl, "ENSEMBLE_APPENDIX_PREFIX", ensemble_appendix_prefix, &
+                 "If set to a non-empty string, this value specifies the substring after which "//&
+                 "the ensemble appendix is inserted in restart, initial conditions, and ocean "//&
+                 "geometry file names. If the specified substring is not found in any of those "//&
+                 "output file names, the model terminates with an error.", &
+                 default="", do_not_log=.true.)
+  call insert_ensemble_appendix(filepath, ensemble_appendix_prefix)
 
   call get_param(param_file, mdl, "PARALLEL_RESTARTFILES", multiple_files, &
                  "If true, the IO layout is used to group processors that write to the same "//&