Skip to content

Commit 7b8f824

Browse files
aakashsaneaakashsane
committed
Machine learning based diffusivity changes, only to the shape function
1 parent 65ef5c7 commit 7b8f824

1 file changed

Lines changed: 239 additions & 12 deletions

File tree

‎src/shared/cvmix_kpp.F90‎

Lines changed: 239 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ module cvmix_kpp
6565
integer, parameter :: LANGMUIR_ENTRAINMENT_LWF16 = 1
6666
integer, parameter :: LANGMUIR_ENTRAINMENT_LF17 = 2
6767
integer, parameter :: LANGMUIR_ENTRAINMENT_RWHGK16 = 3
68+
integer, parameter :: ML_DIFFUSIVITY_SHAPE = 5 ! ML_Diffusivity
6869

6970
! !PUBLIC MEMBER FUNCTIONS:
7071

@@ -88,7 +89,7 @@ module cvmix_kpp
8889
public :: cvmix_kpp_compute_nu_at_OBL_depth_LMD94
8990
public :: cvmix_kpp_EFactor_model
9091
public :: cvmix_kpp_ustokes_SL_model
91-
92+
public :: cvmix_kpp_compute_ER_depth
9293

9394
interface cvmix_coeffs_kpp
9495
module procedure cvmix_coeffs_kpp_low
@@ -210,6 +211,12 @@ module cvmix_kpp
210211

211212
real(cvmix_r8) :: CVt2 ! Tunable parameter for convection entrainment
212213
! (Only used with StokesMOST)
214+
real(cvmix_r8) :: ER_Cb ! Entrainment Rule TKE buoyancy production weight [nondim]
215+
real(cvmix_r8) :: ER_Cs ! Entrainment Rule TKE Stokes production weight [nondim]
216+
real(cvmix_r8) :: ER_Cu ! Entrainment Rule TKE shear production weight [nondim]
217+
218+
logical :: ML_diffusivity ! uses the Machine Learned equations to set diffusivity
219+
! from Sane et al. (2025) https://doi.org/10.31219/osf.io/uab7v_v2
213220

214221
end type cvmix_kpp_params_type
215222

@@ -229,7 +236,8 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
229236
old_vals, lEkman, lStokesMOST, lMonOb, lnoDGat1, &
230237
lenhanced_diff, lnonzero_surf_nonlocal, &
231238
Langmuir_mixing_str, Langmuir_entrainment_str, &
232-
l_LMD_ws, CVmix_kpp_params_user)
239+
l_LMD_ws, ML_diffusivity, &
240+
ER_Cb, ER_Cs, ER_Cu, CVmix_kpp_params_user)
233241

234242
! !DESCRIPTION:
235243
! Initialization routine for KPP mixing.
@@ -247,7 +255,10 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
247255
zeta_s, &
248256
surf_layer_ext, &
249257
CVt2, &
250-
Cv
258+
Cv, &
259+
ER_Cb, &
260+
ER_Cs, &
261+
ER_Cu
251262
character(len=*), optional, intent(in) :: interp_type, &
252263
interp_type2, &
253264
MatchTechnique, &
@@ -260,10 +271,9 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
260271
lnoDGat1, &
261272
lenhanced_diff, &
262273
lnonzero_surf_nonlocal, &
263-
l_LMD_ws
264-
265-
! !OUTPUT PARAMETERS:
266-
type(cvmix_kpp_params_type), intent(inout), target, optional :: &
274+
l_LMD_ws, &
275+
ML_diffusivity
276+
type(cvmix_kpp_params_type), optional, target, intent(inout) :: &
267277
CVmix_kpp_params_user
268278

269279
!EOP
@@ -272,6 +282,7 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
272282
real(cvmix_r8) :: zm, zs, a_m, a_s, c_m, c_s
273283
real(cvmix_r8) :: Cstar_loc, vonkar_loc, surf_layer_ext_loc
274284
real(cvmix_r8) :: nonlocal_coeff
285+
real(cvmix_r8) :: ER_Cb_loc, ER_Cs_loc, ER_Cu_loc
275286

