Coverage for jaxquantum/core/qp_distributions.py: 95%
108 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1import jax.numpy as jnp
2from jax import config, lax, vmap
3from jax.scipy.special import factorial
4import jax
6config.update("jax_enable_x64", True)
8def wigner(psi, xvec, yvec, method="clenshaw", g=2):
9 """Wigner function for a state vector or density matrix at points
10 `xvec + i * yvec`.
12 Parameters
13 ----------
15 state : Qarray
16 A state vector or density matrix.
18 xvec : array_like
19 x-coordinates at which to calculate the Wigner function.
21 yvec : array_like
22 y-coordinates at which to calculate the Wigner function.
24 g : float, default: 2
25 Scaling factor for `a = 0.5 * g * (x + iy)`, default `g = 2`.
26 The value of `g` is related to the value of `hbar` in the commutation
27 relation `[x, y] = i * hbar` via `hbar=2/g^2`.
29 method : string {'clenshaw', 'iterative', 'laguerre', 'fft'}, default: 'clenshaw'
30 Only 'clenshaw' is currently supported.
31 Select method 'clenshaw' 'iterative', 'laguerre', or 'fft', where 'clenshaw'
32 and 'iterative' use an iterative method to evaluate the Wigner functions for density
33 matrices :math:`|m><n|`, while 'laguerre' uses the Laguerre polynomials
34 in scipy for the same task. The 'fft' method evaluates the Fourier
35 transform of the density matrix. The 'iterative' method is default, and
36 in general recommended, but the 'laguerre' method is more efficient for
37 very sparse density matrices (e.g., superpositions of Fock states in a
38 large Hilbert space). The 'clenshaw' method is the preferred method for
39 dealing with density matrices that have a large number of excitations
40 (>~50). 'clenshaw' is a fast and numerically stable method.
42 Returns
43 -------
45 W : array
46 Values representing the Wigner function calculated over the specified
47 range [xvec,yvec].
50 References
51 ----------
53 Ulf Leonhardt,
54 Measuring the Quantum State of Light, (Cambridge University Press, 1997)
56 """
58 if not (psi.is_vec() or psi.is_dm()):
59 raise TypeError("Input state is not a valid operator.")
61 if method == "fft":
62 raise NotImplementedError("Only the 'clenshaw' method is implemented.")
64 if method == "iterative":
65 raise NotImplementedError("Only the 'clenshaw' method is implemented.")
67 elif method == "laguerre":
68 raise NotImplementedError("Only the 'clenshaw' method is implemented.")
70 elif method == "clenshaw":
71 rho = psi.to_dm()
72 bdims = rho.bdims
73 rho = rho.data
75 vmapped_wigner_clenshaw = [_wigner_clenshaw]
77 for _ in bdims:
78 vmapped_wigner_clenshaw.append(
79 vmap(
80 vmapped_wigner_clenshaw[-1],
81 in_axes=(0, None, None, None),
82 out_axes=0,
83 )
84 )
85 return vmapped_wigner_clenshaw[-1](rho, xvec, yvec, g)
87 else:
88 raise TypeError("method must be either 'iterative', 'laguerre', or 'fft'.")
91def _wigner_clenshaw(rho, xvec, yvec, g):
92 r"""
93 Using Clenshaw summation - numerically stable and efficient
94 iterative algorithm to evaluate polynomial series.
96 The Wigner function is calculated as
97 :math:`W = e^(-0.5*x^2)/pi * \sum_{L} c_L (2x)^L / \sqrt(L!)` where
98 :math:`c_L = \sum_n \rho_{n,L+n} LL_n^L` where
99 :math:`LL_n^L = (-1)^n \sqrt(L!n!/(L+n)!) LaguerreL[n,L,x]`
100 Heavily inspired by Qutip and Dynamiqs
101 https://github.com/dynamiqs/dynamiqs
102 https://github.com/qutip/qutip
103 """
105 M = jnp.prod(rho.shape[0])
106 X, Y = jnp.meshgrid(xvec, yvec)
107 A = 0.5 * g * (X + 1.0j * Y)
108 B = jnp.abs(2*A)
110 B *= B
112 w0 = (2 * rho[0, -1]) * jnp.ones_like(A)
114 # calculation of \sum_{L} c_L (2x)^L / \sqrt(L!)
115 # using Horner's method
117 rho = rho * (2 * jnp.ones((M, M)) - jnp.diag(jnp.ones(M)))
118 def loop(i: int, w: jax.Array) -> jax.Array:
119 i = M - 2 - i
120 w = w * (2 * A * (i + 1) ** (-0.5))
121 return w + _wig_laguerre_val(i, B, rho, M)
123 w = jax.lax.fori_loop(0, M - 1, loop, w0)
125 return w.real * jnp.exp(-B * 0.5) * (g * g * 0.5 / jnp.pi)
127def _extract_diag_element(rho: jnp.array, L: int, n:int):
128 """"
129 Extract element at index n from diagonal L of matrix rho.
130 Heavily inspired from https://github.com/dynamiqs/dynamiqs
131 """
132 N = rho.shape[0]
133 n = jax.lax.select(n < 0, N - jnp.abs(L) - jnp.abs(n), n)
134 row = jnp.maximum(-L, 0) + n
135 col = jnp.maximum(L, 0) + n
136 return rho[row, col]
138def _wig_laguerre_val(L, x, rho, N):
139 r"""
140 Evaluate Laguerre polynomials.
141 Implementation in Jax from https://github.com/dynamiqs/dynamiqs
142 """
144 def len_c_1():
145 return _extract_diag_element(rho, L, 0) * jnp.ones_like(x)
147 def len_c_2():
148 c0 = _extract_diag_element(rho, L, 0)
149 c1 = _extract_diag_element(rho, L, 1)
150 return (c0 - c1 * (L + 1 - x) * (L + 1) ** (-0.5)) * jnp.ones_like(x)
152 def len_c_other():
153 cm2 = _extract_diag_element(rho, L, -2)
154 cm1 = _extract_diag_element(rho, L, -1)
155 y0 = cm2 * jnp.ones_like(x)
156 y1 = cm1 * jnp.ones_like(x)
158 def loop(j: int, args: tuple[jax.Array, jax.Array]) -> tuple[
159 jax.Array, jax.Array]:
160 def body() -> tuple[jax.Array, jax.Array]:
161 k = N + 1 - L - j
162 y0, y1 = args
163 ckm1 = _extract_diag_element(rho, L, -j)
164 y0, y1 = (
165 ckm1 - y1 * (k * (L + k) / ((L + k + 1) * (k + 1))) ** 0.5,
166 y0 - y1 * (L + 2 * k - x + 1) * (
167 (L + k + 1) * (k + 1)) ** -0.5,
168 )
170 return y0, y1
172 return jax.lax.cond(j >= N + 1 - L, lambda: args, body)
174 y0, y1 = jax.lax.fori_loop(3, N + 1, loop, (y0, y1))
176 return y0 - y1 * (L + 1 - x) * (L + 1) ** (-0.5)
179 return jax.lax.cond(N - L == 1, len_c_1, lambda: jax.lax.cond(N - L == 2,
180 len_c_2,
181 len_c_other))
184def qfunc(psi, xvec, yvec, g=2):
185 r"""
186 Husimi-Q function of a given state vector or density matrix at phase-space
187 points ``0.5 * g * (xvec + i*yvec)``.
189 Parameters
190 ----------
191 state : Qarray
192 A state vector or density matrix. This cannot have tensor-product
193 structure.
195 xvec, yvec : array_like
196 x- and y-coordinates at which to calculate the Husimi-Q function.
198 g : float, default: 2
199 Scaling factor for ``a = 0.5 * g * (x + iy)``. The value of `g` is
200 related to the value of :math:`\hbar` in the commutation relation
201 :math:`[x,\,y] = i\hbar` via :math:`\hbar=2/g^2`.
203 Returns
204 -------
205 jnp.ndarray
206 Values representing the Husimi-Q function calculated over the specified
207 range ``[xvec, yvec]``.
209 """
211 alpha_grid, prefactor = _qfunc_coherent_grid(xvec, yvec, g)
213 if psi.is_vec():
214 psi = psi.to_ket()
216 def _compute_qfunc(psi, alpha_grid, prefactor, g):
217 out = _qfunc_iterative_single(psi, alpha_grid, prefactor, g)
218 out /= jnp.pi
219 return out
220 else:
222 def _compute_qfunc(psi, alpha_grid, prefactor, g):
223 values, vectors = jnp.linalg.eigh(psi)
224 vectors = vectors.T
225 chunk_size = min(8, values.shape[0])
226 chunks = (values.shape[0] + chunk_size - 1) // chunk_size
227 padding = chunks * chunk_size - values.shape[0]
228 values = jnp.pad(values, (0, padding))
229 vectors = jnp.pad(vectors, ((0, padding), (0, 0)))
231 def add_chunk(index, total):
232 start = index * chunk_size
233 chunk_values = lax.dynamic_slice_in_dim(
234 values, start, chunk_size
235 )
236 chunk_vectors = lax.dynamic_slice_in_dim(
237 vectors, start, chunk_size
238 )
239 components = vmap(
240 lambda vector: _qfunc_iterative_single(
241 vector, alpha_grid, prefactor, g
242 )
243 )(chunk_vectors)
244 return total + jnp.tensordot(chunk_values, components, axes=1)
246 out = lax.fori_loop(
247 0,
248 chunks,
249 add_chunk,
250 jnp.zeros_like(alpha_grid.real),
251 )
252 out /= jnp.pi
254 return out
256 # Use the Qarray batch dims rather than data.shape[:-2]: kets now have a
257 # single trailing space axis, so shape[:-2] would under-count batch dims.
258 bdims = psi.bdims
259 psi = psi.data
261 vmapped_compute_qfunc = [_compute_qfunc]
263 for _ in bdims:
264 vmapped_compute_qfunc.append(
265 vmap(
266 vmapped_compute_qfunc[-1],
267 in_axes=(0, None, None, None),
268 out_axes=0,
269 )
270 )
271 return vmapped_compute_qfunc[-1](psi, alpha_grid, prefactor, g)
274def _qfunc_iterative_single(
275 vector,
276 grid,
277 prefactor,
278 g,
279):
280 r"""
281 Get the Q function (without the :math:`\pi` scaling factor) of a single
282 state vector, using the iterative algorithm which recomputes the powers of
283 the coherent-state matrix.
284 """
285 vector = vector.squeeze()
286 ns = jnp.arange(vector.shape[-1])
287 out = jnp.polyval(
288 (vector / jnp.sqrt(factorial(ns)))[::-1],
289 grid,
290 )
291 out *= prefactor
292 return jnp.abs(out) ** 2
295def _qfunc_coherent_grid(xvec, yvec, g):
296 x, y = jnp.meshgrid(0.5 * g * xvec, 0.5 * g * yvec)
297 # We produce the adjoint of the coherent states to save an operation
298 # later when computing dot products, hence the negative imaginary part.
299 grid = x - 1.0j * y
300 prefactor = jnp.exp(-0.5 * (x * x + y * y)).astype(jnp.complex128)
301 return grid, prefactor