如何在不实际索引数组的情况下获得已知形状的数组的索引元素的数量？

问题描述

我有一个索引IDX(它可以是索引列表、布尔掩码、切片元组等)。索引某个已知形状的抽象数字数组shape(可能很大)。

我知道我可以创建一个伪数组，对其进行索引并对元素进行计数：

A = np.zeros(shape)
print(A[IDX].size)

有没有什么合理的方法可以在不创建任何(可能很大的)数组的情况下获取索引元素的数量？

我需要列出3D空间中某些点的函数列表。这些点是以X、Y、Z列表形式给出的矩形网格的子集，并且IDX正在索引其笛卡尔乘积：

XX, YY, ZZ = [A[IDX] for A in np.meshgrid(X, Y, Z)]

这些函数接受X、Y、Z参数(并返回需要编制索引的笛卡尔乘积的值)或XX、YY、ZZ。 在创建XX、YY和ZZ数组时，无论是否使用它们，然后为函数值分配一个数组：

self.TAB = np.full((len(functions), XX.size),
                   np.nan)

但我只想在必要时创建XX、YY和ZZ。我还想把TAB分配和填充它的行分开，所以我需要提前知道列数。

---

解决方案

只是为了好玩，让我们看看我们是否可以在这里做出一个可以通过的近似。您的输入可以是以下任一项：

• 切片
• 数组(含标量)
  • 整数数组的索引功能很出色
  • 布尔数组进行掩码
• 元组

如果输入从一开始就不是显式的元组，那么就让它显式地成为元组。现在，您可以迭代元组并将其与形状进行匹配。您不能将它们压缩在一起，因为布尔数组会消耗形状的多个元素，而尾轴是批量包含的。

应该是这样的：

def pint(x):
    """ Mimic numpy errors """
    if isinstance(x, bool):
        raise TypeError('an integer is required')
    try:
        y = int(x)
    except TypeError:
        raise TypeError('an integer is required')
    else:
        if y < 0:
            raise ValueError('negative dimensions are not allowed')
    return y


def estimate_size(shape, index):
    # Ensure input is a tuple
    if not isinstance(index, tuple):
        index = (index,)

    # Clean out Nones: they don't change size
    index = tuple(i for i in index if i is not None)

    # Check shape shape and type
    try:
        shape = tuple(shape)
    except TypeError:
        shape = (shape,)
    shape = tuple(pint(s) for s in shape)

    size = 1

    # Check for scalars
    if not shape:
        if index:
            raise IndexError('too many indices for array')
        return size

    # Process index dimensions
    # you could probably use iter(shape) instead of shape[s]
    s = 0

    # fancy indices need to be gathered together and processed as one
    fancy = []

    def get(n):
        nonlocal s
        s += n
        if s > len(shape):
            raise IndexError('too many indices for array')
        return shape[s - n:s]

    for ind in index:
        if isinstance(ind, slice):
            ax, = get(1)
            size *= len(range(*ind.indices(ax)))
        else:
            ind = np.array(ind, ndmin=1, subok=True, copy=False)
            if ind.dtype == np.bool_:
                # Boolean masking
                ax = get(ind.ndim)
                if ind.shape != ax:
                    k = np.not_equal(ind.shape, ax).argmax()
                    IndexError(f'IndexError: boolean index did not match indexed array along dimension {s - n.ndim + k}; dimension is {shape[s - n.ndim + k]} but corresponding boolean dimension is {ind.shape[k]}')
                size *= np.count_nonzero(ind)
            elif np.issubdtype(ind.dtype, np.integer):
                # Fancy indexing
                ax, = get(1)
                if ind.min() < -ax or ind.max() >= ax:
                    k = ind.min() if ind.min() < -ax else ind.max()
                    raise IndexError(f'index {k} is out of bounds for axis {s} with size {ax}')
                fancy.append(ind)
            else:
                raise IndexError('arrays used as indices must be of integer (or boolean) type')

    # Add in trailing dimensions
    size *= np.prod(shape[s:])

    # Add fancy indices
    if fancy:
        size *= np.broadcast(*fancy).size

    return size

这只是一个近似值。您将需要在API更改时随时更改它，并且它已经具有一些不完整的功能。测试、修复和扩展留给读者练习。