Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dev access candidate 20241217 #3

Closed
wants to merge 5 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 149 additions & 4 deletions config_src/drivers/nuopc_cap/mom_cap.F90
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -96,11 +107,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
Expand Down Expand Up @@ -416,6 +430,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
Expand Down Expand Up @@ -447,6 +465,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=:), allocatable :: rpointer_filename
Expand Down Expand Up @@ -525,6 +544,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, &
Expand Down Expand Up @@ -656,13 +677,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)
Expand Down Expand Up @@ -729,6 +781,31 @@ subroutine InitializeAdvertise(gcomp, importState, exportState, clock, rc)
Ice_ocean_boundary%hcond = 0.0
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

call query_ocean_state(ocean_state, use_waves=use_waves, wave_method=wave_method)
if (use_waves) then
if (wave_method == "EFACTOR") then
Expand Down Expand Up @@ -797,6 +874,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")
Expand All @@ -808,6 +894,8 @@ subroutine InitializeAdvertise(gcomp, importState, exportState, clock, rc)
call fld_list_add(fldsFrOcn_num, fldsFrOcn, "Fioo_q" , "will provide")
call fld_list_add(fldsFrOcn_num, fldsFrOcn, "So_bldepth" , "will provide")

! 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
Expand Down Expand Up @@ -1619,11 +1707,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
Expand All @@ -1635,12 +1726,13 @@ 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(240) :: additional_restart_dir
integer :: num_rest_files
real(8) :: MPI_Wtime, timers
logical :: write_restart
Expand Down Expand Up @@ -1683,6 +1775,7 @@ subroutine ModelAdvance(gcomp, rc)

Time_step_coupled = esmf2fms_time(timeStep)
Time = esmf2fms_time(currTime)
Time_import = Time

!---------------
! Apply ocean lag for startup runs:
Expand Down Expand Up @@ -1758,8 +1851,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
Expand Down Expand Up @@ -1831,7 +1955,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
Expand Down Expand Up @@ -1889,6 +2013,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
Expand Down Expand Up @@ -2169,6 +2301,7 @@ subroutine ocean_model_finalize(gcomp, rc)
integer :: alarmCount
character(len=64) :: timestamp
logical :: write_restart
character(240) :: additional_restart_dir
character(len=*),parameter :: subname='(MOM_cap:ocean_model_finalize)'
real(8) :: MPI_Wtime, timefs

Expand Down Expand Up @@ -2202,6 +2335,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()

Expand Down
Loading
Loading