MOM_kappa_shear.F90

!> Shear-dependent mixing following Jackson et al. 2008.
module mom_kappa_shear

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

use mom_cpu_clock, only : cpu_clock_id, cpu_clock_begin, cpu_clock_end
use mom_cpu_clock, only : clock_module_driver, clock_module, clock_routine
use mom_diag_mediator, only : post_data, register_diag_field, safe_alloc_ptr
use mom_diag_mediator, only : diag_ctrl, time_type
use mom_debugging, only : hchksum, bchksum
use mom_error_handler, only : mom_error, is_root_pe, fatal, warning, note
use mom_file_parser, only : get_param, log_version, param_file_type
use mom_grid, only : ocean_grid_type
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_derivs

implicit none ; private

#include <MOM_memory.h>

public calculate_kappa_shear, calc_kappa_shear_vertex, kappa_shear_init
public kappa_shear_is_used, kappa_shear_at_vertex

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

!> This control structure holds the parameters that regulate shear mixing
type, public :: kappa_shear_cs ; private
  real    :: rino_crit       !< The critical shear Richardson number for
                             !! shear-entrainment [nondim]. The theoretical value is 0.25.
                             !! The values found by Jackson et al. are 0.25-0.35.
  real    :: shearmix_rate   !< A nondimensional rate scale for shear-driven
                             !! entrainment [nondim].  The value given by Jackson et al.
                             !! is 0.085-0.089.
  real    :: fri_curvature   !<   A constant giving the curvature of the function
                             !! of the Richardson number that relates shear to
                             !! sources in the kappa equation [nondim].
                             !! The values found by Jackson et al. are -0.97 - -0.89.
  real    :: c_n             !<   The coefficient for the decay of TKE due to
                             !! stratification (i.e. proportional to N*tke) [nondim].
                             !! The values found by Jackson et al. are 0.24-0.28.
  real    :: c_s             !<   The coefficient for the decay of TKE due to
                             !! shear (i.e. proportional to |S|*tke) [nondim].
                             !! The values found by Jackson et al. are 0.14-0.12.
  real    :: lambda          !<   The coefficient for the buoyancy length scale
                             !! in the kappa equation [nondim].
                             !! The values found by Jackson et al. are 0.82-0.81.
  real    :: lambda2_n_s     !<   The square of the ratio of the coefficients of
                             !! the buoyancy and shear scales in the diffusivity
                             !! equation, 0 to eliminate the shear scale [nondim].
  real    :: tke_bg          !<   The background level of TKE [Z2 T-2 ~> m2 s-2].
  real    :: kappa_0         !<   The background diapycnal diffusivity [Z2 T-1 ~> m2 s-1].
  real    :: kappa_trunc     !< Diffusivities smaller than this are rounded to 0 [Z2 T-1 ~> m2 s-1].
  real    :: kappa_tol_err   !<   The fractional error in kappa that is tolerated [nondim].
  real    :: prandtl_turb    !< Prandtl number used to convert Kd_shear into viscosity [nondim].
  integer :: nkml            !<   The number of layers in the mixed layer, as
                             !! treated in this routine.  If the pieces of the
                             !! mixed layer are not to be treated collectively,
                             !! nkml is set to 1.
  integer :: max_rino_it     !< The maximum number of iterations that may be used
                             !! to estimate the instantaneous shear-driven mixing.
  integer :: max_ks_it       !< The maximum number of iterations that may be used
                             !! to estimate the time-averaged diffusivity.
  logical :: dkdq_iteration_bug !< If true. use an older, dimensionally inconsistent estimate of
                             !! the derivative of diffusivity with energy in the Newton's method
                             !! iteration.  The bug causes undercorrections when dz > 1m.
  logical :: ks_at_vertex    !< If true, do the calculations of the shear-driven mixing
                             !! at the cell vertices (i.e., the vorticity points).
  logical :: eliminate_massless !< If true, massless layers are merged with neighboring
                             !! massive layers in this calculation.
                             !  I can think of no good reason why this should be false. - RWH
  real    :: vel_underflow   !< Velocity components smaller than vel_underflow
                             !! are set to 0 [L T-1 ~> m s-1].
  real    :: kappa_src_max_chg !< The maximum permitted increase in the kappa source within an
                             !! iteration relative to the local source [nondim].  This must be
                             !! greater than 1.  The lower limit for the permitted fractional
                             !! decrease is (1 - 0.5/kappa_src_max_chg).  These limits could
                             !! perhaps be made dynamic with an improved iterative solver.
  logical :: psurf_bug       !< If true, do a simple average of the cell surface pressures to get a
                             !! surface pressure at the corner if VERTEX_SHEAR=True.  Otherwise mask
                             !! out any land points in the average.
  logical :: all_layer_tke_bug !< If true, report back the latest estimate of TKE instead of the
                             !! time average TKE when there is mass in all layers.  Otherwise always
                             !! report the time-averaged TKE, as is currently done when there
                             !! are some massless layers.
  logical :: restrictive_tolerance_check !< If false, uses the less restrictive tolerance check to
                             !! determine if a timestep is acceptable for the KS_it outer iteration
                             !! loop, as the code was originally written.  True uses the more
                             !! restrictive check.
!  logical :: layer_stagger = .false. ! If true, do the calculations centered at
                             !  layers, rather than the interfaces.
  logical :: debug = .false. !< If true, write verbose debugging messages.
  type(diag_ctrl), pointer :: diag => null() !< A structure that is used to
                             !! regulate the timing of diagnostic output.
  !>@{ Diagnostic IDs
  integer :: id_kd_shear = -1, id_tke = -1, id_ild2 = -1, id_dz_int = -1
  !>@}
end type kappa_shear_cs

! integer :: id_clock_project, id_clock_KQ, id_clock_avg, id_clock_setup

contains

!> Subroutine for calculating shear-driven diffusivity and TKE in tracer columns
subroutine calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, &
                                 kv_io, dt, G, GV, US, CS, initialize_all)
  type(ocean_grid_type),   intent(in)    :: g      !< The ocean's grid structure.
  type(verticalgrid_type), intent(in)    :: gv     !< The ocean's vertical grid structure.
  type(unit_scale_type),   intent(in)    :: us     !< A dimensional unit scaling type
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)),   &
                           intent(in)    :: u_in   !< Initial zonal velocity [L T-1 ~> m s-1].
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)),   &
                           intent(in)    :: v_in   !< Initial meridional velocity [L T-1 ~> m s-1].
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)),   &
                           intent(in)    :: h      !< Layer thicknesses [H ~> m or kg m-2].
  type(thermo_var_ptrs),   intent(in)    :: tv     !< A structure containing pointers to any
                                                   !! available thermodynamic fields. Absent fields
                                                   !! have NULL ptrs.
  real, dimension(:,:),    pointer       :: p_surf !< The pressure at the ocean surface [R L2 T-2 ~> Pa] (or NULL).
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), &
                           intent(inout) :: kappa_io !< The diapycnal diffusivity at each interface
                                                   !! (not layer!) [Z2 T-1 ~> m2 s-1].  Initially this is the
                                                   !! value from the previous timestep, which may
                                                   !! accelerate the iteration toward convergence.
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), &
                           intent(out) :: tke_io   !< The turbulent kinetic energy per unit mass at
                                                   !! each interface (not layer!) [Z2 T-2 ~> m2 s-2].
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), &
                           intent(inout) :: kv_io  !< The vertical viscosity at each interface
                                                   !! (not layer!) [Z2 T-1 ~> m2 s-1]. This discards any
                                                   !! previous value (i.e. it is intent out) and
                                                   !! simply sets Kv = Prandtl * Kd_shear
  real,                    intent(in)    :: dt     !< Time increment [T ~> s].
  type(kappa_shear_cs),    pointer       :: cs     !< The control structure returned by a previous
                                                   !! call to kappa_shear_init.
  logical,       optional, intent(in)    :: initialize_all !< If present and false, the previous
                                                   !! value of kappa is used to start the iterations

  ! Local variables
  real, dimension(SZI_(G),SZK_(GV)) :: &
    h_2d, &             ! A 2-D version of h, but converted to [Z ~> m].
    u_2d, v_2d, &       ! 2-D versions of u_in and v_in, converted to [L T-1 ~> m s-1].
    t_2d, s_2d, rho_2d  ! 2-D versions of T [degC], S [ppt], and rho [R ~> kg m-3].
  real, dimension(SZI_(G),SZK_(GV)+1) :: &
    kappa_2d, & ! 2-D version of kappa_io [Z2 T-1 ~> m2 s-1].
    tke_2d      ! 2-D version tke_io [Z2 T-2 ~> m2 s-2].
  real, dimension(SZK_(GV)) :: &
    idz, &      ! The inverse of the distance between TKE points [Z-1 ~> m-1].
    dz, &       ! The layer thickness [Z ~> m].
    u0xdz, &    ! The initial zonal velocity times dz [Z L T-1 ~> m2 s-1].
    v0xdz, &    ! The initial meridional velocity times dz [Z L T-1 ~> m2 s-1].
    t0xdz, &    ! The initial temperature times dz [degC Z ~> degC m].
    s0xdz       ! The initial salinity times dz [ppt Z ~> ppt m].
  real, dimension(SZK_(GV)+1) :: &
    kappa, &    ! The shear-driven diapycnal diffusivity at an interface [Z2 T-1 ~> m2 s-1].
    tke, &      ! The Turbulent Kinetic Energy per unit mass at an interface [Z2 T-2 ~> m2 s-2].
    kappa_avg, & ! The time-weighted average of kappa [Z2 T-1 ~> m2 s-1].
    tke_avg     ! The time-weighted average of TKE [Z2 T-2 ~> m2 s-2].
  real :: f2   ! The squared Coriolis parameter of each column [T-2 ~> s-2].
  real :: surface_pres  ! The top surface pressure [R L2 T-2 ~> Pa].

  real :: dz_in_lay     !   The running sum of the thickness in a layer [Z ~> m].
  real :: k0dt          ! The background diffusivity times the timestep [Z2 ~> m2].
  real :: dz_massless   ! A layer thickness that is considered massless [Z ~> m].
  logical :: use_temperature  !  If true, temperature and salinity have been
                        ! allocated and are being used as state variables.
  logical :: new_kappa = .true. ! If true, ignore the value of kappa from the
                        ! last call to this subroutine.

  integer, dimension(SZK_(GV)+1) :: kc ! The index map between the original
                        ! interfaces and the interfaces with massless layers
                        ! merged into nearby massive layers.
  real, dimension(SZK_(GV)+1) :: kf ! The fractional weight of interface kc+1 for
                        ! interpolating back to the original index space [nondim].
  integer :: is, ie, js, je, i, j, k, nz, nzc

  is = g%isc ; ie = g%iec; js = g%jsc ; je = g%jec ; nz = gv%ke

  use_temperature = .false. ; if (associated(tv%T)) use_temperature = .true.
  new_kappa = .true. ; if (present(initialize_all)) new_kappa = initialize_all

  k0dt = dt*cs%kappa_0
  dz_massless = 0.1*sqrt(k0dt)

  !$OMP parallel do default(private) shared(js,je,is,ie,nz,h,u_in,v_in,use_temperature,new_kappa, &
  !$OMP                                tv,G,GV,US,CS,kappa_io,dz_massless,k0dt,p_surf,dt,tke_io,kv_io)
  do j=js,je
    do k=1,nz ; do i=is,ie
      h_2d(i,k) = h(i,j,k)*gv%H_to_Z
      u_2d(i,k) = u_in(i,j,k) ; v_2d(i,k) = v_in(i,j,k)
    enddo ; enddo
    if (use_temperature) then ; do k=1,nz ; do i=is,ie
      t_2d(i,k) = tv%T(i,j,k) ; s_2d(i,k) = tv%S(i,j,k)
    enddo ; enddo ; else ; do k=1,nz ; do i=is,ie
      rho_2d(i,k) = gv%Rlay(k) ! Could be tv%Rho(i,j,k) ?
    enddo ; enddo ; endif
    if (.not.new_kappa) then ; do k=1,nz+1 ; do i=is,ie
      kappa_2d(i,k) = kappa_io(i,j,k)
    enddo ; enddo ; endif

!---------------------------------------
! Work on each column.
!---------------------------------------
    do i=is,ie ; if (g%mask2dT(i,j) > 0.5) then
    ! call cpu_clock_begin(id_clock_setup)
      ! Store a transposed version of the initial arrays.
      ! Any elimination of massless layers would occur here.
      if (cs%eliminate_massless) then
        nzc = 1
        do k=1,nz
          ! Zero out the thicknesses of all layers, even if they are unused.
          dz(k) = 0.0 ; u0xdz(k) = 0.0 ; v0xdz(k) = 0.0
          t0xdz(k) = 0.0 ; s0xdz(k) = 0.0

          ! Add a new layer if this one has mass.
!          if ((dz(nzc) > 0.0) .and. (h_2d(i,k) > dz_massless)) nzc = nzc+1
          if ((k>cs%nkml) .and. (dz(nzc) > 0.0) .and. &
              (h_2d(i,k) > dz_massless)) nzc = nzc+1

          ! Only merge clusters of massless layers.
