From 74ca00a4dbfa5d86f5ebe83aeb5c94a4dddde548 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sat, 8 Nov 2025 12:14:43 -0500 Subject: [PATCH 1/5] add adapt_structure for MoistAirBuoyancy and ReferenceState --- examples/free_convection.jl | 15 +++++++-------- src/MoistAirBuoyancies.jl | 4 ++++ src/Thermodynamics/reference_states.jl | 7 +++++++ 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/examples/free_convection.jl b/examples/free_convection.jl index 01c1fb2..055f389 100644 --- a/examples/free_convection.jl +++ b/examples/free_convection.jl @@ -3,9 +3,8 @@ using Oceananigans.Units using Printf using Breeze -arch = CPU() - -Nx = Nz = 64 +arch = GPU() +Nx = Nz = 128 Lz = 4 * 1024 grid = RectilinearGrid(arch, size=(Nx, Nz), x=(0, 2Lz), z=(0, Lz), topology=(Periodic, Flat, Bounded)) @@ -26,17 +25,15 @@ vapor_flux = FluxBoundaryCondition(1e-2) qᵗ_bcs = FieldBoundaryConditions(bottom=vapor_flux) advection = WENO() #(momentum=WENO(), θ=WENO(), q=WENO(bounds=(0, 1))) -tracers = (:θ, :q) model = NonhydrostaticModel(; grid, advection, buoyancy, tracers = (:θ, :qᵗ), boundary_conditions = (θ=θ_bcs, qᵗ=qᵗ_bcs)) Lz = grid.Lz -Δθ = 5 # K +Δθ = 2 # K Tₛ = buoyancy.reference_state.potential_temperature # K -θᵢ(x, z) = Tₛ + Δθ * z / Lz + 1e-2 * Δθ * randn() -qᵗᵢ(x, z) = 0 # 1e-2 + 1e-5 * rand() -set!(model, θ=θᵢ, qᵗ=qᵗᵢ) +θᵢ(x, z) = Tₛ + Δθ * z / Lz + 1e-2 * Δθ * rand() +set!(model, θ=θᵢ) simulation = Simulation(model, Δt=10, stop_iteration=1000) #stop_time=4hours) conjure_time_step_wizard!(simulation, cfl=0.7) @@ -92,6 +89,7 @@ simulation.output_writers[:jld2] = ow run!(simulation) +#= wt = FieldTimeSeries("free_convection.jld2", "θ") θt = FieldTimeSeries("free_convection.jld2", "θ") Tt = FieldTimeSeries("free_convection.jld2", "T") @@ -154,3 +152,4 @@ record(fig, "free_convection.mp4", 1:Nt, framerate=12) do nn @info "Drawing frame $nn of $Nt..." n[] = nn end +=# \ No newline at end of file diff --git a/src/MoistAirBuoyancies.jl b/src/MoistAirBuoyancies.jl index a282f7c..1478749 100644 --- a/src/MoistAirBuoyancies.jl +++ b/src/MoistAirBuoyancies.jl @@ -42,6 +42,10 @@ struct MoistAirBuoyancy{RS, AT} <: AbstractBuoyancyFormulation{Nothing} thermodynamics :: AT end +Adapt.adapt_structure(to, mb::MoistAirBuoyancy) = + MoistAirBuoyancy(adapt(to, mb.reference_state), + adapt(to, mb.thermodynamics)) + """ MoistAirBuoyancy(grid; base_pressure = 101325, diff --git a/src/Thermodynamics/reference_states.jl b/src/Thermodynamics/reference_states.jl index cd7eb30..1f4d274 100644 --- a/src/Thermodynamics/reference_states.jl +++ b/src/Thermodynamics/reference_states.jl @@ -1,4 +1,5 @@ using Oceananigans: Oceananigans, Center, Field, set!, fill_halo_regions! +using Adapt: Adapt, adapt ##### ##### Reference state computations for Boussinesq and Anelastic models @@ -11,6 +12,12 @@ struct ReferenceState{FT, F} density :: F end +Adapt.adapt_structure(to, ref::ReferenceState) = + ReferenceState(adapt(to, ref.base_pressure), + adapt(to, ref.potential_temperature), + adapt(to, ref.pressure), + adapt(to, ref.density)) + Base.eltype(::ReferenceState{FT}) where FT = FT function Base.summary(ref::ReferenceState) From bdeb8ce58a67ca40e11f05d84e5364d8ead4aba9 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Sat, 8 Nov 2025 10:17:03 -0700 Subject: [PATCH 2/5] Update examples/free_convection.jl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Mosè Giordano <765740+giordano@users.noreply.github.com> --- examples/free_convection.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/free_convection.jl b/examples/free_convection.jl index 055f389..96a4e03 100644 --- a/examples/free_convection.jl +++ b/examples/free_convection.jl @@ -152,4 +152,4 @@ record(fig, "free_convection.mp4", 1:Nt, framerate=12) do nn @info "Drawing frame $nn of $Nt..." n[] = nn end -=# \ No newline at end of file +=# From f1ec689062261ff9a1d2220a7324b40d8cf50029 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sat, 8 Nov 2025 12:17:51 -0500 Subject: [PATCH 3/5] add using CUDA --- examples/free_convection.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/free_convection.jl b/examples/free_convection.jl index 96a4e03..59a8bcc 100644 --- a/examples/free_convection.jl +++ b/examples/free_convection.jl @@ -2,6 +2,7 @@ using Oceananigans using Oceananigans.Units using Printf using Breeze +using CUDA arch = GPU() Nx = Nz = 128 From ffaa72cb5931d4b7bb9c647012cfd82e3624fb04 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sat, 8 Nov 2025 14:19:00 -0500 Subject: [PATCH 4/5] fix --- examples/free_convection.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/free_convection.jl b/examples/free_convection.jl index 59a8bcc..198ee0d 100644 --- a/examples/free_convection.jl +++ b/examples/free_convection.jl @@ -68,7 +68,7 @@ function progress(sim) iteration(sim), prettytime(sim), prettytime(sim.Δt), umax, vmax, wmax) msg *= @sprintf(", max(qˡ): %.2e, min(δ): %.2e, extrema(θ): (%.2e, %.2e)", - qᵗmin, qᵗmax, qˡmax, δmax, θmin, θmax) + qˡmax, δmax, θmin, θmax) @info msg From 01e8f2a74af0a79ca157e0bc1c049aee35ee1441 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sat, 8 Nov 2025 14:36:16 -0500 Subject: [PATCH 5/5] add adapt for AnelasticFormulation --- src/AtmosphereModels/anelastic_formulation.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/AtmosphereModels/anelastic_formulation.jl b/src/AtmosphereModels/anelastic_formulation.jl index 4f16f49..12392c5 100644 --- a/src/AtmosphereModels/anelastic_formulation.jl +++ b/src/AtmosphereModels/anelastic_formulation.jl @@ -14,6 +14,7 @@ using Oceananigans.Operators: Δzᵃᵃᶜ, Δzᵃᵃᶠ, divᶜᶜᶜ, Δzᶜ using Oceananigans.Solvers: solve! using KernelAbstractions: @kernel, @index +using Adapt: Adapt, adapt import Oceananigans.Solvers: tridiagonal_direction, compute_main_diagonal!, compute_lower_diagonal! import Oceananigans.TimeSteppers: compute_pressure_correction!, make_pressure_correction! @@ -26,6 +27,9 @@ struct AnelasticFormulation{R} reference_state :: R end +Adapt.adapt_structure(to, formulation::AnelasticFormulation) = + AnelasticFormulation(adapt(to, formulation.reference_state)) + const AnelasticModel = AtmosphereModel{<:AnelasticFormulation} function Base.summary(formulation::AnelasticFormulation)