Breaking: add combine method for groupby output, fixing similar for AbstractDimStack (#903)

* add combine method

* test groupby and similar

* docs entry

Author: rafaqz

docs/src/api/reference.md

@@ -69,6 +69,7 @@ For transforming DimensionalData objects:
 
 ```@docs
 groupby
+combine
 DimensionalData.DimGroupByArray
 Bins
 ranges

src/DimensionalData.jl

@@ -83,7 +83,7 @@ export dimnum, hasdim, hasselection, otherdims
 export set, rebuild, reorder, modify, broadcast_dims, broadcast_dims!,
     mergedims, unmergedims, maplayers
 
-export groupby, seasons, months, hours, intervals, ranges
+export groupby, combine, seasons, months, hours, intervals, ranges
 
 
 export @d

src/array/methods.jl

@@ -421,7 +421,7 @@ function _check_cat_lookups(D, ::Regular, lookups...)
             @warn _cat_warn_string(D, "step sizes $(step(span(l))) and $s do not match")
             return false
         end
-        if !(lastval + s ≈ first(l))
+        if !(s isa Dates.AbstractTime) && !(lastval + s ≈ first(l))
             @warn _cat_warn_string(D, "`Regular` lookups do not join with the correct step size: $(lastval) + $s ≈ $(first(l)) should hold")
             return false
         end

src/groupby.jl

@@ -249,7 +249,6 @@ Group some data along the time dimension:
 
 ```jldoctest groupby; setup = :(using Random; Random.seed!(123))
 julia> using DimensionalData, Dates
-
 julia> A = rand(X(1:0.1:20), Y(1:20), Ti(DateTime(2000):Day(3):DateTime(2003)));
 
 julia> groups = groupby(A, Ti => month) # Group by month
@@ -356,6 +355,7 @@ end
 function _group_indices(dim::Dimension, f::Base.Callable; labels=nothing)
     orig_lookup = lookup(dim)
     k1 = f(first(orig_lookup))
+    # TODO: using a Dict here is a bit slow
     indices_dict = Dict{typeof(k1),Vector{Int}}()
     for (i, x) in enumerate(orig_lookup)
          k = f(x)
@@ -447,11 +447,56 @@ end
 
 Generate a `Vector` of `UnitRange` with length `step(A)`
 """
-intervals(rng::AbstractRange) = IntervalSets.Interval{:closed,:open}.(rng, rng .+ step(rng))
+intervals(rng::AbstractRange) =
+    IntervalSets.Interval{:closed,:open}.(rng, rng .+ step(rng))
 
 """
     ranges(A::AbstractRange{<:Integer})
 
 Generate a `Vector` of `UnitRange` with length `step(A)`
 """
 ranges(rng::AbstractRange{<:Integer}) = map(x -> x:x+step(rng)-1, rng)
+
+"""
+    combine(f::Function, gb::DimGroupByArray; dims=:)
+
+Combine the `DimGroupByArray` using function `f` over the group dimensions.
+Unlike broadcasting a reducing function over a `DimGroupByArray`, this function
+always returns a new flattened `AbstractDimArray` even where not all dimensions 
+are reduced. It will also work over grouped `AbstractDimStack`.
+
+If `dims` is given, it will combine only the dimensions in `dims`, the 
+others will be present in the final array. Note that all grouped dimensions
+must be reduced and included in `dims`.
+
+The reducing function `f` must also accept a `dims` keyword.
+
+# Example
+
+```jldoctest groupby
+````
+"""
+function combine(f::Function, gb::DimGroupByArray{G}; dims=:) where G
+    targetdims = DD.commondims(first(gb), dims)
+    all(hasdim(first(gb), targetdims)) || throw(ArgumentError("dims must be a subset of the groupby dimensions"))
+    all(hasdim(targetdims, DD.dims(gb))) || throw(ArgumentError("grouped dimensions $(DD.basedims(gb)) must be included in dims"))
+    # This works for both arrays and stacks
+    # Combine the remaining dimensions after reduction and the group dimensions
+    destdims = (otherdims(DD.dims(first(gb)), dims)..., DD.dims(gb)...)
+    # Get the output eltype 
+    T = Base.promote_op(f, G)
+    # Create a output array with the combined dimensions
+    dest = similar(first(gb), T, destdims)
+    for D in DimIndices(gb)
+        if all(hasdim(targetdims, DD.dims(first(gb))))
+            # Assigned reduced scalar to dest
+            dest[D...] = f(gb[D])
+        else
+            # Reduce with `f` and drop length 1 dimensions
+            xs = dropdims(f(gb[D]; dims); dims)
+            # Broadcast the reduced array to dest
+            broadcast_dims!(identity, view(dest, D...), xs)
+        end
+    end
+    return dest
+end

src/stack/indexing.jl

@@ -150,6 +150,9 @@ for f in (:getindex, :view, :dotview)
     end
 end
 
+@generated function _any_dimarray(v::Union{NamedTuple,Tuple})
+    any(T -> T <: AbstractDimArray, v.types)
+end
 
 #### setindex ####
 @propagate_inbounds Base.setindex!(s::AbstractDimStack, xs, I...; kw...) =
@@ -160,22 +163,17 @@ end
     hassamedims(s) ? _map_setindex!(s, xs, i; kw...) : _setindex_mixed!(s, xs, i; kw...)
 @propagate_inbounds Base.setindex!(s::AbstractDimStack, xs::NamedTuple, i::AbstractArray; kw...) =
     hassamedims(s) ? _map_setindex!(s, xs, i; kw...) : _setindex_mixed!(s, xs, i; kw...)
+@propagate_inbounds Base.setindex!(s::AbstractDimStack, xs::NamedTuple, i::DimensionIndsArrays; kw...) =
+    _map_setindex!(s, xs, i; kw...)
+@propagate_inbounds Base.setindex!(s::AbstractDimStack, xs::NamedTuple, I...; kw...) =
+    _map_setindex!(s, xs, I...; kw...)
 
-@propagate_inbounds function Base.setindex!(
-    s::AbstractDimStack, xs::NamedTuple, I...; kw...
-)
-    map((A, x) -> setindex!(A, x, I...; kw...), layers(s), xs)
-end
-
-_map_setindex!(s, xs, i; kw...) = map((A, x) -> setindex!(A, x, i...; kw...), layers(s), xs)
+_map_setindex!(s, xs, i...; kw...) = map((A, x) -> setindex!(A, x, i...; kw...), layers(s), xs)
 
-_setindex_mixed!(s::AbstractDimStack, x, i::AbstractArray) =
-    map(A -> setindex!(A, x, DimIndices(dims(s))[i]), layers(s))
-_setindex_mixed!(s::AbstractDimStack, i::Integer) =
-    map(A -> setindex!(A, x, DimIndices(dims(s))[i]), layers(s))
-function _setindex_mixed!(s::AbstractDimStack, x, i::Colon)
-    map(DimIndices(dims(s))) do D
-        map(A -> setindex!(A, D), x, layers(s))
+function _setindex_mixed!(s::AbstractDimStack, xs::NamedTuple, i)
+    D = DimIndices(dims(s))[i]
+    map(layers(s), xs) do A, x
+        A[D] = x
     end
 end
 

src/stack/stack.jl

@@ -153,7 +153,6 @@ Base.length(s::AbstractDimStack) = prod(size(s))
 Base.axes(s::AbstractDimStack) = map(first ∘ axes, dims(s))
 Base.axes(s::AbstractDimStack, dims::DimOrDimType) = axes(s, dimnum(s, dims))
 Base.axes(s::AbstractDimStack, dims::Integer) = axes(s)[dims]
-Base.similar(s::AbstractDimStack, args...) = maplayers(A -> similar(A, args...), s)
 Base.eltype(::AbstractDimStack{<:Any,T}) where T = T
 Base.ndims(::AbstractDimStack{<:Any,<:Any,N}) where N = N
 Base.CartesianIndices(s::AbstractDimStack) = CartesianIndices(dims(s))
@@ -197,6 +196,36 @@ Base.get(f::Base.Callable, st::AbstractDimStack, k::Symbol) =
 @propagate_inbounds Base.iterate(st::AbstractDimStack, i) =
     i > length(st) ? nothing : (st[DimIndices(st)[i]], i + 1)
 
+Base.similar(s::AbstractDimStack) = similar(s, eltype(s))
+Base.similar(s::AbstractDimStack, dims::Dimension...) = similar(s, dims)
+Base.similar(s::AbstractDimStack, ::Type{T},dims::Dimension...) where T =
+    similar(s, T, dims)
+Base.similar(s::AbstractDimStack, dims::Tuple{Vararg{Dimension}}) = 
+    similar(s, eltype(s), dims)
+Base.similar(s::AbstractDimStack, ::Type{T}) where T = 
+    similar(s, T, dims(s))
+function Base.similar(s::AbstractDimStack, ::Type{T}, dims::Tuple) where T
+    # Any dims not in the stack are added to all layers
+    ods = otherdims(dims, DD.dims(s))
+    maplayers(s) do A
+        # Original layer dims are maintained, other dims are added
+        D = DD.commondims(dims, (DD.dims(A)..., ods...))
+        similar(A, T, D)
+    end
+end
+function Base.similar(s::AbstractDimStack, ::Type{T}, dims::Tuple) where T<:NamedTuple
+    ods = otherdims(dims, DD.dims(s))
+    maplayers(s, _nt_types(T)) do A, Tx 
+        D = DD.commondims(dims, (DD.dims(A)..., ods...))
+        similar(A, Tx, D)
+    end
+end
+
+@generated function _nt_types(::Type{NamedTuple{K,T}}) where {K,T}
+    expr = Expr(:tuple, T.parameters...)
+    return :(NamedTuple{K}($expr))
+end
+
 # `merge` for AbstractDimStack and NamedTuple.
 # One of the first three arguments must be an AbstractDimStack for dispatch to work.
 Base.merge(s::AbstractDimStack) = s

src/utils.jl

@@ -160,6 +160,11 @@ function broadcast_dims!(f, dest::AbstractDimArray{<:Any,N}, As::AbstractBasicDi
     od = map(A -> otherdims(dest, dims(A)), As)
     return _broadcast_dims_inner!(f, dest, As, od)
 end
+function broadcast_dims!(f, dest::AbstractDimStack, stacks::AbstractDimStack...)
+    maplayers(dest, stacks...) do d, layers...
+        broadcast_dims!(f, d, layers...)
+    end
+end
 
 # Function barrier
 function _broadcast_dims_inner!(f, dest, As, od)

test/groupby.jl

@@ -8,7 +8,6 @@ days = DateTime(2000):Day(1):DateTime(2000, 12, 31)
 A = DimArray((1:6) * (1:366)', (X(1:0.2:2), Ti(days)))
 st = DimStack((a=A, b=A, c=A[X=1]))
 
-
 @testset "group eltype matches indexed values" begin
     da = rand(X(1:10), Y(1:10))
     grps = groupby(da, X => isodd)
@@ -22,10 +21,16 @@ end
         mean(A[Ti=dayofyear(m):dayofyear(m)+daysinmonth(m)-1])
     end
     @test mean.(groupby(A, Ti=>month)) == manualmeans
+    combinedmeans = combine(mean, groupby(A, Ti=>month))
+    @test combinedmeans isa DimArray
+    @test combinedmeans == manualmeans
     manualmeans_st = map(months) do m
         mean(st[Ti=dayofyear(m):dayofyear(m)+daysinmonth(m)-1])
     end
     @test mean.(groupby(st, Ti=>month)) == manualmeans_st
+    combinedmeans_st = combine(mean, groupby(st, Ti=>month))
+    @test combinedmeans_st isa DimStack{(:a, :b, :c), @NamedTuple{a::Float64, b::Float64, c::Float64}}
+    @test collect(combinedmeans_st) == manualmeans_st
 
     manualsums = mapreduce(hcat, months) do m
         vcat(sum(A[Ti=dayofyear(m):dayofyear(m)+daysinmonth(m)-1, X=1 .. 1.5]), 
@@ -36,6 +41,8 @@ end
     @test dims(gb_sum, Ti) == Ti(Sampled([1:12...], ForwardOrdered(), Irregular((nothing, nothing)), Points(), NoMetadata()))
     @test typeof(dims(gb_sum, X)) == typeof(X(Sampled(BitVector([false, true]), ForwardOrdered(), Irregular((nothing, nothing)), Points(), NoMetadata())))
     @test gb_sum == manualsums
+    combined_sum = combine(sum, groupby(A, Ti=>month, X => >(1.5)))
+    @test collect(combined_sum) == manualsums
 
     manualsums_st = mapreduce(hcat, months) do m
         vcat(sum(st[Ti=dayofyear(m):dayofyear(m)+daysinmonth(m)-1, X=1 .. 1.5]), 
@@ -46,10 +53,22 @@ end
     @test dims(gb_sum_st, Ti) == Ti(Sampled([1:12...], ForwardOrdered(), Irregular((nothing, nothing)), Points(), NoMetadata()))
     @test typeof(dims(gb_sum_st, X)) == typeof(X(Sampled(BitVector([false, true]), ForwardOrdered(), Irregular((nothing, nothing)), Points(), NoMetadata())))
     @test gb_sum_st == manualsums_st
+    combined_sum_st = combine(sum, groupby(st, Ti=>month, X => >(1.5)))
+    @test collect(combined_sum_st) == manualsums_st
 
     @test_throws ArgumentError groupby(st, Ti=>month, Y=>isodd)
 end
 
+@testset "partial reductions in combine" begin
+    months = DateTime(2000):Month(1):DateTime(2000, 12, 31)
+    using BenchmarkTools
+    manualmeans = cat(map(months) do m
+        mean(A[Ti=dayofyear(m):dayofyear(m)+daysinmonth(m)-1]; dims=Ti)
+    end...; dims=Ti(collect(1:12)))
+    combinedmeans = combine(mean, groupby(A, Ti()=>month); dims=Ti())
+    @test combinedmeans == manualmeans
+end
+
 @testset "bins" begin
     seasons = DateTime(2000):Month(3):DateTime(2000, 12, 31)
     manualmeans = map(seasons) do s
@@ -59,6 +78,7 @@ end
     @test mean.(groupby(A, Ti=>Bins(month, ranges(1:3:12)))) == manualmeans
     @test mean.(groupby(A, Ti=>Bins(month, intervals(1:3:12)))) == manualmeans
     @test mean.(groupby(A, Ti=>Bins(month, 4))) == manualmeans
+    @test combine(mean, groupby(A, Ti=>Bins(month, ranges(1:3:12)))) == manualmeans
 end
 
 @testset "dimension matching groupby" begin
@@ -75,9 +95,10 @@ end
     end
     @test all(collect(mean.(gb)) .=== manualmeans)
     @test all(mean.(gb) .=== manualmeans)
+    @test all(combine(mean, gb) .=== manualmeans)
 end
 
-@testset "broadcastdims runs after groupby" begin
+@testset "broadcast_dims runs after groupby" begin
     dimlist = (
         Ti(Date("2021-12-01"):Day(1):Date("2022-12-31")),
         X(range(1, 10, length=10)),
@@ -87,7 +108,7 @@ end
     data = rand(396, 10, 15, 2)
     A = DimArray(data, dimlist)
     month_length = DimArray(daysinmonth, dims(A, Ti))
-    g_tempo = DimensionalData.groupby(month_length, Ti=>seasons(; start=December))
+    g_tempo = DimensionalData.groupby(month_length, Ti => seasons(; start=December))
     sum_days = sum.(g_tempo, dims=Ti)
     @test sum_days isa DimArray
     weights = map(./, g_tempo, sum_days)

test/stack.jl

@@ -3,7 +3,7 @@ using DimensionalData, Test, LinearAlgebra, Statistics, ConstructionBase, Random
 using DimensionalData: data
 using DimensionalData: Sampled, Categorical, AutoLookup, NoLookup, Transformed,
     Regular, Irregular, Points, Intervals, Start, Center, End,
-    Metadata, NoMetadata, ForwardOrdered, ReverseOrdered, Unordered, layers, basedims
+    Metadata, NoMetadata, ForwardOrdered, ReverseOrdered, Unordered, layers, basedims, layerdims
 
 A = [1.0 2.0 3.0;
      4.0 5.0 6.0]
@@ -94,11 +94,23 @@ end
     @test all(maplayers(similar(mixed), mixed) do s, m
         dims(s) == dims(m) && dims(s) === dims(m) && eltype(s) === eltype(m)
     end)
-    @test eltype(similar(s, Int)) === @NamedTuple{one::Int, two::Int, three::Int}
+    @test eltype(similar(s, Int)) === 
+    @NamedTuple{one::Int, two::Int, three::Int}
+    @test eltype(similar(s, @NamedTuple{one::Int, two::Float32, three::Bool})) === 
+        @NamedTuple{one::Int, two::Float32, three::Bool}
     st2 = similar(mixed, Bool, x, y)
     @test dims(st2) === (x, y)
     @test dims(st2[:one]) === (x, y)
     @test eltype(st2) === @NamedTuple{one::Bool, two::Bool, extradim::Bool}
+    @test eltype(similar(mixed)) == eltype(mixed)
+    @test size(similar(mixed)) == size(mixed)
+    @test keys(similar(mixed)) == keys(mixed)
+    @test layerdims(similar(mixed)) == layerdims(mixed)
+    xy = (X(), Y()) 
+    @test layerdims(similar(mixed, dims(mixed, (X, Y)))) == (one=xy, two=xy, extradim=xy)
+    st3 = similar(mixed, @NamedTuple{one::Int, two::Float32, extradim::Bool}, (Z([:a, :b, :c]), Ti(1:12), X(1:3)))
+    @test layerdims(st3) == (one=(Ti(), X()), two=(Ti(), X()), extradim=(Z(), Ti(), X()))
+    @test eltype(st3) == @NamedTuple{one::Int, two::Float32, extradim::Bool}
 end
 
 @testset "merge" begin