added hcat/vcat/hvcat methods for UniformScaling (JuliaLang#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

@@ -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
 -----------------------------
 

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

base/sparse/sparsematrix.jl

@@ -1374,22 +1374,30 @@ 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)
     if m != n; throw(DimensionMismatch("multiplicative identity only defined for square matrices")); end
     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
 
 """

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...)

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

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