New default size for FaceFields (#644)

New default size for FaceFields

Author: ali-ramadhan

src/AbstractOperations/computations.jl

@@ -1,9 +1,12 @@
+using Oceananigans.Utils: @loop_xyz
+import Oceananigans.Fields: location, total_size
+
 """
-    Computation{R, T, O, G}
+    Computation{T, R, O, G}
 
 Represents an operation performed over the elements of a field.
 """
-struct Computation{R, T, O, G}
+struct Computation{T, R, O, G}
       operation :: O
          result :: R
            grid :: G
@@ -20,8 +23,11 @@ Returns a `Computation` representing an `operation` performed over the elements
 Computation(operation, result; return_type=Array) =
     Computation(operation, result, operation.grid, return_type)
 
+# Utilities for limited field-like behavior
 architecture(comp::Computation) = architecture(comp.result)
 Base.parent(comp::Computation) = comp # this enables taking a "horizontal average" of a computation
+location(comp::Computation) = location(comp.operation)
+total_size(comp::Computation) = total_size(comp.result)
 
 """
     compute!(computation::Computation)
@@ -38,12 +44,8 @@ end
 
 """Compute an `operation` over `grid` and store in `result`."""
 function _compute!(result, grid, operation)
-    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
-        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
-            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
-                @inbounds result[i, j, k] = operation[i, j, k]
-            end
-        end
+    @loop_xyz i j k grid begin
+        @inbounds result[i, j, k] = operation[i, j, k]
     end
     return nothing
 end

src/BoundaryConditions/BoundaryConditions.jl

@@ -30,7 +30,6 @@ include("fill_halo_regions.jl")
 include("zero_halo_regions.jl")
 
 include("apply_flux_bcs.jl")
-include("apply_flux_periodic_no_penetration_bcs.jl")
 include("apply_value_gradient_bcs.jl")
 
 end

src/BoundaryConditions/apply_flux_periodic_no_penetration_bcs.jl

@@ -30,3 +30,55 @@ function fill_halo_regions!(c::AbstractArray, fieldbcs, arch, grid, args...)
        fill_top_halo!(c, fieldbcs.z.top,    arch, grid, args...)
     return nothing
 end
+
+#####
+##### Halo filling for flux, periodic, and no-penetration boundary conditions.
+#####
+
+# For flux boundary conditions we fill halos as for a *no-flux* boundary condition, and add the
+# flux divergence associated with the flux boundary condition in a separate step. Note that
+# ranges are used to reference the data copied into halos, as this produces views of the correct
+# dimension (eg size = (1, Ny, Nz) for the west halos).
+
+  _fill_west_halo!(c, ::FBC, H, N) = @views @. c.parent[1:H, :, :] = c.parent[1+H:1+H,  :, :]
+ _fill_south_halo!(c, ::FBC, H, N) = @views @. c.parent[:, 1:H, :] = c.parent[:, 1+H:1+H,  :]
+_fill_bottom_halo!(c, ::FBC, H, N) = @views @. c.parent[:, :, 1:H] = c.parent[:, :,  1+H:1+H]
+
+ _fill_east_halo!(c, ::FBC, H, N) = @views @. c.parent[N+H+1:N+2H, :, :] = c.parent[N+H:N+H, :, :]
+_fill_north_halo!(c, ::FBC, H, N) = @views @. c.parent[:, N+H+1:N+2H, :] = c.parent[:, N+H:N+H, :]
+  _fill_top_halo!(c, ::FBC, H, N) = @views @. c.parent[:, :, N+H+1:N+2H] = c.parent[:, :, N+H:N+H]
+
+# Periodic boundary conditions
+  _fill_west_halo!(c, ::PBC, H, N) = @views @. c.parent[1:H, :, :] = c.parent[N+1:N+H, :, :]
+ _fill_south_halo!(c, ::PBC, H, N) = @views @. c.parent[:, 1:H, :] = c.parent[:, N+1:N+H, :]
+_fill_bottom_halo!(c, ::PBC, H, N) = @views @. c.parent[:, :, 1:H] = c.parent[:, :, N+1:N+H]
+
+ _fill_east_halo!(c, ::PBC, H, N) = @views @. c.parent[N+H+1:N+2H, :, :] = c.parent[1+H:2H, :, :]
+_fill_north_halo!(c, ::PBC, H, N) = @views @. c.parent[:, N+H+1:N+2H, :] = c.parent[:, 1+H:2H, :]
+  _fill_top_halo!(c, ::PBC, H, N) = @views @. c.parent[:, :, N+H+1:N+2H] = c.parent[:, :, 1+H:2H]
+
+# Recall that, by convention, the first grid point (k=1) in an array with a no-penetration boundary
+# condition lies on the boundary, where as the last grid point (k=Nz) lies in the domain.
+
+  _fill_west_halo!(c, ::NPBC, H, N) = @views @. c.parent[1:1+H, :, :] = 0
+ _fill_south_halo!(c, ::NPBC, H, N) = @views @. c.parent[:, 1:1+H, :] = 0
+_fill_bottom_halo!(c, ::NPBC, H, N) = @views @. c.parent[:, :, 1:1+H] = 0
+
+ _fill_east_halo!(c, ::NPBC, H, N) = @views @. c.parent[N+H+1:N+2H, :, :] = 0
+_fill_north_halo!(c, ::NPBC, H, N) = @views @. c.parent[:, N+H+1:N+2H, :] = 0
+  _fill_top_halo!(c, ::NPBC, H, N) = @views @. c.parent[:, :, N+H+1:N+2H] = 0
+
+# Generate functions that implement flux, periodic, and no-penetration boundary conditions
+sides = (:west, :east, :south, :north, :top, :bottom)
+coords = (:x, :x, :y, :y, :z, :z)
+
+for (x, side) in zip(coords, sides)
+    outername = Symbol(:fill_, side, :_halo!)
+    innername = Symbol(:_fill_, side, :_halo!)
+    H = Symbol(:H, x)
+    N = Symbol(:N, x)
+    @eval begin
+        $outername(c, bc::Union{FBC, PBC, NPBC}, arch, grid, args...) =
+            $innername(c, bc, grid.$(H), grid.$(N))
+    end
+end

src/BoundaryConditions/fill_halo_regions.jl

@@ -1,6 +1,7 @@
 using Oceananigans.Architectures
 using Oceananigans.BoundaryConditions
 using Oceananigans.Utils
+using Oceananigans.Grids: total_size
 
 """
     HorizontalAverage{F, R, P, I, Ω} <: AbstractDiagnostic
@@ -37,7 +38,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.Nz + 2field.grid.Hz)
+    result = zeros(arch, field.grid, 1, 1, total_size(parent(field))[3])
     return HorizontalAverage(field, result, frequency, interval, 0.0, return_type, field.grid)
 end