Merge #833

833: IGW: generalize example w/ non-uniform grids and update hydrostatic balance computation r=valeriabarra a=valeriabarra This PR will close #831 Co-authored-by: Dennis Yatunin <[email protected]> - [x] Code follows the [style guidelines](https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/) OR N/A. - [x] Unit tests are included OR N/A. - [x] Code is exercised in an integration test OR N/A. - [x] Documentation has been added/updated OR N/A. Co-authored-by: Valeria Barra <[email protected]>
CliMA · Jul 31, 2022 · cb952ba · cb952ba
2 parents e32396a + 0b7bb36
commit cb952ba
Show file tree

Hide file tree

Showing 6 changed files with 126 additions and 57 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -871,6 +871,20 @@ steps:
       queue: central
       slurm_ntasks: 1
 
+  - label: ":computer: stretched 2D plane inertial gravity wave"
+    key: "cpu_stretch_inertial_gravity_wave"
+    command:
+      - "julia --color=yes --project=examples examples/hybrid/driver.jl"
+    artifact_paths:
+      - "examples/hybrid/plane/output/stretched_inertial_gravity_wave/Float32/*"
+    env:
+      TEST_NAME: "plane/inertial_gravity_wave"
+      Z_STRETCH: "true"
+    agents:
+      config: cpu
+      queue: central
+      slurm_ntasks: 1
+
   - label: ":rocket::computer: Allocations analysis"
     key: "cpu_allocations"
     command:

diff --git a/examples/common_spaces.jl b/examples/common_spaces.jl
@@ -53,13 +53,13 @@ function make_distributed_horizontal_space(mesh, npoly, comms_ctx)
     return space, comms_ctx
 end
 
-function make_hybrid_spaces(h_space, z_max, z_elem)
+function make_hybrid_spaces(h_space, z_max, z_elem; z_stretch)
     z_domain = Domains.IntervalDomain(
         Geometry.ZPoint(zero(z_max)),
         Geometry.ZPoint(z_max);
         boundary_tags = (:bottom, :top),
     )
-    z_mesh = Meshes.IntervalMesh(z_domain, nelems = z_elem)
+    z_mesh = Meshes.IntervalMesh(z_domain, z_stretch; nelems = z_elem)
     z_topology = Topologies.IntervalTopology(z_mesh)
     z_space = Spaces.CenterFiniteDifferenceSpace(z_topology)
     center_space = Spaces.ExtrudedFiniteDifferenceSpace(h_space, z_space)

diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl
@@ -57,13 +57,26 @@ include("../implicit_solver_debugging_tools.jl")
 include("../ordinary_diff_eq_bug_fixes.jl")
 include("../common_spaces.jl")
 
+if get(ENV, "Z_STRETCH", "false") == "true"
+    z_stretch_scale = FT(7e3)
+    z_stretch = Meshes.ExponentialStretching(z_stretch_scale)
+    z_stretch_string = "stretched"
+else
+    z_stretch = Meshes.Uniform()
+    z_stretch_string = "uniform"
+end
+
 if haskey(ENV, "TEST_NAME")
     test_dir, test_file_name = split(ENV["TEST_NAME"], '/')
 else
     error("ENV[\"TEST_NAME\"] required (e.g., \"sphere/baroclinic_wave_rhoe\")")
 end
 include(joinpath(test_dir, "$test_file_name.jl"))
 
+if z_stretch_string == "stretched"
+    test_file_name = "$(z_stretch_string)_$(test_file_name)"
+end
+
 import ClimaCore: enable_threading
 enable_threading() = false
 
@@ -88,12 +101,13 @@ else
         h_space = make_horizontal_space(horizontal_mesh, npoly)
         comms_ctx = nothing
     end
