Merge branch 'main' into ss/for-drakkar

src/BoundaryConditions/fill_halo_regions_periodic.jl

@@ -7,7 +7,7 @@ using KernelAbstractions.Extras.LoopInfo: @unroll
 @inline parent_size_and_offset(c, dim1, dim2, size, offset)     = (parent(c), size, fix_halo_offsets.(offset, c.offsets[[dim1, dim2]]))
 @inline parent_size_and_offset(c, dim1, dim2, ::Symbol, offset) = (parent(c), size(parent(c))[[dim1, dim2]], (0, 0))
 
-@inline function parent_size_and_offset(c::NTuple, dim1, dim2, ::Symbol, offset)
+@inline function parent_size_and_offset(c::Tuple, dim1, dim2, ::Symbol, offset)
     p = parent.(c)
     p_size = (minimum([size(t, dim1) for t in p]), minimum([size(t, dim2) for t in p]))
     return p, p_size, (0, 0)
@@ -71,9 +71,10 @@ end
 #### Tupled periodic boundary condition 
 ####
 
-@kernel function fill_periodic_west_and_east_halo!(c::NTuple{M}, ::Val{H}, N) where {M, H}
+@kernel function fill_periodic_west_and_east_halo!(c::Tuple, ::Val{H}, N) where {H}
     j, k = @index(Global, NTuple)
-    @unroll for n = 1:M
+    ntuple(Val(length(c))) do n
+        Base.@_inline_meta
         @unroll for i = 1:H
             @inbounds begin
                   c[n][i, j, k]     = c[n][N+i, j, k] # west
@@ -83,9 +84,10 @@ end
     end
 end
 
-@kernel function fill_periodic_south_and_north_halo!(c::NTuple{M}, ::Val{H}, N) where {M, H}
+@kernel function fill_periodic_south_and_north_halo!(c::Tuple, ::Val{H}, N) where {H}
     i, k = @index(Global, NTuple)
-    @unroll for n = 1:M
+    ntuple(Val(length(c))) do n
+        Base.@_inline_meta
         @unroll for j = 1:H
             @inbounds begin
                 c[n][i, j, k]     = c[n][i, N+j, k] # south
@@ -95,9 +97,10 @@ end
     end
 end
 
-@kernel function fill_periodic_bottom_and_top_halo!(c::NTuple{M}, ::Val{H}, N) where {M, H}
+@kernel function fill_periodic_bottom_and_top_halo!(c::Tuple, ::Val{H}, N) where {H}
     i, j = @index(Global, NTuple)
-    @unroll for n = 1:M
+    ntuple(Val(length(c))) do n
+        Base.@_inline_meta
         @unroll for k = 1:H
             @inbounds begin
                 c[n][i, j, k]     = c[n][i, j, N+k] # top

src/Grids/Grids.jl

@@ -11,7 +11,7 @@ export XRegularRG, YRegularRG, ZRegularRG, XYRegularRG, XYZRegularRG
 export LatitudeLongitudeGrid, XRegularLLG, YRegularLLG, ZRegularLLG
 export OrthogonalSphericalShellGrid, ConformalCubedSphereGrid, ZRegOrthogonalSphericalShellGrid
 export conformal_cubed_sphere_panel
-export MutableVerticalCoordinate
+export MutableVerticalDiscretization
 export node, nodes
 export ξnode, ηnode, rnode
 export xnode, ynode, znode, λnode, φnode
@@ -120,7 +120,7 @@ struct ZDirection <: AbstractDirection end
 struct NegativeZDirection <: AbstractDirection end
 
 include("abstract_grid.jl")
-include("vertical_coordinates.jl")
+include("vertical_discretization.jl")
 include("grid_utils.jl")
 include("nodes_and_spacings.jl")
 include("zeros_and_ones.jl")

src/Grids/grid_generation.jl

