Skip to content

Commit eef41d5

Browse files
mtfishmanmtfishman
authored
Implement blockwise LQ (#124)
1 parent bb62929 commit eef41d5

File tree

4 files changed

+259
-3
lines changed

4 files changed

+259
-3
lines changed

‎Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "BlockSparseArrays"
22
uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
33
authors = ["ITensor developers <[email protected]> and contributors"]
4-
version = "0.6.2"
4+
version = "0.6.3"
55

66
[deps]
77
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

‎src/BlockSparseArrays.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,5 +47,6 @@ include("BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl")
4747
include("factorizations/svd.jl")
4848
include("factorizations/truncation.jl")
4949
include("factorizations/qr.jl")
50+
include("factorizations/lq.jl")
5051

5152
end

‎src/factorizations/lq.jl

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
1+
using MatrixAlgebraKit: MatrixAlgebraKit, lq_compact!, lq_full!
2+
3+
# TODO: this is a hardcoded for now to get around this function not being defined in the
4+
# type domain
5+
function default_blocksparse_lq_algorithm(A::AbstractMatrix; kwargs...)
6+
blocktype(A) <: StridedMatrix{<:LinearAlgebra.BLAS.BlasFloat} ||
7+
error("unsupported type: $(blocktype(A))")
8+
alg = MatrixAlgebraKit.LAPACK_HouseholderLQ(; kwargs...)
9+
return BlockPermutedDiagonalAlgorithm(alg)
10+
end
11+
function MatrixAlgebraKit.default_algorithm(
12+
::typeof(lq_compact!), A::AbstractBlockSparseMatrix; kwargs...
13+
)
14+
return default_blocksparse_lq_algorithm(A; kwargs...)
15+
end
16+
function MatrixAlgebraKit.default_algorithm(
17+
::typeof(lq_full!), A::AbstractBlockSparseMatrix; kwargs...
18+
)
19+
return default_blocksparse_lq_algorithm(A; kwargs...)
20+
end
21+
22+
function similar_output(
23+
::typeof(lq_compact!), A, L_axis, alg::MatrixAlgebraKit.AbstractAlgorithm
24+
)
25+
L = similar(A, axes(A, 1), L_axis)
26+
Q = similar(A, L_axis, axes(A, 2))
27+
return L, Q
28+
end
29+
30+
function similar_output(
31+
::typeof(lq_full!), A, L_axis, alg::MatrixAlgebraKit.AbstractAlgorithm
32+
)
33+
L = similar(A, axes(A, 1), L_axis)
34+
Q = similar(A, L_axis, axes(A, 2))
35+
return L, Q
36+
end
37+
38+
function MatrixAlgebraKit.initialize_output(
39+
::typeof(lq_compact!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
40+
)
41+
bm, bn = blocksize(A)
42+
bmn = min(bm, bn)
43+
44+
brows = eachblockaxis(axes(A, 1))
45+
bcols = eachblockaxis(axes(A, 2))
46+
l_axes = similar(brows, bmn)
47+
48+
# fill in values for blocks that are present
49+
bIs = collect(eachblockstoredindex(A))
50+
browIs = Int.(first.(Tuple.(bIs)))
51+
bcolIs = Int.(last.(Tuple.(bIs)))
52+
for bI in eachblockstoredindex(A)
53+
row, col = Int.(Tuple(bI))
54+
len = minimum(length, (brows[row], bcols[col]))
55+
l_axes[row] = bcols[col][Base.OneTo(len)]
56+
end
57+
58+
# fill in values for blocks that aren't present, pairing them in order of occurence
59+
# this is a convention, which at least gives the expected results for blockdiagonal
60+
emptyrows = setdiff(1:bm, browIs)
61+
emptycols = setdiff(1:bn, bcolIs)
62+
for (row, col) in zip(emptyrows, emptycols)
63+
len = minimum(length, (brows[row], bcols[col]))
64+
l_axes[row] = bcols[col][Base.OneTo(len)]
65+
end
66+
67+
l_axis = mortar_axis(l_axes)
68+
L, Q = similar_output(lq_compact!, A, l_axis, alg)
69+
70+
# allocate output
71+
for bI in eachblockstoredindex(A)
72+
brow, bcol = Tuple(bI)
73+
L[brow, brow], Q[brow, bcol] = MatrixAlgebraKit.initialize_output(
74+
lq_compact!, @view!(A[bI]), alg.alg
75+
)
76+
end
77+
78+
# allocate output for blocks that aren't present -- do we also fill identities here?
79+
for (row, col) in zip(emptyrows, emptycols)
80+
@view!(Q[Block(row, col)])
81+
end
82+
83+
return L, Q
84+
end
85+
86+
function MatrixAlgebraKit.initialize_output(
87+
::typeof(lq_full!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
88+
)
89+
bm, bn = blocksize(A)
90+
91+
bcols = eachblockaxis(axes(A, 2))
92+
l_axes = copy(bcols)
93+
94+
# fill in values for blocks that are present
95+
bIs = collect(eachblockstoredindex(A))
96+
browIs = Int.(first.(Tuple.(bIs)))
97+
bcolIs = Int.(last.(Tuple.(bIs)))
98+
for bI in eachblockstoredindex(A)
99+
row, col = Int.(Tuple(bI))
100+
l_axes[row] = bcols[col]
101+
end
102+
103+
# fill in values for blocks that aren't present, pairing them in order of occurence
104+
# this is a convention, which at least gives the expected results for blockdiagonal
105+
emptyrows = setdiff(1:bm, browIs)
106+
emptycols = setdiff(1:bn, bcolIs)
107+
for (row, col) in zip(emptyrows, emptycols)
108+
l_axes[row] = bcols[col]
109+
end
110+
for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows))
111+
l_axes[bn + i] = bcols[emptycols[k]]
112+
end
113+
114+
l_axis = mortar_axis(l_axes)
115+
L, Q = similar_output(lq_full!, A, l_axis, alg)
116+
117+
# allocate output
118+
for bI in eachblockstoredindex(A)
119+
brow, bcol = Tuple(bI)
120+
L[brow, brow], Q[brow, bcol] = MatrixAlgebraKit.initialize_output(
121+
lq_full!, @view!(A[bI]), alg.alg
122+
)
123+
end
124+
125+
# allocate output for blocks that aren't present -- do we also fill identities here?
126+
for (row, col) in zip(emptyrows, emptycols)
127+
@view!(Q[Block(row, col)])
128+
end
129+
# also handle extra rows/cols
130+
for (i, k) in enumerate((length(emptyrows) + 1):length(emptycols))
131+
@view!(Q[Block(bm + i, emptycols[k])])
132+
end
133+
134+
return L, Q
135+
end
136+
137+
function MatrixAlgebraKit.check_input(
138+
::typeof(lq_compact!), A::AbstractBlockSparseMatrix, LQ
139+
)
140+
L, Q = LQ
141+
@assert isa(L, AbstractBlockSparseMatrix) && isa(Q, AbstractBlockSparseMatrix)
142+
@assert eltype(A) == eltype(L) == eltype(Q)
143+
@assert axes(A, 1) == axes(L, 1) && axes(A, 2) == axes(Q, 2)
144+
@assert axes(L, 2) == axes(Q, 1)
145+
146+
return nothing
147+
end
148+
149+
function MatrixAlgebraKit.check_input(::typeof(lq_full!), A::AbstractBlockSparseMatrix, LQ)
150+
L, Q = LQ
151+
@assert isa(L, AbstractBlockSparseMatrix) && isa(Q, AbstractBlockSparseMatrix)
152+
@assert eltype(A) == eltype(L) == eltype(Q)
153+
@assert axes(A, 1) == axes(L, 1) && axes(A, 2) == axes(Q, 2)
154+
@assert axes(L, 2) == axes(Q, 1)
155+
156+
return nothing
157+
end
158+
159+
function MatrixAlgebraKit.lq_compact!(
160+
A::AbstractBlockSparseMatrix, LQ, alg::BlockPermutedDiagonalAlgorithm
161+
)
162+
MatrixAlgebraKit.check_input(lq_compact!, A, LQ)
163+
L, Q = LQ
164+
165+
# do decomposition on each block
166+
for bI in eachblockstoredindex(A)
167+
brow, bcol = Tuple(bI)
168+
lq = (@view!(L[brow, brow]), @view!(Q[brow, bcol]))
169+
lq′ = lq_compact!(@view!(A[bI]), lq, alg.alg)
170+
@assert lq === lq′ "lq_compact! might not be in-place"
171+
end
172+
173+
# fill in identities for blocks that aren't present
174+
bIs = collect(eachblockstoredindex(A))
175+
browIs = Int.(first.(Tuple.(bIs)))
176+
bcolIs = Int.(last.(Tuple.(bIs)))
177+
emptyrows = setdiff(1:blocksize(A, 1), browIs)
178+
emptycols = setdiff(1:blocksize(A, 2), bcolIs)
179+
# needs copyto! instead because size(::LinearAlgebra.I) doesn't work
180+
# Q[Block(row, col)] = LinearAlgebra.I
181+
for (row, col) in zip(emptyrows, emptycols)
182+
copyto!(@view!(Q[Block(row, col)]), LinearAlgebra.I)
183+
end
184+
185+
return LQ
186+
end
187+
188+
function MatrixAlgebraKit.lq_full!(
189+
A::AbstractBlockSparseMatrix, LQ, alg::BlockPermutedDiagonalAlgorithm
190+
)
191+
MatrixAlgebraKit.check_input(lq_full!, A, LQ)
192+
L, Q = LQ
193+
194+
# do decomposition on each block
195+
for bI in eachblockstoredindex(A)
196+
brow, bcol = Tuple(bI)
197+
lq = (@view!(L[brow, brow]), @view!(Q[brow, bcol]))
198+
lq′ = lq_full!(@view!(A[bI]), lq, alg.alg)
199+
@assert lq === lq′ "lq_full! might not be in-place"
200+
end
201+
202+
# fill in identities for blocks that aren't present
203+
bIs = collect(eachblockstoredindex(A))
204+
browIs = Int.(first.(Tuple.(bIs)))
205+
bcolIs = Int.(last.(Tuple.(bIs)))
206+
emptyrows = setdiff(1:blocksize(A, 1), browIs)
207+
emptycols = setdiff(1:blocksize(A, 2), bcolIs)
208+
# needs copyto! instead because size(::LinearAlgebra.I) doesn't work
209+
# Q[Block(row, col)] = LinearAlgebra.I
210+
for (row, col) in zip(emptyrows, emptycols)
211+
copyto!(@view!(Q[Block(row, col)]), LinearAlgebra.I)
212+
end
213+
214+
# also handle extra rows/cols
215+
bm = blocksize(A, 1)
216+
for (i, k) in enumerate((length(emptyrows) + 1):length(emptycols))
217+
copyto!(@view!(Q[Block(bm + i, emptycols[k])]), LinearAlgebra.I)
218+
end
219+
220+
return LQ
221+
end

‎test/test_factorizations.jl

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,15 @@
11
using BlockArrays: Block, BlockedMatrix, BlockedVector, blocks, mortar
22
using BlockSparseArrays: BlockSparseArray, BlockDiagonal, eachblockstoredindex
33
using MatrixAlgebraKit:
4-
qr_compact, qr_full, svd_compact, svd_full, svd_trunc, truncrank, trunctol
4+
lq_compact,
5+
lq_full,
6+
qr_compact,
7+
qr_full,
8+
svd_compact,
9+
svd_full,
10+
svd_trunc,
11+
truncrank,
12+
trunctol
513
using LinearAlgebra: LinearAlgebra
614
using Random: Random
715
using Test: @inferred, @testset, @test
@@ -157,7 +165,7 @@ end
157165
end
158166

159167
@testset "qr_compact" for T in (Float32, Float64, ComplexF32, ComplexF64)
160-
for i in [1, 2], j in [1, 2], k in [1, 2], l in [1, 2]
168+
for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3]
161169
A = BlockSparseArray{T}(undef, ([i, j], [k, l]))
162170
A[Block(1, 1)] = randn(T, i, k)
163171
A[Block(2, 2)] = randn(T, j, l)
@@ -181,3 +189,29 @@ end
181189
@test A ≈ Q * R
182190
end
183191
end
192+
193+
@testset "lq_compact" for T in (Float32, Float64, ComplexF32, ComplexF64)
194+
for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3]
195+
A = BlockSparseArray{T}(undef, ([i, j], [k, l]))
196+
A[Block(1, 1)] = randn(T, i, k)
197+
A[Block(2, 2)] = randn(T, j, l)
198+
L, Q = lq_compact(A)
199+
@test Matrix(Q * Q') ≈ LinearAlgebra.I
200+
@test A ≈ L * Q
201+
end
202+
end
203+
204+
@testset "lq_full" for T in (Float32, Float64, ComplexF32, ComplexF64)
205+
for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3]
206+
A = BlockSparseArray{T}(undef, ([i, j], [k, l]))
207+
A[Block(1, 1)] = randn(T, i, k)
208+
A[Block(2, 2)] = randn(T, j, l)
209+
L, Q = lq_full(A)
210+
L′, Q′ = lq_full(Matrix(A))
211+
@test size(L) == size(L′)
212+
@test size(Q) == size(Q′)
213+
@test Matrix(Q * Q') ≈ LinearAlgebra.I
214+
@test Matrix(Q'Q) ≈ LinearAlgebra.I
215+
@test A ≈ L * Q
216+
end
217+
end

0 commit comments

Comments
 (0)