Unit tests for `simulate_one_shot_native_mcm (#6124)

Implement a unit test for `simulate_one_shot_native_mcm`

[sc-71560]

---------

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>

pennylane/devices/qubit/simulate.py

@@ -955,6 +955,8 @@ def simulate_one_shot_native_mcm(
 ) -> Result:
  """Simulate a single shot of a single quantum script with native mid-circuit measurements.
 
+ Assumes that the circuit has been transformed by `dynamic_one_shot`.
+
  Args:
  circuit (QuantumTape): The single circuit to simulate
  debugger (_Debugger): The debugger to use
@@ -968,8 +970,8 @@ def simulate_one_shot_native_mcm(
  keep the same number of shots. Default is ``None``.
 
  Returns:
- tuple(TensorLike): The results of the simulation
- dict: The mid-circuit measurement results of the simulation
+ Result: The results of the simulation
+
  """
  prng_key = execution_kwargs.pop("prng_key", None)
 

tests/devices/qubit/test_simulate.py

@@ -16,10 +16,12 @@
 import numpy as np
 import pytest
 from dummy_debugger import Debugger
+from flaky import flaky
+from stat_utils import fisher_exact_test
 
 import pennylane as qml
 from pennylane.devices.qubit import get_final_state, measure_final_state, simulate
-from pennylane.devices.qubit.simulate import _FlexShots
+from pennylane.devices.qubit.simulate import _FlexShots, simulate_one_shot_native_mcm
 
 
 class TestCurrentlyUnsupportedCases:
@@ -1178,3 +1180,74 @@ def test_qinfo_tf(self):
 
  grad5 = grad_tape.jacobian(results[5], phi)
  assert qml.math.allclose(grad5, expected_grads[5])
+
+
+ml_frameworks_list = [
+ "numpy",
+ pytest.param("autograd", marks=pytest.mark.autograd),
+ pytest.param("jax", marks=pytest.mark.jax),
+ pytest.param("torch", marks=pytest.mark.torch),
+ pytest.param("tensorflow", marks=pytest.mark.tf),
+]
+
+
+# pylint:disable=too-few-public-methods
+@pytest.mark.unit
+class TestMidCircuitMeasurements:
+ """Unit tests for simulating mid-circuit measurements."""
+
+ @flaky(max_runs=3, min_passes=2)
+ @pytest.mark.parametrize("ml_framework", ml_frameworks_list)
+ @pytest.mark.parametrize(
+ "postselect_mode", [None, "hw-like", "pad-invalid-samples", "fill-shots"]
+ )
+ def test_simulate_one_shot_native_mcm(self, ml_framework, postselect_mode):
+ """Unit tests for simulate_one_shot_native_mcm"""
+
+ with qml.queuing.AnnotatedQueue() as q:
+ qml.RX(np.pi / 4, wires=0)
+ m = qml.measure(wires=0, postselect=0)
+ qml.RX(np.pi / 4, wires=0)
+
+ circuit = qml.tape.QuantumScript(q.queue, [qml.expval(qml.Z(0)), qml.sample(m)], shots=[1])
+
+ n_shots = 200
+ results = [
+ simulate_one_shot_native_mcm(
+ circuit,
+ n_shots,
+ interface=ml_framework,
+ postselect_mode=postselect_mode,
+ )
+ for _ in range(n_shots)
+ ]
+ terminal_results, mcm_results = zip(*results)
+
+ if postselect_mode == "fill-shots":
+ assert all(ms == 0 for ms in mcm_results)
+ equivalent_tape = qml.tape.QuantumScript(
+ [qml.RX(np.pi / 4, wires=0)], [qml.expval(qml.Z(0))], shots=n_shots
+ )
+ expected_sample = simulate(equivalent_tape)
+ fisher_exact_test(terminal_results, expected_sample, outcomes=(-1, 1))
+
+ else:
+ equivalent_tape = qml.tape.QuantumScript(
+ [qml.RX(np.pi / 4, wires=0)], [qml.sample(wires=0)], shots=n_shots
+ )
+ expected_result = simulate(equivalent_tape)
+ fisher_exact_test(mcm_results, expected_result)
+
+ subset = [ts for ms, ts in zip(mcm_results, terminal_results) if ms == 0]
+ equivalent_tape = qml.tape.QuantumScript(
+ [qml.RX(np.pi / 4, wires=0)], [qml.expval(qml.Z(0))], shots=n_shots
+ )
+ expected_sample = simulate(equivalent_tape)
+ fisher_exact_test(subset, expected_sample, outcomes=(-1, 1))
+
+ subset = [ts for ms, ts in zip(mcm_results, terminal_results) if ms == 1]
+ equivalent_tape = qml.tape.QuantumScript(
+ [qml.X(0), qml.RX(np.pi / 4, wires=0)], [qml.expval(qml.Z(0))], shots=n_shots
+ )
+ expected_sample = simulate(equivalent_tape)
+ fisher_exact_test(subset, expected_sample, outcomes=(-1, 1))

tests/helpers/stat_utils.py

@@ -0,0 +1,32 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Helper functions for testing stochastic processes.
+"""
+import numpy as np
+from scipy.stats import fisher_exact
+
+
+def fisher_exact_test(actual, expected, outcomes=(0, 1), threshold=0.1):
+ """Checks that a binary sample matches the expected distribution using the Fisher exact test."""
+
+ actual, expected = np.asarray(actual), np.asarray(expected)
+ contingency_table = np.array(
+ [
+ [np.sum(actual == outcomes[0]), np.sum(actual == outcomes[1])],
+ [np.sum(expected == outcomes[0]), np.sum(expected == outcomes[1])],
+ ]
+ )
+ _, p_value = fisher_exact(contingency_table)
+ assert p_value > threshold, "The sample does not match the expected distribution."