Let's make sure this works as expected:
state = qt.rand_dm(12)
dims = [2,3,2]
state.dims = [dims, dims]
assert np.allclose(apply_kraus(state, partial_trace_kraus(0, dims)), state.ptrace(0))
assert np.allclose(apply_kraus(state, partial_trace_kraus(1, dims)), state.ptrace(1))
assert np.allclose(apply_kraus(state, partial_trace_kraus([0,1], dims)), state.ptrace([0,1]))
For example, let's examine the partial trace.
from qbism.povm import *
from qbism.sics import *
from qbism.weyl_heisenberg import *
entangled = qt.rand_dm(4)
entangled.dims = [[2,2],[2,2]]
povm2 = sic_povm(2)
tpovm = tensor_povm(povm2, povm2)
tphi = povm_phi(tpovm)
tp = dm_probs(entangled, tpovm)
ptrA = povm_map(partial_trace_kraus(0, [2,2]), tpovm, povm2)
ptrB = povm_map(partial_trace_kraus(1, [2,2]), tpovm, povm2)
assert np.allclose(dm_probs(entangled.ptrace(0), povm2), ptrA @ tphi @ tp)
assert np.allclose(dm_probs(entangled.ptrace(1), povm2), ptrB @ tphi @ tp)
Above we use a tensor product POVM, consisting of the same POVM on each qubit. This has the nice property:
W = tp.reshape(4,4)
assert np.allclose(np.sum(W, axis=1), dm_probs(entangled.ptrace(0), povm2))
assert np.allclose(np.sum(W, axis=0), dm_probs(entangled.ptrace(1), povm2))
Indeed:
povm2 = sic_povm(2)
tpovm = tensor_povm(povm2, povm2)
A, B = qt.rand_dm(2), qt.rand_dm(2)
AB = qt.tensor(A, B)
assert np.allclose(np.kron(dm_probs(A, povm2), dm_probs(B, povm2)), dm_probs(AB, tpovm))
A more elaborate example:
state = qt.rand_dm(12)
state.dims = [[2,3,2],[2,3,2]]
povm3 = weyl_heisenberg_povm(qt.rand_dm(3))
povm6 = apply_dims(weyl_heisenberg_povm(qt.rand_dm(6)), [2,3])
povm12 = apply_dims(weyl_heisenberg_povm(qt.rand_dm(12)), [2,3,2])
phi = povm_phi(povm12)
p = dm_probs(state, povm12)
ptr0 = povm_map(partial_trace_kraus(0, [2,3,2]), povm12, povm2)
ptr1 = povm_map(partial_trace_kraus(1, [2,3,2]), povm12, povm3)
ptr01 = povm_map(partial_trace_kraus([0,1], [2,3,2]), povm12, povm6)
assert np.allclose(dm_probs(state.ptrace(0), povm2), ptr0 @ phi @ p)
assert np.allclose(dm_probs(state.ptrace(1), povm3), ptr1 @ phi @ p)
assert np.allclose(dm_probs(state.ptrace([0,1]), povm6), ptr01 @ phi @ p)
Compositionality:
d = 2
U = qt.rand_unitary(d)
U2 = qt.tensor(U, U)
povm = sic_povm(d)
tpovm = tensor_povm(povm, povm)
tphi = povm_phi(tpovm)
rho = qt.rand_dm(d**2)
rho.dims = [[d,d],[d,d]]
p = dm_probs(rho, tpovm)
assert np.allclose(dm_probs(U2*rho*U2.dag(), tpovm), povm_map([U2], tpovm, tpovm) @ tphi @ p)
assert np.allclose(dm_probs(U2*rho*U2.dag(), tpovm), np.kron(povm_map([U], povm, povm), povm_map([U], povm, povm)) @ tphi @ p)