_stim_to_cirq.py

import collections
import functools
from typing import (
    Any,
    Callable,
    cast,
    DefaultDict,
    Dict,
    Iterable,
    List,
    Tuple,
    Union,
)

import cirq
import stim

from ._cx_swap_gate import CXSwapGate
from ._cz_swap_gate import CZSwapGate
from ._det_annotation import DetAnnotation
from ._ii_error_gate import IIErrorGate
from ._i_error_gate import IErrorGate
from ._ii_gate import IIGate
from ._measure_and_or_reset_gate import MeasureAndOrResetGate
from ._obs_annotation import CumulativeObservableAnnotation
from ._shift_coords_annotation import ShiftCoordsAnnotation
from ._sweep_pauli import SweepPauli


def _stim_targets_to_dense_pauli_string(
    targets: List[stim.GateTarget],
) -> cirq.BaseDensePauliString:
    obs = cirq.MutableDensePauliString("I" * len(targets))
    for k, target in enumerate(targets):
        if target.is_inverted_result_target:
            obs *= -1
        if target.is_x_target:
            obs.pauli_mask[k] = 1
        elif target.is_y_target:
            obs.pauli_mask[k] = 2
        elif target.is_z_target:
            obs.pauli_mask[k] = 3
        else:
            raise NotImplementedError(f"target={target!r}")
    return obs.frozen()


def _proper_transform_circuit_qubits(circuit: cirq.AbstractCircuit, remap: Dict[cirq.Qid, cirq.Qid]) -> cirq.Circuit:
    # Note: doing this the hard way because cirq.CircuitOperation otherwise remembers the old indices in
    # its `remap` entry, instead of completely expunging those indices.
    return cirq.Circuit(
        cirq.Moment(
            cirq.CircuitOperation(
                circuit=_proper_transform_circuit_qubits(op.circuit, remap).freeze(),
                repetitions=op.repetitions,
                use_repetition_ids=False,
            )
            if isinstance(op, cirq.CircuitOperation)
            else op.with_qubits(*[remap[q] for q in op.qubits])
            for op in moment
        )
        for moment in circuit
    )


