Merge pull request #33 from noajshu/codex/refactor-tesseractconfig-and-simplexconfig-for-verbose-callb

feat: add verbose callback for decoders

Author: noajshu

src/py/common_test.py

@@ -15,7 +15,7 @@
 import pytest
 import stim
 
-from src import tesseract_decoder
+import tesseract_decoder
 
 def get_set_bits(n):
     """

src/py/simplex_test.py

@@ -16,7 +16,7 @@
 import pytest
 import stim
 
-from src import tesseract_decoder
+import tesseract_decoder
 
 _DETECTOR_ERROR_MODEL = stim.DetectorErrorModel(
     """
@@ -33,7 +33,7 @@ def test_create_simplex_config():
     assert sc.window_length == 5
     assert (
         str(sc)
-        == "SimplexConfig(dem=DetectorErrorModel_Object, window_length=5, window_slide_length=0, verbose=0)"
+        == "SimplexConfig(dem=DetectorErrorModel_Object, window_length=5, window_slide_length=0)"
     )
 
 
@@ -46,6 +46,20 @@ def test_create_simplex_decoder():
     assert decoder.cost_from_errors([2]) == pytest.approx(1.0986123)
     assert decoder.decode([1]) == []
 
+
+def test_simplex_verbose_callback_receives_output():
+    lines = []
+
+    def cb(s: str) -> None:
+        lines.append(s)
+
+    config = tesseract_decoder.simplex.SimplexConfig(
+        _DETECTOR_ERROR_MODEL, window_length=5, verbose_callback=cb
+    )
+    decoder = tesseract_decoder.simplex.SimplexDecoder(config)
+    decoder.decode_to_errors([1])
+    assert any(lines)
+
 def test_simplex_decoder_predicts_various_observable_flips():
     """
     Tests that the Simplex decoder correctly predicts a logical observable

src/py/tesseract_test.py

@@ -16,7 +16,7 @@
 import pytest
 import stim
 
-from src import tesseract_decoder
+import tesseract_decoder
 
 _DETECTOR_ERROR_MODEL = stim.DetectorErrorModel(
     """
@@ -30,7 +30,7 @@
 def test_create_config():
     assert (
         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)"
