CliMA · glwagner · Apr 27, 2021 · Apr 23, 2021 · Apr 23, 2021 · Apr 23, 2021
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
-version = "0.56.0"
+version = "0.56.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/benchmark/benchmark_hydrostatic_model.jl b/benchmark/benchmark_hydrostatic_model.jl
@@ -19,22 +19,26 @@ DataDeps.register(dd)
 
 # All grids have 6 * 32^2 = 6144 grid points.
 grids = Dict(
-    :RegularRectilinearGrid => RegularRectilinearGrid(size=(128, 48, 1), extent=(1, 1, 1)),
-    :RegularLatitudeLongitudeGrid => RegularLatitudeLongitudeGrid(size=(128, 48, 1), longitude=(-180, 180), latitude=(-80, 80), z=(-1, 0)),
-    :ConformalCubedSphereFaceGrid => ConformalCubedSphereFaceGrid(size=(128, 48, 1), z=(-1, 0)),
-    :ConformalCubedSphereGrid => ConformalCubedSphereGrid(datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2", Nz=1, z=(-1, 0))
+     (CPU, :RegularRectilinearGrid)       => RegularRectilinearGrid(size=(128, 48, 1), extent=(1, 1, 1)),
+     (CPU, :RegularLatitudeLongitudeGrid) => RegularLatitudeLongitudeGrid(size=(128, 48, 1), longitude=(-180, 180), latitude=(-80, 80), z=(-1, 0)),
+     (CPU, :ConformalCubedSphereFaceGrid) => ConformalCubedSphereFaceGrid(size=(128, 48, 1), z=(-1, 0)),
+     (CPU, :ConformalCubedSphereGrid)     => ConformalCubedSphereGrid(datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2", Nz=1, z=(-1, 0)),
+     (GPU, :RegularRectilinearGrid)       => RegularRectilinearGrid(size=(128, 48, 1), extent=(1, 1, 1)),
+     (GPU, :RegularLatitudeLongitudeGrid) => RegularLatitudeLongitudeGrid(size=(128, 48, 1), longitude=(-180, 180), latitude=(-80, 80), z=(-1, 0)),
+     #(GPU, :ConformalCubedSphereFaceGrid) => ConformalCubedSphereFaceGrid(size=(128, 48, 1), z=(-1, 0), architecture=GPU()),
+     (GPU, :ConformalCubedSphereGrid)     => ConformalCubedSphereGrid(datadep"cubed_sphere_32_grid/cubed_sphere_32_grid.jld2", Nz=1, z=(-1, 0), architecture=GPU()),
 )
 
 free_surfaces = Dict(
     :ExplicitFreeSurface => ExplicitFreeSurface(),
-    :ImplicitFreeSurface => ImplicitFreeSurface(maximum_iterations=1, tolerance=-Inf)  # Force it to take exactly 1 iteration.
+    :ImplicitFreeSurface => ImplicitFreeSurface(maximum_iterations=1, tolerance=-Inf) # Force it to take exactly 1 iteration.
 )
 
 function benchmark_hydrostatic_model(Arch, grid_type, free_surface_type)
 
     model = HydrostaticFreeSurfaceModel(
               architecture = Arch(),
-                      grid = grids[grid_type],
+                      grid = grids[(Arch, grid_type)],
         momentum_advection = VectorInvariant(),
               free_surface = free_surfaces[free_surface_type]
     )
@@ -55,13 +59,13 @@ Architectures = has_cuda() ? [CPU, GPU] : [CPU]
 grid_types = [
     :RegularRectilinearGrid,
     :RegularLatitudeLongitudeGrid,
-    :ConformalCubedSphereFaceGrid,
+    #:ConformalCubedSphereFaceGrid,
     :ConformalCubedSphereGrid
 ]
 
 free_surface_types = [
     :ExplicitFreeSurface,
-    :ImplicitFreeSurface
+    #:ImplicitFreeSurface
 ]
 
 # Run and summarize benchmarks

diff --git a/examples/geostrophic_adjustment.jl b/examples/geostrophic_adjustment.jl
@@ -57,13 +57,13 @@ g = model.free_surface.gravitational_acceleration
 
 η₀ = coriolis.f * U * L / g # geostrohpic free surface amplitude
 
-ηᵍ(x, y, z) = η₀ * Gaussian(x - x₀, L)
+ηᵍ(x) = η₀ * Gaussian(x - x₀, L)
 
 # We use an initial height field that's twice the geostrophic solution,
 # thus superimposing a geostrophic and ageostrophic component in the free
 # surface displacement field:
 
-ηⁱ(x, y, z) = 2 * ηᵍ(x, y, z)
+ηⁱ(x, y) = 2 * ηᵍ(x)
 
 # We set the initial condition to ``vᵍ`` and ``ηⁱ``,
 

diff --git a/examples/shallow_water_Bickley_jet.jl b/examples/shallow_water_Bickley_jet.jl
@@ -133,7 +133,7 @@ using LinearAlgebra: norm
 
 function perturbation_norm(model)
     compute!(v)
-    return norm(interiorparent(v))
+    return norm(interior(v))
 end
 nothing # hide
 

diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl
@@ -51,7 +51,6 @@ include("derivatives.jl")
 include("at.jl")
 include("broadcasting_abstract_operations.jl")
 include("show_abstract_operations.jl")
-include("averages_of_operations.jl")
 
 # Make some operators!
 

diff --git a/src/AbstractOperations/averages_of_operations.jl b/src/AbstractOperations/averages_of_operations.jl
diff --git a/src/Architectures.jl b/src/Architectures.jl
@@ -1,12 +1,13 @@
 module Architectures
 
-export
-    AbstractArchitecture, AbstractCPUArchitecture, AbstractGPUArchitecture, CPU, GPU,
-    device, architecture, array_type, arch_array
+export AbstractArchitecture, AbstractCPUArchitecture, AbstractGPUArchitecture
+export CPU, GPU
+export device, architecture, array_type, arch_array
 
 using CUDA
-
 using KernelAbstractions
+using Adapt
+using OffsetArrays
 
 """
     AbstractArchitecture
@@ -60,4 +61,13 @@ arch_array(::AbstractCPUArchitecture, A::CuArray) = Array(A)
 arch_array(::AbstractGPUArchitecture, A::Array) = CuArray(A)
 arch_array(::AbstractGPUArchitecture, A::CuArray) = A
 
+const OffsetCPUArray = OffsetArray{FT, N, <:Array} where {FT, N}
+const OffsetGPUArray = OffsetArray{FT, N, <:CuArray} where {FT, N}
+
+Adapt.adapt_structure(::CPU, a::OffsetCPUArray) = a
+Adapt.adapt_structure(::GPU, a::OffsetGPUArray) = a
+
+Adapt.adapt_structure(::GPU, a::OffsetCPUArray) = OffsetArray(CuArray(a.parent), a.offsets...)
+Adapt.adapt_structure(::CPU, a::OffsetGPUArray) = OffsetArray(Array(a.parent), a.offsets...)
+
 end
diff --git a/src/BoundaryConditions/continuous_boundary_function.jl b/src/BoundaryConditions/continuous_boundary_function.jl
@@ -76,14 +76,14 @@ The regularization of `bc.condition::ContinuousBoundaryFunction` requries
    of the boundary.
 """
 function regularize_boundary_condition(bc::BoundaryCondition{C, <:ContinuousBoundaryFunction},
-                                       X, Y, Z, I, model_field_names) where C
+                                       LX, LY, LZ, I, model_field_names) where C
     boundary_func = bc.condition
 
-    indices, interps = index_and_interp_dependencies(X, Y, Z,
+    indices, interps = index_and_interp_dependencies(LX, LY, LZ,
                                                      boundary_func.field_dependencies,
                                                      model_field_names)
 
-    regularized_boundary_func = ContinuousBoundaryFunction{X, Y, Z, I}(boundary_func.func,
+    regularized_boundary_func = ContinuousBoundaryFunction{LX, LY, LZ, I}(boundary_func.func,
                                                                        boundary_func.parameters,
                                                                        boundary_func.field_dependencies,
                                                                        indices, interps)
@@ -95,19 +95,19 @@ end
 ##### Kernel functions
 #####
 
-@inline function (bc::ContinuousBoundaryFunction{Nothing, Y, Z, i})(j, k, grid, clock, model_fields) where {Y, Z, i}
+@inline function (bc::ContinuousBoundaryFunction{Nothing, LY, LZ, i})(j, k, grid, clock, model_fields) where {LY, LZ, i}
     args = user_function_arguments(i, j, k, grid, model_fields, bc.parameters, bc)
-    return bc.func(ynode(Y, j, grid), znode(Z, k, grid), clock.time, args...)
+    return bc.func(ynode(LY(), j, grid), znode(LZ(), k, grid), clock.time, args...)
 end
 
-@inline function (bc::ContinuousBoundaryFunction{X, Nothing, Z, j})(i, k, grid, clock, model_fields) where {X, Z, j}
+@inline function (bc::ContinuousBoundaryFunction{LX, Nothing, LZ, j})(i, k, grid, clock, model_fields) where {LX, LZ, j}
     args = user_function_arguments(i, j, k, grid, model_fields, bc.parameters, bc)
-    return bc.func(xnode(X, i, grid), znode(Z, k, grid), clock.time, args...)
+    return bc.func(xnode(LX(), i, grid), znode(LZ(), k, grid), clock.time, args...)
 end
 
-@inline function (bc::ContinuousBoundaryFunction{X, Y, Nothing, k})(i, j, grid, clock, model_fields) where {X, Y, k}
+@inline function (bc::ContinuousBoundaryFunction{LX, LY, Nothing, k})(i, j, grid, clock, model_fields) where {LX, LY, k}
     args = user_function_arguments(i, j, k, grid, model_fields, bc.parameters, bc)
-    return bc.func(xnode(X, i, grid), ynode(Y, j, grid), clock.time, args...)
+    return bc.func(xnode(LX(), i, grid), ynode(LY(), j, grid), clock.time, args...)
 end
 
 # Don't re-convert ContinuousBoundaryFunctions passed to BoundaryCondition constructor
@@ -117,8 +117,8 @@ BoundaryCondition(TBC, condition::ContinuousBoundaryFunction) =
 Adapt.adapt_structure(to, bc::BoundaryCondition{C, <:ContinuousBoundaryFunction}) where C =
     BoundaryCondition(C, Adapt.adapt(to, bc.condition))
 
-Adapt.adapt_structure(to, bf::ContinuousBoundaryFunction{X, Y, Z, I}) where {X, Y, Z, I} =
-    ContinuousBoundaryFunction{X, Y, Z, I}(Adapt.adapt(to, bf.func),
+Adapt.adapt_structure(to, bf::ContinuousBoundaryFunction{LX, LY, LZ, I}) where {LX, LY, LZ, I} =
+    ContinuousBoundaryFunction{LX, LY, LZ, I}(Adapt.adapt(to, bf.func),
                                            Adapt.adapt(to, bf.parameters),
                                            nothing,
                                            Adapt.adapt(to, bf.field_dependencies_indices),

diff --git a/src/BoundaryConditions/fill_halo_regions.jl b/src/BoundaryConditions/fill_halo_regions.jl
@@ -21,6 +21,8 @@ function fill_halo_regions!(fields::Union{Tuple, NamedTuple}, arch, args...)
     return nothing
 end
 
+fill_halo_regions!(c::OffsetArray, ::Nothing, args...; kwargs...) = nothing
+
 "Fill halo regions in x, y, and z for a given field's data."
 function fill_halo_regions!(c::OffsetArray, fieldbcs, arch, grid, args...; kwargs...)
 

diff --git a/src/BuoyancyModels/buoyancy_field.jl b/src/BuoyancyModels/buoyancy_field.jl
@@ -1,9 +1,11 @@
 using Adapt
 using KernelAbstractions
-using Oceananigans.Fields: AbstractDataField, FieldStatus, validate_field_data, new_data, conditional_compute!
+
+using Oceananigans.Fields: AbstractDataField, FieldStatus, validate_field_data, conditional_compute!
 using Oceananigans.Fields: architecture, tracernames
 using Oceananigans.Architectures: device
-using Oceananigans.Utils: work_layout
+using Oceananigans.Utils: work_layout, new_data
+using Oceananigans.Grids: new_data
 
 import Oceananigans.Fields: compute!, compute_at!
 

diff --git a/src/BuoyancyModels/nonlinear_equation_of_state.jl b/src/BuoyancyModels/nonlinear_equation_of_state.jl
@@ -1,26 +1,26 @@
 using Oceananigans.Fields: AbstractField
-using Oceananigans.Grids: zF, zC
+using Oceananigans.Grids: znode
 using Oceananigans.Operators: Δzᵃᵃᶠ, Δzᵃᵃᶜ
 
 """ Return the geopotential height at `i, j, k` at cell centers. """
 @inline function Zᵃᵃᶜ(i, j, k, grid::AbstractGrid{FT}) where FT
     if k < 1
-        return zC(1, grid) + (1 - k) * Δzᵃᵃᶠ(i, j, 1, grid)
+        return znode(Center(), 1, grid) + (1 - k) * Δzᵃᵃᶠ(i, j, 1, grid)
     elseif k > grid.Nz
-        return zC(grid.Nz, grid) - (k - grid.Nz) * Δzᵃᵃᶠ(i, j, grid.Nz, grid)
+        return znode(Center(), grid.Nz, grid) - (k - grid.Nz) * Δzᵃᵃᶠ(i, j, grid.Nz, grid)
     else
-        return zC(k, grid)
+        return znode(Center(), k, grid)
     end
 end
 
 """ Return the geopotential height at `i, j, k` at cell z-interfaces. """
 @inline function Zᵃᵃᶠ(i, j, k, grid::AbstractGrid{FT}) where FT
     if k < 1
-        return zF(1, grid) + (1 - k) * Δzᵃᵃᶜ(i, j, 1, grid)
+        return znode(Face(), 1, grid) + (1 - k) * Δzᵃᵃᶜ(i, j, 1, grid)
     elseif k > grid.Nz + 1
-        return zF(grid.Nz + 1, grid) - (k - grid.Nz + 1) * Δzᵃᵃᶜ(i, j, k, grid)
+        return znode(Face(), grid.Nz + 1, grid) - (k - grid.Nz + 1) * Δzᵃᵃᶜ(i, j, k, grid)
     else
-        return zF(k, grid)
+        return znode(Face(), k, grid)
     end
 end
 

diff --git a/src/CubedSpheres/CubedSpheres.jl b/src/CubedSpheres/CubedSpheres.jl
@@ -35,10 +35,15 @@ validate_vertical_velocity_boundary_conditions(w::AbstractCubedSphereField) =
 
 import Oceananigans.Models.HydrostaticFreeSurfaceModels: apply_flux_bcs!
 
-apply_flux_bcs!(Gcⁿ::AbstractCubedSphereField, events, c::AbstractCubedSphereField, arch, barrier, clock, model_fields) = [
-    apply_flux_bcs!(get_face(Gcⁿ, face_index), events, get_face(c, face_index), arch, barrier, clock, model_fields)
-    for face_index in 1:length(Gcⁿ.data.faces)
-]
+function apply_flux_bcs!(Gcⁿ::AbstractCubedSphereField, events, c::AbstractCubedSphereField, arch, barrier, clock, model_fields)
+
+    for (face_index, Gcⁿ_face) in enumerate(faces(Gcⁿ))
+        apply_flux_bcs!(Gcⁿ_face, events, get_face(c, face_index), arch, barrier,
+                        clock, get_face(model_fields, face_index))
+    end
+
+    return nothing
+end
 
 #####
 ##### NaN checker for cubed sphere fields

diff --git a/src/CubedSpheres/conformal_cubed_sphere_grid.jl b/src/CubedSpheres/conformal_cubed_sphere_grid.jl
@@ -156,11 +156,11 @@ function ConformalCubedSphereGrid(FT=Float64; face_size, z, radius=R_Earth)
     return ConformalCubedSphereGrid{FT, typeof(faces), typeof(face_connectivity)}(faces, face_connectivity)
 end
 
-function ConformalCubedSphereGrid(filepath::AbstractString, FT=Float64; Nz, z, radius = R_Earth, halo = (1, 1, 1))
+function ConformalCubedSphereGrid(filepath::AbstractString, FT=Float64; Nz, z, architecture = CPU(), radius = R_Earth, halo = (1, 1, 1))
     @warn "ConformalCubedSphereGrid is experimental: use with caution!"
 
     face_topo = (Connected, Connected, Bounded)
-    face_kwargs = (Nz=Nz, z=z, topology=face_topo, radius=radius, halo=halo)
+    face_kwargs = (Nz=Nz, z=z, topology=face_topo, radius=radius, halo=halo, architecture=architecture)
 
     faces = Tuple(ConformalCubedSphereFaceGrid(filepath, FT; face=n, face_kwargs...) for n in 1:6)
 

diff --git a/src/CubedSpheres/cubed_sphere_exchange_bcs.jl b/src/CubedSpheres/cubed_sphere_exchange_bcs.jl
@@ -1,3 +1,5 @@
+using Adapt
+
 using Oceananigans.BoundaryConditions
 using Oceananigans.BoundaryConditions: BCType
 
@@ -68,3 +70,6 @@ function inject_cubed_sphere_exchange_boundary_conditions(field_bcs, face_number
 
     return FieldBoundaryConditions(x_bcs, y_bcs, field_bcs.z)
 end
+
+Adapt.adapt_structure(to, ::CubedSphereExchangeInformation) = nothing
+Adapt.adapt_structure(to, ::CubedSphereExchangeBC) = nothing
diff --git a/src/CubedSpheres/cubed_sphere_faces.jl b/src/CubedSpheres/cubed_sphere_faces.jl
@@ -6,7 +6,8 @@ using OffsetArrays: OffsetArray
 import Base: getindex, size, show, minimum, maximum
 import Statistics: mean
 
-import Oceananigans.Fields: AbstractField, AbstractDataField, AbstractReducedField, Field, new_data, minimum, maximum, mean, location
+import Oceananigans.Fields: AbstractField, AbstractDataField, AbstractReducedField, Field, minimum, maximum, mean, location
+import Oceananigans.Grids: new_data
 import Oceananigans.BoundaryConditions: FieldBoundaryConditions
 
 struct CubedSphereFaces{E, F}

diff --git a/src/CubedSpheres/cubed_sphere_halo_filling.jl b/src/CubedSpheres/cubed_sphere_halo_filling.jl
@@ -1,5 +1,5 @@
 import Oceananigans.BoundaryConditions:
-    fill_halo_regions!, fill_west_halo!, fill_east_halo!, fill_south_halo!, fill_north_halo!
+    fill_halo_regions!, fill_top_halo!, fill_bottom_halo!, fill_west_halo!, fill_east_halo!, fill_south_halo!, fill_north_halo!
 
 import Oceananigans.Models.HydrostaticFreeSurfaceModels: fill_horizontal_velocity_halos!
 
@@ -11,10 +11,10 @@ fill_north_halo!(c, bc::CubedSphereExchangeBC, args...; kwargs...) = nothing
 
 function fill_halo_regions!(field::AbstractCubedSphereField, arch, args...; kwargs...)
 
-    for face_field in faces(field)
-        # Fill the top and bottom halos the usual way.
-        fill_halo_regions!(face_field, arch, args...; kwargs...)
+    # Fill the top and bottom halos the usual way.
+    #fill_halo_regions!(field, arch, args...; kwargs...)
 
+    for face_field in faces(field)
         # Deal with halo exchanges.
         fill_west_halo!(face_field, field)
         fill_east_halo!(face_field, field)

diff --git a/src/CubedSpheres/cubed_sphere_set!.jl b/src/CubedSpheres/cubed_sphere_set!.jl
@@ -4,6 +4,7 @@ using Oceananigans.Fields: AbstractField, ReducedField
 import Oceananigans.Fields: set!
 
 const CubedSphereCPUField = CubedSphereField{X, Y, Z, <:AbstractCPUArchitecture} where {X, Y, Z}
+const CubedSphereGPUField = CubedSphereField{X, Y, Z, <:AbstractGPUArchitecture} where {X, Y, Z}
 
 function set!(u::CubedSphereCPUField , v::CubedSphereCPUField)
     for (u_face, v_face) in zip(faces(u), faces(v))
@@ -14,9 +15,12 @@ end
 
 set!(field::CubedSphereCPUField, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)]
 set!(field::CubedSphereCPUField, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)]
+set!(field::CubedSphereGPUField, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)]
+set!(field::CubedSphereGPUField, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)]
 
 set_face_field!(field, a) = set!(field, a)
 
+#=
 function set_face_field!(field, f::Function)
     LX, LY, LZ = location(field)
     grid = field.grid
@@ -42,3 +46,4 @@ function set_face_field!(field::ReducedField{LX, LY, Nothing}, f::Function) wher
 
     return nothing
 end
+=#
diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl
@@ -5,7 +5,7 @@ export
     AbstractField, Field,
     CenterField, XFaceField, YFaceField, ZFaceField,
     ReducedField, AveragedField, ComputedField, KernelComputedField, BackgroundField,
-    interior, interiorparent, data,
+    interior, data,
     xnode, ynode, znode, location,
     set!, compute!, @compute,
     VelocityFields, TracerFields, tracernames, PressureFields, TendencyFields,
@@ -15,13 +15,6 @@ using Oceananigans.Architectures
 using Oceananigans.Grids
 using Oceananigans.BoundaryConditions
 
-import Base: zeros
-
-zeros(FT, ::CPU, Nx, Ny, Nz) = zeros(FT, Nx, Ny, Nz)
-zeros(FT, ::GPU, Nx, Ny, Nz) = zeros(FT, Nx, Ny, Nz) |> CuArray
-zeros(arch, grid, Nx, Ny, Nz) = zeros(eltype(grid), arch, Nx, Ny, Nz)
-
-include("new_data.jl")
 include("abstract_field.jl")
 include("field.jl")
 include("zero_field.jl")
@@ -40,5 +33,6 @@ include("interpolate.jl")
 include("field_slicer.jl")
 include("show_fields.jl")
 include("broadcasting_abstract_fields.jl")
+include("mapreduce_abstract_fields.jl")
 
 end