Machine Learning based vertical diffusivity in EPBL mixing scheme used for ocean surface boundary layer

src/parameterizations/vertical/MOM_energetic_PBL.F90

@@ -155,6 +155,17 @@ module MOM_energetic_PBL
                              !! the Ekman depth over the Obukhov depth with destabilizing forcing [nondim].
   real :: Max_Enhance_M = 5. !< The maximum allowed LT enhancement to the mixing [nondim].
 
+  !/ Machine learned equation discovery model paramters ! eqdisc
+  logical :: eqdisc, eqdisc_v0  ! Machine Learned Equation discovery - shape function and velocity-scale
+  real :: v0 ! <v0 velocity scale from machine learning (GLSscheme) used for diffusivity [Z T-1 ~> m s-1]
+  real :: v0_lower_cap ! Lower cap to prevent v0 from attaining anomlously low values [Z T-1 ~> m s-1]
+  real :: f_lower ! Lower cap of |f| i.e. absolute of Coriolis parameter.
+                  ! Used in v0 subroutines. Default at 0.1deg Lat
+  real :: bflux_lower_cap, bflux_upper_cap ! Lower and upper cap for capping blfux while setting v0. 
+  real, allocatable, dimension(:) :: shape_function ! shape function used in machine learned diffusivity [nondim]
+  !/ Coefficients used in Machine learned diffusivity, Equations 6,7,10,11 in Sane et al. 2024
+  real :: c1, c2, c3, c4, c5, c6, c7, c8, c9  ! Non-dimensional constants
+  real :: c10, c11, c12, c13, c14, c15, c16   ! used in getting v0 and shape function [nondim]
   !/ Others
   type(time_type), pointer :: Time=>NULL() !< A pointer to the ocean model's clock.
 
@@ -202,6 +213,7 @@ module MOM_energetic_PBL
   integer :: id_TKE_mech_decay = -1, id_TKE_conv_decay = -1
   integer :: id_Mixing_Length = -1, id_Velocity_Scale = -1
   integer :: id_MSTAR_mix = -1, id_LA_mod = -1, id_LA = -1, id_MSTAR_LT = -1
+  integer :: id_v0_diag = -1 ! used for diagnostic purposes, part of ML-eqdisc 
   !>@}
 end type energetic_PBL_CS
 
@@ -832,7 +844,10 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing,
   integer, dimension(SZK_(GV)) :: num_itts
 
   integer :: k, nz, itt, max_itt
-
+
+  ! variables for ML based diffusivity
+  real :: v0_dummy ! Variable which get recycled, set equal to v_0
+
   nz = GV%ke
 
   debug = .false.  ! Change this hard-coded value for debugging.
@@ -978,21 +993,15 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing,
           MixLen_shape(K) = 1.0
         enddo ; endif
       else
-        ! Reduce the mixing length based on MLD, with a quadratic
-        ! expression that follows KPP.
-        I_MLD = 1.0 / MLD_guess
-        dz_rsum = 0.0
-        MixLen_shape(1) = 1.0
-        do K=2,nz+1
-          dz_rsum = dz_rsum + dz(k-1)
-          if (CS%MixLenExponent==2.0) then
-            MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * &
-                 (max(0.0, (MLD_guess - dz_rsum)*I_MLD) )**2 ! CS%MixLenExponent
+
+      call getshapefunction(CS,GV,h,absf,B_flux,u_star,MLD_guess,MixLen_shape) ! used for ML-eqdisc
+          v0_dummy = 0.0 ! a variable that gets passed on to subroutine get_eqdisc_v0
+          CS%v0 = 0.0
+          if (CS%eqdisc_v0 .eqv. .true.) then
+                  call get_eqdisc_v0(CS,absf,B_flux,u_star,MLD_guess,v0_dummy)
+                  CS%v0 = v0_dummy
           else
-            MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * &
-                 (max(0.0, (MLD_guess - dz_rsum)*I_MLD) )**CS%MixLenExponent
           endif
-        enddo
       endif
 
       Kd(1) = 0.0 ; Kddt_h(1) = 0.0
@@ -1212,6 +1221,8 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing,
             if (.not.CS%Use_MLD_iteration) then
               Kd_guess0 = (h_dz_int(K)*vstar) * CS%vonKar * ((dz_tt*hbs_here)*vstar) / &
                 ((CS%Ekman_scale_coef * absf) * (dz_tt*hbs_here) + vstar)
+            elseif (CS%eqdisc .eqv. .true.)  then  ! ML-eqdisc line1/2 
+              Kd_guess0 = MixLen_shape(K) * CS%v0 * MLD_guess ! ML-eqdisc
             else
               Kd_guess0 = (h_dz_int(K)*vstar) * CS%vonKar * mixlen(K)
             endif
