tensor compress: add reduced={'left'|'right'} modes

quimb/tensor/tensor_core.py

@@ -669,7 +669,7 @@ def tensor_compress_bond(
     **compress_opts,
 ):
     r"""Inplace compress between the two single tensors. It follows the
-    following steps to minimize the size of SVD performed::
+    following steps (by default) to minimize the size of SVD performed::
 
         a)│   │        b)│        │        c)│       │
         ━━●━━━●━━  ->  ━━>━━○━━○━━<━━  ->  ━━>━━━M━━━<━━
@@ -693,8 +693,12 @@ def tensor_compress_bond(
         The maxmimum bond dimension.
     cutoff : float, optional
         The singular value cutoff to use.
-    reduced : bool, optional
-        Whether to perform the QR reduction as above or not.
+    reduced : {True, False, "left", "right"}, optional
+        Whether to perform the QR reduction as above or not. If False, contract
+        both tensors together and perform a single SVD. If 'left' or 'right'
+        then just perform the svd on the left or right tensor respectively.
+        This can still be optimal if the other tensor is already isometric,
+        i.e. the pair are right or left canonical respectively.
     absorb : {'both', 'left', 'right', None}, optional
         Where to absorb the singular values after decomposition.
     info : None or dict, optional
@@ -735,21 +739,61 @@ def tensor_compress_bond(
         T1C = T1_L.contract(M_L, output_inds=T1.inds)
         T2C = M_R.contract(T2_R, output_inds=T2.inds)
 
-    elif reduced == 'lazy':
-        compress_opts.setdefault('method', 'isvd')
+    elif reduced == "right":
+        # if left canonical, just do svd on right tensor
+        M, *s, T2C = T2.split(
+            left_inds=bix,
+            right_inds=rix,
+            get="tensors",
+            absorb=absorb,
+            **compress_opts,
+        )
+        T1C = T1 @ M
+        T1C.transpose_like_(T1)
+        T2C.transpose_like_(T2)
+
+        if absorb == "right":
+            # can't mark left tensor as isometric if absorbed into right tensor
+            absorb = "both"
+
+    elif reduced == "left":
+        # if right canonical, just do svd on left tensor
+        T1C, *s, M = T1.split(
+            left_inds=lix,
+            right_inds=bix,
+            get="tensors",
+            absorb=absorb,
+            **compress_opts,
+        )
+        T2C = M @ T2
+        T1C.transpose_like_(T1)
+        T2C.transpose_like_(T2)
+
+        if absorb == "left":
+            # can't mark right tensor as isometric if absorbed into left tensor
+            absorb = "both"
+
+    elif reduced == "lazy":
+        compress_opts.setdefault("method", "isvd")
         T12 = TNLinearOperator((T1, T2), lix, rix)
         T1C, *s, T2C = T12.split(get="tensors", absorb=absorb, **compress_opts)
         T1C.transpose_like_(T1)
         T2C.transpose_like_(T2)
 
-    else:
+    elif reduced is False:
         T12 = T1 @ T2
         T1C, *s, T2C = T12.split(
             left_inds=lix, get="tensors", absorb=absorb, **compress_opts
         )
         T1C.transpose_like_(T1)
         T2C.transpose_like_(T2)
 
+    else:
+        raise ValueError(
+            f"Unrecognized value for `reduced` argument: {reduced}."
+            "Valid options are {True, False, 'left', 'right', 'lazy'}."
+        )
+
     # update with the new compressed data
     T1.modify(data=T1C.data)
     T2.modify(data=T2C.data)

tests/test_tensor/test_tensor_core.py

@@ -1343,6 +1343,89 @@ def test_compress_multibond(self):
         x1 = (A & B).trace("a", "d")
         assert x1 == pytest.approx(x0)
 
+    @pytest.mark.parametrize("absorb", ["both", "left", "right"])
+    def test_tensor_compress_bond_reduced_modes(self, absorb):
+        kws = dict(max_bond=4, absorb=absorb)
+
+        A = rand_tensor((3, 4, 5), "abc", tags="A")
+        B = rand_tensor((5, 6), "cd", tags="B")
+        AB = A @ B
+
+        # naive contract and compress should be optimal
+        A1, B1 = A.copy(), B.copy()
+        qtn.tensor_compress_bond(A1, B1, reduced=False, **kws)
+        assert A1.shape == (3, 4, 4)
+        assert B1.shape == (4, 6)
+        if absorb == "right":
+            # assert A is isometric
+            assert A1.norm()**2 == pytest.approx(4)
+        if absorb == "left":
+            # assert B is isometric
+            assert B1.norm()**2 == pytest.approx(4)
+        # compute the optimal fidelity
+        d1 = (A1 @ B1).distance(AB)
+        assert 0 < d1
+
+        A2, B2 = A.copy(), B.copy()
+        qtn.tensor_compress_bond(A2, B2, reduced=True, **kws)
+        assert A2.shape == (3, 4, 4)
+        assert B2.shape == (4, 6)
+        if absorb == "right":
+            assert A2.norm()**2 == pytest.approx(4)
+        if absorb == "left":
+            assert B2.norm()**2 == pytest.approx(4)
+        d2 = (A2 @ B2).distance(AB)
+        # reduced mode should also be optimal
+        assert d2 == pytest.approx(d1)
+
+        Ar, Br = A.copy(), B.copy()
+        qtn.tensor_compress_bond(Ar, Br, reduced="right", **kws)
+        assert Ar.shape == (3, 4, 4)
+        assert Br.shape == (4, 6)
+        if absorb == "right":
+            # A won't be canonical
+            assert Ar.left_inds is None
+            assert Ar.norm()**2 != pytest.approx(4)
+        if absorb == "left":
+            assert Br.norm()**2 == pytest.approx(4)
+        dr = (Ar @ Br).distance(AB)
+        # right reduced mode should not be optimal
+        assert dr > d1
+        # unless we canonicalize first
+        Ar, Br = A.copy(), B.copy()
+        qtn.tensor_canonize_bond(Ar, Br, absorb="right")
+        qtn.tensor_compress_bond(Ar, Br, reduced="right", **kws)
+        if absorb == "right":
+            assert Ar.norm()**2 == pytest.approx(4)
+        if absorb == "left":
+            assert Br.norm()**2 == pytest.approx(4)
+        dr = (Ar @ Br).distance(AB)
+        assert dr == pytest.approx(d1)
+
+        Al, Bl = A.copy(), B.copy()
+        qtn.tensor_compress_bond(Al, Bl, reduced="left", **kws)
+        assert Al.shape == (3, 4, 4)
+        assert Bl.shape == (4, 6)
+        if absorb == "right":
+            assert Al.norm()**2 == pytest.approx(4)
+        if absorb == "left":
+            # B won't be canonical
+            assert Bl.left_inds is None
+            assert Bl.norm()**2 != pytest.approx(4)
+        dl = (Al @ Bl).distance(AB)
+        # left reduced mode should not be optimal
+        assert dl > d1
+        # unless we canonicalize first
+        Al, Bl = A.copy(), B.copy()
+        qtn.tensor_canonize_bond(Al, Bl, absorb="left")
+        qtn.tensor_compress_bond(Al, Bl, reduced="left", **kws)
+        if absorb == "right":
+            assert Al.norm()**2 == pytest.approx(4)
+        if absorb == "left":
+            assert Bl.norm()**2 == pytest.approx(4)
+        dl = (Al @ Bl).distance(AB)
+        assert dl == pytest.approx(d1)
+
     def test_canonize_multibond(self):
         A = rand_tensor((3, 4, 5), "abc", tags="A")
         assert A.H @ A != pytest.approx(3)