Skip to content

Commit 4227837

Browse files
JanLucaJanLuca
committedMay 7, 2024
Implement structure factor calculation
·
1.0.10.6.0
1 parent b8aff48 commit 4227837

File tree

8 files changed

+3018
-33
lines changed

8 files changed

+3018
-33
lines changed
 

‎varipeps/contractions/definitions.py

+784-32
Large diffs are not rendered by default.

‎varipeps/ctmrg/__init__.py

+4
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
from . import absorption
22
from . import projectors
33
from . import routine
4+
from . import structure_factor_absorption
5+
from . import structure_factor_routine
46

57
from .routine import (
68
calc_ctmrg_env,
79
calc_ctmrg_env_custom_rule,
810
CTMRGNotConvergedError,
911
CTMRGGradientNotConvergedError,
1012
)
13+
14+
from .structure_factor_routine import calc_ctmrg_env_structure_factor

‎varipeps/ctmrg/structure_factor_absorption.py

+909
Large diffs are not rendered by default.

‎varipeps/ctmrg/structure_factor_routine.py

+407
Large diffs are not rendered by default.

‎varipeps/expectation/__init__.py

+1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from . import spiral_helpers
2+
from . import structure_factor
23
from . import model
34
from . import one_site
45
from . import two_sites

‎varipeps/expectation/structure_factor.py

