Move early stopping before batch estimation in `qec_util.performance.…

…sample_failures` (#32)

qec_util/performance/sampler.py

@@ -79,6 +79,9 @@ def sample_failures(
         # update the maximum limits based on the already calculated samples
         max_samples -= num_samples
         max_failures -= num_failures
+        # check if desired samples/failures have been reached
+        if (max_samples <= 0) or (max_failures <= 0):
+            return num_failures, num_samples
 
     # estimate the batch size for decoding
     defects, log_flips, _ = sampler.sample(shots=100)
@@ -98,7 +101,7 @@ def sample_failures(
             max_failures / log_err_prob if log_err_prob != 0 else np.inf,
         ]
     )
-    batch_size = estimated_max_samples / 10
+    batch_size = estimated_max_samples / 5
     batch_size = max([batch_size, 1])  # avoid batch_size = 0
     batch_size = batch_size if batch_size != np.inf else 200_000
     batch_size = int(batch_size)  # int(np.inf) raises an error

tests/performance/test_sampler.py

@@ -33,6 +33,31 @@ def test_sampler_to_file(tmp_path):
     return
 
 
+def test_sample_early_stopping(failures_file):
+    circuit = stim.Circuit.generated(
+        code_task="repetition_code:memory",
+        distance=3,
+        rounds=3,
+        after_clifford_depolarization=0.01,
+    )
+    dem = circuit.detector_error_model()
+    mwpm = Matching(dem)
+
+    num_failures, num_samples = sample_failures(
+        dem,
+        mwpm,
+        max_samples=1,
+        max_time=np.inf,
+        max_failures=1,
+        file_name=failures_file,
+    )
+
+    assert num_failures == 21
+    assert num_samples == 100
+
+    return
+
+
 def test_sampler_from_file(failures_file):
     circuit = stim.Circuit.generated(
         code_task="repetition_code:memory",