diff --git a/NEWS.md b/NEWS.md index d2b43582..e52f59e0 100644 --- a/NEWS.md +++ b/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 diff --git a/docs/src/api.md b/docs/src/api.md index c1b8808d..5824c44a 100644 --- a/docs/src/api.md +++ b/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 diff --git a/src/Var.jl b/src/Var.jl index 13d20243..af92d169 100644 --- a/src/Var.jl +++ b/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, diff --git a/test/test_Var.jl b/test/test_Var.jl index 976bf382..b6e6215d 100644 --- a/test/test_Var.jl +++ b/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