From 344c9216f1c9b17c7174bb31d8a540fb0167bc39 Mon Sep 17 00:00:00 2001
From: peter-janderks <peter-janderks@hotmail.com>
Date: Wed, 20 Dec 2023 11:18:39 +0100
Subject: [PATCH] implement example sinter bposd decoder

---
 README.md                                     |  2 +-
 examples/sinter_example.py                    | 67 +++++++++++++++++++
 pyproject.toml                                |  2 +
 .../sinter_decoders/sinter_bposd_decoder.py   | 14 ++--
 4 files changed, 76 insertions(+), 9 deletions(-)
 create mode 100644 examples/sinter_example.py

diff --git a/README.md b/README.md
index d85bd03..814b1e3 100644
--- a/README.md
+++ b/README.md
@@ -12,6 +12,7 @@ A C++ rewrite of the `LDPC` package for decoding low density parity check checks
 - Belief-find with inversion solving (first implementation I am aware of).
 - An implementation of the Kuo and Lai memory belief propagation decoder (https://arxiv.org/abs/2104.13659)
 - Flip and P-flip decoders (https://aps.arxiv.org/abs/2212.06985)
+- STIM integration for circuit level noise.
 
 ## ToDos
 
@@ -22,7 +23,6 @@ A C++ rewrite of the `LDPC` package for decoding low density parity check checks
 - Functions need to be properly documented (in progress)
 - Proper test coverage is required (C++ has 100%, Python tests still need to expanded).
 - The Peeling version of union-find only works for the Toric code. A routine for matching to the boundary needs to be implemented.
-- STIM integration for circuit level noise.
 
 ## Python - Installation from source
 
diff --git a/examples/sinter_example.py b/examples/sinter_example.py
new file mode 100644
index 0000000..dc33bf8
--- /dev/null
+++ b/examples/sinter_example.py
@@ -0,0 +1,67 @@
+import sinter
+from ldpc.sinter_decoders import SinterBpOsdDecoder
+import stim 
+from matplotlib import pyplot as plt
+import numpy as np
+
+def generate_example_tasks():
+    for p in np.arange(0.001,0.01,0.001):
+        for d in [5,7,9]:
+            sc_circuit = stim.Circuit.generated(
+                    rounds=d,
+                    distance=d,
+                    after_clifford_depolarization=p,
+                    after_reset_flip_probability=p,
+                    before_measure_flip_probability=p,
+                    before_round_data_depolarization=p,
+                    code_task=f'surface_code:rotated_memory_z',   
+                )
+            yield sinter.Task(
+                circuit=sc_circuit,
+                json_metadata={
+                    'p': p,
+                    'd': d,
+                    'rounds':d,
+                },
+            )
+
+def main():
+    samples = sinter.collect(
+        num_workers=10,
+        max_shots=10_000,
+        max_errors=200,
+        tasks=generate_example_tasks(),
+        decoders=['bposd' ],
+        custom_decoders={'bposd': SinterBpOsdDecoder()},
+        print_progress=True,
+        save_resume_filepath=f'bposd_surface_code.csv',
+        )
+    
+
+    # Print samples as CSV data.
+    print(sinter.CSV_HEADER)
+    for sample in samples:
+        print(sample.to_csv_line())
+
+    # Render a matplotlib plot of the data.
+    fig, ax = plt.subplots(1, 1)
+    sinter.plot_error_rate(
+        ax=ax,
+        stats=samples,
+        group_func=lambda stat: f"Rotated Surface Code d={stat.json_metadata['d']} dec={stat.decoder}",
+        x_func=lambda stat: stat.json_metadata['p'],
+    )
+    ax.loglog()
+    ax.grid()
+    ax.set_title('Logical Error Rate vs Physical Error Rate')
+    ax.set_ylabel('Logical Error Probability (per shot)')
+    ax.set_xlabel('Physical Error Rate')
+    ax.legend()
+
+    # Save to file and also open in a window.
+    fig.savefig('plot.png')
+    plt.show()
+
+
+if __name__ == '__main__':
+    main()
diff --git a/pyproject.toml b/pyproject.toml
index eb78636..cb64dcd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,6 +21,8 @@ dependencies = [
     "scipy>=1.9.3",
     "tqdm",
     "pytest"
+    "stim"
+    "sinter"
 ]
 version = "2.0.12"
 
diff --git a/src_python/ldpc/sinter_decoders/sinter_bposd_decoder.py b/src_python/ldpc/sinter_decoders/sinter_bposd_decoder.py
index e429a79..13dbefa 100644
--- a/src_python/ldpc/sinter_decoders/sinter_bposd_decoder.py
+++ b/src_python/ldpc/sinter_decoders/sinter_bposd_decoder.py
@@ -1,12 +1,11 @@
 import stim
 import numpy as np
 import pathlib
-from BeliefMatching import detector_error_model_to_check_matrices
-from ldpc.sinter_decoders import SinterDecoder
+from beliefmatching import detector_error_model_to_check_matrices
 from ldpc.bposd_decoder import BpOsdDecoder
+from sinter import Decoder
 
-
-class SinterBpOsdDecoder(SinterDecoder):
+class SinterBpOsdDecoder(Decoder):
     def __init__(self, 
                  bp_method="ms", 
                  ms_scaling_factor=0.625, 
@@ -30,7 +29,7 @@ def decode_via_files(
     ) -> None:
         self.dem = stim.DetectorErrorModel.from_file(dem_path)
         self.matrices = detector_error_model_to_check_matrices(self.dem)
-        self.bposd = BpOsdDecoder(self.matrices.check_matrix, error_channel=self.matrices.priors, max_iter=5, bp_method="ms", ms_scaling_factor = 0.625, schedule="parallel", osd_method = "osd0")
+        self.bposd = BpOsdDecoder(self.matrices.check_matrix, error_channel=list(self.matrices.priors), max_iter=5, bp_method="ms", ms_scaling_factor = 0.625, schedule="parallel", osd_method = "osd0")
 
         shots = stim.read_shot_data_file(
             path=dets_b8_in_path, format="b8", num_detectors=self.dem.num_detectors
@@ -48,6 +47,5 @@ def decode_via_files(
 
 
     def decode(self, syndrome: np.ndarray) -> np.ndarray:
-        corr = self.matrices.decode(syndrome)
-        if self._bpd.converge:
-            return (self.matrices.observables_matrix @ corr) % 2
+        corr = self.bposd.decode(syndrome)
+        return (self.matrices.observables_matrix @ corr) % 2