diff --git a/NEWS.md b/NEWS.md index a38f4dfb53660..e5d3619986d08 100644 --- a/NEWS.md +++ b/NEWS.md @@ -28,6 +28,7 @@ Standard library changes #### LinearAlgebra +* `diagm` and `spdiagm` now accept optional `m,n` initial arguments to specify a size ([#31654]). #### SparseArrays diff --git a/stdlib/LinearAlgebra/src/dense.jl b/stdlib/LinearAlgebra/src/dense.jl index 9326ff8a6155b..359530ecb7041 100644 --- a/stdlib/LinearAlgebra/src/dense.jl +++ b/stdlib/LinearAlgebra/src/dense.jl @@ -247,9 +247,14 @@ diag(A::AbstractMatrix, k::Integer=0) = A[diagind(A,k)] """ diagm(kv::Pair{<:Integer,<:AbstractVector}...) + diagm(m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...) -Construct a square matrix from `Pair`s of diagonals and vectors. +Construct a matrix from `Pair`s of diagonals and vectors. Vector `kv.second` will be placed on the `kv.first` diagonal. +By default the matrix is square and its size is inferred +from `kv`, but a non-square size `m`×`n` (padded with zeros as needed) +can be specified by passing `m,n` as the first arguments. + `diagm` constructs a full matrix; if you want storage-efficient versions with fast arithmetic, see [`Diagonal`](@ref), [`Bidiagonal`](@ref) [`Tridiagonal`](@ref) and [`SymTridiagonal`](@ref). @@ -271,8 +276,10 @@ julia> diagm(1 => [1,2,3], -1 => [4,5]) 0 0 0 0 ``` """ -function diagm(kv::Pair{<:Integer,<:AbstractVector}...) - A = diagm_container(kv...) +diagm(kv::Pair{<:Integer,<:AbstractVector}...) = _diagm(nothing, kv...) +diagm(m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...) = _diagm((Int(m),Int(n)), kv...) +function _diagm(size, kv::Pair{<:Integer,<:AbstractVector}...) + A = diagm_container(size, kv...) for p in kv inds = diagind(A, p.first) for (i, val) in enumerate(p.second) @@ -281,20 +288,32 @@ function diagm(kv::Pair{<:Integer,<:AbstractVector}...) end return A end -function diagm_container(kv::Pair{<:Integer,<:AbstractVector}...) - T = promote_type(map(x -> eltype(x.second), kv)...) - n = mapreduce(x -> length(x.second) + abs(x.first), max, kv) - return zeros(T, n, n) +function diagm_size(size::Nothing, kv::Pair{<:Integer,<:AbstractVector}...) + mnmax = mapreduce(x -> length(x.second) + abs(Int(x.first)), max, kv; init=0) + return mnmax, mnmax end -function diagm_container(kv::Pair{<:Integer,<:BitVector}...) - n = mapreduce(x -> length(x.second) + abs(x.first), max, kv) - return falses(n, n) +function diagm_size(size::Tuple{Int,Int}, kv::Pair{<:Integer,<:AbstractVector}...) + mmax = mapreduce(x -> length(x.second) - min(0,Int(x.first)), max, kv; init=0) + nmax = mapreduce(x -> length(x.second) + max(0,Int(x.first)), max, kv; init=0) + m, n = size + (m ≥ mmax && n ≥ nmax) || throw(DimensionMismatch("invalid size=$size")) + return m, n +end +function diagm_container(size, kv::Pair{<:Integer,<:AbstractVector}...) + T = promote_type(map(x -> eltype(x.second), kv)...) + return zeros(T, diagm_size(size, kv...)...) end +diagm_container(size, kv::Pair{<:Integer,<:BitVector}...) = + falses(diagm_size(size, kv...)...) """ diagm(v::AbstractVector) + diagm(m::Integer, n::Integer, v::AbstractVector) -Construct a square matrix with elements of the vector as diagonal elements. +Construct a matrix with elements of the vector as diagonal elements. +By default (if `size=nothing`), the matrix is square and its size is given by +`length(v)`, but a non-square size `m`×`n` can be specified +by passing `m,n` as the first arguments. # Examples ```jldoctest @@ -306,6 +325,7 @@ julia> diagm([1,2,3]) ``` """ diagm(v::AbstractVector) = diagm(0 => v) +diagm(m::Integer, n::Integer, v::AbstractVector) = diagm(m, n, 0 => v) function tr(A::Matrix{T}) where T n = checksquare(A) diff --git a/stdlib/LinearAlgebra/test/dense.jl b/stdlib/LinearAlgebra/test/dense.jl index 496c84fccbe49..de284f24a5997 100644 --- a/stdlib/LinearAlgebra/test/dense.jl +++ b/stdlib/LinearAlgebra/test/dense.jl @@ -91,6 +91,21 @@ bimg = randn(n,2)/2 @test diagm(v) == diagm(0 => v) end +@testset "Non-square diagm" begin + x = [7, 8] + for m=1:4, n=2:4 + if m < 2 || n < 3 + @test_throws DimensionMismatch diagm(m,n, 0 => x, 1 => x) + @test_throws DimensionMismatch diagm(n,m, 0 => x, -1 => x) + else + M = zeros(m,n) + M[1:2,1:3] = [7 7 0; 0 8 8] + @test diagm(m,n, 0 => x, 1 => x) == M + @test diagm(n,m, 0 => x, -1 => x) == M' + end + end +end + @testset "Test pinv (rtol, atol)" begin M = [1 0 0; 0 1 0; 0 0 0] @test pinv(M,atol=1)== zeros(3,3) diff --git a/stdlib/SparseArrays/src/sparsematrix.jl b/stdlib/SparseArrays/src/sparsematrix.jl index b1dfb9dfc10fe..650fecb706eff 100644 --- a/stdlib/SparseArrays/src/sparsematrix.jl +++ b/stdlib/SparseArrays/src/sparsematrix.jl @@ -3357,9 +3357,13 @@ end """ spdiagm(kv::Pair{<:Integer,<:AbstractVector}...) + spdiagm(m::Integer, n::Ingeger, kv::Pair{<:Integer,<:AbstractVector}...) -Construct a square sparse diagonal matrix from `Pair`s of vectors and diagonals. -Vector `kv.second` will be placed on the `kv.first` diagonal. +Construct a sparse diagonal matrix from `Pair`s of vectors and diagonals. +Each vector `kv.second` will be placed on the `kv.first` diagonal. By +default (if `size=nothing`), the matrix is square and its size is inferred +from `kv`, but a non-square size `m`×`n` (padded with zeros as needed) +can be specified by passing `m,n` as the first arguments. # Examples ```jldoctest @@ -3375,10 +3379,15 @@ julia> spdiagm(-1 => [1,2,3,4], 1 => [4,3,2,1]) [4, 5] = 1 ``` """ -function spdiagm(kv::Pair{<:Integer,<:AbstractVector}...) +spdiagm(kv::Pair{<:Integer,<:AbstractVector}...) = _spdiagm(nothing, kv...) +spdiagm(m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...) = _spdiagm((Int(m),Int(n)), kv...) +function _spdiagm(size, kv::Pair{<:Integer,<:AbstractVector}...) I, J, V = spdiagm_internal(kv...) - n = max(dimlub(I), dimlub(J)) - return sparse(I, J, V, n, n) + mmax, nmax = dimlub(I), dimlub(J) + mnmax = max(mmax, nmax) + m, n = something(size, (mnmax,mnmax)) + (m ≥ mmax && n ≥ nmax) || throw(DimensionMismatch("invalid size=$size")) + return sparse(I, J, V, m, n) end ## expand a colptr or rowptr into a dense index vector diff --git a/stdlib/SparseArrays/test/sparse.jl b/stdlib/SparseArrays/test/sparse.jl index f752982437a7e..3a8ee5a635bcf 100644 --- a/stdlib/SparseArrays/test/sparse.jl +++ b/stdlib/SparseArrays/test/sparse.jl @@ -1583,6 +1583,17 @@ end end # promotion @test spdiagm(0 => [1,2], 1 => [3.5], -1 => [4+5im]) == [1 3.5; 4+5im 2] + + # non-square: + for m=1:4, n=2:4 + if m < 2 || n < 3 + @test_throws DimensionMismatch spdiagm(m,n, 0 => x, 1 => x) + else + M = zeros(m,n) + M[1:2,1:3] = [1 1 0; 0 1 1] + @test spdiagm(m,n, 0 => x, 1 => x) == M + end + end end @testset "diag" begin