State tomography

state-tomo.py

import functools
import itertools
from operator import mul

import numpy as np
from numpy.linalg import pinv
from pyquil import get_qc, Program
from pyquil.gates import *
from pyquil.operator_estimation import ExperimentSetting, TomographyExperiment, measure_observables
from pyquil.paulis import sI, sX, sY, sZ
from pyquil.unitary_tools import lifted_pauli

def vec(arr):
    """Flatten a matrix"""
    return np.reshape(arr.T, (-1, 1))

def state_tomo_settings(qubits):
    n_qubits = len(qubits)
    for o_ops in itertools.product([sI, sX, sY, sZ], repeat=n_qubits):
        o_op = functools.reduce(mul, (op(q) for op, q in zip(o_ops, qubits)), sI())
        yield ExperimentSetting(in_operator=sI(), out_operator=o_op)

def linear_inv_state_estimate(results, qubits):
    measurement_matrix = np.vstack([
        vec(lifted_pauli(result.setting.out_operator, qubits=qubits)).T.conj()
        for result in results
    ])
    expectations = np.array([result.expectation for result in results])
    rho = pinv(measurement_matrix) @ expectations
    dim = 2 ** len(qubits)
    rho = rho.reshape((dim, dim)).T
    return rho

if __name__ == '__main__':
    qc = get_qc('Aspen-1-2Q-B')
    qubits = qc.qubits()
    program = Program([H(qubits[0]), CNOT(qubits[0], qubits[1])])
    experiment = TomographyExperiment(settings=list(state_tomo_settings(qubits)),
                                      program=program, qubits=qubits)
    results = list(measure_observables(qc=qc, tomo_experiment=experiment, n_shots=100_000))
    rho = linear_inv_state_estimate(results, qubits=qubits)
    print(np.round(rho, 3))