Trajectree 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

trajectree/quimb/tests/test_linalg/test_slepc_linalg.py

@@ -0,0 +1,283 @@
+from pytest import fixture, mark
+import numpy as np
+import scipy.sparse as sp
+from numpy.testing import assert_allclose
+
+from quimb import (
+    qu,
+    rand_uni,
+    ldmul,
+    rand_matrix,
+    rand_herm,
+    rand_pos,
+    rand_ket,
+    eigh,
+    eye,
+    norm,
+)
+from quimb.linalg import SLEPC4PY_FOUND
+from quimb.linalg.scipy_linalg import svds_scipy
+
+if SLEPC4PY_FOUND:
+    from quimb.linalg.slepc_linalg import (
+        eigs_slepc,
+        svds_slepc,
+        convert_mat_to_petsc,
+        new_petsc_vec,
+        mfn_multiply_slepc,
+        ssolve_slepc,
+    )
+
+
+slepc4py_notfound_msg = "No SLEPc4py installation"
+slepc4py_test = mark.skipif(not SLEPC4PY_FOUND, reason=slepc4py_notfound_msg)
+
+
+@fixture
+def prematsparse():
+    u = rand_uni(4)
+    a = u @ ldmul(np.array([-1, 2, 4, -3]), u.H)
+    a = qu(a, sparse=True)
+    return u, a
+
+
+@fixture
+def bigsparsemat():
+    return rand_matrix(100, sparse=True, density=0.1)
+
+
+# --------------------------------------------------------------------------- #
+# TESTS                                                                       #
+# --------------------------------------------------------------------------- #
+
+
+@slepc4py_test
+class TestConvertToPETScConversion:
+    def test_csr(self):
+        a = rand_matrix(2, sparse=True, density=0.5)
+        b = convert_mat_to_petsc(a)
+        assert b.getType() == "seqaij"
+
+    def test_bsr(self):
+        a = sp.kron(rand_matrix(2), eye(2, sparse=True), format="bsr")
+        b = convert_mat_to_petsc(a)
+        assert b.getType() == "seqbaij"
+        assert b.getBlockSize() == 2
+
+    # def test_vec(self):
+    #     a = np.array([1, 2, 3, 4])
+    #     b = convert_mat_to_petsc(a)
+    #     assert_allclose(b.getArray(), a)
+
+    def test_dense(self):
+        a = rand_matrix(3)
+        b = convert_mat_to_petsc(a)
+        assert b.getType() == "seqdense"
+
+    def test_new_petsc_vector(self):
+        a = new_petsc_vec(4)
+        assert a.getArray() is not None
+
+
+@slepc4py_test
+class TestSlepceigs:
+    def test_internal_eigvals(self, prematsparse):
+        u, a = prematsparse
+        lk = eigs_slepc(a, k=2, sigma=0.5, return_vecs=False)
+        assert_allclose(lk, [-1, 2])
+
+    @mark.parametrize(
+        "which, output",
+        [
+            ("lm", 4),
+            ("sa", -3),
+        ],
+    )
+    def test_eigs_slepc_groundenergy(self, prematsparse, which, output):
+        u, a = prematsparse
+        lk = eigs_slepc(a, k=1, which=which, return_vecs=False)
+        assert_allclose(lk, output)
+
+    @mark.parametrize("dtype", ["real", "complex"])
+    def test_eigs_slepc_eigvecs(self, dtype):
+        h = rand_herm(32, sparse=True, density=0.5)
+        if dtype == "real":
+            h = h.real
+        lks, vks = eigs_slepc(h, k=5)
+        lka, vka = eigh(h, k=5, backend="scipy")
+        assert vks.shape == vka.shape
+        assert h.dtype == vks.dtype
+
+        assert_allclose(lks, lka)
+        assert_allclose(abs(vka.H @ vks), np.eye(5), atol=1e-8)
+
+
+@slepc4py_test
+class TestSlepcSvds:
+    def test_simple(self, prematsparse):
+        u, a = prematsparse
+        lk = svds_slepc(a, k=1, return_vecs=False)
+        assert_allclose(lk, 4)
+
+    @mark.parametrize("SVDType", ["cross", "lanczos"])
+    def test_random_compare_scipy(self, bigsparsemat, SVDType):
+        a = bigsparsemat
+        lk = svds_slepc(a, k=5, return_vecs=False, SVDType=SVDType)
+        ls = svds_scipy(a, k=5, return_vecs=False)
+        assert_allclose(lk, ls)
+
+    @mark.parametrize("SVDType", ["cross", "lanczos"])
+    def test_unitary_vectors(self, bigsparsemat, SVDType):
+        a = bigsparsemat
+        uk, sk, vk = svds_slepc(a, k=10, return_vecs=True, SVDType=SVDType)
+        assert_allclose(uk.H @ uk, eye(10), atol=1e-6)
+        assert_allclose(vk @ vk.H, eye(10), atol=1e-6)
+        pk, lk, qk = svds_scipy(a, k=10, return_vecs=True)
+        assert_allclose(sk, lk)
+        assert pk.shape == uk.shape
+        assert vk.shape == qk.shape
+        assert_allclose(abs(uk.H @ pk), eye(10), atol=1e-6)
+        assert_allclose(abs(qk @ vk.H), eye(10), atol=1e-6)
+
+
+@slepc4py_test
+class TestSlepcMfnMultiply:
+    def test_exp_sparse(self):
+        a = rand_herm(100, sparse=True, density=0.1)
+        k = rand_ket(100)
+
+        out = mfn_multiply_slepc(a, k)
+
+        al, av = eigh(a.toarray())
+        expected = av @ np.diag(np.exp(al)) @ av.conj().T @ k
+
+        assert_allclose(out, expected)
+
+    def test_sqrt_sparse(self):
+        import scipy.sparse as sp
+
+        a = rand_pos(32, sparse=True, density=0.1)
+        a = a + 0.001 * sp.eye(32)
+        k = rand_ket(32)
+
+        out = mfn_multiply_slepc(a, k, fntype="sqrt", isherm=True)
+
+        al, av = eigh(a.toarray())
+        al[al < 0] = 0.0  # very small neg values spoil sqrt
+        expected = av @ np.diag(np.sqrt(al)) @ av.conj().T @ k
+
+        assert_allclose(out, expected, rtol=1e-6)
+
+
+@slepc4py_test
+class TestShellMatrix:
+    def test_extermal_eigen(self):
+        a = rand_herm(100, sparse=True)
+        alo = sp.linalg.aslinearoperator(a)
+
+        el_us, ev_us = sp.linalg.eigsh(alo, k=1, which="LA")
+        el_u, ev_u = eigs_slepc(alo, k=1, which="LA")
+
+        assert_allclose(el_us, el_u)
+        assert_allclose(np.abs(ev_us.conj().T @ ev_u), 1.0)
+
+        el_ls, ev_ls = sp.linalg.eigsh(alo, k=1, which="SA")
+        el_l, ev_l = eigs_slepc(alo, k=1, which="SA")
+
+        assert_allclose(el_ls, el_l)
+        assert_allclose(np.abs(ev_ls.conj().T @ ev_l), 1.0)
+
+    def test_internal_interior_default(self):
+        a = rand_herm(100, sparse=True)
+        alo = sp.linalg.aslinearoperator(a)
+        el, ev = eigs_slepc(alo, k=1, which="TR", sigma=0.0)
+        el_s, ev_s = sp.linalg.eigsh(a.tocsc(), k=1, which="LM", sigma=0.0)
+
+        assert_allclose(el_s, el, rtol=1e-5)
+        assert_allclose(np.abs(ev_s.conj().T @ ev), 1.0)
+
+    def test_internal_shift_invert_linear_solver(self):
+        a = rand_herm(100, sparse=True, seed=42)
+        alo = sp.linalg.aslinearoperator(a)
+
+        st_opts = {
+            "STType": "sinvert",
+            "KSPType": "bcgs",  # / 'gmres'
+            "PCType": "none",
+        }
+
+        el, ev = eigs_slepc(
+            alo,
+            k=1,
+            which="TR",
+            sigma=0.0,
+            tol=1e-6,
+            st_opts=st_opts,
+            EPSType="krylovschur",
+        )
+        el_s, ev_s = sp.linalg.eigsh(a.tocsc(), k=1, which="LM", sigma=0.0)
+
+        assert_allclose(el_s, el, rtol=1e-5)
+        assert_allclose(np.abs(ev_s.conj().T @ ev), 1.0)
+
+    def test_internal_shift_invert_precond(self):
+        a = rand_herm(20, sparse=True, seed=42)
+        alo = sp.linalg.aslinearoperator(a)
+
+        st_opts = {
+            "STType": "precond",
+            "KSPType": "preonly",
+            "PCType": "none",
+        }
+
+        el, ev = eigs_slepc(
+            alo, k=1, which="TR", sigma=0.0, st_opts=st_opts, EPSType="gd"
+        )
+        el_s, ev_s = sp.linalg.eigsh(a.tocsc(), k=1, which="LM", sigma=0.0)
+
+        assert_allclose(el_s, el, rtol=1e-6)
+        assert_allclose(np.abs(ev_s.conj().T @ ev), 1.0)
+
+    def test_internal_shift_invert_precond_jd(self):
+        a = rand_herm(20, sparse=True, seed=42)
+        alo = sp.linalg.aslinearoperator(a)
+
+        st_opts = {
+            "STType": "precond",
+            "KSPType": "bcgs",  # / 'gmres'
+            "PCType": "none",
+        }
+
+        el, ev = eigs_slepc(
+            alo, k=1, which="TR", sigma=0.0, st_opts=st_opts, EPSType="jd"
+        )
+        el_s, ev_s = sp.linalg.eigsh(a.tocsc(), k=1, which="LM", sigma=0.0)
+
+        assert_allclose(el_s, el, rtol=1e-6)
+        assert_allclose(np.abs(ev_s.conj().T @ ev), 1.0)
+
+
+@slepc4py_test
+class TestCISS:
+    def test_1(self):
+        a = rand_herm(100, sparse=True)
+        el, ev = eigh(a.toarray())
+        which = abs(el) < 0.2
+        el, ev = el[which], ev[:, which]
+
+        offset = norm(a, "fro")
+        a = a + offset * sp.eye(a.shape[0])
+
+        sl, sv = eigs_slepc(a, k=6, l_win=(-0.2 + offset, 0.2 + offset))
+        sl -= offset
+        assert_allclose(el, sl)
+        assert_allclose(np.abs(sv.H @ ev), np.eye(el.size), atol=1e-11)
+
+
+@slepc4py_test
+class TestSSolve:
+    def test_simple_dense(self):
+        a = rand_herm(2**4, sparse=True)
+        y = rand_ket(2**4)
+        x = ssolve_slepc(a, y)
+        assert_allclose(a @ x, y, atol=1e-12, rtol=1e-6)

