Skip to content

Commit

Permalink
MEKE: Move damping diagnostics outside MEKE loop
Browse files Browse the repository at this point in the history 
The damping MEKE loop also included updates to multiple diagnostics,
even if they were not registered.  This would presumably have a negative
impact on performance.

This patch moves each diagnostic into a separate loop.  It also
conditionally precomputes the damping and damp_rate parameters, which
are now stored as 2d arrays rather than in-loop scalars.

As before, the MEKE calculation is left unchanged in order to preserve
bit reproducibility.
  • Loading branch information
@marshallward
marshallward committed Jan 3, 2025
1 parent 0d2d668 commit 345a6fb
Showing 1 changed file with 179 additions and 56 deletions.
235 changes: 179 additions & 56 deletions src/parameterizations/lateral/MOM_MEKE.F90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -212,8 +212,10 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h
barotrFac2, & ! Ratio of EKE_barotropic / EKE [nondim]
bottomFac2, & ! Ratio of EKE_bottom / EKE [nondim]
tmp, & ! Temporary variable for computation of diagnostic velocities [L T-1 ~> m s-1]
equilibrium_value ! The equilibrium value of MEKE to be calculated at each
! time step [L2 T-2 ~> m2 s-2]
equilibrium_value, & ! The equilibrium value of MEKE to be calculated at
! each time step [L2 T-2 ~> m2 s-2]
damp_rate, & ! The MEKE damping rate [T-1 ~> s-1]
damping ! The net damping of a field after sdt_damp [nondim]

real, dimension(SZIB_(G),SZJ_(G)) :: &
MEKE_uflux, & ! The zonal advective and diffusive flux of MEKE with units of [R Z L4 T-3 ~> kg m2 s-3].
Expand All @@ -238,9 +240,9 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h
real :: sdt ! dt to use locally [T ~> s] (could be scaled to accelerate)
real :: sdt_damp ! dt for damping [T ~> s] (sdt could be split).
real :: damp_step ! Size of damping timestep relative to sdt [nondim]
real :: damp_rate ! The MEKE damping rate [T-1 ~> s-1].
real :: damping ! The net damping of a field after sdt_damp [nondim]
logical :: use_drag_rate ! Flag to indicate drag_rate is finite
logical :: any_damping_diags_s1 ! True if any damped diagnostics are enabled in first stage
logical :: any_damping_diags ! True if any damped diagnostics are enabled
integer :: i, j, k, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz
real(kind=real32), dimension(size(MEKE%MEKE),NUM_FEATURES) :: features_array ! The array of features
! needed for the machine learning inference, with different
Expand Down Expand Up @@ -416,16 +418,29 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h
enddo ; enddo

! Initialize diagnostics
! TODO: Remove these after the damping is only computed if the diagnostic
! is registered.
if (CS%id_src_adv > 0) src_adv(is:ie, js:je) = 0.
if (CS%id_src_GM > 0) src_GM(is:ie, js:je) = 0.
if (CS%id_src_mom_lp > 0) src_mom_lp(is:ie, js:je) = 0.
if (CS%id_src_mom_bh > 0) src_mom_bh(is:ie, js:je) = 0.
if (CS%id_src_mom_K4 > 0) src_mom_K4(is:ie, js:je) = 0.
if (CS%id_src_btm_drag > 0) src_btm_drag(is:ie, js:je) = 0.

if (allocated(MEKE%mom_src)) then
! Identify any damped diagnostics in first stage of Strang splitting
any_damping_diags_s1 = any([ &
CS%id_src_GM > 0, &
CS%id_src_mom_lp > 0, &
CS%id_src_mom_bh > 0, &
CS%id_src_btm_drag > 0 &
])

! Identify any damped diagnostics
any_damping_diags = any([ &
any_damping_diags_s1, &
CS%id_src_adv > 0, &
CS%id_src_mom_K4 > 0 &
])

if (CS%MEKE_FrCoeff > 0.) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src(i,j) = src(i,j) - CS%MEKE_FrCoeff * I_mass(i,j) * MEKE%mom_src(i,j)
Expand Down Expand Up @@ -519,36 +534,75 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h
enddo ; enddo
endif

! TODO: All of these diagnostics needs to be pulled out of the MEKE damping
! loop and handled separately, and only if they are registered.

! First stage of Strang splitting

!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
damp_rate = CS%MEKE_damping + drag_rate(i,j) * bottomFac2(i,j)
if (MEKE%MEKE(i,j) < 0.) damp_rate = 0.
damp_rate(i,j) = CS%MEKE_damping + drag_rate(i,j) * bottomFac2(i,j)

if (MEKE%MEKE(i,j) < 0.) damp_rate(i,j) = 0.
! notice that the above line ensures a damping only if MEKE is positive,
! while leaving MEKE unchanged if it is negative
enddo ; enddo

damping = 1. / (1. + sdt_damp * damp_rate)
! NOTE: MEKE%MEKE cannot use `damping` since we must preserve the existing
! bit-reproducible solution.
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
MEKE%MEKE(i,j) = MEKE%MEKE(i,j) / (1. + sdt_damp * damp_rate(i,j))
enddo ; enddo

