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])
