budget_total_ke Function

   pure function budget_total_ke(ms, mask, areaT) result(total)
      !! Total KE: Σ over masked interior of `0.5·h·(u²+v²)·areaT·weight`,
      !! u/v averaged from C-grid faces to centres.  Public only for tests.
      type(multilayer_state_t), intent(in) :: ms
      type(diag_mask_t), intent(in) :: mask
      real(wp), intent(in) :: areaT(:, :)
      real(wp) :: total
      integer :: i, j, k, nx, ny, nz
      real(wp) :: w_dA, uc, vc, ke_col
      total = 0.0_wp
      nx = min(mask%nx, size(ms%h_layer, 1), &
               size(ms%u_face_x_layer, 1) - 1, &
               size(ms%v_face_y_layer, 1))
      ny = min(mask%ny, size(ms%h_layer, 2), &
               size(ms%u_face_x_layer, 2), &
               size(ms%v_face_y_layer, 2) - 1)
      nz = ms%nz_ml
      do j = 1, ny
         do i = 1, nx
            w_dA = mask%weight(i, j)*areaT(i, j)
            if (w_dA <= 0.0_wp) cycle
            ke_col = 0.0_wp
            do k = 1, nz
               uc = 0.5_wp*(ms%u_face_x_layer(i, j, k) + ms%u_face_x_layer(i + 1, j, k))
               vc = 0.5_wp*(ms%v_face_y_layer(i, j, k) + ms%v_face_y_layer(i, j + 1, k))
               ke_col = ke_col + ms%h_layer(i, j, k)*0.5_wp*(uc*uc + vc*vc)
            end do
            total = total + w_dA*ke_col
         end do
      end do
   end function budget_total_ke