trajectree/quimb/tests/test_tensor/test_belief_propagation/test_d1bp.py

@@ -0,0 +1,39 @@
+import pytest
+
+import quimb as qu
+import quimb.tensor as qtn
+from quimb.experimental.belief_propagation import (
+    D1BP,
+    contract_d1bp,
+)
+
+
+def test_contract_tree_exact():
+    tn = qtn.TN_rand_tree(20, 3)
+    Z = tn.contract()
+    info = {}
+    Z_bp = contract_d1bp(tn, info=info, progbar=True)
+    assert info["converged"]
+    assert Z == pytest.approx(Z_bp, rel=1e-12)
+
+
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+def test_contract_normal(damping):
+    tn = qtn.TN2D_from_fill_fn(lambda s: qu.randn(s, dist="uniform"), 6, 6, 2)
+    Z = tn.contract()
+    info = {}
+    Z_bp = contract_d1bp(tn, damping=damping, info=info, progbar=True)
+    assert info["converged"]
+    assert Z == pytest.approx(Z_bp, rel=1e-1)
+
+
+def test_get_gauged_tn():
+    tn = qtn.TN2D_from_fill_fn(lambda s: qu.randn(s, dist="uniform"), 6, 6, 2)
+    Z = tn.contract()
+    bp = D1BP(tn)
+    bp.run()
+    Zbp = bp.contract()
+    assert Z == pytest.approx(Zbp, rel=1e-1)
+    tn_gauged = bp.get_gauged_tn()
+    Zg = qu.prod(array.item(0) for array in tn_gauged.arrays)
+    assert Z == pytest.approx(Zg, rel=1e-1)

