如何索引/切片 3D numpy 数组

Question

I'm relatively new to python/numpy. I have a 3D numpy array of TxNxN. It contains a sequence of symmetrical NxN matrices. I want convert it to a 2D array of TxM (where M = N(N+1)/2). How can I do that? I can certainly use 3 loops, but I thought there probably better ways to do that in python/numpy.

Answer 1

It seems that you want to get the upper triangle or lower triangle of each symmetric matrix. A simple method is to generate a mask array and apply it to each 2D array:

>>> e
array([[[0, 1, 2, 3],
        [1, 2, 3, 0],
        [2, 3, 0, 1],
        [3, 0, 1, 2]],

       [[1, 2, 3, 4],
        [2, 3, 4, 1],
        [3, 4, 1, 2],
        [4, 1, 2, 3]],

       [[2, 3, 4, 5],
        [3, 4, 5, 2],
        [4, 5, 2, 3],
        [5, 2, 3, 4]]])
>>> ii, jj = np.indices(e.shape[1:])
>>> jj >= ii
array([[ True,  True,  True,  True],
       [False,  True,  True,  True],
       [False, False,  True,  True],
       [False, False, False,  True]])
>>> e[:, jj >= ii]
array([[0, 1, 2, 3, 2, 3, 0, 0, 1, 2],
       [1, 2, 3, 4, 3, 4, 1, 1, 2, 3],
       [2, 3, 4, 5, 4, 5, 2, 2, 3, 4]])

Using the numpy.triu_indices function can do better, but you can't put the obtained indices tuple directly between square brackets. You need to unpack them first:

>>> i, j = np.triu_indices(e.shape[1])
>>> e[:, i, j]
array([[0, 1, 2, 3, 2, 3, 0, 0, 1, 2],
       [1, 2, 3, 4, 3, 4, 1, 1, 2, 3],
       [2, 3, 4, 5, 4, 5, 2, 2, 3, 4]])