From e2908fb30f2d5a60b8d15df25b37bbb6e4f17690 Mon Sep 17 00:00:00 2001 From: Johnnie Gray Date: Wed, 3 Apr 2024 18:19:19 -0700 Subject: [PATCH] more 1d and arbgeom compression functionality --- docs/changelog.md | 28 +- quimb/tensor/tensor_1d_compress.py | 55 ++- quimb/tensor/tensor_2d.py | 291 ++++++++++--- quimb/tensor/tensor_arbgeom_compress.py | 532 ++++++++++++++++++++++++ quimb/tensor/tensor_builder.py | 28 +- quimb/tensor/tensor_core.py | 133 +++++- 6 files changed, 962 insertions(+), 105 deletions(-) create mode 100644 quimb/tensor/tensor_arbgeom_compress.py diff --git a/docs/changelog.md b/docs/changelog.md index 190f0cb2..4856859f 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -6,9 +6,16 @@ Release notes for `quimb`. (whats-new-1-7-4)= ## v1.7.4 (unreleased) +**Breaking Changes** + +- all singular value renormalization is turned off by default +- [`TensorNetwork.compress_all`](quimb.tensor.TensorNetwork.compress_all) + now defaults to using some local gauging + + **Enhancements:** -- add `quimb.tensor.tensor_1d_compress` with functions for compressing generic +- add `quimb.tensor.tensor_1d_compress.py` with functions for compressing generic 1D tensor networks (with arbitrary local structure) using various methods. The methods are: @@ -20,6 +27,24 @@ Release notes for `quimb`. - ... and some more niche methods for debugging and testing. And can be accessed via the unified function [`tensor_network_1d_compress`](quimb.tensor.tensor_1d_compress.tensor_network_1d_compress). + Boundary contraction in 2D can now utilize any of these methods. +- add `quimb.tensor.tensor_arbgeom_compress.py` with functions for compressing + arbitrary geometry tensor networks using various methods. The methods are: + + - The **'local-early'** method: + [`tensor_network_ag_compress_local_early`](quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_local_early) + - The **'local-late'** method: + [`tensor_network_ag_compress_local_late`](quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_local_late) + - The **'projector'** method: + [`tensor_network_ag_compress_projector`](quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_projector) + - The **'superorthogonal'** method: + [`tensor_network_ag_compress_superorthogonal`](quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_superorthogonal) + - The **'l2bp'** method: + [`tensor_network_ag_compress_l2bp`](quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_l2bp) + + And can be accessed via the unified function + [`tensor_network_ag_compress`](quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress). + 1D compression can also fall back to these methods. - support PBC in [`tn2d.contract_hotrg`](quimb.tensor.tensor_2d.TensorNetwork2D.contract_hotrg), [`tn2d.contract_ctmrg`](quimb.tensor.tensor_2d.TensorNetwork2D.contract_ctmrg), @@ -34,6 +59,7 @@ Release notes for `quimb`. and [`TN3D_rand_hidden_loop`](quimb.tensor.tensor_builder.TN3D_rand_hidden_loop), with ``cyclic`` kwarg. +- support PBC in the various base PEPS and PEPO construction methods. - add [`tensor_network_apply_op_op`](quimb.tensor.tensor_arbgeom.tensor_network_apply_op_op) for applying 'operator' TNs to 'operator' TNs. - tweak [`tensor_network_apply_op_vec`](quimb.tensor.tensor_arbgeom.tensor_network_apply_op_vec) diff --git a/quimb/tensor/tensor_1d_compress.py b/quimb/tensor/tensor_1d_compress.py index a2b2e0f8..f7692cb9 100644 --- a/quimb/tensor/tensor_1d_compress.py +++ b/quimb/tensor/tensor_1d_compress.py @@ -11,11 +11,13 @@ - [x] the autofit method (via non-1d specific ALS or autodiff) """ + import collections import itertools import warnings from .tensor_arbgeom import tensor_network_apply_op_vec +from .tensor_arbgeom_compress import tensor_network_ag_compress from .tensor_builder import TN1D_matching from .tensor_core import ( TensorNetwork, @@ -1371,19 +1373,46 @@ def tensor_network_1d_compress( """ compress_opts = compress_opts or {} - return _TN1D_COMPRESS_METHODS[method]( - tn, - max_bond=max_bond, - cutoff=cutoff, - site_tags=site_tags, - canonize=canonize, - permute_arrays=permute_arrays, - optimize=optimize, - sweep_reverse=sweep_reverse, - inplace=inplace, - **compress_opts, - **kwargs, - ) + try: + return _TN1D_COMPRESS_METHODS[method]( + tn, + max_bond=max_bond, + cutoff=cutoff, + site_tags=site_tags, + canonize=canonize, + permute_arrays=permute_arrays, + optimize=optimize, + sweep_reverse=sweep_reverse, + inplace=inplace, + **compress_opts, + **kwargs, + ) + except KeyError: + # try arbitrary geometry methods + + if sweep_reverse: + warnings.warn( + "sweep_reverse has no effect for " + "arbitrary geometry (AG) methods." + ) + + tnc = tensor_network_ag_compress( + tn, + max_bond=max_bond, + cutoff=cutoff, + method=method, + site_tags=site_tags, + canonize=canonize, + optimize=optimize, + inplace=inplace, + **compress_opts, + **kwargs, + ) + + if permute_arrays: + possibly_permute_(tnc, permute_arrays) + + return tnc # --------------- MPO-MPS gating using 1D compression methods --------------- # diff --git a/quimb/tensor/tensor_2d.py b/quimb/tensor/tensor_2d.py index 653ea77a..452048b4 100644 --- a/quimb/tensor/tensor_2d.py +++ b/quimb/tensor/tensor_2d.py @@ -1,45 +1,45 @@ """Classes and algorithms related to 2D tensor networks.""" -import re -import random import functools -from operator import add -from numbers import Integral -from itertools import product, cycle, combinations, count, chain +import random +import re from collections import defaultdict +from itertools import chain, combinations, count, cycle, product +from numbers import Integral +from operator import add -from autoray import do, infer_backend, get_dtype_name +from autoray import do, get_dtype_name, infer_backend from ..gen.operators import swap from ..gen.rand import randn, seed_rand from ..utils import ( - deprecated, - print_multi_line, check_opt, - pairwise, + deprecated, ensure_dict, + pairwise, + print_multi_line, ) from ..utils import progbar as Progbar from . import array_ops as ops +from . import decomp +from .tensor_1d import maybe_factor_gate_into_tensor +from .tensor_arbgeom import ( + TensorNetworkGen, + TensorNetworkGenOperator, + TensorNetworkGenVector, + tensor_network_apply_op_vec, +) from .tensor_core import ( - bonds_size, + Tensor, + TensorNetwork, bonds, - oset_union, + bonds_size, oset, + oset_union, rand_uuid, tags_to_oset, tensor_contract, - Tensor, - TensorNetwork, ) -from .tensor_arbgeom import ( - tensor_network_apply_op_vec, - TensorNetworkGen, - TensorNetworkGenOperator, - TensorNetworkGenVector, -) -from .tensor_1d import maybe_factor_gate_into_tensor -from . import decomp def manhattan_distance(coo_a, coo_b): @@ -1075,6 +1075,7 @@ def compress_plane( """ compress_opts = ensure_dict(compress_opts) compress_opts.setdefault("absorb", "right") + compress_opts.setdefault("reduced", "left") compress_opts.setdefault("equalize_norms", equalize_norms) pairs = self.gen_pairs( @@ -1247,6 +1248,62 @@ def compress_column( compress_opts=compress_opts, ) + def _contract_boundary_core_via_1d( + self, + xrange, + yrange, + from_which, + max_bond, + cutoff=1e-10, + method="dm", + layer_tags=None, + **compress_opts, + ): + from quimb.tensor.tensor_1d_compress import tensor_network_1d_compress + + r2d = Rotator2D(self, xrange, yrange, from_which) + site_tag = r2d.site_tag + istep = r2d.istep + + site_tag_tmps = [f"__ST{j}__" for j in r2d.sweep_other] + + if layer_tags is None: + layer_tags = [None] + + # XXX: compress the initial row, which may be multiple tensors? + + for i in r2d.sweep[:-1]: + for layer_tag in layer_tags: + for j, st in zip(r2d.sweep_other, site_tag_tmps): + tag1 = site_tag(i, j) # outer + tag2 = site_tag(i + istep, j) # inner + if layer_tag is None: + # tag and compress any inner tensors + self.select_any((tag1, tag2)).add_tag(st) + else: + # only tag and compress one inner layer + self.select_all((tag1,)).add_tag(st) + self.select_all((tag2, layer_tag)).add_tag(st) + + # split off the boundary network + self, tn_boundary = self.partition(site_tag_tmps, inplace=True) + + # compress it inplace + tensor_network_1d_compress( + tn_boundary, + max_bond=max_bond, + cutoff=cutoff, + method=method, + site_tags=site_tag_tmps, + inplace=True, + **compress_opts, + ) + + # recombine with the main network + self |= tn_boundary + + self.drop_tags(site_tag_tmps) + def _contract_boundary_core( self, xrange, @@ -1611,7 +1668,7 @@ def _contract_boundary_projector( # this handles cyclic boundary conditions jnext = r.get_jnext(j) if jnext is not None: - ltags = (r.site_tag(i, j), r.site_tag(inext, j)) + ltags = (r.site_tag(i, j), r.site_tag(inext, j)) rtags = (r.site_tag(i, jnext), r.site_tag(inext, jnext)) # │ │ # ──O─┐ chi ┌─O── i+1 @@ -1661,7 +1718,7 @@ def contract_boundary_from( """Unified entrypoint for contracting any rectangular patch of tensors from any direction, with any boundary method. """ - check_opt("mode", mode, ("mps", "projector", "full-bond")) + # check_opt("mode", mode, ("mps", "projector", "full-bond")) tn = self if inplace else self.copy() @@ -1678,7 +1735,7 @@ def contract_boundary_from( contract_boundary_opts["cutoff"] = cutoff contract_boundary_opts["compress_opts"] = compress_opts - if mode == "projector": + if mode == "projector2d": tn._contract_boundary_projector(**contract_boundary_opts) return tn @@ -1686,8 +1743,14 @@ def contract_boundary_from( contract_boundary_opts["canonize"] = canonize contract_boundary_opts["layer_tags"] = layer_tags contract_boundary_opts["sweep_reverse"] = sweep_reverse - tn._contract_boundary_core(**contract_boundary_opts) + if mode == "mps": + tn._contract_boundary_core(**contract_boundary_opts) + return tn + + tn._contract_boundary_core_via_1d( + method=mode, **contract_boundary_opts + ) return tn contract_boundary_from_ = functools.partialmethod( @@ -2274,7 +2337,6 @@ def _contract_interleaved_boundary_sequence( if sequence is None: # contract in both sides along short dimension -> less compression if self.Lx >= self.Ly: - if self.is_cyclic_x(): sequence = ("xmin",) else: @@ -2393,6 +2455,7 @@ def contract_boundary( ymin=None, ymax=None, max_separation=1, + max_unfinished=1, around=None, equalize_norms=False, final_contract=True, @@ -2453,6 +2516,9 @@ def contract_boundary( max_separation : int, optional If ``around is None``, when any two sides become this far apart simply contract the remaining tensor network. + max_unfinished : int, optional + If ``around is None``, when this many sides are still not within + ``max_separation`` simply contract the remaining tensor network. around : None or sequence of (int, int), optional If given, don't contract the square of sites bounding these coordinates. @@ -2495,6 +2561,7 @@ def contract_boundary( ymin=ymin, ymax=ymax, max_separation=max_separation, + max_unfinished=max_unfinished, around=around, equalize_norms=equalize_norms, final_contract=final_contract, @@ -4676,7 +4743,7 @@ class PEPS(TensorNetwork2DVector, TensorNetwork2DFlat): Parameters ---------- - arrays : sequence of sequence of array + arrays : sequence of sequence of array_like The core tensor data arrays. shape : str, optional Which order the dimensions of the arrays are stored in, the default @@ -4729,10 +4796,13 @@ def __init__( arrays = tuple(tuple(x for x in xs) for xs in arrays) self._Lx = len(arrays) self._Ly = len(arrays[0]) - tensors = [] + + cyclicx = sum(d > 1 for d in arrays[0][1].shape) == 5 + cyclicy = sum(d > 1 for d in arrays[1][0].shape) == 5 # cache for both creating and retrieving indices ix = defaultdict(rand_uuid) + tensors = [] for i, j in self.gen_site_coos(): array = arrays[i][j] @@ -4740,13 +4810,13 @@ def __init__( # figure out if we need to transpose the arrays from some order # other than up right down left physical array_order = shape - if i == self.Lx - 1: + if (not cyclicx) and (i == self.Lx - 1): array_order = array_order.replace("u", "") - if j == self.Ly - 1: + if (not cyclicy) and (j == self.Ly - 1): array_order = array_order.replace("r", "") - if i == 0: + if (not cyclicx) and (i == 0): array_order = array_order.replace("d", "") - if j == 0: + if (not cyclicy) and (j == 0): array_order = array_order.replace("l", "") # allow convention of missing bonds to be singlet dimensions @@ -4762,17 +4832,20 @@ def __init__( # get the relevant indices corresponding to neighbours inds = [] if "u" in array_order: - inds.append(ix[(i, j), (i + 1, j)]) + i_u = (i + 1) % self.Lx + inds.append(ix[frozenset(((i, j), (i_u, j)))]) if "r" in array_order: - inds.append(ix[(i, j), (i, j + 1)]) + j_r = (j + 1) % self.Ly + inds.append(ix[frozenset(((i, j), (i, j_r)))]) if "d" in array_order: - inds.append(ix[(i - 1, j), (i, j)]) + i_d = (i - 1) % self.Lx + inds.append(ix[frozenset(((i_d, j), (i, j)))]) if "l" in array_order: - inds.append(ix[(i, j - 1), (i, j)]) + j_l = (j - 1) % self.Ly + inds.append(ix[frozenset(((i, j_l), (i, j)))]) inds.append(self.site_ind(i, j)) # mix site, row, column and global tags - ij_tags = tags | oset( (self.site_tag(i, j), self.x_tag(i), self.y_tag(j)) ) @@ -4790,6 +4863,8 @@ def from_fill_fn( Ly, bond_dim, phys_dim=2, + cyclic=False, + shape="urdlp", **peps_opts, ): """Create a 2D PEPS from a filling function with signature @@ -4803,8 +4878,13 @@ def from_fill_fn( The number of columns. bond_dim : int The bond dimension. - physical : int, optional + phys_dim : int, optional The physical index dimension. + cyclic : bool or tuple[bool, bool], optional + Whether the lattice is cyclic in the x and y directions. + shape : str, optional + How to layout the indices of the tensors, the default is + ``(up, right, down, left bra, ket) == 'urdlbk'``. peps_opts Supplied to :class:`~quimb.tensor.tensor_2d.PEPS`. @@ -4814,21 +4894,29 @@ def from_fill_fn( """ arrays = [[None for _ in range(Ly)] for _ in range(Lx)] + try: + cyclicx, cyclicy = cyclic + except (TypeError, ValueError): + cyclicx = cyclicy = cyclic + for i, j in product(range(Lx), range(Ly)): - shape = [] - if i != Lx - 1: # bond up - shape.append(bond_dim) - if j != Ly - 1: # bond right - shape.append(bond_dim) - if i != 0: # bond down - shape.append(bond_dim) - if j != 0: # bond left - shape.append(bond_dim) - shape.append(phys_dim) + shp = [] + + for which in shape: + if (which == "u") and (cyclicx or (i < Lx - 1)): # bond up + shp.append(bond_dim) + elif (which == "r") and (cyclicy or (j < Ly - 1)): # bond right + shp.append(bond_dim) + elif (which == "d") and (cyclicx or (i > 0)): # bond down + shp.append(bond_dim) + elif (which == "l") and (cyclicy or (j > 0)): # bond left + shp.append(bond_dim) + elif which == "p": + shp.append(phys_dim) - arrays[i][j] = fill_fn(shape) + arrays[i][j] = fill_fn(shp) - return cls(arrays, **peps_opts) + return cls(arrays, shape=shape, **peps_opts) @classmethod def empty(cls, Lx, Ly, bond_dim, phys_dim=2, like="numpy", **peps_opts): @@ -4900,7 +4988,16 @@ def ones(cls, Lx, Ly, bond_dim, phys_dim=2, like="numpy", **peps_opts): @classmethod def rand( - cls, Lx, Ly, bond_dim, phys_dim=2, dtype=float, seed=None, **peps_opts + cls, + Lx, + Ly, + bond_dim, + phys_dim=2, + dist="normal", + loc=0.0, + dtype="float64", + seed=None, + **peps_opts, ): """Create a random (un-normalized) PEPS. @@ -4914,6 +5011,10 @@ def rand( The bond dimension. physical : int, optional The physical index dimension. + dist : {'normal', 'uniform', 'rademacher', 'exp'}, optional + Type of random number to generate, defaults to 'normal'. + loc : float, optional + An additive offset to add to the random numbers. dtype : dtype, optional The dtype to create the arrays with, default is real double. seed : int, optional @@ -4934,7 +5035,9 @@ def rand( def fill_fn(shape): return ops.sensibly_scale( - ops.sensibly_scale(randn(shape, dtype=dtype)) + ops.sensibly_scale( + randn(shape, dist=dist, loc=loc, dtype=dtype) + ) ) return cls.from_fill_fn( @@ -5066,10 +5169,13 @@ def __init__( arrays = tuple(tuple(x for x in xs) for xs in arrays) self._Lx = len(arrays) self._Ly = len(arrays[0]) - tensors = [] # cache for both creating and retrieving indices ix = defaultdict(rand_uuid) + tensors = [] + + cyclicx = sum(d > 1 for d in arrays[0][1].shape) == 6 + cyclicy = sum(d > 1 for d in arrays[1][0].shape) == 6 for i, j in product(range(self.Lx), range(self.Ly)): array = arrays[i][j] @@ -5077,13 +5183,13 @@ def __init__( # figure out if we need to transpose the arrays from some order # other than up right down left physical array_order = shape - if i == self.Lx - 1: + if (not cyclicx) and (i == self.Lx - 1): array_order = array_order.replace("u", "") - if j == self.Ly - 1: + if (not cyclicy) and (j == self.Ly - 1): array_order = array_order.replace("r", "") - if i == 0: + if (not cyclicx) and (i == 0): array_order = array_order.replace("d", "") - if j == 0: + if (not cyclicy) and (j == 0): array_order = array_order.replace("l", "") # allow convention of missing bonds to be singlet dimensions @@ -5099,13 +5205,17 @@ def __init__( # get the relevant indices corresponding to neighbours inds = [] if "u" in array_order: - inds.append(ix[(i + 1, j), (i, j)]) + i_u = (i + 1) % self.Lx + inds.append(ix[frozenset(((i, j), (i_u, j)))]) if "r" in array_order: - inds.append(ix[(i, j), (i, j + 1)]) + j_r = (j + 1) % self.Ly + inds.append(ix[frozenset(((i, j), (i, j_r)))]) if "d" in array_order: - inds.append(ix[(i, j), (i - 1, j)]) + i_d = (i - 1) % self.Lx + inds.append(ix[frozenset(((i_d, j), (i, j)))]) if "l" in array_order: - inds.append(ix[(i, j - 1), (i, j)]) + j_l = (j - 1) % self.Ly + inds.append(ix[frozenset(((i, j_l), (i, j)))]) inds.append(self.lower_ind(i, j)) inds.append(self.upper_ind(i, j)) @@ -5127,6 +5237,7 @@ def from_fill_fn( Ly, bond_dim, phys_dim=2, + cyclic=False, shape="urdlbk", **pepo_opts, ): @@ -5145,6 +5256,8 @@ def from_fill_fn( The bond dimension. phys_dim : int, optional The physical indices dimension. + cyclic : bool or tuple[bool, bool], optional + Whether the lattice is cyclic in the x and y directions. shape : str, optional How to layout the indices of the tensors, the default is ``(up, right, down, left bra, ket) == 'urdlbk'``. @@ -5153,16 +5266,21 @@ def from_fill_fn( """ arrays = [[None for _ in range(Ly)] for _ in range(Lx)] + try: + cyclicx, cyclicy = cyclic + except (TypeError, ValueError): + cyclicx = cyclicy = cyclic + for i, j in product(range(Lx), range(Ly)): shp = [] for which in shape: - if which == "u" and i < Lx - 1: + if (which == "u") and (cyclicx or (i < Lx - 1)): shp.append(bond_dim) - elif which == "r" and j < Ly - 1: + elif (which == "r") and (cyclicy or (j < Ly - 1)): shp.append(bond_dim) - elif which == "d" and i > 0: + elif (which == "d") and (cyclicx or (i > 0)): shp.append(bond_dim) - elif which == "l" and j > 0: + elif (which == "l") and (cyclicy or (j > 0)): shp.append(bond_dim) elif which in ("b", "k"): shp.append(phys_dim) @@ -5233,6 +5351,49 @@ def fill_fn(shape): rand_herm = functools.partialmethod(rand, herm=True) + @classmethod + def zeros( + cls, + Lx, + Ly, + bond_dim, + phys_dim=2, + dtype="float64", + backend="numpy", + **pepo_opts, + ): + """Create a PEPO with all zero entries. + + Parameters + ---------- + Lx : int + The number of rows. + Ly : int + The number of columns. + 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. + backend : str, optional + Which backend to use, default is ``'numpy'``. + pepo_opts + Supplied to :class:`~quimb.tensor.tensor_2d.PEPO`. + """ + + def fill_fn(shape): + return do("zeros", shape, dtype=dtype, like=backend) + + return cls.from_fill_fn( + fill_fn, + Lx=Lx, + Ly=Ly, + bond_dim=bond_dim, + phys_dim=phys_dim, + **pepo_opts, + ) + def add_PEPO(self, other, inplace=False): """Add this PEPO with another.""" if (self.Lx, self.Ly) != (other.Lx, other.Ly): diff --git a/quimb/tensor/tensor_arbgeom_compress.py b/quimb/tensor/tensor_arbgeom_compress.py new file mode 100644 index 00000000..78259ba4 --- /dev/null +++ b/quimb/tensor/tensor_arbgeom_compress.py @@ -0,0 +1,532 @@ +"""Generic methods for compressing arbitrary geometry tensor networks, where +the tensor network can locally have arbitrary structure and outer indices. + +- [x] projector +- [x] l2bp +- [x] local early +- [x] local late +- [x] superorthogonal + +""" + +from ..utils import ensure_dict +from .tensor_core import choose_local_compress_gauge_settings + + +def create_lazy_edge_map(tn, site_tags=None): + """Given a tensor network, where each tensor is in exactly one group or + 'site', compute which sites are connected to each other, without checking + each pair. + """ + if site_tags is None: + site_tags = set(tn.site_tags) + else: + site_tags = set(site_tags) + + edges = {} + neighbors = {} + + for ix in tn.ind_map: + ts = tn._inds_get(ix) + tags = {tag for t in ts for tag in t.tags if tag in site_tags} + if len(tags) >= 2: + # index spans multiple sites + i, j = tuple(sorted(tags)) + + if (i, j) in edges: + # already processed this edge + continue + + # add to neighbor map + neighbors.setdefault(i, []).append(j) + neighbors.setdefault(j, []).append(i) + + edges[(i, j)] = None + + return edges, neighbors + + +def tensor_network_ag_compress_projector( + tn, + max_bond=None, + cutoff=1e-10, + site_tags=None, + canonize=True, + canonize_opts=None, + lazy=False, + optimize="auto-hq", + inplace=False, + **compress_opts, +): + """Compress an arbtrary geometry tensor network, with potentially multiple + tensors per site, using locally computed projectors. + + Very loosely, this is like a generalization HOTRG. + + Parameters + ---------- + tn : TensorNetwork + The tensor network to compress. Every tensor should have exactly one of + the site tags. Each site can have multiple tensors and output indices. + max_bond : int + The maximum bond dimension to compress to. + cutoff : float, optional + A dynamic threshold for discarding singular values when compressing. + site_tags : sequence of str, optional + The tags to use to group the tensors from ``tn``. If not + given, uses ``tn.site_tags``. The tensor network built will have one + tensor per site. + canonize : bool, optional + Whether to pseudo canonicalize the initial tensor network. + canonize_opts + Supplied to :meth:`~quimb.tensor.tensor_core.TensorNetwork.gauge_all`. + lazy : bool, optional + Whether to leave the computed projectors uncontracted, default: False. + optimize : str, optional + The contraction path optimizer to use. + inplace : bool, optional + Whether to perform the compression inplace. + compress_opts + Supplied to :func:`~quimb.tensor.tensor_split`. + + Returns + ------- + TensorNetwork + """ + tn = tn if inplace else tn.copy() + + if site_tags is None: + site_tags = tn.site_tags + + edges, _ = create_lazy_edge_map(tn, site_tags) + + if canonize: + # optionally precondition the uncontracted network + canonize_opts = ensure_dict(canonize_opts) + tn.gauge_all_(**canonize_opts) + + # then compute projectors using local information + tn_calc = tn.copy() + for taga, tagb in edges: + tn_calc.insert_compressor_between_regions_( + [taga], + [tagb], + max_bond=max_bond, + cutoff=cutoff, + insert_into=tn, + new_ltags=[taga], + new_rtags=[tagb], + optimize=optimize, + **compress_opts, + ) + + if not lazy: + # then contract each site with all surrounding projectors + for st in site_tags: + tn.contract_(st, optimize=optimize) + + return tn + + +def tensor_network_ag_compress_local_early( + tn, + max_bond=None, + cutoff=1e-10, + site_tags=None, + canonize=True, + tree_gauge_distance=None, + canonize_distance=None, + canonize_after_distance=None, + mode="auto", + optimize="auto-hq", + inplace=False, + **compress_opts, +): + """Compress an arbtrary geometry tensor network, with potentially multiple + tensors per site, using explicit contraction followed by immediate + ('early') compression. In other words, contractions are interleaved with + compressions. + + Very loosely, this is like a generalization of the 'zip-up' algorithm in + 1D, but for arbitrary geometry. + + + Parameters + ---------- + tn : TensorNetwork + The tensor network to compress. Every tensor should have exactly one of + the site tags. Each site can have multiple tensors and output indices. + max_bond : int + The maximum bond dimension to compress to. + cutoff : float, optional + A dynamic threshold for discarding singular values when compressing. + site_tags : sequence of str, optional + The tags to use to group the tensors from ``tn``. If not + given, uses ``tn.site_tags``. The tensor network built will have one + tensor per site. + canonize : bool, optional + Whether to locally gauge before each compression, defaults to True. + tree_gauge_distance : int, optional + The distance to locally gauge to before each compression. Defaults to + 3. + canonize_distance : int, optional + The distance to canonize to before each compression, by default this + is set by ``tree_gauge_distance``. + canonize_after_distance : int, optional + The distance to canonize to after each compression, by default this + is set by ``tree_gauge_distance``, depending on ``mode``. + mode : {'auto', 'basic', 'virtual-tree', ...}, optional + The mode to use for the local gauging. If 'auto' will default to + virtual tree gauging, or basic if `tree_gauge_distance` is 0. + optimize : str, optional + The contraction path optimizer to use. + inplace : bool, optional + Whether to perform the compression inplace. + compress_opts + Supplied to + :meth:`~quimb.tensor.tensor_core.TensorNetwork.compress_between`. + + Returns + ------- + TensorNetwork + """ + tnc = tn if inplace else tn.copy() + + if site_tags is None: + site_tags = tnc.site_tags + + _, neighbors = create_lazy_edge_map(tnc, site_tags) + + canonize_distance, canonize_after_distance, mode = ( + choose_local_compress_gauge_settings( + canonize, + tree_gauge_distance, + canonize_distance, + canonize_after_distance, + mode, + ) + ) + + seen = set() + queue = [next(iter(site_tags))] + + while queue: + # process sites in a breadth-first manner + taga = queue.pop() + seen.add(taga) + for tagb in neighbors[taga]: + if tagb not in seen: + queue.append(tagb) + + # contract this site + tnc.contract_(taga, optimize=optimize) + + # then immediately compress around it + (tida,) = tnc._get_tids_from_tags(taga) + for tidb in tnc._get_neighbor_tids(tida): + tnc._compress_between_tids( + tida, + tidb, + max_bond=max_bond, + cutoff=cutoff, + canonize_distance=canonize_distance, + canonize_after_distance=canonize_after_distance, + mode=mode, + **compress_opts, + ) + + return tnc + + +def tensor_network_ag_compress_local_late( + tn, + max_bond=None, + cutoff=1e-10, + site_tags=None, + canonize=True, + tree_gauge_distance=None, + canonize_distance=None, + canonize_after_distance=None, + mode="auto", + optimize="auto-hq", + inplace=False, + **compress_opts, +): + """Compress an arbtrary geometry tensor network, with potentially multiple + tensors per site, by explicitly contracting all sites first and then + ('late') locally compressing. In other words, all contractions happen, then + all compressions happen. + + Very loosely, this is like a generalization of the 'direct' algorithm in + 1D, but for arbitrary geometry. + + Parameters + ---------- + tn : TensorNetwork + The tensor network to compress. Every tensor should have exactly one of + the site tags. Each site can have multiple tensors and output indices. + max_bond : int + The maximum bond dimension to compress to. + cutoff : float, optional + A dynamic threshold for discarding singular values when compressing. + site_tags : sequence of str, optional + The tags to use to group the tensors from ``tn``. If not + given, uses ``tn.site_tags``. The tensor network built will have one + tensor per site. + canonize : bool, optional + Whether to locally gauge before each compression, defaults to True. + tree_gauge_distance : int, optional + The distance to locally gauge to before each compression. Defaults to + 3. + canonize_distance : int, optional + The distance to canonize to before each compression, by default this + is set by ``tree_gauge_distance``. + canonize_after_distance : int, optional + The distance to canonize to after each compression, by default this + is set by ``tree_gauge_distance``, depending on ``mode``. + mode : {'auto', 'basic', 'virtual-tree', ...}, optional + The mode to use for the local gauging. If 'auto' will default to + virtual tree gauging, or basic if `tree_gauge_distance` is 0. + optimize : str, optional + The contraction path optimizer to use. + inplace : bool, optional + Whether to perform the compression inplace. + compress_opts + Supplied to + :meth:`~quimb.tensor.tensor_core.TensorNetwork.compress_between`. + + Returns + ------- + TensorNetwork + """ + tnc = tn if inplace else tn.copy() + + if site_tags is None: + site_tags = tnc.site_tags + + for st in site_tags: + tnc.contract_(st, optimize=optimize) + + tnc.compress_all_( + max_bond=max_bond, + cutoff=cutoff, + canonize=canonize, + tree_gauge_distance=tree_gauge_distance, + canonize_distance=canonize_distance, + canonize_after_distance=canonize_after_distance, + mode=mode, + **compress_opts, + ) + + return tnc + + +def tensor_network_ag_compress_superorthogonal( + tn, + max_bond=None, + cutoff=1e-10, + site_tags=None, + canonize=True, + optimize="auto-hq", + inplace=False, + **compress_opts, +): + """Compress an arbtrary geometry tensor network, with potentially multiple + tensors per site, using the 'superorthogonal' / 'Vidal' / quasi-canonical + / 'simple update' gauge for compression. This is the same gauge as used in + L2BP, but the intermediate tensor network is explicitly constructed. + + Parameters + ---------- + tn : TensorNetwork + The tensor network to compress. Every tensor should have exactly one of + the site tags. Each site can have multiple tensors and output indices. + max_bond : int + The maximum bond dimension to compress to. + cutoff : float, optional + A dynamic threshold for discarding singular values when compressing. + site_tags : sequence of str, optional + The tags to use to group the tensors from ``tn``. If not + given, uses ``tn.site_tags``. The tensor network built will have one + tensor per site. + canonize : bool, optional + Whether to locally gauge before each compression, defaults to True. + optimize : str, optional + The contraction path optimizer to use. + inplace : bool, optional + Whether to perform the compression inplace. + compress_opts + Supplied to + :meth:`~quimb.tensor.tensor_core.TensorNetwork.compress_all_simple`. + + Returns + ------- + TensorNetwork + """ + tnc = tn if inplace else tn.copy() + + if site_tags is None: + site_tags = tnc.site_tags + + for st in site_tags: + tnc.contract_(st, optimize=optimize) + + if not canonize: + # turn off gauging effect + compress_opts.setdefault("max_iterations", 1) + compress_opts.setdefault("tol", 0.0) + else: + compress_opts.setdefault("max_iterations", 1000) + compress_opts.setdefault("tol", 5e-6) + + tnc.compress_all_simple_( + max_bond=max_bond, + cutoff=cutoff, + **compress_opts, + ) + + return tnc + + +def tensor_network_ag_compress_l2bp( + tn, + max_bond=None, + cutoff=1e-10, + site_tags=None, + canonize=True, + damping=0.0, + local_convergence=True, + update="parallel", + optimize="auto-hq", + inplace=False, + **compress_opts, +): + """Compress an arbitrary geometry tensor network, with potentially multiple + tensors per site, using lazy 2-norm belief propagation. + + Parameters + ---------- + tn : TensorNetwork + The tensor network to compress. Every tensor should have exactly one of + the site tags. Each site can have multiple tensors and output indices. + max_bond : int + The maximum bond dimension to compress to. + cutoff : float, optional + A dynamic threshold for discarding singular values when compressing. + site_tags : sequence of str, optional + The tags to use to group the tensors from ``tn``. If not + given, uses ``tn.site_tags``. The tensor network built will have one + tensor per site. + canonize : bool, optional + Whether to locally gauge before each compression, defaults to True. + damping : float, optional + How much to dampen message updates, to help convergence, defaults to 0. + local_convergence : bool, optional + Whether to use local convergence criteria, defaults to True. + update : {'parallel', 'sequential'}, optional + Whether to update all messages in parallel or sequentially, defaults to + 'parallel'. + optimize : str, optional + The contraction path optimizer to use. + inplace : bool, optional + Whether to perform the compression inplace. + compress_opts + Supplied to + :func:`~quimb.experimental.belief_propagation.l2bp.compress_l2bp`. + + Returns + ------- + TensorNetwork + """ + from quimb.experimental.belief_propagation.l2bp import compress_l2bp + + if not canonize: + compress_opts.setdefault("max_iterations", 1) + + tnc = compress_l2bp( + tn, + max_bond=max_bond, + cutoff=cutoff, + site_tags=site_tags, + damping=damping, + local_convergence=local_convergence, + update=update, + optimize=optimize, + inplace=inplace, + **compress_opts, + ) + + return tnc + + +_TNAG_COMPRESS_METHODS = { + "local-early": tensor_network_ag_compress_local_early, + "local-late": tensor_network_ag_compress_local_late, + "projector": tensor_network_ag_compress_projector, + "superorthogonal": tensor_network_ag_compress_superorthogonal, + "l2bp": tensor_network_ag_compress_l2bp, +} + + +def tensor_network_ag_compress( + tn, + max_bond, + cutoff=1e-10, + method="local-early", + site_tags=None, + canonize=True, + optimize="auto-hq", + inplace=False, + **kwargs, +): + """Compress an arbitrary geometry tensor network, with potentially multiple + tensors per site. + + Parameters + ---------- + tn : TensorNetwork + The tensor network to compress. Every tensor should have exactly one of + the site tags. Each site can have multiple tensors and output indices. + max_bond : int + The maximum bond dimension to compress to. + cutoff : float, optional + A dynamic threshold for discarding singular values when compressing. + method : {'local-early', 'local-late', 'projector', 'superorthogonal', 'l2bp'}, optional + The compression method to use: + + - 'local-early': explicitly contract each site and interleave with + immediate compression, see + :func:`~quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_local_early`. + - 'local-late': explicitly contract all sites and then compress, see + :func:`~quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_local_late`. + - 'projector': use locally computed projectors, see + :func:`~quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_projector`. + - 'superorthogonal': use the 'superorthogonal' gauge, see + :func:`~quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_superorthogonal`. + - 'l2bp': use lazy 2-norm belief propagation, see + :func:`~quimb.tensor.tensor_arbgeom_compress.tensor_network_ag_compress_l2bp`. + + site_tags : sequence of str, optional + The tags to use to group the tensors from ``tn``. If not + given, uses ``tn.site_tags``. The tensor network built will have one + tensor per site. + canonize : bool, optional + Whether to perform canonicalization, pseudo or otherwise depending on + the method, before compressing. + optimize : str, optional + The contraction path optimizer to use. + inplace : bool, optional + Whether to perform the compression inplace. + kwargs + Supplied to the chosen compression method. + """ + return _TNAG_COMPRESS_METHODS[method]( + tn, + max_bond=max_bond, + cutoff=cutoff, + site_tags=site_tags, + canonize=canonize, + optimize=optimize, + inplace=inplace, + **kwargs, + ) diff --git a/quimb/tensor/tensor_builder.py b/quimb/tensor/tensor_builder.py index a42c1272..5c3962b4 100644 --- a/quimb/tensor/tensor_builder.py +++ b/quimb/tensor/tensor_builder.py @@ -3605,8 +3605,10 @@ def MPS_rand_state( phys_dim=2, normalize=True, cyclic=False, - dtype="float64", dist="normal", + loc=0.0, + scale=1.0, + dtype="float64", trans_invar=False, **mps_opts, ): @@ -3625,6 +3627,12 @@ def MPS_rand_state( cyclic : bool, optional Generate a MPS with periodic boundary conditions or not, default is open boundary conditions. + dist : {'normal', 'uniform', 'rademacher', 'exp'}, optional + Type of random number to generate, defaults to 'normal'. + loc : float, optional + An additive offset to add to the random numbers. + scale : float, optional + A multiplicative factor to scale the random numbers by. dtype : {float, complex} or numpy dtype, optional Data type of the tensor network. trans_invar : bool (optional) @@ -3633,6 +3641,8 @@ def MPS_rand_state( mps_opts Supplied to :class:`~quimb.tensor.tensor_1d.MatrixProductState`. """ + randn_opts = {"dist": dist, "loc": loc, "scale": scale, "dtype": dtype} + if trans_invar: if not cyclic: raise ValueError( @@ -3640,11 +3650,7 @@ def MPS_rand_state( "boundary conditions." ) array = sensibly_scale( - randn( - shape=(bond_dim, bond_dim, phys_dim), - dtype=dtype, - dist=dist, - ) + randn(shape=(bond_dim, bond_dim, phys_dim), **randn_opts) ) def fill_fn(shape): @@ -3653,7 +3659,7 @@ def fill_fn(shape): else: def fill_fn(shape): - return sensibly_scale(randn(shape, dtype=dtype, dist=dist)) + return sensibly_scale(randn(shape, **randn_opts)) mps = MatrixProductState.from_fill_fn( fill_fn, @@ -4034,6 +4040,8 @@ def MPO_rand( herm=False, dtype="float64", dist="normal", + loc=0.0, + scale=1.0, **mpo_opts, ): """Generate a random matrix product state. @@ -4055,6 +4063,10 @@ def MPO_rand( Data type of the tensor network. dist : {'normal', 'uniform', 'rademacher', 'exp'}, optional Type of random number to generate, defaults to 'normal'. + loc : float, optional + An additive offset to add to the random numbers. + scale : float, optional + A multiplicative factor to scale the random numbers by. herm : bool, optional Whether to make the matrix hermitian (or symmetric if real) or not. mpo_opts @@ -4069,7 +4081,7 @@ def MPO_rand( ] def gen_data(shape): - data = randn(shape, dtype=dtype, dist=dist) + data = randn(shape, dtype=dtype, dist=dist, loc=loc, scale=scale) if not herm: return data diff --git a/quimb/tensor/tensor_core.py b/quimb/tensor/tensor_core.py index b84d807d..876c300a 100644 --- a/quimb/tensor/tensor_core.py +++ b/quimb/tensor/tensor_core.py @@ -389,7 +389,7 @@ def _parse_split_opts(method, cutoff, absorb, max_bond, cutoff_mode, renorm): opts["cutoff_mode"] = _CUTOFF_MODES[cutoff_mode] # renorm doubles up as the power used to renormalize - if (method in _FULL_SPLIT_METHODS) and (renorm is None): + if (method in _FULL_SPLIT_METHODS) and (renorm is True): opts["renorm"] = _RENORM_LOOKUP.get(cutoff_mode, 0) else: opts["renorm"] = 0 if renorm is None else renorm @@ -658,6 +658,44 @@ def tensor_canonize_bond( tn.gauge_simple_remove(outer=outer) +def choose_local_compress_gauge_settings( + canonize=True, + tree_gauge_distance=None, + canonize_distance=None, + canonize_after_distance=None, + mode="auto", +): + """Choose default gauge settings for arbitrary geometry local compression. + """ + if tree_gauge_distance is None: + if canonize: + # default to r=3 gauge + tree_gauge_distance = 3 + else: + tree_gauge_distance = 0 + + if mode == "auto": + if tree_gauge_distance == 0: + # equivalent to basic mode anyway + mode = "basic" + else: + mode = "virtual-tree" + + if canonize_distance is None: + # default to the tree gauge distance + canonize_distance = tree_gauge_distance + + if canonize_after_distance is None: + if mode == "virtual-tree": + # can avoid resetting the tree gauge + canonize_after_distance = 0 + elif mode == "basic": + # do an eager tree guage and reset + canonize_after_distance = tree_gauge_distance + + return canonize_distance, canonize_after_distance, mode + + def tensor_compress_bond( T1, T2, @@ -6033,22 +6071,16 @@ def _compress_between_tids( **canonize_opts, ) - if mode == 'basic': - tensor_compress_bond( - ta, tb, **compress_opts - ) - elif mode == 'virtual-tree': + if mode == "basic": + tensor_compress_bond(ta, tb, **compress_opts) + elif mode == "virtual-tree": self._compress_between_virtual_tree_tids( tid1, tid2, **compress_opts ) - elif mode == 'full-bond': - self._compress_between_full_bond_tids( - tid1, tid2, **compress_opts - ) - elif mode == 'local-fit': - self._compress_between_local_fit( - tid1, tid2, **compress_opts - ) + elif mode == "full-bond": + self._compress_between_full_bond_tids(tid1, tid2, **compress_opts) + elif mode == "local-fit": + self._compress_between_local_fit(tid1, tid2, **compress_opts) else: # assume callable mode(self, tid1, tid2, **compress_opts) @@ -6141,18 +6173,83 @@ def compress_between( **compress_opts, ) - def compress_all(self, inplace=False, **compress_opts): - """Inplace compress all bonds in this network.""" + def compress_all( + self, + max_bond=None, + cutoff=1e-10, + canonize=True, + tree_gauge_distance=None, + canonize_distance=None, + canonize_after_distance=None, + mode="auto", + inplace=False, + **compress_opts + ): + """Compress all bonds one by one in this network. + + Parameters + ---------- + max_bond : int or None, optional + The maxmimum bond dimension to compress to. + cutoff : float, optional + The singular value cutoff to use. + tree_gauge_distance : int, optional + How far to include local tree gauge information when compressing. + If the local geometry is a tree, then each compression will be + locally optimal up to this distance. + canonize_distance : int, optional + How far to locally canonize around the target tensors first, this + is set automatically by ``tree_gauge_distance`` if not specified. + canonize_after_distance : int, optional + How far to locally canonize around the target tensors after, this + is set automatically by ``tree_gauge_distance``, depending on + ``mode`` if not specified. + mode : {'auto', 'basic', 'virtual-tree'}, optional + The mode to use for compressing the bonds. If 'auto', will use + 'basic' if ``tree_gauge_distance == 0`` else 'virtual-tree'. + inplace : bool, optional + Whether to perform the compression inplace. + compress_opts + Supplied to + :func:`~quimb.tensor.tensor_core.TensorNetwork.compress_between`. + + Returns + ------- + TensorNetwork + + See Also + -------- + compress_between, canonize_all + """ tn = self if inplace else self.copy() - tn.fuse_multibonds_() + canonize_distance, canonize_after_distance, mode = ( + choose_local_compress_gauge_settings( + canonize, + tree_gauge_distance, + canonize_distance, + canonize_after_distance, + mode, + ) + ) + + tn.fuse_multibonds_() for ix in tuple(tn.ind_map): try: tid1, tid2 = tn.ind_map[ix] except (ValueError, KeyError): # not a bond, or index already compressed away continue - tn._compress_between_tids(tid1, tid2, **compress_opts) + tn._compress_between_tids( + tid1, + tid2, + max_bond=max_bond, + cutoff=cutoff, + mode=mode, + canonize_distance=canonize_distance, + canonize_after_distance=canonize_after_distance, + **compress_opts + ) return tn