[dycore] Add simple dynamics test for momentum conservation

examples/acoustic_wave.jl

@@ -49,8 +49,9 @@ constants = model.thermodynamic_constants
 
 θ₀ = 300      # Reference potential temperature (K)
 p₀ = 101325   # Surface pressure (Pa)
+pˢᵗ = 1e5     # Standard pressure (Pa)
 
-reference = ReferenceState(grid, constants; surface_pressure=p₀, potential_temperature=θ₀)
+reference = ReferenceState(grid, constants; surface_pressure=p₀, potential_temperature=θ₀, standard_pressure=pˢᵗ)
 
 # The sound speed at the surface determines the acoustic wave propagation speed.
 
@@ -78,7 +79,7 @@ Uᵢ(z) = U₀ * log((z + ℓ) / ℓ)
 gaussian(x, z) = exp(-(x^2 + z^2) / 2σ^2)
 ρ₀ = interior(reference.density, 1, 1, 1)[]
 
-ρᵢ(x, z) = adiabatic_hydrostatic_density(z, p₀, θ₀, constants) + δρ * gaussian(x, z)
+ρᵢ(x, z) = adiabatic_hydrostatic_density(z, p₀, θ₀, pˢᵗ, constants) + δρ * gaussian(x, z)
 uᵢ(x, z) = Uᵢ(z) #+ (𝕌ˢⁱ / ρ₀) * δρ * gaussian(x, z)
 
 set!(model, ρ=ρᵢ, θ=θ₀, u=uᵢ)
@@ -116,7 +117,7 @@ u, v, w = model.velocities
 ρᵇᵍ = CenterField(grid)
 uᵇᵍ = XFaceField(grid)
 
-set!(ρᵇᵍ, (x, z) -> adiabatic_hydrostatic_density(z, p₀, θ₀, constants))
+set!(ρᵇᵍ, (x, z) -> adiabatic_hydrostatic_density(z, p₀, θ₀, pˢᵗ, constants))
 set!(uᵇᵍ, (x, z) -> Uᵢ(z))
 
 ρ′ = Field(ρ - ρᵇᵍ)

src/Thermodynamics/reference_states.jl

@@ -51,14 +51,31 @@ function surface_density(ref::ReferenceState)
 end
 
 """
-    surface_density(p₀, θ₀, constants)
+    surface_density(p₀, T₀, constants)
 
-Compute the surface air density from surface pressure `p₀`, surface temperature `θ₀`,
+Compute the surface air density from surface pressure `p₀`, surface temperature `T₀`,
 and thermodynamic `constants` using the ideal gas law for dry air.
 """
-@inline function surface_density(p₀, θ₀, constants)
+@inline function surface_density(p₀, T₀, constants)
     Rᵈ = dry_air_gas_constant(constants)
-    return p₀ / (Rᵈ * θ₀)
+    return p₀ / (Rᵈ * T₀)
+end
+
+"""
+    surface_density(p₀, θ₀, pˢᵗ, constants)
+
+Compute the surface air density from surface pressure `p₀`, potential temperature `θ₀`,
+standard pressure `pˢᵗ`, and thermodynamic `constants` using the ideal gas law for dry air.
+
+The temperature is computed from potential temperature using the Exner function:
+`T₀ = Π₀ * θ₀` where `Π₀ = (p₀ / pˢᵗ)^(Rᵈ/cᵖᵈ)`.
+"""
+@inline function surface_density(p₀, θ₀, pˢᵗ, constants)
+    Rᵈ = dry_air_gas_constant(constants)
+    cᵖᵈ = constants.dry_air.heat_capacity
+    Π₀ = (p₀ / pˢᵗ)^(Rᵈ / cᵖᵈ)
+    T₀ = Π₀ * θ₀
+    return p₀ / (Rᵈ * T₀)
 end
 
 """
@@ -78,15 +95,15 @@ end
 """
 $(TYPEDSIGNATURES)
 
