Coverage for jaxquantum/core/visualization.py: 48%

1""" 

2Visualization utils. 

3""" 

4 

5import matplotlib.pyplot as plt 

6from matplotlib.animation import FuncAnimation, PillowWriter 

7 

8from jaxquantum.core.qp_distributions import wigner, qfunc 

9from jaxquantum.core.cfunctions import cf_wigner 

10import jax.numpy as jnp 

11import numpy as np 

12 

13WIGNER = "wigner" 

14HUSIMI = "husimi" 

15 

16 

17def _render_qp_grid( 

18 axs, 

19 QP, 

20 pts_x, 

21 pts_y, 

22 *, 

23 contour, 

24 cmap, 

25 vmin, 

26 vmax, 

27 x_ticks, 

28 y_ticks, 

29 z_ticks, 

30 cbar_label, 

31 plot_cbar, 

32 subtitles, 

33 decorate=True, 

34): 

35 """Render one quasi-probability frame onto a ``(rows, cols)`` axes grid. 

36 

37 ``QP`` has shape ``(rows, cols, len(pts_y), len(pts_x))``. Used by both 

38 the static ``plot_qp`` path (called once, ``decorate=True``) and the gif 

39 path (called once per frame; ``decorate=True`` on frame 0 to lay out 

40 ticks, gridlines, axhline/axvline, labels, and colorbars, then 

41 ``decorate=False`` thereafter so those non-idempotent artists aren't 

42 duplicated as frames advance). 

43 

44 Returns the last ``contourf`` / ``pcolormesh`` artist created. 

45 """ 

46 rows, cols = QP.shape[0], QP.shape[1] 

47 im = None 

48 for row in range(rows): 

49 for col in range(cols): 

50 ax = axs[row, col] 

51 if contour: 

52 im = ax.contourf( 

53 pts_x, 

54 pts_y, 

55 QP[row, col], 

56 cmap=cmap, 

57 vmin=vmin, 

58 vmax=vmax, 

59 levels=np.linspace(vmin, vmax, 101), 

60 ) 

61 else: 

62 im = ax.pcolormesh( 

63 pts_x, 

64 pts_y, 

65 QP[row, col], 

66 cmap=cmap, 

67 vmin=vmin, 

68 vmax=vmax, 

69 ) 

70 if decorate: 

71 ax.set_xticks(x_ticks) 

72 ax.set_yticks(y_ticks) 

73 ax.axhline(0, linestyle="-", color="black", alpha=0.7) 

74 ax.axvline(0, linestyle="-", color="black", alpha=0.7) 

75 ax.grid() 

76 ax.set_aspect("equal", adjustable="box") 

77 

78 if plot_cbar: 

79 cbar = plt.colorbar( 

80 im, ax=ax, orientation="vertical", ticks=np.linspace(-1, 1, 11) 

81 ) 

82 cbar.ax.set_title(cbar_label) 

83 cbar.set_ticks(z_ticks) 

84 

85 ax.set_xlabel(r"Re[$\alpha$]") 

86 ax.set_ylabel(r"Im[$\alpha$]") 

87 if subtitles is not None: 

88 ax.set_title(subtitles[row, col]) 

89 return im 

90 

91 

92def plot_qp( 

93 state, 

94 pts_x, 

95 pts_y=None, 

96 g=2, 

97 axs=None, 

98 contour=True, 

99 qp_type=WIGNER, 

100 cbar_label="", 

101 axis_scale_factor=1, 

102 plot_cbar=True, 

103 x_ticks=None, 

104 y_ticks=None, 

105 z_ticks=None, 

106 subtitles=None, 

107 figtitle=None, 

108 gif=False, 

109 gif_params=None, 

110): 