trajectree/quimb/tests/test_tensor/test_belief_propagation/test_d2bp.py

@@ -0,0 +1,67 @@
+import pytest
+
+import quimb.tensor as qtn
+from quimb.experimental.belief_propagation.d2bp import (
+    compress_d2bp,
+    contract_d2bp,
+    sample_d2bp,
+)
+
+
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+@pytest.mark.parametrize("dtype", ["float32", "complex64"])
+def test_contract(damping, dtype):
+    peps = qtn.PEPS.rand(3, 4, 3, seed=42, dtype=dtype)
+    # normalize exactly
+    peps /= (peps.H @ peps) ** 0.5
+    info = {}
+    N_ap = contract_d2bp(peps, damping=damping, info=info, progbar=True)
+    assert info["converged"]
+    assert N_ap == pytest.approx(1.0, rel=0.3)
+
+
+@pytest.mark.parametrize("dtype", ["float32", "complex64"])
+def test_tree_exact(dtype):
+    psi = qtn.TN_rand_tree(20, 3, 2, dtype=dtype, seed=42)
+    norm2 = psi.H @ psi
+    info = {}
+    norm2_bp = contract_d2bp(psi, info=info, progbar=True)
+    assert info["converged"]
+    assert norm2_bp == pytest.approx(norm2, rel=1e-4)
+
+
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+@pytest.mark.parametrize("dtype", ["float32", "complex64"])
+def test_compress(damping, dtype):
+    peps = qtn.PEPS.rand(3, 4, 3, seed=42, dtype=dtype)
+    # test that using the BP compression gives better fidelity than purely
+    # local, naive compression scheme
+    peps_c1 = peps.compress_all(max_bond=2)
+    info = {}
+    peps_c2 = compress_d2bp(
+        peps, max_bond=2, damping=damping, info=info, progbar=True
+    )
+    assert info["converged"]
+    fid1 = peps_c1.H @ peps_c2
+    fid2 = peps_c2.H @ peps_c2
+    assert abs(fid2) > abs(fid1)
+
+
+@pytest.mark.parametrize("dtype", ["float32", "complex64"])
+def test_sample(dtype):
+    peps = qtn.PEPS.rand(3, 4, 3, seed=42, dtype=dtype)
+    # normalize exactly
+    peps /= (peps.H @ peps) ** 0.5
+    config, peps_config, omega = sample_d2bp(peps, seed=42, progbar=True)
+    assert all(ix in config for ix in peps.site_inds)
+    assert 0.0 < omega < 1.0
+    assert peps_config.outer_inds() == ()
+
+    ptotal = 0.0
+    nrepeat = 4
+    for _ in range(nrepeat):
+        _, peps_config, _ = sample_d2bp(peps, seed=42, progbar=True)
+        ptotal += abs(peps_config.contract()) ** 2
+
+    # check we are doing better than random guessing
+    assert ptotal > nrepeat * 2**-peps.nsites

