Skip to content

Commit 5489eb7

Browse files
lkdvoslkdvos
committed
small fixes
1 parent 991a032 commit 5489eb7

3 files changed

Lines changed: 208 additions & 24 deletions

File tree

‎src/linalg/factorizations.jl‎

Lines changed: 205 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,87 @@
11
using MatrixAlgebraKit
22
using MatrixAlgebraKit: AbstractAlgorithm, YALAPACK.BlasMat, Algorithm
33
import MatrixAlgebraKit as MAK
4+
using TensorKit.Factorizations: @check_space, @check_scalar
45

56
# Type piracy for defining the MAK rules on BlockArrays!
67
# -----------------------------------------------------
78

8-
MAK.@algdef BlockAlgorithm
9-
109
const BlockBlasMat{T <: MAK.BlasFloat} = BlockMatrix{T}
1110

11+
function MatrixAlgebraKit.one!(A::BlockBlasMat)
12+
_one, _zero = one(eltype(A)), zero(eltype(A))
13+
@inbounds for j in axes(A, 2), i in axes(A, 1)
14+
A[i, j] = ifelse(i == j, _one, _zero)
15+
end
16+
return A
17+
end
18+
1219
for f in (
13-
:svd_compact, :svd_full, :svd_trunc, :svd_vals, :qr_compact, :qr_full, :qr_null,
14-
:lq_compact, :lq_full, :lq_null, :eig_full, :eig_trunc, :eig_vals, :eigh_full,
15-
:eigh_trunc, :eigh_vals, :left_polar, :right_polar,
20+
:svd_compact, :svd_full, :svd_vals, :qr_compact, :qr_full, :qr_null,
21+
:lq_compact, :lq_full, :lq_null, :eig_full, :eig_vals, :eigh_full,
22+
:eigh_vals, :left_polar, :right_polar,
1623
)
1724
f! = Symbol(f, :!)
25+
@eval MAK.$f!(t::BlockBlasMat, F, alg::AbstractAlgorithm) = $f!(Array(t), alg)
26+
@eval MAK.$f!(t::BlockBlasMat, F, alg::MAK.DiagonalAlgorithm) = error("Not diagonal")
27+
end
28+
29+
# disambiguations
30+
for (f!, Alg) in (
31+
(:lq_compact!, :LAPACK_HouseholderLQ), (:lq_full!, :LAPACK_HouseholderLQ), (:lq_null!, :LAPACK_HouseholderLQ),
32+
(:lq_compact!, :LQViaTransposedQR), (:lq_full!, :LQViaTransposedQR), (:lq_null!, :LQViaTransposedQR),
33+
(:qr_compact!, :LAPACK_HouseholderQR), (:qr_full!, :LAPACK_HouseholderQR), (:qr_null!, :LAPACK_HouseholderQR),
34+
(:svd_compact!, :LAPACK_SVDAlgorithm), (:svd_full!, :LAPACK_SVDAlgorithm), (:svd_vals!, :LAPACK_SVDAlgorithm),
35+
(:eig_full!, :LAPACK_EigAlgorithm), (:eig_trunc!, :TruncatedAlgorithm), (:eig_vals!, :LAPACK_EigAlgorithm),
36+
(:eigh_full!, :LAPACK_EighAlgorithm), (:eigh_trunc!, :TruncatedAlgorithm), (:eigh_vals!, :LAPACK_EighAlgorithm),
37+
(:left_polar!, :PolarViaSVD), (:right_polar!, :PolarViaSVD),
38+
)
39+
@eval MAK.$f!(t::BlockBlasMat, F, alg::MAK.$Alg) = $f!(Array(t), alg)
40+
end
1841

19-
@eval MAK.copy_input(::typeof($f), A::BlockBlasMat) = Array(A)
20-
@eval MAK.$f!(t::BlockBlasMat, F, alg::BlockAlgorithm) = $f(t; alg.kwargs...)
42+
const GPU_QRAlgorithm = Union{MAK.CUSOLVER_HouseholderQR, MAK.ROCSOLVER_HouseholderQR}
43+
for f! in (:qr_compact!, :qr_full!, :qr_null!)
44+
@eval MAK.$f!(t::BlockBlasMat, QR, alg::GPU_QRAlgorithm) = error()
2145
end
2246

