From e5a9c41ed45afd0c3bc1947efa6f8a924153cc73 Mon Sep 17 00:00:00 2001 From: Tomas Chor Date: Thu, 9 May 2024 22:18:08 -0700 Subject: [PATCH 1/2] Passes `grid` argument to `NetCDFOutputWriter` (#3576) * Update netcdf_output_writer.jl * Add grid as property to NetCDFOutputWriter and use it during file initialization * paragraph explaining that netcdf writer also accepts an output grid * change wording and add jldoctest * Update docs/src/model_setup/output_writers.md Co-authored-by: Gregory L. Wagner * Do a little better with function output for NetCDF output writer * Fixes interpolate! plus some improvements to NetCDFOutputWriter * add new example for docs * simplify example * simplify example in the docstring * Update docs/src/model_setup/output_writers.md Co-authored-by: Gregory L. Wagner * Make grid a kwarg and update language * Update src/OutputWriters/netcdf_output_writer.jl Co-authored-by: Tomas Chor * Update docs/src/model_setup/output_writers.md Co-authored-by: Tomas Chor * Update netcdf_output_writer.jl * Bugfix plus better error message * Correctly define_output_variable for LagrangianParticles * Fix znodes in output writers docs * fix znodes * fix double semicolon * 14.5 kB --------- Co-authored-by: Gregory Wagner Co-authored-by: iury simoes-sousa Co-authored-by: Iury Simoes-Sousa --- docs/src/model_setup/output_writers.md | 35 ++++- src/Fields/interpolate.jl | 2 +- src/OutputWriters/netcdf_output_writer.jl | 167 ++++++++++++++++------ test/test_output_writers.jl | 14 +- 4 files changed, 166 insertions(+), 52 deletions(-) diff --git a/docs/src/model_setup/output_writers.md b/docs/src/model_setup/output_writers.md index 8d1ee230b3..98f78cbb07 100644 --- a/docs/src/model_setup/output_writers.md +++ b/docs/src/model_setup/output_writers.md @@ -122,7 +122,7 @@ simulation = Simulation(model, Δt=1.25, stop_iteration=3) f(model) = model.clock.time^2; # scalar output -g(model) = model.clock.time .* exp.(znodes(Center, grid)) # single-column profile output (vector) +g(model) = model.clock.time .* exp.(znodes(grid, Center())) # single-column profile output (vector) xC, yF = xnodes(grid, Center()), ynodes(grid, Face()) @@ -158,6 +158,38 @@ NetCDFOutputWriter scheduled on IterationInterval(1): └── file size: 17.8 KiB ``` +`NetCDFOutputWriter` can also be configured for `outputs` that are interpolated or regridded +to a different grid than `model.grid`. To use this functionality, include the keyword argument +`grid = output_grid`. + +```jldoctest +using Oceananigans +using Oceananigans.Fields: interpolate! + +grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1)); +model = NonhydrostaticModel(; grid) + +coarse_grid = RectilinearGrid(size=(grid.Nx, grid.Ny, grid.Nz÷2), extent=(grid.Lx, grid.Ly, grid.Lz)) +coarse_u = Field{Face, Center, Center}(coarse_grid) + +interpolate_u(model) = interpolate!(coarse_u, model.velocities.u) +outputs = (; u = interpolate_u) + +output_writer = NetCDFOutputWriter(model, outputs; + grid = coarse_grid, + filename = "coarse_u.nc", + schedule = IterationInterval(1)) + +# output +NetCDFOutputWriter scheduled on IterationInterval(1): +├── filepath: ./coarse_u.nc +├── dimensions: zC(4), zF(5), xC(1), yF(1), xF(1), yC(1), time(0) +├── 1 outputs: u +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 14.5 KiB +``` + See [`NetCDFOutputWriter`](@ref) for more information. ## JLD2 output writer @@ -282,3 +314,4 @@ JLD2OutputWriter scheduled on TimeInterval(4 days): ├── file_splitting: NoFileSplitting └── file size: 26.5 KiB ``` + diff --git a/src/Fields/interpolate.jl b/src/Fields/interpolate.jl index f34fb73f3f..85a5cdf10b 100644 --- a/src/Fields/interpolate.jl +++ b/src/Fields/interpolate.jl @@ -354,5 +354,5 @@ function interpolate!(to_field::Field, from_field::AbstractField) fill_halo_regions!(to_field) - return nothing + return to_field end diff --git a/src/OutputWriters/netcdf_output_writer.jl b/src/OutputWriters/netcdf_output_writer.jl index 0fb468a117..56d328296b 100644 --- a/src/OutputWriters/netcdf_output_writer.jl +++ b/src/OutputWriters/netcdf_output_writer.jl @@ -10,7 +10,8 @@ using Oceananigans.Utils: versioninfo_with_gpu, oceananigans_versioninfo, pretty using Oceananigans.TimeSteppers: float_or_date_time using Oceananigans.Fields: reduced_dimensions, reduced_location, location, validate_indices -mutable struct NetCDFOutputWriter{D, O, T, A, FS} <: AbstractOutputWriter +mutable struct NetCDFOutputWriter{G, D, O, T, A, FS} <: AbstractOutputWriter + grid :: G filepath :: String dataset :: D outputs :: O @@ -162,29 +163,35 @@ function add_schedule_metadata!(global_attributes, schedule::AveragedTimeInterva end """ - NetCDFOutputWriter(model, outputs; filename, schedule - dir = ".", - array_type = Array{Float64}, - indices = nothing, - with_halos = false, - global_attributes = Dict(), - output_attributes = Dict(), - dimensions = Dict(), - overwrite_existing = false, - deflatelevel = 0, - part = 1, - file_splitting = NoFileSplitting(), - verbose = false) + NetCDFOutputWriter(model, outputs; filename, schedule, + grid = model.grid, + dir = ".", + array_type = Array{Float64}, + indices = nothing, + with_halos = false, + global_attributes = Dict(), + output_attributes = Dict(), + dimensions = Dict(), + overwrite_existing = false, + deflatelevel = 0, + part = 1, + file_splitting = NoFileSplitting(), + verbose = false) Construct a `NetCDFOutputWriter` that writes `(label, output)` pairs in `outputs` (which should be a `Dict`) to a NetCDF file, where `label` is a string that labels the output and `output` is either a `Field` (e.g. `model.velocities.u`) or a function `f(model)` that -returns something to be written to disk. Custom output requires the spatial `dimensions` (a -`Dict`) to be manually specified (see examples). +returns something to be written to disk. + +If any of `outputs` are not `AbstractField`, their spatial `dimensions` must be provided. + +To use `outputs` on a `grid` not equal to `model.grid`, provide the keyword argument `grid.` Keyword arguments ================= +- `grid`: The grid associated with `outputs`. Defaults to `model.grid`. + ## Filenaming - `filename` (required): Descriptive filename. `".nc"` is appended to `filename` if `filename` does @@ -319,7 +326,7 @@ simulation = Simulation(model, Δt=1.25, stop_iteration=3) f(model) = model.clock.time^2; # scalar output -g(model) = model.clock.time .* exp.(znodes(Center, grid)) # vector/profile output +g(model) = model.clock.time .* exp.(znodes(grid, Center())) # vector/profile output xC, yF = xnodes(grid, Center()), ynodes(grid, Face()) @@ -354,20 +361,55 @@ NetCDFOutputWriter scheduled on IterationInterval(1): ├── file_splitting: NoFileSplitting └── file size: 17.8 KiB ``` + +`NetCDFOutputWriter` can also be configured for `outputs` that are interpolated or regridded +to a different grid than `model.grid`. To use this functionality, include the keyword argument +`grid = output_grid`. + +```jldoctest +using Oceananigans +using Oceananigans.Fields: interpolate! + +grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1)); +model = NonhydrostaticModel(; grid) + +coarse_grid = RectilinearGrid(size=(grid.Nx, grid.Ny, grid.Nz÷2), extent=(grid.Lx, grid.Ly, grid.Lz)) +coarse_u = Field{Face, Center, Center}(coarse_grid) + +interpolate_u(model) = interpolate!(coarse_u, model.velocities.u) +outputs = (; u = interpolate_u) + +output_writer = NetCDFOutputWriter(model, outputs; + grid = coarse_grid, + filename = "coarse_u.nc", + schedule = IterationInterval(1)) + +# output +NetCDFOutputWriter scheduled on IterationInterval(1): +├── filepath: ./coarse_u.nc +├── dimensions: zC(4), zF(5), xC(1), yF(1), xF(1), yC(1), time(0) +├── 1 outputs: u +└── array type: Array{Float64} +├── file_splitting: NoFileSplitting +└── file size: 14.5 KiB +``` """ -function NetCDFOutputWriter(model, outputs; filename, schedule, - dir = ".", - array_type = Array{Float64}, - indices = (:, :, :), - with_halos = false, +function NetCDFOutputWriter(model, outputs; + filename, + schedule, + grid = model.grid, + dir = ".", + array_type = Array{Float64}, + indices = (:, :, :), + with_halos = false, global_attributes = Dict(), output_attributes = Dict(), - dimensions = Dict(), - overwrite_existing = nothing, - deflatelevel = 0, - part = 1, - file_splitting = NoFileSplitting(), - verbose = false) + dimensions = Dict(), + overwrite_existing = nothing, + deflatelevel = 0, + part = 1, + file_splitting = NoFileSplitting(), + verbose = false) mkpath(dir) filename = auto_extension(filename, ".nc") filepath = joinpath(dir, filename) @@ -394,7 +436,7 @@ function NetCDFOutputWriter(model, outputs; filename, schedule, # with LagrangianParticles output (see the end of the file). # We shouldn't support this in the future; we should require users to 'name' LagrangianParticles output. outputs = dictify(outputs) - outputs = Dict(string(name) => construct_output(outputs[name], model.grid, indices, with_halos) for name in keys(outputs)) + outputs = Dict(string(name) => construct_output(outputs[name], grid, indices, with_halos) for name in keys(outputs)) output_attributes = dictify(output_attributes) global_attributes = dictify(global_attributes) @@ -414,9 +456,11 @@ function NetCDFOutputWriter(model, outputs; filename, schedule, dimensions, overwrite_existing, deflatelevel, + grid, model) - return NetCDFOutputWriter(filepath, + return NetCDFOutputWriter(grid, + filepath, dataset, outputs, schedule, @@ -446,28 +490,44 @@ get_default_dimension_attributes(grid::ImmersedBoundaryGrid) = ##### Variable definition ##### +materialize_output(func, model) = func(model) +materialize_output(field::AbstractField, model) = field +materialize_output(particles::LagrangianParticles, model) = particles +materialize_output(output::WindowedTimeAverage{<:AbstractField}, model) = output + """ Defines empty variables for 'custom' user-supplied `output`. """ -function define_output_variable!(dataset, output, name, array_type, deflatelevel, output_attributes, dimensions) - name ∉ keys(dimensions) && error("Custom output $name needs dimensions!") +function define_output_variable!(dataset, output, name, array_type, + deflatelevel, attrib, dimensions, filepath) - defVar(dataset, name, eltype(array_type), (dimensions[name]..., "time"), - deflatelevel=deflatelevel, attrib=output_attributes) + if name ∉ keys(dimensions) + msg = string("dimensions[$name] for output $name=", typeof(output), " into ", filepath, '\n', + " must be provided when constructing NetCDFOutputWriter") + throw(ArgumentError(msg)) + end + + dims = dimensions[name] + FT = eltype(array_type) + defVar(dataset, name, FT, (dims..., "time"); deflatelevel, attrib) return nothing end """ Defines empty field variable. """ -define_output_variable!(dataset, output::AbstractField, name, array_type, deflatelevel, output_attributes, dimensions) = - defVar(dataset, name, eltype(array_type), - (netcdf_spatial_dimensions(output)..., "time"), - deflatelevel=deflatelevel, attrib=output_attributes) +function define_output_variable!(dataset, output::AbstractField, name, array_type, + deflatelevel, attrib, dimensions, filepath) + + dims = netcdf_spatial_dimensions(output) + FT = eltype(array_type) + defVar(dataset, name, FT, (dims..., "time"); deflatelevel, attrib) + + return nothing +end """ Defines empty field variable for `WindowedTimeAverage`s over fields. """ define_output_variable!(dataset, output::WindowedTimeAverage{<:AbstractField}, args...) = define_output_variable!(dataset, output.operand, args...) - ##### ##### Write output ##### @@ -580,12 +640,17 @@ end ##### """ Defines empty variable for particle trackting. """ -function define_output_variable!(dataset, output::LagrangianParticles, name, array_type, deflatelevel, output_attributes, dimensions) +function define_output_variable!(dataset, output::LagrangianParticles, name, array_type, + deflatelevel, output_attributes, dimensions, filepath) + particle_fields = eltype(output.properties) |> fieldnames .|> string + T = eltype(array_type) + for particle_field in particle_fields - defVar(dataset, particle_field, eltype(array_type), - ("particle_id", "time"), deflatelevel=deflatelevel) + defVar(dataset, particle_field, T, ("particle_id", "time"); deflatelevel) end + + return nothing end dictify(outputs::LagrangianParticles) = Dict("particles" => outputs) @@ -633,6 +698,7 @@ function initialize_nc_file!(filepath, dimensions, overwrite_existing, deflatelevel, + grid, model) mode = overwrite_existing ? "c" : "a" @@ -648,12 +714,12 @@ function initialize_nc_file!(filepath, # schedule::AveragedTimeInterval schedule, outputs = time_average_outputs(schedule, outputs, model) - dims = default_dimensions(outputs, model.grid, indices, with_halos) + dims = default_dimensions(outputs, grid, indices, with_halos) # Open the NetCDF dataset file dataset = NCDataset(filepath, mode, attrib=global_attributes) - default_dimension_attributes = get_default_dimension_attributes(model.grid) + default_dimension_attributes = get_default_dimension_attributes(grid) # Define variables for each dimension and attributes if this is a new file. if mode == "c" @@ -669,7 +735,7 @@ function initialize_nc_file!(filepath, # Creates an unlimited dimension "time" defDim(dataset, "time", Inf) - defVar(dataset, "time", eltype(model.grid), ("time",), attrib=time_attrib) + defVar(dataset, "time", eltype(grid), ("time",), attrib=time_attrib) # Use default output attributes for known outputs if the user has not specified any. # Unknown outputs get an empty tuple (no output attributes). @@ -681,7 +747,15 @@ function initialize_nc_file!(filepath, for (name, output) in outputs attributes = try output_attributes[name]; catch; Dict(); end - define_output_variable!(dataset, output, name, array_type, deflatelevel, attributes, dimensions) + materialized = materialize_output(output, model) + define_output_variable!(dataset, + materialized, + name, + array_type, + deflatelevel, + attributes, + dimensions, + filepath) # for better error messages end sync(dataset) @@ -704,4 +778,5 @@ initialize_nc_file!(ow::NetCDFOutputWriter, model) = ow.dimensions, ow.overwrite_existing, ow.deflatelevel, + ow.grid, model) diff --git a/test/test_output_writers.jl b/test/test_output_writers.jl index bd9f56efb3..0b4ee23ee6 100644 --- a/test/test_output_writers.jl +++ b/test/test_output_writers.jl @@ -68,7 +68,11 @@ function test_dependency_adding(model) dimensions = Dict("time_average" => ("xF", "yC", "zC")) # JLD2 dependencies test - jld2_output_writer = JLD2OutputWriter(model, output, schedule=TimeInterval(4), dir=".", filename="test.jld2", overwrite_existing=true) + jld2_output_writer = JLD2OutputWriter(model, output, + schedule = TimeInterval(4), + dir = ".", + filename = "test.jld2", + overwrite_existing = true) windowed_time_average = jld2_output_writer.outputs.time_average @test dependencies_added_correctly!(model, windowed_time_average, jld2_output_writer) @@ -148,7 +152,8 @@ function test_windowed_time_averaging_simulation(model) # https://github.com/Alexander-Barth/NCDatasets.jl/issues/105 nc_filepath1 = "windowed_time_average_test1.nc" nc_outputs = Dict(string(name) => field for (name, field) in pairs(model.velocities)) - nc_output_writer = NetCDFOutputWriter(model, nc_outputs, filename=nc_filepath1, + nc_output_writer = NetCDFOutputWriter(model, nc_outputs, + filename = nc_filepath1, schedule = AveragedTimeInterval(π, window=1)) jld2_outputs_are_time_averaged = Tuple(typeof(out) <: WindowedTimeAverage for out in jld2_output_writer.outputs) @@ -202,8 +207,9 @@ function test_windowed_time_averaging_simulation(model) nc_filepath2 = "windowed_time_average_test2.nc" nc_outputs = Dict(string(name) => field for (name, field) in pairs(model.velocities)) - simulation.output_writers[:nc] = NetCDFOutputWriter(model, nc_outputs, filename=nc_filepath2, - schedule=AveragedTimeInterval(π, window=π)) + simulation.output_writers[:nc] = NetCDFOutputWriter(model, nc_outputs, + filename = nc_filepath2, + schedule = AveragedTimeInterval(π, window=π)) run!(simulation) From c91df929ac50b5200b72118424db0ed3478faa9d Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Fri, 10 May 2024 06:39:42 -0600 Subject: [PATCH 2/2] 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 +- .../metric_field_reductions.jl | 121 +++++-- src/Fields/Fields.jl | 4 +- src/Fields/field.jl | 2 - src/Fields/field_reductions.jl | 97 ------ src/Fields/scans.jl | 307 ++++++++++++++++++ src/Oceananigans.jl | 2 +- test/runtests.jl | 2 +- ...ield_reductions.jl => test_field_scans.jl} | 142 +++++++- 9 files changed, 525 insertions(+), 154 deletions(-) delete mode 100644 src/Fields/field_reductions.jl create mode 100644 src/Fields/scans.jl rename test/{test_field_reductions.jl => test_field_scans.jl} (61%) 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