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

trajectree/quimb/tests/test_tensor/test_tensor_2d_tebd.py

@@ -0,0 +1,144 @@
+import itertools
+import importlib
+
+import pytest
+
+import quimb as qu
+import quimb.tensor as qtn
+
+found_torch = importlib.util.find_spec("torch") is not None
+
+pytorch_case = pytest.param(
+    "torch",
+    marks=pytest.mark.skipif(not found_torch, reason="pytorch not installed"),
+)
+
+
+class TestLocalHam2DConstruct:
+    @pytest.mark.parametrize("H2_type", ["default", "manual"])
+    @pytest.mark.parametrize("H1_type", [None, "default", "manual"])
+    @pytest.mark.parametrize("Lx", [3, 4])
+    @pytest.mark.parametrize("Ly", [3, 4])
+    def test_construct(self, Lx, Ly, H2_type, H1_type):
+        import matplotlib
+        from matplotlib import pyplot as plt
+
+        matplotlib.use("Template")
+
+        if H2_type == "default":
+            H2 = qu.rand_herm(4)
+        elif H2_type == "manual":
+            H2 = dict()
+            for i, j in itertools.product(range(Lx), range(Ly)):
+                if i + 1 < Lx:
+                    H2[(i, j), (i + 1, j)] = qu.rand_herm(4)
+                if j + 1 < Ly:
+                    H2[(i, j), (i, j + 1)] = qu.rand_herm(4)
+
+        if H1_type is None:
+            H1 = None
+        elif H1_type == "default":
+            H1 = qu.rand_herm(2)
+        elif H1_type == "manual":
+            H1 = dict()
+            for i, j in itertools.product(range(Lx), range(Ly)):
+                H1[i, j] = qu.rand_herm(2)
+
+        ham = qtn.LocalHam2D(Lx, Ly, H2, H1)
+        assert len(ham.terms) == 2 * Lx * Ly - Lx - Ly
+
+        # check that terms are being cached if possible
+        if (H2_type == "default") and (H1_type is None):
+            assert len({id(x) for x in ham.terms.values()}) == 1
+
+        print(ham)
+        fig, ax = ham.draw()
+        plt.close(fig)
+
+    @pytest.mark.parametrize("Lx", [4, 5])
+    @pytest.mark.parametrize("Ly", [4, 5])
+    @pytest.mark.parametrize(
+        "order", [None, "sort", "random", "smallest_last"]
+    )
+    def test_ordering(self, Lx, Ly, order):
+        ham = qtn.ham_2d_j1j2(Lx, Ly)
+        assert len(ham.terms) == 2 * Lx * Ly - Lx - Ly + 2 * (Lx - 1) * (
+            Ly - 1
+        )
+        ordering = ham.get_auto_ordering(order)
+        assert len(ordering) == len(ham.terms)
+        assert set(ordering) == set(ham.terms)
+        assert tuple(ordering) != tuple(ham.terms)
+
+        # make sure first four pairs are in same commuting group at least
+        first_four_pairs = tuple(itertools.chain(*ordering[:4]))
+        assert len(first_four_pairs) == len(set(first_four_pairs))
+
+
+class TestSimpleUpdate:
+    @pytest.mark.parametrize("backend", ["numpy", pytorch_case])
+    def test_heis_small(self, backend):
+        Lx = 3
+        Ly = 4
+        D = 2
+
+        ham = qtn.ham_2d_heis(Lx, Ly)
+        psi0 = qtn.PEPS.rand(Lx, Ly, D)
+
+        def to_backend(x):
+            import autoray
+
+            return autoray.do("array", x, like=backend)
+
+        psi0.apply_to_arrays(to_backend)
+        ham.apply_to_arrays(to_backend)
+
+        su = qtn.SimpleUpdate(
+            psi0,
+            ham,
+            progbar=True,
+            keep_best=True,
+            compute_energy_every=10,
+            ordering="largest_first",
+        )
+
+        su.evolve(33, tau=0.3)
+        su.state = su.best["state"]
+        su.evolve(33, tau=0.1)
+        su.state = su.best["state"]
+        su.evolve(33, tau=0.03)
+        su.state = su.best["state"]
+
+        assert su.best["energy"] < -6.25
+
+
+class TestFullUpdate:
+    @pytest.mark.parametrize("backend", ["numpy", pytorch_case])
+    def test_heis_small(self, backend):
+        Lx = 3
+        Ly = 4
+        D = 2
+
+        psi0 = qtn.PEPS.rand(Lx, Ly, D)
+        ham = qtn.ham_2d_heis(Lx, Ly)
+
+        def to_backend(x):
+            import autoray
+
+            return autoray.do("array", x, like=backend)
+
+        psi0.apply_to_arrays(to_backend)
+        ham.apply_to_arrays(to_backend)
+
+        su = qtn.FullUpdate(
+            psi0, ham, progbar=True, keep_best=True, compute_energy_every=1
+        )
+
+        su.evolve(33, tau=0.3)
+        su.state = su.best["state"]
+        su.evolve(33, tau=0.1)
+        su.state = su.best["state"]
+        su.evolve(33, tau=0.03)
+        su.state = su.best["state"]
+
+        assert su.best["energy"] < -6.30

