diff --git a/glue/sample/src/sinter/_decoding/_decoding.py b/glue/sample/src/sinter/_decoding/_decoding.py index e45aef72b..33416be0e 100644 --- a/glue/sample/src/sinter/_decoding/_decoding.py +++ b/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( diff --git a/glue/sample/src/sinter/_decoding/_decoding_decoder_class.py b/glue/sample/src/sinter/_decoding/_decoding_decoder_class.py index 1d13f4dfa..567ff28fb 100644 --- a/glue/sample/src/sinter/_decoding/_decoding_decoder_class.py +++ b/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) + diff --git a/glue/sample/src/sinter/_decoding/_decoding_test.py b/glue/sample/src/sinter/_decoding/_decoding_test.py index 7dd08f379..a923b0b60 100644 --- a/glue/sample/src/sinter/_decoding/_decoding_test.py +++ b/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) diff --git a/glue/sample/src/sinter/_decoding/_stim_then_decode_sampler.py b/glue/sample/src/sinter/_decoding/_stim_then_decode_sampler.py index e8c4a7e8c..c3ac2f2bb 100755 --- a/glue/sample/src/sinter/_decoding/_stim_then_decode_sampler.py +++ b/glue/sample/src/sinter/_decoding/_stim_then_decode_sampler.py @@ -115,22 +115,47 @@ def decode_shots_bit_packed( num_shots = bit_packed_detection_event_data.shape[0] with open(self.dets_b8_in_path, 'wb') as f: bit_packed_detection_event_data.tofile(f) - self.decoder.decode_via_files( - num_shots=num_shots, - num_obs=self.num_obs, - num_dets=self.num_dets, - dem_path=self.dem_path, - dets_b8_in_path=self.dets_b8_in_path, - obs_predictions_b8_out_path=self.obs_predictions_b8_out_path, - tmp_dir=self.decoder_tmp_dir, - ) + discards_b8_out_path = self.top_tmp_dir / 'discards.b8' + try: + self.decoder.decode_via_files( + num_shots=num_shots, + num_obs=self.num_obs, + num_dets=self.num_dets, + dem_path=self.dem_path, + dets_b8_in_path=self.dets_b8_in_path, + obs_predictions_b8_out_path=self.obs_predictions_b8_out_path, + tmp_dir=self.decoder_tmp_dir, + discards_b8_out_path=discards_b8_out_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_b8_out_path = None + self.decoder.decode_via_files( + num_shots=num_shots, + num_obs=self.num_obs, + num_dets=self.num_dets, + dem_path=self.dem_path, + dets_b8_in_path=self.dets_b8_in_path, + obs_predictions_b8_out_path=self.obs_predictions_b8_out_path, + tmp_dir=self.decoder_tmp_dir, + ) num_obs_bytes = (self.num_obs + 7) // 8 with open(self.obs_predictions_b8_out_path, 'rb') as f: prediction = np.fromfile(f, dtype=np.uint8, count=num_obs_bytes * num_shots) assert prediction.shape == (num_obs_bytes * num_shots,) self.obs_predictions_b8_out_path.unlink() self.dets_b8_in_path.unlink() - return prediction.reshape((num_shots, num_obs_bytes)) + prediction = prediction.reshape((num_shots, num_obs_bytes)) + if discards_b8_out_path is not None: + with open(discards_b8_out_path, 'rb') as f: + discards = np.fromfile(f, dtype=np.uint8, count=num_shots) + assert discards.shape == (num_shots,) + discards_b8_out_path.unlink() + prediction = np.concatenate([prediction, discards.reshape((num_shots, 1))], axis=1) + return prediction def _compile_decoder_with_disk_fallback( diff --git a/glue/sample/src/sinter/_decoding/_stim_then_decode_sampler_test.py b/glue/sample/src/sinter/_decoding/_stim_then_decode_sampler_test.py index 413015f0d..a72067547 100755 --- a/glue/sample/src/sinter/_decoding/_stim_then_decode_sampler_test.py +++ b/glue/sample/src/sinter/_decoding/_stim_then_decode_sampler_test.py @@ -190,3 +190,64 @@ def test_classify_discards_and_errors(): num_obs=13, ) == (0, 1) assert counter == collections.Counter(["obs_mistake_mask=_________E___"]) + + +def test_disk_decoder_reports_discards(): + # DiskDecoder (the compiled-wrapper for file-based decoders) must forward + # discards_b8_out_path to decode_via_files and return the discard bytes as + # a trailing column of prediction data. + import pathlib + import tempfile + + import stim + + import sinter + from sinter._decoding._stim_then_decode_sampler import DiskDecoder, StimThenDecodeSampler + + class DiscardingFileDecoder(sinter.Decoder): + 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.zeros(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[::2] = 1 + discards.tofile(discards_b8_out_path) + + circuit = stim.Circuit.generated( + "repetition_code:memory", + rounds=3, + distance=3, + before_round_data_depolarization=0, + before_measure_flip_probability=0, + after_clifford_depolarization=0, + ) + dem = circuit.detector_error_model() + task = sinter.Task(circuit=circuit, detector_error_model=dem) + num_det_bytes = -(-dem.num_detectors // 8) + num_obs_bytes = -(-dem.num_observables // 8) + + with tempfile.TemporaryDirectory() as d: + d = pathlib.Path(d) + + # The wrapper returns the discard bytes as a trailing column. + disk_decoder = DiskDecoder(DiscardingFileDecoder(), task, d) + dets = np.zeros((10, num_det_bytes), dtype=np.uint8) + result = disk_decoder.decode_shots_bit_packed(bit_packed_detection_event_data=dets) + assert result.shape == (10, num_obs_bytes + 1) + assert np.all(result[:, :-1] == 0) + assert np.count_nonzero(result[:, -1]) == 5 + + # The end-to-end sampler counts the discarded shots, and excludes them + # from error counting. + sampler = StimThenDecodeSampler( + decoder=DiscardingFileDecoder(), + count_observable_error_combos=False, + count_detection_events=False, + tmp_dir=d, + ) + stats = sampler.compiled_sampler_for_task(task).sample(max_shots=10) + assert stats.shots == 10 + assert stats.discards == 5 + assert stats.errors == 0