Merge branch 'master' into ragged

Author: Michael Abbott

.github/workflows/TagBot.yml

@@ -0,0 +1,11 @@
+name: TagBot
+on:
+  schedule:
+    - cron: 0 * * * *
+jobs:
+  TagBot:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: JuliaRegistries/TagBot@v1
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}

Project.toml

@@ -1,7 +1,7 @@
 name = "LazyStack"
 uuid = "1fad7336-0346-5a1a-a56f-a06ba010965b"
 authors = ["Michael Abbott"]
-version = "0.0.5"
+version = "0.0.7"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -11,8 +11,8 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
-NamedDims = "0.2.12"
-OffsetArrays = "0.11, 0.12, 0.13"
+NamedDims = "0.2.16"
+OffsetArrays = "0.11, 0.12, 0.13, 1.0"
 ZygoteRules = "0.1, 0.2, 0.3"
 julia = "1"
 

README.md

@@ -4,20 +4,29 @@
 
 This package exports one function, `stack`, for turning a list of arrays 
 into one `AbstractArray`. Given several arrays with the same `eltype`, 
-or an array of such arrays, it returns a lazy `Stacked{T,N}` view of these. 
+or an array of such arrays, it returns a lazy `Stacked{T,N}` view of these:
 
 ```julia
 stack([zeros(2,2), ones(2,2)])  # isa Stacked{Float64, 3, <:Vector{<:Matrix}}
 stack([1,2,3], 4:6)             # isa Stacked{Int, 2, <:Tuple{<:Vector, <:UnitRange}}
 ```
 
 Given a generator, it instead iterates through the elements and writes into a new array.
-(This is lazy only in that it need not `collect` the generator first.)
-The same method is also used for any list of arrays of heterogeneous element type.
+Given a function and then some arrays, it behaves like `map(f, A, B)` but immediately writes
+into a new array:
 
 ```julia
 stack([i,2i] for i in 1:5)            # isa Matrix{Int}     # size(ans) == (2, 5)
+stack(*, eachcol(ones(2,4)), 1:4)     # == Matrix(stack(map(*, eachcol(...), 1:4)))
+```
+
+The same `stack_iter` method is also used for any list of arrays of heterogeneous element type,
+and for arrays of tuples. Notice that like `map(identity, Any[1, 1.0, 5im])`, this promotes using 
+`promote_typejoin`, to `Number` here, rather than to `Complex{Float64}`.
+
+```julia
 stack([1,2], [3.0, 4.0], [5im, 6im])  # isa Matrix{Number}  # size(ans) == (2, 3)
+stack([(i,2.0,3//j) for i=1:4, j=1:5])# isa Array{Real, 3}  # size(ans) == (3, 4, 5)
 ```
 
 The slices must all have the same `size`, but they (and the container) 