276287
if (present(ri_crit)) then
277288
if (ri_crit.lt.cvmix_zero) then
@@ -452,6 +463,8 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
452463
case ('ParabolicNonLocal')
453464
call cvmix_put_kpp('MatchTechnique', CVMIX_KPP_PARABOLIC_NONLOCAL, &
454465
CVmix_kpp_params_user)
466+
case ('ML_Diffusivity_Shape') ! ML_Diffusivity shape function
467+
call cvmix_put_kpp('MatchTechnique', ML_DIFFUSIVITY_SHAPE, CVmix_kpp_params_user)
455468
case DEFAULT
456469
print*, "ERROR: ", trim(MatchTechnique), " is not a valid choice ", &
457470
"for MatchTechnique!"
@@ -462,6 +475,16 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
462475
CVmix_kpp_params_user)
463476
end if
464477

478+
! --- Add this block to for ML_diffusivity ---
479+
if (present(ML_diffusivity)) then ! ML_diffusivity
480+
call cvmix_put_kpp('ML_diffusivity', ML_diffusivity, CVmix_kpp_params_user)
481+
if (ML_diffusivity) then
482+
call cvmix_put_kpp('MatchTechnique', ML_DIFFUSIVITY_SHAPE, CVmix_kpp_params_user)
483+
end if
484+
else
485+
call cvmix_put_kpp('ML_diffusivity', .false., CVmix_kpp_params_user)
486+
end if
487+
465488
if (present(old_vals)) then
466489
select case (trim(old_vals))
467490
case ("overwrite")
@@ -586,6 +609,24 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
586609
end if
587610

588611
! Initialize parameters for enhanced entrainment
612+
if (present(ER_Cb)) then
613+
ER_Cb_loc = ER_Cb
614+
else
615+
ER_Cb_loc = 0.96_cvmix_r8
616+
end if
617+
call cvmix_put_kpp('ER_Cb', ER_Cb_loc, CVmix_kpp_params_user)
618+
if (present(ER_Cs)) then
619+
ER_Cs_loc = ER_Cs
620+
else
621+
ER_Cs_loc = 0.038_cvmix_r8
622+
end if
623+
call cvmix_put_kpp('ER_Cs', ER_Cs_loc, CVmix_kpp_params_user)
624+
if (present(ER_Cu)) then
625+
ER_Cu_loc = ER_Cu
626+
else
627+
ER_Cu_loc = 0.023_cvmix_r8
628+
end if
629+
call cvmix_put_kpp('ER_Cu', ER_Cu_loc, CVmix_kpp_params_user)
589630
call cvmix_put_kpp('c_ST', 0.17_cvmix_r8, CVmix_kpp_params_user)
590631
call cvmix_put_kpp('c_CT', 0.15_cvmix_r8, CVmix_kpp_params_user)
591632
call cvmix_put_kpp('c_LT', 0.083_cvmix_r8, CVmix_kpp_params_user)
@@ -654,6 +695,7 @@ subroutine cvmix_coeffs_kpp_wrap(CVmix_vars, CVmix_kpp_params_user)
654695
nlev, max_nlev, &
655696
CVmix_vars%LangmuirEnhancementFactor, &
656697
CVmix_vars%StokesMostXi, &
698+
CVmix_vars%Coriolis, &
657699
CVmix_kpp_params_user)
658700

659701
call cvmix_update_wrap(CVmix_kpp_params_in%handle_old_vals, max_nlev, &
@@ -677,7 +719,7 @@ subroutine cvmix_coeffs_kpp_low(Mdiff_out, Tdiff_out, Sdiff_out, zw, zt, &
677719
old_Mdiff, old_Tdiff, old_Sdiff, OBL_depth, &
678720
kOBL_depth, Tnonlocal, Snonlocal, surf_fric,&
679721
surf_buoy, nlev, max_nlev, Langmuir_EFactor,&
680-
StokesXI,CVmix_kpp_params_user)
722+
StokesXI,Coriolis,CVmix_kpp_params_user)
681723

