Skip to content

Commit cb021da

Browse files
lkdvoslkdvos
committed
Precomputed derivatives IV
Precomputed derivatives V
1 parent b4a9606 commit cb021da

File tree

8 files changed

+498
-90
lines changed

8 files changed

+498
-90
lines changed

‎benchmark/Project.toml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,13 @@ BlockTensorKit = "5f87ffc2-9cf1-4a46-8172-465d160bd8cd"
44
CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
55
DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
66
JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
7+
MKL = "33e6dc65-8f57-5167-99aa-e5a354878fb2"
78
MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502"
89
MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab"
910
TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
1011
TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
12+
TensorKitManifolds = "11fa318c-39cb-4a83-b1ed-cdc7ba1e3684"
13+
14+
[sources]
15+
TensorKit = {rev = "main", url = "https://github.com/QuantumKitHub/TensorKit.jl.git"}
16+
TensorKitManifolds = {rev = "main", url = "https://github.com/QuantumKitHub/TensorKitManifolds.jl"}

‎src/MPSKit.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ include("environments/lazylincocache.jl")
142142
include("algorithms/fixedpoint.jl")
143143
include("algorithms/derivatives/derivatives.jl")
144144
include("algorithms/derivatives/mpo_derivatives.jl")
145+
include("algorithms/derivatives/hamiltonian_derivatives.jl")
145146
include("algorithms/derivatives/projection_derivatives.jl")
146147
include("algorithms/expval.jl")
147148
include("algorithms/toolbox.jl")

‎src/algorithms/derivatives/derivatives.jl‎

Lines changed: 7 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -217,27 +217,15 @@ Base.:*(h::LazySum{<:Union{DerivativeOrMultiplied}}, v) = h(v)
217217
# Operator preparation
218218
# --------------------
219219
"""
220-
prepare_operator!!(O, x, [backend], [allocator]) -> O′, x′
220+
prepare_operator!!(O, [backend], [allocator]) -> O′
221221
222222
Given an operator and vector, try to construct a more efficient representation of that operator for repeated application.
223223
This should always be used in conjunction with [`unprepare_operator!!`](@ref).
224224
"""
225-
function prepare_operator!!(
226-
O, x,
227-
backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
228-
)
229-
return O, x
230-
end
231-
232-
"""
233-
unprepare_operator!!(y, O, x, [backend], [allocator]) -> y′
225+
prepare_operator!!(O, backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) = O
234226

235-
Given the result `y` of applying a prepared operator `O` on vectors like `x`, undo the preparation work on the vector, and clean up the operator.
236-
This should always be used in conjunction with [`prepare_operator!!`](@ref).
237-
"""
238-
function unprepare_operator!!(
239-
y, O, x,
240-
backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
241-
)
242-
return y
243-
end
227+
# to make benchmark scripts run
228+
prepare_operator!!(O, x::AbstractTensorMap, backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) =
229+
prepare_operator!!(O, backend, allocator), x
230+
unprepare_operator!!(y, O, x, backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) =
231+
y

‎src/algorithms/derivatives/hamiltonian_derivatives.jl‎

