def test_np_5d_set_axes(self): x = np.zeros([6, 2, 5, 3, 4]) shape = None axes = [4, 1, 2] shape_expected = np.array([4, 2, 5]) axes_expected = np.array([4, 1, 2]) shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected)
def test_np_5d_set_shape_axes(self): x = np.zeros([6, 2, 5, 3, 4]) shape = [10, -1, 2] axes = [1, 0, 3] shape_expected = np.array([10, 6, 2]) axes_expected = np.array([1, 0, 3]) shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected)
def test_np_2d_defaults(self): x = np.arange(0, 1, .1).reshape(5, 2) shape = None axes = None shape_expected = np.array([5, 2]) axes_expected = np.array([0, 1]) shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected)
def test_np_5d_set_shape(self): x = np.zeros([6, 2, 5, 3, 4]) shape = [10, -1, -1, 1, 4] axes = None shape_expected = np.array([10, 2, 5, 1, 4]) axes_expected = np.array([0, 1, 2, 3, 4]) shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected)
def test_py_2d_defaults(self): x = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]) shape = None axes = None shape_expected = np.array([2, 4]) axes_expected = np.array([0, 1]) shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected)
def test_py_1d_defaults(self): x = np.array([1, 2, 3]) shape = None axes = None shape_expected = np.array([3]) axes_expected = np.array([0]) shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected)
def test_np_5d_defaults(self): x = np.zeros([6, 2, 5, 3, 4]) shape = None axes = None shape_expected = np.array([6, 2, 5, 3, 4]) axes_expected = np.array([0, 1, 2, 3, 4]) shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected) shape_res, axes_res = _init_nd_shape_and_axes_sorted(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected)
def test_np_1d_defaults(self): x = np.arange(0, 1, .1) shape = None axes = None shape_expected = np.array([10]) axes_expected = np.array([0]) shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected) shape_res, axes_res = _init_nd_shape_and_axes_sorted(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected)
def test_py_0d_defaults(self): x = 4 shape = None axes = None shape_expected = np.array([]) axes_expected = np.array([]) shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected) shape_res, axes_res = _init_nd_shape_and_axes_sorted(x, shape, axes) assert_equal(shape_res, shape_expected) assert_equal(axes_res, axes_expected)
def gpu_fftconvolve(in1, in2, axes=None): _, axes = _init_nd_shape_and_axes(in1, shape=None, axes=axes) s1 = in1.shape s2 = in2.shape shape = [ max((s1[i], s2[i])) if i not in axes else s1[i] + s2[i] - 1 for i in range(in1.ndim) ] fshape = [next_fast_len(shape[a], True) for a in axes] sp1 = cp.fft.rfft2(in1, fshape, axes=axes) sp2 = cp.fft.rfft2(in2, fshape, axes=axes) ret = cp.fft.irfft2(sp1 * sp2, fshape, axes=axes) fslice = tuple([slice(sz) for sz in shape]) ret = ret[fslice] return _centered(ret, s1).copy()
def _init_nd_shape_and_axes_sorted(x, shape, axes): """Handle and sort shape and axes arguments for n-dimensional transforms. This is identical to `_init_nd_shape_and_axes`, except the axes are returned in sorted order and the shape is reordered to match. Parameters ---------- x : array_like The input array. shape : int or array_like of ints or None The shape of the result. If both `shape` and `axes` (see below) are None, `shape` is ``x.shape``; if `shape` is None but `axes` is not None, then `shape` is ``scipy.take(x.shape, axes, axis=0)``. If `shape` is -1, the size of the corresponding dimension of `x` is used. axes : int or array_like of ints or None Axes along which the calculation is computed. The default is over all axes. Negative indices are automatically converted to their positive counterpart. Returns ------- shape : array The shape of the result. It is a 1D integer array. axes : array The shape of the result. It is a 1D integer array. """ noaxes = axes is None shape, axes = _init_nd_shape_and_axes(x, shape, axes) if not noaxes: shape = shape[axes.argsort()] axes.sort() return shape, axes
def test_errors(self): x = np.zeros(1) with assert_raises(ValueError, match="axes must be a scalar or " "iterable of integers"): _init_nd_shape_and_axes(x, shape=None, axes=[[1, 2], [3, 4]]) with assert_raises(ValueError, match="axes must be a scalar or " "iterable of integers"): _init_nd_shape_and_axes(x, shape=None, axes=[1., 2., 3., 4.]) with assert_raises(ValueError, match="axes exceeds dimensionality of input"): _init_nd_shape_and_axes(x, shape=None, axes=[1]) with assert_raises(ValueError, match="axes exceeds dimensionality of input"): _init_nd_shape_and_axes(x, shape=None, axes=[-2]) with assert_raises(ValueError, match="all axes must be unique"): _init_nd_shape_and_axes(x, shape=None, axes=[0, 0]) with assert_raises(ValueError, match="shape must be a scalar or " "iterable of integers"): _init_nd_shape_and_axes(x, shape=[[1, 2], [3, 4]], axes=None) with assert_raises(ValueError, match="shape must be a scalar or " "iterable of integers"): _init_nd_shape_and_axes(x, shape=[1., 2., 3., 4.], axes=None) with assert_raises(ValueError, match="when given, axes and shape arguments" " have to be of the same length"): _init_nd_shape_and_axes(np.zeros([1, 1, 1, 1]), shape=[1, 2, 3], axes=[1]) with assert_raises(ValueError, match="invalid number of data points" r" \(\[0\]\) specified"): _init_nd_shape_and_axes(x, shape=[0], axes=None) with assert_raises(ValueError, match="invalid number of data points" r" \(\[-2\]\) specified"): _init_nd_shape_and_axes(x, shape=-2, axes=None)
def test_shape_axes_subset(self): x = np.zeros((2, 3, 4, 5)) shape, axes = _init_nd_shape_and_axes(x, shape=(5, 5, 5), axes=None) assert_array_equal(shape, [5, 5, 5]) assert_array_equal(axes, [1, 2, 3])