Simulating flag qubits and conditional branches using Stim

Question

In a quantum error correcting code using flag qubits, it's common to have flag measurements that tell you it's necessary to do some extra measurements for safety. So, for example, I want to say:

if measurement_result_was_true:
    do_a_different_measurement

How do I do this in Stim? There doesn't seem to be an if.

(This was a question I received by email. Moving it here for posterity.)

Answer 1

This is definitely a place where you'll struggle to use Stim. It's possible but it's not nice. It could still be worth your time just because you don't have to eg. manually write code to turn a circuit into a matching graph, but it'll be more tedious than branchless code.

Part 1: Simulating

The easiest thing to do is to use stim.TableauSimulator. It gives you the ability to apply operations one by one driven by python code, which is more than flexible enough to do what you need.

Here is python code showing the general idea:

import stim

simulator = stim.TableauSimulator()

# Run first part of circuit.
simulator.do(stim.Circuit("""
    H 0
    CNOT 0 1
    M 0
"""))

# Do something depending on the measurement result.
latest_measurement_result = simulator.current_measurement_record()[-1]
if latest_measurement_result:
    simulator.do(stim.Circuit("""
        # ...
    """))

# Run the rest of the circuit.
simulator.do(stim.Circuit("""
    # ...
"""))

shot = simulator.current_measurement_record()

The reason Stim doesn't support branching in its circuit format is because branches break an algorithmic optimization where the tableau simulator is only used once for a reference sample, and then stim switches to using much faster error frame simulation. Consequently, you will find that getting thousands of samples by repeatedly using stim.TableauSimulator is easily 100x slower than using circuit.compile_sampler().sample(shots). Hopefully it's still fast enough for your needs.

If you have less than 10 branches, a possible alternative would be to run simulations for each possible case and then postselect out the inconsistent ones where the hardcoded branch disagreed with the measurement result deciding whether or not it was taken. The performance difference can be large enough to overcome the losses from the postselections.

Part 2: Correcting

I'm assuming you're using PyMatching.

For each shot, you need to build the circuit that was actually run. This circuit will give you access to a detector error model (circuit.detector_error_model(decompose_errors=True)) and the ability to convert the measurement results into detection events (circuit.compile_m2d_converter().convert(shot, append_observables=True)). Convert the measurements into detection events + observable frame changes, and then run pymatching on the detector data, configured with the circuit's detector error model, and check whether or not it predicted the observable frame data.

If you only have a few possible cases, you will get a large benefit from grouping shots that took the same paths together. If you have a lot of possible cases, so that each case is hit less than 3 times on average, grouping isn't worth it.

It's likely that error model extraction will be your main bottleneck here. Stim is pretty quick at this, it takes it a third of a second or so to analyze the error model of a distance 25 surface code, but a third of a second per sample is awful.

I intend for this use case to get better in the future, but haven't got any concrete designs I'm happy with yet.

Answer 2

This might come a bit late, but I am also interested in running stim with branching. To facilitate this, I wrote a piece of code that turns a stim program into an iterator that breaks a program into segment circuits that terminate in a measure statement (or the end of the program), and returns a segment on each iteration. Another helper routine runs the segment and returns the measurement results on the requested qubits. Hope it helps:

from typing import Iterator

import numpy as np
import stim

measure_instructions = ('M', 'MR', 'MX', 'MRX', 'MRY', 'MRZ', 'MX', 'MY', 'MZ')


def to_measure_segments(circuit: stim.Circuit) -> Iterator[stim.Circuit]:
    """
    Create an iterator that iterates over a stim_lib program and returns circuit segments that terminate
    in a measure-like instruction, or the last segment, taking control flow into account.

    This is useful when we want to feedback on measurement results.

    :param circuit: The circuit to transform
    :return: None
    """

    circ_segment = stim.Circuit()
    for inst in circuit:
        if isinstance(inst, stim.CircuitRepeatBlock):
            for _ in range(inst.repeat_count):
                for circ_segment2 in to_measure_segments(inst.body_copy()):
                    circ_segment += circ_segment2
                    if circ_segment[-1].name in measure_instructions:
                        yield circ_segment
                        circ_segment = stim.Circuit()
        elif inst.name not in measure_instructions:
            circ_segment.append_operation(inst)
        else:
            circ_segment.append_operation(inst)
            yield circ_segment
            circ_segment = stim.Circuit()
    if len(circ_segment) > 0:
        yield circ_segment


def do_and_get_measure_results(sim: stim.TableauSimulator,
                               segment: stim.Circuit,
                               qubits_to_return: np.ndarray
                               ) -> np.ndarray:
    sim.do(segment)
    record = np.array(sim.current_measurement_record()[-segment.num_measurements:])
    assert segment[-1].name in measure_instructions, f'bug - segment name is {segment[-1].name}'
    meas_targets = [t.value for t in segment[-1].targets_copy()]
    assert len(meas_targets) == len(record)
    return record[np.isin(meas_targets, qubits_to_return)].astype(np.uint8)


if __name__ == '__main__':

    genc = stim.Circuit.generated('surface_code:rotated_memory_z',
                                  distance=3,
                                  rounds=4,
                                  )
    gen = to_measure_segments(genc)
    while True:
        try:
            print(next(gen))
            print('-' * 30)
        except StopIteration:
            break

```