Yet another ZStar implementation

.buildkite/pipeline.yml

@@ -649,6 +649,42 @@ steps:
           limit: 1
     depends_on: "init_cpu"
 
+#####
+##### Vertical Coordinates tests
+#####
+
+  - label: "🥑 gpu vertical coordinate"
+    env:
+      JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
+      TEST_GROUP: "vertical_coordinate"
+      GPU_TEST: "true"
+    commands:
+      - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      queue: Oceananigans
+      architecture: GPU
+    retry:
+      automatic:
+        - exit_status: 1      
+          limit: 1
+    depends_on: "init_gpu"
+
+  - label: "🥒 cpu vertical coordinate"
+    env:
+      JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
+      TEST_GROUP: "vertical_coordinate"
+      CUDA_VISIBLE_DEVICES: "-1"
+    commands:
+      - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      queue: Oceananigans
+      architecture: CPU
+    retry:
+      automatic:
+        - exit_status: 1      
+          limit: 1
+    depends_on: "init_cpu"
+
 #####
 ##### Enzyme extension tests
 #####

src/Advection/Advection.jl

@@ -34,6 +34,7 @@ using Oceananigans.Grids: with_halo
 using Oceananigans.Architectures: architecture, CPU
 
 using Oceananigans.Operators
+using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, Γᶠᶠᶜ, ∂t_σ
 
 import Base: show, summary
 import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z

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,

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_σ
+# ---------- = - --------------------- - -------------
+#    Δ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_σ
+#
+# which represents the static volume times the time derivative of the vertical grid scaling.
+@inline V_times_∂t_σ(i, j, k, grid) = Azᶜᶜᶜ(i, j, k, grid) * Δrᶜᶜᶜ(i, j, k, grid) * ∂t_σ(i, j, k, grid)
+
 @inline function upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme::VectorInvariantCrossVerticalUpwinding, u, v)
     @inbounds û = u[i, j, k]
     δ_stencil = scheme.upwinding.divergence_stencil
 
-    δᴿ = _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) 
+    δᴿ  =    _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) 
+    ∂ts = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, cross_scheme, V_times_∂t_σ)
 
-    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_σ)
 
-    return v̂ * δᴿ
+    return v̂ * (δᴿ + ∂ts)
 end

src/Advection/vector_invariant_self_upwinding.jl

@@ -2,23 +2,26 @@
 ##### 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_σ(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_σ(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_σ(i, j, k, grid)
 
 @inline function upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme::VectorInvariantSelfVerticalUpwinding, u, v)
 
     δU_stencil   = scheme.upwinding.δU_stencil    
     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ᴿ)

src/Grids/Grids.jl

@@ -18,7 +18,9 @@ export xnodes, ynodes, znodes, λnodes, φnodes, rnodes
 export spacings
 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
 

src/Grids/grid_generation.jl

@@ -150,3 +150,48 @@ function generate_coordinate(FT, ::Flat, N, H, ::Nothing, coordinate_name, arch)
         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ᵃᵃᶠ
+
+    Lr, rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ = generate_coordinate(FT, topo[3](), Nz, Hz, r_faces, :r, arch)
+
+    args = (topo, (Nx, Ny, Nz), (Hx, Hy, Hz))
+
+    σᶜᶜ⁻ = new_data(FT, arch, (Center, Center, Nothing), args...)
+    σᶜᶜⁿ = new_data(FT, arch, (Center, Center, Nothing), args...)
+    σᶠᶜⁿ = new_data(FT, arch, (Face,   Center, Nothing), args...)
+    σᶜᶠⁿ = new_data(FT, arch, (Center, Face,   Nothing), args...)
+    σᶠᶠⁿ = new_data(FT, arch, (Face,   Face,   Nothing), args...)
+    ηⁿ    = new_data(FT, arch, (Center, Center, Nothing), args...)
+    ∂t_σ = new_data(FT, arch, (Center, Center, Nothing), args...)
+
+    # Fill all the scalings with one (at rest coordinate)
+    for σ in (σᶜᶜ⁻, σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ)
+        fill!(σ, 1)
+    end
+
+    return Lr, ZStarVerticalCoordinate(rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ, ηⁿ, σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ∂t_σ)
+end

src/Grids/grid_utils.jl

@@ -317,14 +317,20 @@ coordinate_summary(topo, Δ::Union{AbstractVector, AbstractMatrix}, name) =
              name, prettysummary(maximum(parent(Δ))))
 
 #####