@@ -88,7 +88,7 @@ function generate_coordinate(FT, topo::AT, N, H, node_generator, coordinate_name
     C = OffsetArray(on_architecture(arch, C.parent), C.offsets...)
 
     if coordinate_name == :z
-        return L, StaticVerticalCoordinate(F, C, Δᶠ, Δᶜ)
+        return L, StaticVerticalDiscretization(F, C, Δᶠ, Δᶜ)
     else    
         return L, F, C, Δᶠ, Δᶜ
     end
@@ -125,7 +125,7 @@ function generate_coordinate(FT, topo::AT, N, H, node_interval::Tuple{<:Number,
     C = OffsetArray(C, -H)
 
     if coordinate_name == :z
-        return FT(L), StaticVerticalCoordinate(F, C, FT(Δᶠ), FT(Δᶜ))
+        return FT(L), StaticVerticalDiscretization(F, C, FT(Δᶠ), FT(Δᶜ))
     else    
         return FT(L), F, C, FT(Δᶠ), FT(Δᶜ)
     end
@@ -134,7 +134,7 @@ end
 # Flat domains
 function generate_coordinate(FT, ::Flat, N, H, c::Number, coordinate_name, arch) 
     if coordinate_name == :z
-        return FT(1), StaticVerticalCoordinate(range(FT(c), FT(c), length=N), range(FT(c), FT(c), length=N), FT(1), FT(1))
+        return FT(1), StaticVerticalDiscretization(range(FT(c), FT(c), length=N), range(FT(c), FT(c), length=N), FT(1), FT(1))
     else    
         return FT(1), range(FT(c), FT(c), length=N), range(FT(c), FT(c), length=N), FT(1), FT(1)
     end
@@ -145,34 +145,34 @@ end
 #     FT(1), c, c, FT(1), FT(1)
 function generate_coordinate(FT, ::Flat, N, H, ::Nothing, coordinate_name, arch) 
     if coordinate_name == :z
-        return FT(1), StaticVerticalCoordinate(nothing, nothing, FT(1), FT(1))
+        return FT(1), StaticVerticalDiscretization(nothing, nothing, FT(1), FT(1))
     else    
         return FT(1), nothing, nothing, FT(1), FT(1)
     end
 end    
 
 #####
-##### MutableVerticalCoordinate
+##### MutableVerticalDiscretization
 #####
 
-generate_coordinate(FT, ::Periodic, N, H, ::MutableVerticalCoordinate, coordinate_name, arch, args...) = 
-    throw(ArgumentError("Periodic domains are not supported for MutableVerticalCoordinate"))
+generate_coordinate(FT, ::Periodic, N, H, ::MutableVerticalDiscretization, coordinate_name, arch, args...) = 
+    throw(ArgumentError("Periodic domains are not supported for MutableVerticalDiscretization"))
 
 # Generate a vertical coordinate with a scaling (`σ`) with respect to a reference coordinate `r` with spacing `Δr`.
 # The grid might move with time, so the coordinate includes the time-derivative of the scaling `∂t_σ`.
 # The value of the vertical coordinate at `Nz+1` is saved in `ηⁿ`.
-function generate_coordinate(FT, topo, size, halo, coordinate::MutableVerticalCoordinate, coordinate_name, dim::Int, arch)
+function generate_coordinate(FT, topo, size, halo, coordinate::MutableVerticalDiscretization, coordinate_name, dim::Int, arch)
 
     Nx, Ny, Nz = size
     Hx, Hy, Hz = halo
 
     if dim != 3 
-        msg = "MutableVerticalCoordinate is supported only in the third dimension (z)"
+        msg = "MutableVerticalDiscretization is supported only in the third dimension (z)"
         throw(ArgumentError(msg))
     end
 
     if coordinate_name != :z
-        msg = "MutableVerticalCoordinate is supported only for the z-coordinate"
+        msg = "MutableVerticalDiscretization is supported only for the z-coordinate"
         throw(ArgumentError(msg))
     end
 
@@ -195,5 +195,5 @@ function generate_coordinate(FT, topo, size, halo, coordinate::MutableVerticalCo
         fill!(σ, 1)
     end
 
-    return Lr, MutableVerticalCoordinate(rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ, ηⁿ, σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ∂t_σ)
+    return Lr, MutableVerticalDiscretization(rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ, ηⁿ, σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ∂t_σ)
 end

src/Grids/vertical_coordinates.jl → src/Grids/vertical_discretization.jl

@@ -10,7 +10,7 @@
 abstract type AbstractVerticalCoordinate end
 
 """
-    struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate
+    struct StaticVerticalDiscretization{C, D, E, F} <: AbstractVerticalCoordinate
 
 Represent a static one-dimensional vertical coordinate.
 
@@ -22,14 +22,14 @@ Fields
 - `Δᶜ::E`: Cell-centered grid spacing.
 - `Δᶠ::F`: Face-centered grid spacing.
 """
-struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate
+struct StaticVerticalDiscretization{C, D, E, F} <: AbstractVerticalCoordinate
     cᵃᵃᶠ :: C
     cᵃᵃᶜ :: D
     Δᵃᵃᶠ :: E
     Δᵃᵃᶜ :: F
 end
 
-struct MutableVerticalCoordinate{C, D, E, F, H, CC, FC, CF, FF} <: AbstractVerticalCoordinate
+struct MutableVerticalDiscretization{C, D, E, F, H, CC, FC, CF, FF} <: AbstractVerticalCoordinate
     cᵃᵃᶠ :: C
     cᵃᵃᶜ :: D
     Δᵃᵃᶠ :: E
@@ -44,46 +44,46 @@ struct MutableVerticalCoordinate{C, D, E, F, H, CC, FC, CF, FF} <: AbstractVerti
 end
 
 """
-    MutableVerticalCoordinate(r_faces)
+    MutableVerticalDiscretization(r_faces)
 
-Construct a `MutableVerticalCoordinate` from `r_faces` that can be a `Tuple`, a function of an index `k`,
-or an `AbstractArray`. A `MutableVerticalCoordinate` is a vertical coordinate that might evolve in time
-following certain rules. Examples of `MutableVerticalCoordinate`s are free-surface following coordinates, 
+Construct a `MutableVerticalDiscretization` from `r_faces` that can be a `Tuple`, a function of an index `k`,
+or an `AbstractArray`. A `MutableVerticalDiscretization` defines a vertical coordinate that might evolve in time
+following certain rules. Examples of `MutableVerticalDiscretization`s are free-surface following coordinates, 
 or sigma coordinates.
 """
-MutableVerticalCoordinate(r_faces) = MutableVerticalCoordinate(r_faces, r_faces, [nothing for i in 1:9]...)
+MutableVerticalDiscretization(r_faces) = MutableVerticalDiscretization(r_faces, r_faces, [nothing for i in 1:9]...)
 
 ####
 #### Some useful aliases
 ####
 
-const RegularStaticVerticalCoordinate  = StaticVerticalCoordinate{<:Any, <:Any, <:Number}
-const RegularMutableVerticalCoordinate = MutableVerticalCoordinate{<:Any,  <:Any, <:Number}
+const RegularStaticVerticalDiscretization  = StaticVerticalDiscretization{<:Any, <:Any, <:Number}
+const RegularMutableVerticalDiscretization = MutableVerticalDiscretization{<:Any,  <:Any, <:Number}
 
-const RegularVerticalCoordinate = Union{RegularStaticVerticalCoordinate, RegularMutableVerticalCoordinate}
+const RegularVerticalCoordinate = Union{RegularStaticVerticalDiscretization, RegularMutableVerticalDiscretization}
 
-const AbstractMutableGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Bounded, <:MutableVerticalCoordinate}
-const AbstractStaticGrid  = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any,     <:StaticVerticalCoordinate}
+const AbstractMutableGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Bounded, <:MutableVerticalDiscretization}
+const AbstractStaticGrid  = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any,     <:StaticVerticalDiscretization}
 const RegularVerticalGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any,     <:RegularVerticalCoordinate}
 
 ####
 #### Adapt and on_architecture
 ####
 
