Channel tomography

Quantum Channel Tomography¶

Adapted from QGopt/docs/channel-tomography

In [1]:

# https://qgopt.readthedocs.io/en/latest/channel_tomography.html
import numpy as np
import torch
import scipy.optimize
import matplotlib.pyplot as plt
import opt_einsum

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

# https://qgopt.readthedocs.io/en/latest/channel_tomography.html import numpy as np import torch import scipy.optimize import matplotlib.pyplot as plt import opt_einsum 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

In [2]:

num_qubit = 2
dim_in = 2**num_qubit
dim_out = 2**num_qubit
choi_rank = 2

target_choi = numqi.random.rand_choi_op(dim_in, dim_out, choi_rank)

num_qubit = 2 dim_in = 2**num_qubit dim_out = 2**num_qubit choi_rank = 2 target_choi = numqi.random.rand_choi_op(dim_in, dim_out, choi_rank)

In [3]:

num_measure = 600000
tmp0 = np_rng.normal(size=(num_measure,dim_in*2)).astype(np.float64).view(np.complex128)
pure_state_list = tmp0 / np.linalg.norm(tmp0, axis=1, keepdims=True)
rho_in = np.einsum(pure_state_list, [0,1], pure_state_list.conj(), [0,2], [0,1,2], optimize=True)
rho_out = np.einsum(target_choi.reshape(dim_in,dim_out,dim_in,dim_out), [0,1,2,3], rho_in, [4,0,2], [4,1,3], optimize=True)

num_measure = 600000 tmp0 = np_rng.normal(size=(num_measure,dim_in*2)).astype(np.float64).view(np.complex128) pure_state_list = tmp0 / np.linalg.norm(tmp0, axis=1, keepdims=True) rho_in = np.einsum(pure_state_list, [0,1], pure_state_list.conj(), [0,2], [0,1,2], optimize=True) rho_out = np.einsum(target_choi.reshape(dim_in,dim_out,dim_in,dim_out), [0,1,2,3], rho_in, [4,0,2], [4,1,3], optimize=True)

In [4]:

povm = numqi.utils.get_tetrahedron_POVM(num_qubit)
real_prob = np.einsum(povm, [0,1,2], rho_out, [3,2,1], [3,0], optimize=True).real
# Gumbel trick
tmp0 = -np.log(-np.log(np_rng.uniform(0, 1, size=real_prob.shape)))
index_measure = np.argmax(tmp0 + np.log(real_prob), axis=-1)

povm = numqi.utils.get_tetrahedron_POVM(num_qubit) real_prob = np.einsum(povm, [0,1,2], rho_out, [3,2,1], [3,0], optimize=True).real # Gumbel trick tmp0 = -np.log(-np.log(np_rng.uniform(0, 1, size=real_prob.shape))) index_measure = np.argmax(tmp0 + np.log(real_prob), axis=-1)

In [5]:

class ChannelTomographyModel(torch.nn.Module):
    def __init__(self, dim_in:int, dim_out:int, choi_rank:int):
        super().__init__()
        self.manifold = numqi.manifold.QuantumChannel(dim_in, dim_out, method='sqrtm')
        self.rho_in = None
        self.measure_result_T = None
        self.channel_contract = None

    def set_data(self, rho_in, measure_result):
        dim_in = self.manifold.dim_in
        assert (rho_in.ndim==3) and (rho_in.shape[1]==rho_in.shape[2]) and (rho_in.shape[1]==dim_in)
        assert (measure_result.ndim==3) and (measure_result.shape[1]==measure_result.shape[2]) and (measure_result.shape[1]==dim_out)
        assert rho_in.shape[0]==measure_result.shape[0]
        N0 = rho_in.shape[0]
        self.rho_in = torch.tensor(rho_in, dtype=torch.complex128)
        self.measure_result_T = torch.tensor(measure_result.transpose(0,2,1).reshape(N0,-1), dtype=torch.complex128)
        tmp0 = [self.manifold.choi_rank, self.manifold.dim_out, dim_in]
        self.channel_contract = opt_einsum.contract_expression(tmp0, [0,1,2], tmp0, [0,3,4], [N0,dim_in,dim_in], [5,2,4], [5,1,3])

    def forward(self):
        # TODO time issue
        kraus_op = self.manifold()
        rho_out = self.channel_contract(kraus_op, kraus_op.conj(), self.rho_in)
        probability = opt_einsum.contract(rho_out.reshape(rho_out.shape[0],-1), [0,1], self.measure_result_T, [0,1], [0]).real
        loss = -torch.mean(torch.log(probability))
        return loss

    def get_choi_op(self):
        with torch.no_grad():
            kraus_op = self.manifold()
            choi_op = numqi.channel.kraus_op_to_choi_op(kraus_op.numpy())
        return choi_op