111 """Plot a quasi-probability distribution (Wigner or Husimi-Q). 

112 

113 The state may carry an arbitrary number of batch dimensions; they are 

114 flattened to a 2D ``(rows, cols)`` grid of subplots. With ``gif=True``, 

115 one batch axis is animated instead and the remaining batch dims form 

116 the per-frame subplot grid. 

117 

118 Args: 

119 state: state with arbitrary number of batch dimensions; result will 

120 be flattened to a 2d grid to allow for plotting 

121 pts_x: x points to evaluate the quasi-probability distribution at 

122 pts_y: y points to evaluate the quasi-probability distribution at; 

123 defaults to ``pts_x`` 

124 g: float, default 2. Scaling factor for ``a = 0.5 * g * (x + iy)``. 

125 The value of ``g`` is related to the value of :math:`\\hbar` in 

126 the commutation relation :math:`[x,\,y] = i\\hbar` via 

127 :math:`\\hbar=2/g^2`. 

128 axs: matplotlib axes to plot on (created if None) 

129 contour: use ``contourf`` if True, otherwise ``pcolormesh`` 

130 qp_type: type of quasi-probability distribution 

131 (``"wigner"`` or ``"husimi"``) 

132 cbar_label: label for the cbar (overridden internally based on 

133 ``qp_type``) 

134 axis_scale_factor: multiplicative scale applied to the axis tick 

135 positions and labels 

136 plot_cbar: whether to draw a colorbar on each subplot 

137 x_ticks: tick positions for the x-axis (auto if None) 

138 y_ticks: tick positions for the y-axis (auto if None) 

139 z_ticks: tick positions for the colorbar (auto if None) 

140 subtitles: subtitles for the subplots; shape must match 

141 ``state.bdims`` (or the per-frame batch dims when ``gif=True``) 

142 figtitle: figure title 

143 gif: if True, render an animation over one batch axis instead of a 

144 tiled subplot grid. Returns a 

145 ``matplotlib.animation.FuncAnimation`` that auto-renders inline 

146 in Jupyter (its ``_repr_html_`` is patched to ``to_jshtml``). 

147 gif_params: dict of options for the gif path (ignored if 

148 ``gif=False``). Recognized keys: 

149 

150 - ``save_path`` (default ``None``) — if set, save the animation 

151 to this path via ``matplotlib.animation.PillowWriter``. 

152 - ``interval_ms`` (default ``200``) — milliseconds per frame; 

153 also derives ``fps = round(1000 / interval_ms)`` for the writer. 

154 - ``ts`` (default ``None``) — optional 1D array of timestamps 

155 matching the animation-axis length; when set, each frame's 

156 suptitle gets a ``t = …`` label. 

157 - ``batch_animation_axis`` (default ``0``) — index into 

158 ``state.bdims`` selecting which axis becomes the animation 

159 axis. The remaining batch dims form the per-frame subplot grid. 

160 

161 Returns: 

162 ``(axs, im)`` in the static case, or a ``FuncAnimation`` when 

163 ``gif=True``. 

164 """ 

165 if pts_y is None: 

166 pts_y = pts_x 

167 pts_x = jnp.array(pts_x) 

168 pts_y = jnp.array(pts_y) 

169 

170 if len(state.bdims)==1 and state.bdims[0]==1: 

171 state = state[0] 

172 

173 if gif: 

174 return _plot_qp_gif( 

175 state=state, 

176 pts_x=pts_x, 

177 pts_y=pts_y, 

178 g=g, 

179 axs=axs, 

180 contour=contour, 

181 qp_type=qp_type, 

182 axis_scale_factor=axis_scale_factor, 

183 plot_cbar=plot_cbar, 

184 x_ticks=x_ticks, 

185 y_ticks=y_ticks, 

186 z_ticks=z_ticks, 

187 subtitles=subtitles, 

188 figtitle=figtitle, 

189 gif_params=gif_params or {}, 

190 ) 

191 

192 bdims = state.bdims 

193 added_baxes = 0 

194 

195 if subtitles is not None and subtitles.shape != bdims: 

196 raise ValueError( 

197 f"labels must have same shape as bdims, " 

198 f"got shapes {subtitles.shape} and {bdims}" 

199 ) 

200 

201 if len(bdims) == 0: 

202 bdims = (1,) 

203 added_baxes += 1 

204 if len(bdims) == 1: 

205 bdims = (1, bdims[0]) 

206 added_baxes += 1 

207 

208 extra_dims = bdims[2:] 

209 if extra_dims != (): 

210 state = state.reshape_bdims( 

211 bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1] 

212 ) 

213 if subtitles is not None: 

214 subtitles = subtitles.reshape( 

215 bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1] 

216 ) 

217 bdims = state.bdims 

218 

219 if axs is None: 

220 _, axs = plt.subplots( 

221 bdims[0], 

222 bdims[1], 

223 figsize=(3.3 * bdims[1], 3 * bdims[0]), 

224 dpi=200, 

225 ) 

226 

