Add density matrix and entanglement calculations for qubits

unitary/alpha/quantum_world.py

@@ -636,8 +636,91 @@ def get_binary_probabilities(
             binary_probs.append(1 - one_probs[0])
         return binary_probs
 
+    def density_matrix(
+        self, objects: Optional[Sequence[QuantumObject]] = None, count: int = 1000
+    ) -> np.ndarray:
+        """Simulates the density matrix of the given objects. We assume that the overall state of
+        the quantum world (including all quantum objects in it) could be described by one pure
+        state. To calculate the density matrix of the given quantum objects, we would always measure/
+        peek the quantum world for `count` times, deduce the (pure) state vector based on the results,
+        then the density matrix is its outer product. We will then trace out the un-needed quantum
+        objects before returning the density matrix.
+
+        Parameters:
+            objects:    List of QuantumObjects (currently only qubits are supported). If not specified,
+                        all quantum objects' density matrix will be returned.
+            count:      Number of measurements.
+
+        Returns:
+            The density matrix of the specified objects.
+        """
+        num_all_qubits = len(self.object_name_dict.values())
+        num_shown_qubits = len(objects) if objects is not None else num_all_qubits
+
+        specified_names = (
+            [obj.qubit.name for obj in objects] if objects is not None else []
+        )
+        unspecified_names = set(self.object_name_dict.keys()) - set(specified_names)
+
+        # Make sure we have all objects, starting with the specified ones in the given order.
+        ordered_names = specified_names + list(unspecified_names)
+        ordered_objects = [self.object_name_dict[name] for name in ordered_names]
+
+        # Peek the current world `count` times and get the results.
+        histogram = self.get_correlated_histogram(ordered_objects, count)
+
+        # Get an estimate of the state vector.
+        state_vector = np.array([0.0] * (2**num_all_qubits))
+        for key, val in histogram.items():
+            state_vector += self.__to_state_vector__(key) * np.sqrt(val * 1.0 / count)
+        density_matrix = np.outer(state_vector, state_vector)
+
+        if num_shown_qubits == num_all_qubits:
+            return density_matrix
+        else:
+            # We trace out the unspecified qubits.
+            # The reshape is required by the partial_trace method.
+            traced_density_matrix = cirq.partial_trace(
+                density_matrix.reshape((2, 2) * num_all_qubits),
+                range(num_shown_qubits),
+            )
+            # Reshape back to a 2-d matrix.
+            return traced_density_matrix.reshape(
+                2**num_shown_qubits, 2**num_shown_qubits
+            )
+
+    def measure_entanglement(self, obj1: QuantumObject, obj2: QuantumObject) -> float:
+        """Measures the entanglement (i.e. quantum mutual information) of the two given objects.
+        See https://en.wikipedia.org/wiki/Quantum_mutual_information for the formula.
+
+        Parameters:
+            obj1, obj2:     two quantum objects (currently only qubits are supported)
+
+        Returns:
+            The quantum mutual information defined as S_1 + S_2 - S_12, where S denotes (reduced)
+        von Neumann entropy.
+        """
+        density_matrix_12 = self.density_matrix([obj1, obj2]).reshape(2, 2, 2, 2)
+        density_matrix_1 = cirq.partial_trace(density_matrix_12, [0])
+        density_matrix_2 = cirq.partial_trace(density_matrix_12, [1])
+        return (
+            cirq.von_neumann_entropy(density_matrix_1, validate=False)
+            + cirq.von_neumann_entropy(density_matrix_2, validate=False)
+            - cirq.von_neumann_entropy(density_matrix_12.reshape(4, 4), validate=False)
+        )
+
     def __getitem__(self, name: str) -> QuantumObject:
         quantum_object = self.object_name_dict.get(name, None)
         if not quantum_object:
             raise KeyError(f"{name} did not exist in this world.")
         return quantum_object
+
+    def __to_state_vector__(self, input: tuple) -> np.ndarray:
+        """Converts the given tuple (of length N) to the corresponding state vector (of length 2**N).
+        e.g. (0, 1) -> [0, 1, 0, 0]
+        """
+        num = len(input)
+        index = int("".join([str(i) for i in input]), 2)
+        state_vector = np.array([0.0] * (2**num))
+        state_vector[index] = 1.0
+        return state_vector

unitary/alpha/quantum_world_test.py

@@ -17,6 +17,7 @@
 import pytest
 
 import numpy as np
+import numpy.testing as testing
 
 import cirq
 
@@ -892,3 +893,83 @@ def test_save_and_restore_snapshot(simulator, compile_to_qubits):
     # Further restore would return a value error.
     with pytest.raises(ValueError, match="Unable to restore any more."):
         world.restore_last_snapshot()
+
+
+@pytest.mark.parametrize(
+    ("simulator", "compile_to_qubits"),
+    [
+        (cirq.Simulator, False),
+        (cirq.Simulator, True),
+        # Cannot use SparseSimulator without `compile_to_qubits` due to issue #78.
+        (alpha.SparseSimulator, True),
+    ],
+)
+def test_density_matrix(simulator, compile_to_qubits):
+    rho_green = np.reshape([0, 0, 0, 1], (2, 2))
+    rho_red = np.reshape([1, 0, 0, 0], (2, 2))
+    light1 = alpha.QuantumObject("green", Light.GREEN)
+    light2 = alpha.QuantumObject("red1", Light.RED)
+    light3 = alpha.QuantumObject("red2", Light.RED)
+    board = alpha.QuantumWorld(
+        [light1, light2, light3],
+        sampler=simulator(),
+        compile_to_qubits=compile_to_qubits,
+    )
+
+    testing.assert_array_equal(board.density_matrix(objects=[light1]), rho_green)
+    testing.assert_array_equal(board.density_matrix(objects=[light2]), rho_red)
+    testing.assert_array_equal(board.density_matrix(objects=[light3]), rho_red)
+
+    testing.assert_array_equal(
+        board.density_matrix(objects=[light1, light2]), np.kron(rho_green, rho_red)
+    )
+    testing.assert_array_equal(
+        board.density_matrix(objects=[light2, light3]), np.kron(rho_red, rho_red)
+    )
+    testing.assert_array_equal(
+        board.density_matrix(objects=[light3, light1]), np.kron(rho_red, rho_green)
+    )
+
+    testing.assert_array_equal(
+        board.density_matrix(objects=[light1, light2, light3]),
+        np.kron(rho_green, np.kron(rho_red, rho_red)),
+    )
+
+
+@pytest.mark.parametrize(
+    ("simulator", "compile_to_qubits"),
+    [
+        (cirq.Simulator, False),
+        (cirq.Simulator, True),
+        # Cannot use SparseSimulator without `compile_to_qubits` due to issue #78.
+        (alpha.SparseSimulator, True),
+    ],
+)
+def test_measure_entanglement(simulator, compile_to_qubits):
+    rho_green = np.reshape([0, 0, 0, 1], (2, 2))
+    rho_red = np.reshape([1, 0, 0, 0], (2, 2))
+    light1 = alpha.QuantumObject("red1", Light.RED)
+    light2 = alpha.QuantumObject("green", Light.GREEN)
+    light3 = alpha.QuantumObject("red2", Light.RED)
+    board = alpha.QuantumWorld(
+        [light1, light2, light3],
+        sampler=simulator(),
+        compile_to_qubits=compile_to_qubits,
+    )
+
+    # S_1 + S_2 - S_12 = 0 + 0 - 0 = 0 for all three cases.
+    assert round(board.measure_entanglement(light1, light2)) == 0.0
+    assert round(board.measure_entanglement(light1, light3)) == 0.0
+    assert round(board.measure_entanglement(light2, light3)) == 0.0
+
+    alpha.Superposition()(light2)
+    alpha.quantum_if(light2).apply(alpha.Flip())(light3)
+    results = board.peek([light2, light3], count=100)
+    assert not all(result[0] == 0 for result in results)
+    assert (result[0] == result[1] for result in results)
+    # S_1 + S_2 - S_12 = 0 + 1 - 1 = 0
+    assert round(board.measure_entanglement(light1, light2), 1) == 0.0
+    # S_1 + S_2 - S_12 = 0 + 1 - 1 = 0
+    assert round(board.measure_entanglement(light1, light3), 1) == 0.0
+    # S_1 + S_2 - S_12 = 1 + 1 - 0 = 2
+    assert round(board.measure_entanglement(light2, light3), 1) == 2.0