Adds noise to gates crossing a cut.

add_cut_noise.py

from typing import Optional, Dict

import stim

q2_gates = {
    'CNOT',
    'CX',
    'CY',
    'CZ',
    'ISWAP',
    'ISWAP_DAG',
    'CXSWAP',
    'SWAPCX',
    'SQRT_XX',
    'SQRT_XX_DAG',
    'SQRT_YY',
    'SQRT_YY_DAG',
    'SQRT_ZZ',
    'SQRT_ZZ_DAG',
    'SWAP',
    'XCX',
    'XCY',
    'XCZ',
    'YCX',
    'YCY',
    'YCZ',
    'ZCX',
    'ZCY',
    'ZCZ',
}

def circuit_with_cut_noise(circuit: stim.Circuit, *, p: float, cut: Optional[float] = None, i2q: Optional[Dict[int, complex]] = None) -> stim.Circuit:
    output = stim.Circuit()
    if i2q is None:
        i2q = {q: r + i*1j for q, (r, i) in circuit.get_final_qubit_coordinates().items()}
    if cut is None:
        cut = sum(q.real for q in i2q.values()) / len(i2q)
    for inst in circuit:
        if isinstance(inst, stim.CircuitRepeatBlock):
            output.append(stim.CircuitRepeatBlock(
                repeat_count=inst.repeat_count,
                body=circuit_with_cut_noise(inst.body_copy(), p=p, cut=cut, i2q=i2q),
            ))
        elif isinstance(inst, stim.CircuitInstruction):
            output.append(inst)
            if inst.name in q2_gates:
                ts = inst.targets_copy()
                t2s = []
                for k in range(0, len(ts), 2):
                    a = ts[k]
                    b = ts[k + 1]
                    if a.is_qubit_target and b.is_qubit_target:
                        q1 = i2q[a.value]
                        q2 = i2q[b.value]
                        if (q1.real < cut) != (q2.real < cut):
                            t2s.append(a)
                            t2s.append(b)
                if t2s:
                    output.append("DEPOLARIZE2", t2s, p)
        else:
            raise NotImplementedError(f'{inst=}')
    return output