Add section in the docs for reduced operations (#4235)

* add section for reduced operations

* Apply suggestions from code review

Co-authored-by: Gregory L. Wagner <[email protected]>

* apply review suggestions

* Apply suggestions from code review

Co-authored-by: Gregory L. Wagner <[email protected]>

* Update docs/src/operations.md

Co-authored-by: Gregory L. Wagner <[email protected]>

* Update docs/src/operations.md

Co-authored-by: Gregory L. Wagner <[email protected]>

* Apply suggestions from code review

* few tweaks

* no need for CPU() arg in grid

* some space

* few tweaks

* add a sentence about reductions

* fix latex formatting in docstrings

we might as well, since we are at it..

* fix Reduction ref

* Some typo fixes and added note about face fields

* Apply suggestions from code review

* showcase that condition can get an array of bools

* slight rephrase

* move the comment up

* remove duplicate broadcasting

---------

Co-authored-by: Gregory L. Wagner <[email protected]>
Co-authored-by: Taimoor Sohail <[email protected]>

Author: 3 people

docs/src/operations.md

@@ -67,7 +67,7 @@ BinaryOperation at (Center, Center, Center)
     └── 1
 ```
 
-Like `Field`s, `AbstractOperations` have a location and a grid. 
+Like `Field`s, `AbstractOperations` have a location and a grid.
 In addition to `BinaryOperation`s like the kind above, `UnaryOperation`s and `MultiaryOperation`s are also supported,
 
 ```jldoctest operations
@@ -211,3 +211,117 @@ BinaryOperation at (Center, Center, Center)
         │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
         └── 2
 ```
+
+## Averages and integrals
+
+An operation that reduces one or more dimensions of a field is called a [`Reduction`](@ref); some reductions include [`Average`](@ref), [`Integral`](@ref), and [`CumulativeIntegral`](@ref).
+
+Let's demonstrate now how we can compute the average or the integral of a field over the grid or over some part of the grid.
+We start by creating a latitude-longitude grid that only goes up to 30 degrees latitude.
+Conveniently, with this latitude extent that grid covers half the total area of the sphere, i.e., `2π * grid.radius^2`.
+
+Let's try to estimate this area using `Integral` operation.
+We create a Field, we fill it with ones and we integrate it over the whole grid.
+We use a `CenterField` for the example below, but all location combinations work similarly.
+
+
+```jldoctest operations_avg_int
+using Oceananigans
+
+grid = LatitudeLongitudeGrid(size=(60, 10, 5),
+                             latitude = (-30, 30),
+                             longitude = (0, 360),
+                             z = (-1000, 0))
+
+c = CenterField(grid)
+set!(c, 1)
+
+∫c = Field(Integral(c, dims=(1, 2)))
+
+# output
+1×1×5 Field{Nothing, Nothing, Center} reduced over dims = (1, 2) on LatitudeLongitudeGrid on CPU
+├── data: OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, size: (1, 1, 5)
+├── grid: 60×10×5 LatitudeLongitudeGrid{Float64, Periodic, Bounded, Bounded} on CPU with 3×3×3 halo and with precomputed metrics
+├── operand: Integral of BinaryOperation at (Center, Center, Center) over dims (1, 2)
+├── status: time=0.0
+└── data: 1×1×11 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:8) with eltype Float64 with indices 1:1×1:1×-2:8
+    └── max=0.0, min=0.0, mean=0.0
+```
+
+A few remarks: note that the `∫c` has locations `Nothing, Nothing, Center`; this is because we have integrated in the first two dimensions and thus it's `reduced over dims = (1, 2)`.
+Further note that `∫c` is full of zeros; its max, min, and mean values are all 0.
+No computation has been done yet.
+To compute `∫c`, we call [`compute!`](@ref),
+
+```jldoctest operations_avg_int
+compute!(∫c)
+
+# output
+1×1×5 Field{Nothing, Nothing, Center} reduced over dims = (1, 2) on LatitudeLongitudeGrid on CPU
+├── data: OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, size: (1, 1, 5)
+├── grid: 60×10×5 LatitudeLongitudeGrid{Float64, Periodic, Bounded, Bounded} on CPU with 3×3×3 halo and with precomputed metrics
+├── operand: Integral of BinaryOperation at (Center, Center, Center) over dims (1, 2)
+├── status: time=0.0
+└── data: 1×1×11 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:8) with eltype Float64 with indices 1:1×1:1×-2:8
+    └── max=2.55032e14, min=2.55032e14, mean=2.55032e14
+```
+
+Above we see that the max, min and mean of the field are all the same.
+Let's check that these values are what we expect:
+
+```jldoctest operations_avg_int
+∫c[1, 1, 1] ≈ 2π * grid.radius^2 # area of spherical zone with |φ| ≤ 30ᵒ
+
+# output
+
+true
+```
+
+We can further have conditional reduced operations. Let's compute the above integral but only for North hemisphere.
+
+First we need to define a condition which is a function with arguments `(i, j, k, grid, field)` that returns true or false.
+In this example we use `Oceananigans.Grids.φnode` to check whether the latitude is within the range we'd like.
+
+```jldoctest operations_avg_int
+using Oceananigans.Grids: φnode
+
+cond(i, j, k, grid, c) = φnode(j, grid, Center()) ≥ 0
+
+conditional_∫c = Field(Integral(c, dims=(1, 2), condition=cond)) # integrate only when condition is true
+
+# output
+1×1×5 Field{Nothing, Nothing, Center} reduced over dims = (1, 2) on LatitudeLongitudeGrid on CPU
+├── data: OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, size: (1, 1, 5)
+├── grid: 60×10×5 LatitudeLongitudeGrid{Float64, Periodic, Bounded, Bounded} on CPU with 3×3×3 halo and with precomputed metrics
+├── operand: Integral of ConditionalOperation of BinaryOperation at (Center, Center, Center) with condition cond (generic function with 1 method) over dims (1, 2)
+├── status: time=0.0
+└── data: 1×1×11 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:8) with eltype Float64 with indices 1:1×1:1×-2:8
+    └── max=0.0, min=0.0, mean=0.0
+```
+
+Above we have attached a condition to the operand.
+Now the operand is applied only when the condition is true.
+Let's compute and see if we get 1/4 of the area of the sphere
+
+```jldoctest operations_avg_int
+compute!(conditional_∫c)
+conditional_∫c[1, 1, 1] ≈ π * grid.radius^2 # area of spherical zone with 0ᵒ ≤ φ ≤ 30ᵒ
+
+# output
+true
+```
+
+Another way to do the above is to provide the `condition` keyword argument with an array of booleans.
+
+```jldoctest operations_avg_int
+cond_array = trues(size(grid))
+cond_array[:, 1:5, :] .= false # set the first half of the latitude range to false
+
+conditional_∫c = Field(Integral(c, dims=(1, 2), condition=cond_array))
+
+compute!(conditional_∫c)
+conditional_∫c[1, 1, 1] ≈ π * grid.radius^2 # area of spherical zone with 0ᵒ ≤ φ ≤ 30ᵒ
+
+# output
+true
+```

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl

@@ -4,7 +4,7 @@ import Oceananigans.Grids: on_architecture
 
 struct SplitExplicitFreeSurface{H, U, M, FT, K , S, T} <: AbstractFreeSurface{H, FT}
     η :: H
-    barotropic_velocities :: U # A namedtuple with U, V 
+    barotropic_velocities :: U # A namedtuple with U, V
     filtered_state :: M # A namedtuple with η, U, V averaged throughout the substepping
     gravitational_acceleration :: FT
     kernel_parameters :: K
@@ -21,20 +21,20 @@ end
                              averaging_kernel = averaging_shape_function,
                              timestepper = ForwardBackwardScheme())
 
-Return a `SplitExplicitFreeSurface` representing an explicit time discretization of 
+Return a `SplitExplicitFreeSurface` representing an explicit time discretization of
 a free surface dynamics with `gravitational_acceleration`. The free surface dynamics are solved by discretizing:
 ```math
