Merge branch 'main' into krylov-integration

Author: amontoison

.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
 #####

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.95.6"
+version = "0.95.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -65,8 +65,8 @@ Krylov = "0.9.9"
 LinearAlgebra = "1.9"
 Logging = "1.9"
 MPI = "0.16, 0.17, 0.18, 0.19, 0.20"
-Makie = "0.21"
-MakieCore = "0.7, 0.8"
+Makie = "0.21, 0.22"
+MakieCore = "0.7, 0.8, 0.9"
 NCDatasets = "0.12.10, 0.13.1, 0.14"
 OffsetArrays = "1.4"
 OrderedCollections = "1.1"

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,38 @@
 ##### Cross and Self Upwinding of the Divergence flux
 #####
 
+# If the grid is moving, the discrete continuity equation is calculated as:
+#
+#  ωᵏ⁺¹ - ωᵏ      δx(Ax u) + δy(Ay v)     Δrᶜᶜᶜ ∂t_σ
+# ---------- = - --------------------- - -------------
+#    Δzᶜᶜᶜ                Vᶜᶜᶜ              Δzᶜᶜᶜ
+#
+# Where ω is the vertical velocity with respect to a moving grid.
+# We upwind the discrete divergence `δx(Ax u) + δy(Ay v)` and then divide by the volume, 
+# therefore, the correct term to be added to the divergence transport due to the moving grid is:
+#
+#  Azᶜᶜᶜ Δrᶜᶜᶜ ∂t_σ
+#
+# which represents the static volume times the time derivative of the vertical grid scaling.
+# If the grid is stationary, ∂t_σ evaluates to zero, so this term disappears from the divergence flux.
+@inline Az_Δr_∂t_σ(i, j, k, grid) = Azᶜᶜᶜ(i, j, k, grid) * Δrᶜᶜᶜ(i, j, k, grid) * ∂t_σ(i, j, k, grid)
+
 @inline function upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme::VectorInvariantCrossVerticalUpwinding, u, v)
     @inbounds û = u[i, j, k]
     δ_stencil = scheme.upwinding.divergence_stencil
 
-    δᴿ = _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) 
+    δᴿ   =    _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) 
+    ∂t_σ = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, cross_scheme, Az_Δr_∂t_σ)
 
-    return û * δᴿ
+    return û * (δᴿ + ∂t_σ) # For static grids, ∂t_σ == 0
 end
 
 @inline function upwinded_divergence_flux_Vᶜᶠᶜ(i, j, k, grid, scheme::VectorInvariantCrossVerticalUpwinding, u, v)
     @inbounds v̂ = v[i, j, k]
     δ_stencil = scheme.upwinding.divergence_stencil
 
-    δᴿ = _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(v̂), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) 
+    δᴿ   =    _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(v̂), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) 
+    ∂t_σ = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, Az_Δr_∂t_σ)
 
-    return v̂ * δᴿ
+    return v̂ * (δᴿ + ∂t_σ) # For static grids, ∂t_σ == 0
 end

src/Advection/vector_invariant_self_upwinding.jl

@@ -2,23 +2,28 @@
 ##### Self Upwinding of Divergence Flux, the best option!
 #####
 
-@inline δx_U(i, j, k, grid, u, v) =  δxᶜᵃᵃ(i, j, k, grid, Ax_qᶠᶜᶜ, u)
-@inline δy_V(i, j, k, grid, u, v) =  δyᵃᶜᵃ(i, j, k, grid, Ay_qᶜᶠᶜ, v)
+@inline δx_U(i, j, k, grid, u, v) = δxᶜᶜᶜ(i, j, k, grid, Ax_qᶠᶜᶜ, u)
+@inline δy_V(i, j, k, grid, u, v) = δyᶜᶜᶜ(i, j, k, grid, Ay_qᶜᶠᶜ, v)
+
+# For moving grids, we include the time-derivative of the grid scaling in the divergence flux. 
+# If the grid is stationary, `Az_Δr_∂t_σ` evaluates to zero.
+@inline δx_U_plus_∂t_σ(i, j, k, grid, u, v) = δxᶜᶜᶜ(i, j, k, grid, Ax_qᶠᶜᶜ, u) + Az_Δr_∂t_σ(i, j, k, grid)
+@inline δy_V_plus_∂t_σ(i, j, k, grid, u, v) = δyᶜᶜᶜ(i, j, k, grid, Ay_qᶜᶠᶜ, v) + Az_Δr_∂t_σ(i, j, k, grid)
 
 # Velocity smoothness for divergence upwinding
 @inline U_smoothness(i, j, k, grid, u, v) = ℑxᶜᵃᵃ(i, j, k, grid, Ax_qᶠᶜᶜ, u)
 @inline V_smoothness(i, j, k, grid, u, v) = ℑyᵃᶜᵃ(i, j, k, grid, Ay_qᶜᶠᶜ, v)
 
 # Divergence smoothness for divergence upwinding
