Skip to content

Commit 408813c

Browse files
Yue-ZhengyuanYue-Zhengyuan
committed
Merge remote-tracking branch 'upstream/master' into su-finiteT
2 parents 85784b3 + 76344b9 commit 408813c

File tree

5 files changed

+154
-80
lines changed

5 files changed

+154
-80
lines changed

‎src/algorithms/contractions/localoperator.jl‎

Lines changed: 98 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ function contract_local_operator(
3333
env::CTMRGEnv,
3434
) where {T, S, N}
3535
static_inds = Val.(inds)
36-
return _contract_local_operator(static_inds, O, ket, bra, env)
36+
return _contract_local_operator(static_inds, O, (ket, bra), env)
3737
end
3838
function contract_local_operator(
3939
inds::NTuple{N, Tuple{Int, Int}},
@@ -67,7 +67,7 @@ function _contract_corner_expr(rowrange, colrange)
6767
return corner_NW, corner_NE, corner_SE, corner_SW
6868
end
6969

70-
function _contract_edge_expr(rowrange, colrange)
70+
function _contract_edge_expr(rowrange, colrange, height)
7171
rmin, rmax = extrema(rowrange)
7272
cmin, cmax = extrema(colrange)
7373
gridsize = (rmax - rmin + 1, cmax - cmin + 1)
@@ -76,11 +76,7 @@ function _contract_edge_expr(rowrange, colrange)
7676
E_N = :(env.edges[NORTH, mod1($(rmin - 1), end), mod1($(cmin + i - 1), end)])
7777
return tensorexpr(
7878
E_N,
79-
(
80-
envlabel(NORTH, i - 1),
81-
virtuallabel(NORTH, :ket, i),
82-
virtuallabel(NORTH, :bra, i),
83-
),
79+
(envlabel(NORTH, i - 1), virtuallabel.(NORTH, ntuple(identity, height), i)...),
8480
envlabel(NORTH, i),
8581
)
8682
end
@@ -89,11 +85,7 @@ function _contract_edge_expr(rowrange, colrange)
8985
E_E = :(env.edges[EAST, mod1($(rmin + i - 1), end), mod1($(cmax + 1), end)])
9086
return tensorexpr(
9187
E_E,
92-
(
93-
envlabel(EAST, i - 1),
94-
virtuallabel(EAST, :ket, i),
95-
virtuallabel(EAST, :bra, i),
96-
),
88+
(envlabel(EAST, i - 1), virtuallabel.(EAST, ntuple(identity, height), i)...),
9789
envlabel(EAST, i),
9890
)
9991
end
@@ -102,11 +94,7 @@ function _contract_edge_expr(rowrange, colrange)
10294
E_S = :(env.edges[SOUTH, mod1($(rmax + 1), end), mod1($(cmin + i - 1), end)])
10395
return tensorexpr(
10496
E_S,
105-
(
106-
envlabel(SOUTH, i),
107-
virtuallabel(SOUTH, :ket, i),
108-
virtuallabel(SOUTH, :bra, i),
109-
),
97+
(envlabel(SOUTH, i), virtuallabel.(SOUTH, ntuple(identity, height), i)...),
11098
envlabel(SOUTH, i - 1),
11199
)
112100
end
@@ -115,20 +103,20 @@ function _contract_edge_expr(rowrange, colrange)
115103
E_W = :(env.edges[WEST, mod1($(rmin + i - 1), end), mod1($(cmin - 1), end)])
116104
return tensorexpr(
117105
E_W,
118-
(envlabel(WEST, i), virtuallabel(WEST, :ket, i), virtuallabel(WEST, :bra, i)),
106+
(envlabel(WEST, i), virtuallabel.(WEST, ntuple(identity, height), i)...),
119107
envlabel(WEST, i - 1),
120108
)
121109
end
122110

123111
return edges_N, edges_E, edges_S, edges_W
124112
end
125113

126-
function _contract_state_expr(rowrange, colrange, cartesian_inds = nothing)
114+
function _contract_state_expr(rowrange, colrange, height, cartesian_inds = nothing)
127115
rmin, rmax = extrema(rowrange)
128116
cmin, cmax = extrema(colrange)
129117
gridsize = (rmax - rmin + 1, cmax - cmin + 1)
130118

131-
return map((:bra, :ket)) do side
119+
return map(1:height) do side
132120
return map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j)
133121
inds_id = if isnothing(cartesian_inds)
134122
nothing
@@ -141,7 +129,7 @@ function _contract_state_expr(rowrange, colrange, cartesian_inds = nothing)
141129
physicallabel(:O, side, inds_id)
142130
end
143131
return tensorexpr(
144-
:($(side)[mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)]),
132+
:(state[$(side)][mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)]),
145133
(physical_label,),
146134
(
147135
if i == 1
@@ -170,35 +158,50 @@ function _contract_state_expr(rowrange, colrange, cartesian_inds = nothing)
170158
end
171159
end
172160

173-
function _contract_pepo_state_expr(rowrange, colrange, cartesian_inds = nothing)
161+
function _contract_pepo_state_expr(rowrange, colrange, height, cartesian_inds = nothing)
162+
@assert height == 1 || height == 2 "Contraction with more than 2 layers of PEPO is unintended."
174163
rmin, rmax = extrema(rowrange)
175164
cmin, cmax = extrema(colrange)
176165
gridsize = (rmax - rmin + 1, cmax - cmin + 1)
177-
178-
return map((:bra, :ket)) do side
166+
return map(1:height) do side
179167
return map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j)
180168
inds_id = if isnothing(cartesian_inds)
181169
nothing
182170
else
183171
findfirst(==(CartesianIndex(rmin + i - 1, cmin + j - 1)), cartesian_inds)
184172
end
185-
physical_label_in = if isnothing(inds_id)
186-
physicallabel(:in, i, j)
187-
else
188-
physicallabel(:Oopen, inds_id)
189-
end
190-
physical_label_out = if isnothing(inds_id)
191-
physicallabel(:out, i, j)
173+
if height == 1
174+
physical_label_in = if isnothing(inds_id)
175+
physicallabel(i, j)
176+
else
177+
physicallabel(:O, 2, inds_id)
178+
end
179+
physical_label_out = if isnothing(inds_id)
180+
physicallabel(i, j)
181+
else
182+
physicallabel(:O, 1, inds_id)
183+
end
184+
tensor_name = :(twistdual(state[1][mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)], 2))
192185
else
193-
physicallabel(:O, side, inds_id)
186+
physical_label_in = if isnothing(inds_id)
187+
physicallabel(:in, i, j)
188+
else
189+
physicallabel(:Oopen, inds_id)
190+
end
191+
physical_label_out = if isnothing(inds_id)
192+
physicallabel(:out, i, j)
193+
else
194+
physicallabel(:O, side, inds_id)
195+
end
196+
tensor_name = if side == 2
197+
:(state[2][mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)])
198+
else
199+
:(twistdual(state[1][mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)], (1, 2)))
200+
end
194201
end
202+
195203
return tensorexpr(
196-
if side == :ket
197-
:(twistdual(ket[mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)], (1, 2)))
198-
else
199-
:(bra[mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)])
200-
end,
201-
(physical_label_out, physical_label_in),
204+
tensor_name, (physical_label_out, physical_label_in),
202205
(
203206
if i == 1
204207
virtuallabel(NORTH, side, j)
@@ -229,7 +232,7 @@ end
229232
@generated function _contract_local_operator(
230233
inds::NTuple{N, Val},
231234
O::AbstractTensorMap{T, S, N, N},
232-
ket::InfinitePEPS, bra::InfinitePEPS,
235+
state::Tuple{InfinitePEPS, InfinitePEPS},
233236
env::CTMRGEnv,
234237
) where {T, S, N}
235238
cartesian_inds = collect(CartesianIndex{2}, map(x -> x.parameters[1], inds.parameters)) # weird hack to extract information from Val
@@ -239,13 +242,13 @@ end
239242
colrange = getindex.(cartesian_inds, 2)
240243

241244
corner_NW, corner_NE, corner_SE, corner_SW = _contract_corner_expr(rowrange, colrange)
242-
edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange)
245+
edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange, 2)
243246
operator = tensorexpr(
244247
:O,
245-
ntuple(i -> physicallabel(:O, :bra, i), N),
246-
ntuple(i -> physicallabel(:O, :ket, i), N),
248+
ntuple(i -> physicallabel(:O, 2, i), N),
249+
ntuple(i -> physicallabel(:O, 1, i), N),
247250
)
248-
bra, ket = _contract_state_expr(rowrange, colrange, cartesian_inds)
251+
ket, bra = _contract_state_expr(rowrange, colrange, 2, cartesian_inds)
249252