!         if ((dz(nzc) > dz_massless) .or. &
!             ((dz(nzc) > 0.0) .and. (h_2d(i,k) > dz_massless))) nzc = nzc+1

          kc(k) = nzc
          dz(nzc) = dz(nzc) + h_2d(i,k)
          u0xdz(nzc) = u0xdz(nzc) + u_2d(i,k)*h_2d(i,k)
          v0xdz(nzc) = v0xdz(nzc) + v_2d(i,k)*h_2d(i,k)
          if (use_temperature) then
            t0xdz(nzc) = t0xdz(nzc) + t_2d(i,k)*h_2d(i,k)
            s0xdz(nzc) = s0xdz(nzc) + s_2d(i,k)*h_2d(i,k)
          else
            t0xdz(nzc) = t0xdz(nzc) + rho_2d(i,k)*h_2d(i,k)
            s0xdz(nzc) = s0xdz(nzc) + rho_2d(i,k)*h_2d(i,k)
          endif
        enddo
        kc(nz+1) = nzc+1

        ! Set up Idz as the inverse of layer thicknesses.
        do k=1,nzc ; idz(k) = 1.0 / dz(k) ; enddo

        !   Now determine kf, the fractional weight of interface kc when
        ! interpolating between interfaces kc and kc+1.
        kf(1) = 0.0 ; dz_in_lay = h_2d(i,1)
        do k=2,nz
          if (kc(k) > kc(k-1)) then
            kf(k) = 0.0 ; dz_in_lay = h_2d(i,k)
          else
            kf(k) = dz_in_lay*idz(kc(k)) ; dz_in_lay = dz_in_lay + h_2d(i,k)
          endif
        enddo
        kf(nz+1) = 0.0
      else
        do k=1,nz
          dz(k) = h_2d(i,k)
          u0xdz(k) = u_2d(i,k)*dz(k) ; v0xdz(k) = v_2d(i,k)*dz(k)
        enddo
        if (use_temperature) then
          do k=1,nz
            t0xdz(k) = t_2d(i,k)*dz(k) ; s0xdz(k) = s_2d(i,k)*dz(k)
          enddo
        else
          do k=1,nz
            t0xdz(k) = rho_2d(i,k)*dz(k) ; s0xdz(k) = rho_2d(i,k)*dz(k)
          enddo
        endif
        nzc = nz
        do k=1,nzc+1 ; kc(k) = k ; kf(k) = 0.0 ; enddo
      endif
      f2 = 0.25 * ((g%CoriolisBu(i,j)**2 + g%CoriolisBu(i-1,j-1)**2) + &
                   (g%CoriolisBu(i,j-1)**2 + g%CoriolisBu(i-1,j)**2))
      surface_pres = 0.0 ; if (associated(p_surf)) surface_pres = p_surf(i,j)

    ! ----------------------------------------------------    I_Ld2_1d, dz_Int_1d

    ! Set the initial guess for kappa, here defined at interfaces.
    ! ----------------------------------------------------
      if (new_kappa) then
        do k=1,nzc+1 ; kappa(k) = us%m2_s_to_Z2_T*1.0 ; enddo
      else
        do k=1,nzc+1 ; kappa(k) = kappa_2d(i,k) ; enddo
      endif

      call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, &
                              dz, u0xdz, v0xdz, t0xdz, s0xdz, kappa_avg, &
                              tke_avg, tv, cs, gv, us)

    ! call cpu_clock_begin(id_clock_setup)
    ! Extrapolate from the vertically reduced grid back to the original layers.
      if (nz == nzc) then
        do k=1,nz+1
          kappa_2d(i,k) = kappa_avg(k)
          if (cs%all_layer_TKE_bug) then
            tke_2d(i,k) = tke(k)
          else
            tke_2d(i,k) = tke_avg(k)
          endif
        enddo
      else
        do k=1,nz+1
          if (kf(k) == 0.0) then
            kappa_2d(i,k) = kappa_avg(kc(k))
            tke_2d(i,k) = tke_avg(kc(k))
          else
            kappa_2d(i,k) = (1.0-kf(k)) * kappa_avg(kc(k)) + &
                             kf(k) * kappa_avg(kc(k)+1)
            tke_2d(i,k) = (1.0-kf(k)) * tke_avg(kc(k)) + &
                           kf(k) * tke_avg(kc(k)+1)
          endif
        enddo
      endif
    ! call cpu_clock_end(id_clock_setup)
    else  ! Land points, still inside the i-loop.
      do k=1,nz+1
        kappa_2d(i,k) = 0.0 ; tke_2d(i,k) = 0.0
      enddo
    endif ; enddo ! i-loop

    do k=1,nz+1 ; do i=is,ie
      kappa_io(i,j,k) = g%mask2dT(i,j) * kappa_2d(i,k)
      tke_io(i,j,k) = g%mask2dT(i,j) * tke_2d(i,k)
      kv_io(i,j,k) = ( g%mask2dT(i,j) * kappa_2d(i,k) ) * cs%Prandtl_turb
    enddo ; enddo

  enddo ! end of j-loop

  if (cs%debug) then
    call hchksum(kappa_io, "kappa", g%HI, scale=us%Z2_T_to_m2_s)
    call hchksum(tke_io, "tke", g%HI, scale=us%Z_to_m**2*us%s_to_T**2)
  endif

  if (cs%id_Kd_shear > 0) call post_data(cs%id_Kd_shear, kappa_io, cs%diag)
  if (cs%id_TKE > 0) call post_data(cs%id_TKE, tke_io, cs%diag)

end subroutine calculate_kappa_shear


