加速对给定数组的每个索引值应用转换

Question

我需要对给定numpy数组的所有索引值的转换结果应用一个函数。下面的代码是这样做的:

import numpy as np
from matplotlib.transforms import IdentityTransform
# some 2D array
a = np.empty((2,3))
# some affine transformation, identity is just an example here
trans = IdentityTransform()
# some function taking a 2D index and returning some value depending
# on that index, again just an example
def f(idx):
return (idx[0]+idx[1])/2
# apply f to the result of transforming each index of a
b=np.empty_like(a)
for idx in np.ndindex(a.shape):
b[idx] = f(trans.transform(idx))
print(b)

打印以下正确的结果:

[[0。0.5 - 1。］(0.5 - 1。1.5]]

现在的问题是，当a的形状变大时，比如2000x3000，代码太慢了。有办法加快速度吗?

我的想法是创建一个idx = [[0,0], [0,1], ..., [1,2]]的索引数组，然后用tmp = trans.transform(idx)这样的东西一次变换这个数组，最后用np.vectorize(f)(tmp)将f应用于每个元素。

这是一个合理的方法吗?如果是的话，这到底会是什么样子?如果没有，还有其他选择吗?

编辑:我设法得到在tmp通过以下代码:

tmp=trans.transform(np.asarray([idx for idx in np.ndindex(a.shape)]))

所以现在我有一个数组，包含对a的每个索引值进行仿射变换的结果。但这似乎占用了大量内存。

Answer 1

我会把我现在发现的答案贴出来。也许它对某人有用。

为了回答我问题的第一部分，我发现了一种快速有效的方法来创建转换索引值的结果，使用np.indices()的结果，然后对结果进行处理，直到它符合t.transform()的期望。

给定一个数组a = np.empty((2,3))，通过np.indices(a.shape)可以得到该数组的索引。这将返回两个2D数组(实际上，a的每个维度对应一个数组)。我不明白的是如何把这些结果变成transform()理解的东西。

这里的关键是将np.ravel()应用于每个数组的结果，np.indices()返回:

>>> a=np.empty((2,3))
>>> list(map(np.ravel, np.indices(a.shape)))
[array([0, 0, 0, 1, 1, 1]), array([0, 1, 2, 0, 1, 2])]

现在我有一个包含所有x和y索引的数组列表，它只需要与np.vstack()放在一起，然后转置以获得所有(x, y)索引的数组，这是transform()将接受的形式。

>>> l=list(map(np.ravel, np.indices(a.shape)))
>>> np.vstack(l).transpose()
array([[0, 0],
[0, 1],
[0, 2],
[1, 0],
[1, 1],
[1, 2]])

最后，对于任意仿射变换:

>>> from matplotlib.transforms import Affine2D
>>> t = Affine2D().translate(10, 20).scale(0.5)
>>> t.transform(np.vstack(l).transpose())
array([[ 5. , 10. ],
[ 5. , 10.5],
[ 5. , 11. ],
[ 5.5, 10. ],
[ 5.5, 10.5],
[ 5.5, 11. ]])

这是相当快的，即使对于较大的数组大小。如果形状变得足够大(比如20000x30000)，我就会耗尽内存，但对于形状10000x10000，它仍然是惊人的快。

>>> timeit.timeit("t.transform(np.vstack(list(map(np.ravel, np.indices(a.shape, dtype=np.uint16)))).transpose())",
...    "import numpy as np ; from matplotlib.transforms import Affine2D ; a = np.empty((20, 10)) ; t = Affine2D().translate(10, 20).scale(0.5)", number=10)
0.0003051299718208611
>>> timeit.timeit("t.transform(np.vstack(list(map(np.ravel, np.indices(a.shape, dtype=np.uint16)))).transpose())",
...    "import numpy as np ; from matplotlib.transforms import Affine2D ; a = np.empty((200, 100)) ; t = Affine2D().translate(10, 20).scale(0.5)", number=10)
0.0026413939776830375
>>> timeit.timeit("t.transform(np.vstack(list(map(np.ravel, np.indices(a.shape, dtype=np.uint16)))).transpose())",
...    "import numpy as np ; from matplotlib.transforms import Affine2D ; a = np.empty((2000, 1000)) ; t = Affine2D().translate(10, 20).scale(0.5)", number=10)
0.35055489401565865
>>> timeit.timeit("t.transform(np.vstack(list(map(np.ravel, np.indices(a.shape, dtype=np.uint16)))).transpose())",
...    "import numpy as np ; from matplotlib.transforms import Affine2D ; a = np.empty((20000, 10000)) ; t = Affine2D().translate(10, 20).scale(0.5)", number=10)
43.62860555597581

现在，对于第二部分，为了将函数应用于每个转换后的索引值，我现在使用以下代码，在我的情况下，这已经足够快了。

xxyy = t.transform(np.vstack(...).transpose())
np.fromiter((f(*xy) for xy in xxyy), dtype=np.short, count=len(xxyy))