"""
Convolution in n-dimensions.
:copyright: 2022 Agile Geoscience
:license: Apache 2.0
"""
import numpy as np
from bruges.util import apply_along_axis
[docs]def convolve(arr, v, axis=-1, verbose=False):
"""
Convolve n-dimensional arr with a 1D wavelet or 2D wavelet bank.
Args:
arr (ndarray): The trace, or 2D section, or volume.
v (ndarray): The wavelet, must be 1D function or a 2D wavelet 'bank'.
If a wavelet bank, time should be on the last axis.
axis (int): The time axis of the arr data. In other words, the axis
corresponding to a single 'trace'. If you index into this axis,
you will get a single 'trace'.
verbose (bool): If True, print out the shapes of the inputs and output.
Returns:
ndarray: Discrete, linear convolution of `arr` and `v`.
"""
if v.shape[-1] > arr.shape[axis]:
raise TypeError("v (e.g. wavelet) must be shorter in time than arr (e.g. reflectivity).")
arr_ = np.moveaxis(np.asanyarray(arr), axis, -1)
# Compute the target shape of the final synthetic.
outshape = v.shape[:-1] + arr_.shape
# Force v and arr to both be 2D.
bank = np.atleast_2d(v)
arr_2d = arr_.reshape((-1, arr_.shape[-1]))
# Compute synthetic.
syn = np.array([apply_along_axis(np.convolve, arr_2d, w, mode='same') for w in bank])
pos = axis + v.ndim - 1
out = np.moveaxis(syn.reshape(outshape), -1, pos)
# Show the shapes of the data we're handling.
if verbose:
print(arr.shape, ' * ', v.shape, ' -> ', out.shape)
return out