Fix precomputed quantities for Rosenbrock timestepper

.buildkite/pipeline.yml

@@ -981,6 +981,17 @@ 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 Schar Mountain Test (2D, Higher Horizontal Resolution)"
         command: >
           julia --color=yes --project=examples examples/hybrid/driver.jl
@@ -1003,6 +1014,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: "Float32"
+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")
 

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,10 @@ 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))
+    if var_delta == 0
+        var_delta = eps(Float64) # Colorbar throws an error when var_delta is 0
+    end
     levels = range(var_mean - var_delta, var_mean + var_delta, n_levels + 1)
     spectral_colors = Makie.to_colormap(:Spectral)
     plot_kwargs = (;
@@ -713,8 +716,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,8 +743,10 @@ 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"]

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,
 )