!> Subroutine for calculating shear-driven diffusivity and TKE in corner columns
subroutine calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_io, tke_io, &
                                   kv_io, dt, G, GV, US, CS, initialize_all)
  type(ocean_grid_type),   intent(in)    :: g      !< The ocean's grid structure.
  type(verticalgrid_type), intent(in)    :: gv     !< The ocean's vertical grid structure.
  type(unit_scale_type),    intent(in)   :: us     !< A dimensional unit scaling type
  real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)),   &
                           intent(in)    :: u_in   !< Initial zonal velocity [L T-1 ~> m s-1].
  real, dimension(SZI_(G),SZJB_(G),SZK_(GV)),   &
                           intent(in)    :: v_in   !< Initial meridional velocity [L T-1 ~> m s-1].
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)),   &
                           intent(in)    :: h      !< Layer thicknesses [H ~> m or kg m-2].
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)),   &
                           intent(in)    :: t_in   !< Layer potential temperatures [degC]
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)),   &
                           intent(in)    :: s_in   !< Layer salinities in ppt.
  type(thermo_var_ptrs),   intent(in)    :: tv     !< A structure containing pointers to any
                                                   !! available thermodynamic fields. Absent fields
                                                   !! have NULL ptrs.
  real, dimension(:,:),    pointer       :: p_surf !< The pressure at the ocean surface [R L2 T-2 ~> Pa]
                                                   !! (or NULL).
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1), &
                           intent(out)   :: kappa_io !< The diapycnal diffusivity at each interface
                                                   !! (not layer!) [Z2 T-1 ~> m2 s-1].
  real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)+1), &
                           intent(out)   :: tke_io !< The turbulent kinetic energy per unit mass at
                                                   !! each interface (not layer!) [Z2 T-2 ~> m2 s-2].
  real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)+1), &
                           intent(inout) :: kv_io  !< The vertical viscosity at each interface [Z2 T-1 ~> m2 s-1].
                                                   !! The previous value is used to initialize kappa
                                                   !! in the vertex columes as Kappa = Kv/Prandtl
                                                   !! to accelerate the iteration toward covergence.
  real,                    intent(in)    :: dt     !< Time increment [T ~> s].
  type(kappa_shear_cs),    pointer       :: cs     !< The control structure returned by a previous
                                                   !! call to kappa_shear_init.
  logical,       optional, intent(in)    :: initialize_all !< If present and false, the previous
                                                   !! value of kappa is used to start the iterations

  ! Local variables
  real, dimension(SZIB_(G),SZK_(GV)) :: &
    h_2d, &             ! A 2-D version of h, but converted to [Z ~> m].
    u_2d, v_2d, &       ! 2-D versions of u_in and v_in, converted to [L T-1 ~> m s-1].
    t_2d, s_2d, rho_2d  ! 2-D versions of T [degC], S [ppt], and rho [R ~> kg m-3].
  real, dimension(SZIB_(G),SZK_(GV)+1,2) :: &
    kappa_2d    ! Quasi 2-D versions of kappa_io [Z2 T-1 ~> m2 s-1].
  real, dimension(SZIB_(G),SZK_(GV)+1) :: &
    tke_2d      ! 2-D version tke_io [Z2 T-2 ~> m2 s-2].
  real, dimension(SZK_(GV)) :: &
    idz, &      ! The inverse of the distance between TKE points [Z-1 ~> m-1].
    dz, &       ! The layer thickness [Z ~> m].
    u0xdz, &    ! The initial zonal velocity times dz [L Z T-1 ~> m2 s-1].
    v0xdz, &    ! The initial meridional velocity times dz [L Z T-1 ~> m2 s-1].
    t0xdz, &    ! The initial temperature times dz [degC Z ~> degC m].
    s0xdz       ! The initial salinity times dz [ppt Z ~> ppt m].
  real, dimension(SZK_(GV)+1) :: &
    kappa, &    ! The shear-driven diapycnal diffusivity at an interface [Z2 T-1 ~> m2 s-1].
    tke, &      ! The Turbulent Kinetic Energy per unit mass at an interface [Z2 T-2 ~> m2 s-2].
    kappa_avg, & ! The time-weighted average of kappa [Z2 T-1 ~> m2 s-1].
    tke_avg     ! The time-weighted average of TKE [Z2 T-2 ~> m2 s-2].
  real :: f2   ! The squared Coriolis parameter of each column [T-2 ~> s-2].
  real :: surface_pres  ! The top surface pressure [R L2 T-2 ~> Pa].

  real :: dz_in_lay     !   The running sum of the thickness in a layer [Z ~> m].
  real :: k0dt          ! The background diffusivity times the timestep [Z2 ~> m2].
  real :: dz_massless   ! A layer thickness that is considered massless [Z ~> m].
  real :: i_hwt         ! The inverse of the masked thickness weights [H-1 ~> m-1 or m2 kg-1].
  real :: i_prandtl     ! The inverse of the turbulent Prandtl number [nondim].
  logical :: use_temperature  !  If true, temperature and salinity have been
                        ! allocated and are being used as state variables.
  logical :: new_kappa = .true. ! If true, ignore the value of kappa from the
                        ! last call to this subroutine.
  logical :: do_i       ! If true, work on this column.

  integer, dimension(SZK_(GV)+1) :: kc ! The index map between the original
                        ! interfaces and the interfaces with massless layers
                        ! merged into nearby massive layers.
  real, dimension(SZK_(GV)+1) :: kf ! The fractional weight of interface kc+1 for
                        ! interpolating back to the original index space [nondim].
  integer :: isb, ieb, jsb, jeb, i, j, k, nz, nzc, j2, j2m1

  ! Diagnostics that should be deleted?
  isb = g%isc-1 ; ieb = g%iecB ; jsb = g%jsc-1 ; jeb = g%jecB ; nz = gv%ke

  use_temperature = .false. ; if (associated(tv%T)) use_temperature = .true.
  new_kappa = .true. ; if (present(initialize_all)) new_kappa = initialize_all

  k0dt =  dt*cs%kappa_0
  dz_massless = 0.1*sqrt(k0dt)
  i_prandtl = 0.0 ; if (cs%Prandtl_turb > 0.0) i_prandtl = 1.0 / cs%Prandtl_turb

  !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,use_temperature,new_kappa, &
  !$OMP                                tv,G,GV,US,CS,kappa_io,dz_massless,k0dt,p_surf,dt,tke_io,kv_io,I_Prandtl)
  do j=jsb,jeb
    j2 = mod(j,2)+1 ; j2m1 = 3-j2 ! = mod(J-1,2)+1

    ! Interpolate the various quantities to the corners, using masks.
    do k=1,nz ; do i=isb,ieb
      u_2d(i,k) = (u_in(i,j,k)   * (g%mask2dCu(i,j)   * (h(i,j,k)   + h(i+1,j,k))) + &
                   u_in(i,j+1,k) * (g%mask2dCu(i,j+1) * (h(i,j+1,k) + h(i+1,j+1,k))) ) / &
                  ((g%mask2dCu(i,j)   * (h(i,j,k)   + h(i+1,j,k)) + &
                    g%mask2dCu(i,j+1) * (h(i,j+1,k) + h(i+1,j+1,k))) + gv%H_subroundoff)
      v_2d(i,k) = (v_in(i,j,k)   * (g%mask2dCv(i,j)   * (h(i,j,k)   + h(i,j+1,k))) + &
                   v_in(i+1,j,k) * (g%mask2dCv(i+1,j) * (h(i+1,j,k) + h(i+1,j+1,k))) ) / &
                  ((g%mask2dCv(i,j)   * (h(i,j,k)   + h(i,j+1,k)) + &
                    g%mask2dCv(i+1,j) * (h(i+1,j,k) + h(i+1,j+1,k))) + gv%H_subroundoff)
      i_hwt = 1.0 / (((g%mask2dT(i,j) * h(i,j,k) + g%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + &
                      (g%mask2dT(i+1,j) * h(i+1,j,k) + g%mask2dT(i,j+1) * h(i,j+1,k))) + &
                     gv%H_subroundoff)
      if (use_temperature) then
        t_2d(i,k) = ( ((g%mask2dT(i,j) * h(i,j,k)) * t_in(i,j,k) + &
                       (g%mask2dT(i+1,j+1) * h(i+1,j+1,k)) * t_in(i+1,j+1,k)) + &
                      ((g%mask2dT(i+1,j) * h(i+1,j,k)) * t_in(i+1,j,k) + &
                       (g%mask2dT(i,j+1) * h(i,j+1,k)) * t_in(i,j+1,k)) ) * i_hwt
        s_2d(i,k) = ( ((g%mask2dT(i,j) * h(i,j,k)) * s_in(i,j,k) + &
                       (g%mask2dT(i+1,j+1) * h(i+1,j+1,k)) * s_in(i+1,j+1,k)) + &
                      ((g%mask2dT(i+1,j) * h(i+1,j,k)) * s_in(i+1,j,k) + &
                       (g%mask2dT(i,j+1) * h(i,j+1,k)) * s_in(i,j+1,k)) ) * i_hwt
      endif
      h_2d(i,k) = gv%H_to_Z * ((g%mask2dT(i,j) * h(i,j,k) + g%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + &
                               (g%mask2dT(i+1,j) * h(i+1,j,k) + g%mask2dT(i,j+1) * h(i,j+1,k)) ) / &
                              ((g%mask2dT(i,j) + g%mask2dT(i+1,j+1)) + &
                               (g%mask2dT(i+1,j) + g%mask2dT(i,j+1)) + 1.0e-36 )
!      h_2d(I,k) = 0.25*((h(i,j,k) + h(i+1,j+1,k)) + (h(i+1,j,k) + h(i,j+1,k)))*GV%H_to_Z
!      h_2d(I,k) = ((h(i,j,k)**2 + h(i+1,j+1,k)**2) + &
!                   (h(i+1,j,k)**2 + h(i,j+1,k)**2))*GV%H_to_Z * I_hwt
    enddo ; enddo
    if (.not.use_temperature) then ; do k=1,nz ; do i=isb,ieb
      rho_2d(i,k) = gv%Rlay(k)
    enddo ; enddo ; endif
    if (.not.new_kappa) then ; do k=1,nz+1 ; do i=isb,ieb
      kappa_2d(i,k,j2) = kv_io(i,j,k) * i_prandtl
    enddo ; enddo ; endif

!---------------------------------------
! Work on each column.
!---------------------------------------
    do i=isb,ieb ; if ((g%mask2dCu(i,j) + g%mask2dCu(i,j+1)) + &
                       (g%mask2dCv(i,j) + g%mask2dCv(i+1,j)) > 0.0) then
    ! call cpu_clock_begin(Id_clock_setup)
      ! Store a transposed version of the initial arrays.
      ! Any elimination of massless layers would occur here.
      if (cs%eliminate_massless) then
        nzc = 1
        do k=1,nz
          ! Zero out the thicknesses of all layers, even if they are unused.
          dz(k) = 0.0 ; u0xdz(k) = 0.0 ; v0xdz(k) = 0.0
          t0xdz(k) = 0.0 ; s0xdz(k) = 0.0

          ! Add a new layer if this one has mass.
!          if ((dz(nzc) > 0.0) .and. (h_2d(I,k) > dz_massless)) nzc = nzc+1
          if ((k>cs%nkml) .and. (dz(nzc) > 0.0) .and. &
              (h_2d(i,k) > dz_massless)) nzc = nzc+1

          ! Only merge clusters of massless layers.
!         if ((dz(nzc) > dz_massless) .or. &
!             ((dz(nzc) > 0.0) .and. (h_2d(I,k) > dz_massless))) nzc = nzc+1

          kc(k) = nzc
          dz(nzc) = dz(nzc) + h_2d(i,k)
          u0xdz(nzc) = u0xdz(nzc) + u_2d(i,k)*h_2d(i,k)
          v0xdz(nzc) = v0xdz(nzc) + v_2d(i,k)*h_2d(i,k)
          if (use_temperature) then
            t0xdz(nzc) = t0xdz(nzc) + t_2d(i,k)*h_2d(i,k)
            s0xdz(nzc) = s0xdz(nzc) + s_2d(i,k)*h_2d(i,k)
          else
            t0xdz(nzc) = t0xdz(nzc) + rho_2d(i,k)*h_2d(i,k)
            s0xdz(nzc) = s0xdz(nzc) + rho_2d(i,k)*h_2d(i,k)
          endif
        enddo
        kc(nz+1) = nzc+1

        ! Set up Idz as the inverse of layer thicknesses.
        do k=1,nzc ; idz(k) = 1.0 / dz(k) ; enddo

        !   Now determine kf, the fractional weight of interface kc when
        ! interpolating between interfaces kc and kc+1.
        kf(1) = 0.0 ; dz_in_lay = h_2d(i,1)
        do k=2,nz
          if (kc(k) > kc(k-1)) then
            kf(k) = 0.0 ; dz_in_lay = h_2d(i,k)
          else
            kf(k) = dz_in_lay*idz(kc(k)) ; dz_in_lay = dz_in_lay + h_2d(i,k)
          endif
        enddo
        kf(nz+1) = 0.0
      else
        do k=1,nz
          dz(k) = h_2d(i,k)
          u0xdz(k) = u_2d(i,k)*dz(k) ; v0xdz(k) = v_2d(i,k)*dz(k)
        enddo
        if (use_temperature) then
          do k=1,nz
            t0xdz(k) = t_2d(i,k)*dz(k) ; s0xdz(k) = s_2d(i,k)*dz(k)
          enddo
        else
          do k=1,nz
            t0xdz(k) = rho_2d(i,k)*dz(k) ; s0xdz(k) = rho_2d(i,k)*dz(k)
          enddo
        endif
        nzc = nz
        do k=1,nzc+1 ; kc(k) = k ; kf(k) = 0.0 ; enddo
      endif
      f2 = g%CoriolisBu(i,j)**2
      surface_pres = 0.0
      if (associated(p_surf)) then
        if (cs%psurf_bug) then
          ! This is wrong because it is averaging values from land in some places.
          surface_pres = 0.25 * ((p_surf(i,j) + p_surf(i+1,j+1)) + &
                                 (p_surf(i+1,j) + p_surf(i,j+1)))
        else
          surface_pres = ((g%mask2dT(i,j) * p_surf(i,j) + g%mask2dT(i+1,j+1) * p_surf(i+1,j+1)) + &
                          (g%mask2dT(i+1,j) * p_surf(i+1,j) + g%mask2dT(i,j+1) * p_surf(i,j+1)) ) / &
                         ((g%mask2dT(i,j) + g%mask2dT(i+1,j+1)) + &
                          (g%mask2dT(i+1,j) + g%mask2dT(i,j+1)) + 1.0e-36 )
        endif
      endif

    ! ----------------------------------------------------
    ! Set the initial guess for kappa, here defined at interfaces.
    ! ----------------------------------------------------
      if (new_kappa) then
        do k=1,nzc+1 ; kappa(k) = us%m2_s_to_Z2_T*1.0 ; enddo
      else
        do k=1,nzc+1 ; kappa(k) = kappa_2d(i,k,j2) ; enddo
      endif

      call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, &
                              dz, u0xdz, v0xdz, t0xdz, s0xdz, kappa_avg, &
                              tke_avg, tv, cs, gv, us)
    ! call cpu_clock_begin(Id_clock_setup)
    ! Extrapolate from the vertically reduced grid back to the original layers.
      if (nz == nzc) then
        do k=1,nz+1
          kappa_2d(i,k,j2) = kappa_avg(k)
          if (cs%all_layer_TKE_bug) then
            tke_2d(i,k) = tke(k)
          else
            tke_2d(i,k) = tke_avg(k)
          endif
        enddo
      else
        do k=1,nz+1
          if (kf(k) == 0.0) then
            kappa_2d(i,k,j2) = kappa_avg(kc(k))
            tke_2d(i,k) = tke_avg(kc(k))
          else
            kappa_2d(i,k,j2) = (1.0-kf(k)) * kappa_avg(kc(k)) + kf(k) * kappa_avg(kc(k)+1)
            tke_2d(i,k) = (1.0-kf(k)) * tke_avg(kc(k)) + kf(k) * tke_avg(kc(k)+1)
          endif
        enddo
      endif
    ! call cpu_clock_end(Id_clock_setup)
    else  ! Land points, still inside the i-loop.
      do k=1,nz+1
        kappa_2d(i,k,j2) = 0.0 ; tke_2d(i,k) = 0.0
      enddo
    endif ; enddo ! i-loop

    do k=1,nz+1 ; do i=isb,ieb
      tke_io(i,j,k) = g%mask2dBu(i,j) * tke_2d(i,k)
      kv_io(i,j,k) = ( g%mask2dBu(i,j) * kappa_2d(i,k,j2) ) * cs%Prandtl_turb
    enddo ; enddo
    if (j>=g%jsc) then ; do k=1,nz+1 ; do i=g%isc,g%iec
      ! Set the diffusivities in tracer columns from the values at vertices.
      kappa_io(i,j,k) = g%mask2dT(i,j) * 0.25 * &
                        ((kappa_2d(i-1,k,j2m1) + kappa_2d(i,k,j2)) + &
                         (kappa_2d(i-1,k,j2)   + kappa_2d(i,k,j2m1)))
    enddo ; enddo ; endif

  enddo ! end of J-loop

  if (cs%debug) then
    call hchksum(kappa_io, "kappa", g%HI, scale=us%Z2_T_to_m2_s)
    call bchksum(tke_io, "tke", g%HI, scale=us%Z_to_m**2*us%s_to_T**2)
  endif

  if (cs%id_Kd_shear > 0) call post_data(cs%id_Kd_shear, kappa_io, cs%diag)
  if (cs%id_TKE > 0) call post_data(cs%id_TKE, tke_io, cs%diag)

end subroutine calc_kappa_shear_vertex


!> This subroutine calculates shear-driven diffusivity and TKE in a single column
subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, &
                              dz, u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, &
                              tke_avg, tv, CS, GV, US, I_Ld2_1d, dz_Int_1d)
  type(verticalgrid_type), intent(in)    :: GV !< The ocean's vertical grid structure.
  real, dimension(SZK_(GV)+1), &
                     intent(inout) :: kappa !< The time-weighted average of kappa [Z2 T-1 ~> m2 s-1].
  real, dimension(SZK_(GV)+1), &
                     intent(out)   :: tke  !< The Turbulent Kinetic Energy per unit mass at
                                           !! an interface [Z2 T-2 ~> m2 s-2].
  integer,           intent(in)    :: nzc  !< The number of active layers in the column.
  real,              intent(in)    :: f2   !< The square of the Coriolis parameter [T-2 ~> s-2].
  real,              intent(in)    :: surface_pres  !< The surface pressure [R L2 T-2 ~> Pa].
  real, dimension(SZK_(GV)), &
                     intent(in)    :: dz   !< The layer thickness [Z ~> m].
  real, dimension(SZK_(GV)), &
                     intent(in)    :: u0xdz !< The initial zonal velocity times dz [Z L T-1 ~> m2 s-1].
  real, dimension(SZK_(GV)), &
                     intent(in)    :: v0xdz !< The initial meridional velocity times dz [Z L T-1 ~> m2 s-1].
  real, dimension(SZK_(GV)), &
                     intent(in)    :: T0xdz !< The initial temperature times dz [degC Z ~> degC m].
  real, dimension(SZK_(GV)), &
                     intent(in)    :: S0xdz !< The initial salinity times dz [ppt Z ~> ppt m].
  real, dimension(SZK_(GV)+1), &
                     intent(out)   :: kappa_avg !< The time-weighted average of kappa [Z2 T-1 ~> m2 s-1].
  real, dimension(SZK_(GV)+1), &
                     intent(out)   :: tke_avg  !< The time-weighted average of TKE [Z2 T-2 ~> m2 s-2].
  real,                    intent(in)    :: dt !< Time increment [T ~> s].
  type(thermo_var_ptrs),   intent(in)    :: tv !< A structure containing pointers to any
                                               !! available thermodynamic fields. Absent fields
                                               !! have NULL ptrs.
  type(kappa_shear_cs),    pointer       :: CS !< The control structure returned by a previous
                                               !! call to kappa_shear_init.
  type(unit_scale_type),   intent(in)    :: US !< A dimensional unit scaling type
  real,  dimension(SZK_(GV)+1), &
           optional, intent(out)   :: I_Ld2_1d !< The inverse of the squared mixing length [Z-2 ~> m-2].
  real,  dimension(SZK_(GV)+1), &
           optional, intent(out)   :: dz_Int_1d !< The extent of a finite-volume space surrounding an interface,
                                               !! as used in calculating kappa and TKE [Z ~> m].

  ! Local variables
  real, dimension(nzc) :: &
    u, &        ! The zonal velocity after a timestep of mixing [L T-1 ~> m s-1].
    v, &        ! The meridional velocity after a timestep of mixing [L T-1 ~> m s-1].
    Idz, &      ! The inverse of the distance between TKE points [Z-1 ~> m-1].
    T, &        ! The potential temperature after a timestep of mixing [degC].
    Sal, &      ! The salinity after a timestep of mixing [ppt].
    u_test, v_test, & ! Temporary velocities [L T-1 ~> m s-1].
    T_test, S_test ! Temporary temperatures [degC] and salinities [ppt].

  real, dimension(nzc+1) :: &
    N2, &       ! The squared buoyancy frequency at an interface [T-2 ~> s-2].
    dz_Int, &   ! The extent of a finite-volume space surrounding an interface,
                ! as used in calculating kappa and TKE [Z ~> m].
    i_dz_int, & ! The inverse of the distance between velocity & density points
                ! above and below an interface [Z-1 ~> m-1].  This is used to
                ! calculate N2, shear, and fluxes, and it might differ from
                ! 1/dz_Int, as they have different uses.
    s2, &       ! The squared shear at an interface [T-2 ~> s-2].
    a1, &       ! a1 is the coupling between adjacent interfaces in the TKE,
                ! velocity, and density equations [Z s-1 ~> m s-1] or [Z ~> m]
    c1, &       ! c1 is used in the tridiagonal (and similar) solvers.
    k_src, &    ! The shear-dependent source term in the kappa equation [T-1 ~> s-1].
    kappa_src, & ! The shear-dependent source term in the kappa equation [T-1 ~> s-1].
    kappa_out, & ! The kappa that results from the kappa equation [Z2 T-1 ~> m2 s-1].
    kappa_mid, & ! The average of the initial and predictor estimates of kappa [Z2 T-1 ~> m2 s-1].
    tke_pred, & ! The value of TKE from a predictor step [Z2 T-2 ~> m2 s-2].
    kappa_pred, & ! The value of kappa from a predictor step [Z2 T-1 ~> m2 s-1].
    pressure, & ! The pressure at an interface [R L2 T-2 ~> Pa].
    t_int, &    ! The temperature interpolated to an interface [degC].
    sal_int, &  ! The salinity interpolated to an interface [ppt].
    dbuoy_dt, & ! The partial derivatives of buoyancy with changes in temperature
    dbuoy_ds, & ! and salinity, [Z T-2 degC-1 ~> m s-2 degC-1] and [Z T-2 ppt-1 ~> m s-2 ppt-1].
    i_l2_bdry, &   ! The inverse of the square of twice the harmonic mean
                   ! distance to the top and bottom boundaries [Z-2 ~> m-2].
    k_q, &         ! Diffusivity divided by TKE [T ~> s].
    k_q_tmp, &     ! A temporary copy of diffusivity divided by TKE [T ~> s].
    local_src_avg, & ! The time-integral of the local source [nondim].
    tol_min, & ! Minimum tolerated ksrc for the corrector step [T-1 ~> s-1].
    tol_max, & ! Maximum tolerated ksrc for the corrector step [T-1 ~> s-1].
    tol_chg, & ! The tolerated kappa change integrated over a timestep [nondim].
    dist_from_top, &  ! The distance from the top surface [Z ~> m].
    local_src     ! The sum of all sources of kappa, including kappa_src and
                  ! sources from the elliptic term [T-1 ~> s-1].

  real :: dist_from_bot ! The distance from the bottom surface [Z ~> m].
  real :: b1            ! The inverse of the pivot in the tridiagonal equations.
  real :: bd1           ! A term in the denominator of b1.
  real :: d1            ! 1 - c1 in the tridiagonal equations.
  real :: gR0           ! A conversion factor from Z to pressure, given by Rho_0 times g
                        ! [R L2 T-2 Z-1 ~> kg m-2 s-2].
  real :: g_R0          ! g_R0 is a rescaled version of g/Rho [Z R-1 T-2 ~> m4 kg-1 s-2].
  real :: Norm          ! A factor that normalizes two weights to 1 [Z-2 ~> m-2].
  real :: tol_dksrc     ! Tolerance for the change in the kappa source within an iteration
                        ! relative to the local source [nondim].  This must be greater than 1.
  real :: tol2          ! The tolerance for the change in the kappa source within an iteration
                        ! relative to the average local source over previous iterations [nondim].
  real :: tol_dksrc_low ! The tolerance for the fractional decrease in ksrc
                        ! within an iteration [nondim].  0 < tol_dksrc_low < 1.
  real :: Ri_crit       !   The critical shear Richardson number for shear-
                        ! driven mixing [nondim]. The theoretical value is 0.25.
  real :: dt_rem        !   The remaining time to advance the solution [T ~> s].
  real :: dt_now        !   The time step used in the current iteration [T ~> s].
  real :: dt_wt         !   The fractional weight of the current iteration [nondim].
  real :: dt_test       !   A time-step that is being tested for whether it
                        ! gives acceptably small changes in k_src [T ~> s].
  real :: Idtt          !   Idtt = 1 / dt_test [T-1 ~> s-1].
  real :: dt_inc        !   An increment to dt_test that is being tested [T ~> s].

  real :: k0dt          ! The background diffusivity times the timestep [Z2 ~> m2].
  logical :: valid_dt   ! If true, all levels so far exhibit acceptably small changes in k_src.
  logical :: use_temperature  !  If true, temperature and salinity have been
                        ! allocated and are being used as state variables.
  integer :: ks_kappa, ke_kappa  ! The k-range with nonzero kappas.
  integer :: dt_halvings   ! The number of times that the time-step is halved
                           ! in seeking an acceptable timestep.  If none is
                           ! found, dt_rem*0.5^dt_halvings is used.
  integer :: dt_refinements ! The number of 2-fold refinements that will be used
                           ! to estimate the maximum permitted time step.  I.e.,
                           ! the resolution is 1/2^dt_refinements.
  integer :: k, itt, itt_dt

  ! This calculation of N2 is for debugging only.
  ! real, dimension(SZK_(GV)+1) :: &
  !   N2_debug, & ! A version of N2 for debugging [T-2 ~> s-2]

  ri_crit = cs%Rino_crit
  gr0 = gv%Rho0 * gv%g_Earth
  g_r0 = (us%L_to_Z**2 * gv%g_Earth) / (gv%Rho0)
  k0dt = dt*cs%kappa_0

  tol_dksrc = cs%kappa_src_max_chg
  if (tol_dksrc == 10.0) then
    ! This is equivalent to the expression below, but avoids changes at roundoff for the default value.
    tol_dksrc_low = 0.95
  else
    tol_dksrc_low = (tol_dksrc - 0.5)/tol_dksrc
  endif
  tol2 = 2.0*cs%kappa_tol_err
  dt_refinements = 5 ! Selected so that 1/2^dt_refinements < 1-tol_dksrc_low
  use_temperature = .false. ; if (associated(tv%T)) use_temperature = .true.


  ! Set up Idz as the inverse of layer thicknesses.
  do k=1,nzc ; idz(k) = 1.0 / dz(k) ; enddo
  !   Set up I_dz_int as the inverse of the distance between
  ! adjacent layer centers.
  i_dz_int(1) = 2.0 / dz(1)
  dist_from_top(1) = 0.0
  do k=2,nzc
    i_dz_int(k) = 2.0 / (dz(k-1) + dz(k))
    dist_from_top(k) = dist_from_top(k-1) + dz(k-1)
  enddo
  i_dz_int(nzc+1) = 2.0 / dz(nzc)

  !   Determine the velocities and thicknesses after eliminating massless
  ! layers and applying a time-step of background diffusion.
  if (nzc > 1) then
    a1(2) = k0dt*i_dz_int(2)
    b1 = 1.0 / (dz(1) + a1(2))
    u(1) = b1 * u0xdz(1) ; v(1) = b1 * v0xdz(1)
    t(1) = b1 * t0xdz(1) ; sal(1) = b1 * s0xdz(1)
    c1(2) = a1(2) * b1 ; d1 = dz(1) * b1 ! = 1 - c1
    do k=2,nzc-1
      bd1 = dz(k) + d1*a1(k)
      a1(k+1) = k0dt*i_dz_int(k+1)
      b1 = 1.0 / (bd1 + a1(k+1))
      u(k) = b1 * (u0xdz(k) + a1(k)*u(k-1))
      v(k) = b1 * (v0xdz(k) + a1(k)*v(k-1))
      t(k) = b1 * (t0xdz(k) + a1(k)*t(k-1))
      sal(k) = b1 * (s0xdz(k) + a1(k)*sal(k-1))
      c1(k+1) = a1(k+1) * b1 ; d1 = bd1 * b1 ! d1 = 1 - c1
    enddo
    ! rho or T and S have insulating boundary conditions, u & v use no-slip
    ! bottom boundary conditions (if kappa0 > 0).
    ! For no-slip bottom boundary conditions
    b1 = 1.0 / ((dz(nzc) + d1*a1(nzc)) + k0dt*i_dz_int(nzc+1))
    u(nzc) = b1 * (u0xdz(nzc) + a1(nzc)*u(nzc-1))
    v(nzc) = b1 * (v0xdz(nzc) + a1(nzc)*v(nzc-1))
    ! For insulating boundary conditions
    b1 = 1.0 / (dz(nzc) + d1*a1(nzc))
    t(nzc) = b1 * (t0xdz(nzc) + a1(nzc)*t(nzc-1))
    sal(nzc) = b1 * (s0xdz(nzc) + a1(nzc)*sal(nzc-1))
    do k=nzc-1,1,-1
      u(k) = u(k) + c1(k+1)*u(k+1) ; v(k) = v(k) + c1(k+1)*v(k+1)
      t(k) = t(k) + c1(k+1)*t(k+1) ; sal(k) = sal(k) + c1(k+1)*sal(k+1)
    enddo
  else
    ! This is correct, but probably unnecessary.
    b1 = 1.0 / (dz(1) + k0dt*i_dz_int(2))
    u(1) = b1 * u0xdz(1) ; v(1) = b1 * v0xdz(1)
    b1 = 1.0 / dz(1)
    t(1) = b1 * t0xdz(1) ; sal(1) = b1 * s0xdz(1)
  endif

  ! This uses half the harmonic mean of thicknesses to provide two estimates
  ! of the boundary between cells, and the inverse of the harmonic mean to
  ! weight the two estimates.  The net effect is that interfaces around thin
  ! layers have thin cells, and the total thickness adds up properly.
  ! The top- and bottom- interfaces have zero thickness, consistent with
  ! adding additional zero thickness layers.
  dz_int(1) = 0.0 ; dz_int(2) = dz(1)
  do k=2,nzc-1
    norm = 1.0 / (dz(k)*(dz(k-1)+dz(k+1)) + 2.0*dz(k-1)*dz(k+1))
    dz_int(k) = dz_int(k) + dz(k) * ( ((dz(k)+dz(k+1)) * dz(k-1)) * norm)
    dz_int(k+1) = dz(k) * ( ((dz(k-1)+dz(k)) * dz(k+1)) * norm)
  enddo
  dz_int(nzc) = dz_int(nzc) + dz(nzc) ; dz_int(nzc+1) = 0.0

  dist_from_bot = 0.0
  do k=nzc,2,-1
    dist_from_bot = dist_from_bot + dz(k)
    i_l2_bdry(k) = (dist_from_top(k) + dist_from_bot)**2 / &
                   (dist_from_top(k) * dist_from_bot)**2
  enddo

  ! Calculate thermodynamic coefficients and an initial estimate of N2.
  if (use_temperature) then
    pressure(1) = surface_pres
    do k=2,nzc
      pressure(k) = pressure(k-1) + gr0*dz(k-1)
      t_int(k) = 0.5*(t(k-1) + t(k))
      sal_int(k) = 0.5*(sal(k-1) + sal(k))
    enddo
    call calculate_density_derivs(t_int, sal_int, pressure, dbuoy_dt, dbuoy_ds, &
                                  tv%eqn_of_state, (/2,nzc/), scale=-g_r0 )
  else
    do k=1,nzc+1 ; dbuoy_dt(k) = -g_r0 ; dbuoy_ds(k) = 0.0 ; enddo
  endif

  ! N2_debug(1) = 0.0 ; N2_debug(nzc+1) = 0.0
  ! do K=2,nzc
  !   N2_debug(K) = max((dbuoy_dT(K) * (T0xdz(k-1)*Idz(k-1) - T0xdz(k)*Idz(k)) + &
  !                      dbuoy_dS(K) * (S0xdz(k-1)*Idz(k-1) - S0xdz(k)*Idz(k))) * &
  !                      I_dz_int(K), 0.0)
  ! enddo

  ! This call just calculates N2 and S2.
  call calculate_projected_state(kappa, u, v, t, sal, 0.0, nzc, dz, i_dz_int, &
                                 dbuoy_dt, dbuoy_ds, u, v, t, sal, gv, us, &
                                 n2=n2, s2=s2, vel_underflow=cs%vel_underflow)
