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

trajectree/quimb/tests/test_gen/test_states.py

@@ -0,0 +1,261 @@
+from pytest import mark
+import numpy as np
+from numpy.testing import assert_allclose
+from quimb import (
+    tr,
+    eye,
+    chop,
+    ikron,
+    expec,
+    ptr,
+    eigh,
+    groundstate,
+    concurrence,
+    basis_vec,
+    up,
+    down,
+    plus,
+    minus,
+    yplus,
+    yminus,
+    thermal_state,
+    neel_state,
+    ham_j1j2,
+    rand_herm,
+    graph_state_1d,
+    pauli,
+    levi_civita,
+    bloch_state,
+    bell_state,
+    singlet,
+    singlet_pairs,
+    werner_state,
+    ghz_state,
+    w_state,
+    perm_state,
+)
+
+
+class TestBasisVec:
+    def test_basis_vec(self):
+        x = basis_vec(1, 2)
+        assert_allclose(x, [[0.0], [1.0]])
+        x = basis_vec(1, 2, qtype="b")
+        assert_allclose(x, [[0.0, 1.0]])
+
+
+class TestBasicStates:
+    def test_up(self):
+        p = up(qtype="dop")
+        assert_allclose(tr(p @ pauli("z")), 1.0)
+
+    def test_down(self):
+        p = down(qtype="dop")
+        assert_allclose(tr(p @ pauli("z")), -1.0)
+
+    def test_plus(self):
+        p = plus(qtype="dop")
+        assert_allclose(tr(p @ pauli("x")), 1.0)
+
+    def test_minus(self):
+        p = minus(qtype="dop")
+        assert_allclose(tr(p @ pauli("x")), -1.0)
+
+    def test_yplus(self):
+        p = yplus(qtype="dop")
+        assert_allclose(tr(p @ pauli("y")), 1.0)
+
+    def test_yminus(self):
+        p = yminus(qtype="dop")
+        assert_allclose(tr(p @ pauli("y")), -1.0)
+
+
+class TestBlochState:
+    def test_pure(self):
+        for vec, op, val in zip(
+            (
+                (1, 0, 0),
+                (0, 1, 0),
+                (0, 0, 1),
+                (-1, 0, 0),
+                (0, -1, 0),
+                (0, 0, -1),
+            ),
+            ("x", "y", "z", "x", "y", "z"),
+            (1, 1, 1, -1, -1, -1),
+        ):
+            x = tr(bloch_state(*vec) @ pauli(op))
+            assert_allclose(x, val)
+
+    def test_mixed(self):
+        for vec, op, val in zip(
+            (
+                (0.5, 0, 0),
+                (0, 0.5, 0),
+                (0, 0, 0.5),
+                (-0.5, 0, 0),
+                (0, -0.5, 0),
+                (0, 0, -0.5),
+            ),
+            ("x", "y", "z", "x", "y", "z"),
+            (0.5, 0.5, 0.5, -0.5, -0.5, -0.5),
+        ):
+            x = tr(bloch_state(*vec) @ pauli(op))
+            assert_allclose(x, val)
+
+    def test_purify(self):
+        for vec, op, val in zip(
+            (
+                (0.5, 0, 0),
+                (0, 0.5, 0),
+                (0, 0, 0.5),
+                (-0.5, 0, 0),
+                (0, -0.5, 0),
+                (0, 0, -0.5),
+            ),
+            ("x", "y", "z", "x", "y", "z"),
+            (1, 1, 1, -1, -1, -1),
+        ):
+            x = tr(bloch_state(*vec, purified=True) @ pauli(op))
+            assert_allclose(x, val)
+
+
+class TestBellStates:
+    def test_bell_states(self):
+        for s, dic in zip(
+            ("psi-", "psi+", "phi+", "phi-"),
+            ({"qtype": "dop"}, {}, {"sparse": True}, {}),
+        ):
+            p = bell_state(s, **dic)
+            assert_allclose(expec(p, p), 1.0)
+            pa = ptr(p, [2, 2], 0)
+            assert_allclose(expec(pa, pa), 0.5)
+
+    def test_bell_state_singlet(self):
+        p = singlet(qtype="dop", sparse=True)
+        assert_allclose(expec(p, p), 1.0)
+        pa = ptr(p, [2, 2], 0)
+        assert_allclose(expec(pa, pa), 0.5)
+
+
+class TestThermalState:
+    def test_thermal_state_normalization(self):
+        full = rand_herm(2**4)
+        for beta in (0, 0.5, 1, 10):
+            rhoth = thermal_state(full, beta)
+            assert_allclose(tr(rhoth), 1)
+
+    def test_thermal_state_tuple(self):
+        full = rand_herm(2**4)
+        l, v = eigh(full)
+        for beta in (0, 0.5, 1, 10):
+            rhoth1 = thermal_state(full, beta)
+            rhoth2 = thermal_state((l, v), beta)
+            assert_allclose(rhoth1, rhoth2)
+
+    def test_thermal_state_hot(self):
+        full = rand_herm(2**4)
+        rhoth = chop(thermal_state(full, 0.0))
+        assert_allclose(rhoth, eye(2**4) / 2**4)
+
+    def test_thermal_state_cold(self):
+        full = ham_j1j2(4, j2=0.1253)
+        rhoth = thermal_state(full, 100)
+        gs = groundstate(full)
+        assert_allclose(tr(gs.H @ rhoth @ gs), 1.0, rtol=1e-4)
+
+    def test_thermal_state_precomp(self):
+        full = rand_herm(2**4)
+        beta = 0.624
+        rhoth1 = thermal_state(full, beta)
+        func = thermal_state(full, None, precomp_func=True)
+        rhoth2 = func(beta)
+        assert_allclose(rhoth1, rhoth2)
+
+
+class TestNeelState:
+    def test_simple(self):
+        p = neel_state(1)
+        assert_allclose(p, up())
+        p = neel_state(2)
+        assert_allclose(p, up() & down())
+        p = neel_state(3)
+        assert_allclose(p, up() & down() & up())
+
+
+class TestSingletPairs:
+    def test_n2(self):
+        p = singlet_pairs(2)
+        assert_allclose(p, bell_state("psi-"))
+        p = singlet_pairs(4)
+        assert_allclose(p, bell_state("psi-") & bell_state("psi-"))
+
+
+class TestWernerState:
+    def test_extremes(self):
+        p = werner_state(1)
+        assert_allclose(p, bell_state("psi-", qtype="dop"))
+        p = werner_state(0)
+        assert_allclose(p, eye(4) / 4, atol=1e-14)
+
+    @mark.parametrize("p", np.linspace(0, 1 / 3, 10))
+    def test_no_entanglement(self, p):
+        rho = werner_state(p)
+        assert concurrence(rho) < 1e-14
+
+    @mark.parametrize("p", np.linspace(1 / 3 + 0.001, 1.0, 10))
+    def test_entanglement(self, p):
+        rho = werner_state(p)
+        assert concurrence(rho) > 1e-14
+
+
+class TestGHZState:
+    def test_n2(self):
+        p = ghz_state(2)
+        assert_allclose(p, bell_state("phi+"))
+
+
+class TestWState:
+    def test_n2(self):
+        p = w_state(2)
+        assert_allclose(p, bell_state("psi+"))
+
+
+class TestLeviCivita:
+    def test_levi_civita_pos(self):
+        perm = [0, 1, 2, 3]
+        assert levi_civita(perm) == 1
+        perm = [2, 3, 0, 1]
+        assert levi_civita(perm) == 1
+
+    def test_levi_civita_neg(self):
+        perm = [0, 2, 1, 3]
+        assert levi_civita(perm) == -1
+        perm = [2, 3, 1, 0]
+        assert levi_civita(perm) == -1
+
+    def test_levi_civita_nzero(self):
+        perm = [2, 3, 1, 1]
+        assert levi_civita(perm) == 0
+        perm = [0, 0, 1, 1]
+        assert levi_civita(perm) == 0
+
+
+class TestPermState:
+    def test_n2(self):
+        p = perm_state([up(), down()])
+        assert_allclose(p, bell_state("psi-"))
+
+
+class TestGraphState:
+    def test_graph_state_1d(self):
+        n = 5
+        p = graph_state_1d(n, cyclic=True)
+        for j in range(n):
+            k = ikron(
+                [pauli("x"), pauli("z"), pauli("z")],
+                dims=[2] * n,
+                inds=(j, (j - 1) % n, (j + 1) % n),
+            )
+            o = p.H @ k @ p
+            np.testing.assert_allclose(o, 1)

