Fix type inference failure in AbstractOperations

docs/src/operations.md

@@ -53,17 +53,16 @@ and even by chaining expressions together, which may themselves include `Abstrac
 quadratic = c^2 + two_c + 1
 
 # output
-BinaryOperation at (Center, Center, Center)
+MultiaryOperation at (Center, Center, Center)
 ├── grid: 4×1×4 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
 └── tree:
     + at (Center, Center, Center)
-    ├── + at (Center, Center, Center)
-    │   ├── ^ at (Center, Center, Center)
-    │   │   ├── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
-    │   │   └── 2
-    │   └── * at (Center, Center, Center)
-        │   ├── 2
-        │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
+    ├── ^ at (Center, Center, Center)
+    │   ├── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
+    │   └── 2
+    ├── * at (Center, Center, Center)
+    │   ├── 2
+    │   └── 4×1×4 Field{Center, Center, Center} on RectilinearGrid on CPU
     └── 1
 ```
 

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
@@ -32,6 +33,8 @@ abstract type AbstractOperation{LX, LY, LZ, G, T} <: AbstractField{LX, LY, LZ, G
 
 const AF = AbstractField # used in unary_operations.jl, binary_operations.jl, etc
 
+const Location = Union{Face, Center, Nothing}
+
 # We have no halos to fill
 @inline fill_halo_regions!(::AbstractOperation, args...; kwargs...) = nothing
 

src/AbstractOperations/at.jl

@@ -22,6 +22,20 @@ function insert_location!(ex::Expr, location)
     return nothing
 end
 
+# Transform expressions like :((LX, LY, LZ)) into :((LX(), LY(), LZ()))
+# when LX, LY, LZ are Face, Center, or Nothing. Leave eveything else unchanged.
+function instantiate_location_expression(exp::Expr)
+    new_exp = deepcopy(exp)
+    for (i, arg) in enumerate(new_exp.args)
+        new_exp.args[i] = instantiate_location_expression(arg) 
+    end
+    return new_exp
+end
+
+instantiate_location_expression(arg) = arg == :Center  ? :(Center()) :
+                                       arg == :Face    ? :(Face()  ) :
+                                       arg == :Nothing ? :(nothing ) : arg
+
 "Fallback for when `insert_location` is called on objects other than expressions."
 insert_location!(anything, location) = nothing
 
@@ -32,7 +46,9 @@ insert_location!(anything, location) = nothing
 interpolate_operation(L, x) = x
 
 function interpolate_operation(L, x::AbstractField)
-    L == location(x) && return x # Don't interpolate unnecessarily
+    if L == instantiated_location(x) || L == location(x)
+        return x # Don't interpolate unnecessarily
+    end
     return interpolate_identity(L, x)
 end
 
@@ -43,6 +59,7 @@ Modify the `abstract_operation` so that it returns values at
 `location`, where `location` is a 3-tuple of `Face`s and `Center`s.
 """
 macro at(location, abstract_operation)
+    location = instantiate_location_expression(location)
     insert_location!(abstract_operation, location)
 
     # We wrap it all in an interpolator to help "stubborn" binary operations

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)
-     ▶a = interpolation_operator(La, Lc)
-     ▶b = interpolation_operator(Lb, Lc)
+function _binary_operation(Lc::Tuple{LX, LY, LZ}, op, a, b, La, Lb, grid) where {LX<:Location, LY<:Location, LZ<:Location}
+    ▶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]
@@ -106,9 +106,9 @@ function define_binary_operator(op)
         the location of the dimension in question is supplied either by `location(b)` or
         if that is also Nothing, `Lc`.
         """
-        function $op(Lc::Tuple, a, b)
-            La = location(a)
-            Lb = location(b)
+        function $op(Lc::Tuple{<:$Location, <:$Location, <:$Location}, a, b)
+            La = Oceananigans.Fields.instantiated_location(a)
+            Lb = Oceananigans.Fields.instantiated_location(b)
             Lab = choose_location.(La, Lb, Lc)
 
             grid = Oceananigans.AbstractOperations.validate_grid(a, b)
@@ -120,22 +120,31 @@ function define_binary_operator(op)
         $op(Lc::Tuple, a::Number, b::Number) = $op(a, b)
 
         # Sugar for mixing in functions of (x, y, z)
-        $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{<:$Location, <:$Location, <:$Location}, f::Function, b::AbstractField) = $op(Lc, FunctionField(location(b), f, b.grid), b)
+        $op(Lc::Tuple{<:$Location, <:$Location, <:$Location}, 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{<:$Location, <:$Location, <:$Location}, m::GridMetric, b::AbstractField) = $op(Lc, grid_metric_operation(Oceananigans.Fields.instantiated_location(b), m, b.grid), b)
+        $op(Lc::Tuple{<:$Location, <:$Location, <:$Location}, a::AbstractField, m::GridMetric) = $op(Lc, a, grid_metric_operation(Oceananigans.Fields.instantiated_location(a), m, a.grid))
+
+        # instantiate location if types are passed
+        $op(Lc::Tuple, a, b) = $op((Lc[1](), Lc[2](), Lc[3]()), a, b)
+
+        $op(Lc::Tuple, f::Function, b::AbstractField) = $op((Lc[1](), Lc[2](), Lc[3]()), a, b)
+        $op(Lc::Tuple, a::AbstractField, f::Function) = $op((Lc[1](), Lc[2](), Lc[3]()), a, b)
+
+        $op(Lc::Tuple, m::GridMetric, b::AbstractField) = $op((Lc[1](), Lc[2](), Lc[3]()), a, b)
+        $op(Lc::Tuple, a::AbstractField, m::GridMetric) = $op((Lc[1](), Lc[2](), Lc[3]()), a, b)
 
         # 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(Oceananigans.Fields.instantiated_location(a), a, b)
