Skip to content

added ML_diffusivity shape function #111

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.

Sign up for GitHub

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

Open
aakashsane wants to merge 4 commits into
base: master
Choose a base branch
from
Open

added ML_diffusivity shape function #111

Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7b8f824
Machine learning based diffusivity changes, only to the shape function
aakashsane Jan 28, 2026
cf4c5b9
Removed trailing white spaces
aakashsane Jan 28, 2026
96a02c6
Merge branch 'master' into ML_diffusivity_3
aakashsane Jan 28, 2026
aa778fb
Clean up following merge of master
mnlevy1981 Jan 30, 2026
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
217 changes: 207 additions & 10 deletions src/shared/cvmix_kpp.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ module cvmix_kpp
integer, parameter :: LANGMUIR_ENTRAINMENT_LWF16 = 1
integer, parameter :: LANGMUIR_ENTRAINMENT_LF17 = 2
integer, parameter :: LANGMUIR_ENTRAINMENT_RWHGK16 = 3
integer, parameter :: ML_DIFFUSIVITY_SHAPE = 5 ! ML_Diffusivity

! !PUBLIC MEMBER FUNCTIONS:

Expand Down Expand Up @@ -214,6 +215,9 @@ module cvmix_kpp
real(cvmix_r8) :: ER_Cs ! Entrainment Rule TKE Stokes production weight [nondim]
real(cvmix_r8) :: ER_Cu ! Entrainment Rule TKE shear production weight [nondim]

logical :: ML_diffusivity ! uses the Machine Learned equations to set diffusivity
! from Sane et al. (2025) https://doi.org/10.31219/osf.io/uab7v_v2

end type cvmix_kpp_params_type

!EOP
Expand All @@ -232,7 +236,8 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
old_vals, lEkman, lStokesMOST, lMonOb, lnoDGat1, &
lenhanced_diff, lnonzero_surf_nonlocal, &
Langmuir_mixing_str, Langmuir_entrainment_str, &
l_LMD_ws, ER_Cb, ER_Cs, ER_Cu, CVmix_kpp_params_user)
l_LMD_ws, ML_diffusivity, &
ER_Cb, ER_Cs, ER_Cu, CVmix_kpp_params_user)

! !DESCRIPTION:
! Initialization routine for KPP mixing.
Expand Down Expand Up @@ -266,10 +271,9 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
lnoDGat1, &
lenhanced_diff, &
lnonzero_surf_nonlocal, &
l_LMD_ws

! !OUTPUT PARAMETERS:
type(cvmix_kpp_params_type), intent(inout), target, optional :: &
l_LMD_ws, &
ML_diffusivity
type(cvmix_kpp_params_type), optional, target, intent(inout) :: &
CVmix_kpp_params_user

!EOP
Expand All @@ -279,6 +283,7 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
real(cvmix_r8) :: Cstar_loc, vonkar_loc, surf_layer_ext_loc
real(cvmix_r8) :: ER_Cb_loc, ER_Cs_loc, ER_Cu_loc
real(cvmix_r8) :: nonlocal_coeff
real(cvmix_r8) :: ER_Cb_loc, ER_Cs_loc, ER_Cu_loc

if (present(ri_crit)) then
if (ri_crit.lt.cvmix_zero) then
Expand Down Expand Up @@ -459,6 +464,8 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
case ('ParabolicNonLocal')
call cvmix_put_kpp('MatchTechnique', CVMIX_KPP_PARABOLIC_NONLOCAL, &
CVmix_kpp_params_user)
case ('ML_Diffusivity_Shape') ! ML_Diffusivity shape function
call cvmix_put_kpp('MatchTechnique', ML_DIFFUSIVITY_SHAPE, CVmix_kpp_params_user)
Comment on lines +468 to +469
Copy link
Contributor

@mnlevy1981 mnlevy1981 Feb 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens if MatchTechnique = 'ML_Diffusivity_Shape' but ML_diffusivity = .false.? If that is not a reasonable configuration, then I think you should just set ML_diffusivity = .true. when using 'ML_Diffusivity_Shape' and .false. in all other configurations

case DEFAULT
print*, "ERROR: ", trim(MatchTechnique), " is not a valid choice ", &
"for MatchTechnique!"
Expand All @@ -469,6 +476,16 @@ subroutine cvmix_init_kpp(ri_crit, minOBLdepth, maxOBLdepth, minVtsqr, &
CVmix_kpp_params_user)
end if

