BroadcastStyle for lazy triangular or HermOrSym, improved broadcast * (#360)

* BroadcastStyle for lazy triangular or HermOrSym

* Update inv.jl

* Update broadcasttests.jl

* broadcast mul with / and \

* Update LazyArraysBandedMatricesExt.jl

* Add tests and back compat

* Update lazybroadcasting.jl

* simplifiable for broadcast mul

* Update LazyArraysBandedMatricesExt.jl

* increase coverage

* Update broadcasttests.jl

* Update cachetests.jl

* increase coverage

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "2.3.2"
+version = "2.4"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

ext/LazyArraysBandedMatricesExt.jl

@@ -10,7 +10,7 @@ import LazyArrays: sublayout, symmetriclayout, hermitianlayout, applylayout, cac
                    AbstractPaddedLayout, PaddedLayout, AbstractLazyBandedLayout, LazyBandedLayout, PaddedRows,
                    PaddedColumns, CachedArray, CachedMatrix, LazyLayout, BroadcastLayout, ApplyLayout,
                    paddeddata, resizedata!, broadcastlayout, _broadcastarray2broadcasted, _broadcast_sub_arguments,
-                   arguments, call, applybroadcaststyle, simplify, simplifiable, islazy_layout, lazymaterialize, _broadcast_mul_mul,
+                   arguments, call, applybroadcaststyle, simplify, simplifiable, islazy_layout, lazymaterialize, _broadcast_mul_mul, _broadcast_mul_simplifiable,
                    triangularlayout, AbstractCachedMatrix, _mulbanded_copyto!, ApplyBandedLayout, BroadcastBandedLayout
 import Base: BroadcastStyle, similar, copy, broadcasted, getindex, OneTo, oneto, tail, sign, abs
 import BandedMatrices: bandedbroadcaststyle, bandwidths, isbanded, bandedcolumns, bandeddata, BandedStyle,
@@ -539,6 +539,7 @@ bandeddata(R::ApplyMatrix{<:Any,typeof(rot180)}) = @view(bandeddata(arguments(R)
 const BandedLazyLayouts = Union{AbstractLazyBandedLayout, BandedColumns{LazyLayout}, BandedRows{LazyLayout},
     TriangularLayout{UPLO,UNIT,BandedRows{LazyLayout}} where {UPLO,UNIT},
     TriangularLayout{UPLO,UNIT,BandedColumns{LazyLayout}} where {UPLO,UNIT},
+    TriangularLayout{UPLO,UNIT,LazyBandedLayout} where {UPLO,UNIT},
     SymTridiagonalLayout{LazyLayout}, BidiagonalLayout{LazyLayout}, TridiagonalLayout{LazyLayout},
     SymmetricLayout{BandedColumns{LazyLayout}}, HermitianLayout{BandedColumns{LazyLayout}}}
 
@@ -551,7 +552,12 @@ copy(M::Mul{<:LazyLayouts, <:BandedLazyLayouts}) = simplify(M)
 copy(M::Mul{<:Any, <:BandedLazyLayouts}) = simplify(M)
 copy(M::Mul{<:BandedLazyLayouts, <:AbstractLazyLayout}) = simplify(M)
 copy(M::Mul{<:AbstractLazyLayout, <:BandedLazyLayouts}) = simplify(M)
-copy(M::Mul{BroadcastLayout{typeof(*)}, <:BandedLazyLayouts}) = _broadcast_mul_mul(arguments(BroadcastLayout{typeof(*)}(), M.A), M.B)
+for op in (:*, :/, :\)
+    @eval begin
+        simplifiable(M::Mul{BroadcastLayout{typeof($op)}, <:BandedLazyLayouts}) = _broadcast_mul_simplifiable($op, arguments(BroadcastLayout{typeof($op)}(), M.A), M.B)
+        copy(M::Mul{BroadcastLayout{typeof($op)}, <:BandedLazyLayouts}) = _broadcast_mul_mul($op, arguments(BroadcastLayout{typeof($op)}(), M.A), M.B)
+    end
+end
 copy(M::Mul{<:BandedLazyLayouts, <:DiagonalLayout}) = simplify(M)
 copy(M::Mul{<:DiagonalLayout, <:BandedLazyLayouts}) = simplify(M)
 

src/lazybroadcasting.jl

@@ -137,6 +137,11 @@ BroadcastStyle(::Type{<:Transpose{<:Any,<:LazyVector}}) = LazyArrayStyle{2}()
 BroadcastStyle(::Type{<:Adjoint{<:Any,<:LazyMatrix}}) = LazyArrayStyle{2}()
 BroadcastStyle(::Type{<:Transpose{<:Any,<:LazyMatrix}}) = LazyArrayStyle{2}()
 BroadcastStyle(::Type{<:SubArray{<:Any,1,<:LazyMatrix,<:Tuple{Slice,Any}}}) = LazyArrayStyle{1}()
+
+BroadcastStyle(::Type{<:UpperOrLowerTriangular{<:Any,<:LazyMatrix}}) = LazyArrayStyle{2}()
+BroadcastStyle(::Type{<:LinearAlgebra.HermOrSym{<:Any,<:LazyMatrix}}) = LazyArrayStyle{2}()
+
+
 BroadcastStyle(L::LazyArrayStyle{N}, ::StructuredMatrixStyle)  where N = L
 
 
@@ -397,11 +402,37 @@ end
 ###
 
 _broadcast_mul_mul(A, B) = simplify(Mul(broadcast(*, A...), B))
-_broadcast_mul_mul((a,B)::Tuple{AbstractVector,AbstractMatrix}, C) = a .* (B*C)
-_broadcast_mul_mul((A,b)::Tuple{AbstractMatrix,AbstractVector}, C) = b .* (A*C)
-@inline copy(M::Mul{BroadcastLayout{typeof(*)}}) = _broadcast_mul_mul(arguments(BroadcastLayout{typeof(*)}(), M.A), M.B)
-@inline copy(M::Mul{BroadcastLayout{typeof(*)},<:AbstractLazyLayout}) = _broadcast_mul_mul(arguments(BroadcastLayout{typeof(*)}(), M.A), M.B)
-@inline copy(M::Mul{BroadcastLayout{typeof(*)},ApplyLayout{typeof(*)}}) = _broadcast_mul_mul(arguments(BroadcastLayout{typeof(*)}(), M.A), M.B)
+_broadcast_mul_mul(::typeof(*), A, B) = _broadcast_mul_mul(A, B) # maintain back-compatibility with Quasi/ContiuumArrays.jl
+_broadcast_mul_simplifiable(op, A, B) = Val(false)
+_broadcast_mul_mul(op, A, B) = simplify(Mul(broadcast(op, A...), B))
+
+for op in (:*, :\)
+    @eval begin
+        _broadcast_mul_simplifiable(::typeof($op), (a,B)::Tuple{Union{AbstractVector,Number},AbstractMatrix}, C) = simplifiable(*, B, C)
+        _broadcast_mul_mul(::typeof($op), (a,B)::Tuple{Union{AbstractVector,Number},AbstractMatrix}, C) = broadcast($op, a, (B*C))
+    end
+end
+
+for op in (:*, :/)
+    @eval begin
+        _broadcast_mul_simplifiable(::typeof($op), (A,b)::Tuple{AbstractMatrix,Union{AbstractVector,Number}}, C) = simplifiable(*, A, C)
+        _broadcast_mul_mul(::typeof($op), (A,b)::Tuple{AbstractMatrix,Union{AbstractVector,Number}}, C) = broadcast($op, (A*C), b)
+    end
+end
+
+
+
+for op in (:*, :/, :\)
+    @eval begin
+        @inline simplifiable(M::Mul{BroadcastLayout{typeof($op)}}) = _broadcast_mul_simplifiable($op, arguments(BroadcastLayout{typeof($op)}(), M.A), M.B)
+        @inline simplifiable(M::Mul{BroadcastLayout{typeof($op)},<:LazyLayouts}) = _broadcast_mul_simplifiable($op, arguments(BroadcastLayout{typeof($op)}(), M.A), M.B)
+        @inline simplifiable(M::Mul{BroadcastLayout{typeof($op)},ApplyLayout{typeof(*)}}) = _broadcast_mul_simplifiable($op, arguments(BroadcastLayout{typeof($op)}(), M.A), M.B)
+        @inline copy(M::Mul{BroadcastLayout{typeof($op)}}) = _broadcast_mul_mul($op, arguments(BroadcastLayout{typeof($op)}(), M.A), M.B)
+        @inline copy(M::Mul{BroadcastLayout{typeof($op)},<:LazyLayouts}) = _broadcast_mul_mul($op, arguments(BroadcastLayout{typeof($op)}(), M.A), M.B)
+        @inline copy(M::Mul{BroadcastLayout{typeof($op)},ApplyLayout{typeof(*)}}) = _broadcast_mul_mul($op, arguments(BroadcastLayout{typeof($op)}(), M.A), M.B)
+    end
+end
+
 
 for op in (:*, :\, :/)
     @eval begin

src/linalg/inv.jl

@@ -88,10 +88,13 @@ applylayout(::Type{typeof(pinv)}, ::A) where A = PInvLayout{A}()
 simplifiable(::Mul{<:AbstractInvLayout}) = Val(true)
 
 copy(M::Mul{<:AbstractInvLayout}) = ArrayLayouts.ldiv(pinv(M.A), M.B)
-copy(M::Mul{<:AbstractInvLayout,<:AbstractLazyLayout}) = ArrayLayouts.ldiv(pinv(M.A), M.B)
-@inline copy(M::Mul{<:AbstractInvLayout,<:DiagonalLayout{<:AbstractFillLayout}}) = copy(mulreduce(M))
-@inline copy(M::Mul{<:AbstractInvLayout,ApplyLayout{typeof(*)}}) = simplify(M)
+copy(M::Mul{<:AbstractInvLayout, <:AbstractLazyLayout}) = ArrayLayouts.ldiv(pinv(M.A), M.B)
+@inline copy(M::Mul{<:AbstractInvLayout, <:DiagonalLayout{<:AbstractFillLayout}}) = copy(mulreduce(M))
+@inline copy(M::Mul{<:AbstractInvLayout, ApplyLayout{typeof(*)}}) = simplify(M)
 copy(L::Ldiv{<:AbstractInvLayout}) = pinv(L.A) * L.B
+copy(L::Ldiv{<:AbstractInvLayout, <:AbstractLazyLayout}) = pinv(L.A) * L.B
+copy(L::Ldiv{<:AbstractInvLayout, <:AbstractInvLayout}) = pinv(L.A) * L.B
+copy(L::Ldiv{<:AbstractInvLayout, ApplyLayout{typeof(*)}}) = pinv(L.A) * L.B
 Ldiv(A::Applied{<:Any,typeof(\)}) = Ldiv(A.args...)
 
 

src/padded.jl

@@ -510,6 +510,9 @@ copy(M::Mul{<:Union{TriangularLayout{'U', 'N', <:AbstractLazyLayout}, Triangular
 simplifiable(::Mul{<:Union{TriangularLayout{'U', 'N', <:AbstractLazyLayout}, TriangularLayout{'U', 'U', <:AbstractLazyLayout}}, <:Union{PaddedColumns,PaddedLayout}}) = Val(true)
 
 
+@inline simplifiable(M::Mul{BroadcastLayout{typeof(*)},<:Union{PaddedColumns,PaddedLayout}}) = simplifiable(Mul{BroadcastLayout{typeof(*)},UnknownLayout}(M.A,M.B))
+@inline copy(M::Mul{BroadcastLayout{typeof(*)},<:Union{PaddedColumns,PaddedLayout}}) = copy(Mul{BroadcastLayout{typeof(*)},UnknownLayout}(M.A,M.B))
+
 simplifiable(::Mul{<:DualLayout{<:AbstractLazyLayout}, <:Union{PaddedColumns,PaddedLayout}}) = Val(true)
 copy(M::Mul{<:DualLayout{<:AbstractLazyLayout}, <:Union{PaddedColumns,PaddedLayout}}) = copy(mulreduce(M))
 simplifiable(::Mul{<:DiagonalLayout{<:AbstractFillLayout}, <:Union{PaddedColumns,PaddedLayout}}) = Val(true)

test/bandedtests.jl

@@ -425,6 +425,13 @@ LinearAlgebra.lmul!(β::Number, A::PseudoBandedMatrix) = (lmul!(β, A.data); A)
                 @test MemoryLayout(BroadcastMatrix(cos, A)) isa BroadcastLayout
             end
         end
+
+        @testset "broadcast_mul_mul" begin
+            A = BroadcastMatrix(*, randn(5,5), randn(5,5))
+            B = ApplyArray(*, brand(5,5,1,2), brand(5,5,2,1))
+            @test A * UpperTriangular(B) ≈ Matrix(A) * UpperTriangular(B)
+            @test simplifiable(*, A, UpperTriangular(B)) == Val(false) # TODO: probably should be true
+        end
     end
 
     @testset "Cache" begin

test/broadcasttests.jl

@@ -2,7 +2,7 @@ module BroadcastTests
 
 using LazyArrays, ArrayLayouts, LinearAlgebra, FillArrays, Base64, Test
 using StaticArrays, Tracker
-import LazyArrays: BroadcastLayout, arguments, LazyArrayStyle, sub_materialize
+import LazyArrays: BroadcastLayout, arguments, LazyArrayStyle, sub_materialize, simplifiable
 import Base: broadcasted
 
 using ..InfiniteArrays
@@ -224,6 +224,9 @@ using Infinities
         @test A[:,2] ≈ Ã[:,2] ≈ Matrix(A)[:,2]
         @test C*b ≈ Matrix(C)*b
 
+        @test simplifiable(*, A, B) == Val(true)
+        @test simplifiable(*, Ã, B) == Val(true)
+
         D = Diagonal(Fill(2,4))
         @test A*D ≈ Matrix(A)*D
     end
@@ -396,7 +399,7 @@ using Infinities
         @test a[:,1:3] isa Adjoint{Int,Vector{Int}}
     end
 
-    @testset "broadcast with adjtrans" begin
+    @testset "broadcast with adjtrans/triangular/hermsym" begin
         a = BroadcastArray(real, ((1:5) .+ im))
         b = BroadcastArray(exp, ((1:5) .+ im))
         @test exp.(transpose(a)) isa Transpose{<:Any,<:BroadcastVector}
@@ -407,8 +410,17 @@ using Infinities
         @test exp.(b') isa BroadcastMatrix
         @test exp.(transpose(b)) == transpose(exp.(b))
         @test exp.(b') == exp.(b)'
+
+        A = BroadcastArray(*, ((1:5) .+ im), (1:5)')
+        @test exp.(UpperTriangular(A)) isa BroadcastArray
+        @test exp.(Symmetric(A)) isa BroadcastArray
+        @test exp.(Hermitian(A)) isa BroadcastArray
+        @test exp.(UpperTriangular(A)) == exp.(UpperTriangular(Matrix(A)))
+        @test exp.(Symmetric(A)) == exp.(Symmetric(Matrix(A)))
+        @test exp.(Hermitian(A)) == exp.(Hermitian(Matrix(A)))
     end
 
+
     @testset "linear indexing" begin
         a = BroadcastArray(real, ((1:5) .+ im))
         b = BroadcastArray(exp, ((1:5) .+ im))
@@ -421,6 +433,22 @@ using Infinities
         a = BroadcastArray(Base.literal_pow, ^, 1:5, Val(2))
         @test last(a) == 25
     end
+
+    @testset "BroadcastArray(*) * MulArray" begin
+        A = BroadcastArray(*, 1:3, randn(3,4))
+        B = ApplyArray(*, randn(4,3), randn(3,4))
+        @test A*B ≈ Matrix(A)*Matrix(B)
+        @test A*UpperTriangular(B) ≈ Matrix(A)*UpperTriangular(Matrix(B))
+        @test simplifiable(*,A,B) == Val(false) # TODO: Why False?
+        @test simplifiable(*,A,UpperTriangular(B)) == Val(false) # TODO: Why False?
+    end
+
+    @testset "/" begin
+        A = BroadcastArray(/, randn(3,4), randn(3,4))
+        B = randn(4,3)
+        @test A*B ≈ Matrix(A)*Matrix(B)
+        @test simplifiable(*,A,B) == Val(false) # TODO: Why False?
+    end
 end
 
 end #module

test/cachetests.jl

@@ -189,7 +189,7 @@ using Infinities
     @testset "linalg" begin
         c = cache(Fill(3,3,3))
         @test fill(2,1,3) * c == fill(18,1,3)
-        @test ApplyMatrix(exp,fill(3,3,3)) * c == exp(fill(3,3,3)) * fill(3,3,3)
+        @test ApplyMatrix(exp,fill(3,3,3)) * c ≈ exp(fill(3,3,3)) * fill(3,3,3)
         @test BroadcastMatrix(exp,fill(3,3,3)) * c == exp.(fill(3,3,3)) * fill(3,3,3)
         @test fill(2,3)' * c == fill(18,1,3)
         @test fill(2,3,1)' * c == fill(18,1,3)

test/ldivtests.jl

@@ -205,4 +205,12 @@ end
     @test rowsupport(invL, ()) == 1:0
 end
 
+@testset "Inv \\ Lazy" begin
+    A = randn(5,5)
+    Ai = InvMatrix(A)
+    @test Ai \ Ai ≈ I
+    @test Ai \ BroadcastArray(exp, A) ≈ Ai \ exp.(A) ≈ A*exp.(A)
+    @test Ai \ ApplyArray(*, A, A) ≈ A^3
+end
+
 end # module

test/paddedtests.jl

@@ -430,5 +430,12 @@ paddeddata(a::PaddedPadded) = a
         @test_throws SingularException ArrayLayouts.ldiv!(Bidiagonal(-1:3, 1:4, :L), c)
         @test_throws SingularException ArrayLayouts.ldiv!(Bidiagonal(-4:0, 1:4, :L), c)
     end
+
+    @testset "Broadcast * Padded" begin
+        B = BroadcastArray(*, 1:8, (2:9)')
+        p = Vcat(1:2, Zeros(6))
+        @test B*p == Matrix(B)*p
+        @test simplifiable(*,B,p) == Val(true)
+    end
 end
 end # module