def _correlate(in1, in2, mode='full', method='auto', convolution=False):
quick_out = _st_core._check_conv_inputs(in1, in2, mode, convolution)
if quick_out is not None:
return quick_out
if method not in ('auto', 'direct', 'fft'):
raise ValueError('acceptable methods are "auto", "direct", or "fft"')
if method == 'auto':
method = choose_conv_method(in1, in2, mode=mode)
if method == 'direct':
return _st_core._direct_correlate(in1, in2, mode, in1.dtype,
convolution)
# if method == 'fft':
inputs_swapped = _st_core._inputs_swap_needed(mode, in1.shape, in2.shape)
if inputs_swapped:
in1, in2 = in2, in1
if not convolution:
in2 = _st_core._reverse_and_conj(in2)
out = fftconvolve(in1, in2, mode)
result_type = cupy.result_type(in1, in2)
if result_type.kind in 'ui':
out = out.round()
out = out.astype(result_type, copy=False)
if not convolution and inputs_swapped:
out = cupy.ascontiguousarray(_st_core._reverse_and_conj(out))
return out
def fftconvolve(in1, in2, mode='full', axes=None):
"""Convolve two N-dimensional arrays using FFT.
Convolve ``in1`` and ``in2`` using the fast Fourier transform method, with
the output size determined by the ``mode`` argument.
This is generally much faster than the ``'direct'`` method of ``convolve``
for large arrays, but can be slower when only a few output values are
needed, and can only output float arrays (int or object array inputs will
be cast to float).
Args:
in1 (cupy.ndarray): First input.
in2 (cupy.ndarray): Second input. Should have the same number of
dimensions as ``in1``.
mode (str): Indicates the size of the output:
- ``'full'``: output is the full discrete linear \
cross-correlation (default)
- ``'valid'``: output consists only of those elements that do \
not rely on the zero-padding. Either ``in1`` or \
``in2`` must be at least as large as the other in \
every dimension.
- ``'same'``: output is the same size as ``in1``, centered \
with respect to the 'full' output
axes (scalar or tuple of scalar or None): Axes over which to compute
the convolution. The default is over all axes.
Returns:
cupy.ndarray: the result of convolution
.. seealso:: :func:`cupyx.scipy.signal.choose_conv_method`
.. seealso:: :func:`cupyx.scipy.signal.correlation`
.. seealso:: :func:`cupyx.scipy.signal.convolve`
.. seealso:: :func:`cupyx.scipy.signal.oaconvolve`
.. seealso:: :func:`cupyx.scipy.ndimage.convolve`
.. seealso:: :func:`scipy.signal.correlation`
"""
out = _st_core._check_conv_inputs(in1, in2, mode)
if out is not None:
return out
in1, in2, axes = _st_core._init_freq_conv_axes(in1, in2, mode, axes, False)
shape = [
max(x1, x2) if a not in axes else x1 + x2 - 1
for a, (x1, x2) in enumerate(zip(in1.shape, in2.shape))
]
out = _st_core._freq_domain_conv(in1, in2, axes, shape, calc_fast_len=True)
return _st_core._apply_conv_mode(out, in1.shape, in2.shape, mode, axes)
def _correlate(in1, in2, mode='full', method='auto', convolution=False):
quick_out = _st_core._check_conv_inputs(in1, in2, mode, convolution)
if quick_out is not None:
return quick_out
if method not in ('auto', 'direct', 'fft'):
raise ValueError("acceptable methods are 'auto', 'direct', or 'fft'")
if method == 'auto':
method = choose_conv_method(in1, in2, mode=mode)
if method == 'direct':
return _st_core._direct_correlate(in1, in2, mode, in1.dtype,
convolution)
# if method == 'fft':
raise ValueError('fftconvolve currently not supported')
def _correlate2d(in1, in2, mode, boundary, fillvalue, convolution=False):
if not (in1.ndim == in2.ndim == 2):
raise ValueError('{} inputs must both be 2-D arrays'.format(
'convolve2d' if convolution else 'correlate2d'))
_boundaries = {
'fill': 'constant', 'pad': 'constant',
'wrap': 'wrap', 'circular': 'wrap',
'symm': 'reflect', 'symmetric': 'reflect',
}
boundary = _boundaries.get(boundary)
if boundary is None:
raise ValueError('Acceptable boundary flags are "fill" (or "pad"), '
'"circular" (or "wrap"), and '
'"symmetric" (or "symm").')
quick_out = _st_core._check_conv_inputs(in1, in2, mode, convolution)
if quick_out is not None:
return quick_out
return _st_core._direct_correlate(in1, in2, mode, in1.dtype, convolution,
boundary, fillvalue, not convolution)
def oaconvolve(in1, in2, mode="full", axes=None):
"""Convolve two N-dimensional arrays using the overlap-add method.
Convolve ``in1`` and ``in2`` using the overlap-add method, with the output
size determined by the ``mode`` argument. This is generally faster than
``convolve`` for large arrays, and generally faster than ``fftconvolve``
when one array is much larger than the other, but can be slower when only a
few output values are needed or when the arrays are very similar in shape,
and can only output float arrays (int or object array inputs will be cast
to float).
Args:
in1 (cupy.ndarray): First input.
in2 (cupy.ndarray): Second input. Should have the same number of
dimensions as ``in1``.
mode (str): Indicates the size of the output:
- ``'full'``: output is the full discrete linear \
cross-correlation (default)
- ``'valid'``: output consists only of those elements that do \
not rely on the zero-padding. Either ``in1`` or \
``in2`` must be at least as large as the other in \
every dimension.
- ``'same'``: output is the same size as ``in1``, centered \
with respect to the ``'full'`` output
axes (scalar or tuple of scalar or None): Axes over which to compute
the convolution. The default is over all axes.
Returns:
cupy.ndarray: the result of convolution
.. seealso:: :func:`cupyx.scipy.signal.convolve`
.. seealso:: :func:`cupyx.scipy.signal.fftconvolve`
.. seealso:: :func:`cupyx.scipy.ndimage.convolve`
.. seealso:: :func:`scipy.signal.oaconvolve`
"""
out = _st_core._check_conv_inputs(in1, in2, mode)
if out is not None:
return out
if in1.shape == in2.shape: # Equivalent to fftconvolve
return fftconvolve(in1, in2, mode=mode, axes=axes)
in1, in2, axes = _st_core._init_freq_conv_axes(in1,
in2,
mode,
axes,
sorted_axes=True)
s1, s2 = in1.shape, in2.shape
if not axes:
return _st_core._apply_conv_mode(in1 * in2, s1, s2, mode, axes)
# Calculate the block sizes for the output, steps, first and second inputs.
# It is simpler to calculate them all together than doing them in separate
# loops due to all the special cases that need to be handled.
optimal_sizes = (_st_core._calc_oa_lens(s1[i], s2[i]) if i in axes else
(-1, -1, s1[i], s2[i]) for i in range(in1.ndim))
block_size, overlaps, in1_step, in2_step = zip(*optimal_sizes)
# Fall back to fftconvolve if there is only one block in every dimension
if in1_step == s1 and in2_step == s2:
return fftconvolve(in1, in2, mode=mode, axes=axes)
# Pad and reshape the inputs for overlapping and adding
shape_final = [
s1[i] + s2[i] - 1 if i in axes else None for i in range(in1.ndim)
]
in1, in2 = _st_core._oa_reshape_inputs(in1, in2, axes, shape_final,
block_size, overlaps, in1_step,
in2_step)
# Reshape the overlap-add parts to input block sizes
split_axes = [iax + i for i, iax in enumerate(axes)]
fft_axes = [iax + 1 for iax in split_axes]
# Do the convolution
fft_shape = [block_size[i] for i in axes]
ret = _st_core._freq_domain_conv(in1,
in2,
fft_axes,
fft_shape,
calc_fast_len=False)
# Do the overlap-add
for ax, ax_fft, ax_split in zip(axes, fft_axes, split_axes):
overlap = overlaps[ax]
if overlap is None:
continue
ret, overpart = cupy.split(ret, [-overlap], ax_fft)
overpart = cupy.split(overpart, [-1], ax_split)[0]
ret_overpart = cupy.split(ret, [overlap], ax_fft)[0]
ret_overpart = cupy.split(ret_overpart, [1], ax_split)[1]
ret_overpart += overpart
# Reshape back to the correct dimensionality
shape_ret = [
ret.shape[i] if i not in fft_axes else ret.shape[i] * ret.shape[i - 1]
for i in range(ret.ndim) if i not in split_axes
]
ret = ret.reshape(*shape_ret)
# Slice to the correct size
ret = ret[tuple([slice(islice) for islice in shape_final])]
return _st_core._apply_conv_mode(ret, s1, s2, mode, axes)
