diff --git a/NEWS.md b/NEWS.md index 4d7aafa..25c227c 100644 --- a/NEWS.md +++ b/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 diff --git a/docs/src/api.md b/docs/src/api.md index 36c4618..39a8e4d 100644 --- a/docs/src/api.md +++ b/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 diff --git a/docs/src/var.md b/docs/src/var.md index 52e0451..cecf55c 100644 --- a/docs/src/var.md +++ b/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 diff --git a/src/Var.jl b/src/Var.jl index 76590cd..b1bab4c 100644 --- a/src/Var.jl +++ b/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) diff --git a/test/test_Var.jl b/test/test_Var.jl index d4988d6..532c921 100644 --- a/test/test_Var.jl +++ b/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