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