Search code examples
pythonfor-loopfor-comprehension

applying function to array


Python newbie here, I'm struggling to find the right idiom for applying a function over an array.

I want to compute an associated Legendre polynomial from -1 to 1,

import scipy.special as sc
m = 0
n = 2
z = 0.5
sc.lpmn(m, n, z)

That's all good, but the function isn't vectorised over z. How do I apply it (efficiently) to an array of values? Here's my attempt at using a comprehension, assuming I need to loop one way or another

import numpy as np
z = np.linspace(-1,1,20)
result = [sc.lpmn(0, 2, z[i])[0][0] for i in enumerate(z)]

Solution

  • It's simpler than that, and I don't think there's a reason to use enumerate. The following code should suffice:

    import scipy.special as sc
    z = np.linspace(-1,1,20)
    result = [sc.lpmn(0, 2, i)[0][0] for i in z]
    

    The output is:

    [array([ 1., -1.,  1.]),
     array([ 1.        , -0.89473684,  0.70083102]),
     array([ 1.        , -0.78947368,  0.43490305]),
     array([ 1.        , -0.68421053,  0.20221607]),
     array([ 1.        , -0.57894737,  0.00277008]),
     array([ 1.        , -0.47368421, -0.1634349 ]),
     array([ 1.        , -0.36842105, -0.29639889]),
     array([ 1.        , -0.26315789, -0.39612188]),
     array([ 1.        , -0.15789474, -0.46260388]),
     array([ 1.        , -0.05263158, -0.49584488]),
     array([ 1.        ,  0.05263158, -0.49584488]),
     array([ 1.        ,  0.15789474, -0.46260388]),
     array([ 1.        ,  0.26315789, -0.39612188]),
     array([ 1.        ,  0.36842105, -0.29639889]),
     array([ 1.        ,  0.47368421, -0.1634349 ]),
     array([1.        , 0.57894737, 0.00277008]),
     array([1.        , 0.68421053, 0.20221607]),
     array([1.        , 0.78947368, 0.43490305]),
     array([1.        , 0.89473684, 0.70083102]),
     array([1., 1., 1.])]