Coverage for python / aubellhop / plotutils.py: 49%

335 statements  

« prev     ^ index     » next       coverage.py v7.14.0, created at 2026-05-24 14:11 +0000

1"""Easy-to-use plotting utilities for aubellhop based on `Bokeh <http://bokeh.pydata.org>`_.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any, List, Tuple 

6import numpy as np 

7import os as _os 

8import warnings as _warnings 

9from tempfile import mkstemp as _mkstemp 

10import bokeh.plotting as _bplt 

11import bokeh.models as _bmodels 

12import bokeh.resources as _bres 

13import bokeh.io as _bio 

14import scipy.signal as _sig 

15 

16 

17light_palette = ['mediumblue', 'crimson', 'forestgreen', 'gold', 'darkmagenta', 'olive', 'palevioletred', 'yellowgreen', 

18 'deepskyblue', 'dimgray', 'indianred', 'mediumaquamarine', 'orange', 'saddlebrown', 'teal', 'mediumorchid'] 

19dark_palette = ['lightskyblue', 'red', 'limegreen', 'salmon', 'magenta', 'forestgreen', 'silver', 'teal'] 

20 

21_figure = None 

22_figures = None 

23_hold = False 

24_figsize = (600, 400) 

25_color = 0 

26_notebook = False 

27_disable_js = False 

28_using_js = False 

29_interactive = True 

30_static_images = False 

31_colors = light_palette 

32 

33# Detect Jupyter notebook environment 

34try: 

35 from IPython import get_ipython 

36 ipython = get_ipython() 

37 _notebook = ipython is not None and "ZMQInteractiveShell" in ipython.__class__.__name__ 

38 if _notebook: 

39 _bplt.output_notebook(resources=_bres.INLINE, hide_banner=True) 

40except (ImportError, AttributeError): 

41 _notebook = False 

42 

43def _new_figure(title: str | None, width: int | None, height: int | None, xlabel: str | None, ylabel: str | None, xlim: Tuple[float, float | None], ylim: Tuple[float, float | None], xtype: str | None, ytype: str | None, interactive: bool | None) -> Any: 

44 global _color, _figure 

45 

46 width = width or _figsize[0] 

47 height = height or _figsize[1] 

48 _color = 0 

49 tools: List[str] | str = [] 

50 interactive = interactive or _interactive 

51 if interactive: 

52 tools = 'pan,box_zoom,wheel_zoom,reset,save' 

53 args = dict(title=title, width=width, height=height, x_range=xlim, y_range=ylim, x_axis_label=xlabel, y_axis_label=ylabel, x_axis_type=xtype, y_axis_type=ytype, tools=tools) 

54 

55 if _figure is not None: 

56 f = _figure 

57 if title: 

58 f.title.text = title 

59 if width: 

60 f.width = width 

61 if height: 

62 f.height = height 

63 if xlabel and f.xaxis: 

64 f.xaxis[0].axis_label = xlabel 

65 if ylabel and f.yaxis: 

66 f.yaxis[0].axis_label = ylabel 

67 if xlim and hasattr(f, "x_range"): 

68 f.x_range.start, f.x_range.end = xlim 

69 if ylim and hasattr(f, "y_range"): 

70 f.y_range.start, f.y_range.end = ylim 

71 return f 

72 

73 f = _bplt.figure(**{k:v for (k,v) in args.items() if v is not None}) 

74 f.toolbar.logo = None 

75 _figure = f 

76 return f 

77 

78def _process_canvas(figures: List[Any]) -> None: 

79 """Replace non-interactive Bokeh canvases with static images in Jupyter notebooks. 

80 

81 This optimization converts non-interactive plots to static images to reduce 

82 JavaScript overhead in notebooks. Only runs if JavaScript is enabled and there 

83 are figures without interactive tools. 

84 """ 

85 global _using_js 

86 

87 if _disable_js or (not figures and _using_js): 

88 return 

89 

90 # Find indices of non-interactive figures 

91 disable_indices = [i + 1 for i, f in enumerate(figures) if f is not None and not f.tools] 

92 

93 if not disable_indices and not _using_js: 

94 return 

95 

96 _using_js = True 

97 

98 # JavaScript to convert non-interactive canvases to static images 

99 js_code = f""" 

100 var disable = {disable_indices}; 

101 var clist = document.getElementsByClassName('bk-canvas'); 

102 var j = 0; 

103 for (var i = 0; i < clist.length; i++) {{ 

104 if (clist[i].id == '') {{ 

105 j++; 

106 clist[i].id = 'bkc-' + String(i) + '-' + String(+new Date()); 

107 if (disable.indexOf(j) >= 0) {{ 

108 var png = clist[i].toDataURL(); 

109 var img = document.createElement('img'); 

110 img.src = png; 

111 clist[i].parentNode.replaceChild(img, clist[i]); 

112 }} 

113 }} 

114 }} 

