假设我有一个函数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,但我不知道该如何使用它。
itertools
中的product
了。它被隐藏在第二个函数的返回值中:D。但感谢更紧凑的循环! - Gnarflordrange(*r)
替换range(r[0], r[1])
。同时,你正在使用一个立即解包的列表推导式。你可以使用生成器表达式代替(即将[]
替换为()
),避免创建临时列表。 - Bakuriuit.product(*my_vector_list)
,并且放弃均匀间距的生成器/列表。这可以推广到非均匀间距向量的情况。 - aeolus