Add support for unitful

NEWS.md

@@ -71,6 +71,29 @@ OrderedDict{String, Vector{Float64}} with 2 entries:
   "latitude" => [0.0, 1.0, 2.0]
 ```
 
+### Add support for converting units
+
+`ClimaAnalysis` now uses
+[Unitful](https://painterqubits.github.io/Unitful.jl/stable) to handle variable
+units, when possible.
+
+When a `OutputVar` has `units` among its `attributes`, `ClimaAnalysis` will try
+to use `Unitful` to parse it. If successful, `OutputVar` can be directly
+converted to other compatible units. For example, if `var` has units of `m/s`,
+```julia-repl
+julia> ClimaAnalysis.convert_units(var, "cm/s")
+```
+will convert to `cm/s`.
+
+Some units are not recognized by `Unitful`. Please, open an issue about that:
+we can add more units.
+
+In those cases, or when units are incompatible, you can also pass a
+`conversion_function` that specify how to transform units.
+```julia-repl
+julia> ClimaAnalysis.convert_units(var, "kg/s", conversion_function = (x) - 1000x)
+```
+
 ## Bug fixes
 
 - Increased the default value for `warp_string` to 72.

Project.toml

@@ -11,14 +11,15 @@ NaNStatistics = "b946abbf-3ea7-4610-9019-9858bfdeaf2d"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [weakdeps]
 GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [extensions]
-GeoMakieExt = "GeoMakie"
-MakieExt = "Makie"
+ClimaAnalysisGeoMakieExt = "GeoMakie"
+ClimaAnalysisMakieExt = "Makie"
 
 [compat]
 Aqua = "0.8"
@@ -37,6 +38,7 @@ Reexport = "1"
 SafeTestsets = "0.1"
 Statistics = "1"
 Test = "1"
+Unitful = "1"
 julia = "1.9"
 
 [extras]

docs/make.jl

@@ -16,8 +16,8 @@ DocMeta.setdocmeta!(
 makedocs(;
     modules = [
         ClimaAnalysis,
-        Base.get_extension(ClimaAnalysis, :MakieExt),
-        Base.get_extension(ClimaAnalysis, :GeoMakieExt),
+        Base.get_extension(ClimaAnalysis, :ClimaAnalysisMakieExt),
+        Base.get_extension(ClimaAnalysis, :ClimaAnalysisGeoMakieExt),
     ],
     authors = "Climate Modelling Alliance",
     sitename = "ClimaAnalysis.jl",

docs/src/api.md

@@ -23,6 +23,7 @@ Var.is_z_1D
 Var.short_name
 Var.long_name
 Var.units
+Var.has_units
 Var.slice
 Var.average_lat
 Var.weighted_average_lat
@@ -53,6 +54,7 @@ Var.conventional_dim_name
 Var.dim_units
 Var.range_dim
 Var.resampled_as
+Var.convert_units
 ```
 
 ## Utilities

docs/src/var.md

