added ML_diffusivity shape function #111

aakashsane · 2025-12-17T07:00:02Z

Adds option to activate shape function obtained from machine learning (genetic programming and least squares fitting).
The relevant flag is 'CVmix_kpp_params_in%ML_diffusivity', when set to True uses the mixing coefficients.
The mixing coefficients are obtained using equations similar to the ones in the GFDL's ePBL scheme.

The corresponding article preprint is: Sane et al. (2025) https://doi.org/10.31219/osf.io/uab7v_v2

mnlevy1981 · 2026-01-29T23:55:18Z

src/shared/cvmix_kpp.F90

+
+!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
+        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
+


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...

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

mnlevy1981 · 2026-01-30T00:33:04Z

src/shared/cvmix_kpp.F90

-        call cvmix_math_poly_interp(coeffs, sigma(iz:iz+2), Bflux(iz:iz+2))
+        ! 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


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.

mnlevy1981 · 2026-02-06T17:40:05Z

src/shared/cvmix_kpp.F90

+        case ('ML_Diffusivity_Shape') ! ML_Diffusivity shape function
+          call cvmix_put_kpp('MatchTechnique', ML_DIFFUSIVITY_SHAPE, CVmix_kpp_params_user)


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

mnlevy1981

Summarizing a productive meeting with @aakashsane and @YakelynRJ

The new ML_DIFFUSIVITY_SHAPE option needs to be documented in https://github.com/CVMix/CVMix-src/blob/master/src/shared/cvmix_kpp.F90#L158-L173 It should definitely be defined in the line after CVMIX_KPP_PARABOLIC_NONLOCAL instead of at the bottom of the parameter block, and could maybe be renamed CVMIX_KPP_ML_SHAPE?
Use MatchTechnique == ML_DIFFUSIVITY_SHAPE instead of introducing a new ML_diffusivity logical flag

There are a few in-line comments about moving code into the select case block that ends at line 1225 on this branch. This will allow us to do something like

if (CVmix_kpp_params_in%MatchTechnique == ML_DIFFUSIVITY_SHAPE) then ! ML_diffusivity
  if (sigma(kw) .le. sigma_max) then
    MshapeAtS = cvmix_math_evaluate_cubic(MLshape, sigma(kw))
    TshapeAtS = cvmix_math_evaluate_cubic(MLshape, sigma(kw))
    SshapeAtS = cvmix_math_evaluate_cubic(MLshape, sigma(kw))
  else
    MshapeAtS = cvmix_math_evaluate_cubic(MLshape2, sigma(kw))
    TshapeAtS = cvmix_math_evaluate_cubic(MLshape2, sigma(kw))
    SshapeAtS = cvmix_math_evaluate_cubic(MLshape2, sigma(kw))
  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

For (3), I could see using ways to reduce the logical checks, but for starters this should be good.

Also, we talked about doing some off-line testing to make sure the parameters passed to cvmix_evaluate_cubic() match the polynomials computed in

          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

and

          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

mnlevy1981 · 2026-02-09T20:46:20Z

src/shared/cvmix_kpp.F90

+          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


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.

mnlevy1981 · 2026-02-09T20:48:04Z

src/shared/cvmix_kpp.F90

+          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


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?

Machine learning based diffusivity changes, only to the shape function

7b8f824

aakashsane force-pushed the ML_diffusivity_3 branch from b213d2c to 7b8f824 Compare January 28, 2026 05:34

aakashsane added 2 commits January 28, 2026 00:41

Removed trailing white spaces

cf4c5b9

Merge branch 'master' into ML_diffusivity_3

96a02c6

mnlevy1981 reviewed Jan 29, 2026

View reviewed changes

Clean up following merge of master

aa778fb

mnlevy1981 requested changes Jan 30, 2026

View reviewed changes

mnlevy1981 reviewed Feb 6, 2026

View reviewed changes

mnlevy1981 requested changes Feb 9, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

added ML_diffusivity shape function #111

added ML_diffusivity shape function #111

Uh oh!

aakashsane commented Dec 17, 2025 •

edited

Loading

Uh oh!

mnlevy1981 Jan 29, 2026

Uh oh!

aakashsane Jan 31, 2026

Uh oh!

mnlevy1981 Jan 30, 2026

Uh oh!

mnlevy1981 Feb 6, 2026

Uh oh!

mnlevy1981 left a comment

Uh oh!

mnlevy1981 Feb 9, 2026

Uh oh!

mnlevy1981 Feb 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		case ('ML_Diffusivity_Shape') ! ML_Diffusivity shape function
		call cvmix_put_kpp('MatchTechnique', ML_DIFFUSIVITY_SHAPE, CVmix_kpp_params_user)

		g_sigma = (2.0sigma(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

added ML_diffusivity shape function #111

Are you sure you want to change the base?

added ML_diffusivity shape function #111

Uh oh!

Conversation

aakashsane commented Dec 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mnlevy1981 Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

aakashsane Jan 31, 2026

Choose a reason for hiding this comment

Uh oh!

mnlevy1981 Jan 30, 2026

Choose a reason for hiding this comment

Uh oh!

mnlevy1981 Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

mnlevy1981 left a comment

Choose a reason for hiding this comment

Uh oh!

mnlevy1981 Feb 9, 2026

Choose a reason for hiding this comment

Uh oh!

mnlevy1981 Feb 9, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

aakashsane commented Dec 17, 2025 •

edited

Loading