MOM_wave_structure.F90

!> Vertical structure functions for first baroclinic mode wave speed
module mom_wave_structure

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

!  By Benjamin Mater & Robert Hallberg, 2015

! The subroutine in this module calculates the vertical structure
! functions of the first baroclinic mode internal wave speed.
! Calculation of interface values is the same as done in
! MOM_wave_speed by Hallberg, 2008.

use mom_debugging,     only : isnan => is_nan
use mom_diag_mediator, only : post_data, query_averaging_enabled, diag_ctrl
use mom_diag_mediator, only : register_diag_field, safe_alloc_ptr, time_type
use mom_eos,           only : calculate_density_derivs
use mom_error_handler, only : mom_error, fatal, warning
use mom_file_parser,   only : log_version, param_file_type, get_param
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

implicit none ; private

#include <MOM_memory.h>

public wave_structure, wave_structure_init

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

!> The control structure for the MOM_wave_structure module
type, public :: wave_structure_cs ; !private
  type(diag_ctrl), pointer :: diag => null() !< A structure that is used to
                                   !! regulate the timing of diagnostic output.
  real, allocatable, dimension(:,:,:) :: w_strct
                                   !< Vertical structure of vertical velocity (normalized) [m s-1].
  real, allocatable, dimension(:,:,:) :: u_strct
                                   !< Vertical structure of horizontal velocity (normalized) [m s-1].
  real, allocatable, dimension(:,:,:) :: w_profile
                                   !< Vertical profile of w_hat(z), where
                                   !! w(x,y,z,t) = w_hat(z)*exp(i(kx+ly-freq*t)) is the full time-
                                   !! varying vertical velocity with w_hat(z) = W0*w_strct(z) [Z T-1 ~> m s-1].
  real, allocatable, dimension(:,:,:) :: uavg_profile
                                   !< Vertical profile of the magnitude of horizontal velocity,
                                   !! (u^2+v^2)^0.5, averaged over a period [L T-1 ~> m s-1].
  real, allocatable, dimension(:,:,:) :: z_depths
                                   !< Depths of layer interfaces [m].
  real, allocatable, dimension(:,:,:) :: n2
                                   !< Squared buoyancy frequency at each interface [s-2].
  integer, allocatable, dimension(:,:):: num_intfaces
                                   !< Number of layer interfaces (including surface and bottom)
  real    :: int_tide_source_x     !< X Location of generation site
                                   !! for internal tide for testing (BDM)
  real    :: int_tide_source_y     !< Y Location of generation site
                                   !! for internal tide for testing (BDM)

end type wave_structure_cs

contains

!>  This subroutine determines the internal wave velocity structure for any mode.
!!
!! This subroutine solves for the eigen vector [vertical structure, e(k)] associated with
!! the first baroclinic mode speed [i.e., smallest eigen value (lam = 1/c^2)] of the
!! system d2e/dz2 = -(N2/cn2)e, or (A-lam*I)e = 0, where A = -(1/N2)(d2/dz2), lam = 1/c^2,
!! and I is the identity matrix. 2nd order discretization in the vertical lets this system
!! be represented as
!!
!!   -Igu(k)*e(k-1) + (Igu(k)+Igl(k)-lam)*e(k) - Igl(k)*e(k+1) = 0.0
!!
!! with rigid lid boundary conditions e(1) = e(nz+1) = 0.0 giving
!!
!!   (Igu(2)+Igl(2)-lam)*e(2) - Igl(2)*e(3) = 0.0
!!   -Igu(nz)*e(nz-1) + (Igu(nz)+Igl(nz)-lam)*e(nz) = 0.0
!!
!! where, upon noting N2 = reduced gravity/layer thickness, we get
!!    Igl(k) = 1.0/(gprime(k)*H(k)) ; Igu(k) = 1.0/(gprime(k)*H(k-1))
!!
!! The eigen value for this system is approximated using "wave_speed." This subroutine uses
!! these eigen values (mode speeds) to estimate the corresponding eigen vectors (velocity
!! structure) using the "inverse iteration with shift" method. The algorithm is
!!
!!   Pick a starting vector reasonably close to mode structure and with unit magnitude, b_guess
!!   For n=1,2,3,...
!!     Solve (A-lam*I)e = e_guess for e
!!     Set e_guess=e/|e| and repeat, with each iteration refining the estimate of e
subroutine wave_structure(h, tv, G, GV, US, cn, ModeNum, freq, CS, En, full_halos)
  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_(G)), intent(in)  :: h  !< Layer thicknesses [H ~> m or kg m-2]
  type(thermo_var_ptrs),                    intent(in)  :: tv !< A structure pointing to various
                                                              !! thermodynamic variables.
  real, dimension(SZI_(G),SZJ_(G)),         intent(in)  :: cn !< The (non-rotational) mode internal
                                                              !! gravity wave speed [L T-1 ~> m s-1].
  integer,                                  intent(in)  :: modenum !< Mode number
  real,                                     intent(in)  :: freq !< Intrinsic wave frequency [T-1 ~> s-1].
  type(wave_structure_cs),                  pointer     :: cs !< The control structure returned by a
                                                              !! previous call to wave_structure_init.
  real, dimension(SZI_(G),SZJ_(G)), &
                                  optional, intent(in)  :: en !< Internal wave energy density [R Z3 T-2 ~> J m-2]
  logical,                        optional, intent(in)  :: full_halos !< If true, do the calculation
                                                              !! over the entire computational domain.
  ! Local variables
  real, dimension(SZK_(G)+1) :: &
    drho_dt, &    ! Partial derivative of density with temperature [R degC-1 ~> kg m-3 degC-1]
    drho_ds, &    ! Partial derivative of density with salinity [R ppt-1 ~> kg m-3 ppt-1]
    pres, &       ! Interface pressure [R L2 T-2 ~> Pa]
    t_int, &      ! Temperature interpolated to interfaces [degC]
    s_int, &      ! Salinity interpolated to interfaces [ppt]
    gprime        ! The reduced gravity across each interface [m2 Z-1 s-2 ~> m s-2].
  real, dimension(SZK_(G)) :: &
    igl, igu      ! The inverse of the reduced gravity across an interface times
                  ! the thickness of the layer below (Igl) or above (Igu) it [s2 m-2].
  real, dimension(SZK_(G),SZI_(G)) :: &
    hf, &         ! Layer thicknesses after very thin layers are combined [Z ~> m]
    tf, &         ! Layer temperatures after very thin layers are combined [degC]
    sf, &         ! Layer salinities after very thin layers are combined [ppt]
    rf            ! Layer densities after very thin layers are combined [R ~> kg m-3]
  real, dimension(SZK_(G)) :: &
    hc, &         ! A column of layer thicknesses after convective istabilities are removed [Z ~> m]
    tc, &         ! A column of layer temperatures after convective istabilities are removed [degC]
    sc, &         ! A column of layer salinites after convective istabilities are removed [ppt]
    rc, &         ! A column of layer densities after convective istabilities are removed [R ~> kg m-3]
    det, ddet
  real, dimension(SZI_(G),SZJ_(G)) :: &
    htot          ! The vertical sum of the thicknesses [Z ~> m]
  real :: lam
  real :: min_h_frac
  real :: z_to_pres ! A conversion factor from thicknesses to pressure [R L2 T-2 Z-1 ~> Pa m-1]
  real, dimension(SZI_(G)) :: &
    hmin, &        ! Thicknesses [Z ~> m]
    h_here, &      ! A thickness [Z ~> m]
    hxt_here, &    ! A layer integrated temperature [degC Z ~> degC m]
    hxs_here, &    ! A layer integrated salinity [ppt Z ~> ppt m]
    hxr_here       ! A layer integrated density [R Z ~> kg m-2]
  real :: speed2_tot
  real :: i_hnew   ! The inverse of a new layer thickness [Z-1 ~> m-1]
  real :: drxh_sum ! The sum of density diffrences across interfaces times thicknesses [R Z ~> kg m-2]
  real, parameter :: tol1  = 0.0001, tol2 = 0.001
  real, pointer, dimension(:,:,:) :: t => null(), s => null()
  real :: g_rho0  ! G_Earth/Rho0 in [m2 s-2 Z-1 R-1 ~> m4 s-2 kg-1].
  ! real :: rescale, I_rescale
  integer :: kf(szi_(g))
  integer, parameter :: max_itt = 1 ! number of times to iterate in solving for eigenvector
  real :: cg_subro        ! A tiny wave speed to prevent division by zero [L T-1 ~> m s-1]
  real, parameter    :: a_int = 0.5 ! value of normalized integral: \int(w_strct^2)dz = a_int
  real               :: i_a_int     ! inverse of a_int
  real               :: f2          ! squared Coriolis frequency [T-2 ~> s-2]
  real               :: kmag2       ! magnitude of horizontal wave number squared
  logical            :: use_eos     ! If true, density is calculated from T & S using an
                                    ! equation of state.
  real, dimension(SZK_(G)+1) :: w_strct, u_strct, w_profile, uavg_profile, z_int, n2
                                        ! local representations of variables in CS; note,
                                        ! not all rows will be filled if layers get merged!
  real, dimension(SZK_(G)+1) :: w_strct2, u_strct2
                                        ! squared values
  real, dimension(SZK_(G))   :: dz      ! thicknesses of merged layers (same as Hc I hope)
  ! real, dimension(SZK_(G)+1) :: dWdz_profile ! profile of dW/dz
  real                       :: w2avg   ! average of squared vertical velocity structure funtion
  real                       :: int_dwdz2
  real                       :: int_w2
  real                       :: int_n2w2
  real                       :: ke_term ! terms in vertically averaged energy equation
  real                       :: pe_term ! terms in vertically averaged energy equation
  real                       :: w0      ! A vertical velocity magnitude [Z T-1 ~> m s-1]
  real                       :: gp_unscaled ! A version of gprime rescaled to [m s-2].
  real, dimension(SZK_(G)-1) :: lam_z   ! product of eigen value and gprime(k); one value for each
                                        ! interface (excluding surface and bottom)
  real, dimension(SZK_(G)-1) :: a_diag, b_diag, c_diag
                                        ! diagonals of tridiagonal matrix; one value for each
                                        ! interface (excluding surface and bottom)
  real, dimension(SZK_(G)-1) :: e_guess ! guess at eigen vector with unit amplitde (for TDMA)
  real, dimension(SZK_(G)-1) :: e_itt   ! improved guess at eigen vector (from TDMA)
  real    :: pi
  integer :: kc
  integer :: i, j, k, k2, itt, is, ie, js, je, nz, nzm, row, ig, jg, ig_stop, jg_stop

  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = g%ke
  i_a_int = 1/a_int

  !if (present(CS)) then
    if (.not. associated(cs)) call mom_error(fatal, "MOM_wave_structure: "// &
           "Module must be initialized before it is used.")
  !endif

  if (present(full_halos)) then ; if (full_halos) then
    is = g%isd ; ie = g%ied ; js = g%jsd ; je = g%jed
  endif ; endif

  pi = (4.0*atan(1.0))

  s => tv%S ; t => tv%T
  g_rho0 = us%L_T_to_m_s**2 * gv%g_Earth / gv%Rho0
  cg_subro = 1e-100*us%m_s_to_L_T  ! The hard-coded value here might need to increase.
  use_eos = associated(tv%eqn_of_state)

  ! Simplifying the following could change answers at roundoff.
  z_to_pres = gv%Z_to_H * (gv%H_to_RZ * gv%g_Earth)
  ! rescale = 1024.0**4 ; I_rescale = 1.0/rescale

  min_h_frac = tol1 / real(nz)

  do j=js,je
    !   First merge very thin layers with the one above (or below if they are
    ! at the top).  This also transposes the row order so that columns can
    ! be worked upon one at a time.
    do i=is,ie ; htot(i,j) = 0.0 ; enddo
    do k=1,nz ; do i=is,ie ; htot(i,j) = htot(i,j) + h(i,j,k)*gv%H_to_Z ; enddo ; enddo

    do i=is,ie
      hmin(i) = htot(i,j)*min_h_frac ; kf(i) = 1 ; h_here(i) = 0.0
      hxt_here(i) = 0.0 ; hxs_here(i) = 0.0 ; hxr_here(i) = 0.0
    enddo
    if (use_eos) then
      do k=1,nz ; do i=is,ie
        if ((h_here(i) > hmin(i)) .and. (h(i,j,k)*gv%H_to_Z > hmin(i))) then
          hf(kf(i),i) = h_here(i)
          tf(kf(i),i) = hxt_here(i) / h_here(i)
          sf(kf(i),i) = hxs_here(i) / h_here(i)
          kf(i) = kf(i) + 1

          ! Start a new layer
          h_here(i) = h(i,j,k)*gv%H_to_Z
          hxt_here(i) = (h(i,j,k)*gv%H_to_Z)*t(i,j,k)
          hxs_here(i) = (h(i,j,k)*gv%H_to_Z)*s(i,j,k)
        else
          h_here(i) = h_here(i) + h(i,j,k)*gv%H_to_Z
          hxt_here(i) = hxt_here(i) + (h(i,j,k)*gv%H_to_Z)*t(i,j,k)
          hxs_here(i) = hxs_here(i) + (h(i,j,k)*gv%H_to_Z)*s(i,j,k)
        endif
      enddo ; enddo
      do i=is,ie ; if (h_here(i) > 0.0) then
        hf(kf(i),i) = h_here(i)
        tf(kf(i),i) = hxt_here(i) / h_here(i)
        sf(kf(i),i) = hxs_here(i) / h_here(i)
      endif ; enddo
    else
      do k=1,nz ; do i=is,ie
        if ((h_here(i) > hmin(i)) .and. (h(i,j,k)*gv%H_to_Z > hmin(i))) then
          hf(kf(i),i) = h_here(i) ; rf(kf(i),i) = hxr_here(i) / h_here(i)
          kf(i) = kf(i) + 1

          ! Start a new layer
          h_here(i) = h(i,j,k)*gv%H_to_Z
          hxr_here(i) = (h(i,j,k)*gv%H_to_Z)*gv%Rlay(k)
        else
          h_here(i) = h_here(i) + h(i,j,k)*gv%H_to_Z
          hxr_here(i) = hxr_here(i) + (h(i,j,k)*gv%H_to_Z)*gv%Rlay(k)
        endif
      enddo ; enddo
      do i=is,ie ; if (h_here(i) > 0.0) then
        hf(kf(i),i) = h_here(i) ; rf(kf(i),i) = hxr_here(i) / h_here(i)
      endif ; enddo
    endif ! use_EOS?

    ! From this point, we can work on individual columns without causing memory
    ! to have page faults.
    do i=is,ie ; if (cn(i,j)>0.0)then
      !----for debugging, remove later----
      ig = i + g%idg_offset ; jg = j + g%jdg_offset
      !if (ig == CS%int_tide_source_x .and. jg == CS%int_tide_source_y) then
      !-----------------------------------
      if (g%mask2dT(i,j) > 0.5) then

        lam = 1/(us%L_T_to_m_s**2 * cn(i,j)**2)

        ! Calculate drxh_sum
        if (use_eos) then
          pres(1) = 0.0
          do k=2,kf(i)
            pres(k) = pres(k-1) + z_to_pres*hf(k-1,i)
            t_int(k) = 0.5*(tf(k,i)+tf(k-1,i))
            s_int(k) = 0.5*(sf(k,i)+sf(k-1,i))
          enddo
          call calculate_density_derivs(t_int, s_int, pres, drho_dt, drho_ds, &
                                        tv%eqn_of_state, (/2,kf(i)/) )

          ! Sum the reduced gravities to find out how small a density difference
          ! is negligibly small.
          drxh_sum = 0.0
          do k=2,kf(i)
            drxh_sum = drxh_sum + 0.5*(hf(k-1,i)+hf(k,i)) * &
                max(0.0,drho_dt(k)*(tf(k,i)-tf(k-1,i)) + &
                        drho_ds(k)*(sf(k,i)-sf(k-1,i)))
          enddo
        else
          drxh_sum = 0.0
          do k=2,kf(i)
            drxh_sum = drxh_sum + 0.5*(hf(k-1,i)+hf(k,i)) * &
                              max(0.0,rf(k,i)-rf(k-1,i))
          enddo
        endif ! use_EOS?

        !   Find gprime across each internal interface, taking care of convective
        ! instabilities by merging layers.
        if (drxh_sum >= 0.0) then
          ! Merge layers to eliminate convective instabilities or exceedingly
          ! small reduced gravities.
          if (use_eos) then
            kc = 1
            hc(1) = hf(1,i) ; tc(1) = tf(1,i) ; sc(1) = sf(1,i)
            do k=2,kf(i)
              if ((drho_dt(k)*(tf(k,i)-tc(kc)) + drho_ds(k)*(sf(k,i)-sc(kc))) * &
                  (hc(kc) + hf(k,i)) < 2.0 * tol2*drxh_sum) then
                ! Merge this layer with the one above and backtrack.
                i_hnew = 1.0 / (hc(kc) + hf(k,i))
                tc(kc) = (hc(kc)*tc(kc) + hf(k,i)*tf(k,i)) * i_hnew
                sc(kc) = (hc(kc)*sc(kc) + hf(k,i)*sf(k,i)) * i_hnew
                hc(kc) = (hc(kc) + hf(k,i))
                ! Backtrack to remove any convective instabilities above...  Note
                ! that the tolerance is a factor of two larger, to avoid limit how
                ! far back we go.
                do k2=kc,2,-1
                  if ((drho_dt(k2)*(tc(k2)-tc(k2-1)) + drho_ds(k2)*(sc(k2)-sc(k2-1))) * &
                      (hc(k2) + hc(k2-1)) < tol2*drxh_sum) then
                    ! Merge the two bottommost layers.  At this point kc = k2.
                    i_hnew = 1.0 / (hc(kc) + hc(kc-1))
                    tc(kc-1) = (hc(kc)*tc(kc) + hc(kc-1)*tc(kc-1)) * i_hnew
                    sc(kc-1) = (hc(kc)*sc(kc) + hc(kc-1)*sc(kc-1)) * i_hnew
                    hc(kc-1) = (hc(kc) + hc(kc-1))
                    kc = kc - 1
                  else ; exit ; endif
                enddo
              else
                ! Add a new layer to the column.
                kc = kc + 1
                drho_ds(kc) = drho_ds(k) ; drho_dt(kc) = drho_dt(k)
                tc(kc) = tf(k,i) ; sc(kc) = sf(k,i) ; hc(kc) = hf(k,i)
              endif
            enddo
            ! At this point there are kc layers and the gprimes should be positive.
            do k=2,kc ! Revisit this if non-Boussinesq.
              gprime(k) = g_rho0 * (drho_dt(k)*(tc(k)-tc(k-1)) + &
                                    drho_ds(k)*(sc(k)-sc(k-1)))
            enddo
          else  ! .not.use_EOS
            ! Do the same with density directly...
            kc = 1
            hc(1) = hf(1,i) ; rc(1) = rf(1,i)
            do k=2,kf(i)
              if ((rf(k,i) - rc(kc)) * (hc(kc) + hf(k,i)) < 2.0*tol2*drxh_sum) then
                ! Merge this layer with the one above and backtrack.
                rc(kc) = (hc(kc)*rc(kc) + hf(k,i)*rf(k,i)) / (hc(kc) + hf(k,i))
                hc(kc) = (hc(kc) + hf(k,i))
                ! Backtrack to remove any convective instabilities above...  Note
                ! that the tolerance is a factor of two larger, to avoid limit how
                ! far back we go.
                do k2=kc,2,-1
                  if ((rc(k2)-rc(k2-1)) * (hc(k2)+hc(k2-1)) < tol2*drxh_sum) then
                    ! Merge the two bottommost layers.  At this point kc = k2.
                    rc(kc-1) = (hc(kc)*rc(kc) + hc(kc-1)*rc(kc-1)) / (hc(kc) + hc(kc-1))
                    hc(kc-1) = (hc(kc) + hc(kc-1))
                    kc = kc - 1
                  else ; exit ; endif
                enddo
              else
                ! Add a new layer to the column.
                kc = kc + 1
                rc(kc) = rf(k,i) ; hc(kc) = hf(k,i)
              endif
            enddo
            ! At this point there are kc layers and the gprimes should be positive.
            do k=2,kc ! Revisit this if non-Boussinesq.
              gprime(k) = g_rho0 * (rc(k)-rc(k-1))
            enddo
          endif  ! use_EOS?

          !-----------------NOW FIND WAVE STRUCTURE-------------------------------------
          ! Construct and solve tridiagonal system for the interior interfaces
          ! Note that kc   = number of layers,
          !           kc+1 = nzm = number of interfaces,
          !           kc-1 = number of interior interfaces (excluding surface and bottom)
          ! Also, note that "K" refers to an interface, while "k" refers to the layer below.
          ! Need at least 3 layers (2 internal interfaces) to generate a matrix, also
          ! need number of layers to be greater than the mode number
          if (kc >= max(3, modenum + 1)) then
            ! Set depth at surface
            z_int(1) = 0.0
            ! Calculate Igu, Igl, depth, and N2 at each interior interface
            ! [excludes surface (K=1) and bottom (K=kc+1)]
            do k=2,kc
              igl(k) = 1.0/(gprime(k)*hc(k)) ; igu(k) = 1.0/(gprime(k)*hc(k-1))
              z_int(k) = z_int(k-1) + hc(k-1)
              n2(k) = us%m_to_Z**2*gprime(k)/(0.5*(hc(k)+hc(k-1)))
            enddo
            ! Set stratification for surface and bottom (setting equal to nearest interface for now)
            n2(1) = n2(2) ; n2(kc+1) = n2(kc)
            ! Calcualte depth at bottom
            z_int(kc+1) = z_int(kc)+hc(kc)
            ! check that thicknesses sum to total depth
            if (abs(z_int(kc+1)-htot(i,j)) > 1.e-14*htot(i,j)) then
              call mom_error(fatal, "wave_structure: mismatch in total depths")
            endif

            ! Note that many of the calcluation from here on revert to using vertical
            ! distances in m, not Z.

            ! Populate interior rows of tridiagonal matrix; must multiply through by
            ! gprime to get tridiagonal matrix to the symmetrical form:
            ! [-1/H(k-1)]e(k-1) + [1/H(k-1)+1/H(k)-lam_z]e(k) + [-1/H(k)]e(k+1) = 0,
            ! where lam_z = lam*gprime is now a function of depth.
            ! Frist, populate interior rows
            do k=3,kc-1
              row = k-1 ! indexing for TD matrix rows
              gp_unscaled = us%m_to_Z*gprime(k)
              lam_z(row) = lam*gp_unscaled
              a_diag(row) = gp_unscaled*(-igu(k))
              b_diag(row) = gp_unscaled*(igu(k)+igl(k)) - lam_z(row)
              c_diag(row) = gp_unscaled*(-igl(k))
              if (isnan(lam_z(row)))then  ; print *, "Wave_structure: lam_z(row) is NAN" ; endif
              if (isnan(a_diag(row)))then ; print *, "Wave_structure: a(k) is NAN" ; endif
              if (isnan(b_diag(row)))then ; print *, "Wave_structure: b(k) is NAN" ; endif
              if (isnan(c_diag(row)))then ; print *, "Wave_structure: c(k) is NAN" ; endif
            enddo
            ! Populate top row of tridiagonal matrix
            k=2 ; row = k-1 ;
            gp_unscaled = us%m_to_Z*gprime(k)
            lam_z(row) = lam*gp_unscaled
            a_diag(row) = 0.0
            b_diag(row) = gp_unscaled*(igu(k)+igl(k)) - lam_z(row)
            c_diag(row) = gp_unscaled*(-igl(k))
            ! Populate bottom row of tridiagonal matrix
            k=kc ; row = k-1
            gp_unscaled = us%m_to_Z*gprime(k)
            lam_z(row) = lam*gp_unscaled
            a_diag(row) = gp_unscaled*(-igu(k))
            b_diag(row) = gp_unscaled*(igu(k)+igl(k)) - lam_z(row)
            c_diag(row) = 0.0

            ! Guess a vector shape to start with (excludes surface and bottom)
            e_guess(1:kc-1) = sin((z_int(2:kc)/htot(i,j)) *pi)
            e_guess(1:kc-1) = e_guess(1:kc-1)/sqrt(sum(e_guess(1:kc-1)**2))

            ! Perform inverse iteration with tri-diag solver
            do itt=1,max_itt
              call tridiag_solver(a_diag(1:kc-1),b_diag(1:kc-1),c_diag(1:kc-1), &
                                  -lam_z(1:kc-1),e_guess(1:kc-1),"TDMA_H",e_itt)
              e_guess(1:kc-1) = e_itt(1:kc-1) / sqrt(sum(e_itt(1:kc-1)**2))
            enddo ! itt-loop
            w_strct(2:kc) = e_guess(1:kc-1)
            w_strct(1)    = 0.0 ! rigid lid at surface
            w_strct(kc+1) = 0.0 ! zero-flux at bottom

            ! Check to see if solver worked
            ig_stop = 0 ; jg_stop = 0
            if (isnan(sum(w_strct(1:kc+1))))then
              print *, "Wave_structure: w_strct has a NAN at ig=", ig, ", jg=", jg
              if (i<g%isc .or. i>g%iec .or. j<g%jsc .or. j>g%jec)then
                print *, "This is occuring at a halo point."
              endif
              ig_stop = ig ; jg_stop = jg
            endif

            ! Normalize vertical structure function of w such that
            ! \int(w_strct)^2dz = a_int (a_int could be any value, e.g., 0.5)
            nzm = kc+1 ! number of layer interfaces after merging
                       !(including surface and bottom)
            w2avg = 0.0
            do k=1,nzm-1
              dz(k) = us%Z_to_m*hc(k)
              w2avg = w2avg + 0.5*(w_strct(k)**2+w_strct(k+1)**2)*dz(k)
            enddo
            !### Some mathematical cancellations could occur in the next two lines.
            w2avg = w2avg / htot(i,j)
            w_strct(:) = w_strct(:) / sqrt(htot(i,j)*w2avg*i_a_int)

            ! Calculate vertical structure function of u (i.e. dw/dz)
            do k=2,nzm-1
              u_strct(k) = 0.5*((w_strct(k-1) - w_strct(k)  )/dz(k-1) + &
                                (w_strct(k)   - w_strct(k+1))/dz(k))
            enddo
            u_strct(1)   = (w_strct(1)   -  w_strct(2) )/dz(1)
            u_strct(nzm) = (w_strct(nzm-1)-  w_strct(nzm))/dz(nzm-1)

            ! Calculate wavenumber magnitude
            f2 = g%CoriolisBu(i,j)**2
            !f2 = 0.25*US%s_to_T**2 *((G%CoriolisBu(I,J)**2 + G%CoriolisBu(I-1,J-1)**2) + &
            !    (G%CoriolisBu(I,J-1)**2 + G%CoriolisBu(I-1,J)**2))
            kmag2 = us%m_to_L**2 * (freq**2 - f2) / (cn(i,j)**2 + cg_subro**2)

            ! Calculate terms in vertically integrated energy equation
            int_dwdz2 = 0.0 ; int_w2 = 0.0 ; int_n2w2 = 0.0
            u_strct2(1:nzm) = u_strct(1:nzm)**2
            w_strct2(1:nzm) = w_strct(1:nzm)**2
            ! vertical integration with Trapezoidal rule
            do k=1,nzm-1
              int_dwdz2 = int_dwdz2 + 0.5*(u_strct2(k)+u_strct2(k+1)) * us%m_to_Z*dz(k)
              int_w2    = int_w2    + 0.5*(w_strct2(k)+w_strct2(k+1)) * us%m_to_Z*dz(k)
              int_n2w2  = int_n2w2  + 0.5*(w_strct2(k)*n2(k)+w_strct2(k+1)*n2(k+1)) * us%m_to_Z*dz(k)
            enddo

            ! Back-calculate amplitude from energy equation
            if (present(en) .and. (freq**2*kmag2 > 0.0)) then
              ! Units here are [R
              ke_term = 0.25*gv%Rho0*( ((freq**2 + f2) / (freq**2*kmag2))*int_dwdz2 + int_w2 )
              pe_term = 0.25*gv%Rho0*( int_n2w2 / (us%s_to_T*freq)**2 )
              if (en(i,j) >= 0.0) then
                w0 = sqrt( en(i,j) / (ke_term + pe_term) )
              else
                call mom_error(warning, "wave_structure: En < 0.0; setting to W0 to 0.0")
                print *, "En(i,j)=", en(i,j), " at ig=", ig, ", jg=", jg
                w0 = 0.0
              endif
              ! Calculate actual vertical velocity profile and derivative
              w_profile(:)    = w0*w_strct(:)
              ! dWdz_profile(:) = W0*u_strct(:)
              ! Calculate average magnitude of actual horizontal velocity over a period
              uavg_profile(:) = us%Z_to_L*abs(w0*u_strct(:)) * sqrt((freq**2 + f2) / (2.0*freq**2*kmag2))
            else
              w_profile(:)    = 0.0
              ! dWdz_profile(:) = 0.0
              uavg_profile(:) = 0.0
            endif

            ! Store values in control structure
            cs%w_strct(i,j,1:nzm)     = w_strct(1:nzm)
            cs%u_strct(i,j,1:nzm)     = u_strct(1:nzm)
            cs%W_profile(i,j,1:nzm)   = w_profile(1:nzm)
            cs%Uavg_profile(i,j,1:nzm)= uavg_profile(1:nzm)
            cs%z_depths(i,j,1:nzm)    = us%Z_to_m*z_int(1:nzm)
            cs%N2(i,j,1:nzm)          = n2(1:nzm)
            cs%num_intfaces(i,j)      = nzm
          else
            ! If not enough layers, default to zero
            nzm = kc+1
            cs%w_strct(i,j,1:nzm)     = 0.0
            cs%u_strct(i,j,1:nzm)     = 0.0
            cs%W_profile(i,j,1:nzm)   = 0.0
            cs%Uavg_profile(i,j,1:nzm)= 0.0
            cs%z_depths(i,j,1:nzm)    = 0.0 ! could use actual values
            cs%N2(i,j,1:nzm)          = 0.0 ! could use with actual values
            cs%num_intfaces(i,j)       = nzm
          endif  ! kc >= 3 and kc > ModeNum + 1?
        endif ! drxh_sum >= 0?
      !else     ! if at test point - delete later
      !  return ! if at test point - delete later
      !endif    ! if at test point - delete later
      endif ! mask2dT > 0.5?
    else
      ! if cn=0.0, default to zero
      nzm                       = nz+1! could use actual values
      cs%w_strct(i,j,1:nzm)     = 0.0
      cs%u_strct(i,j,1:nzm)     = 0.0
      cs%W_profile(i,j,1:nzm)   = 0.0
      cs%Uavg_profile(i,j,1:nzm)= 0.0
      cs%z_depths(i,j,1:nzm)    = 0.0 ! could use actual values
      cs%N2(i,j,1:nzm)          = 0.0 ! could use with actual values
      cs%num_intfaces(i,j)       = nzm
    endif ; enddo ! if cn>0.0? ; i-loop
  enddo ! j-loop

end subroutine wave_structure

!> Solves a tri-diagonal system Ax=y using either the standard
!! Thomas algorithm (TDMA_T) or its more stable variant that invokes the
!! "Hallberg substitution" (TDMA_H).
subroutine tridiag_solver(a, b, c, h, y, method, x)
  real, dimension(:), intent(in)  :: a !< lower diagonal with first entry equal to zero.
  real, dimension(:), intent(in)  :: b !< middle diagonal.
  real, dimension(:), intent(in)  :: c !< upper diagonal with last entry equal to zero.
  real, dimension(:), intent(in)  :: h !< vector of values that have already been added to b; used
           !! for systems of the form (e.g. average layer thickness in vertical diffusion case):
           !! [ -alpha(k-1/2) ]                       * e(k-1) +
           !! [  alpha(k-1/2) + alpha(k+1/2) + h(k) ] * e(k)   +
           !! [ -alpha(k+1/2) ]                       * e(k+1) = y(k)
           !! where a(k)=[-alpha(k-1/2)], b(k)=[alpha(k-1/2)+alpha(k+1/2) + h(k)],
           !! and c(k)=[-alpha(k+1/2)]. Only used with TDMA_H method.
  real, dimension(:), intent(in)  :: y !< vector of known values on right hand side.
  character(len=*),   intent(in)  :: method !< A string describing the algorithm to use
  real, dimension(:), intent(out) :: x !< vector of unknown values to solve for.
  ! Local variables
  integer :: nrow                         ! number of rows in A matrix
!  real, allocatable, dimension(:,:) :: A_check ! for solution checking
!  real, allocatable, dimension(:)   :: y_check ! for solution checking
  real, allocatable, dimension(:) :: c_prime, y_prime, q, alpha
                                          ! intermediate values for solvers
  real    :: Q_prime, beta                ! intermediate values for solver
  integer :: k                            ! row (e.g. interface) index
  integer :: i,j

  nrow = size(y)
  allocate(c_prime(nrow))
  allocate(y_prime(nrow))
  allocate(q(nrow))
  allocate(alpha(nrow))
!  allocate(A_check(nrow,nrow))
!  allocate(y_check(nrow))

  if (method == 'TDMA_T') then
    ! Standard Thomas algoritim (4th variant).
    ! Note: Requires A to be non-singular for accuracy/stability
    c_prime(:) = 0.0       ; y_prime(:) = 0.0
    c_prime(1) = c(1)/b(1) ; y_prime(1) = y(1)/b(1)

    ! Forward sweep
    do k=2,nrow-1
      c_prime(k) = c(k)/(b(k)-a(k)*c_prime(k-1))
    enddo
    !print *, 'c_prime=', c_prime(1:nrow)
    do k=2,nrow
      y_prime(k) = (y(k)-a(k)*y_prime(k-1))/(b(k)-a(k)*c_prime(k-1))
    enddo
    !print *, 'y_prime=', y_prime(1:nrow)
    x(nrow) = y_prime(nrow)

    ! Backward sweep
    do k=nrow-1,1,-1
      x(k) = y_prime(k)-c_prime(k)*x(k+1)
    enddo
    !print *, 'x=',x(1:nrow)

    ! Check results - delete later
    !do j=1,nrow ; do i=1,nrow
    !  if (i==j)then ;       A_check(i,j) = b(i)
    !  elseif (i==j+1)then ; A_check(i,j) = a(i)
    !  elseif (i==j-1)then ; A_check(i,j) = c(i)
    !  endif
    !enddo ; enddo
    !print *, 'A(2,1),A(2,2),A(1,2)=', A_check(2,1), A_check(2,2), A_check(1,2)
    !y_check = matmul(A_check,x)
    !if (all(y_check /= y))then
    !  print *, "tridiag_solver: Uh oh, something's not right!"
    !  print *, "y=", y
    !  print *, "y_check=", y_check
    !endif

  elseif (method == 'TDMA_H') then
    ! Thomas algoritim (4th variant) w/ Hallberg substitution.
    ! For a layered system where k is at interfaces, alpha{k+1/2} refers to
    ! some property (e.g. inverse thickness for mode-structure problem) of the
    ! layer below and alpha{k-1/2} refers to the layer above.
    ! Here, alpha(k)=alpha{k+1/2} and alpha(k-1)=alpha{k-1/2}.
    ! Strictly speaking, this formulation requires A to be a non-singular,
    ! symmetric, diagonally dominant matrix, with h>0.
    ! Need to add a check for these conditions.
    do k=1,nrow-1
      if (abs(a(k+1)-c(k)) > 1.e-10*(abs(a(k+1))+abs(c(k)))) then
        call mom_error(warning, "tridiag_solver: matrix not symmetric; need symmetry when invoking TDMA_H")
      endif
    enddo
    alpha = -c
    ! Alpha of the bottom-most layer is not necessarily zero. Therefore,
    ! back out the value from the provided b(nrow and h(nrow) values
    alpha(nrow)   = b(nrow)-h(nrow)-alpha(nrow-1)
    ! Prime other variables
    beta       = 1/b(1)
    y_prime(:) = 0.0       ; q(:) = 0.0
    y_prime(1) = beta*y(1) ; q(1) = beta*alpha(1)
    q_prime    = 1-q(1)

    ! Forward sweep
    do k=2,nrow-1
      beta = 1/(h(k)+alpha(k-1)*q_prime+alpha(k))
      if (isnan(beta))then ; print *, "Tridiag_solver: beta is NAN" ; endif
      q(k) = beta*alpha(k)
      y_prime(k) = beta*(y(k)+alpha(k-1)*y_prime(k-1))
      q_prime = beta*(h(k)+alpha(k-1)*q_prime)
    enddo
    if ((h(nrow)+alpha(nrow-1)*q_prime+alpha(nrow)) == 0.0)then
      call mom_error(fatal, "Tridiag_solver: this system is not stable.") ! ; overriding beta(nrow)
      ! This has hard-coded dimensions: beta = 1/(1e-15) ! place holder for unstable systems - delete later
    else
      beta = 1/(h(nrow)+alpha(nrow-1)*q_prime+alpha(nrow))
    endif
    y_prime(nrow) = beta*(y(nrow)+alpha(nrow-1)*y_prime(nrow-1))
    x(nrow) = y_prime(nrow)
    ! Backward sweep
    do k=nrow-1,1,-1
      x(k) = y_prime(k)+q(k)*x(k+1)
    enddo
    !print *, 'yprime=',y_prime(1:nrow)
    !print *, 'x=',x(1:nrow)
  endif

  deallocate(c_prime,y_prime,q,alpha)
! deallocate(A_check,y_check)

end subroutine tridiag_solver

!> Allocate memory associated with the wave structure module and read parameters.
subroutine wave_structure_init(Time, G, 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(param_file_type),   intent(in) :: param_file !< A structure to parse for run-time
                                              !! parameters.
  type(diag_ctrl), target, intent(in) :: diag !< A structure that is used to regulate
                                              !! diagnostic output.
  type(wave_structure_cs), pointer    :: cs   !< A pointer that is set to point to the
                                              !! control structure for this module.
! This include declares and sets the variable "version".
#include "version_variable.h"
  character(len=40)  :: mdl = "MOM_wave_structure"  ! This module's name.
  integer :: isd, ied, jsd, jed, nz

  isd = g%isd ; ied = g%ied ; jsd = g%jsd ; jed = g%jed ; nz = g%ke

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

  call get_param(param_file, mdl, "INTERNAL_TIDE_SOURCE_X", cs%int_tide_source_x, &
                 "X Location of generation site for internal tide", default=1.)
  call get_param(param_file, mdl, "INTERNAL_TIDE_SOURCE_Y", cs%int_tide_source_y, &
                 "Y Location of generation site for internal tide", default=1.)

  cs%diag => diag

  ! Allocate memory for variable in control structure; note,
  ! not all rows will be filled if layers get merged!
  allocate(cs%w_strct(isd:ied,jsd:jed,nz+1))
  allocate(cs%u_strct(isd:ied,jsd:jed,nz+1))
  allocate(cs%W_profile(isd:ied,jsd:jed,nz+1))
  allocate(cs%Uavg_profile(isd:ied,jsd:jed,nz+1))
  allocate(cs%z_depths(isd:ied,jsd:jed,nz+1))
  allocate(cs%N2(isd:ied,jsd:jed,nz+1))
  allocate(cs%num_intfaces(isd:ied,jsd:jed))

  ! Write all relevant parameters to the model log.
  call log_version(param_file, mdl, version, "")

end subroutine wave_structure_init

end module mom_wave_structure