Skip to content

Commit

Permalink
Merge branch 'main' into glw/last-dt
Browse files Browse the repository at this point in the history
  • Loading branch information
glwagner committed May 13, 2024
2 parents 845bc07 + c91df92 commit 0254952
Show file tree
Hide file tree
Showing 13 changed files with 691 additions and 206 deletions.
35 changes: 34 additions & 1 deletion docs/src/model_setup/output_writers.md
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -282,3 +314,4 @@ JLD2OutputWriter scheduled on TimeInterval(4 days):
├── file_splitting: NoFileSplitting
└── file size: 26.5 KiB
```

2 changes: 1 addition & 1 deletion src/AbstractOperations/AbstractOperations.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
121 changes: 88 additions & 33 deletions src/AbstractOperations/metric_field_reductions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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`.
Expand Down Expand Up @@ -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
Expand All @@ -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

4 changes: 2 additions & 2 deletions src/Fields/Fields.jl
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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")
Expand Down
2 changes: 0 additions & 2 deletions src/Fields/field.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
97 changes: 0 additions & 97 deletions src/Fields/field_reductions.jl

This file was deleted.

2 changes: 1 addition & 1 deletion src/Fields/interpolate.jl
Original file line number Diff line number Diff line change
Expand Up @@ -354,5 +354,5 @@ function interpolate!(to_field::Field, from_field::AbstractField)

fill_halo_regions!(to_field)

return nothing
return to_field
end
Loading

0 comments on commit 0254952

Please sign in to comment.