115 """ 

116 

117 import IPython.display as _ipyd 

118 _ipyd.display(_ipyd.Javascript(js_code)) 

119 

120def _show_static_images(f: Any) -> None: 

121 fh, fname = _mkstemp(suffix='.png') 

122 _os.close(fh) 

123 with _warnings.catch_warnings(): # to avoid displaying deprecation warning 

124 _warnings.simplefilter('ignore') # from bokeh 0.12.16 

125 _bio.export_png(f, fname) 

126 import IPython.display as _ipyd 

127 _ipyd.display(_ipyd.Image(filename=fname, embed=True)) 

128 _os.unlink(fname) 

129 

130def _show(f: Any) -> None: 

131 if _figures is None: 

132 if _static_images: 

133 _show_static_images(f) 

134 else: 

135 _process_canvas([]) 

136 _bplt.show(f) 

137 _process_canvas([f]) 

138 else: 

139 _figures[-1].append(f) 

140 

141def _hold_enable(enable: bool) -> bool: 

142 global _hold, _figure 

143 ohold = _hold 

144 _hold = enable 

145 if not _hold and _figure is not None: 

146 _show(_figure) 

147 _figure = None 

148 return ohold 

149 

150def theme(name: str) -> None: 

151 """Set color theme. 

152 

153 Parameters 

154 ---------- 

155 name : str 

156 Name of theme 

157 

158 Examples 

159 -------- 

160 >>> import arlpy.plot 

161 >>> arlpy.plot.theme('dark') 

162 """ 

163 if name == 'dark': 

164 name = 'dark_minimal' 

165 set_colors(dark_palette) 

166 elif name == 'light': 

167 name = 'light_minimal' 

168 set_colors(light_palette) 

169 _bio.curdoc().theme = name 

170 

171def figsize(x: int, y: int) -> None: 

172 """Set the default figure size in pixels. 

173 

174 Parameters 

175 ---------- 

176 x : int 

177 Figure width 

178 y : int 

179 Figure height 

180 """ 

181 global _figsize 

182 _figsize = (x, y) 

183 

184def interactive(b: bool) -> None: 

185 """Set default interactivity for plots. 

186 

187 Parameters 

188 ---------- 

189 b : bool 

190 True to enable interactivity, False to disable it 

191 """ 

192 global _interactive 

193 _interactive = b 

194 

195def enable_javascript(b: bool) -> None: 

196 """Enable/disable Javascript. 

197 

198 Parameters 

199 ---------- 

200 b : bool 

201 True to use Javascript, False to avoid use of Javascript 

202 

203 Notes 

204 ----- 

205 Jupyterlab does not support Javascript output. To avoid error messages, 

206 Javascript can be disabled using this call. This removes an optimization 

207 to replace non-interactive plots with static images, but other than that 

208 does not affect functionality. 

209 """ 

210 global _disable_js 

211 _disable_js = not b 

212 

213def use_static_images(b: bool = True) -> None: 

214 """Use static images instead of dynamic HTML/Javascript in Jupyter notebook. 

215 

216 Parameters 

217 ---------- 

218 b : bool, default=True 

219 True to use static images, False to use HTML/Javascript 

220 

221 Notes 

222 ----- 

223 Static images are useful when the notebook is to be exported as a markdown, 

224 LaTeX or PDF document, since dynamic HTML/Javascript is not rendered in these 

225 formats. When static images are used, all interactive functionality is disabled. 

226 

227 To use static images, you must have the following packages installed: 

228 selenium, pillow, phantomjs. 

229 """ 

230 global _static_images, _interactive 

231 if not b: 

232 _static_images = False 

233 return 

234 if not _notebook: 

235 _warnings.warn('Not running in a Jupyter notebook, static png support disabled') 

236 return 

237 _interactive = False 

238 _static_images = True 

239 

240def hold(enable: bool = True) -> bool | None: 

241 """Combine multiple plots into one. 

242 

243 Parameters 

244 ---------- 

245 enable : bool, default=True 

246 True to hold plot, False to release hold 

247 

248 Returns 

249 ------- 

250 bool or None 

251 Old state of hold if enable is True 

252 

253 Examples 

254 -------- 

255 >>> import arlpy.plot 

256 >>> oh = arlpy.plot.hold() 

257 >>> arlpy.plot.plot([0,10], [0,10], color='blue', legend='A') 

258 >>> arlpy.plot.plot([10,0], [0,10], marker='o', color='green', legend='B') 

259 >>> arlpy.plot.hold(oh) 

260 """ 

261 rv = _hold_enable(enable) 

262 return rv if enable else None 

263 

264class figure: 

265 """Create a new figure, and optionally automatically display it. 

266 

267 Parameters 

268 ---------- 

269 title : str, optional 

270 Figure title 

271 xlabel : str, optional 

272 X-axis label 

273 ylabel : str, optional 

274 Y-axis label 

275 xlim : tuple of float, optional 

276 X-axis limits (min, max) 

277 ylim : tuple of float, optional 

278 Y-axis limits (min, max) 

279 xtype : str, default='auto' 

280 X-axis type ('auto', 'linear', 'log', etc) 

281 ytype : str, default='auto' 

282 Y-axis type ('auto', 'linear', 'log', etc) 

283 width : int, optional 

284 Figure width in pixels 

285 height : int, optional 

