Add support for ignored nodes and returning a weight (#6)

- 4th detector coordinate set to -1 now means 'ignore this detector' - Improve some of the decomposition logic - Add `chromobius.CompiledDecoder.predict_weighted_obs_flips_from_dets_bit_packed` - Update some of the `gen` utilities for making test circuits - Fix bit rotted CI configuration
quantumlib · Sep 24, 2024 · 16a3593 · 16a3593
1 parent 3b4220b
commit 16a3593
Show file tree

Hide file tree

Showing 171 changed files with 5,790 additions and 9,075 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -53,25 +53,41 @@ jobs:
       MACOSX_DEPLOYMENT_TARGET: "10.15"
       CIBW_BUILD: "${{ matrix.os_dist.dist }}"
       CIBW_ARCHS_MACOS: "${{ matrix.os_dist.macosarch }}"
-      CIBW_TEST_REQUIRES: pytest stim sinter pygltflib
+      CIBW_TEST_REQUIRES: pytest stim~=1.14 sinter pygltflib
       CIBW_TEST_COMMAND: pytest {project}/src
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v3
       - run: python tools/overwrite_dev_versions_with_date.py
-      - run: mkdir -p output/stim
-      - run: mkdir -p output/stimcirq
-      - run: mkdir -p output/sinter
       - run: python -m pip install pybind11~=2.11.1 cibuildwheel~=2.16.2 setuptools
       - run: python -m cibuildwheel --print-build-identifiers
       - run: python -m cibuildwheel --output-dir output/chromobius
+      - uses: actions/upload-artifact@v4.4.0
+        with:
+          name: "dist-chromobius-${{ matrix.os_dist.os }}-${{ matrix.os_dist.dist }}-${{ matrix.os_dist.macosarch }}"
+          path: dist/*
+  build_sdist:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v3
+      - run: python -m pip install setuptools pybind11~=2.11.1
+      - run: python tools/overwrite_dev_versions_with_date.py
+      - run: mkdir output
       - run: python setup.py sdist
-      - run: mv dist/* output/chromobius
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4.4.0
         with:
-          name: dist
-          path: |
-            ./output/chromobius/*
+          name: "dist-chromobius-sdist"
+          path: dist/*.tar.gz
+  merge_upload_artifacts:
+    needs: ["build_dist", "build_sdist"]
+    runs-on: ubuntu-latest
+    steps:
+      - name: Merge Artifacts
+        uses: actions/upload-artifact/merge@v4
+        with:
+          name: dist-chromobius
+          pattern: dist-chromobius-*
   check_sdist_installs:
     runs-on: ubuntu-latest
     steps:
@@ -90,14 +106,19 @@ jobs:
   build_bazel:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v1
-    - uses: bazelbuild/setup-bazelisk@v1
+    - uses: actions/checkout@v3
+    - uses: bazel-contrib/setup-bazel@0.8.5
+      with:
+        bazelisk-cache: true
+        disk-cache: ${{ github.workflow }}
+        repository-cache: true
+        bazelisk-version: 1.x
     - run: bazel build :all
     - run: bazel test :chromobius_test
   build_clang:
     runs-on: ubuntu-20.04
     steps:
-    - uses: actions/checkout@v1
+    - uses: actions/checkout@v3
     - run: |
         cd ..
         git clone https://github.com/google/googletest.git -b release-1.12.1
@@ -115,14 +136,14 @@ jobs:
   perf:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v1
+    - uses: actions/checkout@v3
     - run: cmake .
     - run: make chromobius_perf -j 2
     - run: out/chromobius_perf
   test:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v1
+    - uses: actions/checkout@v3
     - run: |
         cd ..
         git clone https://github.com/google/googletest.git -b release-1.12.1
@@ -136,7 +157,7 @@ jobs:
   test_o3:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v1
+    - uses: actions/checkout@v3
     - run: |
         cd ..
         git clone https://github.com/google/googletest.git -b release-1.12.1
@@ -152,7 +173,12 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v3
-      - uses: bazelbuild/setup-bazelisk@v1
+      - uses: bazel-contrib/setup-bazel@0.8.5
+        with:
+          bazelisk-cache: true
+          disk-cache: ${{ github.workflow }}
+          repository-cache: true
+          bazelisk-version: 1.x
       - run: bazel build :chromobius_dev_wheel
       - run: pip install bazel-bin/chromobius-0.0.dev0-py3-none-any.whl
       - run: diff <(python tools/gen_chromobius_api_reference.py -dev) doc/chromobius_api_reference.md
@@ -172,10 +198,15 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v3
-      - uses: bazelbuild/setup-bazelisk@v1
+      - uses: bazel-contrib/setup-bazel@0.8.5
+        with:
+          bazelisk-cache: true
+          disk-cache: ${{ github.workflow }}
+          repository-cache: true
+          bazelisk-version: 1.x
       - run: bazel build :chromobius_dev_wheel
       - run: pip install bazel-bin/chromobius-0.0.dev0-py3-none-any.whl
-      - run: pip install pytest stim sinter pygltflib
+      - run: pip install pytest stim~=1.14 sinter pygltflib
       - run: pytest src
       - run: tools/doctest_proper.py --module chromobius
   upload_dev_release_to_pypi:

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -43,7 +43,7 @@ set(SIMD_WIDTH 128)
 include(FetchContent)
 FetchContent_Declare(stim
         GIT_REPOSITORY https://github.com/quantumlib/stim.git
-        GIT_TAG 3e38d12d0a0fb3022646b694137b733a4700d300)
+        GIT_TAG da4594c5ede00a063ec2b84bd830f846b5d097dd)
 FetchContent_GetProperties(stim)
 if(NOT stim_POPULATED)
   FetchContent_Populate(stim)

diff --git a/WORKSPACE b/WORKSPACE
@@ -3,7 +3,7 @@ load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 
 git_repository(
     name = "stim",
-    commit = "3e38d12d0a0fb3022646b694137b733a4700d300",
+    commit = "da4594c5ede00a063ec2b84bd830f846b5d097dd",
     remote = "https://github.com/quantumlib/stim.git",
 )
 

diff --git a/doc/chromobius.pyi b/doc/chromobius.pyi
@@ -48,6 +48,7 @@ class CompiledDecoder:
                         3 = Red Z
                         4 = Green Z
                         5 = Blue Z
+                        -1 = Ignore this Detector
                 2. Rainbow triplets. Bulk errors with three symptoms in one basis should
                     have one symptom of each color. Errors with three symptoms that
                     repeat a color will cause an exception unless they can be decomposed
@@ -161,6 +162,87 @@ class CompiledDecoder:
             >>> mistakes = np.count_nonzero(differences)
             >>> assert mistakes < shots / 5
         """
+    @staticmethod
+    def predict_weighted_obs_flips_from_dets_bit_packed(
+        dets: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Predicts observable flips and weights from detection events.
+
+        The returned weight comes directly from the underlying call to pymatching, not
+        accounting for the lifting process.
+
+        Args:
+            dets: A bit packed numpy array of detection event data. The array can either
+                be 1-dimensional (a single shot to decode) or 2-dimensional (multiple
+                shots to decode, with the first axis being the shot axis and the second
+                axis being the detection event byte axis).
+
+                The array's dtype must be np.uint8. If you have an array of dtype
+                np.bool_, you have data that's not bit packed. You can pack it by
+                using `np.packbits(array, bitorder='little')`. But ideally you
+                should attempt to never have unpacked data in the first place,
+                since it's 8x larger which can be a large performance loss. For
+                example, stim's sampler methods all have a `bit_packed=True` argument
+                that cause them to return bit packed data.
+
+        Returns:
+            A tuple (obs, weights).
+            Obs is a bit packed numpy array of observable flip data.
+            Weights is a numpy array (or scalar) of floats.
+
+            If dets is a 1D array, then the result has:
+                obs.shape = (math.ceil(num_obs / 8),)
+                obs.dtype = np.uint8
+                weights.shape = ()
+                weights.dtype = np.float32
+            If dets is a 2D array, then the result has:
+                shape = (dets.shape[0], math.ceil(num_obs / 8),)
+                dtype = np.uint8
+                weights.shape = (dets.shape[0],)
+                weights.dtype = np.float32
+
+            To determine if the observable with index k was flipped in shot s, compute:
+                `bool((obs[s, k // 8] >> (k % 8)) & 1)`
+
+        Example:
+            >>> import stim
+            >>> import chromobius
+            >>> import numpy as np
+
+            >>> repetition_color_code = stim.Circuit('''
+            ...     # Apply noise.
+            ...     X_ERROR(0.1) 0 1 2 3 4 5 6 7
+            ...     # Measure three-body stabilizers to catch errors.
+            ...     MPP Z0*Z1*Z2 Z1*Z2*Z3 Z2*Z3*Z4 Z3*Z4*Z5 Z4*Z5*Z6 Z5*Z6*Z7
+            ...
+            ...     # Annotate detectors, with a coloring in the 4th coordinate.
+            ...     DETECTOR(0, 0, 0, 2) rec[-6]
+            ...     DETECTOR(1, 0, 0, 0) rec[-5]
+            ...     DETECTOR(2, 0, 0, 1) rec[-4]
+            ...     DETECTOR(3, 0, 0, 2) rec[-3]
+            ...     DETECTOR(4, 0, 0, 0) rec[-2]
+            ...     DETECTOR(5, 0, 0, 1) rec[-1]
+            ...
+            ...     # Check on the message.
+            ...     M 0
+            ...     OBSERVABLE_INCLUDE(0) rec[-1]
+            ... ''')
+
+            >>> # Sample the circuit.
+            >>> shots = 4096
+            >>> sampler = repetition_color_code.compile_detector_sampler()
+            >>> dets, actual_obs_flips = sampler.sample(
+            ...     shots=shots,
+            ...     separate_observables=True,
+            ...     bit_packed=True,
+            ... )
+
+            >>> # Decode with Chromobius.
+            >>> dem = repetition_color_code.detector_error_model()
+            >>> decoder = chromobius.compile_decoder_for_dem(dem)
+            >>> result = decoder.predict_weighted_obs_flips_from_dets_bit_packed(dets)
+            >>> pred, weights = result
+        """
 def compile_decoder_for_dem(
     dem: stim.DetectorErrorModel,
 ) -> chromobius.CompiledDecoder:
@@ -178,6 +260,7 @@ def compile_decoder_for_dem(
                     3 = Red Z
                     4 = Green Z
                     5 = Blue Z
+                    -1 = Ignore this Detector
             2. Rainbow triplets. Bulk errors with three symptoms in one basis should
                 have one symptom of each color. Errors with three symptoms that
                 repeat a color will cause an exception unless they can be decomposed

diff --git a/doc/chromobius_api_reference.md b/doc/chromobius_api_reference.md
@@ -7,6 +7,7 @@
 - [`chromobius.CompiledDecoder`](#chromobius.CompiledDecoder)
     - [`chromobius.CompiledDecoder.from_dem`](#chromobius.CompiledDecoder.from_dem)
     - [`chromobius.CompiledDecoder.predict_obs_flips_from_dets_bit_packed`](#chromobius.CompiledDecoder.predict_obs_flips_from_dets_bit_packed)
+    - [`chromobius.CompiledDecoder.predict_weighted_obs_flips_from_dets_bit_packed`](#chromobius.CompiledDecoder.predict_weighted_obs_flips_from_dets_bit_packed)
 ```python
 # Types used by the method definitions.
 from typing import overload, TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union
@@ -37,6 +38,7 @@ def compile_decoder_for_dem(
                     3 = Red Z
                     4 = Green Z
                     5 = Blue Z
+                    -1 = Ignore this Detector
             2. Rainbow triplets. Bulk errors with three symptoms in one basis should
                 have one symptom of each color. Errors with three symptoms that
                 repeat a color will cause an exception unless they can be decomposed
@@ -143,6 +145,7 @@ def from_dem(
                     3 = Red Z
                     4 = Green Z
                     5 = Blue Z
+                    -1 = Ignore this Detector
             2. Rainbow triplets. Bulk errors with three symptoms in one basis should
                 have one symptom of each color. Errors with three symptoms that
                 repeat a color will cause an exception unless they can be decomposed
@@ -264,3 +267,91 @@ def predict_obs_flips_from_dets_bit_packed(
         >>> assert mistakes < shots / 5
     """
 ```
+
+<a name="chromobius.CompiledDecoder.predict_weighted_obs_flips_from_dets_bit_packed"></a>
+```python
+# chromobius.CompiledDecoder.predict_weighted_obs_flips_from_dets_bit_packed
+
+# (in class chromobius.CompiledDecoder)
+@staticmethod
+def predict_weighted_obs_flips_from_dets_bit_packed(
+    dets: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Predicts observable flips and weights from detection events.
+
+    The returned weight comes directly from the underlying call to pymatching, not
+    accounting for the lifting process.
+
+    Args:
+        dets: A bit packed numpy array of detection event data. The array can either
+            be 1-dimensional (a single shot to decode) or 2-dimensional (multiple
+            shots to decode, with the first axis being the shot axis and the second
+            axis being the detection event byte axis).
+
+            The array's dtype must be np.uint8. If you have an array of dtype
+            np.bool_, you have data that's not bit packed. You can pack it by
+            using `np.packbits(array, bitorder='little')`. But ideally you
+            should attempt to never have unpacked data in the first place,
+            since it's 8x larger which can be a large performance loss. For
+            example, stim's sampler methods all have a `bit_packed=True` argument
+            that cause them to return bit packed data.
+
+    Returns:
+        A tuple (obs, weights).
+        Obs is a bit packed numpy array of observable flip data.
+        Weights is a numpy array (or scalar) of floats.
+
+        If dets is a 1D array, then the result has:
+            obs.shape = (math.ceil(num_obs / 8),)
+            obs.dtype = np.uint8
+            weights.shape = ()
+            weights.dtype = np.float32
+        If dets is a 2D array, then the result has:
+            shape = (dets.shape[0], math.ceil(num_obs / 8),)
+            dtype = np.uint8
+            weights.shape = (dets.shape[0],)
+            weights.dtype = np.float32
+
+        To determine if the observable with index k was flipped in shot s, compute:
+            `bool((obs[s, k // 8] >> (k % 8)) & 1)`
+
+    Example:
+        >>> import stim
+        >>> import chromobius
+        >>> import numpy as np
+
+        >>> repetition_color_code = stim.Circuit('''
+        ...     # Apply noise.
+        ...     X_ERROR(0.1) 0 1 2 3 4 5 6 7
+        ...     # Measure three-body stabilizers to catch errors.
+        ...     MPP Z0*Z1*Z2 Z1*Z2*Z3 Z2*Z3*Z4 Z3*Z4*Z5 Z4*Z5*Z6 Z5*Z6*Z7
+        ...
+        ...     # Annotate detectors, with a coloring in the 4th coordinate.
+        ...     DETECTOR(0, 0, 0, 2) rec[-6]
+        ...     DETECTOR(1, 0, 0, 0) rec[-5]
+        ...     DETECTOR(2, 0, 0, 1) rec[-4]
+        ...     DETECTOR(3, 0, 0, 2) rec[-3]
+        ...     DETECTOR(4, 0, 0, 0) rec[-2]
+        ...     DETECTOR(5, 0, 0, 1) rec[-1]
+        ...
+        ...     # Check on the message.
+        ...     M 0
+        ...     OBSERVABLE_INCLUDE(0) rec[-1]
+        ... ''')
+
+        >>> # Sample the circuit.
+        >>> shots = 4096
+        >>> sampler = repetition_color_code.compile_detector_sampler()
+        >>> dets, actual_obs_flips = sampler.sample(
+        ...     shots=shots,
+        ...     separate_observables=True,
+        ...     bit_packed=True,
+        ... )
+
+        >>> # Decode with Chromobius.
+        >>> dem = repetition_color_code.detector_error_model()
+        >>> decoder = chromobius.compile_decoder_for_dem(dem)
+        >>> result = decoder.predict_weighted_obs_flips_from_dets_bit_packed(dets)
+        >>> pred, weights = result
+    """
+```
diff --git a/doc/getting_started.ipynb b/doc/getting_started.ipynb
@@ -346,6 +346,7 @@
     "- 3: red Z\n",
     "- 4: green Z\n",
     "- 5: blue Z.\n",
+    "- (-1): ignore this detector (e.g. it's from something beyond the color code).\n",
     "\n",
     "As you can maybe tell, making a circuit can be a very involved process!\n",
     "This notebook isn't really about making circuits, and you probably don't want to get too bogged down in this detail, so let's not get stuck on this.\n",

diff --git a/file_lists/source_files_no_main b/file_lists/source_files_no_main
@@ -14,6 +14,9 @@ src/chromobius/datatypes/color_basis.h
 src/chromobius/datatypes/conf.h
 src/chromobius/datatypes/rgb_edge.cc
 src/chromobius/datatypes/rgb_edge.h
+src/chromobius/datatypes/stim_integration.cc
+src/chromobius/datatypes/stim_integration.h
+src/chromobius/datatypes/xor_vec.h
 src/chromobius/decode/decoder.cc
 src/chromobius/decode/decoder.h
 src/chromobius/decode/matcher_interface.h

diff --git a/file_lists/test_files b/file_lists/test_files
@@ -6,6 +6,8 @@ src/chromobius/commands/main_predict.test.cc
 src/chromobius/datatypes/atomic_error.test.cc
 src/chromobius/datatypes/color_basis.test.cc
 src/chromobius/datatypes/rgb_edge.test.cc
+src/chromobius/datatypes/stim_integration.test.cc
+src/chromobius/datatypes/xor_vec.test.cc
 src/chromobius/decode/decoder.test.cc
 src/chromobius/decode/decoder_integration.test.cc
 src/chromobius/graph/charge_graph.test.cc