Add indexing to RMSEVariable

This commit adds indexing to the `RMSEVariable`. In other words, the
struct behaves like an array. We allow indexing by both integers and
strings. We also made the decision to support indexing using only a
single string. For instance, `rmse_var["model_name"]` will return the
row of RMSEs across the different categories for that model. However, we
do not support `rmse_var[1]` because it is unclear whether this should
return a row (similar to the string case) or the first value of the
entry (similar to if the struct behaves like an array).

Author: ph-kev

NEWS.md

@@ -343,6 +343,18 @@ rmse_var = ClimaAnalysis.read_rmses(
 )
 ```
 
+#### Indexing
+`RMSEVariable` supports indexing by integer or string. See the example for indexing into
+a `RMSEVariable`.
+
+```julia
+rmse_var["ACCESS-CM2"]
+rmse_var[:, "MAM"]
+rmse_var["ACCESS-CM2", ["ANN", "DJF", "MAM"]]
+rmse_var[2,5] = 11.2
+rmse_var[:, :]
+```
+
 ## Bug fixes
 
 - Increased the default value for `warp_string` to 72.

docs/src/api.md

@@ -81,6 +81,10 @@ Leaderboard.model_names
 Leaderboard.category_names
 Leaderboard.rmse_units
 Leaderboard.read_rmses
+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)
 ```
 
 ## Utilities

docs/src/rmse_var.md

@@ -93,4 +93,21 @@ rmse_var = ClimaAnalysis.read_rmses(
 )
 
 nothing # hide