-@inline divergence_smoothness(i, j, k, grid, u, v) = δx_U(i, j, k, grid, u, v) + δy_V(i, j, k, grid, u, v)
+@inline divergence_smoothness(i, j, k, grid, u, v) = δx_U(i, j, k, grid, u, v) + δy_V(i, j, k, grid, u, v) 
 
 @inline function upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme::VectorInvariantSelfVerticalUpwinding, u, v)
 
     δ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_∂t_σ, u, v) 
     δuᴿ =    _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), δx_U, δU_stencil, u, v) 
 
     return û * (δvˢ + δuᴿ)
@@ -30,7 +35,7 @@ end
     cross_scheme = scheme.upwinding.cross_scheme
 
     @inbounds v̂ = v[i, j, k]
-    δuˢ = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, δx_U, u, v)
+    δuˢ = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, δx_U_plus_∂t_σ, u, v)
     δvᴿ =    _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(v̂), δy_V, δV_stencil, u, v) 
 
     return v̂ * (δuˢ + δvᴿ)

src/Grids/Grids.jl

@@ -11,6 +11,7 @@ export XRegularRG, YRegularRG, ZRegularRG, XYRegularRG, XYZRegularRG
 export LatitudeLongitudeGrid, XRegularLLG, YRegularLLG, ZRegularLLG
 export OrthogonalSphericalShellGrid, ConformalCubedSphereGrid, ZRegOrthogonalSphericalShellGrid
 export conformal_cubed_sphere_panel
+export MutableVerticalDiscretization
 export node, nodes
 export ξnode, ηnode, rnode
 export xnode, ynode, znode, λnode, φnode
@@ -19,6 +20,7 @@ export spacings
 export xspacings, yspacings, zspacings, λspacings, φspacings, rspacings
 export minimum_xspacing, minimum_yspacing, minimum_zspacing
 export static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ
+export column_depthᶜᶜᵃ, column_depthᶠᶜᵃ, column_depthᶜᶠᵃ, column_depthᶠᶠᵃ
 export offset_data, new_data
 export on_architecture
 
@@ -118,7 +120,7 @@ struct ZDirection <: AbstractDirection end
 struct NegativeZDirection <: AbstractDirection end
 
 include("abstract_grid.jl")
-include("vertical_coordinates.jl")
+include("vertical_discretization.jl")
 include("grid_utils.jl")
 include("nodes_and_spacings.jl")
 include("zeros_and_ones.jl")

src/Grids/grid_generation.jl

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

src/Grids/grid_utils.jl

@@ -299,7 +299,7 @@ end
 function dimension_summary(topo, name, dom, z::AbstractVerticalCoordinate, pad_domain=0)
     prefix = domain_summary(topo, name, dom)
     padding = " "^(pad_domain+1) 
-    return string(prefix, padding, coordinate_summary(topo, z.Δᵃᵃᶜ, name))
+    return string(prefix, padding, coordinate_summary(topo, z, name))
 end
 
 function dimension_summary(topo, name, dom, spacing, pad_domain=0)
@@ -317,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 a ``mutable'' grid type
+@inline column_depthᶜᶜᵃ(i, j, k, grid, η) = static_column_depthᶜᶜᵃ(i, j, grid) 
+@inline column_depthᶠᶜᵃ(i, j, k, grid, η) = static_column_depthᶠᶜᵃ(i, j, grid) 
+@inline column_depthᶜᶠᵃ(i, j, k, grid, η) = static_column_depthᶜᶠᵃ(i, j, grid) 
+@inline column_depthᶠᶠᵃ(i, j, k, grid, η) = static_column_depthᶠᶠᵃ(i, j, grid) 
+
 #####
 ##### Spherical geometry
 #####