class CircuitTranslationTracker:
    def __init__(self, flatten: bool):
        self.qubit_coords: Dict[int, cirq.Qid] = {}
        self.origin: DefaultDict[float] = collections.defaultdict(float)
        self.num_measurements_seen = 0
        self.full_circuit = cirq.Circuit()
        self.tick_circuit = cirq.Circuit()
        self.flatten = flatten
        self.have_seen_loop = False

    def get_next_measure_id(self) -> int:
        self.num_measurements_seen += 1
        return self.num_measurements_seen - 1

    def append_operation(self, op: cirq.Operation) -> None:
        self.tick_circuit.append(op, strategy=cirq.InsertStrategy.INLINE)

    def process_gate_instruction(
        self, gate: cirq.Gate, instruction: stim.CircuitInstruction
    ) -> None:
        targets: List[stim.GateTarget] = instruction.targets_copy()
        m = cirq.num_qubits(gate)
        if not all(t.is_qubit_target for t in targets) or len(targets) % m != 0:
            raise NotImplementedError(f"instruction={instruction!r}")
        if instruction.tag:
            tags = [instruction.tag]
        else:
            tags = ()
        for k in range(0, len(targets), m):
            self.append_operation(gate(*[cirq.LineQubit(t.value) for t in targets[k : k + m]]).with_tags(*tags))

    def process_tick(self, instruction: stim.CircuitInstruction) -> None:
        self.full_circuit += self.tick_circuit or cirq.Moment()
        self.tick_circuit = cirq.Circuit()

    def process_pauli_channel_1(self, instruction: stim.CircuitInstruction) -> None:
        args = instruction.gate_args_copy()
        if len(args) != 3:
            raise ValueError(f"len(args={args!r}) != 3")
        self.process_gate_instruction(
            cirq.AsymmetricDepolarizingChannel(p_x=args[0], p_y=args[1], p_z=args[2]), instruction
        )

    def process_pauli_channel_2(self, instruction: stim.CircuitInstruction) -> None:
        args = instruction.gate_args_copy()
        if len(args) != 15:
            raise ValueError(f"len(args={args!r}) != 15")

        ps = {
            'IX': args[0],
            'IY': args[1],
            'IZ': args[2],
            'XI': args[3],
            'XX': args[4],
            'XY': args[5],
            'XZ': args[6],
            'YI': args[7],
            'YX': args[8],
            'YY': args[9],
            'YZ': args[10],
            'ZI': args[11],
            'ZX': args[12],
            'ZY': args[13],
            'ZZ': args[14],
        }
        ps = {k: v for k, v in ps.items() if v}
        if not ps:
            ps['II'] = 1
        gate = cirq.asymmetric_depolarize(error_probabilities=ps)
        self.process_gate_instruction(gate, instruction)

    def process_repeat_block(self, block: stim.CircuitRepeatBlock):
        if self.flatten or (block.repeat_count == 1 and block.tag == ""):
            self.process_circuit(block.repeat_count, block.body_copy())
            return

        self.have_seen_loop = True
        child = CircuitTranslationTracker(flatten=self.flatten)
        child.origin = self.origin.copy()
        child.num_measurements_seen = self.num_measurements_seen
        child.qubit_coords = self.qubit_coords.copy()
        child.have_seen_loop = True
        child.process_circuit(1, block.body_copy())

        # Circuit operation will always be in their own cirq.Moment
        if block.tag == "":
            tags = ()
        else:
            tags = (block.tag,)
        if len(self.tick_circuit):
            self.full_circuit += self.tick_circuit
        self.full_circuit += cirq.Moment(
            cirq.CircuitOperation(
                cirq.FrozenCircuit(child.full_circuit + child.tick_circuit),
                repetitions=block.repeat_count,
                use_repetition_ids=False,
            ).with_tags(*tags)
        )
        self.tick_circuit = cirq.Circuit()

        self.qubit_coords = child.qubit_coords
        self.num_measurements_seen += (
            child.num_measurements_seen - self.num_measurements_seen
        ) * block.repeat_count
        for k, v in child.origin.items():
            self.origin[k] += (v - self.origin[k]) * block.repeat_count

    def process_measurement_instruction(
        self, instruction: stim.CircuitInstruction, measure: bool, reset: bool, basis: str
    ) -> None:
        args = instruction.gate_args_copy()
        flip_probability = 0
        if args:
            flip_probability = args[0]

        targets: List[stim.GateTarget] = instruction.targets_copy()
        if instruction.tag:
            tags = [instruction.tag]
        else:
            tags = ()
        for t in targets:
            if not t.is_qubit_target:
                raise NotImplementedError(f"instruction={instruction!r}")
            key = str(self.get_next_measure_id())
            self.append_operation(
                MeasureAndOrResetGate(
                    measure=measure,
                    reset=reset,
                    basis=basis,
                    invert_measure=t.is_inverted_result_target,
                    key=key,
                    measure_flip_probability=flip_probability,
                ).resolve(cirq.LineQubit(t.value)).with_tags(*tags)
            )

    def process_circuit(self, repetitions: int, circuit: stim.Circuit) -> None:
        handler_table = CircuitTranslationTracker.get_handler_table()
        for _ in range(repetitions):
            for instruction in circuit:
                if isinstance(instruction, stim.CircuitInstruction):
                    handler = handler_table.get(instruction.name)
                    if handler is None:
                        raise NotImplementedError(f"{instruction!r}")
                    handler(self, instruction)
                elif isinstance(instruction, stim.CircuitRepeatBlock):
                    self.process_repeat_block(instruction)
                else:
                    raise NotImplementedError(f"instruction={instruction!r}")

    def output(self) -> cirq.Circuit:
        out = self.full_circuit + self.tick_circuit

        if self.qubit_coords:
            remap: Dict[cirq.Qid, cirq.Qid] = {
                q: self.qubit_coords.get(cast(cirq.LineQubit, q).x, q) for q in out.all_qubits()
            }

            # Only remap if there are no collisions.
            if len(set(remap.values())) == len(remap):
                # Note: doing this the hard way because cirq.CircuitOperation otherwise remembers the old indices in
                # its `remap` entry, instead of completely expunging those indices.
                out = _proper_transform_circuit_qubits(out, remap)

        return out

    def process_mpp(self, instruction: stim.CircuitInstruction) -> None:
        args = instruction.gate_args_copy()
        if args and args[0]:
            raise NotImplementedError("Noisy MPP")

        targets: List[stim.GateTarget] = instruction.targets_copy()
        if instruction.tag:
            tags = [instruction.tag]
        else:
            tags = ()
        start = 0
        while start < len(targets):
            next_start = start + 1
            while next_start < len(targets) and targets[next_start].is_combiner:
                next_start += 2
            group = targets[start:next_start:2]
            start = next_start

            obs = _stim_targets_to_dense_pauli_string(group)
            qubits = [cirq.LineQubit(t.value) for t in group]
            key = str(self.get_next_measure_id())
            self.append_operation(cirq.PauliMeasurementGate(obs, key=key).on(*qubits).with_tags(*tags))

    def process_spp_dag(self, instruction: stim.CircuitInstruction) -> None:
        self.process_spp(instruction, dag=True)

    def process_spp(self, instruction: stim.CircuitInstruction, dag: bool = False) -> None:
        targets: List[stim.GateTarget] = instruction.targets_copy()
        if instruction.tag:
            tags = [instruction.tag]
        else:
            tags = ()
        start = 0
        while start < len(targets):
            next_start = start + 1
            while next_start < len(targets) and targets[next_start].is_combiner:
                next_start += 2
            group = targets[start:next_start:2]
            start = next_start

            obs = _stim_targets_to_dense_pauli_string(group)
            qubits = [cirq.LineQubit(t.value) for t in group]
            self.append_operation(cirq.PauliStringPhasorGate(
                obs,
                exponent_neg=-0.5 if dag else 0.5,
            ).on(*qubits).with_tags(*tags))

    def process_m_pair(self, instruction: stim.CircuitInstruction, basis: str) -> None:
        args = instruction.gate_args_copy()
        if args and args[0]:
            raise NotImplementedError("Noisy M" + basis*2)

        if instruction.tag:
            tags = [instruction.tag]
        else:
            tags = ()
        targets: List[stim.GateTarget] = instruction.targets_copy()
        for k in range(0, len(targets), 2):
            obs = cirq.DensePauliString(basis * 2)
            if targets[0].is_inverted_result_target ^ targets[1].is_inverted_result_target:
                obs *= -1
            qubits = [cirq.LineQubit(targets[0].value), cirq.LineQubit(targets[1].value)]
            key = str(self.get_next_measure_id())
            self.append_operation(cirq.PauliMeasurementGate(obs, key=key).on(*qubits).with_tags(*tags))

    def process_mxx(self, instruction: stim.CircuitInstruction) -> None:
        self.process_m_pair(instruction, "X")

    def process_myy(self, instruction: stim.CircuitInstruction) -> None:
        self.process_m_pair(instruction, "Y")

    def process_mzz(self, instruction: stim.CircuitInstruction) -> None:
        self.process_m_pair(instruction, "Z")

    def process_mpad(self, instruction: stim.CircuitInstruction) -> None:
        targets: List[stim.GateTarget] = instruction.targets_copy()
        for t in targets:
            obs = cirq.DensePauliString("")
            if t.value == 1:
                obs *= -1
            qubits = []
            key = str(self.get_next_measure_id())
            self.append_operation(cirq.PauliMeasurementGate(obs, key=key).on(*qubits))

    def process_correlated_error(self, instruction: stim.CircuitInstruction) -> None:
        if instruction.tag:
            tags = [instruction.tag]
        else:
            tags = ()
        args = instruction.gate_args_copy()
        probability = args[0] if args else 0
        targets = instruction.targets_copy()
        qubits = [cirq.LineQubit(t.value) for t in targets]
        self.append_operation(
            _stim_targets_to_dense_pauli_string(targets).on(*qubits).with_probability(probability).with_tags(*tags)
        )

    def coords_after_offset(
        self, relative_coords: List[float], *, even_if_flattening: bool = False
    ) -> List[Union[float, int]]:
        if not self.flatten and not even_if_flattening:
            return list(relative_coords)
        result = []
        for k in range(len(relative_coords)):
            t = relative_coords[k] + self.origin[k]
            if t == int(t):
                t = int(t)
            result.append(t)
        return result

    def resolve_measurement_record_keys(
        self, targets: Iterable[stim.GateTarget]
    ) -> Tuple[List[str], List[int], List[tuple[cirq.Qid, str]]]:
        pauli_targets, meas_targets = [], []
        for t in targets:
            if t.is_measurement_record_target:
                meas_targets.append(t)
            else:
                pauli_targets.append((cirq.LineQubit(t.value), t.pauli_type))

        if self.have_seen_loop:
            return [], [t.value for t in meas_targets], pauli_targets
        else:
            return [str(self.num_measurements_seen + t.value) for t in meas_targets], [], pauli_targets

    def process_detector(self, instruction: stim.CircuitInstruction) -> None:
        if instruction.tag:
            tags = [instruction.tag]
        else:
            tags = ()
        coords = self.coords_after_offset(instruction.gate_args_copy())
        keys, rels, _ = self.resolve_measurement_record_keys(instruction.targets_copy())
        self.append_operation(
            DetAnnotation(parity_keys=keys, relative_keys=rels, coordinate_metadata=coords).with_tags(*tags)
        )

    def process_observable_include(self, instruction: stim.CircuitInstruction) -> None:
        if instruction.tag:
            tags = [instruction.tag]
        else:
            tags = ()
        args = instruction.gate_args_copy()
        index = 0 if not args else int(args[0])
        keys, rels, paulis = self.resolve_measurement_record_keys(instruction.targets_copy())
        self.append_operation(
            CumulativeObservableAnnotation(
                parity_keys=keys, relative_keys=rels, pauli_keys=paulis, observable_index=index
            ).with_tags(*tags)
        )

    def process_qubit_coords(self, instruction: stim.CircuitInstruction) -> None:
        coords = self.coords_after_offset(instruction.gate_args_copy(), even_if_flattening=True)
        for t in instruction.targets_copy():
            if len(coords) == 1:
                self.qubit_coords[t.value] = cirq.LineQubit(*coords)
            elif len(coords) == 2:
                self.qubit_coords[t.value] = cirq.GridQubit(*coords)

    def process_shift_coords(self, instruction: stim.CircuitInstruction) -> None:
        if instruction.tag:
            tags = [instruction.tag]
        else:
            tags = ()
        args = instruction.gate_args_copy()
        if not self.flatten:
            self.append_operation(ShiftCoordsAnnotation(args).with_tags(*tags))
        for k, a in enumerate(args):
            self.origin[k] += a

    class OneToOneGateHandler:
        def __init__(self, gate: cirq.Gate):
            self.gate = gate

        def __call__(
            self, tracker: 'CircuitTranslationTracker', instruction: stim.CircuitInstruction
        ) -> None:
            tracker.process_gate_instruction(gate=self.gate, instruction=instruction)

    class SweepableGateHandler:
        def __init__(self, pauli_gate: cirq.Pauli, gate: cirq.Gate):
            self.pauli_gate = pauli_gate
            self.gate = gate

        def __call__(
            self, tracker: 'CircuitTranslationTracker', instruction: stim.CircuitInstruction
        ) -> None:
            if instruction.tag:
                tags = [instruction.tag]
            else:
                tags = ()
            targets: List[stim.GateTarget] = instruction.targets_copy()
            for k in range(0, len(targets), 2):
                a = targets[k]
                b = targets[k + 1]
                if not a.is_qubit_target and not b.is_qubit_target:
                    raise NotImplementedError(f"instruction={instruction!r}")
                if a.is_sweep_bit_target or b.is_sweep_bit_target:
                    if b.is_sweep_bit_target:
                        a, b = b, a
                    assert not a.is_inverted_result_target
                    tracker.append_operation(
                        SweepPauli(
                            stim_sweep_bit_index=a.value,
                            cirq_sweep_symbol=f'sweep[{a.value}]',
                            pauli=self.pauli_gate,
                        ).on(cirq.LineQubit(b.value)).with_tags(*tags)
                    )
                else:
                    if not a.is_qubit_target or not b.is_qubit_target:
                        raise NotImplementedError(f"instruction={instruction!r}")
                    tracker.append_operation(
                        self.gate(cirq.LineQubit(a.value), cirq.LineQubit(b.value)).with_tags(*tags)
                    )

    class OneToOneMeasurementHandler:
        def __init__(self, *, reset: bool, measure: bool, basis: str):
            self.reset = reset
            self.measure = measure
            self.basis = basis

        def __call__(
            self, tracker: 'CircuitTranslationTracker', instruction: stim.CircuitInstruction
        ) -> None:
            tracker.process_measurement_instruction(
                measure=self.measure, reset=self.reset, basis=self.basis, instruction=instruction
            )

    class OneToOneNoisyGateHandler:
        def __init__(self, prob_to_gate: Callable[[float], cirq.Gate]):
            self.prob_to_gate = prob_to_gate

        def __call__(
            self, tracker: 'CircuitTranslationTracker', instruction: stim.CircuitInstruction
        ) -> None:
            tracker.process_gate_instruction(
                self.prob_to_gate(instruction.gate_args_copy()[0]), instruction
            )

    class MultiArgumentGateHandler:
        def __init__(self, args_to_gate: Callable[[List[float]], cirq.Gate]):
            self.args_to_gate = args_to_gate

        def __call__(
            self, tracker: 'CircuitTranslationTracker', instruction: stim.CircuitInstruction
        ) -> None:
            tracker.process_gate_instruction(
                self.args_to_gate(instruction.gate_args_copy()), instruction
            )

    @staticmethod
    @functools.lru_cache(maxsize=1)
    def get_handler_table() -> Dict[
        str, Callable[['CircuitTranslationTracker', stim.Circuit], None]
    ]:
        gate = CircuitTranslationTracker.OneToOneGateHandler
        measure_gate = CircuitTranslationTracker.OneToOneMeasurementHandler
        noise = CircuitTranslationTracker.OneToOneNoisyGateHandler
        multi_arg = CircuitTranslationTracker.MultiArgumentGateHandler
        sweep_gate = CircuitTranslationTracker.SweepableGateHandler

        def not_impl(message) -> Callable[[Any, Any], None]:
            def handler(
                tracker: CircuitTranslationTracker, instruction: stim.CircuitInstruction
            ) -> None:
                raise NotImplementedError(message)

            return handler

        return {
            "M": measure_gate(measure=True, reset=False, basis='Z'),
            "MX": measure_gate(measure=True, reset=False, basis='X'),
            "MY": measure_gate(measure=True, reset=False, basis='Y'),
            "MZ": measure_gate(measure=True, reset=False, basis='Z'),
            "MR": measure_gate(measure=True, reset=True, basis='Z'),
            "MRX": measure_gate(measure=True, reset=True, basis='X'),
            "MRY": measure_gate(measure=True, reset=True, basis='Y'),
            "MRZ": measure_gate(measure=True, reset=True, basis='Z'),
            "R": gate(cirq.ResetChannel()),
            "RX": gate(
                MeasureAndOrResetGate(
                    measure=False, reset=True, basis='X', invert_measure=False, key=''
                )
            ),
            "CZSWAP": gate(CZSwapGate()),
            "CXSWAP": gate(CXSwapGate(inverted=False)),
            "SWAPCX": gate(CXSwapGate(inverted=True)),
            "RY": gate(
                MeasureAndOrResetGate(
                    measure=False, reset=True, basis='Y', invert_measure=False, key=''
                )
            ),
            "RZ": gate(cirq.ResetChannel()),
            "I": gate(cirq.I),
            "II": gate(IIGate()),
            "X": gate(cirq.X),
            "Y": gate(cirq.Y),
            "Z": gate(cirq.Z),
            "H_XY": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Y, False), z_to=(cirq.Z, True))
            ),
            "H": gate(cirq.H),
            "H_XZ": gate(cirq.H),
            "H_NXZ": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Z, True), z_to=(cirq.X, True))
            ),
            "H_YZ": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.X, True), z_to=(cirq.Y, False))
            ),
            "H_NXY": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Y, True), z_to=(cirq.Z, True))
            ),
            "H_NYZ": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.X, True), z_to=(cirq.Y, True))
            ),
            "SQRT_X": gate(cirq.X ** 0.5),
            "SQRT_X_DAG": gate(cirq.X ** -0.5),
            "SQRT_Y": gate(cirq.Y ** 0.5),
            "SQRT_Y_DAG": gate(cirq.Y ** -0.5),
            "C_XYZ": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Y, False), z_to=(cirq.X, False))
            ),
            "C_NXYZ": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Y, True), z_to=(cirq.X, True))
            ),
            "C_XNYZ": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Y, True), z_to=(cirq.X, False))
            ),
            "C_XYNZ": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Y, False), z_to=(cirq.X, True))
            ),
            "C_ZYX": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Z, False), z_to=(cirq.Y, False))
            ),
            "C_NZYX": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Z, True), z_to=(cirq.Y, True))
            ),
            "C_ZNYX": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Z, False), z_to=(cirq.Y, True))
            ),
            "C_ZYNX": gate(
                cirq.SingleQubitCliffordGate.from_xz_map(x_to=(cirq.Z, True), z_to=(cirq.Y, False))
            ),
            "SQRT_XX": gate(cirq.XX ** 0.5),
            "SQRT_YY": gate(cirq.YY ** 0.5),
            "SQRT_ZZ": gate(cirq.ZZ ** 0.5),
            "SQRT_XX_DAG": gate(cirq.XX ** -0.5),
            "SQRT_YY_DAG": gate(cirq.YY ** -0.5),
            "SQRT_ZZ_DAG": gate(cirq.ZZ ** -0.5),
            "S": gate(cirq.S),
            "S_DAG": gate(cirq.S ** -1),
            "SQRT_Z": gate(cirq.S),
            "SQRT_Z_DAG": gate(cirq.S ** -1),
            "SWAP": gate(cirq.SWAP),
            "ISWAP": gate(cirq.ISWAP),
            "ISWAP_DAG": gate(cirq.ISWAP ** -1),
            "XCX": gate(cirq.PauliInteractionGate(cirq.X, False, cirq.X, False)),
            "XCY": gate(cirq.PauliInteractionGate(cirq.X, False, cirq.Y, False)),
            "XCZ": sweep_gate(cirq.X, cirq.PauliInteractionGate(cirq.X, False, cirq.Z, False)),
            "YCX": gate(cirq.PauliInteractionGate(cirq.Y, False, cirq.X, False)),
            "YCY": gate(cirq.PauliInteractionGate(cirq.Y, False, cirq.Y, False)),
            "YCZ": sweep_gate(cirq.Y, cirq.PauliInteractionGate(cirq.Y, False, cirq.Z, False)),
            "CX": sweep_gate(cirq.X, cirq.CNOT),
            "CNOT": sweep_gate(cirq.X, cirq.CNOT),
            "ZCX": sweep_gate(cirq.X, cirq.CNOT),
            "CY": sweep_gate(cirq.Y, cirq.Y.controlled(1)),
            "ZCY": sweep_gate(cirq.Y, cirq.Y.controlled(1)),
            "CZ": sweep_gate(cirq.Z, cirq.CZ),
            "ZCZ": sweep_gate(cirq.Z, cirq.CZ),
            "DEPOLARIZE1": noise(lambda p: cirq.DepolarizingChannel(p, 1)),
            "DEPOLARIZE2": noise(lambda p: cirq.DepolarizingChannel(p, 2)),
            "X_ERROR": noise(cirq.X.with_probability),
            "Y_ERROR": noise(cirq.Y.with_probability),
            "Z_ERROR": noise(cirq.Z.with_probability),
            "I_ERROR": multi_arg(IErrorGate),
            "II_ERROR": multi_arg(IIErrorGate),
            "PAULI_CHANNEL_1": CircuitTranslationTracker.process_pauli_channel_1,
            "PAULI_CHANNEL_2": CircuitTranslationTracker.process_pauli_channel_2,
            "ELSE_CORRELATED_ERROR": not_impl(
                "Converting ELSE_CORRELATED_ERROR to cirq is not supported."
            ),
            "REPEAT": not_impl("[handled special]"),
            "TICK": CircuitTranslationTracker.process_tick,
            "SHIFT_COORDS": CircuitTranslationTracker.process_shift_coords,
            "QUBIT_COORDS": CircuitTranslationTracker.process_qubit_coords,
            "E": CircuitTranslationTracker.process_correlated_error,
            "CORRELATED_ERROR": CircuitTranslationTracker.process_correlated_error,
            "MPP": CircuitTranslationTracker.process_mpp,
            "SPP": CircuitTranslationTracker.process_spp,
            "SPP_DAG": CircuitTranslationTracker.process_spp_dag,
            "MXX": CircuitTranslationTracker.process_mxx,
            "MYY": CircuitTranslationTracker.process_myy,
            "MZZ": CircuitTranslationTracker.process_mzz,
            "MPAD": CircuitTranslationTracker.process_mpad,
            "DETECTOR": CircuitTranslationTracker.process_detector,
            "OBSERVABLE_INCLUDE": CircuitTranslationTracker.process_observable_include,
            "HERALDED_ERASE": not_impl(
                "Converting HERALDED_ERASE to cirq is not supported."
            ),
            "HERALDED_PAULI_CHANNEL_1": not_impl(
                "Converting HERALDED_PAULI_CHANNEL_1 to cirq is not supported."
            ),
        }


