Grids now have a topology (#614)

Grids now have a topology

Author: ali-ramadhan

Project.toml

@@ -1,6 +1,6 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
-version = "0.20.0"
+version = "0.21.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/AbstractOperations/AbstractOperations.jl

@@ -18,7 +18,6 @@ using Oceananigans.Fields
 using Oceananigans.Operators
 using Oceananigans.BoundaryConditions
 
-using Oceananigans: AbstractGrid
 using Oceananigans.Architectures: device
 using Oceananigans.Models: AbstractModel
 using Oceananigans.Diagnostics: HorizontalAverage, normalize_horizontal_sum!

src/AbstractOperations/binary_operations.jl

@@ -39,7 +39,7 @@ end
 """Return an expression that defines an abstract `BinaryOperator` named `op` for `AbstractLocatedField`."""
 function define_binary_operator(op)
     return quote
-        import Oceananigans: AbstractGrid
+        import Oceananigans.Grids: AbstractGrid
         import Oceananigans.Fields: AbstractLocatedField
 
         local location = Oceananigans.Fields.location

src/AbstractOperations/unary_operations.jl

@@ -73,7 +73,7 @@ macro unary(ops...)
 
     for op in ops
         define_unary_operator = quote
-            import Oceananigans: AbstractGrid
+            import Oceananigans.Grids: AbstractGrid
             import Oceananigans.Fields: AbstractLocatedField
 
             local location = Oceananigans.Fields.location

src/BoundaryConditions/BoundaryConditions.jl

@@ -1,7 +1,7 @@
 module BoundaryConditions
 
 export
-    BCType, Periodic, Flux, Gradient, Value, NoPenetration,
+    BCType, Flux, Gradient, Value, NoPenetration,
     BoundaryCondition, bctype, getbc, setbc!,
     CoordinateBoundaryConditions,
     FieldBoundaryConditions, HorizontallyPeriodicBCs, ChannelBCs,

src/Buoyancy/Buoyancy.jl

@@ -5,11 +5,9 @@ export
     LinearEquationOfState, RoquetIdealizedNonlinearEquationOfState,
     ∂x_b, ∂y_b, ∂z_b, buoyancy_perturbation, buoyancy_frequency_squared
 
-using 
-    Printf,
-    Oceananigans.Operators
-
-using Oceananigans: AbstractGrid
+using Printf
+using Oceananigans.Grids
+using Oceananigans.Operators
 
 # Physical constants
 # https://en.wikipedia.org/wiki/Gravitational_acceleration#Gravity_model_for_Earth (30 Oct 2019)

src/Coriolis/Coriolis.jl

@@ -4,6 +4,10 @@ export
     FPlane, NonTraditionalFPlane, BetaPlane,
     x_f_cross_U, y_f_cross_U, z_f_cross_U
 
+using Printf
+using Oceananigans.Grids
+using Oceananigans.Operators
+
 # Physical constants for FPlane and BetaPlane constructors.
 const Ω_Earth = 7.292115e-5 # [s⁻¹] https://en.wikipedia.org/wiki/Earth%27s_rotation#Angular_speed
 const R_Earth = 6371.0e3    # Mean radius of the Earth [m] https://en.wikipedia.org/wiki/Earth

src/Coriolis/beta_plane.jl

@@ -1,8 +1,3 @@
-using Printf
-using Oceananigans.Operators
-
-using Oceananigans: AbstractGrid
-
 """
     BetaPlane{T} <: AbstractRotation
 

src/Coriolis/f_plane.jl

@@ -1,8 +1,3 @@
-using Printf
-using Oceananigans.Operators
-
-using Oceananigans: AbstractGrid
-
 """
     FPlane{FT} <: AbstractRotation
 

src/Coriolis/no_rotation.jl

@@ -1,5 +1,3 @@
-using Oceananigans: AbstractGrid
-
 #####
 ##### Functions for non-rotating models
 #####

src/Diagnostics/horizontal_average.jl

@@ -37,7 +37,7 @@ model and you want to save the output to disk by passing it to an output writer.
 """
 function HorizontalAverage(field; frequency=nothing, interval=nothing, return_type=Array)
     arch = architecture(field)
-    result = zeros(arch, field.grid, 1, 1, field.grid.Tz)
+    result = zeros(arch, field.grid, 1, 1, field.grid.Nz + 2field.grid.Hz)
     return HorizontalAverage(field, result, frequency, interval, 0.0, return_type, field.grid)
 end
 

src/Fields/field.jl

@@ -7,8 +7,9 @@ using OffsetArrays
 import Oceananigans.Architectures: architecture
 import Oceananigans.Utils: datatuple
 
-using Oceananigans: AbstractGrid
-using Oceananigans.Architectures: @hascuda, CPU, GPU
+using Oceananigans.Architectures
+using Oceananigans.Grids
+using Oceananigans.Utils
 
 @hascuda using CuArrays
 
@@ -197,12 +198,12 @@ function OffsetArray(underlying_data, grid)
 end
 
 function Base.zeros(T, ::CPU, grid)
-    underlying_data = zeros(T, grid.Tx, grid.Ty, grid.Tz)
+    underlying_data = zeros(T, grid.Nx + 2grid.Hx, grid.Ny + 2grid.Hy, grid.Nz + 2grid.Hz)
     return OffsetArray(underlying_data, grid)
 end
 
 function Base.zeros(T, ::GPU, grid)
-    underlying_data = CuArray{T}(undef, grid.Tx, grid.Ty, grid.Tz)
+    underlying_data = CuArray{T}(undef, grid.Nx + 2grid.Hx, grid.Ny + 2grid.Hy, grid.Nz + 2grid.Hz)
     underlying_data .= 0  # Gotta do this otherwise you might end up with a few NaN values!
     return OffsetArray(underlying_data, grid)
 end

src/Grids/Grids.jl

@@ -1,9 +1,56 @@
 module Grids
 
-export RegularCartesianGrid, VerticallyStretchedCartesianGrid
+export
+    AbstractTopology, Periodic, Bounded, Flat, topology,
+    AbstractGrid, RegularCartesianGrid, VerticallyStretchedCartesianGrid
+
+import Base: size, length, eltype, show
 
 using Oceananigans
 
+"""
+    AbstractTopology
+
+Abstract supertype for grid topologies.
+"""
+abstract type AbstractTopology end
+
+"""
+    Periodic
+
+Grid topology for periodic dimensions.
+"""
+struct Periodic <: AbstractTopology end
+
+"""
+    Bounded
+
+Grid topology for bounded dimensions. These could be wall-bounded dimensions
+or dimensions
+"""
+struct Bounded <: AbstractTopology end
+
+"""
+    Flat
+
+Grid topology for flat dimensions, generally with one grid point, along which the solution
+is uniform and does not vary.
+"""
+struct Flat <: AbstractTopology end
+
+"""
+    AbstractGrid{FT, TX, TY, TZ}
+
+Abstract supertype for grids with elements of type `FT` and topology `{TX, TY, TZ}`.
+"""
+abstract type AbstractGrid{FT, TX, TY, TZ} end
+
+eltype(::AbstractGrid{FT}) where FT = FT
+topology(::AbstractGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = (TX, TY, TZ)
+
+size(grid::AbstractGrid)   = (grid.Nx, grid.Ny, grid.Nz)
+length(grid::AbstractGrid) = (grid.Lx, grid.Ly, grid.Lz)
+
 include("grid_utils.jl")
 include("regular_cartesian_grid.jl")
 include("vertically_stretched_cartesian_grid.jl")

src/Grids/grid_utils.jl

@@ -8,6 +8,26 @@ unpack_grid(grid) = grid.Nx, grid.Ny, grid.Nz, grid.Lx, grid.Ly, grid.Lz
 ##### Input validation
 #####
 
+instantiate_datatype(t::DataType) = t()
+instantiate_datatype(t) = t
+
+function validate_topology(topology)
+    TX, TY, TZ = topology
+    TX = instantiate_datatype(TX)
+    TY = instantiate_datatype(TY)
+    TZ = instantiate_datatype(TZ)
+
+    for t in (TX, TY, TZ)
+        if !isa(t, AbstractTopology)
+            e = "$(typeof(t)) is not a valid topology! " *
+                "Valid topologies are: Periodic, Bounded, Flat."
+            throw(ArgumentError(e))
+        end
+    end
+
+    return TX, TY, TZ
+end
+
 """Validate that an argument tuple is the right length and has elements of type `argtype`."""
 function validate_tupled_argument(arg, argtype, argname)
     length(arg) == 3        || throw(ArgumentError("length($argname) must be 3."))

src/Grids/regular_cartesian_grid.jl

@@ -1,14 +1,11 @@
-import Base: size, length, eltype, show
-
-using Oceananigans: AbstractGrid
-
 """
-    RegularCartesianGrid{FT<:AbstractFloat, R<:AbstractRange} <: AbstractGrid{FT}
+    RegularCartesianGrid{FT, TX, TY, TZ, R} <: AbstractGrid{FT, TX, TY, TZ}
 
 A Cartesian grid with with constant grid spacings `Δx`, `Δy`, and `Δz` between cell centers
-and cell faces.
+and cell faces, elements of type `FT`, topology `{TX, TY, TZ}`, and coordinate ranges
+of type `R`.
 """
-struct RegularCartesianGrid{FT<:AbstractFloat, R<:AbstractRange} <: AbstractGrid{FT}
+struct RegularCartesianGrid{FT, TX, TY, TZ, R} <: AbstractGrid{FT, TX, TY, TZ}
     # Number of grid points in (x,y,z).
     Nx::Int
     Ny::Int
@@ -17,10 +14,6 @@ struct RegularCartesianGrid{FT<:AbstractFloat, R<:AbstractRange} <: AbstractGrid
     Hx::Int
     Hy::Int
     Hz::Int
-    # Total number of grid points (including halo regions).
-    Tx::Int
-    Ty::Int
-    Tz::Int
     # Domain size [m].
     Lx::FT
     Ly::FT
@@ -41,7 +34,8 @@ struct RegularCartesianGrid{FT<:AbstractFloat, R<:AbstractRange} <: AbstractGrid
 end
 
 """
-    RegularCartesianGrid([FT=Float64]; size, length, x, y, z)
+    RegularCartesianGrid([FT=Float64]; size, length, topology=(Periodic, Periodic, Bounded),
+                         x=nothing, y=nothing, z=nothing)
 
 Creates a `RegularCartesianGrid` with `size = (Nx, Ny, Nz)` grid points.
 
@@ -50,6 +44,10 @@ indicating the left and right endpoints of each dimensions, e.g. `x=(-π, π)` o
 the `length` argument, e.g. `length=(Lx, Ly, Lz)` which specifies the length of each dimension
 in which case 0 ≤ x ≤ Lx, 0 ≤ y ≤ Ly, and -Lz ≤ z ≤ 0.
 
+A grid topology may be specified via a tuple assigning one of `Periodic`, `Bounded, and `Flat`
+to each dimension. By default, a horizontally periodic grid topology `(Periodic, Periodic, Flat)`
+is assumed.
+
 Constants are stored using floating point values of type `FT`. By default this is `Float64`.
 Make sure to specify the desired `FT` if not using `Float64`.
 
@@ -79,23 +77,20 @@ domain: x ∈ [0.0, 8.0], y ∈ [-10.0, 10.0], z ∈ [3.141592653589793, -3.1415
 grid spacing (Δx, Δy, Δz) = (0.25f0, 0.625f0, 0.3926991f0)
 ```
 """
-function RegularCartesianGrid(FT=Float64; size, halo=(1, 1, 1),
+function RegularCartesianGrid(FT=Float64; size, halo=(1, 1, 1), topology=(Periodic, Periodic, Bounded),
                               length=nothing, x=nothing, y=nothing, z=nothing)
 
     # Hack that allows us to use `size` and `length` as keyword arguments but then also
     # use the `length` function.
     sz, len = size, length
     length = Base.length
 
+    TX, TY, TZ = validate_topology(topology)
     Lx, Ly, Lz, x, y, z = validate_grid_size_and_length(sz, len, halo, x, y, z)
 
     Nx, Ny, Nz = sz
     Hx, Hy, Hz = halo
 
-    Tx = Nx + 2Hx
-    Ty = Ny + 2Hy
-    Tz = Nz + 2Hz
-
     Δx = convert(FT, Lx / Nx)
     Δy = convert(FT, Ly / Ny)
     Δz = convert(FT, Lz / Nz)
@@ -112,23 +107,19 @@ function RegularCartesianGrid(FT=Float64; size, halo=(1, 1, 1),
     yF = range(y₁, y₂; length=Ny+1)
     zF = range(z₁, z₂; length=Nz+1)
 
-    RegularCartesianGrid{FT, typeof(xC)}(Nx, Ny, Nz, Hx, Hy, Hz, Tx, Ty, Tz,
-                                         Lx, Ly, Lz, Δx, Δy, Δz, xC, yC, zC, xF, yF, zF)
+    RegularCartesianGrid{FT, typeof(TX), typeof(TY), typeof(TZ), typeof(xC)}(
+        Nx, Ny, Nz, Hx, Hy, Hz, Lx, Ly, Lz, Δx, Δy, Δz, xC, yC, zC, xF, yF, zF)
 end
 
-size(grid::RegularCartesianGrid)   = (grid.Nx, grid.Ny, grid.Nz)
-length(grid::RegularCartesianGrid) = (grid.Lx, grid.Ly, grid.Lz)
-eltype(grid::RegularCartesianGrid{FT}) where FT = FT
-
-short_show(grid::RegularCartesianGrid{T}) where T =
-    "RegularCartesianGrid{$T}(Nx=$(grid.Nx), Ny=$(grid.Ny), Nz=$(grid.Nz))"
+short_show(grid::RegularCartesianGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} =
+    "RegularCartesianGrid{$FT, $TX, $TY, $TZ}(Nx=$(grid.Nx), Ny=$(grid.Ny), Nz=$(grid.Nz))"
 
 show_domain(grid) = string("x ∈ [", grid.xF[1], ", ", grid.xF[end], "], ",
                            "y ∈ [", grid.yF[1], ", ", grid.yF[end], "], ",
                            "z ∈ [", grid.zF[1], ", ", grid.zF[end], "]")
 
-show(io::IO, g::RegularCartesianGrid) =
-    print(io, "RegularCartesianGrid{$(eltype(g))}\n",
+show(io::IO, g::RegularCartesianGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} =
+    print(io, "RegularCartesianGrid{$FT, $TX, $TY, $TZ}\n",
               "domain: x ∈ [$(g.xF[1]), $(g.xF[end])], y ∈ [$(g.yF[1]), $(g.yF[end])], z ∈ [$(g.zF[1]), $(g.zF[end])]", '\n',
               "  resolution (Nx, Ny, Nz) = ", (g.Nx, g.Ny, g.Nz), '\n',
               "   halo size (Hx, Hy, Hz) = ", (g.Hx, g.Hy, g.Hz), '\n',

src/Grids/vertically_stretched_cartesian_grid.jl

@@ -1,14 +1,12 @@
-import Base: size, length, eltype, show
-
-using Oceananigans: AbstractGrid
-
 """
-    VerticallyStretchedCartesianGrid{FT<:AbstractFloat, R<:AbstractRange, A<:AbstractArray} <: AbstractGrid{FT}
+    VerticallyStretchedCartesianGrid{FT, TX, TY, TZ, R, A} <: AbstractGrid{FT, TX, TY, TZ}
 
 A Cartesian grid with with constant horizontal grid spacings `Δx` and `Δy`, and
-non-uniform or stretched vertical grid spacing `Δz` between cell centers and cell faces.
+non-uniform or stretched vertical grid spacing `Δz` between cell centers and cell faces,
+topology `{TX, TY, TZ}`, and coordinate ranges of type `R` (where a range can be used) and
+`A` (where an array is needed).
 """
-struct VerticallyStretchedCartesianGrid{FT<:AbstractFloat, R<:AbstractRange, A<:AbstractArray} <: AbstractGrid{FT}
+struct VerticallyStretchedCartesianGrid{FT, TX, TY, TZ, R, A} <: AbstractGrid{FT, TX, TY, TZ}
     # Number of grid points in (x,y,z).
     Nx::Int
     Ny::Int
@@ -17,10 +15,6 @@ struct VerticallyStretchedCartesianGrid{FT<:AbstractFloat, R<:AbstractRange, A<:
     Hx::Int
     Hy::Int
     Hz::Int
-    # Total number of grid points (including halo regions).
-    Tx::Int
-    Ty::Int
-    Tz::Int
     # Domain size [m].
     Lx::FT
     Ly::FT
@@ -42,22 +36,20 @@ struct VerticallyStretchedCartesianGrid{FT<:AbstractFloat, R<:AbstractRange, A<:
 end
 
 function VerticallyStretchedCartesianGrid(FT=Float64, arch=CPU();
-        size, halo=(1, 1, 1), length=nothing, x=nothing, y=nothing, z=nothing, zF=nothing)
+        size, halo=(1, 1, 1), topology=(Periodic, Periodic, Bounded),
+        length=nothing, x=nothing, y=nothing, z=nothing, zF=nothing)
 
     # Hack that allows us to use `size` and `length` as keyword arguments but then also
     # use the `size` and `length` functions.
     sz, len = size, length
     length = Base.length
 
+    TX, TY, TZ = validate_topology(topology)
     Lx, Ly, Lz, x, y, z = validate_grid_size_and_length(sz, len, halo, x, y, z)
 
     Nx, Ny, Nz = sz
     Hx, Hy, Hz = halo
 
-    Tx = Nx + 2*Hx
-    Ty = Ny + 2*Hy
-    Tz = Nz + 2*Hz
-
     Δx = convert(FT, Lx / Nx)
     Δy = convert(FT, Ly / Ny)
 
@@ -73,6 +65,6 @@ function VerticallyStretchedCartesianGrid(FT=Float64, arch=CPU();
 
     zF, zC, ΔzF, ΔzC = validate_and_generate_variable_grid_spacing(FT, zF, Nz, z₁, z₂)
 
-    VerticallyStretchedCartesianGrid{FT, typeof(xF), typeof(zF)}(Nx, Ny, Nz, Hx, Hy, Hz, Tx, Ty, Tz, Lx, Ly, Lz,
-                                                                 Δx, Δy, ΔzF, ΔzC, xC, yC, zC, xF, yF, zF)
+    VerticallyStretchedCartesianGrid{FT, typeof(TX), typeof(TY), typeof(TZ), typeof(xF), typeof(zF)}(
+        Nx, Ny, Nz, Hx, Hy, Hz, Lx, Ly, Lz, Δx, Δy, ΔzF, ΔzC, xC, yC, zC, xF, yF, zF)
 end