def test_3qb(self):
from_np = self.from_numpy
to_np = self.to_numpy
wf = from_np(get_random_wf(np, 3))
rho = make_density_matrix(wf)
o1 = partial_trace_wf_keep_first(wf, 1)
s1 = partial_trace(rho, retain_qubits=[0])
assert np.allclose(to_np(o1), to_np(s1))
o1 = partial_trace_wf_keep_first(wf, 2)
s1 = partial_trace(rho, retain_qubits=[0, 1])
assert np.allclose(to_np(o1), to_np(s1))
def test_work_for_mixed_case_ndim(self):
mep = self.mep
mep_reduced = self.mep_reduced
from_np = self.from_numpy
to_np = self.to_numpy
rho1 = to_np(partial_trace(from_np(mep), (0,)))
rho2 = to_np(partial_trace(from_np(mep), (1,)))
assert np.array_equal(rho1, mep_reduced)
assert np.array_equal(rho2, mep_reduced)
rho_cpx = self.rho_cpx
rho_0 = self.rho_cpx_0
rho_1 = self.rho_cpx_1
assert np.array_equal(rho_0, to_np(partial_trace(from_np(rho_cpx), (0,))))
assert np.array_equal(rho_1, to_np(partial_trace(from_np(rho_cpx), (1,))))
def test_nd_retain_qbs_order_does_not_matter(self):
from_np = self.from_numpy
to_np = self.to_numpy
nqb = 4
dim = 2 ** nqb
qbs_tokeep = np.array([1, 2, 3])
qbs_tokeep_shuffled = qbs_tokeep.copy()
while np.array_equal(qbs_tokeep_shuffled, qbs_tokeep):
np.random.shuffle(qbs_tokeep_shuffled)
rho = np.random.rand(dim, dim) + 1j * np.random.rand(dim, dim)
rho = from_np(rho)
assert np.allclose(
to_np(partial_trace(rho, retain_qubits=qbs_tokeep)),
to_np(partial_trace(rho, retain_qubits=qbs_tokeep_shuffled)),
)
def test_cupy_thread_dim_limit(self):
# It is import to test when the thread dim is larger then 32 for cuda
# code. This happens when there are more than 6 qubits to keep.
# Note: 2 ** 5 = 32, 2 ** 6 = 64.
# See the implementation for cuda for details.
import numpy
to_np = self.to_numpy
wf = self.from_numpy(get_random_wf(numpy, 7))
rho = make_density_matrix(wf)
# Before boundary
o1 = partial_trace_wf_keep_first(wf, 5)
s1 = partial_trace(rho, retain_qubits=list(range(5)))
assert numpy.allclose(to_np(o1), to_np(s1))
# After boundary
o1 = partial_trace_wf_keep_first(wf, 6)
s1 = partial_trace(rho, retain_qubits=list(range(6)))
assert numpy.allclose(to_np(o1), to_np(s1))
def test_trivial_case_ndim(self):
rand = np.random.rand
from_np = self.from_numpy
to_np = self.to_numpy
rho = rand(2, 2) + 1j * rand(2, 2)
assert np.array_equal(
rho, to_np(partial_trace(from_np(rho), retain_qubits=[0]))
)
rho2 = from_np(rand(4, 4) + 1j * rand(4, 4))
assert np.allclose(
to_np(rho2), to_np(partial_trace(rho2, retain_qubits=[1, 0]))
)
# lastly, test if all qubits are killed
assert np.allclose(
np.trace(to_np(rho2)), to_np(partial_trace(rho2, retain_qubits=[]))
)
def test_1d_and_ndim_work_the_same(self):
from_np = self.from_numpy
to_np = self.to_numpy
nqb = 4
dim = 2 ** nqb
qb_to_keep = 2
rho = np.random.rand(dim, dim) + 1j * np.random.rand(dim, dim)
rho = from_np(rho)
assert np.allclose(
to_np(partial_trace_1d(rho, retain_qubit=qb_to_keep)),
to_np(partial_trace(rho, retain_qubits=(qb_to_keep,))),
)
def test_work_for_randommat_ndim(self):
from_np = self.from_numpy
to_np = self.to_numpy
nqb = 4
dim = 2 ** nqb
qbs_tokeep = [1, 2]
rho = np.random.rand(dim, dim) + 1j * np.random.rand(dim, dim)
rho_tmp = from_np(rho)
def get_shouldbe1(arho):
arho = arho.reshape([2] * nqb * 2)
# retains [1, 2]
# [0,1,2,3, 4,5,6,7] -> [0,1,2, 4,5,6] (trace 3, 7)
arho = np.trace(arho, axis1=3, axis2=7)
# [0,1,2, 4,5,6] -> [1,2, 5,6] (trace 0, 3)
arho = np.trace(arho, axis1=0, axis2=3)
return arho.reshape(4, 4)
assert np.allclose(
to_np(partial_trace(rho_tmp, retain_qubits=qbs_tokeep)), get_shouldbe1(rho)
)
qbs_tokeep = [0, 3]
def get_shouldbe2(arho):
arho = arho.reshape([2] * nqb * 2)
# retains [0, 3]
# [0,1,2,3, 4,5,6,7] -> [0,1,3, 4,5,7] (trace 2, 6)
arho = np.trace(arho, axis1=2, axis2=6)
# [0,1,3, 4,5,7] -> [0,3, 4,7] (trace 1, 4)
arho = np.trace(arho, axis1=1, axis2=4)
return arho.reshape(4, 4)
assert np.allclose(
to_np(partial_trace(rho_tmp, retain_qubits=qbs_tokeep)), get_shouldbe2(rho)
)
def test_work_for_pure_states_ndim(self):
from_np = self.from_numpy
to_np = self.to_numpy
plus_minus = self.plus_minus
plus = self.plus
minus = self.minus
plus_to_test = partial_trace(from_np(plus_minus), retain_qubits=(0,))
minus_to_test = partial_trace(from_np(plus_minus), retain_qubits=(1,))
assert np.array_equal(plus, to_np(plus_to_test))
assert np.array_equal(minus, to_np(minus_to_test))
plus_plus_minus = self.ppmm
p_1 = partial_trace(from_np(plus_plus_minus), retain_qubits=(0,))
p_2 = partial_trace(from_np(plus_plus_minus), retain_qubits=(1,))
m_1 = partial_trace(from_np(plus_plus_minus), retain_qubits=(2,))
m_2 = partial_trace(from_np(plus_plus_minus), retain_qubits=(3,))
assert np.array_equal(to_np(p_1), plus)
assert np.array_equal(to_np(p_2), plus)
assert np.array_equal(to_np(m_1), minus)
assert np.array_equal(to_np(m_2), minus)
def get_shouldbe(arho, retain):
arho = from_np(arho)
arho = partial_trace(arho, retain_qubits=retain)
return to_np(arho)
def ptwfs(iwf, retain_qubits):
iwf = make_density_matrix(iwf)
out = [partial_trace(iwf[i], retain_qubits) for i in range(iwf.shape[0])]
out = [to_np(o) for o in out]
return np.asarray(out)