Lines changed: 337 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,337 @@
1+
"""
2+
JordanMPO_AC_Hamiltonian{O1,O2,O3}
3+
4+
Efficient operator for representing the single-site derivative of a `MPOHamiltonian` sandwiched between two MPSs.
5+
In particular, this operator aims to make maximal use of the structure of the `MPOHamiltonian` to reduce the number of operations required to apply the operator to a tensor.
6+
"""
7+
struct JordanMPO_AC_Hamiltonian{O1, O2, O3} <: DerivativeOperator
8+
D::Union{O1, Missing} # onsite
9+
I::Union{O1, Missing} # not started
10+
E::Union{O1, Missing} # finished
11+
C::Union{O2, Missing} # starting
12+
B::Union{O2, Missing} # ending
13+
A::Union{O3, Missing} # continuing
14+
end
15+
16+
"""
17+
JordanMPO_AC2_Hamiltonian{O1,O2,O3,O4}
18+
19+
Efficient operator for representing the single-site derivative of a `MPOHamiltonian` sandwiched between two MPSs.
20+
In particular, this operator aims to make maximal use of the structure of the `MPOHamiltonian` to reduce the number of operations required to apply the operator to a tensor.
21+
"""
22+
struct JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4} <: DerivativeOperator
23+
II::Union{O1, Missing} # not_started
24+
IC::Union{O2, Missing} # starting right
25+
ID::Union{O1, Missing} # onsite right
26+
CB::Union{O2, Missing} # starting left - ending right
27+
CA::Union{O3, Missing} # starting left - continuing right
28+
AB::Union{O3, Missing} # continuing left - ending right
29+
AA::Union{O4, Missing} # continuing left - continuing right
30+
BE::Union{O2, Missing} # ending left
31+
DE::Union{O1, Missing} # onsite left
32+
EE::Union{O1, Missing} # finished
33+
end
34+
35+
# Constructors
36+
# ------------
37+
const _HAM_MPS_TYPES = Union{
38+
FiniteMPS{<:MPSTensor},
39+
WindowMPS{<:MPSTensor},
40+
InfiniteMPS{<:MPSTensor},
41+
}
42+
43+
function AC_hamiltonian(
44+
site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian, above::_HAM_MPS_TYPES, envs
45+
)
46+
GL = leftenv(envs, site, below)
47+
GR = rightenv(envs, site, below)
48+
W = operator[site]
49+
50+
# starting
51+
if nonzero_length(W.C) > 0
52+
GR_2 = GR[2:(end - 1)]
53+
@plansor starting_[-1 -2; -3 -4] ≔ W.C[-1; -3 1] * GR_2[-4 1; -2]
54+
starting = only(starting_)
55+
else
56+
starting = missing
57+
end
58+
59+
# ending
60+
if nonzero_length(W.B) > 0
61+
GL_2 = GL[2:(end - 1)]
62+
@plansor ending_[-1 -2; -3 -4] ≔ GL_2[-1 1; -3] * W.B[1 -2; -4]
63+
ending = only(ending_)
64+
else
65+
ending = missing
66+
end
67+
68+
# onsite
69+
if nonzero_length(W.D) > 0
70+
if !ismissing(starting)
71+
@plansor starting[-1 -2; -3 -4] += W.D[-1; -3] * removeunit(GR[end], 2)[-4; -2]
72+
onsite = missing
73+
elseif !ismissing(ending)
74+
@plansor ending[-1 -2; -3 -4] += removeunit(GL[1], 2)[-1; -3] * W.D[-2; -4]
75+
onsite = missing
76+
else
77+
onsite = W.D
78+
end
79+
else
80+
onsite = missing
81+
end
82+
83+
# not_started
84+
if size(W, 4) == 1
85+
not_started = missing
86+
elseif !ismissing(starting)
87+
I = id(storagetype(GR[1]), physicalspace(W))
88+
@plansor starting[-1 -2; -3 -4] += I[-1; -3] * removeunit(GR[1], 2)[-4; -2]
89+
not_started = missing
90+
else
91+
not_started = removeunit(GR[1], 2)
92+
end
93+
94+
# finished
95+
if size(W, 1) == 1
96+
finished = missing
97+
elseif !ismissing(ending)
98+
I = id(storagetype(GL[end]), physicalspace(W))
99+
@plansor ending[-1 -2; -3 -4] += removeunit(GL[end], 2)[-1; -3] * I[-2; -4]
100+
finished = missing
101+
else
102+
finished = removeunit(GL[end], 2)
103+
end
104+
105+
# continuing
106+
continuing = MPO_AC_Hamiltonian(GL[2:(end - 1)], W.A, GR[2:(end - 1)])
107+
108+
# obtaining storagetype of environments since these should have already mixed
109+
# the types of the operator and state
110+
S = spacetype(GL)
111+
M = storagetype(GL)
112+
O1 = tensormaptype(S, 1, 1, M)
113+
O2 = tensormaptype(S, 2, 2, M)
114+
O3 = typeof(continuing)
115+
116+
if nonzero_length(W.A) == 0
117+
continuing = missing
118+
end
119+
120+
return JordanMPO_AC_Hamiltonian{O1, O2, O3}(
121+
onsite, not_started, finished, starting, ending, continuing
122+
)
123+
end
124+
125+
function prepare_operator!!(
126+
H::JordanMPO_AC_Hamiltonian{O1, O2, O3}, backend::AbstractBackend, allocator
127+
) where {O1, O2, O3}
128+
O3′ = Core.Compiler.return_type(prepare_operator!!, Tuple{O3, typeof(backend), typeof(allocator)})
129+
A = prepare_operator!!(H.A, backend, allocator)
130+
return JordanMPO_AC_Hamiltonian{O1, O2, O3}(H.D, H.I, H.E, H.C, H.B, A)
131+
end
132+
133+
function AC2_hamiltonian(site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian, above, envs)
134+
GL = leftenv(envs, site, below)
135+
GR = rightenv(envs, site + 1, below)
136+
W1 = operator[site]
137+
W2 = operator[site + 1]
138+
139+
# starting left - continuing right
140+
if nonzero_length(W1.C) > 0 && nonzero_length(W2.A) > 0
141+
@plansor CA_[-1 -2 -3; -4 -5 -6] ≔ W1.C[-1; -4 2] * W2.A[2 -2; -5 1] *
142+
GR[2:(end - 1)][-6 1; -3]
143+
CA = only(CA_)
144+
else
145+
CA = missing
146+
end
147+
148+
# continuing left - ending right
149+
if nonzero_length(W1.A) > 0 && nonzero_length(W2.B) > 0
150+
@plansor AB_[-1 -2 -3; -4 -5 -6] ≔ GL[2:(end - 1)][-1 2; -4] * W1.A[2 -2; -5 1] *
151+
W2.B[1 -3; -6]
152+
AB = only(AB_)
153+
else
154+
AB = missing
155+
end
156+
157+
# middle
158+
if nonzero_length(W1.C) > 0 && nonzero_length(W2.B) > 0
159+
if !ismissing(CA)
160+
@plansor CA[-1 -2 -3; -4 -5 -6] += W1.C[-1; -4 1] * W2.B[1 -2; -5] *
161+
removeunit(GR[end], 2)[-6; -3]
162+
CB = missing
163+
elseif !ismissing(AB)
164+
@plansor AB[-1 -2 -3; -4 -5 -6] += removeunit(GL[1], 2)[-1; -4] *
165+
W1.C[-2; -5 1] * W2.B[1 -3; -6]
166+
CB = missing
167+
else
168+
@plansor CB_[-1 -2; -3 -4] ≔ W1.C[-1; -3 1] * W2.B[1 -2; -4]
169+
CB = only(CB_)
170+
end
171+
else
172+
CB = missing
173+
end
174+
175+
# starting right
176+
if nonzero_length(W2.C) > 0
177+
if !ismissing(CA)
178+
I = id(storagetype(GR[1]), physicalspace(W1))
179+
@plansor CA[-1 -2 -3; -4 -5 -6] += (I[-1; -4] * W2.C[-2; -5 1]) *
180+
GR[2:(end - 1)][-6 1; -3]
181+
IC = missing
182+
else
183+
@plansor IC[-1 -2; -3 -4] ≔ W2.C[-1; -3 1] * GR[2:(end - 1)][-4 1; -2]
184+
end
185+
else
186+
IC = missing
187+
end
188+
189+
# ending left
190+
if nonzero_length(W1.B) > 0
191+
if !ismissing(AB)
192+
I = id(storagetype(GL[end]), physicalspace(W2))
193+
@plansor AB[-1 -2 -3; -4 -5 -6] += GL[2:(end - 1)][-1 1; -4] *
194+
(W1.B[1 -2; -5] * I[-3; -6])
195+
BE = missing
196+
else
197+
@plansor BE[-1 -2; -3 -4] ≔ GL[2:(end - 1)][-1 2; -3] * W1.B[2 -2; -4]
198+
end
199+
else
200+
BE = missing
201+
end
202+
203+
# onsite left
204+
if nonzero_length(W1.D) > 0
205+
if !ismissing(BE)
206+
@plansor BE[-1 -2; -3 -4] += removeunit(GL[1], 2)[-1; -3] * W1.D[-2; -4]
207+
DE = missing
208+
elseif !ismissing(AB)
209+
I = id(storagetype(GL[end]), physicalspace(W2))
210+
@plansor AB[-1 -2 -3; -4 -5 -6] += removeunit(GL[1], 2)[-1; -4] *
211+
(W1.D[-2; -5] * I[-3; -6])
212+
DE = missing
213+
# TODO: could also try in CA?
214+
else
215+
DE = only(W1.D)
216+
end
217+
else
218+
DE = missing
219+
end
220+
221+
# onsite right
222+
if nonzero_length(W2.D) > 0
223+
if !ismissing(IC)
224+
@plansor IC[-1 -2; -3 -4] += W2.D[-1; -3] * removeunit(GR[end], 2)[-4; -2]
225+
ID = missing
226+
elseif !ismissing(CA)
227+
I = id(storagetype(GR[1]), physicalspace(W1))
228+
@plansor CA[-1 -2 -3; -4 -5 -6] += (I[-1; -4] * W2.D[-2; -5]) *
229+
removeunit(GR[end], 2)[-6; -3]
230+
ID = missing
231+
else
232+
ID = only(W2.D)
233+
end
234+
else
235+
ID = missing
236+
end
237+
238+
# finished
239+
if size(W2, 4) == 1
240+
II = missing
241+
elseif !ismissing(IC)
242+
I = id(storagetype(GR[1]), physicalspace(W2))
243+
@plansor IC[-1 -2; -3 -4] += I[-1; -3] * removeunit(GR[1], 2)[-4; -2]
244+
II = missing
245+
elseif !ismissing(CA)
246+
I = id(storagetype(GR[1]), physicalspace(W1) ⊗ physicalspace(W2))
247+
@plansor CA[-1 -2 -3; -4 -5 -6] += I[-1 -2; -4 -5] * removeunit(GR[1], 2)[-6; -3]
248+
II = missing
249+
else
250+
II = transpose(removeunit(GR[1], 2))
251+
end
252+
253+
# unstarted
254+
if size(W1, 1) == 1
255+
EE = missing
256+
elseif !ismissing(BE)
257+
I = id(storagetype(GL[end]), physicalspace(W1))
258+
@plansor BE[-1 -2; -3 -4] += removeunit(GL[end], 2)[-1; -3] * I[-2; -4]
259+
EE = missing
260+
elseif !ismissing(AB)
261+
I = id(storagetype(GL[end]), physicalspace(W1) ⊗ physicalspace(W2))
262+
@plansor AB[-1 -2 -3; -4 -5 -6] += removeunit(GL[end], 2)[-1; -4] * I[-2 -3; -5 -6]
263+
EE = missing
264+
else
265+
EE = removeunit(GL[end], 2)
266+
end
267+
268+
# continuing - continuing
269+
AA = MPO_AC2_Hamiltonian(GL[2:(end - 1)], W1.A, W2.A, GR[2:(end - 1)])
270+
271+
S = spacetype(GL)
272+
M = storagetype(GL)
273+
O1 = tensormaptype(S, 1, 1, M)
274+
O2 = tensormaptype(S, 2, 2, M)
275+
O3 = tensormaptype(S, 3, 3, M)
276+
O4 = typeof(AA)
277+
278+
if nonzero_length(W1.A) == 0 && nonzero_length(W2.A) == 0
279+
AA = missing
280+
else
281+
mask1 = falses(size(W1.A, 1), size(W1.A, 4))
282+
for I in nonzero_keys(W1.A)
283+
mask1[I] = true
284+
end
285+
286+
mask2 = falses(size(W2.A, 1), size(W2.A, 4))
287+
for I in nonzero_keys(W2.A)
288+
mask2[I] = true
289+
end
290+
291+
mask_left = transpose(mask1) & trues(size(mask1, 1))
292+
mask_right = mask2 * trues(size(mask2, 2))
293+
294+
if !any(mask1 * mask2)
295+
AA = missing
296+
end
297+
end
298+
299+
return JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4}(II, IC, ID, CB, CA, AB, AA, BE, DE, EE)
300+
end
301+
302+
function prepare_operator!!(
303+
H::JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4}, backend::AbstractBackend, allocator
304+
) where {O1, O2, O3, O4}
305+
O4′ = Core.Compiler.return_type(prepare_operator!!, Tuple{O4, typeof(backend), typeof(allocator)})
306+
AA = prepare_operator!!(H.AA, backend, allocator)
307+
return JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4′}(H.II, H.IC, H.ID, H.CB, H.CA, H.AB, AA, H.BE, H.DE, H.EE)
308+
end
309+
# Actions
310+
# -------
311+
function (H::JordanMPO_AC_Hamiltonian)(x::MPSTensor)
312+
y = ismissing(H.A) ? zerovector(x) : H.A(x)
313+
314+
ismissing(H.D) || @plansor y[-1 -2; -3] += x[-1 1; -3] * H.D[-2; 1]
315+
ismissing(H.E) || @plansor y[-1 -2; -3] += H.E[-1; 1] * x[1 -2; -3]
316+
ismissing(H.I) || @plansor y[-1 -2; -3] += x[-1 -2; 1] * H.I[1; -3]
317+
ismissing(H.C) || @plansor y[-1 -2; -3] += x[-1 2; 1] * H.C[-2 -3; 2 1]
318+
ismissing(H.B) || @plansor y[-1 -2; -3] += H.B[-1 -2; 1 2] * x[1 2; -3]
319+
320+
return y
321+
end
322+
323+
function (H::JordanMPO_AC2_Hamiltonian)(x::MPOTensor)
324+
y = ismissing(H.AA) ? zerovector(x) : H.AA(x)
325+
326+
ismissing(H.II) || @plansor y[-1 -2; -3 -4] += x[-1 -2; 1 -4] * H.II[-3; 1]
327+
ismissing(H.IC) || @plansor y[-1 -2; -3 -4] += x[-1 -2; 1 2] * H.IC[-4 -3; 2 1]
328+
ismissing(H.ID) || @plansor y[-1 -2; -3 -4] += x[-1 -2; -3 1] * H.ID[-4; 1]
329+
ismissing(H.CB) || @plansor y[-1 -2; -3 -4] += x[-1 1; -3 2] * H.CB[-2 -4; 1 2]
330+
ismissing(H.CA) || @plansor y[-1 -2; -3 -4] += x[-1 1; 3 2] * H.CA[-2 -4 -3; 1 2 3]
331+
ismissing(H.AB) || @plansor y[-1 -2; -3 -4] += x[1 2; -3 3] * H.AB[-1 -2 -4; 1 2 3]
332+
ismissing(H.BE) || @plansor y[-1 -2; -3 -4] += x[1 2; -3 -4] * H.BE[-1 -2; 1 2]
333+
ismissing(H.DE) || @plansor y[-1 -2; -3 -4] += x[-1 1; -3 -4] * H.DE[-2; 1]
334+
ismissing(H.EE) || @plansor y[-1 -2; -3 -4] += x[1 -2; -3 -4] * H.EE[-1; 1]
335+
336+
return y
337+
end

0 commit comments

Comments
 (0)