qec 0.3.0__py3-none-any.whl → 0.3.3__py3-none-any.whl

qec/code_constructions/stabilizer_code.py

@@ -0,0 +1,654 @@
+from qec.utils.sparse_binary_utils import (
+    convert_to_binary_scipy_sparse,
+    binary_csr_matrix_to_dict,
+)
+from qec.utils.binary_pauli_utils import (
+    symplectic_product,
+    check_binary_pauli_matrices_commute,
+    pauli_str_to_binary_pcm,
+    binary_pcm_to_pauli_str,
+    binary_pauli_hamming_weight,
+)
+import re
+import numpy.typing
+import numpy as np
+import scipy.sparse
+import json
+from tqdm import tqdm
+from ldpc import BpOsdDecoder
+import ldpc.mod2
+import time
+from typing import Tuple, Optional, Union, Sequence
+from pathlib import Path
+import logging
+
+logging.basicConfig(level=logging.DEBUG)
+
+
+class StabilizerCode(object):
+    """
+    A quantum stabilizer code, which defines and manipulates stabilizer generators,
+    computes logical operators, and stores parameters such as the number of physical qubits
+    and the number of logical qubits.
+
+    Parameters
+    ----------
+    stabilizers : numpy.typing.ArrayLike or scipy.sparse.spmatrix or list
+        Either a binary parity check matrix (with an even number of columns),
+        or a list of Pauli strings that specify the stabilizers of the code.
+    name : str, optional
+        A name for the code. Defaults to "stabilizer code".
+
+    Attributes
+    ----------
+    name : str
+        The name of the code.
+    stabilizer_matrix : scipy.sparse.spmatrix
+        The binary parity check matrix representation of the stabilizers.
+    phyical_qubit_count : int
+        The number of physical qubits in the code.
+    logical_qubit_count : int
+        The number of logical qubits in the code.
+    code_distance : int
+        (Not computed by default) The distance of the code, if known or computed.
+    logical_operator_basis : scipy.sparse.spmatrix or None
+        A basis for the logical operators of the code.
+    """
+
+    def __init__(
+        self,
+        stabilizers: Union[np.ndarray, scipy.sparse.spmatrix, list],
+        name: str = None,
+    ):
+        """
+        Construct a StabilizerCode instance from either a parity check matrix or a list of
+        Pauli stabilizers.
+
+        Parameters
+        ----------
+        stabilizers : numpy.typing.ArrayLike or scipy.sparse.spmatrix or list
+            Either a binary parity check matrix (with an even number of columns),
+            or a list of Pauli strings that specify the stabilizers of the code.
+        name : str, optional
+            A name for the code. If None, it defaults to "stabilizer code".
+
+        Raises
+        ------
+        TypeError
+            If `stabilizers` is not an array-like, sparse matrix, or list of Pauli strings.
+        ValueError
+            If the parity check matrix does not have an even number of columns,
+            or the stabilizers do not mutually commute.
+        """
+        self.name = name if name else "stabilizer code"
+
+        self.stabilizer_matrix = None
+        self.physical_qubit_count = None
+        self._logical_qubit_count = None
+        self.code_distance = None
+        self.logical_operator_basis = None
+
+        if isinstance(stabilizers, list):
+            stabilizers = np.array(stabilizers)
+
+        if not isinstance(stabilizers, (np.ndarray, scipy.sparse.spmatrix)):
+            raise TypeError(
+                "Please provide either a parity check matrix or a list of Pauli stabilizers."
+            )
+
+        if isinstance(stabilizers, np.ndarray) and stabilizers.dtype.kind in {"U", "S"}:
+            self.stabilizer_matrix = pauli_str_to_binary_pcm(stabilizers)
+        else:
+            if stabilizers.shape[1] % 2 == 0:
+                self.stabilizer_matrix = convert_to_binary_scipy_sparse(stabilizers)
+            else:
+                raise ValueError(
+                    "The parity check matrix must have an even number of columns."
+                )
+
+        self.physical_qubit_count = self.stabilizer_matrix.shape[1] // 2
+
+        # Check that stabilizers commute
+        if not self.check_stabilizers_commute():
+            raise ValueError("The stabilizers do not commute.")
+
+    @property
+    def logical_qubit_count(self):
+        if self._logical_qubit_count is None:
+            self._logical_qubit_count = self.physical_qubit_count - ldpc.mod2.rank(
+                self.stabilizer_matrix, method="dense"
+            )
+        return self._logical_qubit_count
+
+    @logical_qubit_count.setter
+    def logical_qubit_count(self, value):
+        self._logical_qubit_count = value
+
+    @property
+    def pauli_stabilizers(self):
+        """
+        Get or set the stabilizers in Pauli string format.
+
+        Returns
+        -------
+        np.ndarray
+            An array of Pauli strings representing the stabilizers.
+        """
+        return binary_pcm_to_pauli_str(self.stabilizer_matrix)
+
+    @pauli_stabilizers.setter
+    def pauli_stabilizers(self, pauli_stabilizers: np.ndarray):
+        """
+        Set the stabilizers using Pauli strings.
+
+        Parameters
+        ----------
+        pauli_stabilizers : np.ndarray
+            An array of Pauli strings representing the stabilizers.
+
+        Raises
+        ------
+        AssertionError
+            If the newly set stabilizers do not commute.
+        """
+        self.stabilizer_matrix = pauli_str_to_binary_pcm(pauli_stabilizers)
+        if not self.check_stabilizers_commute():
+            raise ValueError("The stabilizers do not commute.")
+
+    def check_stabilizers_commute(self) -> bool:
+        """
+        Check whether the current set of stabilizers mutually commute.
+
+        Returns
+        -------
+        bool
+            True if all stabilizers commute, otherwise False.
+        """
+        return check_binary_pauli_matrices_commute(
+            self.stabilizer_matrix, self.stabilizer_matrix
+        )
+
+    def compute_logical_basis(self) -> scipy.sparse.spmatrix:
+        """
+        Compute a basis for the logical operators of the code by extending the parity check
+        matrix. The resulting basis operators are stored in `self.logicals`.
+
+        Returns
+        -------
+        scipy.sparse.spmatrix
+            A basis for the logical operators in binary representation.
+
+        Notes
+        -----
+        This method uses the kernel of the parity check matrix to find operators that
+        commute with all stabilizers, and then identifies a subset that spans the space
+        of logical operators.
+        """
+        kernel_h = ldpc.mod2.kernel(self.stabilizer_matrix)
+
+        # Sort the rows of the kernel by weight
+        row_weights = kernel_h.getnnz(axis=1)
+        sorted_rows = np.argsort(row_weights)
+        kernel_h = kernel_h[sorted_rows, :]
+
+        swapped_kernel = scipy.sparse.hstack(
+            [
+                kernel_h[:, self.physical_qubit_count :],
+                kernel_h[:, : self.physical_qubit_count],
+            ]
+        )
+
+        logical_stack = scipy.sparse.vstack([self.stabilizer_matrix, swapped_kernel])
+        p_rows = ldpc.mod2.pivot_rows(logical_stack)
+
+        self.logical_operator_basis = logical_stack[
+            p_rows[self.stabilizer_matrix.shape[0] :]
+        ]
+
+        if self.logical_operator_basis.nnz == 0:
+            self.code_distance = np.inf
+            return self.logical_operator_basis
+
+        basis_minimum_hamming_weight = np.min(
+            binary_pauli_hamming_weight(self.logical_operator_basis).flatten()
+        )
+
+        # Update distance based on the minimum hamming weight of the logical operators in this basis
+        if self.code_distance is None:
+            self.code_distance = basis_minimum_hamming_weight
+        elif basis_minimum_hamming_weight < self.code_distance:
+            self.code_distance = basis_minimum_hamming_weight
+        else:
+            pass
+
+        return logical_stack[p_rows[self.stabilizer_matrix.shape[0] :]]
+
+    @property
+    def logical_operator_basis(self):
+        if self._logical_operator_basis is None:
+            self._logical_operator_basis = self.compute_logical_basis()
+
+        return self._logical_operator_basis
+
+    @logical_operator_basis.setter
+    def logical_operator_basis(self, basis: scipy.sparse.spmatrix):
+        self._logical_operator_basis = basis
+
+    def check_valid_logical_basis(self) -> bool:
+        """
+        Validate that the stored logical operators form a proper logical basis for the code.
+
+        Checks that they commute with the stabilizers, pairwise anti-commute (in the symplectic
+        sense), and have full rank.
+
+        Returns
+        -------
+        bool
+            True if the logical operators form a valid basis, otherwise False.
+        """
+        try:
+            assert check_binary_pauli_matrices_commute(
+                self.stabilizer_matrix, self.logical_operator_basis
+            ), "Logical operators do not commute with stabilizers."
+
+            logical_product = symplectic_product(
+                self.logical_operator_basis, self.logical_operator_basis
+            )
+            logical_product.eliminate_zeros()
+            assert (
+                logical_product.nnz != 0
+            ), "The logical operators do not anti-commute with one another."
+
+            assert (
+                ldpc.mod2.rank(self.logical_operator_basis, method="dense")
+                == 2 * self.logical_qubit_count
+            ), "The logical operators do not form a basis for the code."
+
+            assert (
+                self.logical_operator_basis.shape[0] == 2 * self.logical_qubit_count
+            ), "The logical operators are not linearly independent."
+
+        except AssertionError as e:
+            logging.error(e)
+            return False
+
+        return True
+
+    def compute_exact_code_distance(
+        self, timeout: float = 0.5
+    ) -> Tuple[Optional[int], float]:
+        """
+        Compute the distance of the code by searching through linear combinations of
+        logical operators and stabilizers, returning a tuple of the minimal Hamming weight
+        found and the fraction of logical operators considered before timing out.
+
+        Parameters
+        ----------
+        timeout : float, optional
+            The time limit (in seconds) for the exhaustive search. Default is 0.5 seconds. To obtain the exact distance, set to `np.inf`.
+
+        Returns
+        -------
+        Tuple[Optional[int], float]
+            A tuple containing:
+            - The best-known distance of the code as an integer (or `None` if no distance was found).
+            - The fraction of logical combinations considered before the search ended.
+
+        Notes
+        -----
+        - We compute the row span of both the stabilizers and the logical operators.
+        - For every logical operator in the logical span, we add (mod 2) each stabilizer
+          in the stabilizer span to form candidate logical operators.
+        - We compute the Hamming weight of each candidate operator (i.e. how many qubits
+          are acted upon by the operator).
+        - We track the minimal Hamming weight encountered. If `timeout` is exceeded,
+          we immediately return the best distance found so far.
+
+        Examples
+        --------
+        >>> code = StabilizerCode(["XZZX", "ZZXX"])
+        >>> dist, fraction = code.compute_exact_code_distance(timeout=1.0)
+        >>> print(dist, fraction)
+        """
+        start_time = time.time()
+
+        stabilizer_span = ldpc.mod2.row_span(self.stabilizer_matrix)[1:]
+        logical_span = ldpc.mod2.row_span(self.logical_operator_basis)[1:]
+
+        if self.code_distance is None:
+            distance = np.inf
+        else:
+            distance = self.code_distance
+
+        logicals_considered = 0
+        total_logical_operators = stabilizer_span.shape[0] * logical_span.shape[0]
+
+        for logical in logical_span:
+            if time.time() - start_time > timeout:
+                break
+            for stabilizer in stabilizer_span:
+                if time.time() - start_time > timeout:
+                    break
+                candidate_logical = logical + stabilizer
+                candidate_logical.data %= 2
+
+                hamming_weight = binary_pauli_hamming_weight(candidate_logical)[0]
+                if hamming_weight < distance:
+                    distance = hamming_weight
+                logicals_considered += 1
+
+        self.code_distance = distance
+        fraction_considered = logicals_considered / total_logical_operators
+
+        return (
+            (int(distance), fraction_considered)
+            if distance != np.inf
+            else (None, fraction_considered)
+        )
+
+    def get_code_parameters(self) -> tuple:
+        """
+        Return the parameters of the code as a tuple: (n, k, d).
+
+        Returns
+        -------
+        tuple
+            A tuple of integers representing the number of physical qubits, logical qubits,
+            and the distance of the code.
+        """
+        return self.physical_qubit_count, self.logical_qubit_count, self.code_distance
+
+    def __repr__(self):
+        """
+        Return an unambiguous string representation of the StabilizerCode instance.
+
+        Returns
+        -------
+        str
+            An unambiguous representation for debugging and development.
+        """
+        return f"Name: {self.name}, Class: Stabilizer Code"
+
+    def __str__(self):
+        """
+        Return a string describing the stabilizer code, including its parameters.
+
+        Returns
+        -------
+        str
+            A human-readable string with the name, n, k, and d parameters of the code.
+        """
+        return f"< Stabilizer Code, Name: {self.name}, Parameters: [[{self.physical_qubit_count}, {self.logical_qubit_count}, {self.code_distance}]] >"
+
+    def estimate_min_distance(
+        self,
+        timeout_seconds: float = 0.25,
+        p: float = 0.25,
+        reduce_logical_basis: bool = False,
+        decoder: Optional[BpOsdDecoder] = None,
+    ) -> int:
+        """
+        Estimate the minimum distance of the stabilizer code using a BP+OSD decoder-based search.
+
+        Parameters
+        ----------
+        timeout_seconds : float
+            Time limit in seconds for the search. Default: 0.25
+        p : float
+            Probability for including each logical operator in trial combinations. Default: 0.25
+        reduce_logical_basis : bool
+            Whether to attempt reducing logical operator basis. Default: False
+        decoder : Optional[BpOsdDecoder]
+            Pre-configured BP+OSD decoder. If None, initialises with default settings.
+
+        Returns
+        -------
+        int
+            Best estimate of code distance found within time limit
+        """
+        if self.logical_operator_basis is None:
+            self.logical_operator_basis = self.compute_logical_basis()
+
+        # Initial setup of decoder and parameters
+        bp_osd, stack, full_rank_stabilizer_matrix, min_distance, max_distance = (
+            self._setup_distance_estimation_decoder(decoder)
+        )
+
+        # Initialize storage for candidate logicals and tracking
+        candidate_logicals = []
+        weight_one_syndromes_searched = 0
+
+        # Main search loop
+        start_time = time.time()
+        with tqdm(total=timeout_seconds, desc="Estimating distance") as pbar:
+            while time.time() - start_time < timeout_seconds:
+                # Update progress bar
+                elapsed = time.time() - start_time
+                pbar.update(elapsed - pbar.n)
+
+                # Initialize an empty dummy syndrome
+                dummy_syndrome = np.zeros(stack.shape[0], dtype=np.uint8)
+
+                if weight_one_syndromes_searched < self.logical_operator_basis.shape[0]:
+                    # Try each logical operator individually first
+                    dummy_syndrome[
+                        full_rank_stabilizer_matrix.shape[0]
+                        + weight_one_syndromes_searched
+                    ] = 1
+                    weight_one_syndromes_searched += 1
+                else:
+                    # Randomly pick a combination of logical rows
+                    while True:
+                        random_mask = np.random.choice(
+                            [0, 1],
+                            size=self.logical_operator_basis.shape[0],
+                            p=[1 - p, p],
+                        )
+                        if np.any(random_mask):
+                            break
+                    for idx, bit in enumerate(random_mask):
+                        if bit == 1:
+                            dummy_syndrome[full_rank_stabilizer_matrix.shape[0] + idx] = 1
+
+                candidate = bp_osd.decode(dummy_syndrome)
+                w = np.count_nonzero(
+                    candidate[: self.physical_qubit_count]
+                    | candidate[self.physical_qubit_count :]
+                )
+
+                if w < min_distance:
+                    min_distance = w
+                if w < max_distance and reduce_logical_basis:
+                    lc = np.hstack(
+                        [
+                            candidate[self.physical_qubit_count :],
+                            candidate[: self.physical_qubit_count],
+                        ]
+                    )
+                    candidate_logicals.append(lc)
+
+                # Reduce logical operator basis if we have enough candidates
+                if (
+                    len(candidate_logicals) >= self.logical_qubit_count
+                    and reduce_logical_basis
+                ):
+                    self._reduce_logical_operator_basis(candidate_logicals)
+                    (
+                        bp_osd,
+                        stack,
+                        full_rank_stabilizer_matrix,
+                        min_distance,
+                        max_distance,
+                    ) = self._setup_distance_estimation_decoder(decoder)
+                    candidate_logicals = []
+                    weight_one_syndromes_searched = 0
+
+                pbar.set_description(
+                    f"Estimating distance: min-weight found <= {min_distance}, "
+                    f"basis weights: {self.logical_basis_weights()}"
+                )
+
+        # Final basis reduction if needed
+        if reduce_logical_basis and candidate_logicals:
+            self._reduce_logical_operator_basis(candidate_logicals)
+
+        self.code_distance = min_distance
+        return min_distance
+
+    def _setup_distance_estimation_decoder(
+        self, decoder: Optional[BpOsdDecoder] = None
+    ) -> Tuple[BpOsdDecoder, scipy.sparse.spmatrix, scipy.sparse.spmatrix, int, int]:
+        """
+        Set up decoder and initial parameters.
+
+        Parameters
+        ----------
+        decoder : Optional[BpOsdDecoder]
+            Pre-configured decoder. If None, initialises with default settings.
+
+        Returns
+        -------
+        Tuple[BpOsdDecoder, scipy.sparse.spmatrix, scipy.sparse.spmatrix, int, int]
+            Returns (decoder, stack matrix, full rank stabilizer matrix, min distance, max distance)
+        """
+        # Remove redundant rows from stabilizer matrix
+        p_rows = ldpc.mod2.pivot_rows(self.stabilizer_matrix)
+        full_rank_stabilizer_matrix = self.stabilizer_matrix[p_rows]
+
+        # Build a stacked matrix of stabilizers and logicals
+        stack = scipy.sparse.vstack(
+            [full_rank_stabilizer_matrix, self.logical_operator_basis]
+        ).tocsr()
+
+        # Initial distance estimate from current logicals
+        min_distance = np.min(binary_pauli_hamming_weight(self.logical_operator_basis))
+        max_distance = np.max(self.logical_basis_weights())
+
+        # Set up BP+OSD decoder if not provided
+        if decoder is None:
+            decoder = BpOsdDecoder(
+                stack,
+                error_rate=0.1,
+                max_iter=10,
+                bp_method="ms",
+                schedule="parallel",
+                ms_scaling_factor=1.0,
+                osd_method="osd_0",
+                osd_order=0,
+            )
+
+        return decoder, stack, full_rank_stabilizer_matrix, min_distance, max_distance
+
+    def _reduce_logical_operator_basis(
+        self,
+        candidate_logicals: Union[Sequence, np.ndarray, scipy.sparse.spmatrix] = [],
+    ):
+        """
+        Reduce the logical operator basis to include lower-weight logicals.
+
+        Parameters
+        ----------
+        candidate_logicals : Union[Sequence, np.ndarray, scipy.sparse.spmatrix], optional
+            A list or array of candidate logical operators to be considered for reducing the basis.
+            Defaults to an empty list.
+        """
+        if len(candidate_logicals) != 0:
+            # Convert candidates to a sparse matrix if they aren't already
+            if not isinstance(candidate_logicals, scipy.sparse.spmatrix):
+                candidate_logicals = scipy.sparse.csr_matrix(
+                    scipy.sparse.csr_matrix(candidate_logicals)
+                )
+
+            # Stack the candidate logicals with the existing logicals
+            temp1 = scipy.sparse.vstack(
+                [candidate_logicals, self.logical_operator_basis]
+            ).tocsr()
+
+            # Compute the Hamming weight over GF4 (number of qubits with non-identity operators)
+            # Split into X and Z parts
+            row_weights = binary_pauli_hamming_weight(temp1).flatten()
+
+            # Sort the rows by Hamming weight (ascending)
+            sorted_rows = np.argsort(row_weights)
+            temp1 = temp1[sorted_rows, :]
+
+            # Add the stabilizer matrix to the top of the stack
+            temp1 = scipy.sparse.vstack([self.stabilizer_matrix, temp1])
+
+            # Calculate the rank of the stabilizer matrix (todo: find way of removing this step)
+            stabilizer_rank = ldpc.mod2.rank(self.stabilizer_matrix)
+
+            # Perform row reduction to find a new logical basis
+            p_rows = ldpc.mod2.pivot_rows(temp1)
+            self.logical_operator_basis = temp1[p_rows[stabilizer_rank:]]
+
+    def logical_basis_weights(self):
+        """
+        Return the Hamming weights of the logical operators in the current basis.
+
+        Returns
+        -------
+        np.ndarray
+            An array of integers representing the Hamming weights of the logical operators.
+        """
+        return binary_pauli_hamming_weight(self.logical_operator_basis).flatten()
+
+    def _class_specific_save(self):
+        class_specific_data = {
+            "stabilizers": binary_csr_matrix_to_dict(self.stabilizer_matrix),
+            "logical_operator_basis": binary_csr_matrix_to_dict(
+                self.logical_operator_basis
+            )
+            if self.logical_operator_basis is not None
+            else "?",
+        }
+        return class_specific_data
+
+    def save_code(self, filepath: Union[str, Path], notes: str = "") -> None:
+        """
+        Save the stabilizer code as a JSON file.
+
+        Parameters
+        ----------
+        filepath : Union[str, Path]
+            Path where the generated JSON is saved.
+        notes : str
+            Additional notes to be saved.
+        """
+
+        def repl_function(match):
+            return " ".join(match.group().split())
+
+        filepath = Path(filepath)
+        if not filepath.parent.exists():
+            filepath.parent.mkdir(parents=True)
+
+        general_data = {
+            "notes": notes,
+            "class_name": self.__class__.__name__,
+            "name": self.name,
+            "physical_qubit_count": self.physical_qubit_count,
+            "logical_qubit_count": self.logical_qubit_count,
+            "code_distance": int(self.code_distance)
+            if hasattr(self, "code_distance") and self.code_distance is not None
+            else "?",
+        }
+
+        class_specific_data = self._class_specific_save()
+        merged_data = general_data.copy()
+
+        for key, value in class_specific_data.items():
+            if (
+                key in merged_data
+                and isinstance(merged_data[key], dict)
+                and isinstance(value, dict)
+            ):
+                merged_data[key].update(value)
+            else:
+                merged_data[key] = value
+
+        output = json.dumps(merged_data, indent=4)
+        formatted_output = re.sub(r"(?<=\[)[^\[\]]+(?=\])", repl_function, output)
+
+        with open(filepath, "w") as f:
+            f.write(formatted_output)