ks_solve_column Subroutine

   pure subroutine ks_solve_column(nz, dt, f2_val, rho0, &
                                   ri_crit, shearmix_rate, fri_curvature, &
                                   c_n, c_s, lambda, kappa_0, kappa_seed_in, &
                                   kappa_trunc, tke_bg, tol_err, max_inner_it, &
                                   max_substep_it, src_max_chg, vel_underflow, &
                                   eos, &
                                   h_sd, u_sd, v_sd, t_sd, s_sd, &
                                   idz_s, idz_int_s, hint_s, il2_s, &
                                   kappa_avg_sd, tke_avg_sd)
      !! Full JHL08 column solve in surface-down indices (design doc
      !! section 5.1-5.2): background kappa_0 pre-step (no-slip bed for
      !! u,v; insulating T,S), frozen interface buoyancy derivatives,
      !! e1 tail recursion, then the adaptive predictor-corrector outer
      !! loop driving the Picard inner solve.  Returns the time-mean
      !! diffusivity `kappa_avg_sd` and TKE `tke_avg_sd` over dt.
      !$acc routine seq
      type(eos_t), intent(in) :: eos
         !! Shared EOS handle (by value) for the buoyancy derivatives.
      integer, intent(in) :: nz, max_inner_it, max_substep_it
      real(wp), intent(in) :: dt, f2_val, rho0
      real(wp), intent(in) :: ri_crit, shearmix_rate, fri_curvature
      real(wp), intent(in) :: c_n, c_s, lambda, kappa_0, kappa_seed_in
      real(wp), intent(in) :: kappa_trunc, tke_bg, tol_err, src_max_chg
      real(wp), intent(in) :: vel_underflow
      real(wp), intent(in) :: h_sd(NZL), u_sd(NZL), v_sd(NZL)
      real(wp), intent(in) :: t_sd(NZL), s_sd(NZL)
      real(wp), intent(in) :: idz_s(NZL), idz_int_s(NZLI)
      real(wp), intent(in) :: hint_s(NZLI), il2_s(NZLI)
      real(wp), intent(out) :: kappa_avg_sd(NZLI), tke_avg_sd(NZLI)

      ! Per-thread work arrays (fixed size, surface-down).
      real(wp) :: u_c(NZL), v_c(NZL), t_c(NZL), s_c(NZL)
      real(wp) :: dbuoy_t(NZLI), dbuoy_s(NZLI)
      real(wp) :: e1(NZLI)
      real(wp) :: kappa(NZLI), k_q(NZLI), kappa_avg(NZLI), tke_avg(NZLI)
      real(wp) :: n2(NZLI), s2(NZLI), tke(NZLI)
      real(wp) :: ksrc(NZLI), tkedec(NZLI)
      real(wp) :: aq(NZL), dq(NZLI), cq(NZLI)
      real(wp) :: dk(NZLI), ck(NZLI), ild2(NZLI)
      real(wp) :: cqsav(NZLI)
      real(wp) :: u_ps(NZL), v_ps(NZL), t_ps(NZL), s_ps(NZL), c1_ps(NZLI)
      real(wp) :: n2p(NZLI), s2p(NZLI), n2c(NZLI), s2c(NZLI)
      real(wp) :: kappa_out(NZLI), kq_tmp(NZLI)
      real(wp) :: tke_pred(NZLI), kappa_pred(NZLI), kappa_mid(NZLI)
      real(wp) :: tke_fin(NZLI), kappa_pred2(NZLI)
      real(wp) :: local_src_avg(NZLI), kappa_src(NZLI), local_src(NZLI)

      integer :: k, kk, io, ii
      real(wp) :: k0dt, tke_min, c_n2, c_s2, ilambda2
      real(wp) :: ome_l, eden1_l, eden2_l, i_eden_l
      real(wp) :: a1n, a1c, b1_l, d1_l, bd1_l, base_l, b1ns, b1in
      integer :: ks_src, ke_src, ks_kap, ke_kap
      real(wp) :: dsv_dt_k, dsv_ds_k, t_int, s_int, p_int, dpres
      real(wp) :: n2v, dt_rem, dt_now, dt_wt
      integer :: ks_ps, ke_ps, ks_mid, ke_mid
      logical :: no_mixing

      k0dt = dt*kappa_0
      tke_min = max(tke_bg, 1.0e-20_wp)
      c_n2 = c_n*c_n
      c_s2 = c_s*c_s
      ilambda2 = 1.0_wp/(lambda*lambda)

      ! ---- Background-diffusion pre-step (kappa_0) ----
      if (nz == 1) then
         b1_l = 1.0_wp/(h_sd(1) + k0dt*idz_int_s(2))
         u_c(1) = b1_l*h_sd(1)*u_sd(1)
         v_c(1) = b1_l*h_sd(1)*v_sd(1)
         t_c(1) = t_sd(1)
         s_c(1) = s_sd(1)
      else
         a1n = k0dt*idz_int_s(2)
         b1_l = 1.0_wp/(h_sd(1) + a1n)
         cqsav(2) = a1n*b1_l
         d1_l = h_sd(1)*b1_l
         u_c(1) = b1_l*h_sd(1)*u_sd(1)
         v_c(1) = b1_l*h_sd(1)*v_sd(1)
         t_c(1) = b1_l*h_sd(1)*t_sd(1)
         s_c(1) = b1_l*h_sd(1)*s_sd(1)
         do k = 2, nz - 1
            a1c = a1n
            a1n = k0dt*idz_int_s(k + 1)
            bd1_l = h_sd(k) + d1_l*a1c
            b1_l = 1.0_wp/(bd1_l + a1n)
            u_c(k) = b1_l*(h_sd(k)*u_sd(k) + a1c*u_c(k - 1))
            v_c(k) = b1_l*(h_sd(k)*v_sd(k) + a1c*v_c(k - 1))
            t_c(k) = b1_l*(h_sd(k)*t_sd(k) + a1c*t_c(k - 1))
            s_c(k) = b1_l*(h_sd(k)*s_sd(k) + a1c*s_c(k - 1))
            cqsav(k + 1) = a1n*b1_l
            d1_l = bd1_l*b1_l
         end do
         a1c = a1n
         base_l = h_sd(nz) + d1_l*a1c
         b1ns = 1.0_wp/(base_l + k0dt*idz_int_s(nz + 1))
         b1in = 1.0_wp/base_l
         u_c(nz) = b1ns*(h_sd(nz)*u_sd(nz) + a1c*u_c(nz - 1))
         v_c(nz) = b1ns*(h_sd(nz)*v_sd(nz) + a1c*v_c(nz - 1))
         t_c(nz) = b1in*(h_sd(nz)*t_sd(nz) + a1c*t_c(nz - 1))
         s_c(nz) = b1in*(h_sd(nz)*s_sd(nz) + a1c*s_c(nz - 1))
         cqsav(nz + 1) = 0.0_wp
         do k = nz - 1, 1, -1
            u_c(k) = u_c(k) + cqsav(k + 1)*u_c(k + 1)
            v_c(k) = v_c(k) + cqsav(k + 1)*v_c(k + 1)
            t_c(k) = t_c(k) + cqsav(k + 1)*t_c(k + 1)
            s_c(k) = s_c(k) + cqsav(k + 1)*s_c(k + 1)
         end do
      end if

      ! ---- Frozen interface buoyancy derivatives (design 5.1) ----
      ! Pressure accumulates downward (surface-relative, p(1)=0 -> D10);
      ! dbuoy_dX = -(g/rho0) drho_dX = g*rho0*dSV_dX.
      dbuoy_t(1) = 0.0_wp
      dbuoy_s(1) = 0.0_wp
      dbuoy_t(nz + 1) = 0.0_wp
      dbuoy_s(nz + 1) = 0.0_wp
      p_int = 0.0_wp
      do kk = 2, nz
         dpres = GRAVITY*rho0*h_sd(kk - 1)
         p_int = p_int + dpres
         t_int = 0.5_wp*(t_c(kk - 1) + t_c(kk))
         s_int = 0.5_wp*(s_c(kk - 1) + s_c(kk))
         call eos_specvol_derivs(eos, t_int, s_int, p_int, &
                                 dsv_dt_k, dsv_ds_k)
         dbuoy_t(kk) = GRAVITY*rho0*dsv_dt_k
         dbuoy_s(kk) = GRAVITY*rho0*dsv_ds_k
      end do

      ! ---- Initial N^2, S^2 ----
      n2(1) = 0.0_wp
      n2(nz + 1) = 0.0_wp
      s2(1) = 0.0_wp
      s2(nz + 1) = 0.0_wp
      do kk = 2, nz
         n2v = idz_int_s(kk)*(dbuoy_t(kk)*(t_c(kk - 1) - t_c(kk)) + &
                              dbuoy_s(kk)*(s_c(kk - 1) - s_c(kk)))
         if (n2v < 0.0_wp) n2v = 0.0_wp
         n2(kk) = n2v
         s2(kk) = ((u_c(kk - 1) - u_c(kk))**2 + (v_c(kk - 1) - v_c(kk))**2)* &
                  idz_int_s(kk)**2
      end do

      ! ---- e1 tail recursion ----
      e1(nz + 1) = 0.0_wp
      ome_l = 1.0_wp
      eden2_l = kappa_0*idz_s(nz)
      do kk = nz, 2, -1
         eden1_l = hint_s(kk)*sqrt(c_n2*n2(kk) + c_s2*s2(kk)) + ome_l*eden2_l
         eden2_l = kappa_0*idz_s(kk - 1)
         i_eden_l = 1.0_wp/(eden2_l + eden1_l)
         e1(kk) = eden2_l*i_eden_l
         ome_l = eden1_l*i_eden_l
      end do
      e1(1) = 0.0_wp

      ! ---- Outer-loop init ----
      do kk = 1, nz + 1
         k_q(kk) = 0.0_wp
         kappa_avg(kk) = 0.0_wp
         tke_avg(kk) = 0.0_wp
         kappa(kk) = kappa_seed_in
      end do
      kappa(1) = 0.0_wp
      kappa(nz + 1) = 0.0_wp
      dt_rem = dt
      local_src_avg(1) = 0.0_wp
      local_src_avg(nz + 1) = 0.0_wp
      do kk = 2, nz
         if (hint_s(kk) > 0.0_wp) then
            local_src_avg(kk) = 0.1_wp*k0dt*idz_int_s(kk)/hint_s(kk)
         else
            local_src_avg(kk) = 0.0_wp
         end if
      end do

      ! ---- Adaptive outer substepping ----
      do io = 1, max_substep_it
         ! Step 1: K_src + seed TKE from previous kappa/K_Q.
         ks_src = nz + 2
         ke_src = 0
         ksrc(1) = 0.0_wp
         ksrc(nz + 1) = 0.0_wp
         do kk = 2, nz
            ksrc(kk) = ks_src_func(ri_crit, shearmix_rate, fri_curvature, &
                                   n2(kk), s2(kk))
            if (ksrc(kk) > 0.0_wp) then
               if (ks_src > kk) ks_src = kk
               ke_src = kk
            end if
         end do
         do kk = 1, nz + 1
            kappa_src(kk) = ksrc(kk)
         end do

         do kk = 2, nz
            tkedec(kk) = sqrt(c_n2*n2(kk) + c_s2*s2(kk))
         end do
         tke(1) = tke_bg
         do kk = 2, nz
            if (kappa(kk) > 0.0_wp .and. k_q(kk) > 0.0_wp) then
               tke(kk) = kappa(kk)/k_q(kk)
            else
               tke(kk) = tke_min
            end if
         end do
         tke(nz + 1) = tke_min

         kappa_out = kappa

         if (ks_src > ke_src) then
            do kk = 1, nz + 1
               kappa_out(kk) = 0.0_wp
            end do
            do kk = 2, nz
               ild2(kk) = 0.0_wp
            end do
         else
            kq_tmp = k_q
            call ks_find_kappa_tke(nz, tke_min, f2_val, ri_crit, &
                                   shearmix_rate, fri_curvature, c_n2, c_s2, &
                                   ilambda2, kappa_0, kappa_trunc, tke_bg, &
                                   tol_err, max_inner_it, n2, s2, kappa, &
                                   kq_tmp, idz_s, hint_s, il2_s, e1, tke, &
                                   kappa_out, ksrc, tkedec, aq, dq, cq, dk, &
                                   ck, ild2)
            do kk = 1, nz + 1
               k_q(kk) = kq_tmp(kk)
            end do
         end if

         ! local_src for the adaptive-dt bands (design doc 5.4d): the
         ! K_src term, the kappa_0 background-change term, and the
         ! kappa diffusive-spreading term (added only when it is
         ! positive — a net inflow into the interface).
         local_src(1) = 0.0_wp
         local_src(nz + 1) = 0.0_wp
         do kk = 2, nz
            if (hint_s(kk) > 0.0_wp) then
               local_src(kk) = ksrc(kk) + kappa_0* &
                               ((idz_s(kk - 1) + idz_s(kk))/ &
                                max(hint_s(kk), 1.0e-30_wp) + ild2(kk))
               n2v = idz_s(kk - 1)*(kappa_out(kk - 1) - kappa_out(kk)) + &
                     idz_s(kk)*(kappa_out(kk + 1) - kappa_out(kk))
               if (n2v > 0.0_wp) then
                  local_src(kk) = local_src(kk) + n2v/max(hint_s(kk), 1.0e-30_wp)
               end if
            else
               local_src(kk) = ksrc(kk)
            end if
         end do

         ! Step 2: active range of kappa_out.
         ks_kap = nz + 2
         ke_kap = 0
         do kk = 2, nz
            if (kappa_out(kk) > 0.0_wp) then
               if (ks_kap > kk) ks_kap = kk
               ke_kap = kk
            end if
         end do
         if (ke_kap == nz) kappa_out(nz + 1) = 0.0_wp
         no_mixing = (ke_kap < ks_kap)

         ! Step 3: choose dt_now.
         if (no_mixing .or. io == max_substep_it) then
            dt_now = dt_rem
         else
            dt_now = ks_adaptive_dt(nz, dt_rem, io, max_substep_it, ri_crit, &
                                    shearmix_rate, fri_curvature, src_max_chg, &
                                    tol_err, vel_underflow, dbuoy_t, dbuoy_s, &
                                    h_sd, u_c, v_c, t_c, s_c, kappa_out, &
                                    kappa_src, local_src, local_src_avg, &
                                    ks_kap, ke_kap, idz_int_s)
         end if
         do kk = 2, nz
            local_src_avg(kk) = local_src_avg(kk) + dt_now*local_src(kk)
         end do
         dt_wt = dt_now/dt

         if (no_mixing) then
            ! No source reappears; remaining kappa stays 0.
            do kk = 1, nz + 1
               tke_avg(kk) = tke_avg(kk) + dt_wt*tke(kk)
            end do
            dt_rem = 0.0_wp
         else
            ! Predictor.
            ks_ps = max(ks_kap - 1, 1)
            ke_ps = min(ke_kap, nz)
            call ks_projected_state(nz, dt_now, ks_ps, ke_ps, vel_underflow, &
                                    dbuoy_t, dbuoy_s, h_sd, idz_int_s, u_c, &
                                    v_c, t_c, s_c, kappa_out, u_ps, v_ps, &
                                    t_ps, s_ps, c1_ps, n2p, s2p)
            kq_tmp = k_q
            call ks_find_kappa_tke(nz, tke_min, f2_val, ri_crit, &
                                   shearmix_rate, fri_curvature, c_n2, c_s2, &
                                   ilambda2, kappa_0, kappa_trunc, tke_bg, &
                                   tol_err, max_inner_it, n2p, s2p, kappa_out, &
                                   kq_tmp, idz_s, hint_s, il2_s, e1, tke_pred, &
                                   kappa_pred, ksrc, tkedec, aq, dq, cq, dk, &
                                   ck, ild2)
            do kk = 1, nz + 1
               kappa_mid(kk) = 0.5_wp*(kappa_out(kk) + kappa_pred(kk))
            end do
            ks_mid = nz + 2
            ke_mid = 0
            do kk = 1, nz + 1
               if (kappa_mid(kk) > 0.0_wp) then
                  if (ks_mid > kk) ks_mid = kk
                  ke_mid = kk
               end if
            end do
            ks_ps = max(ks_mid - 1, 1)
            ke_ps = min(ke_mid, nz)

            ! Corrector (real K_Q now).
            call ks_projected_state(nz, dt_now, ks_ps, ke_ps, vel_underflow, &
                                    dbuoy_t, dbuoy_s, h_sd, idz_int_s, u_c, &
                                    v_c, t_c, s_c, kappa_mid, u_ps, v_ps, &
                                    t_ps, s_ps, c1_ps, n2c, s2c)
            call ks_find_kappa_tke(nz, tke_min, f2_val, ri_crit, &
                                   shearmix_rate, fri_curvature, c_n2, c_s2, &
                                   ilambda2, kappa_0, kappa_trunc, tke_bg, &
                                   tol_err, max_inner_it, n2c, s2c, kappa_out, &
                                   k_q, idz_s, hint_s, il2_s, e1, tke_fin, &
                                   kappa_pred2, ksrc, tkedec, aq, dq, cq, dk, &
                                   ck, ild2)
            dt_rem = dt_rem - dt_now
            do kk = 1, nz + 1
               kappa_avg(kk) = kappa_avg(kk) + &
                               dt_wt*0.5_wp*(kappa_out(kk) + kappa_pred2(kk))
               tke_avg(kk) = tke_avg(kk) + &
                             dt_wt*0.5_wp*(tke_pred(kk) + tke_fin(kk))
               kappa(kk) = kappa_pred2(kk)
            end do
            kappa(1) = 0.0_wp
            kappa(nz + 1) = 0.0_wp

            ! Step 5: full-column real-state advance for the next substep.
            if (dt_rem > 0.0_wp) then
               do kk = 1, nz + 1
                  kappa_mid(kk) = 0.5_wp*(kappa_out(kk) + kappa_pred2(kk))
               end do
               call ks_projected_state(nz, dt_now, 1, nz, vel_underflow, &
                                       dbuoy_t, dbuoy_s, h_sd, idz_int_s, u_c, &
                                       v_c, t_c, s_c, kappa_mid, u_ps, v_ps, &
                                       t_ps, s_ps, c1_ps, n2, s2)
               do k = 1, nz
                  u_c(k) = u_ps(k)
                  v_c(k) = v_ps(k)
                  t_c(k) = t_ps(k)
                  s_c(k) = s_ps(k)
               end do
            end if
         end if

         if (dt_rem <= 0.0_wp) exit
      end do

      do kk = 1, nz + 1
         kappa_avg_sd(kk) = kappa_avg(kk)
         tke_avg_sd(kk) = tke_avg(kk)
      end do
      kappa_avg_sd(1) = 0.0_wp
      kappa_avg_sd(nz + 1) = 0.0_wp
   end subroutine ks_solve_column