MOM_spatial_means.F90

!> Functions and routines to take area, volume, mass-weighted, layerwise, zonal or meridional means
module mom_spatial_means

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

use mom_coms, only : efp_type, operator(+), operator(-), assignment(=)
use mom_coms, only : efp_to_real, real_to_efp, efp_sum_across_pes
use mom_coms, only : reproducing_sum, reproducing_sum_efp, efp_to_real
use mom_coms, only : query_efp_overflow_error, reset_efp_overflow_error
use mom_error_handler, only : mom_error, note, warning, fatal, is_root_pe
use mom_file_parser, only : get_param, log_version, param_file_type
use mom_grid, only : ocean_grid_type
use mom_verticalgrid, only : verticalgrid_type

implicit none ; private

#include <MOM_memory.h>

public :: global_i_mean, global_j_mean
public :: global_area_mean, global_layer_mean
public :: global_area_integral
public :: global_volume_mean, global_mass_integral
public :: adjust_area_mean_to_zero

contains

!> Return the global area mean of a variable. This uses reproducing sums.
function global_area_mean(var, G, scale)
  type(ocean_grid_type),             intent(in)  :: g    !< The ocean's grid structure
  real, dimension(SZI_(G), SZJ_(G)), intent(in)  :: var  !< The variable to average
  real,                    optional, intent(in)  :: scale !< A rescaling factor for the variable

  real, dimension(SZI_(G), SZJ_(G))              :: tmpforsumming
  real :: global_area_mean

  real :: scalefac  ! An overall scaling factor for the areas and variable.
  integer :: i, j, is, ie, js, je
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec

  scalefac = g%US%L_to_m**2 ; if (present(scale)) scalefac = g%US%L_to_m**2*scale

  tmpforsumming(:,:) = 0.
  do j=js,je ; do i=is,ie
    tmpforsumming(i,j) = var(i,j) * (scalefac * g%areaT(i,j) * g%mask2dT(i,j))
  enddo ; enddo
  global_area_mean = reproducing_sum(tmpforsumming) * g%IareaT_global

end function global_area_mean

!> Return the global area integral of a variable, by default using the masked area from the
!! grid, but an alternate could be used instead.  This uses reproducing sums.
function global_area_integral(var, G, scale, area)
  type(ocean_grid_type),            intent(in)  :: g     !< The ocean's grid structure
  real, dimension(SZI_(G),SZJ_(G)), intent(in)  :: var   !< The variable to integrate
  real,                   optional, intent(in)  :: scale !< A rescaling factor for the variable
  real, dimension(SZI_(G),SZJ_(G)), optional, intent(in) :: area !< The alternate area to use, including
                                                          !! any required masking [L2 ~> m2].
  real :: global_area_integral !< The returned area integral, usually in the units of var times [m2].

  ! Local variables
  real, dimension(SZI_(G),SZJ_(G)) :: tmpforsumming
  real :: scalefac  ! An overall scaling factor for the areas and variable.
  integer :: i, j, is, ie, js, je
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec

  scalefac = g%US%L_to_m**2 ; if (present(scale)) scalefac = g%US%L_to_m**2*scale

  tmpforsumming(:,:) = 0.
  if (present(area)) then
    do j=js,je ; do i=is,ie
      tmpforsumming(i,j) = var(i,j) * (scalefac * area(i,j))
    enddo ; enddo
  else
    do j=js,je ; do i=is,ie
      tmpforsumming(i,j) = var(i,j) * (scalefac * g%areaT(i,j) * g%mask2dT(i,j))
    enddo ; enddo
  endif
  global_area_integral = reproducing_sum(tmpforsumming)

end function global_area_integral