! ----------------------------------------------------
! Iterate
! ----------------------------------------------------
  dt_rem = dt
  do k=1,nzc+1
    k_q(k) = 0.0
    kappa_avg(k) = 0.0 ; tke_avg(k) = 0.0
    local_src_avg(k) = 0.0
    ! Use the grid spacings to scale errors in the source.
    if ( dz_int(k) > 0.0 ) &
      local_src_avg(k) = 0.1 * k0dt * i_dz_int(k) / dz_int(k)
  enddo

! call cpu_clock_end(id_clock_setup)

! do itt=1,CS%max_RiNo_it
  do itt=1,cs%max_KS_it

! ----------------------------------------------------
! Calculate new values of u, v, rho, N^2 and S.
! ----------------------------------------------------

  ! call cpu_clock_begin(id_clock_KQ)
    call find_kappa_tke(n2, s2, kappa, idz, dz_int, i_l2_bdry, f2, &
                        nzc, cs, gv, us, k_q, tke, kappa_out, kappa_src, local_src)
  ! call cpu_clock_end(id_clock_KQ)

  ! call cpu_clock_begin(id_clock_avg)
    ! Determine the range of non-zero values of kappa_out.
    ks_kappa = gv%ke+1 ; ke_kappa = 0
    do k=2,nzc ; if (kappa_out(k) > 0.0) then
      ks_kappa = k ; exit
    endif ; enddo
    do k=nzc,ks_kappa,-1 ; if (kappa_out(k) > 0.0) then
      ke_kappa = k ; exit
    endif ; enddo
    if (ke_kappa == nzc) kappa_out(nzc+1) = 0.0
  ! call cpu_clock_end(id_clock_avg)

    ! Determine how long to use this value of kappa (dt_now).

  ! call cpu_clock_begin(id_clock_project)
    if ((ke_kappa < ks_kappa) .or. (itt==cs%max_KS_it)) then
      dt_now = dt_rem
    else
      ! Limit dt_now so that |k_src(k)-kappa_src(k)| < tol * local_src(k)
      dt_test = dt_rem
      do k=2,nzc
        tol_max(k) = kappa_src(k) + tol_dksrc * local_src(k)
        tol_min(k) = kappa_src(k) - tol_dksrc_low * local_src(k)
        tol_chg(k) = tol2 * local_src_avg(k)
      enddo

      do itt_dt=1,(cs%max_KS_it+1-itt)/2
        !   The maximum number of times that the time-step is halved in
        ! seeking an acceptable timestep is reduced with each iteration,
        ! so that as the maximum number of iterations is approached, the
        ! whole remaining timestep is used.  Typically, an acceptable
        ! timestep is found long before the minimum is reached, so the
        ! value of max_KS_it may be unimportant, especially if it is large
        ! enough.
        call calculate_projected_state(kappa_out, u, v, t, sal, 0.5*dt_test, nzc, dz, i_dz_int, &
                                       dbuoy_dt, dbuoy_ds, u_test, v_test, t_test, s_test, &
                                       gv, us, n2, s2, ks_int=ks_kappa, ke_int=ke_kappa, &
                                       vel_underflow=cs%vel_underflow)
        valid_dt = .true.
        idtt = 1.0 / dt_test
        do k=max(ks_kappa-1,2),min(ke_kappa+1,nzc)
          if (n2(k) < ri_crit * s2(k)) then ! Equivalent to Ri < Ri_crit.
            k_src(k) = (2.0 * cs%Shearmix_rate * sqrt(s2(k))) * &
                       ((ri_crit*s2(k) - n2(k)) / (ri_crit*s2(k) + cs%FRi_curvature*n2(k)))
            if (cs%restrictive_tolerance_check) then
              if ((k_src(k) > min(tol_max(k), kappa_src(k) + idtt*tol_chg(k))) .or. &
                  (k_src(k) < max(tol_min(k), kappa_src(k) - idtt*tol_chg(k)))) then
                valid_dt = .false. ; exit
              endif
            else
              if ((k_src(k) > max(tol_max(k), kappa_src(k) + idtt*tol_chg(k))) .or. &
                  (k_src(k) < min(tol_min(k), kappa_src(k) - idtt*tol_chg(k)))) then
                valid_dt = .false. ; exit
              endif
            endif
          else
            if (0.0 < min(tol_min(k), kappa_src(k) - idtt*tol_chg(k))) then
              valid_dt = .false. ; k_src(k) = 0.0 ; exit
            endif
          endif
        enddo

        if (valid_dt) exit
        dt_test = 0.5*dt_test
      enddo
      if ((dt_test < dt_rem) .and. valid_dt) then
        dt_inc = 0.5*dt_test
        do itt_dt=1,dt_refinements
          call calculate_projected_state(kappa_out, u, v, t, sal, 0.5*(dt_test+dt_inc), &
                   nzc, dz, i_dz_int, dbuoy_dt, dbuoy_ds, u_test, v_test, t_test, s_test, &
                   gv, us, n2, s2, ks_int=ks_kappa, ke_int=ke_kappa, vel_underflow=cs%vel_underflow)
          valid_dt = .true.
          idtt = 1.0 / (dt_test+dt_inc)
          do k=max(ks_kappa-1,2),min(ke_kappa+1,nzc)
            if (n2(k) < ri_crit * s2(k)) then ! Equivalent to Ri < Ri_crit.
              k_src(k) = (2.0 * cs%Shearmix_rate * sqrt(s2(k))) * &
                         ((ri_crit*s2(k) - n2(k)) / (ri_crit*s2(k) + cs%FRi_curvature*n2(k)))
              if ((k_src(k) > max(tol_max(k), kappa_src(k) + idtt*tol_chg(k))) .or. &
                  (k_src(k) < min(tol_min(k), kappa_src(k) - idtt*tol_chg(k)))) then
                valid_dt = .false. ; exit
              endif
            else
              if (0.0 < min(tol_min(k), kappa_src(k) - idtt*tol_chg(k))) then
                valid_dt = .false. ; k_src(k) = 0.0 ; exit
              endif
            endif
          enddo

          if (valid_dt) dt_test = dt_test + dt_inc
          dt_inc = 0.5*dt_inc
        enddo
      else
        dt_inc = 0.0
      endif

      dt_now = min(dt_test*(1.0+cs%kappa_tol_err)+dt_inc, dt_rem)
      do k=2,nzc
        local_src_avg(k) = local_src_avg(k) + dt_now * local_src(k)
      enddo
    endif  ! Are all the values of kappa_out 0?
  ! call cpu_clock_end(id_clock_project)

    ! The state has already been projected forward. Now find new values of kappa.

    if (ke_kappa < ks_kappa) then
      ! There is no mixing now, and will not be again.
    ! call cpu_clock_begin(id_clock_avg)
      dt_wt = dt_rem / dt ; dt_rem = 0.0
      do k=1,nzc+1
        kappa_mid(k) = 0.0
        ! This would be here but does nothing.
        ! kappa_avg(K) = kappa_avg(K) + kappa_mid(K)*dt_wt
        tke_avg(k) = tke_avg(k) + dt_wt*tke(k)
      enddo
    ! call cpu_clock_end(id_clock_avg)
    else
    ! call cpu_clock_begin(id_clock_project)
      call calculate_projected_state(kappa_out, u, v, t, sal, dt_now, nzc, dz, i_dz_int, &
                                     dbuoy_dt, dbuoy_ds, u_test, v_test, t_test, s_test, &
                                     gv, us, n2=n2, s2=s2, ks_int=ks_kappa, ke_int=ke_kappa, &
                                     vel_underflow=cs%vel_underflow)
    ! call cpu_clock_end(id_clock_project)

    ! call cpu_clock_begin(id_clock_KQ)
      do k=1,nzc+1 ; k_q_tmp(k) = k_q(k) ; enddo
      call find_kappa_tke(n2, s2, kappa_out, idz, dz_int, i_l2_bdry, f2, &
                          nzc, cs, gv, us, k_q_tmp, tke_pred, kappa_pred)
    ! call cpu_clock_end(id_clock_KQ)

      ks_kappa = gv%ke+1 ; ke_kappa = 0
      do k=1,nzc+1
        kappa_mid(k) = 0.5*(kappa_out(k) + kappa_pred(k))
        if ((kappa_mid(k) > 0.0) .and. (k<ks_kappa)) ks_kappa = k
        if (kappa_mid(k) > 0.0) ke_kappa = k
      enddo

    ! call cpu_clock_begin(id_clock_project)
      call calculate_projected_state(kappa_mid, u, v, t, sal, dt_now, nzc, dz, i_dz_int, &
                                     dbuoy_dt, dbuoy_ds, u_test, v_test, t_test, s_test, &
                                     gv, us, n2=n2, s2=s2, ks_int=ks_kappa, ke_int=ke_kappa, &
                                     vel_underflow=cs%vel_underflow)
    ! call cpu_clock_end(id_clock_project)

    ! call cpu_clock_begin(id_clock_KQ)
      call find_kappa_tke(n2, s2, kappa_out, idz, dz_int, i_l2_bdry, f2, &
                          nzc, cs, gv, us, k_q, tke_pred, kappa_pred)
    ! call cpu_clock_end(id_clock_KQ)

    ! call cpu_clock_begin(id_clock_avg)
      dt_wt = dt_now / dt ; dt_rem = dt_rem - dt_now
      do k=1,nzc+1
        kappa_mid(k) = 0.5*(kappa_out(k) + kappa_pred(k))
        kappa_avg(k) = kappa_avg(k) + kappa_mid(k)*dt_wt
        tke_avg(k) = tke_avg(k) + dt_wt*0.5*(tke_pred(k) + tke(k))
        kappa(k) = kappa_pred(k) ! First guess for the next iteration.
      enddo
    ! call cpu_clock_end(id_clock_avg)
    endif

    if (dt_rem > 0.0) then
      ! Update the values of u, v, T, Sal, N2, and S2 for the next iteration.
    ! call cpu_clock_begin(id_clock_project)
      call calculate_projected_state(kappa_mid, u, v, t, sal, dt_now, nzc, &
                                     dz, i_dz_int, dbuoy_dt, dbuoy_ds, u, v, t, sal, &
                                     gv, us, n2, s2, vel_underflow=cs%vel_underflow)
    ! call cpu_clock_end(id_clock_project)
    endif

    if (dt_rem <= 0.0) exit

  enddo ! end itt loop

  if (present(i_ld2_1d)) then
    do k=1,gv%ke+1 ; i_ld2_1d(k) = 0.0 ; enddo
    do k=2,nzc ; if (tke(k) > 0.0) &
      i_ld2_1d(k) = i_l2_bdry(k) + (n2(k) / cs%lambda**2 + f2) / tke(k)
    enddo
  endif
  if (present(dz_int_1d)) then
    do k=1,nzc+1 ; dz_int_1d(k) = dz_int(k) ; enddo
    do k=nzc+2,gv%ke ; dz_int_1d(k) = 0.0 ; enddo
  endif

