From f49a7ae18bc65db60d5563d4df21404aad591c01 Mon Sep 17 00:00:00 2001 From: Daniel Karrasch Date: Sun, 4 Jul 2021 11:36:06 +0200 Subject: [PATCH] Concatenation with UniformScaling and numbers (#41394) Co-authored-by: Jameson Nash --- stdlib/LinearAlgebra/src/LinearAlgebra.jl | 4 ++-- stdlib/LinearAlgebra/src/uniformscaling.jl | 26 +++++++++++++++------ stdlib/LinearAlgebra/test/uniformscaling.jl | 12 ++++++++++ stdlib/SparseArrays/src/sparsevector.jl | 17 ++++++++------ 4 files changed, 43 insertions(+), 16 deletions(-) diff --git a/stdlib/LinearAlgebra/src/LinearAlgebra.jl b/stdlib/LinearAlgebra/src/LinearAlgebra.jl index 55caad1f82bee9..f0f13776146d1a 100644 --- a/stdlib/LinearAlgebra/src/LinearAlgebra.jl +++ b/stdlib/LinearAlgebra/src/LinearAlgebra.jl @@ -15,8 +15,8 @@ import Base: USE_BLAS64, abs, acos, acosh, acot, acoth, acsc, acsch, adjoint, as oneunit, parent, power_by_squaring, print_matrix, promote_rule, real, round, sec, sech, setindex!, show, similar, sin, sincos, sinh, size, sqrt, strides, stride, tan, tanh, transpose, trunc, typed_hcat, vec -using Base: IndexLinear, promote_op, promote_typeof, - @propagate_inbounds, @pure, reduce, typed_vcat, require_one_based_indexing, +using Base: IndexLinear, promote_eltype, promote_op, promote_typeof, + @propagate_inbounds, @pure, reduce, typed_hvcat, typed_vcat, require_one_based_indexing, splat using Base.Broadcast: Broadcasted, broadcasted import Libdl diff --git a/stdlib/LinearAlgebra/src/uniformscaling.jl b/stdlib/LinearAlgebra/src/uniformscaling.jl index 44ad5d98e99a41..bbcade43c5569d 100644 --- a/stdlib/LinearAlgebra/src/uniformscaling.jl +++ b/stdlib/LinearAlgebra/src/uniformscaling.jl @@ -391,6 +391,7 @@ end # in A to matrices of type T and sizes given by n[k:end]. n is an array # so that the same promotion code can be used for hvcat. We pass the type T # so that we can re-use this code for sparse-matrix hcat etcetera. +promote_to_arrays_(n::Int, ::Type, a::Number) = a promote_to_arrays_(n::Int, ::Type{Matrix}, J::UniformScaling{T}) where {T} = copyto!(Matrix{T}(undef, n,n), J) promote_to_arrays_(n::Int, ::Type, A::AbstractVecOrMat) = A promote_to_arrays(n,k, ::Type) = () @@ -401,11 +402,11 @@ promote_to_arrays(n,k, ::Type{T}, A, B, C) where {T} = (promote_to_arrays_(n[k], T, A), promote_to_arrays_(n[k+1], T, B), promote_to_arrays_(n[k+2], T, C)) promote_to_arrays(n,k, ::Type{T}, A, B, Cs...) where {T} = (promote_to_arrays_(n[k], T, A), promote_to_arrays_(n[k+1], T, B), promote_to_arrays(n,k+2, T, Cs...)...) -promote_to_array_type(A::Tuple{Vararg{Union{AbstractVecOrMat,UniformScaling}}}) = Matrix +promote_to_array_type(A::Tuple{Vararg{Union{AbstractVecOrMat,UniformScaling,Number}}}) = Matrix for (f,dim,name) in ((:hcat,1,"rows"), (:vcat,2,"cols")) @eval begin - function $f(A::Union{AbstractVecOrMat,UniformScaling}...) + function $f(A::Union{AbstractVecOrMat,UniformScaling,Number}...) n = -1 for a in A if !isa(a, UniformScaling) @@ -418,13 +419,13 @@ for (f,dim,name) in ((:hcat,1,"rows"), (:vcat,2,"cols")) end end n == -1 && throw(ArgumentError($("$f of only UniformScaling objects cannot determine the matrix size"))) - return $f(promote_to_arrays(fill(n,length(A)),1, promote_to_array_type(A), A...)...) + return cat(promote_to_arrays(fill(n, length(A)), 1, promote_to_array_type(A), A...)..., dims=Val(3-$dim)) end end end -function hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractVecOrMat,UniformScaling}...) +function hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractVecOrMat,UniformScaling,Number}...) require_one_based_indexing(A...) nr = length(rows) sum(rows) == length(A) || throw(ArgumentError("mismatch between row sizes and number of arguments")) @@ -467,8 +468,8 @@ function hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractVecOrMat,UniformScalin j = 0 for i = 1:nr if rows[i] > 0 && n[j+1] == -1 # this row consists entirely of UniformScalings - nci = nc ÷ rows[i] - nci * rows[i] != nc && throw(DimensionMismatch("indivisible UniformScaling sizes")) + nci, r = divrem(nc, rows[i]) + r != 0 && throw(DimensionMismatch("indivisible UniformScaling sizes")) for k = 1:rows[i] n[j+k] = nci end @@ -476,7 +477,18 @@ function hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractVecOrMat,UniformScalin j += rows[i] end end - return hvcat(rows, promote_to_arrays(n,1, promote_to_array_type(A), A...)...) + Atyp = promote_to_array_type(A) + Amat = promote_to_arrays(n, 1, Atyp, A...) + # We have two methods for promote_to_array_type, one returning Matrix and + # another one returning SparseMatrixCSC (in SparseArrays.jl). In the dense + # case, we cannot call hvcat for the promoted UniformScalings because this + # causes a stack overflow. In the sparse case, however, we cannot call + # typed_hvcat because we need a sparse output. + if Atyp == Matrix + return typed_hvcat(promote_eltype(Amat...), rows, Amat...) + else + return hvcat(rows, Amat...) + end end ## Matrix construction from UniformScaling diff --git a/stdlib/LinearAlgebra/test/uniformscaling.jl b/stdlib/LinearAlgebra/test/uniformscaling.jl index 04d4d7d7e0c9ea..edbb1789a366bc 100644 --- a/stdlib/LinearAlgebra/test/uniformscaling.jl +++ b/stdlib/LinearAlgebra/test/uniformscaling.jl @@ -335,10 +335,19 @@ end B = T(rand(3,3)) C = T(rand(0,3)) D = T(rand(2,0)) + E = T(rand(1,3)) + F = T(rand(3,1)) + α = rand() @test (hcat(A, 2I))::T == hcat(A, Matrix(2I, 3, 3)) + @test (hcat(E, α))::T == hcat(E, [α]) + @test (hcat(E, α, 2I))::T == hcat(E, [α], fill(2, 1, 1)) @test (vcat(A, 2I))::T == vcat(A, Matrix(2I, 4, 4)) + @test (vcat(F, α))::T == vcat(F, [α]) + @test (vcat(F, α, 2I))::T == vcat(F, [α], fill(2, 1, 1)) @test (hcat(C, 2I))::T == C + @test_throws DimensionMismatch hcat(C, α) @test (vcat(D, 2I))::T == D + @test_throws DimensionMismatch vcat(D, α) @test (hcat(I, 3I, A, 2I))::T == hcat(Matrix(I, 3, 3), Matrix(3I, 3, 3), A, Matrix(2I, 3, 3)) @test (vcat(I, 3I, A, 2I))::T == vcat(Matrix(I, 4, 4), Matrix(3I, 4, 4), A, Matrix(2I, 4, 4)) @test (hvcat((2,1,2), B, 2I, I, 3I, 4I))::T == @@ -353,6 +362,9 @@ end hvcat((2,2,2), B, Matrix(2I, 3, 3), C, C, Matrix(3I, 3, 3), Matrix(4I, 3, 3)) @test hvcat((3,2,1), C, C, I, B ,3I, 2I)::T == hvcat((2,2,1), C, C, B, Matrix(3I,3,3), Matrix(2I,6,6)) + @test (hvcat((1,2), A, E, α))::T == hvcat((1,2), A, E, [α]) == hvcat((1,2), A, E, α*I) + @test (hvcat((2,2), α, E, F, 3I))::T == hvcat((2,2), [α], E, F, Matrix(3I, 3, 3)) + @test (hvcat((2,2), 3I, F, E, α))::T == hvcat((2,2), Matrix(3I, 3, 3), F, E, [α]) end end diff --git a/stdlib/SparseArrays/src/sparsevector.jl b/stdlib/SparseArrays/src/sparsevector.jl index dae26a77626b44..55ad738a7eb770 100644 --- a/stdlib/SparseArrays/src/sparsevector.jl +++ b/stdlib/SparseArrays/src/sparsevector.jl @@ -1083,23 +1083,26 @@ const _Triangular_DenseArrays{T,A<:Matrix} = LinearAlgebra.AbstractTriangular{T, const _Annotated_DenseArrays = Union{_Triangular_DenseArrays, _Symmetric_DenseArrays, _Hermitian_DenseArrays} const _Annotated_Typed_DenseArrays{T} = Union{_Triangular_DenseArrays{T}, _Symmetric_DenseArrays{T}, _Hermitian_DenseArrays{T}} -const _SparseConcatGroup = Union{Vector, Adjoint{<:Any,<:Vector}, Transpose{<:Any,<:Vector}, Matrix, _SparseConcatArrays, _Annotated_SparseConcatArrays, _Annotated_DenseArrays} -const _DenseConcatGroup = Union{Vector, Adjoint{<:Any,<:Vector}, Transpose{<:Any,<:Vector}, Matrix, _Annotated_DenseArrays} +const _SparseConcatGroup = Union{Number, Vector, Adjoint{<:Any,<:Vector}, Transpose{<:Any,<:Vector}, Matrix, _SparseConcatArrays, _Annotated_SparseConcatArrays, _Annotated_DenseArrays} +const _DenseConcatGroup = Union{Number, Vector, Adjoint{<:Any,<:Vector}, Transpose{<:Any,<:Vector}, Matrix, _Annotated_DenseArrays} const _TypedDenseConcatGroup{T} = Union{Vector{T}, Adjoint{T,Vector{T}}, Transpose{T,Vector{T}}, Matrix{T}, _Annotated_Typed_DenseArrays{T}} # Concatenations involving un/annotated sparse/special matrices/vectors should yield sparse arrays +_makesparse(x::Number) = x +_makesparse(x::AbstractArray) = SparseMatrixCSC(issparse(x) ? x : sparse(x)) + function Base._cat(dims, Xin::_SparseConcatGroup...) - X = map(x -> SparseMatrixCSC(issparse(x) ? x : sparse(x)), Xin) + X = map(_makesparse, Xin) T = promote_eltype(Xin...) Base.cat_t(T, X...; dims=dims) end function hcat(Xin::_SparseConcatGroup...) - X = map(x -> SparseMatrixCSC(issparse(x) ? x : sparse(x)), Xin) - hcat(X...) + X = map(_makesparse, Xin) + return cat(X..., dims=Val(2)) end function vcat(Xin::_SparseConcatGroup...) - X = map(x -> SparseMatrixCSC(issparse(x) ? x : sparse(x)), Xin) - vcat(X...) + X = map(_makesparse, Xin) + return cat(X..., dims=Val(1)) end hvcat(rows::Tuple{Vararg{Int}}, X::_SparseConcatGroup...) = vcat(_hvcat_rows(rows, X...)...)