+        == "TesseractConfig(dem=DetectorErrorModel_Object, det_beam=65535, no_revisit_dets=0, at_most_two_errors_per_detector=0, pqlimit=18446744073709551615, det_orders=[], det_penalty=0)"
     )
     assert (
         tesseract_decoder.tesseract.TesseractConfig(_DETECTOR_ERROR_MODEL).dem
@@ -52,6 +52,20 @@ def test_create_decoder():
     assert decoder.cost_from_errors([1]) == pytest.approx(0.5108256237659907)
     assert decoder.decode([0]) == []
 
+
+def test_tesseract_verbose_callback_receives_output():
+    lines = []
+
+    def cb(s: str) -> None:
+        lines.append(s)
+
+    config = tesseract_decoder.tesseract.TesseractConfig(
+        _DETECTOR_ERROR_MODEL, verbose_callback=cb
+    )
+    decoder = tesseract_decoder.tesseract.TesseractDecoder(config)
+    decoder.decode_to_errors([0])
+    assert any(lines)
+
 def test_tesseract_decoder_predicts_various_observable_flips():
     """
     Tests that the Tesseract decoder correctly predicts a logical observable

src/py/utils_test.py

@@ -16,7 +16,7 @@
 import pytest
 import stim
 
-from src import tesseract_decoder
+import tesseract_decoder
 
 
 _DETECTOR_ERROR_MODEL = stim.DetectorErrorModel(

src/simplex.cc

@@ -15,24 +15,36 @@
 #include "simplex.h"
 
 #include <cassert>
+#include <iostream>
 
 #include "Highs.h"
 #include "io/HMPSIO.h"
 
 constexpr size_t T_COORD = 2;
 
+namespace {
+void highs_log_cb(HighsLogType, const char* msg, void* user_data) {
+  CallbackStream* stream = static_cast<CallbackStream*>(user_data);
+  (*stream) << msg;
+  stream->flush();
+}
+}  // namespace
+
 std::string SimplexConfig::str() {
   auto& self = *this;
   std::stringstream ss;
   ss << "SimplexConfig(";
   ss << "dem=" << "DetectorErrorModel_Object" << ", ";
   ss << "window_length=" << self.window_length << ", ";
-  ss << "window_slide_length=" << self.window_slide_length << ", ";
-  ss << "verbose=" << self.verbose << ")";
+  ss << "window_slide_length=" << self.window_slide_length << ")";
   return ss.str();
 }
 
 SimplexDecoder::SimplexDecoder(SimplexConfig _config) : config(_config) {
+  if (!config.verbose_callback) {
+    config.verbose_callback = [](const std::string& s) { std::cout << s; };
+  }
+  config.log_stream.callback = config.verbose_callback;
   config.dem = common::remove_zero_probability_errors(config.dem);
   std::vector<double> detector_t_coords(config.dem.count_detectors());
   for (const stim::DemInstruction& instruction : config.dem.flattened().instructions) {
@@ -152,7 +164,11 @@ void SimplexDecoder::init_ilp() {
   // Disabled presolve entirely after encountering bugs similar to this one:
   // https://github.com/ERGO-Code/HiGHS/issues/1273
   highs->setOptionValue("presolve", "off");
-  highs->setOptionValue("output_flag", config.verbose);
+  highs->setOptionValue("output_flag", config.log_stream.active);
+  highs->setOptionValue("log_to_console", false);
+  if (config.log_stream.active) {
+    highs->setLogCallback(highs_log_cb, &config.log_stream);
+  }
 }
 
 void SimplexDecoder::decode_to_errors(const std::vector<uint64_t>& detections) {
@@ -197,9 +213,7 @@ void SimplexDecoder::decode_to_errors(const std::vector<uint64_t>& detections) {
         add_costs_for_time(t1);
         ++t1;
       }
-      if (config.verbose) {
-        std::cout << "t0 = " << t0 << " t1 = " << t1 << std::endl;
-      }
+      config.log_stream << "t0 = " << t0 << " t1 = " << t1 << std::endl;
 
       // Pass the model to HiGHS
       *return_status = highs->passModel(*model);
@@ -235,16 +249,20 @@ void SimplexDecoder::decode_to_errors(const std::vector<uint64_t>& detections) {
       }
       assert(*return_status == HighsStatus::kOk);
 
-      if (config.verbose) {
-        // Get the solution information
+      if (config.log_stream.active) {
         const HighsInfo& info = highs->getInfo();
-        std::cout << "Simplex iteration count: " << info.simplex_iteration_count << std::endl;
-        std::cout << "Objective function value: " << info.objective_function_value << std::endl;
-        std::cout << "Primal  solution status: "
-                  << highs->solutionStatusToString(info.primal_solution_status) << std::endl;
-        std::cout << "Dual    solution status: "
-                  << highs->solutionStatusToString(info.dual_solution_status) << std::endl;
-        std::cout << "Basis: " << highs->basisValidityToString(info.basis_validity) << std::endl;
+        config.log_stream << "Simplex iteration count: " << info.simplex_iteration_count
+                          << std::endl;
+        config.log_stream << "Objective function value: " << info.objective_function_value
+                          << std::endl;
+        config.log_stream << "Primal  solution status: "
+                          << highs->solutionStatusToString(info.primal_solution_status)
+                          << std::endl;
+        config.log_stream << "Dual    solution status: "
+                          << highs->solutionStatusToString(info.dual_solution_status)
+                          << std::endl;
+        config.log_stream << "Basis: "
+                           << highs->basisValidityToString(info.basis_validity) << std::endl;
       }
 
       // Get the model status
@@ -286,16 +304,20 @@ void SimplexDecoder::decode_to_errors(const std::vector<uint64_t>& detections) {
     *return_status = highs->run();
     assert(*return_status == HighsStatus::kOk);
 
-    if (config.verbose) {
-      // Get the solution information
+    if (config.log_stream.active) {
       const HighsInfo& info = highs->getInfo();
-      std::cout << "Simplex iteration count: " << info.simplex_iteration_count << std::endl;
-      std::cout << "Objective function value: " << info.objective_function_value << std::endl;
-      std::cout << "Primal  solution status: "
-                << highs->solutionStatusToString(info.primal_solution_status) << std::endl;
-      std::cout << "Dual    solution status: "
-                << highs->solutionStatusToString(info.dual_solution_status) << std::endl;
-      std::cout << "Basis: " << highs->basisValidityToString(info.basis_validity) << std::endl;
+      config.log_stream << "Simplex iteration count: " << info.simplex_iteration_count
+                        << std::endl;
+      config.log_stream << "Objective function value: " << info.objective_function_value
+                        << std::endl;
+      config.log_stream << "Primal  solution status: "
+                        << highs->solutionStatusToString(info.primal_solution_status)
+                        << std::endl;
+      config.log_stream << "Dual    solution status: "
+                        << highs->solutionStatusToString(info.dual_solution_status)
+                        << std::endl;
+      config.log_stream << "Basis: "
+                         << highs->basisValidityToString(info.basis_validity) << std::endl;
     }
 
     // Get the model status

src/simplex.h

@@ -14,11 +14,13 @@
 
 #ifndef SIMPLEX_HPP
 #define SIMPLEX_HPP
+#include <functional>
 #include <unordered_set>
 #include <vector>
 
 #include "common.h"
 #include "stim.h"
+#include "utils.h"
 
 struct HighsModel;
 struct Highs;
@@ -29,7 +31,8 @@ struct SimplexConfig {
   bool parallelize = false;
   size_t window_length = 0;
   size_t window_slide_length = 0;
-  bool verbose = false;
+  std::function<void(const std::string&)> verbose_callback;
+  CallbackStream log_stream;
   bool windowing_enabled() {
     return (window_length != 0);
   }

src/simplex.pybind.h

@@ -18,6 +18,7 @@
 #include <pybind11/operators.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <iostream>
 
 #include "common.h"
 #include "simplex.h"
@@ -28,9 +29,26 @@ 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) {
+                                   py::object verbose_callback = py::none()) {
   stim::DetectorErrorModel input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
-  return SimplexConfig({input_dem, parallelize, window_length, window_slide_length, verbose});
+  SimplexConfig cfg;
+  cfg.dem = input_dem;
+  cfg.parallelize = parallelize;
+  cfg.window_length = window_length;
+  cfg.window_slide_length = window_slide_length;
+  std::function<void(const std::string&)> cb;
+  bool active = false;
+  if (!verbose_callback.is_none()) {
+    py::function f = verbose_callback;
+    cb = [f](const std::string& s) {
+      py::gil_scoped_acquire gil;
+      f(s);
+    };
+    active = true;
+  }
+  cfg.verbose_callback = cb;
+  cfg.log_stream = CallbackStream(active, cfg.verbose_callback);
+  return cfg;
 }
 
 };  // namespace
@@ -42,12 +60,11 @@ void add_simplex_module(py::module& root) {
   py::class_<SimplexConfig>(m, "SimplexConfig")
       .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)
+           py::arg("verbose_callback") = py::none())
       .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)
-      .def_readwrite("verbose", &SimplexConfig::verbose)
       .def("windowing_enabled", &SimplexConfig::windowing_enabled)
       .def("__str__", &SimplexConfig::str);
 

src/simplex_main.cc

@@ -15,6 +15,7 @@
 #include <argparse/argparse.hpp>
 #include <atomic>
 #include <fstream>
+#include <iostream>
 #include <nlohmann/json.hpp>
 #include <thread>
 
@@ -134,6 +135,8 @@ struct Args {
 
   void extract(SimplexConfig& config, std::vector<stim::SparseShot>& shots,
                std::unique_ptr<stim::MeasureRecordWriter>& writer) {
+    config.verbose_callback = [](const std::string& s) { std::cout << s; };
+    config.log_stream = CallbackStream(verbose, config.verbose_callback);
     // Get a circuit, if available
     stim::Circuit circuit;
     if (!circuit_path.empty()) {
@@ -261,7 +264,8 @@ struct Args {
     config.parallelize = enable_ilp_solver_parallelism;
     config.window_length = window_length;
     config.window_slide_length = window_slide_length;
-    config.verbose = verbose;
+    config.log_stream.active = verbose;
+    config.log_stream.callback = config.verbose_callback;
   }
 };