!> Return the layerwise global thickness-weighted mean of a variable. This uses reproducing sums.
function global_layer_mean(var, h, G, GV, scale)
  type(ocean_grid_type),                     intent(in)  :: g    !< The ocean's grid structure
  type(verticalgrid_type),                   intent(in)  :: gv   !< The ocean's vertical grid structure
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in)  :: var  !< The variable to average
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in)  :: h    !< Layer thicknesses [H ~> m or kg m-2]
  real,                            optional, intent(in)  :: scale !< A rescaling factor for the variable
  real, dimension(SZK_(GV))                   :: global_layer_mean

  real, dimension(G%isc:G%iec, G%jsc:G%jec, SZK_(GV)) :: tmpforsumming, weight
  type(efp_type), dimension(2*SZK_(GV)) :: laysums
  real, dimension(SZK_(GV)) :: global_temp_scalar, global_weight_scalar
  real :: scalefac  ! A scaling factor for the variable.
  integer :: i, j, k, is, ie, js, je, nz
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke

  scalefac = 1.0 ; if (present(scale)) scalefac = scale
  tmpforsumming(:,:,:) = 0. ; weight(:,:,:) = 0.

  do k=1,nz ; do j=js,je ; do i=is,ie
    weight(i,j,k)  =  (gv%H_to_m * h(i,j,k)) * (g%US%L_to_m**2*g%areaT(i,j) * g%mask2dT(i,j))
    tmpforsumming(i,j,k) =  scalefac * var(i,j,k) * weight(i,j,k)
  enddo ; enddo ; enddo

  global_temp_scalar = reproducing_sum(tmpforsumming, efp_lay_sums=laysums(1:nz), only_on_pe=.true.)
  global_weight_scalar = reproducing_sum(weight, efp_lay_sums=laysums(nz+1:2*nz), only_on_pe=.true.)
  call efp_sum_across_pes(laysums, 2*nz)

  do k=1,nz
    global_layer_mean(k) = efp_to_real(laysums(k)) / efp_to_real(laysums(nz+k))
  enddo

end function global_layer_mean

!> Find the global thickness-weighted mean of a variable. This uses reproducing sums.
function global_volume_mean(var, h, G, GV, scale)
  type(ocean_grid_type),   intent(in)  :: g    !< The ocean's grid structure
  type(verticalgrid_type), intent(in)  :: gv   !< The ocean's vertical grid structure
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), &
                           intent(in)  :: var  !< The variable being averaged
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), &
                           intent(in)  :: h    !< Layer thicknesses [H ~> m or kg m-2]
  real,          optional, intent(in)  :: scale !< A rescaling factor for the variable
  real :: global_volume_mean  !< The thickness-weighted average of var

  real :: scalefac  ! A scaling factor for the variable.
  real :: weight_here
  real, dimension(SZI_(G), SZJ_(G)) :: tmpforsumming, sum_weight
  integer :: i, j, k, is, ie, js, je, nz
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke

  scalefac = 1.0 ; if (present(scale)) scalefac = scale
  tmpforsumming(:,:) = 0. ; sum_weight(:,:) = 0.

  do k=1,nz ; do j=js,je ; do i=is,ie
    weight_here  =  (gv%H_to_m * h(i,j,k)) * (g%US%L_to_m**2*g%areaT(i,j) * g%mask2dT(i,j))
    tmpforsumming(i,j) = tmpforsumming(i,j) + scalefac * var(i,j,k) * weight_here
    sum_weight(i,j) = sum_weight(i,j) + weight_here
  enddo ; enddo ; enddo
  global_volume_mean = (reproducing_sum(tmpforsumming)) / &
                       (reproducing_sum(sum_weight))

end function global_volume_mean


