Improve Lmul (#61)

* Move Lmul to separate file

* axis, similar, copy for Lmul

* use Lmul for Diagonal

* Add Rmul

* Update multests.jl

* Update lmul.jl

* Add Rmul for Triangular

* Fix UPLO in Triangular mul

* Update multests.jl

* v0.12, BroadcsatMatrix macro

* improvements for broadcast array

* Simplify ApplyStyle

* Update memorylayouttests.jl

* Update lazybroadcasting.jl

* Update memorylayouttests.jl

* improve coverage

* Improve coverage of Rmul

* fix broadcast row/colsupport

* don't instantiate F in apply/broadcastlayout

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "0.11.1"
+version = "0.12"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

src/LazyArrays.jl

@@ -52,7 +52,7 @@ else
 end             
 
 export Mul, Applied, MulArray, MulVector, MulMatrix, InvMatrix, PInvMatrix,
-        Hcat, Vcat, Kron, BroadcastArray, cache, Ldiv, Inv, PInv, Diff, Cumsum,
+        Hcat, Vcat, Kron, BroadcastArray, BroadcastMatrix, BroadcastVector, cache, Ldiv, Inv, PInv, Diff, Cumsum,
         applied, materialize, materialize!, ApplyArray, ApplyMatrix, ApplyVector, apply, ⋆, @~, LazyArray
 
 include("memorylayout.jl")

src/lazyapplying.jl

@@ -141,6 +141,9 @@ array.
 """
 abstract type LazyArray{T,N} <: AbstractArray{T,N} end
 
+const LazyMatrix{T} = LazyArray{T,2}
+const LazyVector{T} = LazyArray{T,1}
+
 struct ApplyArray{T, N, F, Args<:Tuple} <: LazyArray{T,N}
     f::F
     args::Args
@@ -214,10 +217,14 @@ result_mul_style(::LazyArrayApplyStyle, _) = LazyArrayApplyStyle()
 result_mul_style(_, ::LazyArrayApplyStyle) = LazyArrayApplyStyle()
 
 
-struct  ApplyLayout{F, LAY} <: MemoryLayout end
+struct  ApplyLayout{F} <: MemoryLayout end
+
+applylayout(::Type{F}, args...) where F = ApplyLayout{F}()
 
-MemoryLayout(M::Type{Applied{Style,F,Args}}) where {Style,F,Args} = ApplyLayout{F,tuple_type_memorylayouts(Args)}()
-MemoryLayout(M::Type{ApplyArray{T,N,F,Args}}) where {T,N,F,Args} = ApplyLayout{F,tuple_type_memorylayouts(Args)}()
+MemoryLayout(::Type{Applied{Style,F,Args}}) where {Style,F,Args} = 
+    applylayout(F, tuple_type_memorylayouts(Args)...)
+MemoryLayout(::Type{ApplyArray{T,N,F,Args}}) where {T,N,F,Args} = 
+    applylayout(F, tuple_type_memorylayouts(Args)...)
 
 function show(io::IO, A::Applied) 
     print(io, "Applied(", A.f)
@@ -230,7 +237,7 @@ end
 applybroadcaststyle(_1, _2) = DefaultArrayStyle{2}()
 BroadcastStyle(M::Type{<:ApplyArray}) = applybroadcaststyle(M, MemoryLayout(M))
 
-Base.replace_in_print_matrix(A::ApplyMatrix, i::Integer, j::Integer, s::AbstractString) =
+Base.replace_in_print_matrix(A::LazyMatrix, i::Integer, j::Integer, s::AbstractString) =
     i in colsupport(A,j) ? s : Base.replace_with_centered_mark(s)
 
 ### 
@@ -263,5 +270,6 @@ end
 @inline getindex(A::ApplyMatrix, kr::AbstractUnitRange, jr::AbstractUnitRange) = lazy_getindex(A, kr, jr)
 
 
-diagonallayout(::LazyLayout) = LazyLayout()
-diagonallayout(::ApplyLayout) = DiagonalLayout{LazyLayout}()
+diagonallayout(::LazyLayout) = DiagonalLayout{LazyLayout}()
+diagonallayout(::ApplyLayout) = DiagonalLayout{LazyLayout}()
+

src/lazybroadcasting.jl

@@ -18,17 +18,25 @@ BroadcastArray{T,N}(bc::Broadcasted{Style,Axes,F,Args}) where {T,N,Style,Axes,F,
 BroadcastArray{T}(bc::Broadcasted{<:Union{Nothing,BroadcastStyle},<:Tuple{Vararg{Any,N}},<:Any,<:Tuple}) where {T,N} =
     BroadcastArray{T,N}(bc)
 
+BroadcastVector(bc::Broadcasted) = BroadcastVector{combine_eltypes(bc.f, bc.args)}(bc)  
+BroadcastMatrix(bc::Broadcasted) = BroadcastMatrix{combine_eltypes(bc.f, bc.args)}(bc)
+
 _broadcast2broadcastarray(a, b...) = tuple(a, b...)
 _broadcast2broadcastarray(a::Broadcasted, b...) = tuple(BroadcastArray(a), b...)
 
 _BroadcastArray(bc::Broadcasted) = BroadcastArray{combine_eltypes(bc.f, bc.args)}(bc)
 BroadcastArray(bc::Broadcasted{S}) where S =
     _BroadcastArray(instantiate(Broadcasted{S}(bc.f, _broadcast2broadcastarray(bc.args...))))
-BroadcastArray(b::BroadcastArray) = b
+
 BroadcastArray(f, A, As...) = BroadcastArray(broadcasted(f, A, As...))
+BroadcastMatrix(f, A...) = BroadcastMatrix(broadcasted(f, A...))
+BroadcastVector(f, A...) = BroadcastVector(broadcasted(f, A...))
 
-Broadcasted(A::BroadcastArray) = instantiate(broadcasted(A.f, A.args...))
+BroadcastArray(b::BroadcastArray) = b
+BroadcastVector(A::BroadcastVector) = A
+BroadcastMatrix(A::BroadcastMatrix) = A
 
+Broadcasted(A::BroadcastArray) = instantiate(broadcasted(A.f, A.args...))
 
 axes(A::BroadcastArray) = axes(Broadcasted(A))
 size(A::BroadcastArray) = map(length, axes(A))
@@ -73,26 +81,33 @@ end
 
 
 BroadcastStyle(::Type{<:BroadcastArray{<:Any,N}}) where N = LazyArrayStyle{N}()
+BroadcastStyle(::Type{<:Adjoint{<:Any,<:BroadcastVector{<:Any}}}) where N = LazyArrayStyle{2}()
+BroadcastStyle(::Type{<:Transpose{<:Any,<:BroadcastVector{<:Any}}}) where N = LazyArrayStyle{2}()
+BroadcastStyle(::Type{<:Adjoint{<:Any,<:BroadcastMatrix{<:Any}}}) where N = LazyArrayStyle{2}()
+BroadcastStyle(::Type{<:Transpose{<:Any,<:BroadcastMatrix{<:Any}}}) where N = LazyArrayStyle{2}()
 BroadcastStyle(L::LazyArrayStyle{N}, ::StaticArrayStyle{N}) where N = L
 BroadcastStyle(::StaticArrayStyle{N}, L::LazyArrayStyle{N})  where N = L
 
 """
-    BroadcastLayout(f, layouts)
+    BroadcastLayout{F}()
 
 is returned by `MemoryLayout(A)` if a matrix `A` is a `BroadcastArray`.
-`f` is a function that broadcast operation is applied and `layouts` is
-a tuple of `MemoryLayout` of the broadcasted arguments.
+`F` is the typeof function that broadcast operation is applied.
 """
-struct BroadcastLayout{F, LAY} <: MemoryLayout end
-
-tuple_type_memorylayouts(::Type{I}) where I<:Tuple = Tuple{typeof.(MemoryLayout.(I.parameters))...}
-tuple_type_memorylayouts(::Type{Tuple{A}}) where {A} = Tuple{typeof(MemoryLayout(A))}
-tuple_type_memorylayouts(::Type{Tuple{A,B}}) where {A,B} = Tuple{typeof(MemoryLayout(A)),typeof(MemoryLayout(B))}
-tuple_type_memorylayouts(::Type{Tuple{A,B,C}}) where {A,B,C} = Tuple{typeof(MemoryLayout(A)),typeof(MemoryLayout(B)),typeof(MemoryLayout(C))}
-
+struct BroadcastLayout{F} <: MemoryLayout end
+
+tuple_type_memorylayouts(::Type{I}) where I<:Tuple = MemoryLayout.(I.parameters)
+tuple_type_memorylayouts(::Type{Tuple{A}}) where {A} = (MemoryLayout(A),)
+tuple_type_memorylayouts(::Type{Tuple{A,B}}) where {A,B} = (MemoryLayout(A),MemoryLayout(B))
+tuple_type_memorylayouts(::Type{Tuple{A,B,C}}) where {A,B,C} = (MemoryLayout(A),MemoryLayout(B),MemoryLayout(C))
+
+broadcastlayout(::Type{F}, _...) where F = BroadcastLayout{F}()
+broadcastlayout(::Type, ::LazyLayout...) = LazyLayout()
+broadcastlayout(::Type, _, ::LazyLayout) = LazyLayout()
+broadcastlayout(::Type, _, _, ::LazyLayout) = LazyLayout()
+broadcastlayout(::Type, _, _, _, ::LazyLayout) = LazyLayout()
 MemoryLayout(::Type{BroadcastArray{T,N,F,Args}}) where {T,N,F,Args} = 
-    BroadcastLayout{F,tuple_type_memorylayouts(Args)}()
-
+    broadcastlayout(F, tuple_type_memorylayouts(Args)...)
 ## scalar-range broadcast operations ##
 # Ranges already support smart broadcasting
 for op in (+, -, big)
@@ -134,3 +149,18 @@ broadcasted(::LazyArrayStyle{N}, ::typeof(*), a::Zeros{T,N}, b::AbstractArray{V,
     broadcast(DefaultArrayStyle{N}(), *, a, b)
 
 diagonallayout(::BroadcastLayout) = DiagonalLayout{LazyLayout}()    
+
+
+###
+# support
+###
+
+_broadcast_colsupport(sz, A::Number, j) = OneTo(sz[1])
+_broadcast_colsupport(sz, A::AbstractVector, j) = colsupport(A,j)
+_broadcast_colsupport(sz, A::AbstractMatrix, j) = size(A,1) == 1 ? OneTo(sz[1]) : colsupport(A,j)
+_broadcast_rowsupport(sz, A::Number, j) = OneTo(sz[2])
+_broadcast_rowsupport(sz, A::AbstractVector, j) = OneTo(sz[2])
+_broadcast_rowsupport(sz, A::AbstractMatrix, j) = size(A,2) == 1 ? OneTo(sz[2]) : rowsupport(A,j)
+
+colsupport(::BroadcastLayout{typeof(*)}, A, j) = intersect(_broadcast_colsupport.(Ref(size(A)), A.args, j)...)
+rowsupport(::BroadcastLayout{typeof(*)}, A, j) = intersect(_broadcast_rowsupport.(Ref(size(A)), A.args, j)...)

src/lazyconcat.jl

@@ -436,7 +436,7 @@ end
 # *
 ###
 
-function materialize!(M::MatMulVecAdd{<:ApplyLayout{typeof(hcat)},<:ApplyLayout{typeof(vcat)}})
+function materialize!(M::MatMulVecAdd{ApplyLayout{typeof(hcat)},ApplyLayout{typeof(vcat)}})
     α,A,B,β,C =  M.α,M.A,M.B,M.β,M.C
     T = eltype(C)
     _fill_lmul!(β,C) # this is temporary until strong β = false is supported
@@ -446,7 +446,7 @@ function materialize!(M::MatMulVecAdd{<:ApplyLayout{typeof(hcat)},<:ApplyLayout{
     C
  end
 
- function materialize!(M::MatMulMatAdd{<:ApplyLayout{typeof(hcat)},<:ApplyLayout{typeof(vcat)}})
+ function materialize!(M::MatMulMatAdd{ApplyLayout{typeof(hcat)},ApplyLayout{typeof(vcat)}})
     α,A,B,β,C =  M.α,M.A,M.B,M.β,M.C
     T = eltype(C)
     _fill_lmul!(β,C) # this is temporary until strong β = false is supported
@@ -462,4 +462,13 @@ function materialize!(M::MatMulVecAdd{<:ApplyLayout{typeof(hcat)},<:ApplyLayout{
 
 
  most(a) = reverse(tail(reverse(a)))
-colsupport(M::Vcat, j) = first(colsupport(first(M.args),j)):(size(Vcat(most(M.args)...),1)+last(colsupport(last(M.args),j)))
+colsupport(M::Vcat, j) = first(colsupport(first(M.args),j)):(size(Vcat(most(M.args)...),1)+last(colsupport(last(M.args),j)))
+
+
+
+###
+# padded
+####
+
+struct PaddedLayout{L} <: MemoryLayout end
+applylayout(::Type{typeof(vcat)}, ::A, ::ZerosLayout) where A = PaddedLayout{A}()

src/linalg/add.jl

@@ -47,9 +47,9 @@ end
 _fill_lmul!(β, A::AbstractArray{T}) where T = iszero(β) ? zero!(A) : lmul!(β, A)
 combine_mul_styles(::ApplyLayout{typeof(+)}) = IdentityMulStyle()
 for MulAdd_ in [MatMulMatAdd, MatMulVecAdd]
-    # `MulAdd{<:ApplyLayout{typeof(+)}}` cannot "win" against
+    # `MulAdd{ApplyLayout{typeof(+)}}` cannot "win" against
     # `MatMulMatAdd` and `MatMulVecAdd` hence `@eval`:
-    @eval function materialize!(M::$MulAdd_{<:ApplyLayout{typeof(+)}})
+    @eval function materialize!(M::$MulAdd_{ApplyLayout{typeof(+)}})
         α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
         if C ≡ B
             B = copy(B)

src/linalg/factorizations.jl

@@ -14,14 +14,15 @@ function check_mul_axes(A::AbstractQ, B, C...)
     check_mul_axes(B, C...)
 end
 
+copy(M::Lmul{QLayout}) = copyto!(similar(M), M)
 
 function copyto!(dest::AbstractArray{T}, M::Lmul{QLayout}) where T
     A,B = M.A,M.B
     if size(dest,1) == size(B,1) 
         copyto!(dest, B)
     else
         copyto!(view(dest,1:size(B,1),:), B)
-        fill!(@view(dest[size(B,1)+1:end,:]), zero(T))
+        zero!(@view(dest[size(B,1)+1:end,:]))
     end
     materialize!(Lmul(A,dest))
 end

src/linalg/inv.jl

@@ -178,8 +178,14 @@ size(A::PInvMatrix) = map(length, axes(A))
 @propagate_inbounds getindex(A::InvMatrix{T}, k::Int, j::Int) where T =
     (parent(A)\[Zeros(j-1); one(T); Zeros(size(A,2) - j)])[k]
 
-mulapplystyle(::ApplyLayout{typeof(inv),Tuple{A}}, B) where A = ldivapplystyle(A, B)
-mulapplystyle(::ApplyLayout{typeof(pinv),Tuple{A}}, B) where A = ldivapplystyle(A, B)
+struct InvLayout{L} <: MemoryLayout end
+struct PInvLayout{L} <: MemoryLayout end
+
+applylayout(::Type{typeof(inv)}, ::A) where A = InvLayout{A}()
+applylayout(::Type{typeof(pinv)}, ::A) where A = PInvLayout{A}()
+
+mulapplystyle(::InvLayout{A}, B) where A = ldivapplystyle(A, B)
+mulapplystyle(::PInvLayout{A}, B) where A = ldivapplystyle(A, B)
 
 
 @inline function Ldiv(M::Mul)

src/linalg/lazymul.jl

@@ -28,6 +28,13 @@ macro lazymul(Typ)
             Base.:*(A::LazyArrays.ApplyMatrix, B::$Typ, C...) = LazyArrays.apply(*,A,B,C...)
         end
     end
+    if Typ ≠ :BroadcastMatrix && Typ ≠ :ApplyMatrix
+        ret = quote
+            $ret
+            Base.:*(A::$Typ, B::LazyArrays.BroadcastMatrix, C...) = LazyArrays.apply(*,A,B,C...)
+            Base.:*(A::LazyArrays.BroadcastMatrix, B::$Typ, C...) = LazyArrays.apply(*,A,B,C...)
+        end
+    end
     for Struc in (:AbstractTriangular, :Diagonal)
         ret = quote
             $ret
@@ -60,6 +67,18 @@ macro lazymul(Typ)
             Base.:*(A::LinearAlgebra.AbstractTriangular, B::$Mod{<:Any,<:$Typ}, C...) = LazyArrays.apply(*,A,B, C...)
             Base.:*(A::$Mod{<:Any,<:$Typ}, B::LinearAlgebra.AbstractTriangular, C...) = LazyArrays.apply(*,A,B, C...)
         end
+        if Typ ≠ :ApplyMatrix
+            ret = quote
+                $ret
+                Base.:*(A::$Mod{<:Any,<:$Typ}, B::ApplyMatrix, C...) = LazyArrays.apply(*,A,B, C...)
+            end
+        end
+        if Typ ≠ :BroadcastMatrix && Typ ≠ :ApplyMatrix
+            ret = quote
+                $ret
+                Base.:*(A::$Mod{<:Any,<:$Typ}, B::BroadcastMatrix, C...) = LazyArrays.apply(*,A,B, C...)
+            end
+        end
     end
 
     esc(ret)

src/linalg/linalg.jl

@@ -1,6 +1,7 @@
 include("mul.jl")
 include("lazymul.jl")
 include("muladd.jl")
+include("lmul.jl")
 include("inv.jl")
 include("add.jl")
 include("factorizations.jl")