-    center_space, face_space = make_hybrid_spaces(h_space, z_max, z_elem)
+    center_space, face_space =
+        make_hybrid_spaces(h_space, z_max, z_elem; z_stretch)
     ᶜlocal_geometry = Fields.local_geometry_field(center_space)
     ᶠlocal_geometry = Fields.local_geometry_field(face_space)
     Y = Fields.FieldVector(
-        c = center_initial_condition.(ᶜlocal_geometry),
-        f = face_initial_condition.(ᶠlocal_geometry),
+        c = center_initial_condition(ᶜlocal_geometry),
+        f = face_initial_condition(ᶠlocal_geometry),
     )
 end
 p = get_cache(ᶜlocal_geometry, ᶠlocal_geometry, Y, dt, upwinding_mode)
@@ -187,14 +201,15 @@ if haskey(ENV, "CI_PERF_SKIP_RUN") # for performance analysis
 end
 
 @info "Running `$test_dir/$test_file_name` test case"
+@info "on a vertical $z_stretch_string grid"
 
 walltime = @elapsed sol = OrdinaryDiffEq.solve!(integrator)
 
 if is_distributed # replace sol.u on the root processor with the global sol.u
     if ClimaComms.iamroot(comms_ctx)
         global_h_space = make_horizontal_space(horizontal_mesh, npoly)
         global_center_space, global_face_space =
