Allow generic tracers in NUOPC cap (#8)

These changes provide an initial implementation to allow use of coupled generic tracers with NUOPC-coupled MOM6. This is a bit of a hack as coupling in generic tracers effectively assumes the use of FMScoupler. Broadly, the changes to the NUOPC cap in this commit do the following:

- Initialisation phase: Initialise the FMS coupler type data structures used by generic tracers for handling surface fluxes. Register with NUOPC the additional import fields required by the generic tracer package being used. Export of additional fields is not currently handled.
- Advance phase: Populate the relevant FMS coupler type data structure with fields received from the atmosphere. Calculate the surface fluxes using a modified version of the routine used by FMScoupler.

Restart read/write for the FMS coupler types is also handled. Fields in the FMS coupler types can be overridden using the data_table

Author: dougiesquire

config_src/drivers/nuopc_cap/mom_cap.F90

@@ -2,8 +2,9 @@
 
 module MOM_cap_mod
 
+use field_manager_mod,        only: field_manager_init, field_manager_end
 use MOM_domains,              only: get_domain_extent
-use MOM_io,                   only: stdout, io_infra_end
+use MOM_io,                   only: stdout, io_infra_end, slasher
 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
@@ -30,13 +31,23 @@ module MOM_cap_mod
 use MOM_cap_methods,          only: ChkErr
 use MOM_ensemble_manager,     only: ensemble_manager_init
 use MOM_coms,                 only: sum_across_PEs
+use MOM_coupler_types,        only: coupler_1d_bc_type, coupler_2d_bc_type
 
 #ifdef CESMCOUPLED
 use shr_log_mod,             only: shr_log_setLogUnit
 use nuopc_shr_methods,       only: get_component_instance
 #endif
 use time_utils_mod,           only: esmf2fms_time
 
+#ifdef _USE_GENERIC_TRACER
+use MOM_coupler_types,        only: coupler_type_spawn, coupler_type_destructor
+use MOM_coupler_types,        only: coupler_type_set_diags, coupler_type_send_data, coupler_type_data_override
+use MOM_data_override,        only: data_override_init, data_override
+use MOM_cap_gtracer_flux,     only: gas_exchange_init, gas_fields_restore, gas_fields_restart
+use MOM_cap_gtracer_flux,     only: get_coupled_field_name, add_gas_fluxes_param, UNKNOWN_CMEPS_FIELD
+use MOM_cap_gtracer_flux,     only: atmos_ocean_fluxes_calc
+#endif
+
 use, intrinsic :: iso_fortran_env, only: output_unit
 
 use ESMF,  only: ESMF_ClockAdvance, ESMF_ClockGet, ESMF_ClockPrint, ESMF_VMget
@@ -100,11 +111,14 @@ module MOM_cap_mod
 public SetServices
 public SetVM
 
-!> Internal state type with pointers to three types defined by MOM.
+!> Internal state type with pointers to types defined by MOM.
 type ocean_internalstate_type
   type(ocean_public_type),       pointer :: ocean_public_type_ptr
   type(ocean_state_type),        pointer :: ocean_state_type_ptr
   type(ice_ocean_boundary_type), pointer :: ice_ocean_boundary_type_ptr
+#ifdef _USE_GENERIC_TRACER
+  type(coupler_2d_bc_type),      pointer :: coupler_2d_bc_type_ptr
+#endif
 end type
 
 !>  Wrapper-derived type required to associate an internal state instance
@@ -420,6 +434,10 @@ subroutine InitializeAdvertise(gcomp, importState, exportState, clock, rc)
   type (ocean_public_type),      pointer :: ocean_public => NULL()
   type (ocean_state_type),       pointer :: ocean_state => NULL()
   type(ice_ocean_boundary_type), pointer :: Ice_ocean_boundary => NULL()
+  type(coupler_1d_bc_type),      pointer :: gas_fields_atm => NULL()
+  type(coupler_1d_bc_type),      pointer :: gas_fields_ocn => NULL()
+  type(coupler_1d_bc_type),      pointer :: gas_fluxes => NULL()
+  type(coupler_2d_bc_type),      pointer :: atm_fields => NULL()
   type(ocean_internalstate_wrapper)      :: ocean_internalstate
   type(ocean_grid_type),         pointer :: ocean_grid => NULL()
   type(directories)                      :: dirs
@@ -452,6 +470,7 @@ subroutine InitializeAdvertise(gcomp, importState, exportState, clock, rc)
   character(len=512)                     :: restartfile          ! Path/Name of restart file
   character(len=2048)                    :: restartfiles         ! Path/Name of restart files
                                                                  ! (same as restartfile if single restart file)
+  character(240)                         :: additional_restart_dir
   character(len=*), parameter            :: subname='(MOM_cap:InitializeAdvertise)'
   character(len=32)                      :: calendar
   character(len=17)                      :: timestamp
@@ -548,6 +567,8 @@ subroutine InitializeAdvertise(gcomp, importState, exportState, clock, rc)
   if (chkerr(rc,__LINE__,u_FILE_u)) return
   call MOM_infra_init(mpi_comm_mom)
 
+  call field_manager_init
+
   ! determine the calendar
   if (cesm_coupled) then
     call NUOPC_CompAttributeGet(gcomp, name="calendar", value=cvalue, &
@@ -706,13 +727,44 @@ subroutine InitializeAdvertise(gcomp, importState, exportState, clock, rc)
 
   endif
 
+  ! Set NUOPC attribute additional_restart_dir to RESTART/ if not defined
+  additional_restart_dir = "RESTART/"
+  call NUOPC_CompAttributeGet(gcomp, name="additional_restart_dir", value=cvalue, &
+        isPresent=isPresent, isSet=isSet, rc=rc)
+  if (ChkErr(rc,__LINE__,u_FILE_u)) return
+  if (isPresent .and. isSet) then
+    additional_restart_dir = slasher(cvalue)
+  else
+    call ESMF_LogWrite('MOM_cap:additional_restart_dir unset. Defaulting to '//trim(additional_restart_dir), &
+          ESMF_LOGMSG_INFO)
+  endif
+  call NUOPC_CompAttributeSet(gcomp, name="additional_restart_dir", value=additional_restart_dir, rc=rc)
+  if (chkerr(rc,__LINE__,u_FILE_u)) return
+
   ocean_public%is_ocean_pe = .true.
+#ifdef _USE_GENERIC_TRACER
+  ! Initialise structures for extra tracer fluxes
+  call gas_exchange_init(gas_fields_atm=gas_fields_atm, gas_fields_ocn=gas_fields_ocn, gas_fluxes=gas_fluxes)
+
+  if (cesm_coupled .and. len_trim(inst_suffix)>0) then
+    call ocean_model_init(ocean_public, ocean_state, time0, time_start, gas_fields_ocn=gas_fields_ocn, &
+          input_restart_file=trim(adjustl(restartfiles)), inst_index=inst_index)
+  else
+    call ocean_model_init(ocean_public, ocean_state, time0, time_start, gas_fields_ocn=gas_fields_ocn, &
+          input_restart_file=trim(adjustl(restartfiles)))
+  endif
+
+  ! Enable data override via the data_table using the component name 'OCN'
+  call get_ocean_grid(ocean_state, ocean_grid)
+  call data_override_init(ocean_grid%Domain)
+#else
   if (cesm_coupled .and. len_trim(inst_suffix)>0) then
     call ocean_model_init(ocean_public, ocean_state, time0, time_start, &
       input_restart_file=trim(adjustl(restartfiles)), inst_index=inst_index)
   else
     call ocean_model_init(ocean_public, ocean_state, time0, time_start, input_restart_file=trim(adjustl(restartfiles)))
   endif
+#endif
 
   ! GMM, this call is not needed in CESM. Check with EMC if it can be deleted.
   call ocean_model_flux_init(ocean_state)
@@ -780,6 +832,31 @@ subroutine InitializeAdvertise(gcomp, importState, exportState, clock, rc)
     endif
   endif
 
+#ifdef _USE_GENERIC_TRACER
+  ! Allocate fields for extra tracer fluxes in Ice_ocean_boundary
+  ! Annoyingly, spawning doesn't copy param array, so add manually
+  call coupler_type_spawn(gas_fluxes, Ice_ocean_boundary%fluxes, (/isc,isc,iec,iec/), &
+        (/jsc,jsc,jec,jec/), suffix='_ice_ocn')
+  call add_gas_fluxes_param(Ice_ocean_boundary%fluxes)
+
+  ! Initialise structure for atmos fields related to extra tracer fluxes
+  ! This is set in the ESMF Internal State to be accessed elsewhere
+  ! TODO: should we deallocate atm_fields in a finalise step? Ice_ocean_boundary is handled
+  ! in a similar way and does not appear to be deallocated.
+  allocate(atm_fields)
+  ocean_internalstate%ptr%coupler_2d_bc_type_ptr => atm_fields
+  call coupler_type_spawn(gas_fields_atm, atm_fields, (/isc,isc,iec,iec/), &
+        (/jsc,jsc,jec,jec/), suffix='_atm')
+
+  ! Register diagnosics for extra tracer flux structures
+  call coupler_type_set_diags(Ice_ocean_boundary%fluxes, "ocean_flux", ocean_public%axes(1:2), time_start)
+  call coupler_type_set_diags(atm_fields, "atmos_sfc", ocean_public%axes(1:2), time_start)
+
+  ! Restore ocean fields related to extra tracer fluxes from restart files
+  call get_MOM_input(dirs=dirs)
+  call gas_fields_restore(ocean_public%fields, ocean_public%domain, dirs%restart_input_dir)
+#endif
+
   if (use_waves) then
     if (wave_method == "EFACTOR") then
       allocate( Ice_ocean_boundary%lamult(isc:iec,jsc:jec), source=0.0)
@@ -881,6 +958,15 @@ subroutine InitializeAdvertise(gcomp, importState, exportState, clock, rc)
     endif
   endif
 
+#ifdef _USE_GENERIC_TRACER
+  ! Add import fields required for extra tracer fluxes
+  do n = 1, gas_fluxes%num_bcs
+    stdname = get_coupled_field_name(gas_fluxes%bc(n)%name)
+    if (stdname /= UNKNOWN_CMEPS_FIELD) &
+      call fld_list_add(fldsToOcn_num, fldsToOcn, stdname, "will provide")
+  enddo
+#endif
+
   !--------- export fields -------------
   call fld_list_add(fldsFrOcn_num, fldsFrOcn, "So_omask"   , "will provide")
   call fld_list_add(fldsFrOcn_num, fldsFrOcn, "So_t"       , "will provide")
@@ -895,6 +981,8 @@ subroutine InitializeAdvertise(gcomp, importState, exportState, clock, rc)
     call fld_list_add(fldsFrOcn_num, fldsFrOcn, "Faoo_fco2_ocn", "will provide")
   endif
 
+  ! TODO: dts: How to handle export fields from generic tracers?
+
   do n = 1,fldsToOcn_num
     call NUOPC_Advertise(importState, standardName=fldsToOcn(n)%stdname, name=fldsToOcn(n)%shortname, rc=rc)
     if (ChkErr(rc,__LINE__,u_FILE_u)) return
@@ -1706,11 +1794,14 @@ subroutine ModelAdvance(gcomp, rc)
   type (ocean_public_type),      pointer :: ocean_public       => NULL()
   type (ocean_state_type),       pointer :: ocean_state        => NULL()
   type(ice_ocean_boundary_type), pointer :: Ice_ocean_boundary => NULL()
+  type(coupler_2d_bc_type),      pointer :: atm_fields         => NULL()
   type(ocean_internalstate_wrapper)      :: ocean_internalstate
   type(ocean_grid_type)        , pointer :: ocean_grid
   type(time_type)                        :: Time
+  type(time_type)                        :: Time_import
   type(time_type)                        :: Time_step_coupled
   type(time_type)                        :: Time_restart_current
+  integer                                :: isc,iec,jsc,jec
   integer                                :: dth, dtm, dts
   integer                                :: nc
   type(ESMF_Time)                        :: MyTime
@@ -1722,13 +1813,14 @@ subroutine ModelAdvance(gcomp, rc)
   integer                                :: writeunit
   integer                                :: localPet
   type(ESMF_VM)                          :: vm
-  integer                                :: n, i
+  integer                                :: m, n, i
   character(240)                         :: import_timestr, export_timestr
   character(len=128)                     :: fldname
   character(len=*),parameter             :: subname='(MOM_cap:ModelAdvance)'
   character(len=8)                       :: suffix
   character(len=:), allocatable          :: rpointer_filename
   character(len=17)                      :: timestamp
+  character(240)                         :: additional_restart_dir
   integer                                :: num_rest_files
   real(8)                                :: MPI_Wtime, timers
   logical                                :: write_restart, write_restartfh
@@ -1769,6 +1861,7 @@ subroutine ModelAdvance(gcomp, rc)
 
   Time_step_coupled = esmf2fms_time(timeStep)
   Time = esmf2fms_time(currTime)
+  Time_import = Time
 
   !---------------
   ! Apply ocean lag for startup runs:
@@ -1844,8 +1937,39 @@ subroutine ModelAdvance(gcomp, rc)
     ! Import data
     !---------------
 
+#ifdef _USE_GENERIC_TRACER
+    atm_fields => ocean_internalstate%ptr%coupler_2d_bc_type_ptr
+
+    call mom_import(ocean_public, ocean_grid, importState, ice_ocean_boundary, atm_fields=atm_fields, rc=rc)
+    if (ChkErr(rc,__LINE__,u_FILE_u)) return
+
+    ! Potentially override atm_fields from data_table.
+    call coupler_type_data_override('OCN', atm_fields, Time_import)
+
+    ! Potentially override ice_ocean_boundary%fluxes from data_table.
+    ! Doing this before atmos_ocean_fluxes_calc call avoids unnecessary calculation of overridden fluxes.
+    ! However, we cannot use coupler_type_data_override here since it does not set the override flag on
+    ! overridden fields
+    do n = 1, ice_ocean_boundary%fluxes%num_bcs
+      do m = 1, ice_ocean_boundary%fluxes%bc(n)%num_fields
+        call data_override('OCN', ice_ocean_boundary%fluxes%bc(n)%field(m)%name, &
+              ice_ocean_boundary%fluxes%bc(n)%field(m)%values, Time_import, &
+              override=ice_ocean_boundary%fluxes%bc(n)%field(m)%override)
+      enddo
+    enddo
+
+    ! Calculate the extra tracer fluxes
+    call get_domain_extent(ocean_public%domain, isc, iec, jsc, jec)
+    call atmos_ocean_fluxes_calc(atm_fields, ocean_public%fields, ice_ocean_boundary%fluxes, &
+          ice_ocean_boundary%ice_fraction, isc, iec, jsc, jec)
+
+    ! Send diagnostics
+    call coupler_type_send_data(atm_fields, Time_import)
+    call coupler_type_send_data(ice_ocean_boundary%fluxes, Time_import)
+#else
     call mom_import(ocean_public, ocean_grid, importState, ice_ocean_boundary, rc=rc)
     if (ChkErr(rc,__LINE__,u_FILE_u)) return
+#endif
 
     !---------------
     ! Update MOM6
@@ -1912,7 +2036,7 @@ subroutine ModelAdvance(gcomp, rc)
       ! determine restart filename
       call ESMF_ClockGetNextTime(clock, MyTime, rc=rc)
       if (ChkErr(rc,__LINE__,u_FILE_u)) return
-      call ESMF_TimeGet (MyTime, yy=year, mm=month, dd=day, h=hour, m=minute, s=seconds, rc=rc )
+      call ESMF_TimeGet (MyTime, yy=year, mm=month, dd=day, s=seconds, rc=rc )
       if (ChkErr(rc,__LINE__,u_FILE_u)) return
 
       if (cesm_coupled) then
@@ -1975,6 +2099,14 @@ subroutine ModelAdvance(gcomp, rc)
 
       endif
 
+#ifdef _USE_GENERIC_TRACER
+      ! Write fields for extra tracer fluxes to their internally defined ocean restart file
+      call NUOPC_CompAttributeGet(gcomp, name="additional_restart_dir", value=additional_restart_dir, rc=rc)
+      if (ChkErr(rc,__LINE__,u_FILE_u)) return
+
+      call gas_fields_restart(ocean_public%fields, ocean_public%domain, additional_restart_dir)
+#endif
+
       if (is_root_pe()) then
         write(stdout,*) subname//' writing restart file ',trim(restartname)
       endif
@@ -2213,6 +2345,7 @@ subroutine ocean_model_finalize(gcomp, rc)
   type(ESMF_Alarm), allocatable          :: alarmList(:)
   integer                                :: alarmCount
   logical                                :: write_restart
+  character(240)                         :: additional_restart_dir
   character(len=*),parameter  :: subname='(MOM_cap:ocean_model_finalize)'
   real(8)                                :: MPI_Wtime, timefs
 
@@ -2246,6 +2379,18 @@ subroutine ocean_model_finalize(gcomp, rc)
 
   call ocean_model_end(ocean_public, ocean_State, Time, write_restart=write_restart)
 
+#ifdef _USE_GENERIC_TRACER
+  if (write_restart) then
+    ! Write fields for extra tracer fluxes to their internally defined ocean restart file
+    call NUOPC_CompAttributeGet(gcomp, name="additional_restart_dir", value=additional_restart_dir, rc=rc)
+    if (ChkErr(rc,__LINE__,u_FILE_u)) return
+
+    call gas_fields_restart(ocean_public%fields, ocean_public%domain, additional_restart_dir)
+  endif
+#endif
+
+  call field_manager_end()
+
   call io_infra_end()
   call MOM_infra_end()