Coverage for python/bellhop/pyplot.py: 81%

1##############################################################################

6# This file was originally part of arlpy, released under Simplified BSD License.

7# It has been relicensed in this repository to be compatible with the Bellhop licence (GPL).

9##############################################################################

11"""Plotting functions for the underwater acoustic propagation modeling toolbox.

12"""

14from typing import Any, Dict, Optional

15from sys import float_info as _fi

17import numpy as _np

18import scipy.interpolate as _interp

19import pandas as _pd

21import matplotlib.pyplot as _pyplt

22import matplotlib.colors as _mplc

24from bellhop.constants import _Strings

26def pyplot_env(env: Dict[str, Any], surface_color: str = 'dodgerblue', bottom_color: str = 'peru', source_color: str = 'orangered', receiver_color: str = 'midnightblue',

27 receiver_plot: Optional[bool] = None, **kwargs: Any) -> None:

28 """Plots a visual representation of the environment with matplotlib.

30 Parameters

31 ----------

32 env : dict

33 Environment description

34 surface_color : str, default='dodgerblue'

35 Color of the surface (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)

36 bottom_color : str, default='peru'

37 Color of the bottom (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)

38 source_color : str, default='orangered'

39 Color of transmitters (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)

40 receiver_color : str, default='midnightblue'

41 Color of receivers (see `Bokeh colors <https://bokeh.pydata.org/en/latest/docs/reference/colors.html>`_)

42 receiver_plot : bool, optional

43 True to plot all receivers, False to not plot any receivers, None to automatically decide

44 **kwargs

45 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported

47 Notes

48 -----

49 The surface, bottom, transmitters (marker: '*') and receivers (marker: 'o')

50 are plotted in the environment. If `receiver_plot` is set to None and there are

51 more than 2000 receivers, they are not plotted.

53 Examples

54 --------

55 >>> import bellhop as bh

56 >>> env = bh.create_env(depth=[[0, 40], [100, 30], [500, 35], [700, 20], [1000,45]])

57 >>> bh.plot_env(env)

58 """

60 if _np.array(env['receiver_range']).size > 1:

61 min_x = _np.min(env['receiver_range'])

62 else:

63 min_x = 0

64 max_x = _np.max(env['receiver_range'])

65 if max_x - min_x > 10000:

66 divisor = 1000

67 min_x /= divisor

68 max_x /= divisor

69 xlabel = 'Range (km)'

70 else:

71 divisor = 1

72 xlabel = 'Range (m)'

73 if env['surface'] is None:

74 min_y = 0

75 else:

76 min_y = _np.min(env['surface'][:, 1])

77 if _np.size(env['depth']) > 1:

78 max_y = _np.max(env['depth'][:, 1])

79 else:

80 max_y = env['depth']

81 mgn_x = 0.01 * (max_x - min_x)

82 mgn_y = 0.1 * (max_y - min_y)

83 if env['surface'] is None:

84 _pyplt.plot([min_x, max_x], [0, 0], color=surface_color, **kwargs)

85 _pyplt.xlabel(xlabel)

86 _pyplt.ylabel('Depth (m)')

87 print(min_x, mgn_x, max_x, mgn_x)

88 _pyplt.xlim([min_x - mgn_x, max_x + mgn_x])

89 _pyplt.ylim([-max_y - mgn_y, -min_y + mgn_y])

90 else:

91 # linear and curvilinear options use the same altimetry, just with different normals

92 s = env['surface']

93 _pyplt.plot(s[:, 0] / divisor, -s[:, 1], color=surface_color, **kwargs)

94 _pyplt.xlabel(xlabel)

95 _pyplt.ylabel('Depth (m)')

96 _pyplt.xlim([min_x - mgn_x, max_x + mgn_x])

97 _pyplt.ylim([-max_y - mgn_y, -min_y + mgn_y])

98 if _np.size(env['depth']) == 1:

99 _pyplt.plot([min_x, max_x], [-env['depth'], -env['depth']], color=bottom_color, **kwargs)

100 else:

101 # linear and curvilinear options use the same bathymetry, just with different normals

102 s = env['depth']

103 _pyplt.plot(s[:, 0] / divisor, -s[:, 1], color=bottom_color, **kwargs)

104 txd = env['source_depth']

105 # print(txd, [0]*_np.size(txd))

