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: 

196 if subtitles.shape != bdims: 

197 raise ValueError( 

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

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

200 ) 

201 

202 if len(bdims) == 0: 

203 bdims = (1,) 

204 added_baxes += 1 

205 if len(bdims) == 1: 

206 bdims = (1, bdims[0]) 

207 added_baxes += 1 

208 

209 extra_dims = bdims[2:] 

210 if extra_dims != (): 

211 state = state.reshape_bdims( 

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

213 ) 

214 if subtitles is not None: 

215 subtitles = subtitles.reshape( 

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

217 ) 

218 bdims = state.bdims 

219 

220 if axs is None: 

221 _, axs = plt.subplots( 

222 bdims[0], 

223 bdims[1], 

224 figsize=(4 * bdims[1], 3 * bdims[0]), 

225 dpi=200, 

226 ) 

227 

228 if qp_type == WIGNER: 

229 vmin = -1 

230 vmax = 1 

231 scale = np.pi / 2 

232 cmap = "seismic" 

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

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

235 

236 elif qp_type == HUSIMI: 

237 vmin = 0 

238 vmax = 1 

239 scale = np.pi 

240 cmap = "jet" 

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

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

243 

244 

245 

246 for _ in range(added_baxes): 

247 QP = jnp.array([QP]) 

248 axs = np.array([axs]) 

249 if subtitles is not None: 

250 subtitles = np.array([subtitles]) 

251 

252 

253 

254 

255 pts_x = pts_x * axis_scale_factor 

256 pts_y = pts_y * axis_scale_factor 

257 

258 x_ticks = ( 

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

260 ) 

261 y_ticks = ( 

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

263 ) 

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

265 

266 im = _render_qp_grid( 

267 axs, 

268 QP, 

269 pts_x, 

270 pts_y, 

271 contour=contour, 

272 cmap=cmap, 

273 vmin=vmin, 

274 vmax=vmax, 

275 x_ticks=x_ticks, 

276 y_ticks=y_ticks, 

277 z_ticks=z_ticks, 

278 cbar_label=cbar_label, 

279 plot_cbar=plot_cbar, 

280 subtitles=subtitles, 

281 decorate=True, 

282 ) 

283 

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

285 fig.tight_layout() 

286 if figtitle is not None: 

287 fig.suptitle(figtitle, y=1.04) 

288 return axs, im 

289 

290 

291def _plot_qp_gif( 

292 state, 

293 pts_x, 

294 pts_y, 

295 *, 

296 g, 

297 axs, 

298 contour, 

299 qp_type, 

300 axis_scale_factor, 

301 plot_cbar, 

302 x_ticks, 

303 y_ticks, 

304 z_ticks, 

305 subtitles, 

306 figtitle, 

307 gif_params, 

308): 

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

310 

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

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

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

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

315 frame 0 are preserved). 

316 

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

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

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

320 without an extra static last-frame image. 

