fix: add observable targets not in instructions to circuit qubit coun…

…t and qubits (#167)

src/braket/circuits/circuit.py

@@ -111,6 +111,7 @@ def __init__(self, addable: AddableTypes = None, *args, **kwargs):
         self._result_types: List[ResultType] = []
         self._qubit_observable_mapping: Dict[Union[int, Circuit._ALL_QUBITS], Observable] = {}
         self._qubit_target_mapping: Dict[int, List[int]] = {}
+        self._qubit_observable_set = set()
 
         if addable is not None:
             self.add(addable, *args, **kwargs)
@@ -179,18 +180,19 @@ def _observable_to_instruction(observable: Observable, target_list: List[int]):
 
     @property
     def moments(self) -> Moments:
-        """Moments: Get the `moments` for this circuit."""
+        """Moments: Get the `moments` for this circuit. Note that this includes observables."""
         return self._moments
 
     @property
     def qubit_count(self) -> int:
-        """Get the qubit count for this circuit."""
-        return self._moments.qubit_count
+        """Get the qubit count for this circuit. Note that this includes observables."""
+        all_qubits = self._moments.qubits.union(self._qubit_observable_set)
+        return len(all_qubits)
 
     @property
     def qubits(self) -> QubitSet:
         """QubitSet: Get a copy of the qubits for this circuit."""
-        return QubitSet(self._moments.qubits)
+        return QubitSet(self._moments.qubits.union(self._qubit_observable_set))
 
     def add_result_type(
         self,
@@ -260,8 +262,9 @@ def add_result_type(
             result_type_to_add = result_type.copy(target=target)
 
         if result_type_to_add not in self._result_types:
-            self._add_to_qubit_observable_mapping(result_type)
+            self._add_to_qubit_observable_mapping(result_type_to_add)
             self._result_types.append(result_type_to_add)
+            self._add_to_qubit_observable_set(result_type_to_add)
         return self
 
     def _add_to_qubit_observable_mapping(self, result_type: ResultType) -> None:
@@ -298,6 +301,10 @@ def _add_to_qubit_observable_mapping(self, result_type: ResultType) -> None:
         if not result_type.target:
             self._qubit_observable_mapping[Circuit._ALL_QUBITS] = observable
 
+    def _add_to_qubit_observable_set(self, result_type: ResultType) -> None:
+        if isinstance(result_type, ObservableResultType) and result_type.target:
+            self._qubit_observable_set.update(result_type.target)
+
     def add_instruction(
         self,
         instruction: Instruction,

test/integ_tests/gate_model_device_testing_utils.py

@@ -52,6 +52,21 @@ def no_result_types_bell_pair_testing(device: Device, run_kwargs: Dict[str, Any]
     assert len(result.measurements) == shots
 
 
+def result_types_observable_not_in_instructions(device: Device, run_kwargs: Dict[str, Any]):
+    shots = run_kwargs["shots"]
+    tol = get_tol(shots)
+    bell = (
+        Circuit()
+        .h(0)
+        .cnot(0, 1)
+        .expectation(observable=Observable.X(), target=[2])
+        .variance(observable=Observable.Y(), target=[3])
+    )
+    result = device.run(bell, **run_kwargs).result()
+    assert np.allclose(result.values[0], 0, **tol)
+    assert np.allclose(result.values[1], 1, **tol)
+
+
 def result_types_zero_shots_bell_pair_testing(
     device: Device, include_state_vector: bool, run_kwargs: Dict[str, Any]
 ):

test/integ_tests/test_local_braket_simulator.py

@@ -21,6 +21,7 @@
     result_types_bell_pair_marginal_probability_testing,
     result_types_hermitian_testing,
     result_types_nonzero_shots_bell_pair_testing,
+    result_types_observable_not_in_instructions,
     result_types_tensor_hermitian_hermitian_testing,
     result_types_tensor_x_y_testing,
     result_types_tensor_y_hermitian_testing,
@@ -104,3 +105,8 @@ def test_result_types_tensor_y_hermitian(shots):
 @pytest.mark.parametrize("shots", [0, SHOTS])
 def test_result_types_all_selected(shots):
     result_types_all_selected_testing(DEVICE, {"shots": shots})
+
+
+@pytest.mark.parametrize("shots", [0, SHOTS])
+def test_result_types_observable_not_in_instructions(shots):
+    result_types_observable_not_in_instructions(DEVICE, {"shots": shots})

test/integ_tests/test_simulator_quantum_task.py

@@ -21,6 +21,7 @@
     result_types_bell_pair_marginal_probability_testing,
     result_types_hermitian_testing,
     result_types_nonzero_shots_bell_pair_testing,
+    result_types_observable_not_in_instructions,
     result_types_tensor_hermitian_hermitian_testing,
     result_types_tensor_x_y_testing,
     result_types_tensor_y_hermitian_testing,
@@ -156,3 +157,13 @@ def test_result_types_all_selected(simulator_arn, shots, aws_session, s3_destina
     result_types_all_selected_testing(
         device, {"shots": shots, "s3_destination_folder": s3_destination_folder}
     )
+
+
+@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS)])
+def test_result_types_observable_not_in_instructions(
+    simulator_arn, shots, aws_session, s3_destination_folder
+):
+    device = AwsDevice(simulator_arn, aws_session)
+    result_types_observable_not_in_instructions(
+        device, {"shots": shots, "s3_destination_folder": s3_destination_folder}
+    )

test/unit_tests/braket/circuits/test_circuit.py

@@ -608,6 +608,64 @@ def test_qubit_count_setter(h):
     h.qubit_count = 1
 
 
+@pytest.mark.parametrize(
+    "circuit,expected_qubit_count",
+    [
+        (Circuit().h(0).h(1).h(2), 3),
+        (
+            Circuit()
+            .h(0)
+            .expectation(observable=Observable.H() @ Observable.X(), target=[0, 1])
+            .sample(observable=Observable.H() @ Observable.X(), target=[0, 1]),
+            2,
+        ),
+        (
+            Circuit().h(0).probability([1, 2]).state_vector(),
+            1,
+        ),
+        (
+            Circuit()
+            .h(0)
+            .variance(observable=Observable.H(), target=1)
+            .state_vector()
+            .amplitude(["01"]),
+            2,
+        ),
+    ],
+)
+def test_qubit_count(circuit, expected_qubit_count):
+    assert circuit.qubit_count == expected_qubit_count
+
+
+@pytest.mark.parametrize(
+    "circuit,expected_qubits",
+    [
+        (Circuit().h(0).h(1).h(2), QubitSet([0, 1, 2])),
+        (
+            Circuit()
+            .h(0)
+            .expectation(observable=Observable.H() @ Observable.X(), target=[0, 1])
+            .sample(observable=Observable.H() @ Observable.X(), target=[0, 1]),
+            QubitSet([0, 1]),
+        ),
+        (
+            Circuit().h(0).probability([1, 2]).state_vector(),
+            QubitSet([0]),
+        ),
+        (
+            Circuit()
+            .h(0)
+            .variance(observable=Observable.H(), target=1)
+            .state_vector()
+            .amplitude(["01"]),
+            QubitSet([0, 1]),
+        ),
+    ],
+)
+def test_circuit_qubits(circuit, expected_qubits):
+    assert circuit.qubits == expected_qubits
+
+
 def test_qubits_getter(h):
     assert h.qubits == h._moments.qubits
     assert h.qubits is not h._moments.qubits