How to convert Numpy index to coordinates?-CodePudding

I wanted to customize getitem for my array class. But I don't know how to convert the indices to coordinates. For example, for an array with shape (2,2,2), if one requests A[...,1], I'd like to get the coordinates as [[0,0,1],[0,1,1],[1,0,1],[1,1,1]]. Are there any functions in Numpy that can achieve this? Thanks.

CodePudding user response：

Take a look at np.indices:

In [199]: idx = np.indices((2,2,2))

In [200]: idx.shape
Out[200]: (3, 2, 2, 2)

In [201]: idx
Out[201]: 
array([[[[0, 0],
         [0, 0]],

        [[1, 1],
         [1, 1]]],


       [[[0, 0],
         [1, 1]],

        [[0, 0],
         [1, 1]]],


       [[[0, 1],
         [0, 1]],

        [[0, 1],
         [0, 1]]]])

That's like your A but with an added size 3 dimension.

In [202]: idx[:,:,:,1]
Out[202]: 
array([[[0, 0],
        [1, 1]],

       [[0, 1],
        [0, 1]],

       [[1, 1],
        [1, 1]]])

Which can be turned into your desired array with:

In [203]: idx[:,:,:,1].reshape(3,4).T
Out[203]: 
array([[0, 0, 1],
       [0, 1, 1],
       [1, 0, 1],
       [1, 1, 1]])

CodePudding user response：

class Class(np.ndarray):
    def __new__(cls, data, time=None):
        obj = np.asarray(data).view(cls)
        obj.time = time
        return obj
    def __array_finalize__(self, obj):
        setattr(self, 'time', obj.time)
        try:
            self.time = self.time[obj._new_time_index]
        except:
            pass

    def __getitem__(self, item):
        try:
            if isinstance(item, (slice, int)):
                self._new_time_index = item
            else:
                self._new_time_index = item[0]
        except: 
            pass
        return super().__getitem__(item)