Add Leaderboard.add_category, add_model, add_unit

This commit adds functionality to add a single or multiple categories,
models, or units.

Author: ph-kev

NEWS.md

@@ -355,6 +355,19 @@ rmse_var[2,5] = 11.2
 rmse_var[:, :]
 ```
 
+#### Adding categories, models, and units
+Adding categories (e.g., seasons, months, years, etc.), models, and units to a `RMSEVariable`
+can be done using `add_category`, `add_model`, and `add_units!`.
+
+See the example below for how to use this functionality.
+
+Adding categories (e.g., seasons, months, years, etc.), models, and units to a `RMSEVariable`
+rmse_var2 = ClimaAnalysis.add_category(rmse_var, "Jan") # can take in mode than one category
+rmse_var = ClimaAnalysis.add_model(rmse_var, "CliMA") # can take in more than one model name
+ClimaAnalysis.add_unit!(rmse_var, "CliMA", "K")
+ClimaAnalysis.add_unit!(rmse_var, Dict("CliMA" => "K")) # for adding multiple units
+```
+
 ## Bug fixes
 
 - Increased the default value for `warp_string` to 72.

docs/src/api.md

@@ -85,6 +85,9 @@ Base.getindex(rmse_var::RMSEVariable, model_name, category)
 Base.getindex(rmse_var::RMSEVariable, model_name::String)
 Base.setindex!(rmse_var::RMSEVariable, rmse, model_name, category)
 Base.setindex!(rmse_var::RMSEVariable, rmse, model_name::String)
+Leaderboard.add_category
+Leaderboard.add_model
+Leaderboard.add_unit!
 ```
 
 ## Utilities

docs/src/rmse_var.md

@@ -110,4 +110,28 @@ rmse_var[:, "MAM"]
 rmse_var["ACCESS-CM2", ["ANN", "DJF", "MAM"]]
 rmse_var[2,5] = 11.2;
 rmse_var[:, :]
+```
+
+## Adding categories, models, and units
+
+It may be the case that the CSV file does not contain all the models you want to analyze, or
+you want to consider another category but do not want to go in and manually edit the CSV
+file to add it. `ClimaAnalysis` provides `add_category`, `add_model`, and `add_unit!` for
+adding categories, models, and units respectively. Multiple model or categories can be
+provided (e.g., `add_model(rmse_var, "model1", "model2")`) in the functions. For adding
+multiple units, one can pass in a dictionary mapping model names to units. See the example
+below using this functionality.
+
+```@julia rmse_var
+rmse_var2 = ClimaAnalysis.add_category(rmse_var, "Jan") # can take in more than one category
+rmse_var = ClimaAnalysis.add_model(rmse_var, "CliMA") # can take in more than one model name
+ClimaAnalysis.add_unit!(rmse_var, "CliMA", "K")
+ClimaAnalysis.add_unit!(rmse_var, Dict("CliMA" => "K")) # for adding multiple units
+```
+
+```@repl rmse_var
+ClimaAnalysis.category_names(rmse_var2)
+ClimaAnalysis.model_names(rmse_var)
+ClimaAnalysis.rmse_units(rmse_var)
+rmse_var[:,:]
 ```

src/Leaderboard.jl

@@ -9,7 +9,10 @@ export RMSEVariable,
     rmse_units,
     read_rmses,
     getindex,