-\begin{gather}
-∂_t η = - \nabla ⋅ U, \\
-∂_t U = - g H \nabla η + G^U,
-\end{gather}
+\\begin{gather}
+∂_t η = - \\nabla ⋅ U, \\\\
+∂_t U = - g H \\nabla η + G^U,
+\\end{gather}
 ```
-where ``η`` is the free surface displacement, ``U`` is the barotropic velocity vector, calculated as the vertical integral 
-of the velocity field ``u`` and ``v``, ``H`` is the column depth, ``G^U`` is the slow forcing calculated as the integral of the 
-tendency of ``u`` and ``v``, and ``g`` is the gravitational acceleration. 
+where ``η`` is the free surface displacement, ``U`` is the barotropic velocity vector, calculated as the vertical integral
+of the velocity field ``u`` and ``v``, ``H`` is the column depth, ``G^U`` is the slow forcing calculated as the integral of the
+tendency of ``u`` and ``v``, and ``g`` is the gravitational acceleration.
 
-The discretized equations are solved within a baroclinic timestep (``Δt``) by substepping with a ``Δτ < Δt``. 
-The barotropic velocities are filtered throughout the substepping and, finally, 
+The discretized equations are solved within a baroclinic timestep (``Δt``) by substepping with a ``Δτ < Δt``.
+The barotropic velocities are filtered throughout the substepping and, finally,
 the barotropic mode of the velocities at the new time step is corrected with the filtered velocities.
 
 Keyword Arguments
@@ -55,7 +55,7 @@ Keyword Arguments
 
 !!! info "Needed keyword arguments"
     Either `substeps` _or_ `cfl` need to be prescribed.
-    
+
     When `cfl` is prescribed then `grid` is also required as a positional argument.
 
 - `fixed_Δt`: The maximum baroclinic timestep allowed. If `fixed_Δt` is a `nothing` and a cfl is provided,
@@ -124,21 +124,21 @@ function split_explicit_substepping(::Nothing, substeps, fixed_Δt, grid, averag
 end
 
 # The substeps are calculated dynamically when a cfl without a fixed_Δt is provided
-function split_explicit_substepping(cfl, ::Nothing, ::Nothing, grid, averaging_kernel, gravitational_acceleration)  
+function split_explicit_substepping(cfl, ::Nothing, ::Nothing, grid, averaging_kernel, gravitational_acceleration)
     cfl = convert(eltype(grid), cfl)
     return FixedTimeStepSize(grid; cfl, averaging_kernel)
 end
 
 # The number of substeps are calculated based on the cfl and the fixed_Δt
 function split_explicit_substepping(cfl, ::Nothing, fixed_Δt, grid, averaging_kernel, gravitational_acceleration)
-    substepping = split_explicit_substepping(cfl, nothing, nothing, grid, averaging_kernel, gravitational_acceleration)    
+    substepping = split_explicit_substepping(cfl, nothing, nothing, grid, averaging_kernel, gravitational_acceleration)
     substeps    = ceil(Int, 2 * fixed_Δt / substepping.Δt_barotropic)
-    substepping = split_explicit_substepping(nothing, substeps, nothing, grid, averaging_kernel, gravitational_acceleration)        
+    substepping = split_explicit_substepping(nothing, substeps, nothing, grid, averaging_kernel, gravitational_acceleration)
     return substepping
 end
 
 # TODO: When open boundary conditions are online
-# We need to calculate the barotropic boundary conditions 
+# We need to calculate the barotropic boundary conditions
 # from the baroclinic boundary conditions by integrating the BC upwards
 @inline  west_barotropic_bc(baroclinic_velocity) = baroclinic_velocity.boundary_conditions.west
 @inline  east_barotropic_bc(baroclinic_velocity) = baroclinic_velocity.boundary_conditions.east
@@ -202,7 +202,7 @@ function materialize_free_surface(free_surface::SplitExplicitFreeSurface, veloci
                                     timestepper)
 end
 
-# (p = 2, q = 4, r = 0.18927) minimize dispersion error from Shchepetkin and McWilliams (2005): https://doi.org/10.1016/j.ocemod.2004.08.002 
+# (p = 2, q = 4, r = 0.18927) minimize dispersion error from Shchepetkin and McWilliams (2005): https://doi.org/10.1016/j.ocemod.2004.08.002
 @inline function averaging_shape_function(τ::FT; p = 2, q = 4, r = FT(0.18927)) where FT
     τ₀ = (p + 2) * (p + q + 2) / (p + 1) / (p + q + 1)
 
@@ -274,7 +274,7 @@ Base.show(io::IO, sefs::SplitExplicitFreeSurface) = print(io, "$(summary(sefs))\
 maybe_extend_halos(TX, TY, grid, ::FixedTimeStepSize) = grid
 
 function maybe_extend_halos(TX, TY, grid, substepping::FixedSubstepNumber)
-    
+
     old_halos = halo_size(grid)
     Nsubsteps = length(substepping.averaging_weights)
     step_halo = Nsubsteps + 1

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl

@@ -2,14 +2,14 @@
     struct ForwardBackwardScheme
 
 A timestepping scheme used for substepping in the split-explicit free surface solver.
-    
+
 The equations are evolved as follows:
 ```math
