add tensor_2d_compress and more equalize_norms funcs

quimb/tensor/decomp.py

@@ -23,9 +23,7 @@
 )
 
 from ..core import njit
-from ..linalg import base_linalg
-from ..linalg import rand_linalg
-
+from ..linalg import base_linalg, rand_linalg
 
 _CUTOFF_MODE_MAP = {
     "abs": 1,

quimb/tensor/tensor_1d_compress.py

@@ -188,6 +188,7 @@ def tensor_network_1d_compress_dm(
     optimize="auto-hq",
     sweep_reverse=False,
     canonize=True,
+    equalize_norms=False,
     inplace=False,
     **compress_opts,
 ):
@@ -232,6 +233,10 @@ def tensor_network_1d_compress_dm(
         canonical form instead of right canonical.
     canonize : bool, optional
         Dummy argument to match the signature of other compression methods.
+    equalize_norms : bool or float, optional
+        Whether to equalize the norms of the tensors after compression. If an
+        explicit value is give, then the norms will be set to that value, and
+        the overall scaling factor will be accumulated into `.exponent`.
     inplace : bool, optional
         Whether to perform the compression inplace or not.
     compress_opts
@@ -376,6 +381,12 @@ def tensor_network_1d_compress_dm(
     # possibly put the array indices in canonical order (e.g. when MPS or MPO)
     possibly_permute_(new, permute_arrays)
 
+    # XXX: do better than simply waiting til the end to equalize norms
+    if equalize_norms is True:
+        new.equalize_norms_()
+    elif equalize_norms:
+        new.equalize_norms_(value=equalize_norms)
+
     return new
 
 
@@ -390,6 +401,7 @@ def tensor_network_1d_compress_zipup(
     permute_arrays=True,
     optimize="auto-hq",
     sweep_reverse=False,
+    equalize_norms=False,
     inplace=False,
     **compress_opts,
 ):
@@ -436,6 +448,10 @@ def tensor_network_1d_compress_zipup(
     sweep_reverse : bool, optional
         Whether to sweep in the reverse direction, resulting in a left
         canonical form instead of right canonical.
+    equalize_norms : bool or float, optional
+        Whether to equalize the norms of the tensors after compression. If an
+        explicit value is give, then the norms will be set to that value, and
+        the overall scaling factor will be accumulated into `.exponent`.
     inplace : bool, optional
         Whether to perform the compression inplace or not.
     compress_opts
@@ -541,6 +557,12 @@ def tensor_network_1d_compress_zipup(
     # possibly put the array indices in canonical order (e.g. when MPS or MPO)
     possibly_permute_(new, permute_arrays)
 
+    # XXX: do better than simply waiting til the end to equalize norms
+    if equalize_norms is True:
+        new.equalize_norms_()
+    elif equalize_norms:
+        new.equalize_norms_(value=equalize_norms)
+
     return new
 
 
@@ -557,6 +579,7 @@ def tensor_network_1d_compress_zipup_first(
     permute_arrays=True,
     optimize="auto-hq",
     sweep_reverse=False,
+    equalize_norms=False,
     inplace=False,
     **compress_opts,
 ):
@@ -611,6 +634,10 @@ def tensor_network_1d_compress_zipup_first(
     sweep_reverse : bool, optional
         Whether to sweep in the reverse direction, resulting in a left
         canonical form instead of right canonical.
+    equalize_norms : bool or float, optional
+        Whether to equalize the norms of the tensors after compression. If an
+        explicit value is give, then the norms will be set to that value, and
+        the overall scaling factor will be accumulated into `.exponent`.
     inplace : bool, optional
         Whether to perform the compression inplace or not.
     compress_opts
@@ -650,6 +677,7 @@ def tensor_network_1d_compress_zipup_first(
         cutoff_mode=cutoff_mode,
         optimize=optimize,
         sweep_reverse=True,
+        equalize_norms=equalize_norms,
         inplace=inplace,
         **compress_opts,
     )
@@ -664,6 +692,7 @@ def tensor_network_1d_compress_zipup_first(
             max_bond=max_bond,
             cutoff=cutoff,
             cutoff_mode=cutoff_mode,
+            equalize_norms=equalize_norms,
             **compress_opts,
         )
 
@@ -674,6 +703,11 @@ def tensor_network_1d_compress_zipup_first(
     # possibly put the array indices in canonical order (e.g. when MPS or MPO)
     possibly_permute_(tn, permute_arrays)
 
+    if equalize_norms is True:
+        tn.equalize_norms_()
+    elif equalize_norms:
+        tn.equalize_norms_(value=equalize_norms)
+
     return tn
 
 
@@ -946,6 +980,7 @@ def tensor_network_1d_compress_fit(
     optimize="auto-hq",
     canonize=True,
     sweep_reverse=False,
+    equalize_norms=False,
     inplace_fit=False,
     inplace=False,
     progbar=False,
@@ -1021,12 +1056,16 @@ def tensor_network_1d_compress_fit(
         string this specifies a custom order.
     optimize : str, optional
         The contraction path optimizer to use.
+    canonize : bool, optional
+        Dummy argument to match the signature of other compression methods.
     sweep_reverse : bool, optional
         Whether to sweep in the reverse direction, swapping whether the final
         tensor network is in right or left canonical form, which also depends
         on the last sweep direction.
-    canonize : bool, optional
-        Dummy argument to match the signature of other compression methods.
+    equalize_norms : bool or float, optional
+        Whether to equalize the norms of the tensors after compression. If an
+        explicit value is give, then the norms will be set to that value, and
+        the overall scaling factor will be accumulated into `.exponent`.
     inplace_fit : bool, optional
         Whether to perform the compression inplace on the initial guess tensor
         network, ``tn_fit``, if supplied.
@@ -1210,6 +1249,12 @@ def tensor_network_1d_compress_fit(
     # possibly put the array indices in canonical order (e.g. when MPS or MPO)
     possibly_permute_(tn_fit, permute_arrays)
 
+    # XXX: do better than simply waiting til the end to equalize norms
+    if equalize_norms is True:
+        tn_fit.equalize_norms_()
+    elif equalize_norms:
+        tn_fit.equalize_norms_(value=equalize_norms)
+
     return tn_fit
 
 
@@ -1268,6 +1313,10 @@ def tensor_network_1d_compress(
         Whether to sweep in the reverse direction, resulting in a left
         canonical form instead of right canonical (for the fit method, this
         also depends on the last sweep direction).
+    equalize_norms : bool or float, optional
+        Whether to equalize the norms of the tensors after compression. If an
+        explicit value is give, then the norms will be set to that value, and
+        the overall scaling factor will be accumulated into `.exponent`.
     inplace : bool, optional
         Whether to perform the compression inplace.
     kwargs
@@ -1289,6 +1338,7 @@ def tensor_network_1d_compress(
             permute_arrays=permute_arrays,
             optimize=optimize,
             sweep_reverse=sweep_reverse,
+            equalize_norms=equalize_norms,
             inplace=inplace,
             **compress_opts,
             **kwargs,
@@ -1310,6 +1360,7 @@ def tensor_network_1d_compress(
             site_tags=site_tags,
             canonize=canonize,
             optimize=optimize,
+            equalize_norms=equalize_norms,
             inplace=inplace,
             **compress_opts,
             **kwargs,

quimb/tensor/tensor_2d_compress.py

@@ -0,0 +1,96 @@
+from .tensor_1d_compress import possibly_permute_
+from .tensor_arbgeom_compress import tensor_network_ag_compress
+
+_TN2D_COMPRESS_METHODS = {}
+
+
+def tensor_network_2d_compress(
+    tn,
+    max_bond=None,
+    cutoff=1e-10,
+    method="local-early",
+    site_tags=None,
+    canonize=True,
+    permute_arrays=True,
+    optimize="auto-hq",
+    equalize_norms=False,
+    compress_opts=None,
+    inplace=False,
+    **kwargs,
+):
+    """Compress a 2D-like tensor network using the specified method.
+
+    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 : {"direct", "dm", "zipup", "zipup-first", "fit", "projector"}
+        The compression method to use.
+    site_tags : sequence of sequence of str, optional
+        The tags to use to group and order the tensors from ``tn``. If not
+        given, uses ``tn.site_tags``. The tensor network built will have one
+        tensor per site, in the order given by ``site_tags``.
+    canonize : bool, optional
+        Whether to perform canonicalization, pseudo or otherwise depending on
+        the method, before compressing. Ignored for ``method='dm'`` and
+        ``method='fit'``.
+    permute_arrays : bool or str, optional
+        Whether to permute the array indices of the final tensor network into
+        canonical order. If ``True`` will use the default order, otherwise if a
+        string this specifies a custom order.
+    optimize : str, optional
+        The contraction path optimizer to use.
+    sweep_reverse : bool, optional
+        Whether to sweep in the reverse direction, resulting in a left
+        canonical form instead of right canonical (for the fit method, this
+        also depends on the last sweep direction).
+    inplace : bool, optional
+        Whether to perform the compression inplace.
+    kwargs
+        Supplied to the chosen compression method.
+
+    Returns
+    -------
+    TensorNetwork
+    """
+    compress_opts = compress_opts or {}
+
+    try:
+        return _TN2D_COMPRESS_METHODS[method](
+            tn,
+            max_bond=max_bond,
+            cutoff=cutoff,
+            site_tags=site_tags,
+            canonize=canonize,
+            permute_arrays=permute_arrays,
+            optimize=optimize,
+            equalize_norms=equalize_norms,
+            inplace=inplace,
+            **compress_opts,
+            **kwargs,
+        )
+    except KeyError:
+        # try arbitrary geometry methods
+        tnc = tensor_network_ag_compress(
+            tn,
+            max_bond=max_bond,
+            cutoff=cutoff,
+            method=method,
+            site_tags=site_tags,
+            canonize=canonize,
+            optimize=optimize,
+            equalize_norms=equalize_norms,
+            inplace=inplace,
+            **compress_opts,
+            **kwargs,
+        )
+
+        if permute_arrays:
+            possibly_permute_(tnc, permute_arrays)
+
+        return tnc

quimb/tensor/tensor_3d.py

@@ -985,6 +985,75 @@ def compress_plane(
                 tag_a, tag_b, max_bond=max_bond, cutoff=cutoff, **compress_opts
             )
 
+    def _contract_boundary_core_via_2d(
+        self,
+        xrange,
+        yrange,
+        zrange,
+        from_which,
+        max_bond,
+        cutoff=1e-10,
+        method="local-early",
+        layer_tags=None,
+        **compress_opts,
+    ):
+        from quimb.tensor.tensor_2d_compress import tensor_network_2d_compress
+
+        r2d = Rotator3D(self, xrange, yrange, zrange, from_which)
+        site_tag = r2d.site_tag
+        istep = r2d.istep
+
+        def _do_compress(site_tag_tmps):
+            nonlocal self
+
+            # split off the boundary network
+            self, tn_boundary = self.partition(site_tag_tmps, inplace=True)
+
+            # compress it inplace
+            tensor_network_2d_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
+
+        # maybe compress the initial row, which may be multiple layers
+        # and have effective bond dimension > max_bond already
+        site_tag_tmps = [
+            site_tag(r2d.sweep[0], j, k) for j, k in r2d.sweep_other
+        ]
+        if any(len(self.tag_map[st]) > 1 for st in site_tag_tmps):
+            _do_compress(site_tag_tmps)
+
+        site_tag_tmps = [f"__ST{j},{k}__" for j, k in r2d.sweep_other]
+
+        if layer_tags is None:
+            layer_tags = [None]
+
+        for i in r2d.sweep[:-1]:
+            for layer_tag in layer_tags:
+                for (j, k), st in zip(r2d.sweep_other, site_tag_tmps):
+                    # group outer single tensor with inner tensor(s)
+                    tag1 = site_tag(i, j, k)  # outer
+                    tag2 = site_tag(i + istep, j, k)  # 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)
+
+                _do_compress(site_tag_tmps)
+
+        self.drop_tags(site_tag_tmps)
+
     def _contract_boundary_core(
         self,
         xrange,
@@ -1187,8 +1256,6 @@ def contract_boundary_from(
         """Unified entrypoint for contracting any rectangular patch of tensors
         from any direction, with any boundary method.
         """
-        check_opt("mode", mode, ("peps", "projector"))
-
         tn = self if inplace else self.copy()
 
         # universal options
@@ -1201,15 +1268,22 @@ def contract_boundary_from(
         contract_boundary_opts["equalize_norms"] = equalize_norms
         contract_boundary_opts["compress_opts"] = compress_opts
 
-        if mode == "projector":
+        if mode == "peps":
+            return tn._contract_boundary_core(
+                canonize_opts=canonize_opts,
+                **contract_boundary_opts,
+            )
+
+        if mode == "l2bp3d":
+            return tn._contract_boundary_l2bp(**contract_boundary_opts)
+
+        if mode == "projector3d":
             return tn._contract_boundary_projector(
                 canonize_opts=canonize_opts, **contract_boundary_opts
             )
 
-        # mode == 'peps' options
-        return tn._contract_boundary_core(
-            canonize_opts=canonize_opts,
-            **contract_boundary_opts,
+        return tn._contract_boundary_core_via_2d(
+            method=mode, **contract_boundary_opts
         )
 
     contract_boundary_from_ = functools.partialmethod(