Skip to content
Fetching contributors… Cannot retrieve contributors at this time
181 lines (129 sloc) 5.08 KB
 import itertools import numpy import cupy from cupy import core from cupy.core._kernel import _get_axis def flip(a, axis): """Reverse the order of elements in an array along the given axis. Note that ``flip`` function has been introduced since NumPy v1.12. The contents of this document is the same as the original one. Args: a (~cupy.ndarray): Input array. axis (int): Axis in array, which entries are reversed. Returns: ~cupy.ndarray: Output array. .. seealso:: :func:`numpy.flip` """ a_ndim = a.ndim if a_ndim < 1: raise core._AxisError('Input must be >= 1-d') axis = int(axis) if not -a_ndim <= axis < a_ndim: raise core._AxisError( 'axis must be >= %d and < %d' % (-a_ndim, a_ndim)) return _flip(a, axis) def fliplr(a): """Flip array in the left/right direction. Flip the entries in each row in the left/right direction. Columns are preserved, but appear in a different order than before. Args: a (~cupy.ndarray): Input array. Returns: ~cupy.ndarray: Output array. .. seealso:: :func:`numpy.fliplr` """ if a.ndim < 2: raise ValueError('Input must be >= 2-d') return a[::, ::-1] def flipud(a): """Flip array in the up/down direction. Flip the entries in each column in the up/down direction. Rows are preserved, but appear in a different order than before. Args: a (~cupy.ndarray): Input array. Returns: ~cupy.ndarray: Output array. .. seealso:: :func:`numpy.flipud` """ if a.ndim < 1: raise ValueError('Input must be >= 1-d') return a[::-1] 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) else: axis = _get_axis(axis, a.ndim) broadcasted = numpy.broadcast(shift, axis) 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 rot90(a, k=1, axes=(0, 1)): """Rotate an array by 90 degrees in the plane specified by axes. Note that ``axes`` argument has been introduced since NumPy v1.12. The contents of this document is the same as the original one. Args: a (~cupy.ndarray): Array of two or more dimensions. k (int): Number of times the array is rotated by 90 degrees. axes: (tuple of ints): The array is rotated in the plane defined by the axes. Axes must be different. Returns: ~cupy.ndarray: Output array. .. seealso:: :func:`numpy.rot90` """ a_ndim = a.ndim if a_ndim < 2: raise ValueError('Input must be >= 2-d') axes = tuple(axes) if len(axes) != 2: raise ValueError('len(axes) must be 2') if axes == axes or abs(axes - axes) == a_ndim: raise ValueError('axes must be different') if not (-a_ndim <= axes < a_ndim and -a_ndim <= axes < a_ndim): raise ValueError('axes must be >= %d and < %d' % (-a_ndim, a_ndim)) k = k % 4 if k == 0: return a[:] if k == 2: return _flip(_flip(a, axes), axes) axes_t = list(range(0, a_ndim)) axes_t[axes], axes_t[axes] = axes_t[axes], axes_t[axes] if k == 1: return cupy.transpose(_flip(a, axes), axes_t) else: return _flip(cupy.transpose(a, axes_t), axes) def _flip(a, axis): # This function flips array without checking args. indexer = [slice(None)] * a.ndim indexer[axis] = slice(None, None, -1) return a[tuple(indexer)]
You can’t perform that action at this time.