Trajectree 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (122) hide show
  1. trajectree/__init__.py +0 -3
  2. trajectree/fock_optics/devices.py +1 -1
  3. trajectree/fock_optics/light_sources.py +2 -2
  4. trajectree/fock_optics/measurement.py +3 -3
  5. trajectree/fock_optics/utils.py +6 -6
  6. trajectree/trajectory.py +2 -2
  7. {trajectree-0.0.1.dist-info → trajectree-0.0.2.dist-info}/METADATA +2 -3
  8. trajectree-0.0.2.dist-info/RECORD +16 -0
  9. trajectree/quimb/docs/_pygments/_pygments_dark.py +0 -118
  10. trajectree/quimb/docs/_pygments/_pygments_light.py +0 -118
  11. trajectree/quimb/docs/conf.py +0 -158
  12. trajectree/quimb/docs/examples/ex_mpi_expm_evo.py +0 -62
  13. trajectree/quimb/quimb/__init__.py +0 -507
  14. trajectree/quimb/quimb/calc.py +0 -1491
  15. trajectree/quimb/quimb/core.py +0 -2279
  16. trajectree/quimb/quimb/evo.py +0 -712
  17. trajectree/quimb/quimb/experimental/__init__.py +0 -0
  18. trajectree/quimb/quimb/experimental/autojittn.py +0 -129
  19. trajectree/quimb/quimb/experimental/belief_propagation/__init__.py +0 -109
  20. trajectree/quimb/quimb/experimental/belief_propagation/bp_common.py +0 -397
  21. trajectree/quimb/quimb/experimental/belief_propagation/d1bp.py +0 -316
  22. trajectree/quimb/quimb/experimental/belief_propagation/d2bp.py +0 -653
  23. trajectree/quimb/quimb/experimental/belief_propagation/hd1bp.py +0 -571
  24. trajectree/quimb/quimb/experimental/belief_propagation/hv1bp.py +0 -775
  25. trajectree/quimb/quimb/experimental/belief_propagation/l1bp.py +0 -316
  26. trajectree/quimb/quimb/experimental/belief_propagation/l2bp.py +0 -537
  27. trajectree/quimb/quimb/experimental/belief_propagation/regions.py +0 -194
  28. trajectree/quimb/quimb/experimental/cluster_update.py +0 -286
  29. trajectree/quimb/quimb/experimental/merabuilder.py +0 -865
  30. trajectree/quimb/quimb/experimental/operatorbuilder/__init__.py +0 -15
  31. trajectree/quimb/quimb/experimental/operatorbuilder/operatorbuilder.py +0 -1631
  32. trajectree/quimb/quimb/experimental/schematic.py +0 -7
  33. trajectree/quimb/quimb/experimental/tn_marginals.py +0 -130
  34. trajectree/quimb/quimb/experimental/tnvmc.py +0 -1483
  35. trajectree/quimb/quimb/gates.py +0 -36
  36. trajectree/quimb/quimb/gen/__init__.py +0 -2
  37. trajectree/quimb/quimb/gen/operators.py +0 -1167
  38. trajectree/quimb/quimb/gen/rand.py +0 -713
  39. trajectree/quimb/quimb/gen/states.py +0 -479
  40. trajectree/quimb/quimb/linalg/__init__.py +0 -6
  41. trajectree/quimb/quimb/linalg/approx_spectral.py +0 -1109
  42. trajectree/quimb/quimb/linalg/autoblock.py +0 -258
  43. trajectree/quimb/quimb/linalg/base_linalg.py +0 -719
  44. trajectree/quimb/quimb/linalg/mpi_launcher.py +0 -397
  45. trajectree/quimb/quimb/linalg/numpy_linalg.py +0 -244
  46. trajectree/quimb/quimb/linalg/rand_linalg.py +0 -514
  47. trajectree/quimb/quimb/linalg/scipy_linalg.py +0 -293
  48. trajectree/quimb/quimb/linalg/slepc_linalg.py +0 -892
  49. trajectree/quimb/quimb/schematic.py +0 -1518
  50. trajectree/quimb/quimb/tensor/__init__.py +0 -401
  51. trajectree/quimb/quimb/tensor/array_ops.py +0 -610
  52. trajectree/quimb/quimb/tensor/circuit.py +0 -4824
  53. trajectree/quimb/quimb/tensor/circuit_gen.py +0 -411
  54. trajectree/quimb/quimb/tensor/contraction.py +0 -336
  55. trajectree/quimb/quimb/tensor/decomp.py +0 -1255
  56. trajectree/quimb/quimb/tensor/drawing.py +0 -1646
  57. trajectree/quimb/quimb/tensor/fitting.py +0 -385
  58. trajectree/quimb/quimb/tensor/geometry.py +0 -583
  59. trajectree/quimb/quimb/tensor/interface.py +0 -114
  60. trajectree/quimb/quimb/tensor/networking.py +0 -1058
  61. trajectree/quimb/quimb/tensor/optimize.py +0 -1818
  62. trajectree/quimb/quimb/tensor/tensor_1d.py +0 -4778
  63. trajectree/quimb/quimb/tensor/tensor_1d_compress.py +0 -1854
  64. trajectree/quimb/quimb/tensor/tensor_1d_tebd.py +0 -662
  65. trajectree/quimb/quimb/tensor/tensor_2d.py +0 -5954
  66. trajectree/quimb/quimb/tensor/tensor_2d_compress.py +0 -96
  67. trajectree/quimb/quimb/tensor/tensor_2d_tebd.py +0 -1230
  68. trajectree/quimb/quimb/tensor/tensor_3d.py +0 -2869
  69. trajectree/quimb/quimb/tensor/tensor_3d_tebd.py +0 -46
  70. trajectree/quimb/quimb/tensor/tensor_approx_spectral.py +0 -60
  71. trajectree/quimb/quimb/tensor/tensor_arbgeom.py +0 -3237
  72. trajectree/quimb/quimb/tensor/tensor_arbgeom_compress.py +0 -565
  73. trajectree/quimb/quimb/tensor/tensor_arbgeom_tebd.py +0 -1138
  74. trajectree/quimb/quimb/tensor/tensor_builder.py +0 -5411
  75. trajectree/quimb/quimb/tensor/tensor_core.py +0 -11179
  76. trajectree/quimb/quimb/tensor/tensor_dmrg.py +0 -1472
  77. trajectree/quimb/quimb/tensor/tensor_mera.py +0 -204
  78. trajectree/quimb/quimb/utils.py +0 -892
  79. trajectree/quimb/tests/__init__.py +0 -0
  80. trajectree/quimb/tests/test_accel.py +0 -501
  81. trajectree/quimb/tests/test_calc.py +0 -788
  82. trajectree/quimb/tests/test_core.py +0 -847
  83. trajectree/quimb/tests/test_evo.py +0 -565
  84. trajectree/quimb/tests/test_gen/__init__.py +0 -0
  85. trajectree/quimb/tests/test_gen/test_operators.py +0 -361
  86. trajectree/quimb/tests/test_gen/test_rand.py +0 -296
  87. trajectree/quimb/tests/test_gen/test_states.py +0 -261
  88. trajectree/quimb/tests/test_linalg/__init__.py +0 -0
  89. trajectree/quimb/tests/test_linalg/test_approx_spectral.py +0 -368
  90. trajectree/quimb/tests/test_linalg/test_base_linalg.py +0 -351
  91. trajectree/quimb/tests/test_linalg/test_mpi_linalg.py +0 -127
  92. trajectree/quimb/tests/test_linalg/test_numpy_linalg.py +0 -84
  93. trajectree/quimb/tests/test_linalg/test_rand_linalg.py +0 -134
  94. trajectree/quimb/tests/test_linalg/test_slepc_linalg.py +0 -283
  95. trajectree/quimb/tests/test_tensor/__init__.py +0 -0
  96. trajectree/quimb/tests/test_tensor/test_belief_propagation/__init__.py +0 -0
  97. trajectree/quimb/tests/test_tensor/test_belief_propagation/test_d1bp.py +0 -39
  98. trajectree/quimb/tests/test_tensor/test_belief_propagation/test_d2bp.py +0 -67
  99. trajectree/quimb/tests/test_tensor/test_belief_propagation/test_hd1bp.py +0 -64
  100. trajectree/quimb/tests/test_tensor/test_belief_propagation/test_hv1bp.py +0 -51
  101. trajectree/quimb/tests/test_tensor/test_belief_propagation/test_l1bp.py +0 -142
  102. trajectree/quimb/tests/test_tensor/test_belief_propagation/test_l2bp.py +0 -101
  103. trajectree/quimb/tests/test_tensor/test_circuit.py +0 -816
  104. trajectree/quimb/tests/test_tensor/test_contract.py +0 -67
  105. trajectree/quimb/tests/test_tensor/test_decomp.py +0 -40
  106. trajectree/quimb/tests/test_tensor/test_mera.py +0 -52
  107. trajectree/quimb/tests/test_tensor/test_optimizers.py +0 -488
  108. trajectree/quimb/tests/test_tensor/test_tensor_1d.py +0 -1171
  109. trajectree/quimb/tests/test_tensor/test_tensor_2d.py +0 -606
  110. trajectree/quimb/tests/test_tensor/test_tensor_2d_tebd.py +0 -144
  111. trajectree/quimb/tests/test_tensor/test_tensor_3d.py +0 -123
  112. trajectree/quimb/tests/test_tensor/test_tensor_arbgeom.py +0 -226
  113. trajectree/quimb/tests/test_tensor/test_tensor_builder.py +0 -441
  114. trajectree/quimb/tests/test_tensor/test_tensor_core.py +0 -2066
  115. trajectree/quimb/tests/test_tensor/test_tensor_dmrg.py +0 -388
  116. trajectree/quimb/tests/test_tensor/test_tensor_spectral_approx.py +0 -63
  117. trajectree/quimb/tests/test_tensor/test_tensor_tebd.py +0 -270
  118. trajectree/quimb/tests/test_utils.py +0 -85
  119. trajectree-0.0.1.dist-info/RECORD +0 -126
  120. {trajectree-0.0.1.dist-info → trajectree-0.0.2.dist-info}/WHEEL +0 -0
  121. {trajectree-0.0.1.dist-info → trajectree-0.0.2.dist-info}/licenses/LICENSE +0 -0
  122. {trajectree-0.0.1.dist-info → trajectree-0.0.2.dist-info}/top_level.txt +0 -0