227 if qp_type == WIGNER: 

228 vmin = -1 

229 vmax = 1 

230 scale = np.pi / 2 

231 cmap = "seismic" 

232 cbar_label = r"$\mathcal{W}(\alpha)$" 

233 QP = scale * wigner(state, pts_x, pts_y, g=g) 

234 

235 elif qp_type == HUSIMI: 

236 vmin = 0 

237 vmax = 1 

238 scale = np.pi 

239 cmap = "jet" 

240 cbar_label = r"$\mathcal{Q}(\alpha)$" 

241 QP = scale * qfunc(state, pts_x, pts_y, g=g) 

242 

243 

244 

245 for _ in range(added_baxes): 

246 QP = jnp.array([QP]) 

247 axs = np.array([axs]) 

248 if subtitles is not None: 

249 subtitles = np.array([subtitles]) 

250 

251 

252 

253 

254 pts_x = pts_x * axis_scale_factor 

255 pts_y = pts_y * axis_scale_factor 

256 

257 x_ticks = ( 

258 jnp.linspace(jnp.min(pts_x), jnp.max(pts_x), 5) if x_ticks is None else x_ticks 

259 ) 

260 y_ticks = ( 

261 jnp.linspace(jnp.min(pts_y), jnp.max(pts_y), 5) if y_ticks is None else y_ticks 

262 ) 

263 z_ticks = jnp.linspace(vmin, vmax, 3) if z_ticks is None else z_ticks 

264 

265 im = _render_qp_grid( 

266 axs, 

267 QP, 

268 pts_x, 

269 pts_y, 

270 contour=contour, 

271 cmap=cmap, 

272 vmin=vmin, 

273 vmax=vmax, 

274 x_ticks=x_ticks, 

275 y_ticks=y_ticks, 

276 z_ticks=z_ticks, 

277 cbar_label=cbar_label, 

278 plot_cbar=plot_cbar, 

279 subtitles=subtitles, 

280 decorate=True, 

281 ) 

282 

283 fig = axs[bdims[0] - 1, bdims[1] - 1].get_figure() 

284 fig.tight_layout(w_pad=0.3, h_pad=0.3) 

285 if figtitle is not None: 

286 fig.suptitle(figtitle, y=1.04) 

287 return axs, im 

288 

289 

290def _plot_qp_gif( 

291 state, 

292 pts_x, 

293 pts_y, 

294 *, 

295 g, 

296 axs, 

297 contour, 

298 qp_type, 

299 axis_scale_factor, 

300 plot_cbar, 

301 x_ticks, 

302 y_ticks, 

303 z_ticks, 

304 subtitles, 

305 figtitle, 

306 gif_params, 

307): 

308 """Build the ``FuncAnimation`` for ``plot_qp(gif=True)``. 

309 

310 Moves ``state.bdims[batch_animation_axis]`` to the front, tiles the 

311 remaining batch dims as a ``(rows, cols)`` per-frame subplot grid, and 

312 reuses ``_render_qp_grid`` per frame (clearing prior ``contourf`` / 

313 ``pcolormesh`` collections each update so the colorbars laid out on 

314 frame 0 are preserved). 

315 

316 Optionally saves the animation to ``gif_params['save_path']`` via 

317 ``PillowWriter``. Patches ``anim._repr_html_`` to ``anim.to_jshtml`` and 

318 closes the figure so the animation auto-renders inline in Jupyter 

319 without an extra static last-frame image. 

320 """ 

321 save_path = gif_params.get("save_path", None) 

322 interval_ms = gif_params.get("interval_ms", 200) 

323 ts = gif_params.get("ts", None) 

324 batch_animation_axis = gif_params.get("batch_animation_axis", 0) 

325 

326 bdims = state.bdims 

327 if len(bdims) < 1: 

328 raise ValueError( 

329 "gif=True requires the state to have at least one batch dimension" 

330 ) 

331 if not 0 <= batch_animation_axis < len(bdims): 

332 raise ValueError( 

333 f"batch_animation_axis={batch_animation_axis} is out of range " 

334 f"for state.bdims={bdims}" 

335 ) 

336 N = bdims[batch_animation_axis] 

337 if ts is not None and len(ts) != N: 

338 raise ValueError( 

339 f"ts has length {len(ts)} but animation axis has length {N}" 

340 ) 

341 

342 if qp_type == WIGNER: 