def stim_circuit_to_cirq_circuit(circuit: stim.Circuit, *, flatten: bool = False) -> cirq.Circuit:
    """Converts a stim circuit into an equivalent cirq circuit.

    Qubit indices are turned into cirq.LineQubit instances. Measurements are
    keyed by their ordering (e.g. the first measurement is keyed "0", the second
    is keyed "1", etc).

    Not all circuits can be converted:
        - ELSE_CORRELATED_ERROR instructions are not supported.

    Not all circuits can be converted with perfect 1:1 fidelity:
        - DETECTOR annotations are discarded.
        - OBSERVABLE_INCLUDE annotations are discarded.

    Args:
        circuit: The stim circuit to convert into a cirq circuit.
        flatten: Defaults to False. When set to True, REPEAT blocks are removed by
            explicitly repeating their instructions multiple times. Also,
            SHIFT_COORDS instructions are removed by appropriately adjusting the
            coordinate metadata of later instructions.

    Returns:
        The converted circuit.

    Examples:

        >>> import stimcirq
        >>> import stim
        >>> print(stimcirq.stim_circuit_to_cirq_circuit(stim.Circuit('''
        ...     H 0
        ...     CNOT 0 1
        ...     X_ERROR(0.25) 0
        ...     TICK
        ...     M !1 0
        ... ''')))
        0: ───H───@───X[prob=0.25]───M('1')────
                  │
        1: ───────X──────────────────!M('0')───
    """
    tracker = CircuitTranslationTracker(flatten=flatten)
    tracker.process_circuit(repetitions=1, circuit=circuit)
    return tracker.output()