Coverage for python / aubellhop / plot.py: 100%
127 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-24 14:11 +0000
1"""Plotting functions 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
16from .environment import Environment
17from .constants import BHStrings
18from .plotutils import figure as figure
20import aubellhop.plotutils as _plt
23def plot_env(env: Environment,
24 surface_color: str = 'dodgerblue',
25 bottom_color: str = 'peru',
26 source_color: str = 'orangered',
27 receiver_color: str = 'midnightblue',
28 receiver_plot: bool | None = None,
29 **kwargs: Any
30 ) -> None:
31 """Plots a visual representation of the environment.
33 Parameters
34 ----------
35 env : dict
36 Environment description
37 surface_color : str, default='dodgerblue'
38 Color of the surface (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
39 bottom_color : str, default='peru'
40 Color of the bottom (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
41 source_color : str, default='orangered'
42 Color of transmitters (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
43 receiver_color : str, default='midnightblue'
44 Color of receivers (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
45 receiver_plot : bool, optional
46 True to plot all receivers, False to not plot any receivers, None to automatically decide
47 **kwargs
48 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
50 Notes
51 -----
52 The surface, bottom, transmitters (marker: '*') and receivers (marker: 'o')
53 are plotted in the environment. If `receiver_plot` is set to None and there are
54 more than 2000 receivers, they are not plotted.
56 Examples
57 --------
58 >>> import aubellhop as bh
59 >>> env = bh.Environment(bottom_depth=[[0, 40], [100, 30], [500, 35], [700, 20], [1000,45]])
60 >>> bh.plot_env(env)
61 """
63 if env is not None:
64 env.check()
66 min_x = 0.0
67 max_x = float(np.max(env['receiver_range']))
68 if max_x-min_x > 10000:
69 divisor = 1000.0
70 min_x /= divisor
71 max_x /= divisor
72 xlabel = 'Range (km)'
73 else:
74 divisor = 1.0
75 xlabel = 'Range (m)'
76 if np.size(env['surface_depth']) == 1:
77 min_y = 0
78 else:
79 min_y = np.min(env['surface_depth'][:,1])
80 max_y = env['_depth_max']
81 mgn_x = 0.01*(max_x-min_x)
82 mgn_y = 0.1*(max_y-min_y)
84 oh = _plt.hold()
85 if np.size(env['surface_depth']) == 1:
86 xx = [min_x, max_x]
87 yy = [0, 0]
88 else:
89 # linear and curvilinear options use the same altimetry, just with different normals
90 s = env['surface_depth']
91 xx = s[:,0]/divisor
92 yy = -s[:,1]
93 _plt.plot(xx, yy, xlabel=xlabel, ylabel='Depth (m)', xlim=(min_x-mgn_x, max_x+mgn_x), ylim=(-max_y-mgn_y, -min_y+mgn_y), color=surface_color, **kwargs)
95 if np.size(env['bottom_depth']) == 1:
96 xx = [min_x, max_x]
97 yy = [-env['bottom_depth'], -env['bottom_depth']]
98 else:
99 # linear and curvilinear options use the same bathymetry, just with different normals
100 s = env['bottom_depth']
101 xx = s[:,0]/divisor
102 yy = -s[:,1]
103 _plt.plot(xx, yy, color=bottom_color)
105 txd = env['source_depth']
106 _plt.plot([0]*np.size(txd), -txd, marker='*', style='solid', color=source_color)
108 if receiver_plot is None:
109 receiver_plot = np.size(env['receiver_depth'])*np.size(env['receiver_range']) < 2000
110 if receiver_plot:
111 rxr = env['receiver_range']
112 if np.size(rxr) == 1:
113 rxr = [rxr]
114 for r in np.array(rxr):
115 rxd = env['receiver_depth']
116 _plt.plot([r/divisor]*np.size(rxd), -rxd, marker='o', style='solid', color=receiver_color)
118 _plt.hold(oh if oh is not None else False)
120def plot_ssp(env: Environment, **kwargs: Any) -> None:
121 """Plots the sound speed profile.
123 Parameters
124 ----------
125 env : dict
126 Environment description
127 **kwargs
128 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
130 Notes
131 -----
132 If the sound speed profile is range-dependent, this function only plots the first profile.
134 Examples
135 --------
136 >>> import aubellhop as bh
137 >>> env = bh.Environment(soundspeed=[[ 0, 1540], [10, 1530], [20, 1532], [25, 1533], [30, 1535]])
138 >>> bh.plot_ssp(env)
139 """
141 if env is not None:
142 env.check()
144 oh = _plt.hold()
145 svp = env['soundspeed']
146 if isinstance(svp, pd.DataFrame):
147 svp = np.hstack((np.array([svp.index]).T, np.asarray(svp)))
148 if env['soundspeed_interp'] == BHStrings.spline:
149 ynew = np.linspace(np.min(svp[:,0]), np.max(svp[:,0]), 100)
150 tck = _interp.splrep(svp[:,0], svp[:,1], s=0)
151 xnew = _interp.splev(ynew, tck, der=0)
152 _plt.plot(xnew, -ynew, xlabel='Soundspeed (m/s)', ylabel='Depth (m)', hold=True, **kwargs)
153 _plt.scatter(svp[:,1], -svp[:,0], **kwargs)
154 else:
155 for rr in range(1,svp.shape[1]):
156 _plt.plot(svp[:,rr], -svp[:,0], xlabel='Soundspeed (m/s)', ylabel='Depth (m)', legend=f'Range {rr}', **kwargs)
157 _plt.hold(oh if oh is not None else False)
159def plot_arrivals(arrivals: Any, dB: bool = False, color: str = 'blue', **kwargs: Any) -> None:
160 """Plots the arrival times and amplitudes.
162 Parameters
163 ----------
164 arrivals : pandas.DataFrame
165 Arrivals times (s) and coefficients
166 dB : bool, default=False
167 True to plot in dB, False for linear scale
168 color : str, default='blue'
169 Line color (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)
170 **kwargs
171 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
173 Examples
174 --------
175 >>> import aubellhop as bh
176 >>> env = bh.Environment()
177 >>> arrivals = bh.compute_arrivals(env)
178 >>> bh.plot_arrivals(arrivals)
179 """
181 t0 = min(arrivals.time_of_arrival)
182 t1 = max(arrivals.time_of_arrival)
183 oh = _plt.hold()
184 if dB:
185 min_y = 20*np.log10(np.max(np.abs(arrivals.arrival_amplitude)))-60
186 ylabel = 'Amplitude (dB)'
187 else:
188 ylabel = 'Amplitude'
189 min_y = 0
190 _plt.plot([t0, t1], [min_y, min_y], xlabel='Arrival time (s)', ylabel=ylabel, color=color, **kwargs)
191 for _, row in arrivals.iterrows():
192 t = row.time_of_arrival.real
193 y = np.abs(row.arrival_amplitude)
194 if dB:
195 y = max(20*np.log10(_fi.epsilon+y), min_y)
196 _plt.plot([t, t], [min_y, y], color=color, **kwargs)
197 _plt.hold(oh if oh is not None else False)
199def plot_rays(rays: pd.DataFrame, env: Environment | None = None, invert_colors: bool = False, **kwargs: Any) -> None:
200 """Plots ray paths.
202 Parameters
203 ----------
204 rays : pandas.DataFrame
205 Ray paths
206 env : dict, optional
207 Environment definition
208 invert_colors : bool, default=False
209 False to use black for high intensity rays, True to use white
210 **kwargs
211 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported
213 Notes
214 -----
215 If environment definition is provided, it is overlayed over this plot using default
216 parameters for `bellhop.plot_env()`.
218 Examples
219 --------
220 >>> import aubellhop as bh
221 >>> env = bh.Environment()
222 >>> rays = bh.compute_eigenrays(env)
223 >>> bh.plot_rays(rays, width=1000)
224 """
225 rays = rays.sort_values('bottom_bounces', ascending=False)
226 dim = rays["ray"].iloc[0][0].shape[0]
228 # some edge cases to worry about here: rays.bottom_bounces could be all zeros?
229 max_amp = np.max(np.abs(rays.bottom_bounces)) if len(rays.bottom_bounces) > 0 else 0
230 max_amp = max_amp or 1
232 divisor = 1
233 xunits = '(m)'
234 r = []
235 for _, row in rays.iterrows():
236 r += list(row.ray[:,0])
237 if max(r)-min(r) > 10000:
238 divisor = 1000
239 xunits = '(km)'
241 oh = _plt.hold()
242 for _, row in rays.iterrows():
243 rr = float( row.bottom_bounces / (max_amp + 1) ) # avoid rr = 1 == 100% white
244 c = 1.0 - rr if invert_colors else rr
245 cmap = _pyplt.get_cmap("gray")
246 col_str = _mplc.to_hex(cmap(c))
247 ax_ind = dim - 1
248 _plt.plot(row.ray[:,0]/divisor, -row.ray[:,ax_ind], color=col_str, xlabel='Range '+xunits, ylabel='Depth (m)', **kwargs)
249 if env is not None:
250 plot_env(env,title=None)
251 _plt.hold(oh if oh is not None else False)
253def plot_transmission_loss(tloss: Any, env: Environment | None = None, **kwargs: Any) -> None:
254 """Plots transmission loss.
256 Parameters
257 ----------
258 tloss : pandas.DataFrame
259 Complex transmission loss
260 env : dict, optional
261 Environment definition
262 **kwargs
263 Other keyword arguments applicable for `bellhop.plot.image()` are also supported
265 Notes
266 -----
267 If environment definition is provided, it is overlayed over this plot using default
268 parameters for `bellhop.plot_env()`.
270 Examples
271 --------
272 >>> import aubellhop as bh
273 >>> import numpy as np
274 >>> env = bh.Environment(
275 receiver_depth=np.arange(0, 25),
276 receiver_range=np.arange(0, 1000),
277 beam_angle_min=-45,
278 beam_angle_max=45
279 )
280 >>> tloss = bh.compute_transmission_loss(env)
281 >>> bh.plot_transmission_loss(tloss, width=1000)
282 """
284 xr = (min(tloss.columns), max(tloss.columns))
285 yr = (-max(tloss.index), -min(tloss.index))
286 xlabel = 'Range (m)'
287 if xr[1]-xr[0] > 10000:
288 xr = (min(tloss.columns)/1000, max(tloss.columns)/1000)
289 xlabel = 'Range (km)'
290 oh = _plt.hold()
291 _plt.image(20*np.log10(_fi.epsilon+np.abs(np.flipud(np.array(tloss)))), x=xr, y=yr, xlabel=xlabel, ylabel='Depth (m)', xlim=xr, ylim=yr, **kwargs)
292 if env is not None:
293 plot_env(env, receiver_plot=False, title=None)
294 _plt.hold(oh if oh is not None else False)
297### Export module names for auto-importing in __init__.py
299__all__ = [
300 name for name in globals() if not name.startswith("_") # ignore private names
301]