From 176534d701cdbb5565fef2ebcdea3744dff4ce72 Mon Sep 17 00:00:00 2001 From: cesar-hernando Date: Wed, 21 Oct 2026 19:23:01 +0200 Subject: [PATCH] Add reweighting_strength to correlated matching --- src/pymatching/matching.py | 24 +++++-- .../sparse_blossom/driver/implied_weights.h | 1 + .../sparse_blossom/driver/user_graph.cc | 6 +- .../driver/user_graph.pybind.cc | 33 ++++++++-- .../sparse_blossom/flooder/graph.cc | 17 ++++- src/pymatching/sparse_blossom/flooder/graph.h | 2 + .../sparse_blossom/search/search_graph.cc | 18 +++++- .../sparse_blossom/search/search_graph.h | 2 + tests/matching/reweighting_strength_test.py | 63 +++++++++++++++++++ 9 files changed, 151 insertions(+), 15 deletions(-) create mode 100644 tests/matching/reweighting_strength_test.py diff --git a/src/pymatching/matching.py b/src/pymatching/matching.py index 3ac8fed0e..d786817f9 100644 --- a/src/pymatching/matching.py +++ b/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, diff --git a/src/pymatching/sparse_blossom/driver/implied_weights.h b/src/pymatching/sparse_blossom/driver/implied_weights.h index f652c6d82..63368353e 100644 --- a/src/pymatching/sparse_blossom/driver/implied_weights.h +++ b/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); diff --git a/src/pymatching/sparse_blossom/driver/user_graph.cc b/src/pymatching/sparse_blossom/driver/user_graph.cc index 313343978..a4a24b4e0 100644 --- a/src/pymatching/sparse_blossom/driver/user_graph.cc +++ b/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); } } diff --git a/src/pymatching/sparse_blossom/driver/user_graph.pybind.cc b/src/pymatching/sparse_blossom/driver/user_graph.pybind.cc index 7b227e5cc..3f21848a2 100644 --- a/src/pymatching/sparse_blossom/driver/user_graph.pybind.cc +++ b/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_ &g) { g.def(py::init<>()); g.def(py::init(), "num_nodes"_a); @@ -178,10 +183,15 @@ void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_ &detection_events, bool enable_correlations) { + [](pm::UserGraph &self, + const py::array_t &detection_events, + bool enable_correlations, + double reweighting_strength) { std::vector 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(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_ &detection_events, bool enable_correlations) { + [](pm::UserGraph &self, + const py::array_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 detection_events_vec( detection_events.data(), detection_events.data() + detection_events.size()); auto edges = new std::vector(); @@ -216,7 +232,8 @@ void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_ &detection_events) { @@ -250,7 +267,8 @@ void pm_pybind::pybind_user_graph_methods(py::module &m, py::class_ &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_(); auto &mwpm = enable_correlations ? self.get_mwpm_with_search_graph() : self.get_mwpm(); + if (enable_correlations) + set_reweighting_strength(mwpm, reweighting_strength); std::vector 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_ &detection_events) { diff --git a/src/pymatching/sparse_blossom/flooder/graph.cc b/src/pymatching/sparse_blossom/flooder/graph.cc index 5c1b8850d..44322a308 100644 --- a/src/pymatching/sparse_blossom/flooder/graph.cc +++ b/src/pymatching/sparse_blossom/flooder/graph.cc @@ -14,6 +14,7 @@ #include "pymatching/sparse_blossom/flooder/graph.h" +#include #include #include #include @@ -188,7 +189,12 @@ void MatchingGraph::convert_implied_weights(double normalising_constant) { const std::vector>& 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); } @@ -196,6 +202,15 @@ void MatchingGraph::convert_implied_weights(double normalising_constant) { } } +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) { diff --git a/src/pymatching/sparse_blossom/flooder/graph.h b/src/pymatching/sparse_blossom/flooder/graph.h index abfddfc25..c2c2c11a6 100644 --- a/src/pymatching/sparse_blossom/flooder/graph.h +++ b/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& 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& implied_weights); diff --git a/src/pymatching/sparse_blossom/search/search_graph.cc b/src/pymatching/sparse_blossom/search/search_graph.cc index d4be3ae5d..a2020d8b8 100644 --- a/src/pymatching/sparse_blossom/search/search_graph.cc +++ b/src/pymatching/sparse_blossom/search/search_graph.cc @@ -14,6 +14,7 @@ #include "search_graph.h" +#include #include #include @@ -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>& 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); + } +} diff --git a/src/pymatching/sparse_blossom/search/search_graph.h b/src/pymatching/sparse_blossom/search/search_graph.h index e57006c66..e30f64cfa 100644 --- a/src/pymatching/sparse_blossom/search/search_graph.h +++ b/src/pymatching/sparse_blossom/search/search_graph.h @@ -37,6 +37,7 @@ class SearchGraph { // Used to restore weights after reweighting. std::vector previous_weights; std::map>> 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& observables, const std::vector& implied_weights = {}); void convert_implied_weights(const double normalizing_constant); + void set_reweighting_strength(double strength, double normalising_constant); void reweight(std::vector& implied_weights); void reweight_for_edge(const int64_t& u, const int64_t& v); void reweight_for_edges(const std::vector& edges); diff --git a/tests/matching/reweighting_strength_test.py b/tests/matching/reweighting_strength_test.py new file mode 100644 index 000000000..a91443f1f --- /dev/null +++ b/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)