From 847e0e908f8b5817479dee70ee348c2abf52620a Mon Sep 17 00:00:00 2001 From: Johnnie Gray Date: Wed, 12 Jun 2024 11:24:21 -0700 Subject: [PATCH] abstract some SU related functionality (for block/fermionic) --- docs/changelog.md | 7 +- pyproject.toml | 1 - quimb/tensor/array_ops.py | 25 +- quimb/tensor/tensor_2d.py | 355 ++++++++++++++-------------- quimb/tensor/tensor_2d_tebd.py | 4 +- quimb/tensor/tensor_arbgeom.py | 2 +- quimb/tensor/tensor_arbgeom_tebd.py | 5 +- quimb/tensor/tensor_core.py | 20 +- 8 files changed, 224 insertions(+), 195 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index 576b07fe..b1cf5827 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -3,7 +3,7 @@ Release notes for `quimb`. (whats-new-1-8-2)= -## v1.8.2 (unreleased) +## v1.8.2 (2024-06-12) **Enhancements:** @@ -13,7 +13,10 @@ Release notes for `quimb`. - Update generic TN optimizer docs. - add [`tn.gen_inds_loops`](quimb.tensor.tensor_core.TensorNetwork.gen_inds_loops) for generating all loops of indices in a TN. - add [`tn.gen_inds_connected`](quimb.tensor.tensor_core.TensorNetwork.gen_inds_connected) for generating all connected sets of indices in a TN. - +- make SVD fallback error catching more generic ({pull}`#238`) +- fix some windows + numba CI issues. +- [`approx_spectral_function`](quimb.linalg.approx_spectral.approx_spectral_function) add plotting and tracking +- add dispatching to various tensor primitives to allow overriding (whats-new-1-8-1)= ## v1.8.1 (2024-05-06) diff --git a/pyproject.toml b/pyproject.toml index 626d79b8..a4f7587c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,6 @@ write_to = "quimb/_version.py" [tool.pytest.ini_options] testpaths = "tests" -addopts = "--cov=quimb --cov-report term-missing --cov-report xml:coverage.xml --verbose --durations=10" filterwarnings = "once" [tool.coverage.run] diff --git a/quimb/tensor/array_ops.py b/quimb/tensor/array_ops.py index 5da3399d..04aa57b4 100644 --- a/quimb/tensor/array_ops.py +++ b/quimb/tensor/array_ops.py @@ -165,6 +165,25 @@ def ndim(array): return len(array.shape) +@compose +def multiply_diagonal(x, v, axis, backend=None): + """Multiply v into x as if contracting in a diagonal matrix.""" + newshape = tuple((-1 if i == axis else 1) for i in range(ndim(x))) + v_broadcast = do("reshape", v, newshape, like=backend) + return x * v_broadcast + + +@compose +def align_axes(*arrays, axes, backend=None): + """Prepare a set of arrays that should be contractible along ``axes``. + + For example, block symmetric arrays need to have aligned sectors prior to + fusing. + """ + # default implementation is nothing + return arrays + + # ------------- miscelleneous other backend agnostic functions -------------- # @@ -281,7 +300,11 @@ def find_diag_axes(x, atol=1e-12): if shape[i] != shape[j]: continue if do( - "allclose", x[indxrs[i] != indxrs[j]], zero, atol=atol, like=backend + "allclose", + x[indxrs[i] != indxrs[j]], + zero, + atol=atol, + like=backend, ): return (i, j) return None diff --git a/quimb/tensor/tensor_2d.py b/quimb/tensor/tensor_2d.py index e45f860c..3b1615ec 100644 --- a/quimb/tensor/tensor_2d.py +++ b/quimb/tensor/tensor_2d.py @@ -8,7 +8,7 @@ from numbers import Integral from operator import add -from autoray import do, get_dtype_name, infer_backend +import autoray as ar from ..gen.operators import swap from ..gen.rand import randn, seed_rand @@ -3992,6 +3992,140 @@ def gate_string_reduce_split_( to.modify(data=t.data) +def gate_2d_long_range( + self, + G, + where, + contract=False, + tags=None, + inplace=False, + info=None, + long_range_use_swaps=False, + long_range_path_sequence=None, + **compress_opts, +): + psi = self if inplace else self.copy() + + ng = len(where) + + dp = psi.phys_dim(*where[0]) + tags = tags_to_oset(tags) + + # allow a matrix to be reshaped into a tensor if it factorizes + # i.e. (4, 4) assumed to be two qubit gate -> (2, 2, 2, 2) + G = maybe_factor_gate_into_tensor(G, dp, ng, where) + + site_ix = [psi.site_ind(i, j) for i, j in where] + # new indices to join old physical sites to new gate + bnds = [rand_uuid() for _ in range(ng)] + reindex_map = dict(zip(site_ix, bnds)) + + TG = Tensor(G, inds=site_ix + bnds, tags=tags, left_inds=bnds) + + # following are all based on splitting tensors to maintain structure + ij_a, ij_b = where + + # parse the argument specifying how to find the path between + # non-nearest neighbours + if long_range_path_sequence is not None: + # make sure we can index + long_range_path_sequence = tuple(long_range_path_sequence) + # if the first element is a str specifying move sequence, e.g. + # ('v', 'h') + # ('av', 'bv', 'ah', 'bh') # using swaps + manual_lr_path = not isinstance(long_range_path_sequence[0], str) + # otherwise assume a path has been manually specified, e.g. + # ((1, 2), (2, 2), (2, 3), ... ) + # (((1, 1), (1, 2)), ((4, 3), (3, 3)), ...) # using swaps + else: + manual_lr_path = False + + # check if we are not nearest neighbour and need to swap first + if long_range_use_swaps: + if manual_lr_path: + *swaps, final = long_range_path_sequence + else: + # find a swap path + *swaps, final = gen_long_range_swap_path( + ij_a, ij_b, sequence=long_range_path_sequence + ) + + # move the sites together + SWAP = get_swap( + dp, dtype=ar.get_dtype_name(G), backend=ar.infer_backend(G) + ) + for pair in swaps: + psi.gate_(SWAP, pair, contract=contract, absorb="right") + + compress_opts["info"] = info + compress_opts["contract"] = contract + + # perform actual gate also compressing etc on 'way back' + psi.gate_(G, final, **compress_opts) + + compress_opts.setdefault("absorb", "both") + for pair in reversed(swaps): + psi.gate_(SWAP, pair, **compress_opts) + + return psi + + if manual_lr_path: + string = long_range_path_sequence + else: + string = tuple( + gen_long_range_path(*where, sequence=long_range_path_sequence) + ) + + # the tensors along this string, which will be updated + original_ts = [psi[coo] for coo in string] + + # the len(string) - 1 indices connecting the string + bonds_along = [ + next(iter(bonds(t1, t2))) for t1, t2 in pairwise(original_ts) + ] + + if contract == "split": + # + # │╱ │╱ │╱ │╱ + # ──GGGGG── ==> ──G┄┄┄G── + # ╱ ╱ ╱ ╱ + # + gate_string_split_( + TG, + where, + string, + original_ts, + bonds_along, + reindex_map, + site_ix, + info, + **compress_opts, + ) + + elif contract == "reduce-split": + # + # │ │ │ │ + # GGGGG GGG │ │ + # │╱ │╱ ==> ╱│ │ ╱ ==> ╱│ │ ╱ │╱ │╱ + # ──●───●── ──>─●─●─<── ──>─GGG─<── ==> ──G┄┄┄G── + # ╱ ╱ ╱ ╱ ╱ ╱ ╱ ╱ + # + # + gate_string_reduce_split_( + TG, + where, + string, + original_ts, + bonds_along, + reindex_map, + site_ix, + info, + **compress_opts, + ) + + return psi + + class TensorNetwork2DVector(TensorNetwork2D, TensorNetworkGenVector): """Mixin class for a 2D square lattice vector TN, i.e. one with a single physical index per site. @@ -4147,179 +4281,43 @@ def gate( For one site gates when one of the 'split' methods is supplied ``contract=True`` is assumed. """ - check_opt("contract", contract, (False, True, "split", "reduce-split")) - - psi = self if inplace else self.copy() - if is_lone_coo(where): where = (where,) else: where = tuple(where) - ng = len(where) - - dp = psi.phys_dim(*where[0]) - tags = tags_to_oset(tags) - - # allow a matrix to be reshaped into a tensor if it factorizes - # i.e. (4, 4) assumed to be two qubit gate -> (2, 2, 2, 2) - G = maybe_factor_gate_into_tensor(G, dp, ng, where) - - site_ix = [psi.site_ind(i, j) for i, j in where] - # new indices to join old physical sites to new gate - bnds = [rand_uuid() for _ in range(ng)] - reindex_map = dict(zip(site_ix, bnds)) - - TG = Tensor(G, inds=site_ix + bnds, tags=tags, left_inds=bnds) - if contract is False: - # - # │ │ <- site_ix - # GGGGG - # │╱ │╱ <- bnds - # ──●───●── - # ╱ ╱ - # - if propagate_tags: - if propagate_tags == "register": - old_tags = oset(map(psi.site_tag, where)) - else: - old_tags = oset_union( - psi.tensor_map[tid].tags - for ind in site_ix - for tid in psi.ind_map[ind] - ) - - if propagate_tags == "sites": - # use regex to take tags only matching e.g. 'I4,3' - rex = re.compile(psi.site_tag_id.format(r"\d+", r"\d+")) - old_tags = oset(filter(rex.match, old_tags)) - - TG.modify(tags=TG.tags | old_tags) - - psi.reindex_(reindex_map) - psi |= TG - return psi - - if (contract is True) or (ng == 1): - # - # │╱ │╱ - # ──GGGGG── - # ╱ ╱ - # - psi.reindex_(reindex_map) - - # get the sites that used to have the physical indices - site_tids = psi._get_tids_from_inds(bnds, which="any") - - # pop the sites, contract, then re-add - pts = [psi.pop_tensor(tid) for tid in site_tids] - psi |= tensor_contract(*pts, TG) - - return psi - - # following are all based on splitting tensors to maintain structure - ij_a, ij_b = where - - # parse the argument specifying how to find the path between - # non-nearest neighbours - if long_range_path_sequence is not None: - # make sure we can index - long_range_path_sequence = tuple(long_range_path_sequence) - # if the first element is a str specifying move sequence, e.g. - # ('v', 'h') - # ('av', 'bv', 'ah', 'bh') # using swaps - manual_lr_path = not isinstance(long_range_path_sequence[0], str) - # otherwise assume a path has been manually specified, e.g. - # ((1, 2), (2, 2), (2, 3), ... ) - # (((1, 1), (1, 2)), ((4, 3), (3, 3)), ...) # using swaps - else: - manual_lr_path = False - - # check if we are not nearest neighbour and need to swap first - if long_range_use_swaps: - if manual_lr_path: - *swaps, final = long_range_path_sequence - else: - # find a swap path - *swaps, final = gen_long_range_swap_path( - ij_a, ij_b, sequence=long_range_path_sequence - ) - - # move the sites together - SWAP = get_swap( - dp, dtype=get_dtype_name(G), backend=infer_backend(G) - ) - for pair in swaps: - psi.gate_(SWAP, pair, contract=contract, absorb="right") - - compress_opts["info"] = info - compress_opts["contract"] = contract - - # perform actual gate also compressing etc on 'way back' - psi.gate_(G, final, **compress_opts) - - compress_opts.setdefault("absorb", "both") - for pair in reversed(swaps): - psi.gate_(SWAP, pair, **compress_opts) - - return psi - - if manual_lr_path: - string = long_range_path_sequence - else: - string = tuple( - gen_long_range_path(*where, sequence=long_range_path_sequence) - ) - - # the tensors along this string, which will be updated - original_ts = [psi[coo] for coo in string] - - # the len(string) - 1 indices connecting the string - bonds_along = [ - next(iter(bonds(t1, t2))) for t1, t2 in pairwise(original_ts) - ] + use_2d_long_range = ( + (len(where) == 2) + and (contract in {"split", "reduce-split"}) + and (manhattan_distance(*where) > 1) + ) - if contract == "split": - # - # │╱ │╱ │╱ │╱ - # ──GGGGG── ==> ──G┄┄┄G── - # ╱ ╱ ╱ ╱ - # - gate_string_split_( - TG, - where, - string, - original_ts, - bonds_along, - reindex_map, - site_ix, - info, + if use_2d_long_range: + return gate_2d_long_range( + self, + G=G, + where=where, + contract=contract, + tags=tags, + inplace=inplace, + info=info, + long_range_use_swaps=long_range_use_swaps, + long_range_path_sequence=long_range_path_sequence, **compress_opts, ) - - elif contract == "reduce-split": - # - # │ │ │ │ - # GGGGG GGG │ │ - # │╱ │╱ ==> ╱│ │ ╱ ==> ╱│ │ ╱ │╱ │╱ - # ──●───●── ──>─●─●─<── ──>─GGG─<── ==> ──G┄┄┄G── - # ╱ ╱ ╱ ╱ ╱ ╱ ╱ ╱ - # - # - gate_string_reduce_split_( - TG, - where, - string, - original_ts, - bonds_along, - reindex_map, - site_ix, - info, + else: + # can just use generic arbgeom methods + return super().gate( + G=G, + where=where, + contract=contract, + tags=tags, + propagate_tags=propagate_tags, + inplace=inplace, + info=info, **compress_opts, ) - return psi - gate_ = functools.partialmethod(gate, inplace=True) def compute_norm( @@ -4808,8 +4806,8 @@ def __init__( self._Lx = len(arrays) self._Ly = len(arrays[0]) - cyclicx = sum(d > 1 for d in arrays[0][1].shape) == 5 - cyclicy = sum(d > 1 for d in arrays[1][0].shape) == 5 + cyclicx = sum(d > 1 for d in ar.shape(arrays[0][1])) == 5 + cyclicy = sum(d > 1 for d in ar.shape(arrays[1][0])) == 5 # cache for both creating and retrieving indices ix = defaultdict(rand_uuid) @@ -4831,14 +4829,14 @@ def __init__( array_order = array_order.replace("l", "") # allow convention of missing bonds to be singlet dimensions - if len(array.shape) != len(array_order): - array = do("squeeze", array) + if ar.ndim(array) != len(array_order): + array = ar.do("squeeze", array) transpose_order = tuple( array_order.find(x) for x in "urdlp" if x in array_order ) if transpose_order != tuple(range(len(array_order))): - array = do("transpose", array, transpose_order) + array = ar.do("transpose", array, transpose_order) # get the relevant indices corresponding to neighbours inds = [] @@ -4958,7 +4956,7 @@ def empty(cls, Lx, Ly, bond_dim, phys_dim=2, like="numpy", **peps_opts): PEPS.from_fill_fn """ return cls.from_fill_fn( - lambda shape: do("zeros", shape, like=like), + lambda shape: ar.do("zeros", shape, like=like), Lx, Ly, bond_dim, @@ -4992,7 +4990,7 @@ def ones(cls, Lx, Ly, bond_dim, phys_dim=2, like="numpy", **peps_opts): PEPS.from_fill_fn """ return cls.from_fill_fn( - lambda shape: do("ones", shape, like=like), + lambda shape: ar.do("ones", shape, like=like), Lx, Ly, bond_dim, @@ -5151,8 +5149,8 @@ def __init__( 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 + cyclicx = sum(d > 1 for d in ar.shape(arrays[0][1])) == 6 + cyclicy = sum(d > 1 for d in ar.shape(arrays[1][0])) == 6 for i, j in product(range(self.Lx), range(self.Ly)): array = arrays[i][j] @@ -5170,14 +5168,14 @@ def __init__( array_order = array_order.replace("l", "") # allow convention of missing bonds to be singlet dimensions - if len(array.shape) != len(array_order): - array = do("squeeze", array) + if ar.ndim(array) != len(array_order): + array = ar.do("squeeze", array) transpose_order = tuple( array_order.find(x) for x in "urdlbk" if x in array_order ) if transpose_order != tuple(range(len(array_order))): - array = do("transpose", array, transpose_order) + array = ar.do("transpose", array, transpose_order) # get the relevant indices corresponding to neighbours inds = [] @@ -5318,7 +5316,7 @@ def fill_fn(shape): if herm: new_order = list(range(len(shape))) new_order[-2], new_order[-1] = new_order[-1], new_order[-2] - X = (do("conj", X) + do("transpose", X, new_order)) / 2 + X = (ar.do("conj", X) + ar.do("transpose", X, new_order)) / 2 return X return cls.from_fill_fn( @@ -5364,7 +5362,7 @@ def zeros( """ def fill_fn(shape): - return do("zeros", shape, dtype=dtype, like=backend) + return ar.do("zeros", shape, dtype=dtype, like=backend) return cls.from_fill_fn( fill_fn, @@ -5764,4 +5762,5 @@ def swap_path_to_long_range_path(swap_path, ij_a): @functools.lru_cache(8) def get_swap(dp, dtype, backend): SWAP = swap(dp, dtype=dtype) - return do("array", SWAP, like=backend) + return ar.do("array", SWAP, like=backend) + diff --git a/quimb/tensor/tensor_2d_tebd.py b/quimb/tensor/tensor_2d_tebd.py index aff190d8..a0a9772a 100644 --- a/quimb/tensor/tensor_2d_tebd.py +++ b/quimb/tensor/tensor_2d_tebd.py @@ -604,10 +604,10 @@ def env_neighbours(i, j): # set the new singualar values all along the chain for site_a, site_b in pairwise(string): bond_pair = tuple(sorted((site_a, site_b))) - s = info['singular_values', bond_pair] + s, = info.values() if self.gauge_renorm: # keep the singular values from blowing up - s = s / s[0] + s = s / do("max", s) Tsval = self.gauges[bond_pair] Tsval.modify(data=s) diff --git a/quimb/tensor/tensor_arbgeom.py b/quimb/tensor/tensor_arbgeom.py index 2ffbea8c..700674d6 100644 --- a/quimb/tensor/tensor_arbgeom.py +++ b/quimb/tensor/tensor_arbgeom.py @@ -1038,7 +1038,7 @@ def gate_simple_(self, G, where, gauges, renorm=True, **gate_opts): # inner ungauging is performed by tracking the new singular values (((_, ix), s),) = info.items() if renorm: - s = s / s[0] + s = s / do("max", s) gauges[ix] = s return self diff --git a/quimb/tensor/tensor_arbgeom_tebd.py b/quimb/tensor/tensor_arbgeom_tebd.py index a79768e4..844eeb84 100644 --- a/quimb/tensor/tensor_arbgeom_tebd.py +++ b/quimb/tensor/tensor_arbgeom_tebd.py @@ -192,8 +192,9 @@ def _expm_cached(self, x, y): cache = self._op_cache["expm"] key = (id(x), y) if key not in cache: - el, ev = do("linalg.eigh", x) - cache[key] = ev @ do("diag", do("exp", el * y)) @ dag(ev) + cache[key] = do("scipy.linalg.expm", x * y) + # el, ev = do("linalg.eigh", x) + # cache[key] = ev @ do("diag", do("exp", el * y)) @ dag(ev) return cache[key] def get_gate(self, where): diff --git a/quimb/tensor/tensor_core.py b/quimb/tensor/tensor_core.py index b8b8deee..39f6d7a6 100644 --- a/quimb/tensor/tensor_core.py +++ b/quimb/tensor/tensor_core.py @@ -561,10 +561,11 @@ def tensor_split( # ``s`` itself will be None unless ``absorb=None`` is specified left, s, right = _SPLIT_FNS[method](array, **opts) + if len(left_dims) != 1: - left = do("reshape", left, (*left_dims, -1)) + left = do("reshape", left, (*left_dims, shape(left)[-1])) if len(right_dims) != 1: - right = do("reshape", right, (-1, *right_dims)) + right = do("reshape", right, (shape(right)[0], *right_dims)) if get == "arrays": if absorb is None: @@ -2789,10 +2790,8 @@ def multiply_index_diagonal(self, ind, x, inplace=False): tensor being contracted into index ``ind``. """ t = self if inplace else self.copy() - x_broadcast = do( - "reshape", x, tuple((-1 if i == ind else 1) for i in t.inds) - ) - t.modify(data=t.data * x_broadcast) + ax = t.inds.index(ind) + t.modify(data=do("multiply_diagonal", t.data, x, axis=ax)) return t multiply_index_diagonal_ = functools.partialmethod( @@ -6966,7 +6965,7 @@ def gauge_all_simple( ) s = info["singular_values"] - smax = s[0] + smax = do("max", s) new_gauge = s / smax nfact = do("log10", smax) + nfact @@ -7193,7 +7192,12 @@ def gauge_simple_insert( g = _get(ix, None) if g is None: continue - g = (g + smudge * g[0]) ** power + + if smudge != 0.0: + g = g + smudge * do("max", g) + if power != 1.0: + g = g**power + (t,) = self._inds_get(ix) t.multiply_index_diagonal_(ix, g) outer.append((t, ix, g))