Skip to content

Commit

Permalink
Handle AbstractQ in concatenation (#51132)
Browse files Browse the repository at this point in the history 
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 50146b2)
  • Loading branch information
@dkarrasch
dkarrasch authored and KristofferC committed Oct 23, 2023
1 parent f2c447d commit 4d3022e
Show file tree
Hide file tree
Showing 5 changed files with 128 additions and 115 deletions.
9 changes: 4 additions & 5 deletions base/abstractarray.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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...)
Expand Down
11 changes: 11 additions & 0 deletions stdlib/LinearAlgebra/src/abstractq.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand Down
110 changes: 110 additions & 0 deletions stdlib/LinearAlgebra/src/special.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
108 changes: 0 additions & 108 deletions stdlib/LinearAlgebra/src/uniformscaling.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
5 changes: 3 additions & 2 deletions stdlib/LinearAlgebra/test/special.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 4d3022e

Please sign in to comment.