fix: measure target qubits are required (#940)

* fix: measure target qubits are required

* removed commented out code

* fix test name that changed during merge conflict resolution

* convert non iterable targets to QubitSet, reformat conditions, add test for target_mapping to increase coverage

* move check for already measured target qubits to check_if_qubit_measured function

* simplify previously measured qubit error message

* pull in the latest default-simulator version

* Apply suggestions from code review

Co-authored-by: Cody Wang <[email protected]>

* Update src/braket/circuits/circuit.py

* Apply suggestions from code review

---------

Co-authored-by: Cody Wang <[email protected]>

setup.py

@@ -28,7 +28,7 @@
     package_dir={"": "src"},
     install_requires=[
         "amazon-braket-schemas>=1.21.3",
-        "amazon-braket-default-simulator>=1.21.2",
+        "amazon-braket-default-simulator>=1.21.4",
         "oqpy~=0.3.5",
         "backoff",
         "boltons",

src/braket/circuits/circuit.py

@@ -440,22 +440,17 @@ def _check_if_qubit_measured(
         Raises:
             ValueError: If adding a gate or noise operation after a measure instruction.
         """
-        if (
-            # check if there is a measure instruction on the target qubit
-            target
-            and target in self._measure_targets
-            # check if there is a measure instruction on any qubits in the target_mapping
-            or (target_mapping and any(targ in self._measure_targets for targ in target_mapping))
-            # If no target or target_mapping is supplied, check if there is a measure
-            # instruction on the current instructions target qubit
-            or (
-                instruction.target
-                and any(targ in self._measure_targets for targ in instruction.target)
-            )
-        ):
-            raise ValueError(
-                "cannot add a gate or noise operation on a qubit after a measure instruction."
+        if self._measure_targets:
+            measure_on_target_mapping = target_mapping and any(
+                targ in self._measure_targets for targ in target_mapping.values()
             )
+            if (
+                # check if there is a measure instruction on the targeted qubit(s)
+                measure_on_target_mapping
+                or any(tar in self._measure_targets for tar in QubitSet(target))
+                or any(tar in self._measure_targets for tar in QubitSet(instruction.target))
+            ):
+                raise ValueError("cannot apply instruction to measured qubits.")
 
     def add_instruction(
         self,
@@ -510,8 +505,7 @@ def add_instruction(
             raise TypeError("Only one of 'target_mapping' or 'target' can be supplied.")
 
         # Check if there is a measure instruction on the circuit
-        if not isinstance(instruction.operator, Measure) and self._measure_targets:
-            self._check_if_qubit_measured(instruction, target, target_mapping)
+        self._check_if_qubit_measured(instruction, target, target_mapping)
 
         if not target_mapping and not target:
             # Nothing has been supplied, add instruction
@@ -724,13 +718,12 @@ def _add_measure(self, target_qubits: QubitSetInput) -> None:
             else:
                 self._measure_targets = [target]
 
-    def measure(self, target_qubits: QubitSetInput | None = None) -> Circuit:
+    def measure(self, target_qubits: QubitSetInput) -> Circuit:
         """
         Add a `measure` operator to `self` ensuring only the target qubits are measured.
 
         Args:
-            target_qubits (QubitSetInput | None): target qubits to measure.
-                Default=None
+            target_qubits (QubitSetInput): target qubits to measure.
 
         Returns:
             Circuit: self
@@ -750,47 +743,21 @@ def measure(self, target_qubits: QubitSetInput | None = None) -> Circuit:
             Instruction('operator': H('qubit_count': 1), 'target': QubitSet([Qubit(2)]),
             Instruction('operator': Measure, 'target': QubitSet([Qubit(0)])]
         """
-        # check whether measuring an empty circuit
-        if not self.qubits:
-            raise IndexError("cannot measure an empty circuit.")
-
-        if isinstance(target_qubits, int):
-            target_qubits = [target_qubits]
+        if not isinstance(target_qubits, Iterable):
+            target_qubits = QubitSet(target_qubits)
 
         # Check if result types are added on the circuit
         if self.result_types:
             raise ValueError("a circuit cannot contain both measure instructions and result types.")
 
-        if target_qubits:
-            # Check if the target_qubits are already measured
-            if self._measure_targets and any(
-                target in self._measure_targets for target in target_qubits
-            ):
-                intersection = set(target_qubits) & set(self._measure_targets)
-                raise ValueError(
-                    f"cannot measure the same qubit(s) {', '.join(map(str, intersection))} "
-                    "more than once."
-                )
-            # Check if there are repeated qubits in the same measurement
-            if len(target_qubits) != len(set(target_qubits)):
-                intersection = [
-                    qubit for qubit, count in Counter(target_qubits).items() if count > 1
-                ]
-                raise ValueError(
-                    f"cannot repeat qubit(s) {', '.join(map(str, intersection))} "
-                    "in the same measurement."
-                )
-            self._add_measure(target_qubits=target_qubits)
-        else:
-            # Check if any qubits are already measured
-            if self._measure_targets:
-                intersection = set(self.qubits) & set(self._measure_targets)
-                raise ValueError(
-                    f"cannot measure the same qubit(s) {', '.join(map(str, intersection))} "
-                    "more than once."
-                )
-            # Measure all the qubits
-            self._add_measure(target_qubits=self.qubits)
+        # Check if there are repeated qubits in the same measurement
+        if len(target_qubits) != len(set(target_qubits)):
+            intersection = [qubit for qubit, count in Counter(target_qubits).items() if count > 1]
+            raise ValueError(
+                f"cannot repeat qubit(s) {', '.join(map(str, intersection))} "
+                "in the same measurement."
+            )
+        self._add_measure(target_qubits=target_qubits)
 
         return self
 
@@ -916,6 +883,9 @@ def apply_gate_noise(
         if not all(qubit in self.qubits for qubit in target_qubits):
             raise IndexError("target_qubits must be within the range of the current circuit.")
 
+        # Check if there is a measure instruction on the circuit
+        self._check_if_qubit_measured(instruction=noise, target=target_qubits)
+
         # make noise a list
         noise = wrap_with_list(noise)
 

src/braket/circuits/moments.py

@@ -257,13 +257,17 @@ def sort_moments(self) -> None:
         key_readout_noise = []
         moment_copy = OrderedDict()
         sorted_moment = OrderedDict()
+        last_measure = self._depth
 
         for key, instruction in self._moments.items():
             moment_copy[key] = instruction
             if key.moment_type == MomentType.READOUT_NOISE:
                 key_readout_noise.append(key)
             elif key.moment_type == MomentType.INITIALIZATION_NOISE:
                 key_initialization_noise.append(key)
+            elif key.moment_type == MomentType.MEASURE:
+                last_measure = key.time
+                key_noise.append(key)
             else:
                 key_noise.append(key)
 
@@ -272,7 +276,7 @@ def sort_moments(self) -> None:
         for key in key_noise:
             sorted_moment[key] = moment_copy[key]
         # find the max time in the circuit and make it the time for readout noise
-        max_time = max(self._depth - 1, 0)
+        max_time = max(last_measure - 1, 0)
 
         for key in key_readout_noise:
             sorted_moment[

test/unit_tests/braket/circuits/test_ascii_circuit_diagram.py

@@ -20,6 +20,7 @@
     FreeParameter,
     Gate,
     Instruction,
+    Noise,
     Observable,
     Operator,
 )
@@ -935,3 +936,23 @@ def test_measure_multiple_instructions_after():
         "T  : |0|1|2|3|4|5|6|",
     )
     _assert_correct_diagram(circ, expected)
+
+
+def test_measure_with_readout_noise():
+    circ = (
+        Circuit()
+        .h(0)
+        .cnot(0, 1)
+        .apply_readout_noise(Noise.BitFlip(probability=0.1), target_qubits=1)
+        .measure([0, 1])
+    )
+    expected = (
+        "T  : |0|    1    |2|",
+        "                    ",
+        "q0 : -H-C---------M-",
+        "        |           ",
+        "q1 : ---X-BF(0.1)-M-",
+        "",
+        "T  : |0|    1    |2|",
+    )
+    _assert_correct_diagram(circ, expected)

test/unit_tests/braket/circuits/test_circuit.py

@@ -24,6 +24,7 @@
     Gate,
     Instruction,
     Moments,
+    Noise,
     Observable,
     QubitSet,
     ResultType,
@@ -670,22 +671,26 @@ def test_measure_qubits_out_of_range():
 
 
 def test_measure_empty_circuit():
-    with pytest.raises(IndexError):
-        Circuit().measure()
-
-
-def test_measure_no_target():
-    circ = Circuit().h(0).cnot(0, 1).measure()
+    circ = Circuit().measure([0, 1, 2])
     expected = (
         Circuit()
-        .add_instruction(Instruction(Gate.H(), 0))
-        .add_instruction(Instruction(Gate.CNot(), [0, 1]))
         .add_instruction(Instruction(Measure(), 0))
         .add_instruction(Instruction(Measure(), 1))
+        .add_instruction(Instruction(Measure(), 2))
     )
     assert circ == expected
 
 
+def test_measure_target_input():
+    message = "Supplied qubit index, 1.1, must be an integer."
+    with pytest.raises(TypeError, match=message):
+        Circuit().h(0).cnot(0, 1).measure(1.1)
+
+    message = "Supplied qubit index, a, must be an integer."
+    with pytest.raises(TypeError, match=message):
+        Circuit().h(0).cnot(0, 1).measure(FreeParameter("a"))
+
+
 def test_measure_with_result_types():
     message = "a circuit cannot contain both measure instructions and result types."
     with pytest.raises(ValueError, match=message):
@@ -713,13 +718,15 @@ def test_measure_with_multiple_measures():
 
 
 def test_measure_same_qubit_twice():
-    message = "cannot measure the same qubit\\(s\\) 0 more than once."
+    # message = "cannot measure the same qubit\\(s\\) Qubit\\(0\\) more than once."
+    message = "cannot apply instruction to measured qubits."
     with pytest.raises(ValueError, match=message):
         Circuit().h(0).cnot(0, 1).measure(0).measure(1).measure(0)
 
 
 def test_measure_same_qubit_twice_with_list():
-    message = "cannot measure the same qubit\\(s\\) 0 more than once."
+    # message = "cannot measure the same qubit\\(s\\) Qubit\\(0\\) more than once."
+    message = "cannot apply instruction to measured qubits."
     with pytest.raises(ValueError, match=message):
         Circuit().h(0).cnot(0, 1).measure(0).measure([0, 1])
 
@@ -730,20 +737,56 @@ def test_measure_same_qubit_twice_with_one_measure():
         Circuit().h(0).cnot(0, 1).measure([0, 0, 0])
 
 
-def test_measure_empty_measure_after_measure_with_targets():
-    message = "cannot measure the same qubit\\(s\\) 0, 1 more than once."
+def test_measure_gate_after():
+    # message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    message = "cannot apply instruction to measured qubits."
     with pytest.raises(ValueError, match=message):
-        Circuit().h(0).cnot(0, 1).cnot(1, 2).measure(0).measure(1).measure()
+        Circuit().h(0).measure(0).h([0, 1])
 
+    # message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    message = "cannot apply instruction to measured qubits."
+    with pytest.raises(ValueError, match=message):
+        instr = Instruction(Gate.CNot(), [0, 1])
+        Circuit().measure([0, 1]).add_instruction(instr, target_mapping={0: 0, 1: 1})
 
-def test_measure_gate_after():
-    message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    # message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    message = "cannot apply instruction to measured qubits."
     with pytest.raises(ValueError, match=message):
-        Circuit().h(0).measure(0).h([0, 1])
+        instr = Instruction(Gate.CNot(), [0, 1])
+        Circuit().h(0).measure(0).add_instruction(instr, target=[0, 1])
+
+
+def test_measure_noise_after():
+    # message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    message = "cannot apply instruction to measured qubits."
+    with pytest.raises(ValueError, match=message):
+        Circuit().h(1).h(1).h(2).h(5).h(4).h(3).cnot(1, 2).measure([0, 1, 2, 3, 4]).kraus(
+            targets=[0], matrices=[np.array([[1, 0], [0, 1]])]
+        )
+
+
+def test_measure_with_readout_noise():
+    circ = (
+        Circuit()
+        .h(0)
+        .cnot(0, 1)
+        .apply_readout_noise(Noise.BitFlip(probability=0.1), target_qubits=1)
+        .measure([0, 1])
+    )
+    expected = (
+        Circuit()
+        .add_instruction(Instruction(Gate.H(), 0))
+        .add_instruction(Instruction(Gate.CNot(), [0, 1]))
+        .apply_readout_noise(Noise.BitFlip(probability=0.1), target_qubits=1)
+        .add_instruction(Instruction(Measure(), 0))
+        .add_instruction(Instruction(Measure(), 1))
+    )
+    assert circ == expected
 
 
 def test_measure_gate_after_with_target_mapping():
-    message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    # message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    message = "cannot apply instruction to measured qubits."
     instr = Instruction(Gate.CNot(), [0, 1])
     with pytest.raises(ValueError, match=message):
         Circuit().h(0).cnot(0, 1).cnot(1, 2).measure([0, 1]).add_instruction(
@@ -752,7 +795,8 @@ def test_measure_gate_after_with_target_mapping():
 
 
 def test_measure_gate_after_with_target():
-    message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    # message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    message = "cannot apply instruction to measured qubits."
     instr = Instruction(Gate.CNot(), [0, 1])
     with pytest.raises(ValueError, match=message):
         Circuit().h(0).cnot(0, 1).cnot(1, 2).measure([0, 1]).add_instruction(instr, target=[10, 11])

test/unit_tests/braket/circuits/test_unicode_circuit_diagram.py

@@ -19,6 +19,7 @@
     FreeParameter,
     Gate,
     Instruction,
+    Noise,
     Observable,
     Operator,
     UnicodeCircuitDiagram,
@@ -1097,3 +1098,24 @@ def test_measure_multiple_instructions_after():
         "T  : │  0  │  1  │  2  │  3  │  4  │  5  │  6  │",
     )
     _assert_correct_diagram(circ, expected)
+
+
+def test_measure_with_readout_noise():
+    circ = (
+        Circuit()
+        .h(0)
+        .cnot(0, 1)
+        .apply_readout_noise(Noise.BitFlip(probability=0.1), target_qubits=1)
+        .measure([0, 1])
+    )
+    expected = (
+        "T  : │  0  │        1        │  2  │",
+        "      ┌───┐                   ┌───┐ ",
+        "q0 : ─┤ H ├───●───────────────┤ M ├─",
+        "      └───┘   │               └───┘ ",
+        "            ┌─┴─┐ ┌─────────┐ ┌───┐ ",
+        "q1 : ───────┤ X ├─┤ BF(0.1) ├─┤ M ├─",
+        "            └───┘ └─────────┘ └───┘ ",
+        "T  : │  0  │        1        │  2  │",
+    )
+    _assert_correct_diagram(circ, expected)