diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 8a162418cd..910d23d885 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -649,6 +649,42 @@ steps: limit: 1 depends_on: "init_cpu" +##### +##### Vertical Coordinates tests +##### + + - label: "🥑 gpu vertical coordinate" + env: + JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" + TEST_GROUP: "vertical_coordinate" + GPU_TEST: "true" + commands: + - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" + agents: + queue: Oceananigans + architecture: GPU + retry: + automatic: + - exit_status: 1 + limit: 1 + depends_on: "init_gpu" + + - label: "🥒 cpu vertical coordinate" + env: + JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" + TEST_GROUP: "vertical_coordinate" + CUDA_VISIBLE_DEVICES: "-1" + commands: + - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" + agents: + queue: Oceananigans + architecture: CPU + retry: + automatic: + - exit_status: 1 + limit: 1 + depends_on: "init_cpu" + ##### ##### Enzyme extension tests ##### diff --git a/Project.toml b/Project.toml index abebca1b28..a1eb9fd78c 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Oceananigans" uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" authors = ["Climate Modeling Alliance and contributors"] -version = "0.95.6" +version = "0.95.7" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl index 2a46564dd8..d528aa8610 100644 --- a/src/Advection/Advection.jl +++ b/src/Advection/Advection.jl @@ -34,6 +34,7 @@ using Oceananigans.Grids: with_halo using Oceananigans.Architectures: architecture, CPU using Oceananigans.Operators +using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, Γᶠᶠᶜ, ∂t_σ import Base: show, summary import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z diff --git a/src/Advection/vector_invariant_advection.jl b/src/Advection/vector_invariant_advection.jl index 5d5becb2f5..7f02532200 100644 --- a/src/Advection/vector_invariant_advection.jl +++ b/src/Advection/vector_invariant_advection.jl @@ -1,6 +1,3 @@ -using Oceananigans.Operators -using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, Γᶠᶠᶜ - # These are also used in Coriolis/hydrostatic_spherical_coriolis.jl struct EnergyConserving{FT} <: AbstractAdvectionScheme{1, FT} end struct EnstrophyConserving{FT} <: AbstractAdvectionScheme{1, FT} end @@ -168,7 +165,7 @@ Base.show(io::IO, a::VectorInvariant{N, FT}) where {N, FT} = ##### Convenience for WENO Vector Invariant ##### -nothing_to_default(user_value; default) = isnothing(user_value) ? default : user_value +nothing_to_default(user_value; default = nothing) = isnothing(user_value) ? default : user_value """ WENOVectorInvariant(FT = Float64; @@ -221,14 +218,14 @@ function WENOVectorInvariant(FT::DataType = Float64; default_upwinding = OnlySelfUpwinding(cross_scheme = divergence_scheme) upwinding = nothing_to_default(upwinding; default = default_upwinding) - N = max(required_halo_size_x(vorticity_scheme), - required_halo_size_y(vorticity_scheme), - required_halo_size_x(divergence_scheme), - required_halo_size_y(divergence_scheme), - required_halo_size_x(kinetic_energy_gradient_scheme), - required_halo_size_y(kinetic_energy_gradient_scheme), - required_halo_size_z(vertical_scheme)) + schemes = (vorticity_scheme, vertical_scheme, kinetic_energy_gradient_scheme, divergence_scheme) + + NX = maximum(required_halo_size_x(s) for s in schemes) + NY = maximum(required_halo_size_y(s) for s in schemes) + NZ = maximum(required_halo_size_z(s) for s in schemes) + N = max(NX, NY, NZ) + FT = eltype(vorticity_scheme) # assumption return VectorInvariant{N, FT, multi_dimensional_stencil}(vorticity_scheme, diff --git a/src/Advection/vector_invariant_cross_upwinding.jl b/src/Advection/vector_invariant_cross_upwinding.jl index c3f2f88fd9..045094bad3 100644 --- a/src/Advection/vector_invariant_cross_upwinding.jl +++ b/src/Advection/vector_invariant_cross_upwinding.jl @@ -17,20 +17,38 @@ ##### Cross and Self Upwinding of the Divergence flux ##### +# If the grid is moving, the discrete continuity equation is calculated as: +# +# ωᵏ⁺¹ - ωᵏ δx(Ax u) + δy(Ay v) Δrᶜᶜᶜ ∂t_σ +# ---------- = - --------------------- - ------------- +# Δzᶜᶜᶜ Vᶜᶜᶜ Δzᶜᶜᶜ +# +# Where ω is the vertical velocity with respect to a moving grid. +# We upwind the discrete divergence `δx(Ax u) + δy(Ay v)` and then divide by the volume, +# therefore, the correct term to be added to the divergence transport due to the moving grid is: +# +# Azᶜᶜᶜ Δrᶜᶜᶜ ∂t_σ +# +# which represents the static volume times the time derivative of the vertical grid scaling. +# If the grid is stationary, ∂t_σ evaluates to zero, so this term disappears from the divergence flux. +@inline Az_Δr_∂t_σ(i, j, k, grid) = Azᶜᶜᶜ(i, j, k, grid) * Δrᶜᶜᶜ(i, j, k, grid) * ∂t_σ(i, j, k, grid) + @inline function upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme::VectorInvariantCrossVerticalUpwinding, u, v) @inbounds û = u[i, j, k] δ_stencil = scheme.upwinding.divergence_stencil - δᴿ = _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) + δᴿ = _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) + ∂t_σ = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, cross_scheme, Az_Δr_∂t_σ) - return û * δᴿ + return û * (δᴿ + ∂t_σ) # For static grids, ∂t_σ == 0 end @inline function upwinded_divergence_flux_Vᶜᶠᶜ(i, j, k, grid, scheme::VectorInvariantCrossVerticalUpwinding, u, v) @inbounds v̂ = v[i, j, k] δ_stencil = scheme.upwinding.divergence_stencil - δᴿ = _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(v̂), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) + δᴿ = _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(v̂), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) + ∂t_σ = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, Az_Δr_∂t_σ) - return v̂ * δᴿ + return v̂ * (δᴿ + ∂t_σ) # For static grids, ∂t_σ == 0 end diff --git a/src/Advection/vector_invariant_self_upwinding.jl b/src/Advection/vector_invariant_self_upwinding.jl index e8ee512cfd..ad1a253832 100644 --- a/src/Advection/vector_invariant_self_upwinding.jl +++ b/src/Advection/vector_invariant_self_upwinding.jl @@ -2,15 +2,20 @@ ##### Self Upwinding of Divergence Flux, the best option! ##### -@inline δx_U(i, j, k, grid, u, v) = δxᶜᵃᵃ(i, j, k, grid, Ax_qᶠᶜᶜ, u) -@inline δy_V(i, j, k, grid, u, v) = δyᵃᶜᵃ(i, j, k, grid, Ay_qᶜᶠᶜ, v) +@inline δx_U(i, j, k, grid, u, v) = δxᶜᶜᶜ(i, j, k, grid, Ax_qᶠᶜᶜ, u) +@inline δy_V(i, j, k, grid, u, v) = δyᶜᶜᶜ(i, j, k, grid, Ay_qᶜᶠᶜ, v) + +# For moving grids, we include the time-derivative of the grid scaling in the divergence flux. +# If the grid is stationary, `Az_Δr_∂t_σ` evaluates to zero. +@inline δx_U_plus_∂t_σ(i, j, k, grid, u, v) = δxᶜᶜᶜ(i, j, k, grid, Ax_qᶠᶜᶜ, u) + Az_Δr_∂t_σ(i, j, k, grid) +@inline δy_V_plus_∂t_σ(i, j, k, grid, u, v) = δyᶜᶜᶜ(i, j, k, grid, Ay_qᶜᶠᶜ, v) + Az_Δr_∂t_σ(i, j, k, grid) # Velocity smoothness for divergence upwinding @inline U_smoothness(i, j, k, grid, u, v) = ℑxᶜᵃᵃ(i, j, k, grid, Ax_qᶠᶜᶜ, u) @inline V_smoothness(i, j, k, grid, u, v) = ℑyᵃᶜᵃ(i, j, k, grid, Ay_qᶜᶠᶜ, v) # Divergence smoothness for divergence upwinding -@inline divergence_smoothness(i, j, k, grid, u, v) = δx_U(i, j, k, grid, u, v) + δy_V(i, j, k, grid, u, v) +@inline divergence_smoothness(i, j, k, grid, u, v) = δx_U(i, j, k, grid, u, v) + δy_V(i, j, k, grid, u, v) @inline function upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme::VectorInvariantSelfVerticalUpwinding, u, v) @@ -18,7 +23,7 @@ cross_scheme = scheme.upwinding.cross_scheme @inbounds û = u[i, j, k] - δvˢ = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, cross_scheme, δy_V, u, v) + δvˢ = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, cross_scheme, δy_V_plus_∂t_σ, u, v) δuᴿ = _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), δx_U, δU_stencil, u, v) return û * (δvˢ + δuᴿ) @@ -30,7 +35,7 @@ end cross_scheme = scheme.upwinding.cross_scheme @inbounds v̂ = v[i, j, k] - δuˢ = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, δx_U, u, v) + δuˢ = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, δx_U_plus_∂t_σ, u, v) δvᴿ = _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(v̂), δy_V, δV_stencil, u, v) return v̂ * (δuˢ + δvᴿ) diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 320e9035cc..4b331048d3 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -11,6 +11,7 @@ export XRegularRG, YRegularRG, ZRegularRG, XYRegularRG, XYZRegularRG export LatitudeLongitudeGrid, XRegularLLG, YRegularLLG, ZRegularLLG export OrthogonalSphericalShellGrid, ConformalCubedSphereGrid, ZRegOrthogonalSphericalShellGrid export conformal_cubed_sphere_panel +export MutableVerticalDiscretization export node, nodes export ξnode, ηnode, rnode export xnode, ynode, znode, λnode, φnode @@ -19,6 +20,7 @@ export spacings export xspacings, yspacings, zspacings, λspacings, φspacings, rspacings export minimum_xspacing, minimum_yspacing, minimum_zspacing export static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ +export column_depthᶜᶜᵃ, column_depthᶠᶜᵃ, column_depthᶜᶠᵃ, column_depthᶠᶠᵃ export offset_data, new_data export on_architecture @@ -118,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") diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl index fc90eeba92..e3c54e323b 100644 --- a/src/Grids/grid_generation.jl +++ b/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,8 +145,55 @@ 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 + +##### +##### MutableVerticalDiscretization +##### + +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::MutableVerticalDiscretization, coordinate_name, dim::Int, arch) + + Nx, Ny, Nz = size + Hx, Hy, Hz = halo + + if dim != 3 + msg = "MutableVerticalDiscretization is supported only in the third dimension (z)" + throw(ArgumentError(msg)) + end + + if coordinate_name != :z + msg = "MutableVerticalDiscretization is supported only for the z-coordinate" + throw(ArgumentError(msg)) + end + + r_faces = coordinate.cᵃᵃᶠ + + Lr, rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ = generate_coordinate(FT, topo[3](), Nz, Hz, r_faces, :r, arch) + + args = (topo, (Nx, Ny, Nz), (Hx, Hy, Hz)) + + σᶜᶜ⁻ = new_data(FT, arch, (Center, Center, Nothing), args...) + σᶜᶜⁿ = new_data(FT, arch, (Center, Center, Nothing), args...) + σᶠᶜⁿ = new_data(FT, arch, (Face, Center, Nothing), args...) + σᶜᶠⁿ = new_data(FT, arch, (Center, Face, Nothing), args...) + σᶠᶠⁿ = new_data(FT, arch, (Face, Face, Nothing), args...) + ηⁿ = new_data(FT, arch, (Center, Center, Nothing), args...) + ∂t_σ = new_data(FT, arch, (Center, Center, Nothing), args...) + + # Fill all the scalings with one for now (i.e. z == r) + for σ in (σᶜᶜ⁻, σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ) + fill!(σ, 1) + end + + return Lr, MutableVerticalDiscretization(rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ, ηⁿ, σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ∂t_σ) +end diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index 92c3aec5b8..7bf74fbdeb 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -299,7 +299,7 @@ end function dimension_summary(topo, name, dom, z::AbstractVerticalCoordinate, pad_domain=0) prefix = domain_summary(topo, name, dom) padding = " "^(pad_domain+1) - return string(prefix, padding, coordinate_summary(topo, z.Δᵃᵃᶜ, name)) + return string(prefix, padding, coordinate_summary(topo, z, name)) end function dimension_summary(topo, name, dom, spacing, pad_domain=0) @@ -317,7 +317,7 @@ coordinate_summary(topo, Δ::Union{AbstractVector, AbstractMatrix}, name) = name, prettysummary(maximum(parent(Δ)))) ##### -##### Static column depths +##### Static and Dynamic column depths ##### @inline static_column_depthᶜᶜᵃ(i, j, grid) = grid.Lz @@ -325,6 +325,12 @@ coordinate_summary(topo, Δ::Union{AbstractVector, AbstractMatrix}, name) = @inline static_column_depthᶠᶜᵃ(i, j, grid) = grid.Lz @inline static_column_depthᶠᶠᵃ(i, j, grid) = grid.Lz +# Will be extended in the `ImmersedBoundaries` module for a ``mutable'' grid type +@inline column_depthᶜᶜᵃ(i, j, k, grid, η) = static_column_depthᶜᶜᵃ(i, j, grid) +@inline column_depthᶠᶜᵃ(i, j, k, grid, η) = static_column_depthᶠᶜᵃ(i, j, grid) +@inline column_depthᶜᶠᵃ(i, j, k, grid, η) = static_column_depthᶜᶠᵃ(i, j, grid) +@inline column_depthᶠᶠᵃ(i, j, k, grid, η) = static_column_depthᶠᶠᵃ(i, j, grid) + ##### ##### Spherical geometry ##### diff --git a/src/Grids/vertical_coordinates.jl b/src/Grids/vertical_coordinates.jl deleted file mode 100644 index 3028fb013b..0000000000 --- a/src/Grids/vertical_coordinates.jl +++ /dev/null @@ -1,109 +0,0 @@ -#### -#### Vertical coordinates -#### - -# This file implements everything related to vertical coordinates in Oceananigans. -# Vertical coordinates are independent of the underlying grid type since only grids that are -# "unstructured" or "curvilinear" in the horizontal directions are supported in Oceananigans. -# Thus the vertical coordinate is _special_, and it can be implemented once for all grid types. - -abstract type AbstractVerticalCoordinate end - -""" - struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate - -Represent a static one-dimensional vertical coordinate. - -Fields -====== - -- `cᶜ::C`: Cell-centered coordinate. -- `cᶠ::D`: Face-centered coordinate. -- `Δᶜ::E`: Cell-centered grid spacing. -- `Δᶠ::F`: Face-centered grid spacing. -""" -struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate - cᵃᵃᶠ :: C - cᵃᵃᶜ :: D - Δᵃᵃᶠ :: E - Δᵃᵃᶜ :: F -end - -#### -#### Some useful aliases -#### - -const RegularVerticalCoordinate = StaticVerticalCoordinate{<:Any, <:Any, <:Number, <:Number} -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(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.cᵃᵃᶜ), - on_architecture(arch, coord.Δᵃᵃᶠ), - on_architecture(arch, coord.Δᵃᵃᶜ)) - -##### -##### Nodes and spacings (common to every grid)... -##### - -AUG = AbstractUnderlyingGrid - -@inline rnode(i, j, k, grid, ℓx, ℓy, ℓz) = rnode(k, grid, ℓz) -@inline rnode(k, grid, ::Center) = getnode(grid.z.cᵃᵃᶜ, k) -@inline rnode(k, grid, ::Face) = getnode(grid.z.cᵃᵃᶠ, k) - -# These will be extended in the Operators module -@inline znode(k, grid, ℓz) = rnode(k, grid, ℓz) -@inline znode(i, j, k, grid, ℓx, ℓy, ℓz) = rnode(i, j, k, grid, ℓx, ℓy, ℓz) - -@inline rnodes(grid::AUG, ℓz::Face; with_halos=false) = _property(grid.z.cᵃᵃᶠ, ℓz, topology(grid, 3), size(grid, 3), with_halos) -@inline rnodes(grid::AUG, ℓz::Center; with_halos=false) = _property(grid.z.cᵃᵃᶜ, ℓz, topology(grid, 3), size(grid, 3), with_halos) -@inline rnodes(grid::AUG, ℓx, ℓy, ℓz; with_halos=false) = rnodes(grid, ℓz; with_halos) - -rnodes(grid::AUG, ::Nothing; kwargs...) = 1:1 -znodes(grid::AUG, ::Nothing; kwargs...) = 1:1 - -# TODO: extend in the Operators module -@inline znodes(grid::AUG, ℓz; kwargs...) = rnodes(grid, ℓz; kwargs...) -@inline znodes(grid::AUG, ℓx, ℓy, ℓz; kwargs...) = rnodes(grid, ℓx, ℓy, ℓz; kwargs...) - -function rspacings end -function zspacings end - -@inline rspacings(grid, ℓz) = rspacings(grid, nothing, nothing, ℓz) -@inline zspacings(grid, ℓz) = zspacings(grid, nothing, nothing, ℓz) - -#### -#### `z_domain` and `cpu_face_constructor_z` -#### - -z_domain(grid) = domain(topology(grid, 3)(), grid.Nz, grid.z.cᵃᵃᶠ) - -@inline cpu_face_constructor_r(grid::RegularVerticalGrid) = z_domain(grid) - -@inline function cpu_face_constructor_r(grid) - Nz = size(grid, 3) - nodes = rnodes(grid, Face(); with_halos=true) - cpu_nodes = on_architecture(CPU(), nodes) - return cpu_nodes[1:Nz+1] -end - -@inline cpu_face_constructor_z(grid) = cpu_face_constructor_r(grid) - -#### -#### Utilities -#### - -# Summaries -coordinate_summary(::Bounded, z::AbstractVerticalCoordinate, name) = - @sprintf("Free-surface following with Δ%s=%s", name, prettysummary(z.Δᵃᵃᶜ)) diff --git a/src/Grids/vertical_discretization.jl b/src/Grids/vertical_discretization.jl new file mode 100644 index 0000000000..d622d58ce8 --- /dev/null +++ b/src/Grids/vertical_discretization.jl @@ -0,0 +1,181 @@ +#### +#### Vertical coordinates +#### + +# This file implements everything related to vertical coordinates in Oceananigans. +# Vertical coordinates are independent of the underlying grid type since only grids that are +# "unstructured" or "curvilinear" in the horizontal directions are supported in Oceananigans. +# Thus the vertical coordinate is _special_, and it can be implemented once for all grid types. + +abstract type AbstractVerticalCoordinate end + +""" + struct StaticVerticalDiscretization{C, D, E, F} <: AbstractVerticalCoordinate + +Represent a static one-dimensional vertical coordinate. + +Fields +====== + +- `cᶜ::C`: Cell-centered coordinate. +- `cᶠ::D`: Face-centered coordinate. +- `Δᶜ::E`: Cell-centered grid spacing. +- `Δᶠ::F`: Face-centered grid spacing. +""" +struct StaticVerticalDiscretization{C, D, E, F} <: AbstractVerticalCoordinate + cᵃᵃᶠ :: C + cᵃᵃᶜ :: D + Δᵃᵃᶠ :: E + Δᵃᵃᶜ :: F +end + +struct MutableVerticalDiscretization{C, D, E, F, H, CC, FC, CF, FF} <: AbstractVerticalCoordinate + cᵃᵃᶠ :: C + cᵃᵃᶜ :: D + Δᵃᵃᶠ :: E + Δᵃᵃᶜ :: F + ηⁿ :: H + σᶜᶜⁿ :: CC + σᶠᶜⁿ :: FC + σᶜᶠⁿ :: CF + σᶠᶠⁿ :: FF + σᶜᶜ⁻ :: CC + ∂t_σ :: CC +end + +""" + MutableVerticalDiscretization(r_faces) + +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. +""" +MutableVerticalDiscretization(r_faces) = MutableVerticalDiscretization(r_faces, r_faces, [nothing for i in 1:9]...) + +#### +#### Some useful aliases +#### + +const RegularStaticVerticalDiscretization = StaticVerticalDiscretization{<:Any, <:Any, <:Number} +const RegularMutableVerticalDiscretization = MutableVerticalDiscretization{<:Any, <:Any, <:Number} + +const RegularVerticalCoordinate = Union{RegularStaticVerticalDiscretization, RegularMutableVerticalDiscretization} + +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::StaticVerticalDiscretization) = + StaticVerticalDiscretization(Adapt.adapt(to, coord.cᵃᵃᶠ), + Adapt.adapt(to, coord.cᵃᵃᶜ), + Adapt.adapt(to, coord.Δᵃᵃᶠ), + Adapt.adapt(to, coord.Δᵃᵃᶜ)) + +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::MutableVerticalDiscretization) = + MutableVerticalDiscretization(Adapt.adapt(to, coord.cᵃᵃᶠ), + Adapt.adapt(to, coord.cᵃᵃᶜ), + Adapt.adapt(to, coord.Δᵃᵃᶠ), + Adapt.adapt(to, coord.Δᵃᵃᶜ), + Adapt.adapt(to, coord.ηⁿ), + Adapt.adapt(to, coord.σᶜᶜⁿ), + Adapt.adapt(to, coord.σᶠᶜⁿ), + Adapt.adapt(to, coord.σᶜᶠⁿ), + Adapt.adapt(to, coord.σᶠᶠⁿ), + Adapt.adapt(to, coord.σᶜᶜ⁻), + Adapt.adapt(to, coord.∂t_σ)) + +on_architecture(arch, coord::MutableVerticalDiscretization) = + MutableVerticalDiscretization(on_architecture(arch, coord.cᵃᵃᶠ), + on_architecture(arch, coord.cᵃᵃᶜ), + on_architecture(arch, coord.Δᵃᵃᶠ), + on_architecture(arch, coord.Δᵃᵃᶜ), + on_architecture(arch, coord.ηⁿ), + on_architecture(arch, coord.σᶜᶜⁿ), + on_architecture(arch, coord.σᶠᶜⁿ), + on_architecture(arch, coord.σᶜᶠⁿ), + on_architecture(arch, coord.σᶠᶠⁿ), + on_architecture(arch, coord.σᶜᶜ⁻), + on_architecture(arch, coord.∂t_σ)) + +##### +##### Nodes and spacings (common to every grid)... +##### + +AUG = AbstractUnderlyingGrid + +@inline rnode(i, j, k, grid, ℓx, ℓy, ℓz) = rnode(k, grid, ℓz) +@inline rnode(k, grid, ::Center) = getnode(grid.z.cᵃᵃᶜ, k) +@inline rnode(k, grid, ::Face) = getnode(grid.z.cᵃᵃᶠ, k) + +# These will be extended in the Operators module +@inline znode(k, grid, ℓz) = rnode(k, grid, ℓz) +@inline znode(i, j, k, grid, ℓx, ℓy, ℓz) = rnode(i, j, k, grid, ℓx, ℓy, ℓz) + +@inline rnodes(grid::AUG, ℓz::Face; with_halos=false) = _property(grid.z.cᵃᵃᶠ, ℓz, topology(grid, 3), size(grid, 3), with_halos) +@inline rnodes(grid::AUG, ℓz::Center; with_halos=false) = _property(grid.z.cᵃᵃᶜ, ℓz, topology(grid, 3), size(grid, 3), with_halos) +@inline rnodes(grid::AUG, ℓx, ℓy, ℓz; with_halos=false) = rnodes(grid, ℓz; with_halos) + +rnodes(grid::AUG, ::Nothing; kwargs...) = 1:1 +znodes(grid::AUG, ::Nothing; kwargs...) = 1:1 + +# TODO: extend in the Operators module +@inline znodes(grid::AUG, ℓz; kwargs...) = rnodes(grid, ℓz; kwargs...) +@inline znodes(grid::AUG, ℓx, ℓy, ℓz; kwargs...) = rnodes(grid, ℓx, ℓy, ℓz; kwargs...) + +function rspacings end +function zspacings end + +@inline rspacings(grid, ℓz) = rspacings(grid, nothing, nothing, ℓz) +@inline zspacings(grid, ℓz) = zspacings(grid, nothing, nothing, ℓz) + +#### +#### `z_domain` and `cpu_face_constructor_z` +#### + +z_domain(grid) = domain(topology(grid, 3)(), grid.Nz, grid.z.cᵃᵃᶠ) + +@inline cpu_face_constructor_r(grid::RegularVerticalGrid) = z_domain(grid) + +@inline function cpu_face_constructor_r(grid) + Nz = size(grid, 3) + nodes = rnodes(grid, Face(); with_halos=true) + cpu_nodes = on_architecture(CPU(), nodes) + return cpu_nodes[1:Nz+1] +end + +@inline cpu_face_constructor_z(grid) = cpu_face_constructor_r(grid) +@inline cpu_face_constructor_z(grid::AbstractMutableGrid) = MutableVerticalDiscretization(cpu_face_constructor_r(grid)) + +#### +#### Utilities +#### + +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 MutableVerticalDiscretization(cᶠ, cᶜ, args[3:end]...) +end + +# Summaries +coordinate_summary(topo, z::StaticVerticalDiscretization, name) = coordinate_summary(topo, z.Δᵃᵃᶜ, name) + +coordinate_summary(::Bounded, z::RegularMutableVerticalDiscretization, name) = + @sprintf("regularly spaced with Δr=%s (mutable)", prettysummary(z.Δᵃᵃᶜ)) + +coordinate_summary(::Bounded, z::MutableVerticalDiscretization, name) = + @sprintf("variably spaced with min(Δr)=%s, max(Δr)=%s (mutable)", + prettysummary(minimum(z.Δᵃᵃᶜ)), + prettysummary(maximum(z.Δᵃᵃᶜ))) diff --git a/src/ImmersedBoundaries/ImmersedBoundaries.jl b/src/ImmersedBoundaries/ImmersedBoundaries.jl index 2b45f4af9a..87b655861f 100644 --- a/src/ImmersedBoundaries/ImmersedBoundaries.jl +++ b/src/ImmersedBoundaries/ImmersedBoundaries.jl @@ -44,5 +44,6 @@ include("immersed_boundary_condition.jl") include("conditional_differences.jl") include("mask_immersed_field.jl") include("immersed_reductions.jl") +include("mutable_immersed_grid.jl") end # module diff --git a/src/ImmersedBoundaries/immersed_grid_metrics.jl b/src/ImmersedBoundaries/immersed_grid_metrics.jl index 54cd4b4eb9..7857d20016 100644 --- a/src/ImmersedBoundaries/immersed_grid_metrics.jl +++ b/src/ImmersedBoundaries/immersed_grid_metrics.jl @@ -1,6 +1,6 @@ using Oceananigans.AbstractOperations: GridMetricOperation -import Oceananigans.Operators: Δrᵃᵃᶠ, Δrᵃᵃᶜ +import Oceananigans.Operators: Δrᵃᵃᶠ, Δrᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ import Oceananigans.Operators: Δxᶠᵃᵃ, Δxᶜᵃᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ import Oceananigans.Operators: Δyᵃᶠᵃ, Δyᵃᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ, Δyᶜᶜᵃ import Oceananigans.Operators: Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ @@ -22,8 +22,8 @@ import Oceananigans.Operators: intrinsic_vector, extrinsic_vector @inline Δrᵃᵃᶠ(i, j, k, ibg::IBG) = Δrᵃᵃᶠ(i, j, k, ibg.underlying_grid) @inline Δrᵃᵃᶜ(i, j, k, ibg::IBG) = Δrᵃᵃᶜ(i, j, k, ibg.underlying_grid) -@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.underlying_grid) @inline Δzᵃᵃᶜ(i, j, k, ibg::IBG) = Δzᵃᵃᶜ(i, j, k, ibg.underlying_grid) +@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.underlying_grid) # 1D Horizontal spacings diff --git a/src/ImmersedBoundaries/mutable_immersed_grid.jl b/src/ImmersedBoundaries/mutable_immersed_grid.jl new file mode 100644 index 0000000000..c4b9f73b57 --- /dev/null +++ b/src/ImmersedBoundaries/mutable_immersed_grid.jl @@ -0,0 +1,58 @@ +using Oceananigans.Grids: AbstractMutableGrid +using Oceananigans.Operators +using Oceananigans.Operators: MRG, MLLG, MOSG, superscript_location + +import Oceananigans.Grids: column_depthᶜᶜᵃ, + column_depthᶜᶠᵃ, + column_depthᶠᶜᵃ, + column_depthᶠᶠᵃ + +import Oceananigans.Operators: σⁿ, σ⁻, ∂t_σ + +const MutableImmersedGrid = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:AbstractMutableGrid} +const MutableGridOfSomeKind = Union{MutableImmersedGrid, AbstractMutableGrid} + +@inline column_depthᶜᶜᵃ(i, j, k, grid::MutableGridOfSomeKind, η) = static_column_depthᶜᶜᵃ(i, j, grid) + @inbounds η[i, j, k] +@inline column_depthᶠᶜᵃ(i, j, k, grid::MutableGridOfSomeKind, η) = static_column_depthᶠᶜᵃ(i, j, grid) + ℑxᶠᵃᵃ(i, j, k, grid, η) +@inline column_depthᶜᶠᵃ(i, j, k, grid::MutableGridOfSomeKind, η) = static_column_depthᶜᶠᵃ(i, j, grid) + ℑyᵃᶠᵃ(i, j, k, grid, η) +@inline column_depthᶠᶠᵃ(i, j, k, grid::MutableGridOfSomeKind, η) = static_column_depthᶠᶠᵃ(i, j, grid) + ℑxyᶠᶠᵃ(i, j, k, grid, η) + +# Convenience methods +@inline column_depthᶜᶜᵃ(i, j, grid) = static_column_depthᶜᶜᵃ(i, j, grid) +@inline column_depthᶜᶠᵃ(i, j, grid) = static_column_depthᶜᶠᵃ(i, j, grid) +@inline column_depthᶠᶜᵃ(i, j, grid) = static_column_depthᶠᶜᵃ(i, j, grid) +@inline column_depthᶠᶠᵃ(i, j, grid) = static_column_depthᶠᶠᵃ(i, j, grid) + +@inline column_depthᶜᶜᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶜᶜᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline column_depthᶜᶠᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶜᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline column_depthᶠᶜᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶠᶜᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline column_depthᶠᶠᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶠᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) + +# Fallbacks +@inline σⁿ(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = σⁿ(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) +@inline σ⁻(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = σ⁻(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) + +@inline ∂t_σ(i, j, k, ibg::IBG) = ∂t_σ(i, j, k, ibg.underlying_grid) + +# Extend the 3D vertical spacing operators on an Immersed Mutable grid +const IMRG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:MRG} +const IMLLG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:MLLG} +const IMOSG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:MOSG} + +for LX in (:ᶠ, :ᶜ), LY in (:ᶠ, :ᶜ), LZ in (:ᶠ, :ᶜ) + zspacing = Symbol(:Δz, LX, LY, LZ) + rspacing = Symbol(:Δr, LX, LY, LZ) + + ℓx = superscript_location(LX) + ℓy = superscript_location(LY) + ℓz = superscript_location(LZ) + + @eval begin + using Oceananigans.Operators: $rspacing + import Oceananigans.Operators: $zspacing + + @inline $zspacing(i, j, k, grid::IMRG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::IMLLG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::IMOSG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + end +end \ No newline at end of file diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl index 2c17fc47d5..7ab1fc8d08 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -3,7 +3,7 @@ module HydrostaticFreeSurfaceModels export HydrostaticFreeSurfaceModel, ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface, - PrescribedVelocityFields + PrescribedVelocityFields, ZStar, ZCoordinate using KernelAbstractions: @index, @kernel using KernelAbstractions.Extras.LoopInfo: @unroll @@ -23,6 +23,9 @@ using Oceananigans.TimeSteppers: SplitRungeKutta3TimeStepper, QuasiAdamsBashfort abstract type AbstractFreeSurface{E, G} end +struct ZCoordinate end +struct ZStar end + # This is only used by the cubed sphere for now. fill_horizontal_velocity_halos!(args...) = nothing @@ -57,6 +60,10 @@ include("SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl") using .SplitExplicitFreeSurfaces +# ZStar implementation +include("z_star_vertical_spacing.jl") + +# Hydrostatic model implementation include("hydrostatic_free_surface_field_tuples.jl") include("hydrostatic_free_surface_model.jl") include("show_hydrostatic_free_surface_model.jl") diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl index 6ce20fd0b7..b2408e695e 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl @@ -11,6 +11,7 @@ using Oceananigans.Utils using Oceananigans.Grids using Oceananigans.Operators using Oceananigans.BoundaryConditions +using Oceananigans.ImmersedBoundaries using Oceananigans.Grids: AbstractGrid, topology using Oceananigans.ImmersedBoundaries: active_linear_index_to_tuple, mask_immersed_field! using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface, @@ -21,6 +22,11 @@ using Base using KernelAbstractions: @index, @kernel using KernelAbstractions.Extras.LoopInfo: @unroll +using Oceananigans.Grids: column_depthᶜᶜᵃ, + column_depthᶜᶠᵃ, + column_depthᶠᶜᵃ, + column_depthᶠᶠᵃ + import Oceananigans.Models.HydrostaticFreeSurfaceModels: initialize_free_surface!, materialize_free_surface, step_free_surface!, diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl index 2b7f7845f6..5b31e5f98c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl @@ -1,39 +1,54 @@ # Kernels to compute the vertical integral of the velocities -@kernel function _barotropic_mode_kernel!(U, V, grid, ::Nothing, u, v) +@kernel function _barotropic_mode_kernel!(U̅, V̅, grid, ::Nothing, u, v, η) i, j = @index(Global, NTuple) - barotropic_mode_kernel!(U, V, i, j, grid, u, v) + barotropic_mode_kernel!(U̅, V̅, i, j, grid, u, v, η) end -@kernel function _barotropic_mode_kernel!(U, V, grid, active_cells_map, u, v) +@kernel function _barotropic_mode_kernel!(U̅, V̅, grid, active_cells_map, u, v, η) idx = @index(Global, Linear) i, j = active_linear_index_to_tuple(idx, active_cells_map) - barotropic_mode_kernel!(U, V, i, j, grid, u, v) + barotropic_mode_kernel!(U̅, V̅, i, j, grid, u, v, η) end -@inline function barotropic_mode_kernel!(U, V, i, j, grid, u, v) - @inbounds U[i, j, 1] = Δzᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1] - @inbounds V[i, j, 1] = Δzᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1] +@inline function barotropic_mode_kernel!(U̅, V̅, i, j, grid, u, v, η) + k_top = size(grid, 3) + 1 + + hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) + hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) + + Hᶠᶜ = column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = column_depthᶜᶠᵃ(i, j, k_top, grid, η) + + # If the static depths are zero (i.e. the column is immersed), + # we set the grid scaling factor to 1 + # (There is no free surface on an immersed column (η == 0)) + σᶠᶜ = ifelse(hᶠᶜ == 0, one(grid), Hᶠᶜ / hᶠᶜ) + σᶜᶠ = ifelse(hᶜᶠ == 0, one(grid), Hᶜᶠ / hᶜᶠ) + + @inbounds U̅[i, j, 1] = Δrᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1] * σᶠᶜ + @inbounds V̅[i, j, 1] = Δrᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1] * σᶜᶠ for k in 2:grid.Nz - @inbounds U[i, j, 1] += Δzᶠᶜᶜ(i, j, k, grid) * u[i, j, k] - @inbounds V[i, j, 1] += Δzᶜᶠᶜ(i, j, k, grid) * v[i, j, k] + @inbounds U̅[i, j, 1] += Δrᶠᶜᶜ(i, j, k, grid) * u[i, j, k] * σᶠᶜ + @inbounds V̅[i, j, 1] += Δrᶜᶠᶜ(i, j, k, grid) * v[i, j, k] * σᶜᶠ end return nothing end -@inline function compute_barotropic_mode!(U, V, grid, u, v) +@inline function compute_barotropic_mode!(U̅, V̅, grid, u, v, η) active_cells_map = retrieve_surface_active_cells_map(grid) - launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, active_cells_map, u, v; active_cells_map) + launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U̅, V̅, grid, active_cells_map, u, v, η; active_cells_map) return nothing end -@kernel function _barotropic_split_explicit_corrector!(u, v, U, V, U̅, V̅, grid) +@kernel function _barotropic_split_explicit_corrector!(u, v, U, V, U̅, V̅, η, grid) i, j, k = @index(Global, NTuple) + k_top = size(grid, 3) + 1 @inbounds begin - Hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) - Hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) + Hᶠᶜ = column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = column_depthᶜᶠᵃ(i, j, k_top, grid, η) u[i, j, k] = u[i, j, k] + (U[i, j, 1] - U̅[i, j, 1]) / Hᶠᶜ v[i, j, k] = v[i, j, k] + (V[i, j, 1] - V̅[i, j, 1]) / Hᶜᶠ @@ -43,6 +58,7 @@ end # Correcting `u` and `v` with the barotropic mode computed in `free_surface` function barotropic_split_explicit_corrector!(u, v, free_surface, grid) state = free_surface.filtered_state + η = free_surface.η U, V = free_surface.barotropic_velocities U̅, V̅ = state.U, state.V arch = architecture(grid) @@ -50,11 +66,11 @@ function barotropic_split_explicit_corrector!(u, v, free_surface, grid) # NOTE: the filtered `U̅` and `V̅` have been copied in the instantaneous `U` and `V`, # so we use the filtered velocities as "work arrays" to store the vertical integrals # of the instantaneous velocities `u` and `v`. - compute_barotropic_mode!(U̅, V̅, grid, u, v) + compute_barotropic_mode!(U̅, V̅, grid, u, v, η) # add in "good" barotropic mode launch!(arch, grid, :xyz, _barotropic_split_explicit_corrector!, - u, v, U, V, U̅, V̅, grid) + u, v, U, V, U̅, V̅, η, grid) return nothing end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl index 4c2976591f..ea49c3591c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl @@ -14,7 +14,7 @@ using Oceananigans.Operators: Δz # from the initial velocity conditions. function initialize_free_surface!(sefs::SplitExplicitFreeSurface, grid, velocities) barotropic_velocities = sefs.barotropic_velocities - @apply_regionally compute_barotropic_mode!(barotropic_velocities.U, barotropic_velocities.V, grid, velocities.u, velocities.v) + @apply_regionally compute_barotropic_mode!(barotropic_velocities.U, barotropic_velocities.V, grid, velocities.u, velocities.v, sefs.η) fill_halo_regions!((barotropic_velocities.U, barotropic_velocities.V)) return nothing end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl index e624054d5e..a450e99de6 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl @@ -135,10 +135,10 @@ function initialize_free_surface_timestepper!(timestepper::AdamsBashforth3Scheme end # The functions `η★` `U★` and `V★` represent the value of free surface, barotropic zonal and meridional velocity at time step m+1/2 -@inline U★(i, j, k, grid, t::ForwardBackwardScheme, Uᵐ) = @inbounds Uᵐ[i, j, k] +@inline U★(i, j, k, grid, ::ForwardBackwardScheme, Uᵐ) = @inbounds Uᵐ[i, j, k] @inline U★(i, j, k, grid, t::AdamsBashforth3Scheme, Uᵐ) = @inbounds t.α * Uᵐ[i, j, k] + t.θ * t.Uᵐ⁻¹[i, j, k] + t.β * t.Uᵐ⁻²[i, j, k] -@inline η★(i, j, k, grid, t::ForwardBackwardScheme, ηᵐ⁺¹) = @inbounds ηᵐ⁺¹[i, j, k] +@inline η★(i, j, k, grid, ::ForwardBackwardScheme, ηᵐ⁺¹) = @inbounds ηᵐ⁺¹[i, j, k] @inline η★(i, j, k, grid, t::AdamsBashforth3Scheme, ηᵐ⁺¹) = @inbounds t.δ * ηᵐ⁺¹[i, j, k] + t.μ * t.ηᵐ[i, j, k] + t.γ * t.ηᵐ⁻¹[i, j, k] + t.ϵ * t.ηᵐ⁻²[i, j, k] @inline store_previous_velocities!(::ForwardBackwardScheme, i, j, k, U) = nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl index 655e2e9068..21c3b11a40 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl @@ -1,3 +1,5 @@ +using Oceananigans.ImmersedBoundaries: MutableGridOfSomeKind + # Evolution Kernels # # ∂t(η) = -∇⋅U @@ -26,8 +28,8 @@ end store_previous_velocities!(timestepper, i, j, 1, U) store_previous_velocities!(timestepper, i, j, 1, V) - Hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) - Hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) + Hᶠᶜ = column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = column_depthᶜᶠᵃ(i, j, k_top, grid, η) @inbounds begin # ∂τ(U) = - ∇η + G @@ -164,5 +166,11 @@ function step_free_surface!(free_surface::SplitExplicitFreeSurface, model, baroc mask_immersed_field!(model.velocities.v) end + # Needed for Mutable to compute the barotropic correction. + # TODO: Would it be possible to remove it in some way? + if model.grid isa MutableGridOfSomeKind + fill_halo_regions!(η) + end + return nothing end diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl index 29303f6581..481b962df0 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl @@ -140,6 +140,7 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para model.diffusivity_fields, model.pressure.pHY′, model.auxiliary_fields, + model.vertical_coordinate, model.clock) u_kernel_args = tuple(start_momentum_kernel_args..., u_immersed_bc, end_momentum_kernel_args..., u_forcing) diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl index 927489cf50..4e633d9499 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl @@ -1,7 +1,8 @@ using Oceananigans.Architectures: device using Oceananigans.Grids: halo_size, topology using Oceananigans.Grids: XFlatGrid, YFlatGrid -using Oceananigans.Operators: div_xyᶜᶜᶜ, Δzᶜᶜᶜ +using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, div_xyᶜᶜᶜ, Δzᶜᶜᶜ +using Oceananigans.ImmersedBoundaries: immersed_cell """ compute_w_from_continuity!(model) @@ -18,12 +19,37 @@ compute_w_from_continuity!(model; kwargs...) = compute_w_from_continuity!(velocities, arch, grid; parameters = w_kernel_parameters(grid)) = launch!(arch, grid, parameters, _compute_w_from_continuity!, velocities, grid) +# If the grid is following the free surface, then the derivative of the moving grid is: +# +# δx(Δy U) + δy(Δx V) ∇ ⋅ U +# ∂t_σ = - --------------------- = - -------- +# Az ⋅ H H +# +# The discrete divergence is then calculated as: +# +# wᵏ⁺¹ - wᵏ δx(Ax u) + δy(Ay v) Δr ∂t_σ +# ---------- = - --------------------- - ---------- +# Δz vol Δz +# +# This makes sure that summing up till the top of the domain, results in: +# +# ∇ ⋅ U +# wᴺᶻ⁺¹ = w⁰ - ------- - ∂t_σ ≈ 0 (if w⁰ == 0) +# H +# +# If the grid is static, then ∂t_σ = 0 and the moving grid contribution is equal to zero @kernel function _compute_w_from_continuity!(U, grid) i, j = @index(Global, NTuple) @inbounds U.w[i, j, 1] = 0 for k in 2:grid.Nz+1 - @inbounds U.w[i, j, k] = U.w[i, j, k-1] - Δzᶜᶜᶜ(i, j, k-1, grid) * div_xyᶜᶜᶜ(i, j, k-1, grid, U.u, U.v) + δh_u = flux_div_xyᶜᶜᶜ(i, j, k-1, grid, U.u, U.v) / Azᶜᶜᶜ(i, j, k-1, grid) + ∂tσ = Δrᶜᶜᶜ(i, j, k-1, grid) * ∂t_σ(i, j, k-1, grid) + + immersed = immersed_cell(i, j, k-1, grid) + Δw = δh_u + ifelse(immersed, zero(grid), ∂tσ) # We do not account for grid changes in immersed cells + + @inbounds U.w[i, j, k] = U.w[i, j, k-1] - Δw end end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl index ad0dbe784a..1f7fa9a40f 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl @@ -90,8 +90,7 @@ function ab2_step_tracers!(tracers, model, Δt, χ) tracer_field = tracers[tracer_name] closure = model.closure - launch!(model.architecture, model.grid, :xyz, - ab2_step_field!, tracer_field, Δt, χ, Gⁿ, G⁻) + ab2_step_tracer_field!(tracer_field, model.grid, Δt, χ, Gⁿ, G⁻) implicit_step!(tracer_field, model.timestepper.implicit_solver, @@ -106,3 +105,34 @@ function ab2_step_tracers!(tracers, model, Δt, χ) return nothing end +ab2_step_tracer_field!(tracer_field, grid, Δt, χ, Gⁿ, G⁻) = + launch!(architecture(grid), grid, :xyz, _ab2_step_tracer_field!, tracer_field, grid, Δt, χ, Gⁿ, G⁻) + +##### +##### Tracer update in mutable vertical coordinates +##### + +# σθ is the evolved quantity. Once σⁿ⁺¹ is known we can retrieve θⁿ⁺¹ +# with the `unscale_tracers!` function +@kernel function _ab2_step_tracer_field!(θ, grid, Δt, χ, Gⁿ, G⁻) + i, j, k = @index(Global, NTuple) + + FT = eltype(χ) + α = convert(FT, 1.5) + χ + β = convert(FT, 0.5) + χ + + σᶜᶜⁿ = σⁿ(i, j, k, grid, Center(), Center(), Center()) + σᶜᶜ⁻ = σ⁻(i, j, k, grid, Center(), Center(), Center()) + + @inbounds begin + ∂t_σθ = α * σᶜᶜⁿ * Gⁿ[i, j, k] - β * σᶜᶜ⁻ * G⁻[i, j, k] + + # We store temporarily σθ in θ. + # The unscaled θ will be retrieved with `unscale_tracers!` + θ[i, j, k] = σᶜᶜⁿ * θ[i, j, k] + convert(FT, Δt) * ∂t_σθ + end +end + +# Fallback! We need to unscale the tracers only in case of +# a grid with a mutable vertical coordinate, i.e. where `σ != 1` +unscale_tracers!(tracers, grid; kwargs...) = nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index 354b273017..675db7f6ba 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -26,25 +26,26 @@ const ParticlesOrNothing = Union{Nothing, AbstractLagrangianParticles} const AbstractBGCOrNothing = Union{Nothing, AbstractBiogeochemistry} mutable struct HydrostaticFreeSurfaceModel{TS, E, A<:AbstractArchitecture, S, - G, T, V, B, R, F, P, BGC, U, C, Φ, K, AF} <: AbstractModel{TS} - - architecture :: A # Computer `Architecture` on which `Model` is run - grid :: G # Grid of physical points on which `Model` is solved - clock :: Clock{T} # Tracks iteration number and simulation time of `Model` - advection :: V # Advection scheme for tracers - buoyancy :: B # Set of parameters for buoyancy model - coriolis :: R # Set of parameters for the background rotation rate of `Model` - free_surface :: S # Free surface parameters and fields - forcing :: F # Container for forcing functions defined by the user - closure :: E # Diffusive 'turbulence closure' for all model fields - particles :: P # Particle set for Lagrangian tracking - biogeochemistry :: BGC # Biogeochemistry for Oceananigans tracers - velocities :: U # Container for velocity fields `u`, `v`, and `w` - tracers :: C # Container for tracer fields - pressure :: Φ # Container for hydrostatic pressure - diffusivity_fields :: K # Container for turbulent diffusivities - timestepper :: TS # Object containing timestepper fields and parameters - auxiliary_fields :: AF # User-specified auxiliary fields for forcing functions and boundary conditions + G, T, V, B, R, F, P, BGC, U, C, Φ, K, AF, Z} <: AbstractModel{TS} + + architecture :: A # Computer `Architecture` on which `Model` is run + grid :: G # Grid of physical points on which `Model` is solved + clock :: Clock{T} # Tracks iteration number and simulation time of `Model` + advection :: V # Advection scheme for tracers + buoyancy :: B # Set of parameters for buoyancy model + coriolis :: R # Set of parameters for the background rotation rate of `Model` + free_surface :: S # Free surface parameters and fields + forcing :: F # Container for forcing functions defined by the user + closure :: E # Diffusive 'turbulence closure' for all model fields + particles :: P # Particle set for Lagrangian tracking + biogeochemistry :: BGC # Biogeochemistry for Oceananigans tracers + velocities :: U # Container for velocity fields `u`, `v`, and `w` + tracers :: C # Container for tracer fields + pressure :: Φ # Container for hydrostatic pressure + diffusivity_fields :: K # Container for turbulent diffusivities + timestepper :: TS # Object containing timestepper fields and parameters + auxiliary_fields :: AF # User-specified auxiliary fields for forcing functions and boundary conditions + vertical_coordinate :: Z # Rulesets that define the time-evolution of the grid end default_free_surface(grid::XYRegularRG; gravitational_acceleration=g_Earth) = @@ -71,7 +72,8 @@ default_free_surface(grid; gravitational_acceleration=g_Earth) = velocities = nothing, pressure = nothing, diffusivity_fields = nothing, - auxiliary_fields = NamedTuple()) + auxiliary_fields = NamedTuple(), + vertical_coordinate = ZCoordinate()) Construct a hydrostatic model with a free surface on `grid`. @@ -103,6 +105,7 @@ Keyword arguments - `pressure`: Hydrostatic pressure field. Default: `nothing`. - `diffusivity_fields`: Diffusivity fields. Default: `nothing`. - `auxiliary_fields`: `NamedTuple` of auxiliary fields. Default: `nothing`. + - `vertical_coordinate`: Rulesets that define the time-evolution of the grid (ZStar/ZCoordinate). Default: `ZCoordinate()`. """ function HydrostaticFreeSurfaceModel(; grid, clock = Clock{eltype(grid)}(time = 0), @@ -121,11 +124,16 @@ function HydrostaticFreeSurfaceModel(; grid, velocities = nothing, pressure = nothing, diffusivity_fields = nothing, - auxiliary_fields = NamedTuple()) + auxiliary_fields = NamedTuple(), + vertical_coordinate = ZCoordinate()) # Check halos and throw an error if the grid's halo is too small @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 `MutableVerticalDiscretization` to allow the use of ZStar (see `MutableVerticalDiscretization`).") + end + # Validate biogeochemistry (add biogeochemical tracers automagically) tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example) biogeochemical_fields = merge(auxiliary_fields, biogeochemical_auxiliary_fields(biogeochemistry)) @@ -208,7 +216,7 @@ function HydrostaticFreeSurfaceModel(; grid, model = HydrostaticFreeSurfaceModel(arch, grid, clock, advection, buoyancy, coriolis, free_surface, forcing, closure, particles, biogeochemistry, velocities, tracers, - pressure, diffusivity_fields, timestepper, auxiliary_fields) + pressure, diffusivity_fields, timestepper, auxiliary_fields, vertical_coordinate) update_state!(model; compute_tendencies = false) diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl index 28cb289e58..9447b3effd 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl @@ -38,6 +38,7 @@ implicitly during time-stepping. diffusivities, hydrostatic_pressure_anomaly, auxiliary_fields, + ztype, clock, forcing) @@ -47,6 +48,7 @@ implicitly during time-stepping. - explicit_barotropic_pressure_x_gradient(i, j, k, grid, free_surface) - x_f_cross_U(i, j, k, grid, coriolis, velocities) - ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure_anomaly) + - grid_slope_contribution_x(i, j, k, grid, buoyancy, ztype, model_fields) - ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy) - immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields) + forcing(i, j, k, grid, clock, model_fields)) @@ -77,6 +79,7 @@ implicitly during time-stepping. diffusivities, hydrostatic_pressure_anomaly, auxiliary_fields, + ztype, clock, forcing) @@ -86,6 +89,7 @@ implicitly during time-stepping. - explicit_barotropic_pressure_y_gradient(i, j, k, grid, free_surface) - y_f_cross_U(i, j, k, grid, coriolis, velocities) - ∂yᶜᶠᶜ(i, j, k, grid, hydrostatic_pressure_anomaly) + - grid_slope_contribution_y(i, j, k, grid, buoyancy, ztype, model_fields) - ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy) - immersed_∂ⱼ_τ₂ⱼ(i, j, k, grid, velocities, v_immersed_bc, closure, diffusivities, clock, model_fields) + forcing(i, j, k, grid, clock, model_fields)) diff --git a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl index 1e44b9bea1..b63b797321 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl @@ -100,6 +100,7 @@ extract_boundary_conditions(::PrescribedVelocityFields) = NamedTuple() free_surface_displacement_field(::PrescribedVelocityFields, ::Nothing, grid) = nothing HorizontalVelocityFields(::PrescribedVelocityFields, grid) = nothing, nothing +materialize_free_surface(::Nothing, velocities, grid) = nothing materialize_free_surface(::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing materialize_free_surface(::ImplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing materialize_free_surface(::SplitExplicitFreeSurface, ::PrescribedVelocityFields, grid) = nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index 00a7dedb47..7ab34ff634 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -84,6 +84,10 @@ function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters = P = model.pressure.pHY′ arch = architecture(grid) + # Update the grid and unscale the tracers + update_grid!(model, grid, model.vertical_coordinate; parameters = w_parameters) + unscale_tracers!(tracers, grid; parameters = w_parameters) + # Advance diagnostic quantities compute_w_from_continuity!(model; parameters = w_parameters) update_hydrostatic_pressure!(P, arch, grid, buoyancy, tracers; parameters = p_parameters) diff --git a/src/Models/HydrostaticFreeSurfaceModels/z_star_vertical_spacing.jl b/src/Models/HydrostaticFreeSurfaceModels/z_star_vertical_spacing.jl new file mode 100644 index 0000000000..9f69f148ca --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/z_star_vertical_spacing.jl @@ -0,0 +1,158 @@ +using Oceananigans.Grids +using Oceananigans.ImmersedBoundaries: MutableGridOfSomeKind + +##### +##### Mutable-specific vertical spacings update +##### + +# The easy case +barotropic_velocities(free_surface::SplitExplicitFreeSurface) = free_surface.barotropic_velocities + +# The "harder" case, barotropic velocities are computed on the fly +barotropic_velocities(free_surface) = nothing, nothing + +# Fallback +update_grid!(model, grid, ztype; parameters) = nothing + +function update_grid!(model, grid::MutableGridOfSomeKind, ::ZStar; parameters = :xy) + + # Scalings and free surface + σᶜᶜ⁻ = grid.z.σᶜᶜ⁻ + σᶜᶜⁿ = grid.z.σᶜᶜⁿ + σᶠᶜⁿ = grid.z.σᶠᶜⁿ + σᶜᶠⁿ = grid.z.σᶜᶠⁿ + σᶠᶠⁿ = grid.z.σᶠᶠⁿ + ∂t_σ = grid.z.∂t_σ + ηⁿ = grid.z.ηⁿ + η = model.free_surface.η + + launch!(architecture(grid), grid, parameters, _update_grid_scaling!, + σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ηⁿ, grid, η) + + # the barotropic velocities are retrieved from the free surface model for a + # SplitExplicitFreeSurface and are calculated for other free surface models + U, V = barotropic_velocities(model.free_surface) + u, v, _ = model.velocities + + # Update the time derivative of the vertical spacing, + # No need to fill the halo as the scaling is updated _IN_ the halos + launch!(architecture(grid), grid, parameters, _update_grid_vertical_velocity!, ∂t_σ, grid, U, V, u, v) + + return nothing +end + +@kernel function _update_grid_scaling!(σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ηⁿ, grid, η) + i, j = @index(Global, NTuple) + k_top = size(grid, 3) + 1 + + hᶜᶜ = static_column_depthᶜᶜᵃ(i, j, grid) + hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) + hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) + hᶠᶠ = static_column_depthᶠᶠᵃ(i, j, grid) + + Hᶜᶜ = column_depthᶜᶜᵃ(i, j, k_top, grid, η) + Hᶠᶜ = column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = column_depthᶜᶠᵃ(i, j, k_top, grid, η) + Hᶠᶠ = column_depthᶠᶠᵃ(i, j, k_top, grid, η) + + @inbounds begin + σᶜᶜ = ifelse(hᶜᶜ == 0, one(grid), Hᶜᶜ / hᶜᶜ) + σᶠᶜ = ifelse(hᶠᶜ == 0, one(grid), Hᶠᶜ / hᶠᶜ) + σᶜᶠ = ifelse(hᶜᶠ == 0, one(grid), Hᶜᶠ / hᶜᶠ) + σᶠᶠ = ifelse(hᶠᶠ == 0, one(grid), Hᶠᶠ / hᶠᶠ) + + # Update previous scaling + σᶜᶜ⁻[i, j, 1] = σᶜᶜⁿ[i, j, 1] + + # update current scaling + σᶜᶜⁿ[i, j, 1] = σᶜᶜ + σᶠᶜⁿ[i, j, 1] = σᶠᶜ + σᶜᶠⁿ[i, j, 1] = σᶜᶠ + σᶠᶠⁿ[i, j, 1] = σᶠᶠ + + # Update η in the grid + ηⁿ[i, j, 1] = η[i, j, k_top] + end +end + +@kernel function _update_grid_vertical_velocity!(∂t_σ, grid, U, V, u, v) + i, j = @index(Global, NTuple) + kᴺ = size(grid, 3) + + hᶜᶜ = static_column_depthᶜᶜᵃ(i, j, grid) + + # ∂(η / H)/∂t = - ∇ ⋅ ∫udz / H + δx_U = δxᶜᶜᶜ(i, j, kᴺ, grid, Δy_qᶠᶜᶜ, barotropic_U, U, u) + δy_V = δyᶜᶜᶜ(i, j, kᴺ, grid, Δx_qᶜᶠᶜ, barotropic_V, V, v) + + δh_U = (δx_U + δy_V) / Azᶜᶜᶜ(i, j, kᴺ, grid) + + @inbounds ∂t_σ[i, j, 1] = ifelse(hᶜᶜ == 0, zero(grid), - δh_U / hᶜᶜ) +end + +# If U and V exist, we just take them +@inline barotropic_U(i, j, k, grid, U, u) = @inbounds U[i, j, k] +@inline barotropic_V(i, j, k, grid, V, v) = @inbounds V[i, j, k] + +# If U and V are not available, we compute them +@inline function barotropic_U(i, j, k, grid, ::Nothing, u) + U = 0 + for k in 1:size(grid, 3) + U += u[i, j, k] * Δzᶠᶜᶜ(i, j, k, grid) + end + return U +end + +@inline function barotropic_V(i, j, k, grid, ::Nothing, v) + V = 0 + for k in 1:size(grid, 3) + V += v[i, j, k] * Δzᶜᶠᶜ(i, j, k, grid) + end + return V +end + +##### +##### ZStar-specific implementation of the additional terms to be included in the momentum equations +##### + +# Fallbacks +@inline grid_slope_contribution_x(i, j, k, grid, buoyancy, ztype, model_fields) = zero(grid) +@inline grid_slope_contribution_y(i, j, k, grid, buoyancy, ztype, model_fields) = zero(grid) + +@inline grid_slope_contribution_x(i, j, k, grid::MutableGridOfSomeKind, ::Nothing, ::ZStar, model_fields) = zero(grid) +@inline grid_slope_contribution_y(i, j, k, grid::MutableGridOfSomeKind, ::Nothing, ::ZStar, model_fields) = zero(grid) + +@inline ∂x_z(i, j, k, grid) = @inbounds ∂xᶠᶜᶜ(i, j, k, grid, znode, Center(), Center(), Center()) +@inline ∂y_z(i, j, k, grid) = @inbounds ∂yᶜᶠᶜ(i, j, k, grid, znode, Center(), Center(), Center()) + +@inline grid_slope_contribution_x(i, j, k, grid::MutableGridOfSomeKind, buoyancy, ::ZStar, model_fields) = + ℑxᶠᵃᵃ(i, j, k, grid, buoyancy_perturbationᶜᶜᶜ, buoyancy.formulation, model_fields) * ∂x_z(i, j, k, grid) + +@inline grid_slope_contribution_y(i, j, k, grid::MutableGridOfSomeKind, buoyancy, ::ZStar, model_fields) = + ℑyᵃᶠᵃ(i, j, k, grid, buoyancy_perturbationᶜᶜᶜ, buoyancy.formulation, model_fields) * ∂y_z(i, j, k, grid) + +#### +#### Removing the scaling of the vertical coordinate from the tracer fields +#### + +const EmptyTuples = Union{NamedTuple{(), Tuple{}}, Tuple{}} + +unscale_tracers!(::EmptyTuples, ::MutableGridOfSomeKind; kwargs...) = nothing + +function unscale_tracers!(tracers, grid::MutableGridOfSomeKind; parameters = :xy) + + for tracer in tracers + launch!(architecture(grid), grid, parameters, _unscale_tracer!, + tracer, grid, Val(grid.Hz), Val(grid.Nz)) + end + + return nothing +end + +@kernel function _unscale_tracer!(tracer, grid, ::Val{Hz}, ::Val{Nz}) where {Hz, Nz} + i, j = @index(Global, NTuple) + + @unroll for k in -Hz+1:Nz+Hz + tracer[i, j, k] /= σⁿ(i, j, k, grid, Center(), Center(), Center()) + end +end diff --git a/src/Models/Models.jl b/src/Models/Models.jl index b7dbdf0702..268927dbe9 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -3,7 +3,7 @@ module Models export NonhydrostaticModel, ShallowWaterModel, ConservativeFormulation, VectorInvariantFormulation, - HydrostaticFreeSurfaceModel, + HydrostaticFreeSurfaceModel, ZStar, ZCoordinate, ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface, PrescribedVelocityFields, PressureField, LagrangianParticles, @@ -105,7 +105,7 @@ using .NonhydrostaticModels: NonhydrostaticModel, PressureField using .HydrostaticFreeSurfaceModels: HydrostaticFreeSurfaceModel, ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface, - PrescribedVelocityFields + PrescribedVelocityFields, ZStar, ZCoordinate using .ShallowWaterModels: ShallowWaterModel, ConservativeFormulation, VectorInvariantFormulation diff --git a/src/Operators/Operators.jl b/src/Operators/Operators.jl index ffd9eeb1e2..7be1572ccf 100644 --- a/src/Operators/Operators.jl +++ b/src/Operators/Operators.jl @@ -77,6 +77,9 @@ export ∂xTᶠᶜᶠ, ∂yTᶜᶠᶠ # Reference frame conversion export intrinsic_vector, extrinsic_vector +# Variable grid operators +export σⁿ, σ⁻, ∂t_σ + using Oceananigans.Grids ##### @@ -119,6 +122,7 @@ include("topology_aware_operators.jl") include("vorticity_operators.jl") include("laplacian_operators.jl") +include("time_variable_grid_operators.jl") include("vector_rotation_operators.jl") @inline xspacing(args...) = Δx(args...) diff --git a/src/Operators/spacings_and_areas_and_volumes.jl b/src/Operators/spacings_and_areas_and_volumes.jl index a7df404aca..9f8f866a3e 100644 --- a/src/Operators/spacings_and_areas_and_volumes.jl +++ b/src/Operators/spacings_and_areas_and_volumes.jl @@ -35,17 +35,6 @@ The operators in this file fall into three categories: ##### ##### -const ZRG = Union{LLGZ, RGZ, OSSGZ} - -@inline getspacing(k, Δz::AbstractVector) = @inbounds Δz[k] -@inline getspacing(k, Δz::Number) = @inbounds Δz - -@inline Δrᵃᵃᶜ(i, j, k, grid::AbstractGrid) = getspacing(k, grid.z.Δᵃᵃᶜ) -@inline Δrᵃᵃᶠ(i, j, k, grid::AbstractGrid) = getspacing(k, grid.z.Δᵃᵃᶠ) - -@inline Δzᵃᵃᶜ(i, j, k, grid::AbstractGrid) = Δrᵃᵃᶜ(i, j, k, grid) -@inline Δzᵃᵃᶠ(i, j, k, grid::AbstractGrid) = Δrᵃᵃᶠ(i, j, k, grid) - # This metaprogramming loop defines all possible combinations of locations for spacings. # The general 2D and 3D spacings are reconducted to their one - dimensional counterparts. # Grids that do not have a specific one - dimensional spacing for a given location need to @@ -114,11 +103,14 @@ end ##### One - dimensional Vertical spacing (same for all grids) ##### -@inline Δzᵃᵃᶜ(i, j, k, grid) = @inbounds grid.z.Δᵃᵃᶜ[k] -@inline Δzᵃᵃᶠ(i, j, k, grid) = @inbounds grid.z.Δᵃᵃᶠ[k] +@inline getspacing(k, Δz::Number) = Δz +@inline getspacing(k, Δz::AbstractVector) = @inbounds Δz[k] + +@inline Δrᵃᵃᶜ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶜ) +@inline Δrᵃᵃᶠ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶠ) -@inline Δzᵃᵃᶜ(i, j, k, grid::ZRG) = grid.z.Δᵃᵃᶜ -@inline Δzᵃᵃᶠ(i, j, k, grid::ZRG) = grid.z.Δᵃᵃᶠ +@inline Δzᵃᵃᶜ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶜ) +@inline Δzᵃᵃᶠ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶠ) ##### ##### diff --git a/src/Operators/time_variable_grid_operators.jl b/src/Operators/time_variable_grid_operators.jl new file mode 100644 index 0000000000..0dbdf1a959 --- /dev/null +++ b/src/Operators/time_variable_grid_operators.jl @@ -0,0 +1,57 @@ +import Oceananigans.Grids: znode, AbstractMutableGrid + +##### +##### MutableVerticalDiscretization-specific vertical spacing functions +##### + +const C = Center +const F = Face + +const AMG = AbstractMutableGrid + +# Fallbacks +@inline σⁿ(i, j, k, grid, ℓx, ℓy, ℓz) = one(grid) +@inline σ⁻(i, j, k, grid, ℓx, ℓy, ℓz) = one(grid) + +@inline σⁿ(i, j, k, grid::AMG, ::C, ::C, ℓz) = @inbounds grid.z.σᶜᶜⁿ[i, j, 1] +@inline σⁿ(i, j, k, grid::AMG, ::F, ::C, ℓz) = @inbounds grid.z.σᶠᶜⁿ[i, j, 1] +@inline σⁿ(i, j, k, grid::AMG, ::C, ::F, ℓz) = @inbounds grid.z.σᶜᶠⁿ[i, j, 1] +@inline σⁿ(i, j, k, grid::AMG, ::F, ::F, ℓz) = @inbounds grid.z.σᶠᶠⁿ[i, j, 1] + +# σ⁻ is needed only at centers +@inline σ⁻(i, j, k, grid::AMG, ::C, ::C, ℓz) = @inbounds grid.z.σᶜᶜ⁻[i, j, 1] + +@inline ∂t_σ(i, j, k, grid) = zero(grid) +@inline ∂t_σ(i, j, k, grid::AMG) = @inbounds grid.z.∂t_σ[i, j, 1] + +#### +#### Vertical spacing functions +#### + +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 + +for LX in (:ᶠ, :ᶜ), LY in (:ᶠ, :ᶜ), LZ in (:ᶠ, :ᶜ) + zspacing = Symbol(:Δz, LX, LY, LZ) + rspacing = Symbol(:Δr, LX, LY, LZ) + + ℓx = superscript_location(LX) + ℓy = superscript_location(LY) + ℓz = superscript_location(LZ) + + @eval begin + @inline $zspacing(i, j, k, grid::MRG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::MLLG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::MOSG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + end +end + +# znode for an AbstractMutableGrid grid is the reference node (`rnode`) scaled by the derivative with respect to the reference (σⁿ) +# added to the surface value of `z` (which we here call ηⁿ) +@inline znode(i, j, k, grid::AMG, ::C, ::C, ℓz) = rnode(i, j, k, grid, C(), C(), ℓz) * σⁿ(i, j, k, grid, C(), C(), ℓz) + @inbounds grid.z.ηⁿ[i, j, 1] +@inline znode(i, j, k, grid::AMG, ::F, ::C, ℓz) = rnode(i, j, k, grid, F(), C(), ℓz) * σⁿ(i, j, k, grid, F(), C(), ℓz) + ℑxᶠᵃᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline znode(i, j, k, grid::AMG, ::C, ::F, ℓz) = rnode(i, j, k, grid, C(), F(), ℓz) * σⁿ(i, j, k, grid, C(), F(), ℓz) + ℑyᵃᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline znode(i, j, k, grid::AMG, ::F, ::F, ℓz) = rnode(i, j, k, grid, F(), F(), ℓz) * σⁿ(i, j, k, grid, F(), F(), ℓz) + ℑxyᶠᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) diff --git a/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl b/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl index 86795b9e24..3aaa551622 100644 --- a/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl +++ b/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl @@ -46,7 +46,7 @@ const c = Center() const f = Face() # The vertical spacing used here is Δz for velocities and Δr for tracers, since the -# implicit solver operator is applied to the scaled tracer e₃θ instead of just θ +# implicit solver operator is applied to the scaled tracer σθ instead of just θ @inline vertical_spacing(i, j, k, grid, ℓx, ℓy, ℓz) = Δz(i, j, k, grid, ℓx, ℓy, ℓz) @inline vertical_spacing(i, j, k, grid, ::Center, ::Center, ℓz) = Δr(i, j, k, grid, c, c, ℓz) diff --git a/test/runtests.jl b/test/runtests.jl index a47f0e3d3d..fbfbe6dac7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -41,7 +41,7 @@ CUDA.allowscalar() do CUDA.versioninfo() catch; end end - + # Core Oceananigans if group == :unit || group == :all @testset "Unit tests" begin @@ -120,6 +120,7 @@ CUDA.allowscalar() do @testset "Model and time stepping tests (part 1)" begin include("test_nonhydrostatic_models.jl") include("test_time_stepping.jl") + include("test_active_cells_map.jl") end end @@ -219,6 +220,12 @@ CUDA.allowscalar() do end end + if group == :vertical_coordinate || group == :all + @testset "Vertical coordinate tests" begin + include("test_zstar_coordinate.jl") + end + end + # Tests for Enzyme extension if group == :enzyme || group == :all @testset "Enzyme extension tests" begin diff --git a/test/test_active_cells_map.jl b/test/test_active_cells_map.jl new file mode 100644 index 0000000000..d1533b9451 --- /dev/null +++ b/test/test_active_cells_map.jl @@ -0,0 +1,124 @@ +include("dependencies_for_runtests.jl") + +using Oceananigans.Operators: hack_cosd +using Oceananigans.ImmersedBoundaries: retrieve_surface_active_cells_map, + retrieve_interior_active_cells_map, + immersed_cell + +function Δ_min(grid) + Δx_min = minimum_xspacing(grid, Center(), Center(), Center()) + Δy_min = minimum_yspacing(grid, Center(), Center(), Center()) + return min(Δx_min, Δy_min) +end + +@inline Gaussian(x, y, L) = exp(-(x^2 + y^2) / L^2) + +function solid_body_rotation_test(grid) + + free_surface = SplitExplicitFreeSurface(grid; substeps = 10, gravitational_acceleration = 1) + coriolis = HydrostaticSphericalCoriolis(rotation_rate = 1) + + model = HydrostaticFreeSurfaceModel(; grid, + momentum_advection = VectorInvariant(), + free_surface = free_surface, + coriolis = coriolis, + tracers = :c, + tracer_advection = WENO(), + buoyancy = nothing, + closure = nothing) + + g = model.free_surface.gravitational_acceleration + R = grid.radius + Ω = model.coriolis.rotation_rate + + uᵢ(λ, φ, z) = 0.1 * cosd(φ) * sind(λ) + ηᵢ(λ, φ, z) = (R * Ω * 0.1 + 0.1^2 / 2) * sind(φ)^2 / g * sind(λ) + + # Gaussian leads to values with O(1e-60), + # too small for repetible testing. We cap it at 0.1 + cᵢ(λ, φ, z) = max(Gaussian(λ, φ - 5, 10), 0.1) + vᵢ(λ, φ, z) = 0.1 + + set!(model, u=uᵢ, η=ηᵢ, c=cᵢ) + + Δt = 0.1 * Δ_min(grid) / sqrt(g * grid.Lz) + + simulation = Simulation(model; Δt, stop_iteration = 10) + run!(simulation) + + return merge(model.velocities, model.tracers, (; η = model.free_surface.η)) +end + +Nx = 16 +Ny = 16 +Nz = 10 + +@testset "Active cells map" begin + for arch in archs + underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, Nz), + halo = (4, 4, 4), + latitude = (-80, 80), + longitude = (-160, 160), + z = (-10, 0), + radius = 1, + topology = (Bounded, Bounded, Bounded)) + + # Make sure the bottom is the same + bottom_height = zeros(Nx, Ny) + for i in 1:Nx, j in 1:Ny + bottom_height[i, j] = - rand() * 5 - 5 + end + + bottom_height = on_architecture(arch, bottom_height) + immersed_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height)) + immersed_active_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height); active_cells_map = true) + + @testset "Active cells map construction" begin + surface_active_cells_map = retrieve_surface_active_cells_map(immersed_active_grid) + interior_active_cells_map = retrieve_interior_active_cells_map(immersed_active_grid, Val(:interior)) + + surface_active_cells_map = on_architecture(CPU(), surface_active_cells_map) + interior_active_cells_map = on_architecture(CPU(), interior_active_cells_map) + grid = on_architecture(CPU(), immersed_grid) + + for i in 1:Nx, j in 1:Ny, k in 1:Nz + immersed = immersed_cell(i, j, k, grid) + active = (i, j, k) ∈ interior_active_cells_map + @test immersed ⊻ active + end + + for i in 1:Nx, j in 1:Ny + immersed = all(immersed_cell(i, j, k, grid) for k in 1:Nz) + active = (i, j) ∈ surface_active_cells_map + @test immersed ⊻ active + end + end + + @testset "Active cells map solid body rotation" begin + + ua, va, wa, ca, ηa = solid_body_rotation_test(immersed_active_grid) + u, v, w, c, η = solid_body_rotation_test(immersed_grid) + + ua = interior(on_architecture(CPU(), ua)) + va = interior(on_architecture(CPU(), va)) + wa = interior(on_architecture(CPU(), wa)) + ca = interior(on_architecture(CPU(), ca)) + ηa = interior(on_architecture(CPU(), ηa)) + + u = interior(on_architecture(CPU(), u)) + v = interior(on_architecture(CPU(), v)) + w = interior(on_architecture(CPU(), w)) + c = interior(on_architecture(CPU(), c)) + η = interior(on_architecture(CPU(), η)) + + atol = eps(eltype(immersed_grid)) + rtol = sqrt(eps(eltype(immersed_grid))) + + @test all(isapprox(u, ua; atol, rtol)) + @test all(isapprox(v, va; atol, rtol)) + @test all(isapprox(w, wa; atol, rtol)) + @test all(isapprox(c, ca; atol, rtol)) + @test all(isapprox(η, ηa; atol, rtol)) + end + end +end \ No newline at end of file diff --git a/test/test_split_explicit_vertical_integrals.jl b/test/test_split_explicit_vertical_integrals.jl index 045531ed67..62f9acf6f6 100644 --- a/test/test_split_explicit_vertical_integrals.jl +++ b/test/test_split_explicit_vertical_integrals.jl @@ -35,7 +35,9 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces @testset "Average to zero" begin # set equal to something else - η̅ .= U̅ .= V̅ .= 1.0 + η̅ .= 1 + U̅ .= 1 + V̅ .= 1 # now set equal to zero initialize_free_surface_state!(sefs, sefs.timestepper, sefs.timestepper, Val(1)) @@ -47,8 +49,8 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces # check @test all(Array(η̅.data.parent) .== 0.0) - @test all(Array(U̅.data.parent .== 0.0)) - @test all(Array(V̅.data.parent .== 0.0)) + @test all(Array(U̅.data.parent) .== 0.0) + @test all(Array(V̅.data.parent) .== 0.0) end @testset "Inexact integration" begin @@ -61,7 +63,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces set!(u, set_u_check) exact_U = similar(U) set!(exact_U, set_U_check) - compute_barotropic_mode!(U, V, grid, u, v) + compute_barotropic_mode!(U, V, grid, u, v, η̅) tolerance = 1e-3 @test all((Array(interior(U) .- interior(exact_U))) .< tolerance) @@ -70,7 +72,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces set!(v, set_v_check) exact_V = similar(V) set!(exact_V, set_V_check) - compute_barotropic_mode!(U, V, grid, u, v) + compute_barotropic_mode!(U, V, grid, u, v, η̅) @test all((Array(interior(V) .- interior(exact_V))) .< tolerance) end @@ -78,14 +80,14 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces Δz = zeros(Nz) Δz .= grid.z.Δᵃᵃᶜ - u .= 0.0 - U .= 1.0 - compute_barotropic_mode!(U, V, grid, u, v) - @test all(Array(U.data.parent) .== 0.0) + set!(u, 0) + set!(U, 1) + compute_barotropic_mode!(U, V, grid, u, v, η̅) + @test all(Array(interior(U)) .== 0.0) - u .= 1.0 - U .= 1.0 - compute_barotropic_mode!(U, V, grid, u, v) + set!(u, 1) + set!(U, 1) + compute_barotropic_mode!(U, V, grid, u, v, η̅) @test all(Array(interior(U)) .≈ Lz) set_u_check(x, y, z) = sin(x) @@ -93,7 +95,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces set!(u, set_u_check) exact_U = similar(U) set!(exact_U, set_U_check) - compute_barotropic_mode!(U, V, grid, u, v) + compute_barotropic_mode!(U, V, grid, u, v, η̅) @test all(Array(interior(U)) .≈ Array(interior(exact_U))) set_v_check(x, y, z) = sin(x) * z * cos(y) @@ -101,7 +103,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces set!(v, set_v_check) exact_V = similar(V) set!(exact_V, set_V_check) - compute_barotropic_mode!(U, V, grid, u, v) + compute_barotropic_mode!(U, V, grid, u, v, η̅) @test all(Array(interior(V)) .≈ Array(interior(exact_V))) end diff --git a/test/test_zstar_coordinate.jl b/test/test_zstar_coordinate.jl new file mode 100644 index 0000000000..d60f75fb88 --- /dev/null +++ b/test/test_zstar_coordinate.jl @@ -0,0 +1,201 @@ +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) + + bᵢ = deepcopy(model.tracers.b) + cᵢ = deepcopy(model.tracers.c) + + ∫bᵢ = Field(Integral(bᵢ)) + ∫cᵢ = Field(Integral(cᵢ)) + w = model.velocities.w + Nz = model.grid.Nz + + for _ in 1:Ni + time_step!(model, Δt) + end + + ∫b = Field(Integral(model.tracers.b)) + ∫c = Field(Integral(model.tracers.c)) + + # Testing that: + # (1) tracers are conserved down to machine precision + # (2) vertical velocities are zero at the top surface + @test interior(∫b, 1, 1, 1) ≈ interior(∫bᵢ, 1, 1, 1) + @test interior(∫c, 1, 1, 1) ≈ interior(∫cᵢ, 1, 1, 1) + @test maximum(abs, interior(w, :, :, Nz+1)) < eps(eltype(w)) + + return nothing +end + +function info_message(grid, free_surface) + msg1 = "$(architecture(grid)) " + msg2 = string(getnamewrapper(grid)) + msg3 = grid isa ImmersedBoundaryGrid ? " on a " * string(getnamewrapper(grid.underlying_grid)) : "" + msg4 = grid.z.Δᵃᵃᶠ isa Number ? " with uniform spacing" : " with stretched spacing" + msg5 = grid isa ImmersedBoundaryGrid ? " and $(string(getnamewrapper(grid.immersed_boundary))) immersed boundary" : "" + msg6 = " using a " * string(getnamewrapper(free_surface)) + return msg1 * msg2 * msg3 * msg4 * msg5 * msg6 +end + +const C = Center +const F = Face + +@testset "ZStar coordinate scaling tests" begin + @info "testing the ZStar coordinate scalings" + + z = MutableVerticalDiscretization((-20, 0)) + + grid = RectilinearGrid(size = (2, 2, 20), + x = (0, 2), + y = (0, 1), + z = z, + topology = (Bounded, Periodic, Bounded)) + + grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y) -> -10)) + + model = HydrostaticFreeSurfaceModel(; grid, + free_surface = SplitExplicitFreeSurface(grid; substeps = 20), + vertical_coordinate = ZStar()) + + @test znode(1, 1, 21, grid, C(), C(), F()) == 0 + @test column_depthᶜᶜᵃ(1, 1, grid) == 10 + @test static_column_depthᶜᶜᵃ(1, 1, grid) == 10 + + set!(model, η = [1 1; 2 2]) + set!(model, u = (x, y, z) -> x, v = (x, y, z) -> y) + update_state!(model) + + @test σⁿ(1, 1, 1, grid, C(), C(), C()) == 11 / 10 + @test σⁿ(2, 1, 1, grid, C(), C(), C()) == 12 / 10 + + @test znode(1, 1, 21, grid, C(), C(), F()) == 1 + @test znode(2, 1, 21, grid, C(), C(), F()) == 2 + @test rnode(1, 1, 21, grid, C(), C(), F()) == 0 + @test column_depthᶜᶜᵃ(1, 1, grid) == 11 + @test column_depthᶜᶜᵃ(2, 1, grid) == 12 + @test static_column_depthᶜᶜᵃ(1, 1, grid) == 10 + @test static_column_depthᶜᶜᵃ(2, 1, grid) == 10 +end + +@testset "MutableVerticalDiscretization tests" begin + @info "testing the MutableVerticalDiscretization in ZCoordinate mode" + + 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 = MutableVerticalDiscretization((-20, 0)) + z_stretched = MutableVerticalDiscretization(collect(-20:0)) + topologies = ((Periodic, Periodic, Bounded), + (Periodic, Bounded, Bounded), + (Bounded, Periodic, Bounded), + (Bounded, Bounded, Bounded)) + + for arch in archs + for topology in topologies + Random.seed!(1234) + + rtg = RectilinearGrid(arch; size = (10, 10, 20), x = (0, 100kilometers), y = (-10kilometers, 10kilometers), topology, z = z_uniform) + rtgv = RectilinearGrid(arch; size = (10, 10, 20), x = (0, 100kilometers), y = (-10kilometers, 10kilometers), topology, z = z_stretched) + + irtg = ImmersedBoundaryGrid(rtg, GridFittedBottom((x, y) -> rand() - 10)) + irtgv = ImmersedBoundaryGrid(rtgv, GridFittedBottom((x, y) -> rand() - 10)) + prtg = ImmersedBoundaryGrid(rtg, PartialCellBottom((x, y) -> rand() - 10)) + prtgv = ImmersedBoundaryGrid(rtgv, PartialCellBottom((x, y) -> rand() - 10)) + + if topology[2] == Bounded + llg = LatitudeLongitudeGrid(arch; size = (10, 10, 20), latitude = (0, 1), longitude = (0, 1), topology, z = z_uniform) + llgv = LatitudeLongitudeGrid(arch; size = (10, 10, 20), latitude = (0, 1), longitude = (0, 1), topology, z = z_stretched) + + illg = ImmersedBoundaryGrid(llg, GridFittedBottom((x, y) -> rand() - 10)) + illgv = ImmersedBoundaryGrid(llgv, GridFittedBottom((x, y) -> rand() - 10)) + pllg = ImmersedBoundaryGrid(llg, PartialCellBottom((x, y) -> rand() - 10)) + pllgv = ImmersedBoundaryGrid(llgv, PartialCellBottom((x, y) -> rand() - 10)) + + # TODO: Partial cell bottom are broken at the moment and do not account for the Δz in the volumes + # and vertical areas (see https://github.com/CliMA/Oceananigans.jl/issues/3958) + # When this is issue is fixed we can add the partial cells to the testing. + grids = [llg, rtg, llgv, rtgv, illg, irtg, illgv, irtgv] # , pllg, prtg, pllgv, prtgv] + else + grids = [rtg, rtgv, irtg, irtgv] #, prtg, prtgv] + end + + for grid in grids + split_free_surface = SplitExplicitFreeSurface(grid; cfl = 0.75) + implicit_free_surface = ImplicitFreeSurface() + explicit_free_surface = ExplicitFreeSurface() + + for free_surface in [split_free_surface, implicit_free_surface, explicit_free_surface] + + # TODO: There are parameter space issues with ImplicitFreeSurface and a immersed LatitudeLongitudeGrid + # For the moment we are skipping these tests. + if (arch isa GPU) && + (free_surface isa ImplicitFreeSurface) && + (grid isa ImmersedBoundaryGrid) && + (grid.underlying_grid isa LatitudeLongitudeGrid) + + @info " Skipping $(info_message(grid, free_surface)) because of parameter space issues" + continue + end + + info_msg = info_message(grid, free_surface) + @testset "$info_msg" begin + @info " Testing a $info_msg" + model = HydrostaticFreeSurfaceModel(; grid, + free_surface, + tracers = (:b, :c), + buoyancy = BuoyancyTracer(), + vertical_coordinate = ZStar()) + + bᵢ(x, y, z) = x < grid.Lx / 2 ? 0.06 : 0.01 + + set!(model, c = (x, y, z) -> rand(), b = bᵢ) + + Δt = free_surface isa ExplicitFreeSurface ? 10 : 2minutes + test_zstar_coordinate(model, 100, Δt) + end + end + end + end + end +end \ No newline at end of file diff --git a/validation/z_star_coordinate/lock_release.jl b/validation/z_star_coordinate/lock_release.jl new file mode 100644 index 0000000000..d4c625ff1f --- /dev/null +++ b/validation/z_star_coordinate/lock_release.jl @@ -0,0 +1,80 @@ +using Oceananigans +using Oceananigans.Grids +using Oceananigans.Units +using Oceananigans.Utils: prettytime +using Oceananigans.Advection: WENOVectorInvariant +using Oceananigans.AbstractOperations: GridMetricOperation +using Printf + +z_faces = MutableVerticalDiscretization(-20, 0) + +grid = RectilinearGrid(size = (128, 20), + x = (0, 64kilometers), + z = z_faces, + halo = (6, 6), + topology = (Bounded, Flat, Bounded)) + +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(x -> - (64kilometers - x) / 64kilometers * 20)) + +model = HydrostaticFreeSurfaceModel(; grid, + momentum_advection = WENO(order = 5), + tracer_advection = WENO(order = 5), + buoyancy = BuoyancyTracer(), + closure = nothing, + tracers = :b, + free_surface = SplitExplicitFreeSurface(grid; substeps = 10)) + +g = model.free_surface.gravitational_acceleration + +bᵢ(x, z) = x < 32kilometers ? 0.06 : 0.01 + +set!(model, b = bᵢ) + +Δt = 10 + +@info "the time step is $Δt" + +simulation = Simulation(model; Δt, stop_iteration = 10000, stop_time = 17hours) + +Δz = zspacings(grid, Center(), Center(), Center()) +∫b_init = sum(model.tracers.b * Δz) / sum(Δz) + +field_outputs = merge(model.velocities, model.tracers, (; Δz)) + +simulation.output_writers[:other_variables] = JLD2OutputWriter(model, field_outputs, + overwrite_existing = true, + schedule = IterationInterval(100), + filename = "zstar_model") + +function progress(sim) + w = interior(sim.model.velocities.w, :, :, sim.model.grid.Nz+1) + u = sim.model.velocities.u + ∫b = sum(model.tracers.b * Δz) / sum(Δz) + + msg0 = @sprintf("Time: %s iteration %d ", prettytime(sim.model.clock.time), sim.model.clock.iteration) + msg1 = @sprintf("extrema w: %.2e %.2e ", maximum(w), minimum(w)) + msg2 = @sprintf("extrema u: %.2e %.2e ", maximum(u), minimum(u)) + msg3 = @sprintf("drift b: %6.3e ", ∫b - ∫b_init) + msg4 = @sprintf("extrema Δz: %.2e %.2e ", maximum(Δz), minimum(Δz)) + @info msg0 * msg1 * msg2 * msg3 * msg4 + + return nothing +end + +simulation.callbacks[:progress] = Callback(progress, IterationInterval(100)) + +run!(simulation) + +using Oceananigans.Fields: OneField + +# # Check tracer conservation +b = FieldTimeSeries("zstar_model.jld2", "b") +dz = FieldTimeSeries("zstar_model.jld2", "Δz") + +init = sum(dz[1] * b[1]) / sum(dz[1]) +drift = [] + +for t in 1:length(b.times) + push!(drift, sum(dz[t] * b[t]) / sum(dz[t]) - init) +end +