How can I speed up this Python function?

user_12848149 0
Viewed 92

Is there any way to vectorize this expression?

import numpy as np
def phase(f, dt=1):
    c = [f[0]] + [fi*dt for fi in f[1:]]
    
    s=0
    x = []
    for ci in c:
        s = (s+ci)%(2*np.pi)
        x.append(s)
    return x

f is an array of arbitrary length, dt is a constant. The code does:

  1. multiply the whole vector f but the first instance by dt.
  2. return the partial sum i.e. output[i] = sum(c[0:i]) modulo 2pi.

Any idea on how to speed it up? Thank you!

vectorization python loops performance
2 Answers

You can use Numpy's cumsum() function to do this in a vectorized way. This should be quite a bit faster:

import numpy as np

def phase(f, dt=1):
    pi_2 = 2 * np.pi
    c = np.concatenate((f[:1], f[1:] * dt))
    
    return np.cumsum(c) % pi_2

Timings:

n = np.arange(5000)

%timeit original_phase(n, 3)
# 3.2 ms ± 120 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit phase(n, 3)
# 272 µs ± 12.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

np.testing.assert_allclose(original_phase(n, 3), phase(n, 3))
user_38746230

Only one for loop is needed.

import numpy as np


def phase(f, dt=1):
    s = 0
    x = []
    if f:
        s = (s + f[0]) % (2 * np.pi)
        x.append(s)
    
        for fi in f[1:]:
            s = (s + fi * dt) % (2 * np.pi)
            x.append(s)
    return x
user_161821500
Related