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

trajectree/quimb/tests/test_core.py

@@ -0,0 +1,847 @@
+import itertools
+
+from pytest import fixture, raises, mark
+import scipy.sparse as sp
+import numpy as np
+from numpy.testing import assert_allclose, assert_almost_equal
+
+import quimb as qu
+
+stypes = ("csr", "csc", "bsr", "coo")
+
+
+@fixture
+def od1():
+    return qu.rand_matrix(3)
+
+
+@fixture
+def os1():
+    return qu.rand_matrix(3, sparse=True, density=0.5)
+
+
+class TestSparseMatrix:
+    @mark.parametrize("stype", stypes)
+    def test_simple(self, stype):
+        a = qu.core.sparse_matrix([[0, 3], [1, 2]], stype)
+        assert a.format == stype
+        assert a.dtype == complex
+
+
+class TestQuimbify:
+    def test_vector_create(self):
+        x = [1, 2, 3j]
+        p = qu.qu(x, qtype="ket")
+        assert type(p) == qu.qarray
+        assert p.dtype == complex
+        assert p.shape == (3, 1)
+        p = qu.qu(x, qtype="bra")
+        assert p.shape == (1, 3)
+        assert_almost_equal(p[0, 2], -3.0j)
+
+    def test_dop_create(self):
+        x = np.random.randn(3, 3)
+        p = qu.qu(x, qtype="dop")
+        assert type(p) == qu.qarray
+        assert p.dtype == complex
+        assert p.shape == (3, 3)
+
+    def test_convert_vector_to_dop(self):
+        x = [1, 2, 3j]
+        p = qu.qu(x, qtype="r")
+        assert_allclose(
+            p,
+            qu.qarray(
+                [
+                    [1.0 + 0.0j, 2.0 + 0.0j, 0.0 - 3.0j],
+                    [2.0 + 0.0j, 4.0 + 0.0j, 0.0 - 6.0j],
+                    [0.0 + 3.0j, 0.0 + 6.0j, 9.0 + 0.0j],
+                ]
+            ),
+        )
+
+    def test_chopped(self):
+        x = [9e-16, 1]
+        p = qu.qu(x, "k", chopped=False)
+        assert p[0, 0] != 0.0
+        p = qu.qu(x, "k", chopped=True)
+        assert p[0, 0] == 0.0
+
+    def test_normalized(self):
+        x = [3j, 4j]
+        p = qu.qu(x, "k", normalized=False)
+        assert_almost_equal(qu.tr(p.H @ p), 25.0)
+        p = qu.qu(x, "k", normalized=True)
+        assert_almost_equal(qu.tr(p.H @ p), 1.0)
+        p = qu.qu(x, "dop", normalized=True)
+        assert_almost_equal(qu.tr(p), 1.0)
+
+    def test_sparse_create(self):
+        x = [[1, 0], [3, 0]]
+        p = qu.qu(x, "dop", sparse=False)
+        assert type(p) == qu.qarray
+        p = qu.qu(x, "dop", sparse=True)
+        assert type(p) == sp.csr_matrix
+        assert p.dtype == complex
+        assert p.nnz == 2
+
+    def test_sparse_convert_to_dop(self):
+        x = [1, 0, 9e-16, 0, 3j]
+        p = qu.qu(x, "ket", sparse=True)
+        q = qu.qu(p, "dop", sparse=True)
+        assert q.shape == (5, 5)
+        assert q.nnz == 9
+        assert_almost_equal(q[4, 4], 9.0)
+        q = qu.qu(p, "dop", sparse=True, normalized=True)
+        assert_almost_equal(qu.tr(q), 1.0)
+
+    @mark.parametrize(
+        "qtype, shape, out",
+        (
+            ("bra", (1, 4), [[1, 0, 2, -3j]]),
+            ("ket", (4, 1), [[1], [0], [2], [3j]]),
+            (
+                "dop",
+                (4, 4),
+                [[1, 0, 2, -3j], [0, 0, 0, 0], [2, 0, 4, -6j], [3j, 0, 6j, 9]],
+            ),
+        ),
+    )
+    @mark.parametrize("format_in", stypes)
+    @mark.parametrize("format_out", (None,) + stypes)
+    @mark.parametrize("dtype", [float, complex, np.float64, np.complex128])
+    def test_reshape_sparse(
+        self, qtype, shape, out, format_in, format_out, dtype
+    ):
+        import warnings
+
+        in_ = [[1], [0], [2], [3j]]
+        x = qu.core.sparse_matrix(in_, stype=format_in)
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            y = qu.qu(x, qtype=qtype, stype=format_out, dtype=dtype)
+
+        assert y.shape == shape
+        assert y.dtype == dtype
+        if format_out is None:
+            format_out = format_in
+        assert y.format == format_out
+
+        if np.issubdtype(dtype, np.floating):
+            assert_allclose(y.toarray(), np.real(out), atol=1e-12)
+        else:
+            assert_allclose(y.toarray(), out)
+
+    @mark.parametrize(
+        "qtype, shape, out",
+        (
+            ("bra", (1, 4), [[1, 0, 2, -3j]]),
+            ("ket", (4, 1), [[1], [0], [2], [3j]]),
+            (
+                "dop",
+                (4, 4),
+                [[1, 0, 2, -3j], [0, 0, 0, 0], [2, 0, 4, -6j], [3j, 0, 6j, 9]],
+            ),
+        ),
+    )
+    @mark.parametrize("format_out", (None,) + stypes)
+    def test_dense_to_sparse_format(self, qtype, shape, out, format_out):
+        x = [[1], [0], [2], [3j]]
+        y = qu.qu(x, qtype=qtype, stype=format_out, sparse=True)
+        assert y.shape == shape
+        assert y.dtype == complex
+        if format_out is None:
+            format_out = "csr"
+        assert y.format == format_out
+        assert_allclose(y.toarray(), out)
+
+    @mark.parametrize(
+        "qtype, shape", (["bra", (1, 4)], ["ket", (4, 1)], ["dop", (4, 4)])
+    )
+    @mark.parametrize("format_out", stypes)
+    def test_give_sformat_only(self, qtype, shape, format_out):
+        x = [[1], [0], [2], [3j]]
+        y = qu.qu(x, qtype=qtype, stype=format_out)
+        assert qu.issparse(y)
+        assert y.shape == shape
+        assert y.format == format_out
+
+
+class TestInferSize:
+    @mark.parametrize(
+        "d,base,n", ([8, 2, 3], [16, 2, 4], [9, 3, 2], [81, 3, 4])
+    )
+    def test_infer_size(self, d, base, n):
+        p = qu.rand_ket(d)
+        assert qu.infer_size(p, base) == n
+
+    def test_raises(self):
+        p = qu.rand_ket(2) & qu.rand_ket(3)
+        with raises(ValueError):
+            qu.infer_size(p, base=2)
+
+
+class TestTrace:
+    @mark.parametrize(
+        "inpt, outpt", ([[[2, 1], [4, 5]], 7], [[[2, 1], [4, 5j]], 2 + 5j])
+    )
+    @mark.parametrize(
+        "sparse, func",
+        ([False, qu.core._trace_dense], [True, qu.core._trace_sparse]),
+    )
+    def test_simple(self, inpt, outpt, sparse, func):
+        a = qu.qu(inpt, sparse=sparse)
+        assert qu.trace(a) == outpt
+
+
+class TestITrace:
+    @mark.parametrize("axes", [(0, 1), ((0,), (1,))])
+    def test_axes_types(self, axes):
+        a = qu.rand_matrix(4)
+        b = qu.itrace(a, axes)
+        assert_allclose(b, np.trace(a))
+
+    def test_complex_dims(self):
+        a = np.random.rand(4, 3, 2, 2, 4, 3)
+        atr = qu.itrace(a, ((0, 1, 2), (4, 5, 3)))
+        btr = np.trace(
+            np.trace(np.trace(a, axis1=1, axis2=5), axis1=1, axis2=2)
+        )
+        assert_allclose(atr, btr)
+
+
+class TestNormalize:
+    def test_normalize_ket(self):
+        a = qu.qu([1, -1j], "ket")
+        b = qu.nmlz(a, inplace=False)
+        assert_almost_equal(qu.trace(b.H @ b), 1.0)
+        assert_almost_equal(qu.trace(a.H @ a), 2.0)
+
+    def test_normalize_bra(self):
+        a = qu.qu([1, -1j], "bra")
+        b = qu.nmlz(a, inplace=False)
+        assert_almost_equal(qu.trace(b @ b.H), 1.0)
+
+    def test_normalize_dop(self):
+        a = qu.qu([1, -1j], "dop")
+        b = qu.nmlz(a, inplace=False)
+        assert_almost_equal(qu.trace(b), 1.0)
+
+    def test_normalize_inplace_ket(self):
+        a = qu.qu([1, -1j], "ket")
+        a.nmlz(inplace=True)
+        assert_almost_equal(qu.trace(a.H @ a), 1.0)
+
+    def test_normalize_inplace_bra(self):
+        a = qu.qu([1, -1j], "bra")
+        a.nmlz(inplace=True)
+        assert_almost_equal(qu.trace(a @ a.H), 1.0)
+
+    def test_normalize_inplace_dop(self):
+        a = qu.qu([1, -1j], "dop")
+        b = qu.nmlz(a, inplace=True)
+        assert_almost_equal(qu.trace(a), 1.0)
+        assert_almost_equal(qu.trace(b), 1.0)
+
+
+class TestDimMap:
+    @mark.parametrize("numpy", [False, True])
+    def test_1d(self, numpy):
+        dims = [10, 11, 12, 13]
+        coos = (1, 2, 3)
+        if numpy:
+            dims, coos = np.asarray(dims), np.asarray(coos)
+        ndims, ncoos = qu.dim_map(dims, coos)
+        assert_allclose([ndims[c] for c in ncoos], (11, 12, 13))
+        coos = ([-1], [2], [5])
+        with raises(ValueError):
+            ndims, ncoos = qu.dim_map(dims, coos)
+        ndims, ncoos = qu.dim_map(dims, coos, cyclic=True)
+        assert_allclose([ndims[c] for c in ncoos], (13, 12, 11))
+        ndims, ncoos = qu.dim_map(dims, coos, trim=True)
+        assert_allclose([ndims[c] for c in ncoos], [12])
+
+    def test_2d(self):
+        dims = [[200, 201, 202, 203], [210, 211, 212, 213]]
+        coos = ((1, 2), (1, 3), (0, 3))
+        ndims, ncoos = qu.dim_map(dims, coos)
+        assert_allclose([ndims[c] for c in ncoos], (212, 213, 203))
+        coos = ((-1, 1), (1, 2), (3, 4))
+        with raises(ValueError):
+            ndims, ncoos = qu.dim_map(dims, coos)
+        ndims, ncoos = qu.dim_map(dims, coos, cyclic=True)
+        assert_allclose([ndims[c] for c in ncoos], (211, 212, 210))
+        ndims, ncoos = qu.dim_map(dims, coos, trim=True)
+        assert_allclose([ndims[c] for c in ncoos], [212])
+
+    def test_3d(self):
+        dims = [
+            [[3000, 3001, 3002], [3010, 3011, 3012], [3020, 3021, 3022]],
+            [[3100, 3101, 3102], [3110, 3111, 3112], [3120, 3121, 3122]],
+        ]
+        coos = ((0, 0, 2), (1, 1, 2), (1, 2, 0))
+        ndims, ncoos = qu.dim_map(dims, coos)
+        assert_allclose([ndims[c] for c in ncoos], (3002, 3112, 3120))
+        coos = ((0, -1, 2), (1, 2, 2), (4, -1, 3))
+        with raises(ValueError):
+            ndims, ncoos = qu.dim_map(dims, coos)
+        ndims, ncoos = qu.dim_map(dims, coos, cyclic=True)
+        assert_allclose([ndims[c] for c in ncoos], (3022, 3122, 3020))
+        ndims, ncoos = qu.dim_map(dims, coos, trim=True)
+        assert_allclose([ndims[c] for c in ncoos], [3122])
+
+
+class TestDimCompress:
+    def test_edge(self):
+        dims = [2, 3, 2, 4, 5]
+        coos = [0, 4]
+        ndims, ncoos = qu.dim_compress(dims, coos)
+        assert ndims == (2, 24, 5)
+        assert ncoos == (0, 2)
+
+    def test_middle(self):
+        dims = [5, 3, 2, 5, 4, 3, 2]
+        coos = [1, 2, 3, 5]
+        ndims, ncoos = qu.dim_compress(dims, coos)
+        assert ndims == (5, 30, 4, 3, 2)
+        assert ncoos == (1, 3)
+
+    def test_single(self):
+        dims = [5, 3, 2, 5, 4, 3, 2]
+        coos = 3
+        ndims, ncoos = qu.dim_compress(dims, coos)
+        assert ndims == (30, 5, 24)
+        assert ncoos == (1,)
+
+    @mark.parametrize(
+        "dims, inds, ndims, ninds",
+        [([2, 2], [0, 1], (4,), (0,)), ([4], [0], (4,), (0,))],
+    )
+    def test_tiny(self, dims, inds, ndims, ninds):
+        dims, inds = qu.dim_compress(dims, inds)
+        assert dims == ndims
+        assert inds == ninds
+
+
+class TestEye:
+    def test_eye_dense(self):
+        a = qu.eye(3, sparse=False)
+        assert a.shape == (3, 3)
+        assert isinstance(a, qu.qarray)
+        assert a.dtype == complex
+
+    def test_eye_sparse(self):
+        a = qu.eye(3, sparse=True)
+        assert a.shape == (3, 3)
+        assert isinstance(a, sp.csr_matrix)
+        assert a.dtype == complex
+
+
+class TestKron:
+    @mark.parametrize("parallel", [True, False])
+    def test_kron_basic(self, parallel):
+        a = qu.rand_ket(2)
+        b = qu.rand_ket(4)
+        c = qu.rand_ket(4)
+        d = qu.rand_ket(5)
+        t = qu.kron(a, b, c, d, parallel=parallel)
+        assert_allclose(t, a & b & c & d)
+
+    @mark.parametrize("sparse", [True, False])
+    @mark.parametrize("ri,rf", ([0, 4], [75, 89], [150, 168], [0, 168]))
+    def test_kron_ownership(self, sparse, ri, rf):
+        dims = [7, 2, 4, 3]
+        ops = [qu.rand_matrix(d, sparse=sparse) for d in dims]
+        X1 = qu.kron(*ops)[ri:rf, :]
+        X2 = qu.kron(*ops, ownership=(ri, rf))
+        assert_allclose(X1.toarray(), X2.toarray())
+
+
+class Testikron:
+    def test_basic(self):
+        a = qu.rand_matrix(2)
+        i = qu.eye(2)
+        dims = [2, 2, 2]
+        b = qu.ikron([a], dims, [0])
+        assert_allclose(b, a & i & i)
+        b = qu.ikron([a], dims, [1])
+        assert_allclose(b, i & a & i)
+        b = qu.ikron([a], dims, [2])
+        assert_allclose(b, i & i & a)
+        b = qu.ikron([a], dims, [0, 2])
+        assert_allclose(b, a & i & a)
+        b = qu.ikron([a], dims, [0, 1, 2])
+        assert_allclose(b, a & a & a)
+
+    def test_mid_multi(self):
+        a = [qu.rand_matrix(2) for i in range(3)]
+        i = qu.eye(2)
+        dims = [2, 2, 2, 2, 2, 2]
+        inds = [1, 2, 4]
+        b = qu.ikron(a, dims, inds)
+        assert_allclose(b, i & a[0] & a[1] & i & a[2] & i)
+
+    def test_mid_multi_reverse(self):
+        a = [qu.rand_matrix(2) for i in range(3)]
+        i = qu.eye(2)
+        dims = [2, 2, 2, 2, 2, 2]
+        inds = [5, 4, 1]
+        b = qu.ikron(a, dims, inds)
+        assert_allclose(b, i & a[2] & i & i & a[1] & a[0])
+
+    def test_auto(self):
+        a = qu.rand_matrix(2)
+        i = qu.eye(2)
+        b = qu.ikron([a], (2, -1, 2), [1])
+        assert_allclose(b, i & a & i)
+
+    def test_ndarrays(self):
+        a = qu.rand_matrix(2)
+        i = qu.eye(2)
+        b = qu.ikron([a], np.array([2, 2, 2]), [0, 2])
+        assert_allclose(b, a & i & a)
+        b = qu.ikron([a], [2, 2, 2], np.array([0, 2]))
+        assert_allclose(b, a & i & a)
+
+    def test_overlap(self):
+        a = [qu.rand_matrix(4) for i in range(2)]
+        dims1 = [2, 2, 2, 2, 2, 2]
+        dims2 = [2, 4, 4, 2]
+        b = qu.ikron(a, dims1, [1, 2, 3, 4])
+        c = qu.ikron(a, dims2, [1, 2])
+        assert_allclose(c, b)
+        dims2 = [4, 2, 2, 4]
+        b = qu.ikron(a, dims1, [0, 1, 4, 5])
+        c = qu.ikron(a, dims2, [0, 3])
+        assert_allclose(c, b)
+
+    def test_holey_overlap(self):
+        a = qu.rand_matrix(8)
+        dims1 = (2, 2, 2, 2, 2)
+        dims2 = (2, 8, 2)
+        b = qu.ikron(a, dims1, (1, 3))
+        c = qu.ikron(a, dims2, 1)
+        assert_allclose(b, c)
+        dims1 = (2, 2, 2, 2, 2)
+        dims2 = (2, 2, 8)
+        b = qu.ikron(a, dims1, (2, 4))
+        c = qu.ikron(a, dims2, 2)
+        assert_allclose(b, c)
+        dims1 = (2, 2, 2, 2, 2)
+        dims2 = (8, 2, 2)
+        b = qu.ikron(a, dims1, (0, 2))
+        c = qu.ikron(a, dims2, 0)
+        assert_allclose(b, c)
+
+    def test_sparse(self):
+        i = qu.eye(2, sparse=True)
+        a = qu.qu(qu.rand_matrix(2), sparse=True)
+        b = qu.ikron(a, [2, 2, 2], 1)  # infer sparse
+        assert qu.issparse(b)
+        assert_allclose(b.toarray(), (i & a & i).toarray())
+        a = qu.rand_matrix(2)
+        b = qu.ikron(a, [2, 2, 2], 1, sparse=True)  # explicit sparse
+        assert qu.issparse(b)
+        assert_allclose(b.toarray(), (i & a & i).toarray())
+
+    def test_2d_simple(self):
+        a = (qu.rand_matrix(2), qu.rand_matrix(2))
+        dims = ((2, 3), (3, 2))
+        inds = ((0, 0), (1, 1))
+        b = qu.ikron(a, dims, inds)
+        assert b.shape == (36, 36)
+        assert_allclose(b, a[0] & qu.eye(9) & a[1])
+
+    @mark.parametrize("stype", (None,) + stypes)
+    @mark.parametrize(
+        "pos", [0, 1, 2, (0,), (1,), (2,), (0, 1), (1, 2), (0, 2)]
+    )
+    @mark.parametrize("coo_build", [False, True])
+    def test_sparse_format_outputs(self, os1, stype, pos, coo_build):
+        x = qu.ikron(os1, [3, 3, 3], pos, stype=stype, coo_build=coo_build)
+        assert x.format == "csr" if stype is None else stype
+
+    @mark.parametrize("stype", (None,) + stypes)
+    @mark.parametrize(
+        "pos", [0, 1, 2, (0,), (1,), (2,), (0, 1), (1, 2), (0, 2)]
+    )
+    @mark.parametrize("coo_build", [False, True])
+    def test_sparse_format_outputs_with_dense(
+        self, od1, stype, pos, coo_build
+    ):
+        x = qu.ikron(
+            od1, [3, 3, 3], pos, sparse=True, stype=stype, coo_build=coo_build
+        )
+        try:
+            default = "bsr" if (2 in pos and not coo_build) else "csr"
+        except TypeError:
+            default = "bsr" if (pos == 2 and not coo_build) else "csr"
+        assert x.format == default if stype is None else stype
+
+    @mark.parametrize("sparse", [True, False])
+    @mark.parametrize("ri,rf", ([0, 4], [75, 89], [150, 168], [0, 168]))
+    def test_ikron_ownership(self, sparse, ri, rf):
+        dims = [7, 2, 4, 3]
+        X = qu.rand_matrix(2, sparse=sparse)
+        Y = qu.rand_matrix(3, sparse=sparse)
+        X1 = qu.ikron((X, Y), dims, (1, 3))[ri:rf, :]
+        X2 = qu.ikron((X, Y), dims, (1, 3), ownership=(ri, rf))
+        assert_allclose(X1.toarray(), X2.toarray())
+
+
+class TestPermikron:
+    def test_dop_spread(self):
+        a = qu.rand_rho(4)
+        b = qu.pkron(a, [2, 2, 2], [0, 2])
+        c = (
+            (a & qu.eye(2))
+            .toarray().reshape([2, 2, 2, 2, 2, 2])
+            .transpose([0, 2, 1, 3, 5, 4])
+            .reshape([8, 8])
+        )
+        assert_allclose(b, c)
+
+    def test_dop_reverse(self):
+        a = qu.rand_rho(4)
+        b = qu.pkron(a, np.array([2, 2, 2]), [2, 0])
+        c = (
+            (a & qu.eye(2))
+            .toarray().reshape([2, 2, 2, 2, 2, 2])
+            .transpose([1, 2, 0, 4, 5, 3])
+            .reshape([8, 8])
+        )
+        assert_allclose(b, c)
+
+    def test_dop_reverse_sparse(self):
+        a = qu.rand_rho(4, sparse=True, density=0.5)
+        b = qu.pkron(a, np.array([2, 2, 2]), [2, 0])
+        c = (
+            (a & qu.eye(2))
+            .toarray().reshape([2, 2, 2, 2, 2, 2])
+            .transpose([1, 2, 0, 4, 5, 3])
+            .reshape([8, 8])
+        )
+        assert_allclose(b.toarray(), c)
+
+
+class TestPermute:
+    def test_permute_ket(self):
+        a = qu.up() & qu.plus() & qu.yplus()
+        b = qu.permute(a, [2, 2, 2], [2, 0, 1])
+        assert_allclose(b, qu.yplus() & qu.up() & qu.plus())
+
+    def test_permute_op(self):
+        a = qu.pauli("x") & qu.pauli("y") & qu.pauli("z")
+        b = qu.permute(a, [2, 2, 2], [2, 0, 1])
+        assert_allclose(b, qu.pauli("z") & qu.pauli("x") & qu.pauli("y"))
+
+    def test_entangled_permute(self):
+        dims = [2, 2, 2]
+        a = qu.bell_state(0) & qu.up()
+        assert_allclose(qu.mutinf_subsys(a, dims, 0, 1), 2.0)
+        b = qu.permute(a, dims, [1, 2, 0])
+        assert_allclose(qu.mutinf_subsys(b, dims, 0, 1), 0.0, atol=1e-12)
+        assert_allclose(qu.mutinf_subsys(b, dims, 0, 2), 2.0)
+
+    def test_permute_sparse_ket(self):
+        dims = [3, 2, 5, 4]
+        a = qu.rand_ket(qu.prod(dims), sparse=True, density=0.5)
+        b = qu.permute(a, dims, [3, 1, 2, 0])
+        c = qu.permute(a.toarray(), dims, [3, 1, 2, 0])
+        assert_allclose(b.toarray(), c)
+
+    def test_permute_sparse_op(self):
+        dims = [3, 2, 5, 4]
+        a = qu.rand_rho(qu.prod(dims), sparse=True, density=0.5)
+        b = qu.permute(a, dims, [3, 1, 2, 0])
+        c = qu.permute(a.toarray(), dims, [3, 1, 2, 0])
+        assert_allclose(b.toarray(), c)
+
+
+class TestPartialTraceDense:
+    def test_partial_trace_basic(self):
+        a = qu.rand_rho(2**2)
+        b = qu.partial_trace(a, [2, 2], 0)
+        assert isinstance(b, qu.qarray)
+        assert qu.isherm(b)
+        assert_allclose(qu.tr(b), 1.0)
+
+    def test_ptr_compare_to_manual(self):
+        a = qu.rand_rho(2**2)
+        b = qu.partial_trace(a, [2, 2], 0)
+        c = a.toarray().reshape([2, 2, 2, 2]).trace(axis1=1, axis2=3)
+        assert_allclose(b, c)
+        b = qu.partial_trace(a, [2, 2], 1)
+        c = a.toarray().reshape([2, 2, 2, 2]).trace(axis1=0, axis2=2)
+        assert_allclose(b, c)
+
+    def test_partial_trace_early_return(self):
+        a = qu.qu([0.5, 0.5, 0.5, 0.5], "ket")
+        b = qu.partial_trace(a, [2, 2], [0, 1])
+        assert_allclose(a @ a.H, b)
+        a = qu.qu([0.5, 0.5, 0.5, 0.5], "dop")
+        b = qu.partial_trace(a, [2, 2], [0, 1])
+        assert_allclose(a, b)
+
+    def test_partial_trace_return_type(self):
+        a = qu.qu([0, 2**-0.5, 2**-0.5, 0], "ket")
+        b = qu.partial_trace(a, [2, 2], 1)
+        assert type(b) == qu.qarray
+        a = qu.qu([0, 2**-0.5, 2**-0.5, 0], "dop")
+        b = qu.partial_trace(a, [2, 2], 1)
+        assert type(b) == qu.qarray
+
+    def test_partial_trace_single_ket(self):
+        dims = [2, 3, 4]
+        a = np.random.randn(qu.prod(dims), 1)
+        for i, dim in enumerate(dims):
+            b = qu.partial_trace(a, dims, i)
+            assert b.shape[0] == dim
+
+    def test_partial_trace_multi_ket(self):
+        dims = [2, 3, 4]
+        a = np.random.randn(qu.prod(dims), 1)
+        for i1, i2 in itertools.combinations([0, 1, 2], 2):
+            b = qu.partial_trace(a, dims, [i1, i2])
+            assert b.shape[1] == dims[i1] * dims[i2]
+
+    def test_partial_trace_dop_product_state(self):
+        dims = [3, 2, 4, 2, 3]
+        ps = [qu.rand_rho(dim) for dim in dims]
+        pt = qu.kron(*ps)
+        for i, dim in enumerate(dims):
+            p = qu.partial_trace(pt, dims, i)
+            assert_allclose(p, ps[i])
+
+    def test_partial_trace_bell_states(self):
+        for lab in ("psi-", "psi+", "phi-", "phi+"):
+            psi = qu.bell_state(lab, qtype="dop")
+            rhoa = qu.partial_trace(psi, [2, 2], 0)
+            assert_allclose(rhoa, qu.eye(2) / 2)
+
+    def test_partial_trace_supply_ndarray(self):
+        a = qu.rand_rho(2**3)
+        dims = np.array([2, 2, 2])
+        keep = np.array(1)
+        b = qu.partial_trace(a, dims, keep)
+        assert b.shape[0] == 2
+
+    def test_partial_trace_order_doesnt_matter(self):
+        a = qu.rand_rho(2**3)
+        dims = np.array([2, 2, 2])
+        b1 = qu.partial_trace(a, dims, [0, 2])
+        b2 = qu.partial_trace(a, dims, [2, 0])
+        assert_allclose(b1, b2)
+
+
+class TestTraceLose:
+    def test_rps(self):
+        a, b, c = (
+            qu.rand_rho(2, sparse=True, density=0.5),
+            qu.rand_rho(3, sparse=True, density=0.5),
+            qu.rand_rho(2, sparse=True, density=0.5),
+        )
+        abc = a & b & c
+        pab = qu.core._trace_lose(abc, [2, 3, 2], 2)
+        assert_allclose(pab, (a & b).toarray())
+        pac = qu.core._trace_lose(abc, [2, 3, 2], 1)
+        assert_allclose(pac, (a & c).toarray())
+        pbc = qu.core._trace_lose(abc, [2, 3, 2], 0)
+        assert_allclose(pbc, (b & c).toarray())
+
+    def test_bell_state(self):
+        a = qu.bell_state("psi-", sparse=True)
+        b = qu.core._trace_lose(a @ a.H, [2, 2], 0)
+        assert_allclose(b, qu.eye(2) / 2)
+        b = qu.core._trace_lose(a @ a.H, [2, 2], 1)
+        assert_allclose(b, qu.eye(2) / 2)
+
+    def test_vs_ptr(self):
+        a = qu.rand_rho(6, sparse=True, density=0.5)
+        b = qu.core._trace_lose(a, [2, 3], 1)
+        c = qu.partial_trace(a.toarray(), [2, 3], 0)
+        assert_allclose(b, c)
+        b = qu.core._trace_lose(a, [2, 3], 0)
+        c = qu.partial_trace(a.toarray(), [2, 3], 1)
+        assert_allclose(b, c)
+
+    def test_vec_dense(self):
+        a = qu.rand_ket(4)
+        b = qu.core._trace_lose(a, [2, 2], 1)
+        c = qu.partial_trace(a.toarray(), [2, 2], 0)
+        assert_allclose(b, c)
+        b = qu.core._trace_lose(a, [2, 2], 0)
+        c = qu.partial_trace(a.toarray(), [2, 2], 1)
+        assert_allclose(b, c)
+
+
+class TestTraceKeep:
+    def test_rps(self):
+        a, b, c = (
+            qu.rand_rho(2, sparse=True, density=0.5),
+            qu.rand_rho(3, sparse=True, density=0.5),
+            qu.rand_rho(2, sparse=True, density=0.5),
+        )
+        abc = a & b & c
+        pc = qu.core._trace_keep(abc, [2, 3, 2], 2)
+        assert_allclose(pc, c.toarray())
+        pb = qu.core._trace_keep(abc, [2, 3, 2], 1)
+        assert_allclose(pb, b.toarray())
+        pa = qu.core._trace_keep(abc, [2, 3, 2], 0)
+        assert_allclose(pa, a.toarray())
+
+    def test_bell_state(self):
+        a = qu.bell_state("psi-", sparse=True)
+        b = qu.core._trace_keep(a @ a.H, [2, 2], 0)
+        assert_allclose(b, qu.eye(2) / 2)
+        b = qu.core._trace_keep(a @ a.H, [2, 2], 1)
+        assert_allclose(b, qu.eye(2) / 2)
+
+    def test_vs_ptr(self):
+        a = qu.rand_rho(6, sparse=True, density=0.5)
+        b = qu.core._trace_keep(a, [2, 3], 0)
+        c = qu.partial_trace(a.toarray(), [2, 3], 0)
+        assert_allclose(b, c)
+        b = qu.core._trace_keep(a, [2, 3], 1)
+        c = qu.partial_trace(a.toarray(), [2, 3], 1)
+        assert_allclose(b, c)
+
+    def test_vec_dense(self):
+        a = qu.rand_ket(4)
+        b = qu.core._trace_keep(a, [2, 2], 0)
+        c = qu.partial_trace(a.toarray(), [2, 2], 0)
+        assert_allclose(b, c)
+        b = qu.core._trace_keep(a, [2, 2], 1)
+        c = qu.partial_trace(a.toarray(), [2, 2], 1)
+        assert_allclose(b, c)
+
+
+class TestPartialTraceSparse:
+    def test_partial_trace_sparse_basic(self):
+        a = qu.rand_rho(4)
+        b = qu.partial_trace(a, [2, 2], 0)
+        assert type(b) == qu.qarray
+        assert qu.isherm(b)
+        assert_allclose(qu.tr(b), 1.0)
+
+    def test_partial_trace_simple_single(self):
+        a = qu.rand_rho(12, sparse=True, density=0.5)
+        dims = [2, 3, 2]
+        b = qu.partial_trace(a, dims, 1)
+        c = (
+            a.toarray().reshape([*dims, *dims])
+            .trace(axis1=2, axis2=5)
+            .trace(axis1=0, axis2=2)
+        )
+        assert_allclose(c, b)
+
+    def test_partial_trace_simple_double(self):
+        a = qu.rand_rho(12, sparse=True, density=0.5)
+        dims = [2, 3, 2]
+        b = qu.partial_trace(a, dims, [0, 2])
+        c = qu.partial_trace(a.toarray(), dims, [0, 2])
+        assert_allclose(b, c)
+        b = qu.partial_trace(a, dims, [1, 2])
+        c = qu.partial_trace(a.toarray(), dims, [1, 2])
+        assert_allclose(b, c)
+
+    def test_partial_trace_simple_ket(self):
+        a = qu.rand_ket(12, sparse=True, density=0.5)
+        dims = [2, 3, 2]
+        b = qu.partial_trace(a, dims, [0, 1])
+        c = qu.partial_trace(a.toarray(), dims, [0, 1])
+        assert_allclose(b, c)
+
+
+class TestChop:
+    def test_chop_inplace(self):
+        a = qu.qu([-1j, 0.1 + 0.2j])
+        qu.chop(a, tol=0.11, inplace=True)
+        assert_allclose(a, qu.qu([-1j, 0.2j]))
+        # Sparse
+        a = qu.qu([-1j, 0.1 + 0.2j], sparse=True)
+        qu.chop(a, tol=0.11, inplace=True)
+        b = qu.qu([-1j, 0.2j], sparse=True)
+        assert (a != b).nnz == 0
+
+    def test_chop_inplace_dop(self):
+        a = qu.qu([1, 0.1], "dop")
+        qu.chop(a, tol=0.11, inplace=True)
+        assert_allclose(a, qu.qu([1, 0], "dop"))
+        a = qu.qu([1, 0.1], "dop", sparse=True)
+        qu.chop(a, tol=0.11, inplace=True)
+        b = qu.qu([1, 0.0], "dop", sparse=True)
+        assert (a != b).nnz == 0
+
+    def test_chop_copy(self):
+        a = qu.qu([-1j, 0.1 + 0.2j])
+        b = qu.chop(a, tol=0.11, inplace=False)
+        assert_allclose(a, qu.qu([-1j, 0.1 + 0.2j]))
+        assert_allclose(b, qu.qu([-1j, 0.2j]))
+        # Sparse
+        a = qu.qu([-1j, 0.1 + 0.2j], sparse=True)
+        b = qu.chop(a, tol=0.11, inplace=False)
+        ao = qu.qu([-1j, 0.1 + 0.2j], sparse=True)
+        bo = qu.qu([-1j, 0.2j], sparse=True)
+        assert (a != ao).nnz == 0
+        assert (b != bo).nnz == 0
+
+
+class TestExpec:
+    @mark.parametrize("qtype1", ["ket", "dop"])
+    @mark.parametrize("spars1", [True, False])
+    @mark.parametrize("qtype2", ["ket", "dop"])
+    @mark.parametrize("spars2", [True, False])
+    def test_all(self, qtype1, spars1, qtype2, spars2):
+        a = qu.qu([[1], [2j], [3]], qtype=qtype1, sparse=spars1)
+        b = qu.qu([[1j], [2], [3j]], qtype=qtype2, sparse=spars2)
+        c = qu.expec(a, b)
+        assert not isinstance(c, complex)
+        assert_allclose(c, 36)
+
+    @mark.parametrize("qtype", ["ket", "dop"])
+    @mark.parametrize("sparse", [True, False])
+    @mark.parametrize("s", ["x", "y", "z"])
+    def test_negative_expec(self, qtype, sparse, s):
+        a = qu.singlet(qtype=qtype)
+        b = qu.pauli(s, sparse=sparse) & qu.pauli(s, sparse=sparse)
+        assert_allclose(qu.expec(a, b), -1)
+
+
+class TestNumbaFuncs:
+    @mark.parametrize("size", [300, 3000, (300, 5), (3000, 5)])
+    @mark.parametrize(
+        "X_dtype", ["float32", "float64", "complex64", "complex128"]
+    )
+    @mark.parametrize("c_dtype", ["float32", "float64"])
+    def test_subtract_update(
+        self,
+        size,
+        X_dtype,
+        c_dtype,
+    ):
+        X = qu.randn(size, dtype=X_dtype)
+        Y = qu.randn(size, dtype=X_dtype)
+        c = qu.randn(1, dtype=c_dtype)[0]
+        res = X - c * Y
+        qu.core.subtract_update_(X, c, Y)
+        assert_allclose(res, X, rtol=1e-6)
+
+    @mark.parametrize("size", [300, 3000, (300, 5), (3000, 5)])
+    @mark.parametrize(
+        "X_dtype", ["float32", "float64", "complex64", "complex128"]
+    )
+    @mark.parametrize("c_dtype", ["float32", "float64"])
+    def test_divide_update(
+        self,
+        size,
+        X_dtype,
+        c_dtype,
+    ):
+        X = qu.randn(size, dtype=X_dtype)
+        Y = np.empty_like(X)
+        c = qu.randn(1, dtype=c_dtype).item()
+        res = X / c
+        qu.core.divide_update_(X, c, Y)
+        assert_allclose(res, Y, rtol=1e-6)