def _normalize_axis_indices(axes, ndim): # NOQA """Normalize axis indices. Args: axis (int, tuple of int or None): The un-normalized indices of the axis. Can be negative. ndim (int): The number of dimensions of the array that ``axis`` should be normalized against Returns: tuple of int: The tuple of normalized axis indices. """ if axes is None: axes = tuple(range(ndim)) elif not isinstance(axes, tuple): axes = (axes, ) res = [] for axis in axes: axis = _normalize_axis_index(axis, ndim) if axis in res: raise ValueError("Duplicate value in 'axis'") res.append(axis) return tuple(sorted(res))
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_along_axis` """ 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') _core.elementwise_copy(input, output) 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') _core.elementwise_copy(input, output) 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 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)
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 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 roll(a, shift, axis=None): """Roll array elements along a given axis. Elements that roll beyond the last position are re-introduced at the first. Args: a (~cupy.ndarray): Array to be rolled. shift (int or tuple of int): The number of places by which elements are shifted. If a tuple, then `axis` must be a tuple of the same size, and each of the given axes is shifted by the corresponding number. If an int while `axis` is a tuple of ints, then the same value is used for all given axes. axis (int or tuple of int or None): The axis along which elements are shifted. By default, the array is flattened before shifting, after which the original shape is restored. Returns: ~cupy.ndarray: Output array. .. seealso:: :func:`numpy.roll` """ if axis is None: return roll(a.ravel(), shift, 0).reshape(a.shape) axes = (axis,) if numpy.isscalar(axis) else axis axes = tuple([ # allow_duplicate internal._normalize_axis_index(ax, a.ndim) for ax in axes ]) if isinstance(shift, cupy.ndarray): shift = shift.ravel() n_axes = max(len(axes), shift.size) axes = numpy.broadcast_to(axes, (n_axes,)) shift = cupy.broadcast_to(shift, (n_axes,)) # TODO(asi1024): Improve after issue #4799 is resolved. indices = [] for ax in range(a.ndim): ind_shape = [1] * a.ndim ind_shape[ax] = a.shape[ax] indices.append(cupy.arange(a.shape[ax]).reshape(ind_shape)) for ax, s in zip(axes, shift): indices[ax] -= s indices[ax] %= a.shape[ax] for ax in range(a.ndim): indices[ax] = cupy.broadcast_to(indices[ax], a.shape) return a[tuple(indices)] else: broadcasted = numpy.broadcast(shift, axes) if broadcasted.nd > 1: raise ValueError( '\'shift\' and \'axis\' should be scalars or 1D sequences') shifts = {ax: 0 for ax in range(a.ndim)} for sh, ax in broadcasted: shifts[ax] += sh rolls = [((slice(None), slice(None)),)] * a.ndim for ax, offset in shifts.items(): offset %= a.shape[ax] or 1 # If `a` is empty, nothing matters. if offset: # (original, result), (original, result) rolls[ax] = ((slice(None, -offset), slice(offset, None)), (slice(-offset, None), slice(None, offset))) result = cupy.empty_like(a) for indices in itertools.product(*rolls): arr_index, res_index = zip(*indices) result[res_index] = a[arr_index] return result
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') _core.elementwise_copy(input, output) # 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