From 020be439cddc3e895a3c6a386ed8b3b96369ef88 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Wed, 4 Sep 2024 19:59:10 +0530 Subject: [PATCH] Broadcast for Bidiagonal/Tridiagonal/SymTridiaognal --- src/infarrays.jl | 8 ++++++-- test/runtests.jl | 10 ++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/src/infarrays.jl b/src/infarrays.jl index 6576282..541b30b 100644 --- a/src/infarrays.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 650a5ba..3d39749 100644 --- a/test/runtests.jl +++ b/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