286 Figure height in pixels 

287 interactive : bool, optional 

288 Enable interactive tools (pan, zoom, etc) for plot 

289 

290 Notes 

291 ----- 

292 This function can be used in standalone mode to create a figure: 

293 

294 >>> import arlpy.plot 

295 >>> arlpy.plot.figure(title='Demo 1', width=500) 

296 >>> arlpy.plot.plot([0,10], [0,10]) 

297 

298 Or it can be used as a context manager to create, hold and display a figure: 

299 

300 >>> import arlpy.plot 

301 >>> with arlpy.plot.figure(title='Demo 2', width=500): 

302 >>> arlpy.plot.plot([0,10], [0,10], color='blue', legend='A') 

303 >>> arlpy.plot.plot([10,0], [0,10], marker='o', color='green', legend='B') 

304 

305 It can even be used as a context manager to work with Bokeh functions directly: 

306 

307 >>> import arlpy.plot 

308 >>> with arlpy.plot.figure(title='Demo 3', width=500) as f: 

309 >>> f.line([0,10], [0,10], line_color='blue') 

310 >>> f.square([3,7], [4,5], line_color='green', fill_color='yellow', size=10) 

311 """ 

312 

313 def __init__(self, title=None, xlabel=None, ylabel=None, xlim=None, ylim=None, xtype='auto', ytype='auto', width=None, height=None, interactive=None): 

314 global _figure 

315 _figure = _new_figure(title, width, height, xlabel, ylabel, xlim, ylim, xtype, ytype, interactive) 

316 

317 def __enter__(self): 

318 global _hold 

319 _hold = True 

320 return _figure 

321 

322 def __exit__(self, *args): 

323 global _hold, _figure 

324 _hold = False 

325 _show(_figure) 

326 _figure = None 

327 

328class many_figures: 

329 """Create a grid of many figures. 

330 

331 Parameters 

332 ---------- 

333 figsize : tuple of int, optional 

334 Default size of figure in grid as (width, height) 

335 

336 Examples 

337 -------- 

338 >>> import arlpy.plot 

339 >>> with arlpy.plot.many_figures(figsize=(300,200)): 

340 >>> arlpy.plot.plot([0,10], [0,10]) 

341 >>> arlpy.plot.plot([0,10], [0,10]) 

342 >>> arlpy.plot.next_row() 

343 >>> arlpy.plot.next_column() 

344 >>> arlpy.plot.plot([0,10], [0,10]) 

345 """ 

346 

347 def __init__(self, figsize: Tuple[int, int | None] = None): 

348 self.figsize = figsize 

349 self.old_figsize: Tuple[int, int | None] = None 

350 

351 def __enter__(self) -> None: 

352 global _figures, _figsize 

353 _figures = [[]] 

354 self.old_figsize = _figsize 

355 if self.figsize is not None: 

356 _figsize = self.figsize 

357 

358 def __exit__(self, *args: Any) -> None: 

359 global _figures, _figsize 

360 if _figures and (len(_figures) > 1 or _figures[0]): 

361 # Flatten nested list to get all figures 

362 all_figures = [fig for row in _figures for fig in row if fig is not None] 

363 gridplot = _bplt.gridplot(_figures, merge_tools=False) 

364 

365 if _static_images: 

366 _show_static_images(gridplot) 

367 else: 

368 _process_canvas([]) 

369 _bplt.show(gridplot) 

370 _process_canvas(all_figures) 

371 

372 _figures = None 

373 _figsize = self.old_figsize 

374 

375def next_row() -> None: 

376 """Move to the next row in a grid of many figures.""" 

377 global _figures 

378 if _figures is not None: 

379 _figures.append([]) 

380 

381def next_column() -> None: 

382 """Move to the next column in a grid of many figures.""" 

383 global _figures 

384 if _figures is not None: 

385 _figures[-1].append(None) 

386 

387def gcf() -> Any: 

388 """Get the current figure. 

389 

390 :returns: handle to the current figure 

391 """ 

392 return _figure 

393 

394def plot(x: Any, 

395 y: Any = None, 

396 fs: float | None = None, 

397 maxpts: int = 10000, 

398 pooling: str | None = None, 

399 color: str | None = None, 

400 style: str = 'solid', 

401 thickness: int = 1, 

402 marker: str | None = None, 

403 filled: bool = False, 

404 size: int = 6, 

405 mskip: int = 0, 

406 title: str | None = None, 

407 xlabel: str | None = None, 

408 ylabel: str | None = None, 

409 xlim: Tuple[float, float | None] = None, 

410 ylim: Tuple[float, float | None] = None, 

411 xtype: str = 'auto', 

412 ytype: str = 'auto', 

413 width: int | None = None, 

414 height: int | None = None, 

415 legend: str | None = None, 

416 interactive: bool | None = None, 

417 hold: bool = False, 

418 ) -> None: 

419 """Plot a line graph or time series. 

420 

421 :param x: x data or time series data (if y is None) 

422 :param y: y data or None (if time series) 

423 :param fs: sampling rate for time series data 

424 :param maxpts: maximum number of points to plot (downsampled if more points provided) 

