Add Var.push_dim

NEWS.md

@@ -8,6 +8,19 @@ main
 With this release, all resulting `OutputVar` from binary operations with
 `OutputVar`s and real numbers keep their units when appropriate.
 
+## Add singleton dimension
+
+This release introduces `push_dim` which allows the user to add a singleton
+dimension to end of the dimensions of a `OutputVar`.
+
+See the example below where a time dimension with the value 0.0 is added to
+a `OutputVar`.
+
+```julia
+# var is a OutputVar
+push_dim(var, "time", 0.0, "axis" => "T")
+```
+
 v0.5.20
 -------
 This release introduces the following features and bug fixes

docs/src/api.md

@@ -81,6 +81,7 @@ Var.find_corresponding_dim_name_in_var
 Var.dim_units
 Var.range_dim
 Base.permutedims(var::OutputVar, perm)
+Var.push_dim
 Var.reordered_as
 Var.resampled_as
 Var.convert_units

docs/src/howdoi.md

@@ -426,4 +426,32 @@ ClimaAnalysis.shift_to_previous_day(var) |> ClimaAnalysis.dates
 ```
 
 These functions are helpful with aligning the dates of observational and
-simulational data.
+simulation data.
+
+## How do I add a singleton dimension to a `OutputVar`?
+
+You can add a singleton dimension to a `OutputVar` with [`push_dim`](@ref).
+
+See the example below where a time dimension with the value `0.0` and an
+attribute about the axis is added to a `OutputVar`.
+
+```@setup push_dim
+import ClimaAnalysis
+import ClimaAnalysis.Template:
+    TemplateVar,
+    add_attribs,
+    add_dim,
+    initialize
+lat = [-90.0, 0.0, 90.0]
+var =
+    TemplateVar() |>
+    add_dim("lat", lat, units = "degrees") |>
+    add_attribs(short_name = "pr", start_date = "2010-1-1") |>
+    initialize
+```
+
+```@example push_dim
+# var is a OutputVar
+var = ClimaAnalysis.push_dim(var, "time", 0.0, "axis" => "T")
+var.dims
+```

src/Var.jl

@@ -75,7 +75,8 @@ export OutputVar,
     reverse_dim,
     reverse_dim!,
     remake,
-    permutedims
+    permutedims,
+    push_dim
 
 """
     Representing an output variable
@@ -2649,6 +2650,38 @@ function Base.cat(vars::OutputVar...; dims::String)
     )
 end
 
+"""
+    push_dim(var::OutputVar, dim_name, dim_value, dim_attrib_pairs...)
+
+Push a singleton dimension to the end of the dimensions in `var`.
+
+If you want to change the order of the dimensions, then call `permutedims` after
+this function.
+"""
+function push_dim(var::OutputVar, dim_name, dim_value, dim_attrib_pairs...;)
+    dims = deepcopy(var.dims)
+    dim_attribs = deepcopy(var.dim_attributes)
+    data = deepcopy(var.data)
+
+    conventional_dim_names = conventional_dim_name.(keys(dims))
+    if conventional_dim_name(dim_name) in conventional_dim_names
+        error(
+            "Dimension ($dim_name) cannot be added as it already corresponds to an existing dimension ($conventional_dim_names)",
+        )
+    end
+
+    dims[dim_name] = [dim_value]
+    dim_attrib =
+        isempty(dim_attrib_pairs) ? valtype(dim_attribs)() :
+        Dict(dim_attrib_pairs)
+    dim_attribs[dim_name] = dim_attrib
+
+    new_dims = (size(data)..., 1)
+    data = reshape(data, new_dims)
+
+    return remake(var, dims = dims, data = data, dim_attributes = dim_attribs)
+end
+
 """
     Base.show(io::IO, var::OutputVar)
 

test/test_Var.jl

@@ -3544,6 +3544,46 @@ end
     )
 end
 
+@testset "Push dim" begin
+    lon = [-60.0, 45.0]
+    var =
+        TemplateVar() |>
+        add_dim("lon", lon, units = "degrees") |>
+        add_attribs(long_name = "hi") |>
+        initialize
+
+    time_var = ClimaAnalysis.push_dim(var, "time", 2.0)
+
+    @test time_var.attributes == var.attributes
+    @test time_var.dims["lon"] == var.dims["lon"]
+    @test time_var.dims["time"] == [2.0]
+    @test length(time_var.dim_attributes) == 2
+    @test time_var.dim_attributes["lon"] == var.dim_attributes["lon"]
+    @test time_var.dim_attributes["time"] == Dict{String, String}()
+    @test isequal(vec(time_var.data), vec(var.data))
+    @test size(time_var.data) == (2, 1)
+
+    lat_var = ClimaAnalysis.push_dim(
+        var,
+        "lat",
+        0.0,
+        "units" => "s",
+        "is_lat" => "yes",
+    )
+    @test lat_var.attributes == var.attributes
+    @test lat_var.dims["lon"] == var.dims["lon"]
+    @test lat_var.dims["lat"] == [0.0]
+    @test length(lat_var.dim_attributes) == 2
+    @test lat_var.dim_attributes["lon"] == var.dim_attributes["lon"]
+    @test lat_var.dim_attributes["lat"] ==
+          Dict("units" => "s", "is_lat" => "yes")
+    @test isequal(vec(lat_var.data), vec(var.data))
+    @test size(lat_var.data) == (2, 1)
+
+    # Error handling
+    @test_throws ErrorException ClimaAnalysis.push_dim(var, "lon", 42.0)
+end
+
 @testset "Show" begin
     lat = collect(range(-89.5, 89.5, 180))
     lon = collect(range(-179.5, 179.5, 360))