Distill a T state using on average 22 CS gates

distill_t_using_cs.py

import random

import cirq
import numpy as np


def fresh_attempt() -> int:
    a, b = cirq.LineQubit.range(2)
    h_cx_cs_h = cirq.unitary(cirq.Circuit(cirq.H(b), cirq.X(a).controlled_by(b), cirq.S(a).controlled_by(b), cirq.H(b)))
    check_output_tdag_h = cirq.unitary(cirq.Circuit(cirq.T(a) ** -1, cirq.H(a)))
    initial_state = np.array([np.sqrt(0.5), 0, np.sqrt(0.5), 0], dtype=np.complex128)

    state = np.copy(initial_state)
    q0_given_q1_off = state[0::2]
    q0_given_q1_on = state[1::2]

    step = 0
    while True:
        step += 1

        # Kickback XS onto ancilla (~15% if T, ~85% if ZT).
        state[:] = h_cx_cs_h @ state

        # Perform projective measurement of the ancilla.
        p1 = np.abs(np.sum(q0_given_q1_on * np.conj(q0_given_q1_on)))
        measure_result = random.random() < p1

        if measure_result:
            q0_given_q1_off[:] = q0_given_q1_on[:]
            q0_given_q1_on[:] = 0
            q0_given_q1_off /= np.sqrt(p1)
        else:
            q0_given_q1_on[:] = 0
            q0_given_q1_off /= np.sqrt(1 - p1)

        # Check stopping condition where the actual state differs from the predicted T-or-ZT by less than 1e-12.
        infidelity = abs((check_output_tdag_h @ q0_given_q1_off)[1])**2
        if infidelity > 0.13:
            # We're further than the initial state! Start over!
            state[:] = np.copy(initial_state)
        if infidelity < 1e-12:
            return step

def main():
    total = 0
    n = 10000
    for attempts in range(1, n + 1):
        step = fresh_attempt()
        total += step
        if attempts % 100 == 0:
            print(attempts, "expectation", total / attempts)
    print("expectation", total / n)


if __name__ == '__main__':
    main()