From 349dcdfafa69f8b562da1203a70ecc50d86ad158 Mon Sep 17 00:00:00 2001 From: Johnnie Gray Date: Fri, 15 Nov 2024 22:24:34 -0800 Subject: [PATCH] PEPS.product_state and other constructor options --- docs/changelog.md | 2 + quimb/gen/rand.py | 40 +++++- quimb/tensor/tensor_2d.py | 216 +++++++++++++++++++++++++++------ quimb/tensor/tensor_builder.py | 50 ++------ quimb/tensor/tensor_core.py | 1 - 5 files changed, 228 insertions(+), 81 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index 385be185..55167e20 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -22,6 +22,8 @@ Release notes for `quimb`. - add `fit-zipup` and `fit-projector` shorthand methods to the general 1d tensor network compression function - add [`MatrixProductState.compute_local_expectation`](quimb.tensor.tensor_1d.MatrixProductState.compute_local_expectation) for computing many local expectations for a MPS at once, to match the interface for this method elsewhere. These can either be computed via canonicalization (`method="canonical"`), or via explicit left and right environment contraction (`method="envs"`) - specialize [`CircuitMPS.local_expectation`](quimb.tensor.circuit.CircuitMPS.local_expectation) to make use of the MPS form. +- add [`PEPS.product_state`](quimb.tensor.tensor_2d.PEPS.product_state) for constructing a PEPS representing a product state. +- add [`PEPS.vacuum`](quimb.tensor.tensor_2d.PEPS.vacuum) for constructing a PEPS representing the vacuum state $|000\ldots0\rangle$. --- diff --git a/quimb/gen/rand.py b/quimb/gen/rand.py index 77e6877d..6bc07309 100644 --- a/quimb/gen/rand.py +++ b/quimb/gen/rand.py @@ -1,5 +1,5 @@ -"""Functions for generating random quantum objects and states. -""" +"""Functions for generating random quantum objects and states.""" + import math import random from functools import wraps @@ -340,6 +340,42 @@ def rand_phase(shape, scale=1, dtype=complex): return z +def get_rand_fill_fn( + dist="normal", + loc=0.0, + scale=1.0, + seed=None, + dtype="float64", +): + """Get a callable with the given random distribution and parameters, that + has signature ``fill_fn(shape) -> array``. + + Parameters + ---------- + 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. + seed : int, optional + A random seed. + dtype : {'float64', 'complex128', 'float32', 'complex64'}, optional + The underlying data type. + + Returns + ------- + callable + """ + if seed is not None: + seed_rand(seed) + + def fill_fn(shape=()): + return randn(shape, dtype=dtype, dist=dist, loc=loc, scale=scale) + + return fill_fn + + def rand_matrix( d, scaled=True, diff --git a/quimb/tensor/tensor_2d.py b/quimb/tensor/tensor_2d.py index 1d501e17..cba0841c 100644 --- a/quimb/tensor/tensor_2d.py +++ b/quimb/tensor/tensor_2d.py @@ -10,7 +10,7 @@ import autoray as ar from ..gen.operators import swap -from ..gen.rand import randn, seed_rand +from ..gen.rand import get_rand_fill_fn from ..utils import ( check_opt, deprecated, @@ -19,7 +19,6 @@ 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 ( @@ -2501,8 +2500,9 @@ def contract_boundary( contraction. canonize : bool, optional Whether to sweep one way with canonization before compressing. - mode : {'mps', 'full-bond'}, optional - How to perform the compression on the boundary. + mode : {'mps', 'full-bond', ...}, optional + How to perform the compression on the boundary, can also be any of + the generic 1D or arbgeom methods. layer_tags : None or sequence of str, optional If given, perform a multilayer contraction, contracting the inner sites in each layer into the boundary individually. @@ -4320,12 +4320,72 @@ def gate( def compute_norm( self, + max_bond=None, + *, + cutoff=1e-10, + canonize=True, + mode="mps", layer_tags=("KET", "BRA"), + compress_opts=None, + sequence=None, + equalize_norms=False, + progbar=None, **contract_opts, ): - """Compute the norm of this vector via boundary contraction.""" + """Compute the norm of this vector via boundary contraction. + + Parameters + ---------- + max_bond : int, optional + The maximum boundary dimension, AKA 'chi'. The default of ``None`` + means truncation is left purely to ``cutoff`` and is not + recommended in 2D. + cutoff : float, optional + Cut-off value to used to truncate singular values in the boundary + contraction. + canonize : bool, optional + Whether to sweep one way with canonization before compressing. + mode : {'mps', 'full-bond', ...}, optional + How to perform the compression on the boundary, can also be any of + the generic 1D or arbgeom methods. + layer_tags : None or sequence of str, optional + If given, perform a multilayer contraction, contracting the inner + sites in each layer into the boundary individually. + compress_opts : None or dict, optional + Other low level options to pass to + :meth:`~quimb.tensor.tensor_core.TensorNetwork.compress_between`. + sequence : sequence of {'xmin', 'xmax', 'ymin', 'ymax'}, optional + Which directions to cycle throught when performing the inwards + contractions, i.e. *from* that direction. If ``around`` is + specified you will likely need all of these! Default is to contract + from the two shortest opposing sides. + equalize_norms : bool or float, optional + Whether to equalize the norms of the boundary tensors after each + contraction, gathering the overall scaling coefficient, log10, in + ``tn.exponent``. + progbar : bool, optional + Whether to show a progress bar. + contract_opts + Additional options to pass to :meth:`contract_boundary`. + + Returns + ------- + scalar + """ norm = self.make_norm(layer_tags=layer_tags) - return norm.contract_boundary(layer_tags=layer_tags, **contract_opts) + return norm.contract_boundary( + max_bond=max_bond, + cutoff=cutoff, + canonize=canonize, + mode=mode, + layer_tags=layer_tags, + compress_opts=compress_opts, + sequence=sequence, + equalize_norms=equalize_norms, + progbar=progbar, + inplace=True, + **contract_opts + ) def compute_local_expectation( self, @@ -4804,19 +4864,20 @@ def __init__( self._Lx = len(arrays) self._Ly = len(arrays[0]) - cyclicx = ( - sum(d > 1 for d in ar.shape(arrays[0][1])) == 5 - ) or ( + shape_on_xmin_edge = ar.shape(arrays[0][self._Ly // 2]) + ndim_xmin_edge = len(shape_on_xmin_edge) + shape_on_ymin_edge = ar.shape(arrays[self._Lx // 2][0]) + ndim_ymin_edge = len(shape_on_ymin_edge) + + cyclicx = (sum(d > 1 for d in shape_on_xmin_edge) == 5) or ( # handle D=1 PBC case - (ar.ndim(arrays[0][1]) == 5) and - (sum(d == 1 for d in ar.shape(arrays[0][1])) == 4) + (ndim_xmin_edge == 5) + and (sum(d == 1 for d in shape_on_xmin_edge) == 4) ) - cyclicy = ( - sum(d > 1 for d in ar.shape(arrays[1][0])) == 5 - ) or ( + cyclicy = (sum(d > 1 for d in shape_on_ymin_edge) == 5) or ( # handle D=1 PBC case - (ar.ndim(arrays[1][0]) == 5) and - (sum(d == 1 for d in ar.shape(arrays[1][0])) == 4) + (ndim_ymin_edge == 5) + and (sum(d == 1 for d in shape_on_ymin_edge) == 4) ) # cache for both creating and retrieving indices @@ -5051,6 +5112,7 @@ def rand( phys_dim=2, dist="normal", loc=0.0, + scale=1.0, dtype="float64", seed=None, **peps_opts, @@ -5086,20 +5148,97 @@ def rand( -------- PEPS.from_fill_fn """ - if seed is not None: - seed_rand(seed) - - def fill_fn(shape): - return ops.sensibly_scale( - ops.sensibly_scale( - randn(shape, dist=dist, loc=loc, dtype=dtype) - ) - ) + fill_fn = get_rand_fill_fn( + dist=dist, + loc=loc, + scale=scale, + dtype=dtype, + seed=seed, + ) return cls.from_fill_fn( fill_fn, Lx, Ly, bond_dim, phys_dim, **peps_opts ) + @classmethod + def product_state(cls, site_map, cyclic=False, **peps_opts): + """Create a PEPS representing a product state, with explicit bonds of + dimension 1 between sites. + + Parameters + ---------- + site_map : dict[tuple[int, int], array] or Sequence[Sequence[array]] + A mapping of site coordinates to physical vectors, or a 2D array of + physical vectors. Each vector being a single site state. + cyclic : bool or tuple[bool, bool], optional + Whether the lattice is cyclic in the x and y directions. + peps_opts + Supplied to :class:`~quimb.tensor.tensor_2d.PEPS`. + + Returns + ------- + PEPS + """ + try: + cyclicx, cyclicy = cyclic + except (TypeError, ValueError): + cyclicx = cyclicy = cyclic + + if isinstance(site_map, dict): + getarray = site_map.get + Lx = max(i for i, j in site_map.keys()) + 1 + Ly = max(j for i, j in site_map.keys()) + 1 + else: + + def getarray(ij): + i, j = ij + return site_map[i][j] + + Lx = len(site_map) + Ly = len(site_map[0]) + + arrays = [[None for _ in range(Ly)] for _ in range(Lx)] + + for i in range(Lx): + for j in range(Ly): + bond_shape = [] + if cyclicx or (i < Lx - 1): # bond up + bond_shape.append(1) + if cyclicy or (j < Ly - 1): # bond right + bond_shape.append(1) + if cyclicx or (i > 0): # bond down + bond_shape.append(1) + if cyclicy or (j > 0): # bond left + bond_shape.append(1) + + ary = getarray((i, j)) + new_shape = (*bond_shape, *ar.shape(ary)) + arrays[i][j] = ar.do("reshape", ary, new_shape) + + return cls(arrays, **peps_opts) + + @classmethod + def vacuum(cls, Lx, Ly, phys_dim=2, **peps_opts): + """Create the 'vaccum' state PEPS, i.e. |00...0>. + + Parameters + ---------- + Lx : int + The number of rows. + Ly : int + The number of columns. + phys_dim : int, optional + The physical index dimension. + peps_opts + Supplied to :class:`~quimb.tensor.tensor_2d.PEPS.product_state`. + + Returns + ------- + """ + data = ar.do("array", [1.0] + [0.0] * (phys_dim - 1)) + site_map = {(i, j): data for i in range(Lx) for j in range(Ly)} + return cls.product_state(site_map, **peps_opts) + def add_PEPS(self, other, inplace=False): return tensor_network_ag_sum(self, other, inplace=inplace) @@ -5318,6 +5457,7 @@ def rand( herm=False, dist="normal", loc=0.0, + scale=1.0, dtype="float64", seed=None, **pepo_opts, @@ -5348,20 +5488,24 @@ def rand( ------- X : PEPO """ - if seed is not None: - seed_rand(seed) + fill_fn = get_rand_fill_fn( + dist=dist, + loc=loc, + scale=scale, + dtype=dtype, + seed=seed, + ) - def fill_fn(shape): - X = ops.sensibly_scale( - ops.sensibly_scale( - randn(shape, dist=dist, loc=loc, dtype=dtype) - ) - ) - if herm: + if herm: + _fill_fn_orig = fill_fn + + def fill_fn(shape): + X = _fill_fn_orig(shape) new_order = list(range(len(shape))) new_order[-2], new_order[-1] = new_order[-1], new_order[-2] - X = (ar.do("conj", X) + ar.do("transpose", X, new_order)) / 2 - return X + return ( + ar.do("conj", X) + ar.do("transpose", X, new_order) + ) / 2 return cls.from_fill_fn( fill_fn, diff --git a/quimb/tensor/tensor_builder.py b/quimb/tensor/tensor_builder.py index 53ffdd09..0f131573 100644 --- a/quimb/tensor/tensor_builder.py +++ b/quimb/tensor/tensor_builder.py @@ -16,7 +16,13 @@ ham_heis, spin_operator, ) -from ..gen.rand import choice, rand_phase, randn, random_seed_fn, seed_rand +from ..gen.rand import ( + choice, + rand_phase, + randn, + random_seed_fn, + get_rand_fill_fn, +) from ..utils import concat, deprecated, unique from .array_ops import asarray, do, reshape, sensibly_scale from .contraction import array_contract @@ -68,42 +74,6 @@ def delta_array(shape, dtype="float64"): return x -def get_rand_fill_fn( - dist="normal", - loc=0.0, - scale=1.0, - seed=None, - dtype="float64", -): - """Get a callable with the given random distribution and parameters, that - has signature ``fill_fn(shape) -> array``. - - Parameters - ---------- - 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. - seed : int, optional - A random seed. - dtype : {'float64', 'complex128', 'float32', 'complex64'}, optional - The underlying data type. - - Returns - ------- - callable - """ - if seed is not None: - seed_rand(seed) - - def fill_fn(shape=()): - return randn(shape, dtype=dtype, dist=dist, loc=loc, scale=scale) - - return fill_fn - - @random_seed_fn def rand_tensor( shape, @@ -600,11 +570,7 @@ def TN_from_edges_rand( ------- TensorNetworkGen, TensorNetworkGenVector or TensorNetworkGenOperator """ - if seed is not None: - seed_rand(seed) - - def fill_fn(shape): - return randn(shape, dtype=dtype, **randn_opts) + fill_fn = get_rand_fill_fn(seed=seed, dtype=dtype, **randn_opts) return TN_from_edges_and_fill_fn( fill_fn=fill_fn, diff --git a/quimb/tensor/tensor_core.py b/quimb/tensor/tensor_core.py index 5e6615d2..7d9f65f1 100644 --- a/quimb/tensor/tensor_core.py +++ b/quimb/tensor/tensor_core.py @@ -63,7 +63,6 @@ array_contract_pathinfo, array_contract_tree, get_contract_backend, - get_contract_strategy, get_symbol, get_tensor_linop_backend, inds_to_eq,