106 _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)

116

117def pyplot_ssp(env: Dict[str, Any], **kwargs: Any) -> None:

118 """Plots the sound speed profile with matplotlib.

119

120 Parameters

121 ----------

122 env : dict

123 Environment description

124 **kwargs

125 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported

126

127 Notes

128 -----

129 If the sound speed profile is range-dependent, this function only plots the first profile.

130

131 Examples

132 --------

133 >>> import bellhop as bh

134 >>> env = bh.create_env(soundspeed=[[ 0, 1540], [10, 1530], [20, 1532], [25, 1533], [30, 1535]])

135 >>> bh.plot_ssp(env)

136 """

137

138 svp = env['soundspeed']

139 if isinstance(svp, _pd.DataFrame):

140 svp = _np.hstack((_np.array([svp.index]).T, _np.asarray(svp)))

141 if _np.size(svp) == 1:

142 if _np.size(env['depth']) > 1:

143 max_y = _np.max(env['depth'][:, 1])

144 else:

145 max_y = env['depth']

146 _pyplt.plot([svp, svp], [0, -max_y], **kwargs)

147 _pyplt.xlabel('Soundspeed (m/s)')

148 _pyplt.ylabel('Depth (m)')

149 elif env['soundspeed_interp'] == _Strings.spline:

150 ynew = _np.linspace(_np.min(svp[:, 0]), _np.max(svp[:, 0]), 100)

151 tck = _interp.splrep(svp[:, 0], svp[:, 1], s=0)

152 xnew = _interp.splev(ynew, tck, der=0)

153 _pyplt.plot(xnew, -ynew, **kwargs)

154 _pyplt.xlabel('Soundspeed (m/s)')

155 _pyplt.ylabel('Depth (m)')

156 _pyplt.plot(svp[:, 1], -svp[:, 0], marker='.', **kwargs)

157 else:

158 _pyplt.plot(svp[:, 1], -svp[:, 0], **kwargs)

159 _pyplt.xlabel('Soundspeed (m/s)')

160 _pyplt.ylabel('Depth (m)')

161

162def pyplot_arrivals(arrivals: Any, dB: bool = False, color: str = 'blue', **kwargs: Any) -> None:

163 """Plots the arrival times and amplitudes with matplotlib.

164

165 Parameters

166 ----------

167 arrivals : pandas.DataFrame

168 Arrivals times (s) and coefficients

169 dB : bool, default=False

170 True to plot in dB, False for linear scale

171 color : str, default='blue'

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

173 **kwargs

174 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported

175

176 Examples

177 --------

178 >>> import bellhop as bh

179 >>> env = bh.create_env()

180 >>> arrivals = bh.compute_arrivals(env)

181 >>> bh.plot_arrivals(arrivals)

182 """

183 t0 = min(arrivals.time_of_arrival)

184 t1 = max(arrivals.time_of_arrival)

185 if dB:

186 min_y = 20 * _np.log10(_np.max(_np.abs(arrivals.arrival_amplitude))) - 60

187 ylabel = 'Amplitude (dB)'

188 else:

189 ylabel = 'Amplitude'

190 _pyplt.plot([t0, t1], [0, 0], color=color, **kwargs)

191 _pyplt.xlabel('Arrival time (s)')

192 _pyplt.ylabel(ylabel)

193 min_y = 0

194 for _, row in arrivals.iterrows():

195 t = row.time_of_arrival.real

196 y = _np.abs(row.arrival_amplitude)

197 if dB:

198 y = max(20 * _np.log10(_fi.epsilon + y), min_y)

199 _pyplt.plot([t, t], [min_y, y], color=color, **kwargs)

200 _pyplt.xlabel('Arrival time (s)')

201 _pyplt.ylabel(ylabel)

202

203def pyplot_rays(rays: Any, env: Optional[Dict[str, Any]] = None, invert_colors: bool = False, **kwargs: Any) -> None:

204 """Plots ray paths with matplotlib

205

206 Parameters

207 ----------

208 rays : pandas.DataFrame

209 Ray paths

210 env : dict, optional

211 Environment definition

212 invert_colors : bool, default=False

213 False to use black for high intensity rays, True to use white

214 **kwargs

215 Other keyword arguments applicable for `bellhop.plot.plot()` are also supported

216

217 Notes

218 -----

219 If environment definition is provided, it is overlayed over this plot using default

220 parameters for `bellhop.plot_env()`.

221

222 Examples

223 --------

224 >>> import bellhop as bh

225 >>> env = bh.create_env()

226 >>> rays = bh.compute_eigenrays(env)

227 >>> bh.plot_rays(rays, width=1000)

228 """