250253
multiplication_ex = Expr(
251254
:call, :*,
@@ -274,24 +277,24 @@ function contract_local_norm(
274277
inds::NTuple{N, CartesianIndex{2}}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
275278
) where {N}
276279
static_inds = Val.(inds)
277-
return _contract_local_norm(static_inds, ket, bra, env)
280+
return _contract_local_norm(static_inds, (ket, bra), env)
278281
end
279282
function contract_local_norm(
280283
inds::NTuple{N, Tuple{Int, Int}}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
281284
) where {N}
282285
return contract_local_norm(CartesianIndex.(inds), ket, bra, env)
283286
end
284287
@generated function _contract_local_norm(
285-
inds::NTuple{N, Val}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
288+
inds::NTuple{N, Val}, state::Tuple{InfinitePEPS, InfinitePEPS}, env::CTMRGEnv
286289
) where {N}
287290
cartesian_inds = collect(CartesianIndex{2}, map(x -> x.parameters[1], inds.parameters)) # weird hack to extract information from Val
288291
allunique(cartesian_inds) || throw(ArgumentError("Indices should not overlap: $cartesian_inds."))
289292
rowrange = getindex.(cartesian_inds, 1)
290293
colrange = getindex.(cartesian_inds, 2)
291294

292295
corner_NW, corner_NE, corner_SE, corner_SW = _contract_corner_expr(rowrange, colrange)
293-
edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange)
294-
bra, ket = _contract_state_expr(rowrange, colrange)
296+
edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange, 2)
297+
ket, bra = _contract_state_expr(rowrange, colrange, 2)
295298

