Refactor of spacings and correction of spacings-connected bugs

src/AbstractOperations/binary_operations.jl

@@ -1,3 +1,5 @@
+using Oceananigans.Operators
+
 const binary_operators = Set()
 
 struct BinaryOperation{LX, LY, LZ, O, A, B, IA, IB, G, T} <: AbstractOperation{LX, LY, LZ, G, T}

src/AbstractOperations/grid_metrics.jl

@@ -2,51 +2,69 @@ using Adapt
 using Oceananigans.Operators
 using Oceananigans.Grids: AbstractGrid
 using Oceananigans.Fields: AbstractField, default_indices, location
+using Oceananigans.Operators: Δx, Δy, Δz, Ax, Δλ, Δφ, Ay, Az, volume
 
 import Oceananigans.Grids: xspacings, yspacings, zspacings, λspacings, φspacings
 
-abstract type AbstractGridMetric end
+const AbstractGridMetric = Union{typeof(Δx), 
+                                 typeof(Δy), 
+                                 typeof(Δz), 
+                                 typeof(Δλ),
+                                 typeof(Δφ),
+                                 typeof(Ax), 
+                                 typeof(Ay), 
+                                 typeof(Az), 
+                                 typeof(volume)} # Do we want it to be `volume` or just `V` like in the Operators module?
 
-struct XSpacingMetric <: AbstractGridMetric end
-struct YSpacingMetric <: AbstractGridMetric end
-struct ZSpacingMetric <: AbstractGridMetric end
+"""
+    metric_function(loc, metric::AbstractGridMetric)
 
-metric_function_prefix(::XSpacingMetric) = :Δx
-metric_function_prefix(::YSpacingMetric) = :Δy
-metric_function_prefix(::ZSpacingMetric) = :Δz
+Return the function associated with `metric::AbstractGridMetric` at `loc`ation.
+"""
+function metric_function(loc, metric)
+    code = Tuple(interpolation_code(ℓ) for ℓ in loc)
+    if metric isa typeof(volume)
+        metric_function_symbol = Symbol(:V, code...)
+    else
+        metric_function_symbol = Symbol(metric, code...)
+    end
+    return getglobal(@__MODULE__, metric_function_symbol)
+end
 
-struct XAreaMetric <: AbstractGridMetric end
-struct YAreaMetric <: AbstractGridMetric end
-struct ZAreaMetric <: AbstractGridMetric end
+struct GridMetricOperation{LX, LY, LZ, G, T, M} <: AbstractOperation{LX, LY, LZ, G, T}
+          metric :: M
+            grid :: G
+    function GridMetricOperation{LX, LY, LZ}(metric::M, grid::G) where {LX, LY, LZ, M, G}
+        T = eltype(grid)
+        return new{LX, LY, LZ, G, T, M}(metric, grid)
+    end
+end
 
-metric_function_prefix(::XAreaMetric) = :Ax
-metric_function_prefix(::YAreaMetric) = :Ay
-metric_function_prefix(::ZAreaMetric) = :Az
+Adapt.adapt_structure(to, gm::GridMetricOperation{LX, LY, LZ}) where {LX, LY, LZ} =
+         GridMetricOperation{LX, LY, LZ}(Adapt.adapt(to, gm.metric),
+                                         Adapt.adapt(to, gm.grid))
 
-struct VolumeMetric <: AbstractGridMetric end
+on_architecture(to, gm::GridMetricOperation{LX, LY, LZ}) where {LX, LY, LZ} =
+    GridMetricOperation{LX, LY, LZ}(on_architecture(to, gm.metric),
+                                    on_architecture(to, gm.grid))
 
-metric_function_prefix(::VolumeMetric) = :V
+@inline Base.getindex(gm::GridMetricOperation, i, j, k) = gm.metric(i, j, k, gm.grid)
 
-# Convenient instances for users
-const Δx = XSpacingMetric()
-const Δy = YSpacingMetric()
+indices(::GridMetricOperation) = default_indices(3)
 
 """
-    Δz = ZSpacingMetric()
+    GridMetricOperation(L, metric, grid)
 
-Instance of `ZSpacingMetric` that generates `BinaryOperation`s
-between `AbstractField`s and the vertical grid spacing evaluated
+Instance of `GridMetricOperation` that generates `BinaryOperation`s between `AbstractField`s and the metric `metric`
 at the same location as the `AbstractField`.
 
-`Δx` and `Δy` play a similar role for horizontal grid spacings.
-
 Example
 =======
-
 ```jldoctest
+
 julia> using Oceananigans
 
-julia> using Oceananigans.AbstractOperations: Δz
+julia> using Oceananigans.Operators: Δz
 
 julia> c = CenterField(RectilinearGrid(size=(1, 1, 1), extent=(1, 2, 3)));
 
@@ -55,94 +73,15 @@ BinaryOperation at (Center, Center, Center)
 ├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
 └── tree:
     * at (Center, Center, Center)
-    ├── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
-    └── Δzᶜᶜᶜ at (Center, Center, Center)
+   ├── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
+   └── Δzᶜᶜᶜ at (Center, Center, Center)
 
 julia> c .= 1;
 
 julia> c_dz[1, 1, 1]
 3.0
 ```
 """