+        $op(a::AF, b) = $op(Oceananigans.Fields.instantiated_location(a), a, b)
+        $op(a, b::AF) = $op(Oceananigans.Fields.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(Oceananigans.Fields.instantiated_location(a), a, b)
+        $op(a::Number, b::AF) = $op(Oceananigans.Fields.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(Oceananigans.Fields.instantiated_location(a), a, b.constant)
+        $op(a::ConstantField, b::AF) = $op(Oceananigans.Fields.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,17 +41,17 @@ indices(d::Derivative) = indices(d.arg)
 """Return `Center` if given `Face` or `Face` if given `Center`."""
 flip(::Type{Face}) = Center
 flip(::Type{Center}) = Face
-
-const LocationType = Union{Type{Face}, Type{Center}, Type{Nothing}}
+flip(::Face) = Center()
+flip(::Center) = Face()
 
 """Return the ``x``-derivative function acting at (`X`, `Y`, `Any`)."""
-∂x(X::LocationType, Y::LocationType, Z::LocationType) = eval(Symbol(:∂x, interpolation_code(flip(X)), interpolation_code(Y), interpolation_code(Z)))
+∂x(X::Location, Y::Location, Z::Location) = eval(Symbol(:∂x, interpolation_code(flip(X)), interpolation_code(Y), interpolation_code(Z)))
 
 """Return the ``y``-derivative function acting at (`X`, `Y`, `Any`)."""
-∂y(X::LocationType, Y::LocationType, Z::LocationType) = eval(Symbol(:∂y, interpolation_code(X), interpolation_code(flip(Y)), interpolation_code(Z)))
+∂y(X::Location, Y::Location, Z::Location) = eval(Symbol(:∂y, interpolation_code(X), interpolation_code(flip(Y)), interpolation_code(Z)))
 
 """Return the ``z``-derivative function acting at (`Any`, `Any`, `Z`)."""
-∂z(X::LocationType, Y::LocationType, Z::LocationType) = eval(Symbol(:∂z, interpolation_code(X), interpolation_code(Y), interpolation_code(flip(Z))))
+∂z(X::Location, Y::Location, Z::Location) = eval(Symbol(:∂z, interpolation_code(X), interpolation_code(Y), interpolation_code(flip(Z))))
 
 const derivative_operators = Set([:∂x, :∂y, :∂z])
 push!(operators, derivative_operators...)
@@ -60,28 +60,33 @@ push!(operators, derivative_operators...)
     ∂x(L::Tuple, arg::AbstractField)
 
 Return an abstract representation of an ``x``-derivative acting on field `arg` followed
-by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s.
+by interpolation to `L`, where `L` is a 3-tuple of instantiated `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)
+∂x(L::Tuple{<:Location, <:Location, <:Location}, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} =
+    _derivative(L, ∂x(LX(), LY(), LZ()), arg, (flip(LX()), LY(), LZ()), ∂x, arg.grid)
 
 """
     ∂y(L::Tuple, arg::AbstractField)
 
 Return an abstract representation of a ``y``-derivative acting on field `arg` followed
-by interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s.
+by interpolation to `L`, where `L` is a 3-tuple of instantiated `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)
+∂y(L::Tuple{<:Location, <:Location, <:Location}, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} =
+    _derivative(L, ∂y(LX(), LY(), LZ()), arg, (LX(), flip(LY()), LZ()), ∂y, arg.grid)
 
 """
     ∂z(L::Tuple, arg::AbstractField)
 
 Return an abstract representation of a ``z``-derivative acting on field `arg` followed
-by  interpolation to `L`, where `L` is a 3-tuple of `Face`s and `Center`s.
+by  interpolation to `L`, where `L` is a 3-tuple of instantiated `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)
+∂z(L::Tuple{<:Location, <:Location, <:Location}, arg::AF{LX, LY, LZ}) where {LX, LY, LZ} =
+    _derivative(L, ∂z(LX(), LY(), LZ()), arg, (LX(), LY(), flip(LZ())), ∂z, arg.grid)
+
+# Instantiate location if types are passed
+∂x(L::Tuple, arg::AF) = ∂x((L[1](), L[2](), L[3]()), arg)
+∂y(L::Tuple, arg::AF) = ∂y((L[1](), L[2](), L[3]()), arg)  
+∂z(L::Tuple, arg::AF) = ∂z((L[1](), L[2](), L[3]()), arg)   
 
 # Defaults
 
@@ -90,21 +95,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,16 @@ 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<:Location, LY<:Location, LZ<:Location} =
+    KernelFunctionOperation{LX, LY, LZ}(metric_function(loc, metric), grid)
+
+# Instantiated location if location types are passed as values
+grid_metric_operation(Loc::Tuple, metric, grid) = grid_metric_operation((Loc[1](), Loc[2](), Loc[3]()), 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<:Location, LY<:Location, LZ<:Location} 
+    ℓ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
@@ -33,7 +33,7 @@ end
 """Return an expression that defines an abstract `MultiaryOperator` named `op` for `AbstractField`."""
 function define_multiary_operator(op)
     return quote
-        function $op(Lop::Tuple,
+        function $op(Lop::Tuple{<:$Location, <:$Location, <:$Location},
                      a::Union{Function, Number, Oceananigans.Fields.AbstractField},
                      b::Union{Function, Number, Oceananigans.Fields.AbstractField},
                      c::Union{Function, Number, Oceananigans.Fields.AbstractField},
@@ -49,10 +49,17 @@ function define_multiary_operator(op)
             return Oceananigans.AbstractOperations._multiary_operation(Lop, $op, args, Largs, grid)
         end
 
+        # Instantiate location if types are passed
+        $op(Lop::Tuple, 
+            a::Union{Function, Number, Oceananigans.Fields.AbstractField}, 
+            b::Union{Function, Number, Oceananigans.Fields.AbstractField}, 
+            c::Union{Function, Number, Oceananigans.Fields.AbstractField}, 
+            d::Union{Function, Number, Oceananigans.Fields.AbstractField}...) = $op((Lop[1](), Lop[2](), Lop[3]()), a, b, c, d...)
+
         $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...)
+            b::Union{Function, Number, Oceananigans.Fields.AbstractField},
+            c::Union{Function, Number, Oceananigans.Fields.AbstractField},
+            d::Union{Function, Number, 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)
@@ -102,12 +102,15 @@ macro unary(ops...)
             Returns an abstract representation of the operator `$($op)` acting on the Oceananigans `Field`
             `a`, and subsequently interpolated to the location indicated by `Lop`.
             """
-            function $op(Lop::Tuple, a::AbstractField)
-                L = location(a)
+            function $op(Lop::Tuple{<:$Location, <:$Location, <:$Location}, a::AbstractField)
+                L = instantiated_location(a)
                 return Oceananigans.AbstractOperations._unary_operation(Lop, $op, a, L, a.grid)
             end
 
-            $op(a::AbstractField) = $op(location(a), a)
+            # instantiate location if types are passed
+            $op(Lc::Tuple, a::AbstractField) = $op((Lc[1](), Lc[2](), Lc[3]()), a)
+
+            $op(a::AbstractField) = $op(instantiated_location(a), a)
 
             push!(Oceananigans.AbstractOperations.operators, Symbol($op))
             push!(Oceananigans.AbstractOperations.unary_operators, Symbol($op))

src/Fields/Fields.jl

@@ -25,7 +25,6 @@ import Oceananigans: location, instantiated_location
     return (LX(), LY(), LZ())
 end
 
-
 include("abstract_field.jl")
 include("constant_field.jl")
 include("function_field.jl")