StructuredMatrixStyle{OffDiagonal} broadcasting

Project.toml

@@ -1,6 +1,6 @@
 name = "FillArrays"
 uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
-version = "1.18.0"
+version = "1.18.1"
 
 
 [deps]

src/offdiagonal.jl

@@ -133,3 +133,50 @@ Base.@propagate_inbounds _singleband_getindex(_, A, kr, jr) = invoke(getindex, T
     @inbounds band = _band(A)[st:st+len-1]
     OffDiagonal(band, k, shape)
 end
+
+##
+# Broadcasting a zero-preserving function over OffDiagonals with the same band returns an OffDiagonal
+##
+
+BroadcastStyle(::Type{<:OffDiagonal}) = LinearAlgebra.StructuredMatrixStyle{OffDiagonal}()
+@inline LinearAlgebra.fzero(::OffDiagonal{T}) where T = zero(T)
+
+# the band index shared by all the OffDiagonal arguments of a broadcast, or nothing if they lie on different
+# bands or are broadcast to a different size, in which case the result is not supported on a single band
+function _offdiagonal_bcindex(bc::Broadcasted)
+    k = nothing
+    for x in Broadcast.cat_nested(bc)
+        if x isa OffDiagonal
+            axes(x) == axes(bc) || return nothing
+            k === nothing || k == x.k || return nothing
+            k = x.k
+        end
+    end
+    k
+end
+
+# replace the OffDiagonal arguments by their bands, so that e.g. a Fill band stays a Fill
+_bcband(A::OffDiagonal) = A.band
+_bcband(bc::Broadcasted) = broadcasted(bc.f, map(_bcband, bc.args)...)
+_bcband(x) = x
+
+_isoffdiagonalbc(bc) = _offdiagonal_bcindex(bc) !== nothing && LinearAlgebra.fzeropreserving(bc)
+
+function Base.copy(bc::Broadcasted{LinearAlgebra.StructuredMatrixStyle{OffDiagonal}})
+    _isoffdiagonalbc(bc) || return copy(convert(Broadcasted{DefaultArrayStyle{2}}, bc))
+    OffDiagonal(Broadcast.materialize(_bcband(bc)), _offdiagonal_bcindex(bc), axes(bc))
+end
+
+function Base.copyto!(dest::OffDiagonal, bc::Broadcasted{LinearAlgebra.StructuredMatrixStyle{OffDiagonal}})
+    if axes(dest) == axes(bc) && _isoffdiagonalbc(bc) && _offdiagonal_bcindex(bc) == dest.k
+        Broadcast.materialize!(dest.band, _bcband(bc))
+        return dest
+    end
+    copyto!(dest, convert(Broadcasted{Nothing}, bc))
+end
+
+function LinearAlgebra.structured_broadcast_alloc(bc, ::Type{<:OffDiagonal}, ::Type{ElType}, _) where ElType
+    k = _offdiagonal_bcindex(bc)
+    k === nothing && return similar(convert(Broadcasted{DefaultArrayStyle{2}}, bc), ElType)
+    OffDiagonal(Vector{ElType}(undef, _bandlength(map(length, axes(bc))..., k)), k, axes(bc))
+end

test/runtests.jl

@@ -443,6 +443,58 @@ end
         @test rot180(Eye(5)[2:5, 1:4]) ≡ OffDiagonal(Ones(3), -1, 4, 4)
     end
 
+    @testset "broadcasting" begin
+        @test Base.BroadcastStyle(typeof(OffDiagonal([1, 2, 3], 1))) == LinearAlgebra.StructuredMatrixStyle{OffDiagonal}()
+
+        A = Eye(5)[2:end, :]
+        @test sin.(A) isa OffDiagonal
+        @test sin.(A) == sin.(Matrix(A))
+        # lazy bands stay lazy
+        @test sin.(A) ≡ OffDiagonal(Fill(sin(1.0), 4), 1, 4, 5)
+        @test 2 .* A ≡ OffDiagonal(Fill(2.0, 4), 1, 4, 5)
+
+        for A in (OffDiagonal([1.0, 2, 3], 1), OffDiagonal([1.0, 2, 3], -2, 5, 4), OffDiagonal(Fill(2.0, 3), 2, 3, 5),
+                  OffDiagonal(1.0:3, -1), OffDiagonal(Float64[], 4, 2, 3), Eye(5)[2:end, :], Eye(4, 6)[:, 2:5])
+            M = Matrix(A)
+            # zero-preserving operations keep the band
+            for (B, BM) in ((sin.(A), sin.(M)), (abs.(A), abs.(M)), (2 .* A, 2 .* M), (A ./ 2, M ./ 2), (.-A, .-M),
+                            (A .^ 2, M .^ 2), (A .+ A, M .+ M), (A .* A, M .* M), (round.(Int, A), round.(Int, M)),
+                            (sin.(A) .+ 2 .* A, sin.(M) .+ 2 .* M), (A .* sin.(A), M .* sin.(M)))
+                @test B isa OffDiagonal
+                @test B.k == A.k
+                @test B == BM
+                @test eltype(B) == eltype(BM)
+            end
+            # other operations are dense
+            for (B, BM) in ((cos.(A), cos.(M)), (A .+ 1, M .+ 1), (A .* M, M .* M), (A .+ M, M .+ M), (A .* (1:size(A,1)), M .* (1:size(A,1))))
+                @test B isa Matrix
+                @test B == BM
+            end
+            # different bands are dense
+            O = OffDiagonal(ones(FillArrays._bandlength(size(A)..., A.k + 1)), A.k + 1, axes(A))
+            @test A .+ O isa Matrix
+            @test A .+ O == M .+ Matrix(O)
+
+            # in-place
+            B = OffDiagonal(zeros(length(A.band)), A.k, axes(A))
+            B .= sin.(A)
+            @test B == sin.(M)
+            B .= A .* B
+            @test B == M .* sin.(M)
+            @test similar(Broadcast.instantiate(Broadcast.broadcasted(sin, A)), Float64) isa OffDiagonal{Float64}
+            if !isempty(A.band)
+                @test_throws ArgumentError B .= A .+ 1
+            end
+        end
+
+        # broadcasting a 1×1 OffDiagonal expands it, so is dense
+        @test OffDiagonal([1.0], 0) .* OffDiagonal([1.0, 2, 3], 0) == Diagonal([1.0, 2, 3])
+        @test OffDiagonal([1.0], 0) .* OffDiagonal([1.0, 2, 3], 0) isa Matrix
+        # mixing with other structured matrices is dense
+        A = OffDiagonal([1.0, 2, 3], 0)
+        @test A .+ Diagonal([1.0, 2, 3]) == Diagonal([2.0, 4, 6])
+    end
+
     @testset "slicing returns OffDiagonal" begin
         for M in (Eye(5), Eye{Int}(4), Diagonal(Fill(2, 6)), Diagonal(Zeros(3)), Eye(4, 6), Eye(6, 4),
                   RectDiagonal([1, 2, 3]), RectDiagonal([1, 2, 3], (1:3, 1:5)), RectDiagonal([1, 2, 3], (1:5, 1:3)),