-const Δz = ZSpacingMetric()
-
-const Ax = XAreaMetric()
-const Ay = YAreaMetric()
-const Az = ZAreaMetric()
-
-"""
-    volume = VolumeMetric()
-
-Instance of `VolumeMetric` that generates `BinaryOperation`s
-between `AbstractField`s and their cell volumes. Summing
-this `BinaryOperation` yields an integral of `AbstractField`
-over the domain.
-
-Example
-=======
-
-```jldoctest
-julia> using Oceananigans
-
-julia> using Oceananigans.AbstractOperations: volume
-
-julia> c = CenterField(RectilinearGrid(size=(2, 2, 2), extent=(1, 2, 3)));
-
-julia> c .= 1;
-
-julia> c_dV = c * volume
-BinaryOperation at (Center, Center, Center)
-├── grid: 2×2×2 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×2×2 halo
-└── tree:
-    * at (Center, Center, Center)
-    ├── 2×2×2 Field{Center, Center, Center} on RectilinearGrid on CPU
-    └── Vᶜᶜᶜ at (Center, Center, Center)
-
-julia> c_dV[1, 1, 1]
-0.75
-
-julia> sum(c_dV)
-6.0
-```
-"""
-const volume = VolumeMetric()
-
-"""
-    metric_function(loc, metric::AbstractGridMetric)
-
-Return the function associated with `metric::AbstractGridMetric`
-at `loc`ation.
-"""
-function metric_function(loc, metric::AbstractGridMetric)
-    code = Tuple(interpolation_code(ℓ) for ℓ in loc)
-    prefix = metric_function_prefix(metric)
-    metric_function_symbol = Symbol(prefix, code...)
-    return getglobal(@__MODULE__, metric_function_symbol)
-end
-
-struct GridMetricOperation{LX, LY, LZ, G, T, M} <: AbstractOperation{LX, LY, LZ, G, T}
-          metric :: M
-            grid :: G
-    function GridMetricOperation{LX, LY, LZ}(metric::M, grid::G) where {LX, LY, LZ, M, G}
-        T = eltype(grid)
-        return new{LX, LY, LZ, G, T, M}(metric, grid)
-    end
-end
-
-Adapt.adapt_structure(to, gm::GridMetricOperation{LX, LY, LZ}) where {LX, LY, LZ} =
-         GridMetricOperation{LX, LY, LZ}(Adapt.adapt(to, gm.metric),
-                                         Adapt.adapt(to, gm.grid))
-
-on_architecture(to, gm::GridMetricOperation{LX, LY, LZ}) where {LX, LY, LZ} =
-    GridMetricOperation{LX, LY, LZ}(on_architecture(to, gm.metric),
-                                    on_architecture(to, gm.grid))
-
-
-@inline Base.getindex(gm::GridMetricOperation, i, j, k) = gm.metric(i, j, k, gm.grid)
-
-indices(::GridMetricOperation) = default_indices(3)
-
-# Special constructor for BinaryOperation
 GridMetricOperation(L, metric, grid) = GridMetricOperation{L[1], L[2], L[3]}(metric_function(L, metric), grid)
 
 #####
@@ -165,13 +104,13 @@ julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1));
 julia> xspacings(grid, Center(), Center(), Center())
 KernelFunctionOperation at (Center, Center, Center)
 ├── grid: 2×4×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×3×3 halo