! --- Add this block to for ML_diffusivity ---
if (present(ML_diffusivity)) then ! ML_diffusivity
call cvmix_put_kpp('ML_diffusivity', ML_diffusivity, CVmix_kpp_params_user)
if (ML_diffusivity) then
call cvmix_put_kpp('MatchTechnique', ML_DIFFUSIVITY_SHAPE, CVmix_kpp_params_user)
end if
else
call cvmix_put_kpp('ML_diffusivity', .false., CVmix_kpp_params_user)
end if

if (present(old_vals)) then
select case (trim(old_vals))
case ("overwrite")
Expand Down Expand Up @@ -679,6 +696,7 @@ subroutine cvmix_coeffs_kpp_wrap(CVmix_vars, CVmix_kpp_params_user)
nlev, max_nlev, &
CVmix_vars%LangmuirEnhancementFactor, &
CVmix_vars%StokesMostXi, &
CVmix_vars%Coriolis, &
CVmix_kpp_params_user)

call cvmix_update_wrap(CVmix_kpp_params_in%handle_old_vals, max_nlev, &
Expand All @@ -702,7 +720,7 @@ subroutine cvmix_coeffs_kpp_low(Mdiff_out, Tdiff_out, Sdiff_out, zw, zt, &
old_Mdiff, old_Tdiff, old_Sdiff, OBL_depth, &
kOBL_depth, Tnonlocal, Snonlocal, surf_fric,&
surf_buoy, nlev, max_nlev, Langmuir_EFactor,&
StokesXI,CVmix_kpp_params_user)
StokesXI,Coriolis,CVmix_kpp_params_user)

