diff --git a/doc/python_api_reference_vDev.md b/doc/python_api_reference_vDev.md index 01ce7fe4b..243c4dd7c 100644 --- a/doc/python_api_reference_vDev.md +++ b/doc/python_api_reference_vDev.md @@ -6124,10 +6124,9 @@ def sample( self, shots: int, *, - prepend_observables: bool = False, - append_observables: bool = False, separate_observables: bool = False, bit_packed: bool = False, + transposed: bool = False, dets_out: Optional[np.ndarray] = None, obs_out: Optional[np.ndarray] = None, ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: @@ -6142,12 +6141,16 @@ def sample( separate_observables: Defaults to False. When set to True, the return value is a (detection_events, observable_flips) tuple instead of a flat detection_events array. - prepend_observables: Defaults to false. When set, observables are included - with the detectors and are placed at the start of the results. - append_observables: Defaults to false. When set, observables are included - with the detectors and are placed at the end of the results. bit_packed: Returns a uint8 numpy array with 8 bits per byte, instead of a bool_ numpy array with 1 bit per byte. Uses little endian packing. + transposed: Defaults to false. When set to true, outputs are transposed + (i.e. the first axis is detectors/observables and the second axis is + shots instead of vice versa). + + Using transposed=True has better performance than transposed=False, + because the output more closely corresponds to stim's internal + representation. But data in this transposed form is often more difficult + to work with. dets_out: Defaults to None. Specifies a pre-allocated numpy array to write the detection event data into. This array must have the correct shape and dtype. @@ -6155,65 +6158,100 @@ def sample( the observable flip data into. This array must have the correct shape and dtype. - Returns: - A numpy array or tuple of numpy arrays containing the samples. - - if separate_observables=False and bit_packed=False: - A single numpy array. - dtype=bool_ - shape=( - shots, - num_detectors + num_observables * ( - append_observables + prepend_observables), - ) - The bit for detection event `m` in shot `s` is at - result[s, m] + Deprecated Arguments: + (These arguments still work but are hidden and not recommended.) + + prepend_observables: Prefer `separate_observables`. Defaults to false. + When set, observables are included with the detectors and are placed at + the start of the results. Using this argument is generally a bad idea, + compared to using `separate_observables=True`, because you will + inevitably just end up immediately separating the observable data + from the detector data anyways. + append_observables: Prefer `separate_observables`. Defaults to false. + When set, observables are included with the detectors and are placed at + the end of the results. Using this argument is generally a bad idea, + compared to using `separate_observables=True`, because you will + inevitably just end up immediately separating the observable data + from the detector data anyways. - if separate_observables=False and bit_packed=True: - A single numpy array. - dtype=uint8 - shape=( - shots, - math.ceil((num_detectors + num_observables * ( - append_observables + prepend_observables)) / 8), - ) - The bit for detection event `m` in shot `s` is at - (result[s, m // 8] >> (m % 8)) & 1 + Returns: + If separate_observables: + A tuple `(dets, obs)` of two numpy arrays. + else: + A single numpy array `dets`. - if separate_observables=True and bit_packed=False: - A (dets, obs) tuple. + if bit_packed=False and transposed=False: dets.dtype=bool_ dets.shape=(shots, num_detectors) + Detection event `m` in shot `s` is dets[s, m] + if bit_packed=False and transposed=True: + dets.dtype=bool_ + dets.shape=(num_detectors, shots) + Detection event `m` in shot `s` is dets[m, s] + if bit_packed=True and transposed=False: + dets.dtype=uint8 + dets.shape=(shots, (num_detectors + 7) // 8) + Detection event `m` in shot `s` is `(dets[s, m // 8] >> (m % 8)) & 1` + if bit_packed=True and transposed=True: + dets.dtype=uint8 + dets.shape=(num_detectors, (shots + 7) // 8) + Detection event `m` in shot `s` is `(dets[m, s // 8] >> (s % 8)) & 1` + + if separate_observables and bit_packed=False and transposed=False: obs.dtype=bool_ obs.shape=(shots, num_observables) - The bit for detection event `m` in shot `s` is at - dets[s, m] - The bit for observable `m` in shot `s` is at - obs[s, m] - - if separate_observables=True and bit_packed=True: - A (dets, obs) tuple. - dets.dtype=uint8 - dets.shape=(shots, math.ceil(num_detectors / 8)) + Observable `m` in shot `s` is obs[s, m] + if separate_observables and bit_packed=False and transposed=True: + obs.dtype=bool_ + obs.shape=(num_observables, shots) + Observable `m` in shot `s` is obs[m, s] + if separate_observables and bit_packed=True and transposed=False: obs.dtype=uint8 - obs.shape=(shots, math.ceil(num_observables / 8)) - The bit for detection event `m` in shot `s` is at - (dets[s, m // 8] >> (m % 8)) & 1 - The bit for observable `m` in shot `s` is at - (obs[s, m // 8] >> (m % 8)) & 1 + obs.shape=(shots, (num_observables + 7) // 8) + Observable `m` in shot `s` is `(obs[s, m // 8] >> (m % 8)) & 1` + if separate_observables and bit_packed=True and transposed=True: + obs.dtype=uint8 + obs.shape=(num_observables, (shots + 7) // 8) + Observable `m` in shot `s` is `(obs[m, s // 8] >> (s % 8)) & 1` Examples: >>> import stim >>> c = stim.Circuit(''' - ... H 0 - ... CNOT 0 1 + ... R 0 1 2 + ... CX 0 1 2 1 ... X_ERROR(1.0) 0 - ... M 0 1 - ... DETECTOR rec[-1] rec[-2] + ... M 0 1 2 + ... DETECTOR rec[-2] + ... DETECTOR rec[-1] rec[-2] rec[-3] + ... OBSERVABLE_INCLUDE(0) rec[-1] ... ''') >>> s = c.compile_detector_sampler() - >>> s.sample(shots=1) - array([[ True]]) + >>> dets, obs = s.sample(shots=5, separate_observables=True) + >>> dets + array([[False, True], + [False, True], + [False, True], + [False, True], + [False, True]]) + >>> obs + array([[False], + [False], + [False], + [False], + [False]]) + + >>> dets, obs = s.sample( + ... shots=90, + ... bit_packed=True, + ... transposed=True, + ... separate_observables=True, + ... ) + >>> dets + array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 3]], + dtype=uint8) + >>> obs + array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8) """ ``` diff --git a/doc/stim.pyi b/doc/stim.pyi index d71851394..195a124a8 100644 --- a/doc/stim.pyi +++ b/doc/stim.pyi @@ -4763,10 +4763,9 @@ class CompiledDetectorSampler: self, shots: int, *, - prepend_observables: bool = False, - append_observables: bool = False, separate_observables: bool = False, bit_packed: bool = False, + transposed: bool = False, dets_out: Optional[np.ndarray] = None, obs_out: Optional[np.ndarray] = None, ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: @@ -4781,12 +4780,16 @@ class CompiledDetectorSampler: separate_observables: Defaults to False. When set to True, the return value is a (detection_events, observable_flips) tuple instead of a flat detection_events array. - prepend_observables: Defaults to false. When set, observables are included - with the detectors and are placed at the start of the results. - append_observables: Defaults to false. When set, observables are included - with the detectors and are placed at the end of the results. bit_packed: Returns a uint8 numpy array with 8 bits per byte, instead of a bool_ numpy array with 1 bit per byte. Uses little endian packing. + transposed: Defaults to false. When set to true, outputs are transposed + (i.e. the first axis is detectors/observables and the second axis is + shots instead of vice versa). + + Using transposed=True has better performance than transposed=False, + because the output more closely corresponds to stim's internal + representation. But data in this transposed form is often more difficult + to work with. dets_out: Defaults to None. Specifies a pre-allocated numpy array to write the detection event data into. This array must have the correct shape and dtype. @@ -4794,65 +4797,100 @@ class CompiledDetectorSampler: the observable flip data into. This array must have the correct shape and dtype. - Returns: - A numpy array or tuple of numpy arrays containing the samples. - - if separate_observables=False and bit_packed=False: - A single numpy array. - dtype=bool_ - shape=( - shots, - num_detectors + num_observables * ( - append_observables + prepend_observables), - ) - The bit for detection event `m` in shot `s` is at - result[s, m] - - if separate_observables=False and bit_packed=True: - A single numpy array. - dtype=uint8 - shape=( - shots, - math.ceil((num_detectors + num_observables * ( - append_observables + prepend_observables)) / 8), - ) - The bit for detection event `m` in shot `s` is at - (result[s, m // 8] >> (m % 8)) & 1 + Deprecated Arguments: + (These arguments still work but are hidden and not recommended.) + + prepend_observables: Prefer `separate_observables`. Defaults to false. + When set, observables are included with the detectors and are placed at + the start of the results. Using this argument is generally a bad idea, + compared to using `separate_observables=True`, because you will + inevitably just end up immediately separating the observable data + from the detector data anyways. + append_observables: Prefer `separate_observables`. Defaults to false. + When set, observables are included with the detectors and are placed at + the end of the results. Using this argument is generally a bad idea, + compared to using `separate_observables=True`, because you will + inevitably just end up immediately separating the observable data + from the detector data anyways. + + Returns: + If separate_observables: + A tuple `(dets, obs)` of two numpy arrays. + else: + A single numpy array `dets`. - if separate_observables=True and bit_packed=False: - A (dets, obs) tuple. + if bit_packed=False and transposed=False: dets.dtype=bool_ dets.shape=(shots, num_detectors) + Detection event `m` in shot `s` is dets[s, m] + if bit_packed=False and transposed=True: + dets.dtype=bool_ + dets.shape=(num_detectors, shots) + Detection event `m` in shot `s` is dets[m, s] + if bit_packed=True and transposed=False: + dets.dtype=uint8 + dets.shape=(shots, (num_detectors + 7) // 8) + Detection event `m` in shot `s` is `(dets[s, m // 8] >> (m % 8)) & 1` + if bit_packed=True and transposed=True: + dets.dtype=uint8 + dets.shape=(num_detectors, (shots + 7) // 8) + Detection event `m` in shot `s` is `(dets[m, s // 8] >> (s % 8)) & 1` + + if separate_observables and bit_packed=False and transposed=False: obs.dtype=bool_ obs.shape=(shots, num_observables) - The bit for detection event `m` in shot `s` is at - dets[s, m] - The bit for observable `m` in shot `s` is at - obs[s, m] - - if separate_observables=True and bit_packed=True: - A (dets, obs) tuple. - dets.dtype=uint8 - dets.shape=(shots, math.ceil(num_detectors / 8)) + Observable `m` in shot `s` is obs[s, m] + if separate_observables and bit_packed=False and transposed=True: + obs.dtype=bool_ + obs.shape=(num_observables, shots) + Observable `m` in shot `s` is obs[m, s] + if separate_observables and bit_packed=True and transposed=False: obs.dtype=uint8 - obs.shape=(shots, math.ceil(num_observables / 8)) - The bit for detection event `m` in shot `s` is at - (dets[s, m // 8] >> (m % 8)) & 1 - The bit for observable `m` in shot `s` is at - (obs[s, m // 8] >> (m % 8)) & 1 + obs.shape=(shots, (num_observables + 7) // 8) + Observable `m` in shot `s` is `(obs[s, m // 8] >> (m % 8)) & 1` + if separate_observables and bit_packed=True and transposed=True: + obs.dtype=uint8 + obs.shape=(num_observables, (shots + 7) // 8) + Observable `m` in shot `s` is `(obs[m, s // 8] >> (s % 8)) & 1` Examples: >>> import stim >>> c = stim.Circuit(''' - ... H 0 - ... CNOT 0 1 + ... R 0 1 2 + ... CX 0 1 2 1 ... X_ERROR(1.0) 0 - ... M 0 1 - ... DETECTOR rec[-1] rec[-2] + ... M 0 1 2 + ... DETECTOR rec[-2] + ... DETECTOR rec[-1] rec[-2] rec[-3] + ... OBSERVABLE_INCLUDE(0) rec[-1] ... ''') >>> s = c.compile_detector_sampler() - >>> s.sample(shots=1) - array([[ True]]) + >>> dets, obs = s.sample(shots=5, separate_observables=True) + >>> dets + array([[False, True], + [False, True], + [False, True], + [False, True], + [False, True]]) + >>> obs + array([[False], + [False], + [False], + [False], + [False]]) + + >>> dets, obs = s.sample( + ... shots=90, + ... bit_packed=True, + ... transposed=True, + ... separate_observables=True, + ... ) + >>> dets + array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 3]], + dtype=uint8) + >>> obs + array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8) """ def sample_bit_packed( self, diff --git a/glue/python/src/stim/__init__.pyi b/glue/python/src/stim/__init__.pyi index d71851394..195a124a8 100644 --- a/glue/python/src/stim/__init__.pyi +++ b/glue/python/src/stim/__init__.pyi @@ -4763,10 +4763,9 @@ class CompiledDetectorSampler: self, shots: int, *, - prepend_observables: bool = False, - append_observables: bool = False, separate_observables: bool = False, bit_packed: bool = False, + transposed: bool = False, dets_out: Optional[np.ndarray] = None, obs_out: Optional[np.ndarray] = None, ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: @@ -4781,12 +4780,16 @@ class CompiledDetectorSampler: separate_observables: Defaults to False. When set to True, the return value is a (detection_events, observable_flips) tuple instead of a flat detection_events array. - prepend_observables: Defaults to false. When set, observables are included - with the detectors and are placed at the start of the results. - append_observables: Defaults to false. When set, observables are included - with the detectors and are placed at the end of the results. bit_packed: Returns a uint8 numpy array with 8 bits per byte, instead of a bool_ numpy array with 1 bit per byte. Uses little endian packing. + transposed: Defaults to false. When set to true, outputs are transposed + (i.e. the first axis is detectors/observables and the second axis is + shots instead of vice versa). + + Using transposed=True has better performance than transposed=False, + because the output more closely corresponds to stim's internal + representation. But data in this transposed form is often more difficult + to work with. dets_out: Defaults to None. Specifies a pre-allocated numpy array to write the detection event data into. This array must have the correct shape and dtype. @@ -4794,65 +4797,100 @@ class CompiledDetectorSampler: the observable flip data into. This array must have the correct shape and dtype. - Returns: - A numpy array or tuple of numpy arrays containing the samples. - - if separate_observables=False and bit_packed=False: - A single numpy array. - dtype=bool_ - shape=( - shots, - num_detectors + num_observables * ( - append_observables + prepend_observables), - ) - The bit for detection event `m` in shot `s` is at - result[s, m] - - if separate_observables=False and bit_packed=True: - A single numpy array. - dtype=uint8 - shape=( - shots, - math.ceil((num_detectors + num_observables * ( - append_observables + prepend_observables)) / 8), - ) - The bit for detection event `m` in shot `s` is at - (result[s, m // 8] >> (m % 8)) & 1 + Deprecated Arguments: + (These arguments still work but are hidden and not recommended.) + + prepend_observables: Prefer `separate_observables`. Defaults to false. + When set, observables are included with the detectors and are placed at + the start of the results. Using this argument is generally a bad idea, + compared to using `separate_observables=True`, because you will + inevitably just end up immediately separating the observable data + from the detector data anyways. + append_observables: Prefer `separate_observables`. Defaults to false. + When set, observables are included with the detectors and are placed at + the end of the results. Using this argument is generally a bad idea, + compared to using `separate_observables=True`, because you will + inevitably just end up immediately separating the observable data + from the detector data anyways. + + Returns: + If separate_observables: + A tuple `(dets, obs)` of two numpy arrays. + else: + A single numpy array `dets`. - if separate_observables=True and bit_packed=False: - A (dets, obs) tuple. + if bit_packed=False and transposed=False: dets.dtype=bool_ dets.shape=(shots, num_detectors) + Detection event `m` in shot `s` is dets[s, m] + if bit_packed=False and transposed=True: + dets.dtype=bool_ + dets.shape=(num_detectors, shots) + Detection event `m` in shot `s` is dets[m, s] + if bit_packed=True and transposed=False: + dets.dtype=uint8 + dets.shape=(shots, (num_detectors + 7) // 8) + Detection event `m` in shot `s` is `(dets[s, m // 8] >> (m % 8)) & 1` + if bit_packed=True and transposed=True: + dets.dtype=uint8 + dets.shape=(num_detectors, (shots + 7) // 8) + Detection event `m` in shot `s` is `(dets[m, s // 8] >> (s % 8)) & 1` + + if separate_observables and bit_packed=False and transposed=False: obs.dtype=bool_ obs.shape=(shots, num_observables) - The bit for detection event `m` in shot `s` is at - dets[s, m] - The bit for observable `m` in shot `s` is at - obs[s, m] - - if separate_observables=True and bit_packed=True: - A (dets, obs) tuple. - dets.dtype=uint8 - dets.shape=(shots, math.ceil(num_detectors / 8)) + Observable `m` in shot `s` is obs[s, m] + if separate_observables and bit_packed=False and transposed=True: + obs.dtype=bool_ + obs.shape=(num_observables, shots) + Observable `m` in shot `s` is obs[m, s] + if separate_observables and bit_packed=True and transposed=False: obs.dtype=uint8 - obs.shape=(shots, math.ceil(num_observables / 8)) - The bit for detection event `m` in shot `s` is at - (dets[s, m // 8] >> (m % 8)) & 1 - The bit for observable `m` in shot `s` is at - (obs[s, m // 8] >> (m % 8)) & 1 + obs.shape=(shots, (num_observables + 7) // 8) + Observable `m` in shot `s` is `(obs[s, m // 8] >> (m % 8)) & 1` + if separate_observables and bit_packed=True and transposed=True: + obs.dtype=uint8 + obs.shape=(num_observables, (shots + 7) // 8) + Observable `m` in shot `s` is `(obs[m, s // 8] >> (s % 8)) & 1` Examples: >>> import stim >>> c = stim.Circuit(''' - ... H 0 - ... CNOT 0 1 + ... R 0 1 2 + ... CX 0 1 2 1 ... X_ERROR(1.0) 0 - ... M 0 1 - ... DETECTOR rec[-1] rec[-2] + ... M 0 1 2 + ... DETECTOR rec[-2] + ... DETECTOR rec[-1] rec[-2] rec[-3] + ... OBSERVABLE_INCLUDE(0) rec[-1] ... ''') >>> s = c.compile_detector_sampler() - >>> s.sample(shots=1) - array([[ True]]) + >>> dets, obs = s.sample(shots=5, separate_observables=True) + >>> dets + array([[False, True], + [False, True], + [False, True], + [False, True], + [False, True]]) + >>> obs + array([[False], + [False], + [False], + [False], + [False]]) + + >>> dets, obs = s.sample( + ... shots=90, + ... bit_packed=True, + ... transposed=True, + ... separate_observables=True, + ... ) + >>> dets + array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 3]], + dtype=uint8) + >>> obs + array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8) """ def sample_bit_packed( self, diff --git a/src/stim/circuit/circuit_pybind_perf.py b/src/stim/circuit/circuit_pybind_perf.py index f015c476f..5d8cf6d76 100644 --- a/src/stim/circuit/circuit_pybind_perf.py +++ b/src/stim/circuit/circuit_pybind_perf.py @@ -2,6 +2,7 @@ import time from typing import Callable, Iterable +import numpy as np import stim @@ -185,6 +186,28 @@ def run(): return wrap(original_method) +@benchmark(goal_micros=79, units={"shots": 1024, "detectors": 1024*336, "ops": 1024*5858}) +def benchmark_sample_detectors(): + c = stim.Circuit.generated("surface_code:rotated_memory_x", distance=7, rounds=7, after_clifford_depolarization=1e-3) + sampler = c.compile_detector_sampler() + + def run(): + sampler.sample(shots=1024, bit_packed=True, separate_observables=True) + return benchmark_go(run) + + +@benchmark(goal_micros=59, units={"shots": 1024, "detectors": 1024*336, "ops": 1024*5858}) +def benchmark_sample_detectors_transposed_buffered(): + c = stim.Circuit.generated("surface_code:rotated_memory_x", distance=7, rounds=7, after_clifford_depolarization=1e-3) + sampler = c.compile_detector_sampler() + dets_out = np.zeros(dtype=np.uint8, shape=(c.num_detectors, 1024 // 8)) + obs_out = np.zeros(dtype=np.uint8, shape=(c.num_observables, 1024 // 8)) + + def run(): + sampler.sample(shots=1024, bit_packed=True, transposed=True, separate_observables=True, dets_out=dets_out, obs_out=obs_out) + return benchmark_go(run) + + @benchmark(goal_micros=27, units={"targets": 2000}) def benchmark_circuit_append_int_in_large_chunks(): c = stim.Circuit() diff --git a/src/stim/py/compiled_detector_sampler.pybind.cc b/src/stim/py/compiled_detector_sampler.pybind.cc index 05bf5ff64..747e91014 100644 --- a/src/stim/py/compiled_detector_sampler.pybind.cc +++ b/src/stim/py/compiled_detector_sampler.pybind.cc @@ -37,6 +37,8 @@ pybind11::object CompiledDetectorSampler::sample_to_numpy( bool append_observables, bool separate_observables, bool bit_packed, + bool transposed_dets, + bool transposed_obs, pybind11::object dets_out, pybind11::object obs_out) { if (separate_observables && (append_observables || prepend_observables)) { @@ -57,7 +59,7 @@ pybind11::object CompiledDetectorSampler::sample_to_numpy( pybind11::object py_obs_data = pybind11::none(); if (separate_observables || !obs_out.is_none()) { py_obs_data = - simd_bit_table_to_numpy(obs_data, circuit_stats.num_observables, num_shots, bit_packed, true, obs_out); + simd_bit_table_to_numpy(obs_data, circuit_stats.num_observables, num_shots, bit_packed, transposed_obs, obs_out); } pybind11::object py_det_data = pybind11::none(); @@ -72,10 +74,10 @@ pybind11::object CompiledDetectorSampler::sample_to_numpy( concat_data = obs_data.concat_major(concat_data, circuit_stats.num_observables, num_concat); num_concat += circuit_stats.num_observables; } - py_det_data = simd_bit_table_to_numpy(concat_data, num_concat, num_shots, bit_packed, true, dets_out); + py_det_data = simd_bit_table_to_numpy(concat_data, num_concat, num_shots, bit_packed, transposed_dets, dets_out); } else { py_det_data = - simd_bit_table_to_numpy(det_data, circuit_stats.num_detectors, num_shots, bit_packed, true, dets_out); + simd_bit_table_to_numpy(det_data, circuit_stats.num_detectors, num_shots, bit_packed, transposed_dets, dets_out); } if (separate_observables) { @@ -221,9 +223,14 @@ void stim_pybind::pybind_compiled_detector_sampler_methods( bool append, bool separate_observables, bool bit_packed, + bool python_api_transposed, pybind11::object dets_out, pybind11::object obs_out) { - return self.sample_to_numpy(shots, prepend, append, separate_observables, bit_packed, dets_out, obs_out); + // Caution: the python API has an inverted meaning for transposed! + // Users expect (shot, detector) indexing meaning transposed is (detector, shot) indexing. + // But the internal implementation bit stripes over shots for performance so uses (detector, shot indexing) + // with (shot, detector) being the transpose. Thus we invert the transposed value here. + return self.sample_to_numpy(shots, prepend, append, separate_observables, bit_packed, !python_api_transposed, !python_api_transposed, dets_out, obs_out); }, pybind11::arg("shots"), pybind11::kw_only(), @@ -231,10 +238,11 @@ void stim_pybind::pybind_compiled_detector_sampler_methods( pybind11::arg("append_observables") = false, pybind11::arg("separate_observables") = false, pybind11::arg("bit_packed") = false, + pybind11::arg("transposed") = false, pybind11::arg("dets_out") = pybind11::none(), pybind11::arg("obs_out") = pybind11::none(), clean_doc_string(R"DOC( - @signature def sample(self, shots: int, *, prepend_observables: bool = False, append_observables: bool = False, separate_observables: bool = False, bit_packed: bool = False, dets_out: Optional[np.ndarray] = None, obs_out: Optional[np.ndarray] = None) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: + @signature def sample(self, shots: int, *, separate_observables: bool = False, bit_packed: bool = False, transposed: bool = False, dets_out: Optional[np.ndarray] = None, obs_out: Optional[np.ndarray] = None) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: Returns a numpy array containing a batch of detector samples from the circuit. The circuit must define the detectors using DETECTOR instructions. Observables @@ -246,12 +254,16 @@ void stim_pybind::pybind_compiled_detector_sampler_methods( separate_observables: Defaults to False. When set to True, the return value is a (detection_events, observable_flips) tuple instead of a flat detection_events array. - prepend_observables: Defaults to false. When set, observables are included - with the detectors and are placed at the start of the results. - append_observables: Defaults to false. When set, observables are included - with the detectors and are placed at the end of the results. bit_packed: Returns a uint8 numpy array with 8 bits per byte, instead of a bool_ numpy array with 1 bit per byte. Uses little endian packing. + transposed: Defaults to false. When set to true, outputs are transposed + (i.e. the first axis is detectors/observables and the second axis is + shots instead of vice versa). + + Using transposed=True has better performance than transposed=False, + because the output more closely corresponds to stim's internal + representation. But data in this transposed form is often more difficult + to work with. dets_out: Defaults to None. Specifies a pre-allocated numpy array to write the detection event data into. This array must have the correct shape and dtype. @@ -259,72 +271,107 @@ void stim_pybind::pybind_compiled_detector_sampler_methods( the observable flip data into. This array must have the correct shape and dtype. + Deprecated Arguments: + (These arguments still work but are hidden and not recommended.) + + prepend_observables: Prefer `separate_observables`. Defaults to false. + When set, observables are included with the detectors and are placed at + the start of the results. Using this argument is generally a bad idea, + compared to using `separate_observables=True`, because you will + inevitably just end up immediately separating the observable data + from the detector data anyways. + append_observables: Prefer `separate_observables`. Defaults to false. + When set, observables are included with the detectors and are placed at + the end of the results. Using this argument is generally a bad idea, + compared to using `separate_observables=True`, because you will + inevitably just end up immediately separating the observable data + from the detector data anyways. + Returns: - A numpy array or tuple of numpy arrays containing the samples. - - if separate_observables=False and bit_packed=False: - A single numpy array. - dtype=bool_ - shape=( - shots, - num_detectors + num_observables * ( - append_observables + prepend_observables), - ) - The bit for detection event `m` in shot `s` is at - result[s, m] - - if separate_observables=False and bit_packed=True: - A single numpy array. - dtype=uint8 - shape=( - shots, - math.ceil((num_detectors + num_observables * ( - append_observables + prepend_observables)) / 8), - ) - The bit for detection event `m` in shot `s` is at - (result[s, m // 8] >> (m % 8)) & 1 - - if separate_observables=True and bit_packed=False: - A (dets, obs) tuple. + If separate_observables: + A tuple `(dets, obs)` of two numpy arrays. + else: + A single numpy array `dets`. + + if bit_packed=False and transposed=False: dets.dtype=bool_ dets.shape=(shots, num_detectors) + Detection event `m` in shot `s` is dets[s, m] + if bit_packed=False and transposed=True: + dets.dtype=bool_ + dets.shape=(num_detectors, shots) + Detection event `m` in shot `s` is dets[m, s] + if bit_packed=True and transposed=False: + dets.dtype=uint8 + dets.shape=(shots, (num_detectors + 7) // 8) + Detection event `m` in shot `s` is `(dets[s, m // 8] >> (m % 8)) & 1` + if bit_packed=True and transposed=True: + dets.dtype=uint8 + dets.shape=(num_detectors, (shots + 7) // 8) + Detection event `m` in shot `s` is `(dets[m, s // 8] >> (s % 8)) & 1` + + if separate_observables and bit_packed=False and transposed=False: obs.dtype=bool_ obs.shape=(shots, num_observables) - The bit for detection event `m` in shot `s` is at - dets[s, m] - The bit for observable `m` in shot `s` is at - obs[s, m] - - if separate_observables=True and bit_packed=True: - A (dets, obs) tuple. - dets.dtype=uint8 - dets.shape=(shots, math.ceil(num_detectors / 8)) + Observable `m` in shot `s` is obs[s, m] + if separate_observables and bit_packed=False and transposed=True: + obs.dtype=bool_ + obs.shape=(num_observables, shots) + Observable `m` in shot `s` is obs[m, s] + if separate_observables and bit_packed=True and transposed=False: obs.dtype=uint8 - obs.shape=(shots, math.ceil(num_observables / 8)) - The bit for detection event `m` in shot `s` is at - (dets[s, m // 8] >> (m % 8)) & 1 - The bit for observable `m` in shot `s` is at - (obs[s, m // 8] >> (m % 8)) & 1 + obs.shape=(shots, (num_observables + 7) // 8) + Observable `m` in shot `s` is `(obs[s, m // 8] >> (m % 8)) & 1` + if separate_observables and bit_packed=True and transposed=True: + obs.dtype=uint8 + obs.shape=(num_observables, (shots + 7) // 8) + Observable `m` in shot `s` is `(obs[m, s // 8] >> (s % 8)) & 1` Examples: >>> import stim >>> c = stim.Circuit(''' - ... H 0 - ... CNOT 0 1 + ... R 0 1 2 + ... CX 0 1 2 1 ... X_ERROR(1.0) 0 - ... M 0 1 - ... DETECTOR rec[-1] rec[-2] + ... M 0 1 2 + ... DETECTOR rec[-2] + ... DETECTOR rec[-1] rec[-2] rec[-3] + ... OBSERVABLE_INCLUDE(0) rec[-1] ... ''') >>> s = c.compile_detector_sampler() - >>> s.sample(shots=1) - array([[ True]]) + >>> dets, obs = s.sample(shots=5, separate_observables=True) + >>> dets + array([[False, True], + [False, True], + [False, True], + [False, True], + [False, True]]) + >>> obs + array([[False], + [False], + [False], + [False], + [False]]) + + >>> dets, obs = s.sample( + ... shots=90, + ... bit_packed=True, + ... transposed=True, + ... separate_observables=True, + ... ) + >>> dets + array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 3]], + dtype=uint8) + >>> obs + array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8) )DOC") .data()); c.def( "sample_bit_packed", [](CompiledDetectorSampler &self, size_t shots, bool prepend, bool append) { - return self.sample_to_numpy(shots, prepend, append, false, true, pybind11::none(), pybind11::none()); + return self.sample_to_numpy(shots, prepend, append, false, true, true, true, pybind11::none(), pybind11::none()); }, pybind11::arg("shots"), pybind11::kw_only(), diff --git a/src/stim/py/compiled_detector_sampler.pybind.h b/src/stim/py/compiled_detector_sampler.pybind.h index 53912bb15..8529bc54f 100644 --- a/src/stim/py/compiled_detector_sampler.pybind.h +++ b/src/stim/py/compiled_detector_sampler.pybind.h @@ -40,6 +40,8 @@ struct CompiledDetectorSampler { bool append_observables, bool separate_observables, bool bit_packed, + bool transposed_dets, + bool transposed_obs, pybind11::object dets_out, pybind11::object obs_out); void sample_write( diff --git a/src/stim/py/compiled_detector_sampler_pybind_test.py b/src/stim/py/compiled_detector_sampler_pybind_test.py index 3430f5ace..36f2c63dd 100644 --- a/src/stim/py/compiled_detector_sampler_pybind_test.py +++ b/src/stim/py/compiled_detector_sampler_pybind_test.py @@ -312,3 +312,148 @@ def test_manual_output_buffer(): assert ret is buf assert np.array_equal(buf, [[0, 0, 1, 1, 1, 0, 0, 1, 1, 1]] * 17) assert np.array_equal(buf2, [[1]] * 17) + + +def test_transposed(): + circuit = stim.Circuit(''' + X_ERROR(1) 0 + M 0 + DETECTOR + DETECTOR + DETECTOR rec[-1] + DETECTOR rec[-1] + + DETECTOR rec[-1] + DETECTOR + DETECTOR + DETECTOR rec[-1] + + DETECTOR rec[-1] + OBSERVABLE_INCLUDE(0) rec[-1] + ''') + sampler = circuit.compile_detector_sampler() + + with pytest.raises(ValueError, match="shape"): + sampler.sample(shots=17, transposed=True, dets_out=np.zeros(shape=(17, 9), dtype=np.bool_)) + with pytest.raises(ValueError, match="shape"): + sampler.sample(shots=17, transposed=True, dets_out=np.zeros(shape=(8, 17), dtype=np.bool_)) + with pytest.raises(ValueError, match="shape"): + sampler.sample(shots=17, transposed=True, dets_out=np.zeros(shape=(10, 17), dtype=np.bool_)) + with pytest.raises(ValueError, match="shape"): + sampler.sample(shots=17, transposed=True, dets_out=np.zeros(shape=(9, 18), dtype=np.bool_)) + with pytest.raises(ValueError, match="shape"): + sampler.sample(shots=17, transposed=True, dets_out=np.zeros(shape=(9, 16), dtype=np.bool_)) + with pytest.raises(ValueError, match="bool"): + sampler.sample(shots=17, transposed=True, dets_out=np.zeros(shape=(9, 17), dtype=np.uint8)) + + buf = np.zeros(shape=(9, 17), dtype=np.bool_) + ret = sampler.sample(shots=17, dets_out=buf, transposed=True) + assert ret is buf + assert np.array_equal(buf, [[e] * 17 for e in [0, 0, 1, 1, 1, 0, 0, 1, 1]]) + + buf = np.zeros(shape=(9, 3), dtype=np.uint8) + ret = sampler.sample( + shots=17, + dets_out=buf, + bit_packed=True, + transposed=True, + ) + assert ret is buf + assert np.array_equal(buf, [([255, 255, 1] if e else [0, 0, 0]) for e in [0, 0, 1, 1, 1, 0, 0, 1, 1]]) + + buf = np.zeros(shape=(3, 9), dtype=np.uint8).transpose() + ret = sampler.sample( + shots=17, + dets_out=buf, + bit_packed=True, + transposed=True, + ) + assert ret is buf + assert np.array_equal(buf, [([255, 255, 1] if e else [0, 0, 0]) for e in [0, 0, 1, 1, 1, 0, 0, 1, 1]]) + + buf = np.zeros(shape=(9, 17), dtype=np.bool_) + buf2 = np.zeros(shape=(1, 17), dtype=np.bool_) + ret = sampler.sample( + shots=17, + dets_out=buf, + obs_out=buf2, + transposed=True, + ) + assert ret is buf + assert np.array_equal(buf, [[e] * 17 for e in [0, 0, 1, 1, 1, 0, 0, 1, 1]]) + assert np.array_equal(buf2, [[True] * 17]) + + buf = np.zeros(shape=(9, 17), dtype=np.bool_) + buf2 = np.zeros(shape=(1, 17), dtype=np.bool_) + ret, ret2 = sampler.sample( + shots=17, + dets_out=buf, + obs_out=buf2, + separate_observables=True, + transposed=True, + ) + assert ret is buf + assert ret2 is buf2 + assert np.array_equal(buf, [[e] * 17 for e in [0, 0, 1, 1, 1, 0, 0, 1, 1]]) + assert np.array_equal(buf2, [[True] * 17]) + + buf = np.zeros(shape=(10, 17), dtype=np.bool_) + buf2 = np.zeros(shape=(1, 17), dtype=np.bool_) + ret = sampler.sample( + shots=17, + dets_out=buf, + obs_out=buf2, + append_observables=True, + transposed=True, + ) + assert ret is buf + assert np.array_equal(buf, [[e] * 17 for e in [0, 0, 1, 1, 1, 0, 0, 1, 1, 1]]) + assert np.array_equal(buf2, [[1] * 17]) + + buf = np.zeros(shape=(17, 10), dtype=np.bool_).transpose() + buf2 = np.zeros(shape=(1, 17), dtype=np.bool_) + ret = sampler.sample( + shots=17, + dets_out=buf, + obs_out=buf2, + append_observables=True, + transposed=True, + ) + assert ret is buf + assert np.array_equal(buf, [[e] * 17 for e in [0, 0, 1, 1, 1, 0, 0, 1, 1, 1]]) + assert np.array_equal(buf2, [[1] * 17]) + + +def test_consistent_data_across_detector_sampling_options(): + circuit = stim.Circuit(''' + REPEAT 501 { + M(0.5) 0 + DETECTOR rec[-1] + } + ''') + sampler = circuit.compile_detector_sampler(seed=42) + a = sampler.sample(shots=311, bit_packed=True, transposed=True) + sampler = circuit.compile_detector_sampler(seed=42) + b = sampler.sample(shots=311, bit_packed=True, transposed=False) + sampler = circuit.compile_detector_sampler(seed=42) + c = sampler.sample(shots=311, bit_packed=False, transposed=True) + sampler = circuit.compile_detector_sampler(seed=42) + d = sampler.sample(shots=311, bit_packed=False, transposed=False) + + np.testing.assert_array_equal(c, np.transpose(d)) + np.testing.assert_array_equal(b, np.packbits(d, bitorder='little', axis=1)) + np.testing.assert_array_equal(a, np.packbits(c, bitorder='little', axis=1)) + + +def test_consistent_data_across_measurement_sampling_options(): + circuit = stim.Circuit(''' + REPEAT 501 { + M(0.5) 0 + } + ''') + sampler = circuit.compile_sampler(seed=42) + a = sampler.sample(shots=311, bit_packed=True) + sampler = circuit.compile_sampler(seed=42) + b = sampler.sample(shots=311, bit_packed=False) + + np.testing.assert_array_equal(a, np.packbits(b, bitorder='little', axis=1)) diff --git a/src/stim/py/numpy.pybind.cc b/src/stim/py/numpy.pybind.cc index 9f5e75fd2..64553852d 100644 --- a/src/stim/py/numpy.pybind.cc +++ b/src/stim/py/numpy.pybind.cc @@ -19,14 +19,14 @@ using namespace stim_pybind; static pybind11::object transposed_simd_bit_table_to_numpy_uint8( const simd_bit_table &table, - size_t num_major_in, - size_t num_minor_in, + size_t table_shape0, + size_t table_shape1, pybind11::object out_buffer) { - size_t num_major_bytes_in = (num_major_in + 7) / 8; + size_t num_major_bytes_in = (table_shape0 + 7) / 8; if (out_buffer.is_none()) { auto numpy = pybind11::module::import("numpy"); - out_buffer = numpy.attr("empty")(pybind11::make_tuple(num_minor_in, num_major_bytes_in), numpy.attr("uint8")); + out_buffer = numpy.attr("empty")(pybind11::make_tuple(table_shape1, num_major_bytes_in), numpy.attr("uint8")); } if (!pybind11::isinstance>(out_buffer)) { @@ -36,25 +36,45 @@ static pybind11::object transposed_simd_bit_table_to_numpy_uint8( if (buf.ndim() != 2) { throw std::invalid_argument("Output buffer wasn't two dimensional."); } - if ((size_t)buf.shape(0) != num_minor_in || (size_t)buf.shape(1) != num_major_bytes_in) { + if ((size_t)buf.shape(0) != table_shape1 || (size_t)buf.shape(1) != num_major_bytes_in) { std::stringstream ss; - ss << "Expected output buffer to have shape=(" << num_minor_in << ", " << num_major_bytes_in << ")"; + ss << "Expected output buffer to have shape=(" << table_shape1 << ", " << num_major_bytes_in << ")"; ss << " but its shape is (" << buf.shape(0) << ", " << buf.shape(1) << ")."; throw std::invalid_argument(ss.str()); } - if (num_major_in && num_minor_in) { - auto stride = buf.strides(1); - for (size_t minor_in = 0; minor_in < num_minor_in; minor_in++) { - auto ptr = buf.mutable_data(minor_in, 0); - for (size_t major_in = 0; major_in < num_major_in; major_in += 8) { - uint8_t v = 0; - for (size_t b = 0; b < 8 && major_in + b < num_major_in; b++) { - bool bit = table[major_in + b][minor_in]; - v |= bit << b; + if (table_shape0 && table_shape1) { + std::array transpose_buffer; + + uint8_t *base = buf.mutable_data(0, 0); + auto stride0 = buf.strides(0); + auto stride1 = buf.strides(1); + for (size_t b0 = 0; b0 < table_shape0; b0 += 64) { + for (size_t b1 = 0; b1 < table_shape1; b1 += 64) { + size_t n0 = std::min(size_t{64}, table_shape0 - b0); + size_t n1 = std::min(size_t{64}, table_shape1 - b1); + + // Read a 64x64 block of bits out of the table. + for (size_t d0 = 0; d0 < n0; d0++) { + transpose_buffer[d0] = table[b0 + d0].u64[b1 / 64]; + } + for (size_t d0 = n0; d0 < 64; d0++) { + transpose_buffer[d0] = 0; + } + + // Transpose the block of bits. + inplace_transpose_64x64(transpose_buffer.data(), 1); + + // Write the transposed 64x64 block of bits into the numpy array. + uint8_t *row_out = base + stride0 * b1 + stride1 * (b0 / 8); + for (size_t d1 = 0; d1 < n1; d1++) { + uint8_t *cell_out = row_out; + for (size_t d0 = 0; d0 < n0; d0 += 8) { + *cell_out = (transpose_buffer[d1] >> d0) & 0xFF; + cell_out += stride1; + } + row_out += stride0; } - *ptr = v; - ptr += stride; } } }