编写访问参数维度numpy数组的方法

Question

我正在使用 python 2.7 和 numpy 1.9。
我有 3 种方法可以将转换应用于几个 numpy 数组。

def sum_arrays2(a, b):
    c = np.zeros(a.shape)
    c[:, 0:-1] = (a[:, 1:] + b[:, 0:-1]) ** 2
    return a[0:-1, 1:] + c[1:, 0:-1]

def sum_arrays3(a, b):
    c = np.zeros(a.shape)
    c[:, :, 0:-1] = (a[:, :, 1:] + b[:, :, 0:-1]) ** 2
    return a[0:-1, :, 1:] + c[1:, :, 0:-1]

def sum_arrays4(a, b):
    c = np.zeros(a.shape)
    c[:, :, :, 0:-1] = (a[:, :, :, 1:] + b[:, :, :, 0:-1]) ** 2
    return a[0:-1, :, :, 1:] + c[1:, :, :, 0:-1]

如您所见，它们非常相似。唯一的区别是所需的输入数组大小。
根据我的数据大小，我必须调用第一个、第二个或第三个。

实际上我必须这样做：

if a.ndims == 2:
  result = sum_arrays2(a, b)
elif a.ndims == 3:
  result = sum_arrays3(a, b)
elif a.ndims == 4:
  result = sum_arrays4(a, b)

如何制定更通用的方法来计算 n 维输入？

我找到的唯一解决方案是这样的：

def n_size_sum_arrays(a, b):
    c = np.zeros(a.shape)
    c[(Ellipsis, np.r_[0:c.shape[-1]-1])] = (a[(Ellipsis, np.r_[0:a.shape[-1]])] + b[(Ellipsis, np.r_[0:b.shape[-1]])]) ** 2
    return a[(r_[0:a.shape[0]-1], Ellipsis, np.r_[1:a.shape[-1]])] + c[(np.r_[1:c.shape[0]], Ellipsis, np.r_[0:c.shape[-1]-1])]

但它绝对不清楚，我不确定它是否正确。

有更好的方法吗？

Answer 1

您可以执行以下操作：

def sum_arrays(a, b):
    c = np.zeros(a.shape)
    c[..., :-1] = (a[..., 1:] + b[..., :-1]) ** 2
    return a[:-1, ..., 1:] + c[1:, ..., :-1]