Merge pull request #8 from mcabbott/ragged

Ragged stack

Author: mcabbott

src/LazyStack.jl

@@ -1,13 +1,20 @@
 module LazyStack
 
-export stack
+export stack, rstack
 
 #===== Tuples =====#
 
 ndims(A) = Base.ndims(A)
 ndims(::Tuple) = 1
 ndims(::NamedTuple) = 1
 
+axes(A) = Base.axes(A)
+axes(A, d) = Base.axes(A, d)
+if VERSION < v"1.1" # because, on Julia 1.0, axes((1,2)) === Base.OneTo(2)
+    axes(t::Tuple) = tuple(Base.axes(t))
+end
+axes(nt::NamedTuple) = tuple(Base.OneTo(length(nt)))
+
 size(A) = Base.size(A)
 size(t::Tuple) = tuple(length(t))
 size(t::NamedTuple) = tuple(length(t))
@@ -80,9 +87,6 @@ Base.size(x::Stacked) = (size(first(x.slices))..., size(x.slices)...)
 Base.size(x::Stacked{T,N,<:Tuple}) where {T,N} = (size(first(x.slices))..., length(x.slices))
 
 Base.axes(x::Stacked) = (axes(first(x.slices))..., axes(x.slices)...)
-if VERSION < v"1.1" # axes((1:9, 1:9)) == Base.OneTo(2) # on Julia 1.0
-    Base.axes(x::Stacked{T,N,<:Tuple}) where {T,N} = (axes(first(x.slices))..., axes(x.slices))
-end
 
 Base.parent(x::Stacked) = x.slices
 
@@ -368,4 +372,100 @@ function storage_type(x::AbstractArray)
     typeof(x) === typeof(p) ? typeof(x) : storage_type(p)
 end
 
+#===== Ragged =====#
+
+"""
+    rstack(arrays; fill=0)
+
+Ragged `stack`, which allows slices of varying size, and fills the gaps with zero
+or the given `fill`. Always returns an `Array`.
+
+```
+julia> rstack(1:n for n in 1:5)
+5×5 Array{Int64,2}:
+ 1  1  1  1  1
+ 0  2  2  2  2
+ 0  0  3  3  3
+ 0  0  0  4  4
+ 0  0  0  0  5
+
+julia> rstack([[1,2,3], [10,20.0]], fill=missing)
+3×2 Array{Union{Missing, Float64},2}:
+ 1.0  10.0
+ 2.0  20.0
+ 3.0    missing
+
+julia> using OffsetArrays
+
+julia> rstack(1:3, OffsetArray([2.0,2.1,2.2], -1), OffsetArray([3.2,3.3,3.4], +1))
+5×3 OffsetArray(::Array{Real,2}, 0:4, 1:3) with eltype Real with indices 0:4×1:3:
+ 0  2.0  0
+ 1  2.1  0
+ 2  2.2  3.2
+ 3  0    3.3
+ 0  0    3.4
+```
+
+"""
+rstack(x::AbstractArray, ys::AbstractArray...; kw...) = rstack((x, ys...); kw...)
+rstack(g::Base.Generator; kw...) = rstack(collect(g); kw...)
+rstack(f::Function, ABC...; kw...) = rstack(map(f, ABC...); kw...)
+rstack(list::AbstractArray{<:AbstractArray}; fill=zero(eltype(first(list)))) = rstack_iter(list; fill=fill)
+rstack(list::Tuple{Vararg{<:AbstractArray}}; fill=zero(eltype(first(list)))) = rstack_iter(list; fill=fill)
+
+function rstack_iter(list; fill)
+    T = mapreduce(eltype, Base.promote_typejoin, list, init=typeof(fill))
+    # T = mapreduce(eltype, Base.promote_type, list, init=typeof(fill))
+    N = maximum(ndims, list)
+    ax = ntuple(N) do d
+        hi = maximum(x -> last(axes(x,d)), list)
+        if all(x -> axes(x,d) isa Base.OneTo, list)
+            Base.OneTo(hi)
+        else
+            lo = minimum(x -> first(axes(x,d)), list)
+            lo:hi
+        end
+    end
+    arr = Array{T}(undef, map(length, ax)..., size(list)...)
+    fill!(arr, fill)
+    out = if ax isa Tuple{Vararg{Base.OneTo}}
+        arr
+    else
+        OffsetArray(arr, (ax..., axes(list)...))
+    end
+    z = rstack_copyto!(out, list, Val(N))
+
+    rewrap_names(z, first(list)) # now I want to separate names & offsets again!
+end
+
+function rstack_copyto!(out, list, ::Val{N}) where {N}
+    for i in tupleindices(list)
+        item = list[i...]
+        o = ntuple(_->1, N - ndims(item))
+        out[CartesianIndices(axes(item)), o..., i...] .= item
+
+        # https://github.com/JuliaArrays/OffsetArrays.jl/issues/100
+        # view(out, axes(item)..., o..., i...) .= item
+
+        # for I in CartesianIndices(item)
+        #     out[Tuple(I)..., o..., i...] = item[I]
+        # end
+    end
+    out
+end
+
+tupleindices(t::Tuple) = ((i,) for i in 1:length(t))
+tupleindices(A::AbstractArray) = (Tuple(I) for I in CartesianIndices(A))
+
+rewrap_names(A, a) = A
+function rewrap_names(A, a::NamedDimsArray{L}) where {L}
+    B = rewrap_names(A, parent(a))
+    ensure_named(B, (L..., ntuple(_ -> :_, ndims(A) - ndims(a))...))
+end
+function rstack(s::Symbol, args...)
+    data = rstack(args...)
+    name_last = ntuple(d -> d==ndims(data) ? s : :_, ndims(data))
+    ensure_named(data, name_last)
+end
+
 end # module

