def test_np_5d_set_shape_axes_sorted(self):
x = np.zeros([6, 2, 5, 3, 4])
shape = [10, -1, 2]
axes = [1, 0, 3]
shape_expected = np.array([6, 10, 2])
axes_expected = np.array([0, 1, 3])
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_5d_set_axes_sorted(self):
x = np.zeros([6, 2, 5, 3, 4])
shape = None
axes = [4, 1, 2]
shape_expected = np.array([2, 5, 4])
axes_expected = np.array([1, 2, 4])
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_5d_set_axes_sorted(self):
x = np.zeros([6, 2, 5, 3, 4])
shape = None
axes = [4, 1, 2]
shape_expected = np.array([2, 5, 4])
axes_expected = np.array([1, 2, 4])
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_5d_set_shape_axes_sorted(self):
x = np.zeros([6, 2, 5, 3, 4])
shape = [10, -1, 2]
axes = [1, 0, 3]
shape_expected = np.array([6, 10, 2])
axes_expected = np.array([0, 1, 3])
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_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_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)
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_1d_defaults(self):
x = [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)
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_0d_defaults(self):
x = np.array(7.)
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 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_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)
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_2d_defaults(self):
x = [[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)
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_1d_defaults(self):
x = [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)
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_0d_defaults(self):
x = np.array(7.)
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 test_py_2d_defaults(self):
x = [[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)
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 _raw_fftnd(x, s, axes, direction, overwrite_x, work_function):
"""Internal auxiliary function for fftnd, ifftnd."""
noaxes = axes is None
s, axes = _init_nd_shape_and_axes_sorted(x, s, axes)
# No need to swap axes, array is in C order
if noaxes:
for ax in axes:
x, copy_made = _fix_shape(x, s[ax], ax)
overwrite_x = overwrite_x or copy_made
return work_function(x, s, direction, overwrite_x=overwrite_x)
# Swap the request axes, last first (i.e. First swap the axis which ends up
# at -1, then at -2, etc...), such as the request axes on which the
# operation is carried become the last ones
for i in range(1, axes.size + 1):
x = numpy.swapaxes(x, axes[-i], -i)
# We can now operate on the axes waxes, the p last axes (p = len(axes)), by
# fixing the shape of the input array to 1 for any axis the fft is not
# carried upon.
waxes = list(range(x.ndim - axes.size, x.ndim))
shape = numpy.ones(x.ndim)
shape[waxes] = s
for i in range(len(waxes)):
x, copy_made = _fix_shape(x, s[i], waxes[i])
overwrite_x = overwrite_x or copy_made
r = work_function(x, shape, direction, overwrite_x=overwrite_x)
# reswap in the reverse order (first axis first, etc...) to get original
# order
for i in range(len(axes), 0, -1):
r = numpy.swapaxes(r, -i, axes[-i])
return r
def _raw_fftnd(x, s, axes, direction, overwrite_x, work_function):
"""Internal auxiliary function for fftnd, ifftnd."""
noaxes = axes is None
s, axes = _init_nd_shape_and_axes_sorted(x, s, axes)
# No need to swap axes, array is in C order
if noaxes:
for ax in axes:
x, copy_made = _fix_shape(x, s[ax], ax)
overwrite_x = overwrite_x or copy_made
return work_function(x, s, direction, overwrite_x=overwrite_x)
# Swap the request axes, last first (i.e. First swap the axis which ends up
# at -1, then at -2, etc...), such as the request axes on which the
# operation is carried become the last ones
for i in range(1, axes.size+1):
x = numpy.swapaxes(x, axes[-i], -i)
# We can now operate on the axes waxes, the p last axes (p = len(axes)), by
# fixing the shape of the input array to 1 for any axis the fft is not
# carried upon.
waxes = list(range(x.ndim - axes.size, x.ndim))
shape = numpy.ones(x.ndim)
shape[waxes] = s
for i in range(len(waxes)):
x, copy_made = _fix_shape(x, s[i], waxes[i])
overwrite_x = overwrite_x or copy_made
r = work_function(x, shape, direction, overwrite_x=overwrite_x)
# reswap in the reverse order (first axis first, etc...) to get original
# order
for i in range(len(axes), 0, -1):
r = numpy.swapaxes(r, -i, axes[-i])
return r
