Add analytic mountain wave tests

.buildkite/pipeline.yml

@@ -208,6 +208,12 @@ steps:
           --job_id plane_density_current_test
         artifact_paths: "plane_density_current_test/output_active/*"
 
+      - label: ":computer: Analytic Cosine Mountain Test (2D)"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/plane_analytic_cosine_mountain_test.yml
+          --job_id plane_analytic_cosine_mountain_test
+        artifact_paths: "plane_analytic_cosine_mountain_test/output_active/*"
 
   - group: "Sphere Examples (Dycore)"
     steps:
@@ -887,6 +893,39 @@ steps:
           slurm_gpus: 1
           slurm_mem: 20G
 
+      - label: "GPU: Analytic Schar Mountain Test (2D)"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/plane_analytic_schar_mountain_test.yml
+          --job_id plane_analytic_schar_mountain_test
+        artifact_paths: "plane_analytic_schar_mountain_test/output_active/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+
+      - label: "GPU: Analytic Cosine Mountain Test (2D)"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/plane_analytic_cosine_mountain_test.yml
+          --job_id plane_analytic_cosine_mountain_test
+        artifact_paths: "plane_analytic_cosine_mountain_test/output_active/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+
+      - label: "GPU: Analytic Cosine Mountain Test (3D)"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/box_analytic_cosine_mountain_test.yml
+          --job_id box_analytic_cosine_mountain_test
+        artifact_paths: "box_analytic_cosine_mountain_test/output_active/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+
   - group: "Benchmarks"
     steps:
 

Project.toml

@@ -32,6 +32,7 @@ NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 RRTMGP = "a01a1ee8-cea4-48fc-987c-fc7878d79da1"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
@@ -72,6 +73,7 @@ NCDatasets = "0.14.2"
 NVTX = "0.3"
 Pkg = "1.8"
 Printf = "1"
+QuadGK = "2"
 RRTMGP = "0.17"
 Random = "1"
 RootSolvers = "0.2, 0.3, 0.4"

config/default_configs/default_config.yml

@@ -154,6 +154,9 @@ start_date:
 forcing:
   help: "Forcing [`nothing` (default), `held_suarez`]"
   value: ~
+analytic_check:
+  help: "Test final state against analytic reference, if one is available [`false` (default), `true`]"
+  value: false
 test_dycore_consistency:
   help: "Test dycore consistency [`false` (default), `true`]"
   value: false

config/model_configs/box_analytic_cosine_mountain_test.yml

@@ -0,0 +1,21 @@
+config: "box"
+FLOAT_TYPE: "Float64"
+initial_condition: "ConstantBuoyancyFrequencyProfile"
+topography: "Cosine3D"
+x_max: 300e3
+y_max: 300e3
+z_max: 21e3
+x_elem: 30
+y_elem: 30
+z_elem: 100
+z_stretch: false
+dt: "8secs"
+t_end: "24hours"
+use_reference_state: false
+rayleigh_sponge: true
+toml: [toml/analytic_mountain_test.toml]
+analytic_check: true
+output_default_diagnostics: false
+diagnostics:
+  - short_name: [orog, ua, wa, uapredicted, wapredicted, uaerror, waerror]
+    period: 1hours

config/model_configs/plane_analytic_cosine_mountain_test.yml

@@ -0,0 +1,19 @@
+config: "plane"
+FLOAT_TYPE: "Float64"
+initial_condition: "ConstantBuoyancyFrequencyProfile"
+topography: "Cosine2D"
+x_max: 300e3
+z_max: 21e3
+x_elem: 30
+z_elem: 100
+z_stretch: false
+dt: "8secs"
+t_end: "24hours"
+use_reference_state: false
+rayleigh_sponge: true
+toml: [toml/analytic_mountain_test.toml]
+analytic_check: true
+output_default_diagnostics: false
+diagnostics:
+  - short_name: [orog, ua, wa, uapredicted, wapredicted, uaerror, waerror]
+    period: 1hours

config/model_configs/plane_analytic_schar_mountain_test.yml

@@ -0,0 +1,19 @@
+config: "plane"
+FLOAT_TYPE: "Float64"
+initial_condition: "ConstantBuoyancyFrequencyProfile"
+topography: "Schar"
+x_max: 300e3
+z_max: 21e3
+x_elem: 300
+z_elem: 100
+z_stretch: false
+dt: "0.5secs"
+t_end: "24hours"
+use_reference_state: false
+rayleigh_sponge: true
+toml: [toml/analytic_mountain_test.toml]
+analytic_check: true
+output_default_diagnostics: false
+diagnostics:
+  - short_name: [orog, ua, wa, uapredicted, wapredicted, uaerror, waerror]
+    period: 1hours

docs/Manifest.toml

@@ -292,7 +292,7 @@ weakdeps = ["SparseArrays"]
     ChainRulesCoreSparseArraysExt = "SparseArrays"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
+deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "QuadGK", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.27.0"

examples/Manifest.toml

@@ -320,7 +320,7 @@ version = "0.5.5"
     GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
+deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "QuadGK", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.27.0"

examples/hybrid/driver.jl

@@ -10,6 +10,14 @@ import ClimaAtmos as CA
 import Random
 Random.seed!(1234)
 
+# TODO: Fix the bug that's making this necessary for 2D topography runs.
+import ClimaCore.DataLayouts: promote_parent_array_type
+import CUDA
+promote_parent_array_type(
+    ::Type{CUDA.CuArray{T1, N, B} where {N}},
+    ::Type{Array{T2, N} where {N}},
+) where {T1, T2, B} = CUDA.CuArray{promote_type(T1, T2), N, B} where {N}
+
 if !(@isdefined config)
     (; config_file, job_id) = CA.commandline_kwargs()
     config = CA.AtmosConfig(config_file; job_id)
@@ -22,9 +30,10 @@ sol_res = CA.solve_atmos!(simulation)
 (; p) = integrator
 
 import ClimaCore
-import ClimaCore: Topologies, Quadratures, Spaces
+import ClimaCore: Topologies, Quadratures, Spaces, Fields
 import ClimaAtmos.InitialConditions as ICs
 using Statistics: mean
+import LinearAlgebra: norm_sqr
 import ClimaAtmos.Parameters as CAP
 import Thermodynamics as TD
 import ClimaComms
@@ -114,6 +123,42 @@ end
 @info "Callback verification, n_expected_calls: $(CA.n_expected_calls(integrator))"
 @info "Callback verification, n_measured_calls: $(CA.n_measured_calls(integrator))"
 
+if config.parsed_args["analytic_check"]
+    @info "Comparing final state against predicted steady-state solution"
+
+    Y_end = integrator.sol.u[end]
+    (; predicted_steady_state) = integrator.p
+    @assert !isnothing(predicted_steady_state)
+
+    ᶜuₕ_error = Y_end.c.uₕ .- CA.C12.(integrator.p.predicted_steady_state.ᶜu)
+    ᶠu₃_error = Y_end.f.u₃ .- CA.C3.(integrator.p.predicted_steady_state.ᶠu)
+
+    ᶜsponge_mask = ifelse.(Fields.coordinate_field(Y_end.c).z .< 13e3, 1.0, 0.0)
+    ᶠsponge_mask = ifelse.(Fields.coordinate_field(Y_end.f).z .< 13e3, 1.0, 0.0)
+    ᶠsponge_and_surface_mask = copy(ᶠsponge_mask)
+    Fields.level(ᶠsponge_and_surface_mask, Fields.half) .= 0.0
+
+    uₕ_rmse = sqrt(sum(norm_sqr, ᶜuₕ_error .* ᶜsponge_mask) / sum(ᶜsponge_mask))
+    u₃_rmse = sqrt(sum(norm_sqr, ᶠu₃_error .* ᶠsponge_mask) / sum(ᶠsponge_mask))
+    u₃_rmse_no_sfc = sqrt(
+        sum(norm_sqr, ᶠu₃_error .* ᶠsponge_and_surface_mask) /
+        sum(ᶠsponge_and_surface_mask),
+    )
+
+    uₕ_rmse_by_level = map(1:5) do level
+        sqrt(mean(norm_sqr, Fields.level(ᶜuₕ_error, level)))
+    end
+    u₃_rmse_by_level = map(1:5) do level
+        sqrt(mean(norm_sqr, Fields.level(ᶠu₃_error, level - Fields.half)))
+    end
+
+    @info "    RMSE of uₕ below sponge layer: $uₕ_rmse"
+    @info "    RMSE of u₃ below sponge layer: $u₃_rmse"
+    @info "    RMSE of u₃ below sponge layer and above surface: $u₃_rmse_no_sfc"
+    @info "    RMSE of uₕ on first 5 levels: $uₕ_rmse_by_level"
+    @info "    RMSE of u₃ on first 5 levels: $u₃_rmse_by_level"
+end
+
 # Conservation checks
 if config.parsed_args["check_conservation"]
     FT = Spaces.undertype(axes(sol.u[end].c.ρ))

perf/Manifest.toml

@@ -331,7 +331,7 @@ version = "0.5.5"
     GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
+deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "QuadGK", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.27.0"