这是我的代码,使用dim = 3,但是我希望它可以在任何维度上工作,而无需手动编辑代码。
我希望能够最终在3到20之间更改尺寸,而无需手动添加循环。
我当时在看itertools,但不知道如何从itertools.product()
创建的元组中选择正确的值以平方并加到if语句中。
arrayshape = (width * 2 + 1,) * dim
funcspace = np.zeros(shape=arrayshape, dtype='b')
x1 = list(range(-int(width), int(width + 1)))
x2 = x1
x3 = x1
for i in range(len(x1)):
for j in range(len(x2)):
for k in range(len(x3)):
if round(np.sqrt(x1[i] ** 2 + x2[j] ** 2 + x3[k] ** 2)) in ranges:
funcspace[i][j][k] = 1
答案 0 :(得分:4)
您可以在product
个向量上使用enumerate
,这将产生值和索引:
for ((i,v1),(j,v2),(k,v3)) in itertools.product(enumerate(x1),enumerate(x2),enumerate(x3)):
if round(np.sqrt(v1**2+v2**2+v3**2)) in ranges:
funcspace[i][j][k]=1
作为奖励,您摆脱了非Python的range(len())
结构。
当您拥有向量载体时,我已经做了一个更一般的案例。读起来有点难,因为解包不是在for
循环中完成的。
平方和是使用1个索引(值)上的sum
完成的,如果条件匹配,我们将循环运行直到找到“更深”列表以将值设置为1
。
for t in itertools.product(*(enumerate(x) for x in x_list)):
# compute the squared sum of values
sqsum = sum(v[1]**2 for v in t)
if round(sqsum) in ranges:
# traverse the dimensions except the last one
deeper_list = funcspace
for i in range(len(t)-1):
deeper_list = deeper_list[t[i][0]]
# set the flag using the last dimension list
deeper_list[t[-1][0]] = 1
如评论中所述,由于似乎重复了x1
,因此您可以将第一条语句替换为:
for t in itertools.product(enumerate(x1), repeat=dim):
另一条评论指出,由于funcspace
是一个小数ndarray
,我们可以通过传递索引列表来简化“设置为1”循环:
funcspace[[x[0] for x in t]] = 1