Source code for quimb.tensor.tensor_3d

import functools
import itertools
from operator import add
from numbers import Integral
from collections import defaultdict
from itertools import product, starmap, cycle, combinations

from autoray import do, dag

from ..utils import check_opt, ensure_dict
from ..utils import progbar as Progbar
from ..gen.rand import randn, seed_rand
from . import array_ops as ops
from .tensor_core import (
    Tensor,
    TensorNetwork,
    bonds_size,
    oset,
    tags_to_oset,
    rand_uuid,
)


[docs]def gen_3d_bonds(Lx, Ly, Lz, steppers, coo_filter=None): """Convenience function for tiling pairs of bond coordinates on a 3D lattice given a function like ``lambda i, j, k: (i + 1, j + 1, k + 1)``. Parameters ---------- Lx : int The number of x-slices. Ly : int The number of y-slices. Lz : int The number of z-slices. steppers : callable or sequence of callable Function(s) that take args ``(i, j, k)`` and generate another coordinate, thus defining a bond. coo_filter : callable Function that takes args ``(i, j, k)`` and only returns ``True`` if this is to be a valid starting coordinate. Yields ------ bond : tuple[tuple[int, int, int], tuple[int, int, int]] A pair of coordinates. Examples -------- Generate nearest neighbor bonds: >>> for bond in gen_3d_bonds(2, 2, 2, [lambda i, j, k: (i + 1, j, k), ... lambda i, j, k: (i, j + 1, k), ... lambda i, j, k: (i, j, k + 1)]): ... print(bond) ((0, 0, 0), (1, 0, 0)) ((0, 0, 0), (0, 1, 0)) ((0, 0, 0), (0, 0, 1)) ((0, 0, 1), (1, 0, 1)) ((0, 0, 1), (0, 1, 1)) ((0, 1, 0), (1, 1, 0)) ((0, 1, 0), (0, 1, 1)) ((0, 1, 1), (1, 1, 1)) ((1, 0, 0), (1, 1, 0)) ((1, 0, 0), (1, 0, 1)) ((1, 0, 1), (1, 1, 1)) ((1, 1, 0), (1, 1, 1)) """ if callable(steppers): steppers = (steppers,) for i, j, k in product(range(Lx), range(Ly), range(Lz)): if (coo_filter is None) or coo_filter(i, j, k): for stepper in steppers: i2, j2, k2 = stepper(i, j, k) if (0 <= i2 < Lx) and (0 <= j2 < Ly) and (0 <= k2 < Lz): yield (i, j, k), (i2, j2, k2)
[docs]class Rotator3D: """Object for rotating coordinates and various contraction functions so that the core algorithms only have to written once, but nor does the actual TN have to be modified. """ def __init__(self, tn, xrange, yrange, zrange, from_which): check_opt('from_which', from_which, {'xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax'}) if xrange is None: xrange = (0, tn.Lx - 1) if yrange is None: yrange = (0, tn.Ly - 1) if zrange is None: zrange = (0, tn.Lz - 1) self.xrange = xrange self.yrange = yrange self.zrange = zrange self.from_which = from_which self.plane = from_which[0] if self.plane == 'x': # -> no rotation needed self.imin, self.imax = sorted(xrange) self.jmin, self.jmax = sorted(yrange) self.kmin, self.kmax = sorted(zrange) self.x_tag = tn.x_tag self.y_tag = tn.y_tag self.z_tag = tn.z_tag self.site_tag = tn.site_tag elif self.plane == 'y': # -> (y, z, x) self.imin, self.imax = sorted(yrange) self.jmin, self.jmax = sorted(zrange) self.kmin, self.kmax = sorted(xrange) self.x_tag = tn.y_tag self.y_tag = tn.z_tag self.z_tag = tn.x_tag self.site_tag = lambda i, j, k: tn.site_tag(k, i, j) else: # self.plane == 'z' # -> (z, x, y) self.imin, self.imax = sorted(zrange) self.jmin, self.jmax = sorted(xrange) self.kmin, self.kmax = sorted(yrange) self.x_tag = tn.z_tag self.y_tag = tn.x_tag self.z_tag = tn.y_tag self.site_tag = lambda i, j, k: tn.site_tag(j, k, i) if 'min' in from_which: # -> sweeps are increasing self.sweep = range(self.imin, self.imax + 1, + 1) self.istep = +1 else: # 'max' # -> sweeps are decreasing self.sweep = range(self.imax, self.imin - 1, -1) self.istep = -1
# reference for viewing a cube from each direction # # ┌──┐ ┌──┐ ┌──┐ ┌──┐ ┌──┐ ┌──┐ # │y+│ │z+│ │x-│ │y-│ │z-│ │x+│ # ┌──┼──┼──┐ ┌──┼──┼──┐ ┌──┼──┼──┐ ┌──┼──┼──┐ ┌──┼──┼──┐ ┌──┼──┼──┐ # │z-│x-│z+│ │x-│y-│x+│ │y+│z-│y-│ │z-│x+│z+│ │x-│y+│x+│ │y+│z+│y-│ # └──┼──┼──┘, └──┼──┼──┘, └──┼──┼──┘, └──┼──┼──┘, └──┼──┼──┘, └──┼──┼──┘ # │y-│ │z-│ │x+│ │y+│ │z+│ │x-│ # └──┘ └──┘ └──┘ └──┘ └──┘ └──┘ _canonize_plane_opts = { 'xmin': { 'yreverse': False, 'zreverse': False, 'coordinate_order': 'yz', 'stepping_order': 'zy', }, 'ymin': { 'zreverse': False, 'xreverse': True, 'coordinate_order': 'zx', 'stepping_order': 'xz', }, 'zmin': { 'xreverse': True, 'yreverse': True, 'coordinate_order': 'xy', 'stepping_order': 'yx', }, 'xmax': { 'yreverse': True, 'zreverse': True, 'coordinate_order': 'yz', 'stepping_order': 'zy', }, 'ymax': { 'zreverse': True, 'xreverse': False, 'coordinate_order': 'zx', 'stepping_order': 'xz', }, 'zmax': { 'xreverse': False, 'yreverse': False, 'coordinate_order': 'xy', 'stepping_order': 'yx', }, } _compress_plane_opts = { 'xmin': { 'yreverse': True, 'zreverse': True, 'coordinate_order': 'yz', 'stepping_order': 'zy', }, 'ymin': { 'zreverse': True, 'xreverse': False, 'coordinate_order': 'zx', 'stepping_order': 'xz', }, 'zmin': { 'xreverse': False, 'yreverse': False, 'coordinate_order': 'xy', 'stepping_order': 'yx', }, 'xmax': { 'yreverse': False, 'zreverse': False, 'coordinate_order': 'yz', 'stepping_order': 'zy', }, 'ymax': { 'zreverse': False, 'xreverse': True, 'coordinate_order': 'zx', 'stepping_order': 'xz', }, 'zmax': { 'xreverse': True, 'yreverse': True, 'coordinate_order': 'xy', 'stepping_order': 'yx', }, }
[docs]class TensorNetwork3D(TensorNetwork): _NDIMS = 3 _EXTRA_PROPS = ( '_site_tag_id', '_x_tag_id', '_y_tag_id', '_z_tag_id', '_Lx', '_Ly', '_Lz', ) def _compatible_3d(self, other): """Check whether ``self`` and ``other`` are compatible 3D tensor networks such that they can remain a 3D tensor network when combined. """ return ( isinstance(other, TensorNetwork3D) and all(getattr(self, e) == getattr(other, e) for e in TensorNetwork3D._EXTRA_PROPS) ) def __and__(self, other): new = super().__and__(other) if self._compatible_3d(other): new.view_as_(TensorNetwork3D, like=self) return new def __or__(self, other): new = super().__or__(other) if self._compatible_3d(other): new.view_as_(TensorNetwork3D, like=self) return new @property def Lx(self): """The number of x-slices. """ return self._Lx @property def Ly(self): """The number of y-slices. """ return self._Ly @property def Lz(self): """The number of z-slices. """ return self._Lz @property def nsites(self): """The total number of sites. """ return self._Lx * self._Ly * self._Lz @property def site_tag_id(self): """The string specifier for tagging each site of this 3D TN. """ return self._site_tag_id
[docs] def site_tag(self, i, j, k): """The name of the tag specifiying the tensor at site ``(i, j, k)``. """ if not isinstance(i, str): i = i % self.Lx if not isinstance(j, str): j = j % self.Ly if not isinstance(k, str): k = k % self.Lz return self.site_tag_id.format(i, j, k)
@property def x_tag_id(self): """The string specifier for tagging each x-slice of this 3D TN. """ return self._x_tag_id def x_tag(self, i): if not isinstance(i, str): i = i % self.Lx return self.x_tag_id.format(i) @property def x_tags(self): """A tuple of all of the ``Lx`` different x-slice tags. """ return tuple(map(self.x_tag, range(self.Lx))) @property def y_tag_id(self): """The string specifier for tagging each y-slice of this 3D TN. """ return self._y_tag_id def y_tag(self, j): if not isinstance(j, str): j = j % self.Ly return self.y_tag_id.format(j) @property def y_tags(self): """A tuple of all of the ``Ly`` different y-slice tags. """ return tuple(map(self.y_tag, range(self.Ly))) @property def z_tag_id(self): """The string specifier for tagging each z-slice of this 3D TN. """ return self._z_tag_id def z_tag(self, k): if not isinstance(k, str): k = k % self.Lz return self.z_tag_id.format(k) @property def z_tags(self): """A tuple of all of the ``Lz`` different z-slice tags. """ return tuple(map(self.z_tag, range(self.Lz))) @property def site_tags(self): """All of the ``Lx * Ly`` site tags. """ return tuple(starmap(self.site_tag, self.gen_site_coos()))
[docs] def maybe_convert_coo(self, coo): """Check if ``coo`` is a tuple of three ints and convert to the corresponding site tag if so. """ if not isinstance(coo, str): try: i, j, k = map(int, coo) return self.site_tag(i, j, k) except (ValueError, TypeError): pass return coo
def _get_tids_from_tags(self, tags, which='all'): """This is the function that lets coordinates such as ``(i, j, k)`` be used for many 'tag' based functions. """ tags = self.maybe_convert_coo(tags) return super()._get_tids_from_tags(tags, which=which)
[docs] def gen_site_coos(self): """Generate coordinates for all the sites in this 3D TN. """ return product(range(self.Lx), range(self.Ly), range(self.Lz))
[docs] def gen_bond_coos(self): """Generate pairs of coordinates for all the bonds in this 3D TN. """ return gen_3d_bonds(self.Lx, self.Ly, self.Lz, steppers=[ lambda i, j, k: (i + 1, j, k), lambda i, j, k: (i, j + 1, k), lambda i, j, k: (i, j, k + 1), ])
[docs] def valid_coo(self, coo, xrange=None, yrange=None, zrange=None): """Check whether ``coo`` is in-bounds. Parameters ---------- coo : (int, int, int), optional The coordinates to check. xrange, yrange, zrange : (int, int), optional The range of allowed values for the x, y, and z coordinates. Returns ------- bool """ if xrange is None: xrange = (0, self.Lx - 1) if yrange is None: yrange = (0, self.Ly - 1) if zrange is None: zrange = (0, self.Lz - 1) return all(mn <= u <= mx for u, (mn, mx) in zip(coo, (xrange, yrange, zrange)))
[docs] def gen_site_coos_present(self): """Return only the sites that are present in this TN. Examples -------- >>> tn = qtn.TN3D_rand(4, 4, 4, 2) >>> tn_sub = tn.select_local('I1,2,3', max_distance=1) >>> list(tn_sub.gen_site_coos_present()) [(0, 2, 3), (1, 1, 3), (1, 2, 2), (1, 2, 3), (1, 3, 3), (2, 2, 3)] """ return ( coo for coo in self.gen_site_coos() if self.site_tag(*coo) in self.tag_map )
[docs] def get_ranges_present(self): """Return the range of site coordinates present in this TN. Returns ------- xrange, yrange, zrange : tuple[tuple[int, int]] The minimum and maximum site coordinates present in each direction. Examples -------- >>> tn = qtn.TN3D_rand(4, 4, 4, 2) >>> tn_sub = tn.select_local('I1,2,3', max_distance=1) >>> tn_sub.get_ranges_present() ((0, 2), (1, 3), (2, 3)) """ xmin = ymin = zmin = float('inf') xmax = ymax = zmax = float('-inf') for i, j, k in self.gen_site_coos_present(): xmin = min(i, xmin) ymin = min(j, ymin) zmin = min(k, zmin) xmax = max(i, xmax) ymax = max(j, ymax) zmax = max(k, zmax) return (xmin, xmax), (ymin, ymax), (zmin, zmax)
def __getitem__(self, key): """Key based tensor selection, checking for integer based shortcut. """ return super().__getitem__(self.maybe_convert_coo(key)) def __repr__(self): """Insert number of slices into standard print. """ s = super().__repr__() extra = (f', Lx={self.Lx}, Ly={self.Ly}, Lz={self.Lz}, ' f'max_bond={self.max_bond()}') s = f'{s[:-2]}{extra}{s[-2:]}' return s def __str__(self): """Insert number of slices into standard print. """ s = super().__repr__() extra = (f', Lx={self.Lx}, Ly={self.Ly}, Lz={self.Lz}' f'max_bond={self.max_bond()}') s = f'{s[:-2]}{extra}{s[-2:]}' return s
[docs] def flatten(self, fuse_multibonds=True, inplace=False): """Contract all tensors corresponding to each site into one. """ tn = self if inplace else self.copy() for i, j, k in self.gen_site_coos(): tn ^= (i, j, k) if fuse_multibonds: tn.fuse_multibonds_() return tn.view_as_(TensorNetwork3DFlat, like=self)
flatten_ = functools.partialmethod(flatten, inplace=True)
[docs] def gen_pairs( self, xrange=None, yrange=None, zrange=None, xreverse=False, yreverse=False, zreverse=False, coordinate_order='xyz', xstep=None, ystep=None, zstep=None, stepping_order='xyz', step_only=None, ): """Helper function for generating pairs of cooordinates for all bonds within a certain range, optionally specifying an order. Parameters ---------- xrange, yrange, zrange : (int, int), optional The range of allowed values for the x, y, and z coordinates. xreverse, yreverse, zreverse : bool, optional Whether to reverse the order of the x, y, and z sweeps. coordinate_order : str, optional The order in which to sweep the x, y, and z coordinates. Earlier dimensions will change slower. If the corresponding range has size 1 then that dimension doesn't need to be specified. xstep, ystep, zstep : int, optional When generating a bond, step in this direction to yield the neighboring coordinate. By default, these follow ``xreverse``, ``yreverse``, and ``zreverse`` respectively. stepping_order : str, optional The order in which to step the x, y, and z coordinates to generate bonds. Does not need to include all dimensions. step_only : int, optional Only perform the ith steps in ``stepping_order``, used to interleave canonizing and compressing for example. Yields ------ coo_a, coo_b : ((int, int, int), (int, int, int)) """ if xrange is None: xrange = (0, self.Lx - 1) if yrange is None: yrange = (0, self.Ly - 1) if zrange is None: zrange = (0, self.Lz - 1) # generate the sites and order we will visit them in sweeps = { 'x': ( range(min(xrange), max(xrange) + 1, +1) if not xreverse else range(max(xrange), min(xrange) - 1, -1) ), 'y': ( range(min(yrange), max(yrange) + 1, +1) if not yreverse else range(max(yrange), min(yrange) - 1, -1) ), 'z': ( range(min(zrange), max(zrange) + 1, +1) if not zreverse else range(max(zrange), min(zrange) - 1, -1) ), } # for convenience, allow subselecting part of stepping_order only if step_only is not None: stepping_order = stepping_order[step_only] # stepping_order = stepping_order[::-1] # at each step generate the bonds if xstep is None: xstep = -1 if xreverse else +1 if ystep is None: ystep = -1 if yreverse else +1 if zstep is None: zstep = -1 if zreverse else +1 steps = { 'x': lambda i, j, k: (i + xstep, j, k), 'y': lambda i, j, k: (i, j + ystep, k), 'z': lambda i, j, k: (i, j, k + zstep), } # make sure all coordinates exist - only allow them not to be specified # if their range is a unit slice for w in 'xyz': if w not in coordinate_order: if len(sweeps[w]) > 1: raise ValueError( f'{w} not in coordinate_order and is not size 1.') else: # just append -> it won't change order as coord is constant coordinate_order += w xi, yi, zi = map(coordinate_order.index, 'xyz') # generate the pairs for perm_coo_a in product(*(sweeps[xyz] for xyz in coordinate_order)): coo_a = perm_coo_a[xi], perm_coo_a[yi], perm_coo_a[zi] for xyz in stepping_order: coo_b = steps[xyz](*coo_a) # filter out bonds with are out of bounds if self.valid_coo(coo_b, xrange, yrange, zrange): yield coo_a, coo_b
[docs] def canonize_plane( self, xrange, yrange, zrange, equalize_norms=False, canonize_opts=None, **gen_pair_opts ): """Canonize every pair of tensors within a subrange, optionally specifying a order to visit those pairs in. """ canonize_opts = ensure_dict(canonize_opts) canonize_opts.setdefault('equalize_norms', equalize_norms) pairs = self.gen_pairs( xrange=xrange, yrange=yrange, zrange=zrange, **gen_pair_opts, ) for coo_a, coo_b in pairs: tag_a = self.site_tag(*coo_a) tag_b = self.site_tag(*coo_b) try: num_a = len(self.tag_map[tag_a]) if num_a > 1: self ^= tag_a except KeyError: continue try: num_b = len(self.tag_map[tag_b]) if num_b > 1: self ^= tag_b except KeyError: continue self.canonize_between(tag_a, tag_b, **canonize_opts)
[docs] def compress_plane( self, xrange, yrange, zrange, max_bond=None, cutoff=1e-10, equalize_norms=False, compress_opts=None, **gen_pair_opts ): """Compress every pair of tensors within a subrange, optionally specifying a order to visit those pairs in. """ compress_opts = ensure_dict(compress_opts) compress_opts.setdefault('absorb', 'both') compress_opts.setdefault('equalize_norms', equalize_norms) pairs = list(self.gen_pairs( xrange=xrange, yrange=yrange, zrange=zrange, **gen_pair_opts, )) for coo_a, coo_b in pairs: tag_a = self.site_tag(*coo_a) tag_b = self.site_tag(*coo_b) try: num_a = len(self.tag_map[tag_a]) if num_a > 1: self ^= tag_a except KeyError: continue try: num_b = len(self.tag_map[tag_b]) if num_b > 1: self ^= tag_b except KeyError: continue self.compress_between(tag_a, tag_b, max_bond=max_bond, cutoff=cutoff, **compress_opts)
def _contract_boundary_core( self, xrange, yrange, zrange, from_which, max_bond, cutoff=1e-10, canonize=True, canonize_interleave=True, layer_tags=None, compress_late=True, equalize_norms=False, compress_opts=None, canonize_opts=None, ): canonize_opts = ensure_dict(canonize_opts) canonize_opts.setdefault('absorb', 'right') compress_opts = ensure_dict(compress_opts) compress_opts.setdefault('absorb', 'both') r3d = Rotator3D(self, xrange, yrange, zrange, from_which) site_tag = r3d.site_tag plane, istep = r3d.plane, r3d.istep jmin, jmax = r3d.jmin, r3d.jmax kmin, kmax = r3d.kmin, r3d.kmax if canonize_interleave: # interleave canonizing and compressing in each direction step_onlys = [0, 1] else: # perform whole sweep of canonizing before compressing step_onlys = [None] if layer_tags is None: layer_tags = [None] for i in r3d.sweep[:-1]: for layer_tag in layer_tags: for j in range(jmin, jmax + 1): for k in range(kmin, kmax + 1): tag1 = site_tag(i, j, k) # outer tag2 = site_tag(i + istep, j, k) # inner if ( (tag1 not in self.tag_map) or (tag2 not in self.tag_map) ): # allow completely missing sites continue if (layer_tag is None) or len(self.tag_map[tag2]) == 1: # contract *any* tensors with pair of coordinates self.contract_((tag1, tag2), which='any') else: # ensure boundary is single tensor per site if len(self.tag_map[tag1]) > 1: self ^= tag1 # contract specific pair (i.e. one 'inner' layer) self.contract_between( tag1, (tag2, layer_tag), equalize_norms=equalize_norms ) # drop inner site tag merged into outer boundary so # we can still uniquely identify inner tensors if layer_tag != layer_tags[-1]: self[tag1].drop_tags(tag2) if not compress_late: tid1, = self.tag_map[tag1] for tidn in self._get_neighbor_tids(tid1): t1, tn = self._tids_get(tid1, tidn) if bonds_size(t1, tn) > max_bond: self._compress_between_tids( tid1, tidn, max_bond=max_bond, cutoff=cutoff, equalize_norms=equalize_norms, **compress_opts ) if compress_late: for step_only in step_onlys: if canonize: self.canonize_plane( xrange=xrange if plane != 'x' else (i, i), yrange=yrange if plane != 'y' else (i, i), zrange=zrange if plane != 'z' else (i, i), equalize_norms=equalize_norms, canonize_opts=canonize_opts, step_only=step_only, **_canonize_plane_opts[from_which] ) self.compress_plane( xrange=xrange if plane != 'x' else (i, i), yrange=yrange if plane != 'y' else (i, i), zrange=zrange if plane != 'z' else (i, i), max_bond=max_bond, cutoff=cutoff, equalize_norms=equalize_norms, compress_opts=compress_opts, step_only=step_only, **_compress_plane_opts[from_which] )
[docs] def contract_boundary( self, max_bond=None, *, cutoff=1e-10, canonize=True, max_separation=1, sequence=None, xmin=None, xmax=None, ymin=None, ymax=None, zmin=None, zmax=None, optimize='auto-hq', equalize_norms=False, compress_opts=None, inplace=False, **contract_boundary_opts, ): """ """ tn = self if inplace else self.copy() contract_boundary_opts['max_bond'] = max_bond # contract_boundary_opts['mode'] = mode contract_boundary_opts['cutoff'] = cutoff contract_boundary_opts['canonize'] = canonize # contract_boundary_opts['layer_tags'] = layer_tags contract_boundary_opts['compress_opts'] = compress_opts contract_boundary_opts['equalize_norms'] = equalize_norms # set default starting borders if any(d is None for d in (xmin, xmax, ymin, ymax, zmin, zmax)): ( (auto_xmin, auto_xmax), (auto_ymin, auto_ymax), (auto_zmin, auto_zmax), ) = self.get_ranges_present() if xmin is None: xmin = auto_xmin if xmax is None: xmax = auto_xmax if ymin is None: ymin = auto_ymin if ymax is None: ymax = auto_ymax if zmin is None: zmin = auto_zmin if zmax is None: zmax = auto_zmax if sequence is None: sequence = ['xmin', 'xmax', 'ymin', 'ymax', 'zmin', 'zmax'] finished = {'x': abs(xmax - xmin) <= max_separation, 'y': abs(ymax - ymin) <= max_separation, 'z': abs(zmax - zmin) <= max_separation} for direction in cycle(sequence): if sum(finished.values()) >= 2: # have reached 'tube' we should contract exactly if equalize_norms is True: tn.equalize_norms_() return tn.contract_(..., optimize=optimize) xyz, minmax = direction[0], direction[1:] if finished[xyz]: # we have already finished this direction continue # prepare the sub-cube we will contract and compress if xyz == 'x': if minmax == 'min': xrange = (xmin, xmin + 1) xmin += 1 elif minmax == 'max': xrange = (xmax - 1, xmax) xmax -= 1 finished['x'] = abs(xmax - xmin) <= max_separation else: xrange = (xmin, xmax) if xyz == 'y': if minmax == 'min': yrange = (ymin, ymin + 1) ymin += 1 elif minmax == 'max': yrange = (ymax - 1, ymax) ymax -= 1 finished['y'] = abs(ymax - ymin) <= max_separation else: yrange = (ymin, ymax) if xyz == 'z': if minmax == 'min': zrange = (zmin, zmin + 1) zmin += 1 elif minmax == 'max': zrange = (zmax - 1, zmax) zmax -= 1 finished['z'] = abs(zmax - zmin) <= max_separation else: zrange = (zmin, zmax) tn._contract_boundary_core( xrange=xrange, yrange=yrange, zrange=zrange, from_which=direction, **contract_boundary_opts)
contract_boundary_ = functools.partialmethod( contract_boundary, inplace=True) def _compute_plane_envs( self, xrange, yrange, zrange, from_which, envs=None, storage_factory=None, **contract_boundary_opts, ): """Compute all 'plane' environments for the cube given by ``xrange``, ``yrange``, ``zrange``, with direction given by ``from_which``. """ tn = self.copy() # rotate virtually r3d = Rotator3D(tn, xrange, yrange, zrange, from_which) plane, p_tag, istep = r3d.plane, r3d.x_tag, r3d.istep # 0th plane has no environment, 1st plane's environment is the 0th if envs is None: if storage_factory is not None: envs = storage_factory() else: envs = {} envs[r3d.sweep[1]] = tn.select_any( p_tag(r3d.sweep[0]), virtual=False ) for i in r3d.sweep[:-2]: # contract the boundary in one step tn._contract_boundary_core( xrange=xrange if plane != 'x' else (i, i + istep), yrange=yrange if plane != 'y' else (i, i + istep), zrange=zrange if plane != 'z' else (i, i + istep), from_which=from_which, **contract_boundary_opts ) # set the boundary as the environment for the next plane beyond envs[i + 2 * istep] = tn.select_any( p_tag(i + istep), virtual=False ) return envs def _maybe_compute_cell_env( self, key, envs=None, storage_factory=None, boundary_order=None, **contract_boundary_opts, ): """Recursively compute the necessary 2D, 1D, and 0D environments. """ if not key: # env is the whole TN return self.copy() if envs is None: if storage_factory is not None: envs = storage_factory() else: envs = {} if key in envs: return envs[key].copy() if boundary_order is None: scores = {'x': -self.Lx, 'y': -self.Ly, 'z': -self.Lz} boundary_order = sorted(scores, key=scores.__getitem__) # check already available parent environments for parent_key in itertools.combinations(key, len(key) - 1): if parent_key in envs: parent_tn = envs[parent_key].copy() break else: # choose which to compute next based on `boundary_order` ranked = sorted(key, key=lambda x: boundary_order.index(x[0]))[:-1] parent_key = tuple(sorted(ranked)) # else choose a parent to compute parent_tn = self._maybe_compute_cell_env( key=parent_key, envs=envs, storage_factory=storage_factory, boundary_order=boundary_order, **contract_boundary_opts) # need to compute parent first - first set the parents range Ls = {'x': self.Lx, 'y': self.Ly, 'z': self.Lz} plane_tags = {'x': self.x_tag, 'y': self.y_tag, 'z': self.z_tag} ranges = {'xrange': None, 'yrange': None, 'zrange': None} for d, s, bsz in parent_key: ranges[f'{d}range'] = (max(0, s - 1), min(s + bsz, Ls[d] - 1)) # then compute the envs for the new direction ``d`` (d, _, bsz), = (x for x in key if x not in parent_key) envs_s_min = parent_tn._compute_plane_envs( from_which=f'{d}min', storage_factory=storage_factory, **ranges, **contract_boundary_opts) envs_s_max = parent_tn._compute_plane_envs( from_which=f'{d}max', storage_factory=storage_factory, **ranges, **contract_boundary_opts) for s in range(0, Ls[d] - bsz + 1): # the central non-boundary slice of tensors tags_s = tuple(map(plane_tags[d], range(s, s + bsz))) tn_s = parent_tn.select_any(tags_s, virtual=False) # the min boundary if s in envs_s_min: tn_s &= envs_s_min[s] # the max boundary imax = s + bsz - 1 if imax in envs_s_max: tn_s &= envs_s_max[imax] # store the newly created cell along with env key_s = tuple(sorted((*parent_key, (d, s, bsz)))) envs[key_s] = tn_s # one of key_s == key return envs[key].copy()
[docs]def is_lone_coo(where): """Check if ``where`` has been specified as a single coordinate triplet. """ return (len(where) == 3) and (isinstance(where[0], Integral))
def calc_cell_sizes(coo_groups, autogroup=True): # get the rectangular size of each coordinate pair bszs = set() for coos in coo_groups: if is_lone_coo(coos): bszs.add((1, 1, 1)) continue xs, ys, zs = zip(*coos) xsz = max(xs) - min(xs) + 1 ysz = max(ys) - min(ys) + 1 zsz = max(zs) - min(zs) + 1 bszs.add((xsz, ysz, zsz)) # remove block size pairs that can be contained in another block pair size bszs = tuple(sorted( b for b in bszs if not any( (b[0] <= b2[0]) and (b[1] <= b2[1]) and (b[2] <= b2[2]) for b2 in bszs - {b} ) )) # return each cell size separately if autogroup: return bszs # else choose a single blocksize that will cover all terms return (tuple(map(max, zip(*bszs))),)
[docs]def cell_to_sites(p): """Turn a cell ``((i0, j0, k0), (di, dj, dk))`` into the sites it contains. Examples -------- >>> cell_to_sites([(3, 4), (2, 2)]) ((3, 4), (3, 5), (4, 4), (4, 5)) """ (i0, j0, k0), (di, dj, dk) = p return tuple((i, j, k) for i in range(i0, i0 + di) for j in range(j0, j0 + dj) for k in range(k0, k0 + dk))
[docs]def sites_to_cell(sites): """Get the minimum covering cell for ``sites``. Examples -------- >>> sites_to_cell([(1, 3, 3), (2, 2, 4)]) ((1, 2, 3) , (2, 2, 2)) """ imin = jmin = kmin = float('inf') imax = jmax = kmax = float('-inf') for i, j, k in sites: imin, jmin, kmin = min(imin, i), min(jmin, j), min(kmin, k) imax, jmax, kmax = max(imax, i), max(jmax, j), max(kmax, k) x_bsz, y_bsz, z_bsz = imax - imin + 1, jmax - jmin + 1, kmax - kmin + 1 return (imin, jmin, kmin), (x_bsz, y_bsz, z_bsz)
def calc_cell_map(cells): # sort in descending total cell size cs = sorted(cells, key=lambda c: (-c[1][0] * c[1][1] * c[1][2], c)) mapping = dict() for c in cs: sites = cell_to_sites(c) for site in sites: mapping[site] = c # this will generate all coordinate pairs with ijk_a < ijk_b for ijk_a, ijk_b in combinations(sites, 2): mapping[ijk_a, ijk_b] = c return mapping
[docs]class TensorNetwork3DVector(TensorNetwork3D, TensorNetwork): """Mixin class for a 3D square lattice vector TN, i.e. one with a single physical index per site. """ _EXTRA_PROPS = ( '_site_tag_id', '_x_tag_id', '_y_tag_id', '_z_tag_id', '_Lx', '_Ly', '_Lz', '_site_ind_id', ) @property def site_ind_id(self): return self._site_ind_id def site_ind(self, i, j, k): if not isinstance(i, str): i = i % self.Lx if not isinstance(j, str): j = j % self.Ly if not isinstance(k, str): k = k % self.Lz return self.site_ind_id.format(i, j, k) def reindex_sites(self, new_id, where=None, inplace=False): if where is None: where = self.gen_site_coos() return self.reindex( { self.site_ind(*coo): new_id.format(*coo) for coo in where }, inplace=inplace ) @site_ind_id.setter def site_ind_id(self, new_id): if self._site_ind_id != new_id: self.reindex_sites(new_id, inplace=True) self._site_ind_id = new_id @property def site_inds(self): """All of the site inds. """ return tuple(starmap(self.site_ind, self.gen_site_coos()))
[docs] def to_dense(self, *inds_seq, **contract_opts): """Return the dense ket version of this 3D vector, i.e. a ``qarray`` with shape (-1, 1). """ if not inds_seq: # just use list of site indices return do('reshape', TensorNetwork.to_dense( self, self.site_inds, **contract_opts ), (-1, 1)) return TensorNetwork.to_dense(self, *inds_seq, **contract_opts)
[docs] def phys_dim(self, i=None, j=None, k=None): """Get the size of the physical indices / a specific physical index. """ if (i is not None) and (j is not None) and (k is not None): pix = self.site_ind(i, j, k) else: # allow for when some physical indices might have been contracted pix = next(iter( ix for ix in self.site_inds if ix in self.ind_map )) return self.ind_size(pix)
[docs] def gate( self, G, where, contract=False, tags=None, info=None, inplace=False, **compress_opts, ): """Apply a gate ``G`` to sites ``where``, preserving the outer site inds. """ if is_lone_coo(where): where = (where,) else: where = tuple(where) inds = tuple(starmap(self.site_ind, where)) return super().gate_inds( G, inds, contract=contract, tags=tags, info=info, inplace=inplace, **compress_opts )
gate_ = functools.partialmethod(gate, inplace=True)
[docs]class TensorNetwork3DFlat(TensorNetwork3D, TensorNetwork): """Mixin class for a 3D square lattice tensor network with a single tensor per site, for example, both PEPS and PEPOs. """ _EXTRA_PROPS = ( '_site_tag_id', '_x_tag_id', '_y_tag_id', '_z_tag_id', '_Lx', '_Ly', '_Lz', )
[docs] def bond(self, coo1, coo2): """Get the name of the index defining the bond between sites at ``coo1`` and ``coo2``. """ b_ix, = self[coo1].bonds(self[coo2]) return b_ix
[docs] def bond_size(self, coo1, coo2): """Return the size of the bond between sites at ``coo1`` and ``coo2``. """ b_ix = self.bond(coo1, coo2) return self[coo1].ind_size(b_ix)
[docs]class PEPS3D(TensorNetwork3DVector, TensorNetwork3DFlat, TensorNetwork3D, TensorNetwork): r"""Projected Entangled Pair States object (3D). Parameters ---------- arrays : sequence of sequence of sequence of array The core tensor data arrays. shape : str, optional Which order the dimensions of the arrays are stored in, the default ``'urfdlbp'`` stands for ('up', 'right', 'front', down', 'left', 'behind', 'physical') meaning (x+, y+, z+, x-, y-, z-, physical) respectively. Arrays on the edge of lattice are assumed to be missing the corresponding dimension. tags : set[str], optional Extra global tags to add to the tensor network. site_ind_id : str, optional String specifier for naming convention of site indices. site_tag_id : str, optional String specifier for naming convention of site tags. x_tag_id : str, optional String specifier for naming convention of x-slice tags. y_tag_id : str, optional String specifier for naming convention of y-slice tags. z_tag_id : str, optional String specifier for naming convention of z-slice tags. """ _EXTRA_PROPS = ( '_site_tag_id', '_x_tag_id', '_y_tag_id', '_z_tag_id', '_Lx', '_Ly', '_Lz', '_site_ind_id', ) def __init__( self, arrays, *, shape='urfdlbp', tags=None, site_ind_id='k{},{},{}', site_tag_id='I{},{},{}', x_tag_id='X{}', y_tag_id='Y{}', z_tag_id='Z{}', **tn_opts ): if isinstance(arrays, PEPS3D): super().__init__(arrays) return tags = tags_to_oset(tags) self._site_ind_id = site_ind_id self._site_tag_id = site_tag_id self._x_tag_id = x_tag_id self._y_tag_id = y_tag_id self._z_tag_id = z_tag_id arrays = tuple(tuple(tuple(z for z in y) for y in x) for x in arrays) self._Lx = len(arrays) self._Ly = len(arrays[0]) self._Lz = len(arrays[0][0]) tensors = [] # cache for both creating and retrieving indices ix = defaultdict(rand_uuid) for i, j, k in self.gen_site_coos(): array = arrays[i][j][k] # figure out if we need to transpose the arrays from some order # other than up right front down left behind physical array_order = shape if i == self.Lx - 1: array_order = array_order.replace('u', '') if j == self.Ly - 1: array_order = array_order.replace('r', '') if k == self.Lz - 1: array_order = array_order.replace('f', '') if i == 0: array_order = array_order.replace('d', '') if j == 0: array_order = array_order.replace('l', '') if k == 0: array_order = array_order.replace('b', '') # allow convention of missing bonds to be singlet dimensions if len(array.shape) != len(array_order): array = do('squeeze', array) transpose_order = tuple( array_order.find(x) for x in 'urfdlbp' if x in array_order ) if transpose_order != tuple(range(len(array_order))): array = do('transpose', array, transpose_order) # get the relevant indices corresponding to neighbours inds = [] if 'u' in array_order: inds.append(ix[(i, j, k), (i + 1, j, k)]) if 'r' in array_order: inds.append(ix[(i, j, k), (i, j + 1, k)]) if 'f' in array_order: inds.append(ix[(i, j, k), (i, j, k + 1)]) if 'd' in array_order: inds.append(ix[(i - 1, j, k), (i, j, k)]) if 'l' in array_order: inds.append(ix[(i, j - 1, k), (i, j, k)]) if 'b' in array_order: inds.append(ix[(i, j, k - 1), (i, j, k)]) inds.append(self.site_ind(i, j, k)) # mix site, slice and global tags ijk_tags = tags | oset((self.site_tag(i, j, k), self.x_tag(i), self.y_tag(j), self.z_tag(k))) # create the site tensor! tensors.append(Tensor(data=array, inds=inds, tags=ijk_tags)) super().__init__(tensors, virtual=True, **tn_opts)
[docs] def permute_arrays(self, shape='urfdlbp'): """Permute the indices of each tensor in this PEPS3D to match ``shape``. This doesn't change how the overall object interacts with other tensor networks but may be useful for extracting the underlying arrays consistently. This is an inplace operation. Parameters ---------- shape : str, optional A permutation of ``'lrp'`` specifying the desired order of the left, right, and physical indices respectively. """ steps = { 'u': lambda i, j, k: (i + 1, j, k), 'r': lambda i, j, k: (i, j + 1, k), 'f': lambda i, j, k: (i, j, k + 1), 'd': lambda i, j, k: (i - 1, j, k), 'l': lambda i, j, k: (i, j - 1, k), 'b': lambda i, j, k: (i, j, k - 1), } for i, j, k in self.gen_site_coos(): t = self[i, j, k] inds = [] for s in shape: if s == 'p': inds.append(self.site_ind(i, j, k)) else: coo2 = steps[s](i, j, k) if self.valid_coo(coo2): t2 = self[coo2] bix, = t.bonds(t2) inds.append(bix) t.transpose_(*inds)
[docs] @classmethod def from_fill_fn( cls, fill_fn, Lx, Ly, Lz, bond_dim, phys_dim=2, **peps3d_opts ): """Create a 3D PEPS from a filling function with signature ``fill_fn(shape)``. Parameters ---------- Lx : int The number of x-slices. Ly : int The number of y-slices. Lz : int The number of z-slices. bond_dim : int The bond dimension. physical : int, optional The physical index dimension. peps_opts Supplied to :class:`~quimb.tensor.tensor_3d.PEPS3D`. Returns ------- psi : PEPS3D """ arrays = [[[None for _ in range(Lz)] for _ in range(Ly)] for _ in range(Lx)] for i, j, k in product(range(Lx), range(Ly), range(Lz)): shape = [] if i != Lx - 1: # bond up shape.append(bond_dim) if j != Ly - 1: # bond right shape.append(bond_dim) if k != Lz - 1: # bond front shape.append(bond_dim) if i != 0: # bond down shape.append(bond_dim) if j != 0: # bond left shape.append(bond_dim) if k != 0: # bond behind shape.append(bond_dim) shape.append(phys_dim) arrays[i][j][k] = fill_fn(shape) return cls(arrays, **peps3d_opts)
[docs] @classmethod def empty( self, Lx, Ly, Lz, bond_dim, phys_dim=2, like='numpy', **peps3d_opts, ): """Create an empty 3D PEPS. Parameters ---------- Lx : int The number of x-slices. Ly : int The number of y-slices. Lz : int The number of z-slices. bond_dim : int The bond dimension. physical : int, optional The physical index dimension. peps3d_opts Supplied to :class:`~quimb.tensor.tensor_3d.PEPS3D`. Returns ------- psi : PEPS3D See Also -------- PEPS3D.from_fill_fn """ return self.from_fill_fn( lambda shape: do("zeros", shape, like=like), Lx, Ly, Lz, bond_dim, phys_dim, **peps3d_opts )
[docs] @classmethod def ones( self, Lx, Ly, Lz, bond_dim, phys_dim=2, like='numpy', **peps3d_opts, ): """Create a 3D PEPS whose tensors are filled with ones. Parameters ---------- Lx : int The number of x-slices. Ly : int The number of y-slices. Lz : int The number of z-slices. bond_dim : int The bond dimension. physical : int, optional The physical index dimension. peps3d_opts Supplied to :class:`~quimb.tensor.tensor_3d.PEPS3D`. Returns ------- psi : PEPS3D See Also -------- PEPS3D.from_fill_fn """ return self.from_fill_fn( lambda shape: do("ones", shape, like=like), Lx, Ly, Lz, bond_dim, phys_dim, **peps3d_opts )
[docs] @classmethod def rand( cls, Lx, Ly, Lz, bond_dim, phys_dim=2, dtype='float64', seed=None, **peps3d_opts ): """Create a random (un-normalized) 3D PEPS. Parameters ---------- Lx : int The number of x-slices. Ly : int The number of y-slices. Lz : int The number of z-slices. bond_dim : int The bond dimension. physical : int, optional The physical index dimension. dtype : dtype, optional The dtype to create the arrays with, default is real double. seed : int, optional A random seed. peps_opts Supplied to :class:`~quimb.tensor.tensor_3d.PEPS3D`. Returns ------- psi : PEPS3D See Also -------- PEPS3D.from_fill_fn """ if seed is not None: seed_rand(seed) def fill_fn(shape): return ops.sensibly_scale(ops.sensibly_scale( randn(shape, dtype=dtype))) return cls.from_fill_fn( fill_fn, Lx, Ly, Lz, bond_dim, phys_dim, **peps3d_opts )
def partial_trace_cluster( self, keep, max_bond=None, *, cutoff=1e-10, max_distance=0, fillin=0, gauges=False, flatten=False, normalized=True, symmetrized='auto', get=None, **contract_boundary_opts ): if is_lone_coo(keep): keep = (keep,) tags = [self.site_tag(i, j, k) for i, j, k in keep] k = self.select_local( tags, 'any', max_distance=max_distance, fillin=fillin, virtual=False) k.add_tag('KET') if gauges: k.gauge_simple_insert(gauges) kix = [self.site_ind(i, j, k) for i, j, k in keep] bix = [rand_uuid() for _ in kix] b = k.H.reindex_(dict(zip(kix, bix))).retag_({'KET': 'BRA'}) rho_tn = k | b rho_tn.fuse_multibonds_() if get == 'tn': return rho_tn # contract boundaries largest dimension last ri, rj, rk = zip(*keep) imin, imax = min(ri), max(ri) jmin, jmax = min(rj), max(rj) kmin, kmax = min(rk), max(rk) sequence = sorted([ ((imax - imin), ('xmin', 'xmax')), ((jmax - jmin), ('ymin', 'ymax')), ((kmax - kmin), ('zmin', 'zmax')), ]) sequence = [x for s in sequence for x in s[1]] rho_t = rho_tn.contract_boundary( max_bond=max_bond, cutoff=cutoff, sequence=sequence, layer_tags=None if flatten else ('KET', 'BRA'), **contract_boundary_opts ) if symmetrized == 'auto': symmetrized = not flatten rho = rho_t.to_dense(kix, bix) # maybe fix up if symmetrized: rho = (rho + dag(rho)) / 2 if normalized: rho = rho / do("trace", rho) return rho def partial_trace( self, keep, max_bond=None, *, cutoff=1e-10, canonize=True, flatten=False, normalized=True, symmetrized='auto', envs=None, storage_factory=None, boundary_order=None, contract_cell_optimize='auto-hq', contract_cell_method='boundary', contract_cell_opts=None, get=None, **contract_boundary_opts, ): contract_cell_opts = ensure_dict(contract_cell_opts) norm = self.make_norm() contract_boundary_opts['max_bond'] = max_bond contract_boundary_opts['cutoff'] = cutoff contract_boundary_opts['canonize'] = canonize contract_boundary_opts['layer_tags'] = ( None if flatten else ('KET', 'BRA') ) if symmetrized == 'auto': symmetrized = not flatten # get minimal covering cell, allow single coordinate if is_lone_coo(keep): keep = (keep,) cell = sites_to_cell(keep) # get the environment surrounding the cell, allowing reuse via ``envs`` (i, j, k), (x_bsz, y_bsz, z_bsz) = cell key = (('x', i, x_bsz), ('y', j, y_bsz), ('z', k, z_bsz)) tn_cell = norm._maybe_compute_cell_env( key=key, envs=envs, storage_factory=storage_factory, boundary_order=boundary_order, **contract_boundary_opts) # cut the bonds between target norm sites to make density matrix tags = [tn_cell.site_tag(*site) for site in keep] kix = [f"k{i},{j},{k}" for i, j, k in keep] bix = [f"b{i},{j},{k}" for i, j, k in keep] for tag, ind_k, ind_b in zip(tags, kix, bix): tn_cell.cut_between((tag, 'KET'), (tag, 'BRA'), ind_k, ind_b) if get == 'tn': return tn_cell if contract_cell_method == 'boundary': # perform the contract to single tensor as boundary contraction # -> still likely far too expensive to contract exactly xmin, xmax = max(0, i - 1), min(i + x_bsz, self.Lx - 1) ymin, ymax = max(0, j - 1), min(j + y_bsz, self.Ly - 1) zmin, zmax = max(0, k - 1), min(k + z_bsz, self.Lz - 1) sequence = [] if i > 0: sequence.append('xmin') if i < self.Lx - 1: sequence.append('xmax') if j > 0: sequence.append('ymin') if j < self.Ly - 1: sequence.append('ymax') if k > 0: sequence.append('zmin') if k < self.Lz - 1: sequence.append('zmax') # contract longest boundary first scores = {'x': xmax - xmin, 'y': ymax - ymin, 'z': zmax - zmin} sequence.sort(key=lambda s: scores[s[0]], reverse=True) rho_tn = tn_cell.contract_boundary_( xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, zmin=zmin, zmax=zmax, sequence=sequence, optimize=contract_cell_optimize, **contract_boundary_opts) else: contract_cell_opts.setdefault('optimize', contract_cell_optimize) contract_cell_opts.setdefault('max_bond', max_bond) contract_cell_opts.setdefault('cutoff', cutoff) rho_tn = tn_cell.contract_compressed_( output_inds=kix + bix, **contract_cell_opts) # turn into raw array rho = rho_tn.to_dense(kix, bix) # maybe fix up if symmetrized: rho = (rho + dag(rho)) / 2 if normalized: rho = rho / do("trace", rho) return rho def compute_local_expectation( self, terms, max_bond=None, *, cutoff=1e-10, canonize=True, flatten=False, normalized=True, symmetrized='auto', return_all=False, envs=None, storage_factory=None, progbar=False, **contract_boundary_opts ): if envs is None: if storage_factory is not None: envs = storage_factory() else: envs = {} if progbar: items = Progbar(terms.items()) else: items = terms.items() expecs = dict() for where, G in items: rho = self.partial_trace( where, max_bond=max_bond, cutoff=cutoff, canonize=canonize, flatten=flatten, symmetrized=symmetrized, normalized=normalized, envs=envs, storage_factory=storage_factory, **contract_boundary_opts, ) expecs[where] = do("trace", G @ rho) if return_all: return expecs return functools.reduce(add, expecs.values())