425 :param pooling: pooling for downsampling (None, 'max', 'min', 'mean', 'median') 

426 :param color: line color (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_) 

427 :param style: line style ('solid', 'dashed', 'dotted', 'dotdash', 'dashdot', None) 

428 :param thickness: line width in pixels 

429 :param marker: point markers ('.', 'o', 's', '*', 'x', '+', 'd', '^') 

430 :param filled: filled markers or outlined ones 

431 :param size: marker size 

432 :param mskip: number of points to skip marking (to avoid too many markers) 

433 :param title: figure title 

434 :param xlabel: x-axis label 

435 :param ylabel: y-axis label 

436 :param xlim: x-axis limits (min, max) 

437 :param ylim: y-axis limits (min, max) 

438 :param xtype: x-axis type ('auto', 'linear', 'log', etc) 

439 :param ytype: y-axis type ('auto', 'linear', 'log', etc) 

440 :param width: figure width in pixels 

441 :param height: figure height in pixels 

442 :param legend: legend text 

443 :param interactive: enable interactive tools (pan, zoom, etc) for plot 

444 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

445 

446 >>> import arlpy.plot 

447 >>> import numpy as np 

448 >>> arlpy.plot.plot([0,10], [1,-1], color='blue', marker='o', filled=True, legend='A', hold=True) 

449 >>> arlpy.plot.plot(np.random.normal(size=1000), fs=100, color='green', legend='B') 

450 """ 

451 global _figure, _color 

452 x = np.asarray(x, dtype=np.float64) 

453 if y is None: 

454 y = x 

455 x = np.arange(x.size) 

456 if fs is not None: 

457 x = x/fs 

458 if xlabel is None: 

459 xlabel = 'Time (s)' 

460 if xlim is None: 

461 xlim = (x[0], x[-1]) 

462 else: 

463 y = np.asarray(y, dtype=np.float64) 

464 

465 if x.ndim == 0: # 0-dimensional array (scalar) 

466 x = np.array([x]) 

467 if y.ndim == 0: # 0-dimensional array (scalar) 

468 y = np.array([y]) 

469 

470 _figure = _new_figure(title, width, height, xlabel, ylabel, xlim, ylim, xtype, ytype, interactive) 

471 if color is None: 

472 color = _colors[_color % len(_colors)] 

473 _color += 1 

474 if x.size > maxpts: 

475 n = int(np.ceil(x.size / maxpts)) 

476 x = x[::n] 

477 desc = f'Downsampled by {n}' 

478 

479 # Apply pooling to reduce data 

480 if pooling is None: 

481 y = y[::n] 

482 else: 

483 # Trim data to fit evenly into bins 

484 trimmed_size = n * (y.size // n) 

485 y_trimmed = y[:trimmed_size].reshape(-1, n) 

486 

487 pooling_funcs = { 

488 'max': np.amax, 

489 'min': np.amin, 

490 'mean': np.mean, 

491 'median': np.median 

492 } 

493 

494 if pooling in pooling_funcs: 

495 y = pooling_funcs[pooling](y_trimmed, axis=1) 

496 desc += f', {pooling} pooled' 

497 else: 

498 _warnings.warn(f'Unknown pooling: {pooling}') 

499 y = y[::n] 

500 

501 # Ensure x and y have the same length 

502 if len(x) > len(y): 

503 x = x[:len(y)] 

504 

505 _figure.add_layout(_bmodels.Label( 

506 x=5, y=5, x_units='screen', y_units='screen', 

507 text=desc, text_font_size="8pt", text_alpha=0.5 

508 )) 

509 if style is not None: 

510 if legend is None: 

511 _figure.line(x, y, line_color=color, line_dash=style, line_width=thickness) 

512 else: 

513 _figure.line(x, y, line_color=color, line_dash=style, line_width=thickness, legend_label=legend) 

514 if marker is not None: 

515 scatter(x[::(mskip+1)], y[::(mskip+1)], marker=marker, filled=filled, size=size, color=color, legend=legend, hold=True) 

516 if not hold and not _hold: 

517 _show(_figure) 

518 _figure = None 

519 

520def scatter(x: Any, y: Any, marker: str = '.', filled: bool = False, size: int = 6, color: str | None = None, title: str | None = None, xlabel: str | None = None, ylabel: str | None = None, xlim: Tuple[float, float | None] = None, ylim: Tuple[float, float | None] = None, xtype: str = 'auto', ytype: str = 'auto', width: int | None = None, height: int | None = None, legend: str | None = None, hold: bool = False, interactive: bool | None = None) -> None: 

521 """Plot a scatter plot. 

522 

523 :param x: x data 

524 :param y: y data 

525 :param color: marker color (see `Bokeh colors`_) 

526 :param marker: point markers ('.', 'o', 's', '*', 'x', '+', 'd', '^') 

527 :param filled: filled markers or outlined ones 

528 :param size: marker size 

529 :param title: figure title 

530 :param xlabel: x-axis label 

531 :param ylabel: y-axis label 

532 :param xlim: x-axis limits (min, max) 