682724
! !DESCRIPTION:
683725
! Computes vertical diffusion coefficients for the KPP boundary layer mixing
@@ -702,6 +744,7 @@ subroutine cvmix_coeffs_kpp_low(Mdiff_out, Tdiff_out, Sdiff_out, zw, zt, &
702744
! Langmuir enhancement factor
703745
real(cvmix_r8), intent(in), optional :: Langmuir_EFactor
704746
real(cvmix_r8), intent(in), optional :: StokesXI
747+
real(cvmix_r8), intent(in), optional :: Coriolis ! required for ML_diffusivity
705748
! !INPUT/OUTPUT PARAMETERS:
706749
real(cvmix_r8), dimension(max_nlev+1), intent(inout) :: Mdiff_out, &
707750
Tdiff_out, &
@@ -769,6 +812,14 @@ subroutine cvmix_coeffs_kpp_low(Mdiff_out, Tdiff_out, Sdiff_out, zw, zt, &
769812
! Parameters for Stokes_MOST
770813
real(cvmix_r8) :: Gcomposite, Hsigma, sigh, T_NLenhance , S_NLenhance , XIone
771814

815+
! Parameters for the machine learning part, ML_diffusivity
816+
real(cvmix_r8) :: sigma_max ! sigma coordinate of location of maximum diffusivity
817+
real(cvmix_r8) :: g_sigma ! shape function at a sigma coordinate.
818+
real(cvmix_r8) :: L_h ! Non-dimensional L_h = B*OBL_depth/u_*^3
819+
real(cvmix_r8) :: E_h ! Non-dimensional E_h = OBL_depth * Coriolis /u_*
820+
real(cvmix_r8) :: F_inter_func ! Stands for F_intermediate_function,
821+
! Non-dimensional intermediate function used to calculate sigma_max
822+
772823
! Constant from params
773824
integer :: interp_type2, MatchTechnique
774825

@@ -1183,11 +1234,46 @@ subroutine cvmix_coeffs_kpp_low(Mdiff_out, Tdiff_out, Sdiff_out, zw, zt, &
11831234
call cvmix_kpp_compute_turbulent_scales(sigma, OBL_depth, surf_buoy, &
11841235
surf_fric, XIone, w_m, w_s, &
11851236
CVmix_kpp_params_user)
1237+
1238+
1239+
if (CVmix_kpp_params_in%ML_diffusivity) then ! ML_diffusivity
1240+
!!! Evaluating L_h and E_h
1241+
L_h = -surf_buoy * OBL_depth / (surf_fric ** 3.0) ! OBL / Monin-Obukhov-Depth
1242+
E_h = OBL_depth * Coriolis / surf_fric ! OBL / Ekman Depth i.e. hf/u*
1243+
F_inter_func = ( cvmix_one / ( 0.0712 + 0.4380 * exp(-1.0*(2.6821 * L_h)) ) ) + 1.5845
1244+
sigma_max = (F_inter_func * E_h) / ( 1.7908*(F_inter_func * E_h) + 0.6904)
1245+
!!! capping sigma_max between 0.1 and 0.7
1246+
sigma_max = min( max(sigma_max, 0.1), 0.7)
1247+
endif
11861248
do kw=2,kwup
11871249
! (3b) Evaluate G(sigma) at each cell interface
1188-
MshapeAtS = cvmix_math_evaluate_cubic(Mshape, sigma(kw))
1189-
TshapeAtS = cvmix_math_evaluate_cubic(Tshape, sigma(kw))
1190-
SshapeAtS = cvmix_math_evaluate_cubic(Sshape, sigma(kw))
1250+
1251+
if (CVmix_kpp_params_in%ML_diffusivity) then ! ML_diffusivity
1252+
1253+
! ML-diffusivity modification:
1254+
if (sigma(kw) .le. sigma_max ) then ! ML based shape function
1255+
! quadratic part above sigma_max
1256+
g_sigma = (2.0*sigma(kw)/sigma_max) - (sigma(kw)/sigma_max)**2.0
1257+
g_sigma = g_sigma * (4.0/27.0) ! reduces the amplitude from 1 to 4/27 to match G(\sigma) of KPP
1258+
MshapeAtS = g_sigma
1259+
TshapeAtS = g_sigma
1260+
SshapeAtS = g_sigma
1261+
else
1262+
! cubic part below sigma_max
1263+
g_sigma = 2.0*((sigma(kw) - sigma_max)/(cvmix_one - sigma_max))**3.0 &
1264+
- 3.0*((sigma(kw) - sigma_max)/(cvmix_one - sigma_max))**2.0 &
1265+
+ cvmix_one
1266+
g_sigma = g_sigma * (4.0/27.0) ! reduces the amplitude from 1 to 4/27 to match G(\sigma) of KPP
1267+
MshapeAtS = g_sigma
1268+
TshapeAtS = g_sigma
1269+
SshapeAtS = g_sigma
1270+
end if
1271+
1272+
else
1273+
MshapeAtS = cvmix_math_evaluate_cubic(Mshape, sigma(kw))
1274+
TshapeAtS = cvmix_math_evaluate_cubic(Tshape, sigma(kw))
1275+
SshapeAtS = cvmix_math_evaluate_cubic(Sshape, sigma(kw))
1276+
end if
11911277
! The RWHGK16 Langmuir uses the shape function to shape the
11921278
! enhancement to the mixing coefficient.
11931279
ShapeNoMatchAtS = cvmix_math_evaluate_cubic(NMshape, sigma(kw))
@@ -1277,7 +1363,7 @@ end subroutine cvmix_coeffs_kpp_low
12771363
! !IROUTINE: cvmix_put_kpp_real
12781364
! !INTERFACE:
12791365

1280-
subroutine cvmix_put_kpp_real(varname, val, CVmix_kpp_params_user)
1366+
subroutine cvmix_put_kpp_real(varname, val, CVmix_kpp_params_user)
12811367

12821368
! !DESCRIPTION:
12831369
! Write a real value into a cvmix\_kpp\_params\_type variable.
@@ -1335,6 +1421,12 @@ subroutine cvmix_put_kpp_real(varname, val, CVmix_kpp_params_user)
13351421
CVmix_kpp_params_out%CVt2 = val
13361422
case ('nonlocal_coeff')
13371423
CVmix_kpp_params_out%nonlocal_coeff = val
1424+
case ('ER_Cb')
1425+
CVmix_kpp_params_out%ER_Cb = val
1426+
case ('ER_Cs')
1427+
CVmix_kpp_params_out%ER_Cs = val
1428+
case ('ER_Cu')
1429+
CVmix_kpp_params_out%ER_Cu = val
13381430
case ('c_CT')
13391431
CVmix_kpp_params_out%c_CT = val
13401432
case ('c_ST')
@@ -1452,6 +1544,8 @@ subroutine cvmix_put_kpp_logical(varname, val, CVmix_kpp_params_user)
14521544
CVmix_kpp_params_out%lenhanced_diff = val
14531545
case ('l_LMD_ws')
14541546
CVmix_kpp_params_out%l_LMD_ws = val
1547+
case ('ML_diffusivity')
1548+
CVmix_kpp_params_out%ML_diffusivity = val
14551549
case DEFAULT
14561550
print*, "ERROR: ", trim(varname), " is not a boolean variable!"
14571551
stop 1
@@ -3508,4 +3602,137 @@ subroutine cvmix_kpp_compute_StokesXi (zi, zk, kSL, SLDepth, &
35083602

35093603
end subroutine cvmix_kpp_compute_StokesXi
35103604

3605+
3606+
!BOP
3607+
3608+
! !DESCRIPTION:
3609+
! Use Entrainment Rule, BEdE_ER, To Find Entrainment Flux and Depth
3610+
! for each assumed OBL_depth = cell centers,
3611+
! until the boundary layer depth, ERdepth > 0 kER_depth are determined, OR
3612+
! if there is no viable solution ERdepth = -1 , kER_depth=-1
3613+
3614+
subroutine cvmix_kpp_compute_ER_depth( z_inter,Nsq,OBL_depth, &
3615+
uStar,Bsfc_ns,surfBuoy,StokesXI,BEdE_ER,ERdepth, &
3616+
CVMix_kpp_params_user)
3617+
3618+
! !INPUT PARAMETERS:
3619+
real(cvmix_r8), dimension(:), intent(in) :: &
3620+
z_inter, & ! Interface heights <= 0 [m]
3621+
Nsq ! Column of Buoyancy Gradients at interfaces
3622+
real(cvmix_r8), dimension(:), intent(in) :: &
3623+
OBL_depth, & ! Array of assumed OBL depths >0 at cell centers [m]
3624+
surfBuoy, & ! surface Buoyancy flux surface to OBL_depth
3625+
StokesXI, & ! Stokes similarity parameter given OBL_depth
3626+
BEdE_ER ! Parameterized Entrainment Rule given OBL_depth
3627+
real(cvmix_r8), intent(in) :: uStar ! surface friction velocity [m s-1]
3628+
real(cvmix_r8), intent(in) :: Bsfc_ns ! non-solar surface buoyancy flux boundary condition
3629+
3630+
type(cvmix_kpp_params_type), optional, target, intent(in) :: CVmix_kpp_params_user
3631+
3632+
! !OUTPUT PARAMETERS:
3633+
real(cvmix_r8), intent(out) :: ERdepth ! ER Boundary Layer Depth [m]
3634+
3635+
!EOP
3636+
!BOC
3637+
3638+
! Local variables
3639+
integer :: iz, nlev , kbl , kinv
3640+
real(cvmix_r8), dimension(size(OBL_depth)+1) :: zdepth, BEdE ! surface then cell-centers<0
3641+
real(cvmix_r8), dimension(size(z_inter)+1) :: sigma, Bflux ! interface values
3642+
real(cvmix_r8) :: ws_i, Cemp_Rs, Gsig_i, Fdelrho, Bnonlocal, sigE, maxNsq, BEnt
3643+
real(kind=cvmix_r8), dimension(4) :: coeffs
3644+
type(cvmix_kpp_params_type), pointer :: CVmix_kpp_params_in
3645+
3646+
CVmix_kpp_params_in => CVmix_kpp_params_saved
3647+
if (present(CVmix_kpp_params_user)) then
3648+
CVmix_kpp_params_in => CVmix_kpp_params_user
3649+
end if
3650+
3651+
nlev = size(OBL_depth)
3652+
Cemp_Rs = 4.7_cvmix_r8
3653+
Fdelrho = cvmix_one
3654+
maxNsq = 0.0
3655+
do kbl = 2, nlev+1
3656+
if ( Nsq(kbl) > maxNsq ) then
3657+
kinv = kbl
3658+
maxNsq = Nsq(kbl)
3659+
endif
3660+
enddo
3661+
3662+
! Set default output values (no solution)
3663+
ERdepth = -cvmix_one
3664+
! Set surface values
3665+
zdepth(1) = cvmix_zero
3666+
BEdE(1) = cvmix_zero
3667+
sigma(:) = cvmix_zero
3668+
Bflux(1) = Bsfc_ns ! non-solar surface buoyancy boundary condition for all kbl
3669+
! Set OBL_depth(1)=top cell center values
3670+
zdepth(2) = -OBL_depth(1)
3671+
BEdE(2) = cvmix_zero
3672+
3673+
do kbl = 2, nlev
3674+
zdepth(kbl+1)= -OBL_depth(kbl)
3675+
BEdE(kbl+1) = cvmix_zero
3676+
3677+
sigma(kbl+1) = cvmix_one
3678+
Bflux(kbl+1) = cvmix_zero
3679+
sigma(kbl+2) = -z_inter(kbl+1)/OBL_depth(kbl) ! > 1.0
3680+
Bflux(kbl+2) = cvmix_zero
3681+
Bnonlocal = 0.5_cvmix_r8 * Cemp_Rs * (abs(surfBuoy(kbl)) - surfBuoy(kbl))
3682+
3683+
do iz = kbl,1,-1
3684+
if (iz .gt. 1) then
3685+
sigma(iz) = -z_inter(iz)/OBL_depth(kbl) ! < 1.0
3686+
call cvmix_kpp_compute_turbulent_scales(sigma(iz), OBL_depth(kbl), surfBuoy(kbl), uStar, StokesXI(kbl), & !0d
3687+
w_s=ws_i , CVmix_kpp_params_user=CVmix_kpp_params_user)
3688+
Gsig_i = cvmix_kpp_composite_shape(sigma(iz))
3689+
Bflux(iz) = Gsig_i * (OBL_depth(kbl) * ws_i * Nsq(iz) - Bnonlocal)
3690+
endif
3691+
3692+
! find the peak
3693+
if ( (Bflux(iz+1) .gt. Bflux(iz+2)) .and. (Bflux(iz+1) .ge. Bflux(iz)) .and. &
3694+
(Bflux(iz+1) .gt. cvmix_zero) ) then
3695+
! Find sigE (the root of the derivative of the quadratic polynomial
3696+
! interpolating (sigma(i), Bflux(i)) for i in [iz, iz+1, iz+2])
3697+
! Also find BEnt (value of quadratic at sigE)
3698+
! call cvmix_math_poly_interp(coeffs, sigma(iz:iz+2), Bflux(iz:iz+2))
3699+
! comment by Aakash: the above is the original line,
3700+
! it gives me errors so I changed it to below call. not sure if it is correct
3701+
call cvmix_math_poly_interp(coeffs, CVMIX_MATH_INTERP_QUAD, &
3702+
sigma(iz:iz+1), Bflux(iz:iz+1), &
3703+
sigma(iz+2), Bflux(iz+2))
3704+
! coeffs(3) = 0 => sigma(iz), sigma(iz+1), and sigma(iz+2) are not unique values
3705+
! so there interpolation returned a linear equation. In this case we select
3706+
! (sigma(iz+1), Bflux(iz+1)) as the maximum.
3707+
if (coeffs(3) .eq. cvmix_zero) then
3708+
sigE = sigma(iz+1)
3709+
Bent = Bflux(iz+1)
3710+
else
3711+
sigE = -0.5_cvmix_r8 * (coeffs(2) / coeffs(3))
3712+
Bent = cvmix_math_evaluate_cubic(coeffs, sigE)
3713+
endif
3714+
Fdelrho = cvmix_one
3715+
BEdE(kbl+1) = Fdelrho*BEnt*sigE*OBL_depth(kbl)
3716+
exit ! No exit leaves BEdE(kbl+1) = cvmix_zero
3717+
endif
3718+
enddo
3719+
3720+
if ( (BEdE(kbl+1) > BEdE_ER(kbl)) .and. (Bsfc_ns < cvmix_zero) ) then
3721+
call cvmix_math_poly_interp(coeffs, CVmix_kpp_params_in%interp_type, &
3722+
zdepth(kbl:kbl+1) , BEdE(kbl:kbl+1), &
3723+
zdepth(kbl-1) , BEdE(kbl-1) ) ! surface values if kbl=2;
3724+
coeffs(1) = coeffs(1) - BEdE_ER(kbl)
3725+
ERdepth = -cvmix_math_cubic_root_find(coeffs, 0.5_cvmix_r8 * &
3726+
(zdepth(kbl)+zdepth(kbl+1)))
3727+
3728+
exit
3729+
endif
3730+
3731+
enddo
3732+
3733+
!EOC
3734+
3735+
end subroutine cvmix_kpp_compute_ER_depth
3736+
3737+
35113738
end module cvmix_kpp

0 commit comments

Comments
 (0)