!> Find the global mass-weighted integral of a variable. This uses reproducing sums.
function global_mass_integral(h, G, GV, var, on_PE_only, scale)
  type(ocean_grid_type),   intent(in)  :: g    !< The ocean's grid structure
  type(verticalgrid_type), intent(in)  :: gv   !< The ocean's vertical grid structure
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), &
                           intent(in)  :: h    !< Layer thicknesses [H ~> m or kg m-2]
  real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), &
                 optional, intent(in)  :: var  !< The variable being integrated
  logical,       optional, intent(in)  :: on_pe_only  !< If present and true, the sum is only
                                !! done on the local PE, and it is _not_ order invariant.
  real,          optional, intent(in)  :: scale !< A rescaling factor for the variable
  real :: global_mass_integral  !< The mass-weighted integral of var (or 1) in
                                !! kg times the units of var

  real, dimension(SZI_(G), SZJ_(G)) :: tmpforsumming
  real :: scalefac  ! An overall scaling factor for the areas and variable.
  logical :: global_sum
  integer :: i, j, k, is, ie, js, je, nz
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke

  scalefac = g%US%L_to_m**2 ; if (present(scale)) scalefac = g%US%L_to_m**2*scale
  tmpforsumming(:,:) = 0.0

  if (present(var)) then
    do k=1,nz ; do j=js,je ; do i=is,ie
      tmpforsumming(i,j) = tmpforsumming(i,j) + var(i,j,k) * &
                ((gv%H_to_kg_m2 * h(i,j,k)) * (scalefac*g%areaT(i,j) * g%mask2dT(i,j)))
    enddo ; enddo ; enddo
  else
    do k=1,nz ; do j=js,je ; do i=is,ie
      tmpforsumming(i,j) = tmpforsumming(i,j) + &
                ((gv%H_to_kg_m2 * h(i,j,k)) * (scalefac*g%areaT(i,j) * g%mask2dT(i,j)))
    enddo ; enddo ; enddo
  endif
  global_sum = .true. ; if (present(on_pe_only)) global_sum = .not.on_pe_only
  if (global_sum) then
    global_mass_integral = reproducing_sum(tmpforsumming)
  else
    global_mass_integral = 0.0
    do j=js,je ; do i=is,ie
      global_mass_integral = global_mass_integral + tmpforsumming(i,j)
    enddo ; enddo
  endif

end function global_mass_integral


!> Determine the global mean of a field along rows of constant i, returning it
!! in a 1-d array using the local indexing. This uses reproducing sums.
subroutine global_i_mean(array, i_mean, G, mask, scale, tmp_scale)
  type(ocean_grid_type),            intent(inout) :: g    !< The ocean's grid structure
  real, dimension(SZI_(G),SZJ_(G)), intent(in)    :: array  !< The variable being averaged
  real, dimension(SZJ_(G)),         intent(out)   :: i_mean !< Global mean of array along its i-axis
  real, dimension(SZI_(G),SZJ_(G)), &
                          optional, intent(in)    :: mask  !< An array used for weighting the i-mean
  real,                   optional, intent(in)    :: scale !< A rescaling factor for the output variable
  real,                   optional, intent(in)    :: tmp_scale !< A rescaling factor for the internal
                                                           !! calculations that is removed from the output

  ! Local variables
  type(efp_type), allocatable, dimension(:) :: asum, mask_sum
  real :: scalefac  ! A scaling factor for the variable.
  real :: unscale   ! A factor for undoing any internal rescaling before output.
  real :: mask_sum_r
  integer :: is, ie, js, je, idg_off, jdg_off
  integer :: i, j

  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec
  idg_off = g%idg_offset ; jdg_off = g%jdg_offset

  scalefac = 1.0 ; if (present(scale)) scalefac = scale
  unscale = 1.0
  if (present(tmp_scale)) then ; if (tmp_scale /= 0.0) then
    scalefac = scalefac * tmp_scale ; unscale = 1.0 / tmp_scale
  endif ; endif
  call reset_efp_overflow_error()

  allocate(asum(g%jsg:g%jeg))
  if (present(mask)) then
    allocate(mask_sum(g%jsg:g%jeg))

    do j=g%jsg,g%jeg
      asum(j) = real_to_efp(0.0) ; mask_sum(j) = real_to_efp(0.0)
    enddo

    do i=is,ie ; do j=js,je
      asum(j+jdg_off) = asum(j+jdg_off) + real_to_efp(scalefac*array(i,j)*mask(i,j))
      mask_sum(j+jdg_off) = mask_sum(j+jdg_off) + real_to_efp(mask(i,j))
    enddo ; enddo

    if (query_efp_overflow_error()) call mom_error(fatal, &
      "global_i_mean overflow error occurred before sums across PEs.")

    call efp_sum_across_pes(asum(g%jsg:g%jeg), g%jeg-g%jsg+1)
    call efp_sum_across_pes(mask_sum(g%jsg:g%jeg), g%jeg-g%jsg+1)

    if (query_efp_overflow_error()) call mom_error(fatal, &
      "global_i_mean overflow error occurred during sums across PEs.")

    do j=js,je
      mask_sum_r = efp_to_real(mask_sum(j+jdg_off))
      if (mask_sum_r == 0.0 ) then ; i_mean(j) = 0.0 ; else
        i_mean(j) = efp_to_real(asum(j+jdg_off)) / mask_sum_r
      endif
    enddo

    deallocate(mask_sum)
  else
    do j=g%jsg,g%jeg ; asum(j) = real_to_efp(0.0) ; enddo

    do i=is,ie ; do j=js,je
      asum(j+jdg_off) = asum(j+jdg_off) + real_to_efp(scalefac*array(i,j))
    enddo ; enddo

    if (query_efp_overflow_error()) call mom_error(fatal, &
      "global_i_mean overflow error occurred before sum across PEs.")

    call efp_sum_across_pes(asum(g%jsg:g%jeg), g%jeg-g%jsg+1)

    if (query_efp_overflow_error()) call mom_error(fatal, &
      "global_i_mean overflow error occurred during sum across PEs.")

    do j=js,je
      i_mean(j) = efp_to_real(asum(j+jdg_off)) / real(g%ieg-g%isg+1)
    enddo
  endif

  if (unscale /= 1.0) then ; do j=js,je ; i_mean(j) = unscale*i_mean(j) ; enddo ; endif

  deallocate(asum)

