Add pickling support to stim.CompiledMeasurementsToDetectionEventsConverter

src/stim/simulators/measurements_to_detection_events.pybind.cc

@@ -161,6 +161,7 @@ CompiledMeasurementsToDetectionEventsConverter stim_pybind::py_init_compiled_mea
 
 void stim_pybind::pybind_compiled_measurements_to_detection_events_converter_methods(
     pybind11::module &m, pybind11::class_<CompiledMeasurementsToDetectionEventsConverter> &c) {
+    using SerializationTuple = std::tuple<stim::Circuit, bool, pybind11::object, size_t>;
     c.def(
         pybind11::init(&py_init_compiled_measurements_to_detection_events_converter),
         pybind11::arg("circuit"),
@@ -371,4 +372,29 @@ void stim_pybind::pybind_compiled_measurements_to_detection_events_converter_met
         &CompiledMeasurementsToDetectionEventsConverter::repr,
         "Returns text that is a valid python expression evaluating to an equivalent "
         "`stim.CompiledMeasurementsToDetectionEventsConverter`.");
+
+    c.def(pybind11::pickle(
+        // __getstate__ function: returns a tuple to be pickled.
+        [](const CompiledMeasurementsToDetectionEventsConverter &self) -> 
+        SerializationTuple {
+            size_t num_ref_bits = self.circuit_stats.num_measurements;
+            pybind11::object ref_sample_numpy =
+                stim_pybind::simd_bits_to_numpy(self.ref_sample, num_ref_bits,
+                                            /*bit_packed=*/true);
+
+            return SerializationTuple(self.circuit, self.skip_reference_sample,
+                                  ref_sample_numpy, num_ref_bits);
+        },
+        // __setstate__ function: reconstructs the object from the Python tuple.
+        [](SerializationTuple t_py) {
+            const auto &[circuit, skip_ref, ref_bits_npy, num_ref_bits] = t_py;
+
+            stim::simd_bits<stim::MAX_BITWORD_WIDTH> reconstructed_ref_sample(
+                num_ref_bits);
+            stim_pybind::memcpy_bits_from_numpy_to_simd(num_ref_bits, ref_bits_npy,
+                                                    reconstructed_ref_sample);
+
+            return CompiledMeasurementsToDetectionEventsConverter(
+                std::move(reconstructed_ref_sample), circuit, skip_ref);
+      }));
 }

src/stim/simulators/measurements_to_detection_events_test.py

@@ -297,3 +297,42 @@ def test_anticommuting_pieces_combining_into_deterministic_observable():
     m = c.compile_sampler().sample_bit_packed(shots=1000)
     det, obs = c.compile_m2d_converter().convert(measurements=m, separate_observables=True)
     np.testing.assert_array_equal(obs, obs * 0)
+
+
+def test_converter_pickle():
+    converter = stim.Circuit('''
+       X_ERROR(0.1) 0
+       X 0
+       CNOT sweep[0] 0
+       M 0
+       DETECTOR rec[-1]
+       OBSERVABLE_INCLUDE(0) rec[-1]
+    ''').compile_m2d_converter()
+
+    roundtripped = pickle.loads(pickle.dumps(converter))
+    assert str(converter) == str(roundtripped)
+
+    result = roundtripped.convert(
+        measurements=np.array([[0], [1]], dtype=np.bool_),
+        append_observables=False,
+    )
+    assert result.dtype == np.bool_
+    assert result.shape == (2, 1)
+    np.testing.assert_array_equal(result, [[1], [0]])
+
+    result = roundtripped.convert(
+        measurements=np.array([[0], [1]], dtype=np.bool_),
+        append_observables=True,
+    )
+    assert result.dtype == np.bool_
+    assert result.shape == (2, 2)
+    np.testing.assert_array_equal(result, [[1, 1], [0, 0]])
+
+    result = roundtripped.convert(
+        measurements=np.array([[0], [1], [0], [1]], dtype=np.bool_),
+        sweep_bits=np.array([[0], [0], [1], [1]], dtype=np.bool_),
+        append_observables=True,
+    )
+    assert result.dtype == np.bool_
+    assert result.shape == (4, 2)
+    np.testing.assert_array_equal(result, [[1, 1], [0, 0], [0, 0], [1, 1]])