Quantum State Tomography¶

Adapted from QGopt/docs/state-tomography

In [1]:

import numpy as np
import collections
import torch
import scipy.optimize
import matplotlib.pyplot as plt

try:
    import numqi
except ImportError:
    %pip install numqi
    import numqi

np_rng = np.random.default_rng()

hf_trace_distance = lambda x,y: np.abs(np.linalg.eigvalsh(x-y)).sum()
# https://en.wikipedia.org/wiki/Trace_distance

import numpy as np import collections import torch import scipy.optimize import matplotlib.pyplot as plt try: import numqi except ImportError: %pip install numqi import numqi np_rng = np.random.default_rng() hf_trace_distance = lambda x,y: np.abs(np.linalg.eigvalsh(x-y)).sum() # https://en.wikipedia.org/wiki/Trace_distance

see wiki-link for tetrahedron POVM

In [2]:

num_qubit = 2
povm_list = numqi.utils.get_tetrahedron_POVM(num_qubit)

target_rho = numqi.random.rand_density_matrix(2**num_qubit)

real_prob = np.trace(povm_list @ target_rho, axis1=1, axis2=2).real
print(f'real probability:\n{np.around(real_prob,3).tolist()}')

num_qubit = 2 povm_list = numqi.utils.get_tetrahedron_POVM(num_qubit) target_rho = numqi.random.rand_density_matrix(2**num_qubit) real_prob = np.trace(povm_list @ target_rho, axis1=1, axis2=2).real print(f'real probability:\n{np.around(real_prob,3).tolist()}')

real probability:
[0.067, 0.032, 0.029, 0.055, 0.06, 0.061, 0.092, 0.031, 0.073, 0.088, 0.077, 0.088, 0.034, 0.044, 0.106, 0.062]

In [3]:

# Gumbel trick in the original tutorial is interesting, but it's not used here for simplicity
# https://qgopt.readthedocs.io/en/latest/state_tomography.html
num_measurement = 600000
index_measure = np_rng.choice(len(real_prob), size=num_measurement, p=real_prob)
tmp0 = collections.Counter(index_measure.tolist())
measure_prob = np.array([tmp0[i] for i in range(len(real_prob))])/num_measurement
print(f'measured probability:\n{np.around(measure_prob,3).tolist()}')

# Gumbel trick in the original tutorial is interesting, but it's not used here for simplicity # https://qgopt.readthedocs.io/en/latest/state_tomography.html num_measurement = 600000 index_measure = np_rng.choice(len(real_prob), size=num_measurement, p=real_prob) tmp0 = collections.Counter(index_measure.tolist()) measure_prob = np.array([tmp0[i] for i in range(len(real_prob))])/num_measurement print(f'measured probability:\n{np.around(measure_prob,3).tolist()}')

measured probability:
[0.067, 0.032, 0.029, 0.055, 0.061, 0.061, 0.092, 0.031, 0.073, 0.089, 0.077, 0.088, 0.034, 0.044, 0.106, 0.062]

In [4]:

class StateTomographyModel(torch.nn.Module):
    def __init__(self, povm_list):
        super().__init__()
        assert (povm_list.ndim==3) and (povm_list.shape[1]==povm_list.shape[2])
        dim = povm_list.shape[1]
        assert np.abs(povm_list.sum(axis=0) - np.eye(dim)).max() < 1e-10
        assert np.linalg.eigvalsh(povm_list)[:,0].min() > -1e-10
        self.manifold = numqi.manifold.Trace1PSD(dim, dim, dtype=torch.complex128)
        self.povm_T = torch.tensor(povm_list.transpose(0,2,1).reshape(-1,dim*dim).copy(), dtype=torch.complex128)
        self.target_prob = None

    def set_probability(self, np0):
        assert np0.shape==(self.povm_T.shape[0],)
        assert (np0.min() >= 0) and (abs(np0.sum()-1) < 1e-10)
        self.target_prob = torch.tensor(np0, dtype=torch.float64)

    def get_rho_and_probability(self, return_numpy=True):
        rho = self.manifold()
        prob = (self.povm_T @ rho.reshape(-1)).real
        ret = rho,prob
        if return_numpy:
            ret = ret[0].detach().numpy(), ret[1].detach().numpy()
        return ret

    def forward(self):
        _, prob = self.get_rho_and_probability(return_numpy=False)
        # loss = torch.sum((prob - self.target_prob)**2)
        loss = -torch.dot(self.target_prob, torch.log(prob)) # K-L divergence
        return loss

