python - 设置为 numpy 数组切片时,如何禁用改变大小的广播?
问题描述
让NoBroadcastArray
成为 的子类np.ndarray
。如果x
是一个实例NoBroadcastArray
并且arr
是一个np.ndarray
,那么我想要
x[slice] = arr
当且仅当arr.size
匹配切片的大小时才能成功。
x[1] = 1 # should succeed
x[1:2] = 1 # should fail - scalar doesn't have size 2
x[1:2] = [1,2] # should succeed
x[1:2] = np.array([[1,2]]) # should succeed - shapes don't match but sizes do.
x[1:2, 3:4] = np.array([1,2]) # should fail - 1x2 array doesn't have same size as 2x2 array
换句话说,只有当 RHS 不必更改大小以适应 LHS 切片时,分配才应该成功。我不介意它是否会改变形状,例如,如果它从形状为 1x2 的数组变为形状为 2x1x1 的数组。
我怎样才能实现这一目标?我现在尝试的路径是覆盖 __setitem__NoBroadcastArray
以匹配切片的大小与要设置的项目的大小。事实证明这很棘手,所以我想知道是否有人有更好的想法,可能使用 __array_wrap__ 或 __array_finalize__。
解决方案
这是我想出的实现:
import numpy as np
class NoBroadcastArray(np.ndarray):
def __new__(cls, input_array):
return np.asarray(input_array).view(cls)
def __setitem__(self, args, value):
value = np.asarray(value, dtype=self.dtype)
expected_size = self._compute_expected_size(args)
if expected_size != value.size:
raise ValueError(("assigned value size {} does not match expected size {} "
"in non-broadcasting assignment".format(value.size, expected_size)))
return super(NoBroadcastArray, self).__setitem__(args, value)
def _compute_expected_size(self, args):
if not isinstance(args, tuple):
args = (args,)
# Iterate through indexing arguments
arr_dim = 0
ellipsis_dim = len(args)
i_arg = 0
size = 1
adv_idx_shapes = []
for i_arg, arg in enumerate(args):
if isinstance(arg, slice):
size *= self._compute_slice_size(arg, arr_dim)
arr_dim += 1
elif arg is Ellipsis:
ellipsis_dim = arr_dim
break
elif arg is np.newaxis:
pass
else:
adv_idx_shapes.append(np.shape(arg))
arr_dim += 1
# Go backwards from end after ellipsis if necessary
arr_dim = -1
for arg in args[:i_arg:-1]:
if isinstance(arg, slice):
size *= self._compute_slice_size(arg, arr_dim)
arr_dim -= 1
elif arg is Ellipsis:
raise IndexError("an index can only have a single ellipsis ('...')")
elif arg is np.newaxis:
pass
else:
adv_idx_shapes.append(np.shape(arg))
arr_dim -= 1
# Include dimensions under ellipsis
ellipsis_end_dim = arr_dim + self.ndim + 1
if ellipsis_dim > ellipsis_end_dim:
raise IndexError("too many indices for array")
for i_dim in range(ellipsis_dim, ellipsis_end_dim):
size *= self.shape[i_dim]
size *= NoBroadcastArray._advanced_index_size(adv_idx_shapes)
return size
def _compute_slice_size(self, slice, axis):
if axis >= self.ndim or axis < -self.ndim:
raise IndexError("too many indices for array")
size = self.shape[axis]
start = slice.start
stop = slice.stop
step = slice.step if slice.step is not None else 1
if step == 0:
raise ValueError("slice step cannot be zero")
if start is not None:
start = start if start >= 0 else start + size
start = min(max(start, 0), size - 1)
else:
start = 0 if step > 0 else size - 1
if stop is not None:
stop = stop if stop >= 0 else stop + size
stop = min(max(stop, 0), size)
else:
stop = size if step > 0 else -1
slice_size = stop - start
if step < 0:
slice_size = -slice_size
step = -step
slice_size = ((slice_size - 1) // step + 1 if slice_size > 0 else 0)
return slice_size
@staticmethod
def _advanced_index_size(shapes):
size = 1
if not shapes:
return size
dims = max(len(s) for s in shapes)
for dim_sizes in zip(*(s[::-1] + (1,) * (dims - len(s)) for s in shapes)):
d = 1
for dim_size in dim_sizes:
if dim_size != 1:
if d != 1 and dim_size != d:
raise IndexError("shape mismatch: indexing arrays could not be "
"broadcast together with shapes " + " ".join(map(str, shapes)))
d = dim_size
size *= d
return size
你会像这样使用它:
import numpy as np
a = NoBroadcastArray(np.arange(24).reshape(4, 3, 2, 1))
a[:] = 1
# ValueError: assigned value size 1 does not match expected size 24 in non-broadcasting assignment
a[:, ..., [0, 1], :] = 1
# ValueError: assigned value size 1 does not match expected size 16 in non-broadcasting assignment
a[[[0, 1], [2, 3]], :, [1, 0]] = 1
# ValueError: assigned value size 1 does not match expected size 12 in non-broadcasting assignment
这仅检查给定值的大小是否与索引匹配,但它不会对值进行任何重塑,因此仍然可以像往常一样使用 NumPy(即可以添加额外的外部尺寸)。
推荐阅读
- c# - 恢复sql数据库C#时找不到方法
- odoo-11 - 我在哪里可以下载适用于企业的 odoo 11 插件?
- python - TypeError: 无法连接类型为 ' 的对象
'; 只有 Series 和 DataFrame obj 是有效的 - regex - 在 Bash 的正则表达式中添加下划线
- sql - SQL将一张表的多条记录合并成另一张表的多列
- javascript - 如何在点击时使弹出 div 隐藏/显示
- import - 从外部 wiki 读取和导入数据
- php - PHP中的单个到多个数组
- java - SWT 浏览器 - 谷歌账户登录
- crystal-reports - 是否有任何 SQL 查询来订购 NVARCHR 文件作为那些(852-2010 ,14-2015 ,20-2011 ,