-Adapt.adapt_structure(to, coord::StaticVerticalCoordinate) =
-    StaticVerticalCoordinate(Adapt.adapt(to, coord.cᵃᵃᶠ),
+Adapt.adapt_structure(to, coord::StaticVerticalDiscretization) =
+    StaticVerticalDiscretization(Adapt.adapt(to, coord.cᵃᵃᶠ),
                              Adapt.adapt(to, coord.cᵃᵃᶜ),
                              Adapt.adapt(to, coord.Δᵃᵃᶠ),
                              Adapt.adapt(to, coord.Δᵃᵃᶜ))
 
-on_architecture(arch, coord::StaticVerticalCoordinate) = 
-    StaticVerticalCoordinate(on_architecture(arch, coord.cᵃᵃᶠ),
+on_architecture(arch, coord::StaticVerticalDiscretization) = 
+    StaticVerticalDiscretization(on_architecture(arch, coord.cᵃᵃᶠ),
                              on_architecture(arch, coord.cᵃᵃᶜ),
                              on_architecture(arch, coord.Δᵃᵃᶠ),
                              on_architecture(arch, coord.Δᵃᵃᶜ))
 
-Adapt.adapt_structure(to, coord::MutableVerticalCoordinate) = 
-    MutableVerticalCoordinate(Adapt.adapt(to, coord.cᵃᵃᶠ),
+Adapt.adapt_structure(to, coord::MutableVerticalDiscretization) = 
+    MutableVerticalDiscretization(Adapt.adapt(to, coord.cᵃᵃᶠ),
                               Adapt.adapt(to, coord.cᵃᵃᶜ),
                               Adapt.adapt(to, coord.Δᵃᵃᶠ),
                               Adapt.adapt(to, coord.Δᵃᵃᶜ),
@@ -95,8 +95,8 @@ Adapt.adapt_structure(to, coord::MutableVerticalCoordinate) =
                               Adapt.adapt(to, coord.σᶜᶜ⁻),
                               Adapt.adapt(to, coord.∂t_σ))
 
-on_architecture(arch, coord::MutableVerticalCoordinate) = 
-    MutableVerticalCoordinate(on_architecture(arch, coord.cᵃᵃᶠ),
+on_architecture(arch, coord::MutableVerticalDiscretization) = 
+    MutableVerticalDiscretization(on_architecture(arch, coord.cᵃᵃᶠ),
                               on_architecture(arch, coord.cᵃᵃᶜ),
                               on_architecture(arch, coord.Δᵃᵃᶠ),
                               on_architecture(arch, coord.Δᵃᵃᶜ),
@@ -155,27 +155,27 @@ z_domain(grid) = domain(topology(grid, 3)(), grid.Nz, grid.z.cᵃᵃᶠ)
 end
 
 @inline cpu_face_constructor_z(grid) = cpu_face_constructor_r(grid)
-@inline cpu_face_constructor_z(grid::AbstractMutableGrid) = MutableVerticalCoordinate(cpu_face_constructor_r(grid))
+@inline cpu_face_constructor_z(grid::AbstractMutableGrid) = MutableVerticalDiscretization(cpu_face_constructor_r(grid))
 
 ####
 #### Utilities
 ####
 
-function validate_dimension_specification(T, ξ::MutableVerticalCoordinate, dir, N, FT)
+function validate_dimension_specification(T, ξ::MutableVerticalDiscretization, dir, N, FT)
     cᶠ = validate_dimension_specification(T, ξ.cᵃᵃᶠ, dir, N, FT)
     cᶜ = validate_dimension_specification(T, ξ.cᵃᵃᶜ, dir, N, FT)
     args = Tuple(getproperty(ξ, prop) for prop in propertynames(ξ))
 
-    return MutableVerticalCoordinate(cᶠ, cᶜ, args[3:end]...)
+    return MutableVerticalDiscretization(cᶠ, cᶜ, args[3:end]...)
 end
 
 # Summaries
-coordinate_summary(topo, z::StaticVerticalCoordinate, name) = coordinate_summary(topo, z.Δᵃᵃᶜ, name)
+coordinate_summary(topo, z::StaticVerticalDiscretization, name) = coordinate_summary(topo, z.Δᵃᵃᶜ, name)
 
-coordinate_summary(::Bounded, z::RegularMutableVerticalCoordinate, name) = 
+coordinate_summary(::Bounded, z::RegularMutableVerticalDiscretization, name) = 
     @sprintf("regularly spaced with Δr=%s (mutable)", prettysummary(z.Δᵃᵃᶜ))
 
-coordinate_summary(::Bounded, z::MutableVerticalCoordinate, name) = 
+coordinate_summary(::Bounded, z::MutableVerticalDiscretization, name) = 
     @sprintf("variably spaced with min(Δr)=%s, max(Δr)=%s (mutable)", 
              prettysummary(minimum(z.Δᵃᵃᶜ)), 
              prettysummary(maximum(z.Δᵃᵃᶜ)))

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

@@ -131,7 +131,7 @@ function HydrostaticFreeSurfaceModel(; grid,
     @apply_regionally validate_model_halo(grid, momentum_advection, tracer_advection, closure)
 
     if !(grid isa MutableGridOfSomeKind) && (vertical_coordinate isa ZStar)
-        error("The grid does not support ZStar vertical coordinates. Use a `MutableVerticalCoordinate` to allow the use of ZStar (see `MutableVerticalCoordinate`).")
+        error("The grid does not support ZStar vertical coordinates. Use a `MutableVerticalDiscretization` to allow the use of ZStar (see `MutableVerticalDiscretization`).")
     end
 
     # Validate biogeochemistry (add biogeochemical tracers automagically)

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl

@@ -137,4 +137,4 @@ where `c = C[tracer_index]`.
              - immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
              + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields)
              + forcing(i, j, k, grid, clock, model_fields))
-end
+end

src/Operators/time_variable_grid_operators.jl

@@ -1,7 +1,7 @@
 import Oceananigans.Grids: znode, AbstractMutableGrid
 
 #####
-##### MutableVerticalCoordinate-specific vertical spacing functions
+##### MutableVerticalDiscretization-specific vertical spacing functions
 #####
 
 const C = Center
@@ -28,9 +28,9 @@ const AMG = AbstractMutableGrid
 #### Vertical spacing functions
 ####
 
-const MRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:MutableVerticalCoordinate}
-const MLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:MutableVerticalCoordinate}
-const MOSG = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:MutableVerticalCoordinate}
+const MRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:MutableVerticalDiscretization}
+const MLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:MutableVerticalDiscretization}
+const MOSG = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:MutableVerticalDiscretization}
 
 superscript_location(s::Symbol) = s == :ᶜ ? :Center : :Face
 

