MOM_mixed_layer_restrat.F90

!> \brief Parameterization of mixed layer restratification by unresolved mixed-layer eddies.
module mom_mixed_layer_restrat

! This file is part of MOM6. See LICENSE.md for the license.

use mom_debugging,     only : hchksum
use mom_diag_mediator, only : post_data, query_averaging_enabled, diag_ctrl
use mom_diag_mediator, only : register_diag_field, safe_alloc_ptr, time_type
use mom_diag_mediator, only : diag_update_remap_grids
use mom_domains,       only : pass_var, to_west, to_south, omit_corners
use mom_error_handler, only : mom_error, fatal, warning
use mom_file_parser,   only : get_param, log_version, param_file_type
use mom_forcing_type,  only : mech_forcing
use mom_grid,          only : ocean_grid_type
use mom_hor_index,     only : hor_index_type
use mom_io,            only : vardesc, var_desc
use mom_lateral_mixing_coeffs, only : varmix_cs
use mom_restart,       only : register_restart_field, query_initialized, mom_restart_cs
use mom_unit_scaling,  only : unit_scale_type
use mom_variables,     only : thermo_var_ptrs
use mom_verticalgrid,  only : verticalgrid_type
use mom_eos,           only : calculate_density, eos_domain

implicit none ; private

#include <MOM_memory.h>

public mixedlayer_restrat
public mixedlayer_restrat_init
public mixedlayer_restrat_register_restarts

! A note on unit descriptions in comments: MOM6 uses units that can be rescaled for dimensional
! consistency testing. These are noted in comments with units like Z, H, L, and T, along with
! their mks counterparts with notation like "a velocity [Z T-1 ~> m s-1]".  If the units
! vary with the Boussinesq approximation, the Boussinesq variant is given first.

!> Control structure for mom_mixed_layer_restrat
type, public :: mixedlayer_restrat_cs ; private
  real    :: ml_restrat_coef       !< A non-dimensional factor by which the instability is enhanced
                                   !! over what would be predicted based on the resolved gradients
                                   !! [nondim].  This increases with grid spacing^2, up to something
                                   !! of order 500.
  real    :: ml_restrat_coef2      !< As for ml_restrat_coef but using the slow filtered MLD [nondim].
  real    :: front_length          !< If non-zero, is the frontal-length scale [L ~> m] used to calculate the
                                   !! upscaling of buoyancy gradients that is otherwise represented
                                   !! by the parameter FOX_KEMPER_ML_RESTRAT_COEF. If MLE_FRONT_LENGTH is
                                   !! non-zero, it is recommended to set FOX_KEMPER_ML_RESTRAT_COEF=1.0.
  logical :: mle_use_pbl_mld       !< If true, use the MLD provided by the PBL parameterization.
                                   !! if false, MLE will calculate a MLD based on a density difference
                                   !! based on the parameter MLE_DENSITY_DIFF.
  real    :: mle_mld_decay_time    !< Time-scale to use in a running-mean when MLD is retreating [T ~> s].
  real    :: mle_mld_decay_time2   !< Time-scale to use in a running-mean when filtered MLD is retreating [T ~> s].
  real    :: mle_density_diff      !< Density difference used in detecting mixed-layer depth [R ~> kg m-3].
  real    :: mle_tail_dh           !< Fraction by which to extend the mixed-layer restratification
                                   !! depth used for a smoother stream function at the base of
                                   !! the mixed-layer [nondim].
  real    :: mle_mld_stretch       !< A scaling coefficient for stretching/shrinking the MLD used in
                                   !! the MLE scheme [nondim]. This simply multiplies MLD wherever used.
  logical :: debug = .false.       !< If true, calculate checksums of fields for debugging.
  type(diag_ctrl), pointer :: diag !< A structure that is used to regulate the
                                   !! timing of diagnostic output.

  real, dimension(:,:), pointer :: &
         mld_filtered => null(), &   !< Time-filtered MLD [H ~> m or kg m-2]
         mld_filtered_slow => null() !< Slower time-filtered MLD [H ~> m or kg m-2]

  !>@{
  !! Diagnostic identifier
  integer :: id_urestrat_time = -1
  integer :: id_vrestrat_time = -1
  integer :: id_uhml = -1
  integer :: id_vhml = -1
  integer :: id_mld  = -1
  integer :: id_rml  = -1
  integer :: id_udml = -1
  integer :: id_vdml = -1
  integer :: id_uml  = -1
  integer :: id_vml  = -1
  !>@}

end type mixedlayer_restrat_cs

character(len=40)  :: mdl = "MOM_mixed_layer_restrat" !< This module's name.

contains

!> Driver for the mixed-layer restratification parameterization.
!! The code branches between two different implementations depending
!! on whether the bulk-mixed layer or a general coordinate are in use.
subroutine mixedlayer_restrat(h, uhtr, vhtr, tv, forces, dt, MLD, VarMix, G, GV, US, CS)
  type(ocean_grid_type),                     intent(inout) :: g      !< Ocean grid structure
  type(verticalgrid_type),                   intent(in)    :: gv     !< Ocean vertical grid structure
  type(unit_scale_type),                     intent(in)    :: us     !< A dimensional unit scaling type
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)),  intent(inout) :: h      !< Layer thickness [H ~> m or kg m-2]
  real, dimension(SZIB_(G),SZJ_(G),SZK_(G)), intent(inout) :: uhtr   !< Accumulated zonal mass flux
                                                                     !! [H L2 ~> m3 or kg]
  real, dimension(SZI_(G),SZJB_(G),SZK_(G)), intent(inout) :: vhtr   !< Accumulated meridional mass flux
                                                                     !! [H L2 ~> m3 or kg]
  type(thermo_var_ptrs),                     intent(in)    :: tv     !< Thermodynamic variables structure
  type(mech_forcing),                        intent(in)    :: forces !< A structure with the driving mechanical forces
  real,                                      intent(in)    :: dt     !< Time increment [T ~> s]
  real, dimension(:,:),                      pointer       :: mld    !< Mixed layer depth provided by the
                                                                     !! PBL scheme [Z ~> m]
  type(varmix_cs),                           pointer       :: varmix !< Container for derived fields
  type(mixedlayer_restrat_cs),               pointer       :: cs     !< Module control structure

  if (.not. associated(cs)) call mom_error(fatal, "MOM_mixedlayer_restrat: "// &
         "Module must be initialized before it is used.")

  if (gv%nkml>0) then
    call mixedlayer_restrat_bml(h, uhtr, vhtr, tv, forces, dt, g, gv, us, cs)
  else
    call mixedlayer_restrat_general(h, uhtr, vhtr, tv, forces, dt, mld, varmix, g, gv, us, cs)
  endif

end subroutine mixedlayer_restrat

