Add reweighting strength option to allow regularized correlated matching

src/pymatching/matching.py

@@ -220,6 +220,7 @@ def decode(self,
                *,
                return_weight: bool = False,
                enable_correlations: bool = False,
+               reweighting_strength: float = 1.0,
                **kwargs
                ) -> Union[np.ndarray, Tuple[np.ndarray, int]]:
         r"""
@@ -254,6 +255,12 @@ def decode(self,
             `stim.DetectorErrorModel` with `enable_correlations=True`. For a description
             of the correlated matching algorithm, see https://arxiv.org/abs/1310.0863.
             By default, False
+        reweighting_strength : float, optional
+            Only used when `enable_correlations==True`. Multiplies the conditional probabilities that
+            correlated matching uses to reweight edges after the first pass. The default, 1.0, is
+            standard correlated matching. Values below 1.0 weaken the reweighting, which has been shown
+            to be beneficial in the presence of circuit-level noise at low physical error rates, see
+            https://arxiv.org/abs/2610.26785 for more details.
 
         Returns
         -------
@@ -338,7 +345,7 @@ def decode(self,
         """
         detection_events = self._syndrome_array_to_detection_events(z)
         correction, weight = self._matching_graph.decode(
-            detection_events, enable_correlations=enable_correlations
+            detection_events, enable_correlations=enable_correlations, reweighting_strength=reweighting_strength
         )
         if return_weight:
             return correction, weight