-Compute the reference density at height `z` that associated with the reference pressure and
-potential temperature `θ₀`. The reference density is defined as the density of dry air at the
-reference pressure and temperature.
+Compute the reference density at height `z` that associated with the reference pressure `p₀`,
+potential temperature `θ₀`, and standard pressure `pˢᵗ`. The reference density is defined as
+the density of dry air at the reference pressure and temperature.
 """
-@inline function adiabatic_hydrostatic_density(z, p₀, θ₀, constants)
+@inline function adiabatic_hydrostatic_density(z, p₀, θ₀, pˢᵗ, constants)
     Rᵈ = dry_air_gas_constant(constants)
     cᵖᵈ = constants.dry_air.heat_capacity
     pᵣ = adiabatic_hydrostatic_pressure(z, p₀, θ₀, constants)
-    ρ₀ = surface_density(p₀, θ₀, constants)
+    ρ₀ = surface_density(p₀, θ₀, pˢᵗ, constants)
     return ρ₀ * (pᵣ / p₀)^(1 - Rᵈ / cᵖᵈ)
 end
 
@@ -119,10 +136,10 @@ function ReferenceState(grid, constants=ThermodynamicConstants(eltype(grid));
     pˢᵗ = convert(FT, standard_pressure)
     loc = (nothing, nothing, Center())
 
-    ρ₀ = surface_density(p₀, θ₀, constants)
+    ρ₀ = 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₀, θ₀, constants))
+    set!(ρᵣ, z -> adiabatic_hydrostatic_density(z, p₀, θ₀, pˢᵗ, constants))
     fill_halo_regions!(ρᵣ)
 
     p_bcs = FieldBoundaryConditions(grid, loc, bottom=ValueBoundaryCondition(p₀))

test/atmosphere_model_construction.jl

@@ -49,8 +49,8 @@ end
             reference_state = ReferenceState(grid, constants, surface_pressure=p₀, potential_temperature=θ₀)
 
             # Check that interpolating to the first face (k=1) recovers surface values
-            q₀ = Breeze.Thermodynamics.MoistureMassFractions{FT} |> zero
-            ρ₀ = Breeze.Thermodynamics.density(θ₀, p₀, q₀, constants)
+            # Note: surface_density correctly converts potential temperature to temperature using the Exner function
+            ρ₀ = 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/dynamics.jl

@@ -0,0 +1,114 @@
+using Breeze
+using Breeze.Microphysics
+using Oceananigans
+using Test
+using GPUArraysCore: @allowscalar
+
+"""
+    thermal_bubble_model(grid; Δθ=10, uᵢ=0, vᵢ=0, wᵢ=0)
+
+Set up a thermal bubble initial condition for an `AtmosphereModel` on the provided grid.
+
+The thermal bubble is a spherical potential temperature perturbation centered in the domain.
+The background state has a stable stratification with Brunt-Väisälä frequency squared N² = 1e-6.
+
+# Arguments
+- `grid`: An `Oceananigans.AbstractGrid` or compatible grid for the model domain.
+- `uᵢ`: Optional initial x-velocity (default: 0).
+- `vᵢ`: Optional initial y-velocity (default: 0).
+- `wᵢ`: Optional initial z-velocity (default: 0).
+```
+"""
+function thermal_bubble_model(grid; Δθ=10, N²=1e-6, uᵢ=0, vᵢ=0, wᵢ=0, qᵗ=0, microphysics=nothing)
+    model = AtmosphereModel(grid; advection=WENO(), microphysics)
+    r₀ = 2e3
+    θ₀ = model.dynamics.reference_state.potential_temperature
+    g = model.thermodynamic_constants.gravitational_acceleration
+
+    function θᵢ(x, y, z)
+        θ̄ = θ₀ * exp(N² * z / g)
+        r = sqrt(x^2 + y^2 + z^2)
+        θ′ = Δθ * max(0, 1 - r / r₀)
+        return θ̄ + θ′
+    end
+
+    set!(model, θ=θᵢ, u=uᵢ, v=vᵢ, w=wᵢ, qᵗ=qᵗ)
+
+    return model
+end
+
+Δt = 1e-3
+tol = 1e-6
+
+@testset "Horizontal momentum conservation with spherical thermal bubble [$(FT)]" for FT in (Float32, Float64)
+    Oceananigans.defaults.FloatType = FT
+    grid = RectilinearGrid(default_arch;
+                           size = (16, 16, 16),
+                           x = (-10e3, 10e3),
+                           y = (-10e3, 10e3),
+                           z = (-3e3, 7e3),
+                           topology = (Periodic, Periodic, Bounded),
+                           halo = (5, 5, 5))
+
+    # Set spherical thermal bubble initial condition with sheared horizontal velocities
+    uᵢ(x, y, z) = 5 # * (z + 3e3) / 10e3
+    vᵢ(x, y, z) = 3 # * (z + 3e3) / 10e3
+    model = thermal_bubble_model(grid; uᵢ, vᵢ)
+
+    # Compute initial total u-momentum
+    ∫ρu = Field(Integral(model.momentum.ρu))
+    ∫ρv = Field(Integral(model.momentum.ρv))
+    Px₀ = @allowscalar first(∫ρu)
+    Py₀ = @allowscalar first(∫ρv)
+
+    # Time step the model
+    Nt = 10
+
+    for step in 1:Nt
+        time_step!(model, Δt)
+        compute!(∫ρu)
+        compute!(∫ρv)
+        Px = @allowscalar first(∫ρu)
+        Py = @allowscalar first(∫ρv)
+        @test Px ≈ Px₀
+        @test Py ≈ Py₀
+    end
+end
+
+@testset "Vertical momentum conservation for neutral initial condition [$(FT)]" for FT in (Float32, Float64)
+    Oceananigans.defaults.FloatType = FT
+    grid = RectilinearGrid(default_arch;
+                           size = (16, 5, 16),
+                           x = (-10e3, 10e3),
+                           y = (-10e3, 10e3),
+                           z = (-5e3, 5e3),
+                           topology = (Bounded, Periodic, Bounded),
+                           halo = (5, 5, 5))
+
+    # Set spherical thermal bubble initial condition with u velocity only
+    wᵢ(x, y, z) = 0 # * x / 20e3 * exp(-z^2 / (2 * 1e3^2))
+    model = thermal_bubble_model(grid; wᵢ, Δθ=0, N²=0)
+
+    # Compute initial total u-momentum
+    ∫ρu = Field(Integral(model.momentum.ρu))
+    ∫ρw = Field(Integral(model.momentum.ρw))
+    Px₀ = @allowscalar first(∫ρu)
+    Pz₀ = @allowscalar first(∫ρw)
+
+    # Time step the model
+    Nt = 10
+
+    # Use absolute tolerance for comparisons with zero.
+    # Float32 has larger roundoff errors due to limited precision.
+    atol = FT == Float32 ? FT(1e-2) : FT(1e-6)
+
+    for step in 1:Nt
+        time_step!(model, Δt)
+        compute!(∫ρu)
+        compute!(∫ρw)
+        Px = @allowscalar first(∫ρu)
+        Pz = @allowscalar first(∫ρw)
+        @test Px ≈ Px₀
+        @test isapprox(Pz, Pz₀, atol=atol)
+    end
+end