!> Calculates a restratifying flow in the mixed layer.
subroutine mixedlayer_restrat_general(h, uhtr, vhtr, tv, forces, dt, MLD_in, VarMix, G, GV, US, CS)
  ! Arguments
  type(ocean_grid_type),                     intent(inout) :: G      !< Ocean grid structure
  type(verticalGrid_type),                   intent(in)    :: GV     !< Ocean vertical grid structure
  type(unit_scale_type),                     intent(in)    :: US     !< A dimensional unit scaling type
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)),  intent(inout) :: h      !< Layer thickness [H ~> m or kg m-2]
  real, dimension(SZIB_(G),SZJ_(G),SZK_(G)), intent(inout) :: uhtr   !< Accumulated zonal mass flux
                                                                     !!   [H L2 ~> m3 or kg]
  real, dimension(SZI_(G),SZJB_(G),SZK_(G)), intent(inout) :: vhtr   !< Accumulated meridional mass flux
                                                                     !!   [H L2 ~> m3 or kg]
  type(thermo_var_ptrs),                     intent(in)    :: tv     !< Thermodynamic variables structure
  type(mech_forcing),                        intent(in)    :: forces !< A structure with the driving mechanical forces
  real,                                      intent(in)    :: dt     !< Time increment [T ~> s]
  real, dimension(:,:),                      pointer       :: MLD_in !< Mixed layer depth provided by the
                                                                     !! PBL scheme [Z ~> m] (not H)
  type(VarMix_CS),                           pointer       :: VarMix !< Container for derived fields
  type(mixedlayer_restrat_CS),               pointer       :: CS     !< Module control structure
  ! Local variables
  real :: uhml(SZIB_(G),SZJ_(G),SZK_(G)) ! zonal mixed layer transport [H L2 T-1 ~> m3 s-1 or kg s-1]
  real :: vhml(SZI_(G),SZJB_(G),SZK_(G)) ! merid mixed layer transport [H L2 T-1 ~> m3 s-1 or kg s-1]
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)) :: &
    h_avail               ! The volume available for diffusion out of each face of each
                          ! sublayer of the mixed layer, divided by dt [H L2 T-1 ~> m3 s-1 or kg s-1].
  real, dimension(SZI_(G),SZJ_(G)) :: &
    MLD_fast, &           ! Mixed layer depth actually used in MLE restratification parameterization [H ~> m or kg m-2]
    htot_fast, &          ! The sum of the thicknesses of layers in the mixed layer [H ~> m or kg m-2]
    Rml_av_fast, &        ! g_Rho0 times the average mixed layer density [L2 Z-1 T-2 ~> m s-2]
    MLD_slow, &           ! Mixed layer depth actually used in MLE restratification parameterization [H ~> m or kg m-2]
    htot_slow, &          ! The sum of the thicknesses of layers in the mixed layer [H ~> m or kg m-2]
    Rml_av_slow           ! g_Rho0 times the average mixed layer density [L2 Z-1 T-2 ~> m s-2]
  real :: g_Rho0          ! G_Earth/Rho0 [L2 Z-1 T-2 R-1 ~> m4 s-2 kg-1]
  real :: rho_ml(SZI_(G)) ! Potential density relative to the surface [R ~> kg m-3]
  real :: p0(SZI_(G))     ! A pressure of 0 [R L2 T-2 ~> Pa]

  real :: h_vel           ! htot interpolated onto velocity points [Z ~> m] (not H).
  real :: absf            ! absolute value of f, interpolated to velocity points [T-1 ~> s-1]
  real :: u_star          ! surface friction velocity, interpolated to velocity points [Z T-1 ~> m s-1].
  real :: mom_mixrate     ! rate at which momentum is homogenized within mixed layer [T-1 ~> s-1]
  real :: timescale       ! mixing growth timescale [T ~> s]
  real :: h_neglect       ! tiny thickness usually lost in roundoff so can be neglected [H ~> m or kg m-2]
  real :: dz_neglect      ! A tiny thickness that is usually lost in roundoff so can be neglected [Z ~> m]
  real :: I4dt            ! 1/(4 dt) [T-1 ~> s-1]
  real :: Ihtot,Ihtot_slow! Inverses of the total mixed layer thickness [H-1 ~> m-1 or m2 kg-1]
  real :: a(SZK_(G))      ! A non-dimensional value relating the overall flux
                          ! magnitudes (uDml & vDml) to the realized flux in a
                          ! layer.  The vertical sum of a() through the pieces of
                          ! the mixed layer must be 0.
  real :: b(SZK_(G))      ! As for a(k) but for the slow-filtered MLD
  real :: uDml(SZIB_(G))  ! The zonal and meridional volume fluxes in the upper
  real :: vDml(SZI_(G))   ! half of the mixed layer [H L2 T-1 ~> m3 s-1 or kg s-1].
  real :: uDml_slow(SZIB_(G))  ! The zonal and meridional volume fluxes in the upper
  real :: vDml_slow(SZI_(G))   ! half of the mixed layer [H L2 T-1 ~> m3 s-1 or kg s-1].
  real :: utimescale_diag(SZIB_(G),SZJ_(G)) ! restratification timescales in the zonal and
  real :: vtimescale_diag(SZI_(G),SZJB_(G)) ! meridional directions [T ~> s], stored in 2-D arrays
                                            ! for diagnostic purposes.
  real :: uDml_diag(SZIB_(G),SZJ_(G)), vDml_diag(SZI_(G),SZJB_(G))
  real, dimension(SZI_(G)) :: rhoSurf, deltaRhoAtKm1, deltaRhoAtK ! Densities [R ~> kg m-3]
  real, dimension(SZI_(G)) :: dK, dKm1 ! Depths of layer centers [H ~> m or kg m-2].
  real, dimension(SZI_(G)) :: pRef_MLD ! A reference pressure for calculating the mixed layer
                                       ! densities [R L2 T-2 ~> Pa].
  real, dimension(SZI_(G)) :: rhoAtK, rho1, d1, pRef_N2 ! Used for N2
  real :: aFac, bFac ! Nondimensional ratios [nondim]
  real :: ddRho    ! A density difference [R ~> kg m-3]
  real :: hAtVel, zpa, zpb, dh, res_scaling_fac
  real :: I_LFront ! The inverse of the frontal length scale [L-1 ~> m-1]
  logical :: line_is_empty, keep_going, res_upscale
  integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state
  integer :: i, j, k, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz

  real :: PSI, PSI1, z, BOTTOP, XP, DD ! For the following statement functions
  ! Stream function as a function of non-dimensional position within mixed-layer (F77 statement function)
  !PSI1(z) = max(0., (1. - (2.*z+1.)**2 ) )
  psi1(z) = max(0., (1. - (2.*z+1.)**2 ) * (1. + (5./21.)*(2.*z+1.)**2) )
  bottop(z) = 0.5*(1.-sign(1.,z+0.5)) ! =0 for z>-0.5, =1 for z<-0.5
  xp(z) = max(0., min(1., (-z-0.5)*2./(1.+2.*cs%MLE_tail_dh) ) )
  dd(z) = (1.-3.*(xp(z)**2)+2.*(xp(z)**3))**(1.+2.*cs%MLE_tail_dh)
  psi(z) = max( psi1(z), dd(z)*bottop(z) ) ! Combines original PSI1 with tail

  is  = g%isc  ; ie  = g%iec  ; js  = g%jsc  ; je  = g%jec ; nz = g%ke
  isq = g%IscB ; ieq = g%IecB ; jsq = g%JscB ; jeq = g%JecB

  if (.not.associated(tv%eqn_of_state)) call mom_error(fatal, "MOM_mixedlayer_restrat: "// &
         "An equation of state must be used with this module.")
  if (.not.associated(varmix) .and. cs%front_length>0.) call mom_error(fatal, "MOM_mixedlayer_restrat: "// &
         "The resolution argument, Rd/dx, was not associated.")

  if (cs%MLE_density_diff > 0.) then ! We need to calculate a mixed layer depth, MLD.
    !! TODO: use derivatives and mid-MLD pressure. Currently this is sigma-0. -AJA
    pref_mld(:) = 0.
    eosdom(:) = eos_domain(g%HI, halo=1)
    do j = js-1, je+1
      dk(:) = 0.5 * h(:,j,1) ! Depth of center of surface layer
      call calculate_density(tv%T(:,j,1), tv%S(:,j,1), pref_mld, rhosurf, tv%eqn_of_state, eosdom)
      deltarhoatk(:) = 0.
      mld_fast(:,j) = 0.
      do k = 2, nz
        dkm1(:) = dk(:) ! Depth of center of layer K-1
        dk(:) = dk(:) + 0.5 * ( h(:,j,k) + h(:,j,k-1) ) ! Depth of center of layer K
        ! Mixed-layer depth, using sigma-0 (surface reference pressure)
        deltarhoatkm1(:) = deltarhoatk(:) ! Store value from previous iteration of K
        call calculate_density(tv%T(:,j,k), tv%S(:,j,k), pref_mld, deltarhoatk, tv%eqn_of_state, eosdom)
        do i = is-1,ie+1
          deltarhoatk(i) = deltarhoatk(i) - rhosurf(i) ! Density difference between layer K and surface
        enddo
        do i = is-1, ie+1
          ddrho = deltarhoatk(i) - deltarhoatkm1(i)
          if ((mld_fast(i,j)==0.) .and. (ddrho>0.) .and. &
              (deltarhoatkm1(i)<cs%MLE_density_diff) .and. (deltarhoatk(i)>=cs%MLE_density_diff)) then
            afac = ( cs%MLE_density_diff - deltarhoatkm1(i) ) / ddrho
            mld_fast(i,j) = dk(i) * afac + dkm1(i) * (1. - afac)
          endif
        enddo ! i-loop
      enddo ! k-loop
      do i = is-1, ie+1
        mld_fast(i,j) = cs%MLE_MLD_stretch * mld_fast(i,j)
        if ((mld_fast(i,j)==0.) .and. (deltarhoatk(i)<cs%MLE_density_diff)) &
          mld_fast(i,j) = dk(i) ! Assume mixing to the bottom
      enddo
    enddo ! j-loop
  elseif (cs%MLE_use_PBL_MLD) then
    if (.not. associated(mld_in)) call mom_error(fatal, "MOM_mixedlayer_restrat: "// &
         "Argument MLD_in was not associated!")
    do j = js-1, je+1 ; do i = is-1, ie+1
      mld_fast(i,j) = (cs%MLE_MLD_stretch * gv%Z_to_H) * mld_in(i,j)
    enddo ; enddo
  else
    call mom_error(fatal, "MOM_mixedlayer_restrat: "// &
         "No MLD to use for MLE parameterization.")
  endif

  ! Apply time filter (to remove diurnal cycle)
  if (cs%MLE_MLD_decay_time>0.) then
    if (cs%debug) then
      call hchksum(cs%MLD_filtered, 'mixed_layer_restrat: MLD_filtered', g%HI, haloshift=1, scale=gv%H_to_m)
      call hchksum(mld_in, 'mixed_layer_restrat: MLD in', g%HI, haloshift=1, scale=us%Z_to_m)
    endif
    afac = cs%MLE_MLD_decay_time / ( dt + cs%MLE_MLD_decay_time )
    bfac = dt / ( dt + cs%MLE_MLD_decay_time )
    do j = js-1, je+1 ; do i = is-1, ie+1
      ! Expression bFac*MLD_fast(i,j) + aFac*CS%MLD_filtered(i,j) is the time-filtered
      ! (running mean) of MLD. The max() allows the "running mean" to be reset
      ! instantly to a deeper MLD.
      cs%MLD_filtered(i,j) = max( mld_fast(i,j), bfac*mld_fast(i,j) + afac*cs%MLD_filtered(i,j) )
      mld_fast(i,j) = cs%MLD_filtered(i,j)
    enddo ; enddo
  endif

  ! Apply slower time filter (to remove seasonal cycle) on already filtered MLD_fast
  if (cs%MLE_MLD_decay_time2>0.) then
    if (cs%debug) then
      call hchksum(cs%MLD_filtered_slow,'mixed_layer_restrat: MLD_filtered_slow',g%HI,haloshift=1,scale=gv%H_to_m)
      call hchksum(mld_fast,'mixed_layer_restrat: MLD fast',g%HI,haloshift=1,scale=gv%H_to_m)
    endif
    afac = cs%MLE_MLD_decay_time2 / ( dt + cs%MLE_MLD_decay_time2 )
    bfac = dt / ( dt + cs%MLE_MLD_decay_time2 )
    do j = js-1, je+1 ; do i = is-1, ie+1
      ! Expression bFac*MLD_fast(i,j) + aFac*CS%MLD_filtered(i,j) is the time-filtered
      ! (running mean) of MLD. The max() allows the "running mean" to be reset
      ! instantly to a deeper MLD.
      cs%MLD_filtered_slow(i,j) = max( mld_fast(i,j), bfac*mld_fast(i,j) + afac*cs%MLD_filtered_slow(i,j) )
      mld_slow(i,j) = cs%MLD_filtered_slow(i,j)
    enddo ; enddo
  else
    do j = js-1, je+1 ; do i = is-1, ie+1
      mld_slow(i,j) = mld_fast(i,j)
    enddo ; enddo
  endif

  udml(:) = 0.0 ; vdml(:) = 0.0
  udml_slow(:) = 0.0 ; vdml_slow(:) = 0.0
  i4dt = 0.25 / dt
  g_rho0 = gv%g_Earth / gv%Rho0
  h_neglect = gv%H_subroundoff
  dz_neglect = gv%H_subroundoff*gv%H_to_Z
  if (cs%front_length>0.) then
    res_upscale = .true.
    i_lfront = 1. / cs%front_length
  else
    res_upscale = .false.
  endif

  p0(:) = 0.0
  eosdom(:) = eos_domain(g%HI, halo=1)
