Surface code memory experiment

try:
    import cirq
    assert hasattr(cirq.transformers, 'apply_lazy_args_on_circuit_operation')
except:
    !pip install --upgrade cirq~=1.0.dev

try:
    import recirq.qec_memory.decoding as qec_decoding
    import wget
except:
    !pip install "recirq[qec_memory] @ git+https://github.com/quantumlib/ReCirq.git"

import copy
import time
import zipfile
from typing import Literal

import cirq_google
import matplotlib.pyplot as plt
import numpy as np
import stim
import wget
from google.colab import auth
from IPython.display import display

import recirq.qec_memory.analysis as qec_analysis
import recirq.qec_memory.circuits as qec_circuits
import recirq.qec_memory.decoding as qec_decoding

auth.authenticate_user(clear_output=False)

# download the circuits
wget.download('https://zenodo.org/records/23004813/files/ZXXZ circuits.zip')

# unzip them:
with zipfile.ZipFile('ZXXZ circuits.zip', 'r') as zip_ref:
    zip_ref.extractall()

config_name = 'd5' # choose between d5 or d7

if config_name == 'd5':
    processor_id = '' # fill this in
    distance_to_shifts = {
        3: [(0,0), (0,4), (2,2), (-2,2)],
        5: [(0,0)]
    }
    directory = 'd3_d5/'
elif config_name == 'd7':
    processor_id = '' # fill this in
    distance_to_shifts = {
        3: [(0, 0),
        (-2, 2),
        (-4, 4),
        (2, 2),
        (0, 4),
        (-2, 6),
        (4, 4),
        (2, 6),
        (0, 8)],
        5: [(0, 0), (-2, 2), (2, 2), (0, 4)],
        7: [(0, 0)]
        }
    directory = 'd3_d5_d7/'

def get_zxxz_circuit(
    distance: int, observable: Literal["H", "V"], shift: tuple[int, int]
) -> stim.Circuit:
    """Load a XZZX surface code circuit from a file.

    Args:
        distance: The code distance.
        observable: Which basis to measure in.
        shift: Determines which qubits are used.

    Returns:
        A stim circuit for the memory experiment.
    """
    circuit = stim.Circuit.from_file(
        directory
        + f"distance_{distance}_cycles_10_observable_{observable}_shift_{shift}.stim"
    )
    # remove the sweep bits and subsequent TICK
    nq = 2 * distance**2 - 1
    return circuit[:nq] + circuit[(nq + 2) :]

circuit = get_zxxz_circuit(3, 'H', (0,0))
print(circuit.diagram())

print(circuit.to_crumble_url())

circuit = get_zxxz_circuit(3, 'V', (0,0))
print(circuit.diagram())

engine = cirq_google.get_engine("")  ## fill in your GCP project
sampler = engine.get_sampler(
    processor_id, circuits_per_job=200, device_config_name=config_name
)

qec_recipe = [
    "ADEPT",
    "ECHO",
    "DQLR_MSMLR",
    "DD",
    "RESET_AFTER_MEAS",
    "READOUT_RESET_TAG",
]
num_sweep_bit_choices = 10
repetitions = 1000
cycles_list = np.arange(10, 251, 20)

# Repeat the cycle numbers num_sweep_bit_choices times.
# For each one we will add random sweep bits.
cycles_times_sweep_bits = np.repeat(cycles_list, num_sweep_bit_choices)

# shuffle the order of cycle numbers
rng = np.random.default_rng()
rng.shuffle(cycles_times_sweep_bits)

# set up some variables that we will use to keep track of the iterations
bases = ["H", "V"]
total_param_choices = len(bases) * sum(
    len(shifts) for shifts in distance_to_shifts.values()
)
done_params = set()

# we will store the results in here:
lambda_results = qec_analysis.LambdaExperimentResults(
    cycles_list, repetitions, num_sweep_bit_choices
)

# and plot them in here:
fig, ax = plt.subplots(dpi=150, facecolor="white")

# we will print out how long it is taking:
start_time = time.time()
time_string = (
    lambda: f"{(time.time()-start_time)//60:.0f}:{(time.time()-start_time)%60:02.0f}"
)

# start data taking
while len(done_params) < total_param_choices:
    # first choose the code distance
    distance = rng.choice(list(distance_to_shifts.keys()))

    # the following if block helps us get a lambda estimate faster but can be
    # deleted to take data in a random order
    if len(done_params) == 1:
        if next(iter(done_params)).distance == 3:
            distance = 5
        elif next(iter(done_params)).distance == 5:
            distance = 3
        elif next(iter(done_params)).distance == 7:
            distance = 5

    # next choose from the allowed placements on the device (shifts)
    shift_idx = rng.integers(0, len(distance_to_shifts[distance]))
    shift = distance_to_shifts[distance][shift_idx]

    # next choose the observable
    observable = rng.choice(bases)

    # skip if we have already taken data for these params
    params = qec_analysis.SurfaceCodeParams(distance, observable, shift)
    if params in done_params:
        continue
    print(f"{time_string()}: Starting d={distance}, shift={shift}, basis={observable}")

    # get the Zenodo stim circuit
    stim_circuit = get_zxxz_circuit(distance, observable, shift)

    # add the hardware techniques
    print(f"{time_string()}: Compiling hardware circuit")
    hardware_circuit = engine.compile_circuit(
        stim_circuit,
        qec_recipe,
        processor_id=processor_id,
        config_name=config_name,
    )

    # calibrate the ADEPT phases. (Refer to the paper, arXiv:1603.03082.)
    print(f"{time_string()}: Calibrating ADEPT")
    resolver = engine.calibrate_for_circuit(
        hardware_circuit, processor_id, config_name=config_name
    )

    # Generate the circuits for the various cycle numbers and sweep bits.
    # Uncomment the following line to use a different random order of cycles for
    # each choice of params; otherwise the same order is reused.
    # rng.shuffle(cycles_times_sweep_bits) # uncomment to
    print(f"Order of cycles: {cycles_times_sweep_bits}")
    circuits = [
        qec_circuits.add_sweep_bits(
            qec_circuits.replace_loop_repetitions(
                hardware_circuit, cycles, original_cycles=10
            ),
            rng,
        )
        for cycles in cycles_times_sweep_bits
    ]

    # Take the data.
    print(f"{time_string()}: Taking data")
    result = sampler.run_batch(
        circuits,
        repetitions=repetitions,
        params_list=[resolver] * len(cycles_times_sweep_bits),
    )

    # Extract the logical error probability.
    print(f"{time_string()}: Analyzing data")
    lep = np.array(
        [
            qec_decoding.get_logical_error_probability(res_i[0].data, circuit, 0.001)
            for circuit, res_i in zip(circuits, result)
        ]
    )

    # Fit the logical error rate and plot the results. Extract Λ.
    lambda_results.add_result(params, lep, copy.deepcopy(cycles_times_sweep_bits))
    ax.clear()
    ax = lambda_results.plot(ax)
    display(fig)
    done_params.add(params)