python - 可重用代码沿不同的数组维度进行迭代

Question

假设我有一个 N 维数组 my_array[D1][D2]...[DN]

对于特定应用，例如敏感性分析，我需要固定一个点 p=(d1, d2, ..., dN) 并一次沿每个维度进行迭代。结果行为是

for x1 in range(0, D1):
    do_something(my_array[x1][d2][d3]...[dN])
for x2 in range(0, D2):
    do_something(my_array[d1][x2][d3]...[dN])
.
.
.
for xN in range(0, DN):
    do_something(my_array[d1][d2][d3]...[xN])

正如你所看到的，这里有很多重复的代码。如何减少工作量并编写一些优雅的代码？

例如，有没有类似下面的代码生成方法？

for d in range(0, N):
    iterate along the (d+1)th dimension of my_array, denoting the element as x:
        do_something(x)

Answer 1

您可以使用 numpy.take 并执行如下操作。浏览文档以供引用。

https://docs.scipy.org/doc/numpy/reference/generated/numpy.take.html

N = len(my_array)
for i in range(N):
  n = len(my_array(i))
  indices = p
  indices[i] = x[i]
  for j in range(n):
     do_something(np.take(my_array,indices))