321 """ 

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

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

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

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

326 

327 bdims = state.bdims 

328 if len(bdims) < 1: 

329 raise ValueError( 

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

331 ) 

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

333 raise ValueError( 

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

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

336 ) 

337 N = bdims[batch_animation_axis] 

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

339 raise ValueError( 

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

341 ) 

342 

343 if qp_type == WIGNER: 

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

345 cmap = "seismic" 

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

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

348 elif qp_type == HUSIMI: 

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

350 cmap = "jet" 

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

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

353 

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

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

356 

357 grid_dims = list(rest_bdims) 

358 added_baxes = 0 

359 if len(grid_dims) == 0: 

360 grid_dims = [1] 

361 added_baxes += 1 

362 if len(grid_dims) == 1: 

363 grid_dims = [1, grid_dims[0]] 

364 added_baxes += 1 

365 extras = grid_dims[2:] 

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

367 cols = grid_dims[1] 

368 

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

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

371 

372 if subtitles is not None: 

373 subtitles = np.asarray(subtitles) 

374 if subtitles.shape != rest_bdims: 

375 raise ValueError( 

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

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

378 ) 

379 subtitles = subtitles.reshape(rows, cols) 

380 

381 pts_x_scaled = pts_x * axis_scale_factor 

382 pts_y_scaled = pts_y * axis_scale_factor 

383 x_ticks = ( 

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

385 if x_ticks is None 

386 else x_ticks 

387 ) 

388 y_ticks = ( 

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

390 if y_ticks is None 

391 else y_ticks 

392 ) 

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

394 

395 if axs is None: 

396 _, axs = plt.subplots( 

397 rows, cols, figsize=(4 * cols, 3 * rows), dpi=200 

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, y=0.98) 

412 elif figtitle is not None: 

413 fig.suptitle(figtitle, y=0.98) 

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 if has_suptitle: 

434 # Reserve top strip of the figure so the suptitle isn't clipped in the 

435 # rendered animation (PillowWriter uses the figure bbox as-is). 

436 fig.tight_layout(rect=[0, 0, 1, 0.92]) 

437 else: 

438 fig.tight_layout() 

439 

440 def update(k): 

441 for r in range(rows): 

442 for c in range(cols): 

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

444 coll.remove() 

445 _render_qp_grid( 

446 axs_arr, 

447 QP_anim[k], 

448 pts_x_scaled, 

449 pts_y_scaled, 

450 contour=contour, 

451 cmap=cmap, 

452 vmin=vmin, 

453 vmax=vmax, 

454 x_ticks=x_ticks, 

455 y_ticks=y_ticks, 

456 z_ticks=z_ticks, 

457 cbar_label=cbar_label, 

458 plot_cbar=plot_cbar, 

459 subtitles=subtitles, 

460 decorate=False, 

461 ) 

462 _set_suptitle(k) 

463 return [] 

464 

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

466 if save_path is not None: 

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

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

469 

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

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

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

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

474 plt.close(fig) 

475 return anim 

476 

477 

478def plot_wigner( 

479 state, 

480 pts_x, 

481 pts_y=None, 

482 g=2, 

483 axs=None, 

484 contour=True, 

485 cbar_label="", 

486 axis_scale_factor=1, 

487 plot_cbar=True, 

488 x_ticks=None, 

489 y_ticks=None, 

490 z_ticks=None, 

491 subtitles=None, 

492 figtitle=None, 

493 gif=False, 

494 gif_params=None, 

495): 

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

497 

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

499 

500 Args: 

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

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

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

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

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

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

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

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

509 axs: matplotlib axes to plot on 

510 contour: make the plot use contouring 

511 cbar_label: label for the cbar 

512 axis_scale_factor: scale of the axes labels relative 

513 plot_cbar: whether to plot cbar 

514 x_ticks: tick position for the x-axis 

515 y_ticks: tick position for the y-axis 

516 z_ticks: tick position for the z-axis 

517 subtitles: subtitles for the subplots 

518 figtitle: figure title 

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

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

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

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

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

524 ``batch_animation_axis`` (default 0). 

525 

526 Returns: 

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

528 when ``gif=True``. 

529 """ 

530 return plot_qp( 

531 state=state, 

532 pts_x=pts_x, 

533 pts_y=pts_y, 

534 g=g, 

535 axs=axs, 

536 contour=contour, 

537 qp_type=WIGNER, 

538 cbar_label=cbar_label, 

539 axis_scale_factor=axis_scale_factor, 

540 plot_cbar=plot_cbar, 

541 x_ticks=x_ticks, 

542 y_ticks=y_ticks, 

543 z_ticks=z_ticks, 

544 subtitles=subtitles, 

545 figtitle=figtitle, 

546 gif=gif, 

547 gif_params=gif_params, 

548 ) 

549 

550 

551def plot_qfunc( 

552 state, 

553 pts_x, 

554 pts_y=None, 

555 g=2, 

556 axs=None, 

557 contour=True, 

558 cbar_label="", 

559 axis_scale_factor=1, 

560 plot_cbar=True, 

561 x_ticks=None, 

562 y_ticks=None, 

563 z_ticks=None, 

564 subtitles=None, 

565 figtitle=None, 

566 gif=False, 

567 gif_params=None, 

568): 

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

570 

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

572 

573 Args: 

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

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

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

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

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

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

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

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

582 axs: matplotlib axes to plot on 

583 contour: make the plot use contouring 

584 cbar_label: label for the cbar 

585 axis_scale_factor: scale of the axes labels relative 

586 plot_cbar: whether to plot cbar 

587 x_ticks: tick position for the x-axis 

588 y_ticks: tick position for the y-axis 