533 :param ylim: y-axis limits (min, max) 

534 :param xtype: x-axis type ('auto', 'linear', 'log', etc) 

535 :param ytype: y-axis type ('auto', 'linear', 'log', etc) 

536 :param width: figure width in pixels 

537 :param height: figure height in pixels 

538 :param legend: legend text 

539 :param interactive: enable interactive tools (pan, zoom, etc) for plot 

540 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

541 

542 >>> import arlpy.plot 

543 >>> import numpy as np 

544 >>> arlpy.plot.scatter(np.random.normal(size=100), np.random.normal(size=100), color='blue', marker='o') 

545 """ 

546 global _figure, _color 

547 if _figure is None: 

548 _figure = _new_figure(title, width, height, xlabel, ylabel, xlim, ylim, xtype, ytype, interactive) 

549 x = np.asarray(x, dtype=np.float64) 

550 y = np.asarray(y, dtype=np.float64) 

551 if color is None: 

552 color = _colors[_color % len(_colors)] 

553 _color += 1 

554 # Build kwargs for marker rendering 

555 kwargs = {'size': size, 'line_color': color} 

556 if filled: 

557 kwargs['fill_color'] = color 

558 if legend is not None: 

559 kwargs['legend_label'] = legend 

560 

561 # Map marker types to Bokeh scatter marker names (using modern Bokeh 3.4+ API) 

562 marker_map = { 

563 '.': 'circle', 

564 'o': 'circle', 

565 's': 'square', 

566 '*': 'star', 

567 'x': 'x', 

568 '+': 'cross', 

569 'd': 'diamond', 

570 '^': 'triangle', 

571 } 

572 

573 if marker in marker_map: 

574 bokeh_marker = marker_map[marker] 

575 # Small dots use smaller size and always filled 

576 if marker == '.': 

577 _figure.scatter(x, y, marker=bokeh_marker, **{**kwargs, 'size': size/2, 'fill_color': color}) 

578 else: 

579 _figure.scatter(x, y, marker=bokeh_marker, **kwargs) 

580 elif marker is not None: 

581 _warnings.warn(f'Bad marker type: {marker}') 

582 if not hold and not _hold: 

583 _show(_figure) 

584 _figure = None 

585 

586def image(img: Any, x: Any | None = None, y: Any | None = None, colormap: str = 'Plasma256', clim: Tuple[float, float | None] = None, clabel: str | None = None, title: str | None = None, xlabel: str | None = None, ylabel: str | None = None, xlim: Tuple[float, float | None] = None, ylim: Tuple[float, float | None] = None, xtype: str = 'auto', ytype: str = 'auto', width: int | None = None, height: int | None = None, hold: bool = False, interactive: bool | None = None) -> None: 

587 """Plot a heatmap of 2D scalar data. 

588 

589 :param img: 2D image data 

590 :param x: x-axis range for image data (min, max) 

591 :param y: y-axis range for image data (min, max) 

592 :param colormap: named color palette or Bokeh ColorMapper (see `Bokeh palettes <https://bokeh.pydata.org/en/latest/docs/reference/palettes.html>`_) 

593 :param clim: color axis limits (min, max) 

594 :param clabel: color axis label 

595 :param title: figure title 

596 :param xlabel: x-axis label 

597 :param ylabel: y-axis label 

598 :param xlim: x-axis limits (min, max) 

599 :param ylim: y-axis limits (min, max) 

600 :param xtype: x-axis type ('auto', 'linear', 'log', etc) 

601 :param ytype: y-axis type ('auto', 'linear', 'log', etc) 

602 :param width: figure width in pixels 

603 :param height: figure height in pixels 

604 :param interactive: enable interactive tools (pan, zoom, etc) for plot 

605 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

606 

607 >>> import arlpy.plot 

608 >>> import numpy as np 

609 >>> arlpy.plot.image(np.random.normal(size=(100,100)), colormap='Inferno256') 

610 """ 

611 global _figure 

612 if x is None: 

613 x = (0, img.shape[1]-1) 

614 if y is None: 

615 y = (0, img.shape[0]-1) 

616 if xlim is None: 

617 xlim = x 

618 if ylim is None: 

619 ylim = y 

620 if _figure is None: 

621 _figure = _new_figure(title, width, height, xlabel, ylabel, xlim, ylim, xtype, ytype, interactive) 

622 if clim is None: 

623 clim = [np.amin(img), np.amax(img)] 

624 if not isinstance(colormap, _bmodels.ColorMapper): 

625 colormap = _bmodels.LinearColorMapper(palette=colormap, low=clim[0], high=clim[1]) 

626 _figure.image([img], x=x[0], y=y[0], dw=x[-1]-x[0], dh=y[-1]-y[0], color_mapper=colormap) 

627 cbar = _bmodels.ColorBar(color_mapper=colormap, location=(0,0), title=clabel) 

628 _figure.add_layout(cbar, 'right') 

629 if not hold and not _hold: 

630 _show(_figure) 

631 _figure = None 

632 

633def vlines(x: Any, color: str = 'gray', style: str = 'dashed', thickness: int = 1, hold: bool = False) -> None: 

634 """Draw vertical lines on a plot. 

635 

636 :param x: x location of lines 

637 :param color: line color (see `Bokeh colors`_) 

638 :param style: line style ('solid', 'dashed', 'dotted', 'dotdash', 'dashdot') 

639 :param thickness: line width in pixels 

640 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

641 

642 >>> import arlpy.plot 

643 >>> arlpy.plot.plot([0, 20], [0, 10], hold=True) 

644 >>> arlpy.plot.vlines([7, 12]) 

645 """ 

646 global _figure 

647 if _figure is None: 

648 return 

649 x = np.asarray(x, dtype=np.float64) 

650 for j in range(x.size): 

651 _figure.add_layout(_bmodels.Span(location=x[j], dimension='height', line_color=color, line_dash=style, line_width=thickness)) 

652 if not hold and not _hold: 

653 _show(_figure) 

654 _figure = None 

655 

656def hlines(y: Any, color: str = 'gray', style: str = 'dashed', thickness: int = 1, hold: bool = False) -> None: 

657 """Draw horizontal lines on a plot. 

658 

659 :param y: y location of lines 

660 :param color: line color (see `Bokeh colors`_) 

661 :param style: line style ('solid', 'dashed', 'dotted', 'dotdash', 'dashdot') 

662 :param thickness: line width in pixels 

663 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

664 

665 >>> import arlpy.plot 

666 >>> arlpy.plot.plot([0, 20], [0, 10], hold=True) 

667 >>> arlpy.plot.hlines(3, color='red', style='dotted') 

668 """ 

669 global _figure 

670 if _figure is None: 

671 return 

672 y = np.asarray(y, dtype=np.float64) 

673 for j in range(y.size): 

674 _figure.add_layout(_bmodels.Span(location=y[j], dimension='width', line_color=color, line_dash=style, line_width=thickness)) 

675 if not hold and not _hold: 

676 _show(_figure) 

677 _figure = None 

678 

679def text(x: float, y: float, s: str, color: str = 'gray', size: str = '8pt', hold: bool = False) -> None: 

680 """Add text annotation to a plot. 

681 

682 :param x: x location of left of text 

683 :param y: y location of bottom of text 

684 :param s: text to add 

685 :param color: text color (see `Bokeh colors`_) 

686 :param size: text size (e.g. '12pt', '3em') 

687 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

688 

689 >>> import arlpy.plot 

690 >>> arlpy.plot.plot([0, 20], [0, 10], hold=True) 

691 >>> arlpy.plot.text(7, 3, 'demo', color='orange') 

692 """ 

693 global _figure 

694 if _figure is None: 

695 return 

696 _figure.add_layout(_bmodels.Label(x=x, y=y, text=s, text_font_size=size, text_color=color)) 

697 if not hold and not _hold: 

698 _show(_figure) 

699 _figure = None 

700 

701def box(left: float | None = None, right: float | None = None, top: float | None = None, bottom: float | None = None, color: str = 'yellow', alpha: float = 0.1, hold: bool = False) -> None: 

702 """Add a highlight box to a plot. 

703 

704 :param left: x location of left of box 

705 :param right: x location of right of box 

706 :param top: y location of top of box 

707 :param bottom: y location of bottom of box 

708 :param color: text color (see `Bokeh colors`_) 

709 :param alpha: transparency (0-1) 

710 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

711 

712 >>> import arlpy.plot 

713 >>> arlpy.plot.plot([0, 20], [0, 10], hold=True) 

714 >>> arlpy.plot.box(left=5, right=10, top=8) 

715 """ 

716 global _figure 

717 if _figure is None: 

718 return 

719 _figure.add_layout(_bmodels.BoxAnnotation(left=left, right=right, top=top, bottom=bottom, fill_color=color, fill_alpha=alpha)) 

720 if not hold and not _hold: 

721 _show(_figure) 

722 _figure = None 

723 

724def color(n: int) -> str: 

725 """Get a numbered color to cycle over a set of colors. 

726 

727 >>> import arlpy.plot 

728 >>> arlpy.plot.color(0) 

729 'blue' 

730 >>> arlpy.plot.color(1) 

731 'red' 

732 >>> arlpy.plot.plot([0, 20], [0, 10], color=arlpy.plot.color(3)) 

733 """ 

734 return _colors[n % len(_colors)] 

735 

736def set_colors(c: List[str]) -> None: 

737 """Provide a list of named colors to cycle over. 

738 

739 >>> import arlpy.plot 

740 >>> arlpy.plot.set_colors(['red', 'blue', 'green', 'black']) 

741 >>> arlpy.plot.color(2) 

742 'green' 

743 """ 

744 global _colors 

745 _colors = c 

746 

747def specgram(x: Any, fs: float = 2, nfft: int | None = None, noverlap: int | None = None, colormap: str = 'Plasma256', clim: Tuple[float, float | None] = None, clabel: str = 'dB', title: str | None = None, xlabel: str = 'Time (s)', ylabel: str = 'Frequency (Hz)', xlim: Tuple[float, float | None] = None, ylim: Tuple[float, float | None] = None, width: int | None = None, height: int | None = None, hold: bool = False, interactive: bool | None = None) -> None: 

748 """Plot spectrogram of a given time series signal. 

749 

750 :param x: time series signal 

751 :param fs: sampling rate 

752 :param nfft: FFT size (see `scipy.signal.spectrogram <https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.spectrogram.html>`_) 

753 :param noverlap: overlap size (see `scipy.signal.spectrogram`_) 

754 :param colormap: named color palette or Bokeh ColorMapper (see `Bokeh palettes`_) 

755 :param clim: color axis limits (min, max), or dynamic range with respect to maximum 

756 :param clabel: color axis label 

757 :param title: figure title 

758 :param xlabel: x-axis label 

759 :param ylabel: y-axis label 

760 :param xlim: x-axis limits (min, max) 

761 :param ylim: y-axis limits (min, max) 

762 :param width: figure width in pixels 

763 :param height: figure height in pixels 

764 :param interactive: enable interactive tools (pan, zoom, etc) for plot 

765 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

766 

767 >>> import arlpy.plot 

768 >>> import numpy as np 

769 >>> arlpy.plot.specgram(np.random.normal(size=(10000)), fs=10000, clim=30) 

770 """ 

771 f, t, Sxx = _sig.spectrogram(x, fs=fs, nperseg=nfft, noverlap=noverlap) 

772 Sxx = 10 * np.log10(Sxx + np.finfo(float).eps) 

773 

774 # Convert scalar clim to range (for dynamic range specification) 

775 if isinstance(clim, (int, float)): 

776 max_val = np.max(Sxx) 

777 clim = (max_val - clim, max_val) 

778 

779 image(Sxx, x=(t[0], t[-1]), y=(f[0], f[-1]), title=title, colormap=colormap, 

780 clim=clim, clabel=clabel, xlabel=xlabel, ylabel=ylabel, xlim=xlim, 

781 ylim=ylim, width=width, height=height, hold=hold, interactive=interactive) 

782 

783def psd(x: Any, fs: float = 2, nfft: int = 512, noverlap: int | None = None, window: str = 'hann', color: str | None = None, style: str = 'solid', thickness: int = 1, marker: str | None = None, filled: bool = False, size: int = 6, title: str | None = None, xlabel: str = 'Frequency (Hz)', ylabel: str = 'Power spectral density (dB/Hz)', xlim: Tuple[float, float | None] = None, ylim: Tuple[float, float | None] = None, width: int | None = None, height: int | None = None, legend: str | None = None, hold: bool = False, interactive: bool | None = None) -> None: 

784 """Plot power spectral density of a given time series signal. 

785 

786 :param x: time series signal 

787 :param fs: sampling rate 

788 :param nfft: segment size (see `scipy.signal.welch <https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.welch.html>`_) 

789 :param noverlap: overlap size (see `scipy.signal.welch`_) 

790 :param window: window to use (see `scipy.signal.welch`_) 

791 :param color: line color (see `Bokeh colors`_) 

792 :param style: line style ('solid', 'dashed', 'dotted', 'dotdash', 'dashdot') 

793 :param thickness: line width in pixels 

794 :param marker: point markers ('.', 'o', 's', '*', 'x', '+', 'd', '^') 

795 :param filled: filled markers or outlined ones 

796 :param size: marker size 

797 :param title: figure title 

798 :param xlabel: x-axis label 

799 :param ylabel: y-axis label 

800 :param xlim: x-axis limits (min, max) 

801 :param ylim: y-axis limits (min, max) 

802 :param width: figure width in pixels 

803 :param height: figure height in pixels 

804 :param legend: legend text 

805 :param interactive: enable interactive tools (pan, zoom, etc) for plot 

806 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

807 

808 >>> import arlpy.plot 

809 >>> import numpy as np 

810 >>> arlpy.plot.psd(np.random.normal(size=(10000)), fs=10000) 

811 """ 

812 f, Pxx = _sig.welch(x, fs=fs, nperseg=nfft, noverlap=noverlap, window=window) 

813 Pxx = 10 * np.log10(Pxx + np.finfo(float).eps) 

814 

815 # Set default axis limits if not specified 

816 xlim = xlim or (0, fs / 2) 

817 if ylim is None: 

818 max_pxx = np.max(Pxx) 

819 ylim = (max_pxx - 50, max_pxx + 10) 

820 

821 plot(f, Pxx, color=color, style=style, thickness=thickness, marker=marker, 

822 filled=filled, size=size, title=title, xlabel=xlabel, ylabel=ylabel, 

823 xlim=xlim, ylim=ylim, maxpts=len(f), width=width, height=height, 

824 hold=hold, legend=legend, interactive=interactive) 

825 

826def iqplot(data: Any, marker: str = '.', color: str | None = None, labels: Any | None = None, filled: bool = False, size: int | None = None, title: str | None = None, xlabel: str | None = None, ylabel: str | None = None, xlim: List[float] = [-2, 2], ylim: List[float] = [-2, 2], width: int | None = None, height: int | None = None, hold: bool = False, interactive: bool | None = None) -> None: 

827 """Plot signal points. 

