Coverage for qml_essentials / math.py: 97%
105 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-06-22 12:56 +0000
1import jax
2import jax.numpy as jnp
3from qml_essentials.operations import _cdtype
4from scipy.linalg import logm
7def logm_v(A: jnp.ndarray, **kwargs) -> jnp.ndarray:
8 """
9 Compute the logarithm of a matrix. If the provided matrix has an additional
10 batch dimension, the logarithm of each matrix is computed.
12 Args:
13 A (jnp.ndarray): The (potentially batched) matrices of which to compute
14 the logarithm.
16 Returns:
17 jnp.ndarray: The log matrices
18 """
19 # TODO: check warnings
20 if len(A.shape) == 2:
21 return logm(A, **kwargs)
22 elif len(A.shape) == 3:
23 AV = jnp.zeros(A.shape, dtype=_cdtype())
24 for i in range(A.shape[0]):
25 AV = AV.at[i].set(logm(A[i], **kwargs))
26 return AV
27 else:
28 raise NotImplementedError("Unsupported shape of input matrix")
31def _sqrt_matrix(density_matrix: jnp.ndarray) -> jnp.ndarray:
32 r"""Compute the matrix square root of a density matrix.
34 Uses eigendecomposition: if :math:`\rho = V \Lambda V^\dagger`, then
35 :math:`\sqrt{\rho} = V \sqrt{\Lambda} V^\dagger`.
37 Negative eigenvalues (numerical noise) are clamped to zero.
39 Args:
40 density_matrix: Density matrix of shape ``(d, d)`` or ``(B, d, d)``.
42 Returns:
43 The matrix square root with the same shape as the input.
44 """
45 evs, vecs = jnp.linalg.eigh(density_matrix)
46 evs = jnp.real(evs)
47 evs = jnp.where(evs > 0.0, evs, 0.0)
49 if density_matrix.ndim == 3:
50 # batched: (B, d, d)
51 sqrt_evs = jnp.sqrt(evs)[:, :, None] * jnp.eye(
52 density_matrix.shape[-1], dtype=_cdtype()
53 )
54 return vecs @ sqrt_evs @ jnp.conj(jnp.transpose(vecs, (0, 2, 1)))
56 # single: (d, d)
57 return vecs @ jnp.diag(jnp.sqrt(evs)) @ jnp.conj(vecs.T)
60def _fidelity_statevector(
61 state0: jnp.ndarray,
62 state1: jnp.ndarray,
63) -> jnp.ndarray:
64 r"""Fidelity between two pure states (state vectors).
66 .. math::
68 F(\ket{\psi}, \ket{\phi}) = \left|\braket{\psi | \phi}\right|^2
69 The inputs are normalised before the overlap is computed so that
70 the result is always in :math:`[0, 1]`.
71 """
72 # Normalise so that unnormalised inputs don't produce F > 1.
73 norm0 = jnp.linalg.norm(state0, axis=-1, keepdims=True)
74 norm1 = jnp.linalg.norm(state1, axis=-1, keepdims=True)
75 state0 = state0 / jnp.where(norm0 > 0, norm0, 1.0)
76 state1 = state1 / jnp.where(norm1 > 0, norm1, 1.0)
78 batched0 = state0.ndim > 1
79 batched1 = state1.ndim > 1
81 idx0 = "ab" if batched0 else "b"
82 idx1 = "ab" if batched1 else "b"
83 target = "a" if (batched0 or batched1) else ""
85 overlap = jnp.einsum(f"{idx0},{idx1}->{target}", jnp.conj(state0), state1)
86 return jnp.abs(overlap) ** 2
89def _fidelity_dm(
90 state0: jnp.ndarray,
91 state1: jnp.ndarray,
92) -> jnp.ndarray:
93 r"""Fidelity between two mixed states (density matrices)."""
94 sqrt_state0 = _sqrt_matrix(state0)
95 product = sqrt_state0 @ state1 @ sqrt_state0
97 evs = jnp.linalg.eigvalsh(product)
98 evs = jnp.real(evs)
99 evs = jnp.where(evs > 0.0, evs, 0.0)
101 return jnp.sum(jnp.sqrt(evs), axis=-1) ** 2
104def fidelity(
105 state0: jnp.ndarray,
106 state1: jnp.ndarray,
107) -> jnp.ndarray:
108 r"""Compute the fidelity between two quantum states.
110 Accepts either state vectors or density matrices.
112 Args:
113 state0: State vector or density matrix.
114 state1: State vector or density matrix (same kind as *state0*).
116 Returns:
117 Fidelity (scalar or shape ``(B,)``).
119 Raises:
120 ValueError: If the two states have incompatible shapes or
121 different representations (vector vs. matrix).
122 """
123 state0 = jnp.asarray(state0, dtype=_cdtype())
124 state1 = jnp.asarray(state1, dtype=_cdtype())
126 if state0.shape[-1] != state1.shape[-1]:
127 raise ValueError("The two states must have the same number of wires.")
129 is_sv0 = state0.ndim <= 2 and (
130 state0.ndim == 1 or state0.shape[-2] != state0.shape[-1]
131 )
132 is_sv1 = state1.ndim <= 2 and (
133 state1.ndim == 1 or state1.shape[-2] != state1.shape[-1]
134 )
136 if is_sv0 != is_sv1:
137 raise ValueError(
138 "Both states must be of the same kind "
139 "(both state vectors or both density matrices)."
140 )
142 if is_sv0:
143 return _fidelity_statevector(state0, state1)
144 return _fidelity_dm(state0, state1)
147def trace_distance(
148 state0: jnp.ndarray,
149 state1: jnp.ndarray,
150) -> jnp.ndarray:
151 r"""Compute the trace distance between two quantum states.
153 Supports single density matrices of shape ``(2**N, 2**N)`` and batched
154 density matrices of shape ``(B, 2**N, 2**N)``.
156 Args:
157 state0: Density matrix of shape ``(2**N, 2**N)`` or ``(B, 2**N, 2**N)``.
158 state1: Density matrix of shape ``(2**N, 2**N)`` or ``(B, 2**N, 2**N)``.
160 Returns:
161 Trace distance (scalar or shape ``(B,)``).
162 """
163 state0 = jnp.asarray(state0, dtype=_cdtype())
164 state1 = jnp.asarray(state1, dtype=_cdtype())
166 if state0.shape[-1] != state1.shape[-1]:
167 raise ValueError("The two states must have the same number of wires.")
169 eigvals = jnp.abs(jnp.linalg.eigvalsh(state0 - state1))
170 return jnp.sum(eigvals, axis=-1) / 2
173def phase_difference(
174 state0: jnp.ndarray,
175 state1: jnp.ndarray,
176) -> jnp.ndarray:
177 r"""Compute the phase difference between two state vectors.
179 A value of zero indicates the two states are related by at most a
180 real global factor (i.e. no relative phase). The result lies in
181 :math:`[-\pi, 1 + \pi]`.
183 Supports single state vectors of shape ``(2**N,)`` and batched state
184 vectors of shape ``(B, 2**N)``.
186 Args:
187 state0: State vector of shape ``(2**N,)`` or ``(B, 2**N)``.
188 state1: State vector of shape ``(2**N,)`` or ``(B, 2**N)``.
190 Returns:
191 Phase difference (scalar or shape ``(B,)``).
192 """
193 state0 = jnp.asarray(state0, dtype=_cdtype())
194 state1 = jnp.asarray(state1, dtype=_cdtype())
196 if state0.shape[-1] != state1.shape[-1]:
197 raise ValueError("The two states must have the same number of wires.")
199 batched0 = state0.ndim > 1
200 batched1 = state1.ndim > 1
202 idx0 = "ab" if batched0 else "b"
203 idx1 = "ab" if batched1 else "b"
204 target = "a" if (batched0 or batched1) else ""
206 inner = jnp.einsum(f"{idx0},{idx1}->{target}", jnp.conj(state0), state1)
207 return jnp.angle(inner)
210def _fubini_study_statevector(
211 jac: jnp.ndarray,
212 state: jnp.ndarray,
213) -> jnp.ndarray:
214 r"""Fubini-Study metric of a pure state.
216 The Fubini-Study metric is the real part of the quantum geometric tensor:
218 .. math::
220 g_{ij} = \mathrm{Re}\left[
221 \braket{\partial_i\psi | \partial_j\psi}
222 - \braket{\partial_i\psi | \psi}\braket{\psi | \partial_j\psi}
223 \right]
225 It relates to the pure-state quantum Fisher information by
226 :math:`F_{ij} = 4\,g_{ij}`.
228 Args:
229 jac: Jacobian :math:`\partial\ket{\psi}/\partial\theta` of shape
230 ``(d, P)`` with ``d = 2**N`` and ``P`` the number of parameters.
231 state: State vector of shape ``(d,)``.
233 Returns:
234 Real, symmetric metric of shape ``(P, P)``.
235 """
236 A = jnp.conj(jac.T) @ jac # A_ij = <∂_i ψ | ∂_j ψ>
237 v = jnp.conj(jac.T) @ state # v_i = <∂_i ψ | ψ>
238 return jnp.real(A - jnp.outer(v, jnp.conj(v)))
241def _qfi_statevector(
242 jac: jnp.ndarray,
243 state: jnp.ndarray,
244) -> jnp.ndarray:
245 r"""Quantum Fisher Information of a pure state.
247 For a normalised state :math:`\ket{\psi(\theta)}` the QFI is four times the
248 Fubini-Study metric (see :func:`_fubini_study_statevector`):
250 .. math::
252 F_{ij} = 4\,\mathrm{Re}\left[
253 \braket{\partial_i\psi | \partial_j\psi}
254 - \braket{\partial_i\psi | \psi}\braket{\psi | \partial_j\psi}
255 \right]
257 Args:
258 jac: Jacobian :math:`\partial\ket{\psi}/\partial\theta` of shape
259 ``(d, P)`` with ``d = 2**N`` and ``P`` the number of parameters.
260 state: State vector of shape ``(d,)``.
262 Returns:
263 Real, symmetric QFI matrix of shape ``(P, P)``.
264 """
265 return 4.0 * _fubini_study_statevector(jac, state)
268def _qfi_density(
269 jac: jnp.ndarray,
270 state: jnp.ndarray,
271 eps: float = 1e-12,
272) -> jnp.ndarray:
273 r"""Quantum Fisher Information of a mixed state.
275 Using the symmetric logarithmic derivative, the QFI of a density matrix
276 :math:`\rho = \sum_k p_k \ket{k}\bra{k}` reads
278 .. math::
280 F_{ij} = 2 \sum_{k, l : p_k + p_l > 0}
281 \frac{\mathrm{Re}\left(
282 \braket{k | \partial_i\rho | l}\braket{l | \partial_j\rho | k}
283 \right)}{p_k + p_l}
285 Eigenvalue pairs with :math:`p_k + p_l \le` ``eps`` are excluded from the
286 sum. Negative eigenvalues (numerical noise) are clamped to zero.
288 Args:
289 jac: Jacobian :math:`\partial\rho/\partial\theta` of shape
290 ``(d, d, P)`` with ``d = 2**N`` and ``P`` the number of parameters.
291 state: Density matrix of shape ``(d, d)``.
292 eps: Threshold below which an eigenvalue pair is masked out.
294 Returns:
295 Real, symmetric QFI matrix of shape ``(P, P)``.
296 """
297 evals, evecs = jnp.linalg.eigh(state)
298 evals = jnp.where(jnp.real(evals) > 0.0, jnp.real(evals), 0.0)
300 # ∂_i ρ in the eigenbasis: M[i]_kl = <k| ∂_i ρ |l>.
301 drho = jnp.moveaxis(jac, -1, 0) # (P, d, d)
302 M = jnp.conj(evecs.T) @ drho @ evecs # broadcast (d, d) over P
304 s = evals[:, None] + evals[None, :] # (d, d)
305 weights = jnp.where(s > eps, 2.0 / s, 0.0)
307 # ∂_i ρ is Hermitian, so <l| ∂_j ρ |k> = conj(<k| ∂_j ρ |l>).
308 F = jnp.einsum("ikl,jkl->ij", M * weights[None], jnp.conj(M))
309 return jnp.real(F)
312def _state_and_jacobian(state_fn, params: jnp.ndarray):
313 r"""Evaluate *state_fn* and its Jacobian at *params*.
315 The Jacobian is obtained with forward-mode automatic differentiation
316 (:func:`jax.jacfwd`), which yields the complex Jacobian directly for the
317 real-valued parameters.
319 Args:
320 state_fn: Callable mapping *params* to a quantum state.
321 params: Parameters at which to evaluate.
323 Returns:
324 Tuple ``(state, jac)`` of the state and its Jacobian, both cast to the
325 complex working dtype.
326 """
327 state = jnp.asarray(state_fn(params), dtype=_cdtype())
328 jac = jnp.asarray(jax.jacfwd(state_fn)(params), dtype=_cdtype())
329 return state, jac
332def quantum_fisher_information(
333 state_fn,
334 params: jnp.ndarray,
335) -> jnp.ndarray:
336 r"""Compute the Quantum Fisher Information (QFI) at a parameter point.
338 The QFI is the metric tensor of the state manifold evaluated at
339 ``params``. It therefore requires the state as a *function* of the
340 parameters rather than a single state; the Jacobian is obtained with
341 forward-mode automatic differentiation (:func:`jax.jacfwd`), which yields
342 the complex Jacobian directly for real-valued parameters.
344 Both pure and mixed states are supported and dispatched on the kind of
345 object returned by *state_fn* (state vector vs. density matrix), mirroring
346 :func:`fidelity`:
348 - state vector of shape ``(d,)`` -> Fubini-Study formula
349 (see :func:`_qfi_statevector`),
350 - density matrix of shape ``(d, d)`` -> symmetric logarithmic derivative
351 formula (see :func:`_qfi_density`).
353 The returned matrix has shape ``(P, P)`` where ``P`` is the total number of
354 parameters (the parameter axes of *params* are flattened).
356 Args:
357 state_fn: Callable mapping *params* to a normalised quantum state.
358 Typically ``lambda p: model(params=p, inputs=x)`` with the model's
359 ``execution_type`` set to ``"state"`` (pure) or ``"density"``
360 (mixed).
361 params: Parameters at which the QFI is evaluated. Must be passed in the
362 shape expected by *state_fn* (e.g. the model's batched
363 ``model.params``).
365 Returns:
366 Real, symmetric QFI matrix of shape ``(P, P)``.
368 Raises:
369 ValueError: If *state_fn* returns neither a state vector nor a square
370 density matrix.
371 """
372 state, jac = _state_and_jacobian(state_fn, params)
374 if state.ndim == 1:
375 jac = jac.reshape(state.shape[0], -1)
376 return _qfi_statevector(jac, state)
377 elif state.ndim == 2 and state.shape[-1] == state.shape[-2]:
378 jac = jac.reshape(state.shape[0], state.shape[1], -1)
379 return _qfi_density(jac, state)
380 else:
381 raise ValueError(
382 "state_fn must return a state vector of shape (d,) or a density "
383 f"matrix of shape (d, d), got shape {state.shape}."
384 )
387def fubini_study_metric(
388 state_fn,
389 params: jnp.ndarray,
390) -> jnp.ndarray:
391 r"""Compute the Fubini-Study metric tensor at a parameter point.
393 The Fubini-Study metric is the real part of the quantum geometric tensor on
394 the manifold of pure states and equals the pure-state quantum Fisher
395 information up to a factor of four, :math:`F_{ij} = 4\,g_{ij}`:
397 .. math::
399 g_{ij} = \mathrm{Re}\left[
400 \braket{\partial_i\psi | \partial_j\psi}
401 - \braket{\partial_i\psi | \psi}\braket{\psi | \partial_j\psi}
402 \right]
404 It is only defined for pure states; *state_fn* must therefore return a
405 normalised state vector. See :func:`quantum_fisher_information` for the
406 calling convention.
408 Args:
409 state_fn: Callable mapping *params* to a normalised state vector.
410 Typically ``lambda p: model(params=p, inputs=x)`` with the model's
411 ``execution_type`` set to ``"state"``.
412 params: Parameters at which the metric is evaluated.
414 Returns:
415 Real, symmetric metric of shape ``(P, P)`` where ``P`` is the total
416 number of parameters.
418 Raises:
419 ValueError: If *state_fn* does not return a state vector.
420 """
421 state, jac = _state_and_jacobian(state_fn, params)
423 if state.ndim != 1:
424 raise ValueError(
425 "The Fubini-Study metric is only defined for pure states; "
426 f"state_fn must return a state vector of shape (d,), got shape "
427 f"{state.shape}."
428 )
430 jac = jac.reshape(state.shape[0], -1)
431 return _fubini_study_statevector(jac, state)