test/test_zstar_coordinate.jl

@@ -3,6 +3,8 @@ include("dependencies_for_runtests.jl")
 using Random
 using Oceananigans: initialize!
 using Oceananigans.ImmersedBoundaries: PartialCellBottom
+using Oceananigans.Grids: MutableVerticalDiscretization
+using Oceananigans.Models: ZStar
 
 function test_zstar_coordinate(model, Ni, Δt)
 
@@ -47,7 +49,7 @@ const F = Face
 @testset "ZStar coordinate scaling tests" begin
     @info "testing the ZStar coordinate scalings"
 
-    z = MutableVerticalCoordinate((-20, 0))
+    z = MutableVerticalDiscretization((-20, 0))
 
     grid = RectilinearGrid(size = (2, 2, 20), 
                               x = (0, 2), 
@@ -79,19 +81,51 @@ const F = Face
     @test column_depthᶜᶜᵃ(2, 1, grid) == 12
     @test  static_column_depthᶜᶜᵃ(1, 1, grid) == 10
     @test  static_column_depthᶜᶜᵃ(2, 1, grid) == 10
+end
 
-    # Make sure a model with a MutableVerticalCoordinate but ZCoordinate still runs
-    model = HydrostaticFreeSurfaceModel(; grid, free_surface = SplitExplicitFreeSurface(grid; substeps = 20))
+@testset "MutableVerticalDiscretization tests" begin
+    @info "testing the MutableVerticalDiscretization in ZCoordinate mode"
 
-    @test begin
-        time_step!(model, 1.0)
-        true
-    end
+    z = MutableVerticalDiscretization((-20, 0))
+
+    # A mutable immersed grid
+    mutable_grid = RectilinearGrid(size=(2, 2, 20), x=(0, 2), y=(0, 1), z=z)
+    mutable_grid = ImmersedBoundaryGrid(mutable_grid, GridFittedBottom((x, y) -> -10))
+
+    # A static immersed grid
+    static_grid = RectilinearGrid(size=(2, 2, 20), x=(0, 2), y=(0, 1), z=(-20, 0))
+    static_grid = ImmersedBoundaryGrid(static_grid, GridFittedBottom((x, y) -> -10))
+
+    # Make sure a model with a MutableVerticalDiscretization but ZCoordinate still runs and 
+    # the results are the same as a model with a static vertical discretization.
+    mutable_model = HydrostaticFreeSurfaceModel(; grid=mutable_grid, free_surface=ImplicitFreeSurface())
+    static_model  = HydrostaticFreeSurfaceModel(; grid=static_grid,  free_surface=ImplicitFreeSurface())
+
+    uᵢ = rand(size(mutable_model.velocities.u)...)
+    vᵢ = rand(size(mutable_model.velocities.v)...)
+
+    set!(mutable_model; u=uᵢ, v=vᵢ)
+    set!(static_model;  u=uᵢ, v=vᵢ)
+
+    static_sim  = Simulation(static_model;  Δt=1e-3, stop_iteration=100)
+    mutable_sim = Simulation(mutable_model; Δt=1e-3, stop_iteration=100)         
+
+    run!(mutable_sim)
+    run!(static_sim)
+
+    # Check that fields are the same
+    um, vm, wm = mutable_model.velocities
+    us, vs, ws = static_model.velocities
+
+    @test all(um.data .≈ us.data)
+    @test all(vm.data .≈ vs.data)
+    @test all(wm.data .≈ ws.data)
+    @test all(um.data .≈ us.data)
 end
 
 @testset "ZStar coordinate simulation testset" begin
-    z_uniform   = MutableVerticalCoordinate((-20, 0))
-    z_stretched = MutableVerticalCoordinate(collect(-20:0))
+    z_uniform   = MutableVerticalDiscretization((-20, 0))
+    z_stretched = MutableVerticalDiscretization(collect(-20:0))
     topologies  = ((Periodic, Periodic, Bounded), 
                    (Periodic, Bounded, Bounded),
                    (Bounded, Periodic, Bounded),