!$OMP parallel default(none) shared(is,ie,js,je,G,GV,US,htot_fast,Rml_av_fast,tv,p0,h,h_avail,&
!$OMP                               h_neglect,g_Rho0,I4dt,CS,uhml,uhtr,dt,vhml,vhtr,EOSdom,   &
!$OMP                               utimescale_diag,vtimescale_diag,forces,dz_neglect, &
!$OMP                               htot_slow,MLD_slow,Rml_av_slow,VarMix,I_LFront,    &
!$OMP                               res_upscale, nz,MLD_fast,uDml_diag,vDml_diag)      &
!$OMP                       private(rho_ml,h_vel,u_star,absf,mom_mixrate,timescale,    &
!$OMP                               line_is_empty, keep_going,res_scaling_fac,         &
!$OMP                               a,IhTot,b,Ihtot_slow,zpb,hAtVel,zpa,dh)            &
!$OMP                       firstprivate(uDml,vDml,uDml_slow,vDml_slow)
!$OMP do
  do j=js-1,je+1
    do i=is-1,ie+1
      htot_fast(i,j) = 0.0 ; rml_av_fast(i,j) = 0.0
      htot_slow(i,j) = 0.0 ; rml_av_slow(i,j) = 0.0
    enddo
    keep_going = .true.
    do k=1,nz
      do i=is-1,ie+1
        h_avail(i,j,k) = max(i4dt*g%areaT(i,j)*(h(i,j,k)-gv%Angstrom_H),0.0)
      enddo
      if (keep_going) then
        call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p0, rho_ml(:), tv%eqn_of_state, eosdom)
        line_is_empty = .true.
        do i=is-1,ie+1
          if (htot_fast(i,j) < mld_fast(i,j)) then
            dh = min( h(i,j,k), mld_fast(i,j)-htot_fast(i,j) )
            rml_av_fast(i,j) = rml_av_fast(i,j) + dh*rho_ml(i)
            htot_fast(i,j) = htot_fast(i,j) + dh
            line_is_empty = .false.
          endif
          if (htot_slow(i,j) < mld_slow(i,j)) then
            dh = min( h(i,j,k), mld_slow(i,j)-htot_slow(i,j) )
            rml_av_slow(i,j) = rml_av_slow(i,j) + dh*rho_ml(i)
            htot_slow(i,j) = htot_slow(i,j) + dh
            line_is_empty = .false.
          endif
        enddo
        if (line_is_empty) keep_going=.false.
      endif
    enddo

    do i=is-1,ie+1
      rml_av_fast(i,j) = -(g_rho0*rml_av_fast(i,j)) / (htot_fast(i,j) + h_neglect)
      rml_av_slow(i,j) = -(g_rho0*rml_av_slow(i,j)) / (htot_slow(i,j) + h_neglect)
    enddo
  enddo

  if (cs%debug) then
    call hchksum(h,'mixed_layer_restrat: h', g%HI, haloshift=1, scale=gv%H_to_m)
    call hchksum(forces%ustar,'mixed_layer_restrat: u*', g%HI, haloshift=1, scale=us%Z_to_m*us%s_to_T)
    call hchksum(mld_fast,'mixed_layer_restrat: MLD', g%HI, haloshift=1, scale=gv%H_to_m)
    call hchksum(rml_av_fast,'mixed_layer_restrat: rml', g%HI, haloshift=1, &
                 scale=us%m_to_Z*us%L_T_to_m_s**2)
  endif

! TO DO:
!   1. Mixing extends below the mixing layer to the mixed layer.  Find it!
!   2. Add exponential tail to stream-function?

!   U - Component
!$OMP do
  do j=js,je ; do i=is-1,ie
    u_star = 0.5*(forces%ustar(i,j) + forces%ustar(i+1,j))
    absf = 0.5*(abs(g%CoriolisBu(i,j-1)) + abs(g%CoriolisBu(i,j)))
    ! If needed, res_scaling_fac = min( ds, L_d ) / l_f
    if (res_upscale) res_scaling_fac = &
          ( sqrt( 0.5 * ( g%dxCu(i,j)**2 + g%dyCu(i,j)**2 ) ) * i_lfront ) &
          * min( 1., 0.5*( varmix%Rd_dx_h(i,j) + varmix%Rd_dx_h(i+1,j) ) )

    ! peak ML visc: u_star * 0.41 * (h_ml*u_star)/(absf*h_ml + 4.0*u_star)
    ! momentum mixing rate: pi^2*visc/h_ml^2
    ! 0.41 is the von Karmen constant, 9.8696 = pi^2.
    h_vel = 0.5*((htot_fast(i,j) + htot_fast(i+1,j)) + h_neglect) * gv%H_to_Z
    mom_mixrate = (0.41*9.8696)*u_star**2 / &
                  (absf*h_vel**2 + 4.0*(h_vel+dz_neglect)*u_star)
    timescale = 0.0625 * (absf + 2.0*mom_mixrate) / (absf**2 + mom_mixrate**2)
    timescale = timescale * cs%ml_restrat_coef
    if (res_upscale) timescale = timescale * res_scaling_fac
    udml(i) = timescale * g%mask2dCu(i,j)*g%dyCu(i,j)*g%IdxCu(i,j) * &
        (rml_av_fast(i+1,j)-rml_av_fast(i,j)) * (h_vel**2 * gv%Z_to_H)
    ! As above but using the slow filtered MLD
    h_vel = 0.5*((htot_slow(i,j) + htot_slow(i+1,j)) + h_neglect) * gv%H_to_Z
    mom_mixrate = (0.41*9.8696)*u_star**2 / &
                  (absf*h_vel**2 + 4.0*(h_vel+dz_neglect)*u_star)
    timescale = 0.0625 * (absf + 2.0*mom_mixrate) / (absf**2 + mom_mixrate**2)
    timescale = timescale * cs%ml_restrat_coef2
    if (res_upscale) timescale = timescale * res_scaling_fac
    udml_slow(i) = timescale * g%mask2dCu(i,j)*g%dyCu(i,j)*g%IdxCu(i,j) * &
        (rml_av_slow(i+1,j)-rml_av_slow(i,j)) * (h_vel**2 * gv%Z_to_H)

    if (udml(i) + udml_slow(i) == 0.) then
      do k=1,nz ; uhml(i,j,k) = 0.0 ; enddo
    else
      ihtot = 2.0 / ((htot_fast(i,j) + htot_fast(i+1,j)) + h_neglect)
      ihtot_slow = 2.0 / ((htot_slow(i,j) + htot_slow(i+1,j)) + h_neglect)
      zpa = 0.0 ; zpb = 0.0
      ! a(k) relates the sublayer transport to uDml with a linear profile.
      ! The sum of a(k) through the mixed layers must be 0.
      do k=1,nz
        hatvel = 0.5*(h(i,j,k) + h(i+1,j,k))
        a(k) = psi(zpa)                     ! Psi(z/MLD) for upper interface
        zpa = zpa - (hatvel * ihtot)        ! z/H for lower interface
        a(k) = a(k) - psi(zpa)              ! Transport profile
        ! Limit magnitude (uDml) if it would violate CFL
        if (a(k)*udml(i) > 0.0) then
          if (a(k)*udml(i) > h_avail(i,j,k)) udml(i) = h_avail(i,j,k) / a(k)
        elseif (a(k)*udml(i) < 0.0) then
          if (-a(k)*udml(i) > h_avail(i+1,j,k)) udml(i) = -h_avail(i+1,j,k) / a(k)
        endif
      enddo
      do k=1,nz
        ! Transport for slow-filtered MLD
        hatvel = 0.5*(h(i,j,k) + h(i+1,j,k))
        b(k) = psi(zpb)                     ! Psi(z/MLD) for upper interface
        zpb = zpb - (hatvel * ihtot_slow)   ! z/H for lower interface
        b(k) = b(k) - psi(zpb)              ! Transport profile
        ! Limit magnitude (uDml_slow) if it would violate CFL when added to uDml
        if (b(k)*udml_slow(i) > 0.0) then
          if (b(k)*udml_slow(i) > h_avail(i,j,k) - a(k)*udml(i)) &
             udml_slow(i) = max( 0., h_avail(i,j,k) - a(k)*udml(i) ) / b(k)
        elseif (b(k)*udml_slow(i) < 0.0) then
          if (-b(k)*udml_slow(i) > h_avail(i+1,j,k) + a(k)*udml(i)) &
             udml_slow(i) = -max( 0., h_avail(i+1,j,k) + a(k)*udml(i) ) / b(k)
        endif
      enddo
      do k=1,nz
        uhml(i,j,k) = a(k)*udml(i) + b(k)*udml_slow(i)
        uhtr(i,j,k) = uhtr(i,j,k) + uhml(i,j,k)*dt
      enddo
    endif

    utimescale_diag(i,j) = timescale
    udml_diag(i,j) = udml(i)
  enddo ; enddo