-            make_hybrid_spaces(global_h_space, z_max, z_elem)
+            make_hybrid_spaces(global_h_space, z_max, z_elem; z_stretch)
         global_Y_c_type = Fields.Field{
             typeof(Fields.field_values(Y.c)),
             typeof(global_center_space),

diff --git a/examples/hybrid/plane/inertial_gravity_wave.jl b/examples/hybrid/plane/inertial_gravity_wave.jl
@@ -20,6 +20,7 @@ include("../staggered_nonhydrostatic_model.jl")
 
 # Additional constants required for inertial gravity wave initial condition
 z_max = FT(10e3)
+z_stretch_scale = FT(7e3)
 const x_max = is_small_scale ? FT(300e3) : FT(6000e3)
 const x_mid = is_small_scale ? FT(100e3) : FT(3000e3)
 const d = is_small_scale ? FT(5e3) : FT(100e3)
@@ -29,7 +30,7 @@ const T₀ = FT(250)
 const ΔT = FT(0.01)
 
 # Additional values required for driver
-upwinding_mode = :third_order
+upwinding_mode = :third_order # :none to switch to centered diff
 
 # Other convenient constants used in reference paper
 const δ = grav / (R_d * T₀)        # Bretherton height parameter
@@ -39,18 +40,20 @@ const ρₛ = p_0 / (R_d * T₀)        # air density at surface
 # TODO: Loop over all domain setups used in reference paper
 const Δx = is_small_scale ? FT(1e3) : FT(20e3)
 const Δz = is_small_scale ? Δx / 2 : Δx / 40
-z_elem = Int(z_max / Δz)
+z_elem = Int(z_max / Δz) # default 20 vertical elements
 npoly, x_elem = 1, Int(x_max / Δx) # max small-scale dt = 1.5
 # npoly, x_elem = 4, Int(x_max / (Δx * (4 + 1))) # max small-scale dt = 0.8
 
 # Animation-related values
 animation_duration = FT(5)
 fps = 2
 
+# Set up mesh
+horizontal_mesh = periodic_line_mesh(; x_max, x_elem = x_elem)
+
 # Additional values required for driver
-horizontal_mesh = periodic_line_mesh(; x_max, x_elem)
-t_end = is_small_scale ? FT(60 * 60 * 0.5) : FT(60 * 60 * 8)
 dt = is_small_scale ? FT(1.5) : FT(20)
+t_end = is_small_scale ? FT(60 * 60 * 0.5) : FT(60 * 60 * 8)
 dt_save_to_sol = t_end / (animation_duration * fps)
 ode_algorithm = OrdinaryDiffEq.Rosenbrock23
 jacobian_flags = (;
@@ -59,42 +62,69 @@ jacobian_flags = (;
 )
 show_progress_bar = true
 
-if is_discrete_hydrostatic_balance
+function discrete_hydrostatic_balance!(ᶠΔz, ᶜΔz, grav)
+
     # Yₜ.f.w = 0 in implicit tendency                                        ==>
     # -(ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) + ᶠgradᵥᶜΦ) = 0                             ==>
     # ᶠgradᵥ(ᶜp) = -grav * ᶠinterp(ᶜρ)                                       ==>
     # (p(z + Δz) - p(z)) / Δz = -grav * (ρ(z + Δz) + ρ(z)) / 2               ==>
     # p(z + Δz) + grav * Δz * ρ(z + Δz) / 2 = p(z) - grav * Δz * ρ(z) / 2    ==>
     # p(z + Δz) * (1 + δ * Δz / 2) = p(z) * (1 - δ * Δz / 2)                 ==>
     # p(z + Δz) / p(z) = (1 - δ * Δz / 2) / (1 + δ * Δz / 2)                 ==>
-    # p(z) = p(0) * ((1 - δ * Δz / 2) / (1 + δ * Δz / 2))^(z / Δz)
-    p₀(z) = p_0 * ((1 - δ * Δz / 2) / (1 + δ * Δz / 2))^(z / Δz)
-else
-    p₀(z) = p_0 * exp(-δ * z)
+    # p(z + Δz) = p(z) * (1 - δ * Δz / 2) / (1 + δ * Δz / 2)
+    ᶜp = similar(ᶜΔz)
+    ᶜp1 = Fields.level(ᶜp, 1)
+    ᶜΔz1 = Fields.level(ᶜΔz, 1)
+    @. ᶜp1 = p_0 * (1 - δ * ᶜΔz1 / 4) / (1 + δ * ᶜΔz1 / 4)
+    for i in 1:(Spaces.nlevels(axes(ᶜp)) - 1)
+        ᶜpi = parent(Fields.level(ᶜp, i))
+        ᶜpi1 = parent(Fields.level(ᶜp, i + 1))
+        ᶠΔzi1 = parent(Fields.level(ᶠΔz, Spaces.PlusHalf(i)))
+        @. ᶜpi1 = ᶜpi * (1 - δ * ᶠΔzi1 / 2) / (1 + δ * ᶠΔzi1 / 2)
+    end
+    return ᶜp
 end
-Tb_init(x, z) = ΔT * exp(-(x - x_mid)^2 / d^2) * sin(π * z / z_max::FT)
-T′_init(x, z) = Tb_init(x, z) * exp(δ * z / 2)
-
-function center_initial_condition(local_geometry)
-    (; x, z) = local_geometry.coordinates
-    p = p₀(z)
-    T = T₀ + T′_init(x, z)
-    ρ = p / (R_d * T)
-    uₕ_local = Geometry.UVVector(u₀, v₀)
-    uₕ = Geometry.Covariant12Vector(uₕ_local, local_geometry)
+Tb_init(x, z, ΔT, x_mid, d, z_max) =
+    ΔT * exp(-(x - x_mid)^2 / d^2) * sin(π * z / z_max)
+T′_init(x, z, ΔT, x_mid, d, z_max, δ) =
+    Tb_init(x, z, ΔT, x_mid, d, z_max) * exp(δ * z / 2)
+# Pressure definition, when not in discrete hydrostatic balance state
+p₀(z) = @. p_0 * exp(-δ * z)
+
+function center_initial_condition(ᶜlocal_geometry)
+    ᶜx = ᶜlocal_geometry.coordinates.x
+    ᶜz = ᶜlocal_geometry.coordinates.z
+    # Correct pressure and density if in hydrostatic balance state
+    if is_discrete_hydrostatic_balance
+        face_space =
+            Spaces.FaceExtrudedFiniteDifferenceSpace(axes(ᶜlocal_geometry))
+        ᶠΔz = Fields.local_geometry_field(face_space).∂x∂ξ.components.data.:4
+        ᶜΔz = ᶜlocal_geometry.∂x∂ξ.components.data.:4
+        ᶜp = discrete_hydrostatic_balance!(ᶠΔz, ᶜΔz, grav)
+    else
+        ᶜp = @. p₀(ᶜz)
+    end
+    T = @. T₀ + T′_init(ᶜx, ᶜz, ΔT, x_mid, d, z_max, δ)
+    ᶜρ = @. ᶜp / (R_d * T)
+    ᶜuₕ_local = @. Geometry.UVVector(u₀ * one(ᶜz), v₀ * one(ᶜz))
+    ᶜuₕ = @. Geometry.Covariant12Vector(ᶜuₕ_local)
     if ᶜ𝔼_name == :ρθ
-        ρθ = ρ * T * (p_0 / p)^(R_d / cp_d)
-        return (; ρ, ρθ, uₕ)
+        ᶜρθ = @. ᶜρ * T * (p_0 / ᶜp)^(R_d / cp_d)
+        return NamedTuple{(:ρ, :ρθ, :uₕ)}.(tuple.(ᶜρ, ᶜρθ, uₕ))
     elseif ᶜ𝔼_name == :ρe
-        ρe = ρ * (cv_d * (T - T_tri) + norm_sqr(uₕ_local) / 2 + grav * z)
-        return (; ρ, ρe, uₕ)
+        ᶜρe = @. ᶜρ * (cv_d * (T - T_tri) + norm_sqr(ᶜuₕ_local) / 2 + grav * ᶜz)
+        return NamedTuple{(:ρ, :ρe, :uₕ)}.(tuple.(ᶜρ, ᶜρe, ᶜuₕ))
     elseif ᶜ𝔼_name == :ρe_int
-        ρe_int = ρ * cv_d * (T - T_tri)
-        return (; ρ, ρe_int, uₕ)
+        ᶜρe_int = @. ᶜρ * cv_d * (T - T_tri)
+        return NamedTuple{(:ρ, :ρe_int, :uₕ)}.(tuple.(ᶜρ, ᶜρe_int, ᶜuₕ))
     end
 end
-face_initial_condition(local_geometry) =
-    (; w = Geometry.Covariant3Vector(FT(0)))
+
+function face_initial_condition(local_geometry)
+    (; x, z) = local_geometry.coordinates
+    w = @. Geometry.Covariant3Vector(zero(z))
+    return NamedTuple{(:w,)}.(tuple.(w))
+end
 
 function postprocessing(sol, output_dir)
     ᶜlocal_geometry = Fields.local_geometry_field(sol.u[1].c)
@@ -189,20 +219,26 @@ function ρfb_init_coefs(
     horizontal_mesh =
         periodic_line_mesh(; x_max, x_elem = upsampling_factor * x_elem)
     h_space = make_horizontal_space(horizontal_mesh, npoly)
-    center_space, _ =
-        make_hybrid_spaces(h_space, z_max, upsampling_factor * z_elem)
+    center_space, _ = make_hybrid_spaces(
+        h_space,
+        z_max,
+        upsampling_factor * z_elem;
+        z_stretch,
+    )
     ᶜlocal_geometry = Fields.local_geometry_field(center_space)
     ᶜx = ᶜlocal_geometry.coordinates.x
     ᶜz = ᶜlocal_geometry.coordinates.z
 
     # Bretherton transform of initial perturbation
     linearize_density_perturbation = false
     if linearize_density_perturbation
-        ᶜρb_init = @. -ρₛ * Tb_init(ᶜx, ᶜz) / T₀
+        ᶜρb_init = @. -ρₛ * Tb_init(ᶜx, ᶜz, ΔT, x_mid, d, z_max) / T₀
     else
         ᶜp₀ = @. p₀(ᶜz)
         ᶜρ₀ = @. ᶜp₀ / (R_d * T₀)
-        ᶜρ′_init = @. ᶜp₀ / (R_d * (T₀ + T′_init(ᶜx, ᶜz))) - ᶜρ₀
+        ᶜρ′_init =
+            @. ᶜp₀ / (R_d * (T₀ + T′_init(ᶜx, ᶜz, ΔT, x_mid, d, z_max, δ))) -
+               ᶜρ₀
         ᶜbretherton_factor_pρ = @. exp(-δ * ᶜz / 2)
         ᶜρb_init = @. ᶜρ′_init / ᶜbretherton_factor_pρ
     end

diff --git a/examples/hybrid/sphere/baroclinic_wave_utilities.jl b/examples/hybrid/sphere/baroclinic_wave_utilities.jl
@@ -80,35 +80,39 @@ cond(λ, ϕ) = (0 < r(λ, ϕ) < d_0) * (r(λ, ϕ) != R * pi)
     sind(λ - λ_c) / sin(r(λ, ϕ) / R) * cond(λ, ϕ)
 
 function center_initial_condition(
-    local_geometry,
+    ᶜlocal_geometry,
     ᶜ𝔼_name;
     is_balanced_flow = false,
 )
-    (; lat, long, z) = local_geometry.coordinates
-    ρ = pres(lat, z) / R_d / temp(lat, z)
-    u₀ = u(lat, z)
-    v₀ = v(lat, z)
+    (; lat, long, z) = ᶜlocal_geometry.coordinates
+    ᶜρ = @. pres(lat, z) / R_d / temp(lat, z)
+    u₀ = @. u(lat, z)
+    v₀ = @. v(lat, z)
     if !is_balanced_flow
-        u₀ += δu(long, lat, z)
-        v₀ += δv(long, lat, z)
+        @. u₀ += δu(long, lat, z)
+        @. v₀ += δv(long, lat, z)
     end
-    uₕ_local = Geometry.UVVector(u₀, v₀)
-    uₕ = Geometry.Covariant12Vector(uₕ_local, local_geometry)
+    ᶜuₕ_local = @. Geometry.UVVector(u₀, v₀)
+    ᶜuₕ = @. Geometry.Covariant12Vector(ᶜuₕ_local, ᶜlocal_geometry)
     if ᶜ𝔼_name === Val(:ρθ)
-        ρθ = ρ * θ(lat, z)
-        return (; ρ, ρθ, uₕ)
+        ᶜρθ = @. ᶜρ * θ(lat, z)
+        return NamedTuple{(:ρ, :ρθ, :uₕ)}.(tuple.(ᶜρ, ᶜρθ, ᶜuₕ))
     elseif ᶜ𝔼_name === Val(:ρe)
-        ρe =
-            ρ *
-            (cv_d * (temp(lat, z) - T_tri) + norm_sqr(uₕ_local) / 2 + grav * z)
-        return (; ρ, ρe, uₕ)
+        ᶜρe = @. ᶜρ * (
+            cv_d * (temp(lat, z) - T_tri) + norm_sqr(ᶜuₕ_local) / 2 + grav * z
+        )
+        return NamedTuple{(:ρ, :ρe, :uₕ)}.(tuple.(ᶜρ, ᶜρe, ᶜuₕ))
     elseif ᶜ𝔼_name === Val(:ρe_int)
-        ρe_int = ρ * cv_d * (temp(lat, z) - T_tri)
-        return (; ρ, ρe_int, uₕ)
+        ᶜρe_int = @. ᶜρ * cv_d * (temp(lat, z) - T_tri)
+        return NamedTuple{(:ρ, :ρe_int, :uₕ)}.(tuple.(ᶜρ, ᶜρe_int, ᶜuₕ))
     end
 end
-face_initial_condition(local_geometry) =
-    (; w = Geometry.Covariant3Vector(FT(0)))
+
+function face_initial_condition(local_geometry)
+    (; lat, long, z) = local_geometry.coordinates
+    w = @. Geometry.Covariant3Vector(zero(z))
+    return NamedTuple{(:w,)}.(tuple.(w))
+end
 
 ##
 ## Additional tendencies

diff --git a/src/Meshes/interval.jl b/src/Meshes/interval.jl
@@ -9,7 +9,7 @@ A 1D mesh on an `IntervalDomain`.
 
 Construct a 1D mesh with face locations at `faces`.
 
-    IntervalMesh(domain::IntervalDomain[, stetching=Uniform()]; nelems=)
+    IntervalMesh(domain::IntervalDomain[, stretching=Uniform()]; nelems=)
 
 Constuct a 1D mesh on `domain` with `nelems` elements, using `stretching`. Possible values of `stretching` are: