Modified reset_registers to take in tuples. Modified respective tests.

src/qibo/tomography/gate_set_tomography.py

@@ -70,6 +70,92 @@ def measurement_basis(j, circ):
     return new_circ
 
 
+def reset_register(circuit, invert_register):
+    """Returns an inverse circuit of the selected register to prepare the zero state \\(|0\rangle\\).
+        One can then add inverse_circuit to the original circuit by addition:
+            circ_with_inverse = circ.copy()
+            circ_with_inverse.add(inverse_circuit.on_qubits(invert_register))
+        where register_to_reset = (0,), (1,) , or (0, 1).
+
+        Args:
+        circuit (:class:`qibo.models.Circuit`): original circuit
+        invert_register (tuple): Qubit(s) to reset: Use a tuple to specify which qubit(s) to reset:
+            - (0,) to reset qubit 0;
+            - (1,) to reset qubit 1; or
+            - (0,1) to reset both qubits.
+    Returns:
+        inverse_circuit (:class:`qibo.models.Circuit`): Inverse of the input circuit's register.
+    """
+    valid_registers = [(0,), (1,), (0, 1)]
+    if invert_register is not None:
+        if not isinstance(invert_register, tuple) or invert_register not in valid_registers:
+            raise_error(
+                NameError,
+                f"{invert_register} not recognized.",
+            )
+
+        elif invert_register == (0,) or invert_register == (1,):
+            register_to_reset = invert_register[0]
+            new_circ = Circuit(1)
+            for data in circuit.raw["queue"]:
+                init_kwargs = data.get("init_kwargs", {})
+                if data["_target_qubits"][0] == register_to_reset:
+                    new_circ.add(getattr(gates, data["_class"])(0, **init_kwargs))
+
+        else:
+            new_circ = circuit.copy()
+
+    return new_circ.invert()
+
+
+# def reset_register(circuit, invert_register):
+#     """Returns an inverse circuit of the selected register to prepare the zero state \\(|0\rangle\\).
+#         One can then add inverse_circuit to the original circuit by addition:
+#             circ_with_inverse = circ.copy()
+#             circ_with_inverse.add(inverse_circuit.on_qubits(invert_register))
+#         where register_to_reset = 0, 1, or [0,1].
+
+#         Args:
+#         circuit (:class:`qibo.models.Circuit`): original circuit
+#         invert_register (string): Qubit(s) to reset:
+#             'sp_0' (qubit 0);
+#             'sp_1' (qubit 1); or
+#             'sp_t' (both qubits)
+#             where 'sp' is an abbreviation for state_preparation.
+#     Returns:
+#         inverse_circuit (:class:`qibo.models.Circuit`): Inverse of the input circuit's register.
+#     """
+
+#     if invert_register == "sp_0" or invert_register == "sp_1":
+#         if invert_register == "sp_0":
+#             register_to_reset = 0
+#         elif invert_register == "sp_1":
+#             register_to_reset = 1
+
+#         new_circ = Circuit(1)
+#         for data in circuit.raw["queue"]:
+#             init_kwargs = data.get("init_kwargs", {})
+#             if data["_target_qubits"][0] == register_to_reset:
+#                 new_circ.add(getattr(gates, data["_class"])(0, **init_kwargs))
+
+#     elif invert_register == "sp_t":
+#         new_circ = circuit.copy()
+
+#     else:
+#         raise_error(
+#             NameError,
+#             f"{invert_register} not recognized. Input "
+#             "sp_0"
+#             " to reset qubit 0, "
+#             "sp_1"
+#             " to reset qubit 1, or "
+#             "sp_t"
+#             " to reset both qubits.",
+#         )
+
+#     return new_circ.invert()
+
+
 def GST_execute_circuit(circuit, k, j, nshots=int(1e4), backend=None):
     """Executes a circuit used in gate set tomography and processes the
         measurement outcomes for the Pauli Transfer Matrix notation. The circuit
@@ -142,24 +228,33 @@ def execute_GST(
             f"nqubits given as {nqubits}. nqubits needs to be either 1 or 2.",
         )
 
-    # Check if invert_register has the correct string.
+    # Check if invert_register has the correct register(s).
+    valid_registers = [(0,), (1,), (0, 1)]
     if invert_register is not None:
-        if (
-            invert_register != "sp_0"
-            and invert_register != "sp_1"
-            and invert_register != "sp_t"
-        ):
+        if not isinstance(invert_register, tuple) or invert_register not in valid_registers:
             raise_error(
                 NameError,
-                f"{invert_register} not recognized. Input "
-                "sp_0"
-                " to reset qubit 0, "
-                "sp_1"
-                " to reset qubit 1, or "
-                "sp_t"
-                " to reset both qubits.",
+                f"{invert_register} not recognized.",
             )
 
+    # # Check if invert_register has the correct string.
+    # if invert_register is not None:
+    #     if (
+    #         invert_register != "sp_0"
+    #         and invert_register != "sp_1"
+    #         and invert_register != "sp_t"
+    #     ):
+    #         raise_error(
+    #             NameError,
+    #             f"{invert_register} not recognized. Input "
+    #             "sp_0"
+    #             " to reset qubit 0, "
+    #             "sp_1"
+    #             " to reset qubit 1, or "
+    #             "sp_t"
+    #             " to reset both qubits.",
+    #         )
+
     if backend is None:  # pragma: no cover
         backend = GlobalBackend()
 
@@ -177,12 +272,13 @@ def execute_GST(
         circ = prepare_states(k, nqubits)
         if invert_register is not None:
             inverted_circuit = reset_register(circ, invert_register)
-            if invert_register == "sp_0":
-                circ.add(inverted_circuit.on_qubits(0))
-            elif invert_register == "sp_1":
-                circ.add(inverted_circuit.on_qubits(1))
-            elif invert_register == "sp_t":
-                circ.add(inverted_circuit.on_qubits(0, 1))
+            # if invert_register == "sp_0":
+            #     circ.add(inverted_circuit.on_qubits(0))
+            # elif invert_register == "sp_1":
+            #     circ.add(inverted_circuit.on_qubits(1))
+            # elif invert_register == "sp_t":
+            #     circ.add(inverted_circuit.on_qubits(0, 1))
+            circ.add(inverted_circuit.on_qubits(*invert_register))
 
         if gate is not None:
             circ.add(gate)

tests/test_tomography_gate_set_tomography.py

@@ -9,6 +9,7 @@
     execute_GST,
     measurement_basis,
     prepare_states,
+    reset_register,
 )
 
 
@@ -151,6 +152,120 @@ def test_measurement_basis_invalid_j_valid_nqubits(j, nqubits):
         new_circuit = measurement_basis(j, test_circuit)
 
 
+def test_reset_register_valid_tuple_1qb():
+    # Test for valid tuple
+    nqubits = 1
+    test_circuit = qibo.models.Circuit(nqubits)
+    test_circuit.add(gates.H(0))
+
+    invert_register = (0,)
+    inverse_circuit = reset_register(test_circuit, invert_register)
+
+    correct_gates = [
+        [gates.H(0)],
+    ]
+    for groundtruth, gate in zip(correct_gates, inverse_circuit.queue):
+        assert isinstance(gate, type(groundtruth[0]))
+
+
+def test_reset_register0_twoqubitcircuit():
+    # Test resetting qubit 0
+    nqubits = 2
+    test_circuit = qibo.models.Circuit(nqubits)
+    test_circuit.add(gates.H(0))
+    test_circuit.add(gates.S(1))
+
+    inverse_circuit = reset_register(test_circuit, (0,))
+
+    correct_gates = [
+        [gates.H(0)],
+    ]
+
+    for groundtruth, gate in zip(correct_gates, inverse_circuit.queue):
+        assert isinstance(gate, type(groundtruth[0]))
+
+
+def test_reset_register0_singlequbitcircuit():
+    # Test resetting qubit 0
+
+    nqubits = 1
+    test_circuit = qibo.models.Circuit(nqubits)
+    test_circuit.add(gates.H(0))
+    test_circuit.add(gates.RX(0, np.pi / 3))
+
+    inverse_circuit = reset_register(test_circuit, (0,))
+
+    correct_gates = [
+        [gates.RX(0, np.pi / 3).dagger()],
+        [gates.H(0)],
+    ]
+
+    for groundtruth, gate in zip(correct_gates, inverse_circuit.queue):
+        assert isinstance(gate, type(groundtruth[0]))
+
+
+def test_reset_register1_twoqubitcircuit():
+    # Test resetting qubit 1
+    nqubits = 2
+    test_circuit = qibo.models.Circuit(nqubits)
+    test_circuit.add(gates.H(0))
+    test_circuit.add(gates.S(1))
+
+    inverse_circuit = reset_register(test_circuit, (1,))
+
+    correct_gates = [[gates.S(0).dagger()]]
+
+    for groundtruth, gate in zip(correct_gates, inverse_circuit.queue):
+        assert isinstance(gate, type(groundtruth[0]))
+
+
+def test_reset_register1_singlequbitcircuit():
+    # Test resetting qubit 1
+    nqubits = 2
+    test_circuit = qibo.models.Circuit(nqubits)
+    test_circuit.add(gates.H(0))
+    test_circuit.add(gates.S(0))
+    test_circuit.add(gates.RX(1, np.pi / 3))
+
+    inverse_circuit = reset_register(test_circuit, (1,))
+
+    correct_gates = [[gates.RX(1, np.pi / 3).dagger()]]
+
+    for groundtruth, gate in zip(correct_gates, inverse_circuit.queue):
+        assert isinstance(gate, type(groundtruth[0]))
+
+
+def test_reset_register_2():
+    # Test resetting both qubits
+    nqubits = 2
+    test_circuit = qibo.models.Circuit(nqubits)
+    test_circuit.add(gates.H(0))
+    test_circuit.add(gates.CNOT(0, 1))
+
+    inverse_circuit = reset_register(test_circuit, (0, 1))
+
+    correct_gates = [
+        [gates.CNOT(0, 1)],
+        [gates.H(0)],
+    ]
+    for groundtruth, gate in zip(correct_gates, inverse_circuit.queue):
+        assert isinstance(gate, type(groundtruth[0]))
+
+
+@pytest.mark.parametrize("a, b", [(0, 2), (1, 2), (2, 3)])
+def test_reset_register_invalid_tuple(a, b):
+    # Test resetting both qubits
+
+    nqubits = 2
+    test_circuit = qibo.models.Circuit(nqubits)
+    test_circuit.add(gates.H(0))
+    test_circuit.add(gates.CNOT(0, 1))
+
+    # Check if NameError is raised
+    with pytest.raises(NameError):
+        inverse_circuit = reset_register(test_circuit, (a, b))
+
+
 def test_GST_execute_circuit_1qb_j0():
     np.random.seed(42)
     nqubits = 1
@@ -277,6 +392,95 @@ def test_GST_two_qubit_with_CRXgate(backend):
     backend.assert_allclose(test_result, control_result, rtol=5e-2, atol=5e-2)
 
 
+def test_GST_one_qubit_with_gate_with_valid_reset_register_tuple(backend):
+    nqubits = 1
+    invert_register = (0,)
+    control_result = execute_GST(
+        nqubits=nqubits, gate=None, invert_register=invert_register
+    )
+    test_result = execute_GST(
+        nqubits=nqubits, gate=None, invert_register=invert_register
+    )
+    backend.assert_allclose(test_result, control_result, rtol=5e-2, atol=5e-2)
+
+
+def test_GST_two_qubit_with_gate_with_valid_reset_register_tuple(backend):
+    nqubits = 2
+    invert_register = (1,)
+    np.random.seed(42)
+    control_result = execute_GST(
+        nqubits=nqubits, gate=None, invert_register=invert_register, backend=backend
+    )
+    np.random.seed(42)
+    test_result = execute_GST(
+        nqubits=nqubits, gate=None, invert_register=invert_register, backend=backend
+    )
+
+    backend.assert_allclose(test_result, control_result, rtol=5e-2, atol=5e-2)
+
+
+def test_GST_one_qubit_with_param_gate_with_valid_reset_register_tuple(backend):
+    nqubits = 1
+    test_gate = gates.RX(0, np.pi / 7)
+    invert_register = (0,)
+    control_result = execute_GST(
+        nqubits=nqubits, gate=None, invert_register=invert_register
+    )
+    test_result = execute_GST(
+        nqubits=nqubits, gate=None, invert_register=invert_register
+    )
+    backend.assert_allclose(test_result, control_result, rtol=5e-2, atol=5e-2)
+
+
+def test_GST_two_qubit_with_param_gate_with_valid_reset_register_tuple(backend):
+    nqubits = 2
+    test_gate = gates.CNOT(0, 1)
+    invert_register = (1,)
+    np.random.seed(42)
+    control_result = execute_GST(
+        nqubits=nqubits, gate=None, invert_register=invert_register, backend=backend
+    )
+    np.random.seed(42)
+    test_result = execute_GST(
+        nqubits=nqubits, gate=None, invert_register=invert_register, backend=backend
+    )
+
+    backend.assert_allclose(test_result, control_result, rtol=5e-2, atol=5e-2)
+
+
+def test_GST_two_qubit_with_gate_with_valid_reset_register_tuple(backend):
+    nqubits = 2
+    test_gate = gates.CZ(0, 1)
+    invert_register = (0, 1)
+    np.random.seed(42)
+    control_result = execute_GST(
+        nqubits=nqubits,
+        gate=test_gate,
+        invert_register=invert_register,
+        backend=backend,
+    )
+    np.random.seed(42)
+    test_result = execute_GST(
+        nqubits=nqubits,
+        gate=test_gate,
+        invert_register=invert_register,
+        backend=backend,
+    )
+
+    backend.assert_allclose(test_result, control_result, rtol=5e-2, atol=5e-2)
+
+
+@pytest.mark.parametrize("a, b", [(0, 2), (1, 2), (2, 3)])
+def test_GST_two_qubit_with_gate_with_invalid_reset_register_tuple(a, b):
+    nqubits = 2
+    test_gate = gates.CZ(0, 1)
+    invert_register = (a, b)
+    with pytest.raises(NameError):
+        result = execute_GST(
+            nqubits=nqubits, gate=test_gate, invert_register=invert_register
+        )
+
+
 def test_GST_empty_circuit_with_invalid_qb(backend):
     nqubits = 3
     # Check if ValueError is raised