diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index 47ac657a6a..7696712e21 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/src/AbstractOperations/AbstractOperations.jl @@ -16,6 +16,7 @@ using Oceananigans.Fields using Oceananigans.Utils using Oceananigans: location +using Oceananigans.Fields: instantiated_location using Oceananigans.Operators: interpolation_operator import Adapt diff --git a/src/AbstractOperations/binary_operations.jl b/src/AbstractOperations/binary_operations.jl index 41b869e76f..38ef095e96 100644 --- a/src/AbstractOperations/binary_operations.jl +++ b/src/AbstractOperations/binary_operations.jl @@ -35,21 +35,21 @@ indices(β::BinaryOperation) = construct_regionally(intersect_indices, location( """Create a binary operation for `op` acting on `a` and `b` at `Lc`, where `a` and `b` have location `La` and `Lb`.""" -function _binary_operation(Lc, op, a, b, La, Lb, grid) +function _binary_operation(Lc::Tuple{LX, LY, LZ}, op, a, b, La, Lb, grid) where {LX, LY, LZ} ▶a = interpolation_operator(La, Lc) ▶b = interpolation_operator(Lb, Lc) - return BinaryOperation{Lc[1], Lc[2], Lc[3]}(op, a, b, ▶a, ▶b, grid) + return BinaryOperation{LX, LY, LZ}(op, a, b, ▶a, ▶b, grid) end -const ConcreteLocationType = Union{Type{Face}, Type{Center}} +const ConcreteLocationType = Union{Face, Center} # Precedence rules for choosing operation location: -choose_location(La, Lb, Lc) = Lc # Fallback to the specification Lc, but also... -choose_location(::Type{Face}, ::Type{Face}, Lc) = Face # keep common locations; and -choose_location(::Type{Center}, ::Type{Center}, Lc) = Center # -choose_location(La::ConcreteLocationType, ::Type{Nothing}, Lc) = La # don't interpolate unspecified locations. -choose_location(::Type{Nothing}, Lb::ConcreteLocationType, Lc) = Lb # +choose_location(La, Lb, Lc) = Lc # Fallback to the specification Lc, but also... +choose_location(::Face, ::Face, Lc) = Face() # keep common locations; and +choose_location(::Center, ::Center, Lc) = Center() # +choose_location(La::ConcreteLocationType, ::Nothing, Lc) = La # don't interpolate unspecified locations. +choose_location(::Nothing, Lb::ConcreteLocationType, Lc) = Lb # # Apply the function if the inputs are scalars, otherwise broadcast it over the inputs # This can occur in the binary operator code if we index into with an array, e.g. array[1:10] @@ -107,8 +107,8 @@ function define_binary_operator(op) if that is also Nothing, `Lc`. """ function $op(Lc::Tuple, a, b) - La = location(a) - Lb = location(b) + La = instantiated_location(a) + Lb = instantiated_location(b) Lab = choose_location.(La, Lb, Lc) grid = Oceananigans.AbstractOperations.validate_grid(a, b) @@ -123,19 +123,19 @@ function define_binary_operator(op) $op(Lc::Tuple, f::Function, b::AbstractField) = $op(Lc, FunctionField(location(b), f, b.grid), b) $op(Lc::Tuple, a::AbstractField, f::Function) = $op(Lc, a, FunctionField(location(a), f, a.grid)) - $op(Lc::Tuple, m::GridMetric, b::AbstractField) = $op(Lc, grid_metric_operation(location(b), m, b.grid), b) - $op(Lc::Tuple, a::AbstractField, m::GridMetric) = $op(Lc, a, grid_metric_operation(location(a), m, a.grid)) + $op(Lc::Tuple, m::GridMetric, b::AbstractField) = $op(Lc, grid_metric_operation(instantiated_location(b), m, b.grid), b) + $op(Lc::Tuple, a::AbstractField, m::GridMetric) = $op(Lc, a, grid_metric_operation(instantiated_location(a), m, a.grid)) # Sugary versions with default locations - $op(a::AF, b::AF) = $op(location(a), a, b) - $op(a::AF, b) = $op(location(a), a, b) - $op(a, b::AF) = $op(location(b), a, b) + $op(a::AF, b::AF) = $op(instantiated_location(a), a, b) + $op(a::AF, b) = $op(instantiated_location(a), a, b) + $op(a, b::AF) = $op(instantiated_location(b), a, b) - $op(a::AF, b::Number) = $op(location(a), a, b) - $op(a::Number, b::AF) = $op(location(b), a, b) + $op(a::AF, b::Number) = $op(instantiated_location(a), a, b) + $op(a::Number, b::AF) = $op(instantiated_location(b), a, b) - $op(a::AF, b::ConstantField) = $op(location(a), a, b.constant) - $op(a::ConstantField, b::AF) = $op(location(b), a.constant, b) + $op(a::AF, b::ConstantField) = $op(instantiated_location(a), a, b.constant) + $op(a::ConstantField, b::AF) = $op(instantiated_location(b), a.constant, b) $op(a::Number, b::ConstantField) = ConstantField($op(a, b.constant)) $op(a::ConstantField, b::Number) = ConstantField($op(a.constant, b)) diff --git a/src/AbstractOperations/derivatives.jl b/src/AbstractOperations/derivatives.jl index 64dddc356b..ef6b04c6bc 100644 --- a/src/AbstractOperations/derivatives.jl +++ b/src/AbstractOperations/derivatives.jl @@ -28,9 +28,9 @@ end """Create a derivative operator `∂` acting on `arg` at `L∂`, followed by interpolation to `L` on `grid`.""" -function _derivative(L, ∂, arg, L∂, abstract_∂, grid) +function _derivative(L::Tuple{LX, LY, LZ}, ∂, arg, L∂, abstract_∂, grid) where {LX, LY, LZ} ▶ = interpolation_operator(L∂, L) - return Derivative{L[1], L[2], L[3]}(∂, arg, ▶, abstract_∂, grid) + return Derivative{LX, LY, LZ}(∂, arg, ▶, abstract_∂, grid) end indices(d::Derivative) = indices(d.arg) @@ -41,6 +41,8 @@ indices(d::Derivative) = indices(d.arg) """Return `Center` if given `Face` or `Face` if given `Center`.""" flip(::Type{Face}) = Center flip(::Type{Center}) = Face +flip(::Face) = Center() +flip(::Center) = Face() const LocationType = Union{Type{Face}, Type{Center}, Type{Nothing}} @@ -63,7 +65,7 @@ Return an abstract representation of an ``x``-derivative acting on field `arg` f by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. """ ∂x(L::Tuple, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = - _derivative(L, ∂x(LX, LY, LZ), arg, (flip(LX), LY, LZ), ∂x, arg.grid) + _derivative(L, ∂x(LX(), LY(), LZ()), arg, (flip(LX()), LY(), LZ()), ∂x, arg.grid) """ ∂y(L::Tuple, arg::AbstractField) @@ -72,7 +74,7 @@ Return an abstract representation of a ``y``-derivative acting on field `arg` fo by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. """ ∂y(L::Tuple, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = - _derivative(L, ∂y(LX, LY, LZ), arg, (LX, flip(LY), LZ), ∂y, arg.grid) + _derivative(L, ∂y(LX(), LY(), LZ()), arg, (LX(), flip(LY()), LZ()), ∂y, arg.grid) """ ∂z(L::Tuple, arg::AbstractField) @@ -81,7 +83,7 @@ Return an abstract representation of a ``z``-derivative acting on field `arg` fo by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. """ ∂z(L::Tuple, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = - _derivative(L, ∂z(LX, LY, LZ), arg, (LX, LY, flip(LZ)), ∂z, arg.grid) + _derivative(L, ∂z(LX(), LY(), LZ()), arg, (LX(), LY(), flip(LZ())), ∂z, arg.grid) # Defaults @@ -90,21 +92,21 @@ by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s. Return an abstract representation of a ``x``-derivative acting on field `arg`. """ -∂x(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂x((flip(LX), LY, LZ), arg) +∂x(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂x((flip(LX()), LY(), LZ()), arg) """ ∂y(arg::AbstractField) Return an abstract representation of a ``y``-derivative acting on field `arg`. """ -∂y(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂y((LX, flip(LY), LZ), arg) +∂y(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂y((LX(), flip(LY()), LZ()), arg) """ ∂z(arg::AbstractField) Return an abstract representation of a ``z``-derivative acting on field `arg`. """ -∂z(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂z((LX, LY, flip(LZ)), arg) +∂z(arg::AF{LX, LY, LZ}) where {LX, LY, LZ} = ∂z((LX(), LY(), flip(LZ())), arg) ##### ##### Nested computations diff --git a/src/AbstractOperations/grid_metrics.jl b/src/AbstractOperations/grid_metrics.jl index 0dc32a59e8..4a60255cf7 100644 --- a/src/AbstractOperations/grid_metrics.jl +++ b/src/AbstractOperations/grid_metrics.jl @@ -79,14 +79,13 @@ julia> c_dz[1, 1, 1] 3.0 ``` """ -grid_metric_operation(loc, metric, grid) = - KernelFunctionOperation{loc[1], loc[2], loc[3]}(metric_function(loc, metric), grid) +grid_metric_operation(loc::Tuple{LX, LY, LZ}, metric, grid) where {LX, LY, LZ} = + KernelFunctionOperation{LX, LY, LZ}(metric_function(loc, metric), grid) const NodeMetric = Union{XNode, YNode, ZNode, ΛNode, ΦNode, RNode} -function grid_metric_operation(loc, metric::NodeMetric, grid) - LX, LY, LZ = loc - ℓx, ℓy, ℓz = LX(), LY(), LZ() +function grid_metric_operation(loc::Tuple{LX, LY, LZ}, metric::NodeMetric, grid) where {LX, LY, LZ} + ℓx, ℓy, ℓz = loc ξnode = metric_function(loc, metric) return KernelFunctionOperation{LX, LY, LZ}(ξnode, grid, ℓx, ℓy, ℓz) end diff --git a/src/AbstractOperations/multiary_operations.jl b/src/AbstractOperations/multiary_operations.jl index dc9c41d972..eb9dbac5e4 100644 --- a/src/AbstractOperations/multiary_operations.jl +++ b/src/AbstractOperations/multiary_operations.jl @@ -22,9 +22,9 @@ end indices(Π::MultiaryOperation) = construct_regionally(intersect_indices, location(Π), Π.args...) -function _multiary_operation(L, op, args, Largs, grid) +function _multiary_operation(L::Tuple{LX, LY, LZ}, op, args, Largs, grid) where {LX, LY, LZ} ▶ = Tuple(interpolation_operator(La, L) for La in Largs) - return MultiaryOperation{L[1], L[2], L[3]}(op, Tuple(a for a in args), ▶, grid) + return MultiaryOperation{LX, LY, LZ}(op, Tuple(a for a in args), ▶, grid) end # Recompute location of multiary operation @@ -52,7 +52,7 @@ function define_multiary_operator(op) $op(a::Oceananigans.Fields.AbstractField, b::Union{Function, Oceananigans.Fields.AbstractField}, c::Union{Function, Oceananigans.Fields.AbstractField}, - d::Union{Function, Oceananigans.Fields.AbstractField}...) = $op(Oceananigans.Fields.location(a), a, b, c, d...) + d::Union{Function, Oceananigans.Fields.AbstractField}...) = $op(Oceananigans.Fields.instantiated_location(a), a, b, c, d...) end end diff --git a/src/AbstractOperations/unary_operations.jl b/src/AbstractOperations/unary_operations.jl index d6c05d058e..8831558428 100644 --- a/src/AbstractOperations/unary_operations.jl +++ b/src/AbstractOperations/unary_operations.jl @@ -28,9 +28,9 @@ indices(υ::UnaryOperation) = indices(υ.arg) """Create a unary operation for `operator` acting on `arg` which interpolates the result from `Larg` to `L`.""" -function _unary_operation(L, operator, arg, Larg, grid) +function _unary_operation(L::Tuple{LX, LY, LZ}, operator, arg, Larg, grid) where {LX, LY, LZ} ▶ = interpolation_operator(Larg, L) - return UnaryOperation{L[1], L[2], L[3]}(operator, arg, ▶, grid) + return UnaryOperation{LX, LY, LZ}(operator, arg, ▶, grid) end # Recompute location of unary operation @@ -91,7 +91,7 @@ macro unary(ops...) import Oceananigans.Grids: AbstractGrid import Oceananigans.Fields: AbstractField - local location = Oceananigans.Fields.location + local instantiated_location = Oceananigans.Fields.instantiated_location @inline $op(i, j, k, grid::AbstractGrid, a) = @inbounds $op(a[i, j, k]) @inline $op(i, j, k, grid::AbstractGrid, a::Number) = $op(a) @@ -103,11 +103,11 @@ macro unary(ops...) `a`, and subsequently interpolated to the location indicated by `Lop`. """ function $op(Lop::Tuple, a::AbstractField) - L = location(a) + L = instantiated_location(a) return Oceananigans.AbstractOperations._unary_operation(Lop, $op, a, L, a.grid) end - $op(a::AbstractField) = $op(location(a), a) + $op(a::AbstractField) = $op(instantiated_location(a), a) push!(Oceananigans.AbstractOperations.operators, Symbol($op)) push!(Oceananigans.AbstractOperations.unary_operators, Symbol($op)) diff --git a/src/Fields/scans.jl b/src/Fields/scans.jl index d31d557a02..7dad4cc577 100644 --- a/src/Fields/scans.jl +++ b/src/Fields/scans.jl @@ -34,8 +34,12 @@ 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 +@inline location(s::Scan) = location(s.operand) +@inline instantiated_location(s::Scan) = instantiated_location(s.operand) + +scan_indices(::AbstractReducing, indices, dims) = Tuple(i ∈ dims ? Colon() : indices[i] for i in 1:3) +scan_indices(::AbstractAccumulating, indices, dims) = indices +scan_indices(::AbstractReducing, ::Tuple{Colon, Colon, Colon}, dims) = (:, :, :) Base.summary(s::Scan) = string(summary(s.type), " ", s.scan!, @@ -52,7 +56,7 @@ function Field(scan::Scan; grid = operand.grid LX, LY, LZ = loc = instantiated_location(scan) dims = filter_nothing_dims(scan.dims, loc) - indices = scan_indices(scan.type, indices; dims) + indices = scan_indices(scan.type, indices, dims) if isnothing(data) data = new_data(grid, loc, indices) diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_vertically_integrated_variables.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_vertically_integrated_variables.jl index d9bfe26e3b..d3bf08f4f2 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_vertically_integrated_variables.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_vertically_integrated_variables.jl @@ -10,8 +10,8 @@ function compute_vertically_integrated_lateral_areas!(∫ᶻ_A) field_grid = ∫ᶻ_A.xᶠᶜᶜ.grid - Axᶠᶜᶜ = grid_metric_operation((Face, Center, Center), Ax, field_grid) - Ayᶜᶠᶜ = grid_metric_operation((Center, Face, Center), Ay, field_grid) + Axᶠᶜᶜ = grid_metric_operation((Face(), Center(), Center()), Ax, field_grid) + Ayᶜᶠᶜ = grid_metric_operation((Center(), Face(), Center()), Ay, field_grid) sum!(∫ᶻ_A.xᶠᶜᶜ, Axᶠᶜᶜ) sum!(∫ᶻ_A.yᶜᶠᶜ, Ayᶜᶠᶜ)