support combine_mul for QuasiArrays (#75)

dlfivefifty · web-flow · commit ea10e2ee4373 · 2019-11-01T13:39:31.000Z
* support combine_mul for QuasiArrays

* fill layouts

* Update multests.jl

* Cache broadcasting

* increase cov

* Update cache.jl

* Increase coverage

* Update runtests.jl

* v0.13
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "0.12.4"
+version = "0.13"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -9,7 +9,7 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
-FillArrays = "0.7"
+FillArrays = "0.7,0.8"
 MacroTools = "0.4.5,0.5"
 StaticArrays = "0.8,0.9,0.10,0.11"
 julia = "1"
diff --git a/src/LazyArrays.jl b/src/LazyArrays.jl
@@ -57,10 +57,10 @@ export Mul, Applied, MulArray, MulVector, MulMatrix, InvMatrix, PInvMatrix,
         applied, materialize, materialize!, ApplyArray, ApplyMatrix, ApplyVector, apply, ⋆, @~, LazyArray
 
 include("memorylayout.jl")
-include("cache.jl")
 include("lazyapplying.jl")
 include("lazybroadcasting.jl")
 include("linalg/linalg.jl")
+include("cache.jl")
 include("lazyconcat.jl")
 include("lazysetoperations.jl")
 include("lazyoperations.jl")
diff --git a/src/cache.jl b/src/cache.jl
@@ -73,6 +73,9 @@ function getindex(A::CachedArray, I...)
     A.data[I...]
 end
 
+getindex(A::CachedVector, ::Colon) = copy(A)
+getindex(A::CachedVector, ::Slice) = copy(A)
+
 function getindex(A::CachedVector, I, J...)
     @boundscheck checkbounds(A, I, J...)
     resizedata!(A, _maximum(axes(A,1), I))
@@ -166,4 +169,52 @@ zero!(A::CachedArray{<:Any,N,<:Any,<:Zeros}) where N = zero!(A.data)
 ####
 
 cachedlayout(_, _) = UnknownLayout()
-MemoryLayout(C::Type{CachedArray{T,N,DAT,ARR}}) where {T,N,DAT,ARR} = cachedlayout(MemoryLayout(DAT), MemoryLayout(ARR))
+MemoryLayout(C::Type{CachedArray{T,N,DAT,ARR}}) where {T,N,DAT,ARR} = cachedlayout(MemoryLayout(DAT), MemoryLayout(ARR))
+
+
+
+#####
+# broadcasting
+#
+# We want broadcasting for numbers with concaenations to pass through
+# to take advantage of special implementations of the sub-components
+######
+
+BroadcastStyle(::Type{<:CachedArray{<:Any,N}}) where N = LazyArrayStyle{N}()
+
+broadcasted(::LazyArrayStyle, op, A::CachedArray) =
+    CachedArray(broadcast(op, paddeddata(A)), broadcast(op, A.array))
+
+broadcasted(::LazyArrayStyle, op, A::CachedArray, c::Number) =
+    CachedArray(broadcast(op, paddeddata(A), c), broadcast(op, A.array, c))
+broadcasted(::LazyArrayStyle, op, c::Number, A::CachedArray) =
+CachedArray(broadcast(op, c, paddeddata(A)), broadcast(op, c, A.array))
+broadcasted(::LazyArrayStyle, op, A::CachedArray, c::Ref) =
+    CachedArray(broadcast(op, paddeddata(A), c), broadcast(op, A.array, c))
+broadcasted(::LazyArrayStyle, op, c::Ref, A::CachedArray) =
+    CachedArray(broadcast(op, c, paddeddata(A)), broadcast(op, c, A.array)) 
+
+
+function broadcasted(::LazyArrayStyle, op, A::CachedVector, B::AbstractVector)
+    dat = paddeddata(A)
+    n = length(dat)
+    m = length(B)
+    CachedArray(broadcast(op, dat, view(B,1:n)), broadcast(op, A.array, B))
+end
+
+function broadcasted(::LazyArrayStyle, op, A::AbstractVector, B::CachedVector)
+    dat = paddeddata(B)
+    n = length(dat)
+    m = length(A)
+    CachedArray(broadcast(op, view(A,1:n), dat), broadcast(op, A, B.array))
+end
+
+function broadcasted(::LazyArrayStyle, op, A::CachedVector, B::CachedVector)
+    n = max(A.datasize[1],B.datasize[1])
+    resizedata!(A,n)
+    resizedata!(B,n)
+    Adat = view(paddeddata(A),1:n)
+    Bdat = view(paddeddata(B),1:n)
+    CachedArray(broadcast(op, Adat, Bdat), broadcast(op, A.array, B.array))
+end
+
diff --git a/src/lazyapplying.jl b/src/lazyapplying.jl
@@ -209,13 +209,16 @@ for F in (:exp, :log, :sqrt, :cos, :sin, :tan, :csc, :sec, :cot,
     end
 end
 
-struct LazyLayout <: MemoryLayout end
+abstract type AbstractLazyLayout <: MemoryLayout end
+struct LazyLayout <: AbstractLazyLayout end
 
 
-MemoryLayout(::Type{<:LazyArray}) = LazyLayout()
+MemoryLayout(::Type{<:LazyArray}) = LazyArrayLayout()
 
-transposelayout(::LazyLayout) = LazyLayout()
-conjlayout(::LazyLayout) = LazyLayout()
+transposelayout(L::LazyLayout) = L
+conjlayout(L::LazyLayout) = L
+subarraylayout(L::LazyLayout, _) = L
+reshapedlayout(::LazyLayout, _) = LazyLayout()
 
 combine_mul_styles(::LazyLayout) = LazyArrayApplyStyle()
 result_mul_style(::LazyArrayApplyStyle, ::LazyArrayApplyStyle) = LazyArrayApplyStyle()
diff --git a/src/lazybroadcasting.jl b/src/lazybroadcasting.jl
@@ -41,7 +41,6 @@ Broadcasted(A::BroadcastArray) = instantiate(broadcasted(A.f, A.args...))
 axes(A::BroadcastArray) = axes(Broadcasted(A))
 size(A::BroadcastArray) = map(length, axes(A))
 
-IndexStyle(::BroadcastArray{<:Any,1}) = IndexLinear()
 
 @propagate_inbounds getindex(A::BroadcastArray, kj::Int...) = Broadcasted(A)[kj...]
 
diff --git a/src/linalg/ldiv.jl b/src/linalg/ldiv.jl
@@ -35,6 +35,9 @@ struct Ldiv{StyleA, StyleB, AType, BType}
     B::BType
 end
 
+Ldiv{StyleA, StyleB}(A::AType, B::BType) where {StyleA,StyleB,AType,BType} = 
+    Ldiv{StyleA,StyleB,AType,BType}(A,B)
+
 Ldiv(A::AType, B::BType) where {AType,BType} = 
     Ldiv{typeof(MemoryLayout(AType)),typeof(MemoryLayout(BType)),AType,BType}(A, B)
 
@@ -167,4 +170,11 @@ copy(M::Applied{LdivApplyStyle}) = copy(Ldiv(M))
     materialize(Ldiv(A))[kj...]
 
 
+###
+# * layout
+###
 
+function copy(L::Ldiv{<:Any,ApplyLayout{typeof(*)}}) 
+    args = arguments(L.B)
+    apply(*, L.A \  first(args),  tail(args)...)
+end
diff --git a/src/linalg/mul.jl b/src/linalg/mul.jl
@@ -223,29 +223,48 @@ subarraylayout(::ApplyLayout{typeof(*)}, _...) = ApplyLayout{typeof(*)}()
 
 call(::ApplyLayout{typeof(*)}, V::SubArray) = *
 
-function arguments(::ApplyLayout{typeof(*)}, V::SubArray{<:Any,2})
+function _mat_mul_arguments(V)
     P = parent(V)
     kr, jr = parentindices(V)
     as = arguments(P)
     kjr = intersect.(_mul_args_rows(kr, as...), _mul_args_cols(jr, reverse(as)...))
     view.(as, (kr, kjr...), (kjr..., jr))
 end
 
+
 _vec_mul_view(a...) = view(a...)
 _vec_mul_view(a::AbstractVector, kr, ::Colon) = view(a, kr)
 
-function arguments(::ApplyLayout{typeof(*)}, V::SubArray{<:Any,1})
+function _vec_mul_arguments(V)
     P = parent(V)
     kr, = parentindices(V)
     as = arguments(P)
     kjr = intersect.(_mul_args_rows(kr, as...), _mul_args_cols(Base.OneTo(1), reverse(as)...))
     _vec_mul_view.(as, (kr, kjr...), (kjr..., :))
 end
 
+arguments(::ApplyLayout{typeof(*)}, V::SubArray{<:Any,2}) = _mat_mul_arguments(V)
+arguments(::ApplyLayout{typeof(*)}, V::SubArray{<:Any,1}) = _vec_mul_arguments(V)
+
 @inline sub_materialize(::ApplyLayout{typeof(*)}, V) = apply(*, arguments(V)...)
 @inline copyto!(dest::AbstractArray{T,N}, src::SubArray{T,N,<:ApplyArray{T,N,typeof(*)}}) where {T,N} = 
     copyto!(dest, Applied(src))
 
+##
+# adoint Mul
+##
+
+adjointlayout(::Type, ::ApplyLayout{typeof(*)}) = ApplyLayout{typeof(*)}()
+transposelayout(::ApplyLayout{typeof(*)}) = ApplyLayout{typeof(*)}()
+
+call(::ApplyLayout{typeof(*)}, V::Adjoint) = *
+call(::ApplyLayout{typeof(*)}, V::Transpose) = *
+
+arguments(::ApplyLayout{typeof(*)}, V::Adjoint) = reverse(adjoint.(arguments(V')))
+arguments(::ApplyLayout{typeof(*)}, V::Transpose) = reverse(transpose.(arguments(V')))
+
+
+
 ## 
 # * specialcase
 ##    
diff --git a/src/linalg/muladd.jl b/src/linalg/muladd.jl
@@ -97,10 +97,11 @@ scalarone(::Type{<:AbstractArray{T}}) where T = scalarone(T)
 scalarzero(::Type{T}) where T = zero(T)
 scalarzero(::Type{<:AbstractArray{T}}) where T = scalarzero(T)
 
+fillzeros(::Type{T}, ax) where T = Zeros{T}(ax)
 
 _αAB(M::Mul{MulAddStyle,<:Tuple{<:AbstractArray,<:AbstractArray}}, ::Type{T}) where T = tuple(scalarone(T), M.args...)
 _αAB(M::Mul{MulAddStyle,<:Tuple{<:Number,<:AbstractArray,<:AbstractArray}}, ::Type{T}) where T = M.args
-_αABβC(M::Mul, ::Type{T}) where T = tuple(_αAB(M, T)..., scalarzero(T), Zeros{T}(axes(M)))
+_αABβC(M::Mul, ::Type{T}) where T = tuple(_αAB(M, T)..., scalarzero(T), fillzeros(T,axes(M)))
 
 _βC(M::Mul, ::Type{T}) where T = M.args
 _βC(M::AbstractArray, ::Type{T}) where T = (scalarone(T), M)
@@ -249,7 +250,7 @@ function default_blasmul!(α, A::AbstractMatrix, B::AbstractVector, β, C::Abstr
     z = zero(A[1]*B[1] + A[1]*B[1])
     Astride = size(A, 1) # use size, not stride, since its not pointer arithmetic
 
-    @inbounds for k = 1:mB
+    @inbounds for k in colsupport(B,1)
         aoffs = (k-1)*Astride
         b = B[k]
         for i = 1:mA
diff --git a/src/memorylayout.jl b/src/memorylayout.jl
@@ -145,13 +145,14 @@ dispatch to BLAS and LAPACK routines if the memory layout is compatible.
 @inline MemoryLayout(::Type{<:Number}) = ScalarLayout()
 @inline MemoryLayout(::Type{<:DenseArray}) = DenseColumnMajor()
 
-@inline MemoryLayout(::Type{<:ReinterpretArray{T,N,S,P}}) where {T,N,S,P} = reinterpretedmemorylayout(MemoryLayout(P))
-@inline reinterpretedmemorylayout(::MemoryLayout) = UnknownLayout()
-@inline reinterpretedmemorylayout(::DenseColumnMajor) = DenseColumnMajor()
+@inline MemoryLayout(::Type{<:ReinterpretArray{T,N,S,P}}) where {T,N,S,P} = reinterpretedlayout(MemoryLayout(P))
+@inline reinterpretedlayout(::MemoryLayout) = UnknownLayout()
+@inline reinterpretedlayout(::DenseColumnMajor) = DenseColumnMajor()
 
-@inline MemoryLayout(A::Type{<:ReshapedArray{T,N,P}}) where {T,N,P} = reshapedmemorylayout(MemoryLayout(P))
-@inline reshapedmemorylayout(::MemoryLayout) = UnknownLayout()
-@inline reshapedmemorylayout(::DenseColumnMajor) = DenseColumnMajor()
+
+MemoryLayout(::Type{<:ReshapedArray{T,N,A,DIMS}}) where {T,N,A,DIMS} = reshapedlayout(MemoryLayout(A), DIMS)
+@inline reshapedlayout(_, _) = UnknownLayout()
+@inline reshapedlayout(::DenseColumnMajor, _) = DenseColumnMajor()
 
 
 @inline MemoryLayout(A::Type{<:SubArray{T,N,P,I}}) where {T,N,P,I} = 
@@ -442,4 +443,10 @@ struct EyeLayout <: MemoryLayout end
 
 MemoryLayout(::Type{<:AbstractFill}) = FillLayout()
 MemoryLayout(::Type{<:Zeros}) = ZerosLayout()
-diagonallayout(::ML) where ML<:AbstractFillLayout = DiagonalLayout{ML}()
+diagonallayout(::ML) where ML<:AbstractFillLayout = DiagonalLayout{ML}()
+# all sub arrays are same
+subarraylayout(L::AbstractFillLayout, inds::Type) = L
+reshapedlayout(L::AbstractFillLayout, _) = L
+adjointlayout(::Type, L::AbstractFillLayout) = L
+transposelayout(L::AbstractFillLayout) = L
+
diff --git a/test/lazymultests.jl b/test/lazymultests.jl
@@ -187,5 +187,10 @@ LinearAlgebra.factorize(A::MyLazyArray) = factorize(A.data)
         @test MemoryLayout(typeof(Diagonal(x))) isa DiagonalLayout{LazyLayout}
         @test MemoryLayout(typeof(Diagonal(ApplyArray(+,x,x)))) isa DiagonalLayout{LazyLayout}
         @test MemoryLayout(typeof(Diagonal(1:6))) isa DiagonalLayout{UnknownLayout}
+
+        @test MemoryLayout(typeof(A')) isa LazyLayout
+        @test MemoryLayout(typeof(transpose(A))) isa LazyLayout
+        @test MemoryLayout(typeof(view(A,1:2,1:2))) isa LazyLayout
+        @test MemoryLayout(typeof(reshape(A,4))) isa LazyLayout
     end
 end
diff --git a/test/ldivtests.jl b/test/ldivtests.jl
@@ -155,4 +155,11 @@ import Base.Broadcast: materialize
         dest = fill(NaN,5)
         @test copyto!(dest, applied(\,Q,b)) == apply(\,Q,b)
     end
+
+    @testset "Lazy *" begin
+        A = randn(5,5)
+        B = randn(5,5)
+        b = 1:5
+        @test apply(\,B,ApplyArray(*,A,b)) == B\A*b
+    end
 end
diff --git a/test/memorylayouttests.jl b/test/memorylayouttests.jl
@@ -171,8 +171,8 @@ struct FooNumber <: Number end
         @test @inferred(MemoryLayout(typeof(Vcat(Ones(10),Zeros(10))))) == PaddedLayout{FillLayout}()
         @test @inferred(MemoryLayout(typeof(Vcat([1.],Zeros(10))))) == PaddedLayout{DenseColumnMajor}()
 
-        @test MemoryLayout(typeof(view(Fill(1,10),1:3))) == UnknownLayout()
-        @test MemoryLayout(typeof(view(Fill(1,10),1:3,1))) == UnknownLayout()
+        @test MemoryLayout(typeof(view(Fill(1,10),1:3))) == FillLayout()
+        @test MemoryLayout(typeof(view(Fill(1,10),1:3,1))) == FillLayout()
     end
 
     @testset "BroadcastArray" begin
diff --git a/test/multests.jl b/test/multests.jl
@@ -560,7 +560,8 @@ end
         A = randn(Float64,20,22)
         B = randn(ComplexF64,22,24)
         C = similar(B,20,24)
-        @test all((C .= Mul(A,B)  ) .=== copyto!(similar(C), MulAdd(1.0, A, B, 0.0, C)) .=== A*B)
+        @test all((C .= Mul(A,B)  ) .=== copyto!(similar(C), MulAdd(1.0, A, B, 0.0, C)))
+        @test copyto!(similar(C), MulAdd(1.0, A, B, 0.0, C)) ≈ A*B
     end
 
     @testset "no allocation" begin
@@ -1063,4 +1064,15 @@ end
         @test (2.0\A) == 2.0\Array(A)
         @test A/2.0 == Array(A)/2.0
     end
+
+    @testset "* trans" begin
+        A = ApplyArray(*,[1 2; 3 4], Vcat(Fill(1,1,3),Fill(2,1,3)))
+        V = view(A,1:2,1:2)
+        @test MemoryLayout(typeof(V)) isa ApplyLayout{typeof(*)}
+        @test V == ApplyArray(V)
+        @test MemoryLayout(typeof(V')) isa ApplyLayout{typeof(*)}
+        @test V' == ApplyArray(V')
+        @test MemoryLayout(typeof(transpose(V))) isa ApplyLayout{typeof(*)}
+        @test transpose(V) == ApplyArray(transpose(V))
+    end
 end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -72,6 +72,7 @@ end
     @test cache(C) isa CachedArray{Int,1,Vector{Int},UnitRange{Int}}
     C2 = cache(C)
     @test C2.data !== C.data
+    @test C[:] == C[Base.Slice(Base.OneTo(10))] == C
 
     @test cache(A)[2,1] == 2
     @test_throws BoundsError cache(A)[2,2]
@@ -135,7 +136,31 @@ end
 
     @testset "colsupport past size" begin
         C = cache(Zeros(5,5)); C[5,1]; 
+        @test colsupport(C,1) == Base.OneTo(5)
         @test colsupport(C,3) == 1:0
+        @test rowsupport(C,1) == Base.OneTo(1)
+    end
+
+    @testset "broadcast" begin
+        x = CachedArray([1,2,3],1:8);
+        y = 1:8;
+        f = (x,y) -> cos(x*y)
+        @test f.(x,y) isa CachedArray
+        @test @inferred(broadcast(f,x,y)) == f.(Vector(x), Vector(y))
+
+        @test (x + y) isa CachedArray
+        @test (x + y).array isa AbstractRange
+        @test (x + y) == Vector(x) + Vector(y)
+
+        z = CachedArray([1,4],Zeros{Int}(8));
+        @test (x .+ z) isa CachedArray
+        @test (x + z) isa CachedArray
+        @test Vector( x .+ z) == Vector( x + z) == Vector(x) + Vector(z)
+
+        # Lazy mixed with Static treats as Lazy
+        s = SVector(1,2,3,4,5,6,7,8)
+        @test f.(x , s) isa CachedArray
+        @test f.(x , s) == f.(Vector(x), Vector(s))
     end
 end
 
@@ -202,4 +227,23 @@ end
     A = ApplyArray(tril,randn(2,2),-1)
     @test A isa ApplyArray{Float64}
     @test A == tril(A.args[1],-1)
+end
+
+@testset "BroadcastArray" begin
+    bc = broadcasted(exp,[1,2,3])
+    v = BroadcastArray(exp, [1,2,3])
+    @test BroadcastArray(bc) == BroadcastVector(bc) == BroadcastVector{Float64,typeof(exp),typeof(bc.args)}(bc) ==
+        v == BroadcastVector(exp, [1,2,3]) == exp.([1,2,3])
+
+    Base.IndexStyle(typeof(BroadcastVector(exp, [1,2,3]))) == IndexLinear()
+    
+    bc = broadcasted(exp,[1 2; 3 4])
+    M = BroadcastArray(exp, [1 2; 3 4])
+    @test BroadcastArray(bc) == BroadcastMatrix(bc) == BroadcastMatrix{Float64,typeof(exp),typeof(bc.args)}(bc) ==
+        M == BroadcastMatrix(BroadcastMatrix(bc)) == BroadcastMatrix(exp,[1 2; 3 4]) == exp.([1 2; 3 4])
+    
+    @test exp.(v') isa BroadcastMatrix
+    @test exp.(transpose(v)) isa BroadcastMatrix
+    @test exp.(M') isa BroadcastMatrix
+    @test exp.(transpose(M)) isa BroadcastMatrix
 end
