Expose det_order building in the python module (#65)

- The flag `--det-order-bfs` significantly improves the performance of
tesseract for decoding e.g. color code circuit-level noise
- There is no way to use it from python currently because the logic is
contained in `tesseract_main.cc`
- Here, we move it into utils and expose it in the python module

Author: noajshu

src/py/utils_test.py

@@ -44,6 +44,18 @@ def test_build_detector_graph():
     ]
 
 
+def test_build_det_orders():
+    assert tesseract_decoder.utils.build_det_orders(
+        _DETECTOR_ERROR_MODEL, num_det_orders=1, seed=0
+    ) == [[0, 1]]
+
+
+def test_build_det_orders_no_bfs():
+    assert tesseract_decoder.utils.build_det_orders(
+        _DETECTOR_ERROR_MODEL, num_det_orders=1, det_order_bfs=False, seed=0
+    ) == [[0, 1]]
+
+
 def test_get_errors_from_dem():
     expected = "Error{cost=1.945910, symptom=Symptom{D0 }}, Error{cost=0.510826, symptom=Symptom{D0 D1 }}, Error{cost=1.098612, symptom=Symptom{D1 }}"
     assert (

src/tesseract_main.cc

@@ -171,12 +171,8 @@ struct Args {
 
     // Sample orientations of the error model to use for the det priority
     {
-      config.det_orders.resize(num_det_orders);
-      std::mt19937_64 rng(det_order_seed);
-      std::normal_distribution<double> dist(/*mean=*/0, /*stddev=*/1);
-
-      std::vector<std::vector<double>> detector_coords = get_detector_coords(config.dem);
       if (verbose) {
+        auto detector_coords = get_detector_coords(config.dem);
         for (size_t d = 0; d < detector_coords.size(); ++d) {
           std::cout << "Detector D" << d << " coordinate (";
           size_t e = std::min(3ul, detector_coords[d].size());
@@ -187,88 +183,8 @@ struct Args {
           std::cout << ")" << std::endl;
         }
       }
-
-      if (det_order_bfs) {
-        auto graph = build_detector_graph(config.dem);
-        std::uniform_int_distribution<size_t> dist_det(0, graph.size() - 1);
-        for (size_t det_order = 0; det_order < num_det_orders; ++det_order) {
-          std::vector<size_t> perm;
-          perm.reserve(graph.size());
-          std::vector<char> visited(graph.size(), false);
-          std::queue<size_t> q;
-          size_t start = dist_det(rng);
-          while (perm.size() < graph.size()) {
-            if (!visited[start]) {
-              visited[start] = true;
-              q.push(start);
-              perm.push_back(start);
-            }
-            while (!q.empty()) {
-              size_t cur = q.front();
-              q.pop();
-              auto neigh = graph[cur];
-              std::shuffle(neigh.begin(), neigh.end(), rng);
-              for (size_t n : neigh) {
-                if (!visited[n]) {
-                  visited[n] = true;
-                  q.push(n);
-                  perm.push_back(n);
-                }
-              }
-            }
-            if (perm.size() < graph.size()) {
-              do {
-                start = dist_det(rng);
-              } while (visited[start]);
-            }
-          }
-          std::vector<size_t> inv_perm(graph.size());
-          for (size_t i = 0; i < perm.size(); ++i) {
-            inv_perm[perm[i]] = i;
-          }
-          config.det_orders[det_order] = inv_perm;
-        }
-      } else {
-        std::vector<double> inner_products(config.dem.count_detectors());
-
-        if (!detector_coords.size() || !detector_coords.at(0).size()) {
-          // If there are no detector coordinates, just use the standard
-          // ordering of the indices.
-          for (size_t det_order = 0; det_order < num_det_orders; ++det_order) {
-            config.det_orders[det_order].resize(config.dem.count_detectors());
-            std::iota(config.det_orders[det_order].begin(), config.det_orders[det_order].end(), 0);
-          }
-
-        } else {
-          // Use the coordinates to order the detectors based on a random
-          // orientation
-          for (size_t det_order = 0; det_order < num_det_orders; ++det_order) {
-            // Sample a direction
-            std::vector<double> orientation_vector;
-            for (size_t i = 0; i < detector_coords.at(0).size(); ++i) {
-              orientation_vector.push_back(dist(rng));
-            }
-
-            for (size_t i = 0; i < detector_coords.size(); ++i) {
-              inner_products[i] = 0;
-              for (size_t j = 0; j < orientation_vector.size(); ++j) {
-                inner_products[i] += detector_coords[i][j] * orientation_vector[j];
-              }
-            }
-            std::vector<size_t> perm(config.dem.count_detectors());
-            std::iota(perm.begin(), perm.end(), 0);
-            std::sort(perm.begin(), perm.end(), [&](const size_t& i, const size_t& j) {
-              return inner_products[i] > inner_products[j];
-            });
-            // Invert the permutation
-            std::vector<size_t> inv_perm(config.dem.count_detectors());
-            for (size_t i = 0; i < perm.size(); ++i) {
-              inv_perm[perm[i]] = i;
-            }
-            config.det_orders[det_order] = inv_perm;
-          }
-        }
-      }
+      config.det_orders =
+          build_det_orders(config.dem, num_det_orders, det_order_bfs, det_order_seed);
     }
 
     if (sample_num_shots > 0) {

src/utils.cc

@@ -19,13 +19,15 @@
 #include <filesystem>
 #include <fstream>
 #include <iostream>
+#include <numeric>
+#include <queue>
 #include <random>
 #include <string>
 
 #include "common.h"
 #include "stim.h"
 
-std::vector<std::vector<double>> get_detector_coords(stim::DetectorErrorModel& dem) {
+std::vector<std::vector<double>> get_detector_coords(const stim::DetectorErrorModel& dem) {
   std::vector<std::vector<double>> detector_coords;
   for (const stim::DemInstruction& instruction : dem.flattened().instructions) {
     switch (instruction.type) {
@@ -79,6 +81,91 @@ std::vector<std::vector<size_t>> build_detector_graph(const stim::DetectorErrorM
   return neighbors;
 }
 
+std::vector<std::vector<size_t>> build_det_orders(const stim::DetectorErrorModel& dem,
+                                                  size_t num_det_orders, bool det_order_bfs,
+                                                  uint64_t seed) {
+  std::vector<std::vector<size_t>> det_orders(num_det_orders);
+  std::mt19937_64 rng(seed);
+  std::normal_distribution<double> dist(0, 1);
+
+  auto detector_coords = get_detector_coords(dem);
+
+  if (det_order_bfs) {
+    auto graph = build_detector_graph(dem);
+    std::uniform_int_distribution<size_t> dist_det(0, graph.size() - 1);
+    for (size_t det_order = 0; det_order < num_det_orders; ++det_order) {
+      std::vector<size_t> perm;
+      perm.reserve(graph.size());
+      std::vector<char> visited(graph.size(), false);
+      std::queue<size_t> q;
+      size_t start = dist_det(rng);
+      while (perm.size() < graph.size()) {
+        if (!visited[start]) {
+          visited[start] = true;
+          q.push(start);
+          perm.push_back(start);
+        }
+        while (!q.empty()) {
+          size_t cur = q.front();
+          q.pop();
+          auto neigh = graph[cur];
+          std::shuffle(neigh.begin(), neigh.end(), rng);
+          for (size_t n : neigh) {
+            if (!visited[n]) {
+              visited[n] = true;
+              q.push(n);
+              perm.push_back(n);
+            }
+          }
+        }
+        if (perm.size() < graph.size()) {
+          do {
+            start = dist_det(rng);
+          } while (visited[start]);
+        }
+      }
+      std::vector<size_t> inv_perm(graph.size());
+      for (size_t i = 0; i < perm.size(); ++i) {
+        inv_perm[perm[i]] = i;
+      }
+      det_orders[det_order] = inv_perm;
+    }
+  } else {
+    std::vector<double> inner_products(dem.count_detectors());
+    if (!detector_coords.size() || !detector_coords.at(0).size()) {
+      for (size_t det_order = 0; det_order < num_det_orders; ++det_order) {
+        det_orders[det_order].resize(dem.count_detectors());
+        std::iota(det_orders[det_order].begin(), det_orders[det_order].end(), 0);
+      }
+    } else {
+      for (size_t det_order = 0; det_order < num_det_orders; ++det_order) {
+        std::vector<double> orientation_vector;
+        for (size_t i = 0; i < detector_coords.at(0).size(); ++i) {
+          orientation_vector.push_back(dist(rng));
+        }
+
+        for (size_t i = 0; i < detector_coords.size(); ++i) {
+          inner_products[i] = 0;
+          for (size_t j = 0; j < orientation_vector.size(); ++j) {
+            inner_products[i] += detector_coords[i][j] * orientation_vector[j];
+          }
+        }
+        std::vector<size_t> perm(dem.count_detectors());
+        std::iota(perm.begin(), perm.end(), 0);
+        std::sort(perm.begin(), perm.end(), [&](const size_t& i, const size_t& j) {
+          return inner_products[i] > inner_products[j];
+        });
+        std::vector<size_t> inv_perm(dem.count_detectors());
+        for (size_t i = 0; i < perm.size(); ++i) {
+          inv_perm[perm[i]] = i;
+        }
+        det_orders[det_order] = inv_perm;
+      }
+    }
+  }
+  return det_orders;
+}
+
 bool sampling_from_dem(uint64_t seed, size_t num_shots, stim::DetectorErrorModel dem,
                        std::vector<stim::SparseShot>& shots) {
   stim::DemSampler<stim::MAX_BITWORD_WIDTH> sampler(dem, std::mt19937_64{seed}, num_shots);

src/utils.h

@@ -28,12 +28,16 @@
 
 constexpr const double EPSILON = 1e-7;
 
-std::vector<std::vector<double>> get_detector_coords(stim::DetectorErrorModel& dem);
+std::vector<std::vector<double>> get_detector_coords(const stim::DetectorErrorModel& dem);
 
 // Builds an adjacency list graph where two detectors share an edge iff an error
 // in the model activates them both.
 std::vector<std::vector<size_t>> build_detector_graph(const stim::DetectorErrorModel& dem);
 
+std::vector<std::vector<size_t>> build_det_orders(const stim::DetectorErrorModel& dem,
+                                                  size_t num_det_orders, bool det_order_bfs = true,
+                                                  uint64_t seed = 0);
+
 const double INF = std::numeric_limits<double>::infinity();
 
 bool sampling_from_dem(uint64_t seed, size_t num_shots, stim::DetectorErrorModel dem,

src/utils.pybind.h

@@ -42,6 +42,14 @@ void add_utils_module(py::module &root) {
         return build_detector_graph(input_dem);
       },
       py::arg("dem"));
+  m.def(
+      "build_det_orders",
+      [](py::object dem, size_t num_det_orders, bool det_order_bfs, uint64_t seed) {
+        auto input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
+        return build_det_orders(input_dem, num_det_orders, det_order_bfs, seed);
+      },
+      py::arg("dem"), py::arg("num_det_orders"), py::arg("det_order_bfs") = true,
+      py::arg("seed") = 0);
   m.def(
       "get_errors_from_dem",
       [](py::object dem) {