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

Question

我有一个索引 IDX（可能是索引列表、布尔掩码、切片元组等）索引一些已知形状的抽象 numpy 数组 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 分配与填充其行分开，因此我需要提前知道列数。

Answer 1

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

• 切片
• 类数组（包括标量）
  • 整数数组可以做漂亮的索引
  • 布尔数组进行屏蔽
• 元组

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

应该这样做：

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