-\begin{gather}
-η^{m+1} = η^m - Δτ (∂_x U^m + ∂_y V^m), \\
-U^{m+1} = U^m - Δτ (∂_x η^{m+1} - G^U), \\
+\\begin{gather}
+η^{m+1} = η^m - Δτ (∂_x U^m + ∂_y V^m), \\\\
+U^{m+1} = U^m - Δτ (∂_x η^{m+1} - G^U), \\\\
 V^{m+1} = V^m - Δτ (∂_y η^{m+1} - G^V).
-\end{gather}
+\\end{gather}
 ```
 """
 struct ForwardBackwardScheme end
@@ -50,31 +50,31 @@ free surface at time-step `m + 1/2`:
 The equations are evolved as follows:
 
 ```math
-\begin{gather}
-η^{m+1} = η^m - Δτ g H (∂_x Ũ + ∂y Ṽ), \\
-U^{m+1} = U^m - Δτ (∂_x η̃ - G^U), \\
+\\begin{gather}
+η^{m+1} = η^m - Δτ g H (∂_x Ũ + ∂y Ṽ), \\\\
+U^{m+1} = U^m - Δτ (∂_x η̃ - G^U), \\\\
 V^{m+1} = V^m - Δτ (∂_y η̃ - G^V),
-\end{gather}
-```    
+\\end{gather}
+```
 
