diff --git a/BUILD b/BUILD index 0d6a8b6d..a02f45be 100644 --- a/BUILD +++ b/BUILD @@ -26,6 +26,8 @@ py_wheel( ], version = "$(VERSION)", requires=[ + "numpy", + "scipy", "stim", ], python_tag="$(TARGET_VERSION)", diff --git a/docs/tutorial.ipynb b/docs/tutorial.ipynb index 53c10ce5..e00ef674 100644 --- a/docs/tutorial.ipynb +++ b/docs/tutorial.ipynb @@ -906,6 +906,99 @@ "print_results(results)" ] }, + { + "cell_type": "markdown", + "id": "2b9319e5", + "metadata": { + "id": "gari-correlated-decoding" + }, + "source": [ + "# Faster Methods of Decoding Correlated Errors with Tesseract\n", + "\n", + "## GARI\n", + "\n", + "Graph augmentation and rewiring for inference (GARI) transforms a correlated\n", + "CSS detector matrix into a block form for Tesseract; see [Decoding correlated\n", + "errors in quantum LDPC codes](https://doi.org/10.1038/s41467-026-70556-3).\n", + "This example reuses the superdense color-code memory-Z circuit and sampled\n", + "data from above, and applies the XOR prior policy." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd1644de", + "metadata": { + "id": "gari-transform-example" + }, + "outputs": [], + "source": [ + "gari = tesseract_decoder.demutil.gari\n", + "\n", + "gari_dem, gari_layout = gari.circuit_to_gari(\n", + " circuit,\n", + " prior_function=gari.tesseract_xor_prior_probabilities,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "34e07e3a", + "metadata": { + "id": "gari-syndrome-layout" + }, + "source": [ + "Sample detection events only from the original circuit. The GARI matrix DEM\n", + "stores the transformed matrices for decoding and is not sampled. Copy the\n", + "source syndrome into its physical rows; the added virtual entries stay zero." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d22c7d2c", + "metadata": { + "id": "gari-sample-example" + }, + "outputs": [], + "source": [ + "num_shots = 10\n", + "gari_dets = np.zeros((num_shots, gari_dem.num_detectors), dtype=bool)\n", + "gari_dets[:, gari_layout[\"source_to_gari\"]] = dets[:num_shots]" + ] + }, + { + "cell_type": "markdown", + "id": "abc39564", + "metadata": { + "id": "gari-detector-order" + }, + "source": [ + "The layout is physical-then-virtual. Setting `num_det_orders=0` selects one\n", + "ascending detector order, so Tesseract processes the rows in that order." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bc0c802c", + "metadata": { + "id": "gari-decode-example" + }, + "outputs": [], + "source": [ + "short_beam = tesseract_decoder.make_tesseract_sinter_decoders_dict()[\n", + " \"tesseract-short-beam\"\n", + "]\n", + "short_beam.num_det_orders = 0\n", + "gari_decoder = short_beam.compile_decoder_for_dem(dem=gari_dem).decoder\n", + "predicted_observables = gari_decoder.decode_batch(gari_dets)\n", + "logical_failures = np.count_nonzero(\n", + " np.any(predicted_observables != obs[:num_shots], axis=1)\n", + ")\n", + "print(f\"Logical failures: {logical_failures}/{num_shots}\")" + ] + }, { "cell_type": "markdown", "metadata": { diff --git a/docs/tutorial.py b/docs/tutorial.py index 5f418ebf..5f3028e0 100644 --- a/docs/tutorial.py +++ b/docs/tutorial.py @@ -344,6 +344,51 @@ def run_tesseract_decoder(decoder, dets, obs): results = run_tesseract_decoder(tesseract_config2.compile_decoder(), dets, obs) print_results(results) +# %% [markdown] id="gari-correlated-decoding" +# # Faster Methods of Decoding Correlated Errors with Tesseract +# +# ## GARI +# +# Graph augmentation and rewiring for inference (GARI) transforms a correlated +# CSS detector matrix into a block form for Tesseract; see [Decoding correlated +# errors in quantum LDPC codes](https://doi.org/10.1038/s41467-026-70556-3). +# This example reuses the superdense color-code memory-Z circuit and sampled +# data from above, and applies the XOR prior policy. + +# %% id="gari-transform-example" +gari = tesseract_decoder.demutil.gari + +gari_dem, gari_layout = gari.circuit_to_gari( + circuit, + prior_function=gari.tesseract_xor_prior_probabilities, +) + +# %% [markdown] id="gari-syndrome-layout" +# Sample detection events only from the original circuit. The GARI matrix DEM +# stores the transformed matrices for decoding and is not sampled. Copy the +# source syndrome into its physical rows; the added virtual entries stay zero. + +# %% id="gari-sample-example" +num_shots = 10 +gari_dets = np.zeros((num_shots, gari_dem.num_detectors), dtype=bool) +gari_dets[:, gari_layout["source_to_gari"]] = dets[:num_shots] + +# %% [markdown] id="gari-detector-order" +# The layout is physical-then-virtual. Setting `num_det_orders=0` selects one +# ascending detector order, so Tesseract processes the rows in that order. + +# %% id="gari-decode-example" +short_beam = tesseract_decoder.make_tesseract_sinter_decoders_dict()[ + "tesseract-short-beam" +] +short_beam.num_det_orders = 0 +gari_decoder = short_beam.compile_decoder_for_dem(dem=gari_dem).decoder +predicted_observables = gari_decoder.decode_batch(gari_dets) +logical_failures = np.count_nonzero( + np.any(predicted_observables != obs[:num_shots], axis=1) +) +print(f"Logical failures: {logical_failures}/{num_shots}") + # %% [markdown] id="BoEALeo3OYGp" # # Decoding Wild Stabilizer Codes under Code Capacity Noise with Tesseract # diff --git a/src/py/README.md b/src/py/README.md index 8b566e29..0e9ada0e 100644 --- a/src/py/README.md +++ b/src/py/README.md @@ -678,3 +678,48 @@ nice_calibrated_dem = demutil.regeneralize_spatial_dem( ) # Result will have error probability (0.1 + 0.2) / 2 = 0.15 ``` + +#### GARI transformed matrices + +`demutil.gari.circuit_to_gari` converts a supported correlated CSS Stim +circuit into a GARI matrix DEM and companion layout for Tesseract. It +generates a flattened source DEM with `decompose_errors=False`. Detectors must +follow the repository's fourth-coordinate convention: values `0`–`2` identify +X detectors and `3`–`5` identify Z detectors. + +```python +import stim +from tesseract_decoder import demutil + +circuit = stim.Circuit.from_file("circuitFile.stim") +gari_dem, gari_layout = demutil.gari.circuit_to_gari( + circuit, + prior_function=demutil.gari.tesseract_xor_prior_probabilities, +) +``` + +`circuit_to_gari` returns: + +* `gari_dem`: the augmented detector and logical matrices stored using Stim + DEM syntax. +* `gari_layout`: a `tesseract.gari_layout.v1` dictionary containing the source + and GARI detector counts, the `source_to_gari` detector mapping, and the + `physical_then_virtual` detector order. + +Related public APIs: + +* `demutil.gari.dem_to_matrices(dem)` returns the sparse detector matrix, + sparse logical matrix, and one probability per source error column. +* `demutil.gari.GariTransform` is passed to prior-policy callbacks. It exposes + the transformed detector and logical matrices, the `U` and `V` projection + matrices, the source `e_Z`, `e_X`, and `e_Y` column indices, and the source + detector mapping. +* `paper_prior_probabilities`, `tesseract_xor_prior_probabilities`, and + `tesseract_lp_max_barred_cost_prior_probabilities` return one probability for + each transformed GARI column. A user-defined prior can follow the same + callable interface. + +The returned GARI matrix DEM stores transformed matrices for decoding and must +not be sampled. Sample from the original circuit and use the companion layout +to place its physical syndrome. See the +[GARI tutorial](../../docs/tutorial.ipynb) for a complete decoding example. diff --git a/src/py/_tesseract_py_util/BUILD b/src/py/_tesseract_py_util/BUILD index 59f2131f..284bcd67 100644 --- a/src/py/_tesseract_py_util/BUILD +++ b/src/py/_tesseract_py_util/BUILD @@ -8,6 +8,21 @@ py_library( deps = [ "@pypi//stim", "@pypi//numpy", + "@pypi//scipy", + ], +) + +py_test( + name = "gari_test", + srcs = ["gari_test.py"], + imports = ["..", ".", "../.."], + visibility = ["//:__subpackages__"], + deps = [ + ":_tesseract_py_util", + "@pypi//numpy", + "@pypi//pytest", + "@pypi//stim", + "//src:lib_tesseract_decoder", ], ) diff --git a/src/py/_tesseract_py_util/__init__.py b/src/py/_tesseract_py_util/__init__.py index fe103fec..1cee2fbf 100644 --- a/src/py/_tesseract_py_util/__init__.py +++ b/src/py/_tesseract_py_util/__init__.py @@ -17,6 +17,7 @@ and related utilities, in `decompose_errors.py` and `generalize_dem.py`. """ +from _tesseract_py_util import gari as gari from _tesseract_py_util.demutil import decompose_errors from _tesseract_py_util.generalize_dem import \ generalize as regeneralize_spatial_dem diff --git a/src/py/_tesseract_py_util/demutil_test.py b/src/py/_tesseract_py_util/demutil_test.py index 7aee8897..0935b7c3 100644 --- a/src/py/_tesseract_py_util/demutil_test.py +++ b/src/py/_tesseract_py_util/demutil_test.py @@ -32,6 +32,7 @@ def test_import_exposes_demutil_submodule(): assert hasattr(tesseract_decoder, "demutil") assert hasattr(demutil, "regeneralize_spatial_dem") assert hasattr(demutil, "decompose_errors") + assert hasattr(demutil.gari, "circuit_to_gari") def test_decompose_errors_rejects_unknown_method(): diff --git a/src/py/_tesseract_py_util/gari.py b/src/py/_tesseract_py_util/gari.py new file mode 100644 index 00000000..e86fd322 --- /dev/null +++ b/src/py/_tesseract_py_util/gari.py @@ -0,0 +1,738 @@ +# Copyright 2025 Google LLC +# +# 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. + +"""Graph augmentation and rewiring for inference (GARI). + +This module implements the matrix construction from A. S. Maan et al., +"Decoding correlated errors in quantum LDPC codes," Nature Communications 17, +3965 (2026), https://doi.org/10.1038/s41467-026-70556-3. + +For source columns ``e_Z``, ``e_X``, and ``e_Y``, the supported CSS check +matrix has the form + +:: + + e_Z e_X e_Y + +------------+------------+------------+ + X syndrome | D_X | 0 | D_X U | + +------------+------------+------------+ + Z syndrome | 0 | D_Z | D_Z V | + +------------+------------+------------+ + +over GF(2). GARI substitutes + +``bar(e)_Z = e_Z XOR U e_Y`` and ``bar(e)_X = e_X XOR V e_Y``. + +Columns are emitted as ``[e_Z, e_X, e_Y, bar(e)_Z, bar(e)_X]`` and rows as +``[physical X, physical Z, virtual Z, virtual X]``: + +:: + + e_Z e_X e_Y bar(e)_Z bar(e)_X + +----+----+----+---------+---------+ + physical X syndrome | 0 | 0 | 0 | D_X | 0 | + physical Z syndrome | 0 | 0 | 0 | 0 | D_Z | + virtual Z constraint | I | 0 | U | I | 0 | + virtual X constraint | 0 | I | V | 0 | I | + +----+----+----+---------+---------+ + +The corresponding decoder syndrome is ``[s_X, s_Z, 0, 0]``. The logical map +stays on the original physical variables: +``[L_eZ, L_eX, L_eY, 0, 0]``. These are the GARI transformed matrices. They can +be stored using Stim's DEM syntax, but the resulting GARI DEM is only a matrix +storage and decoding representation. It is not a physical detector error model +and must not be sampled. + +GARI source DEMs must be undecomposed. Circuit conversion generates them with +``decompose_errors=False`` and ``flatten_loops=True``. Matrix extraction also +flattens its input before treating each Stim ``error`` instruction as one +source matrix column. Instructions containing Stim's ``^`` decomposition +separator are not supported. Repeated detector or logical targets are reduced +modulo two, following Stim's GF(2) parity semantics. + +For certain single-basis CSS memory experiments, the paper instead evaluates +the logical observable on ``bar(e)_X`` or ``bar(e)_Z``. That placement is +experiment-specific and is not implemented by this generic transform. + +Every pure ``e_Z`` and ``e_X`` column receives a barred counterpart, including +columns that are not the projection of any ``e_Y`` column. Such an unused pure +column has an all-zero row in ``U`` or ``V``; its virtual identity constraint +therefore only copies ``e`` to ``bar(e)``. This deliberate redundancy keeps the +five-block structure uniform and the physical top-left blocks zero. +""" + +from __future__ import annotations + +import dataclasses +import json +from collections.abc import Callable, Sequence +from pathlib import Path + +import numpy as np +import scipy.optimize +import scipy.sparse +import stim + + +@dataclasses.dataclass +class GariTransform: + """GARI matrices and their source-column and detector mappings.""" + + checks: scipy.sparse.csc_matrix + logicals: scipy.sparse.csc_matrix + u: scipy.sparse.csc_matrix + v: scipy.sparse.csc_matrix + e_z_columns: np.ndarray + e_x_columns: np.ndarray + e_y_columns: np.ndarray + source_to_gari_detectors: np.ndarray + + +def _circuit_to_gari_source_dem( + circuit: stim.Circuit, +) -> stim.DetectorErrorModel: + """Creates the flattened, undecomposed source DEM required by GARI.""" + return circuit.detector_error_model( + decompose_errors=False, + flatten_loops=True, + ).flattened() + + +def _nonzero_column_rows( + matrix: scipy.sparse.csc_matrix, column: int +) -> tuple[int, ...]: + start = matrix.indptr[column] + stop = matrix.indptr[column + 1] + return tuple(int(v) for v in matrix.indices[start:stop]) + + +def _projection_matrix( + pure_columns: scipy.sparse.csc_matrix, + mixed_projections: scipy.sparse.csc_matrix, + mixed_source_columns: np.ndarray, + *, + matrix_name: str, + pure_name: str, +) -> scipy.sparse.csc_matrix: + lookup: dict[tuple[int, ...], int] = {} + for local_column in range(pure_columns.shape[1]): + support = _nonzero_column_rows(pure_columns, local_column) + if support in lookup: + raise ValueError(f"{pure_name} has duplicate columns.") + lookup[support] = local_column + + rows: list[int] = [] + for local_column, source_column in enumerate(mixed_source_columns): + support = _nonzero_column_rows(mixed_projections, local_column) + if support not in lookup: + raise ValueError( + f"{matrix_name} mixed source column {int(source_column)} has " + f"no corresponding {pure_name} pure column." + ) + rows.append(lookup[support]) + + column_count = len(mixed_source_columns) + return scipy.sparse.csc_matrix( + ( + np.ones(column_count, dtype=np.uint8), + ( + np.asarray(rows, dtype=np.int64), + np.arange(column_count, dtype=np.int64), + ), + ), + shape=(pure_columns.shape[1], column_count), + dtype=np.uint8, + ) + + +def dem_to_matrices( + dem: stim.DetectorErrorModel, +) -> tuple[ + scipy.sparse.csc_matrix, scipy.sparse.csc_matrix, np.ndarray +]: + """Extracts matrices from an undecomposed source DEM. + + Repeat blocks and detector shifts are flattened first. Each resulting Stim + ``error`` instruction becomes exactly one source matrix column. The input + should be generated with ``decompose_errors=False``. This function does + not merge duplicate instructions or reconstruct correlations split across + instructions. A Stim ``^`` decomposition separator is rejected. Repeated + detector or logical targets within an instruction are reduced modulo two. + Resulting no-op errors are omitted, and logical-only errors are rejected. + """ + dem = dem.flattened() + detector_rows: list[int] = [] + detector_columns: list[int] = [] + logical_rows: list[int] = [] + logical_columns: list[int] = [] + probabilities: list[float] = [] + + for instruction in dem: + if instruction.type != "error": + continue + targets = instruction.targets_copy() + if any(target.is_separator() for target in targets): + raise ValueError( + "GARI requires a DEM generated with decompose_errors=False." + ) + detectors: set[int] = set() + logicals: set[int] = set() + for target in targets: + if target.is_relative_detector_id(): + detectors ^= {target.val} + elif target.is_logical_observable_id(): + logicals ^= {target.val} + if not detectors and not logicals: + continue + if not detectors: + raise ValueError("GARI does not support logical-only source errors.") + column = len(probabilities) + detector_rows.extend(sorted(detectors)) + detector_columns.extend([column] * len(detectors)) + logical_rows.extend(sorted(logicals)) + logical_columns.extend([column] * len(logicals)) + probabilities.append(float(instruction.args_copy()[0])) + + source_column_count = len(probabilities) + checks = scipy.sparse.csc_matrix( + ( + np.ones(len(detector_rows), dtype=np.uint8), + (detector_rows, detector_columns), + ), + shape=(dem.num_detectors, source_column_count), + dtype=np.uint8, + ) + logicals = scipy.sparse.csc_matrix( + ( + np.ones(len(logical_rows), dtype=np.uint8), + (logical_rows, logical_columns), + ), + shape=(dem.num_observables, source_column_count), + dtype=np.uint8, + ) + return checks, logicals, np.asarray(probabilities, dtype=np.float64) + + +def _matrices_to_gari_dem( + checks: scipy.sparse.csc_matrix, + logicals: scipy.sparse.csc_matrix, + probabilities: np.ndarray, +) -> stim.DetectorErrorModel: + """Stores GARI transformed matrices using Stim's DEM syntax.""" + detector_target = stim.target_relative_detector_id + logical_target = stim.target_logical_observable_id + gari_dem = stim.DetectorErrorModel() + for column, probability in enumerate(probabilities): + targets = [ + detector_target(detector) + for detector in _nonzero_column_rows(checks, column) + ] + targets.extend( + logical_target(observable) + for observable in _nonzero_column_rows(logicals, column) + ) + gari_dem.append("error", float(probability), targets) + + # Declare only dimensions not already implied by the error targets. + if gari_dem.num_detectors < checks.shape[0]: + gari_dem.append( + "detector", [], [detector_target(checks.shape[0] - 1)] + ) + if gari_dem.num_observables < logicals.shape[0]: + gari_dem.append( + "logical_observable", + [], + [logical_target(logicals.shape[0] - 1)], + ) + return gari_dem + + +def _detector_partition_from_fourth_coordinate( + dem: stim.DetectorErrorModel, +) -> tuple[np.ndarray, np.ndarray]: + """Partitions detectors using the repository's fourth-coordinate rule. + + This is the color-code-style convention followed by the test-data circuits + associated with this repository, not a universal Stim convention. The + fourth-coordinate values ``0``, ``1``, or ``2`` identify X detectors, + while values ``3``, ``4``, or ``5`` identify Z detectors. + """ + coordinates = dem.get_detector_coordinates() + x_detectors: list[int] = [] + z_detectors: list[int] = [] + for detector in range(dem.num_detectors): + coordinate = coordinates.get(detector) + if coordinate is None or len(coordinate) < 4: + raise ValueError( + f"Detector {detector} is missing the fourth coordinate " + "required by GARI's color-code-style convention (0, 1, or 2 " + "for X detectors; 3, 4, or 5 for Z detectors)." + ) + basis = coordinate[3] + if basis in (0, 1, 2): + x_detectors.append(detector) + elif basis in (3, 4, 5): + z_detectors.append(detector) + else: + raise ValueError( + f"Detector {detector} has fourth coordinate {basis}; GARI's " + "color-code-style convention requires an integer from 0 to 2 " + "for X detectors or 3 to 5 for Z detectors." + ) + return np.asarray(x_detectors, dtype=np.int64), np.asarray( + z_detectors, dtype=np.int64 + ) + + +def _gari_transform( + checks: scipy.sparse.csc_matrix, + logicals: scipy.sparse.csc_matrix, + *, + x_detectors: Sequence[int], + z_detectors: Sequence[int], +) -> GariTransform: + """Constructs validated GARI transformed matrices over GF(2). + + ``x_detectors`` and ``z_detectors`` partition the source detector rows. + Their sequence order determines the order within the physical X and + physical Z row blocks, respectively. + + Args: + checks: Binary source detector-by-error matrix. + logicals: Binary source observable-by-error matrix. + x_detectors: Source rows containing X-type checks. + z_detectors: Source rows containing Z-type checks. + + Returns: + The transformed checks, physical logical map, projection matrices, + source column classes, and detector mapping. + """ + source_checks = checks.tocsc() + source_logicals = logicals.tocsc() + if source_checks.shape[1] != source_logicals.shape[1]: + raise ValueError( + "checks and logicals must have the same source column count; " + f"found {source_checks.shape[1]} and {source_logicals.shape[1]}." + ) + + detector_count = source_checks.shape[0] + x_rows = np.asarray(x_detectors, dtype=np.int64) + z_rows = np.asarray(z_detectors, dtype=np.int64) + partition = np.concatenate([x_rows, z_rows]) + if not np.array_equal(np.sort(partition), np.arange(detector_count)): + raise ValueError( + "x_detectors and z_detectors must partition all detector rows." + ) + + x_checks = source_checks[x_rows, :] + z_checks = source_checks[z_rows, :] + x_support_counts = np.diff(x_checks.indptr) + z_support_counts = np.diff(z_checks.indptr) + + e_z_columns = np.flatnonzero( + (x_support_counts > 0) & (z_support_counts == 0) + ) + e_x_columns = np.flatnonzero( + (x_support_counts == 0) & (z_support_counts > 0) + ) + e_y_columns = np.flatnonzero( + (x_support_counts > 0) & (z_support_counts > 0) + ) + detectorless_columns = np.flatnonzero( + (x_support_counts == 0) & (z_support_counts == 0) + ) + if detectorless_columns.size: + raise ValueError( + f"Source column {int(detectorless_columns[0])} is detectorless." + ) + + d_x = x_checks[:, e_z_columns] + d_z = z_checks[:, e_x_columns] + d_x_prime = x_checks[:, e_y_columns] + d_z_prime = z_checks[:, e_y_columns] + u = _projection_matrix( + d_x, + d_x_prime, + e_y_columns, + matrix_name="U", + pure_name="D_X", + ) + v = _projection_matrix( + d_z, + d_z_prime, + e_y_columns, + matrix_name="V", + pure_name="D_Z", + ) + e_z_count = len(e_z_columns) + e_x_count = len(e_x_columns) + zero = scipy.sparse.csc_matrix + + # Keep a barred variable for every pure column, even when its row in U or V + # is zero. In that case the identity blocks add the redundant constraint + # e = bar(e), preserving the same block form for every supported model. + identity_z = scipy.sparse.identity(e_z_count, dtype=np.uint8, format="csc") + identity_x = scipy.sparse.identity(e_x_count, dtype=np.uint8, format="csc") + augmented_checks = scipy.sparse.bmat( + [ + [None, None, None, d_x, None], + [None, None, None, None, d_z], + [identity_z, None, u, identity_z, None], + [None, identity_x, v, None, identity_x], + ], + format="csc", + dtype=np.uint8, + ) + + augmented_logicals = scipy.sparse.hstack( + [ + source_logicals[:, e_z_columns], + source_logicals[:, e_x_columns], + source_logicals[:, e_y_columns], + zero((source_logicals.shape[0], e_z_count), dtype=np.uint8), + zero((source_logicals.shape[0], e_x_count), dtype=np.uint8), + ], + format="csc", + ).astype(np.uint8) + + source_to_gari = np.empty(detector_count, dtype=np.int64) + source_to_gari[partition] = np.arange(detector_count, dtype=np.int64) + return GariTransform( + checks=augmented_checks, + logicals=augmented_logicals, + u=u, + v=v, + e_z_columns=e_z_columns, + e_x_columns=e_x_columns, + e_y_columns=e_y_columns, + source_to_gari_detectors=source_to_gari, + ) + + +def _physical_probability_blocks( + transform: GariTransform, source_probabilities: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + probabilities = np.asarray(source_probabilities, dtype=np.float64) + return ( + probabilities[transform.e_z_columns], + probabilities[transform.e_x_columns], + probabilities[transform.e_y_columns], + ) + + +def paper_prior_probabilities( + transform: GariTransform, source_probabilities: np.ndarray +) -> np.ndarray: + """Returns the published GARI initialization in GARI column order. + + Physical ``e_Z``, ``e_X``, and ``e_Y`` variables retain their source + probabilities. Every auxiliary variable is assigned probability exactly + ``0.5``. In Tesseract this gives the auxiliary variable zero search cost, + which can produce a very large search space. + """ + p_e_z, p_e_x, p_e_y = _physical_probability_blocks( + transform, source_probabilities + ) + return np.concatenate( + [ + p_e_z, + p_e_x, + p_e_y, + np.full(len(p_e_z), 0.5), + np.full(len(p_e_x), 0.5), + ] + ) + + +def _barred_xor_probabilities( + base_probabilities: np.ndarray, + y_probabilities: np.ndarray, + projection_matrix: scipy.sparse.csc_matrix, +) -> np.ndarray: + """Returns marginals of ``base XOR projection_matrix @ e_Y``.""" + with np.errstate(divide="ignore"): + log_even_bias = np.log1p(-2 * base_probabilities) + ( + projection_matrix @ np.log1p(-2 * y_probabilities) + ) + return -0.5 * np.expm1(log_even_bias) + + +def tesseract_xor_prior_probabilities( + transform: GariTransform, source_probabilities: np.ndarray +) -> np.ndarray: + """Returns experimental independent-XOR marginals for Tesseract. + + Each auxiliary probability is the independent Bernoulli parity marginal + implied by ``bar(e)_Z = e_Z XOR U e_Y`` or + ``bar(e)_X = e_X XOR V e_Y``. The computation uses log-domain products + for numerical stability and does not clip invalid inputs. + + This is a Tesseract-specific experimental heuristic, not the published + GARI prior. It only defines auxiliary search weights and makes no claim + about decoding optimality. + """ + p_e_z, p_e_x, p_e_y = _physical_probability_blocks( + transform, source_probabilities + ) + p_bar_e_z = _barred_xor_probabilities( + p_e_z, p_e_y, transform.u + ) + p_bar_e_x = _barred_xor_probabilities( + p_e_x, p_e_y, transform.v + ) + return np.concatenate([p_e_z, p_e_x, p_e_y, p_bar_e_z, p_bar_e_x]) + + +def tesseract_lp_max_barred_cost_prior_probabilities( + transform: GariTransform, source_probabilities: np.ndarray +) -> np.ndarray: + """Balances physical and barred-variable search costs with two LPs. + + The source costs ``c = log((1-p)/p)`` are ordered as ``[e_Z, e_X, e_Y]``. + The auxiliary costs ``g`` are ordered as ``[bar(e)_Z, bar(e)_X]``. The + incidence matrix is ``A = [[I, 0], [0, I], [U.T, V.T]]``, so the residual + physical costs are ``r = c - A g``. + + Maximizing only ``sum(g)`` can put most of the cost on a few variables and + leave many physical or auxiliary costs at zero. A zero cost becomes + probability ``0.5`` and gives Tesseract no search preference. The first LP + instead maximizes a common floor for every entry of ``r`` and ``g``. The + second LP keeps that floor and then maximizes ``sum(g)``. This gives a more + balanced set of search costs while still favoring the barred variables. + + The result is ``[r, g]`` converted back to probabilities in GARI column + order. This experimental policy is not part of the GARI paper. It only + defines search costs and makes no claim about decoding optimality. Solver + failure is a hard error; there is no fallback. + """ + p_e_z, p_e_x, p_e_y = _physical_probability_blocks( + transform, source_probabilities + ) + physical_probabilities = np.concatenate([p_e_z, p_e_x, p_e_y]) + source_costs = np.log1p(-physical_probabilities) - np.log( + physical_probabilities + ) + # A maps [bar(e)_Z, bar(e)_X] costs into [e_Z, e_X, e_Y] costs. + identity = scipy.sparse.identity + cost_matrix = scipy.sparse.bmat( + [ + [identity(len(p_e_z), format="csc"), None], + [None, identity(len(p_e_x), format="csc")], + [transform.u.T, transform.v.T], + ], + format="csc", + ) + auxiliary_count = cost_matrix.shape[1] + if auxiliary_count == 0: + return physical_probabilities + + # First maximize t subject to every physical and auxiliary cost being at + # least t: c - A g >= t and g >= t. + floor_constraints = scipy.sparse.bmat( + [ + [cost_matrix, np.ones((len(source_costs), 1))], + [ + -identity(auxiliary_count, format="csc"), + np.ones((auxiliary_count, 1)), + ], + ], + format="csc", + ) + floor_objective = np.zeros(auxiliary_count + 1, dtype=np.float64) + floor_objective[-1] = -1 + floor_result = scipy.optimize.linprog( + floor_objective, + A_ub=floor_constraints, + b_ub=np.concatenate([source_costs, np.zeros(auxiliary_count)]), + bounds=(0, None), + method="highs", + ) + if not floor_result.success: + raise RuntimeError( + "LP prior floor solver failed: " + str(floor_result.message) + ) + tolerance = 1e-7 * max( + 1.0, float(np.max(source_costs, initial=0.0)) + ) + cost_floor = max(0.0, float(floor_result.x[-1]) - tolerance) + + # Then maximize the total barred cost without lowering the common floor. + objective = -np.ones(auxiliary_count, dtype=np.float64) + result = scipy.optimize.linprog( + objective, + A_ub=cost_matrix, + b_ub=source_costs - cost_floor, + bounds=(cost_floor, None), + method="highs", + ) + if not result.success: + raise RuntimeError( + "LP prior barred-cost solver failed: " + str(result.message) + ) + auxiliary_costs = np.asarray(result.x) + if np.min(auxiliary_costs) < cost_floor - tolerance: + raise RuntimeError( + "LP max-barred-cost solver returned an infeasible solution." + ) + auxiliary_costs = np.maximum(auxiliary_costs, 0.0) + residual_costs = source_costs - np.asarray( + cost_matrix @ auxiliary_costs + ).reshape(-1) + if np.min(residual_costs) < cost_floor - tolerance: + raise RuntimeError( + "LP max-barred-cost solver returned an infeasible solution." + ) + # Normalize only active-constraint noise accepted by the LP solver. + residual_costs = np.maximum(residual_costs, 0.0) + gari_costs = np.concatenate([residual_costs, auxiliary_costs]) + return np.exp(-np.logaddexp(0, gari_costs)) + + +def _build_gari_dem( + transform: GariTransform, + source_probabilities: np.ndarray, + *, + prior_function: Callable[[GariTransform, np.ndarray], np.ndarray], +) -> stim.DetectorErrorModel: + """Builds a GARI DEM using an explicit prior policy. + + ``prior_function`` may be one of this module's three built-in policies or + a user-defined callable. Its output is validated before serialization. + Stim's DEM syntax is used only to store the GARI transformed matrices. The + result is not a physical detector error model and must not be sampled. + """ + def validated_probabilities( + values: np.ndarray, expected_count: int, name: str + ) -> np.ndarray: + result = np.asarray(values, dtype=np.float64) + if result.shape != (expected_count,) or not np.all( + (result > 0) & (result <= 0.5) + ): + raise ValueError( + f"{name} must contain {expected_count} finite values in " + "(0, 0.5]." + ) + return result + + source_count = ( + len(transform.e_z_columns) + + len(transform.e_x_columns) + + len(transform.e_y_columns) + ) + probabilities = validated_probabilities( + source_probabilities, + source_count, + "source_probabilities", + ) + gari_probabilities = validated_probabilities( + prior_function(transform, probabilities), + transform.checks.shape[1], + "prior_function probabilities", + ) + return _matrices_to_gari_dem( + transform.checks, transform.logicals, gari_probabilities + ) + + +def circuit_to_gari( + circuit: stim.Circuit, + *, + prior_function: Callable[[GariTransform, np.ndarray], np.ndarray], +) -> tuple[stim.DetectorErrorModel, dict[str, object]]: + """Converts a supported CSS circuit into a GARI matrix DEM and v1 layout. + + The source DEM is generated undecomposed (``decompose_errors=False``) and + flattened. Every detector must follow the repository's fourth-coordinate + convention: integer values 0–2 identify X detectors and 3–5 identify Z + detectors. The returned DEM stores transformed matrices for decoding and + must not be sampled. + """ + source_dem = _circuit_to_gari_source_dem(circuit) + checks, logicals, probabilities = dem_to_matrices(source_dem) + x_detectors, z_detectors = _detector_partition_from_fourth_coordinate( + source_dem + ) + transform = _gari_transform( + checks, + logicals, + x_detectors=x_detectors, + z_detectors=z_detectors, + ) + gari_dem = _build_gari_dem( + transform, probabilities, prior_function=prior_function + ) + layout = { + "schema": "tesseract.gari_layout.v1", + "source_detector_count": len(transform.source_to_gari_detectors), + "gari_detector_count": transform.checks.shape[0], + "source_to_gari": transform.source_to_gari_detectors.tolist(), + "detector_order": "physical_then_virtual", + } + return gari_dem, layout + + +def call_gari(circuit_fname: str, prior_name: str, output_dir: str) -> None: + """Converts one circuit and writes its GARI DEM and layout files.""" + prior_function = { + "paper": paper_prior_probabilities, + "xor": tesseract_xor_prior_probabilities, + "lp-max-barred-cost": tesseract_lp_max_barred_cost_prior_probabilities, + }[prior_name] + gari_dem, layout = circuit_to_gari( + stim.Circuit.from_file(circuit_fname), + prior_function=prior_function, + ) + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + output_name = f"{Path(circuit_fname).stem}_gari_{prior_name.replace('-', '_')}" + gari_dem.to_file(output_path / f"{output_name}.dem") + (output_path / f"{output_name}_layout.json").write_text( + json.dumps(layout, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def main() -> None: + import argparse + + parser = argparse.ArgumentParser( + description=( + "Convert one Stim circuit into a GARI matrix DEM and " + "detector-layout JSON file." + ) + ) + parser.add_argument( + "--circuit", required=True, help="Input Stim circuit file." + ) + parser.add_argument( + "--prior", + choices=("paper", "xor", "lp-max-barred-cost"), + required=True, + help="Prior policy used for the GARI matrix probabilities.", + ) + parser.add_argument( + "--out-dir", + required=True, + help=( + "Output directory, created if needed. Files are named " + "_gari_.dem and " + "_gari__layout.json." + ), + ) + args = parser.parse_args() + call_gari(args.circuit, args.prior, args.out_dir) + + +if __name__ == "__main__": + main() diff --git a/src/py/_tesseract_py_util/gari_test.py b/src/py/_tesseract_py_util/gari_test.py new file mode 100644 index 00000000..530adaf8 --- /dev/null +++ b/src/py/_tesseract_py_util/gari_test.py @@ -0,0 +1,195 @@ +# Copyright 2025 Google LLC +# +# 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 json + +import numpy as np +import pytest +import stim + +from _tesseract_py_util import gari +from tesseract_decoder import demutil + + +def _tiny_circuit(): + return stim.Circuit(""" + R 0 1 2 3 4 5 + CORRELATED_ERROR(0.01) X0 X2 X4 + CORRELATED_ERROR(0.02) X1 X3 X5 + CORRELATED_ERROR(0.04) X0 X1 X2 X3 X4 + M 0 1 2 3 4 5 + DETECTOR(0, 0, 0, 0) rec[-6] + DETECTOR(0, 0, 0, 3) rec[-5] + DETECTOR(0, 0, 0, 2) rec[-4] + DETECTOR(0, 0, 0, 4) rec[-3] + OBSERVABLE_INCLUDE(0) rec[-2] + OBSERVABLE_INCLUDE(1) rec[-1] + """) + + +def _tiny_model(): + source_dem = gari._circuit_to_gari_source_dem(_tiny_circuit()) + checks, logicals, probabilities = gari.dem_to_matrices(source_dem) + x_detectors, z_detectors = gari._detector_partition_from_fourth_coordinate( + source_dem + ) + transform = gari._gari_transform( + checks, + logicals, + x_detectors=x_detectors, + z_detectors=z_detectors, + ) + return probabilities, transform + + +def test_tiny_transform(): + folded_dem = stim.DetectorErrorModel(""" + repeat 2 { + error(0.1) D0 L0 + shift_detectors 1 + } + """) + checks, logicals, probabilities = gari.dem_to_matrices(folded_dem) + np.testing.assert_array_equal(checks.toarray(), np.eye(2, dtype=np.uint8)) + np.testing.assert_array_equal(logicals.toarray(), [[1, 1]]) + np.testing.assert_allclose(probabilities, [0.1, 0.1]) + + with pytest.raises(ValueError, match="integer from 0 to 2"): + gari._detector_partition_from_fourth_coordinate( + stim.DetectorErrorModel("detector(0, 0, 0, 2.5) D0") + ) + + with pytest.raises(ValueError, match="decompose_errors=False"): + gari.dem_to_matrices(stim.DetectorErrorModel("error(0.1) D0 ^ D1")) + + checks, logicals, probabilities = gari.dem_to_matrices( + stim.DetectorErrorModel(""" + error(0.1) D0 D0 + error(0.2) D0 D0 D1 L0 L1 L1 + """) + ) + np.testing.assert_array_equal(checks.toarray(), [[0], [1]]) + np.testing.assert_array_equal(logicals.toarray(), [[1], [0]]) + np.testing.assert_allclose(probabilities, [0.2]) + + with pytest.raises(ValueError, match="logical-only source errors"): + gari.dem_to_matrices( + stim.DetectorErrorModel("error(0.1) D0 D0 L0") + ) + + _, transform = _tiny_model() + np.testing.assert_array_equal( + transform.checks.toarray(), + [ + [0, 0, 0, 1, 0], + [0, 0, 0, 1, 0], + [0, 0, 0, 0, 1], + [0, 0, 0, 0, 1], + [1, 0, 1, 1, 0], + [0, 1, 1, 0, 1], + ], + ) + np.testing.assert_array_equal( + transform.logicals.toarray(), + [[1, 0, 1, 0, 0], [0, 1, 0, 0, 0]], + ) + np.testing.assert_array_equal( + transform.source_to_gari_detectors, [0, 2, 1, 3] + ) + + +def test_prior_probabilities_and_gari_dem_round_trip(): + source_probabilities, transform = _tiny_model() + paper_probabilities = gari.paper_prior_probabilities( + transform, source_probabilities + ) + np.testing.assert_array_equal( + paper_probabilities, + [0.01, 0.02, 0.04, 0.5, 0.5], + ) + xor_probabilities = gari.tesseract_xor_prior_probabilities( + transform, source_probabilities + ) + np.testing.assert_allclose( + xor_probabilities, [0.01, 0.02, 0.04, 0.0492, 0.0584] + ) + lp_probabilities = gari.tesseract_lp_max_barred_cost_prior_probabilities( + transform, source_probabilities + ) + source_costs = np.log1p(-paper_probabilities[:3]) - np.log( + paper_probabilities[:3] + ) + lp_costs = np.log1p(-lp_probabilities) - np.log(lp_probabilities) + assert np.all(lp_costs > 0) + np.testing.assert_allclose( + lp_costs[2:], source_costs[2] / 3, rtol=1e-6 + ) + np.testing.assert_allclose( + [ + lp_costs[0] + lp_costs[3], + lp_costs[1] + lp_costs[4], + lp_costs[2] + lp_costs[3] + lp_costs[4], + ], + source_costs, + ) + + gari_dem = gari._build_gari_dem( + transform, + source_probabilities, + prior_function=gari.tesseract_xor_prior_probabilities, + ) + checks, logicals, probabilities = gari.dem_to_matrices(gari_dem) + assert gari_dem.num_detectors == transform.checks.shape[0] + assert gari_dem.num_observables == transform.logicals.shape[0] + assert (checks != transform.checks).nnz == 0 + assert (logicals != transform.logicals).nnz == 0 + np.testing.assert_allclose(probabilities, xor_probabilities) + + +def test_public_circuit_conversion_and_file_output(tmp_path): + public_gari = demutil.gari + circuit = _tiny_circuit() + gari_dem, layout = public_gari.circuit_to_gari( + circuit, + prior_function=public_gari.tesseract_xor_prior_probabilities, + ) + assert layout == { + "schema": "tesseract.gari_layout.v1", + "source_detector_count": 4, + "gari_detector_count": 6, + "source_to_gari": [0, 2, 1, 3], + "detector_order": "physical_then_virtual", + } + assert gari_dem.num_detectors == 6 + assert gari_dem.num_observables == 2 + + circuit_path = tmp_path / "tiny.stim" + circuit.to_file(circuit_path) + output_dir = tmp_path / "gari" + public_gari.call_gari(str(circuit_path), "xor", str(output_dir)) + output_name = "tiny_gari_xor" + written_dem = stim.DetectorErrorModel.from_file( + output_dir / f"{output_name}.dem" + ) + written_layout = json.loads( + (output_dir / f"{output_name}_layout.json").read_text( + encoding="utf-8" + ) + ) + assert str(written_dem) == str(gari_dem) + assert written_layout == layout + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))