Coverage for python / aubellhop / pyplot.py: 62%
231 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-24 14:11 +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
21def pyplot_env2d(
22 env: Environment,
23 surface_color: str = 'dodgerblue',
24 bottom_color: str = 'peru',
25 source_color: str = 'orangered',
26 receiver_color: str = 'midnightblue',
27 receiver_plot: bool | None = None,
28 fill: bool | None = None,
29 ax: Any | None = None,
30 **kwargs: Any
31 ) -> None:
32 """Plots a visual representation of the environment with matplotlib.
34 Parameters
35 ----------
36 env : dict
37 Environment description
38 surface_color : str, default='dodgerblue'
39 Color of the surface (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
40 bottom_color : str, default='peru'
41 Color of the bottom (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
42 source_color : str, default='orangered'
43 Color of transmitters (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
44 receiver_color : str, default='midnightblue'
45 Color of receivers (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
46 receiver_plot : bool, optional
47 True to plot all receivers, False to not plot any receivers, None to automatically decide
48 **kwargs
49 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
51 Notes
52 -----
53 The surface, bottom, transmitters (marker: '*') and receivers (marker: 'o')
54 are plotted in the environment. If `receiver_plot` is set to None and there are
55 more than 2000 receivers, they are not plotted.
57 Examples
58 --------
59 >>> import aubellhop as bh
60 >>> env = bh.Environment(bottom_depth=[[0, 40], [100, 30], [500, 35], [700, 20], [1000,45]])
61 >>> bh.plot_env(env)
62 """
64 env.check()
66 if ax is None:
67 fig = _pyplt.figure()
68 ax = fig.add_subplot()
70 if np.array(env['receiver_range']).size > 1:
71 min_x = np.min(env['receiver_range'])
72 else:
73 min_x = 0
74 max_x = np.max(env['receiver_range'])
75 if max_x - min_x > 10000:
76 divisor = 1000
77 min_x /= divisor
78 max_x /= divisor
79 range_unit = ' (km)'
80 else:
81 divisor = 1
82 range_unit = ' (m)'
83 if np.size(env['surface_depth']) == 1:
84 min_y = 0
85 else:
86 min_y = np.min(env['surface_depth'][:, 1])
87 max_y = env['_depth_max']
88 mgn_x = 0.01 * (max_x - min_x)
89 mgn_y = 0.1 * (max_y - min_y)
91 if np.size(env['surface_depth']) == 1:
92 surface_x = [min_x, max_x]
93 surface_y = [0, 0]
94 else:
95 surface_x = env['surface_depth'][:, 0] / divisor
96 surface_y = env['surface_depth'][:, 1]
97 _pyplt.plot(surface_x, surface_y, color=surface_color, **kwargs)
99 if np.size(env['bottom_depth']) == 1:
100 _pyplt.plot([min_x, max_x], [env['bottom_depth'], env['bottom_depth']], color=bottom_color, **kwargs)
101 else:
102 _pyplt.plot(env['bottom_depth'][:, 0] / divisor, env['bottom_depth'][:, 1], color=bottom_color, **kwargs)
104 txd = env['source_depth']
105 _pyplt.plot([0] * np.size(txd), txd, marker='*', markersize=6, color=source_color, **kwargs)
107 if receiver_plot is None:
108 receiver_plot = np.size(env['receiver_depth']) * np.size(env['receiver_range']) < 2000
109 if receiver_plot:
110 rxr = env['receiver_range']
111 if np.size(rxr) == 1:
112 rxr = [rxr]
113 for r in np.array(rxr):
114 rxd = env['receiver_depth']
115 _pyplt.plot([r / divisor] * np.size(rxd), rxd, marker='o', color=receiver_color, **kwargs)
117 if fill:
118 y0 = 0.0
119 _pyplt.axhline(y0, color="w", linestyle="-")
120 _pyplt.fill_between(surface_x, surface_y, y0, color="w")
122 _pyplt.xlabel('Range'+range_unit)
123 _pyplt.ylabel('Depth (m)')
124 ax.yaxis.set_inverted(True)
125 _pyplt.xlim([min_x - mgn_x, max_x + mgn_x])
126 _pyplt.ylim([max_y + mgn_y, min_y - mgn_y])
128def pyplot_env3d(env: Environment, surface_color: str = 'dodgerblue', bottom_color: str = 'peru', source_color: str = 'orangered', receiver_color: str = 'midnightblue',
129 receiver_plot: bool | None = None, ax: Any | None = None, **kwargs: Any) -> None:
130 """Plots a visual representation of the environment with matplotlib.
131 """
133 env.check()
135 if ax is None:
136 fig = _pyplt.figure()
137 ax = fig.add_subplot(projection='3d')
139 if np.array(env['receiver_range']).size > 1:
140 min_x = np.min(env['receiver_range'])
141 else:
142 min_x = 0
143 max_x = env['simulation_range']
144 min_y = -env['simulation_cross_range']
145 max_y = +env['simulation_cross_range']
146 xdivisor = 1
147 ydivisor = 1
148 xrange_unit = ' (m)'
149 yrange_unit = ' (m)'
150 if max_x - min_x > 10000:
151 xdivisor = 1000
152 min_x /= xdivisor
153 max_x /= xdivisor
154 xrange_unit = ' (km)'
155 if max_y - min_y > 10000:
156 ydivisor = 1000
157 min_y /= ydivisor
158 max_y /= ydivisor
159 yrange_unit = ' (km)'
160 if np.size(env['surface_depth']) == 1:
161 min_z = 0
162 else:
163 min_z = np.min(env['surface_depth'][:, 1])
164 max_z = env['simulation_depth']
165 mgn_x = 0.01 * (max_x - min_x)
166 mgn_z = 0.1 * (max_z - min_z)
168 if np.size(env['surface_depth']) == 1:
169 z = float(env['surface_depth'])
170 X, Y = np.meshgrid([min_x, max_x], [min_y, max_y])
171 Z = np.full_like(X, z)
172 ax.plot_surface(X, Y, Z, color=surface_color, alpha=0.3, **kwargs)
173 else:
174 _pyplt.plot(env['surface_depth'][:, 0] / xdivisor, env['surface_depth'][:, 1], color=surface_color, **kwargs)
176 if np.size(env['bottom_depth']) == 1:
177 z = float(env['bottom_depth'])
178 X, Y = np.meshgrid([min_x, max_x], [min_y, max_y])
179 Z = np.full_like(X, z)
180 ax.plot_surface(X, Y, Z, color=bottom_color, alpha=0.3, **kwargs)
181 else:
182 _pyplt.plot(env['bottom_depth'][:, 0] / xdivisor, env['bottom_depth'][:, 1], color=bottom_color, **kwargs)
184 if env._source_num == 1:
185 _pyplt.plot(
186 env['source_range'] / xdivisor,
187 env['source_cross_range'] / ydivisor,
188 env['source_depth'],
189 marker='*',
190 markersize=6,
191 color=source_color,
192 **kwargs,
193 )
194 else:
195 print("MULTIPLE SOURCES NOT IMPLEMENTED YET")
197 if env._source_num == 1:
198 _pyplt.plot(
199 env['receiver_range'] * np.cos(env['receiver_bearing']) / xdivisor,
200 env['receiver_range'] * np.sin(env['receiver_bearing']) / ydivisor,
201 env['receiver_depth'],
202 marker='o',
203 markersize=6,
204 color=receiver_color,
205 **kwargs,
206 )
207 else:
208 print("MULTIPLE RECEIVERS NOT IMPLEMENTED YET")
210 ax.set_xlabel('Range'+xrange_unit)
211 ax.set_ylabel('Cross range'+yrange_unit)
212 ax.set_zlabel('Depth (m)')
213 ax.yaxis.set_inverted(True)
214 ax.set_xlim([min_x - mgn_x, max_x + mgn_x])
215 ax.set_ylim([min_y, max_y])
216 ax.set_zlim([max_z + mgn_z, min_z - mgn_z])
218def pyplot_ssp(env: Environment, ax: Any | None = None, **kwargs: Any) -> None:
219 """Plots the sound speed profile with matplotlib.
221 Parameters
222 ----------
223 env : Environment
224 Environment description
225 **kwargs
226 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
228 Notes
229 -----
230 If the sound speed profile is range-dependent, this function only plots the first profile.
232 Examples
233 --------
234 >>> import aubellhop as bh
235 >>> env = bh.Environment(soundspeed=[[ 0, 1540], [10, 1530], [20, 1532], [25, 1533], [30, 1535]])
236 >>> bh.plot_ssp(env)
237 """
239 if ax is None:
240 fig = _pyplt.figure()
241 ax = fig.add_subplot()
243 assert(isinstance(ax, Axes))
245 env.check()
246 svp = env['soundspeed']
247 if isinstance(svp, pd.DataFrame):
248 svp = np.hstack((np.array([svp.index]).T, np.asarray(svp)))
249 if np.size(svp) == 1:
250 if np.size(env['bottom_depth']) > 1:
251 max_y = np.max(env['bottom_depth'][:, 1])
252 else:
253 max_y = env['bottom_depth']
254 _pyplt.plot([svp, svp], [0, -max_y], **kwargs)
255 _pyplt.xlabel('Soundspeed (m/s)')
256 _pyplt.ylabel('Depth (m)')
257 elif env['soundspeed_interp'] == BHStrings.spline:
258 ynew = np.linspace(np.min(svp[:, 0]), np.max(svp[:, 0]), 100)
259 tck = _interp.splrep(svp[:, 0], svp[:, 1], s=0)
260 xnew = _interp.splev(ynew, tck, der=0)
261 _pyplt.plot(xnew, -ynew, **kwargs)
262 _pyplt.xlabel('Soundspeed (m/s)')
263 _pyplt.ylabel('Depth (m)')
264 _pyplt.plot(svp[:, 1], -svp[:, 0], marker='.', **kwargs)
265 else:
266 for i in range(svp.shape[1]-1):
267 _pyplt.plot(svp[:, i+1], -svp[:, 0], **kwargs)
268 _pyplt.xlabel('Soundspeed (m/s)')
269 _pyplt.ylabel('Depth (m)')
271def pyplot_arrivals(arrivals: Any, dB: bool = False, color: str = 'blue', **kwargs: Any) -> None:
272 """Plots the arrival times and amplitudes with matplotlib.
274 Parameters
275 ----------
276 arrivals : pandas.DataFrame
277 Arrivals times (s) and coefficients
278 dB : bool, default=False
279 True to plot in dB, False for linear scale
280 color : str, default='blue'
281 Line color (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
282 **kwargs
283 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
285 Examples
286 --------
287 >>> import aubellhop as bh
288 >>> env = bh.Environment()
289 >>> arrivals = bh.compute_arrivals(env)
290 >>> bh.plot_arrivals(arrivals)
291 """
292 t0 = min(arrivals.time_of_arrival)
293 t1 = max(arrivals.time_of_arrival)
294 if dB:
295 min_y = 20 * np.log10(np.max(np.abs(arrivals.arrival_amplitude))) - 60
296 ylabel = 'Amplitude (dB)'
297 else:
298 ylabel = 'Amplitude'
299 _pyplt.plot([t0, t1], [0, 0], color=color, **kwargs)
300 _pyplt.xlabel('Arrival time (s)')
301 _pyplt.ylabel(ylabel)
302 min_y = 0
303 for _, row in arrivals.iterrows():
304 t = row.time_of_arrival.real
305 y = np.abs(row.arrival_amplitude)
306 if dB:
307 y = max(20 * np.log10(_fi.epsilon + y), min_y)
308 _pyplt.plot([t, t], [min_y, y], color=color, **kwargs)
309 _pyplt.xlabel('Arrival time (s)')
310 _pyplt.ylabel(ylabel)
312def pyplot_rays(
313 rays: Any,
314 env: Environment | None = None,
315 invert_colors: bool = False,
316 ax: Any | None = None,
317 **kwargs: Any
318 ) -> Axes:
319 """Plots ray paths with matplotlib
321 Parameters
322 ----------
323 rays : pandas.DataFrame
324 Ray paths
325 env : Environment, optional
326 Environment definition
327 invert_colors : bool, default=False
328 False to use black for high intensity rays, True to use white
329 **kwargs
330 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
332 Notes
333 -----
334 If environment definition is provided, it is overlayed over this plot using default
335 parameters for `bellhop.plot_env()`. Without an environment file, no axis labels etc
336 are provided, you are in charge of that.
338 Examples
339 --------
340 >>> import aubellhop as bh
341 >>> env = bh.Environment()
342 >>> rays = bh.compute_eigenrays(env)
343 >>> bh.plot_rays(rays, width=1000)
344 """
345 if env is not None:
346 env.check()
348 rays = rays.sort_values('bottom_bounces', ascending=False)
349 dim = rays["ray"].iloc[0][0].shape[0]
351 if ax is None:
352 fig = _pyplt.figure()
353 if dim == 2:
354 ax = fig.add_subplot()
355 elif dim == 3:
356 ax = fig.add_subplot(projection='3d')
357 assert(isinstance(ax, Axes))
359 max_amp = np.max(np.abs(rays.bottom_bounces)) if len(rays.bottom_bounces) > 0 else 0
360 if max_amp <= 0:
361 max_amp = 1
362 divisor = 1
363 r = []
364 for _, row in rays.iterrows():
365 r += list(row.ray[:, 0])
366 if max(r) - min(r) > 10000:
367 divisor = 1000
368 for _, row in rays.iterrows():
369 rr = float( row.bottom_bounces / (max_amp + 1) ) # avoid rr = 1 == 100% white
370 c = 1.0 - rr if invert_colors else rr
371 cmap = _pyplt.get_cmap("gray")
372 col_str = _mplc.to_hex(cmap(c))
373 if dim == 2:
374 if "color" in kwargs.keys():
375 ax.plot(row.ray[:, 0] / divisor, row.ray[:, 1], **kwargs)
376 else:
377 ax.plot(row.ray[:, 0] / divisor, row.ray[:, 1], color=col_str, **kwargs)
378 if dim == 3:
379 if "color" in kwargs.keys():
380 ax.plot(row.ray[:, 0] / divisor, row.ray[:, 1], row.ray[:, 2], **kwargs)
381 else:
382 ax.plot(row.ray[:, 0] / divisor, row.ray[:, 1], row.ray[:, 2], color=col_str, **kwargs)
383 if env is not None:
384 if dim == 2:
385 pyplot_env2d(env,ax=ax,receiver_plot=False)
386 elif dim == 3:
387 pyplot_env3d(env,ax=ax)
389 return ax
391def pyplot_transmission_loss(
392 tloss: Any,
393 env: Environment | None = None,
394 ax: Any | None = None,
395 **kwargs: Any
396 ) -> Axes:
397 """Plots transmission loss with matplotlib.
399 Parameters
400 ----------
401 tloss : pandas.DataFrame
402 Complex transmission loss
403 env : Environment, optional
404 Environment definition
405 **kwargs
406 Other keyword arguments applicable for `bellhop.plot.image()` are also supported
408 Notes
409 -----
410 If environment definition is provided, it is overlayed over this plot using default
411 parameters for `bellhop.plot_env()`.
413 Examples
414 --------
415 >>> import aubellhop as bh
416 >>> import numpy as np
417 >>> env = bh.Environment(
418 receiver_depth=np.arange(0, 25),
419 receiver_range=np.arange(0, 1000),
420 beam_angle_min=-45,
421 beam_angle_max=45
422 )
423 >>> tloss = bh.compute_transmission_loss(env)
424 >>> bh.plot_transmission_loss(tloss, width=1000)
425 """
426 if env is not None:
427 env.check()
429 if ax is None:
430 fig = _pyplt.figure()
431 ax = fig.add_subplot()
432 assert(isinstance(ax, Axes))
434 xr = (min(tloss.columns), max(tloss.columns))
435 yr = (max(tloss.index), min(tloss.index))
436 xlabel = 'Range (m)'
437 if xr[1] - xr[0] > 10000:
438 xr = (min(tloss.columns) / 1000, max(tloss.columns) / 1000)
439 xlabel = 'Range (km)'
441 trans_loss = 20 * np.log10(_fi.epsilon + np.abs(np.flipud(np.array(tloss))))
442 x_mesh, y_mesh = np.meshgrid(np.linspace(xr[0], xr[1], trans_loss.shape[1]),
443 np.linspace(yr[0], yr[1], trans_loss.shape[0]))
445 vmin = kwargs.get("vmin", None)
446 vmax = kwargs.get("vmax", None)
447 trans_loss = np.clip(trans_loss, vmin, vmax)
449 _pyplt.contourf(x_mesh, y_mesh, trans_loss, cmap="jet", **kwargs)
450 _pyplt.xlabel(xlabel)
451 _pyplt.ylabel('Depth (m)')
452 _pyplt.colorbar(label="Transmission loss (dB)")
453 if env is not None:
454 pyplot_env2d(env, ax=ax, receiver_plot=False, fill=True)
456 return ax
459### Export module names for auto-importing in __init__.py
461__all__ = [
462 name for name in globals() if not name.startswith("_") # ignore private names
463]