class ChannelTomographyModel(torch.nn.Module): def __init__(self, dim_in:int, dim_out:int, choi_rank:int): super().__init__() self.manifold = numqi.manifold.QuantumChannel(dim_in, dim_out, method='sqrtm') self.rho_in = None self.measure_result_T = None self.channel_contract = None def set_data(self, rho_in, measure_result): dim_in = self.manifold.dim_in assert (rho_in.ndim==3) and (rho_in.shape[1]==rho_in.shape[2]) and (rho_in.shape[1]==dim_in) assert (measure_result.ndim==3) and (measure_result.shape[1]==measure_result.shape[2]) and (measure_result.shape[1]==dim_out) assert rho_in.shape[0]==measure_result.shape[0] N0 = rho_in.shape[0] self.rho_in = torch.tensor(rho_in, dtype=torch.complex128) self.measure_result_T = torch.tensor(measure_result.transpose(0,2,1).reshape(N0,-1), dtype=torch.complex128) tmp0 = [self.manifold.choi_rank, self.manifold.dim_out, dim_in] self.channel_contract = opt_einsum.contract_expression(tmp0, [0,1,2], tmp0, [0,3,4], [N0,dim_in,dim_in], [5,2,4], [5,1,3]) def forward(self): # TODO time issue kraus_op = self.manifold() rho_out = self.channel_contract(kraus_op, kraus_op.conj(), self.rho_in) probability = opt_einsum.contract(rho_out.reshape(rho_out.shape[0],-1), [0,1], self.measure_result_T, [0,1], [0]).real loss = -torch.mean(torch.log(probability)) return loss def get_choi_op(self): with torch.no_grad(): kraus_op = self.manifold() choi_op = numqi.channel.kraus_op_to_choi_op(kraus_op.numpy()) return choi_op

In [6]:

model = ChannelTomographyModel(dim_in, dim_out, choi_rank)
model.set_data(rho_in, povm[index_measure])
theta_optim = numqi.optimize.minimize(model, theta0='uniform', num_repeat=1, tol=1e-7, print_freq=1)
# theta_optim = numqi.optimize.minimize_adam(model, num_step=100, theta0='uniform', optim_args=('adam',0.03,0.01), tqdm_update_freq=1)

choi_restored = model.get_choi_op()
# also known as J-distance https://doi.org/10.1103/PhysRevA.98.062336
print('trace distance of Choi operator:', hf_trace_distance(target_choi, choi_restored))

model = ChannelTomographyModel(dim_in, dim_out, choi_rank) model.set_data(rho_in, povm[index_measure]) theta_optim = numqi.optimize.minimize(model, theta0='uniform', num_repeat=1, tol=1e-7, print_freq=1) # theta_optim = numqi.optimize.minimize_adam(model, num_step=100, theta0='uniform', optim_args=('adam',0.03,0.01), tqdm_update_freq=1) choi_restored = model.get_choi_op() # also known as J-distance https://doi.org/10.1103/PhysRevA.98.062336 print('trace distance of Choi operator:', hf_trace_distance(target_choi, choi_restored))

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[6], line 1
----> 1 model = ChannelTomographyModel(dim_in, dim_out, choi_rank)
      2 model.set_data(rho_in, povm[index_measure])
      3 theta_optim = numqi.optimize.minimize(model, theta0='uniform', num_repeat=1, tol=1e-7, print_freq=1)

Cell In[5], line 4, in ChannelTomographyModel.__init__(self, dim_in, dim_out, choi_rank)
      2 def __init__(self, dim_in:int, dim_out:int, choi_rank:int):
      3     super().__init__()