end subroutine global_i_mean

!> Determine the global mean of a field along rows of constant j, returning it
!! in a 1-d array using the local indexing. This uses reproducing sums.
subroutine global_j_mean(array, j_mean, G, mask, scale, tmp_scale)
  type(ocean_grid_type),            intent(inout) :: g    !< The ocean's grid structure
  real, dimension(SZI_(G),SZJ_(G)), intent(in)    :: array  !< The variable being averaged
  real, dimension(SZI_(G)),         intent(out)   :: j_mean !<  Global mean of array along its j-axis
  real, dimension(SZI_(G),SZJ_(G)), &
                          optional, intent(in)    :: mask  !< An array used for weighting the j-mean
  real,                   optional, intent(in)    :: scale !< A rescaling factor for the output variable
  real,                   optional, intent(in)    :: tmp_scale !< A rescaling factor for the internal
                                                           !! calculations that is removed from the output

  ! Local variables
  type(efp_type), allocatable, dimension(:) :: asum, mask_sum
  real :: mask_sum_r
  real :: scalefac  ! A scaling factor for the variable.
  real :: unscale   ! A factor for undoing any internal rescaling before output.
  integer :: is, ie, js, je, idg_off, jdg_off
  integer :: i, j

  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec
  idg_off = g%idg_offset ; jdg_off = g%jdg_offset

  scalefac = 1.0 ; if (present(scale)) scalefac = scale
  unscale = 1.0
  if (present(tmp_scale)) then ; if (tmp_scale /= 0.0) then
    scalefac = scalefac * tmp_scale ; unscale = 1.0 / tmp_scale
  endif ; endif
  call reset_efp_overflow_error()

  allocate(asum(g%isg:g%ieg))
  if (present(mask)) then
    allocate (mask_sum(g%isg:g%ieg))

    do i=g%isg,g%ieg
      asum(i) = real_to_efp(0.0) ; mask_sum(i) = real_to_efp(0.0)
    enddo

    do i=is,ie ; do j=js,je
      asum(i+idg_off) = asum(i+idg_off) + real_to_efp(scalefac*array(i,j)*mask(i,j))
      mask_sum(i+idg_off) = mask_sum(i+idg_off) + real_to_efp(mask(i,j))
    enddo ; enddo

    if (query_efp_overflow_error()) call mom_error(fatal, &
      "global_j_mean overflow error occurred before sums across PEs.")

    call efp_sum_across_pes(asum(g%isg:g%ieg), g%ieg-g%isg+1)
    call efp_sum_across_pes(mask_sum(g%isg:g%ieg), g%ieg-g%isg+1)

    if (query_efp_overflow_error()) call mom_error(fatal, &
      "global_j_mean overflow error occurred during sums across PEs.")

    do i=is,ie
      mask_sum_r = efp_to_real(mask_sum(i+idg_off))
      if (mask_sum_r == 0.0 ) then ; j_mean(i) = 0.0 ; else
        j_mean(i) = efp_to_real(asum(i+idg_off)) / mask_sum_r
      endif
    enddo

    deallocate(mask_sum)
  else
    do i=g%isg,g%ieg ; asum(i) = real_to_efp(0.0) ; enddo

    do i=is,ie ; do j=js,je
      asum(i+idg_off) = asum(i+idg_off) + real_to_efp(scalefac*array(i,j))
    enddo ; enddo

    if (query_efp_overflow_error()) call mom_error(fatal, &
      "global_j_mean overflow error occurred before sum across PEs.")

    call efp_sum_across_pes(asum(g%isg:g%ieg), g%ieg-g%isg+1)

    if (query_efp_overflow_error()) call mom_error(fatal, &
      "global_j_mean overflow error occurred during sum across PEs.")

    do i=is,ie
      j_mean(i) = efp_to_real(asum(i+idg_off)) / real(g%jeg-g%jsg+1)
    enddo
  endif

  if (unscale /= 1.0) then ; do i=is,ie ; j_mean(i) = unscale*j_mean(i) ; enddo ; endif

  deallocate(asum)

