diff --git a/NEWS.md b/NEWS.md index bb22c9f940a78..9aebf5d42d954 100644 --- a/NEWS.md +++ b/NEWS.md @@ -138,6 +138,8 @@ Standard library changes (callable via `cholesky[!](A, RowMaximum())`) ([#54619]). * The number of default BLAS threads now respects process affinity, instead of using total number of logical threads available on the system ([#55574]). +* A new function `zeroslike` is added that is used to generate the zero elements for matrix-valued banded matrices. + Custom array types may specialize this function to return an appropriate result. ([#55252]) #### Logging diff --git a/stdlib/LinearAlgebra/src/LinearAlgebra.jl b/stdlib/LinearAlgebra/src/LinearAlgebra.jl index 49d73127ba7ba..15354603943c2 100644 --- a/stdlib/LinearAlgebra/src/LinearAlgebra.jl +++ b/stdlib/LinearAlgebra/src/LinearAlgebra.jl @@ -175,6 +175,7 @@ public AbstractTriangular, peakflops, symmetric, symmetric_type, + zeroslike, matprod_dest const BlasFloat = Union{Float64,Float32,ComplexF64,ComplexF32} diff --git a/stdlib/LinearAlgebra/src/bidiag.jl b/stdlib/LinearAlgebra/src/bidiag.jl index e5482cbba5595..abce41995647a 100644 --- a/stdlib/LinearAlgebra/src/bidiag.jl +++ b/stdlib/LinearAlgebra/src/bidiag.jl @@ -118,15 +118,14 @@ Bidiagonal(A::Bidiagonal) = A Bidiagonal{T}(A::Bidiagonal{T}) where {T} = A Bidiagonal{T}(A::Bidiagonal) where {T} = Bidiagonal{T}(A.dv, A.ev, A.uplo) -bidiagzero(::Bidiagonal{T}, i, j) where {T} = zero(T) -function bidiagzero(A::Bidiagonal{<:AbstractMatrix}, i, j) - Tel = eltype(eltype(A.dv)) +function diagzero(A::Bidiagonal{<:AbstractMatrix}, i, j) + Tel = eltype(A) if i < j && A.uplo == 'U' #= top right zeros =# - return zeros(Tel, size(A.ev[i], 1), size(A.ev[j-1], 2)) + return zeroslike(Tel, axes(A.ev[i], 1), axes(A.ev[j-1], 2)) elseif j < i && A.uplo == 'L' #= bottom left zeros =# - return zeros(Tel, size(A.ev[i-1], 1), size(A.ev[j], 2)) + return zeroslike(Tel, axes(A.ev[i-1], 1), axes(A.ev[j], 2)) else - return zeros(Tel, size(A.dv[i], 1), size(A.dv[j], 2)) + return zeroslike(Tel, axes(A.dv[i], 1), axes(A.dv[j], 2)) end end @@ -161,7 +160,7 @@ end elseif i == j - _offdiagind(A.uplo) return @inbounds A.ev[A.uplo == 'U' ? i : j] else - return bidiagzero(A, i, j) + return diagzero(A, i, j) end end @@ -173,7 +172,7 @@ end # we explicitly compare the possible bands as b.band may be constant-propagated return @inbounds A.ev[b.index] else - return bidiagzero(A, Tuple(_cartinds(b))...) + return diagzero(A, Tuple(_cartinds(b))...) end end diff --git a/stdlib/LinearAlgebra/src/diagonal.jl b/stdlib/LinearAlgebra/src/diagonal.jl index 0a95bac5ffb93..be253c03ab3f6 100644 --- a/stdlib/LinearAlgebra/src/diagonal.jl +++ b/stdlib/LinearAlgebra/src/diagonal.jl @@ -185,8 +185,27 @@ end end r end -diagzero(::Diagonal{T}, i, j) where {T} = zero(T) -diagzero(D::Diagonal{<:AbstractMatrix{T}}, i, j) where {T} = zeros(T, size(D.diag[i], 1), size(D.diag[j], 2)) +""" + diagzero(A::AbstractMatrix, i, j) + +Return the appropriate zero element `A[i, j]` corresponding to a banded matrix `A`. +""" +diagzero(A::AbstractMatrix, i, j) = zero(eltype(A)) +diagzero(D::Diagonal{M}, i, j) where {M<:AbstractMatrix} = + zeroslike(M, axes(D.diag[i], 1), axes(D.diag[j], 2)) +# dispatching on the axes permits specializing on the axis types to return something other than an Array +zeroslike(M::Type, ax::Vararg{Union{AbstractUnitRange, Integer}}) = zeroslike(M, ax) +""" + zeroslike(::Type{M}, ax::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) where {M<:AbstractMatrix} + zeroslike(::Type{M}, sz::Tuple{Integer, Vararg{Integer}}) where {M<:AbstractMatrix} + +Return an appropriate zero-ed array similar to `M`, with either the axes `ax` or the size `sz`. +This will be used as a structural zero element of a matrix-valued banded matrix. +By default, `zeroslike` falls back to using the size along each axis to construct the array. +""" +zeroslike(M::Type, ax::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) = zeroslike(M, map(length, ax)) +zeroslike(M::Type, sz::Tuple{Integer, Vararg{Integer}}) = zeros(M, sz) +zeroslike(::Type{M}, sz::Tuple{Integer, Vararg{Integer}}) where {M<:AbstractMatrix} = zeros(eltype(M), sz) @inline function getindex(D::Diagonal, b::BandIndex) @boundscheck checkbounds(D, b) diff --git a/stdlib/LinearAlgebra/test/bidiag.jl b/stdlib/LinearAlgebra/test/bidiag.jl index d633a99a2390e..628e59debe8b7 100644 --- a/stdlib/LinearAlgebra/test/bidiag.jl +++ b/stdlib/LinearAlgebra/test/bidiag.jl @@ -839,6 +839,16 @@ end B = Bidiagonal(dv, ev, :U) @test B == Matrix{eltype(B)}(B) end + + @testset "non-standard axes" begin + LinearAlgebra.diagzero(T::Type, ax::Tuple{SizedArrays.SOneTo, Vararg{SizedArrays.SOneTo}}) = + zeros(T, ax) + + s = SizedArrays.SizedArray{(2,2)}([1 2; 3 4]) + B = Bidiagonal(fill(s,4), fill(s,3), :U) + @test @inferred(B[2,1]) isa typeof(s) + @test all(iszero, B[2,1]) + end end @testset "copyto!" begin diff --git a/stdlib/LinearAlgebra/test/diagonal.jl b/stdlib/LinearAlgebra/test/diagonal.jl index 98f5498c71033..85fe963e3592b 100644 --- a/stdlib/LinearAlgebra/test/diagonal.jl +++ b/stdlib/LinearAlgebra/test/diagonal.jl @@ -815,6 +815,13 @@ end D = Diagonal(fill(S,3)) @test D * fill(S,2,3)' == fill(S * S', 3, 2) @test fill(S,3,2)' * D == fill(S' * S, 2, 3) + + @testset "indexing with non-standard-axes" begin + s = SizedArrays.SizedArray{(2,2)}([1 2; 3 4]) + D = Diagonal(fill(s,3)) + @test @inferred(D[1,2]) isa typeof(s) + @test all(iszero, D[1,2]) + end end @testset "Eigensystem for block diagonal (issue #30681)" begin diff --git a/test/testhelpers/SizedArrays.jl b/test/testhelpers/SizedArrays.jl index 495fe03347ee7..e52e965a64859 100644 --- a/test/testhelpers/SizedArrays.jl +++ b/test/testhelpers/SizedArrays.jl @@ -99,4 +99,7 @@ mul!(dest::AbstractMatrix, S1::SizedMatrix, S2::SizedMatrix, α::Number, β::Num mul!(dest::AbstractVector, M::AbstractMatrix, v::SizedVector, α::Number, β::Number) = mul!(dest, M, _data(v), α, β) +LinearAlgebra.zeroslike(::Type{S}, ax::Tuple{SizedArrays.SOneTo, Vararg{SizedArrays.SOneTo}}) where {S<:SizedArray} = + zeros(eltype(S), ax) + end