Adding custom anticommuting observables to chunks in stimflow fails verification

[kaavyas99]

Goal: to get stimflow to recognize and track observables that anticommute with resets and measurements, similar to recent functionality added in Stim 1.15

Issues: Direct implementation is not supported. Via creation of a custom chunk, we attempted to create a function that mimics the above. However, an initialization chunk (code sample below) does not successfully compile the Z_wrong observable added after RX resets on qubits. Note that if the new observable is changed to an X observable, verification succeeds.

import stimflow as sf

def make_surface_code(diameter: int) -> sf.StabilizerCode:
    #helper fn to make surface code of distance = diameter, from getting started nb
    tiles = []

    for x in range(-1, diameter):
        for y in range(-1, diameter):
            m = x + 1j * y + 0.5 + 0.5j
            potential_data = [m + 1j**k * (0.5 + 0.5j) for k in range(4)]
            data = [d for d in potential_data if 0 <= d.real < diameter if 0 <= d.imag < diameter]
            if len(data) not in [2, 4]:
                continue

            basis = "XZ"[(x.real + y.real) % 2 == 0]
            if not (0 <= m.real < diameter - 1) and basis != "Z":
                continue
            if not (0 <= m.imag < diameter - 1) and basis != "X":
                continue
            tiles.append(sf.Tile(measure_qubit=m, data_qubits=data, bases=basis))

    return sf.Patch(tiles)


def create_anti_obs(patch: sf.Patch, obsstr: sf.PauliMap):
    #custom fn to make obs that anticommutes with reset
    chunk = sf.ChunkBuilder()

    chunk.append("RX" , patch.data_set)
    chunk.append("OBSERVABLE_INCLUDE", obsstr, arg=0)

    for tile in patch:
            if tile.basis == 'X':
                chunk.add_flow(end=tile)
            else:
                chunk.add_discarded_flow_output(tile)

    chunk.add_flow(end=obsstr)
    hfchunk = chunk.finish_chunk()
    return sf.Chunk(
        circuit = hfchunk.circuit,
        flows = hfchunk.flows,
        discarded_inputs = hfchunk.discarded_inputs,
        discarded_outputs = hfchunk.discarded_outputs,
        q2i = hfchunk.q2i,
        o2i = {obsstr: 0}
    )


#driver fxns
sc5patch = make_surface_code(5) 
vcirc = create_anti_obs(sc5patch, sf.PauliMap({i:'Z' for i in range(5)}, obs_name='Z_wrong')) 
vcirc.verify() #verification fails

[Strilanc]

The issue here was that we were passing the observable (instead of its name) as the key of the o2i value. I'll add a safety check to raise an error when this occurs. Also I clarified the documentation around automatic indexing of observables: simply deleting the , arg=0 and replacing hfchunk = chunk.finish_chunk() with return chunk.finish_chunk() would have done the correct thing.