-    setindex!
+    setindex!,
+    add_category,
+    add_model,
+    add_unit!
 
 """
     Holding root mean squared errors over multiple categories and models for a single
@@ -422,4 +425,97 @@ function Base.setindex!(rmse_var::RMSEVariable, rmse, model_name::String)
     rmse_var.RMSEs[model_idx, :] = rmse
 end
 
+"""
+    add_category(rmse_var::RMSEVariable, categories::String...)
+
+Add one or more categories named `categories` to `rmse_var`.
+"""
+function add_category(rmse_var::RMSEVariable, categories::String...)
+    # Add new category
+    categ_names = category_names(rmse_var)
+    push!(categ_names, categories...)
+
+    # Add new column
+    mdl_names = model_names(rmse_var)
+    num_mdl_names = length(mdl_names)
+    nan_vecs = (fill(NaN, num_mdl_names) for _ in categories)
+    rmses = hcat(rmse_var.RMSEs, nan_vecs...)
+    return RMSEVariable(
+        rmse_var.short_name,
+        mdl_names,
+        categ_names,
+        rmses,
+        rmse_var.units |> deepcopy,
+    )
+end
+
+"""
+    add_model(rmse_var::RMSEVariable, models::String...)
+
+Add one or more models named `models` to `rmse_var`.
+"""
+function add_model(rmse_var::RMSEVariable, models::String...)
+    # Add new model name
+    mdl_names = model_names(rmse_var)
+    push!(mdl_names, models...)
+
+    # Add new row
+    categ_names = category_names(rmse_var)
+    num_categ_names = length(categ_names)
+    nan_vecs = (fill(NaN, num_categ_names)' for _ in models)
+    rmses = vcat(rmse_var.RMSEs, nan_vecs...)
+
+    # Add missing units for model
+    units = rmse_var.units |> deepcopy
+    for name in models
+        units[name] = ""
+    end
+    return RMSEVariable(
+        rmse_var.short_name,
+        mdl_names,
+        categ_names,
+        rmses,
+        units,
+    )
+end
+
+"""
+    _model_name_check(rmse_var::RMSEVariable, model_name)
+
+Check if `model_name` is present in the model names of `rmse_var`.
+
+Return nothing if `model_name` is present in the model names of `rmse_var`. Otherwise,
+return an error.
+"""
+function _model_name_check(rmse_var::RMSEVariable, model_name)
+    mdl_names = model_names(rmse_var)
+    (model_name in mdl_names) ||
+        error("Model name ($model_name) is not in $mdl_names")
+    return nothing
+end
+
+"""
+    add_unit!(rmse_var::RMSEVariable, model_name, unit)
+
+Add a unit named `unit` to a model named `model_name` in `rmse_var`.
+"""
+function add_unit!(rmse_var::RMSEVariable, model_name, unit)
+    _model_name_check(rmse_var, model_name)
+    rmse_var.units[model_name] = unit
+    return nothing
+end
+
+"""
+    add_unit!(rmse_var::RMSEVariable, model_name2unit::Dict)
+
+Add all model name and unit pairs in the dictionary `model_name2unit` to `rmse_var`.
+"""
+function add_unit!(rmse_var::RMSEVariable, model_name2unit::Dict)
+    for (model_name, unit) in model_name2unit
+        _model_name_check(rmse_var, model_name)
+        rmse_var.units[model_name] = unit
+    end
+    return nothing
+end
+
 end

test/test_Leaderboard.jl

@@ -188,3 +188,78 @@ end
         [120.0, 130.0, 140.0]
     @test_throws DimensionMismatch rmse_var[1, :] = [120.0, 130.0, 140.0]
 end
+
+@testset "Adding model and category" begin
+    # Add single model
+    csv_file_path = joinpath(@__DIR__, "sample_data/test_csv.csv")
+    rmse_var = ClimaAnalysis.read_rmses(csv_file_path, "ta")
+    rmse_var = ClimaAnalysis.add_model(rmse_var, "new model")
+    @test ClimaAnalysis.model_names(rmse_var) ==
+          ["ACCESS-CM2", "ACCESS-ESM1-5", "new model"]
+    @test ClimaAnalysis.category_names(rmse_var) ==
+          ["DJF", "MAM", "JJA", "SON", "ANN"]
+    @test ClimaAnalysis.rmse_units(rmse_var) ==
+          Dict("ACCESS-CM2" => "", "ACCESS-ESM1-5" => "", "new model" => "")
+    @test all(isnan.(rmse_var["new model"]))
+
+    # Add single category
+    rmse_var = ClimaAnalysis.add_category(rmse_var, "new cat")
+    @test ClimaAnalysis.model_names(rmse_var) ==
+          ["ACCESS-CM2", "ACCESS-ESM1-5", "new model"]
+    @test ClimaAnalysis.category_names(rmse_var) ==
+          ["DJF", "MAM", "JJA", "SON", "ANN", "new cat"]
+    @test ClimaAnalysis.rmse_units(rmse_var) ==
+          Dict("ACCESS-CM2" => "", "ACCESS-ESM1-5" => "", "new model" => "")
+    @test all(isnan.(rmse_var[:, "new cat"]))
+
+    # Add multiple models
+    csv_file_path = joinpath(@__DIR__, "sample_data/test_csv.csv")
+    rmse_var = ClimaAnalysis.read_rmses(csv_file_path, "ta")
+    rmse_var = ClimaAnalysis.add_model(rmse_var, "new model", "new model 2")
+    @test ClimaAnalysis.model_names(rmse_var) ==
+          ["ACCESS-CM2", "ACCESS-ESM1-5", "new model", "new model 2"]
+    @test ClimaAnalysis.category_names(rmse_var) ==
+          ["DJF", "MAM", "JJA", "SON", "ANN"]
+    @test ClimaAnalysis.rmse_units(rmse_var) == Dict(
+        "ACCESS-CM2" => "",
+        "ACCESS-ESM1-5" => "",
+        "new model" => "",
+        "new model 2" => "",
+    )
+    @test all(isnan.(rmse_var["new model"]))
+    @test all(isnan.(rmse_var["new model 2"]))
+
+    # Add multiple categories
+    rmse_var = ClimaAnalysis.add_category(rmse_var, "new cat", "new cat 2")
+    @test ClimaAnalysis.model_names(rmse_var) ==
+          ["ACCESS-CM2", "ACCESS-ESM1-5", "new model", "new model 2"]
+    @test ClimaAnalysis.category_names(rmse_var) ==
+          ["DJF", "MAM", "JJA", "SON", "ANN", "new cat", "new cat 2"]
+    @test ClimaAnalysis.rmse_units(rmse_var) == Dict(
+        "ACCESS-CM2" => "",
+        "ACCESS-ESM1-5" => "",
+        "new model" => "",
+        "new model 2" => "",
+    )
+    @test all(isnan.(rmse_var[:, "new cat"]))
+    @test all(isnan.(rmse_var[:, "new cat 2"]))
+end
+
+@testset "Adding units" begin
+    csv_file_path = joinpath(@__DIR__, "sample_data/test_csv.csv")
+    rmse_var = ClimaAnalysis.read_rmses(csv_file_path, "ta")
+
+    # Adding a single unit
+    rmse_var = ClimaAnalysis.add_model(rmse_var, "hi")
+    ClimaAnalysis.add_unit!(rmse_var, "hi", "test")
+    @test ClimaAnalysis.rmse_units(rmse_var)["hi"] == "test"
+
+    # Adding multiple units
+    rmse_var = ClimaAnalysis.add_model(rmse_var, "hello1", "hello2")
+    ClimaAnalysis.add_unit!(
+        rmse_var,
+        Dict("hello1" => "units1", "hello2" => "units2"),
+    )
+    @test ClimaAnalysis.rmse_units(rmse_var)["hello1"] == "units1"
+    @test ClimaAnalysis.rmse_units(rmse_var)["hello2"] == "units2"
+end