(0.90.12) Bugfix for SpitExplicitFreeSurface on MultiRegionGrids and make SpitExplicitFreeSurface the default + don't allow more halos than interior points

src/Grids/input_validation.jl

@@ -4,7 +4,9 @@ using Oceananigans: tupleit
 ##### Some validation tools
 #####
 
-# Tuple inflation for topologies with Flat dimensions
+# Tuple inflation/deflation for topologies with Flat dimensions
+
+"adds tup element with `default` value for flat dimensions"
 inflate_tuple(TX, TY, TZ, tup; default) = tup
 
 inflate_tuple(::Type{Flat}, TY, TZ, tup; default) = tuple(default, tup[1], tup[2])
@@ -17,13 +19,26 @@ inflate_tuple(::Type{Flat}, ::Type{Flat}, TZ, tup; default) = (default, default,
 
 inflate_tuple(::Type{Flat}, ::Type{Flat}, ::Type{Flat}, tup; default) = (default, default, default)
 
+"removes tup elements that correspond to flat dimensions"
+deflate_tuple(TX, TY, TZ, tup) = tup
+
+deflate_tuple(::Type{Flat}, TY, TZ, tup) = tuple(tup[2], tup[3])
+deflate_tuple(TY, ::Type{Flat}, TZ, tup) = tuple(tup[1], tup[3])
+deflate_tuple(TY, TZ, ::Type{Flat}, tup) = tuple(tup[1], tup[2])
+
+deflate_tuple(TX, ::Type{Flat}, ::Type{Flat}, tup) = (tup[1],)
+deflate_tuple(::Type{Flat}, TY, ::Type{Flat}, tup) = (tup[2],)
+deflate_tuple(::Type{Flat}, ::Type{Flat}, TZ, tup) = (tup[3],)
+
+deflate_tuple(::Type{Flat}, ::Type{Flat}, ::Type{Flat}, tup) = ()
+
 topological_tuple_length(TX, TY, TZ) = sum(T === Flat ? 0 : 1 for T in (TX, TY, TZ))
 
 """Validate that an argument tuple is the right length and has elements of type `argtype`."""
 function validate_tupled_argument(arg, argtype, argname, len=3; greater_than=0)
-    length(arg) == len        || throw(ArgumentError("length($argname) must be $len."))
-    all(isa.(arg, argtype))   || throw(ArgumentError("$argname=$arg must contain $(argtype)s."))
-    all(arg .> greater_than)  || throw(ArgumentError("Elements of $argname=$arg must be > $(greater_than)!"))
+    length(arg) == len       || throw(ArgumentError("length($argname) must be $len."))
+    all(isa.(arg, argtype))  || throw(ArgumentError("$argname=$arg must contain $(argtype)s."))
+    all(arg .> greater_than) || throw(ArgumentError("Elements of $argname=$arg must be > $(greater_than)!"))
     return nothing
 end
 
@@ -49,21 +64,31 @@ function validate_size(TX, TY, TZ, sz)
     return inflate_tuple(TX, TY, TZ, sz, default=1)
 end
 
-# Note that the default halo size is specified to be 1 in the following function.
-# This is easily changed but many of the tests will fail so this situation needs to be 
+# Note that if provided with halo=nothing, the default halo size for coord i
+# is the min(default_halo_size, size[i]).
+# While this is easy to change, many of tests might fail so this situation needs to be
 # cleaned up.
-function validate_halo(TX, TY, TZ, ::Nothing)
-    halo = Tuple(3 for i = 1:topological_tuple_length(TX, TY, TZ))
-    return validate_halo(TX, TY, TZ, halo)
+function validate_halo(TX, TY, TZ, size, ::Nothing)
+    maximum_halo = size
+    default_halo = (3, 3, 3)
+    halo = map(min, default_halo, maximum_halo)
+    halo = deflate_tuple(TX, TY, TZ, halo)
+    return validate_halo(TX, TY, TZ, size, halo)
 end
 
-function validate_halo(TX, TY, TZ, halo)
+coordinate_name(i) = i == 1 ? "x" : i == 2 ? "y" : "z"
+
+function validate_halo(TX, TY, TZ, size, halo)
     halo = tupleit(halo)
     validate_tupled_argument(halo, Integer, "halo", topological_tuple_length(TX, TY, TZ))
-    return inflate_tuple(TX, TY, TZ, halo, default=0)
-end
+    halo = inflate_tuple(TX, TY, TZ, halo, default=0)
 
-coordinate_name(i) = i == 1 ? "x" : i == 2 ? "y" : "z"
+    for i in 1:3
+        !(halo[i] ≤ size[i]) && throw(ArgumentError("halo must be ≤ size for coordinate $(coordinate_name(i))"))
+    end
+
+    return halo
+end
 
 function validate_dimension_specification(T, ξ, dir, N, FT)
 

src/Grids/latitude_longitude_grid.jl

@@ -272,7 +272,7 @@ function validate_lat_lon_grid_args(topology, size, halo, FT, latitude, longitud
     latitude  = validate_dimension_specification(TY, latitude,  :latitude,  Nφ, FT)
     z         = validate_dimension_specification(TZ, z,         :z,         Nz, FT)
 
-    halo = validate_halo(TX, TY, TZ, halo)
+    halo = validate_halo(TX, TY, TZ, size, halo)
     topology = (TX, TY, TZ)
 
     return topology, size, halo, latitude, longitude, z, precompute_metrics

src/Grids/rectilinear_grid.jl

@@ -289,7 +289,7 @@ end
 function validate_rectilinear_grid_args(topology, size, halo, FT, extent, x, y, z)
     TX, TY, TZ = topology = validate_topology(topology)
     size = validate_size(TX, TY, TZ, size)
-    halo = validate_halo(TX, TY, TZ, halo)
+    halo = validate_halo(TX, TY, TZ, size, halo)
 
     # Validate the rectilinear domain
     x, y, z = validate_rectilinear_domain(TX, TY, TZ, FT, size, extent, x, y, z)

src/Models/HydrostaticFreeSurfaceModels/distributed_split_explicit_free_surface.jl

@@ -84,8 +84,8 @@ end
 
     Rx, Ry, _ = architecture(grid).ranks
 
-    Ax = Rx == 1 ? old_halos[1] : step_halo
-    Ay = Ry == 1 ? old_halos[2] : step_halo
+    Ax = Rx == 1 ? old_halos[1] : max(step_halo, old_halos[1])
+    Ay = Ry == 1 ? old_halos[2] : max(step_halo, old_halos[2])
 
     return (Ax, Ay, old_halos[3])
 end

src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl

@@ -28,10 +28,7 @@ on_architecture(to, free_surface::ExplicitFreeSurface) =
     ExplicitFreeSurface(on_architecture(to, free_surface.η),
                         on_architecture(to, free_surface.gravitational_acceleration))
 
-#####
-##### Interface to HydrostaticFreeSurfaceModel
-#####
-
+# Internal function for HydrostaticFreeSurfaceModel
 function materialize_free_surface(free_surface::ExplicitFreeSurface{Nothing}, velocities, grid)
     η = free_surface_displacement_field(velocities, free_surface, grid)
     g = convert(eltype(grid), free_surface.gravitational_acceleration)

src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl

@@ -108,7 +108,7 @@ end
 ColumnEnsembleSize(; Nz, ensemble=(0, 0), Hz=1) = ColumnEnsembleSize(ensemble, Nz, Hz)
 
 validate_size(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(e.ensemble[1], e.ensemble[2], e.Nz)
-validate_halo(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(0, 0, e.Hz)
+validate_halo(TX, TY, TZ, size, e::ColumnEnsembleSize) = tuple(0, 0, e.Hz)
 
 @inline function time_discretization(closure_array::AbstractArray)
     first_closure = @allowscalar first(closure_array) # assumes all closures have same time-discretization

src/Models/HydrostaticFreeSurfaceModels/slice_ensemble_model_mode.jl

@@ -41,7 +41,7 @@ end
 SliceEnsembleSize(; size, ensemble=0, halo=(1, 1)) = SliceEnsembleSize(ensemble, size[1], size[2], halo[1], halo[2])
 
 validate_size(TX, TY, TZ, e::SliceEnsembleSize) = tuple(e.ensemble, e.Ny, e.Nz)
-validate_halo(TX, TY, TZ, e::SliceEnsembleSize) = tuple(0, e.Hy, e.Hz)
+validate_halo(TX, TY, TZ, size, e::SliceEnsembleSize) = tuple(0, e.Hy, e.Hz)
 
 #####
 ##### CATKEVerticalDiffusivity helpers

src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl

@@ -347,7 +347,7 @@ const FTS = FixedTimeStepSize
 # we set a minimum number of substeps to execute to avoid numerical issues
 const MINIMUM_SUBSTEPS = 5
 
-@inline calculate_substeps(substepping::FNS, Δt) = length(substepping.averaging_weights)
+@inline calculate_substeps(substepping::FNS, Δt=nothing) = length(substepping.averaging_weights)
 @inline calculate_substeps(substepping::FTS, Δt) = max(MINIMUM_SUBSTEPS, ceil(Int, 2 * Δt / substepping.Δt_barotropic))
 
 @inline calculate_adaptive_settings(substepping::FNS, substeps) = substepping.fractional_step_size, substepping.averaging_weights

src/MultiRegion/multi_region_grid.jl

@@ -51,7 +51,7 @@ minimum_spacing(dir, grid::MultiRegionGrid, ℓx, ℓy, ℓz) =
 
 # the default free surface solver; see Models.HydrostaticFreeSurfaceModels
 default_free_surface(grid::MultiRegionGrid; gravitational_acceleration=g_Earth) =
-    ImplicitFreeSurface(; gravitational_acceleration)
+    SplitExplicitFreeSurface(; substeps=50, gravitational_acceleration)
 
 """
     MultiRegionGrid(global_grid; partition = XPartition(2),

src/MultiRegion/multi_region_split_explicit_free_surface.jl

@@ -2,7 +2,9 @@ using Oceananigans.Utils
 using Oceananigans.AbstractOperations: GridMetricOperation, Δz
 using Oceananigans.Models.HydrostaticFreeSurfaceModels: SplitExplicitFreeSurface,
                                                         SplitExplicitSettings,
-                                                        SplitExplicitState
+                                                        SplitExplicitState,
+                                                        FixedSubstepNumber, FixedTimeStepSize,
+                                                        calculate_substeps
 
 import Oceananigans.Models.HydrostaticFreeSurfaceModels: materialize_free_surface, SplitExplicitAuxiliaryFields
 
@@ -44,14 +46,17 @@ end
 
 # Internal function for HydrostaticFreeSurfaceModel
 function materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, grid::MultiRegionGrids)
-
     settings = SplitExplicitSettings(grid; free_surface.settings.settings_kwargs...)
 
+    settings.substepping isa FixedTimeStepSize &&
+        throw(ArgumentError("SplitExplicitFreeSurface on MultiRegionGrids only suports FixedSubstepNumber; re-initialized SplitExplicitFreeSurface using substeps kwarg"))
+
     switch_device!(grid.devices[1])
+
     old_halos = halo_size(getregion(grid, 1))
-    Nsubsteps = length(settings.substepping.averaging_weights)
+    Nsubsteps = calculate_substeps(settings.substepping)
 
-    new_halos = multiregion_split_explicit_halos(old_halos, Nsubsteps+1, grid.partition)         
+    new_halos = multiregion_split_explicit_halos(old_halos, Nsubsteps+1, grid.partition)
     new_grid  = with_halo(new_halos, grid)
 
     η = ZFaceField(new_grid, indices = (:, :, size(new_grid, 3)+1))
@@ -63,5 +68,5 @@ function materialize_free_surface(free_surface::SplitExplicitFreeSurface, veloci
                                     free_surface.settings)
 end
 
-@inline multiregion_split_explicit_halos(old_halos, step_halo, ::XPartition) = (step_halo, old_halos[2], old_halos[3])
-@inline multiregion_split_explicit_halos(old_halos, step_halo, ::YPartition) = (old_halos[1], step_halo, old_halos[3])
+@inline multiregion_split_explicit_halos(old_halos, step_halo, ::XPartition) = (max(step_halo, old_halos[1]), old_halos[2], old_halos[3])
+@inline multiregion_split_explicit_halos(old_halos, step_halo, ::YPartition) = (old_halos[1], max(step_halo, old_halo[2]), old_halos[3])

test/test_multi_region_advection_diffusion.jl

@@ -99,9 +99,7 @@ function solid_body_rotation_test(grid; P = XPartition, regions = 1)
     return merge(model.velocities, model.tracers, (; η = model.free_surface.η))
 end
 
-function diffusion_cosine_test(grid;  P = XPartition, regions = 1, closure, field_name = :c) 
-    κ, m = 1, 2 # diffusivity and cosine wavenumber
-
+function diffusion_cosine_test(grid; P = XPartition, regions = 1, closure, field_name = :c)
     if architecture(grid) isa GPU
         devices = (0, 0)
     else
@@ -110,12 +108,21 @@ function diffusion_cosine_test(grid;  P = XPartition, regions = 1, closure, fiel
 
     mrg = MultiRegionGrid(grid, partition = P(regions), devices = devices)
 
-    model = HydrostaticFreeSurfaceModel(grid = mrg,
-                                        coriolis = nothing,
-                                        closure = closure,
+    # For MultiRegionGrids with regions > 1, the SplitExplicitFreeSurface extends the
+    # halo region in the horizontal. Because the extented halo regio size cannot exceed
+    # the grid's interior size we pick here the number of substeps taking the grid's
+    # size into consideration.
+    substeps = Int(round(maximum(size(grid)) .* 0.8))
+    free_surface = SplitExplicitFreeSurface(; substeps)
+
+    model = HydrostaticFreeSurfaceModel(; grid = mrg,
+                                        free_surface,
+                                        closure,
                                         tracers = :c,
+                                        coriolis = nothing,
                                         buoyancy=nothing)
 
+    m = 2 # cosine wavenumber
     initial_condition(x, y, z) = cos(m * x)
 
     f = fields(model)[field_name]
@@ -125,16 +132,15 @@ function diffusion_cosine_test(grid;  P = XPartition, regions = 1, closure, fiel
     update_state!(model)
 
     # Step forward with small time-step relative to diffusive time-scale
-    Δt = 1e-6 * grid.Lz^2 / κ
+    Δt = 1e-6 * grid.Lz^2 / closure.κ
     for _ = 1:10
         time_step!(model, Δt)
     end
 
     return f
 end
 
-Nx = 32
-Ny = 32
+Nx = Ny = 32
 
 partitioning = [XPartition]
 
@@ -223,16 +229,16 @@ for arch in archs
         diff₂ = ScalarDiffusivity(ν = 1, κ = 1)
         diff₄ = ScalarBiharmonicDiffusivity(ν = 1e-5, κ = 1e-5)
 
-        for fieldname in [:u, :v, :c]
-            for closure in [diff₂, diff₄]
+        for field_name in (:u, :v, :c)
+            for closure in (diff₂, diff₄)
 
-                fs = diffusion_cosine_test(grid; closure, field_name = fieldname)
-                fs = Array(interior(fs));
+                fs = diffusion_cosine_test(grid; closure, field_name, regions = 1)
+                fs = Array(interior(fs))
 
-                for regions in [2], P in partitioning
-                    @info "  Testing diffusion of $fieldname on $regions $(P)s with $(typeof(closure).name.wrapper) on the $arch"
+                for regions in (2,), P in partitioning
+                    @info "  Testing diffusion of $field_name on $regions $(P)s with $(typeof(closure).name.wrapper) on the $arch"
 
-                    f = diffusion_cosine_test(grid; closure, P = P, field_name = fieldname, regions = regions)
+                    f = diffusion_cosine_test(grid; closure, P, field_name, regions)
                     f = interior(reconstruct_global_field(f))
 
                     @test all(f .≈ fs)