Multi-dimensional array notation in Python-CodePudding

I have two arrays A and i with dimensions (1, 3, 3) and (1, 2, 2) respectively. I want to define a new array I which gives the elements of A based on i. The current and desired outputs are attached.

import numpy as np
i=np.array([[[0,0],[1,2],[2,2]]])
A = np.array([[[1,2,3],[4,5,6],[7,8,9]]], dtype=float)
I=A[0,i]
print([I])

The current output is

[array([[[[1.000000000, 2.000000000, 3.000000000],
         [1.000000000, 2.000000000, 3.000000000]],

        [[4.000000000, 5.000000000, 6.000000000],
         [7.000000000, 8.000000000, 9.000000000]],

        [[7.000000000, 8.000000000, 9.000000000],
         [7.000000000, 8.000000000, 9.000000000]]]])]

The desired output is

[array(([[[1],[6],[9]]]))

CodePudding user response：

In [131]: A.shape, i.shape
Out[131]: ((1, 3, 3), (1, 3, 2))

That leading size 1 dimension just adds a [] layer, and complicates indexing (a bit):

In [132]: A[0]
Out[132]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

This is the indexing that I think you want:

In [133]: A[0,i[0,:,0],i[0,:,1]]
Out[133]: array([1, 6, 9])

If you really need a trailing size 1 dimension, add it after:

In [134]: A[0,i[0,:,0],i[0,:,1]][:,None]
Out[134]: 
array([[1],
       [6],
       [9]])

From the desired numbers, I deduced that you wanted to use the 2 columns of i as indices to two different dimensions of A:

In [135]: i[0]
Out[135]: 
array([[0, 0],
       [1, 2],
       [2, 2]])

Another way to do the same thing:

In [139]: tuple(i.T)
Out[139]: 
(array([[0],
        [1],
        [2]]),
 array([[0],
        [2],
        [2]]))

In [140]: A[0][tuple(i.T)]
Out[140]: 
array([[1],
       [6],
       [9]])

CodePudding user response：

You must enter

I=A[0,:1,i[:,1]]

CodePudding user response：

You can use numpy's take for that. However, take works with a flat index, so you will need to use [0, 5, 8] for your indexes instead.

Here is an example:

>>> I = [A.shape[2] * x   y for x,y in i[0]] # Convert to flat indexes
>>> I = np.expand_dims(I, axis=(1,2))
>>> A.take(I)
array([[[1.]],
       [[6.]],
       [[9.]]])