Compute reference density with derivative

src/Thermodynamics/reference_states.jl

@@ -1,5 +1,5 @@
-using Oceananigans: Oceananigans, Center, Field, set!, fill_halo_regions!
-using Oceananigans.BoundaryConditions: FieldBoundaryConditions, ValueBoundaryCondition
+using Oceananigans: Oceananigans, Center, Field, set!, fill_halo_regions!, ∂z, znodes
+using Oceananigans.BoundaryConditions: FieldBoundaryConditions, GradientBoundaryCondition
 using Oceananigans.Operators: ℑzᵃᵃᶠ
 
 using Adapt: Adapt, adapt
@@ -9,12 +9,12 @@ using GPUArraysCore: @allowscalar
 ##### Reference state computations for Boussinesq and Anelastic models
 #####
 
-struct ReferenceState{FT, F}
+struct ReferenceState{FT, P, R}
     surface_pressure :: FT # base pressure: reference pressure at z=0
     potential_temperature :: FT  # constant reference potential temperature
     standard_pressure :: FT # pˢᵗ: reference pressure for potential temperature (default 1e5)
-    pressure :: F
-    density :: F
+    pressure :: P
+    density :: R
 end
 
 Adapt.adapt_structure(to, ref::ReferenceState) =
@@ -134,18 +134,36 @@ function ReferenceState(grid, constants=ThermodynamicConstants(eltype(grid));
     p₀ = convert(FT, surface_pressure)
     θ₀ = convert(FT, potential_temperature)
     pˢᵗ = convert(FT, standard_pressure)
+    g = constants.gravitational_acceleration
     loc = (nothing, nothing, Center())
 
-    ρ₀ = surface_density(p₀, θ₀, pˢᵗ, constants)
-    ρ_bcs = FieldBoundaryConditions(grid, loc, bottom=ValueBoundaryCondition(ρ₀))
-    ρᵣ = Field{Nothing, Nothing, Center}(grid, boundary_conditions=ρ_bcs)
-    set!(ρᵣ, z -> adiabatic_hydrostatic_density(z, p₀, θ₀, pˢᵗ, constants))
-    fill_halo_regions!(ρᵣ)
-
-    p_bcs = FieldBoundaryConditions(grid, loc, bottom=ValueBoundaryCondition(p₀))
+    # Use GradientBoundaryCondition at both top and bottom boundaries to ensure
+    # correct discrete hydrostatic balance: ρ = -∂z(p)/g. The gradient is set
+    # using the analytical hydrostatic density at each boundary cell center.
+    z = znodes(grid, Center())
+    z_bottom = first(z)
+    z_top = last(z)
+
+    ρ_bottom = adiabatic_hydrostatic_density(z_bottom, p₀, θ₀, pˢᵗ, constants)
+    ρ_top = adiabatic_hydrostatic_density(z_top, p₀, θ₀, pˢᵗ, constants)
+    ∂p∂z_bottom = -ρ_bottom * g
+    ∂p∂z_top = -ρ_top * g
+
+    p_bcs = FieldBoundaryConditions(grid, loc,
+        bottom = GradientBoundaryCondition(∂p∂z_bottom),
+        top = GradientBoundaryCondition(∂p∂z_top))
     pᵣ = Field{Nothing, Nothing, Center}(grid, boundary_conditions=p_bcs)
     set!(pᵣ, z -> adiabatic_hydrostatic_pressure(z, p₀, θ₀, constants))
     fill_halo_regions!(pᵣ)
 
+    # Compute density from discrete pressure gradient for discrete hydrostatic balance.
+    # Use gradient BC based on analytical density at each boundary.
+    ρ_bcs = FieldBoundaryConditions(grid, loc,
+        bottom = GradientBoundaryCondition(zero(FT)),
+        top = GradientBoundaryCondition(zero(FT)))
+    ρᵣ = Field{Nothing, Nothing, Center}(grid, boundary_conditions=ρ_bcs)
+    set!(ρᵣ, - ∂z(pᵣ) / g)
+    fill_halo_regions!(ρᵣ)
+
     return ReferenceState(p₀, θ₀, pˢᵗ, pᵣ, ρᵣ)
 end

test/atmosphere_model_construction.jl

@@ -48,12 +48,13 @@ end
         @testset let p₀ = p₀, θ₀ = θ₀, formulation = formulation
             reference_state = ReferenceState(grid, constants, surface_pressure=p₀, potential_temperature=θ₀)
 
-            # Check that interpolating to the first face (k=1) recovers surface values
-            # Note: surface_density correctly converts potential temperature to temperature using the Exner function
+            # Check that interpolating to the first face (k=1) approximately recovers surface values.
+            # Uses rtol=1e-4 because GradientBoundaryCondition is used for discrete hydrostatic balance,
+            # which introduces small interpolation errors at the surface.
             ρ₀ = surface_density(reference_state)
             for i = 1:Nx, j = 1:Ny
-                @test p₀ ≈ @allowscalar ℑzᵃᵃᶠ(i, j, 1, grid, reference_state.pressure)
-                @test ρ₀ ≈ @allowscalar ℑzᵃᵃᶠ(i, j, 1, grid, reference_state.density)
+                @test isapprox(p₀, @allowscalar(ℑzᵃᵃᶠ(i, j, 1, grid, reference_state.pressure)), rtol=FT(1e-4))
+                @test isapprox(ρ₀, @allowscalar(ℑzᵃᵃᶠ(i, j, 1, grid, reference_state.density)), rtol=FT(1e-4))
             end
 
             dynamics = AnelasticDynamics(reference_state)

test/diagnostics.jl

@@ -7,7 +7,7 @@ using GPUArraysCore: @allowscalar
 
 @testset "Potential temperature diagnostics [$(FT)]" for FT in test_float_types()
     Oceananigans.defaults.FloatType = FT
-    grid = RectilinearGrid(default_arch; size=(2, 2, 8), extent=(100, 100, 1000))
+    grid = RectilinearGrid(default_arch; size=(2, 2, 8), x=(0, 100), y=(0, 100), z=(0, 1000))
     model = AtmosphereModel(grid)
     set!(model, θ=300, qᵗ=0.01)
 
@@ -80,7 +80,7 @@ end
 
 @testset "Static energy diagnostics [$(FT)]" for FT in test_float_types()
     Oceananigans.defaults.FloatType = FT
-    grid = RectilinearGrid(default_arch; size=(2, 2, 8), extent=(100, 100, 1000))
+    grid = RectilinearGrid(default_arch; size=(2, 2, 8), x=(0, 100), y=(0, 100), z=(0, 1000))
     model = AtmosphereModel(grid)
     set!(model, θ=300, qᵗ=0.01)
 
@@ -100,7 +100,7 @@ end
 
 @testset "Relative humidity diagnostics [$(FT)]" for FT in test_float_types()
     Oceananigans.defaults.FloatType = FT
-    grid = RectilinearGrid(default_arch; size=(2, 2, 8), extent=(100, 100, 1000))
+    grid = RectilinearGrid(default_arch; size=(2, 2, 8), x=(0, 100), y=(0, 100), z=(0, 1000))
     microphysics = SaturationAdjustment()
     model = AtmosphereModel(grid; microphysics)
 

test/dynamics.jl

@@ -75,7 +75,9 @@ tol = 1e-6
     end
 end
 
-@testset "Vertical momentum conservation for neutral initial condition [$(FT)]" for FT in (Float32, Float64)
+# Float32 is excluded because the integrated momentum over the large domain (4e12 m³)
+# accumulates roundoff errors to O(100), making the test meaningless for Float32.
+@testset "Vertical momentum conservation for neutral initial condition [$(FT)]" for FT in (Float64,)
     Oceananigans.defaults.FloatType = FT
     grid = RectilinearGrid(default_arch;
                            size = (16, 5, 16),