Fix precomputed quantities for Rosenbrock timestepper

Author: dennisYatunin

.buildkite/pipeline.yml

@@ -981,6 +981,28 @@ steps:
         agents:
           slurm_gpus: 1
 
+      - label: "GPU: Analytic Schar Mountain Test (2D, Float32)"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/plane_analytic_schar_mountain_float32_test.yml
+          --job_id gpu_plane_analytic_schar_mountain_float32_test
+        artifact_paths: "gpu_plane_analytic_schar_mountain_float32_test/output_active/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+
+      - label: "GPU: Analytic Teeny Teeny Teeny Tiny Schar Mountain Test (2D, Float32)"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/plane_analytic_teeny_teeny_teeny_tiny_schar_mountain_float32_test.yml
+          --job_id gpu_plane_analytic_teeny_teeny_teeny_tiny_schar_mountain_float32_test
+        artifact_paths: "gpu_plane_analytic_teeny_teeny_teeny_tiny_schar_mountain_float32_test/output_active/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+
       - label: "GPU: Analytic Schar Mountain Test (2D, Higher Horizontal Resolution)"
         command: >
           julia --color=yes --project=examples examples/hybrid/driver.jl
@@ -1003,6 +1025,17 @@ steps:
         agents:
           slurm_gpus: 1
 
+      - label: "GPU: Analytic Schar Mountain Test (2D, Stretched Grid)"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/plane_analytic_schar_mountain_stretched_grid_test.yml
+          --job_id gpu_plane_analytic_schar_mountain_stretched_grid_test
+        artifact_paths: "gpu_plane_analytic_schar_mountain_stretched_grid_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/model_configs/plane_analytic_schar_mountain_float32_test.yml

@@ -0,0 +1,18 @@
+config: "plane"
+FLOAT_TYPE: "Float32"
+initial_condition: "ConstantBuoyancyFrequencyProfile"
+topography: "Schar"
+x_max: 100e3
+z_max: 21e3
+x_elem: 100
+z_elem: 100
+z_stretch: false
+dt: "0.5secs"
+t_end: "24hours"
+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, c1error, c3error]
+    period: 1hours

config/model_configs/plane_analytic_schar_mountain_stretched_grid_test.yml

@@ -0,0 +1,19 @@
+config: "plane"
+FLOAT_TYPE: "Float64"
+initial_condition: "ConstantBuoyancyFrequencyProfile"
+topography: "Schar"
+x_max: 100e3
+z_max: 21e3
+x_elem: 100
+z_elem: 100
+dz_bottom: 50
+dz_top: 1000
+dt: "0.5secs"
+t_end: "24hours"
+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, c1error, c3error]
+    period: 1hours

examples/hybrid/driver.jl

@@ -10,6 +10,7 @@ import ClimaAtmos as CA
 import Random
 Random.seed!(1234)
 
+include("fix_rosenbrock_precomputed_quantities.jl")
 include("fix_1d_spectral_space_on_gpu.jl")
 # include("disable_topography_dss.jl")
 
@@ -122,25 +123,21 @@ 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
+    (; predicted_steady_state, params) = integrator.p
     @assert !isnothing(predicted_steady_state)
 
