使用相同的功能或操作来改变3D和4D阵列

Question

我想要改变一些3D和4D阵列。索引数组如下：

• 3D: array[time][x][y]
• 4D: array[time][z][x][y]

我想做多种类型的数据变异，并且已经编写了很多for-in循环。数据突变的例子：

1. 将0.值替换为NaN
2. 创建累积数组（时间= 0，x = x0，y = y0 - ＆gt;时间= 1，x = x0 + x1，y = y0 + y1）等。
3. 使用2个数组并将它们组合在一起（sqrt(xa * xa + xb * xb))）

Numpy可以用来做（1）使用：

for i in range(len(data)): 
    if len(data[i].shape) == 3:
        for z in range(len(data[i]):
            data[i][z][data[i][z] == 0.] = np.NaN
    else:
        data[i][data[i] == 0.] = np.NaN

但是每次遇到这样的问题时我都要编写for循环，并且必须将数据突变写入2次。一次用于4D阵列，一次用于3D阵列。 For循环是一个需要克服的功能，不必编写if else语句，并且几次写入变异几乎是必须的。如果我以某种方式或其他人更改第一部分但忘记更改第二部分，则代码会被窃听。

有没有办法例如枚举[z][x][y]，但如果数组只是[x][y]，那么只需执行一次，就好像有一个像[0][x][y]这样的z-index？