Tensor: improve ind expansion options

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
 
 

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

quimb/tensor/tensor_core.py

@@ -1898,28 +1898,101 @@ 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
         ----------
         name : str
             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,20 +2005,44 @@ 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)
 
         # list.insert has different behavior to expand_dims for -ve. axis
         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

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)