Coverage for python/aubellhop/pyplot.py: 61%
272 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-08 07:28 +0000
1"""Plotting functions using Matplotlib for aubellhop.
2"""
4from __future__ import annotations
6from typing import Any
7from sys import float_info as _fi
9import numpy as np
10import scipy.interpolate as _interp
11import pandas as pd
13import matplotlib.pyplot as _pyplt
14import matplotlib.colors as _mplc
15from matplotlib.axes import Axes
17from .constants import BHStrings
18from .environment import Environment
21# Scale factors from seconds, for the `time_units` argument of `pyplot_arrivals`
22_TIME_UNITS = {'s': 1.0, 'ms': 1e3, 'us': 1e6, 'µs': 1e6, 'ns': 1e9}
25def _resolve_time_units(time_units: str, tmax: float) -> tuple[str, float]:
26 """Resolve a time unit name to its name and scale factor from seconds.
28 'auto' selects the largest unit that keeps the maximum time at or above one.
29 """
30 if time_units == 'auto':
31 if not np.isfinite(tmax) or tmax <= 0:
32 return 's', _TIME_UNITS['s']
33 for unit in ('s', 'ms', 'us'):
34 if tmax * _TIME_UNITS[unit] >= 1:
35 return unit, _TIME_UNITS[unit]
36 return 'ns', _TIME_UNITS['ns']
37 if time_units not in _TIME_UNITS:
38 raise ValueError(f"Unknown time_units {time_units!r}; expected 'auto' or one of {sorted(_TIME_UNITS)}")
39 return time_units, _TIME_UNITS[time_units]
42def _bounce_shades(color: Any, nbounce: Any, lightest: float, scale: str) -> Any:
43 """Tint `color` towards white in proportion to bounce count, one RGB row per arrival.
45 The tint is normalised against the largest bounce count present, so the palest
46 arrival sits `lightest` of the way to white whether the maximum is 3 or 30.
47 """
48 base = np.asarray(_mplc.to_rgb(color))
49 nb = np.asarray(nbounce, dtype=float)
50 nmax = float(nb.max()) if nb.size else 0.0
51 if nmax <= 0:
52 frac = np.zeros_like(nb)
53 elif scale == 'linear':
54 frac = nb / nmax
55 else:
56 frac = np.log1p(nb) / np.log1p(nmax)
57 return base + lightest * frac[:, None] * (1.0 - base)
60def pyplot_env2d(
61 env: Environment,
62 surface_color: str = 'dodgerblue',
63 bottom_color: str = 'peru',
64 source_color: str = 'orangered',
65 receiver_color: str = 'midnightblue',
66 receiver_plot: bool | None = None,
67 fill: bool | None = None,
68 ax: Any | None = None,
69 **kwargs: Any
70 ) -> None:
71 """Plots a visual representation of the environment with matplotlib.
73 Parameters
74 ----------
75 env : dict
76 Environment description
77 surface_color : str, default='dodgerblue'
78 Color of the surface (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
79 bottom_color : str, default='peru'
80 Color of the bottom (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
81 source_color : str, default='orangered'
82 Color of transmitters (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
83 receiver_color : str, default='midnightblue'
84 Color of receivers (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
85 receiver_plot : bool, optional
86 True to plot all receivers, False to not plot any receivers, None to automatically decide
87 **kwargs
88 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
90 Notes
91 -----
92 The surface, bottom, transmitters (marker: '*') and receivers (marker: 'o')
93 are plotted in the environment. If `receiver_plot` is set to None and there are
94 more than 2000 receivers, they are not plotted.
96 Examples
97 --------
98 >>> import aubellhop as bh
99 >>> env = bh.Environment(bottom_depth=[[0, 40], [100, 30], [500, 35], [700, 20], [1000,45]])
100 >>> bh.plot_env(env)
101 """
103 env.check()
105 if ax is None:
106 fig = _pyplt.figure()
107 ax = fig.add_subplot()
109 if np.array(env['receiver_range']).size > 1:
110 min_x = np.min(env['receiver_range'])
111 else:
112 min_x = 0
113 max_x = np.max(env['receiver_range'])
114 if max_x - min_x > 10000:
115 divisor = 1000
116 min_x /= divisor
117 max_x /= divisor
118 range_unit = ' (km)'
119 else:
120 divisor = 1
121 range_unit = ' (m)'
122 if np.size(env['surface_depth']) == 1:
123 min_y = 0
124 else:
125 min_y = np.min(env['surface_depth'][:, 1])
126 max_y = env['_depth_max']
127 mgn_x = 0.01 * (max_x - min_x)
128 mgn_y = 0.1 * (max_y - min_y)
130 if np.size(env['surface_depth']) == 1:
131 surface_x = [min_x, max_x]
132 surface_y = [0, 0]
133 else:
134 surface_x = env['surface_depth'][:, 0] / divisor
135 surface_y = env['surface_depth'][:, 1]
136 _pyplt.plot(surface_x, surface_y, color=surface_color, **kwargs)
138 if np.size(env['bottom_depth']) == 1:
139 _pyplt.plot([min_x, max_x], [env['bottom_depth'], env['bottom_depth']], color=bottom_color, **kwargs)
140 else:
141 _pyplt.plot(env['bottom_depth'][:, 0] / divisor, env['bottom_depth'][:, 1], color=bottom_color, **kwargs)
143 txd = env['source_depth']
144 _pyplt.plot([0] * np.size(txd), txd, marker='*', markersize=6, color=source_color, **kwargs)
146 if receiver_plot is None:
147 receiver_plot = np.size(env['receiver_depth']) * np.size(env['receiver_range']) < 2000
148 if receiver_plot:
149 rxr = env['receiver_range']
150 if np.size(rxr) == 1:
151 rxr = [rxr]
152 for r in np.array(rxr):
153 rxd = env['receiver_depth']
154 _pyplt.plot([r / divisor] * np.size(rxd), rxd, marker='o', color=receiver_color, **kwargs)
156 if fill:
157 y0 = 0.0
158 _pyplt.axhline(y0, color="w", linestyle="-")
159 _pyplt.fill_between(surface_x, surface_y, y0, color="w")
161 _pyplt.xlabel('Range'+range_unit)
162 _pyplt.ylabel('Depth (m)')
163 ax.yaxis.set_inverted(True)
164 _pyplt.xlim((min_x - mgn_x, max_x + mgn_x))
165 _pyplt.ylim((max_y + mgn_y, min_y - mgn_y))
167def pyplot_env3d(env: Environment, surface_color: str = 'dodgerblue', bottom_color: str = 'peru', source_color: str = 'orangered', receiver_color: str = 'midnightblue',
168 receiver_plot: bool | None = None, ax: Any | None = None, **kwargs: Any) -> None:
169 """Plots a visual representation of the environment with matplotlib.
170 """
172 env.check()
174 if ax is None:
175 fig = _pyplt.figure()
176 ax = fig.add_subplot(projection='3d')
178 if np.array(env['receiver_range']).size > 1:
179 min_x = np.min(env['receiver_range'])
180 else:
181 min_x = 0
182 max_x = env['simulation_range']
183 min_y = -env['simulation_cross_range']
184 max_y = +env['simulation_cross_range']
185 xdivisor = 1
186 ydivisor = 1
187 xrange_unit = ' (m)'
188 yrange_unit = ' (m)'
189 if max_x - min_x > 10000:
190 xdivisor = 1000
191 min_x /= xdivisor
192 max_x /= xdivisor
193 xrange_unit = ' (km)'
194 if max_y - min_y > 10000:
195 ydivisor = 1000
196 min_y /= ydivisor
197 max_y /= ydivisor
198 yrange_unit = ' (km)'
199 if np.size(env['surface_depth']) == 1:
200 min_z = 0
201 else:
202 min_z = np.min(env['surface_depth'][:, 1])
203 max_z = env['simulation_depth']
204 mgn_x = 0.01 * (max_x - min_x)
205 mgn_z = 0.1 * (max_z - min_z)
207 if np.size(env['surface_depth']) == 1:
208 z = float(env['surface_depth'])
209 X, Y = np.meshgrid([min_x, max_x], [min_y, max_y])
210 Z = np.full_like(X, z)
211 ax.plot_surface(X, Y, Z, color=surface_color, alpha=0.3, **kwargs)
212 else:
213 _pyplt.plot(env['surface_depth'][:, 0] / xdivisor, env['surface_depth'][:, 1], color=surface_color, **kwargs)
215 if np.size(env['bottom_depth']) == 1:
216 z = float(env['bottom_depth'])
217 X, Y = np.meshgrid([min_x, max_x], [min_y, max_y])
218 Z = np.full_like(X, z)
219 ax.plot_surface(X, Y, Z, color=bottom_color, alpha=0.3, **kwargs)
220 else:
221 _pyplt.plot(env['bottom_depth'][:, 0] / xdivisor, env['bottom_depth'][:, 1], color=bottom_color, **kwargs)
223 if env._source_num == 1:
224 _pyplt.plot(
225 env['source_range'] / xdivisor,
226 env['source_cross_range'] / ydivisor,
227 env['source_depth'],
228 marker='*',
229 markersize=6,
230 color=source_color,
231 **kwargs,
232 )
233 else:
234 print("MULTIPLE SOURCES NOT IMPLEMENTED YET")
236 if env._source_num == 1:
237 _pyplt.plot(
238 env['receiver_range'] * np.cos(env['receiver_bearing']) / xdivisor,
239 env['receiver_range'] * np.sin(env['receiver_bearing']) / ydivisor,
240 env['receiver_depth'],
241 marker='o',
242 markersize=6,
243 color=receiver_color,
244 **kwargs,
245 )
246 else:
247 print("MULTIPLE RECEIVERS NOT IMPLEMENTED YET")
249 ax.set_xlabel('Range'+xrange_unit)
250 ax.set_ylabel('Cross range'+yrange_unit)
251 ax.set_zlabel('Depth (m)')
252 ax.yaxis.set_inverted(True)
253 ax.set_xlim([min_x - mgn_x, max_x + mgn_x])
254 ax.set_ylim([min_y, max_y])
255 ax.set_zlim([max_z + mgn_z, min_z - mgn_z])
257def pyplot_ssp(env: Environment, ax: Any | None = None, **kwargs: Any) -> None:
258 """Plots the sound speed profile with matplotlib.
260 Parameters
261 ----------
262 env : Environment
263 Environment description
264 **kwargs
265 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
267 Notes
268 -----
269 If the sound speed profile is range-dependent, this function only plots the first profile.
271 Examples
272 --------
273 >>> import aubellhop as bh
274 >>> env = bh.Environment(soundspeed=[[ 0, 1540], [10, 1530], [20, 1532], [25, 1533], [30, 1535]])
275 >>> bh.plot_ssp(env)
276 """
278 if ax is None:
279 fig = _pyplt.figure()
280 ax = fig.add_subplot()
282 assert(isinstance(ax, Axes))
284 env.check()
285 svp = env['soundspeed']
286 if isinstance(svp, pd.DataFrame):
287 svp = np.hstack((np.array([svp.index]).T, np.asarray(svp)))
288 if np.size(svp) == 1:
289 if np.size(env['bottom_depth']) > 1:
290 max_y = np.max(env['bottom_depth'][:, 1])
291 else:
292 max_y = env['bottom_depth']
293 _pyplt.plot([svp, svp], [0, -max_y], **kwargs)
294 _pyplt.xlabel('Soundspeed (m/s)')
295 _pyplt.ylabel('Depth (m)')
296 elif env['soundspeed_interp'] == BHStrings.spline:
297 ynew = np.linspace(np.min(svp[:, 0]), np.max(svp[:, 0]), 100)
298 tck = _interp.splrep(svp[:, 0], svp[:, 1], s=0)
299 xnew = _interp.splev(ynew, tck, der=0)
300 _pyplt.plot(xnew, -ynew, **kwargs)
301 _pyplt.xlabel('Soundspeed (m/s)')
302 _pyplt.ylabel('Depth (m)')
303 _pyplt.plot(svp[:, 1], -svp[:, 0], marker='.', **kwargs)
304 else:
305 for i in range(svp.shape[1]-1):
306 _pyplt.plot(svp[:, i+1], -svp[:, 0], **kwargs)
307 _pyplt.xlabel('Soundspeed (m/s)')
308 _pyplt.ylabel('Depth (m)')
310def pyplot_arrivals(
311 arrivals: Any,
312 dB: bool = False,
313 ax: Any | None = None,
314 color: str = 'blue',
315 baseline: float | None = 0.0,
316 time_units: str = 'auto',
317 bounce_shading: str | None = 'linear',
318 lightest: float = 0.7,
319 colorbar: bool = False,
320 **kwargs: Any) -> None:
321 """Plots the arrival times and amplitudes with matplotlib.
323 Parameters
324 ----------
325 arrivals : pandas.DataFrame
326 Arrivals times (s) and coefficients
327 dB : bool, default=False
328 True to plot in dB, False for linear scale
329 color : str, default='blue'
330 Line color (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
331 baseline : float, optional, default=0.0
332 Amplitude at which to draw a horizontal reference line spanning the plot
333 (equivalent to Matlab's `yline`). None to omit the line.
334 time_units : str, default='auto'
335 Units for the time axis: 'auto', or one of 's', 'ms', 'us' (or 'µs'), 'ns'.
336 Both the scaling of the arrival times and the axis label follow from this.
337 'auto' picks the largest unit keeping the latest arrival at or above one;
338 pin it explicitly to keep the axis consistent across several plots.
339 bounce_shading : str, optional, default='linear'
340 Lighten each arrival towards white in proportion to its total number of
341 surface and bottom bounces: 'linear' for an equal step per bounce, 'log'
342 to spread the low-order arrivals out when the bounce count spans a wide
343 range. None to draw every arrival in `color`. The shading is normalised
344 against the largest bounce count present, so the range of shades is the
345 same whether the maximum is 3 bounces or 30.
346 lightest : float, default=0.7
347 How far towards white the most-bounced arrival is drawn, from 0 to 1.
348 colorbar : bool, default=False
349 Draw a discrete colorbar keying each shade to its bounce count. Requires
350 `bounce_shading`.
351 **kwargs
352 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
354 Examples
355 --------
356 >>> import aubellhop as bh
357 >>> env = bh.Environment()
358 >>> arrivals = bh.compute_arrivals(env)
359 >>> bh.plot_arrivals(arrivals)
360 """
361 if bounce_shading is not None and bounce_shading not in ('linear', 'log'):
362 raise ValueError(f"Unknown bounce_shading {bounce_shading!r}; expected 'linear', 'log' or None")
363 if colorbar and bounce_shading is None:
364 raise ValueError("colorbar=True has nothing to key without bounce_shading")
366 times = np.real(np.asarray(arrivals.time_of_arrival))
367 tmax = float(np.max(np.abs(times))) if times.size else 0.0
368 time_units, tscale = _resolve_time_units(time_units, tmax)
370 shades = None
371 nbounce = np.zeros(0)
372 if bounce_shading is not None:
373 nbounce = np.asarray(arrivals.surface_bounces + arrivals.bottom_bounces, dtype=float)
374 shades = _bounce_shades(color, nbounce, lightest, bounce_shading)
376 if ax is None:
377 fig = _pyplt.figure()
378 ax = fig.add_subplot()
380 ylabel = 'Amplitude, dB' if dB else 'Amplitude'
381 if baseline is not None:
382 ax.axhline(baseline, color='black', linewidth=0.8, zorder=0)
384 for j, (_, row) in enumerate(arrivals.iterrows()):
385 t = row.time_of_arrival.real * tscale
386 y = np.abs(row.arrival_amplitude)
387 if dB:
388 y = 20 * np.log10(_fi.epsilon + y)
389 c = color if shades is None else shades[j]
390 ax.plot([t, t], [baseline, y], color=c, **kwargs)
391 ax.plot(t, y, color=c, marker='.', **kwargs)
393 ax.set_xlabel(f'Arrival time, {time_units}')
394 ax.set_ylabel(ylabel)
396 if colorbar and bounce_shading is not None:
397 # One colormap entry per whole bounce count, so the bar matches the stems exactly
398 nmax = int(nbounce.max()) if nbounce.size else 0
399 cmap = _mplc.ListedColormap(_bounce_shades(color, np.arange(nmax + 1), lightest, bounce_shading))
400 norm = _mplc.BoundaryNorm(np.arange(nmax + 2) - 0.5, cmap.N)
401 cbar = ax.figure.colorbar(_pyplt.cm.ScalarMappable(cmap=cmap, norm=norm), ax=ax, label='Bounces')
402 cbar.set_ticks(np.arange(0, nmax + 1, max(1, int(np.ceil((nmax + 1) / 11)))).tolist())
403 cbar.ax.invert_yaxis() # darkest (direct) arrival at the top, matching the tallest stems
405def pyplot_rays(
406 rays: Any,
407 env: Environment | None = None,
408 invert_colors: bool = False,
409 ax: Any | None = None,
410 **kwargs: Any
411 ) -> Axes:
412 """Plots ray paths with matplotlib
414 Parameters
415 ----------
416 rays : pandas.DataFrame
417 Ray paths
418 env : Environment, optional
419 Environment definition
420 invert_colors : bool, default=False
421 False to use black for high intensity rays, True to use white
422 **kwargs
423 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
425 Notes
426 -----
427 If environment definition is provided, it is overlayed over this plot using default
428 parameters for `bellhop.plot_env()`. Without an environment file, no axis labels etc
429 are provided, you are in charge of that.
431 Examples
432 --------
433 >>> import aubellhop as bh
434 >>> env = bh.Environment()
435 >>> rays = bh.compute_eigenrays(env)
436 >>> bh.plot_rays(rays, width=1000)
437 """
438 if env is not None:
439 env.check()
441 rays = rays.sort_values('bottom_bounces', ascending=False)
442 dim = rays["ray"].iloc[0][0].shape[0]
444 if ax is None:
445 fig = _pyplt.figure()
446 if dim == 2:
447 ax = fig.add_subplot()
448 elif dim == 3:
449 ax = fig.add_subplot(projection='3d')
450 assert(isinstance(ax, Axes))
452 max_amp = np.max(np.abs(rays.bottom_bounces)) if len(rays.bottom_bounces) > 0 else 0
453 if max_amp <= 0:
454 max_amp = 1
455 divisor = 1
456 r = []
457 for _, row in rays.iterrows():
458 r += list(row.ray[:, 0])
459 if max(r) - min(r) > 10000:
460 divisor = 1000
461 for _, row in rays.iterrows():
462 rr = float( row.bottom_bounces / (max_amp + 1) ) # avoid rr = 1 == 100% white
463 c = 1.0 - rr if invert_colors else rr
464 cmap = _pyplt.get_cmap("gray")
465 col_str = _mplc.to_hex(cmap(c))
466 if dim == 2:
467 if "color" in kwargs.keys():
468 ax.plot(row.ray[:, 0] / divisor, row.ray[:, 1], **kwargs)
469 else:
470 ax.plot(row.ray[:, 0] / divisor, row.ray[:, 1], color=col_str, **kwargs)
471 if dim == 3:
472 if "color" in kwargs.keys():
473 ax.plot(row.ray[:, 0] / divisor, row.ray[:, 1], row.ray[:, 2], **kwargs)
474 else:
475 ax.plot(row.ray[:, 0] / divisor, row.ray[:, 1], row.ray[:, 2], color=col_str, **kwargs)
476 if env is not None:
477 if dim == 2:
478 pyplot_env2d(env,ax=ax,receiver_plot=False)
479 elif dim == 3:
480 pyplot_env3d(env,ax=ax)
482 return ax
484def pyplot_transmission_loss(
485 tloss: Any,
486 env: Environment | None = None,
487 ax: Any | None = None,
488 vmin: float | None = None,
489 vmax: float | None = None,
490 **kwargs: Any
491 ) -> Axes:
492 """Plots transmission loss with matplotlib.
494 Parameters
495 ----------
496 tloss : pandas.DataFrame
497 Complex transmission loss
498 env : Environment, optional
499 Environment definition
500 vmin, vmax : float, optional
501 Colour limits in dB (equivalent to Matlab's `clim`). Values outside the
502 range saturate at the end colours. Ignored if `levels` is passed explicitly.
503 **kwargs
504 Other keyword arguments applicable for `bellhop.plot.image()` are also supported
506 Notes
507 -----
508 If environment definition is provided, it is overlayed over this plot using default
509 parameters for `bellhop.plot_env()`.
511 Examples
512 --------
513 >>> import aubellhop as bh
514 >>> import numpy as np
515 >>> env = bh.Environment(
516 receiver_depth=np.arange(0, 25),
517 receiver_range=np.arange(0, 1000),
518 beam_angle_min=-45,
519 beam_angle_max=45
520 )
521 >>> tloss = bh.compute_transmission_loss(env)
522 >>> bh.plot_transmission_loss(tloss, width=1000)
523 """
524 if env is not None:
525 env.check()
527 if ax is None:
528 fig = _pyplt.figure()
529 ax = fig.add_subplot()
530 assert(isinstance(ax, Axes))
532 xr = (min(tloss.columns), max(tloss.columns))
533 yr = (max(tloss.index), min(tloss.index))
534 xlabel = 'Range (m)'
535 if xr[1] - xr[0] > 10000:
536 xr = (min(tloss.columns) / 1000, max(tloss.columns) / 1000)
537 xlabel = 'Range (km)'
539 trans_loss = 20 * np.log10(_fi.epsilon + np.abs(np.flipud(np.array(tloss))))
540 x_mesh, y_mesh = np.meshgrid(np.linspace(xr[0], xr[1], trans_loss.shape[1]),
541 np.linspace(yr[0], yr[1], trans_loss.shape[0]))
543 if vmin is not None or vmax is not None:
544 lo = vmin if vmin is not None else trans_loss.min()
545 hi = vmax if vmax is not None else trans_loss.max()
546 kwargs.setdefault("levels", np.linspace(lo, hi, 21))
547 kwargs.setdefault("extend", "both")
549 _pyplt.contourf(x_mesh, y_mesh, trans_loss, cmap="jet", **kwargs)
550 _pyplt.xlabel(xlabel)
551 _pyplt.ylabel('Depth (m)')
552 _pyplt.colorbar(label="Transmission loss (dB)")
553 if env is not None:
554 pyplot_env2d(env, ax=ax, receiver_plot=False, fill=True)
556 return ax
559### Export module names for auto-importing in __init__.py
561__all__ = [
562 name for name in globals() if not name.startswith("_") # ignore private names
563]