@@ -1271,6 +1282,10 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing,
                 !  instead of redoing the computation will change answers...
                   Kd(K) = (h_dz_int(K)*vstar) * CS%vonKar *  ((dz_tt*hbs_here)*vstar) / &
                         ((CS%Ekman_scale_coef * absf) * (dz_tt*hbs_here) + vstar)
+
+                  elseif (CS%eqdisc .eqv. .true.)  then  ! ML-eqdisc line2/2 
+                  Kd(K) = MixLen_shape(K) * CS%v0 * MLD_guess ! ML-eqdisc
+
                 else
                   Kd(K) = (h_dz_int(K)*vstar) * CS%vonKar * mixlen(K)
                 endif
@@ -1561,6 +1576,205 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing,
 
 end subroutine ePBL_column
 
+subroutine getshapefunction(CS, GV, h, absf, B_flux, u_star, MLD_guess, MixLen_shape)
+! gives you shape function
+  type(verticalGrid_type), intent(in)    :: GV     !< The ocean's vertical grid structure.
+  type(energetic_PBL_CS),  intent(in) :: CS     !< Energetic PBL control struct
+  ! integer, intent(in) :: szkgv
+
+  real, dimension(SZK_(GV)), intent(in)  :: h      !< Layer thicknesses [H ~> m or kg m-2].
+  ! Local variables
+  real, dimension(SZK_(GV)+1) :: MixLen_shape !< A nondimensional shape factor for the mixing length that
+                                              !! gives it an appropriate asymptotic value at the bottom of
+                                              !! the boundary layer [nondim].
+  real, intent(in) :: absf      !< The absolute value of f [T-1 ~> s-1].
+  real, intent(in) :: u_star    !< The surface friction velocity [Z T-1 ~> m s-1].
+  real, intent(in) :: B_Flux    !< The surface buoyancy flux [Z2 T-3 ~> m2 s-3]
+  real, dimension(SZK_(GV)+1) :: hz !< depth variable, only used in this routine
+  real, intent(in) :: MLD_guess !< Mixing Layer depth guessed/found for iteration [Z ~> m].
+
+  ! Local variables for this subroutine
+  real :: I_MLD     ! The inverse of the current value of MLD [Z-1 ~> m-1].
+  real :: h_rsum    ! The running sum of h from the top [Z ~> m].
+  integer :: K  ! integer for loopping 
+  integer :: nz ! nz = GV%ke
+
+  nz = GV%ke
+  I_MLD = 1.0 / MLD_guess
+  h_rsum = 0.0
+  MixLen_shape(1) = 1.0
+
+  if (CS%eqdisc) then ! update Kd as per Machine Learning equation discovery
+    call kappa_eqdisc(MixLen_shape, CS, GV, h, absf, B_flux, u_star, MLD_guess)
+  else
+  do K=2,nz+1
+    h_rsum = h_rsum + h(k-1)*GV%H_to_Z
+      if (CS%MixLenExponent==2.0) then
+        MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * &
+          (max(0.0, (MLD_guess - h_rsum)*I_MLD) )**2 ! CS%MixLenExponent
+      else
+        MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * &
+          (max(0.0, (MLD_guess - h_rsum)*I_MLD) )**CS%MixLenExponent
+      endif
+  enddo
+  endif
+end subroutine getshapefunction
+
+subroutine kappa_eqdisc(shape_func, CS, GV, h, absf, B_flux, u_star, MLD_guess)
+! gives shape function from Sane et al. 2024.
+  type(verticalGrid_type), intent(in)    :: GV     !< The ocean's vertical grid structure.
+  type(energetic_PBL_CS),  intent(in) :: CS     !< Energetic PBL control struct
+  real, dimension(SZK_(GV)+1), intent(inout) :: shape_func  !< shape function
+  real, intent(in) :: absf      !< The absolute value of f [T-1 ~> s-1].
+  real, intent(in) :: u_star    !< The surface friction velocity [Z T-1 ~> m s-1].
+  real, intent(in) :: B_Flux    !< The surface buoyancy flux [Z2 T-3 ~> m2 s-3]
+  real, dimension(SZK_(GV)),intent(in) :: h   !<  The layer thickness [H ~> m or kg m-2].
+  real, intent(in) :: MLD_guess !< Mixing Layer depth guessed/found for iteration [Z ~> m].
+  real, dimension(SZK_(GV)+1) :: hz !< depth variable, only used in this routine [H ~> m]
+
+  ! local variables for this subroutine
+  integer :: nz
+  integer :: K ! integer for looping
+  integer :: i,j,n ! integer for looping, local only
+  real :: p1 ! ((B_flux * h))/(u_star^3), boundary layer depth by M-O depth, [nondim]
+  real :: p2 ! ((h f)/u_star ),  boundary layer depth by Ekman depth, [nondim]
+  real :: sm ! sigma_max: location of maximum of shape function in sigma coordinate [nondim]
+  real :: sm_I ! inverse of sm,[nondim]
+  real :: sm_I2 ! An inverse variable given by 1.0/(1.0 - sm), [nondim]
+  real :: hbl ! Boundary layer depth, same as MLD_guess [Z ~> m]
+  real :: F ! function, used in asymptotic model for sm, Equation 7 in Sane et al. 2024 [nondim]
+  real :: F_I ! Inverse of F, [nondim]
+  real :: Fp1 ! = F*p1, [nondim]
+
+  nz = SZK_(GV)+1
+  hz(1) = 0.0
+  do K=2,nz
+    hz(K) = hz(K-1) + h(K-1)*GV%H_to_Z
+  end do
+  hbl = MLD_Guess ! hbl is boundary layer depth.
+  shape_func(:) = 0.0  ! initializing the entire shape_function array
+
+  p1 = (hbl * absf)/(u_star)   ! Boundary layer depth divided by Ekman depth
+  p2 = ((-B_flux * hbl))/(u_star**3) ! Boundary layer depth divided by Monin-Obukhov depth
+  ! B_flux given negative sign to follow convention used in Sane et al. 2023,2024
+  ! p2 < 0 --> surface stabilizing i.e. heating, and p2 > 0 --> surface destabilizing i.e. cooling
+  p2 = p2 * 2.4390 ! dividing by von-Karman constant 0.41 i.e. multiply by 2.4390
+  ! This capping does not matter because these equations have asymptotes. Not sensitive beyond the caps.
+  p1 = min(p1, 2.0) ! capping p1 to less than 2.0. It is always >0.0.
+  p2 = max(p2, -8.0) ! capping p2 to -8.0 if less than -8.0
+  p2 = min(p2,  8.0) ! capping p2 to 8.0 if greater than 8.0
+  ! Empirical model to predict sm:
+  ! F1 is solely function of p2
+  F = exp( (-p2-CS%c6)/ CS%c7 ) ! originally, F=(CS%c4/(CS%c5+exp((-p2-CS%c6)/CS%c7)))+CS%c8
+  F = CS%c5 + F
+  F = CS%c4 / F
+  F = F + CS%c8
+  Fp1 = F*p1
+  Fp1 = max(Fp1, 1E-05) ! an arbitrary small value to cap Fp1, result insensitive below that value
+  F_I = 1.0 / ( Fp1 )
+  sm = CS%c2 + (CS%c3 * F_I)
+  sm = CS%c1 / sm
+  sm = min(sm,0.7) ! makes sure sm is less than 0.7, true sm range is from 0.2 to 0.60
+  sm = max(sm,0.1) ! makes sure sm is more than 0.1
+  sm= sm * hbl ! peak of shape function in model vertical coordinate z, or peak of shape function in z coordinate
+  sm_I = 1.0/sm ! inverse of sm x hbl
+  sm_I2 = 1.0/(hbl-sm)  ! inverse of (hbl-sm), as 0.1<sm<0.7, hbl>sm, hence (hbl-sm) always >0.0
+
+  ! gives the shape, quadratic above sm, cubic below sm.
+  ! see Equation 8 in Sane et al. 2024
+  ! interpolates a quadratic function from z=0 to z=sm*hbl, and then a cubic from z=sm*hbl to z=hbl
+  shape_func(:) = 0.0
+  do n = 2,nz
+    if  (hz(n) <= sm) then
+      shape_func(n) = -(hz(n) * sm_I)**2.0 + 2.0*(hz(n)*sm_I)
+    elseif  ((hz(n) > sm) .and. (hz(n) <= hbl)) then
+      shape_func(n) =  ( (1.99 * ((hz(n)-sm)*sm_I2)**3.0) - ( 2.98 *((hz(n)-sm)*sm_I2)**2.0 ) ) + 1.0
+      ! the coefficients 1.99 and 2.98 are dependent on the below value of 0.01.
+      ! They smoothly make the cubic go towards 0.01 below hbl.
+    elseif ((hz(n) > hbl)) then
+      shape_func(n) = 0.01 ! we set an arbitrary low value to 0.01
+      ! This value should be small such as 0.01, or 0.001, result is not sensitive. It should not be 0.0
+
+    endif
+  end do
+
+  end subroutine kappa_eqdisc
+
+  subroutine get_eqdisc_v0(CS, absf, B_flux, u_star, MLD_guess, v0_dummy)
+    ! gives velocity scale using equations that approximate neural network of Sane et al. 2023
+    type(energetic_PBL_CS),  intent(inout) :: CS     !< Energetic PBL control struct
+    real, intent(in) :: B_flux !< The surface buoyancy flux [Z2 T-3 ~> m2 s-3]
+    real, intent(in) :: u_star !< The surface friction velocity [Z T-1 ~> m s-1]
+    real, intent(in) :: MLD_guess !< boundary layer depth guessed/found for iteration [Z ~> m]
+
+    real, intent(in) :: absf  !< The absolute value of f [T-1 ~> s-1].
+    real, intent(inout) :: v0_dummy   !< velocity scale v0, local variable [Z T-1 ~> m s-1]
+
+    ! local variables for this subroutine
+    real :: bflux_c  ! capped bflux [Z2 T-3 ~> m2 s-3]
+    real :: ust_c    ! capped ustar [Z T-1 ~> m s-1]
+    real :: absf_c   ! capped absf [T-1 ~> s-1]
+    real :: p1, p2, p3, p4  ! nondimensional numbers [nondim]
+    ! from Sane et al. 2024: 
+    ! " p_1 &= \frac{a}{u_*} \sqrt{\frac{|B|}{f}}, \\ %= \sqrt{ \frac{L_{Ek}}{L_{MO}}}  \\
+    !   p_2 &= \frac{f}{\Omega},
+    !   Where $a = -1$ for $B \leq 0$ and $a = 1$ for $B > 0$ to distinguish between 
+    !  surface heating and cooling conditions. " 
+    ! p_3 = (CS%c11**2.0) / (p1-CS%c11), used to simplify calculation
+    ! p4 = (p1+CS%c10) + p3 , used to simplify calculation
+
+
+    if (B_flux <= CS%bflux_lower_cap) then
+      bflux_c = CS%bflux_lower_cap
+    elseif (B_flux >= CS%bflux_upper_cap) then
+      bflux_c = CS%bflux_upper_cap
+    else
+      bflux_c = B_flux
+    endif
+
+    ust_c = u_star
+
+    if (absf <= CS%f_lower) then   ! 
+      absf_c = CS%f_lower    ! 0.1 deg Latitude, cap avoids zero coriolis, 
+                             ! solution is not sensitive below 0.1 Degrees
+    else
+      absf_c = absf
+    endif
+
+    ! setting v0_dummy here:
+
+    if (bflux_c >= 0.0) then ! surface heating and neutral conditions
+    ! Equation 10 in Sane et al. 2024:
+    ! \frac{v}{u_*} = \frac{-c_9}{p_1 + c_{10} + \frac{c_{11}^2}{p_1 - c_{11}} }
+      p1 =  -(1.0/ust_c) * sqrt(abs(bflux_c)/absf_c) 
+      p3 = (CS%c11**2.0) / (p1-CS%c11)
+      p4 = (p1+CS%c10) + p3
+      v0_dummy = -CS%c9/p4
+
+    else ! surface cooling
+    ! Equation 11 in Sane et al. 2024:
+    ! \frac{v}{u_*}=\frac{c_{12} p_1 \cdot \sqrt{p_2} }{1 +  ...
+    ! \frac{(c_{13} e^{(-p_2/c_{14})} + c_{15}) }{p_1 ^2} }
+    ! p1 is merged in p3
+      p2 = absf_c/(CS%omega)
+      p3 = (CS%c12/ust_c)*sqrt(abs(bflux_c)/(CS%omega)) ! 7.2921e-05/s is CS%Omega, Earth rotation
+      p4 = CS%c13 * exp(-p2/ CS%c14) + CS%c15
+      p4 =  1 + (absf_c * p4 * ust_c * ust_c)/abs(bflux_c) 
+      v0_dummy  = (p3 / p4 ) + CS%c16 
+    endif
+
+    v0_dummy = v0_dummy * ust_c ! v0_dummy = v/u*, so it is multiplied by ust_c to get v0
+    v0_dummy = max(v0_dummy,CS%v0_lower_cap)  ! CS%v0_lower_cap
+    v0_dummy = min(v0_dummy,0.1) ! kept for safety, but has never hit this cap. 
+
+    ! v0_lower_cap has been set to 0.0001 as data below that values does not exist in the training
+    ! solution was tested for lower cap of 0.00001 and was found to be insensitive. 
+    ! sensitivity arises when lower cap is 0.0. That is when diffusivity attains extremely low values and 
+    ! they go near molecular diffusivity. Boundary layers might become "sub-grid" i.e. < 1 metre 
+    ! some cause issues such as anomlous surface warming. 
+    ! this needs further investigation, our choices are motivated by practicallity for now.
+end subroutine get_eqdisc_v0
+
 !> This subroutine calculates the change in potential energy and or derivatives
 !! for several changes in an interface's diapycnal diffusivity times a timestep.
 subroutine find_PE_chg(Kddt_h0, dKddt_h, hp_a, hp_b, Th_a, Sh_a, Th_b, Sh_b, &
@@ -2428,7 +2642,79 @@ subroutine energetic_PBL_init(Time, G, GV, US, param_file, diag, CS)
                  units="nondim", default=0.95,  do_not_log=(CS%LT_enhance_form==No_Langmuir))
   endif
 
