add density matrix and entanglement calculations

examples/tic_tac_toe/README.md

@@ -5,3 +5,55 @@ After cloning the Unitary library you can use command line flags to run the game
 ```
 python -m examples.tic_tac_toe.tic_tac_toe
 ```
+
+## Rules of the game
+
+There are nine positions on the 3x3 board, labeled a through h.
+The object is to get three in a row.  Players alternate turns,
+placing either an X or an O on each square.
+
+In the quantum version, there is an additional move for placing
+an X or an O on two squares in superposition, called the split move.
+
+Once the board is full (i.e. each of the squares has either X,
+O, or a superposition of them), the board is measured and the
+result is determined.
+
+## Quantum representation
+
+Each square is represented as a Qutrit.  This square can either be:
+
+*  |0> meaning that the square is empty
+*  |1> meaning that the square has an X
+*  |2> meaning that the square has an O
+
+Placing an X or O will do the qutrit version of an "X" gate
+on that square.  Since these are qutrits (not qubits), the
+definition of this gate is that it is swapping the amplitude
+of the zero state with either the one or two state
+(depending on whether this is an X or O move).
+For example, if someone places an O on a square, then placing
+an X on that same square has no effect, since the zero state
+has zero amplitude.
+
+Note that there are two variants of the rules with regards to
+the split move.  In one rule, the qutrit X gate is applied
+followed by a square root of iSWAP gate (restricted to either the
+|1> or |2> subspace depending on the move).
+
+In the other variant, the corresponding states |01> and |10>
+(or |02> and |20>) are swapped.
+
+## Code organization
+
+The qutrit operations (such as the split move) are defined in
+`tic_tac_split.py`.  These classes can be used as examples of
+how to create your own gates and effects using unitary.
+
+The game itself is defined in `tic_tac_toe.py`.  The `GameInterface`
+class defines the main game loop and input/output.
+
+The `TicTacToe` class keeps track of the game state and also
+allows you to sample (measure) the board.  
+
+Enums used by the example are stored in `enums.py`.

examples/tic_tac_toe/enums.py

@@ -17,12 +17,20 @@
 
 
 class GameMoves(enum.Enum):
+    """Possible inputs for the ASCII version of tic tac toe."""
+
     EXIT = "exit"
     MAP = "map"
     HELP = "help"
 
 
 class TicTacSquare(enum.Enum):
+    """Possible states of one tic tac toe square.
+
+    For the quantum version of tic tac toe, these
+    are represented as qutrits (qubits with three states).
+    """
+
     EMPTY = 0
     X = 1
     O = 2
@@ -35,6 +43,14 @@ def from_result(cls, value: Union[enum.Enum, int]):
 
 
 class TicTacResult(enum.Enum):
+    """End results of a tic tac toe game.
+
+    Either one side has won or it is a draw.
+    If the game has continued past the end state,
+    it is possible both sides have completed a three-in-a-row.
+    If the game is not complete, the result is UNFINISHED.
+    """
+
     UNFINISHED = 0
     X_WINS = 1
     O_WINS = 2

examples/tic_tac_toe/tic_tac_toe.py

@@ -79,6 +79,13 @@ def _result_to_str(result: List[TicTacSquare]) -> str:
 
 
 def eval_board(result: List[TicTacSquare]) -> TicTacResult:
+    """Determines who has won the tic tac toe board.
+
+    This function checks all the possible three-in-a-row positions
+    (all cols, all rows, and the two diagonals.  Depending on which
+    player(s) have a three in a row (X's, O's, both, or neither)
+    returns the result of the tic tac toe board.
+    """
     x_wins = False
     o_wins = False
     still_empty = False
@@ -90,18 +97,12 @@ def eval_board(result: List[TicTacSquare]) -> TicTacResult:
         if any(result[check[idx]] == TicTacSquare.EMPTY for idx in range(3)):
             still_empty = True
     if x_wins:
-        if o_wins:
-            return TicTacResult.BOTH_WIN
-        else:
-            return TicTacResult.X_WINS
+        return TicTacResult.BOTH_WIN if o_wins else TicTacResult.X_WINS
     else:
         if o_wins:
             return TicTacResult.O_WINS
         else:
-            if still_empty:
-                return TicTacResult.UNFINISHED
-            else:
-                return TicTacResult.DRAW
+            return TicTacResult.UNFINISHED if still_empty else TicTacResult.DRAW
 
 
 class TicTacToe:

unitary/alpha/quantum_world.py

@@ -116,9 +116,9 @@ def copy(self) -> "QuantumWorld":
         for remap in self.qubit_remapping_dict:
             new_dict = {}
             for key_obj, value_obj in remap.items():
-                new_dict[new_world.get_object_by_name(key_obj.name)] = (
-                    new_world.get_object_by_name(value_obj.name)
-                )
+                new_dict[
+                    new_world.get_object_by_name(key_obj.name)
+                ] = new_world.get_object_by_name(value_obj.name)
             new_world.qubit_remapping_dict.append(new_dict)
         new_world.qubit_remapping_dict_length = self.qubit_remapping_dict_length.copy()
         return new_world
@@ -636,6 +636,71 @@ def get_binary_probabilities(
             binary_probs.append(1 - one_probs[0])
         return binary_probs
 
+    def density_matrix(
+        self, objects: Optional[Sequence[QuantumObject]] = None
+    ) -> np.ndarray:
+        """Simulates the density matrix of the given objects.
+
+        Parameters:
+            objects:    List of QuantumObjects (currently only qubits are supported)
+
+        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_qubits = [self.object_name_dict[name].qubit for name in ordered_names]
+
+        simulator = cirq.DensityMatrixSimulator()
+        qubit_order = cirq.QubitOrder.explicit(
+            ordered_qubits, fallback=cirq.QubitOrder.DEFAULT
+        )
+        result = simulator.simulate(self.circuit, qubit_order=qubit_order)
+
+        if num_shown_qubits == num_all_qubits:
+            return result.final_density_matrix
+        else:
+            # We trace out the unspecified qubits.
+            # The reshape is required by the partial_trace method.
+            traced_density_matrix = cirq.partial_trace(
+                result.final_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)
+        print(self.density_matrix([obj1, obj2]))
+        density_matrix_1 = cirq.partial_trace(density_matrix_12, [0])
+        print(density_matrix_1)
+        density_matrix_2 = cirq.partial_trace(density_matrix_12, [1])
+        print(density_matrix_2)
+        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:

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,59 @@ 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()
+
+
+def test_density_matrix():
+    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])
+
+    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, rho_green
+    )
+
+    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=[light1, light3]), np.kron(rho_green, rho_red)
+    )
+
+    testing.assert_array_equal(
+        board.density_matrix(objects=[light1, light2, light3]),
+        np.kron(rho_green, np.kron(rho_red, rho_red)),
+    )
+
+
+def test_measure_entanglement():
+    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])
+
+    # 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), 3) == 0.0
+    # S_1 + S_2 - S_12 = 0 + 1 - 1 = 0
+    assert round(board.measure_entanglement(light1, light3), 3) == 0.0
+    # S_1 + S_2 - S_12 = 1 + 1 - 0 = 2
+    assert round(board.measure_entanglement(light2, light3), 3) == 2.0