Skip to content

Commit dba9c26

Browse files
JanLucaJanLuca
committed
Rework the code to calculate correlation lengths
1 parent f5e67d4 commit dba9c26

File tree

1 file changed

+61
-143
lines changed

1 file changed

+61
-143
lines changed

‎varipeps/corrlength/corrlength.py

+61-143
Original file line numberDiff line numberDiff line change
@@ -9,162 +9,80 @@
99

1010

1111
def calculate_correlation_length(unitcell: PEPS_Unit_Cell):
12-
use_full = True
13-
14-
full_transfer_right = None
15-
full_transfer_bottom = None
16-
17-
try:
12+
initial_vector_left = apply_contraction_jitted(
13+
"corrlength_vector_left",
14+
(unitcell[0, 0][0][0].tensor,),
15+
(unitcell[0, 0][0][0],),
16+
(),
17+
)
18+
initial_vector_left = initial_vector_left.reshape(-1)
19+
20+
initial_vector_top = apply_contraction_jitted(
21+
"corrlength_vector_top",
22+
(unitcell[0, 0][0][0].tensor,),
23+
(unitcell[0, 0][0][0],),
24+
(),
25+
)
26+
initial_vector_top = initial_vector_top.reshape(-1)
27+
28+
def left_matvec(vec):
29+
vec = jnp.asarray(vec)
1830
for _, view in unitcell.iter_one_row(0):
19-
transfer_right = apply_contraction_jitted(
20-
"corrlength_transfer_right",
31+
if vec.ndim != 4:
32+
vec = vec.reshape(
33+
view[0, 0][0][0].T1.shape[0],
34+
view[0, 0][0][0].tensor.shape[0],
35+
view[0, 0][0][0].tensor.shape[0],
36+
view[0, 0][0][0].T3.shape[0],
37+
)
38+
vec = apply_contraction_jitted(
39+
"corrlength_absorb_one_column",
2140
(view[0, 0][0][0].tensor,),
2241
(view[0, 0][0][0],),
23-
(),
24-
)
25-
transfer_right = transfer_right.reshape(
26-
np.prod(transfer_right.shape[:4]), np.prod(transfer_right.shape[4:])
42+
(vec,),
2743
)
44+
return vec.reshape(-1)
45+
46+
left_lin_op = LinearOperator(
47+
(initial_vector_left.shape[0], initial_vector_left.shape[0]),
48+
matvec=left_matvec,
49+
)
50+
51+
eig_left, eigvec_left = eigs(left_lin_op, k=5, v0=initial_vector_left, which="LM")
2852

29-
if full_transfer_right is None:
30-
full_transfer_right = transfer_right
31-
else:
32-
full_transfer_right = full_transfer_right @ transfer_right
53+
eig_left = eig_left[np.argsort(np.abs(eig_left))[::-1]]
54+
eig_left /= np.abs(eig_left[0])
3355

34-
initial_vector_right = apply_contraction_jitted(
35-
"corrlength_vector_right",
36-
(unitcell[0, -1][0][0].tensor,),
37-
(view[0, -1][0][0],),
38-
(),
39-
)
40-
initial_vector_right = initial_vector_right.reshape(-1)
56+
corr_len_left = -1 / np.log(np.abs(eig_left[1]))
4157

58+
def top_matvec(vec):
59+
vec = jnp.asarray(vec)
4260
for _, view in unitcell.iter_one_column(0):
43-
transfer_bottom = apply_contraction_jitted(
44-
"corrlength_transfer_bottom",
61+
if vec.ndim != 4:
62+
vec = vec.reshape(
63+
view[0, 0][0][0].T4.shape[3],
64+
view[0, 0][0][0].tensor.shape[4],
65+
view[0, 0][0][0].tensor.shape[4],
66+
view[0, 0][0][0].T2.shape[3],
67+
)
68+
vec = apply_contraction_jitted(
69+
"corrlength_absorb_one_row",
4570
(view[0, 0][0][0].tensor,),
4671
(view[0, 0][0][0],),
47-
(),
48-
)
49-
transfer_bottom = transfer_bottom.reshape(
50-
np.prod(transfer_bottom.shape[:4]), np.prod(transfer_bottom.shape[4:])
72+
(vec,),
5173
)
74+
return vec.reshape(-1)
5275

53-
if full_transfer_bottom is None:
54-
full_transfer_bottom = transfer_bottom
55-
else:
56-
full_transfer_bottom = full_transfer_bottom @ transfer_bottom
57-
58-
initial_vector_bottom = apply_contraction_jitted(
59-
"corrlength_vector_bottom",
60-
(unitcell[-1, 0][0][0].tensor,),
61-
(view[-1, 0][0][0],),
62-
(),
63-
)
64-
initial_vector_bottom = initial_vector_bottom.reshape(-1)
65-
66-
eig_right, eigvec_right = eigs(
67-
np.asarray(full_transfer_right),
68-
k=5,
69-
v0=np.asarray(initial_vector_right),
70-
which="LM",
71-
)
72-
73-
eig_right = np.abs(eig_right)
74-
eig_right /= eig_right[0]
75-
76-
corr_len_right = -1 / np.log(eig_right[1])
77-
78-
eig_bottom, eigvec_bottom = eigs(
79-
np.asarray(full_transfer_bottom),
80-
k=5,
81-
v0=np.asarray(initial_vector_bottom),
82-
which="LM",
83-
)
84-
85-
eig_bottom = np.abs(eig_bottom)
86-
eig_bottom /= eig_bottom[0]
87-
88-
corr_len_bottom = -1 / np.log(eig_bottom[1])
89-
90-
return (corr_len_right, eig_right), (corr_len_bottom, eig_bottom)
91-
except XlaRuntimeError:
92-
initial_vector_left = apply_contraction_jitted(
93-
"corrlength_vector_left",
94-
(unitcell[0, 0][0][0].tensor,),
95-
(view[0, 0][0][0],),
96-
(),
97-
)
98-
initial_vector_left = initial_vector_left.reshape(-1)
99-
100-
initial_vector_top = apply_contraction_jitted(
101-
"corrlength_vector_top",
102-
(unitcell[0, 0][0][0].tensor,),
103-
(view[0, 0][0][0],),
104-
(),
105-
)
106-
initial_vector_top = initial_vector_top.reshape(-1)
107-
108-
def left_matvec(vec):
109-
vec = jnp.asarray(vec)
110-
for _, view in unitcell.iter_one_row(0):
111-
if vec.ndim != 4:
112-
vec = vec.reshape(
113-
view[0, 0][0][0].T1.shape[0],
114-
view[0, 0][0][0].tensor.shape[0],
115-
view[0, 0][0][0].tensor.shape[0],
116-
view[0, 0][0][0].T3.shape[0],
117-
)
118-
vec = apply_contraction_jitted(
119-
"corrlength_absorb_one_column",
120-
(view[0, 0][0][0].tensor,),
121-
(view[0, 0][0][0],),
122-
(vec,),
123-
)
124-
return vec.reshape(-1)
125-
126-
left_lin_op = LinearOperator(
127-
(initial_vector_left.shape[0], initial_vector_left.shape[0]),
128-
matvec=left_matvec,
129-
)
130-
131-
eig_left, eigvec_left = eigs(
132-
left_lin_op, k=5, v0=initial_vector_left, which="LM"
133-
)
134-
135-
eig_left = np.abs(eig_left)
136-
eig_left /= eig_left[0]
137-
138-
corr_len_left = -1 / np.log(eig_left[1])
139-
140-
def top_matvec(vec):
141-
vec = jnp.asarray(vec)
142-
for _, view in unitcell.iter_one_row(0):
143-
if vec.ndim != 4:
144-
vec = vec.reshape(
145-
view[0, 0][0][0].T4.shape[3],
146-
view[0, 0][0][0].tensor.shape[4],
147-
view[0, 0][0][0].tensor.shape[4],
148-
view[0, 0][0][0].T2.shape[3],
149-
)
150-
vec = apply_contraction_jitted(
151-
"corrlength_absorb_one_row",
152-
(view[0, 0][0][0].tensor,),
153-
(view[0, 0][0][0],),
154-
(vec,),
155-
)
156-
return vec.reshape(-1)
157-
158-
top_lin_op = LinearOperator(
159-
(initial_vector_top.shape[0], initial_vector_top.shape[0]),
160-
matvec=top_matvec,
161-
)
76+
top_lin_op = LinearOperator(
77+
(initial_vector_top.shape[0], initial_vector_top.shape[0]),
78+
matvec=top_matvec,
79+
)
16280

163-
eig_top, eigvec_top = eigs(top_lin_op, k=5, v0=initial_vector_top, which="LM")
81+
eig_top, eigvec_top = eigs(top_lin_op, k=5, v0=initial_vector_top, which="LM")
16482

165-
eig_top = np.abs(eig_top)
166-
eig_top /= eig_top[0]
83+
eig_top = eig_top[np.argsort(np.abs(eig_top))[::-1]]
84+
eig_top /= np.abs(eig_top[0])
16785

168-
corr_len_top = -1 / np.log(eig_top[1])
86+
corr_len_top = -1 / np.log(np.abs(eig_top[1]))
16987

170-
return (corr_len_left, eig_left), (corr_len_top, eig_top)
88+
return (corr_len_left, eig_left), (corr_len_top, eig_top)

0 commit comments

Comments
 (0)