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