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

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

我有一个索引 IDX（可以是索引列表、布尔掩码、切片元组等），它为一些已知形状 shape（可能很大）的抽象 numpy 数组建立索引。

>

我知道我可以创建一个虚拟数组，索引它并计算元素数：

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、{{1} },YY。 在我创建 ZZ、XX 和 YY 数组时，无论是否使用它们，然后我为函数值分配一个数组：

ZZ

但我只想在必要时创建 self.TAB = np.full((len(functions),XX.size),np.nan) 、XX 和 YY。我还想将 ZZ 分配与填充行分开，因此我需要提前知道列数。

解决方法

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

• 切片
• 类数组（包括标量）
  • 整数数组做花哨的索引
  • 布尔数组做掩蔽
• 元组

如果输入不是明确的元组开始，则将其设为元组。现在您可以沿元组迭代并将其与形状匹配。您无法将它们完全压缩在一起，因为布尔数组占用了形状的多个元素，并且尾随轴被批发包含在内。

应该这样做：

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 更改时，您都需要更改它，并且它已经有一些不完整的功能。测试、修复和扩展留给读者作为练习。