@@ -18,3 +18,39 @@ import ClimaAnalysis: OutputVar
 
 myfile = OutputVar("my_netcdf_file.nc", "myvar")
 ```
+
+## Physical units
+
+`OutputVar`s can contain information about their physical units. For
+`OutputVar`s read from NetCDF files, this is obtained from the `units` attribute
+(and stored in the `attributes["units"]`).
+
+When possible, `ClimaAnalysis` uses
+[Unitful](https://painterqubits.github.io/Unitful.jl/stable) to handle units.
+This enables automatic unit conversion for `OutputVar`s.
+
+Consider the following example: 
+```julia
+import ClimaAnalysis
+values = 0:100.0 |> collect
+data = copy(values)
+attribs = Dict("long_name" => "speed", "units" => "m/s")
+dim_attribs = Dict{String, Any}()
+var = ClimaAnalysis.OutputVar(attribs, Dict("distance" => values), dim_attribs, data)
+
+var_cms = ClimaAnalysis.convert_units(var, "cm/s")
+```
+In this example, we set up` var`, an `OutputVar` with units of meters per second.
+Then, we called [`ClimaAnalysis.convert_units`](@ref) to convert the units to
+centimeters per second.
+
+Sometimes, this automatic unit conversion is not possible (e.g., when you want
+to transform between incompatible units). In this case, you an pass a function
+that specify how to apply this transformation. For example, in the previous
+case, we can assume that we are talking about water and transform units into a
+mass flux:
+```julia
+new_var = ClimaAnalysis.convert_units(var, "kg m/s", conversion_function = (x) -> 1000x)
+```
+
+!!! note If you find some unparseable units, please open an issue. We can fix them!

ext/GeoMakieExt.jl → ext/ClimaAnalysisGeoMakieExt.jl

@@ -1,4 +1,4 @@
-module GeoMakieExt
+module ClimaAnalysisGeoMakieExt
 
 import GeoMakie
 import GeoMakie: Makie
@@ -39,7 +39,7 @@ function _geomakie_plot_on_globe!(
     lon = var.dims[lon_name]
     lat = var.dims[lat_name]
 
-    units = var.attributes["units"]
+    units = ClimaAnalysis.units(var)
     short_name = var.attributes["short_name"]
     colorbar_label = "$short_name [$units]"
 

ext/MakieExt.jl → ext/ClimaAnalysisMakieExt.jl

@@ -1,4 +1,4 @@
-module MakieExt
+module ClimaAnalysisMakieExt
 
 import Makie
 import ClimaAnalysis
@@ -51,7 +51,7 @@ function Visualize.heatmap2D!(
     dim1 = var.dims[dim1_name]
     dim2 = var.dims[dim2_name]
 
-    units = var.attributes["units"]
+    units = ClimaAnalysis.units(var)
     short_name = var.attributes["short_name"]
     colorbar_label = "$short_name [$units]"
     dim1_units = var.dim_attributes[dim1_name]["units"]
@@ -290,7 +290,7 @@ function Visualize.line_plot1D!(
     dim_name = var.index2dim[]
     dim = var.dims[dim_name]
 
-    units = var.attributes["units"]
+    units = ClimaAnalysis.units(var)
     short_name = var.attributes["short_name"]
 
     dim_units = var.dim_attributes[dim_name]["units"]
@@ -554,7 +554,7 @@ end
 """
     _to_unitrange(x::Number, lo::Number, hi::Number)
 
-Linearly transform x ∈ [lo, hi] to [0, 1]. 
+Linearly transform x ∈ [lo, hi] to [0, 1].
 """
 _to_unitrange(x::Number, lo::Number, hi::Number) = (x - lo) / (hi - lo)
 
@@ -581,18 +581,25 @@ symmetrically around zero maps the same color intensity to the same magnitude.
 # Returns
 - `cmap::Makie.ColorGradient`: a colormap
 """
-function _constrained_cmap(
+function Visualize._constrained_cmap(
     cols::Vector,
     lo,
     hi;
     mid = 0,
     categorical = false,
     rev = false,
 )
-    _constrained_cmap(Makie.ColorScheme(cols), lo, hi; mid, categorical, rev)
+    Visualize._constrained_cmap(
+        Makie.ColorScheme(cols),
+        lo,
+        hi;
+        mid,
+        categorical,
+        rev,
+    )
 end
 
-function _constrained_cmap(
+function Visualize._constrained_cmap(
     cols::Makie.ColorScheme,
     lo,
     hi;

ext/ClimaAnalysisUnitfulExt.jl

@@ -0,0 +1,76 @@
+module ClimaAnalysisUnitfulExt
+
+# ClimaAnalysisUnitfulExt is implemented as extension in case we decide to turn Unitful into
+# an extension in the future, but, right now, everything is included directly in
+# ClimaAnalysis.
+
+import Unitful
+
+import ClimaAnalysis: Var
+
+"""
+    _maybe_convert_to_unitful(value)
+
+Try converting `value` to a `Uniftul` object. If unsuccessful, just return it.
+"""
+function Var._maybe_convert_to_unitful(value)
+    # This function in inherently type-unstable
+    try
+        return Unitful.uparse(value)
+    catch exc
+        # ParseError when it cannot be parsed
+        # ArgumentError when symbols are not available
+        if exc isa Base.Meta.ParseError || exc isa ArgumentError
+            return value
+        else
+            rethrow(exc)
+        end
+    end
+end
+
+function _converted_data(data, conversion_function)
+    return conversion_function.(data)
+end
+
+function _converted_data_unitful(data, old_units, new_units)
+    # We add FT because sometimes convert changes the type and because
+    # ustrip only reinterprets the given array
+    FT = eltype(data)
+    return FT.(Unitful.ustrip(Unitful.uconvert.(new_units, data * old_units)))
+end
+
+function Var.convert_units(
+    var::Var.OutputVar,
+    new_units::AbstractString;
+    conversion_function = nothing,
+)
+    has_unitful_units =
+        Var.has_units(var) && (var.attributes["units"] isa Unitful.Units)
+    new_units_maybe_unitful = Var._maybe_convert_to_unitful(new_units)
+    new_units_are_unitful = new_units_maybe_unitful isa Unitful.Units
+
+    if has_unitful_units && new_units_are_unitful
+        isnothing(conversion_function) ||
+            @warn "Ignoring conversion_function, units are parseable."
+        convert_function =
+            data -> _converted_data_unitful(
+                data,
+                var.attributes["units"],
+                new_units_maybe_unitful,
+            )
+    else
+        isnothing(conversion_function) && error(
+            "Conversion function required for var with non-parseable/absent units.",
+        )
+        convert_function = data -> _converted_data(data, conversion_function)
+    end
+
+    new_data = convert_function(var.data)
+    new_attribs = copy(var.attributes)
+    # The constructor will take care of converting new_units to Unitful
+    new_attribs["units"] = new_units
+
+    return Var.OutputVar(new_attribs, var.dims, var.dim_attributes, new_data)
+end
+
+end

src/ClimaAnalysis.jl

@@ -14,4 +14,7 @@ include("Visualize.jl")
 
 include("Atmos.jl")
 
+# In case we want to turn Unitful into an extension
+include("../ext/ClimaAnalysisUnitfulExt.jl")
+
 end # module ClimaAnalysis