end subroutine kappa_shear_column

!>   This subroutine calculates the velocities, temperature and salinity that
!! the water column will have after mixing for dt with diffusivities kappa.  It
!! may also calculate the projected buoyancy frequency and shear.
subroutine calculate_projected_state(kappa, u0, v0, T0, S0, dt, nz, &
                                     dz, I_dz_int, dbuoy_dT, dbuoy_dS, &
                                     u, v, T, Sal, GV, US, N2, S2, ks_int, ke_int, vel_underflow)
  integer,               intent(in)    :: nz  !< The number of layers (after eliminating massless
                                              !! layers?).
  real, dimension(nz+1), intent(in)    :: kappa !< The diapycnal diffusivity at interfaces,
                                              !! [Z2 T-1 ~> m2 s-1].
  real, dimension(nz),   intent(in)    :: u0  !< The initial zonal velocity [L T-1 ~> m s-1].
  real, dimension(nz),   intent(in)    :: v0  !< The initial meridional velocity [L T-1 ~> m s-1].
  real, dimension(nz),   intent(in)    :: T0  !< The initial temperature [degC].
  real, dimension(nz),   intent(in)    :: S0  !< The initial salinity [ppt].
  real, dimension(nz),   intent(in)    :: dz  !< The grid spacing of layers [Z ~> m].
  real, dimension(nz+1), intent(in)    :: I_dz_int !< The inverse of the layer's thicknesses
                                              !! [Z-1 ~> m-1].
  real, dimension(nz+1), intent(in)    :: dbuoy_dT !< The partial derivative of buoyancy with
                                              !! temperature [Z T-2 degC-1 ~> m s-2 degC-1].
  real, dimension(nz+1), intent(in)    :: dbuoy_dS !< The partial derivative of buoyancy with
                                              !! salinity [Z T-2 ppt-1 ~> m s-2 ppt-1].
  real,                  intent(in)    :: dt  !< The time step [T ~> s].
  real, dimension(nz),   intent(inout) :: u   !< The zonal velocity after dt [L T-1 ~> m s-1].
  real, dimension(nz),   intent(inout) :: v   !< The meridional velocity after dt [L T-1 ~> m s-1].
  real, dimension(nz),   intent(inout) :: T   !< The temperature after dt [degC].
  real, dimension(nz),   intent(inout) :: Sal !< The salinity after dt [ppt].
  type(verticalgrid_type), intent(in)  :: GV  !< The ocean's vertical grid structure.
  type(unit_scale_type), intent(in)    :: US  !< A dimensional unit scaling type
  real, dimension(nz+1), optional, &
                         intent(inout) :: N2  !< The buoyancy frequency squared at interfaces [T-2 ~> s-2].
  real, dimension(nz+1), optional, &
                         intent(inout) :: S2  !< The squared shear at interfaces [T-2 ~> s-2].
  integer, optional,     intent(in)    :: ks_int !< The topmost k-index with a non-zero diffusivity.
  integer, optional,     intent(in)    :: ke_int !< The bottommost k-index with a non-zero
                                              !! diffusivity.
  real,    optional,     intent(in)    :: vel_underflow !< If present and true, any velocities that
                                              !! are smaller in magnitude than this value are
                                              !! set to 0 [L T-1 ~> m s-1].

  ! Local variables
  real, dimension(nz+1) :: c1
  real :: L2_to_Z2       ! A conversion factor from horizontal length units to vertical depth
                         ! units squared [Z2 s2 T-2 m-2 ~> 1].
  real :: underflow_vel  ! Velocities smaller in magnitude than underflow_vel are set to 0 [L T-1 ~> m s-1].
  real :: a_a, a_b, b1, d1, bd1, b1nz_0
  integer :: k, ks, ke

  ks = 1 ; ke = nz
  if (present(ks_int)) ks = max(ks_int-1,1)
  if (present(ke_int)) ke = min(ke_int,nz)
  underflow_vel = 0.0 ; if (present(vel_underflow)) underflow_vel = vel_underflow

  if (ks > ke) return

  if (dt > 0.0) then
    a_b = dt*(kappa(ks+1)*i_dz_int(ks+1))
    b1 = 1.0 / (dz(ks) + a_b)
    c1(ks+1) = a_b * b1 ; d1 = dz(ks) * b1 ! = 1 - c1

    u(ks) = (b1 * dz(ks))*u0(ks) ; v(ks) = (b1 * dz(ks))*v0(ks)
    t(ks) = (b1 * dz(ks))*t0(ks) ; sal(ks) = (b1 * dz(ks))*s0(ks)
    do k=ks+1,ke-1
      a_a = a_b
      a_b = dt*(kappa(k+1)*i_dz_int(k+1))
      bd1 = dz(k) + d1*a_a
      b1 = 1.0 / (bd1 + a_b)
      c1(k+1) = a_b * b1 ; d1 = bd1 * b1 ! d1 = 1 - c1

      u(k) = b1 * (dz(k)*u0(k) + a_a*u(k-1))
      v(k) = b1 * (dz(k)*v0(k) + a_a*v(k-1))
      t(k) = b1 * (dz(k)*t0(k) + a_a*t(k-1))
      sal(k) = b1 * (dz(k)*s0(k) + a_a*sal(k-1))
    enddo
    !   T and S have insulating boundary conditions, u & v use no-slip
    ! bottom boundary conditions at the solid bottom.

    ! For insulating boundary conditions or mixing simply stopping, use...
    a_a = a_b
    b1 = 1.0 / (dz(ke) + d1*a_a)
    t(ke) = b1 * (dz(ke)*t0(ke) + a_a*t(ke-1))
    sal(ke) = b1 * (dz(ke)*s0(ke) + a_a*sal(ke-1))

    !   There is no distinction between the effective boundary conditions for
    ! tracers and velocities if the mixing is separated from the bottom, but if
    ! the mixing goes all the way to the bottom, use no-slip BCs for velocities.
    if (ke == nz) then
      a_b = dt*(kappa(nz+1)*i_dz_int(nz+1))
      b1nz_0 = 1.0 / ((dz(nz) + d1*a_a) + a_b)
    else
      b1nz_0 = b1
    endif
    u(ke) = b1nz_0 * (dz(ke)*u0(ke) + a_a*u(ke-1))
    v(ke) = b1nz_0 * (dz(ke)*v0(ke) + a_a*v(ke-1))
    if (abs(u(ke)) < underflow_vel) u(ke) = 0.0
    if (abs(v(ke)) < underflow_vel) v(ke) = 0.0

    do k=ke-1,ks,-1
      u(k) = u(k) + c1(k+1)*u(k+1)
      v(k) = v(k) + c1(k+1)*v(k+1)
      if (abs(u(k)) < underflow_vel) u(k) = 0.0
      if (abs(v(k)) < underflow_vel) v(k) = 0.0
      t(k) = t(k) + c1(k+1)*t(k+1)
      sal(k) = sal(k) + c1(k+1)*sal(k+1)
    enddo
  else ! dt <= 0.0
    do k=1,nz
      u(k) = u0(k) ; v(k) = v0(k) ; t(k) = t0(k) ; sal(k) = s0(k)
      if (abs(u(k)) < underflow_vel) u(k) = 0.0
      if (abs(v(k)) < underflow_vel) v(k) = 0.0
    enddo
  endif

  if (present(s2)) then
    ! L2_to_Z2 = US%m_to_Z**2 * US%T_to_s**2
    l2_to_z2 = us%L_to_Z**2
    s2(1) = 0.0 ; s2(nz+1) = 0.0
    if (ks > 1) &
      s2(ks) = ((u(ks)-u0(ks-1))**2 + (v(ks)-v0(ks-1))**2) * (l2_to_z2*i_dz_int(ks)**2)
    do k=ks+1,ke
      s2(k) = ((u(k)-u(k-1))**2 + (v(k)-v(k-1))**2) * (l2_to_z2*i_dz_int(k)**2)
    enddo
    if (ke<nz) &
      s2(ke+1) = ((u0(ke+1)-u(ke))**2 + (v0(ke+1)-v(ke))**2) * (l2_to_z2*i_dz_int(ke+1)**2)
  endif

  if (present(n2)) then
    n2(1) = 0.0 ; n2(nz+1) = 0.0
    if (ks > 1) &
      n2(ks) = max(0.0, i_dz_int(ks) * &
        (dbuoy_dt(ks) * (t0(ks-1)-t(ks)) + dbuoy_ds(ks) * (s0(ks-1)-sal(ks))))
    do k=ks+1,ke
      n2(k) = max(0.0, i_dz_int(k) * &
        (dbuoy_dt(k) * (t(k-1)-t(k)) + dbuoy_ds(k) * (sal(k-1)-sal(k))))
    enddo
    if (ke<nz) &
      n2(ke+1) = max(0.0, i_dz_int(ke+1) * &
        (dbuoy_dt(ke+1) * (t(ke)-t0(ke+1)) + dbuoy_ds(ke+1) * (sal(ke)-s0(ke+1))))
  endif

end subroutine calculate_projected_state

