diff --git a/Project.toml b/Project.toml index 2e152e82..1f2e6f0a 100644 --- a/Project.toml +++ b/Project.toml @@ -10,6 +10,7 @@ projects = ["docs", "test"] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -31,6 +32,7 @@ ClimaComms = "0.6" CloudMicrophysics = "0.29.0" Dates = "<0.0.1, 1" DocStringExtensions = "0.9.5" +GPUArraysCore = "0.2" InteractiveUtils = "<0.0.1, 1" JLD2 = "0.5.13, 0.6" KernelAbstractions = "0.9" diff --git a/docs/make.jl b/docs/make.jl index 4d3d6b28..bee842b8 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -38,13 +38,12 @@ examples = [ Example("Single column radiation", "single_column_radiation", true), Example("Stationary parcel model", "stationary_parcel_model", true), Example("Acoustic wave in shear layer", "acoustic_wave", true), + Example("Cloud formation in prescribed updraft", "kinematic_driver", true), ] # Filter out long-running example if necessary filter!(x -> x.build_always || get(ENV, "BREEZE_BUILD_ALL_EXAMPLES", "false") == "true", examples) - example_pages = [ex.title => joinpath("literated", ex.basename * ".md") for ex in examples] - semaphore = Base.Semaphore(Threads.nthreads(:interactive)) @time "literate" @sync for example in examples script_file = example.basename * ".jl" diff --git a/examples/kinematic_driver.jl b/examples/kinematic_driver.jl new file mode 100644 index 00000000..3ee79105 --- /dev/null +++ b/examples/kinematic_driver.jl @@ -0,0 +1,223 @@ +# # Kinematic driver: cloud formation in an idealized updraft +# +# In atmospheric modeling, we sometimes want to isolate microphysics and thermodynamics +# from dynamics. **Kinematic models** prescribe the velocity field rather than solving +# momentum equations, letting us focus purely on tracer transport and phase changes. +# +# This example demonstrates Breeze's [`PrescribedDynamics`](@ref) formulation by simulating +# cloud formation in an idealized updraft. A uniform vertical velocity lifts moist air +# through a realistic temperature profile. As air rises and cools, water vapor +# condenses to form clouds — a fundamental process driving all precipitation on Earth. +# +# ## Physical setup +# +# We simulate a 1D column representing a rising air parcel with: +# - A realistic potential temperature profile (troposphere + stratosphere) +# - Uniform upward velocity `W₀ = 2 m/s` (a gentle cumulus updraft) +# - Moist boundary layer air entering from below +# +# The **divergence correction** option compensates for the non-zero mass flux divergence +# ∇·(ρU) that arises when velocity doesn't vary with the reference density profile. +# This is essential for physically consistent tracer advection in kinematic models. + +using Breeze +using CairoMakie +using Printf + +# ## Grid and reference state +# +# We construct a 20 km tall column extending from the surface through the tropopause +# into the lower stratosphere. The reference state establishes the background +# pressure and density profile based on a hydrostatically-balanced atmosphere. + +Nz = 100 +Lz = 20000 # 20 km domain height +grid = RectilinearGrid(CPU(); size=Nz, x=0, y=0, z=(0, Lz), + topology=(Flat, Flat, Bounded)) + +constants = ThermodynamicConstants() +θ₀ = 300 # Surface potential temperature (K) +p₀ = 1e5 # Surface pressure (Pa) +reference_state = ReferenceState(grid, constants; + surface_pressure=p₀, + potential_temperature=θ₀) + +# ## Prescribing dynamics with divergence correction +# +# The key feature of kinematic models is [`PrescribedDynamics`](@ref), which fixes +# the density and pressure fields from a reference state. We enable +# `divergence_correction=true` because our constant vertical velocity doesn't +# satisfy the anelastic continuity constraint ∇·(ρU) = 0. +# +# Without this correction, the tracer equation would see spurious sources/sinks +# from the non-zero velocity divergence. The correction adds a term `c ∇·(ρU)` +# that compensates for the prescribed velocity field's divergence. + +W₀ = 2 # Vertical velocity (m/s) — a gentle updraft +dynamics = PrescribedDynamics(reference_state; divergence_correction=true) + +# ## Boundary conditions +# +# The key boundary condition is at the surface: we prescribe incoming moist air +# with constant potential temperature and specific humidity. This represents +# the boundary layer air being lifted into the updraft. + +ρ₀ = surface_density(reference_state) + +# Surface boundary conditions for tracers +qᵗ₀ = 0.018 # Incoming specific humidity (18 g/kg) — typical tropical boundary layer +ρθ_bcs = FieldBoundaryConditions(bottom=ValueBoundaryCondition(ρ₀ * θ₀)) +ρqᵗ_bcs = FieldBoundaryConditions(bottom=ValueBoundaryCondition(ρ₀ * qᵗ₀)) +w_bcs = FieldBoundaryConditions(bottom=OpenBoundaryCondition(W₀), top=OpenBoundaryCondition(W₀)) + +# ## Microphysics: warm-phase saturation adjustment +# +# We use [`SaturationAdjustment`](@ref) with [`WarmPhaseEquilibrium`](@ref), which +# instantaneously partitions total water between vapor and liquid based on +# saturation. When air becomes supersaturated, excess vapor condenses to cloud +# liquid, releasing latent heat. This captures the essence of cloud formation +# without explicit condensation timescales. + +microphysics = SaturationAdjustment(equilibrium=WarmPhaseEquilibrium()) + +# ## Building the atmosphere model +# +# We assemble all components into an [`AtmosphereModel`](@ref). The combination +# of `PrescribedDynamics` with microphysics creates a powerful tool for +# understanding cloud processes in isolation from dynamics. + +model = AtmosphereModel(grid; dynamics, microphysics, + advection = WENO(order=5), + boundary_conditions = (ρθ=ρθ_bcs, ρqᵗ=ρqᵗ_bcs, w=w_bcs), + thermodynamic_constants = constants) + +# ## Initial conditions +# +# We initialize with a realistic tropospheric potential temperature profile +# that increases with height (stable stratification). Above the tropopause +# at 12 km, we switch to a stratospheric profile. The initial moisture +# decreases with height, typical of a real atmosphere. + +zᵗʳ = 12000 # Tropopause height (m) +θᵗʳ = 343 # Potential temperature at tropopause (K) +Tᵗʳ = 213 # Temperature at tropopause (K) +g = constants.gravitational_acceleration +cᵖᵈ = constants.dry_air.heat_capacity + +function θ_initial(z) + θ_troposphere = θ₀ + (θᵗʳ - θ₀) * (z / zᵗʳ)^(5/4) + θ_stratosphere = θᵗʳ * exp(g / (cᵖᵈ * Tᵗʳ) * (z - zᵗʳ)) + return ifelse(z <= zᵗʳ, θ_troposphere, θ_stratosphere) +end + +# Moisture profile: high in the boundary layer, decreasing with height +function qᵗ_initial(z) + z_scale = 3000 # Scale height for moisture (m) + return qᵗ₀ * exp(-z / z_scale) +end + +set!(model; θ=θ_initial, qᵗ=qᵗ_initial, w=W₀) + +# ## Running the simulation +# +# We run for 60 minutes, enough time for air parcels to rise several kilometers +# and for a quasi-steady cloud layer to develop. + +simulation = Simulation(model; Δt=1, stop_time=60*60, verbose=false) + +θ = model.formulation.potential_temperature +qˡ = model.microphysical_fields.qˡ +qᵛ = model.microphysical_fields.qᵛ + +times = Float64[] +θ_data, qˡ_data, qᵛ_data = [], [], [] + +function record_profiles(sim) + push!(times, time(sim)) + push!(θ_data, Array(interior(θ, 1, 1, :))) + push!(qˡ_data, Array(interior(qˡ, 1, 1, :))) + push!(qᵛ_data, Array(interior(qᵛ, 1, 1, :))) +end + +add_callback!(simulation, record_profiles, TimeInterval(10*60)) + +function progress(sim) + qˡ_max = maximum(qˡ) + θ_surf = θ[1, 1, 1] + @info @sprintf("t = %s, θ_surface = %.1f K, max(qˡ) = %.2e kg/kg", + prettytime(sim), θ_surf, qˡ_max) +end + +add_callback!(simulation, progress, TimeInterval(10*60)) + +@info "Running kinematic updraft simulation with cloud microphysics..." +run!(simulation) + +# ## Visualization +# +# The results reveal the physics of adiabatic cloud formation. The left panel +# shows how potential temperature evolves — influenced by latent heat release +# where clouds form. The center panel shows cloud liquid mixing ratio, +# clearly revealing the cloud base and cloud layer. The right panel shows +# water vapor, which decreases sharply above the cloud base where condensation occurs. + +z_km = znodes(grid, Center()) ./ 1000 +fig = Figure(size=(1000, 450)) + +ax_θ = Axis(fig[1, 1]; xlabel="θ (K)", ylabel="z (km)", + title="Potential temperature") +ax_qˡ = Axis(fig[1, 2]; xlabel="qˡ (g/kg)", ylabel="z (km)", + title="Cloud liquid", yticklabelsvisible=false) +ax_qᵛ = Axis(fig[1, 3]; xlabel="qᵛ (g/kg)", ylabel="z (km)", + title="Water vapor", yticklabelsvisible=false) + +colors = cgrad(:viridis, length(times), categorical=true) + +for (n, t) in enumerate(times) + t_min = Int(t / 60) + lines!(ax_θ, θ_data[n], z_km; color=colors[n], linewidth=2, label="t = $t_min min") + lines!(ax_qˡ, qˡ_data[n] .* 1000, z_km; color=colors[n], linewidth=2) + lines!(ax_qᵛ, qᵛ_data[n] .* 1000, z_km; color=colors[n], linewidth=2) +end + +# Add tropopause marker +for ax in [ax_θ, ax_qˡ, ax_qᵛ] + hlines!(ax, [zᵗʳ/1000]; color=:gray, linestyle=:dash, linewidth=1.5, label="Tropopause") +end + +Legend(fig[1, 4], ax_θ; framevisible=false) + +Label(fig[0, :], "Kinematic updraft (W₀ = $W₀ m/s) with warm-phase saturation adjustment"; + fontsize=18, tellwidth=false) + +save("kinematic_driver.png", fig) +fig + +# ![](kinematic_driver.png) + +# ## Discussion +# +# The kinematic driver framework enables focused study of cloud microphysics by +# decoupling them from dynamical feedbacks. Key observations from this simulation: +# +# 1. **Cloud base formation**: Moist boundary layer air rises and cools adiabatically. +# When it reaches its Lifting Condensation Level (LCL), condensation begins +# and cloud liquid appears. The sharp transition in qˡ marks the cloud base. +# +# 2. **Moisture partitioning**: Above the cloud base, total water is partitioned +# between vapor (at saturation) and liquid (the excess). Water vapor decreases +# with height because saturation vapor pressure decreases with temperature. +# +# 3. **Potential temperature**: Initially, θ increases with height. As the simulation +# progresses, latent heat release from condensation modifies the temperature +# profile within the cloud layer. +# +# 4. **Divergence correction**: Without `divergence_correction=true`, the constant +# velocity field would create spurious tracer sources because ∇·(ρW) ≠ 0. +# The correction adds a compensating term to the tracer equations. +# +# This setup is analogous to classic parcel theory experiments in cloud physics, +# but resolved on a grid. It's particularly useful for: +# - Testing and validating microphysics schemes in isolation +# - Understanding sensitivities to initial moisture and temperature +# - Pedagogical demonstrations of cloud formation physics diff --git a/src/AnelasticEquations/anelastic_dynamics.jl b/src/AnelasticEquations/anelastic_dynamics.jl index fcc673d9..54c57acd 100644 --- a/src/AnelasticEquations/anelastic_dynamics.jl +++ b/src/AnelasticEquations/anelastic_dynamics.jl @@ -38,7 +38,7 @@ $(TYPEDSIGNATURES) Materialize a stub `AnelasticDynamics` into a full dynamics object with the pressure anomaly field. """ -function AtmosphereModels.materialize_dynamics(dynamics::AnelasticDynamics, grid, boundary_conditions) +function AtmosphereModels.materialize_dynamics(dynamics::AnelasticDynamics, grid, boundary_conditions, thermodynamic_constants) pressure_anomaly = CenterField(grid) return AnelasticDynamics(dynamics.reference_state, pressure_anomaly) end diff --git a/src/AtmosphereModels/atmosphere_model.jl b/src/AtmosphereModels/atmosphere_model.jl index f2edf092..224df2f0 100644 --- a/src/AtmosphereModels/atmosphere_model.jl +++ b/src/AtmosphereModels/atmosphere_model.jl @@ -107,6 +107,7 @@ function AtmosphereModel(grid; thermodynamic_constants = ThermodynamicConstants(eltype(grid)), formulation = :LiquidIcePotentialTemperature, dynamics = nothing, + velocities = nothing, moisture_density = DefaultValue(), tracers = tuple(), coriolis = nothing, @@ -123,6 +124,9 @@ function AtmosphereModel(grid; # Use default dynamics if not specified isnothing(dynamics) && (dynamics = default_dynamics(grid, thermodynamic_constants)) + # Validate that velocity boundary conditions are only provided for dynamics that support them + validate_velocity_boundary_conditions(dynamics, boundary_conditions) + if !(advection isa DefaultValue) # TODO: check that tracer+momentum advection were not independently set. scalar_advection = momentum_advection = advection @@ -144,7 +148,9 @@ function AtmosphereModel(grid; # Get field names from dynamics and formulation prognostic_names = prognostic_field_names(dynamics, formulation, microphysics, tracers) - default_boundary_conditions = NamedTuple{prognostic_names}(FieldBoundaryConditions() for _ in prognostic_names) + velocity_bc_names = velocity_boundary_condition_names(dynamics) + default_bc_names = tuple(prognostic_names..., velocity_bc_names...) + default_boundary_conditions = NamedTuple{default_bc_names}(FieldBoundaryConditions() for _ in default_bc_names) boundary_conditions = merge(default_boundary_conditions, boundary_conditions) # Pre-regularize AtmosphereModel boundary conditions (fill in reference_density, compute saturation humidity, etc.) @@ -155,10 +161,19 @@ function AtmosphereModel(grid; regularized_boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, all_names) # Materialize dynamics and formulation - dynamics = materialize_dynamics(dynamics, grid, regularized_boundary_conditions) + dynamics = materialize_dynamics(dynamics, grid, regularized_boundary_conditions, thermodynamic_constants) formulation = materialize_formulation(formulation, dynamics, grid, regularized_boundary_conditions) - momentum, velocities = materialize_momentum_and_velocities(dynamics, grid, regularized_boundary_conditions) + # Materialize momentum and velocities + # If velocities is provided (e.g., PrescribedVelocityFields), use it + if isnothing(velocities) + momentum, velocities = materialize_momentum_and_velocities(dynamics, grid, regularized_boundary_conditions) + else + # Store velocity specification in dynamics for dispatch (e.g., PrescribedVelocityFields) + dynamics = update_dynamics_with_velocities(dynamics, velocities) + momentum, _ = materialize_momentum_and_velocities(dynamics, grid, regularized_boundary_conditions) + velocities = materialize_velocities(velocities, grid) + end microphysical_fields = materialize_microphysical_fields(microphysics, grid, regularized_boundary_conditions) tracers = NamedTuple(name => CenterField(grid, boundary_conditions=regularized_boundary_conditions[name]) for name in tracer_names) @@ -272,10 +287,10 @@ Advection.cell_advection_timescale(model::AtmosphereModel) = cell_advection_time # Prognostic field names from dynamics + thermodynamic formulation + microphysics + tracers function prognostic_field_names(dynamics, formulation, microphysics, tracer_names) - default_names = (:ρu, :ρv, :ρw, :ρqᵗ) + momentum_names = prognostic_momentum_field_names(dynamics) formulation_names = prognostic_thermodynamic_field_names(formulation) microphysical_names = prognostic_field_names(microphysics) - return tuple(default_names..., formulation_names..., microphysical_names..., tracer_names...) + return tuple(momentum_names..., :ρqᵗ, formulation_names..., microphysical_names..., tracer_names...) end function field_names(dynamics, formulation, microphysics, tracer_names) diff --git a/src/AtmosphereModels/dynamics_interface.jl b/src/AtmosphereModels/dynamics_interface.jl index 73e62786..0201488e 100644 --- a/src/AtmosphereModels/dynamics_interface.jl +++ b/src/AtmosphereModels/dynamics_interface.jl @@ -18,7 +18,7 @@ Return the default dynamics for the given grid and thermodynamic constants. function default_dynamics end """ - materialize_dynamics(dynamics_stub, grid, boundary_conditions) + materialize_dynamics(dynamics_stub, grid, boundary_conditions, thermodynamic_constants) Materialize a dynamics stub into a complete dynamics object with all required fields. """ @@ -31,6 +31,21 @@ Create momentum and velocity fields for the given dynamics. """ function materialize_momentum_and_velocities end +""" + materialize_velocities(velocities, grid) + +Create velocity fields from a velocity specification (e.g., `PrescribedVelocityFields`). +""" +function materialize_velocities end + +""" + update_dynamics_with_velocities(dynamics, velocities) + +Update dynamics with velocity specification. Default is a no-op. +For `PrescribedDynamics`, stores the `PrescribedVelocityFields` for dispatch. +""" +update_dynamics_with_velocities(dynamics, velocities) = dynamics + """ dynamics_pressure_solver(dynamics, grid) @@ -110,6 +125,33 @@ function buoyancy_forceᶜᶜᶜ end ##### Boundary condition interface ##### +""" + validate_velocity_boundary_conditions(dynamics, user_boundary_conditions) + +Validate that velocity boundary conditions are only provided for dynamics that support them. + +By default, throws an error if the user provides boundary conditions for `:u`, `:v`, or `:w`, +since velocity is a diagnostic field for most dynamics (e.g., anelastic, compressible). + +For `PrescribedDynamics`, velocity boundary conditions are allowed since velocities are +regular fields that can be set directly. +""" +function validate_velocity_boundary_conditions(dynamics, user_boundary_conditions) + velocity_names = (:u, :v, :w) + user_bc_names = keys(user_boundary_conditions) + provided_velocity_bcs = filter(name -> name ∈ user_bc_names, velocity_names) + + if !isempty(provided_velocity_bcs) + throw(ArgumentError( + "Boundary conditions for velocity components $(provided_velocity_bcs) are not supported " * + "for $(summary(dynamics)). Velocity boundary conditions are only valid for PrescribedDynamics " * + "(kinematic models) where velocities are regular fields. For prognostic dynamics, " * + "set boundary conditions on momentum fields (:ρu, :ρv, :ρw) instead." + )) + end + return nothing +end + """ surface_pressure(dynamics) @@ -140,6 +182,16 @@ initialize_model_thermodynamics!(model) = nothing # default: do nothing ##### Prognostic fields interface ##### +""" + prognostic_momentum_field_names(dynamics) + +Return a tuple of prognostic momentum field names. + +For prognostic dynamics (anelastic, compressible), returns `(:ρu, :ρv, :ρw)`. +For kinematic dynamics (prescribed velocities), returns an empty tuple. +""" +prognostic_momentum_field_names(::Any) = (:ρu, :ρv, :ρw) + """ prognostic_dynamics_field_names(dynamics) @@ -157,6 +209,19 @@ Return a tuple of additional (diagnostic) field names for the dynamics. """ additional_dynamics_field_names(::Any) = () +""" + velocity_boundary_condition_names(dynamics) + +Return a tuple of velocity field names that need default boundary conditions. + +For most dynamics (anelastic, compressible), velocities are diagnostic and their boundary +conditions are created internally. Returns an empty tuple. + +For `PrescribedDynamics`, velocities are regular fields that can have user-provided +boundary conditions, so this returns `(:u, :v, :w)`. +""" +velocity_boundary_condition_names(::Any) = () + """ dynamics_prognostic_fields(dynamics) diff --git a/src/AtmosphereModels/dynamics_kernel_functions.jl b/src/AtmosphereModels/dynamics_kernel_functions.jl index 719c3a6b..f7dda965 100644 --- a/src/AtmosphereModels/dynamics_kernel_functions.jl +++ b/src/AtmosphereModels/dynamics_kernel_functions.jl @@ -7,6 +7,7 @@ using Oceananigans.Utils: sum_of_velocities @inline ∂ⱼ_𝒯₂ⱼ(i, j, k, grid, args...) = zero(grid) @inline ∂ⱼ_𝒯₃ⱼ(i, j, k, grid, args...) = zero(grid) @inline div_ρUc(i, j, k, grid, args...) = zero(grid) +@inline c_div_ρU(i, j, k, grid, args...) = zero(grid) """ ∇_dot_Jᶜ(i, j, k, grid, ρ, closure::AbstractTurbulenceClosure, closure_fields, @@ -139,6 +140,7 @@ end 𝒰 = diagnose_thermodynamic_state(i, j, k, grid, formulation, dynamics, q) return ( - div_ρUc(i, j, k, grid, advection, ρ_field, Uᵗ, c) + + c_div_ρU(i, j, k, grid, dynamics, velocities, c) # for PrescribedDynamics - ∇_dot_Jᶜ(i, j, k, grid, ρ_field, closure, closure_fields, id, c, clock, model_fields, closure_buoyancy) + microphysical_tendency(i, j, k, grid, microphysics, name, ρ, microphysical_fields, 𝒰, constants) + c_forcing(i, j, k, grid, clock, model_fields)) diff --git a/src/AtmosphereModels/set_atmosphere_model.jl b/src/AtmosphereModels/set_atmosphere_model.jl index f5354a2a..cb7e4bb5 100644 --- a/src/AtmosphereModels/set_atmosphere_model.jl +++ b/src/AtmosphereModels/set_atmosphere_model.jl @@ -28,6 +28,36 @@ end const settable_thermodynamic_variables = (:ρθ, :θ, :ρθˡⁱ, :θˡⁱ, :ρe, :e, :T) function set_thermodynamic_variable! end +##### +##### Velocity and momentum setting (extensible for kinematic models) +##### + +""" + set_velocity!(model, name, value) + +Set the velocity component `name` (`:u`, `:v`, or `:w`) to `value`. +Also updates the corresponding momentum field. +""" +function set_velocity!(model::AtmosphereModel, name::Symbol, value) + u = model.velocities[name] + set!(u, value) + ρ = dynamics_density(model.dynamics) + ϕ = model.momentum[Symbol(:ρ, name)] + set!(ϕ, ρ * u) + return nothing +end + +""" + set_momentum!(model, name, value) + +Set the momentum component `name` (`:ρu`, `:ρv`, or `:ρw`) to `value`. +""" +function set_momentum!(model::AtmosphereModel, name::Symbol, value) + ρu = getproperty(model.momentum, name) + set!(ρu, value) + return nothing +end + """ $(TYPEDSIGNATURES) @@ -130,8 +160,7 @@ function Fields.set!(model::AtmosphereModel; time=nothing, enforce_mass_conserva # Prognostic variables if name ∈ propertynames(model.momentum) - ρu = getproperty(model.momentum, name) - set!(ρu, value) + set_momentum!(model, name, value) elseif name ∈ propertynames(model.tracers) c = getproperty(model.tracers, name) @@ -163,13 +192,7 @@ function Fields.set!(model::AtmosphereModel; time=nothing, enforce_mass_conserva set!(ρqᵗ, ρ * qᵗ) elseif name ∈ (:u, :v, :w) - u = model.velocities[name] - set!(u, value) - - ρ = dynamics_density(model.dynamics) - ϕ = model.momentum[Symbol(:ρ, name)] - value = ρ * u - set!(ϕ, value) + set_velocity!(model, name, value) elseif name ∈ settable_thermodynamic_variables set_thermodynamic_variable!(model, Val(name), value) diff --git a/src/AtmosphereModels/update_atmosphere_model_state.jl b/src/AtmosphereModels/update_atmosphere_model_state.jl index 32c0ee37..9f30b31f 100644 --- a/src/AtmosphereModels/update_atmosphere_model_state.jl +++ b/src/AtmosphereModels/update_atmosphere_model_state.jl @@ -64,19 +64,16 @@ diagnostic_indices(::Bounded, N, H) = 1:N+1 diagnostic_indices(::Periodic, N, H) = -H+1:N+H diagnostic_indices(::Flat, N, H) = 1:N +##### +##### Velocity and momentum computation +##### + """ $(TYPEDSIGNATURES) -Compute auxiliary model variables: - -- velocities from momentum and density (eg ``u = ρu / ρ``) - -- thermodynamic variables from the prognostic thermodynamic state, - * temperature ``T``, possibly involving saturation adjustment - * specific thermodynamic variable (``e = ρe / ρ`` or ``θ = ρθ / ρ``) - * moisture mass fraction ``qᵗ = ρqᵗ / ρ`` +Compute velocities from momentum: `u = ρu / ρ` for each velocity component. """ -function compute_auxiliary_variables!(model) +function compute_velocities!(model::AtmosphereModel) grid = model.grid arch = grid.architecture @@ -106,6 +103,66 @@ function compute_auxiliary_variables!(model) foreach(mask_immersed_field!, model.velocities) fill_halo_regions!(model.velocities) + return nothing +end + +function compute_momentum_tendencies!(model::AtmosphereModel, model_fields) + grid = model.grid + arch = grid.architecture + Gρu = model.timestepper.Gⁿ.ρu + Gρv = model.timestepper.Gⁿ.ρv + Gρw = model.timestepper.Gⁿ.ρw + + momentum_args = ( + dynamics_density(model.dynamics), + model.advection.momentum, + model.velocities, + model.closure, + model.closure_fields, + model.momentum, + model.coriolis, + model.clock, + model_fields) + + u_args = tuple(momentum_args..., model.forcing.ρu, model.dynamics) + v_args = tuple(momentum_args..., model.forcing.ρv, model.dynamics) + + # Extra arguments for vertical velocity are required to compute buoyancy + w_args = tuple(momentum_args..., model.forcing.ρw, + model.dynamics, + model.formulation, + model.temperature, + model.specific_moisture, + model.microphysics, + model.microphysical_fields, + model.thermodynamic_constants) + + launch!(arch, grid, :xyz, compute_x_momentum_tendency!, Gρu, grid, u_args) + launch!(arch, grid, :xyz, compute_y_momentum_tendency!, Gρv, grid, v_args) + launch!(arch, grid, :xyz, compute_z_momentum_tendency!, Gρw, grid, w_args) + + return nothing +end + +""" +$(TYPEDSIGNATURES) + +Compute auxiliary model variables: + +- velocities from momentum and density (eg ``u = ρu / ρ``) + +- thermodynamic variables from the prognostic thermodynamic state, + * temperature ``T``, possibly involving saturation adjustment + * specific thermodynamic variable (``e = ρe / ρ`` or ``θ = ρθ / ρ``) + * moisture mass fraction ``qᵗ = ρqᵗ / ρ`` +""" +function compute_auxiliary_variables!(model) + grid = model.grid + arch = grid.architecture + + # Compute velocities from momentum (skip for kinematic dynamics with prescribed velocities) + compute_velocities!(model) + # Dispatch on thermodynamic formulation type compute_auxiliary_thermodynamic_variables!(model) @@ -205,44 +262,14 @@ end function compute_tendencies!(model::AtmosphereModel) grid = model.grid arch = grid.architecture - Gρu = model.timestepper.Gⁿ.ρu - Gρv = model.timestepper.Gⁿ.ρv - Gρw = model.timestepper.Gⁿ.ρw model_fields = fields(model) ##### - ##### Momentum tendencies + ##### Momentum tendencies (skip for kinematic dynamics) ##### - momentum_args = ( - dynamics_density(model.dynamics), - model.advection.momentum, - model.velocities, - model.closure, - model.closure_fields, - model.momentum, - model.coriolis, - model.clock, - model_fields) - - u_args = tuple(momentum_args..., model.forcing.ρu, model.dynamics) - v_args = tuple(momentum_args..., model.forcing.ρv, model.dynamics) - - # Extra arguments for vertical velocity are required to compute - # buoyancy: - w_args = tuple(momentum_args..., model.forcing.ρw, - model.dynamics, - model.formulation, - model.temperature, - model.specific_moisture, - model.microphysics, - model.microphysical_fields, - model.thermodynamic_constants) - - launch!(arch, grid, :xyz, compute_x_momentum_tendency!, Gρu, grid, u_args) - launch!(arch, grid, :xyz, compute_y_momentum_tendency!, Gρv, grid, v_args) - launch!(arch, grid, :xyz, compute_z_momentum_tendency!, Gρw, grid, w_args) + compute_momentum_tendencies!(model, model_fields) # Arguments common to energy density, moisture density, and tracer density tendencies: common_args = ( diff --git a/src/Breeze.jl b/src/Breeze.jl index 48b3ad3d..bc0199a9 100644 --- a/src/Breeze.jl +++ b/src/Breeze.jl @@ -10,10 +10,14 @@ export MoistAirBuoyancy, ThermodynamicConstants, ReferenceState, + surface_density, AnelasticDynamics, AnelasticModel, CompressibleDynamics, CompressibleModel, + PrescribedDensity, + PrescribedDynamics, + KinematicModel, AtmosphereModel, StaticEnergyFormulation, LiquidIcePotentialTemperatureFormulation, @@ -151,6 +155,9 @@ using .AnelasticEquations: AnelasticDynamics, AnelasticModel include("CompressibleEquations/CompressibleEquations.jl") using .CompressibleEquations: CompressibleDynamics, CompressibleModel +include("KinematicDriver/KinematicDriver.jl") +using .KinematicDriver: PrescribedDensity, PrescribedDynamics, KinematicModel + include("Microphysics/Microphysics.jl") using .Microphysics diff --git a/src/CompressibleEquations/compressible_dynamics.jl b/src/CompressibleEquations/compressible_dynamics.jl index c53551e9..482b7789 100644 --- a/src/CompressibleEquations/compressible_dynamics.jl +++ b/src/CompressibleEquations/compressible_dynamics.jl @@ -35,7 +35,7 @@ $(TYPEDSIGNATURES) Materialize a stub `CompressibleDynamics` into a full dynamics object with density and pressure fields. """ -function AtmosphereModels.materialize_dynamics(dynamics::CompressibleDynamics, grid, boundary_conditions) +function AtmosphereModels.materialize_dynamics(dynamics::CompressibleDynamics, grid, boundary_conditions, thermodynamic_constants) # Get density boundary conditions if provided if haskey(boundary_conditions, :ρ) density = CenterField(grid, boundary_conditions=boundary_conditions.ρ) diff --git a/src/KinematicDriver/KinematicDriver.jl b/src/KinematicDriver/KinematicDriver.jl new file mode 100644 index 00000000..306db506 --- /dev/null +++ b/src/KinematicDriver/KinematicDriver.jl @@ -0,0 +1,46 @@ +""" + KinematicDriver + +Module implementing kinematic dynamics for atmosphere models. + +Kinematic dynamics prescribes the velocity field rather than solving for it, +enabling isolated testing of microphysics, thermodynamics, and other physics +without the complexity of solving the momentum equations. + +This is analogous to the `kin1d` driver in P3-microphysics. +""" +module KinematicDriver + +export + PrescribedDensity, + PrescribedDynamics, + KinematicModel + +using DocStringExtensions: TYPEDSIGNATURES, TYPEDEF +using Adapt: Adapt, adapt + +using Oceananigans: Oceananigans, CenterField, XFaceField, YFaceField, ZFaceField +using Oceananigans.BoundaryConditions: FieldBoundaryConditions, ValueBoundaryCondition, fill_halo_regions! +using Oceananigans.Architectures: on_architecture +using Oceananigans.Fields: AbstractField, FunctionField, ZeroField, field, set! +using Oceananigans.Grids: Face, Center +using Oceananigans.Operators: Δzᶜᶜᶜ +using Oceananigans.TimeSteppers: Clock, TimeSteppers +using Oceananigans.Utils: launch!, prettysummary + +using KernelAbstractions: @kernel, @index + +# Import PrescribedVelocityFields from Oceananigans +using Oceananigans.Models.HydrostaticFreeSurfaceModels: PrescribedVelocityFields + +using Breeze.AtmosphereModels: AtmosphereModels, AtmosphereModel, dynamics_density +using Breeze.Thermodynamics: ReferenceState + +include("prescribed_dynamics.jl") + +# Type alias for kinematic models +const KinematicModel = AtmosphereModel{<:PrescribedDynamics} + +include("kinematic_driver_time_stepping.jl") + +end # module diff --git a/src/KinematicDriver/kinematic_driver_time_stepping.jl b/src/KinematicDriver/kinematic_driver_time_stepping.jl new file mode 100644 index 00000000..a8fe157a --- /dev/null +++ b/src/KinematicDriver/kinematic_driver_time_stepping.jl @@ -0,0 +1,99 @@ +##### +##### Time stepping for PrescribedDynamics (kinematic dynamics) +##### + +using KernelAbstractions: @kernel, @index + +using Oceananigans.BoundaryConditions: fill_halo_regions! +using Oceananigans.Fields: set!, FunctionField +using Oceananigans.Operators: V⁻¹ᶜᶜᶜ, δxᶜᵃᵃ, δyᵃᶜᵃ, δzᵃᵃᶜ, ℑxᶠᵃᵃ, ℑyᵃᶠᵃ, ℑzᵃᵃᶠ +using Oceananigans.Utils: launch! + +##### +##### Model initialization +##### + +AtmosphereModels.initialize_model_thermodynamics!(::KinematicModel) = nothing + +##### +##### Velocity and momentum: no-ops (no momentum, velocities may be FunctionFields) +##### + +AtmosphereModels.compute_velocities!(::KinematicModel) = nothing +AtmosphereModels.compute_momentum_tendencies!(::KinematicModel, model_fields) = nothing + +##### +##### Setting velocities for kinematic models +##### + +# Dispatch on velocity field type +AtmosphereModels.set_velocity!(model::KinematicModel, name::Symbol, value) = + set_velocity!(model.velocities[name], value) + +# Regular velocity fields: just set directly +set_velocity!(velocity::AbstractField, value) = set!(velocity, value) + +# FunctionFields (from PrescribedVelocityFields): cannot be set +set_velocity!(::FunctionField, value) = + throw(ArgumentError("Cannot set velocity component of PrescribedVelocityFields.")) + +# No momentum in kinematic models +AtmosphereModels.set_momentum!(::KinematicModel, name::Symbol, value) = + throw(ArgumentError("Cannot set momentum component '$name' of a KinematicModel.")) + +##### +##### Pressure correction: no-op for kinematic dynamics +##### + +TimeSteppers.compute_pressure_correction!(::KinematicModel, Δt) = nothing +TimeSteppers.make_pressure_correction!(::KinematicModel, Δt) = nothing + +##### +##### Density tendency (prognostic density only) +##### + +@inline mass_flux_x(i, j, k, grid, ρ, u) = ℑxᶠᵃᵃ(i, j, k, grid, ρ) * u[i, j, k] +@inline mass_flux_y(i, j, k, grid, ρ, v) = ℑyᵃᶠᵃ(i, j, k, grid, ρ) * v[i, j, k] +@inline mass_flux_z(i, j, k, grid, ρ, w) = ℑzᵃᵃᶠ(i, j, k, grid, ρ) * w[i, j, k] + +@inline function div_ρU(i, j, k, grid, ρ, velocities) + return V⁻¹ᶜᶜᶜ(i, j, k, grid) * ( + δxᶜᵃᵃ(i, j, k, grid, mass_flux_x, ρ, velocities.u) + + δyᵃᶜᵃ(i, j, k, grid, mass_flux_y, ρ, velocities.v) + + δzᵃᵃᶜ(i, j, k, grid, mass_flux_z, ρ, velocities.w) + ) +end + +# Default: no divergence correction +@inline AtmosphereModels.c_div_ρU(i, j, k, grid, ::PrescribedDynamics, velocities, c) = zero(grid) + +# With divergence correction: c * ∇·(ρU) +@inline function AtmosphereModels.c_div_ρU(i, j, k, grid, dynamics::PrescribedDynamics{true}, velocities, c) + return @inbounds c[i, j, k] * div_ρU(i, j, k, grid, dynamics_density(dynamics), velocities) +end + +# Fixed density: no tendency to compute +AtmosphereModels.compute_dynamics_tendency!(::AtmosphereModel{<:PrescribedDynamics{<:Any, <:PrescribedDensity}}) = nothing + +# Prognostic density: compute tendency from continuity equation +function AtmosphereModels.compute_dynamics_tendency!(model::KinematicModel) + grid = model.grid + arch = grid.architecture + ρ = dynamics_density(model.dynamics) + Gρ = model.timestepper.Gⁿ.ρ + + launch!(arch, grid, :xyz, _compute_density_tendency!, Gρ, grid, ρ, model.velocities) + + return nothing +end + +@kernel function _compute_density_tendency!(Gρ, grid, ρ, velocities) + i, j, k = @index(Global, NTuple) + @inbounds Gρ[i, j, k] = -div_ρU(i, j, k, grid, ρ, velocities) +end + +##### +##### Pressure: always prescribed in kinematic models +##### + +AtmosphereModels.compute_auxiliary_dynamics_variables!(::KinematicModel) = nothing diff --git a/src/KinematicDriver/prescribed_dynamics.jl b/src/KinematicDriver/prescribed_dynamics.jl new file mode 100644 index 00000000..711966e0 --- /dev/null +++ b/src/KinematicDriver/prescribed_dynamics.jl @@ -0,0 +1,238 @@ +##### +##### PrescribedDensity: wrapper for fixed density +##### + +""" +$(TYPEDSIGNATURES) + +Wrapper indicating that density is fixed (not prognostic). +""" +struct PrescribedDensity{D} + density :: D +end + +Base.summary(::PrescribedDensity) = "PrescribedDensity" +Base.eltype(pd::PrescribedDensity) = eltype(pd.density) +Base.show(io::IO, d::PrescribedDensity) = print(io, "PrescribedDensity(", summary(d.density), ")") + +Adapt.adapt_structure(to, pd::PrescribedDensity) = PrescribedDensity(adapt(to, pd.density)) + +Oceananigans.Architectures.on_architecture(to, pd::PrescribedDensity) = + PrescribedDensity(on_architecture(to, pd.density)) + +##### +##### PrescribedDynamics: kinematic model dynamics +##### + +""" +$(TYPEDEF) + +Dynamics for kinematic atmosphere models where velocity is prescribed. +The type parameter `Div` indicates whether divergence correction is applied. +""" +struct PrescribedDynamics{Div, D, P, FT} + density :: D + pressure :: P + surface_pressure :: FT + standard_pressure :: FT +end + +# Convenient method for letting the user specify only the divergence parameter +# and infer the others. +function PrescribedDynamics{Div}(density::D, + pressure::P, + surface_pressure::FT, + standard_pressure::FT) where {Div, D, P, FT} + return PrescribedDynamics{Div, D, P, FT}(density, pressure, surface_pressure, standard_pressure) +end + +""" +$(TYPEDSIGNATURES) + +Construct `PrescribedDynamics` from a [`ReferenceState`](@ref). +Wraps density in `PrescribedDensity` (fixed in time). + +If `divergence_correction=true`, scalar tendencies include `+c∇·(ρU)` to +account for the non-divergent velocity field. + +# Example + +```jldoctest +using Oceananigans +using Breeze + +grid = RectilinearGrid(size=(4, 4, 8), extent=(1000, 1000, 2000)) +reference_state = ReferenceState(grid, ThermodynamicConstants()) +dynamics = PrescribedDynamics(reference_state) + +# output +PrescribedDynamics +├── density: PrescribedDensity +├── pressure: 1×1×8 Field{Nothing, Nothing, Center} reduced over dims = (1, 2) on RectilinearGrid on CPU +├── surface_pressure: 101325.0 +└── standard_pressure: 100000.0 +``` +""" +function PrescribedDynamics(reference_state::ReferenceState; divergence_correction=false) + density = PrescribedDensity(reference_state.density) + pressure = reference_state.pressure + p₀ = reference_state.surface_pressure + pˢᵗ = reference_state.standard_pressure + return PrescribedDynamics{divergence_correction}(density, pressure, p₀, pˢᵗ) +end + +""" +$(TYPEDSIGNATURES) + +Construct `PrescribedDynamics` from a density field or `PrescribedDensity`. +If `pressure=nothing`, hydrostatic pressure is computed during materialization. +""" +function PrescribedDynamics(density; + pressure = nothing, + surface_pressure = 101325, + standard_pressure = 1e5, + divergence_correction = false) + + FT = eltype(density) + return PrescribedDynamics{divergence_correction}(density, pressure, + convert(FT, surface_pressure), + convert(FT, standard_pressure)) +end + +Base.summary(::PrescribedDynamics) = "PrescribedDynamics" + +function Base.show(io::IO, d::PrescribedDynamics) + print(io, "PrescribedDynamics\n") + print(io, "├── density: ", summary(d.density), '\n') + print(io, "├── pressure: ", prettysummary(d.pressure), '\n') + print(io, "├── surface_pressure: ", prettysummary(d.surface_pressure), '\n') + print(io, "└── standard_pressure: ", prettysummary(d.standard_pressure)) +end + +##### +##### Dynamics interface +##### + +# Extract the underlying density field +@inline unwrap_density(pd::PrescribedDensity) = pd.density +@inline unwrap_density(ρ) = ρ # pass-through for regular fields + +@inline AtmosphereModels.dynamics_density(d::PrescribedDynamics) = unwrap_density(d.density) + +AtmosphereModels.prognostic_momentum_field_names(::PrescribedDynamics) = () +AtmosphereModels.additional_dynamics_field_names(::PrescribedDynamics) = () +AtmosphereModels.validate_velocity_boundary_conditions(::PrescribedDynamics, bcs) = nothing +AtmosphereModels.velocity_boundary_condition_names(::PrescribedDynamics) = (:u, :v, :w) + +# Prescribed density → no prognostic density; otherwise ρ is prognostic +AtmosphereModels.prognostic_dynamics_field_names(::PrescribedDynamics{<:Any, <:PrescribedDensity}) = () +AtmosphereModels.prognostic_dynamics_field_names(::PrescribedDynamics) = tuple(:ρ) + +AtmosphereModels.dynamics_prognostic_fields(::PrescribedDynamics{<:Any, <:PrescribedDensity}) = NamedTuple() +AtmosphereModels.dynamics_prognostic_fields(d::PrescribedDynamics) = (; ρ=dynamics_density(d)) + +# Pressure accessors +AtmosphereModels.dynamics_pressure_solver(::PrescribedDynamics, grid) = nothing +AtmosphereModels.dynamics_pressure(d::PrescribedDynamics) = d.pressure +AtmosphereModels.mean_pressure(d::PrescribedDynamics) = d.pressure +AtmosphereModels.pressure_anomaly(::PrescribedDynamics) = ZeroField() +AtmosphereModels.total_pressure(d::PrescribedDynamics) = d.pressure +AtmosphereModels.surface_pressure(d::PrescribedDynamics) = d.surface_pressure +AtmosphereModels.standard_pressure(d::PrescribedDynamics) = d.standard_pressure + +##### +##### Materialization +##### + +function AtmosphereModels.materialize_dynamics(d::PrescribedDynamics{Div}, grid, bcs, constants) where Div + FT = eltype(grid) + p₀ = convert(FT, d.surface_pressure) + pˢᵗ = convert(FT, d.standard_pressure) + g = constants.gravitational_acceleration + density = materialize_density(d.density, grid, bcs) + pressure = materialize_pressure(d.pressure, density, p₀, g, grid) + return PrescribedDynamics{Div}(density, pressure, p₀, pˢᵗ) +end + +materialize_density(density::AbstractField, grid, bcs) = density +materialize_pressure(pressure::AbstractField, args...) = pressure + +function materialize_density(density::PrescribedDensity, grid, bcs) + ρ = materialize_density(density.density, grid, bcs) + return PrescribedDensity(ρ) +end + +function materialize_density(density, grid, bcs) + ρ_bcs = haskey(bcs, :ρ) ? bcs.ρ : FieldBoundaryConditions() + ρ = CenterField(grid, boundary_conditions=ρ_bcs) + if !isnothing(density) + set!(ρ, density) + fill_halo_regions!(ρ) + end + return ρ +end + +function materialize_pressure(pressure, density, p₀, g, grid) + loc = (Center(), Center(), Center()) + p_bcs = FieldBoundaryConditions(grid, loc, bottom=ValueBoundaryCondition(p₀)) + p = CenterField(grid, boundary_conditions=p_bcs) + + if isnothing(pressure) + # Compute hydrostatic pressure: ∂p/∂z = -ρg + ρ = unwrap_density(density) + arch = grid.architecture + launch!(arch, grid, :xy, _hydrostatic_pressure!, p, ρ, p₀, g, grid) + else + set!(p, pressure) + end + + fill_halo_regions!(p) + return p +end + +@kernel function _hydrostatic_pressure!(p, ρ, p₀, g, grid) + i, j = @index(Global, NTuple) + @inbounds begin + pₖ = p₀ + for k in 1:grid.Nz + Δz = Δzᶜᶜᶜ(i, j, k, grid) + p[i, j, k] = pₖ - ρ[i, j, k] * g * Δz / 2 + pₖ = pₖ - ρ[i, j, k] * g * Δz + end + end +end + +##### +##### Velocity materialization +##### + +function AtmosphereModels.materialize_momentum_and_velocities(::PrescribedDynamics, grid, bcs) + u = XFaceField(grid, boundary_conditions=bcs.u) + v = YFaceField(grid, boundary_conditions=bcs.v) + w = ZFaceField(grid, boundary_conditions=bcs.w) + return NamedTuple(), (; u, v, w) +end + +function AtmosphereModels.materialize_velocities(velocities::PrescribedVelocityFields, grid) + clock = Clock{eltype(grid)}(time=0) + params = velocities.parameters + u = wrap_velocity(Face, Center, Center, velocities.u, grid; clock, parameters=params) + v = wrap_velocity(Center, Face, Center, velocities.v, grid; clock, parameters=params) + w = wrap_velocity(Center, Center, Face, velocities.w, grid; clock, parameters=params) + return (; u, v, w) +end + +wrap_velocity(X, Y, Z, f::Function, grid; kwargs...) = FunctionField{X, Y, Z}(f, grid; kwargs...) +wrap_velocity(X, Y, Z, f, grid; kwargs...) = field((X, Y, Z), f, grid) + +##### +##### Adapt and architecture transfer +##### + +Adapt.adapt_structure(to, d::PrescribedDynamics{Div}) where Div = + PrescribedDynamics{Div}(adapt(to, d.density), adapt(to, d.pressure), + d.surface_pressure, d.standard_pressure) + +Oceananigans.Architectures.on_architecture(to, d::PrescribedDynamics{Div}) where Div = + PrescribedDynamics{Div}(on_architecture(to, d.density), on_architecture(to, d.pressure), + d.surface_pressure, d.standard_pressure) diff --git a/src/PotentialTemperatureFormulations/PotentialTemperatureFormulations.jl b/src/PotentialTemperatureFormulations/PotentialTemperatureFormulations.jl index 5c968b5f..6732b065 100644 --- a/src/PotentialTemperatureFormulations/PotentialTemperatureFormulations.jl +++ b/src/PotentialTemperatureFormulations/PotentialTemperatureFormulations.jl @@ -21,7 +21,7 @@ using Oceananigans.Utils: prettysummary, launch! using Breeze.AtmosphereModels: AtmosphereModels, diagnose_thermodynamic_state, set_thermodynamic_variable!, dynamics_density, dynamics_pressure, standard_pressure, dynamics_prognostic_fields, compute_moisture_fractions, maybe_adjust_thermodynamic_state, - div_ρUc, ∇_dot_Jᶜ, AtmosphereModelBuoyancy, microphysical_tendency + div_ρUc, c_div_ρU, ∇_dot_Jᶜ, AtmosphereModelBuoyancy, microphysical_tendency using Breeze.Thermodynamics: LiquidIcePotentialTemperatureState, StaticEnergyState, with_temperature, exner_function, mixture_heat_capacity # The lowercase c is a singleton instance of Center @@ -38,4 +38,3 @@ include("potential_temperature_tendency.jl") end end # module - diff --git a/src/PotentialTemperatureFormulations/potential_temperature_tendency.jl b/src/PotentialTemperatureFormulations/potential_temperature_tendency.jl index 8f418ce1..a2b8543f 100644 --- a/src/PotentialTemperatureFormulations/potential_temperature_tendency.jl +++ b/src/PotentialTemperatureFormulations/potential_temperature_tendency.jl @@ -66,6 +66,7 @@ end closure_buoyancy = AtmosphereModelBuoyancy(dynamics, formulation, constants) return ( - div_ρUc(i, j, k, grid, advection, ρ_field, velocities, potential_temperature) + + c_div_ρU(i, j, k, grid, dynamics, velocities, potential_temperature) - ∇_dot_Jᶜ(i, j, k, grid, ρ_field, closure, closure_fields, id, potential_temperature, clock, model_fields, closure_buoyancy) + microphysical_tendency(i, j, k, grid, microphysics, Val(:ρθ), ρ, microphysical_fields, 𝒰, constants) + ρθ_forcing(i, j, k, grid, clock, model_fields) @@ -219,4 +220,3 @@ end @inbounds potential_temperature[i, j, k] = θ @inbounds potential_temperature_density[i, j, k] = ρᵣ * θ end - diff --git a/src/StaticEnergyFormulations/StaticEnergyFormulations.jl b/src/StaticEnergyFormulations/StaticEnergyFormulations.jl index 68eba400..b22304bf 100644 --- a/src/StaticEnergyFormulations/StaticEnergyFormulations.jl +++ b/src/StaticEnergyFormulations/StaticEnergyFormulations.jl @@ -22,8 +22,9 @@ using Oceananigans.Utils: prettysummary, launch! using Breeze.AtmosphereModels: AtmosphereModels, diagnose_thermodynamic_state, dynamics_density, dynamics_pressure, standard_pressure, dynamics_prognostic_fields, - compute_moisture_fractions, maybe_adjust_thermodynamic_state, div_ρUc, ∇_dot_Jᶜ, - w_buoyancy_forceᶜᶜᶠ, AtmosphereModelBuoyancy, microphysical_tendency + compute_moisture_fractions, maybe_adjust_thermodynamic_state, div_ρUc, + c_div_ρU, ∇_dot_Jᶜ, w_buoyancy_forceᶜᶜᶠ, + AtmosphereModelBuoyancy, microphysical_tendency using Breeze.Thermodynamics: StaticEnergyState, LiquidIcePotentialTemperatureState, with_temperature diff --git a/src/StaticEnergyFormulations/static_energy_tendency.jl b/src/StaticEnergyFormulations/static_energy_tendency.jl index 09261608..31c0ad20 100644 --- a/src/StaticEnergyFormulations/static_energy_tendency.jl +++ b/src/StaticEnergyFormulations/static_energy_tendency.jl @@ -69,6 +69,7 @@ end closure_buoyancy = AtmosphereModelBuoyancy(dynamics, formulation, constants) return ( - div_ρUc(i, j, k, grid, advection, ρ_field, velocities, specific_energy) + + c_div_ρU(i, j, k, grid, dynamics, velocities, specific_energy) + buoyancy_flux - ∇_dot_Jᶜ(i, j, k, grid, ρ_field, closure, closure_fields, id, specific_energy, clock, model_fields, closure_buoyancy) + microphysical_tendency(i, j, k, grid, microphysics, Val(:ρe), ρ, microphysical_fields, 𝒰, constants) @@ -217,4 +218,3 @@ end @inbounds specific_energy[i, j, k] = e @inbounds energy_density[i, j, k] = ρᵣ * e end - diff --git a/src/Thermodynamics/Thermodynamics.jl b/src/Thermodynamics/Thermodynamics.jl index 753d1f5c..1b0a1f08 100644 --- a/src/Thermodynamics/Thermodynamics.jl +++ b/src/Thermodynamics/Thermodynamics.jl @@ -11,7 +11,7 @@ export ThermodynamicConstants, ReferenceState, IdealGas, saturation_vapor_pressure, saturation_specific_humidity, supersaturation, equilibrium_saturation_specific_humidity, adjustment_saturation_specific_humidity, vapor_pressure, relative_humidity, - adiabatic_hydrostatic_pressure, adiabatic_hydrostatic_density, + adiabatic_hydrostatic_pressure, adiabatic_hydrostatic_density, surface_density, PlanarLiquidSurface, PlanarIceSurface, PlanarMixedPhaseSurface, # Phase equilibrium types AbstractPhaseEquilibrium, WarmPhaseEquilibrium, MixedPhaseEquilibrium, diff --git a/src/Thermodynamics/reference_states.jl b/src/Thermodynamics/reference_states.jl index ee0be4b4..92cf4a2e 100644 --- a/src/Thermodynamics/reference_states.jl +++ b/src/Thermodynamics/reference_states.jl @@ -1,7 +1,9 @@ using Oceananigans: Oceananigans, Center, Field, set!, fill_halo_regions! using Oceananigans.BoundaryConditions: FieldBoundaryConditions, ValueBoundaryCondition +using Oceananigans.Operators: ℑzᵃᵃᶠ using Adapt: Adapt, adapt +using GPUArraysCore: @allowscalar ##### ##### Reference state computations for Boussinesq and Anelastic models @@ -37,6 +39,23 @@ Base.show(io::IO, ref::ReferenceState) = print(io, summary(ref)) ##### How to compute the reference state ##### +""" + surface_density(reference_state) + +Return the density at z=0 by interpolating the reference density field to the surface. +""" +function surface_density(ref::ReferenceState) + ρ = ref.density + grid = ρ.grid + return @allowscalar ℑzᵃᵃᶠ(1, 1, 1, grid, ρ) +end + +""" + surface_density(p₀, θ₀, constants) + +Compute the surface air density from surface pressure `p₀`, surface temperature `θ₀`, +and thermodynamic `constants` using the ideal gas law for dry air. +""" @inline function surface_density(p₀, θ₀, constants) Rᵈ = dry_air_gas_constant(constants) return p₀ / (Rᵈ * θ₀) diff --git a/test/Project.toml b/test/Project.toml index 851e81d0..ff808852 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" Breeze = "660aa2fb-d4c8-4359-a52c-9c057bc511da" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" @@ -19,6 +20,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Breeze = {path = ".."} [compat] +Adapt = "4.3.0" Aqua = "0.8" CUDA = "5" ClimaComms = "0.6" diff --git a/test/anelastic_dynamics.jl b/test/anelastic_dynamics.jl index bc03ba78..34a32f23 100644 --- a/test/anelastic_dynamics.jl +++ b/test/anelastic_dynamics.jl @@ -33,7 +33,7 @@ using Test dynamics_stub = AnelasticDynamics(reference_state) boundary_conditions = NamedTuple() - dynamics = materialize_dynamics(dynamics_stub, grid, boundary_conditions) + dynamics = materialize_dynamics(dynamics_stub, grid, boundary_conditions, constants) @test dynamics isa AnelasticDynamics @test dynamics.reference_state === reference_state @@ -43,7 +43,7 @@ using Test @testset "Pressure utilities" begin reference_state = ReferenceState(grid, constants; surface_pressure=101325, potential_temperature=300) dynamics_stub = AnelasticDynamics(reference_state) - dynamics = materialize_dynamics(dynamics_stub, grid, NamedTuple()) + dynamics = materialize_dynamics(dynamics_stub, grid, NamedTuple(), constants) # Test mean_pressure p̄ = mean_pressure(dynamics) diff --git a/test/compressible_dynamics.jl b/test/compressible_dynamics.jl index 6168d968..86f08517 100644 --- a/test/compressible_dynamics.jl +++ b/test/compressible_dynamics.jl @@ -18,7 +18,8 @@ using Test @testset "materialize_dynamics" begin dynamics_stub = CompressibleDynamics() - dynamics = materialize_dynamics(dynamics_stub, grid, NamedTuple()) + constants = ThermodynamicConstants() + dynamics = materialize_dynamics(dynamics_stub, grid, NamedTuple(), constants) @test dynamics isa CompressibleDynamics @test dynamics.density isa Field diff --git a/test/kinematic_driver.jl b/test/kinematic_driver.jl new file mode 100644 index 00000000..80cb090e --- /dev/null +++ b/test/kinematic_driver.jl @@ -0,0 +1,113 @@ +using Adapt: adapt +using Breeze +using Breeze: PrescribedDensity, PrescribedDynamics, KinematicModel +using GPUArraysCore: @allowscalar +using Oceananigans +using Oceananigans.Architectures: on_architecture +using Oceananigans.BoundaryConditions: FieldBoundaryConditions, OpenBoundaryCondition +using Oceananigans.Fields: ZeroField +using Oceananigans.Models.HydrostaticFreeSurfaceModels: PrescribedVelocityFields +using Test + +@testset "KinematicDriver [$(FT)]" for FT in (Float32, Float64) + Oceananigans.defaults.FloatType = FT + grid = RectilinearGrid(default_arch; size=(4, 4, 8), extent=(1000, 1000, 2000)) + constants = ThermodynamicConstants() + reference_state = ReferenceState(grid, constants) + + @testset "PrescribedDynamics construction" begin + dynamics = PrescribedDynamics(reference_state) + @test dynamics.density isa PrescribedDensity + @test dynamics_density(dynamics) === reference_state.density + + dynamics_div = PrescribedDynamics(reference_state; divergence_correction=true) + @test dynamics_div isa PrescribedDynamics{true} + end + + @testset "PrescribedDensity adapt and on_architecture" begin + pd = PrescribedDensity(reference_state.density) + + # Test adapt_structure + adapted_pd = adapt(CPU(), pd) + @test adapted_pd isa PrescribedDensity + + # Test on_architecture + transferred_pd = on_architecture(CPU(), pd) + @test transferred_pd isa PrescribedDensity + end + + @testset "KinematicModel with prognostic density" begin + ρ = CenterField(grid) + set!(ρ, FT(1)) + model = AtmosphereModel(grid; dynamics=PrescribedDynamics(ρ)) + @test haskey(Oceananigans.prognostic_fields(model), :ρ) + end + + @testset "KinematicModel with regular fields" begin + model = AtmosphereModel(grid; dynamics=PrescribedDynamics(reference_state)) + @test model isa KinematicModel + @test model.pressure_solver === nothing + + set!(model, θ=300, qᵗ=0.01, w=1) + @test @allowscalar(model.velocities.w[1, 1, 4]) ≈ FT(1) + + time_step!(model, 1) + @test model.clock.iteration == 1 + end + + @testset "KinematicModel with PrescribedVelocityFields" begin + w_evolving(x, y, z, t) = (1 - exp(-t / 100)) * sin(π * z / 2000) + + model = AtmosphereModel(grid; + dynamics = PrescribedDynamics(reference_state), + velocities = PrescribedVelocityFields(w=w_evolving)) + + @test model isa KinematicModel + set!(model, θ=300, qᵗ=0.01) + @test_throws ArgumentError set!(model, w=1) + + time_step!(model, 10) + @test model.clock.time ≈ 10 + end + + @testset "Velocity boundary conditions" begin + w_inlet(x, y, t) = FT(0.5) + w_bcs = FieldBoundaryConditions(bottom=OpenBoundaryCondition(w_inlet)) + boundary_conditions = (; w = w_bcs) + + model = AtmosphereModel(grid; dynamics=PrescribedDynamics(reference_state), boundary_conditions) + @test model isa KinematicModel + @test model.velocities.w.boundary_conditions.bottom isa Oceananigans.BoundaryConditions.BoundaryCondition + + # AnelasticDynamics does not allow velocity boundary conditions + @test_throws ArgumentError AtmosphereModel(grid; boundary_conditions) + end +end + +@testset "Gaussian advection (analytical solution) [Float64]" begin + FT = Float64 + Oceananigans.defaults.FloatType = FT + + Lz, Nz, w₀ = 4000, 64, 10 # Reduced resolution for faster test + grid = RectilinearGrid(default_arch; size=(4, 4, Nz), x=(0, 100), y=(0, 100), z=(0, Lz)) + + model = AtmosphereModel(grid; + dynamics = PrescribedDynamics(ReferenceState(grid, ThermodynamicConstants())), + tracers = :c, + advection = WENO()) + + z₀, σ = 1000, 100 + c_exact(x, y, z, t) = exp(-(z - z₀ - w₀ * t)^2 / (2 * σ^2)) + + set!(model, θ=300, qᵗ=0, w=w₀, c=(x, y, z) -> c_exact(x, y, z, 0)) + + stop_time = 50 + simulation = Simulation(model; Δt=1, stop_time) + run!(simulation) + + c_truth = CenterField(grid) + set!(c_truth, (x, y, z) -> c_exact(x, y, z, stop_time)) + + error = @allowscalar maximum(abs, interior(model.tracers.c) .- interior(c_truth)) + @test error < FT(0.1) # Relaxed tolerance for reduced resolution test +end diff --git a/test/thermodynamic_formulations.jl b/test/thermodynamic_formulations.jl index 66a94a85..eec02b58 100644 --- a/test/thermodynamic_formulations.jl +++ b/test/thermodynamic_formulations.jl @@ -14,7 +14,7 @@ using Test constants = ThermodynamicConstants() reference_state = ReferenceState(grid, constants) dynamics_stub = AnelasticDynamics(reference_state) - dynamics = materialize_dynamics(dynamics_stub, grid, NamedTuple()) + dynamics = materialize_dynamics(dynamics_stub, grid, NamedTuple(), constants) # Boundary conditions needed for materialization (must pass grid to respect topology) ccc = (Center(), Center(), Center())