! !DESCRIPTION:
! Computes vertical diffusion coefficients for the KPP boundary layer mixing
Expand All @@ -727,6 +745,7 @@ subroutine cvmix_coeffs_kpp_low(Mdiff_out, Tdiff_out, Sdiff_out, zw, zt, &
! Langmuir enhancement factor
real(cvmix_r8), intent(in), optional :: Langmuir_EFactor
real(cvmix_r8), intent(in), optional :: StokesXI
real(cvmix_r8), intent(in), optional :: Coriolis ! required for ML_diffusivity
! !INPUT/OUTPUT PARAMETERS:
real(cvmix_r8), dimension(max_nlev+1), intent(inout) :: Mdiff_out, &
Tdiff_out, &
Expand Down Expand Up @@ -794,6 +813,14 @@ subroutine cvmix_coeffs_kpp_low(Mdiff_out, Tdiff_out, Sdiff_out, zw, zt, &
! Parameters for Stokes_MOST
real(cvmix_r8) :: Gcomposite, Hsigma, sigh, T_NLenhance , S_NLenhance , XIone

! Parameters for the machine learning part, ML_diffusivity
real(cvmix_r8) :: sigma_max ! sigma coordinate of location of maximum diffusivity
real(cvmix_r8) :: g_sigma ! shape function at a sigma coordinate.
real(cvmix_r8) :: L_h ! Non-dimensional L_h = B*OBL_depth/u_*^3
real(cvmix_r8) :: E_h ! Non-dimensional E_h = OBL_depth * Coriolis /u_*
real(cvmix_r8) :: F_inter_func ! Stands for F_intermediate_function,
! Non-dimensional intermediate function used to calculate sigma_max

! Constant from params
integer :: interp_type2, MatchTechnique

Expand Down Expand Up @@ -1208,11 +1235,46 @@ subroutine cvmix_coeffs_kpp_low(Mdiff_out, Tdiff_out, Sdiff_out, zw, zt, &
call cvmix_kpp_compute_turbulent_scales(sigma, OBL_depth, surf_buoy, &
surf_fric, XIone, w_m, w_s, &
CVmix_kpp_params_user)


if (CVmix_kpp_params_in%ML_diffusivity) then ! ML_diffusivity
!!! Evaluating L_h and E_h
L_h = -surf_buoy * OBL_depth / (surf_fric ** 3.0) ! OBL / Monin-Obukhov-Depth
E_h = OBL_depth * Coriolis / surf_fric ! OBL / Ekman Depth i.e. hf/u*
F_inter_func = ( cvmix_one / ( 0.0712 + 0.4380 * exp(-1.0*(2.6821 * L_h)) ) ) + 1.5845
sigma_max = (F_inter_func * E_h) / ( 1.7908*(F_inter_func * E_h) + 0.6904)
!!! capping sigma_max between 0.1 and 0.7
sigma_max = min( max(sigma_max, 0.1), 0.7)
endif
do kw=2,kwup
! (3b) Evaluate G(sigma) at each cell interface
MshapeAtS = cvmix_math_evaluate_cubic(Mshape, sigma(kw))
TshapeAtS = cvmix_math_evaluate_cubic(Tshape, sigma(kw))
SshapeAtS = cvmix_math_evaluate_cubic(Sshape, sigma(kw))

if (CVmix_kpp_params_in%ML_diffusivity) then ! ML_diffusivity

! ML-diffusivity modification:
if (sigma(kw) .le. sigma_max ) then ! ML based shape function
! quadratic part above sigma_max
g_sigma = (2.0*sigma(kw)/sigma_max) - (sigma(kw)/sigma_max)**2.0
g_sigma = g_sigma * (4.0/27.0) ! reduces the amplitude from 1 to 4/27 to match G(\sigma) of KPP
Comment on lines +1258 to +1259
Copy link
Contributor

@mnlevy1981 mnlevy1981 Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should define this shape function in the select case (MatchTechnique) block (part of case DEFAULT; note that we want to compute that same shape function and use it as Tshape2 and Sshape2) but for ML_DIFFUSIVITY_SHAPE we want to encode this quadratic in such a way that cvmix_evaluate_cubic() reconstructs it correctly.

MshapeAtS = g_sigma
TshapeAtS = g_sigma
SshapeAtS = g_sigma
else
! cubic part below sigma_max
g_sigma = 2.0*((sigma(kw) - sigma_max)/(cvmix_one - sigma_max))**3.0 &
- 3.0*((sigma(kw) - sigma_max)/(cvmix_one - sigma_max))**2.0 &
+ cvmix_one
g_sigma = g_sigma * (4.0/27.0) ! reduces the amplitude from 1 to 4/27 to match G(\sigma) of KPP
Comment on lines +1265 to +1268
Copy link
Contributor

@mnlevy1981 mnlevy1981 Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This cubic should also be evaluated with cvmix_evaluate_cubic() -- maybe we need TshapeML2 (along with M and S versions)? Or if we are using the same shape function for M, T, and S, we could go with MLshape and MLshape2?

MshapeAtS = g_sigma
TshapeAtS = g_sigma
SshapeAtS = g_sigma
end if

else
MshapeAtS = cvmix_math_evaluate_cubic(Mshape, sigma(kw))
TshapeAtS = cvmix_math_evaluate_cubic(Tshape, sigma(kw))
SshapeAtS = cvmix_math_evaluate_cubic(Sshape, sigma(kw))
end if
! The RWHGK16 Langmuir uses the shape function to shape the
! enhancement to the mixing coefficient.
ShapeNoMatchAtS = cvmix_math_evaluate_cubic(NMshape, sigma(kw))
Expand Down Expand Up @@ -1302,7 +1364,7 @@ end subroutine cvmix_coeffs_kpp_low
! !IROUTINE: cvmix_put_kpp_real
! !INTERFACE:

subroutine cvmix_put_kpp_real(varname, val, CVmix_kpp_params_user)
subroutine cvmix_put_kpp_real(varname, val, CVmix_kpp_params_user)

! !DESCRIPTION:
! Write a real value into a cvmix\_kpp\_params\_type variable.
Expand Down Expand Up @@ -1483,6 +1545,8 @@ subroutine cvmix_put_kpp_logical(varname, val, CVmix_kpp_params_user)
CVmix_kpp_params_out%lenhanced_diff = val
case ('l_LMD_ws')
CVmix_kpp_params_out%l_LMD_ws = val
case ('ML_diffusivity')
CVmix_kpp_params_out%ML_diffusivity = val
case DEFAULT
print*, "ERROR: ", trim(varname), " is not a boolean variable!"
stop 1
Expand Down Expand Up @@ -3683,4 +3747,137 @@ subroutine cvmix_kpp_compute_ER_depth( z_inter,Nsq,OBL_depth, &
end subroutine cvmix_kpp_compute_ER_depth



!BOP

! !DESCRIPTION:
! Use Entrainment Rule, BEdE_ER, To Find Entrainment Flux and Depth
! for each assumed OBL_depth = cell centers,
! until the boundary layer depth, ERdepth > 0 kER_depth are determined, OR
! if there is no viable solution ERdepth = -1 , kER_depth=-1

subroutine cvmix_kpp_compute_ER_depth( z_inter,Nsq,OBL_depth, &
uStar,Bsfc_ns,surfBuoy,StokesXI,BEdE_ER,ERdepth, &
CVMix_kpp_params_user)

! !INPUT PARAMETERS:
real(cvmix_r8), dimension(:), intent(in) :: &
z_inter, & ! Interface heights <= 0 [m]
Nsq ! Column of Buoyancy Gradients at interfaces
real(cvmix_r8), dimension(:), intent(in) :: &
OBL_depth, & ! Array of assumed OBL depths >0 at cell centers [m]
surfBuoy, & ! surface Buoyancy flux surface to OBL_depth
StokesXI, & ! Stokes similarity parameter given OBL_depth
BEdE_ER ! Parameterized Entrainment Rule given OBL_depth
real(cvmix_r8), intent(in) :: uStar ! surface friction velocity [m s-1]
real(cvmix_r8), intent(in) :: Bsfc_ns ! non-solar surface buoyancy flux boundary condition

type(cvmix_kpp_params_type), optional, target, intent(in) :: CVmix_kpp_params_user

! !OUTPUT PARAMETERS:
real(cvmix_r8), intent(out) :: ERdepth ! ER Boundary Layer Depth [m]

!EOP
!BOC

! Local variables
integer :: iz, nlev , kbl , kinv
real(cvmix_r8), dimension(size(OBL_depth)+1) :: zdepth, BEdE ! surface then cell-centers<0
real(cvmix_r8), dimension(size(z_inter)+1) :: sigma, Bflux ! interface values
real(cvmix_r8) :: ws_i, Cemp_Rs, Gsig_i, Fdelrho, Bnonlocal, sigE, maxNsq, BEnt
real(kind=cvmix_r8), dimension(4) :: coeffs
type(cvmix_kpp_params_type), pointer :: CVmix_kpp_params_in

CVmix_kpp_params_in => CVmix_kpp_params_saved
if (present(CVmix_kpp_params_user)) then
CVmix_kpp_params_in => CVmix_kpp_params_user
end if

nlev = size(OBL_depth)
Cemp_Rs = 4.7_cvmix_r8
Fdelrho = cvmix_one
maxNsq = 0.0
do kbl = 2, nlev+1
if ( Nsq(kbl) > maxNsq ) then
kinv = kbl
maxNsq = Nsq(kbl)
endif
enddo

! Set default output values (no solution)
ERdepth = -cvmix_one
! Set surface values
zdepth(1) = cvmix_zero
BEdE(1) = cvmix_zero
sigma(:) = cvmix_zero
Bflux(1) = Bsfc_ns ! non-solar surface buoyancy boundary condition for all kbl
! Set OBL_depth(1)=top cell center values
zdepth(2) = -OBL_depth(1)
BEdE(2) = cvmix_zero

do kbl = 2, nlev
zdepth(kbl+1)= -OBL_depth(kbl)
BEdE(kbl+1) = cvmix_zero

sigma(kbl+1) = cvmix_one
Bflux(kbl+1) = cvmix_zero
sigma(kbl+2) = -z_inter(kbl+1)/OBL_depth(kbl) ! > 1.0
Bflux(kbl+2) = cvmix_zero
Bnonlocal = 0.5_cvmix_r8 * Cemp_Rs * (abs(surfBuoy(kbl)) - surfBuoy(kbl))

do iz = kbl,1,-1
if (iz .gt. 1) then
sigma(iz) = -z_inter(iz)/OBL_depth(kbl) ! < 1.0
call cvmix_kpp_compute_turbulent_scales(sigma(iz), OBL_depth(kbl), surfBuoy(kbl), uStar, StokesXI(kbl), & !0d
w_s=ws_i , CVmix_kpp_params_user=CVmix_kpp_params_user)
Gsig_i = cvmix_kpp_composite_shape(sigma(iz))
Bflux(iz) = Gsig_i * (OBL_depth(kbl) * ws_i * Nsq(iz) - Bnonlocal)
endif

! find the peak
if ( (Bflux(iz+1) .gt. Bflux(iz+2)) .and. (Bflux(iz+1) .ge. Bflux(iz)) .and. &
(Bflux(iz+1) .gt. cvmix_zero) ) then
! Find sigE (the root of the derivative of the quadratic polynomial
! interpolating (sigma(i), Bflux(i)) for i in [iz, iz+1, iz+2])
! Also find BEnt (value of quadratic at sigE)
! call cvmix_math_poly_interp(coeffs, sigma(iz:iz+2), Bflux(iz:iz+2))
! comment by Aakash: the above is the original line,
! it gives me errors so I changed it to below call. not sure if it is correct
Copy link
Contributor

@mnlevy1981 mnlevy1981 Jan 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What error did you get? These two calls to cvmix_math_poly_interp are not equivalent, and will change answers regardless of the value of ML_diffusivity:

call cvmix_math_poly_interp(coeffs, sigma(iz:iz+2), Bflux(iz:iz+2))

fits a quadratic such that P(sigma(i)) = Bflux(i) for i = iz, iz+1, iz+2 (standard quadratic interpolation)

call cvmix_math_poly_interp(coeffs, CVMIX_MATH_INTERP_QUAD,           &
                            sigma(iz:iz+1), Bflux(iz:iz+1),           &
                            sigma(iz+2), Bflux(iz+2))

fits a quadratic such that P(sigma(i)) = Bflux(i) for i = iz, iz+1, and P'(sigma(iz)) = (Bflux(iz) - Bflux(iz+2)) / (sigma(iz) - sigma(iz+2)). If we were okay with this interpolation change, you'd still need to change your arguments to

call cvmix_math_poly_interp(coeffs, CVMIX_MATH_INTERP_QUAD,           &
                            sigma(iz+1:iz+2), Bflux(iz+1:iz+2),           &
                            sigma(iz), Bflux(iz))

And this would interpolate at iz+1 and iz+2 while having the derivative at iz+1 match the finite difference computed between iz and iz+1...

Anyway, if you could print out the three sigma values and three Bflux values that cause problems in your runs then I'll look into what's going wrong in the interpolation routine we want to use.

call cvmix_math_poly_interp(coeffs, CVMIX_MATH_INTERP_QUAD, &
sigma(iz:iz+1), Bflux(iz:iz+1), &
sigma(iz+2), Bflux(iz+2))
! coeffs(3) = 0 => sigma(iz), sigma(iz+1), and sigma(iz+2) are not unique values
! so there interpolation returned a linear equation. In this case we select
! (sigma(iz+1), Bflux(iz+1)) as the maximum.
if (coeffs(3) .eq. cvmix_zero) then
sigE = sigma(iz+1)
Bent = Bflux(iz+1)
else
sigE = -0.5_cvmix_r8 * (coeffs(2) / coeffs(3))
Bent = cvmix_math_evaluate_cubic(coeffs, sigE)
endif
Fdelrho = cvmix_one
BEdE(kbl+1) = Fdelrho*BEnt*sigE*OBL_depth(kbl)
exit ! No exit leaves BEdE(kbl+1) = cvmix_zero
endif
enddo

if ( (BEdE(kbl+1) > BEdE_ER(kbl)) .and. (Bsfc_ns < cvmix_zero) ) then
call cvmix_math_poly_interp(coeffs, CVmix_kpp_params_in%interp_type, &
zdepth(kbl:kbl+1) , BEdE(kbl:kbl+1), &
zdepth(kbl-1) , BEdE(kbl-1) ) ! surface values if kbl=2;
coeffs(1) = coeffs(1) - BEdE_ER(kbl)
ERdepth = -cvmix_math_cubic_root_find(coeffs, 0.5_cvmix_r8 * &
(zdepth(kbl)+zdepth(kbl+1)))

exit
endif

enddo

!EOC

end subroutine cvmix_kpp_compute_ER_depth

Copy link
Contributor

@mnlevy1981 mnlevy1981 Jan 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Your merge of the latest master introduced a second copy of cvmix_kpp_compute_ER_depth(). Did you change this function in a previous commit? I'd suggest deleting the first instance of this function, in case there are differences in this second copy...

Copy link
Author

@aakashsane aakashsane Jan 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't introduced any changes to this function. I think I copied it in error!


end module cvmix_kpp
Loading