Broadcast for Bidiagonal/Tridiagonal/SymTridiaognal (#194)

Author: jishnub

src/infarrays.jl

@@ -173,7 +173,9 @@ for typ in (:Ones, :Zeros, :Fill)
     end
 end
 
-BroadcastStyle(::Type{<:Diagonal{T,<:AbstractFill{T,1,Tuple{OneToInf{I}}}}}) where {T,I} = LazyArrayStyle{2}()
+for M in (:Diagonal, :Bidiagonal, :Tridiagonal, :SymTridiagonal)
+    @eval BroadcastStyle(::Type{<:$M{T,<:AbstractFill{T,1,Tuple{OneToInf{I}}}}}) where {T,I} = LazyArrayStyle{2}()
+end
 
 ## Support broadcast(*, ::AbstractFill, A)
 
@@ -285,7 +287,9 @@ end
 one(D::Diagonal{T,<:AbstractFill{T,1,Tuple{OneToInf{Int}}}}) where T = Eye{T}(size(D,1))
 copy(D::Diagonal{T,<:AbstractFill{T,1,Tuple{OneToInf{Int}}}}) where T = D
 
-BroadcastStyle(::Type{<:Diagonal{<:Any,<:AbstractInfUnitRange}}) = LazyArrayStyle{2}()
+for M in (:Diagonal, :Bidiagonal, :Tridiagonal, :SymTridiagonal)
+    @eval BroadcastStyle(::Type{<:$M{<:Any,<:AbstractInfUnitRange}}) = LazyArrayStyle{2}()
+end
 sub_materialize(::AbstractBandedLayout, V, ::Tuple{InfAxes,InfAxes}) = V
 sub_materialize(::AbstractBandedLayout, V, ::Tuple{OneTo{Int},InfAxes}) = V
 sub_materialize(::AbstractBandedLayout, V, ::Tuple{InfAxes,OneTo{Int}}) = V

test/runtests.jl

@@ -912,6 +912,16 @@ end
         @test c[1:6] == a[r[2:7]]
         @test Base.BroadcastStyle(typeof(c)) isa LazyArrayStyle
     end
+
+    @testset "structured matrices" begin
+        r = 1:∞
+        f = Fill(2, ∞)
+        for B in (Bidiagonal(r, r, :U), Tridiagonal(r, r, r), SymTridiagonal(r, r),
+                   Bidiagonal(f, f, :U), Tridiagonal(f, f, f), SymTridiagonal(f, f))
+            B2 = B .+ B
+            @test B2[1:10, 1:10] == 2B[1:10, 1:10]
+        end
+    end
 end
 
 @testset "Cumsum and diff" begin