trajectree/quimb/tests/test_tensor/test_belief_propagation/test_hd1bp.py

@@ -0,0 +1,64 @@
+import pytest
+
+import quimb as qu
+import quimb.tensor as qtn
+from quimb.experimental.belief_propagation.hd1bp import (
+    HD1BP,
+    contract_hd1bp,
+    sample_hd1bp,
+)
+
+
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+def test_contract_hyper(damping):
+    htn = qtn.HTN_random_ksat(3, 50, alpha=2.0, seed=42, mode="dense")
+    info = {}
+    num_solutions = contract_hd1bp(
+        htn, damping=damping, info=info, progbar=True
+    )
+    assert info["converged"]
+    assert num_solutions == pytest.approx(309273226, rel=0.1)
+
+
+def test_contract_tree_exact():
+    tn = qtn.TN_rand_tree(20, 3)
+    Z = tn.contract()
+    info = {}
+    Z_bp = contract_hd1bp(tn, info=info, progbar=True)
+    assert info["converged"]
+    assert Z == pytest.approx(Z_bp, rel=1e-12)
+
+
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+def test_contract_normal(damping):
+    tn = qtn.TN2D_from_fill_fn(lambda s: qu.randn(s, dist="uniform"), 6, 6, 2)
+    Z = tn.contract()
+    info = {}
+    Z_bp = contract_hd1bp(tn, damping=damping, info=info, progbar=True)
+    assert info["converged"]
+    assert Z == pytest.approx(Z_bp, rel=1e-1)
+
+
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+def test_sample(damping):
+    nvars = 20
+    htn = qtn.HTN_random_ksat(3, nvars, alpha=2.0, seed=42, mode="dense")
+    config, tn_config, omega = sample_hd1bp(
+        htn, damping=damping, seed=42, progbar=True
+    )
+    assert len(config) == nvars
+    assert tn_config.num_indices == 0
+    assert tn_config.contract() == pytest.approx(1.0)
+    assert 0.0 < omega < 1.0
+
+
+def test_get_gauged_tn():
+    tn = qtn.TN2D_from_fill_fn(lambda s: qu.randn(s, dist="uniform"), 6, 6, 2)
+    Z = tn.contract()
+    bp = HD1BP(tn)
+    bp.run()
+    Zbp = bp.contract()
+    assert Z == pytest.approx(Zbp, rel=1e-1)
+    tn_gauged = bp.get_gauged_tn()
+    Zg = qu.prod(array.item(0) for array in tn_gauged.arrays)
+    assert Z == pytest.approx(Zg, rel=1e-1)

