MOM_diapyc_energy_req.F90

!> Calculates the energy requirements of mixing.
module mom_diapyc_energy_req

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

!! \author By Robert Hallberg, May 2015

use mom_diag_mediator, only : diag_ctrl, time_type, post_data, register_diag_field
use mom_error_handler, only : mom_error, fatal, warning, mom_mesg, is_root_pe
use mom_file_parser, only : get_param, log_version, param_file_type
use mom_grid, only : ocean_grid_type
use mom_unit_scaling, only : unit_scale_type
use mom_variables, only : thermo_var_ptrs
use mom_verticalgrid, only : verticalgrid_type
use mom_eos, only : calculate_specific_vol_derivs, calculate_density

implicit none ; private

public diapyc_energy_req_init, diapyc_energy_req_calc, diapyc_energy_req_test, diapyc_energy_req_end

! 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 parameters for the MOM_diapyc_energy_req module
type, public :: diapyc_energy_req_cs ; private
  logical :: initialized = .false. !< A variable that is here because empty
                               !! structures are not permitted by some compilers.
  real :: test_kh_scaling      !< A scaling factor for the diapycnal diffusivity.
  real :: colht_scaling        !< A scaling factor for the column height change correction term.
  logical :: use_test_kh_profile !< If true, use the internal test diffusivity profile in place of
                               !! any that might be passed in as an argument.
  type(diag_ctrl), pointer :: diag => null() !< A structure that is used to
                               !! regulate the timing of diagnostic output.

  !>@{ Diagnostic IDs
  integer :: id_ert=-1, id_erb=-1, id_erc=-1, id_erh=-1, id_kddt=-1, id_kd=-1
  integer :: id_chct=-1, id_chcb=-1, id_chcc=-1, id_chch=-1
  integer :: id_t0=-1, id_tf=-1, id_s0=-1, id_sf=-1, id_n2_0=-1, id_n2_f=-1
  integer :: id_h=-1, id_zint=-1
  !>@}
end type diapyc_energy_req_cs

contains

!> This subroutine helps test the accuracy of the diapycnal mixing energy requirement code
!! by writing diagnostics, possibly using an intensely mixing test profile of diffusivity
subroutine diapyc_energy_req_test(h_3d, dt, tv, G, GV, US, CS, Kd_int)
  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(G%isd:G%ied,G%jsd:G%jed,GV%ke), &
                                  intent(in)    :: h_3d !< Layer thickness before entrainment [H ~> m or kg m-2].
  type(thermo_var_ptrs),          intent(inout) :: tv   !< A structure containing pointers to any
                                                        !! available thermodynamic fields.
                                                        !! Absent fields have NULL ptrs.
  real,                           intent(in)    :: dt   !< The amount of time covered by this call [T ~> s].
  type(diapyc_energy_req_cs),     pointer       :: cs   !< This module's control structure.
  real, dimension(G%isd:G%ied,G%jsd:G%jed,GV%ke+1), &
                        optional, intent(in)    :: kd_int !< Interface diffusivities [Z2 T-1 ~> m2 s-1].

  ! Local variables
  real, dimension(GV%ke) :: &
    t0, s0, &   ! T0 & S0 are columns of initial temperatures and salinities [degC] and g/kg.
    h_col       ! h_col is a column of thicknesses h at tracer points [H ~> m or kg m-2].
  real, dimension(GV%ke+1) :: &
    kd, &       ! A column of diapycnal diffusivities at interfaces [Z2 T-1 ~> m2 s-1].
    h_top, h_bot ! Distances from the top or bottom [H ~> m or kg m-2].
  real :: ustar, absf, htot
  real :: energy_kd ! The energy used by diapycnal mixing [W m-2].
  real :: tmp1  ! A temporary array.
  integer :: i, j, k, is, ie, js, je, nz, itt
  logical :: may_print
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = g%ke

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

!$OMP do
  do j=js,je ; do i=is,ie ; if (g%mask2dT(i,j) > 0.5) then
    if (present(kd_int) .and. .not.cs%use_test_Kh_profile) then
      do k=1,nz+1 ; kd(k) = cs%test_Kh_scaling*kd_int(i,j,k) ; enddo
    else
      htot = 0.0 ; h_top(1) = 0.0
      do k=1,nz
        t0(k) = tv%T(i,j,k) ; s0(k) = tv%S(i,j,k)
        h_col(k) = h_3d(i,j,k)
        h_top(k+1) = h_top(k) + h_col(k)
      enddo
      htot = h_top(nz+1)
      h_bot(nz+1) = 0.0
      do k=nz,1,-1
        h_bot(k) = h_bot(k+1) + h_col(k)
      enddo

      ustar = 0.01*us%m_to_Z*us%T_to_s ! Change this to being an input parameter?
      absf = 0.25*((abs(g%CoriolisBu(i-1,j-1)) + abs(g%CoriolisBu(i,j))) + &
                             (abs(g%CoriolisBu(i-1,j-1)) + abs(g%CoriolisBu(i,j))))
      kd(1) = 0.0 ; kd(nz+1) = 0.0
      do k=2,nz
        tmp1 = h_top(k) * h_bot(k) * gv%H_to_Z
        kd(k) = cs%test_Kh_scaling *  &
                ustar * 0.41 * (tmp1*ustar) / (absf*tmp1 + htot*ustar)
      enddo
    endif
    may_print = is_root_pe() .and. (i==ie) .and. (j==je)
    call diapyc_energy_req_calc(h_col, t0, s0, kd, energy_kd, dt, tv, g, gv, us, &
                                may_print=may_print, cs=cs)
  endif ; enddo ; enddo

end subroutine diapyc_energy_req_test

