sub–views of Vcat and Hcat (#72)

* sub–views of Vcat and Hcat

* allow memorylayout for arguments

* more subconcat

* Update concattests.jl

Author: dlfivefifty

src/lazyapplying.jl

@@ -17,6 +17,8 @@ end
 @inline Applied{Style}(A::Applied) where Style = Applied{Style}(A.f, A.args)
 
 arguments(a) = a.args
+arguments(_, a) = a.args
+arguments(a::AbstractArray) = arguments(MemoryLayout(typeof(a)), a)
 
 @inline check_applied_axes(A::Applied) = nothing
 

src/lazyconcat.jl

@@ -550,3 +550,83 @@ for Cat in (:Vcat, :Hcat)
     end
     @eval normp(a::$Cat, p) = norm(norm.(a.args, p), p)
 end
+
+
+###
+# subarrays
+###
+
+subarraylayout(::ApplyLayout{typeof(vcat)}, _) = 
+    ApplyLayout{typeof(vcat)}()
+subarraylayout(::ApplyLayout{typeof(hcat)}, _) = 
+    ApplyLayout{typeof(hcat)}()
+
+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}}) = 
+    view.(arguments(parent(V)), Ref(parentindices(V)[1]), Ref(:))
+
+
+_vcat_lastinds(sz) = _vcat_cumsum(sz...)
+_vcat_firstinds(sz) = (1, (1 .+ most(_vcat_lastinds(sz)))...)
+
+_argsindices(sz) = broadcast(:, _vcat_firstinds(sz), _vcat_lastinds(sz))
+
+_view_vcat(a::Number, kr) = Fill(a,length(kr))
+_view_vcat(a::Number, kr, jr) = Fill(a,length(kr), length(jr))
+_view_vcat(a, kr...) = view(a, kr...)
+
+function arguments(::ApplyLayout{typeof(vcat)}, V::SubArray{<:Any,1})
+    A = parent(V)
+    kr = parentindices(V)[1]
+    sz = size.(arguments(A),1)
+    skr = intersect.(_argsindices(sz), Ref(kr))
+    skr2 = broadcast((a,b) -> a .- b .+ 1, skr, _vcat_firstinds(sz))
+    _view_vcat.(arguments(A), skr2)
+end
+
+function arguments(::ApplyLayout{typeof(vcat)}, V::SubArray{<:Any,2})
+    A = parent(V)
+    kr,jr = parentindices(V)
+    sz = size.(arguments(A),1)
+    skr = intersect.(_argsindices(sz), Ref(kr))
+    skr2 = broadcast((a,b) -> a .- b .+ 1, skr, _vcat_firstinds(sz))
+    _view_vcat.(arguments(A), skr2, Ref(jr))
+end
+
+_view_hcat(a::Number, kr, jr) = Fill(a,length(kr),length(jr))
+_view_hcat(a, kr, jr) = view(a, kr, jr)
+
+function arguments(::ApplyLayout{typeof(hcat)}, V::SubArray{<:Any,2})
+    A = parent(V)
+    kr,jr = parentindices(V)
+    sz = size.(arguments(A),2)
+    sjr = intersect.(_argsindices(sz), Ref(jr))
+    sjr2 = broadcast((a,b) -> a .- b .+ 1, sjr, _vcat_firstinds(sz))
+    _view_hcat.(arguments(A), Ref(kr), sjr2)
+end
+
+function sub_materialize(::ApplyLayout{typeof(vcat)}, V)
+    ret = similar(V)
+    n = 0
+    _,jr = parentindices(V)
+    for a in arguments(V)
+        m = size(a,1)
+        view(ret,n+1:n+m,:) .= a
+        n += m
+    end
+    ret
+end
+    
+function sub_materialize(::ApplyLayout{typeof(hcat)}, V)
+    ret = similar(V)
+    n = 0
+    kr,_ = parentindices(V)
+    for a in arguments(V)
+        m = size(a,2)
+        view(ret,:,n+1:n+m) .= a
+        n += m
+    end
+    ret
+end
+    

src/linalg/add.jl

@@ -63,8 +63,8 @@ for MulAdd_ in [MatMulMatAdd, MatMulVecAdd]
                 B = copy(B)
             end
             _fill_lmul!(β, C)
-            for A in arguments(A)
-                C .= applied(+,applied(*,α, A,B), C)
+            for a in arguments(A)
+                C .= applied(+,applied(*,α, a,B), C)
             end
             C
         end
@@ -86,11 +86,11 @@ end
 ### 
 # views
 ####
-_view(a, b::Tuple) = view(a, b...)
+_view_tuple(a, b::Tuple) = view(a, b...)
 for op in (:+, :-)
     @eval begin
         subarraylayout(a::ApplyLayout{typeof($op)}, _) = a
         arguments(a::SubArray{<:Any,N,<:ApplyArray{<:Any,N,typeof($op)}}) where N =
-            _view.(arguments(parent(a)), Ref(parentindices(a)))
+            _view_tuple.(arguments(parent(a)), Ref(parentindices(a)))
     end
 end

test/concattests.jl

@@ -1,5 +1,5 @@
 using LazyArrays, FillArrays, LinearAlgebra, StaticArrays, Test
-import LazyArrays: MemoryLayout, DenseColumnMajor, PaddedLayout, materialize!, MulAdd, Applied, ApplyLayout
+import LazyArrays: MemoryLayout, DenseColumnMajor, PaddedLayout, materialize!, MulAdd, Applied, ApplyLayout, arguments
 
 @testset "concat" begin
     @testset "Vcat" begin
@@ -354,4 +354,30 @@ import LazyArrays: MemoryLayout, DenseColumnMajor, PaddedLayout, materialize!, M
             @test norm(a,p) ≈ norm(Array(a),p)
         end
     end
+
+    @testset "SubVcat" begin
+        A = Vcat(1,[2,3], Fill(5,10))
+        V = view(A,3:5)
+        @test MemoryLayout(typeof(V)) isa ApplyLayout{typeof(vcat)}
+        VERSION ≥ v"1.1" && @inferred(arguments(V))
+        @test arguments(V)[1] ≡ Fill(1,0)
+        @test A[parentindices(V)...] == copy(V) == Array(A)[parentindices(V)...]
+
+        A = Vcat((1:100)', Zeros(1,100),Fill(1,2,100))
+        V = view(A,:,3:5)
+        @test MemoryLayout(typeof(V)) isa ApplyLayout{typeof(vcat)}
+        @test A[parentindices(V)...] == copy(V) == Array(A)[parentindices(V)...]
+        V = view(A,2:3,3:5)
+        @test MemoryLayout(typeof(V)) isa ApplyLayout{typeof(vcat)}
+        @test A[parentindices(V)...] == copy(V) == Array(A)[parentindices(V)...]
+
+        A = Hcat(1:10, Zeros(10,10))
+        V = view(A,3:5,:)
+        @test MemoryLayout(typeof(V)) isa ApplyLayout{typeof(hcat)}
+        @test A[parentindices(V)...] == copy(V) == Array(A)[parentindices(V)...]
+        V = view(A,3:5,1:4)
+        @test MemoryLayout(typeof(V)) isa ApplyLayout{typeof(hcat)}
+        VERSION ≥ v"1.1" && @inferred(arguments(V))
+        @test arguments(V)[1] == reshape(3:5,3,1)
+    end
 end