Add Var.set_dim_units!

NEWS.md

@@ -55,6 +55,14 @@ the data. The function `replace` can be used to replace `missing` or `NaN` value
 ClimaAnalysis.replace(var, NaN => 0.0, missing => 0.0)
 ```
 
+### Set units for dimensions
+Similar to `set_units`, there is the function `set_dim_units!` which one can use to set
+the units of a dimension.
+
+```julia
+new_var = ClimaAnalysis.set_dim_units!(var, "lon", "degrees_east")
+```
+
 ## Bug fixes
 - Masking now affects the colorbar.
 - `Var.shift_to_start_of_previous_month` now checks for duplicate dates and throws an error

docs/src/api.md

@@ -58,6 +58,7 @@ Var.reordered_as
 Var.resampled_as
 Var.convert_units
 Var.set_units
+Var.set_dim_units!
 Var.integrate_lonlat
 Var.integrate_lat
 Var.integrate_lon

docs/src/var.md

@@ -55,13 +55,19 @@ new_var = ClimaAnalysis.convert_units(var, "kg m/s", conversion_function = (x) -
 
 !!! note If you find some unparseable units, please open an issue. We can fix them!
 
-If units do not exist or you want to change the name of the units, then one can uses the
+If units do not exist or you want to change the name of the units, then one can use the
 `set_units` function.
 ```julia
 new_var = ClimaAnalysis.set_units(var, "kg m s^-1")
 ```
+
+Similarly, to set the units of a dimension, one can use the `dim_set_units!` function.
+```julia
+new_var = ClimaAnalysis.set_dim_units!(var, "lon", "degrees_east")
+```
+
 !!! warning "Override existing units"
-    If units already exist, this will override the units for data in `var`.
+    If units already exist, this will override the units for data or the dimension in `var`.
 
 ## Interpolations and extrapolations
 

src/Var.jl

@@ -53,6 +53,7 @@ export OutputVar,
     global_mse,
     global_rmse,
     set_units,
+    set_dim_units!,
     shift_to_start_of_previous_month,
     apply_landmask,
     apply_oceanmask,
@@ -517,6 +518,30 @@ function set_units(var::OutputVar, units::AbstractString)
     return converted_var
 end
 
+"""
+    set_dim_units!(var::OutputVar, dim_name::AbstractString, units::AbstractString)
+
+Set `units` for the `dim_name` dimension in `var`.
+
+!!! warning "Override existing units"
+    If units already exist for the dimension, this will override the units for the dimension
+    in `var`.
+"""
+function set_dim_units!(
+    var::OutputVar,
+    dim_name::AbstractString,
+    units::AbstractString,
+)
+    !haskey(var.dims, dim_name) &&
+        error("Var does not have dimension $dim_name, found $(keys(var.dims))")
+    if haskey(var.dim_attributes, dim_name)
+        push!(var.dim_attributes[dim_name], "units" => units)
+    else
+        var.dim_attributes[dim_name] = Dict("units" => units)
+    end
+    return nothing
+end
+
 """
     is_z_1D(var::OutputVar)
 

test/test_Var.jl

@@ -1754,3 +1754,49 @@ end
     @test var_no_nan.attributes == var.attributes
     @test var_no_nan.dim_attributes == var.dim_attributes
 end
+
+@testset "Set units for dimension" begin
+    # Units exist in dim_attribs
+    lat = collect(range(-89.5, 89.5, 180))
+    lon = collect(range(-179.5, 179.5, 360))
+    data = ones(length(lat), length(lon))
+    dims = OrderedDict(["lat" => lat, "lon" => lon])
+    attribs = Dict("long_name" => "hi")
+    dim_attribs = OrderedDict([
+        "lat" => Dict("units" => "deg"),
+        "lon" => Dict("units" => "deg"),
+    ])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+    ClimaAnalysis.set_dim_units!(var, "lat", "degrees")
+    @test ClimaAnalysis.dim_units(var, "lat") == "degrees"
+
+    # Units do not exist in dim_attribs as a key
+    lat = collect(range(-89.5, 89.5, 180))
+    lon = collect(range(-179.5, 179.5, 360))
+    data = ones(length(lat), length(lon))
+    dims = OrderedDict(["lat" => lat, "lon" => lon])
+    attribs = Dict("long_name" => "hi")
+    dim_attribs =
+        OrderedDict(["lat" => Dict(), "lon" => Dict("units" => "deg")])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+    ClimaAnalysis.set_dim_units!(var, "lat", "degrees")
+    @test ClimaAnalysis.dim_units(var, "lat") == "degrees"
+
+    # Dimension is not present in dim_attribs
+    lat = collect(range(-89.5, 89.5, 180))
+    lon = collect(range(-179.5, 179.5, 360))
+    data = ones(length(lat), length(lon))
+    dims = OrderedDict(["lat" => lat, "lon" => lon])
+    attribs = Dict("long_name" => "hi")
+    dim_attribs = OrderedDict(["lon" => Dict("units" => "deg")])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+    ClimaAnalysis.set_dim_units!(var, "lat", "degrees")
+    @test ClimaAnalysis.dim_units(var, "lat") == "degrees"
+
+    # Error handling
+    @test_throws ErrorException ClimaAnalysis.set_dim_units!(
+        var,
+        "wacky",
+        "idk",
+    )
+end