-├── kernel_function: xspacing (generic function with 27 methods)
+├── kernel_function: Δx (generic function with 27 methods)
 └── arguments: ("Center", "Center", "Center")
 ```
 """
 function xspacings(grid, ℓx, ℓy, ℓz)
     LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz))
-    Δx_op = KernelFunctionOperation{LX, LY, LZ}(xspacing, grid, ℓx, ℓy, ℓz)
+    Δx_op = KernelFunctionOperation{LX, LY, LZ}(Δx, grid, ℓx, ℓy, ℓz)
     return Δx_op
 end
 
@@ -191,13 +130,13 @@ julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1));
 julia> yspacings(grid, Center(), Face(), Center())
 KernelFunctionOperation at (Center, Face, Center)
 ├── grid: 2×4×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×3×3 halo
-├── kernel_function: yspacing (generic function with 27 methods)
+├── kernel_function: Δy (generic function with 27 methods)
 └── arguments: ("Center", "Face", "Center")
 ```
 """
 function yspacings(grid, ℓx, ℓy, ℓz)
     LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz))
-    Δy_op = KernelFunctionOperation{LX, LY, LZ}(yspacing, grid, ℓx, ℓy, ℓz)
+    Δy_op = KernelFunctionOperation{LX, LY, LZ}(Δy, grid, ℓx, ℓy, ℓz)
     return Δy_op
 end
 
@@ -217,21 +156,21 @@ julia> grid = RectilinearGrid(size=(2, 4, 8), extent=(1, 1, 1));
 julia> zspacings(grid, Center(), Center(), Face())
 KernelFunctionOperation at (Center, Center, Face)
 ├── grid: 2×4×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×3×3 halo
-├── kernel_function: zspacing (generic function with 27 methods)
+├── kernel_function: Δz (generic function with 27 methods)
 └── arguments: ("Center", "Center", "Face")
 ```
 """
 function zspacings(grid, ℓx, ℓy, ℓz)
     LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz))
-    Δz_op = KernelFunctionOperation{LX, LY, LZ}(zspacing, grid, ℓx, ℓy, ℓz)
+    Δz_op = KernelFunctionOperation{LX, LY, LZ}(Δz, grid, ℓx, ℓy, ℓz)
     return Δz_op
 end
 
 """
     λspacings(grid, ℓx, ℓy, ℓz)
 
 Return a `KernelFunctionOperation` that computes the grid spacings for `grid`
-in the ``z`` direction at location `ℓx, ℓy, ℓz`.
+in the ``λ`` direction at location `ℓx, ℓy, ℓz`.
 
 Examples
 ========
@@ -246,21 +185,21 @@ julia> grid = LatitudeLongitudeGrid(size=(36, 34, 25),
 julia> λspacings(grid, Center(), Face(), Center())
 KernelFunctionOperation at (Center, Face, Center)
 ├── grid: 36×34×25 LatitudeLongitudeGrid{Float64, Periodic, Bounded, Bounded} on CPU with 3×3×3 halo and with precomputed metrics
-├── kernel_function: λspacing (generic function with 5 methods)
+├── kernel_function: Δλ (generic function with 5 methods)
 └── arguments: ("Center", "Face", "Center")
 ```
 """
 function λspacings(grid, ℓx, ℓy, ℓz)
     LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz))
-    Δλ_op = KernelFunctionOperation{LX, LY, LZ}(λspacing, grid, ℓx, ℓy, ℓz)
+    Δλ_op = KernelFunctionOperation{LX, LY, LZ}(Δλ, grid, ℓx, ℓy, ℓz)
     return Δλ_op
 end
 
 """
     φspacings(grid, ℓx, ℓy, ℓz)
 
 Return a `KernelFunctionOperation` that computes the grid spacings for `grid`
-in the ``z`` direction at location `ℓx, ℓy, ℓz`.
+in the ``φ`` direction at location `ℓx, ℓy, ℓz`.
 
 Examples
 ========
@@ -275,13 +214,13 @@ julia> grid = LatitudeLongitudeGrid(size=(36, 34, 25),
 julia> φspacings(grid, Center(), Face(), Center())
 KernelFunctionOperation at (Center, Face, Center)
 ├── grid: 36×34×25 LatitudeLongitudeGrid{Float64, Periodic, Bounded, Bounded} on CPU with 3×3×3 halo and with precomputed metrics
-├── kernel_function: φspacing (generic function with 5 methods)
+├── kernel_function: Δφ (generic function with 5 methods)
 └── arguments: ("Center", "Face", "Center")
 ```
 """
 function φspacings(grid, ℓx, ℓy, ℓz)
     LX, LY, LZ = map(typeof, (ℓx, ℓy, ℓz))
-    Δφ_op = KernelFunctionOperation{LX, LY, LZ}(φspacing, grid, ℓx, ℓy, ℓz)
+    Δφ_op = KernelFunctionOperation{LX, LY, LZ}(Δφ, grid, ℓx, ℓy, ℓz)
     return Δφ_op
 end
 

src/BoundaryConditions/flat_extrapolation_open_boundary_matching_scheme.jl

