Skip to content

Commit 7c5c010

Browse files
dlfivefiftydlfivefifty
authored
Add QR layouts (#17)
* Add QR layouts * add ArrayLayuots.ldiv! (and etc.) * Update muladd.jl * Update muladd.jl * Degenerate bands * Update factorizations.jl * QR tests, improve rdiv! * add tests * rectangular Q
1 parent 51dc733 commit 7c5c010

File tree

9 files changed

+463
-76
lines changed

9 files changed

+463
-76
lines changed

‎Project.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "ArrayLayouts"
22
uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
33
authors = ["Sheehan Olver <[email protected]>"]
4-
version = "0.2.6"
4+
version = "0.3.0"
55

66
[deps]
77
FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -12,8 +12,8 @@ FillArrays = "0.8"
1212
julia = "1"
1313

1414
[extras]
15-
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
1615
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
16+
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
1717

1818
[targets]
1919
test = ["Test", "Random"]

‎src/factorizations.jl

Lines changed: 268 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,108 @@
1-
abstract type AbstractQLayout <: MemoryLayout end
2-
struct QLayout <: AbstractQLayout end
1+
abstract type AbstractQRLayout <: MemoryLayout end
2+
3+
"""
4+
QRCompactWYLayout{SLAY,TLAY}()
5+
6+
represents a Compact-WY QR factorization whose
7+
factors are stored with layout SLAY and τ stored with layout TLAY
8+
"""
9+
struct QRCompactWYLayout{SLAY,TLAY} <: AbstractQRLayout end
10+
"""
11+
QRPackedLayout{SLAY,TLAY}()
12+
13+
represents a Packed QR factorization whose
14+
factors are stored with layout SLAY and τ stored with layout TLAY
15+
"""
16+
struct QRPackedLayout{SLAY,TLAY} <: AbstractQRLayout end
17+
18+
MemoryLayout(::Type{<:LinearAlgebra.QRCompactWY{<:Any,MAT}}) where MAT =
19+
QRCompactWYLayout{typeof(MemoryLayout(MAT)),DenseColumnMajor}()
20+
MemoryLayout(::Type{<:LinearAlgebra.QR{<:Any,MAT}}) where MAT =
21+
QRPackedLayout{typeof(MemoryLayout(MAT)),DenseColumnMajor}()
22+
23+
function materialize!(L::Ldiv{<:QRCompactWYLayout,<:Any,<:Any,<:AbstractVector})
24+
A,b = L.A, L.B
25+
ldiv!(UpperTriangular(A.R), view(lmul!(adjoint(A.Q), b), 1:size(A, 2)))
26+
b
27+
end
28+
29+
function materialize!(L::Ldiv{<:QRCompactWYLayout,<:Any,<:Any,<:AbstractMatrix})
30+
A,B = L.A, L.B
31+
ldiv!(UpperTriangular(A.R), view(lmul!(adjoint(A.Q), B), 1:size(A, 2), 1:size(B, 2)))
32+
B
33+
end
34+
35+
# Julia implementation similar to xgelsy
36+
function materialize!(Ldv::Ldiv{<:QRPackedLayout,<:Any,<:Any,<:AbstractMatrix{T}}) where T
37+
A,B = Ldv.A,Ldv.B
38+
m, n = size(A)
39+
minmn = min(m,n)
40+
mB, nB = size(B)
41+
lmul!(adjoint(A.Q), view(B, 1:m, :))
42+
mB ≥ n || throw(DimensionMismatch("Number of rows in B must exceed number of columns in A"))
43+
R = A.R
44+
@inbounds begin
45+
if n > m # minimum norm solution
46+
τ = zeros(T,m)
47+
for k = m:-1:1 # Trapezoid to triangular by elementary operation
48+
x = view(R, k, [k; m + 1:n])
49+
τk = reflector!(x)
50+
τ[k] = conj(τk)
51+
for i = 1:k - 1
52+
vRi = R[i,k]
53+
for j = m + 1:n
54+
vRi += R[i,j]*x[j - m + 1]'
55+
end
56+
vRi *= τk
57+
R[i,k] -= vRi
58+
for j = m + 1:n
59+
R[i,j] -= vRi*x[j - m + 1]
60+
end
61+
end
62+
end
63+
end
64+
ldiv!(UpperTriangular(view(R, :, 1:minmn)), view(B, 1:minmn, :))
65+
if n > m # Apply elementary transformation to solution
66+
B[m + 1:mB,1:nB] .= zero(T)
67+
for j = 1:nB
68+
for k = 1:m
69+
vBj = B[k,j]
70+
for i = m + 1:n
71+
vBj += B[i,j]*R[k,i]'
72+
end
73+
vBj *= τ[k]
74+
B[k,j] -= vBj
75+
for i = m + 1:n
76+
B[i,j] -= R[k,i]*vBj
77+
end
78+
end
79+
end
80+
end
81+
end
82+
return B
83+
end
84+
materialize!(Ldv::Ldiv{<:QRPackedLayout,<:Any,<:Any,<:AbstractVector{T}}) where T =
85+
ldiv!(Ldv.A, reshape(Ldv.B, length(Ldv.B), 1))[:]
86+
87+
388

4-
MemoryLayout(::Type{<:AbstractQ}) = QLayout()
89+
abstract type AbstractQLayout <: MemoryLayout end
90+
struct QRPackedQLayout{SLAY,TLAY} <: AbstractQLayout end
91+
struct AdjQRPackedQLayout{SLAY,TLAY} <: AbstractQLayout end
92+
struct QRCompactWYQLayout{SLAY,TLAY} <: AbstractQLayout end
93+
struct AdjQRCompactWYQLayout{SLAY,TLAY} <: AbstractQLayout end
594

6-
adjointlayout(::Type, ::AbstractQLayout) = QLayout()
95+
MemoryLayout(::Type{<:LinearAlgebra.QRPackedQ{<:Any,S}}) where {S,T} =
96+
QRPackedQLayout{typeof(MemoryLayout(S)),DenseColumnMajor}()
97+
MemoryLayout(::Type{<:LinearAlgebra.QRCompactWYQ{<:Any,S}}) where {S,T} =
98+
QRCompactWYQLayout{typeof(MemoryLayout(S)),DenseColumnMajor}()
799

100+
adjointlayout(::Type, ::QRPackedQLayout{SLAY,TLAY}) where {SLAY,TLAY} = AdjQRPackedQLayout{SLAY,TLAY}()
101+
adjointlayout(::Type, ::QRCompactWYQLayout{SLAY,TLAY}) where {SLAY,TLAY} = AdjQRCompactWYQLayout{SLAY,TLAY}()
8102

9103
copy(M::Lmul{<:AbstractQLayout}) = copyto!(similar(M), M)
10104

105+
11106
function copyto!(dest::AbstractArray{T}, M::Lmul{<:AbstractQLayout}) where T
12107
A,B = M.A,M.B
13108
if size(dest,1) == size(B,1)
@@ -16,22 +111,167 @@ function copyto!(dest::AbstractArray{T}, M::Lmul{<:AbstractQLayout}) where T
16111
copyto!(view(dest,1:size(B,1),:), B)
17112
zero!(@view(dest[size(B,1)+1:end,:]))
18113
end
19-
materialize!(Lmul(A,dest))
114+
lmul!(A,dest)
20115
end
21116

22117
function copyto!(dest::AbstractArray, M::Ldiv{<:AbstractQLayout})
23118
A,B = M.A,M.B
24119
copyto!(dest, B)
25-
materialize!(Ldiv(A,dest))
120+
ldiv!(A,dest)
26121
end
27122

28123
materialize!(M::Lmul{LAY}) where LAY<:AbstractQLayout = error("Overload materialize!(::Lmul{$(LAY)})")
29124
materialize!(M::Rmul{LAY}) where LAY<:AbstractQLayout = error("Overload materialize!(::Rmul{$(LAY)})")
30125

31-
materialize!(M::Lmul{QLayout}) = LinearAlgebra.lmul!(M.A, M.B)
126+
materialize!(M::Ldiv{<:AbstractQLayout}) = materialize!(Lmul(M.A',M.B))
32127

128+
materialize!(M::Lmul{<:QRPackedQLayout{<:AbstractColumnMajor,<:AbstractColumnMajor},<:AbstractColumnMajor,<:AbstractMatrix{T},<:AbstractVecOrMat{T}}) where T<:BlasFloat =
129+
LAPACK.ormqr!('L','N',M.A.factors,M.A.τ,M.B)
130+
131+
materialize!(M::Lmul{<:QRCompactWYQLayout{<:AbstractColumnMajor,<:AbstractColumnMajor},<:AbstractColumnMajor,<:AbstractMatrix{T},<:AbstractVecOrMat{T}}) where T<:BlasFloat =
132+
LAPACK.gemqrt!('L','N',M.A.factors,M.A.T,M.B)
133+
134+
function materialize!(M::Lmul{<:QRCompactWYQLayout{<:AbstractColumnMajor,<:AbstractColumnMajor},<:AbstractRowMajor,<:AbstractMatrix{T},<:AbstractMatrix{T}}) where T<:BlasReal
135+
LAPACK.gemqrt!('R','T',M.A.factors,M.A.T,transpose(M.B))
136+
M.B
137+
end
138+
function materialize!(M::Lmul{<:QRCompactWYQLayout{<:AbstractColumnMajor,<:AbstractColumnMajor},<:ConjLayout{<:AbstractRowMajor},<:AbstractMatrix{T},<:AbstractMatrix{T}}) where T<:BlasComplex
139+
LAPACK.gemqrt!('R','C',M.A.factors,M.A.T,(M.B)')
140+
M.B
141+
end
33142

34-
materialize!(M::Ldiv{<:AbstractQLayout}) = materialize!(Lmul(M.A',M.B))
143+
144+
function materialize!(M::Lmul{<:QRPackedQLayout})
145+
A,B = M.A, M.B
146+
require_one_based_indexing(B)
147+
mA, nA = size(A.factors)
148+
mB, nB = size(B,1), size(B,2)
149+
if mA != mB
150+
throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but B has dimensions ($mB, $nB)"))
151+
end
152+
Afactors = A.factors
153+
@inbounds begin
154+
for k = min(mA,nA):-1:1
155+
for j = 1:nB
156+
vBj = B[k,j]
157+
for i = k+1:mB
158+
vBj += conj(Afactors[i,k])*B[i,j]
159+
end
160+
vBj = A.τ[k]*vBj
161+
B[k,j] -= vBj
162+
for i = k+1:mB
163+
B[i,j] -= Afactors[i,k]*vBj
164+
end
165+
end
166+
end
167+
end
168+
B
169+
end
170+
171+
172+
### QcB
173+
materialize!(M::Lmul{<:AdjQRPackedQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractStridedLayout,<:AbstractMatrix{T},<:AbstractVecOrMat{T}}) where T<:BlasFloat =
174+
(A = M.A.parent; LAPACK.ormqr!('L','T',A.factors,A.τ,M.B))
175+
materialize!(M::Lmul{<:AdjQRPackedQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractStridedLayout,<:AbstractMatrix{T},<:AbstractVecOrMat{T}}) where T<:BlasComplex =
176+
(A = M.A.parent; LAPACK.ormqr!('L','C',A.factors,A.τ,M.B))
177+
materialize!(M::Lmul{<:AdjQRCompactWYQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractStridedLayout,<:AbstractMatrix{T},<:AbstractVecOrMat{T}}) where T<:BlasFloat =
178+
(A = M.A.parent; LAPACK.gemqrt!('L','T',A.factors,A.T,M.B))
179+
materialize!(M::Lmul{<:AdjQRCompactWYQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractStridedLayout,<:AbstractMatrix{T},<:AbstractVecOrMat{T}}) where T<:BlasComplex =
180+
(A = M.A.parent; LAPACK.gemqrt!('L','C',A.factors,A.T,M.B))
181+
function materialize!(M::Lmul{<:AdjQRPackedQLayout})
182+
adjA,B = M.A, M.B
183+
require_one_based_indexing(B)
184+
A = adjA.parent
185+
mA, nA = size(A.factors)
186+
mB, nB = size(B,1), size(B,2)
187+
if mA != mB
188+
throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but B has dimensions ($mB, $nB)"))
189+
end
190+
Afactors = A.factors
191+
@inbounds begin
192+
for k = 1:min(mA,nA)
193+
for j = 1:nB
194+
vBj = B[k,j]
195+
for i = k+1:mB
196+
vBj += conj(Afactors[i,k])*B[i,j]
197+
end
198+
vBj = conj(A.τ[k])*vBj
199+
B[k,j] -= vBj
200+
for i = k+1:mB
201+
B[i,j] -= Afactors[i,k]*vBj
202+
end
203+
end
204+
end
205+
end
206+
B
207+
end
208+
209+
## AQ
210+
materialize!(M::Rmul{<:AbstractStridedLayout,<:QRPackedQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractVecOrMat{T},<:AbstractMatrix{T}}) where T<:BlasFloat =
211+
LAPACK.ormqr!('R', 'N', M.B.factors, M.B.τ, M.A)
212+
materialize!(M::Rmul{<:AbstractStridedLayout,<:QRCompactWYQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractVecOrMat{T},<:AbstractMatrix{T}}) where T<:BlasFloat =
213+
LAPACK.gemqrt!('R','N', M.B.factors, M.B.T, M.A)
214+
function materialize!(M::Rmul{<:Any,<:QRPackedQLayout})
215+
A,Q = M.A,M.B
216+
mQ, nQ = size(Q.factors)
217+
mA, nA = size(A,1), size(A,2)
218+
if nA != mQ
219+
throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but matrix Q has dimensions ($mQ, $nQ)"))
220+
end
221+
Qfactors = Q.factors
222+
@inbounds begin
223+
for k = 1:min(mQ,nQ)
224+
for i = 1:mA
225+
vAi = A[i,k]
226+
for j = k+1:mQ
227+
vAi += A[i,j]*Qfactors[j,k]
228+
end
229+
vAi = vAi*Q.τ[k]
230+
A[i,k] -= vAi
231+
for j = k+1:nA
232+
A[i,j] -= vAi*conj(Qfactors[j,k])
233+
end
234+
end
235+
end
236+
end
237+
A
238+
end
239+
240+
### AQc
241+
materialize!(M::Rmul{<:AbstractStridedLayout,<:AdjQRPackedQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractVecOrMat{T},<:AbstractMatrix{T}}) where T<:BlasReal =
242+
(B = M.B.parent; LAPACK.ormqr!('R','T',B.factors,B.τ,M.A))
243+
materialize!(M::Rmul{<:AbstractStridedLayout,<:AdjQRPackedQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractVecOrMat{T},<:AbstractMatrix{T}}) where T<:BlasComplex =
244+
(B = M.B.parent; LAPACK.ormqr!('R','C',B.factors,B.τ,M.A))
245+
materialize!(M::Rmul{<:AbstractStridedLayout,<:AdjQRCompactWYQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractVecOrMat{T},<:AbstractMatrix{T}}) where T<:BlasReal =
246+
(B = M.B.parent; LAPACK.gemqrt!('R','T',B.factors,B.T,M.A))
247+
materialize!(M::Rmul{<:AbstractStridedLayout,<:AdjQRCompactWYQLayout{<:AbstractStridedLayout,<:AbstractStridedLayout},<:AbstractVecOrMat{T},<:AbstractMatrix{T}}) where T<:BlasComplex =
248+
(B = M.B.parent; LAPACK.gemqrt!('R','C',B.factors,B.T,M.A))
249+
function materialize!(M::Rmul{<:Any,<:AdjQRPackedQLayout})
250+
A,adjQ = M.A,M.B
251+
Q = adjQ.parent
252+
mQ, nQ = size(Q.factors)
253+
mA, nA = size(A,1), size(A,2)
254+
if nA != mQ
255+
throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but matrix Q has dimensions ($mQ, $nQ)"))
256+
end
257+
Qfactors = Q.factors
258+
@inbounds begin
259+
for k = min(mQ,nQ):-1:1
260+
for i = 1:mA
261+
vAi = A[i,k]
262+
for j = k+1:mQ
263+
vAi += A[i,j]*Qfactors[j,k]
264+
end
265+
vAi = vAi*conj(Q.τ[k])
266+
A[i,k] -= vAi
267+
for j = k+1:nA
268+
A[i,j] -= vAi*conj(Qfactors[j,k])
269+
end
270+
end
271+
end
272+
end
273+
A
274+
end
35275

36276

37277
_qr(layout, axes, A; kwds...) = Base.invoke(qr, Tuple{AbstractMatrix{eltype(A)}}, A; kwds...)
@@ -40,7 +280,25 @@ _lu(layout, axes, A; kwds...) = Base.invoke(lu, Tuple{AbstractMatrix{eltype(A)}}
40280
_lu(layout, axes, A, pivot::P; kwds...) where P = Base.invoke(lu, Tuple{AbstractMatrix{eltype(A)},P}, A, pivot; kwds...)
41281
_qr!(layout, axes, A, args...; kwds...) = error("Overload _qr!(::$(typeof(layout)), axes, A)")
42282
_lu!(layout, axes, A, args...; kwds...) = error("Overload _lu!(::$(typeof(layout)), axes, A)")
43-
_factorize(layout, axes, A) = Base.invoke(factorize, Tuple{AbstractMatrix{eltype(A)}}, A)
283+
_factorize(layout, axes, A) = qr(A) # Default to QR
284+
285+
_inv_eye(_, ::Type{T}, axs::NTuple{2,Base.OneTo{Int}}) where T = Matrix{T}(I, map(length,axs)...)
286+
function _inv_eye(A, ::Type{T}, (rows,cols)) where T
287+
dest = zero!(similar(A, T, (cols,rows)))
288+
dest[diagind(dest)] .= one(T)
289+
dest
290+
end
291+
292+
function _inv(layout, axes, A)
293+
T = eltype(A)
294+
(rows,cols) = axes
295+
n = checksquare(A)
296+
S = typeof(zero(T)/one(T)) # dimensionful
297+
S0 = typeof(zero(T)/oneunit(T)) # dimensionless
298+
dest = _inv_eye(A, S0, (cols,rows))
299+
ldiv!(factorize(convert(AbstractMatrix{S}, A)), dest)
300+
end
301+
44302

45303
macro _layoutfactorizations(Typ)
46304
esc(quote
@@ -50,6 +308,7 @@ macro _layoutfactorizations(Typ)
50308
LinearAlgebra.lu(A::$Typ{T}; kwds...) where T = ArrayLayouts._lu(ArrayLayouts.MemoryLayout(A), axes(A), A; kwds...)
51309
LinearAlgebra.lu!(A::$Typ, args...; kwds...) = ArrayLayouts._lu!(ArrayLayouts.MemoryLayout(A), axes(A), A, args...; kwds...)
52310
LinearAlgebra.factorize(A::$Typ) = ArrayLayouts._factorize(ArrayLayouts.MemoryLayout(A), axes(A), A)
311+
LinearAlgebra.inv(A::$Typ) = ArrayLayouts._inv(ArrayLayouts.MemoryLayout(A), axes(A), A)
53312
end)
54313
end
55314

‎src/ldiv.jl

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -64,20 +64,28 @@ end
6464
end
6565

6666
__ldiv!(::Mat, ::Mat, B) where Mat = error("Overload materialize!(::Ldiv{$(typeof(MemoryLayout(Mat))),$(typeof(MemoryLayout(B)))})")
67-
__ldiv!(_, F, B) = ldiv!(F, B)
67+
__ldiv!(::Mat, ::Mat, B::LayoutArray) where Mat = error("Overload materialize!(::Ldiv{$(typeof(MemoryLayout(Mat))),$(typeof(MemoryLayout(B)))})")
68+
__ldiv!(_, F, B) = LinearAlgebra.ldiv!(F, B)
6869
@inline _ldiv!(A, B) = __ldiv!(A, factorize(A), B)
69-
@inline _ldiv!(A::Factorization, B) = ldiv!(A, B)
70+
@inline _ldiv!(A::Factorization, B) = LinearAlgebra.ldiv!(A, B)
71+
@inline _ldiv!(A::Factorization, B::LayoutArray) = error("Overload materialize!(::Ldiv{$(typeof(MemoryLayout(A))),$(typeof(MemoryLayout(B)))})")
7072

7173
@inline _ldiv!(dest, A, B) = ldiv!(dest, factorize(A), B)
72-
@inline _ldiv!(dest, A::Factorization, B) = ldiv!(dest, A, B)
73-
@inline _ldiv!(dest, A::Transpose{<:Any,<:Factorization}, B) = ldiv!(dest, A, B)
74-
@inline _ldiv!(dest, A::Adjoint{<:Any,<:Factorization}, B) = ldiv!(dest, A, B)
74+
@inline _ldiv!(dest, A::Factorization, B) = LinearAlgebra.ldiv!(dest, A, B)
75+
@inline _ldiv!(dest, A::Transpose{<:Any,<:Factorization}, B) = LinearAlgebra.ldiv!(dest, A, B)
76+
@inline _ldiv!(dest, A::Adjoint{<:Any,<:Factorization}, B) = LinearAlgebra.ldiv!(dest, A, B)
7577

7678
@inline ldiv(A, B) = materialize(Ldiv(A,B))
7779
@inline rdiv(A, B) = materialize(Rdiv(A,B))
7880

81+
@inline ldiv!(A, B) = materialize!(Ldiv(A,B))
82+
@inline rdiv!(A, B) = materialize!(Rdiv(A,B))
83+
84+
@inline ldiv!(C, A, B) = copyto!(C, Ldiv(A,B))
85+
@inline rdiv!(C, A, B) = copyto!(C, Rdiv(A,B))
86+
7987
@inline materialize!(M::Ldiv) = _ldiv!(M.A, M.B)
80-
@inline materialize!(M::Rdiv) = materialize!(Lmul(M.B', M.A'))'
88+
@inline materialize!(M::Rdiv) = ldiv!(M.B', M.A')'
8189
@inline copyto!(dest::AbstractArray, M::Rdiv) = copyto!(dest', Ldiv(M.B', M.A'))'
8290

8391
if VERSION ≥ v"1.1-pre"
@@ -106,6 +114,8 @@ macro _layoutldiv(Typ)
106114
LinearAlgebra.ldiv!(A::$Typ, x::StridedVector) = ArrayLayouts.materialize!(ArrayLayouts.Ldiv(A,x))
107115
LinearAlgebra.ldiv!(A::$Typ, x::StridedMatrix) = ArrayLayouts.materialize!(ArrayLayouts.Ldiv(A,x))
108116

117+
LinearAlgebra.ldiv!(A::Factorization, x::$Typ) = ArrayLayouts.materialize!(ArrayLayouts.Ldiv(A,x))
118+
109119
Base.:\(A::$Typ, x::AbstractVector) = ArrayLayouts.materialize(ArrayLayouts.Ldiv(A,x))
110120
Base.:\(A::$Typ, x::AbstractMatrix) = ArrayLayouts.materialize(ArrayLayouts.Ldiv(A,x))
111121

@@ -138,4 +148,3 @@ macro layoutldiv(Typ)
138148
ArrayLayouts.@_layoutldiv UnitLowerTriangular{T, <:SubArray{T,2,<:$Typ{T}}} where T
139149
end)
140150
end
141-

0 commit comments

Comments
 (0)