Search code examples
arrayspython-2.7numpyscipybessel-functions

Is there a more pythonic way to write


Learning to be pythonic in 2.7. Is there a way to avoid the explicit loop? answer = [5, 4, 4, 3, 3, 2]

import numpy as np
import scipy.special as spe

nmax = 5     # n = 0, 1 ...5
mmax = 7     # m = 1, 2 ...7
big = 15.
z = np.zeros((nmax+1, mmax))
for i in range(nmax+1):
    z[i] = spe.jn_zeros(i, mmax)
answer = [np.max(np.where(z[i]<big))+1 for i in range(nmax+1)]
print answer # list of the largest m for each n where the mth zero of Jn < big

Solution

  • What does "more Pythonic" really mean here. One of the core tenets of Python is readability, so if there is no real performance reason to get rid of the loops, just keep them.

    If you really wanted to see some other ways to do the same thing, then:

    z = np.zeros((nmax+1, mmax))
    for i in range(nmax+1):
        z[i] = spe.jn_zeros(i, mmax)
    

    could be replaced with:

    func = lambda i:spe.jn_zeros(i,mmax)
    np.vstack(np.vectorize(func, otypes=[np.ndarray])(np.arange(nmax+1)))
    

    which is slightly faster (1.35 ms vs. 1.77 ms) but probably less Pythonic and

    [np.max(np.where(z[i]<big))+1 for i in range(nmax+1)]
    

    could be replaced by

    np.cumsum(z < big,axis=1)[:,-1]
    

    which I would argue is more Pythonic (or numpythonic) and much faster (20 us vs. 212 us).