Add discards_b8_out_path support to sinter file decoders

glue/sample/src/sinter/_decoding/_decoding.py

@@ -88,6 +88,7 @@ def _streaming_count_mistakes(
         predictions_in: pathlib.Path,
         count_detection_events: bool,
         count_observable_error_combos: bool,
+        discards_in: Optional[pathlib.Path] = None,
 ) -> Tuple[int, int, collections.Counter]:
 
     num_det_bytes = math.ceil(num_det / 8)
@@ -108,29 +109,36 @@ def _streaming_count_mistakes(
 
     with open(obs_in, 'rb') as obs_in_f:
         with open(predictions_in, 'rb') as predictions_in_f:
-            num_shots_left = num_shots
-            while num_shots_left:
-                batch_size = min(num_shots_left, math.ceil(10**6 / max(num_obs, 1)))
-
-                obs_batch = np.fromfile(obs_in_f, dtype=np.uint8, count=num_obs_bytes * batch_size)
-                pred_batch = np.fromfile(predictions_in_f, dtype=np.uint8, count=num_obs_bytes * batch_size)
-                obs_batch.shape = (batch_size, num_obs_bytes)
-                pred_batch.shape = (batch_size, num_obs_bytes)
-
-                cmp_table = pred_batch ^ obs_batch
-                err_mask = np.any(cmp_table, axis=1)
-                if postselected_observable_mask is not None:
-                    discard_mask = np.any(cmp_table & postselected_observable_mask, axis=1)
-                    err_mask &= ~discard_mask
-                    num_discards += np.count_nonzero(discard_mask)
-
-                if count_observable_error_combos:
-                    for misprediction_arr in cmp_table[err_mask]:
-                        err_key = "obs_mistake_mask=" + ''.join('_E'[b] for b in np.unpackbits(misprediction_arr, count=num_obs, bitorder='little'))
-                        custom_counts[err_key] += 1
-
-                num_errors += np.count_nonzero(err_mask)
-                num_shots_left -= batch_size
+            with (open(discards_in, 'rb') if discards_in is not None else contextlib.nullcontext()) as discards_in_f:
+                num_shots_left = num_shots
+                while num_shots_left:
+                    batch_size = min(num_shots_left, math.ceil(10**6 / max(num_obs, 1)))
+
+                    obs_batch = np.fromfile(obs_in_f, dtype=np.uint8, count=num_obs_bytes * batch_size)
+                    pred_batch = np.fromfile(predictions_in_f, dtype=np.uint8, count=num_obs_bytes * batch_size)
+                    obs_batch.shape = (batch_size, num_obs_bytes)
+                    pred_batch.shape = (batch_size, num_obs_bytes)
+
+                    cmp_table = pred_batch ^ obs_batch
+                    err_mask = np.any(cmp_table, axis=1)
+                    discard_mask = np.zeros(batch_size, dtype=bool)
+                    if discards_in_f is not None:
+                        disc_batch = np.fromfile(discards_in_f, dtype=np.uint8, count=batch_size)
+                        disc_batch.shape = (batch_size,)
+                        discard_mask |= disc_batch != 0
+                    if postselected_observable_mask is not None:
+                        discard_mask |= np.any(cmp_table & postselected_observable_mask, axis=1)
+                    if np.any(discard_mask):
+                        err_mask &= ~discard_mask
+                        num_discards += np.count_nonzero(discard_mask)
+
+                    if count_observable_error_combos:
+                        for misprediction_arr in cmp_table[err_mask]:
+                            err_key = "obs_mistake_mask=" + ''.join('_E'[b] for b in np.unpackbits(misprediction_arr, count=num_obs, bitorder='little'))
+                            custom_counts[err_key] += 1
+
+                    num_errors += np.count_nonzero(err_mask)
+                    num_shots_left -= batch_size
     return num_discards, num_errors, custom_counts
 
 
@@ -284,14 +292,26 @@ def _sample_decode_helper_using_memory(
         # Have the decoder predict which observables are flipped.
         predict_data = compiled_decoder.decode_shots_bit_packed(bit_packed_detection_event_data=dets_data)
 
+        # A trailing byte of prediction data marks shots the decoder had low
+        # confidence in; these shots are counted as discards.
+        num_obs_bytes = (num_obs + 7) // 8
+        if predict_data.shape[1] == num_obs_bytes + 1:
+            decoder_discarded_flags = predict_data[:, -1] != 0
+            predict_data = predict_data[:, :-1]
+        elif predict_data.shape[1] == num_obs_bytes:
+            decoder_discarded_flags = np.zeros(predict_data.shape[0], dtype=bool)
+        else:
+            raise ValueError(f"Got a numpy array with shape={predict_data.shape} from {type(compiled_decoder).__qualname__}.decode_shots_bit_packed(...). Expected shape={(predict_data.shape[0], num_obs_bytes)} or {(predict_data.shape[0], num_obs_bytes + 1)}.")
+
         # Discard any shots where the decoder predicts a flipped postselected observable.
+        discarded_flags = decoder_discarded_flags
         if postselected_observable_mask is not None:
-            discarded_flags = np.any(postselected_observable_mask & (predict_data ^ obs_data), axis=1)
-            cur_num_discarded_shots = np.count_nonzero(discarded_flags)
-            if cur_num_discarded_shots:
-                out_num_discards += cur_num_discarded_shots
-                obs_data = obs_data[~discarded_flags, :]
-                predict_data = predict_data[~discarded_flags, :]
+            discarded_flags = discarded_flags | np.any(postselected_observable_mask & (predict_data ^ obs_data), axis=1)
+        cur_num_discarded_shots = np.count_nonzero(discarded_flags)
+        if cur_num_discarded_shots:
+            out_num_discards += cur_num_discarded_shots
+            obs_data = obs_data[~discarded_flags, :]
+            predict_data = predict_data[~discarded_flags, :]
 
         # Count how many mistakes the decoder made on non-discarded shots.
         mispredictions = obs_data ^ predict_data
@@ -381,15 +401,34 @@ def _sample_decode_helper_using_disk(
         num_kept_shots = num_shots - num_det_discards
 
         # Perform syndrome decoding to predict observables from detection events.
-        decoder_obj.decode_via_files(
-            num_shots=num_kept_shots,
-            num_dets=num_dets,
-            num_obs=num_obs,
-            dem_path=dem_path,
-            dets_b8_in_path=dets_used_path,
-            obs_predictions_b8_out_path=predictions_path,
-            tmp_dir=tmp_dir,
-        )
+        discards_path = tmp_dir / 'sinter_discards.b8'
+        discards_available = True
+        try:
+            decoder_obj.decode_via_files(
+                num_shots=num_kept_shots,
+                num_dets=num_dets,
+                num_obs=num_obs,
+                dem_path=dem_path,
+                dets_b8_in_path=dets_used_path,
+                obs_predictions_b8_out_path=predictions_path,
+                tmp_dir=tmp_dir,
+                discards_b8_out_path=discards_path,
+            )
+        except TypeError:
+            # Decoders whose decode_via_files doesn't accept the new argument
+            # (e.g. older decoders, or C++-bound methods whose signatures are
+            # not introspectable) reject it during argument parsing, before any
+            # side effects, so it is safe to retry without it.
+            discards_available = False
+            decoder_obj.decode_via_files(
+                num_shots=num_kept_shots,
+                num_dets=num_dets,
+                num_obs=num_obs,
+                dem_path=dem_path,
+                dets_b8_in_path=dets_used_path,
+                obs_predictions_b8_out_path=predictions_path,
+                tmp_dir=tmp_dir,
+            )
 
         # Count how many predictions matched the actual observable data.
         num_obs_discards, num_errors, custom_counts = _streaming_count_mistakes(
@@ -402,6 +441,7 @@ def _sample_decode_helper_using_disk(
             postselected_observable_mask=postselected_observable_mask,
             count_detection_events=count_detection_events,
             count_observable_error_combos=count_observable_error_combos,
+            discards_in=discards_path if discards_available else None,
         )
 
         return AnonTaskStats(

glue/sample/src/sinter/_decoding/_decoding_decoder_class.py

@@ -1,5 +1,6 @@
 import abc
 import pathlib
+from typing import Optional
 
 import numpy as np
 import stim
@@ -51,6 +52,14 @@ def decode_shots_bit_packed(
             where `num_shots` is bit_packed_detection_event_data.shape[0] and
             `dem` is the detector error model this instance was compiled to
             decode.
+
+            The returned array may optionally contain one extra column of data
+            (shape `(num_shots, ceil(dem.num_observables / 8) + 1)`). When the
+            extra column is present, a nonzero value in the final column of a
+            row marks that shot as a "discard", meaning the decoder had low
+            confidence in its prediction for that shot. Sinter counts discarded
+            shots as conservative failures (`errors + discards`) when computing
+            logical error rates.
         """
         pass
 
@@ -108,6 +117,7 @@ def decode_via_files(self,
                          dets_b8_in_path: pathlib.Path,
                          obs_predictions_b8_out_path: pathlib.Path,
                          tmp_dir: pathlib.Path,
+                         discards_b8_out_path: Optional[pathlib.Path] = None,
                        ) -> None:
         """Performs decoding by reading/writing problems and answers from disk.
 
@@ -143,6 +153,16 @@ def decode_via_files(self,
                 process without warning, without giving it time to clean up any
                 temporary objects. All cleanup should be done via sinter
                 deleting this directory after killing the decoder.
+            discards_b8_out_path: If specified, the decoder must additionally
+                write one byte per shot to this file, in b8 format, where a
+                nonzero byte marks the corresponding shot as a discard (e.g.
+                because the decoder had low confidence in its prediction for
+                that shot). Sinter counts discarded shots as conservative
+                failures when computing logical error rates. If not specified,
+                the decoder should not write any discard data and all shots are
+                treated as not discarded. Decoders that do not support
+                reporting discards may ignore this parameter (sinter will only
+                pass it to decoders whose signature accepts it).
         """
         dem = stim.DetectorErrorModel.from_file(dem_path)
 
@@ -156,6 +176,13 @@ def decode_via_files(self,
         dets = np.fromfile(dets_b8_in_path, dtype=np.uint8, count=num_shots * num_det_bytes)
         dets = dets.reshape(num_shots, num_det_bytes)
         obs = compiled.decode_shots_bit_packed(bit_packed_detection_event_data=dets)
-        if obs.dtype != np.uint8 or obs.shape != (num_shots, num_obs_bytes):
-            raise ValueError(f"Got a numpy array with dtype={obs.dtype},shape={obs.shape} instead of dtype={np.uint8},shape={(num_shots, num_obs_bytes)} from {type(self).__qualname__}(...).compile_decoder_for_dem(...).decode_shots_bit_packed(...).")
+        if obs.dtype != np.uint8 or obs.shape not in ((num_shots, num_obs_bytes), (num_shots, num_obs_bytes + 1)):
+            raise ValueError(f"Got a numpy array with dtype={obs.dtype},shape={obs.shape} instead of dtype={np.uint8},shape={(num_shots, num_obs_bytes)} or {(num_shots, num_obs_bytes + 1)} from {type(self).__qualname__}(...).compile_decoder_for_dem(...).decode_shots_bit_packed(...).")
+        if obs.shape[1] > num_obs_bytes:
+            # The extra trailing byte marks shots as discards, as documented in
+            # `CompiledDecoder.decode_shots_bit_packed`.
+            if discards_b8_out_path is not None:
+                obs[:, -1:].tofile(discards_b8_out_path)
+            obs = obs[:, :num_obs_bytes]
         obs.tofile(obs_predictions_b8_out_path)
+

glue/sample/src/sinter/_decoding/_decoding_test.py

@@ -480,3 +480,145 @@ def compile_decoder_for_dem(
         )
         obs = np.fromfile(d / 'obs.b8', dtype=np.uint8, count=10)
         np.testing.assert_array_equal(obs, [0] * 10)
+
+
+class DiscardingFileDecoder(sinter.Decoder):
+    """A file-based decoder that predicts all observables as flipped and reports
+    every other shot as a low-confidence discard."""
+
+    def decode_via_files(self, *, num_shots, num_dets, num_obs, dem_path, dets_b8_in_path, obs_predictions_b8_out_path, tmp_dir, discards_b8_out_path=None):
+        num_det_bytes = -(-num_dets // 8)
+        num_obs_bytes = -(-num_obs // 8)
+        dets = np.fromfile(dets_b8_in_path, dtype=np.uint8, count=num_shots * num_det_bytes).reshape(num_shots, num_det_bytes)
+        np.ones(num_shots * num_obs_bytes, dtype=np.uint8).tofile(obs_predictions_b8_out_path)
+        if discards_b8_out_path is not None:
+            discards = np.zeros(num_shots, dtype=np.uint8)
+            discards[1::2] = 1
+            discards.tofile(discards_b8_out_path)
+
+
+class LegacyFileDecoder(sinter.Decoder):
+    """A file-based decoder with the old signature that doesn't support discards."""
+
+    def decode_via_files(self, *, num_shots, num_dets, num_obs, dem_path, dets_b8_in_path, obs_predictions_b8_out_path, tmp_dir):
+        num_obs_bytes = -(-num_obs // 8)
+        np.zeros(num_shots * num_obs_bytes, dtype=np.uint8).tofile(obs_predictions_b8_out_path)
+
+
+class TrailingByteCompiledDecoder(sinter.CompiledDecoder):
+    def __init__(self, num_obs: int):
+        self.num_obs_bytes = -(-num_obs // 8)
+
+    def decode_shots_bit_packed(
+            self,
+            *,
+            bit_packed_detection_event_data: np.ndarray,
+    ) -> np.ndarray:
+        num_shots = bit_packed_detection_event_data.shape[0]
+        return np.ones((num_shots, self.num_obs_bytes + 1), dtype=np.uint8)
+
+
+class TrailingByteDecoder(sinter.Decoder):
+    def compile_decoder_for_dem(
+        self,
+        *,
+        dem: stim.DetectorErrorModel,
+    ) -> CompiledDecoder:
+        return TrailingByteCompiledDecoder(num_obs=dem.num_observables)
+
+
+def _noiseless_repetition_code() -> stim.Circuit:
+    return stim.Circuit.generated(
+        "repetition_code:memory",
+        rounds=3,
+        distance=3,
+        before_round_data_depolarization=0,
+        before_measure_flip_probability=0,
+        after_clifford_depolarization=0,
+    )
+
+
+def test_file_decoder_discards_supported():
+    # A file-based decoder that supports discards_b8_out_path must have its
+    # low-confidence shots counted as discards and excluded from errors.
+    circuit = _noiseless_repetition_code()
+    result = sample_decode(
+        circuit_obj=circuit,
+        circuit_path=None,
+        dem_obj=circuit.detector_error_model(),
+        dem_path=None,
+        num_shots=10,
+        decoder="discarding_file",
+        custom_decoders={"discarding_file": DiscardingFileDecoder()},
+        __private__unstable__force_decode_on_disk=True,
+    )
+    assert result.shots == 10
+    assert result.discards == 5
+    assert result.errors == 5
+
+
+def test_file_decoder_discards_unsupported():
+    # A file-based decoder with the old signature (no discards_b8_out_path)
+    # must continue to work unchanged, with zero discards.
+    circuit = _noiseless_repetition_code()
+    result = sample_decode(
+        circuit_obj=circuit,
+        circuit_path=None,
+        dem_obj=circuit.detector_error_model(),
+        dem_path=None,
+        num_shots=10,
+        decoder="legacy_file",
+        custom_decoders={"legacy_file": LegacyFileDecoder()},
+        __private__unstable__force_decode_on_disk=True,
+    )
+    assert result.shots == 10
+    assert result.discards == 0
+    assert result.errors == 0
+
+
+def test_compiled_decoder_trailing_discard_byte():
+    # A compiled decoder that returns an extra trailing byte of prediction data
+    # must have its nonzero trailing bytes counted as discards, with those
+    # shots excluded from error counting (memory path).
+    circuit = _noiseless_repetition_code()
+    result = sample_decode(
+        circuit_obj=circuit,
+        circuit_path=None,
+        dem_obj=circuit.detector_error_model(),
+        dem_path=None,
+        num_shots=10,
+        decoder="trailing_byte",
+        custom_decoders={"trailing_byte": TrailingByteDecoder()},
+    )
+    assert result.shots == 10
+    assert result.discards == 10
+    assert result.errors == 0
+
+
+def test_base_class_decode_via_files_writes_discards():
+    # The base-class default implementation of decode_via_files must forward a
+    # trailing discard byte from the compiled decoder into discards_b8_out_path.
+    circuit = _noiseless_repetition_code()
+    dem = circuit.detector_error_model()
+    with tempfile.TemporaryDirectory() as d:
+        d = pathlib.Path(d)
+        dem.to_file(d / 'dem.dem')
+        num_det_bytes = -(-dem.num_detectors // 8)
+        num_obs_bytes = -(-dem.num_observables // 8)
+        np.zeros((10, num_det_bytes), dtype=np.uint8).tofile(d / 'dets.b8')
+        TrailingByteDecoder().decode_via_files(
+            num_shots=10,
+            num_dets=dem.num_detectors,
+            num_obs=dem.num_observables,
+            dem_path=d / 'dem.dem',
+            dets_b8_in_path=d / 'dets.b8',
+            obs_predictions_b8_out_path=d / 'obs.b8',
+            tmp_dir=d,
+            discards_b8_out_path=d / 'discards.b8',
+        )
+        obs = np.fromfile(d / 'obs.b8', dtype=np.uint8)
+        discards = np.fromfile(d / 'discards.b8', dtype=np.uint8)
+        assert obs.shape == (10 * num_obs_bytes,)
+        assert np.all(obs == 1)
+        assert discards.shape == (10,)
+        assert np.all(discards == 1)