获取 NumPy 数组中列数的函数，如果它是一维数组，则返回 1

Question

我已经定义了对3xN NumPy数组的操作，我想遍历数组的每一列。

我试过：

for i in range(nparray.shape[1]):

但是，如果nparray.ndim == 1，则失败。

有没有一种干净的方法来确定NumPy数组的列数，例如，1确定它是否是一个1D数组（就像 MATLAB 的size操作那样）？

否则，我已经实施了：

if nparray.ndim == 1:
    num_points = 1
else:
    num_points = nparray.shape[1]

for i in range(num_points):

Answer 1

如果你只是在寻找不那么冗长的东西，你可以这样做：

num_points = np.atleast_2d(nparray).shape[1]

当然，这将创建一个新的临时数组以使其成形，这有点愚蠢……但它会非常便宜，因为它只是相同内存的视图。

但是，我认为您的显式代码更具可读性，只是我可能会使用try：

try:
    num_points = nparray.shape[1]
except IndexError:
    num_points = 1

如果您重复执行此操作，无论您做什么，都应该将其包装在一个函数中。例如：

def num_points(arr, axis):
    try:
        return arr.shape[axis]
    except IndexError:
        return 1

那么你所要做的就是：

for i in range(num_points(nparray, 1)):

当然，这意味着您只需编辑一个地方就可以更改任何地方的内容，例如：

def num_points(arr, axis):
    return nparray[:,...,np.newaxis].shape[1]

Answer 2

如果你想保留一句台词，可以使用条件表达式：

for i in range(nparray.shape[1] if nparray.ndim > 1 else 1):
    pass