+```
+
+## Indexing
+
+After loading the data, one may want to inspect, change, or manipulate the data. This is
+possible by the indexing functionality that `RMSEVariable` provides. Indexing into a
+`RMSEVariable` is similar, but not the same as indexing into an array. Indexing by
+integer or string is supported, but linear indexing (e.g. `rmse_var[1]`) is not supported.
+integer or string is supported, but linear indexing (e.g., `rmse_var[1]`) is not supported.
+
+```@repl rmse_var
+rmse_var[:, :]
+rmse_var["ACCESS-CM2"]
+rmse_var[:, "MAM"]
+rmse_var["ACCESS-CM2", ["ANN", "DJF", "MAM"]]
+rmse_var[2,5] = 11.2;
+rmse_var[:, :]
 ```

src/Leaderboard.jl

@@ -7,7 +7,9 @@ export RMSEVariable,
     model_names,
     category_names,
     rmse_units,
-    read_rmses
+    read_rmses,
+    getindex,
+    setindex!
 
 """
     Holding root mean squared errors over multiple categories and models for a single
@@ -304,4 +306,120 @@ function read_rmses(csv_file::String, short_name::String; units = nothing)
     )
 end
 
+"""
+    function _index_convert(key2index, key::Colon)
+
+Convert the symbol colon into an index for indexing.
+"""
+function _index_convert(key2index, key::Colon)
+    return collect(values(key2index))
+end
+
+"""
+    function _index_convert(key2index,
+                            indices::AbstractVector{I})
+                            where {I <: Integer}
+
+Convert a string into an index for indexing.
+"""
+function _index_convert(key2index, key::AbstractString)
+    !haskey(key2index, key) &&
+        error("Key ($key) is not present in ($(keys(key2index)))")
+    return key2index[key]
+end
+
+"""
+    function _index_convert(key2index,
+                            keys::AbstractVector{S})
+                            where {S <: AbstractString}
+
+Convert a vector of strings to indices for indexing.
+"""
+function _index_convert(
+    key2index,
+    keys::AbstractVector{S},
+) where {S <: AbstractString}
+    for key in keys
+        !haskey(key2index, key) &&
+            error("Key ($key) is not present in ($(keys(key2index)))")
+    end
+    return [key2index[key] for key in keys]
+end
+
+"""
+    function _index_convert(key2index,
+                            indices::AbstractVector{I})
+                            where {I <: Integer}
+
+Convert an integer to an index for indexing.
+"""
+function _index_convert(key2index, index::Integer)
+    !(index in values(key2index)) &&
+        error("Index ($index) is not present in ($(values(key2index)))")
+    return index
+end
+
+"""
+    function _index_convert(key2index,
+                            indices::AbstractVector{I})
+                            where {I <: Integer}
+
+Convert a vector of integers to indices for indexing.
+"""
+function _index_convert(
+    key2index,
+    indices::AbstractVector{I},
+) where {I <: Integer}
+    for index in indices
+        !(index in values(key2index)) &&
+            error("Index ($index) is not present in ($(values(key2index)))")
+    end
+    return indices
+end
+
+"""
+    Base.getindex(rmse_var::RMSEVariable, model_name, category)
+
+Return a subset of the array holding the root mean square errors as specified by
+`model_name` and `category`. Support indexing by `String` and `Int`.
+"""
+function Base.getindex(rmse_var::RMSEVariable, model_name, category)
+    model_idx = _index_convert(rmse_var.model2index, model_name)
+    cat_idx = _index_convert(rmse_var.category2index, category)
+    return rmse_var.RMSEs[model_idx, cat_idx]
+end
+
+"""
+    Base.getindex(rmse_var::RMSEVariable, model_name::String)
+
+Return a subset of the array holding the root mean square errors as specified by
+`model_name`. Support indexing by `String`. Do not support linear indexing.
+"""
+function Base.getindex(rmse_var::RMSEVariable, model_name::String)
+    return rmse_var[model_name, :]
+end
+
+"""
+    Base.setindex!(rmse_var::RMSEVariable, rmse, model_name, category)
+
+Store a value or values from an array in the array of root mean squared errors in
+`rmse_var`. Support indexing by `String` and `Int`.
+"""
+function Base.setindex!(rmse_var::RMSEVariable, rmse, model_name, category)
+    model_idx = _index_convert(rmse_var.model2index, model_name)
+    cat_idx = _index_convert(rmse_var.category2index, category)
+    rmse_var.RMSEs[model_idx, cat_idx] = rmse
+end
+
+"""
+    Base.setindex!(rmse_var::RMSEVariable, rmse, model_name::String)
+
+Store a value or values from an array into the array of root mean squared errors in
+`rmse_var`. Support indexing by `String`. Do not support linear indexing.
+"""
+function Base.setindex!(rmse_var::RMSEVariable, rmse, model_name::String)
+    model_idx = _index_convert(rmse_var.model2index, model_name)
+    rmse_var.RMSEs[model_idx, :] = rmse
+end
+
 end

test/test_Leaderboard.jl

@@ -155,3 +155,36 @@ end
     @test ClimaAnalysis.rmse_units(rmse_var) ==
           Dict("ACCESS-CM2" => "m", "ACCESS-ESM1-5" => "m")
 end
+
+@testset "Indexing" begin
+    csv_file_path = joinpath(@__DIR__, "sample_data/test_csv.csv")
+    rmse_var = ClimaAnalysis.read_rmses(csv_file_path, "ta")
+
+    rmse_var[1, 1] = 100.0
+    @test rmse_var.RMSEs[1, 1] == 100.0
+
+    rmse_var[1, 1:3] = [100.0, 110.0, 120.0]
+    @test rmse_var.RMSEs[1, 1:3] == [100.0, 110.0, 120.0]
+
+    rmse_var[1:2, 1] = [100.0, 110.0]'
+    @test rmse_var.RMSEs[1:2, 1] == [100.0, 110.0]
+
+    rmse_var["ACCESS-ESM1-5", "DJF"] = 200.0
+    @test rmse_var["ACCESS-ESM1-5", "DJF"] == 200.0
+
+    rmse_var["ACCESS-ESM1-5", ["DJF", "MAM", "ANN"]] = [200.0, 210.0, 220.0]
+    @test rmse_var["ACCESS-ESM1-5", [1, 2, 5]] == [200.0, 210.0, 220.0]
+
+    rmse_var["ACCESS-ESM1-5"] = [120.0, 130.0, 140.0, 150.0, 160.0]
+    @test rmse_var["ACCESS-ESM1-5"] == [120.0, 130.0, 140.0, 150.0, 160.0]
+
+    # Check error handling
+    @test_throws ErrorException rmse_var[5, 5] = 100.0
+    @test_throws ErrorException rmse_var["do not exist"] = 100.0
+    @test_throws ErrorException rmse_var["do not exist", "test"] = 100.0
+    @test_throws ErrorException rmse_var["ACCESS-ESM1-5", "test"] = 100.0
+    @test_throws ErrorException rmse_var["model1", "ANN"] = 100.0
+    @test_throws DimensionMismatch rmse_var["ACCESS-ESM1-5"] =
+        [120.0, 130.0, 140.0]
+    @test_throws DimensionMismatch rmse_var[1, :] = [120.0, 130.0, 140.0]
+end