296299
multiplication_ex = Expr(
297300
:call, :*,
@@ -309,8 +312,8 @@ end
309312
@doc """
310313
$(SIGNATURES)
311314
312-
Construct the reduced density matrix `ρ` of the PEPS `peps` with open indices `inds` using the environment `env`.
313-
Alternatively, construct the reduced density matrix `ρ` of a full density matrix PEPO with open indices `inds` using the environment `env`.
315+
Construct the reduced density matrix `ρ` of PEPS or PEPO (representing PEPS with ancilla legs) `ket`, `bra` with open indices `inds` using the environment `env`.
316+
Alternatively, construct the reduced density matrix `ρ` of a mixed state specified by the density matrix PEPO `state` with open indices `inds` using the environment `env`.
314317
315318
This works by generating the appropriate contraction on a rectangular patch with its corners
316319
specified by `inds`. The result is normalized such that `tr(ρ) = 1`.
@@ -320,15 +323,22 @@ function reduced_densitymatrix(
320323
inds::NTuple{N, CartesianIndex{2}}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
321324
) where {N}
322325
static_inds = Val.(inds)
323-
return _contract_densitymatrix(static_inds, ket, bra, env)
326+
return _contract_densitymatrix(static_inds, (ket, bra), env)
327+
end
328+
function reduced_densitymatrix(
329+
inds::NTuple{N, CartesianIndex{2}}, state::InfinitePEPO, env::CTMRGEnv
330+
) where {N}
331+
size(state, 3) == 1 || throw(DimensionMismatch("only single-layer densitymatrices are supported"))
332+
static_inds = Val.(inds)
333+
return _contract_densitymatrix(static_inds, (state,), env)
324334
end
325335
function reduced_densitymatrix(
326336
inds::NTuple{N, CartesianIndex{2}}, ket::InfinitePEPO, bra::InfinitePEPO, env::CTMRGEnv
327337
) where {N}
328338
size(ket) == size(bra) || throw(DimensionMismatch("incompatible bra and ket dimensions"))
329339
size(ket, 3) == 1 || throw(DimensionMismatch("only single-layer densitymatrices are supported"))
330340
static_inds = Val.(inds)
331-
return _contract_densitymatrix(static_inds, ket, bra, env)
341+
return _contract_densitymatrix(static_inds, (ket, bra), env)
332342
end
333343
function reduced_densitymatrix(
334344
inds::NTuple{N, Tuple{Int, Int}}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
@@ -343,6 +353,11 @@ end
343353
function reduced_densitymatrix(inds, ket::InfinitePEPS, env::CTMRGEnv)
344354
return reduced_densitymatrix(inds, ket, ket, env)
345355
end
356+
function reduced_densitymatrix(
357+
inds::NTuple{N, Tuple{Int, Int}}, state::InfinitePEPO, env::CTMRGEnv
358+
) where {N}
359+
return reduced_densitymatrix(CartesianIndex.(inds), state, env)
360+
end
346361

347362
# Special case 1x1 density matrix:
348363
# Keep contraction order but try to optimize intermediate permutations:
@@ -395,7 +410,7 @@ function reduced_densitymatrix(
395410
return reduced_densitymatrix1x2(inds[1], ket, bra, env)
396411
else
397412
static_inds = Val.(inds)
398-
return _contract_densitymatrix(static_inds, ket, bra, env)
413+
return _contract_densitymatrix(static_inds, (ket, bra), env)
399414
end
400415
end
401416

@@ -502,7 +517,7 @@ function reduced_densitymatrix1x2(
502517
end
503518

504519
@generated function _contract_densitymatrix(
505-
inds::NTuple{N, Val}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
520+
inds::NTuple{N, Val}, state::Tuple{InfinitePEPS, InfinitePEPS}, env::CTMRGEnv
506521
) where {N}
507522
cartesian_inds = collect(CartesianIndex{2}, map(x -> x.parameters[1], inds.parameters)) # weird hack to extract information from Val
508523
allunique(cartesian_inds) ||
@@ -511,13 +526,13 @@ end
511526
colrange = getindex.(cartesian_inds, 2)
512527

513528
corner_NW, corner_NE, corner_SE, corner_SW = _contract_corner_expr(rowrange, colrange)
514-
edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange)
529+
edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange, 2)
515530
result = tensorexpr(
516531
:ρ,
517-
ntuple(i -> physicallabel(:O, :ket, i), N),
518-
ntuple(i -> physicallabel(:O, :bra, i), N),
532+
ntuple(i -> physicallabel(:O, 1, i), N),
533+
ntuple(i -> physicallabel(:O, 2, i), N),
519534
)
520-
bra, ket = _contract_state_expr(rowrange, colrange, cartesian_inds)
535+
ket, bra = _contract_state_expr(rowrange, colrange, 2, cartesian_inds)
521536

522537
multiplication_ex = Expr(
523538
:call, :*,
@@ -533,29 +548,43 @@ end
533548
end
534549

535550
@generated function _contract_densitymatrix(
536-
inds::NTuple{N, Val}, ket::InfinitePEPO, bra::InfinitePEPO, env::CTMRGEnv
537-
) where {N}
551+
inds::NTuple{N, Val}, state::Tuple{InfinitePEPO, Vararg{InfinitePEPO, M}}, env::CTMRGEnv
552+
) where {N, M}
553+
height = M + 1
554+
@assert height == 1 || height == 2 "Contraction with more than 2 layers of PEPO is unintended."
538555
cartesian_inds = collect(CartesianIndex{2}, map(x -> x.parameters[1], inds.parameters)) # weird hack to extract information from Val
539556
allunique(cartesian_inds) ||
540557
throw(ArgumentError("Indices should not overlap: $cartesian_inds."))
541558
rowrange = getindex.(cartesian_inds, 1)
542559
colrange = getindex.(cartesian_inds, 2)
543560

544561
corner_NW, corner_NE, corner_SE, corner_SW = _contract_corner_expr(rowrange, colrange)
545-
edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange)
562+
edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange, height)
546563
result = tensorexpr(
547564
:ρ,
548-
ntuple(i -> physicallabel(:O, :ket, i), N),
549-
ntuple(i -> physicallabel(:O, :bra, i), N),
565+
ntuple(i -> physicallabel(:O, 1, i), N),
566+
ntuple(i -> physicallabel(:O, 2, i), N),
550567
)
551-
bra, ket = _contract_pepo_state_expr(rowrange, colrange, cartesian_inds)
552568

553-
multiplication_ex = Expr(
554-
:call, :*,
555-
corner_NW, corner_NE, corner_SE, corner_SW,
556-
edges_N..., edges_E..., edges_S..., edges_W...,
557-
ket..., map(x -> Expr(:call, :conj, x), bra)...,
558-
)
569+
pepos = _contract_pepo_state_expr(rowrange, colrange, height, cartesian_inds)
570+
571+
if height == 2
572+
ket, bra = pepos
573+
multiplication_ex = Expr(
574+
:call, :*,
575+
corner_NW, corner_NE, corner_SE, corner_SW,
576+
edges_N..., edges_E..., edges_S..., edges_W...,
577+
ket..., map(x -> Expr(:call, :conj, x), bra)...,
578+
)
579+
else
580+
multiplication_ex = Expr(
581+
:call, :*,
582+
corner_NW, corner_NE, corner_SE, corner_SW,
583+
edges_N..., edges_E..., edges_S..., edges_W...,
584+
pepos[1]...
585+
)
586+
end
587+
559588
multex = :(@autoopt @tensor contractcheck = true $result := $multiplication_ex)
560589
return quote
561590
$(macroexpand(@__MODULE__, multex))

‎src/algorithms/toolbox.jl‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,16 @@ function MPSKit.expectation_value(
1919
return sum(term_vals)
2020
end
2121
MPSKit.expectation_value(peps::InfinitePEPS, O::LocalOperator, env::CTMRGEnv) = expectation_value(peps, O, peps, env)
22+
function MPSKit.expectation_value(
23+
state::InfinitePEPO, O::LocalOperator, env::CTMRGEnv
24+
)
25+
checklattice(state, O)
26+
term_vals = dtmap([O.terms...]) do (inds, operator) # OhMyThreads can't iterate over O.terms directly
27+
ρ = reduced_densitymatrix(inds, state, env)
28+
return trmul(operator, ρ)
29+
end
30+
return sum(term_vals)
31+
end
2232

2333
"""
2434
expectation_value(pf::InfinitePartitionFunction, inds => O, env::CTMRGEnv)

‎src/operators/infinitepepo.jl‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -150,10 +150,8 @@ end
150150
TensorKit.spacetype(::Type{P}) where {P <: InfinitePEPO} = spacetype(eltype(P))
151151
virtualspace(T::InfinitePEPO, r::Int, c::Int, h::Int, dir) = virtualspace(T[r, c, h], dir)
152152
domain_physicalspace(T::InfinitePEPO, r::Int, c::Int) = domain_physicalspace(T[r, c, 1])
153-
function codomain_physicalspace(T::InfinitePEPO, r::Int, c::Int)
154-
return codomain_physicalspace(T[r, c, end])
155-
end
156-
physicalspace(T::InfinitePEPO) = physicalspace.(Ref(T), 1:size(T, 1), 1:size(T, 2))
153+
codomain_physicalspace(T::InfinitePEPO, r::Int, c::Int) = codomain_physicalspace(T[r, c, end])
154+
physicalspace(T::InfinitePEPO) = [physicalspace(T, row, col) for row in axes(T, 1), col in axes(T, 2)]
157155
function physicalspace(T::InfinitePEPO, r::Int, c::Int)
158156
codomain_physicalspace(T, r, c) == domain_physicalspace(T, r, c) || throw(
159157
SpaceMismatch(

0 commit comments

Comments
 (0)