codekingpro/portable-devtools

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