From 9df076133b4edaac510aea11a8f309f17feb8592 Mon Sep 17 00:00:00 2001 From: "Steven G. Johnson" Date: Wed, 16 Nov 2016 15:44:02 -0500 Subject: [PATCH] added hcat/vcat/hvcat methods for UniformScaling (#19305) * added hcat/vcat/hvcat methods for UniformScaling * NEWS for 19305 * unified speye and sparse(I) code, and added a note about the latter in the docs of the former --- NEWS.md | 4 ++ base/linalg/uniformscaling.jl | 102 +++++++++++++++++++++++++++++++++- base/sparse/sparsematrix.jl | 24 +++++--- base/sparse/sparsevector.jl | 8 ++- doc/stdlib/arrays.rst | 2 + test/linalg/uniformscaling.jl | 20 +++++++ 6 files changed, 150 insertions(+), 10 deletions(-) diff --git a/NEWS.md b/NEWS.md index 4d9b092e8a0095..4769d77ec04ddd 100644 --- a/NEWS.md +++ b/NEWS.md @@ -56,6 +56,10 @@ Library improvements * BitArrays can now be constructed from arbitrary iterables, in particular from generator expressions, e.g. `BitArray(isodd(x) for x = 1:100)` ([#19018]). + * `hcat`, `vcat`, and `hvcat` now work with `UniformScaling` objects, so + you can now do e.g. `[A I]` and it will concatenate an appropriately sized + identity matrix ([#19305]). + Compiler/Runtime improvements ----------------------------- diff --git a/base/linalg/uniformscaling.jl b/base/linalg/uniformscaling.jl index 4bfc5d67fb6d7e..afec8815e77e3a 100644 --- a/base/linalg/uniformscaling.jl +++ b/base/linalg/uniformscaling.jl @@ -1,6 +1,7 @@ # This file is a part of Julia. License is MIT: http://julialang.org/license -import Base: copy, ctranspose, getindex, show, transpose, one, zero, inv +import Base: copy, ctranspose, getindex, show, transpose, one, zero, inv, + @_pure_meta, hcat, vcat, hvcat import Base.LinAlg: SingularException immutable UniformScaling{T<:Number} @@ -166,3 +167,102 @@ inv(J::UniformScaling) = UniformScaling(inv(J.λ)) isapprox{T<:Number,S<:Number}(J1::UniformScaling{T}, J2::UniformScaling{S}; rtol::Real=Base.rtoldefault(T,S), atol::Real=0) = isapprox(J1.λ, J2.λ, rtol=rtol, atol=atol) + +function copy!(A::AbstractMatrix, J::UniformScaling) + size(A,1)==size(A,2) || throw(DimensionMismatch("a UniformScaling can only be copied to a square matrix")) + fill!(A, 0) + λ = J.λ + for i = 1:size(A,1) + @inbounds A[i,i] = λ + end + return A +end + +# promote_to_arrays(n,k, T, A...) promotes any UniformScaling matrices +# 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_{T}(n::Int, ::Type{Matrix}, J::UniformScaling{T}) = copy!(Matrix{T}(n,n), J) +promote_to_arrays_(n::Int, ::Type, A::AbstractVecOrMat) = A +promote_to_arrays(n,k, ::Type) = () +promote_to_arrays{T}(n,k, ::Type{T}, A) = (promote_to_arrays_(n[k], T, A),) +promote_to_arrays{T}(n,k, ::Type{T}, A, B) = + (promote_to_arrays_(n[k], T, A), promote_to_arrays_(n[k+1], T, B)) +promote_to_arrays{T}(n,k, ::Type{T}, A, B, C) = + (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{T}(n,k, ::Type{T}, A, B, Cs...) = + (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}}}) = (@_pure_meta; Matrix) + +for (f,dim,name) in ((:hcat,1,"rows"), (:vcat,2,"cols")) + @eval begin + function $f(A::Union{AbstractVecOrMat,UniformScaling}...) + n = 0 + for a in A + if !isa(a, UniformScaling) + na = size(a,$dim) + n > 0 && n != na && + throw(DimensionMismatch(string("number of ", $name, + " of each array must match (got ", n, " and ", na, ")"))) + n = na + end + end + n == 0 && 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...)...) + end + end +end + + +function hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractVecOrMat,UniformScaling}...) + nr = length(rows) + sum(rows) == length(A) || throw(ArgumentError("mismatch between row sizes and number of arguments")) + n = zeros(Int, length(A)) + needcols = false # whether we also need to infer some sizes from the column count + j = 0 + for i = 1:nr # infer UniformScaling sizes from row counts, if possible: + ni = 0 # number of rows in this block-row + for k = 1:rows[i] + if !isa(A[j+k], UniformScaling) + na = size(A[j+k], 1) + ni > 0 && ni != na && + throw(DimensionMismatch("mismatch in number of rows")) + ni = na + end + end + if ni > 0 + for k = 1:rows[i] + n[j+k] = ni + end + else # row consisted only of UniformScaling objects + needcols = true + end + j += rows[i] + end + if needcols # some sizes still unknown, try to infer from column count + nc = j = 0 + for i = 1:nr + nci = 0 + rows[i] > 0 && n[j+1] == 0 && continue # column count unknown in this row + for k = 1:rows[i] + nci += isa(A[j+k], UniformScaling) ? n[j+k] : size(A[j+k], 2) + end + nc > 0 && nc != nci && throw(DimensionMismatch("mismatch in number of columns")) + nc = nci + j += rows[i] + end + nc == 0 && throw(ArgumentError("sizes of UniformScalings could not be inferred")) + j = 0 + for i = 1:nr + if rows[i] > 0 && n[j+1] == 0 # this row consists entirely of UniformScalings + nci = nc ÷ rows[i] + nci * rows[i] != nc && throw(DimensionMismatch("indivisible UniformScaling sizes")) + for k = 1:rows[i] + n[j+k] = nci + end + end + j += rows[i] + end + end + return hvcat(rows, promote_to_arrays(n,1, promote_to_array_type(A), A...)...) +end diff --git a/base/sparse/sparsematrix.jl b/base/sparse/sparsematrix.jl index 4a24f4464dc90f..64582d2ada4843 100644 --- a/base/sparse/sparsematrix.jl +++ b/base/sparse/sparsematrix.jl @@ -1374,15 +1374,12 @@ eye(S::SparseMatrixCSC) = speye(S) Create a sparse identity matrix of size `m x m`. When `n` is supplied, create a sparse identity matrix of size `m x n`. The type defaults to `Float64` if not specified. + +`sparse(I, m, n)` is equivalent to `speye(Int, m, n)`, and +`sparse(α*I, m, n)` can be used to efficiently create a sparse +multiple `α` of the identity matrix. """ -function speye(T::Type, m::Integer, n::Integer) - ((m < 0) || (n < 0)) && throw(ArgumentError("invalid Array dimensions")) - x = min(m,n) - rowval = [1:x;] - colptr = [rowval; fill(Int(x+1), n+1-x)] - nzval = ones(T, x) - return SparseMatrixCSC(m, n, colptr, rowval, nzval) -end +speye(T::Type, m::Integer, n::Integer) = speye_scaled(one(T), m, n) function one{T}(S::SparseMatrixCSC{T}) m,n = size(S) @@ -1390,6 +1387,17 @@ function one{T}(S::SparseMatrixCSC{T}) speye(T, m) end +function speye_scaled(diag, m::Integer, n::Integer) + ((m < 0) || (n < 0)) && throw(ArgumentError("invalid array dimensions")) + nnz = min(m,n) + colptr = Array(Int, 1+n) + colptr[1:nnz+1] = 1:nnz+1 + colptr[nnz+2:end] = nnz+1 + SparseMatrixCSC(Int(m), Int(n), colptr, Vector{Int}(1:nnz), fill(diag, nnz)) +end + +sparse(S::UniformScaling, m::Integer, n::Integer=m) = speye_scaled(S.λ, m, n) + ## Unary arithmetic and boolean operators """ diff --git a/base/sparse/sparsevector.jl b/base/sparse/sparsevector.jl index 976cf208db6db4..8ff15081e52e50 100644 --- a/base/sparse/sparsevector.jl +++ b/base/sparse/sparsevector.jl @@ -2,7 +2,8 @@ ### Common definitions -import Base: scalarmax, scalarmin, sort, find, findnz +import Base: scalarmax, scalarmin, sort, find, findnz, @_pure_meta +import Base.LinAlg: promote_to_array_type, promote_to_arrays_ ### The SparseVector @@ -895,6 +896,11 @@ function hvcat(rows::Tuple{Vararg{Int}}, X::_SparseConcatGroup...) vcat(tmp_rows...) end +# make sure UniformScaling objects are converted to sparse matrices for concatenation +promote_to_array_type(A::Tuple{Vararg{Union{_SparseConcatGroup,UniformScaling}}}) = (@_pure_meta; SparseMatrixCSC) +promote_to_array_type(A::Tuple{Vararg{Union{_DenseConcatGroup,UniformScaling}}}) = (@_pure_meta; Matrix) +promote_to_arrays_{T}(n::Int, ::Type{SparseMatrixCSC}, J::UniformScaling{T}) = sparse(J, n,n) + # Concatenations strictly involving un/annotated dense matrices/vectors should yield dense arrays cat(catdims, xs::_DenseConcatGroup...) = Base.cat_t(catdims, promote_eltype(xs...), xs...) vcat(A::_DenseConcatGroup...) = Base.typed_vcat(promote_eltype(A...), A...) diff --git a/doc/stdlib/arrays.rst b/doc/stdlib/arrays.rst index 7bcdfa25f602c3..dad5555044bfa8 100644 --- a/doc/stdlib/arrays.rst +++ b/doc/stdlib/arrays.rst @@ -2237,6 +2237,8 @@ dense counterparts. The following functions are specific to sparse arrays. Create a sparse identity matrix of size ``m x m``\ . When ``n`` is supplied, create a sparse identity matrix of size ``m x n``\ . The type defaults to ``Float64`` if not specified. + ``sparse(I, m, n)`` is equivalent to ``speye(Int, m, n)``\ , and ``sparse(α*I, m, n)`` can be used to efficiently create a sparse multiple ``α`` of the identity matrix. + .. function:: speye(S) .. Docstring generated from Julia source diff --git a/test/linalg/uniformscaling.jl b/test/linalg/uniformscaling.jl index c848ab8de21023..fb67d84d282983 100644 --- a/test/linalg/uniformscaling.jl +++ b/test/linalg/uniformscaling.jl @@ -16,6 +16,8 @@ srand(123) @test one(UniformScaling(rand(Complex128))) == one(UniformScaling{Complex128}) @test eltype(one(UniformScaling(rand(Complex128)))) == Complex128 @test -one(UniformScaling(2)) == UniformScaling(-1) + @test sparse(3I,4,5) == spdiagm(fill(3,4),0,4,5) + @test sparse(3I,5,4) == spdiagm(fill(3,4),0,5,4) end @testset "istriu, istril, issymmetric, ishermitian, isapprox" begin @@ -144,3 +146,21 @@ B = bitrand(2,2) end end end + +@testset "hcat and vcat" begin + @test_throws ArgumentError hcat(I) + @test_throws ArgumentError [I I] + @test_throws ArgumentError vcat(I) + @test_throws ArgumentError [I; I] + @test_throws ArgumentError [I I; I] + for T in (Matrix, SparseMatrixCSC) + A = T(rand(3,4)) + B = T(rand(3,3)) + @test (hcat(A,2I))::T == hcat(A,2eye(3,3)) + @test (vcat(A,2I))::T == vcat(A,2eye(4,4)) + @test (hcat(I,3I,A,2I))::T == hcat(eye(3,3),3eye(3,3),A,2eye(3,3)) + @test (vcat(I,3I,A,2I))::T == vcat(eye(4,4),3eye(4,4),A,2eye(4,4)) + @test (hvcat((2,1,2),B,2I,I,3I,4I))::T == + hvcat((2,1,2),B,2eye(3,3),eye(6,6),3eye(3,3),4eye(3,3)) + end +end