!> This subroutine calculates new, consistent estimates of TKE and kappa.
subroutine find_kappa_tke(N2, S2, kappa_in, Idz, dz_Int, I_L2_bdry, f2, &
                          nz, CS, GV, US, K_Q, tke, kappa, kappa_src, local_src)
  integer,               intent(in)    :: nz  !< The number of layers to work on.
  real, dimension(nz+1), intent(in)    :: N2  !< The buoyancy frequency squared at interfaces [T-2 ~> s-2].
  real, dimension(nz+1), intent(in)    :: S2  !< The squared shear at interfaces [T-2 ~> s-2].
  real, dimension(nz+1), intent(in)    :: kappa_in  !< The initial guess at the diffusivity
                                              !! [Z2 T-1 ~> m2 s-1].
  real, dimension(nz+1), intent(in)    :: dz_Int !< The thicknesses associated with interfaces
                                              !! [Z-1 ~> m-1].
  real, dimension(nz+1), intent(in)    :: I_L2_bdry !< The inverse of the squared distance to
                                              !! boundaries [Z-2 ~> m-2].
  real, dimension(nz),   intent(in)    :: Idz !< The inverse grid spacing of layers [Z-1 ~> m-1].
  real,                  intent(in)    :: f2  !< The squared Coriolis parameter [T-2 ~> s-2].
  type(kappa_shear_cs),  pointer       :: CS  !< A pointer to this module's control structure.
  type(verticalgrid_type), intent(in)  :: GV  !< The ocean's vertical grid structure.
  type(unit_scale_type), intent(in)    :: US  !< A dimensional unit scaling type
  real, dimension(nz+1), intent(inout) :: K_Q !< The shear-driven diapycnal diffusivity divided by
                                              !! the turbulent kinetic energy per unit mass at
                                              !! interfaces [T ~> s].
  real, dimension(nz+1), intent(out)   :: tke !< The turbulent kinetic energy per unit mass at
                                              !! interfaces [Z2 T-2 ~> m2 s-2].
  real, dimension(nz+1), intent(out)   :: kappa  !< The diapycnal diffusivity at interfaces
                                              !! [Z2 T-1 ~> m2 s-1].
  real, dimension(nz+1), optional, &
                         intent(out)   :: kappa_src !< The source term for kappa [T-1 ~> s-1].
  real, dimension(nz+1), optional, &
                         intent(out)   :: local_src !< The sum of all local sources for kappa,
                                              !! [T-1 ~> s-1].
  ! This subroutine calculates new, consistent estimates of TKE and kappa.

  ! Local variables
  real, dimension(nz) :: &
    aQ, &       ! aQ is the coupling between adjacent interfaces in the TKE equations [Z T-1 ~> m s-1].
    dQdz        ! Half the partial derivative of TKE with depth [Z T-2 ~> m s-2].
  real, dimension(nz+1) :: &
    dK, &         ! The change in kappa [Z2 T-1 ~> m2 s-1].
    dQ, &         ! The change in TKE [Z2 T-2 ~> m2 s-2].
    cQ, cK, &     ! cQ and cK are the upward influences in the tridiagonal and
                  ! hexadiagonal solvers for the TKE and kappa equations [nondim].
    i_ld2, &      ! 1/Ld^2, where Ld is the effective decay length scale for kappa [Z-2 ~> m-2].
    tke_decay, &  ! The local TKE decay rate [T-1 ~> s-1].
    k_src, &      ! The source term in the kappa equation [T-1 ~> s-1].
    dqmdk, &      ! With Newton's method the change in dQ(k-1) due to dK(k) [T ~> s].
    dkdq, &       ! With Newton's method the change in dK(k) due to dQ(k) [T-1 ~> s-1].
    e1            ! The fractional change in a layer TKE due to a change in the
                  ! TKE of the layer above when all the kappas below are 0.
                  ! e1 is nondimensional, and 0 < e1 < 1.
  real :: tke_src       ! The net source of TKE due to mixing against the shear
                        ! and stratification [Z2 T-3 ~> m2 s-3].  (For convenience,
                        ! a term involving the non-dissipation of q0 is also
                        ! included here.)
  real :: bQ            ! The inverse of the pivot in the tridiagonal equations [T Z-1 ~> s m-1].
  real :: bK            ! The inverse of the pivot in the tridiagonal equations [Z-1 ~> m-1].
  real :: bQd1          ! A term in the denominator of bQ [Z T-1 ~> m s-1].
  real :: bKd1          ! A term in the denominator of bK [Z ~> m].
  real :: cQcomp, cKcomp ! 1 - cQ or 1 - cK in the tridiagonal equations.
  real :: c_s2          !   The coefficient for the decay of TKE due to
                        ! shear (i.e. proportional to |S|*tke), nondimensional.
  real :: c_n2          !   The coefficient for the decay of TKE due to
                        ! stratification (i.e. proportional to N*tke) [nondim].
  real :: Ri_crit       !   The critical shear Richardson number for shear-
                        ! driven mixing. The theoretical value is 0.25.
  real :: q0            !   The background level of TKE [Z2 T-2 ~> m2 s-2].
  real :: Ilambda2      ! 1.0 / CS%lambda**2 [nondim]
  real :: TKE_min       !   The minimum value of shear-driven TKE that can be
                        ! solved for [Z2 T-2 ~> m2 s-2].
  real :: kappa0        ! The background diapycnal diffusivity [Z2 T-1 ~> m2 s-1].
  real :: kappa_trunc   ! Diffusivities smaller than this are rounded to 0 [Z2 T-1 ~> m2 s-1].

  real :: eden1, eden2, I_eden, ome  ! Variables used in calculating e1.
  real :: diffusive_src ! The diffusive source in the kappa equation [Z T-1 ~> m s-1].
  real :: chg_by_k0     ! The value of k_src that leads to an increase of
                        ! kappa_0 if only the diffusive term is a sink [T-1 ~> s-1].

  real :: kappa_mean    ! A mean value of kappa [Z2 T-1 ~> m2 s-1].
  real :: Newton_test   ! The value of relative error that will cause the next
                        ! iteration to use Newton's method.
  ! Temporary variables used in the Newton's method iterations.
  real :: decay_term_k  ! The decay term in the diffusivity equation
  real :: decay_term_Q  ! The decay term in the TKE equation - proportional to [T-1 ~> s-1]
  real :: I_Q           ! The inverse of TKE [T2 Z-2 ~> s2 m-2]
  real :: kap_src
  real :: v1            ! A temporary variable proportional to [T-1 ~> s-1]
  real :: v2
  real :: tol_err        ! The tolerance for max_err that determines when to
                         ! stop iterating.
  real :: Newton_err     ! The tolerance for max_err that determines when to
                         ! start using Newton's method.  Empirically, an initial
                         ! value of about 0.2 seems to be most efficient.
  real, parameter :: roundoff = 1.0e-16 ! A negligible fractional change in TKE.
                         ! This could be larger but performance gains are small.

  logical :: tke_noflux_bottom_BC = .false. ! Specify the boundary conditions
  logical :: tke_noflux_top_BC = .false.    ! that are applied to the TKE eqns.
  logical :: do_Newton    ! If .true., use Newton's method for the next iteration.
  logical :: abort_Newton ! If .true., an Newton's method has encountered a 0
                          ! pivot, and should not have been used.
  logical :: was_Newton   ! The value of do_Newton before checking convergence.
  logical :: within_tolerance ! If .true., all points are within tolerance to
                          ! enable this subroutine to return.
  integer :: ks_src, ke_src ! The range indices that have nonzero k_src.
  integer :: ks_kappa, ke_kappa, ke_tke   ! The ranges of k-indices that are or
  integer :: ks_kappa_prev, ke_kappa_prev ! were being worked on.
  integer :: itt, k, k2

  ! These variables are used only for debugging.
  logical, parameter :: debug_soln = .false.
  real :: K_err_lin, Q_err_lin, TKE_src_norm
  real, dimension(nz+1) :: &
    I_Ld2_debug, & ! A separate version of I_Ld2 for debugging [Z-2 ~> m-2].
    kappa_prev, & ! The value of kappa at the start of the current iteration [Z2 T-1 ~> m2 s-1].
    TKE_prev   ! The value of TKE at the start of the current iteration [Z2 T-2 ~> m2 s-2].

  c_n2 = cs%C_N**2 ; c_s2 = cs%C_S**2
  q0 = cs%TKE_bg ; kappa0 = cs%kappa_0
  tke_min = max(cs%TKE_bg, 1.0e-20*us%m_to_Z**2*us%T_to_s**2)
  ri_crit = cs%Rino_crit
  ilambda2 = 1.0 / cs%lambda**2
  kappa_trunc = cs%kappa_trunc
  do_newton = .false. ; abort_newton = .false.
  tol_err = cs%kappa_tol_err
  newton_err = 0.2     ! This initial value may be automatically reduced later.

  ks_kappa = 2 ; ke_kappa = nz ; ks_kappa_prev = 2 ; ke_kappa_prev = nz

  ke_src = 0 ; ks_src = nz+1
  do k=2,nz
    if (n2(k) < ri_crit * s2(k)) then ! Equivalent to Ri < Ri_crit.
!       Ri = N2(K) / S2(K)
!       k_src(K) = (2.0 * CS%Shearmix_rate * sqrt(S2(K))) * &
!                  ((Ri_crit - Ri) / (Ri_crit + CS%FRi_curvature*Ri))
      k_src(k) = (2.0 * cs%Shearmix_rate * sqrt(s2(k))) * &
                 ((ri_crit*s2(k) - n2(k)) / (ri_crit*s2(k) + cs%FRi_curvature*n2(k)))
      ke_src = k
      if (ks_src > k) ks_src = k
    else
      k_src(k) = 0.0
    endif
  enddo

  ! If there is no source anywhere, return kappa(K) = 0.
  if (ks_src > ke_src) then
    do k=1,nz+1
      kappa(k) = 0.0 ; k_q(k) = 0.0 ; tke(k) = tke_min
    enddo
    if (present(kappa_src)) then ; do k=1,nz+1 ; kappa_src(k) = 0.0 ; enddo ; endif
    if (present(local_src)) then ; do k=1,nz+1 ; local_src(k) = 0.0 ; enddo ; endif
    return
  endif

  do k=1,nz+1
    kappa(k) = kappa_in(k)
!     TKE_decay(K) = c_n*sqrt(N2(K)) + c_s*sqrt(S2(K)) ! The expression in JHL.
    tke_decay(k) = sqrt(c_n2*n2(k) + c_s2*s2(k))
    if ((kappa(k) > 0.0) .and. (k_q(k) > 0.0)) then
      tke(k) = kappa(k) / k_q(k) ! Perhaps take the max with TKE_min
    else
      tke(k) = tke_min
    endif
  enddo
  ! Apply boundary conditions to kappa.
  kappa(1) = 0.0 ; kappa(nz+1) = 0.0

  ! Calculate the term (e1) that allows changes in TKE to be calculated quickly
  ! below the deepest nonzero value of kappa.  If kappa = 0, below interface
  ! k-1, the final changes in TKE are related by dQ(K+1) = e1(K+1)*dQ(K).
  eden2 = kappa0 * idz(nz)
  if (tke_noflux_bottom_bc) then
    eden1 = dz_int(nz+1)*tke_decay(nz+1)
    i_eden = 1.0 / (eden2 + eden1)
    e1(nz+1) = eden2 * i_eden ; ome = eden1 * i_eden
  else
    e1(nz+1) = 0.0 ; ome = 1.0
  endif
  do k=nz,2,-1
    eden1 = dz_int(k)*tke_decay(k) + ome * eden2
    eden2 = kappa0 * idz(k-1)
    i_eden = 1.0 / (eden2 + eden1)
    e1(k) = eden2 * i_eden ; ome = eden1 * i_eden ! = 1-e1
  enddo
  e1(1) = 0.0


  ! Iterate here to convergence to within some tolerance of order tol_err.
  do itt=1,cs%max_RiNo_it

  ! ----------------------------------------------------
  ! Calculate TKE
  ! ----------------------------------------------------

    if (debug_soln) then ; do k=1,nz+1 ; kappa_prev(k) = kappa(k) ; tke_prev(k) = tke(k) ; enddo ; endif

    if (.not.do_newton) then
      !   Use separate steps of the TKE and kappa equations, that are
      ! explicit in the nonlinear source terms, implicit in a linearized
      ! version of the nonlinear sink terms, and implicit in the linear
      ! terms.

      ke_tke = max(ke_kappa,ke_kappa_prev)+1
      ! aQ is the coupling between adjacent interfaces [Z T-1 ~> m s-1].
      do k=1,min(ke_tke,nz)
        aq(k) = (0.5*(kappa(k)+kappa(k+1)) + kappa0) * idz(k)
      enddo
      dq(1) = -tke(1)
      if (tke_noflux_top_bc) then
        tke_src = kappa0*s2(1) + q0 * tke_decay(1) ! Uses that kappa(1) = 0
        bqd1 = dz_int(1) * tke_decay(1)
        bq = 1.0 / (bqd1 +  aq(1))
        tke(1) = bq * (dz_int(1)*tke_src)
        cq(2) = aq(1) * bq ; cqcomp = bqd1 * bq ! = 1 - cQ
      else
        tke(1) = q0 ; cq(2) = 0.0 ; cqcomp = 1.0
      endif
      do k=2,ke_tke-1
        dq(k) = -tke(k)
        tke_src = (kappa(k) + kappa0)*s2(k) + q0*tke_decay(k)
        bqd1 = dz_int(k)*(tke_decay(k) + n2(k)*k_q(k)) + cqcomp*aq(k-1)
        bq = 1.0 / (bqd1 + aq(k))
        tke(k) = bq * (dz_int(k)*tke_src + aq(k-1)*tke(k-1))
        cq(k+1) = aq(k) * bq ; cqcomp = bqd1 * bq ! = 1 - cQ
      enddo
      if ((ke_tke == nz+1) .and. .not.(tke_noflux_bottom_bc)) then
        tke(nz+1) = tke_min
        dq(nz+1) = 0.0
      else
        k = ke_tke
        tke_src = kappa0*s2(k) + q0*tke_decay(k) ! Uses that kappa(ke_tke) = 0
        if (k == nz+1) then
          dq(k) = -tke(k)
          bq = 1.0 / (dz_int(k)*tke_decay(k) + cqcomp*aq(k-1))
          tke(k) = max(tke_min, bq * (dz_int(k)*tke_src + aq(k-1)*tke(k-1)))
          dq(k) = tke(k) + dq(k)
        else
          bq = 1.0 / ((dz_int(k)*tke_decay(k) + cqcomp*aq(k-1)) + aq(k))
          cq(k+1) = aq(k) * bq
          ! Account for all changes deeper in the water column.
          dq(k) = -tke(k)
          tke(k) = max((bq * (dz_int(k)*tke_src + aq(k-1)*tke(k-1)) + &
                        cq(k+1)*(tke(k+1) - e1(k+1)*tke(k))) / (1.0 - cq(k+1)*e1(k+1)), tke_min)
          dq(k) = tke(k) + dq(k)

          ! Adjust TKE deeper in the water column in case ke_tke increases.
          ! This might not be strictly necessary?
          do k=ke_tke+1,nz+1
            dq(k) = e1(k)*dq(k-1)
            tke(k) = max(tke(k) + dq(k), tke_min)
            if (abs(dq(k)) < roundoff*tke(k)) exit
          enddo
          do k2=k+1,nz
            if (dq(k2) == 0.0) exit
            dq(k2) = 0.0
          enddo
        endif
      endif
      do k=ke_tke-1,1,-1
        tke(k) = max(tke(k) + cq(k+1)*tke(k+1), tke_min)
        dq(k) = tke(k) + dq(k)
      enddo

  ! ----------------------------------------------------
  ! Calculate kappa, here defined at interfaces.
  ! ----------------------------------------------------

      ke_kappa_prev = ke_kappa ; ks_kappa_prev = ks_kappa

      dk(1) = 0.0 ! kappa takes boundary values of 0.
      ck(2) = 0.0 ; ckcomp = 1.0
      if (itt == 1) then ; dO k=2,nz
        i_ld2(k) = (n2(k)*ilambda2 + f2) / tke(k) + i_l2_bdry(k)
      enddo ; endif
      do k=2,nz
        dk(k) = -kappa(k)
        if (itt>1) &
          i_ld2(k) = (n2(k)*ilambda2 + f2) / tke(k) + i_l2_bdry(k)
        bkd1 = dz_int(k)*i_ld2(k) + ckcomp*idz(k-1)
        bk = 1.0 / (bkd1 + idz(k))

        kappa(k) = bk * (idz(k-1)*kappa(k-1) + dz_int(k) * k_src(k))
        ck(k+1) = idz(k) * bk ; ckcomp = bkd1 * bk ! = 1 - cK(K+1)

        ! Neglect values that are smaller than kappa_trunc.
        if (kappa(k) < ckcomp*kappa_trunc) then
          kappa(k) = 0.0
          if (k > ke_src) then ; ke_kappa = k-1 ; k_q(k) = 0.0 ; exit ; endif
        elseif (kappa(k) < 2.0*ckcomp*kappa_trunc) then
          kappa(k) = 2.0 * (kappa(k) - ckcomp*kappa_trunc)
        endif
      enddo
      k_q(ke_kappa) = kappa(ke_kappa) / tke(ke_kappa)
      dk(ke_kappa) = dk(ke_kappa) + kappa(ke_kappa)
      do k=ke_kappa+2,ke_kappa_prev
        dk(k) = -kappa(k) ; kappa(k) = 0.0 ; k_q(k) = 0.0
      enddo
      do k=ke_kappa-1,2,-1
        kappa(k) = kappa(k) + ck(k+1)*kappa(k+1)
        ! Neglect values that are smaller than kappa_trunc.
        if (kappa(k) <= kappa_trunc) then
          kappa(k) = 0.0
          if (k < ks_src) then ; ks_kappa = k+1 ; k_q(k) = 0.0 ; exit ; endif
        elseif (kappa(k) < 2.0*kappa_trunc) then
          kappa(k) = 2.0 * (kappa(k) - kappa_trunc)
        endif

        dk(k) = dk(k) + kappa(k)
        k_q(k) = kappa(k) / tke(k)
      enddo
      do k=ks_kappa_prev,ks_kappa-2 ; kappa(k) = 0.0 ; k_q(k) = 0.0 ; enddo

    else ! do_Newton is .true.
