rstack_iter, names, CartesianIndices

Author: Michael Abbott

src/LazyStack.jl

@@ -405,8 +405,11 @@ julia> rstack(1:3, OffsetArray([2.0,2.1,2.2], -1), OffsetArray([3.2,3.3,3.4], +1
 """
 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(list::Union{AbstractArray{<:AbstractArray}, Tuple{Vararg{<:AbstractArray}}}; fill=zero(eltype(first(list))))
+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)
@@ -426,19 +429,39 @@ function rstack(list::Union{AbstractArray{<:AbstractArray}, Tuple{Vararg{<:Abstr
     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))
-        # view(out, axes(item)..., i...) .= item
-        for I in CartesianIndices(item)
-            out[Tuple(I)..., o..., i...] = item[I]
-        end
+        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

@@ -178,23 +178,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
 
@@ -223,7 +206,26 @@ end
     @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], 2:3)) == [1 0; 2 2; 0 3]
-    @test rstack(1:2, OffsetArray([0.1,1], 0:1)) == OffsetArray([0 0.1; 1 1.0; 2 0],-1,0)
+    @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
+@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