def test_flatten_axes_swapped():
a = seq_array((5,4,3,2))
assert_array_equal(flatten_axes(a, (0, 3)), flatten_axes(a, (3, 0)))
def test_flatten_axes_masked():
arr = seq_array((5, 4, 3, 2))
a = np.ma.array(arr, mask=(arr % 2 == 0))
b = flatten_axes(a, (0, 1, 2, 3))
assert_array_equal(np.arange(arr.size), b.data)
assert_array_equal(np.arange(arr.size) % 2 == 0, b.mask)
def test_flatten_axes_fully():
a = seq_array((5, 4, 3, 2))
assert_array_equal(np.arange(5*4*3*2), flatten_axes(a, (0, 1, 2, 3)))
def test_flatten_axies_2d():
a = seq_array((2, 2, 2))
f = flatten_axes(a, (0, 2))
expected = np.array([[0, 1, 4, 5],
[2, 3, 6, 7]])
assert expected == f
def test_flatten_axes():
a = seq_array((5, 4, 3, 2))
b = flatten_axes(a, (2, 3))
c = seq_array((5, 4, 6))
assert (5,4,6) == b.shape
assert_array_equal(c, b)
def test_flatten_axes_one_dim():
"""Flattening one dimension does not change the array"""
a = seq_array((5, 4, 3, 2))
b = flatten_axes(a, -1)
assert_array_equal(a, b)