589 z_ticks: tick position for the z-axis 

590 subtitles: subtitles for the subplots 

591 figtitle: figure title 

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

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

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

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

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

597 ``batch_animation_axis`` (default 0). 

598 

599 Returns: 

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

601 when ``gif=True``. 

602 """ 

603 return plot_qp( 

604 state=state, 

605 pts_x=pts_x, 

606 pts_y=pts_y, 

607 g=g, 

608 axs=axs, 

609 contour=contour, 

610 qp_type=HUSIMI, 

611 cbar_label=cbar_label, 

612 axis_scale_factor=axis_scale_factor, 

613 plot_cbar=plot_cbar, 

614 x_ticks=x_ticks, 

615 y_ticks=y_ticks, 

616 z_ticks=z_ticks, 

617 subtitles=subtitles, 

618 figtitle=figtitle, 

619 gif=gif, 

620 gif_params=gif_params, 

621 ) 

622 

623 

624def _render_cf_grid( 

625 axs, 

626 QP, 

627 pts_x, 

628 pts_y, 

629 *, 

630 contour, 

631 cmap, 

632 vmin, 

633 vmax, 

634 x_ticks, 

635 y_ticks, 

636 z_ticks, 

637 cbar_label, 

638 plot_cbar, 

639 plot_grid, 

640 subtitles, 

641 decorate=True, 

642): 

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

644 

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

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

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

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

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

650 """ 

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

652 im = None 

653 for row in range(rows): 

654 for col in range(cols): 

655 for subcol in range(2): 

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

657 data = ( 

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

659 if subcol == 0 

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

661 ) 

662 if contour: 

663 im = ax.contourf( 

664 pts_x, 

665 pts_y, 

666 data, 

667 cmap=cmap, 

668 vmin=vmin, 

669 vmax=vmax, 

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

671 ) 

672 else: 

673 im = ax.pcolormesh( 

674 pts_x, 

675 pts_y, 

676 data, 

677 cmap=cmap, 

678 vmin=vmin, 

679 vmax=vmax, 

680 ) 

681 if decorate: 

682 ax.set_xticks(x_ticks) 

683 ax.set_yticks(y_ticks) 

684 if plot_grid: 

685 ax.grid() 

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

687 if plot_cbar: 

688 cbar = plt.colorbar( 

689 im, 

690 ax=ax, 

691 orientation="vertical", 

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

693 ) 

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

695 cbar.set_ticks(z_ticks) 

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

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

698 if subtitles is not None: 

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

700 return im 

701 

702 

703def plot_cf( 

704 state, 

705 pts_x, 

706 pts_y=None, 

707 axs=None, 

708 contour=True, 

709 qp_type=WIGNER, 

710 cbar_label="", 

711 axis_scale_factor=1, 

712 plot_cbar=True, 

713 plot_grid=True, 

714 x_ticks=None, 

715 y_ticks=None, 

716 z_ticks=None, 

717 subtitles=None, 

718 figtitle=None, 

719 gif=False, 

720 gif_params=None, 

721): 

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

723 

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

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

726 

727 Args: 

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

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

730 pts_x: x points to evaluate the characteristic function at 

731 pts_y: y points to evaluate the characteristic function at 

732 axs: matplotlib axes to plot on 

733 contour: make the plot use contouring 

734 qp_type: type of characteristic function. Currently only 

735 ``"wigner"`` is supported. 

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

737 internally based on ``qp_type``) 

738 axis_scale_factor: scale of the axes labels relative 

739 plot_cbar: whether to plot cbar 

740 plot_grid: whether to draw gridlines on each subplot 

741 x_ticks: tick position for the x-axis 

742 y_ticks: tick position for the y-axis 

743 z_ticks: tick position for the z-axis 

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

745 figtitle: figure title 

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

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

748 that auto-renders inline in Jupyter. 

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

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

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

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

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

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

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

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

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

758 

759 Returns: 

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

761 ``gif=True``. 

