Python 以numpy生成一个n维坐标数组

假设我有一个函数f，它可以将坐标作为参数，并返回一个整数（在本例中为f（x））。坐标可以是多维的，并且是列表的形式。我的目标是用两个坐标之间的所有值填充numpy数组。我试图列出所有可能的索引，并将其用作矢量化函数的输入

这是我的二维坐标代码：

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

下面是我对n维的尝试：

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，但我不知道如何使用它。

确实可以通过一些

堆叠来实现，如下所示-

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

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

2） 
3D
案例：

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

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

另一个选项是使用
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)]

呃，我已经在使用
itertools
中的
product
。它隐藏在第二个函数的返回中：D。但是感谢更紧凑的for循环！您可以使用
range（*r）
更改
范围（r[0]，r[1]）
。此外，您正在使用一个列表ComperExtension，它可以立即解包。您可以使用gen exp（即
*（range（*r）…）
将
[]
替换为
（）
），并避免创建临时列表。@Bakuriu工作得很好，感谢您指出这一点。我相信这是一个更有效、更简洁的解决方案。泰，这是解决这个问题的一个如此优雅、简单的解决方案。我要补充一点：如果已经有一个向量列表，只需将该列表替换为I.e.
it.product（*my_vector_list）
，并放弃偶数间距生成器/列表。这推广了非一致向量空间的解。