!  V- component
!$OMP do
  do j=js-1,je ; do i=is,ie
    u_star = 0.5*(forces%ustar(i,j) + forces%ustar(i,j+1))
    absf = 0.5*(abs(g%CoriolisBu(i-1,j)) + abs(g%CoriolisBu(i,j)))
    ! If needed, res_scaling_fac = min( ds, L_d ) / l_f
    if (res_upscale) res_scaling_fac = &
          ( sqrt( 0.5 * ( (g%dxCv(i,j))**2 + (g%dyCv(i,j))**2 ) ) * i_lfront ) &
          * min( 1., 0.5*( varmix%Rd_dx_h(i,j) + varmix%Rd_dx_h(i,j+1) ) )

    ! peak ML visc: u_star * 0.41 * (h_ml*u_star)/(absf*h_ml + 4.0*u_star)
    ! momentum mixing rate: pi^2*visc/h_ml^2
    ! 0.41 is the von Karmen constant, 9.8696 = pi^2.
    h_vel = 0.5*((htot_fast(i,j) + htot_fast(i,j+1)) + h_neglect) * gv%H_to_Z
    mom_mixrate = (0.41*9.8696)*u_star**2 / &
                  (absf*h_vel**2 + 4.0*(h_vel+dz_neglect)*u_star)
    timescale = 0.0625 * (absf + 2.0*mom_mixrate) / (absf**2 + mom_mixrate**2)
    timescale = timescale * cs%ml_restrat_coef
    if (res_upscale) timescale = timescale * res_scaling_fac
    vdml(i) = timescale * g%mask2dCv(i,j)*g%dxCv(i,j)*g%IdyCv(i,j) * &
        (rml_av_fast(i,j+1)-rml_av_fast(i,j)) * (h_vel**2 * gv%Z_to_H)
    ! As above but using the slow filtered MLD
    h_vel = 0.5*((htot_slow(i,j) + htot_slow(i,j+1)) + h_neglect) * gv%H_to_Z
    mom_mixrate = (0.41*9.8696)*u_star**2 / &
                  (absf*h_vel**2 + 4.0*(h_vel+dz_neglect)*u_star)
    timescale = 0.0625 * (absf + 2.0*mom_mixrate) / (absf**2 + mom_mixrate**2)
    timescale = timescale * cs%ml_restrat_coef2
    if (res_upscale) timescale = timescale * res_scaling_fac
    vdml_slow(i) = timescale * g%mask2dCv(i,j)*g%dxCv(i,j)*g%IdyCv(i,j) * &
        (rml_av_slow(i,j+1)-rml_av_slow(i,j)) * (h_vel**2 * gv%Z_to_H)

    if (vdml(i) + vdml_slow(i) == 0.) then
      do k=1,nz ; vhml(i,j,k) = 0.0 ; enddo
    else
      ihtot = 2.0 / ((htot_fast(i,j) + htot_fast(i,j+1)) + h_neglect)
      ihtot_slow = 2.0 / ((htot_slow(i,j) + htot_slow(i,j+1)) + h_neglect)
      zpa = 0.0 ; zpb = 0.0
      ! a(k) relates the sublayer transport to vDml with a linear profile.
      ! The sum of a(k) through the mixed layers must be 0.
      do k=1,nz
        hatvel = 0.5*(h(i,j,k) + h(i,j+1,k))
        a(k) = psi( zpa )                   ! Psi(z/MLD) for upper interface
        zpa = zpa - (hatvel * ihtot)        ! z/H for lower interface
        a(k) = a(k) - psi( zpa )            ! Transport profile
        ! Limit magnitude (vDml) if it would violate CFL
        if (a(k)*vdml(i) > 0.0) then
          if (a(k)*vdml(i) > h_avail(i,j,k)) vdml(i) = h_avail(i,j,k) / a(k)
        elseif (a(k)*vdml(i) < 0.0) then
          if (-a(k)*vdml(i) > h_avail(i,j+1,k)) vdml(i) = -h_avail(i,j+1,k) / a(k)
        endif
      enddo
      do k=1,nz
        ! Transport for slow-filtered MLD
        hatvel = 0.5*(h(i,j,k) + h(i,j+1,k))
        b(k) = psi(zpb)                     ! Psi(z/MLD) for upper interface
        zpb = zpb - (hatvel * ihtot_slow)   ! z/H for lower interface
        b(k) = b(k) - psi(zpb)              ! Transport profile
        ! Limit magnitude (vDml_slow) if it would violate CFL when added to vDml
        if (b(k)*vdml_slow(i) > 0.0) then
          if (b(k)*vdml_slow(i) > h_avail(i,j,k) - a(k)*vdml(i)) &
             vdml_slow(i) = max( 0., h_avail(i,j,k) - a(k)*vdml(i) ) / b(k)
        elseif (b(k)*vdml_slow(i) < 0.0) then
          if (-b(k)*vdml_slow(i) > h_avail(i,j+1,k) + a(k)*vdml(i)) &
             vdml_slow(i) = -max( 0., h_avail(i,j+1,k) + a(k)*vdml(i) ) / b(k)
        endif
      enddo
      do k=1,nz
        vhml(i,j,k) = a(k)*vdml(i) + b(k)*vdml_slow(i)
        vhtr(i,j,k) = vhtr(i,j,k) + vhml(i,j,k)*dt
      enddo
    endif

    vtimescale_diag(i,j) = timescale
    vdml_diag(i,j) = vdml(i)
  enddo ; enddo

!$OMP do
  do j=js,je ; do k=1,nz ; do i=is,ie
    h(i,j,k) = h(i,j,k) - dt*g%IareaT(i,j) * &
        ((uhml(i,j,k) - uhml(i-1,j,k)) + (vhml(i,j,k) - vhml(i,j-1,k)))
  enddo ; enddo ; enddo