class StateTomographyModel(torch.nn.Module): def __init__(self, povm_list): super().__init__() assert (povm_list.ndim==3) and (povm_list.shape[1]==povm_list.shape[2]) dim = povm_list.shape[1] assert np.abs(povm_list.sum(axis=0) - np.eye(dim)).max() < 1e-10 assert np.linalg.eigvalsh(povm_list)[:,0].min() > -1e-10 self.manifold = numqi.manifold.Trace1PSD(dim, dim, dtype=torch.complex128) self.povm_T = torch.tensor(povm_list.transpose(0,2,1).reshape(-1,dim*dim).copy(), dtype=torch.complex128) self.target_prob = None def set_probability(self, np0): assert np0.shape==(self.povm_T.shape[0],) assert (np0.min() >= 0) and (abs(np0.sum()-1) < 1e-10) self.target_prob = torch.tensor(np0, dtype=torch.float64) def get_rho_and_probability(self, return_numpy=True): rho = self.manifold() prob = (self.povm_T @ rho.reshape(-1)).real ret = rho,prob if return_numpy: ret = ret[0].detach().numpy(), ret[1].detach().numpy() return ret def forward(self): _, prob = self.get_rho_and_probability(return_numpy=False) # loss = torch.sum((prob - self.target_prob)**2) loss = -torch.dot(self.target_prob, torch.log(prob)) # K-L divergence return loss

In [5]:

model = StateTomographyModel(povm_list)
model.set_probability(measure_prob)
theta_optim = numqi.optimize.minimize(model, theta0='uniform', num_repeat=1, tol=1e-7, print_freq=10)
rho_restore, prob_restore = model.get_rho_and_probability()
print('trace distance:', hf_trace_distance(rho_restore, target_rho))

model = StateTomographyModel(povm_list) model.set_probability(measure_prob) theta_optim = numqi.optimize.minimize(model, theta0='uniform', num_repeat=1, tol=1e-7, print_freq=10) rho_restore, prob_restore = model.get_rho_and_probability() print('trace distance:', hf_trace_distance(rho_restore, target_rho))

[step=0][time=0.004 seconds] loss=2.7533896391821506
[step=10][time=0.012 seconds] loss=2.702294219773905
[step=20][time=0.013 seconds] loss=2.700830364322484
[step=30][time=0.011 seconds] loss=2.700816591838533
[round=0] min(f)=2.7008150678467295, current(f)=2.7008150678467295
trace distance: 0.006476811515170481

In case one want to check the information at each training step, the following code retrieves the loss value and trace distance at each step.

In [6]:

def hf_callback_wrapper(hf_model, model, target_rho, info_dict):
    info_dict['loss'] = []
    info_dict['trace-distance'] = []
    def hf0(theta):
        info_dict['loss'].append(hf_model(theta, tag_grad=False)) #set parameter at this step
        rho,_ = model.get_rho_and_probability()
        info_dict['trace-distance'].append(hf_trace_distance(target_rho, rho))
    return hf0

hf_model = numqi.optimize.hf_model_wrapper(model)
info_dict = dict()
hf_callback = hf_callback_wrapper(hf_model, model, target_rho, info_dict)
num_parameter = len(numqi.optimize.get_model_flat_parameter(model))
theta0 = np_rng.uniform(-0.1, 0.1, size=num_parameter)
theta_optim = scipy.optimize.minimize(hf_model, theta0, jac=True, method='L-BFGS-B', tol=1e-7, callback=hf_callback)

def hf_callback_wrapper(hf_model, model, target_rho, info_dict): info_dict['loss'] = [] info_dict['trace-distance'] = [] def hf0(theta): info_dict['loss'].append(hf_model(theta, tag_grad=False)) #set parameter at this step rho,_ = model.get_rho_and_probability() info_dict['trace-distance'].append(hf_trace_distance(target_rho, rho)) return hf0 hf_model = numqi.optimize.hf_model_wrapper(model) info_dict = dict() hf_callback = hf_callback_wrapper(hf_model, model, target_rho, info_dict) num_parameter = len(numqi.optimize.get_model_flat_parameter(model)) theta0 = np_rng.uniform(-0.1, 0.1, size=num_parameter) theta_optim = scipy.optimize.minimize(hf_model, theta0, jac=True, method='L-BFGS-B', tol=1e-7, callback=hf_callback)

In [7]:

fig,(ax0,ax1) = plt.subplots(1, 2, figsize=(8,4))
ax0.plot(info_dict['loss'])
ax0.set_xlabel('step')
ax0.set_ylabel('loss')
ax1.plot(info_dict['trace-distance'])
ax1.set_xlabel('step')
ax1.set_ylabel('trace-distance')
fig.tight_layout()

fig,(ax0,ax1) = plt.subplots(1, 2, figsize=(8,4)) ax0.plot(info_dict['loss']) ax0.set_xlabel('step') ax0.set_ylabel('loss') ax1.plot(info_dict['trace-distance']) ax1.set_xlabel('step') ax1.set_ylabel('trace-distance') fig.tight_layout()