MOM_kappa_shear.F90

! This file is part of MOM6, the Modular Ocean Model version 6.

! See the LICENSE file for licensing information.

! SPDX-License-Identifier: Apache-2.0

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

module mom_kappa_shear

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_interface_heights, only : thickness_to_dz

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

use mom_eos,               only : calculate_density, calculate_specific_vol_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    :: lz_rescale      !<   A coefficient to rescale the distance to the nearest

                             !! solid boundary. This adjustment is to account for

                             !! regions where 3 dimensional turbulence prevents the

                             !! growth of shear instabilities [nondim].

  real    :: tke_bg          !<   The background level of TKE [Z2 T-2 ~> m2 s-2].

  real    :: kappa_0         !<   The background diapycnal diffusivity [H Z T-1 ~> m2 s-1 or Pa s]

  real    :: kappa_seed      !<   A moderately large seed value of diapycnal diffusivity that

                             !! is used as a starting turbulent diffusivity in the iterations

                             !! to finding an energetically constrained solution for the

                             !! shear-driven diffusivity [H Z T-1 ~> m2 s-1 or Pa s]

  real    :: kappa_trunc     !< Diffusivities smaller than this are rounded to 0 [H Z T-1 ~> m2 s-1 or Pa s]

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

  real    :: vs_geomean_kdmin !< A minimum diffusivity for computing the horizontal averages

                             !! when using the geometric mean with VERTEX_SHEAR=True.  The model

                             !! is sensitive to this value, which is a drawback of using the

                             !! geometric average as currently implemented.

  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 :: vs_viscosity_bug !< If true, use a bug in the calculation of the viscosity that sets

                             !! it to zero for all vertices that are on a coastline.

  logical :: vertex_shear_obc_bug !< If false, use extra masking when interpolating thicknesses to velocity

                             !! points for setting up the shear velocities at vertices to avoid using

                             !! external thicknesses at open boundaries.  When OBCs are not in use,

                             !! this parameter does not change answers, but true is more efficient.

  logical :: vs_geometricmean !< If true use geometric averaging for Kd from vertices to tracer points

  logical :: vs_thicknessmean !< If true use thickness weighting when averaging Kd from vertices to

                             !! tracer points

  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_kd_vertex = -1, &

             id_s2_init = -1, id_n2_init = -1, id_s2_mean = -1, id_n2_mean = -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)

  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

                                                   !! [H Z T-1 ~> m2 s-1 or kg m-1 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!) [H Z T-1 ~> m2 s-1 or Pa s]. 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.

  ! Local variables

  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1) :: &

    diag_n2_init, & ! Diagnostic of N2 as provided to this routine [T-2 ~> s-2]

    diag_s2_init, & ! Diagnostic of S2 as provided to this routine [T-2 ~> s-2]

    diag_n2_mean, & ! Diagnostic of N2 averaged over the timestep applied [T-2 ~> s-2]

    diag_s2_mean ! Diagnostic of S2 averaged over the timestep applied [T-2 ~> s-2]

  real, dimension(SZI_(G),SZK_(GV)) :: &

    h_2d, &             ! A 2-D version of h [H ~> m or kg m-2].

    dz_2d, &            ! Vertical distance between interface heights [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 [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3].

  real, dimension(SZI_(G),SZK_(GV)+1) :: &

    kappa_2d, & ! 2-D version of kappa_io [H Z T-1 ~> m2 s-1 or Pa s]

    tke_2d      ! 2-D version tke_io [Z2 T-2 ~> m2 s-2].

  real, dimension(SZK_(GV)) :: &

    idz, &      ! The inverse of the thickness of the merged layers [H-1 ~> m2 kg-1].

    h_lay, &    ! The layer thickness [H ~> m or kg m-2]

    dz_lay, &   ! The geometric layer thickness in height units [Z ~> m]

    u0xdz, &    ! The initial zonal velocity times dz [H L T-1 ~> m2 s-1 or kg m-1 s-1]

    v0xdz, &    ! The initial meridional velocity times dz [H L T-1 ~> m2 s-1 or kg m-1 s-1]

    t0xdz, &    ! The initial temperature times thickness [C H ~> degC m or degC kg m-2] or if

                ! temperature is not a state variable, the density times thickness [R H ~> kg m-2 or kg2 m-5]

    s0xdz       ! The initial salinity times dz [S H ~> ppt m or ppt kg m-2].

  real, dimension(SZK_(GV)+1) :: &

    kappa, &    ! The shear-driven diapycnal diffusivity at an interface [H Z T-1 ~> m2 s-1 or Pa s]

    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 [H Z T-1 ~> m2 s-1 or Pa s]

    tke_avg, &  ! The time-weighted average of TKE [Z2 T-2 ~> m2 s-2].

    n2_init, &  ! N2 as provided to this routine [T-2 ~> s-2].

    s2_init, &  ! S2 as provided to this routine [T-2 ~> s-2].

    n2_mean, &  ! The time-weighted average of N2 [T-2 ~> s-2].

    s2_mean     ! The time-weighted average of S2 [T-2 ~> 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 [H ~> m or kg m-2]

  real :: k0dt          ! The background diffusivity times the timestep [H Z ~> m2 or kg m-1]

  real :: dz_massless   ! A layer thickness that is considered massless [H ~> m or kg m-2]

  logical :: use_temperature  !  If true, temperature and salinity have been

                        ! allocated and are being used as state variables.

  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 = associated(tv%T)

  k0dt = dt*cs%kappa_0

  dz_massless = 0.1*sqrt((us%Z_to_m*gv%m_to_H)*k0dt)

  if ((cs%id_N2_init>0) .or. cs%debug) diag_n2_init(:,:,:) = 0.0

  if ((cs%id_S2_init>0) .or. cs%debug) diag_s2_init(:,:,:) = 0.0

  if (cs%id_N2_mean>0) diag_n2_mean(:,:,:) = 0.0

  if (cs%id_S2_mean>0) diag_s2_mean(:,:,:) = 0.0

  !$OMP parallel do default(private) shared(js,je,is,ie,nz,h,u_in,v_in,use_temperature,tv,G,GV,US, &

  !$OMP                                     CS,kappa_io,dz_massless,k0dt,p_surf,dt,tke_io,kv_io, &

  !$OMP                                     diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean)

  do j=js,je

    ! Convert layer thicknesses into geometric thickness in height units.

    call thickness_to_dz(h, tv, dz_2d, j, g, gv)

    do k=1,nz ; do i=is,ie

      h_2d(i,k) = h(i,j,k)

      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

!---------------------------------------

! Work on each column.

!---------------------------------------

    do i=is,ie ; if (g%mask2dT(i,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.

          h_lay(k) = 0.0 ; dz_lay(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 ((h_lay(nzc) > 0.0) .and. (h_2d(i,k) > dz_massless)) nzc = nzc+1

          if ((k>cs%nkml) .and. (h_lay(nzc) > 0.0) .and. &

              (h_2d(i,k) > dz_massless)) nzc = nzc+1

          ! Only merge clusters of massless layers.

!         if ((h_lay(nzc) > dz_massless) .or. &

!             ((h_lay(nzc) > 0.0) .and. (h_2d(i,k) > dz_massless))) nzc = nzc+1

          kc(k) = nzc

          h_lay(nzc) = h_lay(nzc) + h_2d(i,k)

          dz_lay(nzc) = dz_lay(nzc) + dz_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 / h_lay(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

          h_lay(k) = h_2d(i,k)

          dz_lay(k) = dz_2d(i,k)

          u0xdz(k) = u_2d(i,k)*h_lay(k) ; v0xdz(k) = v_2d(i,k)*h_lay(k)

        enddo

        if (use_temperature) then

          do k=1,nz

            t0xdz(k) = t_2d(i,k)*h_lay(k) ; s0xdz(k) = s_2d(i,k)*h_lay(k)

          enddo

        else

          do k=1,nz

            t0xdz(k) = rho_2d(i,k)*h_lay(k) ; s0xdz(k) = rho_2d(i,k)*h_lay(k)

          enddo

        endif

        nzc = nz

        do k=1,nzc+1 ; kc(k) = k ; kf(k) = 0.0 ; enddo

      endif

      f2 = 0.25 * ((g%Coriolis2Bu(i,j) + g%Coriolis2Bu(i-1,j-1)) + &

                   (g%Coriolis2Bu(i,j-1) + g%Coriolis2Bu(i-1,j)))

      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.

    ! ----------------------------------------------------

      do k=1,nzc+1 ; kappa(k) = cs%kappa_seed ; enddo

      call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, &

                              h_lay, dz_lay, u0xdz, v0xdz, t0xdz, s0xdz, kappa_avg, &

                              tke_avg, n2_init, s2_init, n2_mean, s2_mean, &

                              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

        if (cs%id_N2_mean>0) then ; do k=1,nz+1

          diag_n2_mean(i,j,k) = n2_mean(k)

        enddo ; endif

        if (cs%id_S2_mean>0) then ; do k=1,nz+1

          diag_s2_mean(i,j,k) = s2_mean(k)

        enddo ; endif

        if ((cs%id_N2_init>0) .or. cs%debug) then ; do k=1,nz+1

          diag_n2_init(i,j,k) = n2_init(k)

        enddo ; endif

        if ((cs%id_S2_init>0) .or. cs%debug) then ; do k=1,nz+1

          diag_s2_init(i,j,k) = s2_init(k)

        enddo ; endif

      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

        do k=1,nz+1

          if (kf(k) == 0.0) then

            if (cs%id_N2_mean>0) diag_n2_mean(i,j,k) = n2_mean(kc(k))

            if (cs%id_S2_mean>0) diag_s2_mean(i,j,k) = s2_mean(kc(k))

            if ((cs%id_N2_init>0) .or. cs%debug) diag_n2_init(i,j,k) = n2_init(kc(k))

            if ((cs%id_S2_init>0) .or. cs%debug) diag_s2_init(i,j,k) = s2_init(kc(k))

          else

            if (cs%id_N2_mean>0) &

              diag_n2_mean(i,j,k) = (1.0-kf(k)) * n2_mean(kc(k)) + kf(k) * n2_mean(kc(k)+1)

            if (cs%id_S2_mean>0) &

              diag_s2_mean(i,j,k) = (1.0-kf(k)) * s2_mean(kc(k)) + kf(k) * s2_mean(kc(k)+1)

            if ((cs%id_N2_init>0) .or. cs%debug) &

              diag_n2_init(i,j,k) = (1.0-kf(k)) * n2_init(kc(k)) + kf(k) * n2_init(kc(k)+1)

            if ((cs%id_S2_init>0) .or. cs%debug) &

              diag_s2_init(i,j,k) = (1.0-kf(k)) * s2_init(kc(k)) + kf(k) * s2_init(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(diag_n2_init, "kappa_shear N2_init", g%HI, unscale=us%s_to_T**2)

    call hchksum(diag_s2_init, "kappa_shear S2_init", g%HI, unscale=us%s_to_T**2)

    call hchksum(kappa_io, "kappa", g%HI, unscale=gv%HZ_T_to_m2_s)

    call hchksum(tke_io, "tke", g%HI, unscale=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)

  if (cs%id_N2_init > 0) call post_data(cs%id_N2_init, diag_n2_init, cs%diag)

  if (cs%id_S2_init > 0) call post_data(cs%id_S2_init, diag_s2_init, cs%diag)

  if (cs%id_N2_mean > 0) call post_data(cs%id_N2_mean, diag_n2_mean, cs%diag)

  if (cs%id_S2_mean > 0) call post_data(cs%id_S2_mean, diag_s2_mean, 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)

  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 [C ~> degC]

  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)),   &

                           intent(in)    :: s_in   !< Layer salinities [S ~> 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!) [H Z T-1 ~> m2 s-1 or kg m-1 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

                                                   !! [H Z T-1 ~> m2 s-1 or Pa s].

                                                   !! The previous value is used to initialize kappa

                                                   !! in the vertex columns as Kappa = Kv/Prandtl

                                                   !! to accelerate the iteration toward convergence.

  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.

  ! Local variables

  real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)+1) :: &

    diag_n2_init, & ! Diagnostic of N2 as provided to this routine [T-2 ~> s-2]

    diag_s2_init, & ! Diagnostic of S2 as provided to this routine [T-2 ~> s-2]

    diag_n2_mean, & ! Diagnostic of N2 averaged over the timestep applied [T-2 ~> s-2]

    diag_s2_mean    ! Diagnostic of S2 averaged over the timestep applied [T-2 ~> s-2]

  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: &

    dz_3d           ! Vertical distance between interface heights [Z ~> m].

  real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)+1) :: &

    kappa_vertex    ! Diffusivity at interfaces and vertices [H Z T-1 ~> m2 s-1 or Pa s]

  real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)) :: &

    h_vert          ! Thicknesses interpolated to vertices [H ~> m or kg m-2]

  real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)) :: &

    h_at_u          ! A mask-weighted thickness interpolated to u-points [H ~> m or kg m-2]

  real, dimension(SZI_(G),SZJB_(G),SZK_(GV)) :: &

    h_at_v          ! A mask-weighted thickness interpolated to v-points [H ~> m or kg m-2]

  real, dimension(SZIB_(G),SZK_(GV)) :: &

    h_2d, &             ! A 2-D version of h interpolated to vertices [H ~> m or kg m-2].

    dz_2d, &            ! Vertical distance between interface heights [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 [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3].

  real, dimension(SZIB_(G),SZK_(GV)+1) :: &

    kappa_2d    ! 2-D slice of kappa_vert [H Z T-1 ~> m2 s-1 or Pa s]

  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 thickness of the merged layers [H-1 ~> m2 kg-1].

    h_lay, &    ! The layer thickness [H ~> m or kg m-2]

    dz_lay, &   ! The geometric layer thickness in height units [Z ~> m]

    u0xdz, &    ! The initial zonal velocity times dz [L H T-1 ~> m2 s-1 or kg m-1 s-1].

    v0xdz, &    ! The initial meridional velocity times dz [H L T-1 ~> m2 s-1 or kg m-1 s-1]

    t0xdz, &    ! The initial temperature times dz [C H ~> degC m or degC kg m-2]

    s0xdz       ! The initial salinity times dz [S H ~> ppt m or ppt kg m-2]

  real, dimension(SZK_(GV)+1) :: &

    kappa, &    ! The shear-driven diapycnal diffusivity at an interface [H Z T-1 ~> m2 s-1 or Pa s]

    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 [H Z T-1 ~> m2 s-1 or Pa s]

    tke_avg, &  ! The time-weighted average of TKE [Z2 T-2 ~> m2 s-2].

    n2_init, &  ! N2 as provided to this routine [T-2 ~> s-2].

    s2_init, &  ! S2 as provided to this routine [T-2 ~> s-2].

    n2_mean, &  ! The time-weighted average of N2 [T-2 ~> s-2].

    s2_mean     ! The time-weighted average of S2 [T-2 ~> 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 [H ~> m or kg m-2]

  real :: k0dt          ! The background diffusivity times the timestep [H Z ~> m2 or kg m-1]

  real :: dz_massless   ! A layer thickness that is considered massless [H ~> m or kg m-2]

  real :: i_hwt         ! The inverse of the sum of the adjacent masked thickness weights [H-1 ~> m-1 or m2 kg-1]

  real :: i_htot        ! The inverse of the sum of the thicknesses at adjacent vertices [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.

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

  real :: h_sw, h_se, h_nw, h_ne ! Thicknesses at adjacent vertices [H ~> m or kg m-2]

  real :: mks_to_hz_t   ! A factor used to restore dimensional scaling after the geometric mean

                        ! diffusivity is taken using thickness weighted powers [H Z s m-2 T-1 ~> 1]

                        ! or [H Z m s kg-1 T-1 ~> 1]

  real :: h_tiny        ! A sub-roundoff thickness to use in the denominator when calculating

                        ! thickness-weighted averages [H ~> m or kg m-2]

  integer :: isb, ieb, jsb, jeb, i, j, k, nz, nzc

  ! Diagnostics that should be deleted?

  isb = g%isc-1 ; ieb = g%iecB ; jsb = g%jsc-1 ; jeb = g%jecB ; nz = gv%ke

  if ((cs%id_N2_init>0) .or. cs%debug) diag_n2_init(:,:,:) = 0.0

  if ((cs%id_S2_init>0) .or. cs%debug) diag_s2_init(:,:,:) = 0.0

  if (cs%id_N2_mean>0) diag_n2_mean(:,:,:) = 0.0

  if (cs%id_S2_mean>0) diag_s2_mean(:,:,:) = 0.0

  kappa_vertex(:,:,:) = 0.0

  use_temperature = associated(tv%T)

  k0dt =  dt*cs%kappa_0

  dz_massless = 0.1*sqrt((us%Z_to_m*gv%m_to_H)*k0dt)

  i_prandtl = 0.0 ; if (cs%Prandtl_turb > 0.0) i_prandtl = 1.0 / cs%Prandtl_turb

  h_tiny = 0.5 * gv%H_subroundoff

  ! Convert layer thicknesses into geometric thickness in height units.

  call thickness_to_dz(h, tv, dz_3d, g, gv, us, halo_size=1)

  if (cs%vertex_shear_OBC_bug) then

    !$OMP parallel do default(shared)

    do k=1,nz

      do j=jsb,jeb+1 ; do i=isb,ieb

        h_at_u(i,j,k) = g%mask2dCu(i,j) * (h(i,j,k) + h(i+1,j,k)) * 0.5

      enddo ; enddo

      do j=jsb,jeb ; do i=isb,ieb+1

        h_at_v(i,j,k) = g%mask2dCv(i,j) * (h(i,j,k) + h(i,j+1,k)) * 0.5

      enddo ; enddo

    enddo

  else

    ! Because G%mask2dCu(I,j) is zero if either G%mask2dT(i,j) or G%mask2dT(i+1,j) except at OBC

    ! faces, the following form give equivalent answers to those above unless OBCs are in use,

    ! although the former is clearly less complicated and costly.

    !$OMP parallel do default(shared)

    do k=1,nz

      do j=jsb,jeb+1 ; do i=isb,ieb

        h_at_u(i,j,k) = g%mask2dCu(i,j) * (g%mask2dT(i,j) * h(i,j,k) + g%mask2dT(i+1,j) * h(i+1,j,k)) / &

                                          (g%mask2dT(i,j) + g%mask2dT(i+1,j) + 1.0e-36)

      enddo ; enddo

      do j=jsb,jeb ; do i=isb,ieb+1

        h_at_v(i,j,k) = g%mask2dCv(i,j) * (g%mask2dT(i,j) * h(i,j,k) + g%mask2dT(i,j+1) * h(i,j+1,k)) / &

                                          (g%mask2dT(i,j) + g%mask2dT(i,j+1) + 1.0e-36)

      enddo ; enddo

    enddo

  endif

  !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,use_temperature,tv,G,GV,US,CS,kappa_io, &

  !$OMP                                     dz_massless,k0dt,p_surf,dt,tke_io,kv_io,kappa_vertex,h_vert,I_Prandtl, &

  !$OMP                                     diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean)

  do j=jsb,jeb

    ! 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) * h_at_u(i,j,k)) + (u_in(i,j+1,k) * h_at_u(i,j+1,k)) ) / &

                  ( (h_at_u(i,j,k) + h_at_u(i,j+1,k)) + h_tiny )

      v_2d(i,k) = ( (v_in(i,j,k) * h_at_v(i,j,k)) + (v_in(i+1,j,k) * h_at_v(i+1,j,k)) ) / &

                  ( (h_at_v(i,j,k) + h_at_v(i+1,j,k)) + h_tiny )

      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) = ((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 )

      dz_2d(i,k) = ((g%mask2dT(i,j) * dz_3d(i,j,k) + g%mask2dT(i+1,j+1) * dz_3d(i+1,j+1,k)) + &

                    (g%mask2dT(i+1,j) * dz_3d(i+1,j,k) + g%mask2dT(i,j+1) * dz_3d(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)))

!      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))) * 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

