Switch from using stim_py to use py::object in python wrapper (#48)

following the discussion in quantumlib/Stim#963,
this will unblock
#41

Author: NoureldinYosri

WORKSPACE

@@ -66,11 +66,3 @@ http_archive(
     urls = ["https://github.com/bazelbuild/platforms/archive/refs/tags/0.0.6.zip"],
     strip_prefix = "platforms-0.0.6",
 )
-
-http_archive(
-    name = "stim_py",
-    build_file = "//external:stim_py.BUILD",
-    sha256 = "95236006859d6754be99629d4fb44788e742e962ac8c59caad421ca088f7350e",
-    strip_prefix = "stim-1.15.0",
-    urls = ["https://github.com/quantumlib/Stim/releases/download/v1.15.0/stim-1.15.0.tar.gz"],
-)

external/stim_py.BUILD

@@ -67,6 +67,7 @@ cc_library(
 pybind_library(
     name = "tesseract_decoder_pybind",
     srcs = [
+        "stim_utils.pybind.h",
         "common.pybind.h",
         "utils.pybind.h",
         "simplex.pybind.h",
@@ -77,7 +78,6 @@ pybind_library(
         ":libutils",
         ":libsimplex",
         ":libtesseract",
-        "@stim_py//:stim_pybind_lib",
     ],
 )
 
@@ -88,7 +88,6 @@ pybind_extension(
     ],
     deps = [
         ":tesseract_decoder_pybind",
-        "@stim_py//:stim",
     ],
 )
 
@@ -213,4 +212,3 @@ cc_binary(
         "@nlohmann_json//:json",
     ],
 )
-

src/BUILD

@@ -22,8 +22,7 @@
 #include <vector>
 
 #include "common.h"
-#include "src/stim/dem/dem_instruction.pybind.h"
-#include "stim/dem/detector_error_model_target.pybind.h"
+#include "stim_utils.pybind.h"
 
 namespace py = pybind11;
 
@@ -39,8 +38,9 @@ void add_common_module(py::module &root) {
       .def(py::self == py::self)
       .def(py::self != py::self)
       .def("as_dem_instruction_targets", [](common::Symptom s) {
-        std::vector<stim_pybind::ExposedDemTarget> ret;
-        for (auto &t : s.as_dem_instruction_targets()) ret.emplace_back(t);
+        std::vector<py::object> ret;
+        for (auto &t : s.as_dem_instruction_targets())
+          ret.push_back(make_py_object(t, "DemTarget"));
         return ret;
       });
 
@@ -57,15 +57,38 @@ void add_common_module(py::module &root) {
                     std::vector<bool> &>(),
            py::arg("likelihood_cost"), py::arg("probability"), py::arg("detectors"),
            py::arg("observables"), py::arg("dets_array"))
-      .def(py::init([](stim_pybind::ExposedDemInstruction edi) {
-             return new common::Error(edi.as_dem_instruction());
+      .def(py::init([](py::object edi) {
+             std::vector<double> args;
+             std::vector<stim::DemTarget> targets;
+             auto di = parse_py_dem_instruction(edi, args, targets);
+             return new common::Error(di);
            }),
            py::arg("error"));
 
-  m.def("merge_identical_errors", &common::merge_identical_errors, py::arg("dem"));
-  m.def("remove_zero_probability_errors", &common::remove_zero_probability_errors, py::arg("dem"));
-  m.def("dem_from_counts", &common::dem_from_counts, py::arg("orig_dem"), py::arg("error_counts"),
-        py::arg("num_shots"));
+  m.def(
+      "merge_identical_errors",
+      [](py::object dem) {
+        auto input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
+        auto res = common::merge_identical_errors(input_dem);
+        return make_py_object(res, "DetectorErrorModel");
+      },
+      py::arg("dem"));
+  m.def(
+      "remove_zero_probability_errors",
+      [](py::object dem) {
+        return make_py_object(
+            common::remove_zero_probability_errors(parse_py_object<stim::DetectorErrorModel>(dem)),
+            "DetectorErrorModel");
+      },
+      py::arg("dem"));
+  m.def(
+      "dem_from_counts",
+      [](py::object orig_dem, const std::vector<size_t> error_counts, size_t num_shots) {
+        auto dem = parse_py_object<stim::DetectorErrorModel>(orig_dem);
+        return make_py_object(common::dem_from_counts(dem, error_counts, num_shots),
+                              "DetectorErrorModel");
+      },
+      py::arg("orig_dem"), py::arg("error_counts"), py::arg("num_shots"));
 }
 
 #endif

src/common.pybind.h

@@ -7,6 +7,7 @@ py_test(
     visibility = ["//:__subpackages__"],
     deps = [
         "@pypi//pytest",
+        "@pypi//stim",
         "//src:lib_tesseract_decoder",
     ],
 )
@@ -17,6 +18,7 @@ py_test(
     visibility = ["//:__subpackages__"],
     deps = [
         "@pypi//pytest",
+        "@pypi//stim",
         "//src:lib_tesseract_decoder",
     ],
 )
@@ -27,6 +29,7 @@ py_test(
     visibility = ["//:__subpackages__"],
     deps = [
         "@pypi//pytest",
+        "@pypi//stim",
         "//src:lib_tesseract_decoder",
     ],
 )
@@ -37,6 +40,7 @@ py_test(
     visibility = ["//:__subpackages__"],
     deps = [
         "@pypi//pytest",
+        "@pypi//stim",
         "//src:lib_tesseract_decoder",
     ],
 )

src/py/BUILD

@@ -32,6 +32,10 @@ def test_create_config():
         str(tesseract_decoder.tesseract.TesseractConfig(_DETECTOR_ERROR_MODEL))
         == "TesseractConfig(dem=DetectorErrorModel_Object, det_beam=65535, no_revisit_dets=0, at_most_two_errors_per_detector=0, verbose=0, pqlimit=18446744073709551615, det_orders=[], det_penalty=0)"
     )
+    assert (
+        tesseract_decoder.tesseract.TesseractConfig(_DETECTOR_ERROR_MODEL).dem
+        == _DETECTOR_ERROR_MODEL
+    )
 
 
 def test_create_node():

src/py/tesseract_test.py

@@ -21,18 +21,29 @@
 
 #include "common.h"
 #include "simplex.h"
+#include "stim_utils.pybind.h"
 
 namespace py = pybind11;
 
+namespace {
+SimplexConfig simplex_config_maker(py::object dem, bool parallelize = false,
+                                   size_t window_length = 0, size_t window_slide_length = 0,
+                                   bool verbose = false) {
+  stim::DetectorErrorModel input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
+  return SimplexConfig({input_dem, parallelize, window_length, window_slide_length, verbose});
+}
+
+};  // namespace
+
 void add_simplex_module(py::module &root) {
   auto m =
       root.def_submodule("simplex", "Module containing the SimplexDecoder and related methods");
 
   py::class_<SimplexConfig>(m, "SimplexConfig")
-      .def(py::init<stim::DetectorErrorModel, bool, size_t, size_t, bool>(), py::arg("dem"),
-           py::arg("parallelize") = false, py::arg("window_length") = 0,
-           py::arg("window_slide_length") = 0, py::arg("verbose") = false)
-      .def_readwrite("dem", &SimplexConfig::dem)
+      .def(py::init(&simplex_config_maker), py::arg("dem"), py::arg("parallelize") = false,
+           py::arg("window_length") = 0, py::arg("window_slide_length") = 0,
+           py::arg("verbose") = false)
+      .def_property("dem", &dem_getter<SimplexConfig>, &dem_setter<SimplexConfig>)
       .def_readwrite("parallelize", &SimplexConfig::parallelize)
       .def_readwrite("window_length", &SimplexConfig::window_length)
       .def_readwrite("window_slide_length", &SimplexConfig::window_slide_length)

src/simplex.pybind.h

@@ -0,0 +1,69 @@
+#ifndef _STIM_UTILS_PYBIND_H
+#define _STIM_UTILS_PYBIND_H
+
+#include <pybind11/operators.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "stim.h"
+
+namespace {
+namespace py = pybind11;
+}
+
+template <typename T>
+py::object make_py_object(const T cpp_obj, const char* py_name) {
+  auto stim_lib = py::module::import("stim");
+  return stim_lib.attr(py_name)(cpp_obj.str());
+}
+
+template <typename T>
+T parse_py_object(py::object py_obj) {
+  std::string obj_str = py::cast<std::string>(py_obj.attr("__str__")());
+  return T(obj_str);
+}
+
+stim::DemInstructionType parse_dit(std::string dit_str) {
+  if (dit_str == "error") return stim::DemInstructionType::DEM_ERROR;
+  if (dit_str == "detector") return stim::DemInstructionType::DEM_DETECTOR;
+  if (dit_str == "logical_observable") return stim::DemInstructionType::DEM_LOGICAL_OBSERVABLE;
+  if (dit_str == "shift_detectors") return stim::DemInstructionType::DEM_SHIFT_DETECTORS;
+  if (dit_str == "repeat") return stim::DemInstructionType::DEM_REPEAT_BLOCK;
+  throw std::invalid_argument("unknown dem instruction type: " + dit_str);
+  return stim::DemInstructionType::DEM_DETECTOR;
+}
+
+stim::DemTarget parse_py_dem_target(py::object py_obj) {
+  return stim::DemTarget::from_text(py::cast<std::string>(py_obj.attr("__str__")()));
+}
+
+stim::DemInstruction parse_py_dem_instruction(py::object py_obj, std::vector<double>& args,
+                                              std::vector<stim::DemTarget>& targets) {
+  for (auto t : py_obj.attr("args_copy")()) args.push_back(t.cast<double>());
+  stim::SpanRef args_ref(args);
+
+  for (auto t : py_obj.attr("targets_copy")())
+    targets.push_back(parse_py_dem_target(t.cast<py::object>()));
+
+  stim::SpanRef targets_ref(targets);
+  auto ty = parse_dit(py::cast<std::string>(py_obj.attr("type")));
+  std::string tag = py::cast<std::string>(py_obj.attr("tag"));
+
+  auto di = stim::DemInstruction();
+  di.arg_data = args_ref;
+  di.target_data = targets_ref;
+  di.tag = tag;
+  di.type = ty;
+  return di;
+}
+
+template <typename T>
+py::object dem_getter(const T& config) {
+  return make_py_object(config.dem, "DetectorErrorModel");
+}
+template <typename T>
+void dem_setter(T& config, py::object dem) {
+  config.dem = parse_py_object<stim::DetectorErrorModel>(dem);
+}
+
+#endif

src/stim_utils.pybind.h

@@ -19,22 +19,35 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#include "stim_utils.pybind.h"
 #include "tesseract.h"
 
 namespace py = pybind11;
 
+namespace {
+TesseractConfig tesseract_config_maker(
+    py::object dem, int det_beam = INF_DET_BEAM, bool beam_climbing = false,
+    bool no_revisit_dets = false, bool at_most_two_errors_per_detector = false,
+    bool verbose = false, size_t pqlimit = std::numeric_limits<size_t>::max(),
+    std::vector<std::vector<size_t>> det_orders = std::vector<std::vector<size_t>>(),
+    double det_penalty = 0.0) {
+  stim::DetectorErrorModel input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
+  return TesseractConfig({input_dem, det_beam, beam_climbing, no_revisit_dets,
+                          at_most_two_errors_per_detector, verbose, pqlimit, det_orders,
+                          det_penalty});
+}
+};  // namespace
 void add_tesseract_module(py::module &root) {
   auto m = root.def_submodule("tesseract", "Module containing the tesseract algorithm");
 
   m.attr("INF_DET_BEAM") = INF_DET_BEAM;
   py::class_<TesseractConfig>(m, "TesseractConfig")
-      .def(py::init<stim::DetectorErrorModel, int, bool, bool, bool, bool, size_t,
-                    std::vector<std::vector<size_t>>, double>(),
-           py::arg("dem"), py::arg("det_beam") = INF_DET_BEAM, py::arg("beam_climbing") = false,
-           py::arg("no_revisit_dets") = false, py::arg("at_most_two_errors_per_detector") = false,
-           py::arg("verbose") = false, py::arg("pqlimit") = std::numeric_limits<size_t>::max(),
+      .def(py::init(&tesseract_config_maker), py::arg("dem"), py::arg("det_beam") = INF_DET_BEAM,
+           py::arg("beam_climbing") = false, py::arg("no_revisit_dets") = false,
+           py::arg("at_most_two_errors_per_detector") = false, py::arg("verbose") = false,
+           py::arg("pqlimit") = std::numeric_limits<size_t>::max(),
            py::arg("det_orders") = std::vector<std::vector<size_t>>(), py::arg("det_penalty") = 0.0)
-      .def_readwrite("dem", &TesseractConfig::dem)
+      .def_property("dem", &dem_getter<TesseractConfig>, &dem_setter<TesseractConfig>)
       .def_readwrite("det_beam", &TesseractConfig::det_beam)
       .def_readwrite("no_revisit_dets", &TesseractConfig::no_revisit_dets)
       .def_readwrite("at_most_two_errors_per_detector",

src/tesseract.pybind.h

@@ -28,9 +28,27 @@ void add_utils_module(py::module &root) {
 
   m.attr("EPSILON") = EPSILON;
   m.attr("INF") = INF;
-  m.def("get_detector_coords", &get_detector_coords, py::arg("dem"));
-  m.def("build_detector_graph", &build_detector_graph, py::arg("dem"));
-  m.def("get_errors_from_dem", &get_errors_from_dem, py::arg("dem"));
+  m.def(
+      "get_detector_coords",
+      [](py::object dem) {
+        auto input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
+        return get_detector_coords(input_dem);
+      },
+      py::arg("dem"));
+  m.def(
+      "build_detector_graph",
+      [](py::object dem) {
+        auto input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
+        return build_detector_graph(input_dem);
+      },
+      py::arg("dem"));
+  m.def(
+      "get_errors_from_dem",
+      [](py::object dem) {
+        auto input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
+        return get_errors_from_dem(input_dem);
+      },
+      py::arg("dem"));
 
   // Not exposing sampling_from_dem and sample_shots because they depend on
   // stim::SparseShot which stim doesn't expose to python.