!   Once the solutions are close enough, use a Newton's method solver of the
!  whole system to accelerate convergence.
      ks_kappa_prev = ks_kappa ; ke_kappa_prev = ke_kappa ; ke_kappa = nz
      ks_kappa = 2
      dk(1) = 0.0 ; ck(2) = 0.0 ; ckcomp = 1.0 ; dkdq(1) = 0.0
      aq(1) = (0.5*(kappa(1)+kappa(2))+kappa0) * idz(1)
      dqdz(1) = 0.5*(tke(1) - tke(2))*idz(1)
      if (tke_noflux_top_bc) then
        tke_src = dz_int(1) * (kappa0*s2(1) - (tke(1) - q0)*tke_decay(1)) - &
                  aq(1) * (tke(1) - tke(2))

        bq = 1.0 / (aq(1) + dz_int(1)*tke_decay(1))
        cq(2) = aq(1) * bq
        cqcomp = (dz_int(1)*tke_decay(1)) * bq ! = 1 - cQ(2)
        dqmdk(2) = -dqdz(1) * bq
        dq(1) = bq * tke_src
      else
        dq(1) = 0.0 ; cq(2) = 0.0 ; cqcomp = 1.0 ; dqmdk(2) = 0.0
      endif
      do k=2,nz
        i_q = 1.0 / tke(k)
        i_ld2(k) = (n2(k)*ilambda2 + f2) * i_q + i_l2_bdry(k)

        kap_src = dz_int(k) * (k_src(k) - i_ld2(k)*kappa(k)) + &
                            idz(k-1)*(kappa(k-1)-kappa(k)) - idz(k)*(kappa(k)-kappa(k+1))

        ! Ensure that the pivot is always positive, and that 0 <= cK <= 1.
        ! Otherwise do not use Newton's method.
        decay_term_k = -idz(k-1)*dqmdk(k)*dkdq(k-1) + dz_int(k)*i_ld2(k)
        if (decay_term_k < 0.0) then ; abort_newton = .true. ; exit ; endif
        bk = 1.0 / (idz(k) + idz(k-1)*ckcomp + decay_term_k)

        ck(k+1) = bk * idz(k)
        ckcomp = bk * (idz(k-1)*ckcomp + decay_term_k) ! = 1-cK(K+1)
        if (cs%dKdQ_iteration_bug) then
          dkdq(k) = bk * (idz(k-1)*dkdq(k-1)*cq(k) + &
                      us%m_to_Z*(n2(k)*ilambda2 + f2) * i_q**2 * kappa(k) )
        else
          dkdq(k) = bk * (idz(k-1)*dkdq(k-1)*cq(k) + &
                      dz_int(k)*(n2(k)*ilambda2 + f2) * i_q**2 * kappa(k) )
        endif
        dk(k) = bk * (kap_src + idz(k-1)*dk(k-1) + idz(k-1)*dkdq(k-1)*dq(k-1))

        ! Truncate away negligibly small values of kappa.
        if (dk(k) <= ckcomp*(kappa_trunc - kappa(k))) then
          dk(k) = -ckcomp*kappa(k)
!         if (K > ke_src) then ; ke_kappa = k-1 ; K_Q(K) = 0.0 ; exit ; endif
        elseif (dk(k) < ckcomp*(2.0*kappa_trunc - kappa(k))) then
          dk(k) = 2.0 * dk(k) - ckcomp*(2.0*kappa_trunc - kappa(k))
        endif

        ! Solve for dQ(K)...
        aq(k) = (0.5*(kappa(k)+kappa(k+1))+kappa0) * idz(k)
        dqdz(k) = 0.5*(tke(k) - tke(k+1))*idz(k)
        tke_src = dz_int(k) * (((kappa(k) + kappa0)*s2(k) - kappa(k)*n2(k)) - &
                                  (tke(k) - q0)*tke_decay(k)) - &
                  (aq(k) * (tke(k) - tke(k+1)) - aq(k-1) * (tke(k-1) - tke(k)))
        v1 = aq(k-1) + dqdz(k-1)*dkdq(k-1)
        v2 = (v1*dqmdk(k) + dqdz(k-1)*ck(k)) + &
             ((dqdz(k-1) - dqdz(k)) + dz_int(k)*(s2(k) - n2(k)))

        ! Ensure that the pivot is always positive, and that 0 <= cQ <= 1.
        ! Otherwise do not use Newton's method.
        decay_term_q = dz_int(k)*tke_decay(k) - dqdz(k-1)*dkdq(k-1)*cq(k) - v2*dkdq(k)
        if (decay_term_q < 0.0) then ; abort_newton = .true. ; exit ; endif
        bq = 1.0 / (aq(k) + (cqcomp*aq(k-1) + decay_term_q))

        cq(k+1) = aq(k) * bq
        cqcomp = (cqcomp*aq(k-1) + decay_term_q) * bq
        dqmdk(k+1) = (v2 * ck(k+1) - dqdz(k)) * bq

        ! Ensure that TKE+dQ will not drop below 0.5*TKE.
        dq(k) = max(bq * ((v1 * dq(k-1) + dqdz(k-1)*dk(k-1)) + &
                          (v2 * dk(k) + tke_src)), cqcomp*(-0.5*tke(k)))

        ! Check whether the next layer will be affected by any nonzero kappas.
        if ((itt > 1) .and. (k > ke_src) .and. (dk(k) == 0.0) .and. &
            ((kappa(k) + kappa(k+1)) == 0.0)) then
        ! Could also do  .and. (bQ*abs(tke_src) < roundoff*TKE(K)) then
          ke_kappa = k-1 ; exit
        endif
      enddo
      if ((ke_kappa == nz) .and. (.not. abort_newton)) then
        dk(nz+1) = 0.0 ; dkdq(nz+1) = 0.0
        if (tke_noflux_bottom_bc) then
          k = nz+1
          tke_src = dz_int(k) * (kappa0*s2(k) - (tke(k) - q0)*tke_decay(k)) + &
                    aq(k-1) * (tke(k-1) - tke(k))

          v1 = aq(k-1) + dqdz(k-1)*dkdq(k-1)
          decay_term_q = max(0.0, dz_int(k)*tke_decay(k) - dqdz(k-1)*dkdq(k-1)*cq(k))
          if (decay_term_q < 0.0) then
            abort_newton = .true.
          else
            bq = 1.0 / (aq(k) + (cqcomp*aq(k-1) + decay_term_q))
          ! Ensure that TKE+dQ will not drop below 0.5*TKE.
            dq(k) = max(bq * ((v1 * dq(k-1) + dqdz(k-1)*dk(k-1)) + tke_src), -0.5*tke(k))
            tke(k) = max(tke(k) + dq(k), tke_min)
          endif
        else
          dq(nz+1) = 0.0
        endif
      elseif (.not. abort_newton) then
        ! Alter the first-guess determination of dQ(K).
        dq(ke_kappa+1) = dq(ke_kappa+1) / (1.0 - cq(ke_kappa+2)*e1(ke_kappa+2))
        tke(ke_kappa+1) = max(tke(ke_kappa+1) + dq(ke_kappa+1), tke_min)
        do k=ke_kappa+2,nz+1
          if (debug_soln .and. (k < nz+1)) then
          ! Ignore this source?
            aq(k) = (0.5*(kappa(k)+kappa(k+1))+kappa0) * idz(k)
        !    tke_src_norm = (dz_Int(K) * (kappa0*S2(K) - (TKE(K)-q0)*TKE_decay(K)) - &
        !                   (aQ(k) * (TKE(K) - TKE(K+1)) - aQ(k-1) * (TKE(K-1) - TKE(K))) ) / &
        !                   (aQ(k) + (aQ(k-1) + dz_Int(K)*TKE_decay(K)))
          endif
          dk(k) = 0.0
        ! Ensure that TKE+dQ will not drop below 0.5*TKE.
          dq(k) = max(e1(k)*dq(k-1),-0.5*tke(k))
          tke(k) = max(tke(k) + dq(k), tke_min)
          if (abs(dq(k)) < roundoff*tke(k)) exit
        enddo
        if (debug_soln) then ; do k2=k+1,nz+1 ; dq(k2) = 0.0 ; dk(k2) = 0.0 ; enddo ; endif
      endif
      if (.not. abort_newton) then
        do k=ke_kappa,2,-1
          ! Ensure that TKE+dQ will not drop below 0.5*TKE.
          dq(k) = max(dq(k) + (cq(k+1)*dq(k+1) + dqmdk(k+1) * dk(k+1)), -0.5*tke(k))
          tke(k) = max(tke(k) + dq(k), tke_min)
          dk(k) = dk(k) + (ck(k+1)*dk(k+1) + dkdq(k) * dq(k))
          ! Truncate away negligibly small values of kappa.
          if (dk(k) <= kappa_trunc - kappa(k)) then
            dk(k) = -kappa(k)
            kappa(k) = 0.0
            if ((k < ks_src) .and. (k+1 > ks_kappa)) ks_kappa = k+1
          elseif (dk(k) < 2.0*kappa_trunc - kappa(k)) then
            dk(k) =  2.0*dk(k) - (2.0*kappa_trunc - kappa(k))
            kappa(k) = max(kappa(k) + dk(k), 0.0) ! The max is for paranoia.
            if (k<=ks_kappa) ks_kappa = 2
          else
            kappa(k) = kappa(k) + dk(k)
            if (k<=ks_kappa) ks_kappa = 2
          endif
        enddo
        dq(1) = max(dq(1) + cq(2)*dq(2) + dqmdk(2) * dk(2), tke_min - tke(1))
        tke(1) = max(tke(1) + dq(1), tke_min)
        dk(1) = 0.0
      endif

      ! Check these solutions for consistency.
      !  The unit conversions here have not been carefully tested.
      if (debug_soln) then ; do k=2,nz
        ! In these equations, K_err_lin and Q_err_lin should be at round-off levels
        ! compared with the dominant terms, perhaps, dz_Int*I_Ld2*kappa and
        ! dz_Int*TKE_decay*TKE.  The exception is where, either 1) the decay term has been
        ! been increased to ensure a positive pivot, or 2) negative TKEs have been
        ! truncated, or 3) small or negative kappas have been rounded toward 0.
        i_q = 1.0 / tke(k)
        i_ld2_debug(k) = (n2(k)*ilambda2 + f2) * i_q + i_l2_bdry(k)

        kap_src = dz_int(k) * (k_src(k) - i_ld2(k)*kappa_prev(k)) + &
                            (idz(k-1)*(kappa_prev(k-1)-kappa_prev(k)) - &
                             idz(k)*(kappa_prev(k)-kappa_prev(k+1)))
        k_err_lin = -idz(k-1)*(dk(k-1)-dk(k)) + idz(k)*(dk(k)-dk(k+1)) + &
                     dz_int(k)*i_ld2_debug(k)*dk(k) - kap_src - &
                     dz_int(k)*(n2(k)*ilambda2 + f2)*i_q**2*kappa_prev(k) * dq(k)

        tke_src = dz_int(k) * ((kappa_prev(k) + kappa0)*s2(k) - &
                     kappa_prev(k)*n2(k) - (tke_prev(k) - q0)*tke_decay(k)) - &
                  (aq(k) * (tke_prev(k) - tke_prev(k+1)) -  aq(k-1) * (tke_prev(k-1) - tke_prev(k)))
        q_err_lin = tke_src + (aq(k-1) * (dq(k-1)-dq(k)) - aq(k) * (dq(k)-dq(k+1))) - &
                    0.5*(tke_prev(k)-tke_prev(k+1))*idz(k)  * (dk(k) + dk(k+1)) - &
                    0.5*(tke_prev(k)-tke_prev(k-1))*idz(k-1)* (dk(k-1) + dk(k)) + &
                    dz_int(k) * (dk(k) * (s2(k) - n2(k)) - dq(k)*tke_decay(k))
      enddo ; endif

    endif  ! End of the Newton's method solver.

    ! Test kappa for convergence...
    if ((tol_err < newton_err) .and. (.not.abort_newton)) then
      !  A lower tolerance is used to switch to Newton's method than to switch back.
      newton_test = newton_err ; if (do_newton) newton_test = 2.0*newton_err
      was_newton = do_newton
      within_tolerance = .true. ; do_newton = .true.
      do k=min(ks_kappa,ks_kappa_prev),max(ke_kappa,ke_kappa_prev)
        kappa_mean = kappa0 + (kappa(k) - 0.5*dk(k))
        if (abs(dk(k)) > newton_test * kappa_mean) then
          if (do_newton) abort_newton = .true.
          within_tolerance = .false. ; do_newton = .false. ; exit
        elseif (abs(dk(k)) > tol_err * kappa_mean) then
          within_tolerance = .false. ; if (.not.do_newton) exit
        endif
        if (abs(dq(k)) > newton_test*(tke(k) - 0.5*dq(k))) then
          if (do_newton) abort_newton = .true.
          do_newton = .false. ; if (.not.within_tolerance) exit
        endif
      enddo

    else  ! Newton's method will not be used again, so no need to check.
      within_tolerance = .true.
      do k=min(ks_kappa,ks_kappa_prev),max(ke_kappa,ke_kappa_prev)
        if (abs(dk(k)) > tol_err * (kappa0 + (kappa(k) - 0.5*dk(k)))) then
          within_tolerance = .false. ;  exit
        endif
      enddo
    endif

    if (abort_newton) then
      do_newton = .false. ; abort_newton = .false.
      ! We went to Newton too quickly last time, so restrict the tolerance.
      newton_err = 0.5*newton_err
      ke_kappa_prev = nz
      do k=2,nz ; k_q(k) = kappa(k) / max(tke(k), tke_min) ; enddo
    endif

    if (within_tolerance) exit

  enddo

  if (do_newton) then  ! K_Q needs to be calculated.
    do k=1,ks_kappa-1 ;  k_q(k) = 0.0 ; enddo
    do k=ks_kappa,ke_kappa ; k_q(k) = kappa(k) / tke(k) ; enddo
    do k=ke_kappa+1,nz+1 ; k_q(k) = 0.0 ; enddo
  endif

  if (present(local_src)) then
    local_src(1) = 0.0 ; local_src(nz+1) = 0.0
    do k=2,nz
      diffusive_src = idz(k-1)*(kappa(k-1)-kappa(k)) + idz(k)*(kappa(k+1)-kappa(k))
      chg_by_k0 = kappa0 * ((idz(k-1)+idz(k)) / dz_int(k) + i_ld2(k))
      if (diffusive_src <= 0.0) then
        local_src(k) = k_src(k) + chg_by_k0
      else
        local_src(k) = (k_src(k) + chg_by_k0) + diffusive_src / dz_int(k)
      endif
    enddo
  endif
  if (present(kappa_src)) then
    kappa_src(1) = 0.0 ; kappa_src(nz+1) = 0.0
    do k=2,nz
      kappa_src(k) = k_src(k)
    enddo
  endif

end subroutine find_kappa_tke