trajectree/quimb/tests/test_linalg/test_approx_spectral.py

@@ -0,0 +1,368 @@
+import pytest
+import numpy as np
+from numpy.testing import assert_allclose
+
+from quimb import (
+    prod,
+    rand_ket,
+    rand_pos,
+    rand_herm,
+    partial_transpose,
+    eigvalsh,
+    neel_state,
+    ham_heis,
+    logneg,
+    negativity,
+    entropy,
+    can_use_mpi_pool,
+)
+from quimb.utils import last
+
+from quimb.linalg.approx_spectral import (
+    lazy_ptr_ppt_linop,
+    construct_lanczos_tridiag,
+    lanczos_tridiag_eig,
+    approx_spectral_function,
+    lazy_ptr_linop,
+    tr_abs_approx,
+    tr_exp_approx,
+    tr_sqrt_approx,
+    tr_xlogx_approx,
+    entropy_subsys_approx,
+    logneg_subsys_approx,
+    negativity_subsys_approx,
+    norm_fro,
+    norm_fro_approx,
+)
+
+MPI_PARALLEL = [False]
+if can_use_mpi_pool():
+    MPI_PARALLEL.append(True)
+
+
+np.random.seed(42)
+
+DIMS = [5, 6, 7]
+DIMS_MB = [2] * 11
+
+
+@pytest.fixture
+def psi_abc():
+    return rand_ket(prod(DIMS))
+
+
+@pytest.fixture
+def psi_ab():
+    return rand_ket(prod(DIMS[:-1]))
+
+
+@pytest.fixture
+def psi_ab_mat():
+    return np.concatenate(
+        [rand_ket(prod(DIMS[:-1])) for _ in range(6)], axis=1
+    )
+
+
+@pytest.fixture
+def psi_mb_abc():
+    return rand_ket(prod(DIMS_MB))
+
+
+@pytest.fixture
+def psi_mb_ab():
+    return rand_ket(prod(DIMS_MB[:7]))
+
+
+@pytest.fixture
+def psi_mb_ab_mat():
+    return np.concatenate(
+        [rand_ket(prod(DIMS_MB[:7])) for _ in range(6)], axis=1
+    )
+
+
+class TestLazyTensorEval:
+    # ------------------------- Just partial trace -------------------------- #
+
+    def test_lazy_ptr_dot_simple_linear_op(self, psi_abc, psi_ab):
+        rho_ab = psi_abc.ptr(DIMS, [0, 1])
+        psi_out_expected = rho_ab @ psi_ab
+        lo = lazy_ptr_linop(psi_abc, DIMS, [0, 1])
+        assert hasattr(lo, "H")
+        assert lo.dtype == complex
+        assert lo.shape == (DIMS[0] * DIMS[1], DIMS[0] * DIMS[1])
+        psi_out_got = lo @ psi_ab
+        assert_allclose(psi_out_expected, psi_out_got)
+
+    def test_lazy_ptr_dot_mat_simple_linear_op(self, psi_abc, psi_ab_mat):
+        rho_ab = psi_abc.ptr(DIMS, [0, 1])
+        psi_out_expected = rho_ab @ psi_ab_mat
+        lo = lazy_ptr_linop(psi_abc, DIMS, [0, 1])
+        psi_out_got = lo @ psi_ab_mat
+        assert_allclose(psi_out_expected, psi_out_got)
+
+    def test_lazy_ptr_dot_manybody_linear_op(self, psi_mb_abc, psi_mb_ab):
+        sysa = [0, 1, 2, 3, 7, 8, 9]
+        rho_ab = psi_mb_abc.ptr(DIMS_MB, sysa)
+        psi_out_expected = rho_ab @ psi_mb_ab
+        lo = lazy_ptr_linop(psi_mb_abc, DIMS_MB, sysa=sysa)
+        assert hasattr(lo, "H")
+        assert lo.dtype == complex
+        assert lo.shape == (128, 128)
+        psi_out_got = lo @ psi_mb_ab
+        assert_allclose(psi_out_expected, psi_out_got)
+
+    def test_lazy_ptr_dot_mat_manybody_linear_op(
+        self, psi_mb_abc, psi_mb_ab_mat
+    ):
+        sysa = [0, 1, 2, 3, 7, 8, 9]
+        rho_ab = psi_mb_abc.ptr(DIMS_MB, sysa)
+        psi_out_expected = rho_ab @ psi_mb_ab_mat
+        lo = lazy_ptr_linop(psi_mb_abc, DIMS_MB, sysa=sysa)
+        psi_out_got = lo @ psi_mb_ab_mat
+        assert_allclose(psi_out_expected, psi_out_got)
+
+    # ----------------- partial trace and partial transpose ----------------- #
+
+    def test_lazy_ptr_ppt_dot_linear_op(self, psi_abc, psi_ab):
+        rho_ab = psi_abc.ptr(DIMS, [0, 1])
+        rho_ab_pt = partial_transpose(rho_ab, DIMS[:-1])
+        psi_out_expected = rho_ab_pt @ psi_ab
+        lo = lazy_ptr_ppt_linop(psi_abc, DIMS, 0, 1)
+        assert hasattr(lo, "H")
+        assert lo.dtype == complex
+        assert lo.shape == (DIMS[0] * DIMS[1], DIMS[0] * DIMS[1])
+        psi_out_got = lo @ psi_ab
+        assert_allclose(psi_out_expected, psi_out_got)
+
+    def test_lazy_ptr_ppt_dot_mat_linear_op(self, psi_abc, psi_ab_mat):
+        rho_ab = psi_abc.ptr(DIMS, [0, 1])
+        rho_ab_pt = partial_transpose(rho_ab, DIMS[:-1])
+        psi_out_expected = rho_ab_pt @ psi_ab_mat
+        lo = lazy_ptr_ppt_linop(psi_abc, DIMS, 0, 1)
+        psi_out_got = lo @ psi_ab_mat
+        assert_allclose(psi_out_expected, psi_out_got)
+
+    def test_lazy_ptr_ppt_dot_manybody_linear_op(self, psi_mb_abc, psi_mb_ab):
+        sysa = [0, 1, 7, 8]
+        sysb = [2, 3, 9]
+        rho_ab = psi_mb_abc.ptr(DIMS_MB, sysa + sysb)
+        rho_ab = partial_transpose(rho_ab, [2] * 7, sysa=(0, 1, 4, 5))
+        psi_out_expected = rho_ab @ psi_mb_ab
+        lo = lazy_ptr_ppt_linop(psi_mb_abc, DIMS_MB, sysa, sysb)
+        assert hasattr(lo, "H")
+        assert lo.dtype == complex
+        assert lo.shape == (128, 128)
+        psi_out_got = lo.dot(psi_mb_ab)
+        assert_allclose(psi_out_expected, psi_out_got)
+
+    def test_lazy_ptr_ppt_dot_mat_manybody_linear_op(
+        self, psi_mb_abc, psi_mb_ab_mat
+    ):
+        sysa = [0, 1, 7, 8]
+        sysb = [2, 3, 9]
+        rho_ab = psi_mb_abc.ptr(DIMS_MB, sysa + sysb)
+        rho_ab = partial_transpose(rho_ab, [2] * 7, sysa=(0, 1, 4, 5))
+        psi_out_expected = rho_ab @ psi_mb_ab_mat
+        lo = lazy_ptr_ppt_linop(psi_mb_abc, DIMS_MB, sysa, sysb)
+        psi_out_got = lo.dot(psi_mb_ab_mat)
+        assert psi_out_got.shape[1] > 1
+        assert_allclose(psi_out_expected, psi_out_got)
+
+
+# --------------------- Test lanczos spectral functions --------------------- #
+
+
+def np_sqrt(x):
+    out = np.empty_like(x)
+    mtz = x > 0
+    out[~mtz] = 0.0
+    out[mtz] = np.sqrt(x[mtz])
+    return out
+
+
+def np_xlogx(x):
+    out = np.empty_like(x)
+    mtz = x > 0
+    out[~mtz] = 0.0
+    out[mtz] = x[mtz] * np.log2(x[mtz])
+    return out
+
+
+class TestLanczosApprox:
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    def test_construct_lanczos_tridiag(self, bsz):
+        a = rand_herm(2**5)
+        alpha, beta, scaling = last(
+            construct_lanczos_tridiag(a, bsz=bsz, K=20)
+        )
+        assert alpha.shape == (20,)
+        assert beta.shape == (20,)
+
+        el, ev = lanczos_tridiag_eig(alpha, beta)
+        assert el.shape == (20,)
+        assert el.dtype == float
+        assert ev.shape == (20, 20)
+        assert ev.dtype == float
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    def test_construct_lanczos_tridiag_beta_breakdown(self, bsz):
+        a = rand_herm(8)
+        alpha, beta, scaling = last(construct_lanczos_tridiag(a, bsz=bsz, K=9))
+        assert alpha.shape == (8,)
+        assert beta.shape == (8,)
+
+        el, ev = lanczos_tridiag_eig(alpha, beta)
+        assert el.shape == (8,)
+        assert el.dtype == float
+        assert ev.shape == (8, 8)
+        assert ev.dtype == float
+
+    @pytest.mark.parametrize("mpi", MPI_PARALLEL)
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    @pytest.mark.parametrize(
+        "fn_matrix_rtol",
+        [
+            (np.abs, rand_herm, 1e-1),
+            (np.sqrt, rand_pos, 1e-1),
+            (np.log1p, rand_pos, 2e-1),
+            (np.exp, rand_herm, 1e-1),
+        ],
+    )
+    def test_approx_spectral_function(self, fn_matrix_rtol, bsz, mpi):
+        fn, matrix, rtol = fn_matrix_rtol
+        a = matrix(2**7)
+        pos = fn == np.sqrt
+        actual_x = sum(fn(eigvalsh(a)))
+        approx_x = approx_spectral_function(
+            a, fn, mpi=mpi, pos=pos, bsz=bsz, verbosity=2
+        )
+        assert_allclose(actual_x, approx_x, rtol=rtol)
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    @pytest.mark.parametrize("dist", ["gaussian", "phase", "rademacher"])
+    @pytest.mark.parametrize(
+        "fn_matrix_rtol",
+        [
+            (np.abs, rand_herm, 2e-1),
+            (np.sqrt, rand_pos, 2e-1),
+            (np.log1p, rand_pos, 3e-1),
+            (np.exp, rand_herm, 2e-1),
+        ],
+    )
+    def test_approx_spectral_function_with_v0(self, dist, fn_matrix_rtol, bsz):
+        fn, matrix, rtol = fn_matrix_rtol
+        a = matrix(2**7)
+        actual_x = sum(fn(eigvalsh(a)))
+        # check un-normalized state work properly
+        v0 = neel_state(7) + neel_state(7, down_first=True)
+        v0 = v0.toarray().reshape(-1)
+        pos = fn == np.sqrt
+        v0_opts = {"dist": dist, "orthog": True}
+        approx_x = approx_spectral_function(
+            a, fn, v0=v0, pos=pos, bsz=bsz, K=20, verbosity=2, v0_opts=v0_opts
+        )
+        assert_allclose(actual_x, approx_x, rtol=rtol)
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    @pytest.mark.parametrize(
+        "fn_approx_rtol",
+        [
+            (np_sqrt, tr_sqrt_approx, 0.3),
+            (np.exp, tr_exp_approx, 0.03),
+            (np_xlogx, tr_xlogx_approx, 0.3),
+        ],
+    )
+    def test_approx_spectral_function_ptr_lin_op(
+        self, fn_approx_rtol, psi_abc, psi_ab, bsz
+    ):
+        fn, approx, rtol = fn_approx_rtol
+        sysa = [0, 1]
+        rho_ab = psi_abc.ptr(DIMS, sysa)
+        actual_x = sum(fn(eigvalsh(rho_ab)))
+        lo = lazy_ptr_linop(psi_abc, DIMS, sysa)
+        approx_x = approx(lo, R=50, bsz=bsz, verbosity=2)
+        assert_allclose(actual_x, approx_x, rtol=rtol)
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    @pytest.mark.parametrize(
+        "fn_approx_rtol",
+        [(np.exp, tr_exp_approx, 5e-2), (np.abs, tr_abs_approx, 1e-1)],
+    )
+    def test_approx_spectral_function_ptr_ppt_lin_op(
+        self, fn_approx_rtol, psi_abc, psi_ab, bsz
+    ):
+        fn, approx, rtol = fn_approx_rtol
+        rho_ab_ppt = partial_transpose(psi_abc.ptr(DIMS, [0, 1]), DIMS[:-1], 0)
+        actual_x = sum(fn(eigvalsh(rho_ab_ppt)))
+        lo = lazy_ptr_ppt_linop(psi_abc, DIMS, sysa=0, sysb=1)
+        approx_x = approx_spectral_function(lo, fn, K=20, R=20, bsz=bsz)
+        assert_allclose(actual_x, approx_x, rtol=rtol)
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    def test_approx_spectral_subspaces_with_heis_partition(self, bsz):
+        h = ham_heis(10, sparse=True)
+        beta = 0.01
+        actual_Z = sum(np.exp(-beta * eigvalsh(h.toarray())))
+        approx_Z = tr_exp_approx(-beta * h, bsz=bsz)
+        assert_allclose(actual_Z, approx_Z, rtol=3e-2)
+
+    def test_approx_spectral_plot(self):
+        import matplotlib
+
+        matplotlib.use("Template")
+        X = rand_herm(1000, sparse=True)
+        approx_spectral_function(X, lambda x: abs(x), plot=True)
+
+
+# ------------------------ Test specific quantities ------------------------- #
+
+
+class TestSpecificApproxQuantities:
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    def test_entropy_approx_simple(self, psi_abc, bsz):
+        np.random.seed(42)
+        rho_ab = psi_abc.ptr(DIMS, [0, 1])
+        actual_e = entropy(rho_ab)
+        approx_e = entropy_subsys_approx(psi_abc, DIMS, [0, 1], bsz=bsz, R=20)
+        assert_allclose(actual_e, approx_e, rtol=2e-1)
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    def test_entropy_approx_many_body(self, psi_mb_abc, bsz):
+        rho_ab = psi_mb_abc.ptr(DIMS_MB, [0, 1, 7, 8, 2, 3, 9])
+        actual_e = entropy(rho_ab)
+        approx_e = entropy_subsys_approx(
+            psi_mb_abc, DIMS_MB, [0, 1, 7, 8, 2, 3, 9], bsz=bsz
+        )
+        assert_allclose(actual_e, approx_e, rtol=2e-1)
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    def test_logneg_approx_simple(self, psi_abc, bsz):
+        rho_ab = psi_abc.ptr(DIMS, [0, 1])
+        actual_ln = logneg(rho_ab, DIMS[:-1], 0)
+        approx_ln = logneg_subsys_approx(psi_abc, DIMS, 0, 1, bsz=bsz)
+        assert_allclose(actual_ln, approx_ln, rtol=2e-1)
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    def test_logneg_approx_many_body(self, psi_mb_abc, bsz):
+        sysa = [0, 1, 7, 8]
+        sysb = [2, 3, 9]
+        rho_ab = psi_mb_abc.ptr(DIMS_MB, sysa + sysb)
+        actual_ln = logneg(rho_ab, [2] * 7, sysa=(0, 1, 4, 5))
+        approx_ln = logneg_subsys_approx(
+            psi_mb_abc, DIMS_MB, sysa=sysa, sysb=sysb, bsz=bsz
+        )
+        assert_allclose(actual_ln, approx_ln, rtol=1e-1)
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    def test_negativity_approx_simple(self, psi_abc, bsz):
+        rho_ab = psi_abc.ptr(DIMS, [0, 1])
+        actual_neg = negativity(rho_ab, DIMS[:-1], 0)
+        approx_neg = negativity_subsys_approx(psi_abc, DIMS, 0, 1, bsz=bsz)
+        assert_allclose(actual_neg, approx_neg, rtol=2e-1)
+
+    @pytest.mark.parametrize("bsz", [1, 2, 5])
+    def test_norm_fro_approx(self, bsz):
+        A = rand_herm(2**5)
+        actual_norm_fro = norm_fro(A)
+        approx_norm_fro = norm_fro_approx(A, tol=1e-2, bsz=bsz)
+        assert_allclose(actual_norm_fro, approx_norm_fro, rtol=1e-1)