Coverage for jaxquantum/core/measurements.py: 77%
162 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""Measurements."""
3import optax
4import jax.numpy as jnp
6from collections.abc import Callable
7from matplotlib import pyplot as plt
8from tqdm import tqdm
9from typing import Optional, NamedTuple
10from functools import partial, reduce
12from jax import config, Array, jit, value_and_grad, lax, vmap
14from jaxquantum.core.qarray import Qarray, powm
15from jaxquantum.core.operators import identity
17from jax_tqdm import scan_tqdm
19config.update("jax_enable_x64", True)
22# Calculations ----------------------------------------------------------------
25def overlap(rho: Qarray, sigma: Qarray) -> Array:
26 """Overlap between two states or operators.
28 Args:
29 rho: state/operator.
30 sigma: state/operator.
32 Returns:
33 Overlap between rho and sigma.
34 """
36 if rho.is_vec() and sigma.is_vec():
37 # |<a|b>|^2. Compute the inner product directly over the trailing space
38 # axis (vectors are stored as (..., N)); robust for batched states.
39 a = rho.to_ket().data
40 b = sigma.to_ket().data
41 inner = jnp.sum(jnp.conj(a) * b, axis=-1)
42 return jnp.abs(inner) ** 2
43 elif rho.is_vec():
44 # <psi|sigma|psi>
45 psi = rho.to_ket()
46 Opsi = (sigma @ psi).data
47 return jnp.sum(jnp.conj(psi.data) * Opsi, axis=-1)
48 elif sigma.is_vec():
49 # <psi|rho|psi>
50 psi = sigma.to_ket()
51 Opsi = (rho @ psi).data
52 return jnp.sum(jnp.conj(psi.data) * Opsi, axis=-1)
53 else:
54 return (rho.dag() @ sigma).trace()
57def fidelity(rho: Qarray, sigma: Qarray, force_positivity: bool=False) -> (
58 jnp.ndarray):
59 """Fidelity between two states.
61 Args:
62 rho: state.
63 sigma: state.
64 force_positivity: force the states to be positive semidefinite
66 Returns:
67 Fidelity between rho and sigma.
68 """
69 rho = rho.to_dm()
70 sigma = sigma.to_dm()
72 sqrt_rho = powm(rho, 0.5, clip_eigvals=force_positivity)
74 return jnp.real(((powm(sqrt_rho @ sigma @ sqrt_rho, 0.5,
75 clip_eigvals=force_positivity)).tr())
76 ** 2)
79def _reconstruct_density_matrix(params: jnp.ndarray, dim: int) -> jnp.ndarray:
80 """
81 Pure function to parameterize a density matrix.
82 Ensures the resulting matrix is positive semi-definite and has trace 1.
83 """
84 num_real_params = dim * (dim + 1) // 2
86 real_part_flat = params[:num_real_params]
87 imag_part_flat = params[num_real_params:]
89 T = jnp.zeros((dim, dim), dtype=jnp.complex128)
91 # Set the real parts of the lower triangle from the first part of params
92 tril_indices = jnp.tril_indices(dim)
93 T = T.at[tril_indices].set(real_part_flat)
95 # Set the imaginary parts of the strictly lower triangle from the second part of params
96 tril_indices_off_diag = jnp.tril_indices(dim, k=-1)
97 T = T.at[tril_indices_off_diag].add(1j * imag_part_flat)
99 rho_unnormalized = T @ T.conj().T
100 # Enforce trace=1 by dividing by the trace
101 trace = jnp.trace(rho_unnormalized)
102 return rho_unnormalized / jnp.where(trace == 0, 1.0, trace)
105def _parametrize_density_matrix(rho_data: jnp.ndarray, dim: int) -> (
106 jnp.ndarray):
107 """
108 Calculates the parameter vector from a density matrix using Cholesky decomposition.
109 This is the inverse of the _reconstruct_density_matrix function.
110 """
111 # Add a small epsilon for numerical stability, ensuring the matrix is positive definite
112 T = jnp.linalg.cholesky(rho_data + 1e-9 * jnp.eye(dim))
114 # T is lower-triangular with a real, positive diagonal. This matches our
115 # parameterization convention.
117 # Extract the real parts of all lower-triangular elements
118 tril_indices = jnp.tril_indices(dim)
119 real_part = T[tril_indices].real
121 # Extract the imaginary parts of the strictly lower-triangular elements
122 tril_indices_off_diag = jnp.tril_indices(dim, k=-1)
123 imag_part = T[tril_indices_off_diag].imag
125 return jnp.concatenate([real_part, imag_part])
128def _L1_reg(params: jnp.ndarray) -> jnp.ndarray:
129 """Pure function for L1 regularization."""
130 return jnp.sum(jnp.abs(params))
133def _likelihood(
134 params: jnp.ndarray, dim: int, basis: jnp.ndarray, results: jnp.ndarray
135) -> jnp.ndarray:
136 """Compute the log-likelihood."""
137 rho = _reconstruct_density_matrix(params, dim)
138 expected_outcomes = jnp.real(jnp.einsum("ijk,jk->i", basis, rho))
139 return -jnp.sum((expected_outcomes - results) ** 2)
142# This is the core JIT-ted training loop. It is a pure function.
143@partial(
144 jit,
145 static_argnames=[
146 "dim",
147 "epochs",
148 "optimizer",
149 "compute_infidelity",
150 "L1_reg_strength",
151 ],
152)
153def _run_tomography_scan(
154 initial_params,
155 initial_opt_state,
156 true_rho_data,
157 measurement_basis,
158 measurement_results,
159 dim,
160 epochs,
161 optimizer,
162 compute_infidelity,
163 L1_reg_strength,
164):
165 """
166 A pure, JIT-compiled function that runs the entire optimization.
167 Static arguments are those that define the computation graph and don't change during the run.
168 """
170 def loss(params):
171 log_likelihood = _likelihood(
172 params, dim, measurement_basis, measurement_results
173 )
174 regularization = L1_reg_strength * _L1_reg(params)
175 return -log_likelihood + regularization
177 loss_val_grad = value_and_grad(loss)
179 @scan_tqdm(epochs)
180 def train_step(carry, _):
181 params, opt_state = carry
182 loss_val, grads = loss_val_grad(params)
183 updates, new_opt_state = optimizer.update(grads, opt_state, params)
184 new_params = optax.apply_updates(params, updates)
186 # This `if` statement is safe inside a JIT-ted function because
187 # `compute_infidelity` is a "static" argument. JAX compiles a
188 # separate version of the code for each value of this flag.
189 if compute_infidelity:
190 rho = Qarray.create(_reconstruct_density_matrix(params, dim))
191 fid = fidelity(Qarray.create(true_rho_data), rho,
192 force_positivity=True)
193 infidelity = 1.0 - fid
194 else:
195 infidelity = jnp.nan
197 new_carry = (new_params, new_opt_state)
198 history = {
199 "loss": loss_val,
200 "grads": grads,
201 "params": params,
202 "infidelity": infidelity,
203 }
204 return new_carry, history
206 initial_carry = (initial_params, initial_opt_state)
207 final_carry, history = lax.scan(
208 train_step, initial_carry, jnp.arange(epochs), length=epochs
209 )
210 return final_carry, history
213class MLETomographyResult(NamedTuple):
214 rho: Qarray
215 params_history: list
216 loss_history: list
217 grads_history: list
218 infidelity_history: Optional[list]
221class QuantumStateTomography:
222 def __init__(
223 self,
224 rho_guess: Qarray,
225 measurement_basis: Qarray,
226 measurement_results: jnp.ndarray,
227 complete_basis: Optional[Qarray] = None,
228 true_rho: Optional[Qarray] = None,
229 ):
230 """
231 Reconstruct a quantum state from measurement results using quantum state tomography.
232 The tomography can be performed either by direct inversion or by maximum likelihood estimation.
234 Args:
235 rho_guess (Qarray): The initial guess for the quantum state.
236 measurement_basis (Qarray): The basis in which measurements are performed.
237 measurement_results (jnp.ndarray): The results of the measurements.
238 complete_basis (Optional[Qarray]): The complete basis for state
239 reconstruction used when using direct inversion.
240 Defaults to the measurement basis if not provided.
241 true_rho (Optional[Qarray]): The true quantum state, if known.
243 """
244 self.rho_guess = rho_guess.data
245 self.measurement_basis = measurement_basis.data
246 self.measurement_results = measurement_results
247 self.complete_basis = (
248 complete_basis.data
249 if (complete_basis is not None)
250 else measurement_basis.data
251 )
252 self.true_rho = true_rho
253 self._result = None
255 @property
256 def result(self) -> Optional[MLETomographyResult]:
257 return self._result
260 def quantum_state_tomography_mle(
261 self, L1_reg_strength: float = 0.0, epochs: int = 10000, lr: float = 5e-3
262 ) -> MLETomographyResult:
263 """Perform quantum state tomography using maximum likelihood
264 estimation (MLE).
266 This method reconstructs the quantum state from measurement results
267 by optimizing
268 a likelihood function using gradient descent. The optimization
269 ensures the
270 resulting density matrix is positive semi-definite with trace 1.
272 Args:
273 L1_reg_strength (float, optional): Strength of L1
274 regularization. Defaults to 0.0.
275 epochs (int, optional): Number of optimization iterations.
276 Defaults to 10000.
277 lr (float, optional): Learning rate for the Adam optimizer.
278 Defaults to 5e-3.
280 Returns:
281 MLETomographyResult: Named tuple containing:
282 - rho: Reconstructed quantum state as Qarray
283 - params_history: List of parameter values during optimization
284 - loss_history: List of loss values during optimization
285 - grads_history: List of gradient values during optimization
286 - infidelity_history: List of infidelities if true_rho was
287 provided, None otherwise
288 """
290 dim = self.rho_guess.shape[0]
291 optimizer = optax.adam(lr)
293 # Initialize parameters from the initial guess for the density matrix
294 params = _parametrize_density_matrix(self.rho_guess, dim)
295 opt_state = optimizer.init(params)
297 compute_infidelity_flag = self.true_rho is not None
299 # Provide a dummy array if no true_rho is available. It won't be used.
300 true_rho_data_or_dummy = (
301 self.true_rho.data
302 if compute_infidelity_flag
303 else jnp.empty((dim, dim), dtype=jnp.complex64)
304 )
306 final_carry, history = _run_tomography_scan(
307 initial_params=params,
308 initial_opt_state=opt_state,
309 true_rho_data=true_rho_data_or_dummy,
310 measurement_basis=self.measurement_basis,
311 measurement_results=self.measurement_results,
312 dim=dim,
313 epochs=epochs,
314 optimizer=optimizer,
315 compute_infidelity=compute_infidelity_flag,
316 L1_reg_strength=L1_reg_strength,
317 )
319 final_params, _ = final_carry
321 rho = Qarray.create(_reconstruct_density_matrix(final_params, dim))
323 self._result = MLETomographyResult(
324 rho=rho,
325 params_history=history["params"],
326 loss_history=history["loss"],
327 grads_history=history["grads"],
328 infidelity_history=history["infidelity"]
329 if compute_infidelity_flag
330 else None,
331 )
332 return self._result
334 def quantum_state_tomography_direct(
335 self,
336 ) -> Qarray:
338 """Perform quantum state tomography using direct inversion.
340 This method reconstructs the quantum state from measurement results by
341 directly solving a system of linear equations. The method assumes that
342 the measurement basis is complete and the measurement results are
343 noise-free.
345 Returns:
346 Qarray: Reconstructed quantum state.
347 """
349 # Compute overlaps of measurement and complete operator bases
350 A = jnp.einsum("ijk,ljk->il", self.complete_basis, self.measurement_basis)
351 # Solve the linear system to find the coefficients
352 coefficients = jnp.linalg.solve(A, self.measurement_results)
353 # Reconstruct the density matrix
354 rho = jnp.einsum("i, ijk->jk", coefficients, self.complete_basis)
356 return Qarray.create(rho)
358 def plot_results(self):
359 if self._result is None:
360 raise ValueError(
361 "No results to plot. Run quantum_state_tomography_mle first."
362 )
364 _, ax = plt.subplots(1, figsize=(5, 4))
365 if self._result.infidelity_history is not None:
366 ax2 = ax.twinx()
368 ax.plot(self._result.loss_history, color="C0")
369 ax.set_xlabel("Epoch")
370 ax.set_ylabel("$\\mathcal{L}$", color="C0")
371 ax.set_yscale("log")
373 if self._result.infidelity_history is not None:
374 ax2.plot(self._result.infidelity_history, color="C1")
375 ax2.set_yscale("log")
376 ax2.set_ylabel("$1-\\mathcal{F}$", color="C1")
377 plt.grid(False)
379 plt.show()
381def tensor_basis(single_basis: Qarray, n: int) -> Qarray:
382 """Construct n-fold tensor product basis from a single-system basis.
384 Args:
385 single_basis: The single-system operator basis as a Qarray.
386 n: Number of tensor copies to construct.
388 Returns:
389 Qarray containing the n-fold tensor product basis operators.
390 The resulting basis has b^n elements where b is the number
391 of operators in the single-system basis.
392 """
394 dims = single_basis.dims
396 single_basis = single_basis.data
397 b, _, _ = single_basis.shape
398 indices = jnp.stack(jnp.meshgrid(*[jnp.arange(b)] * n, indexing="ij"),
399 axis=-1).reshape(-1, n) # shape (b^n, n)
401 # Select the operators based on indices: shape (b^n, n, d, d)
402 selected = single_basis[indices] # shape: (b^n, n, d, d)
404 # Vectorized Kronecker products
405 full_basis = vmap(lambda ops: reduce(jnp.kron, ops))(selected)
407 new_dims = tuple(tuple(x**n for x in row) for row in dims)
409 return Qarray.create(full_basis, dims=new_dims, bdims=(b**n,))
412def _quantum_process_tomography(
413 map: Callable[[Qarray], Qarray],
414 physical_state_basis: Qarray,
415 physical_operator_basis: Qarray,
416 logical_state_basis: Optional[Qarray] = None,
417 logical_operator_basis: Optional[Qarray] = None,
418) -> Qarray:
420 if logical_state_basis is None:
421 logical_state_basis = physical_state_basis
422 if logical_operator_basis is None:
423 logical_operator_basis = physical_operator_basis
425 dsqr = logical_state_basis.bdims[-1]
427 d = int(jnp.sqrt(dsqr+1))
429 choi = Qarray.create(jnp.zeros((d, d))) ^ Qarray.create(jnp.zeros((d, d)))
431 with (tqdm(total=d * d) as pbar):
432 for k in range(dsqr):
433 rho_k = physical_state_basis[k] @ physical_state_basis[k].dag()
435 E_rho_k = map(rho_k)
437 measurement_results = jnp.real(
438 jnp.einsum('ijk,jk->i', physical_operator_basis.data,
439 E_rho_k.data))
441 QST = QuantumStateTomography(rho_guess=identity(d)/d,
442 measurement_results=measurement_results,
443 measurement_basis=logical_operator_basis,
444 )
445 res = QST.quantum_state_tomography_mle()
446 r = res.rho
447 print("----------------------")
448 print(logical_state_basis[k] @ logical_state_basis[k].dag())
449 print(r)
450 choi += (logical_state_basis[k] @ logical_state_basis[k].dag()
451 ) ^ r
452 pbar.update(1)
453 return choi