broadcast adjoint should return an adjoint when possible

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"

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
 #####

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
 

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)

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