Add Interlace

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "0.14.9"
+version = "0.15"
 
 
 [deps]

src/LazyArrays.jl

@@ -60,7 +60,8 @@ end
 
 export Mul, Applied, MulArray, MulVector, MulMatrix, InvMatrix, PInvMatrix,
         Hcat, Vcat, Kron, BroadcastArray, BroadcastMatrix, BroadcastVector, cache, Ldiv, Inv, PInv, Diff, Cumsum,
-        applied, materialize, materialize!, ApplyArray, ApplyMatrix, ApplyVector, apply, ⋆, @~, LazyArray
+        applied, materialize, materialize!, ApplyArray, ApplyMatrix, ApplyVector, apply, ⋆, @~, LazyArray,
+        Interlace, interlace
 
 
 include("lazyapplying.jl")
@@ -71,5 +72,6 @@ include("lazyconcat.jl")
 include("lazysetoperations.jl")
 include("lazyoperations.jl")
 include("lazymacro.jl")
+include("interlace.jl")
 
 end # module

src/interlace.jl

@@ -0,0 +1,49 @@
+struct Interlace{T, N, AA, INDS} <: AbstractArray{T,N}
+    arrays::AA
+    inds::INDS
+end
+
+
+_sortunion(inds...) = sort!(union(inds...))
+function _sortunion(inds::Vararg{AbstractRange,N}) where N
+    all(isequal(N), map(step, inds)) || throw(ArgumentError("incompatible"))
+    sort([map(first, inds)...]) == OneTo(N) || throw(ArgumentError("incompatible"))
+    n = mapreduce(length, +, inds)
+    maximum(map(last, inds)) == n || throw(ArgumentError("incompatible lengths"))
+    OneTo(n)
+end
+
+
+function check_interlace_inds(a, inds)
+    map(length,a) == map(length,inds) || throw(ArgumentError("Lengths must be compatible"))
+    n = mapreduce(length, +, a)
+    _sortunion(inds...) == OneTo(n) ||  throw(ArgumentError("Every index must be mapped to"))
+end
+
+function Interlace(a::NTuple{M,AbstractVector{T}}, inds::NTuple{M,AbstractVector{Int}}) where {T,M} 
+    check_interlace_inds(a, inds)
+    Interlace{T,1,typeof(a), typeof(inds)}(a, inds)
+end
+
+length(A::Interlace) = sum(map(length,A.arrays))
+size(A::Interlace, m) = sum(size.(A.arrays,m))
+size(A::Interlace{<:Any,1}) = (size(A,1),)
+function getindex(A::Interlace{<:Any,1}, k::Integer)
+    for (a,ind) in zip(A.arrays, A.inds)
+        κ = findfirst(isequal(k), ind)
+        isnothing(κ) || return a[something(κ)]
+    end
+    throw(BoundsError(A, k))
+end
+
+function copyto!(dest::AbstractVector, src::Interlace{<:Any,1})
+    for (a,ind) in zip(src.arrays, src.inds)
+        copyto!(view(dest, ind), a)
+    end
+    dest 
+end
+
+Interlace(a::AbstractVector, b::AbstractVector) = 
+    Interlace((a,b), (1:2:(2length(a)-1), 2:2:2length(b)))
+
+interlace(a...) = Array(Interlace(a...))    

src/lazybroadcasting.jl

@@ -29,6 +29,8 @@
     _BroadcastArray(instantiate(Broadcasted{S}(bc.f, _broadcast2broadcastarray(bc.args...))))
 
 BroadcastArray(f, A, As...) = BroadcastArray(broadcasted(f, A, As...))
+BroadcastArray{T,N}(f, A...) where {T,N} = BroadcastArray{T,N,typeof(f),typeof(A)}(f, A)
+
 BroadcastMatrix(f, A...) = BroadcastMatrix(broadcasted(f, A...))
 BroadcastVector(f, A...) = BroadcastVector(broadcasted(f, A...))
 

src/lazyconcat.jl

@@ -289,7 +289,7 @@
 broadcasted(::LazyArrayStyle, op, A::Vcat) =
     Vcat(broadcast(x -> broadcast(op, x), A.args)...)
 
-for Cat in (:Vcat, :Hcat)    
+for Cat in (:Vcat, :Hcat)
     @eval begin
         broadcasted(::LazyArrayStyle, op, A::$Cat, c::Number) =
             $Cat(broadcast((x,y) -> broadcast(op, x, y), A.args, c)...)
@@ -298,7 +298,7 @@
         broadcasted(::LazyArrayStyle, op, A::$Cat, c::Ref) =
             $Cat(broadcast((x,y) -> broadcast(op, x, Ref(y)), A.args, c)...)
         broadcasted(::LazyArrayStyle, op, c::Ref, A::$Cat) =
-            $Cat(broadcast((x,y) -> broadcast(op, Ref(x), y), c, A.args)...)    
+            $Cat(broadcast((x,y) -> broadcast(op, Ref(x), y), c, A.args)...)
     end
 end
 