@@ -1,4 +1,4 @@
-using Oceananigans.Grids: xspacing, yspacing, zspacing
+using Oceananigans.Operators: Δx, Δy, Δz
 
 """
     FlatExtrapolation
@@ -62,9 +62,9 @@ end
 const c = Center()
 
 @inline function _fill_west_open_halo!(j, k, grid, ϕ, bc::FEOBC, loc, clock, model_fields)
-    Δx₁ = xspacing(1, j, k, grid, c, c, c)
-    Δx₂ = xspacing(2, j, k, grid, c, c, c)
-    Δx₃ = xspacing(3, j, k, grid, c, c, c)
+    Δx₁ = Δx(1, j, k, grid, c, c, c)
+    Δx₂ = Δx(2, j, k, grid, c, c, c)
+    Δx₃ = Δx(3, j, k, grid, c, c, c)
 
     spacing_factor = Δx₁ / (Δx₂ + Δx₃)
 
@@ -78,9 +78,9 @@ end
 @inline function _fill_east_open_halo!(j, k, grid, ϕ, bc::FEOBC, loc, clock, model_fields)
     i = grid.Nx + 1
 
-    Δx₁ = xspacing(i-1, j, k, grid, c, c, c)
-    Δx₂ = xspacing(i-2, j, k, grid, c, c, c)
-    Δx₃ = xspacing(i-3, j, k, grid, c, c, c)
+    Δx₁ = Δx(i-1, j, k, grid, c, c, c)
+    Δx₂ = Δx(i-2, j, k, grid, c, c, c)
+    Δx₃ = Δx(i-3, j, k, grid, c, c, c)
 
     spacing_factor = Δx₁ / (Δx₂ + Δx₃)
 
@@ -92,9 +92,9 @@ end
 end
 
 @inline function _fill_south_open_halo!(i, k, grid, ϕ, bc::FEOBC, loc, clock, model_fields)
-    Δy₁ = yspacing(i, 1, k, grid, c, c, c)
-    Δy₂ = yspacing(i, 2, k, grid, c, c, c)
-    Δy₃ = yspacing(i, 3, k, grid, c, c, c)
+    Δy₁ = Δy(i, 1, k, grid, c, c, c)
+    Δy₂ = Δy(i, 2, k, grid, c, c, c)
+    Δy₃ = Δy(i, 3, k, grid, c, c, c)
 
     spacing_factor = Δy₁ / (Δy₂ + Δy₃)
 
@@ -108,9 +108,9 @@ end
 @inline function _fill_north_open_halo!(i, k, grid, ϕ, bc::FEOBC, loc, clock, model_fields)
     j = grid.Ny + 1
 
-    Δy₁ = yspacing(i, j-1, k, grid, c, c, c)
-    Δy₂ = yspacing(i, j-2, k, grid, c, c, c)
-    Δy₃ = yspacing(i, j-3, k, grid, c, c, c)
+    Δy₁ = Δy(i, j-1, k, grid, c, c, c)
+    Δy₂ = Δy(i, j-2, k, grid, c, c, c)
+    Δy₃ = Δy(i, j-3, k, grid, c, c, c)
 
     spacing_factor = Δy₁ / (Δy₂ + Δy₃)
 
@@ -122,9 +122,9 @@ end
 end
 
 @inline function _fill_bottom_open_halo!(i, j, grid, ϕ, bc::FEOBC, loc, clock, model_fields)
-    Δz₁ = zspacing(i, j, 1, grid, c, c, c)
-    Δz₂ = zspacing(i, j, 2, grid, c, c, c)
-    Δz₃ = zspacing(i, j, 3, grid, c, c, c)
+    Δz₁ = Δz(i, j, 1, grid, c, c, c)
+    Δz₂ = Δz(i, j, 2, grid, c, c, c)
+    Δz₃ = Δz(i, j, 3, grid, c, c, c)
 
     spacing_factor = Δz₁ / (Δz₂ + Δz₃)
 
@@ -138,9 +138,9 @@ end
 @inline function _fill_top_open_halo!(i, j, grid, ϕ, bc::FEOBC, loc, clock, model_fields)
     k = grid.Nz + 1
 
-    Δz₁ = zspacing(i, j, k-1, grid, c, c, c)
-    Δz₂ = zspacing(i, j, k-2, grid, c, c, c)
-    Δz₃ = zspacing(i, j, k-3, grid, c, c, c)
+    Δz₁ = Δz(i, j, k-1, grid, c, c, c)
+    Δz₂ = Δz(i, j, k-2, grid, c, c, c)
+    Δz₃ = Δz(i, j, k-3, grid, c, c, c)
 
     spacing_factor = Δz₁ / (Δz₂ + Δz₃)
 

src/Grids/Grids.jl

@@ -16,7 +16,6 @@ export ξnode, ηnode, rnode
 export xnode, ynode, znode, λnode, φnode
 export xnodes, ynodes, znodes, λnodes, φnodes
 export spacings
-export xspacing, yspacing, zspacing, λspacing, φspacing
 export xspacings, yspacings, zspacings, λspacings, φspacings
 export minimum_xspacing, minimum_yspacing, minimum_zspacing
 export static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