diff --git a/examples/free_convection.jl b/examples/free_convection.jl index 01c1fb2..198ee0d 100644 --- a/examples/free_convection.jl +++ b/examples/free_convection.jl @@ -2,10 +2,10 @@ using Oceananigans using Oceananigans.Units using Printf using Breeze +using CUDA -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 +26,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) @@ -70,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 @@ -92,6 +90,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 +153,4 @@ record(fig, "free_convection.mp4", 1:Nt, framerate=12) do nn @info "Drawing frame $nn of $Nt..." n[] = nn end +=# diff --git a/src/AtmosphereModels/anelastic_formulation.jl b/src/AtmosphereModels/anelastic_formulation.jl index f3b01fe..c6cb223 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) 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)