MaxB - 6 months ago 14

Python Question

Suppose I have a 3D array (tensor)

`print a.shape`

(100, 100, 100)

and want to index its first axis using one array:

`print i.shape`

(20,)

and its last axis using another array

`print j.shape`

(30,)

My intent is to get a (20, 100, 30) array, for example, to be used in assignments like

`a[?!, :, ?!] = b`

However, I can't figure out how.

`print a[i.reshape(20, 1, 1), :, j.reshape(1, 1, 30)].shape`

(20, 1, 30, 100)

print a[i.reshape(20, 1, 1), :, j].shape

(20, 1, 30, 100)

print a[i.reshape(20, 1), :, j].shape

(20, 30, 100)

print a[i.reshape(20, 1, 1), j.reshape(1, 1, 30)].shape

(20, 1, 30, 100)

As I understand the "advanced indexing" rules, the first attempt should have worked, but I didn't even end up with a 3D array, and the full dimension (100) came at the end instead of the middle.

Answer

**Approach #1 :** You can use `np.ix_`

for getting such meshes of indices and simply indexing into the input array must give you the desired output. Thus, an implementation would be like so -

```
a[np.ix_(i,np.arange(a.shape[1]),j)]
```

**Approach #2 :** Simpler way using broadcasted indexing -

```
a[i[:,None],:,j].swapaxes(1,2)
```

Verify with a sample case -

```
In [24]: a = np.random.randint(0,99,(5,3,5))
In [25]: i = np.random.randint(0,5,(2))
In [26]: j = np.random.randint(0,5,(2))
In [27]: a[i[0],:,j[0]]
Out[27]: array([15, 7, 74])
In [28]: a[i[0],:,j[1]]
Out[28]: array([32, 19, 85])
In [29]: a[i[1],:,j[0]]
Out[29]: array([76, 65, 96])
In [30]: a[i[1],:,j[1]]
Out[30]: array([54, 65, 66])
In [31]: a[np.ix_(i,np.arange(a.shape[1]),j)]
Out[31]:
array([[[15, 32],
[ 7, 19],
[74, 85]],
[[76, 54],
[65, 65],
[96, 66]]])
In [50]: a[i[:,None],:,j].swapaxes(1,2)
Out[50]:
array([[[15, 32],
[ 7, 19],
[74, 85]],
[[76, 54],
[65, 65],
[96, 66]]])
```

**Assigning values with the indexing**

For approach #1, it's just straight-forward -

```
a[np.ix_(i,np.arange(a.shape[1]),j)] = b
```

For approach #2, if `b`

is a scalar, it should be straight-forward too -

```
a[i[:,None],:,j] = b
```

For approach #2 again, if you are assigning to a ndarray `b`

of shape `(20,100,30)`

, we need to swap axes of `b`

before assigning, like so -

```
a[i[:,None],:,j] = np.swapaxes(b,1,2)
```