47+
for (f!, Alg) in (
48+
(:eigh_full!, :GPU_EighAlgorithm), (:eigh_vals!, :GPU_EighAlgorithm),
49+
(:eig_full!, :GPU_EigAlgorithm), (:eig_vals!, :GPU_EigAlgorithm),
50+
(:svd_full!, :GPU_SVDAlgorithm), (:svd_compact!, :GPU_SVDAlgorithm), (:svd_vals!, :GPU_SVDAlgorithm),
51+
)
52+
@eval MAK.$f!(t::BlockBlasMat, F, alg::MAK.$Alg) = error()
53+
end
54+
55+
2356
for f in (:qr, :lq, :eig, :eigh, :gen_eig, :svd, :polar)
2457
default_f_algorithm = Symbol(:default_, f, :_algorithm)
2558
@eval MAK.$default_f_algorithm(::Type{<:BlockBlasMat{T}}; kwargs...) where {T} =
26-
BlockAlgorithm(; kwargs...)
59+
MAK.$default_f_algorithm(Matrix{T}; kwargs...)
2760
end
2861

2962
# Make sure sparse blocktensormaps have dense outputs
63+
function MAK.check_input(::typeof(qr_full!), t::AbstractBlockTensorMap, QR, ::AbstractAlgorithm)
64+
Q, R = QR
65+
66+
# type checks
67+
@assert Q isa AbstractTensorMap
68+
@assert R isa AbstractTensorMap
69+
70+
# scalartype checks
71+
@check_scalar Q t
72+
@check_scalar R t
73+
74+
# space checks
75+
V_Q = TK.oplus(fuse(codomain(t)))
76+
@check_space(Q, codomain(t) ← V_Q)
77+
@check_space(R, V_Q ← domain(t))
78+
79+
return nothing
80+
end
81+
MAK.check_input(::typeof(qr_full!), t::AbstractBlockTensorMap, QR, ::DiagonalAlgorithm) = error()
82+
3083
function MAK.initialize_output(::typeof(qr_full!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
31-
V_Q = fuse(codomain(t))
84+
V_Q = TK.oplus(fuse(codomain(t)))
3285
Q = dense_similar(t, codomain(t) ← V_Q)
3386
R = dense_similar(t, V_Q ← domain(t))
3487
return Q, R
@@ -46,8 +99,28 @@ function MAK.initialize_output(::typeof(qr_null!), t::AbstractBlockTensorMap, ::
4699
return N
47100
end
48101

102+
function MAK.check_input(::typeof(lq_full!), t::AbstractBlockTensorMap, LQ, ::AbstractAlgorithm)
103+
L, Q = LQ
104+
105+
# type checks
106+
@assert L isa AbstractTensorMap
107+
@assert Q isa AbstractTensorMap
108+
109+
# scalartype checks
110+
@check_scalar L t
111+
@check_scalar Q t
112+
113+
# space checks
114+
V_Q = TK.oplus(fuse(domain(t)))
115+
@check_space(L, codomain(t) ← V_Q)
116+
@check_space(Q, V_Q ← domain(t))
117+
118+
return nothing
119+
end
120+
MAK.check_input(::typeof(lq_full!), t::AbstractBlockTensorMap, LQ, ::DiagonalAlgorithm) = error()
121+
49122
function MAK.initialize_output(::typeof(lq_full!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
50-
V_Q = fuse(domain(t))
123+
V_Q = TK.oplus(fuse(domain(t)))
51124
L = dense_similar(t, codomain(t) ← V_Q)
52125
Q = dense_similar(t, V_Q ← domain(t))
53126
return L, Q
@@ -65,11 +138,50 @@ function MAK.initialize_output(::typeof(lq_null!), t::AbstractBlockTensorMap, ::
65138
return N
66139
end
67140

141+
function MAK.check_input(::typeof(MAK.left_orth_polar!), t::AbstractBlockTensorMap, WP, ::AbstractAlgorithm)
142+
codomain(t) ≿ domain(t) ||
143+
throw(ArgumentError("Polar decomposition requires `codomain(t) ≿ domain(t)`"))
144+
145+
W, P = WP
146+
@assert W isa AbstractTensorMap
147+
@assert P isa AbstractTensorMap
148+
149+
# scalartype checks
150+
@check_scalar W t
151+
@check_scalar P t
152+
153+
# space checks
154+
VW = TK.oplus(fuse(domain(t)))
155+
@check_space(W, codomain(t) ← VW)
156+
@check_space(P, VW ← domain(t))
157+
158+
return nothing
159+
end
68160
function MAK.initialize_output(::typeof(left_polar!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
69161
W = dense_similar(t, space(t))
70162
P = dense_similar(t, domain(t) ← domain(t))
71163
return W, P
72164
end
165+
166+
function MAK.check_input(::typeof(MAK.right_orth_polar!), t::AbstractBlockTensorMap, PWᴴ, ::AbstractAlgorithm)
167+
codomain(t) ≾ domain(t) ||
168+
throw(ArgumentError("Polar decomposition requires `domain(t) ≿ codomain(t)`"))
169+
170+
P, Wᴴ = PWᴴ
171+
@assert P isa AbstractTensorMap
172+
@assert Wᴴ isa AbstractTensorMap
173+
174+
# scalartype checks
175+
@check_scalar P t
176+
@check_scalar Wᴴ t
177+
178+
# space checks
179+
VW = TK.oplus(fuse(codomain(t)))
180+
@check_space(P, codomain(t) ← VW)
181+
@check_space(Wᴴ, VW ← domain(t))
182+
183+
return nothing
184+
end
73185
function MAK.initialize_output(::typeof(right_polar!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
74186
P = dense_similar(t, codomain(t) ← codomain(t))
75187
Wᴴ = dense_similar(t, space(t))
@@ -87,24 +199,103 @@ function MAK.initialize_output(::typeof(right_null!), t::AbstractBlockTensorMap,
87199
return dense_similar(t, V_N ← domain(t))
88200
end
89201

202+
function MAK.check_input(::typeof(eigh_full!), t::AbstractBlockTensorMap, DV, ::AbstractAlgorithm)
203+
domain(t) == codomain(t) ||
204+
throw(ArgumentError("Eigenvalue decomposition requires square input tensor"))
205+
206+
D, V = DV
207+
208+
# type checks
209+
@assert D isa DiagonalTensorMap
210+
@assert V isa AbstractTensorMap
211+
212+
# scalartype checks
213+
@check_scalar D t real
214+
@check_scalar V t
215+
216+
# space checks
217+
V_D = TK.oplus(fuse(domain(t)))
218+
@check_space(D, V_D ← V_D)
219+
@check_space(V, codomain(t) ← V_D)
220+
221+
return nothing
222+
end
223+
MAK.check_input(::typeof(eigh_full!), t::AbstractBlockTensorMap, DV, ::DiagonalAlgorithm) = error()
224+
225+
function MAK.check_input(::typeof(eigh_vals!), t::AbstractBlockTensorMap, D, ::AbstractAlgorithm)
226+
@check_scalar D t real
227+
@assert D isa DiagonalTensorMap
228+
V_D = TK.oplus(fuse(domain(t)))
229+
@check_space(D, V_D ← V_D)
230+
return nothing
231+
end
232+
MAK.check_input(::typeof(eigh_vals!), t::AbstractBlockTensorMap, D, ::DiagonalAlgorithm) = error()
233+
90234
function MAK.initialize_output(::typeof(eigh_full!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
91-
V_D = fuse(domain(t))
235+
V_D = TK.oplus(fuse(domain(t)))
92236
T = real(scalartype(t))
93237
D = DiagonalTensorMap{T}(undef, V_D)
94238
V = dense_similar(t, codomain(t) ← V_D)
95239
return D, V
96240
end
241+
242+
243+
function MAK.check_input(::typeof(eig_full!), t::AbstractBlockTensorMap, DV, ::AbstractAlgorithm)
244+
domain(t) == codomain(t) ||
245+
throw(ArgumentError("Eigenvalue decomposition requires square input tensor"))
246+
247+
D, V = DV
248+
249+
# type checks
250+
@assert D isa DiagonalTensorMap
251+
@assert V isa AbstractTensorMap
252+
253+
# scalartype checks
254+
@check_scalar D t complex
255+
@check_scalar V t complex
256+
257+
# space checks
258+
V_D = TK.oplus(fuse(domain(t)))
259+
@check_space(D, V_D ← V_D)
260+
@check_space(V, codomain(t) ← V_D)
261+
262+
return nothing
263+
end
264+
MAK.check_input(::typeof(eig_full!), t::AbstractBlockTensorMap, DV, ::DiagonalAlgorithm) = error()
265+
97266
function MAK.initialize_output(::typeof(eig_full!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
98-
V_D = fuse(domain(t))
267+
V_D = TK.oplus(fuse(domain(t)))
99268
Tc = complex(scalartype(t))
100269
D = DiagonalTensorMap{Tc}(undef, V_D)
101270
V = dense_similar(t, Tc, codomain(t) ← V_D)
102271
return D, V
103272
end
104273

274+
function MAK.check_input(::typeof(svd_full!), t::AbstractBlockTensorMap, USVᴴ, ::AbstractAlgorithm)
275+
U, S, Vᴴ = USVᴴ
276+
277+
# type checks
278+
@assert U isa AbstractTensorMap
279+
@assert S isa AbstractTensorMap
280+
@assert Vᴴ isa AbstractTensorMap
281+
282+
# scalartype checks
283+
@check_scalar U t
284+
@check_scalar S t real
285+
@check_scalar Vᴴ t
286+
287+
# space checks
288+
V_cod = TK.oplus(fuse(codomain(t)))
289+
V_dom = TK.oplus(fuse(domain(t)))
290+
@check_space(U, codomain(t) ← V_cod)
291+
@check_space(S, V_cod ← V_dom)
292+
@check_space(Vᴴ, V_dom ← domain(t))
293+
294+
return nothing
295+
end
105296
function MAK.initialize_output(::typeof(svd_full!), t::AbstractBlockTensorMap, ::AbstractAlgorithm)
106-
V_cod = fuse(codomain(t))
107-
V_dom = fuse(domain(t))
297+
V_cod = TK.oplus(fuse(codomain(t)))
298+
V_dom = TK.oplus(fuse(domain(t)))
108299
U = dense_similar(t, codomain(t) ← V_cod)
109300
S = similar(t, real(scalartype(t)), V_cod ← V_dom)
110301
Vᴴ = dense_similar(t, V_dom ← domain(t))

‎src/tensors/abstractblocktensor/abstracttensormap.jl‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -95,14 +95,6 @@ Base.@assume_effects :foldable function TensorKit.blocktype(::Type{TT}) where {T
9595
return BlockMatrix{T, Matrix{B}, BS}
9696
end
9797

98-
# piracy:
99-
function TK.blocktype(
100-
::Type{Union{TM, TB}}
101-
) where {T, S, TM <: TensorMap{T, S}, TB <: BraidingTensor{T, S}}
102-
return Union{TK.blocktype(TM), TK.blocktype(TB)}
103-
end
104-
105-
10698
function Base.iterate(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state...)
10799
next = iterate(iter.structure, state...)
108100
isnothing(next) && return next

‎src/vectorspaces/sumspace.jl‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -159,8 +159,9 @@ TensorKit.:⊕(V::SumSpace{S}) where {S} = reduce(TK.oplus, V.spaces; init = isd
159159
TensorKit.:⊕(V1::SumSpace{S}, V2::SumSpace{S}...) where {S} = TensorKit.oplus(⊕(V1, V2...))
160160
#! format: on
161161

162-
TensorKit.fuse(V1::SumSpace) = fuse(TK.oplus(V1))
163-
TK.fuse(V1::S, V2::S) where {S <: SumSpace} = fuse(TK.oplus(V1), TK.oplus(V2))
162+
function TensorKit.fuse(V1::S, V2::S) where {S <: SumSpace}
163+
return SumSpace(vec([fuse(v1, v2) for (v1, v2) in Base.product(V1.spaces, V2.spaces)]))
164+
end
164165

165166
Base.oneunit(S::Type{<:SumSpace}) = SumSpace(oneunit(eltype(S)))
166167
Base.zero(V::SumSpace{S}) where {S} = SumSpace{S}(; dual = isdual(V))

0 commit comments

Comments
 (0)