-##### Static column depths
+##### Static and Dynamic column depths
 #####
 
 @inline static_column_depthᶜᶜᵃ(i, j, grid) = grid.Lz
 @inline static_column_depthᶜᶠᵃ(i, j, grid) = grid.Lz
 @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
 #####

src/Grids/vertical_coordinate.jl

@@ -23,12 +23,50 @@ struct StaticVerticalCoordinate{C, D} <: AbstractVerticalCoordinate
     Δᵃᵃᶜ :: 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.
+# - `σᶜᶜⁿ::CC`: Vertical grid scaling at center-center at the current time step.
+# - `σᶠᶜⁿ::FC`: Vertical grid scaling at face-center at the current time step.
+# - `σᶜᶠⁿ::CF`: Vertical grid scaling at center-face at the current time step.
+# - `σᶠᶠⁿ::FF`: Vertical grid scaling at face-face at the current time step.
+# - `σᶜᶜ⁻::CC`: Vertical grid scaling at center-center at the previous time step.
+# - `∂t_σ::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
+   σᶜᶜⁿ :: CC
+   σᶠᶜⁿ :: FC
+   σᶜᶠⁿ :: CF
+   σᶠᶠⁿ :: FF
+   σᶜᶜ⁻ :: CC
+   ∂t_σ :: 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 RegularVerticalCoordinate = StaticVerticalCoordinate{<:Any, <:Number}
-const RegularVerticalGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate}
+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
@@ -46,6 +84,32 @@ on_architecture(arch, coord::StaticVerticalCoordinate) =
                             on_architecture(arch, coord.Δᵃᵃᶠ),
                             on_architecture(arch, 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.σᶜᶜⁿ),
+                            Adapt.adapt(to, coord.σᶠᶜⁿ),
+                            Adapt.adapt(to, coord.σᶜᶠⁿ),
+                            Adapt.adapt(to, coord.σᶠᶠⁿ),
+                            Adapt.adapt(to, coord.σᶜᶜ⁻),
+                            Adapt.adapt(to, coord.∂t_σ))
+
+on_architecture(arch, coord::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.σᶜᶜⁿ),
+                            on_architecture(arch, coord.σᶠᶜⁿ),
+                            on_architecture(arch, coord.σᶜᶠⁿ),
+                            on_architecture(arch, coord.σᶠᶠⁿ),
+                            on_architecture(arch, coord.σᶜᶜ⁻),
+                            on_architecture(arch, coord.∂t_σ))
+
 #####
 ##### Nodes and spacings (common to every grid)...
 #####
@@ -93,11 +157,20 @@ z_domain(grid) = domain(topology(grid, 3)(), grid.Nz, grid.z.cᵃᵃᶠ)
 end
 
 @inline cpu_face_constructor_z(grid) = cpu_face_constructor_r(grid)
+@inline cpu_face_constructor_z(grid::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) = 
+coordinate_summary(::Bounded, z::ZStarVerticalCoordinate, name) = 
     @sprintf("Free-surface following with Δ%s=%s", name, prettysummary(z.Δᵃᵃᶜ))

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

src/ImmersedBoundaries/immersed_grid_metrics.jl

@@ -1,7 +1,6 @@
 using Oceananigans.AbstractOperations: GridMetricOperation
 
-import Oceananigans.Grids: coordinates
-import Oceananigans.Operators: Δrᵃᵃᶠ, Δrᵃᵃᶜ
+import Oceananigans.Operators: Δrᵃᵃᶠ, Δrᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ
 import Oceananigans.Operators: Δxᶠᵃᵃ, Δxᶜᵃᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ
 import Oceananigans.Operators: Δyᵃᶠᵃ, Δyᵃᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ, Δyᶜᶜᵃ
 import Oceananigans.Operators: Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ
@@ -23,8 +22,8 @@ import Oceananigans.Operators: intrinsic_vector, extrinsic_vector
 
 @inline Δrᵃᵃᶠ(i, j, k, ibg::IBG) = Δrᵃᵃᶠ(i, j, k, ibg.underlying_grid)
 @inline Δrᵃᵃᶜ(i, j, k, ibg::IBG) = Δrᵃᵃᶜ(i, j, k, ibg.underlying_grid)
-@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.underlying_grid)
 @inline Δzᵃᵃᶜ(i, j, k, ibg::IBG) = Δzᵃᵃᶜ(i, j, k, ibg.underlying_grid)
+@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.underlying_grid)
 
 # 1D Horizontal spacings