343 vmin, vmax, scale = -1, 1, np.pi / 2 

344 cmap = "seismic" 

345 cbar_label = r"$\mathcal{W}(\alpha)$" 

346 QP = scale * wigner(state, pts_x, pts_y, g=g) 

347 elif qp_type == HUSIMI: 

348 vmin, vmax, scale = 0, 1, np.pi 

349 cmap = "jet" 

350 cbar_label = r"$\mathcal{Q}(\alpha)$" 

351 QP = scale * qfunc(state, pts_x, pts_y, g=g) 

352 

353 QP = jnp.moveaxis(QP, batch_animation_axis, 0) 

354 rest_bdims = tuple(d for i, d in enumerate(bdims) if i != batch_animation_axis) 

355 

356 grid_dims = list(rest_bdims) 

357 added_baxes = 0 

358 if len(grid_dims) == 0: 

359 grid_dims = [1] 

360 added_baxes += 1 

361 if len(grid_dims) == 1: 

362 grid_dims = [1, grid_dims[0]] 

363 added_baxes += 1 

364 extras = grid_dims[2:] 

365 rows = grid_dims[0] * int(np.prod(extras)) if extras else grid_dims[0] 

366 cols = grid_dims[1] 

367 

368 h, w = QP.shape[-2], QP.shape[-1] 

369 QP_anim = QP.reshape((N, rows, cols, h, w)) 

370 

371 if subtitles is not None: 

372 subtitles = np.asarray(subtitles) 

373 if subtitles.shape != rest_bdims: 

374 raise ValueError( 

375 f"subtitles shape {subtitles.shape} must match per-frame " 

376 f"batch dims {rest_bdims} (state.bdims minus the animation axis)" 

377 ) 

378 subtitles = subtitles.reshape(rows, cols) 

379 

380 pts_x_scaled = pts_x * axis_scale_factor 

381 pts_y_scaled = pts_y * axis_scale_factor 

382 x_ticks = ( 

383 jnp.linspace(jnp.min(pts_x_scaled), jnp.max(pts_x_scaled), 5) 

384 if x_ticks is None 

385 else x_ticks 

386 ) 

387 y_ticks = ( 

388 jnp.linspace(jnp.min(pts_y_scaled), jnp.max(pts_y_scaled), 5) 

389 if y_ticks is None 

390 else y_ticks 

391 ) 

392 z_ticks = jnp.linspace(vmin, vmax, 3) if z_ticks is None else z_ticks 

393 

394 if axs is None: 

395 _, axs = plt.subplots( 

396 rows, cols, figsize=(3.3 * cols, 3 * rows), dpi=200, 

397 layout="constrained", 

398 ) 

399 axs_arr = axs 

400 for _ in range(added_baxes): 

401 axs_arr = np.array([axs_arr]) 

402 axs_arr = np.asarray(axs_arr).reshape(rows, cols) 

403 fig = axs_arr[0, 0].get_figure() 

404 

405 has_suptitle = figtitle is not None or ts is not None 

406 

407 def _set_suptitle(k): 

408 if ts is not None: 

409 t_str = f"t = {float(ts[k]):.3g}" 

410 title = f"{figtitle} | {t_str}" if figtitle else t_str 

411 fig.suptitle(title) 

412 elif figtitle is not None: 

413 fig.suptitle(figtitle) 

414 

415 _render_qp_grid( 

416 axs_arr, 

417 QP_anim[0], 

418 pts_x_scaled, 

419 pts_y_scaled, 

420 contour=contour, 

421 cmap=cmap, 

422 vmin=vmin, 

423 vmax=vmax, 

424 x_ticks=x_ticks, 

425 y_ticks=y_ticks, 

426 z_ticks=z_ticks, 

427 cbar_label=cbar_label, 

428 plot_cbar=plot_cbar, 

429 subtitles=subtitles, 

430 decorate=True, 

431 ) 

432 _set_suptitle(0) 

433 # PillowWriter saves at the fixed figure bbox (no bbox_inches="tight"), so 

434 # anything spilling past the edges is clipped. tight_layout mis-handles 

435 # equal-aspect axes (the square box is resized at draw time, after layout), 

436 # which pushed the bottom Re[α] label off-frame. constrained_layout instead 

437 # shrinks the axes to fit the suptitle/colorbars/labels, so nothing clips; 

