Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle AbstractQ in concatenation #51132

Merged
merged 4 commits into from
Sep 27, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 4 additions & 5 deletions base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1807,17 +1807,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
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
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
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
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