diff --git a/quimb/tensor/tensor_arbgeom.py b/quimb/tensor/tensor_arbgeom.py index 43c53a60..5657ebc8 100644 --- a/quimb/tensor/tensor_arbgeom.py +++ b/quimb/tensor/tensor_arbgeom.py @@ -10,13 +10,13 @@ from ..utils import check_opt, deprecated, ensure_dict from ..utils import progbar as Progbar from .tensor_core import ( - get_symbol, group_inds, oset, rand_uuid, tags_to_oset, TensorNetwork, ) +from .contraction import get_symbol from . import decomp diff --git a/quimb/tensor/tensor_builder.py b/quimb/tensor/tensor_builder.py index b08280c2..8647f14d 100644 --- a/quimb/tensor/tensor_builder.py +++ b/quimb/tensor/tensor_builder.py @@ -490,6 +490,7 @@ def TN_rand_reg( dtype="float64", site_tag_id="I{}", site_ind_id="k{}", + **randn_opts, ): """Create a random regular tensor network. @@ -527,6 +528,7 @@ def TN_rand_reg( dtype=dtype, site_tag_id=site_tag_id, site_ind_id=site_ind_id, + **randn_opts, ) @@ -538,6 +540,7 @@ def TN_rand_tree( dtype="float64", site_tag_id="I{}", site_ind_id="k{}", + **randn_opts, ): """Create a random tree tensor network. @@ -572,6 +575,7 @@ def TN_rand_tree( dtype=dtype, site_tag_id=site_tag_id, site_ind_id=site_ind_id, + **randn_opts, ) @@ -2758,7 +2762,7 @@ def cnf_file_parse(fname): with open(fname, "r") as f: for line in f: args = line.split() - + # ignore empty lines, other comments and info line if (not args) or (args == ["0"]) or (args[0][0] in "c%"): continue diff --git a/quimb/tensor/tensor_core.py b/quimb/tensor/tensor_core.py index 90c84d71..dee6efa7 100644 --- a/quimb/tensor/tensor_core.py +++ b/quimb/tensor/tensor_core.py @@ -1898,9 +1898,16 @@ def add_tag(self, tag): # TODO: make this more efficient with inplace |= ? self.modify(tags=itertools.chain(self.tags, tags)) - def expand_ind(self, ind, size): + def expand_ind( + self, + ind, + size, + mode=None, + rand_strength=None, + rand_dist="normal", + ): """Inplace increase the size of the dimension of ``ind``, the new array - entries will be filled with zeros. + entries will be filled with zeros by default. Parameters ---------- @@ -1908,18 +1915,84 @@ def expand_ind(self, ind, size): Name of the index to expand. size : int, optional Size of the expanded index. + mode : {None, 'zeros', 'repeat', 'random'}, optional + How to fill any new array entries. If ``'zeros'`` then fill with + zeros, if ``'repeat'`` then repeatedly tile the existing entries. + If ``'random'`` then fill with random entries drawn from + ``rand_dist``, multiplied by ``rand_strength``. If ``None`` then + select from zeros or random depening on non-zero ``rand_strength``. + rand_strength : float, optional + If ``mode='random'``, a multiplicative scale for the random + entries, defaulting to 1.0. If ``mode is None`` then supplying a + non-zero value here triggers ``mode='random'``. + rand_dist : {'normal', 'uniform', 'exp'}, optional + If ``mode='random'``, the distribution to draw the random entries + from. """ if ind not in self.inds: raise ValueError(f"Tensor has no index '{ind}'.") size_current = self.ind_size(ind) - pads = [ - (0, size - size_current) if i == ind else (0, 0) - for i in self.inds - ] - self.modify(data=do('pad', self.data, pads, mode='constant')) + if size_current >= size: + # nothing to do + return - def new_ind(self, name, size=1, axis=0): + # auto select mode + if mode is None: + if (rand_strength is not None) and (rand_strength != 0.0): + mode = "random" + else: + mode = "zeros" + + if mode == "zeros": + pads = [ + (0, size - size_current) if i == ind else (0, 0) + for i in self.inds + ] + new_data = do("pad", self.data, pads, mode="constant") + + elif mode == "repeat": + num_repeats = size // size_current + if size % size_current != 0: + raise ValueError( + f"Cannot expand index '{ind}' to size {size} by repeating " + f"the existing entries as this is not an integer multiple " + f"of the current size {size_current}." + ) + axis = self.inds.index(ind) + new_data = do("concatenate", (self.data,) * num_repeats, axis=axis) + + elif mode == "random": + + if rand_strength is None: + # assume if "random" mode selected then want non-zero strength + rand_strength = 1.0 + + axis = self.inds.index(ind) + rand_shape = list(self.shape) + rand_shape[axis] = size - size_current + rand_data = randn( + shape=tuple(rand_shape), + dtype=self.dtype, + dist=rand_dist, + scale=rand_strength, + ) + new_data = do("concatenate", (self.data, rand_data), axis=axis) + + else: + raise ValueError(f"Invalid mode '{mode}'.") + + self.modify(data=new_data) + + def new_ind( + self, + name, + size=1, + axis=0, + mode=None, + rand_strength=None, + rand_dist="normal", + ): """Inplace add a new index - a named dimension. If ``size`` is specified to be greater than one then the new array entries will be filled with zeros. @@ -1932,6 +2005,23 @@ def new_ind(self, name, size=1, axis=0): Size of the new index. axis : int, optional Position of the new index. + mode : {None, 'zeros', 'repeat', 'random'}, optional + How to fill any new array entries. If ``'zeros'`` then fill with + zeros, if ``'repeat'`` then repeatedly tile the existing entries. + If ``'random'`` then fill with random entries drawn from + ``rand_dist``, multiplied by ``rand_strength``. If ``None`` then + select from zeros or random depening on non-zero ``rand_strength``. + rand_strength : float, optional + If ``mode='random'``, a multiplicative scale for the random + entries, defaulting to 1.0. If ``mode is None`` then supplying a + non-zero value here triggers ``mode='random'``. + rand_dist : {'normal', 'uniform', 'exp'}, optional + If ``mode='random'``, the distribution to draw the random entries + from. + + See Also + -------- + Tensor.expand_ind, new_bond """ new_inds = list(self.inds) @@ -1939,13 +2029,20 @@ def new_ind(self, name, size=1, axis=0): if axis < 0: axis = len(new_inds) + axis + 1 + # initially create size-1 index / dimension new_inds.insert(axis, name) - new_data = do("expand_dims", self.data, axis=axis) - self.modify(data=new_data, inds=new_inds) + if size > 1: - self.expand_ind(name, size, mode=mode) + # tile or pad it to the desired size + self.expand_ind( + ind=name, + size=size, + mode=mode, + rand_strength=rand_strength, + rand_dist=rand_dist, + ) new_bond = new_bond @@ -2048,9 +2145,8 @@ def shared_bond_size(self, other): return bonds_size(self, other) def inner_inds(self): - """ """ - ind_freqs = frequencies(self.inds) - return tuple(i for i in self.inds if ind_freqs[i] == 2) + """Get all indices that appear on two or more tensors.""" + return tuple(self._inner_inds) def transpose(self, *output_inds, inplace=False): """Transpose this tensor - permuting the order of both the data *and* @@ -9089,7 +9185,9 @@ def fuse_multibonds( def expand_bond_dimension( self, new_bond_dim, - rand_strength=0.0, + mode=None, + rand_strength=None, + rand_dist="normal", inds_to_expand=None, inplace=False, ): @@ -9120,33 +9218,41 @@ def expand_bond_dimension( if inds_to_expand is None: # find all 'bonds' - indices connecting two or more tensors - inds_to_expand = set() - for ind, tids in tn.ind_map.items(): - if len(tids) >= 2: - inds_to_expand.add(ind) + inds_to_expand = self._inner_inds else: inds_to_expand = tags_to_oset(inds_to_expand) for t in tn._inds_get(*inds_to_expand): - # perform the array expansions - pads = [ - (0, 0) - if ind not in inds_to_expand - else (0, max(new_bond_dim - d, 0)) - for d, ind in zip(t.shape, t.inds) - ] + for ind in t.inds: + if ind in inds_to_expand: + t.expand_ind( + ind, + new_bond_dim, + mode=mode, + rand_strength=rand_strength, + rand_dist=rand_dist, + ) + ) + ) + else: + edata = do("pad", t.data, pads, mode="constant") - if rand_strength > 0: - edata = do( - "pad", - t.data, - pads, - mode=rand_padder, - rand_strength=rand_strength, + t.modify(data=edata) + ) + else: + edata = do("pad", t.data, pads, mode="constant") + + t.modify(data=edata) + ) ) else: edata = do("pad", t.data, pads, mode="constant") + t.modify(data=edata) + ) + else: + edata = do("pad", t.data, pads, mode="constant") + t.modify(data=edata) return tn diff --git a/tests/test_tensor/test_tensor_core.py b/tests/test_tensor/test_tensor_core.py index d2bf97e4..87ac5dd4 100644 --- a/tests/test_tensor/test_tensor_core.py +++ b/tests/test_tensor/test_tensor_core.py @@ -706,6 +706,33 @@ def test_idxmin(self): assert t.idxmin("abs") == {"a": 0, "b": 2, "c": 3} assert t.idxmax(lambda x: 1 / x) == {"a": 1, "b": 0, "c": 0} + def test_expand_ind(self): + t = Tensor(np.ones((2, 3, 4)), inds=["a", "b", "c"]) + assert t.data.sum() == pytest.approx(24) + + # test zeros mode + t0 = t.copy() + t0.expand_ind("a", size=6, mode="zeros") + assert t0.data.sum() == pytest.approx(24) + + # test tiling mode + tt = t.copy() + tt.expand_ind("a", size=6, mode="repeat") + assert tt.data.sum() == pytest.approx(72) + + # test random mode + tr = t.copy() + tr.expand_ind("a", size=6, mode="random", rand_dist="uniform") + assert 24 <= tr.data.sum() <= 72 + + tr = t.copy() + tr.expand_ind("a", size=6, rand_strength=1.0, rand_dist="uniform") + assert 24 <= tr.data.sum() <= 72 + + tr = t.copy() + tr.expand_ind("a", size=6, rand_strength=-1.0, rand_dist="uniform") + assert -48 <= tr.data.sum() <= 24 + class TestTensorNetwork: def test_combining_tensors(self): @@ -1286,6 +1313,22 @@ def test_subgraphs(_): s1, s2 = tn.subgraphs() assert {s1.num_tensors, s2.num_tensors} == {6, 8} + def test_expand_bond_dimension_zeros(self): + k = MPS_rand_state(10, 7) + k0 = k.copy() + k0.expand_bond_dimension(13) + assert k0.max_bond() == 13 + assert k0.ind_size("k0") == 2 + assert k0.distance(k) == pytest.approx(0.0) + + def test_expand_bond_dimension_random(self): + tn = qtn.TN_rand_reg(6, 3, 2, dist="uniform") + Z = tn ^ ... + # expand w/ positive random entries -> contraction value must increase + tn.expand_bond_dimension_(3, rand_strength=0.1, rand_dist="uniform") + Ze = tn ^ ... + assert Ze > Z + def test_compress_multibond(self): A = rand_tensor((7, 2, 2), "abc", tags="A") A.expand_ind("c", 3)