!$OMP end parallel

  ! Whenever thickness changes let the diag manager know, target grids
  ! for vertical remapping may need to be regenerated.
  if (cs%id_uhml > 0 .or. cs%id_vhml > 0) &
    ! Remapped uhml and vhml require east/north halo updates of h
    call pass_var(h, g%domain, to_west+to_south+omit_corners, halo=1)
  call diag_update_remap_grids(cs%diag)

  ! Offer diagnostic fields for averaging.
  if (query_averaging_enabled(cs%diag)) then
    if (cs%id_urestrat_time > 0) call post_data(cs%id_urestrat_time, utimescale_diag, cs%diag)
    if (cs%id_vrestrat_time > 0) call post_data(cs%id_vrestrat_time, vtimescale_diag, cs%diag)
    if (cs%id_uhml          > 0) call post_data(cs%id_uhml, uhml, cs%diag)
    if (cs%id_vhml          > 0) call post_data(cs%id_vhml, vhml, cs%diag)
    if (cs%id_MLD           > 0) call post_data(cs%id_MLD, mld_fast, cs%diag)
    if (cs%id_Rml           > 0) call post_data(cs%id_Rml, rml_av_fast, cs%diag)
    if (cs%id_uDml          > 0) call post_data(cs%id_uDml, udml_diag, cs%diag)
    if (cs%id_vDml          > 0) call post_data(cs%id_vDml, vdml_diag, cs%diag)

    if (cs%id_uml > 0) then
      do j=js,je ; do i=is-1,ie
        h_vel = 0.5*((htot_fast(i,j) + htot_fast(i+1,j)) + h_neglect)
        udml_diag(i,j) = udml_diag(i,j) / (0.01*h_vel) * g%IdyCu(i,j) * (psi(0.)-psi(-.01))
      enddo ; enddo
      call post_data(cs%id_uml, udml_diag, cs%diag)
    endif
    if (cs%id_vml > 0) then
      do j=js-1,je ; do i=is,ie
        h_vel = 0.5*((htot_fast(i,j) + htot_fast(i,j+1)) + h_neglect)
        vdml_diag(i,j) = vdml_diag(i,j) / (0.01*h_vel) * g%IdxCv(i,j) * (psi(0.)-psi(-.01))
      enddo ; enddo
      call post_data(cs%id_vml, vdml_diag, cs%diag)
    endif
  endif
  ! Whenever thickness changes let the diag manager know, target grids
  ! for vertical remapping may need to be regenerated.
  ! This needs to happen after the H update and before the next post_data.
  call diag_update_remap_grids(cs%diag)

end subroutine mixedlayer_restrat_general


!> Calculates a restratifying flow assuming a 2-layer bulk mixed layer.
subroutine mixedlayer_restrat_bml(h, uhtr, vhtr, tv, forces, dt, G, GV, US, CS)
  type(ocean_grid_type),                     intent(in)    :: G      !< Ocean grid structure
  type(verticalGrid_type),                   intent(in)    :: GV     !< Ocean vertical grid structure
  type(unit_scale_type),                     intent(in)    :: US     !< A dimensional unit scaling type
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)),  intent(inout) :: h      !< Layer thickness [H ~> m or kg m-2]
  real, dimension(SZIB_(G),SZJ_(G),SZK_(G)), intent(inout) :: uhtr   !< Accumulated zonal mass flux
                                                                     !!   [H L2 ~> m3 or kg]
  real, dimension(SZI_(G),SZJB_(G),SZK_(G)), intent(inout) :: vhtr   !< Accumulated meridional mass flux
                                                                     !!   [H L2 ~> m3 or kg]
  type(thermo_var_ptrs),                     intent(in)    :: tv     !< Thermodynamic variables structure
  type(mech_forcing),                        intent(in)    :: forces !< A structure with the driving mechanical forces
  real,                                      intent(in)    :: dt     !< Time increment [T ~> s]
  type(mixedlayer_restrat_CS),               pointer       :: CS     !< Module control structure
  ! Local variables
  real :: uhml(SZIB_(G),SZJ_(G),SZK_(G)) ! zonal mixed layer transport [H L2 T-1 ~> m3 s-1 or kg s-1]
  real :: vhml(SZI_(G),SZJB_(G),SZK_(G)) ! merid mixed layer transport [H L2 T-1 ~> m3 s-1 or kg s-1]
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)) :: &
    h_avail               ! The volume available for diffusion out of each face of each
                          ! sublayer of the mixed layer, divided by dt [H L2 T-1 ~> m3 s-1 or kg s-1].
  real, dimension(SZI_(G),SZJ_(G)) :: &
    htot, &               ! The sum of the thicknesses of layers in the mixed layer [H ~> m or kg m-2]
    Rml_av                ! g_Rho0 times the average mixed layer density [L2 Z-1 T-2 ~> m s-2]
  real :: g_Rho0          ! G_Earth/Rho0 [L2 Z-1 T-2 R-1 ~> m4 s-2 kg-1]
  real :: Rho0(SZI_(G))   ! Potential density relative to the surface [R ~> kg m-3]
  real :: p0(SZI_(G))     ! A pressure of 0 [R L2 T-2 ~> Pa]

  real :: h_vel           ! htot interpolated onto velocity points [Z ~> m]. (The units are not H.)
  real :: absf            ! absolute value of f, interpolated to velocity points [T-1 ~> s-1]
  real :: u_star          ! surface friction velocity, interpolated to velocity points [Z T-1 ~> m s-1].
  real :: mom_mixrate     ! rate at which momentum is homogenized within mixed layer [T-1 ~> s-1]
  real :: timescale       ! mixing growth timescale [T ~> s]
  real :: h_neglect       ! tiny thickness usually lost in roundoff and can be neglected [H ~> m or kg m-2]
  real :: dz_neglect      ! tiny thickness that usually lost in roundoff and can be neglected [Z ~> m]
  real :: I4dt            ! 1/(4 dt) [T-1 ~> s-1]
  real :: I2htot          ! Twice the total mixed layer thickness at velocity points [H ~> m or kg m-2]
  real :: z_topx2         ! depth of the top of a layer at velocity points [H ~> m or kg m-2]
  real :: hx2             ! layer thickness at velocity points [H ~> m or kg m-2]
  real :: a(SZK_(G))      ! A non-dimensional value relating the overall flux
                          ! magnitudes (uDml & vDml) to the realized flux in a
                          ! layer.  The vertical sum of a() through the pieces of
                          ! the mixed layer must be 0.
  real :: uDml(SZIB_(G))  ! The zonal and meridional volume fluxes in the upper
  real :: vDml(SZI_(G))   ! half of the mixed layer [H L2 T-1 ~> m3 s-1 or kg s-1].
  real :: utimescale_diag(SZIB_(G),SZJ_(G)) ! The restratification timescales
  real :: vtimescale_diag(SZI_(G),SZJB_(G)) ! in the zonal and meridional
                                            ! directions [T ~> s], stored in 2-D
                                            ! arrays for diagnostic purposes.
  real :: uDml_diag(SZIB_(G),SZJ_(G)), vDml_diag(SZI_(G),SZJB_(G))
  logical :: use_EOS    ! If true, density is calculated from T & S using an equation of state.

  integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state
  integer :: i, j, k, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz, nkml
  is  = g%isc  ; ie  = g%iec  ; js  = g%jsc  ; je  = g%jec ; nz = g%ke
  isq = g%IscB ; ieq = g%IecB ; jsq = g%JscB ; jeq = g%JecB ; nkml = gv%nkml

  if (.not. associated(cs)) call mom_error(fatal, "MOM_mixedlayer_restrat: "// &
         "Module must be initialized before it is used.")
  if ((nkml<2) .or. (cs%ml_restrat_coef<=0.0)) return

  udml(:)    = 0.0 ; vdml(:) = 0.0
  i4dt       = 0.25 / dt
  g_rho0     = gv%g_Earth / gv%Rho0
  use_eos    = associated(tv%eqn_of_state)
  h_neglect  = gv%H_subroundoff
  dz_neglect = gv%H_subroundoff*gv%H_to_Z

  if (.not.use_eos) call mom_error(fatal, "MOM_mixedlayer_restrat: "// &
         "An equation of state must be used with this module.")

  ! Fix this later for nkml >= 3.

  p0(:) = 0.0
  eosdom(:) = eos_domain(g%HI, halo=1)
!$OMP parallel default(none) shared(is,ie,js,je,G,GV,US,htot,Rml_av,tv,p0,h,h_avail,EOSdom, &
!$OMP                               h_neglect,g_Rho0,I4dt,CS,uhml,uhtr,dt,vhml,vhtr,   &
!$OMP                               utimescale_diag,vtimescale_diag,forces,dz_neglect, &
!$OMP                               uDml_diag,vDml_diag,nkml)                          &
!$OMP                       private(Rho0,h_vel,u_star,absf,mom_mixrate,timescale,      &
!$OMP                               I2htot,z_topx2,hx2,a)                              &
!$OMP                       firstprivate(uDml,vDml)
!$OMP do
  do j=js-1,je+1
    do i=is-1,ie+1
      htot(i,j) = 0.0 ; rml_av(i,j) = 0.0
    enddo
    do k=1,nkml
      call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p0, rho0(:), tv%eqn_of_state, eosdom)
      do i=is-1,ie+1
        rml_av(i,j) = rml_av(i,j) + h(i,j,k)*rho0(i)
        htot(i,j) = htot(i,j) + h(i,j,k)
        h_avail(i,j,k) = max(i4dt*g%areaT(i,j)*(h(i,j,k)-gv%Angstrom_H),0.0)
      enddo
    enddo

    do i=is-1,ie+1
      rml_av(i,j) = (g_rho0*rml_av(i,j)) / (htot(i,j) + h_neglect)
    enddo
  enddo

! TO DO:
!   1. Mixing extends below the mixing layer to the mixed layer.  Find it!
!   2. Add exponential tail to stream-function?

