renormalise_zonal_flux_to_uhbt Subroutine

   pure subroutine renormalise_zonal_flux_to_uhbt(grid, metrics, this, ms, uhbt, dt, skip_walls, &
                                                  has_west, has_east, visc_rem, u_cor, &
                                                  use_por, por, use_open, open_f)
      !! Apply a uniform per-face velocity correction so
      !! `Σ_k mass_flux_x_layer(i, j, k) = uhbt(i, j)` at every face.
      !! Helper for `continuity_zonal_flux`.
      !! `skip_walls` (default true) bypasses the physical-wall faces
      !! where mass_flux was zeroed; pass `.false.` for periodic axes
      !! so the wall faces (which carry real transport) are renormalised.
      type(hgrid_t), intent(in) :: grid
      type(ocean_metrics_t), intent(in) :: metrics
      type(continuity_t), intent(in) :: this
         !! Read-only here (only the PPM edge buffers are consulted).
         !! `intent(in)` is load-bearing, not tidiness: the knob-off call
         !! sites pass one of THESE buffers as the inert `por` stand-in,
         !! and `intent(in)` turns "the callee never defines it" from a
         !! comment into a compiler-enforced invariant, so the argument
         !! association can never become aliasing.
      type(multilayer_state_t), intent(inout) :: ms
      ! assumed-shape-ok: uhbt is a face-sized array (nx+1, ny); explicit-shape
      ! would require a separate nu=nx+1 argument that the caller doesn't pass.
      real(wp), intent(in) :: uhbt(:, :)
      real(wp), intent(in) :: dt
         !! Outer-step dt (s), for the CFL bracket on `du`.
      logical, intent(in), optional :: skip_walls
         !! When .true. (default), cycle the physical-wall faces.
         !! When .false. (periodic axis), include them in the renorm.
      logical, intent(in), optional :: has_west, has_east
         !! Physical-edge flags (default .true. = single-rank behaviour,
         !! bit-identical).  .false. at an MPI seam: the local "wall
         !! position" face `i = nghost+1` (west) / `i = nghost+nx_phys+1`
         !! (east) is a REAL interior face carrying transport, so it MUST
         !! be renormalised like any interior face (O4 seam fix — the
         !! un-renormalised seam face left per-layer fluxes inconsistent
         !! with `uhbt`; the Eulerian-z h-rescale hid it in h while
         !! tracers rode the raw fluxes, breaking hTr/h at the seam).
      ! assumed-shape-ok: mirrors the `uhbt` waiver above — visc_rem is a
      ! (nx+1, ny, nz) face array the caller holds on `bt_work`, and the
      ! routine is cadence-bounded (once per stage, not per substep).
      real(wp), intent(in), optional :: visc_rem(:, :, :)
         !! Per-layer viscous remnant γ_k weighting the barotropic increment
         !! (MOM6 `u_cor = u + du·visc_rem`; Jacobian
         !! `duhdu = dy·h_marg·visc_rem`).  ABSENT ⇒ γ ≡ 1 ⇒ the historical
         !! uniform-`du` form, bit-identical.
      ! assumed-shape-ok: face array, same waiver as `uhbt`.
      real(wp), intent(inout), optional :: u_cor(:, :, :)
         !! MOM6's `u_cor` return: the
         !! transport-consistent velocity `u0 + du*gamma_k`, i.e. the velocity
         !! that yields `uhbt` as the depth-integrated transport.
         !! **This is a SEPARATE time-mean field — it must NEVER be the
         !! prognostic velocity.** MOM6 writes it to `u_av`
         !! and evaluates the slow
         !! tendencies on it; writing it back into the prognostic kills the
         !! run by step 5.  See docs/MOM6_SPLIT_RK2_SPEC.md §5 trap 1.
         !! Absent => flux-only, bit-identical.

      logical, intent(in) :: use_por
         !! Porous barriers active.  `.false.` => `por` is never indexed and
         !! the arithmetic below stays byte-identical to the un-narrowed form.
      real(wp), intent(in) :: por(grid%nx_total + 1, grid%ny_total, ms%nz_ml)
         !! Layer-averaged open-area fraction at this stagger (nondim),
         !! read ONLY when `use_por`.
         !!
         !! `use_por = .false.` callers must still pass a face-sized array
         !! that is genuinely device-present, and NOT the knob-off
         !! `(1,1,1)` placeholder: nvfortran builds the `do concurrent`
         !! data clause from the LOOP BOUNDS, not from the descriptor, so a
         !! placeholder is reported "partially present" and aborts under
         !! `mem:separate` even though the branch that indexes it is never
         !! taken.  The call sites therefore hand over one of this
         !! routine's own read-only PPM edge buffers as an inert stand-in
         !! (right shape, already mapped, never DEFINED here, so no
         !! argument aliasing) -- which keeps the two full-size open-area
         !! fields off the allocation list entirely for a default run.

      logical, intent(in) :: use_open
         !! z-level closed faces active (`&vcoord_nml zfixed_closed_faces`).
         !! `.false.` => `open_f` is never indexed and the arithmetic below
         !! stays byte-identical to the un-masked form.
      real(wp), intent(in) :: open_f(grid%nx_total + 1, grid%ny_total, ms%nz_ml)
         !! Per-layer 0/1 face-open mask at this stagger, read ONLY when
         !! `use_open`.  It enters the SAME weight `wk` the porous fraction
         !! does -- `wk = dy_cu * por * open` -- which is the whole reason
         !! the barotropic transport is distributed over the OPEN layers
         !! only: a closed layer gets `wk = 0`, so it receives no `du` and
         !! contributes nothing to `sum_h`, and `sum_k mass_flux = uhbt`
         !! stays the exact fixed point the Newton solve iterates to.
         !!
         !! `use_open = .false.` callers must still pass a face-sized,
         !! genuinely device-present array and NOT the `(1,1,1)`
         !! placeholder -- see the `por` dummy's docstring for why; the
         !! call sites hand over `h_face_right_x` as the inert stand-in.

      integer :: i, j, k, nx, ny, nz, iter
      real(wp) :: u, h_face, sum_flux, sum_h, du, target, w, wk
      real(wp) :: flux0(NZ_STACK_MAX), u0(NZ_STACK_MAX)
      real(wp) :: u_lim, du_hi, du_lo, vr_k, du_k
      real(wp) :: b_lo, b_hi, du_new
      integer :: maxit
      logical :: consistent
      logical :: skip_w, has_w, has_e, use_vr, upd_u

      nx = grid%nx_total
      ny = grid%ny_total
      nz = ms%nz_ml
      skip_w = .true.
      if (present(skip_walls)) skip_w = skip_walls
      has_w = .true.
      if (present(has_west)) has_w = has_west
      has_e = .true.
      if (present(has_east)) has_e = has_east
      upd_u = present(u_cor)
      ! gamma-weighting and the u_cor write-back are ONE behaviour: together they
      ! make this renormalisation the barotropic correction (MOM6 does both in a
      ! single `continuity(..., uhbt, visc_rem, u_cor)` call).  In the historical
      ! mode `apply_bt_correction` owns the gamma-weighted Delta-u and this routine
      ! must stay the unweighted flux-only renormaliser — so both are off.
      use_vr = present(visc_rem) .and. upd_u
      consistent = this%renorm_consistent_flux
      maxit = RENORM_MAXIT
      if (consistent) maxit = RENORM_MAXIT_CONSISTENT

      ! Skip the array-edge faces (i=1, i=nx+1) and, unless skip_walls
      ! is false (periodic) or the edge is an MPI seam (has_* false),
      ! the two physical walls.
      do concurrent(j=1:ny, i=2:nx) &
         local(k, iter, u, h_face, sum_flux, sum_h, du, target, w, wk, flux0, u0, &
               u_lim, du_hi, du_lo, vr_k, du_k, b_lo, b_hi, du_new)
         ! Bypass physical walls when requested (mass_flux already 0 there
         ! for wall BCs; for periodic, real transport is present).
         if (skip_w .and. &
             ((i == grid%nghost + 1 .and. has_w) .or. &
              (i == grid%nghost + grid%nx_phys + 1 .and. has_e))) cycle
         w = metrics%dy_cu(i, j)
         target = uhbt(i, j)
         do k = 1, nz
            u0(k) = ms%u_face_x_layer(i, j, k)
            flux0(k) = ms%mass_flux_x_layer(i, j, k)
         end do
         ! Legacy single-step form (wet/dry composition — see the
         ! `renorm_legacy_single_step` docstring): one linear correction
         ! with donors picked at the uncorrected velocity, bit-identical
         ! to the pre-Newton renormalisation.
         if (this%renorm_legacy_single_step) then
            sum_flux = 0.0_wp
            sum_h = 0.0_wp
            do k = 1, nz
               wk = w
               if (use_por) wk = w*por(i, j, k)
               if (use_open) wk = wk*open_f(i, j, k)
               if (u0(k) >= 0.0_wp) then
                  h_face = this%h_face_left_x%data(i, j, k)
               else
                  h_face = this%h_face_right_x%data(i, j, k)
               end if
               sum_flux = sum_flux + flux0(k)
               sum_h = sum_h + h_face*wk
            end do
            if (sum_h > 0.0_wp) then
               du = (target - sum_flux)/sum_h
               do k = 1, nz
                  wk = w
                  if (use_por) wk = w*por(i, j, k)
                  if (use_open) wk = wk*open_f(i, j, k)
                  if (u0(k) >= 0.0_wp) then
                     h_face = this%h_face_left_x%data(i, j, k)
                  else
                     h_face = this%h_face_right_x%data(i, j, k)
                  end if
                  ms%mass_flux_x_layer(i, j, k) = flux0(k) + du*h_face*wk
               end do
            end if
            cycle
         end if
         ! CFL bracket: the corrected velocity of EVERY layer must stay within
         ! RENORM_CFL.  A residual `uhbt` mismatch is preferable to handing a
         ! near-massless layer a super-CFL velocity (MOM6 does the same).
         !
         ! `max(..., H_DIV_EPS)` is a 1/0 guard and nothing else.  `idxCu` is
         ! `1/dxCu` MULTIPLIED BY `wet_u` in `metrics_apply_land_mask`, so it
         ! is EXACTLY zero on every land face — and this loop visits land
         ! faces: it skips only the array edges and the two physical walls,
         ! never the interior coastline.  Unguarded that is `0.25/0`, which
         ! traps under `-ffpe-trap=zero` and otherwise makes `u_lim = +Inf`,
         ! poisoning `du_hi`/`du_lo` with infinities on a face where the
         ! answer is not used at all (`sum_h = 0` there, the Newton residual
         ! is identically zero, and the flux stays zero because `h_face·w` is
         ! zero).  On a WET face `dt·idxCu = dt/dxCu` is a physical rate,
         ! decades above `H_DIV_EPS = 1e-20`, so the `max` selects the true
         ! operand and the result is bit-identical.  `max` rather than `+`
         ! deliberately: an additive guard perturbs the last bit once
         ! `dt/dxCu` drops near `1e-16/1e-20`, which a long-dx, short-dt
         ! configuration can reach.
         u_lim = RENORM_CFL/max(dt*metrics%idxCu(i, j), H_DIV_EPS)
         du_hi = huge(1.0_wp)
         du_lo = -huge(1.0_wp)
         do k = 1, nz
            if (use_vr) then
               ! Layer increment is `du·γ_k`, so the CFL bound on `du` is
               ! divided by γ_k.  A layer the implicit friction has fully
               ! damped (γ→0) receives no increment and imposes no bound.
               vr_k = visc_rem(i, j, k)
               if (vr_k > RENORM_VR_MIN) then
                  du_hi = min(du_hi, (u_lim - u0(k))/vr_k)
                  du_lo = max(du_lo, (-u_lim - u0(k))/vr_k)
               end if
            else
               du_hi = min(du_hi, u_lim - u0(k))
               du_lo = max(du_lo, -u_lim - u0(k))
            end if
         end do
         du_hi = max(du_hi, 0.0_wp)
         du_lo = min(du_lo, 0.0_wp)
         ! Newton on `du` with the DONOR RE-PICKED each iteration.  With
         ! `visc_rem` present the per-layer increment is `du·γ_k` and the
         ! Jacobian carries the same weight — MOM6 `duhdu = dy·h_marg·visc_rem`.
         ! `du_k` is assigned (not multiplied
         ! by 1) on the γ-absent path so the expression tree — and therefore
         ! FP contraction — is unchanged: bit-identical by construction.
         du = 0.0_wp
         sum_h = 0.0_wp
         b_lo = -huge(1.0_wp)
         b_hi = huge(1.0_wp)
         do iter = 1, maxit
            sum_flux = 0.0_wp
            sum_h = 0.0_wp
            do k = 1, nz
               wk = w
               if (use_por) wk = w*por(i, j, k)
               if (use_open) wk = wk*open_f(i, j, k)
               du_k = du
               if (use_vr) du_k = du*visc_rem(i, j, k)
               if (u0(k) + du_k >= 0.0_wp) then
                  h_face = this%h_face_left_x%data(i, j, k)
               else
                  h_face = this%h_face_right_x%data(i, j, k)
               end if
               ! `consistent`: a layer whose donor FLIPPED under the
               ! correction carries the flux of its corrected velocity
               ! through its NEW donor, `(u0+du_k)·h_face·wk` — continuous
               ! (→ 0 from both sides) where the historical
               ! `flux0 + du_k·h_face` jumps by `u0·(h_new − h_old)·wk`.
               ! Unflipped layers keep the historical expression.
               if (consistent .and. ((u0(k) + du_k >= 0.0_wp) .neqv. (u0(k) >= 0.0_wp))) then
                  sum_flux = sum_flux + (u0(k) + du_k)*h_face*wk
               else
                  sum_flux = sum_flux + (flux0(k) + du_k*h_face*wk)
               end if
               if (use_vr) then
                  sum_h = sum_h + visc_rem(i, j, k)*h_face*wk
               else
                  sum_h = sum_h + h_face*wk
               end if
            end do
            if (sum_h <= 0.0_wp) exit
            if (abs(target - sum_flux) <= RENORM_TOL*max(1.0_wp, abs(target))) exit
            if (consistent) then
               ! Monotone, continuous F(du): keep a bracket around the root
               ! and bisect whenever the Newton step leaves it (MOM6
               ! `zonal_flux_adjust`: Newton + bisection).
               if (sum_flux < target) then
                  b_lo = max(b_lo, du)
               else
                  b_hi = min(b_hi, du)
               end if
               du_new = du + (target - sum_flux)/sum_h
               if (du_new <= b_lo .or. du_new >= b_hi) then
                  if (b_lo > -huge(1.0_wp) .and. b_hi < huge(1.0_wp)) du_new = 0.5_wp*(b_lo + b_hi)
               end if
               du = min(max(du_new, du_lo), du_hi)
            else
               du = min(max(du + (target - sum_flux)/sum_h, du_lo), du_hi)
            end if
         end do
         if (sum_h > 0.0_wp) then
            do k = 1, nz
               wk = w
               if (use_por) wk = w*por(i, j, k)
               if (use_open) wk = wk*open_f(i, j, k)
               du_k = du
               if (use_vr) du_k = du*visc_rem(i, j, k)
               if (u0(k) + du_k >= 0.0_wp) then
                  h_face = this%h_face_left_x%data(i, j, k)
               else
                  h_face = this%h_face_right_x%data(i, j, k)
               end if
               if (consistent .and. ((u0(k) + du_k >= 0.0_wp) .neqv. (u0(k) >= 0.0_wp))) then
                  ms%mass_flux_x_layer(i, j, k) = (u0(k) + du_k)*h_face*wk
               else
                  ms%mass_flux_x_layer(i, j, k) = flux0(k) + du_k*h_face*wk
               end if
               ! MOM6 `u_cor(I,j,k) = u(I,j,k) + du(I)*visc_rem(I,k)`.
               ! Writing this makes the renormalisation the SOLE barotropic
               ! correction; without it the velocity and the flux carry
               ! differently-weighted corrections that can disagree.
               if (upd_u) u_cor(i, j, k) = u0(k) + du_k
            end do
         end if
      end do
   end subroutine renormalise_zonal_flux_to_uhbt