trajectree/quimb/tests/test_tensor/test_tensor_3d.py

@@ -0,0 +1,123 @@
+import pytest
+import autoray as ar
+
+import quimb as qu
+import quimb.tensor as qtn
+
+
+class TestTensorNetwork3D:
+    def test_cyclic_basic(self):
+        tn = qtn.TN3D_empty(Lx=3, Ly=4, Lz=5, D=2, cyclic=True)
+        assert tn.is_cyclic_x()
+        assert tn.is_cyclic_y()
+        assert tn.is_cyclic_z()
+        assert tn.num_indices == 3 * tn.nsites
+        tn = qtn.TN3D_empty(Lx=3, Ly=4, Lz=5, D=2, cyclic=(False, False, True))
+        assert not tn.is_cyclic_x()
+        assert not tn.is_cyclic_y()
+        assert tn.is_cyclic_z()
+        assert tn.num_indices == 3 * tn.nsites - (tn.Lx * tn.Lz) - (
+            tn.Ly * tn.Lz
+        )
+        tn = qtn.TN3D_empty(Lx=3, Ly=4, Lz=5, D=2, cyclic=(False, True, False))
+        assert not tn.is_cyclic_x()
+        assert tn.is_cyclic_y()
+        assert not tn.is_cyclic_z()
+        assert tn.num_indices == 3 * tn.nsites - (tn.Lx * tn.Ly) - (
+            tn.Ly * tn.Lz
+        )
+        tn = qtn.TN3D_empty(Lx=3, Ly=4, Lz=5, D=2, cyclic=(True, False, False))
+        assert tn.is_cyclic_x()
+        assert not tn.is_cyclic_y()
+        assert not tn.is_cyclic_z()
+        assert tn.num_indices == 3 * tn.nsites - (tn.Lx * tn.Ly) - (
+            tn.Lx * tn.Lz
+        )
+
+
+class Test3DManualContract:
+    @pytest.mark.parametrize("canonize", [False, True])
+    def test_contract_boundary_ising_model(self, canonize):
+        L = 5
+        beta = 0.3
+        fex = -2.7654417752878
+        tn = qtn.TN3D_classical_ising_partition_function(L, L, L, beta=beta)
+        Z = tn.contract_boundary(max_bond=8, canonize=canonize)
+        f = -qu.log(Z) / (L**3 * beta)
+        assert f == pytest.approx(fex, rel=1e-3)
+
+    @pytest.mark.parametrize("dims", [(10, 4, 3), (4, 3, 10), (3, 10, 4)])
+    def test_contract_boundary_stopping_criterion(self, dims):
+        tn = qtn.TN3D_from_fill_fn(
+            lambda shape: ar.lazy.Variable(shape=shape, backend="numpy"),
+            *dims,
+            D=2,
+        )
+        tn.contract_boundary_(
+            4, cutoff=0.0, final_contract=False, progbar=True
+        )
+        assert tn.max_bond() == 4
+        assert 32 <= tn.num_tensors <= 40
+
+    @pytest.mark.parametrize("lazy", [False, True])
+    def test_coarse_grain_basics(self, lazy):
+        tn = qtn.TN3D_from_fill_fn(
+            lambda shape: ar.lazy.Variable(shape, backend="numpy"),
+            Lx=6,
+            Ly=7,
+            Lz=8,
+            D=2,
+        )
+        tncg = tn.coarse_grain_hotrg("x", max_bond=3, cutoff=0.0, lazy=lazy)
+        assert (tncg.Lx, tncg.Ly, tncg.Lz) == (3, 7, 8)
+        assert not tncg.outer_inds()
+        assert tncg.max_bond() == 3
+        assert "I4,0,0" not in tncg.tag_map
+        assert "X3" not in tncg.tag_map
+
+        tncg = tn.coarse_grain_hotrg("y", max_bond=3, cutoff=0.0, lazy=lazy)
+        assert (tncg.Lx, tncg.Ly, tncg.Lz) == (6, 4, 8)
+        assert not tncg.outer_inds()
+        assert tncg.max_bond() == 3
+        assert "I0,5,0" not in tncg.tag_map
+        assert "Y4" not in tncg.tag_map
+
+        tncg = tn.coarse_grain_hotrg("z", max_bond=3, cutoff=0.0, lazy=lazy)
+        assert (tncg.Lx, tncg.Ly, tncg.Lz) == (6, 7, 4)
+        assert "I0,0,5" not in tncg.tag_map
+        assert "Z4" not in tncg.tag_map
+
+    def test_contract_hotrg_ising_model(self):
+        L = 5
+        beta = 0.3
+        fex = -2.7654417752878
+        tn = qtn.TN3D_classical_ising_partition_function(L, L, L, beta=beta)
+        tn.contract_hotrg_(max_bond=4, progbar=True, equalize_norms=1.0)
+        Z = tn.item() * 10**tn.exponent
+        f = -qu.log(Z) / (L**3 * beta)
+        assert f == pytest.approx(fex, rel=1e-2)
+
+    @pytest.mark.parametrize("cyclicx", [False, True])
+    @pytest.mark.parametrize("cyclicy", [False, True])
+    @pytest.mark.parametrize("cyclicz", [False, True])
+    @pytest.mark.parametrize("mode", ["hotrg", "ctmrg"])
+    def test_contract_cyclic(self, cyclicx, cyclicy, cyclicz, mode):
+        Lx, Ly, Lz = 3, 4, 5
+        chi = 3
+        tn = qtn.TN3D_from_fill_fn(
+            lambda shape: ar.lazy.Variable(shape=shape, backend="numpy"),
+            Lx,
+            Ly,
+            Lz,
+            D=2,
+            cyclic=(cyclicx, cyclicy, cyclicz),
+        )
+        if mode == "hotrg":
+            lZ = tn.contract_hotrg(max_bond=chi, cutoff=0.0)
+        elif mode == "ctmrg":
+            lZ = tn.contract_ctmrg(max_bond=chi, cutoff=0.0)
+
+        if any((cyclicx, cyclicy, cyclicz)):
+            assert lZ.history_max_size() < 2**16
+        else:
+            assert lZ.history_max_size() < 2**13