438 # tighten its inter-panel spacing to keep the batched panels close together. 

439 engine = fig.get_layout_engine() 

440 if engine is not None and engine.__class__.__name__ == "ConstrainedLayoutEngine": 

441 engine.set(w_pad=0.04, h_pad=0.04, wspace=0.02, hspace=0.02) 

442 elif has_suptitle: 

443 fig.tight_layout(rect=[0, 0, 1, 0.92], w_pad=0.3, h_pad=0.3) 

444 else: 

445 fig.tight_layout(w_pad=0.3, h_pad=0.3) 

446 

447 def update(k): 

448 for r in range(rows): 

449 for c in range(cols): 

450 for coll in list(axs_arr[r, c].collections): 

451 coll.remove() 

452 _render_qp_grid( 

453 axs_arr, 

454 QP_anim[k], 

455 pts_x_scaled, 

456 pts_y_scaled, 

457 contour=contour, 

458 cmap=cmap, 

459 vmin=vmin, 

460 vmax=vmax, 

461 x_ticks=x_ticks, 

462 y_ticks=y_ticks, 

463 z_ticks=z_ticks, 

464 cbar_label=cbar_label, 

465 plot_cbar=plot_cbar, 

466 subtitles=subtitles, 

467 decorate=False, 

468 ) 

469 _set_suptitle(k) 

470 return [] 

471 

472 anim = FuncAnimation(fig, update, frames=N, interval=interval_ms, blit=False) 

473 if save_path is not None: 

474 fps = max(1, round(1000 / interval_ms)) 

475 anim.save(save_path, writer=PillowWriter(fps=fps)) 

476 

477 # Make Jupyter render the animation inline without needing an explicit 

478 # HTML(anim.to_jshtml()) wrapper, and suppress the static last-frame 

479 # figure that the inline backend would otherwise emit alongside it. 

480 anim._repr_html_ = lambda a=anim: a.to_jshtml() 

481 plt.close(fig) 

482 return anim 

483 

484 

485def plot_wigner( 

486 state, 

487 pts_x, 

488 pts_y=None, 

489 g=2, 

490 axs=None, 

491 contour=True, 

492 cbar_label="", 

493 axis_scale_factor=1, 

494 plot_cbar=True, 

495 x_ticks=None, 

496 y_ticks=None, 

497 z_ticks=None, 

498 subtitles=None, 

499 figtitle=None, 

500 gif=False, 

501 gif_params=None, 

502): 

503 """Plot the wigner function of the state. 

504 

505 Thin wrapper around :func:`plot_qp` with ``qp_type='wigner'``. 

506 

507 Args: 

508 state: state with arbitrary number of batch dimensions, result will 

509 be flattened to a 2d grid to allow for plotting 

510 pts_x: x points to evaluate quasi-probability distribution at 

511 pts_y: y points to evaluate quasi-probability distribution at 

512 g: float, default 2. Scaling factor for ``a = 0.5 * g * (x + iy)``. 

513 The value of ``g`` is related to the value of :math:`\\hbar` in 

514 the commutation relation :math:`[x,\,y] = i\\hbar` via 

515 :math:`\\hbar=2/g^2`. 

516 axs: matplotlib axes to plot on 

517 contour: make the plot use contouring 

518 cbar_label: label for the cbar 

519 axis_scale_factor: scale of the axes labels relative 

520 plot_cbar: whether to plot cbar 

521 x_ticks: tick position for the x-axis 

522 y_ticks: tick position for the y-axis 

523 z_ticks: tick position for the z-axis 

524 subtitles: subtitles for the subplots 

525 figtitle: figure title 

526 gif: if True, render an animation over one batch axis instead of a 

527 tiled subplot grid. See :func:`plot_qp` for details. 

528 gif_params: dict of options for the gif path. Recognized keys: 

529 ``save_path`` (default None), ``interval_ms`` (default 200), 

530 ``ts`` (default None — adds a ``t = …`` label per frame), 

531 ``batch_animation_axis`` (default 0). 

532 

533 Returns: 

534 ``(axs, im)`` in the static case, or a ``matplotlib.animation.FuncAnimation`` 

535 when ``gif=True``. 

536 """ 