end subroutine global_j_mean

!> Adjust 2d array such that area mean is zero without moving the zero contour
subroutine adjust_area_mean_to_zero(array, G, scaling, unit_scale)
  type(ocean_grid_type),            intent(in)    :: g       !< Grid structure
  real, dimension(SZI_(G),SZJ_(G)), intent(inout) :: array   !< 2D array to be adjusted
  real, optional,                   intent(out)   :: scaling !< The scaling factor used
  real,                   optional, intent(in)    :: unit_scale !< A rescaling factor for the variable
  ! Local variables
  real, dimension(G%isc:G%iec, G%jsc:G%jec) :: posvals, negvals, areaxposvals, areaxnegvals
  integer :: i,j
  type(efp_type), dimension(2) :: areaint_efp
  real :: scalefac  ! A scaling factor for the variable.
  real :: i_scalefac ! The Adcroft reciprocal of scalefac
  real :: areaintposvals, areaintnegvals, posscale, negscale

  scalefac = 1.0 ; if (present(unit_scale)) scalefac = unit_scale
  i_scalefac = 0.0 ; if (scalefac /= 0.0) i_scalefac = 1.0 / scalefac

  ! areaXposVals(:,:) = 0.  ! This zeros out halo points.
  ! areaXnegVals(:,:) = 0.  ! This zeros out halo points.

  do j=g%jsc,g%jec ; do i=g%isc,g%iec
    posvals(i,j) = max(0., scalefac*array(i,j))
    areaxposvals(i,j) = g%US%L_to_m**2*g%areaT(i,j) * posvals(i,j)
    negvals(i,j) = min(0., scalefac*array(i,j))
    areaxnegvals(i,j) = g%US%L_to_m**2*g%areaT(i,j) * negvals(i,j)
  enddo ; enddo

  ! Combining the sums like this avoids separate blocking global sums.
  areaint_efp(1) = reproducing_sum_efp( areaxposvals, only_on_pe=.true. )
  areaint_efp(2) = reproducing_sum_efp( areaxnegvals, only_on_pe=.true. )
  call efp_sum_across_pes(areaint_efp, 2)
  areaintposvals = efp_to_real( areaint_efp(1) )
  areaintnegvals = efp_to_real( areaint_efp(2) )

  posscale = 0.0 ; negscale = 0.0
  if ((areaintposvals>0.).and.(areaintnegvals<0.)) then ! Only adjust if possible
    if (areaintposvals>-areaintnegvals) then ! Scale down positive values
      posscale = - areaintnegvals / areaintposvals
      do j=g%jsc,g%jec ; do i=g%isc,g%iec
        array(i,j) = ((posscale * posvals(i,j)) + negvals(i,j)) * i_scalefac
      enddo ; enddo
    elseif (areaintposvals<-areaintnegvals) then ! Scale down negative values
      negscale = - areaintposvals / areaintnegvals
      do j=g%jsc,g%jec ; do i=g%isc,g%iec
        array(i,j) = (posvals(i,j) + (negscale * negvals(i,j))) * i_scalefac
      enddo ; enddo
    endif
  endif
  if (present(scaling)) scaling = posscale - negscale

end subroutine adjust_area_mean_to_zero

end module mom_spatial_means