From 0e6a382ebe0429359f98ec3a09d1ba5244e80035 Mon Sep 17 00:00:00 2001
From: Pengfei Chen <pengfeichen@google.com>
Date: Tue, 17 Oct 2023 16:12:24 -0700
Subject: [PATCH 01/13] initial set

---
 unitary/alpha/quantum_world.py                | 11 +++
 .../examples/quantum_chinese_chess/board.py   |  2 +-
 .../examples/quantum_chinese_chess/chess.py   | 91 +++++++++++++++----
 .../quantum_chinese_chess/chess_test.py       |  4 +-
 .../examples/quantum_chinese_chess/enums.py   | 15 +--
 .../examples/quantum_chinese_chess/move.py    | 77 +++++++++++++++-
 .../examples/quantum_chinese_chess/piece.py   |  5 +-
 7 files changed, 174 insertions(+), 31 deletions(-)

diff --git a/unitary/alpha/quantum_world.py b/unitary/alpha/quantum_world.py
index d5e6390d..f93fd27d 100644
--- a/unitary/alpha/quantum_world.py
+++ b/unitary/alpha/quantum_world.py
@@ -305,6 +305,17 @@ def _interpret_result(self, result: Union[int, Iterable[int]]) -> int:
             return result_list[0]
         return result
 
+    def unhook(self, object: QuantumObject) -> None:
+        """Replaces all usages of the given object in the circuit with a new ancilla with value=0.
+        """
+        new_ancilla = self._add_ancilla(object.name)
+        # Replace operations using the qubit of the given object with the new ancilla.
+        qubit_remapping_dict = {object.qubit: new_ancilla.qubit, new_ancilla.qubit: object.qubit}
+        self.circuit = self.circuit.transform_qubits(
+            lambda q: qubit_remapping_dict.get(q, q)
+        )
+        return
+
     def force_measurement(
         self, obj: QuantumObject, result: Union[enum.Enum, int]
     ) -> None:
diff --git a/unitary/examples/quantum_chinese_chess/board.py b/unitary/examples/quantum_chinese_chess/board.py
index c2401315..97f64395 100644
--- a/unitary/examples/quantum_chinese_chess/board.py
+++ b/unitary/examples/quantum_chinese_chess/board.py
@@ -165,7 +165,7 @@ def path_pieces(self, source: str, target: str) -> Tuple[List[str], List[str]]:
         elif abs(dy) == 2 and abs(dx) == 1:
             pieces.append(f"{chr(x0)}{y0 + dy_sign}")
         else:
-            raise ValueError("Unexpected input to path_pieces().")
+            raise ValueError("The input move is illegal.")
         for piece in pieces:
             if self.board[piece].is_entangled:
                 quantum_pieces.append(piece)
diff --git a/unitary/examples/quantum_chinese_chess/chess.py b/unitary/examples/quantum_chinese_chess/chess.py
index 496ad8cc..1da10a23 100644
--- a/unitary/examples/quantum_chinese_chess/chess.py
+++ b/unitary/examples/quantum_chinese_chess/chess.py
@@ -21,7 +21,8 @@
     MoveType,
     MoveVariant,
 )
-from unitary.examples.quantum_chinese_chess.move import Move
+from unitary.examples.quantum_chinese_chess.move import Jump
+import readline
 
 # List of accepable commands.
 _HELP_TEXT = """
@@ -235,7 +236,7 @@ def classify_move(
             quantum_path_pieces_1: the list of names of quantum pieces from source_0 to target_1 (for split) or
                                      from source_1 to target_0 (for merge) (excluded)
         """
-        move_type = MoveType.UNSPECIFIED_STANDARD
+        move_type = MoveType.UNSPECIFIED
         move_variant = MoveVariant.UNSPECIFIED
 
         source = self.board.board[sources[0]]
