Coverage for qml_essentials / random_sampling.py: 96%
82 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-08-18 14:58 +0000
1import jax
2import jax.numpy as jnp
3from typing import Optional
4from qml_essentials.operations import _cdtype
6_DEFAULT_SEED = 1000
9def _ginibre(key: jax.random.PRNGKey, rows: int, cols: int) -> jnp.ndarray:
10 r"""Draw a complex Ginibre matrix :math:`G \in \mathbb{C}^{rows \times cols}`.
12 Each entry is :math:`G_{jk} = a_{jk} + i\,b_{jk}` with
13 :math:`a_{jk}, b_{jk} \sim \mathcal{N}(0, 1)` i.i.d. standard normal (the
14 Ginibre ensemble). The result is cast to the active complex dtype via
15 :func:`qml_essentials.operations._cdtype`.
17 Args:
18 key (jax.random.PRNGKey): JAX random key, split internally into a real
19 and an imaginary part.
20 rows (int): Number of rows :math:`d`.
21 cols (int): Number of columns :math:`K`.
23 Returns:
24 jnp.ndarray: Complex Ginibre matrix of shape ``(rows, cols)``.
25 """
26 key_re, key_im = jax.random.split(key)
27 real = jax.random.normal(key_re, shape=(rows, cols))
28 imag = jax.random.normal(key_im, shape=(rows, cols))
29 return (real + 1j * imag).astype(_cdtype())
32def _haar_unitary(key: jax.random.PRNGKey, dim: int) -> jnp.ndarray:
33 r"""Draw a Haar-random unitary :math:`U \in U(d)`.
35 Implements the QR construction of Mezzadri (*How to generate random
36 matrices from the classical compact groups*, arXiv:math-ph/0609050): draw a
37 Ginibre matrix :math:`Z`, take its QR decomposition :math:`Z = QR`, and
38 remove the phase ambiguity via
39 :math:`U = Q\,\Lambda`, :math:`\Lambda = \mathrm{diag}(R_{ii}/|R_{ii}|)`,
40 so that :math:`U` is Haar-distributed.
42 Args:
43 key (jax.random.PRNGKey): JAX random key for the Ginibre draw.
44 dim (int): Dimension :math:`d` of the unitary.
46 Returns:
47 jnp.ndarray: Haar-random unitary of shape ``(dim, dim)``.
48 """
49 z = _ginibre(key, dim, dim)
50 q, r = jnp.linalg.qr(z)
51 diag_r = jnp.diagonal(r)
52 # Guard against a zero pivot (|0| -> 1 keeps the phase at 1).
53 phases = diag_r / jnp.where(jnp.abs(diag_r) > 0, jnp.abs(diag_r), 1.0)
54 return q * phases[None, :]
57def _normalize_density(rho: jnp.ndarray) -> jnp.ndarray:
58 r"""Hermitize and trace-normalize a single density matrix.
60 Applies :math:`\rho \leftarrow (\rho + \rho^\dagger)/2` to remove numerical
61 anti-Hermitian noise, then divides by :math:`\mathrm{Tr}\,\rho` so that
62 :math:`\mathrm{Tr}\,\rho = 1`.
64 Args:
65 rho (jnp.ndarray): Unnormalized matrix of shape ``(d, d)``.
67 Returns:
68 jnp.ndarray: Hermitian, unit-trace density matrix of shape ``(d, d)``.
69 """
70 rho = (rho + jnp.conj(rho).T) / 2.0
71 return rho / jnp.trace(rho)
74class DensityMatrix:
75 r"""Random density-matrix samplers.
77 Samplers for the standard random-density-matrix ensembles, built from the
78 Ginibre construction
79 :math:`\rho = G G^\dagger / \mathrm{Tr}(G G^\dagger)` and/or a Haar-random
80 unitary. Every sampler returns a batched array of shape
81 ``(n_samples, 2**n_qubits, 2**n_qubits)``, consistent with
82 :func:`qml_essentials.entanglement.sample_random_separable_states`.
84 The module is structured so that ``StateVector``, ``Hermitian`` and
85 ``Unitary`` sampler classes can be added later, reusing the module-level
86 primitives :func:`_ginibre` and :func:`_haar_unitary`.
87 """
89 @classmethod
90 def induced(
91 cls,
92 n_qubits: int,
93 n_samples: int = 1,
94 rank: Optional[int] = None,
95 random_key: Optional[jax.random.PRNGKey] = None,
96 ) -> jnp.ndarray:
97 r"""Sample from the induced measure via the Ginibre construction.
99 Draws :math:`\rho = G G^\dagger / \mathrm{Tr}(G G^\dagger)` with
100 :math:`G` a :math:`d \times K` complex Ginibre matrix,
101 :math:`d = 2^{n\_qubits}` and :math:`K = \mathrm{rank}`. This is the
102 induced measure of Zyczkowski & Sommers (*Induced measures in the space
103 of mixed quantum states*, arXiv:quant-ph/0012101). The sampled state
104 has rank :math:`\min(d, K)` almost surely: :math:`K = 1` yields (Haar)
105 pure states and :math:`K = d` recovers the Hilbert-Schmidt measure. The
106 mean purity is :math:`\mathbb{E}[\mathrm{Tr}\,\rho^2] = (d + K)/(dK + 1)`.
108 Args:
109 n_qubits (int): Number of qubits; :math:`d = 2^{n\_qubits}`.
110 n_samples (int): Number of density matrices to draw. Defaults to 1.
111 rank (Optional[int]): Number of Ginibre columns :math:`K`. When
112 ``None`` (default) it is set to :math:`d`, recovering the
113 Hilbert-Schmidt measure. Must satisfy ``rank >= 1``.
114 random_key (Optional[jax.random.PRNGKey]): JAX random key. When
115 ``None``, falls back to ``jax.random.key(1000)`` (matching the
116 default ``random_seed`` of :class:`~qml_essentials.model.Model`).
118 Returns:
119 jnp.ndarray: Density matrices of shape
120 ``(n_samples, 2**n_qubits, 2**n_qubits)``.
122 Raises:
123 ValueError: If ``rank`` is given and ``rank < 1``.
124 """
125 d = 2**n_qubits
126 if rank is None:
127 rank = d
128 if rank < 1:
129 raise ValueError(f"rank must be >= 1, got {rank}.")
130 if random_key is None:
131 random_key = jax.random.key(_DEFAULT_SEED)
133 def _sample(key: jax.random.PRNGKey) -> jnp.ndarray:
134 g = _ginibre(key, d, rank)
135 rho = g @ jnp.conj(g).T
136 return _normalize_density(rho)
138 keys = jax.random.split(random_key, n_samples)
139 return jax.vmap(_sample)(keys)
141 @classmethod
142 def hilbert_schmidt(
143 cls,
144 n_qubits: int,
145 n_samples: int = 1,
146 random_key: Optional[jax.random.PRNGKey] = None,
147 ) -> jnp.ndarray:
148 r"""Sample from the Hilbert-Schmidt measure.
150 Special case of the induced measure with :math:`K = d`:
151 :math:`\rho = G G^\dagger / \mathrm{Tr}(G G^\dagger)` with :math:`G` a
152 square :math:`d \times d` complex Ginibre matrix. The Hilbert-Schmidt
153 measure is the flat measure induced by the Hilbert-Schmidt metric
154 (Zyczkowski & Sommers, arXiv:quant-ph/0012101). Delegates to
155 :meth:`induced` with ``rank = d``.
157 Args:
158 n_qubits (int): Number of qubits; :math:`d = 2^{n\_qubits}`.
159 n_samples (int): Number of density matrices to draw. Defaults to 1.
160 random_key (Optional[jax.random.PRNGKey]): JAX random key. When
161 ``None``, falls back to ``jax.random.key(1000)``.
163 Returns:
164 jnp.ndarray: Density matrices of shape
165 ``(n_samples, 2**n_qubits, 2**n_qubits)``.
166 """
167 return cls.induced(
168 n_qubits=n_qubits,
169 n_samples=n_samples,
170 rank=2**n_qubits,
171 random_key=random_key,
172 )
174 @classmethod
175 def bures(
176 cls,
177 n_qubits: int,
178 n_samples: int = 1,
179 random_key: Optional[jax.random.PRNGKey] = None,
180 ) -> jnp.ndarray:
181 r"""Sample from the Bures measure.
183 Uses the construction of Osipov, Sommers & Zyczkowski (*Random Bures
184 mixed states and the distribution of their purity*, arXiv:1004.1655):
186 .. math::
187 \rho = \frac{(\mathbb{1} + U)\, G G^\dagger\, (\mathbb{1} + U)^\dagger}
188 {\mathrm{Tr}\!\left[(\mathbb{1} + U)\, G G^\dagger\,
189 (\mathbb{1} + U)^\dagger\right]},
191 where :math:`G` is a square :math:`d \times d` complex Ginibre matrix
192 and :math:`U` is an independently drawn Haar-random unitary. The Bures
193 measure is induced by the Bures (statistical-distance) metric.
195 Args:
196 n_qubits (int): Number of qubits; :math:`d = 2^{n\_qubits}`.
197 n_samples (int): Number of density matrices to draw. Defaults to 1.
198 random_key (Optional[jax.random.PRNGKey]): JAX random key. When
199 ``None``, falls back to ``jax.random.key(1000)``.
201 Returns:
202 jnp.ndarray: Density matrices of shape
203 ``(n_samples, 2**n_qubits, 2**n_qubits)``.
204 """
205 d = 2**n_qubits
206 if random_key is None:
207 random_key = jax.random.key(_DEFAULT_SEED)
209 eye = jnp.eye(d, dtype=_cdtype())
211 def _sample(key: jax.random.PRNGKey) -> jnp.ndarray:
212 key_g, key_u = jax.random.split(key)
213 g = _ginibre(key_g, d, d)
214 u = _haar_unitary(key_u, d)
215 a = eye + u
216 rho = a @ (g @ jnp.conj(g).T) @ jnp.conj(a).T
217 return _normalize_density(rho)
219 keys = jax.random.split(random_key, n_samples)
220 return jax.vmap(_sample)(keys)
222 @classmethod
223 def eigen(
224 cls,
225 n_qubits: int,
226 n_samples: int = 1,
227 alpha: float = 1.0,
228 eigenvalues: Optional[jnp.ndarray] = None,
229 random_key: Optional[jax.random.PRNGKey] = None,
230 ) -> jnp.ndarray:
231 r"""Sample density matrices with a prescribed or Dirichlet spectrum.
233 Builds :math:`\rho = U\,\mathrm{diag}(\lambda)\,U^\dagger` with :math:`U`
234 Haar-random. The eigenvalue vector :math:`\lambda` is either supplied
235 via ``eigenvalues`` (then identical for every sample), or drawn per
236 sample from a symmetric Dirichlet distribution
237 :math:`\lambda \sim \mathrm{Dir}(\alpha \mathbf{1}_d)`. The default
238 :math:`\alpha = 1` gives the uniform (flat) distribution on the
239 probability simplex.
241 Note that this ensemble is, by construction, different from the
242 Hilbert-Schmidt and induced measures: it lacks their eigenvalue
243 repulsion (the squared Vandermonde factor).
245 Args:
246 n_qubits (int): Number of qubits; :math:`d = 2^{n\_qubits}`.
247 n_samples (int): Number of density matrices to draw. Defaults to 1.
248 alpha (float): Symmetric Dirichlet concentration parameter, used
249 only when ``eigenvalues`` is ``None``. Defaults to ``1.0``.
250 eigenvalues (Optional[jnp.ndarray]): Optional fixed spectrum of
251 length :math:`d`. Entries must be nonnegative and sum to 1. When
252 given, the same spectrum is used for every sample and only
253 :math:`U` is randomized.
254 random_key (Optional[jax.random.PRNGKey]): JAX random key. When
255 ``None``, falls back to ``jax.random.key(1000)``.
257 Returns:
258 jnp.ndarray: Density matrices of shape
259 ``(n_samples, 2**n_qubits, 2**n_qubits)``.
261 Raises:
262 ValueError: If ``eigenvalues`` is supplied and is not a
263 length-:math:`d` vector of nonnegative entries summing to 1.
264 """
265 d = 2**n_qubits
266 if random_key is None:
267 random_key = jax.random.key(_DEFAULT_SEED)
269 fixed_eigs = None
270 if eigenvalues is not None:
271 fixed_eigs = cls._validate_eigenvalues(eigenvalues, d)
273 def _sample(key: jax.random.PRNGKey) -> jnp.ndarray:
274 key_u, key_lam = jax.random.split(key)
275 u = _haar_unitary(key_u, d)
276 if fixed_eigs is None:
277 lam = jax.random.dirichlet(key_lam, alpha * jnp.ones(d))
278 else:
279 lam = fixed_eigs
280 rho = (u * lam[None, :].astype(_cdtype())) @ jnp.conj(u).T
281 return _normalize_density(rho)
283 keys = jax.random.split(random_key, n_samples)
284 return jax.vmap(_sample)(keys)
286 @classmethod
287 def _validate_eigenvalues(cls, eigenvalues: jnp.ndarray, dim: int) -> jnp.ndarray:
288 r"""Validate a user-supplied eigenvalue (spectrum) vector.
290 Checks that the vector has length :math:`d`, is real and nonnegative,
291 and sums to 1 (within tolerance).
293 Args:
294 eigenvalues (jnp.ndarray): Candidate spectrum.
295 dim (int): Expected length :math:`d = 2^{n\_qubits}`.
297 Returns:
298 jnp.ndarray: Validated, real eigenvalue vector of shape ``(dim,)``.
300 Raises:
301 ValueError: If the shape, nonnegativity, or unit-sum constraint is
302 violated.
303 """
304 eigs = jnp.asarray(eigenvalues)
305 if eigs.shape != (dim,):
306 raise ValueError(f"eigenvalues must have shape ({dim},), got {eigs.shape}.")
307 eigs = jnp.real(eigs)
308 if bool(jnp.any(eigs < -1e-10)):
309 raise ValueError("eigenvalues must be nonnegative.")
310 if not bool(jnp.isclose(jnp.sum(eigs), 1.0, atol=1e-8)):
311 raise ValueError(f"eigenvalues must sum to 1, got {float(jnp.sum(eigs))}.")
312 return eigs