From 4d3022ea6098a4382be7775c0de0f1df88b9d1f6 Mon Sep 17 00:00:00 2001 From: Daniel Karrasch Date: Wed, 27 Sep 2023 21:57:54 +0200 Subject: [PATCH] Handle `AbstractQ` in concatenation (#51132) Handling `AbstractQ`s in concatenation got lost in the overhaul. This brings it back (though it seemed to not be used anywhere), and even better: pre-1.9 `Q`s where handled via `AbstractMatrix` fallbacks, so elementwise. Now, it first materializes the matrix, and then copies to the right place in the destination array. (cherry picked from commit 50146b2dd27e8f3bc54367427a70cd96f74ada45) --- base/abstractarray.jl | 9 +- stdlib/LinearAlgebra/src/abstractq.jl | 11 +++ stdlib/LinearAlgebra/src/special.jl | 110 +++++++++++++++++++++ stdlib/LinearAlgebra/src/uniformscaling.jl | 108 -------------------- stdlib/LinearAlgebra/test/special.jl | 5 +- 5 files changed, 128 insertions(+), 115 deletions(-) diff --git a/base/abstractarray.jl b/base/abstractarray.jl index 8f1274d960e41..6f9823a5be5d7 100644 --- a/base/abstractarray.jl +++ b/base/abstractarray.jl @@ -1811,17 +1811,16 @@ function __cat_offset1!(A, shape, catdims, offsets, x) inds = ntuple(length(offsets)) do i (i <= length(catdims) && catdims[i]) ? offsets[i] .+ cat_indices(x, i) : 1:shape[i] end - if x isa AbstractArray - A[inds...] = x - else - fill!(view(A, inds...), x) - end + _copy_or_fill!(A, inds, x) newoffsets = ntuple(length(offsets)) do i (i <= length(catdims) && catdims[i]) ? offsets[i] + cat_size(x, i) : offsets[i] end return newoffsets end +_copy_or_fill!(A, inds, x) = fill!(view(A, inds...), x) +_copy_or_fill!(A, inds, x::AbstractArray) = (A[inds...] = x) + """ vcat(A...) diff --git a/stdlib/LinearAlgebra/src/abstractq.jl b/stdlib/LinearAlgebra/src/abstractq.jl index 93358d052d50b..2aa333beef2b2 100644 --- a/stdlib/LinearAlgebra/src/abstractq.jl +++ b/stdlib/LinearAlgebra/src/abstractq.jl @@ -8,6 +8,7 @@ end parent(adjQ::AdjointQ) = adjQ.Q eltype(::Type{<:AbstractQ{T}}) where {T} = T +Base.eltypeof(Q::AbstractQ) = eltype(Q) ndims(::AbstractQ) = 2 # inversion/adjoint/transpose @@ -129,6 +130,16 @@ function copyto!(dest::PermutedDimsArray{T,2,perm}, src::AbstractQ) where {T,per end return dest end +# used in concatenations: Base.__cat_offset1! +Base._copy_or_fill!(A, inds, Q::AbstractQ) = (A[inds...] = collect(Q)) +# overloads of helper functions +Base.cat_size(A::AbstractQ) = size(A) +Base.cat_size(A::AbstractQ, d) = size(A, d) +Base.cat_length(a::AbstractQ) = prod(size(a)) +Base.cat_ndims(a::AbstractQ) = ndims(a) +Base.cat_indices(A::AbstractQ, d) = axes(A, d) +Base.cat_similar(A::AbstractQ, T::Type, shape::Tuple) = Array{T}(undef, shape) +Base.cat_similar(A::AbstractQ, T::Type, shape::Vector) = Array{T}(undef, shape...) function show(io::IO, ::MIME{Symbol("text/plain")}, Q::AbstractQ) print(io, Base.dims2string(size(Q)), ' ', summary(Q)) diff --git a/stdlib/LinearAlgebra/src/special.jl b/stdlib/LinearAlgebra/src/special.jl index 885f29fa1417b..d028fe43e6338 100644 --- a/stdlib/LinearAlgebra/src/special.jl +++ b/stdlib/LinearAlgebra/src/special.jl @@ -336,6 +336,116 @@ const _SpecialArrays = Union{} promote_to_array_type(::Tuple) = Matrix +# 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_(n::Int, ::Type, a::Number) = a +promote_to_arrays_(n::Int, ::Type{Matrix}, J::UniformScaling{T}) where {T} = Matrix(J, n, n) +promote_to_arrays_(n::Int, ::Type, A::AbstractArray) = A +promote_to_arrays_(n::Int, ::Type, A::AbstractQ) = collect(A) +promote_to_arrays(n,k, ::Type) = () +promote_to_arrays(n,k, ::Type{T}, A) where {T} = (promote_to_arrays_(n[k], T, A),) +promote_to_arrays(n,k, ::Type{T}, A, B) where {T} = + (promote_to_arrays_(n[k], T, A), promote_to_arrays_(n[k+1], T, B)) +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...)...) + +_us2number(A) = A +_us2number(J::UniformScaling) = J.λ + +for (f, _f, dim, name) in ((:hcat, :_hcat, 1, "rows"), (:vcat, :_vcat, 2, "cols")) + @eval begin + @inline $f(A::Union{AbstractArray,AbstractQ,UniformScaling}...) = $_f(A...) + # if there's a Number present, J::UniformScaling must be 1x1-dimensional + @inline $f(A::Union{AbstractArray,AbstractQ,UniformScaling,Number}...) = $f(map(_us2number, A)...) + function $_f(A::Union{AbstractArray,AbstractQ,UniformScaling,Number}...; array_type = promote_to_array_type(A)) + n = -1 + for a in A + if !isa(a, UniformScaling) + require_one_based_indexing(a) + 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 == -1 && throw(ArgumentError($("$f of only UniformScaling objects cannot determine the matrix size"))) + return cat(promote_to_arrays(fill(n, length(A)), 1, array_type, A...)..., dims=Val(3-$dim)) + end + end +end + +hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractArray,AbstractQ,UniformScaling}...) = _hvcat(rows, A...) +hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractArray,AbstractQ,UniformScaling,Number}...) = _hvcat(rows, A...) +function _hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractArray,AbstractQ,UniformScaling,Number}...; array_type = promote_to_array_type(A)) + require_one_based_indexing(A...) + nr = length(rows) + sum(rows) == length(A) || throw(ArgumentError("mismatch between row sizes and number of arguments")) + n = fill(-1, 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 = -1 # number of rows in this block-row, -1 indicates unknown + 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 = -1 + j = 0 + for i = 1:nr + nci = 0 + rows[i] > 0 && n[j+1] == -1 && (j += rows[i]; continue) + 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 == -1 && throw(ArgumentError("sizes of UniformScalings could not be inferred")) + j = 0 + for i = 1:nr + if rows[i] > 0 && n[j+1] == -1 # this row consists entirely of UniformScalings + nci, r = divrem(nc, rows[i]) + r != 0 && throw(DimensionMismatch("indivisible UniformScaling sizes")) + for k = 1:rows[i] + n[j+k] = nci + end + end + j += rows[i] + end + end + Amat = promote_to_arrays(n, 1, array_type, 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 array_type == Matrix + return typed_hvcat(promote_eltype(Amat...), rows, Amat...) + else + return hvcat(rows, Amat...) + end +end + # factorizations function cholesky(S::RealHermSymComplexHerm{<:Real,<:SymTridiagonal}, ::NoPivot = NoPivot(); check::Bool = true) T = choltype(eltype(S)) diff --git a/stdlib/LinearAlgebra/src/uniformscaling.jl b/stdlib/LinearAlgebra/src/uniformscaling.jl index 0b3168113acf7..b014472a9cec2 100644 --- a/stdlib/LinearAlgebra/src/uniformscaling.jl +++ b/stdlib/LinearAlgebra/src/uniformscaling.jl @@ -402,114 +402,6 @@ function cond(J::UniformScaling{T}) where T return J.λ ≠ zero(T) ? onereal : oftype(onereal, Inf) 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_(n::Int, ::Type, a::Number) = a -promote_to_arrays_(n::Int, ::Type{Matrix}, J::UniformScaling{T}) where {T} = Matrix(J, n, n) -promote_to_arrays_(n::Int, ::Type, A::AbstractArray) = A -promote_to_arrays(n,k, ::Type) = () -promote_to_arrays(n,k, ::Type{T}, A) where {T} = (promote_to_arrays_(n[k], T, A),) -promote_to_arrays(n,k, ::Type{T}, A, B) where {T} = - (promote_to_arrays_(n[k], T, A), promote_to_arrays_(n[k+1], T, B)) -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...)...) - -_us2number(A) = A -_us2number(J::UniformScaling) = J.λ - -for (f, _f, dim, name) in ((:hcat, :_hcat, 1, "rows"), (:vcat, :_vcat, 2, "cols")) - @eval begin - @inline $f(A::Union{AbstractArray,UniformScaling}...) = $_f(A...) - # if there's a Number present, J::UniformScaling must be 1x1-dimensional - @inline $f(A::Union{AbstractArray,UniformScaling,Number}...) = $f(map(_us2number, A)...) - function $_f(A::Union{AbstractArray,UniformScaling,Number}...; array_type = promote_to_array_type(A)) - n = -1 - for a in A - if !isa(a, UniformScaling) - require_one_based_indexing(a) - 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 == -1 && throw(ArgumentError($("$f of only UniformScaling objects cannot determine the matrix size"))) - return cat(promote_to_arrays(fill(n, length(A)), 1, array_type, A...)..., dims=Val(3-$dim)) - end - end -end - -hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractArray,UniformScaling,Number}...) = _hvcat(rows, A...) -function _hvcat(rows::Tuple{Vararg{Int}}, A::Union{AbstractArray,UniformScaling,Number}...; array_type = promote_to_array_type(A)) - require_one_based_indexing(A...) - nr = length(rows) - sum(rows) == length(A) || throw(ArgumentError("mismatch between row sizes and number of arguments")) - n = fill(-1, 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 = -1 # number of rows in this block-row, -1 indicates unknown - 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 = -1 - j = 0 - for i = 1:nr - nci = 0 - rows[i] > 0 && n[j+1] == -1 && (j += rows[i]; continue) - 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 == -1 && throw(ArgumentError("sizes of UniformScalings could not be inferred")) - j = 0 - for i = 1:nr - if rows[i] > 0 && n[j+1] == -1 # this row consists entirely of UniformScalings - nci, r = divrem(nc, rows[i]) - r != 0 && throw(DimensionMismatch("indivisible UniformScaling sizes")) - for k = 1:rows[i] - n[j+k] = nci - end - end - j += rows[i] - end - end - Amat = promote_to_arrays(n, 1, array_type, 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 array_type == Matrix - return typed_hvcat(promote_eltype(Amat...), rows, Amat...) - else - return hvcat(rows, Amat...) - end -end - ## Matrix construction from UniformScaling function Matrix{T}(s::UniformScaling, dims::Dims{2}) where {T} A = zeros(T, dims) diff --git a/stdlib/LinearAlgebra/test/special.jl b/stdlib/LinearAlgebra/test/special.jl index eaa297e05d957..7e96af369e310 100644 --- a/stdlib/LinearAlgebra/test/special.jl +++ b/stdlib/LinearAlgebra/test/special.jl @@ -259,9 +259,10 @@ end bidiagmat = Bidiagonal(1:N, 1:(N-1), :U) tridiagmat = Tridiagonal(1:(N-1), 1:N, 1:(N-1)) symtridiagmat = SymTridiagonal(1:N, 1:(N-1)) - specialmats = (diagmat, bidiagmat, tridiagmat, symtridiagmat) + abstractq = qr(tridiagmat).Q + specialmats = (diagmat, bidiagmat, tridiagmat, symtridiagmat, abstractq, zeros(Int,N,N)) for specialmata in specialmats, specialmatb in specialmats - MA = Matrix(specialmata); MB = Matrix(specialmatb) + MA = collect(specialmata); MB = collect(specialmatb) @test hcat(specialmata, specialmatb) == hcat(MA, MB) @test vcat(specialmata, specialmatb) == vcat(MA, MB) @test hvcat((1,1), specialmata, specialmatb) == hvcat((1,1), MA, MB)