File size: 1,099 Bytes
6370773 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import pytest
import numpy as np
from numpy.ma import masked_array
from numpy.testing import assert_array_equal
def test_matrix_transpose_raises_error_for_1d():
msg = "matrix transpose with ndim < 2 is undefined"
ma_arr = masked_array(data=[1, 2, 3, 4, 5, 6],
mask=[1, 0, 1, 1, 1, 0])
with pytest.raises(ValueError, match=msg):
ma_arr.mT
def test_matrix_transpose_equals_transpose_2d():
ma_arr = masked_array(data=[[1, 2, 3], [4, 5, 6]],
mask=[[1, 0, 1], [1, 1, 0]])
assert_array_equal(ma_arr.T, ma_arr.mT)
ARRAY_SHAPES_TO_TEST = (
(5, 2),
(5, 2, 3),
(5, 2, 3, 4),
)
@pytest.mark.parametrize("shape", ARRAY_SHAPES_TO_TEST)
def test_matrix_transpose_equals_swapaxes(shape):
num_of_axes = len(shape)
vec = np.arange(shape[-1])
arr = np.broadcast_to(vec, shape)
rng = np.random.default_rng(42)
mask = rng.choice([0, 1], size=shape)
ma_arr = masked_array(data=arr, mask=mask)
tgt = np.swapaxes(arr, num_of_axes - 2, num_of_axes - 1)
assert_array_equal(tgt, ma_arr.mT)
|