def test_partial_trace():
a = np.reshape(np.arange(4), (2, 2))
b = np.reshape(np.arange(9) + 4, (3, 3))
c = np.reshape(np.arange(16) + 13, (4, 4))
tr_a = np.trace(a)
tr_b = np.trace(b)
tr_c = np.trace(c)
tensor = np.reshape(np.kron(a, np.kron(b, c)), (2, 3, 4, 2, 3, 4))
np.testing.assert_almost_equal(cirq.partial_trace(tensor, []), tr_a * tr_b * tr_c)
np.testing.assert_almost_equal(cirq.partial_trace(tensor, [0]), a * tr_b * tr_c)
np.testing.assert_almost_equal(cirq.partial_trace(tensor, [1]), b * tr_a * tr_c)
np.testing.assert_almost_equal(cirq.partial_trace(tensor, [2]), c * tr_a * tr_b)
np.testing.assert_almost_equal(
cirq.partial_trace(tensor, [0, 1]), np.reshape(np.kron(a, b), (2, 3, 2, 3)) * tr_c
)
np.testing.assert_almost_equal(
cirq.partial_trace(tensor, [1, 2]), np.reshape(np.kron(b, c), (3, 4, 3, 4)) * tr_a
)
np.testing.assert_almost_equal(
cirq.partial_trace(tensor, [0, 2]), np.reshape(np.kron(a, c), (2, 4, 2, 4)) * tr_b
)
np.testing.assert_almost_equal(cirq.partial_trace(tensor, [0, 1, 2]), tensor)
# permutes indices
np.testing.assert_almost_equal(
cirq.partial_trace(tensor, [1, 0]), np.reshape(np.kron(b, a), (3, 2, 3, 2)) * tr_c
)
np.testing.assert_almost_equal(
cirq.partial_trace(tensor, [2, 0, 1]),
np.reshape(np.kron(c, np.kron(a, b)), (4, 2, 3, 4, 2, 3)),
)
def test_partial_trace_invalid_inputs():
with pytest.raises(ValueError, match='2, 3, 2, 2'):
cirq.partial_trace(np.reshape(np.arange(2 * 3 * 2 * 2), (2, 3, 2, 2)),
[1])
with pytest.raises(ValueError, match='2'):
cirq.partial_trace(np.reshape(np.arange(2 * 2 * 2 * 2), (2, ) * 4),
[2])
def test_same_partial_trace():
qubit_order = cirq.LineQubit.range(2)
q0, q1 = qubit_order
mps_simulator = ccq.mps_simulator.MPSSimulator()
for _ in range(50):
for initial_state in range(4):
circuit = cirq.testing.random_circuit(qubit_order, 3, 0.9)
expected_density_matrix = cirq.final_density_matrix(
circuit, qubit_order=qubit_order, initial_state=initial_state)
expected_partial_trace = cirq.partial_trace(
expected_density_matrix.reshape(2, 2, 2, 2), keep_indices=[0])
final_state = mps_simulator.simulate(
circuit, qubit_order=qubit_order,
initial_state=initial_state).final_state
actual_density_matrix = final_state.partial_trace([q0, q1])
actual_partial_trace = final_state.partial_trace([q0])
np.testing.assert_allclose(actual_density_matrix,
expected_density_matrix,
atol=1e-4)
np.testing.assert_allclose(actual_partial_trace,
expected_partial_trace,
atol=1e-4)
def _expectation_from_density_matrix(self,
density_matrix: np.ndarray) -> float:
observable_matrix = self.matrix()
if density_matrix.shape != observable_matrix.shape:
nqubits = int(np.log2(density_matrix.shape[0]))
density_matrix = cirq.partial_trace(
np.reshape(density_matrix, newshape=[2, 2] * nqubits),
keep_indices=self.qubit_indices,
).reshape(observable_matrix.shape)
return cast(float, np.trace(density_matrix @ observable_matrix))
def test_same_partial_trace():
qubit_order = cirq.LineQubit.range(2)
q0, q1 = qubit_order
angles = [0.0, 0.20160913, math.pi / 3.0, math.pi / 2.0, math.pi]
gate_cls = [cirq.rx, cirq.ry, cirq.rz]
for angle_0 in angles:
for gate_0 in gate_cls:
for angle_1 in angles:
for gate_1 in gate_cls:
for use_cnot in [False, True]:
op0 = gate_0(angle_0)
op1 = gate_1(angle_1)
circuit = cirq.Circuit()
circuit.append(op0(q0))
if use_cnot:
circuit.append(cirq.qft(q0, q1))
circuit.append(op1(q1))
if use_cnot:
circuit.append(cirq.qft(q1, q0))
for initial_state in range(4):
expected_density_matrix = cirq.final_density_matrix(
circuit,
qubit_order=qubit_order,
initial_state=initial_state)
expected_partial_trace = cirq.partial_trace(
expected_density_matrix.reshape(2, 2, 2, 2),
keep_indices=[0])
mps_simulator = ccq.mps_simulator.MPSSimulator()
final_state = mps_simulator.simulate(
circuit,
qubit_order=qubit_order,
initial_state=initial_state).final_state
actual_density_matrix = final_state.partial_trace(
[q0, q1])
actual_partial_trace = final_state.partial_trace(
[q0])
np.testing.assert_allclose(actual_density_matrix,
expected_density_matrix,
atol=1e-4)
np.testing.assert_allclose(actual_partial_trace,
expected_partial_trace,
atol=1e-4)
def test_partial_trace_non_kron():
tensor = np.zeros((2, 2, 2, 2))
tensor[0, 0, 0, 0] = 1
tensor[1, 1, 1, 1] = 4
np.testing.assert_almost_equal(cirq.partial_trace(tensor, [0]), np.array([[1, 0], [0, 4]]))