-where `η̃`, `Ũ` and `Ṽ` are the AB3 time-extrapolated values of free surface, 
+where `η̃`, `Ũ` and `Ṽ` are the AB3 time-extrapolated values of free surface,
 barotropic zonal and meridional velocities, respectively:
 
 ```math
-\begin{gather}
-Ũ = α U^m + θ U^{m-1} + β U^{m-2}, \\
-Ṽ = α V^m + θ V^{m-1} + β V^{m-2}, \\
+\\begin{gather}
+Ũ = α U^m + θ U^{m-1} + β U^{m-2}, \\\\
+Ṽ = α V^m + θ V^{m-1} + β V^{m-2}, \\\\
 η̃ = δ η^{m+1} + μ η^m + γ η^{m-1} + ϵ η^{m-2}.
-\end{gather}
+\\end{gather}
 ```
 
-The default values for the time-extrapolation coefficients, described by [Shchepetkin2005](@citet), 
+The default values for the time-extrapolation coefficients, described by [Shchepetkin2005](@citet),
 correspond to the best stability range for the AB3 algorithm.
 """
-AdamsBashforth3Scheme(; β = 0.281105, α = 1.5 + β, θ = - 0.5 - 2β, γ = 0.088, δ = 0.614, ϵ = 0.013, μ = 1 - δ - γ - ϵ) = 
+AdamsBashforth3Scheme(; β = 0.281105, α = 1.5 + β, θ = - 0.5 - 2β, γ = 0.088, δ = 0.614, ϵ = 0.013, μ = 1 - δ - γ - ϵ) =
         AdamsBashforth3Scheme(nothing, nothing, nothing, nothing, nothing, nothing, nothing, β, α, θ, γ, δ, ϵ, μ)
 
-Adapt.adapt_structure(to, t::AdamsBashforth3Scheme) = 
+Adapt.adapt_structure(to, t::AdamsBashforth3Scheme) =
     AdamsBashforth3Scheme(
         Adapt.adapt(to, t.ηᵐ  ),
         Adapt.adapt(to, t.ηᵐ⁻¹),
@@ -104,7 +104,7 @@ function materialize_timestepper(t::AdamsBashforth3Scheme, grid, free_surface, v
     δ = convert(FT, t.δ)
     ϵ = convert(FT, t.ϵ)
     μ = convert(FT, t.μ)
-    
+
     return AdamsBashforth3Scheme(ηᵐ, ηᵐ⁻¹, ηᵐ⁻², Uᵐ⁻¹, Uᵐ⁻², Vᵐ⁻¹, Vᵐ⁻², β, α, θ, γ, δ, ϵ, μ)
 end