From c0b3555039c2be553b2315d502a765dde6dde967 Mon Sep 17 00:00:00 2001
From: Marshall Ward <marshall.ward@noaa.gov>
Date: Mon, 16 Dec 2024 00:48:22 -0500
Subject: [PATCH] Fix up mom_src_bh allocation and MEKE diagnostics

(Redo this commit message)
---
 src/parameterizations/lateral/MOM_MEKE.F90    | 53 +++++++++++++++----
 .../lateral/MOM_hor_visc.F90                  | 26 ++++++---
 2 files changed, 62 insertions(+), 17 deletions(-)

diff --git a/src/parameterizations/lateral/MOM_MEKE.F90 b/src/parameterizations/lateral/MOM_MEKE.F90
index 886ce720d5..2909d7ab2d 100644
--- a/src/parameterizations/lateral/MOM_MEKE.F90
+++ b/src/parameterizations/lateral/MOM_MEKE.F90
@@ -227,6 +227,7 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h
     Kh_v, &         ! The meridional diffusivity that is actually used [L2 T-1 ~> m2 s-1].
     baroHv, &       ! Depth integrated accumulated meridional mass flux [R Z L2 ~> kg].
     drag_vel_v      ! A piston velocity associated with bottom drag at v-points [H T-1 ~> m s-1 or kg m-2 s-1]
+  real :: bh_coeff  ! Biharmonic part of efficiency conversion in total MEKE [nondim]
   real :: Kh_here   ! The local horizontal viscosity [L2 T-1 ~> m2 s-1]
   real :: Inv_Kh_max ! The inverse of the local horizontal viscosity [T L-2 ~> s m-2]
   real :: K4_here   ! The local horizontal biharmonic viscosity [L4 T-1 ~> m4 s-1]
@@ -412,22 +413,51 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h
     !$OMP parallel do default(shared)
     do j=js,je ; do i=is,ie
       src(i,j) = CS%MEKE_BGsrc
-      src_adv(i,j) = 0.
-      src_mom_K4(i,j) = 0.
-      src_btm_drag(i,j) = 0.
-      src_GM(i,j) = 0.
-      src_mom_lp(i,j) = 0.
-      src_mom_bh(i,j) = 0.
     enddo ; enddo
 
+    ! Initialize diagnostics
+    ! TODO: Remove these after the damping is only computed if the diagnostic
+    !   is registered.
+    if (CS%id_src_adv) src_adv(is:ie, js:je) = 0.
+    if (CS%id_src_GM) src_GM(is:ie, js:je) = 0.
+    if (CS%id_src_mom_lp) src_mom_lp(is:ie, js:je) = 0.
+    if (CS%id_src_mom_bh) src_mom_bh(is:ie, js:je) = 0.
+    if (CS%id_src_mom_K4) src_mom_K4(is:ie, js:je) = 0.
+    if (CS%id_src_btm_drag) src_btm_drag(is:ie, js:je) = 0.
+
     if (allocated(MEKE%mom_src)) 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) &
-                   - (CS%MEKE_bhFrCoeff-CS%MEKE_FrCoeff)*I_mass(i,j)*MEKE%mom_src_bh(i,j)
-        src_mom_lp(i,j) = - CS%MEKE_FrCoeff*I_mass(i,j)*(MEKE%mom_src(i,j)-MEKE%mom_src_bh(i,j))
-        src_mom_bh(i,j) = - CS%MEKE_bhFrCoeff*I_mass(i,j)*MEKE%mom_src_bh(i,j)
+        src(i,j) = src(i,j) - CS%MEKE_FrCoeff * I_mass(i,j) * MEKE%mom_src(i,j)
+      enddo ; enddo
+    endif
+
+    if (allocated(MEKE%mom_src_bh)) then
+      if (CS%MEKE_bhFrCoeff > 0. .and. CS%MEKE_FrCoeff > 0.) then
+        bh_coeff = CS%MEKE_bhFrCoeff - CS%MEKE_FrCoeff
+      else
+        bh_coeff = CS%MEKE_bhFrCoeff
+      endif
+
+      !$OMP parallel do default(shared)
+      do j=js,je ; do i=is,ie
+        src(i,j) = src(i,j) - bh_coeff * I_mass(i,j) * MEKE%mom_src_bh(i,j)
       enddo ; enddo