@@ -318,7 +318,7 @@
     B_arrays = _vcat_getindex_eval(B,kr...)    # evaluate B at same chunks as A
     ApplyVector(vcat, broadcast((a,b) -> broadcast(op,a,b), A.args, B_arrays)...)
 end
- 
+
 function broadcasted(::LazyArrayStyle, op, A::AbstractVector, B::Vcat{<:Any,1})
     kr = _vcat_axes(axes.(B.args)...)
     A_arrays = _vcat_getindex_eval(A,kr...)
@@ -440,20 +440,20 @@
     @eval $op(V::Vcat) = $op($op.(V.args))
 end
 
-function in(x, V::Vcat) 
+function in(x, V::Vcat)
     for a in V.args
         in(x, a) && return true
     end
     false
 end
 
 _fill!(a, x) = fill!(a,x)
-function _fill!(a::Number, x) 
+function _fill!(a::Number, x)
     a == x || throw(ArgumentError("Cannot set $a to $x"))
     a
 end
 
-function fill!(V::Union{Vcat,Hcat}, x) 
+function fill!(V::Union{Vcat,Hcat}, x)
     for a in V.args
         _fill!(a, x)
     end
@@ -546,8 +546,8 @@
     convert(promote_type(T,V), dot(view(a,1:m), view(b,1:m)))
 end
 
-dot(a::CachedArray, b::AbstractArray) = materialize(Dot(a,b)) 
-dot(a::LazyArray, b::AbstractArray) = materialize(Dot(a,b)) 
+dot(a::CachedArray, b::AbstractArray) = materialize(Dot(a,b))
+dot(a::LazyArray, b::AbstractArray) = materialize(Dot(a,b))
 
 
 ###
@@ -571,9 +571,9 @@
 # a row-slice of an Hcat is equivalent to a Vcat
 sublayout(::ApplyLayout{typeof(hcat)}, ::Type{<:Tuple{Number,AbstractVector}}) = ApplyLayout{typeof(vcat)}()
 
-arguments(::ApplyLayout{typeof(vcat)}, V::SubArray{<:Any,2,<:Any,<:Tuple{<:Slice,<:Any}}) = 
+arguments(::ApplyLayout{typeof(vcat)}, V::SubArray{<:Any,2,<:Any,<:Tuple{<:Slice,<:Any}}) =
     view.(arguments(parent(V)), Ref(:), Ref(parentindices(V)[2]))
-arguments(::ApplyLayout{typeof(hcat)}, V::SubArray{<:Any,2,<:Any,<:Tuple{<:Any,<:Slice}}) = 
+arguments(::ApplyLayout{typeof(hcat)}, V::SubArray{<:Any,2,<:Any,<:Tuple{<:Any,<:Slice}}) =
     view.(arguments(parent(V)), Ref(parentindices(V)[1]), Ref(:))
 
 copyto!(dest::AbstractArray{T,N}, src::SubArray{T,N,<:Vcat{T,N}}) where {T,N} = 
@@ -638,7 +638,7 @@
     end
     ret
 end
-    
+
 function sub_materialize(::ApplyLayout{typeof(hcat)}, V)
     ret = similar(V)
     n = 0
@@ -652,10 +652,10 @@
 end
 
 # temporarily allocate. In the future, we add a loop over arguments
-materialize!(M::MatMulMatAdd{<:AbstractColumnMajor,<:ApplyLayout{typeof(vcat)}}) = 
+materialize!(M::MatMulMatAdd{<:AbstractColumnMajor,<:ApplyLayout{typeof(vcat)}}) =
+    materialize!(MulAdd(M.α,M.A,Array(M.B),M.β,M.C))
+materialize!(M::MatMulVecAdd{<:AbstractColumnMajor,<:ApplyLayout{typeof(vcat)}}) =
     materialize!(MulAdd(M.α,M.A,Array(M.B),M.β,M.C))
-materialize!(M::MatMulVecAdd{<:AbstractColumnMajor,<:ApplyLayout{typeof(vcat)}}) = 
-    materialize!(MulAdd(M.α,M.A,Array(M.B),M.β,M.C))    
 
 sublayout(::PaddedLayout{L}, ::Type{I}) where {L,I<:Tuple{AbstractUnitRange}} = 
     PaddedLayout{typeof(sublayout(L(), I))}()
@@ -702,4 +702,4 @@
         κ -= n
     end
     throw(BoundsError(f, (k,j)))
-end
+end

test/interlacetests.jl

@@ -0,0 +1,9 @@
+using LazyArrays, Test
+
+@testset "Interlace" begin
+    @test_throws ArgumentError Interlace(1:5, 10:15)
+    @test_throws ArgumentError Interlace(1:5, 10:12)
+    @test eltype(Interlace(1:5, 10:13)) == Int
+    @test Interlace(1:5, 10:13) == interlace(1:5,10:13) == [1,10,2,11,3,12,4,13,5]
+    @test Interlace(1:5, 10:14) == interlace(1:5,10:14) == [1,10,2,11,3,12,4,13,5,14]
+end