Lay the groundwork for prognostic atmospheres (#355)

* Lay the groundwork for prognostic atmospheres

* Fix PrescribedAtmosphere time_step!

* Bugfix in test_surface_fluxes

* Generalize interface for StateExchanger

* Omit sync clock

Author: glwagner

src/OceanSeaIceModels/InterfaceComputations/atmosphere_ocean_fluxes.jl

@@ -1,7 +1,7 @@
 using Oceananigans.Operators: intrinsic_vector
 using Oceananigans.Grids: inactive_node
 using ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: thermodynamics_parameters, 
-                                                          reference_height,
+                                                          surface_layer_height,
                                                           boundary_layer_height
 
 function compute_atmosphere_ocean_fluxes!(coupled_model)
@@ -36,7 +36,7 @@ function compute_atmosphere_ocean_fluxes!(coupled_model)
     interface_properties = coupled_model.interfaces.atmosphere_ocean_interface.properties
     ocean_properties = coupled_model.interfaces.ocean_properties
     atmosphere_properties = (thermodynamics_parameters = thermodynamics_parameters(atmosphere),
-                             reference_height = reference_height(atmosphere))
+                             surface_layer_height = surface_layer_height(atmosphere))
 
     kernel_parameters = interface_kernel_parameters(grid)
 
@@ -96,7 +96,7 @@ end
     #   ⋅ 𝒰 ≡ "dynamic" state vector (thermodynamics + reference height + velocity)
     ℂₐ = atmosphere_properties.thermodynamics_parameters
     𝒬ₐ = thermodynamic_atmospheric_state = AtmosphericThermodynamics.PhaseEquil_pTq(ℂₐ, pₐ, Tₐ, qₐ)
-    zₐ = atmosphere_properties.reference_height # elevation of atmos variables relative to interface
+    zₐ = atmosphere_properties.surface_layer_height # elevation of atmos variables relative to interface
 
     local_atmosphere_state = (z = zₐ,
                               u = uₐ,

src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl

@@ -39,7 +39,7 @@ function compute_atmosphere_sea_ice_fluxes!(coupled_model)
     ocean_properties = coupled_model.interfaces.ocean_properties
 
     atmosphere_properties = (thermodynamics_parameters = atmosphere.thermodynamics_parameters,
-                             reference_height = atmosphere.reference_height)
+                             surface_layer_height = atmosphere.surface_layer_height)
 
     kernel_parameters = interface_kernel_parameters(grid)
 
@@ -108,7 +108,7 @@ end
     #   ⋅ 𝒰 ≡ "dynamic" state vector (thermodynamics + reference height + velocity)
     ℂₐ = atmosphere_properties.thermodynamics_parameters
     𝒬ₐ = thermodynamic_atmospheric_state = AtmosphericThermodynamics.PhaseEquil_pTq(ℂₐ, pₐ, Tₐ, qₐ)
-    zₐ = atmosphere_properties.reference_height # elevation of atmos variables relative to interface
+    zₐ = atmosphere_properties.surface_layer_height # elevation of atmos variables relative to interface
 
     local_atmosphere_state = (z = zₐ,
                               u = uₐ,

src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl

@@ -11,6 +11,10 @@ using ..OceanSeaIceModels: reference_density,
                            freshwater_flux,
                            SeaIceSimulation
 
+using ..OceanSeaIceModels.PrescribedAtmospheres:
+    PrescribedAtmosphere,
+    thermodynamics_parameters
+
 using ClimaSeaIce: SeaIceModel
 
 using Oceananigans: HydrostaticFreeSurfaceModel, architecture
@@ -77,9 +81,12 @@ function StateExchanger(ocean::Simulation, atmosphere)
                                  Qℓ = Field{Center, Center, Nothing}(exchange_grid),
                                  Mp = Field{Center, Center, Nothing}(exchange_grid))
 
-    # TODO: use kernel_parameters not hard code?
-    # kernel_parameters = interface_kernel_parameters(ocean_grid)
-    
+    exchanger = atmosphere_exchanger(atmosphere, exchange_grid)
+
+    return StateExchanger(ocean.model.grid, exchange_atmosphere_state, exchanger)
+end
+
+function atmosphere_exchanger(atmosphere::PrescribedAtmosphere, exchange_grid)
     atmos_grid = atmosphere.grid
     arch = architecture(exchange_grid)
     Nx, Ny, Nz = size(exchange_grid)
@@ -97,7 +104,7 @@ function StateExchanger(ocean::Simulation, atmosphere)
     launch!(arch, exchange_grid, kernel_parameters,
             _compute_fractional_indices!, frac_indices, exchange_grid, atmos_grid)
 
-    return StateExchanger(ocean.model.grid, exchange_atmosphere_state, frac_indices)
+    return frac_indices
 end
 
 @kernel function _compute_fractional_indices!(indices_tuple, exchange_grid, atmos_grid)
