numpy 切片赋值

Question

我一直在寻找一种高效又优雅的方法来解决这个问题。希望下面的例子能说明我的想法。

我们有一个像这样的 np.array：

omega = np.array([1.03415121504, 1.29595060284, 1.55774999064, 1.81954937844,
                  ...
                  2.08134876623, 2.37359445321, -2.11179506541, -1.84999567761])

现在我想对它进行一些操作，比如：

omega[omega < 0.0] = omega + 2 * np.pi
omega[omega >= 2 * np.pi] = omega - 2 * np.pi

第二个操作可能会覆盖第一个操作计算出的值，这样就会产生交集。我发现了 np.piecewise，但它并不支持这样的行为。

我该如何高效地实现这个呢？

目前的做法是这样的（但效率很低且不优雅）：

tmp = []
for o in omega:
    if o < 0.0:
        tmp.append(o + 2 * np.pi)
    elif o >= (2 * np.pi):
        tmp.append(o - 2 * np.pi)
    else:
        tmp.append(o)
omega = np.array(tmp)

所以有没有人尝试过使用 numpy 的 nditer 来实现这样的目的？（特别是关于性能和效率方面）

Answer 1

我建议你在索引上进行计算：

i = omega < 0.0
omega[i] += 2*np.pi
i = (~i) & (omega >= 2 * np.pi)
omega[i] -= 2*np.pi

第三行的位运算确保没有索引被重复使用。根据你给的例子，Sven Marnach 提出的取模方法更有效。不过，如果你有更复杂的使用场景，可能需要更新一下你的问题。

Answer 2

第二个语句可能会覆盖第一个语句计算出来的值。

这根本不可能发生。如果之前的值小于零，加上2*pi后也永远不会变成大于或等于2*pi。

无论如何，可能有更简单的方法来实现你想要的效果。

omega %= 2 * np.pi

Answer 3

你可以试试用 np.select 这个方法：

http://docs.scipy.org/doc/numpy/reference/generated/numpy.select.html

condlist = [omega < 0.0, omega >= 2.0*np.pi]
choicelist = [omega + 2.0*np.pi, omega - 2.0*np.pi]
omega = np.select(condlist,choicelist,default=omega)