trajectree/quimb/tests/test_tensor/test_tensor_arbgeom.py

@@ -0,0 +1,226 @@
+import pytest
+from numpy.testing import assert_allclose
+
+import quimb.tensor as qtn
+
+
+@pytest.mark.parametrize("which_A", ["upper", "lower"])
+@pytest.mark.parametrize("contract", [True, False])
+@pytest.mark.parametrize("inplace", [True, False])
+def test_tensor_network_apply_op_vec(which_A, contract, inplace):
+    A = qtn.TN_from_edges_rand(
+        qtn.edges_2d_square(3, 2),
+        D=2,
+        phys_dim=2,
+        site_ind_id=("k{}", "b{}"),
+        dtype=complex,
+    )
+    x = qtn.TN_from_edges_rand(
+        qtn.edges_2d_square(3, 2),
+        D=3,
+        phys_dim=2,
+        site_ind_id="x{}",
+        dtype=complex,
+    )
+
+    Ad = A.to_dense()
+    if which_A == "upper":
+        Ad = Ad.T
+    xd = x.to_dense()
+    C = Ad @ xd
+
+    Ax = qtn.tensor_network_apply_op_vec(
+        A,
+        x,
+        which_A,
+        inplace=inplace,
+        contract=contract,
+    )
+
+    if contract:
+        # checks fusing
+        assert Ax.num_indices == x.num_indices
+
+    if inplace:
+        assert Ax is x
+    else:
+        assert isinstance(Ax, x.__class__)
+        assert Ax.site_ind_id == x.site_ind_id
+
+    assert_allclose(Ax.to_dense(), C)
+
+
+@pytest.mark.parametrize("which_A", ["upper", "lower"])
+@pytest.mark.parametrize("which_B", ["upper", "lower"])
+@pytest.mark.parametrize("contract", [True, False])
+@pytest.mark.parametrize("inplace", [True, False])
+def test_tensor_network_apply_op_op(which_A, which_B, contract, inplace):
+    A = qtn.TN_from_edges_rand(
+        qtn.edges_2d_square(3, 2),
+        D=2,
+        phys_dim=2,
+        site_ind_id=("k{}", "b{}"),
+        dtype=complex,
+    )
+    B = qtn.TN_from_edges_rand(
+        qtn.edges_2d_square(3, 2),
+        D=3,
+        phys_dim=2,
+        site_ind_id=("x{}", "y{}"),
+        dtype=complex,
+    )
+    Ad = A.to_dense()
+    if which_A == "upper":
+        Ad = Ad.T
+    Bd = B.to_dense()
+    if which_B == "lower":
+        Bd = Bd.T
+    C = Ad @ Bd
+    if which_B == "lower":
+        C = C.T
+
+    AB = qtn.tensor_network_apply_op_op(
+        A,
+        B,
+        which_A,
+        which_B,
+        inplace=inplace,
+        contract=contract,
+    )
+
+    if contract:
+        # checks fusing
+        assert AB.num_indices == B.num_indices
+
+    if inplace:
+        assert AB is B
+    else:
+        assert isinstance(AB, B.__class__)
+        assert AB.upper_ind_id == B.upper_ind_id
+        assert AB.lower_ind_id == B.lower_ind_id
+
+    assert_allclose(AB.to_dense(), C)
+
+
+def test_gate_with_op():
+    A = qtn.MPO_rand(5, 3, dtype=complex)
+    x = qtn.MPS_rand_state(5, 3, dtype=complex)
+    y = A.to_dense() @ x.to_dense()
+    x.gate_with_op_lazy_(A)
+    assert_allclose(x.to_dense(), y)
+
+
+def test_gate_sandwich_with_op():
+    B = qtn.MPO_rand(5, 3, dtype=complex)
+    A = qtn.MPO_rand(5, 3, dtype=complex)
+    y = A.to_dense() @ B.to_dense() @ A.to_dense().conj().T
+    B.gate_sandwich_with_op_lazy_(A)
+    assert_allclose(B.to_dense(), y)
+
+
+def test_normalize_simple():
+    psi = qtn.PEPS.rand(3, 3, 2, dtype=complex)
+    gauges = {}
+    psi.gauge_all_simple_(100, 5e-6, gauges=gauges)
+    psi.normalize_simple(gauges)
+
+    for where in [
+        [(0, 0)],
+        [(1, 1), (1, 2)],
+        [(0, 0), (0, 1), (0, 2), (1, 2), (2, 2), (2, 1)],
+    ]:
+        tags = [psi.site_tag(w) for w in where]
+        k = psi.select_any(tags, virtual=False)
+        k.gauge_simple_insert(gauges)
+
+        assert k.H @ k == pytest.approx(1.0)
+
+
+def test_local_expectation_loop_expansions():
+    import quimb as qu
+
+    edges = [(0, 1), (0, 2), (2, 3), (1, 3), (2, 4), (3, 5), (4, 5)]
+    psi = qtn.TN_from_edges_rand(
+        edges,
+        D=3,
+        phys_dim=2,
+        seed=42,
+        dist="uniform",
+        loc=-0.1,
+    )
+    G = qu.rand_herm(4)
+    where = (0, 2)
+    o_ex = psi.local_expectation_exact(G, where)
+
+    gauges = {}
+    psi.gauge_all_simple_(100, 5e-6, gauges=gauges)
+    psi.normalize_simple(gauges)
+
+    # test loop generation per term
+    o_c0 = psi.local_expectation_loop_expansion(
+        G, where, loops=0, gauges=gauges
+    )
+    assert o_c0 == pytest.approx(
+        psi.local_expectation_cluster(G, where, gauges=gauges)
+    )
+    assert o_ex == pytest.approx(o_c0, rel=0.5, abs=0.01)
+    o_c1 = psi.local_expectation_loop_expansion(
+        G, where, loops=4, gauges=gauges
+    )
+    assert o_ex == pytest.approx(o_c1, rel=0.5, abs=0.01)
+    o_c2 = psi.local_expectation_loop_expansion(
+        G, where, loops=6, gauges=gauges
+    )
+    assert o_ex == pytest.approx(o_c2, rel=0.4, abs=0.01)
+
+    # test manual loops supply
+    loops = tuple(psi.gen_paths_loops(6))
+    o_cl = psi.local_expectation_loop_expansion(
+        G, where, loops=loops, gauges=gauges
+    )
+    assert o_ex == pytest.approx(o_cl, rel=0.4, abs=0.01)
+
+
+def test_local_expectation_cluster_expansions():
+    import quimb as qu
+
+    edges = [(0, 1), (0, 2), (2, 3), (1, 3), (2, 4), (3, 5), (4, 5)]
+    psi = qtn.TN_from_edges_rand(
+        edges,
+        D=3,
+        phys_dim=2,
+        seed=42,
+        dist="uniform",
+        loc=-0.1,
+    )
+    G = qu.rand_herm(4)
+    where = (0, 2)
+    o_ex = psi.local_expectation_exact(G, where)
+
+    gauges = {}
+    psi.gauge_all_simple_(100, 5e-6, gauges=gauges)
+    psi.normalize_simple(gauges)
+
+    # test cluster generation per term
+    o_c0 = psi.local_expectation_cluster_expansion(
+        G, where, clusters=0, gauges=gauges
+    )
+    assert o_c0 == pytest.approx(
+        psi.local_expectation_cluster(G, where, gauges=gauges)
+    )
+    assert o_ex == pytest.approx(o_c0, rel=0.5, abs=0.01)
+    o_c1 = psi.local_expectation_cluster_expansion(
+        G, where, clusters=4, gauges=gauges
+    )
+    assert o_ex == pytest.approx(o_c1, rel=0.5, abs=0.01)
+    o_c2 = psi.local_expectation_cluster_expansion(
+        G, where, clusters=6, gauges=gauges
+    )
+    assert o_ex == pytest.approx(o_c2, rel=0.4, abs=0.01)
+
+    # test manual clusters supply
+    clusters = tuple(psi.gen_regions(4))
+    o_cl = psi.local_expectation_cluster_expansion(
+        G, where, clusters=clusters, gauges=gauges
+    )
+    assert o_ex == pytest.approx(o_cl, rel=0.4, abs=0.01)