Generowanie n-wymiarowej tablicy współrzędnych w numpy

Question

Załóżmy, że mam funkcję f, która może przyjmować współrzędne jako parametr i zwraca liczbę całkowitą (f (x) w tym przypadku). Współrzędne mogą być wielowymiarowe i mają postać listy. Moim celem jest wypełnienie tablicy numpy wszystkimi wartościami pomiędzy dwiema współrzędnymi. Próbowałem utworzyć listę wszystkich możliwych wskaźników i użyć jej jako danych wejściowych dla funkcji wektorowej.

Oto mój kod dla 2 współrzędnych przestrzennych:

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])) 

ta zwróci listę indeksu jak oczekiwano:

[[ 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]] 

I tu jest moja próba dla n wymiarach:

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) 

Chociaż działa to dość wolno i wymaga dużej ilości pamięci. Czy można to napisać bardziej efektywnie? Mógłbym myśleć o używaniu numpy.meshgrid, ale nie wiem, jak bym go użył.

Answer 1

Rzeczywiście np.meshgrid byłoby jednym ze sposobów, aby to zrobić z jakimś stacking, jak pokazano poniżej -

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 

run Próbka

1) 2D koperty:

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 Przypadek:

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

Inną opcją jest użycie product od itertools, to działa także wtedy, gdy rogi są wyższe niż 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)]