!   U - Component
!$OMP do
  do j=js,je; do i=is-1,ie
    h_vel = 0.5*(htot(i,j) + htot(i+1,j)) * gv%H_to_Z

    u_star = 0.5*(forces%ustar(i,j) + forces%ustar(i+1,j))
    absf = 0.5*(abs(g%CoriolisBu(i,j-1)) + abs(g%CoriolisBu(i,j)))
    ! peak ML visc: u_star * 0.41 * (h_ml*u_star)/(absf*h_ml + 4.0*u_star)
    ! momentum mixing rate: pi^2*visc/h_ml^2
    ! 0.41 is the von Karmen constant, 9.8696 = pi^2.
    mom_mixrate = (0.41*9.8696)*u_star**2 / &
                  (absf*h_vel**2 + 4.0*(h_vel+dz_neglect)*u_star)
    timescale = 0.0625 * (absf + 2.0*mom_mixrate) / (absf**2 + mom_mixrate**2)

    timescale = timescale * cs%ml_restrat_coef
!      timescale = timescale*(2?)*(L_def/L_MLI) * min(EKE/MKE,1.0 + (G%dyCv(i,j)/L_def)**2)

    udml(i) = timescale * g%mask2dCu(i,j)*g%dyCu(i,j)*g%IdxCu(i,j) * &
        (rml_av(i+1,j)-rml_av(i,j)) * (h_vel**2 * gv%Z_to_H)

    if (udml(i) == 0) then
      do k=1,nkml ; uhml(i,j,k) = 0.0 ; enddo
    else
      i2htot = 1.0 / (htot(i,j) + htot(i+1,j) + h_neglect)
      z_topx2 = 0.0
      ! a(k) relates the sublayer transport to uDml with a linear profile.
      ! The sum of a(k) through the mixed layers must be 0.
      do k=1,nkml
        hx2 = (h(i,j,k) + h(i+1,j,k) + h_neglect)
        a(k) = (hx2 * i2htot) * (2.0 - 4.0*(z_topx2+0.5*hx2)*i2htot)
        z_topx2 = z_topx2 + hx2
        if (a(k)*udml(i) > 0.0) then
          if (a(k)*udml(i) > h_avail(i,j,k)) udml(i) = h_avail(i,j,k) / a(k)
        else
          if (-a(k)*udml(i) > h_avail(i+1,j,k)) udml(i) = -h_avail(i+1,j,k)/a(k)
        endif
      enddo
      do k=1,nkml
        uhml(i,j,k) = a(k)*udml(i)
        uhtr(i,j,k) = uhtr(i,j,k) + uhml(i,j,k)*dt
      enddo
    endif

    udml_diag(i,j) = udml(i)
    utimescale_diag(i,j) = timescale
  enddo; enddo

!  V- component
!$OMP do
  do j=js-1,je ; do i=is,ie
    h_vel = 0.5*(htot(i,j) + htot(i,j+1)) * gv%H_to_Z

    u_star = 0.5*(forces%ustar(i,j) + forces%ustar(i,j+1))
    absf = 0.5*(abs(g%CoriolisBu(i-1,j)) + abs(g%CoriolisBu(i,j)))
    ! peak ML visc: u_star * 0.41 * (h_ml*u_star)/(absf*h_ml + 4.0*u_star)
    ! momentum mixing rate: pi^2*visc/h_ml^2
    ! 0.41 is the von Karmen constant, 9.8696 = pi^2.
    mom_mixrate = (0.41*9.8696)*u_star**2 / &
                  (absf*h_vel**2 + 4.0*(h_vel+dz_neglect)*u_star)
    timescale = 0.0625 * (absf + 2.0*mom_mixrate) / (absf**2 + mom_mixrate**2)

    timescale = timescale * cs%ml_restrat_coef
!     timescale = timescale*(2?)*(L_def/L_MLI) * min(EKE/MKE,1.0 + (G%dyCv(i,j)/L_def)**2)

    vdml(i) = timescale * g%mask2dCv(i,j)*g%dxCv(i,j)*g%IdyCv(i,j) * &
        (rml_av(i,j+1)-rml_av(i,j)) * (h_vel**2 * gv%Z_to_H)
    if (vdml(i) == 0) then
      do k=1,nkml ; vhml(i,j,k) = 0.0 ; enddo
    else
      i2htot = 1.0 / (htot(i,j) + htot(i,j+1) + h_neglect)
      z_topx2 = 0.0
      ! a(k) relates the sublayer transport to vDml with a linear profile.
      ! The sum of a(k) through the mixed layers must be 0.
      do k=1,nkml
        hx2 = (h(i,j,k) + h(i,j+1,k) + h_neglect)
        a(k) = (hx2 * i2htot) * (2.0 - 4.0*(z_topx2+0.5*hx2)*i2htot)
        z_topx2 = z_topx2 + hx2
        if (a(k)*vdml(i) > 0.0) then
          if (a(k)*vdml(i) > h_avail(i,j,k)) vdml(i) = h_avail(i,j,k) / a(k)
        else
          if (-a(k)*vdml(i) > h_avail(i,j+1,k)) vdml(i) = -h_avail(i,j+1,k)/a(k)
        endif
      enddo
      do k=1,nkml
        vhml(i,j,k) = a(k)*vdml(i)
        vhtr(i,j,k) = vhtr(i,j,k) + vhml(i,j,k)*dt
      enddo
    endif

    vtimescale_diag(i,j) = timescale
    vdml_diag(i,j) = vdml(i)
  enddo; enddo

!$OMP do
  do j=js,je ; do k=1,nkml ; do i=is,ie
    h(i,j,k) = h(i,j,k) - dt*g%IareaT(i,j) * &
        ((uhml(i,j,k) - uhml(i-1,j,k)) + (vhml(i,j,k) - vhml(i,j-1,k)))
  enddo ; enddo ; enddo
!$OMP end parallel

  ! Whenever thickness changes let the diag manager know, target grids
  ! for vertical remapping may need to be regenerated.
  if (cs%id_uhml > 0 .or. cs%id_vhml > 0) &
    ! Remapped uhml and vhml require east/north halo updates of h
    call pass_var(h, g%domain, to_west+to_south+omit_corners, halo=1)
  call diag_update_remap_grids(cs%diag)

  ! Offer diagnostic fields for averaging.
  if (query_averaging_enabled(cs%diag) .and. &
    ((cs%id_urestrat_time > 0)  .or. (cs%id_vrestrat_time > 0))) then
    call post_data(cs%id_urestrat_time, utimescale_diag, cs%diag)
    call post_data(cs%id_vrestrat_time, vtimescale_diag, cs%diag)
  endif
  if (query_averaging_enabled(cs%diag) .and. &
      ((cs%id_uhml>0) .or. (cs%id_vhml>0))) then
    do k=nkml+1,nz
      do j=js,je ; do i=isq,ieq ; uhml(i,j,k) = 0.0 ; enddo ; enddo
      do j=jsq,jeq ; do i=is,ie ; vhml(i,j,k) = 0.0 ; enddo ; enddo
    enddo
    if (cs%id_uhml > 0) call post_data(cs%id_uhml, uhml,      cs%diag)
    if (cs%id_vhml > 0) call post_data(cs%id_vhml, vhml,      cs%diag)
    if (cs%id_MLD  > 0) call post_data(cs%id_MLD,  htot,      cs%diag)
    if (cs%id_Rml  > 0) call post_data(cs%id_Rml,  rml_av,    cs%diag)
    if (cs%id_uDml > 0) call post_data(cs%id_uDml, udml_diag, cs%diag)
    if (cs%id_vDml > 0) call post_data(cs%id_vDml, vdml_diag, cs%diag)
  endif

end subroutine mixedlayer_restrat_bml


!> Initialize the mixed layer restratification module
logical function mixedlayer_restrat_init(Time, G, GV, US, param_file, diag, CS, restart_CS)
  type(time_type),             intent(in)    :: time       !< Current model time
  type(ocean_grid_type),       intent(inout) :: g          !< Ocean grid structure
  type(verticalgrid_type),     intent(in)    :: gv         !< Ocean vertical grid structure
  type(unit_scale_type),       intent(in)    :: us         !< A dimensional unit scaling type
  type(param_file_type),       intent(in)    :: param_file !< Parameter file to parse
  type(diag_ctrl), target,     intent(inout) :: diag       !< Regulate diagnostics
  type(mixedlayer_restrat_cs), pointer       :: cs         !< Module control structure
  type(mom_restart_cs),        pointer       :: restart_cs !< A pointer to the restart control structure

  ! Local variables
  real :: h_rescale  ! A rescaling factor for thicknesses from the representation in
                     ! a restart file to the internal representation in this run.
  real :: flux_to_kg_per_s ! A unit conversion factor for fluxes.
  ! This include declares and sets the variable "version".
