Genera un array n-dimensionale di coordinate in numpy

Question

Diciamo che ho una funzione f che può assumere le coordinate come parametro e restituisce un intero (f (x) in questo caso). Le coordinate possono essere multidimensionali e hanno la forma di una lista. Il mio obiettivo è riempire una matrice numpy con tutti i valori tra due coordinate. Ho cercato di creare un elenco di tutti gli indici possibili e usarlo come input per la funzione vettoriale.

Ecco il mio codice per 2 coordinate tridimensionali:

import itertools 
import numpy 


def index_array(lower_corner, upper_corner): 
    x_range = range(lower_corner[0], upper_corner[0]) 
    y_range = range(lower_corner[1], upper_corner[1]) 
    return numpy.array(list(itertools.product(x_range, y_range))) 


print(index_array([2, -2], [5, 3])) 

Questo restituirà la lista degli indici come previsto:

[[ 2 -2] 
[ 2 -1] 
[ 2 0] 
[ 2 1] 
[ 2 2] 
[ 3 -2] 
[ 3 -1] 
[ 3 0] 
[ 3 1] 
[ 3 2] 
[ 4 -2] 
[ 4 -1] 
[ 4 0] 
[ 4 1] 
[ 4 2]] 

E qui è il mio tentativo per n dimensioni:

import itertools 
import numpy 


def f(x): 
    # dummy function 
    return x + 5 


def index_array(lower_corner, upper_corner): 
    # returns all indices between two n-dimensional points 
    range_list = [] 
    for n in range(len(lower_corner)): 
     range_list.append(range(lower_corner[n], upper_corner[n])) 
    return numpy.array(list(itertools.product(*range_list))) 


lower_corner = numpy.array([2, -2]) 
upper_corner = numpy.array([5, 3]) 
indices = index_array(lower_corner, upper_corner) 
vect_func = numpy.vectorize(f) 
results = vect_func(indices) 
print(results) 

Mentre funziona è piuttosto lento e necessita di enormi quantità di memoria. È possibile scrivere questo in modo più efficiente? Potrei pensare a usare numpy.meshgrid ma non so come lo userei.

Answer 1

Infatti np.meshgrid sarebbe un modo per farlo con un po 'stacking, come mostrato di seguito -

def ndim_grid(start,stop): 
    # Set number of dimensions 
    ndims = len(start) 

    # List of ranges across all dimensions 
    L = [np.arange(start[i],stop[i]) for i in range(ndims)] 

    # Finally use meshgrid to form all combinations corresponding to all 
    # dimensions and stack them as M x ndims array 
    return np.hstack((np.meshgrid(*L))).swapaxes(0,1).reshape(ndims,-1).T 

corsa Esempio

1) 2D Caso:

In [97]: ndim_grid([2, -2],[5, 3]) 
Out[97]: 
array([[ 2, -2], 
     [ 2, -1], 
     [ 2, 0], 
     [ 2, 1], 
     [ 2, 2], 
     [ 3, -2], 
     [ 3, -1], 
     [ 3, 0], 
     [ 3, 1], 
     [ 3, 2], 
     [ 4, -2], 
     [ 4, -1], 
     [ 4, 0], 
     [ 4, 1], 
     [ 4, 2]]) 

2) 3D Caso:

In [98]: ndim_grid([2, -2, 4],[5, 3, 6]) 
Out[98]: 
array([[ 2, -2, 4], 
     [ 2, -2, 5], 
     [ 2, -1, 4], 
     [ 2, -1, 5], 
     [ 2, 0, 4], 
     [ 2, 0, 5], 
     [ 2, 1, 4], 
     [ 2, 1, 5], 
     [ 2, 2, 4], 
     [ 2, 2, 5], 
     [ 3, -2, 4], 
     [ 3, -2, 5], 
     [ 3, -1, 4], 
     [ 3, -1, 5], 
     [ 3, 0, 4], 
     [ 3, 0, 5], 
     [ 3, 1, 4], 
     [ 3, 1, 5], 
     [ 3, 2, 4], 
     [ 3, 2, 5], 
     [ 4, -2, 4], 
     [ 4, -2, 5], 
     [ 4, -1, 4], 
     [ 4, -1, 5], 
     [ 4, 0, 4], 
     [ 4, 0, 5], 
     [ 4, 1, 4], 
     [ 4, 1, 5], 
     [ 4, 2, 4], 
     [ 4, 2, 5]]) 

Answer 2

Un'altra opzione è quella di utilizzare il product da itertools, questo funziona anche se gli angoli sono più alti rispetto 2D:

import itertools as it 
lower_corner = [2, -2] 
upper_corner = [5, 3] 
[coord for coord in it.product(*[range(r[0], r[1]) for r in zip(lower_corner, upper_corner)])] 

[(2, -2), 
(2, -1), 
(2, 0), 
(2, 1), 
(2, 2), 
(3, -2), 
(3, -1), 
(3, 0), 
(3, 1), 
(3, 2), 
(4, -2), 
(4, -1), 
(4, 0), 
(4, 1), 
(4, 2)]