Pure State Decomposition¶
reference: Unified Framework for Calculating Convex Roof Resource Measures arxiv-link
Many quantum phenomena have transitioned from being of purely theoretical interest to serving as invaluable resources in quantum information processing tasks. Convex roof extension is a widely used measures for quantifying the convex resource theories like entanglement, coherence and magic states. Convex roof extension begins with a measure of pure states and then extends to mixed states:
$$ R(\rho)=\min_{\left\{ p_{i},|\psi_{i}\rangle\right\} }\sum_{i}p_{i}R\left(|\psi_{i}\rangle\right), $$
where the minimization is taken over all possible pure state decompositions of the given mixed state $\rho$ satisfying
$$ \rho=\sum_{i}p_{i}|\psi_{i}\rangle\langle\psi_{i}| $$
In this notebook, we will discuss how to perform the pure state decomposition, and how to do the minimization over all possible pure state decompositions.
import numpy as np
import torch
try:
import numqi
except ImportError:
%pip install numqi
import numqi
np_rng = np.random.default_rng()
Stiefel manifold¶
Before we start, let's introduce the Stiefel manifold. The Stiefel manifold $St(n,r)$ is the set of $n$-by-$r$ with complex matrices with orthogonal contraints:
$$ \mathrm{St}(n,r)=\left\{ A\in\mathbb{C}^{n\times r}:A^{\dagger}A=I_{r}\right\}. $$
The element $A$ is called a Stiefel matrix. It's different from the unitary matrices which requires orthogonal contraints on both sides:
$$ A\in U(n)\Rightarrow A^{\dagger}A=AA^{\dagger}=I_n. $$
One can think of the Stiefel matrix as the first $r$ columns of a uniatry matrix. For example, stacking two Bell states makes a Stiefel matrix of size $4\times 2$:
$$ \left[\left|\Phi^{+}\right\rangle ,\left|\Psi^{+}\right\rangle \right]=\frac{1}{\sqrt{2}}\left[\begin{array}{cc} 1&0\\ 0 & 1\\ 0 & 1\\ 1 & 0 \end{array}\right]\in\mathrm{St}(4,2) $$
A = np.stack([numqi.state.Bell(0), numqi.state.Bell(2)], axis=1)
print(A)
[[0.70710678 0. ] [0. 0.70710678] [0. 0.70710678] [0.70710678 0. ]]
Random Stiefel matrix can be generated:
A = numqi.random.rand_Stiefel_matrix(5, 3, iscomplex=True)
print('A\n', np.around(A, 3))
A [[-0.024+0.16j 0.55 +0.317j -0.466+0.209j] [ 0.354+0.082j -0.154+0.268j -0.294-0.215j] [-0.179+0.433j -0.402+0.36j -0.088+0.472j] [ 0.403+0.635j 0.048-0.374j -0.263-0.232j] [-0.039+0.234j -0.241-0.104j 0.11 +0.49j ]]
Let's check $A^{\dagger}A=I_{r}$:
print('A^dag A=\n', np.around(A.conj().T @ A, 3))
A^dag A= [[ 1.+0.j 0.+0.j -0.-0.j] [ 0.-0.j 1.+0.j -0.-0.j] [-0.+0.j -0.+0.j 1.+0.j]]
Only when $n=r$ (unitary matrix), $AA^{\dagger}=I_{n}$ is true.
A = numqi.random.rand_Stiefel_matrix(4, 2, iscomplex=True)
print('A in St(4,2): A^dag A=\n', np.around(A@A.T.conj(), 3))
A = numqi.random.rand_Stiefel_matrix(4, 4, iscomplex=True)
print('\nA in St(4,4): A^dag A=\n', np.around(A@A.T.conj(), 3))
A in St(4,2): A^dag A= [[ 0.183+0.j 0.019+0.274j -0.089+0.035j 0.088-0.239j] [ 0.019-0.274j 0.735+0.j -0.284+0.142j -0.112+0.079j] [-0.089-0.035j -0.284-0.142j 0.382+0.j -0.326-0.141j] [ 0.088+0.239j -0.112-0.079j -0.326+0.141j 0.701+0.j ]] A in St(4,4): A^dag A= [[ 1.+0.j -0.-0.j 0.+0.j 0.-0.j] [-0.+0.j 1.+0.j -0.+0.j 0.+0.j] [ 0.-0.j -0.-0.j 1.+0.j 0.-0.j] [ 0.+0.j 0.-0.j 0.+0.j 1.+0.j]]
Pure state decomposition and Stiefel matrix¶
In this part, we will show how to perform the pure state decomposition with the help of Stiefel matrix. Let's try its inverse problem first to make familiar with this decomposition:
Given the probability distribution $\{p_i\}$ and the corresponding pure states $\{|\psi_i\rangle\}$, how to construct the density matrix $\rho$
The answer is simple: just follow the formula $\rho=\sum_i p_i|\psi_i\rangle\langle\psi_i|$.
dim = 3
num_state = 5
prob_list = numqi.random.rand_discrete_probability(num_state)
psi_list = numqi.random.rand_haar_state(dim, batch_size=num_state)
mat0 = sum(x*y.reshape(-1,1)*y.conj() for x,y in zip(prob_list, psi_list))
# mat0 = np.einsum(prob_list, [0], psi_list, [0,1], psi_list.conj(), [0,2], [1,2], optimize=True)
Let's verify that the output matrix mat0 is a valid density matrix: positive semidefinite and trace 1.
EVL = np.linalg.eigvalsh(mat0)
print('eigenvalues of mat0:', np.around(EVL, 3))
print('trace of mat0:', np.around(np.trace(mat0), 3))
eigenvalues of mat0: [0.061 0.41 0.529] trace of mat0: (1-0j)
However, given the density matrix $\rho\in\mathbb{C}^{r\times r}$, finding its pure state decomposition is not that easy. We need the help of Stiefel matrix. Denote the eigenvalue decomposition of $\rho$ as
$$ \rho=\sum_{j}\lambda_{j}|\lambda_{j}\rangle\langle\lambda_{j}| $$
for any Stiefel matrix $X\in\mathrm{St}(r,n)$, we can construct the auxiliary states as (not normalized)
$$ \left|\tilde{\psi}_{i}\right\rangle =\sum_{j=1}^{r}\sqrt{\lambda_{j}}X_{ij}\left|\lambda_{j}\right\rangle. $$
Then, the following pair of quantities
$$ p_{i}=\left\langle \tilde{\psi}_{i}|\tilde{\psi}_{i}\right\rangle ,\left|\psi_{i}\right\rangle =\frac{1}{\sqrt{p_{i}}}\left|\tilde{\psi}_{i}\right\rangle $$
make a valid pure state decomposition of $\rho$. On the other hand, any valid pure state decomposition with $n$ states can be obtained by choosing $X\in\mathrm{St}(r,n)$. See arxiv-link for more details. Here, let's verify the statement above numerically.
First, let's generate some random density matrix $\rho$ and calculate its pure state decomposition.
dim = 3
num_state = 5
rho = numqi.random.rand_density_matrix(dim)
EVL,EVC = np.linalg.eigh(rho)
print('rho=\n', np.around(rho, 3))
print('\neigenvalues=', np.around(EVL, 3))
rho= [[ 0.595+0.j -0.116-0.004j 0.179-0.183j] [-0.116+0.004j 0.1 +0.j 0.004+0.084j] [ 0.179+0.183j 0.004-0.084j 0.304+0.j ]] eigenvalues= [0.049 0.179 0.772]
Then, we randomly generate a Stiefel matrix $X$ and build the auxiliary states as the formula above.
X = numqi.random.rand_Stiefel_matrix(num_state, dim)
auxiliary_state = X @ (np.sqrt(EVL) * EVC).T
prob_list = np.linalg.norm(auxiliary_state, axis=1)**2
psi_list = auxiliary_state / np.linalg.norm(auxiliary_state, axis=1, keepdims=True)
print('pi:', np.around(prob_list, 3))
print('psi_list.shape:', psi_list.shape)
pi: [0.162 0.076 0.087 0.086 0.589] psi_list.shape: (5, 3)
Let's verify that the decomposition is a valid one $\rho=\sum_i p_i|\psi_i\rangle\langle\psi_i|$.
mat0 = sum(x*y.reshape(-1,1)*y.conj() for x,y in zip(prob_list, psi_list))
print('matO:\n', np.around(mat0, 3))
print('err(mat0-rho)=', np.abs(mat0-rho).max())
# less than 1e-12, which means the reconstruction is perfect
matO: [[ 0.595+0.j -0.116-0.004j 0.179-0.183j] [-0.116+0.004j 0.1 +0.j 0.004+0.084j] [ 0.179+0.183j 0.004-0.084j 0.304+0.j ]] err(mat0-rho)= 4.775249788392736e-16
Minimization over all pure state decompositions¶
The power of connecting the pure state decomposition with the Stiefel matrix is that the whole calculation is differentiable, which means we can use the gradient-based optimization method to minimize the convex roof extension. Instead of generating random Stiefel matrix using numqi.random, we need to use numqi.manifold.Stiefel to support the gradient calculation.
manifold = numqi.manifold.Stiefel(dim=4, rank=2)
A = manifold()
print('A=', A)
print('\nA^dag A=\n', A.conj().T @ A)
A= tensor([[ 0.9060, -0.1453],
[ 0.1714, -0.2821],
[-0.3432, -0.0303],
[ 0.1790, 0.9478]], dtype=torch.float64, grad_fn=<ViewBackward0>)
A^dag A=
tensor([[ 1.0000e+00, -3.6082e-16],
[-3.6082e-16, 1.0000e+00]], dtype=torch.float64,
grad_fn=<MmBackward0>)
Above, A is a Stiefel matirx, also a torch variable with gradient support, which means A.backward() is supported. If one have resource function $R$ for any pure state, then the following DummyModel can be used to calculate the resource for any density matrix $\rho$.
def R_for_pure_state(psi):
# implement the Rourse for a pure state
raise NotImplementedError
class DummyModel(torch.nn.Module):
def __init__(self, rho, num_state):
super().__init__()
EVL,EVC = np.linalg.eigh(rho)
self.EVL = torch.tensor(EVL)
self.EVC = torch.tensor(EVC)
self.manifold = numqi.manifold.Stiefel(num_state, rho.shape[0])
def forward(self):
X = self.manifold()
auxiliary_state = X @ (torch.sqrt(self.EVL) * self.EVC).T
prob_list = np.linalg.norm(auxiliary_state, axis=1)**2
psi_list = auxiliary_state / np.linalg.norm(auxiliary_state, axis=1, keepdims=True)
loss = sum(x*R_for_pure_state for x,y in zip(prob_list, psi_list))
return loss
## write your own "R_for_pure_state" function and "rho"
# model = DummyModel(rho, num_state)
# theta_optim = numqi.optimize.minimize(model, theta0='uniform', num_repeat=3, method='L-BFGS-B')
This provide a unified framework for calculating convex roof resource measures. Various resource measures are presented in the following notebooks: