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=(4 * 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() 

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=(4 * cols, 3 * rows), dpi=200 

397 ) 

398 axs_arr = axs 

399 for _ in range(added_baxes): 

400 axs_arr = np.array([axs_arr]) 

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

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

403 

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

405 

406 def _set_suptitle(k): 

407 if ts is not None: 

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

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

410 fig.suptitle(title, y=0.98) 

411 elif figtitle is not None: 

412 fig.suptitle(figtitle, y=0.98) 

413 

414 _render_qp_grid( 

415 axs_arr, 

416 QP_anim[0], 

417 pts_x_scaled, 

418 pts_y_scaled, 

419 contour=contour, 

420 cmap=cmap, 

421 vmin=vmin, 

422 vmax=vmax, 

423 x_ticks=x_ticks, 

424 y_ticks=y_ticks, 

425 z_ticks=z_ticks, 

426 cbar_label=cbar_label, 

427 plot_cbar=plot_cbar, 

428 subtitles=subtitles, 

429 decorate=True, 

430 ) 

431 _set_suptitle(0) 

432 if has_suptitle: 

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

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

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

436 else: 

437 fig.tight_layout() 

438 

439 def update(k): 

440 for r in range(rows): 

441 for c in range(cols): 

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

443 coll.remove() 

444 _render_qp_grid( 

445 axs_arr, 

446 QP_anim[k], 

447 pts_x_scaled, 

448 pts_y_scaled, 

449 contour=contour, 

450 cmap=cmap, 

451 vmin=vmin, 

452 vmax=vmax, 

453 x_ticks=x_ticks, 

454 y_ticks=y_ticks, 

455 z_ticks=z_ticks, 

456 cbar_label=cbar_label, 

457 plot_cbar=plot_cbar, 

458 subtitles=subtitles, 

459 decorate=False, 

460 ) 

461 _set_suptitle(k) 

462 return [] 

463 

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

465 if save_path is not None: 

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

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

468 

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

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

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

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

473 plt.close(fig) 

474 return anim 

475 

476 

477def plot_wigner( 

478 state, 

479 pts_x, 

480 pts_y=None, 

481 g=2, 

482 axs=None, 

483 contour=True, 

484 cbar_label="", 

485 axis_scale_factor=1, 

486 plot_cbar=True, 

487 x_ticks=None, 

488 y_ticks=None, 

489 z_ticks=None, 

490 subtitles=None, 

491 figtitle=None, 

492 gif=False, 

493 gif_params=None, 

494): 

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

496 

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

498 

499 Args: 

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

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

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

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

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

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

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

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

508 axs: matplotlib axes to plot on 

509 contour: make the plot use contouring 

510 cbar_label: label for the cbar 

511 axis_scale_factor: scale of the axes labels relative 

512 plot_cbar: whether to plot cbar 

513 x_ticks: tick position for the x-axis 

514 y_ticks: tick position for the y-axis 

515 z_ticks: tick position for the z-axis 

516 subtitles: subtitles for the subplots 

517 figtitle: figure title 

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

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

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

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

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

523 ``batch_animation_axis`` (default 0). 

524 

525 Returns: 

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

527 when ``gif=True``. 

528 """ 

529 return plot_qp( 

530 state=state, 

531 pts_x=pts_x, 

532 pts_y=pts_y, 

533 g=g, 

534 axs=axs, 

535 contour=contour, 

536 qp_type=WIGNER, 

537 cbar_label=cbar_label, 

538 axis_scale_factor=axis_scale_factor, 

539 plot_cbar=plot_cbar, 

540 x_ticks=x_ticks, 

541 y_ticks=y_ticks, 

542 z_ticks=z_ticks, 

543 subtitles=subtitles, 

544 figtitle=figtitle, 

545 gif=gif, 

546 gif_params=gif_params, 

547 ) 

548 

549 

550def plot_qfunc( 

551 state, 

552 pts_x, 

553 pts_y=None, 

554 g=2, 

555 axs=None, 

556 contour=True, 

557 cbar_label="", 

558 axis_scale_factor=1, 

559 plot_cbar=True, 

560 x_ticks=None, 

561 y_ticks=None, 

562 z_ticks=None, 

563 subtitles=None, 

564 figtitle=None, 

565 gif=False, 

566 gif_params=None, 

567): 

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

569 

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

571 

572 Args: 

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

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

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

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

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

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

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

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

581 axs: matplotlib axes to plot on 

582 contour: make the plot use contouring 

583 cbar_label: label for the cbar 

584 axis_scale_factor: scale of the axes labels relative 

585 plot_cbar: whether to plot cbar 

586 x_ticks: tick position for the x-axis 

587 y_ticks: tick position for the y-axis 

588 z_ticks: tick position for the z-axis 

589 subtitles: subtitles for the subplots 

590 figtitle: figure title 

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

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

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

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

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

596 ``batch_animation_axis`` (default 0). 

597 

598 Returns: 

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

600 when ``gif=True``. 

601 """ 

602 return plot_qp( 

603 state=state, 

604 pts_x=pts_x, 

605 pts_y=pts_y, 

606 g=g, 

607 axs=axs, 

608 contour=contour, 

609 qp_type=HUSIMI, 

610 cbar_label=cbar_label, 

611 axis_scale_factor=axis_scale_factor, 

612 plot_cbar=plot_cbar, 

613 x_ticks=x_ticks, 

614 y_ticks=y_ticks, 

615 z_ticks=z_ticks, 

616 subtitles=subtitles, 

617 figtitle=figtitle, 

618 gif=gif, 

619 gif_params=gif_params, 

620 ) 

621 

622 

623def _render_cf_grid( 

624 axs, 

625 QP, 

626 pts_x, 

627 pts_y, 

628 *, 

629 contour, 

630 cmap, 

631 vmin, 

632 vmax, 

633 x_ticks, 

634 y_ticks, 

635 z_ticks, 

636 cbar_label, 

637 plot_cbar, 

638 plot_grid, 

639 subtitles, 

640 decorate=True, 

641): 

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

643 

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

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

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

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

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

649 """ 

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

651 im = None 

652 for row in range(rows): 

653 for col in range(cols): 

654 for subcol in range(2): 

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

656 data = ( 

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

658 if subcol == 0 

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

660 ) 

661 if contour: 

662 im = ax.contourf( 

663 pts_x, 

664 pts_y, 

665 data, 

666 cmap=cmap, 

667 vmin=vmin, 

668 vmax=vmax, 

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

670 ) 

671 else: 

672 im = ax.pcolormesh( 

673 pts_x, 

674 pts_y, 

675 data, 

676 cmap=cmap, 

677 vmin=vmin, 

678 vmax=vmax, 

679 ) 

680 if decorate: 

681 ax.set_xticks(x_ticks) 

682 ax.set_yticks(y_ticks) 

683 if plot_grid: 

684 ax.grid() 

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

686 if plot_cbar: 

687 cbar = plt.colorbar( 

688 im, 

689 ax=ax, 

690 orientation="vertical", 

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

692 ) 

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

694 cbar.set_ticks(z_ticks) 

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

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

697 if subtitles is not None: 

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

699 return im 

700 

701 

702def plot_cf( 

703 state, 

704 pts_x, 

705 pts_y=None, 

706 axs=None, 

707 contour=True, 

708 qp_type=WIGNER, 

709 cbar_label="", 

710 axis_scale_factor=1, 

711 plot_cbar=True, 

712 plot_grid=True, 

713 x_ticks=None, 

714 y_ticks=None, 

715 z_ticks=None, 

716 subtitles=None, 

717 figtitle=None, 

718 gif=False, 

719 gif_params=None, 

720): 

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

722 

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

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

725 

726 Args: 

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

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

729 pts_x: x points to evaluate the characteristic function at 

730 pts_y: y points to evaluate the characteristic function at 

731 axs: matplotlib axes to plot on 

732 contour: make the plot use contouring 

733 qp_type: type of characteristic function. Currently only 

734 ``"wigner"`` is supported. 

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

736 internally based on ``qp_type``) 

737 axis_scale_factor: scale of the axes labels relative 

738 plot_cbar: whether to plot cbar 

739 plot_grid: whether to draw gridlines on each subplot 

740 x_ticks: tick position for the x-axis 

741 y_ticks: tick position for the y-axis 

742 z_ticks: tick position for the z-axis 

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

744 figtitle: figure title 

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

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

747 that auto-renders inline in Jupyter. 

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

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

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

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

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

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

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

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

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

757 

758 Returns: 

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

760 ``gif=True``. 

761 """ 

762 if pts_y is None: 

763 pts_y = pts_x 

764 pts_x = jnp.array(pts_x) 

765 pts_y = jnp.array(pts_y) 

766 

767 if gif: 

768 return _plot_cf_gif( 

769 state=state, 

770 pts_x=pts_x, 

771 pts_y=pts_y, 

772 axs=axs, 

773 contour=contour, 

774 qp_type=qp_type, 

775 axis_scale_factor=axis_scale_factor, 

776 plot_cbar=plot_cbar, 

777 plot_grid=plot_grid, 

778 x_ticks=x_ticks, 

779 y_ticks=y_ticks, 

780 z_ticks=z_ticks, 

781 subtitles=subtitles, 

782 figtitle=figtitle, 

783 gif_params=gif_params or {}, 

784 ) 

785 

786 bdims = state.bdims 

787 added_baxes = 0 

788 

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

790 raise ValueError( 

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

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

793 ) 

794 

795 if len(bdims) == 0: 

796 bdims = (1,) 

797 added_baxes += 1 

798 if len(bdims) == 1: 

799 bdims = (1, bdims[0]) 

800 added_baxes += 1 

801 

802 extra_dims = bdims[2:] 

803 if extra_dims != (): 

804 state = state.reshape_bdims( 

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

806 ) 

807 if subtitles is not None: 

808 subtitles = subtitles.reshape( 

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

810 ) 

811 bdims = state.bdims 

812 

813 if axs is None: 

814 _, axs = plt.subplots( 

815 bdims[0], 

816 bdims[1]*2, 

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

818 dpi=200, 

819 ) 

820 

821 

822 if qp_type == WIGNER: 

823 vmin = -1 

824 vmax = 1 

825 scale = 1 

826 cmap = "seismic" 

827 cbar_label = [ 

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

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

830 ] 

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

832 

833 for _ in range(added_baxes): 

834 QP = jnp.array([QP]) 

835 axs = np.array([axs]) 

836 if subtitles is not None: 

837 subtitles = np.array([subtitles]) 

838 

839 if added_baxes==2: 

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

841 # axis that is automatically added due to the subcolumns 

842 

843 

844 pts_x = pts_x * axis_scale_factor 

845 pts_y = pts_y * axis_scale_factor 

846 

847 x_ticks = ( 

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

849 5) if x_ticks is None else x_ticks 

850 ) 

851 y_ticks = ( 

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

853 5) if y_ticks is None else y_ticks 

854 ) 

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

856 

857 im = _render_cf_grid( 

858 axs, 

859 QP, 

860 pts_x, 

861 pts_y, 

862 contour=contour, 

863 cmap=cmap, 

864 vmin=vmin, 

865 vmax=vmax, 

866 x_ticks=x_ticks, 

867 y_ticks=y_ticks, 

868 z_ticks=z_ticks, 

869 cbar_label=cbar_label, 

870 plot_cbar=plot_cbar, 

871 plot_grid=plot_grid, 

872 subtitles=subtitles, 

873 decorate=True, 

874 ) 

875 

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

877 fig.tight_layout() 

878 if figtitle is not None: 

879 fig.suptitle(figtitle, y=1.04) 

880 return axs, im 

881 

882 

883def _plot_cf_gif( 

884 state, 

885 pts_x, 

886 pts_y, 

887 *, 

888 axs, 

889 contour, 

890 qp_type, 

891 axis_scale_factor, 

892 plot_cbar, 

893 plot_grid, 

894 x_ticks, 

895 y_ticks, 

896 z_ticks, 

897 subtitles, 

898 figtitle, 

899 gif_params, 

900): 

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

902 

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

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

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

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

907 per-frame layout, suptitle inside the figure with 

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

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

910 Jupyter rendering. 

911 """ 

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

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