Skip to content

Commit c1b115a

Browse files
JanLucaJanLuca
committed
Add helper to calculate partially traced four site density matrices
1 parent acf453a commit c1b115a

File tree

2 files changed

+298
-0
lines changed

2 files changed

+298
-0
lines changed

‎varipeps/expectation/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from . import helpers
12
from . import spiral_helpers
23
from . import structure_factor
34
from . import model

‎varipeps/expectation/helpers.py‎

Lines changed: 297 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,297 @@
1+
import jax.numpy as jnp
2+
from jax import jit
3+
4+
from varipeps.peps import PEPS_Tensor
5+
from varipeps.contractions import apply_contraction, Definitions
6+
7+
from typing import Sequence, Tuple
8+
9+
10+
def partially_traced_four_site_density_matrices(
11+
peps_tensors: Sequence[jnp.ndarray],
12+
peps_tensor_objs: Sequence[PEPS_Tensor],
13+
real_physical_dimension: int,
14+
num_coarse_grained_physical_indices: int,
15+
open_physical_indices: Tuple[Tuple[int], Tuple[int], Tuple[int], Tuple[int]],
16+
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
17+
t_top_left, t_top_right, t_bottom_left, t_bottom_right = peps_tensors
18+
t_obj_top_left, t_obj_top_right, t_obj_bottom_left, t_obj_bottom_right = (
19+
peps_tensor_objs
20+
)
21+
top_left_i, top_right_i, bottom_left_i, bottom_right_i = open_physical_indices
22+
23+
for t in ("t_top_left", "t_top_right", "t_bottom_left", "t_bottom_right"):
24+
locals()[t] = locals()[t].reshape(
25+
t_left.shape[0],
26+
t_left.shape[1],
27+
*((real_physical_dimension,) * num_coarse_grained_physical_indices),
28+
t_left.shape[3],
29+
t_left.shape[4],
30+
)
31+
32+
if (
33+
not hasattr(
34+
Definitions,
35+
f"partially_traced_four_site_density_matrices_top_left_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{top_left_i}",
36+
)
37+
or not hasattr(
38+
Definitions,
39+
f"partially_traced_four_site_density_matrices_top_right_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{top_right_i}",
40+
)
41+
or not hasattr(
42+
Definitions,
43+
f"partially_traced_four_site_density_matrices_bottom_left_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{bottom_left_i}",
44+
)
45+
or not hasattr(
46+
Definitions,
47+
f"partially_traced_four_site_density_matrices_bottom_right_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{bottom_right_i}",
48+
)
49+
):
50+
phys_contraction_i_top_left = list(
51+
range(7, 7 + num_coarse_grained_physical_indices - len(top_left_i))
52+
)
53+
phys_contraction_i_conj_top_left = list(
54+
range(7, 7 + num_coarse_grained_physical_indices - len(top_left_i))
55+
)
56+
for pos, i in enumerate(top_left_i):
57+
phys_contraction_i_top_left.insert(i - 1, -(pos + 1))
58+
phys_contraction_i_conj_top_left.insert(i - 1, -len(top_left_i) - (pos + 1))
59+
phys_contraction_i_top_left = tuple(phys_contraction_i_top_left)
60+
phys_contraction_i_conj_top_left = tuple(phys_contraction_i_conj_top_left)
61+
62+
contraction_top_left = {
63+
"tensors": [["tensor", "tensor_conj", "T4", "C1", "T1"]],
64+
"network": [
65+
[
66+
(3, -2 * len(top_left_i) - 2)
67+
+ phys_contraction_i_top_left
68+
+ (-2 * len(top_left_i) - 5, 4), # tensor
69+
(5, -2 * len(top_left_i) - 3)
70+
+ phys_contraction_i_conj_top_left
71+
+ (-2 * len(top_left_i) - 6, 6), # tensor_conj
72+
(-2 * len(top_left_i) - 1, 5, 3, 1), # T4
73+
(1, 2), # C1
74+
(2, 4, 6, -2 * len(top_left_i) - 4), # T1
75+
]
76+
],
77+
}
78+
Definitions._process_def(
79+
contraction_top_left,
80+
f"partially_traced_four_site_density_matrices_top_left_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{top_left_i}",
81+
)
82+
setattr(
83+
Definitions,
84+
f"partially_traced_four_site_density_matrices_top_left_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{top_left_i}",
85+
contraction_top_left,
86+
)
87+
88+
phys_contraction_i_top_right = list(
89+
range(7, 7 + num_coarse_grained_physical_indices - len(top_right_i))
90+
)
91+
phys_contraction_i_conj_top_right = list(
92+
range(7, 7 + num_coarse_grained_physical_indices - len(top_right_i))
93+
)
94+
for pos, i in enumerate(top_right_i):
95+
phys_contraction_i_top_right.insert(i - 1, -(pos + 1))
96+
phys_contraction_i_conj_top_right.insert(
97+
i - 1, -len(top_right_i) - (pos + 1)
98+
)
99+
phys_contraction_i_top_right = tuple(phys_contraction_i_top_right)
100+
phys_contraction_i_conj_top_right = tuple(phys_contraction_i_conj_top_right)
101+
102+
contraction_top_right = {
103+
"tensors": [["tensor", "tensor_conj", "T1", "C2", "T2"]],
104+
"network": [
105+
[
106+
(-2 * len(top_right_i) - 2, -2 * len(top_right_i) - 6)
107+
+ phys_contraction_i_top_right
108+
+ (4, 3), # tensor
109+
(-2 * len(top_right_i) - 3, -2 * len(top_right_i) - 5)
110+
+ phys_contraction_i_conj_top_right
111+
+ (6, 5), # tensor_conj
112+
(-2 * len(top_right_i) - 1, 3, 5, 1), # T1
113+
(1, 2), # C2
114+
(4, 6, -2 * len(top_right_i) - 4, 2), # T2
115+
]
116+
],
117+
}
118+
Definitions._process_def(
119+
contraction_top_right,
120+
f"partially_traced_four_site_density_matrices_top_right_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{top_right_i}",
121+
)
122+
setattr(
123+
Definitions,
124+
f"partially_traced_four_site_density_matrices_top_right_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{top_right_i}",
125+
contraction_top_right,
126+
)
127+
128+
phys_contraction_i_bottom_left = list(
129+
range(7, 7 + num_coarse_grained_physical_indices - len(bottom_left_i))
130+
)
131+
phys_contraction_i_conj_bottom_left = list(
132+
range(7, 7 + num_coarse_grained_physical_indices - len(bottom_left_i))
133+
)
134+
for pos, i in enumerate(bottom_left_i):
135+
phys_contraction_i_bottom_left.insert(i - 1, -(pos + 1))
136+
phys_contraction_i_conj_bottom_left.insert(
137+
i - 1, -len(bottom_left_i) - (pos + 1)
138+
)
139+
phys_contraction_i_bottom_left = tuple(phys_contraction_i_bottom_left)
140+
phys_contraction_i_conj_bottom_left = tuple(phys_contraction_i_conj_bottom_left)
141+
142+
contraction_bottom_left = {
143+
"tensors": [["tensor", "tensor_conj", "T3", "C4", "T4"]],
144+
"network": [
145+
[
146+
(3, 4)
147+
+ phys_contraction_i_bottom_left
148+
+ (
149+
-2 * len(bottom_left_i) - 3,
150+
-2 * len(bottom_left_i) - 5,
151+
), # tensor
152+
(5, 6)
153+
+ phys_contraction_i_conj_bottom_left
154+
+ (
155+
-2 * len(bottom_left_i) - 2,
156+
-2 * len(bottom_left_i) - 6,
157+
), # tensor_conj
158+
(2, -2 * len(bottom_left_i) - 1, 6, 4), # T3
159+
(2, 1), # C4
160+
(1, 5, 3, -2 * len(bottom_left_i) - 4), # T4
161+
]
162+
],
163+
}
164+
Definitions._process_def(
165+
contraction_bottom_left,
166+
f"partially_traced_four_site_density_matrices_bottom_left_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{bottom_left_i}",
167+
)
168+
setattr(
169+
Definitions,
170+
f"partially_traced_four_site_density_matrices_bottom_left_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{bottom_left_i}",
171+
contraction_bottom_left,
172+
)
173+
174+
phys_contraction_i_bottom_right = list(
175+
range(7, 7 + num_coarse_grained_physical_indices - len(bottom_right_i))
176+
)
177+
phys_contraction_i_conj_bottom_right = list(
178+
range(7, 7 + num_coarse_grained_physical_indices - len(bottom_right_i))
179+
)
180+
for pos, i in enumerate(bottom_right_i):
181+
phys_contraction_i_bottom_right.insert(i - 1, -(pos + 1))
182+
phys_contraction_i_conj_bottom_right.insert(
183+
i - 1, -len(bottom_right_i) - (pos + 1)
184+
)
185+
phys_contraction_i_bottom_right = tuple(phys_contraction_i_bottom_right)
186+
phys_contraction_i_conj_bottom_right = tuple(
187+
phys_contraction_i_conj_bottom_right
188+
)
189+
190+
contraction_bottom_right = {
191+
"tensors": [["tensor", "tensor_conj", "T2", "T3", "C3"]],
192+
"network": [
193+
[
194+
(-2 * len(bottom_right_i) - 6, 3)
195+
+ phys_contraction_i_bottom_right
196+
+ (4, -2 * len(bottom_right_i) - 3), # tensor
197+
(-2 * len(bottom_right_i) - 5, 5)
198+
+ phys_contraction_i_conj_bottom_right
199+
+ (6, -2 * len(bottom_right_i) - 2), # tensor_conj
200+
(4, 6, 2, -2 * len(bottom_right_i) - 1), # T2
201+
(-2 * len(bottom_right_i) - 4, 1, 5, 3), # T3
202+
(1, 2), # C3
203+
]
204+
],
205+
}
206+
Definitions._process_def(
207+
contraction_bottom_right,
208+
f"partially_traced_four_site_density_matrices_bottom_right_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{bottom_right_i}",
209+
)
210+
setattr(
211+
Definitions,
212+
f"partially_traced_four_site_density_matrices_bottom_right_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{bottom_right_i}",
213+
contraction_bottom_right,
214+
)
215+
216+
density_top_left = apply_contraction(
217+
f"partially_traced_four_site_density_matrices_top_left_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{top_left_i}",
218+
[t_top_left],
219+
[t_obj_top_left],
220+
[],
221+
disable_identity_check=True,
222+
)
223+
if len(top_left_i) > 0:
224+
density_top_left = density_top_left.reshape(
225+
real_physical_dimension ** len(top_left_i),
226+
real_physical_dimension ** len(top_left_i),
227+
density_top_left.shape[-6],
228+
density_top_left.shape[-5],
229+
density_top_left.shape[-4],
230+
density_top_left.shape[-3],
231+
density_top_left.shape[-2],
232+
density_top_left.shape[-1],
233+
)
234+
235+
density_top_right = apply_contraction(
236+
f"partially_traced_four_site_density_matrices_top_right_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{top_right_i}",
237+
[t_top_right],
238+
[t_obj_top_right],
239+
[],
240+
disable_identity_check=True,
241+
)
242+
if len(top_right_i) > 0:
243+
density_top_right = density_top_right.reshape(
244+
real_physical_dimension ** len(top_right_i),
245+
real_physical_dimension ** len(top_right_i),
246+
density_top_right.shape[-6],
247+
density_top_right.shape[-5],
248+
density_top_right.shape[-4],
249+
density_top_right.shape[-3],
250+
density_top_right.shape[-2],
251+
density_top_right.shape[-1],
252+
)
253+
254+
density_bottom_left = apply_contraction(
255+
f"partially_traced_four_site_density_matrices_bottom_left_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{bottom_left_i}",
256+
[t_bottom_left],
257+
[t_obj_bottom_left],
258+
[],
259+
disable_identity_check=True,
260+
)
261+
if len(bottom_left_i) > 0:
262+
density_bottom_left = density_bottom_left.reshape(
263+
real_physical_dimension ** len(bottom_left_i),
264+
real_physical_dimension ** len(bottom_left_i),
265+
density_bottom_left.shape[-6],
266+
density_bottom_left.shape[-5],
267+
density_bottom_left.shape[-4],
268+
density_bottom_left.shape[-3],
269+
density_bottom_left.shape[-2],
270+
density_bottom_left.shape[-1],
271+
)
272+
273+
density_bottom_right = apply_contraction(
274+
f"partially_traced_four_site_density_matrices_bottom_right_{real_physical_dimension}_{num_coarse_grained_physical_indices}_{bottom_right_i}",
275+
[t_bottom_right],
276+
[t_obj_bottom_right],
277+
[],
278+
disable_identity_check=True,
279+
)
280+
if len(bottom_right_i) > 0:
281+
density_bottom_right = density_bottom_right.reshape(
282+
real_physical_dimension ** len(bottom_right_i),
283+
real_physical_dimension ** len(bottom_right_i),
284+
density_bottom_right.shape[-6],
285+
density_bottom_right.shape[-5],
286+
density_bottom_right.shape[-4],
287+
density_bottom_right.shape[-3],
288+
density_bottom_right.shape[-2],
289+
density_bottom_right.shape[-1],
290+
)
291+
292+
return (
293+
density_top_left,
294+
density_top_right,
295+
density_bottom_left,
296+
density_bottom_right,
297+
)

0 commit comments

Comments
 (0)