Diffusive baroclinic wave with thermal slab

.gitignore

@@ -37,4 +37,7 @@ docs/src/generated/
 *.jld2
 
 # Julia System Images
-*.so
+*.so
+
+# internal tests
+testdel.jl

experiments/ClimaCore/bc-wave-slab/Project.toml

@@ -0,0 +1,21 @@
+[deps]
+CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
+ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
+ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
+ClimaCorePlots = "cf7c7e5a-b407-4c48-9047-11a94a308626"
+ClimaCoreTempestRemap = "d934ef94-cdd4-4710-83d6-720549644b70"
+IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
+UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
+
+[extras]
+CPUSummary = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9"

experiments/ClimaCore/bc-wave-slab/README.md

@@ -0,0 +1,101 @@
+# **Baroclinic Wave + Slab**
+
+# Atmosphere
+The momentum equations are in the advective form, and tracers in the consevative form, namely:
+
+- Density:
+$$ \frac{\partial \rho}{\partial t} + \nabla \cdot ({\rho \vec{u}})= 0 $$
+
+- Momentum (flux form):
+$$ \frac{\partial \vec{u_h}}{\partial t} + \vec{u} \cdot \nabla \vec{u_h} = - \frac{1}{\rho}\nabla_h p
++ \frac{\partial}{\partial z} K_v \frac{\partial}{\partial z} \vec{u_h}
+$$
+$$ \frac{\partial w}{\partial t} + \vec{u} \cdot \nabla w= 
+- \frac{1}{\rho}\frac{\partial p}{\partial z}
+- \nabla_z \Phi 
++ \frac{\partial}{\partial z} K_v \frac{\partial}{\partial z} w 
+$$
+
+- Total energy:
+$$ \frac{\partial \rho e_{tot}}{\partial t} + \nabla \cdot (\rho h_{tot} \vec{u}) = \frac{\partial}{\partial z} K_v \frac{\partial}{\partial z} h_{tot}
+$$
+
+where the total specific enthalpy and total specific  energy are
+$$ h_{tot} =  e_{tot} + \frac{p}{\rho}  \,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\, e_{tot} = c_v T + \Phi + \frac{1}{2}\vec{u}^2 
+$$
+(note that $h_{tot} \neq h = c_vT + p/\rho = c_p T$, the specific enthalpy in the thermodynamic sense), $\Phi = gz$ is the geopotential,
+$u_h$ is the horizontal velocity vector, $w$ the vertical velocity, $\rho$ the density, $p$ pressure, $K_v$ the vertical diffusivity (assumed constant here). 
+
+## Boundary conditions (BCs)
+- We implement BCs similarly to other climate models. 
+    - First-order fluxes (i.e., advective fluxes) are always set to zero, corresponding to the *free-slip* and *impenetrable* BC, where:
+    $$
+    w = 0 \,\,\,\,\,\,\, \partial_t w = 0 \,\,\,\,\,\,\, \nabla \times\vec{u_h}=0 \,\,\,\,\,\,\, \nabla \cdot \vec{\rho u_h}=0  \,\,\,\,\,\,\, \nabla \cdot \rho h_{tot} \vec{u_h}=0
+    $$
+    - Second-order fluxes (i.e., diffusive fluxes)
+        - `NoFlux()`: By default we have *impenetrable* or *insulating* BCs (no second-order fluxes) at all boundaries. 
+        - `BulkFormula()`: Applied to tracers (e.g., temperature and moisture), this imposes a boundary fluxes (e.g., sensible and latent heat) calculated using the bulk aerodynamic formulae  using prescribed surface values of ($T_{sfc}$ and $q_{sfc}^{sat}$).  At the surface, the bulk sensible heat flux formula for total enthalpy essentially replaces the above: 
+            $$ (K_v \rho \partial_z h_{tot})_{sfc}$$
+            For **total energy**, we have two choices:
+            - 1. enthalpy flux: 
+                $$ (K_v \rho \partial_z h_{tot})_{sfc} \rightarrow 
+                \hat{n} \cdot  \rho C_H ||u||^{1} (h^1- h_{sfc})   
+                = F_S
+                $$
+            - 2. sensible (and latent) heat flux. The sensible heat flux is: 
+                $$ (K_v \rho \partial_z h_{tot})_{sfc} \rightarrow 
+                \hat{n} \cdot C_H c_{pd} ρ^{1} ||u||^{1} (T^{1} - T_{sfc}) 
+                +  \hat{n} \cdot C_H ρ^{1} ||u||^{1} (\Phi^{1} - \Phi_{sfc})      
+                = F_S
+                $$
+                where $^{1}$ corresponds to the lowest model level, $C_H$ is the dimensionless thermal transfer coefficient, $c_{pd}$ is the specific heat capacity for dry air $||u||$ the wind speed. This is the *bulk turbulent sensible heat flux* parameterization, and $F_S$ is positive when atmosphere receives energy from the surface. 
+                The contribution of the kinetic energy is usually O(1e4) smaller and is neglected, but it can be added to F_S as:
+                $$
+                F_{S_{tot}} = F_S + \hat{n} \cdot C_D ρ^{1} ||u||^{1} (\vec{u_h}^{1})^2 
+                $$        
+         - ` DragLaw()`: essentially the bulk formula for momentum       
+            $$ \frac{\partial}{\partial z} K_v \frac{\partial}{\partial z} \vec{u_h} \rightarrow 
+            \hat{n} \cdot C_D ρ^{1} ||u||^{1} \vec{u_h}^{1}          
+            = F_M
+            $$
+
+
+        - `BulkFormulaCoupled()`: same as `BulkFormula()`, but $T_{sfc}$ is passed from the neighboring model.
+
+    - The diffusive fluxes are applied via the `vertical_diffusion` ClimaAtmos model sub-component. To apply boundary fluxes without diffusion in the atmospheric interior, the viscosity coefficient needs to be set to zero: $ν = FT(0)$
+- Current setup
+    - total energy and momentum:
+        - at z=0, we use `BulkFormulaCoupled()` with the enthalpy flux formulation, and the `DragLaw()`, with $C_D = C_H = 0.001$
+        - interior diffusivity is set to $\nu = 5$ m^2/s
+        - the the top $F_S = F_M = 0$
+    - All other boundary fluxes are set to 0.
+
+## Initial conditions
+- we initialize with a perturbation in a balanced background state, as in:
+https://climate.ucdavis.edu/pubs/UMJS2013QJRMS.pdf
+
+# Heat Slab
+The slab solves for temperature in a single layer, whose tendency is the accumulated fluxes divided by the coupling timestep plus a parameterisation of the internal processes, $G$.
+$$
+\rho c h_s  \, \partial_t T_{sfc} =  - F_{integ}  / \Delta t_{coupler}  
+$$
+
+# Conservation checks
+- this uses the `sum` of `ClimaCore/Fields/mapreduce.jl`, which produces a sum weighted by the area Jacobian. 
+
+# NB:
+- first coupled iteration does not call rhs!s
+- slab `T_sfc` gets huge numbers when using `SSPRK33`. ok with `Euler`
+
+# TODO
+- calculate fluxes at cell faces (not centers) - interp
+- need better CA interface for specifying sensible / latent heat fluxes - i.e. fluxes in terms of temp and q, not enthaplpy fluxes
+- add more precise flux accumulation
+- conservation tests - add error threshold and exception, interval, show option, and make a general interface for it
+
+# References
+- [Kang et al 2021](https://arxiv.org/abs/2101.09263)
+
+
+
+

experiments/ClimaCore/bc-wave-slab/atmos.jl

@@ -0,0 +1,74 @@
+# from ClimaAtmos:test/test_cases/run_3d_baroclinic_wave.jl, removing `step!` 
+
+using OrdinaryDiffEq: SSPRK33
+using ClimaCorePlots, Plots
+using UnPack
+
+import ClimaCore
+import ClimaCore: Fields, Geometry, Operators
+import ClimaCore.Geometry: ⊗
+
+using ClimaAtmos.Utils.InitialConditions: init_3d_baroclinic_wave
+using ClimaAtmos.Domains
+using ClimaAtmos.BoundaryConditions
+using ClimaAtmos.Models: ConstantViscosity
+using ClimaAtmos.Models.Nonhydrostatic3DModels
+using ClimaAtmos.Simulations
+
+using CLIMAParameters
+struct DryBaroclinicWaveParameters <: CLIMAParameters.AbstractEarthParameterSet end
+
+# initiate the lower Atmos boundary Field
+function atmos_bcfield_init(domain)
+    center_space, face_space = make_function_space(domain)
+    ClimaCore.Fields.level(ClimaCore.Fields.zeros(center_space), 1)
+end
+
+function atmos_bcvectorfield_init(domain)
+    center_space, face_space = make_function_space(domain)
+    zero_field = ClimaCore.Fields.zeros(center_space)
+    uv_local = Geometry.UVVector.(zero_field, zero_field)
+    ClimaCore.Fields.level(Geometry.Covariant3Vector.(zero_field) .⊗ Geometry.Covariant12Vector.(uv_local), 1)
+end
+
+function atmos_init(::Type{FT}, tspan; stepper = SSPRK33(), nelements = (6, 10), npolynomial = 4, dt = 0.02) where {FT}
+    params = DryBaroclinicWaveParameters()
+
+    domain = SphericalShell(
+        FT,
+        radius = CLIMAParameters.Planet.planet_radius(params),
+        height = FT(30.0e3),
+        nelements = nelements,
+        npolynomial = npolynomial,
+    )
+
+    # initiate boundary conditions
+    boundary_conditions = (;
+        ρe_tot = (
+            top = NoFlux(),
+            bottom = CouplerEnergyFlux(
+                FT(0.001), # transfer coefficient, default: 1e-3 [dimenisonless]
+                atmos_bcfield_init(domain),
+            ), # surface specific enthalpy ~ (T_sfc * c_p)
+        ),
+        uh = (top = NoVectorFlux(), bottom = CouplerVectorFlux(FT(0.001), atmos_bcvectorfield_init(domain))),
+    )
+
+    model = Nonhydrostatic3DModel(
+        domain = domain,
+        boundary_conditions = boundary_conditions,
+        parameters = params,
+        hyperdiffusivity = FT(1e16),
+        vertical_diffusion = ConstantViscosity(ν = FT(5)), #default: 5 m^2/s
+    )
+
+    # execute (using the regression test from ClimaAtmos)
+    simulation = Simulation(model, stepper, dt = dt, tspan = tspan)
+
+    # test set function
+    @unpack ρ, uh, w, ρe_tot = init_3d_baroclinic_wave(FT, params)
+    set!(simulation, :base, ρ = ρ, uh = uh, w = w)
+    set!(simulation, :thermodynamics, ρe_tot = ρe_tot)
+
+    simulation
+end

experiments/ClimaCore/bc-wave-slab/conservation.jl

@@ -0,0 +1,28 @@
+# generalise this into coupler-specific function
+
+struct ConservationCheck{A}
+    ρe_tot_atmos::A
+    ρe_tot_slab::A
+end
+
+function check_conservation(cs, atmos_sim, slab_sim)
+    atmos_field = atmos_sim.integrator.u.thermodynamics.ρe_tot
+    slab_field = get_slab_energy(slab_sim)
+
+    ρe_tot_atmos = sum(atmos_field)
+    ρe_tot_slab = sum(slab_field)
+
+    push!(cs.ρe_tot_atmos, ρe_tot_atmos)
+    push!(cs.ρe_tot_slab, ρe_tot_slab)
+end
+
+function get_slab_energy(slab_sim)
+    ρe_tot = slab_sim.params.ρ .* slab_sim.params.c .* slab_sim.integrator.u.T_sfc .* slab_sim.params.h
+end
+
+# struct ConservationCheck{F, E, I, S}
+#     fields::F
+#     exception::E
+#     interval::I
+#     show::S
+# end

experiments/ClimaCore/bc-wave-slab/coupled_bc.jl

@@ -0,0 +1,156 @@
+coupler_atmos_boundary_flux(_, ..) = nothing
+
+# momentum
+struct CouplerVectorFlux{C, F} <: AbstractBoundary
+    coefficients::C
+    flux::F
+end
+function get_boundary_flux(model, bc::CouplerVectorFlux, var::Fields.Field, Y, Ya) # TODO: adapt for Field BC (CC#325)
+    flux = Geometry.Covariant3Vector(FT(1)) ⊗ Geometry.Covariant12Vector(FT(1), FT(1)) * parent(bc.flux)[1]
+end
+"""
+get_boundary_flux(
+    model,
+    bc::CouplerVectorFlux,
+    ρc::Fields.Field,
+    Y,
+    Ya,
+)
+
+This is an extenison of ClimaAtmos.get_boundary_flux for momentum
+"""
+function coupler_atmos_boundary_flux(bc::CouplerVectorFlux, atmos_sim, slab_sim, F_a_space, F_s_space)
+
+    Y_atm = atmos_sim.integrator.u
+    FT = eltype(Y_atm)
+
+    uv = Y_atm.base.uh
+
+    uh1_cov = Fields.level(uv.components.data.:1, 1) # TODO: change to Field - pending ClimaCore PR #325 
+    uh2_cov = Fields.level(uv.components.data.:2, 1) # TODO: change to Field - pending ClimaCore PR #325 
+
+    # physical scale (wind * coeff)
+    uv_1 = Fields.level(Geometry.UVVector.(uv), 1)
+    u_wind = norm(uv_1) # TODO: this will need to be spatially variable
+
+    # unit vector in W direction on central space
+    normal_unit_vector = Fields.level(similar(Geometry.WVector.(Y_atm.thermodynamics.ρe_tot)), 1) # TODO: for topography, needs to be perp to surface
+    parent(normal_unit_vector) .= FT(1)
+    local_geometry = Fields.local_geometry_field(axes(normal_unit_vector))
+
+    # contravariant scale
+    scale_con =
+        Geometry.Contravariant3Vector.(
+            Geometry.WVector.(normal_unit_vector .* bc.coefficients .* u_wind),
+            local_geometry,
+        )
+
+    parent(bc.flux) .=
+        parent(Geometry.Contravariant3Vector.(scale_con) .⊗ Geometry.Covariant12Vector.(uh1_cov, uh2_cov))
+end
+
+# energy
+struct CouplerEnergyFlux{C, F} <: AbstractBoundary
+    coefficients::C
+    flux::F
+end
+
+"""
+get_boundary_flux(
+    model,
+    bc::CouplerEnergyFlux,
+    ρc::Fields.Field,
+    Y,
+    Ya,
+)
+
+This is an extenison of ClimaAtmos.get_boundary_flux for enthalpy
+"""
+function get_boundary_flux(model, bc::CouplerEnergyFlux, ρc::ClimaCore.Fields.Field, Y, Ya)
+    bulk_flux = bc.flux
+    ClimaCore.Geometry.WVector.(parent(bulk_flux)[1]) #flux = Geometry.Contravariant3Vector.(bulk_flux) # need CC extension - waiting for PR (https://github.com/CliMA/ClimaCore.jl/pull/325)
+end
+
+"""
+coupler_boundary_flux(
+    model,
+    bc::CouplerEnergyFlux,
+    ρc::Fields.Field,
+    Y,
+    Ya,
+)
+
+Vertical fluxes for arbitrary variables with the bulk aerodynamic turbulent formula,  
+given variable exchange coefficients. (e.g. for energy, or tracers)
+Currently this is a simplified version of the bulk formula for testing. 
+"""
+function coupler_atmos_boundary_flux(bc::CouplerEnergyFlux, atmos_sim, slab_sim, F_a_space, F_s_space)
+
+    Y_atm = atmos_sim.integrator.u
+    FT = eltype(Y_atm)
+
+    T_1 = ClimaCore.Fields.level(calculate_temperature(atmos_sim, FT), 1)
+
+    ρ_1 = ClimaCore.Fields.level(Y_atm.base.ρ, 1)
+
+    # remap slab > atmos
+    T_sfc_atmos_grid = ClimaCore.Fields.zeros(F_a_space)
+    T_sfc_slab_grid = copy(slab_sim.integrator.u.T_sfc)
+    # ClimaCoreTempestRemap.remap!(T_sfc_atmos_grid, T_sfc_slab_grid, R_slab2atm)
+    dummmy_remap!(T_sfc_slab_grid, T_sfc_atmos_grid)
+
+    # save in atmos BCs
+    c_pd = CLIMAParameters.Planet.cp_d(atmos_sim.model.parameters)
+    # parent(bc.flux) .= bc.coefficients .* c_pd .* parent(ρ_1) .* (parent(T_1) .- parent(T_sfc_atmos_grid)) # parent(bc.flux) .= bc.coefficients .* parent(ρ_1 .* u_wind) .* (parent(c_1) .- parent(c_sfc)) # TODO: make neater - same space but different instance error otherwise
+    parent(bc.flux) .= FT(10)
+    # TODO: make neater - same space but different instance error otherwise
+    nothing
+end
+
+# code below should be moved to CA / thermodyn 
+using LinearAlgebra: norm_sqr
+using Thermodynamics
+
+function calculate_temperature(atmos_sim, FT)
+    model = atmos_sim.model
+    Y = atmos_sim.integrator.u
+    p = calculate_pressure(Y, Y, model.thermodynamics, model.moisture, model.moisture, model.parameters, FT)
+    R = FT(287)
+    T = p ./ R ./ Y.base.ρ
+end
+
+@inline function calculate_gravitational_potential(Y, Ya, params, FT)
+    g::FT = CLIMAParameters.Planet.grav(params)
+    ρ = Y.base.ρ
+    z = Fields.coordinate_field(axes(ρ)).z
+
+    return @. g * z
+end
+
+@inline function calculate_pressure(Y, Ya, i, ii, iii, params, FT)
+    ρ = Y.base.ρ
+    uh = Y.base.uh
+    w = Y.base.w
+    ρe_tot = Y.thermodynamics.ρe_tot
+
+    interp_f2c = Operators.InterpolateF2C()
+
+    z = Fields.coordinate_field(axes(ρ)).z
+    K = calculate_kinetic_energy(Y, Y, params, FT)
+    Φ = calculate_gravitational_potential(Y, Y, params, FT)
+
+    e_int = @. ρe_tot / ρ - Φ - K
+    p = Thermodynamics.air_pressure.(Thermodynamics.PhaseDry.(params, e_int, ρ))
+
+    return p
+end
+
+@inline function calculate_kinetic_energy(Y, Ya, params, FT)
+    cuₕ = Y.base.uh # Covariant12Vector on centers
+    fw = Y.base.w # Covariant3Vector on faces
+    If2c = Operators.InterpolateF2C()
+
+    cuvw = Geometry.Covariant123Vector.(cuₕ) .+ Geometry.Covariant123Vector.(If2c.(fw))
+    cK = @. norm_sqr(cuvw) / 2
+    return cK
+end