Implement Scan, generalizing Reduction to accumulating scans like…

… `cumsum!` (#3590)

* Generalize Reduction to include accumulating Scans

* Export Accumulation

* Bugfixes

* Implement reverse_cumsum and CumulativeIntegral

* Export cumulative integral

* Fix Reduction tests

* Working?

* Fix Integral example and add new one for CumulativeIntegral

* Many more tests

* Fix up tests

* Fix what looks like a bug for cumulative integral and add another test

* Add missing method for cumsum?

* Oh, cumsum doesnt take a unary first arg

* We give up and just implement accumulation!

* Delete supfluous cumsum! plus fix doctest

* Fix docstring for CumulativeIntegral

src/AbstractOperations/AbstractOperations.jl

@@ -2,7 +2,7 @@ module AbstractOperations
 
 export ∂x, ∂y, ∂z, @at, @unary, @binary, @multiary
 export Δx, Δy, Δz, Ax, Ay, Az, volume
-export Average, Integral, KernelFunctionOperation
+export Average, Integral, CumulativeIntegral, KernelFunctionOperation
 export UnaryOperation, Derivative, BinaryOperation, MultiaryOperation, ConditionalOperation
 
 using Base: @propagate_inbounds

src/AbstractOperations/metric_field_reductions.jl

@@ -2,8 +2,7 @@ using Statistics: mean!, sum!
 
 using Oceananigans.Utils: tupleit
 using Oceananigans.Grids: regular_dimensions
-using Oceananigans.Fields: condition_operand
-import Oceananigans.Fields: Reduction
+using Oceananigans.Fields: Scan, condition_operand, reverse_cumsum!, AbstractReducing, AbstractAccumulating
 
 ##### 
 ##### Metric inference
@@ -24,15 +23,29 @@ reduction_grid_metric(dims) = dims === tuple(1)  ? Δx :
 ##### Metric reductions
 ##### 
 
-struct Average end
+struct Averaging <: AbstractReducing end
+const Average = Scan{<:Averaging}
+Base.summary(r::Average) = string("Average of ", summary(r.operand), " over dims ", r.dims)
 
-function Reduction(avg::Average, field::AbstractField; condition = nothing, mask = 0, dims)
+"""
+    Average(field::AbstractField; dims=:, condition=nothing, mask=0)
+
+Return `Reduction` representing a spatial average of `field` over `dims`.
+
+Over regularly-spaced dimensions this is equivalent to a numerical `mean!`.
+
+Over dimensions of variable spacing, `field` is multiplied by the
+appropriate grid length, area or volume, and divided by the total
+spatial extent of the interval.
+"""
+function Average(field::AbstractField; dims=:, condition=nothing, mask=0)
     dims = dims isa Colon ? (1, 2, 3) : tupleit(dims)
     dx = reduction_grid_metric(dims)
 
     if all(d in regular_dimensions(field.grid) for d in dims)
         # Dimensions being reduced are regular; just use mean!
-        return Reduction(mean!, condition_operand(field, condition, mask); dims)
+        operand = condition_operand(field, condition, mask)
+        return Scan(Averaging(), mean!, operand, dims)
     else
         # Compute "size" (length, area, or volume) of averaging region
         metric = GridMetricOperation(location(field), dx, field.grid)
@@ -41,33 +54,18 @@ function Reduction(avg::Average, field::AbstractField; condition = nothing, mask
         # Construct summand of the Average
         L⁻¹_field_dx = field * dx / L
 
-        return Reduction(sum!, condition_operand(L⁻¹_field_dx, condition, mask), dims)
+        operand = condition_operand(L⁻¹_field_dx, condition, mask)
+
+        return Scan(Averaging(), sum!, operand, dims)
     end
 end
 
-"""
-    Average(field::AbstractField; condition = nothing, mask = 0, dims=:)
-
-Return `Reduction` representing a spatial average of `field` over `dims`.
-
-Over regularly-spaced dimensions this is equivalent to a numerical `mean!`.
-
-Over dimensions of variable spacing, `field` is multiplied by the
-appropriate grid length, area or volume, and divided by the total
-spatial extent of the interval.
-"""
-Average(field::AbstractField; condition = nothing, mask = 0, dims=:) = Reduction(Average(), field; condition, mask, dims)
-
-struct Integral end
-
-function Reduction(int::Integral, field::AbstractField; condition = nothing, mask = 0, dims)
-    dims = dims isa Colon ? (1, 2, 3) : tupleit(dims)
-    dx = reduction_grid_metric(dims)
-    return Reduction(sum!, condition_operand(field * dx, condition, mask), dims)
-end
+struct Integrating <: AbstractReducing end
+const Integral = Scan{<:Integrating}
+Base.summary(r::Integral) = string("Integral of ", summary(r.operand), " over dims ", r.dims)
 
 """
-    Integral(field::AbstractField; condition = nothing, mask = 0, dims=:)
+    Integral(field::AbstractField; dims=:, condition=nothing, mask=0)
 
 
 Return a `Reduction` representing a spatial integral of `field` over `dims`.
@@ -95,7 +93,7 @@ julia> set!(f, (x, y, z) -> x * y * z)
     └── max=0.823975, min=0.000244141, mean=0.125
 
 julia> ∫f = Integral(f)
-sum! over dims (1, 2, 3) of BinaryOperation at (Center, Center, Center)
+Integral of BinaryOperation at (Center, Center, Center) over dims (1, 2, 3)
 └── operand: BinaryOperation at (Center, Center, Center)
     └── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
 
@@ -107,12 +105,69 @@ julia> ∫f[1, 1, 1]
 0.125
 ```
 """
-Integral(field::AbstractField; condition = nothing, mask = 0, dims=:) =
-    Reduction(Integral(), condition_operand(field, condition, mask); dims)
+function Integral(field::AbstractField; dims=:, condition=nothing, mask=0)
+    dims = dims isa Colon ? (1, 2, 3) : tupleit(dims)
+    dx = reduction_grid_metric(dims)
+    operand = condition_operand(field * dx, condition, mask)
+    return Scan(Integrating(), sum!, operand, dims)
+end
 
 #####
-##### show
+##### CumulativeIntegral
 #####
 
-Base.summary(r::Reduction{<:Average}) = string("Average of ", summary(r.operand), " over dims ", r.dims)
-Base.summary(r::Reduction{<:Integral}) = string("Integral of ", summary(r.operand), " over dims ", r.dims)
+struct CumulativelyIntegrating <: AbstractAccumulating end
+const CumulativeIntegral = Scan{<:CumulativelyIntegrating}
+Base.summary(c::CumulativeIntegral) = string("CumulativeIntegral of ", summary(c.operand), " over dims ", c.dims)
+
+"""
+    CumulativeIntegral(field::AbstractField; dims, reverse=false, condition=nothing, mask=0)
+
+Return an `Accumulation` representing the cumulative spatial integral of `field` over `dims`.
+
+Example
+=======
+
+Compute the cumulative integral of ``f(z) = z`` over z ∈ [0, 1].
+
+```jldoctest
+julia> using Oceananigans
+
+julia> grid = RectilinearGrid(size=8, z=(0, 1), topology=(Flat, Flat, Bounded));
+
+julia> c = CenterField(grid);
+
+julia> set!(c, z -> z)
+1×1×8 Field{Center, Center, Center} on RectilinearGrid on CPU
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── boundary conditions: FieldBoundaryConditions
+│   └── west: Nothing, east: Nothing, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
+└── data: 1×1×14 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:11) with eltype Float64 with indices 1:1×1:1×-2:11
+    └── max=0.9375, min=0.0625, mean=0.5
+
+julia> C_op = CumulativeIntegral(c, dims=3)
+CumulativeIntegral of BinaryOperation at (Center, Center, Center) over dims 3
+└── operand: BinaryOperation at (Center, Center, Center)
+    └── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+
+julia> C = compute!(Field(C_op))
+1×1×8 Field{Center, Center, Center} on RectilinearGrid on CPU
+├── data: OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, size: (1, 1, 8)
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── operand: CumulativeIntegral of BinaryOperation at (Center, Center, Center) over dims 3
+├── status: time=0.0
+└── data: 1×1×14 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:11) with eltype Float64 with indices 1:1×1:1×-2:11
+    └── max=0.5, min=0.0078125, mean=0.199219
+
+julia> C[1, 1, 8]
+0.5
+```
+"""
+function CumulativeIntegral(field::AbstractField; dims, reverse=false, condition=nothing, mask=0)
+    dims ∈ (1, 2, 3) || throw(ArgumentError("CumulativeIntegral only supports dims=1, 2, or 3."))
+    maybe_reverse_cumsum = reverse ? reverse_cumsum! : cumsum!
+    dx = reduction_grid_metric(dims)
+    operand = condition_operand(field * dx, condition, mask)
+    return Scan(CumulativelyIntegrating(), maybe_reverse_cumsum, operand, dims)
+end
+

src/Fields/Fields.jl

@@ -1,7 +1,7 @@
 module Fields
 
 export Face, Center
-export AbstractField, Field, Average, Integral, Reduction, field
+export AbstractField, Field, Average, Integral, Reduction, Accumulation, field
 export CenterField, XFaceField, YFaceField, ZFaceField
 export BackgroundField
 export interior, data, xnode, ynode, znode, location
@@ -21,7 +21,7 @@ include("constant_field.jl")
 include("function_field.jl")
 include("field_boundary_buffers.jl")
 include("field.jl")
-include("field_reductions.jl")
+include("scans.jl")
 include("regridding_fields.jl")
 include("field_tuples.jl")
 include("background_fields.jl")

src/Fields/field.jl

@@ -619,8 +619,6 @@ function reduced_location(loc; dims)
     end
 end
 
-reduced_indices(indices; dims) = Tuple(i ∈ dims ? Colon() : indices[i] for i in 1:3)
-
 function reduced_dimension(loc)
     dims = ()
     for i in 1:3