762 """ 

763 if pts_y is None: 

764 pts_y = pts_x 

765 pts_x = jnp.array(pts_x) 

766 pts_y = jnp.array(pts_y) 

767 

768 if gif: 

769 return _plot_cf_gif( 

770 state=state, 

771 pts_x=pts_x, 

772 pts_y=pts_y, 

773 axs=axs, 

774 contour=contour, 

775 qp_type=qp_type, 

776 axis_scale_factor=axis_scale_factor, 

777 plot_cbar=plot_cbar, 

778 plot_grid=plot_grid, 

779 x_ticks=x_ticks, 

780 y_ticks=y_ticks, 

781 z_ticks=z_ticks, 

782 subtitles=subtitles, 

783 figtitle=figtitle, 

784 gif_params=gif_params or {}, 

785 ) 

786 

787 bdims = state.bdims 

788 added_baxes = 0 

789 

790 if subtitles is not None: 

791 if subtitles.shape != bdims: 

792 raise ValueError( 

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

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

795 ) 

796 

797 if len(bdims) == 0: 

798 bdims = (1,) 

799 added_baxes += 1 

800 if len(bdims) == 1: 

801 bdims = (1, bdims[0]) 

802 added_baxes += 1 

803 

804 extra_dims = bdims[2:] 

805 if extra_dims != (): 

806 state = state.reshape_bdims( 

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

808 ) 

809 if subtitles is not None: 

810 subtitles = subtitles.reshape( 

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

812 ) 

813 bdims = state.bdims 

814 

815 if axs is None: 

816 _, axs = plt.subplots( 

817 bdims[0], 

818 bdims[1]*2, 

819 figsize=(4 * bdims[1]*2, 3 * bdims[0]), 

820 dpi=200, 

821 ) 

822 

823 

824 if qp_type == WIGNER: 

825 vmin = -1 

826 vmax = 1 

827 scale = 1 

828 cmap = "seismic" 

829 cbar_label = [r"$\mathcal{Re}(\chi_W(\alpha))$", r"$\mathcal{" 

830 r"Im}(\chi_W(" 

831 r"\alpha))$"] 

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

833 

834 for _ in range(added_baxes): 

835 QP = jnp.array([QP]) 

836 axs = np.array([axs]) 

837 if subtitles is not None: 

838 subtitles = np.array([subtitles]) 

839 

840 if added_baxes==2: 

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

842 # axis that is automatically added due to the subcolumns 

843 

844 

845 pts_x = pts_x * axis_scale_factor 

846 pts_y = pts_y * axis_scale_factor 

847 

848 x_ticks = ( 

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

850 5) if x_ticks is None else x_ticks 

851 ) 

852 y_ticks = ( 

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

854 5) if y_ticks is None else y_ticks 

855 ) 

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

857 

858 im = _render_cf_grid( 

859 axs, 

860 QP, 

861 pts_x, 

862 pts_y, 

863 contour=contour, 

864 cmap=cmap, 

865 vmin=vmin, 

866 vmax=vmax, 

867 x_ticks=x_ticks, 

868 y_ticks=y_ticks, 

869 z_ticks=z_ticks, 

870 cbar_label=cbar_label, 

871 plot_cbar=plot_cbar, 

872 plot_grid=plot_grid, 

873 subtitles=subtitles, 

874 decorate=True, 

875 ) 

876 

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

878 fig.tight_layout() 

879 if figtitle is not None: 

880 fig.suptitle(figtitle, y=1.04) 

881 return axs, im 

882 

883 

884def _plot_cf_gif( 

885 state, 

886 pts_x, 

887 pts_y, 

888 *, 

889 axs, 

890 contour, 

891 qp_type, 

892 axis_scale_factor, 

893 plot_cbar, 

894 plot_grid, 

895 x_ticks, 

896 y_ticks, 

897 z_ticks, 

898 subtitles, 

899 figtitle, 

900 gif_params, 

901): 

902 """Build the ``FuncAnimation`` for ``plot_cf(gif=True)``. 

903 

904 Counterpart to :func:`_plot_qp_gif` but each frame is a 

905 ``(rows, 2*cols)`` grid of real|imag subplot pairs rendered via 

906 :func:`_render_cf_grid`. Same conventions: animation axis chosen by 

907 ``gif_params['batch_animation_axis']``, remaining batch dims form the 

908 per-frame layout, suptitle inside the figure with 

909 ``tight_layout(rect=[0, 0, 1, 0.92])`` so it doesn't clip in the saved 

910 gif, and ``anim._repr_html_`` patched + figure closed for inline 

911 Jupyter rendering. 

912 """ 

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

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