-
+ !/Options related to Machine Learning Equation Discovery ! eqdisc
+  ! flag for using shape function from equation discovery - machine learning
+
+   call get_param(param_file, mdl, "Equation_Discovery_shape", CS%eqdisc, &
+                 "flag for activating equation discovery for sf", &
+                 units="nondim", default=.false.)
+
+   call get_param(param_file, mdl, "Equation_Discovery_velocity", CS%eqdisc_v0, &
+                   "flag for activating equation discovery for velocity scale", &
+                   units="nondim", default=.false.)
+
+  ! sets a  lower cap for abs_f (Coriolis parameter) required in equation for v_0. 
+  ! Small value, solution not sensitive below 1 deg Latitute
+  ! Default value of 2.5384E-07 corresponds to 0.1 deg. 
+  call get_param(param_file, mdl, "f_lower", CS%f_lower, &
+                       "value of lower limit cap for v0, default is for 0.1 deg, insensitive , & 
+                       below 1deg", units="s-1", default=2.5384E-07, scale=US%T_to_S) 
+
+  call get_param(param_file, mdl, "v0_lower_cap", CS%v0_lower_cap, &
+                       "value of lower limit cap for Coriolis in v0", & 
+                       units="m s-1", default=0.0001, scale=US%m_to_Z*US%T_to_s)
+
+  call get_param(param_file, mdl, "bflux_lower_cap", CS%bflux_lower_cap, &
+                       "value of lower limit cap for Bflux used in setting in v0", & 
+                       units="m2 s-3", default=-7.0E-07, scale=(US%m_to_L**2)*(US%T_to_s**3))
+
+  call get_param(param_file, mdl, "bflux_upper_cap", CS%bflux_upper_cap, &
+                       "value of upper limit cap for Bflux used in setting in v0", & 
+                       units="m2 s-3", default=7.0E-07, scale=(US%m_to_L**2)*(US%T_to_s**3))
+
+  ! The coefficients used for machine learned diffusivity
+  ! c1 to c8 used for sigma_m, 
+  call get_param(param_file, mdl, "c1", CS%c1, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.976)
+  call get_param(param_file, mdl, "c2", CS%c2, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=1.743)
+  call get_param(param_file, mdl, "c3", CS%c3, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=1.551)
+  call get_param(param_file, mdl, "c4", CS%c4, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.943)
+  call get_param(param_file, mdl, "c5", CS%c5, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.016)
+  call get_param(param_file, mdl, "c6", CS%c6, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.693)
+  call get_param(param_file, mdl, "c7", CS%c7, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.379)
+  call get_param(param_file, mdl, "c8", CS%c8, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=2.194)
+
+  ! coefficients related to surface heating v_0, obtained using Genetic Programming
+  call get_param(param_file, mdl, "c9", CS%c9, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.1426)
+  call get_param(param_file, mdl, "c10", CS%c10, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.0434)
+  call get_param(param_file, mdl, "c11", CS%c11, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=1.80)
+
+  ! coefficients related to surface cooling v_0, obtained by empirical fitting (brain)
+  call get_param(param_file, mdl, "c12", CS%c12, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.098)
+  call get_param(param_file, mdl, "c13", CS%c13, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=45.0)
+  call get_param(param_file, mdl, "c14", CS%c14, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.35)
+  call get_param(param_file, mdl, "c15", CS%c15, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=3.29)
+  call get_param(param_file, mdl, "c16", CS%c16, &
+                       "Coefficient used for ML diffusivity,  ", units="nondim", default=0.0764)
+
+  !/ options end for Machine Learning Equation Discovery
+
+
+
 !/ Logging parameters
   ! This gives a minimum decay scale that is typically much less than Angstrom.
   CS%ustar_min = 2e-4*CS%omega*(GV%Angstrom_Z + GV%dZ_subroundoff)