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

trajectree/quimb/tests/test_tensor/test_tensor_dmrg.py

@@ -1,388 +0,0 @@
-import pytest
-
-import numpy as np
-from numpy.testing import assert_allclose
-
-from quimb import (
-    ham_heis,
-    expec,
-    plus,
-    is_eigenvector,
-    eigh,
-    heisenberg_energy,
-)
-
-from quimb.tensor import (
-    MPS_rand_state,
-    MPS_product_state,
-    MPS_computational_state,
-    MPO_ham_ising,
-    MPO_ham_XY,
-    MPO_ham_heis,
-    MPO_ham_mbl,
-    MovingEnvironment,
-    DMRG1,
-    DMRG2,
-    DMRGX,
-    SpinHam1D,
-)
-
-
-np.random.seed(42)
-
-
-class TestMovingEnvironment:
-    def test_bsz1_start_left(self):
-        tn = MPS_rand_state(6, bond_dim=7)
-        env = MovingEnvironment(tn, begin="left", bsz=1)
-        assert env.pos == 0
-        assert len(env().tensors) == 3
-        env.move_right()
-        assert env.pos == 1
-        assert len(env().tensors) == 3
-        env.move_right()
-        assert env.pos == 2
-        assert len(env().tensors) == 3
-        env.move_to(5)
-        assert env.pos == 5
-        assert len(env().tensors) == 3
-
-    def test_bsz1_start_right(self):
-        tn = MPS_rand_state(6, bond_dim=7)
-        env = MovingEnvironment(tn, begin="right", bsz=1)
-        assert env.pos == 5
-        assert len(env().tensors) == 3
-        env.move_left()
-        assert env.pos == 4
-        assert len(env().tensors) == 3
-        env.move_left()
-        assert env.pos == 3
-        assert len(env().tensors) == 3
-        env.move_to(0)
-        assert env.pos == 0
-        assert len(env().tensors) == 3
-
-    def test_bsz2_start_left(self):
-        tn = MPS_rand_state(6, bond_dim=7)
-        env = MovingEnvironment(tn, begin="left", bsz=2)
-        assert len(env().tensors) == 4
-        env.move_right()
-        assert len(env().tensors) == 4
-        env.move_right()
-        assert len(env().tensors) == 4
-        with pytest.raises(ValueError):
-            env.move_to(5)
-        env.move_to(4)
-        assert env.pos == 4
-        assert len(env().tensors) == 4
-
-    def test_bsz2_start_right(self):
-        tn = MPS_rand_state(6, bond_dim=7)
-        env = MovingEnvironment(tn, begin="right", bsz=2)
-        assert env.pos == 4
-        assert len(env().tensors) == 4
-        env.move_left()
-        assert env.pos == 3
-        assert len(env().tensors) == 4
-        env.move_left()
-        assert env.pos == 2
-        assert len(env().tensors) == 4
-        with pytest.raises(ValueError):
-            env.move_to(-1)
-        env.move_to(0)
-        assert env.pos == 0
-        assert len(env().tensors) == 4
-
-    @pytest.mark.parametrize("n", [20, 19])
-    @pytest.mark.parametrize("bsz", [1, 2])
-    @pytest.mark.parametrize("ssz", [1 / 2, 1.0])
-    def test_cyclic_moving_env_init_left(self, n, bsz, ssz):
-        nenv = 2
-        p = MPS_rand_state(n, 4, cyclic=True)
-        norm = p.H & p
-        mes = MovingEnvironment(
-            norm, begin="left", bsz=bsz, cyclic=True, ssz=ssz
-        )
-        assert len(mes.envs) == n // 2 + n % 2
-        assert mes.pos == 0
-        assert len(mes.envs[0].tensors) == 2 * bsz + nenv
-        assert len(mes.envs[n // 2 - 1].tensors) == 2 * bsz + 1
-        assert n // 2 + n % 2 not in mes.envs
-        assert n - 1 not in mes.envs
-
-        for i in range(1, 2 * n):
-            mes.move_right()
-            assert mes.pos == i % n
-            cur_env = mes()
-            assert len(cur_env.tensors) == 2 * bsz + nenv
-            assert (cur_env ^ all) == pytest.approx(1.0)
-
-    @pytest.mark.parametrize("n", [20, 19])
-    @pytest.mark.parametrize("bsz", [1, 2])
-    @pytest.mark.parametrize("ssz", [1 / 2, 1.0])
-    def test_cyclic_moving_env_init_right(self, n, bsz, ssz):
-        p = MPS_rand_state(n, 4, cyclic=True)
-        norm = p.H | p
-        mes = MovingEnvironment(
-            norm, begin="right", bsz=bsz, cyclic=True, ssz=ssz
-        )
-        assert len(mes.envs) == n // 2 + n % 2
-        assert mes.pos == n - 1
-        assert len(mes.envs[n - 1].tensors) == 2 * bsz + 2
-        assert len(mes.envs[n - n // 2].tensors) == 2 * bsz + 1
-        assert 0 not in mes.envs
-        assert n // 2 - 1 not in mes.envs
-
-        for i in reversed(range(-n, n - 1)):
-            mes.move_left()
-            assert mes.pos == i % n
-            cur_env = mes()
-            assert len(cur_env.tensors) == 2 * bsz + 2
-            assert (cur_env ^ all) == pytest.approx(1.0)
-
-
-class TestDMRG1:
-    def test_single_explicit_sweep(self):
-        h = MPO_ham_heis(5)
-        dmrg = DMRG1(h, bond_dims=3)
-        assert dmrg._k[0].dtype == float
-
-        energy_tn = dmrg._b | dmrg.ham | dmrg._k
-
-        e0 = energy_tn ^ ...
-        assert abs(e0.imag) < 1e-13
-
-        de1 = dmrg.sweep_right()
-        e1 = energy_tn ^ ...
-        assert_allclose(de1, e1)
-        assert abs(e1.imag) < 1e-13
-
-        de2 = dmrg.sweep_right()
-        e2 = energy_tn ^ ...
-        assert_allclose(de2, e2)
-        assert abs(e2.imag) < 1e-13
-
-        # state is already left canonized after right sweep
-        de3 = dmrg.sweep_left(canonize=False)
-        e3 = energy_tn ^ ...
-        assert_allclose(de3, e3)
-        assert abs(e2.imag) < 1e-13
-
-        de4 = dmrg.sweep_left()
-        e4 = energy_tn ^ ...
-        assert_allclose(de4, e4)
-        assert abs(e2.imag) < 1e-13
-
-        # test still normalized
-        assert dmrg._k[0].dtype == float
-        dmrg._k.align_(dmrg._b)
-        assert_allclose(abs(dmrg._b @ dmrg._k), 1)
-
-        assert e1.real < e0.real
-        assert e2.real < e1.real
-        assert e3.real < e2.real
-        assert e4.real < e3.real
-
-    @pytest.mark.parametrize("dense", [False, True])
-    @pytest.mark.parametrize("MPO_ham", [MPO_ham_XY, MPO_ham_heis])
-    @pytest.mark.parametrize("cyclic", [False, True])
-    def test_ground_state_matches(self, dense, MPO_ham, cyclic):
-        n = 10
-
-        tol = 3e-2 if cyclic else 1e-4
-
-        h = MPO_ham(n, cyclic=cyclic)
-        dmrg = DMRG1(h, bond_dims=[4, 8, 12])
-        dmrg.opts["local_eig_ham_dense"] = dense
-        dmrg.opts["periodic_segment_size"] = 1.0
-        dmrg.opts["periodic_nullspace_fudge_factor"] = 1e-6
-        assert dmrg.solve(tol=tol / 10, verbosity=1)
-        assert dmrg.state.cyclic == cyclic
-        eff_e, mps_gs = dmrg.energy, dmrg.state
-        mps_gs_dense = mps_gs.to_qarray()
-
-        assert_allclose(mps_gs_dense.H @ mps_gs_dense, 1.0, rtol=tol)
-
-        h_dense = h.to_qarray()
-
-        # check against dense form
-        actual_e, gs = eigh(h_dense, k=1)
-        assert_allclose(actual_e, eff_e, rtol=tol)
-        assert_allclose(abs(expec(mps_gs_dense, gs)), 1.0, rtol=tol)
-
-        # check against actual MPO_ham
-        if MPO_ham is MPO_ham_XY:
-            ham_dense = ham_heis(
-                n, cyclic=cyclic, j=(1.0, 1.0, 0.0), sparse=True
-            )
-        elif MPO_ham is MPO_ham_heis:
-            ham_dense = ham_heis(n, cyclic=cyclic, sparse=True)
-
-        actual_e, gs = eigh(ham_dense, k=1)
-        assert_allclose(actual_e, eff_e, rtol=tol)
-        assert_allclose(abs(expec(mps_gs_dense, gs)), 1.0, rtol=tol)
-
-    def test_ising_and_MPS_product_state(self):
-        h = MPO_ham_ising(6, bx=2.0, j=0.1)
-        dmrg = DMRG1(h, bond_dims=8)
-        assert dmrg.solve(verbosity=1)
-        eff_e, mps_gs = dmrg.energy, dmrg.state
-        mps_gs_dense = mps_gs.to_qarray()
-        assert_allclose(mps_gs_dense.H @ mps_gs_dense, 1.0)
-
-        # check against dense
-        h_dense = h.to_qarray()
-        actual_e, gs = eigh(h_dense, k=1)
-        assert_allclose(actual_e, eff_e)
-        assert_allclose(abs(expec(mps_gs_dense, gs)), 1.0)
-
-        exp_gs = MPS_product_state([plus()] * 6)
-        assert_allclose(abs(exp_gs.H @ mps_gs), 1.0, rtol=1e-3)
-
-
-class TestDMRG2:
-    @pytest.mark.parametrize("dense", [False, True])
-    @pytest.mark.parametrize("MPO_ham", [MPO_ham_XY, MPO_ham_heis])
-    @pytest.mark.parametrize("cyclic", [False, True])
-    def test_matches_exact(self, dense, MPO_ham, cyclic):
-        n = 6
-        h = MPO_ham(n, cyclic=cyclic)
-
-        tol = 3e-2 if cyclic else 1e-4
-
-        dmrg = DMRG2(h, bond_dims=[4, 8, 12])
-        assert dmrg._k[0].dtype == float
-        dmrg.opts["local_eig_ham_dense"] = dense
-        dmrg.opts["periodic_segment_size"] = 1.0
-        dmrg.opts["periodic_nullspace_fudge_factor"] = 1e-6
-
-        assert dmrg.solve(tol=tol / 10, verbosity=1)
-
-        # XXX: need to dispatch SLEPc eigh on real input
-        # assert dmrg._k[0].dtype == float
-
-        eff_e, mps_gs = dmrg.energy, dmrg.state
-        mps_gs_dense = mps_gs.to_qarray()
-
-        assert_allclose(expec(mps_gs_dense, mps_gs_dense), 1.0, rtol=tol)
-
-        h_dense = h.to_qarray()
-
-        # check against dense form
-        actual_e, gs = eigh(h_dense, k=1)
-        assert_allclose(actual_e, eff_e, rtol=tol)
-        assert_allclose(abs(expec(mps_gs_dense, gs)), 1.0, rtol=tol)
-
-        # check against actual MPO_ham
-        if MPO_ham is MPO_ham_XY:
-            ham_dense = ham_heis(n, cyclic=cyclic, j=(1.0, 1.0, 0.0))
-        elif MPO_ham is MPO_ham_heis:
-            ham_dense = ham_heis(n, cyclic=cyclic)
-
-        actual_e, gs = eigh(ham_dense, k=1)
-        assert_allclose(actual_e, eff_e, rtol=tol)
-        assert_allclose(abs(expec(mps_gs_dense, gs)), 1.0, rtol=tol)
-
-    def test_cyclic_solve_big_with_segmenting(self):
-        n = 150
-        ham = MPO_ham_heis(n, cyclic=True)
-        dmrg = DMRG2(ham, bond_dims=range(10, 30, 2))
-        dmrg.opts["periodic_segment_size"] = 1 / 3
-        assert dmrg.solve(tol=1, verbosity=2)
-        assert dmrg.energy == pytest.approx(heisenberg_energy(n), 1e-3)
-
-    @pytest.mark.parametrize(
-        "dtype", [np.float32, np.float64, np.complex64, np.complex128]
-    )
-    def test_dtypes(self, dtype):
-        H = MPO_ham_heis(8).astype(dtype)
-        dmrg = DMRG2(H)
-        dmrg.opts["local_eig_backend"] = "scipy"
-        dmrg.solve(max_sweeps=3)
-        (res_dtype,) = {t.dtype for t in dmrg.state}
-        assert res_dtype == dtype
-
-    def test_total_size_2(self):
-        N = 2
-        builder = SpinHam1D(1 / 2)
-        for i in range(N - 1):
-            builder[i, i + 1] += 1.0, "Z", "Z"
-
-        H = builder.build_mpo(N)
-        dmrg = DMRG2(H)
-        dmrg.solve(verbosity=1)
-        assert dmrg.energy == pytest.approx(-1 / 4)
-
-    def test_variable_bond_ham(self):
-        import quimb as qu
-
-        HB = SpinHam1D(1 / 2)
-        HB[0, 1] += 0.6, "Z", "Z"
-        HB[1, 2] += 0.7, "Z", "Z"
-        HB[1, 2] += 0.8, "X", "X"
-        HB[2, 3] += 0.9, "Y", "Y"
-
-        H_mpo = HB.build_mpo(4)
-        H_sps = HB.build_sparse(4)
-
-        assert H_mpo.bond_sizes() == [3, 4, 3]
-
-        Sx, Sy, Sz = map(qu.spin_operator, "xyz")
-        H_explicit = (
-            qu.ikron(0.6 * Sz & Sz, [2, 2, 2, 2], [0, 1])
-            + qu.ikron(0.7 * Sz & Sz, [2, 2, 2, 2], [1, 2])
-            + qu.ikron(0.8 * Sx & Sx, [2, 2, 2, 2], [1, 2])
-            + qu.ikron(0.9 * Sy & Sy, [2, 2, 2, 2], [2, 3])
-        )
-
-        assert_allclose(H_explicit, H_mpo.to_qarray())
-        assert_allclose(H_explicit, H_sps.toarray())
-
-
-class TestDMRGX:
-    def test_explicit_sweeps(self):
-        n = 8
-        chi = 16
-        ham = MPO_ham_mbl(n, dh=4, seed=42)
-        p0 = MPS_rand_state(n, 2).expand_bond_dimension(chi)
-
-        b0 = p0.H
-        p0.align_(ham, b0)
-        en0 = (p0 & ham & b0) ^ ...
-        dmrgx = DMRGX(ham, p0, chi)
-        dmrgx.sweep_right()
-        en1 = dmrgx.sweep_left(canonize=False)
-        assert en0 != en1
-
-        dmrgx.sweep_right(canonize=False)
-        en = dmrgx.sweep_right(canonize=True)
-
-        # check normalized
-        assert_allclose(dmrgx._k.H @ dmrgx._k, 1.0)
-
-        k = dmrgx._k.to_qarray()
-        h = ham.to_qarray()
-        el, ev = eigh(h)
-
-        # check variance very low
-        assert np.abs((k.H @ h @ h @ k) - (k.H @ h @ k) ** 2) < 1e-12
-
-        # check exactly one eigenvalue matched well
-        assert np.sum(np.abs(el - en) < 1e-12) == 1
-
-        # check exactly one eigenvector is matched with high fidelity
-        ovlps = (ev.H @ k).toarray() ** 2
-        big_ovlps = ovlps[ovlps > 1e-12]
-        assert_allclose(big_ovlps, [1])
-
-        # check fully
-        assert is_eigenvector(k, h, tol=1e-10)
-
-    def test_solve_bigger(self):
-        n = 14
-        chi = 16
-        ham = MPO_ham_mbl(n, dh=8, seed=42)
-        p0 = MPS_computational_state("00110111000101")
-        dmrgx = DMRGX(ham, p0, chi)
-        assert dmrgx.solve(tol=1e-5, sweep_sequence="R")
-        assert dmrgx.state[0].dtype == float

trajectree/quimb/tests/test_tensor/test_tensor_spectral_approx.py

@@ -1,63 +0,0 @@
-from math import log2
-import pytest
-import numpy as np
-from numpy.testing import assert_allclose
-
-from quimb import (
-    seed_rand,
-    approx_spectral_function,
-    eigvalsh,
-    logneg_subsys,
-)
-from quimb.tensor import MPO_rand_herm, MPO_ham_heis, DMRG2
-from quimb.tensor.tensor_approx_spectral import construct_lanczos_tridiag_MPO
-
-
-np.random.seed(42)
-
-
-# XXX: these all need improvement
-
-
-class TestMPOSpectralApprox:
-    def test_constructing_tridiag_works(self):
-        A = MPO_rand_herm(10, 7)
-        for _ in construct_lanczos_tridiag_MPO(A, 5):
-            pass
-
-    @pytest.mark.parametrize("fn", [abs, np.cos, lambda x: np.sin(x) ** 2])
-    def test_approx_fn(self, fn):
-        A = MPO_rand_herm(10, 7, normalize=True)
-        xe = sum(fn(eigvalsh(A.to_dense())))
-        xf = approx_spectral_function(A, fn, tol=0.1, verbosity=2)
-        assert_allclose(xe, xf, rtol=0.5)
-
-    def test_realistic(self):
-        seed_rand(42)
-        ham = MPO_ham_heis(20)
-        dmrg = DMRG2(ham, bond_dims=[2, 4])
-        dmrg.solve()
-        rho_ab = dmrg.state.partial_trace_to_mpo(range(6, 14))
-        xf = approx_spectral_function(
-            rho_ab, lambda x: x, tol=0.1, verbosity=2
-        )
-        assert_allclose(1.0, xf, rtol=0.6, atol=0.001)
-
-    def test_realistic_ent(self):
-        n = 12
-        sysa, sysb = range(3, 6), range(6, 8)
-        sysab = (*sysa, *sysb)
-
-        ham = MPO_ham_heis(n)
-        dmrg = DMRG2(ham, bond_dims=[10])
-        dmrg.solve()
-
-        psi0 = dmrg.state.to_dense()
-        lne = logneg_subsys(psi0, [2] * n, sysa=sysa, sysb=sysb)
-
-        rho_ab = dmrg.state.partial_trace_to_mpo(sysab, rescale_sites=True)
-        rho_ab_pt = rho_ab.partial_transpose(range(3))
-        lnx = log2(
-            approx_spectral_function(rho_ab_pt, abs, tol=0.1, verbosity=2)
-        )
-        assert_allclose(lne, lnx, rtol=0.6, atol=0.1)