! NOTE: MEKE%MEKE should use `damping` but we must preserve the existing
! expression for bit reproducibility
MEKE%MEKE(i,j) = MEKE%MEKE(i,j) / (1. + sdt_damp * damp_rate)
MEKE_decay(i,j) = damp_rate * G%mask2dT(i,j)
if (any_damping_diags_s1) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
damping(i,j) = 1. / (1. + sdt_damp * damp_rate(i,j))
enddo ; enddo

src_GM(i,j) = src_GM(i,j) * damping
src_mom_lp(i,j) = src_mom_lp(i,j) * damping
src_mom_bh(i,j) = src_mom_bh(i,j) * damping
if (CS%id_decay > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
MEKE_decay(i,j) = damp_rate(i,j) * G%mask2dT(i,j)
enddo ; enddo
endif

src_btm_drag(i,j) = - MEKE_current(i,j) * ( &
damp_step * (damp_rate * damping) &
)
if (CS%id_src_GM > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_GM(i,j) = src_GM(i,j) * damping(i,j)
enddo ; enddo
endif

! Store the effective damping rate if sdt is split
if (CS%MEKE_KH >= 0. .or. CS%MEKE_K4 >= 0.) &
damp_rate_s1(i,j) = damp_rate * damping
enddo ; enddo
if (CS%id_src_mom_lp > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_mom_lp(i,j) = src_mom_lp(i,j) * damping(i,j)
enddo ; enddo
endif

if (CS%id_src_mom_bh > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_mom_bh(i,j) = src_mom_bh(i,j) * damping(i,j)
enddo ; enddo
endif

if (CS%id_src_btm_drag > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_btm_drag(i,j) = -MEKE_current(i,j) * ( &
damp_step * (damp_rate(i,j) * damping(i,j)) &
)
enddo ; enddo

! Store the effective damping rate if sdt is split
if (CS%MEKE_KH >= 0. .or. CS%MEKE_K4 >= 0.) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
damp_rate_s1(i,j) = damp_rate(i,j) * damping(i,j)
enddo ; enddo
endif
endif
endif

if (CS%kh_flux_enabled .or. CS%MEKE_K4 >= 0.0) then
! Update MEKE in the halos for lateral or bi-harmonic diffusion
Expand Down Expand Up @@ -687,10 +741,16 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h
MEKE%MEKE(i,j) = MEKE%MEKE(i,j) + (sdt*(G%IareaT(i,j)*I_mass(i,j))) * &
((MEKE_uflux(I-1,j) - MEKE_uflux(I,j)) + &
(MEKE_vflux(i,J-1) - MEKE_vflux(i,J)))
src_adv(i,j) = (G%IareaT(i,j)*I_mass(i,j)) * &
((MEKE_uflux(I-1,j) - MEKE_uflux(I,j)) + &
(MEKE_vflux(i,J-1) - MEKE_vflux(i,J)))
enddo ; enddo

if (CS%id_src_adv > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_adv(i,j) = (G%IareaT(i,j)*I_mass(i,j)) * &
((MEKE_uflux(I-1,j) - MEKE_uflux(I,j)) + &
(MEKE_vflux(i,J-1) - MEKE_vflux(i,J)))
enddo ; enddo
endif
endif ! MEKE_KH>0

! Add on bi-harmonic tendency
Expand All @@ -711,30 +771,80 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h
cdrag2 * ( max(0.0, 2.0*bottomFac2(i,j)*MEKE%MEKE(i,j)) + CS%MEKE_Uscale**2 ) )
enddo ; enddo
endif

!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
damp_rate = CS%MEKE_damping + drag_rate(i,j) * bottomFac2(i,j)
if (MEKE%MEKE(i,j) < 0.) damp_rate = 0.
damp_rate(i,j) = CS%MEKE_damping + drag_rate(i,j) * bottomFac2(i,j)

if (MEKE%MEKE(i,j) < 0.) damp_rate(i,j) = 0.
! notice that the above line ensures a damping only if MEKE is positive,
! while leaving MEKE unchanged if it is negative
enddo ; enddo

damping = 1. / (1. + sdt_damp * damp_rate)
! NOTE: MEKE%MEKE cannot use `damping` since we must preserve the
! existing bit-reproducible solution.
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
MEKE%MEKE(i,j) = MEKE%MEKE(i,j) / (1. + sdt_damp * damp_rate(i,j))
enddo ; enddo

! NOTE: As above, MEKE%MEKE should use `damping` but we must preserve
! the existing expression for bit reproducibility.
MEKE%MEKE(i,j) = MEKE%MEKE(i,j) / (1.0 + sdt_damp*damp_rate)
MEKE_decay(i,j) = damp_rate*G%mask2dT(i,j)
if (any_damping_diags) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
damping(i,j) = 1. / (1. + sdt_damp * damp_rate(i,j))
enddo ; enddo

src_GM(i,j) = src_GM(i,j) * damping
src_mom_lp(i,j) = src_mom_lp(i,j) * damping
src_mom_bh(i,j) = src_mom_bh(i,j) * damping
src_adv(i,j) = src_adv(i,j) * damping
src_mom_K4(i,j) = src_mom_K4(i,j) * damping
if (CS%id_decay > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
MEKE_decay(i,j) = damp_rate(i,j) * G%mask2dT(i,j)
enddo ; enddo
endif

src_btm_drag(i,j) = -MEKE_current(i,j) * ( &
damp_step * damping * (damp_rate + damp_rate_s1(i,j)) &
)
enddo ; enddo
if (CS%id_src_GM > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_GM(i,j) = src_GM(i,j) * damping(i,j)
enddo ; enddo
endif

if (CS%id_src_mom_lp > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_mom_lp(i,j) = src_mom_lp(i,j) * damping(i,j)
enddo ; enddo
endif

if (CS%id_src_mom_bh > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_mom_bh(i,j) = src_mom_bh(i,j) * damping(i,j)
enddo ; enddo
endif

if (CS%id_src_adv > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_adv(i,j) = src_adv(i,j) * damping(i,j)
enddo ; enddo
endif

if (CS%id_src_mom_K4 > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_mom_K4(i,j) = src_mom_K4(i,j) * damping(i,j)
enddo ; enddo
endif

if (CS%id_src_btm_drag > 0) then
!$OMP parallel do default(shared)
do j=js,je ; do i=is,ie
src_btm_drag(i,j) = -MEKE_current(i,j) * (damp_step &
* ((damp_rate(i,j) + damp_rate_s1(i,j)) * damping(i,j)) &
)
enddo ; enddo
endif
endif
endif ! MEKE_KH>=0

if (CS%debug) then
Expand Down Expand Up @@ -1522,20 +1632,33 @@ logical function MEKE_init(Time, G, GV, US, param_file, diag, dbcomms_CS, CS, ME
if (.not. allocated(MEKE%MEKE)) CS%id_Ut = -1
CS%id_src = register_diag_field('ocean_model', 'MEKE_src', diag%axesT1, Time, &
'MEKE energy source', 'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)
!add diagnostics for the terms in the MEKE budget

CS%id_src_adv = register_diag_field('ocean_model', 'MEKE_src_adv', diag%axesT1, Time, &
'MEKE energy source from the horizontal advection of MEKE', 'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)
CS%id_src_mom_K4 = register_diag_field('ocean_model', 'MEKE_src_mom_K4', diag%axesT1, Time, &
'MEKE energy source from the biharmonic of MEKE', 'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)

CS%id_src_btm_drag = register_diag_field('ocean_model', 'MEKE_src_btm_drag', diag%axesT1, Time, &
'MEKE energy source from the bottom drag acting on MEKE', 'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)
CS%id_src_GM = register_diag_field('ocean_model', 'MEKE_src_GM', diag%axesT1, Time, &
'MEKE energy source from the thickness mixing (GM scheme)', 'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)
CS%id_src_mom_lp = register_diag_field('ocean_model', 'MEKE_src_mom_lp', diag%axesT1, Time, &
'MEKE energy source from the Laplacian of resolved flows', 'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)
CS%id_src_mom_bh = register_diag_field('ocean_model', 'MEKE_src_mom_bh', diag%axesT1, Time, &
'MEKE energy source from the biharmonic of resolved flows', 'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)
!end

if (CS%MEKE_K4 >= 0.) &
CS%id_src_mom_K4 = register_diag_field('ocean_model', 'MEKE_src_mom_K4', &
diag%axesT1, Time, 'MEKE energy source from the biharmonic of MEKE', &
'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)

if (CS%MEKE_GMcoeff >= 0.) &
CS%id_src_GM = register_diag_field('ocean_model', 'MEKE_src_GM', &
diag%axesT1, Time, 'MEKE energy source from the thickness mixing (GM scheme)', &
'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)

if (CS%MEKE_FrCoeff >= 0.) &
CS%id_src_mom_lp = register_diag_field('ocean_model', 'MEKE_src_mom_lp', &
diag%axesT1, Time, 'MEKE energy source from the Laplacian of resolved flows', &
'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)

if (CS%MEKE_bhFrCoeff >= 0.) &
CS%id_src_mom_bh = register_diag_field('ocean_model', 'MEKE_src_mom_bh', &
diag%axesT1, Time, 'MEKE energy source from the biharmonic of resolved flows', &
'm2 s-3', conversion=(US%L_T_to_m_s**2)*US%s_to_T)

CS%id_decay = register_diag_field('ocean_model', 'MEKE_decay', diag%axesT1, Time, &
'MEKE decay rate', 's-1', conversion=US%s_to_T)
CS%id_GM_src = register_diag_field('ocean_model', 'MEKE_GM_src', diag%axesT1, Time, &
Expand Down

0 comments on commit 345a6fb

Please sign in to comment.