FfsimSampler: support global depolarizing noise (#325)

* FfsimSampler: support global depolarizing noise

* fix type

python/ffsim/qiskit/sampler.py

@@ -40,16 +40,21 @@ def __init__(
         self,
         *,
         default_shots: int = 1024,
+        global_depolarizing: float = 0.0,
         seed: np.random.Generator | int | None = None,
     ):
         """Initialize the ffsim sampler.
 
         Args:
             default_shots: The default shots to use if not specified during run.
+            global_depolarizing: Depolarizing probability for a noisy simulation.
+                Specifies the probability of sampling from the uniform distribution
+                instead of the state vector.
             seed: A seed to initialize the pseudorandom number generator.
                 Should be a valid input to ``np.random.default_rng``.
         """
         self._default_shots = default_shots
+        self._global_depolarizing = global_depolarizing
         self._rng = np.random.default_rng(seed)
 
     def run(
@@ -81,17 +86,28 @@ def _run_pub(self, pub: SamplerPub) -> SamplerPubResult:
                 norb, nelec = final_state.norb, final_state.nelec
                 if isinstance(nelec, int):
                     orbs = qargs
+                    n_qubits = len(orbs)
                 else:
                     orbs_a = [q for q in qargs if q < norb]
                     orbs_b = [q % norb for q in qargs if q >= norb]
                     orbs = (orbs_a, orbs_b)
-                samples_array = states.sample_state_vector(
+                    n_qubits = len(orbs_a) + len(orbs_b)
+                uniform_shots = self._rng.binomial(pub.shots, self._global_depolarizing)
+                exact_shots = pub.shots - uniform_shots
+                uniform_samples_array = self._rng.integers(
+                    2, size=(uniform_shots, n_qubits), dtype=bool
+                )
+                exact_samples_array = states.sample_state_vector(
                     final_state,
                     orbs=orbs,
-                    shots=pub.shots,
+                    shots=exact_shots,
                     bitstring_type=states.BitstringType.BIT_ARRAY,
                     seed=self._rng,
                 )
+                samples_array = np.concatenate(
+                    [uniform_samples_array, exact_samples_array]
+                )
+                self._rng.shuffle(samples_array)
             else:
                 samples_array = np.empty((pub.shots, 0), dtype=bool)
             for item in meas_info:

tests/python/qiskit/sampler_test.py

@@ -256,6 +256,50 @@ def test_measure_subset_spinless(norb: int, nocc: int):
     assert _fidelity(ffsim_probs, qiskit_probs) > 0.99
 
 
+@pytest.mark.parametrize("norb, nelec", [(5, (3, 2))])
+def test_global_depolarizing(norb: int, nelec: tuple[int, int]):
+    """Test sampler with global depolarizing noise."""
+    rng = np.random.default_rng(12285)
+
+    qubits = QuantumRegister(2 * norb)
+    orbital_rotation = ffsim.random.random_unitary(norb, seed=rng)
+    circuit = QuantumCircuit(qubits)
+    circuit.append(ffsim.qiskit.PrepareHartreeFockJW(norb, nelec), qubits)
+    circuit.append(ffsim.qiskit.OrbitalRotationJW(norb, orbital_rotation), qubits)
+    circuit.measure_all()
+
+    shots = 10_000
+    global_depolarizing = 0.1
+
+    sampler = ffsim.qiskit.FfsimSampler(
+        default_shots=shots, global_depolarizing=global_depolarizing, seed=rng
+    )
+    pub = (circuit,)
+    job = sampler.run([pub])
+    result = job.result()
+    pub_result = result[0]
+    samples = pub_result.data.meas.get_counts()
+
+    vec = ffsim.qiskit.final_state_vector(
+        circuit.remove_final_measurements(inplace=False)
+    ).vec
+    vec = ffsim.qiskit.ffsim_vec_to_qiskit_vec(vec, norb, nelec)
+    exact_probs = np.abs(vec) ** 2
+
+    strings, counts = zip(*samples.items())
+    assert sum(counts) == shots
+    addresses = [int(s, 2) for s in strings]
+    dim = 1 << 2 * norb
+
+    empirical_probs = np.zeros(dim, dtype=float)
+    empirical_probs[addresses] = np.array(counts) / shots
+    expected_probs = (1 - global_depolarizing) * exact_probs + global_depolarizing / dim
+
+    fidelity = np.sum(np.sqrt(exact_probs * empirical_probs))
+    expected_fidelity = np.sum(np.sqrt(exact_probs * expected_probs))
+    assert np.allclose(fidelity, expected_fidelity, rtol=1e-2)
+
+
 # TODO remove after removing UCJOperatorJW
 @pytest.mark.filterwarnings("ignore::DeprecationWarning")
 def test_reproducible_with_seed():