Get part of np array with parameters-CodePudding

I am using python and numpy. I am using n dimensional array. I want to select all elements with index like

arr[a,b,:,c]

but I want to be able to select slice position like parameter. For example if the parameter

#pos =2
arr[a,b,:,c]

#pos =1
arr[a,:,b,c]

CodePudding user response：

I would move the axis of interest (at pos) to the front with numpy.moveaxis(array,pos,0)[1] and then simply slice with [:,a,b,c].

There is also numpy.take[2], but in your case you would still need to loop over each dimension a,b,c, so I think moveaxis is more convenient. Maybe there is an even more direct way to do this.

CodePudding user response：

You could np.transpose the array arr based on your preferences before you try to slice it, since you move your axis of interest (i.e. the :) "to the back". This way, you can rearrange arr, s.t. you can always call arr[a,b,c]. Example with only a and b:

import numpy as np

a = 0
b = 2
target_axis = 1

# Generate some random data
arr = np.random.randint(10, size=[3, 3, 3], dtype=int)
print(arr)
#[[[0 8 2]
#  [3 9 4]
#  [0 3 6]]
#
# [[8 5 4]
#  [9 8 5]
#  [8 6 1]]
#
# [[2 2 5]
#  [5 3 3]
#  [9 1 8]]]

# Define transpose s.t. target_axis is the last axis
transposed_shape = np.arange(arr.ndim)
transposed_shape = np.delete(transposed_shape, target_axis)
transposed_shape = np.append(transposed_shape, target_axis)
print(transposed_shape)
#[0 2 1] 
# Caution! These 0 and 2 above do not come from a or b.
# Instead they are the indices of the axes.

# Transpose arr
arr_T = np.transpose(arr, transposed_shape)
print(arr_T)
#[[[0 3 0]
#  [8 9 3]
#  [2 4 6]]
#
# [[8 9 8]
#  [5 8 6]
#  [4 5 1]]
#
# [[2 5 9]
#  [2 3 1]
#  [5 3 8]]]

print(arr_T[a,b])
#[2 4 6]

CodePudding user response：

The idea of moving the slicing axis to one end is a good one. Various numpy functions use that idea.

In [171]: arr = np.ones((2,3,4,5),int)
In [172]: arr[0,0,:,0].shape
Out[172]: (4,)
In [173]: arr[0,:,0,0].shape
Out[173]: (3,)

Another idea is to build a indexing tuple:

In [176]: idx = (0,0,slice(None),0)
In [177]: arr[idx].shape
Out[177]: (4,)
In [178]: idx = (0,slice(None),0,0)
In [179]: arr[idx].shape
Out[179]: (3,)

To do this programmatically it may be easier to start with a list or array that can be modified, and then convert it to a tuple for indexing. Details will vary depending on how you prefer to specify the axis and variables.

If any of a,b,c are arrays (or lists), you may get some shape surprises, since it's a case of mixing advanced and basic indexing. But as long as they are scalars, that's not an issue.