trajectree/quimb/tests/test_tensor/test_belief_propagation/test_hv1bp.py

@@ -0,0 +1,51 @@
+import pytest
+
+import quimb as qu
+import quimb.tensor as qtn
+from quimb.experimental.belief_propagation.hv1bp import (
+    contract_hv1bp,
+    sample_hv1bp,
+)
+
+
+@pytest.mark.parametrize("damping", [0.0, 0.1, 0.5])
+def test_contract_hyper(damping):
+    htn = qtn.HTN_random_ksat(3, 50, alpha=2.0, seed=42, mode="dense")
+    info = {}
+    num_solutions = contract_hv1bp(
+        htn, damping=damping, info=info, progbar=True
+    )
+    assert info["converged"]
+    assert num_solutions == pytest.approx(309273226, rel=0.1)
+
+
+def test_contract_tree_exact():
+    tn = qtn.TN_rand_tree(20, 3)
+    Z = tn.contract()
+    info = {}
+    Z_bp = contract_hv1bp(tn, info=info, progbar=True)
+    assert info["converged"]
+    assert Z == pytest.approx(Z_bp, rel=1e-12)
+
+
+@pytest.mark.parametrize("damping", [0.0, 0.1, 0.5])
+def test_contract_normal(damping):
+    tn = qtn.TN2D_from_fill_fn(lambda s: qu.randn(s, dist="uniform"), 6, 6, 2)
+    Z = tn.contract()
+    info = {}
+    Z_bp = contract_hv1bp(tn, damping=damping, info=info, progbar=True)
+    assert info["converged"]
+    assert Z == pytest.approx(Z_bp, rel=1e-1)
+
+
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+def test_sample(damping):
+    nvars = 20
+    htn = qtn.HTN_random_ksat(3, nvars, alpha=2.0, seed=42, mode="dense")
+    config, tn_config, omega = sample_hv1bp(
+        htn, damping=damping, seed=42, progbar=True
+    )
+    assert len(config) == nvars
+    assert tn_config.num_indices == 0
+    assert tn_config.contract() == pytest.approx(1.0)
+    assert 0.0 < omega < 1.0

trajectree/quimb/tests/test_tensor/test_belief_propagation/test_l1bp.py

