diff --git a/Project.toml b/Project.toml index cc18e35..55911dd 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "BlockBandedMatrices" uuid = "ffab5731-97b5-5995-9138-79e8c1846df0" -version = "0.13.5" +version = "0.13.6" [deps] ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a" @@ -19,7 +19,7 @@ BlockBandedMatricesSparseArraysExt = "SparseArrays" [compat] Aqua = "0.8" ArrayLayouts = "1" -BandedMatrices = "1" +BandedMatrices = "1.13" BlockArrays = "1" Documenter = "1" FillArrays = "1" diff --git a/src/BandedBlockBandedMatrix.jl b/src/BandedBlockBandedMatrix.jl index 8892187..fc367d0 100644 --- a/src/BandedBlockBandedMatrix.jl +++ b/src/BandedBlockBandedMatrix.jl @@ -616,6 +616,48 @@ rmul!(A::BandedBlockBandedMatrix, x::Number) = (rmul!(A.data, x); A) *(A::BandedBlockBandedMatrix, x::Number) = _BandedBlockBandedMatrix(A.data*x, axes(A,1), blockbandwidths(A), subblockbandwidths(A)) +### +# broadcasting is passed through to the data +# +# The data of a `BandedBlockBandedMatrix` stores every entry inside the bands, so a +# broadcast whose arguments all share the same block structure can be applied to the data +# directly, e.g. `α .* A .+ B` is `α .* A.data .+ B.data` reinterpreted as a +# `BandedBlockBandedMatrix`. This needs the result to vanish outside the bands, where the +# data stores junk that is never read. +### + + +_isdatabroadcast(bc::Broadcasted) = _isbroadcastarith(bc.f) && _isweakzero(bc.f, bc.args...) + +_isdataarg(_, _) = false +_isdataarg(_, ::Number) = true +# the data must also be safe to read, which it need not be for `undef` entries of a +# non-isbits eltype +_isdataarg(A, B::BandedBlockBandedMatrix) = isbitstype(eltype(B)) && + blockisequal(axes(A), axes(B)) && blockbandwidths(A) == blockbandwidths(B) && + subblockbandwidths(A) == subblockbandwidths(B) +_isdataarg(A, bc::Broadcasted) = _isdatabroadcast(bc) && all(map(x -> _isdataarg(A, x), bc.args)) + +_dataarg(x::Number) = x +_dataarg(A::BandedBlockBandedMatrix) = bandedblockbandeddata(A) +_dataarg(bc::Broadcasted) = broadcasted(bc.f, map(_dataarg, bc.args)...) + +# the first BandedBlockBandedMatrix in the tree determines the structure of the result +_firstbandedblockbanded() = nothing +_firstbandedblockbanded(A::BandedBlockBandedMatrix, args...) = A +_firstbandedblockbanded(bc::Broadcasted, args...) = _firstbandedblockbanded(bc.args..., args...) +_firstbandedblockbanded(_, args...) = _firstbandedblockbanded(args...) + +function copy(bc::Broadcasted{BandedBlockBandedStyle}) + A = _firstbandedblockbanded(bc) + if !isnothing(A) && _isdataarg(A, bc) + _BandedBlockBandedMatrix(materialize(_dataarg(bc)), axes(A), blockbandwidths(A), subblockbandwidths(A)) + else # e.g. exp.(A) or matrices whose bands do not line up + invoke(copy, Tuple{Broadcasted}, bc) + end +end + + ##### # summary ##### diff --git a/src/BlockBandedMatrices.jl b/src/BlockBandedMatrices.jl index c9f950e..2e41c13 100644 --- a/src/BlockBandedMatrices.jl +++ b/src/BlockBandedMatrices.jl @@ -8,8 +8,9 @@ import ArrayLayouts: AbstractColumnMajor, AdjQRPackedQLayout, ColumnMajor, layout_replace_in_print_matrix, rowsupport, sub_materialize, sublayout, symmetriclayout, transposelayout, zero!, materialize!, materialize -import BandedMatrices: BandedColumns, BandedLayout, BandedStyle, BlasFloat, - _BandedMatrix, banded_getindex, banded_setindex!, +import BandedMatrices: AdjOrTransStyle, BandedColumns, BandedLayout, BandedStyle, BlasFloat, + _BandedMatrix, _isbroadcastarith, _isweakzero, adjortransstyle, + banded_getindex, banded_setindex!, bandedcolumns, bandeddata, bandshift, bandwidth, bandwidths, colrange, inbands_getindex, inbands_setindex!, isbanded diff --git a/src/broadcast.jl b/src/broadcast.jl index cfc54fb..099e1e1 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -32,6 +32,27 @@ BroadcastStyle(::BlockBandedStyle, ::DefaultArrayStyle{2}) = BlockBandedStyle() BroadcastStyle(::DefaultArrayStyle{2}, ::BandedBlockBandedStyle) = BandedBlockBandedStyle() BroadcastStyle(::BandedBlockBandedStyle, ::DefaultArrayStyle{2}) = BandedBlockBandedStyle() +### +# Adjoints and transposes +# +# The adjoint (transpose) of a block-banded matrix is not stored as a block-banded matrix, +# so we use the `AdjOrTransStyle` machinery of BandedMatrices.jl to broadcast over the +# parents and re-wrap the result, e.g. `α .* A'` returns an `Adjoint` of a block-banded +# matrix instead of a dense matrix. +### + +adjortransstyle(op, sty::AbstractBlockBandedStyle) = AdjOrTransStyle{typeof(op),typeof(sty)}() + +BroadcastStyle(::Type{<:Adjoint{<:Any,Mat}}) where Mat<:AbstractBlockBandedMatrix = + adjortransstyle(adjoint, BroadcastStyle(Mat)) +BroadcastStyle(::Type{<:Transpose{<:Any,Mat}}) where Mat<:AbstractBlockBandedMatrix = + adjortransstyle(transpose, BroadcastStyle(Mat)) + +# mixing with a dense array is not supported by the block-banded machinery, as the block +# bandwidths of an unblocked array are not known, so we materialise densely as before +BroadcastStyle(::AdjOrTransStyle{<:Any,<:AbstractBlockBandedStyle}, T::DefaultArrayStyle{2}) = T + + ### # broadcast blockbandwidths ### @@ -107,8 +128,6 @@ end # zero is preserved. Take the maximum bandwidth -import BandedMatrices: _isweakzero - function blockbandwidths(bc::Broadcasted) (a,b) = size(bc) bnds = (a-1,b-1) diff --git a/test/test_broadcasting.jl b/test/test_broadcasting.jl index 06758fb..296b83b 100644 --- a/test/test_broadcasting.jl +++ b/test/test_broadcasting.jl @@ -297,6 +297,106 @@ import Base: oneto bbw = @inferred blockbandwidths(bc) @test bbw == blockbandwidths(BB) end + + @testset "adjoint/transpose" begin + A = BandedBlockBandedMatrix(randn(ComplexF64,10,10), 1:4,1:4, (1,1), (1,1)) + B = BlockBandedMatrix(randn(ComplexF64,10,10), 1:4,1:4, (1,1)) + α = 2+im + + @testset "$op" for op in (adjoint, transpose) + Wrap = op === adjoint ? Adjoint : Transpose + @testset "$(typeof(M).name.name)" for M in (A, B) + W, Wm = op(M), op(Matrix(M)) + for (R,Rm) in ((α .* W, α .* Wm), (W .* α, Wm .* α), (W ./ α, Wm ./ α), + (α .\ W, α .\ Wm), ((-).(W), (-).(Wm)), + (W .+ W, Wm .+ Wm), (W .- W, Wm .- Wm), + (α .* W .+ W, α .* Wm .+ Wm)) + @test R isa Wrap{ComplexF64} + @test parent(R) isa typeof(M).name.wrapper + @test R ≈ Rm + end + end + @test op(A) .+ op(B) isa Wrap{ComplexF64} + @test parent(op(A) .+ op(B)) isa BlockBandedMatrix + @test op(A) .+ op(B) ≈ op(Matrix(A)) .+ op(Matrix(B)) + end + + # the wrapper is dropped when it cannot be moved through the broadcast + @test exp.(A') ≈ exp.(Matrix(A)') + @test !(exp.(A') isa Adjoint) + @test imag.(A') ≈ imag.(Matrix(A)') + @test !(imag.(A') isa Adjoint) + @test A' .+ transpose(A) ≈ Matrix(A)' .+ transpose(Matrix(A)) + @test !(A' .+ transpose(A) isa Adjoint) + @test A' .+ A ≈ Matrix(A)' .+ Matrix(A) + @test A' .* ones(10,10) ≈ Matrix(A)' + + f(α, A) = α .* A' + @test @inferred(f(α, A)) isa Adjoint{ComplexF64,<:BandedBlockBandedMatrix} + @test f(α, A) ≈ α .* Matrix(A)' + + C = BandedBlockBandedMatrix{ComplexF64}(undef, 1:4,1:4, (1,1), (1,1)) + C .= α .* A' + @test C ≈ α .* Matrix(A)' + D = Matrix{ComplexF64}(undef, 10, 10) + D .= α .* A' + @test D ≈ α .* Matrix(A)' + end + + @testset "data" begin + A = BandedBlockBandedMatrix{Float64}(undef, 1:4,1:4, (1,1),(1,1)) + B = BandedBlockBandedMatrix{Float64}(undef, 1:4,1:4, (1,1),(1,1)) + # the data outside the bands is junk that must not leak into the result + A.data .= NaN; B.data .= NaN + for M in (A,B), J = 1:4, K = max(1,J-1):min(4,J+1) + V = view(M, Block(K), Block(J)) + for j = axes(V,2), k = colrange(V,j) + V[k,j] = randn() + end + end + Am, Bm = Matrix(A), Matrix(B) + α = 2.0 + @test !any(isnan, Am) + + for (R,Rm) in ((α .* A, α .* Am), (A .* α, Am .* α), (A ./ α, Am ./ α), (α .\ A, α .\ Am), + ((-).(A), (-).(Am)), (A .+ B, Am .+ Bm), (A .- B, Am .- Bm), + (A .* B, Am .* Bm), (α .* A .+ B, α .* Am .+ Bm)) + @test R isa BandedBlockBandedMatrix + @test blockbandwidths(R) == (1,1) + @test subblockbandwidths(R) == (1,1) + @test Matrix(R) == Rm + end + @test typeof(α .* A) == typeof(similar(A, Float64)) + + # the data is only used when the structures line up + C = BandedBlockBandedMatrix(randn(10,10), 1:4,1:4, (2,1),(1,0)); Cm = Matrix(C) + @test A .+ C ≈ Am .+ Cm + @test blockbandwidths(A .+ C) == (2,1) + @test subblockbandwidths(A .+ C) == (1,1) + @test A .* C ≈ Am .* Cm + @test blockbandwidths(A .* C) == (1,1) + @test subblockbandwidths(A .* C) == (1,0) + @test_throws DimensionMismatch A .+ BandedBlockBandedMatrix(randn(10,14), 1:4,2:5, (1,1),(1,1)) + + # broadcasts that do not preserve the bands are unchanged + @test A .+ 1 ≈ Am .+ 1 + @test exp.(A) ≈ exp.(Am) + @test Diagonal(1:10) .* A ≈ Diagonal(1:10) .* Am + + f(α, A) = α .* A + g(A, B) = A .+ B + h(α, A, B) = α .* A .+ B + @test @inferred(f(α, A)) ≈ α .* Am + # Following can't be inferred because of issues in type-inferrence + # see Base.Broadcast.axistype overload in + # BlockArrays/src/blockbroadcast.jl:45 + @test g(A, B) ≈ Am .+ Bm + @test h(α, A, B) ≈ α .* Am .+ Bm + + # and the data is used through an adjoint as well + @test parent((2+im) .* A') isa BandedBlockBandedMatrix + @test (2+im) .* A' ≈ (2+im) .* Am' + end end end # module