# include "version_variable.h"
  integer :: i, j

  ! Read all relevant parameters and write them to the model log.
  call get_param(param_file, mdl, "MIXEDLAYER_RESTRAT", mixedlayer_restrat_init, &
             default=.false., do_not_log=.true.)
  call log_version(param_file, mdl, version, "", all_default=.not.mixedlayer_restrat_init)
  call get_param(param_file, mdl, "MIXEDLAYER_RESTRAT", mixedlayer_restrat_init, &
             "If true, a density-gradient dependent re-stratifying "//&
             "flow is imposed in the mixed layer. Can be used in ALE mode "//&
             "without restriction but in layer mode can only be used if "//&
             "BULKMIXEDLAYER is true.", default=.false.)
  if (.not. mixedlayer_restrat_init) return

  if (.not.associated(cs)) then
    call mom_error(fatal, "mixedlayer_restrat_init called without an associated control structure.")
  endif

  ! Nonsense values to cause problems when these parameters are not used
  cs%MLE_MLD_decay_time = -9.e9*us%s_to_T
  cs%MLE_density_diff = -9.e9*us%kg_m3_to_R
  cs%MLE_tail_dh = -9.e9
  cs%MLE_use_PBL_MLD = .false.
  cs%MLE_MLD_stretch = -9.e9

  call get_param(param_file, mdl, "DEBUG", cs%debug, default=.false., do_not_log=.true.)
  call get_param(param_file, mdl, "FOX_KEMPER_ML_RESTRAT_COEF", cs%ml_restrat_coef, &
             "A nondimensional coefficient that is proportional to "//&
             "the ratio of the deformation radius to the dominant "//&
             "lengthscale of the submesoscale mixed layer "//&
             "instabilities, times the minimum of the ratio of the "//&
             "mesoscale eddy kinetic energy to the large-scale "//&
             "geostrophic kinetic energy or 1 plus the square of the "//&
             "grid spacing over the deformation radius, as detailed "//&
             "by Fox-Kemper et al. (2010)", units="nondim", default=0.0)
  ! We use GV%nkml to distinguish between the old and new implementation of MLE.
  ! The old implementation only works for the layer model with nkml>0.
  if (gv%nkml==0) then
    call get_param(param_file, mdl, "FOX_KEMPER_ML_RESTRAT_COEF2", cs%ml_restrat_coef2, &
             "As for FOX_KEMPER_ML_RESTRAT_COEF but used in a second application "//&
             "of the MLE restratification parameterization.", units="nondim", default=0.0)
    call get_param(param_file, mdl, "MLE_FRONT_LENGTH", cs%front_length, &
             "If non-zero, is the frontal-length scale used to calculate the "//&
             "upscaling of buoyancy gradients that is otherwise represented "//&
             "by the parameter FOX_KEMPER_ML_RESTRAT_COEF. If MLE_FRONT_LENGTH is "//&
             "non-zero, it is recommended to set FOX_KEMPER_ML_RESTRAT_COEF=1.0.",&
             units="m", default=0.0, scale=us%m_to_L)
    call get_param(param_file, mdl, "MLE_USE_PBL_MLD", cs%MLE_use_PBL_MLD, &
             "If true, the MLE parameterization will use the mixed-layer "//&
             "depth provided by the active PBL parameterization. If false, "//&
             "MLE will estimate a MLD based on a density difference with the "//&
             "surface using the parameter MLE_DENSITY_DIFF.", default=.false.)
    call get_param(param_file, mdl, "MLE_MLD_DECAY_TIME", cs%MLE_MLD_decay_time, &
             "The time-scale for a running-mean filter applied to the mixed-layer "//&
             "depth used in the MLE restratification parameterization. When "//&
             "the MLD deepens below the current running-mean the running-mean "//&
             "is instantaneously set to the current MLD.", units="s", default=0., scale=us%s_to_T)
    call get_param(param_file, mdl, "MLE_MLD_DECAY_TIME2", cs%MLE_MLD_decay_time2, &
             "The time-scale for a running-mean filter applied to the filtered "//&
             "mixed-layer depth used in a second MLE restratification parameterization. "//&
             "When the MLD deepens below the current running-mean the running-mean "//&
             "is instantaneously set to the current MLD.", units="s", default=0., scale=us%s_to_T)
    if (.not. cs%MLE_use_PBL_MLD) then
      call get_param(param_file, mdl, "MLE_DENSITY_DIFF", cs%MLE_density_diff, &
             "Density difference used to detect the mixed-layer "//&
             "depth used for the mixed-layer eddy parameterization "//&
             "by Fox-Kemper et al. (2010)", units="kg/m3", default=0.03, scale=us%kg_m3_to_R)
    endif
    call get_param(param_file, mdl, "MLE_TAIL_DH", cs%MLE_tail_dh, &
             "Fraction by which to extend the mixed-layer restratification "//&
             "depth used for a smoother stream function at the base of "//&
             "the mixed-layer.", units="nondim", default=0.0)
    call get_param(param_file, mdl, "MLE_MLD_STRETCH", cs%MLE_MLD_stretch, &
             "A scaling coefficient for stretching/shrinking the MLD "//&
             "used in the MLE scheme. This simply multiplies MLD wherever used.",&
             units="nondim", default=1.0)
  endif

  cs%diag => diag

  flux_to_kg_per_s = gv%H_to_kg_m2 * us%L_to_m**2 * us%s_to_T

  cs%id_uhml = register_diag_field('ocean_model', 'uhml', diag%axesCuL, time, &
      'Zonal Thickness Flux to Restratify Mixed Layer', 'kg s-1', &
      conversion=flux_to_kg_per_s, y_cell_method='sum', v_extensive=.true.)
  cs%id_vhml = register_diag_field('ocean_model', 'vhml', diag%axesCvL, time, &
      'Meridional Thickness Flux to Restratify Mixed Layer', 'kg s-1', &
      conversion=flux_to_kg_per_s, x_cell_method='sum', v_extensive=.true.)
  cs%id_urestrat_time = register_diag_field('ocean_model', 'MLu_restrat_time', diag%axesCu1, time, &
      'Mixed Layer Zonal Restratification Timescale', 's', conversion=us%T_to_s)
  cs%id_vrestrat_time = register_diag_field('ocean_model', 'MLv_restrat_time', diag%axesCv1, time, &
      'Mixed Layer Meridional Restratification Timescale', 's', conversion=us%T_to_s)
  cs%id_MLD = register_diag_field('ocean_model', 'MLD_restrat', diag%axesT1, time, &
      'Mixed Layer Depth as used in the mixed-layer restratification parameterization', 'm', &
      conversion=gv%H_to_m)
  cs%id_Rml = register_diag_field('ocean_model', 'ML_buoy_restrat', diag%axesT1, time, &
      'Mixed Layer Buoyancy as used in the mixed-layer restratification parameterization', &
      'm s2', conversion=us%m_to_Z*(us%L_to_m**2)*(us%s_to_T**2))
  cs%id_uDml = register_diag_field('ocean_model', 'udml_restrat', diag%axesCu1, time, &
      'Transport stream function amplitude for zonal restratification of mixed layer', &
      'm3 s-1', conversion=gv%H_to_m*(us%L_to_m**2)*us%s_to_T)
  cs%id_vDml = register_diag_field('ocean_model', 'vdml_restrat', diag%axesCv1, time, &
      'Transport stream function amplitude for meridional restratification of mixed layer', &
      'm3 s-1', conversion=gv%H_to_m*(us%L_to_m**2)*us%s_to_T)
  cs%id_uml = register_diag_field('ocean_model', 'uml_restrat', diag%axesCu1, time, &
      'Surface zonal velocity component of mixed layer restratification', &
      'm s-1', conversion=us%L_T_to_m_s)
  cs%id_vml = register_diag_field('ocean_model', 'vml_restrat', diag%axesCv1, time, &
      'Surface meridional velocity component of mixed layer restratification', &
      'm s-1', conversion=us%L_T_to_m_s)

  ! Rescale variables from restart files if the internal dimensional scalings have changed.
  if (cs%MLE_MLD_decay_time>0. .or. cs%MLE_MLD_decay_time2>0.) then
    if (query_initialized(cs%MLD_filtered, "MLD_MLE_filtered", restart_cs) .and. &
        (gv%m_to_H_restart /= 0.0) .and. (gv%m_to_H_restart /= gv%m_to_H)) then
      h_rescale = gv%m_to_H / gv%m_to_H_restart
      do j=g%jsc,g%jec ; do i=g%isc,g%iec
        cs%MLD_filtered(i,j) = h_rescale * cs%MLD_filtered(i,j)
      enddo ; enddo
    endif
  endif
  if (cs%MLE_MLD_decay_time2>0.) then
    if (query_initialized(cs%MLD_filtered_slow, "MLD_MLE_filtered_slow", restart_cs) .and. &
        (gv%m_to_H_restart /= 0.0) .and. (gv%m_to_H_restart /= gv%m_to_H)) then
      h_rescale = gv%m_to_H / gv%m_to_H_restart
      do j=g%jsc,g%jec ; do i=g%isc,g%iec
        cs%MLD_filtered_slow(i,j) = h_rescale * cs%MLD_filtered_slow(i,j)
      enddo ; enddo
    endif
  endif

  ! If MLD_filtered is being used, we need to update halo regions after a restart
  if (associated(cs%MLD_filtered)) call pass_var(cs%MLD_filtered, g%domain)

end function mixedlayer_restrat_init

