Add Base.dropdims for OutputVars

NEWS.md

@@ -43,6 +43,7 @@ Makie.save("surface_plot.png", fig)
 ## Minor additions
 
 - The number of dimensions of a `OutputVar` can be accessed with `ndims`.
+- You can drop dimensions of size 1 with `dropdims`.
 
 ## Bug fixes
 

docs/src/api.md

@@ -41,6 +41,7 @@ Var.dim_names
 Var.units
 Var.has_units
 Var.remake
+Base.dropdims
 Var.slice
 Var.average_lat
 Var.weighted_average_lat

src/Var.jl

@@ -48,6 +48,7 @@ export OutputVar,
     dim_units,
     range_dim,
     ndims,
+    dropdims,
     reordered_as,
     resampled_as,
     has_units,
@@ -1191,6 +1192,48 @@ function Base.ndims(var::HasDimAndAttribs)
     return length(var.dims)
 end
 
+"""
+    dropdims(var::OutputVar; dims)
+
+Return a `OutputVar` with the same data as `var`, but with the dimensions
+specified by `dims`, an iterable of dimension names, removed. The size of each
+dimension in `dims` must be equal to `1`.
+
+The result shares the same underlying data as `var` which means modifying values of
+`var.data` also results in the same changes to the data of the resulting `OutputVar`.
+"""
+function Base.dropdims(var::OutputVar; dims)
+    dim_names = collect(
+        find_corresponding_dim_name_in_var(dim_name, var) for dim_name in dims
+    )
+
+    # Check dimension names are unique
+    allunique(dim_names) || error("Dimensions ($dim_names) must be unique")
+
+    # Check dimensions are of length one
+    not_size1_dims = filter(
+        dim_name -> size(var.data, var.dim2index[dim_name]) > 1,
+        dim_names,
+    )
+    isempty(not_size1_dims) ||
+        error("Length of dropped dims ($not_size1_dims) must all be size 1")
+
+    dim_indices = Tuple(var.dim2index[dim_name] for dim_name in dim_names)
+    ret_data = dropdims(var.data; dims = dim_indices)
+
+    dim_names = Set(dim_names)
+    ret_dims = deepcopy(filter(kv -> first(kv) ∉ dim_names, var.dims))
+
+    ret_dim_attribs =
+        deepcopy(filter(kv -> first(kv) ∉ dim_names, var.dim_attributes))
+    return remake(
+        var,
+        dims = ret_dims,
+        dim_attributes = ret_dim_attribs,
+        data = ret_data,
+    )
+end
+
 """
     _update_long_name_generic!(
     reduced_var::OutputVar,

test/test_Var.jl

@@ -1593,6 +1593,66 @@ end
     @test_throws ErrorException ClimaAnalysis.reordered_as(src_var, dest_var)
 end
 
+@testset "Drop dims" begin
+    # Reordering the dimensions of a var to match itself
+    lon = [0.0]
+    lat = [1.0]
+    time = [1.0, 2.0, 3.0]
+    var =
+        TemplateVar() |>
+        add_dim("lon", lon, units = "degrees") |>
+        add_dim("lat", lat, units = "degrees") |>
+        add_dim("time", time, units = "s") |>
+        add_attribs(long_name = "hi") |>
+        one_to_n_data(; collected = true) |>
+        initialize
+    drop_lon_var = dropdims(var; dims = ("longitude",))
+    @test size(drop_lon_var.data) == (1, 3)
+    @test drop_lon_var.data == reshape(var.data, (1, 3))
+    @test drop_lon_var.attributes == var.attributes
+    @test drop_lon_var.dims ==
+          filter(kv -> first(kv) in ("lat", "time"), var.dims)
+    @test drop_lon_var.dim_attributes ==
+          filter(kv -> first(kv) in ("lat", "time"), var.dim_attributes)
+    var.data[1] = 42
+    @test drop_lon_var.data[1] == 42
+
+    drop_both_var = dropdims(var; dims = ("lon", "lat"))
+    @test size(drop_both_var.data) == (3,)
+    @test drop_both_var.data == reshape(var.data, (3,))
+    @test drop_lon_var.attributes == var.attributes
+    @test drop_both_var.dims == filter(kv -> first(kv) in ("time",), var.dims)
+    @test drop_both_var.dim_attributes ==
+          filter(kv -> first(kv) in ("time",), var.dim_attributes)
+    var.data[1] = 100
+    @test drop_both_var.data[1] == 100
+
+    # Error handling
+    long = 20.0:30.0 |> collect
+    lat = 30.0:40.0 |> collect
+    var =
+        TemplateVar() |>
+        add_dim("long", long, units = "degrees") |>
+        add_dim("lat", lat, units = "degrees") |>
+        add_attribs(long_name = "hi") |>
+        initialize
+    # Error handling is performed by the call to `dropdims` with
+    # `AbstractArray`s, so the exact error thrown is not guaranteed here
+    @test_throws "Length of dropped dims ([\"long\"]) must all be size 1" dropdims(
+        var,
+        dims = ("long",),
+    )
+    @test_throws "Length of dropped dims ([\"long\"]) must all be size 1" dropdims(
+        var,
+        dims = ("lon",),
+    )
+    @test_throws ErrorException dropdims(var, dims = ("time",))
+    @test_throws "Dimensions ([\"long\", \"long\"]) must be unique" dropdims(
+        var,
+        dims = ("long", "long"),
+    )
+end
+
 @testset "Resampling over all dimensions" begin
     src_long = 0.0:180.0 |> collect
     src_lat = 0.0:90.0 |> collect