@@ -33,6 +42,8 @@ This one plays well with [OffsetArrays.jl](https://github.com/JuliaArrays/Offset
 Besides which, there are several other ways to achieve similar things:
 
 * For an array of arrays, you can also use [`JuliennedArrays.Align`](https://bramtayl.github.io/JuliennedArrays.jl/latest/#JuliennedArrays.Align). This requires (or enables) you to specify which dimensions of the output belong to the sub-arrays, instead of writing `PermutedDimsArray(stack(...), ...)`.
-* There is also [`RecursiveArrayTools.VectorOfArray`](https://github.com/JuliaDiffEq/RecursiveArrayTools.jl#vectorofarray) which as its name hints only allows a one-dimensional container.
+* There is also [`RecursiveArrayTools.VectorOfArray`](https://github.com/JuliaDiffEq/RecursiveArrayTools.jl#vectorofarray) which as its name hints only allows a one-dimensional container. Linear indexing retreives a slice, not an element, which is sometimes surprising.
 * For a tuple of arrays, [`LazyArrays.Hcat`](https://github.com/JuliaArrays/LazyArrays.jl#concatenation) is at present faster to index than `stack`, but doesn't allow arbitrary dimensions.
 * For a generator of arrays, the built-in `reduce(hcat,...)` may work, but it slow compared to `stack`: see [test/speed.jl](test/speed.jl) for some examples.
+
+The package [ArraysOfArrays.jl](https://oschulz.github.io/ArraysOfArrays.jl/stable/#section_ArrayOfSimilarArrays-1) solves the opposite problem, of accessing one large array as if it were many slices. As does [`JuliennedArrays.Slices`](https://bramtayl.github.io/JuliennedArrays.jl/latest/#JuliennedArrays.Slices-Union{Tuple{NumberOfDimensions},%20Tuple{Item},%20Tuple{AbstractArray{Item,NumberOfDimensions},Vararg{Int64,N}%20where%20N}}%20where%20NumberOfDimensions%20where%20Item), and of course [`Base.eachslice`](https://docs.julialang.org/en/v1/base/arrays/#Base.eachslice).

src/LazyStack.jl

@@ -2,7 +2,7 @@ module LazyStack
 
 export stack, rstack
 
-#===== Overloads =====#
+#===== Tuples =====#
 
 ndims(A) = Base.ndims(A)
 ndims(::Tuple) = 1
@@ -12,6 +12,13 @@ size(A) = Base.size(A)
 size(t::Tuple) = tuple(length(t))
 size(t::NamedTuple) = tuple(length(t))
 
+similar_vector(x::AbstractArray, n::Int) = similar(x, n::Int)
+similar_vector(x::Tuple, n::Int) = Vector{eltype(x)}(undef, n::Int)
+similar_vector(x::NamedTuple, n::Int) = Vector{eltype(x)}(undef, n::Int)
+
+# eltype(x::Tuple) = Base.promote_type(x...) # to match choice below
+# eltype(x::NamedTuple) = Base.promote_type(x...)
+
 #===== Slices =====#
 
 """
@@ -56,6 +63,7 @@ stack(xs::AbstractArray{T}...) where {T} = stack(xs)
 
 function stack_slices(xs::AT, ::Val{T}, ::Val{N}) where {T,N,AT}
     length(xs) >= 1 || throw(ArgumentError("stacking an empty collection is not allowed"))
+    storage_type(first(xs)) <: Union{Array, AbstractRange} || return stack_iter(xs)
     s = size(first(xs))
     for x in xs
         size(x) == s || throw(DimensionMismatch(
@@ -83,7 +91,7 @@ outer_ndims(x::Stacked{T,N,<:Tuple}) where {T,N} = 1
 
 inner_ndims(x::Stacked) = ndims(x) - outer_ndims(x)
 
-@inline function Base.getindex(x::Stacked{T}, inds::Integer...) where {T}
+@inline function Base.getindex(x::Stacked{T,N}, inds::Vararg{Integer,N}) where {T,N}
     @boundscheck checkbounds(x, inds...)
     IN, ON = inner_ndims(x), outer_ndims(x)
     outer = @inbounds getindex(x.slices, ntuple(d -> inds[d+IN], ON)...)
@@ -98,6 +106,14 @@ Base.collect(x::Stacked{T,2,<:AbstractArray{<:AbstractArray{T,1}}}) where {T} =
 
 Base.view(x::Stacked{T,2,<:AbstractArray{<:AbstractArray{T,1}}}, ::Colon, i::Int) where {T} = x.slices[i]
 
+function Base.push!(x::Stacked{T,N,<:AbstractVector}, y::AbstractArray) where {T,N}
+    s = size(first(x.slices))
+    isempty(y) || size(y) == s || throw(DimensionMismatch(
+            "slices being stacked must share a common size. Expected $s, got $(size(y))"))
+    push!(x.slices, y)
+    x
+end
+
 function Base.showarg(io::IO, x::Stacked, toplevel)
     print(io, "stack(")
     Base.showarg(io, parent(x), false)
@@ -111,12 +127,20 @@ ITERS = [:Flatten, :Drop, :Filter]
 for iter in ITERS
     @eval ndims(::Iterators.$iter) = 1
 end
+ndims(gen::Base.Generator) = ndims(gen.iter)
+ndims(zed::Iterators.Zip) = maximum(ndims, zed.is)
+if VERSION < v"1.1"
+    ndims(zed::Iterators.Zip2) = max(ndims(zed.a), ndims(zed.b))
+end
+
+similar_vector(x, n::Int) = throw(ArgumentError())
 
 """
     stack(::Generator)
     stack(::Array{T}, ::Array{S}, ...)
 
-This constructs a new array. Can handle inconsistent eltypes, but not inconsistent sizes.
+This constructs a new array, calling `stack_iter`.
+Can handle inconsistent eltypes, but not inconsistent sizes.
 
 ```
 julia> stack([i i; i 10i] for i in 1:2:3)
@@ -166,10 +190,9 @@ function stack_iter(itr)
 
     w, val = vstack_plus(itr)
 
-    z = reshape(w, size(val)..., outsize...)::Array
+    z = reshape(w, size(val)..., outsize...)
 
-    z′ = maybe_add_offsets(z, val)
-    maybe_add_names(z′, val)
+    rewrap_like(z, val)
 end
 
 vstack(itr) = first(vstack_plus(itr))
@@ -178,12 +201,13 @@ function vstack_plus(itr)
     zed = iterate(itr)
     zed === nothing && throw(ArgumentError("stacking an empty collection is not allowed"))
     val, state = zed
+    val isa Union{AbstractArray, Tuple, NamedTuple} || return collect(itr), val
 
     s = size(val)
     n = Base.haslength(itr) ? prod(s)*length(itr) : nothing
 
-    v = Vector{eltype(val)}(undef, something(n, prod(s)))
-    copyto!(v, 1, no_offsets(val), 1, prod(s))
+    v = similar_vector(no_wraps(val), something(n, prod(s)))
+    copyto!(v, 1, no_wraps(val), 1, prod(s))
 
     w = stack_rest(v, 0, n, s, itr, state)::Vector
     w, val
@@ -200,9 +224,9 @@ function stack_rest(v, i, n, s, itr, state)
         i += 1
         if eltype(val) <: eltype(v)
             if n isa Int
-                copyto!(v, i*prod(s)+1, no_offsets(val), 1, prod(s))
+                copyto!(v, i*prod(s)+1, no_wraps(val), 1, prod(s))
             else
-                append!(v, vec(no_offsets(val)))
+                append!(v, vec(no_wraps(val)))
             end
         else
 
@@ -212,9 +236,9 @@ function stack_rest(v, i, n, s, itr, state)
             copyto!(v′, v)
 
             if n isa Int
-                copyto!(v′, i*prod(s)+1, no_offsets(val), 1, prod(s))
+                copyto!(v′, i*prod(s)+1, no_wraps(val), 1, prod(s))
             else
-                append!(v′, vec(no_offsets(val)))
+                append!(v′, vec(no_wraps(val)))
             end
 
             return stack_rest(v′, i, n, s, itr, state)
@@ -223,53 +247,85 @@ function stack_rest(v, i, n, s, itr, state)
     end
 end
 
+"""
+    stack(fun, iters...)
+    stack(eachcol(A), eachslice(B, dims=3)) do a, b
+        f(a,b)
+    end
+
+If the first argument is a function, then this is mapped over the other argumenst.
+The result should be `== stack(map(fun, iters...))`, but done using `stack_iter`
+to return a dense `Array` instead of a `Stacked` object.
+"""
+stack(fun::Function, iter) = stack(fun(arg) for arg in iter)
+function stack(fun::Function, iters...)
+    if all(Base.haslength, iters)
+        sz = size(first(iters))
+        all(a -> size(a)==sz, iters) || throw(DimensionMismatch(
+            "sizes of all argumens must match, in stack(f, A, B, C). " *
+            "This is slightly stricter than map(f, A, B, C), for now."))
+    end
+    stack(fun(args...) for args in zip(iters...))
+end
+
 #===== Offset =====#
 
 using OffsetArrays
 
-no_offsets(a) = a
-no_offsets(a::OffsetArray) = parent(a)
+no_wraps(a) = a
+no_wraps(a::OffsetArray) = parent(a)
 
-maybe_add_offsets(A, a) = A
-maybe_add_offsets(A, a::OffsetArray) = OffsetArray(A, axes(a)..., axes(A, ndims(A)))
+rewrap_like(A, a) = A
+function rewrap_like(A, a::OffsetArray)
+    B = rewrap_like(A, parent(a))
+    OffsetArray(B, axes(a)..., axes(A, ndims(A)))
+end
 
 #===== NamedDims =====#
 
 using NamedDims
 
+ensure_named(a::AbstractArray, L::Tuple) = NamedDimsArray(a, L)
+ensure_named(a::NamedDimsArray, L::Tuple) = refine_names(a, L)
+
 # array of arrays
 stack(xs::NamedDimsArray{<:Any,<:AbstractArray}) =
-    NamedDimsArray(stack(parent(xs)), getnames(xs))
+    ensure_named(stack(parent(xs)), getnames(xs))
 stack(x::AT) where {AT <: AbstractArray{<:NamedDimsArray{L,T,IN},ON}} where {T,IN,ON,L} =
-    NamedDimsArray(Stacked{T, IN+ON, AT}(x), getnames(x))
+    ensure_named(Stacked{T, IN+ON, AT}(x), getnames(x))
 
 getnames(xs::AbstractArray{<:AbstractArray}) =
     (dimnames(eltype(xs))..., dimnames(xs)...)
 
 # tuple of arrays
 stack(x::AT) where {AT <: Tuple{Vararg{NamedDimsArray{L,T,IN}}}} where {T,IN,L} =
-    NamedDimsArray(Stacked{T, IN+1, AT}(x), getnames(x))
+    ensure_named(stack(map(parent, x)), getnames(x))
 
 getnames(xs::Tuple{Vararg{<:NamedDimsArray}}) =
     (dimnames(first(xs))..., :_)
 
 # generators
+#=
 function stack(xs::Base.Generator{<:NamedDimsArray{L}}) where {L}
     w = stack_iter(xs)
     l = (ntuple(_ -> :_, ndims(w)-length(L))..., L...)
-    NamedDimsArray(w, l)
+    ensure_named(w, l)
 end
 
 function stack(xs::Base.Generator{<:Iterators.ProductIterator{<:Tuple{<:NamedDimsArray}}})
     w = stack_iter(xs)
     L = Tuple(Iterators.flatten(map(dimnames, ms.iter.iterators)))
     l = (ntuple(_ -> :_, ndims(w)-length(L))..., L...)
-    NamedDimsArray(w, l)
+    ensure_named(w, l)
 end
+=#
 
-maybe_add_names(A, a) = A
-maybe_add_names(A, a::NamedDimsArray{L}) where {L} =
-    NamedDimsArray(A, (L..., ntuple(_ -> :_, ndims(A) - ndims(a))...))
+function rewrap_like(A, a::NamedDimsArray{L}) where {L}
+    B = rewrap_like(A, parent(a))
+    ensure_named(B, (L..., ntuple(_ -> :_, ndims(A) - ndims(a))...))
+end
+
+no_wraps(a::NamedDimsArray) = no_wraps(parent(a))
 
 """
     stack(name, things...)
@@ -281,7 +337,7 @@ this will be the name of the last dimension of the resulting `NamedDimsArray`.
 function LazyStack.stack(s::Symbol, args...)
     data = stack(args...)
     name_last = ntuple(d -> d==ndims(data) ? s : :_, ndims(data))
-    NamedDimsArray(data, name_last)
+    ensure_named(data, name_last)
 end
 
 #===== Zygote =====#
@@ -300,6 +356,18 @@ end
     stack(gen), Δ -> error("not yet!")
 end
 
+@adjoint function Base.collect(x::Stacked)
+    collect(x), tuple
+end
+
+#===== CuArrays =====#
+# Send these to stack_iter, by testing  storage_type(first(xs)) <: Array
+
+function storage_type(x::AbstractArray)
+    p = parent(x)
+    typeof(x) === typeof(p) ? typeof(x) : storage_type(p)
+end
+
 #===== Ragged =====#
 
 """
@@ -372,4 +440,5 @@ end
 tupleindices(t::Tuple) = ((i,) for i in 1:length(t))
 tupleindices(A::AbstractArray) = (Tuple(I) for I in CartesianIndices(A))
 
+
 end # module