CODEWITHSUNDEEP
pythonnumpy
Themerius
Themerius

Reputation: 1911

numpy slice assignment

I've searching a efficient and elegant way to do this. I'll hope the example will explain my concern.

We've got an np.array like this:

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

And now I want to manipulate it, like

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

The second statement may overwrite the computed values of the fist statement, and then there's an intersection. I found np.piecewise, but this doesn't provides such an behavior.

How can I achieve this (efficient)?

The corrent behavior is like that (but very unefficient/inelegant):

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)

Therefore someone made experiences with numpy's nditer for such purposes? (Especially about performance / efficency)

Upvotes: 1

Views: 1840

Answers (3)

David Zwicker
David Zwicker

Reputation: 24308

I would suggest doing the calculation on the indices:

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

The bitwise logical operation in the third line ensures that no indices are used twice. Judging from the example you gave, the modulo answer by Sven Marnach is more efficient, though. You should probably update your question, if you have a more complex use case.

Upvotes: 2

JoshAdel
JoshAdel

Reputation: 68682

You could try using 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)

Upvotes: 3

Sven Marnach
Sven Marnach

Reputation: 601859

The second statement may overwrite the computed values of the fist statement.

There is no way this can happen. If it was less than zero before, adding 2*pi will never make it greater than or equal to 2*pi

Anyway, an easier way to achieve what you want might be

omega %= 2 * np.pi

Upvotes: 2

Related Questions