def apply_along_axis(func1d, axis, arr, *args, **kwargs):
"""Apply a function to 1-D slices along the given axis.
Args:
func1d (function (M,) -> (Nj...)): This function should accept 1-D
arrays. It is applied to 1-D slices of ``arr`` along the specified
axis. It must return a 1-D ``cupy.ndarray``.
axis (integer): Axis along which ``arr`` is sliced.
arr (cupy.ndarray (Ni..., M, Nk...)): Input array.
args: Additional arguments for ``func1d``.
kwargs: Additional keyword arguments for ``func1d``.
Returns:
cupy.ndarray: The output array. The shape of ``out`` is identical to
the shape of ``arr``, except along the ``axis`` dimension. This
axis is removed, and replaced with new dimensions equal to the
shape of the return value of ``func1d``. So if ``func1d`` returns a
scalar ``out`` will have one fewer dimensions than ``arr``.
.. seealso:: :func:`numpy.apply_over_axes`
"""
ndim = arr.ndim
axis = internal._normalize_axis_index(axis, ndim)
inarr_view = cupy.moveaxis(arr, axis, -1)
# compute indices for the iteration axes, and append a trailing ellipsis to
# prevent 0d arrays decaying to scalars
inds = index_tricks.ndindex(inarr_view.shape[:-1])
inds = (ind + (Ellipsis, ) for ind in inds)
# invoke the function on the first item
try:
ind0 = next(inds)
except StopIteration:
raise ValueError(
'Cannot apply_along_axis when any iteration dimensions are 0')
res = func1d(inarr_view[ind0], *args, **kwargs)
if cupy.isscalar(res):
# scalar outputs need to be transfered to a device ndarray
res = cupy.asarray(res)
# build a buffer for storing evaluations of func1d.
# remove the requested axis, and add the new ones on the end.
# laid out so that each write is contiguous.
# for a tuple index inds, buff[inds] = func1d(inarr_view[inds])
buff = cupy.empty(inarr_view.shape[:-1] + res.shape, res.dtype)
# save the first result, then compute and save all remaining results
buff[ind0] = res
for ind in inds:
buff[ind] = func1d(inarr_view[ind], *args, **kwargs)
# restore the inserted axes back to where they belong
for i in range(res.ndim):
buff = cupy.moveaxis(buff, -1, axis)
return buff
def _convert_1d_args(ndim, weights, origin, axis):
if weights.ndim != 1 or weights.size < 1:
raise RuntimeError('incorrect filter size')
axis = internal._normalize_axis_index(axis, ndim)
w_shape = [1] * ndim
w_shape[axis] = weights.size
weights = weights.reshape(w_shape)
origins = [0] * ndim
origins[axis] = _util._check_origin(origin, weights.size)
return weights, tuple(origins)
def fourier_gaussian(input, sigma, n=-1, axis=-1, output=None):
"""Multidimensional Gaussian shift filter.
The array is multiplied with the Fourier transform of a (separable)
Gaussian kernel.
Args:
input (cupy.ndarray): The input array.
sigma (float or sequence of float): The sigma of the Gaussian kernel.
If a float, `sigma` is the same for all axes. If a sequence,
`sigma` has to contain one value for each axis.
n (int, optional): If `n` is negative (default), then the input is
assumed to be the result of a complex fft. If `n` is larger than or
equal to zero, the input is assumed to be the result of a real fft,
and `n` gives the length of the array before transformation along
the real transform direction.
axis (int, optional): The axis of the real transform (only used when
``n > -1``).
output (cupy.ndarray, optional):
If given, the result of shifting the input is placed in this array.
Returns:
output (cupy.ndarray): The filtered output.
"""
ndim = input.ndim
output = _get_output_fourier(output, input)
axis = internal._normalize_axis_index(axis, ndim)
sigmas = _util._fix_sequence_arg(sigma, ndim, 'sigma')
output[...] = input
for ax, (sigmak, ax_size) in enumerate(zip(sigmas, output.shape)):
# compute the frequency grid in Hz
if ax == axis and n > 0:
arr = cupy.arange(ax_size, dtype=output.real.dtype)
arr /= n
else:
arr = cupy.fft.fftfreq(ax_size)
arr = arr.astype(output.real.dtype, copy=False)
# compute the Gaussian weights
arr *= arr
scale = sigmak * sigmak / -2
arr *= (4 * numpy.pi * numpy.pi) * scale
cupy.exp(arr, out=arr)
# reshape for broadcasting
arr = _reshape_nd(arr, ndim=ndim, axis=ax)
output *= arr
return output
def fourier_shift(input, shift, n=-1, axis=-1, output=None):
"""Multidimensional Fourier shift filter.
The array is multiplied with the Fourier transform of a shift operation.
Args:
input (cupy.ndarray): The input array. This should be in the Fourier
domain.
shift (float or sequence of float): The size of shift. If a float,
`shift` is the same for all axes. If a sequence, `shift` has to
contain one value for each axis.
n (int, optional): If `n` is negative (default), then the input is
assumed to be the result of a complex fft. If `n` is larger than or
equal to zero, the input is assumed to be the result of a real fft,
and `n` gives the length of the array before transformation along
the real transform direction.
axis (int, optional): The axis of the real transform (only used when
``n > -1``).
output (cupy.ndarray, optional):
If given, the result of shifting the input is placed in this array.
Returns:
output (cupy.ndarray): The shifted output (in the Fourier domain).
"""
ndim = input.ndim
output = _get_output_fourier(output, input, complex_only=True)
axis = internal._normalize_axis_index(axis, ndim)
shifts = _util._fix_sequence_arg(shift, ndim, 'shift')
output[...] = input
for ax, (shiftk, ax_size) in enumerate(zip(shifts, output.shape)):
if shiftk == 0:
continue
if ax == axis and n > 0:
# cp.fft.rfftfreq(ax_size) * (-2j * numpy.pi * shiftk * ax_size/n)
arr = cupy.arange(ax_size, dtype=output.dtype)
arr *= -2j * numpy.pi * shiftk / n
else:
arr = cupy.fft.fftfreq(ax_size)
arr = arr * (-2j * numpy.pi * shiftk)
cupy.exp(arr, out=arr)
# reshape for broadcasting
arr = _reshape_nd(arr, ndim=ndim, axis=ax)
output *= arr
return output
def take_along_axis(a, indices, axis):
"""Take values from the input array by matching 1d index and data slices.
Args:
a (cupy.ndarray): Array to extract elements.
indices (cupy.ndarray): Indices to take along each 1d slice of ``a``.
axis (int): The axis to take 1d slices along.
Returns:
cupy.ndarray: The indexed result.
.. seealso:: :func:`numpy.take_along_axis`
"""
if indices.dtype.kind not in ('i', 'u'):
raise IndexError('`indices` must be an integer array')
if axis is None:
a = a.ravel()
axis = 0
ndim = a.ndim
axis = internal._normalize_axis_index(axis, ndim)
if ndim != indices.ndim:
raise ValueError(
'`indices` and `a` must have the same number of dimensions')
fancy_index = []
for i, n in enumerate(a.shape):
if i == axis:
fancy_index.append(indices)
else:
ind_shape = (1, ) * i + (-1, ) + (1, ) * (ndim - i - 1)
fancy_index.append(cupy.arange(n).reshape(ind_shape))
return a[tuple(fancy_index)]
def diff(a, n=1, axis=-1, prepend=None, append=None):
"""Calculate the n-th discrete difference along the given axis.
Args:
a (cupy.ndarray): Input array.
n (int): The number of times values are differenced. If zero, the input
is returned as-is.
axis (int): The axis along which the difference is taken, default is
the last axis.
prepend (int, float, cupy.ndarray): Value to prepend to ``a``.
append (int, float, cupy.ndarray): Value to append to ``a``.
Returns:
cupy.ndarray: The result array.
.. seealso:: :func:`numpy.diff`
"""
if n == 0:
return a
if n < 0:
raise ValueError("order must be non-negative but got " + repr(n))
a = cupy.asanyarray(a)
nd = a.ndim
axis = internal._normalize_axis_index(axis, nd)
combined = []
if prepend is not None:
prepend = cupy.asanyarray(prepend)
if prepend.ndim == 0:
shape = list(a.shape)
shape[axis] = 1
prepend = cupy.broadcast_to(prepend, tuple(shape))
combined.append(prepend)
combined.append(a)
if append is not None:
append = cupy.asanyarray(append)
if append.ndim == 0:
shape = list(a.shape)
shape[axis] = 1
append = cupy.broadcast_to(append, tuple(shape))
combined.append(append)
if len(combined) > 1:
a = cupy.concatenate(combined, axis)
slice1 = [slice(None)] * nd
slice2 = [slice(None)] * nd
slice1[axis] = slice(1, None)
slice2[axis] = slice(None, -1)
slice1 = tuple(slice1)
slice2 = tuple(slice2)
op = cupy.not_equal if a.dtype == numpy.bool_ else cupy.subtract
for _ in range(n):
a = op(a[slice1], a[slice2])
return a
def spline_filter1d(input,
order=3,
axis=-1,
output=cupy.float64,
mode='mirror'):
"""
Calculate a 1-D spline filter along the given axis.
The lines of the array along the given axis are filtered by a
spline filter. The order of the spline must be >= 2 and <= 5.
Args:
input (cupy.ndarray): The input array.
order (int): The order of the spline interpolation, default is 3. Must
be in the range 0-5.
axis (int): The axis along which the spline filter is applied. Default
is the last axis.
output (cupy.ndarray or dtype, optional): The array in which to place
the output, or the dtype of the returned array. Default is
``numpy.float64``.
mode (str): Points outside the boundaries of the input are filled
according to the given mode (``'constant'``, ``'nearest'``,
``'mirror'``, ``'reflect'``, ``'wrap'``, ``'grid-mirror'``,
``'grid-wrap'``, ``'grid-constant'`` or ``'opencv'``).
Returns:
cupy.ndarray: The result of prefiltering the input.
.. seealso:: :func:`scipy.spline_filter1d`
"""
if order < 0 or order > 5:
raise RuntimeError('spline order not supported')
x = input
ndim = x.ndim
axis = internal._normalize_axis_index(axis, ndim)
# order 0, 1 don't require reshaping as no CUDA kernel will be called
# scalar or size 1 arrays also don't need to be filtered
run_kernel = not (order < 2 or x.ndim == 0 or x.shape[axis] == 1)
if not run_kernel:
output = _util._get_output(output, input)
output[...] = x[...]
return output
temp, data_dtype, output_dtype = _get_spline_output(x, output)
data_type = cupy.core._scalar.get_typename(temp.dtype)
pole_type = cupy.core._scalar.get_typename(temp.real.dtype)
index_type = _util._get_inttype(input)
index_dtype = cupy.int32 if index_type == 'int' else cupy.int64
n_samples = x.shape[axis]
n_signals = x.size // n_samples
info = cupy.array((n_signals, n_samples) + x.shape, dtype=index_dtype)
# empirical choice of block size that seemed to work well
block_size = max(2**math.ceil(numpy.log2(n_samples / 32)), 8)
kern = _spline_prefilter_core.get_raw_spline1d_kernel(
axis,
ndim,
mode,
order=order,
index_type=index_type,
data_type=data_type,
pole_type=pole_type,
block_size=block_size,
)
# Due to recursive nature, a given line of data must be processed by a
# single thread. n_signals lines will be processed in total.
block = (block_size, )
grid = ((n_signals + block[0] - 1) // block[0], )
# apply prefilter gain
poles = _spline_prefilter_core.get_poles(order=order)
temp *= _spline_prefilter_core.get_gain(poles)
# apply caual + anti-causal IIR spline filters
kern(grid, block, (temp, info))
if isinstance(output, cupy.ndarray) and temp is not output:
# copy kernel output into the user-provided output array
output[...] = temp[...]
return output
return temp.astype(output_dtype, copy=False)
def fourier_ellipsoid(input, size, n=-1, axis=-1, output=None):
"""Multidimensional ellipsoid Fourier filter.
The array is multiplied with the fourier transform of a ellipsoid of
given sizes.
Args:
input (cupy.ndarray): The input array.
size (float or sequence of float): The size of the box used for
filtering. If a float, `size` is the same for all axes. If a
sequence, `size` has to contain one value for each axis.
n (int, optional): If `n` is negative (default), then the input is
assumed to be the result of a complex fft. If `n` is larger than or
equal to zero, the input is assumed to be the result of a real fft,
and `n` gives the length of the array before transformation along
the real transform direction.
axis (int, optional): The axis of the real transform (only used when
``n > -1``).
output (cupy.ndarray, optional):
If given, the result of shifting the input is placed in this array.
Returns:
output (cupy.ndarray): The filtered output.
"""
ndim = input.ndim
if ndim == 1:
return fourier_uniform(input, size, n, axis, output)
if ndim > 3:
# Note: SciPy currently does not do any filtering on >=4d inputs, but
# does not warn about this!
raise NotImplementedError('Only 1d, 2d and 3d inputs are supported')
output = _get_output_fourier(output, input)
axis = internal._normalize_axis_index(axis, ndim)
sizes = _util._fix_sequence_arg(size, ndim, 'size')
output[...] = input
# compute the distance from the origin for all samples in Fourier space
distance = 0
for ax, (size, ax_size) in enumerate(zip(sizes, output.shape)):
# compute the frequency grid in Hz
if ax == axis and n > 0:
arr = cupy.arange(ax_size, dtype=output.real.dtype)
arr *= numpy.pi * size / n
else:
arr = cupy.fft.fftfreq(ax_size)
arr *= numpy.pi * size
arr = arr.astype(output.real.dtype, copy=False)
arr *= arr
arr = _reshape_nd(arr, ndim=ndim, axis=ax)
distance = distance + arr
cupy.sqrt(distance, out=distance)
if ndim == 2:
special.j1(distance, out=output)
output *= 2
output /= distance
elif ndim == 3:
cupy.sin(distance, out=output)
output -= distance * cupy.cos(distance)
output *= 3
output /= distance**3
output[(0, ) * ndim] = 1.0 # avoid NaN in corner at frequency=0 location
output *= input
return output
def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):
"""Returns the cross product of two vectors.
The cross product of ``a`` and ``b`` in :math:`R^3` is a vector
perpendicular to both ``a`` and ``b``. If ``a`` and ``b`` are arrays
of vectors, the vectors are defined by the last axis of ``a`` and ``b``
by default, and these axes can have dimensions 2 or 3. Where the
dimension of either ``a`` or ``b`` is 2, the third component of the input
vector is assumed to be zero and the cross product calculated accordingly.
In cases where both input vectors have dimension 2, the z-component of
the cross product is returned.
Args:
a (cupy.ndarray): Components of the first vector(s).
b (cupy.ndarray): Components of the second vector(s).
axisa (int, optional):
Axis of ``a`` that defines the vector(s).
By default, the last axis.
axisb (int, optional):
Axis of ``b`` that defines the vector(s).
By default, the last axis.
axisc (int, optional):
Axis of ``c`` containing the cross product vector(s). Ignored if
both input vectors have dimension 2, as the return is scalar.
By default, the last axis.
axis (int, optional):
If defined, the axis of ``a``, ``b`` and ``c``
that defines the vector(s) and cross product(s).
Overrides ``axisa``, ``axisb`` and ``axisc``.
Returns:
cupy.ndarray :
Vector cross product(s).
.. seealso:: :func:`numpy.cross`
"""
if axis is not None:
axisa, axisb, axisc = (axis, ) * 3
a = cupy.asarray(a)
b = cupy.asarray(b)
# Check axisa and axisb are within bounds
axisa = internal._normalize_axis_index(axisa, a.ndim)
axisb = internal._normalize_axis_index(axisb, b.ndim)
# Move working axis to the end of the shape
a = cupy.moveaxis(a, axisa, -1)
b = cupy.moveaxis(b, axisb, -1)
if a.shape[-1] not in (2, 3) or b.shape[-1] not in (2, 3):
msg = ('incompatible dimensions for cross product\n'
'(dimension must be 2 or 3)')
raise ValueError(msg)
# Create the output array
shape = cupy.broadcast(a[..., 0], b[..., 0]).shape
if a.shape[-1] == 3 or b.shape[-1] == 3:
shape += (3, )
# Check axisc is within bounds
axisc = internal._normalize_axis_index(axisc, len(shape))
dtype = cupy.promote_types(a.dtype, b.dtype)
cp = cupy.empty(shape, dtype)
# create local aliases for readability
a0 = a[..., 0]
a1 = a[..., 1]
if a.shape[-1] == 3:
a2 = a[..., 2]
b0 = b[..., 0]
b1 = b[..., 1]
if b.shape[-1] == 3:
b2 = b[..., 2]
if cp.ndim != 0 and cp.shape[-1] == 3:
cp0 = cp[..., 0]
cp1 = cp[..., 1]
cp2 = cp[..., 2]
if a.shape[-1] == 2:
if b.shape[-1] == 2:
# a0 * b1 - a1 * b0
cupy.multiply(a0, b1, out=cp)
cp -= a1 * b0
return cp
else:
assert b.shape[-1] == 3
# cp0 = a1 * b2 - 0 (a2 = 0)
# cp1 = 0 - a0 * b2 (a2 = 0)
# cp2 = a0 * b1 - a1 * b0
cupy.multiply(a1, b2, out=cp0)
cupy.multiply(a0, b2, out=cp1)
cupy.negative(cp1, out=cp1)
cupy.multiply(a0, b1, out=cp2)
cp2 -= a1 * b0
else:
assert a.shape[-1] == 3
if b.shape[-1] == 3:
# cp0 = a1 * b2 - a2 * b1
# cp1 = a2 * b0 - a0 * b2
# cp2 = a0 * b1 - a1 * b0
cupy.multiply(a1, b2, out=cp0)
tmp = a2 * b1
cp0 -= tmp
cupy.multiply(a2, b0, out=cp1)
cupy.multiply(a0, b2, out=tmp)
cp1 -= tmp
cupy.multiply(a0, b1, out=cp2)
cupy.multiply(a1, b0, out=tmp)
cp2 -= tmp
else:
assert b.shape[-1] == 2
# cp0 = 0 - a2 * b1 (b2 = 0)
# cp1 = a2 * b0 - 0 (b2 = 0)
# cp2 = a0 * b1 - a1 * b0
cupy.multiply(a2, b1, out=cp0)
cupy.negative(cp0, out=cp0)
cupy.multiply(a2, b0, out=cp1)
cupy.multiply(a0, b1, out=cp2)
cp2 -= a1 * b0
return cupy.moveaxis(cp, -1, axisc)
