Calculate implied edge weights from joint probabilities

src/pymatching/matching.py

@@ -210,6 +210,7 @@ def decode(self,
                _legacy_return_weight: bool = None,
                *,
                return_weight: bool = False,
+               enable_correlations: bool = False,
                **kwargs
                ) -> Union[np.ndarray, Tuple[np.ndarray, int]]:
         r"""
@@ -322,7 +323,9 @@ def decode(self,
                           "argument.", DeprecationWarning, stacklevel=2)
             return_weight = _legacy_return_weight
         detection_events = self._syndrome_array_to_detection_events(z)
-        correction, weight = self._matching_graph.decode(detection_events)
+        correction, weight = self._matching_graph.decode(
+            detection_events, enable_correlations=enable_correlations
+        )
         if return_weight:
             return correction, weight
         else:
@@ -334,7 +337,8 @@ def decode_batch(
             *,
             return_weights: bool = False,
             bit_packed_shots: bool = False,
-            bit_packed_predictions: bool = False) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+            bit_packed_predictions: bool = False,
+            enable_correlations: bool = False) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
         """
         Decode from a 2D `shots` array containing a batch of syndrome measurements. A faster
         alternative to using `pymatching.Matching.decode` and iterating over the shots in Python.
@@ -420,7 +424,8 @@ def decode_batch(
         predictions, weights = self._matching_graph.decode_batch(
             shots,
             bit_packed_predictions=bit_packed_predictions,
-            bit_packed_shots=bit_packed_shots
+            bit_packed_shots=bit_packed_shots,
+            enable_correlations=enable_correlations
         )
         if return_weights:
             return predictions, weights
@@ -1191,7 +1196,9 @@ def load_from_check_matrix(self,
                                                                                     faults_matrix)
 
     @staticmethod
-    def from_detector_error_model(model: 'stim.DetectorErrorModel') -> 'pymatching.Matching':
+    def from_detector_error_model(
+            model: 'stim.DetectorErrorModel', *, enable_correlations: bool = False
+    ) -> 'pymatching.Matching':
         """
         Constructs a `pymatching.Matching` object by loading from a `stim.DetectorErrorModel`.
 
@@ -1238,7 +1245,7 @@ def from_detector_error_model(model: 'stim.DetectorErrorModel') -> 'pymatching.M
         <pymatching.Matching object with 120 detectors, 0 boundary nodes, and 502 edges>
         """
         m = Matching()
-        m._load_from_detector_error_model(model)
+        m._load_from_detector_error_model(model, enable_correlations=enable_correlations)
         return m
 
     @staticmethod
@@ -1333,7 +1340,7 @@ def from_stim_circuit_file(stim_circuit_path: Union[str, Path]) -> 'pymatching.M
         m._matching_graph = _cpp_pm.stim_circuit_file_to_matching_graph(stim_circuit_path)
         return m
 
-    def _load_from_detector_error_model(self, model: 'stim.DetectorErrorModel') -> None:
+    def _load_from_detector_error_model(self, model: 'stim.DetectorErrorModel', *, enable_correlations: bool = False) -> None:
         try:
             import stim
         except ImportError:  # pragma no cover
@@ -1344,7 +1351,9 @@ def _load_from_detector_error_model(self, model: 'stim.DetectorErrorModel') -> N
             )
         if not isinstance(model, stim.DetectorErrorModel):
             raise TypeError(f"'model' must be `stim.DetectorErrorModel`. Instead, got: {type(model)}")
-        self._matching_graph = _cpp_pm.detector_error_model_to_matching_graph(str(model))
+        self._matching_graph = _cpp_pm.detector_error_model_to_matching_graph(
+            str(model), enable_correlations=enable_correlations
+        )
 
     @staticmethod
     def from_networkx(graph: nx.Graph, *, min_num_fault_ids: int = None) -> 'pymatching.Matching':

src/pymatching/sparse_blossom/driver/implied_weights.h

@@ -0,0 +1,36 @@
+// Copyright 2022 PyMatching Contributors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// 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.
+
+#ifndef PYMATCHING2_IMPLIED_WEIGHT_UNCONVERTED_H
+#define PYMATCHING2_IMPLIED_WEIGHT_UNCONVERTED_H
+
+#include "pymatching/sparse_blossom/ints.h"
+
+namespace pm {
+
+struct ImpliedWeightUnconverted {
+    size_t node1;
+    size_t node2;
+    double implied_weight;
+};
+
+struct ImpliedWeight {
+    weight_int* edge0_ptr;
+    weight_int* edge1_ptr;
+    weight_int implied_weight;
+};
+
+}  // namespace pm
+
+#endif  // PYMATCHING2_IMPLIED_WEIGHT_UNCONVERTED_H

src/pymatching/sparse_blossom/driver/mwpm_decoding.cc

@@ -74,7 +74,8 @@ pm::Mwpm pm::detector_error_model_to_mwpm(
     pm::weight_int num_distinct_weights,
     bool ensure_search_flooder_included,
     bool enable_correlations) {
-    auto user_graph = pm::detector_error_model_to_user_graph(detector_error_model, enable_correlations);
+    auto user_graph =
+        pm::detector_error_model_to_user_graph(detector_error_model, enable_correlations, num_distinct_weights);
     return user_graph.to_mwpm(num_distinct_weights, ensure_search_flooder_included);
 }
 
