Fix type instability in AbstractOperations

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

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))

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

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

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
 

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))

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)

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ᶜᶠᶜ)