537 return plot_qp( 

538 state=state, 

539 pts_x=pts_x, 

540 pts_y=pts_y, 

541 g=g, 

542 axs=axs, 

543 contour=contour, 

544 qp_type=WIGNER, 

545 cbar_label=cbar_label, 

546 axis_scale_factor=axis_scale_factor, 

547 plot_cbar=plot_cbar, 

548 x_ticks=x_ticks, 

549 y_ticks=y_ticks, 

550 z_ticks=z_ticks, 

551 subtitles=subtitles, 

552 figtitle=figtitle, 

553 gif=gif, 

554 gif_params=gif_params, 

555 ) 

556 

557 

558def plot_qfunc( 

559 state, 

560 pts_x, 

561 pts_y=None, 

562 g=2, 

563 axs=None, 

564 contour=True, 

565 cbar_label="", 

566 axis_scale_factor=1, 

567 plot_cbar=True, 

568 x_ticks=None, 

569 y_ticks=None, 

570 z_ticks=None, 

571 subtitles=None, 

572 figtitle=None, 

573 gif=False, 

574 gif_params=None, 

575): 

576 """Plot the husimi (Q) function of the state. 

577 

578 Thin wrapper around :func:`plot_qp` with ``qp_type='husimi'``. 

579 

580 Args: 

581 state: state with arbitrary number of batch dimensions, result will 

582 be flattened to a 2d grid to allow for plotting 

583 pts_x: x points to evaluate quasi-probability distribution at 

584 pts_y: y points to evaluate quasi-probability distribution at 

585 g: float, default 2. Scaling factor for ``a = 0.5 * g * (x + iy)``. 

586 The value of ``g`` is related to the value of :math:`\\hbar` in 

587 the commutation relation :math:`[x,\,y] = i\\hbar` via 

588 :math:`\\hbar=2/g^2`. 

589 axs: matplotlib axes to plot on 

590 contour: make the plot use contouring 

591 cbar_label: label for the cbar 

592 axis_scale_factor: scale of the axes labels relative 

593 plot_cbar: whether to plot cbar 

594 x_ticks: tick position for the x-axis 

595 y_ticks: tick position for the y-axis 

596 z_ticks: tick position for the z-axis 

597 subtitles: subtitles for the subplots 

598 figtitle: figure title 

599 gif: if True, render an animation over one batch axis instead of a 

600 tiled subplot grid. See :func:`plot_qp` for details. 

601 gif_params: dict of options for the gif path. Recognized keys: 

602 ``save_path`` (default None), ``interval_ms`` (default 200), 

603 ``ts`` (default None — adds a ``t = …`` label per frame), 

604 ``batch_animation_axis`` (default 0). 

605 

606 Returns: 

607 ``(axs, im)`` in the static case, or a ``matplotlib.animation.FuncAnimation`` 

608 when ``gif=True``. 

609 """ 

610 return plot_qp( 

611 state=state, 

612 pts_x=pts_x, 

613 pts_y=pts_y, 

614 g=g, 

615 axs=axs, 

616 contour=contour, 

617 qp_type=HUSIMI, 

618 cbar_label=cbar_label, 

619 axis_scale_factor=axis_scale_factor, 

620 plot_cbar=plot_cbar, 

621 x_ticks=x_ticks, 

622 y_ticks=y_ticks, 

623 z_ticks=z_ticks, 

624 subtitles=subtitles, 

625 figtitle=figtitle, 

626 gif=gif, 

627 gif_params=gif_params, 

628 ) 

629 

630 

631def _render_cf_grid( 

632 axs, 

633 QP, 

634 pts_x, 

635 pts_y, 

636 *, 

637 contour, 

638 cmap, 

639 vmin, 

640 vmax, 

641 x_ticks, 

642 y_ticks, 

643 z_ticks, 

644 cbar_label, 

645 plot_cbar, 

646 plot_grid, 

647 subtitles, 

648 decorate=True, 

649): 

650 """Render one characteristic-function frame onto a ``(rows, 2*cols)`` axes grid. 

651 

652 Each batch element ``QP[row, col]`` is drawn as two adjacent subplots: 

653 the real part at column ``2*col``, the imaginary part at ``2*col + 1``. 

654 ``decorate=False`` skips the colorbar/ticks/labels block, used for gif 

655 frames after the first so colorbars laid out on frame 0 aren't 

656 duplicated. Returns the last ``contourf`` / ``pcolormesh`` artist. 

657 """ 

658 rows, cols = QP.shape[0], QP.shape[1] 

659 im = None 

660 for row in range(rows): 

