python - 设置为 numpy 数组切片时如何禁用大小更改广播？

Question

让 NoBroadcastArray 成为 np.ndarray 的子类。如果 x 是 NoBroadcastArray 的一个实例并且 arr 是一个 np.ndarray，那么我想要

x[slice] = arr

当且仅当 arr.size 匹配 slice 的大小时才成功。

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 的数组。

我怎样才能实现这一目标？我现在尝试的方法是重写 NoBroadcastArray 中的 __setitem__ 以使切片的大小与要设置的项目的大小相匹配。这被证明是棘手的，所以我想知道是否有人有更好的主意可以使用 __array_wrap__ 或 __array_finalize__。

Answer 1

这是我想出的实现:

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

这只会检查给定值的大小是否与索引匹配，但它不会对值进行任何 reshape ，因此 NumPy 仍然像往常一样工作(即可以添加额外的外部维度)。