!> Allocate and register fields in the mixed layer restratification structure for restarts
subroutine mixedlayer_restrat_register_restarts(HI, param_file, CS, restart_CS)
  ! Arguments
  type(hor_index_type),        intent(in)    :: hi         !< Horizontal index structure
  type(param_file_type),       intent(in)    :: param_file !< Parameter file to parse
  type(mixedlayer_restrat_cs), pointer       :: cs         !< Module control structure
  type(mom_restart_cs),        pointer       :: restart_cs !< A pointer to the restart control structure
  ! Local variables
  type(vardesc) :: vd
  logical :: mixedlayer_restrat_init

  ! Check to see if this module will be used
  call get_param(param_file, mdl, "MIXEDLAYER_RESTRAT", mixedlayer_restrat_init, &
             default=.false., do_not_log=.true.)
  if (.not. mixedlayer_restrat_init) return

  ! Allocate the control structure. CS will be later populated by mixedlayer_restrat_init()
  if (associated(cs)) call mom_error(fatal, &
       "mixedlayer_restrat_register_restarts called with an associated control structure.")
  allocate(cs)

  call get_param(param_file, mdl, "MLE_MLD_DECAY_TIME", cs%MLE_MLD_decay_time, &
                 default=0., do_not_log=.true.)
  call get_param(param_file, mdl, "MLE_MLD_DECAY_TIME2", cs%MLE_MLD_decay_time2, &
                 default=0., do_not_log=.true.)
  if (cs%MLE_MLD_decay_time>0. .or. cs%MLE_MLD_decay_time2>0.) then
    ! CS%MLD_filtered is used to keep a running mean of the PBL's actively mixed MLD.
    allocate(cs%MLD_filtered(hi%isd:hi%ied,hi%jsd:hi%jed)) ; cs%MLD_filtered(:,:) = 0.
    vd = var_desc("MLD_MLE_filtered","m","Time-filtered MLD for use in MLE", &
                  hor_grid='h', z_grid='1')
    call register_restart_field(cs%MLD_filtered, vd, .false., restart_cs)
  endif
  if (cs%MLE_MLD_decay_time2>0.) then
    ! CS%MLD_filtered_slow is used to keep a running mean of the PBL's seasonal or winter MLD.
    allocate(cs%MLD_filtered_slow(hi%isd:hi%ied,hi%jsd:hi%jed)) ; cs%MLD_filtered_slow(:,:) = 0.
    vd = var_desc("MLD_MLE_filtered_slow","m","c Slower time-filtered MLD for use in MLE", &
                  hor_grid='h', z_grid='1')
    call register_restart_field(cs%MLD_filtered_slow, vd, .false., restart_cs)
  endif

end subroutine mixedlayer_restrat_register_restarts

!> \namespace mom_mixed_layer_restrat
!!
!! \section section_mle Mixed-layer eddy parameterization module
!!
!! The subroutines in this file implement a parameterization of unresolved viscous
!! mixed layer restratification of the mixed layer as described in Fox-Kemper et
!! al., 2008, and whose impacts are described in Fox-Kemper et al., 2011.
!! This is derived in part from the older parameterization that is described in
!! Hallberg (Aha Hulikoa, 2003), which this new parameterization surpasses, which
!! in turn is based on the sub-inertial mixed layer theory of Young (JPO, 1994).
!! There is no net horizontal volume transport due to this parameterization, and
!! no direct effect below the mixed layer.
!!
!! This parameterization sets the restratification timescale to agree with
!! high-resolution studies of mixed layer restratification.
!!
!! The run-time parameter FOX_KEMPER_ML_RESTRAT_COEF is a non-dimensional number of
!! order a few tens, proportional to the ratio of the deformation radius or the
!! grid scale (whichever is smaller to the dominant horizontal length-scale of the
!! sub-meso-scale mixed layer instabilities.
!!
!! \subsection section_mle_nutshell "Sub-meso" in a nutshell
!!
!! The parameterization is colloquially referred to as "sub-meso".
!!
!! The original Fox-Kemper et al., (2008b) paper proposed a quasi-Stokes
!! advection described by the stream function (eq. 5 of Fox-Kemper et al., 2011):
!! \f[
!!    {\bf \Psi}_o = C_e \frac{ H^2 \nabla \bar{b} \times \hat{\bf z} }{ |f| } \mu(z)
!! \f]
!!
!! where the vertical profile function is
!! \f[
!!    \mu(z) = \max \left\{ 0, \left[ 1 - \left(\frac{2z}{H}+1\right)^2 \right]
!!                            \left[ 1 + \frac{5}{21} \left(\frac{2z}{H}+1\right)^2 \right] \right\}
!! \f]
!! and \f$ H \f$ is the mixed-layer depth, \f$ f \f$ is the local Coriolis parameter, \f$ C_e \sim 0.06-0.08 \f$ and
!! \f$ \nabla \bar{b} \f$ is a depth mean buoyancy gradient averaged over the mixed layer.
!!
!! For use in coarse-resolution models, an upscaling of the buoyancy gradients and adaption for the equator
!! leads to the following parameterization (eq. 6 of Fox-Kemper et al., 2011):
!! \f[
!!    {\bf \Psi} = C_e \Gamma_\Delta \frac{\Delta s}{l_f} \frac{ H^2 \nabla \bar{b} \times \hat{\bf z} }
!!                 { \sqrt{ f^2 + \tau^{-2}} } \mu(z)
!! \f]
!! where \f$ \Delta s \f$ is the minimum of grid-scale and deformation radius,
!! \f$ l_f \f$ is the width of the mixed-layer fronts, and \f$ \Gamma_\Delta=1 \f$.
!! \f$ \tau \f$ is a time-scale for mixing momentum across the mixed layer.
!! \f$ l_f \f$ is thought to be of order hundreds of meters.
!!
!! The upscaling factor \f$ \frac{\Delta s}{l_f} \f$ can be a global constant, model parameter FOX_KEMPER_ML_RESTRAT,
!! so that in practice the parameterization is:
!! \f[
!!    {\bf \Psi} = C_e \Gamma_\Delta \frac{ H^2 \nabla \bar{b} \times \hat{\bf z} }{ \sqrt{ f^2 + \tau^{-2}} } \mu(z)
!! \f]
!! with non-unity \f$ \Gamma_\Delta \f$.
!!
!! \f$ C_e \f$ is hard-coded as 0.0625. \f$ \tau \f$ is calculated from the surface friction velocity \f$ u^* \f$.
!! \todo Explain expression for momentum mixing time-scale.
!!
!! \subsection section_mle_filtering Time-filtering of mixed-layer depth
!!
!! Using the instantaneous mixed-layer depth is inconsistent with the finite life-time of
!! mixed-layer instabilities. We provide a one-sided running-mean filter of mixed-layer depth, \f$ H \f$, of the form:
!! \f[
!!    \bar{H} \leftarrow \max \left( H, \frac{ \Delta t H + \tau_h \bar{H} }{ \Delta t + \tau_h } \right)
!! \f]
!! which allows the effective mixed-layer depth seen by the parameterization, \f$\bar{H}\f$, to instantaneously deepen
!! but to decay with time-scale \f$ \tau_h \f$.
!! \f$ \bar{H} \f$ is substituted for \f$ H \f$ in the above equations.
!!
!! \subsection section_mle_mld Defining the mixed-layer-depth
!!
!! If the parameter MLE_USE_PBL_MLD=True then the mixed-layer depth is defined/diagnosed by the
!! boundary-layer parameterization (e.g. ePBL, KPP, etc.).
!!
!! If the parameter MLE_USE_PBL_MLD=False then the mixed-layer depth is diagnosed in this module
!! as the depth of a given density difference, \f$ \Delta \rho \f$, with the surface where the
!! density difference is the parameter MLE_DENSITY_DIFF.
!!
!! \subsection section_mle_ref References
!!
!! Fox-Kemper, B., Ferrari, R. and Hallberg, R., 2008:
!! Parameterization of Mixed Layer Eddies. Part I: Theory and Diagnosis
!! J. Phys. Oceangraphy, 38 (6), p1145-1165.
!! https://doi.org/10.1175/2007JPO3792.1
!!
!! Fox-Kemper, B. and Ferrari, R. 2008:
!! Parameterization of Mixed Layer Eddies. Part II: Prognosis and Impact
!! J. Phys. Oceangraphy, 38 (6), p1166-1179.
!! https://doi.org/10.1175/2007JPO3788.1
!!
!! B. Fox-Kemper, G. Danabasoglu, R. Ferrari, S.M. Griffies, R.W. Hallberg, M.M. Holland, M.E. Maltrud,
!! S. Peacock, and B.L. Samuels, 2011: Parameterization of mixed layer eddies. III: Implementation and impact
!! in global ocean climate simulations. Ocean Modell., 39(1), p61-78.
!! https://doi.org/10.1016/j.ocemod.2010.09.002
!!
!! | Symbol                       | Module parameter      |
!! | ---------------------------- | --------------------- |
!! | \f$ \Gamma_\Delta \f$        | FOX_KEMPER_ML_RESTRAT |
!! | \f$ l_f \f$                  | MLE_FRONT_LENGTH      |
!! | \f$ \tau_h \f$               | MLE_MLD_DECAY_TIME    |
!! | \f$ \Delta \rho \f$          | MLE_DENSITY_DIFF      |

end module mom_mixed_layer_restrat