@@ -352,7 +359,8 @@ def decode_batch(
             return_weights: bool = False,
             bit_packed_shots: bool = False,
             bit_packed_predictions: bool = False,
-            enable_correlations: bool = False) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+            enable_correlations: bool = False,
+            reweighting_strength: float = 1.0) -> 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.
@@ -396,6 +404,8 @@ def decode_batch(
             `stim.DetectorErrorModel` with `enable_correlations=True`. For a description
             of the correlated matching algorithm, see https://arxiv.org/abs/1310.0863.
             By default, False
+        reweighting_strength : float, optional
+            Only used if `enable_correlations==True`. See `Matching.decode`. By default, 1.0
 
         Returns
         -------
@@ -452,7 +462,8 @@ def decode_batch(
             shots,
             bit_packed_predictions=bit_packed_predictions,
             bit_packed_shots=bit_packed_shots,
-            enable_correlations=enable_correlations
+            enable_correlations=enable_correlations,
+            reweighting_strength=reweighting_strength
         )
         if return_weights:
             return predictions, weights
@@ -462,7 +473,8 @@ def decode_batch(
     def decode_to_edges_array(self,
                               syndrome: Union[np.ndarray, List[bool], List[int]],
                               *,
-                              enable_correlations: bool = False
+                              enable_correlations: bool = False,
+                              reweighting_strength: float = 1.0
                               ) -> np.ndarray:
         """
         Decode the syndrome `syndrome` using minimum-weight perfect matching, returning the edges in the
@@ -489,6 +501,8 @@ def decode_to_edges_array(self,
             `stim.DetectorErrorModel` with `enable_correlations=True`. For a description
             of the correlated matching algorithm, see https://arxiv.org/abs/1310.0863.
             By default, False
+        reweighting_strength : float, optional
+            Only used if `enable_correlations==True`. See `Matching.decode`. By default, 1.0
 
         Returns
         -------
@@ -520,7 +534,7 @@ def decode_to_edges_array(self,
         """
         detection_events = self._syndrome_array_to_detection_events(syndrome)
         return self._matching_graph.decode_to_edges_array(
-            detection_events, enable_correlations=enable_correlations
+            detection_events, enable_correlations=enable_correlations, reweighting_strength=reweighting_strength
         )
 
     def decode_to_matched_dets_array(self,

src/pymatching/sparse_blossom/driver/implied_weights.h

@@ -23,6 +23,7 @@ struct ImpliedWeightUnconverted {
     size_t node1;
     size_t node2;
     double implied_weight;
+    double implied_probability = 0.0;
     bool operator==(const ImpliedWeightUnconverted& other) const {
         return (this->node1 == other.node1) && (this->node2 == other.node2) &&
                (this->implied_weight == other.implied_weight);

src/pymatching/sparse_blossom/driver/user_graph.cc

@@ -537,7 +537,11 @@ void pm::UserGraph::populate_implied_edge_weights(
                     double implied_probability_for_other_edge =
                         std::min(0.5, affected_edge_and_probability.second / marginal_probability);
                     double w = pm::to_weight_for_correlations(implied_probability_for_other_edge);
-                    ImpliedWeightUnconverted implied{affected_edge.first, affected_edge.second, w};
+                    ImpliedWeightUnconverted implied{
+                        affected_edge.first,
+                        affected_edge.second,
+                        w,
+                        affected_edge_and_probability.second / marginal_probability};
                     edge.implied_weights_for_other_edges.push_back(implied);
                 }
             }

src/pymatching/sparse_blossom/driver/user_graph.pybind.cc

@@ -84,6 +84,11 @@ pm::MERGE_STRATEGY merge_strategy_from_string(const std::string &merge_strategy)
     }
 }
 
+void set_reweighting_strength(pm::Mwpm &mwpm, double reweighting_strength) {
+    mwpm.flooder.graph.set_reweighting_strength(reweighting_strength);
+    mwpm.search_flooder.graph.set_reweighting_strength(reweighting_strength, mwpm.flooder.graph.normalising_constant);
+}
+
 void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_<pm::UserGraph> &g) {
     g.def(py::init<>());
     g.def(py::init<size_t>(), "num_nodes"_a);
@@ -178,10 +183,15 @@ void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_<pm::UserGrap
     });
     g.def(
         "decode",
-        [](pm::UserGraph &self, const py::array_t<uint64_t> &detection_events, bool enable_correlations) {
+        [](pm::UserGraph &self,
+           const py::array_t<uint64_t> &detection_events,
+           bool enable_correlations,
+           double reweighting_strength) {
             std::vector<uint64_t> detection_events_vec(
                 detection_events.data(), detection_events.data() + detection_events.size());
             auto &mwpm = enable_correlations ? self.get_mwpm_with_search_graph() : self.get_mwpm();
+            if (enable_correlations)
+                set_reweighting_strength(mwpm, reweighting_strength);
             auto obs_crossed = new std::vector<uint8_t>(self.get_num_observables(), 0);
             pm::total_weight_int weight = 0;
             pm::decode_detection_events(mwpm, detection_events_vec, obs_crossed->data(), weight, enable_correlations);
@@ -196,11 +206,17 @@ void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_<pm::UserGrap
             return res;
         },
         "detection_events"_a,
-        "enable_correlations"_a = false);
+        "enable_correlations"_a = false,
+        "reweighting_strength"_a = 1.0);
     g.def(
         "decode_to_edges_array",
-        [](pm::UserGraph &self, const py::array_t<uint64_t> &detection_events, bool enable_correlations) {
+        [](pm::UserGraph &self,
+           const py::array_t<uint64_t> &detection_events,
+           bool enable_correlations,
+           double reweighting_strength) {
             auto &mwpm = self.get_mwpm_with_search_graph();
+            if (enable_correlations)
+                set_reweighting_strength(mwpm, reweighting_strength);
             std::vector<uint64_t> detection_events_vec(
                 detection_events.data(), detection_events.data() + detection_events.size());
             auto edges = new std::vector<int64_t>();
@@ -216,7 +232,8 @@ void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_<pm::UserGrap
             return edges_arr;
         },
         "detection_events"_a,
-        "enable_correlations"_a = false);
+        "enable_correlations"_a = false,
+        "reweighting_strength"_a = 1.0);
     g.def(
         "decode_to_matched_detection_events_array",
         [](pm::UserGraph &self, const py::array_t<uint64_t> &detection_events) {
@@ -250,7 +267,8 @@ void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_<pm::UserGrap
            const py::array_t<uint8_t> &shots,
            bool bit_packed_shots,
            bool bit_packed_predictions,
-           bool enable_correlations) {
+           bool enable_correlations,
+           double reweighting_strength) {
             if (shots.ndim() != 2)
                 throw std::invalid_argument(
                     "`shots` array should have two dimensions, not " + std::to_string(shots.ndim()));
@@ -286,6 +304,8 @@ void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_<pm::UserGrap
             auto ws = weights.mutable_unchecked<1>();
 
             auto &mwpm = enable_correlations ? self.get_mwpm_with_search_graph() : self.get_mwpm();
+            if (enable_correlations)
+                set_reweighting_strength(mwpm, reweighting_strength);
             std::vector<uint64_t> detection_events;
 
             // Vector used to extract predicted observables when decoding if bit_packed_predictions is true
@@ -337,7 +357,8 @@ void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_<pm::UserGrap
         "shots"_a,
         "bit_packed_shots"_a = false,
         "bit_packed_predictions"_a = false,
-        "enable_correlations"_a = false);
+        "enable_correlations"_a = false,
+        "reweighting_strength"_a = 1.0);
     g.def(
         "decode_to_matched_detection_events_dict",
         [](pm::UserGraph &self, const py::array_t<uint64_t> &detection_events) {

src/pymatching/sparse_blossom/flooder/graph.cc

@@ -14,6 +14,7 @@
 
 #include "pymatching/sparse_blossom/flooder/graph.h"
 
+#include <algorithm>
 #include <cassert>
 #include <cmath>
 #include <map>
@@ -188,14 +189,28 @@ void MatchingGraph::convert_implied_weights(double normalising_constant) {
         const std::vector<std::vector<ImpliedWeightUnconverted>>& rules_for_node =
             edges_to_implied_weights_unconverted[u];
         for (size_t v = 0; v < nodes[u].neighbors.size(); v++) {
-            for (const auto& rule : rules_for_node[v]) {
+            nodes[u].neighbor_implied_weights[v].clear();
+            for (auto rule : rules_for_node[v]) {
+                if (reweighting_strength != 1.0) {
+                    double p = std::min(0.5, reweighting_strength * rule.implied_probability);
+                    rule.implied_weight = std::log((1 - p) / p);
+                }
                 ImpliedWeight converted = convert_rule(nodes, rule, normalising_constant);
                 nodes[u].neighbor_implied_weights[v].push_back(converted);
             }
         }
     }
 }
 
+void MatchingGraph::set_reweighting_strength(double strength) {
+    if (strength <= 0)
+        throw std::invalid_argument("reweighting_strength must be positive");
+    if (strength != reweighting_strength) {
+        reweighting_strength = strength;
+        convert_implied_weights(normalising_constant);
+    }
+}
+
 // Reweight assuming an error has occurred on a single edge u, v. When v == -1, assumes an edge from
 // u to the boundary.
 void MatchingGraph::reweight_for_edge(const int64_t& u, const int64_t& v) {

src/pymatching/sparse_blossom/flooder/graph.h

@@ -62,6 +62,7 @@ class MatchingGraph {
     // alert a user if they try to decode with enable_correlations=true, but forgot to load from the
     // dem with enable_correlations=true.
     bool loaded_from_dem_without_correlations = false;
+    double reweighting_strength = 1.0;
 
     MatchingGraph();
     MatchingGraph(size_t num_nodes, size_t num_observables);
@@ -81,6 +82,7 @@ class MatchingGraph {
     void update_negative_weight_observables(const std::vector<size_t>& observables);
     void update_negative_weight_detection_events(size_t node_id);
     void convert_implied_weights(double normalising_constant);
+    void set_reweighting_strength(double strength);
 
     void undo_reweights();
     void reweight(std::vector<ImpliedWeight>& implied_weights);

src/pymatching/sparse_blossom/search/search_graph.cc

@@ -14,6 +14,7 @@
 
 #include "search_graph.h"
 
+#include <algorithm>
 #include <cassert>
 #include <cmath>
 
@@ -27,7 +28,8 @@ pm::SearchGraph::SearchGraph(size_t num_nodes) : num_nodes(num_nodes) {
 pm::SearchGraph::SearchGraph(pm::SearchGraph&& graph) noexcept
     : nodes(std::move(graph.nodes)),
       num_nodes(graph.num_nodes),
-      negative_weight_edges(std::move(graph.negative_weight_edges)) {
+      negative_weight_edges(std::move(graph.negative_weight_edges)),
+      edges_to_implied_weights_unconverted(std::move(graph.edges_to_implied_weights_unconverted)) {
 }
 
 void pm::SearchGraph::add_edge(
@@ -152,10 +154,22 @@ void pm::SearchGraph::convert_implied_weights(const double normalising_constant)
         const std::vector<std::vector<ImpliedWeightUnconverted>>& rules_for_node =
             edges_to_implied_weights_unconverted[u];
         for (size_t v = 0; v < nodes[u].neighbors.size(); v++) {
-            for (const auto& rule : rules_for_node[v]) {
+            nodes[u].neighbor_implied_weights[v].clear();
+            for (auto rule : rules_for_node[v]) {
+                if (reweighting_strength != 1.0) {
+                    double p = std::min(0.5, reweighting_strength * rule.implied_probability);
+                    rule.implied_weight = std::log((1 - p) / p);
+                }
                 ImpliedWeight converted = convert_rule(nodes, rule, normalising_constant);
                 nodes[u].neighbor_implied_weights[v].push_back(converted);
             }
         }
     }
 }
+
+void pm::SearchGraph::set_reweighting_strength(double strength, double normalising_constant) {
+    if (strength != reweighting_strength) {
+        reweighting_strength = strength;
+        convert_implied_weights(normalising_constant);
+    }
+}

src/pymatching/sparse_blossom/search/search_graph.h

@@ -37,6 +37,7 @@ class SearchGraph {
     // Used to restore weights after reweighting.
     std::vector<PreviousWeight> previous_weights;
     std::map<size_t, std::vector<std::vector<ImpliedWeightUnconverted>>> edges_to_implied_weights_unconverted;
+    double reweighting_strength = 1.0;
 
     SearchGraph();
     explicit SearchGraph(size_t num_nodes);
@@ -53,6 +54,7 @@ class SearchGraph {
         const std::vector<size_t>& observables,
         const std::vector<ImpliedWeightUnconverted>& implied_weights = {});
     void convert_implied_weights(const double normalizing_constant);
+    void set_reweighting_strength(double strength, double normalising_constant);
     void reweight(std::vector<ImpliedWeight>& implied_weights);
     void reweight_for_edge(const int64_t& u, const int64_t& v);
     void reweight_for_edges(const std::vector<int64_t>& edges);

tests/matching/reweighting_strength_test.py

@@ -0,0 +1,63 @@
+# 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.
+
+import numpy as np
+import pytest
+import stim
+
+import pymatching
+
+# Edges (0, 1) and (2, 3) each have marginal probability 0.2, and P((2, 3) | (0, 1)) = 0.08 / 0.2 = 0.4.
+DEM = stim.DetectorErrorModel("""
+    error(0.14285714285714285) D0 D1
+    error(0.14285714285714285) D2 D3
+    error(0.08) D0 D1 ^ D2 D3
+""")
+SYNDROME = np.array([1, 1, 1, 1], dtype=np.uint8)
+
+
+def test_default_reweighting_strength_is_one():
+    m = pymatching.Matching.from_detector_error_model(DEM, enable_correlations=True)
+    shots = np.array([[0, 0, 0, 0], [1, 1, 0, 0], [0, 0, 1, 1], [1, 1, 1, 1]], dtype=np.uint8)
+    default = m.decode_batch(shots, return_weights=True, enable_correlations=True)
+    explicit = m.decode_batch(shots, return_weights=True, enable_correlations=True, reweighting_strength=1.0)
+    np.testing.assert_array_equal(default[0], explicit[0])
+    np.testing.assert_array_equal(default[1], explicit[1])
+    np.testing.assert_array_equal(
+        m.decode_to_edges_array(SYNDROME, enable_correlations=True),
+        m.decode_to_edges_array(SYNDROME, enable_correlations=True, reweighting_strength=1.0),
+    )
+
+
+@pytest.mark.parametrize("strength", [0.6, 1.0, 1.5])
+def test_reweighting_strength_scales_conditional_probability(strength):
+    m = pymatching.Matching.from_detector_error_model(DEM, enable_correlations=True)
+    _, weight = m.decode(SYNDROME, return_weight=True, enable_correlations=True, reweighting_strength=strength)
+    p = min(0.5, strength * 0.4)
+    assert weight == pytest.approx(2 * np.log((1 - p) / p), abs=1e-3)
+
+
+def test_changing_reweighting_strength_between_calls():
+    m = pymatching.Matching.from_detector_error_model(DEM, enable_correlations=True)
+    weights = [
+        m.decode(SYNDROME, return_weight=True, enable_correlations=True, reweighting_strength=s)[1]
+        for s in (1.0, 0.6, 1.0)
+    ]
+    assert weights[0] == weights[2] != weights[1]
+
+
+def test_reweighting_strength_must_be_positive():
+    m = pymatching.Matching.from_detector_error_model(DEM, enable_correlations=True)
+    with pytest.raises(ValueError):
+        m.decode(SYNDROME, enable_correlations=True, reweighting_strength=0)