----> 4     self.manifold = numqi.manifold.QuantumChannel(dim_in, dim_out, method='sqrtm')
      5     self.rho_in = None
      6     self.measure_result_T = None

File ~/work/numqi/numqi/python/numqi/manifold/_compose.py:140, in QuantumChannel.__init__(self, dim_in, dim_out, choi_rank, batch_size, method, euler_with_phase, return_kind, requires_grad, dtype, device)
    138 assert dtype in {torch.complex64,torch.complex128}
    139 assert return_kind in {'kraus','choi'}
--> 140 self.manifold = Stiefel(choi_rank*dim_out, dim_in, batch_size, method=method, euler_with_phase=euler_with_phase,
    141                         requires_grad=requires_grad, dtype=dtype, device=device)
    142 self.dim_in = dim_in
    143 self.dim_out = dim_out

File ~/work/numqi/numqi/python/numqi/manifold/_stiefel.py:39, in Stiefel.__init__(self, dim, rank, batch_size, method, euler_with_phase, requires_grad, dtype, device)
     37 assert isinstance(device, torch.device)
     38 # choleskyL is really bad
---> 39 assert method in {'choleskyL','qr','so-exp','so-cayley','polar','euler'}
     40 if method in {'qr','polar'}:
     41     tmp0 = dim*rank if (dtype in {torch.float32,torch.float64}) else 2*dim*rank

AssertionError: 

In [7]:

def hf_callback_wrapper(hf_model, model, target_choi, info_dict):
    info_dict['loss'] = []
    info_dict['trace-distance'] = []
    def hf0(theta):
        loss = hf_model(theta, tag_grad=False) #set parameter at this step
        print('[{}] loss={:.7f}'.format(len(info_dict['loss']), loss))
        info_dict['loss'].append(loss)
        choi_op = model.get_choi_op()
        info_dict['trace-distance'].append(hf_trace_distance(target_choi, choi_op))
    return hf0
hf_model = numqi.optimize.hf_model_wrapper(model)
info_dict = dict()
hf_callback = hf_callback_wrapper(hf_model, model, target_choi, info_dict)
theta0 = np_rng.uniform(-0.1, 0.1, size=numqi.optimize.get_model_flat_parameter(model).size)
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_choi, info_dict): info_dict['loss'] = [] info_dict['trace-distance'] = [] def hf0(theta): loss = hf_model(theta, tag_grad=False) #set parameter at this step print('[{}] loss={:.7f}'.format(len(info_dict['loss']), loss)) info_dict['loss'].append(loss) choi_op = model.get_choi_op() info_dict['trace-distance'].append(hf_trace_distance(target_choi, choi_op)) return hf0 hf_model = numqi.optimize.hf_model_wrapper(model) info_dict = dict() hf_callback = hf_callback_wrapper(hf_model, model, target_choi, info_dict) theta0 = np_rng.uniform(-0.1, 0.1, size=numqi.optimize.get_model_flat_parameter(model).size) theta_optim = scipy.optimize.minimize(hf_model, theta0, jac=True, method='L-BFGS-B', tol=1e-7, callback=hf_callback)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[7], line 11
      9         info_dict['trace-distance'].append(hf_trace_distance(target_choi, choi_op))
     10     return hf0
---> 11 hf_model = numqi.optimize.hf_model_wrapper(model)
     12 info_dict = dict()
     13 hf_callback = hf_callback_wrapper(hf_model, model, target_choi, info_dict)

NameError: name 'model' is not defined

In [8]:

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(r'$|| C_{\mathrm{true}}-C_{\mathrm{restore}} ||_{tr}$')
ax1.set_yscale('log')
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(r'$|| C_{\mathrm{true}}-C_{\mathrm{restore}} ||_{tr}$') ax1.set_yscale('log') fig.tight_layout()

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[8], line 2
      1 fig,(ax0,ax1) = plt.subplots(1, 2, figsize=(8,4))
----> 2 ax0.plot(info_dict['loss'])
      3 ax0.set_xlabel('step')
      4 ax0.set_ylabel('loss')

NameError: name 'info_dict' is not defined