+177
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
from functools import partial
2+
3+
import jax.numpy as jnp
4+
from jax import jit
5+
6+
from varipeps.peps import PEPS_Tensor, PEPS_Unit_Cell
7+
from varipeps.contractions import apply_contraction_jitted
8+
9+
from typing import Sequence
10+
11+
12+
@partial(jit, static_argnums=(4, 5))
13+
def calc_structure_factor_expectation(
14+
peps_tensor_obj: PEPS_Tensor,
15+
alpha_gate: jnp.ndarray,
16+
beta_gate: jnp.array,
17+
structure_factor_inner_factors: Sequence[float],
18+
real_d: int,
19+
num_sites: int,
20+
):
21+
if num_sites > 1:
22+
Id = jnp.eye(real_d ** (num_sites - 1))
23+
24+
full_alpha_gate = [jnp.kron(alpha_gate, Id)]
25+
26+
alpha_gate_tmp = full_alpha_gate[0].reshape((real_d,) * 2 * num_sites)
27+
for i in range(1, num_sites):
28+
trans_order = list(range(1, num_sites)) + list(
29+
range(num_sites + 1, 2 * num_sites)
30+
)
31+
trans_order.insert(i, 0)
32+
trans_order.insert(num_sites + i, num_sites)
33+
34+
tmp_gate = alpha_gate_tmp.transpose(trans_order).reshape(
35+
real_d**num_sites, real_d**num_sites
36+
)
37+
38+
full_alpha_gate.append(tmp_gate * structure_factor_inner_factors[i].conj())
39+
40+
alpha_beta_gate = 0
41+
# alpha_beta_gate_tmp_1 = alpha_beta_gate.reshape((real_d,) * 2 * num_sites)
42+
alpha_beta_gate_tmp_2 = jnp.kron(
43+
jnp.kron(alpha_gate, beta_gate), jnp.eye(real_d ** (num_sites - 2))
44+
).reshape((real_d,) * 2 * num_sites)
45+
46+
for i in range(num_sites):
47+
for j in range(num_sites):
48+
if i == j:
49+
continue
50+
51+
# if i == j:
52+
# trans_order = list(range(1, num_sites)) + list(
53+
# range(num_sites + 1, 2 * num_sites)
54+
# )
55+
# trans_order.insert(i, 0)
56+
# trans_order.insert(num_sites + i, num_sites)
57+
#
58+
# tmp_gate = alpha_beta_gate_tmp_1.transpose(trans_order)
59+
# tmp_gate = tmp_gate.reshape(real_d**num_sites, real_d**num_sites)
60+
#
61+
# alpha_beta_gate += tmp_gate
62+
# else:
63+
trans_order = list(range(2, num_sites)) + list(
64+
range(num_sites + 2, 2 * num_sites)
65+
)
66+
if i <= j:
67+
trans_order.insert(i, 0)
68+
trans_order.insert(j, 1)
69+
trans_order.insert(num_sites + i, num_sites)
70+
trans_order.insert(num_sites + j, num_sites + 1)
71+
else:
72+
trans_order.insert(j, 1)
73+
trans_order.insert(i, 0)
74+
trans_order.insert(num_sites + j, num_sites + 1)
75+
trans_order.insert(num_sites + i, num_sites)
76+
77+
tmp_gate = alpha_beta_gate_tmp_2.transpose(trans_order)
78+
tmp_gate = tmp_gate.reshape(real_d**num_sites, real_d**num_sites)
79+
80+
alpha_beta_gate += (
81+
tmp_gate
82+
* structure_factor_inner_factors[i].conj()
83+
* structure_factor_inner_factors[j]
84+
)
85+
else:
86+
full_alpha_gate = [alpha_gate]
87+
alpha_beta_gate = alpha_gate @ beta_gate
88+
89+
density_matrix = apply_contraction_jitted(
90+
"density_matrix_one_site", [peps_tensor_obj.tensor], [peps_tensor_obj], []
91+
)
92+
93+
norm = jnp.trace(density_matrix)
94+
95+
result = jnp.tensordot(density_matrix, alpha_beta_gate, ((0, 1), (0, 1))) / norm
96+
97+
density_matrix = apply_contraction_jitted(
98+
"density_matrix_one_site_C1_phase",
99+
[peps_tensor_obj.tensor],
100+
[peps_tensor_obj],
101+
[],
102+
)
103+
104+
for g in full_alpha_gate:
105+
result += jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm
106+
107+
density_matrix = apply_contraction_jitted(
108+
"density_matrix_one_site_C2_phase",
109+
[peps_tensor_obj.tensor],
110+
[peps_tensor_obj],
111+
[],
112+
)
113+
114+
for g in full_alpha_gate:
115+
result += jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm
116+
117+
density_matrix = apply_contraction_jitted(
118+
"density_matrix_one_site_C3_phase",
119+
[peps_tensor_obj.tensor],
120+
[peps_tensor_obj],
121+
[],
122+
)
123+
124+
for g in full_alpha_gate:
125+
result += jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm
126+
127+
density_matrix = apply_contraction_jitted(
128+
"density_matrix_one_site_C4_phase",
129+
[peps_tensor_obj.tensor],
130+
[peps_tensor_obj],
131+
[],
132+
)
133+
134+
for g in full_alpha_gate:
135+
result += jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm
136+
137+
density_matrix = apply_contraction_jitted(
138+
"density_matrix_one_site_T1_phase",
139+
[peps_tensor_obj.tensor],
140+
[peps_tensor_obj],
141+
[],
142+
)
143+
144+
for g in full_alpha_gate:
145+
result += jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm
146+
147+
density_matrix = apply_contraction_jitted(
148+
"density_matrix_one_site_T2_phase",
149+
[peps_tensor_obj.tensor],
150+
[peps_tensor_obj],
151+
[],
152+
)
153+
154+
for g in full_alpha_gate:
155+
result += jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm
156+
157+
density_matrix = apply_contraction_jitted(
158+
"density_matrix_one_site_T3_phase",
159+
[peps_tensor_obj.tensor],
160+
[peps_tensor_obj],
161+
[],
162+
)
163+
164+
for g in full_alpha_gate:
165+
result += jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm
166+
167+
density_matrix = apply_contraction_jitted(
168+
"density_matrix_one_site_T4_phase",
169+
[peps_tensor_obj.tensor],
170+
[peps_tensor_obj],
171+
[],
172+
)
173+
174+
for g in full_alpha_gate:
175+
result += jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm
176+
177+
return result

‎varipeps/peps/__init__.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
from . import tensor
22
from . import unitcell
3-
from .tensor import PEPS_Tensor
3+
from .tensor import PEPS_Tensor, PEPS_Tensor_Structure_Factor
44
from .unitcell import PEPS_Unit_Cell

‎varipeps/peps/tensor.py

+735
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)
Please sign in to comment.