Report errors with original DEM indexing (#186)

A big flaw in Tesseract's API is that error indices are always in terms
of the internal `errors` vector, which does not correspond directly to
the errors in the (flattened) input DEM given by the user. This causes a
variety of problems. One such problem is that if we generate a DEM with
`--dem-out` (or e.g. similar manual processing steps in python) it may
bear little resemblance to the input DEM. For example all the targets
will be stripped of separators etc. This makes it annoying to use
tesseract-calibrated error models for downstream tasks like
matching-based decoding.

Here we adopt the principle that the user interface to Tesseract/Simplex
decoders should always be in terms of the error indices from the
original flattened DEM as provided by the user. This is now true across
C++, CLI, and Python APIs.

- Added index-mapping support to DEM preprocessing in common:
  - merge_indistinguishable_errors(..., error_index_map)
  - remove_zero_probability_errors(..., error_index_map)
- `error_index_map` maps original error index to new preprocessed index
- `error_index_map` maps removed / redundant errors to
std::numeric_limits<size_t>::max()
- Update both decoders (TesseractDecoder, SimplexDecoder) to maintain:
  - dem_error_to_error (original flattened DEM index -> internal index)
- error_to_dem_error (internal error index -> original flattened DEM
index)
- predicted_errors_buffer reports errors back with original flattened
DEM error indices.
- cost_from_errors and observables-from-errors methods now:
  - accept original flattened DEM indices
  - throw on unmapped/removed indices (size_t::max())
- Updated Python bindings to use the new helpers.
- Updated pybind common wrappers to pass required map args to common
preprocessing functions.
- Updated --dem-out in both CLI binaries:
  - keep original flattened DEM in scope
- emit updated probabilities by iterating original DEM instruction order
- preserve original error instruction tags and arbitrary formatting
(e.g. `D0 ^ D0 D1`) when writing estimated DEM output.
- Updated tests
- Added AGENTS guidance to run Python Bazel tests

Author: noajshu

AGENTS.md

@@ -19,6 +19,7 @@ bazel build src:tesseract src:simplex
 
 ```bash
 bazel test src:all
+bazel test //src/py:all
 ```
 
 ## Building with CMake

src/common.cc

@@ -11,11 +11,11 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-
 #include "common.h"
 
 #include <iomanip>
 #include <iostream>
+#include <limits>
 #include <sstream>
 #include <string>
 #include <vector>
@@ -117,11 +117,15 @@ double common::merge_weights(double a, double b) {
 }
 
 stim::DetectorErrorModel common::merge_indistinguishable_errors(
-    const stim::DetectorErrorModel& dem) {
+    const stim::DetectorErrorModel& dem, std::vector<size_t>& error_index_map) {
   stim::DetectorErrorModel out_dem;
 
-  // Map to track the distinct symptoms
-  std::unordered_map<Symptom, Error, Symptom::hash> errors_by_symptom;
+  error_index_map.clear();
+
+  // Map to track first-seen distinct symptoms.
+  std::unordered_map<Symptom, size_t, Symptom::hash> merged_index_by_symptom;
+  std::vector<Error> merged_errors;
+
   for (const stim::DemInstruction& instruction : dem.flattened().instructions) {
     switch (instruction.type) {
       case stim::DemInstructionType::DEM_ERROR: {
@@ -131,11 +135,18 @@ stim::DetectorErrorModel common::merge_indistinguishable_errors(
           std::cout << "Warning: the circuit has errors that do not flip any detectors \n";
         }
 
-        if (errors_by_symptom.find(error.symptom) != errors_by_symptom.end()) {
-          error.likelihood_cost = merge_weights(error.likelihood_cost,
-                                                errors_by_symptom[error.symptom].likelihood_cost);
+        auto it = merged_index_by_symptom.find(error.symptom);
+        if (it != merged_index_by_symptom.end()) {
+          size_t merged_error_index = it->second;
+          merged_errors[merged_error_index].likelihood_cost = merge_weights(
+              error.likelihood_cost, merged_errors[merged_error_index].likelihood_cost);
+          error_index_map.push_back(merged_error_index);
+        } else {
+          size_t merged_error_index = merged_errors.size();
+          merged_index_by_symptom[error.symptom] = merged_error_index;
+          merged_errors.push_back(error);
+          error_index_map.push_back(merged_error_index);
         }
-        errors_by_symptom[error.symptom] = error;
         break;
       }
       case stim::DemInstructionType::DEM_DETECTOR: {
@@ -150,22 +161,27 @@ stim::DetectorErrorModel common::merge_indistinguishable_errors(
         throw std::invalid_argument("Unrecognized instruction type: " + instruction.str());
     }
   }
-  for (const auto& it : errors_by_symptom) {
-    out_dem.append_error_instruction(it.second.get_probability(),
-                                     it.second.symptom.as_dem_instruction_targets(),
+  for (const auto& error : merged_errors) {
+    out_dem.append_error_instruction(error.get_probability(),
+                                     error.symptom.as_dem_instruction_targets(),
                                      /*tag=*/"");
   }
   return out_dem;
 }
 
 stim::DetectorErrorModel common::remove_zero_probability_errors(
-    const stim::DetectorErrorModel& dem) {
+    const stim::DetectorErrorModel& dem, std::vector<size_t>& error_index_map) {
   stim::DetectorErrorModel out_dem;
+  error_index_map.clear();
+  size_t output_error_index = 0;
   for (const stim::DemInstruction& instruction : dem.flattened().instructions) {
     switch (instruction.type) {
       case stim::DemInstructionType::DEM_ERROR:
         if (instruction.arg_data[0] > 0) {
           out_dem.append_dem_instruction(instruction);
+          error_index_map.push_back(output_error_index++);
+        } else {
+          error_index_map.push_back(std::numeric_limits<size_t>::max());
         }
         break;
       case stim::DemInstructionType::DEM_DETECTOR:
@@ -181,43 +197,46 @@ stim::DetectorErrorModel common::remove_zero_probability_errors(
   return out_dem;
 }
 
-stim::DetectorErrorModel common::dem_from_counts(stim::DetectorErrorModel& orig_dem,
+void common::chain_error_maps(std::vector<size_t>& base_map, const std::vector<size_t>& next_map) {
+  for (size_t& ei : base_map) {
+    if (ei != std::numeric_limits<size_t>::max()) {
+      ei = next_map[ei];
+    }
+  }
+}
+
+std::vector<size_t> common::invert_error_map(const std::vector<size_t>& error_map,
+                                             size_t num_output_errors) {
+  std::vector<size_t> inverted_map(num_output_errors, std::numeric_limits<size_t>::max());
+  for (size_t i = 0; i < error_map.size(); ++i) {
+    size_t mapped_index = error_map[i];
+    if (mapped_index != std::numeric_limits<size_t>::max() &&
+        inverted_map[mapped_index] == std::numeric_limits<size_t>::max()) {
+      inverted_map[mapped_index] = i;
+    }
+  }
+  return inverted_map;
+}
+
+stim::DetectorErrorModel common::dem_from_counts(const stim::DetectorErrorModel& orig_dem,
                                                  const std::vector<size_t>& error_counts,
                                                  size_t num_shots) {
-  if (orig_dem.count_errors() != error_counts.size()) {
+  stim::DetectorErrorModel flat_dem = orig_dem.flattened();
+  if (flat_dem.count_errors() != error_counts.size()) {
     throw std::invalid_argument(
         "Error hits array must be the same size as the number of errors in the "
         "original DEM.");
   }
 
-  for (const stim::DemInstruction& instruction : orig_dem.flattened().instructions) {
-    if (instruction.type == stim::DemInstructionType::DEM_ERROR && instruction.arg_data[0] == 0) {
-      throw std::invalid_argument(
-          "dem_from_counts requires DEMs without zero-probability errors. Use"
-          " remove_zero_probability_errors first.");
-    }
-  }
-
   stim::DetectorErrorModel out_dem;
-  size_t ei = 0;
-  for (const stim::DemInstruction& instruction : orig_dem.flattened().instructions) {
-    switch (instruction.type) {
-      case stim::DemInstructionType::DEM_ERROR: {
-        double est_probability = double(error_counts.at(ei)) / double(num_shots);
-        out_dem.append_error_instruction(est_probability, instruction.target_data, /*tag=*/"");
-        ++ei;
-        break;
-      }
-      case stim::DemInstructionType::DEM_DETECTOR: {
-        out_dem.append_dem_instruction(instruction);
-        break;
-      }
-      case stim::DemInstructionType::DEM_LOGICAL_OBSERVABLE: {
-        out_dem.append_dem_instruction(instruction);
-        break;
-      }
-      default:
-        throw std::invalid_argument("Unrecognized instruction type: " + instruction.str());
+  size_t error_index = 0;
+  for (const stim::DemInstruction& instruction : flat_dem.instructions) {
+    if (instruction.type == stim::DemInstructionType::DEM_ERROR) {
+      double est_probability = double(error_counts.at(error_index)) / double(num_shots);
+      out_dem.append_error_instruction(est_probability, instruction.target_data, instruction.tag);
+      ++error_index;
+    } else {
+      out_dem.append_dem_instruction(instruction);
     }
   }
   return out_dem;

src/common.h

@@ -73,17 +73,31 @@ struct Error {
 
 // Makes a new (flattened) dem where identical error mechanisms have been
 // merged.
-stim::DetectorErrorModel merge_indistinguishable_errors(const stim::DetectorErrorModel& dem);
+// `error_index_map[old_error_index]` gives the corresponding merged DEM error
+// index in the returned DEM.
+stim::DetectorErrorModel merge_indistinguishable_errors(const stim::DetectorErrorModel& dem,
+                                                        std::vector<size_t>& error_index_map);
 
 // Returns a copy of the given error model with any zero-probability DEM_ERROR
 // instructions removed.
-stim::DetectorErrorModel remove_zero_probability_errors(const stim::DetectorErrorModel& dem);
+// `error_index_map[old_error_index]` gives the corresponding retained DEM error
+// index in the returned DEM, or `std::numeric_limits<size_t>::max()` if the
+// error was removed.
+stim::DetectorErrorModel remove_zero_probability_errors(const stim::DetectorErrorModel& dem,
+                                                        std::vector<size_t>& error_index_map);
+
+// Updates the base_map by chaining it with next_map.
+// base_map[i] = next_map[base_map[i]]
+void chain_error_maps(std::vector<size_t>& base_map, const std::vector<size_t>& next_map);
+
+// Inverts the error_map to create a mapping from output error indices back to
+// the first original error index that maps to it.
+std::vector<size_t> invert_error_map(const std::vector<size_t>& error_map,
+                                     size_t num_output_errors);
 
 // Makes a new dem where the probabilities of errors are estimated from the
 // fraction of shots they were used in.
-// Throws std::invalid_argument if `orig_dem` contains zero-probability errors;
-// call remove_zero_probability_errors first.
-stim::DetectorErrorModel dem_from_counts(stim::DetectorErrorModel& orig_dem,
+stim::DetectorErrorModel dem_from_counts(const stim::DetectorErrorModel& orig_dem,
                                          const std::vector<size_t>& error_counts, size_t num_shots);
 
 /// Computes the weight of an edge resulting from merging edges with weight `a' and weight `b',

src/common.pybind.h

@@ -142,7 +142,8 @@ void add_common_module(py::module& root) {
       "merge_indistinguishable_errors",
       [](py::object dem) {
         auto input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
-        auto res = common::merge_indistinguishable_errors(input_dem);
+        std::vector<size_t> error_index_map;
+        auto res = common::merge_indistinguishable_errors(input_dem, error_index_map);
         return make_py_object(res, "DetectorErrorModel");
       },
       py::arg("dem"), R"pbdoc(
@@ -166,9 +167,13 @@ void add_common_module(py::module& root) {
   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");
+        return make_py_object(([&]() {
+                                std::vector<size_t> error_index_map;
+                                return common::remove_zero_probability_errors(
+                                    parse_py_object<stim::DetectorErrorModel>(dem),
+                                    error_index_map);
+                              })(),
+                              "DetectorErrorModel");
       },
       py::arg("dem"), R"pbdoc(
         Removes errors with a probability of 0 from a `stim.DetectorErrorModel`.

src/common.test.cc

@@ -26,7 +26,7 @@ TEST(common, ErrorsStructFromDemInstruction) {
   EXPECT_EQ(ES.symptom.observables, std::vector<int>{0});
 }
 
-TEST(common, DemFromCountsRejectsZeroProbabilityErrors) {
+TEST(common, DemFromCountsHandlesZeroProbabilityErrors) {
   stim::DetectorErrorModel dem(R"DEM(
         error(0.1) D0
         error(0) D1
@@ -38,20 +38,32 @@ TEST(common, DemFromCountsRejectsZeroProbabilityErrors) {
 
   std::vector<size_t> counts{1, 7, 4};
   size_t num_shots = 10;
-  EXPECT_THROW({ common::dem_from_counts(dem, counts, num_shots); }, std::invalid_argument);
-
-  stim::DetectorErrorModel cleaned = common::remove_zero_probability_errors(dem);
-  stim::DetectorErrorModel out_dem =
-      common::dem_from_counts(cleaned, std::vector<size_t>{1, 4}, num_shots);
+  stim::DetectorErrorModel out_dem = common::dem_from_counts(dem, counts, num_shots);
 
   auto flat = out_dem.flattened();
-  ASSERT_EQ(out_dem.count_errors(), 2);
-  ASSERT_GE(flat.instructions.size(), 2);
+  ASSERT_EQ(out_dem.count_errors(), 3);
+  ASSERT_GE(flat.instructions.size(), 3);
 
   EXPECT_EQ(flat.instructions[0].type, stim::DemInstructionType::DEM_ERROR);
   EXPECT_NEAR(flat.instructions[0].arg_data[0], 0.1, 1e-9);
-  ASSERT_EQ(flat.instructions[1].type, stim::DemInstructionType::DEM_ERROR);
-  EXPECT_NEAR(flat.instructions[1].arg_data[0], 0.4, 1e-9);
+  EXPECT_EQ(flat.instructions[1].type, stim::DemInstructionType::DEM_ERROR);
+  EXPECT_NEAR(flat.instructions[1].arg_data[0], 0.7, 1e-9);
+  EXPECT_EQ(flat.instructions[2].type, stim::DemInstructionType::DEM_ERROR);
+  EXPECT_NEAR(flat.instructions[2].arg_data[0], 0.4, 1e-9);
+
+  std::vector<size_t> error_index_map;
+  stim::DetectorErrorModel cleaned = common::remove_zero_probability_errors(dem, error_index_map);
+  stim::DetectorErrorModel out_dem_cleaned =
+      common::dem_from_counts(cleaned, std::vector<size_t>{1, 4}, num_shots);
+
+  auto flat_cleaned = out_dem_cleaned.flattened();
+  ASSERT_EQ(out_dem_cleaned.count_errors(), 2);
+  ASSERT_GE(flat_cleaned.instructions.size(), 2);
+
+  EXPECT_EQ(flat_cleaned.instructions[0].type, stim::DemInstructionType::DEM_ERROR);
+  EXPECT_NEAR(flat_cleaned.instructions[0].arg_data[0], 0.1, 1e-9);
+  ASSERT_EQ(flat_cleaned.instructions[1].type, stim::DemInstructionType::DEM_ERROR);
+  EXPECT_NEAR(flat_cleaned.instructions[1].arg_data[0], 0.4, 1e-9);
 }
 
 TEST(common, DemFromCountsSimpleTwoErrors) {
@@ -86,7 +98,8 @@ TEST(common, RemoveZeroProbabilityErrors) {
         detector(0, 0, 0) D2
       )DEM");
 
-  stim::DetectorErrorModel cleaned = common::remove_zero_probability_errors(dem);
+  std::vector<size_t> error_index_map;
+  stim::DetectorErrorModel cleaned = common::remove_zero_probability_errors(dem, error_index_map);
 
   EXPECT_EQ(cleaned.count_errors(), 2);
   auto flat = cleaned.flattened();
@@ -153,30 +166,31 @@ TEST(CommonTest, merge_indistinguishable_errors_two_errors) {
   double p2 = 0.2;
   double expected_merged_p = p1 * (1 - p2) + p2 * (1 - p1);
   auto dem1 = create_dem_with_two_errors(p1, p2);
-  auto merged_dem1 = common::merge_indistinguishable_errors(dem1);
+  std::vector<size_t> error_index_map;
+  auto merged_dem1 = common::merge_indistinguishable_errors(dem1, error_index_map);
   ASSERT_NEAR(get_merged_probability(merged_dem1), expected_merged_p, 1e-9);
 
   // Case 2: One low, one high probability.
   p1 = 0.1;
   p2 = 0.8;
   expected_merged_p = p1 * (1 - p2) + p2 * (1 - p1);
   auto dem2 = create_dem_with_two_errors(p1, p2);
-  auto merged_dem2 = common::merge_indistinguishable_errors(dem2);
+  auto merged_dem2 = common::merge_indistinguishable_errors(dem2, error_index_map);
   ASSERT_NEAR(get_merged_probability(merged_dem2), expected_merged_p, 1e-9);
 
   // Case 3: One high, one low probability.
   p1 = 0.8;
   p2 = 0.1;
   expected_merged_p = p1 * (1 - p2) + p2 * (1 - p1);
   auto dem3 = create_dem_with_two_errors(p1, p2);
-  auto merged_dem3 = common::merge_indistinguishable_errors(dem3);
+  auto merged_dem3 = common::merge_indistinguishable_errors(dem3, error_index_map);
   ASSERT_NEAR(get_merged_probability(merged_dem3), expected_merged_p, 1e-9);
 
   // Case 4: Both probabilities are high.
   p1 = 0.8;
   p2 = 0.9;
   expected_merged_p = p1 * (1 - p2) + p2 * (1 - p1);
   auto dem4 = create_dem_with_two_errors(p1, p2);
-  auto merged_dem4 = common::merge_indistinguishable_errors(dem4);
+  auto merged_dem4 = common::merge_indistinguishable_errors(dem4, error_index_map);
   ASSERT_NEAR(get_merged_probability(merged_dem4), expected_merged_p, 1e-9);
 }