229 rays = rays.sort_values('bottom_bounces', ascending=False)

230 max_amp = _np.max(_np.abs(rays.bottom_bounces)) if len(rays.bottom_bounces) > 0 else 0

231 if max_amp <= 0:

232 max_amp = 1

233 divisor = 1

234 xlabel = 'Range (m)'

235 r = []

236 for _, row in rays.iterrows():

237 r += list(row.ray[:, 0])

238 if max(r) - min(r) > 10000:

239 divisor = 1000

240 xlabel = 'Range (km)'

241 for _, row in rays.iterrows():

242 c = float(_np.abs(row.bottom_bounces) / max_amp)

243 if invert_colors:

244 c = 1.0 - c

245 cmap = _pyplt.get_cmap("gray")

246 col_str = _mplc.to_hex(cmap(c))

247 if "color" in kwargs.keys():

248 _pyplt.plot(row.ray[:, 0] / divisor, -row.ray[:, 1], **kwargs)

249 else:

250 _pyplt.plot(row.ray[:, 0] / divisor, -row.ray[:, 1], color=col_str, **kwargs)

251 _pyplt.xlabel(xlabel)

252 _pyplt.ylabel('Depth (m)')

253 if env is not None:

254 pyplot_env(env)

255

256def pyplot_transmission_loss(tloss: Any, env: Optional[Dict[str, Any]] = None, **kwargs: Any) -> None:

257 """Plots transmission loss with matplotlib.

258

259 Parameters

260 ----------

261 tloss : pandas.DataFrame

262 Complex transmission loss

263 env : dict, optional

264 Environment definition

265 **kwargs

266 Other keyword arguments applicable for `bellhop.plot.image()` are also supported

267

268 Notes

269 -----

270 If environment definition is provided, it is overlayed over this plot using default

271 parameters for `bellhop.plot_env()`.

272

273 Examples

274 --------

275 >>> import bellhop as bh

276 >>> import numpy as np

277 >>> env = bh.create_env(

278 receiver_depth=np.arange(0, 25),

279 receiver_range=np.arange(0, 1000),

280 beam_angle_min=-45,

281 beam_angle_max=45

282 )

283 >>> tloss = bh.compute_transmission_loss(env)

284 >>> bh.plot_transmission_loss(tloss, width=1000)

285 """

286 xr = (min(tloss.columns), max(tloss.columns))

287 yr = (-max(tloss.index), -min(tloss.index))

288 xlabel = 'Range (m)'

289 if xr[1] - xr[0] > 10000:

290 xr = (min(tloss.columns) / 1000, max(tloss.columns) / 1000)

291 xlabel = 'Range (km)'

292 trans_loss = 20 * _np.log10(_fi.epsilon + _np.abs(_np.flipud(_np.array(tloss))))

293 x_mesh, ymesh = _np.meshgrid(_np.linspace(xr[0], xr[1], trans_loss.shape[1]),

294 _np.linspace(yr[0], yr[1], trans_loss.shape[0]))

295 trans_loss = trans_loss.reshape(-1)

296 # print(trans_loss.shape)

297 if "vmin" in kwargs.keys():

298 trans_loss[trans_loss < kwargs["vmin"]] = kwargs["vmin"]

299 if "vmax" in kwargs.keys():

300 trans_loss[trans_loss > kwargs["vmax"]] = kwargs["vmax"]

301 trans_loss = trans_loss.reshape((x_mesh.shape[0], -1))

302 _pyplt.contourf(x_mesh, ymesh, trans_loss, cmap="jet", **kwargs)

303 _pyplt.xlabel(xlabel)

304 _pyplt.ylabel('Depth (m)')

305 _pyplt.colorbar(label="Transmission Loss(dB)")

306 if env is not None:

307 pyplot_env(env, receiver_plot=False)

308

309

310### Export module names for auto-importing in __init__.py

311

312__all__ = [

313 name for name in globals() if not name.startswith("_") # ignore private names

314]