@@ -257,7 +258,7 @@ def classify_move(
                     # This handles all classical cases, where no quantum piece is envolved.
                     # We don't need to further classify MoveVariant types since all classical cases
                     # will be handled in a similar way.
-                    return MoveType.CLASSICAL, MoveVariant.UNSPECIFIED
+                    return MoveType.CLASSICAL, MoveVariant.CLASSICAL
                 else:
                     # If any of the source or target is entangled, this move is a JUMP.
                     move_type = MoveType.JUMP
@@ -265,8 +266,12 @@ def classify_move(
                 # If there is any quantum path pieces, this move is a SLIDE.
                 move_type = MoveType.SLIDE
 
-            if source.type_ == Type.CANNON and (
-                len(classical_path_pieces_0) == 1 or len(quantum_path_pieces_0) > 0
+            if (
+                source.type_ == Type.CANNON
+                and (
+                    len(classical_path_pieces_0) == 1 or len(quantum_path_pieces_0) > 0
+                )
+                and target.color.value == 1 - source.color.value
             ):
                 # By this time the classical cannon fire has been identified as CLASSICAL move,
                 # so the current case has quantum piece(s) envolved.
@@ -381,6 +386,9 @@ def apply_move(self, str_to_parse: str) -> None:
             quantum_pieces_1,
         )
 
+        print(move_type, " ", move_variant)
+
+        # Apply the move accoding to its type.
         if move_type == MoveType.CLASSICAL:
             if source_0.type_ == Type.KING:
                 # Update the locations of KING.
@@ -390,34 +398,74 @@ def apply_move(self, str_to_parse: str) -> None:
             if target_0.type_ == Type.KING:
                 # King is captured, then the game is over.
                 self.game_state = GameState(self.current_player)
-            target_0.reset(source_0)
-            source_0.reset()
-            # TODO(): only make such prints for a certain debug level.
-            print("Classical move.")
+            Jump(move_variant)(source_0, target_0)
+        elif move_type == MoveType.JUMP:
+            Jump(move_variant)(source_0, target_0)
         # TODO(): apply other move types.
 
-    def next_move(self) -> bool:
+    def next_move(self) -> Tuple[bool, str]:
         """Check if the player wants to exit or needs help message. Otherwise parse and apply the move.
-        Returns True if the move was made, otherwise returns False.
+        Returns True + output string if the move was made, otherwise returns False + output string.
         """
         input_str = input(
             f"\nIt is {self.players_name[self.current_player]}'s turn to move: "
         )
+        output = ""
         if input_str.lower() == "help":
-            print(_HELP_TEXT)
+            output = _HELP_TEXT
         elif input_str.lower() == "exit":
             # The other player wins if the current player quits.
             self.game_state = GameState(1 - self.current_player)
-            print("Exiting.")
+            output = "Exiting."
+        elif input_str.lower() == "peek":
+            # TODO(): make it look like the normal board. Right now it's only for debugging purposes.
+            print(self.board.board.peek(convert_to_enum=False))
+        elif input_str.lower() == "undo":
+            output = "Undo last quantum effect."
+            # Right now it's only for debugging purposes, since it has following problems:
+            # TODO(): there are several problems here:
+            # 1) last move is quantum but classical piece information is not reversed back.
+            # ==> we may need to save the change of classical piece information of each step.
+            # 2) last move might not be quantum.
+            # ==> we may need to save all classical moves and figure out how to undo each kind of move;
+            # 3) last move is quantum but involved multiple effects.
+            # ==> we may need to save number of effects per move, and undo that number of times.
+            self.board.board.undo_last_effect()
+            return True, output
         else:
             try:
                 # The move is success if no ValueError is raised.
                 self.apply_move(input_str.lower())
-                return True
+                return True, output
             except ValueError as e:
-                print("Invalid move.")
-                print(e)
-        return False
+                output = f"Invalid move. {e}"
+        return False, output
+
+    def update_board_by_sampling(self) -> List[float]:
+        """After quantum moves, there might be pieces that:
+        - is actually empty, but their classical properties is not cleared; or
+        - is actually classically occupied, but their is_entangled state is not updated.
+        This method is called after each quantum move, and runs (100x) sampling of the board
+        to identify and fix those cases.
+        """
+        # TODO(): return the sampled probabilities and pass it into the print method
+        # of the board to print it together with the board.
+        probs = self.board.board.get_binary_probabilities()
+        num_rows = 10
+        num_cols = 9
+        for row in range(num_rows):
+            for col in "abcdefghi":
+                piece = self.board.board[f"{col}{row}"]
+                # We need to do the following range() conversion since the sequence of
+                # qubits returned from get_binary_probabilities() is
+                # a9 b9 ... i9, a8 b8 ... i8, ..., a0 b0 ... i0
+                prob = probs[(num_rows - row - 1) * num_cols + ord(col) - ord("a")]
+                # TODO(): This threshold does not actually work right now since we have 100 sampling.
+                # Change it to be more meaningful values maybe when we do error mitigation.
+                if prob < 1e-3:
+                    piece.reset()
+                elif prob > 1 - 1e-3:
+                    piece.is_entangled = False
 
     def game_over(self) -> None:
         """Checks if the game is over, and update self.game_state accordingly."""
@@ -433,15 +481,20 @@ def game_over(self) -> None:
     def play(self) -> None:
         """The loop where each player takes turn to play."""
         while True:
-            move_success = self.next_move()
-            print(self.board)
+            move_success, output = self.next_move()
             if not move_success:
                 # Continue if the player does not quit.
                 if self.game_state == GameState.CONTINUES:
+                    print(output)
                     print("\nPlease re-enter your move.")
                     continue
+            print(output)
+            # TODO(): maybe we should not check game_over() when an undo is made.
             # Check if the game is over.
             self.game_over()
+            # TODO(): no need to do sampling if the last move was CLASSICAL.
+            self.update_board_by_sampling()
+            print(self.board)
             if self.game_state == GameState.CONTINUES:
                 # If the game continues, switch the player.
                 self.current_player = 1 - self.current_player
diff --git a/unitary/examples/quantum_chinese_chess/chess_test.py b/unitary/examples/quantum_chinese_chess/chess_test.py
index 67278ee9..11f84f7c 100644
--- a/unitary/examples/quantum_chinese_chess/chess_test.py
+++ b/unitary/examples/quantum_chinese_chess/chess_test.py
@@ -260,11 +260,11 @@ def test_classify_move_success(monkeypatch):
     # classical
     assert game.classify_move(["h9"], ["g7"], [], [], [], []) == (
         MoveType.CLASSICAL,
-        MoveVariant.UNSPECIFIED,
+        MoveVariant.CLASSICAL,
     )
     assert game.classify_move(["b2"], ["b9"], ["b7"], [], [], []) == (
         MoveType.CLASSICAL,
-        MoveVariant.UNSPECIFIED,
+        MoveVariant.CLASSICAL,
     )
 
     # jump basic
diff --git a/unitary/examples/quantum_chinese_chess/enums.py b/unitary/examples/quantum_chinese_chess/enums.py
index 485f2b91..f017cec2 100644
--- a/unitary/examples/quantum_chinese_chess/enums.py
+++ b/unitary/examples/quantum_chinese_chess/enums.py
@@ -45,8 +45,8 @@ class GameState(enum.Enum):
 class MoveType(enum.Enum):
     """Each valid move will be classfied into one of the following MoveTypes."""
 
-    CLASSICAL = 0
-    UNSPECIFIED_STANDARD = 1
+    UNSPECIFIED = 0
+    CLASSICAL = 1
     JUMP = 2
     SLIDE = 3
     SPLIT_JUMP = 4
@@ -61,10 +61,11 @@ class MoveVariant(enum.Enum):
     the MoveType above.
     """
 
-    UNSPECIFIED = 0  # Used together with MoveType = CLASSICAL.
-    BASIC = 1  # The target piece is empty.
-    EXCLUDED = 2  # The target piece has the same color.
-    CAPTURE = 3  # The target piece has the opposite color.
+    UNSPECIFIED = 0
+    CLASSICAL = 1  # Used together with MoveType = CLASSICAL.
+    BASIC = 2  # The target piece is empty.
+    EXCLUDED = 3  # The target piece has the same color.
+    CAPTURE = 4  # The target piece has the opposite color.
 
 
 class Color(enum.Enum):
@@ -110,7 +111,7 @@ def type_of(c: str) -> Optional["Type"]:
     def symbol(type_: "Type", color: Color, lang: Language = Language.EN) -> str:
         """Returns symbol of the given piece according to its color and desired language."""
         if type_ == Type.EMPTY:
-            return "."
+            return type_.value[0]
         if lang == Language.EN:  # Return English symbols
             if color == Color.RED:
                 return type_.value[0]
diff --git a/unitary/examples/quantum_chinese_chess/move.py b/unitary/examples/quantum_chinese_chess/move.py
index ec0fb6bf..32e74d72 100644
--- a/unitary/examples/quantum_chinese_chess/move.py
+++ b/unitary/examples/quantum_chinese_chess/move.py
@@ -11,12 +11,17 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Optional, List, Tuple
+from typing import Optional, List, Tuple, Iterator
+import cirq
+from unitary import alpha
 from unitary.alpha.quantum_effect import QuantumEffect
 from unitary.examples.quantum_chinese_chess.board import Board
+from unitary.examples.quantum_chinese_chess.piece import Piece
 from unitary.examples.quantum_chinese_chess.enums import MoveType, MoveVariant, Type
 
 
+# TODO(): now the class is no longer the base class of all chess moves. Maybe convert this class
+# to a helper class to save each move (with its pop results) in a string form into move history.
 class Move(QuantumEffect):
     """The base class of all chess moves."""
 
@@ -101,3 +106,73 @@ def to_str(self, verbose_level: int = 1) -> str:
 
     def __str__(self):
         return self.to_str()
+
+
+class Jump(QuantumEffect):
+    """Jump from source_0 to target_0. The accepted move_variant includes
+    - CLASSICAL (where all classical moves will be handled here)
+    - CAPTURE
+    - EXCLUDED
+    - BASIC
+    """
+
+    def __init__(
+        self,
+        move_variant: MoveVariant,
+    ):
+        self.move_variant = move_variant
+
+    def num_dimension(self) -> Optional[int]:
+        return 2
+
+    def num_objects(self) -> Optional[int]:
+        return 2
+
+    def effect(self, *objects) -> Iterator[cirq.Operation]:
+        # TODO(): currently pawn capture is a same as jump capture, while in quantum chess it's different,
+        # i.e. pawn would move only if the target is there, i.e. CNOT(t, s), and an entanglement could be
+        # created. This could be a general game setting, i.e. we could allow players to choose if they
+        # want the source piece to move (in case of capture) if the target piece is not there.
+        source_0, target_0 = objects
+        world = source_0.world
+        if self.move_variant == MoveVariant.CAPTURE:
+            # We peek and force measure source_0.
+            source_is_occupied = world.pop([source_0])[0]
+            # For move_variant==CAPTURE, we require source_0 to be occupied before further actions.
+            # This is to prevent a piece of the board containing two types of different pieces.
+            if not source_is_occupied:
+                # If source_0 turns out to be not there, we clear set it to be EMPTY, and the jump
+                # could not be made.
+                source_0.reset()
+                print("Jump move: source turns out to be empty.")
+                return iter(())
+            source_0.is_entangled = False
+            # We replace the qubit of target_0 with a new ancilla, and set its classical properties to be EMPTY.
+            world.unhook(target_0)
+            target_0.reset()
+        elif self.move_variant == MoveVariant.EXCLUDED:
+            # We peek and force measure target_0.
+            target_is_occupied = world.pop([target_0])[0]
+            # For move_variant==EXCLUDED, we require target_0 to be empty before further actions.
+            # This is to prevent a piece of the board containing two types of different pieces.
+            if target_is_occupied:
+                # If target_0 turns out to be there, we set it to be a classically OCCUPIED, and
+                # the jump could not be made.
+                print("Jump move: target turns out to be occupied.")
+                target_0.is_entangled = False
+                return iter(())
+            # Otherwise we set target_0 to be classically EMPTY.
+            target_0.reset()
+        elif self.move_variant == MoveVariant.CLASSICAL:
+            if target_0.type_ != Type.EMPTY:
+                # For classical moves with target_0 occupied, we replace the qubit of target_0 with
+                # a new ancilla, and set its classical properties to be EMPTY.
+                world.unhook(target_0)
+                target_0.reset()
+
+        # Make the jump move.
+        alpha.PhasedMove()(source_0, target_0)
+        # Move the classical properties of the source piece to the target piece.
+        target_0.reset(source_0)
+        source_0.reset()
+        return iter(())
diff --git a/unitary/examples/quantum_chinese_chess/piece.py b/unitary/examples/quantum_chinese_chess/piece.py
index e8d4b79a..9dd6e763 100644
--- a/unitary/examples/quantum_chinese_chess/piece.py
+++ b/unitary/examples/quantum_chinese_chess/piece.py
@@ -40,11 +40,14 @@ def __str__(self):
 
     def reset(self, piece: "Piece" = None) -> None:
         """Modifies the classical attributes of the piece.
-        If piece is provided, then its type_ and color is copied, otherwise set the current piece to be an empty piece.
+        If piece is provided, then its type_, color, and is_entangled is copied,
+        otherwise set the current piece to be a classically empty piece.
         """
         if piece is not None:
             self.type_ = piece.type_
             self.color = piece.color
+            self.is_entangled = piece.is_entangled
         else:
             self.type_ = Type.EMPTY
             self.color = Color.NA
+            self.is_entangled = False

From 1b69d6fafcfaa175a3cc6394685e81d5f7b62579 Mon Sep 17 00:00:00 2001
From: Pengfei Chen <pengfeichen@google.com>
Date: Tue, 17 Oct 2023 16:26:41 -0700
Subject: [PATCH 02/13] update

---
 unitary/alpha/quantum_world.py                       | 9 ++++++---
 unitary/examples/quantum_chinese_chess/chess.py      | 4 +++-
 unitary/examples/quantum_chinese_chess/piece_test.py | 3 +++
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/unitary/alpha/quantum_world.py b/unitary/alpha/quantum_world.py
index f93fd27d..f5f74bb3 100644
--- a/unitary/alpha/quantum_world.py
+++ b/unitary/alpha/quantum_world.py
@@ -306,11 +306,14 @@ def _interpret_result(self, result: Union[int, Iterable[int]]) -> int:
         return result
 
     def unhook(self, object: QuantumObject) -> None:
-        """Replaces all usages of the given object in the circuit with a new ancilla with value=0.
-        """
+        """Replaces all usages of the given object in the circuit with a new ancilla with value=0."""
+        # Creates a new ancilla.
         new_ancilla = self._add_ancilla(object.name)
         # Replace operations using the qubit of the given object with the new ancilla.
-        qubit_remapping_dict = {object.qubit: new_ancilla.qubit, new_ancilla.qubit: object.qubit}
+        qubit_remapping_dict = {
+            object.qubit: new_ancilla.qubit,
+            new_ancilla.qubit: object.qubit,
+        }
         self.circuit = self.circuit.transform_qubits(
             lambda q: qubit_remapping_dict.get(q, q)
         )
diff --git a/unitary/examples/quantum_chinese_chess/chess.py b/unitary/examples/quantum_chinese_chess/chess.py
index 1da10a23..ee30cf7c 100644
--- a/unitary/examples/quantum_chinese_chess/chess.py
+++ b/unitary/examples/quantum_chinese_chess/chess.py
@@ -449,7 +449,9 @@ def update_board_by_sampling(self) -> List[float]:
         to identify and fix those cases.
         """
         # TODO(): return the sampled probabilities and pass it into the print method
-        # of the board to print it together with the board.
+        # of the board to print it together with the board, or better use mathemetical
+        # matrix calculations to determine the probability, and use it (with some error
+        # threshold) to update the piece infos.
         probs = self.board.board.get_binary_probabilities()
         num_rows = 10
         num_cols = 9
diff --git a/unitary/examples/quantum_chinese_chess/piece_test.py b/unitary/examples/quantum_chinese_chess/piece_test.py
index c15ac942..291a2691 100644
--- a/unitary/examples/quantum_chinese_chess/piece_test.py
+++ b/unitary/examples/quantum_chinese_chess/piece_test.py
@@ -52,11 +52,14 @@ def test_enum():
 def test_reset():
     p0 = Piece("a0", SquareState.OCCUPIED, Type.CANNON, Color.RED)
     p1 = Piece("b1", SquareState.OCCUPIED, Type.HORSE, Color.BLACK)
+    p1.is_entangled = True
 
     p0.reset()
     assert p0.type_ == Type.EMPTY
     assert p0.color == Color.NA
+    assert p0.is_entangled == False
 
     p0.reset(p1)
     assert p0.type_ == p1.type_
     assert p0.color == p1.color
+    assert p0.is_entangled == p1.is_entangled

From 8c18f5fd87d6051ba444bc49a4fb17dc0647ffd7 Mon Sep 17 00:00:00 2001
From: madcpf <madcpf@gmail.com>
Date: Tue, 17 Oct 2023 23:54:47 -0700
Subject: [PATCH 03/13] update test

---
 .../examples/quantum_chinese_chess/board.py   |  10 +-
 .../quantum_chinese_chess/board_test.py       |  76 +++++-----
 .../examples/quantum_chinese_chess/chess.py   |  14 +-
 .../quantum_chinese_chess/chess_test.py       |  46 +++---
 .../examples/quantum_chinese_chess/move.py    |   8 +-
 .../quantum_chinese_chess/move_test.py        |  83 ++++++++++-
 .../quantum_chinese_chess/test_utils.py       | 141 +++++++++++++++++-
 7 files changed, 295 insertions(+), 83 deletions(-)

diff --git a/unitary/examples/quantum_chinese_chess/board.py b/unitary/examples/quantum_chinese_chess/board.py
index 97f64395..efd3fbf4 100644
--- a/unitary/examples/quantum_chinese_chess/board.py
+++ b/unitary/examples/quantum_chinese_chess/board.py
@@ -49,7 +49,7 @@ def from_fen(cls, fen: str = _INITIAL_FEN) -> "Board":
         FEN rule for Chinese Chess could be found at https://www.wxf-xiangqi.org/images/computer-xiangqi/fen-for-xiangqi-chinese-chess.pdf
         """
         chess_board = {}
-        row_index = 9
+        row_index = 0
         king_locations = []
         pieces, turns = fen.split(" ", 1)
         for row in pieces.split("/"):
@@ -74,15 +74,11 @@ def from_fen(cls, fen: str = _INITIAL_FEN) -> "Board":
                         name, SquareState.OCCUPIED, piece_type, color
                     )
                     col += 1
-            row_index -= 1
+            row_index += 1
         board = alpha.QuantumWorld(chess_board.values())
         # Here 0 means the player RED while 1 the player BLACK.
         current_player = 0 if "w" in turns else 1
         # TODO(): maybe add check to make sure the input fen itself is correct.
-        if len(king_locations) != 2:
-            raise ValueError(
-                f"We expect two KINGs on the board, but got {len(king_locations)}."
-            )
         return cls(board, current_player, king_locations)
 
     def __str__(self):
@@ -92,7 +88,7 @@ def __str__(self):
         for col in "abcdefghi":
             board_string.append(f" {col}")
         board_string.append("\n")
-        for row in range(num_rows):
+        for row in range(num_rows - 1, -1, -1):
             # Print the row index on the left.
             board_string.append(f"{row} ")
             for col in "abcdefghi":
diff --git a/unitary/examples/quantum_chinese_chess/board_test.py b/unitary/examples/quantum_chinese_chess/board_test.py
index 6aec090b..220c5439 100644
--- a/unitary/examples/quantum_chinese_chess/board_test.py
+++ b/unitary/examples/quantum_chinese_chess/board_test.py
@@ -28,16 +28,16 @@ def test_init_with_default_fen():
         board.__str__()
         == """
   a b c d e f g h i
-0 r h e a k a e h r  0
-1 . . . . . . . . .  1
-2 . c . . . . . c .  2
-3 p . p . p . p . p  3
-4 . . . . . . . . .  4
-5 . . . . . . . . .  5
-6 P . P . P . P . P  6
-7 . C . . . . . C .  7
+9 r h e a k a e h r  9
 8 . . . . . . . . .  8
-9 R H E A K A E H R  9
+7 . c . . . . . c .  7
+6 p . p . p . p . p  6
+5 . . . . . . . . .  5
+4 . . . . . . . . .  4
+3 P . P . P . P . P  3
+2 . C . . . . . C .  2
+1 . . . . . . . . .  1
+0 R H E A K A E H R  0
   a b c d e f g h i
 """
     )
@@ -47,21 +47,21 @@ def test_init_with_default_fen():
         board.__str__()
         == """
 　ａｂｃｄｅｆｇｈｉ
-０車馬相仕帥仕相馬車０
-１．．．．．．．．．１
-２．砲．．．．．砲．２
-３卒．卒．卒．卒．卒３
-４．．．．．．．．．４
-５．．．．．．．．．５
-６兵．兵．兵．兵．兵６
-７．炮．．．．．炮．７
+９車馬相仕帥仕相馬車９
 ８．．．．．．．．．８
-９车马象士将士象马车９
+７．砲．．．．．砲．７
+６卒．卒．卒．卒．卒６
+５．．．．．．．．．５
+４．．．．．．．．．４
+３兵．兵．兵．兵．兵３
+２．炮．．．．．炮．２
+１．．．．．．．．．１
+０车马象士将士象马车０
 　ａｂｃｄｅｆｇｈｉ
 """
     )
 
-    assert board.king_locations == ["e9", "e0"]
+    assert board.king_locations == ["e0", "e9"]
 
 
 def test_init_with_specified_fen():
@@ -71,16 +71,16 @@ def test_init_with_specified_fen():
         board.__str__()
         == """
   a b c d e f g h i
-0 . . . A K . . . c  0
-1 . . . . A p . r .  1
-2 . . . . . . . . .  2
-3 . . . . . . . . .  3
-4 . . . . . . . . .  4
+9 . . . A K . . . c  9
+8 . . . . A p . r .  8
+7 . . . . . . . . .  7
+6 . . . . . . . . .  6
 5 . . . . . . . . .  5
-6 . . . . . . . H .  6
-7 . . . h R . . . .  7
-8 . . . . a . . . .  8
-9 . . . . k a R . .  9
+4 . . . . . . . . .  4
+3 . . . . . . . H .  3
+2 . . . h R . . . .  2
+1 . . . . a . . . .  1
+0 . . . . k a R . .  0
   a b c d e f g h i
 """
     )
@@ -90,21 +90,21 @@ def test_init_with_specified_fen():
         board.__str__()
         == """
 　ａｂｃｄｅｆｇｈｉ
-０．．．士将．．．砲０
-１．．．．士卒．車．１
-２．．．．．．．．．２
-３．．．．．．．．．３
-４．．．．．．．．．４
+９．．．士将．．．砲９
+８．．．．士卒．車．８
+７．．．．．．．．．７
+６．．．．．．．．．６
 ５．．．．．．．．．５
-６．．．．．．．马．６
-７．．．馬车．．．．７
-８．．．．仕．．．．８
-９．．．．帥仕车．．９
+４．．．．．．．．．４
+３．．．．．．．马．３
+２．．．馬车．．．．２
+１．．．．仕．．．．１
+０．．．．帥仕车．．０
 　ａｂｃｄｅｆｇｈｉ
 """
     )
 
-    assert board.king_locations == ["e9", "e0"]
+    assert board.king_locations == ["e0", "e9"]
 
 
 def test_path_pieces():
diff --git a/unitary/examples/quantum_chinese_chess/chess.py b/unitary/examples/quantum_chinese_chess/chess.py
index ee30cf7c..976f68d7 100644
--- a/unitary/examples/quantum_chinese_chess/chess.py
+++ b/unitary/examples/quantum_chinese_chess/chess.py
@@ -125,7 +125,7 @@ def _is_in_palace(self, color: Color, x: int, y: int) -> bool:
         return (
             x <= ord("f")
             and x >= ord("d")
-            and ((color == Color.RED and y >= 7) or (color == Color.BLACK and y <= 2))
+            and ((color == Color.RED and y <= 2) or (color == Color.BLACK and y >= 7))
         )
 
     def check_classical_rule(
@@ -168,8 +168,8 @@ def check_classical_rule(
         elif source_piece.type_ == Type.ELEPHANT:
             if not (abs(dx) == 2 and abs(dy) == 2):
                 raise ValueError("ELEPHANT cannot move like this.")
-            if (source_piece.color == Color.RED and y1 < 5) or (
-                source_piece.color == Color.BLACK and y1 > 4
+            if (source_piece.color == Color.RED and y1 > 4) or (
+                source_piece.color == Color.BLACK and y1 < 5
             ):
                 raise ValueError(
                     "ELEPHANT cannot cross the river (i.e. the middle line)."
@@ -200,16 +200,16 @@ def check_classical_rule(
             if abs(dx) + abs(dy) != 1:
                 raise ValueError("PAWN cannot move like this.")
             if source_piece.color == Color.RED:
-                if dy == 1:
+                if dy == -1:
                     raise ValueError("PAWN can not move backward.")
-                if y0 > 4 and dy != -1:
+                if y0 <= 4 and dy != 1:
                     raise ValueError(
                         "PAWN can only go forward before crossing the river (i.e. the middle line)."
                     )
             else:
-                if dy == -1:
+                if dy == 1:
                     raise ValueError("PAWN can not move backward.")
-                if y0 <= 4 and dy != 1:
+                if y0 > 4 and dy != -1:
                     raise ValueError(
                         "PAWN can only go forward before crossing the river (i.e. the middle line)."
                     )
diff --git a/unitary/examples/quantum_chinese_chess/chess_test.py b/unitary/examples/quantum_chinese_chess/chess_test.py
index 11f84f7c..bd00ec65 100644
--- a/unitary/examples/quantum_chinese_chess/chess_test.py
+++ b/unitary/examples/quantum_chinese_chess/chess_test.py
@@ -73,15 +73,15 @@ def test_apply_move_fail(monkeypatch):
     monkeypatch.setattr("builtins.input", lambda _: next(inputs))
     game = QuantumChineseChess()
     with pytest.raises(ValueError, match="Could not move empty piece."):
-        game.apply_move("a1b1")
+        game.apply_move("a8b8")
     with pytest.raises(ValueError, match="Could not move the other player's piece."):
-        game.apply_move("a0b1")
+        game.apply_move("a9b8")
     with pytest.raises(ValueError, match="Two sources need to be the same type."):
-        game.apply_move("a9a6^a5")
+        game.apply_move("a0a3^a4")
     with pytest.raises(ValueError, match="Two targets need to be the same type."):
-        game.apply_move("b7^a7h7")
+        game.apply_move("b2^a2h2")
     with pytest.raises(ValueError, match="Two targets need to be the same color."):
-        game.apply_move("b7^b2h7")
+        game.apply_move("b2^b7h2")
 
 
 def test_game_invalid_move(monkeypatch):
@@ -128,35 +128,35 @@ def test_check_classical_rule(monkeypatch):
     game.check_classical_rule("c0", "e2", [])
     with pytest.raises(ValueError, match="ELEPHANT cannot move like this."):
         game.check_classical_rule("c0", "e1", [])
-    game.board.board["g4"].reset(
-        Piece("g4", SquareState.OCCUPIED, Type.ELEPHANT, Color.BLACK)
+    game.board.board["g5"].reset(
+        Piece("g5", SquareState.OCCUPIED, Type.ELEPHANT, Color.BLACK)
     )
     with pytest.raises(ValueError, match="ELEPHANT cannot cross the river"):
-        game.check_classical_rule("g4", "i6", [])
-    game.board.board["c5"].reset(
-        Piece("c5", SquareState.OCCUPIED, Type.ELEPHANT, Color.RED)
+        game.check_classical_rule("g5", "i3", [])
+    game.board.board["c4"].reset(
+        Piece("c4", SquareState.OCCUPIED, Type.ELEPHANT, Color.RED)
     )
     with pytest.raises(ValueError, match="ELEPHANT cannot cross the river"):
-        game.check_classical_rule("c5", "e3", [])
+        game.check_classical_rule("c4", "e6", [])
 
     # ADVISOR
     game.check_classical_rule("d9", "e8", [])
     with pytest.raises(ValueError, match="ADVISOR cannot move like this."):
         game.check_classical_rule("d9", "d8", [])
     with pytest.raises(ValueError, match="ADVISOR cannot leave the palace."):
-        game.check_classical_rule("d0", "c1", [])
+        game.check_classical_rule("d9", "c8", [])
     with pytest.raises(ValueError, match="ADVISOR cannot leave the palace."):
-        game.check_classical_rule("f9", "g8", [])
+        game.check_classical_rule("f0", "g1", [])
 
     # KING
     game.check_classical_rule("e9", "e8", [])
     with pytest.raises(ValueError, match="KING cannot move like this."):
         game.check_classical_rule("e9", "d8", [])
-    game.board.board["c9"].reset()
-    game.board.board["d9"].reset(game.board.board["e9"])
-    game.board.board["e9"].reset()
+    game.board.board["c0"].reset()
+    game.board.board["d0"].reset(game.board.board["e0"])
+    game.board.board["e0"].reset()
     with pytest.raises(ValueError, match="KING cannot leave the palace."):
-        game.check_classical_rule("d9", "c9", [])
+        game.check_classical_rule("d0", "c0", [])
 
     # CANNON
     game.check_classical_rule("b7", "b4", [])
@@ -178,21 +178,21 @@ def test_check_classical_rule(monkeypatch):
         game.check_classical_rule("b3", "d3", ["c3"])
 
     # PAWN
-    game.check_classical_rule("a6", "a5", [])
+    game.check_classical_rule("a3", "a4", [])
     with pytest.raises(ValueError, match="PAWN cannot move like this."):
-        game.check_classical_rule("a6", "a4", [])
+        game.check_classical_rule("a3", "a5", [])
     with pytest.raises(
         ValueError, match="PAWN can only go forward before crossing the river"
     ):
-        game.check_classical_rule("a6", "b6", [])
+        game.check_classical_rule("e3", "f3", [])
     with pytest.raises(
         ValueError, match="PAWN can only go forward before crossing the river"
     ):
-        game.check_classical_rule("g3", "h3", [])
+        game.check_classical_rule("g6", "h6", [])
     with pytest.raises(ValueError, match="PAWN can not move backward."):
-        game.check_classical_rule("a6", "a7", [])
+        game.check_classical_rule("a3", "a2", [])
     with pytest.raises(ValueError, match="PAWN can not move backward."):
-        game.check_classical_rule("g3", "g2", [])
+        game.check_classical_rule("g6", "g7", [])
     # After crossing the rive the pawn could move horizontally.
     game.board.board["c4"].reset(game.board.board["c6"])
     game.board.board["c6"].reset()
diff --git a/unitary/examples/quantum_chinese_chess/move.py b/unitary/examples/quantum_chinese_chess/move.py
index 32e74d72..91da5d40 100644
--- a/unitary/examples/quantum_chinese_chess/move.py
+++ b/unitary/examples/quantum_chinese_chess/move.py
@@ -137,10 +137,13 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
         world = source_0.world
         if self.move_variant == MoveVariant.CAPTURE:
             # We peek and force measure source_0.
-            source_is_occupied = world.pop([source_0])[0]
+            print("Here 1 ")
+            source_is_occupied = world.pop([source_0])[0].value
+            print(source_is_occupied)
             # For move_variant==CAPTURE, we require source_0 to be occupied before further actions.
             # This is to prevent a piece of the board containing two types of different pieces.
             if not source_is_occupied:
+                print("Here 2 ")
                 # If source_0 turns out to be not there, we clear set it to be EMPTY, and the jump
                 # could not be made.
                 source_0.reset()
@@ -152,7 +155,7 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
             target_0.reset()
         elif self.move_variant == MoveVariant.EXCLUDED:
             # We peek and force measure target_0.
-            target_is_occupied = world.pop([target_0])[0]
+            target_is_occupied = world.pop([target_0])[0].value
             # For move_variant==EXCLUDED, we require target_0 to be empty before further actions.
             # This is to prevent a piece of the board containing two types of different pieces.
             if target_is_occupied:
@@ -171,6 +174,7 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
                 target_0.reset()
 
         # Make the jump move.
+        print("Here 3 ")        
         alpha.PhasedMove()(source_0, target_0)
         # Move the classical properties of the source piece to the target piece.
         target_0.reset(source_0)
diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index 4e787794..5232fcf4 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -11,10 +11,50 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from unitary.examples.quantum_chinese_chess.move import Move
+from unitary.examples.quantum_chinese_chess.move import Move, Jump
 from unitary.examples.quantum_chinese_chess.board import Board
-from unitary.examples.quantum_chinese_chess.enums import MoveType, MoveVariant
+from unitary.examples.quantum_chinese_chess.piece import Piece
+from unitary.examples.quantum_chinese_chess.enums import (
+    MoveType,
+    MoveVariant,
+    SquareState,
+    Type,
+    Color,
+)
+from unitary.examples.quantum_chinese_chess.test_utils import (
+    locations_to_bitboard,
+    assert_samples_in,
+    assert_sample_distribution,
+    assert_this_or_that,
+    assert_prob_about,
+    assert_fifty_fifty,
+    sample_board,
+    get_board_probability_distribution,
+    print_samples,
+)
 import pytest
+from unitary import alpha
+from typing import List
+from string import ascii_lowercase, digits
+
+
+_EMPTY_FEN = "9/9/9/9/9/9/9/9/9/9 w---1"
+
+
+def global_names():
+    global board
+    board = Board.from_fen(_EMPTY_FEN)
+    for col in ascii_lowercase[:9]:
+        for row in digits:
+            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
+
+
+def set_board(positions: List[str]):
+    for position in positions:
+        board.board[position].reset(
+            Piece(position, SquareState.OCCUPIED, Type.ROOK, Color.RED)
+        )
+        alpha.Flip()(board.board[position])
 
 
 def test_move_eq():
@@ -98,7 +138,7 @@ def test_to_str():
     assert move1.to_str(0) == ""
     assert move1.to_str(1) == "a0c1^a6"
     assert move1.to_str(2) == "a0c1^a6:MERGE_JUMP:CAPTURE"
-    assert move1.to_str(3) == "a0c1^a6:MERGE_JUMP:CAPTURE:BLACK_ROOK->RED_PAWN"
+    assert move1.to_str(3) == "a0c1^a6:MERGE_JUMP:CAPTURE:RED_ROOK->BLACK_PAWN"
 
     move2 = Move(
         "a0",
@@ -111,7 +151,7 @@ def test_to_str():
     assert move2.to_str(0) == ""
     assert move2.to_str(1) == "a0^b3c1"
     assert move2.to_str(2) == "a0^b3c1:SPLIT_JUMP:BASIC"
-    assert move2.to_str(3) == "a0^b3c1:SPLIT_JUMP:BASIC:BLACK_ROOK->NA_EMPTY"
+    assert move2.to_str(3) == "a0^b3c1:SPLIT_JUMP:BASIC:RED_ROOK->NA_EMPTY"
 
     move3 = Move(
         "a0", "a6", board, move_type=MoveType.SLIDE, move_variant=MoveVariant.CAPTURE
@@ -119,4 +159,37 @@ def test_to_str():
     assert move3.to_str(0) == ""
     assert move3.to_str(1) == "a0a6"
     assert move3.to_str(2) == "a0a6:SLIDE:CAPTURE"
-    assert move3.to_str(3) == "a0a6:SLIDE:CAPTURE:BLACK_ROOK->RED_PAWN"
+    assert move3.to_str(3) == "a0a6:SLIDE:CAPTURE:RED_ROOK->BLACK_PAWN"
+
+
+def test_jump_classical():
+    global_names()
+
+    # basic case
+    set_board(["a1", "a3"])
+    Jump(MoveVariant.CLASSICAL)(a1, b2)
+    assert_samples_in(board, [locations_to_bitboard(["b2", "a3"])])
+
+    # capture case
+    Jump(MoveVariant.CLASSICAL)(b2, a3)
+    assert_samples_in(board, [locations_to_bitboard(["a3"])])
+
+
+def test_jump_capture():
+    global_names()
+    set_board(["a1", "a3"])
+    alpha.PhasedSplit()(a1, b1, b2)
+    board_probabilities = get_board_probability_distribution(board, 5000)
+    assert len(board_probabilities) == 2
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b1", "a3"]))
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "a3"]))
+
+    Jump(MoveVariant.CAPTURE)(b1, a3)
+    samples = sample_board(board, 100)
+    assert_this_or_that(
+        samples, locations_to_bitboard(["a3"]), locations_to_bitboard(["b2", "a3"])
+    )
+    board_probabilities = get_board_probability_distribution(board, 5000)
+    assert len(board_probabilities) == 2
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3"]))
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "a3"]))
diff --git a/unitary/examples/quantum_chinese_chess/test_utils.py b/unitary/examples/quantum_chinese_chess/test_utils.py
index 98184187..d7f74f94 100644
--- a/unitary/examples/quantum_chinese_chess/test_utils.py
+++ b/unitary/examples/quantum_chinese_chess/test_utils.py
@@ -13,8 +13,11 @@
 # limitations under the License.
 from unitary.alpha import QuantumObject, QuantumWorld
 from unitary.examples.quantum_chinese_chess.enums import SquareState
+from unitary.examples.quantum_chinese_chess.board import Board
 from string import ascii_lowercase, digits
-
+from typing import List, Dict
+from collections import defaultdict
+from scipy.stats import chisquare
 
 # Build quantum objects a0 to i9, and add them to a quantum world.
 def init_board() -> QuantumWorld:
@@ -23,3 +26,139 @@ def init_board() -> QuantumWorld:
         for row in digits:
             board[col + row] = QuantumObject(col + row, SquareState.EMPTY)
     return QuantumWorld(list(board.values()))
+
+
+def location_to_bit(location: str) -> int:
+    """Transform location notation (e.g. "a3") into a bitboard bit number."""
+    x = ord(location[0]) - ord("a")
+    y = int(location[1])
+    return y * 9 + x
+
+
+def locations_to_bitboard(locations: List[str]) -> int:
+    """Transform a list of locations into a 90-bit board bitstring."""
+    bitboard = 0
+    for location in locations:
+        bitboard += 1 << location_to_bit(location)
+    return bitboard
+
+
+def nth_bit_of(n: int, bit_board: int) -> bool:
+    """Returns the n-th bit of a 90-bit bitstring."""
+    return (bit_board >> n) % 2 == 1
+
+
+def bit_to_location(bit: int) -> str:
+    """Transform a bitboard bit number into a location (e.g. "a3")."""
+    y = bit // 9
+    x = chr(bit % 9 + ord("a"))
+    return f"{x}{y}"
+
+
+def bitboard_to_locations(bitboard: int) -> List[str]:
+    """Transform a 90-bit bitstring into a list of locations."""
+    locations = []
+    for n in range(90):
+        if nth_bit_of(n, bitboard):
+            locations.append(bit_to_location(n))
+    return locations
+
+def sample_board(board: Board, repetitions: int) -> List[int]:
+    samples = board.board.peek(count = repetitions, convert_to_enum = False)
+    # Convert peek results (in List[List[int]]) into bitstring.
+    samples = [int("0b" + "".join([str(i) for i in sample[::-1]]), base=2) for sample in samples]
+    return samples
+
+
+def print_samples(samples):
+    """Prints all the samples as lists of locations."""
+    sample_dict = {}
+    for sample in samples:
+        if sample not in sample_dict:
+            sample_dict[sample] = 0
+        sample_dict[sample] += 1
+    for key in sample_dict:
+        print(f"{bitboard_to_locations(key)}: {sample_dict[key]}")
+
+
+def get_board_probability_distribution(board: Board, repetitions: int = 1000) -> Dict[int, float]:
+    """Returns the probability distribution for each board found in the sample.
+
+    The values are returned as a dict{bitboard(int): probability(float)}.
+    """
+    board_probabilities: Dict[int, float] = {}
+
+    samples = board.board.peek(count = repetitions, convert_to_enum = False)
+    # Convert peek results (in List[List[int]]) into bitstring.
+    samples = [int("0b" + "".join([str(i) for i in sample[::-1]]), base=2) for sample in samples]
+    for sample in samples:
+        if sample not in board_probabilities:
+            board_probabilities[sample] = 0.0
+        board_probabilities[sample] += 1.0
+
+    for board in board_probabilities:
+        board_probabilities[board] /= repetitions
+
+    return board_probabilities
+
+
+def assert_samples_in(board: Board, possibilities):
+    samples = sample_board(board, 500)
+    assert len(samples) == 500
+    all_in = all(sample in possibilities for sample in samples)
+    print(possibilities)
+    print(set(samples))
+    assert all_in, print_samples(samples)
+    # Make sure each possibility is represented at least once.
+    for possibility in possibilities:
+        any_in = any(sample == possibility for sample in samples)
+        assert any_in, print_samples(samples)
+
+
+def assert_sample_distribution(board: Board, probability_map, p_significant=1e-6):
+    """Performs a chi-squared test that samples follow an expected distribution.
+
+    probability_map is a map from bitboards to expected probability. An
+    assertion is raised if one of the samples is not in the map, or if the
+    probability that the samples are at least as different from the expected
+    ones as the observed sampless is less than p_significant.
+    """
+    assert abs(sum(probability_map.values()) - 1) < 1e-9
+    samples = sample_board(board, 500)
+    assert len(samples) == 500
+    counts = defaultdict(int)
+    for sample in samples:
+        assert sample in probability_map, bitboard_to_locations(sample)
+        counts[sample] += 1
+    observed = []
+    expected = []
+    for position, probability in probability_map.items():
+        observed.append(counts[position])
+        expected.append(500 * probability)
+    p = chisquare(observed, expected).pvalue
+    assert (
+        p > p_significant
+    ), f"Observed {observed} far from expected {expected} (p = {p})"
+
+
+def assert_this_or_that(samples, this, that):
+    """Asserts all the samples are either equal to this or that,
+    and that one of each exists in the samples.
+    """
+    assert any(sample == this for sample in samples), print_samples(samples)
+    assert any(sample == that for sample in samples), print_samples(samples)
+    assert all(sample == this or sample == that for sample in samples), print_samples(
+        samples
+    )
+
+
+def assert_prob_about(probs, that, expected, atol=0.04):
+    """Checks that the probability is within atol of the expected value."""
+    assert probs[that] > expected - atol, print_samples([that])
+    assert probs[that] < expected + atol, print_samples([that])
+
+
+def assert_fifty_fifty(probs, that):
+    """Checks that the probability is close to 50%."""
+    assert_prob_about(probs, that, 0.5), print_samples([that])
+

From 46348d8d3c386853aae401e5654903f06d8f0644 Mon Sep 17 00:00:00 2001
From: madcpf <madcpf@gmail.com>
Date: Wed, 18 Oct 2023 00:02:02 -0700
Subject: [PATCH 04/13] update

---
 .../examples/quantum_chinese_chess/move.py    |  3 ---
 .../quantum_chinese_chess/move_test.py        |  6 ++---
 .../quantum_chinese_chess/test_utils.py       | 25 +++++++++++++------
 3 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/unitary/examples/quantum_chinese_chess/move.py b/unitary/examples/quantum_chinese_chess/move.py
index 91da5d40..7bbcab43 100644
--- a/unitary/examples/quantum_chinese_chess/move.py
+++ b/unitary/examples/quantum_chinese_chess/move.py
@@ -137,13 +137,11 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
         world = source_0.world
         if self.move_variant == MoveVariant.CAPTURE:
             # We peek and force measure source_0.
-            print("Here 1 ")
             source_is_occupied = world.pop([source_0])[0].value
             print(source_is_occupied)
             # For move_variant==CAPTURE, we require source_0 to be occupied before further actions.
             # This is to prevent a piece of the board containing two types of different pieces.
             if not source_is_occupied:
-                print("Here 2 ")
                 # If source_0 turns out to be not there, we clear set it to be EMPTY, and the jump
                 # could not be made.
                 source_0.reset()
@@ -174,7 +172,6 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
                 target_0.reset()
 
         # Make the jump move.
-        print("Here 3 ")        
         alpha.PhasedMove()(source_0, target_0)
         # Move the classical properties of the source piece to the target piece.
         target_0.reset(source_0)
diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index 5232fcf4..7f12a9f0 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -185,11 +185,9 @@ def test_jump_capture():
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "a3"]))
 
     Jump(MoveVariant.CAPTURE)(b1, a3)
+    # pop() will break the supersition and only one of the following two states are possible.
     samples = sample_board(board, 100)
+    assert len(set(samples)) == 1
     assert_this_or_that(
         samples, locations_to_bitboard(["a3"]), locations_to_bitboard(["b2", "a3"])
     )
-    board_probabilities = get_board_probability_distribution(board, 5000)
-    assert len(board_probabilities) == 2
-    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3"]))
-    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "a3"]))
diff --git a/unitary/examples/quantum_chinese_chess/test_utils.py b/unitary/examples/quantum_chinese_chess/test_utils.py
index d7f74f94..f676c46e 100644
--- a/unitary/examples/quantum_chinese_chess/test_utils.py
+++ b/unitary/examples/quantum_chinese_chess/test_utils.py
@@ -19,6 +19,7 @@
 from collections import defaultdict
 from scipy.stats import chisquare
 
+
 # Build quantum objects a0 to i9, and add them to a quantum world.
 def init_board() -> QuantumWorld:
     board = {}
@@ -63,10 +64,14 @@ def bitboard_to_locations(bitboard: int) -> List[str]:
             locations.append(bit_to_location(n))
     return locations
 
+
 def sample_board(board: Board, repetitions: int) -> List[int]:
-    samples = board.board.peek(count = repetitions, convert_to_enum = False)
+    samples = board.board.peek(count=repetitions, convert_to_enum=False)
     # Convert peek results (in List[List[int]]) into bitstring.
-    samples = [int("0b" + "".join([str(i) for i in sample[::-1]]), base=2) for sample in samples]
+    samples = [
+        int("0b" + "".join([str(i) for i in sample[::-1]]), base=2)
+        for sample in samples
+    ]
     return samples
 
 
@@ -81,16 +86,21 @@ def print_samples(samples):
         print(f"{bitboard_to_locations(key)}: {sample_dict[key]}")
 
 
-def get_board_probability_distribution(board: Board, repetitions: int = 1000) -> Dict[int, float]:
+def get_board_probability_distribution(
+    board: Board, repetitions: int = 1000
+) -> Dict[int, float]:
     """Returns the probability distribution for each board found in the sample.
 
     The values are returned as a dict{bitboard(int): probability(float)}.
     """
     board_probabilities: Dict[int, float] = {}
 
-    samples = board.board.peek(count = repetitions, convert_to_enum = False)
+    samples = board.board.peek(count=repetitions, convert_to_enum=False)
     # Convert peek results (in List[List[int]]) into bitstring.
-    samples = [int("0b" + "".join([str(i) for i in sample[::-1]]), base=2) for sample in samples]
+    samples = [
+        int("0b" + "".join([str(i) for i in sample[::-1]]), base=2)
+        for sample in samples
+    ]
     for sample in samples:
         if sample not in board_probabilities:
             board_probabilities[sample] = 0.0
@@ -145,8 +155,8 @@ def assert_this_or_that(samples, this, that):
     """Asserts all the samples are either equal to this or that,
     and that one of each exists in the samples.
     """
-    assert any(sample == this for sample in samples), print_samples(samples)
-    assert any(sample == that for sample in samples), print_samples(samples)
+    # assert any(sample == this for sample in samples), print_samples(samples)
+    # assert any(sample == that for sample in samples), print_samples(samples)
     assert all(sample == this or sample == that for sample in samples), print_samples(
         samples
     )
@@ -161,4 +171,3 @@ def assert_prob_about(probs, that, expected, atol=0.04):
 def assert_fifty_fifty(probs, that):
     """Checks that the probability is close to 50%."""
     assert_prob_about(probs, that, 0.5), print_samples([that])
-

From bcdb75a5ee224821b94332ffec157477a0a1b0e7 Mon Sep 17 00:00:00 2001
From: madcpf <madcpf@gmail.com>
Date: Wed, 18 Oct 2023 10:19:42 -0700
Subject: [PATCH 05/13] test

---
 .../examples/quantum_chinese_chess/move.py    |  4 +-
 .../quantum_chinese_chess/move_test.py        | 61 ++++++++++++++++++-
 .../quantum_chinese_chess/test_utils.py       |  7 ++-
 3 files changed, 66 insertions(+), 6 deletions(-)

diff --git a/unitary/examples/quantum_chinese_chess/move.py b/unitary/examples/quantum_chinese_chess/move.py
index 7bbcab43..36e491fc 100644
--- a/unitary/examples/quantum_chinese_chess/move.py
+++ b/unitary/examples/quantum_chinese_chess/move.py
@@ -145,7 +145,7 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
                 # If source_0 turns out to be not there, we clear set it to be EMPTY, and the jump
                 # could not be made.
                 source_0.reset()
-                print("Jump move: source turns out to be empty.")
+                print("Jump move not applied: source turns out to be empty.")
                 return iter(())
             source_0.is_entangled = False
             # We replace the qubit of target_0 with a new ancilla, and set its classical properties to be EMPTY.
@@ -159,7 +159,7 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
             if target_is_occupied:
                 # If target_0 turns out to be there, we set it to be a classically OCCUPIED, and
                 # the jump could not be made.
-                print("Jump move: target turns out to be occupied.")
+                print("Jump move not applied: target turns out to be occupied.")
                 target_0.is_entangled = False
                 return iter(())
             # Otherwise we set target_0 to be classically EMPTY.
diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index 7f12a9f0..ba8a2174 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -176,10 +176,11 @@ def test_jump_classical():
 
 
 def test_jump_capture():
+    # Source is in quantum state
     global_names()
     set_board(["a1", "a3"])
     alpha.PhasedSplit()(a1, b1, b2)
-    board_probabilities = get_board_probability_distribution(board, 5000)
+    board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b1", "a3"]))
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "a3"]))
@@ -191,3 +192,61 @@ def test_jump_capture():
     assert_this_or_that(
         samples, locations_to_bitboard(["a3"]), locations_to_bitboard(["b2", "a3"])
     )
+
+    # Target is in quantum state
+    global_names()
+    set_board(["a1", "a3"])
+    alpha.PhasedSplit()(a1, b1, b2)
+    Jump(MoveVariant.CAPTURE)(a3, b1)
+    board_probabilities = get_board_probability_distribution(board, 1000)
+    assert len(board_probabilities) == 2
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b1"]))
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b1", "b2"]))
+
+    # Both source and target are in quantum state
+    global_names()
+    set_board(["a1", "b1"])
+    alpha.PhasedSplit()(a1, a2, a3)
+    alpha.PhasedSplit()(b1, b2, b3)
+    assert_sample_distribution(
+        board,
+        {
+            locations_to_bitboard(["a2", "b2"]): 1 / 4.0,
+            locations_to_bitboard(["a2", "b3"]): 1 / 4.0,
+            locations_to_bitboard(["a3", "b2"]): 1 / 4.0,
+            locations_to_bitboard(["a3", "b3"]): 1 / 4.0,
+        },
+    )
+    Jump(MoveVariant.CAPTURE)(a2, b2)
+    board_probabilities = get_board_probability_distribution(board, 1000)
+    assert len(board_probabilities) == 2
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b2"]))
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b3"]))
+
+
+def test_jump_excluded():
+    global_names()
+    set_board(["a1", "a3"])
+    alpha.PhasedSplit()(a1, b1, b2)
+
+    Jump(MoveVariant.EXCLUDED)(a3, b1)
+    # pop() will break the supersition and only one of the following two states are possible.
+    samples = sample_board(board, 100)
+    assert len(set(samples)) == 1
+    assert_this_or_that(
+        samples,
+        locations_to_bitboard(["a3", "b1"]),
+        locations_to_bitboard(["b1", "b2"]),
+    )
+
+
+def test_jump_basic():
+    global_names()
+    set_board(["a1"])
+    alpha.PhasedSplit()(a1, b1, b2)
+
+    Jump(MoveVariant.BASIC)(b1, d1)
+    board_probabilities = get_board_probability_distribution(board, 1000)
+    assert len(board_probabilities) == 2
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2"]))
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["d1"]))
diff --git a/unitary/examples/quantum_chinese_chess/test_utils.py b/unitary/examples/quantum_chinese_chess/test_utils.py
index f676c46e..03e00685 100644
--- a/unitary/examples/quantum_chinese_chess/test_utils.py
+++ b/unitary/examples/quantum_chinese_chess/test_utils.py
@@ -164,10 +164,11 @@ def assert_this_or_that(samples, this, that):
 
 def assert_prob_about(probs, that, expected, atol=0.04):
     """Checks that the probability is within atol of the expected value."""
-    assert probs[that] > expected - atol, print_samples([that])
-    assert probs[that] < expected + atol, print_samples([that])
+    assert that in probs, print_samples(list(probs.keys()))
+    assert probs[that] > expected - atol, print_samples(list(probs.keys()))
+    assert probs[that] < expected + atol, print_samples(list(probs.keys()))
 
 
 def assert_fifty_fifty(probs, that):
     """Checks that the probability is close to 50%."""
-    assert_prob_about(probs, that, 0.5), print_samples([that])
+    assert_prob_about(probs, that, 0.5), print_samples(list(probs.keys()))

From cf35e05cf9bab7acc42b64544468de06212fc7bc Mon Sep 17 00:00:00 2001
From: Pengfei Chen <pengfeichen@google.com>
Date: Wed, 18 Oct 2023 11:13:04 -0700
Subject: [PATCH 06/13] update test

---
 .../examples/quantum_chinese_chess/move.py    |   1 -
 .../quantum_chinese_chess/move_test.py        | 111 +++++++++++-------
 .../quantum_chinese_chess/test_utils.py       |   2 +-
 3 files changed, 69 insertions(+), 45 deletions(-)

diff --git a/unitary/examples/quantum_chinese_chess/move.py b/unitary/examples/quantum_chinese_chess/move.py
index 36e491fc..8b34139a 100644
--- a/unitary/examples/quantum_chinese_chess/move.py
+++ b/unitary/examples/quantum_chinese_chess/move.py
@@ -138,7 +138,6 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
         if self.move_variant == MoveVariant.CAPTURE:
             # We peek and force measure source_0.
             source_is_occupied = world.pop([source_0])[0].value
-            print(source_is_occupied)
             # For move_variant==CAPTURE, we require source_0 to be occupied before further actions.
             # This is to prevent a piece of the board containing two types of different pieces.
             if not source_is_occupied:
diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index ba8a2174..c86f13b8 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -165,45 +165,45 @@ def test_to_str():
 def test_jump_classical():
     global_names()
 
-    # basic case
-    set_board(["a1", "a3"])
+    # Target is empty.
+    set_board(["a1", "b1"])
     Jump(MoveVariant.CLASSICAL)(a1, b2)
-    assert_samples_in(board, [locations_to_bitboard(["b2", "a3"])])
+    assert_samples_in(board, [locations_to_bitboard(["b2", "b1"])])
 
-    # capture case
-    Jump(MoveVariant.CLASSICAL)(b2, a3)
-    assert_samples_in(board, [locations_to_bitboard(["a3"])])
+    # Target is occupied.
+    Jump(MoveVariant.CLASSICAL)(b2, b1)
+    assert_samples_in(board, [locations_to_bitboard(["b1"])])
 
 
 def test_jump_capture():
-    # Source is in quantum state
+    # Source is in quantum state.
     global_names()
-    set_board(["a1", "a3"])
-    alpha.PhasedSplit()(a1, b1, b2)
+    set_board(["a1", "b1"])
+    alpha.PhasedSplit()(a1, a2, a3)
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
-    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b1", "a3"]))
-    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "a3"]))
-
-    Jump(MoveVariant.CAPTURE)(b1, a3)
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a2", "b1"]))
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b1"]))
+    Jump(MoveVariant.CAPTURE)(a2, b1)
     # pop() will break the supersition and only one of the following two states are possible.
-    samples = sample_board(board, 100)
-    assert len(set(samples)) == 1
-    assert_this_or_that(
-        samples, locations_to_bitboard(["a3"]), locations_to_bitboard(["b2", "a3"])
-    )
-
-    # Target is in quantum state
+    # We check the ancilla to learn if the jump was applied or not.
+    source_is_occupied = board.board.post_selection[board.board["ancilla_a2_0"]]
+    if source_is_occupied:
+        assert_samples_in(board, [locations_to_bitboard(["b1"])])
+    else:
+        assert_samples_in(board, [locations_to_bitboard(["a3", "b1"])])
+
+    # Target is in quantum state.
     global_names()
-    set_board(["a1", "a3"])
-    alpha.PhasedSplit()(a1, b1, b2)
-    Jump(MoveVariant.CAPTURE)(a3, b1)
+    set_board(["a1", "b1"])
+    alpha.PhasedSplit()(b1, b2, b3)
+    Jump(MoveVariant.CAPTURE)(a1, b2)
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
-    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b1"]))
-    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b1", "b2"]))
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2"]))
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "b3"]))
 
-    # Both source and target are in quantum state
+    # Both source and target are in quantum state.
     global_names()
     set_board(["a1", "b1"])
     alpha.PhasedSplit()(a1, a2, a3)
@@ -220,33 +220,58 @@ def test_jump_capture():
     Jump(MoveVariant.CAPTURE)(a2, b2)
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
-    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b2"]))
-    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b3"]))
+    # We check the ancilla to learn if the jump was applied or not.
+    source_is_occupied = board.board.post_selection[board.board["ancilla_a2_0"]]
+    print(source_is_occupied)
+    if source_is_occupied:
+        assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2"]))
+        assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "b3"]))
+    else:
+        assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b2"]))
+        assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b3"]))
 
 
 def test_jump_excluded():
+    # Target is in quantum state.
     global_names()
-    set_board(["a1", "a3"])
-    alpha.PhasedSplit()(a1, b1, b2)
-
-    Jump(MoveVariant.EXCLUDED)(a3, b1)
+    set_board(["a1", "b1"])
+    alpha.PhasedSplit()(b1, b2, b3)
+    Jump(MoveVariant.EXCLUDED)(a1, b2)
     # pop() will break the supersition and only one of the following two states are possible.
-    samples = sample_board(board, 100)
-    assert len(set(samples)) == 1
-    assert_this_or_that(
-        samples,
-        locations_to_bitboard(["a3", "b1"]),
-        locations_to_bitboard(["b1", "b2"]),
-    )
+    # We check the ancilla to learn if the jump was applied or not.
+    target_is_occupied = board.board.post_selection[board.board["ancilla_b2_0"]]
+    print(target_is_occupied)
+    if target_is_occupied:
+        assert_samples_in(board, [locations_to_bitboard(["a1", "b2"])])
+    else:
+        assert_samples_in(board, [locations_to_bitboard(["b2", "b3"])])
+
+    # Both source and target are in quantum state.
+    global_names()
+    set_board(["a1", "b1"])
+    alpha.PhasedSplit()(a1, a2, a3)
+    alpha.PhasedSplit()(b1, b2, b3)
+    Jump(MoveVariant.EXCLUDED)(a2, b2)
+    board_probabilities = get_board_probability_distribution(board, 1000)
+    assert len(board_probabilities) == 2
+    # We check the ancilla to learn if the jump was applied or not.
+    target_is_occupied = board.board.post_selection[board.board["ancilla_b2_0"]]
+    print(target_is_occupied)
+    if target_is_occupied:
+        assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a2", "b2"]))
+        assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b2"]))
+    else:
+        assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "b3"]))
+        assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b3"]))
 
 
 def test_jump_basic():
+    # Souce is in quantum state.
     global_names()
     set_board(["a1"])
-    alpha.PhasedSplit()(a1, b1, b2)
-
-    Jump(MoveVariant.BASIC)(b1, d1)
+    alpha.PhasedSplit()(a1, a2, a3)
+    Jump(MoveVariant.BASIC)(a2, d1)
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
-    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2"]))
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["d1"]))
+    assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3"]))
diff --git a/unitary/examples/quantum_chinese_chess/test_utils.py b/unitary/examples/quantum_chinese_chess/test_utils.py
index 03e00685..c24761ce 100644
--- a/unitary/examples/quantum_chinese_chess/test_utils.py
+++ b/unitary/examples/quantum_chinese_chess/test_utils.py
@@ -162,7 +162,7 @@ def assert_this_or_that(samples, this, that):
     )
 
 
-def assert_prob_about(probs, that, expected, atol=0.04):
+def assert_prob_about(probs, that, expected, atol=0.05):
     """Checks that the probability is within atol of the expected value."""
     assert that in probs, print_samples(list(probs.keys()))
     assert probs[that] > expected - atol, print_samples(list(probs.keys()))

From 640b6857529c5c643e0b9659f05811d2b059ef8a Mon Sep 17 00:00:00 2001
From: Pengfei Chen <pengfeichen@google.com>
Date: Wed, 18 Oct 2023 11:59:55 -0700
Subject: [PATCH 07/13] update

---
 unitary/alpha/quantum_world_test.py           |  22 +++-
 .../examples/quantum_chinese_chess/chess.py   |  11 +-
 .../quantum_chinese_chess/chess_test.py       | 114 +++++++++++-------
 .../examples/quantum_chinese_chess/move.py    |   6 +-
 .../quantum_chinese_chess/test_utils.py       |   9 --
 5 files changed, 96 insertions(+), 66 deletions(-)

diff --git a/unitary/alpha/quantum_world_test.py b/unitary/alpha/quantum_world_test.py
index b656fcfb..75dc9b56 100644
--- a/unitary/alpha/quantum_world_test.py
+++ b/unitary/alpha/quantum_world_test.py
@@ -176,7 +176,7 @@ def test_pop(simulator, compile_to_qubits):
         compile_to_qubits=compile_to_qubits,
     )
     alpha.Split()(light, light2, light3)
-    results = board.peek([light2, light3], count=200)
+    results = board.peek([light2, light3], count=200, convert_to_enum=False)
     assert all(result[0] != result[1] for result in results)
     assert not all(result[0] == 0 for result in results)
     assert not all(result[0] == 1 for result in results)
@@ -189,6 +189,26 @@ def test_pop(simulator, compile_to_qubits):
     assert all(result[1] != popped for result in results)
 
 
+@pytest.mark.parametrize("compile_to_qubits", [False, True])
+@pytest.mark.parametrize("simulator", [cirq.Simulator, alpha.SparseSimulator])
+def test_unhook(simulator, compile_to_qubits):
+    light = alpha.QuantumObject("l1", Light.GREEN)
+    light2 = alpha.QuantumObject("l2", Light.RED)
+    light3 = alpha.QuantumObject("l3", Light.RED)
+    board = alpha.QuantumWorld(
+        [light, light2, light3],
+        sampler=simulator(),
+        compile_to_qubits=compile_to_qubits,
+    )
+    alpha.Split()(light, light2, light3)
+    board.unhook(light2)
+    results = board.peek([light2, light3], count=200, convert_to_enum=False)
+    print(results)
+    assert all(result[0] == 0 for result in results)
+    assert not all(result[1] == 0 for result in results)
+    assert not all(result[1] == 1 for result in results)
+
+
 # TODO: Consider moving to qudit_effects.py if this can be broadly useful.
 class QuditSwapEffect(alpha.QuantumEffect):
     def __init__(self, dimension):
diff --git a/unitary/examples/quantum_chinese_chess/chess.py b/unitary/examples/quantum_chinese_chess/chess.py
index 976f68d7..27aa8bc6 100644
--- a/unitary/examples/quantum_chinese_chess/chess.py
+++ b/unitary/examples/quantum_chinese_chess/chess.py
@@ -424,11 +424,9 @@ def next_move(self) -> Tuple[bool, str]:
             output = "Undo last quantum effect."
             # Right now it's only for debugging purposes, since it has following problems:
             # TODO(): there are several problems here:
-            # 1) last move is quantum but classical piece information is not reversed back.
+            # 1) the classical piece information is not reversed back.
             # ==> we may need to save the change of classical piece information of each step.
-            # 2) last move might not be quantum.
-            # ==> we may need to save all classical moves and figure out how to undo each kind of move;
-            # 3) last move is quantum but involved multiple effects.
+            # 2) last move involved multiple effects.
             # ==> we may need to save number of effects per move, and undo that number of times.
             self.board.board.undo_last_effect()
             return True, output
@@ -458,10 +456,7 @@ def update_board_by_sampling(self) -> List[float]:
         for row in range(num_rows):
             for col in "abcdefghi":
                 piece = self.board.board[f"{col}{row}"]
-                # We need to do the following range() conversion since the sequence of
-                # qubits returned from get_binary_probabilities() is
-                # a9 b9 ... i9, a8 b8 ... i8, ..., a0 b0 ... i0
-                prob = probs[(num_rows - row - 1) * num_cols + ord(col) - ord("a")]
+                prob = probs[row * num_cols + ord(col) - ord("a")]
                 # TODO(): This threshold does not actually work right now since we have 100 sampling.
                 # Change it to be more meaningful values maybe when we do error mitigation.
                 if prob < 1e-3:
diff --git a/unitary/examples/quantum_chinese_chess/chess_test.py b/unitary/examples/quantum_chinese_chess/chess_test.py
index bd00ec65..b278fe35 100644
--- a/unitary/examples/quantum_chinese_chess/chess_test.py
+++ b/unitary/examples/quantum_chinese_chess/chess_test.py
@@ -104,6 +104,7 @@ def test_check_classical_rule(monkeypatch):
     inputs = iter(["y", "Bob", "Ben"])
     monkeypatch.setattr("builtins.input", lambda _: next(inputs))
     game = QuantumChineseChess()
+    board = game.board.board
     # The move is blocked by classical path piece.
     with pytest.raises(ValueError, match="The path is blocked."):
         game.check_classical_rule("a0", "a4", ["a3"])
@@ -128,14 +129,10 @@ def test_check_classical_rule(monkeypatch):
     game.check_classical_rule("c0", "e2", [])
     with pytest.raises(ValueError, match="ELEPHANT cannot move like this."):
         game.check_classical_rule("c0", "e1", [])
-    game.board.board["g5"].reset(
-        Piece("g5", SquareState.OCCUPIED, Type.ELEPHANT, Color.BLACK)
-    )
+    board["g5"].reset(Piece("g5", SquareState.OCCUPIED, Type.ELEPHANT, Color.BLACK))
     with pytest.raises(ValueError, match="ELEPHANT cannot cross the river"):
         game.check_classical_rule("g5", "i3", [])
-    game.board.board["c4"].reset(
-        Piece("c4", SquareState.OCCUPIED, Type.ELEPHANT, Color.RED)
-    )
+    board["c4"].reset(Piece("c4", SquareState.OCCUPIED, Type.ELEPHANT, Color.RED))
     with pytest.raises(ValueError, match="ELEPHANT cannot cross the river"):
         game.check_classical_rule("c4", "e6", [])
 
@@ -152,9 +149,9 @@ def test_check_classical_rule(monkeypatch):
     game.check_classical_rule("e9", "e8", [])
     with pytest.raises(ValueError, match="KING cannot move like this."):
         game.check_classical_rule("e9", "d8", [])
-    game.board.board["c0"].reset()
-    game.board.board["d0"].reset(game.board.board["e0"])
-    game.board.board["e0"].reset()
+    board["c0"].reset()
+    board["d0"].reset(board["e0"])
+    board["e0"].reset()
     with pytest.raises(ValueError, match="KING cannot leave the palace."):
         game.check_classical_rule("d0", "c0", [])
 
@@ -167,9 +164,9 @@ def test_check_classical_rule(monkeypatch):
     with pytest.raises(ValueError, match="CANNON cannot fire like this."):
         game.check_classical_rule("b2", "b9", ["b5", "b7"])
     # Cannon cannot fire to a piece with same color.
-    game.board.board["b3"].reset(game.board.board["b2"])
-    game.board.board["b2"].reset()
-    game.board.board["e3"].is_entangled = True
+    board["b3"].reset(board["b2"])
+    board["b2"].reset()
+    board["e3"].is_entangled = True
     with pytest.raises(
         ValueError, match="CANNON cannot fire to a piece with same color."
     ):
@@ -194,8 +191,8 @@ def test_check_classical_rule(monkeypatch):
     with pytest.raises(ValueError, match="PAWN can not move backward."):
         game.check_classical_rule("g6", "g7", [])
     # After crossing the rive the pawn could move horizontally.
-    game.board.board["c4"].reset(game.board.board["c6"])
-    game.board.board["c6"].reset()
+    board["c4"].reset(board["c6"])
+    board["c6"].reset()
     game.check_classical_rule("c4", "b4", [])
     game.check_classical_rule("c4", "d4", [])
 
@@ -206,6 +203,7 @@ def test_classify_move_fail(monkeypatch):
     inputs = iter(["y", "Bob", "Ben"])
     monkeypatch.setattr("builtins.input", lambda _: next(inputs))
     game = QuantumChineseChess()
+    board = game.board.board
     with pytest.raises(
         ValueError, match="CANNON could not fire/capture without a cannon platform."
     ):
@@ -216,21 +214,21 @@ def test_classify_move_fail(monkeypatch):
     ):
         game.classify_move(["b2", "h2"], ["e2"], [], [], [], [])
 
-    game.board.board["c0"].reset(game.board.board["b7"])
-    game.board.board["c0"].is_entangled = True
-    game.board.board["b7"].reset()
-    game.board.board["g0"].reset(game.board.board["h7"])
-    game.board.board["g0"].is_entangled = True
-    game.board.board["h7"].reset()
+    board["c0"].reset(board["b7"])
+    board["c0"].is_entangled = True
+    board["b7"].reset()
+    board["g0"].reset(board["h7"])
+    board["g0"].is_entangled = True
+    board["h7"].reset()
     with pytest.raises(ValueError, match="Currently CANNON cannot merge while firing."):
         game.classify_move(["c0", "g0"], ["e0"], ["d0"], [], ["f0"], [])
 
-    game.board.board["b3"].reset(game.board.board["b2"])
-    game.board.board["b3"].is_entangled = True
-    game.board.board["b2"].reset()
-    game.board.board["d3"].reset(game.board.board["h2"])
-    game.board.board["d3"].is_entangled = True
-    game.board.board["h2"].reset()
+    board["b3"].reset(board["b2"])
+    board["b3"].is_entangled = True
+    board["b2"].reset()
+    board["d3"].reset(board["h2"])
+    board["d3"].is_entangled = True
+    board["h2"].reset()
     with pytest.raises(
         ValueError, match="Currently we could only merge into an empty piece."
     ):
@@ -239,14 +237,14 @@ def test_classify_move_fail(monkeypatch):
     with pytest.raises(ValueError, match="Currently CANNON cannot split while firing."):
         game.classify_move(["g0"], ["e0", "i0"], ["f0"], [], ["h0"], [])
 
-    game.board.board["d0"].is_entangled = True
+    board["d0"].is_entangled = True
     with pytest.raises(
         ValueError, match="Currently we could only split into empty pieces."
     ):
         game.classify_move(["d3"], ["d0", "d4"], [], [], [], [])
 
-    game.board.board["d0"].reset()
-    game.board.board["f0"].reset()
+    board["d0"].reset()
+    board["f0"].reset()
     with pytest.raises(ValueError, match="King split is not supported currently."):
         game.classify_move(["e0"], ["d0", "f0"], [], [], [], [])
 
@@ -257,6 +255,7 @@ def test_classify_move_success(monkeypatch):
     inputs = iter(["y", "Bob", "Ben"])
     monkeypatch.setattr("builtins.input", lambda _: next(inputs))
     game = QuantumChineseChess()
+    board = game.board.board
     # classical
     assert game.classify_move(["h9"], ["g7"], [], [], [], []) == (
         MoveType.CLASSICAL,
@@ -268,28 +267,28 @@ def test_classify_move_success(monkeypatch):
     )
 
     # jump basic
-    game.board.board["c9"].is_entangled = True
+    board["c9"].is_entangled = True
     assert game.classify_move(["c9"], ["e7"], [], [], [], []) == (
         MoveType.JUMP,
         MoveVariant.BASIC,
     )
-    game.board.board["b2"].is_entangled = True
+    board["b2"].is_entangled = True
     assert game.classify_move(["b2"], ["e2"], [], [], [], []) == (
         MoveType.JUMP,
         MoveVariant.BASIC,
     )
 
     # jump excluded
-    game.board.board["a3"].is_entangled = True
+    board["a3"].is_entangled = True
     assert game.classify_move(["a0"], ["a3"], [], [], [], []) == (
         MoveType.JUMP,
         MoveVariant.EXCLUDED,
     )
 
     # jump capture
-    game.board.board["g4"].reset(game.board.board["g6"])
-    game.board.board["g4"].is_entangled = True
-    game.board.board["g6"].reset()
+    board["g4"].reset(board["g6"])
+    board["g4"].is_entangled = True
+    board["g6"].reset()
     assert game.classify_move(["g4"], ["g3"], [], [], [], []) == (
         MoveType.JUMP,
         MoveVariant.CAPTURE,
@@ -302,10 +301,10 @@ def test_classify_move_success(monkeypatch):
     )
 
     # slide excluded
-    game.board.board["i7"].reset(game.board.board["h7"])
-    game.board.board["i7"].is_entangled = True
-    game.board.board["h7"].reset()
-    game.board.board["i6"].is_entangled = True
+    board["i7"].reset(board["h7"])
+    board["i7"].is_entangled = True
+    board["h7"].reset()
+    board["i6"].is_entangled = True
     assert game.classify_move(["i9"], ["i6"], [], ["i7"], [], []) == (
         MoveType.SLIDE,
         MoveVariant.EXCLUDED,
@@ -324,17 +323,17 @@ def test_classify_move_success(monkeypatch):
     )
 
     # split_slide basic
-    game.board.board["d3"].reset(game.board.board["h2"])
-    game.board.board["h2"].reset()
-    game.board.board["c3"].is_entangled = True
-    game.board.board["e3"].is_entangled = True
+    board["d3"].reset(board["h2"])
+    board["h2"].reset()
+    board["c3"].is_entangled = True
+    board["e3"].is_entangled = True
     assert game.classify_move(["d3"], ["b3", "f3"], [], ["c3"], [], ["e3"]) == (
         MoveType.SPLIT_SLIDE,
         MoveVariant.BASIC,
     )
 
     # merge_jump basic
-    game.board.board["b7"].is_entangled = True
+    board["b7"].is_entangled = True
     assert game.classify_move(["b7", "i7"], ["e7"], [], [], [], []) == (
         MoveType.MERGE_JUMP,
         MoveVariant.BASIC,
@@ -351,8 +350,33 @@ def test_classify_move_success(monkeypatch):
         MoveType.CANNON_FIRE,
         MoveVariant.CAPTURE,
     )
-    game.board.board["i6"].is_entangled = False
+    board["i6"].is_entangled = False
     assert game.classify_move(["i7"], ["i3"], ["i6"], [], [], []) == (
         MoveType.CANNON_FIRE,
         MoveVariant.CAPTURE,
     )
+
+
+def test_update_board_by_sampling(monkeypatch):
+    output = io.StringIO()
+    sys.stdout = output
+    inputs = iter(["y", "Bob", "Ben"])
+    monkeypatch.setattr("builtins.input", lambda _: next(inputs))
+    game = QuantumChineseChess()
+    board = game.board.board
+
+    board.unhook(board["a0"])
+    board["a0"].type_ = Type.ROOK
+    board["a0"].color = Color.RED
+    board["a0"].is_entangled = True
+
+    # Verify that the method would set a0 to classically empty.
+    game.update_board_by_sampling()
+    assert board["a0"].type_ == Type.EMPTY
+    assert board["a0"].color == Color.NA
+    assert board["a0"].is_entangled == False
+
+    board["a1"].is_entangled = True
+    # Verify that the method would set a1 to classically occupied.
+    game.update_board_by_sampling()
+    assert board["a1"].is_entangled == False
diff --git a/unitary/examples/quantum_chinese_chess/move.py b/unitary/examples/quantum_chinese_chess/move.py
index 8b34139a..6f8bddce 100644
--- a/unitary/examples/quantum_chinese_chess/move.py
+++ b/unitary/examples/quantum_chinese_chess/move.py
@@ -129,7 +129,7 @@ def num_objects(self) -> Optional[int]:
         return 2
 
     def effect(self, *objects) -> Iterator[cirq.Operation]:
-        # TODO(): currently pawn capture is a same as jump capture, while in quantum chess it's different,
+        # TODO(): currently pawn capture is the same as jump capture, while in quantum chess it's different,
         # i.e. pawn would move only if the target is there, i.e. CNOT(t, s), and an entanglement could be
         # created. This could be a general game setting, i.e. we could allow players to choose if they
         # want the source piece to move (in case of capture) if the target piece is not there.
@@ -141,7 +141,7 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
             # For move_variant==CAPTURE, we require source_0 to be occupied before further actions.
             # This is to prevent a piece of the board containing two types of different pieces.
             if not source_is_occupied:
-                # If source_0 turns out to be not there, we clear set it to be EMPTY, and the jump
+                # If source_0 turns out to be not there, we set it to be EMPTY, and the jump
                 # could not be made.
                 source_0.reset()
                 print("Jump move not applied: source turns out to be empty.")
@@ -156,7 +156,7 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
             # For move_variant==EXCLUDED, we require target_0 to be empty before further actions.
             # This is to prevent a piece of the board containing two types of different pieces.
             if target_is_occupied:
-                # If target_0 turns out to be there, we set it to be a classically OCCUPIED, and
+                # If target_0 turns out to be there, we set it to be classically OCCUPIED, and
                 # the jump could not be made.
                 print("Jump move not applied: target turns out to be occupied.")
                 target_0.is_entangled = False
diff --git a/unitary/examples/quantum_chinese_chess/test_utils.py b/unitary/examples/quantum_chinese_chess/test_utils.py
index c24761ce..6804558b 100644
--- a/unitary/examples/quantum_chinese_chess/test_utils.py
+++ b/unitary/examples/quantum_chinese_chess/test_utils.py
@@ -20,15 +20,6 @@
 from scipy.stats import chisquare
 
 
-# Build quantum objects a0 to i9, and add them to a quantum world.
-def init_board() -> QuantumWorld:
-    board = {}
-    for col in ascii_lowercase[:9]:
-        for row in digits:
-            board[col + row] = QuantumObject(col + row, SquareState.EMPTY)
-    return QuantumWorld(list(board.values()))
-
-
 def location_to_bit(location: str) -> int:
     """Transform location notation (e.g. "a3") into a bitboard bit number."""
     x = ord(location[0]) - ord("a")

From e8cce5a48304193c27d11ca6bbd2d789af3fc190 Mon Sep 17 00:00:00 2001
From: Pengfei Chen <pengfeichen@google.com>
Date: Wed, 18 Oct 2023 13:09:55 -0700
Subject: [PATCH 08/13] update

---
 unitary/examples/quantum_chinese_chess/move_test.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index c86f13b8..4f1ef918 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -42,14 +42,19 @@
 
 
 def global_names():
+    pass
+
+
+# global board
+# board = Board.from_fen(_EMPTY_FEN)
+
+
+def set_board(positions: List[str]):
     global board
     board = Board.from_fen(_EMPTY_FEN)
     for col in ascii_lowercase[:9]:
         for row in digits:
             globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
-
-
-def set_board(positions: List[str]):
     for position in positions:
         board.board[position].reset(
             Piece(position, SquareState.OCCUPIED, Type.ROOK, Color.RED)

From 576b13ca92cc88dbdcfd26e6d0c26f761a590c84 Mon Sep 17 00:00:00 2001
From: Pengfei Chen <pengfeichen@google.com>
Date: Wed, 18 Oct 2023 14:21:26 -0700
Subject: [PATCH 09/13] update

---
 .../quantum_chinese_chess/move_test.py        |  77 +++++------
 .../quantum_chinese_chess/test_utils.py       | 124 +++++++++++-------
 2 files changed, 111 insertions(+), 90 deletions(-)

diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index 4f1ef918..c375e1bb 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -14,6 +14,10 @@
 from unitary.examples.quantum_chinese_chess.move import Move, Jump
 from unitary.examples.quantum_chinese_chess.board import Board
 from unitary.examples.quantum_chinese_chess.piece import Piece
+import pytest
+from unitary import alpha
+from typing import List
+from string import ascii_lowercase, digits
 from unitary.examples.quantum_chinese_chess.enums import (
     MoveType,
     MoveVariant,
@@ -31,35 +35,8 @@
     sample_board,
     get_board_probability_distribution,
     print_samples,
+    set_board,
 )
-import pytest
-from unitary import alpha
-from typing import List
-from string import ascii_lowercase, digits
-
-
-_EMPTY_FEN = "9/9/9/9/9/9/9/9/9/9 w---1"
-
-
-def global_names():
-    pass
-
-
-# global board
-# board = Board.from_fen(_EMPTY_FEN)
-
-
-def set_board(positions: List[str]):
-    global board
-    board = Board.from_fen(_EMPTY_FEN)
-    for col in ascii_lowercase[:9]:
-        for row in digits:
-            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
-    for position in positions:
-        board.board[position].reset(
-            Piece(position, SquareState.OCCUPIED, Type.ROOK, Color.RED)
-        )
-        alpha.Flip()(board.board[position])
 
 
 def test_move_eq():
@@ -168,10 +145,12 @@ def test_to_str():
 
 
 def test_jump_classical():
-    global_names()
-
     # Target is empty.
-    set_board(["a1", "b1"])
+    board = set_board(["a1", "b1"])
+    # TODO(): try move the following varaibles declarations into a function.
+    for col in ascii_lowercase[:9]:
+        for row in digits:
+            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
     Jump(MoveVariant.CLASSICAL)(a1, b2)
     assert_samples_in(board, [locations_to_bitboard(["b2", "b1"])])
 
@@ -182,8 +161,10 @@ def test_jump_classical():
 
 def test_jump_capture():
     # Source is in quantum state.
-    global_names()
-    set_board(["a1", "b1"])
+    board = set_board(["a1", "b1"])
+    for col in ascii_lowercase[:9]:
+        for row in digits:
+            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
     alpha.PhasedSplit()(a1, a2, a3)
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
@@ -199,8 +180,10 @@ def test_jump_capture():
         assert_samples_in(board, [locations_to_bitboard(["a3", "b1"])])
 
     # Target is in quantum state.
-    global_names()
-    set_board(["a1", "b1"])
+    board = set_board(["a1", "b1"])
+    for col in ascii_lowercase[:9]:
+        for row in digits:
+            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
     alpha.PhasedSplit()(b1, b2, b3)
     Jump(MoveVariant.CAPTURE)(a1, b2)
     board_probabilities = get_board_probability_distribution(board, 1000)
@@ -209,8 +192,10 @@ def test_jump_capture():
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "b3"]))
 
     # Both source and target are in quantum state.
-    global_names()
-    set_board(["a1", "b1"])
+    board = set_board(["a1", "b1"])
+    for col in ascii_lowercase[:9]:
+        for row in digits:
+            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
     alpha.PhasedSplit()(a1, a2, a3)
     alpha.PhasedSplit()(b1, b2, b3)
     assert_sample_distribution(
@@ -238,8 +223,10 @@ def test_jump_capture():
 
 def test_jump_excluded():
     # Target is in quantum state.
-    global_names()
-    set_board(["a1", "b1"])
+    board = set_board(["a1", "b1"])
+    for col in ascii_lowercase[:9]:
+        for row in digits:
+            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
     alpha.PhasedSplit()(b1, b2, b3)
     Jump(MoveVariant.EXCLUDED)(a1, b2)
     # pop() will break the supersition and only one of the following two states are possible.
@@ -252,8 +239,10 @@ def test_jump_excluded():
         assert_samples_in(board, [locations_to_bitboard(["b2", "b3"])])
 
     # Both source and target are in quantum state.
-    global_names()
-    set_board(["a1", "b1"])
+    board = set_board(["a1", "b1"])
+    for col in ascii_lowercase[:9]:
+        for row in digits:
+            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
     alpha.PhasedSplit()(a1, a2, a3)
     alpha.PhasedSplit()(b1, b2, b3)
     Jump(MoveVariant.EXCLUDED)(a2, b2)
@@ -272,8 +261,10 @@ def test_jump_excluded():
 
 def test_jump_basic():
     # Souce is in quantum state.
-    global_names()
-    set_board(["a1"])
+    board = set_board(["a1"])
+    for col in ascii_lowercase[:9]:
+        for row in digits:
+            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
     alpha.PhasedSplit()(a1, a2, a3)
     Jump(MoveVariant.BASIC)(a2, d1)
     board_probabilities = get_board_probability_distribution(board, 1000)
diff --git a/unitary/examples/quantum_chinese_chess/test_utils.py b/unitary/examples/quantum_chinese_chess/test_utils.py
index 6804558b..1c467a0b 100644
--- a/unitary/examples/quantum_chinese_chess/test_utils.py
+++ b/unitary/examples/quantum_chinese_chess/test_utils.py
@@ -12,23 +12,45 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from unitary.alpha import QuantumObject, QuantumWorld
-from unitary.examples.quantum_chinese_chess.enums import SquareState
+from unitary.examples.quantum_chinese_chess.enums import SquareState, Type, Color
 from unitary.examples.quantum_chinese_chess.board import Board
-from string import ascii_lowercase, digits
+from unitary.examples.quantum_chinese_chess.piece import Piece
+from unitary import alpha
 from typing import List, Dict
 from collections import defaultdict
 from scipy.stats import chisquare
 
 
+_EMPTY_FEN = "9/9/9/9/9/9/9/9/9/9 w---1"
+
+
+def set_board(positions: List[str]) -> Board:
+    """Returns a board with the specified positions filled with
+    RED ROOKs.
+    """
+    board = Board.from_fen(_EMPTY_FEN)
+    for position in positions:
+        board.board[position].reset(
+            Piece(position, SquareState.OCCUPIED, Type.ROOK, Color.RED)
+        )
+        alpha.Flip()(board.board[position])
+    return board
+
+
 def location_to_bit(location: str) -> int:
-    """Transform location notation (e.g. "a3") into a bitboard bit number."""
+    """Transform location notation (e.g. "a3") into a bitboard bit number.
+    The return value ranges from 0 to 89.
+    """
     x = ord(location[0]) - ord("a")
     y = int(location[1])
     return y * 9 + x
 
 
 def locations_to_bitboard(locations: List[str]) -> int:
-    """Transform a list of locations into a 90-bit board bitstring."""
+    """Transform a list of locations into a 90-bit board bitstring.
+    Each nonzero bit of the bitstring indicates that the corresponding
+    piece is occupied.
+    """
     bitboard = 0
     for location in locations:
         bitboard += 1 << location_to_bit(location)
@@ -36,7 +58,7 @@ def locations_to_bitboard(locations: List[str]) -> int:
 
 
 def nth_bit_of(n: int, bit_board: int) -> bool:
-    """Returns the n-th bit of a 90-bit bitstring."""
+    """Returns the `n`-th (zero-based) bit of a 90-bit bitstring `bit_board`."""
     return (bit_board >> n) % 2 == 1
 
 
@@ -48,7 +70,7 @@ def bit_to_location(bit: int) -> str:
 
 
 def bitboard_to_locations(bitboard: int) -> List[str]:
-    """Transform a 90-bit bitstring into a list of locations."""
+    """Transform a 90-bit bitstring `bitboard` into a list of locations."""
     locations = []
     for n in range(90):
         if nth_bit_of(n, bitboard):
@@ -57,8 +79,11 @@ def bitboard_to_locations(bitboard: int) -> List[str]:
 
 
 def sample_board(board: Board, repetitions: int) -> List[int]:
+    """Sample the given `board` by the given `repetitions`.
+    Returns a list of 90-bit bitstring, each corresponding to one sample.
+    """
     samples = board.board.peek(count=repetitions, convert_to_enum=False)
-    # Convert peek results (in List[List[int]]) into bitstring.
+    # Convert peek results (in List[List[int]]) into List[int].
     samples = [
         int("0b" + "".join([str(i) for i in sample[::-1]]), base=2)
         for sample in samples
@@ -66,13 +91,14 @@ def sample_board(board: Board, repetitions: int) -> List[int]:
     return samples
 
 
-def print_samples(samples):
-    """Prints all the samples as lists of locations."""
+def print_samples(samples: List[int]) -> None:
+    """Aggregate all the samples and print the dictionary of {locations: count}."""
     sample_dict = {}
     for sample in samples:
         if sample not in sample_dict:
             sample_dict[sample] = 0
         sample_dict[sample] += 1
+    print("Actual samples:")
     for key in sample_dict:
         print(f"{bitboard_to_locations(key)}: {sample_dict[key]}")
 
@@ -81,17 +107,11 @@ def get_board_probability_distribution(
     board: Board, repetitions: int = 1000
 ) -> Dict[int, float]:
     """Returns the probability distribution for each board found in the sample.
-
     The values are returned as a dict{bitboard(int): probability(float)}.
     """
     board_probabilities: Dict[int, float] = {}
 
-    samples = board.board.peek(count=repetitions, convert_to_enum=False)
-    # Convert peek results (in List[List[int]]) into bitstring.
-    samples = [
-        int("0b" + "".join([str(i) for i in sample[::-1]]), base=2)
-        for sample in samples
-    ]
+    samples = sample_board(board, repetitions)
     for sample in samples:
         if sample not in board_probabilities:
             board_probabilities[sample] = 0.0
@@ -103,63 +123,73 @@ def get_board_probability_distribution(
     return board_probabilities
 
 
-def assert_samples_in(board: Board, possibilities):
+def assert_samples_in(board: Board, probabilities: Dict[int, float]) -> None:
+    """Samples the given `board` and asserts that all samples are within
+    the given `probabilities` (i.e. a map from bitstring into its possibility),
+    and that each possibility is represented at least once in the samples.
+    """
     samples = sample_board(board, 500)
     assert len(samples) == 500
-    all_in = all(sample in possibilities for sample in samples)
-    print(possibilities)
-    print(set(samples))
+    all_in = all(sample in probabilities for sample in samples)
     assert all_in, print_samples(samples)
     # Make sure each possibility is represented at least once.
-    for possibility in possibilities:
+    for possibility in probabilities:
         any_in = any(sample == possibility for sample in samples)
         assert any_in, print_samples(samples)
 
 
-def assert_sample_distribution(board: Board, probability_map, p_significant=1e-6):
+def assert_sample_distribution(
+    board: Board, probabilities: Dict[int, float], p_significant: float = 1e-6
+) -> None:
     """Performs a chi-squared test that samples follow an expected distribution.
-
-    probability_map is a map from bitboards to expected probability. An
-    assertion is raised if one of the samples is not in the map, or if the
-    probability that the samples are at least as different from the expected
-    ones as the observed sampless is less than p_significant.
+    `probabilities` is a map from bitboards to expected probability. An
+    AssertionError is raised if any of the samples is not in the map, or if the
+    expected versus observed samples fails the chi-squared test.
     """
-    assert abs(sum(probability_map.values()) - 1) < 1e-9
-    samples = sample_board(board, 500)
-    assert len(samples) == 500
+    n_samples = 500
+    assert abs(sum(probabilities.values()) - 1) < 1e-9
+    samples = sample_board(board, n_samples)
     counts = defaultdict(int)
     for sample in samples:
-        assert sample in probability_map, bitboard_to_locations(sample)
+        assert sample in probabilities, bitboard_to_locations(sample)
         counts[sample] += 1
     observed = []
     expected = []
-    for position, probability in probability_map.items():
+    for position, probability in probabilities.items():
         observed.append(counts[position])
-        expected.append(500 * probability)
+        expected.append(n_samples * probability)
     p = chisquare(observed, expected).pvalue
     assert (
         p > p_significant
-    ), f"Observed {observed} far from expected {expected} (p = {p})"
+    ), f"Observed {observed} is far from expected {expected} (p = {p})"
 
 
-def assert_this_or_that(samples, this, that):
-    """Asserts all the samples are either equal to this or that,
-    and that one of each exists in the samples.
+def assert_this_or_that(samples: List[int], this: int, that: int) -> None:
+    """Asserts all the samples are either equal to `this` or `that`,
+    and that at least one of them exists in the samples.
     """
-    # assert any(sample == this for sample in samples), print_samples(samples)
-    # assert any(sample == that for sample in samples), print_samples(samples)
+    assert any(sample == this for sample in samples), print_samples(samples)
+    assert any(sample == that for sample in samples), print_samples(samples)
     assert all(sample == this or sample == that for sample in samples), print_samples(
         samples
     )
 
 
-def assert_prob_about(probs, that, expected, atol=0.05):
-    """Checks that the probability is within atol of the expected value."""
-    assert that in probs, print_samples(list(probs.keys()))
-    assert probs[that] > expected - atol, print_samples(list(probs.keys()))
-    assert probs[that] < expected + atol, print_samples(list(probs.keys()))
+def assert_prob_about(
+    probabilities: Dict[int, float], that: int, expected: float, atol: float = 0.05
+) -> None:
+    """Checks that the probability of `that` is within `atol` of the value of `expected`."""
+    assert that in probabilities, print_samples(list(probabilities.keys()))
+    assert probabilities[that] > expected - atol, print_samples(
+        list(probabilities.keys())
+    )
+    assert probabilities[that] < expected + atol, print_samples(
+        list(probabilities.keys())
+    )
 
 
-def assert_fifty_fifty(probs, that):
-    """Checks that the probability is close to 50%."""
-    assert_prob_about(probs, that, 0.5), print_samples(list(probs.keys()))
+def assert_fifty_fifty(probabilities, that):
+    """Checks that the probability of `that` is close to 50%."""
+    assert_prob_about(probabilities, that, 0.5), print_samples(
+        list(probabilities.keys())
+    )

From 0ee28c1a340624af8755d4da28fe8cd9531110fb Mon Sep 17 00:00:00 2001
From: Pengfei Chen <pengfeichen@google.com>
Date: Wed, 18 Oct 2023 14:37:40 -0700
Subject: [PATCH 10/13] update

---
 .../quantum_chinese_chess/move_test.py        | 71 ++++++++-----------
 1 file changed, 28 insertions(+), 43 deletions(-)

diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index c375e1bb..e44804ba 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -17,7 +17,6 @@
 import pytest
 from unitary import alpha
 from typing import List
-from string import ascii_lowercase, digits
 from unitary.examples.quantum_chinese_chess.enums import (
     MoveType,
     MoveVariant,
@@ -147,33 +146,29 @@ def test_to_str():
 def test_jump_classical():
     # Target is empty.
     board = set_board(["a1", "b1"])
-    # TODO(): try move the following varaibles declarations into a function.
-    for col in ascii_lowercase[:9]:
-        for row in digits:
-            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
-    Jump(MoveVariant.CLASSICAL)(a1, b2)
+    world = board.board
+    # TODO(): try move all varaibles declarations of a1 = world["a1"] into a function.
+    Jump(MoveVariant.CLASSICAL)(world["a1"], world["b2"])
     assert_samples_in(board, [locations_to_bitboard(["b2", "b1"])])
 
     # Target is occupied.
-    Jump(MoveVariant.CLASSICAL)(b2, b1)
+    Jump(MoveVariant.CLASSICAL)(world["b2"], world["b1"])
     assert_samples_in(board, [locations_to_bitboard(["b1"])])
 
 
 def test_jump_capture():
     # Source is in quantum state.
     board = set_board(["a1", "b1"])
-    for col in ascii_lowercase[:9]:
-        for row in digits:
-            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
-    alpha.PhasedSplit()(a1, a2, a3)
+    world = board.board
+    alpha.PhasedSplit()(world["a1"], world["a2"], world["a3"])
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a2", "b1"]))
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b1"]))
-    Jump(MoveVariant.CAPTURE)(a2, b1)
+    Jump(MoveVariant.CAPTURE)(world["a2"], world["b1"])
     # pop() will break the supersition and only one of the following two states are possible.
     # We check the ancilla to learn if the jump was applied or not.
-    source_is_occupied = board.board.post_selection[board.board["ancilla_a2_0"]]
+    source_is_occupied = world.post_selection[world["ancilla_a2_0"]]
     if source_is_occupied:
         assert_samples_in(board, [locations_to_bitboard(["b1"])])
     else:
@@ -181,11 +176,9 @@ def test_jump_capture():
 
     # Target is in quantum state.
     board = set_board(["a1", "b1"])
-    for col in ascii_lowercase[:9]:
-        for row in digits:
-            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
-    alpha.PhasedSplit()(b1, b2, b3)
-    Jump(MoveVariant.CAPTURE)(a1, b2)
+    world = board.board
+    alpha.PhasedSplit()(world["b1"], world["b2"], world["b3"])
+    Jump(MoveVariant.CAPTURE)(world["a1"], world["b2"])
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2"]))
@@ -193,11 +186,9 @@ def test_jump_capture():
 
     # Both source and target are in quantum state.
     board = set_board(["a1", "b1"])
-    for col in ascii_lowercase[:9]:
-        for row in digits:
-            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
-    alpha.PhasedSplit()(a1, a2, a3)
-    alpha.PhasedSplit()(b1, b2, b3)
+    world = board.board
+    alpha.PhasedSplit()(world["a1"], world["a2"], world["a3"])
+    alpha.PhasedSplit()(world["b1"], world["b2"], world["b3"])
     assert_sample_distribution(
         board,
         {
@@ -207,11 +198,11 @@ def test_jump_capture():
             locations_to_bitboard(["a3", "b3"]): 1 / 4.0,
         },
     )
-    Jump(MoveVariant.CAPTURE)(a2, b2)
+    Jump(MoveVariant.CAPTURE)(world["a2"], world["b2"])
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
     # We check the ancilla to learn if the jump was applied or not.
-    source_is_occupied = board.board.post_selection[board.board["ancilla_a2_0"]]
+    source_is_occupied = world.post_selection[world["ancilla_a2_0"]]
     print(source_is_occupied)
     if source_is_occupied:
         assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2"]))
@@ -224,14 +215,12 @@ def test_jump_capture():
 def test_jump_excluded():
     # Target is in quantum state.
     board = set_board(["a1", "b1"])
-    for col in ascii_lowercase[:9]:
-        for row in digits:
-            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
-    alpha.PhasedSplit()(b1, b2, b3)
-    Jump(MoveVariant.EXCLUDED)(a1, b2)
+    world = board.board
+    alpha.PhasedSplit()(world["b1"], world["b2"], world["b3"])
+    Jump(MoveVariant.EXCLUDED)(world["a1"], world["b2"])
     # pop() will break the supersition and only one of the following two states are possible.
     # We check the ancilla to learn if the jump was applied or not.
-    target_is_occupied = board.board.post_selection[board.board["ancilla_b2_0"]]
+    target_is_occupied = world.post_selection[world["ancilla_b2_0"]]
     print(target_is_occupied)
     if target_is_occupied:
         assert_samples_in(board, [locations_to_bitboard(["a1", "b2"])])
@@ -240,16 +229,14 @@ def test_jump_excluded():
 
     # Both source and target are in quantum state.
     board = set_board(["a1", "b1"])
-    for col in ascii_lowercase[:9]:
-        for row in digits:
-            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
-    alpha.PhasedSplit()(a1, a2, a3)
-    alpha.PhasedSplit()(b1, b2, b3)
-    Jump(MoveVariant.EXCLUDED)(a2, b2)
+    world = board.board
+    alpha.PhasedSplit()(world["a1"], world["a2"], world["a3"])
+    alpha.PhasedSplit()(world["b1"], world["b2"], world["b3"])
+    Jump(MoveVariant.EXCLUDED)(world["a2"], world["b2"])
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
     # We check the ancilla to learn if the jump was applied or not.
-    target_is_occupied = board.board.post_selection[board.board["ancilla_b2_0"]]
+    target_is_occupied = world.post_selection[world["ancilla_b2_0"]]
     print(target_is_occupied)
     if target_is_occupied:
         assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a2", "b2"]))
@@ -262,11 +249,9 @@ def test_jump_excluded():
 def test_jump_basic():
     # Souce is in quantum state.
     board = set_board(["a1"])
-    for col in ascii_lowercase[:9]:
-        for row in digits:
-            globals()[f"{col}{row}"] = board.board[f"{col}{row}"]
-    alpha.PhasedSplit()(a1, a2, a3)
-    Jump(MoveVariant.BASIC)(a2, d1)
+    world = board.board
+    alpha.PhasedSplit()(world["a1"], world["a2"], world["a3"])
+    Jump(MoveVariant.BASIC)(world["a2"], world["d1"])
     board_probabilities = get_board_probability_distribution(board, 1000)
     assert len(board_probabilities) == 2
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["d1"]))

From 40463a1b9cb6f262e9146217b154984adc871729 Mon Sep 17 00:00:00 2001
From: madcpf <madcpf@gmail.com>
Date: Thu, 19 Oct 2023 21:55:05 -0700
Subject: [PATCH 11/13] update

---
 unitary/examples/quantum_chinese_chess/move_test.py | 2 +-
 unitary/examples/quantum_chinese_chess/piece.py     | 7 ++++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index e44804ba..0c8e156f 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -247,7 +247,7 @@ def test_jump_excluded():
 
 
 def test_jump_basic():
-    # Souce is in quantum state.
+    # Source is in quantum state.
     board = set_board(["a1"])
     world = board.board
     alpha.PhasedSplit()(world["a1"], world["a2"], world["a3"])
diff --git a/unitary/examples/quantum_chinese_chess/piece.py b/unitary/examples/quantum_chinese_chess/piece.py
index 9dd6e763..fc6736e8 100644
--- a/unitary/examples/quantum_chinese_chess/piece.py
+++ b/unitary/examples/quantum_chinese_chess/piece.py
@@ -22,6 +22,11 @@
 
 
 class Piece(QuantumObject):
+    """Each Piece stands for a position in the board, it could hold different types/colors
+    of chess pieces either in quantum state or classical state. It could also be classically
+    EMPTY.
+    """
+
     def __init__(self, name: str, state: SquareState, type_: Type, color: Color):
         QuantumObject.__init__(self, name, state)
         self.type_ = type_
@@ -40,7 +45,7 @@ def __str__(self):
 
     def reset(self, piece: "Piece" = None) -> None:
         """Modifies the classical attributes of the piece.
-        If piece is provided, then its type_, color, and is_entangled is copied,
+        If `piece` is provided, then its type_, color, and is_entangled is copied,
         otherwise set the current piece to be a classically empty piece.
         """
         if piece is not None:

From e3a7fbcb33fb0429079b0a164dd0270fc16c44bb Mon Sep 17 00:00:00 2001
From: Pengfei Chen <pengfeichen@google.com>
Date: Wed, 25 Oct 2023 12:48:28 -0700
Subject: [PATCH 12/13] up

---
 unitary/alpha/quantum_world.py                | 13 ++++++++---
 .../examples/quantum_chinese_chess/board.py   | 12 ++++++++++
 .../examples/quantum_chinese_chess/chess.py   |  7 +++---
 .../examples/quantum_chinese_chess/move.py    |  6 ++---
 .../quantum_chinese_chess/move_test.py        | 15 ++++++++-----
 .../quantum_chinese_chess/test_utils.py       | 22 ++++---------------
 6 files changed, 43 insertions(+), 32 deletions(-)

diff --git a/unitary/alpha/quantum_world.py b/unitary/alpha/quantum_world.py
index f5f74bb3..f72b1c6e 100644
--- a/unitary/alpha/quantum_world.py
+++ b/unitary/alpha/quantum_world.py
@@ -306,10 +306,17 @@ def _interpret_result(self, result: Union[int, Iterable[int]]) -> int:
         return result
 
     def unhook(self, object: QuantumObject) -> None:
-        """Replaces all usages of the given object in the circuit with a new ancilla with value=0."""
-        # Creates a new ancilla.
+        """Replace all usages of the given `object` in the circuit with a new ancilla,
+        so that
+         - all former operations on `object` will be applied on the new ancilla; 
+         - future operations on `object` start with its new reset value.
+        
+        Note that we don't do force measurement on it, since we don't care about its
+        current value but just want to reset it.
+        """
+        # Create a new ancilla.
         new_ancilla = self._add_ancilla(object.name)
-        # Replace operations using the qubit of the given object with the new ancilla.
+        # Replace operations of the given `object` with the new ancilla.
         qubit_remapping_dict = {
             object.qubit: new_ancilla.qubit,
             new_ancilla.qubit: object.qubit,
diff --git a/unitary/examples/quantum_chinese_chess/board.py b/unitary/examples/quantum_chinese_chess/board.py
index efd3fbf4..4c824c55 100644
--- a/unitary/examples/quantum_chinese_chess/board.py
+++ b/unitary/examples/quantum_chinese_chess/board.py
@@ -184,3 +184,15 @@ def flying_general_check(self) -> bool:
             # If there are no pieces between two KINGs, the check successes. Game ends.
             return True
         # TODO(): add check when there are quantum pieces in between.
+
+    def sample(self, repetitions: int) -> List[int]:
+        """Sample the current board by the given `repetitions`.
+        Returns a list of 90-bit bitstring, each corresponding to one sample.
+        """
+        samples = self.board.peek(count=repetitions, convert_to_enum=False)
+        # Convert peek results (in List[List[int]]) into List[int].
+        samples = [
+            int("0b" + "".join([str(i) for i in sample[::-1]]), base=2)
+            for sample in samples
+        ]
+        return samples
diff --git a/unitary/examples/quantum_chinese_chess/chess.py b/unitary/examples/quantum_chinese_chess/chess.py
index 27aa8bc6..fad477e5 100644
--- a/unitary/examples/quantum_chinese_chess/chess.py
+++ b/unitary/examples/quantum_chinese_chess/chess.py
@@ -386,15 +386,16 @@ def apply_move(self, str_to_parse: str) -> None:
             quantum_pieces_1,
         )
 
-        print(move_type, " ", move_variant)
+        if self.debug_level > 1:
+            print(move_type, " ", move_variant)
 
         # Apply the move accoding to its type.
         if move_type == MoveType.CLASSICAL:
             if source_0.type_ == Type.KING:
                 # Update the locations of KING.
                 self.board.king_locations[self.current_player] = targets[0]
-                # TODO(): only make such prints for a certain debug level.
-                print(f"Updated king locations: {self.board.king_locations}.")
+                if self.debug_level > 1:
+                    print(f"Updated king locations: {self.board.king_locations}.")
             if target_0.type_ == Type.KING:
                 # King is captured, then the game is over.
                 self.game_state = GameState(self.current_player)
diff --git a/unitary/examples/quantum_chinese_chess/move.py b/unitary/examples/quantum_chinese_chess/move.py
index 6f8bddce..1b2cdd59 100644
--- a/unitary/examples/quantum_chinese_chess/move.py
+++ b/unitary/examples/quantum_chinese_chess/move.py
@@ -165,9 +165,9 @@ def effect(self, *objects) -> Iterator[cirq.Operation]:
             target_0.reset()
         elif self.move_variant == MoveVariant.CLASSICAL:
             if target_0.type_ != Type.EMPTY:
-                # For classical moves with target_0 occupied, we replace the qubit of target_0 with
-                # a new ancilla, and set its classical properties to be EMPTY.
-                world.unhook(target_0)
+                # For classical moves with target_0 occupied, we flip target_0 to be empty,
+                # and set its classical properties to be EMPTY.
+                alpha.Flip()(target_0)
                 target_0.reset()
 
         # Make the jump move.
diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index 0c8e156f..d6ae9a6d 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -31,7 +31,6 @@
     assert_this_or_that,
     assert_prob_about,
     assert_fifty_fifty,
-    sample_board,
     get_board_probability_distribution,
     print_samples,
     set_board,
@@ -156,7 +155,7 @@ def test_jump_classical():
     assert_samples_in(board, [locations_to_bitboard(["b1"])])
 
 
-def test_jump_capture():
+def test_jump_capture_quantum_source():
     # Source is in quantum state.
     board = set_board(["a1", "b1"])
     world = board.board
@@ -166,7 +165,7 @@ def test_jump_capture():
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a2", "b1"]))
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b1"]))
     Jump(MoveVariant.CAPTURE)(world["a2"], world["b1"])
-    # pop() will break the supersition and only one of the following two states are possible.
+    # pop() will break the superposition and only one of the following two states are possible.
     # We check the ancilla to learn if the jump was applied or not.
     source_is_occupied = world.post_selection[world["ancilla_a2_0"]]
     if source_is_occupied:
@@ -174,6 +173,8 @@ def test_jump_capture():
     else:
         assert_samples_in(board, [locations_to_bitboard(["a3", "b1"])])
 
+
+def test_jump_capture_quantum_target():
     # Target is in quantum state.
     board = set_board(["a1", "b1"])
     world = board.board
@@ -184,6 +185,8 @@ def test_jump_capture():
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2"]))
     assert_fifty_fifty(board_probabilities, locations_to_bitboard(["b2", "b3"]))
 
+
+def test_jump_capture_quantum_source_and_target():
     # Both source and target are in quantum state.
     board = set_board(["a1", "b1"])
     world = board.board
@@ -212,13 +215,13 @@ def test_jump_capture():
         assert_fifty_fifty(board_probabilities, locations_to_bitboard(["a3", "b3"]))
 
 
-def test_jump_excluded():
+def test_jump_excluded_quantum_target():
     # Target is in quantum state.
     board = set_board(["a1", "b1"])
     world = board.board
     alpha.PhasedSplit()(world["b1"], world["b2"], world["b3"])
     Jump(MoveVariant.EXCLUDED)(world["a1"], world["b2"])
-    # pop() will break the supersition and only one of the following two states are possible.
+    # pop() will break the superposition and only one of the following two states are possible.
     # We check the ancilla to learn if the jump was applied or not.
     target_is_occupied = world.post_selection[world["ancilla_b2_0"]]
     print(target_is_occupied)
@@ -227,6 +230,8 @@ def test_jump_excluded():
     else:
         assert_samples_in(board, [locations_to_bitboard(["b2", "b3"])])
 
+
+def test_jump_excluded_quantum_source_and_target():
     # Both source and target are in quantum state.
     board = set_board(["a1", "b1"])
     world = board.board
diff --git a/unitary/examples/quantum_chinese_chess/test_utils.py b/unitary/examples/quantum_chinese_chess/test_utils.py
index 1c467a0b..8f4c6f24 100644
--- a/unitary/examples/quantum_chinese_chess/test_utils.py
+++ b/unitary/examples/quantum_chinese_chess/test_utils.py
@@ -78,19 +78,6 @@ def bitboard_to_locations(bitboard: int) -> List[str]:
     return locations
 
 
-def sample_board(board: Board, repetitions: int) -> List[int]:
-    """Sample the given `board` by the given `repetitions`.
-    Returns a list of 90-bit bitstring, each corresponding to one sample.
-    """
-    samples = board.board.peek(count=repetitions, convert_to_enum=False)
-    # Convert peek results (in List[List[int]]) into List[int].
-    samples = [
-        int("0b" + "".join([str(i) for i in sample[::-1]]), base=2)
-        for sample in samples
-    ]
-    return samples
-
-
 def print_samples(samples: List[int]) -> None:
     """Aggregate all the samples and print the dictionary of {locations: count}."""
     sample_dict = {}
@@ -111,7 +98,7 @@ def get_board_probability_distribution(
     """
     board_probabilities: Dict[int, float] = {}
 
-    samples = sample_board(board, repetitions)
+    samples = board.sample(repetitions)
     for sample in samples:
         if sample not in board_probabilities:
             board_probabilities[sample] = 0.0
@@ -128,8 +115,7 @@ def assert_samples_in(board: Board, probabilities: Dict[int, float]) -> None:
     the given `probabilities` (i.e. a map from bitstring into its possibility),
     and that each possibility is represented at least once in the samples.
     """
-    samples = sample_board(board, 500)
-    assert len(samples) == 500
+    samples = board.sample(500)
     all_in = all(sample in probabilities for sample in samples)
     assert all_in, print_samples(samples)
     # Make sure each possibility is represented at least once.
@@ -148,7 +134,7 @@ def assert_sample_distribution(
     """
     n_samples = 500
     assert abs(sum(probabilities.values()) - 1) < 1e-9
-    samples = sample_board(board, n_samples)
+    samples = board.sample(n_samples)
     counts = defaultdict(int)
     for sample in samples:
         assert sample in probabilities, bitboard_to_locations(sample)
@@ -176,7 +162,7 @@ def assert_this_or_that(samples: List[int], this: int, that: int) -> None:
 
 
 def assert_prob_about(
-    probabilities: Dict[int, float], that: int, expected: float, atol: float = 0.05
+    probabilities: Dict[int, float], that: int, expected: float, atol: float = 0.06
 ) -> None:
     """Checks that the probability of `that` is within `atol` of the value of `expected`."""
     assert that in probabilities, print_samples(list(probabilities.keys()))

From 8d5bee57f4d3258fedd394225b62dedb56738e43 Mon Sep 17 00:00:00 2001
From: Pengfei Chen <pengfeichen@google.com>
Date: Wed, 25 Oct 2023 12:49:45 -0700
Subject: [PATCH 13/13] update

---
 unitary/alpha/quantum_world.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/unitary/alpha/quantum_world.py b/unitary/alpha/quantum_world.py
index f72b1c6e..d1b34c70 100644
--- a/unitary/alpha/quantum_world.py
+++ b/unitary/alpha/quantum_world.py
@@ -308,9 +308,9 @@ def _interpret_result(self, result: Union[int, Iterable[int]]) -> int:
     def unhook(self, object: QuantumObject) -> None:
         """Replace all usages of the given `object` in the circuit with a new ancilla,
         so that
-         - all former operations on `object` will be applied on the new ancilla; 
+         - all former operations on `object` will be applied on the new ancilla;
          - future operations on `object` start with its new reset value.
-        
+
         Note that we don't do force measurement on it, since we don't care about its
         current value but just want to reset it.
         """