Add missing docstrings and tests inside Python interface (#80)

Address issue #54

I have added docstrings where they were missing for each class we have
in Python Interface. Also, inside the `shared_decoding_tests` module
that I created, I have moved tests that are the same for `Simplex` and
`Tesseract`, so we have it more generic.

Our Python Interface/wrapper should look much prettier now.

---------

Co-authored-by: oscarhiggott <29460323+oscarhiggott@users.noreply.github.com>

Author: draganaurosgrbic

src/common.cc

@@ -25,9 +25,16 @@ std::string common::Symptom::str() {
 }
 
 common::Error::Error(const stim::DemInstruction& error) {
+  if (error.type != stim::DemInstructionType::DEM_ERROR) {
+    throw std::invalid_argument(
+        "Error must be loaded from an error dem instruction, but received: " + error.str());
+  }
   assert(error.type == stim::DemInstructionType::DEM_ERROR);
-  probability = error.arg_data[0];
-  assert(probability >= 0 && probability <= 1);
+  double probability = error.arg_data[0];
+  if (probability < 0 || probability > 1) {
+    throw std::invalid_argument("Probability must be between 0 and 1, but received: " +
+                                std::to_string(probability));
+  }
 
   std::set<int> detectors_set;
   std::set<int> observables_set;
@@ -60,6 +67,17 @@ std::string common::Error::str() {
   return "Error{cost=" + std::to_string(likelihood_cost) + ", symptom=" + symptom.str() + "}";
 }
 
+double common::Error::get_probability() const {
+  return 1.0 / (1.0 + std::exp(likelihood_cost));
+}
+
+void common::Error::set_with_probability(double p) {
+  if (p <= 0 || p >= 1) {
+    throw std::invalid_argument("Probability must be between 0 and 1.");
+  }
+  likelihood_cost = -std::log(p / (1.0 - p));
+}
+
 std::vector<stim::DemTarget> common::Symptom::as_dem_instruction_targets() const {
   std::vector<stim::DemTarget> targets;
   for (int d : detectors) {
@@ -71,7 +89,15 @@ std::vector<stim::DemTarget> common::Symptom::as_dem_instruction_targets() const
   return targets;
 }
 
-stim::DetectorErrorModel common::merge_identical_errors(const stim::DetectorErrorModel& dem) {
+double common::merge_weights(double a, double b) {
+  auto sgn = std::copysign(1, a) * std::copysign(1, b);
+  auto signed_min = sgn * std::min(std::abs(a), std::abs(b));
+  return signed_min + std::log(1 + std::exp(-std::abs(a + b))) -
+         std::log(1 + std::exp(-std::abs(a - b)));
+}
+
+stim::DetectorErrorModel common::merge_indistinguishable_errors(
+    const stim::DetectorErrorModel& dem) {
   stim::DetectorErrorModel out_dem;
 
   // Map to track the distinct symptoms
@@ -80,13 +106,14 @@ stim::DetectorErrorModel common::merge_identical_errors(const stim::DetectorErro
     switch (instruction.type) {
       case stim::DemInstructionType::DEM_ERROR: {
         Error error(instruction);
-        assert(error.symptom.detectors.size());
-        // Merge with existing error with the same symptom (if applicable)
+        if (error.symptom.detectors.size() == 0) {
+          throw std::invalid_argument("Errors that do not flip any detectors are not supported.");
+        }
+
         if (errors_by_symptom.find(error.symptom) != errors_by_symptom.end()) {
-          double p0 = errors_by_symptom[error.symptom].probability;
-          error.probability = p0 * (1 - error.probability) + (1 - p0) * error.probability;
+          error.likelihood_cost = merge_weights(error.likelihood_cost,
+                                                errors_by_symptom[error.symptom].likelihood_cost);
         }
-        error.likelihood_cost = -1 * std::log(error.probability / (1 - error.probability));
         errors_by_symptom[error.symptom] = error;
         break;
       }
@@ -103,7 +130,7 @@ stim::DetectorErrorModel common::merge_identical_errors(const stim::DetectorErro
     }
   }
   for (const auto& it : errors_by_symptom) {
-    out_dem.append_error_instruction(it.second.probability,
+    out_dem.append_error_instruction(it.second.get_probability(),
                                      it.second.symptom.as_dem_instruction_targets(),
                                      /*tag=*/"");
   }

src/common.h

@@ -48,26 +48,23 @@ struct Symptom {
 // Represents an error / weighted hyperedge
 struct Error {
   double likelihood_cost;
-  double probability;
   Symptom symptom;
-  std::vector<bool> dets_array;
   Error() = default;
-  Error(double likelihood_cost, std::vector<int>& detectors, std::vector<int> observables,
-        std::vector<bool>& dets_array)
-      : likelihood_cost(likelihood_cost), symptom{detectors, observables}, dets_array(dets_array) {}
-  Error(double likelihood_cost, double probability, std::vector<int>& detectors,
-        std::vector<int> observables, std::vector<bool>& dets_array)
-      : likelihood_cost(likelihood_cost),
-        probability(probability),
-        symptom{detectors, observables},
-        dets_array(dets_array) {}
+  Error(double likelihood_cost, std::vector<int>& detectors, std::vector<int> observables)
+      : likelihood_cost(likelihood_cost), symptom{detectors, observables} {}
   Error(const stim::DemInstruction& error);
   std::string str();
+
+  // Get/calculate the probability from the likelihood cost.
+  double get_probability() const;
+
+  // Set/calculate the likelihood cost from a probability.
+  void set_with_probability(double p);
 };
 
 // Makes a new (flattened) dem where identical error mechanisms have been
 // merged.
-stim::DetectorErrorModel merge_identical_errors(const stim::DetectorErrorModel& dem);
+stim::DetectorErrorModel merge_indistinguishable_errors(const stim::DetectorErrorModel& dem);
 
 // Returns a copy of the given error model with any zero-probability DEM_ERROR
 // instructions removed.
@@ -80,6 +77,12 @@ stim::DetectorErrorModel remove_zero_probability_errors(const stim::DetectorErro
 stim::DetectorErrorModel dem_from_counts(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',
+/// assuming each edge weight is a log-likelihood ratio log((1-p)/p) associated with the probability
+/// p of an error occurring on the edge, and that the error mechanisms associated with the two edges
+/// being merged are independent.
+double merge_weights(double a, double b);
+
 }  // namespace common
 
 #endif

src/common.pybind.h

@@ -29,65 +29,187 @@ namespace py = pybind11;
 void add_common_module(py::module &root) {
   auto m = root.def_submodule("common", "classes commonly used by the decoder");
 
-  py::class_<common::Symptom>(m, "Symptom")
+  py::class_<common::Symptom>(m, "Symptom", R"pbdoc(
+        Represents a symptom of an error, which is a list of detectors and a list of observables
+
+        A symptom is defined by a list of detectors that are flipped and a list of
+        observables that are flipped.
+    )pbdoc")
       .def(py::init<std::vector<int>, std::vector<int>>(),
-           py::arg("detectors") = std::vector<int>(), py::arg("observables") = std::vector<int>())
-      .def_readwrite("detectors", &common::Symptom::detectors)
-      .def_readwrite("observables", &common::Symptom::observables)
+           py::arg("detectors") = std::vector<int>(), py::arg("observables") = std::vector<int>(),
+           R"pbdoc(
+            The constructor for the `Symptom` class.
+
+            Parameters
+            ----------
+            detectors : list[int], default=[]
+                The indices of the detectors in this symptom.
+            observables : list[int], default=[]
+                The indices of the flipped observables.
+           )pbdoc")
+      .def_readwrite("detectors", &common::Symptom::detectors,
+                     "A list of the detector indices that are flipped in this symptom.")
+      .def_readwrite("observables", &common::Symptom::observables,
+                     "A list of observable indices that are flipped in this symptom.")
       .def("__str__", &common::Symptom::str)
       .def(py::self == py::self)
       .def(py::self != py::self)
-      .def("as_dem_instruction_targets", [](common::Symptom s) {
-        std::vector<py::object> ret;
-        for (auto &t : s.as_dem_instruction_targets())
-          ret.push_back(make_py_object(t, "DemTarget"));
-        return ret;
-      });
-
-  py::class_<common::Error>(m, "Error")
-      .def_readwrite("likelihood_cost", &common::Error::likelihood_cost)
-      .def_readwrite("probability", &common::Error::probability)
-      .def_readwrite("symptom", &common::Error::symptom)
+      .def(
+          "as_dem_instruction_targets",
+          [](common::Symptom s) {
+            std::vector<py::object> ret;
+            for (auto &t : s.as_dem_instruction_targets())
+              ret.push_back(make_py_object(t, "DemTarget"));
+            return ret;
+          },
+          R"pbdoc(
+        Converts the symptom into a list of `stim.DemTarget` objects.
+        
+        Returns
+        -------
+        list[stim.DemTarget]
+            A list of `stim.DemTarget` objects representing the detectors and observables.
+      )pbdoc");
+
+  py::class_<common::Error>(m, "Error", R"pbdoc(
+        Represents an error, including its cost, and symptom.
+
+        An error is a physical event (or set of indistinguishable physical events)
+        defined by the detectors and observables that it flips in the circuit.
+    )pbdoc")
+      .def_readwrite("likelihood_cost", &common::Error::likelihood_cost,
+                     "The cost of this error (often log((1 - probability) / probability)).")
+      .def_readwrite("symptom", &common::Error::symptom, "The symptom associated with this error.")
       .def("__str__", &common::Error::str)
-      .def(py::init<>())
-      .def(py::init<double, std::vector<int> &, std::vector<int>, std::vector<bool> &>(),
-           py::arg("likelihood_cost"), py::arg("detectors"), py::arg("observables"),
-           py::arg("dets_array"))
-      .def(py::init<double, double, std::vector<int> &, std::vector<int>, std::vector<bool> &>(),
-           py::arg("likelihood_cost"), py::arg("probability"), py::arg("detectors"),
-           py::arg("observables"), py::arg("dets_array"))
+      .def(py::init<>(), R"pbdoc(
+        Default constructor for the `Error` class.
+      )pbdoc")
+      .def(py::init<double, std::vector<int> &, std::vector<int>>(), py::arg("likelihood_cost"),
+           py::arg("detectors"), py::arg("observables"), R"pbdoc(
+            Constructor for the `Error` class.
+
+            Parameters
+            ----------
+            likelihood_cost : float
+                The cost of this error. 
+                This is often `log((1 - probability) / probability)`.
+            detectors : list[int]
+                A list of indices of the detectors flipped by this error.
+            observables : list[int]
+                A list of indices of the observables flipped by this error.
+           )pbdoc")
+
       .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"));
+           py::arg("error"), R"pbdoc(
+            Constructor that creates an `Error` from a `stim.DemInstruction`.
+
+            Parameters
+            ----------
+            error : stim.DemInstruction
+                The instruction to convert into an `Error` object.
+           )pbdoc")
+      .def("get_probability", &common::Error::get_probability,
+           R"pbdoc(
+            Gets the probability associated with the likelihood cost.
+
+            Returns
+            -------
+            float
+                The probability of the error, calculated from the likelihood cost.
+           )pbdoc")
+      .def("set_with_probability", &common::Error::set_with_probability, py::arg("probability"),
+           R"pbdoc(
+            Sets the likelihood cost based on a given probability.
+
+            Parameters
+            ----------
+            probability : float
+                The probability to use for setting the likelihood cost.
+                Must be between 0 and 1 (exclusive).
+
+            Raises
+            ------
+            ValueError
+                If the provided probability is not between 0 and 1.
+           )pbdoc");
 
   m.def(
-      "merge_identical_errors",
+      "merge_indistinguishable_errors",
       [](py::object dem) {
         auto input_dem = parse_py_object<stim::DetectorErrorModel>(dem);
-        auto res = common::merge_identical_errors(input_dem);
+        auto res = common::merge_indistinguishable_errors(input_dem);
         return make_py_object(res, "DetectorErrorModel");
       },
-      py::arg("dem"));
+      py::arg("dem"), R"pbdoc(
+        Merges identical errors in a `stim.DetectorErrorModel`.
+        
+        Errors are identical if they flip the same set of detectors and observables (the same symptom).
+        For example, two identical errors with probabilities p1 and p2
+        would be merged into a single error with the same symptom,
+        but with probability `p1 * (1 - p2) + p2 * (1 - p1)`
+
+        Parameters
+        ----------
+        dem : stim.DetectorErrorModel
+            The detector error model to process.
+
+        Returns
+        -------
+        stim.DetectorErrorModel
+            A new `DetectorErrorModel` with identical errors merged.
+      )pbdoc");
   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"));
+      py::arg("dem"), R"pbdoc(
+        Removes errors with a probability of 0 from a `stim.DetectorErrorModel`.
+
+        Parameters
+        ----------
+        dem : stim.DetectorErrorModel
+            The detector error model to process.
+
+        Returns
+        -------
+        stim.DetectorErrorModel
+            A new `DetectorErrorModel` with zero-probability errors removed.
+      )pbdoc");
   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"));
+      py::arg("orig_dem"), py::arg("error_counts"), py::arg("num_shots"), R"pbdoc(
+        Re-weights errors in a `stim.DetectorErrorModel` based on observed counts.
+
+        This function re-calculates the probability of each error based on a list of
+        observed counts and the total number of shots.
+
+        Parameters
+        ----------
+        orig_dem : stim.DetectorErrorModel
+            The original detector error model.
+        error_counts : list[int]
+            A list of counts for each error in the DEM.
+        num_shots : int
+            The total number of shots in the experiment.
+
+        Returns
+        -------
+        stim.DetectorErrorModel
+            A new `DetectorErrorModel` with updated error probabilities.
+      )pbdoc");
 }
 
 #endif