""" This module contains methods for numerically solving the parabolic wave
equation of Thomson and Chapman.
"""
import numpy as np
from numpy.lib import scimath
from kadlu.utils import toarray, deg2rad
from tqdm import tqdm
from kadlu.plot_util import plot_transm_loss_horiz, plot_transm_loss_vert
[docs]def thomson_starter(k0, kz, dz, zs, theta1):
""" Compute Thomson starter field :math:`\psi (0, k_z)` as defined in
Jensen Sec. 6.4.2.3.
Args:
k0: float
Reference wavenumber in inverse meters
kz: numpy.array
Vertical wavenumber in inverse meters
dz: float
Vertical bin size in meters
zs: array-like
Source depth(s) in meters
theta1: float
Half-beamwidth in radians
Returns:
psi: numpy array
Starter field, :math:`\psi (0, k_z)`.
Has shape (len(zs), len(kz), 1)
Example:
>>> import numpy as np
>>> from kadlu.sound.parabolic_equation import Grid, thomson_starter
>>>
>>> # Create a grid with depth of 500 m, range of 1 km,
>>> # and grid spacings of 100 meters and 10 degrees
>>> grid = Grid(100., 1000., 100., 500., 10.*np.pi/180.)
>>>
>>> # compute Thomson starter field with an aperture of 86 degrees,
>>> # for a reference wavenumber of 0.04 m^-1 and a source depth of 9 meters,
>>> psi = thomson_starter(k0=0.04, kz=grid.kz, dz=grid.dz, zs=9, theta1=86*np.pi/180.)
>>> psi = np.round(psi, 4) # round to 4 decimals
>>> print(psi[0])
[[ 0. +0.j ]
[ 0.0101-0.0101j]
[ 0.0205-0.0205j]
[ 0.0319-0.0319j]
[ 0.0451-0.0451j]
[ 0. +0.j ]
[-0.0451+0.0451j]
[-0.0319+0.0319j]
[-0.0205+0.0205j]
[-0.0101+0.0101j]]
"""
nz = len(kz)
kz = kz[np.newaxis, :]
zs = toarray(zs)[:, np.newaxis]
# compute [scimath.sqrt(scimath.sqrt(k0**2 - kz**2))]^-1
idx = np.nonzero(np.abs(kz) >= np.abs(k0))
a = np.empty(kz.shape)
a[idx] = 0
idx = np.nonzero(np.abs(kz) < np.abs(k0))
a[idx] = 1. / scimath.sqrt(scimath.sqrt(k0**2 - kz[idx]**2))
psi = np.exp(-1j * np.pi / 4.) * scimath.sqrt(8 * np.pi) * np.sin(
kz * zs) * a
psi = psi / dz # normalize
psi[:, int(nz / 2)] = 0 # set to 0 at the bottom of the ocean
# taper the spectrum to obtain desired angle using Turkey window
kcut1 = k0 * np.sin(theta1)
kcut0 = k0 * np.sin(theta1 - 1.5 * deg2rad)
W = 0.5 * (1 + np.cos(np.pi / (kcut1 - kcut0) * (np.abs(kz) - kcut0)))
W[np.abs(kz) >= kcut1] = 0
W[np.abs(kz) <= kcut0] = 1
psi = psi * W
psi = psi[:, :, np.newaxis]
return psi
[docs]class Grid():
""" Regular grid in cylindrical coordinates r,q,z, where
* r: radial coordinate (distance) in meters
* z: vertical coordinate (depth) in meters
* q: azimuthal coordinate (angle) in radians
Note: The sea surface is at z = 0. We use positive z values below
the sea surface and negative values above.
Args:
dr: float
Radial grid spacing in meters
r_max: float
Radial range in meters
dz: float
Vertical grid spacing in meters
z_max: float
Vertical range in meters
dq: float
Angular grid spacing in radians
q_max: float
Angular range in radians
q: float, list or numpy.array
Angular coordinates in radians. If specified, dq and q_max are ignored.
Attrs:
dr: float
Radial grid spacing in meters
r: numpy.array
Radial coordinates in meters
nr: int
Number of radial grid points
dz: float
Vertical grid spacing in meters
z: numpy.array
Vertical coordinates in meters
nz: int
Number of vertical grid points
dq: float
Angular grid spacing in radians
q: numpy.array
Angular coordinates in radians
nq: int
Number of angular grid points
kz: numpy.array
Vertical wavenumber coordinates in inverse meters
below: numpy.array
Indices of z axis corresponding to values below (z>=0) the surface
q_qz, z_qz: numpy.array
Coordinates of q-z meshgrid
below_qz: numpy.array
Indices of q-z meshgrid corresponding to values below (z>=0) the surface
"""
def __init__(self, dr, r_max, dz, z_max, dq=None, q_max=2 * np.pi, q=None):
assert dq is not None or q is not None, "either dq or q must be specified"
self.r, self.dr = self._radial_coords(dr, r_max)
self.nr = len(self.r)
self.dq = None
if q is None: self.q, self.dq = self._azimuth_coords(dq, q_max)
elif isinstance(q, list): self.q = np.array(q)
elif isinstance(q, (int, float)): self.q = np.array([q])
self.nq = len(self.q)
self.z, self.dz, self.below, self._above, self._mirror = self._vertical_coords(
dz, z_max)
self.nz = len(self.z)
self.kz = self.z * 2 * np.pi / (len(self.z) * self.dz**2) #wave number
self.q_qz, self.z_qz = np.meshgrid(self.q, self.z) #q-z meshgrid
self.below_qz = np.nonzero(self.z_qz >= 0)
[docs] def mirror(self, x, z_axis=0):
""" Replace values above sea surface with mirrored values from below surface.
Args:
x: numpy.array
Array containing the data to be mirrored
z_axis: int
z axis of array
Returns:
x: numpy.array
Mirrored array
"""
x = np.swapaxes(x, 0, z_axis)
x[self._above] = x[self._mirror]
x = np.swapaxes(x, 0, z_axis)
return x
def _radial_coords(self, dr, r_max):
""" Compute radial coordinates"""
N = int(round(r_max / dr)) + 1
r = np.arange(N, dtype=float) * dr
return r, dr
def _azimuth_coords(self, dq, q_max):
""" Compute azimuthal coordinates"""
N = int(np.ceil(q_max / dq))
N += N % 2 # ensure even number of angular bins
q_pos = np.arange(start=0, stop=N / 2, step=1, dtype=float)
q_neg = np.arange(start=-N / 2, stop=0, step=1, dtype=float)
q = np.concatenate((q_pos, q_neg)) * dq
return q, dq
def _vertical_coords(self, dz, z_max):
""" Compute vertical coordinates"""
N = int(round(z_max / dz) * 2) # ensure even number of vertical bins
z_pos = np.arange(start=0, stop=N / 2 + 1, step=1, dtype=float)
z_neg = np.arange(start=-N / 2 + 1, stop=0, step=1, dtype=float)
z = np.concatenate((z_pos, z_neg)) * dz
idx_pos = np.nonzero(z >= 0)[0]
idx_neg = np.nonzero(z < 0)[0]
idx_mirror = idx_pos[-2:0:-1]
return z, dz, idx_pos, idx_neg, idx_mirror
[docs]def prop_defr(x, k0, kz, nq):
""" Compute the defractive propagation matrix, :math:`U_D`.
Args:
x: float
Step size in meters
k0: float
Reference wavenumber in inverse meters
kz: array-like
Vertical wavenumber in inverse meters
nq: int
Number of angular bins
Returns:
U: numpy.array
Defractive propagation matrix, has shape (1,nz,nq)
"""
U = np.exp(1j * (scimath.sqrt(k0**2 - kz**2) - k0) * x)
U = U[:, np.newaxis] * np.ones(shape=(1, nq))
U = U[np.newaxis, :, :]
return U
[docs]def prop_refr(x, k0, n):
""" Compute the refractive propagation matrix, :math:`U_R`.
Args:
x: float
Step size in meters
k0: float
Reference wavenumber in inverse meters
n: numpy.array
Index of refraction, has shape (nz,nq)
Returns:
U: numpy.array
Refractive propagation matrix, has shape (1,nz,nq)
"""
U = np.exp(1j * x * k0 * (n - 1))
U = U[np.newaxis, :, :]
return U
[docs]def artificial_absorp(z, z_max, D, alpha):
""" Compute complex term, to be added to the square of the index of
refraction in order to simulate the effect of an artificial
absorption layer at the bottom of the physical domain.
Args:
z: numpy.array
Depth values in meters
z_max: float
Maximum depth in meters
D: float
Characteristic width of the absorption layer in meters
alpha: float
Attenuation coefficient
Returns:
: numpy.array
Absorption term, to be added to the square of the physical
index of refraction.
"""
return 1j * alpha * np.exp(-(np.abs(z) - z_max)**2 / D**2)
[docs]def complex_sound_speed(c, alpha):
""" Computes complex sound speed representing the effect
of volume attenuation.
Args:
c: array-like
Sound speed in m/s
alpha: float
Attenuation coefficient in dB/lambda
Returns:
ci: array-like
Complex sound speed term
"""
if alpha == 0: return 0
beta = alpha / (40. * np.pi * np.log10(np.e) * c)
if isinstance(c, (int, float)):
ci = np.roots([beta, -1,
beta * c**2]) # roots of polynomial p[0]*x^n+...p[n]
ci = ci[np.imag(ci) == 0]
ci = ci[np.logical_and(ci >= 0, ci < c)]
ci = -1j * ci[0]
else:
ci = -1j * beta * c**2
return ci
[docs]def index_refr_sq(c0, c, alpha=0):
""" Compute index of refraction.
If alpha > 0, the index of refraction receives a complex term,
which represents the effect of volume attenuation.
Args:
c0: float
Reference sound speed in m/s
c: array-like
Sound speed in m/s
alpha: float
Attenuation coefficient in dB/lambda
Returns:
n2: array-like
Index of refraction squared
"""
c += complex_sound_speed(c, alpha)
n2 = (c0 / c)**2
return n2
[docs]def smooth_index_refr(z, zb, n2, n2b, L):
""" Smoothen the index of refraction from the water column to the
bottom layer.
Args:
z: numpy.array
Depth values in meters
zb: numpy.array
Seafloor depth in meters
n2: numpy.array
Index of refraction squared in water
n2b: numpy.array
Index of refraction squared in bottom
L: float
Characteristic smoothing length in meters
Returns:
: numpy.array
Index of refraction squared, with smooth transition from water to bottom
"""
x = (np.abs(z) - zb) / L
return n2 + 0.5 * (n2b - n2) * (1 + np.tanh(x))
[docs]def smooth_density(z, zb, dzb2, L, r, rb):
""" Smoothen the density transition from the water column to the
bottom layer and compute first and second derivatives.
Args:
z: numpy.array
Depth values in meters
zb: numpy.array
Seafloor depth in meters
dzb2: numpy.array
Gradient of the seafloor depth squared
L: float
Characteristic smoothing length in meters
r: numpy.array
Water density in g/cm^3
rb: numpy.array
Bottom density in g/cm^3
Returns:
rs: numpy.array
Density in g/cm^3, with smooth transition from water to bottom
dr2: numpy.array
Square of density gradient in (g/cm^3/m)^2
d2r: numpy.array
Divergence of density gradient (g/cm^3/m^2)
"""
x = (np.abs(z) - zb) / L
tanh = np.tanh(x)
sech2 = np.empty(x.shape)
idx = np.nonzero(np.abs(x) < 100)
sech2[idx] = 1. / np.cosh(x[idx])**2
idx = np.nonzero(np.abs(x) >= 100)
sech2[idx] = 0
rs = r + 0.5 * (rb - r) * (
1 + tanh) # density (rho), smooth transition from water to bottom
dr2 = 1. / (4 * L**2) * (rb - r)**2 * sech2**2 * (1 + dzb2
) # gradient squared
d2r = -1. / L**2 * (rb - r) * sech2 * tanh * (1 + dzb2
) # divergence of gradient
return rs, dr2, d2r
[docs]def eff_index_refr_sq(z,
zb,
dzb2,
k0,
L,
n2,
n2b,
r,
rb,
return_density=False):
""" Compute the effective index of refraction squared.
Args:
z: numpy.array
Depth values in meters
zb: numpy.array
Seafloor depth in meters
dzb2: numpy.array
Gradient of the seafloor depth squared
k0: float
Reference wavenumber in inverse meters
L: float
Characteristic smoothing length in meters
n2: numpy.array
Water refractive index squared
n2b: numpy.array
Bottom refractive index squared
r: numpy.array
Water density in g/cm^3
rb: numpy.array
Bottom density in g/cm^3
return_density: bool
Return the smoothened density array. Default is False.
Returns:
n2e: numpy.array
Effective index of refraction squared
rs: numpy.array
Smoothened density. Only returned if return_density is True.
"""
n2s = smooth_index_refr(z, zb, n2, n2b, np.finfo(float).eps)
rs, dr2, d2r = smooth_density(z, zb, dzb2, L, r, rb)
n2e = n2s + 1 / (2 * k0**2) * (1 / rs * d2r - 3 / (2 * rs**2) * dr2)
if return_density: return n2e, rs
else: return n2e
[docs]class TransmissionLoss():
""" Compute the transmission loss by solving the parabolic wave
equation.
The seafloor acoustic properties are specified in a dictionary
with the following keys,
* sound_speed (bottom sound speed in m/s)
* density (bottom density in g/cm3)
* attenuation (bottom attenuation in dB/lambda)
Keyword arguments can be used to specify the computational
grid :class:`kadlu.sound.parabolic_equation.Grid`.
Args:
freq: float
Frequency in Hz
bathy_func: function
Bathymetry interpolation function in variables x,y
bathy_deriv_func: function
Bathymetry derivative interpolation function in variables x,y
sound_speed_func: function
Sound speed interpolation function in variables x,y,z
bottom: dict
Seafloor acoutic properties. Must contain the keys 'sound_speed',
'density', and 'attenuation'.
propagation_range: float
Propagation range in km. Default is 50 km.
angular_bin: float
Angular bin size in degrees. Default is 10 degrees.
c0: float
Reference sound speed in m/s
source_depth: array-like
Source depths in meters.
Attributes:
Example:
"""
def __init__(self,
freq,
bathy_func,
bathy_deriv_func,
sound_speed_func,
bottom,
propagation_range=50,
angular_bin=10,
c0=1500,
**kwargs):
self.c0 = c0
self.k0 = 2 * np.pi / c0 * freq
self.water_density = 1.0
r_max = 1e3 * propagation_range
dq = angular_bin * deg2rad
H = self._max_depth(bathy_func,
r_max) + 3 * c0 / freq #depth of physical domain
z_max = 4. / 3 * H
# default grid
grid_kwargs = {
'dz': c0 / freq / 2.,
'dr': c0 / freq,
'dq': dq,
'q': None,
'r_max': r_max,
'z_max': z_max
}
# replace defaults with input args, if any
for key in grid_kwargs:
if key in kwargs.keys(): grid_kwargs[key] = kwargs[key]
self.grid = Grid(**grid_kwargs)
self.bottom = bottom
self.costheta = np.cos(self.grid.q)
self.sintheta = np.sin(self.grid.q)
# bottom refractive index
self.n2b = index_refr_sq(c0=c0,
c=bottom['sound_speed'],
alpha=bottom['attenuation'])
# artificial bottom absorption layer
alpha = 1. / np.log10(np.e) / np.pi
z_max = np.max(self.grid.z)
D = (z_max - H) / 3
self.absorp = artificial_absorp(z=self.grid.z_qz,
z_max=z_max,
D=D,
alpha=alpha)
# interpolation functions
self._bathy = bathy_func
self._bathy_deriv = bathy_deriv_func
self._sound_speed = sound_speed_func
self._do_vertical = False # compute transmission loss on vertical plane
# source depth
self._source_depth = kwargs[
'source_depth'] if 'source_depth' in kwargs.keys() else None
[docs] def calc(self,
source_depth=None,
rec_depth=[.1],
vertical=False,
aperture=86,
progress_bar=True,
nz_max=250,
nr_max=250,
size_limit=True,
return_field=False):
""" Calculate the transmission loss in the horizontal plane at
the specified depth(s).
Args:
source_depth: array-like
Source depths in meters
rec_depth: array-like
Receiver depths in meters
vertical: array-like
Compute the transmission loss in the vertical plane.
Note: This will slow down the computation.
aperture: float
Half-beamwidth of the starter field in degrees.
progress_bar: bool
Show progress bar. Default is True
nz_max: int
Maximum number of vertical bins for output arrays
nr_max: int
Maximum number of radial bins for output arrays
size_limit: bool
Limit of the size of the output arrays to the dimensions
specified in nz_max and nr_max. If False, nz_max and nr_max
have no effect
return_field: bool
Instead of returning the transmission loss, TL, return
the complex pressure field, f, where TL = 20 * np.log10(np.abs(f)).
Default is False.
Returns:
tl_h: numpy.array
Transmission loss in dB in the horizontal plane; has shape
(len(source_depth), len(rec_depth), len(q), len(r)).
ax_h: dict
Axes of the horizontal transmission loss array, (source_depth, rec_depth, q, r).
q and r are the azimuthal and radial coordinate axes.
tl_v: numpy.array
Transmission loss in dB in the vertical plane; has shape
(len(source_depth), len(z), len(r), len(q)).
Only returned if vertical is True.
ax_v: dict
Axes of the vertical transmission loss array, (source_depth, z, r, q).
z and r are the vertical and radial coordinate axes.
q is the azimuthal axes.
Only returned if vertical is True.
"""
# source depth
if source_depth is None: source_depth = self._source_depth
assert source_depth is not None, 'source depth must be specified'
self._do_vertical = vertical
source_depth = toarray(source_depth)
rec_depth = toarray(rec_depth)
self._init_output(num_sources=len(source_depth),
rec_depth=rec_depth,
nz_max=nz_max,
nr_max=nr_max,
size_limit=size_limit)
psi = self._solve_pe(source_depth=source_depth,
rec_depth=rec_depth,
aperture=aperture,
progress_bar=progress_bar)
# transmission loss, horizontal plane
field_h = np.fft.fftshift(self._field_horiz[:, :, :, 1:],
axes=2) #re-order q axis
tl_h = -20 * np.log10(np.abs(field_h))
q = self.grid.q
if len(q) > 1: q = np.squeeze(q)
q = np.fft.fftshift(q)
ax_h = {
'source_depth': source_depth,
'receiver_depth': rec_depth,
'azimuthal_axis': q,
'radial_axis': self._field_r_ax[1:]
}
res_h = field_h if return_field else tl_h
self.tl_h, self.ax_h = tl_h, ax_h
if not self._do_vertical:
return res_h, ax_h
else: # transmission loss, vertical plane
field_v = np.fft.fftshift(self._field_vert[:, :, :, :],
axes=3) #re-order q axis
tl_v = -20 * np.ma.log10(
np.abs(field_v)) # OBS: this computation is rather slow
ax_v = {
'source_depth': source_depth,
'vertical_axis': self._field_z_ax,
'radial_axis': self._field_r_ax,
'azimuthal_axis': q
}
self.tl_v, self.ax_v = tl_v, ax_v
res_v = field_v if return_field else tl_v
return res_h, ax_h, res_v, ax_v
[docs] def plot_horiz(self, source_depth_idx=0, rec_depth_idx=0):
""" Plot the transmission loss on a horizontal plane in polar coordinates.
Args:
source_depth_idx: int
Source depth index
rec_depth_idx: int
Receive depth index
Returns:
fig: matplotlib.figure.Figure
A figure object.
"""
tl = self.tl_h[source_depth_idx, rec_depth_idx]
r = self.ax_h['radial_axis']
q = self.ax_h['azimuthal_axis']
fig = plot_transm_loss_horiz(tl, r, q)
return fig
[docs] def plot_vert(self,
angle=0,
source_depth_idx=0,
show_bathy=True,
max_depth=None,
show_ssp=False):
""" Plot the transmission loss on a vertical plane in carthesian coordinates.
Args:
angle: float
Angle in degrees
source_depth_idx: int
Source depth index
show_bathy: bool
Superimpose bathymetry on plot
show_ssp: bool
Superimpose sound speed profile on plot
Returns:
fig: matplotlib.figure.Figure
A figure object.
"""
z = self.ax_v['vertical_axis']
r = self.ax_v['radial_axis']
q = self.ax_v['azimuthal_axis']
# truncate z axis
if max_depth is not None: z = z[z < max_depth]
# find nearest angular bin
if angle <= 180: q0 = angle * np.pi / 180
else: q0 = (angle - 360) * np.pi / 180
angle_idx = np.abs(q - q0).argmin()
tl = self.tl_v[source_depth_idx, 1:len(z), :, angle_idx]
z = z[:-1] #drop last
# bathymetry
if show_bathy:
def bathy_r(r):
angle_rad = angle * deg2rad
x = np.cos(angle_rad) * r
y = np.sin(angle_rad) * r
return self._bathy(x, y)
else:
bathy_r = None
# sound speed profile
if show_ssp:
x = np.cos(angle * deg2rad) * np.linspace(0, np.max(r), 1000)
y = np.sin(angle * deg2rad) * np.linspace(0, np.max(r), 1000)
idx = np.argmax(self._bathy(x, y))
x0, y0 = x[idx], y[idx]
def ssp_z(z):
return self._sound_speed(x0, y0, z, grid=True)
else:
ssp_z = None
fig = plot_transm_loss_vert(tl, z, r, bathy_r, ssp_z)
return fig
def _solve_pe(self,
source_depth,
rec_depth,
aperture,
progress_bar,
return_field=False):
""" Solve the parabolic wave equation of Thomson and Chapman by means
of a split-step Fourier algorithm.
The displacement field is computed on one or several horizontal planes
at the specified receiver depth(s).
Optionally, the field may also be computed on a vertical plane intersecting
the source position.
Args:
source_depth: array-like
Source depths in meters
rec_depth: array-like
Receiver depths in meters
aperture: float
Half-beamwidth of the starter field in degrees.
progress_bar: bool
Show progress bar. Default is True
return_field: bool
Return the displacement field :math:`\psi (r, k_z)` after propagation to
:math:`r=r_{max}`. Primarily for diagnostics/debugging purposes.
Default is False.
Returns:
psi: numpy.array
Displacement field :math:`\psi (r, k_z)` at :math:`r=r_{max}`.
Only returned if return_field is True.
"""
dr = self.grid.dr
dz = self.grid.dz
k0 = self.k0
kz = self.grid.kz
nq = self.grid.nq
nr = self.grid.nr
rec_depth = toarray(rec_depth)
UD = prop_defr(x=dr / 2, k0=k0, kz=kz,
nq=nq) #diffractive propagation matrix
psi = thomson_starter(k0=k0,
kz=kz,
dz=dz,
zs=source_depth,
theta1=aperture * deg2rad) # starter field
psi = psi * np.ones((1, 1, nq))
self._save_output(step_no=0,
r=0,
psi=psi,
sqrt_rho=0,
rec_depth=rec_depth) #save output
r = 0
for i in tqdm(range(nr - 1), disable=not progress_bar): # PE marching
psi = UD * psi # diffractive propagation, half-step
n, sqrt_rho = self._update_env(
r + dr / 2) # update acoustic environment
UR = prop_refr(x=dr, k0=k0, n=n) # refractive propagation
psi = np.fft.fft(UR * np.fft.ifft(psi, axis=1),
axis=1) # refractive propagation, full step
psi = UD * psi # diffractive propagation, half-step
r += dr # increment distance
self._save_output(step_no=i + 1,
r=r,
psi=psi,
sqrt_rho=sqrt_rho,
rec_depth=rec_depth) # collect output
if return_field: return psi
def _max_depth(self, bathy, r_max, return_xy=False):
""" Find the maximum depth in the computational domain.
Args:
bathy:
Bathymetry interpolation function in x,y variables
r_max: float
Range in meters
return_xy: bool
Return also the coordinates of the location of maximum depth
Returns:
b0: float
Maximum depth in meters
x0,y0: float,float
x,y coordinates of location of max depth. Only return if return_xy is True.
"""
x = np.linspace(-r_max, r_max, 2000)
b = bathy(x, x, grid=True)
b0 = np.max(b)
if return_xy:
idx = np.unravel_index(np.argmax(b, axis=None), b.shape)
x0 = x[idx[0]]
y0 = x[idx[1]]
return b0, x0, y0
else:
return b0
def _bathy_grad_sq(self, r):
""" Compute seafloor gradient squared at a specified distance
from the source in the radial direction.
Args:
r: float
Radial coordinate in meters
Returns:
: numpy.array
Seafloor gradient squared in all angular bin
"""
x = self.costheta * r
y = self.sintheta * r
dzdx = self._bathy_deriv(x=x, y=y, axis='x')
dzdy = self._bathy_deriv(x=x, y=y, axis='y')
return (self.costheta * dzdx)**2 + (self.sintheta * dzdy)**2
def _update_bathy(self, r):
""" Update the bathymetry to reflect conditions at a
specified distance from the source.
Args:
r: float
Radial coordinate in meters
Returns:
zb: numpy.array
Seafloor depth in meters
dzb2: numpy.array
Seafloor gradient squared
"""
x = self.costheta * r
y = self.sintheta * r
zb = self._bathy(x=x, y=y) #bathymetry
# temp fix to ensure that bathy returns array even when x and y have length 1
if np.ndim(zb) == 0: zb = np.array([zb])
dzb2 = self._bathy_grad_sq(r=r) #gradient squared
nz = self.grid.nz
zb = np.ones((nz, 1)) * zb[np.newaxis, :]
dzb2 = np.ones((nz, 1)) * dzb2[np.newaxis, :]
return zb, dzb2
def _update_sound_speed(self, r):
""" Update the sound speed to reflect conditions at a
specified distance from the source.
Args:
r: float
Radial coordinate in meters
Returns:
c: numpy.array
Sound speed in m/s
"""
x = self.costheta * r
y = self.sintheta * r
z = self.grid.z[self.grid.below]
x, _ = np.meshgrid(x, z)
y, z = np.meshgrid(y, z)
x = x.flatten()
y = y.flatten()
z = z.flatten()
c = self._sound_speed(x=x, y=y, z=z)
return c
def _eff_index_refr(self, c, zb, dzb2):
""" Compute the effective index of refraction,
including the artificial bottom absorption term.
Args:
c: numpy.array
Sound speed in m/s
zb: numpy.array
Seafloor depth in meters
dzb2: numpy.array
Seafloor gradient squared
Returns:
: numpy.array
Effective index of refraction
: numpy.array
Square root of density
"""
# refractive index squared
n2 = np.zeros((self.grid.nz, self.grid.nq))
n2[self.grid.below_qz] = index_refr_sq(self.c0, c)
n2 = self.grid.mirror(n2)
# effective refr. index squared
n2, rho = eff_index_refr_sq(z=self.grid.z_qz,
zb=zb,
dzb2=dzb2,
k0=self.k0,
L=np.pi / self.k0,
n2=n2,
n2b=self.n2b,
r=self.water_density,
rb=self.bottom['density'],
return_density=True)
# add absorption
n2 += self.absorp
return scimath.sqrt(n2), scimath.sqrt(rho)
def _update_env(self, r):
""" Update the acoustic environment to reflect
the conditions at a specified distance from the source.
Args:
r: float
Radial coordinate in meters
Returns:
n: numpy.array
Refractive index
sqrt_rho: numpy.array
Square root of density
"""
zb, dzb2 = self._update_bathy(r)
c = self._update_sound_speed(r)
n, sqrt_rho = self._eff_index_refr(c, zb, dzb2)
return n, sqrt_rho
def _save_output(self, step_no, r, psi, sqrt_rho, rec_depth):
""" Post-processe and save output data at specified distance from the source.
Args:
step_no: int
Step number
r: float
Radial distance from the source in meters
psi: numpy.array
Displacement field at the specified radial distance; has shape (ns,nz,nq).
sqrt_rho: numpy.array
Square root of density
rec_depth: array-like
Receiver depths in meters
"""
if step_no % self._r_step != 0: return
bin_no = int(step_no / self._r_step)
if step_no > 0:
dz = self.grid.dz
idx = np.squeeze(np.round(rec_depth / dz).astype(int))
if np.ndim(idx) == 0: idx = np.array([idx])
F = np.matmul(self._ifft_kernel, psi)
G = sqrt_rho[idx][np.newaxis, :, :]
# inverse fourier transform and multiply by sqrt(density)
self._field_horiz[:, :, :, bin_no] = F * G * np.exp(
1j * self.k0 * r) / np.sqrt(r)
if self._do_vertical:
# inverse fourier transform and multiply by sqrt(density)
psi_z = np.fft.ifft(psi, axis=1)
if r > 0: psi_z *= np.exp(1j * self.k0 * r) / np.sqrt(r) * sqrt_rho
# only save values below sea surface
n = int(self.grid.nz / 2)
self._field_vert[:, :, bin_no, :] = psi_z[:, :n:self._z_step, :]
def _init_output(self,
num_sources,
rec_depth,
nz_max=250,
nr_max=250,
size_limit=False):
""" Initialize output containers and compute inverse fourier transform kernel.
Args:
num_sources: int
Number of source depths
rec_depth: array-like
Receiver depths in meters
nz_max: int
Maximum number of vertical bins for output arrays
nr_max: int
Maximum number of radial bins for output arrays
size_limit: bool
Limit of the size of the output arrays to the dimensions
specified in nz_max and nr_max. If False, nz_max and nr_max
have no effect
"""
nr = self.grid.nr
nq = self.grid.nq
nz = self.grid.nz
ns = num_sources
nd = len(rec_depth)
kz = self.grid.kz
rec_depth = rec_depth[:, np.newaxis]
# ifft kernel
self._ifft_kernel = np.exp(
1j * np.matmul(rec_depth, kz[np.newaxis, :])) / len(kz)
self._ifft_kernel = self._ifft_kernel[np.newaxis, :, :]
# output arrays
if size_limit:
self._r_step = int(nr / nr_max) + 1
self._z_step = int(nz / 2 / nz_max) + 1
else:
self._r_step = 1
self._z_step = 1
self._field_r_ax = self.grid.r[::self._r_step] #output r axis
self._field_z_ax = self.grid.z[:int(nz /
2):self._z_step] #output z axis
nr = len(self._field_r_ax)
nz = len(self._field_z_ax)
self._field_horiz = np.empty(shape=(ns, nd, nq, nr), dtype=complex)
self._field_vert = np.empty(shape=(ns, nz, nr, nq), dtype=complex)