@@ -1,2869 +0,0 @@
1
- """Classes and algorithms related to 3D tensor networks."""
2
-
3
- import functools
4
- import itertools
5
- from operator import add
6
- from numbers import Integral
7
- from collections import defaultdict
8
- from itertools import product, combinations
9
-
10
- from autoray import do, dag
11
-
12
- from ..utils import check_opt, ensure_dict, pairwise
13
- from ..utils import progbar as Progbar
14
- from ..gen.rand import randn, seed_rand
15
- from . import array_ops as ops
16
- from .tensor_core import (
17
- bonds,
18
- bonds_size,
19
- oset,
20
- rand_uuid,
21
- tags_to_oset,
22
- Tensor,
23
- )
24
- from .tensor_arbgeom import (
25
- TensorNetworkGen,
26
- TensorNetworkGenVector,
27
- )
28
-
29
-
30
- def gen_3d_bonds(Lx, Ly, Lz, steppers=None, coo_filter=None, cyclic=False):
31
- """Convenience function for tiling pairs of bond coordinates on a 3D
32
- lattice given a function like ``lambda i, j, k: (i + 1, j + 1, k + 1)``.
33
-
34
- Parameters
35
- ----------
36
- Lx : int
37
- The number of x-slices.
38
- Ly : int
39
- The number of y-slices.
40
- Lz : int
41
- The number of z-slices.
42
- steppers : callable or sequence of callable
43
- Function(s) that take args ``(i, j, k)`` and generate another
44
- coordinate, thus defining a bond. Only valid steps are taken. If not
45
- given, defaults to nearest neighbor bonds.
46
- coo_filter : callable
47
- Function that takes args ``(i, j, k)`` and only returns ``True`` if
48
- this is to be a valid starting coordinate.
49
-
50
- Yields
51
- ------
52
- bond : tuple[tuple[int, int, int], tuple[int, int, int]]
53
- A pair of coordinates.
54
-
55
- Examples
56
- --------
57
-
58
- Generate nearest neighbor bonds:
59
-
60
- >>> for bond in gen_3d_bonds(2, 2, 2, [lambda i, j, k: (i + 1, j, k),
61
- ... lambda i, j, k: (i, j + 1, k),
62
- ... lambda i, j, k: (i, j, k + 1)]):
63
- ... print(bond)
64
- ((0, 0, 0), (1, 0, 0))
65
- ((0, 0, 0), (0, 1, 0))
66
- ((0, 0, 0), (0, 0, 1))
67
- ((0, 0, 1), (1, 0, 1))
68
- ((0, 0, 1), (0, 1, 1))
69
- ((0, 1, 0), (1, 1, 0))
70
- ((0, 1, 0), (0, 1, 1))
71
- ((0, 1, 1), (1, 1, 1))
72
- ((1, 0, 0), (1, 1, 0))
73
- ((1, 0, 0), (1, 0, 1))
74
- ((1, 0, 1), (1, 1, 1))
75
- ((1, 1, 0), (1, 1, 1))
76
-
77
- """
78
- if steppers is None:
79
- steppers = [
80
- lambda i, j, k: (i, j, k + 1),
81
- lambda i, j, k: (i, j + 1, k),
82
- lambda i, j, k: (i + 1, j, k),
83
- ]
84
-
85
- if callable(steppers):
86
- steppers = (steppers,)
87
-
88
- try:
89
- cyclic_x, cyclic_y, cyclic_z = cyclic
90
- except (TypeError, ValueError):
91
- cyclic_x = cyclic_y = cyclic_z = cyclic
92
-
93
- def _maybe_wrap_coo(w, Lw, cyclic):
94
- if 0 <= w < Lw:
95
- return w
96
- if cyclic:
97
- return w % Lw
98
- return None
99
-
100
- for i, j, k in product(range(Lx), range(Ly), range(Lz)):
101
- if (coo_filter is None) or coo_filter(i, j, k):
102
- for stepper in steppers:
103
- i2, j2, k2 = stepper(i, j, k)
104
-
105
- i2 = _maybe_wrap_coo(i2, Lx, cyclic_x)
106
- j2 = _maybe_wrap_coo(j2, Ly, cyclic_y)
107
- k2 = _maybe_wrap_coo(k2, Lz, cyclic_z)
108
-
109
- if all(x is not None for x in (i2, j2, k2)):
110
- yield (i, j, k), (i2, j2, k2)
111
-
112
-
113
- def gen_3d_plaquette(coo0, steps):
114
- """Generate a plaquette at site ``coo0`` by stepping first in ``steps`` and
115
- then the reverse steps.
116
-
117
- Parameters
118
- ----------
119
- coo0 : tuple
120
- The coordinate of the first site in the plaquette.
121
- steps : tuple
122
- The steps to take to generate the plaquette. Each element should be
123
- one of ``('x+', 'x-', 'y+', 'y-', 'z+', 'z-')``.
124
-
125
- Yields
126
- ------
127
- coo : tuple
128
- The coordinates of the sites in the plaquette, including the last
129
- site which will be the same as the first.
130
- """
131
- x, y, z = coo0
132
- smap = {"+": +1, "-": -1}
133
- step_backs = []
134
- yield x, y, z
135
- for step in steps:
136
- d, s = step
137
- x, y, z = {
138
- "x": (x + smap[s], y, z),
139
- "y": (x, y + smap[s], z),
140
- "z": (x, y, z + smap[s]),
141
- }[d]
142
- yield x, y, z
143
- step_backs.append(d + "-" if s == "+" else "-")
144
- for step in step_backs:
145
- d, s = step
146
- x, y, z = {
147
- "x": (x + smap[s], y, z),
148
- "y": (x, y + smap[s], z),
149
- "z": (x, y, z + smap[s]),
150
- }[d]
151
- yield x, y, z
152
-
153
-
154
- def gen_3d_plaquettes(Lx, Ly, Lz, tiling="1"):
155
- """Generate a tiling of plaquettes in a cubic 3D lattice.
156
-
157
- Parameters
158
- ----------
159
- Lx : int
160
- The length of the lattice in the x direction.
161
- Ly : int
162
- The length of the lattice in the y direction.
163
- Lz : int
164
- The length of the lattice in the z direction.
165
- tiling : {'1', '2', '4', 'full'}
166
- The tiling to use:
167
-
168
- - '1': plaquettes in a sparse checkerboard pattern, such that each edge
169
- is covered by a maximum of one plaquette.
170
- - '2': less sparse checkerboard pattern, such that each edge is
171
- covered by a maximum of two plaquettes.
172
- - '4' or 'full': dense tiling of plaquettes. All bulk edges will
173
- be covered four times.
174
-
175
- Yields
176
- ------
177
- plaquette : tuple[tuple[int]]
178
- The coordinates of the sites in each plaquette, including the last
179
- site which will be the same as the first.
180
- """
181
- if isinstance(tiling, int):
182
- tiling = str(tiling)
183
-
184
- if tiling == "1":
185
- for x, y, z in itertools.product(range(Lx), range(Ly), range(Lz)):
186
- if (x % 2 == 0) and (y % 2 == 0) and (x < Lx - 1 and y < Ly - 1):
187
- yield tuple(gen_3d_plaquette((x, y, z), ("x+", "y+")))
188
- if (y % 2 == 1) and (z % 2 == 0) and (y < Ly - 1 and z < Lz - 1):
189
- yield tuple(gen_3d_plaquette((x, y, z), ("y+", "z+")))
190
- if (z % 2 == 1) and (x % 2 == 1) and (z < Lz - 1 and x < Lx - 1):
191
- yield tuple(gen_3d_plaquette((x, y, z), ("z+", "x+")))
192
- elif tiling == "2":
193
- for x, y, z in itertools.product(range(Lx), range(Ly), range(Lz)):
194
- if ((x + y) % 2 == 0) and (x < Lx - 1 and y < Ly - 1):
195
- yield tuple(gen_3d_plaquette((x, y, z), ("x+", "y+")))
196
- if ((y + z) % 2 == 0) and (y < Ly - 1 and z < Lz - 1):
197
- yield tuple(gen_3d_plaquette((x, y, z), ("y+", "z+")))
198
- if ((x + z) % 2 == 1) and (z < Lz - 1 and x < Lx - 1):
199
- yield tuple(gen_3d_plaquette((x, y, z), ("z+", "x+")))
200
- elif tiling in ("4", "full"):
201
- for x, y, z in itertools.product(range(Lx), range(Ly), range(Lz)):
202
- if x < Lx - 1 and y < Ly - 1:
203
- yield tuple(gen_3d_plaquette((x, y, z), ("x+", "y+")))
204
- if y < Ly - 1 and z < Lz - 1:
205
- yield tuple(gen_3d_plaquette((x, y, z), ("y+", "z+")))
206
- if z < Lz - 1 and x < Lx - 1:
207
- yield tuple(gen_3d_plaquette((x, y, z), ("z+", "x+")))
208
- else:
209
- raise ValueError(
210
- "Invalid tiling: {}. Must be one of '1', '2', '4', 'full'."
211
- )
212
-
213
-
214
- def gen_3d_strings(Lx, Ly, Lz):
215
- """Generate all length-wise strings in a cubic 3D lattice."""
216
- for x, y in itertools.product(range(Lx), range(Ly)):
217
- yield tuple((x, y, z) for z in range(Lz))
218
- for y, z in itertools.product(range(Ly), range(Lz)):
219
- yield tuple((x, y, z) for x in range(Lx))
220
- for x, z in itertools.product(range(Lx), range(Lz)):
221
- yield tuple((x, y, z) for y in range(Ly))
222
-
223
-
224
- class Rotator3D:
225
- """Object for rotating coordinates and various contraction functions so
226
- that the core algorithms only have to written once, but nor does the actual
227
- TN have to be modified.
228
- """
229
-
230
- def __init__(self, tn, xrange, yrange, zrange, from_which):
231
- check_opt(
232
- "from_which",
233
- from_which,
234
- {"xmin", "xmax", "ymin", "ymax", "zmin", "zmax"},
235
- )
236
-
237
- if xrange is None:
238
- xrange = (0, tn.Lx - 1)
239
- if yrange is None:
240
- yrange = (0, tn.Ly - 1)
241
- if zrange is None:
242
- zrange = (0, tn.Lz - 1)
243
-
244
- self.xrange = xrange
245
- self.yrange = yrange
246
- self.zrange = zrange
247
- self.from_which = from_which
248
- self.plane = from_which[0]
249
-
250
- if self.plane == "x":
251
- # -> no rotation needed
252
- self.imin, self.imax = sorted(xrange)
253
- self.jmin, self.jmax = sorted(yrange)
254
- self.kmin, self.kmax = sorted(zrange)
255
- self.x_tag = tn.x_tag
256
- self.y_tag = tn.y_tag
257
- self.z_tag = tn.z_tag
258
- self.site_tag = tn.site_tag
259
- self.is_cyclic_x = tn.is_cyclic_x
260
- self.is_cyclic_y = tn.is_cyclic_y
261
- self.is_cyclic_z = tn.is_cyclic_z
262
-
263
- elif self.plane == "y":
264
- # -> (y, z, x)
265
- self.imin, self.imax = sorted(yrange)
266
- self.jmin, self.jmax = sorted(zrange)
267
- self.kmin, self.kmax = sorted(xrange)
268
- self.x_tag = tn.y_tag
269
- self.y_tag = tn.z_tag
270
- self.z_tag = tn.x_tag
271
- self.site_tag = lambda i, j, k: tn.site_tag(k, i, j)
272
- self.is_cyclic_x = tn.is_cyclic_y
273
- self.is_cyclic_y = tn.is_cyclic_z
274
- self.is_cyclic_z = tn.is_cyclic_x
275
-
276
- else: # self.plane == 'z'
277
- # -> (z, x, y)
278
- self.imin, self.imax = sorted(zrange)
279
- self.jmin, self.jmax = sorted(xrange)
280
- self.kmin, self.kmax = sorted(yrange)
281
- self.x_tag = tn.z_tag
282
- self.y_tag = tn.x_tag
283
- self.z_tag = tn.y_tag
284
- self.site_tag = lambda i, j, k: tn.site_tag(j, k, i)
285
- self.is_cyclic_x = tn.is_cyclic_z
286
- self.is_cyclic_y = tn.is_cyclic_x
287
- self.is_cyclic_z = tn.is_cyclic_y
288
-
289
- if "min" in self.from_which:
290
- # -> sweeps are increasing
291
- self.sweep = range(self.imin, self.imax + 1, +1)
292
- self.istep = +1
293
- else: # 'max'
294
- # -> sweeps are decreasing
295
- self.sweep = range(self.imax, self.imin - 1, -1)
296
- self.istep = -1
297
-
298
- @property
299
- def sweep_other(self):
300
- return itertools.product(
301
- range(self.jmin, self.jmax + 1),
302
- range(self.kmin, self.kmax + 1),
303
- )
304
-
305
- @functools.cached_property
306
- def cyclic_x(self):
307
- return self.is_cyclic_x(
308
- (self.jmin + self.jmax) // 2,
309
- (self.kmin + self.kmax) // 2,
310
- self.imin,
311
- self.imax,
312
- )
313
-
314
- @functools.cached_property
315
- def cyclic_y(self):
316
- return self.is_cyclic_y(
317
- (self.kmin + self.kmax) // 2,
318
- (self.imin + self.imax) // 2,
319
- self.jmin,
320
- self.jmax,
321
- )
322
-
323
- @functools.cached_property
324
- def cyclic_z(self):
325
- return self.is_cyclic_z(
326
- (self.imin + self.imax) // 2,
327
- (self.jmin + self.jmax) // 2,
328
- self.kmin,
329
- self.kmax,
330
- )
331
-
332
- def get_jnext(self, j):
333
- if j == self.jmax:
334
- if self.cyclic_y:
335
- # wrap around
336
- return self.jmin
337
- # no more steps
338
- return None
339
- # normal step
340
- return j + 1
341
-
342
- def get_knext(self, k):
343
- if k == self.kmax:
344
- if self.cyclic_z:
345
- # wrap around
346
- return self.kmin
347
- # no more steps
348
- return None
349
- # normal step
350
- return k + 1
351
-
352
-
353
- # reference for viewing a cube from each direction
354
- #
355
- # ┌──┐ ┌──┐ ┌──┐ ┌──┐ ┌──┐ ┌──┐
356
- # │y+│ │z+│ │x-│ │y-│ │z-│ │x+│
357
- # ┌──┼──┼──┐ ┌──┼──┼──┐ ┌──┼──┼──┐ ┌──┼──┼──┐ ┌──┼──┼──┐ ┌──┼──┼──┐
358
- # │z-│x-│z+│ │x-│y-│x+│ │y+│z-│y-│ │z-│x+│z+│ │x-│y+│x+│ │y+│z+│y-│
359
- # └──┼──┼──┘, └──┼──┼──┘, └──┼──┼──┘, └──┼──┼──┘, └──┼──┼──┘, └──┼──┼──┘
360
- # │y-│ │z-│ │x+│ │y+│ │z+│ │x-│
361
- # └──┘ └──┘ └──┘ └──┘ └──┘ └──┘
362
-
363
-
364
- _canonize_plane_opts = {
365
- "xmin": {
366
- "yreverse": False,
367
- "zreverse": False,
368
- "coordinate_order": "yz",
369
- "stepping_order": "zy",
370
- },
371
- "ymin": {
372
- "zreverse": False,
373
- "xreverse": True,
374
- "coordinate_order": "zx",
375
- "stepping_order": "xz",
376
- },
377
- "zmin": {
378
- "xreverse": True,
379
- "yreverse": True,
380
- "coordinate_order": "xy",
381
- "stepping_order": "yx",
382
- },
383
- "xmax": {
384
- "yreverse": True,
385
- "zreverse": True,
386
- "coordinate_order": "yz",
387
- "stepping_order": "zy",
388
- },
389
- "ymax": {
390
- "zreverse": True,
391
- "xreverse": False,
392
- "coordinate_order": "zx",
393
- "stepping_order": "xz",
394
- },
395
- "zmax": {
396
- "xreverse": False,
397
- "yreverse": False,
398
- "coordinate_order": "xy",
399
- "stepping_order": "yx",
400
- },
401
- }
402
-
403
-
404
- _compress_plane_opts = {
405
- "xmin": {
406
- "yreverse": True,
407
- "zreverse": True,
408
- "coordinate_order": "yz",
409
- "stepping_order": "zy",
410
- },
411
- "ymin": {
412
- "zreverse": True,
413
- "xreverse": False,
414
- "coordinate_order": "zx",
415
- "stepping_order": "xz",
416
- },
417
- "zmin": {
418
- "xreverse": False,
419
- "yreverse": False,
420
- "coordinate_order": "xy",
421
- "stepping_order": "yx",
422
- },
423
- "xmax": {
424
- "yreverse": False,
425
- "zreverse": False,
426
- "coordinate_order": "yz",
427
- "stepping_order": "zy",
428
- },
429
- "ymax": {
430
- "zreverse": False,
431
- "xreverse": True,
432
- "coordinate_order": "zx",
433
- "stepping_order": "xz",
434
- },
435
- "zmax": {
436
- "xreverse": True,
437
- "yreverse": True,
438
- "coordinate_order": "xy",
439
- "stepping_order": "yx",
440
- },
441
- }
442
-
443
- BOUNDARY_SEQUENCE_MAP = {
444
- "xmin": "xmin",
445
- "xmax": "xmax",
446
- "ymin": "ymin",
447
- "ymax": "ymax",
448
- "zmin": "zmin",
449
- "zmax": "zmax",
450
- }
451
-
452
-
453
- def parse_boundary_sequence(sequence):
454
- if isinstance(sequence, str):
455
- if sequence in BOUNDARY_SEQUENCE_MAP:
456
- return (sequence,)
457
- return tuple(BOUNDARY_SEQUENCE_MAP[s] for s in sequence)
458
-
459
-
460
- class TensorNetwork3D(TensorNetworkGen):
461
- """Mixin class for tensor networks with a cubic lattice three-dimensional
462
- structure.
463
- """
464
-
465
- _NDIMS = 3
466
- _EXTRA_PROPS = (
467
- "_site_tag_id",
468
- "_x_tag_id",
469
- "_y_tag_id",
470
- "_z_tag_id",
471
- "_Lx",
472
- "_Ly",
473
- "_Lz",
474
- )
475
-
476
- def _compatible_3d(self, other):
477
- """Check whether ``self`` and ``other`` are compatible 3D tensor
478
- networks such that they can remain a 3D tensor network when combined.
479
- """
480
- return isinstance(other, TensorNetwork3D) and all(
481
- getattr(self, e) == getattr(other, e)
482
- for e in TensorNetwork3D._EXTRA_PROPS
483
- )
484
-
485
- def combine(self, other, *, virtual=False, check_collisions=True):
486
- """Combine this tensor network with another, returning a new tensor
487
- network. If the two are compatible, cast the resulting tensor network
488
- to a :class:`TensorNetwork3D` instance.
489
-
490
- Parameters
491
- ----------
492
- other : TensorNetwork3D or TensorNetwork
493
- The other tensor network to combine with.
494
- virtual : bool, optional
495
- Whether the new tensor network should copy all the incoming tensors
496
- (``False``, the default), or view them as virtual (``True``).
497
- check_collisions : bool, optional
498
- Whether to check for index collisions between the two tensor
499
- networks before combining them. If ``True`` (the default), any
500
- inner indices that clash will be mangled.
501
-
502
- Returns
503
- -------
504
- TensorNetwork3D or TensorNetwork
505
- """
506
- new = super().combine(
507
- other, virtual=virtual, check_collisions=check_collisions
508
- )
509
- if self._compatible_3d(other):
510
- new.view_as_(TensorNetwork3D, like=self)
511
- return new
512
-
513
- @property
514
- def Lx(self):
515
- """The number of x-slices."""
516
- return self._Lx
517
-
518
- @property
519
- def Ly(self):
520
- """The number of y-slices."""
521
- return self._Ly
522
-
523
- @property
524
- def Lz(self):
525
- """The number of z-slices."""
526
- return self._Lz
527
-
528
- @property
529
- def nsites(self):
530
- """The total number of sites."""
531
- return self._Lx * self._Ly * self._Lz
532
-
533
- def site_tag(self, i, j=None, k=None):
534
- """The name of the tag specifiying the tensor at site ``(i, j, k)``."""
535
- if j is None:
536
- i, j, k = i
537
- if not isinstance(i, str):
538
- i = i % self.Lx
539
- if not isinstance(j, str):
540
- j = j % self.Ly
541
- if not isinstance(k, str):
542
- k = k % self.Lz
543
- return self.site_tag_id.format(i, j, k)
544
-
545
- @property
546
- def x_tag_id(self):
547
- """The string specifier for tagging each x-slice of this 3D TN."""
548
- return self._x_tag_id
549
-
550
- def x_tag(self, i):
551
- if not isinstance(i, str):
552
- i = i % self.Lx
553
- return self.x_tag_id.format(i)
554
-
555
- @property
556
- def x_tags(self):
557
- """A tuple of all of the ``Lx`` different x-slice tags."""
558
- return tuple(map(self.x_tag, range(self.Lx)))
559
-
560
- @property
561
- def y_tag_id(self):
562
- """The string specifier for tagging each y-slice of this 3D TN."""
563
- return self._y_tag_id
564
-
565
- def y_tag(self, j):
566
- if not isinstance(j, str):
567
- j = j % self.Ly
568
- return self.y_tag_id.format(j)
569
-
570
- @property
571
- def y_tags(self):
572
- """A tuple of all of the ``Ly`` different y-slice tags."""
573
- return tuple(map(self.y_tag, range(self.Ly)))
574
-
575
- @property
576
- def z_tag_id(self):
577
- """The string specifier for tagging each z-slice of this 3D TN."""
578
- return self._z_tag_id
579
-
580
- def z_tag(self, k):
581
- if not isinstance(k, str):
582
- k = k % self.Lz
583
- return self.z_tag_id.format(k)
584
-
585
- @property
586
- def z_tags(self):
587
- """A tuple of all of the ``Lz`` different z-slice tags."""
588
- return tuple(map(self.z_tag, range(self.Lz)))
589
-
590
- def maybe_convert_coo(self, coo):
591
- """Check if ``coo`` is a tuple of three ints and convert to the
592
- corresponding site tag if so.
593
- """
594
- if not isinstance(coo, str):
595
- try:
596
- i, j, k = map(int, coo)
597
- return self.site_tag(i, j, k)
598
- except (ValueError, TypeError):
599
- pass
600
- return coo
601
-
602
- def _get_tids_from_tags(self, tags, which="all"):
603
- """This is the function that lets coordinates such as ``(i, j, k)`` be
604
- used for many 'tag' based functions.
605
- """
606
- tags = self.maybe_convert_coo(tags)
607
- return super()._get_tids_from_tags(tags, which=which)
608
-
609
- def gen_site_coos(self):
610
- """Generate coordinates for all the sites in this 3D TN."""
611
- return product(range(self.Lx), range(self.Ly), range(self.Lz))
612
-
613
- def gen_bond_coos(self):
614
- """Generate pairs of coordinates for all the bonds in this 3D TN."""
615
- return gen_3d_bonds(
616
- self.Lx,
617
- self.Ly,
618
- self.Lz,
619
- steppers=[
620
- lambda i, j, k: (i + 1, j, k),
621
- lambda i, j, k: (i, j + 1, k),
622
- lambda i, j, k: (i, j, k + 1),
623
- ],
624
- cyclic=(
625
- self.is_cyclic_x(),
626
- self.is_cyclic_y(),
627
- self.is_cyclic_z(),
628
- ),
629
- )
630
-
631
- def valid_coo(self, coo, xrange=None, yrange=None, zrange=None):
632
- """Check whether ``coo`` is in-bounds.
633
-
634
- Parameters
635
- ----------
636
- coo : (int, int, int), optional
637
- The coordinates to check.
638
- xrange, yrange, zrange : (int, int), optional
639
- The range of allowed values for the x, y, and z coordinates.
640
-
641
- Returns
642
- -------
643
- bool
644
- """
645
- if xrange is None:
646
- xrange = (0, self.Lx - 1)
647
- if yrange is None:
648
- yrange = (0, self.Ly - 1)
649
- if zrange is None:
650
- zrange = (0, self.Lz - 1)
651
- return all(
652
- mn <= u <= mx for u, (mn, mx) in zip(coo, (xrange, yrange, zrange))
653
- )
654
-
655
- def get_ranges_present(self):
656
- """Return the range of site coordinates present in this TN.
657
-
658
- Returns
659
- -------
660
- xrange, yrange, zrange : tuple[tuple[int, int]]
661
- The minimum and maximum site coordinates present in each direction.
662
-
663
- Examples
664
- --------
665
-
666
- >>> tn = qtn.TN3D_rand(4, 4, 4, 2)
667
- >>> tn_sub = tn.select_local('I1,2,3', max_distance=1)
668
- >>> tn_sub.get_ranges_present()
669
- ((0, 2), (1, 3), (2, 3))
670
-
671
- """
672
- xmin = ymin = zmin = float("inf")
673
- xmax = ymax = zmax = float("-inf")
674
- for i, j, k in self.gen_sites_present():
675
- xmin = min(i, xmin)
676
- ymin = min(j, ymin)
677
- zmin = min(k, zmin)
678
- xmax = max(i, xmax)
679
- ymax = max(j, ymax)
680
- zmax = max(k, zmax)
681
- return (xmin, xmax), (ymin, ymax), (zmin, zmax)
682
-
683
- def is_cyclic_x(self, j=None, k=None, imin=None, imax=None):
684
- """Check if the x dimension is cyclic (periodic), specifically whether
685
- a bond exists between ``(imin, j, k)`` and ``(imax, j, k)``, with
686
- default values of ``imin = 0`` and ``imax = Lx - 1``, and ``j`` and
687
- ``k`` the center of the lattice. If ``imin`` and ``imax`` are adjacent
688
- then this is considered False, since there is no 'extra' connectivity.
689
- """
690
- if imin is None:
691
- imin = 0
692
- if imax is None:
693
- imax = self.Lx - 1
694
-
695
- if abs(imax - imin) <= 1:
696
- # first and last sites already connected -> a bit undefined
697
- return False
698
-
699
- if j is None:
700
- j = self.Ly // 2
701
- if k is None:
702
- k = self.Lz // 2
703
-
704
- return bool(
705
- bonds(
706
- self[self.site_tag(imin, j, k)],
707
- self[self.site_tag(imax, j, k)],
708
- )
709
- )
710
-
711
- def is_cyclic_y(self, k=None, i=None, jmin=None, jmax=None):
712
- """Check if the y dimension is cyclic (periodic), specifically whether
713
- a bond exists between ``(i, jmin, k)`` and ``(i, jmax, k)``, with
714
- default values of ``jmin = 0`` and ``jmax = Ly - 1``, and ``i`` and
715
- ``k`` the center of the lattice. If ``jmin`` and ``jmax`` are adjacent
716
- then this is considered False, since there is no 'extra' connectivity.
717
- """
718
- if jmin is None:
719
- jmin = 0
720
- if jmax is None:
721
- jmax = self.Ly - 1
722
-
723
- if abs(jmax - jmin) <= 1:
724
- # first and last sites already connected -> a bit undefined
725
- return False
726
-
727
- if i is None:
728
- i = self.Lx // 2
729
- if k is None:
730
- k = self.Lz // 2
731
- return bool(
732
- bonds(
733
- self[self.site_tag(i, jmin, k)],
734
- self[self.site_tag(i, jmax, k)],
735
- )
736
- )
737
-
738
- def is_cyclic_z(self, i=None, j=None, kmin=None, kmax=None):
739
- """Check if the z dimension is cyclic (periodic), specifically whether
740
- a bond exists between ``(i, j, kmin)`` and ``(i, j, kmax)``, with
741
- default values of ``kmin = 0`` and ``kmax = Lz - 1``, and ``i`` and
742
- ``j`` the center of the lattice. If ``kmin`` and ``kmax`` are adjacent
743
- then this is considered False, since there is no 'extra' connectivity.
744
- """
745
- if kmin is None:
746
- kmin = 0
747
- if kmax is None:
748
- kmax = self.Lz - 1
749
-
750
- if abs(kmax - kmin) <= 1:
751
- # first and last sites already connected -> a bit undefined
752
- return False
753
-
754
- if i is None:
755
- i = self.Lx // 2
756
- if j is None:
757
- j = self.Ly // 2
758
- return bool(
759
- bonds(
760
- self[self.site_tag(i, j, kmin)],
761
- self[self.site_tag(i, j, kmax)],
762
- )
763
- )
764
-
765
- def __getitem__(self, key):
766
- """Key based tensor selection, checking for integer based shortcut."""
767
- return super().__getitem__(self.maybe_convert_coo(key))
768
-
769
- def _repr_info(self):
770
- info = super()._repr_info()
771
- info["Lx"] = self.Lx
772
- info["Ly"] = self.Ly
773
- info["Lz"] = self.Lz
774
- info["max_bond"] = self.max_bond()
775
- return info
776
-
777
- def flatten(self, fuse_multibonds=True, inplace=False):
778
- """Contract all tensors corresponding to each site into one."""
779
- tn = self if inplace else self.copy()
780
-
781
- for i, j, k in self.gen_site_coos():
782
- tn ^= (i, j, k)
783
-
784
- if fuse_multibonds:
785
- tn.fuse_multibonds_()
786
-
787
- return tn.view_as_(TensorNetwork3DFlat, like=self)
788
-
789
- flatten_ = functools.partialmethod(flatten, inplace=True)
790
-
791
- def gen_pairs(
792
- self,
793
- xrange=None,
794
- yrange=None,
795
- zrange=None,
796
- xreverse=False,
797
- yreverse=False,
798
- zreverse=False,
799
- coordinate_order="xyz",
800
- xstep=None,
801
- ystep=None,
802
- zstep=None,
803
- stepping_order="xyz",
804
- step_only=None,
805
- ):
806
- """Helper function for generating pairs of cooordinates for all bonds
807
- within a certain range, optionally specifying an order.
808
-
809
- Parameters
810
- ----------
811
- xrange, yrange, zrange : (int, int), optional
812
- The range of allowed values for the x, y, and z coordinates.
813
- xreverse, yreverse, zreverse : bool, optional
814
- Whether to reverse the order of the x, y, and z sweeps.
815
- coordinate_order : str, optional
816
- The order in which to sweep the x, y, and z coordinates. Earlier
817
- dimensions will change slower. If the corresponding range has
818
- size 1 then that dimension doesn't need to be specified.
819
- xstep, ystep, zstep : int, optional
820
- When generating a bond, step in this direction to yield the
821
- neighboring coordinate. By default, these follow ``xreverse``,
822
- ``yreverse``, and ``zreverse`` respectively.
823
- stepping_order : str, optional
824
- The order in which to step the x, y, and z coordinates to generate
825
- bonds. Does not need to include all dimensions.
826
- step_only : int, optional
827
- Only perform the ith steps in ``stepping_order``, used to
828
- interleave canonizing and compressing for example.
829
-
830
- Yields
831
- ------
832
- coo_a, coo_b : ((int, int, int), (int, int, int))
833
- """
834
- if xrange is None:
835
- xrange = (0, self.Lx - 1)
836
- if yrange is None:
837
- yrange = (0, self.Ly - 1)
838
- if zrange is None:
839
- zrange = (0, self.Lz - 1)
840
-
841
- # generate the sites and order we will visit them in
842
- sweeps = {
843
- "x": (
844
- range(min(xrange), max(xrange) + 1, +1)
845
- if not xreverse
846
- else range(max(xrange), min(xrange) - 1, -1)
847
- ),
848
- "y": (
849
- range(min(yrange), max(yrange) + 1, +1)
850
- if not yreverse
851
- else range(max(yrange), min(yrange) - 1, -1)
852
- ),
853
- "z": (
854
- range(min(zrange), max(zrange) + 1, +1)
855
- if not zreverse
856
- else range(max(zrange), min(zrange) - 1, -1)
857
- ),
858
- }
859
-
860
- # for convenience, allow subselecting part of stepping_order only
861
- if step_only is not None:
862
- stepping_order = stepping_order[step_only]
863
-
864
- # at each step generate the bonds
865
- if xstep is None:
866
- xstep = -1 if xreverse else +1
867
- if ystep is None:
868
- ystep = -1 if yreverse else +1
869
- if zstep is None:
870
- zstep = -1 if zreverse else +1
871
- steps = {
872
- "x": lambda i, j, k: (i + xstep, j, k),
873
- "y": lambda i, j, k: (i, j + ystep, k),
874
- "z": lambda i, j, k: (i, j, k + zstep),
875
- }
876
-
877
- # make sure all coordinates exist - only allow them not to be specified
878
- # if their range is a unit slice
879
- for w in "xyz":
880
- if w not in coordinate_order:
881
- if len(sweeps[w]) > 1:
882
- raise ValueError(
883
- f"{w} not in coordinate_order and is not size 1."
884
- )
885
- else:
886
- # just append -> it won't change order as coord is constant
887
- coordinate_order += w
888
- xi, yi, zi = map(coordinate_order.index, "xyz")
889
-
890
- # generate the pairs
891
- for perm_coo_a in product(*(sweeps[xyz] for xyz in coordinate_order)):
892
- coo_a = perm_coo_a[xi], perm_coo_a[yi], perm_coo_a[zi]
893
- for xyz in stepping_order:
894
- coo_b = steps[xyz](*coo_a)
895
- # filter out bonds with are out of bounds
896
- if self.valid_coo(coo_b, xrange, yrange, zrange):
897
- yield coo_a, coo_b
898
-
899
- def canonize_plane(
900
- self,
901
- xrange,
902
- yrange,
903
- zrange,
904
- equalize_norms=False,
905
- canonize_opts=None,
906
- **gen_pair_opts,
907
- ):
908
- """Canonize every pair of tensors within a subrange, optionally
909
- specifying a order to visit those pairs in.
910
- """
911
- canonize_opts = ensure_dict(canonize_opts)
912
- canonize_opts.setdefault("equalize_norms", equalize_norms)
913
-
914
- pairs = self.gen_pairs(
915
- xrange=xrange,
916
- yrange=yrange,
917
- zrange=zrange,
918
- **gen_pair_opts,
919
- )
920
-
921
- for coo_a, coo_b in pairs:
922
- tag_a = self.site_tag(*coo_a)
923
- tag_b = self.site_tag(*coo_b)
924
-
925
- # make sure single tensor at each site, skip if none
926
- try:
927
- num_a = len(self.tag_map[tag_a])
928
- if num_a > 1:
929
- self ^= tag_a
930
- except KeyError:
931
- continue
932
- try:
933
- num_b = len(self.tag_map[tag_b])
934
- if num_b > 1:
935
- self ^= tag_b
936
- except KeyError:
937
- continue
938
-
939
- self.canonize_between(tag_a, tag_b, **canonize_opts)
940
-
941
- def compress_plane(
942
- self,
943
- xrange,
944
- yrange,
945
- zrange,
946
- max_bond=None,
947
- cutoff=1e-10,
948
- equalize_norms=False,
949
- compress_opts=None,
950
- **gen_pair_opts,
951
- ):
952
- """Compress every pair of tensors within a subrange, optionally
953
- specifying a order to visit those pairs in.
954
- """
955
- compress_opts = ensure_dict(compress_opts)
956
- compress_opts.setdefault("absorb", "both")
957
- compress_opts.setdefault("equalize_norms", equalize_norms)
958
-
959
- pairs = self.gen_pairs(
960
- xrange=xrange,
961
- yrange=yrange,
962
- zrange=zrange,
963
- **gen_pair_opts,
964
- )
965
-
966
- for coo_a, coo_b in pairs:
967
- tag_a = self.site_tag(*coo_a)
968
- tag_b = self.site_tag(*coo_b)
969
-
970
- # make sure single tensor at each site, skip if none
971
- try:
972
- num_a = len(self.tag_map[tag_a])
973
- if num_a > 1:
974
- self ^= tag_a
975
- except KeyError:
976
- continue
977
- try:
978
- num_b = len(self.tag_map[tag_b])
979
- if num_b > 1:
980
- self ^= tag_b
981
- except KeyError:
982
- continue
983
-
984
- self.compress_between(
985
- tag_a, tag_b, max_bond=max_bond, cutoff=cutoff, **compress_opts
986
- )
987
-
988
- def _contract_boundary_core_via_2d(
989
- self,
990
- xrange,
991
- yrange,
992
- zrange,
993
- from_which,
994
- max_bond,
995
- cutoff=1e-10,
996
- method="local-early",
997
- layer_tags=None,
998
- **compress_opts,
999
- ):
1000
- from quimb.tensor.tensor_2d_compress import tensor_network_2d_compress
1001
-
1002
- r2d = Rotator3D(self, xrange, yrange, zrange, from_which)
1003
- site_tag = r2d.site_tag
1004
- istep = r2d.istep
1005
-
1006
- def _do_compress(site_tag_tmps):
1007
- nonlocal self
1008
-
1009
- # split off the boundary network
1010
- self, tn_boundary = self.partition(site_tag_tmps, inplace=True)
1011
-
1012
- # compress it inplace
1013
- tensor_network_2d_compress(
1014
- tn_boundary,
1015
- max_bond=max_bond,
1016
- cutoff=cutoff,
1017
- method=method,
1018
- site_tags=site_tag_tmps,
1019
- inplace=True,
1020
- **compress_opts,
1021
- )
1022
-
1023
- # recombine with the main network
1024
- self |= tn_boundary
1025
-
1026
- # maybe compress the initial row, which may be multiple layers
1027
- # and have effective bond dimension > max_bond already
1028
- site_tag_tmps = [
1029
- site_tag(r2d.sweep[0], j, k) for j, k in r2d.sweep_other
1030
- ]
1031
- if any(len(self.tag_map[st]) > 1 for st in site_tag_tmps):
1032
- _do_compress(site_tag_tmps)
1033
-
1034
- site_tag_tmps = [f"__ST{j},{k}__" for j, k in r2d.sweep_other]
1035
-
1036
- if layer_tags is None:
1037
- layer_tags = [None]
1038
-
1039
- for i in r2d.sweep[:-1]:
1040
- for layer_tag in layer_tags:
1041
- for (j, k), st in zip(r2d.sweep_other, site_tag_tmps):
1042
- # group outer single tensor with inner tensor(s)
1043
- tag1 = site_tag(i, j, k) # outer
1044
- tag2 = site_tag(i + istep, j, k) # inner
1045
- if layer_tag is None:
1046
- # tag and compress any inner tensors
1047
- self.select_any((tag1, tag2)).add_tag(st)
1048
- else:
1049
- # only tag and compress one inner layer
1050
- self.select_all((tag1,)).add_tag(st)
1051
- self.select_all((tag2, layer_tag)).add_tag(st)
1052
-
1053
- _do_compress(site_tag_tmps)
1054
-
1055
- self.drop_tags(site_tag_tmps)
1056
-
1057
- def _contract_boundary_core(
1058
- self,
1059
- xrange,
1060
- yrange,
1061
- zrange,
1062
- from_which,
1063
- max_bond,
1064
- cutoff=1e-10,
1065
- canonize=True,
1066
- canonize_interleave=True,
1067
- layer_tags=None,
1068
- compress_late=True,
1069
- equalize_norms=False,
1070
- compress_opts=None,
1071
- canonize_opts=None,
1072
- ):
1073
- canonize_opts = ensure_dict(canonize_opts)
1074
- canonize_opts.setdefault("absorb", "right")
1075
- compress_opts = ensure_dict(compress_opts)
1076
- compress_opts.setdefault("absorb", "both")
1077
-
1078
- r3d = Rotator3D(self, xrange, yrange, zrange, from_which)
1079
- site_tag = r3d.site_tag
1080
- plane, istep = r3d.plane, r3d.istep
1081
- jmin, jmax = r3d.jmin, r3d.jmax
1082
- kmin, kmax = r3d.kmin, r3d.kmax
1083
-
1084
- if canonize_interleave:
1085
- # interleave canonizing and compressing in each direction
1086
- step_onlys = [0, 1]
1087
- else:
1088
- # perform whole sweep of canonizing before compressing
1089
- step_onlys = [None]
1090
-
1091
- if layer_tags is None:
1092
- layer_tags = [None]
1093
-
1094
- for i in r3d.sweep[:-1]:
1095
- for layer_tag in layer_tags:
1096
- for j in range(jmin, jmax + 1):
1097
- for k in range(kmin, kmax + 1):
1098
- tag1 = site_tag(i, j, k) # outer
1099
- tag2 = site_tag(i + istep, j, k) # inner
1100
-
1101
- if (tag1 not in self.tag_map) or (
1102
- tag2 not in self.tag_map
1103
- ):
1104
- # allow completely missing sites
1105
- continue
1106
-
1107
- if (layer_tag is None) or len(self.tag_map[tag2]) == 1:
1108
- # contract *any* tensors with pair of coordinates
1109
- self.contract_((tag1, tag2), which="any")
1110
- else:
1111
- # ensure boundary is single tensor per site
1112
- if len(self.tag_map[tag1]) > 1:
1113
- self ^= tag1
1114
-
1115
- # contract specific pair (i.e. one 'inner' layer)
1116
- self.contract_between(
1117
- tag1,
1118
- (tag2, layer_tag),
1119
- equalize_norms=equalize_norms,
1120
- )
1121
-
1122
- # drop inner site tag merged into outer boundary so
1123
- # we can still uniquely identify inner tensors
1124
- if layer_tag != layer_tags[-1]:
1125
- self[tag1].drop_tags(tag2)
1126
-
1127
- if not compress_late:
1128
- (tid1,) = self.tag_map[tag1]
1129
- for tidn in self._get_neighbor_tids(tid1):
1130
- t1, tn = self._tids_get(tid1, tidn)
1131
- if bonds_size(t1, tn) > max_bond:
1132
- self._compress_between_tids(
1133
- tid1,
1134
- tidn,
1135
- max_bond=max_bond,
1136
- cutoff=cutoff,
1137
- equalize_norms=equalize_norms,
1138
- **compress_opts,
1139
- )
1140
-
1141
- if compress_late:
1142
- for step_only in step_onlys:
1143
- if canonize:
1144
- self.canonize_plane(
1145
- xrange=xrange if plane != "x" else (i, i),
1146
- yrange=yrange if plane != "y" else (i, i),
1147
- zrange=zrange if plane != "z" else (i, i),
1148
- equalize_norms=equalize_norms,
1149
- canonize_opts=canonize_opts,
1150
- step_only=step_only,
1151
- **_canonize_plane_opts[from_which],
1152
- )
1153
- self.compress_plane(
1154
- xrange=xrange if plane != "x" else (i, i),
1155
- yrange=yrange if plane != "y" else (i, i),
1156
- zrange=zrange if plane != "z" else (i, i),
1157
- max_bond=max_bond,
1158
- cutoff=cutoff,
1159
- equalize_norms=equalize_norms,
1160
- compress_opts=compress_opts,
1161
- step_only=step_only,
1162
- **_compress_plane_opts[from_which],
1163
- )
1164
-
1165
- def _contract_boundary_projector(
1166
- self,
1167
- xrange,
1168
- yrange,
1169
- zrange,
1170
- from_which,
1171
- max_bond=None,
1172
- cutoff=1e-10,
1173
- canonize=False,
1174
- layer_tags=None,
1175
- lazy=False,
1176
- equalize_norms=False,
1177
- optimize="auto-hq",
1178
- compress_opts=None,
1179
- canonize_opts=None,
1180
- ):
1181
- compress_opts = ensure_dict(compress_opts)
1182
-
1183
- r = Rotator3D(self, xrange, yrange, zrange, from_which)
1184
-
1185
- if canonize:
1186
- canonize_opts = ensure_dict(canonize_opts)
1187
- canonize_opts.setdefault("max_iterations", 2)
1188
- self.select(r.x_tag(r.sweep[0])).gauge_all_(**canonize_opts)
1189
-
1190
- for i, inext in pairwise(r.sweep):
1191
- # we compute the projectors from an untouched copy
1192
- tn_calc = self.copy()
1193
-
1194
- for j, k in r.sweep_other:
1195
- # get next neighboring sites, wrapping around if necessary
1196
- jnext = r.get_jnext(j)
1197
- knext = r.get_knext(k)
1198
-
1199
- tag_ijk = r.site_tag(i, j, k)
1200
- tag_ip1jk = r.site_tag(inext, j, k)
1201
- rtags = (tag_ijk, tag_ip1jk)
1202
-
1203
- poss_ltags = []
1204
- if jnext is not None:
1205
- poss_ltags.append(
1206
- (r.site_tag(i, jnext, k), r.site_tag(inext, jnext, k))
1207
- )
1208
- if knext is not None:
1209
- poss_ltags.append(
1210
- (r.site_tag(i, j, knext), r.site_tag(inext, j, knext))
1211
- )
1212
-
1213
- for ltags in poss_ltags:
1214
- tn_calc.insert_compressor_between_regions(
1215
- ltags,
1216
- rtags,
1217
- new_ltags=ltags,
1218
- new_rtags=rtags,
1219
- insert_into=self,
1220
- max_bond=max_bond,
1221
- cutoff=cutoff,
1222
- **compress_opts,
1223
- )
1224
-
1225
- if not lazy:
1226
- # contract each pair of boundary tensors with their projectors
1227
- for j, k in r.sweep_other:
1228
- self.contract_tags_(
1229
- (r.site_tag(i, j, k), r.site_tag(inext, j, k)),
1230
- optimize=optimize,
1231
- )
1232
-
1233
- if equalize_norms:
1234
- for t in self.select_tensors(r.x_tag(inext)):
1235
- self.strip_exponent(t, equalize_norms)
1236
- if equalize_norms:
1237
- for t in self.select_tensors(r.x_tag(i1)):
1238
- self.strip_exponent(t, equalize_norms)
1239
-
1240
- def contract_boundary_from(
1241
- self,
1242
- xrange,
1243
- yrange,
1244
- zrange,
1245
- from_which,
1246
- max_bond=None,
1247
- *,
1248
- cutoff=1e-10,
1249
- mode="peps",
1250
- equalize_norms=False,
1251
- compress_opts=None,
1252
- canonize_opts=None,
1253
- inplace=False,
1254
- **contract_boundary_opts,
1255
- ):
1256
- """Unified entrypoint for contracting any rectangular patch of tensors
1257
- from any direction, with any boundary method.
1258
- """
1259
- tn = self if inplace else self.copy()
1260
-
1261
- # universal options
1262
- contract_boundary_opts["xrange"] = xrange
1263
- contract_boundary_opts["yrange"] = yrange
1264
- contract_boundary_opts["zrange"] = zrange
1265
- contract_boundary_opts["from_which"] = from_which
1266
- contract_boundary_opts["max_bond"] = max_bond
1267
- contract_boundary_opts["cutoff"] = cutoff
1268
- contract_boundary_opts["equalize_norms"] = equalize_norms
1269
- contract_boundary_opts["compress_opts"] = compress_opts
1270
-
1271
- if mode == "peps":
1272
- return tn._contract_boundary_core(
1273
- canonize_opts=canonize_opts,
1274
- **contract_boundary_opts,
1275
- )
1276
-
1277
- if mode == "l2bp3d":
1278
- return tn._contract_boundary_l2bp(**contract_boundary_opts)
1279
-
1280
- if mode == "projector3d":
1281
- return tn._contract_boundary_projector(
1282
- canonize_opts=canonize_opts, **contract_boundary_opts
1283
- )
1284
-
1285
- return tn._contract_boundary_core_via_2d(
1286
- method=mode, **contract_boundary_opts
1287
- )
1288
-
1289
- contract_boundary_from_ = functools.partialmethod(
1290
- contract_boundary_from, inplace=True
1291
- )
1292
-
1293
- def _contract_interleaved_boundary_sequence(
1294
- self,
1295
- *,
1296
- contract_boundary_opts=None,
1297
- sequence=None,
1298
- xmin=None,
1299
- xmax=None,
1300
- ymin=None,
1301
- ymax=None,
1302
- zmin=None,
1303
- zmax=None,
1304
- max_separation=1,
1305
- max_unfinished=1,
1306
- around=None,
1307
- equalize_norms=False,
1308
- canonize=False,
1309
- canonize_opts=None,
1310
- final_contract=True,
1311
- final_contract_opts=None,
1312
- optimize="auto-hq",
1313
- progbar=False,
1314
- inplace=False,
1315
- ):
1316
- """Unified handler for performing iterleaved contractions in a
1317
- sequence of inwards boundary directions.
1318
- """
1319
- tn = self if inplace else self.copy()
1320
-
1321
- contract_boundary_opts = ensure_dict(contract_boundary_opts)
1322
- if canonize:
1323
- canonize_opts = ensure_dict(canonize_opts)
1324
- canonize_opts.setdefault("max_iterations", 2)
1325
-
1326
- if progbar:
1327
- pbar = Progbar()
1328
- pbar.set_description(
1329
- f"contracting boundary, Lx={tn.Lx}, Ly={tn.Ly}, Lz={tn.Lz}"
1330
- )
1331
- else:
1332
- pbar = None
1333
-
1334
- # set default starting borders
1335
- if any(d is None for d in (xmin, xmax, ymin, ymax, zmin, zmax)):
1336
- (
1337
- (auto_xmin, auto_xmax),
1338
- (auto_ymin, auto_ymax),
1339
- (auto_zmin, auto_zmax),
1340
- ) = self.get_ranges_present()
1341
-
1342
- # location of current boundaries
1343
- boundaries = {
1344
- "xmin": auto_xmin if xmin is None else xmin,
1345
- "xmax": auto_xmax if xmax is None else xmax,
1346
- "ymin": auto_ymin if ymin is None else ymin,
1347
- "ymax": auto_ymax if ymax is None else ymax,
1348
- "zmin": auto_zmin if zmin is None else zmin,
1349
- "zmax": auto_zmax if zmax is None else zmax,
1350
- }
1351
- separations = {
1352
- d: boundaries[f"{d}max"] - boundaries[f"{d}min"] for d in "xyz"
1353
- }
1354
- boundary_tags = {
1355
- "xmin": tn.x_tag(boundaries["xmin"]),
1356
- "xmax": tn.x_tag(boundaries["xmax"]),
1357
- "ymin": tn.y_tag(boundaries["ymin"]),
1358
- "ymax": tn.y_tag(boundaries["ymax"]),
1359
- "zmin": tn.z_tag(boundaries["zmin"]),
1360
- "zmax": tn.z_tag(boundaries["zmax"]),
1361
- }
1362
- if around is not None:
1363
- target_xmin = min(x[0] for x in around)
1364
- target_xmax = max(x[0] for x in around)
1365
- target_ymin = min(x[1] for x in around)
1366
- target_ymax = max(x[1] for x in around)
1367
- target_zmin = min(x[2] for x in around)
1368
- target_zmax = max(x[2] for x in around)
1369
- target_check = {
1370
- "xmin": lambda x: x >= target_xmin - 1,
1371
- "xmax": lambda x: x <= target_xmax + 1,
1372
- "ymin": lambda y: y >= target_ymin - 1,
1373
- "ymax": lambda y: y <= target_ymax + 1,
1374
- "zmin": lambda z: z >= target_zmin - 1,
1375
- "zmax": lambda z: z <= target_zmax + 1,
1376
- }
1377
-
1378
- if sequence is None:
1379
- sequence = []
1380
- if tn.is_cyclic_x():
1381
- sequence.append("xmin")
1382
- else:
1383
- sequence.extend(["xmin", "xmax"])
1384
-
1385
- if tn.is_cyclic_y():
1386
- sequence.append("ymin")
1387
- else:
1388
- sequence.extend(["ymin", "ymax"])
1389
-
1390
- if tn.is_cyclic_z():
1391
- sequence.append("zmin")
1392
- else:
1393
- sequence.extend(["zmin", "zmax"])
1394
-
1395
- else:
1396
- sequence = parse_boundary_sequence(sequence)
1397
-
1398
- def _is_finished(direction):
1399
- return (
1400
- # two opposing sides have got sufficiently close
1401
- (separations[direction[0]] <= max_separation)
1402
- or (
1403
- # there is a target region
1404
- (around is not None)
1405
- and
1406
- # and we have reached it
1407
- target_check[direction](boundaries[direction])
1408
- )
1409
- )
1410
-
1411
- sequence = [d for d in sequence if not _is_finished(d)]
1412
- while sequence:
1413
- direction = sequence.pop(0)
1414
- if _is_finished(direction):
1415
- # just remove direction from sequence
1416
- continue
1417
- # do a contraction, and keep direction in sequence to try again
1418
- sequence.append(direction)
1419
-
1420
- if pbar is not None:
1421
- pbar.set_description(
1422
- f"contracting {direction}, "
1423
- f"Lx={separations['x'] + 1}, "
1424
- f"Ly={separations['y'] + 1}, "
1425
- f"Lz={separations['z'] + 1}"
1426
- )
1427
-
1428
- if canonize:
1429
- tn.select(boundary_tags[direction]).gauge_all_(**canonize_opts)
1430
-
1431
- if direction[0] == "x":
1432
- if direction[1:] == "min":
1433
- xrange = (boundaries["xmin"], boundaries["xmin"] + 1)
1434
- else: # xmax
1435
- xrange = (boundaries["xmax"] - 1, boundaries["xmax"])
1436
- yrange = (boundaries["ymin"], boundaries["ymax"])
1437
- zrange = (boundaries["zmin"], boundaries["zmax"])
1438
- elif direction[0] == "y":
1439
- if direction[1:] == "min":
1440
- yrange = (boundaries["ymin"], boundaries["ymin"] + 1)
1441
- else: # ymax
1442
- yrange = (boundaries["ymax"] - 1, boundaries["ymax"])
1443
- xrange = (boundaries["xmin"], boundaries["xmax"])
1444
- zrange = (boundaries["zmin"], boundaries["zmax"])
1445
- else: # z
1446
- if direction[1:] == "min":
1447
- zrange = (boundaries["zmin"], boundaries["zmin"] + 1)
1448
- else: # zmax
1449
- zrange = (boundaries["zmax"] - 1, boundaries["zmax"])
1450
- xrange = (boundaries["xmin"], boundaries["xmax"])
1451
- yrange = (boundaries["ymin"], boundaries["ymax"])
1452
-
1453
- # do the contractions!
1454
- tn.contract_boundary_from_(
1455
- xrange=xrange,
1456
- yrange=yrange,
1457
- zrange=zrange,
1458
- from_which=direction,
1459
- equalize_norms=equalize_norms,
1460
- **contract_boundary_opts,
1461
- )
1462
-
1463
- # update the boundaries and separations
1464
- xyz, minmax = direction[0], direction[1:]
1465
- separations[xyz] -= 1
1466
- if minmax == "min":
1467
- boundaries[direction] += 1
1468
- else:
1469
- boundaries[direction] -= 1
1470
-
1471
- if pbar is not None:
1472
- pbar.update()
1473
-
1474
- # check if enough directions are finished -> reached max separation
1475
- if (
1476
- sum(separations[d] > max_separation for d in "xyz")
1477
- <= max_unfinished
1478
- ):
1479
- break
1480
-
1481
- if equalize_norms is True:
1482
- tn.equalize_norms_()
1483
-
1484
- if pbar is not None:
1485
- pbar.set_description(
1486
- f"contracted boundary, "
1487
- f"Lx={separations['x'] + 1}, "
1488
- f"Ly={separations['y'] + 1}, "
1489
- f"Lz={separations['z'] + 1}"
1490
- )
1491
- pbar.close()
1492
-
1493
- if final_contract and (around is None):
1494
- final_contract_opts = ensure_dict(final_contract_opts)
1495
- final_contract_opts.setdefault("optimize", optimize)
1496
- final_contract_opts.setdefault("inplace", inplace)
1497
- return tn.contract(**final_contract_opts)
1498
-
1499
- return tn
1500
-
1501
- def contract_boundary(
1502
- self,
1503
- max_bond=None,
1504
- *,
1505
- cutoff=1e-10,
1506
- mode="peps",
1507
- canonize=True,
1508
- compress_opts=None,
1509
- sequence=None,
1510
- xmin=None,
1511
- xmax=None,
1512
- ymin=None,
1513
- ymax=None,
1514
- zmin=None,
1515
- zmax=None,
1516
- max_separation=1,
1517
- max_unfinished=1,
1518
- around=None,
1519
- equalize_norms=False,
1520
- final_contract=True,
1521
- final_contract_opts=None,
1522
- optimize="auto-hq",
1523
- progbar=False,
1524
- inplace=False,
1525
- **contract_boundary_opts,
1526
- ):
1527
- """ """
1528
- contract_boundary_opts["max_bond"] = max_bond
1529
- contract_boundary_opts["cutoff"] = cutoff
1530
- contract_boundary_opts["mode"] = mode
1531
- contract_boundary_opts["canonize"] = canonize
1532
- contract_boundary_opts["compress_opts"] = compress_opts
1533
-
1534
- return self._contract_interleaved_boundary_sequence(
1535
- contract_boundary_opts=contract_boundary_opts,
1536
- sequence=sequence,
1537
- xmin=xmin,
1538
- xmax=xmax,
1539
- ymin=ymin,
1540
- ymax=ymax,
1541
- zmin=zmin,
1542
- zmax=zmax,
1543
- max_separation=max_separation,
1544
- max_unfinished=max_unfinished,
1545
- around=around,
1546
- equalize_norms=equalize_norms,
1547
- final_contract=final_contract,
1548
- final_contract_opts=final_contract_opts,
1549
- optimize=optimize,
1550
- progbar=progbar,
1551
- inplace=inplace,
1552
- )
1553
-
1554
- contract_boundary_ = functools.partialmethod(
1555
- contract_boundary, inplace=True
1556
- )
1557
-
1558
- def contract_ctmrg(
1559
- self,
1560
- max_bond=None,
1561
- *,
1562
- cutoff=1e-10,
1563
- canonize=False,
1564
- canonize_opts=None,
1565
- lazy=False,
1566
- mode="projector",
1567
- compress_opts=None,
1568
- sequence=None,
1569
- xmin=None,
1570
- xmax=None,
1571
- ymin=None,
1572
- ymax=None,
1573
- zmin=None,
1574
- zmax=None,
1575
- max_separation=1,
1576
- max_unfinished=1,
1577
- around=None,
1578
- equalize_norms=False,
1579
- final_contract=True,
1580
- final_contract_opts=None,
1581
- progbar=False,
1582
- inplace=False,
1583
- **contract_boundary_opts,
1584
- ):
1585
- contract_boundary_opts["max_bond"] = max_bond
1586
- contract_boundary_opts["cutoff"] = cutoff
1587
- contract_boundary_opts["mode"] = mode
1588
- contract_boundary_opts["compress_opts"] = compress_opts
1589
- contract_boundary_opts["lazy"] = lazy
1590
-
1591
- if lazy:
1592
- # we are implicitly asking for the tensor network
1593
- final_contract = False
1594
-
1595
- if sequence is None:
1596
- sequence = []
1597
- if self.is_cyclic_x():
1598
- sequence.append("xmin")
1599
- else:
1600
- sequence.extend(["xmin", "xmax"])
1601
-
1602
- if self.is_cyclic_y():
1603
- sequence.append("ymin")
1604
- else:
1605
- sequence.extend(["ymin", "ymax"])
1606
-
1607
- if self.is_cyclic_z():
1608
- sequence.append("zmin")
1609
- else:
1610
- sequence.extend(["zmin", "zmax"])
1611
-
1612
- return self._contract_interleaved_boundary_sequence(
1613
- contract_boundary_opts=contract_boundary_opts,
1614
- canonize=canonize,
1615
- canonize_opts=canonize_opts,
1616
- sequence=sequence,
1617
- xmin=xmin,
1618
- xmax=xmax,
1619
- ymin=ymin,
1620
- ymax=ymax,
1621
- zmin=zmin,
1622
- zmax=zmax,
1623
- max_separation=max_separation,
1624
- max_unfinished=max_unfinished,
1625
- around=around,
1626
- equalize_norms=equalize_norms,
1627
- final_contract=final_contract,
1628
- final_contract_opts=final_contract_opts,
1629
- progbar=progbar,
1630
- inplace=inplace,
1631
- )
1632
-
1633
- contract_ctmrg_ = functools.partialmethod(contract_ctmrg, inplace=True)
1634
-
1635
- def _compute_plane_envs(
1636
- self,
1637
- xrange,
1638
- yrange,
1639
- zrange,
1640
- from_which,
1641
- envs=None,
1642
- storage_factory=None,
1643
- **contract_boundary_opts,
1644
- ):
1645
- """Compute all 'plane' environments for the cube given by
1646
- ``xrange``, ``yrange``, ``zrange``, with direction given by
1647
- ``from_which``.
1648
- """
1649
- tn = self.copy()
1650
-
1651
- # rotate virtually
1652
- r3d = Rotator3D(tn, xrange, yrange, zrange, from_which)
1653
- plane, p_tag, istep = r3d.plane, r3d.x_tag, r3d.istep
1654
-
1655
- # 0th plane has no environment, 1st plane's environment is the 0th
1656
- if envs is None:
1657
- if storage_factory is not None:
1658
- envs = storage_factory()
1659
- else:
1660
- envs = {}
1661
-
1662
- envs[r3d.sweep[1]] = tn.select_any(p_tag(r3d.sweep[0]), virtual=False)
1663
-
1664
- for i in r3d.sweep[:-2]:
1665
- # contract the boundary in one step
1666
- tn.contract_boundary_from_(
1667
- xrange=xrange if plane != "x" else (i, i + istep),
1668
- yrange=yrange if plane != "y" else (i, i + istep),
1669
- zrange=zrange if plane != "z" else (i, i + istep),
1670
- from_which=from_which,
1671
- **contract_boundary_opts,
1672
- )
1673
- # set the boundary as the environment for the next plane beyond
1674
- envs[i + 2 * istep] = tn.select_any(
1675
- p_tag(i + istep), virtual=False
1676
- )
1677
-
1678
- return envs
1679
-
1680
- def _maybe_compute_cell_env(
1681
- self,
1682
- key,
1683
- envs=None,
1684
- storage_factory=None,
1685
- boundary_order=None,
1686
- **contract_boundary_opts,
1687
- ):
1688
- """Recursively compute the necessary 2D, 1D, and 0D environments."""
1689
- if not key:
1690
- # env is the whole TN
1691
- return self.copy()
1692
-
1693
- if envs is None:
1694
- if storage_factory is not None:
1695
- envs = storage_factory()
1696
- else:
1697
- envs = {}
1698
-
1699
- if key in envs:
1700
- return envs[key].copy()
1701
-
1702
- if boundary_order is None:
1703
- scores = {"x": -self.Lx, "y": -self.Ly, "z": -self.Lz}
1704
- boundary_order = sorted(scores, key=scores.__getitem__)
1705
-
1706
- # check already available parent environments
1707
- for parent_key in itertools.combinations(key, len(key) - 1):
1708
- if parent_key in envs:
1709
- parent_tn = envs[parent_key].copy()
1710
- break
1711
- else:
1712
- # choose which to compute next based on `boundary_order`
1713
- ranked = sorted(key, key=lambda x: boundary_order.index(x[0]))[:-1]
1714
- parent_key = tuple(sorted(ranked))
1715
-
1716
- # else choose a parent to compute
1717
- parent_tn = self._maybe_compute_cell_env(
1718
- key=parent_key,
1719
- envs=envs,
1720
- storage_factory=storage_factory,
1721
- boundary_order=boundary_order,
1722
- **contract_boundary_opts,
1723
- )
1724
-
1725
- # need to compute parent first - first set the parents range
1726
- Ls = {"x": self.Lx, "y": self.Ly, "z": self.Lz}
1727
- plane_tags = {"x": self.x_tag, "y": self.y_tag, "z": self.z_tag}
1728
- ranges = {"xrange": None, "yrange": None, "zrange": None}
1729
- for d, s, bsz in parent_key:
1730
- ranges[f"{d}range"] = (max(0, s - 1), min(s + bsz, Ls[d] - 1))
1731
-
1732
- # then compute the envs for the new direction ``d``
1733
- ((d, _, bsz),) = (x for x in key if x not in parent_key)
1734
-
1735
- envs_s_min = parent_tn._compute_plane_envs(
1736
- from_which=f"{d}min",
1737
- storage_factory=storage_factory,
1738
- **ranges,
1739
- **contract_boundary_opts,
1740
- )
1741
- envs_s_max = parent_tn._compute_plane_envs(
1742
- from_which=f"{d}max",
1743
- storage_factory=storage_factory,
1744
- **ranges,
1745
- **contract_boundary_opts,
1746
- )
1747
-
1748
- for s in range(0, Ls[d] - bsz + 1):
1749
- # the central non-boundary slice of tensors
1750
- tags_s = tuple(map(plane_tags[d], range(s, s + bsz)))
1751
- tn_s = parent_tn.select_any(tags_s, virtual=False)
1752
-
1753
- # the min boundary
1754
- if s in envs_s_min:
1755
- tn_s &= envs_s_min[s]
1756
- # the max boundary
1757
- imax = s + bsz - 1
1758
- if imax in envs_s_max:
1759
- tn_s &= envs_s_max[imax]
1760
-
1761
- # store the newly created cell along with env
1762
- key_s = tuple(sorted((*parent_key, (d, s, bsz))))
1763
- envs[key_s] = tn_s
1764
-
1765
- # one of key_s == key
1766
- return envs[key].copy()
1767
-
1768
- def coarse_grain_hotrg(
1769
- self,
1770
- direction,
1771
- max_bond=None,
1772
- cutoff=1e-10,
1773
- lazy=False,
1774
- equalize_norms=False,
1775
- optimize="auto-hq",
1776
- compress_opts=None,
1777
- inplace=False,
1778
- ):
1779
- """Coarse grain this tensor network in ``direction`` using HOTRG. This
1780
- inserts oblique projectors between tensor pairs and then optionally
1781
- contracts them into new sites for form a lattice half the size.
1782
-
1783
- Parameters
1784
- ----------
1785
- direction : {'x', 'y', 'z'}
1786
- The direction to coarse grain in.
1787
- max_bond : int, optional
1788
- The maximum bond dimension of the projector pairs inserted.
1789
- cutoff : float, optional
1790
- The cutoff for the singular values of the projector pairs.
1791
- lazy : bool, optional
1792
- Whether to contract the coarse graining projectors or leave them
1793
- in the tensor network lazily. Default is to contract them.
1794
- equalize_norms : bool, optional
1795
- Whether to equalize the norms of the tensors in the coarse grained
1796
- lattice.
1797
- optimize : str, optional
1798
- The optimization method to use when contracting the coarse grained
1799
- lattice, if ``lazy=False``.
1800
- compress_opts : None or dict, optional
1801
- Supplied to
1802
- :meth:`~quimb.tensor.tensor_core.TensorNetwork.insert_compressor_between_regions`.
1803
- inplace : bool, optional
1804
- Whether to perform the coarse graining in place.
1805
-
1806
- Returns
1807
- -------
1808
- TensorNetwork2D
1809
- The coarse grained tensor network, with size halved in
1810
- ``direction``.
1811
-
1812
- See Also
1813
- --------
1814
- contract_hotrg, insert_compressor_between_regions
1815
- """
1816
- compress_opts = ensure_dict(compress_opts)
1817
- check_opt("direction", direction, ("x", "y", "z"))
1818
-
1819
- tn = self if inplace else self.copy()
1820
- tn_calc = tn.copy()
1821
-
1822
- r = Rotator3D(tn, None, None, None, f"{direction}min")
1823
-
1824
- retag_map = {}
1825
-
1826
- for i in range(r.imin, r.imax + 1, 2):
1827
- # since we are partitioning into pairs in the `x` direction, we
1828
- # don't have to worry about PBC here - the boundary is never coarse
1829
- # grained over
1830
- next_i_in_lattice = i + 1 <= r.imax
1831
-
1832
- for j, k in r.sweep_other:
1833
- # however when computing neighboring sites in the orthogonal
1834
- # direction, we do need to consider PBC
1835
- jnext = r.get_jnext(j)
1836
- knext = r.get_knext(k)
1837
-
1838
- tag_ijk = r.site_tag(i, j, k)
1839
- tag_ip1jk = r.site_tag(i + 1, j, k)
1840
- new_tag = r.site_tag(i // 2, j, k)
1841
- retag_map[tag_ijk] = new_tag
1842
-
1843
- if next_i_in_lattice:
1844
- retag_map[tag_ip1jk] = new_tag
1845
-
1846
- rtags = (tag_ijk, tag_ip1jk)
1847
- poss_ltags = []
1848
- if jnext is not None:
1849
- poss_ltags.append(
1850
- (
1851
- r.site_tag(i, jnext, k),
1852
- r.site_tag(i + 1, jnext, k),
1853
- )
1854
- )
1855
- if knext is not None:
1856
- poss_ltags.append(
1857
- (
1858
- r.site_tag(i, j, knext),
1859
- r.site_tag(i + 1, j, knext),
1860
- )
1861
- )
1862
- for ltags in poss_ltags:
1863
- tn_calc.insert_compressor_between_regions(
1864
- ltags,
1865
- rtags,
1866
- new_ltags=ltags,
1867
- new_rtags=rtags,
1868
- insert_into=tn,
1869
- max_bond=max_bond,
1870
- cutoff=cutoff,
1871
- **compress_opts,
1872
- )
1873
-
1874
- retag_map[r.x_tag(i)] = r.x_tag(i // 2)
1875
- if next_i_in_lattice:
1876
- retag_map[r.x_tag(i + 1)] = r.x_tag(i // 2)
1877
-
1878
- # then we retag the tensor network and adjust its size
1879
- tn.retag_(retag_map)
1880
- if direction == "x":
1881
- tn._Lx = tn.Lx // 2 + tn.Lx % 2
1882
- elif direction == "y":
1883
- tn._Ly = tn.Ly // 2 + tn.Ly % 2
1884
- else: # direction == "z":
1885
- tn._Lz = tn.Lz // 2 + tn.Lz % 2
1886
-
1887
- # need this since we've fundamentally changed the geometry
1888
- tn.reset_cached_properties()
1889
-
1890
- if not lazy:
1891
- # contract each pair of tensors with their projectors
1892
- for st in tn.site_tags:
1893
- tn.contract_tags_(st, optimize=optimize)
1894
-
1895
- if equalize_norms:
1896
- tn.equalize_norms_(value=equalize_norms)
1897
-
1898
- return tn
1899
-
1900
- coarse_grain_hotrg_ = functools.partialmethod(
1901
- coarse_grain_hotrg, inplace=True
1902
- )
1903
-
1904
- def contract_hotrg(
1905
- self,
1906
- max_bond=None,
1907
- cutoff=1e-10,
1908
- canonize=False,
1909
- canonize_opts=None,
1910
- sequence=("x", "y", "z"),
1911
- max_separation=1,
1912
- max_unfinished=1,
1913
- lazy=False,
1914
- equalize_norms=False,
1915
- final_contract=True,
1916
- final_contract_opts=None,
1917
- progbar=False,
1918
- inplace=False,
1919
- **coarse_grain_opts,
1920
- ):
1921
- """Contract this tensor network using the finite version of HOTRG.
1922
- See https://arxiv.org/abs/1201.1144v4 and
1923
- https://arxiv.org/abs/1905.02351 for the more optimal computaton of the
1924
- projectors used here. The TN is contracted sequentially in
1925
- ``sequence`` directions by inserting oblique projectors between tensor
1926
- pairs, and then optionally contracting these new effective sites. The
1927
- algorithm stops when only one direction has a length larger than 2, and
1928
- thus exact contraction can be used.
1929
-
1930
- Parameters
1931
- ----------
1932
- max_bond : int, optional
1933
- The maximum bond dimension of the projector pairs inserted.
1934
- cutoff : float, optional
1935
- The cutoff for the singular values of the projector pairs.
1936
- sequence : tuple of str, optional
1937
- The directions to contract in. Default is to contract in all
1938
- directions.
1939
- max_separation : int, optional
1940
- The maximum distance between sides (i.e. length - 1) of the tensor
1941
- network before that direction is considered finished.
1942
- max_unfinished : int, optional
1943
- The maximum number of directions that can be unfinished (i.e. are
1944
- still longer than max_separation + 1) before the coarse graining
1945
- terminates.
1946
- lazy : bool, optional
1947
- Whether to contract the coarse graining projectors or leave them
1948
- in the tensor network lazily. Default is to contract them.
1949
- equalize_norms : bool or float, optional
1950
- Whether to equalize the norms of the tensors in the tensor network
1951
- after each coarse graining step.
1952
- final_contract : bool, optional
1953
- Whether to exactly contract the remaining tensor network after the
1954
- coarse graining contractions.
1955
- final_contract_opts : None or dict, optional
1956
- Options to pass to :meth:`contract`, ``optimize`` defaults to
1957
- ``'auto-hq'``.
1958
- inplace : bool, optional
1959
- Whether to perform the coarse graining in place.
1960
- coarse_grain_opts
1961
- Additional options to pass to :meth:`coarse_grain_hotrg`.
1962
-
1963
- Returns
1964
- -------
1965
- TensorNetwork3D
1966
- The contracted tensor network, which will have no more than one
1967
- directino of length > 2.
1968
-
1969
- See Also
1970
- --------
1971
- coarse_grain_hotrg, insert_compressor_between_regions
1972
- """
1973
- tn = self if inplace else self.copy()
1974
-
1975
- if lazy:
1976
- # we are implicitly asking for the tensor network
1977
- final_contract = False
1978
-
1979
- if canonize:
1980
- canonize_opts = ensure_dict(canonize_opts)
1981
- canonize_opts.setdefault("max_iterations", 2)
1982
-
1983
- if progbar:
1984
- pbar = Progbar(
1985
- desc=f"contracting HOTRG, Lx={tn.Lx}, Ly={tn.Ly}, Lz={tn.Lz}"
1986
- )
1987
- else:
1988
- pbar = None
1989
-
1990
- def _is_finished(direction):
1991
- return getattr(tn, "L" + direction) <= max_separation + 1
1992
-
1993
- sequence = [d for d in sequence if not _is_finished(d)]
1994
- while sequence:
1995
- direction = sequence.pop(0)
1996
- if _is_finished(direction):
1997
- # just remove direction from sequence
1998
- continue
1999
- # do a contraction, and keep direction in sequence to try again
2000
- sequence.append(direction)
2001
-
2002
- if pbar is not None:
2003
- pbar.set_description(
2004
- f"contracting {direction}, "
2005
- f"Lx={tn.Lx}, Ly={tn.Ly}, Lz={tn.Lz}"
2006
- )
2007
-
2008
- if canonize:
2009
- tn.gauge_all_(**canonize_opts)
2010
-
2011
- # do the contractions!
2012
- tn.coarse_grain_hotrg_(
2013
- direction=direction,
2014
- max_bond=max_bond,
2015
- cutoff=cutoff,
2016
- lazy=lazy,
2017
- equalize_norms=equalize_norms,
2018
- **coarse_grain_opts,
2019
- )
2020
-
2021
- if pbar is not None:
2022
- pbar.update()
2023
-
2024
- # check if enough directions are finished -> reached max separation
2025
- if sum(not _is_finished(d) for d in "xyz") <= max_unfinished:
2026
- break
2027
-
2028
- if equalize_norms is True:
2029
- # redistribute the exponent equally among all tensors
2030
- tn.equalize_norms_()
2031
-
2032
- if pbar is not None:
2033
- pbar.set_description(
2034
- f"contracted HOTRG, " f"Lx={tn.Lx}, Ly={tn.Ly}, Lz={tn.Lz}"
2035
- )
2036
- pbar.close()
2037
-
2038
- if final_contract:
2039
- final_contract_opts = ensure_dict(final_contract_opts)
2040
- final_contract_opts.setdefault("optimize", "auto-hq")
2041
- final_contract_opts.setdefault("inplace", inplace)
2042
- return tn.contract(**final_contract_opts)
2043
-
2044
- return tn
2045
-
2046
- contract_hotrg_ = functools.partialmethod(contract_hotrg, inplace=True)
2047
-
2048
-
2049
- def is_lone_coo(where):
2050
- """Check if ``where`` has been specified as a single coordinate triplet."""
2051
- return (len(where) == 3) and (isinstance(where[0], Integral))
2052
-
2053
-
2054
- def calc_cell_sizes(coo_groups, autogroup=True):
2055
- # get the rectangular size of each coordinate pair
2056
- bszs = set()
2057
- for coos in coo_groups:
2058
- if is_lone_coo(coos):
2059
- bszs.add((1, 1, 1))
2060
- continue
2061
- xs, ys, zs = zip(*coos)
2062
- xsz = max(xs) - min(xs) + 1
2063
- ysz = max(ys) - min(ys) + 1
2064
- zsz = max(zs) - min(zs) + 1
2065
- bszs.add((xsz, ysz, zsz))
2066
-
2067
- # remove block size pairs that can be contained in another block pair size
2068
- bszs = tuple(
2069
- sorted(
2070
- b
2071
- for b in bszs
2072
- if not any(
2073
- (b[0] <= b2[0]) and (b[1] <= b2[1]) and (b[2] <= b2[2])
2074
- for b2 in bszs - {b}
2075
- )
2076
- )
2077
- )
2078
-
2079
- # return each cell size separately
2080
- if autogroup:
2081
- return bszs
2082
-
2083
- # else choose a single blocksize that will cover all terms
2084
- return (tuple(map(max, zip(*bszs))),)
2085
-
2086
-
2087
- def cell_to_sites(p):
2088
- """Turn a cell ``((i0, j0, k0), (di, dj, dk))`` into the sites it contains.
2089
-
2090
- Examples
2091
- --------
2092
-
2093
- >>> cell_to_sites([(3, 4), (2, 2)])
2094
- ((3, 4), (3, 5), (4, 4), (4, 5))
2095
- """
2096
- (i0, j0, k0), (di, dj, dk) = p
2097
- return tuple(
2098
- (i, j, k)
2099
- for i in range(i0, i0 + di)
2100
- for j in range(j0, j0 + dj)
2101
- for k in range(k0, k0 + dk)
2102
- )
2103
-
2104
-
2105
- def sites_to_cell(sites):
2106
- """Get the minimum covering cell for ``sites``.
2107
-
2108
- Examples
2109
- --------
2110
-
2111
- >>> sites_to_cell([(1, 3, 3), (2, 2, 4)])
2112
- ((1, 2, 3) , (2, 2, 2))
2113
- """
2114
- imin = jmin = kmin = float("inf")
2115
- imax = jmax = kmax = float("-inf")
2116
- for i, j, k in sites:
2117
- imin, jmin, kmin = min(imin, i), min(jmin, j), min(kmin, k)
2118
- imax, jmax, kmax = max(imax, i), max(jmax, j), max(kmax, k)
2119
- x_bsz, y_bsz, z_bsz = imax - imin + 1, jmax - jmin + 1, kmax - kmin + 1
2120
- return (imin, jmin, kmin), (x_bsz, y_bsz, z_bsz)
2121
-
2122
-
2123
- def calc_cell_map(cells):
2124
- # sort in descending total cell size
2125
- cs = sorted(cells, key=lambda c: (-c[1][0] * c[1][1] * c[1][2], c))
2126
-
2127
- mapping = dict()
2128
- for c in cs:
2129
- sites = cell_to_sites(c)
2130
- for site in sites:
2131
- mapping[site] = c
2132
- # this will generate all coordinate pairs with ijk_a < ijk_b
2133
- for ijk_a, ijk_b in combinations(sites, 2):
2134
- mapping[ijk_a, ijk_b] = c
2135
-
2136
- return mapping
2137
-
2138
-
2139
- class TensorNetwork3DVector(TensorNetwork3D, TensorNetworkGenVector):
2140
- """Mixin class for a 3D square lattice vector TN, i.e. one with a single
2141
- physical index per site.
2142
- """
2143
-
2144
- _EXTRA_PROPS = (
2145
- "_site_tag_id",
2146
- "_x_tag_id",
2147
- "_y_tag_id",
2148
- "_z_tag_id",
2149
- "_Lx",
2150
- "_Ly",
2151
- "_Lz",
2152
- "_site_ind_id",
2153
- )
2154
-
2155
- def site_ind(self, i, j=None, k=None):
2156
- if j is None:
2157
- i, j, k = i
2158
- if not isinstance(i, str):
2159
- i = i % self.Lx
2160
- if not isinstance(j, str):
2161
- j = j % self.Ly
2162
- if not isinstance(k, str):
2163
- k = k % self.Lz
2164
- return self.site_ind_id.format(i, j, k)
2165
-
2166
- def reindex_sites(self, new_id, where=None, inplace=False):
2167
- if where is None:
2168
- where = self.gen_sites_present()
2169
-
2170
- return self.reindex(
2171
- {self.site_ind(*coo): new_id.format(*coo) for coo in where},
2172
- inplace=inplace,
2173
- )
2174
-
2175
- def phys_dim(self, i=None, j=None, k=None):
2176
- """Get the size of the physical indices / a specific physical index."""
2177
- if (i is not None) and (j is not None) and (k is not None):
2178
- pix = self.site_ind(i, j, k)
2179
- else:
2180
- # allow for when some physical indices might have been contracted
2181
- pix = next(iter(ix for ix in self.site_inds if ix in self.ind_map))
2182
- return self.ind_size(pix)
2183
-
2184
- def gate(
2185
- self,
2186
- G,
2187
- where,
2188
- contract=False,
2189
- tags=None,
2190
- info=None,
2191
- inplace=False,
2192
- **compress_opts,
2193
- ):
2194
- """Apply a gate ``G`` to sites ``where``, preserving the outer site
2195
- inds.
2196
- """
2197
- if is_lone_coo(where):
2198
- where = (where,)
2199
- else:
2200
- where = tuple(where)
2201
-
2202
- inds = tuple(map(self.site_ind, where))
2203
- return super().gate_inds(
2204
- G,
2205
- inds,
2206
- contract=contract,
2207
- tags=tags,
2208
- info=info,
2209
- inplace=inplace,
2210
- **compress_opts,
2211
- )
2212
-
2213
- gate_ = functools.partialmethod(gate, inplace=True)
2214
-
2215
-
2216
- class TensorNetwork3DFlat(TensorNetwork3D):
2217
- """Mixin class for a 3D square lattice tensor network with a single tensor
2218
- per site, for example, both PEPS and PEPOs.
2219
- """
2220
-
2221
- _EXTRA_PROPS = (
2222
- "_site_tag_id",
2223
- "_x_tag_id",
2224
- "_y_tag_id",
2225
- "_z_tag_id",
2226
- "_Lx",
2227
- "_Ly",
2228
- "_Lz",
2229
- )
2230
-
2231
- def bond(self, coo1, coo2):
2232
- """Get the name of the index defining the bond between sites at
2233
- ``coo1`` and ``coo2``.
2234
- """
2235
- (b_ix,) = self[coo1].bonds(self[coo2])
2236
- return b_ix
2237
-
2238
- def bond_size(self, coo1, coo2):
2239
- """Return the size of the bond between sites at ``coo1`` and ``coo2``."""
2240
- b_ix = self.bond(coo1, coo2)
2241
- return self[coo1].ind_size(b_ix)
2242
-
2243
-
2244
- class PEPS3D(TensorNetwork3DVector, TensorNetwork3DFlat):
2245
- r"""Projected Entangled Pair States object (3D).
2246
-
2247
- Parameters
2248
- ----------
2249
- arrays : sequence of sequence of sequence of array
2250
- The core tensor data arrays.
2251
- shape : str, optional
2252
- Which order the dimensions of the arrays are stored in, the default
2253
- ``'urfdlbp'`` stands for ('up', 'right', 'front', down', 'left',
2254
- 'behind', 'physical') meaning (x+, y+, z+, x-, y-, z-, physical)
2255
- respectively. Arrays on the edge of lattice are assumed to be missing
2256
- the corresponding dimension.
2257
- tags : set[str], optional
2258
- Extra global tags to add to the tensor network.
2259
- site_ind_id : str, optional
2260
- String specifier for naming convention of site indices.
2261
- site_tag_id : str, optional
2262
- String specifier for naming convention of site tags.
2263
- x_tag_id : str, optional
2264
- String specifier for naming convention of x-slice tags.
2265
- y_tag_id : str, optional
2266
- String specifier for naming convention of y-slice tags.
2267
- z_tag_id : str, optional
2268
- String specifier for naming convention of z-slice tags.
2269
- """
2270
-
2271
- _EXTRA_PROPS = (
2272
- "_site_tag_id",
2273
- "_x_tag_id",
2274
- "_y_tag_id",
2275
- "_z_tag_id",
2276
- "_Lx",
2277
- "_Ly",
2278
- "_Lz",
2279
- "_site_ind_id",
2280
- )
2281
-
2282
- def __init__(
2283
- self,
2284
- arrays,
2285
- *,
2286
- shape="urfdlbp",
2287
- tags=None,
2288
- site_ind_id="k{},{},{}",
2289
- site_tag_id="I{},{},{}",
2290
- x_tag_id="X{}",
2291
- y_tag_id="Y{}",
2292
- z_tag_id="Z{}",
2293
- **tn_opts,
2294
- ):
2295
- if isinstance(arrays, PEPS3D):
2296
- super().__init__(arrays)
2297
- return
2298
-
2299
- tags = tags_to_oset(tags)
2300
- self._site_ind_id = site_ind_id
2301
- self._site_tag_id = site_tag_id
2302
- self._x_tag_id = x_tag_id
2303
- self._y_tag_id = y_tag_id
2304
- self._z_tag_id = z_tag_id
2305
-
2306
- arrays = tuple(tuple(tuple(z for z in y) for y in x) for x in arrays)
2307
- self._Lx = len(arrays)
2308
- self._Ly = len(arrays[0])
2309
- self._Lz = len(arrays[0][0])
2310
-
2311
- cyclicx = sum(d > 1 for d in arrays[0][1][1].shape) == 7
2312
- cyclicy = sum(d > 1 for d in arrays[1][0][1].shape) == 7
2313
- cyclicz = sum(d > 1 for d in arrays[0][1][0].shape) == 7
2314
-
2315
- # cache for both creating and retrieving indices
2316
- ix = defaultdict(rand_uuid)
2317
- tensors = []
2318
-
2319
- for i, j, k in self.gen_site_coos():
2320
- array = arrays[i][j][k]
2321
-
2322
- # figure out if we need to transpose the arrays from some order
2323
- # other than up right front down left behind physical
2324
- array_order = shape
2325
- if (not cyclicx) and (i == self.Lx - 1):
2326
- array_order = array_order.replace("u", "")
2327
- if (not cyclicy) and (j == self.Ly - 1):
2328
- array_order = array_order.replace("r", "")
2329
- if (not cyclicz) and (k == self.Lz - 1):
2330
- array_order = array_order.replace("f", "")
2331
- if (not cyclicx) and (i == 0):
2332
- array_order = array_order.replace("d", "")
2333
- if (not cyclicy) and (j == 0):
2334
- array_order = array_order.replace("l", "")
2335
- if (not cyclicz) and (k == 0):
2336
- array_order = array_order.replace("b", "")
2337
-
2338
- # allow convention of missing bonds to be singlet dimensions
2339
- if len(array.shape) != len(array_order):
2340
- array = do("squeeze", array)
2341
-
2342
- transpose_order = tuple(
2343
- array_order.find(x) for x in "urfdlbp" if x in array_order
2344
- )
2345
- if transpose_order != tuple(range(len(array_order))):
2346
- array = do("transpose", array, transpose_order)
2347
-
2348
- # get the relevant indices corresponding to neighbours
2349
- inds = []
2350
- if "u" in array_order:
2351
- i_u = (i + 1) % self.Lx
2352
- inds.append(ix[frozenset(((i, j, k), (i_u, j, k)))])
2353
- if "r" in array_order:
2354
- j_r = (j + 1) % self.Ly
2355
- inds.append(ix[frozenset(((i, j, k), (i, j_r, k)))])
2356
- if "f" in array_order:
2357
- k_f = (k + 1) % self.Lz
2358
- inds.append(ix[frozenset(((i, j, k), (i, j, k_f)))])
2359
- if "d" in array_order:
2360
- i_d = (i - 1) % self.Lx
2361
- inds.append(ix[frozenset(((i_d, j, k), (i, j, k)))])
2362
- if "l" in array_order:
2363
- j_l = (j - 1) % self.Ly
2364
- inds.append(ix[frozenset(((i, j_l, k), (i, j, k)))])
2365
- if "b" in array_order:
2366
- k_b = (k - 1) % self.Lz
2367
- inds.append(ix[frozenset(((i, j, k_b), (i, j, k)))])
2368
- inds.append(self.site_ind(i, j, k))
2369
-
2370
- # mix site, slice and global tags
2371
- ijk_tags = tags | oset(
2372
- (
2373
- self.site_tag(i, j, k),
2374
- self.x_tag(i),
2375
- self.y_tag(j),
2376
- self.z_tag(k),
2377
- )
2378
- )
2379
-
2380
- # create the site tensor!
2381
- tensors.append(Tensor(data=array, inds=inds, tags=ijk_tags))
2382
-
2383
- super().__init__(tensors, virtual=True, **tn_opts)
2384
-
2385
- def permute_arrays(self, shape="urfdlbp"):
2386
- """Permute the indices of each tensor in this PEPS3D to match
2387
- ``shape``. This doesn't change how the overall object interacts with
2388
- other tensor networks but may be useful for extracting the underlying
2389
- arrays consistently. This is an inplace operation.
2390
-
2391
- Parameters
2392
- ----------
2393
- shape : str, optional
2394
- A permutation of ``'lrp'`` specifying the desired order of the
2395
- left, right, and physical indices respectively.
2396
- """
2397
- steps = {
2398
- "u": lambda i, j, k: (i + 1, j, k),
2399
- "r": lambda i, j, k: (i, j + 1, k),
2400
- "f": lambda i, j, k: (i, j, k + 1),
2401
- "d": lambda i, j, k: (i - 1, j, k),
2402
- "l": lambda i, j, k: (i, j - 1, k),
2403
- "b": lambda i, j, k: (i, j, k - 1),
2404
- }
2405
- for i, j, k in self.gen_site_coos():
2406
- t = self[i, j, k]
2407
- inds = []
2408
- for s in shape:
2409
- if s == "p":
2410
- inds.append(self.site_ind(i, j, k))
2411
- else:
2412
- coo2 = steps[s](i, j, k)
2413
- if self.valid_coo(coo2):
2414
- t2 = self[coo2]
2415
- (bix,) = t.bonds(t2)
2416
- inds.append(bix)
2417
- t.transpose_(*inds)
2418
-
2419
- @classmethod
2420
- def from_fill_fn(
2421
- cls, fill_fn, Lx, Ly, Lz, bond_dim, phys_dim=2,
2422
- cyclic=False,
2423
- shape="urfdlbp",
2424
- **peps3d_opts
2425
- ):
2426
- """Create a 3D PEPS from a filling function with signature
2427
- ``fill_fn(shape)``.
2428
-
2429
- Parameters
2430
- ----------
2431
- Lx : int
2432
- The number of x-slices.
2433
- Ly : int
2434
- The number of y-slices.
2435
- Lz : int
2436
- The number of z-slices.
2437
- bond_dim : int
2438
- The bond dimension.
2439
- phys_dim : int, optional
2440
- The physical index dimension.
2441
- shape : str, optional
2442
- How to layout the indices of the tensors, the default is
2443
- ``(up, right, front, down, left, back, phys) == 'urfdlbp'``. This
2444
- is the order of the shape supplied to the filling function.
2445
- peps_opts
2446
- Supplied to :class:`~quimb.tensor.tensor_3d.PEPS3D`.
2447
-
2448
- Returns
2449
- -------
2450
- psi : PEPS3D
2451
- """
2452
- arrays = [
2453
- [[None for _ in range(Lz)] for _ in range(Ly)] for _ in range(Lx)
2454
- ]
2455
-
2456
- try:
2457
- cyclicx, cyclicy, cyclicz = cyclic
2458
- except (TypeError, ValueError):
2459
- cyclicx = cyclicy = cyclicz = cyclic
2460
-
2461
- for i, j, k in product(range(Lx), range(Ly), range(Lz)):
2462
- shp = []
2463
-
2464
- for which in shape:
2465
- if (which == "u") and (cyclicx or (i != Lx - 1)): # up
2466
- shp.append(bond_dim)
2467
- elif (which == "r") and (cyclicy or (j != Ly - 1)): # right
2468
- shp.append(bond_dim)
2469
- elif (which == "f") and (cyclicz or (k != Lz - 1)): # front
2470
- shp.append(bond_dim)
2471
- elif (which == "d") and (cyclicx or (i != 0)): # down
2472
- shp.append(bond_dim)
2473
- elif (which == "l") and (cyclicy or (j != 0)): # left
2474
- shp.append(bond_dim)
2475
- elif (which == "b") and (cyclicz or (k != 0)): # back
2476
- shp.append(bond_dim)
2477
-
2478
- shp.append(phys_dim)
2479
-
2480
- arrays[i][j][k] = fill_fn(shp)
2481
-
2482
- return cls(arrays, **peps3d_opts)
2483
-
2484
- @classmethod
2485
- def empty(
2486
- self,
2487
- Lx,
2488
- Ly,
2489
- Lz,
2490
- bond_dim,
2491
- phys_dim=2,
2492
- like="numpy",
2493
- **peps3d_opts,
2494
- ):
2495
- """Create an empty 3D PEPS.
2496
-
2497
- Parameters
2498
- ----------
2499
- Lx : int
2500
- The number of x-slices.
2501
- Ly : int
2502
- The number of y-slices.
2503
- Lz : int
2504
- The number of z-slices.
2505
- bond_dim : int
2506
- The bond dimension.
2507
- physical : int, optional
2508
- The physical index dimension.
2509
- peps3d_opts
2510
- Supplied to :class:`~quimb.tensor.tensor_3d.PEPS3D`.
2511
-
2512
- Returns
2513
- -------
2514
- psi : PEPS3D
2515
-
2516
- See Also
2517
- --------
2518
- PEPS3D.from_fill_fn
2519
- """
2520
- return self.from_fill_fn(
2521
- lambda shape: do("zeros", shape, like=like),
2522
- Lx,
2523
- Ly,
2524
- Lz,
2525
- bond_dim,
2526
- phys_dim,
2527
- **peps3d_opts,
2528
- )
2529
-
2530
- @classmethod
2531
- def ones(
2532
- self,
2533
- Lx,
2534
- Ly,
2535
- Lz,
2536
- bond_dim,
2537
- phys_dim=2,
2538
- like="numpy",
2539
- **peps3d_opts,
2540
- ):
2541
- """Create a 3D PEPS whose tensors are filled with ones.
2542
-
2543
- Parameters
2544
- ----------
2545
- Lx : int
2546
- The number of x-slices.
2547
- Ly : int
2548
- The number of y-slices.
2549
- Lz : int
2550
- The number of z-slices.
2551
- bond_dim : int
2552
- The bond dimension.
2553
- physical : int, optional
2554
- The physical index dimension.
2555
- peps3d_opts
2556
- Supplied to :class:`~quimb.tensor.tensor_3d.PEPS3D`.
2557
-
2558
- Returns
2559
- -------
2560
- psi : PEPS3D
2561
-
2562
- See Also
2563
- --------
2564
- PEPS3D.from_fill_fn
2565
- """
2566
- return self.from_fill_fn(
2567
- lambda shape: do("ones", shape, like=like),
2568
- Lx,
2569
- Ly,
2570
- Lz,
2571
- bond_dim,
2572
- phys_dim,
2573
- **peps3d_opts,
2574
- )
2575
-
2576
- @classmethod
2577
- def rand(
2578
- cls,
2579
- Lx,
2580
- Ly,
2581
- Lz,
2582
- bond_dim,
2583
- phys_dim=2,
2584
- dist="normal",
2585
- loc=0.0,
2586
- dtype="float64",
2587
- seed=None,
2588
- **peps3d_opts,
2589
- ):
2590
- """Create a random (un-normalized) 3D PEPS.
2591
-
2592
- Parameters
2593
- ----------
2594
- Lx : int
2595
- The number of x-slices.
2596
- Ly : int
2597
- The number of y-slices.
2598
- Lz : int
2599
- The number of z-slices.
2600
- bond_dim : int
2601
- The bond dimension.
2602
- physical : int, optional
2603
- The physical index dimension.
2604
- dtype : dtype, optional
2605
- The dtype to create the arrays with, default is real double.
2606
- seed : int, optional
2607
- A random seed.
2608
- peps_opts
2609
- Supplied to :class:`~quimb.tensor.tensor_3d.PEPS3D`.
2610
-
2611
- Returns
2612
- -------
2613
- psi : PEPS3D
2614
-
2615
- See Also
2616
- --------
2617
- PEPS3D.from_fill_fn
2618
- """
2619
- if seed is not None:
2620
- seed_rand(seed)
2621
-
2622
- def fill_fn(shape):
2623
- return ops.sensibly_scale(
2624
- ops.sensibly_scale(
2625
- randn(shape, dist=dist, loc=loc, dtype=dtype)
2626
- )
2627
- )
2628
-
2629
- return cls.from_fill_fn(
2630
- fill_fn, Lx, Ly, Lz, bond_dim, phys_dim, **peps3d_opts
2631
- )
2632
-
2633
- def partial_trace_cluster(
2634
- self,
2635
- keep,
2636
- max_bond=None,
2637
- *,
2638
- cutoff=1e-10,
2639
- max_distance=0,
2640
- fillin=0,
2641
- gauges=False,
2642
- flatten=False,
2643
- normalized=True,
2644
- symmetrized="auto",
2645
- get=None,
2646
- **contract_boundary_opts,
2647
- ):
2648
- if is_lone_coo(keep):
2649
- keep = (keep,)
2650
-
2651
- tags = [self.site_tag(i, j, k) for i, j, k in keep]
2652
-
2653
- k = self.select_local(
2654
- tags,
2655
- "any",
2656
- max_distance=max_distance,
2657
- fillin=fillin,
2658
- virtual=False,
2659
- )
2660
-
2661
- k.add_tag("KET")
2662
- if gauges:
2663
- k.gauge_simple_insert(gauges)
2664
-
2665
- kix = [self.site_ind(i, j, k) for i, j, k in keep]
2666
- bix = [rand_uuid() for _ in kix]
2667
-
2668
- b = k.H.reindex_(dict(zip(kix, bix))).retag_({"KET": "BRA"})
2669
- rho_tn = k | b
2670
- rho_tn.fuse_multibonds_()
2671
-
2672
- if get == "tn":
2673
- return rho_tn
2674
-
2675
- # contract boundaries largest dimension last
2676
- ri, rj, rk = zip(*keep)
2677
- imin, imax = min(ri), max(ri)
2678
- jmin, jmax = min(rj), max(rj)
2679
- kmin, kmax = min(rk), max(rk)
2680
- sequence = sorted(
2681
- [
2682
- ((imax - imin), ("xmin", "xmax")),
2683
- ((jmax - jmin), ("ymin", "ymax")),
2684
- ((kmax - kmin), ("zmin", "zmax")),
2685
- ]
2686
- )
2687
- sequence = [x for s in sequence for x in s[1]]
2688
-
2689
- rho_t = rho_tn.contract_boundary(
2690
- max_bond=max_bond,
2691
- cutoff=cutoff,
2692
- sequence=sequence,
2693
- layer_tags=None if flatten else ("KET", "BRA"),
2694
- **contract_boundary_opts,
2695
- )
2696
-
2697
- if symmetrized == "auto":
2698
- symmetrized = not flatten
2699
-
2700
- rho = rho_t.to_dense(kix, bix)
2701
-
2702
- # maybe fix up
2703
- if symmetrized:
2704
- rho = (rho + dag(rho)) / 2
2705
- if normalized:
2706
- rho = rho / do("trace", rho)
2707
-
2708
- return rho
2709
-
2710
- def partial_trace(
2711
- self,
2712
- keep,
2713
- max_bond=None,
2714
- *,
2715
- cutoff=1e-10,
2716
- canonize=True,
2717
- flatten=False,
2718
- normalized=True,
2719
- symmetrized="auto",
2720
- envs=None,
2721
- storage_factory=None,
2722
- boundary_order=None,
2723
- contract_cell_optimize="auto-hq",
2724
- contract_cell_method="boundary",
2725
- contract_cell_opts=None,
2726
- get=None,
2727
- **contract_boundary_opts,
2728
- ):
2729
- contract_cell_opts = ensure_dict(contract_cell_opts)
2730
-
2731
- norm = self.make_norm()
2732
-
2733
- contract_boundary_opts["max_bond"] = max_bond
2734
- contract_boundary_opts["cutoff"] = cutoff
2735
- contract_boundary_opts["canonize"] = canonize
2736
- contract_boundary_opts["layer_tags"] = (
2737
- None if flatten else ("KET", "BRA")
2738
- )
2739
- if symmetrized == "auto":
2740
- symmetrized = not flatten
2741
-
2742
- # get minimal covering cell, allow single coordinate
2743
- if is_lone_coo(keep):
2744
- keep = (keep,)
2745
- cell = sites_to_cell(keep)
2746
-
2747
- # get the environment surrounding the cell, allowing reuse via ``envs``
2748
- (i, j, k), (x_bsz, y_bsz, z_bsz) = cell
2749
- key = (("x", i, x_bsz), ("y", j, y_bsz), ("z", k, z_bsz))
2750
- tn_cell = norm._maybe_compute_cell_env(
2751
- key=key,
2752
- envs=envs,
2753
- storage_factory=storage_factory,
2754
- boundary_order=boundary_order,
2755
- **contract_boundary_opts,
2756
- )
2757
-
2758
- # cut the bonds between target norm sites to make density matrix
2759
- tags = [tn_cell.site_tag(*site) for site in keep]
2760
- kix = [f"k{i},{j},{k}" for i, j, k in keep]
2761
- bix = [f"b{i},{j},{k}" for i, j, k in keep]
2762
- for tag, ind_k, ind_b in zip(tags, kix, bix):
2763
- tn_cell.cut_between((tag, "KET"), (tag, "BRA"), ind_k, ind_b)
2764
-
2765
- if get == "tn":
2766
- return tn_cell
2767
-
2768
- if contract_cell_method == "boundary":
2769
- # perform the contract to single tensor as boundary contraction
2770
- # -> still likely far too expensive to contract exactly
2771
- xmin, xmax = max(0, i - 1), min(i + x_bsz, self.Lx - 1)
2772
- ymin, ymax = max(0, j - 1), min(j + y_bsz, self.Ly - 1)
2773
- zmin, zmax = max(0, k - 1), min(k + z_bsz, self.Lz - 1)
2774
-
2775
- sequence = []
2776
- if i > 0:
2777
- sequence.append("xmin")
2778
- if i < self.Lx - 1:
2779
- sequence.append("xmax")
2780
- if j > 0:
2781
- sequence.append("ymin")
2782
- if j < self.Ly - 1:
2783
- sequence.append("ymax")
2784
- if k > 0:
2785
- sequence.append("zmin")
2786
- if k < self.Lz - 1:
2787
- sequence.append("zmax")
2788
-
2789
- # contract longest boundary first
2790
- scores = {"x": xmax - xmin, "y": ymax - ymin, "z": zmax - zmin}
2791
- sequence.sort(key=lambda s: scores[s[0]], reverse=True)
2792
-
2793
- rho_tn = tn_cell.contract_boundary_(
2794
- xmin=xmin,
2795
- xmax=xmax,
2796
- ymin=ymin,
2797
- ymax=ymax,
2798
- zmin=zmin,
2799
- zmax=zmax,
2800
- sequence=sequence,
2801
- optimize=contract_cell_optimize,
2802
- **contract_boundary_opts,
2803
- )
2804
- else:
2805
- contract_cell_opts.setdefault("optimize", contract_cell_optimize)
2806
- contract_cell_opts.setdefault("max_bond", max_bond)
2807
- contract_cell_opts.setdefault("cutoff", cutoff)
2808
- rho_tn = tn_cell.contract_compressed_(
2809
- output_inds=kix + bix, **contract_cell_opts
2810
- )
2811
-
2812
- # turn into raw array
2813
- rho = rho_tn.to_dense(kix, bix)
2814
-
2815
- # maybe fix up
2816
- if symmetrized:
2817
- rho = (rho + dag(rho)) / 2
2818
- if normalized:
2819
- rho = rho / do("trace", rho)
2820
-
2821
- return rho
2822
-
2823
- def compute_local_expectation(
2824
- self,
2825
- terms,
2826
- max_bond=None,
2827
- *,
2828
- cutoff=1e-10,
2829
- canonize=True,
2830
- flatten=False,
2831
- normalized=True,
2832
- symmetrized="auto",
2833
- return_all=False,
2834
- envs=None,
2835
- storage_factory=None,
2836
- progbar=False,
2837
- **contract_boundary_opts,
2838
- ):
2839
- if envs is None:
2840
- if storage_factory is not None:
2841
- envs = storage_factory()
2842
- else:
2843
- envs = {}
2844
-
2845
- if progbar:
2846
- items = Progbar(terms.items())
2847
- else:
2848
- items = terms.items()
2849
-
2850
- expecs = dict()
2851
- for where, G in items:
2852
- rho = self.partial_trace(
2853
- where,
2854
- max_bond=max_bond,
2855
- cutoff=cutoff,
2856
- canonize=canonize,
2857
- flatten=flatten,
2858
- symmetrized=symmetrized,
2859
- normalized=normalized,
2860
- envs=envs,
2861
- storage_factory=storage_factory,
2862
- **contract_boundary_opts,
2863
- )
2864
- expecs[where] = do("tensordot", G, rho, ((0, 1), (1, 0)))
2865
-
2866
- if return_all:
2867
- return expecs
2868
-
2869
- return functools.reduce(add, expecs.values())