661 for col in range(cols): 

662 for subcol in range(2): 

663 ax = axs[row, 2 * col + subcol] 

664 data = ( 

665 jnp.real(QP[row, col]) 

666 if subcol == 0 

667 else jnp.imag(QP[row, col]) 

668 ) 

669 if contour: 

670 im = ax.contourf( 

671 pts_x, 

672 pts_y, 

673 data, 

674 cmap=cmap, 

675 vmin=vmin, 

676 vmax=vmax, 

677 levels=np.linspace(vmin, vmax, 101), 

678 ) 

679 else: 

680 im = ax.pcolormesh( 

681 pts_x, 

682 pts_y, 

683 data, 

684 cmap=cmap, 

685 vmin=vmin, 

686 vmax=vmax, 

687 ) 

688 if decorate: 

689 ax.set_xticks(x_ticks) 

690 ax.set_yticks(y_ticks) 

691 if plot_grid: 

692 ax.grid() 

693 ax.set_aspect("equal", adjustable="box") 

694 if plot_cbar: 

695 cbar = plt.colorbar( 

696 im, 

697 ax=ax, 

698 orientation="vertical", 

699 ticks=np.linspace(-1, 1, 11), 

700 ) 

701 cbar.ax.set_title(cbar_label[subcol]) 

702 cbar.set_ticks(z_ticks) 

703 ax.set_xlabel(r"Re[$\alpha$]") 

704 ax.set_ylabel(r"Im[$\alpha$]") 

705 if subtitles is not None: 

706 ax.set_title(subtitles[row, col]) 

707 return im 

708 

709 

710def plot_cf( 

711 state, 

712 pts_x, 

713 pts_y=None, 

714 axs=None, 

715 contour=True, 

716 qp_type=WIGNER, 

717 cbar_label="", 

718 axis_scale_factor=1, 

719 plot_cbar=True, 

720 plot_grid=True, 

721 x_ticks=None, 

722 y_ticks=None, 

723 z_ticks=None, 

724 subtitles=None, 

725 figtitle=None, 

726 gif=False, 

727 gif_params=None, 

728): 

729 """Plot a characteristic function as paired real/imag subplots. 

730 

731 Each batch element produces two adjacent subplots — real part followed 

732 by imaginary part — so the rendered grid has shape ``(rows, 2 * cols)``. 

733 

734 Args: 

735 state: state with arbitrary number of batch dimensions, result will 

736 be flattened to a 2d grid to allow for plotting 

737 pts_x: x points to evaluate the characteristic function at 

738 pts_y: y points to evaluate the characteristic function at 

739 axs: matplotlib axes to plot on 

740 contour: make the plot use contouring 

741 qp_type: type of characteristic function. Currently only 

742 ``"wigner"`` is supported. 

743 cbar_label: labels for the real and imaginary cbar (overridden 

744 internally based on ``qp_type``) 

745 axis_scale_factor: scale of the axes labels relative 

746 plot_cbar: whether to plot cbar 

747 plot_grid: whether to draw gridlines on each subplot 

748 x_ticks: tick position for the x-axis 

749 y_ticks: tick position for the y-axis 

750 z_ticks: tick position for the z-axis 

751 subtitles: subtitles for the subplots (shape must match ``state.bdims``) 

752 figtitle: figure title 

753 gif: if True, render an animation over one batch axis instead of a 

754 tiled grid. Returns a ``matplotlib.animation.FuncAnimation`` 

755 that auto-renders inline in Jupyter. 

756 gif_params: dict of options for the gif path. Recognized keys: 

757 ``save_path`` (default None) — if set, save the animation here 

758 via PillowWriter; ``interval_ms`` (default 200) — milliseconds 

759 per frame; ``ts`` (default None) — optional 1D array of 

760 timestamps matching the animation-axis length; when set, each 

761 frame's suptitle gets a ``t = …`` label; 

762 ``batch_animation_axis`` (default 0) — index into 

763 ``state.bdims`` selecting which axis becomes the animation/time 

764 axis (the remaining batch dims form the per-frame subplot grid). 

765 

766 Returns: 

767 ``(axs, im)`` in the static case, or a ``FuncAnimation`` when 

768 ``gif=True``. 

769 """ 

770 if pts_y is None: 

771 pts_y = pts_x 

772 pts_x = jnp.array(pts_x) 

