Allow asking for psi+ EPR pairs

QuTech-Delft · Feb 15, 2024 · ed7791c · ed7791c
1 parent 216eeb3
commit ed7791c
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 26 deletions.
diff --git a/netqasm/sdk/build_epr.py b/netqasm/sdk/build_epr.py
@@ -28,6 +28,7 @@ class EntRequestParams:
     time_unit: TimeUnit = TimeUnit.MICRO_SECONDS
     max_time: int = 0
     expect_phi_plus: bool = True
+    expect_psi_plus: bool = False
     min_fidelity_all_at_end: Optional[int] = None
     max_tries: Optional[int] = None
     random_basis_local: Optional[RandomBasis] = None

diff --git a/netqasm/sdk/builder.py b/netqasm/sdk/builder.py
@@ -373,7 +373,7 @@ def post_loop(conn: BaseNetQASMConnection, loop_reg: RegFuture):
                 pair=pair,
             )
 
-            if params.expect_phi_plus and role == EPRRole.RECV:
+            if (params.expect_phi_plus or params.expect_psi_plus) and role == EPRRole.RECV:
                 # Perform Bell corrections
                 bell_state = self._get_raw_bell_state(
                     ent_results_array, loop_reg, bell_state_reg
@@ -382,7 +382,7 @@ def post_loop(conn: BaseNetQASMConnection, loop_reg: RegFuture):
                     instruction=GenericInstr.SET, operands=[qubit_reg, 0]
                 )
                 self.subrt_add_pending_command(set_qubit_reg_cmd)  # type: ignore
-                self._build_cmds_epr_keep_corrections_single_pair(bell_state, qubit_reg)
+                self._build_cmds_epr_keep_corrections_single_pair(bell_state, qubit_reg, params)
 
             # If it's the last pair, don't move it to a mem qubit
             with loop_reg.if_ne(params.number - 1):
@@ -444,15 +444,15 @@ def post_loop(conn: BaseNetQASMConnection, loop_reg: RegFuture):
             )
             assert tp == EPRType.K or tp == EPRType.R
 
-            if params.expect_phi_plus and role == EPRRole.RECV:
+            if (params.expect_phi_plus or params.expect_psi_plus) and role == EPRRole.RECV:
                 bell_state = self._get_raw_bell_state(
                     ent_results_array, loop_reg, bell_state_reg
                 )
                 set_qubit_reg_cmd = ICmd(
                     instruction=GenericInstr.SET, operands=[qubit_reg, 0]
                 )
                 self.subrt_add_pending_command(set_qubit_reg_cmd)  # type: ignore
-                self._build_cmds_epr_keep_corrections_single_pair(bell_state, qubit_reg)
+                self._build_cmds_epr_keep_corrections_single_pair(bell_state, qubit_reg, params)
 
             q_id = qubit_ids.get_future_index(loop_register)
             q = FutureQubit(conn=conn, future_id=q_id)
@@ -1353,26 +1353,47 @@ def undef_result_element(conn: BaseNetQASMConnection, _: RegFuture):
         )
 
     def _build_cmds_epr_keep_corrections_single_pair(
-        self, bell_state: RegFuture, qubit_reg: operand.Register
+        self, bell_state: RegFuture, qubit_reg: operand.Register, params: EntRequestParams
     ) -> None:
