transform n-dimensional numpy array to a 2D-array based on a variable in Python

Question

I am trying to find the minimum value in an N-dimensional array spanned by (N-Parameters of varying values) and take out a 2-dimensional array spanned by 2 of the (N-Parameters) around the minimum value to make a contour plot. I can do this by hard coding the different cases, but it should preferably be done using a variable list of which axis should be extracted (contour_param).

Please see the code below for some clarification.

import numpy as np

np.random.seed(10) # seed random for reproducebility

#Example for a 3D input array (my_data)
param_sizes = [2, 3, 4]

#Generate a data_cube
my_data = np.random.normal(size=np.prod(param_sizes)).reshape(param_sizes)

#find minimum
min_pos = np.where(my_data == my_data.min())    


#what I want:
#define a parameter with the indexs of the axis to be used for the contour plot: e.i. : contour_param = [0, 1] 
               
#for contour_param = [0, 1] i would need the the 2D array:
result = my_data[:, :, min_pos[2][0]]

#for contour_param = [1, 2] i would need the the 2D array: 
result = my_data[min_pos[0][0], :, :]



#What I have tried is to convert min_pos to a list and change the entries to arrays:
contour_param = [0, 1] 
min_pos = list(np.where(my_data == my_data.min()))
min_pos[contour_param[0]] = np.arange(param_sizes[contour_param[0]])
min_pos[contour_param[1]] = np.arange(param_sizes[contour_param[1]])

result = my_data[min_pos] #This throws an error


#In an attempt to clarify - I have included a sample for a 4D array
#Example for a 4D array
param_sizes = [2, 3, 4, 3]


#Generate a data_cube
my_data = np.random.normal(size=np.prod(param_sizes)).reshape(param_sizes)

#find minimum
min_pos = np.where(my_data == my_data.min())    

                
#for contour_param = [0, 1] i would need the the 2D array:
result = my_data[:, :, min_pos[2][0], min_pos[3][0]]

#for contour_param = [1, 2] i would need the the 2D array 
result = my_data[min_pos[0][0], :, :, min_pos[3][0]]

Answer 1

Great Question...

You can make use of np.s_ for that, as you can build up your slicer with that.

For instance the function:

def build_slicer(contour_param,min_pos):
    
    assert len(contur_param) + 1 == min_pos.shape[0]

    output = []  # init a emtpy output list
    
    for main_index in range(min_pos.shape[0]):
        
        if main_index in contour_param:  
            output.append(np.s_[:])
             
        else:
            output.append(np.s_[min_pos[main_index][0]])
        
    return tuple(output)

would return:

import numpy as np
np.random.seed(10)

param_sizes = [2, 3, 4]
my_data = np.random.normal(size=np.prod(param_sizes)).reshape(param_sizes)
min_pos = np.where(my_data == my_data.min())    

contour_param = [0,2]

build_slicer(contour_param,min_pos)  

>>> (slice(None, None, None), 2, slice(None, None, None))

you can then use this to just slice your array

slice = build_slicer(contour_param,min_pos)  
my_data[slice]