@@ -263,7 +270,7 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing;
     sea_ice_heat_capacity     = convert(FT, sea_ice_heat_capacity)
     freshwater_density        = convert(FT, freshwater_density)
 
-    atmosphere_properties = atmosphere.thermodynamics_parameters
+    atmosphere_properties = thermodynamics_parameters(atmosphere)
 
     ocean_properties = (reference_density  = ocean_reference_density,
                         heat_capacity      = ocean_heat_capacity,

src/OceanSeaIceModels/InterfaceComputations/interpolate_atmospheric_state.jl

@@ -5,10 +5,13 @@ using Oceananigans.OutputReaders: TimeInterpolator
 
 using ...OceanSimulations: forcing_barotropic_potential
 
+using ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: PrescribedAtmosphere
+import ClimaOcean.OceanSeaIceModels: interpolate_atmosphere_state!
+
+# TODO: move to PrescribedAtmospheres
 """Interpolate the atmospheric state onto the ocean / sea-ice grid."""
-function interpolate_atmospheric_state!(coupled_model)
+function interpolate_atmosphere_state!(interfaces, atmosphere::PrescribedAtmosphere, coupled_model)
     ocean = coupled_model.ocean
-    atmosphere = coupled_model.atmosphere
     atmosphere_grid = atmosphere.grid
 
     # Basic model properties

src/OceanSeaIceModels/OceanSeaIceModels.jl

@@ -46,6 +46,11 @@ sea_ice_concentration(sea_ice::SeaIceSimulation) = sea_ice.model.ice_concentrati
 ##### Some implementation
 #####
 
+# Atmosphere interface
+interpolate_atmosphere_state!(interfaces, atmosphere, coupled_model) = nothing
+compute_net_atmosphere_fluxes!(coupled_model) = nothing
+
+# TODO: import this last
 include("PrescribedAtmospheres.jl")
 
 using .PrescribedAtmospheres:
@@ -73,7 +78,6 @@ compute_atmosphere_ocean_fluxes!(::NoAtmosphereModel) = nothing
 compute_atmosphere_sea_ice_fluxes!(::NoAtmosphereModel) = nothing
 
 const PrescribedAtmosphereModel = OceanSeaIceModel{<:Any, <:PrescribedAtmosphere}
-
 compute_net_atmosphere_fluxes!(::PrescribedAtmosphereModel) = nothing
 
 # "No sea ice" implementation

src/OceanSeaIceModels/PrescribedAtmospheres.jl

@@ -4,7 +4,7 @@ using Oceananigans.Grids: grid_name
 using Oceananigans.Utils: prettysummary, Time
 using Oceananigans.Fields: Center
 using Oceananigans.OutputReaders: FieldTimeSeries, update_field_time_series!, extract_field_time_series
-using Oceananigans.TimeSteppers: tick!, Clock
+using Oceananigans.TimeSteppers: Clock, tick!
 
 using Adapt
 using Thermodynamics.Parameters: AbstractThermodynamicsParameters
@@ -289,7 +289,7 @@ const PATP = PrescribedAtmosphereThermodynamicsParameters
 ##### Prescribed atmosphere (as opposed to dynamically evolving / prognostic)
 #####
 
-struct PrescribedAtmosphere{FT, M, G, T, U, P, C, F, I, R, TP, TI}
+mutable struct PrescribedAtmosphere{FT, M, G, T, U, P, C, F, I, R, TP, TI}
     grid :: G
     clock :: Clock{T}
     metadata :: M
@@ -301,7 +301,7 @@ struct PrescribedAtmosphere{FT, M, G, T, U, P, C, F, I, R, TP, TI}
     downwelling_radiation :: R
     thermodynamics_parameters :: TP
     times :: TI
-    reference_height :: FT
+    surface_layer_height :: FT
     boundary_layer_height :: FT
 end
 
@@ -315,7 +315,7 @@ end
 function Base.show(io::IO, pa::PrescribedAtmosphere)
     print(io, summary(pa), " on ", grid_name(pa.grid), ":", '\n')
     print(io, "├── times: ", prettysummary(pa.times), '\n')
-    print(io, "├── reference_height: ", prettysummary(pa.reference_height), '\n')
+    print(io, "├── surface_layer_height: ", prettysummary(pa.surface_layer_height), '\n')
     print(io, "└── boundary_layer_height: ", prettysummary(pa.boundary_layer_height))
 end
 
@@ -350,9 +350,9 @@ function default_atmosphere_pressure(grid, times)
     return pa
 end
 
-@inline function time_step!(atmos::PrescribedAtmosphere, Δt) 
+@inline function time_step!(atmos::PrescribedAtmosphere, Δt)
     tick!(atmos.clock, Δt)
-    
+
     time = Time(atmos.clock.time)
     ftses = extract_field_time_series(atmos)
 
@@ -364,14 +364,14 @@ end
 end
 
 @inline thermodynamics_parameters(atmos::PrescribedAtmosphere) = atmos.thermodynamics_parameters
-@inline reference_height(atmos::PrescribedAtmosphere) = atmos.reference_height
+@inline surface_layer_height(atmos::PrescribedAtmosphere) = atmos.surface_layer_height
 @inline boundary_layer_height(atmos::PrescribedAtmosphere) = atmos.boundary_layer_height    
 
 """
     PrescribedAtmosphere(grid, times;
                          clock = Clock{Float64}(time = 0),
                          metadata = nothing,
-                         reference_height = 10, # meters
+                         surface_layer_height = 10, # meters
                          boundary_layer_height = 600 # meters,
                          thermodynamics_parameters = PrescribedAtmosphereThermodynamicsParameters(FT),
                          auxiliary_freshwater_flux = nothing,
@@ -387,7 +387,7 @@ state with data given at `times`.
 function PrescribedAtmosphere(grid, times;
                               clock = Clock{Float64}(time = 0),
                               metadata = nothing,  
-                              reference_height = convert(eltype(grid), 10),
+                              surface_layer_height = convert(eltype(grid), 10),
                               boundary_layer_height = convert(eltype(grid), 600),
                               thermodynamics_parameters = nothing,
                               auxiliary_freshwater_flux = nothing,
@@ -413,7 +413,7 @@ function PrescribedAtmosphere(grid, times;
                                 downwelling_radiation,
                                 thermodynamics_parameters,
                                 times,
-                                convert(FT, reference_height),
+                                convert(FT, surface_layer_height),
                                 convert(FT, boundary_layer_height))
 end
 

src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl

@@ -3,7 +3,7 @@ using .InterfaceComputations:
     compute_sea_ice_ocean_fluxes!,
     compute_net_ocean_fluxes!,
     compute_net_sea_ice_fluxes!,
-    interpolate_atmospheric_state!
+    interpolate_atmosphere_state!
 
 using ClimaSeaIce: SeaIceModel, SeaIceThermodynamics
 using Oceananigans.Grids: φnode
@@ -14,6 +14,7 @@ function time_step!(coupled_model::OceanSeaIceModel, Δt; callbacks=[], compute_
     ocean = coupled_model.ocean
     sea_ice = coupled_model.sea_ice
     atmosphere = coupled_model.atmosphere
+    clock = coupled_model.clock
 
     # Be paranoid and update state at iteration 0
     coupled_model.clock.iteration == 0 && update_state!(coupled_model, callbacks)
@@ -31,14 +32,12 @@ function time_step!(coupled_model::OceanSeaIceModel, Δt; callbacks=[], compute_
             parent(h⁻) .= parent(hⁿ)
         end
 
-        sea_ice.Δt = Δt
-        time_step!(sea_ice)
+        time_step!(sea_ice, Δt)
     end
 
     # TODO after ice time-step:
     #  - Adjust ocean heat flux if the ice completely melts?
-    ocean.Δt = Δt
-    time_step!(ocean)
+    time_step!(ocean, Δt)
 
     # Time step the atmosphere
     time_step!(atmosphere, Δt)
@@ -55,7 +54,7 @@ end
 function update_state!(coupled_model::OceanSeaIceModel, callbacks=[]; compute_tendencies=true)
 
     # This function needs to be specialized to allow different atmospheric models
-    interpolate_atmospheric_state!(coupled_model) 
+    interpolate_atmosphere_state!(coupled_model.interfaces, coupled_model.atmosphere, coupled_model) 
 
     # Compute interface states
     compute_atmosphere_ocean_fluxes!(coupled_model)

src/OceanSimulations/ocean_simulation.jl

@@ -105,6 +105,7 @@ function ocean_simulation(grid;
                           boundary_conditions::NamedTuple = NamedTuple(),
                           tracer_advection = default_tracer_advection(),
                           vertical_coordinate = default_vertical_coordinate(grid),
+                          warn = true,
                           verbose = false)
 
     FT = eltype(grid)
@@ -131,7 +132,7 @@ function ocean_simulation(grid;
         u_immersed_bc = DefaultBoundaryCondition()
         v_immersed_bc = DefaultBoundaryCondition()
     else
-        if !(grid isa ImmersedBoundaryGrid)
+        if warn && !(grid isa ImmersedBoundaryGrid)
             msg = """Are you totally, 100% sure that you want to build a simulation on
 
                    $(summary(grid))

test/test_surface_fluxes.jl

@@ -45,7 +45,7 @@ end
         atmosphere = JRA55PrescribedAtmosphere(1:2; grid, architecture=arch, backend=InMemory()) 
 
         CUDA.@allowscalar begin
-            h  = atmosphere.reference_height
+            h  = atmosphere.surface_layer_height
             pₐ = atmosphere.pressure[1][1, 1, 1]
 
             Tₐ = atmosphere.tracers.T[1][1, 1, 1]