Как я могу отключить трансляцию с изменяющимся размером при настройке на кусочек массива?

Question

Пусть 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.Я не против, если он изменит форму, например, если он превратится из массива формы 1х2 в массив формы 2х1х1.

Как я могу добиться этого?Путь, который я сейчас пытаюсь - переопределить __setitem__ в NoBroadcastArray, чтобы размер среза соответствовал размеру элемента, который нужно установить.Это оказывается сложно, поэтому мне интересно, если у кого-нибудь есть идея получше, которая, возможно, использует __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

Это только проверяет, что размер данного значения соответствуетиндекс, но он не изменяет значение, так что он все еще работает как обычно с NumPy (т. е. могут быть добавлены дополнительные внешние измерения).