Merge pull request #1402 from CliMA/kp/rmse-regression

ph-kev · web-flow · commit acdc4e708f74 · 2025-07-03T10:48:02.000-07:00
Add check for RMSEs
diff --git a/experiments/ClimaEarth/leaderboard/data_sources.jl b/experiments/ClimaEarth/leaderboard/data_sources.jl
@@ -78,14 +78,18 @@ includes unit conversion and shifting the dates.
 The variable should have only three dimensions: latitude, longitude, and time.
 """
 function get_sim_var_dict(diagnostics_folder_path)
-    # List of short names
-    available_short_names = ClimaAnalysis.available_vars(ClimaAnalysis.SimDir(diagnostics_folder_path))
+    available_short_names = get_short_names_monthly_averages(diagnostics_folder_path)
     sim_var_dict = Dict{String, Any}()
     # Dict for loading in simulation data
     "pr" in available_short_names && (
         sim_var_dict["pr"] =
             () -> begin
-                sim_var = get(ClimaAnalysis.SimDir(diagnostics_folder_path), short_name = "pr")
+                sim_var = get(
+                    ClimaAnalysis.SimDir(diagnostics_folder_path),
+                    short_name = "pr",
+                    reduction = "average",
+                    period = "1M",
+                )
                 sim_var = ClimaAnalysis.convert_units(
                     sim_var,
                     "mm/day",
@@ -101,7 +105,12 @@ function get_sim_var_dict(diagnostics_folder_path)
         short_name in available_short_names && (
             sim_var_dict[short_name] =
                 () -> begin
-                    sim_var = get(ClimaAnalysis.SimDir(diagnostics_folder_path), short_name = short_name)
+                    sim_var = get(
+                        ClimaAnalysis.SimDir(diagnostics_folder_path),
+                        short_name = short_name,
+                        reduction = "average",
+                        period = "1M",
+                    )
                     sim_var = ClimaAnalysis.shift_to_start_of_previous_month(sim_var)
                     return sim_var
                 end
@@ -224,13 +233,24 @@ function get_sim_var_in_pfull_dict(diagnostics_folder_path)
     available_short_names = ClimaAnalysis.available_vars(ClimaAnalysis.SimDir(diagnostics_folder_path))
     sim_var_pfull_dict = Dict{String, Any}()
 
-    short_names = ["ta", "hur", "hus"]
+    short_names = get_short_names_monthly_averages(diagnostics_folder_path)
+    available_short_names = intersect(short_names, Set(["ta", "hur", "hus"]))
     for short_name in short_names
         short_name in available_short_names && (
             sim_var_pfull_dict[short_name] =
                 () -> begin
-                    sim_var = get(ClimaAnalysis.SimDir(diagnostics_folder_path), short_name = short_name)
-                    pfull_var = get(ClimaAnalysis.SimDir(diagnostics_folder_path), short_name = "pfull")
+                    sim_var = get(
+                        ClimaAnalysis.SimDir(diagnostics_folder_path),
+                        short_name = short_name,
+                        reduction = "average",
+                        period = "1M",
+                    )
+                    pfull_var = get(
+                        ClimaAnalysis.SimDir(diagnostics_folder_path),
+                        short_name = "pfull",
+                        reduction = "average",
+                        period = "1M",
+                    )
 
                     (ClimaAnalysis.units(sim_var) == "") &&
                         (sim_var = ClimaAnalysis.set_units(sim_var, "unitless"))
@@ -360,3 +380,23 @@ function get_compare_vars_biases_heatmap_extrema_pfull()
     compare_vars_biases_heatmap_extrema = Dict("ta" => (-10.0, 10.0), "hur" => (-0.4, 0.4), "hus" => (-0.001, 0.001))
     return compare_vars_biases_heatmap_extrema
 end
+
+"""
+    get_short_names_of_monthly_averages(diagnostics_folder_path)
+
+Get all the short names of the monthly averages.
+"""
+function get_short_names_monthly_averages(diagnostics_folder_path)
+    available_short_names = Set{String}()
+    simdir = ClimaAnalysis.SimDir(diagnostics_folder_path)
+    for short_name in ClimaAnalysis.available_vars(simdir)
+        for reduction in ClimaAnalysis.available_reductions(simdir, short_name = short_name)
+            for period in ClimaAnalysis.available_periods(simdir, short_name = short_name, reduction = reduction)
+                if reduction == "average" && period == "1M"
+                    push!(available_short_names, short_name)
+                end
+            end
+        end
+    end
+    return available_short_names
+end
diff --git a/experiments/ClimaEarth/leaderboard/test_rmses.jl b/experiments/ClimaEarth/leaderboard/test_rmses.jl
@@ -0,0 +1,78 @@
+import ClimaAnalysis
+import Dates
+import Test: @test, @testset
+
+include("data_sources.jl")
+
+"""
+    test_rmse_thresholds(diagnostics_folder_path, spinup)
+
+Test that the annual RMSE values for specific variables have not increased
+beyond acceptable thresholds. The variables tested are:
+- pr (precipitation)
+- rsut (top of atmosphere outgoing shortwave radiation)
+- rsutcs (clear-sky top of atmosphere outgoing shortwave radiation)
+
+The spinup is the number of months to remove from the beginning of the
+simulation.
+
+More variables can be added by adding the short name and RMSE pair to the
+dictionary returned by `get_rmse_thresholds`.
+
+If this test fails, it indicates a regression in the model's physics, resulting
+in a higher RMSE. If this increased RMSE is considered acceptable, then the
+thresholds should be updated accordingly.
+"""
+function test_rmse_thresholds(diagnostics_folder_path, spinup)
+    sim_var_dict = get_sim_var_dict(diagnostics_folder_path)
+    obs_var_dict = get_obs_var_dict()
+    rmse_thresholds = get_rmse_thresholds()
+
+    sim_vars = (sim_var_dict[short_name]() for short_name in keys(rmse_thresholds))
+    obs_vars =
+        (obs_var_dict[ClimaAnalysis.short_name(sim_var)](sim_var.attributes["start_date"]) for sim_var in sim_vars)
+    short_names = (ClimaAnalysis.short_name(var) for var in sim_vars)
+
+    rmses = map(sim_vars, obs_vars) do sim_var, obs_var
+        # Remove first spin_up_months from simulation
+        spinup_cutoff = spinup * 31 * 86400.0
+        ClimaAnalysis.times(sim_var)[end] >= spinup_cutoff &&
+            (sim_var = ClimaAnalysis.window(sim_var, "time", left = spinup_cutoff))
+
+        obs_var = ClimaAnalysis.resampled_as(obs_var, sim_var)
+        obs_var = ClimaAnalysis.average_time(obs_var)
+        sim_var = ClimaAnalysis.average_time(sim_var)
+
+        ClimaAnalysis.global_rmse(sim_var, obs_var)
+    end
+
+    @testset "RMSE thresholds" begin
+        for (short_name, rmse) in zip(short_names, rmses)
+            @info "RMSE for $short_name: $rmse"
+            @test rmse < rmse_thresholds[short_name]
+        end
+    end
+end
+
+"""
+    get_rmse_thresholds()
+
+Return a dictionary mapping short names to maximum acceptable RMSE values.
+"""
+function get_rmse_thresholds()
+    rmse_thresholds = Dict(
+        "pr" => 3.0,      # mm/day
+        "rsut" => 20.0,   # W/m²
+        "rsutcs" => 7.0,  # W/m²
+    )
+    return rmse_thresholds
+end
+
+if abspath(PROGRAM_FILE) == @__FILE__
+    if length(ARGS) != 1
+        error("Usage: julia test_rmses.jl <Filepath to simulation data>")
+    end
+    leaderboard_base_path = ARGS[begin]
+    spinup = 3
+    test_rmse_thresholds(leaderboard_base_path, spinup)
+end
diff --git a/experiments/ClimaEarth/user_io/postprocessing.jl b/experiments/ClimaEarth/user_io/postprocessing.jl
@@ -2,6 +2,7 @@
 include("debug_plots.jl")
 include("diagnostics_plots.jl")
 include("../leaderboard/leaderboard.jl")
+include("../leaderboard/test_rmses.jl")
 
 """
     postprocess_sim(cs, postprocessing_vars)
@@ -32,14 +33,16 @@ function postprocess_sim(cs, postprocessing_vars)
 
     # If we have enough data (in time, but also enough variables), plot the leaderboard.
     # We need pressure to compute the leaderboard.
-    pressure_in_output = "pfull" in CAN.available_vars(CAN.SimDir(atmos_output_dir))
-    if pressure_in_output
-        times = CAN.times(get(CAN.SimDir(atmos_output_dir), "pfull"))
+    simdir = CAN.SimDir(atmos_output_dir)
+    if !isempty(simdir)
+        pressure_in_output = "pfull" in CAN.available_vars(simdir)
+        times = CAN.times(get(simdir, first(CAN.available_vars(simdir))))
         t_end = times[end]
         if t_end > 84600 * 31 * 3 # 3 months for spin up
             leaderboard_base_path = artifact_dir
             compute_leaderboard(leaderboard_base_path, atmos_output_dir, 3)
-            compute_pfull_leaderboard(leaderboard_base_path, atmos_output_dir, 6)
+            test_rmse_thresholds(atmos_output_dir, 3)
+            pressure_in_output && compute_pfull_leaderboard(leaderboard_base_path, atmos_output_dir, 6)
         end
     end
 
