Создайте n-мерный массив координат в numpy

Question

Предположим, что у меня есть функция f, которая может принимать координаты как параметр и возвращает целое число (f (x) в этом случае). Координаты могут быть многомерными и представлены в виде списка. Моя цель - заполнить массив numpy всеми значениями между двумя координатами. Я попытался составить список всех возможных индексов и использовать его в качестве ввода для векторизованной функции.

Вот мой код для 2-х мерных координат:

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

Это возвращает список индексов, как ожидалось:

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

А вот моя попытка для п измерений:

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) 

Хотя это работает, это довольно медленно и требует огромного объема памяти. Можно ли написать это более эффективным способом? Я мог подумать об использовании numpy.meshgrid, но я не знаю, как я буду использовать его.

Answer 1

Действительно np.meshgrid бы один из способов сделать это с некоторым stacking, как показано на рисунке ниже -

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 

Пример запуска

1) 2D Корпус:

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 Корпус:

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

Другой вариант заключается в использовании product из itertools, это также работает, если углы выше 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)]