-        with bell_state.if_eq(BellState.PHI_MINUS.value):  # Phi- -> apply Z-gate
-            correction_cmds = [
-                ICmd(instruction=GenericInstr.ROT_Z, operands=[qubit_reg, 16, 4])
-            ]
-            self.subrt_add_pending_commands(correction_cmds)  # type: ignore
-        with bell_state.if_eq(BellState.PSI_PLUS.value):  # Psi+ -> apply X-gate
-            correction_cmds = [
-                ICmd(instruction=GenericInstr.ROT_X, operands=[qubit_reg, 16, 4])
-            ]
-            self.subrt_add_pending_commands(correction_cmds)  # type: ignore
-        with bell_state.if_eq(
-            BellState.PSI_MINUS.value
-        ):  # Psi- -> apply X-gate and Z-gate
-            correction_cmds = [
-                ICmd(instruction=GenericInstr.ROT_X, operands=[qubit_reg, 16, 4]),
-                ICmd(instruction=GenericInstr.ROT_Z, operands=[qubit_reg, 16, 4]),
-            ]
-            self.subrt_add_pending_commands(correction_cmds)  # type: ignore
+        if params.expect_phi_plus:
+            with bell_state.if_eq(BellState.PHI_MINUS.value):  # Phi- -> apply Z-gate
+                correction_cmds = [
+                    ICmd(instruction=GenericInstr.ROT_Z, operands=[qubit_reg, 16, 4])
+                ]
+                self.subrt_add_pending_commands(correction_cmds)  # type: ignore
+            with bell_state.if_eq(BellState.PSI_PLUS.value):  # Psi+ -> apply X-gate
+                correction_cmds = [
+                    ICmd(instruction=GenericInstr.ROT_X, operands=[qubit_reg, 16, 4])
+                ]
+                self.subrt_add_pending_commands(correction_cmds)  # type: ignore
+            with bell_state.if_eq(
+                BellState.PSI_MINUS.value
+            ):  # Psi- -> apply X-gate and Z-gate
+                correction_cmds = [
+                    ICmd(instruction=GenericInstr.ROT_X, operands=[qubit_reg, 16, 4]),
+                    ICmd(instruction=GenericInstr.ROT_Z, operands=[qubit_reg, 16, 4]),
+                ]
+                self.subrt_add_pending_commands(correction_cmds)  # type: ignore
+        else:
+            assert params.expect_psi_plus
+            with bell_state.if_eq(BellState.PHI_PLUS.value):  # Phi+ -> apply X-gate
+                correction_cmds = [
+                    ICmd(instruction=GenericInstr.ROT_X, operands=[qubit_reg, 16, 4])
+                ]
+                self.subrt_add_pending_commands(correction_cmds)  # type: ignore
+            with bell_state.if_eq(BellState.PHI_MINUS.value):  # Phi- -> apply X-gate and Z-gate
+                correction_cmds = [
+                    ICmd(instruction=GenericInstr.ROT_X, operands=[qubit_reg, 16, 4]),
+                    ICmd(instruction=GenericInstr.ROT_Z, operands=[qubit_reg, 16, 4])
+                ]
+                self.subrt_add_pending_commands(correction_cmds)  # type: ignore
+            with bell_state.if_eq(
+                BellState.PSI_MINUS.value
+            ):  # Psi- -> apply Z-gate
+                correction_cmds = [
+                    ICmd(instruction=GenericInstr.ROT_Z, operands=[qubit_reg, 16, 4]),
+                ]
+                self.subrt_add_pending_commands(correction_cmds)  # type: ignore
 
     def _build_cmds_epr_keep_corrections(
         self, qubit_ids_array: Array, ent_results_array: Array, params: EntRequestParams
@@ -1394,7 +1415,7 @@ def loop(conn: BaseNetQASMConnection, index: RegFuture):
                 instruction=GenericInstr.SET, operands=[qubit_reg, 0]
             )
             self.subrt_add_pending_command(set_qubit_reg_cmd)  # type: ignore
-            self._build_cmds_epr_keep_corrections_single_pair(bell_state, qubit_reg)
+            self._build_cmds_epr_keep_corrections_single_pair(bell_state, qubit_reg, params)
 
         self._build_cmds_loop_body(
             loop, stop=params.number, loop_register=loop_register
@@ -1501,7 +1522,7 @@ def _build_cmds_epr_recv_keep(
 
         self.subrt_add_pending_commands(wait_cmds)  # type: ignore
 
-        if wait_all and params.expect_phi_plus:
+        if wait_all and (params.expect_phi_plus or params.expect_psi_plus):
             self._build_cmds_epr_keep_corrections(
                 qubit_ids_array, ent_results_array, params
             )
@@ -1641,7 +1662,7 @@ def _build_cmds_epr_recv_rsp(
 
         self.subrt_add_pending_commands(wait_cmds)  # type: ignore
 
-        if wait_all and params.expect_phi_plus:
+        if wait_all and (params.expect_phi_plus or params.expect_psi_plus):
             self._build_cmds_epr_keep_corrections(
                 qubit_ids_array, ent_results_array, params
             )

diff --git a/netqasm/sdk/epr_socket.py b/netqasm/sdk/epr_socket.py
@@ -644,6 +644,7 @@ def recv_keep(
         post_routine: Optional[Callable] = None,
         sequential: bool = False,
         expect_phi_plus: bool = True,
+        expect_psi_plus: bool = False,
         min_fidelity_all_at_end: Optional[int] = None,
         max_tries: Optional[int] = None,
     ) -> List[Qubit]:
@@ -680,6 +681,8 @@ def recv_keep(
 
         if self.conn is None:
             raise RuntimeError("EPRSocket does not have an open connection")
+
+        assert not (expect_phi_plus and expect_psi_plus), "cannot ask for both phi+ and psi+"
 
         qubits, _ = self.conn._builder.sdk_recv_epr_keep(
             params=EntRequestParams(
@@ -689,6 +692,7 @@ def recv_keep(
                 post_routine=post_routine,
                 sequential=sequential,
                 expect_phi_plus=expect_phi_plus,
+                expect_psi_plus=expect_psi_plus,
                 min_fidelity_all_at_end=min_fidelity_all_at_end,
                 max_tries=max_tries,
             ),