[GPU] Add adapt_structure for MoistAirBuoyancy, AnelasticFormulation, and ReferenceState (#94)

* add adapt_structure for MoistAirBuoyancy and ReferenceState

* Update examples/free_convection.jl

Co-authored-by: Mosè Giordano <[email protected]>

* add using CUDA

* fix

* add adapt for AnelasticFormulation

---------

Co-authored-by: Mosè Giordano <[email protected]>

Author: glwagner

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
+=#

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)

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,

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)