Quantum State Tomography in Python
A quantum state cannot be read off a single measurement. Each measurement returns only a random outcome, and the state is destroyed in the process. What we can do is prepare the same state many times, measure it in several different bases, and reconstruct the density matrix statistically from the outcome frequencies. This procedure is called quantum state tomography.
In this article we work through a concrete example: a noisy two-qubit entangled state measured in all nine Pauli-basis combinations. We reconstruct it with two methods, linear inversion and maximum likelihood estimation (MLE), and compare them with a 3D bar chart and a shot-number scaling study.
1. The Problem Setup
A density matrix $\rho$ of an $n$-qubit system is a $2^n \times 2^n$ matrix satisfying
$$
\rho = \rho^\dagger, \qquad \rho \succeq 0, \qquad \mathrm{Tr},\rho = 1 .
$$
A two-qubit state has $4^2 - 1 = 15$ free real parameters. Any such $\rho$ can be expanded in Pauli strings:
$$
\rho = \frac{1}{4} \sum_{i,j \in {I,X,Y,Z}} \langle \sigma_i \otimes \sigma_j \rangle , \sigma_i \otimes \sigma_j .
$$
Measuring in a basis with projectors $E_k$ yields outcome $k$ with probability given by the Born rule:
$$
p_k = \mathrm{Tr}\left( E_k \rho \right).
$$
The example state. We use a Bell-type state with a relative phase, mixed with white noise:
$$
\rho_{\mathrm{true}} = p , |\psi\rangle\langle\psi| + (1-p),\frac{I}{4}, \qquad |\psi\rangle = \frac{|00\rangle + e^{i\pi/3}|11\rangle}{\sqrt{2}}, \qquad p = 0.85 .
$$
The phase $e^{i\pi/3}$ makes the off-diagonal elements genuinely complex, so both the real and imaginary parts of $\rho$ matter.
The measurement. Each qubit is measured in the $X$, $Y$, or $Z$ basis, giving $3 \times 3 = 9$ settings. Each setting has 4 outcomes, so there are 36 projectors $E_k$ in total. Every setting is repeated $N$ times, and we record the relative frequencies $f_k$.
2. Two Reconstruction Methods
Linear inversion
Since $p_k = \mathrm{Tr}(E_k \rho)$ is linear in $\rho$, replacing $p_k$ with the observed frequency $f_k$ gives a linear least-squares problem:
It is fast and simple, but statistical noise often pushes the result outside the set of valid states. The reconstructed matrix can have negative eigenvalues, which is unphysical.
Maximum likelihood estimation
MLE maximizes the log-likelihood over the set of valid density matrices:
$$
\mathcal{L}(\rho) = \sum_k f_k \log \mathrm{Tr}(E_k \rho) .
$$
A classic fixed-point algorithm, the $R\rho R$ iteration, does this without any constraint handling. Define
$$
R(\rho) = \sum_k \frac{f_k}{\mathrm{Tr}(E_k \rho)} , E_k ,
$$
and iterate
$$
\rho \leftarrow \frac{R(\rho),\rho,R(\rho)}{\mathrm{Tr}\left[R(\rho),\rho,R(\rho)\right]} , \qquad \rho_0 = \frac{I}{4} .
$$
Because $R$ is Hermitian and positive semidefinite, every iterate is automatically a valid density matrix.
Evaluation metrics
We measure the quality of a reconstruction $\sigma$ against the truth $\rho$ with the Uhlmann fidelity and the trace distance:
$$
F(\rho,\sigma) = \left( \mathrm{Tr}\sqrt{\sqrt{\rho},\sigma,\sqrt{\rho}} \right)^2, \qquad T(\rho,\sigma) = \frac{1}{2}\left| \rho - \sigma \right|_1 .
$$
3. The Complete Code
The code below is a single script that runs everything: data simulation, both reconstructions, the scaling study, and one combined figure.
The MLE loop is the only expensive part. To speed it up, all 36 projectors are flattened into a single $36 \times 16$ matrix. Each iteration then needs only two matrix-vector products instead of Python loops over projectors.
1 | import numpy as np |
4. Code Walkthrough
Measurement operators. The function eig_projectors diagonalizes a Pauli matrix and returns the projectors onto its $+1$ and $-1$ eigenspaces. For each of the nine settings $(a, b)$ we build four two-qubit projectors with a Kronecker product, $E = P_a^{(\pm)} \otimes P_b^{(\pm)}$. Stacking them yields the array E of shape $(36, 4, 4)$.
The vectorization trick. Because $\mathrm{Tr}(E\rho) = \sum_{mn} E_{nm}\rho_{mn}$, the Born probabilities for all 36 outcomes come from a single matrix-vector product, A @ rho.reshape(-1). Here A holds the transposed, flattened projectors. This is the basis of the fast implementation: no Python loop over projectors appears anywhere in the hot path. The same matrix also serves as the design matrix of the linear inversion.
Simulating data. simulate computes the Born probabilities for each setting, clips tiny negative values caused by floating-point error, renormalizes, and draws one multinomial sample of size shots per setting. Dividing by shots gives the relative frequencies $f_k$.
Linear inversion. linear_inversion solves the complex least-squares problem with np.linalg.lstsq. The 36 equations for 16 unknowns are overdetermined but full rank, because the Pauli-basis projectors span the whole operator space. The result is symmetrized to enforce Hermiticity and normalized to unit trace. Nothing forces positivity, which is exactly the weakness of this method.
MLE. In mle, each iteration computes the probabilities p, then builds $R = \sum_k (f_k/p_k) E_k$ as the single product (freq / p) @ Ef. The update $R\rho R$ followed by trace normalization keeps $\rho$ positive semidefinite at every step. A floor of $10^{-12}$ on p prevents division by zero, and the loop stops once successive iterates differ by less than the tolerance.
Metrics. psd_sqrt computes a matrix square root through an eigendecomposition with negative eigenvalues clipped to zero. This is more robust than a general-purpose matrix square root routine, and it lets the fidelity be evaluated even for the non-physical linear-inversion estimate. trace_distance sums the absolute eigenvalues of $\rho - \sigma$.
Scaling study. For shot counts from 10 to 30000 we repeat the whole experiment 15 times, reconstruct with both methods, and record mean and standard deviation of fidelity and trace distance.
5. Console Output
=== Experiment with 1000 shots per setting ===
Purity of true state : 0.7919
Eigenvalues (true) : [0.887 0.038 0.038 0.037]
Eigenvalues (linear inv.) : [0.888 0.06 0.043 0.009]
Eigenvalues (MLE) : [0.888 0.059 0.04 0.013]
MLE iterations : 601
Fidelity LI : 0.9868 MLE : 0.9910
Trace distance LI : 0.0339 MLE : 0.0304
True rho (real part):
[[0.462 0. 0. 0.212]
[0. 0.038 0. 0. ]
[0. 0. 0.038 0. ]
[0.212 0. 0. 0.462]]
True rho (imaginary part):
[[ 0. 0. 0. -0.368]
[ 0. 0. 0. 0. ]
[ 0. 0. 0. 0. ]
[ 0.368 0. 0. -0. ]]
MLE rho (real part):
[[ 0.46 -0. 0.008 0.211]
[-0. 0.036 0.017 0. ]
[ 0.008 0.017 0.036 0.01 ]
[ 0.211 0. 0.01 0.468]]
MLE rho (imaginary part):
[[ 0. 0.008 0.003 -0.367]
[-0.008 0. 0.012 -0.008]
[-0.003 -0.012 0. -0.003]
[ 0.367 0.008 0.003 0. ]]
=== Scaling study (mean over 15 trials) ===
shots | F (LI) F (MLE) | T (LI) T (MLE)
10 | 1.0250 0.8615 | 0.3810 0.2334
30 | 0.9862 0.9132 | 0.2496 0.1832
100 | 0.9664 0.9446 | 0.1272 0.1025
300 | 0.9688 0.9607 | 0.0782 0.0691
1000 | 0.9902 0.9902 | 0.0415 0.0398
3000 | 0.9979 0.9980 | 0.0210 0.0201
10000 | 0.9994 0.9995 | 0.0112 0.0109
30000 | 0.9997 0.9997 | 0.0076 0.0075
6. Result Image
7. Reading the Results
The figure has six panels. The top row shows the magnitudes $|\rho_{ij}|$ of the density matrix as 3D bar charts, and the bottom row summarizes spectra and accuracy.
Top row: the density matrices. The true state has a distinctive pattern. The four diagonal bars sit at the populations, with $|00\rangle$ and $|11\rangle$ tall and $|01\rangle$, $|10\rangle$ low. The two corner bars at positions $(00, 11)$ and $(11, 00)$ are the coherences, which encode the entanglement. Both reconstructions reproduce this pattern. Linear inversion shows small spurious bars in elements that should be zero, which is statistical noise from finite sampling. MLE suppresses much of this, because the positivity constraint links elements together.
Bottom left: the eigenvalue spectrum. The true state has one large eigenvalue and three equal small ones, reflecting a pure-state component plus a uniform noise floor. The spectrum of the linear-inversion estimate is spread apart, with the largest eigenvalue pushed up and the smallest pushed down, often below zero. A negative eigenvalue means the estimate is not a legitimate quantum state. The MLE spectrum stays non-negative by construction, and the smallest eigenvalues are pulled toward the boundary.
Bottom middle: fidelity versus shots. Fidelity rises toward 1 as the number of shots grows. At small $N$ the two methods differ visibly, and the MLE curve is generally higher and has a narrower band, indicating a more stable estimator. At large $N$ the statistical noise becomes small enough that the two curves merge.
Bottom right: trace distance versus shots. On log-log axes the error falls roughly along the dashed $N^{-1/2}$ reference line. This is the standard statistical limit: the noise in each estimated frequency shrinks as $1/\sqrt{N}$, and the reconstruction error follows it. Linear inversion and MLE approach the same asymptotic behavior, but MLE has a smaller prefactor in the low-shot regime, where positivity carries real information.
Takeaways. Linear inversion is a one-line solve and is perfectly adequate when data are plentiful. Its estimate can be unphysical when data are scarce. MLE costs an iterative loop, but the vectorized implementation above keeps the cost small, and it always returns a valid density matrix. For real experiments with limited shot budgets, MLE or a related constrained estimator is the standard choice.
8. Extending the Experiment
The same code structure scales to more qubits. Only the list of measurement settings and the projector construction change, while the linear-inversion and MLE routines remain untouched. Since the number of settings grows as $3^n$ and the matrix dimension as $2^n$, the vectorized formulation becomes even more valuable. Natural next steps include replacing MLE with a Bayesian estimator to obtain error bars on every matrix element, or using compressed-sensing methods that exploit the low rank of nearly pure states to reduce the number of required settings.
first and are then pushed further by the logarithmic growth of $x^{*}$. Expensive resources appear mainly at the ridge peaks, and cheap resources also cover the quiet hours.