+    FT = eltype(Y_end)
+    (; zd_rayleigh) = params
+
     ᶜuₕ_error_squared =
         norm_sqr.(Y_end.c.uₕ .- CA.C12.(predicted_steady_state.ᶜu))
     ᶠu₃_error_squared =
         norm_sqr.(Y_end.f.u₃ .- CA.C3.(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
-
+    ᶜsponge_mask = FT.(Fields.coordinate_field(Y_end.c).z .< zd_rayleigh)
+    ᶠsponge_mask = FT.(Fields.coordinate_field(Y_end.f).z .< zd_rayleigh)
     uₕ_rmse = sqrt(sum(ᶜuₕ_error_squared .* ᶜsponge_mask) / sum(ᶜsponge_mask))
     u₃_rmse = sqrt(sum(ᶠu₃_error_squared .* ᶠsponge_mask) / sum(ᶠsponge_mask))
-    u₃_rmse_no_sfc = sqrt(
-        sum(ᶠu₃_error_squared .* ᶠsponge_and_surface_mask) /
-        sum(ᶠsponge_and_surface_mask),
-    )
 
     uₕ_rmse_by_level = map(1:5) do level
         sqrt(mean(Fields.level(ᶜuₕ_error_squared, level)))
@@ -151,7 +148,6 @@ if config.parsed_args["analytic_check"]
 
     @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

examples/hybrid/fix_rosenbrock_precomputed_quantities.jl

@@ -0,0 +1,93 @@
+import ClimaTimeSteppers as CTS
+import LinearAlgebra: ldiv!
+
+function CTS.step_u!(int, cache::CTS.RosenbrockCache{Nstages}) where {Nstages}
+    (; m, Γ, A, α, C) = int.alg.tableau
+    (; u, p, t, dt) = int
+    (; W, U, fU, fU_imp, fU_exp, fU_lim, k, ∂Y∂t) = cache
+    T_imp! = int.sol.prob.f.T_imp!
+    T_exp! = int.sol.prob.f.T_exp!
+    T_exp_lim! = int.sol.prob.f.T_exp_T_lim!
+    tgrad! = isnothing(T_imp!) ? nothing : T_imp!.tgrad
+
+    (; post_explicit!, post_implicit!, dss!) = int.sol.prob.f
+
+    # TODO: This is only valid when Γ[i, i] is constant, otherwise we have to
+    # move this in the for loop
+    dtγ = dt * Γ[1, 1]
+
+    if !isnothing(T_imp!)
+        Wfact! = int.sol.prob.f.T_imp!.Wfact
+        Wfact!(W, u, p, dtγ, t)
+    end
+
+    if !isnothing(tgrad!)
+        tgrad!(∂Y∂t, u, p, t)
+    end
+
+    for i in 1:Nstages
+        # Reset tendency
+        fill!(fU, 0)
+
+        αi = sum(α[i, 1:(i - 1)])
+        γi = sum(Γ[i, 1:i])
+
+        U .= u
+        for j in 1:(i - 1)
+            U .+= A[i, j] .* k[j]
+        end
+
+        if !isnothing(post_implicit!) && i != 1
+            # NOTE: post_implicit! is a misnomer
+            post_implicit!(U, p, t + αi * dt)
+        end
+
+        if !isnothing(T_imp!)
+            T_imp!(fU_imp, U, p, t + αi * dt)
+            fU .+= fU_imp
+        end
+
+        if !isnothing(T_exp!)
+            T_exp!(fU_exp, U, p, t + αi * dt)
+            fU .+= fU_exp
+        end
+
+        if !isnothing(T_exp_lim!)
+            T_exp_lim!(fU_exp, fU_lim, U, p, t + αi * dt)
+            fU .+= fU_exp
+            fU .+= fU_lim
+        end
+
+        if !isnothing(tgrad!)
+            fU .+= γi .* dt .* ∂Y∂t
+        end
+
+        # We dss the tendency at every stage but the last. At the last stage, we
+        # dss the incremented state
+        (i != Nstages) && dss!(fU, p, t + αi * dt)
+
+        for j in 1:(i - 1)
+            fU .+= (C[i, j] / dt) .* k[j]
+        end
+
+        fU .*= -dtγ
+
+        if !isnothing(T_imp!)
+            if W isa Matrix
+                ldiv!(k[i], lu(W), fU)
+            else
+                ldiv!(k[i], W, fU)
+            end
+        else
+            k[i] .= .-fU
+        end
+    end
+
+    for i in 1:Nstages
+        u .+= m[i] .* k[i]
+    end
+
+    dss!(u, p, t + dt)
+    post_explicit!(u, p, t + dt)
+    return nothing
+end

post_processing/ci_plots.jl

@@ -348,7 +348,7 @@ function plot_contours!(place, var)
 
     n_levels = 21
     var_mean = round(mean(var.data))
-    var_delta = minimum(abs, extrema(value -> value - var_mean, var.data))
+    var_delta = maximum(abs, extrema(value -> value - var_mean, var.data))
     levels = range(var_mean - var_delta, var_mean + var_delta, n_levels + 1)
     spectral_colors = Makie.to_colormap(:Spectral)
     plot_kwargs = (;
@@ -367,7 +367,13 @@ function plot_contours!(place, var)
     # end # TODO: This is taking too long to render.
     plot = CairoMakie.contourf!(dim1, dim2, var.data; plot_kwargs...)
 
-    CairoMakie.Colorbar(place[1, 2], plot, label = "$var_name [$var_units]")
+    try
+        CairoMakie.Colorbar(place[1, 2], plot, label = "$var_name [$var_units]")
+    catch _
+        @show mean(var.data)
+        @show var_mean
+        @show var_delta
+    end
 end
 
 """
@@ -713,8 +719,10 @@ const AnalyticMountainTest = Union{
     Val{:gpu_plane_analytic_teeny_tiny_schar_mountain_test},
     Val{:gpu_plane_analytic_teeny_teeny_tiny_schar_mountain_test},
     Val{:gpu_plane_analytic_teeny_teeny_teeny_tiny_schar_mountain_test},
+    Val{:gpu_plane_analytic_schar_mountain_float32_test},
     Val{:gpu_plane_analytic_schar_mountain_high_horz_res_test},
     Val{:gpu_plane_analytic_schar_mountain_high_vert_res_test},
+    Val{:gpu_plane_analytic_schar_mountain_stretched_grid_test},
     Val{:gpu_plane_analytic_cosine_mountain_test},
     Val{:gpu_extruded_plane_analytic_cosine_mountain_test},
     Val{:gpu_box_analytic_cosine_mountain_test},
@@ -738,17 +746,14 @@ function make_plots(
         Val(:gpu_plane_analytic_teeny_tiny_schar_mountain_test),
         Val(:gpu_plane_analytic_teeny_teeny_tiny_schar_mountain_test),
         Val(:gpu_plane_analytic_teeny_teeny_teeny_tiny_schar_mountain_test),
+        Val(:gpu_plane_analytic_schar_mountain_float32_test),
         Val(:gpu_plane_analytic_schar_mountain_high_horz_res_test),
         Val(:gpu_plane_analytic_schar_mountain_high_vert_res_test),
+        Val(:gpu_plane_analytic_schar_mountain_stretched_grid_test),
     )
 
-    error_short_names = ["uaerror", "waerror", "c1error", "c3error"]
-    summary_short_names =
-        ["ua", "wa", "uapredicted", "wapredicted", error_short_names...]
-    closeup_summary_short_names = ["wa", "wapredicted", "waerror", "c3error"]
-
     rms_error_vars = Iterators.flatmap((level_rms, column_rms)) do rms_func
-        Iterators.flatmap(error_short_names) do short_name
+        Iterators.flatmap(["uaerror", "waerror"]) do short_name
             map(simdirs) do simdir
                 data = get(simdir; short_name)
                 data = window(data, "z_reference"; left = 0, right = 13e3)
@@ -764,7 +769,17 @@ function make_plots(
         output_name = "rms_errors_24h",
     )
 
-    slice_summary_vars = Iterators.flatmap(summary_short_names) do short_name
+    all_short_names = [
+        "ua",
+        "wa",
+        "uapredicted",
+        "wapredicted",
+        "uaerror",
+        "waerror",
+        "c1error",
+        "c3error",
+    ]
+    slice_summary_vars = Iterators.flatmap(all_short_names) do short_name
         hours_to_plot = short_name in ("ua", "wa") ? (1, 2, 5, 24) : (24,)
         Iterators.flatmap(hours_to_plot) do n_hours
             map(simdirs) do simdir
@@ -783,8 +798,9 @@ function make_plots(
     )
 
     if is_schar_mountain
+        closeup_short_names = ["wa", "wapredicted", "waerror", "c3error"]
         closeup_slice_summary_vars =
-            Iterators.flatmap(closeup_summary_short_names) do short_name
+            Iterators.flatmap(closeup_short_names) do short_name
                 map(simdirs) do simdir
                     data = get(simdir; short_name)
                     data = window(data, "x"; left = 40e3, right = 60e3)

src/diagnostics/core_diagnostics.jl

@@ -881,16 +881,12 @@ add_diagnostic_variable!(
     comments = "Error of third covariant wind component against predicted value",
     compute! = (out, state, cache, time) -> begin
         if isnothing(out)
-            out =
-                state.f.u₃.components.data.:1 .-
-                C3.(cache.predicted_steady_state.ᶠu).components.data.:1
+            state.f.u₃.components.data.:1 .-
+            C3.(cache.predicted_steady_state.ᶠu).components.data.:1
         else
             out .=
                 state.f.u₃.components.data.:1 .-
                 C3.(cache.predicted_steady_state.ᶠu).components.data.:1
         end
-        # TODO: The surface level is currently off by an absurd amount.
-        Fields.level(out, Fields.half) .= 0
-        out
     end,
 )