diff --git a/src/braket/default_simulator/openqasm/program_context.py b/src/braket/default_simulator/openqasm/program_context.py index 0fa6964c..c9ea2408 100644 --- a/src/braket/default_simulator/openqasm/program_context.py +++ b/src/braket/default_simulator/openqasm/program_context.py @@ -1999,6 +1999,11 @@ def _branch_measurement( classical_targets, allow_remeasure=self.supports_midcircuit_measurement, ) + classical_targets_list = list(classical_targets) + for path_idx in self._active_path_indices: + path = self._paths[path_idx] + for qubit_idx, classical_idx in zip(target, classical_targets_list): + path._mcm_outcomes[classical_idx] = path._measurements[qubit_idx][-1] def _measure_and_branch(self, target: tuple[int]) -> None: """Sample outcomes per active path and branch with proportional shot diff --git a/src/braket/default_simulator/openqasm/simulation_path.py b/src/braket/default_simulator/openqasm/simulation_path.py index e3571dfb..253949b4 100644 --- a/src/braket/default_simulator/openqasm/simulation_path.py +++ b/src/braket/default_simulator/openqasm/simulation_path.py @@ -74,12 +74,14 @@ def __init__( variables: dict[str, FramedVariable] | None = None, measurements: dict[int, list[int]] | None = None, frame_number: int = 0, + mcm_outcomes: dict[int, int] | None = None, ): self._instructions = instructions if instructions is not None else [] self._shots = shots self._variables = variables if variables is not None else {} self._measurements = measurements if measurements is not None else {} self._frame_number = frame_number + self._mcm_outcomes = mcm_outcomes if mcm_outcomes is not None else {} @property def instructions(self) -> list[GateOperation]: @@ -122,6 +124,7 @@ def branch(self) -> SimulationPath: variables=deepcopy(self._variables), measurements=deepcopy(self._measurements), frame_number=self._frame_number, + mcm_outcomes=dict(self._mcm_outcomes), ) def enter_frame(self) -> int: diff --git a/src/braket/default_simulator/simulator.py b/src/braket/default_simulator/simulator.py index de61c738..f59f2993 100644 --- a/src/braket/default_simulator/simulator.py +++ b/src/braket/default_simulator/simulator.py @@ -917,6 +917,17 @@ def _run_branched( qubits_not_in_circuit = mapped_arr[~in_circuit_mask] measurements_array = np.array(measurements) selected = measurements_array[:, qubits_in_circuit] + # Overlay per-classical-bit MCM outcomes over the final-state + # samples so repeat measurements aren't reported as the qubit's + # final resampled state. + shot_offset = 0 + for path in context.active_paths: + mcm_outcomes = path._mcm_outcomes + for shot_idx in range(shot_offset, shot_offset + path.shots): + for col, classical_idx in enumerate(sorted(circuit.target_classical_indices)): + if classical_idx in mcm_outcomes: + selected[shot_idx, col] = str(mcm_outcomes[classical_idx]) + shot_offset += path.shots measurements = np.pad(selected, ((0, 0), (0, len(qubits_not_in_circuit)))).tolist() return GateModelTaskResult.construct( diff --git a/test/unit_tests/braket/default_simulator/test_mcm.py b/test/unit_tests/braket/default_simulator/test_mcm.py index 50080a83..e444f7fc 100644 --- a/test/unit_tests/braket/default_simulator/test_mcm.py +++ b/test/unit_tests/braket/default_simulator/test_mcm.py @@ -2634,13 +2634,12 @@ def test_multiple_measurements_and_branching(self, simulator): """ result = simulator.run_openqasm(OpenQASMProgram(source=qasm, inputs={}), shots=1000) assert len(result.measurements) == 1000 - # First two columns mirror q[0] (always 1 after the conditional flip); - # the third column samples q[1] and is 1 only when b[0]==1. + # 50% b[0]=0 → no q[1] flip → "010"; 50% b[0]=1 → q[1] flip → "111". counter = Counter(["".join(m) for m in result.measurements]) + assert "010" in counter assert "111" in counter - assert "110" in counter + assert 400 < counter["010"] < 600 assert 400 < counter["111"] < 600 - assert 400 < counter["110"] < 600 def test_complex_conditional_logic(self, simulator): """Complex conditional with if/else blocks.""" @@ -3292,7 +3291,7 @@ class TestMCMNonContiguousClassicalIndices: """MCMs that write non-contiguous classical indices.""" def test_minimal_unassigned_low_bit_stays_zero(self, simulator): - """Untouched ``q[0]`` stays 0 when only ``c[1]`` is assigned.""" + """``c[1]`` preserves its 50/50 MCM outcome; ``q[0]`` never touched.""" qasm = """ OPENQASM 3.0; qubit[2] q; @@ -3304,7 +3303,7 @@ def test_minimal_unassigned_low_bit_stays_zero(self, simulator): """ result = simulator.run_openqasm(OpenQASMProgram(source=qasm, inputs={}), shots=2000) counter = Counter(["".join(m) for m in result.measurements]) - assert counter == {"00": 2000} + assert set(counter) == {"00", "01"}, counter def test_if_branch_preserves_captured_bit_position(self, simulator): """``if (c[N])`` keeps the captured value at ``q[N]``'s position.""" @@ -5113,9 +5112,8 @@ def test_repeated_mcm_with_classical_feedforward(self): x q[1]; } """ - # After the conditional, q[0] is always |1>. b[1] is always 1, so q[1] - # only flips when b[0]==1: "10" and "11" each ~50%. - self._assert_distributions_match(qasm, expected_keys={"10", "11"}) + # b[0] captures the 50/50 MCM outcome; b[1] always 1: "01" and "11". + self._assert_distributions_match(qasm, expected_keys={"01", "11"}) def test_bell_pair_mcm_decoupling(self): """Bell-state MCM: measuring one half of an entangled pair must @@ -5196,3 +5194,45 @@ def test_sparse_qubit_register_does_not_blow_up(self): # Mirrors a verbatim ``measure_ff(q[13], 0); cc_prx(q[17], pi, 0, 0)`` # emission: q[17] flips iff b[0]==1 → outcomes "00" and "11" each ~50%. self._assert_distributions_match(qasm, expected_keys={"00", "11"}) + + +class TestBranchedResultRepeatedMeasurements: + """Repeated MCMs on the same qubit must report their historical outcomes.""" + + def test_repeated_measurement_of_same_qubit_reports_history(self): + qasm = """ + OPENQASM 3.0; + qubit[1] q; + bit[2] c; + x q[0]; + c[0] = measure q[0]; + reset q[0]; + c[1] = measure q[0]; + if (c[0]) { } + """ + result = StateVectorSimulator().run(OpenQASMProgram(source=qasm), shots=100) + keys = {"".join(m) for m in result.measurements} + assert keys == {"10"}, keys + + +class TestBranchedResultShotIndexing: + """The per-path overlay must map each path's outcomes to that path's shot slice.""" + + def test_asymmetric_path_shot_counts_preserve_alignment(self, simulator): + # Verifies that measurements explicitly assigned to a classical register and + # also end-of-circuit measurements are returned in the expected order. + # Circuit returns anti-correlated c[0] and c[1] with asymmetric probabilities. + qasm = """ + OPENQASM 3.0; + qubit[2] q; + bit[1] c; + ry(2.0943951023931953) q[0]; // 2π/3 → P(|1>) = 0.75 + c[0] = measure q[0]; + if (c[0] == 0) x q[1]; + measure q[1]; + """ + result = simulator.run_openqasm(OpenQASMProgram(source=qasm, inputs={}), shots=4000) + counter = Counter(["".join(m) for m in result.measurements]) + assert set(counter) == {"01", "10"}, counter + assert 800 < counter["01"] < 1200, counter + assert 2800 < counter["10"] < 3200, counter