!> This subroutine initializes the parameters that regulate shear-driven mixing
function kappa_shear_init(Time, G, GV, US, param_file, diag, CS)
  type(time_type),         intent(in)    :: time !< The current model time.
  type(ocean_grid_type),   intent(in)    :: g    !< The ocean's grid structure.
  type(verticalgrid_type), intent(in)    :: gv   !< The ocean's vertical grid structure.
  type(unit_scale_type),   intent(in)    :: us   !< A dimensional unit scaling type
  type(param_file_type),   intent(in)    :: param_file !< A structure to parse for run-time
                                                 !! parameters.
  type(diag_ctrl), target, intent(inout) :: diag !< A structure that is used to regulate diagnostic
                                                 !! output.
  type(kappa_shear_cs),    pointer       :: cs   !< A pointer that is set to point to the control
                                                 !! structure for this module
  logical :: kappa_shear_init !< True if module is to be used, False otherwise

  ! Local variables
  logical :: merge_mixedlayer
  logical :: debug_shear
  logical :: just_read ! If true, this module is not used, so only read the parameters.
  ! This include declares and sets the variable "version".
# include "version_variable.h"
  character(len=40)  :: mdl = "MOM_kappa_shear"  ! This module's name.
  real :: kappa_0_unscaled  ! The value of kappa_0 in MKS units [m2 s-1]
  real :: kd_normal ! The KD of the main model, read here only as a parameter
                    ! for setting the default of KD_SMOOTH in MKS units [m2 s-1]

  if (associated(cs)) then
    call mom_error(warning, "kappa_shear_init called with an associated "// &
                            "control structure.")
    return
  endif
  allocate(cs)

  !   The Jackson-Hallberg-Legg shear mixing parameterization uses the following
  ! 6 nondimensional coefficients.  That paper gives 3 best fit parameter sets.
  !    Ri_Crit  Rate    FRi_Curv  K_buoy  TKE_N  TKE_Shear
  ! p1: 0.25    0.089    -0.97     0.82    0.24    0.14
  ! p2: 0.30    0.085    -0.94     0.86    0.26    0.13
  ! p3: 0.35    0.088    -0.89     0.81    0.28    0.12
  !   Future research will reveal how these should be modified to take
  ! subgridscale inhomogeneity into account.

! Set default, read and log parameters
  call get_param(param_file, mdl, "USE_JACKSON_PARAM", kappa_shear_init, default=.false., do_not_log=.true.)
  call log_version(param_file, mdl, version, &
    "Parameterization of shear-driven turbulence following Jackson, Hallberg and Legg, JPO 2008", &
    log_to_all=.true., debugging=kappa_shear_init, all_default=.not.kappa_shear_init)
  call get_param(param_file, mdl, "USE_JACKSON_PARAM", kappa_shear_init, &
                 "If true, use the Jackson-Hallberg-Legg (JPO 2008) "//&
                 "shear mixing parameterization.", default=.false.)
  just_read = .not.kappa_shear_init
  call get_param(param_file, mdl, "VERTEX_SHEAR", cs%KS_at_vertex, &
                 "If true, do the calculations of the shear-driven mixing "//&
                 "at the cell vertices (i.e., the vorticity points).", &
                 default=.false., do_not_log=just_read)
  call get_param(param_file, mdl, "RINO_CRIT", cs%RiNo_crit, &
                 "The critical Richardson number for shear mixing.", &
                 units="nondim", default=0.25, do_not_log=just_read)
  call get_param(param_file, mdl, "SHEARMIX_RATE", cs%Shearmix_rate, &
                 "A nondimensional rate scale for shear-driven entrainment. "//&
                 "Jackson et al find values in the range of 0.085-0.089.", &
                 units="nondim", default=0.089, do_not_log=just_read)
  call get_param(param_file, mdl, "MAX_RINO_IT", cs%max_RiNo_it, &
                 "The maximum number of iterations that may be used to "//&
                 "estimate the Richardson number driven mixing.", &
                 units="nondim", default=50, do_not_log=just_read)
  call get_param(param_file, mdl, "KD", kd_normal, default=0.0, do_not_log=.true.)
  call get_param(param_file, mdl, "KD_KAPPA_SHEAR_0", cs%kappa_0, &
                 "The background diffusivity that is used to smooth the "//&
                 "density and shear profiles before solving for the "//&
                 "diffusivities.  The default is the greater of KD and 1e-7 m2 s-1.", &
                 units="m2 s-1", default=max(kd_normal, 1.0e-7), scale=us%m2_s_to_Z2_T, &
                 unscaled=kappa_0_unscaled, do_not_log=just_read)
  call get_param(param_file, mdl, "KD_TRUNC_KAPPA_SHEAR", cs%kappa_trunc, &
                 "The value of shear-driven diffusivity that is considered negligible "//&
                 "and is rounded down to 0. The default is 1% of KD_KAPPA_SHEAR_0.", &
                 units="m2 s-1", default=0.01*kappa_0_unscaled, scale=us%m2_s_to_Z2_T, do_not_log=just_read)
  call get_param(param_file, mdl, "FRI_CURVATURE", cs%FRi_curvature, &
                 "The nondimensional curvature of the function of the "//&
                 "Richardson number in the kappa source term in the "//&
                 "Jackson et al. scheme.", units="nondim", default=-0.97, do_not_log=just_read)
  call get_param(param_file, mdl, "TKE_N_DECAY_CONST", cs%C_N, &
                 "The coefficient for the decay of TKE due to "//&
                 "stratification (i.e. proportional to N*tke). "//&
                 "The values found by Jackson et al. are 0.24-0.28.", &
                 units="nondim", default=0.24, do_not_log=just_read)
!  call get_param(param_file, mdl, "LAYER_KAPPA_STAGGER", CS%layer_stagger, &
!                 default=.false., do_not_log=just_read)
  call get_param(param_file, mdl, "TKE_SHEAR_DECAY_CONST", cs%C_S, &
                 "The coefficient for the decay of TKE due to shear (i.e. "//&
                 "proportional to |S|*tke). The values found by Jackson "//&
                 "et al. are 0.14-0.12.", units="nondim", default=0.14, do_not_log=just_read)
  call get_param(param_file, mdl, "KAPPA_BUOY_SCALE_COEF", cs%lambda, &
                 "The coefficient for the buoyancy length scale in the "//&
                 "kappa equation.  The values found by Jackson et al. are "//&
                 "in the range of 0.81-0.86.", units="nondim", default=0.82, do_not_log=just_read)
  call get_param(param_file, mdl, "KAPPA_N_OVER_S_SCALE_COEF2", cs%lambda2_N_S, &
                 "The square of the ratio of the coefficients of the "//&
                 "buoyancy and shear scales in the diffusivity equation, "//&
                 "Set this to 0 (the default) to eliminate the shear scale. "//&
                 "This is only used if USE_JACKSON_PARAM is true.", &
                 units="nondim", default=0.0, do_not_log=just_read)
  call get_param(param_file, mdl, "KAPPA_SHEAR_TOL_ERR", cs%kappa_tol_err, &
                 "The fractional error in kappa that is tolerated. "//&
                 "Iteration stops when changes between subsequent "//&
                 "iterations are smaller than this everywhere in a "//&
                 "column.  The peak diffusivities usually converge most "//&
                 "rapidly, and have much smaller errors than this.", &
                 units="nondim", default=0.1, do_not_log=just_read)
  call get_param(param_file, mdl, "TKE_BACKGROUND", cs%TKE_bg, &
                 "A background level of TKE used in the first iteration "//&
                 "of the kappa equation.  TKE_BACKGROUND could be 0.", &
                 units="m2 s-2", default=0.0, scale=us%m_to_Z**2*us%T_to_s**2)
  call get_param(param_file, mdl, "KAPPA_SHEAR_ELIM_MASSLESS", cs%eliminate_massless, &
                 "If true, massless layers are merged with neighboring "//&
                 "massive layers in this calculation.  The default is "//&
                 "true and I can think of no good reason why it should "//&
                 "be false. This is only used if USE_JACKSON_PARAM is true.", &
                 default=.true., do_not_log=just_read)
  call get_param(param_file, mdl, "MAX_KAPPA_SHEAR_IT", cs%max_KS_it, &
                 "The maximum number of iterations that may be used to "//&
                 "estimate the time-averaged diffusivity.", units="nondim", &
                 default=13, do_not_log=just_read)
  call get_param(param_file, mdl, "PRANDTL_TURB", cs%Prandtl_turb, &
                 "The turbulent Prandtl number applied to shear instability.", &
                 units="nondim", default=1.0, do_not_log=.true.)
  call get_param(param_file, mdl, "VEL_UNDERFLOW", cs%vel_underflow, &
                 "A negligibly small velocity magnitude below which velocity components are set "//&
                 "to 0.  A reasonable value might be 1e-30 m/s, which is less than an "//&
                 "Angstrom divided by the age of the universe.", &
                 units="m s-1", default=0.0, scale=us%m_s_to_L_T, do_not_log=just_read)
  call get_param(param_file, mdl, "KAPPA_SHEAR_MAX_KAP_SRC_CHG", cs%kappa_src_max_chg, &
                 "The maximum permitted increase in the kappa source within an iteration relative "//&
                 "to the local source; this must be greater than 1.  The lower limit for the "//&
                 "permitted fractional decrease is (1 - 0.5/kappa_src_max_chg).  These limits "//&
                 "could perhaps be made dynamic with an improved iterative solver.", &
                 default=10.0, units="nondim", do_not_log=just_read)

  call get_param(param_file, mdl, "DEBUG", cs%debug, &
                 "If true, write out verbose debugging data.", &
                 default=.false., debuggingparam=.true., do_not_log=just_read)
  call get_param(param_file, mdl, "DEBUG_KAPPA_SHEAR", debug_shear, &
                 "If true, write debugging data for the kappa-shear code.", &
                 default=.false., debuggingparam=.true., do_not_log=.true.)
  if (debug_shear) cs%debug = .true.
  call get_param(param_file, mdl, "KAPPA_SHEAR_VERTEX_PSURF_BUG", cs%psurf_bug, &
                 "If true, do a simple average of the cell surface pressures to get a pressure "//&
                 "at the corner if VERTEX_SHEAR=True.  Otherwise mask out any land points in "//&
                 "the average.", default=.true., do_not_log=(just_read .or. (.not.cs%KS_at_vertex)))

  call get_param(param_file, mdl, "KAPPA_SHEAR_ITER_BUG", cs%dKdQ_iteration_bug, &
                 "If true, use an older, dimensionally inconsistent estimate of the "//&
                 "derivative of diffusivity with energy in the Newton's method iteration.  "//&
                 "The bug causes undercorrections when dz > 1 m.", default=.false., do_not_log=just_read)
  call get_param(param_file, mdl, "KAPPA_SHEAR_ALL_LAYER_TKE_BUG", cs%all_layer_TKE_bug, &
                 "If true, report back the latest estimate of TKE instead of the time average "//&
                 "TKE when there is mass in all layers.  Otherwise always report the time "//&
                 "averaged TKE, as is currently done when there are some massless layers.", &
                 default=.false., do_not_log=just_read)
  call get_param(param_file, mdl, "USE_RESTRICTIVE_TOLERANCE_CHECK", cs%restrictive_tolerance_check, &
                 "If true, uses the more restrictive tolerance check to determine if a timestep "//&
                 "is acceptable for the KS_it outer iteration loop.  False uses the original less "//&
                 "restrictive check.", default=.false., do_not_log=just_read)
!    id_clock_KQ = cpu_clock_id('Ocean KS kappa_shear', grain=CLOCK_ROUTINE)
!    id_clock_avg = cpu_clock_id('Ocean KS avg', grain=CLOCK_ROUTINE)
!    id_clock_project = cpu_clock_id('Ocean KS project', grain=CLOCK_ROUTINE)
!    id_clock_setup = cpu_clock_id('Ocean KS setup', grain=CLOCK_ROUTINE)

  cs%nkml = 1
  if (gv%nkml>0) then
    call get_param(param_file, mdl, "KAPPA_SHEAR_MERGE_ML",merge_mixedlayer, &
                 "If true, combine the mixed layers together before solving the "//&
                 "kappa-shear equations.", default=.true., do_not_log=just_read)
    if (merge_mixedlayer) cs%nkml = gv%nkml
  endif

! Forego remainder of initialization if not using this scheme
  if (.not. kappa_shear_init) return

  cs%diag => diag

  cs%id_Kd_shear = register_diag_field('ocean_model','Kd_shear', diag%axesTi, time, &
      'Shear-driven Diapycnal Diffusivity', 'm2 s-1', conversion=us%Z2_T_to_m2_s)
  cs%id_TKE = register_diag_field('ocean_model','TKE_shear', diag%axesTi, time, &
      'Shear-driven Turbulent Kinetic Energy', 'm2 s-2', conversion=us%Z_to_m**2*us%s_to_T**2)

end function kappa_shear_init

!> This function indicates to other modules whether the Jackson et al shear mixing
!! parameterization will be used without needing to duplicate the log entry.
logical function kappa_shear_is_used(param_file)
  type(param_file_type), intent(in) :: param_file !< A structure to parse for run-time parameters

  ! Local variables
  character(len=40)  :: mdl = "MOM_kappa_shear"  ! This module's name.
  ! This function reads the parameter "USE_JACKSON_PARAM" and returns its value.

  call get_param(param_file, mdl, "USE_JACKSON_PARAM", kappa_shear_is_used, &
                 default=.false., do_not_log=.true.)
end function kappa_shear_is_used

!> This function indicates to other modules whether the Jackson et al shear mixing parameterization
!! will be used at the vertices without needing to duplicate the log entry.  It returns false if
!! the Jackson et al scheme is not used or if it is used via calculations at the tracer points.
logical function kappa_shear_at_vertex(param_file)
  type(param_file_type), intent(in) :: param_file !< A structure to parse for run-time parameters

  ! Local variables
  character(len=40)  :: mdl = "MOM_kappa_shear"  ! This module's name.
  logical :: do_kappa_shear
  ! This function returns true only if the parameters "USE_JACKSON_PARAM" and "VERTEX_SHEAR" are both true.

  kappa_shear_at_vertex = .false.

  call get_param(param_file, mdl, "USE_JACKSON_PARAM", do_kappa_shear, &
                 default=.false., do_not_log=.true.)
  if (do_kappa_shear) &
    call get_param(param_file, mdl, "VERTEX_SHEAR", kappa_shear_at_vertex, &
                 "If true, do the calculations of the shear-driven mixing "//&
                 "at the cell vertices (i.e., the vorticity points).", &
                 default=.false., do_not_log=.true.)

end function kappa_shear_at_vertex

!> \namespace mom_kappa_shear
!!
!! By Laura Jackson and Robert Hallberg, 2006-2008
!!
!!   This file contains the subroutines that determine the diapycnal
!! diffusivity driven by resolved shears, as specified by the
!! parameterizations described in Jackson and Hallberg (JPO, 2008).
!!
!!   The technique by which the 6 equations (for kappa, TKE, u, v, T,
!! and S) are solved simultaneously has been dramatically revised
!! from the previous version. The previous version was not converging
!! in some cases, especially near the surface mixed layer, while the
!! revised version does.  The revised version solves for kappa and
!! TKE with shear and stratification fixed, then marches the density
!! and velocities forward with an adaptive (and aggressive) time step
!! in a predictor-corrector-corrector emulation of a trapezoidal
!! scheme.  Run-time-settable parameters determine the tolerence to
!! which the kappa and TKE equations are solved and the minimum time
!! step that can be taken.

end module mom_kappa_shear