@@ -164,22 +165,51 @@ void shatter_blossoms_for_all_detection_events_and_extract_match_edges(
 }
 
 pm::MatchingResult pm::decode_detection_events_for_up_to_64_observables(
-    pm::Mwpm& mwpm, const std::vector<uint64_t>& detection_events) {
+    pm::Mwpm& mwpm, const std::vector<uint64_t>& detection_events, bool edge_correlations) {
+    if (edge_correlations) {
+        // Edge correlations might also be called 2-pass matching. This is the slowest and
+        // highest-accuracy reweighting rule for correlated decoding in which we first decode to edges,
+        // and then reweight the associated edges conditioned on the assumption that an error occurred
+        // at that edge.
+        std::vector<int64_t> edges;
+        decode_detection_events_to_edges(mwpm, detection_events, edges);
+        mwpm.flooder.graph.reweight_for_edges(edges);
+        mwpm.search_flooder.graph.reweight_for_edges(edges);
+    }
+
     process_timeline_until_completion(mwpm, detection_events);
     auto res = shatter_blossoms_for_all_detection_events_and_extract_obs_mask_and_weight(mwpm, detection_events);
     if (!mwpm.flooder.negative_weight_detection_events.empty())
         res += shatter_blossoms_for_all_detection_events_and_extract_obs_mask_and_weight(
             mwpm, mwpm.flooder.negative_weight_detection_events);
     res.obs_mask ^= mwpm.flooder.negative_weight_obs_mask;
     res.weight += mwpm.flooder.negative_weight_sum;
+
+    if (edge_correlations) {
+        mwpm.flooder.graph.undo_reweights();
+        mwpm.search_flooder.graph.undo_reweights();
+    }
+
     return res;
 }
 
 void pm::decode_detection_events(
     pm::Mwpm& mwpm,
     const std::vector<uint64_t>& detection_events,
     uint8_t* obs_begin_ptr,
-    pm::total_weight_int& weight) {
+    pm::total_weight_int& weight,
+    bool edge_correlations) {
+    if (edge_correlations) {
+        // Edge correlations might also be called 2-pass matching. This is the slowest and
+        // highest-accuracy reweighting rule for correlated decoding in which we first decode to edges,
+        // and then reweight the associated edges conditioned on the assumption that an error occurred
+        // at that edge.
+        std::vector<int64_t> edges;
+        decode_detection_events_to_edges(mwpm, detection_events, edges);
+        mwpm.flooder.graph.reweight_for_edges(edges);
+        mwpm.search_flooder.graph.reweight_for_edges(edges);
+    }
+
     size_t num_observables = mwpm.flooder.graph.num_observables;
     process_timeline_until_completion(mwpm, detection_events);
 
@@ -210,6 +240,11 @@ void pm::decode_detection_events(
         // Add negative weight sum to blossom solution weight
         weight = bit_packed_res.weight + mwpm.flooder.negative_weight_sum;
     }
+
+    if (edge_correlations) {
+        mwpm.flooder.graph.undo_reweights();
+        mwpm.search_flooder.graph.undo_reweights();
+    }
 }
 
 void pm::decode_detection_events_to_match_edges(pm::Mwpm& mwpm, const std::vector<uint64_t>& detection_events) {
@@ -289,3 +324,16 @@ void pm::decode_detection_events_to_edges(
         }
     }
 }
+
+void pm::decode_detection_events_to_edges_with_edge_correlations(
+    pm::Mwpm& mwpm, const std::vector<uint64_t>& detection_events, std::vector<int64_t>& edges) {
+    decode_detection_events_to_edges(mwpm, detection_events, edges);
+    mwpm.flooder.graph.reweight_for_edges(edges);
+    mwpm.search_flooder.graph.reweight_for_edges(edges);
+    edges.clear();
+
+    decode_detection_events_to_edges(mwpm, detection_events, edges);
+
+    mwpm.flooder.graph.undo_reweights();
+    mwpm.search_flooder.graph.undo_reweights();
+}

src/pymatching/sparse_blossom/driver/mwpm_decoding.h

@@ -53,7 +53,7 @@ Mwpm detector_error_model_to_mwpm(
     bool enable_correlations = false);
 
 MatchingResult decode_detection_events_for_up_to_64_observables(
-    pm::Mwpm& mwpm, const std::vector<uint64_t>& detection_events);
+    pm::Mwpm& mwpm, const std::vector<uint64_t>& detection_events, bool edge_correlations);
 
 /// Used to decode detection events for an existing Mwpm object `mwpm', and a vector of
 /// detection event indices `detection_events'. The predicted observables are XOR-ed into an
@@ -64,7 +64,8 @@ void decode_detection_events(
     pm::Mwpm& mwpm,
     const std::vector<uint64_t>& detection_events,
     uint8_t* obs_begin_ptr,
-    pm::total_weight_int& weight);
+    pm::total_weight_int& weight,
+    bool edge_correlations);
 
 /// Decode detection events using a Mwpm object and vector of detection event indices
 /// Returns the compressed edges in the matching: the pairs of detection events that are
@@ -78,6 +79,9 @@ void decode_detection_events_to_match_edges(pm::Mwpm& mwpm, const std::vector<ui
 void decode_detection_events_to_edges(
     pm::Mwpm& mwpm, const std::vector<uint64_t>& detection_events, std::vector<int64_t>& edges);
 
+void decode_detection_events_to_edges_with_edge_correlations(
+    pm::Mwpm& mwpm, const std::vector<uint64_t>& detection_events, std::vector<int64_t>& edges);
+
 }  // namespace pm
 
 #endif  // PYMATCHING2_MWPM_DECODING_H