From 0b7bb36ac84c154dd7b1bfe55af504aa6d22da6d Mon Sep 17 00:00:00 2001 From: Valeria Barra Date: Sun, 17 Jul 2022 18:06:35 -0700 Subject: [PATCH] IGW: Add stetched grid option and generalize hydrostatic balance computation Co-authored-by: Dennis Yatunin --- .buildkite/pipeline.yml | 14 +++ examples/common_spaces.jl | 4 +- examples/hybrid/driver.jl | 23 ++++- .../hybrid/plane/inertial_gravity_wave.jl | 98 +++++++++++++------ .../sphere/baroclinic_wave_utilities.jl | 42 ++++---- src/Meshes/interval.jl | 2 +- 6 files changed, 126 insertions(+), 57 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index cbdd37b82c..28732581df 100755 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -835,6 +835,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 index 4d0d10b978..d8d2b755b7 100644 --- 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 index a484d8e152..560286575e 100644 --- a/examples/hybrid/driver.jl +++ b/examples/hybrid/driver.jl @@ -57,6 +57,15 @@ 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 @@ -64,6 +73,10 @@ else 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,6 +201,7 @@ 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) @@ -194,7 +209,7 @@ 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 index 7bd120f30d..e41a0bc2fd 100644 --- 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,7 +40,7 @@ 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 @@ -47,10 +48,12 @@ npoly, x_elem = 1, Int(x_max / Δx) # max small-scale dt = 1.5 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,7 +62,8 @@ 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(ᶜρ) ==> @@ -67,34 +71,60 @@ if is_discrete_hydrostatic_balance # 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,8 +219,12 @@ 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 @@ -198,11 +232,13 @@ function ρfb_init_coefs( # 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 index c9b30b7468..53c00fd462 100644 --- 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 index 450eb2c19a..19ac65c730 100644 --- 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: