diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index d8dc9987da..00167e1a10 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -661,6 +661,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 cec2f85661..093b3e7a42 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.94.3" +version = "0.95.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/examples/langmuir_turbulence.jl b/examples/langmuir_turbulence.jl index 9886c2fca9..6824ddf484 100644 --- a/examples/langmuir_turbulence.jl +++ b/examples/langmuir_turbulence.jl @@ -313,7 +313,7 @@ Vₙ = @lift view(time_series.V[$n], 1, 1, :) wuₙ = @lift view(time_series.wu[$n], 1, 1, :) wvₙ = @lift view(time_series.wv[$n], 1, 1, :) -k = searchsortedfirst(grid.zᵃᵃᶠ[:], -8) +k = searchsortedfirst(znodes(grid, Center(); with_halos=true), -8) wxyₙ = @lift view(time_series.w[$n], :, :, k) wxzₙ = @lift view(time_series.w[$n], :, 1, :) uxzₙ = @lift view(time_series.u[$n], :, 1, :) diff --git a/src/AbstractOperations/grid_metrics.jl b/src/AbstractOperations/grid_metrics.jl index 4205135ef0..9fbfb79570 100644 --- a/src/AbstractOperations/grid_metrics.jl +++ b/src/AbstractOperations/grid_metrics.jl @@ -290,6 +290,7 @@ end @inline zspacings(field::AbstractField) = zspacings(field.grid, location(field)...) @inline λspacings(field::AbstractField) = λspacings(field.grid, location(field)...) @inline φspacings(field::AbstractField) = φspacings(field.grid, location(field)...) +@inline rspacings(field::AbstractField) = rspacings(field.grid, location(field)...) # Some defaults for e.g. easy CFL computations. @inline xspacings(grid::AbstractGrid) = xspacings(grid, Center(), Center(), Center()) @@ -297,3 +298,4 @@ end @inline zspacings(grid::AbstractGrid) = zspacings(grid, Center(), Center(), Center()) @inline λspacings(grid::AbstractGrid) = λspacings(grid, Center(), Center(), Center()) @inline φspacings(grid::AbstractGrid) = φspacings(grid, Center(), Center(), Center()) +@inline rspacings(grid::AbstractGrid) = rspacings(grid, Center(), Center(), Center()) diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl index 6473915aed..d813b6d7b9 100644 --- a/src/Advection/Advection.jl +++ b/src/Advection/Advection.jl @@ -34,6 +34,7 @@ using Oceananigans.Grids: with_halo, coordinates using Oceananigans.Architectures: architecture, CPU using Oceananigans.Operators +using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, Γᶠᶠᶜ, ∂t_e₃ 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..7f82a509a7 100644 --- a/src/Advection/vector_invariant_cross_upwinding.jl +++ b/src/Advection/vector_invariant_cross_upwinding.jl @@ -17,20 +17,37 @@ ##### Cross and Self Upwinding of the Divergence flux ##### +# This is the additional term appearing in the continuity equation due to a moving grid. +# The discrete divergence is calculated as: +# +# wᵏ⁺¹ - wᵏ δx(Ax u) + δy(Ay v) Δrᶜᶜᶜ ∂t_e₃ +# ---------- = - --------------------- - ------------- +# Δzᶜᶜᶜ Vᶜᶜᶜ Δzᶜᶜᶜ +# +# We upwind with 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_e₃ +# +# which represents the static volume times the time derivative of the vertical grid scaling. +@inline V_times_∂t_e₃(i, j, k, grid) = Azᶜᶜᶜ(i, j, k, grid) * Δrᶜᶜᶜ(i, j, k, grid) * ∂t_e₃(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) + ∂ts = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, cross_scheme, V_times_∂t_e₃) - return û * δᴿ + return û * (δᴿ + ∂ts) 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) + ∂ts = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, V_times_∂t_e₃) - return v̂ * δᴿ + return v̂ * (δᴿ + ∂ts) end diff --git a/src/Advection/vector_invariant_self_upwinding.jl b/src/Advection/vector_invariant_self_upwinding.jl index e8ee512cfd..c530fbc4b1 100644 --- a/src/Advection/vector_invariant_self_upwinding.jl +++ b/src/Advection/vector_invariant_self_upwinding.jl @@ -2,15 +2,18 @@ ##### 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) + +@inline δx_U_plus_metric(i, j, k, grid, u, v) = δxᶜᶜᶜ(i, j, k, grid, Ax_qᶠᶜᶜ, u) + V_times_∂t_e₃(i, j, k, grid) +@inline δy_V_plus_metric(i, j, k, grid, u, v) = δyᶜᶜᶜ(i, j, k, grid, Ay_qᶜᶠᶜ, v) + V_times_∂t_e₃(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) + V_times_∂t_e₃(i, j, k, grid) @inline function upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme::VectorInvariantSelfVerticalUpwinding, u, v) @@ -18,7 +21,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_metric, 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 +33,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_metric, 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/DistributedComputations/distributed_grids.jl b/src/DistributedComputations/distributed_grids.jl index 1b52477149..4fae2d30f4 100644 --- a/src/DistributedComputations/distributed_grids.jl +++ b/src/DistributedComputations/distributed_grids.jl @@ -13,11 +13,11 @@ import Oceananigans.Grids: RectilinearGrid, LatitudeLongitudeGrid, with_halo const DistributedGrid{FT, TX, TY, TZ} = AbstractGrid{FT, TX, TY, TZ, <:Distributed} -const DistributedRectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ} = - RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, <:Distributed} where {FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ} +const DistributedRectilinearGrid{FT, TX, TY, TZ, CZ, FX, FY, VX, VY} = + RectilinearGrid{FT, TX, TY, TZ, CZ, FX, FY, VX, VY, <:Distributed} where {FT, TX, TY, TZ, CZ, FX, FY, VX, VY} -const DistributedLatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ} = - LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, <:Distributed} where {FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ} +const DistributedLatitudeLongitudeGrid{FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY} = + LatitudeLongitudeGrid{FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY, <:Distributed} where {FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY} # Local size from global size and architecture local_size(arch::Distributed, global_sz) = (local_size(global_sz[1], arch.partition.x, arch.local_index[1]), @@ -81,12 +81,15 @@ function RectilinearGrid(arch::Distributed, TY = insert_connected_topology(topology[2], Ry, rj) TZ = insert_connected_topology(topology[3], Rz, rk) + local_topo = (TX, TY, TZ) + xl = Rx == 1 ? x : partition_coordinate(x, nx, arch, 1) yl = Ry == 1 ? y : partition_coordinate(y, ny, arch, 2) zl = Rz == 1 ? z : partition_coordinate(z, nz, arch, 3) - Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology[1](), nx, Hx, xl, :x, child_architecture(arch)) - Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology[2](), ny, Hy, yl, :y, child_architecture(arch)) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topology[3](), nz, Hz, zl, :z, child_architecture(arch)) + + Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, local_topo, local_sz, halo, xl, :x, 1, child_architecture(arch)) + Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, local_topo, local_sz, halo, yl, :y, 2, child_architecture(arch)) + Lz, z = generate_coordinate(FT, local_topo, local_sz, halo, zl, :z, 3, child_architecture(arch)) return RectilinearGrid{TX, TY, TZ}(arch, nx, ny, nz, @@ -94,7 +97,7 @@ function RectilinearGrid(arch::Distributed, Lx, Ly, Lz, Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, Δyᵃᶜᵃ, Δyᵃᶠᵃ, yᵃᶠᵃ, yᵃᶜᵃ, - Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) + z) end """ @@ -127,6 +130,8 @@ function LatitudeLongitudeGrid(arch::Distributed, TY = insert_connected_topology(topology[2], Ry, rj) TZ = insert_connected_topology(topology[3], Rz, rk) + local_topo = (TX, TY, TZ) + λl = Rx == 1 ? longitude : partition_coordinate(longitude, nλ, arch, 1) φl = Ry == 1 ? latitude : partition_coordinate(latitude, nφ, arch, 2) zl = Rz == 1 ? z : partition_coordinate(z, nz, arch, 3) @@ -134,8 +139,8 @@ function LatitudeLongitudeGrid(arch::Distributed, # Calculate all direction (which might be stretched) # A direction is regular if the domain passed is a Tuple{<:Real, <:Real}, # it is stretched if being passed is a function or vector (as for the VerticallyStretchedRectilinearGrid) - Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), nλ, Hλ, λl, :longitude, arch.child_architecture) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), nz, Hz, zl, :z, arch.child_architecture) + Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, local_topo, local_sz, halo, λl, :longitude, 1, child_architecture(arch)) + Lz, z = generate_coordinate(FT, local_topo, local_sz, halo, zl, :z, 3, child_architecture(arch)) # The Latitudinal direction is _special_: # precompute_metrics assumes that `length(φᵃᶠᵃ) = length(φᵃᶜᵃ) + 1`, which is always the case in a @@ -145,7 +150,8 @@ function LatitudeLongitudeGrid(arch::Distributed, # we disregard the topology when constructing the metrics and add a halo point! # Furthermore, the `LatitudeLongitudeGrid` requires an extra halo on it's latitudinal coordinate to allow calculating # the z-area on halo cells. (see: Az = R^2 * Δλ * (sin(φ[j]) - sin(φ[j-1])) - Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, Bounded(), nφ, Hφ + 1, φl, :latitude, arch.child_architecture) + Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, Bounded(), nφ, Hφ + 1, φl, :latitude, child_architecture(arch)) + preliminary_grid = LatitudeLongitudeGrid{TX, TY, TZ}(arch, nλ, nφ, nz, @@ -153,7 +159,7 @@ function LatitudeLongitudeGrid(arch::Distributed, Lλ, Lφ, Lz, Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, - Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, + z, (nothing for i=1:10)..., convert(FT, radius)) return !precompute_metrics ? preliminary_grid : with_precomputed_metrics(preliminary_grid) @@ -199,7 +205,7 @@ function reconstruct_global_grid(grid::DistributedRectilinearGrid) Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, TX(), Nx, Hx, xG, :x, child_arch) Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, TY(), Ny, Hy, yG, :y, child_arch) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, zG, :z, child_arch) + Lz, z = generate_coordinate(FT, (TX, TY, TZ), (Nx, Ny, Nz), H, zG, :z, 3, child_arch) return RectilinearGrid{TX, TY, TZ}(child_arch, Nx, Ny, Nz, @@ -207,7 +213,7 @@ function reconstruct_global_grid(grid::DistributedRectilinearGrid) Lx, Ly, Lz, Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ, yᵃᶠᵃ, yᵃᶜᵃ, - Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) + z) end function reconstruct_global_grid(grid::DistributedLatitudeLongitudeGrid) @@ -245,7 +251,7 @@ function reconstruct_global_grid(grid::DistributedLatitudeLongitudeGrid) # it is stretched if being passed is a function or vector Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), Nλ, Hλ, λG, :longitude, child_arch) Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY(), Nφ, Hφ, φG, :latitude, child_arch) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, zG, :z, child_arch) + Lz, z = generate_coordinate(FT, (TX, TY, TZ), (Nλ, Nφ, Nz), H, zG, :z, 3, child_arch) precompute_metrics = metrics_precomputed(grid) @@ -255,7 +261,7 @@ function reconstruct_global_grid(grid::DistributedLatitudeLongitudeGrid) Lλ, Lφ, Lz, Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, - Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, + z, (nothing for i=1:10)..., grid.radius) return !precompute_metrics ? preliminary_grid : with_precomputed_metrics(preliminary_grid) diff --git a/src/DistributedComputations/distributed_on_architecture.jl b/src/DistributedComputations/distributed_on_architecture.jl index f4229b59f2..75b936fb9f 100644 --- a/src/DistributedComputations/distributed_on_architecture.jl +++ b/src/DistributedComputations/distributed_on_architecture.jl @@ -23,7 +23,7 @@ function on_architecture(new_arch::Distributed, old_grid::LatitudeLongitudeGrid) child_arch = child_architecture(new_arch) old_properties = (old_grid.Δλᶠᵃᵃ, old_grid.Δλᶜᵃᵃ, old_grid.λᶠᵃᵃ, old_grid.λᶜᵃᵃ, old_grid.Δφᵃᶠᵃ, old_grid.Δφᵃᶜᵃ, old_grid.φᵃᶠᵃ, old_grid.φᵃᶜᵃ, - old_grid.Δzᵃᵃᶠ, old_grid.Δzᵃᵃᶜ, old_grid.zᵃᵃᶠ, old_grid.zᵃᵃᶜ, + old_grid.z, old_grid.Δxᶠᶜᵃ, old_grid.Δxᶜᶠᵃ, old_grid.Δxᶠᶠᵃ, old_grid.Δxᶜᶜᵃ, old_grid.Δyᶠᶜᵃ, old_grid.Δyᶜᶠᵃ, old_grid.Azᶠᶜᵃ, old_grid.Azᶜᶠᵃ, old_grid.Azᶠᶠᵃ, old_grid.Azᶜᶜᵃ) @@ -44,7 +44,7 @@ function on_architecture(new_arch::Distributed, old_grid::RectilinearGrid) child_arch = child_architecture(new_arch) old_properties = (old_grid.Δxᶠᵃᵃ, old_grid.Δxᶜᵃᵃ, old_grid.xᶠᵃᵃ, old_grid.xᶜᵃᵃ, old_grid.Δyᵃᶠᵃ, old_grid.Δyᵃᶜᵃ, old_grid.yᵃᶠᵃ, old_grid.yᵃᶜᵃ, - old_grid.Δzᵃᵃᶠ, old_grid.Δzᵃᵃᶜ, old_grid.zᵃᵃᶠ, old_grid.zᵃᵃᶜ) + old_grid.z) new_properties = Tuple(on_architecture(child_arch, p) for p in old_properties) diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index 8f470663fb..72a8e78540 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -11,7 +11,7 @@ using Oceananigans.Grids: interior_indices, interior_parent_indices import Base: minimum, maximum, extrema import Oceananigans.Architectures: architecture, child_architecture import Oceananigans.Grids: interior_x_indices, interior_y_indices, interior_z_indices -import Oceananigans.Grids: total_size, topology, nodes, xnodes, ynodes, znodes, node, xnode, ynode, znode +import Oceananigans.Grids: total_size, topology, nodes, xnodes, ynodes, znodes, rnodes, node, xnode, ynode, znode, rnode import Oceananigans.Utils: datatuple const ArchOrNothing = Union{AbstractArchitecture, Nothing} @@ -106,10 +106,12 @@ interior(f::AbstractField) = f @propagate_inbounds xnode(i, j, k, ψ::AbstractField) = xnode(i, j, k, ψ.grid, instantiated_location(ψ)...) @propagate_inbounds ynode(i, j, k, ψ::AbstractField) = ynode(i, j, k, ψ.grid, instantiated_location(ψ)...) @propagate_inbounds znode(i, j, k, ψ::AbstractField) = znode(i, j, k, ψ.grid, instantiated_location(ψ)...) +@propagate_inbounds rnode(i, j, k, ψ::AbstractField) = rnode(i, j, k, ψ.grid, instantiated_location(ψ)...) xnodes(ψ::AbstractField; kwargs...) = xnodes(ψ.grid, instantiated_location(ψ)...; kwargs...) ynodes(ψ::AbstractField; kwargs...) = ynodes(ψ.grid, instantiated_location(ψ)...; kwargs...) znodes(ψ::AbstractField; kwargs...) = znodes(ψ.grid, instantiated_location(ψ)...; kwargs...) +rnodes(ψ::AbstractField; kwargs...) = rnodes(ψ.grid, instantiated_location(ψ)...; kwargs...) nodes(ψ::AbstractField; kwargs...) = nodes(ψ.grid, instantiated_location(ψ); kwargs...) diff --git a/src/Forcings/forcing.jl b/src/Forcings/forcing.jl index c3f405be62..24c2eeb838 100644 --- a/src/Forcings/forcing.jl +++ b/src/Forcings/forcing.jl @@ -143,7 +143,7 @@ DiscreteForcing{Nothing} ```jldoctest forcing # Discrete-form forcing function with parameters masked_damping(i, j, k, grid, clock, model_fields, parameters) = - @inbounds - parameters.μ * exp(grid.zᵃᵃᶜ[k] / parameters.λ) * model_fields.u[i, j, k] + @inbounds - parameters.μ * exp(grid.z.cᶜ[k] / parameters.λ) * model_fields.u[i, j, k] masked_damping_forcing = Forcing(masked_damping, parameters=(μ=42, λ=π), discrete_form=true) diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 519233561b..decf316586 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -14,12 +14,14 @@ export conformal_cubed_sphere_panel export node, nodes export ξnode, ηnode, rnode export xnode, ynode, znode, λnode, φnode -export xnodes, ynodes, znodes, λnodes, φnodes +export xnodes, ynodes, znodes, λnodes, φnodes, rnodes export spacings -export xspacing, yspacing, zspacing, λspacing, φspacing -export xspacings, yspacings, zspacings, λspacings, φspacings +export xspacing, yspacing, zspacing, λspacing, φspacing, rspacing +export xspacings, yspacings, zspacings, λspacings, φspacings, rspacings export minimum_xspacing, minimum_yspacing, minimum_zspacing +export ZStarVerticalCoordinate, vertical_scaling, previous_vertical_scaling, reference_zspacings export static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ +export dynamic_column_depthᶜᶜᵃ, dynamic_column_depthᶠᶜᵃ, dynamic_column_depthᶜᶠᵃ, dynamic_column_depthᶠᶠᵃ export offset_data, new_data export on_architecture @@ -27,6 +29,7 @@ using CUDA using CUDA: has_cuda using Adapt using OffsetArrays +using Printf using Oceananigans using Oceananigans.Architectures @@ -118,6 +121,7 @@ struct ZDirection <: AbstractDirection end struct NegativeZDirection <: AbstractDirection end include("abstract_grid.jl") +include("vertical_coordinate.jl") include("grid_utils.jl") include("nodes_and_spacings.jl") include("zeros_and_ones.jl") diff --git a/src/Grids/abstract_grid.jl b/src/Grids/abstract_grid.jl index 92b0c3f21b..4a2398cb4f 100644 --- a/src/Grids/abstract_grid.jl +++ b/src/Grids/abstract_grid.jl @@ -9,23 +9,25 @@ abstract type AbstractGrid{FT, TX, TY, TZ, Arch} end AbstractUnderlyingGrid{FT, TX, TY, TZ} Abstract supertype for "primary" grids (as opposed to grids with immersed boundaries) -with elements of type `FT` and topology `{TX, TY, TZ}`. +with elements of type `FT`, topology `{TX, TY, TZ}` and vertical coordinate `CZ`. """ -abstract type AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} <: AbstractGrid{FT, TX, TY, TZ, Arch} end +abstract type AbstractUnderlyingGrid{FT, TX, TY, TZ, CZ, Arch} <: AbstractGrid{FT, TX, TY, TZ, Arch} end """ AbstractCurvilinearGrid{FT, TX, TY, TZ} -Abstract supertype for curvilinear grids with elements of type `FT` and topology `{TX, TY, TZ}`. +Abstract supertype for curvilinear grids with elements of type `FT`, +topology `{TX, TY, TZ}`, and vertical coordinate `CZ`. """ -abstract type AbstractCurvilinearGrid{FT, TX, TY, TZ, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} end +abstract type AbstractCurvilinearGrid{FT, TX, TY, TZ, CZ, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, CZ, Arch} end """ AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ} -Abstract supertype for horizontally-curvilinear grids with elements of type `FT` and topology `{TX, TY, TZ}`. +Abstract supertype for horizontally-curvilinear grids with elements of type `FT`, +topology `{TX, TY, TZ}` and vertical coordinate `CZ`. """ -abstract type AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch} <: AbstractCurvilinearGrid{FT, TX, TY, TZ, Arch} end +abstract type AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, CZ, Arch} <: AbstractCurvilinearGrid{FT, TX, TY, TZ, CZ, Arch} end const XFlatGrid = AbstractGrid{<:Any, Flat} const YFlatGrid = AbstractGrid{<:Any, <:Any, Flat} diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl index 84212fb580..51831b4313 100644 --- a/src/Grids/grid_generation.jl +++ b/src/Grids/grid_generation.jl @@ -26,6 +26,10 @@ total_interior_length(::BoundedTopology, N) = N + 1 bad_coordinate_message(ξ::Function, name) = "The values of $name(index) must increase as the index increases!" bad_coordinate_message(ξ::AbstractArray, name) = "The elements of $name must be increasing!" +# General generate_coordinate +generate_coordinate(FT, topology, size, halo, nodes, coordinate_name, dim::Int, arch) = + generate_coordinate(FT, topology[dim](), size[dim], halo[dim], nodes, coordinate_name, arch) + # generate a variably-spaced coordinate passing the explicit coord faces as vector or function function generate_coordinate(FT, topo::AT, N, H, node_generator, coordinate_name, arch) @@ -83,7 +87,11 @@ function generate_coordinate(FT, topo::AT, N, H, node_generator, coordinate_name F = OffsetArray(on_architecture(arch, F.parent), F.offsets...) C = OffsetArray(on_architecture(arch, C.parent), C.offsets...) - return L, F, C, Δᶠ, Δᶜ + if coordinate_name == :z + return L, StaticVerticalCoordinate(F, C, Δᶠ, Δᶜ) + else + return L, F, C, Δᶠ, Δᶜ + end end # Generate a regularly-spaced coordinate passing the domain extent (2-tuple) and number of points @@ -116,16 +124,74 @@ function generate_coordinate(FT, topo::AT, N, H, node_interval::Tuple{<:Number, F = OffsetArray(F, -H) C = OffsetArray(C, -H) - return FT(L), F, C, FT(Δᶠ), FT(Δᶜ) + if coordinate_name == :z + return FT(L), StaticVerticalCoordinate(F, C, FT(Δᶠ), FT(Δᶜ)) + else + return FT(L), F, C, FT(Δᶠ), FT(Δᶜ) + end end # Flat domains -generate_coordinate(FT, ::Flat, N, H, c::Number, coordinate_name, arch) = - FT(1), range(FT(c), FT(c), length=N), range(FT(c), FT(c), length=N), FT(1), FT(1) +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)) + else + return FT(1), range(FT(c), FT(c), length=N), range(FT(c), FT(c), length=N), FT(1), FT(1) + end +end # What's the use case for this? # generate_coordinate(FT, ::Flat, N, H, c::Tuple{Number, Number}, coordinate_name, arch) = # 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)) + else + return FT(1), nothing, nothing, FT(1), FT(1) + end +end + +#### +#### ZStarVerticalCoordinate +#### + +generate_coordinate(FT, ::Periodic, N, H, ::ZStarVerticalCoordinate, coordinate_name, arch, args...) = + throw(ArgumentError("Periodic domains are not supported for ZStarVerticalCoordinate")) + +# Generate a moving coordinate with evolving scaling (`s`) for spacings and znodes +function generate_coordinate(FT, topo, size, halo, coordinate::ZStarVerticalCoordinate, coordinate_name, dim::Int, arch) + + Nx, Ny, Nz = size + Hx, Hy, Hz = halo + + if dim != 3 + msg = "ZStarVerticalCoordinate is supported only in the third dimension (z)" + throw(ArgumentError(msg)) + end + + if coordinate_name != :z + msg = "Only z-coordinate is supported for ZStarVerticalCoordinate" + throw(ArgumentError(msg)) + end + + r_faces = coordinate.cᶠ -generate_coordinate(FT, ::Flat, N, H, ::Nothing, coordinate_name, arch) = - FT(1), nothing, nothing, FT(1), FT(1) + Lr, rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ = generate_coordinate(FT, topo[3](), Nz, Hz, r_faces, :r, arch) + + args = (topo, (Nx, Ny, Nz), (Hx, Hy, Hz)) + + e₃ᶜᶜ⁻ = new_data(FT, arch, (Center, Center, Nothing), args...) + e₃ᶜᶜⁿ = new_data(FT, arch, (Center, Center, Nothing), args...) + e₃ᶠᶜⁿ = new_data(FT, arch, (Face, Center, Nothing), args...) + e₃ᶜᶠⁿ = new_data(FT, arch, (Center, Face, Nothing), args...) + e₃ᶠᶠⁿ = new_data(FT, arch, (Face, Face, Nothing), args...) + ηⁿ = new_data(FT, arch, (Center, Center, Nothing), args...) + ∂t_e₃ = new_data(FT, arch, (Center, Center, Nothing), args...) + + # Fill all the scalings with one (at rest coordinate) + for e₃ in (e₃ᶜᶜ⁻, e₃ᶜᶜⁿ, e₃ᶠᶜⁿ, e₃ᶜᶠⁿ, e₃ᶠᶠⁿ) + fill!(e₃, 1) + end + + return Lr, ZStarVerticalCoordinate(rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ, ηⁿ, e₃ᶜᶜⁿ, e₃ᶠᶜⁿ, e₃ᶜᶠⁿ, e₃ᶠᶠⁿ, e₃ᶜᶜ⁻, ∂t_e₃) +end diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index 178d7ab87c..d6cc7c87d1 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -1,5 +1,3 @@ -using CUDA -using Printf using Base.Ryu: writeshortest using LinearAlgebra: dot, cross using OffsetArrays: IdOffsetRange @@ -126,7 +124,6 @@ constant grid spacing `Δ`, and interior extent `L`. @inline x_domain(grid) = domain(topology(grid, 1)(), grid.Nx, grid.xᶠᵃᵃ) @inline y_domain(grid) = domain(topology(grid, 2)(), grid.Ny, grid.yᵃᶠᵃ) -@inline z_domain(grid) = domain(topology(grid, 3)(), grid.Nz, grid.zᵃᵃᶠ) regular_dimensions(grid) = () @@ -234,7 +231,6 @@ const f = Face() # What's going on here? @inline cpu_face_constructor_x(grid) = Array(getindex(nodes(grid, f, c, c; with_halos=true), 1)[1:size(grid, 1)+1]) @inline cpu_face_constructor_y(grid) = Array(getindex(nodes(grid, c, f, c; with_halos=true), 2)[1:size(grid, 2)+1]) -@inline cpu_face_constructor_z(grid) = Array(getindex(nodes(grid, c, c, f; with_halos=true), 3)[1:size(grid, 3)+1]) ##### ##### Convenience functions @@ -300,6 +296,12 @@ function domain_summary(topo, name, (left, right)) prettysummary(right), interval) 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)) +end + function dimension_summary(topo, name, dom, spacing, pad_domain=0) prefix = domain_summary(topo, name, dom) padding = " "^(pad_domain+1) @@ -315,7 +317,7 @@ coordinate_summary(topo, Δ::Union{AbstractVector, AbstractMatrix}, name) = name, prettysummary(maximum(parent(Δ)))) ##### -##### Static column depth +##### Static and Dynamic column depths ##### @inline static_column_depthᶜᶜᵃ(i, j, grid) = grid.Lz @@ -323,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 any ZStar grid type +@inline dynamic_column_depthᶜᶜᵃ(i, j, k, grid, η) = static_column_depthᶜᶜᵃ(i, j, grid) +@inline dynamic_column_depthᶜᶠᵃ(i, j, k, grid, η) = static_column_depthᶜᶠᵃ(i, j, grid) +@inline dynamic_column_depthᶠᶜᵃ(i, j, k, grid, η) = static_column_depthᶠᶜᵃ(i, j, grid) +@inline dynamic_column_depthᶠᶠᵃ(i, j, k, grid, η) = static_column_depthᶠᶠᵃ(i, j, grid) + ##### ##### Spherical geometry ##### diff --git a/src/Grids/input_validation.jl b/src/Grids/input_validation.jl index 9624086bba..0ef5a9e018 100644 --- a/src/Grids/input_validation.jl +++ b/src/Grids/input_validation.jl @@ -90,16 +90,6 @@ function validate_halo(TX, TY, TZ, size, halo) return halo end -function validate_dimension_specification(T, ξ, dir, N, FT) - - isnothing(ξ) && throw(ArgumentError("Must supply extent or $dir keyword when $dir-direction is $T")) - length(ξ) == 2 || throw(ArgumentError("$dir length($ξ) must be 2.")) - all(isa.(ξ, Number)) || throw(ArgumentError("$dir=$ξ should contain numbers.")) - ξ[2] ≥ ξ[1] || throw(ArgumentError("$dir=$ξ should be an increasing interval.")) - - return FT.(ξ) -end - function validate_rectilinear_domain(TX, TY, TZ, FT, size, extent, x, y, z) # Find domain endpoints or domain extent, depending on user input: @@ -126,6 +116,16 @@ function validate_rectilinear_domain(TX, TY, TZ, FT, size, extent, x, y, z) return x, y, z end +function validate_dimension_specification(T, ξ, dir, N, FT) + + isnothing(ξ) && throw(ArgumentError("Must supply extent or $dir keyword when $dir-direction is $T")) + length(ξ) == 2 || throw(ArgumentError("$dir length($ξ) must be 2.")) + all(isa.(ξ, Number)) || throw(ArgumentError("$dir=$ξ should contain numbers.")) + ξ[2] ≥ ξ[1] || throw(ArgumentError("$dir=$ξ should be an increasing interval.")) + + return FT.(ξ) +end + function validate_dimension_specification(T, ξ::AbstractVector, dir, N, FT) ξ = FT.(ξ) diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index bd411ea5d1..6c12486e6e 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -1,6 +1,6 @@ using KernelAbstractions: @kernel, @index -struct LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch} +struct LatitudeLongitudeGrid{FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, CZ, Arch} architecture :: Arch Nx :: Int Ny :: Int @@ -21,10 +21,7 @@ struct LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch Δφᵃᶜᵃ :: FY φᵃᶠᵃ :: VY φᵃᶜᵃ :: VY - Δzᵃᵃᶠ :: FZ - Δzᵃᵃᶜ :: FZ - zᵃᵃᶠ :: VZ - zᵃᵃᶜ :: VZ + z :: CZ # Precomputed metrics M <: Nothing means metrics will be computed on the fly Δxᶠᶜᵃ :: M Δxᶜᶠᵃ :: M @@ -47,35 +44,34 @@ struct LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch λᶠᵃᵃ :: VX, λᶜᵃᵃ :: VX, Δφᵃᶠᵃ :: FY, Δφᵃᶜᵃ :: FY, φᵃᶠᵃ :: VY, φᵃᶜᵃ :: VY, - Δzᵃᵃᶠ :: FZ, Δzᵃᵃᶜ :: FZ, - zᵃᵃᶠ :: VZ, zᵃᵃᶜ :: VZ, + z :: CZ, Δxᶠᶜᵃ :: M, Δxᶜᶠᵃ :: M, Δxᶠᶠᵃ :: M, Δxᶜᶜᵃ :: M, Δyᶠᶜᵃ :: MY, Δyᶜᶠᵃ :: MY, - Azᶠᶜᵃ :: M, Azᶜᶠᵃ :: M, + Azᶠᶜᵃ :: M, Azᶜᶠᵃ :: M, Azᶠᶠᵃ :: M, Azᶜᶜᵃ :: M, - radius :: FT) where {Arch, FT, TX, TY, TZ, - FX, FY, FZ, VX, VY, VZ, + radius :: FT) where {Arch, FT, TX, TY, TZ, + FX, FY, CZ, VX, VY, M, MY} = - new{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch}(architecture, - Nλ, Nφ, Nz, - Hλ, Hφ, Hz, - Lλ, Lφ, Lz, - Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, - Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, - Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, - Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ, - Δyᶠᶜᵃ, Δyᶜᶠᵃ, - Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ, radius) + new{FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY, Arch}(architecture, + Nλ, Nφ, Nz, + Hλ, Hφ, Hz, + Lλ, Lφ, Lz, + Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, + Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, + z, + Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ, + Δyᶠᶜᵃ, Δyᶜᶠᵃ, + Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ, radius) end const LLG = LatitudeLongitudeGrid -const XRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} -const YRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} -const ZRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} -const HRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:Number} -const HNonRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractArray, <:AbstractArray} -const YNonRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:AbstractArray} +const XRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} +const YRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} +const ZRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate} +const HRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:Number} +const HNonRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractArray, <:AbstractArray} +const YNonRegularLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:AbstractArray} regular_dimensions(::ZRegularLLG) = tuple(3) @@ -197,9 +193,9 @@ function LatitudeLongitudeGrid(architecture::AbstractArchitecture = CPU(), # it is stretched if being passed is a function or vector (as for the VerticallyStretchedRectilinearGrid) TX, TY, TZ = topology - Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), Nλ, Hλ, longitude, :longitude, architecture) - Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY(), Nφ, Hφ, latitude, :latitude, architecture) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, z, :z, architecture) + Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, topology, size, halo, longitude, :longitude, 1, architecture) + Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, topology, size, halo, latitude, :latitude, 2, architecture) + Lz, z = generate_coordinate(FT, topology, size, halo, z, :z, 3, architecture) preliminary_grid = LatitudeLongitudeGrid{TX, TY, TZ}(architecture, Nλ, Nφ, Nz, @@ -207,7 +203,7 @@ function LatitudeLongitudeGrid(architecture::AbstractArchitecture = CPU(), Lλ, Lφ, Lz, Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ, - Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ, + z, (nothing for i=1:10)..., FT(radius)) if !precompute_metrics @@ -248,7 +244,7 @@ function with_precomputed_metrics(grid) grid.Lx, grid.Ly, grid.Lz, grid.Δλᶠᵃᵃ, grid.Δλᶜᵃᵃ, grid.λᶠᵃᵃ, grid.λᶜᵃᵃ, grid.Δφᵃᶠᵃ, grid.Δφᵃᶜᵃ, grid.φᵃᶠᵃ, grid.φᵃᶜᵃ, - grid.Δzᵃᵃᶠ, grid.Δzᵃᵃᶜ, grid.zᵃᵃᶠ, grid.zᵃᵃᶜ, + grid.z, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ, grid.radius) end @@ -311,7 +307,7 @@ function Base.show(io::IO, grid::LatitudeLongitudeGrid, withsummary=true) Ωλ = domain(TX(), size(grid, 1), grid.λᶠᵃᵃ) Ωφ = domain(TY(), size(grid, 2), grid.φᵃᶠᵃ) - Ωz = domain(TZ(), size(grid, 3), grid.zᵃᵃᶠ) + Ωz = domain(TZ(), size(grid, 3), grid.z.cᶠ) x_summary = domain_summary(TX(), "λ", Ωλ) y_summary = domain_summary(TY(), "φ", Ωφ) @@ -321,7 +317,7 @@ function Base.show(io::IO, grid::LatitudeLongitudeGrid, withsummary=true) x_summary = "longitude: " * dimension_summary(TX(), "λ", Ωλ, grid.Δλᶜᵃᵃ, longest - length(x_summary)) y_summary = "latitude: " * dimension_summary(TY(), "φ", Ωφ, grid.Δφᵃᶜᵃ, longest - length(y_summary)) - z_summary = "z: " * dimension_summary(TZ(), "z", Ωz, grid.Δzᵃᵃᶜ, longest - length(z_summary)) + z_summary = "z: " * dimension_summary(TZ(), "z", Ωz, grid.z, longest - length(z_summary)) if withsummary print(io, summary(grid), "\n") @@ -334,11 +330,9 @@ end @inline x_domain(grid::LLG) = domain(topology(grid, 1)(), grid.Nx, grid.λᶠᵃᵃ) @inline y_domain(grid::LLG) = domain(topology(grid, 2)(), grid.Ny, grid.φᵃᶠᵃ) -@inline z_domain(grid::LLG) = domain(topology(grid, 3)(), grid.Nz, grid.zᵃᵃᶠ) @inline cpu_face_constructor_x(grid::XRegularLLG) = x_domain(grid) @inline cpu_face_constructor_y(grid::YRegularLLG) = y_domain(grid) -@inline cpu_face_constructor_z(grid::ZRegularLLG) = z_domain(grid) function constructor_arguments(grid::LatitudeLongitudeGrid) arch = architecture(grid) @@ -410,10 +404,7 @@ function Adapt.adapt_structure(to, grid::LatitudeLongitudeGrid) Adapt.adapt(to, grid.Δφᵃᶜᵃ), Adapt.adapt(to, grid.φᵃᶠᵃ), Adapt.adapt(to, grid.φᵃᶜᵃ), - Adapt.adapt(to, grid.Δzᵃᵃᶠ), - Adapt.adapt(to, grid.Δzᵃᵃᶜ), - Adapt.adapt(to, grid.zᵃᵃᶠ), - Adapt.adapt(to, grid.zᵃᵃᶜ), + Adapt.adapt(to, grid.z), Adapt.adapt(to, grid.Δxᶠᶜᵃ), Adapt.adapt(to, grid.Δxᶜᶠᵃ), Adapt.adapt(to, grid.Δxᶠᶠᵃ), @@ -583,14 +574,10 @@ rname(::LLG) = :z @inline λnode(i, grid::LLG, ::Face) = getnode(grid.λᶠᵃᵃ, i) @inline φnode(j, grid::LLG, ::Center) = getnode(grid.φᵃᶜᵃ, j) @inline φnode(j, grid::LLG, ::Face) = getnode(grid.φᵃᶠᵃ, j) -@inline znode(k, grid::LLG, ::Center) = getnode(grid.zᵃᵃᶜ, k) -@inline znode(k, grid::LLG, ::Face) = getnode(grid.zᵃᵃᶠ, k) # Definitions for node @inline ξnode(i, j, k, grid::LLG, ℓx, ℓy, ℓz) = λnode(i, grid, ℓx) @inline ηnode(i, j, k, grid::LLG, ℓx, ℓy, ℓz) = φnode(j, grid, ℓy) -@inline rnode(i, j, k, grid::LLG, ℓx, ℓy, ℓz) = znode(k, grid, ℓz) - @inline xnode(i, j, grid::LLG, ℓx, ℓy) = grid.radius * deg2rad(λnode(i, grid, ℓx)) * hack_cosd((φnode(j, grid, ℓy))) @inline ynode(j, grid::LLG, ℓy) = grid.radius * deg2rad(φnode(j, grid, ℓy)) @@ -599,7 +586,6 @@ rname(::LLG) = :z @inline φnode(i, j, k, grid::LLG, ℓx, ℓy, ℓz) = φnode(j, grid, ℓy) @inline xnode(i, j, k, grid::LLG, ℓx, ℓy, ℓz) = xnode(i, j, grid, ℓx, ℓy) @inline ynode(i, j, k, grid::LLG, ℓx, ℓy, ℓz) = ynode(j, grid, ℓy) -@inline znode(i, j, k, grid::LLG, ℓx, ℓy, ℓz) = znode(k, grid, ℓz) function nodes(grid::LLG, ℓx, ℓy, ℓz; reshape=false, with_halos=false) λ = λnodes(grid, ℓx, ℓy, ℓz; with_halos) @@ -645,13 +631,9 @@ end return @. R * deg2rad(φ) end -@inline znodes(grid::LLG, ℓz::F; with_halos=false) = _property(grid.zᵃᵃᶠ, ℓz, topology(grid, 3), size(grid, 3), with_halos) -@inline znodes(grid::LLG, ℓz::C; with_halos=false) = _property(grid.zᵃᵃᶜ, ℓz, topology(grid, 3), size(grid, 3), with_halos) - # Convenience @inline λnodes(grid::LLG, ℓx, ℓy, ℓz; with_halos=false) = λnodes(grid, ℓx; with_halos) @inline φnodes(grid::LLG, ℓx, ℓy, ℓz; with_halos=false) = φnodes(grid, ℓy; with_halos) -@inline znodes(grid::LLG, ℓx, ℓy, ℓz; with_halos=false) = znodes(grid, ℓz; with_halos) @inline xnodes(grid::LLG, ℓx, ℓy, ℓz; with_halos=false) = xnodes(grid, ℓx, ℓy; with_halos) @inline ynodes(grid::LLG, ℓx, ℓy, ℓz; with_halos=false) = ynodes(grid, ℓy; with_halos) @@ -663,11 +645,9 @@ end # Generalized coordinates @inline ξnodes(grid::LLG, ℓx; kwargs...) = λnodes(grid, ℓx; kwargs...) @inline ηnodes(grid::LLG, ℓy; kwargs...) = φnodes(grid, ℓy; kwargs...) -@inline rnodes(grid::LLG, ℓz; kwargs...) = znodes(grid, ℓz; kwargs...) @inline ξnodes(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = λnodes(grid, ℓx; kwargs...) @inline ηnodes(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = φnodes(grid, ℓy; kwargs...) -@inline rnodes(grid::LLG, ℓx, ℓy, ℓz; kwargs...) = znodes(grid, ℓz; kwargs...) ##### ##### Grid spacings @@ -688,7 +668,6 @@ end @inline xspacings(grid::LLG, ℓx, ℓy) = xspacings(grid, ℓx, ℓy, nothing) @inline yspacings(grid::LLG, ℓx, ℓy) = yspacings(grid, ℓx, ℓy, nothing) -@inline zspacings(grid::LLG, ℓz) = zspacings(grid, nothing, nothing, ℓz) @inline λspacings(grid::LLG, ℓx) = λspacings(grid, ℓx, nothing, nothing) @inline φspacings(grid::LLG, ℓy) = φspacings(grid, nothing, ℓy, nothing) diff --git a/src/Grids/nodes_and_spacings.jl b/src/Grids/nodes_and_spacings.jl index fd814d94e5..0d0056c09b 100644 --- a/src/Grids/nodes_and_spacings.jl +++ b/src/Grids/nodes_and_spacings.jl @@ -66,7 +66,6 @@ _node_names(grid, ::Nothing, ::Nothing, ::Nothing) = tuple() xnodes(grid, ::Nothing; kwargs...) = 1:1 ynodes(grid, ::Nothing; kwargs...) = 1:1 -znodes(grid, ::Nothing; kwargs...) = 1:1 """ xnodes(grid, ℓx, ℓy, ℓz, with_halos=false) @@ -88,36 +87,6 @@ See [`znodes`](@ref) for examples. """ @inline ynodes(grid, ℓx, ℓy, ℓz; kwargs...) = ynodes(grid, ℓy; kwargs...) -""" - znodes(grid, ℓx, ℓy, ℓz; with_halos=false) - -Return the positions over the interior nodes on `grid` in the ``z``-direction for the location `ℓx`, -`ℓy`, `ℓz`. For `Bounded` directions, `Face` nodes include the boundary points. - -```jldoctest znodes -julia> using Oceananigans - -julia> horz_periodic_grid = RectilinearGrid(size=(3, 3, 3), extent=(2π, 2π, 1), halo=(1, 1, 1), - topology=(Periodic, Periodic, Bounded)); - -julia> zC = znodes(horz_periodic_grid, Center()) -3-element view(OffsetArray(::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, 0:4), 1:3) with eltype Float64: - -0.8333333333333334 - -0.5 - -0.16666666666666666 - -julia> zC = znodes(horz_periodic_grid, Center(), Center(), Center()) -3-element view(OffsetArray(::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, 0:4), 1:3) with eltype Float64: - -0.8333333333333334 - -0.5 - -0.16666666666666666 - -julia> zC = znodes(horz_periodic_grid, Center(), Center(), Center(), with_halos=true) --1.1666666666666667:0.3333333333333333:0.16666666666666666 with indices 0:4 -``` -""" -@inline znodes(grid, ℓx, ℓy, ℓz; kwargs...) = znodes(grid, ℓz; kwargs...) - """ λnodes(grid::AbstractCurvilinearGrid, ℓx, ℓy, ℓz, with_halos=false) @@ -171,7 +140,6 @@ function φspacing end function xspacings end function yspacings end -function zspacings end function λspacings end function φspacings end diff --git a/src/Grids/orthogonal_spherical_shell_grid.jl b/src/Grids/orthogonal_spherical_shell_grid.jl index 7a30c2b7a4..93db54f105 100644 --- a/src/Grids/orthogonal_spherical_shell_grid.jl +++ b/src/Grids/orthogonal_spherical_shell_grid.jl @@ -15,7 +15,7 @@ struct CubedSphereConformalMapping{FT, Rotation} rotation :: Rotation end -struct OrthogonalSphericalShellGrid{FT, TX, TY, TZ, A, R, FR, C, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch} +struct OrthogonalSphericalShellGrid{FT, TX, TY, TZ, CZ, A, C, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, CZ, Arch} architecture :: Arch Nx :: Int Ny :: Int @@ -32,8 +32,7 @@ struct OrthogonalSphericalShellGrid{FT, TX, TY, TZ, A, R, FR, C, Arch} <: Abstra φᶠᶜᵃ :: A φᶜᶠᵃ :: A φᶠᶠᵃ :: A - zᵃᵃᶜ :: R - zᵃᵃᶠ :: R + z :: CZ Δxᶜᶜᵃ :: A Δxᶠᶜᵃ :: A Δxᶜᶠᵃ :: A @@ -42,8 +41,6 @@ struct OrthogonalSphericalShellGrid{FT, TX, TY, TZ, A, R, FR, C, Arch} <: Abstra Δyᶜᶠᵃ :: A Δyᶠᶜᵃ :: A Δyᶠᶠᵃ :: A - Δzᵃᵃᶜ :: FR - Δzᵃᵃᶠ :: FR Azᶜᶜᵃ :: A Azᶠᶜᵃ :: A Azᶜᶠᵃ :: A @@ -56,36 +53,36 @@ struct OrthogonalSphericalShellGrid{FT, TX, TY, TZ, A, R, FR, C, Arch} <: Abstra Hx, Hy, Hz, Lz :: FT, λᶜᶜᵃ :: A, λᶠᶜᵃ :: A, λᶜᶠᵃ :: A, λᶠᶠᵃ :: A, - φᶜᶜᵃ :: A, φᶠᶜᵃ :: A, φᶜᶠᵃ :: A, φᶠᶠᵃ :: A, zᵃᵃᶜ :: R, zᵃᵃᶠ :: R, + φᶜᶜᵃ :: A, φᶠᶜᵃ :: A, φᶜᶠᵃ :: A, φᶠᶠᵃ :: A, z :: CZ, Δxᶜᶜᵃ :: A, Δxᶠᶜᵃ :: A, Δxᶜᶠᵃ :: A, Δxᶠᶠᵃ :: A, - Δyᶜᶜᵃ :: A, Δyᶜᶠᵃ :: A, Δyᶠᶜᵃ :: A, Δyᶠᶠᵃ :: A, Δzᵃᵃᶜ :: FR, Δzᵃᵃᶠ :: FR, + Δyᶜᶜᵃ :: A, Δyᶜᶠᵃ :: A, Δyᶠᶜᵃ :: A, Δyᶠᶠᵃ :: A, Azᶜᶜᵃ :: A, Azᶠᶜᵃ :: A, Azᶜᶠᵃ :: A, Azᶠᶠᵃ :: A, radius :: FT, - conformal_mapping :: C) where {TX, TY, TZ, FT, A, R, FR, C, Arch} = - new{FT, TX, TY, TZ, A, R, FR, C, Arch}(architecture, + conformal_mapping :: C) where {TX, TY, TZ, FT, CZ, A, C, Arch} = + new{FT, TX, TY, TZ, CZ, A, C, Arch}(architecture, Nx, Ny, Nz, Hx, Hy, Hz, Lz, λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, - φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ, + φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, z, Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, - Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δzᵃᵃᶜ, Δzᵃᵃᶠ, + Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius, conformal_mapping) end const OSSG = OrthogonalSphericalShellGrid -const ZRegOSSG = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} +const ZRegOSSG = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate} const ZRegOrthogonalSphericalShellGrid = ZRegOSSG -const ConformalCubedSpherePanel = OrthogonalSphericalShellGrid{<:Any, FullyConnected, FullyConnected, <:Any, <:Any, <:Any, <:Any, <:CubedSphereConformalMapping} +const ConformalCubedSpherePanel = OrthogonalSphericalShellGrid{<:Any, FullyConnected, FullyConnected, <:Any, <:Any, <:Any, <:CubedSphereConformalMapping} # convenience constructor for OSSG without any conformal_mapping properties OrthogonalSphericalShellGrid(architecture, Nx, Ny, Nz, Hx, Hy, Hz, Lz, - λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ, - Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δzᵃᵃᶜ, Δzᵃᵃᶠ, + λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, z, + Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius) = OrthogonalSphericalShellGrid(architecture, Nx, Ny, Nz, Hx, Hy, Hz, Lz, - λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, zᵃᵃᶜ, zᵃᵃᶠ, - Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Δzᵃᵃᶜ, Δzᵃᵃᶠ, + λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, z, + Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius, nothing) """ @@ -202,11 +199,8 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU() ηᵃᶜᵃ = ynodes(ξη_grid, Center()) ## The vertical coordinates and metrics can come out of the regular rectilinear grid! - zᵃᵃᶠ = ξη_grid.zᵃᵃᶠ - zᵃᵃᶜ = ξη_grid.zᵃᵃᶜ - Δzᵃᵃᶜ = ξη_grid.Δzᵃᵃᶜ - Δzᵃᵃᶠ = ξη_grid.Δzᵃᵃᶠ - Lz = ξη_grid.Lz + zc = ξη_grid.z + Lz = ξη_grid.Lz ## Compute staggered grid latitude-longitude (φ, λ) coordinates. @@ -603,11 +597,10 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU() coordinate_arrays = (λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, - zᵃᵃᶜ, zᵃᵃᶠ) + zc) metric_arrays = (Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, - Δzᵃᵃᶜ, Δzᵃᵃᶠ, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ) conformal_mapping = CubedSphereConformalMapping(ξ, η, rotation) @@ -624,11 +617,10 @@ function conformal_cubed_sphere_panel(architecture::AbstractArchitecture = CPU() coordinate_arrays = (grid.λᶜᶜᵃ, grid.λᶠᶜᵃ, grid.λᶜᶠᵃ, grid.λᶠᶠᵃ, grid.φᶜᶜᵃ, grid.φᶠᶜᵃ, grid.φᶜᶠᵃ, grid.φᶠᶠᵃ, - grid.zᵃᵃᶜ, grid.zᵃᵃᶠ) + grid.z) metric_arrays = (grid.Δxᶜᶜᵃ, grid.Δxᶠᶜᵃ, grid.Δxᶜᶠᵃ, grid.Δxᶠᶠᵃ, grid.Δyᶜᶜᵃ, grid.Δyᶜᶠᵃ, grid.Δyᶠᶜᵃ, grid.Δyᶠᶠᵃ, - grid.Δzᵃᵃᶜ, grid.Δzᵃᵃᶠ, grid.Azᶜᶜᵃ, grid.Azᶠᶜᵃ, grid.Azᶜᶠᵃ, grid.Azᶠᶠᵃ) coordinate_arrays = map(a -> on_architecture(architecture, a), coordinate_arrays) @@ -824,17 +816,7 @@ function conformal_cubed_sphere_panel(filepath::AbstractString, architecture = C TX, TY, TZ = topology Hx, Hy, Hz = halo - ## The vertical coordinates can come out of the regular rectilinear grid! - - z_grid = RectilinearGrid(architecture, FT; size = Nz, z, topology=(Flat, Flat, topology[3]), halo=halo[3]) - - zᵃᵃᶠ = z_grid.zᵃᵃᶠ - zᵃᵃᶜ = z_grid.zᵃᵃᶜ - Δzᵃᵃᶜ = z_grid.Δzᵃᵃᶜ - Δzᵃᵃᶠ = z_grid.Δzᵃᵃᶠ - Lz = z_grid.Lz - - ## Read everything else from the file + ## Read everything from the file except the z-coordinates file = jldopen(filepath, "r")["panel$panel"] @@ -883,16 +865,18 @@ function conformal_cubed_sphere_panel(filepath::AbstractString, architecture = C φᶠᶜᵃ = offset_data(zeros(FT, architecture, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2]) φᶜᶠᵃ = offset_data(zeros(FT, architecture, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2]) + ## The vertical coordinates can come out of the regular rectilinear grid! + Lz, z = generate_coordinate(FT, topology, (Nξ, Nη, Nz), halo, z, :z, 3, architecture) + ξ, η = (-1, 1), (-1, 1) conformal_mapping = CubedSphereConformalMapping(ξ, η, rotation) return OrthogonalSphericalShellGrid{TX, TY, TZ}(architecture, Nξ, Nη, Nz, Hx, Hy, Hz, Lz, λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ, φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ, - zᵃᵃᶜ, zᵃᵃᶠ, + z, Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δyᶜᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶜᵃ, Δyᶠᶠᵃ, - Δzᵃᵃᶜ, Δzᵃᵃᶠ, Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, radius, conformal_mapping) @@ -908,8 +892,7 @@ function on_architecture(arch::AbstractSerialArchitecture, grid::OrthogonalSpher :φᶠᶜᵃ, :φᶜᶠᵃ, :φᶠᶠᵃ, - :zᵃᵃᶜ, - :zᵃᵃᶠ) + :z) grid_spacings = (:Δxᶜᶜᵃ, :Δxᶠᶜᵃ, @@ -918,9 +901,7 @@ function on_architecture(arch::AbstractSerialArchitecture, grid::OrthogonalSpher :Δyᶜᶜᵃ, :Δyᶜᶠᵃ, :Δyᶠᶜᵃ, - :Δyᶠᶠᵃ, - :Δzᵃᵃᶜ, - :Δzᵃᵃᶜ) + :Δyᶠᶠᵃ) horizontal_areas = (:Azᶜᶜᵃ, :Azᶠᶜᵃ, @@ -961,8 +942,7 @@ function Adapt.adapt_structure(to, grid::OrthogonalSphericalShellGrid) adapt(to, grid.φᶠᶜᵃ), adapt(to, grid.φᶜᶠᵃ), adapt(to, grid.φᶠᶠᵃ), - adapt(to, grid.zᵃᵃᶜ), - adapt(to, grid.zᵃᵃᶠ), + adapt(to, grid.z), adapt(to, grid.Δxᶜᶜᵃ), adapt(to, grid.Δxᶠᶜᵃ), adapt(to, grid.Δxᶜᶠᵃ), @@ -971,8 +951,6 @@ function Adapt.adapt_structure(to, grid::OrthogonalSphericalShellGrid) adapt(to, grid.Δyᶜᶠᵃ), adapt(to, grid.Δyᶠᶜᵃ), adapt(to, grid.Δyᶠᶠᵃ), - adapt(to, grid.Δzᵃᵃᶜ), - adapt(to, grid.Δzᵃᵃᶠ), adapt(to, grid.Azᶜᶜᵃ), adapt(to, grid.Azᶠᶜᵃ), adapt(to, grid.Azᶜᶠᵃ), @@ -1046,7 +1024,7 @@ function Base.show(io::IO, grid::OrthogonalSphericalShellGrid, withsummary=true) λ₁, λ₂ = minimum(grid.λᶠᶠᵃ[1:Nx_face, 1:Ny_face]), maximum(grid.λᶠᶠᵃ[1:Nx_face, 1:Ny_face]) φ₁, φ₂ = minimum(grid.φᶠᶠᵃ[1:Nx_face, 1:Ny_face]), maximum(grid.φᶠᶠᵃ[1:Nx_face, 1:Ny_face]) - Ωz = domain(topology(grid, 3)(), Nz, grid.zᵃᵃᶠ) + Ωz = domain(topology(grid, 3)(), Nz, grid.z.cᶠ) (λ_center, φ_center), (extent_λ, extent_φ) = get_center_and_extents_of_shell(grid) @@ -1081,7 +1059,7 @@ function Base.show(io::IO, grid::OrthogonalSphericalShellGrid, withsummary=true) φ_summary = "latitude: $(TY) extent $(prettysummary(extent_φ)) degrees" * padding_φ * " " * coordinate_summary(TY, rad2deg.(grid.Δyᶠᶠᵃ[1:Nx_face, 1:Ny_face] ./ grid.radius), "φ") - z_summary = "z: " * dimension_summary(TZ(), "z", Ωz, grid.Δzᵃᵃᶜ, longest - length(z_summary)) + z_summary = "z: " * dimension_summary(TZ(), "z", Ωz, grid.z, longest - length(z_summary)) if withsummary print(io, summary(grid), "\n") @@ -1093,9 +1071,6 @@ function Base.show(io::IO, grid::OrthogonalSphericalShellGrid, withsummary=true) "└── ", z_summary) end -@inline z_domain(grid::OrthogonalSphericalShellGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TZ, grid.Nz, grid.zᵃᵃᶠ) -@inline cpu_face_constructor_z(grid::ZRegOrthogonalSphericalShellGrid) = z_domain(grid) - function with_halo(new_halo, old_grid::OrthogonalSphericalShellGrid; rotation=nothing) size = (old_grid.Nx, old_grid.Ny, old_grid.Nz) @@ -1153,15 +1128,9 @@ end @inline xnodes(grid::OSSG, ℓx, ℓy; with_halos=false) = grid.radius * deg2rad.(λnodes(grid, ℓx, ℓy; with_halos=with_halos)) .* hack_cosd.(φnodes(grid, ℓx, ℓy; with_halos=with_halos)) @inline ynodes(grid::OSSG, ℓx, ℓy; with_halos=false) = grid.radius * deg2rad.(φnodes(grid, ℓx, ℓy; with_halos=with_halos)) -@inline znodes(grid::OSSG, ℓz::Face ; with_halos=false) = with_halos ? grid.zᵃᵃᶠ : - view(grid.zᵃᵃᶠ, interior_indices(ℓz, topology(grid, 3)(), grid.Nz)) -@inline znodes(grid::OSSG, ℓz::Center; with_halos=false) = with_halos ? grid.zᵃᵃᶜ : - view(grid.zᵃᵃᶜ, interior_indices(ℓz, topology(grid, 3)(), grid.Nz)) - # convenience @inline λnodes(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = λnodes(grid, ℓx, ℓy; with_halos) @inline φnodes(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = φnodes(grid, ℓx, ℓy; with_halos) -@inline znodes(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = znodes(grid, ℓz ; with_halos) @inline xnodes(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = xnodes(grid, ℓx, ℓy; with_halos) @inline ynodes(grid::OSSG, ℓx, ℓy, ℓz; with_halos=false) = ynodes(grid, ℓx, ℓy; with_halos) @@ -1178,20 +1147,15 @@ end @inline xnode(i, j, grid::OSSG, ℓx, ℓy) = grid.radius * deg2rad(λnode(i, j, grid, ℓx, ℓy)) * hack_cosd((φnode(i, j, grid, ℓx, ℓy))) @inline ynode(i, j, grid::OSSG, ℓx, ℓy) = grid.radius * deg2rad(φnode(i, j, grid, ℓx, ℓy)) -@inline znode(k, grid::OSSG, ::Center) = @inbounds grid.zᵃᵃᶜ[k] -@inline znode(k, grid::OSSG, ::Face ) = @inbounds grid.zᵃᵃᶠ[k] - # convenience @inline λnode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = λnode(i, j, grid, ℓx, ℓy) @inline φnode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = φnode(i, j, grid, ℓx, ℓy) -@inline znode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = znode(k, grid, ℓz) @inline xnode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = xnode(i, j, grid, ℓx, ℓy) @inline ynode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = ynode(i, j, grid, ℓx, ℓy) # Definitions for node @inline ξnode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = λnode(i, j, grid, ℓx, ℓy) @inline ηnode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = φnode(i, j, grid, ℓx, ℓy) -@inline rnode(i, j, k, grid::OSSG, ℓx, ℓy, ℓz) = znode(k, grid, ℓz) ξname(::OSSG) = :λ ηname(::OSSG) = :φ @@ -1203,7 +1167,6 @@ rname(::OSSG) = :z @inline xspacings(grid::OSSG, ℓx, ℓy) = xspacings(grid, ℓx, ℓy, nothing) @inline yspacings(grid::OSSG, ℓx, ℓy) = yspacings(grid, ℓx, ℓy, nothing) -@inline zspacings(grid::OSSG, ℓz) = zspacings(grid, nothing, nothing, ℓz) @inline λspacings(grid::OSSG, ℓx, ℓy) = λspacings(grid, ℓx, ℓy, nothing) @inline φspacings(grid::OSSG, ℓx, ℓy) = φspacings(grid, ℓx, ℓy, nothing) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index f148625f15..a6bb01cc4b 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -1,4 +1,4 @@ -struct RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} +struct RectilinearGrid{FT, TX, TY, TZ, CZ, FX, FY, VX, VY, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, CZ, Arch} architecture :: Arch Nx :: Int Ny :: Int @@ -19,10 +19,7 @@ struct RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch} <: Abstract Δyᵃᶜᵃ :: FY yᵃᶠᵃ :: VY yᵃᶜᵃ :: VY - Δzᵃᵃᶠ :: FZ - Δzᵃᵃᶜ :: FZ - zᵃᵃᶠ :: VZ - zᵃᵃᶜ :: VZ + z :: CZ RectilinearGrid{TX, TY, TZ}(arch::Arch, Nx, Ny, Nz, @@ -32,27 +29,26 @@ struct RectilinearGrid{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch} <: Abstract xᶠᵃᵃ :: VX, xᶜᵃᵃ :: VX, Δyᵃᶠᵃ :: FY, Δyᵃᶜᵃ :: FY, yᵃᶠᵃ :: VY, yᵃᶜᵃ :: VY, - Δzᵃᵃᶠ :: FZ, Δzᵃᵃᶜ :: FZ, - zᵃᵃᶠ :: VZ, zᵃᵃᶜ :: VZ) where {Arch, FT, - TX, TY, TZ, - FX, VX, FY, - VY, FZ, VZ} = - new{FT, TX, TY, TZ, FX, FY, FZ, VX, VY, VZ, Arch}(arch, Nx, Ny, Nz, - Hx, Hy, Hz, Lx, Ly, Lz, - Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, - Δyᵃᶠᵃ, Δyᵃᶜᵃ, yᵃᶠᵃ, yᵃᶜᵃ, - Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) + z :: CZ) where {Arch, FT, + TX, TY, TZ, + FX, VX, FY, + VY, CZ} = + new{FT, TX, TY, TZ, CZ, FX, FY, VX, VY, Arch}(arch, Nx, Ny, Nz, + Hx, Hy, Hz, Lx, Ly, Lz, + Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, + Δyᵃᶠᵃ, Δyᵃᶜᵃ, yᵃᶠᵃ, yᵃᶜᵃ, + z) end const RG = RectilinearGrid -const XRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number} -const YRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Number} -const ZRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} -const XYRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number} -const XZRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Any, <:Number} -const YZRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:Number} -const XYZRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number, <:Number} +const XRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Number} +const YRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number} +const ZRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate} +const XYRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:Number} +const XZRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate, <:Number} +const YZRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate, <:Any, <:Number} +const XYZRegularRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate, <:Number, <:Number} regular_dimensions(::XRegularRG) = tuple(1) regular_dimensions(::YRegularRG) = tuple(2) @@ -271,9 +267,9 @@ function RectilinearGrid(architecture::AbstractArchitecture = CPU(), Nx, Ny, Nz = size Hx, Hy, Hz = halo - Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, TX(), Nx, Hx, x, :x, architecture) - Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, TY(), Ny, Hy, y, :y, architecture) - Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, z, :z, architecture) + Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology, size, halo, x, :x, 1, architecture) + Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology, size, halo, y, :y, 2, architecture) + Lz, z = generate_coordinate(FT, topology, size, halo, z, :z, 3, architecture) return RectilinearGrid{TX, TY, TZ}(architecture, Nx, Ny, Nz, @@ -281,7 +277,7 @@ function RectilinearGrid(architecture::AbstractArchitecture = CPU(), Lx, Ly, Lz, Δxᶠᵃᵃ, Δxᶜᵃᵃ, xᶠᵃᵃ, xᶜᵃᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ, yᵃᶠᵃ, yᵃᶜᵃ, - Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ) + z) end """ Validate user input arguments to the `RectilinearGrid` constructor. """ @@ -302,7 +298,6 @@ end x_domain(grid::RectilinearGrid) = domain(topology(grid, 1)(), grid.Nx, grid.xᶠᵃᵃ) y_domain(grid::RectilinearGrid) = domain(topology(grid, 2)(), grid.Ny, grid.yᵃᶠᵃ) -z_domain(grid::RectilinearGrid) = domain(topology(grid, 3)(), grid.Nz, grid.zᵃᵃᶠ) # architecture = CPU() default, assuming that a DataType positional arg # is specifying the floating point type. @@ -322,7 +317,7 @@ function Base.show(io::IO, grid::RectilinearGrid, withsummary=true) Ωx = domain(TX(), grid.Nx, grid.xᶠᵃᵃ) Ωy = domain(TY(), grid.Ny, grid.yᵃᶠᵃ) - Ωz = domain(TZ(), grid.Nz, grid.zᵃᵃᶠ) + Ωz = domain(TZ(), grid.Nz, grid.z.cᶠ) x_summary = domain_summary(TX(), "x", Ωx) y_summary = domain_summary(TY(), "y", Ωy) @@ -332,7 +327,7 @@ function Base.show(io::IO, grid::RectilinearGrid, withsummary=true) x_summary = dimension_summary(TX(), "x", Ωx, grid.Δxᶜᵃᵃ, longest - length(x_summary)) y_summary = dimension_summary(TY(), "y", Ωy, grid.Δyᵃᶜᵃ, longest - length(y_summary)) - z_summary = dimension_summary(TZ(), "z", Ωz, grid.Δzᵃᵃᶜ, longest - length(z_summary)) + z_summary = dimension_summary(TZ(), "z", Ωz, grid.z, longest - length(z_summary)) if withsummary print(io, summary(grid), "\n") @@ -361,15 +356,11 @@ function Adapt.adapt_structure(to, grid::RectilinearGrid) Adapt.adapt(to, grid.Δyᵃᶜᵃ), Adapt.adapt(to, grid.yᵃᶠᵃ), Adapt.adapt(to, grid.yᵃᶜᵃ), - Adapt.adapt(to, grid.Δzᵃᵃᶠ), - Adapt.adapt(to, grid.Δzᵃᵃᶜ), - Adapt.adapt(to, grid.zᵃᵃᶠ), - Adapt.adapt(to, grid.zᵃᵃᶜ)) + Adapt.adapt(to, grid.z)) end cpu_face_constructor_x(grid::XRegularRG) = x_domain(grid) cpu_face_constructor_y(grid::YRegularRG) = y_domain(grid) -cpu_face_constructor_z(grid::ZRegularRG) = z_domain(grid) function constructor_arguments(grid::RectilinearGrid) arch = architecture(grid) @@ -421,7 +412,7 @@ function with_halo(halo, grid::RectilinearGrid) halo = pop_flat_elements(halo, topology(grid)) kwargs[:halo] = halo arch = args[:architecture] - FT = args[:number_type] + FT = args[:number_type] return RectilinearGrid(arch, FT; kwargs...) end @@ -454,22 +445,18 @@ rname(::RG) = :z @inline xnode(i, grid::RG, ::Face) = getnode(grid.xᶠᵃᵃ, i) @inline ynode(j, grid::RG, ::Center) = getnode(grid.yᵃᶜᵃ, j) @inline ynode(j, grid::RG, ::Face) = getnode(grid.yᵃᶠᵃ, j) -@inline znode(k, grid::RG, ::Center) = getnode(grid.zᵃᵃᶜ, k) -@inline znode(k, grid::RG, ::Face) = getnode(grid.zᵃᵃᶠ, k) @inline ξnode(i, j, k, grid::RG, ℓx, ℓy, ℓz) = xnode(i, grid, ℓx) @inline ηnode(i, j, k, grid::RG, ℓx, ℓy, ℓz) = ynode(j, grid, ℓy) -@inline rnode(i, j, k, grid::RG, ℓx, ℓy, ℓz) = znode(k, grid, ℓz) # Convenience definitions for x, y, znode @inline xnode(i, j, k, grid::RG, ℓx, ℓy, ℓz) = xnode(i, grid, ℓx) @inline ynode(i, j, k, grid::RG, ℓx, ℓy, ℓz) = ynode(j, grid, ℓy) -@inline znode(i, j, k, grid::RG, ℓx, ℓy, ℓz) = znode(k, grid, ℓz) function nodes(grid::RectilinearGrid, ℓx, ℓy, ℓz; reshape=false, with_halos=false) x = xnodes(grid, ℓx, ℓy, ℓz; with_halos) y = ynodes(grid, ℓx, ℓy, ℓz; with_halos) - z = znodes(grid, ℓx, ℓy, ℓz; with_halos) + z = znodes(grid, ℓx, ℓy, ℓz; with_halos) if reshape # Here we have to deal with the fact that Flat directions may have @@ -501,22 +488,17 @@ const C = Center @inline xnodes(grid::RG, ℓx::C; with_halos=false) = _property(grid.xᶜᵃᵃ, ℓx, topology(grid, 1), size(grid, 1), with_halos) @inline ynodes(grid::RG, ℓy::F; with_halos=false) = _property(grid.yᵃᶠᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) @inline ynodes(grid::RG, ℓy::C; with_halos=false) = _property(grid.yᵃᶜᵃ, ℓy, topology(grid, 2), size(grid, 2), with_halos) -@inline znodes(grid::RG, ℓz::F; with_halos=false) = _property(grid.zᵃᵃᶠ, ℓz, topology(grid, 3), size(grid, 3), with_halos) -@inline znodes(grid::RG, ℓz::C; with_halos=false) = _property(grid.zᵃᵃᶜ, ℓz, topology(grid, 3), size(grid, 3), with_halos) # convenience @inline xnodes(grid::RG, ℓx, ℓy, ℓz; with_halos=false) = xnodes(grid, ℓx; with_halos) @inline ynodes(grid::RG, ℓx, ℓy, ℓz; with_halos=false) = ynodes(grid, ℓy; with_halos) -@inline znodes(grid::RG, ℓx, ℓy, ℓz; with_halos=false) = znodes(grid, ℓz; with_halos) # Generalized coordinates @inline ξnodes(grid::RG, ℓx; kwargs...) = xnodes(grid, ℓx; kwargs...) @inline ηnodes(grid::RG, ℓy; kwargs...) = ynodes(grid, ℓy; kwargs...) -@inline rnodes(grid::RG, ℓz; kwargs...) = znodes(grid, ℓz; kwargs...) @inline ξnodes(grid::RG, ℓx, ℓy, ℓz; kwargs...) = xnodes(grid, ℓx; kwargs...) @inline ηnodes(grid::RG, ℓx, ℓy, ℓz; kwargs...) = ynodes(grid, ℓy; kwargs...) -@inline rnodes(grid::RG, ℓx, ℓy, ℓz; kwargs...) = znodes(grid, ℓz; kwargs...) @inline isrectilinear(::RG) = true @@ -526,4 +508,3 @@ const C = Center @inline xspacings(grid::RG, ℓx) = xspacings(grid, ℓx, nothing, nothing) @inline yspacings(grid::RG, ℓy) = yspacings(grid, nothing, ℓy, nothing) -@inline zspacings(grid::RG, ℓz) = zspacings(grid, nothing, nothing, ℓz) diff --git a/src/Grids/vertical_coordinate.jl b/src/Grids/vertical_coordinate.jl new file mode 100644 index 0000000000..0167f0d376 --- /dev/null +++ b/src/Grids/vertical_coordinate.jl @@ -0,0 +1,176 @@ +#### +#### Vertical coordinates +#### + +# This file implements everything related to vertical coordinates in Oceananigans. +# Vertical coordinates are independent of the underlying grid type as we support grids that are +# "unstructured" or "curvilinear" only in the horizontal direction. +# For this reason the vertical coodinate is _special_, and it can be implemented once for all grid types. + +abstract type AbstractVerticalCoordinate end + +# Represents a static one-dimensional vertical coordinate. +# +# # Fields +# - `cᶜ::C`: Cell-centered coordinate. +# - `cᶠ::C`: Face-centered coordinate. +# - `Δᶜ::D`: Cell-centered grid spacing. +# - `Δᶠ::D`: Face-centered grid spacing. +struct StaticVerticalCoordinate{C, D} <: AbstractVerticalCoordinate + cᶠ :: C + cᶜ :: C + Δᶠ :: D + Δᶜ :: D +end + +# Represents a z-star three-dimensional vertical coordinate. +# +# # Fields +# - `cᶜ::C`: Cell-centered coordinate. +# - `cᶠ::C`: Face-centered coordinate. +# - `Δᶜ::D`: Cell-centered grid spacing. +# - `Δᶠ::D`: Face-centered grid spacing. +# - `ηⁿ::E`: Surface elevation at the current time step. +# - `e₃ᶜᶜⁿ::CC`: Vertical grid scaling at center-center at the current time step. +# - `e₃ᶠᶜⁿ::FC`: Vertical grid scaling at face-center at the current time step. +# - `e₃ᶜᶠⁿ::CF`: Vertical grid scaling at center-face at the current time step. +# - `e₃ᶠᶠⁿ::FF`: Vertical grid scaling at face-face at the current time step. +# - `e₃ᶜᶜ⁻::CC`: Vertical grid scaling at center-center at the previous time step. +# - `∂t_e₃::CC`: Time derivative of the vertical grid scaling at cell centers. +struct ZStarVerticalCoordinate{C, D, E, CC, FC, CF, FF} <: AbstractVerticalCoordinate + cᶠ :: C + cᶜ :: C + Δᶠ :: D + Δᶜ :: D + ηⁿ :: E + e₃ᶜᶜⁿ :: CC + e₃ᶠᶜⁿ :: FC + e₃ᶜᶠⁿ :: CF + e₃ᶠᶠⁿ :: FF + e₃ᶜᶜ⁻ :: CC + ∂t_e₃ :: CC +end + +# Convenience constructors for Zstar vertical coordinate +ZStarVerticalCoordinate(r_faces::Union{Tuple, AbstractVector}) = ZStarVerticalCoordinate(r_faces, r_faces, [nothing for i in 1:9]...) +ZStarVerticalCoordinate(r⁻::Number, r⁺::Number) = ZStarVerticalCoordinate((r⁻, r⁺), (r⁻, r⁺), [nothing for i in 1:9]...) + +#### +#### Some usefull aliases +#### + +const RegularStaticVerticalCoordinate = StaticVerticalCoordinate{<:Any, <:Number} +const RegularZStarVerticalCoordinate = ZStarVerticalCoordinate{<:Any, <:Number} + +const RegularVerticalCoordinate = Union{RegularStaticVerticalCoordinate, RegularZStarVerticalCoordinate} + +const AbstractZStarGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Bounded, <:ZStarVerticalCoordinate} +const AbstractStaticGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any, <:StaticVerticalCoordinate} +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.Δᶜ)) + +Adapt.adapt_structure(to, coord::ZStarVerticalCoordinate) = + ZStarVerticalCoordinate(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.e₃ᶜᶜⁿ), + Adapt.adapt(to, coord.e₃ᶠᶜⁿ), + Adapt.adapt(to, coord.e₃ᶜᶠⁿ), + Adapt.adapt(to, coord.e₃ᶠᶠⁿ), + Adapt.adapt(to, coord.e₃ᶜᶜ⁻), + Adapt.adapt(to, coord.∂t_e₃)) + +on_architecture(arch, coord::StaticVerticalCoordinate) = + StaticVerticalCoordinate(on_architecture(arch, coord.cᶠ), + on_architecture(arch, coord.cᶜ), + on_architecture(arch, coord.Δᶠ), + on_architecture(arch, coord.Δᶜ)) + +on_architecture(arch, coord::ZStarVerticalCoordinate) = + ZStarVerticalCoordinate(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.e₃ᶜᶜⁿ), + on_architecture(arch, coord.e₃ᶠᶜⁿ), + on_architecture(arch, coord.e₃ᶜᶠⁿ), + on_architecture(arch, coord.e₃ᶠᶠⁿ), + on_architecture(arch, coord.e₃ᶜᶜ⁻), + on_architecture(arch, coord.∂t_e₃)) + +##### +##### 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` (independent of ZStar or not) 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::AbstractZStarGrid) = ZStarVerticalCoordinate(cpu_face_constructor_r(grid)) + +#### +#### Utilities +#### + +function validate_dimension_specification(T, ξ::ZStarVerticalCoordinate, 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 ZStarVerticalCoordinate(cᶠ, cᶜ, args[3:end]...) +end + +# Summaries +coordinate_summary(::Bounded, z::AbstractVerticalCoordinate, name) = + @sprintf("Free-surface following with Δ%s=%s", name, prettysummary(z.Δᶜ)) diff --git a/src/ImmersedBoundaries/ImmersedBoundaries.jl b/src/ImmersedBoundaries/ImmersedBoundaries.jl index 2b45f4af9a..de714053e2 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("zstar_immersed_grid.jl") end # module diff --git a/src/ImmersedBoundaries/active_cells_map.jl b/src/ImmersedBoundaries/active_cells_map.jl index ce1d186eb8..2225937528 100644 --- a/src/ImmersedBoundaries/active_cells_map.jl +++ b/src/ImmersedBoundaries/active_cells_map.jl @@ -72,9 +72,6 @@ function ImmersedBoundaryGrid(grid, ib; active_cells_map::Bool = true) column_map) end -with_halo(halo, ibg::ActiveCellsIBG) = - ImmersedBoundaryGrid(with_halo(halo, ibg.underlying_grid), ibg.immersed_boundary; active_cells_map = true) - @inline active_cell(i, j, k, ibg) = !immersed_cell(i, j, k, ibg) @inline active_column(i, j, k, grid, column) = column[i, j, k] != 0 diff --git a/src/ImmersedBoundaries/grid_fitted_bottom.jl b/src/ImmersedBoundaries/grid_fitted_bottom.jl index 9364bba3e3..3c70dfbd97 100644 --- a/src/ImmersedBoundaries/grid_fitted_bottom.jl +++ b/src/ImmersedBoundaries/grid_fitted_bottom.jl @@ -2,7 +2,7 @@ using Adapt using CUDA: CuArray using OffsetArrays: OffsetArray using Oceananigans.Utils: getnamewrapper -using Oceananigans.Grids: total_size +using Oceananigans.Grids: total_size, rnode using Oceananigans.Fields: fill_halo_regions! using Oceananigans.BoundaryConditions: FBC using Printf @@ -112,18 +112,20 @@ compute_numerical_bottom_height!(bottom_field, grid, ib) = @kernel function _compute_numerical_bottom_height!(bottom_field, grid, ib::GridFittedBottom) i, j = @index(Global, NTuple) zb = @inbounds bottom_field[i, j, 1] - @inbounds bottom_field[i, j, 1] = znode(i, j, 1, grid, c, c, f) + @inbounds bottom_field[i, j, 1] = rnode(i, j, 1, grid, c, c, f) condition = ib.immersed_condition for k in 1:grid.Nz - z⁺ = znode(i, j, k+1, grid, c, c, f) - z = znode(i, j, k, grid, c, c, c) + z⁺ = rnode(i, j, k+1, grid, c, c, f) + z = rnode(i, j, k, grid, c, c, c) bottom_cell = ifelse(condition isa CenterImmersedCondition, z ≤ zb, z⁺ ≤ zb) @inbounds bottom_field[i, j, 1] = ifelse(bottom_cell, z⁺, bottom_field[i, j, 1]) end end @inline function _immersed_cell(i, j, k, underlying_grid, ib::GridFittedBottom) - z = znode(i, j, k, underlying_grid, c, c, c) + # We use `rnode` for the `immersed_cell` because we do not want to have + # wetting or drying that could happen for a moving grid if we use znode + z = rnode(i, j, k, underlying_grid, c, c, c) zb = @inbounds ib.bottom_height[i, j, 1] return z ≤ zb end @@ -134,7 +136,7 @@ end const AGFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom} -@inline static_column_depthᶜᶜᵃ(i, j, ibg::AGFBIBG) = @inbounds znode(i, j, ibg.Nz+1, ibg, c, c, f) - ibg.immersed_boundary.bottom_height[i, j, 1] +@inline static_column_depthᶜᶜᵃ(i, j, ibg::AGFBIBG) = @inbounds rnode(i, j, ibg.Nz+1, ibg, c, c, f) - ibg.immersed_boundary.bottom_height[i, j, 1] @inline static_column_depthᶜᶠᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶜᶜᵃ(i, j-1, ibg), static_column_depthᶜᶜᵃ(i, j, ibg)) @inline static_column_depthᶠᶜᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶜᶜᵃ(i-1, j, ibg), static_column_depthᶜᶜᵃ(i, j, ibg)) @inline static_column_depthᶠᶠᵃ(i, j, ibg::AGFBIBG) = min(static_column_depthᶠᶜᵃ(i, j-1, ibg), static_column_depthᶠᶜᵃ(i, j, ibg)) diff --git a/src/ImmersedBoundaries/immersed_boundary_nodes.jl b/src/ImmersedBoundaries/immersed_boundary_nodes.jl index 62f39894f8..9149dd4d55 100644 --- a/src/ImmersedBoundaries/immersed_boundary_nodes.jl +++ b/src/ImmersedBoundaries/immersed_boundary_nodes.jl @@ -1,3 +1,4 @@ +import Oceananigans.Grids: cpu_face_constructor_x, cpu_face_constructor_y, cpu_face_constructor_z import Oceananigans.Grids: xspacings, yspacings, zspacings const c = Center() diff --git a/src/ImmersedBoundaries/immersed_grid_metrics.jl b/src/ImmersedBoundaries/immersed_grid_metrics.jl index 48cb239768..a5d8248115 100644 --- a/src/ImmersedBoundaries/immersed_grid_metrics.jl +++ b/src/ImmersedBoundaries/immersed_grid_metrics.jl @@ -1,7 +1,8 @@ using Oceananigans.AbstractOperations: GridMetricOperation import Oceananigans.Grids: coordinates -import Oceananigans.Operators: Δzᵃᵃᶠ, Δzᵃᵃᶜ, intrinsic_vector, extrinsic_vector +import Oceananigans.Operators: Δrᵃᵃᶜ, Δrᵃᵃᶠ, Δzᵃᵃᶜ, Δzᵃᵃᶠ +import Oceananigans.Operators: intrinsic_vector, extrinsic_vector # Grid metrics for ImmersedBoundaryGrid # @@ -12,7 +13,6 @@ import Oceananigans.Operators: Δzᵃᵃᶠ, Δzᵃᵃᶜ, intrinsic_vector, ext for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ) for dir in (:x, :y, :z), operator in (:Δ, :A) - metric = Symbol(operator, dir, LX, LY, LZ) @eval begin import Oceananigans.Operators: $metric @@ -20,6 +20,12 @@ for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ) end end + metric = Symbol(:Δr, LX, LY, LZ) + @eval begin + import Oceananigans.Operators: $metric + @inline $metric(i, j, k, ibg::IBG) = $metric(i, j, k, ibg.underlying_grid) + end + volume = Symbol(:V, LX, LY, LZ) @eval begin import Oceananigans.Operators: $volume @@ -27,14 +33,16 @@ for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ) end end -coordinates(grid::IBG) = coordinates(grid.underlying_grid) - -@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.underlying_grid) +@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) + +coordinates(grid::IBG) = coordinates(grid.underlying_grid) # Extend both 2D and 3D methods @inline intrinsic_vector(i, j, k, ibg::IBG, u, v) = intrinsic_vector(i, j, k, ibg.underlying_grid, u, v) @inline extrinsic_vector(i, j, k, ibg::IBG, u, v) = extrinsic_vector(i, j, k, ibg.underlying_grid, u, v) @inline intrinsic_vector(i, j, k, ibg::IBG, u, v, w) = intrinsic_vector(i, j, k, ibg.underlying_grid, u, v, w) -@inline extrinsic_vector(i, j, k, ibg::IBG, u, v, w) = extrinsic_vector(i, j, k, ibg.underlying_grid, u, v, w) \ No newline at end of file +@inline extrinsic_vector(i, j, k, ibg::IBG, u, v, w) = extrinsic_vector(i, j, k, ibg.underlying_grid, u, v, w) diff --git a/src/ImmersedBoundaries/partial_cell_bottom.jl b/src/ImmersedBoundaries/partial_cell_bottom.jl index 36cc32487b..81f9872f6f 100644 --- a/src/ImmersedBoundaries/partial_cell_bottom.jl +++ b/src/ImmersedBoundaries/partial_cell_bottom.jl @@ -2,6 +2,8 @@ using Oceananigans.Utils: prettysummary using Oceananigans.Fields: fill_halo_regions! using Printf +import Oceananigans.Architectures: on_architecture + ##### ##### PartialCellBottom ##### @@ -77,17 +79,18 @@ end zb = @inbounds bottom_field[i, j, 1] # Cap bottom height at Lz and at rnode(i, j, grid.Nz+1, grid, c, c, f) - domain_bottom = znode(i, j, 1, grid, c, c, f) - domain_top = znode(i, j, grid.Nz+1, grid, c, c, f) + + domain_bottom = rnode(i, j, 1, grid, c, c, f) + domain_top = rnode(i, j, grid.Nz+1, grid, c, c, f) @inbounds bottom_field[i, j, 1] = clamp(zb, domain_bottom, domain_top) adjusted_zb = bottom_field[i, j, 1] ϵ = ib.minimum_fractional_cell_height for k in 1:grid.Nz - z⁻ = znode(i, j, k, grid, c, c, f) - z⁺ = znode(i, j, k+1, grid, c, c, f) - Δz = Δzᶜᶜᶜ(i, j, k, grid) + z⁻ = rnode(i, j, k, grid, c, c, f) + z⁺ = rnode(i, j, k+1, grid, c, c, f) + Δz = Δrᶜᶜᶜ(i, j, k, grid) bottom_cell = (z⁻ ≤ zb) & (z⁺ ≥ zb) capped_zb = min(z⁺ - ϵ * Δz, zb) @@ -135,9 +138,9 @@ Criterion is zb ≥ z - ϵ Δz """ @inline function _immersed_cell(i, j, k, underlying_grid, ib::PartialCellBottom) - z⁺ = znode(i, j, k+1, underlying_grid, c, c, f) + z⁺ = rnode(i, j, k+1, underlying_grid, c, c, f) ϵ = ib.minimum_fractional_cell_height - Δz = Δzᶜᶜᶜ(i, j, k, underlying_grid) + Δz = Δrᶜᶜᶜ(i, j, k, underlying_grid) z★ = z⁺ - Δz * ϵ zb = @inbounds ib.bottom_height[i, j, 1] return z★ < zb @@ -150,12 +153,12 @@ end return !immersed_cell(i, j, k, grid, ib) & immersed_cell(i, j, k-1, grid, ib) end -@inline function Δzᶜᶜᶜ(i, j, k, ibg::PCBIBG) +@inline function Δrᶜᶜᶜ(i, j, k, ibg::PCBIBG) underlying_grid = ibg.underlying_grid ib = ibg.immersed_boundary # Get node at face above and defining nodes on c,c,f - z = znode(i, j, k+1, underlying_grid, c, c, f) + z = rnode(i, j, k+1, underlying_grid, c, c, f) # Get bottom z-coordinate and fractional Δz parameter zb = @inbounds ib.bottom_height[i, j, 1] @@ -163,42 +166,42 @@ end # Are we in a bottom cell? at_the_bottom = bottom_cell(i, j, k, ibg) - full_Δz = Δzᶜᶜᶜ(i, j, k, ibg.underlying_grid) + full_Δz = Δrᶜᶜᶜ(i, j, k, ibg.underlying_grid) partial_Δz = z - zb return ifelse(at_the_bottom, partial_Δz, full_Δz) end -@inline function Δzᶜᶜᶠ(i, j, k, ibg::PCBIBG) +@inline function Δrᶜᶜᶠ(i, j, k, ibg::PCBIBG) just_above_bottom = bottom_cell(i, j, k-1, ibg) - zc = znode(i, j, k, ibg.underlying_grid, c, c, c) - zf = znode(i, j, k, ibg.underlying_grid, c, c, f) + zc = rnode(i, j, k, ibg.underlying_grid, c, c, c) + zf = rnode(i, j, k, ibg.underlying_grid, c, c, f) - full_Δz = Δzᶜᶜᶠ(i, j, k, ibg.underlying_grid) - partial_Δz = zc - zf + Δzᶜᶜᶜ(i, j, k-1, ibg) / 2 + full_Δz = Δrᶜᶜᶠ(i, j, k, ibg.underlying_grid) + partial_Δz = zc - zf + Δrᶜᶜᶜ(i, j, k-1, ibg) / 2 Δz = ifelse(just_above_bottom, partial_Δz, full_Δz) return Δz end -@inline Δzᶠᶜᶜ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶜ(i-1, j, k, ibg), Δzᶜᶜᶜ(i, j, k, ibg)) -@inline Δzᶜᶠᶜ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶜ(i, j-1, k, ibg), Δzᶜᶜᶜ(i, j, k, ibg)) -@inline Δzᶠᶠᶜ(i, j, k, ibg::PCBIBG) = min(Δzᶠᶜᶜ(i, j-1, k, ibg), Δzᶠᶜᶜ(i, j, k, ibg)) - -@inline Δzᶠᶜᶠ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶠ(i-1, j, k, ibg), Δzᶜᶜᶠ(i, j, k, ibg)) -@inline Δzᶜᶠᶠ(i, j, k, ibg::PCBIBG) = min(Δzᶜᶜᶠ(i, j-1, k, ibg), Δzᶜᶜᶠ(i, j, k, ibg)) -@inline Δzᶠᶠᶠ(i, j, k, ibg::PCBIBG) = min(Δzᶠᶜᶠ(i, j-1, k, ibg), Δzᶠᶜᶠ(i, j, k, ibg)) +@inline Δrᶠᶜᶜ(i, j, k, ibg::PCBIBG) = min(Δrᶜᶜᶜ(i-1, j, k, ibg), Δrᶜᶜᶜ(i, j, k, ibg)) +@inline Δrᶜᶠᶜ(i, j, k, ibg::PCBIBG) = min(Δrᶜᶜᶜ(i, j-1, k, ibg), Δrᶜᶜᶜ(i, j, k, ibg)) +@inline Δrᶠᶠᶜ(i, j, k, ibg::PCBIBG) = min(Δrᶠᶜᶜ(i, j-1, k, ibg), Δrᶠᶜᶜ(i, j, k, ibg)) + +@inline Δrᶠᶜᶠ(i, j, k, ibg::PCBIBG) = min(Δrᶜᶜᶠ(i-1, j, k, ibg), Δrᶜᶜᶠ(i, j, k, ibg)) +@inline Δrᶜᶠᶠ(i, j, k, ibg::PCBIBG) = min(Δrᶜᶜᶠ(i, j-1, k, ibg), Δrᶜᶜᶠ(i, j, k, ibg)) +@inline Δrᶠᶠᶠ(i, j, k, ibg::PCBIBG) = min(Δrᶠᶜᶠ(i, j-1, k, ibg), Δrᶠᶜᶠ(i, j, k, ibg)) # Make sure Δz works for horizontally-Flat topologies. # (There's no point in using z-Flat with PartialCellBottom). XFlatPCBIBG = ImmersedBoundaryGrid{<:Any, <:Flat, <:Any, <:Any, <:Any, <:PartialCellBottom} YFlatPCBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Flat, <:Any, <:Any, <:PartialCellBottom} -@inline Δzᶠᶜᶜ(i, j, k, ibg::XFlatPCBIBG) = Δzᶜᶜᶜ(i, j, k, ibg) -@inline Δzᶠᶜᶠ(i, j, k, ibg::XFlatPCBIBG) = Δzᶜᶜᶠ(i, j, k, ibg) -@inline Δzᶜᶠᶜ(i, j, k, ibg::YFlatPCBIBG) = Δzᶜᶜᶜ(i, j, k, ibg) +@inline Δrᶠᶜᶜ(i, j, k, ibg::XFlatPCBIBG) = Δrᶜᶜᶜ(i, j, k, ibg) +@inline Δrᶠᶜᶠ(i, j, k, ibg::XFlatPCBIBG) = Δrᶜᶜᶠ(i, j, k, ibg) +@inline Δrᶜᶠᶜ(i, j, k, ibg::YFlatPCBIBG) = Δrᶜᶜᶜ(i, j, k, ibg) -@inline Δzᶜᶠᶠ(i, j, k, ibg::YFlatPCBIBG) = Δzᶜᶜᶠ(i, j, k, ibg) -@inline Δzᶠᶠᶜ(i, j, k, ibg::XFlatPCBIBG) = Δzᶜᶠᶜ(i, j, k, ibg) -@inline Δzᶠᶠᶜ(i, j, k, ibg::YFlatPCBIBG) = Δzᶠᶜᶜ(i, j, k, ibg) +@inline Δrᶜᶠᶠ(i, j, k, ibg::YFlatPCBIBG) = Δrᶜᶜᶠ(i, j, k, ibg) +@inline Δrᶠᶠᶜ(i, j, k, ibg::XFlatPCBIBG) = Δrᶜᶠᶜ(i, j, k, ibg) +@inline Δrᶠᶠᶜ(i, j, k, ibg::YFlatPCBIBG) = Δrᶠᶜᶜ(i, j, k, ibg) \ No newline at end of file diff --git a/src/ImmersedBoundaries/zstar_immersed_grid.jl b/src/ImmersedBoundaries/zstar_immersed_grid.jl new file mode 100644 index 0000000000..385db35497 --- /dev/null +++ b/src/ImmersedBoundaries/zstar_immersed_grid.jl @@ -0,0 +1,34 @@ +using Oceananigans.Grids: AbstractZStarGrid +using Oceananigans.Operators + +import Oceananigans.Grids: dynamic_column_depthᶜᶜᵃ, + dynamic_column_depthᶜᶠᵃ, + dynamic_column_depthᶠᶜᵃ, + dynamic_column_depthᶠᶠᵃ + +import Oceananigans.Operators: e₃ⁿ, e₃⁻, ∂t_e₃ + +const ZStarImmersedGrid = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:AbstractZStarGrid} +const ZStarGridOfSomeKind = Union{ZStarImmersedGrid, AbstractZStarGrid} + +@inline dynamic_column_depthᶜᶜᵃ(i, j, k, grid::ZStarGridOfSomeKind, η) = @inbounds static_column_depthᶜᶜᵃ(i, j, grid) + η[i, j, k] +@inline dynamic_column_depthᶜᶠᵃ(i, j, k, grid::ZStarGridOfSomeKind, η) = @inbounds static_column_depthᶠᶜᵃ(i, j, grid) + ℑyᵃᶠᵃ(i, j, k, grid, η) +@inline dynamic_column_depthᶠᶜᵃ(i, j, k, grid::ZStarGridOfSomeKind, η) = @inbounds static_column_depthᶜᶠᵃ(i, j, grid) + ℑxᶠᵃᵃ(i, j, k, grid, η) +@inline dynamic_column_depthᶠᶠᵃ(i, j, k, grid::ZStarGridOfSomeKind, η) = @inbounds static_column_depthᶠᶠᵃ(i, j, grid) + ℑxyᶠᶠᵃ(i, j, k, grid, η) + +# Convenience methods +@inline dynamic_column_depthᶜᶜᵃ(i, j, grid) = static_column_depthᶜᶜᵃ(i, j, grid) +@inline dynamic_column_depthᶜᶠᵃ(i, j, grid) = static_column_depthᶜᶠᵃ(i, j, grid) +@inline dynamic_column_depthᶠᶜᵃ(i, j, grid) = static_column_depthᶠᶜᵃ(i, j, grid) +@inline dynamic_column_depthᶠᶠᵃ(i, j, grid) = static_column_depthᶠᶠᵃ(i, j, grid) + +@inline dynamic_column_depthᶜᶜᵃ(i, j, grid::ZStarGridOfSomeKind) = dynamic_column_depthᶜᶜᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline dynamic_column_depthᶜᶠᵃ(i, j, grid::ZStarGridOfSomeKind) = dynamic_column_depthᶜᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline dynamic_column_depthᶠᶜᵃ(i, j, grid::ZStarGridOfSomeKind) = dynamic_column_depthᶠᶜᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline dynamic_column_depthᶠᶠᵃ(i, j, grid::ZStarGridOfSomeKind) = dynamic_column_depthᶠᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) + +# Fallbacks +@inline e₃ⁿ(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = e₃ⁿ(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) +@inline e₃⁻(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = e₃⁻(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) + +@inline ∂t_e₃(i, j, k, ibg::IBG) = ∂t_e₃(i, j, k, ibg.underlying_grid) diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl index 033e304998..262af463cd 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -62,6 +62,9 @@ include("hydrostatic_free_surface_model.jl") include("show_hydrostatic_free_surface_model.jl") include("set_hydrostatic_free_surface_model.jl") +# ZStar implementation +include("z_star_vertical_spacing.jl") + ##### ##### AbstractModel interface ##### @@ -134,7 +137,6 @@ include("compute_hydrostatic_free_surface_tendencies.jl") include("compute_hydrostatic_free_surface_buffers.jl") include("update_hydrostatic_free_surface_model_state.jl") include("hydrostatic_free_surface_ab2_step.jl") -include("hydrostatic_free_surface_rk3_step.jl") include("store_hydrostatic_free_surface_tendencies.jl") include("prescribed_hydrostatic_velocity_fields.jl") include("single_column_model_mode.jl") diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl index 6ce20fd0b7..01807d4f14 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: dynamic_column_depthᶜᶜᵃ, + dynamic_column_depthᶜᶠᵃ, + dynamic_column_depthᶠᶜᵃ, + dynamic_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..20a6dca582 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl @@ -1,39 +1,45 @@ # 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 + + sᶠᶜ = dynamic_column_depthᶠᶜᵃ(i, j, k_top, grid, η) / static_column_depthᶠᶜᵃ(i, j, grid) + sᶜᶠ = dynamic_column_depthᶜᶠᵃ(i, j, k_top, grid, η) / static_column_depthᶜᶠᵃ(i, j, grid) + + @inbounds U[i, j, 1] = Δrᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1] * sᶠᶜ + @inbounds V[i, j, 1] = Δrᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1] * sᶜᶠ 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] * sᶠᶜ + @inbounds V[i, j, 1] += Δrᶜᶠᶜ(i, j, k, grid) * v[i, j, k] * sᶜᶠ 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ᶠᶜ = dynamic_column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = dynamic_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 +49,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 +57,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/step_split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl index 655e2e9068..65ecd3fff6 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl @@ -26,8 +26,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ᶠᶜ = dynamic_column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = dynamic_column_depthᶜᶠᵃ(i, j, k_top, grid, η) @inbounds begin # ∂τ(U) = - ∇η + G @@ -164,5 +164,9 @@ function step_free_surface!(free_surface::SplitExplicitFreeSurface, model, baroc mask_immersed_field!(model.velocities.v) end + # Needed for ZStar to compute the barotropic correction? + # Can we remove it? + fill_halo_regions!(η) + return nothing end diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl index 927489cf50..4fcd297b7c 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) + +# Since the derivative of the moving grid is: +# +# δx(Δy U) + δy(Δx V) ∇ ⋅ U +# ∂t_e₃ = - --------------------- = - -------- +# Az ⋅ H H +# +# The discrete divergence is calculated as: +# +# wᵏ⁺¹ - wᵏ δx(Ax u) + δy(Ay v) Δr ∂t_e₃ +# ---------- = - --------------------- - ---------- +# Δz V Δz +# +# This makes sure that if we sum up till the top of the domain, we get +# +# ∇ ⋅ U +# wᴺᶻ⁺¹ = w⁰ - ------- - ∂t_e₃ ≈ 0 (if w⁰ == 0) +# H +# @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) + ∂te₃ = Δrᶜᶜᶜ(i, j, k-1, grid) * ∂t_e₃(i, j, k-1, grid) + + immersed = immersed_cell(i, j, k-1, grid) + Δw = δh_u + ifelse(immersed, 0, ∂te₃) # 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 6d53944ebd..7cc936db07 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl @@ -85,8 +85,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, @@ -101,3 +100,32 @@ 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 generalized vertical coordinates +##### We advance e₃θ but store θ once e₃ⁿ⁺¹ is known +##### + +@kernel function _ab2_step_tracer_field!(θ, grid, Δt, χ, Gⁿ, G⁻) + i, j, k = @index(Global, NTuple) + + FT = eltype(χ) + C₁ = convert(FT, 1.5) + χ + C₂ = convert(FT, 0.5) + χ + + eⁿ = e₃ⁿ(i, j, k, grid, Center(), Center(), Center()) + e⁻ = e₃⁻(i, j, k, grid, Center(), Center(), Center()) + + @inbounds begin + ∂t_sθ = C₁ * eⁿ * Gⁿ[i, j, k] - C₂ * e⁻ * G⁻[i, j, k] + + # We store temporarily sθ in θ. the unscaled θ will be retrived later on with `unscale_tracers!` + θ[i, j, k] = eⁿ * θ[i, j, k] + convert(FT, Δt) * ∂t_sθ + end +end + +# Fallback! We need to unscale the tracers only in case of +# a grid with a moving vertical cocrdinate, i.e. where e₃ is not constant +unscale_tracers!(tracers, grid; kwargs...) = nothing 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 605849dbee..81fac3bf19 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl @@ -47,6 +47,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, 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, hydrostatic_prognostic_fields(velocities, free_surface, tracers))) @@ -82,10 +83,11 @@ implicitly during time-stepping. model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields) - return ( - U_dot_∇v(i, j, k, grid, advection, velocities) + return ( - U_dot_∇v(i, j, k, grid, advection, velocities) - 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, 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 a78b837017..85e919b464 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl @@ -74,7 +74,7 @@ end function HydrostaticFreeSurfaceTendencyFields(::PrescribedVelocityFields, free_surface, grid, tracer_names) tracer_tendencies = TracerFields(tracer_names, grid) - momentum_tendencies = (u = nothing, v = nothing) + momentum_tendencies = (u = nothing, v = nothing, η = nothing) return merge(momentum_tendencies, tracer_tendencies) end @@ -89,7 +89,6 @@ end @inline datatuple(obj::PrescribedVelocityFields) = (; u = datatuple(obj.u), v = datatuple(obj.v), w = datatuple(obj.w)) ab2_step_velocities!(::PrescribedVelocityFields, args...) = nothing -rk3_substep_velocities!(::PrescribedVelocityFields, args...) = nothing step_free_surface!(::Nothing, model, timestepper, Δt) = nothing compute_w_from_continuity!(::PrescribedVelocityFields, args...; kwargs...) = nothing @@ -99,6 +98,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 89aa0d6ab8..5ecaead713 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -27,9 +27,10 @@ they are called in the end. Finally, the tendencies for the new time-step are co `compute_tendencies = true`. """ update_state!(model::HydrostaticFreeSurfaceModel, callbacks=[]; compute_tendencies = true) = - update_state!(model, model.grid, callbacks; compute_tendencies) + update_state!(model, model.grid, callbacks; compute_tendencies) function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; compute_tendencies = true) + @apply_regionally mask_immersed_model_fields!(model, grid) # Update possible FieldTimeSeries used in the model @@ -49,7 +50,7 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp update_biogeochemical_state!(model.biogeochemistry, model) - compute_tendencies && + compute_tendencies && @apply_regionally compute_tendencies!(model, callbacks) return nothing @@ -70,20 +71,29 @@ function mask_immersed_model_fields!(model, grid) return nothing end -function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters = tuple(w_kernel_parameters(model.grid)), - p_parameters = tuple(p_kernel_parameters(model.grid)), - κ_parameters = tuple(:xyz)) - - grid = model.grid - closure = model.closure +function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters = w_kernel_parameters(model.grid), + p_parameters = p_kernel_parameters(model.grid), + κ_parameters = :xyz) + + grid = model.grid + closure = model.closure + tracers = model.tracers diffusivity = model.diffusivity_fields - - for (wpar, ppar, κpar) in zip(w_parameters, p_parameters, κ_parameters) - compute_w_from_continuity!(model; parameters = wpar) - compute_diffusivities!(diffusivity, closure, model; parameters = κpar) - update_hydrostatic_pressure!(model.pressure.pHY′, architecture(grid), - grid, model.buoyancy, model.tracers; - parameters = ppar) - end + buoyancy = model.buoyancy + + P = model.pressure.pHY′ + arch = architecture(grid) + + # Update the grid and unscale the tracers + update_grid!(model, grid; 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) + + # Update closure diffusivities + compute_diffusivities!(diffusivity, closure, model; parameters = κ_parameters) + return nothing end 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..e7917b1f16 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/z_star_vertical_spacing.jl @@ -0,0 +1,166 @@ +using Oceananigans.Grids +using Oceananigans.ImmersedBoundaries: ZStarGridOfSomeKind + +# TODO: ZStar is not currently working for an ImplicitFreeSurface on +# - `LatitudeLongitudeGrid`s +# - `ImmersedBoundaryGrid`s +# and for an ExplicitFreeSurface on +# - `ImmersedBoundaryGrid`s + +##### +##### ZStar-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; parameters) = nothing + +function update_grid!(model, grid::ZStarGridOfSomeKind; parameters = :xy) + + # Scalings and free surface + e₃ᶜᶜ⁻ = grid.z.e₃ᶜᶜ⁻ + e₃ᶜᶜⁿ = grid.z.e₃ᶜᶜⁿ + e₃ᶠᶜⁿ = grid.z.e₃ᶠᶜⁿ + e₃ᶜᶠⁿ = grid.z.e₃ᶜᶠⁿ + e₃ᶠᶠⁿ = grid.z.e₃ᶠᶠⁿ + ∂t_e₃ = grid.z.∂t_e₃ + ηⁿ = grid.z.ηⁿ + η = model.free_surface.η + + launch!(architecture(grid), grid, parameters, _update_grid_scaling!, + e₃ᶜᶜⁿ, e₃ᶠᶜⁿ, e₃ᶜᶠⁿ, e₃ᶠᶠⁿ, e₃ᶜᶜ⁻, ηⁿ, 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 vertical spacing with available parameters + # No need to fill the halo as the scaling is updated _IN_ the halos + launch!(architecture(grid), grid, parameters, _update_grid_vertical_velocity!, ∂t_e₃, grid, U, V, u, v) + + return nothing +end + +@kernel function _update_grid_scaling!(e₃ᶜᶜⁿ, e₃ᶠᶜⁿ, e₃ᶜᶠⁿ, e₃ᶠᶠⁿ, e₃ᶜᶜ⁻, ηⁿ, 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ᶜᶜ = dynamic_column_depthᶜᶜᵃ(i, j, k_top, grid, η) + Hᶠᶜ = dynamic_column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = dynamic_column_depthᶜᶠᵃ(i, j, k_top, grid, η) + Hᶠᶠ = dynamic_column_depthᶠᶠᵃ(i, j, k_top, grid, η) + + @inbounds begin + e₃ᶜᶜ = ifelse(hᶜᶜ == 0, one(grid), Hᶜᶜ / hᶜᶜ) + e₃ᶠᶜ = ifelse(hᶠᶜ == 0, one(grid), Hᶠᶜ / hᶠᶜ) + e₃ᶜᶠ = ifelse(hᶜᶠ == 0, one(grid), Hᶜᶠ / hᶜᶠ) + e₃ᶠᶠ = ifelse(hᶠᶠ == 0, one(grid), Hᶠᶠ / hᶠᶠ) + + # Update previous scaling + e₃ᶜᶜ⁻[i, j, 1] = e₃ᶜᶜⁿ[i, j, 1] + + # update current scaling + e₃ᶜᶜⁿ[i, j, 1] = e₃ᶜᶜ + e₃ᶠᶜⁿ[i, j, 1] = e₃ᶠᶜ + e₃ᶜᶠⁿ[i, j, 1] = e₃ᶜᶠ + e₃ᶠᶠⁿ[i, j, 1] = e₃ᶠᶠ + + # Update η in the grid + ηⁿ[i, j, 1] = η[i, j, k_top] + end +end + +@kernel function _update_grid_vertical_velocity!(∂t_e₃, 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_e₃[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, model_fields) = zero(grid) +@inline grid_slope_contribution_y(i, j, k, grid, buoyancy, model_fields) = zero(grid) + +@inline grid_slope_contribution_x(i, j, k, grid::ZStarGridOfSomeKind, ::Nothing, model_fields) = zero(grid) +@inline grid_slope_contribution_y(i, j, k, grid::ZStarGridOfSomeKind, ::Nothing, 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::ZStarGridOfSomeKind, buoyancy, model_fields) = + ℑxᶠᵃᵃ(i, j, k, grid, buoyancy_perturbationᶜᶜᶜ, buoyancy.model, model_fields) * ∂x_z(i, j, k, grid) + +@inline grid_slope_contribution_y(i, j, k, grid::ZStarGridOfSomeKind, buoyancy, model_fields) = + ℑyᵃᶠᵃ(i, j, k, grid, buoyancy_perturbationᶜᶜᶜ, buoyancy.model, 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, ::ZStarGridOfSomeKind; kwargs...) = nothing + +tracer_scaling_parameters(param::Symbol, tracers, grid) = KernelParameters((size(grid, 1), size(grid, 2), length(tracers)), (0, 0, 0)) +tracer_scaling_parameters(param::KernelParameters{S, O}, tracers, grid) where {S, O} = KernelParameters((S..., length(tracers)), (O..., 0)) + +function unscale_tracers!(tracers, grid::ZStarGridOfSomeKind; parameters = :xy) + parameters = tracer_scaling_parameters(parameters, tracers, grid) + + launch!(architecture(grid), grid, parameters, _unscale_tracers!, tracers, grid, + Val(grid.Hz), Val(grid.Nz)) + + return nothing +end + +@kernel function _unscale_tracers!(tracers, grid, ::Val{Hz}, ::Val{Nz}) where {Hz, Nz} + i, j, n = @index(Global, NTuple) + + @unroll for k in -Hz+1:Nz+Hz + tracers[n][i, j, k] /= e₃ⁿ(i, j, k, grid, Center(), Center(), Center()) + end +end \ No newline at end of file diff --git a/src/MultiRegion/MultiRegion.jl b/src/MultiRegion/MultiRegion.jl index 4bb0d2e7eb..671eacc12e 100644 --- a/src/MultiRegion/MultiRegion.jl +++ b/src/MultiRegion/MultiRegion.jl @@ -35,7 +35,7 @@ import Oceananigans.Utils: _getregion, sync_all_devices! -abstract type AbstractMultiRegionGrid{FT, TX, TY, TZ, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} end +abstract type AbstractMultiRegionGrid{FT, TX, TY, TZ, Arch} <: AbstractGrid{FT, TX, TY, TZ, Arch} end abstract type AbstractPartition end diff --git a/src/MultiRegion/cubed_sphere_grid.jl b/src/MultiRegion/cubed_sphere_grid.jl index 0cfa3db30e..e68a2621f3 100644 --- a/src/MultiRegion/cubed_sphere_grid.jl +++ b/src/MultiRegion/cubed_sphere_grid.jl @@ -11,7 +11,7 @@ using Distances import Oceananigans.Grids: grid_name -const ConformalCubedSphereGrid{FT, TX, TY, TZ} = MultiRegionGrid{FT, TX, TY, TZ, <:CubedSpherePartition} +const ConformalCubedSphereGrid{FT, TX, TY, TZ, CZ} = MultiRegionGrid{FT, TX, TY, TZ, CZ, <:CubedSpherePartition} """ ConformalCubedSphereGrid(arch=CPU(), FT=Float64; @@ -343,11 +343,14 @@ function ConformalCubedSphereGrid(arch::AbstractArchitecture=CPU(), FT=Float64; new_region_grids = MultiRegionObject(new_region_grids.regional_objects, new_devices) - new_grid = MultiRegionGrid{FT, region_topology...}(arch, - partition, - connectivity, - new_region_grids, - new_devices) + # Propagate the vertical coordinate type in the `MultiRegionGrid` + CZ = typeof(getregion(region_grids, 1).z) + + new_grid = MultiRegionGrid{FT, region_topology..., CZ}(arch, + partition, + connectivity, + new_region_grids, + new_devices) return new_grid end @@ -390,7 +393,9 @@ function ConformalCubedSphereGrid(filepath::AbstractString, arch::AbstractArchit connectivity = CubedSphereConnectivity(devices, partition) - return MultiRegionGrid{FT, panel_topology...}(arch, partition, connectivity, region_grids, devices) + CZ = typeof(getregion(region_grids, 1).z) + + return MultiRegionGrid{FT, panel_topology..., CZ}(arch, partition, connectivity, region_grids, devices) end function with_halo(new_halo, csg::ConformalCubedSphereGrid) diff --git a/src/MultiRegion/multi_region_grid.jl b/src/MultiRegion/multi_region_grid.jl index 264b6a29be..ce7c54d5e2 100644 --- a/src/MultiRegion/multi_region_grid.jl +++ b/src/MultiRegion/multi_region_grid.jl @@ -8,16 +8,16 @@ import Oceananigans.Grids: minimum_xspacing, minimum_yspacing, minimum_zspacing import Oceananigans.Models.HydrostaticFreeSurfaceModels: default_free_surface import Oceananigans.DistributedComputations: reconstruct_global_grid -struct MultiRegionGrid{FT, TX, TY, TZ, P, C, G, D, Arch} <: AbstractMultiRegionGrid{FT, TX, TY, TZ, Arch} +struct MultiRegionGrid{FT, TX, TY, TZ, CZ, P, C, G, D, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, CZ, Arch} architecture :: Arch partition :: P connectivity :: C region_grids :: G devices :: D - MultiRegionGrid{FT, TX, TY, TZ}(arch::A, partition::P, connectivity::C, - region_grids::G, devices::D) where {FT, TX, TY, TZ, P, C, G, D, A} = - new{FT, TX, TY, TZ, P, C, G, D, A}(arch, partition, connectivity, region_grids, devices) + MultiRegionGrid{FT, TX, TY, TZ, CZ}(arch::A, partition::P, connectivity::C, + region_grids::G, devices::D) where {FT, TX, TY, TZ, CZ, P, C, G, D, A} = + new{FT, TX, TY, TZ, CZ, P, C, G, D, A}(arch, partition, connectivity, region_grids, devices) end const ImmersedMultiRegionGrid = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:MultiRegionGrid} @@ -140,10 +140,13 @@ function MultiRegionGrid(global_grid; partition = XPartition(2), region_grids = construct_regionally(construct_grid, args...) + # Propagate the vertical coordinate type in the `MultiRegionGrid` + CZ = typeof(global_grid.z) + ## If we are on GPUs we want to enable peer access, which we do by just copying fake arrays between all devices maybe_enable_peer_access!(devices) - return MultiRegionGrid{FT, global_topo[1], global_topo[2], global_topo[3]}(arch, partition, connectivity, region_grids, devices) + return MultiRegionGrid{FT, global_topo[1], global_topo[2], global_topo[3], CZ}(arch, partition, connectivity, region_grids, devices) end function construct_grid(grid::RectilinearGrid, child_arch, topo, size, extent, args...) @@ -232,10 +235,10 @@ function with_halo(new_halo, mrg::MultiRegionGrid) return MultiRegionGrid(new_global; partition, devices, validate = false) end -function on_architecture(::CPU, mrg::MultiRegionGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} +function on_architecture(::CPU, mrg::MultiRegionGrid{FT, TX, TY, TZ, CZ}) where {FT, TX, TY, TZ, CZ} new_grids = construct_regionally(on_architecture, CPU(), mrg) devices = Tuple(CPU() for i in 1:length(mrg)) - return MultiRegionGrid{FT, TX, TY, TZ}(CPU(), mrg.partition, mrg.connectivity, new_grids, devices) + return MultiRegionGrid{FT, TX, TY, TZ, CZ}(CPU(), mrg.partition, mrg.connectivity, new_grids, devices) end Base.summary(mrg::MultiRegionGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = diff --git a/src/MultiRegion/unified_implicit_free_surface_solver.jl b/src/MultiRegion/unified_implicit_free_surface_solver.jl index a6370ca97c..45080b3572 100644 --- a/src/MultiRegion/unified_implicit_free_surface_solver.jl +++ b/src/MultiRegion/unified_implicit_free_surface_solver.jl @@ -1,7 +1,6 @@ using Oceananigans.Solvers using Oceananigans.Operators using Oceananigans.Architectures -using Oceananigans.Grids: on_architecture using Oceananigans.Fields: Field using Oceananigans.Models.HydrostaticFreeSurfaceModels: compute_vertically_integrated_lateral_areas!, @@ -12,7 +11,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: compute_vertically_integ import Oceananigans.Models.HydrostaticFreeSurfaceModels: build_implicit_step_solver, compute_implicit_free_surface_right_hand_side! -import Oceananigans.Architectures: architecture +import Oceananigans.Architectures: architecture, on_architecture import Oceananigans.Solvers: solve! struct UnifiedImplicitFreeSurfaceSolver{S, R, T} diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 39701328c5..990fa3e3ec 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -11,9 +11,10 @@ export # Grids Center, Face, Periodic, Bounded, Flat, + ZStarVerticalCoordinate, RectilinearGrid, LatitudeLongitudeGrid, OrthogonalSphericalShellGrid, - nodes, xnodes, ynodes, znodes, λnodes, φnodes, - xspacings, yspacings, zspacings, λspacings, φspacings, + nodes, xnodes, ynodes, rnodes, znodes, λnodes, φnodes, + xspacings, yspacings, rspacings, zspacings, λspacings, φspacings, minimum_xspacing, minimum_yspacing, minimum_zspacing, # Immersed boundaries diff --git a/src/Operators/Operators.jl b/src/Operators/Operators.jl index a033f62d9e..6c5ff40c26 100644 --- a/src/Operators/Operators.jl +++ b/src/Operators/Operators.jl @@ -4,6 +4,7 @@ module Operators export Δxᶠᶠᶠ, Δxᶠᶠᶜ, Δxᶠᶜᶠ, Δxᶠᶜᶜ, Δxᶜᶠᶠ, Δxᶜᶠᶜ, Δxᶜᶜᶠ, Δxᶜᶜᶜ export Δyᶠᶠᶠ, Δyᶠᶠᶜ, Δyᶠᶜᶠ, Δyᶠᶜᶜ, Δyᶜᶠᶠ, Δyᶜᶠᶜ, Δyᶜᶜᶠ, Δyᶜᶜᶜ export Δzᶠᶠᶠ, Δzᶠᶠᶜ, Δzᶠᶜᶠ, Δzᶠᶜᶜ, Δzᶜᶠᶠ, Δzᶜᶠᶜ, Δzᶜᶜᶠ, Δzᶜᶜᶜ +export Δrᶠᶠᶠ, Δrᶠᶠᶜ, Δrᶠᶜᶠ, Δrᶠᶜᶜ, Δrᶜᶠᶠ, Δrᶜᶠᶜ, Δrᶜᶜᶠ, Δrᶜᶜᶜ # Areas export Axᶠᶠᶠ, Axᶠᶠᶜ, Axᶠᶜᶠ, Axᶠᶜᶜ, Axᶜᶠᶠ, Axᶜᶠᶜ, Axᶜᶜᶠ, Axᶜᶜᶜ @@ -73,6 +74,9 @@ export ∂xTᶠᶜᶠ, ∂yTᶜᶠᶠ # Reference frame conversion export intrinsic_vector, extrinsic_vector +# Variable grid operators +export e₃ⁿ, e₃⁻, ∂t_e₃ + using Oceananigans.Grids import Oceananigans.Grids: xspacing, yspacing, zspacing @@ -101,9 +105,9 @@ const LLGY = YRegularLLG const LLGZ = ZRegularLLG # On the fly calculations of metrics -const LLGF = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Nothing} -const LLGFX = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Nothing, <:Any, <:Number} -const LLGFY = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Nothing, <:Any, <:Any, <:Number} +const LLGF = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Nothing} +const LLGFX = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Nothing, <:Any, <:Number} +const LLGFY = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Nothing, <:Any, <:Any, <:Number} include("difference_operators.jl") include("interpolation_operators.jl") @@ -118,6 +122,7 @@ include("topology_aware_operators.jl") include("vorticity_operators.jl") include("laplacian_operators.jl") +include("variable_grid_operators.jl") include("vector_rotation_operators.jl") end # module diff --git a/src/Operators/spacings_and_areas_and_volumes.jl b/src/Operators/spacings_and_areas_and_volumes.jl index 649be89975..9d8b7026bd 100644 --- a/src/Operators/spacings_and_areas_and_volumes.jl +++ b/src/Operators/spacings_and_areas_and_volumes.jl @@ -1,4 +1,5 @@ -using Oceananigans.Grids: Center, Face +using Oceananigans.Grids: Center, Face, AbstractGrid +using Oceananigans.Grids: AbstractZStarGrid @inline hack_cosd(φ) = cos(π * φ / 180) @inline hack_sind(φ) = sin(π * φ / 180) @@ -35,6 +36,17 @@ 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) = getspacing(k, grid.z.Δᶜ) +@inline Δzᵃᵃᶠ(i, j, k, grid::AbstractGrid) = getspacing(k, grid.z.Δᶠ) + # Convenience Functions for all grids for LX in (:ᶜ, :ᶠ, :ᵃ), LY in (:ᶜ, :ᶠ, :ᵃ) @@ -56,10 +68,14 @@ for LX in (:ᶜ, :ᶠ, :ᵃ), LY in (:ᶜ, :ᶠ, :ᵃ) z_spacing_1D = Symbol(:Δz, :ᵃ, :ᵃ, LZ) z_spacing_3D = Symbol(:Δz, LX, LY, LZ) + r_spacing_1D = Symbol(:Δr, :ᵃ, :ᵃ, LZ) + r_spacing_3D = Symbol(:Δr, LX, LY, LZ) + @eval begin @inline $x_spacing_3D(i, j, k, grid) = $x_spacing_2D(i, j, k, grid) @inline $y_spacing_3D(i, j, k, grid) = $y_spacing_2D(i, j, k, grid) @inline $z_spacing_3D(i, j, k, grid) = $z_spacing_1D(i, j, k, grid) + @inline $r_spacing_3D(i, j, k, grid) = $r_spacing_1D(i, j, k, grid) end end end @@ -80,12 +96,6 @@ end @inline Δyᶠᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶜᵃ[j] @inline Δyᶜᵃᶠ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶜᵃ[j] -@inline Δzᵃᵃᶠ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶠ[k] -@inline Δzᵃᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶜ[k] -@inline Δzᶜᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶜ[k] -@inline Δzᶠᵃᶜ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶜ[k] -@inline Δzᶜᵃᶠ(i, j, k, grid::RG) = @inbounds grid.Δzᵃᵃᶠ[k] - ## XRegularRG @inline Δxᶠᵃᵃ(i, j, k, grid::RGX) = grid.Δxᶠᵃᵃ @@ -104,15 +114,6 @@ end @inline Δyᶠᵃᶜ(i, j, k, grid::RGY) = grid.Δyᵃᶜᵃ @inline Δyᶜᵃᶠ(i, j, k, grid::RGY) = grid.Δyᵃᶜᵃ -## ZRegularRG - -@inline Δzᵃᵃᶠ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶠ -@inline Δzᵃᵃᶜ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶜ - -@inline Δzᶜᵃᶜ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶜ -@inline Δzᶠᵃᶜ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶜ -@inline Δzᶜᵃᶠ(i, j, k, grid::RGZ) = grid.Δzᵃᵃᶠ - ##### ##### LatitudeLongitudeGrid ##### @@ -129,9 +130,6 @@ end @inline Δyᶜᶜᵃ(i, j, k, grid::LLG) = Δyᶠᶜᵃ(i, j, k, grid) @inline Δyᶠᶠᵃ(i, j, k, grid::LLG) = Δyᶜᶠᵃ(i, j, k, grid) -@inline Δzᵃᵃᶠ(i, j, k, grid::LLG) = @inbounds grid.Δzᵃᵃᶠ[k] -@inline Δzᵃᵃᶜ(i, j, k, grid::LLG) = @inbounds grid.Δzᵃᵃᶜ[k] - ### XRegularLLG with pre-computed metrics @inline Δxᶠᶜᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δxᶠᶜᵃ[j] @@ -144,11 +142,6 @@ end @inline Δyᶜᶠᵃ(i, j, k, grid::LLGY) = grid.Δyᶜᶠᵃ @inline Δyᶠᶜᵃ(i, j, k, grid::LLGY) = grid.Δyᶠᶜᵃ -### ZRegularLLG with pre-computed metrics - -@inline Δzᵃᵃᶠ(i, j, k, grid::LLGZ) = grid.Δzᵃᵃᶠ -@inline Δzᵃᵃᶜ(i, j, k, grid::LLGZ) = grid.Δzᵃᵃᶜ - ## On the fly metrics @inline Δxᶠᶜᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ[i]) * hack_cosd(grid.φᵃᶜᵃ[j]) @@ -185,12 +178,6 @@ end @inline Δyᶜᶠᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δyᶜᶠᵃ[i, j] @inline Δyᶠᶠᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δyᶠᶠᵃ[i, j] -@inline Δzᵃᵃᶜ(i, j, k, grid::OSSG) = @inbounds grid.Δzᵃᵃᶜ[k] -@inline Δzᵃᵃᶠ(i, j, k, grid::OSSG) = @inbounds grid.Δzᵃᵃᶠ[k] - -@inline Δzᵃᵃᶜ(i, j, k, grid::OSSGZ) = grid.Δzᵃᵃᶜ -@inline Δzᵃᵃᶠ(i, j, k, grid::OSSGZ) = grid.Δzᵃᵃᶠ - ##### ##### ##### Areas!! @@ -294,6 +281,11 @@ for LX in (:Center, :Face, :Nothing) @inline $func(i, j, k, grid, ::$LX, ::$LY, ::$LZ) = $metric(i, j, k, grid) end end + + metric_function = Symbol(:Δr, location_code(LXe, LYe, LZe)) + @eval begin + @inline Δr(i, j, k, grid, ::$LX, ::$LY, ::$LZ) = $metric_function(i, j, k, grid) + end end end end diff --git a/src/Operators/variable_grid_operators.jl b/src/Operators/variable_grid_operators.jl new file mode 100644 index 0000000000..24c3eda599 --- /dev/null +++ b/src/Operators/variable_grid_operators.jl @@ -0,0 +1,56 @@ +import Oceananigans.Grids: znode, AbstractZStarGrid + +##### +##### ZStar-specific vertical spacing functions +##### + +const C = Center +const F = Face + +const ZSG = AbstractZStarGrid + +# Fallbacks +@inline e₃ⁿ(i, j, k, grid, ℓx, ℓy, ℓz) = one(grid) +@inline e₃⁻(i, j, k, grid, ℓx, ℓy, ℓz) = one(grid) + +@inline e₃ⁿ(i, j, k, grid::ZSG, ::C, ::C, ℓz) = @inbounds grid.z.e₃ᶜᶜⁿ[i, j, 1] +@inline e₃ⁿ(i, j, k, grid::ZSG, ::F, ::C, ℓz) = @inbounds grid.z.e₃ᶠᶜⁿ[i, j, 1] +@inline e₃ⁿ(i, j, k, grid::ZSG, ::C, ::F, ℓz) = @inbounds grid.z.e₃ᶜᶠⁿ[i, j, 1] +@inline e₃ⁿ(i, j, k, grid::ZSG, ::F, ::F, ℓz) = @inbounds grid.z.e₃ᶠᶠⁿ[i, j, 1] + +# e₃⁻ is needed only at centers +@inline e₃⁻(i, j, k, grid::ZSG, ::C, ::C, ℓz) = @inbounds grid.z.e₃ᶜᶜ⁻[i, j, 1] + +@inline ∂t_e₃(i, j, k, grid) = zero(grid) +@inline ∂t_e₃(i, j, k, grid::ZSG) = @inbounds grid.z.∂t_e₃[i, j, 1] + +#### +#### Vertical spacing functions +#### + +const ZSRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:ZStarVerticalCoordinate} +const ZSLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:ZStarVerticalCoordinate} +const ZSOSG = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:ZStarVerticalCoordinate} + +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::ZSRG) = $rspacing(i, j, k, grid) * e₃ⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::ZSLLG) = $rspacing(i, j, k, grid) * e₃ⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::ZSOSG) = $rspacing(i, j, k, grid) * e₃ⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + end +end + +# znode for an AbstractZStarGrid grid is scaled by the free surface +@inline znode(i, j, k, grid::ZSG, ::C, ::C, ℓz) = @inbounds rnode(i, j, k, grid, C(), C(), ℓz) * e₃ⁿ(i, j, k, grid, C(), C(), ℓz) + grid.z.ηⁿ[i, j, 1] +@inline znode(i, j, k, grid::ZSG, ::F, ::C, ℓz) = @inbounds rnode(i, j, k, grid, F(), C(), ℓz) * e₃ⁿ(i, j, k, grid, F(), C(), ℓz) + ℑxᶠᵃᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline znode(i, j, k, grid::ZSG, ::C, ::F, ℓz) = @inbounds rnode(i, j, k, grid, C(), F(), ℓz) * e₃ⁿ(i, j, k, grid, C(), F(), ℓz) + ℑyᵃᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline znode(i, j, k, grid::ZSG, ::F, ::F, ℓz) = @inbounds rnode(i, j, k, grid, F(), F(), ℓz) * e₃ⁿ(i, j, k, grid, F(), F(), ℓz) + ℑxyᶠᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) diff --git a/src/OutputWriters/output_writer_utils.jl b/src/OutputWriters/output_writer_utils.jl index 7b79755b3b..8c5dc6b596 100644 --- a/src/OutputWriters/output_writer_utils.jl +++ b/src/OutputWriters/output_writer_utils.jl @@ -1,6 +1,6 @@ +using Oceananigans.DistributedComputations using StructArrays: StructArray, replace_storage using Oceananigans.Grids: on_architecture, architecture -using Oceananigans.DistributedComputations using Oceananigans.DistributedComputations: DistributedGrid, Partition using Oceananigans.Fields: AbstractField, indices, boundary_conditions, instantiated_location using Oceananigans.BoundaryConditions: bc_str, FieldBoundaryConditions, ContinuousBoundaryFunction, DiscreteBoundaryFunction diff --git a/src/Solvers/batched_tridiagonal_solver.jl b/src/Solvers/batched_tridiagonal_solver.jl index 2886847938..bebc7c74ce 100644 --- a/src/Solvers/batched_tridiagonal_solver.jl +++ b/src/Solvers/batched_tridiagonal_solver.jl @@ -1,6 +1,5 @@ using Oceananigans.Architectures: on_architecture using Oceananigans.Grids: XDirection, YDirection, ZDirection - import Oceananigans.Architectures: architecture """ diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index 6d8d7ba24a..5c78640839 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -115,6 +115,7 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt ab2_step!(model, Δt) # full step for tracers, fractional step for velocities. tick!(model.clock, Δt) + model.clock.last_Δt = Δt model.clock.last_stage_Δt = Δt # just one stage diff --git a/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl b/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl index 01e322645d..b3d6f95222 100644 --- a/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl +++ b/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl @@ -259,7 +259,6 @@ end ##### Products of viscosity and stress, divergence, vorticity ##### - @inline κ_σᶠᶜᶜ(i, j, k, grid, closure, K, id, clock, fields, σᶠᶜᶜ, args...) = κᶠᶜᶜ(i, j, k, grid, closure, K, id, clock, fields) * σᶠᶜᶜ(i, j, k, grid, args...) @inline κ_σᶜᶠᶜ(i, j, k, grid, closure, K, id, clock, fields, σᶜᶠᶜ, args...) = κᶜᶠᶜ(i, j, k, grid, closure, K, id, clock, fields) * σᶜᶠᶜ(i, j, k, grid, args...) @inline κ_σᶜᶜᶠ(i, j, k, grid, closure, K, id, clock, fields, σᶜᶜᶠ, args...) = κᶜᶜᶠ(i, j, k, grid, closure, K, id, clock, fields) * σᶜᶜᶠ(i, j, k, grid, args...) diff --git a/src/TurbulenceClosures/discrete_diffusion_function.jl b/src/TurbulenceClosures/discrete_diffusion_function.jl index 19d8203783..8d5bc29d7b 100644 --- a/src/TurbulenceClosures/discrete_diffusion_function.jl +++ b/src/TurbulenceClosures/discrete_diffusion_function.jl @@ -1,5 +1,6 @@ using Oceananigans.Operators: ℑxyz using Oceananigans.Utils: instantiate +import Oceananigans.Architectures: on_architecture """ struct DiscreteDiffusionFunction{LX, LY, LZ, P, F} diff --git a/src/TurbulenceClosures/implicit_explicit_time_discretization.jl b/src/TurbulenceClosures/implicit_explicit_time_discretization.jl index 3fa9d7a54e..603ef75d55 100644 --- a/src/TurbulenceClosures/implicit_explicit_time_discretization.jl +++ b/src/TurbulenceClosures/implicit_explicit_time_discretization.jl @@ -1,5 +1,7 @@ using Oceananigans.Grids: AbstractGrid +import Oceananigans.Architectures: on_architecture + abstract type AbstractTimeDiscretization end """ diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl index cfdd407523..53df44fec9 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl @@ -3,6 +3,8 @@ using Oceananigans.AbstractOperations: KernelFunctionOperation using Oceananigans.BuoyancyModels: ∂z_b using Oceananigans.Operators: ℑzᵃᵃᶜ +import Oceananigans.Architectures: on_architecture + struct ConvectiveAdjustmentVerticalDiffusivity{TD, CK, CN, BK, BN} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 1} convective_κz :: CK convective_νz :: CN diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl index 2d4b7bb2eb..9706c49843 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl @@ -4,6 +4,8 @@ using Oceananigans.Operators using Oceananigans.Grids: inactive_node using Oceananigans.Operators: ℑzᵃᵃᶜ +import Oceananigans.Architectures: on_architecture + struct RiBasedVerticalDiffusivity{TD, FT, R, HR} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 1} ν₀ :: FT κ₀ :: FT diff --git a/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl b/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl index 04582652ee..86795b9e24 100644 --- a/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl +++ b/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl @@ -1,4 +1,4 @@ -using Oceananigans.Operators: Δz +using Oceananigans.Operators: Δz, Δr using Oceananigans.Solvers: BatchedTridiagonalSolver, solve! using Oceananigans.ImmersedBoundaries: immersed_peripheral_node, ImmersedBoundaryGrid using Oceananigans.Grids: ZDirection @@ -45,12 +45,18 @@ implicit_diffusion_solver(::ExplicitTimeDiscretization, args...; kwargs...) = no 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 θ + +@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) + # Tracers and horizontal velocities at cell centers in z @inline function ivd_upper_diagonal(i, j, k, grid, closure, K, id, ℓx, ℓy, ::Center, clock, Δt, κz) closure_ij = getclosure(i, j, closure) κᵏ⁺¹ = κz(i, j, k+1, grid, closure_ij, K, id, clock) - Δzᶜₖ = Δz(i, j, k, grid, ℓx, ℓy, c) - Δzᶠₖ₊₁ = Δz(i, j, k+1, grid, ℓx, ℓy, f) + Δzᶜₖ = vertical_spacing(i, j, k, grid, ℓx, ℓy, c) + Δzᶠₖ₊₁ = vertical_spacing(i, j, k+1, grid, ℓx, ℓy, f) du = - Δt * κᵏ⁺¹ / (Δzᶜₖ * Δzᶠₖ₊₁) # This conditional ensures the diagonal is correct @@ -61,8 +67,8 @@ end k = k′ + 1 # Shift index to match LinearAlgebra.Tridiagonal indexing convenction closure_ij = getclosure(i, j, closure) κᵏ = κz(i, j, k, grid, closure_ij, K, id, clock) - Δzᶜₖ = Δz(i, j, k, grid, ℓx, ℓy, c) - Δzᶠₖ = Δz(i, j, k, grid, ℓx, ℓy, f) + Δzᶜₖ = vertical_spacing(i, j, k, grid, ℓx, ℓy, c) + Δzᶠₖ = vertical_spacing(i, j, k, grid, ℓx, ℓy, f) dl = - Δt * κᵏ / (Δzᶜₖ * Δzᶠₖ) # This conditional ensures the diagonal is correct: the lower diagonal does not @@ -80,8 +86,8 @@ end @inline function ivd_upper_diagonal(i, j, k, grid, closure, K, id, ℓx, ℓy, ::Face, clock, Δt, νzᶜᶜᶜ) closure_ij = getclosure(i, j, closure) νᵏ = νzᶜᶜᶜ(i, j, k, grid, closure_ij, K, clock) - Δzᶜₖ = Δz(i, j, k, grid, ℓx, ℓy, c) - Δzᶠₖ = Δz(i, j, k, grid, ℓx, ℓy, f) + Δzᶜₖ = vertical_spacing(i, j, k, grid, ℓx, ℓy, c) + Δzᶠₖ = vertical_spacing(i, j, k, grid, ℓx, ℓy, f) du = - Δt * νᵏ / (Δzᶜₖ * Δzᶠₖ) return ifelse(k < 1, zero(grid), du) end @@ -90,8 +96,8 @@ end k′ = k + 2 # Shift to adjust for Tridiagonal indexing convention closure_ij = getclosure(i, j, closure) νᵏ⁻¹ = νzᶜᶜᶜ(i, j, k′-1, grid, closure_ij, K, clock) - Δzᶜₖ = Δz(i, j, k′, grid, ℓx, ℓy, c) - Δzᶠₖ₋₁ = Δz(i, j, k′-1, grid, ℓx, ℓy, f) + Δzᶜₖ = vertical_spacing(i, j, k′, grid, ℓx, ℓy, c) + Δzᶠₖ₋₁ = vertical_spacing(i, j, k′-1, grid, ℓx, ℓy, f) dl = - Δt * νᵏ⁻¹ / (Δzᶜₖ * Δzᶠₖ₋₁) return ifelse(k < 1, zero(grid), dl) end diff --git a/test/runtests.jl b/test/runtests.jl index a47f0e3d3d..0146132dbe 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 @@ -219,6 +219,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_diagnostics.jl b/test/test_diagnostics.jl index 91a37c532e..d402e29c48 100644 --- a/test/test_diagnostics.jl +++ b/test/test_diagnostics.jl @@ -67,7 +67,7 @@ function advective_timescale_cfl_on_regular_grid(arch, FT) Δx = model.grid.Δxᶜᵃᵃ Δy = model.grid.Δyᵃᶜᵃ - Δz = model.grid.Δzᵃᵃᶜ + Δz = model.grid.z.Δᶜ u₀ = FT(1.2) v₀ = FT(-2.5) @@ -121,7 +121,7 @@ function advective_timescale_cfl_on_lat_lon_grid(arch, FT) # Will be the same at every grid point. Δy_min = CUDA.@allowscalar Oceananigans.Operators.Δyᶜᶠᵃ(1, 1, 1, grid) - Δz = model.grid.Δzᵃᵃᶠ + Δz = model.grid.z.Δᶠ u₀ = FT(1.2) v₀ = FT(-2.5) diff --git a/test/test_distributed_hydrostatic_model.jl b/test/test_distributed_hydrostatic_model.jl index c11d3c5248..9f530d0702 100644 --- a/test/test_distributed_hydrostatic_model.jl +++ b/test/test_distributed_hydrostatic_model.jl @@ -76,6 +76,9 @@ end Nx = 32 Ny = 32 +active_cells_map(grid) = false +active_cells_map(grid::ImmersedBoundaryGrid) = !isnothing(grid.interior_active_cells) + for arch in archs @testset "Testing distributed solid body rotation" begin underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 1), @@ -95,41 +98,42 @@ for arch in archs global_immersed_grid = ImmersedBoundaryGrid(global_underlying_grid, GridFittedBottom(bottom)) for (grid, global_grid) in zip((underlying_grid, immersed_grid, immersed_active_grid), (global_underlying_grid, global_immersed_grid, global_immersed_grid)) - - # "s" for "serial" computation - us, vs, ws, cs, ηs = solid_body_rotation_test(global_grid) - - us = interior(on_architecture(CPU(), us)) - vs = interior(on_architecture(CPU(), vs)) - ws = interior(on_architecture(CPU(), ws)) - cs = interior(on_architecture(CPU(), cs)) - ηs = interior(on_architecture(CPU(), ηs)) - - @info " Testing distributed solid body rotation with architecture $arch on $(typeof(grid).name.wrapper)" - u, v, w, c, η = solid_body_rotation_test(grid) - - cpu_arch = cpu_architecture(arch) - - u = interior(on_architecture(cpu_arch, u)) - v = interior(on_architecture(cpu_arch, v)) - w = interior(on_architecture(cpu_arch, w)) - c = interior(on_architecture(cpu_arch, c)) - η = interior(on_architecture(cpu_arch, η)) - - us = partition(us, cpu_arch, size(u)) - vs = partition(vs, cpu_arch, size(v)) - ws = partition(ws, cpu_arch, size(w)) - cs = partition(cs, cpu_arch, size(c)) - ηs = partition(ηs, cpu_arch, size(η)) - - atol = eps(eltype(grid)) - rtol = sqrt(eps(eltype(grid))) - - @test all(isapprox(u, us; atol, rtol)) - @test all(isapprox(v, vs; atol, rtol)) - @test all(isapprox(w, ws; atol, rtol)) - @test all(isapprox(c, cs; atol, rtol)) - @test all(isapprox(η, ηs; atol, rtol)) + @testset "Test solid rotation on $(summary(grid)) with $(active_cells_map(grid)) on $arch)" begin + # "s" for "serial" computation + us, vs, ws, cs, ηs = solid_body_rotation_test(global_grid) + + us = interior(on_architecture(CPU(), us)) + vs = interior(on_architecture(CPU(), vs)) + ws = interior(on_architecture(CPU(), ws)) + cs = interior(on_architecture(CPU(), cs)) + ηs = interior(on_architecture(CPU(), ηs)) + + @info " Testing distributed solid body rotation with architecture $arch on $(typeof(grid).name.wrapper)" + u, v, w, c, η = solid_body_rotation_test(grid) + + cpu_arch = cpu_architecture(arch) + + u = interior(on_architecture(cpu_arch, u)) + v = interior(on_architecture(cpu_arch, v)) + w = interior(on_architecture(cpu_arch, w)) + c = interior(on_architecture(cpu_arch, c)) + η = interior(on_architecture(cpu_arch, η)) + + us = partition(us, cpu_arch, size(u)) + vs = partition(vs, cpu_arch, size(v)) + ws = partition(ws, cpu_arch, size(w)) + cs = partition(cs, cpu_arch, size(c)) + ηs = partition(ηs, cpu_arch, size(η)) + + atol = eps(eltype(grid)) + rtol = sqrt(eps(eltype(grid))) + + @test all(isapprox(u, us; atol, rtol)) + @test all(isapprox(v, vs; atol, rtol)) + @test all(isapprox(w, ws; atol, rtol)) + @test all(isapprox(c, cs; atol, rtol)) + @test all(isapprox(η, ηs; atol, rtol)) + end end end end diff --git a/test/test_distributed_models.jl b/test/test_distributed_models.jl index 12d45191c9..7d291b2eba 100644 --- a/test/test_distributed_models.jl +++ b/test/test_distributed_models.jl @@ -237,8 +237,8 @@ function test_triply_periodic_local_grid_with_411_ranks() @test local_grid.xᶠᵃᵃ[nx+1] == 0.25*(local_rank+1) @test local_grid.yᵃᶠᵃ[1] == 0 @test local_grid.yᵃᶠᵃ[ny+1] == 2 - @test local_grid.zᵃᵃᶠ[1] == -3 - @test local_grid.zᵃᵃᶠ[nz+1] == 0 + @test local_grid.z.cᶠ[1] == -3 + @test local_grid.z.cᶠ[nz+1] == 0 return nothing end @@ -254,8 +254,8 @@ function test_triply_periodic_local_grid_with_141_ranks() @test local_grid.xᶠᵃᵃ[nx+1] == 1 @test local_grid.yᵃᶠᵃ[1] == 0.5*local_rank @test local_grid.yᵃᶠᵃ[ny+1] == 0.5*(local_rank+1) - @test local_grid.zᵃᵃᶠ[1] == -3 - @test local_grid.zᵃᵃᶠ[nz+1] == 0 + @test local_grid.z.cᶠ[1] == -3 + @test local_grid.z.cᶠ[nz+1] == 0 return nothing end @@ -271,8 +271,8 @@ function test_triply_periodic_local_grid_with_221_ranks() @test local_grid.xᶠᵃᵃ[nx+1] == 0.5*i @test local_grid.yᵃᶠᵃ[1] == j-1 @test local_grid.yᵃᶠᵃ[ny+1] == j - @test local_grid.zᵃᵃᶠ[1] == -3 - @test local_grid.zᵃᵃᶠ[nz+1] == 0 + @test local_grid.z.cᶠ[1] == -3 + @test local_grid.z.cᶠ[nz+1] == 0 return nothing end diff --git a/test/test_dynamics.jl b/test/test_dynamics.jl index 9097150946..974c0d6453 100644 --- a/test/test_dynamics.jl +++ b/test/test_dynamics.jl @@ -52,14 +52,14 @@ function test_ScalarDiffusivity_budget(fieldname, model) set!(model; Dict(fieldname => (x, y, z) -> rand())...) field = fields(model)[fieldname] ν = viscosity(model.closure, nothing) - return test_diffusion_budget(fieldname, field, model, ν, model.grid.Δzᵃᵃᶜ) + return test_diffusion_budget(fieldname, field, model, ν, model.grid.z.Δᶜ) end function test_ScalarBiharmonicDiffusivity_budget(fieldname, model) set!(model; u=0, v=0, w=0, c=0) set!(model; Dict(fieldname => (x, y, z) -> rand())...) field = fields(model)[fieldname] - return test_diffusion_budget(fieldname, field, model, model.closure.ν, model.grid.Δzᵃᵃᶜ, 4) + return test_diffusion_budget(fieldname, field, model, model.closure.ν, model.grid.z.Δᶜ, 4) end function test_diffusion_cosine(fieldname, grid, closure, ξ, tracers=:c) @@ -90,7 +90,7 @@ function test_immersed_diffusion(Nz, z, time_discretization) underlying_grid = RectilinearGrid(size=Nz, z=z, topology=(Flat, Flat, Bounded)) grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(() -> 0); active_cells_map = true) - Δz_min = minimum(underlying_grid.Δzᵃᵃᶜ) + Δz_min = minimum(underlying_grid.z.Δᶜ) model_kwargs = (tracers=:c, buoyancy=nothing, velocities=PrescribedVelocityFields()) full_model = HydrostaticFreeSurfaceModel(; grid=underlying_grid, closure, model_kwargs...) @@ -133,7 +133,7 @@ function test_3D_immersed_diffusion(Nz, z, time_discretization) underlying_grid = RectilinearGrid(size=(9, 9, Nz), x=(0, 1), y=(0, 1), z=z, topology=(Periodic, Periodic, Bounded)) grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bathymetry); active_cells_map = true) - Δz_min = minimum(grid.underlying_grid.Δzᵃᵃᶜ) + Δz_min = minimum(grid.underlying_grid.z.Δᶜ) model_kwargs = (tracers=:c, buoyancy=nothing, velocities=PrescribedVelocityFields()) full_model = HydrostaticFreeSurfaceModel(; grid=underlying_grid, closure, model_kwargs...) diff --git a/test/test_ensemble_hydrostatic_free_surface_models.jl b/test/test_ensemble_hydrostatic_free_surface_models.jl index 95ab1b701c..d0ea1561e6 100644 --- a/test/test_ensemble_hydrostatic_free_surface_models.jl +++ b/test/test_ensemble_hydrostatic_free_surface_models.jl @@ -55,7 +55,7 @@ end @test size(closures) == (3, 2) @test closures[2, 1].background_κz == 1.2 - Δt = 0.01 * grid.Δzᵃᵃᶜ^2 + Δt = 0.01 * grid.z.Δᶜ^2 model_kwargs = (; tracers=:c, buoyancy=nothing, coriolis=nothing) simulation_kwargs = (; Δt, stop_iteration=100) diff --git a/test/test_grids.jl b/test/test_grids.jl index b30b48c1bb..ad14d96192 100644 --- a/test/test_grids.jl +++ b/test/test_grids.jl @@ -46,10 +46,11 @@ function test_regular_rectilinear_correct_coordinate_lengths(FT) @test length(grid.xᶜᵃᵃ) == Nx + 2Hx @test length(grid.yᵃᶜᵃ) == Ny + 2Hy - @test length(grid.zᵃᵃᶜ) == Nz + 2Hz @test length(grid.xᶠᵃᵃ) == Nx + 2Hx @test length(grid.yᵃᶠᵃ) == Ny + 2Hy + 1 - @test length(grid.zᵃᵃᶠ) == Nz + 2Hz + 1 + + @test length(grid.z.cᶜ) == Nz + 2Hz + @test length(grid.z.cᶠ) == Nz + 2Hz + 1 return nothing end @@ -75,11 +76,11 @@ function test_regular_rectilinear_correct_halo_faces(FT) @test grid.xᶠᵃᵃ[0] == - H * Δ @test grid.yᵃᶠᵃ[0] == - H * Δ - @test grid.zᵃᵃᶠ[0] == - H * Δ + @test grid.z.cᶠ[0] == - H * Δ @test grid.xᶠᵃᵃ[N+1] == L # Periodic @test grid.yᵃᶠᵃ[N+2] == L + H * Δ - @test grid.zᵃᵃᶠ[N+2] == L + H * Δ + @test grid.z.cᶠ[N+2] == L + H * Δ return nothing end @@ -94,7 +95,7 @@ function test_regular_rectilinear_correct_first_cells(FT) @test grid.xᶜᵃᵃ[1] == Δ/2 @test grid.yᵃᶜᵃ[1] == Δ/2 - @test grid.zᵃᵃᶜ[1] == Δ/2 + @test grid.z.cᶜ[1] == Δ/2 return nothing end @@ -109,7 +110,7 @@ function test_regular_rectilinear_correct_end_faces(FT) @test grid.xᶠᵃᵃ[N+1] == L @test grid.yᵃᶠᵃ[N+2] == L + Δ - @test grid.zᵃᵃᶠ[N+2] == L + Δ + @test grid.z.cᶠ[N+2] == L + Δ return nothing end @@ -123,10 +124,11 @@ function test_regular_rectilinear_ranges_have_correct_length(FT) @test length(grid.xᶜᵃᵃ) == Nx + 2Hx @test length(grid.yᵃᶜᵃ) == Ny + 2Hy - @test length(grid.zᵃᵃᶜ) == Nz + 2Hz @test length(grid.xᶠᵃᵃ) == Nx + 1 + 2Hx @test length(grid.yᵃᶠᵃ) == Ny + 1 + 2Hy - @test length(grid.zᵃᵃᶠ) == Nz + 1 + 2Hz + + @test length(grid.z.cᶜ) == Nz + 2Hz + @test length(grid.z.cᶠ) == Nz + 1 + 2Hz return nothing end @@ -139,8 +141,8 @@ function test_regular_rectilinear_no_roundoff_error_in_ranges(FT) grid = RectilinearGrid(CPU(), FT, size=(Nx, Ny, Nz), extent=(1, 1, π/2), halo=(1, 1, Hz)) - @test length(grid.zᵃᵃᶜ) == Nz + 2Hz - @test length(grid.zᵃᵃᶠ) == Nz + 2Hz + 1 + @test length(grid.z.cᶜ) == Nz + 2Hz + @test length(grid.z.cᶠ) == Nz + 2Hz + 1 return nothing end @@ -153,14 +155,14 @@ function test_regular_rectilinear_grid_properties_are_same_type(FT) @test grid.Lz isa FT @test grid.Δxᶠᵃᵃ isa FT @test grid.Δyᵃᶠᵃ isa FT - @test grid.Δzᵃᵃᶠ isa FT + @test grid.z.Δᶠ isa FT @test eltype(grid.xᶠᵃᵃ) == FT @test eltype(grid.yᵃᶠᵃ) == FT - @test eltype(grid.zᵃᵃᶠ) == FT + @test eltype(grid.z.cᶠ) == FT @test eltype(grid.xᶜᵃᵃ) == FT @test eltype(grid.yᵃᶜᵃ) == FT - @test eltype(grid.zᵃᵃᶜ) == FT + @test eltype(grid.z.cᶜ) == FT return nothing end @@ -359,11 +361,11 @@ function test_vertically_stretched_grid_properties_are_same_type(FT, arch) @test eltype(grid.xᶜᵃᵃ) == FT @test eltype(grid.yᵃᶠᵃ) == FT @test eltype(grid.yᵃᶜᵃ) == FT - @test eltype(grid.zᵃᵃᶠ) == FT - @test eltype(grid.zᵃᵃᶜ) == FT + @test eltype(grid.z.cᶠ) == FT + @test eltype(grid.z.cᶜ) == FT - @test eltype(grid.Δzᵃᵃᶜ) == FT - @test eltype(grid.Δzᵃᵃᶠ) == FT + @test eltype(grid.z.Δᶜ) == FT + @test eltype(grid.z.Δᶠ) == FT return nothing end @@ -372,10 +374,10 @@ function test_architecturally_correct_stretched_grid(FT, arch, zᵃᵃᶠ) grid = RectilinearGrid(arch, FT, size=(1, 1, length(zᵃᵃᶠ)-1), x=(0, 1), y=(0, 1), z=zᵃᵃᶠ) ArrayType = array_type(arch) - @test grid.zᵃᵃᶠ isa OffsetArray{FT, 1, <:ArrayType} - @test grid.zᵃᵃᶜ isa OffsetArray{FT, 1, <:ArrayType} - @test grid.Δzᵃᵃᶠ isa OffsetArray{FT, 1, <:ArrayType} - @test grid.Δzᵃᵃᶜ isa OffsetArray{FT, 1, <:ArrayType} + @test grid.z.cᶠ isa OffsetArray{FT, 1, <:ArrayType} + @test grid.z.cᶜ isa OffsetArray{FT, 1, <:ArrayType} + @test grid.z.Δᶠ isa OffsetArray{FT, 1, <:ArrayType} + @test grid.z.Δᶜ isa OffsetArray{FT, 1, <:ArrayType} return nothing end @@ -407,20 +409,20 @@ function test_rectilinear_grid_correct_spacings(FT, N) Δzᵃᵃᶜ(k) = zᵃᵃᶠ(k+1) - zᵃᵃᶠ(k) Δzᵃᵃᶠ(k) = zᵃᵃᶜ(k) - zᵃᵃᶜ(k-1) - @test all(isapprox.( grid.zᵃᵃᶠ[1:N+1], zᵃᵃᶠ.(1:N+1) )) - @test all(isapprox.( grid.zᵃᵃᶜ[1:N], zᵃᵃᶜ.(1:N) )) - @test all(isapprox.( grid.Δzᵃᵃᶜ[1:N], Δzᵃᵃᶜ.(1:N) )) + @test all(isapprox.( grid.z.cᶠ[1:N+1], zᵃᵃᶠ.(1:N+1) )) + @test all(isapprox.( grid.z.cᶜ[1:N], zᵃᵃᶜ.(1:N) )) + @test all(isapprox.( grid.z.Δᶜ[1:N], Δzᵃᵃᶜ.(1:N) )) - @test all(isapprox.(zspacings(grid, Face()), reshape(grid.Δzᵃᵃᶠ[1:N+1], 1, 1, N+1))) - @test all(isapprox.(zspacings(grid, Center()), reshape(grid.Δzᵃᵃᶜ[1:N], 1, 1, N))) + @test all(isapprox.(zspacings(grid, Face()), reshape(grid.z.Δᶠ[1:N+1], 1, 1, N+1))) + @test all(isapprox.(zspacings(grid, Center()), reshape(grid.z.Δᶜ[1:N], 1, 1, N))) - @test zspacing(1, 1, 2, grid, Center(), Center(), Face()) == grid.Δzᵃᵃᶠ[2] + @test zspacing(1, 1, 2, grid, Center(), Center(), Face()) == grid.z.Δᶠ[2] - @test minimum_zspacing(grid, Center(), Center(), Center()) ≈ minimum(grid.Δzᵃᵃᶜ[1:grid.Nz]) + @test minimum_zspacing(grid, Center(), Center(), Center()) ≈ minimum(grid.z.Δᶜ[1:grid.Nz]) # Note that Δzᵃᵃᶠ[1] involves a halo point, which is not directly determined by # the user-supplied zᵃᵃᶠ - @test all(isapprox.( grid.Δzᵃᵃᶠ[2:N], Δzᵃᵃᶠ.(2:N) )) + @test all(isapprox.( grid.z.Δᶠ[2:N], Δzᵃᵃᶠ.(2:N) )) return nothing end @@ -447,8 +449,8 @@ function test_basic_lat_lon_bounded_domain(FT) @test grid.Δλᶠᵃᵃ == 10 @test grid.Δφᵃᶠᵃ == 5 - @test grid.Δzᵃᵃᶜ == 1 - @test grid.Δzᵃᵃᶠ == 1 + @test grid.z.Δᶜ == 1 + @test grid.z.Δᶠ == 1 @test length(grid.λᶠᵃᵃ) == Nλ + 2Hλ + 1 @test length(grid.λᶜᵃᵃ) == Nλ + 2Hλ @@ -496,8 +498,8 @@ function test_basic_lat_lon_periodic_domain(FT) @test grid.Δλᶠᵃᵃ == 10 @test grid.Δφᵃᶠᵃ == 5 - @test grid.Δzᵃᵃᶜ == 1 - @test grid.Δzᵃᵃᶠ == 1 + @test grid.z.Δᶜ == 1 + @test grid.z.Δᶠ == 1 @test length(grid.λᶠᵃᵃ) == Nλ + 2Hλ @test length(grid.λᶜᵃᵃ) == Nλ + 2Hλ @@ -542,7 +544,7 @@ function test_basic_lat_lon_general_grid(FT) grid_reg = LatitudeLongitudeGrid(CPU(), FT, size=grid_size, halo=halo, latitude=lat, longitude=lon, z=zᵣ) - @test typeof(grid_reg.Δzᵃᵃᶜ) == typeof(grid_reg.Δzᵃᵃᶠ) == FT + @test typeof(grid_reg.z.Δᶜ) == typeof(grid_reg.z.Δᶠ) == FT @test all(xspacings(grid_reg, Center(), Center()) .== reshape(grid_reg.Δxᶜᶜᵃ[1:Nφ], 1, Nφ, 1)) @test all(xspacings(grid_reg, Center(), Face() ) .== reshape(grid_reg.Δxᶜᶠᵃ[1:Nφ+1], 1, Nφ+1, 1)) @@ -550,8 +552,8 @@ function test_basic_lat_lon_general_grid(FT) @test all(xspacings(grid_reg, Face(), Face()) .== reshape(grid_reg.Δxᶠᶠᵃ[1:Nφ+1], 1, Nφ+1, 1)) @test all(yspacings(grid_reg, Center(), Face()) .== grid_reg.Δyᶜᶠᵃ) @test all(yspacings(grid_reg, Face(), Center()) .== grid_reg.Δyᶠᶜᵃ) - @test all(zspacings(grid_reg, Center()) .== grid_reg.Δzᵃᵃᶜ) - @test all(zspacings(grid_reg, Face()) .== grid_reg.Δzᵃᵃᶠ) + @test all(zspacings(grid_reg, Center()) .== grid_reg.z.Δᶜ) + @test all(zspacings(grid_reg, Face()) .== grid_reg.z.Δᶠ) @test all(xspacings(grid_reg, Center(), Center(), Center()) .== xspacings(grid_reg, Center(), Center())) @test all(xspacings(grid_reg, Face(), Face(), Center()) .== xspacings(grid_reg, Face(), Face())) @@ -564,8 +566,8 @@ function test_basic_lat_lon_general_grid(FT) @test xspacing(1, 2, 3, grid_reg, Center(), Face(), Center()) == grid_reg.Δxᶜᶠᵃ[2] @test yspacing(1, 2, 3, grid_reg, Center(), Face(), Center()) == grid_reg.Δyᶜᶠᵃ @test yspacing(1, 2, 3, grid_reg, Face(), Center(), Center()) == grid_reg.Δyᶠᶜᵃ - @test zspacing(1, 2, 3, grid_reg, Center(), Center(), Face() ) == grid_reg.Δzᵃᵃᶠ - @test zspacing(1, 2, 3, grid_reg, Center(), Center(), Center()) == grid_reg.Δzᵃᵃᶜ + @test zspacing(1, 2, 3, grid_reg, Center(), Center(), Face() ) == grid_reg.z.Δᶠ + @test zspacing(1, 2, 3, grid_reg, Center(), Center(), Center()) == grid_reg.z.Δᶜ @test all(λspacings(grid_reg, Center()) .== grid_reg.Δλᶜᵃᵃ) @test all(λspacings(grid_reg, Face()) .== grid_reg.Δλᶠᵃᵃ) @@ -578,7 +580,7 @@ function test_basic_lat_lon_general_grid(FT) Δλ = grid_reg.Δλᶠᵃᵃ λₛ = (-grid_reg.Lx/2):Δλ:(grid_reg.Lx/2) - Δz = grid_reg.Δzᵃᵃᶜ + Δz = grid_reg.z.Δᶜ zₛ = -Lz:Δz:0 grid_str = LatitudeLongitudeGrid(CPU(), FT, size=grid_size, halo=halo, latitude=lat, longitude=λₛ, z=zₛ) @@ -589,21 +591,21 @@ function test_basic_lat_lon_general_grid(FT) @test length(grid_str.φᵃᶠᵃ) == length(grid_reg.φᵃᶠᵃ) == Nφ + 2Hφ + 1 @test length(grid_str.φᵃᶜᵃ) == length(grid_reg.φᵃᶜᵃ) == Nφ + 2Hφ - @test length(grid_str.zᵃᵃᶠ) == length(grid_reg.zᵃᵃᶠ) == Nz + 2Hz + 1 - @test length(grid_str.zᵃᵃᶜ) == length(grid_reg.zᵃᵃᶜ) == Nz + 2Hz + @test length(grid_str.z.cᶠ) == length(grid_reg.z.cᶠ) == Nz + 2Hz + 1 + @test length(grid_str.z.cᶜ) == length(grid_reg.z.cᶜ) == Nz + 2Hz - @test length(grid_str.Δzᵃᵃᶠ) == Nz + 2Hz + 1 - @test length(grid_str.Δzᵃᵃᶜ) == Nz + 2Hz + @test length(grid_str.z.Δᶠ) == Nz + 2Hz + 1 + @test length(grid_str.z.Δᶜ) == Nz + 2Hz @test all(grid_str.λᶜᵃᵃ == grid_reg.λᶜᵃᵃ) @test all(grid_str.λᶠᵃᵃ == grid_reg.λᶠᵃᵃ) @test all(grid_str.φᵃᶜᵃ == grid_reg.φᵃᶜᵃ) @test all(grid_str.φᵃᶠᵃ == grid_reg.φᵃᶠᵃ) - @test all(grid_str.zᵃᵃᶜ == grid_reg.zᵃᵃᶜ) - @test all(grid_str.zᵃᵃᶠ == grid_reg.zᵃᵃᶠ) + @test all(grid_str.z.cᶜ == grid_reg.z.cᶜ) + @test all(grid_str.z.cᶠ == grid_reg.z.cᶠ) - @test sum(grid_str.Δzᵃᵃᶜ) == grid_reg.Δzᵃᵃᶜ * length(grid_str.Δzᵃᵃᶜ) - @test sum(grid_str.Δzᵃᵃᶠ) == grid_reg.Δzᵃᵃᶠ * length(grid_str.Δzᵃᵃᶠ) + @test sum(grid_str.z.Δᶜ) == grid_reg.z.Δᶜ * length(grid_str.z.Δᶜ) + @test sum(grid_str.z.Δᶠ) == grid_reg.z.Δᶠ * length(grid_str.z.Δᶠ) @test all(xspacings(grid_str, Center(), Center()) .== reshape(grid_str.Δxᶜᶜᵃ[1:Nλ, 1:Nφ], Nλ, Nφ, 1)) @test all(xspacings(grid_str, Center(), Face()) .== reshape(grid_str.Δxᶜᶠᵃ[1:Nλ, 1:Nφ+1], Nλ, Nφ+1, 1)) @@ -613,8 +615,8 @@ function test_basic_lat_lon_general_grid(FT) @test all(yspacings(grid_str, Center(), Face()) .== grid_str.Δyᶜᶠᵃ) @test all(yspacings(grid_str, Face(), Center()) .== grid_str.Δyᶠᶜᵃ) - @test all(zspacings(grid_str, Center()) .== reshape(grid_str.Δzᵃᵃᶜ[1:Nz], 1, 1, Nz)) - @test all(zspacings(grid_str, Face()) .== reshape(grid_str.Δzᵃᵃᶠ[1:Nz+1], 1, 1, Nz+1)) + @test all(zspacings(grid_str, Center()) .== reshape(grid_str.z.Δᶜ[1:Nz], 1, 1, Nz)) + @test all(zspacings(grid_str, Face()) .== reshape(grid_str.z.Δᶠ[1:Nz+1], 1, 1, Nz+1)) @test all(zspacings(grid_str, Center()) .== zspacings(grid_str, Center(), Center(), Center())) @test all(zspacings(grid_str, Face()) .== zspacings(grid_str, Face(), Center(), Face())) @@ -749,8 +751,8 @@ function test_orthogonal_shell_grid_array_sizes_and_spacings(FT) @test all(yspacings(grid, Face(), Center(), Face()) .== yspacings(grid, Face(), Center()) .== grid.Δyᶠᶜᵃ[1:Nx+1, 1:Ny]) @test all(yspacings(grid, Face(), Face(), Face()) .== yspacings(grid, Face(), Face() ) .== grid.Δyᶠᶠᵃ[1:Nx+1, 1:Ny+1]) - @test all(zspacings(grid, Center(), Center(), Face() ) .== zspacings(grid, Face() ) .== grid.Δzᵃᵃᶠ) - @test all(zspacings(grid, Center(), Center(), Center()) .== zspacings(grid, Center()) .== grid.Δzᵃᵃᶜ) + @test all(zspacings(grid, Center(), Center(), Face() ) .== zspacings(grid, Face() ) .== grid.z.Δᶠ) + @test all(zspacings(grid, Center(), Center(), Center()) .== zspacings(grid, Center()) .== grid.z.Δᶜ) return nothing end diff --git a/test/test_hydrostatic_regression.jl b/test/test_hydrostatic_regression.jl index d989ba6c27..7e3b7e7d61 100644 --- a/test/test_hydrostatic_regression.jl +++ b/test/test_hydrostatic_regression.jl @@ -1,4 +1,4 @@ -# include("dependencies_for_runtests.jl") +include("dependencies_for_runtests.jl") include("data_dependencies.jl") using Oceananigans.Grids: topology, XRegularLLG, YRegularLLG, ZRegularLLG diff --git a/test/test_lagrangian_particle_tracking.jl b/test/test_lagrangian_particle_tracking.jl index 6ee6492a3f..e813c9d79f 100644 --- a/test/test_lagrangian_particle_tracking.jl +++ b/test/test_lagrangian_particle_tracking.jl @@ -75,8 +75,8 @@ function run_simple_particle_tracking_tests(grid, timestepper=:QuasiAdamsBashfor ##### Test Boundary restitution ##### - initial_z = CUDA.@allowscalar grid.zᵃᵃᶜ[grid.Nz-1] - top_boundary = CUDA.@allowscalar grid.zᵃᵃᶠ[grid.Nz+1] + initial_z = CUDA.@allowscalar grid.z.cᶜ[grid.Nz-1] + top_boundary = CUDA.@allowscalar grid.z.cᶠ[grid.Nz+1] x, y, z = on_architecture.(Ref(arch), ([0.0], [0.0], [initial_z])) diff --git a/test/test_netcdf_output_writer.jl b/test/test_netcdf_output_writer.jl index a354766176..7d95282f2b 100644 --- a/test/test_netcdf_output_writer.jl +++ b/test/test_netcdf_output_writer.jl @@ -241,15 +241,15 @@ function test_thermal_bubble_netcdf_output(arch) @test ds3["xF"][1] == grid.xᶠᵃᵃ[1] @test ds3["yC"][1] == grid.yᵃᶜᵃ[1] @test ds3["yF"][1] == grid.yᵃᶠᵃ[1] - @test ds3["zC"][1] == grid.zᵃᵃᶜ[1] - @test ds3["zF"][1] == grid.zᵃᵃᶠ[1] + @test ds3["zC"][1] == grid.z.cᶜ[1] + @test ds3["zF"][1] == grid.z.cᶠ[1] @test ds3["xC"][end] == grid.xᶜᵃᵃ[Nx] @test ds3["xF"][end] == grid.xᶠᵃᵃ[Nx] @test ds3["yC"][end] == grid.yᵃᶜᵃ[Ny] @test ds3["yF"][end] == grid.yᵃᶠᵃ[Ny] - @test ds3["zC"][end] == grid.zᵃᵃᶜ[Nz] - @test ds3["zF"][end] == grid.zᵃᵃᶠ[Nz+1] # z is Bounded + @test ds3["zC"][end] == grid.z.cᶜ[Nz] + @test ds3["zF"][end] == grid.z.cᶠ[Nz+1] # z is Bounded @test eltype(ds3["u"]) == Float64 @test eltype(ds3["v"]) == Float64 @@ -300,15 +300,15 @@ function test_thermal_bubble_netcdf_output(arch) @test ds2["xF"][1] == grid.xᶠᵃᵃ[i_slice[1]] @test ds2["yC"][1] == grid.yᵃᶜᵃ[j_slice[1]] @test ds2["yF"][1] == grid.yᵃᶠᵃ[j_slice[1]] - @test ds2["zC"][1] == grid.zᵃᵃᶜ[k_slice[1]] - @test ds2["zF"][1] == grid.zᵃᵃᶠ[k_slice[1]] + @test ds2["zC"][1] == grid.z.cᶜ[k_slice[1]] + @test ds2["zF"][1] == grid.z.cᶠ[k_slice[1]] @test ds2["xC"][end] == grid.xᶜᵃᵃ[i_slice[end]] @test ds2["xF"][end] == grid.xᶠᵃᵃ[i_slice[end]] @test ds2["yC"][end] == grid.yᵃᶜᵃ[j_slice[end]] @test ds2["yF"][end] == grid.yᵃᶠᵃ[j_slice[end]] - @test ds2["zC"][end] == grid.zᵃᵃᶜ[k_slice[end]] - @test ds2["zF"][end] == grid.zᵃᵃᶠ[k_slice[end]] + @test ds2["zC"][end] == grid.z.cᶜ[k_slice[end]] + @test ds2["zF"][end] == grid.z.cᶠ[k_slice[end]] @test eltype(ds2["u"]) == Float32 @test eltype(ds2["v"]) == Float32 @@ -395,15 +395,15 @@ function test_thermal_bubble_netcdf_output_with_halos(arch) @test ds["xF"][1] == grid.xᶠᵃᵃ[1-Hx] @test ds["yC"][1] == grid.yᵃᶜᵃ[1-Hy] @test ds["yF"][1] == grid.yᵃᶠᵃ[1-Hy] - @test ds["zC"][1] == grid.zᵃᵃᶜ[1-Hz] - @test ds["zF"][1] == grid.zᵃᵃᶠ[1-Hz] + @test ds["zC"][1] == grid.z.cᶜ[1-Hz] + @test ds["zF"][1] == grid.z.cᶠ[1-Hz] @test ds["xC"][end] == grid.xᶜᵃᵃ[Nx+Hx] @test ds["xF"][end] == grid.xᶠᵃᵃ[Nx+Hx] @test ds["yC"][end] == grid.yᵃᶜᵃ[Ny+Hy] @test ds["yF"][end] == grid.yᵃᶠᵃ[Ny+Hy] - @test ds["zC"][end] == grid.zᵃᵃᶜ[Nz+Hz] - @test ds["zF"][end] == grid.zᵃᵃᶠ[Nz+Hz+1] # z is Bounded + @test ds["zC"][end] == grid.z.cᶜ[Nz+Hz] + @test ds["zF"][end] == grid.z.cᶠ[Nz+Hz+1] # z is Bounded @test eltype(ds["u"]) == Float64 @test eltype(ds["v"]) == Float64 @@ -507,15 +507,15 @@ function test_netcdf_function_output(arch) @test ds["xF"][1] == grid.xᶠᵃᵃ[1] @test ds["yC"][1] == grid.yᵃᶜᵃ[1] @test ds["yF"][1] == grid.yᵃᶠᵃ[1] - @test ds["zC"][1] == grid.zᵃᵃᶜ[1] - @test ds["zF"][1] == grid.zᵃᵃᶠ[1] + @test ds["zC"][1] == grid.z.cᶜ[1] + @test ds["zF"][1] == grid.z.cᶠ[1] @test ds["xC"][end] == grid.xᶜᵃᵃ[N] @test ds["yC"][end] == grid.yᵃᶜᵃ[N] - @test ds["zC"][end] == grid.zᵃᵃᶜ[N] @test ds["xF"][end] == grid.xᶠᵃᵃ[N] @test ds["yF"][end] == grid.yᵃᶠᵃ[N] - @test ds["zF"][end] == grid.zᵃᵃᶠ[N+1] # z is Bounded + @test ds["zC"][end] == grid.z.cᶜ[N] + @test ds["zF"][end] == grid.z.cᶠ[N+1] # z is Bounded @test ds.attrib["location"] == "Bay of Fundy" @test ds.attrib["onions"] == 7 @@ -849,16 +849,16 @@ function test_netcdf_vertically_stretched_grid_output(arch) @test ds["yC"][1] == grid.yᵃᶜᵃ[1] @test ds["yF"][1] == grid.yᵃᶠᵃ[1] - @test CUDA.@allowscalar ds["zC"][1] == grid.zᵃᵃᶜ[1] - @test CUDA.@allowscalar ds["zF"][1] == grid.zᵃᵃᶠ[1] + @test CUDA.@allowscalar ds["zC"][1] == grid.z.cᶜ[1] + @test CUDA.@allowscalar ds["zF"][1] == grid.z.cᶠ[1] @test ds["xC"][end] == grid.xᶜᵃᵃ[Nx] @test ds["xF"][end] == grid.xᶠᵃᵃ[Nx] @test ds["yC"][end] == grid.yᵃᶜᵃ[Ny] @test ds["yF"][end] == grid.yᵃᶠᵃ[Ny] - @test CUDA.@allowscalar ds["zC"][end] == grid.zᵃᵃᶜ[Nz] - @test CUDA.@allowscalar ds["zF"][end] == grid.zᵃᵃᶠ[Nz+1] # z is Bounded + @test CUDA.@allowscalar ds["zC"][end] == grid.z.cᶜ[Nz] + @test CUDA.@allowscalar ds["zF"][end] == grid.z.cᶠ[Nz+1] # z is Bounded close(ds) rm(nc_filepath) @@ -907,15 +907,15 @@ function test_netcdf_regular_lat_lon_grid_output(arch; immersed = false) @test ds["xF"][1] == grid.λᶠᵃᵃ[1] @test ds["yC"][1] == grid.φᵃᶜᵃ[1] @test ds["yF"][1] == grid.φᵃᶠᵃ[1] - @test ds["zC"][1] == grid.zᵃᵃᶜ[1] - @test ds["zF"][1] == grid.zᵃᵃᶠ[1] + @test ds["zC"][1] == grid.z.cᶜ[1] + @test ds["zF"][1] == grid.z.cᶠ[1] @test ds["xC"][end] == grid.λᶜᵃᵃ[Nx] @test ds["xF"][end] == grid.λᶠᵃᵃ[Nx] @test ds["yC"][end] == grid.φᵃᶜᵃ[Ny] @test ds["yF"][end] == grid.φᵃᶠᵃ[Ny+1] # y is Bounded - @test ds["zC"][end] == grid.zᵃᵃᶜ[Nz] - @test ds["zF"][end] == grid.zᵃᵃᶠ[Nz+1] # z is Bounded + @test ds["zC"][end] == grid.z.cᶜ[Nz] + @test ds["zF"][end] == grid.z.cᶠ[Nz+1] # z is Bounded close(ds) rm(nc_filepath) diff --git a/test/test_split_explicit_vertical_integrals.jl b/test/test_split_explicit_vertical_integrals.jl index 406f31c057..d5a6d6a7ea 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,21 +49,18 @@ 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 # Test 2: Check that vertical integrals work on the CPU(). The following should be "inexact" - Δz = zeros(Nz) - Δz .= grid.Δzᵃᵃᶠ - set_u_check(x, y, z) = cos((π / 2) * z / Lz) set_U_check(x, y) = (sin(0) - (-2 * Lz / (π))) 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,22 +69,20 @@ 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 @testset "Vertical Integral " begin - Δ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.0) + set!(U, 1.0) + 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.0) + set!(U, 1.0) + compute_barotropic_mode!(U, V, grid, u, v, η̅) @test all(Array(interior(U)) .≈ Lz) set_u_check(x, y, z) = sin(x) @@ -93,7 +90,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 +98,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..e47ef65830 --- /dev/null +++ b/test/test_zstar_coordinate.jl @@ -0,0 +1,145 @@ +include("dependencies_for_runtests.jl") + +using Random +using Oceananigans: initialize! +using Oceananigans.ImmersedBoundaries: PartialCellBottom + +function test_zstar_coordinate(model, Ni, Δt) + + ∫bᵢ = Field(Integral(model.tracers.b)) + ∫cᵢ = Field(Integral(model.tracers.c)) + w = model.velocities.w + Nz = model.grid.Nz + + # Testing that at each timestep + # (1) tracers are conserved down to machine precision + # (2) vertical velocities are zero at the top surface + for _ in 1:Ni + time_step!(model, Δt) + end + + ∫b = Field(Integral(model.tracers.b)) + ∫c = Field(Integral(model.tracers.c)) + + @test interior(∫b, 1, 1, 1) ≈ interior(∫bᵢ, 1, 1, 1) + @test interior(∫c, 1, 1, 1) ≈ interior(∫cᵢ, 1, 1, 1) + @test maximum(interior(w, :, :, Nz+1)) < model.grid.Nz * eps(eltype(w)) + + return nothing +end + +function info_message(grid) + msg1 = "Testing z-star coordinates on $(architecture(grid)) on a " + 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" : "" + + return msg1 * msg2 * msg3 * msg4 * msg5 +end + +const C = Center +const F = Face + +@testset "ZStar coordinate scaling tests" begin + @info "testing the ZStar coordinate scalings" + + z = ZStarVerticalCoordinate(-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)) + + @test znode(1, 1, 21, grid, C(), C(), F()) == 0 + @test dynamic_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 e₃ⁿ(1, 1, 1, grid, C(), C(), C()) == 11 / 10 + @test e₃ⁿ(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 dynamic_column_depthᶜᶜᵃ(1, 1, grid) == 11 + @test dynamic_column_depthᶜᶜᵃ(2, 1, grid) == 12 + @test static_column_depthᶜᶜᵃ(1, 1, grid) == 10 + @test static_column_depthᶜᶜᵃ(2, 1, grid) == 10 +end + +@testset "ZStar coordinate simulation testset" begin + z_uniform = ZStarVerticalCoordinate(-20, 0) + z_stretched = ZStarVerticalCoordinate(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 + info_msg = info_message(grid) + + split_free_surface = SplitExplicitFreeSurface(grid; substeps = 20) + + # TODO: Implicit and Explicit free surfaces are not fully supported yet + implicit_free_surface = ImplicitFreeSurface() + explicit_free_surface = ExplicitFreeSurface() + + for free_surface in [split_free_surface] #, implicit_free_surface, explicit_free_surface] + @testset "$info_msg on $(free_surface)" begin + @info " $info_msg of $(free_surface)" + # TODO: minimum_xspacing(grid) on a Immersed GPU grid with ZStarVerticalCoordinate + # fails because it uses too much parameter space. Figure out a way to reduce it + model = HydrostaticFreeSurfaceModel(; grid, + free_surface, + tracers = (:b, :c), + buoyancy = BuoyancyTracer()) + + bᵢ(x, y, z) = x < grid.Lx / 2 ? 0.06 : 0.01 + + set!(model, c = (x, y, z) -> rand(), b = bᵢ) + + test_zstar_coordinate(model, 100, 10) + end + end + end + end + end +end \ No newline at end of file diff --git a/validation/distributed_simulations/distributed_hydrostatic_turbulence.jl b/validation/distributed_simulations/distributed_hydrostatic_turbulence.jl index 6e88bdd049..0156951505 100644 --- a/validation/distributed_simulations/distributed_hydrostatic_turbulence.jl +++ b/validation/distributed_simulations/distributed_hydrostatic_turbulence.jl @@ -65,7 +65,7 @@ end Nx = 32 Ny = 32 -arch = Distributed(CPU(), partition = Partition(2, 2)) +arch = Distributed(CPU(), partition = Partition(2)) # Run the simulation run_simulation(Nx, Ny, arch) diff --git a/validation/implicit_free_surface/geostrophic_adjustment_test.jl b/validation/implicit_free_surface/geostrophic_adjustment_test.jl index 45d38d8987..a3290a8f83 100644 --- a/validation/implicit_free_surface/geostrophic_adjustment_test.jl +++ b/validation/implicit_free_surface/geostrophic_adjustment_test.jl @@ -47,7 +47,7 @@ function geostrophic_adjustment_simulation(free_surface, topology, multi_region; L = grid.Lx / 40 # gaussian width x₀ = grid.Lx / 4 # gaussian center - vᵍ(x, y, z) = -U * (x - x₀) / L * gaussian(x - x₀, L) + vᵍ(x) = -U * (x - x₀) / L * gaussian(x - x₀, L) g = model.free_surface.gravitational_acceleration η = model.free_surface.η @@ -56,7 +56,7 @@ function geostrophic_adjustment_simulation(free_surface, topology, multi_region; ηᵍ(x) = η₀ * gaussian(x - x₀, L) - ηⁱ(x, y, z) = 2 * ηᵍ(x) + ηⁱ(x) = 2 * ηᵍ(x) set!(model, v = vᵍ) set!(model.free_surface.η, ηⁱ) @@ -112,7 +112,12 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: averaging_shape_function AdamsBashforth3Scheme # @inline new_function(t) = averaging_shape_function(t; r = 0.15766, p = 2, q = 2) -@inline new_function(t) = t ≥ 1/2 && t ≤ 3/2 ? 1.0 : 0.0 +@inline new_function(t) = t >= 1/2 && t <= 3/2 ? 1.0 : 0.0 + +splitexplicit_free_surface = SplitExplicitFreeSurface(substeps = 10, + averaging_weighting_function = averaging_shape_function, + timestepper = AdamsBashforth3Scheme()) +explicit_free_surface = ExplicitFreeSurface() topology_types = [(Bounded, Periodic, Bounded), (Periodic, Periodic, Bounded)] topology_types = [topology_types[2]] diff --git a/validation/stratified_couette_flow/stratified_couette_flow.jl b/validation/stratified_couette_flow/stratified_couette_flow.jl index ca8a7ca600..82558572b1 100644 --- a/validation/stratified_couette_flow/stratified_couette_flow.jl +++ b/validation/stratified_couette_flow/stratified_couette_flow.jl @@ -11,7 +11,7 @@ using Oceananigans.Advection: cell_advection_timescale """ Friction velocity. See equation (16) of Vreugdenhil & Taylor (2018). """ function uτ(model, Uavg, U_wall, n) - Nz, Hz, Δz = model.grid.Nz, model.grid.Hz, model.grid.Δzᵃᵃᶜ + Nz, Hz, Δz = model.grid.Nz, model.grid.Hz, model.grid.z.Δᶜ ν = model.closure[n].ν compute!(Uavg) @@ -30,7 +30,7 @@ end """ Heat flux at the wall. See equation (16) of Vreugdenhil & Taylor (2018). """ function q_wall(model, Tavg, Θ_wall, n) - Nz, Hz, Δz = model.grid.Nz, model.grid.Hz, model.grid.Δzᵃᵃᶜ + Nz, Hz, Δz = model.grid.Nz, model.grid.Hz, model.grid.z.Δᶜ # TODO: interface function for extracting diffusivity? κ = model.closure[n].κ.T @@ -261,7 +261,7 @@ function simulate_stratified_couette_flow(; Nxy, Nz, arch=GPU(), h=1, U_wall=1, wmax = maximum(abs, model.velocities.w.data.parent) CFL = simulation.Δt / cell_advection_timescale(model) - Δ = min(model.grid.Δxᶜᵃᵃ, model.grid.Δyᵃᶜᵃ, model.grid.Δzᵃᵃᶜ) + Δ = min(model.grid.Δxᶜᵃᵃ, model.grid.Δyᵃᶜᵃ, model.grid.z.Δᶜ) νmax = maximum(model.diffusivity_fields[n_amd].νₑ.data.parent) κmax = maximum(model.diffusivity_fields[n_amd].κₑ.T.data.parent) νCFL = simulation.Δt / (Δ^2 / νmax) diff --git a/validation/thermal_bubble/thermal_bubble.jl b/validation/thermal_bubble/thermal_bubble.jl index 5358786ae0..eb747510dc 100644 --- a/validation/thermal_bubble/thermal_bubble.jl +++ b/validation/thermal_bubble/thermal_bubble.jl @@ -62,7 +62,7 @@ function print_progress(simulation) u_max = maximum(abs, interior(model.velocities.u)) w_max = maximum(abs, interior(model.velocities.w)) T_min, T_max = extrema(interior(model.tracers.T)) - CFL = max(u_max, w_max) * simulation.Δt / min(model.grid.Δxᶜᵃᵃ, model.grid.Δzᵃᵃᶜ) + CFL = max(u_max, w_max) * simulation.Δt / min(model.grid.Δxᶜᵃᵃ, model.grid.z.Δᶜ) i, t = model.clock.iteration, model.clock.time @info @sprintf("[%06.2f%%] i: %d, t: %.4f, U_max: (%.2e, %.2e), T: (min=%.5f, max=%.5f), CFL: %.4f", diff --git a/validation/z_star_coordinate/baroclinic_double_gyre.jl b/validation/z_star_coordinate/baroclinic_double_gyre.jl new file mode 100644 index 0000000000..bf85e64ed6 --- /dev/null +++ b/validation/z_star_coordinate/baroclinic_double_gyre.jl @@ -0,0 +1,160 @@ +using Oceananigans +using Oceananigans.Units +using Oceananigans.Operators +using Oceananigans.Grids: φnode +using Oceananigans.AbstractOperations: GridMetricOperation +using Printf + +arch = CPU() +z_faces = ZStarVerticalCoordinate((-1800, 0)) +grid = LatitudeLongitudeGrid(arch; size = (60, 60, 2), + latitude = (15, 75), + longitude = (0, 60), + halo = (5, 5, 5), + z = z_faces) + +##### +##### Parameters +##### + +θ⁺ = 30 # ᵒC maximum temperature +θ⁻ = 0 # ᵒC maximum temperature +α = 2e-4 # ᵒC⁻¹ thermal expansion coefficient +ρ₀ = 1000 # kg m⁻³ reference density +g = 9.80665 # m s⁻² gravitational acceleration +λ = 30days # time scale for restoring + +##### +##### Numerics +##### + +Δt = 20minutes + +Δx = minimum_xspacing(grid) +Δy = minimum_yspacing(grid) + +Δs = sqrt(1 / (1 / Δx^2 + 1 / Δy^2)) +sp = sqrt(g * grid.Lz) +CFL = 0.75 +Δτ = Δs / sp * CFL + +substeps = ceil(Int, 3 * Δt / Δτ) + +coriolis = HydrostaticSphericalCoriolis() +momentum_advection = WENOVectorInvariant(vorticity_order = 5) +tracer_advection = WENO(order = 5) +free_surface = SplitExplicitFreeSurface(grid; substeps) + +numerics = (; coriolis, free_surface, momentum_advection, tracer_advection) + +##### +##### Closure +##### + +closure = ConvectiveAdjustmentVerticalDiffusivity(convective_κz = 1.0, + background_κz = 1e-5, + convective_νz = 1e-2, + background_νz = 1e-2) + +##### +##### Boundary Conditions +##### + +@inline function wind_stress(i, j, grid, clock, fields, p) + τ₀ = p.τ₀ + y = (φnode(j, grid, Center()) - p.φ₀) / grid.Ly + + return τ₀ * cos(2π * y) +end + +@inline function buoyancy_restoring(i, j, grid, clock, fields, p) + b = @inbounds fields.b[i, j, 1] + y = (φnode(j, grid, Center()) - p.φ₀) / grid.Ly + b★ = p.Δb * y + + return p.𝓋 * (b - b★) +end + +Δz₀ = 10 # Surface layer thickness [m] + +Δb = α * g * (θ⁺ - θ⁻) # Buoyancy difference + +parameters = (; τ₀ = 0.1 / ρ₀, # Wind stress + φ₀ = 15, # Latitude of southern edge + Δb, # Buoyancy difference + 𝓋 = Δz₀ / λ) # Pumping velocity for restoring + +u_boundary = FluxBoundaryCondition(wind_stress; discrete_form = true, parameters) +b_boundary = FluxBoundaryCondition(buoyancy_restoring; discrete_form = true, parameters) + +no_slip = ValueBoundaryCondition(0.0) + +u_bcs = FieldBoundaryConditions(north = no_slip, south = no_slip, top = u_boundary) +v_bcs = FieldBoundaryConditions(west = no_slip, east = no_slip) +b_bcs = FieldBoundaryConditions(top = b_boundary) + +##### +##### Model +##### + +model = HydrostaticFreeSurfaceModel(; grid, + boundary_conditions = (u = u_bcs, v = v_bcs, b = b_bcs), + buoyancy = BuoyancyTracer(), + tracers = :b, + numerics..., + closure) + + +N² = Δb / grid.Lz +bᵢ(x, y, z) = N² * (grid.Lz + z) + +set!(model, b = bᵢ) + +##### +##### Simulation +##### + +simulation = Simulation(model; Δt, stop_time = 2000days) + +##### +##### Output +##### + +Δzᶜᶜ = GridMetricOperation((Center, Center, Center), Oceananigans.AbstractOperations.Δz, model.grid) +Δzᶠᶜ = GridMetricOperation((Face, Center, Center), Oceananigans.AbstractOperations.Δz, model.grid) +Δzᶜᶠ = GridMetricOperation((Center, Face, Center), Oceananigans.AbstractOperations.Δz, model.grid) + +field_outputs = merge(model.velocities, + model.tracers, + model.pressure, + (; Δzᶜᶜ, Δzᶠᶜ, Δzᶜᶠ)) + +function progress(sim) + w = interior(sim.model.velocities.w, :, :, sim.model.grid.Nz+1) + u = sim.model.velocities.u + b = sim.model.tracers.b + + 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("extrema b: %.2e %.2e ", maximum(b), minimum(b)) + 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)) + +simulation.output_writers[:snapshots] = JLD2OutputWriter(model, field_outputs, + overwrite_existing = true, + schedule = TimeInterval(60days), + filename = "baroclinic_double_gyre_new") + +simulation.output_writers[:free_surface] = JLD2OutputWriter(model, (; η = model.free_surface.η), + overwrite_existing = true, + indices = (:, :, grid.Nz+1), + schedule = TimeInterval(60days), + filename = "baroclinic_double_gyre_free_surface_new") + +run!(simulation) \ No newline at end of file diff --git a/validation/z_star_coordinate/internal_tide.jl b/validation/z_star_coordinate/internal_tide.jl new file mode 100644 index 0000000000..dac756bdae --- /dev/null +++ b/validation/z_star_coordinate/internal_tide.jl @@ -0,0 +1,203 @@ +using CairoMakie +using Oceananigans +using Oceananigans.Units +using Oceananigans.ImmersedBoundaries: PartialCellBottom +using Oceananigans.AbstractOperations: GridMetricOperation + +Nx, Nz = 250, 125 + +H = 2kilometers +z_faces = ZStarVerticalCoordinate(-H, 0) + +underlying_grid = RectilinearGrid(size = (Nx, Nz), + x = (-1000kilometers, 1000kilometers), + z = z_faces, + halo = (4, 4), + topology = (Periodic, Flat, Bounded)) + +h₀ = 250meters +width = 20kilometers +hill(x) = h₀ * exp(-x^2 / 2width^2) +bottom(x) = - H + hill(x) + +grid = ImmersedBoundaryGrid(underlying_grid, PartialCellBottom(bottom)) + +coriolis = FPlane(latitude = -45) + +# Now we have everything we require to construct the tidal forcing given a value of the +# excursion parameter. + +T₂ = 12.421hours +ω₂ = 2π / T₂ # radians/sec + +ϵ = 0.1 # excursion parameter + +U_tidal = ϵ * ω₂ * width + +tidal_forcing_amplitude = U_tidal * (ω₂^2 - coriolis.f^2) / ω₂ + +@inline tidal_forcing(x, z, t, p) = p.tidal_forcing_amplitude * sin(p.ω₂ * t) + +u_forcing = Forcing(tidal_forcing, parameters=(; tidal_forcing_amplitude, ω₂)) + +# ## Model + +# We built a `HydrostaticFreeSurfaceModel`: + +model = HydrostaticFreeSurfaceModel(; grid, coriolis, + buoyancy = BuoyancyTracer(), + tracers = :b, + free_surface = SplitExplicitFreeSurface(grid; substeps = 20), + momentum_advection = WENO(), + tracer_advection = WENO(), + forcing = (; u = u_forcing)) + +# We initialize the model with the tidal flow and a linear stratification. + +uᵢ(x, z) = 0 + +Nᵢ² = 1e-4 # [s⁻²] initial buoyancy frequency / stratification +bᵢ(x, z) = Nᵢ² * z + +set!(model, u=uᵢ, b=bᵢ) + +# Now let's build a `Simulation`. + +Δt = 5minutes +stop_time = 4days + +simulation = Simulation(model; Δt, stop_time, stop_iteration = 10) + +# We add a callback to print a message about how the simulation is going, + +using Printf + +wall_clock = Ref(time_ns()) + +dz = GridMetricOperation((Center, Center, Center), Oceananigans.AbstractOperations.Δz, model.grid) +∫b_init = sum(model.tracers.b * dz) / sum(dz) + +function progress(sim) + elapsed = 1e-9 * (time_ns() - wall_clock[]) + + ∫b = sum(model.tracers.b * dz) / sum(dz) + + msg = @sprintf("iteration: %d, time: %s, wall time: %s, max|w|: %6.3e, max|u|: %6.3e, drift: %6.3e\n", + iteration(sim), prettytime(sim), prettytime(elapsed), + maximum(abs, interior(w, :, :, grid.Nz+1)), maximum(abs, simulation.model.velocities.u), + ∫b - ∫b_init) + + wall_clock[] = time_ns() + + @info msg + + return nothing +end + +simulation.callbacks[:progress] = Callback(progress, IterationInterval(100)) + +b = model.tracers.b +u, v, w = model.velocities + +U = Field(Average(u)) + +u′ = u - U + +N² = ∂z(b) + +filename = "internal_tide" +save_fields_interval = 30minutes + +simulation.output_writers[:fields] = JLD2OutputWriter(model, (; u, u′, w, b, N²); + filename, + schedule = TimeInterval(save_fields_interval), + overwrite_existing = true) + +# We are ready -- let's run! +run!(simulation) + +# # ## Load output + +# # # First, we load the saved velocities and stratification output as `FieldTimeSeries`es. + +# saved_output_filename = filename * ".jld2" + +# u′_t = FieldTimeSeries(saved_output_filename, "u′") +# w_t = FieldTimeSeries(saved_output_filename, "w") +# N²_t = FieldTimeSeries(saved_output_filename, "N²") + +# umax = maximum(abs, u′_t[end]) +# wmax = maximum(abs, w_t[end]) + +# times = u′_t.times +# nothing #hide + +# # We retrieve each field's coordinates and convert from meters to kilometers. + +# xu, _, zu = nodes(u′_t[1]) +# xw, _, zw = nodes(w_t[1]) +# xN², _, zN² = nodes(N²_t[1]) + +# xu = xu ./ 1e3 +# xw = xw ./ 1e3 +# xN² = xN² ./ 1e3 +# zu = zu ./ 1e3 +# zw = zw ./ 1e3 +# zN² = zN² ./ 1e3 +# nothing #hide + +# # ## Visualize + +# # Now we can visualize our resutls! We use `CairoMakie` here. On a system with OpenGL +# # `using GLMakie` is more convenient as figures will be displayed on the screen. +# # +# # We use Makie's `Observable` to animate the data. To dive into how `Observable`s work we +# # refer to [Makie.jl's Documentation](https://makie.juliaplots.org/stable/documentation/nodes/index.html). + +# using CairoMakie + +# n = Observable(1) + +# title = @lift @sprintf("t = %1.2f days = %1.2f T₂", +# round(times[$n] / day, digits=2) , round(times[$n] / T₂, digits=2)) + +# u′n = @lift u′_t[$n] +# wn = @lift w_t[$n] +# N²n = @lift N²_t[$n] + +# axis_kwargs = (xlabel = "x [km]", +# ylabel = "z [km]", +# limits = ((-grid.Lx/2e3, grid.Lx/2e3), (-grid.Lz/1e3, 0)), # note conversion to kilometers +# titlesize = 20) + +# fig = Figure(size = (700, 900)) + +# fig[1, :] = Label(fig, title, fontsize=24, tellwidth=false) + +# ax_u = Axis(fig[2, 1]; title = "u'-velocity", axis_kwargs...) +# hm_u = heatmap!(ax_u, xu, zu, u′n; nan_color=:gray, colorrange=(-umax, umax), colormap=:balance) +# Colorbar(fig[2, 2], hm_u, label = "m s⁻¹") + +# ax_w = Axis(fig[3, 1]; title = "w-velocity", axis_kwargs...) +# hm_w = heatmap!(ax_w, xw, zw, wn; nan_color=:gray, colorrange=(-wmax, wmax), colormap=:balance) +# Colorbar(fig[3, 2], hm_w, label = "m s⁻¹") + +# ax_N² = Axis(fig[4, 1]; title = "stratification N²", axis_kwargs...) +# hm_N² = heatmap!(ax_N², xN², zN², N²n; nan_color=:gray, colorrange=(0.9Nᵢ², 1.1Nᵢ²), colormap=:magma) +# Colorbar(fig[4, 2], hm_N², label = "s⁻²") + +# fig + +# # Finally, we can record a movie. + +# @info "Making an animation from saved data..." + +# frames = 1:length(times) + +# record(fig, filename * ".mp4", frames, framerate=16) do i +# @info string("Plotting frame ", i, " of ", frames[end]) +# n[] = i +# end +# nothing #hide + +# # ![](internal_tide.mp4) diff --git a/validation/z_star_coordinate/lock_release.jl b/validation/z_star_coordinate/lock_release.jl new file mode 100644 index 0000000000..ec4da9a51b --- /dev/null +++ b/validation/z_star_coordinate/lock_release.jl @@ -0,0 +1,82 @@ +using Oceananigans +using Oceananigans.Grids +using Oceananigans.Units +using Oceananigans.Utils: prettytime +using Oceananigans.Advection: WENOVectorInvariant +using Oceananigans.AbstractOperations: GridMetricOperation +using Printf + +z_faces = ZStarVerticalCoordinate(-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 = FluxFormAdvection(WENO(; order = 5), nothing, WENO(; order = 5)), + buoyancy = BuoyancyTracer(), + closure = nothing, + tracers = :b, + free_surface = SplitExplicitFreeSurface(; substeps = 10)) + +g = model.free_surface.gravitational_acceleration + +model.timestepper.χ = 0.1 + +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 = GridMetricOperation((Center, Center, Center), Oceananigans.AbstractOperations.Δz, model.grid) +∫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(1)) + +run!(simulation) + +using Oceananigans.Fields: OneField + +# # Check tracer conservation +b = FieldTimeSeries("zstar_model.jld2", "b") +Δz = FieldTimeSeries("zstar_model.jld2", "Δz") + +init = sum(Δz[1] * b[1]) / sum(Δz[1]) +drift = [] + +for t in 1:length(b.times) + push!(drift, sum(Δz[t] * b[t]) / sum(Δz[t]) - init) +end +