+
+      if (CS%id_src_mom_lp) then
+        !$OMP parallel do default(shared)
+        do j=js,je ; do i=is,ie
+          src_mom_lp(i,j) = -CS%MEKE_FrCoeff * I_mass(i,j) &
+              * (MEKE%mom_src(i,j) - MEKE%mom_src_bh(i,j))
+        enddo ; enddo
+      endif
+
+      if (CS%id_src_mom_bh) then
+        !$OMP parallel do default(shared)
+        do j=js,je ; do i=is,ie
+          src_mom_bh(i,j) = -CS%MEKE_bhFrCoeff * I_mass(i,j) * MEKE%mom_src_bh(i,j)
+        enddo ; enddo
+      endif
     endif
 
     if (allocated(MEKE%GME_snk)) then
@@ -489,6 +519,9 @@ 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
diff --git a/src/parameterizations/lateral/MOM_hor_visc.F90 b/src/parameterizations/lateral/MOM_hor_visc.F90
index 92794c54e7..c2496e0d3a 100644
--- a/src/parameterizations/lateral/MOM_hor_visc.F90
+++ b/src/parameterizations/lateral/MOM_hor_visc.F90
@@ -2126,8 +2126,14 @@ subroutine horizontal_viscosity(u, v, h, uh, vh, diffu, diffv, MEKE, VarMix, G,
       if (k==1) then
         do j=js,je ; do i=is,ie
           MEKE%mom_src(i,j) = 0.
-          MEKE%mom_src_bh(i,j) = 0.
         enddo ; enddo
+
+        if (allocated(MEKE%mom_src_bh)) then
+          do j=js,je ; do i=is,ie
+            MEKE%mom_src_bh(i,j) = 0.
+          enddo ; enddo
+        endif
+
         if (allocated(MEKE%GME_snk)) then
           do j=js,je ; do i=is,ie
             MEKE%GME_snk(i,j) = 0.
@@ -2160,15 +2166,21 @@ subroutine horizontal_viscosity(u, v, h, uh, vh, diffu, diffv, MEKE, VarMix, G,
           endif
 
           MEKE%mom_src(i,j) = MEKE%mom_src(i,j) + (FrictWork(i,j,k) - RoScl*FrictWork_bh(i,j,k))
-          MEKE%mom_src_bh(i,j) = MEKE%mom_src_bh(i,j) + &
-                                 (FrictWork_bh(i,j,k) - RoScl*FrictWork_bh(i,j,k))
+
+          if (allocated(MEKE%mom_src_bh)) &
+            MEKE%mom_src_bh(i,j) = MEKE%mom_src_bh(i,j) &
+                + (FrictWork_bh(i,j,k) - RoScl * FrictWork_bh(i,j,k))
         enddo ; enddo
       else
+        do j=js,je ; do i=is,ie
+          MEKE%mom_src(i,j) = MEKE%mom_src(i,j) + FrictWork(i,j,k)
+        enddo ; enddo
 
-      do j=js,je ; do i=is,ie
-        MEKE%mom_src(i,j) = MEKE%mom_src(i,j) + FrictWork(i,j,k)
-        MEKE%mom_src_bh(i,j) = MEKE%mom_src_bh(i,j) + FrictWork_bh(i,j,k)
-      enddo ; enddo
+        if (allocated(MEKE%mom_src_bh)) then
+          do j=js,je ; do i=is,ie
+            MEKE%mom_src_bh(i,j) = MEKE%mom_src_bh(i,j) + FrictWork_bh(i,j,k)
+          enddo ; enddo
+        endif
       endif ! MEKE%backscatter_Ro_c
 
       if (CS%use_GME .and. allocated(MEKE%GME_snk)) then