828 

829 :param data: complex baseband signal points 

830 :param marker: point markers ('.', 'o', 's', '*', 'x', '+', 'd', '^') 

831 :param color: marker/text color (see `Bokeh colors`_) 

832 :param labels: label for each signal point, or True to auto-generate labels 

833 :param filled: filled markers or outlined ones 

834 :param size: marker/text size (e.g. 5, '8pt') 

835 :param title: figure title 

836 :param xlabel: x-axis label 

837 :param ylabel: y-axis label 

838 :param xlim: x-axis limits (min, max) 

839 :param ylim: y-axis limits (min, max) 

840 :param width: figure width in pixels 

841 :param height: figure height in pixels 

842 :param interactive: enable interactive tools (pan, zoom, etc) for plot 

843 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

844 

845 >>> import arlpy 

846 >>> import arlpy.plot 

847 >>> arlpy.plot.iqplot(arlpy.comms.psk(8)) 

848 >>> arlpy.plot.iqplot(arlpy.comms.qam(16), color='red', marker='x') 

849 >>> arlpy.plot.iqplot(arlpy.comms.psk(4), labels=['00', '01', '11', '10']) 

850 """ 

851 data = np.asarray(data, dtype=np.complex128) 

852 if not _hold: 

853 figure(title=title, xlabel=xlabel, ylabel=ylabel, xlim=xlim, ylim=ylim, width=width, height=height, interactive=interactive) 

854 if labels is None: 

855 if size is None: 

856 size = 5 

857 scatter(data.real, data.imag, marker=marker, filled=filled, color=color, size=size, hold=hold) 

858 else: 

859 if labels: 

860 labels = range(len(data)) 

861 if color is None: 

862 color = 'black' 

863 plot([0], [0], hold=True) 

864 for i in range(len(data)): 

865 text(data[i].real, data[i].imag, str(labels[i]), color=color, size=size, hold=True if i < len(data)-1 else hold) 

866 

867def freqz(b: Any, a: Any = 1, fs: float = 2.0, worN: int | None = None, whole: bool = False, degrees: bool = True, style: str = 'solid', thickness: int = 1, title: str | None = None, xlabel: str = 'Frequency (Hz)', xlim: Tuple[float, float | None] = None, ylim: Tuple[float, float | None] = None, width: int | None = None, height: int | None = None, hold: bool = False, interactive: bool | None = None) -> None: 

868 """Plot frequency response of a filter. 

869 

870 This is a convenience function to plot frequency response, and internally uses 

871 :func:`scipy.signal.freqz` to estimate the response. For further details, see the 

872 documentation for :func:`scipy.signal.freqz`. 

873 

874 :param b: numerator of a linear filter 

875 :param a: denominator of a linear filter 

876 :param fs: sampling rate in Hz (optional, normalized frequency if not specified) 

877 :param worN: see :func:`scipy.signal.freqz` 

878 :param whole: see :func:`scipy.signal.freqz` 

879 :param degrees: True to display phase in degrees, False for radians 

880 :param style: line style ('solid', 'dashed', 'dotted', 'dotdash', 'dashdot') 

881 :param thickness: line width in pixels 

882 :param title: figure title 

883 :param xlabel: x-axis label 

884 :param ylabel1: y-axis label for magnitude 

885 :param ylabel2: y-axis label for phase 

886 :param xlim: x-axis limits (min, max) 

887 :param ylim: y-axis limits (min, max) 

888 :param width: figure width in pixels 

889 :param height: figure height in pixels 

890 :param interactive: enable interactive tools (pan, zoom, etc) for plot 

891 :param hold: if set to True, output is not plotted immediately, but combined with the next plot 

892 

893 >>> import arlpy 

894 >>> arlpy.plot.freqz([1,1,1,1,1], fs=120000); 

895 """ 

896 w, h = _sig.freqz(b, a, worN, whole) 

897 Hxx = 20*np.log10(abs(h)+np.finfo(float).eps) 

898 f = w*fs/(2*np.pi) 

899 if xlim is None: 

900 xlim = (0, fs/2) 

901 if ylim is None: 

902 ylim = (np.max(Hxx)-50, np.max(Hxx)+10) 

903 figure(title=title, xlabel=xlabel, ylabel='Amplitude (dB)', xlim=xlim, ylim=ylim, width=width, height=height, interactive=interactive) 

904 _hold_enable(True) 

905 plot(f, Hxx, color=color(0), style=style, thickness=thickness, legend='Magnitude') 

906 fig = gcf() 

907 units = 180/np.pi if degrees else 1 

908 fig.extra_y_ranges = {'phase': _bmodels.Range1d(start=-np.pi*units, end=np.pi*units)} 

909 fig.add_layout(_bmodels.LinearAxis(y_range_name='phase', axis_label='Phase (degrees)' if degrees else 'Phase (radians)'), 'right') 

910 phase = np.angle(h)*units 

911 fig.line(f, phase, line_color=color(1), line_dash=style, line_width=thickness, legend_label='Phase', y_range_name='phase') 

912 _hold_enable(hold)