!>   This subroutine uses a substantially refactored tridiagonal equation for
!! diapycnal mixing of temperature and salinity to estimate the potential energy
!! change due to diapycnal mixing in a column of water.  It does this estimate
!! 4 different ways, all of which should be equivalent, but reports only one.
!! The various estimates are taken because they will later be used as templates
!! for other bits of code
subroutine diapyc_energy_req_calc(h_in, T_in, S_in, Kd, energy_Kd, dt, tv, &
                                  G, GV, US, may_print, 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(GV%ke),   intent(in)    :: h_in !< Layer thickness before entrainment,
                                                  !! [H ~> m or kg m-2].
  real, dimension(GV%ke),   intent(in)    :: t_in !< The layer temperatures [degC].
  real, dimension(GV%ke),   intent(in)    :: s_in !< The layer salinities [ppt].
  real, dimension(GV%ke+1), intent(in)    :: kd   !< The interfaces diapycnal diffusivities
                                                  !! [Z2 T-1 ~> m2 s-1].
  real,                     intent(in)    :: dt   !< The amount of time covered by this call [T ~> s].
  real,                     intent(out)   :: energy_kd !< The column-integrated rate of energy
                                                  !! consumption by diapycnal diffusion [R Z L2 T-3 ~> W m-2].
  type(thermo_var_ptrs),    intent(inout) :: tv   !< A structure containing pointers to any
                                                  !! available thermodynamic fields.
                                                  !! Absent fields have NULL ptrs.
  logical,        optional, intent(in)    :: may_print !< If present and true, write out diagnostics
                                                  !! of energy use.
  type(diapyc_energy_req_cs), &
                  optional, pointer       :: cs   !< This module's control structure.

!   This subroutine uses a substantially refactored tridiagonal equation for
! diapycnal mixing of temperature and salinity to estimate the potential energy
! change due to diapycnal mixing in a column of water.  It does this estimate
! 4 different ways, all of which should be equivalent, but reports only one.
! The various estimates are taken because they will later be used as templates
! for other bits of code.

  real, dimension(GV%ke) :: &
    p_lay, &    ! Average pressure of a layer [R L2 T-2 ~> Pa].
    dsv_dt, &   ! Partial derivative of specific volume with temperature [R-1 degC-1 ~> m3 kg-1 degC-1].
    dsv_ds, &   ! Partial derivative of specific volume with salinity [R-1 ppt-1 ~> m3 kg-1 ppt-1].
    t0, s0, &   ! Initial temperatures and salinities [degC] and [ppt].
    te, se, &   ! Running incomplete estimates of the new temperatures and salinities [degC] and [ppt].
    te_a, se_a, & ! Running incomplete estimates of the new temperatures and salinities [degC] and [ppt].
    te_b, se_b, & ! Running incomplete estimates of the new temperatures and salinities [degC] and [ppt].
    tf, sf, &   ! New final values of the temperatures and salinities [degC] and [ppt].
    dte, dse, & ! Running (1-way) estimates of temperature and salinity change [degC] and [ppt].
    dte_a, dse_a, & ! Running (1-way) estimates of temperature and salinity change [degC] and [ppt].
    dte_b, dse_b, & ! Running (1-way) estimates of temperature and salinity change [degC] and [ppt].
    th_a, &     ! An effective temperature times a thickness in the layer above, including implicit
                ! mixing effects with other yet higher layers [degC H ~> degC m or degC kg m-2].
    sh_a, &     ! An effective salinity times a thickness in the layer above, including implicit
                ! mixing effects with other yet higher layers [ppt H ~> ppt m or ppt kg m-2].
    th_b, &     ! An effective temperature times a thickness in the layer below, including implicit
                ! mixing effects with other yet lower layers [degC H ~> degC m or degC kg m-2].
    sh_b, &     ! An effective salinity times a thickness in the layer below, including implicit
                ! mixing effects with other yet lower layers [ppt H ~> ppt m or ppt kg m-2].
    dt_to_dpe, & ! Partial derivative of column potential energy with the temperature and salinity
    ds_to_dpe, & ! changes within a layer [R Z L2 T-2 degC-1 ~> J m-2 degC-1] and [R Z L2 T-2 ppt-1 ~> J m-2 ppt-1].
    dt_to_dcolht, & ! Partial derivatives of the total column height with the temperature
    ds_to_dcolht, & ! and salinity changes within a layer [Z degC-1 ~> m degC-1] and [Z ppt-1 ~> m ppt-1].
    dt_to_dcolht_a, & ! Partial derivatives of the total column height with the temperature
    ds_to_dcolht_a, & ! and salinity changes within a layer, including the implicit effects
                    ! of mixing with layers higher in the water colun [Z degC-1 ~> m degC-1] and [Z ppt-1 ~> m ppt-1].
    dt_to_dcolht_b, & ! Partial derivatives of the total column height with the temperature
    ds_to_dcolht_b, & ! and salinity changes within a layer, including the implicit effects
                    ! of mixing with layers lower in the water colun [Z degC-1 ~> m degC-1] and [Z ppt-1 ~> m ppt-1].
    dt_to_dpe_a, &  ! Partial derivatives of column potential energy with the temperature
    ds_to_dpe_a, &  ! and salinity changes within a layer, including the implicit effects
                    ! of mixing with layers higher in the water column, in
                    ! units of [R Z L2 T-2 degC-1 ~> J m-2 degC-1] and [R Z L2 T-2 ppt-1 ~> J m-2 ppt-1].
    dt_to_dpe_b, &  ! Partial derivatives of column potential energy with the temperature
    ds_to_dpe_b, &  ! and salinity changes within a layer, including the implicit effects
                    ! of mixing with layers lower in the water column, in
                    ! units of [R Z L2 T-2 degC-1 ~> J m-2 degC-1] and [R Z L2 T-2 ppt-1 ~> J m-2 ppt-1].
    hp_a, &     ! An effective pivot thickness of the layer including the effects
                ! of coupling with layers above [H ~> m or kg m-2].  This is the first term
                ! in the denominator of b1 in a downward-oriented tridiagonal solver.
    hp_b, &     ! An effective pivot thickness of the layer including the effects
                ! of coupling with layers below [H ~> m or kg m-2].  This is the first term
                ! in the denominator of b1 in an upward-oriented tridiagonal solver.
    c1_a, &     ! c1_a is used by a downward-oriented tridiagonal solver [nondim].
    c1_b, &     ! c1_b is used by an upward-oriented tridiagonal solver [nondim].
    h_tr        ! h_tr is h at tracer points with a h_neglect added to
                ! ensure positive definiteness [H ~> m or kg m-2].
  real, dimension(GV%ke+1) :: &
    pres, &     ! Interface pressures [R L2 T-2 ~> Pa].
    pres_z, &   ! Interface pressures with a rescaling factor to convert interface height
                ! movements into changes in column potential energy [R L2 T-2 m Z-1 ~> J m-3].
    z_int, &    ! Interface heights relative to the surface [H ~> m or kg m-2].
    n2, &       ! An estimate of the buoyancy frequency [T-2 ~> s-2].
    kddt_h, &   ! The diapycnal diffusivity times a timestep divided by the
                ! average thicknesses around a layer [H ~> m or kg m-2].
    kddt_h_a, & ! The value of Kddt_h for layers above the central point in the
                ! tridiagonal solver [H ~> m or kg m-2].
    kddt_h_b, & ! The value of Kddt_h for layers below the central point in the
                ! tridiagonal solver [H ~> m or kg m-2].
    kd_so_far   ! The value of Kddt_h that has been applied already in
                ! calculating the energy changes [H ~> m or kg m-2].
  real, dimension(GV%ke+1,4) :: &
    pe_chg_k, &     ! The integrated potential energy change within a timestep due
                    ! to the diffusivity at interface K for 4 different orders of
                    ! accumulating the diffusivities [R Z L2 T-2 ~> J m-2].
    colht_cor_k     ! The correction to the potential energy change due to
                    ! changes in the net column height [R Z L2 T-2 ~> J m-2].
  real :: &
    b1              ! b1 is used by the tridiagonal solver [H-1 ~> m-1 or m2 kg-1].
  real :: &
    i_b1            ! The inverse of b1 [H ~> m or kg m-2].
  real :: kd0       ! The value of Kddt_h that has already been applied [H ~> m or kg m-2].
  real :: dkd       ! The change in the value of Kddt_h [H ~> m or kg m-2].
  real :: h_neglect ! A thickness that is so small it is usually lost
                    ! in roundoff and can be neglected [H ~> m or kg m-2].
  real :: dte_term  ! A diffusivity-independent term related to the temperature
                    ! change in the layer below the interface [degC H ~> degC m or degC kg m-2].
  real :: dse_term  ! A diffusivity-independent term related to the salinity
                    ! change in the layer below the interface [ppt H ~> ppt m or ppt kg m-2].
  real :: kddt_h_guess ! A guess of the final value of Kddt_h [H ~> m or kg m-2].
  real :: dmass     ! The mass per unit area within a layer [R Z ~> kg m-2].
  real :: dpres     ! The hydrostatic pressure change across a layer [R L2 T-2 ~> Pa].
  real :: dmke_max  ! The maximum amount of mean kinetic energy that could be
                    ! converted to turbulent kinetic energy if the velocity in
                    ! the layer below an interface were homogenized with all of
                    ! the water above the interface [R Z L2 T-2 ~> J m-2 = kg s-2].
  real :: rho_here  ! The in-situ density [R ~> kg m-3].
  real :: pe_change ! The change in column potential energy from applying Kddt_h at the
                    ! present interface [R L2 Z T-2 ~> J m-2].
  real :: colht_cor ! The correction to PE_chg that is made due to a net
                    ! change in the column height [R L2 Z T-2 ~> J m-2].
  real :: htot      ! A running sum of thicknesses [H ~> m or kg m-2].
  real :: dte_t2, dt_km1_t2, dt_k_t2  ! Temporary arrays describing temperature changes [degC].
  real :: dse_t2, ds_km1_t2, ds_k_t2  ! Temporary arrays describing salinity changes [ppt].
  logical :: do_print

  ! The following are a bunch of diagnostic arrays for debugging purposes.
  real, dimension(GV%ke) :: &
    ta, sa, tb, sb
  real, dimension(GV%ke+1) :: &
    dpea_dkd, dpea_dkd_est, dpea_dkd_err, dpea_dkd_trunc, dpea_dkd_err_norm, &
    dpeb_dkd, dpeb_dkd_est, dpeb_dkd_err, dpeb_dkd_trunc, dpeb_dkd_err_norm
  real :: pe_chg_tot1a, pe_chg_tot2a, t_chg_tota
  real :: pe_chg_tot1b, pe_chg_tot2b, t_chg_totb
  real :: pe_chg_tot1c, pe_chg_tot2c, t_chg_totc
  real :: pe_chg_tot1d, pe_chg_tot2d, t_chg_totd
  real, dimension(GV%ke+1)  :: dpechg_dkd
  real :: pe_chg(6)
  real, dimension(6) :: dt_k_itt, ds_k_itt, dt_km1_itt, ds_km1_itt

  integer :: k, nz, itt, max_itt, k_cent
  logical :: surface_bl, bottom_bl, central, halves, debug
  logical :: old_pe_calc
  nz = g%ke
  h_neglect = gv%H_subroundoff

  debug = .true.

  surface_bl = .true. ; bottom_bl = .true. ; halves = .true.
  central = .true. ; k_cent = nz/2

  do_print = .false. ; if (present(may_print) .and. present(cs)) do_print = may_print

  dpea_dkd(:) = 0.0 ; dpea_dkd_est(:) = 0.0 ; dpea_dkd_err(:) = 0.0
  dpea_dkd_err_norm(:) = 0.0 ; dpea_dkd_trunc(:) = 0.0
  dpeb_dkd(:) = 0.0 ; dpeb_dkd_est(:) = 0.0 ; dpeb_dkd_err(:) = 0.0
  dpeb_dkd_err_norm(:) = 0.0 ; dpeb_dkd_trunc(:) = 0.0

  htot = 0.0 ; pres(1) = 0.0 ; pres_z(1) = 0.0 ; z_int(1) = 0.0
  do k=1,nz
    t0(k) = t_in(k) ; s0(k) = s_in(k)
    h_tr(k) = h_in(k)
    htot = htot + h_tr(k)
    pres(k+1) = pres(k) + (gv%g_Earth * gv%H_to_RZ) * h_tr(k)
    pres_z(k+1) = pres(k+1)
    p_lay(k) = 0.5*(pres(k) + pres(k+1))
    z_int(k+1) = z_int(k) - h_tr(k)
  enddo
  do k=1,nz
    h_tr(k) = max(h_tr(k), 1e-15*htot)
  enddo

  ! Introduce a diffusive flux variable, Kddt_h(K) = ea(k) = eb(k-1)

  kddt_h(1) = 0.0 ; kddt_h(nz+1) = 0.0
  do k=2,nz
    kddt_h(k) = min((gv%Z_to_H**2*dt)*kd(k) / (0.5*(h_tr(k-1) + h_tr(k))), 1e3*htot)
  enddo

  ! Solve the tridiagonal equations for new temperatures.

  call calculate_specific_vol_derivs(t0, s0, p_lay, dsv_dt, dsv_ds, tv%eqn_of_state)

  do k=1,nz
    dmass = gv%H_to_RZ * h_tr(k)
    dpres = (gv%g_Earth * gv%H_to_RZ) * h_tr(k)
    dt_to_dpe(k) = (dmass * (pres(k) + 0.5*dpres)) * dsv_dt(k)
    ds_to_dpe(k) = (dmass * (pres(k) + 0.5*dpres)) * dsv_ds(k)
    dt_to_dcolht(k) = dmass * dsv_dt(k) * cs%ColHt_scaling
    ds_to_dcolht(k) = dmass * dsv_ds(k) * cs%ColHt_scaling
  enddo

!  PE_chg_k(1) = 0.0 ; PE_chg_k(nz+1) = 0.0
  ! PEchg(:) = 0.0
  pe_chg_k(:,:) = 0.0 ; colht_cor_k(:,:) = 0.0
  dpechg_dkd(:) = 0.0

  if (surface_bl) then  ! This version is appropriate for a surface boundary layer.
    old_pe_calc = .false.

    ! Set up values appropriate for no diffusivity.
    do k=1,nz
      hp_a(k) = h_tr(k) ; hp_b(k) = h_tr(k)
      dt_to_dpe_a(k) = dt_to_dpe(k) ; ds_to_dpe_a(k) = ds_to_dpe(k)
      dt_to_dpe_b(k) = dt_to_dpe(k) ; ds_to_dpe_b(k) = ds_to_dpe(k)
      dt_to_dcolht_a(k) = dt_to_dcolht(k) ; ds_to_dcolht_a(k) = ds_to_dcolht(k)
      dt_to_dcolht_b(k) = dt_to_dcolht(k) ; ds_to_dcolht_b(k) = ds_to_dcolht(k)
    enddo

    do k=2,nz
      ! At this point, Kddt_h(K) will be unknown because its value may depend
      ! on how much energy is available.

      ! Precalculate some temporary expressions that are independent of Kddt_h_guess.
      if (old_pe_calc) then
        if (k==2) then
          dt_km1_t2 = (t0(k)-t0(k-1))
          ds_km1_t2 = (s0(k)-s0(k-1))
          dte_t2 = 0.0 ; dse_t2 = 0.0
        else
          dte_t2 = kddt_h(k-1) * ((t0(k-2) - t0(k-1)) + dte(k-2))
          dse_t2 = kddt_h(k-1) * ((s0(k-2) - s0(k-1)) + dse(k-2))
          dt_km1_t2 = (t0(k)-t0(k-1)) - &
                (kddt_h(k-1) / hp_a(k-1)) * ((t0(k-2) - t0(k-1)) + dte(k-2))
          ds_km1_t2 = (s0(k)-s0(k-1)) - &
                (kddt_h(k-1) / hp_a(k-1)) * ((s0(k-2) - s0(k-1)) + dse(k-2))
        endif
        dte_term = dte_t2 + hp_a(k-1) * (t0(k-1)-t0(k))
        dse_term = dse_t2 + hp_a(k-1) * (s0(k-1)-s0(k))
      else
        if (k<=2) then
          th_a(k-1) = h_tr(k-1) * t0(k-1) ; sh_a(k-1) = h_tr(k-1) * s0(k-1)
        else
          th_a(k-1) = h_tr(k-1) * t0(k-1) + kddt_h(k-1) * te(k-2)
          sh_a(k-1) = h_tr(k-1) * s0(k-1) + kddt_h(k-1) * se(k-2)
        endif
        th_b(k) = h_tr(k) * t0(k) ; sh_b(k) = h_tr(k) * s0(k)
      endif


      ! Find the energy change due to a guess at the strength of diffusion at interface K.

      kddt_h_guess = kddt_h(k)
      if (old_pe_calc) then
        call find_pe_chg_orig(kddt_h_guess, h_tr(k), hp_a(k-1), &
                         dte_term, dse_term, dt_km1_t2, ds_km1_t2, &
                         dt_to_dpe(k), ds_to_dpe(k), dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), &
                         pres_z(k), dt_to_dcolht(k), ds_to_dcolht(k), &
                         dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                         pe_chg_k(k,1), dpea_dkd(k))
      else
        call find_pe_chg(0.0, kddt_h_guess, hp_a(k-1), hp_b(k), &
                         th_a(k-1), sh_a(k-1), th_b(k), sh_b(k), &
                         dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                         pres_z(k), dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                         dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                         pe_chg=pe_chg_k(k,1), dpec_dkd=dpea_dkd(k), &
                         colht_cor=colht_cor_k(k,1))
      endif

      if (debug) then
        do itt=1,5
          kddt_h_guess = (1.0+0.01*(itt-3))*kddt_h(k)

          if (old_pe_calc) then
            call find_pe_chg_orig(kddt_h_guess, h_tr(k), hp_a(k-1), &
                             dte_term, dse_term, dt_km1_t2, ds_km1_t2, &
                             dt_to_dpe(k), ds_to_dpe(k), dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), &
                             pres_z(k), dt_to_dcolht(k), ds_to_dcolht(k), &
                             dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                             pe_chg=pe_chg(itt))
          else
            call find_pe_chg(0.0, kddt_h_guess, hp_a(k-1), hp_b(k), &
                             th_a(k-1), sh_a(k-1), th_b(k), sh_b(k), &
                             dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                             pres_z(k), dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                             dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                             pe_chg=pe_chg(itt))
          endif
        enddo
        ! Compare with a 4th-order finite difference estimate.
        dpea_dkd_est(k) = (4.0*(pe_chg(4)-pe_chg(2))/(0.02*kddt_h(k)) - &
                               (pe_chg(5)-pe_chg(1))/(0.04*kddt_h(k))) / 3.0
        dpea_dkd_trunc(k) = (pe_chg(4)-pe_chg(2))/(0.02*kddt_h(k)) - &
                            (pe_chg(5)-pe_chg(1))/(0.04*kddt_h(k))
        dpea_dkd_err(k) = (dpea_dkd_est(k) - dpea_dkd(k))
        dpea_dkd_err_norm(k) = (dpea_dkd_est(k) - dpea_dkd(k)) / &
                              (abs(dpea_dkd_est(k)) + abs(dpea_dkd(k)) + 1e-100*us%RZ_to_kg_m2*us%L_T_to_m_s**2)
      endif

      !   At this point, the final value of Kddt_h(K) is known, so the estimated
      ! properties for layer k-1 can be calculated.

      b1 = 1.0 / (hp_a(k-1) + kddt_h(k))
      c1_a(k) = kddt_h(k) * b1
      if (k==2) then
        te(1) = b1*(h_tr(1)*t0(1))
        se(1) = b1*(h_tr(1)*s0(1))
      else
        te(k-1) = b1 * (h_tr(k-1) * t0(k-1) + kddt_h(k-1) * te(k-2))
        se(k-1) = b1 * (h_tr(k-1) * s0(k-1) + kddt_h(k-1) * se(k-2))
      endif
      if (old_pe_calc) then
        dte(k-1) = b1 * ( kddt_h(k)*(t0(k)-t0(k-1)) + dte_t2 )
        dse(k-1) = b1 * ( kddt_h(k)*(s0(k)-s0(k-1)) + dse_t2 )
      endif

      hp_a(k) = h_tr(k) + (hp_a(k-1) * b1) * kddt_h(k)
      dt_to_dpe_a(k) = dt_to_dpe(k) + c1_a(k)*dt_to_dpe_a(k-1)
      ds_to_dpe_a(k) = ds_to_dpe(k) + c1_a(k)*ds_to_dpe_a(k-1)
      dt_to_dcolht_a(k) = dt_to_dcolht(k) + c1_a(k)*dt_to_dcolht_a(k-1)
      ds_to_dcolht_a(k) = ds_to_dcolht(k) + c1_a(k)*ds_to_dcolht_a(k-1)

    enddo

    b1 = 1.0 / (hp_a(nz))
    tf(nz) = b1 * (h_tr(nz) * t0(nz) + kddt_h(nz) * te(nz-1))
    sf(nz) = b1 * (h_tr(nz) * s0(nz) + kddt_h(nz) * se(nz-1))
    if (old_pe_calc) then
      dte(nz) = b1 * kddt_h(nz) * ((t0(nz-1)-t0(nz)) + dte(nz-1))
      dse(nz) = b1 * kddt_h(nz) * ((s0(nz-1)-s0(nz)) + dse(nz-1))
    endif

    do k=nz-1,1,-1
      tf(k) = te(k) + c1_a(k+1)*tf(k+1)
      sf(k) = se(k) + c1_a(k+1)*sf(k+1)
    enddo

    if (debug) then
      do k=1,nz ; ta(k) = tf(k) ; sa(k) = sf(k) ; enddo
      pe_chg_tot1a = 0.0 ; pe_chg_tot2a = 0.0 ; t_chg_tota = 0.0
      do k=1,nz
        pe_chg_tot1a = pe_chg_tot1a + (dt_to_dpe(k) * (tf(k) - t0(k)) + &
                                       ds_to_dpe(k) * (sf(k) - s0(k)))
        t_chg_tota = t_chg_tota + h_tr(k) * (tf(k) - t0(k))
      enddo
      do k=2,nz
        pe_chg_tot2a = pe_chg_tot2a + (pe_chg_k(k,1) - colht_cor_k(k,1))
      enddo
    endif

  endif

  if (bottom_bl) then  ! This version is appropriate for a bottom boundary layer.
    old_pe_calc = .false.

    ! Set up values appropriate for no diffusivity.
    do k=1,nz
      hp_a(k) = h_tr(k) ; hp_b(k) = h_tr(k)
      dt_to_dpe_a(k) = dt_to_dpe(k) ; ds_to_dpe_a(k) = ds_to_dpe(k)
      dt_to_dpe_b(k) = dt_to_dpe(k) ; ds_to_dpe_b(k) = ds_to_dpe(k)
      dt_to_dcolht_a(k) = dt_to_dcolht(k) ; ds_to_dcolht_a(k) = ds_to_dcolht(k)
      dt_to_dcolht_b(k) = dt_to_dcolht(k) ; ds_to_dcolht_b(k) = ds_to_dcolht(k)
    enddo

    do k=nz,2,-1  ! Loop over interior interfaces.
      ! At this point, Kddt_h(K) will be unknown because its value may depend
      ! on how much energy is available.

      ! Precalculate some temporary expressions that are independent of Kddt_h_guess.
      if (old_pe_calc) then
        if (k==nz) then
          dt_k_t2 = (t0(k-1)-t0(k))
          ds_k_t2 = (s0(k-1)-s0(k))
          dte_t2 = 0.0 ; dse_t2 = 0.0
        else
          dte_t2 = kddt_h(k+1) * ((t0(k+1) - t0(k)) + dte(k+1))
          dse_t2 = kddt_h(k+1) * ((s0(k+1) - s0(k)) + dse(k+1))
          dt_k_t2 = (t0(k-1)-t0(k)) - &
                  (kddt_h(k+1)/ hp_b(k)) * ((t0(k+1) - t0(k)) + dte(k+1))
          ds_k_t2 = (s0(k-1)-s0(k)) - &
                  (kddt_h(k+1)/ hp_b(k)) * ((s0(k+1) - s0(k)) + dse(k+1))
        endif
        dte_term = dte_t2 + hp_b(k) * (t0(k)-t0(k-1))
        dse_term = dse_t2 + hp_b(k) * (s0(k)-s0(k-1))
      else
        th_a(k-1) = h_tr(k-1) * t0(k-1) ; sh_a(k-1) = h_tr(k-1) * s0(k-1)
        if (k>=nz) then
          th_b(k) = h_tr(k) * t0(k) ; sh_b(k) = h_tr(k) * s0(k)
        else
          th_b(k) = h_tr(k) * t0(k) + kddt_h(k+1) * te(k+1)
          sh_b(k) = h_tr(k) * s0(k) + kddt_h(k+1) * se(k+1)
        endif
      endif

      ! Find the energy change due to a guess at the strength of diffusion at interface K.
      kddt_h_guess = kddt_h(k)

      if (old_pe_calc) then
        call find_pe_chg_orig(kddt_h_guess, h_tr(k-1), hp_b(k), &
                         dte_term, dse_term, dt_k_t2, ds_k_t2, &
                         dt_to_dpe(k-1), ds_to_dpe(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                         pres_z(k), dt_to_dcolht(k-1), ds_to_dcolht(k-1), &
                         dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                         pe_chg=pe_chg_k(k,2), dpec_dkd=dpeb_dkd(k))
      else
        call find_pe_chg(0.0, kddt_h_guess, hp_a(k-1), hp_b(k), &
                         th_a(k-1), sh_a(k-1), th_b(k), sh_b(k), &
                         dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                         pres_z(k), dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                         dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                         pe_chg=pe_chg_k(k,2), dpec_dkd=dpeb_dkd(k), &
                         colht_cor=colht_cor_k(k,2))
      endif

      if (debug) then
        ! Compare with a 4th-order finite difference estimate.
        do itt=1,5
          kddt_h_guess = (1.0+0.01*(itt-3))*kddt_h(k)

          if (old_pe_calc) then
            call find_pe_chg_orig(kddt_h_guess, h_tr(k-1), hp_b(k), &
                           dte_term, dse_term, dt_k_t2, ds_k_t2, &
                           dt_to_dpe(k-1), ds_to_dpe(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                           pres_z(k), dt_to_dcolht(k-1), ds_to_dcolht(k-1), &
                           dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                           pe_chg=pe_chg(itt))
          else
            call find_pe_chg(0.0, kddt_h_guess, hp_a(k-1), hp_b(k), &
                             th_a(k-1), sh_a(k-1), th_b(k), sh_b(k), &
                             dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                             pres_z(k), dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                             dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                             pe_chg=pe_chg(itt))
          endif
        enddo

        dpeb_dkd_est(k) = (4.0*(pe_chg(4)-pe_chg(2))/(0.02*kddt_h(k)) - &
                               (pe_chg(5)-pe_chg(1))/(0.04*kddt_h(k))) / 3.0
        dpeb_dkd_trunc(k) = (pe_chg(4)-pe_chg(2))/(0.02*kddt_h(k)) - &
                            (pe_chg(5)-pe_chg(1))/(0.04*kddt_h(k))
        dpeb_dkd_err(k) = (dpeb_dkd_est(k) - dpeb_dkd(k))
        dpeb_dkd_err_norm(k) = (dpeb_dkd_est(k) - dpeb_dkd(k)) / &
                              (abs(dpeb_dkd_est(k)) + abs(dpeb_dkd(k)) + 1e-100*us%RZ_to_kg_m2*us%L_T_to_m_s**2)
      endif

      !   At this point, the final value of Kddt_h(K) is known, so the estimated
      ! properties for layer k can be calculated.

      b1 = 1.0 / (hp_b(k) + kddt_h(k))
      c1_b(k) = kddt_h(k) * b1
      if (k==nz) then
        te(nz) = b1* (h_tr(nz)*t0(nz))
        se(nz) = b1* (h_tr(nz)*s0(nz))
      else
        te(k) = b1 * (h_tr(k) * t0(k) + kddt_h(k+1) * te(k+1))
        se(k) = b1 * (h_tr(k) * s0(k) + kddt_h(k+1) * se(k+1))
      endif
      if (old_pe_calc) then
        dte(k) = b1 * ( kddt_h(k)*(t0(k-1)-t0(k)) + dte_t2 )
        dse(k) = b1 * ( kddt_h(k)*(s0(k-1)-s0(k)) + dse_t2 )
      endif

      hp_b(k-1) = h_tr(k-1) + (hp_b(k) * b1) * kddt_h(k)
      dt_to_dpe_b(k-1) = dt_to_dpe(k-1) + c1_b(k)*dt_to_dpe_b(k)
      ds_to_dpe_b(k-1) = ds_to_dpe(k-1) + c1_b(k)*ds_to_dpe_b(k)
      dt_to_dcolht_b(k-1) = dt_to_dcolht(k-1) + c1_b(k)*dt_to_dcolht_b(k)
      ds_to_dcolht_b(k-1) = ds_to_dcolht(k-1) + c1_b(k)*ds_to_dcolht_b(k)

    enddo

    b1 = 1.0 / (hp_b(1))
    tf(1) = b1 * (h_tr(1) * t0(1) + kddt_h(2) * te(2))
    sf(1) = b1 * (h_tr(1) * s0(1) + kddt_h(2) * se(2))
    if (old_pe_calc) then
      dte(1) = b1 * kddt_h(2) * ((t0(2)-t0(1)) + dte(2))
      dse(1) = b1 * kddt_h(2) * ((s0(2)-s0(1)) + dse(2))
    endif

    do k=2,nz
      tf(k) = te(k) + c1_b(k)*tf(k-1)
      sf(k) = se(k) + c1_b(k)*sf(k-1)
    enddo

    if (debug) then
      do k=1,nz ; tb(k) = tf(k) ; sb(k) = sf(k) ; enddo
      pe_chg_tot1b = 0.0 ; pe_chg_tot2b = 0.0 ; t_chg_totb = 0.0
      do k=1,nz
        pe_chg_tot1b = pe_chg_tot1b + (dt_to_dpe(k) * (tf(k) - t0(k)) + &
                                       ds_to_dpe(k) * (sf(k) - s0(k)))
        t_chg_totb = t_chg_totb + h_tr(k) * (tf(k) - t0(k))
      enddo
      do k=2,nz
        pe_chg_tot2b = pe_chg_tot2b + (pe_chg_k(k,2) - colht_cor_k(k,2))
      enddo
    endif

  endif

  if (central) then

    ! Set up values appropriate for no diffusivity.
    do k=1,nz
      hp_a(k) = h_tr(k) ; hp_b(k) = h_tr(k)
      dt_to_dpe_a(k) = dt_to_dpe(k) ; ds_to_dpe_a(k) = ds_to_dpe(k)
      dt_to_dpe_b(k) = dt_to_dpe(k) ; ds_to_dpe_b(k) = ds_to_dpe(k)
      dt_to_dcolht_a(k) = dt_to_dcolht(k) ; ds_to_dcolht_a(k) = ds_to_dcolht(k)
      dt_to_dcolht_b(k) = dt_to_dcolht(k) ; ds_to_dcolht_b(k) = ds_to_dcolht(k)
    enddo

    ! Calculate the dependencies on layers above.
    kddt_h_a(1) = 0.0
    do k=2,nz  ! Loop over interior interfaces.
      ! First calculate some terms that are independent of the change in Kddt_h(K).
      kd0 = 0.0  ! This might need to be changed - it is the already applied value of Kddt_h(K).
      if (k<=2) then
        th_a(k-1) = h_tr(k-1) * t0(k-1) ; sh_a(k-1) = h_tr(k-1) * s0(k-1)
      else
        th_a(k-1) = h_tr(k-1) * t0(k-1) + kddt_h(k-1) * te_a(k-2)
        sh_a(k-1) = h_tr(k-1) * s0(k-1) + kddt_h(k-1) * se_a(k-2)
      endif
      th_b(k) = h_tr(k) * t0(k) ; sh_b(k) = h_tr(k) * s0(k)

      kddt_h_a(k) = 0.0 ; if (k < k_cent) kddt_h_a(k) = kddt_h(k)
      dkd = kddt_h_a(k)

      call find_pe_chg(kd0, dkd, hp_a(k-1), hp_b(k), &
                       th_a(k-1), sh_a(k-1), th_b(k), sh_b(k), &
                       dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                       pres_z(k), dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                       dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                       pe_chg=pe_change, colht_cor=colht_cor)
      pe_chg_k(k,3) = pe_change
      colht_cor_k(k,3) = colht_cor

      b1 = 1.0 / (hp_a(k-1) + kddt_h_a(k))
      c1_a(k) = kddt_h_a(k) * b1
      if (k==2) then
        te_a(1) = b1*(h_tr(1)*t0(1))
        se_a(1) = b1*(h_tr(1)*s0(1))
      else
        te_a(k-1) = b1 * (h_tr(k-1) * t0(k-1) + kddt_h_a(k-1) * te_a(k-2))
        se_a(k-1) = b1 * (h_tr(k-1) * s0(k-1) + kddt_h_a(k-1) * se_a(k-2))
      endif

      hp_a(k) = h_tr(k) + (hp_a(k-1) * b1) * kddt_h_a(k)
      dt_to_dpe_a(k) = dt_to_dpe(k) + c1_a(k)*dt_to_dpe_a(k-1)
      ds_to_dpe_a(k) = ds_to_dpe(k) + c1_a(k)*ds_to_dpe_a(k-1)
      dt_to_dcolht_a(k) = dt_to_dcolht(k) + c1_a(k)*dt_to_dcolht_a(k-1)
      ds_to_dcolht_a(k) = ds_to_dcolht(k) + c1_a(k)*ds_to_dcolht_a(k-1)
    enddo

    ! Calculate the dependencies on layers below.
    kddt_h_b(nz+1) = 0.0
    do k=nz,2,-1  ! Loop over interior interfaces.
      ! First calculate some terms that are independent of the change in Kddt_h(K).
      kd0 = 0.0  ! This might need to be changed - it is the already applied value of Kddt_h(K).
!     if (K<=2) then
        th_a(k-1) = h_tr(k-1) * t0(k-1) ; sh_a(k-1) = h_tr(k-1) * s0(k-1)
!     else
!       Th_a(k-1) = h_tr(k-1) * T0(k-1) + Kddt_h(K-1) * Te_a(k-2)
!       Sh_a(k-1) = h_tr(k-1) * S0(k-1) + Kddt_h(K-1) * Se_a(k-2)
!     endif
      if (k>=nz) then
        th_b(k) = h_tr(k) * t0(k) ; sh_b(k) = h_tr(k) * s0(k)
      else
        th_b(k) = h_tr(k) * t0(k) + kddt_h(k+1) * te_b(k+1)
        sh_b(k) = h_tr(k) * s0(k) + kddt_h(k+1) * se_b(k+1)
      endif

      kddt_h_b(k) = 0.0 ; if (k > k_cent) kddt_h_b(k) = kddt_h(k)
      dkd = kddt_h_b(k)

      call find_pe_chg(kd0, dkd, hp_a(k-1), hp_b(k), &
                       th_a(k-1), sh_a(k-1), th_b(k), sh_b(k), &
                       dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                       pres_z(k), dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                       dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                       pe_chg=pe_change, colht_cor=colht_cor)
      pe_chg_k(k,3) = pe_chg_k(k,3) + pe_change
      colht_cor_k(k,3) = colht_cor_k(k,3) + colht_cor

      b1 = 1.0 / (hp_b(k) + kddt_h_b(k))
      c1_b(k) = kddt_h_b(k) * b1
      if (k==nz) then
        te_b(k) = b1 * (h_tr(k)*t0(k))
        se_b(k) = b1 * (h_tr(k)*s0(k))
      else
        te_b(k) = b1 * (h_tr(k) * t0(k) + kddt_h_b(k+1) * te_b(k+1))
        se_b(k) = b1 * (h_tr(k) * s0(k) + kddt_h_b(k+1) * se_b(k+1))
      endif

      hp_b(k-1) = h_tr(k-1) + (hp_b(k) * b1) * kddt_h_b(k)
      dt_to_dpe_b(k-1) = dt_to_dpe(k-1) + c1_b(k)*dt_to_dpe_b(k)
      ds_to_dpe_b(k-1) = ds_to_dpe(k-1) + c1_b(k)*ds_to_dpe_b(k)
      dt_to_dcolht_b(k-1) = dt_to_dcolht(k-1) + c1_b(k)*dt_to_dcolht_b(k)
      ds_to_dcolht_b(k-1) = ds_to_dcolht(k-1) + c1_b(k)*ds_to_dcolht_b(k)

    enddo

    ! Calculate the final solution for the layers surrounding interface K_cent
    k = k_cent

    ! First calculate some terms that are independent of the change in Kddt_h(K).
    kd0 = 0.0  ! This might need to be changed - it is the already applied value of Kddt_h(K).
    if (k<=2) then
      th_a(k-1) = h_tr(k-1) * t0(k-1) ; sh_a(k-1) = h_tr(k-1) * s0(k-1)
    else
      th_a(k-1) = h_tr(k-1) * t0(k-1) + kddt_h(k-1) * te_a(k-2)
      sh_a(k-1) = h_tr(k-1) * s0(k-1) + kddt_h(k-1) * se_a(k-2)
    endif
    if (k>=nz) then
      th_b(k) = h_tr(k) * t0(k) ; sh_b(k) = h_tr(k) * s0(k)
    else
      th_b(k) = h_tr(k) * t0(k) + kddt_h(k+1) * te_b(k+1)
      sh_b(k) = h_tr(k) * s0(k) + kddt_h(k+1) * se_b(k+1)
    endif

    dkd = kddt_h(k)

    call find_pe_chg(kd0, dkd, hp_a(k-1), hp_b(k), &
                     th_a(k-1), sh_a(k-1), th_b(k), sh_b(k), &
                     dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                     pres_z(k), dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                     dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                     pe_chg=pe_change, colht_cor=colht_cor)
    pe_chg_k(k,3) = pe_chg_k(k,3) + pe_change
    colht_cor_k(k,3) = colht_cor_k(k,3) + colht_cor


    !   To derive the following, first to a partial update for the estimated
    ! temperatures and salinities in the layers around this interface:
    !   b1_a = 1.0 / (hp_a(k-1) + Kddt_h(K))
    !   b1_b = 1.0 / (hp_b(k) + Kddt_h(K))
    !   Te_up(k) = Th_b(k) * b1_b ; Se_up(k) = Sh_b(k) * b1_b
    !   Te_up(k-1) = Th_a(k-1) * b1_a ; Se_up(k-1) = Sh_a(k-1) * b1_a
    ! Find the final values of T & S for both layers around K_cent, using that
    !   c1_a(K) = Kddt_h(K) * b1_a ; c1_b(K) = Kddt_h(K) * b1_b
    !   Tf(K_cent-1) = Te_up(K_cent-1) + c1_a(K_cent)*Tf(K_cent)
    !   Tf(K_cent)   = Te_up(K_cent) + c1_b(K_cent)*Tf(K_cent-1)
    ! but further exploiting the expressions for c1_a and c1_b to avoid
    ! subtraction in the denominator, and use only a single division.
    b1 = 1.0 / (hp_a(k-1)*hp_b(k) + kddt_h(k)*(hp_a(k-1) + hp_b(k)))
    tf(k-1) = ((hp_b(k) + kddt_h(k)) * th_a(k-1) + kddt_h(k) * th_b(k) ) * b1
    sf(k-1) = ((hp_b(k) + kddt_h(k)) * sh_a(k-1) + kddt_h(k) * sh_b(k) ) * b1
    tf(k) = (kddt_h(k) * th_a(k-1) + (hp_a(k-1) + kddt_h(k)) * th_b(k) ) * b1
    sf(k) = (kddt_h(k) * sh_a(k-1) + (hp_a(k-1) + kddt_h(k)) * sh_b(k) ) * b1

    c1_a(k) = kddt_h(k) / (hp_a(k-1) + kddt_h(k))
    c1_b(k) = kddt_h(k) / (hp_b(k) + kddt_h(k))

    ! Now update the other layer working outward from k_cent to determine the final
    ! temperatures and salinities.
    do k=k_cent-2,1,-1
      tf(k) = te_a(k) + c1_a(k+1)*tf(k+1)
      sf(k) = se_a(k) + c1_a(k+1)*sf(k+1)
    enddo
    do k=k_cent+1,nz
      tf(k) = te_b(k) + c1_b(k)*tf(k-1)
      sf(k) = se_b(k) + c1_b(k)*sf(k-1)
    enddo

    if (debug) then
      pe_chg_tot1c = 0.0 ; pe_chg_tot2c = 0.0 ; t_chg_totc = 0.0
      do k=1,nz
        pe_chg_tot1c = pe_chg_tot1c + (dt_to_dpe(k) * (tf(k) - t0(k)) + &
                                       ds_to_dpe(k) * (sf(k) - s0(k)))
        t_chg_totc = t_chg_totc + h_tr(k) * (tf(k) - t0(k))
      enddo
      do k=2,nz
        pe_chg_tot2c = pe_chg_tot2c + (pe_chg_k(k,3) - colht_cor_k(k,3))
      enddo
    endif

  endif

  if (halves) then

    ! Set up values appropriate for no diffusivity.
    do k=1,nz
      hp_a(k) = h_tr(k) ; hp_b(k) = h_tr(k)
      dt_to_dpe_a(k) = dt_to_dpe(k) ; ds_to_dpe_a(k) = ds_to_dpe(k)
      dt_to_dpe_b(k) = dt_to_dpe(k) ; ds_to_dpe_b(k) = ds_to_dpe(k)
      dt_to_dcolht_a(k) = dt_to_dcolht(k) ; ds_to_dcolht_a(k) = ds_to_dcolht(k)
      dt_to_dcolht_b(k) = dt_to_dcolht(k) ; ds_to_dcolht_b(k) = ds_to_dcolht(k)
    enddo
    do k=1,nz+1
      kd_so_far(k) = 0.0
    enddo

    ! Calculate the dependencies on layers above.
    kddt_h_a(1) = 0.0
    do k=2,nz  ! Loop over interior interfaces.
      ! First calculate some terms that are independent of the change in Kddt_h(K).
      kd0 = kd_so_far(k)
      if (k<=2) then
        th_a(k-1) = h_tr(k-1) * t0(k-1) ; sh_a(k-1) = h_tr(k-1) * s0(k-1)
      else
        th_a(k-1) = h_tr(k-1) * t0(k-1) + kd_so_far(k-1) * te(k-2)
        sh_a(k-1) = h_tr(k-1) * s0(k-1) + kd_so_far(k-1) * se(k-2)
      endif
      th_b(k) = h_tr(k) * t0(k) ; sh_b(k) = h_tr(k) * s0(k)

      dkd = 0.5 * kddt_h(k) - kd_so_far(k)

      call find_pe_chg(kd0, dkd, hp_a(k-1), hp_b(k), &
                       th_a(k-1), sh_a(k-1), th_b(k), sh_b(k), &
                       dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                       pres_z(k), dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                       dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                       pe_chg=pe_change, colht_cor=colht_cor)

      pe_chg_k(k,4) = pe_change
      colht_cor_k(k,4) = colht_cor

      kd_so_far(k) = kd_so_far(k) + dkd

      b1 = 1.0 / (hp_a(k-1) + kd_so_far(k))
      c1_a(k) = kd_so_far(k) * b1
      if (k==2) then
        te(1) = b1*(h_tr(1)*t0(1))
        se(1) = b1*(h_tr(1)*s0(1))
      else
        te(k-1) = b1 * (h_tr(k-1) * t0(k-1) + kd_so_far(k-1) * te(k-2))
        se(k-1) = b1 * (h_tr(k-1) * s0(k-1) + kd_so_far(k-1) * se(k-2))
      endif

      hp_a(k) = h_tr(k) + (hp_a(k-1) * b1) * kd_so_far(k)
      dt_to_dpe_a(k) = dt_to_dpe(k) + c1_a(k)*dt_to_dpe_a(k-1)
      ds_to_dpe_a(k) = ds_to_dpe(k) + c1_a(k)*ds_to_dpe_a(k-1)
      dt_to_dcolht_a(k) = dt_to_dcolht(k) + c1_a(k)*dt_to_dcolht_a(k-1)
      ds_to_dcolht_a(k) = ds_to_dcolht(k) + c1_a(k)*ds_to_dcolht_a(k-1)
    enddo

    ! Calculate the dependencies on layers below.
    do k=nz,2,-1  ! Loop over interior interfaces.
      ! First calculate some terms that are independent of the change in Kddt_h(K).
      kd0 = kd_so_far(k)
      if (k<=2) then
        th_a(k-1) = h_tr(k-1) * t0(k-1) ; sh_a(k-1) = h_tr(k-1) * s0(k-1)
      else
        th_a(k-1) = h_tr(k-1) * t0(k-1) + kd_so_far(k-1) * te(k-2)
        sh_a(k-1) = h_tr(k-1) * s0(k-1) + kd_so_far(k-1) * se(k-2)
      endif
      if (k>=nz) then
        th_b(k) = h_tr(k) * t0(k) ; sh_b(k) = h_tr(k) * s0(k)
      else
        th_b(k) = h_tr(k) * t0(k) + kd_so_far(k+1) * te(k+1)
        sh_b(k) = h_tr(k) * s0(k) + kd_so_far(k+1) * se(k+1)
      endif

      dkd = kddt_h(k) - kd_so_far(k)

      call find_pe_chg(kd0, dkd, hp_a(k-1), hp_b(k), &
                       th_a(k-1), sh_a(k-1), th_b(k), sh_b(k), &
                       dt_to_dpe_a(k-1), ds_to_dpe_a(k-1), dt_to_dpe_b(k), ds_to_dpe_b(k), &
                       pres_z(k), dt_to_dcolht_a(k-1), ds_to_dcolht_a(k-1), &
                       dt_to_dcolht_b(k), ds_to_dcolht_b(k), &
                       pe_chg=pe_change, colht_cor=colht_cor)

      pe_chg_k(k,4) = pe_chg_k(k,4) + pe_change
      colht_cor_k(k,4) = colht_cor_k(k,4) + colht_cor


      kd_so_far(k) = kd_so_far(k) + dkd

      b1 = 1.0 / (hp_b(k) + kd_so_far(k))
      c1_b(k) = kd_so_far(k) * b1
      if (k==nz) then
        te(k) = b1 * (h_tr(k)*t0(k))
        se(k) = b1 * (h_tr(k)*s0(k))
      else
        te(k) = b1 * (h_tr(k) * t0(k) + kd_so_far(k+1) * te(k+1))
        se(k) = b1 * (h_tr(k) * s0(k) + kd_so_far(k+1) * se(k+1))
      endif

      hp_b(k-1) = h_tr(k-1) + (hp_b(k) * b1) * kd_so_far(k)
      dt_to_dpe_b(k-1) = dt_to_dpe(k-1) + c1_b(k)*dt_to_dpe_b(k)
      ds_to_dpe_b(k-1) = ds_to_dpe(k-1) + c1_b(k)*ds_to_dpe_b(k)
      dt_to_dcolht_b(k-1) = dt_to_dcolht(k-1) + c1_b(k)*dt_to_dcolht_b(k)
      ds_to_dcolht_b(k-1) = ds_to_dcolht(k-1) + c1_b(k)*ds_to_dcolht_b(k)

    enddo

    ! Now update the other layer working down from the top to determine the
    ! final temperatures and salinities.
    b1 = 1.0 / (hp_b(1))
    tf(1) = b1 * (h_tr(1) * t0(1) + kddt_h(2) * te(2))
    sf(1) = b1 * (h_tr(1) * s0(1) + kddt_h(2) * se(2))
    do k=2,nz
      tf(k) = te(k) + c1_b(k)*tf(k-1)
      sf(k) = se(k) + c1_b(k)*sf(k-1)
    enddo

    if (debug) then
      pe_chg_tot1d = 0.0 ; pe_chg_tot2d = 0.0 ; t_chg_totd = 0.0
      do k=1,nz
        pe_chg_tot1d = pe_chg_tot1d + (dt_to_dpe(k) * (tf(k) - t0(k)) + &
                                       ds_to_dpe(k) * (sf(k) - s0(k)))
        t_chg_totd = t_chg_totd + h_tr(k) * (tf(k) - t0(k))
      enddo
      do k=2,nz
        pe_chg_tot2d = pe_chg_tot2d + (pe_chg_k(k,4) - colht_cor_k(k,4))
      enddo
    endif

  endif

  energy_kd = 0.0 ; do k=2,nz ; energy_kd = energy_kd + pe_chg_k(k,1) ; enddo
  energy_kd = energy_kd / dt

  if (do_print) then
    if (cs%id_ERt>0) call post_data(cs%id_ERt, pe_chg_k(:,1), cs%diag)
    if (cs%id_ERb>0) call post_data(cs%id_ERb, pe_chg_k(:,2), cs%diag)
    if (cs%id_ERc>0) call post_data(cs%id_ERc, pe_chg_k(:,3), cs%diag)
    if (cs%id_ERh>0) call post_data(cs%id_ERh, pe_chg_k(:,4), cs%diag)
    if (cs%id_Kddt>0) call post_data(cs%id_Kddt, kddt_h, cs%diag)
    if (cs%id_Kd>0)   call post_data(cs%id_Kd,   kd, cs%diag)
    if (cs%id_h>0)    call post_data(cs%id_h,    h_tr, cs%diag)
    if (cs%id_zInt>0) call post_data(cs%id_zInt, z_int, cs%diag)
    if (cs%id_CHCt>0) call post_data(cs%id_CHCt, colht_cor_k(:,1), cs%diag)
    if (cs%id_CHCb>0) call post_data(cs%id_CHCb, colht_cor_k(:,2), cs%diag)
    if (cs%id_CHCc>0) call post_data(cs%id_CHCc, colht_cor_k(:,3), cs%diag)
    if (cs%id_CHCh>0) call post_data(cs%id_CHCh, colht_cor_k(:,4), cs%diag)
    if (cs%id_T0>0) call post_data(cs%id_T0, t0, cs%diag)
    if (cs%id_Tf>0) call post_data(cs%id_Tf, tf, cs%diag)
    if (cs%id_S0>0) call post_data(cs%id_S0, s0, cs%diag)
    if (cs%id_Sf>0) call post_data(cs%id_Sf, sf, cs%diag)
    if (cs%id_N2_0>0) then
      n2(1) = 0.0 ; n2(nz+1) = 0.0
      do k=2,nz
        call calculate_density(0.5*(t0(k-1) + t0(k)), 0.5*(s0(k-1) + s0(k)), &
                               pres(k), rho_here, tv%eqn_of_state)
        n2(k) = ((us%L_to_Z**2*gv%g_Earth) * rho_here / (0.5*gv%H_to_Z*(h_tr(k-1) + h_tr(k)))) * &
                ( 0.5*(dsv_dt(k-1) + dsv_dt(k)) * (t0(k-1) - t0(k)) + &
                  0.5*(dsv_ds(k-1) + dsv_ds(k)) * (s0(k-1) - s0(k)) )
      enddo
      call post_data(cs%id_N2_0, n2, cs%diag)
    endif
    if (cs%id_N2_f>0) then
      n2(1) = 0.0 ; n2(nz+1) = 0.0
      do k=2,nz
        call calculate_density(0.5*(tf(k-1) + tf(k)), 0.5*(sf(k-1) + sf(k)), &
                               pres(k), rho_here, tv%eqn_of_state)
        n2(k) = ((us%L_to_Z**2*gv%g_Earth) * rho_here / (0.5*gv%H_to_Z*(h_tr(k-1) + h_tr(k)))) * &
                ( 0.5*(dsv_dt(k-1) + dsv_dt(k)) * (tf(k-1) - tf(k)) + &
                  0.5*(dsv_ds(k-1) + dsv_ds(k)) * (sf(k-1) - sf(k)) )
      enddo
      call post_data(cs%id_N2_f, n2, cs%diag)
    endif
  endif

end subroutine diapyc_energy_req_calc

!> This subroutine calculates the change in potential energy and or derivatives
!! for several changes in an interfaces's diapycnal diffusivity times a timestep.
subroutine find_pe_chg(Kddt_h0, dKddt_h, hp_a, hp_b, Th_a, Sh_a, Th_b, Sh_b, &
                       dT_to_dPE_a, dS_to_dPE_a, dT_to_dPE_b, dS_to_dPE_b, &
                       pres_Z, dT_to_dColHt_a, dS_to_dColHt_a, dT_to_dColHt_b, dS_to_dColHt_b, &
                       PE_chg, dPEc_dKd, dPE_max, dPEc_dKd_0, ColHt_cor)
  real, intent(in)  :: Kddt_h0  !< The previously used diffusivity at an interface times
                                !! the time step and  divided by the average of the
                                !! thicknesses around the interface [H ~> m or kg m-2].
  real, intent(in)  :: dKddt_h  !< The trial change in the diffusivity at an interface times
                                !! the time step and  divided by the average of the
                                !! thicknesses around the interface [H ~> m or kg m-2].
  real, intent(in)  :: hp_a     !< The effective pivot thickness of the layer above the
                                !! interface, given by h_k plus a term that
                                !! is a fraction (determined from the tridiagonal solver) of
                                !! Kddt_h for the interface above [H ~> m or kg m-2].
  real, intent(in)  :: hp_b     !< The effective pivot thickness of the layer below the
                                !! interface, given by h_k plus a term that
                                !! is a fraction (determined from the tridiagonal solver) of
                                !! Kddt_h for the interface above [H ~> m or kg m-2].
  real, intent(in)  :: Th_a     !< An effective temperature times a thickness in the layer
                                !! above, including implicit mixing effects with other
                                !! yet higher layers [degC H ~> degC m or degC kg m-2].
  real, intent(in)  :: Sh_a     !< An effective salinity times a thickness in the layer
                                !! above, including implicit mixing effects with other
                                !! yet higher layers [ppt H ~> ppt m or ppt kg m-2].
  real, intent(in)  :: Th_b     !< An effective temperature times a thickness in the layer
                                !! below, including implicit mixing effects with other
                                !! yet lower layers [degC H ~> degC m or degC kg m-2].
  real, intent(in)  :: Sh_b     !< An effective salinity times a thickness in the layer
                                !! below, including implicit mixing effects with other
                                !! yet lower layers [ppt H ~> ppt m or ppt kg m-2].
  real, intent(in)  :: dT_to_dPE_a !< A factor (pres_lay*mass_lay*dSpec_vol/dT) relating
                                !! a layer's temperature change to the change in column
                                !! potential energy, including all implicit diffusive changes
                                !! in the temperatures of all the layers above [R Z L2 T-2 degC-1 ~> J m-2 degC-1].
  real, intent(in)  :: dS_to_dPE_a !< A factor (pres_lay*mass_lay*dSpec_vol/dS) relating
                                !! a layer's salinity change to the change in column
                                !! potential energy, including all implicit diffusive changes
                                !! in the salinities of all the layers above [R Z L2 T-2 ppt-1 ~> J m-2 ppt-1].
  real, intent(in)  :: dT_to_dPE_b !< A factor (pres_lay*mass_lay*dSpec_vol/dT) relating
                                !! a layer's temperature change to the change in column
                                !! potential energy, including all implicit diffusive changes
                                !! in the temperatures of all the layers below [R Z L2 T-2 degC-1 ~> J m-2 degC-1].
  real, intent(in)  :: dS_to_dPE_b !< A factor (pres_lay*mass_lay*dSpec_vol/dS) relating
                                !! a layer's salinity change to the change in column
                                !! potential energy, including all implicit diffusive changes
                                !! in the salinities of all the layers below [R Z L2 T-2 ppt-1 ~> J m-2 ppt-1].
  real, intent(in)  :: pres_Z   !< The hydrostatic interface pressure, which is used to relate
                                !! the changes in column thickness to the energy that is radiated
                                !! as gravity waves and unavailable to drive mixing [R L2 T-2 m Z-1 ~> J m-3].
  real, intent(in)  :: dT_to_dColHt_a !< A factor (mass_lay*dSColHtc_vol/dT) relating
                                !! a layer's temperature change to the change in column
                                !! height, including all implicit diffusive changes
                                !! in the temperatures of all the layers above [Z degC-1 ~> m degC-1].
  real, intent(in)  :: dS_to_dColHt_a !< A factor (mass_lay*dSColHtc_vol/dS) relating
                                !! a layer's salinity change to the change in column
                                !! height, including all implicit diffusive changes
                                !! in the salinities of all the layers above [Z ppt-1 ~> m ppt-1].
  real, intent(in)  :: dT_to_dColHt_b !< A factor (mass_lay*dSColHtc_vol/dT) relating
                                !! a layer's temperature change to the change in column
                                !! height, including all implicit diffusive changes
                                !! in the temperatures of all the layers below [Z degC-1 ~> m degC-1].
  real, intent(in)  :: dS_to_dColHt_b !< A factor (mass_lay*dSColHtc_vol/dS) relating
                                !! a layer's salinity change to the change in column
                                !! height, including all implicit diffusive changes
                                !! in the salinities of all the layers below [Z ppt-1 ~> m ppt-1].

  real, optional, intent(out) :: PE_chg   !< The change in column potential energy from applying
                                          !! Kddt_h at the present interface [R Z L2 T-2 ~> J m-2].
  real, optional, intent(out) :: dPEc_dKd !< The partial derivative of PE_chg with Kddt_h,
                                          !! [R Z L2 T-2 H-1 ~> J m-3 or J kg-1].
  real, optional, intent(out) :: dPE_max  !< The maximum change in column potential energy that could
                                          !! be realizedd by applying a huge value of Kddt_h at the
                                          !! present interface [R Z L2 T-2 ~> J m-2].
  real, optional, intent(out) :: dPEc_dKd_0 !< The partial derivative of PE_chg with Kddt_h in the
                                            !! limit where Kddt_h = 0 [R Z L2 T-2 H-1 ~> J m-3 or J kg-1].
  real, optional, intent(out) :: ColHt_cor  !< The correction to PE_chg that is made due to a net
                                            !! change in the column height [R Z L2 T-2 ~> J m-2].

  real :: hps  ! The sum of the two effective pivot thicknesses [H ~> m or kg m-2].
  real :: bdt1 ! A product of the two pivot thicknesses plus a diffusive term [H2 ~> m2 or kg2 m-4].
  real :: dT_c ! The core term in the expressions for the temperature changes [degC H2 ~> degC m2 or degC kg2 m-4].
  real :: dS_c ! The core term in the expressions for the salinity changes [psu H2 ~> psu m2 or psu kg2 m-4].
  real :: PEc_core ! The diffusivity-independent core term in the expressions
                   ! for the potential energy changes [R L2 T-2 ~> J m-3].
  real :: ColHt_core ! The diffusivity-independent core term in the expressions
                     ! for the column height changes [R L2 T-2 ~> J m-3].
  real :: ColHt_chg  ! The change in the column height [Z ~> m].
  real :: y1   ! A local temporary term, in [H-3] or [H-4] in various contexts.

  !   The expression for the change in potential energy used here is derived
  ! from the expression for the final estimates of the changes in temperature
  ! and salinities, and then extensively manipulated to get it into its most
  ! succint form. The derivation is not necessarily obvious, but it demonstrably
  ! works by comparison with separate calculations of the energy changes after
  ! the tridiagonal solver for the final changes in temperature and salinity are
  ! applied.

  hps = hp_a + hp_b
  bdt1 = hp_a * hp_b + kddt_h0 * hps
  dt_c = hp_a * th_b - hp_b * th_a
  ds_c = hp_a * sh_b - hp_b * sh_a
  pec_core = hp_b * (dt_to_dpe_a * dt_c + ds_to_dpe_a * ds_c) - &
             hp_a * (dt_to_dpe_b * dt_c + ds_to_dpe_b * ds_c)
  colht_core = hp_b * (dt_to_dcolht_a * dt_c + ds_to_dcolht_a * ds_c) - &
               hp_a * (dt_to_dcolht_b * dt_c + ds_to_dcolht_b * ds_c)

  if (present(pe_chg)) then
    ! Find the change in column potential energy due to the change in the
    ! diffusivity at this interface by dKddt_h.
    y1 = dkddt_h / (bdt1 * (bdt1 + dkddt_h * hps))
    pe_chg = pec_core * y1
    colht_chg = colht_core * y1
    if (colht_chg < 0.0) pe_chg = pe_chg - pres_z * colht_chg
    if (present(colht_cor)) colht_cor = -pres_z * min(colht_chg, 0.0)
  elseif (present(colht_cor)) then
    y1 = dkddt_h / (bdt1 * (bdt1 + dkddt_h * hps))
    colht_cor = -pres_z * min(colht_core * y1, 0.0)
  endif

  if (present(dpec_dkd)) then
    ! Find the derivative of the potential energy change with dKddt_h.
    y1 = 1.0 / (bdt1 + dkddt_h * hps)**2
    dpec_dkd = pec_core * y1
    colht_chg = colht_core * y1
    if (colht_chg < 0.0) dpec_dkd = dpec_dkd - pres_z * colht_chg
  endif

  if (present(dpe_max)) then
    ! This expression is the limit of PE_chg for infinite dKddt_h.
    y1 = 1.0 / (bdt1 * hps)
    dpe_max = pec_core * y1
    colht_chg = colht_core * y1
    if (colht_chg < 0.0) dpe_max = dpe_max - pres_z * colht_chg
  endif

  if (present(dpec_dkd_0)) then
    ! This expression is the limit of dPEc_dKd for dKddt_h = 0.
    y1 = 1.0 / bdt1**2
    dpec_dkd_0 = pec_core * y1
    colht_chg = colht_core * y1
    if (colht_chg < 0.0) dpec_dkd_0 = dpec_dkd_0 - pres_z * colht_chg
  endif

end subroutine find_pe_chg


!> This subroutine calculates the change in potential energy and or derivatives
!! for several changes in an interfaces's diapycnal diffusivity times a timestep
!! using the original form used in the first version of ePBL.
subroutine find_pe_chg_orig(Kddt_h, h_k, b_den_1, dTe_term, dSe_term, &
                            dT_km1_t2, dS_km1_t2, dT_to_dPE_k, dS_to_dPE_k, &
                            dT_to_dPEa, dS_to_dPEa, pres_Z, dT_to_dColHt_k, &
                            dS_to_dColHt_k, dT_to_dColHta, dS_to_dColHta, &
                            PE_chg, dPEc_dKd, dPE_max, dPEc_dKd_0)
  real, intent(in)  :: Kddt_h   !< The diffusivity at an interface times the time step and
                                !! divided by the average of the thicknesses around the
                                !! interface [H ~> m or kg m-2].
  real, intent(in)  :: h_k      !< The thickness of the layer below the interface [H ~> m or kg m-2].
  real, intent(in)  :: b_den_1  !< The first term in the denominator of the pivot
                                !! for the tridiagonal solver, given by h_k plus a term that
                                !! is a fraction (determined from the tridiagonal solver) of
                                !! Kddt_h for the interface above [H ~> m or kg m-2].
  real, intent(in)  :: dTe_term !< A diffusivity-independent term related to the temperature change
                                !! in the layer below the interface [degC H ~> degC m or degC kg m-2].
  real, intent(in)  :: dSe_term !< A diffusivity-independent term related to the salinity change
                                !! in the layer below the interface [ppt H ~> ppt m or ppt kg m-2].
  real, intent(in)  :: dT_km1_t2 !< A diffusivity-independent term related to the
                                 !! temperature change in the layer above the interface [degC].
  real, intent(in)  :: dS_km1_t2 !< A diffusivity-independent term related to the
                                 !! salinity change in the layer above the interface [ppt].
  real, intent(in)  :: pres_Z    !< The hydrostatic interface pressure, which is used to relate
                                 !! the changes in column thickness to the energy that is radiated
                                 !! as gravity waves and unavailable to drive mixing [R L2 T-2 ~> J m-3].
  real, intent(in)  :: dT_to_dPE_k !< A factor (pres_lay*mass_lay*dSpec_vol/dT) relating
                                 !! a layer's temperature change to the change in column
                                 !! potential energy, including all implicit diffusive changes
                                 !! in the temperatures of all the layers below [R Z L2 T-2 degC-1 ~> J m-2 degC-1].
  real, intent(in)  :: dS_to_dPE_k !< A factor (pres_lay*mass_lay*dSpec_vol/dS) relating
                                 !! a layer's salinity change to the change in column
                                 !! potential energy, including all implicit diffusive changes
                                 !! in the salinities of all the layers below [R Z L2 T-2 ppt-1 ~> J m-2 ppt-1].
  real, intent(in)  :: dT_to_dPEa !< A factor (pres_lay*mass_lay*dSpec_vol/dT) relating
                                 !! a layer's temperature change to the change in column
                                 !! potential energy, including all implicit diffusive changes
                                 !! in the temperatures of all the layers above [R Z L2 T-2 degC-1 ~> J m-2 degC-1].
  real, intent(in)  :: dS_to_dPEa !< A factor (pres_lay*mass_lay*dSpec_vol/dS) relating
                                 !! a layer's salinity change to the change in column
                                 !! potential energy, including all implicit diffusive changes
                                 !! in the salinities of all the layers above [R Z L2 T-2 ppt-1 ~> J m-2 ppt-1].
  real, intent(in)  :: dT_to_dColHt_k !< A factor (mass_lay*dSColHtc_vol/dT) relating
                                 !! a layer's temperature change to the change in column
                                 !! height, including all implicit diffusive changes
                                 !! in the temperatures of all the layers below [Z degC-1 ~> m degC-1].
  real, intent(in)  :: dS_to_dColHt_k !< A factor (mass_lay*dSColHtc_vol/dS) relating
                                 !! a layer's salinity change to the change in column
                                 !! height, including all implicit diffusive changes
                                 !! in the salinities of all the layers below [Z ppt-1 ~> m ppt-1].
  real, intent(in)  :: dT_to_dColHta !< A factor (mass_lay*dSColHtc_vol/dT) relating
                                 !! a layer's temperature change to the change in column
                                 !! height, including all implicit diffusive changes
                                 !! in the temperatures of all the layers above [Z degC-1 ~> m degC-1].
  real, intent(in)  :: dS_to_dColHta !< A factor (mass_lay*dSColHtc_vol/dS) relating
                                 !! a layer's salinity change to the change in column
                                 !! height, including all implicit diffusive changes
                                 !! in the salinities of all the layers above [Z ppt-1 ~> m ppt-1].

  real, optional, intent(out) :: PE_chg   !< The change in column potential energy from applying
                                          !! Kddt_h at the present interface [R Z L2 T-2 ~> J m-2].
  real, optional, intent(out) :: dPEc_dKd !< The partial derivative of PE_chg with Kddt_h,
                                          !! [R Z L2 T-2 H-1 ~> J m-3 or J kg-1].
  real, optional, intent(out) :: dPE_max  !< The maximum change in column potential energy that could
                                          !! be realized by applying a huge value of Kddt_h at the
                                          !! present interface [R Z L2 T-2 ~> J m-2].
  real, optional, intent(out) :: dPEc_dKd_0 !< The partial derivative of PE_chg with Kddt_h in the
                                            !! limit where Kddt_h = 0 [R Z L2 T-2 H-1 ~> J m-3 or J kg-1].

!   This subroutine determines the total potential energy change due to mixing
! at an interface, including all of the implicit effects of the prescribed
! mixing at interfaces above.  Everything here is derived by careful manipulation
! of the robust tridiagonal solvers used for tracers by MOM6.  The results are
! positive for mixing in a stably stratified environment.
!   The comments describing these arguments are for a downward mixing pass, but
! this routine can also be used for an upward pass with the sense of direction
! reversed.

  real :: b1            ! b1 is used by the tridiagonal solver [H-1 ~> m-1 or m2 kg-1].
  real :: b1Kd          ! Temporary array [nondim]
  real :: ColHt_chg     ! The change in column thickness [Z ~> m].
  real :: dColHt_max    ! The change in column thickness for infinite diffusivity [Z ~> m].
  real :: dColHt_dKd    ! The partial derivative of column thickness with Kddt_h [Z H-1 ~> 1 or m3 kg-1].
  real :: dT_k, dT_km1  ! Temporary arrays [degC].
  real :: dS_k, dS_km1  ! Temporary arrays [ppt].
  real :: I_Kr_denom    ! Temporary arrays [H-2 ~> m-2 or m4 kg-2].
  real :: dKr_dKd       ! Nondimensional temporary array [nondim].
  real :: ddT_k_dKd, ddT_km1_dKd ! Temporary arrays [degC H-1 ~> m-1 or m2 kg-1].
  real :: ddS_k_dKd, ddS_km1_dKd ! Temporary arrays [ppt H-1 ~> ppt m-1 or ppt m2 kg-1].

  b1 = 1.0 / (b_den_1 + kddt_h)
  b1kd = kddt_h*b1

  ! Start with the temperature change in layer k-1 due to the diffusivity at
  ! interface K without considering the effects of changes in layer k.

  ! Calculate the change in PE due to the diffusion at interface K
  ! if Kddt_h(K+1) = 0.
  i_kr_denom = 1.0 / (h_k*b_den_1 + (b_den_1 + h_k)*kddt_h)

  dt_k = (kddt_h*i_kr_denom) * dte_term
  ds_k = (kddt_h*i_kr_denom) * dse_term

  if (present(pe_chg)) then
    ! Find the change in energy due to diffusion with strength Kddt_h at this interface.
    ! Increment the temperature changes in layer k-1 due the changes in layer k.
    dt_km1 = b1kd * ( dt_k + dt_km1_t2 )
    ds_km1 = b1kd * ( ds_k + ds_km1_t2 )

    pe_chg = (dt_to_dpe_k * dt_k + dt_to_dpea * dt_km1) + &
             (ds_to_dpe_k * ds_k + ds_to_dpea * ds_km1)
    colht_chg = (dt_to_dcolht_k * dt_k + dt_to_dcolhta * dt_km1) + &
                (ds_to_dcolht_k * ds_k + ds_to_dcolhta * ds_km1)
    if (colht_chg < 0.0) pe_chg = pe_chg - pres_z * colht_chg
  endif

  if (present(dpec_dkd)) then
    ! Find the derivatives of the temperature and salinity changes with Kddt_h.
    dkr_dkd = (h_k*b_den_1) * i_kr_denom**2

    ddt_k_dkd = dkr_dkd * dte_term
    dds_k_dkd = dkr_dkd * dse_term
    ddt_km1_dkd = (b1**2 * b_den_1) * ( dt_k + dt_km1_t2 ) + b1kd * ddt_k_dkd
    dds_km1_dkd = (b1**2 * b_den_1) * ( ds_k + ds_km1_t2 ) + b1kd * dds_k_dkd

    ! Calculate the partial derivative of Pe_chg with Kddt_h.
    dpec_dkd = (dt_to_dpe_k * ddt_k_dkd + dt_to_dpea * ddt_km1_dkd) + &
               (ds_to_dpe_k * dds_k_dkd + ds_to_dpea * dds_km1_dkd)
    dcolht_dkd = (dt_to_dcolht_k * ddt_k_dkd + dt_to_dcolhta * ddt_km1_dkd) + &
                 (ds_to_dcolht_k * dds_k_dkd + ds_to_dcolhta * dds_km1_dkd)
    if (dcolht_dkd < 0.0) dpec_dkd = dpec_dkd - pres_z * dcolht_dkd
  endif

  if (present(dpe_max)) then
    ! This expression is the limit of PE_chg for infinite Kddt_h.
    dpe_max = (dt_to_dpea * dt_km1_t2 + ds_to_dpea * ds_km1_t2) + &
              ((dt_to_dpe_k + dt_to_dpea) * dte_term + &
               (ds_to_dpe_k + ds_to_dpea) * dse_term) / (b_den_1 + h_k)
    dcolht_max = (dt_to_dcolhta * dt_km1_t2 + ds_to_dcolhta * ds_km1_t2) + &
              ((dt_to_dcolht_k + dt_to_dcolhta) * dte_term + &
               (ds_to_dcolht_k + ds_to_dcolhta) * dse_term) / (b_den_1 + h_k)
    if (dcolht_max < 0.0) dpe_max = dpe_max - pres_z*dcolht_max
  endif

  if (present(dpec_dkd_0)) then
    ! This expression is the limit of dPEc_dKd for Kddt_h = 0.
    dpec_dkd_0 = (dt_to_dpea * dt_km1_t2 + ds_to_dpea * ds_km1_t2) / (b_den_1) + &
                 (dt_to_dpe_k * dte_term + ds_to_dpe_k * dse_term) / (h_k*b_den_1)
    dcolht_dkd = (dt_to_dcolhta * dt_km1_t2 + ds_to_dcolhta * ds_km1_t2) / (b_den_1) + &
                 (dt_to_dcolht_k * dte_term + ds_to_dcolht_k * dse_term) / (h_k*b_den_1)
    if (dcolht_dkd < 0.0) dpec_dkd_0 = dpec_dkd_0 - pres_z*dcolht_dkd
  endif

end subroutine find_pe_chg_orig

!> Initialize parameters and allocate memory associated with the diapycnal energy requirement module.
subroutine diapyc_energy_req_init(Time, G, GV, US, param_file, diag, CS)
  type(time_type),            intent(in)    :: time        !< model time
  type(ocean_grid_type),      intent(in)    :: g           !< model grid structure
  type(verticalgrid_type),    intent(in)    :: gv          !< ocean vertical grid structure
  type(unit_scale_type),      intent(in)    :: us          !< A dimensional unit scaling type
  type(param_file_type),      intent(in)    :: param_file  !< file to parse for parameter values
  type(diag_ctrl),    target, intent(inout) :: diag        !< structure to regulate diagnostic output
  type(diapyc_energy_req_cs), pointer       :: cs          !< module control structure

  integer, save :: init_calls = 0
! This include declares and sets the variable "version".
#include "version_variable.h"
  character(len=40)  :: mdl = "MOM_diapyc_energy_req" ! This module's name.
  character(len=256) :: mesg    ! Message for error messages.

  if (.not.associated(cs)) then ; allocate(cs)
  else ; return ; endif

  cs%initialized = .true.
  cs%diag => diag

  ! Read all relevant parameters and write them to the model log.
  call log_version(param_file, mdl, version, "")
  call get_param(param_file, mdl, "ENERGY_REQ_KH_SCALING", cs%test_Kh_scaling, &
                 "A scaling factor for the diapycnal diffusivity used in "//&
                 "testing the energy requirements.", default=1.0, units="nondim")
  call get_param(param_file, mdl, "ENERGY_REQ_COL_HT_SCALING", cs%ColHt_scaling, &
                 "A scaling factor for the column height change correction "//&
                 "used in testing the energy requirements.", default=1.0, units="nondim")
  call get_param(param_file, mdl, "ENERGY_REQ_USE_TEST_PROFILE", &
                 cs%use_test_Kh_profile, &
                 "If true, use the internal test diffusivity profile in "//&
                 "place of any that might be passed in as an argument.", default=.false.)

  cs%id_ERt = register_diag_field('ocean_model', 'EnReqTest_ERt', diag%axesZi, time, &
                 "Diffusivity Energy Requirements, top-down", &
                 "J m-2", conversion=us%RZ_to_kg_m2*us%L_T_to_m_s**2)
  cs%id_ERb = register_diag_field('ocean_model', 'EnReqTest_ERb', diag%axesZi, time, &
                 "Diffusivity Energy Requirements, bottom-up", &
                 "J m-2", conversion=us%RZ_to_kg_m2*us%L_T_to_m_s**2)
  cs%id_ERc = register_diag_field('ocean_model', 'EnReqTest_ERc', diag%axesZi, time, &
                 "Diffusivity Energy Requirements, center-last", &
                 "J m-2", conversion=us%RZ_to_kg_m2*us%L_T_to_m_s**2)
  cs%id_ERh = register_diag_field('ocean_model', 'EnReqTest_ERh', diag%axesZi, time, &
                 "Diffusivity Energy Requirements, halves", &
                 "J m-2", conversion=us%RZ_to_kg_m2*us%L_T_to_m_s**2)
  cs%id_Kddt = register_diag_field('ocean_model', 'EnReqTest_Kddt', diag%axesZi, time, &
                 "Implicit diffusive coupling coefficient", "m", conversion=gv%H_to_m)
  cs%id_Kd = register_diag_field('ocean_model', 'EnReqTest_Kd', diag%axesZi, time, &
                 "Diffusivity in test", "m2 s-1", conversion=us%Z_to_m**2)
  cs%id_h   = register_diag_field('ocean_model', 'EnReqTest_h_lay', diag%axesZL, time, &
                 "Test column layer thicknesses", "m", conversion=gv%H_to_m)
  cs%id_zInt = register_diag_field('ocean_model', 'EnReqTest_z_int', diag%axesZi, time, &
                 "Test column layer interface heights", "m", conversion=gv%H_to_m)
  cs%id_CHCt = register_diag_field('ocean_model', 'EnReqTest_CHCt', diag%axesZi, time, &
                 "Column Height Correction to Energy Requirements, top-down", &
                 "J m-2", conversion=us%RZ_to_kg_m2*us%L_T_to_m_s**2)
  cs%id_CHCb = register_diag_field('ocean_model', 'EnReqTest_CHCb', diag%axesZi, time, &
                 "Column Height Correction to Energy Requirements, bottom-up", &
                 "J m-2", conversion=us%RZ_to_kg_m2*us%L_T_to_m_s**2)
  cs%id_CHCc = register_diag_field('ocean_model', 'EnReqTest_CHCc', diag%axesZi, time, &
                 "Column Height Correction to Energy Requirements, center-last", &
                 "J m-2", conversion=us%RZ_to_kg_m2*us%L_T_to_m_s**2)
  cs%id_CHCh = register_diag_field('ocean_model', 'EnReqTest_CHCh', diag%axesZi, time, &
                 "Column Height Correction to Energy Requirements, halves", &
                 "J m-2", conversion=us%RZ_to_kg_m2*us%L_T_to_m_s**2)
  cs%id_T0 = register_diag_field('ocean_model', 'EnReqTest_T0', diag%axesZL, time, &
                 "Temperature before mixing", "deg C")
  cs%id_Tf = register_diag_field('ocean_model', 'EnReqTest_Tf', diag%axesZL, time, &
                 "Temperature after mixing", "deg C")
  cs%id_S0 = register_diag_field('ocean_model', 'EnReqTest_S0', diag%axesZL, time, &
                 "Salinity before mixing", "g kg-1")
  cs%id_Sf = register_diag_field('ocean_model', 'EnReqTest_Sf', diag%axesZL, time, &
                 "Salinity after mixing", "g kg-1")
  cs%id_N2_0 = register_diag_field('ocean_model', 'EnReqTest_N2_0', diag%axesZi, time, &
                 "Squared buoyancy frequency before mixing", "second-2", conversion=us%s_to_T**2)
  cs%id_N2_f = register_diag_field('ocean_model', 'EnReqTest_N2_f', diag%axesZi, time, &
                 "Squared buoyancy frequency after mixing", "second-2", conversion=us%s_to_T**2)

end subroutine diapyc_energy_req_init

!> Clean up and deallocate memory associated with the diapycnal energy requirement module.
subroutine diapyc_energy_req_end(CS)
  type(diapyc_energy_req_cs), pointer :: cs !< Diapycnal energy requirement control structure that
                                            !! will be deallocated in this subroutine.
  if (associated(cs)) deallocate(cs)
end subroutine diapyc_energy_req_end

end module mom_diapyc_energy_req