python - 如何仅使用数据模型 dunder 方法执行(高级)索引就地(增强)分配给 numpy 向量？

Question

我正在使用 python 实现一个狭窄且有限的脚本 DSL，我希望能够在功能上执行以下等效操作:

import numpy as np
a = np.arange(10)
a[ a > 5 ] += 42

=> array([ 0,  1,  2,  3,  4,  5, 48, 49, 50, 51])

上面的代码可以正常工作。如果我开始扩展上面的代码，我会得到以下第一层内部结构:

a[a>5].__iadd__(42)

这也按预期工作。但是，我找不到允许我在数组本身而不是数组副本上操作 __iadd__ 的索引器方法。因此，不出所料，以下代码无法满足我的要求:

import numpy as np
a = np.arange(10)
a.__getitem__(a>5).__iadd__(42)

=> array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

仅当我这样做时:

a.__setitem__(a>5, a.__getitem__(a>5).__iadd__(42))

我似乎得到了我正在寻找的行为，但在这一点上，这根本不再是一个合适的就地赋值运算符，更重要的是，我索引了两次(一次用于读取，一次用于写)。

Numpy 的索引 page似乎暗示高级索引(即下标列表是 ndarray 的索引)总是 返回一个副本。这是否真的意味着 a[a>5].__iadd__(42) 实际上总是使用回退方法实现的？有没有我遗漏的东西，或者这根本不可能，或者至少没有解释器魔法是不可能的？

---

编辑:

因此根据@donkopotamus 的回答，数据模型不允许我们一次完成此操作。 这回答了问题。

但是，numpy 作为一个矢量化库，索引绝对不能承受非矢量化和多次执行。

这是一个“证明”:

import cython
import numpy as np

@cython.locals(arr="float[:]",
               mask="bint[:]",
               val=float,
               i=int)
@cython.boundscheck(False)
def func(arr,mask,val):
       for i in range(len(mask)):
               if mask[i]:
                        arr[i] += val

此代码在编译和计时时，比 numpy 慢:

a = np.arange(1e6)

%%timeit
a[a%3==0] += 42

=> 40.5 ms ± 376 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

a = np.arange(1e6)

%%timeit
func(a, (a%3==0), 42)

=> 116 ms ± 2.76 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

因此 REPL 解释语句比 3 行 cython 函数运行得更快，后者几乎可以像 CPU 允许的那样快地撕裂内存 View 。

在这个阶段，一切都不再有意义了。我知道 numpy 是手工制作的，用于优化矢量化操作，但我不了解它如何以有意义的方式与 python 解释器集成。它是否缓存 BINARY_SUBSCR/STORE_SUBSCR 对？

@donkopotamus 请注意，虽然索引操作不是计算 两次，但在 python 代码中，它是 解释 两次，因为对读取，然后对写入执行整个第二次扫描和掩码。在上面的 cython 代码中，该操作只发生一次读取和写入)。

欢迎任何见解。

Answer 1

您面临的问题不是特定于 numpy 本身，也不是 numpy 中的高级索引，也不是它是否创建副本。相反，它完全是由围绕是否:

• 索引保证返回容器“内部”的值(实际上不是)；以及是否

• 就地添加保证返回原始值的修改版本(不是)

考虑表达式:

x[a] += 100

其中 x 是一个列表 say。 x[a] 的结果是一个值 y，它不知道它恰好包含在列表中，表达式 y += 100 不能保证修改 y 的原始值...因此我们永远不能确保 x.__getitem__(x).__iadd__(100)< 形式的表达式 影响原始的 x。

因此表达式 x[a] += 100 必须由编译器使用以下步骤计算:

1. y = x[a]
2. y += 100
3. x[a] = y

或者在高级索引的情况下，我们可以期望 a[ a > 5 ] += 42 实现为:

1. b = a > 5
2. c = a[b]
3. c += 42
4. a[b] = c

这可以通过查看示例函数的反汇编来证明

def f(a):
    a[a > 5] += 42

然后

>>> dis.dis(f)
 0 LOAD_FAST                0 (a)
 2 LOAD_FAST                0 (a)
 4 LOAD_CONST               1 (5)
 6 COMPARE_OP               4 (>)  # 1. b = a > 5
 8 DUP_TOP_TWO
10 BINARY_SUBSCR                   # 2. c = a[b]
12 LOAD_CONST               2 (42)
14 INPLACE_ADD                     # 3. c += 42
16 ROT_THREE
18 STORE_SUBSCR                    # 4. a[b] = c
20 LOAD_CONST               0 (None)
22 RETURN_VALUE

请注意，在此实现中，索引 a > 5 不会执行两次。但是，如果您必须作为一组链式方法来实现，那么您别无选择，只能按照您的建议实现。