test/runtests.jl

@@ -18,6 +18,7 @@ using OffsetArrays, NamedDims
     @test stack(v34)[1,1,1] == v34[1][1] # trailing dims
     @test stack(v34) * ones(4) ≈ hcat(v34...) * ones(4) # issue #6
     @test stack(v34) * ones(4,2) ≈ hcat(v34...) * ones(4,2)
+    @test axes(stack(v34)) === axes(stack(v34...)) === axes(stack(v34[i] for i in 1:4))
 
 end
 @testset "tuples" begin
@@ -132,6 +133,9 @@ end
     @test axes(stack(oout)) == (1:3, 11:14)
     @test axes(copy(stack(oout))) ==  (1:3, 11:14)
 
+    oboth = OffsetArray(oin, 11:14)
+    @test axes(stack(oboth)) == (3:5, 11:14)
+
     ogen = (OffsetArray([3,4,5], 3:5) for i in 1:4)
     @test axes(stack(ogen)) == (3:5, 1:4)
 
@@ -178,23 +182,6 @@ end
 
     @test_throws DimensionMismatch push!(stack([rand(2)]), rand(3))
 
-end
-@info "loading Zygote"
-using Zygote
-@testset "zygote" begin
-
-    @test Zygote.gradient((x,y) -> sum(stack(x,y)), ones(2), ones(2)) == ([1,1], [1,1])
-    @test Zygote.gradient((x,y) -> sum(stack([x,y])), ones(2), ones(2)) == ([1,1], [1,1])
-
-    f399(x) = sum(stack(x) * sum(x))
-    f399c(x) = sum(collect(stack(x)) * sum(x))
-    @test Zygote.gradient(f399, [ones(2), ones(2)]) == ([[4,4], [4,4]],)
-    @test Zygote.gradient(f399c, [ones(2), ones(2)]) == ([[4,4], [4,4]],)
-    ftup(x) = sum(stack(x...) * sum(x))
-    ftupc(x) = sum(collect(stack(x...)) * sum(x))
-    @test Zygote.gradient(ftup, (ones(2), ones(2))) == (([4,4], [4,4]),)
-    @test Zygote.gradient(ftupc, (ones(2), ones(2))) == (([4,4], [4,4]),)
-
 end
 @testset "readme" begin
 
@@ -218,3 +205,46 @@ end
     @test LazyStack.vstack(g234) == reduce(vcat, collect(g234))
 
 end
+@testset "ragged" begin
+
+    @test rstack([1,2], 1:3) == [1 1; 2 2; 0 3]
+    @test rstack([[1,2], 1:3], fill=99) == [1 1; 2 2; 99 3]
+
+    @test rstack(1:2, OffsetArray([2,3], +1)) == [1 0; 2 2; 0 3]
+    @test rstack(1:2, OffsetArray([0.1,1], -1)) == OffsetArray([0 0.1; 1 1.0; 2 0],-1,0)
+
+    @test dimnames(rstack(:b, NamedDimsArray(1:2, :a), OffsetArray([2,3], +1))) == (:a, :b)
+
+end
+@testset "tuple functions" begin
+
+    @test LazyStack.ndims([1,2]) == 1
+    @test LazyStack.ndims((1,2)) == 1
+    @test LazyStack.ndims((a=1,b=2)) == 1
+
+    @test LazyStack.size([1,2]) == (2,)
+    @test LazyStack.size((1,2)) == (2,)
+    @test LazyStack.size((a=1,b=2)) == (2,)
+
+    @test LazyStack.axes([1,2]) == (1:2,)
+    @test LazyStack.axes((1,2)) == (1:2,)
+    @test LazyStack.axes((a=1,b=2)) == (1:2,)
+
+end
+@info "loading Zygote"
+using Zygote
+@testset "zygote" begin
+
+    @test Zygote.gradient((x,y) -> sum(stack(x,y)), ones(2), ones(2)) == ([1,1], [1,1])
+    @test Zygote.gradient((x,y) -> sum(stack([x,y])), ones(2), ones(2)) == ([1,1], [1,1])
+
+    f399(x) = sum(stack(x) * sum(x))
+    f399c(x) = sum(collect(stack(x)) * sum(x))
+    @test Zygote.gradient(f399, [ones(2), ones(2)]) == ([[4,4], [4,4]],)
+    @test Zygote.gradient(f399c, [ones(2), ones(2)]) == ([[4,4], [4,4]],)
+    ftup(x) = sum(stack(x...) * sum(x))
+    ftupc(x) = sum(collect(stack(x...)) * sum(x))
+    @test Zygote.gradient(ftup, (ones(2), ones(2))) == (([4,4], [4,4]),)
+    @test Zygote.gradient(ftupc, (ones(2), ones(2))) == (([4,4], [4,4]),)
+
+end