diff --git a/Project.toml b/Project.toml index a4a70cba..2f3fac62 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FillArrays" uuid = "1a297f60-69ca-5386-bcde-b61e274b549b" -version = "1.18.0" +version = "1.18.1" [deps] diff --git a/src/offdiagonal.jl b/src/offdiagonal.jl index 68759512..32a97baa 100644 --- a/src/offdiagonal.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 0e70b97e..407395ca 100644 --- a/test/runtests.jl +++ b/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)),