773 pts_y = jnp.array(pts_y) 

774 

775 if gif: 

776 return _plot_cf_gif( 

777 state=state, 

778 pts_x=pts_x, 

779 pts_y=pts_y, 

780 axs=axs, 

781 contour=contour, 

782 qp_type=qp_type, 

783 axis_scale_factor=axis_scale_factor, 

784 plot_cbar=plot_cbar, 

785 plot_grid=plot_grid, 

786 x_ticks=x_ticks, 

787 y_ticks=y_ticks, 

788 z_ticks=z_ticks, 

789 subtitles=subtitles, 

790 figtitle=figtitle, 

791 gif_params=gif_params or {}, 

792 ) 

793 

794 bdims = state.bdims 

795 added_baxes = 0 

796 

797 if subtitles is not None and subtitles.shape != bdims: 

798 raise ValueError( 

799 f"labels must have same shape as bdims, " 

800 f"got shapes {subtitles.shape} and {bdims}" 

801 ) 

802 

803 if len(bdims) == 0: 

804 bdims = (1,) 

805 added_baxes += 1 

806 if len(bdims) == 1: 

807 bdims = (1, bdims[0]) 

808 added_baxes += 1 

809 

810 extra_dims = bdims[2:] 

811 if extra_dims != (): 

812 state = state.reshape_bdims( 

813 bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1] 

814 ) 

815 if subtitles is not None: 

816 subtitles = subtitles.reshape( 

817 bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1] 

818 ) 

819 bdims = state.bdims 

820 

821 if axs is None: 

822 _, axs = plt.subplots( 

823 bdims[0], 

824 bdims[1]*2, 

825 figsize=(3.3 * bdims[1]*2, 3 * bdims[0]), 

826 dpi=200, 

827 ) 

828 

829 

830 if qp_type == WIGNER: 

831 vmin = -1 

832 vmax = 1 

833 scale = 1 

834 cmap = "seismic" 

835 cbar_label = [ 

836 r"$\mathcal{Re}(\chi_W(\alpha))$", 

837 r"$\mathcal{Im}(\chi_W(\alpha))$", 

838 ] 

839 QP = scale * cf_wigner(state, pts_x, pts_y) 

840 

841 for _ in range(added_baxes): 

842 QP = jnp.array([QP]) 

843 axs = np.array([axs]) 

844 if subtitles is not None: 

845 subtitles = np.array([subtitles]) 

846 

847 if added_baxes==2: 

848 axs = axs[0] # When the input state is zero-dimensional, remove an 

849 # axis that is automatically added due to the subcolumns 

850 

851 

852 pts_x = pts_x * axis_scale_factor 

853 pts_y = pts_y * axis_scale_factor 

854 

855 x_ticks = ( 

856 jnp.linspace(jnp.min(pts_x), jnp.max(pts_x), 

857 5) if x_ticks is None else x_ticks 

858 ) 

859 y_ticks = ( 

860 jnp.linspace(jnp.min(pts_y), jnp.max(pts_y), 

861 5) if y_ticks is None else y_ticks 

862 ) 

863 z_ticks = jnp.linspace(vmin, vmax, 11) if z_ticks is None else z_ticks 

864 

865 im = _render_cf_grid( 

866 axs, 

867 QP, 

868 pts_x, 

869 pts_y, 

870 contour=contour, 

871 cmap=cmap, 

872 vmin=vmin, 

873 vmax=vmax, 

874 x_ticks=x_ticks, 

875 y_ticks=y_ticks, 

876 z_ticks=z_ticks, 

877 cbar_label=cbar_label, 

878 plot_cbar=plot_cbar, 

879 plot_grid=plot_grid, 

880 subtitles=subtitles, 

881 decorate=True, 

882 ) 

883 

884 fig = axs[0, 0].get_figure() 

885 fig.tight_layout(w_pad=0.3, h_pad=0.3) 

886 if figtitle is not None: 

887 fig.suptitle(figtitle, y=1.04) 

888 return axs, im 

889 

890 

891def _plot_cf_gif( 

892 state, 

893 pts_x, 

894 pts_y, 

895 *, 

896 axs, 

897 contour, 

898 qp_type, 

899 axis_scale_factor, 

900 plot_cbar, 

901 plot_grid, 

902 x_ticks, 

903 y_ticks, 

904 z_ticks, 

905 subtitles, 

906 figtitle,