@@ -0,0 +1,142 @@
+import pytest
+
+import quimb as qu
+import quimb.tensor as qtn
+from quimb.experimental.belief_propagation.l1bp import contract_l1bp
+from quimb.experimental.belief_propagation.d2bp import contract_d2bp
+
+
+@pytest.mark.parametrize("dtype", ["float32", "complex64"])
+def test_contract_tree_exact(dtype):
+    tn = qtn.TN_rand_tree(10, 3, seed=42, dtype=dtype)
+    Z_ex = tn.contract()
+    info = {}
+    Z_bp = contract_l1bp(tn, info=info, progbar=True)
+    assert info["converged"]
+    assert Z_ex == pytest.approx(Z_bp, rel=5e-6)
+
+
+@pytest.mark.parametrize("dtype", ["float32", "complex64"])
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+def test_contract_loopy_approx(dtype, damping):
+    tn = qtn.TN2D_rand(3, 4, 5, dtype=dtype, dist="uniform")
+    Z_ex = tn.contract()
+    info = {}
+    Z_bp = contract_l1bp(tn, damping=damping, info=info, progbar=True)
+    assert info["converged"]
+    assert Z_ex == pytest.approx(Z_bp, rel=0.1)
+
+
+@pytest.mark.parametrize("dtype", ["float32", "complex64"])
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+@pytest.mark.parametrize("update", ("parallel", "sequential"))
+def test_contract_double_loopy_approx(dtype, damping, update):
+    peps = qtn.PEPS.rand(4, 3, 2, seed=42, dtype=dtype)
+    tn = peps.H & peps
+    Z_ex = tn.contract()
+    info = {}
+    Z_bp1 = contract_l1bp(
+        tn, damping=damping, update=update, info=info, progbar=True
+    )
+    assert info["converged"]
+    assert Z_bp1 == pytest.approx(Z_ex, rel=0.3)
+    # compare with 2-norm BP on the peps directly
+    Z_bp2 = contract_d2bp(peps)
+    assert Z_bp1 == pytest.approx(Z_bp2, rel=5e-6)
+
+
+@pytest.mark.parametrize("dtype", ["float32", "complex64"])
+def test_contract_tree_triple_sandwich_exact(dtype):
+    edges = qtn.edges_tree_rand(20, 3, seed=42)
+    ket = qtn.TN_from_edges_rand(
+        edges,
+        3,
+        phys_dim=2,
+        seed=42,
+        site_ind_id="k{}",
+        dtype=dtype,
+    )
+    op = qtn.TN_from_edges_rand(
+        edges,
+        2,
+        phys_dim=2,
+        seed=42,
+        site_ind_id=("k{}", "b{}"),
+        dtype=dtype,
+    )
+    bra = qtn.TN_from_edges_rand(
+        edges,
+        3,
+        phys_dim=2,
+        seed=42,
+        site_ind_id="b{}",
+        dtype=dtype,
+    )
+    tn = bra.H | op | ket
+    Z_ex = tn.contract()
+    info = {}
+    Z_bp = contract_l1bp(tn, info=info, progbar=True)
+    assert info["converged"]
+    assert Z_ex == pytest.approx(Z_bp, rel=5e-6)
+
+
+@pytest.mark.parametrize("dtype", ["float32", "complex64"])
+@pytest.mark.parametrize("damping", [0.0, 0.1])
+def test_contract_tree_triple_sandwich_loopy_approx(dtype, damping):
+    edges = qtn.edges_2d_hexagonal(2, 3)
+    ket = qtn.TN_from_edges_rand(
+        edges,
+        3,
+        phys_dim=2,
+        seed=42,
+        site_ind_id="k{}",
+        dtype=dtype,
+        # make the wavefunction postive to make easier
+        dist="uniform",
+    )
+    ket /= (ket.H @ ket) ** 0.5
+
+    G_ket = ket.gate(qu.pauli("Z"), [(1, 1, "A")], propagate_tags="sites")
+    tn = ket.H | G_ket
+    Z_ex = tn.contract()
+    info = {}
+    Z_bp = contract_l1bp(tn, damping=damping, info=info, progbar=True)
+    assert info["converged"]
+    assert Z_bp == pytest.approx(Z_ex, rel=0.5)
+
+
+def test_contract_cluster_approx():
+    tn = qtn.TN2D_classical_ising_partition_function(8, 8, 0.4, h=0.2)
+    f_ex = qu.log(tn.contract())
+    f_bp = qu.log(contract_l1bp(tn))
+    assert f_bp == pytest.approx(f_ex, rel=0.3)
+    cluster_tags = []
+    for i in range(0, 8, 2):
+        for j in range(0, 8, 2):
+            cluster_tag = f"C{i},{j}"
+            tn[i, j].add_tag(cluster_tag)
+            tn[i, j + 1].add_tag(cluster_tag)
+            tn[i + 1, j].add_tag(cluster_tag)
+            tn[i + 1, j + 1].add_tag(cluster_tag)
+            cluster_tags.append(cluster_tag)
+    info = {}
+    f_bp2 = qu.log(
+        contract_l1bp(tn, site_tags=cluster_tags, info=info, progbar=True)
+    )
+    assert info["converged"]
+    assert f_bp == pytest.approx(f_ex, rel=0.1)
+    assert abs(1 - f_ex / f_bp2) < abs(1 - f_ex / f_bp)
+
+
+def test_mps():
+    # catch bug to do with structured contract and output inds
+    L = 6
+    psi = qtn.MPS_rand_state(L=L, seed=20, bond_dim=3)
+    psiG = psi.copy()
+    psiG.gate_(qu.pauli("X"), 5, contract=True)
+    expec = psi.H & psiG
+    O = contract_l1bp(
+        expec,
+        site_tags=[f"I{i}" for i in range(L)],
+    )
+    assert O == pytest.approx(expec ^ ..., abs=1e-6)