!---------------------------------------

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

          h_lay(k) = 0.0 ; dz_lay(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 ((h_lay(nzc) > 0.0) .and. (h_2d(I,k) > dz_massless)) nzc = nzc+1

          if ((k>cs%nkml) .and. (h_lay(nzc) > 0.0) .and. &

              (h_2d(i,k) > dz_massless)) nzc = nzc+1

          ! Only merge clusters of massless layers.

!         if ((h_lay(nzc) > dz_massless) .or. &

!             ((h_lay(nzc) > 0.0) .and. (h_2d(I,k) > dz_massless))) nzc = nzc+1

          kc(k) = nzc

          h_lay(nzc) = h_lay(nzc) + h_2d(i,k)

          dz_lay(nzc) = dz_lay(nzc) + dz_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 / h_lay(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

          h_lay(k) = h_2d(i,k)

          dz_lay(k) = dz_2d(i,k)

          u0xdz(k) = u_2d(i,k)*h_lay(k) ; v0xdz(k) = v_2d(i,k)*h_lay(k)

        enddo

        if (use_temperature) then

          do k=1,nz

            t0xdz(k) = t_2d(i,k)*h_lay(k) ; s0xdz(k) = s_2d(i,k)*h_lay(k)

          enddo

        else

          do k=1,nz

            t0xdz(k) = rho_2d(i,k)*h_lay(k) ; s0xdz(k) = rho_2d(i,k)*h_lay(k)

          enddo

        endif

        nzc = nz

        do k=1,nzc+1 ; kc(k) = k ; kf(k) = 0.0 ; enddo

      endif

      f2 = g%Coriolis2Bu(i,j)

      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.

    ! ----------------------------------------------------

      do k=1,nzc+1 ; kappa(k) = cs%kappa_seed ; enddo

      call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, &

                              h_lay, dz_lay, u0xdz, v0xdz, t0xdz, s0xdz, kappa_avg, &

                              tke_avg, n2_init, s2_init, n2_mean, s2_mean, 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

        if (cs%id_N2_mean>0) then ; do k=1,nz+1

          diag_n2_mean(i,j,k) = n2_mean(k)

        enddo ; endif

        if (cs%id_S2_mean>0) then ; do k=1,nz+1

          diag_s2_mean(i,j,k) = s2_mean(k)

        enddo ; endif

        if ((cs%id_N2_init>0) .or. cs%debug) then ; do k=1,nz+1

          diag_n2_init(i,j,k) = n2_init(k)

        enddo ; endif

        if ((cs%id_S2_init>0) .or. cs%debug) then ; do k=1,nz+1

          diag_s2_init(i,j,k) = s2_init(k)

        enddo ; endif

      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

        do k=1,nz+1

          if (kf(k) == 0.0) then

            if (cs%id_N2_mean>0) diag_n2_mean(i,j,k) = n2_mean(kc(k))

            if (cs%id_S2_mean>0) diag_s2_mean(i,j,k) = s2_mean(kc(k))

            if ((cs%id_N2_init>0) .or. cs%debug) diag_n2_init(i,j,k) = n2_init(kc(k))

            if ((cs%id_S2_init>0) .or. cs%debug) diag_s2_init(i,j,k) = s2_init(kc(k))

          else

            if (cs%id_N2_mean>0) &

              diag_n2_mean(i,j,k) = (1.0-kf(k)) * n2_mean(kc(k)) + kf(k) * n2_mean(kc(k)+1)

            if (cs%id_S2_mean>0) &

              diag_s2_mean(i,j,k) = (1.0-kf(k)) * s2_mean(kc(k)) + kf(k) * s2_mean(kc(k)+1)

            if ((cs%id_N2_init>0) .or. cs%debug) &

              diag_n2_init(i,j,k) = (1.0-kf(k)) * n2_init(kc(k)) + kf(k) * n2_init(kc(k)+1)

            if ((cs%id_S2_init>0) .or. cs%debug) &

              diag_s2_init(i,j,k) = (1.0-kf(k)) * s2_init(kc(k)) + kf(k) * s2_init(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

    ! Store the 2-d slices back in the 3-d arrays for restarts or interpolation back to tracer points.

    if (cs%VS_ThicknessMean) then

      do k=1,nz+1 ; do i=isb,ieb

        h_vert(i,j,k) = h_2d(i,k)

      enddo ; enddo

    endif

    if (cs%VS_viscosity_bug) then

      do k=1,nz+1 ; do i=isb,ieb

        kappa_vertex(i,j,k) = kappa_2d(i,k)

        tke_io(i,j,k) = g%mask2dBu(i,j) * tke_2d(i,k)

        kv_io(i,j,k) = ( g%mask2dBu(i,j) * kappa_vertex(i,j,k) ) * cs%Prandtl_turb

      enddo ; enddo

    else

      do k=1,nz+1 ; do i=isb,ieb

        kappa_vertex(i,j,k) = kappa_2d(i,k)

        tke_io(i,j,k) = tke_2d(i,k)

        kv_io(i,j,k) = kappa_vertex(i,j,k) * cs%Prandtl_turb

      enddo ; enddo

    endif

  enddo ! end of J-loop

  ! Set the diffusivities in tracer columns from the values at vertices.

  !$OMP parallel do default(private) shared(G,kappa_io)

  do j=g%jsc,g%jec ; do i=g%isc,g%iec

    ! The turbulent length scales (and hence turbulent diffusivity) should always go to 0 at the top and bottom.

    kappa_io(i,j,1) = 0.0

    kappa_io(i,j,nz+1) = 0.0

  enddo ; enddo

  if (cs%VS_ThicknessMean .and. cs%VS_GeometricMean) then

    ! This conversion factor is required to allow for arbitrary fractional powers of the diffusivities.

    mks_to_hz_t = 1.0 /  gv%HZ_T_to_MKS

    !$OMP parallel do default(private) shared(nz,G,GV,CS,kappa_io,kappa_vertex,h_vert)

    do k=2,nz ; do j=g%jsc,g%jec ; do i=g%isc,g%iec

      h_sw = 0.5 * (h_vert(i-1,j-1,k) + h_vert(i-1,j-1,k-1))

      h_ne = 0.5 * (h_vert(i,j,k) + h_vert(i,j,k-1))

      h_nw = 0.5 * (h_vert(i-1,j,k) + h_vert(i-1,j,k-1))

      h_se = 0.5 * (h_vert(i,j-1,k) + h_vert(i,j-1,k-1))

      if ((h_sw + h_ne) + (h_nw + h_se) > 0.0) then

        !  The geometric mean is zero if any component is zero, hence the need to use a floor

        !  on the value of kappa_trunc in regions on boundaries of shear zones.

        i_htot = 1.0 / ((h_sw + h_ne) + (h_nw + h_se))

        kappa_io(i,j,k) = g%mask2dT(i,j) * mks_to_hz_t * &

                            ( ((gv%HZ_T_to_MKS * max(kappa_vertex(i-1,j-1,k), cs%VS_GeoMean_Kdmin))**(h_sw*i_htot) * &

                               (gv%HZ_T_to_MKS * max(kappa_vertex(i,j,k), cs%VS_GeoMean_Kdmin))**(h_ne*i_htot)) * &

                              ((gv%HZ_T_to_MKS * max(kappa_vertex(i-1,j,k), cs%VS_GeoMean_Kdmin))**(h_nw*i_htot) * &

                               (gv%HZ_T_to_MKS * max(kappa_vertex(i,j-1,k), cs%VS_GeoMean_Kdmin))**(h_se*i_htot)) )

      else

        ! If all points have zero thickness, the thickness-weighted geometric mean is undefined, so use

        ! the non-thickness weighted geometric mean instead.

        kappa_io(i,j,k) = g%mask2dT(i,j) * sqrt(sqrt( &

              (max(kappa_vertex(i-1,j-1,k),cs%VS_GeoMean_Kdmin) * max(kappa_vertex(i,j,k),cs%VS_GeoMean_Kdmin)) * &

              (max(kappa_vertex(i-1,j,k),cs%VS_GeoMean_Kdmin) * max(kappa_vertex(i,j-1,k),cs%VS_GeoMean_Kdmin)) ))

      endif

    enddo ; enddo ; enddo

  elseif (cs%VS_ThicknessMean) then   ! Use thickness-weighted arithmetic mean diffusivities.

    !$OMP parallel do default(private) shared(nz,G,GV,CS,kappa_io,kappa_vertex,h_vert)

    do k=2,nz ; do j=g%jsc,g%jec ; do i=g%isc,g%iec

      h_sw = 0.5 * (h_vert(i-1,j-1,k) + h_vert(i-1,j-1,k-1))

      h_ne = 0.5 * (h_vert(i,j,k) + h_vert(i,j,k-1))

      h_nw = 0.5 * (h_vert(i-1,j,k) + h_vert(i-1,j,k-1))

      h_se = 0.5 * (h_vert(i,j-1,k) + h_vert(i,j-1,k-1))

      ! The following expression is a thickness weighted arithmetic mean at tracer points:

      i_htot = 1.0 / (((h_sw + h_ne) + (h_nw + h_se)) + gv%H_subroundoff)

      kappa_io(i,j,k) = g%mask2dT(i,j) * &

            (((kappa_vertex(i-1,j-1,k)*h_sw) + (kappa_vertex(i,j,k)*h_ne)) + &

             ((kappa_vertex(i-1,j,k)*h_nw) + (kappa_vertex(i,j-1,k)*h_se))) * i_htot

    enddo ; enddo ; enddo

  elseif (cs%VS_GeometricMean) then   ! The geometic mean diffusivities are not thickness weighted.

    !$OMP parallel do default(private) shared(nz,G,CS,kappa_io,kappa_vertex)

    do k=2,nz ; do j=g%jsc,g%jec ; do i=g%isc,g%iec

      kappa_io(i,j,k) = g%mask2dT(i,j) * sqrt(sqrt( &

            (max(kappa_vertex(i-1,j-1,k),cs%VS_GeoMean_Kdmin) * max(kappa_vertex(i,j,k),cs%VS_GeoMean_Kdmin)) * &

            (max(kappa_vertex(i-1,j,k),cs%VS_GeoMean_Kdmin) * max(kappa_vertex(i,j-1,k),cs%VS_GeoMean_Kdmin)) ))

    enddo ; enddo ; enddo

  else   ! Use a non-thickness weighted arithmetic mean.

    !$OMP parallel do default(private) shared(nz,G,CS,kappa_io,kappa_vertex)

    do k=2,nz ; do j=g%jsc,g%jec ; do i=g%isc,g%iec

      kappa_io(i,j,k) = g%mask2dT(i,j) * 0.25 * &

                       ((kappa_vertex(i-1,j-1,k) + kappa_vertex(i,j,k)) +&

                        (kappa_vertex(i-1,j,k) + kappa_vertex(i,j-1,k)))

    enddo ; enddo ; enddo

  endif

  if (cs%debug) then

    call bchksum(diag_n2_init, "shear_vertex N2_init", g%HI, unscale=us%s_to_T**2)

    call bchksum(diag_s2_init, "shear_vertex S2_init", g%HI, unscale=us%s_to_T**2)

    call hchksum(kappa_io, "kappa", g%HI, unscale=gv%HZ_T_to_m2_s)

    call bchksum(tke_io, "tke", g%HI, unscale=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)

  if (cs%id_Kd_vertex > 0) call post_data(cs%id_Kd_vertex, kappa_vertex, cs%diag)

  if (cs%id_N2_init > 0) call post_data(cs%id_N2_init, diag_n2_init, cs%diag)

  if (cs%id_S2_init > 0) call post_data(cs%id_S2_init, diag_s2_init, cs%diag)

  if (cs%id_N2_mean > 0) call post_data(cs%id_N2_mean, diag_n2_mean, cs%diag)

  if (cs%id_S2_mean > 0) call post_data(cs%id_S2_mean, diag_s2_mean, 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, hlay, dz_lay, &

                              u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, tke_avg, N2_init, S2_init, &

                              N2_mean, S2_mean, tv, CS, GV, US )

  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 [H Z T-1 ~> m2 s-1 or Pa s]

  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)    :: hlay  !< The layer thickness [H ~> m or kg m-2]

  real, dimension(SZK_(GV)), &

                     intent(in)    :: dz_lay !< The geometric layer thickness in height units [Z ~> m]

  real, dimension(SZK_(GV)), &

                     intent(in)    :: u0xdz !< The initial zonal velocity times hlay [H L T-1 ~> m2 s-1 or kg m-1 s-1]

  real, dimension(SZK_(GV)), &

                     intent(in)    :: v0xdz !< The initial meridional velocity times the

                                            !! layer thickness [H L T-1 ~> m2 s-1 or kg m-1 s-1]

  real, dimension(SZK_(GV)), &

                     intent(in)    :: T0xdz !< The initial temperature times hlay [C H ~> degC m or degC kg m-2]

  real, dimension(SZK_(GV)), &

                     intent(in)    :: S0xdz !< The initial salinity times hlay [S H ~> ppt m or ppt kg m-2]

  real, dimension(SZK_(GV)+1), &

                     intent(out)   :: kappa_avg !< The time-weighted average of kappa [H Z T-1 ~> m2 s-1 or Pa s]

  real, dimension(SZK_(GV)+1), &

                     intent(out)   :: tke_avg  !< The time-weighted average of TKE [Z2 T-2 ~> m2 s-2].

  real, dimension(SZK_(GV)+1), &

                     intent(out)   :: N2_mean  !< The time-weighted average of N2 [Z2 T-2 ~> m2 s-2].

  real, dimension(SZK_(GV)+1), &

                     intent(out)   :: S2_mean  !< The time-weighted average of S2 [Z2 T-2 ~> m2 s-2].

  real, dimension(SZK_(GV)+1), &

                     intent(out)   :: N2_init  !< The initial value of N2 [Z2 T-2 ~> m2 s-2].

  real, dimension(SZK_(GV)+1), &

                     intent(out)   :: S2_init  !< The initial value of S2 [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

  ! 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 [C ~> degC].

    Sal, &      ! The salinity after a timestep of mixing [S ~> ppt].

    u_test, v_test, & ! Temporary velocities [L T-1 ~> m s-1].

    T_test, S_test ! Temporary temperatures [C ~> degC] and salinities [S ~> ppt].

  real, dimension(nzc+1) :: &

    N2, &       ! The squared buoyancy frequency at an interface [T-2 ~> s-2].

    h_Int, &    ! The extent of a finite-volume space surrounding an interface,

                ! as used in calculating kappa and TKE [H ~> m or kg m-2]

    dz_int, &   ! The vertical distance with the space surrounding an interface,

                ! as used in calculating kappa and TKE [Z ~> m]

    dz_h_int, & ! The ratio of the vertical distances to the thickness around an

                ! interface [Z H-1 ~> nondim or m3 kg-1].  In non-Boussinesq mode

                ! this is the specific volume, otherwise it is a scaling factor.

    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.

    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 [H ~> m or kg m-2]

    c1, &       ! c1 is used in the tridiagonal (and similar) solvers [nondim].

    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 [H Z T-1 ~> m2 s-1 or Pa s]

    kappa_mid, & ! The average of the initial and predictor estimates of kappa [H Z T-1 ~> m2 s-1 or Pa s]

    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 [H Z T-1 ~> m2 s-1 or Pa s]

    pressure, & ! The pressure at an interface [R L2 T-2 ~> Pa].

    t_int, &    ! The temperature interpolated to an interface [C ~> degC].

    sal_int, &  ! The salinity interpolated to an interface [S ~> ppt].

    dbuoy_dt, & ! The partial derivative of buoyancy with changes in temperature [Z T-2 C-1 ~> m s-2 degC-1]

    dbuoy_ds, & ! The partial derivative of buoyancy with changes in salinity [Z T-2 S-1 ~> m s-2 ppt-1]

    dspv_dt, &  ! The partial derivative of specific volume with changes in temperature [R-1 C-1 ~> m3 kg-1 degC-1]

    dspv_ds, &  ! The partial derivative of specific volume with changes in salinity [R-1 S-1 ~> m3 kg-1 ppt-1]

    rho_int, &  ! The in situ density interpolated to an interface [R ~> kg m-3]

    i_l2_bdry, &   ! The inverse of the square of twice the harmonic mean

                   ! distance to the top and bottom boundaries [H-1 Z-1 ~> m-2 or m kg-1].

    k_q, &         ! Diffusivity divided by TKE [H T Z-1 ~> s or kg s m-3]

    k_q_tmp, &     ! A temporary copy of diffusivity divided by TKE [H T Z-1 ~> s or kg s m-3]

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

    h_from_top, & ! The total thickness above an interface [H ~> m or kg m-2]

    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 :: h_from_bot    ! The total thickness below and interface [H ~> m or kg m-2]

  real :: b1            ! The inverse of the pivot in the tridiagonal equations [H-1 ~> m-1 or m2 kg-1].

  real :: bd1           ! A term in the denominator of b1 [H ~> m or kg m-2].

  real :: d1            ! 1 - c1 in the tridiagonal equations [nondim]

  real :: gR0           ! A conversion factor from H to pressure, Rho_0 times g in Boussinesq

                        ! mode, or just g when non-Boussinesq [R L2 T-2 H-1 ~> kg m-2 s-2 or m 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 [H-2 ~> m-2 or m4 kg-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 :: wt_a          ! The fraction of a layer thickness identified with the interface

                        ! above a layer [nondim]

  real :: wt_b          ! The fraction of a layer thickness identified with the interface

                        ! below a layer [nondim]

  real :: k0dt          ! The background diffusivity times the timestep [H Z ~> m2 or kg m-1].

  real :: I_lz_rescale_sqr ! The inverse of a rescaling factor for L2_bdry (Lz) squared [nondim].

  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_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%H_to_RZ * gv%g_Earth

  g_r0 = gv%g_Earth_Z_T2 / gv%Rho0

  k0dt = dt*cs%kappa_0

  i_lz_rescale_sqr = 1.0 ; if (cs%lz_rescale > 0) i_lz_rescale_sqr = 1/(cs%lz_rescale*cs%lz_rescale)

  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_lay(k) ; enddo

  !   Set up I_dz_int as the inverse of the distance between

  ! adjacent layer centers.

  i_dz_int(1) = 2.0 / dz_lay(1)

  dist_from_top(1) = 0.0 ; h_from_top(1) = 0.0

  do k=2,nzc

    i_dz_int(k) = 2.0 / (dz_lay(k-1) + dz_lay(k))

    dist_from_top(k) = dist_from_top(k-1) + dz_lay(k-1)

    h_from_top(k) = h_from_top(k-1) + hlay(k-1)

  enddo

  i_dz_int(nzc+1) = 2.0 / dz_lay(nzc)

  ! Find the inverse of the squared distances from the boundaries.

  dist_from_bot = 0.0 ; h_from_bot = 0.0

  do k=nzc,2,-1

    dist_from_bot = dist_from_bot + dz_lay(k)

    h_from_bot = h_from_bot + hlay(k)

    ! Find the inverse of the squared distances from the boundaries,

    i_l2_bdry(k) = ((dist_from_top(k) + dist_from_bot) * (h_from_top(k) + h_from_bot)) / &

                   ((dist_from_top(k) * dist_from_bot) * (h_from_top(k) * h_from_bot))

    ! reduce the distance by a factor of "lz_rescale"

    i_l2_bdry(k) = i_lz_rescale_sqr*i_l2_bdry(k)

  enddo

  !   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 / (hlay(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 = hlay(1) * b1 ! = 1 - c1

    do k=2,nzc-1

      bd1 = hlay(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 / ((hlay(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 / (hlay(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 / (hlay(1) + k0dt*i_dz_int(2))

    u(1) = b1 * u0xdz(1) ; v(1) = b1 * v0xdz(1)

    b1 = 1.0 / hlay(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.

  h_int(1) = 0.0 ; h_int(2) = hlay(1)

  dz_int(1) = 0.0 ; dz_int(2) = dz_lay(1)

  do k=2,nzc-1

    norm = 1.0 / (hlay(k)*(hlay(k-1)+hlay(k+1)) + 2.0*hlay(k-1)*hlay(k+1))

    wt_a = ((hlay(k)+hlay(k+1)) * hlay(k-1)) * norm

    wt_b = ((hlay(k-1)+hlay(k)) * hlay(k+1)) * norm

    h_int(k) = h_int(k) + hlay(k) * wt_a

    h_int(k+1) = hlay(k) * wt_b

    dz_int(k) = dz_int(k) + dz_lay(k) * wt_a

    dz_int(k+1) = dz_lay(k) * wt_b

  enddo

  h_int(nzc) = h_int(nzc) + hlay(nzc) ; h_int(nzc+1) = 0.0

  dz_int(nzc) = dz_int(nzc) + dz_lay(nzc) ; dz_int(nzc+1) = 0.0

  if (gv%Boussinesq) then

    do k=1,nzc+1 ; dz_h_int(k) = gv%H_to_Z ; enddo

  else

    ! Find an effective average specific volume around an interface.

    dz_h_int(1:nzc+1) = 0.0

    if (hlay(1) > 0.0) dz_h_int(1) = dz_lay(1) / hlay(1)

    do k=2,nzc+1

      if (h_int(k) > 0.0) then

        dz_h_int(k) = dz_int(k) / h_int(k)

      else

        dz_h_int(k) = dz_h_int(k-1)

      endif

    enddo

  endif

  ! 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*hlay(k-1)

      t_int(k) = 0.5*(t(k-1) + t(k))

      sal_int(k) = 0.5*(sal(k-1) + sal(k))

    enddo

    if (gv%Boussinesq .or. gv%semi_Boussinesq) then

      call calculate_density_derivs(t_int, sal_int, pressure, dbuoy_dt, dbuoy_ds, &

                                    tv%eqn_of_state, (/2,nzc/), scale=-g_r0 )

    else

      ! These should perhaps be combined into a single call to calculate the thermal expansion

      ! and haline contraction coefficients?

      call calculate_specific_vol_derivs(t_int, sal_int, pressure, dspv_dt, dspv_ds, &

                                    tv%eqn_of_state, (/2,nzc/) )

      call calculate_density(t_int, sal_int, pressure, rho_int, tv%eqn_of_state, (/2,nzc/) )

      do k=2,nzc

        dbuoy_dt(k) = gv%g_Earth_Z_T2 * (rho_int(k) * dspv_dt(k))

        dbuoy_ds(k) = gv%g_Earth_Z_T2 * (rho_int(k) * dspv_ds(k))

      enddo

    endif

  elseif (gv%Boussinesq .or. gv%semi_Boussinesq) then

    do k=1,nzc+1 ; dbuoy_dt(k) = -g_r0 ; dbuoy_ds(k) = 0.0 ; enddo

  else

    do k=1,nzc+1 ; dbuoy_ds(k) = 0.0 ; enddo

    dbuoy_dt(1) = -gv%g_Earth_Z_T2 / gv%Rlay(1)

    do k=2,nzc

      dbuoy_dt(k) = -gv%g_Earth_Z_T2 / (0.5*(gv%Rlay(k-1) + gv%Rlay(k)))

    enddo

    dbuoy_dt(nzc+1) = -gv%g_Earth_Z_T2 / gv%Rlay(nzc)

  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, hlay, i_dz_int, dbuoy_dt, dbuoy_ds, &

                                 cs%vel_underflow, u, v, t, sal, n2, s2, gv, us)

  do k=1,nzc+1

    n2_init(k) = n2(k)

    s2_init(k) = s2(k)

  enddo

! ----------------------------------------------------

! Iterate

! ----------------------------------------------------

  dt_rem = dt

  do k=1,nzc+1

    k_q(k) = 0.0

    kappa_avg(k) = 0.0 ; tke_avg(k) = 0.0 ; n2_mean(k) = 0.0 ; s2_mean(k) = 0.0

    local_src_avg(k) = 0.0

    ! Use the grid spacings to scale errors in the source.

    if ( h_int(k) > 0.0 ) &

      local_src_avg(k) = 0.1 * k0dt * i_dz_int(k) / h_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, h_int, dz_int, dz_h_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, hlay, i_dz_int, &

                                       dbuoy_dt, dbuoy_ds, cs%vel_underflow, u_test, v_test, &

                                       t_test, s_test, n2, s2, gv, us, ks_int=ks_kappa, ke_int=ke_kappa)

        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, hlay, &

                   i_dz_int, dbuoy_dt, dbuoy_ds, cs%vel_underflow, u_test, v_test, t_test, s_test, &

                   n2, s2, gv, us, ks_int=ks_kappa, ke_int=ke_kappa)

          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, hlay, i_dz_int, &

                                     dbuoy_dt, dbuoy_ds, cs%vel_underflow, u_test, v_test, &

                                     t_test, s_test, n2, s2, gv, us, ks_int=ks_kappa, ke_int=ke_kappa)

    ! 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, h_int, dz_int, dz_h_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, hlay, i_dz_int, &

                                     dbuoy_dt, dbuoy_ds, cs%vel_underflow, u_test, v_test, &

                                     t_test, s_test, n2, s2, gv, us, ks_int=ks_kappa, ke_int=ke_kappa)

    ! call cpu_clock_end(id_clock_project)

    ! call cpu_clock_begin(id_clock_KQ)

      call find_kappa_tke(n2, s2, kappa_out, idz, h_int, dz_int, dz_h_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))

        n2_mean(k) = n2_mean(k) + dt_wt*n2(k)

        s2_mean(k) = s2_mean(k) + dt_wt*s2(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, hlay, i_dz_int, &

                                     dbuoy_dt, dbuoy_ds, cs%vel_underflow, u, v, t, sal, n2, s2, &

                                     gv, us)

    ! call cpu_clock_end(id_clock_project)

    endif

    if (dt_rem <= 0.0) exit

  enddo ! end itt loop

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, &

                                     vel_under, u, v, T, Sal, N2, S2, GV, US, ks_int, ke_int)

  integer,               intent(in)    :: nz  !< The number of layers (after eliminating massless

                                              !! layers?).

  real, dimension(nz+1), intent(in)    :: kappa !< The diapycnal diffusivity at interfaces,

                                              !! [H Z T-1 ~> m2 s-1 or Pa s].

  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 [C ~> degC].

  real, dimension(nz),   intent(in)    :: S0  !< The initial salinity [S ~> ppt].

  real,                  intent(in)    :: dt  !< The time step [T ~> s].

  real, dimension(nz),   intent(in)    :: dz  !< The layer thicknesses [H ~> m or kg m-2]

  real, dimension(nz+1), intent(in)    :: I_dz_int !< The inverse of the distance between successive

                                              !! layer centers [Z-1 ~> m-1].

  real, dimension(nz+1), intent(in)    :: dbuoy_dT !< The partial derivative of buoyancy with

                                              !! temperature [Z T-2 C-1 ~> m s-2 degC-1].

  real, dimension(nz+1), intent(in)    :: dbuoy_dS !< The partial derivative of buoyancy with

                                              !! salinity [Z T-2 S-1 ~> m s-2 ppt-1].

  real,                  intent(in)    :: vel_under !< Any velocities that are smaller in magnitude

                                              !! than this value are set to 0 [L T-1 ~> m s-1].

  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 [C ~> degC].

  real, dimension(nz),   intent(inout) :: Sal !< The salinity after dt [S ~> ppt].

  real, dimension(nz+1), intent(inout) :: N2  !< The buoyancy frequency squared at interfaces [T-2 ~> s-2].

  real, dimension(nz+1), intent(inout) :: S2  !< The squared shear at interfaces [T-2 ~> s-2].

  type(verticalgrid_type), intent(in)  :: GV  !< The ocean's vertical grid structure.

  type(unit_scale_type), intent(in)    :: US  !< A dimensional unit scaling type

  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.

  ! Local variables

  real, dimension(nz+1) :: c1 ! A tridiagonal variable [nondim]

  real :: a_a, a_b   ! Tridiagonal coupling coefficients [H ~> m or kg m-2]

  real :: b1, b1nz_0 ! Tridiagonal variables [H-1 ~> m-1 or m2 kg-1]

  real :: bd1        ! A term in the denominator of b1 [H ~> m or kg m-2]

  real :: d1         ! A tridiagonal variable [nondim]

  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)

  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)) < vel_under) u(ke) = 0.0

    if (abs(v(ke)) < vel_under) 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)) < vel_under) u(k) = 0.0

      if (abs(v(k)) < vel_under) 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)) < vel_under) u(k) = 0.0

      if (abs(v(k)) < vel_under) v(k) = 0.0

    enddo

  endif

  ! Store the squared shear at interfaces

  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)) * (us%L_to_Z*i_dz_int(ks))**2

  do k=ks+1,ke

    s2(k) = (((u(k)-u(k-1))**2) + ((v(k)-v(k-1))**2)) * (us%L_to_Z*i_dz_int(k))**2

  enddo

  if (ke<nz) &

    s2(ke+1) = (((u0(ke+1)-u(ke))**2) + ((v0(ke+1)-v(ke))**2)) * (us%L_to_Z*i_dz_int(ke+1))**2

  ! Store the buoyancy frequency at interfaces

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

