diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index 9656716307..390597cc6c 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/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 diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 9663260de3..024d8c2877 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/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 + diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index 0c749009b9..56b2d11594 100644 --- a/src/Fields/Fields.jl +++ b/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") diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 13f3b21b35..0f5f476a0b 100644 --- a/src/Fields/field.jl +++ b/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 diff --git a/src/Fields/field_reductions.jl b/src/Fields/field_reductions.jl deleted file mode 100644 index a97e17d510..0000000000 --- a/src/Fields/field_reductions.jl +++ /dev/null @@ -1,97 +0,0 @@ -##### -##### Reductions of AbstractField -##### - -struct Reduction{R, O, D} - reduce! :: R - operand :: O - dims :: D -end - -""" - Reduction(reduce!, operand; dims) - -Return a `Reduction` of `operand` with `reduce!`, along `dims`. Note that `Reduction` -expects `reduce!` to operate in-place. - -Example -======= - -```jldoctest -using Oceananigans - -Nx, Ny, Nz = 3, 3, 3 - -grid = RectilinearGrid(size=(Nx, Ny, Nz), x=(0, 1), y=(0, 1), z=(0, 1), - topology=(Periodic, Periodic, Periodic)) - -c = CenterField(grid) - -set!(c, (x, y, z) -> x + y + z) - -max_c² = Field(Reduction(maximum!, c^2, dims=3)) - -compute!(max_c²) - -max_c²[1:Nx, 1:Ny] - -# output -3×3 Matrix{Float64}: - 1.36111 2.25 3.36111 - 2.25 3.36111 4.69444 - 3.36111 4.69444 6.25 -``` -""" -Reduction(reduce!, operand; dims) = Reduction(reduce!, operand, dims) - -location(r::Reduction) = reduced_location(location(r.operand); dims=r.dims) - -function Field(reduction::Reduction; - data = nothing, - indices = indices(reduction.operand), - recompute_safely = true) - - operand = reduction.operand - grid = operand.grid - LX, LY, LZ = loc = location(reduction) - indices = reduced_indices(indices; dims=reduction.dims) - - if isnothing(data) - data = new_data(grid, loc, indices) - recompute_safely = false - end - - boundary_conditions = FieldBoundaryConditions(grid, loc, indices) - status = recompute_safely ? nothing : FieldStatus() - - return Field(loc, grid, data, boundary_conditions, indices, reduction, status) -end - -const ReducedComputedField = Field{<:Any, <:Any, <:Any, <:Reduction} - -function compute!(field::ReducedComputedField, time=nothing) - reduction = field.operand - compute_at!(reduction.operand, time) - reduction.reduce!(field, reduction.operand) - return field -end - -##### -##### show -##### - -Base.show(io::IO, field::ReducedComputedField) = - print(io, "$(summary(field))\n", - "├── data: $(typeof(field.data)), size: $(size(field))\n", - "├── grid: $(summary(field.grid))\n", - "├── operand: $(summary(field.operand))\n", - "└── status: $(summary(field.status))") - -Base.summary(r::Reduction) = string(r.reduce!, - " over dims ", r.dims, - " of ", summary(r.operand)) - -Base.show(io::IO, r::Reduction) = - print(io, "$(summary(r))\n", - "└── operand: $(summary(r.operand))\n", - " └── grid: $(summary(r.operand.grid))") diff --git a/src/Fields/scans.jl b/src/Fields/scans.jl new file mode 100644 index 0000000000..7ebf6fc20d --- /dev/null +++ b/src/Fields/scans.jl @@ -0,0 +1,307 @@ +using KernelAbstractions: @kernel, @index + +##### +##### "Scans" of AbstractField. +##### + +""" + Scan{T, R, O, D} + +An operand for `Field` that is computed by traversing a dimension. +This includes `Reducing` `Scan`s like `sum!` and `maximum!`, as well as +`Accumulating` `Scan`s like `cumsum!`. +""" +struct Scan{T, R, O, D} + type :: T + scan! :: R + operand :: O + dims :: D +end + +abstract type AbstractReducing end +abstract type AbstractAccumulating end + +struct Reducing <: AbstractReducing end +struct Accumulating <: AbstractAccumulating end + +Base.summary(::Reducing) = "Reducing" +Base.summary(::Accumulating) = "Accumulating" + +const Reduction = Scan{<:AbstractReducing} +const Accumulation = Scan{<:AbstractAccumulating} + +scan_indices(::AbstractReducing, indices; dims) = Tuple(i ∈ dims ? Colon() : indices[i] for i in 1:3) +scan_indices(::AbstractAccumulating, indices; dims) = indices + +Base.summary(s::Scan) = string(summary(s.type), " ", + s.scan!, + " over dims ", s.dims, + " of ", summary(s.operand)) + +function Field(scan::Scan; + data = nothing, + indices = indices(scan.operand), + recompute_safely = true) + + operand = scan.operand + grid = operand.grid + LX, LY, LZ = loc = location(scan) + indices = scan_indices(scan.type, indices; dims=scan.dims) + + if isnothing(data) + data = new_data(grid, loc, indices) + recompute_safely = false + end + + boundary_conditions = FieldBoundaryConditions(grid, loc, indices) + status = recompute_safely ? nothing : FieldStatus() + + return Field(loc, grid, data, boundary_conditions, indices, scan, status) +end + +const ScannedComputedField = Field{<:Any, <:Any, <:Any, <:Scan} + +function compute!(field::ScannedComputedField, time=nothing) + s = field.operand + compute_at!(s.operand, time) + + if s.type isa AbstractReducing + s.scan!(field, s.operand) + elseif s.type isa AbstractAccumulating + s.scan!(field, s.operand; dims=s.dims) + end + + return field +end + +##### +##### show +##### + +function Base.show(io::IO, field::ScannedComputedField) + print(io, summary(field), '\n', + "├── data: ", typeof(field.data), ", size: ", size(field), '\n', + "├── grid: ", summary(field.grid), '\n', + "├── operand: ", summary(field.operand), '\n', + "├── status: ", summary(field.status), '\n') + + data_str = string("└── data: ", summary(field.data), '\n', + " └── ", data_summary(field)) + + print(io, data_str) +end + +Base.show(io::IO, s::Scan) = + print(io, "$(summary(s))\n", + "└── operand: $(summary(s.operand))\n", + " └── grid: $(summary(s.operand.grid))") + +##### +##### Reductions (where the output has fewer dimensions than the input) +##### + +""" + Reduction(reduce!, operand; dims) + +Return a `Reduction` of `operand` with `reduce!`, where `reduce!` can be called with + +``` +reduce!(field, operand) +``` + +to reduce `operand` along `dims` and store in `field`. + +Example +======= + +```jldoctest +using Oceananigans + +Nx, Ny, Nz = 3, 3, 3 + +grid = RectilinearGrid(size=(Nx, Ny, Nz), x=(0, 1), y=(0, 1), z=(0, 1), + topology=(Periodic, Periodic, Periodic)) + +c = CenterField(grid) + +set!(c, (x, y, z) -> x + y + z) + +max_c² = Field(Reduction(maximum!, c^2, dims=3)) + +compute!(max_c²) + +max_c²[1:Nx, 1:Ny] + +# output +3×3 Matrix{Float64}: + 1.36111 2.25 3.36111 + 2.25 3.36111 4.69444 + 3.36111 4.69444 6.25 +``` +""" +Reduction(reduce!, operand; dims) = Scan(Reducing(), reduce!, operand, dims) +location(r::Reduction) = reduced_location(location(r.operand); dims=r.dims) + +##### +##### Accumulations (where the output has the same dimensions as the input) +##### + +""" + Accumulation(accumulate!, operand; dims) + +Return a `Accumulation` of `operand` with `accumulate!`, where `accumulate!` can be called with + +``` +accumulate!(field, operand; dims) +``` + +to accumulate `operand` along `dims` and store in `field`. + +Example +======= + +```jldoctest +using Oceananigans + +Nx, Ny, Nz = 3, 3, 3 + +grid = RectilinearGrid(size=(Nx, Ny, Nz), x=(0, 1), y=(0, 1), z=(0, 1), + topology=(Periodic, Periodic, Periodic)) + +c = CenterField(grid) + +set!(c, (x, y, z) -> x + y + z) + +cumsum_c² = Field(Accumulation(cumsum!, c^2, dims=3)) + +compute!(cumsum_c²) + +cumsum_c²[1:Nx, 1:Ny, 1:Nz] + +# output +3×3×3 Array{Float64, 3}: +[:, :, 1] = + 0.25 0.694444 1.36111 + 0.694444 1.36111 2.25 + 1.36111 2.25 3.36111 + +[:, :, 2] = + 0.944444 2.05556 3.61111 + 2.05556 3.61111 5.61111 + 3.61111 5.61111 8.05556 + +[:, :, 3] = + 2.30556 4.30556 6.97222 + 4.30556 6.97222 10.3056 + 6.97222 10.3056 14.3056 +``` +""" +Accumulation(accumulate!, operand; dims) = Scan(Accumulating(), accumulate!, operand, dims) +location(a::Accumulation) = location(a.operand) + +##### +##### Some custom scans +##### + +# directions +struct Forward end +struct Reverse end + +@inline increment(::Forward, idx) = idx + 1 +@inline decrement(::Forward, idx) = idx - 1 + +@inline increment(::Reverse, idx) = idx - 1 +@inline decrement(::Reverse, idx) = idx + 1 + +# TODO: if there is a way to re-use Base stuff rather than write ourselves, that might be preferred. +# Also: these don't work on ImmersedBoundaryGrid out of the box, +# we need to somehow find the neutral element for `op` +Base.accumulate!(op, B::Field, A::AbstractField; dims::Integer) = + directional_accumulate!(op, B, A, dims, Forward()) + +reverse_accumulate!(op, B::Field, A::AbstractField; dims::Integer) = + directional_accumulate!(op, B, A, dims, Reverse()) + +function Base.cumsum!(B::Field, A::AbstractField; dims, condition=nothing, mask=get_neutral_mask(Base.sum!)) + Ac = condition_operand(A, condition, mask) + return directional_accumulate!(Base.add_sum, B, Ac, dims, Forward()) +end + +function reverse_cumsum!(B::Field, A::AbstractField; dims, condition=nothing, mask=get_neutral_mask(Base.sum!)) + Ac = condition_operand(A, condition, mask) + return directional_accumulate!(Base.add_sum, B, Ac, dims, Reverse()) +end + +function directional_accumulate!(op, B, A, dim, direction) + + grid = B.grid + arch = architecture(B) + + # TODO: this won't work on windowed fields + # To fix this we can change config, start, and finish. + if dim == 1 + config = :yz + kernel = accumulate_x + elseif dim == 2 + config = :xz + kernel = accumulate_y + elseif dim == 3 + config = :xy + kernel = accumulate_z + end + + if direction isa Forward + start = 1 + finish = size(B, dim) + elseif direction isa Reverse + start = size(B, dim) + finish = 1 + end + + launch!(arch, grid, config, kernel, op, B, A, start, finish, direction) + + return B +end + +@inline function accumulation_range(dir, start, finish) + by = increment(dir, 0) + from = increment(dir, start) + return StepRange(from, by, finish) +end + +@kernel function accumulate_x(op, B, A, start, finish, dir) + j, k = @index(Global, NTuple) + + # Initialize + @inbounds B[start, j, k] = Base.reduce_first(op, A[start, j, k]) + + for i in accumulation_range(dir, start, finish) + pr = decrement(dir, i) + @inbounds B[i, j, k] = op(B[pr, j, k], A[i, j, k]) + end +end + +@kernel function accumulate_y(op, B, A, start, finish, dir) + i, k = @index(Global, NTuple) + + # Initialize + @inbounds B[i, start, k] = Base.reduce_first(op, A[i, start, k]) + + for j in accumulation_range(dir, start, finish) + pr = decrement(dir, j) + @inbounds B[i, j, k] = op(B[i, pr, k], A[i, j, k]) + end +end + +@kernel function accumulate_z(op, B, A, start, finish, dir) + i, j = @index(Global, NTuple) + + # Initialize + @inbounds B[i, j, start] = Base.reduce_first(op, A[i, j, start]) + + for k in accumulation_range(dir, start, finish) + pr = decrement(dir, k) + @inbounds B[i, j, k] = op(B[i, j, pr], A[i, j, k]) + end +end + diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 743f53a512..0faca4afa9 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -41,7 +41,7 @@ export # Fields and field manipulation Field, CenterField, XFaceField, YFaceField, ZFaceField, - Average, Integral, Reduction, BackgroundField, + Average, Integral, CumulativeIntegral, Reduction, Accumulation, BackgroundField, interior, set!, compute!, regrid!, location, # Forcing functions diff --git a/test/runtests.jl b/test/runtests.jl index 95e9020214..b954e1d14b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -31,7 +31,7 @@ CUDA.allowscalar() do include("test_boundary_conditions.jl") include("test_field.jl") include("test_regrid.jl") - include("test_field_reductions.jl") + include("test_field_scans.jl") include("test_halo_regions.jl") include("test_coriolis.jl") include("test_buoyancy.jl") diff --git a/test/test_field_reductions.jl b/test/test_field_scans.jl similarity index 61% rename from test/test_field_reductions.jl rename to test/test_field_scans.jl index f44712ab68..8fb864639a 100644 --- a/test/test_field_reductions.jl +++ b/test/test_field_scans.jl @@ -3,11 +3,12 @@ include("dependencies_for_runtests.jl") using Statistics using Oceananigans.Architectures: on_architecture using Oceananigans.AbstractOperations: BinaryOperation -using Oceananigans.Fields: ReducedField, CenterField, ZFaceField, compute_at!, @compute +using Oceananigans.Fields: ReducedField, CenterField, ZFaceField, compute_at!, @compute, reverse_cumsum! using Oceananigans.BoundaryConditions: fill_halo_regions! using Oceananigans.Grids: halo_size trilinear(x, y, z) = x + y + z +interior_array(a, i, j, k) = Array(interior(a, i, j, k)) @testset "Fields computed by Reduction" begin @info "Testing Fields computed by reductions..." @@ -23,8 +24,8 @@ trilinear(x, y, z) = x + y + z x=(0, 2), y=(0, 2), z=[0, 1, 2], topology = (Periodic, Periodic, Bounded)) - @testset "Averaged fields [$arch_str]" begin - @info " Testing averaged Fields [$arch_str]" + @testset "Averaged and integrated fields [$arch_str]" begin + @info " Testing averaged and integrated Fields [$arch_str]" for grid in (regular_grid, xy_regular_grid) @@ -56,6 +57,38 @@ trilinear(x, y, z) = x + y + z @compute ζxy = Field(Average(ζ, dims=(1, 2))) @compute ζx = Field(Average(ζ, dims=1)) + @compute Θxyz = Field(Integral(T, dims=(1, 2, 3))) + @compute Θxy = Field(Integral(T, dims=(1, 2))) + @compute Θx = Field(Integral(T, dims=1)) + + @compute Wxyz = Field(Integral(w, dims=(1, 2, 3))) + @compute Wxy = Field(Integral(w, dims=(1, 2))) + @compute Wx = Field(Integral(w, dims=1)) + + @compute Zxyz = Field(Integral(ζ, dims=(1, 2, 3))) + @compute Zxy = Field(Integral(ζ, dims=(1, 2))) + @compute Zx = Field(Integral(ζ, dims=1)) + + @compute Tcx = Field(CumulativeIntegral(T, dims=1)) + @compute Tcy = Field(CumulativeIntegral(T, dims=2)) + @compute Tcz = Field(CumulativeIntegral(T, dims=3)) + @compute wcx = Field(CumulativeIntegral(w, dims=1)) + @compute wcy = Field(CumulativeIntegral(w, dims=2)) + @compute wcz = Field(CumulativeIntegral(w, dims=3)) + @compute ζcx = Field(CumulativeIntegral(ζ, dims=1)) + @compute ζcy = Field(CumulativeIntegral(ζ, dims=2)) + @compute ζcz = Field(CumulativeIntegral(ζ, dims=3)) + + @compute Trx = Field(CumulativeIntegral(T, dims=1, reverse=true)) + @compute Try = Field(CumulativeIntegral(T, dims=2, reverse=true)) + @compute Trz = Field(CumulativeIntegral(T, dims=3, reverse=true)) + @compute wrx = Field(CumulativeIntegral(w, dims=1, reverse=true)) + @compute wry = Field(CumulativeIntegral(w, dims=2, reverse=true)) + @compute wrz = Field(CumulativeIntegral(w, dims=3, reverse=true)) + @compute ζrx = Field(CumulativeIntegral(ζ, dims=1, reverse=true)) + @compute ζry = Field(CumulativeIntegral(ζ, dims=2, reverse=true)) + @compute ζrz = Field(CumulativeIntegral(ζ, dims=3, reverse=true)) + for T′ in (Tx, Txy) @test T′.operand.operand === T end @@ -68,9 +101,30 @@ trilinear(x, y, z) = x + y + z @test ζ′.operand.operand === ζ end - for f in (wx, wxy, Tx, Txy, ζx, ζxy) + for f in (wx, wxy, Tx, Txy, ζx, ζxy, Wx, Wxy, Θx, Θxy, Zx, Zxy) @test f.operand isa Reduction - @test f.operand.reduce! === mean! + end + + for f in (Tcx, Tcy, Tcz, Trx, Try, Trz, + wcx, wcy, wcz, wrx, wry, wrz, + ζcx, ζcy, ζcz, ζrx, ζry, ζrz) + @test f.operand isa Accumulation + end + + for f in (wx, wxy, Tx, Txy, ζx, ζxy) + @test f.operand.scan! === mean! + end + + for f in (wx, wxy, Tx, Txy, ζx, ζxy) + @test f.operand.scan! === mean! + end + + for f in (Tcx, Tcy, Tcz, wcx, wcy, wcz, ζcx, ζcy, ζcz) + @test f.operand.scan! === cumsum! + end + + for f in (Trx, Try, Trz, wrx, wry, wrz, ζrx, ζry, ζrz) + @test f.operand.scan! === reverse_cumsum! end @test Txyz.operand isa Reduction @@ -79,13 +133,13 @@ trilinear(x, y, z) = x + y + z # Different behavior for regular grid z vs not. if grid === regular_grid - @test Txyz.operand.reduce! === mean! - @test wxyz.operand.reduce! === mean! + @test Txyz.operand.scan! === mean! + @test wxyz.operand.scan! === mean! @test Txyz.operand.operand === T @test wxyz.operand.operand === w else - @test Txyz.operand.reduce! === sum! - @test wxyz.operand.reduce! === sum! + @test Txyz.operand.scan! === sum! + @test wxyz.operand.scan! === sum! @test Txyz.operand.operand isa BinaryOperation @test wxyz.operand.operand isa BinaryOperation end @@ -98,28 +152,82 @@ trilinear(x, y, z) = x + y + z @test wxyz.operand.dims === (1, 2, 3) @test CUDA.@allowscalar Txyz[1, 1, 1] ≈ 3 - @test Array(interior(Txy))[1, 1, :] ≈ [2.5, 3.5] - @test Array(interior(Tx))[1, :, :] ≈ [[2, 3] [3, 4]] + @test interior_array(Txy, 1, 1, :) ≈ [2.5, 3.5] + @test interior_array(Tx, 1, :, :) ≈ [[2, 3] [3, 4]] @test CUDA.@allowscalar wxyz[1, 1, 1] ≈ 3 - @test Array(interior(wxy))[1, 1, :] ≈ [2, 3, 4] - @test Array(interior(wx))[1, :, :] ≈ [[1.5, 2.5] [2.5, 3.5] [3.5, 4.5]] + @test interior_array(wxy, 1, 1, :) ≈ [2, 3, 4] + @test interior_array(wx, 1, :, :) ≈ [[1.5, 2.5] [2.5, 3.5] [3.5, 4.5]] + + averages_1d = (Tx, wx, ζx) + integrals_1d = (Θx, Wx, Zx) + + for (a, i) in zip(averages_1d, integrals_1d) + @test interior(i) == 2 .* interior(a) + end + + averages_2d = (Txy, wxy, ζxy) + integrals_2d = (Θxy, Wxy, Zxy) + + for (a, i) in zip(averages_2d, integrals_2d) + @test interior(i) == 4 .* interior(a) + end + + # T(x, y, z) = x + y + z + # T(0.5, 0.5, z) = [1.5, 2.5] + @test interior_array(Tcx, :, 1, 1) ≈ [1.5, 4] + @test interior_array(Tcy, 1, :, 1) ≈ [1.5, 4] + @test interior_array(Tcz, 1, 1, :) ≈ [1.5, 4] + + @test interior_array(Trx, :, 1, 1) ≈ [4, 2.5] + @test interior_array(Try, 1, :, 1) ≈ [4, 2.5] + @test interior_array(Trz, 1, 1, :) ≈ [4, 2.5] + + # w(x, y, z) = x + y + z + # w(0.5, 0.5, z) = [1, 2, 3] + # w(x, 0.5, 0) = w(0.5, y, 0) = [1, 2] + @test interior_array(wcx, :, 1, 1) ≈ [1, 3] + @test interior_array(wcy, 1, :, 1) ≈ [1, 3] + @test interior_array(wcz, 1, 1, :) ≈ [1, 3, 6] + + @test interior_array(wrx, :, 1, 1) ≈ [3, 2] + @test interior_array(wry, 1, :, 1) ≈ [3, 2] + @test interior_array(wrz, 1, 1, :) ≈ [6, 5, 3] @compute Txyz = CUDA.@allowscalar Field(Average(T, condition=T.>3)) @compute Txy = CUDA.@allowscalar Field(Average(T, dims=(1, 2), condition=T.>3)) @compute Tx = CUDA.@allowscalar Field(Average(T, dims=1, condition=T.>2)) @test CUDA.@allowscalar Txyz[1, 1, 1] ≈ 3.75 - @test Array(interior(Txy))[1, 1, :] ≈ [3.5, 11.5/3] - @test Array(interior(Tx))[1, :, :] ≈ [[2.5, 3] [3, 4]] + @test interior_array(Txy, 1, 1, :) ≈ [3.5, 11.5/3] + @test interior_array(Tx, 1, :, :) ≈ [[2.5, 3] [3, 4]] @compute wxyz = CUDA.@allowscalar Field(Average(w, condition=w.>3)) @compute wxy = CUDA.@allowscalar Field(Average(w, dims=(1, 2), condition=w.>2)) @compute wx = CUDA.@allowscalar Field(Average(w, dims=1, condition=w.>1)) @test CUDA.@allowscalar wxyz[1, 1, 1] ≈ 4.25 - @test Array(interior(wxy))[1, 1, :] ≈ [3, 10/3, 4] - @test Array(interior(wx))[1, :, :] ≈ [[2, 2.5] [2.5, 3.5] [3.5, 4.5]] + @test interior_array(wxy, 1, 1, :) ≈ [3, 10/3, 4] + @test interior_array(wx, 1, :, :) ≈ [[2, 2.5] [2.5, 3.5] [3.5, 4.5]] + + # A bit more for cumulative integral + @compute T2cx = Field(CumulativeIntegral(2 * T, dims=1)) + @compute T2cy = Field(CumulativeIntegral(2 * T, dims=2)) + @compute T2cz = Field(CumulativeIntegral(2 * T, dims=3)) + + @compute T2rx = Field(CumulativeIntegral(2 * T, dims=1, reverse=true)) + @compute T2ry = Field(CumulativeIntegral(2 * T, dims=2, reverse=true)) + @compute T2rz = Field(CumulativeIntegral(2 * T, dims=3, reverse=true)) + + # T(x, y, z) = x + y + z + # 2 * T(0.5, 0.5, z) = [3, 5] + @test interior_array(T2cx, :, 1, 1) ≈ [3, 8] + @test interior_array(T2cy, 1, :, 1) ≈ [3, 8] + @test interior_array(T2cz, 1, 1, :) ≈ [3, 8] + + @test interior_array(T2rx, :, 1, 1) ≈ [8, 5] + @test interior_array(T2ry, 1, :, 1) ≈ [8, 5] + @test interior_array(T2rz, 1, 1, :) ≈ [8, 5] end # Test whether a race condition gets hit for averages over large fields