end subroutine calculate_projected_state

!> This subroutine calculates new, consistent estimates of TKE and kappa.

subroutine find_kappa_tke(N2, S2, kappa_in, Idz, h_Int, dz_Int, dz_h_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

                                              !! [H Z T-1 ~> m2 s-1 or Pa s]

  real, dimension(nz+1), intent(in)    :: h_Int !< The thicknesses associated with interfaces

                                              !! [H ~> m or kg m-2]

  real, dimension(nz+1), intent(in)    :: dz_Int !< The vertical distances around interfaces [Z ~> m]

  real, dimension(nz+1), intent(in)    :: dz_h_Int !< The ratio of the vertical distances to the

                                              !! thickness around an interface [Z H-1 ~> nondim or m3 kg-1].

                                              !! In non-Boussinesq mode this is the specific volume.

  real, dimension(nz+1), intent(in)    :: I_L2_bdry !< The inverse of the squared distance to

                                              !! boundaries [H-1 Z-1 ~> m-2 or m kg-1].

  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 [H T Z-1 ~> s or kg s m-3].

  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

                                              !! [H Z T-1 ~> m2 s-1 or Pa s]

  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 [H T-1 ~> m s-1 or kg m-2 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 [H Z T-1 ~> m2 s-1 or Pa s].

    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 [H-1 Z-1 ~> m-2 or m kg-1]

    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) [Z T H-1 ~> s or m3 s kg-1]

    dkdq, &       ! With Newton's method the change in dK(k) due to dQ(k) [H Z-1 T-1 ~> s-1 or kg m-3 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 [nondim].

                  ! 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] or [H Z T-3 ~> m2 s-3 or kg m-1 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 H-1 ~> s m-1 or m2 s kg-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 [H T-1 ~> m s-1 or kg m-2 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 [nondim].

  real :: c_s2          !   The coefficient for the decay of TKE due to

                        ! shear (i.e. proportional to |S|*tke) [nondim].

  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 [nondim]. 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 [H Z T-1 ~> m2 s-1 or Pa s]

  real :: kappa_trunc   ! Diffusivities smaller than this are rounded to 0 [H Z T-1 ~> m2 s-1 or Pa s]

  real :: eden1, eden2  ! Variables used in calculating e1 [H Z-2 ~> m-1 or kg m-4]

  real :: I_eden        ! The inverse of the denominator in e1 [Z2 H-1 ~> m or m4 kg-1]

  real :: ome           ! Variables used in calculating e1 [nondim]

  real :: diffusive_src ! The diffusive source in the kappa equation [H T-1 ~> m s-1 or kg m-2 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 :: h_dz_here     ! The ratio of the thicknesses to the vertical distances around an interface

                        ! [H Z-1 ~> nondim or kg m-3].  In non-Boussinesq mode this is the density.

  real :: kappa_mean    ! A mean value of kappa [H Z T-1 ~> m2 s-1 or Pa s]

  real :: Newton_test   ! The value of relative error that will cause the next

                        ! iteration to use Newton's method [nondim].

  ! Temporary variables used in the Newton's method iterations.

  real :: decay_term_k  ! The decay term in the diffusivity equation [Z-1 ~> m-1]

  real :: decay_term_Q  ! The decay term in the TKE equation - proportional to [H Z-1 T-1 ~> s-1 or kg m-3 s-1]

  real :: I_Q           ! The inverse of TKE [T2 Z-2 ~> s2 m-2]

  real :: kap_src       ! A source term in the kappa equation [H T-1 ~> m s-1 or kg m-2 s-1]

  real :: v1            ! A temporary variable proportional to [H Z-1 T-1 ~> s-1 or kg m-3 s-1]

  real :: v2            ! A temporary variable in [Z T-2 ~> m s-2]

  real :: tol_err       ! The tolerance for max_err that determines when to

                        ! stop iterating [nondim].

  real :: Newton_err    ! The tolerance for max_err that determines when to

                        ! start using Newton's method [nondim].  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 [nondim].

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

  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 ! The imbalance in the K equation [H T-1 ~> m s-1 or kg m-2 s-1]

  real :: Q_err_lin ! The imbalance in the Q equation [H Z T-3 ~> m2 s-3 or kg m-1 s-3]

  real, dimension(nz+1) :: &

    I_Ld2_debug, & ! A separate version of I_Ld2 for debugging [H-1 Z-1 ~> m-2 or m kg-1].

    kappa_prev, & ! The value of kappa at the start of the current iteration [H Z T-1 ~> m2 s-1 or Pa s]

    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 = h_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 = h_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 = dz_h_int(1)*kappa0*s2(1) + q0 * tke_decay(1) ! Uses that kappa(1) = 0

        bqd1 = h_int(1) * tke_decay(1)

        bq = 1.0 / (bqd1 +  aq(1))

        tke(1) = bq * (h_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 = dz_h_int(k)*(kappa(k) + kappa0)*s2(k) + q0*tke_decay(k)

        bqd1 = h_int(k)*(tke_decay(k) + dz_h_int(k)*n2(k)*k_q(k)) + cqcomp*aq(k-1)

        bq = 1.0 / (bqd1 + aq(k))

        tke(k) = bq * (h_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 = dz_h_int(k)*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 / (h_int(k)*tke_decay(k) + cqcomp*aq(k-1))

          tke(k) = max(tke_min, bq * (h_int(k)*tke_src + aq(k-1)*tke(k-1)))

          dq(k) = tke(k) + dq(k)

        else

          bq = 1.0 / ((h_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 * (h_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) = dz_h_int(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) = dz_h_int(k)*(n2(k)*ilambda2 + f2) / tke(k) + i_l2_bdry(k)

        bkd1 = h_int(k)*i_ld2(k) + ckcomp*idz(k-1)

        bk = 1.0 / (bkd1 + idz(k))

        kappa(k) = bk * (idz(k-1)*kappa(k-1) + h_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 = h_int(1) * (kappa0*dz_h_int(1)*s2(1) - (tke(1) - q0)*tke_decay(1)) - &

                  aq(1) * (tke(1) - tke(2))

        bq = 1.0 / (aq(1) + h_int(1)*tke_decay(1))

        cq(2) = aq(1) * bq

        cqcomp = (h_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) * dz_h_int(k)*i_q + i_l2_bdry(k)

        kap_src = h_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) + h_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 = h_int(k) * ((dz_h_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 = h_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 = h_int(k) * (kappa0*dz_h_int(k)*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, h_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 = ((kappa0*dz_Int(K)*S2(K) - h_Int(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) + h_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, h_Int*I_Ld2*kappa and

        ! h_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) * dz_h_int(k)*i_q + i_l2_bdry(k)

        kap_src = h_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)) + &

                     h_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)

        h_dz_here = 0.0 ; if (abs(dz_h_int(k)) > 0.0) h_dz_here = 1.0 / dz_h_int(k)

        tke_src = h_int(k) * ((kappa_prev(k) + kappa0)*s2(k) - &

                     kappa_prev(k)*n2(k) - (tke_prev(k) - q0)*h_dz_here*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)) / h_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 / h_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

  real :: kd_normal ! The KD of the main model, read here only as a parameter

                    ! for setting the default of KD_SMOOTH [Z2 T-1 ~> m2 s-1]

  real :: kappa_0_default ! The default value for KD_KAPPA_SHEAR_0 [Z2 T-1 ~> m2 s-1]

  logical :: merge_mixedlayer

  integer :: number_of_obc_segments

  logical :: debug_shear

  logical :: enable_bugs  ! If true, the defaults for recently added bug-fix flags are set to

                          ! recreate the bugs, or if false bugs are only used if actively selected.

  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.

  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, "ENABLE_BUGS_BY_DEFAULT", enable_bugs, &

                 default=.true., do_not_log=.true.)  ! This is logged from MOM.F90.

  call get_param(param_file, mdl, "VERTEX_SHEAR_VISCOSITY_BUG", cs%VS_viscosity_bug, &

                 "If true, use a bug in vertex shear that zeros out viscosities at "//&

                 "vertices on coastlines.", &

                 default=enable_bugs, do_not_log=just_read.or.(.not.cs%KS_at_vertex))

  call get_param(param_file, mdl, "OBC_NUMBER_OF_SEGMENTS", number_of_obc_segments, &

                 default=0, do_not_log=.true.)

  call get_param(param_file, mdl, "VERTEX_SHEAR_OBC_BUG", cs%vertex_shear_OBC_bug, &

                 "If false, use extra masking when interpolating thicknesses to velocity "//&

                 "points for setting up the shear velocities at vertices to avoid using "//&

                 "external thicknesses at open boundaries.  When OBCs are not in use, "//&

                 "this parameter does not change answers, but true is more efficient.", &

                 default=enable_bugs, &

                 do_not_log=just_read.or.(.not.cs%KS_at_vertex).or.(number_of_obc_segments<=0))

                 ! Use OBC settings to set the default for VERTEX_SHEAR_OBC_BUG?

  call get_param(param_file, mdl, "VERTEX_SHEAR_GEOMETRIC_MEAN", cs%VS_GeometricMean, &

                 "If true, use a geometric mean for moving diffusivity from "//&

                 "vertices to tracer points.  False uses algebraic mean.", &

                 default=.false., do_not_log=just_read.or.(.not.cs%KS_at_vertex))

  call get_param(param_file, mdl, "VERTEX_SHEAR_THICKNESS_MEAN", cs%VS_ThicknessMean, &

                 "If true, apply thickness weighting to horizontal averagings of diffusivity "//&

                 "to tracer points in the kappa shear solver.", &

                 default=.false.)

  if (cs%VS_GeometricMean) then

    call get_param(param_file, mdl, "VERTEX_SHEAR_GEOMETRIC_MEAN_KDMIN", &

                   cs%VS_GeoMean_Kdmin, "If using the geometric mean in vertex shear, "//&

                   "use this minimum value for Kd. This is an ad-hoc parameter, the "//&

                   "diffusivities on the edge of shear regions are sensitive to the choice.",&

                   units="m2 s-1", default=0.0, scale=gv%m2_s_to_HZ_T, do_not_log=just_read)

  endif

  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, &

                 units="m2 s-1", scale=us%m2_s_to_Z2_T, default=0.0, do_not_log=.true.)

  kappa_0_default = max(kd_normal, 1.0e-7*us%m2_s_to_Z2_T)

  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=kappa_0_default*us%Z2_T_to_m2_s, scale=gv%m2_s_to_HZ_T, &

                 do_not_log=just_read)

  call get_param(param_file, mdl, "KD_SEED_KAPPA_SHEAR", cs%kappa_seed, &

                 "A moderately large seed value of diapycnal diffusivity that is used as a "//&

                 "starting turbulent diffusivity in the iterations to find an energetically "//&

                 "constrained solution for the shear-driven diffusivity.", &

                 units="m2 s-1", default=1.0, scale=gv%m2_s_to_HZ_T)

  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*cs%kappa_0*gv%HZ_T_to_m2_s, scale=gv%m2_s_to_HZ_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, "LZ_RESCALE", cs%lz_rescale, &

                 "A coefficient to rescale the distance to the nearest solid boundary. "//&

                 "This adjustment is to account for regions where 3 dimensional turbulence "//&

                 "prevents the growth of shear instabilities [nondim].", &

                 units="nondim", default=1.0)

  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.", &

                 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=just_read)

  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=.false., 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 at horizontal tracer points', 'm2 s-1', conversion=gv%HZ_T_to_m2_s)

  if (cs%KS_at_vertex) then

    cs%id_TKE = register_diag_field('ocean_model','TKE_shear', diag%axesBi, time, &

       'Shear-driven Turbulent Kinetic Energy at horizontal vertices', 'm2 s-2', conversion=us%Z_to_m**2*us%s_to_T**2)

    cs%id_Kd_vertex = register_diag_field('ocean_model','Kd_shear_vertex', diag%axesBi, time, &

         'Shear-driven Diapycnal Diffusivity at horizontal vertices', 'm2 s-1', conversion=gv%HZ_T_to_m2_s)

    cs%id_S2_init = register_diag_field('ocean_model','S2_shear_in', diag%axesBi, time, &

         'Interface shear squared at horizontal vertices, as input to kappa-shear', 's-2', conversion=us%s_to_T**2)

    cs%id_N2_init = register_diag_field('ocean_model','N2_shear_in', diag%axesBi, time, &

         'Interface stratification at horizontal vertices, as input to kappa-shear', 's-2', conversion=us%s_to_T**2)

    cs%id_S2_mean = register_diag_field('ocean_model','S2_shear_mean', diag%axesBi, time, &

         'Interface shear squared at horizontal vertices, averaged over timestep in kappa-shear', &

         's-2', conversion=us%s_to_T**2)

    cs%id_N2_mean = register_diag_field('ocean_model','N2_shear_mean', diag%axesBi, time, &

         'Interface stratification at horizontal vertices, averaged over timestep in kappa-shear', &

         's-2', conversion=us%s_to_T**2)

  else

    cs%id_TKE = register_diag_field('ocean_model','TKE_shear', diag%axesTi, time, &

         'Shear-driven Turbulent Kinetic Energy at horizontal tracer points', &

         'm2 s-2', conversion=us%Z_to_m**2*us%s_to_T**2)

    cs%id_S2_init = register_diag_field('ocean_model','S2_shear_in', diag%axesTi, time, &

         'Interface shear squared at horizontal tracer points, as input to kappa-shear', 's-2', conversion=us%s_to_T**2)

    cs%id_N2_init = register_diag_field('ocean_model','N2_shear_in', diag%axesTi, time, &

         'Interface stratification at horizontal tracer points, as input to kappa-shear', &

         's-2', conversion=us%s_to_T**2)

    cs%id_S2_mean = register_diag_field('ocean_model','S2_shear_mean', diag%axesTi, time, &

         'Interface shear squared at horizontal tracer points, averaged over timestep in kappa-shear', &

         's-2', conversion=us%s_to_T**2)

    cs%id_N2_mean = register_diag_field('ocean_model','N2_shear_mean', diag%axesTi, time, &

         'Interface stratification at horizontal tracer points, averaged ove timestep in kappa-shear', &

         's-2', conversion=us%s_to_T**2)

  endif

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 tolerance to

!! which the kappa and TKE equations are solved and the minimum time

!! step that can be taken.

end module mom_kappa_shear